fix(aura_flow): cast latents to VAE dtype on every decode#14219
Open
sergioperezcheco wants to merge 1 commit into
Open
fix(aura_flow): cast latents to VAE dtype on every decode#14219sergioperezcheco wants to merge 1 commit into
sergioperezcheco wants to merge 1 commit into
Conversation
github-actions
Bot
added
tests
pipelines
fixes-issue
size/S
sergioperezcheco
force-pushed
the
fix/auraflow-vae-dtype-reuse
branch
from
July 18, 2026 04:19
1f82193 to
8821e8d
Compare
AuraFlowPipeline guarded the latents dtype cast behind needs_upcasting, so after the first call upcast_vae() left the VAE in fp32 and the second call fed fp16 latents into it, raising RuntimeError. Cast latents to the VAE dtype inline at the decode call (matching the pixart_sigma fix in huggingface#8391) so the cast runs unconditionally. Added a regression test. Fixes huggingface#14183
sergioperezcheco
force-pushed
the
fix/auraflow-vae-dtype-reuse
branch
from
July 18, 2026 17:43
8821e8d to
86dc582
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Calling AuraFlowPipeline a second time on the same instance raises RuntimeError: Input type (c10::Half) and bias type (float) should be the same. The decode block only cast latents to the VAE dtype when needs_upcasting was True, but upcast_vae() moves the whole VAE to fp32 in place on the first call, so on the second call needs_upcasting is False and fp16 latents go straight into the now-fp32 VAE. I moved the latents.to(self.vae.dtype) inline at the decode call so it always runs, matching what pixart_sigma already does (the fix from #8391 for the same class of bug). Added a regression test that simulates the post-upcast state and confirms the cast runs unconditionally. Fixes #14183.