VAEs do Kandinsky 6

O Kandinsky 6 usa um K-VAE 3D causal para super-resolução de vídeo e o VAE de mel-espectrograma do MMAudio, emparelhado com um MMAudioVocoder BigVGAN separado, para geração de áudio sincronizado.

Kandinsky6SRVAE

O K-VAE 3D causal usado pelo Kandinsky6SRPipeline. Ele processa vídeos de comprimento arbitrário em segmentos de memória limitada, reproduzindo a saída exata de uma única passagem não segmentada.

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 )

Parâmetros

in_channels (int, padrão 3) — Número de canais de pixels.
out_channels (int, padrão 3) — Número de canais de pixels reconstruídos.
latent_channels (int, padrão 64) — Número de canais latentes.
encoder_block_out_channels (tuple[int, ...], padrão (16, 128, 256, 512, 1024)) — Largura de saída dos blocos residuais em cada nível do codificador; todos os níveis, exceto o último, reduzem pela metade o tamanho espacial e dobram a largura ao avançar para o próximo nível.
decoder_block_out_channels (tuple[int, ...], padrão (16, 256, 512, 1024, 2048)) — Largura de saída dos blocos residuais em cada nível do decodificador.
layers_per_block (int, padrão 2) — Número de blocos residuais por nível do codificador; o decodificador usa um a mais por nível.
temporal_compression_ratio (int, padrão 4) — Fator de compressão temporal; log2 de seus níveis consecutivos também comprime o tempo.
temporal_compression_start_level (int, padrão 1) — Primeiro nível que comprime o tempo.
scaling_factor (float, padrão 0.910344) — Escala aplicada aos latentes antes de entrarem no difusor transformador.

Causal 3D K-VAE usado pelo Kandinsky6SRPipeline para codificar e decodificar vídeo.

Os vídeos são processados em segmentos temporais de 16 quadros de pixels (mais o quadro inicial). As convoluções causais mantêm seu estado de preenchimento entre os segmentos, de modo que a segmentação apenas limita a memória de pico e não altera o resultado.

encode

< source >

( x: torch.Tensorreturn_dict: bool = True )

Parâmetros

x (torch.Tensor de forma (batch_size, channels, num_frames, height, width)) — Vídeo em pixels em [-1, 1]. num_frames deve ser 1 + k * temporal_compression_ratio.
return_dict (bool, padrão True) — Se deve retornar um AutoencoderKLOutput em vez de uma tupla simples.

Codifique um vídeo em sua distribuição latente.

decode

< source >

( z: torch.Tensorreturn_dict: bool = True )

Parâmetros

z (torch.Tensor de forma (batch_size, latent_channels, num_latent_frames, height, width)) — Latentes, já divididos por scaling_factor.
return_dict (bool, padrão True) — Se deve retornar um ~models.autoencoder_kl.DecoderOutput em vez de uma tupla simples.

Decodifique os latentes em um vídeo.

forward

< source >

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

Parâmetros

sample (torch.Tensor de forma (batch_size, channels, num_frames, height, width)) — Vídeo em pixels em [-1, 1]. num_frames deve ser 1 + k * temporal_compression_ratio.
sample_posterior (bool, opcional, padrão False) — Se deve amostrar da posterior latente em vez de usar seu modo.
return_dict (bool, opcional, padrão True) — Se deve retornar um ~models.autoencoder_kl.DecoderOutput em vez de uma tupla simples.
generator (torch.Generator, opcional) — Um torch.Generator para tornar a amostragem determinística.

Retorna

~models.autoencoder_kl.DecoderOutput ou tuple

Se return_dict for True, um ~models.autoencoder_kl.DecoderOutput é retornado; caso contrário, uma tuple simples é retornada. Seu sample é o vídeo reconstruído.

MMAudioVAE

O VAE de espectrograma mel usado por Kandinsky6TI2VAPipeline quando sample_audio=True. A saída de decode é um espectrograma mel; passe-o por MMAudioVocoder para obter uma forma de onda.

A implementação de referência pode ser encontrada em hkchengrex/MMAudio (licença 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 )

Parâmetros

mel_bins (int, padrão 128) — Número de bins mel.
latent_channels (int, padrão 40) — Número de canais latentes.
hidden_channels (int, padrão 512) — Largura base do autoencoder.
channel_multipliers (tuple[int, ...], padrão (1, 2, 4)) — Multiplicadores de largura dos níveis do autoencoder.
layers_per_block (int, padrão 2) — Blocos residuais por nível do codificador; o decodificador usa um a mais por nível.
sample_rate (int, padrão 44100) — Taxa de amostragem da forma de onda.
n_fft (int, padrão 2048) — Tamanho da FFT do front end mel.
hop_length (int, padrão 512) — Comprimento do salto do front end mel. Deve corresponder ao fator total de upsampling do MMAudioVocoder com o qual este VAE é pareado.
scaling_factor (float, padrão 0.417) — Escala aplicada aos latentes antes de entrarem no difusor transformador.

