[feat] add KVAE-Audio autoencoder

[ivnvalex]

What does this PR do?

Adds KVAE-Audio as AutoencoderKLKVAEAudio, a 1D convolutional audio VAE, completing the KVAE tokenizers family from the Kandinsky Lab team in Diffusers alongside the already-merged AutoencoderKLKVAE (image) and AutoencoderKLKVAEVideo (video).
It loads the checkpoint from kandinskylab/KVAE-Audio and compresses/reconstructs raw full-band (48 kHz) waveforms into continuous latents.

We have also released a technical report: KVAE: Family of Tokenizers for Multimodal Generative Models.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@yiyixuxu, @asomoza hello guys, could you please check the PR? Thank you!

[ivnvalex]

Self-review — AutoencoderKLKVAEAudio

Ran the self-review skill against the full main...HEAD diff (9 files, 747 insertions), re-reading
.ai/review-rules.md, .ai/models.md, .ai/testing.md, and .ai/skills/model-integration/pitfalls.md
fresh rather than from memory.

Blocking issues

None found.

Non-blocking issues

1. Non-contiguous tensor feeds into in_proj after attention — performance, not correctness — FIXED
src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py, _encode(). Per .ai/models.md
gotcha #7: "prefer calling contiguous() on the output tensor to maintain performance" — the
transpose after self.attn(...) left a non-contiguous tensor feeding directly into in_proj (a
Conv1d). Live code, not a hypothetical: the real kandinskylab/KVAE-Audio checkpoint has
use_attn=True, so every real inference call hit this path. Fixed by adding .contiguous() after the
final transpose. Correctness was unaffected either way (verified via parity check before and after:
2.289e-05 / 1.492e-04, bit-identical to the pre-fix numbers) — this is a pure memory-layout fix, no
behavior change. Re-ran the full test suite (28 passed), check_copies, and make quality after
applying it — all clean.

2. Reuse of the generic Attention class — worth reviewer sign-off
Same file, __init__. .ai/models.md's "Attention pattern" section says attention "must" get a bespoke
Attention+AttnProcessor pair defined in-file. I reused the generic Attention class from
attention_processor.py instead, matching AutoencoderKL's own mid-block self-attention (no
RoPE/masking, so the bespoke-processor rule's rationale doesn't apply) — but the rule's text doesn't
explicitly carve out that exception, so flagging for a maintainer to confirm rather than assuming.

3. sample_rate mismatch-raises path is untested
Same file, encode()/forward(). The ValueError branch when sample_rate != self.config.sample_rate
is reachable but never exercised by test_models_autoencoder_kl_kvae_audio.py (get_dummy_inputs()
never passes sample_rate). Low risk (simple guard), but a one-line pytest.raises test would close
the gap.

Dead code (advisory)

No pipeline is introduced in this PR (standalone VAE addition), so the usual "trace from pipeline
__call__" step doesn't apply — checked reachability against the real checkpoint's published
config.json instead.

path	Status	Reason
use_attn=True / self.attn block	Used	Real checkpoint's config.json sets use_attn: true — confirmed live
attn_num_heads (defaults to 8)	Used	Not overridden by the real checkpoint, but actively consumed since use_attn=True
latent_dim is None branch (__init__)	Partially used	Real checkpoint's config sets latent_dim: 2048 explicitly, so loading it skips this branch; only exercised via direct instantiation (which the test suite does cover)

Summary — READY

No blocking issues. Numerics were independently verified against the reference kvae_1d.py (not
shipped — internal parity check) at ~1.5e-4 on the deterministic path, and strict=True state-dict
loading round-trips against the real checkpoint.

• Fixed: item Add glide modeling files #1 (missing .contiguous() after attention).
• Leave for actual review: items + cosine schedule and unet config #2 (Attention-reuse convention) and make from hub import work #3 (untested validation
  branch) above.

[ivnvalex]

Hi @yiyixuxu, @asomoza, just a friendly reminder about the PR 👉👈

[ivnvalex]

Hey guys @yiyixuxu @asomoza any advice on merging the PR? Thx

[leffff]

Dear
@yiyixuxu @osanseviero the KVAE model line has a huge amount of downloads

We are among a small circle of labs that develops their own VAEs and ours are truely good.
This PR ads audio VAE. Please Help us popularize our Audio VAE!
I appreciate your work. Please help us get this PR done after 2 months

[DN6]

LGTM. Requesting a small change to the attention module and I have a question about the sample_rate arg.

[DN6]

@ivnvalex Could we just define a dedicated attention class and processor for this model please.

[DN6]

@ivnvalex Could we just define a dedicated attention class and processor for this model please. Similar to

diffusers/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py

Line 341 in 916a579

class MiniMaxH3AudioCausalAttention(nn.Module, AttentionModuleMixin):

[DN6]

Is sample_rate necessary? It doesn't seem to be used anywhere beyond this raise?