298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-01-11 02:15:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix case where a memory leak could cause crash.

Browse Source 
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.
This commit is contained in:
comfyanonymous 2024-12-02 19:49:49 -05:00
parent 0ee322ec5f
commit 57e8bf6a9f
1 changed files with 6 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
comfy/model_management.py
View File

@ -369,6 +369,9 @@ class LoadedModel:
if self._patcher_finalizer is not None: if self._patcher_finalizer is not None:
self._patcher_finalizer.detach() 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): def use_more_memory(extra_memory, loaded_models, device):
for m in loaded_models: 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): for i in range(len(current_loaded_models) -1, -1, -1):
shift_model = current_loaded_models[i] shift_model = current_loaded_models[i]
if shift_model.device == device: 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)) can_unload.append((-shift_model.model_offloaded_memory(), sys.getrefcount(shift_model.model), shift_model.model_memory(), i))
shift_model.currently_used = False shift_model.currently_used = False
@ -531,7 +534,7 @@ def cleanup_models_gc():
do_gc = False do_gc = False
for i in range(len(current_loaded_models)): for i in range(len(current_loaded_models)):
cur = current_loaded_models[i] 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__)) 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 do_gc = True
break break
@ -542,7 +545,7 @@ def cleanup_models_gc():
for i in range(len(current_loaded_models)): for i in range(len(current_loaded_models)):
cur = current_loaded_models[i] 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__)) logging.warning("WARNING, memory leak with model {}. Please make sure it is not being referenced from somewhere.".format(cur.real_model().__class__.__name__))