From cfb145187dca05e293ee6b0c9a216dd7b2df1286 Mon Sep 17 00:00:00 2001 From: "kosinkadink1@gmail.com" Date: Tue, 17 Sep 2024 09:55:14 +0900 Subject: [PATCH] Made Set Clip Hooks node work with hooks from Create Hook nodes, began work on better Create Hook Model As LoRA node --- comfy/hooks.py | 35 +++++++++++++++++++++++--- comfy/lora.py | 9 ++++--- comfy_extras/nodes_hooks.py | 50 ++++++++++++++++++++++++++++++++++++- 3 files changed, 85 insertions(+), 9 deletions(-) diff --git a/comfy/hooks.py b/comfy/hooks.py index 570628668..f8b8b685b 100644 --- a/comfy/hooks.py +++ b/comfy/hooks.py @@ -66,6 +66,7 @@ class WeightHook(Hook): super().__init__(hook_type=EnumHookType.Weight) self.weights: Dict = None self.weights_clip: Dict = None + self.is_diff = False self.need_weight_init = True self._strength_model = strength_model self._strength_clip = strength_clip @@ -91,13 +92,13 @@ class WeightHook(Hook): key_map = comfy.lora.model_lora_keys_unet(model.model, key_map) else: key_map = comfy.lora.model_lora_keys_clip(model.model, key_map) - weights = comfy.lora.load_lora(self.weights, key_map) + weights = comfy.lora.load_lora(self.weights, key_map, log_missing=False) else: if target == EnumWeightTarget.Model: weights = self.weights else: weights = self.weights_clip - k = model.add_hook_patches(hook=self, patches=weights, strength_patch=strength) + k = model.add_hook_patches(hook=self, patches=weights, strength_patch=strength, is_diff=self.is_diff) # TODO: add logs about any keys that were not applied def clone(self, subtype: Callable=None): @@ -107,6 +108,7 @@ class WeightHook(Hook): c.weights = self.weights c.weights_clip = self.weights_clip c.need_weight_init = self.need_weight_init + c.is_diff = self.is_diff c._strength_model = self._strength_model c._strength_clip = self._strength_clip return c @@ -151,6 +153,12 @@ class HookGroup: for hook in self.hooks: hook.hook_keyframe = hook_kf + def get_dict_repr(self): + d = {} + for hook in self.hooks: + d[hook] = None + return d + @staticmethod def combine_all_hooks(hooks_list: List['HookGroup'], require_count=1) -> 'HookGroup': actual: List[HookGroup] = [] @@ -323,10 +331,29 @@ def create_hook_lora(lora: Dict[str, torch.Tensor], strength_model: float, stren hook = WeightHook(strength_model=strength_model, strength_clip=strength_clip) hook_group.add(hook) hook.weights = lora - hook.need_weight_init = True return hook_group -def create_hook_model_as_lora(model: 'ModelPatcher', clip: 'CLIP', +def create_hook_model_as_lora(weights_model, weights_clip, strength_model: float, strength_clip: 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 + return hook_group + +def get_patch_weights_from_model(model: 'ModelPatcher', discard_model_sampling=False): + if model is None: + return None + patches_model: Dict[str, torch.Tensor] = model.model.state_dict() + if discard_model_sampling: + # do not include ANY model_sampling components of the model that should act as a patch + for key in list(patches_model.keys()): + if key.startswith("model_sampling"): + patches_model.pop(key, None) + return patches_model + +def create_hook_model_as_lora_precalc(model: 'ModelPatcher', clip: 'CLIP', model_loaded: 'ModelPatcher', clip_loaded: 'CLIP', strength_model: float, strength_clip: float): hook_group = HookGroup() diff --git a/comfy/lora.py b/comfy/lora.py index 61979e500..5ea8d62a0 100644 --- a/comfy/lora.py +++ b/comfy/lora.py @@ -33,7 +33,7 @@ LORA_CLIP_MAP = { } -def load_lora(lora, to_load): +def load_lora(lora, to_load, log_missing=True): patch_dict = {} loaded_keys = set() for x in to_load: @@ -193,9 +193,10 @@ def load_lora(lora, to_load): patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = ("diff", (diff_bias,)) loaded_keys.add(diff_bias_name) - for x in lora.keys(): - if x not in loaded_keys: - logging.warning("lora key not loaded: {}".format(x)) + if log_missing: + for x in lora.keys(): + if x not in loaded_keys: + logging.warning("lora key not loaded: {}".format(x)) return patch_dict diff --git a/comfy_extras/nodes_hooks.py b/comfy_extras/nodes_hooks.py index db219c25c..534484628 100644 --- a/comfy_extras/nodes_hooks.py +++ b/comfy_extras/nodes_hooks.py @@ -175,6 +175,7 @@ class SetClipHooks: if hooks is not None: clip = clip.clone() clip.patcher.forced_hooks = hooks + clip.patcher.register_all_hook_patches(hooks.get_dict_repr(), comfy.hooks.EnumWeightTarget.Clip) return (clip,) class ConditioningTimestepsRange: @@ -288,7 +289,7 @@ class CreateHookModelAsLora: model_loaded = out[0] clip_loaded = out[1] - hooks = comfy.hooks.create_hook_model_as_lora(model=model, clip=clip, + 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,) @@ -313,6 +314,52 @@ class CreateHookModelAsLoraModelOnly: def create_hook_model_only(self, model: 'ModelPatcher', ckpt_name: str, strength_model: float): return CreateHookModelAsLora.create_hook(self, model=model, clip=None, ckpt_name=ckpt_name, strength_model=strength_model, strength_clip=0) + +class CreateHookModelAsLoraTest: + NodeId = 'CreateHookModelAsLoraTest' + NodeName = 'Create Hook Model as LoRA (TEST)' + + def __init__(self): + # when not None, will be in following format: + # (ckpt_path: str, weights_model: dict, weights_clip: dict) + self.loaded_weights = None + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "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}), + } + } + + RETURN_TYPES = ("HOOKS",) + CATEGORY = "advanced/hooks/create" + FUNCTION = "create_hook" + + def create_hook(self, ckpt_name: str, strength_model: float, strength_clip: float): + ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + weights_model = None + weights_clip = None + if self.loaded_weights is not None: + if self.loaded_weights[0] == ckpt_path: + weights_model = self.loaded_weights[1] + weights_clip = self.loaded_weights[2] + else: + temp = self.loaded_weights + self.loaded_weights = None + del temp + + if weights_model is None: + out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) + weights_model = comfy.hooks.get_patch_weights_from_model(out[0]) + weights_clip = comfy.hooks.get_patch_weights_from_model(out[1].patcher if out[1] else out[1]) + self.loaded_weights = (ckpt_path, weights_model, weights_clip) + + 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,) #------------------------------------------ ########################################### @@ -665,6 +712,7 @@ node_list = [ CreateHookLora, CreateHookLoraModelOnly, CreateHookModelAsLora, + CreateHookModelAsLoraTest, CreateHookModelAsLoraModelOnly, # Register RegisterHookLora,