diff --git a/comfy/model_base.py b/comfy/model_base.py index 9bfdb3b3e..fbd11b489 100644 --- a/comfy/model_base.py +++ b/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: diff --git a/comfy/sampler_helpers.py b/comfy/sampler_helpers.py index 4a2ec123b..25a0f5ccb 100644 --- a/comfy/sampler_helpers.py +++ b/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 diff --git a/comfy/samplers.py b/comfy/samplers.py index 578e29dc8..cbf01bff9 100644 --- a/comfy/samplers.py +++ b/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 = []