Same issue affects image2image, so do the same again

invoke-ai · Oct 28, 2024 · 81e3001 · 81e3001
1 parent fa0612f
commit 81e3001
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/invokeai/backend/flux/modules/autoencoder.py b/invokeai/backend/flux/modules/autoencoder.py
@@ -312,6 +312,12 @@ def encode(self, x: Tensor, sample: bool = True, generator: torch.Generator | No
             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