nodes/kijai-ComfyUI-SUPIR
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix using fp32 VAE

Browse Source
This commit is contained in:
kijai
2024-03-04 11:42:52 +02:00
parent 637a3f487a
commit a3794f3a38
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
SUPIR/models/SUPIR_model.py
+1 -1
View File
@@ -54,7 +54,7 @@ class SUPIRModel(DiffusionEngine):
def encode_first_stage_with_denoise(self, x, use_sample=True, is_stage1=False): def encode_first_stage_with_denoise(self, x, use_sample=True, is_stage1=False):
#with torch.autocast(device, dtype=self.ae_dtype): #with torch.autocast(device, dtype=self.ae_dtype):
self.first_stage_model.to(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(): with torch.autocast(comfy.model_management.get_autocast_device(device), dtype=self.ae_dtype) if autocast_condition else nullcontext():
if is_stage1: if is_stage1:
h = self.first_stage_model.denoise_encoder_s1(x) h = self.first_stage_model.denoise_encoder_s1(x)