AudioVAE encode/decode crashes with dtype mismatch: "Input type (float) and bias type (c10::BFloat16) should be the same"
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
## Description
`AudioVAE.encode()` in `comfy/ldm/lightricks/vae/audio_vae.py` computes a mel spectrogram via `torchaudio.transforms.MelSpectrogram`, which always returns float32, then feeds it directly into `self.autoencoder.encode(mel_spec)` without casting to the model's dtype. When the audio VAE is loaded with bfloat16 weights (which is the default — see the log line `VAE load device: cuda:0, offload device: cpu, dtype: torch.bfloat16`), the conv2d inside the encoder crashes because the input is float32 but the weight/bias is bfloat16.
The same issue exists on the decode path — incoming latents may be float32 while the decoder weights are bf16.
## Error
```
RuntimeError: Input type (float) and bias type (c10::BFloat16) should be the same
```
Full traceback:
```
File "/ComfyUI/comfy_extras/nodes_lt_audio.py", line 57, in execute
return super().execute(audio_vae, audio)
File "/ComfyUI/comfy_extras/nodes_audio.py", line 92, in execute
t = vae.encode(waveform.movedim(1, -1))
File "/ComfyUI/comfy/sd.py", line 1169, in encode
model_management.raise_non_oom(e)
File "/ComfyUI/comfy/sd.py", line 1162, in encode
out = self.first_stage_model.encode(pixels_in)
File "/ComfyUI/comfy/ldm/lightricks/vae/audio_vae.py", line 157, in encode
latents = self.autoencoder.encode(mel_spec)
File "/ComfyUI/comfy/ldm/lightricks/vae/causal_audio_autoencoder.py", line 896, in encode
return self.encoder(x)
...
File "/torch/nn/modules/conv.py", line 548, in _conv_forward
return F.conv2d(
RuntimeError: Input type (float) and bias type (c10::BFloat16) should be the same
```
## Root cause
`AudioPreprocessor.waveform_to_mel()` creates a `MelSpectrogram` transform and runs it on the waveform. The output is always float32 regardless of the input dtype or model dtype. This float32 mel spec is then passed directly to `self.autoencoder.encode()` whose conv layers have bfloat16 weights.
The video VAE doesn't hit this because `comfy.sd.VAE.encode()` has dtype casting logic via `cast_bias_weight`, but the audio VAE's internal mel computation bypasses that path entirely.
## Fix
Cast the mel spectrogram to the autoencoder's dtype before encoding, and cast latents before decoding:
```python
# In AudioVAE.encode(), after mel computation:
mel_spec = mel_spec.to(self.autoencoder.dtype if hasattr(self.autoencoder, 'dtype') else next(self.autoencoder.parameters()).dtype)
# In AudioVAE.decode(), at the start:
model_dtype = self.autoencoder.dtype if hasattr(self.autoencoder, 'dtype') else next(self.autoencoder.parameters()).dtype
latents = latents.to(model_dtype)
```
## Environment
- ComfyUI: v0.27.0-25-g51bf508a
- PyTorch: 2.11.0+cu130
- CUDA: 13.0
- Python: 3.12.11
- GPU: RTX 5000 Ada (SM 8.9)
- comfy-kitchen: 0.2.16
- comfy-aimdo: 0.4.10
- OS: Linux (Kali 6.19.14)
Reproduced with LTX-2.3 audio VAE encode using both the stock `LTXVAudioVAEEncode` node and `VAEEncodeAudio` with an LTX audio VAE checkpoint loaded via `LTXVAudioVAELoader`.
## Additional note
The `ComfyUI-LTXVideo` custom node's `LowVRAMAudioVAELoader` (`low_vram_loaders.py:96`) also has a stale constructor call — `AudioVAE(sd, metadata)` — but the core `AudioVAE.__init__` signature was changed to `AudioVAE(metadata)` (state dict is now loaded via the standard `comfy.sd.VAE` wrapper). This was noted by fredbliss on the Banodoco Discord back in April. The fix is to use `comfy.sd.VAE(sd=sd, metadata=metadata)` like the core `LTXVAudioVAELoader` does.
Contributor guide
Research direction
Start in comfy/ldm/lightricks/vae/audio_vae.py by reading AudioVAE.encode() and decode(), then trace the mel spectrogram and latent values into the autoencoder. Ensure both paths use the autoencoder's dtype, and verify that bfloat16 audio VAE encode/decode no longer raises the reported mismatch.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 78/100