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

Updated different hook types to reflect actual progress of implementation, initial scaffolding for working WrapperHook functionality

Browse Source
This commit is contained in:
Jedrzej Kosinski 2024-11-11 10:13:02 -06:00
parent b12cc83c5b
commit 66b3386ee7
3 changed files with 45 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

73
comfy/hooks.py
View File

@ -12,6 +12,7 @@ if TYPE_CHECKING:
from comfy.sd import CLIP from comfy.sd import CLIP
import comfy.lora import comfy.lora
import comfy.model_management import comfy.model_management
import comfy.patcher_extension
from node_helpers import conditioning_set_values from node_helpers import conditioning_set_values
class EnumHookMode(enum.Enum): class EnumHookMode(enum.Enum):
@ -23,9 +24,9 @@ class EnumHookType(enum.Enum):
Patch = "patch" Patch = "patch"
ObjectPatch = "object_patch" ObjectPatch = "object_patch"
AddModels = "add_models" AddModels = "add_models"
AddCallback = "add_callback" Callbacks = "add_callback"
Wrappers = "add_wrapper"
SetInjections = "add_injections" SetInjections = "add_injections"
AddWrapper = "add_wrapper"
class EnumWeightTarget(enum.Enum): class EnumWeightTarget(enum.Enum):
Model = "model" Model = "model"
@ -35,15 +36,16 @@ class _HookRef:
pass pass
# NOTE: this is an example of how the should_register function should look # NOTE: this is an example of how the should_register function should look
def default_should_register(hook: 'Hook', model: 'ModelPatcher', target: EnumWeightTarget, registered: list[Hook]): def default_should_register(hook: 'Hook', model: 'ModelPatcher', model_options: dict, target: EnumWeightTarget, registered: list[Hook]):
return True return True
class Hook: class Hook:
def __init__(self, hook_type: EnumHookType=None, hook_ref: _HookRef=None, def __init__(self, hook_type: EnumHookType=None, hook_ref: _HookRef=None, hook_id: str=None,
hook_keyframe: 'HookKeyframeGroup'=None): hook_keyframe: 'HookKeyframeGroup'=None):
self.hook_type = hook_type self.hook_type = hook_type
self.hook_ref = hook_ref if hook_ref else _HookRef() self.hook_ref = hook_ref if hook_ref else _HookRef()
self.hook_id = hook_id
self.hook_keyframe = hook_keyframe if hook_keyframe else HookKeyframeGroup() self.hook_keyframe = hook_keyframe if hook_keyframe else HookKeyframeGroup()
self.custom_should_register = default_should_register self.custom_should_register = default_should_register
self.auto_apply_to_nonpositive = False self.auto_apply_to_nonpositive = False
@ -65,16 +67,17 @@ class Hook:
c: Hook = subtype() c: Hook = subtype()
c.hook_type = self.hook_type c.hook_type = self.hook_type
c.hook_ref = self.hook_ref c.hook_ref = self.hook_ref
c.hook_id = self.hook_id
c.hook_keyframe = self.hook_keyframe c.hook_keyframe = self.hook_keyframe
c.custom_should_register = self.custom_should_register c.custom_should_register = self.custom_should_register
# TODO: make this do something # TODO: make this do something
c.auto_apply_to_nonpositive = self.auto_apply_to_nonpositive c.auto_apply_to_nonpositive = self.auto_apply_to_nonpositive
return c return c
def should_register(self, model: 'ModelPatcher', target: EnumWeightTarget, registered: list[Hook]): def should_register(self, model: 'ModelPatcher', model_options: dict, target: EnumWeightTarget, registered: list[Hook]):
return self.custom_should_register(self, model, target, registered) return self.custom_should_register(self, model, model_options, target, registered)
def add_hook_patches(self, model: 'ModelPatcher', target: EnumWeightTarget, registered: list[Hook]): def add_hook_patches(self, model: 'ModelPatcher', model_options: dict, target: EnumWeightTarget, registered: list[Hook]):
raise NotImplementedError("add_hook_patches should be defined for Hook subclasses") raise NotImplementedError("add_hook_patches should be defined for Hook subclasses")
def on_apply(self, model: 'ModelPatcher', transformer_options: dict[str]): def on_apply(self, model: 'ModelPatcher', transformer_options: dict[str]):
@ -107,8 +110,8 @@ class WeightHook(Hook):
def strength_clip(self): def strength_clip(self):
return self._strength_clip * self.strength return self._strength_clip * self.strength
def add_hook_patches(self, model: 'ModelPatcher', target: EnumWeightTarget, registered: list[Hook]): def add_hook_patches(self, model: 'ModelPatcher', model_options: dict, target: EnumWeightTarget, registered: list[Hook]):
if not self.should_register(model, target, registered): if not self.should_register(model, model_options, target, registered):
return False return False
weights = None weights = None
if target == EnumWeightTarget.Model: if target == EnumWeightTarget.Model:
@ -155,9 +158,7 @@ class PatchHook(Hook):
c: PatchHook = super().clone(subtype) c: PatchHook = super().clone(subtype)
c.patches = self.patches c.patches = self.patches
return c return c
# TODO: add functionality
def add_hook_patches(self, model: 'ModelPatcher'):
pass
class ObjectPatchHook(Hook): class ObjectPatchHook(Hook):
def __init__(self): def __init__(self):
@ -170,9 +171,7 @@ class ObjectPatchHook(Hook):
c: ObjectPatchHook = super().clone(subtype) c: ObjectPatchHook = super().clone(subtype)
c.object_patches = self.object_patches c.object_patches = self.object_patches
return c return c
# TODO: add functionality
def add_hook_object_patches(self, model: 'ModelPatcher'):
pass
class AddModelsHook(Hook): class AddModelsHook(Hook):
def __init__(self, key: str=None, models: list['ModelPatcher']=None): def __init__(self, key: str=None, models: list['ModelPatcher']=None):
@ -189,13 +188,11 @@ class AddModelsHook(Hook):
c.models = self.models.copy() if self.models else self.models c.models = self.models.copy() if self.models else self.models
c.append_when_same = self.append_when_same c.append_when_same = self.append_when_same
return c return c
# TODO: add functionality
def add_hook_models(self, model: 'ModelPatcher'):
pass
class CallbackHook(Hook): class CallbackHook(Hook):
def __init__(self, key: str=None, callback: Callable=None): def __init__(self, key: str=None, callback: Callable=None):
super().__init__(hook_type=EnumHookType.AddCallback) super().__init__(hook_type=EnumHookType.Callbacks)
self.key = key self.key = key
self.callback = callback self.callback = callback
@ -206,9 +203,25 @@ class CallbackHook(Hook):
c.key = self.key c.key = self.key
c.callback = self.callback c.callback = self.callback
return c return c
# TODO: add functionality
class WrapperHook(Hook):
def __init__(self, wrappers_dict: dict[str, dict[str, list[Callable]]]=None):
super().__init__(hook_type=EnumHookType.Wrappers)
self.wrappers_dict = wrappers_dict
def clone(self, subtype: Callable=None):
if subtype is None:
subtype = type(self)
c: WrapperHook = super().clone(subtype)
c.wrappers_dict = self.wrappers_dict
return c
def add_hook_callback(self, model: 'ModelPatcher'): def add_hook_wrapper(self, model: 'ModelPatcher', model_options: dict, target: EnumWeightTarget, registered: list[Hook]):
pass if not self.should_register(model, model_options, target, registered):
return False
add_model_options = {"transformer_options": {"wrappers": self.wrappers_dict}}
comfy.patcher_extension.merge_nested_dicts(model_options, add_model_options, copy_dict1=False)
class SetInjectionsHook(Hook): class SetInjectionsHook(Hook):
def __init__(self, key: str=None, injections: list['PatcherInjection']=None): def __init__(self, key: str=None, injections: list['PatcherInjection']=None):
@ -225,23 +238,7 @@ class SetInjectionsHook(Hook):
return c return c
def add_hook_injections(self, model: 'ModelPatcher'): def add_hook_injections(self, model: 'ModelPatcher'):
pass # TODO: add functionality
class WrapperHook(Hook):
def __init__(self, key: str=None, wrapper: Callable=None):
super().__init__(hook_type=EnumHookType.AddWrapper)
self.key = key
self.wrapper = wrapper
def clone(self, subtype: Callable=None):
if subtype is None:
subtype = type(self)
c: WrapperHook = super().clone(subtype)
c.key = self.key
c.wrapper = self.wrapper
return c
def add_hook_wrapper(self, model: 'ModelPatcher'):
pass pass
class HookGroup: class HookGroup:

11
comfy/model_patcher.py
View File

@ -897,10 +897,15 @@ class ModelPatcher:
if cached_group.contains(hook): if cached_group.contains(hook):
self.cached_hook_patches.pop(cached_group) self.cached_hook_patches.pop(cached_group)
def register_all_hook_patches(self, hooks_dict: dict[comfy.hooks.EnumHookType, dict[comfy.hooks.Hook, None]], target: comfy.hooks.EnumWeightTarget): def register_all_hook_patches(self, hooks_dict: dict[comfy.hooks.EnumHookType, dict[comfy.hooks.Hook, None]], target: comfy.hooks.EnumWeightTarget, model_options: dict=None):
self.restore_hook_patches() self.restore_hook_patches()
weight_hooks_to_register: list[comfy.hooks.WeightHook] = []
registered_hooks: list[comfy.hooks.Hook] = [] registered_hooks: list[comfy.hooks.Hook] = []
# handle WrapperHooks, if model_options provided
if model_options is not None:
for hook in hooks_dict.get(comfy.hooks.EnumHookType.Wrappers, {}):
hook.add_hook_patches(self, model_options, target, registered_hooks)
# handle WeightHooks
weight_hooks_to_register: list[comfy.hooks.WeightHook] = []
for hook in hooks_dict.get(comfy.hooks.EnumHookType.Weight, {}): for hook in hooks_dict.get(comfy.hooks.EnumHookType.Weight, {}):
if hook.hook_ref not in self.hook_patches: if hook.hook_ref not in self.hook_patches:
weight_hooks_to_register.append(hook) weight_hooks_to_register.append(hook)
@ -908,7 +913,7 @@ class ModelPatcher:
# clone hook_patches to become backup so that any non-dynamic hooks will return to their original state # clone hook_patches to become backup so that any non-dynamic hooks will return to their original state
self.hook_patches_backup = create_hook_patches_clone(self.hook_patches) self.hook_patches_backup = create_hook_patches_clone(self.hook_patches)
for hook in weight_hooks_to_register: for hook in weight_hooks_to_register:
hook.add_hook_patches(self, target, registered_hooks) hook.add_hook_patches(self, model_options, target, registered_hooks)
for callback in self.get_all_callbacks(CallbacksMP.ON_REGISTER_ALL_HOOK_PATCHES): for callback in self.get_all_callbacks(CallbacksMP.ON_REGISTER_ALL_HOOK_PATCHES):
callback(self, hooks_dict, target) callback(self, hooks_dict, target)

4
comfy/sampler_helpers.py
View File

@ -133,8 +133,8 @@ def prepare_model_patcher(model: 'ModelPatcher', conds, model_options: dict):
hooks = {} hooks = {}
for k in conds: for k in conds:
get_hooks_from_cond(conds[k], hooks) get_hooks_from_cond(conds[k], hooks)
model.register_all_hook_patches(hooks, comfy.hooks.EnumWeightTarget.Model)
# add wrappers and callbacks from ModelPatcher to transformer_options # add wrappers and callbacks from ModelPatcher to transformer_options
model_options["transformer_options"]["wrappers"] = comfy.patcher_extension.copy_nested_dicts(model.wrappers) model_options["transformer_options"]["wrappers"] = comfy.patcher_extension.copy_nested_dicts(model.wrappers)
model_options["transformer_options"]["callbacks"] = comfy.patcher_extension.copy_nested_dicts(model.callbacks) model_options["transformer_options"]["callbacks"] = comfy.patcher_extension.copy_nested_dicts(model.callbacks)
# TODO: add wrappers and callbacks from registered hooks for functions called prior to calc_batch_conds # register hooks on model/model_options
model.register_all_hook_patches(hooks, comfy.hooks.EnumWeightTarget.Model, model_options)