From 51e8d5554c1b4d35ca739763cf1d4f8d989c4480 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sat, 2 Nov 2024 22:21:16 -0500 Subject: [PATCH] Moved WrappersMP/CallbacksMP/WrapperExecutor to patcher_extension.py --- comfy/model_patcher.py | 82 +---------------------------------- comfy/patcher_extension.py | 87 ++++++++++++++++++++++++++++++++++++++ comfy/samplers.py | 13 +++--- 3 files changed, 95 insertions(+), 87 deletions(-) create mode 100644 comfy/patcher_extension.py diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index dfd578371..f6e5bab83 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -31,6 +31,7 @@ import comfy.float import comfy.model_management import comfy.lora import comfy.hooks +from comfy.patcher_extension import CallbacksMP, WrappersMP, PatcherInjection from comfy.comfy_types import UnetWrapperFunction def string_to_seed(data): @@ -141,82 +142,6 @@ def get_key_weight(model, key): return weight, set_func, convert_func -class CallbacksMP: - ON_CLONE = "on_clone" - ON_LOAD = "on_load_after" - ON_CLEANUP = "on_cleanup" - ON_PRE_RUN = "on_pre_run" - ON_PREPARE_STATE = "on_prepare_state" - ON_APPLY_HOOKS = "on_apply_hooks" - ON_REGISTER_ALL_HOOK_PATCHES = "on_register_all_hook_patches" - ON_INJECT_MODEL = "on_inject_model" - ON_EJECT_MODEL = "on_eject_model" - - @classmethod - def init_callbacks(cls): - return { - cls.ON_CLONE: {None: []}, - cls.ON_LOAD: {None: []}, - cls.ON_CLEANUP: {None: []}, - cls.ON_PRE_RUN: {None: []}, - cls.ON_PREPARE_STATE: {None: []}, - cls.ON_APPLY_HOOKS: {None: []}, - cls.ON_REGISTER_ALL_HOOK_PATCHES: {None: []}, - cls.ON_INJECT_MODEL: {None: []}, - cls.ON_EJECT_MODEL: {None: []}, - } - -class WrappersMP: - OUTER_SAMPLE = "outer_sample" - CALC_COND_BATCH = "calc_cond_batch" - SAMPLER_SAMPLE = "sampler_sample" - - @classmethod - def init_wrappers(cls): - return { - cls.OUTER_SAMPLE: {None: []}, - cls.SAMPLER_SAMPLE: {None: []}, - cls.CALC_COND_BATCH: {None: []}, - } - -class WrapperExecutor: - """Handles call stack of wrappers around a function in an ordered manner.""" - def __init__(self, original: Callable, class_obj: object, wrappers: list[Callable], idx: int): - self.original = original - self.class_obj = class_obj - self.wrappers = wrappers.copy() - self.idx = idx - self.is_last = idx == len(wrappers) - - def __call__(self, *args, **kwargs): - """Calls the next wrapper in line or original function, whichever is appropriate.""" - new_executor = self._create_next_executor() - return new_executor.execute(*args, **kwargs) - - def execute(self, *args, **kwargs): - """Used to initiate executor internally - DO NOT use this if you received executor in wrapper.""" - args = list(args) - kwargs = dict(kwargs) - if self.is_last: - return self.original(*args, **kwargs) - return self.wrappers[self.idx](self, *args, **kwargs) - - def _create_next_executor(self) -> 'WrapperExecutor': - new_idx = self.idx + 1 - if new_idx > len(self.wrappers): - raise Exception(f"Wrapper idx exceeded available wrappers; something went very wrong.") - if self.class_obj is None: - return WrapperExecutor.new_executor(self.original, self.wrappers, new_idx) - return WrapperExecutor.new_class_executor(self.original, self.class_obj, self.wrappers, new_idx) - - @classmethod - def new_executor(cls, original: Callable, wrappers: list[Callable], idx=0): - return cls(original, class_obj=None, wrappers=wrappers, idx=idx) - - @classmethod - def new_class_executor(cls, original: Callable, class_obj: object, wrappers: list[Callable], idx=0): - return cls(original, class_obj, wrappers, idx=idx) - class AutoPatcherEjector: def __init__(self, model: 'ModelPatcher', skip_and_inject_on_exit_only=False): self.model = model @@ -241,11 +166,6 @@ class AutoPatcherEjector: self.model.inject_model() self.model.skip_injection = self.prev_skip_injection -class PatcherInjection: - def __init__(self, inject: Callable, eject: Callable): - self.inject = inject - self.eject = eject - class MemoryCounter: def __init__(self, initial: int, minimum=0): self.value = initial diff --git a/comfy/patcher_extension.py b/comfy/patcher_extension.py new file mode 100644 index 000000000..980fb8eb6 --- /dev/null +++ b/comfy/patcher_extension.py @@ -0,0 +1,87 @@ +from __future__ import annotations +from typing import Callable + +class CallbacksMP: + ON_CLONE = "on_clone" + ON_LOAD = "on_load_after" + ON_CLEANUP = "on_cleanup" + ON_PRE_RUN = "on_pre_run" + ON_PREPARE_STATE = "on_prepare_state" + ON_APPLY_HOOKS = "on_apply_hooks" + ON_REGISTER_ALL_HOOK_PATCHES = "on_register_all_hook_patches" + ON_INJECT_MODEL = "on_inject_model" + ON_EJECT_MODEL = "on_eject_model" + + @classmethod + def init_callbacks(cls): + return { + cls.ON_CLONE: {None: []}, + cls.ON_LOAD: {None: []}, + cls.ON_CLEANUP: {None: []}, + cls.ON_PRE_RUN: {None: []}, + cls.ON_PREPARE_STATE: {None: []}, + cls.ON_APPLY_HOOKS: {None: []}, + cls.ON_REGISTER_ALL_HOOK_PATCHES: {None: []}, + cls.ON_INJECT_MODEL: {None: []}, + cls.ON_EJECT_MODEL: {None: []}, + } + +class WrappersMP: + OUTER_SAMPLE = "outer_sample" + SAMPLER_SAMPLE = "sampler_sample" + CALC_COND_BATCH = "calc_cond_batch" + APPLY_MODEL = "apply_model" + DIFFUSION_MODEL = "diffusion_model" + + @classmethod + def init_wrappers(cls): + return { + cls.OUTER_SAMPLE: {None: []}, + cls.SAMPLER_SAMPLE: {None: []}, + cls.CALC_COND_BATCH: {None: []}, + cls.APPLY_MODEL: {None: []}, + cls.DIFFUSION_MODEL: {None: []}, + } + +class WrapperExecutor: + """Handles call stack of wrappers around a function in an ordered manner.""" + def __init__(self, original: Callable, class_obj: object, wrappers: list[Callable], idx: int): + self.original = original + self.class_obj = class_obj + self.wrappers = wrappers.copy() + self.idx = idx + self.is_last = idx == len(wrappers) + + def __call__(self, *args, **kwargs): + """Calls the next wrapper in line or original function, whichever is appropriate.""" + new_executor = self._create_next_executor() + return new_executor.execute(*args, **kwargs) + + def execute(self, *args, **kwargs): + """Used to initiate executor internally - DO NOT use this if you received executor in wrapper.""" + args = list(args) + kwargs = dict(kwargs) + if self.is_last: + return self.original(*args, **kwargs) + return self.wrappers[self.idx](self, *args, **kwargs) + + def _create_next_executor(self) -> 'WrapperExecutor': + new_idx = self.idx + 1 + if new_idx > len(self.wrappers): + raise Exception(f"Wrapper idx exceeded available wrappers; something went very wrong.") + if self.class_obj is None: + return WrapperExecutor.new_executor(self.original, self.wrappers, new_idx) + return WrapperExecutor.new_class_executor(self.original, self.class_obj, self.wrappers, new_idx) + + @classmethod + def new_executor(cls, original: Callable, wrappers: list[Callable], idx=0): + return cls(original, class_obj=None, wrappers=wrappers, idx=idx) + + @classmethod + def new_class_executor(cls, original: Callable, class_obj: object, wrappers: list[Callable], idx=0): + return cls(original, class_obj, wrappers, idx=idx) + +class PatcherInjection: + def __init__(self, inject: Callable, eject: Callable): + self.inject = inject + self.eject = eject diff --git a/comfy/samplers.py b/comfy/samplers.py index b616cfe75..943b7cbb9 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -11,6 +11,7 @@ import math import logging import comfy.sampler_helpers import comfy.model_patcher +import comfy.patcher_extension import comfy.hooks import scipy.stats import numpy @@ -185,9 +186,9 @@ def finalize_default_conds(model: 'BaseModel', hooked_to_run: dict[comfy.hooks.H 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( + executor = comfy.patcher_extension.WrapperExecutor.new_executor( outer_calc_cond_batch, - model.current_patcher.get_all_wrappers(comfy.model_patcher.WrappersMP.CALC_COND_BATCH) + model.current_patcher.get_all_wrappers(comfy.patcher_extension.WrappersMP.CALC_COND_BATCH) ) return executor.execute(model, conds, x_in, timestep, model_options) @@ -804,10 +805,10 @@ class CFGGuider: extra_args = {"model_options": comfy.model_patcher.create_model_options_clone(self.model_options), "seed": seed} - executor = comfy.model_patcher.WrapperExecutor.new_class_executor( + executor = comfy.patcher_extension.WrapperExecutor.new_class_executor( sampler.sample, sampler, - self.model_patcher.get_all_wrappers(comfy.model_patcher.WrappersMP.SAMPLER_SAMPLE) + self.model_patcher.get_all_wrappers(comfy.patcher_extension.WrappersMP.SAMPLER_SAMPLE) ) samples = executor.execute(self, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar) return self.inner_model.process_latent_out(samples.to(torch.float32)) @@ -844,10 +845,10 @@ class CFGGuider: try: comfy.sampler_helpers.prepare_model_patcher(self.model_patcher, self.conds) - executor = comfy.model_patcher.WrapperExecutor.new_class_executor( + executor = comfy.patcher_extension.WrapperExecutor.new_class_executor( self.outer_sample, self, - self.model_patcher.get_all_wrappers(comfy.model_patcher.WrappersMP.OUTER_SAMPLE) + self.model_patcher.get_all_wrappers(comfy.patcher_extension.WrappersMP.OUTER_SAMPLE) ) output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed) finally: