fix using fp32 VAE

This commit is contained in:
kijai
2024-03-04 11:42:52 +02:00
parent 637a3f487a
commit a3794f3a38
+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):
#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)