Added injections support to ModelPatcher + necessary bookkeeping, added additional_models support in ModelPatcher, conds, and hooks

comfy/hooks.py

@@ -19,6 +19,7 @@ class EnumHookMode(enum.Enum):
 class EnumHookType(enum.Enum):
     Weight = "weight"
     Patch = "patch"
+    AddModel = "addmodel"
 
 class EnumWeightTarget(enum.Enum):
     Model = "model"
@@ -121,10 +122,22 @@ class PatchHook(Hook):
     def clone(self, subtype: Callable=None):
         if subtype is None:
             subtype = type(self)
-        c: PatchHook = super().clone(type(self))
+        c: PatchHook = super().clone(subtype)
         c.patches = self.patches
         return c
 
+class AddModelHook(Hook):
+    def __init__(self, model: 'ModelPatcher'):
+        super().__init__(hook_type=EnumHookType.AddModel)
+        self.model = model
+    
+    def clone(self, subtype: Callable=None):
+        if subtype is None:
+            subtype = type(self)
+        c: AddModelHook = super().clone(subtype)
+        c.model = self.model
+        return c
+
 class HookGroup:
     def __init__(self):
         self.hooks: List[Hook] = []

comfy/model_patcher.py

@@ -108,6 +108,8 @@ class CallbacksMP:
     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):
@@ -119,8 +121,37 @@ class CallbacksMP:
             cls.ON_PREPARE_STATE: [],
             cls.ON_APPLY_HOOKS: [],
             cls.ON_REGISTER_ALL_HOOK_PATCHES: [],
+            cls.ON_INJECT_MODEL: [],
+            cls.ON_EJECT_MODEL: [],
         }
 
+class AutoPatcherEjector:
+    def __init__(self, model: 'ModelPatcher', skip_until_exit=False):
+        self.model = model
+        self.was_injected = False
+        self.prev_skip_injection = False
+        self.skip_until_exit = skip_until_exit
+    
+    def __enter__(self):
+        self.was_injected = False
+        self.prev_skip_injection = self.model.skip_injection
+        if self.skip_until_exit:
+            self.model.skip_injection = True
+        if self.model.is_injected:
+            self.model.eject_model()
+            self.was_injected = True
+
+    def __exit__(self, *args):
+        if self.was_injected:
+            if self.skip_until_exit or not self.model.skip_injection:
+                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 ModelPatcher:
     def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False):
         self.size = size
@@ -143,9 +174,13 @@ class ModelPatcher:
         self.patches_uuid = uuid.uuid4()
 
         self.attachments: Dict[str] = {}
-        self.additional_models: list[ModelPatcher] = []
+        self.additional_models: Dict[str, List[ModelPatcher]] = {}
         self.callbacks: Dict[str, List[Callable]] = CallbacksMP.init_callbacks()
 
+        self.is_injected = False
+        self.skip_injection = False
+        self.injections: Dict[str, List[PatcherInjection]] = {}
+
         self.hook_patches: Dict[comfy.hooks._HookRef] = {}
         self.hook_patches_backup: Dict[comfy.hooks._HookRef] = {}
         self.hook_backup: Dict[str, Tuple[torch.Tensor, torch.device]] = {}
@@ -196,11 +231,16 @@ class ModelPatcher:
             else:
                 n.attachments[k] = self.attachments[k]
         # additional models
-        for m in self.additional_models:
+        for k, c in self.additional_models.items():
-            n.additional_models.append(m.clone())
+            n.additional_models[k] = [x.clone() for x in c]
         # callbacks
         for k, c in self.callbacks.items():
             n.callbacks[k] = c.copy()
+        # injection
+        n.is_injected = self.is_injected
+        n.skip_injection = self.skip_injection
+        for k, i in self.injections.items():
+            n.injections[k] = i.copy()
         # hooks
         n.hook_patches = create_hook_patches_clone(self.hook_patches)
         n.hook_patches_backup = create_hook_patches_clone(self.hook_patches_backup)
@@ -342,6 +382,7 @@ class ModelPatcher:
             return self.model.get_dtype()
 
     def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
+        with self.use_ejected():
             p = set()
             model_sd = self.model.state_dict()
             for k in patches:
@@ -386,6 +427,7 @@ class ModelPatcher:
         return p
 
     def model_state_dict(self, filter_prefix=None):
+        with self.use_ejected():
             sd = self.model.state_dict()
             keys = list(sd.keys())
             if filter_prefix is not None:
