diff --git a/README.md b/README.md index 96cde4c4b..5e32a74f3 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,8 @@ This ui will let you design and execute advanced stable diffusion pipelines usin - [Upscale Models (ESRGAN, ESRGAN variants, SwinIR, Swin2SR, etc...)](https://comfyanonymous.github.io/ComfyUI_examples/upscale_models/) - [unCLIP Models](https://comfyanonymous.github.io/ComfyUI_examples/unclip/) - [GLIGEN](https://comfyanonymous.github.io/ComfyUI_examples/gligen/) -- Latent previews with [TAESD](https://github.com/madebyollin/taesd) +- [Model Merging](https://comfyanonymous.github.io/ComfyUI_examples/model_merging/) +- Latent previews with [TAESD](#how-to-show-high-quality-previews) - Starts up very fast. - Works fully offline: will never download anything. - [Config file](extra_model_paths.yaml.example) to set the search paths for models. diff --git a/comfy/sd.py b/comfy/sd.py index 542f704a6..8eac1f8ed 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -89,8 +89,7 @@ LORA_UNET_MAP_RESNET = { "skip_connection": "resnets_{}_conv_shortcut" } -def load_lora(path, to_load): - lora = utils.load_torch_file(path, safe_load=True) +def load_lora(lora, to_load): patch_dict = {} loaded_keys = set() for x in to_load: @@ -501,10 +500,10 @@ class ModelPatcher: self.backup = {} -def load_lora_for_models(model, clip, lora_path, strength_model, strength_clip): +def load_lora_for_models(model, clip, lora, strength_model, strength_clip): key_map = model_lora_keys(model.model) key_map = model_lora_keys(clip.cond_stage_model, key_map) - loaded = load_lora(lora_path, key_map) + loaded = load_lora(lora, key_map) new_modelpatcher = model.clone() k = new_modelpatcher.add_patches(loaded, strength_model) new_clip = clip.clone() diff --git a/comfy_extras/nodes_model_merging.py b/comfy_extras/nodes_model_merging.py index 4f71b2031..72eeffb39 100644 --- a/comfy_extras/nodes_model_merging.py +++ b/comfy_extras/nodes_model_merging.py @@ -14,7 +14,7 @@ class ModelMergeSimple: RETURN_TYPES = ("MODEL",) FUNCTION = "merge" - CATEGORY = "_for_testing/model_merging" + CATEGORY = "advanced/model_merging" def merge(self, model1, model2, ratio): m = model1.clone() @@ -35,7 +35,7 @@ class ModelMergeBlocks: RETURN_TYPES = ("MODEL",) FUNCTION = "merge" - CATEGORY = "_for_testing/model_merging" + CATEGORY = "advanced/model_merging" def merge(self, model1, model2, **kwargs): m = model1.clone() @@ -68,7 +68,7 @@ class CheckpointSave: FUNCTION = "save" OUTPUT_NODE = True - CATEGORY = "_for_testing/model_merging" + CATEGORY = "advanced/model_merging" def save(self, model, clip, vae, filename_prefix, prompt=None, extra_pnginfo=None): full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) diff --git a/execution.py b/execution.py index f93de8465..a40b1dd36 100644 --- a/execution.py +++ b/execution.py @@ -110,7 +110,7 @@ def format_value(x): else: return str(x) -def recursive_execute(server, prompt, outputs, current_item, extra_data, executed, prompt_id, outputs_ui): +def recursive_execute(server, prompt, outputs, current_item, extra_data, executed, prompt_id, outputs_ui, object_storage): unique_id = current_item inputs = prompt[unique_id]['inputs'] class_type = prompt[unique_id]['class_type'] @@ -125,7 +125,7 @@ def recursive_execute(server, prompt, outputs, current_item, extra_data, execute input_unique_id = input_data[0] output_index = input_data[1] if input_unique_id not in outputs: - result = recursive_execute(server, prompt, outputs, input_unique_id, extra_data, executed, prompt_id, outputs_ui) + result = recursive_execute(server, prompt, outputs, input_unique_id, extra_data, executed, prompt_id, outputs_ui, object_storage) if result[0] is not True: # Another node failed further upstream return result @@ -136,7 +136,11 @@ def recursive_execute(server, prompt, outputs, current_item, extra_data, execute if server.client_id is not None: server.last_node_id = unique_id server.send_sync("executing", { "node": unique_id, "prompt_id": prompt_id }, server.client_id) - obj = class_def() + + obj = object_storage.get((unique_id, class_type), None) + if obj is None: + obj = class_def() + object_storage[(unique_id, class_type)] = obj output_data, output_ui = get_output_data(obj, input_data_all) outputs[unique_id] = output_data @@ -256,6 +260,7 @@ def recursive_output_delete_if_changed(prompt, old_prompt, outputs, current_item class PromptExecutor: def __init__(self, server): self.outputs = {} + self.object_storage = {} self.outputs_ui = {} self.old_prompt = {} self.server = server @@ -322,6 +327,17 @@ class PromptExecutor: for o in to_delete: d = self.outputs.pop(o) del d + to_delete = [] + for o in self.object_storage: + if o[0] not in prompt: + to_delete += [o] + else: + p = prompt[o[0]] + if o[1] != p['class_type']: + to_delete += [o] + for o in to_delete: + d = self.object_storage.pop(o) + del d for x in prompt: recursive_output_delete_if_changed(prompt, self.old_prompt, self.outputs, x) @@ -349,7 +365,7 @@ class PromptExecutor: # This call shouldn't raise anything if there's an error deep in # the actual SD code, instead it will report the node where the # error was raised - success, error, ex = recursive_execute(self.server, prompt, self.outputs, output_node_id, extra_data, executed, prompt_id, self.outputs_ui) + success, error, ex = recursive_execute(self.server, prompt, self.outputs, output_node_id, extra_data, executed, prompt_id, self.outputs_ui, self.object_storage) if success is not True: self.handle_execution_error(prompt_id, prompt, current_outputs, executed, error, ex) break diff --git a/nodes.py b/nodes.py index f10515f89..a9f2e962e 100644 --- a/nodes.py +++ b/nodes.py @@ -434,6 +434,9 @@ class CLIPSetLastLayer: return (clip,) class LoraLoader: + def __init__(self): + self.loaded_lora = None + @classmethod def INPUT_TYPES(s): return {"required": { "model": ("MODEL",), @@ -452,7 +455,18 @@ class LoraLoader: return (model, clip) lora_path = folder_paths.get_full_path("loras", lora_name) - model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora_path, strength_model, strength_clip) + lora = None + if self.loaded_lora is not None: + if self.loaded_lora[0] == lora_path: + lora = self.loaded_lora[1] + else: + del self.loaded_lora + + if lora is None: + lora = comfy.utils.load_torch_file(lora_path, safe_load=True) + self.loaded_lora = (lora_path, lora) + + model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip) return (model_lora, clip_lora) class VAELoader: