Update nodes.py

This commit is contained in:
Kijai 2024-03-15 16:51:05 +02:00
parent 656feb2c14
commit 95ae8b067a

View File

@ -3707,13 +3707,17 @@ class LoadResAdapterNormalization:
FUNCTION = "load_res_adapter"
def load_res_adapter(self, model, resadapter_path):
resadapter_path = f"{folder_paths.get_folder_paths('checkpoints')[0]}/{resadapter_path}"
if os.path.exists(resadapter_path):
print("ResAdapter: Checking ResAdapter path")
resadapter_full_path = folder_paths.get_full_path("checkpoints", resadapter_path)
if not os.path.exists(resadapter_full_path):
raise Exception("Invalid model path")
else:
print("ResAdapter: Loading ResAdapter normalization weights")
prefix_to_remove = 'diffusion_model.'
model_clone = model.clone()
norm_state_dict = comfy.utils.load_torch_file(resadapter_path)
norm_state_dict = comfy.utils.load_torch_file(resadapter_full_path)
new_values = {key[len(prefix_to_remove):]: value for key, value in norm_state_dict.items() if key.startswith(prefix_to_remove)}
print("ResAdapter: Attempting to add patches with ResAdapter weights")
try:
for key in model.model.diffusion_model.state_dict().keys():
if key in new_values:
@ -3721,9 +3725,11 @@ class LoadResAdapterNormalization:
new_tensor = new_values[key].to(model.model.diffusion_model.dtype)
if original_tensor.shape == new_tensor.shape:
model_clone.add_object_patch(f"diffusion_model.{key}.data", new_tensor)
else:
print("ResAdapter: No match for key: ",key)
except:
raise Exception("Could not patch model, this way of patching was added to ComfyUI on March 3rd 2024, is your ComfyUI up to date?")
print("Added resnet normalization patch.")
print("ResAdapter: Added resnet normalization patches")
return (model_clone, )
class Superprompt: