Spaces:
Running
on
Zero
Running
on
Zero
| #---------------------------------------------------------------------------------------------------------------------# | |
| # Comfyroll Studio custom nodes by RockOfFire and Akatsuzi https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes | |
| # for ComfyUI https://github.com/comfyanonymous/ComfyUI | |
| #---------------------------------------------------------------------------------------------------------------------# | |
| import comfy.sd | |
| import comfy.model_management | |
| import folder_paths | |
| from ..categories import icons | |
| #---------------------------------------------------------------------------------------------------------------------# | |
| # Model Merge Nodes | |
| #---------------------------------------------------------------------------------------------------------------------# | |
| class CR_ModelMergeStack: | |
| def INPUT_TYPES(cls): | |
| checkpoint_files = ["None"] + folder_paths.get_filename_list("checkpoints") | |
| return {"required": {"switch_1": (["Off","On"],), | |
| "ckpt_name1": (checkpoint_files,), | |
| "model_ratio1": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), | |
| "clip_ratio1": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), | |
| # | |
| "switch_2": (["Off","On"],), | |
| "ckpt_name2": (checkpoint_files,), | |
| "model_ratio2": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), | |
| "clip_ratio2": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), | |
| # | |
| "switch_3": (["Off","On"],), | |
| "ckpt_name3": (checkpoint_files,), | |
| "model_ratio3": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), | |
| "clip_ratio3": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), | |
| }, | |
| "optional":{ | |
| "model_stack": ("MODEL_STACK",), | |
| }, | |
| } | |
| RETURN_TYPES = ("MODEL_STACK", "STRING", ) | |
| RETURN_NAMES = ("MODEL_STACK", "show_help", ) | |
| FUNCTION = "list_checkpoints" | |
| CATEGORY = icons.get("Comfyroll/Model Merge") | |
| def list_checkpoints(self, switch_1, ckpt_name1, model_ratio1, clip_ratio1, switch_2, ckpt_name2, model_ratio2, clip_ratio2, switch_3, ckpt_name3, model_ratio3, clip_ratio3, model_stack=None): | |
| # Initialise the list | |
| model_list = list() | |
| if model_stack is not None: | |
| model_list.extend([l for l in model_stack if l[0] != "None"]) | |
| if ckpt_name1 != "None" and switch_1 == "On": | |
| model_list.extend([(ckpt_name1, model_ratio1, clip_ratio1)]), | |
| if ckpt_name2 != "None" and switch_2 == "On": | |
| model_list.extend([(ckpt_name2, model_ratio2, clip_ratio2)]), | |
| if ckpt_name3 != "None" and switch_3 == "On": | |
| model_list.extend([(ckpt_name3, model_ratio3, clip_ratio3)]), | |
| show_help = "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes/wiki/Model-Merge-Nodes#cr-model-stack" | |
| return (model_list, show_help, ) | |
| #---------------------------------------------------------------------------------------------------------------------# | |
| class CR_ApplyModelMerge: | |
| def INPUT_TYPES(s): | |
| merge_methods = ["Recursive", "Weighted"] | |
| return {"required": {"model_stack": ("MODEL_STACK",), | |
| "merge_method": (merge_methods,), | |
| "normalise_ratios": (["Yes","No"],), | |
| "weight_factor":("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), | |
| } | |
| } | |
| RETURN_TYPES = ("MODEL", "CLIP", "STRING", "STRING", ) | |
| RETURN_NAMES = ("MODEL", "CLIP", "model_mix_info", "show_help", ) | |
| FUNCTION = "merge" | |
| CATEGORY = icons.get("Comfyroll/Model Merge") | |
| def merge(self, model_stack, merge_method, normalise_ratios, weight_factor): | |
| # Initialise | |
| sum_clip_ratio = 0 | |
| sum_model_ratio = 0 | |
| model_mix_info = str("Merge Info:\n") | |
| # If no models | |
| if len(model_stack) == 0: | |
| print(f"[Warning] Apply Model Merge: No active models selected in the model merge stack") | |
| return() | |
| # If only one model | |
| if len(model_stack) == 1: | |
| print(f"[Warning] Apply Model Merge: Only one active model found in the model merge stack. At least 2 models are normally needed for merging. The active model will be output.") | |
| model_name, model_ratio, clip_ratio = model_stack[0] | |
| ckpt_path = folder_paths.get_full_path("checkpoints", model_name) | |
| return comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) | |
| # Calculate ratio sums for normalisation | |
| for i, model_tuple in enumerate(model_stack): | |
| model_name, model_ratio, clip_ratio = model_tuple | |
| sum_model_ratio += model_ratio | |
| sum_clip_ratio += clip_ratio | |
| # Do recursive merge loops | |
| model_mix_info = model_mix_info + "Ratios are applied using the Recursive method\n\n" | |
| # Loop through the models and compile the merged model | |
| for i, model_tuple in enumerate(model_stack): | |
| model_name, model_ratio, clip_ratio = model_tuple | |
| ckpt_path = folder_paths.get_full_path("checkpoints", model_name) | |
| merge_model = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) | |
| print(f"Apply Model Merge: Model Name {model_name}, Model Ratio {model_ratio}, CLIP Ratio {clip_ratio}") | |
| if sum_model_ratio != 1 and normalise_ratios == "Yes": | |
| print(f"[Warning] Apply Model Merge: Sum of model ratios != 1. Ratios will be normalised") | |
| # Normalise the ratios | |
| model_ratio = round(model_ratio / sum_model_ratio, 2) | |
| clip_ratio = round(clip_ratio / sum_clip_ratio, 2) | |
| # Weighted merge method | |
| if merge_method == "Weighted": | |
| if i == 1: | |
| # Reassign extra weight to the second model | |
| model_ratio = 1 - weight_factor + (weight_factor * model_ratio) | |
| clip_ratio = 1 - weight_factor + (weight_factor * clip_ratio) | |
| #Clone the first model | |
| if i == 0: | |
| model1 = merge_model[0].clone() | |
| clip1 = merge_model[1].clone() | |
| model_mix_info = model_mix_info + "Base Model Name: " + model_name | |
| else: | |
| # Merge next model | |
| # Comfy merge logic is flipped for stacked nodes. This is because the first model is effectively model1 and all subsequent models are model2. | |
| model2 = merge_model[0].clone() | |
| kp = model2.get_key_patches("diffusion_model.") | |
| for k in kp: | |
| #model1.add_patches({k: kp[k]}, 1.0 - model_ratio, model_ratio) #original logic | |
| model1.add_patches({k: kp[k]}, model_ratio, 1.0 - model_ratio) #flipped logic | |
| # Merge next clip | |
| clip2 = merge_model[1].clone() | |
| kp = clip2.get_key_patches() | |
| for k in kp: | |
| if k.endswith(".position_ids") or k.endswith(".logit_scale"): | |
| continue | |
| #clip1.add_patches({k: kp[k]}, 1.0 - clip_ratio, clip_ratio) #original logic | |
| clip1.add_patches({k: kp[k]}, clip_ratio, 1.0 - clip_ratio) #flipped logic | |
| # Update model info | |
| model_mix_info = model_mix_info + "\nModel Name: " + model_name + "\nModel Ratio: " + str(model_ratio) + "\nCLIP Ratio: " + str(clip_ratio) + "\n" | |
| show_help = "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes/wiki/Model-Merge-Nodes#cr-apply-model-merge" | |
| return (model1, clip1, model_mix_info, show_help, ) | |
| #---------------------------------------------------------------------------------------------------------------------# | |
| # MAPPINGS | |
| #---------------------------------------------------------------------------------------------------------------------# | |
| # For reference only, actual mappings are in __init__.py | |
| ''' | |
| NODE_CLASS_MAPPINGS = { | |
| "CR Apply Model Merge": CR_ApplyModelMerge, | |
| "CR Model Merge Stack": CR_ModelMergeStack, | |
| } | |
| ''' | |