From 787ef34842ec58d831da99a384997532dc8e396c Mon Sep 17 00:00:00 2001 From: "kosinkadink1@gmail.com" Date: Thu, 19 Sep 2024 11:47:25 +0900 Subject: [PATCH] Continued work on simpler Create Hook Model As LoRA node, started to implement ModelPatcher callbacks, attachments, and additional_models --- comfy/hooks.py | 16 ++++-- comfy/lora.py | 3 +- comfy/model_base.py | 5 +- comfy/model_patcher.py | 98 +++++++++++++++++++++++++++++++------ comfy/samplers.py | 14 ++++-- comfy_extras/nodes_hooks.py | 52 +++++++++++++++----- 6 files changed, 153 insertions(+), 35 deletions(-) diff --git a/comfy/hooks.py b/comfy/hooks.py index f8b8b685b..3fc8c8339 100644 --- a/comfy/hooks.py +++ b/comfy/hooks.py @@ -337,9 +337,19 @@ def create_hook_model_as_lora(weights_model, weights_clip, strength_model: float hook_group = HookGroup() hook = WeightHook(strength_model=strength_model, strength_clip=strength_clip) hook_group.add(hook) - hook.weights = weights_model - hook.weights_clip = weights_clip - hook.is_diff = True + patches_model = None + patches_clip = None + if weights_model is not None: + patches_model = {} + for key in weights_model: + patches_model[key] = ("model_as_lora", (weights_model[key],)) + if weights_clip is not None: + patches_clip = {} + for key in weights_clip: + patches_clip[key] = ("model_as_lora", (weights_clip[key],)) + hook.weights = patches_model + hook.weights_clip = patches_clip + hook.need_weight_init = False return hook_group def get_patch_weights_from_model(model: 'ModelPatcher', discard_model_sampling=False): diff --git a/comfy/lora.py b/comfy/lora.py index 38061ad31..4c01a37b2 100644 --- a/comfy/lora.py +++ b/comfy/lora.py @@ -438,7 +438,8 @@ def calculate_weight(patches, weight, key, intermediate_dtype=torch.float32, ori weight += function(strength * comfy.model_management.cast_to_device(diff, weight.device, weight.dtype)) elif patch_type == "model_as_lora": target_weight: torch.Tensor = v[0] - diff_weight = target_weight.to(intermediate_dtype) - original_weights[key].to(intermediate_dtype) + diff_weight = comfy.model_management.cast_to_device(target_weight, weight.device, intermediate_dtype) - \ + comfy.model_management.cast_to_device(original_weights[key][0], weight.device, intermediate_dtype) weight += function(strength * comfy.model_management.cast_to_device(diff_weight, weight.device, weight.dtype)) elif patch_type == "lora": #lora/locon mat1 = comfy.model_management.cast_to_device(v[0], weight.device, intermediate_dtype) diff --git a/comfy/model_base.py b/comfy/model_base.py index fbd11b489..01316400c 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -37,6 +37,9 @@ from enum import Enum from . import utils import comfy.latent_formats import math +from typing import TYPE_CHECKING +if TYPE_CHECKING: + from model_patcher import ModelPatcher class ModelType(Enum): EPS = 1 @@ -93,7 +96,7 @@ class BaseModel(torch.nn.Module): self.model_config = model_config self.manual_cast_dtype = model_config.manual_cast_dtype self.device = device - self.current_patcher = None + self.current_patcher: 'ModelPatcher' = None if not unet_config.get("disable_unet_model_creation", False): if model_config.custom_operations is None: diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 0d4360977..80b4651e2 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -78,6 +78,14 @@ def set_model_options_pre_cfg_function(model_options, pre_cfg_function, disable_ model_options["disable_cfg1_optimization"] = True return model_options +def create_hook_patches_clone(orig_hook_patches): + new_hook_patches = {} + for hook_ref in orig_hook_patches: + new_hook_patches[hook_ref] = {} + for k in orig_hook_patches[hook_ref]: + new_hook_patches[hook_ref][k] = orig_hook_patches[hook_ref][k][:] + return new_hook_patches + def wipe_lowvram_weight(m): if hasattr(m, "prev_comfy_cast_weights"): m.comfy_cast_weights = m.prev_comfy_cast_weights @@ -92,6 +100,27 @@ class LowVramPatch: def __call__(self, weight): return comfy.lora.calculate_weight(self.patches[self.key], weight, self.key, intermediate_dtype=weight.dtype) +class CallbacksMP: + ON_CLONE = "on_clone" + ON_LOAD = "on_load_after" + ON_CLEANUP = "on_cleanup" + ON_PRE_RUN = "on_pre_run" + ON_PREPARE_STATE = "on_prepare_state" + ON_APPLY_HOOKS = "on_apply_hooks" + ON_REGISTER_ALL_HOOK_PATCHES = "on_register_all_hook_patches" + + @classmethod + def init_callbacks(cls): + return { + cls.ON_CLONE: [], + cls.ON_LOAD: [], + cls.ON_CLEANUP: [], + cls.ON_PRE_RUN: [], + cls.ON_PREPARE_STATE: [], + cls.ON_APPLY_HOOKS: [], + cls.ON_REGISTER_ALL_HOOK_PATCHES: [], + } + class ModelPatcher: def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False): self.size = size @@ -113,6 +142,10 @@ class ModelPatcher: self.weight_inplace_update = weight_inplace_update self.patches_uuid = uuid.uuid4() + self.attachments: Dict[str] = {} + self.additional_models: list[ModelPatcher] = [] + self.callbacks: Dict[str, List[Callable]] = CallbacksMP.init_callbacks() + self.hook_patches: Dict[comfy.hooks._HookRef] = {} self.hook_patches_backup: Dict[comfy.hooks._HookRef] = {} self.hook_backup: Dict[str, Tuple[torch.Tensor, torch.device]] = {} @@ -155,6 +188,19 @@ class ModelPatcher: n.backup = self.backup n.object_patches_backup = self.object_patches_backup + # attachments + n.attachments = {} + for k in self.attachments: + if hasattr(self.attachments[k], "on_model_patcher_clone"): + n.attachments[k] = self.attachments[k].on_model_patcher_clone() + else: + n.attachments[k] = self.attachments[k] + # additional models + for m in self.additional_models: + n.additional_models.append(m.clone()) + # callbacks + for k, c in self.callbacks.items(): + n.callbacks[k] = c.copy() # hooks n.hook_patches = self.create_hook_patches_clone(self.hook_patches) n.hook_patches_backup = self.create_hook_patches_clone(self.hook_patches_backup) @@ -167,16 +213,10 @@ class ModelPatcher: n.current_hooks = self.current_hooks.clone() if self.current_hooks else self.current_hooks n.forced_hooks = self.forced_hooks.clone() if self.forced_hooks else self.forced_hooks n.hook_mode = self.hook_mode - return n - @staticmethod - def create_hook_patches_clone(orig_hook_patches): - new_hook_patches = {} - for hook_ref in orig_hook_patches: - new_hook_patches[hook_ref] = {} - for k in orig_hook_patches[hook_ref]: - new_hook_patches[hook_ref][k] = orig_hook_patches[hook_ref][k][:] - return new_hook_patches + for callback in self.callbacks[CallbacksMP.ON_CLONE]: + callback(self, n) + return n def is_clone(self, other): if hasattr(other, 'model') and self.model is other.model: @@ -332,8 +372,11 @@ class ModelPatcher: if not k.startswith(filter_prefix): continue bk = self.backup.get(k, None) + hbk = self.hook_backup.get(k, None) if bk is not None: weight = bk.weight + if hbk is not None: + weight = hbk[0] else: weight = model_sd[k] if k in self.patches: @@ -580,9 +623,27 @@ class ModelPatcher: print("WARNING the ModelPatcher.calculate_weight function is deprecated, please use: comfy.lora.calculate_weight instead") return comfy.lora.calculate_weight(patches, weight, key, intermediate_dtype=intermediate_dtype) - def clean(self): + def cleanup(self): self.clean_hooks() self.restore_hook_patches() + for callback in self.callbacks[CallbacksMP.ON_CLEANUP]: + callback(self) + + def add_callback(self, key, callback: Callable): + if key not in self.callbacks: + raise Exception(f"Callback '{key}' is not recognized.") + self.callbacks[key].append(callback) + + def add_attachment(self, attachment): + self.attachments.append(attachment) + + def pre_run(self): + for callback in self.callbacks[CallbacksMP.ON_PRE_RUN]: + callback(self) + + def prepare_state(self, timestep): + for callback in self.callbacks[CallbacksMP.ON_PREPARE_STATE]: + callback(self, timestep) def restore_hook_patches(self): if len(self.hook_patches_backup) > 0: @@ -596,10 +657,10 @@ class ModelPatcher: curr_t = t[0] for hook in hook_group.hooks: changed = hook.hook_keyframe.prepare_current_keyframe(curr_t=curr_t) - # if keyframe changed, remove any cached LoraHookGroups that contain hook with the same hook_ref; + # if keyframe changed, remove any cached HookGroups that contain hook with the same hook_ref; # this will cause the weights to be recalculated when sampling if changed: - # reset current_lora_hooks if contains lora hook that changed + # reset current_hooks if contains hook that changed if self.current_hooks is not None: for current_hook in self.current_hooks.hooks: if current_hook == hook: @@ -620,6 +681,8 @@ class ModelPatcher: self.hook_patches_backup = self.create_hook_patches_clone(self.hook_patches) for hook in weight_hooks_to_register: hook.add_hook_patches(self, target) + for callback in self.callbacks[CallbacksMP.ON_REGISTER_ALL_HOOK_PATCHES]: + callback(self, hooks_dict, target) def add_hook_patches(self, hook: comfy.hooks.WeightHook, patches, strength_patch=1.0, strength_model=1.0, is_diff=False): # NOTE: this mirrors behavior of add_patches func @@ -698,6 +761,8 @@ class ModelPatcher: if self.current_hooks == hooks: return self.patch_hooks(hooks=hooks) + for callback in self.callbacks[CallbacksMP.ON_APPLY_HOOKS]: + callback(self, hooks) def patch_hooks(self, hooks: comfy.hooks.HookGroup): self.unpatch_hooks() @@ -712,11 +777,14 @@ class ModelPatcher: self.patch_cached_hook_weights(cached_weights=cached_weights, key=key) else: relevant_patches = self.get_combined_hook_patches(hooks=hooks) + original_weights = None + if len(relevant_patches) > 0: + original_weights = self.get_key_patches() for key in relevant_patches: if key not in model_sd: print(f"WARNING cached hook would not patch. key does not exist in model: {key}") continue - self.patch_hook_weight_to_device(hooks=hooks, combined_patches=relevant_patches, key=key) + self.patch_hook_weight_to_device(hooks=hooks, combined_patches=relevant_patches, key=key, original_weights=original_weights) self.current_hooks = hooks def patch_cached_hook_weights(self, cached_weights: Dict, key: str): @@ -735,7 +803,7 @@ class ModelPatcher: self.cached_hook_patches.clear() self.current_hooks = None - def patch_hook_weight_to_device(self, hooks: comfy.hooks.HookGroup, combined_patches: dict, key: str): + def patch_hook_weight_to_device(self, hooks: comfy.hooks.HookGroup, combined_patches: dict, key: str, original_weights: dict): if key not in combined_patches: return weight: torch.Tensor = comfy.utils.get_attr(self.model, key) @@ -747,7 +815,7 @@ class ModelPatcher: # TODO: properly handle lowvram situations for cached hook patches temp_weight = comfy.model_management.cast_to_device(weight, weight.device, torch.float32, copy=True) - out_weight = comfy.lora.calculate_weight(combined_patches[key], temp_weight, key).to(weight.dtype) + out_weight = comfy.lora.calculate_weight(combined_patches[key], temp_weight, key, original_weights=original_weights).to(weight.dtype) if self.hook_mode == comfy.hooks.EnumHookMode.MaxSpeed: self.cached_hook_patches.setdefault(hooks, {}) self.cached_hook_patches[hooks][key] = out_weight diff --git a/comfy/samplers.py b/comfy/samplers.py index 537acd523..e3891ba25 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -1,6 +1,9 @@ from .k_diffusion import sampling as k_diffusion_sampling from .extra_samplers import uni_pc -from typing import Dict, List, Tuple +from typing import TYPE_CHECKING, Dict, List, Tuple +if TYPE_CHECKING: + from comfy.model_patcher import ModelPatcher + from comfy.model_base import BaseModel import torch import collections from comfy import model_management @@ -178,7 +181,7 @@ def finalize_default_conds(hooked_to_run: Dict[comfy.hooks.HookGroup,List[Tuple[ hooked_to_run.setdefault(hook, list()) hooked_to_run[hook] += [(p, i)] -def calc_cond_batch(model, conds: List[List[Dict]], x_in, timestep, model_options): +def calc_cond_batch(model: 'BaseModel', conds: List[List[Dict]], x_in: torch.Tensor, timestep, model_options): out_conds = [] out_counts = [] # separate conds by matching hooks @@ -212,6 +215,8 @@ def calc_cond_batch(model, conds: List[List[Dict]], x_in, timestep, model_option if has_default_conds: finalize_default_conds(hooked_to_run, default_conds, x_in, timestep) + model.current_patcher.prepare_state(timestep) + # run every hooked_to_run separately for hooks, to_run in hooked_to_run.items(): while len(to_run) > 0: @@ -729,7 +734,7 @@ def process_conds(model, noise, conds, device, latent_image=None, denoise_mask=N class CFGGuider: def __init__(self, model_patcher): - self.model_patcher = model_patcher + self.model_patcher: 'ModelPatcher' = model_patcher self.model_options = model_patcher.model_options self.original_conds = {} self.cfg = 1.0 @@ -780,10 +785,11 @@ class CFGGuider: sigmas = sigmas.to(device) try: + self.model_patcher.pre_run() comfy.sampler_helpers.prepare_model_patcher(self.model_patcher, self.conds) output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed) finally: - self.model_patcher.clean() + self.model_patcher.cleanup() comfy.sampler_helpers.cleanup_models(self.conds, self.loaded_models) del self.inner_model diff --git a/comfy_extras/nodes_hooks.py b/comfy_extras/nodes_hooks.py index 534484628..048b37663 100644 --- a/comfy_extras/nodes_hooks.py +++ b/comfy_extras/nodes_hooks.py @@ -216,6 +216,9 @@ class CreateHookLora: "lora_name": (folder_paths.get_filename_list("loras"), ), "strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), "strength_clip": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), + }, + "optional": { + "prev_hooks": ("HOOKS",) } } @@ -223,9 +226,13 @@ class CreateHookLora: CATEGORY = "advanced/hooks/create" FUNCTION = "create_hook" - def create_hook(self, lora_name: str, strength_model: float, strength_clip: float): + def create_hook(self, lora_name: str, strength_model: float, strength_clip: float, prev_hooks: comfy.hooks.HookGroup=None): + if prev_hooks is None: + prev_hooks = comfy.hooks.HookGroup() + prev_hooks.clone() + if strength_model == 0 and strength_clip == 0: - return (None,) + return (prev_hooks,) lora_path = folder_paths.get_full_path("loras", lora_name) lora = None @@ -242,7 +249,7 @@ class CreateHookLora: self.loaded_lora = (lora_path, lora) hooks = comfy.hooks.create_hook_lora(lora=lora, strength_model=strength_model, strength_clip=strength_clip) - return (hooks,) + return (prev_hooks.clone_and_combine(hooks),) class CreateHookLoraModelOnly(CreateHookLora): NodeId = 'CreateHookLoraModelOnly' @@ -253,6 +260,9 @@ class CreateHookLoraModelOnly(CreateHookLora): "required": { "lora_name": (folder_paths.get_filename_list("loras"), ), "strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), + }, + "optional": { + "prev_hooks": ("HOOKS",) } } @@ -260,8 +270,8 @@ class CreateHookLoraModelOnly(CreateHookLora): CATEGORY = "advanced/hooks/create" FUNCTION = "create_hook_model_only" - def create_hook_model_only(self, lora_name: str, strength_model: float): - return self.create_hook(lora_name=lora_name, strength_model=strength_model, strength_clip=0) + def create_hook_model_only(self, lora_name: str, strength_model: float, prev_hooks: comfy.hooks.HookGroup=None): + return self.create_hook(lora_name=lora_name, strength_model=strength_model, strength_clip=0, prev_hooks=prev_hooks) class CreateHookModelAsLora: NodeId = 'CreateHookModelAsLora' @@ -275,6 +285,9 @@ class CreateHookModelAsLora: "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), "strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), "strength_clip": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), + }, + "optional": { + "prev_hooks": ("HOOKS",) } } @@ -283,7 +296,12 @@ class CreateHookModelAsLora: FUNCTION = "create_hook" def create_hook(self, model: 'ModelPatcher', clip: 'CLIP', ckpt_name: str, - strength_model: float, strength_clip: float): + strength_model: float, strength_clip: float, + prev_hooks: comfy.hooks.HookGroup=None): + if prev_hooks is None: + prev_hooks = comfy.hooks.HookGroup() + prev_hooks.clone() + ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) model_loaded = out[0] @@ -292,7 +310,7 @@ class CreateHookModelAsLora: hooks = comfy.hooks.create_hook_model_as_lora_precalc(model=model, clip=clip, model_loaded=model_loaded, clip_loaded=clip_loaded, strength_model=strength_model, strength_clip=strength_clip) - return (hooks,) + return (prev_hooks.clone_and_combine(hooks),) class CreateHookModelAsLoraModelOnly: NodeId = 'CreateHookModelAsLoraModelOnly' @@ -304,6 +322,9 @@ class CreateHookModelAsLoraModelOnly: "model": ("MODEL",), "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), "strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), + }, + "optional": { + "prev_hooks": ("HOOKS",) } } @@ -311,9 +332,10 @@ class CreateHookModelAsLoraModelOnly: CATEGORY = "advanced/hooks/create" FUNCTION = "create_hook_model_only" - def create_hook_model_only(self, model: 'ModelPatcher', ckpt_name: str, strength_model: float): + def create_hook_model_only(self, model: 'ModelPatcher', ckpt_name: str, strength_model: float, + prev_hooks: comfy.hooks.HookGroup=None): return CreateHookModelAsLora.create_hook(self, model=model, clip=None, ckpt_name=ckpt_name, - strength_model=strength_model, strength_clip=0) + strength_model=strength_model, strength_clip=0, prev_hooks=prev_hooks) class CreateHookModelAsLoraTest: NodeId = 'CreateHookModelAsLoraTest' @@ -331,6 +353,9 @@ class CreateHookModelAsLoraTest: "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), "strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), "strength_clip": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), + }, + "optional": { + "prev_hooks": ("HOOKS",) } } @@ -338,7 +363,12 @@ class CreateHookModelAsLoraTest: CATEGORY = "advanced/hooks/create" FUNCTION = "create_hook" - def create_hook(self, ckpt_name: str, strength_model: float, strength_clip: float): + def create_hook(self, ckpt_name: str, strength_model: float, strength_clip: float, + prev_hooks: comfy.hooks.HookGroup=None): + if prev_hooks is None: + prev_hooks = comfy.hooks.HookGroup() + prev_hooks.clone() + ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) weights_model = None weights_clip = None @@ -359,7 +389,7 @@ class CreateHookModelAsLoraTest: hooks = comfy.hooks.create_hook_model_as_lora(weights_model=weights_model, weights_clip=weights_clip, strength_model=strength_model, strength_clip=strength_clip) - return (hooks,) + return (prev_hooks.clone_and_combine(hooks),) #------------------------------------------ ###########################################