Kandinsky 6 VAE

Kandinsky 6 使用因果 3D K-VAE 进行视频超分辨率，并结合 MMAudio 梅尔频谱 VAE 与独立的 BigVGAN MMAudioVocoder，以实现同步音频生成。

Kandinsky6SRVAE

Kandinsky6SRPipeline 所使用的因果 3D K-VAE。它能够在有界内存的分段中处理任意长度的视频，同时复现单次非分段处理的确切输出。

import torch
from diffusers import Kandinsky6SRVAE

vae = Kandinsky6SRVAE.from_pretrained(
    "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers", subfolder="vae", torch_dtype=torch.bfloat16
)

class diffusers.Kandinsky6SRVAE

< source >

( in_channels: int = 3out_channels: int = 3latent_channels: int = 64encoder_block_out_channels: tuple[int, ...] = (16, 128, 256, 512, 1024)decoder_block_out_channels: tuple[int, ...] = (16, 256, 512, 1024, 2048)layers_per_block: int = 2temporal_compression_ratio: int = 4temporal_compression_start_level: int = 1scaling_factor: float = 0.910344004631042 )

参数

in_channels (int，默认为 3) — 像素通道数。
out_channels (int，默认为 3) — 重建像素的通道数。
latent_channels (int，默认为 64) — 潜在通道的数量。
encoder_block_out_channels (tuple[int, ...]，默认为 (16, 128, 256, 512, 1024)) — 编码器每一层中残差块的输出宽度；除最后一层外，每一层在过渡到下一层时都会将空间尺寸减半，并将宽度加倍。
decoder_block_out_channels (tuple[int, ...]，默认为 (16, 256, 512, 1024, 2048)) — 解码器各层级中残差块的输出宽度。
layers_per_block (int，默认为 2) — 每个编码器层级的残差块数量；解码器在每个层级上多使用一个。
temporal_compression_ratio (int，默认为 4) — 时间压缩因子；其 log2 个连续层级也会压缩时间。
temporal_compression_start_level (int，默认为 1) — 执行时间压缩的起始层级。
scaling_factor (float，默认为 0.910344) — 在潜变量进入扩散 Transformer 之前对其应用的缩放系数。

Kandinsky6SRPipeline 用于对视频进行编码和解码的因果 3D K-VAE。

视频按 16 个像素帧（外加起始帧）的时间片段进行处理。因果卷积会在片段之间传递其填充状态，因此分段仅用于限制峰值内存，不会改变结果。

encode

< source >

( x: torch.Tensorreturn_dict: bool = True )

参数

x (torch.Tensor，形状为 (batch_size, channels, num_frames, height, width)) — 像素值位于 [-1, 1] 范围内的视频。num_frames 应为 1 + k * temporal_compression_ratio。
return_dict (bool，默认为 True) — 是否返回 AutoencoderKLOutput 而不是普通元组。

将视频编码为其潜在分布。

decode

< source >

( z: torch.Tensorreturn_dict: bool = True )

参数

z（形状为 (batch_size, latent_channels, num_latent_frames, height, width) 的 torch.Tensor）— 潜变量，已除以 scaling_factor。
return_dict (bool，默认为 True) — 是否返回 ~models.autoencoder_kl.DecoderOutput 而不是普通元组。

将潜在变量解码为视频。

forward

< source >

( sample: torch.Tensorsample_posterior: bool = Falsereturn_dict: bool = Truegenerator: torch.Generator | None = None ) → ~models.autoencoder_kl.DecoderOutput or tuple

参数

sample (torch.Tensor，形状为 (batch_size, channels, num_frames, height, width)) — 像素视频，取值范围为 [-1, 1]。num_frames 应为 1 + k * temporal_compression_ratio。
sample_posterior (bool，可选，默认为 False) — 是否从潜在后验分布中采样，而不是使用其众数。
return_dict (bool，可选，默认为 True) — 是否返回 ~models.autoencoder_kl.DecoderOutput 而不是普通元组。
generator (torch.Generator，可选) — 用于使采样过程确定性的 torch.Generator。

返回值

~models.autoencoder_kl.DecoderOutput 或 tuple

如果 return_dict 为 True，则返回 ~models.autoencoder_kl.DecoderOutput，否则返回普通的 tuple。其 sample 为重建后的视频。

MMAudioVAE

当 sample_audio=True 时，Kandinsky6TI2VAPipeline 所使用的 mel 频谱图 VAE。其 decode 输出为 mel 频谱图；将其通过 MMAudioVocoder 处理即可得到波形。

参考实现可在 hkchengrex/MMAudio 中找到（MIT 许可证）。

import torch
from diffusers import MMAudioVAE

audio_vae = MMAudioVAE.from_pretrained(
    "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers", subfolder="audio_vae", torch_dtype=torch.bfloat16
)

类 diffusers.MMAudioVAE

< source >

( mel_bins: int = 128latent_channels: int = 40hidden_channels: int = 512channel_multipliers: tuple[int, ...] = (1, 2, 4)layers_per_block: int = 2sample_rate: int = 44100n_fft: int = 2048hop_length: int = 512scaling_factor: float = 0.417 )

参数