VAE de áudio de Kandinsky6TI2VAPipeline: um autoencoder que preserva a magnitude sobre espectrogramas log-mel (MMAudio, https://arxiv.org/abs/2412.15322).

encode converte uma forma de onda em uma distribuição latente; decode converte os latentes de volta em um espectrograma mel, que o MMAudioVocoder então converte em uma forma de onda. Um quadro latente cobre hop_length * 2 amostras.

encode

< source >

( audio: torch.Tensorreturn_dict: bool = True )

Parâmetros

audio (torch.Tensor de forma (batch_size, num_samples)) — Forma de onda mono em [-1, 1] em sample_rate.
return_dict (bool, padrão True) — Se deve retornar um AutoencoderKLOutput em vez de uma tupla simples.

Codifique uma forma de onda em sua distribuição latente.

decode

< source >

( z: torch.Tensorreturn_dict: bool = True )

Parâmetros

z (torch.Tensor de forma (batch_size, latent_channels, num_latent_frames)) — Latentes, já divididos por scaling_factor.
return_dict (bool, padrão True) — Se deve retornar um ~models.autoencoder_kl.DecoderOutput em vez de uma tupla simples.

Decodifique os latentes em um espectrograma mel.

forward

< source >

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

Parâmetros

sample (torch.Tensor de forma (batch_size, num_samples)) — Forma de onda mono em [-1, 1] em sample_rate para codificar e reconstruir como um espectrograma mel.
sample_posterior (bool, opcional, padrão False) — Se deve amostrar da posterior latente em vez de usar seu modo.
return_dict (bool, opcional, padrão True) — Se deve retornar um ~models.autoencoder_kl.DecoderOutput em vez de uma tupla simples.
generator (torch.Generator, opcional) — Um torch.Generator para tornar a amostragem determinística.

Retorna

~models.autoencoder_kl.DecoderOutput ou tuple

Se return_dict for True, um ~models.autoencoder_kl.DecoderOutput é retornado; caso contrário, uma tuple simples é retornada. Seu sample é o espectrograma mel reconstruído, com forma (batch_size, mel_bins, num_mel_frames).

MMAudioVocoder

Adaptado dos bundles do vocoder BigVGAN-v2 do MMAudio, que por sua vez é derivado de NVIDIA/BigVGAN (licença MIT), com as ativações Snake anti-alias 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)) )

Parâmetros

num_mels (int, padrão 128) — Número de bins mel do espectrograma de entrada. Deve corresponder ao mel_bins do MMAudioVAE correspondente.
upsample_initial_channel (int, padrão 1536) — Largura da primeira camada.
upsample_rates (tuple[int, ...], padrão (8, 4, 2, 2, 2, 2)) — Fatores de upsampling dos estágios do vocoder. O produto deles é o fator total de upsampling e deve corresponder ao hop_length do MMAudioVAE correspondente.
upsample_kernel_sizes (tuple[int, ...], padrão (16, 8, 4, 4, 4, 4)) — Tamanhos dos kernels de convolução transposta dos estágios do vocoder.
resblock_kernel_sizes (tuple[int, ...], padrão (3, 7, 11)) — Tamanhos dos kernels dos blocos residuais.
resblock_dilation_sizes (tuple[tuple[int, ...], ...], padrão ((1, 3, 5), (1, 3, 5), (1, 3, 5))) — Dilatações dos blocos residuais.

Vocoder BigVGAN-v2 (https://github.com/NVIDIA/BigVGAN, licença MIT) com as ativações Snake anti-alias de https://github.com/junjun3518/alias-free-torch (Apache 2.0), convertendo os espectrogramas mel decodificados por MMAudioVAE em formas de onda para Kandinsky6TI2VAPipeline.

forward

< source >

( mel: torch.Tensorreturn_dict: bool = True )

Parâmetros

mel (torch.Tensor de forma (batch_size, num_mels, num_mel_frames)) — Espectrograma Mel, conforme decodificado por MMAudioVAE.
return_dict (bool, padrão True) — Se deve retornar um ~models.autoencoder_kl.DecoderOutput em vez de uma tupla simples.

Update on GitHub