Added current_patcher property to BaseModel

comfy/model_base.py

@@ -93,6 +93,7 @@ class BaseModel(torch.nn.Module):
         self.model_config = model_config
         self.manual_cast_dtype = model_config.manual_cast_dtype
         self.device = device
+        self.current_patcher = None
 
         if not unet_config.get("disable_unet_model_creation", False):
             if model_config.custom_operations is None:

comfy/sampler_helpers.py

@@ -65,6 +65,7 @@ def prepare_sampling(model, noise_shape, conds):
     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)
     real_model = model.model
+    real_model.current_patcher = model
 
     return real_model, conds, models
 

comfy/samplers.py

@@ -181,7 +181,8 @@ def calc_cond_batch(model, conds, x_in, timestep, model_options):
                 if model.memory_required(input_shape) * 1.5 < free_memory:
                     to_batch = batch_amount
                     break
-            # TODO: add apply_hooks call here, once a ModelPatcher ref is added to BaseModel
+
+            model.current_patcher.apply_hooks(hooks=hooks)
 
             input_x = []
             mult = []