diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 8466cebe2a8c..25aca3ed5f70 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -449,6 +449,8 @@ title: AutoencoderKLHunyuanVideo15 - local: api/models/autoencoder_kl_kvae title: AutoencoderKLKVAE + - local: api/models/autoencoder_kl_kvae_audio + title: AutoencoderKLKVAEAudio - local: api/models/autoencoder_kl_kvae_video title: AutoencoderKLKVAEVideo - local: api/models/autoencoderkl_audio_ltx_2 diff --git a/docs/source/en/api/models/autoencoder_kl_kvae_audio.md b/docs/source/en/api/models/autoencoder_kl_kvae_audio.md new file mode 100644 index 000000000000..5b0eda5d2688 --- /dev/null +++ b/docs/source/en/api/models/autoencoder_kl_kvae_audio.md @@ -0,0 +1,36 @@ + + +# AutoencoderKLKVAEAudio + +A 1D convolutional variational autoencoder (VAE) with KL loss for audio, introduced by the Kandinsky Lab +Team in [KVAE-Audio](https://huggingface.co/kandinskylab/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. + +```python +import torch +from diffusers import AutoencoderKLKVAEAudio + +vae = AutoencoderKLKVAEAudio.from_pretrained("kandinskylab/KVAE-Audio", subfolder="diffusers", dtype=torch.float32) +``` + +## AutoencoderKLKVAEAudio + +[[autodoc]] AutoencoderKLKVAEAudio + - decode + - encode + - all diff --git a/scripts/convert_kvae_audio_to_diffusers.py b/scripts/convert_kvae_audio_to_diffusers.py new file mode 100644 index 000000000000..09b5f7a3b862 --- /dev/null +++ b/scripts/convert_kvae_audio_to_diffusers.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Usage: +# python scripts/convert_kvae_audio_to_diffusers.py \ +# --checkpoint_path /path/to/kvae-1/release_checkpoints/KVAE-Audio \ +# --output_path /path/to/output + +""" +Converts a `kandinskylab/KVAE-Audio` checkpoint (the `kvae.models.kvae_1d.KVAEAudio` architecture) into an +`AutoencoderKLKVAEAudio`. + +Key mapping, `encoder`/`decoder` (reference `Encoder1D`/`Decoder1D` -> `KVAEAudioEncoder`/`KVAEAudioDecoder`): + {enc,dec}.model.0 -> {encoder,decoder}.conv1 (stem conv) + {enc,dec}.model.{1..5}.block..block. + -> {encoder,decoder}.block..res_unit.{snake1,conv1,snake2,conv2} + (encoder: i in {0,1,2}; decoder: i in {2,3,4}, since decoder's block.0/block.1 are its own + snake1/conv_t1, not a res_unit) + encoder.model.{1..5}.block.3 -> encoder.block..snake1 (pre-downsample activation) + encoder.model.{1..5}.block.4 -> encoder.block..conv1 (strided downsample conv) + decoder.model.{1..5}.block.0 -> decoder.block..snake1 + decoder.model.{1..5}.block.1 -> decoder.block..conv_t1 (strided upsample conv) + {enc,dec}.model.6 -> {encoder,decoder}.snake1 (final activation) + {enc,dec}.model.7 -> {encoder,decoder}.conv2 (final conv) + in_proj / out_proj -> unchanged (same attribute names on both sides) + attn.in_proj_{weight,bias} -> attn.to_{q,k,v}.{weight,bias} (equal 3-way chunk, nn.MultiheadAttention layout) + attn.out_proj.{weight,bias} -> attn.to_out.0.{weight,bias} + +`bias` / `alpha` / `weight_g` / `weight_v` suffixes are copied unchanged. The reference uses +`torch.nn.utils.weight_norm`; the diffusers module uses `torch.nn.utils.parametrizations.weight_norm`. +The two store the reparametrization under different keys (`weight_g`/`weight_v` vs. +`parametrizations.weight.original0`/`original1`), but `load_state_dict` on the new-style module +accepts the legacy key names directly and reconstructs the same weights, so no explicit rename is +needed here. +""" + +import argparse +import json +from pathlib import Path + +import torch +from safetensors.torch import load_file + +from diffusers import AutoencoderKLKVAEAudio + + +RES_UNIT_NAMES = {0: "res_unit1", 1: "res_unit2", 2: "res_unit3"} +RES_UNIT_SUB_NAMES = {0: "snake1", 1: "conv1", 2: "snake2", 3: "conv2"} + + +def convert_encoder_decoder_key(prefix: str, idx: int, rest: list[str]) -> str: + if idx == 0: + return f"{prefix}.conv1." + ".".join(rest) + if idx == 6: + return f"{prefix}.snake1." + ".".join(rest) + if idx == 7: + return f"{prefix}.conv2." + ".".join(rest) + + # idx in {1..5}: one per-stride block, nested under an extra "block" level in the reference + # (an `OrderedDict([("block", block)])` wrapper used to keep the reference's own state-dict layout). + block_idx = idx - 1 + assert rest[0] == "block", f"expected 'block', got {rest}" + inner = int(rest[1]) + remaining = rest[2:] + + if prefix == "encoder": + if inner in RES_UNIT_NAMES: + res_unit = RES_UNIT_NAMES[inner] + assert remaining[0] == "block", f"expected 'block', got {remaining}" + sub_name = RES_UNIT_SUB_NAMES[int(remaining[1])] + param = ".".join(remaining[2:]) + return f"encoder.block.{block_idx}.{res_unit}.{sub_name}.{param}" + if inner == 3: + return f"encoder.block.{block_idx}.snake1." + ".".join(remaining) + if inner == 4: + return f"encoder.block.{block_idx}.conv1." + ".".join(remaining) + else: + if inner == 0: + return f"decoder.block.{block_idx}.snake1." + ".".join(remaining) + if inner == 1: + return f"decoder.block.{block_idx}.conv_t1." + ".".join(remaining) + if inner - 2 in RES_UNIT_NAMES: + res_unit = RES_UNIT_NAMES[inner - 2] + assert remaining[0] == "block", f"expected 'block', got {remaining}" + sub_name = RES_UNIT_SUB_NAMES[int(remaining[1])] + param = ".".join(remaining[2:]) + return f"decoder.block.{block_idx}.{res_unit}.{sub_name}.{param}" + + raise ValueError(f"Unhandled {prefix}.model.{idx}.{'.'.join(rest)}") + + +def convert_kvae_audio_state_dict(original_state_dict: dict) -> dict: + converted_state_dict = {} + + for key, value in original_state_dict.items(): + parts = key.split(".") + + if parts[0] in ("encoder", "decoder") and parts[1] == "model": + new_key = convert_encoder_decoder_key(parts[0], int(parts[2]), parts[3:]) + elif key.startswith("in_proj.") or key.startswith("out_proj."): + new_key = key + elif key == "attn.in_proj_weight": + query, key_, value_ = torch.chunk(value, 3, dim=0) + converted_state_dict["attn.to_q.weight"] = query + converted_state_dict["attn.to_k.weight"] = key_ + converted_state_dict["attn.to_v.weight"] = value_ + continue + elif key == "attn.in_proj_bias": + query, key_, value_ = torch.chunk(value, 3, dim=0) + converted_state_dict["attn.to_q.bias"] = query + converted_state_dict["attn.to_k.bias"] = key_ + converted_state_dict["attn.to_v.bias"] = value_ + continue + elif key == "attn.out_proj.weight": + new_key = "attn.to_out.0.weight" + elif key == "attn.out_proj.bias": + new_key = "attn.to_out.0.bias" + else: + raise ValueError(f"Unhandled key: {key}") + + converted_state_dict[new_key] = value + + return converted_state_dict + + +def convert_kvae_audio(checkpoint_path: str, output_path: str, dtype: str = "fp32"): + dtype_map = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16} + torch_dtype = dtype_map[dtype] + + checkpoint_dir = Path(checkpoint_path) + with open(checkpoint_dir / "config.json") as f: + config = json.load(f) + config.pop("model_type", None) + + original_state_dict = load_file(checkpoint_dir / "model.safetensors") + converted_state_dict = convert_kvae_audio_state_dict(original_state_dict) + + model = AutoencoderKLKVAEAudio(**config) + model.load_state_dict(converted_state_dict, strict=True) + model = model.to(dtype=torch_dtype) + + output_path = Path(output_path) + output_path.mkdir(parents=True, exist_ok=True) + model.save_pretrained(output_path) + + # round-trip check + AutoencoderKLKVAEAudio.from_pretrained(output_path, torch_dtype=torch_dtype) + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--checkpoint_path", + type=str, + required=True, + help="Path to a local directory containing config.json and model.safetensors", + ) + parser.add_argument("--output_path", type=str, required=True, help="Output directory") + parser.add_argument( + "--dtype", type=str, default="fp32", choices=["fp32", "fp16", "bf16"], help="Data type for converted weights" + ) + return parser.parse_args() + + +if __name__ == "__main__": + args = get_args() + convert_kvae_audio( + checkpoint_path=args.checkpoint_path, + output_path=args.output_path, + dtype=args.dtype, + ) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 7a8d727aefea..c92bd0ea2140 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -243,6 +243,7 @@ "AutoencoderKLHunyuanVideo", "AutoencoderKLHunyuanVideo15", "AutoencoderKLKVAE", + "AutoencoderKLKVAEAudio", "AutoencoderKLKVAEVideo", "AutoencoderKLLTX2Audio", "AutoencoderKLLTX2Video", @@ -1084,6 +1085,7 @@ AutoencoderKLHunyuanVideo, AutoencoderKLHunyuanVideo15, AutoencoderKLKVAE, + AutoencoderKLKVAEAudio, AutoencoderKLKVAEVideo, AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 167ee7a534de..689467b815dd 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -41,6 +41,7 @@ _import_structure["autoencoders.autoencoder_kl_hunyuanimage_refiner"] = ["AutoencoderKLHunyuanImageRefiner"] _import_structure["autoencoders.autoencoder_kl_hunyuanvideo15"] = ["AutoencoderKLHunyuanVideo15"] _import_structure["autoencoders.autoencoder_kl_kvae"] = ["AutoencoderKLKVAE"] + _import_structure["autoencoders.autoencoder_kl_kvae_audio"] = ["AutoencoderKLKVAEAudio"] _import_structure["autoencoders.autoencoder_kl_kvae_video"] = ["AutoencoderKLKVAEVideo"] _import_structure["autoencoders.autoencoder_kl_ltx"] = ["AutoencoderKLLTXVideo"] _import_structure["autoencoders.autoencoder_kl_ltx2"] = ["AutoencoderKLLTX2Video"] @@ -175,6 +176,7 @@ AutoencoderKLHunyuanVideo, AutoencoderKLHunyuanVideo15, AutoencoderKLKVAE, + AutoencoderKLKVAEAudio, AutoencoderKLKVAEVideo, AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index 145f62376192..34c74bad6278 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -11,6 +11,7 @@ from .autoencoder_kl_hunyuanimage_refiner import AutoencoderKLHunyuanImageRefiner from .autoencoder_kl_hunyuanvideo15 import AutoencoderKLHunyuanVideo15 from .autoencoder_kl_kvae import AutoencoderKLKVAE +from .autoencoder_kl_kvae_audio import AutoencoderKLKVAEAudio from .autoencoder_kl_kvae_video import AutoencoderKLKVAEVideo from .autoencoder_kl_ltx import AutoencoderKLLTXVideo from .autoencoder_kl_ltx2 import AutoencoderKLLTX2Video diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py new file mode 100644 index 000000000000..3fa1df7fdd90 --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py @@ -0,0 +1,399 @@ +# Copyright 2026 The Kandinsky Lab Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import math +from typing import Optional + +import numpy as np +import torch +import torch.nn as nn +from torch.nn.utils.parametrizations import weight_norm + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils.accelerate_utils import apply_forward_hook +from ..attention_processor import Attention +from ..modeling_outputs import AutoencoderKLOutput +from ..modeling_utils import ModelMixin +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution + + +class Snake1d(nn.Module): + """ + A 1-dimensional Snake activation function module. + """ + + def __init__(self, channels): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, hidden_states): + shape = hidden_states.shape + hidden_states = hidden_states.reshape(shape[0], shape[1], -1) + hidden_states = hidden_states + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * hidden_states).pow(2) + hidden_states = hidden_states.reshape(shape) + return hidden_states + + +# Copied from diffusers.models.autoencoders.autoencoder_oobleck.OobleckResidualUnit with Oobleck->KVAEAudio +class KVAEAudioResidualUnit(nn.Module): + """ + A residual unit composed of Snake1d and weight-normalized Conv1d layers with dilations. + """ + + def __init__(self, dimension: int = 16, dilation: int = 1): + super().__init__() + pad = ((7 - 1) * dilation) // 2 + + self.snake1 = Snake1d(dimension) + self.conv1 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=7, dilation=dilation, padding=pad)) + self.snake2 = Snake1d(dimension) + self.conv2 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=1)) + + def forward(self, hidden_state): + """ + Forward pass through the residual unit. + + Args: + hidden_state (`torch.Tensor` of shape `(batch_size, channels, time_steps)`): + Input tensor . + + Returns: + output_tensor (`torch.Tensor` of shape `(batch_size, channels, time_steps)`) + Input tensor after passing through the residual unit. + """ + output_tensor = hidden_state + output_tensor = self.conv1(self.snake1(output_tensor)) + output_tensor = self.conv2(self.snake2(output_tensor)) + + padding = (hidden_state.shape[-1] - output_tensor.shape[-1]) // 2 + if padding > 0: + hidden_state = hidden_state[..., padding:-padding] + output_tensor = hidden_state + output_tensor + return output_tensor + + +# Copied from diffusers.models.autoencoders.autoencoder_oobleck.OobleckEncoderBlock with Oobleck->KVAEAudio +class KVAEAudioEncoderBlock(nn.Module): + """Encoder block used in KVAEAudio encoder.""" + + def __init__(self, input_dim, output_dim, stride: int = 1): + super().__init__() + + self.res_unit1 = KVAEAudioResidualUnit(input_dim, dilation=1) + self.res_unit2 = KVAEAudioResidualUnit(input_dim, dilation=3) + self.res_unit3 = KVAEAudioResidualUnit(input_dim, dilation=9) + self.snake1 = Snake1d(input_dim) + self.conv1 = weight_norm( + nn.Conv1d(input_dim, output_dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)) + ) + + def forward(self, hidden_state): + hidden_state = self.res_unit1(hidden_state) + hidden_state = self.res_unit2(hidden_state) + hidden_state = self.snake1(self.res_unit3(hidden_state)) + hidden_state = self.conv1(hidden_state) + + return hidden_state + + +class KVAEAudioDecoderBlock(nn.Module): + """Decoder block used in KVAEAudio decoder.""" + + def __init__(self, input_dim, output_dim, stride: int = 1): + super().__init__() + + self.snake1 = Snake1d(input_dim) + # odd strides need output_padding=1 to invert the encoder's downsampling exactly, or decode() silently truncates + self.conv_t1 = weight_norm( + nn.ConvTranspose1d( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + output_padding=stride % 2, + ) + ) + self.res_unit1 = KVAEAudioResidualUnit(output_dim, dilation=1) + self.res_unit2 = KVAEAudioResidualUnit(output_dim, dilation=3) + self.res_unit3 = KVAEAudioResidualUnit(output_dim, dilation=9) + + def forward(self, hidden_state): + hidden_state = self.snake1(hidden_state) + hidden_state = self.conv_t1(hidden_state) + hidden_state = self.res_unit1(hidden_state) + hidden_state = self.res_unit2(hidden_state) + hidden_state = self.res_unit3(hidden_state) + + return hidden_state + + +class KVAEAudioEncoder(nn.Module): + """KVAEAudio encoder: strided Conv1d downsampling with dilated residual blocks.""" + + def __init__(self, encoder_dim: int, encoder_rates: list[int], latent_dim: int, num_channels: int): + super().__init__() + + self.conv1 = weight_norm(nn.Conv1d(num_channels, encoder_dim, kernel_size=7, padding=3)) + + d_model = encoder_dim + blocks = [] + for stride in encoder_rates: + input_dim = d_model + d_model *= 2 + blocks.append(KVAEAudioEncoderBlock(input_dim, d_model, stride=stride)) + self.block = nn.ModuleList(blocks) + + self.snake1 = Snake1d(d_model) + self.conv2 = weight_norm(nn.Conv1d(d_model, latent_dim, kernel_size=3, padding=1)) + + def forward(self, hidden_states): + hidden_states = self.conv1(hidden_states) + + for block in self.block: + hidden_states = block(hidden_states) + + hidden_states = self.snake1(hidden_states) + hidden_states = self.conv2(hidden_states) + return hidden_states + + +class KVAEAudioDecoder(nn.Module): + """KVAEAudio decoder: strided ConvTranspose1d upsampling with dilated residual blocks.""" + + def __init__(self, latent_dim: int, decoder_dim: int, decoder_rates: list[int], num_channels: int): + super().__init__() + + self.conv1 = weight_norm(nn.Conv1d(latent_dim, decoder_dim, kernel_size=7, padding=3)) + + d_model = decoder_dim + blocks = [] + for stride in decoder_rates: + input_dim = d_model + d_model = d_model // 2 + blocks.append(KVAEAudioDecoderBlock(input_dim, d_model, stride=stride)) + self.block = nn.ModuleList(blocks) + + self.snake1 = Snake1d(d_model) + self.conv2 = weight_norm(nn.Conv1d(d_model, num_channels, kernel_size=7, padding=3)) + self.tanh = nn.Tanh() + + def forward(self, hidden_states): + hidden_states = self.conv1(hidden_states) + + for block in self.block: + hidden_states = block(hidden_states) + + hidden_states = self.snake1(hidden_states) + hidden_states = self.conv2(hidden_states) + hidden_states = self.tanh(hidden_states) + return hidden_states + + +class AutoencoderKLKVAEAudio(ModelMixin, AutoencoderMixin, ConfigMixin): + r""" + A 1D convolutional autoencoder for encoding raw audio waveforms into continuous latents and decoding them back into + waveforms. Introduced in [KVAE-Audio](https://huggingface.co/kandinskylab/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). + + 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. Used in reverse order for upsampling in the decoder. + latent_dim (`int`, *optional*): + Channel dimension of the encoder output (before the `in_proj` bottleneck). Defaults to `encoder_dim * 2 ** + len(encoder_rates)` when not provided. + 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). + use_attn (`bool`, *optional*, defaults to `False`): + Whether to apply a global self-attention block to the encoder output before the `in_proj` bottleneck. + attn_num_heads (`int`, *optional*, defaults to 8): + Number of attention heads used when `use_attn=True`. + """ + + _supports_gradient_checkpointing = False + _supports_group_offloading = False + + @register_to_config + def __init__( + self, + encoder_dim: int = 64, + encoder_rates: list[int] = [2, 3, 4, 5, 8], + latent_dim: Optional[int] = None, + codebook_dim: int = 64, + decoder_dim: int = 1536, + decoder_rates: list[int] = [8, 5, 4, 3, 2], + sample_rate: int = 48000, + num_channels: int = 1, + use_attn: bool = False, + attn_num_heads: int = 8, + ): + super().__init__() + + if latent_dim is None: + latent_dim = encoder_dim * 2 ** len(encoder_rates) + self.register_to_config(latent_dim=latent_dim) + + self.hop_length = int(np.prod(encoder_rates)) + + self.encoder = KVAEAudioEncoder(encoder_dim, encoder_rates, latent_dim, num_channels) + self.in_proj = weight_norm(nn.Conv1d(latent_dim, codebook_dim * 2, kernel_size=1)) + self.out_proj = weight_norm(nn.Conv1d(codebook_dim, latent_dim, kernel_size=1)) + self.decoder = KVAEAudioDecoder(latent_dim, decoder_dim, decoder_rates, num_channels) + + self.use_attn = use_attn + if use_attn: + self.attn = Attention( + query_dim=latent_dim, + heads=attn_num_heads, + dim_head=latent_dim // attn_num_heads, + bias=True, + out_bias=True, + ) + + self.use_slicing = False + + def _pad_to_hop_length(self, audio_data: torch.Tensor) -> torch.Tensor: + length = audio_data.shape[-1] + right_pad = math.ceil(length / self.hop_length) * self.hop_length - length + return nn.functional.pad(audio_data, (0, right_pad)) + + def _encode(self, audio_data: torch.Tensor) -> torch.Tensor: + audio_data = self._pad_to_hop_length(audio_data) + hidden_states = self.encoder(audio_data) + + if self.use_attn: + hidden_states = hidden_states.transpose(1, 2) + hidden_states = self.attn(hidden_states) + hidden_states = hidden_states.transpose(1, 2).contiguous() + + return self.in_proj(hidden_states) + + @apply_forward_hook + def encode( + self, audio_data: torch.Tensor, sample_rate: Optional[int] = None, return_dict: bool = True + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: + """ + Encode a batch of audio waveforms into latents. + + Args: + audio_data (`torch.Tensor` of shape `(batch_size, num_channels, num_samples)`): + Input batch of raw audio waveforms. + sample_rate (`int`, *optional*): + Sample rate of `audio_data`, in Hz. If given, it must match `self.config.sample_rate`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + The latent representations of the encoded audio. If `return_dict` is True, a + [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. + """ + if sample_rate is not None and sample_rate != self.config.sample_rate: + raise ValueError( + f"`sample_rate` ({sample_rate}) does not match the model's configured sample rate " + f"({self.config.sample_rate})." + ) + + if self.use_slicing and audio_data.shape[0] > 1: + encoded_slices = [self._encode(x_slice) for x_slice in audio_data.split(1)] + moments = torch.cat(encoded_slices) + else: + moments = self._encode(audio_data) + + posterior = DiagonalGaussianDistribution(moments) + + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) + + def _decode(self, z: torch.Tensor) -> torch.Tensor: + z = self.out_proj(z) + return self.decoder(z) + + @apply_forward_hook + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: + """ + Decode a batch of latents into audio waveforms. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoders.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.autoencoders.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.autoencoders.vae.DecoderOutput`] is returned, otherwise a plain + `tuple` is returned. + """ + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [self._decode(z_slice) for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z) + + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) + + def forward( + self, + sample: torch.Tensor, + sample_rate: Optional[int] = None, + sample_posterior: bool = False, + return_dict: bool = True, + generator: Optional[torch.Generator] = None, + ) -> DecoderOutput | torch.Tensor: + r""" + Args: + sample (`torch.Tensor` of shape `(batch_size, num_channels, num_samples)`): + Input batch of raw audio waveforms. + sample_rate (`int`, *optional*): + Sample rate of `sample`, in Hz. If given, it must match `self.config.sample_rate`. + 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 [`~models.autoencoders.vae.DecoderOutput`] instead of a plain tuple. + generator (`torch.Generator`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling + deterministic. + + Returns: + [`~models.autoencoders.vae.DecoderOutput`] or `tuple`: + If `return_dict` is True, a [`~models.autoencoders.vae.DecoderOutput`] is returned, otherwise a plain + `tuple` is returned. + """ + length = sample.shape[-1] + + posterior = self.encode(sample, sample_rate=sample_rate).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + decoded = self.decode(z).sample[..., :length] + + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 9035efb3e6e2..432c25112cc5 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -690,6 +690,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class AutoencoderKLKVAEAudio(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class AutoencoderKLKVAEVideo(metaclass=DummyObject): _backends = ["torch"] diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py new file mode 100644 index 000000000000..9277dfb1f941 --- /dev/null +++ b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py @@ -0,0 +1,107 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from diffusers import AutoencoderKLKVAEAudio +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TorchCompileTesterMixin +from .testing_utils import NewAutoencoderTesterMixin + + +enable_full_determinism() + + +class AutoencoderKLKVAEAudioTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return AutoencoderKLKVAEAudio + + @property + def main_input_name(self) -> str: + return "sample" + + @property + def output_shape(self) -> tuple: + return (1, 37) + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict: + return { + "encoder_dim": 4, + # Includes an odd stride (3) so the odd-stride `output_padding` fix in + # `KVAEAudioDecoderBlock` (needed for the decoder to invert the encoder's downsampling + # exactly) is actually exercised by the test suite. + "encoder_rates": [2, 3], + "codebook_dim": 4, + "decoder_dim": 16, + "decoder_rates": [3, 2], + "sample_rate": 48000, + "num_channels": 1, + } + + def get_dummy_inputs(self) -> dict: + batch_size = 2 + num_channels = 1 + seq_len = 37 + waveform = randn_tensor((batch_size, num_channels, seq_len), generator=self.generator, device=torch_device) + return {"sample": waveform, "sample_posterior": False} + + +class TestAutoencoderKLKVAEAudio(AutoencoderKLKVAEAudioTesterConfig, ModelTesterMixin): + def test_forward_output_length_matches_input_length(self): + # Regression test: the decoder must invert the encoder's downsampling exactly, for any + # input length, including ones not a multiple of `hop_length`. Without `output_padding` on + # the odd-stride `ConvTranspose1d`s, decode() undershoots the (hop-length-padded) input + # length by a fixed amount per odd stride, so forward()'s trailing `[..., :length]` slice + # becomes a no-op and silently returns short audio. + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device).eval() + + for seq_len in (1, 5, 6, 7, 37, 100): + waveform = randn_tensor((1, 1, seq_len), generator=self.generator, device=torch_device) + with torch.no_grad(): + output = model(waveform, sample_posterior=False) + assert output.sample.shape[-1] == seq_len, ( + f"expected output length {seq_len}, got {output.sample.shape[-1]}" + ) + + +class TestAutoencoderKLKVAEAudioMemory(AutoencoderKLKVAEAudioTesterConfig, MemoryTesterMixin): + pass + + +class TestAutoencoderKLKVAEAudioTorchCompile(AutoencoderKLKVAEAudioTesterConfig, TorchCompileTesterMixin): + pass + + +class TestAutoencoderKLKVAEAudioSlicing(AutoencoderKLKVAEAudioTesterConfig, NewAutoencoderTesterMixin): + pass + + +class AutoencoderKLKVAEAudioAttnTesterConfig(AutoencoderKLKVAEAudioTesterConfig): + def get_init_dict(self) -> dict: + init_dict = super().get_init_dict() + init_dict.update({"use_attn": True, "attn_num_heads": 2}) + return init_dict + + +class TestAutoencoderKLKVAEAudioAttn(AutoencoderKLKVAEAudioAttnTesterConfig, ModelTesterMixin): + pass