Fixed default conds not respecting hook keyframes, made keyframes not reset cache when strength is unchanged, fixed Cond Set Default Combine throwing error, fixed model-as-lora throwing error during calculate_weight after a recent ComfyUI update, small refactoring/scaffolding changes for hooks

comfy/hooks.py

@@ -34,7 +34,7 @@ class _HookRef:
     pass
 
 # NOTE: this is an example of how the should_register function should look
-def default_should_register(hook: 'Hook', model: 'ModelPatcher', target: EnumWeightTarget):
+def default_should_register(hook: 'Hook', model: 'ModelPatcher', target: EnumWeightTarget, registered: list[Hook]):
     return True
 
 
@@ -45,6 +45,7 @@ class Hook:
         self.hook_ref = hook_ref if hook_ref else _HookRef()
         self.hook_keyframe = hook_keyframe if hook_keyframe else HookKeyframeGroup()
         self.custom_should_register = default_should_register
+        self.auto_apply_to_nonpositive = False
 
     @property
     def strength(self):
@@ -65,10 +66,21 @@ class Hook:
         c.hook_ref = self.hook_ref
         c.hook_keyframe = self.hook_keyframe
         c.custom_should_register = self.custom_should_register
+        # TODO: make this do something
+        c.auto_apply_to_nonpositive = self.auto_apply_to_nonpositive
         return c
 
-    def should_register(self, model: 'ModelPatcher', target: EnumWeightTarget):
+    def should_register(self, model: 'ModelPatcher', target: EnumWeightTarget, registered: list[Hook]):
-        return self.custom_should_register(self, model, target)
+        return self.custom_should_register(self, model, target, registered)
+
+    def add_hook_patches(self, model: 'ModelPatcher', target: EnumWeightTarget, registered: list[Hook]):
+        raise NotImplementedError("add_hook_patches should be defined for Hook subclasses")
+
+    def on_apply(self, model: 'ModelPatcher', transformer_options: dict[str]):
+        pass
+
+    def on_unapply(self, model: 'ModelPatcher', transformer_options: dict[str]):
+        pass
 
     def __eq__(self, other: 'Hook'):
         return self.__class__ == other.__class__ and self.hook_ref == other.hook_ref
@@ -94,9 +106,9 @@ class WeightHook(Hook):
     def strength_clip(self):
         return self._strength_clip * self.strength
 
-    def add_hook_patches(self, model: 'ModelPatcher', target: EnumWeightTarget):
+    def add_hook_patches(self, model: 'ModelPatcher', target: EnumWeightTarget, registered: list[Hook]):
-        if not self.should_register(model, target):
+        if not self.should_register(model, target, registered):
-            return
+            return False
         weights = None
         if target == EnumWeightTarget.Model:
             strength = self._strength_model
@@ -116,6 +128,7 @@ class WeightHook(Hook):
             else:
                 weights = self.weights_clip
         k = model.add_hook_patches(hook=self, patches=weights, strength_patch=strength, is_diff=self.is_diff)
+        return True
         # TODO: add logs about any keys that were not applied
 
     def clone(self, subtype: Callable=None):
@@ -308,6 +321,7 @@ class HookKeyframeGroup:
         self._current_keyframe: HookKeyframe = None
         self._current_used_steps = 0
         self._current_index = 0
+        self._current_strength = None
         self._curr_t = -1.
 
     # properties shadow those of HookWeightsKeyframe
@@ -321,6 +335,7 @@ class HookKeyframeGroup:
         self._current_keyframe = None
         self._current_used_steps = 0
         self._current_index = 0
+        self._current_strength = None
         self.curr_t = -1.
         self._set_first_as_current()
     
@@ -359,6 +374,7 @@ class HookKeyframeGroup:
         if curr_t == self._curr_t:
             return False
         prev_index = self._current_index
+        prev_strength = self._current_strength
         # if met guaranteed steps, look for next keyframe in case need to switch
         if self._current_used_steps >= self._current_keyframe.guarantee_steps:
             # if has next index, loop through and see if need to switch
@@ -369,6 +385,7 @@ class HookKeyframeGroup:
                     # NOTE: t is in terms of sigmas, not percent, so bigger number = earlier step in sampling
                     if eval_c.start_t >= curr_t:
                         self._current_index = i
+                        self._current_strength = eval_c.strength
                         self._current_keyframe = eval_c
                         self._current_used_steps = 0
                         # if guarantee_steps greater than zero, stop searching for other keyframes
@@ -381,7 +398,7 @@ class HookKeyframeGroup:
         # update current timestep this was performed on
         self._curr_t = curr_t
         # return True if keyframe changed, False if no change
-        return prev_index != self._current_index
+        return prev_index != self._current_index and prev_strength != self._current_strength
 
 
 class InterpolationMethod:

comfy/lora.py

@@ -440,7 +440,7 @@ def calculate_weight(patches, weight, key, intermediate_dtype=torch.float32, ori
         elif patch_type == "model_as_lora":
             target_weight: torch.Tensor = v[0]
             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)