@@ -417,6 +459,7 @@ class ModelPatcher:
             comfy.utils.set_attr_param(self.model, key, out_weight)
 
     def load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False, full_load=False):
+        with self.use_ejected():
             self.unpatch_hooks()
             mem_counter = 0
             patch_counter = 0
@@ -501,9 +544,14 @@ class ModelPatcher:
             self.model.lowvram_patch_counter += patch_counter
             self.model.device = device_to
             self.model.model_loaded_weight_memory = mem_counter
+
+            for callback in self.callbacks[CallbacksMP.ON_LOAD]:
+                callback(self, device_to, lowvram_model_memory, force_patch_weights, full_load)
+
             self.apply_hooks(self.forced_hooks)
 
     def patch_model(self, device_to=None, lowvram_model_memory=0, load_weights=True, force_patch_weights=False):
+        with self.use_ejected():
             for k in self.object_patches:
                 old = comfy.utils.set_attr(self.model, k, self.object_patches[k])
                 if k not in self.object_patches_backup:
@@ -516,9 +564,11 @@ class ModelPatcher:
 
             if load_weights:
                 self.load(device_to, lowvram_model_memory=lowvram_model_memory, force_patch_weights=force_patch_weights, full_load=full_load)
+        self.inject_model()
         return self.model
 
     def unpatch_model(self, device_to=None, unpatch_weights=True):
+        self.eject_model()
         if unpatch_weights:
             self.unpatch_hooks()
             if self.model.model_lowvram:
@@ -555,6 +605,7 @@ class ModelPatcher:
         self.object_patches_backup.clear()
 
     def partially_unload(self, device_to, memory_to_free=0):
+        with self.use_ejected():
             memory_freed = 0
             patch_counter = 0
             unload_list = []
@@ -605,6 +656,7 @@ class ModelPatcher:
             return memory_freed
 
     def partially_load(self, device_to, extra_memory=0):
+        with self.use_ejected(skip_injection=True):
             self.unpatch_model(unpatch_weights=False)
             self.patch_model(load_weights=False)
             full_load = False
@@ -629,13 +681,49 @@ class ModelPatcher:
         for callback in self.callbacks[CallbacksMP.ON_CLEANUP]:
             callback(self)
 
-    def add_callback(self, key, callback: Callable):
+    def get_all_additional_models(self):
+        all_models = []
+        for models in self.additional_models.values():
+            all_models.extend(models)
+        return all_models
+
+    def add_callback(self, key: str, callback: Callable):
         if key not in self.callbacks:
             raise Exception(f"Callback '{key}' is not recognized.")
         self.callbacks[key].append(callback)
     
-    def add_attachment(self, attachment):
+    def set_attachments(self, key: str, attachment):
-        self.attachments.append(attachment)
+        self.attachments[key] = attachment
+    
+    def set_injections(self, key: str, injections: List[PatcherInjection]):
+        self.injections[key] = injections
+
+    def set_additional_models(self, key: str, models: List['ModelPatcher']):
+        self.additional_models[key] = models
+
+    def use_ejected(self, skip_injection=False):
+        return AutoPatcherEjector(self, skip_until_exit=skip_injection)
+
+    def inject_model(self):
+        if self.is_injected or self.skip_injection:
+            return
+        for injections in self.injections.values():
+            for inj in injections:
+                inj.inject(self)
+                self.is_injected = True
+        if self.is_injected:
+            for callback in self.callbacks[CallbacksMP.ON_INJECT_MODEL]:
+                callback(self)
+
+    def eject_model(self):
+        if not self.is_injected:
+            return
+        for injections in self.injections.values():
+            for inj in injections:
+                inj.eject(self)
+        self.is_injected = False
+        for callback in self.callbacks[CallbacksMP.ON_EJECT_MODEL]:
+            callback(self)
 
     def pre_run(self):
         for callback in self.callbacks[CallbacksMP.ON_PRE_RUN]:
@@ -685,6 +773,7 @@ class ModelPatcher:
             callback(self, hooks_dict, target)
 
     def add_hook_patches(self, hook: comfy.hooks.WeightHook, patches, strength_patch=1.0, strength_model=1.0, is_diff=False):
+        with self.use_ejected():
             # NOTE: this mirrors behavior of add_patches func
             if is_diff:
                 comfy.model_management.unload_model_clones(self)
@@ -723,6 +812,7 @@ class ModelPatcher:
             return list(p)
 
     def get_weight_diffs(self, patches):
+        with self.use_ejected():
             comfy.model_management.unload_model_clones(self)
             weights: Dict[str, Tuple] = {}
             p = set()
