Fix dtype mismatch in AudioVAE.encode for bf16/fp16 audio VAE

[Vek-Snap]

Greetings! First time contributor here; grateful for the work your team does.
While building some experimental workflows, I encountered a minor enhancement
recommendation which may assist those with hardware that is a little aged. I'd be
happy to adjust anything to match your conventions. Just let me know.

What this does:

AudioVAE.encode() can raise a dtype mismatch when the audio VAE's weights are not
float32. The mel spectrogram produced by AudioPreprocessor.waveform_to_mel() is
always float32 (coming from torchaudio's MelSpectrogram + log), but if the
encoder is loaded in bf16/fp16, feeding the float32 mel into the encoder throws:

RuntimeError: Input type (c10::Float) and weight type (c10::BFloat16) should be the same

Tiny Fix:

Cast the mel to the encoder's actual parameter dtype right before encoding. This is a
no-op when the VAE is float32, and resolves the mismatch when it isn't:

encoder_dtype = next(self.autoencoder.encoder.parameters()).dtype
mel_spec = mel_spec.to(dtype=encoder_dtype)

Why encode() instead of decode():

decode()'s input (latents) originates from inside the model, so it already carries the
model dtype. encode()'s input (mel) originates outside the model from torchaudio
DSP, which is fixed at float32. So encode() is the only path exposed to the mismatch.

How this surfaced for me:

Running the LTX-2.3 audio VAE in bf16 (for VRAM headroom on a 16 GB card).

Option:

If you'd prefer the cast handled at the VAE dtype-loading layer rather than inside
encode(), I can move it there, just let me know what fits your codebase best.

[github-actions]

✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged.
_{Posted by the CLA Assistant Lite bot.}

[Vek-Snap]

I have read and agree to the Contributor License Agreement

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 8b6f6b30-fdee-4cde-aece-d5341934c1da

📥 Commits

Reviewing files that changed from the base of the PR and between 7cf4e78 and ad3797a.

📒 Files selected for processing (1)

• comfy/ldm/lightricks/vae/audio_vae.py

📜 Recent review details

⚠️ CI failures not shown inline (2)

GitHub Actions: CLA Assistant / cla-assistant: Fix dtype mismatch in AudioVAE.encode for bf16/fp16 audio VAE

Conclusion: failure

View job details

##[group]Run contributor-assistant/github-action@ca4a40a7d1004f18d9960b404b97e5f30a505a08
 with:
   path-to-document: https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md
   remote-organization-name: comfy-org
   remote-repository-name: comfy-cla
   path-to-signatures: signatures/cla.json
   branch: main
   allowlist: action@github.com,actions-user,ampagent,claude,comfy-pr-bot,GitHub Action,github-actions,github-actions[bot],Glary Bot,Glary-Bot,*[bot],web-flow
   custom-notsigned-prcomment: 🎉 Thank you for your contribution, we really appreciate it! 🎉
Like many open source projects, we require contributors to sign our [Contributor License Agreement (CLA)](https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md). A CLA makes the ownership of contributions explicit, so contributors and the project share a clear understanding of how the code can be used. By signing, you:
- Confirm that you own your contribution.
- Keep the right to reuse your own code.
- Grant us a copyright license to include and share it within our projects.
CLAs are standard practice across major open source projects including those under the Apache Software Foundation and the Linux Foundation. Ours is based on the Apache Software Foundation's CLA. Most importantly, it would enable us to relicense the project under a more permissive license in the future, giving the project and its community greater flexibility.
✍ **To sign, please post a new comment on this PR with exactly the following text:** ✍
   custom-pr-sign-comment: I have read and agree to the Contributor License Agreement
   custom-allsigned-prcomment: ✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged.
   use-dco-flag: false
   lock-pullrequest-aftermerge: true
   suggest-recheck: true
 env:
   GITHUB_***REDACTED***
   PERSONAL_ACCESS_***REDACTED***
 ##[endgroup]
 CLA Assistant GitHub Action bot has started the process
 (node:2157) [DEP0040] DeprecationWarning: The `punycode` module ...

---

GitHub Actions: CLA Assistant / 0_cla-assistant.txt: Fix dtype mismatch in AudioVAE.encode for bf16/fp16 audio VAE

Conclusion: failure

View job details

##[group]Run contributor-assistant/github-action@ca4a40a7d1004f18d9960b404b97e5f30a505a08
 with:
   path-to-document: https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md
   remote-organization-name: comfy-org
   remote-repository-name: comfy-cla
   path-to-signatures: signatures/cla.json
   branch: main
   allowlist: action@github.com,actions-user,ampagent,claude,comfy-pr-bot,GitHub Action,github-actions,github-actions[bot],Glary Bot,Glary-Bot,*[bot],web-flow
   custom-notsigned-prcomment: 🎉 Thank you for your contribution, we really appreciate it! 🎉
Like many open source projects, we require contributors to sign our [Contributor License Agreement (CLA)](https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md). A CLA makes the ownership of contributions explicit, so contributors and the project share a clear understanding of how the code can be used. By signing, you:
- Confirm that you own your contribution.
- Keep the right to reuse your own code.
- Grant us a copyright license to include and share it within our projects.
CLAs are standard practice across major open source projects including those under the Apache Software Foundation and the Linux Foundation. Ours is based on the Apache Software Foundation's CLA. Most importantly, it would enable us to relicense the project under a more permissive license in the future, giving the project and its community greater flexibility.
✍ **To sign, please post a new comment on this PR with exactly the following text:** ✍
   custom-pr-sign-comment: I have read and agree to the Contributor License Agreement
   custom-allsigned-prcomment: ✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged.
   use-dco-flag: false
   lock-pullrequest-aftermerge: true
   suggest-recheck: true
 env:
   GITHUB_***REDACTED***
   PERSONAL_ACCESS_***REDACTED***
 ##[endgroup]
 CLA Assistant GitHub Action bot has started the process
 (node:2157) [DEP0040] DeprecationWarning: The `punycode` module ...

🧰 Additional context used

📓 Path-based instructions (3)

**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

**/*.py: Keep imports at module scope; avoid inline imports unless they are already part of an established optional-backend probe or are needed to avoid an import cycle.
Do not add unnecessary try/except blocks; use them only for optional dependency, platform, or backend capability detection when the program has a useful fallback, and prefer specific exception types when changing new code.
Remove workarounds for PyTorch versions that ComfyUI no longer officially supports; if a workaround does not have a comment naming the exact supported PyTorch version(s), remove it.
Let unsupported model formats, invalid quantization metadata, and bad states fail with clear errors instead of silently producing lower quality output.
Match the existing local style in the file you edit; this codebase tolerates long lines, simple helper functions, module-level state, and direct tensor operations when they make the code easier to follow.
Keep comments sparse and useful; strip useless comments that restate the code or describe obvious behavior, while allowing short TODOs that name the concrete missing follow-up.
Treat dtype, device placement, VRAM usage, and offloading behavior as core correctness concerns; check CPU, CUDA, ROCm, MPS, DirectML, XPU, NPU, and low VRAM implications when touching shared execution or loading code.
Prefer native ComfyUI formats and existing quantization/offload helpers over adding parallel code paths; use comfy.quant_ops, comfy.model_management, comfy.memory_management, comfy.pinned_memory, comfy_aimdo, and comfy-kitchen helpers where they already solve the problem.
Use optimized comfy-kitchen ops where they improve performance without changing expected dtype, device, memory, or interface behavior.
All models should use the optimized attention function selected by ComfyUI; higher-level code must not inspect function identity, names, modules, or implementation details to decide behavior.
Apply the same opacity rule to similar patterns beyond...

Files:

• comfy/ldm/lightricks/vae/audio_vae.py

**

⚙️ CodeRabbit configuration file

**: ## Engineering Style

• Keep changes small and direct. Most fixes should touch the narrowest code path
  that explains the bug, performance issue, dtype issue, model-format issue, or
  user-facing behavior.
• Change the least amount of files possible. A change that touches many files is
  more likely to be a bad change than a good one unless the broader scope is
  directly required.
• Prefer practical fixes over broad architecture work. Add abstractions only
  when they remove real repeated logic or match an existing ComfyUI pattern.
• Prefer fewer dependencies. Do not add new dependencies to ComfyUI unless they
  are absolutely necessary.
• Delete obsolete code aggressively when newer infrastructure makes it useless.
  Remove dead fallbacks, migration paths, unused options, debug prints, and
  compatibility branches that are no longer needed. Do not leave dead branches,
  unreachable code, or functions that are never called. If code is not
  necessary for the current behavior, remove it.
• Revert or disable problematic behavior quickly when it breaks users. It is
  better to remove a broken feature path than keep a complicated partial fix.
• Preserve existing APIs, node names, model-loading behavior, file layout, and
  workflow compatibility unless the change is explicitly about replacing them.
• Code must look hand-written for this repository. Changes that read like
  generic AI-generated code will be rejected automatically: unnecessary helper
  layers, vague names, boilerplate comments, defensive branches without a real
  failure mode, broad rewrites, or code that ignores the local style.

Architecture Boundaries

• Keep each layer focused on the concepts it owns. Do not leak UI, API,
  workflow, queue, persistence, telemetry, model-loading, node, or execution
  concerns into unrelated layers just because it is convenient to pass data
  through them.
• Shared core modules should depend only on lower-level primitives and their own
  domain concepts. Highe...

Files:

• comfy/ldm/lightricks/vae/audio_vae.py

---

⚙️ CodeRabbit configuration file

**: IMPORTANT: Only comment on issues directly introduced by this PR's code changes.
Treat AGENTS.md as mandatory repository policy, not optional style guidance.
Flag PR changes that violate AGENTS.md even when the code is otherwise functional.
In particular, enforce architecture boundaries, dtype/device/memory rules,
interface contracts, import style, no unnecessary try/except blocks, no inline
imports, no outbound internet paths in core ComfyUI, and narrow scoped fixes.
Prefer direct findings over suggestions when a rule is violated. Only ignore
AGENTS.md when it clearly conflicts with a newer explicit maintainer instruction
in the PR.
Do NOT flag pre-existing issues in code that was merely moved, re-indented,
de-indented, or reformatted without logic changes. If code appears in the diff
only due to whitespace or structural reformatting (e.g., removing a with: block),
treat it as unchanged. Contributors should not feel obligated to address
pre-existing issues outside the scope of their contribution.

Files:

• comfy/ldm/lightricks/vae/audio_vae.py

comfy/**

⚙️ CodeRabbit configuration file

comfy/**: Core ML/diffusion engine. Focus on:

• Backward compatibility (breaking changes affect all custom nodes)
• Memory management and GPU resource handling
• Performance implications in hot paths
• Thread safety for concurrent execution

Files:

• comfy/ldm/lightricks/vae/audio_vae.py

🧠 Learnings (2)

📚 Learning: 2026-02-21T14:01:41.482Z

Learnt from: pythongosssss
Repo: Comfy-Org/ComfyUI PR: 12555
File: comfy_extras/nodes_glsl.py:719-724
Timestamp: 2026-02-21T14:01:41.482Z
Learning: In PyOpenGL, bare Python scalars can be accepted for 1-element array parameters by NumberHandler. This means you can pass an int/float directly to OpenGL texture deletion (e.g., glDeleteTextures(tex)) without wrapping in a list. Verify function-specific expectations and ensure types match what the OpenGL call expects; use explicit lists only when the API requires an array.

Applied to files:

• comfy/ldm/lightricks/vae/audio_vae.py

📚 Learning: 2026-05-13T12:31:45.069Z

Learnt from: rattus128
Repo: Comfy-Org/ComfyUI PR: 13802
File: comfy/pinned_memory.py:19-30
Timestamp: 2026-05-13T12:31:45.069Z
Learning: When reviewing code that uses comfy/pinned_memory.py’s `HostBuffer.extend(size=..., reallocate=...)`: by default (`reallocate` is not True / False), `extend(size=...)` is a *relative increment* that grows the buffer by `size` bytes—so slicing like `[offset:offset+size]` after `hostbuf.extend(size=size)` is correct and the argument should not be rewritten to `offset + size`. Only in the single-segment reallocation mode (`reallocate=True`, e.g., as used by `resize_pin_buffer()` in `comfy/model_management.py`) should `size` be treated as an *absolute target* and the call/arguments should be checked accordingly.

Applied to files:

• comfy/ldm/lightricks/vae/audio_vae.py

🔇 Additional comments (1)

comfy/ldm/lightricks/vae/audio_vae.py (1)

157-160: LGTM!

---

📝 Walkthrough

Walkthrough

This change modifies AudioVAE.encode in comfy/ldm/lightricks/vae/audio_vae.py to cast the mel spectrogram tensor to the dtype of the encoder's parameters before passing it to the encoder, rather than passing the mel spectrogram unchanged.

Changes

File	Change Summary
comfy/ldm/lightricks/vae/audio_vae.py	Cast mel_spec to the encoder's parameter dtype prior to encoding

Sequence Diagram(s)

Not applicable — the change is a small dtype-casting adjustment within an existing function and does not alter control flow or component interactions.

Estimated code review effort: 1 (~2 minutes)

Related issues: None provided.

Related PRs: None provided.

Suggested labels: None provided.

Suggested reviewers: None provided.

Poem: A rabbit hopped by the audio stream, cast a dtype to match the dream, mel spectrograms now align, with encoder weights, all in a line, a tiny fix, a tidy gleam.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the main change: fixing a dtype mismatch in AudioVAE.encode for bf16/fp16 models.
Description check	✅ Passed	The description matches the implemented fix and explains the dtype mismatch and casting change accurately.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}