From 57e8bf6a9f9f550613d00ce36ac60841145a2b7b Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 2 Dec 2024 19:49:49 -0500 Subject: [PATCH] Fix case where a memory leak could cause crash. Now the only symptom of code messing up and keeping references to a model object when it should not will be endless prints in the log instead of the next workflow crashing ComfyUI. --- comfy/model_management.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 2a500956..a16cf380 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -369,6 +369,9 @@ class LoadedModel: if self._patcher_finalizer is not None: self._patcher_finalizer.detach() + def is_dead(self): + return self.real_model() is not None and self.model is None + def use_more_memory(extra_memory, loaded_models, device): for m in loaded_models: @@ -409,7 +412,7 @@ def free_memory(memory_required, device, keep_loaded=[]): for i in range(len(current_loaded_models) -1, -1, -1): shift_model = current_loaded_models[i] if shift_model.device == device: - if shift_model not in keep_loaded: + if shift_model not in keep_loaded and not shift_model.is_dead(): can_unload.append((-shift_model.model_offloaded_memory(), sys.getrefcount(shift_model.model), shift_model.model_memory(), i)) shift_model.currently_used = False @@ -531,7 +534,7 @@ def cleanup_models_gc(): do_gc = False for i in range(len(current_loaded_models)): cur = current_loaded_models[i] - if cur.real_model() is not None and cur.model is None: + if cur.is_dead(): logging.info("Potential memory leak detected with model {}, doing a full garbage collect, for maximum performance avoid circular references in the model code.".format(cur.real_model().__class__.__name__)) do_gc = True break @@ -542,7 +545,7 @@ def cleanup_models_gc(): for i in range(len(current_loaded_models)): cur = current_loaded_models[i] - if cur.real_model() is not None and cur.model is None: + if cur.is_dead(): logging.warning("WARNING, memory leak with model {}. Please make sure it is not being referenced from somewhere.".format(cur.real_model().__class__.__name__))