diff --git a/comfy/hooks.py b/comfy/hooks.py index 145b55819..e5fd07599 100644 --- a/comfy/hooks.py +++ b/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): - return self.custom_should_register(self, model, target) + def should_register(self, model: 'ModelPatcher', target: EnumWeightTarget, registered: list[Hook]): + 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): - if not self.should_register(model, target): - return + def add_hook_patches(self, model: 'ModelPatcher', target: EnumWeightTarget, registered: list[Hook]): + if not self.should_register(model, target, registered): + 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: diff --git a/comfy/lora.py b/comfy/lora.py index a728fd85b..460b4a388 100644 --- a/comfy/lora.py +++ b/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) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index abf87965a..3a9242010 100644 --- a/comfy/model_patcher.py +++ b/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: diff --git a/comfy/samplers.py b/comfy/samplers.py index 9ae1f01a5..cbcd350c5 100644 --- a/comfy/samplers.py +++ b/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) - hooked_to_run.setdefault(hook, list()) - hooked_to_run[hook] += [(p, i)] + hooks: comfy.hooks.HookGroup = x.get('hooks', None) + if hooks is not None: + 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) diff --git a/comfy_extras/nodes_hooks.py b/comfy_extras/nodes_hooks.py index 23cfb63b5..33bc2c653 100644 --- a/comfy_extras/nodes_hooks.py +++ b/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)