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

Moved WrappersMP/CallbacksMP/WrapperExecutor to patcher_extension.py

Browse Source
This commit is contained in:
Jedrzej Kosinski 2024-11-02 22:21:16 -05:00
parent e3c3722ec2
commit 51e8d5554c
3 changed files with 95 additions and 87 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

82
comfy/model_patcher.py
View File

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

87
comfy/patcher_extension.py Normal file
View File

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

13
comfy/samplers.py
View File

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