VAE Kandinsky 6

Kandinsky 6 menggunakan K-VAE 3D kausal untuk super-resolusi video dan VAE mel-spektrogram MMAudio, yang dipasangkan dengan MMAudioVocoder BigVGAN terpisah, untuk menghasilkan audio yang tersinkronisasi.

Kandinsky6SRVAE

K-VAE 3D kausal yang digunakan oleh Kandinsky6SRPipeline. Memproses video dengan panjang arbitrer dalam segmen dengan memori terbatas sambil mereproduksi output yang persis sama dengan satu kali proses tanpa segmentasi.

import torch
from diffusers import Kandinsky6SRVAE

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

kelas 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 )

Parameter

in_channels (int, defaultnya 3) — Jumlah kanal piksel.
out_channels (int, defaultnya 3) — Jumlah kanal piksel yang direkonstruksi.
latent_channels (int, defaultnya 64) — Jumlah kanal laten.
encoder_block_out_channels (tuple[int, ...], defaultnya (16, 128, 256, 512, 1024)) — Lebar output dari blok residual pada setiap level encoder; setiap level kecuali yang terakhir mengurangi ukuran spasial menjadi setengahnya dan menggandakan lebar menuju level berikutnya.
decoder_block_out_channels (tuple[int, ...], defaultnya (16, 256, 512, 1024, 2048)) — Lebar output dari blok residual di setiap level decoder.
layers_per_block (int, defaultnya 2) — Jumlah blok residual per level encoder; decoder menggunakan satu blok tambahan per level.
temporal_compression_ratio (int, defaultnya 4) — Faktor kompresi temporal; log2 dari tingkat berurutan yang juga mengompresi waktu.
temporal_compression_start_level (int, defaultnya 1) — Level pertama yang melakukan kompresi waktu.
scaling_factor (float, defaultnya 0.910344) — Skala yang diterapkan pada latents sebelum mereka masuk ke diffusion transformer.

Causal 3D K-VAE yang digunakan oleh Kandinsky6SRPipeline untuk mengodekan dan mendekodekan video.

Video diproses dalam segmen temporal yang terdiri dari 16 bingkai piksel (ditambah bingkai pembuka). Konvolusi kausal membawa status padding mereka antar segmen, sehingga segmentasi hanya membatasi memori puncak dan tidak mengubah hasil.

encode

< source >

( x: torch.Tensorreturn_dict: bool = True )

Parameter

x (torch.Tensor dengan bentuk (batch_size, channels, num_frames, height, width)) — Video piksel dalam [-1, 1]. num_frames harus berupa 1 + k * temporal_compression_ratio.
return_dict (bool, defaultnya True) — Apakah akan mengembalikan AutoencoderKLOutput alih-alih tuple biasa.

Enkodekan video ke dalam distribusi latennya.

decode

< source >

( z: torch.Tensorreturn_dict: bool = True )

Parameter

z (torch.Tensor dengan bentuk (batch_size, latent_channels, num_latent_frames, height, width)) — Latents, yang sudah dibagi dengan scaling_factor.
return_dict (bool, defaultnya True) — Apakah akan mengembalikan ~models.autoencoder_kl.DecoderOutput alih-alih tuple biasa.

Dekode latens menjadi video.

forward

< source >

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

Parameter

sample (torch.Tensor dengan bentuk (batch_size, channels, num_frames, height, width)) — Video piksel dalam [-1, 1]. num_frames harus berupa 1 + k * temporal_compression_ratio.
sample_posterior (bool, opsional, default False) — Apakah akan mengambil sampel dari posterior laten alih-alih menggunakan modusnya.
return_dict (bool, opsional, default True) — Apakah akan mengembalikan ~models.autoencoder_kl.DecoderOutput alih-alih tuple biasa.
generator (torch.Generator, opsional) — torch.Generator untuk membuat sampling menjadi deterministik.

Mengembalikan

~models.autoencoder_kl.DecoderOutput atau tuple

Jika return_dict bernilai True, sebuah ~models.autoencoder_kl.DecoderOutput akan dikembalikan, selain itu sebuah tuple polos akan dikembalikan. sample di dalamnya adalah video yang direkonstruksi.

MMAudioVAE

VAE mel-spectrogram yang digunakan oleh Kandinsky6TI2VAPipeline ketika sample_audio=True. Output decode berupa mel spectrogram; lewati melalui MMAudioVocoder untuk mendapatkan waveform.

Implementasi referensi dapat ditemukan di hkchengrex/MMAudio (lisensi 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
)

kelas 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 )

Parameter

