From a3794f3a384e3f502597170a2c3a3ece93a17b35 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Mon, 4 Mar 2024 11:42:52 +0200 Subject: [PATCH] fix using fp32 VAE --- SUPIR/models/SUPIR_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/SUPIR/models/SUPIR_model.py b/SUPIR/models/SUPIR_model.py index 78c46b2..23ff9e6 100644 --- a/SUPIR/models/SUPIR_model.py +++ b/SUPIR/models/SUPIR_model.py @@ -54,7 +54,7 @@ class SUPIRModel(DiffusionEngine): 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) + autocast_condition = (self.model.dtype == torch.float16 or self.model.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: h = self.first_stage_model.denoise_encoder_s1(x)