Diffusers documentation
Kandinsky 6 VAE
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, defaults toTrue) — Whether to return an AutoencoderKLOutput instead of a plain tuple.
動画をその潜在分布にエンコードします。
decode
< source >( z: torch.Tensorreturn_dict: bool = True )
潜在変数から動画をデコードします。
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, defaults toTrue) — Whether to return an AutoencoderKLOutput instead of a plain tuple.
波形をその潜在分布にエンコードします。
decode
< source >( z: torch.Tensorreturn_dict: bool = True )
潜在変数からメロスペクトログラムをデコードします。
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を返すかどうか。