Kandinsky 6 VAE

Kandinsky 6는 비디오 초해상도에 인과적 3D K-VAE를, 동기화된 오디오 생성에는 MMAudio 멜 스펙트로그램 VAE와 별도의 BigVGAN MMAudioVocoder를 함께 사용합니다.

Kandinsky6SRVAE

Kandinsky6SRPipeline에서 사용하는 인과적 3D 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)) — 각 인코더 레벨의 잔차 블록 출력 폭입니다. 마지막 레벨을 제외한 모든 레벨은 다음 레벨로 넘어갈 때 공간 크기를 절반으로 줄이고 폭을 두 배로 늘립니다.
decoder_block_out_channels (tuple[int, ...], 기본값 (16, 256, 512, 1024, 2048)) — 각 디코더 레벨의 잔차 블록 출력 폭.
layers_per_block (int, 기본값 2) — 인코더 레벨당 잔차 블록의 수입니다. 디코더는 레벨당 하나 더 사용합니다.
temporal_compression_ratio (int, 기본값 4) — 시간 압축 계수입니다. log2 값만큼의 연속된 레벨도 시간을 압축합니다.
temporal_compression_start_level (int, 기본값 1) — 시간을 압축하는 첫 번째 레벨입니다.
scaling_factor (float, 기본값 0.910344) — diffusion transformer에 입력되기 전에 latents에 적용되는 스케일입니다.

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을 반환할지 여부.

Encode a video into its latent distribution.

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를 반환할지 여부.

latents를 비디오로 디코딩합니다.

forward

< source >

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

매개변수

sample (torch.Tensor, shape (batch_size, channels, num_frames, height, width)) — [-1, 1] 범위의 픽셀 비디오. num_frames은 1 + k * temporal_compression_ratio여야 합니다.
sample_posterior (bool, 선택 사항, 기본값 False) — latent posterior의 mode를 사용하는 대신 posterior에서 샘플링할지 여부.
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
)

클래스 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 bin의 수입니다.
latent_channels (int, 기본값 40) — 잠재 채널의 수입니다.
hidden_channels (int, 기본값 512) — autoencoder의 기본 폭입니다.
channel_multipliers (tuple[int, ...], 기본값 (1, 2, 4)) — 오토인코더 레벨의 폭 배율입니다.
layers_per_block (int, 기본값 2) — 인코더 레벨당 잔차 블록 수입니다. 디코더는 레벨당 하나 더 사용합니다.
sample_rate (int, 기본값 44100) — Waveform 샘플링 레이트입니다.
n_fft (int, 기본값 2048) — 멜 프론트 엔드의 FFT 크기입니다.
hop_length (int, 기본값 512) — Mel 프론트 엔드의 hop length입니다. 이 VAE와 페어링된 MMAudioVocoder의 총 업샘플링 계수와 일치해야 합니다.
scaling_factor (float, 기본값 0.417) — Diffusion transformer에 입력되기 전에 latents에 적용되는 스케일입니다.

Kandinsky6TI2VAPipeline의 Audio VAE: 로그 멜 스펙트로그램에 대한 크기 보존 오토인코더(MMAudio, https://arxiv.org/abs/2412.15322).

encode는 웨이브폼을 잠재 분포로 변환하고, decode는 잠재 변수를 다시 멜 스펙트로그램으로 변환합니다. 이후 MMAudioVocoder가 이를 웨이브폼으로 변환합니다. 하나의 프레임은 hop_length * 2개의 샘플을 포함합니다.

encode

< source >

( audio: torch.Tensorreturn_dict: bool = True )

매개변수

audio (torch.Tensor, shape (batch_size, num_samples)) — sample_rate에서 [-1, 1] 범위의 모노 웨이브폼.
return_dict (bool, 기본값 True) — 단순 튜플 대신 AutoencoderKLOutput을 반환할지 여부.

Encode a waveform into its latent distribution.

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를 반환할지 여부.

latents를 mel 스펙트로그램으로 디코딩합니다.

forward

< source >

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

매개변수

sample (torch.Tensor, shape (batch_size, num_samples)) — 멜 스펙트로그램으로 인코딩 및 재구성할 sample_rate의 [-1, 1] 범위를 가진 모노 웨이브폼.
sample_posterior (bool, 선택 사항, 기본값 False) — latent posterior의 mode를 사용하는 대신 posterior에서 샘플링할지 여부.
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) 모양의 재구성된 mel spectrogram입니다.

MMAudioVocoder

NVIDIA/BigVGAN(MIT 라이선스)의 BigVGAN-v2 보코더 MMAudio 번들에서 가져왔으며, alias-free-torch(Apache License 2.0)의 anti-aliased Snake activations를 포함합니다.

클래스 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) — 입력 스펙트로그램의 멜 빈(mel bins) 수입니다. 페어링된 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)) — 보코더 단계의 전치 합성곱 커널 크기입니다.
resblock_kernel_sizes (tuple[int, ...], 기본값 (3, 7, 11)) — 잔차 블록의 커널 크기입니다.
resblock_dilation_sizes (tuple[tuple[int, ...], ...], 기본값 ((1, 3, 5), (1, 3, 5), (1, 3, 5))) — 잔차 블록의 dilations입니다.

https://github.com/NVIDIA/BigVGAN (MIT 라이선스)의 BigVGAN-v2 보코더와 https://github.com/junjun3518/alias-free-torch (Apache 2.0)의 anti-aliased Snake activations를 사용하여, MMAudioVAE가 디코딩한 멜 스펙트로그램을 Kandinsky6TI2VAPipeline의 웨이브폼으로 변환합니다.

forward

< source >

( mel: torch.Tensorreturn_dict: bool = True )

매개변수

mel (torch.Tensor, shape (batch_size, num_mels, num_mel_frames)) — MMAudioVAE로 디코딩된 Mel spectrogram.
return_dict (bool, 기본값 True) — 단순 튜플 대신 ~models.autoencoder_kl.DecoderOutput를 반환할지 여부.

Update on GitHub