xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-13 11:57:12 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jedrzej Kosinski 2024-10-25 18:32:22 -05:00
parent 4bbdf2bfe5
commit daeb2624a9
5 changed files with 38 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
comfy/hooks.py
View File

@ -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:

2
comfy/lora.py
View File

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

9
comfy/model_patcher.py
View File

@ -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:

12
comfy/samplers.py
View File

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

2
comfy_extras/nodes_hooks.py
View File

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