From 093baa095745606e815ceaf551def258e9968e29 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Fri, 1 Mar 2024 15:18:01 +0200
Subject: [PATCH] Attempt to fix fp32 VAE

---
 SUPIR/models/SUPIR_model.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/SUPIR/models/SUPIR_model.py b/SUPIR/models/SUPIR_model.py
index f76eeca..2c18a57 100644
--- a/SUPIR/models/SUPIR_model.py
+++ b/SUPIR/models/SUPIR_model.py
@@ -53,6 +53,7 @@ class SUPIRModel(DiffusionEngine):
     @torch.no_grad()
     def encode_first_stage_with_denoise(self, x, use_sample=True, is_stage1=False):
         #with torch.autocast(device, dtype=self.ae_dtype):
+        self.first_stage_model.to(self.ae_dtype)
         autocast_condition = (self.ae_dtype == torch.float16 or self.ae_dtype == torch.bfloat16) and not comfy.model_management.is_device_mps(device)
         with torch.autocast(comfy.model_management.get_autocast_device(device), dtype=self.ae_dtype) if autocast_condition else nullcontext():
             if is_stage1: