diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 7540f697b..f69a7596f 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -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 diff --git a/comfy/samplers.py b/comfy/samplers.py index a7b59f433..31ffb8dde 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -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()