Comfy-Org / Comfy-Org/ComfyUI

AudioVAE encode/decode crashes with dtype mismatch: "Input type (float) and bias type (c10::BFloat16) should be the same"

Open Beginner friendly
#14,811 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.