wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-06-23 16:29:25 +08:00
Code Issues Actions 16 Packages Projects Releases Wiki Activity

Merge 3e8ef6355b into 6978a466b8

Browse Source
This commit is contained in:
liminfei-amd 2026-06-22 11:58:50 +00:00 committed by GitHub
commit 06a008d0a6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 36 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

51
comfy/model_management.py
View File

@ -767,7 +767,9 @@ class LoadedModel:
self._patcher_finalizer.detach() self._patcher_finalizer.detach()
def is_dead(self): def is_dead(self):
return self.real_model() is not None and self.model is None # real_model is plain None until model_load(); guard so a pre-load entry can't raise.
rm = self.real_model
return rm is not None and rm() 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):
@ -810,28 +812,42 @@ def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins
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]
# Pin the patcher; LoadedModel.model is a weakref deref a mid-loop finalizer can None.
model = shift_model.model
if model is None:
continue
if device is None or shift_model.device == device: if device is None or shift_model.device == device:
if shift_model not in keep_loaded and not shift_model.is_dead(): 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)) # Carry the object so a reentrant cleanup_models pop can't stale the index; i is just a sort tiebreaker.
can_unload.append((-shift_model.model_offloaded_memory(), sys.getrefcount(model), shift_model.model_memory(), i, shift_model))
shift_model.currently_used = False shift_model.currently_used = False
can_unload_sorted = sorted(can_unload) can_unload_sorted = sorted(can_unload)
for x in can_unload_sorted: for x in can_unload_sorted:
i = x[-1] shift_model = x[-1]
# Pin: shift_model.model still re-derefs the weakref a reentrant cleanup_models can None.
model = shift_model.model
if model is None:
continue
memory_to_free = 1e32 memory_to_free = 1e32
if not DISABLE_SMART_MEMORY or device is None: if not DISABLE_SMART_MEMORY or device is None:
memory_to_free = 0 if device is None else memory_required - get_free_memory(device) memory_to_free = 0 if device is None else memory_required - get_free_memory(device)
if current_loaded_models[i].model.is_dynamic() and for_dynamic: if model.is_dynamic() and for_dynamic:
#don't actually unload dynamic models for the sake of other dynamic models #don't actually unload dynamic models for the sake of other dynamic models
#as that works on-demand. #as that works on-demand.
memory_required -= current_loaded_models[i].model.loaded_size() memory_required -= model.loaded_size()
memory_to_free = 0 memory_to_free = 0
if memory_to_free > 0 and current_loaded_models[i].model_unload(memory_to_free): if memory_to_free > 0 and shift_model.model_unload(memory_to_free):
logging.debug(f"Unloading {current_loaded_models[i].model.model.__class__.__name__}") logging.debug(f"Unloading {model.model.__class__.__name__}")
unloaded_model.append(i) unloaded_model.append(shift_model)
for i in sorted(unloaded_model, reverse=True): for shift_model in unloaded_model:
unloaded_models.append(current_loaded_models.pop(i)) unloaded_models.append(shift_model)
# Remove by identity (model is None post-unload, so __eq__ is unsafe); tolerate a reentrant pop.
for idx in range(len(current_loaded_models) - 1, -1, -1):
if current_loaded_models[idx] is shift_model:
current_loaded_models.pop(idx)
break
if not for_dynamic and pins_required > 0: if not for_dynamic and pins_required > 0:
ensure_pin_budget(pins_required) ensure_pin_budget(pins_required)
@ -990,12 +1006,17 @@ def archive_model_dtypes(model):
def cleanup_models(): def cleanup_models():
to_delete = [] to_delete = []
for i in range(len(current_loaded_models)): for i in range(len(current_loaded_models)):
if current_loaded_models[i].real_model() is None: # real_model can still be the plain None from __init__ here; guard the call.
to_delete = [i] + to_delete rm = current_loaded_models[i].real_model
if rm is not None and rm() is None:
to_delete.append(current_loaded_models[i])
for i in to_delete: for lm in to_delete:
x = current_loaded_models.pop(i) # Remove by identity, mirroring free_memory; tolerate a reentrant pop.
del x for idx in range(len(current_loaded_models) - 1, -1, -1):
if current_loaded_models[idx] is lm:
current_loaded_models.pop(idx)
break
def dtype_size(dtype): def dtype_size(dtype):
dtype_size = 4 dtype_size = 4