From c0c5a0d1b41073fa2cc2fa87850542036be53e04 Mon Sep 17 00:00:00 2001 From: Alexander Ivanov Date: Tue, 4 Aug 2026 19:13:08 +0300 Subject: [PATCH 01/10] [feat] add KVAE-Audio autoencoder --- docs/source/en/_toctree.yml | 2 + .../api/models/autoencoder_kl_kvae_audio.md | 36 ++ scripts/convert_kvae_audio_to_diffusers.py | 183 ++++++++ src/diffusers/__init__.py | 2 + src/diffusers/models/__init__.py | 2 + src/diffusers/models/autoencoders/__init__.py | 1 + .../autoencoders/autoencoder_kl_kvae_audio.py | 399 ++++++++++++++++++ src/diffusers/utils/dummy_pt_objects.py | 15 + .../test_models_autoencoder_kl_kvae_audio.py | 107 +++++ 9 files changed, 747 insertions(+) create mode 100644 docs/source/en/api/models/autoencoder_kl_kvae_audio.md create mode 100644 scripts/convert_kvae_audio_to_diffusers.py create mode 100644 src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py create mode 100644 tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py 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..7f9f66a1327d --- /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) + + 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 From 7c25e2288e3cc916a5ecdbafd6ace0e6f966a69e Mon Sep 17 00:00:00 2001 From: Alexander Ivanov Date: Tue, 4 Aug 2026 20:25:48 +0300 Subject: [PATCH 02/10] [fix] non-contiguous tensor into in_proj in AutoencoderKLKVAEAudio --- src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py index 7f9f66a1327d..3fa1df7fdd90 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py @@ -288,7 +288,7 @@ def _encode(self, audio_data: torch.Tensor) -> torch.Tensor: if self.use_attn: hidden_states = hidden_states.transpose(1, 2) hidden_states = self.attn(hidden_states) - hidden_states = hidden_states.transpose(1, 2) + hidden_states = hidden_states.transpose(1, 2).contiguous() return self.in_proj(hidden_states) From af16f7a19a3be5d3967664a8bb73bb51f4c5c6a4 Mon Sep 17 00:00:00 2001 From: Alexander Ivanov Date: Mon, 19 Oct 2026 13:43:50 +0300 Subject: [PATCH 03/10] [chore] sync KVAE-Audio with upstream test and annotation changes --- .../models/autoencoders/autoencoder_kl_kvae_audio.py | 12 ++++++------ .../test_models_autoencoder_kl_kvae_audio.py | 4 ++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py index 3fa1df7fdd90..92b0e1d8351b 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py @@ -38,7 +38,7 @@ def __init__(self, channels): super().__init__() self.alpha = nn.Parameter(torch.ones(1, channels, 1)) - def forward(self, hidden_states): + def forward(self, hidden_states) -> torch.Tensor: 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) @@ -61,7 +61,7 @@ def __init__(self, dimension: int = 16, dilation: int = 1): self.snake2 = Snake1d(dimension) self.conv2 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=1)) - def forward(self, hidden_state): + def forward(self, hidden_state) -> torch.Tensor: """ Forward pass through the residual unit. @@ -99,7 +99,7 @@ def __init__(self, input_dim, output_dim, stride: int = 1): nn.Conv1d(input_dim, output_dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)) ) - def forward(self, hidden_state): + def forward(self, hidden_state) -> torch.Tensor: hidden_state = self.res_unit1(hidden_state) hidden_state = self.res_unit2(hidden_state) hidden_state = self.snake1(self.res_unit3(hidden_state)) @@ -130,7 +130,7 @@ def __init__(self, input_dim, output_dim, stride: int = 1): self.res_unit2 = KVAEAudioResidualUnit(output_dim, dilation=3) self.res_unit3 = KVAEAudioResidualUnit(output_dim, dilation=9) - def forward(self, hidden_state): + def forward(self, hidden_state) -> torch.Tensor: hidden_state = self.snake1(hidden_state) hidden_state = self.conv_t1(hidden_state) hidden_state = self.res_unit1(hidden_state) @@ -159,7 +159,7 @@ def __init__(self, encoder_dim: int, encoder_rates: list[int], latent_dim: int, 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): + def forward(self, hidden_states) -> torch.Tensor: hidden_states = self.conv1(hidden_states) for block in self.block: @@ -190,7 +190,7 @@ def __init__(self, latent_dim: int, decoder_dim: int, decoder_rates: list[int], self.conv2 = weight_norm(nn.Conv1d(d_model, num_channels, kernel_size=7, padding=3)) self.tanh = nn.Tanh() - def forward(self, hidden_states): + def forward(self, hidden_states) -> torch.Tensor: hidden_states = self.conv1(hidden_states) for block in self.block: diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py index 9277dfb1f941..0bc534e03fd5 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py @@ -20,7 +20,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TorchCompileTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -92,7 +92,7 @@ class TestAutoencoderKLKVAEAudioTorchCompile(AutoencoderKLKVAEAudioTesterConfig, pass -class TestAutoencoderKLKVAEAudioSlicing(AutoencoderKLKVAEAudioTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderKLKVAEAudioSlicing(AutoencoderKLKVAEAudioTesterConfig, AutoencoderTesterMixin): pass From ff73335c5711aaf6d27e1a1bc3c5fa2033f7d5cf Mon Sep 17 00:00:00 2001 From: Alexander Ivanov Date: Mon, 19 Oct 2026 13:43:50 +0300 Subject: [PATCH 04/10] [refactor] dedicated attention class and processor for AutoencoderKLKVAEAudio --- .../autoencoders/autoencoder_kl_kvae_audio.py | 56 ++++++++++++++++--- .../test_models_autoencoder_kl_kvae_audio.py | 18 +++++- 2 files changed, 64 insertions(+), 10 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py index 92b0e1d8351b..16faf03b67d0 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py @@ -23,7 +23,8 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...utils.accelerate_utils import apply_forward_hook -from ..attention_processor import Attention +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution @@ -202,7 +203,50 @@ def forward(self, hidden_states) -> torch.Tensor: return hidden_states -class AutoencoderKLKVAEAudio(ModelMixin, AutoencoderMixin, ConfigMixin): +class KVAEAudioAttnProcessor: + """Attention processor used by [`KVAEAudioAttention`].""" + + _attention_backend = None + _parallel_config = None + + def __call__(self, attn: "KVAEAudioAttention", hidden_states: torch.Tensor) -> torch.Tensor: + query = attn.to_q(hidden_states).unflatten(2, (attn.heads, -1)) + key = attn.to_k(hidden_states).unflatten(2, (attn.heads, -1)) + value = attn.to_v(hidden_states).unflatten(2, (attn.heads, -1)) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3) + return attn.to_out[0](hidden_states) + + +class KVAEAudioAttention(nn.Module, AttentionModuleMixin): + """Global multi-head self-attention applied to the encoder output before the `in_proj` bottleneck.""" + + _default_processor_cls = KVAEAudioAttnProcessor + _available_processors = [KVAEAudioAttnProcessor] + _supports_qkv_fusion = False + + def __init__(self, dim: int, heads: int): + super().__init__() + self.heads = heads + self.to_q = nn.Linear(dim, dim) + self.to_k = nn.Linear(dim, dim) + self.to_v = nn.Linear(dim, dim) + self.to_out = nn.ModuleList([nn.Linear(dim, dim), nn.Dropout(0.0)]) + + self.set_processor(KVAEAudioAttnProcessor()) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.processor(self, hidden_states) + + +class AutoencoderKLKVAEAudio(ModelMixin, AttentionMixin, 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). @@ -266,13 +310,7 @@ def __init__( 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.attn = KVAEAudioAttention(latent_dim, attn_num_heads) self.use_slicing = False diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py index 0bc534e03fd5..4f71573b819e 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py @@ -13,13 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest 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 ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, +) from .testing_utils import AutoencoderTesterMixin @@ -105,3 +112,12 @@ def get_init_dict(self) -> dict: class TestAutoencoderKLKVAEAudioAttn(AutoencoderKLKVAEAudioAttnTesterConfig, ModelTesterMixin): pass + + +class TestAutoencoderKLKVAEAudioAttention(AutoencoderKLKVAEAudioAttnTesterConfig, AttentionTesterMixin): + @pytest.mark.skip( + "The model has a single attention module, so a dict of processors can never mismatch the number of " + "attention layers." + ) + def test_attention_processor_count_mismatch_raises_error(self): + pass From b306555db4e0778b38f31124901646c242d66ba9 Mon Sep 17 00:00:00 2001 From: Alexander Ivanov Date: Mon, 19 Oct 2026 13:43:50 +0300 Subject: [PATCH 05/10] [refactor] remove the runtime sample_rate argument from AutoencoderKLKVAEAudio --- .../autoencoders/autoencoder_kl_kvae_audio.py | 15 ++------------- 1 file changed, 2 insertions(+), 13 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py index 16faf03b67d0..a4a0cd14dbdb 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py @@ -332,7 +332,7 @@ def _encode(self, audio_data: torch.Tensor) -> torch.Tensor: @apply_forward_hook def encode( - self, audio_data: torch.Tensor, sample_rate: Optional[int] = None, return_dict: bool = True + self, audio_data: torch.Tensor, return_dict: bool = True ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: """ Encode a batch of audio waveforms into latents. @@ -340,8 +340,6 @@ def encode( 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. @@ -349,12 +347,6 @@ def encode( 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) @@ -399,7 +391,6 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | t def forward( self, sample: torch.Tensor, - sample_rate: Optional[int] = None, sample_posterior: bool = False, return_dict: bool = True, generator: Optional[torch.Generator] = None, @@ -408,8 +399,6 @@ def forward( 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`): @@ -425,7 +414,7 @@ def forward( """ length = sample.shape[-1] - posterior = self.encode(sample, sample_rate=sample_rate).latent_dist + posterior = self.encode(sample).latent_dist if sample_posterior: z = posterior.sample(generator=generator) else: From d811f062dce0a5ca97d82fbae27d974ea99dc713 Mon Sep 17 00:00:00 2001 From: Alexander Ivanov Date: Mon, 19 Oct 2026 13:43:50 +0300 Subject: [PATCH 06/10] [refactor] always apply attention in AutoencoderKLKVAEAudio, drop use_attn --- scripts/convert_kvae_audio_to_diffusers.py | 1 + .../autoencoders/autoencoder_kl_kvae_audio.py | 17 +++++------------ .../test_models_autoencoder_kl_kvae_audio.py | 14 ++------------ 3 files changed, 8 insertions(+), 24 deletions(-) diff --git a/scripts/convert_kvae_audio_to_diffusers.py b/scripts/convert_kvae_audio_to_diffusers.py index 09b5f7a3b862..8c40adc6992d 100644 --- a/scripts/convert_kvae_audio_to_diffusers.py +++ b/scripts/convert_kvae_audio_to_diffusers.py @@ -143,6 +143,7 @@ def convert_kvae_audio(checkpoint_path: str, output_path: str, dtype: str = "fp3 with open(checkpoint_dir / "config.json") as f: config = json.load(f) config.pop("model_type", None) + config.pop("use_attn", None) original_state_dict = load_file(checkpoint_dir / "model.safetensors") converted_state_dict = convert_kvae_audio_state_dict(original_state_dict) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py index a4a0cd14dbdb..14bd782a7905 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py @@ -272,10 +272,8 @@ class AutoencoderKLKVAEAudio(ModelMixin, AttentionMixin, AutoencoderMixin, Confi 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`. + Number of heads in the self-attention block applied to the encoder output. """ _supports_gradient_checkpointing = False @@ -292,7 +290,6 @@ def __init__( 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__() @@ -307,10 +304,7 @@ def __init__( 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 = KVAEAudioAttention(latent_dim, attn_num_heads) + self.attn = KVAEAudioAttention(latent_dim, attn_num_heads) self.use_slicing = False @@ -323,10 +317,9 @@ 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() + 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) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py index 4f71573b819e..99797c236f5e 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py @@ -62,6 +62,7 @@ def get_init_dict(self) -> dict: "decoder_rates": [3, 2], "sample_rate": 48000, "num_channels": 1, + "attn_num_heads": 2, } def get_dummy_inputs(self) -> dict: @@ -103,18 +104,7 @@ class TestAutoencoderKLKVAEAudioSlicing(AutoencoderKLKVAEAudioTesterConfig, Auto 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 - - -class TestAutoencoderKLKVAEAudioAttention(AutoencoderKLKVAEAudioAttnTesterConfig, AttentionTesterMixin): +class TestAutoencoderKLKVAEAudioAttention(AutoencoderKLKVAEAudioTesterConfig, AttentionTesterMixin): @pytest.mark.skip( "The model has a single attention module, so a dict of processors can never mismatch the number of " "attention layers." From ec0d694ac5e956fcf04992d22196f15af485780e Mon Sep 17 00:00:00 2001 From: Alexander Ivanov Date: Mon, 19 Oct 2026 13:43:50 +0300 Subject: [PATCH 07/10] [refactor] drop the derived latent_dim default from AutoencoderKLKVAEAudio --- .../models/autoencoders/autoencoder_kl_kvae_audio.py | 11 +++-------- .../test_models_autoencoder_kl_kvae_audio.py | 1 + 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py index 14bd782a7905..c4f484e5ae9e 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py @@ -259,9 +259,8 @@ class AutoencoderKLKVAEAudio(ModelMixin, AttentionMixin, AutoencoderMixin, Confi 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. + 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): @@ -284,7 +283,7 @@ def __init__( self, encoder_dim: int = 64, encoder_rates: list[int] = [2, 3, 4, 5, 8], - latent_dim: Optional[int] = None, + latent_dim: int = 2048, codebook_dim: int = 64, decoder_dim: int = 1536, decoder_rates: list[int] = [8, 5, 4, 3, 2], @@ -294,10 +293,6 @@ def __init__( ): 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) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py index 99797c236f5e..fce39f2be2d1 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py @@ -57,6 +57,7 @@ def get_init_dict(self) -> dict: # `KVAEAudioDecoderBlock` (needed for the decoder to invert the encoder's downsampling # exactly) is actually exercised by the test suite. "encoder_rates": [2, 3], + "latent_dim": 16, "codebook_dim": 4, "decoder_dim": 16, "decoder_rates": [3, 2], From cf6c47f3e02ae086955824c33964a3a9540eed3c Mon Sep 17 00:00:00 2001 From: Alexander Ivanov Date: Mon, 19 Oct 2026 14:30:50 +0300 Subject: [PATCH 08/10] [docs] fix encoder_rates docstring in AutoencoderKLKVAEAudio --- src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py index c4f484e5ae9e..555184298a2a 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py @@ -258,7 +258,7 @@ class AutoencoderKLKVAEAudio(ModelMixin, AttentionMixin, AutoencoderMixin, Confi 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. + 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): From 3dcde8b6691f4f328b672469a741d48fbf945307 Mon Sep 17 00:00:00 2001 From: Alexander Ivanov Date: Mon, 19 Oct 2026 14:34:22 +0300 Subject: [PATCH 09/10] [refactor] rename Snake1d to KVAEAudioSnake1d in AutoencoderKLKVAEAudio --- .../autoencoders/autoencoder_kl_kvae_audio.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py index 555184298a2a..ad99f19e2a82 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py @@ -30,7 +30,7 @@ from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution -class Snake1d(nn.Module): +class KVAEAudioSnake1d(nn.Module): """ A 1-dimensional Snake activation function module. """ @@ -47,19 +47,19 @@ def forward(self, hidden_states) -> torch.Tensor: return hidden_states -# Copied from diffusers.models.autoencoders.autoencoder_oobleck.OobleckResidualUnit with Oobleck->KVAEAudio +# Copied from diffusers.models.autoencoders.autoencoder_oobleck.OobleckResidualUnit with Oobleck->KVAEAudio, Snake1d->KVAEAudioSnake1d class KVAEAudioResidualUnit(nn.Module): """ - A residual unit composed of Snake1d and weight-normalized Conv1d layers with dilations. + A residual unit composed of KVAEAudioSnake1d 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.snake1 = KVAEAudioSnake1d(dimension) self.conv1 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=7, dilation=dilation, padding=pad)) - self.snake2 = Snake1d(dimension) + self.snake2 = KVAEAudioSnake1d(dimension) self.conv2 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=1)) def forward(self, hidden_state) -> torch.Tensor: @@ -85,7 +85,7 @@ def forward(self, hidden_state) -> torch.Tensor: return output_tensor -# Copied from diffusers.models.autoencoders.autoencoder_oobleck.OobleckEncoderBlock with Oobleck->KVAEAudio +# Copied from diffusers.models.autoencoders.autoencoder_oobleck.OobleckEncoderBlock with Oobleck->KVAEAudio, Snake1d->KVAEAudioSnake1d class KVAEAudioEncoderBlock(nn.Module): """Encoder block used in KVAEAudio encoder.""" @@ -95,7 +95,7 @@ def __init__(self, input_dim, output_dim, stride: int = 1): 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.snake1 = KVAEAudioSnake1d(input_dim) self.conv1 = weight_norm( nn.Conv1d(input_dim, output_dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)) ) @@ -115,7 +115,7 @@ class KVAEAudioDecoderBlock(nn.Module): def __init__(self, input_dim, output_dim, stride: int = 1): super().__init__() - self.snake1 = Snake1d(input_dim) + self.snake1 = KVAEAudioSnake1d(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( @@ -157,7 +157,7 @@ def __init__(self, encoder_dim: int, encoder_rates: list[int], latent_dim: int, blocks.append(KVAEAudioEncoderBlock(input_dim, d_model, stride=stride)) self.block = nn.ModuleList(blocks) - self.snake1 = Snake1d(d_model) + self.snake1 = KVAEAudioSnake1d(d_model) self.conv2 = weight_norm(nn.Conv1d(d_model, latent_dim, kernel_size=3, padding=1)) def forward(self, hidden_states) -> torch.Tensor: @@ -187,7 +187,7 @@ def __init__(self, latent_dim: int, decoder_dim: int, decoder_rates: list[int], blocks.append(KVAEAudioDecoderBlock(input_dim, d_model, stride=stride)) self.block = nn.ModuleList(blocks) - self.snake1 = Snake1d(d_model) + self.snake1 = KVAEAudioSnake1d(d_model) self.conv2 = weight_norm(nn.Conv1d(d_model, num_channels, kernel_size=7, padding=3)) self.tanh = nn.Tanh() From ba2964be59d9163e239c4639d51858a27f043987 Mon Sep 17 00:00:00 2001 From: Alexander Ivanov Date: Mon, 19 Oct 2026 15:56:33 +0300 Subject: [PATCH 10/10] [tests] skip layerwise casting training for AutoencoderKLKVAEAudio --- .../autoencoders/test_models_autoencoder_kl_kvae_audio.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py index fce39f2be2d1..5a4d61fdecfd 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py @@ -94,7 +94,12 @@ def test_forward_output_length_matches_input_length(self): class TestAutoencoderKLKVAEAudioMemory(AutoencoderKLKVAEAudioTesterConfig, MemoryTesterMixin): - pass + @pytest.mark.skip( + "Not supported: training through weight_norm fails when layerwise casting stores its weights in float16 or " + "float8." + ) + def test_layerwise_casting_training(self): + pass class TestAutoencoderKLKVAEAudioTorchCompile(AutoencoderKLKVAEAudioTesterConfig, TorchCompileTesterMixin):