github/InvokeAI
1
1
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI.git synced 2026-02-10 17:35:03 -05:00
Code Issues Packages Projects Releases Wiki Activity

Same issue affects image2image, so do the same again

Browse Source
This commit is contained in:
David Burnett
2024-10-28 22:32:55 +00:00
committed by Kent Keirsey
parent 7b5efc2203
commit 496b02a3bc
1 changed files with 6 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
invokeai/backend/flux/modules/autoencoder.py
View File

@@ -312,6 +312,12 @@ class AutoEncoder(nn.Module):
Tensor: Encoded latent tensor. Shape: (batch_size, z_channels, latent_height, latent_width).
"""
# 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)
z = self.reg(self.encoder(x), sample=sample, generator=generator)
z = self.scale_factor * (z - self.shift_factor)
return z