You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
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:
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Adds
KVAE-AudioasAutoencoderKLKVAEAudio, a 1D convolutional audio VAE, completing the KVAE tokenizers family from the Kandinsky Lab team in Diffusers alongside the already-mergedAutoencoderKLKVAE(image) andAutoencoderKLKVAEVideo(video).It loads the checkpoint from
kandinskylab/KVAE-Audioand compresses/reconstructs raw full-band (48 kHz) waveforms into continuous latents.Before submitting
self-reviewskill on the diff?documentation guidelines, and
here are tips on formatting docstrings.
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!