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

@@ -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()

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
 

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,