Skip to content

[feat] add KVAE-Audio autoencoder - #14384

Open
ivnvalex wants to merge 2 commits into
huggingface:mainfrom
ivnvalex:feat/autoencoder/kvae-audio
Open

[feat] add KVAE-Audio autoencoder#14384
ivnvalex wants to merge 2 commits into
huggingface:mainfrom
ivnvalex:feat/autoencoder/kvae-audio

Conversation

@ivnvalex

@ivnvalex ivnvalex commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

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.

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

ivnvalex commented Aug 4, 2026

Copy link
Copy Markdown
Contributor Author

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.

@github-actions github-actions Bot added documentation Improvements or additions to documentation models tests utils size/L PR with diff > 200 LOC labels Aug 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation models size/L PR with diff > 200 LOC tests utils

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant