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

Update nodes_v2.py

Browse Source
This commit is contained in:
Kijai
2024-03-25 18:06:11 +02:00
parent eb7d1e31ae
commit 203bc93933
1 changed files with 14 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes_v2.py
+14 -16
View File
@@ -645,19 +645,8 @@ class SUPIR_model_loader:
pbar = comfy.utils.ProgressBar(5)
self.model = instantiate_from_config(config.model).cpu()
self.model.model.dtype = dtype
pbar.update(1)
try:
print(f'Attempting to load SUPIR model: [{SUPIR_MODEL_PATH}]')
supir_state_dict = load_state_dict(SUPIR_MODEL_PATH)
self.model.load_state_dict(supir_state_dict, strict=False)
if fp8_unet:
self.model.model.to(torch.float8_e4m3fn)
else:
self.model.model.to(dtype)
del supir_state_dict
pbar.update(1)
except:
raise Exception("Failed to load SUPIR model")
try:
print(f"Attempting to load SDXL model: [{SDXL_MODEL_PATH}]")
sdxl_state_dict = load_state_dict(SDXL_MODEL_PATH)
@@ -706,11 +695,20 @@ class SUPIR_model_loader:
raise Exception("Failed to load second clip model from SDXL checkpoint")
del sd, clip_g
mm.soft_empty_cache()
#only unets and/or vae to fp8
if fp8_unet:
self.model.model.to(torch.float8_e4m3fn)
try:
print(f'Attempting to load SUPIR model: [{SUPIR_MODEL_PATH}]')
supir_state_dict = load_state_dict(SUPIR_MODEL_PATH)
self.model.load_state_dict(supir_state_dict, strict=False)
if fp8_unet:
self.model.model.to(torch.float8_e4m3fn)
else:
self.model.model.to(dtype)
del supir_state_dict
pbar.update(1)
except:
raise Exception("Failed to load SUPIR model")
mm.soft_empty_cache()
return (self.model, self.model.first_stage_model,)