+                          comfy.model_management.cast_to_device(original_weights[key][0][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)

comfy/model_patcher.py

@@ -198,8 +198,6 @@ class WrapperExecutor:
         args = list(args)
         kwargs = dict(kwargs)
         if self.is_last:
-            if self.class_obj is None:
-                return self.original(*args, **kwargs)
             return self.original(*args, **kwargs)
         return self.wrappers[self.idx](self, *args, **kwargs)
 
@@ -968,13 +966,14 @@ class ModelPatcher:
     def register_all_hook_patches(self, hooks_dict: dict[comfy.hooks.EnumHookType, dict[comfy.hooks.Hook, None]], target: comfy.hooks.EnumWeightTarget):
         self.restore_hook_patches()
         weight_hooks_to_register: list[comfy.hooks.WeightHook] = []
+        registered_hooks: list[comfy.hooks.Hook] = []
         for hook in hooks_dict.get(comfy.hooks.EnumHookType.Weight, {}):
             if hook.hook_ref not in self.hook_patches:
                 weight_hooks_to_register.append(hook)
         if len(weight_hooks_to_register) > 0:
             self.hook_patches_backup = create_hook_patches_clone(self.hook_patches)
             for hook in weight_hooks_to_register:
-                hook.add_hook_patches(self, target)
+                hook.add_hook_patches(self, target, registered_hooks)
         for callback in self.get_all_callbacks(CallbacksMP.ON_REGISTER_ALL_HOOK_PATCHES):
             callback(self, hooks_dict, target)
 
@@ -1066,7 +1065,9 @@ class ModelPatcher:
             model_sd = self.model_state_dict()
             memory_counter = None
             if self.hook_mode == comfy.hooks.EnumHookMode.MaxSpeed:
-                memory_counter = MemoryCounter(comfy.model_management.get_free_memory(self.load_device))
+                # TODO: minimum_counter should have a minimum that conforms to loaded model requirements
+                memory_counter = MemoryCounter(initial=comfy.model_management.get_free_memory(self.load_device),
+                                               minimum=comfy.model_management.minimum_inference_memory())
             # if have cached weights for hooks, use it
             cached_weights = self.cached_hook_patches.get(hooks, None)
             if cached_weights is not None:

comfy/samplers.py

@@ -144,7 +144,7 @@ def cond_cat(c_list):
 
     return out
 
-def finalize_default_conds(hooked_to_run: dict[comfy.hooks.HookGroup,list[tuple[tuple,int]]], default_conds: list[list[dict]], x_in, timestep):
+def finalize_default_conds(model: 'BaseModel', hooked_to_run: dict[comfy.hooks.HookGroup,list[tuple[tuple,int]]], default_conds: list[list[dict]], x_in, timestep):
     # need to figure out remaining unmasked area for conds
     default_mults = []
     for _ in default_conds:
@@ -178,9 +178,11 @@ def finalize_default_conds(hooked_to_run: dict[comfy.hooks.HookGroup,list[tuple[
                 continue
             # replace p's mult with calculated mult
             p = p._replace(mult=mult)
-            hook: comfy.hooks.HookGroup = x.get('hooks', None)
+            hooks: comfy.hooks.HookGroup = x.get('hooks', None)
-            hooked_to_run.setdefault(hook, list())
+            if hooks is not None:
-            hooked_to_run[hook] += [(p, i)]
+                model.current_patcher.prepare_hook_patches_current_keyframe(timestep, hooks)
+            hooked_to_run.setdefault(hooks, list())
+            hooked_to_run[hooks] += [(p, i)]
 
 def calc_cond_batch(model: 'BaseModel', conds: list[list[dict]], x_in: torch.Tensor, timestep, model_options):
     executor = comfy.model_patcher.WrapperExecutor.new_executor(
@@ -221,7 +223,7 @@ def outer_calc_cond_batch(model: 'BaseModel', conds: list[list[dict]], x_in: tor
         default_conds.append(default_c)
 
     if has_default_conds:
-        finalize_default_conds(hooked_to_run, default_conds, x_in, timestep)
+        finalize_default_conds(model, hooked_to_run, default_conds, x_in, timestep)
 
     model.current_patcher.prepare_state(timestep)
 

comfy_extras/nodes_hooks.py

@@ -214,7 +214,7 @@ class ConditioningSetDefaultAndCombine:
     CATEGORY = "advanced/hooks/cond single"
     FUNCTION = "set_default_and_combine"
 
-    def append_and_combine(self, cond, cond_DEFAULT,
+    def set_default_and_combine(self, cond, cond_DEFAULT,
                            opt_hooks: comfy.hooks.HookGroup=None):
         (final_conditioning,) = comfy.hooks.set_default_and_combine_conds(conds=[cond], new_conds=[cond_DEFAULT],
                                                                         opt_hooks=opt_hooks)