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

2024-08-30 20:32:17 +00:00 · 2024-08-22 18:16:43 +00:00
parent 14ab339b33
commit a0bf20bcee
2 changed files with 2 additions and 4 deletions
--- a/invokeai/backend/model_manager/load/model_loaders/flux.py
+++ b/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