AutoencoderKLKVAEAudio

A 1D convolutional variational autoencoder (VAE) with KL loss for audio, introduced by the Kandinsky Lab Team in KVAE-Audio. It compresses raw, full-band (48 kHz) waveforms into compact continuous latents and reconstructs them with high fidelity across speech, music, and general audio.

The model can be loaded with the following code snippet.

import torch
from diffusers import AutoencoderKLKVAEAudio

vae = AutoencoderKLKVAEAudio.from_pretrained("kandinskylab/KVAE-Audio", subfolder="diffusers", dtype=torch.float32)

AutoencoderKLKVAEAudio

class diffusers.AutoencoderKLKVAEAudio

< source >

( encoder_dim: int = 64encoder_rates: list = [2, 3, 4, 5, 8]latent_dim: int = 2048codebook_dim: int = 64decoder_dim: int = 1536decoder_rates: list = [8, 5, 4, 3, 2]sample_rate: int = 48000num_channels: int = 1attn_num_heads: int = 8 )

Parameters

encoder_dim (int, optional, defaults to 64) — Base channel dimension for the encoder; doubled at every downsampling stage.
encoder_rates (list[int], optional, defaults to [2, 3, 4, 5, 8]) — Strides for downsampling in the encoder.
latent_dim (int, optional, defaults to 2048) — Channel dimension of the encoder output (before the in_proj bottleneck).
codebook_dim (int, optional, defaults to 64) — Channel dimension of the posterior distribution’s latent space (after the in_proj bottleneck).
decoder_dim (int, optional, defaults to 1536) — Base channel dimension for the decoder; halved at every upsampling stage.
decoder_rates (list[int], optional, defaults to [8, 5, 4, 3, 2]) — Strides for upsampling in the decoder.
sample_rate (int, optional, defaults to 48000) — The sampling rate, in Hz, that the model was trained on.
num_channels (int, optional, defaults to 1) — Number of channels in the audio data (1 for mono).
attn_num_heads (int, optional, defaults to 8) — Number of heads in the self-attention block applied to the encoder output.

A 1D convolutional autoencoder for encoding raw audio waveforms into continuous latents and decoding them back into waveforms. Introduced in KVAE-Audio.

This model inherits from ModelMixin. Check the superclass documentation for it’s generic methods implemented for all models (such as downloading or saving).

decode

< source >

( z: Tensorreturn_dict: bool = True ) → DecoderOutput or tuple

Parameters

z (torch.Tensor) — Input batch of latent vectors.
return_dict (bool, optional, defaults to True) — Whether to return a DecoderOutput instead of a plain tuple.

Returns

DecoderOutput or tuple

If return_dict is True, a DecoderOutput is returned, otherwise a plain tuple is returned.

Decode a batch of latents into audio waveforms.

encode

< source >

( audio_data: Tensorreturn_dict: bool = True )

Parameters

audio_data (torch.Tensor of shape (batch_size, num_channels, num_samples)) — Input batch of raw audio waveforms.
return_dict (bool, optional, defaults to True) — Whether to return a ~models.autoencoder_kl.AutoencoderKLOutput instead of a plain tuple.

Encode a batch of audio waveforms into latents.

forward

< source >

( sample: Tensorsample_posterior: bool = Falsereturn_dict: bool = Truegenerator: typing.Optional[torch.Generator] = None ) → DecoderOutput or tuple

Parameters

sample (torch.Tensor of shape (batch_size, num_channels, num_samples)) — Input batch of raw audio waveforms.
sample_posterior (bool, optional, defaults to False) — Whether to sample from the posterior.
return_dict (bool, optional, defaults to True) — Whether or not to return a DecoderOutput instead of a plain tuple.
generator (torch.Generator, optional) — A torch.Generator to make sampling deterministic.

Returns

DecoderOutput or tuple

If return_dict is True, a DecoderOutput is returned, otherwise a plain tuple is returned.

Update on GitHub

Diffusers

AutoencoderKLKVAEAudio