@@ -765,6 +855,7 @@ class ModelPatcher:
             callback(self, hooks)
 
     def patch_hooks(self, hooks: comfy.hooks.HookGroup):
+        with self.use_ejected():
             self.unpatch_hooks()
             model_sd = self.model_state_dict()
             # if have cached weights for hooks, use it
@@ -825,6 +916,7 @@ class ModelPatcher:
             comfy.utils.set_attr_param(self.model, key, out_weight)
     
     def unpatch_hooks(self) -> None:
+        with self.use_ejected():
             if len(self.hook_backup) == 0:
                 self.current_hooks = None
                 return

comfy/sampler_helpers.py

@@ -2,6 +2,11 @@ import torch
 import comfy.model_management
 import comfy.conds
 import comfy.hooks
+from typing import TYPE_CHECKING, Dict, List
+if TYPE_CHECKING:
+    from comfy.model_patcher import ModelPatcher
+    from comfy.model_base import BaseModel
+    from comfy.controlnet import ControlBase
 
 def prepare_mask(noise_mask, shape, device):
     """ensures noise mask is of proper dimensions"""
@@ -15,9 +20,22 @@ def get_models_from_cond(cond, model_type):
     models = []
     for c in cond:
         if model_type in c:
+            if isinstance(c[model_type], list):
+                models += c[model_type]
+            else:
                 models += [c[model_type]]
     return models
 
+def get_hooks_from_cond(cond, filter_types: List[comfy.hooks.EnumHookType]=None):
+    hooks: Dict[comfy.hooks.Hook, None] = {}
+    for c in cond:
+        if 'hooks' in c:
+            for hook in c['hooks'].hooks:
+                hook: comfy.hooks.Hook
+                if not filter_types or hook.hook_type in filter_types:
+                    hooks[hook] = None
+    return hooks
+
 def convert_cond(cond):
     out = []
     for c in cond:
@@ -32,12 +50,16 @@ def convert_cond(cond):
 
 def get_additional_models(conds, dtype):
     """loads additional models in conditioning"""
-    cnets = []
+    cnets: List[ControlBase] = []
     gligen = []
+    add_models = []
+    hooks: Dict[comfy.hooks.AddModelHook, None] = {}
 
     for k in conds:
         cnets += get_models_from_cond(conds[k], "control")
         gligen += get_models_from_cond(conds[k], "gligen")
+        add_models += get_models_from_cond(conds[k], "additional_models")
+        hooks.update(get_hooks_from_cond(conds[k], [comfy.hooks.EnumHookType.AddModel]))
 
     control_nets = set(cnets)
 
@@ -48,7 +70,9 @@ def get_additional_models(conds, dtype):
         inference_memory += m.inference_memory_requirements(dtype)
 
     gligen = [x[1] for x in gligen]
-    models = control_models + gligen
+    hook_models = [x.model for x in hooks]
+    models = control_models + gligen + add_models + hook_models
+
     return models, inference_memory
 
 def cleanup_additional_models(models):
@@ -58,10 +82,11 @@ def cleanup_additional_models(models):
             m.cleanup()
 
 
-def prepare_sampling(model, noise_shape, conds):
+def prepare_sampling(model: 'ModelPatcher', noise_shape, conds):
     device = model.load_device
-    real_model = None
+    real_model: 'BaseModel' = None
     models, inference_memory = get_additional_models(conds, model.model_dtype())
+    models += model.get_all_additional_models()  # TODO: does this require inference_memory update?
     memory_required = model.memory_required([noise_shape[0] * 2] + list(noise_shape[1:])) + inference_memory
     minimum_memory_required = model.memory_required([noise_shape[0]] + list(noise_shape[1:])) + inference_memory
     comfy.model_management.load_models_gpu([model] + models, memory_required=memory_required, minimum_memory_required=minimum_memory_required)
@@ -79,12 +104,9 @@ def cleanup_models(conds, models):
 
     cleanup_additional_models(set(control_cleanup))
 
-def prepare_model_patcher(model, conds):
+def prepare_model_patcher(model: 'ModelPatcher', conds):
     # check for hooks in conds - if not registered, see if can be applied
     hooks = {}
     for k in conds:
-        for cond in conds[k]:
+        hooks.update(get_hooks_from_cond(conds[k]))
-            if 'hooks' in cond:
-                for hook in cond['hooks'].hooks:
-                    hooks[hook] = None
     model.register_all_hook_patches(hooks, comfy.hooks.EnumWeightTarget.Model)