From d297be4443a883916db7c0f8dd7219950a075c37 Mon Sep 17 00:00:00 2001 From: IvenHsu01 <228416682+IvenHsu01@users.noreply.github.com> Date: Tue, 14 Jul 2026 11:03:42 +0800 Subject: [PATCH] Fix AuraFlow VAE dtype mismatch on pipeline reuse upcast_vae() upcasts the whole VAE to float32 in place, so on a second __call__ needs_upcasting is False and the latents cast guarded by that branch is skipped, feeding fp16 latents to the fp32 VAE. Cast the latents to the VAE dtype inline at the decode call so it always runs, matching the fix applied to pixart_sigma in #8391. --- src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py b/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py index e6f1c2430143..30a12c2ba41f 100644 --- a/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py +++ b/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py @@ -659,8 +659,7 @@ def __call__( needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast if needs_upcasting: self.upcast_vae() - latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype) - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] + image = self.vae.decode(latents.to(self.vae.dtype) / self.vae.config.scaling_factor, return_dict=False)[0] image = self.image_processor.postprocess(image, output_type=output_type) # Offload all models