VAEs de Kandinsky 6

Kandinsky 6 utilise un K-VAE 3D causal pour la super-résolution vidéo et le VAE de spectrogramme mel de MMAudio, associé à un MMAudioVocoder BigVGAN distinct, pour la génération d’audio synchronisé.

Kandinsky6SRVAE

Le K-VAE causal 3D utilisé par Kandinsky6SRPipeline. Il traite des vidéos de longueur arbitraire en segments à mémoire bornée, tout en reproduisant exactement la sortie d'un passage unique, non segmenté.

import torch
from diffusers import Kandinsky6SRVAE

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

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

Paramètres

in_channels (int, par défaut 3) — Nombre de canaux de pixels.
out_channels (int, par défaut 3) — Nombre de canaux de pixels reconstruits.
latent_channels (int, par défaut 64) — Nombre de canaux latents.
encoder_block_out_channels (tuple[int, ...], par défaut (16, 128, 256, 512, 1024)) — Largeur de sortie des blocs résiduels à chaque niveau de l'encodeur ; à chaque niveau, sauf le dernier, la taille spatiale est divisée par deux et la largeur est doublée pour passer au niveau suivant.
decoder_block_out_channels (tuple[int, ...], par défaut (16, 256, 512, 1024, 2048)) — Largeur de sortie des blocs résiduels à chaque niveau du décodeur.
layers_per_block (int, par défaut 2) — Nombre de blocs résiduels par niveau d'encodeur ; le décodeur en utilise un de plus par niveau.
temporal_compression_ratio (int, par défaut 4) — Facteur de compression temporelle ; log2 de niveaux consécutifs le compriment également dans le temps.
temporal_compression_start_level (int, par défaut 1) — Premier niveau qui compresse le temps.
scaling_factor (float, par défaut 0.910344) — Échelle appliquée aux latents avant leur entrée dans le transformateur de diffusion.

Causal 3D K-VAE utilisé par Kandinsky6SRPipeline pour encoder et décoder la vidéo.

Les vidéos sont traitées par segments temporels de 16 trames de pixels (plus la trame initiale). Les convolutions causales transmettent leur état de bourrage entre les segments, de sorte que la segmentation ne fait que limiter la mémoire de pointe sans modifier le résultat.

encode

< source >

( x: torch.Tensorreturn_dict: bool = True )

Paramètres

x (torch.Tensor de forme (batch_size, channels, num_frames, height, width)) — Vidéo en pixels dans [-1, 1]. num_frames doit être 1 + k * temporal_compression_ratio.
return_dict (bool, par défaut True) — Indique s'il faut retourner un AutoencoderKLOutput plutôt qu'un simple tuple.

Encoder une vidéo dans sa distribution latente.

décoder

< source >

( z: torch.Tensorreturn_dict: bool = True )

Paramètres

z (torch.Tensor de forme (batch_size, latent_channels, num_latent_frames, height, width)) — Latents, déjà divisés par scaling_factor.
return_dict (bool, par défaut True) — Indique s'il faut retourner un ~models.autoencoder_kl.DecoderOutput plutôt qu'un simple tuple.

Décodez les latents en vidéo.

forward

< source >

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

Paramètres

sample (torch.Tensor de forme (batch_size, channels, num_frames, height, width)) — Vidéo en pixels dans [-1, 1]. num_frames doit être 1 + k * temporal_compression_ratio.
sample_posterior (bool, facultatif, par défaut False) — Indique s'il faut échantillonner à partir de la postérieure latente plutôt que d'utiliser son mode.
return_dict (bool, facultatif, par défaut True) — Indique s'il faut retourner un ~models.autoencoder_kl.DecoderOutput plutôt qu'un simple tuple.
generator (torch.Generator, facultatif) — Un torch.Generator pour rendre l'échantillonnage déterministe.

Retourne

~models.autoencoder_kl.DecoderOutput ou tuple

Si return_dict est True, un ~models.autoencoder_kl.DecoderOutput est renvoyé, sinon un simple tuple est renvoyé. Son sample est la vidéo reconstruite.

MMAudioVAE

Le VAE de spectrogramme mel utilisé par Kandinsky6TI2VAPipeline lorsque sample_audio=True. Sa sortie decode est un spectrogramme mel ; passez-le à travers MMAudioVocoder pour obtenir une forme d’onde.

L'implémentation de référence se trouve sur hkchengrex/MMAudio (licence 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
)

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

Paramètres

