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

This commit is contained in:
Jedrzej Kosinski 2024-10-11 06:50:55 -05:00
parent 1e2777bab1
commit 1f8d9c040b
2 changed files with 4 additions and 1 deletions

View File

@ -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)

View File

@ -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