mel_bins (int, defaultnya 128) — Jumlah mel bins.
latent_channels (int, defaultnya 40) — Jumlah kanal laten.
hidden_channels (int, defaultnya 512) — Lebar dasar dari autoencoder.
channel_multipliers (tuple[int, ...], defaultnya (1, 2, 4)) — Pengali lebar dari level autoencoder.
layers_per_block (int, defaultnya 2) — Blok residual per level encoder; decoder menggunakan satu blok tambahan per level.
sample_rate (int, defaultnya 44100) — Laju sampel waveform.
n_fft (int, defaultnya 2048) — Ukuran FFT dari mel front end.
hop_length (int, defaultnya 512) — Panjang hop dari mel front end. Harus sesuai dengan faktor upsampling total dari MMAudioVocoder yang dipasangkan dengan VAE ini.
scaling_factor (float, defaultnya 0.417) — Skala yang diterapkan pada latents sebelum mereka masuk ke diffusion transformer.

Audio VAE dari Kandinsky6TI2VAPipeline: autoencoder yang mempertahankan magnitudo pada spektrogram log-mel (MMAudio, https://arxiv.org/abs/2412.15322).

encode mengubah waveform menjadi distribusi laten; decode mengubah latens kembali menjadi spekrogram mel, yang kemudian MMAudioVocoder ubah menjadi waveform. Satu frame laten mencakup hop_length * 2 sampel.

encode

< source >

( audio: torch.Tensorreturn_dict: bool = True )

Parameter

audio (torch.Tensor dengan bentuk (batch_size, num_samples)) — Gelombang mono dalam [-1, 1] pada sample_rate.
return_dict (bool, defaultnya True) — Apakah akan mengembalikan AutoencoderKLOutput alih-alih tuple biasa.

Enkodekan gelombang suara ke dalam distribusi latennya.

decode

< source >

( z: torch.Tensorreturn_dict: bool = True )

Parameter

z (torch.Tensor dengan bentuk (batch_size, latent_channels, num_latent_frames)) — Latents, yang sudah dibagi dengan scaling_factor.
return_dict (bool, defaultnya True) — Apakah akan mengembalikan ~models.autoencoder_kl.DecoderOutput alih-alih tuple biasa.

Decode latents menjadi mel spectrogram.

forward

< source >

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

Parameter

sample (torch.Tensor dengan bentuk (batch_size, num_samples)) — Gelombang mono dalam [-1, 1] pada sample_rate untuk dikodekan dan direkonstruksi sebagai spektrogram mel.
sample_posterior (bool, opsional, default False) — Apakah akan mengambil sampel dari posterior laten alih-alih menggunakan modusnya.
return_dict (bool, opsional, default True) — Apakah akan mengembalikan ~models.autoencoder_kl.DecoderOutput alih-alih tuple biasa.
generator (torch.Generator, opsional) — torch.Generator untuk membuat sampling menjadi deterministik.

Mengembalikan

~models.autoencoder_kl.DecoderOutput atau tuple

Jika return_dict bernilai True, sebuah ~models.autoencoder_kl.DecoderOutput dikembalikan, selain itu sebuah tuple polos dikembalikan. sample di dalamnya adalah spektrogram mel yang direkonstruksi dengan bentuk (batch_size, mel_bins, num_mel_frames).

MMAudioVocoder

Diadaptasi dari bundle vocoder MMAudio BigVGAN-v2, yang sendiri berasal dari NVIDIA/BigVGAN (lisensi MIT), dengan aktivasi Snake anti-alias dari alias-free-torch (Apache License 2.0).

kelas 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)) )

Parameter

num_mels (int, defaultnya 128) — Jumlah mel bins dari spektrogram input. Harus cocok dengan mel_bins dari MMAudioVAE yang berpasangan.
upsample_initial_channel (int, defaultnya 1536) — Lebar dari lapisan pertama.
upsample_rates (tuple[int, ...], defaultnya (8, 4, 2, 2, 2, 2)) — Faktor upsampling dari tahap-tahap vocoder. Hasil perkalian mereka adalah faktor upsampling total dan harus cocok dengan hop_length dari MMAudioVAE yang berpasangan.
upsample_kernel_sizes (tuple[int, ...], defaultnya (16, 8, 4, 4, 4, 4)) — Ukuran kernel transposed-convolution dari tahap-tahap vocoder.
resblock_kernel_sizes (tuple[int, ...], defaultnya (3, 7, 11)) — Ukuran kernel dari blok residual.
resblock_dilation_sizes (tuple[tuple[int, ...], ...], defaultnya ((1, 3, 5), (1, 3, 5), (1, 3, 5))) — Dilasi dari blok residual.

Vocoder BigVGAN-v2 (https://github.com/NVIDIA/BigVGAN, lisensi MIT) dengan aktivasi Snake anti-alias dari https://github.com/junjun3518/alias-free-torch (Apache 2.0), mengubah spektrogram mel yang didekode oleh MMAudioVAE menjadi waveform untuk Kandinsky6TI2VAPipeline.

forward

< source >

( mel: torch.Tensorreturn_dict: bool = True )

Parameter

mel (torch.Tensor dengan bentuk (batch_size, num_mels, num_mel_frames)) — Spektrogram Mel, sebagaimana didekodekan oleh MMAudioVAE.
return_dict (bool, defaultnya True) — Apakah akan mengembalikan ~models.autoencoder_kl.DecoderOutput alih-alih tuple biasa.

Update on GitHub