From ad3797ad4f757852c2ba9ac8929faf0efa761845 Mon Sep 17 00:00:00 2001 From: Refuelko <270010670+Refuelko@users.noreply.github.com> Date: Mon, 6 Jul 2026 23:08:30 -0700 Subject: [PATCH] Fix dtype mismatch in AudioVAE.encode for bf16/fp16 audio VAE --- comfy/ldm/lightricks/vae/audio_vae.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/comfy/ldm/lightricks/vae/audio_vae.py b/comfy/ldm/lightricks/vae/audio_vae.py index dd5320c8f..2e4e91701 100644 --- a/comfy/ldm/lightricks/vae/audio_vae.py +++ b/comfy/ldm/lightricks/vae/audio_vae.py @@ -154,6 +154,10 @@ class AudioVAE(torch.nn.Module): waveform, waveform_sample_rate, device=waveform.device ) + # Cast mel spectrogram to match encoder weight dtype (bf16 VAE + float32 audio input) + encoder_dtype = next(self.autoencoder.encoder.parameters()).dtype + mel_spec = mel_spec.to(dtype=encoder_dtype) + latents = self.autoencoder.encode(mel_spec) posterior = DiagonalGaussianDistribution(latents) latent_mode = posterior.mode()