Skip to content

Fix AuraFlow VAE dtype mismatch on pipeline reuse#14184

Open
IvenHsu01 wants to merge 1 commit into
huggingface:mainfrom
IvenHsu01:fix-auraflow-vae-dtype-pipeline-reuse
Open

Fix AuraFlow VAE dtype mismatch on pipeline reuse#14184
IvenHsu01 wants to merge 1 commit into
huggingface:mainfrom
IvenHsu01:fix-auraflow-vae-dtype-pipeline-reuse

Conversation

@IvenHsu01

Copy link
Copy Markdown

What does this PR do?

Fixes a VAE dtype-mismatch RuntimeError in AuraFlowPipeline that occurs when
the same pipeline instance is called more than once.

upcast_vae() upcasts the whole VAE to float32 in place, so on the second call
needs_upcasting is False and the latents.to(...) cast (guarded by that
if) is skipped, feeding fp16 latents to the fp32 VAE. This applies the same
pattern already merged for pixart_sigma in #8391 — cast the latents to the VAE
dtype inline at the decode call so it always runs.

             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]

Companion PR to #14183, which has the full analysis and verification.
Happy to adjust the scope or approach based on maintainer feedback there.

Coordination

Tests run

Applied this exact change, no other patches, ran warmup + 1 run (the 2nd call
triggers the bug) at 512x512, 50 steps, guidance 3.5, seed 42:

  • NVIDIA RTX 5090 (CUDA): warmup 18.51s, 2nd call 7.31s — Pass, coherent image
  • AMD gfx1151 (ROCm): warmup 64.23s, 2nd call 63.25s — Pass, coherent image

Before submitting

Who can review?

@yiyixuxu @sayakpaul

@github-actions github-actions Bot added pipelines size/S PR with diff < 50 LOC labels Jul 14, 2026
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 huggingface#8391.
@sayakpaul sayakpaul requested a review from dg845 July 14, 2026 03:12
@github-actions

Copy link
Copy Markdown
Contributor

Hi @IvenHsu01, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pipelines size/S PR with diff < 50 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant