diff --git a/comfy/controlnet.py b/comfy/controlnet.py index 3d0e25e89..d13c34b01 100644 --- a/comfy/controlnet.py +++ b/comfy/controlnet.py @@ -118,6 +118,14 @@ class ControlBase: if self.previous_controlnet is not None: out += self.previous_controlnet.get_models() return out + + def get_extra_hooks(self): + out = [] + if self.extra_hooks is not None: + out.append(self.extra_hooks) + if self.previous_controlnet is not None: + out += self.previous_controlnet.get_extra_hooks() + return out def copy_to(self, c): c.cond_hint_original = self.cond_hint_original diff --git a/comfy/hooks.py b/comfy/hooks.py index de6a38c5b..2913f7812 100644 --- a/comfy/hooks.py +++ b/comfy/hooks.py @@ -54,7 +54,7 @@ class Hook: def initialize_timesteps(self, model: 'BaseModel'): self.reset() - self.hook_keyframe.initalize_timesteps(model) + self.hook_keyframe.initialize_timesteps(model) def reset(self): self.hook_keyframe.reset() @@ -193,7 +193,7 @@ class AddModelsHook(Hook): def add_hook_models(self, model: 'ModelPatcher'): pass -class AddCallbackHook(Hook): +class CallbackHook(Hook): def __init__(self, key: str=None, callback: Callable=None): super().__init__(hook_type=EnumHookType.AddCallback) self.key = key @@ -202,7 +202,7 @@ class AddCallbackHook(Hook): def clone(self, subtype: Callable=None): if subtype is None: subtype = type(self) - c: AddCallbackHook = super().clone(subtype) + c: CallbackHook = super().clone(subtype) c.key = self.key c.callback = self.callback return c @@ -227,7 +227,7 @@ class SetInjectionsHook(Hook): def add_hook_injections(self, model: 'ModelPatcher'): pass -class AddWrapperHook(Hook): +class WrapperHook(Hook): def __init__(self, key: str=None, wrapper: Callable=None): super().__init__(hook_type=EnumHookType.AddWrapper) self.key = key @@ -236,7 +236,7 @@ class AddWrapperHook(Hook): def clone(self, subtype: Callable=None): if subtype is None: subtype = type(self) - c: AddWrapperHook = super().clone(subtype) + c: WrapperHook = super().clone(subtype) c.key = self.key c.wrapper = self.wrapper return c @@ -268,7 +268,10 @@ class HookGroup: return c def set_keyframes_on_hooks(self, hook_kf: 'HookKeyframeGroup'): - hook_kf = hook_kf.clone() + if hook_kf is None: + hook_kf = HookKeyframeGroup() + else: + hook_kf = hook_kf.clone() for hook in self.hooks: hook.hook_keyframe = hook_kf @@ -330,7 +333,7 @@ class HookGroup: def reset(self): for hook in self.hooks: - hook.hook_keyframe.reset() + hook.reset() @staticmethod def combine_all_hooks(hooks_list: list['HookGroup'], require_count=0) -> 'HookGroup': @@ -414,11 +417,11 @@ class HookKeyframeGroup: def clone(self): c = HookKeyframeGroup() for keyframe in self.keyframes: - c.keyframes.append(keyframe) + c.keyframes.append(keyframe.clone()) c._set_first_as_current() return c - def initalize_timesteps(self, model: 'BaseModel'): + def initialize_timesteps(self, model: 'BaseModel'): for keyframe in self.keyframes: keyframe.start_t = model.model_sampling.percent_to_sigma(keyframe.start_percent) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 21c758858..d0474f39a 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -990,12 +990,14 @@ class ModelPatcher: combined_patches[key] = current_patches return combined_patches - def apply_hooks(self, hooks: comfy.hooks.HookGroup): + def apply_hooks(self, hooks: comfy.hooks.HookGroup, model_options: dict=None): + # TODO: return transformer_options dict with any additions from hooks if self.current_hooks == hooks: - return + return {} self.patch_hooks(hooks=hooks) for callback in self.get_all_callbacks(CallbacksMP.ON_APPLY_HOOKS): callback(self, hooks) + return {} def patch_hooks(self, hooks: comfy.hooks.HookGroup): with self.use_ejected(): diff --git a/comfy/patcher_extension.py b/comfy/patcher_extension.py index 096f556bb..2a31e57ff 100644 --- a/comfy/patcher_extension.py +++ b/comfy/patcher_extension.py @@ -139,8 +139,11 @@ def copy_nested_dicts(input_dict: dict): new_dict[key] = value.copy() return new_dict -def merge_nested_dicts(dict1: dict, dict2: dict): - merged_dict = copy_nested_dicts(dict1) +def merge_nested_dicts(dict1: dict, dict2: dict, copy_dict1=True): + if copy_dict1: + merged_dict = copy_nested_dicts(dict1) + else: + merged_dict = dict1 for key, value in dict2.items(): if isinstance(value, dict): curr_value = merged_dict.setdefault(key, {}) diff --git a/comfy/samplers.py b/comfy/samplers.py index 5ae56ce85..c94082058 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -1,9 +1,11 @@ +from __future__ import annotations from .k_diffusion import sampling as k_diffusion_sampling from .extra_samplers import uni_pc from typing import TYPE_CHECKING if TYPE_CHECKING: from comfy.model_patcher import ModelPatcher from comfy.model_base import BaseModel + from comfy.controlnet import ControlBase import torch import collections from comfy import model_management @@ -248,8 +250,6 @@ def _calc_cond_batch(model: 'BaseModel', conds: list[list[dict]], x_in: torch.Te to_batch = batch_amount break - model.current_patcher.apply_hooks(hooks=hooks) - input_x = [] mult = [] c = [] @@ -275,11 +275,14 @@ def _calc_cond_batch(model: 'BaseModel', conds: list[list[dict]], x_in: torch.Te c = cond_cat(c) timestep_ = torch.cat([timestep] * batch_chunks) - transformer_options = {} + transformer_options = model.current_patcher.apply_hooks(hooks=hooks) if 'transformer_options' in model_options: - transformer_options = model_options['transformer_options'].copy() + transformer_options = comfy.patcher_extension.merge_nested_dicts(transformer_options, + model_options['transformer_options'], + copy_dict1=False) if patches is not None: + # TODO: replace with merge_nested_dicts function if "patches" in transformer_options: cur_patches = transformer_options["patches"].copy() for p in patches: @@ -774,6 +777,34 @@ def process_conds(model, noise, conds, device, latent_image=None, denoise_mask=N return conds + +def preprocess_conds_hooks(conds: dict[str, list[dict[str]]]): + # determine which ControlNets have extra_hooks that should be combined with normal hooks + hook_replacement: dict[tuple[ControlBase, comfy.hooks.HookGroup], list[dict]] = {} + for k in conds: + for kk in conds[k]: + if 'control' in kk: + control: 'ControlBase' = kk['control'] + extra_hooks = control.get_extra_hooks() + if len(extra_hooks) > 0: + hooks: comfy.hooks.HookGroup = kk.get('hooks', None) + to_replace = hook_replacement.setdefault((control, hooks), []) + to_replace.append(kk) + # if nothing to replace, do nothing + if len(hook_replacement) == 0: + return + + # for optimal sampling performance, common ControlNets + hook combos should have identical hooks + # on the cond dicts + for key, conds_to_modify in hook_replacement.items(): + control = key[0] + hooks = key[1] + hooks = comfy.hooks.HookGroup.combine_all_hooks(control.get_extra_hooks() + [hooks]) + # if combined hooks are not None, set as new hooks for all relevant conds + if hooks is not None: + for cond in conds_to_modify: + cond['hooks'] = hooks + class CFGGuider: def __init__(self, model_patcher): self.model_patcher: 'ModelPatcher' = model_patcher @@ -842,6 +873,7 @@ class CFGGuider: self.conds = {} for k in self.original_conds: self.conds[k] = list(map(lambda a: a.copy(), self.original_conds[k])) + preprocess_conds_hooks(self.conds) try: orig_model_options = self.model_options diff --git a/comfy/sd.py b/comfy/sd.py index 04b108fb3..a09e7d84f 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -4,6 +4,7 @@ from enum import Enum import logging from comfy import model_management +from comfy.utils import ProgressBar from .ldm.models.autoencoder import AutoencoderKL, AutoencodingEngine from .ldm.cascade.stage_a import StageA from .ldm.cascade.stage_c_coder import StageC_coder @@ -119,7 +120,7 @@ class CLIP: def tokenize(self, text, return_word_ids=False): return self.tokenizer.tokenize_with_weights(text, return_word_ids) - def encode_from_tokens_scheduled(self, tokens, unprojected=False, add_dict: dict[str]=None): + def encode_from_tokens_scheduled(self, tokens, unprojected=False, add_dict: dict[str]=None, show_pbar=True): all_cond_pooled: list[tuple[torch.Tensor, dict[str]]] = [] all_hooks = self.patcher.forced_hooks scheduled_keyframes = all_hooks.get_hooks_for_clip_schedule() @@ -132,6 +133,10 @@ class CLIP: self.load_model() all_hooks.reset() + self.patcher.patch_hooks(None) + if show_pbar: + pbar = ProgressBar(len(scheduled_keyframes)) + for scheduled_opts in scheduled_keyframes: t_range = scheduled_opts[0] hooks_keyframes = scheduled_opts[1] @@ -150,6 +155,9 @@ class CLIP: if add_dict is not None: pooled_dict.update(add_dict) all_cond_pooled.append([cond, pooled_dict]) + if show_pbar: + pbar.update(1) + model_management.throw_exception_if_processing_interrupted() all_hooks.reset() return all_cond_pooled diff --git a/comfy_extras/nodes_hooks.py b/comfy_extras/nodes_hooks.py index 8d933ba08..65eb08eff 100644 --- a/comfy_extras/nodes_hooks.py +++ b/comfy_extras/nodes_hooks.py @@ -243,7 +243,9 @@ class SetClipHooks: if hooks is not None: clip = clip.clone() clip.use_clip_schedule = schedule_clip - clip.patcher.forced_hooks = hooks + clip.patcher.forced_hooks = hooks.clone() + if not clip.use_clip_schedule: + clip.patcher.forced_hooks.set_keyframes_on_hooks(None) clip.patcher.register_all_hook_patches(hooks.get_dict_repr(), comfy.hooks.EnumWeightTarget.Clip) return (clip,)