Kandinsky 6 VAE

Kandinsky 6 は、動画スーパーレゾリューションに因果的な 3D K-VAE を、同期した音声生成には MMAudio のメルスペクトログラム VAE と独立した BigVGAN MMAudioVocoder を組み合わせて使用しています。

Kandinsky6SRVAE

Kandinsky6SRPipelineで使用される因果的3次元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)) — 各エンコーダレベルにおける残差ブロックの出力幅です。最終レベルを除き、次のレベルへ進むにつれて空間サイズが半分になり、幅が2倍になります。
decoder_block_out_channels (tuple[int, ...]、デフォルトは (16, 256, 512, 1024, 2048)) — 各デコーダレベルにおける残差ブロックの出力幅です。
layers_per_block (int、デフォルトは 2) — エンコーダの各レベルあたりの残差ブロックの数。デコーダは各レベルで1つ多く使用します。
temporal_compression_ratio (int、デフォルトは 4) — 時間方向の圧縮係数です。log2 だけ連続したレベルでも、時間方向を圧縮します。
temporal_compression_start_level (int、デフォルトは 1) — 時間を圧縮する最初のレベルです。
scaling_factor (float、デフォルト値は0.910344) — 拡散トランスフォーマーに入力される前に、潜在変数に適用されるスケーリング係数です。

Kandinsky6SRPipelineが動画のエンコードとデコードに使用するCausal 3D K-VAE。

動画は、16ピクセルフレーム（先頭フレームを含む）の時間的セグメントごとに処理されます。因果的畳み込みはセグメント間でパディング状態を引き継ぐため、セグメント化はピークメモリを制限するだけで、結果には影響しません。

encode

< source >

( x: torch.Tensorreturn_dict: bool = True )

パラメータ

x (torch.Tensor、shape (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 (torch.Tensor、shape (batch_size, latent_channels, num_latent_frames, height, width)) — scaling_factor で割られた後の Latents。
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 が使用するメロスペクトログラム VAE です。その decode の出力はメロスペクトログラムであり、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
)

class 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) — Mel バインの数。
latent_channels (int、デフォルトは 40) — 潜在チャネルの数。
hidden_channels (int、デフォルトは 512) — オートエンコーダーのベース幅です。
channel_multipliers (tuple[int, ...]、デフォルトは(1, 2, 4)) — オートエンコーダの各レベルの幅の倍率です。
layers_per_block (int、デフォルトは 2) — エンコーダの各レベルあたりの残差ブロック数。デコーダは各レベルで1つ多く使用します。
sample_rate (int、デフォルトは44100) — 波形のサンプリングレート。
n_fft (int, デフォルトは 2048) — MelフロントエンドのFFTサイズ。
hop_length (int、デフォルトは512) — Melフロントエンドのホップ長です。このVAEと組み合わせて使用されるMMAudioVocoderの総アップサンプリング係数と一致する必要があります。
scaling_factor (float、デフォルト値は0.417) — 拡散トランスフォーマーに入力される前に、潜在変数に適用されるスケーリング係数です。

Kandinsky6TI2VAPipelineのAudio VAE: log-melスペクトログラムに対する、振幅を保持するオートエンコーダー(MMAudio、https://arxiv.org/abs/2412.15322)。

encode は波形を潜在分布に変換し、decode は潜在変数をメルスペクトログラムに戻します。その後、MMAudioVocoder がそれを波形に変換します。1つの潜在フレームは hop_length * 2 サンプルをカバーします。

encode

< source >

( audio: torch.Tensorreturn_dict: bool = True )

パラメータ

audio (torch.Tensor、形状 (batch_size, num_samples)) — sample_rate でサンプリングされた、[-1, 1] 範囲のモノラル波形。
return_dict (bool、デフォルトはTrue) — プレーンなタプルの代わりにAutoencoderKLOutputを返すかどうかを指定します。

波形をその潜在分布にエンコードします。

decode

< source >

( z: torch.Tensorreturn_dict: bool = True )

パラメータ

z (torch.Tensor、shape (batch_size, latent_channels, num_latent_frames)) — scaling_factor で割られた状態の Latents。
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, num_samples)) — sample_rate で、mel スペクトログラムとしてエンコードおよび再構成するための、[-1, 1] 範囲のモノラル波形。
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 は、形状が (batch_size, mel_bins, num_mel_frames) の再構成されたメルスペクトログラムです。

MMAudioVocoder

NVIDIA/BigVGAN(MITライセンス)のBigVGAN-v2ボコーダーMMAudioバンドルから派生し、さらに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) — 入力スペクトログラムのメルビンの数です。対応する MMAudioVAE の mel_bins と一致する必要があります。
upsample_initial_channel (int、デフォルト値は1536) — 最初の層の幅です。
upsample_rates (tuple[int, ...]、デフォルトは (8, 4, 2, 2, 2, 2)) — Vocoder の各ステージにおけるアップサンプリング係数です。これらの積が総アップサンプリング係数となり、対応する MMAudioVAE の hop_length と一致する必要があります。
upsample_kernel_sizes (tuple[int, ...]、デフォルトは (16, 8, 4, 4, 4, 4)) — Vocoder の各ステージにおける転置畳み込みカーネルサイズです。
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 (torch.Tensor、形状 (batch_size, num_mels, num_mel_frames)) — MMAudioVAE によってデコードされた Mel スペクトログラム。
return_dict (bool、デフォルトはTrue) — ~models.autoencoder_kl.DecoderOutputを返すかどうか。

Update on GitHub