mel_bins (int，默认为 128) — 梅尔频率带的数量。
latent_channels (int，默认为 40) — 潜在通道的数量。
hidden_channels (int，默认为 512) — 自编码器的基础宽度。
channel_multipliers (tuple[int, ...]，默认为 (1, 2, 4)) — 自编码器各层级的宽度乘数。
layers_per_block (int，默认为 2) — 每个编码器层级的残差块数量；解码器在每个层级上多使用一个。
sample_rate (int，默认为 44100) — 波形采样率。
n_fft (int，默认为 2048) — 梅尔前端（mel front end）的 FFT 大小。
hop_length (int，默认为 512) — Mel 前端的跳跃长度。必须与此 VAE 所配对的 MMAudioVocoder 的总上采样因子相匹配。
scaling_factor (float，默认为 0.417) — 在潜变量进入扩散 Transformer 之前对其应用的缩放系数。

Kandinsky6TI2VAPipeline 的音频 VAE：一种在 log-mel 频谱上保持幅度的自编码器（MMAudio，https://arxiv.org/abs/2412.15322）。

encode 将波形转换为潜在分布；decode 将潜在表示还原为梅尔频谱图，随后由 MMAudioVocoder 将其转换为波形。一个潜在帧覆盖 hop_length * 2 个采样点。

encode

< source >

( audio: torch.Tensorreturn_dict: bool = True )

参数

audio（形状为 (batch_size, num_samples) 的 torch.Tensor）— 采样率为 sample_rate 的 [-1, 1] 范围内的单声道波形。
return_dict (bool，默认为 True) — 是否返回 AutoencoderKLOutput 而不是普通元组。

将波形编码为其潜在分布。

decode

< source >

( z: torch.Tensorreturn_dict: bool = True )

参数

z（形状为 (batch_size, latent_channels, num_latent_frames) 的 torch.Tensor）— 潜变量，已除以 scaling_factor。
return_dict (bool，默认为 True) — 是否返回 ~models.autoencoder_kl.DecoderOutput 而不是普通元组。

将潜在变量解码为梅尔频谱图。

forward

< source >

( sample: torch.Tensorsample_posterior: bool = Falsereturn_dict: bool = Truegenerator: torch.Generator | None = None ) → ~models.autoencoder_kl.DecoderOutput or tuple

参数

sample（形状为 (batch_size, num_samples) 的 torch.Tensor）— 位于 [-1, 1] 区间、采样率为 sample_rate 的单声道波形，用于编码并重建为梅尔频谱图。
sample_posterior (bool，可选，默认为 False) — 是否从潜在后验分布中采样，而不是使用其众数。
return_dict (bool，可选，默认为 True) — 是否返回 ~models.autoencoder_kl.DecoderOutput 而不是普通元组。
generator (torch.Generator，可选) — 用于使采样过程确定性的 torch.Generator。

返回值

~models.autoencoder_kl.DecoderOutput 或 tuple

如果 return_dict 为 True，则返回 ~models.autoencoder_kl.DecoderOutput，否则返回普通的 tuple。其 sample 是重建的 mel 频谱图，形状为 (batch_size, mel_bins, num_mel_frames)。

MMAudioVocoder

改编自 BigVGAN-v2 声码器的 MMAudio 捆绑包，其本身源自 NVIDIA/BigVGAN（MIT 许可证），并采用了 alias-free-torch（Apache License 2.0）中的抗锯齿 Snake 激活函数。

class diffusers.MMAudioVocoder

< source >

( num_mels: int = 128upsample_initial_channel: int = 1536upsample_rates: tuple[int, ...] = (8, 4, 2, 2, 2, 2)upsample_kernel_sizes: tuple[int, ...] = (16, 8, 4, 4, 4, 4)resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11)resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ((1, 3, 5), (1, 3, 5), (1, 3, 5)) )

参数

num_mels (int，默认为 128) — 输入频谱图的 mel 频带数量。必须与配对的 MMAudioVAE 的 mel_bins 相匹配。
upsample_initial_channel (int，默认为 1536) — 第一层的宽度。
upsample_rates (tuple[int, ...]，默认为 (8, 4, 2, 2, 2, 2)) — 声码器各阶段的升采样因子。它们的乘积即为总升采样因子，必须与配对的 MMAudioVAE 的 hop_length 相匹配。
upsample_kernel_sizes (tuple[int, ...]，默认为 (16, 8, 4, 4, 4, 4)) — 声码器各阶段的转置卷积核大小。
resblock_kernel_sizes (tuple[int, ...]，默认为 (3, 7, 11)) — 残差块的卷积核大小。
resblock_dilation_sizes (tuple[tuple[int, ...], ...]，默认为 ((1, 3, 5), (1, 3, 5), (1, 3, 5))) — 残差块的膨胀率。

采用 https://github.com/NVIDIA/BigVGAN（MIT 许可证）的 BigVGAN-v2 声码器，并结合 https://github.com/junjun3518/alias-free-torch（Apache 2.0）中的抗锯齿 Snake 激活函数，将 MMAudioVAE 解码出的梅尔频谱转换为 Kandinsky6TI2VAPipeline 的波形。

forward

< source >

( mel: torch.Tensorreturn_dict: bool = True )

参数

mel（形状为 (batch_size, num_mels, num_mel_frames) 的 torch.Tensor）— 由 MMAudioVAE 解码得到的梅尔频谱图。
return_dict (bool，默认为 True) — 是否返回 ~models.autoencoder_kl.DecoderOutput 而不是普通元组。

Update on GitHub