Moved WrappersMP/CallbacksMP/WrapperExecutor to patcher_extension.py

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

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

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: