Run FLUX VAE decoding in the user's preferred dtype rather than float32. Tested, and seems to work well at float16.

invokeai/app/invocations/flux_text_to_image.py

@@ -131,10 +131,7 @@ class FluxTextToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
         vae_info = context.models.load(self.vae.vae)
         with vae_info as vae:
             assert isinstance(vae, AutoEncoder)
-            # TODO(ryand): Test that this works with both float16 and bfloat16.
-            # with torch.autocast(device_type=latents.device.type, dtype=torch.float32):
-            vae.to(torch.float32)
-            latents.to(torch.float32)
+            latents = latents.to(dtype=TorchDevice.choose_torch_dtype())
             img = vae.decode(latents)
 
         img = img.clamp(-1, 1)

invokeai/backend/model_manager/load/model_loaders/flux.py

@@ -65,6 +65,7 @@ class FluxVAELoader(ModelLoader):
             model = AutoEncoder(params)
             sd = load_file(model_path)
             model.load_state_dict(sd, assign=True)
+            model.to(dtype=self._torch_dtype)
 
         return model