mel_bins (int, par défaut 128) — Nombre de bins mel.
latent_channels (int, par défaut 40) — Nombre de canaux latents.
hidden_channels (int, par défaut 512) — Largeur de base de l'autoencodeur.
channel_multipliers (tuple[int, ...], par défaut (1, 2, 4)) — Multiplicateurs de largeur des niveaux de l'autoencodeur.
layers_per_block (int, par défaut 2) — Nombre de blocs résiduels par niveau d'encodeur ; le décodeur en utilise un de plus par niveau.
sample_rate (int, par défaut 44100) — Taux d'échantillonnage de la forme d'onde.
n_fft (int, par défaut 2048) — Taille de la FFT du front-end mél.
hop_length (int, par défaut 512) — Longueur de saut de la partie avant mélo. Doit correspondre au facteur de suréchantillonnage total du MMAudioVocoder avec lequel ce VAE est associé.
scaling_factor (float, par défaut 0.417) — Échelle appliquée aux latents avant leur entrée dans le transformeur de diffusion.

VAE audio de Kandinsky6TI2VAPipeline : un auto-encodeur préservant l'amplitude sur des spectrogrammes log-mel (MMAudio, https://arxiv.org/abs/2412.15322).

encode convertit une forme d’onde en une distribution latente ; decode convertit les latents en un spectrogramme mel, que MMAudioVocoder transforme ensuite en forme d’onde. Un cadre latent couvre hop_length * 2 échantillons.

encode

< source >

( audio: torch.Tensorreturn_dict: bool = True )

Paramètres

audio (torch.Tensor de forme (batch_size, num_samples)) — Forme d’onde mono dans [-1, 1] à sample_rate.
return_dict (bool, par défaut True) — Indique s'il faut retourner un AutoencoderKLOutput plutôt qu'un simple tuple.

Encoder une forme d'onde dans sa distribution latente.

décoder

< source >

( z: torch.Tensorreturn_dict: bool = True )

Paramètres

z (torch.Tensor de forme (batch_size, latent_channels, num_latent_frames)) — Latents, déjà divisés par scaling_factor.
return_dict (bool, par défaut True) — Indique s'il faut retourner un ~models.autoencoder_kl.DecoderOutput plutôt qu'un simple tuple.

Décode les latents en spectrogramme mel.

forward

< source >

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

Paramètres

sample (torch.Tensor de forme (batch_size, num_samples)) — Forme d'onde mono dans [-1, 1] à sample_rate à encoder et reconstruire en spectrogramme mel.
sample_posterior (bool, facultatif, par défaut False) — Indique s'il faut échantillonner à partir de la postérieure latente plutôt que d'utiliser son mode.
return_dict (bool, facultatif, par défaut True) — Indique s'il faut retourner un ~models.autoencoder_kl.DecoderOutput plutôt qu'un simple tuple.
generator (torch.Generator, facultatif) — Un torch.Generator pour rendre l'échantillonnage déterministe.

Retourne

~models.autoencoder_kl.DecoderOutput ou tuple

Si return_dict est True, un ~models.autoencoder_kl.DecoderOutput est renvoyé, sinon un simple tuple est renvoyé. Son sample est le spectrogramme mel reconstruit de forme (batch_size, mel_bins, num_mel_frames).

MMAudioVocoder

Adapté des bundles de vocoder MMAudio issus de BigVGAN-v2, lui-même dérivé de NVIDIA/BigVGAN (licence MIT), avec les activations Snake anti-aliasing de alias-free-torch (Apache License 2.0).

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

Paramètres

num_mels (int, par défaut 128) — Nombre de bins mel du spectrogramme d’entrée. Doit correspondre à mel_bins du MMAudioVAE associé.
upsample_initial_channel (int, par défaut 1536) — Largeur de la première couche.
upsample_rates (tuple[int, ...], par défaut (8, 4, 2, 2, 2, 2)) — Facteurs de suréchantillonnage des étapes du vocodeur. Leur produit correspond au facteur de suréchantillonnage total et doit correspondre à la hop_length du MMAudioVAE associé.
upsample_kernel_sizes (tuple[int, ...], par défaut (16, 8, 4, 4, 4, 4)) — Tailles des noyaux des convolutions transposées des étapes du vocodeur.
resblock_kernel_sizes (tuple[int, ...], par défaut (3, 7, 11)) — Tailles des noyaux des blocs résiduels.
resblock_dilation_sizes (tuple[tuple[int, ...], ...], par défaut ((1, 3, 5), (1, 3, 5), (1, 3, 5))) — Dilatations des blocs résiduels.

Vocodeur BigVGAN-v2 (https://github.com/NVIDIA/BigVGAN, licence MIT) avec les activations Snake anti-aliasées de https://github.com/junjun3518/alias-free-torch (Apache 2.0), convertissant les spectrogrammes mélo que MMAudioVAE décode en formes d’onde pour Kandinsky6TI2VAPipeline.

forward

< source >

( mel: torch.Tensorreturn_dict: bool = True )

Paramètres

mel (torch.Tensor de forme (batch_size, num_mels, num_mel_frames)) — Spectrogramme Mel, tel que décodé par MMAudioVAE.
return_dict (bool, par défaut True) — Indique s'il faut retourner un ~models.autoencoder_kl.DecoderOutput plutôt qu'un simple tuple.

Update on GitHub