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

Added WrapperExecutor for non-classbound functions, added calc_cond_batch wrappers

Browse Source
This commit is contained in:
kosinkadink1@gmail.com 2024-09-22 11:52:35 +09:00
parent 298397d198
commit 5052a78be2
2 changed files with 44 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
comfy/model_patcher.py
View File

@ -127,11 +127,13 @@ class CallbacksMP:
class WrappersMP:
OUTER_SAMPLE = "outer_sample"
CALC_COND_BATCH = "calc_cond_batch"
@classmethod
def init_wrappers(cls):
return {
cls.OUTER_SAMPLE: [],
cls.CALC_COND_BATCH: [],
}
class WrapperExecutor:
@ -141,16 +143,16 @@ class WrapperExecutor:
self.idx = idx
self.is_last = idx == len(wrappers)
def __call__(self, guider, *args, **kwargs):
def __call__(self, *args, **kwargs):
new_executor = self._create_next_executor()
return new_executor._execute(guider, *args, **kwargs)
return new_executor._execute(*args, **kwargs)
def _execute(self, guider, *args, **kwargs):
def _execute(self, *args, **kwargs):
args = list(args)
kwargs = dict(kwargs)
if self.is_last:
return self.original(*args, **kwargs)
return self.wrappers[self.idx](self, guider, *args, **kwargs)
return self.wrappers[self.idx](self, *args, **kwargs)
def _create_next_executor(self):
new_idx = self.idx + 1
@ -162,6 +164,34 @@ class WrapperExecutor:
def new_executor(cls, original: Callable, wrappers: List[Callable]):
return cls(original, wrappers, idx=0)
class WrapperClassExecutor:
def __init__(self, original: Callable, wrappers: List[Callable], idx: int):
self.original = original
self.wrappers = wrappers.copy()
self.idx = idx
self.is_last = idx == len(wrappers)
def __call__(self, class_inst, *args, **kwargs):
new_executor = self._create_next_executor()
return new_executor._execute(class_inst, *args, **kwargs)
def _execute(self, class_inst, *args, **kwargs):
args = list(args)
kwargs = dict(kwargs)
if self.is_last:
return self.original(*args, **kwargs)
return self.wrappers[self.idx](self, class_inst, *args, **kwargs)
def _create_next_executor(self):
new_idx = self.idx + 1
if new_idx > len(self.wrappers):
raise Exception(f"Wrapper idx exceeded available wrappers; something went very wrong.")
return WrapperClassExecutor(self.original, self.wrappers, new_idx)
@classmethod
def new_executor(cls, original: Callable, wrappers: List[Callable]):
return cls(original, wrappers, idx=0)
class AutoPatcherEjector:
def __init__(self, model: 'ModelPatcher', skip_and_inject_on_exit_only=False):
self.model = model

12
comfy/samplers.py
View File

@ -183,6 +183,13 @@ def finalize_default_conds(hooked_to_run: Dict[comfy.hooks.HookGroup,List[Tuple[
hooked_to_run[hook] += [(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(
outer_calc_cond_batch,
model.current_patcher.get_wrappers(comfy.model_patcher.WrappersMP.CALC_COND_BATCH)
)
return executor._execute(model, conds, x_in, timestep, model_options)
def outer_calc_cond_batch(model: 'BaseModel', conds: List[List[Dict]], x_in: torch.Tensor, timestep, model_options):
out_conds = []
out_counts = []
# separate conds by matching hooks
@ -799,9 +806,10 @@ class CFGGuider:
try:
comfy.sampler_helpers.prepare_model_patcher(self.model_patcher, self.conds)
executor = comfy.model_patcher.WrapperExecutor.new_executor(
executor = comfy.model_patcher.WrapperClassExecutor.new_executor(
self.outer_sample,
self.model_patcher.get_wrappers(comfy.model_patcher.WrappersMP.OUTER_SAMPLE))
self.model_patcher.get_wrappers(comfy.model_patcher.WrappersMP.OUTER_SAMPLE)
)
output = executor._execute(self, noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
finally:
self.model_patcher.restore_hook_patches()