Diffusers documentation

VAE для Kandinsky 6

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

VAE для Kandinsky 6

Kandinsky 6 использует каузальный 3D K-VAE для супер-разрешения видео и VAE мел-спектрограммы MMAudio, в паре с отдельным MMAudioVocoder BigVGAN, для синхронизированной генерации аудио.

Kandinsky6SRVAE

Причинный 3D K-VAE, используемый в Kandinsky6SRPipeline. Он обрабатывает видео произвольной длины в сегментах с ограниченной памятью, воспроизводя точный результат одного прохода без сегментации.

import torch
from diffusers import Kandinsky6SRVAE

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

класс diffusers.Kandinsky6SRVAE

< >

( 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) — Масштаб, применяемый к латентам перед их передачей в диффузионный трансформер.

Причинный 3D K-VAE, используемый Kandinsky6SRPipeline для кодирования и декодирования видео.

Видео обрабатываются во временных сегментах по 16 пиксельных кадров (плюс начальный кадр). Причинные свёртки переносят состояние своей отсечки между сегментами, поэтому сегментация ограничивает только пиковое потребление памяти и не изменяет результат.

encode

< >

( x: torch.Tensorreturn_dict: bool = True )

Параметры

  • x (torch.Tensor формы (batch_size, channels, num_frames, height, width)) — Пиксельное видео в диапазоне [-1, 1]. num_frames должно быть равно 1 + k * temporal_compression_ratio.
  • return_dict (bool, по умолчанию True) — Определяет, следует ли возвращать AutoencoderKLOutput вместо обычного кортежа.

Кодирует видео в его латентное распределение.

decode

< >

( z: torch.Tensorreturn_dict: bool = True )

Параметры

  • z (torch.Tensor формы (batch_size, latent_channels, num_latent_frames, height, width)) — Латенты, уже разделённые на scaling_factor.
  • return_dict (bool, по умолчанию True) — Определяет, следует ли возвращать ~models.autoencoder_kl.DecoderOutput вместо обычного кортежа.

Декодировать латенты в видео.

forward

< >

( 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

Мел-спектрограммный VAE, используемый Kandinsky6TI2VAPipeline при sample_audio=True. Выход его 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

< >

( 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) — Количество остаточных блоков на каждом уровне кодировщика; декодер использует на один блок больше на каждом уровне.
  • sample_rate (int, по умолчанию 44100) — Частота дискретизации аудиосигнала.
  • n_fft (int, по умолчанию 2048) — размер FFT для мел-фронтенда.
  • hop_length (int, по умолчанию 512) — Шаг (hop length) мел-фронтенда. Должен совпадать с общим коэффициентом апсемплинга MMAudioVocoder, с которым используется данный VAE.
  • scaling_factor (float, по умолчанию 0.417) — Масштаб, применяемый к латентам перед их передачей в диффузионный трансформер.

Аудио-VAE для Kandinsky6TI2VAPipeline: автоэнкодер, сохраняющий амплитуду, работающий с лог-мел-спектрограммами (MMAudio, https://arxiv.org/abs/2412.15322).

encode преобразует звуковой сигнал в латентное распределение; decode возвращает латенты в мел-спектрограмму, которую MMAudioVocoder затем преобразует в звуковой сигнал. Одна латентная рамка покрывает hop_length * 2 сэмпла.

encode

< >

( audio: torch.Tensorreturn_dict: bool = True )

Параметры

  • audio (torch.Tensor формы (batch_size, num_samples)) — Моно-волна в диапазоне [-1, 1] с частотой дискретизации sample_rate.
  • return_dict (bool, по умолчанию True) — Определяет, следует ли возвращать AutoencoderKLOutput вместо обычного кортежа.

Кодирует звуковую волну в её латентное распределение.

decode

< >

( z: torch.Tensorreturn_dict: bool = True )

Параметры

  • z (torch.Tensor формы (batch_size, latent_channels, num_latent_frames)) — Латенты, уже разделённые на scaling_factor.
  • return_dict (bool, по умолчанию True) — Определяет, следует ли возвращать ~models.autoencoder_kl.DecoderOutput вместо обычного кортежа.

Декодирование латентов в мел-спектрограмму.

forward

< >

( 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)) — Монофоническая волна в диапазоне [-1, 1] с частотой дискретизации sample_rate для кодирования и реконструкции в виде мел-спектрограммы.
  • 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

Адаптировано из вокодерных пакетов MMAudio, основанных на BigVGAN-v2, который, в свою очередь, взят из NVIDIA/BigVGAN (лицензия MIT), с использованием антиалиасинговых активаций Snake из alias-free-torch (Apache License 2.0).

класс diffusers.MMAudioVocoder

< >

( 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-полос входной спектрограммы. Должно совпадать с mel_bins соответствующей модели MMAudioVAE.
  • upsample_initial_channel (int, по умолчанию 1536) — Ширина первого слоя.
  • upsample_rates (tuple[int, ...], по умолчанию (8, 4, 2, 2, 2, 2)) — Коэффициенты апсемплинга для стадий вокодера. Их произведение является общим коэффициентом апсемплинга и должно совпадать с hop_length соответствующей модели MMAudioVAE.
  • 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))) — Размеры дилейшнов в остаточных блоках.

Вокодер BigVGAN-v2 (https://github.com/NVIDIA/BigVGAN, лицензия MIT) с антиалиасинговыми активациями Snake из https://github.com/junjun3518/alias-free-torch (Apache 2.0), преобразующий мел-спектрограммы, декодированные MMAudioVAE, в звуковые волны для Kandinsky6TI2VAPipeline.

forward

< >

( mel: torch.Tensorreturn_dict: bool = True )

Параметры

  • mel (torch.Tensor формы (batch_size, num_mels, num_mel_frames)) — Мел-спектрограмма, декодированная MMAudioVAE.
  • return_dict (bool, по умолчанию True) — Определяет, следует ли возвращать ~models.autoencoder_kl.DecoderOutput вместо обычного кортежа.
Update on GitHub