From 1f8d9c040b8d5cd029d9db649ffd87e2299f9363 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Fri, 11 Oct 2024 06:50:55 -0500 Subject: [PATCH] Fixed models not being unloaded properly due to current_patcher reference; the current ComfyUI model cleanup code requires that nothing else has a reference to the ModelPatcher instances --- comfy/model_patcher.py | 4 ++++ comfy/sampler_helpers.py | 1 - 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 77f7d9bb7..fb0410d1f 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -793,6 +793,8 @@ class ModelPatcher: def cleanup(self): self.clean_hooks() + if hasattr(self.model, "current_patcher"): + self.model.current_patcher = None for callback in self.get_all_callbacks(CallbacksMP.ON_CLEANUP): callback(self) @@ -897,6 +899,8 @@ class ModelPatcher: callback(self) def pre_run(self): + if hasattr(self.model, "current_patcher"): + self.model.current_patcher = self for callback in self.get_all_callbacks(CallbacksMP.ON_PRE_RUN): callback(self) diff --git a/comfy/sampler_helpers.py b/comfy/sampler_helpers.py index a84778507..e1c0d9db9 100644 --- a/comfy/sampler_helpers.py +++ b/comfy/sampler_helpers.py @@ -93,7 +93,6 @@ def prepare_sampling(model: 'ModelPatcher', 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