[feat] add KVAE-Audio autoencoder (#14384)

* [feat] add KVAE-Audio autoencoder

* [fix] non-contiguous tensor into in_proj in AutoencoderKLKVAEAudio

* [chore] sync KVAE-Audio with upstream test and annotation changes

* [refactor] dedicated attention class and processor for AutoencoderKLKVAEAudio

* [refactor] remove the runtime sample_rate argument from AutoencoderKLKVAEAudio

* [refactor] always apply attention in AutoencoderKLKVAEAudio, drop use_attn

* [refactor] drop the derived latent_dim default from AutoencoderKLKVAEAudio

* [docs] fix encoder_rates docstring in AutoencoderKLKVAEAudio

* [refactor] rename Snake1d to KVAEAudioSnake1d in AutoencoderKLKVAEAudio

* [tests] skip layerwise casting training for AutoencoderKLKVAEAudio

docs/source/en/_toctree.yml

@@ -481,6 +481,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

docs/source/en/api/models/autoencoder_kl_kvae_audio.md

@@ -0,0 +1,36 @@
+<!-- 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. -->
+
+# 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

scripts/convert_kvae_audio_to_diffusers.py

@@ -0,0 +1,184 @@
+#!/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.<i>.block.<j>
+                                -> {encoder,decoder}.block.<stride-1>.res_unit<i+1>.{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.<stride-1>.snake1      (pre-downsample activation)
+  encoder.model.{1..5}.block.4 -> encoder.block.<stride-1>.conv1       (strided downsample conv)
+  decoder.model.{1..5}.block.0 -> decoder.block.<stride-1>.snake1
+  decoder.model.{1..5}.block.1 -> decoder.block.<stride-1>.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)
+    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)
+
+    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,
+    )

src/diffusers/__init__.py

@@ -258,6 +258,7 @@ else:
             "AutoencoderKLHunyuanVideo",
             "AutoencoderKLHunyuanVideo15",
             "AutoencoderKLKVAE",
+            "AutoencoderKLKVAEAudio",
             "AutoencoderKLKVAEVideo",
             "AutoencoderKLLTX2Audio",
             "AutoencoderKLLTX2Video",
@@ -1173,6 +1174,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
             AutoencoderKLHunyuanVideo,
             AutoencoderKLHunyuanVideo15,
             AutoencoderKLKVAE,
+            AutoencoderKLKVAEAudio,
             AutoencoderKLKVAEVideo,
             AutoencoderKLLTX2Audio,
             AutoencoderKLLTX2Video,

src/diffusers/models/__init__.py

@@ -42,6 +42,7 @@ if is_torch_available():
     _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"]
@@ -197,6 +198,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
             AutoencoderKLHunyuanVideo,
             AutoencoderKLHunyuanVideo15,
             AutoencoderKLKVAE,
+            AutoencoderKLKVAEAudio,
             AutoencoderKLKVAEVideo,
             AutoencoderKLLTX2Audio,
             AutoencoderKLLTX2Video,

src/diffusers/models/autoencoders/__init__.py

@@ -12,6 +12,7 @@ from .autoencoder_kl_hunyuanimage import AutoencoderKLHunyuanImage
 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

src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py

@@ -0,0 +1,414 @@
+# 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 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
+
+
+class KVAEAudioSnake1d(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) -> 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)
+        hidden_states = hidden_states.reshape(shape)
+        return hidden_states
+
+
+# Copied from diffusers.models.autoencoders.autoencoder_oobleck.OobleckResidualUnit with Oobleck->KVAEAudio, Snake1d->KVAEAudioSnake1d
+class KVAEAudioResidualUnit(nn.Module):
+    """
+    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 = KVAEAudioSnake1d(dimension)
+        self.conv1 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=7, dilation=dilation, padding=pad))
+        self.snake2 = KVAEAudioSnake1d(dimension)
+        self.conv2 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=1))
+
+    def forward(self, hidden_state) -> torch.Tensor:
+        """
+        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, Snake1d->KVAEAudioSnake1d
+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 = KVAEAudioSnake1d(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) -> 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))
+        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 = 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(
+                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) -> torch.Tensor:
+        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 = 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:
+        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 = KVAEAudioSnake1d(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) -> torch.Tensor:
+        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 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).
+
+    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.
+        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.
+    """
+
+    _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: int = 2048,
+        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,
+        attn_num_heads: int = 8,
+    ):
+        super().__init__()
+
+        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.attn = KVAEAudioAttention(latent_dim, attn_num_heads)
+
+        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)
+
+        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, 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.
+            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 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_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_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).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)

src/diffusers/utils/dummy_pt_objects.py

@@ -724,6 +724,21 @@ class AutoencoderKLKVAE(metaclass=DummyObject):
         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"]
 

tests/models/autoencoders/test_models_autoencoder_kl_kvae_audio.py

@@ -0,0 +1,119 @@
+# 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 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 (
+    AttentionTesterMixin,
+    BaseModelTesterConfig,
+    MemoryTesterMixin,
+    ModelTesterMixin,
+    TorchCompileTesterMixin,
+)
+from .testing_utils import AutoencoderTesterMixin
+
+
+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],
+            "latent_dim": 16,
+            "codebook_dim": 4,
+            "decoder_dim": 16,
+            "decoder_rates": [3, 2],
+            "sample_rate": 48000,
+            "num_channels": 1,
+            "attn_num_heads": 2,
+        }
+
+    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):
+    @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):
+    pass
+
+
+class TestAutoencoderKLKVAEAudioSlicing(AutoencoderKLKVAEAudioTesterConfig, AutoencoderTesterMixin):
+    pass
+
+
+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."
+    )
+    def test_attention_processor_count_mismatch_raises_error(self):
+        pass