Flux Vae broke for float16, force bfloat16 or float32 were compatible #7213
Conversation
Vargol
requested review from
lstein,
blessedcoolant,
brandonrising,
RyanJDick and
hipsterusername
as code owners
github-actions
bot
added
python
hipsterusername
force-pushed
the
flux_vae_fp16_broke
branch
from
81e3001
to
496b02a
Compare
| # VAE is broken in float16, use same logic in model loading to pick bfloat16 or float32 | ||
| if x.dtype == torch.float16: | ||
| try: | ||
| x = x.to(torch.bfloat16) | ||
| except TypeError: | ||
| x = x.to(torch.float32) |
There was a problem hiding this comment.
The AutoEncoder class should not be responsible for casting input dtypes - it should simply raise if an incompatible dtype is passed, which it already does.
The casting should happen in the caller (flux_vae_encode.py / flux_vae_decode.py). And rather than mirroring the model loading logic, we should simply cast the input to match the dtype of the model.
There was a problem hiding this comment.
Is there a sensible way to get the vae dtype, as the original code just grabs the default / config dtype rather than grabbing the models and I can't see a sensible variable in AutoEncoder.
At the moment I'm looking at crap code like
vae_dtype = next(iter(vae.state_dict().items()))[1].dtypeor picking a state_dict key and getting the value, I'm assuming just grabbing the first value is more efficient ATM .
Added a bool to allow the node user to add noise in to initial latents (default) or to leave them alone.
github-actions
bot
added
invocations
Summary
The Flux VAE, like many VAEs, is broken if run using float16 inputs returning black images due to NaNs
This will fix the issue by forcing the VAE to run in bfloat16 or float32 were compatible
Related Issues / Discussions
Fix for issue #7208
QA Instructions
Tested on MacOS, VAE works with float16 in the invoke.yaml and left to default.
I also briefly forced it down the float32 route to check that to.
Needs testing on CUDA / ROCm
Merge Plan
It should be a straight forward merge,