mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-03 13:19:23 +08:00
Merge 3711b4d01f into 694815f498
This commit is contained in:
commit
9f0fb0e111
@ -808,30 +808,33 @@ def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins
|
|||||||
can_unload = []
|
can_unload = []
|
||||||
unloaded_models = []
|
unloaded_models = []
|
||||||
|
|
||||||
for i in range(len(current_loaded_models) -1, -1, -1):
|
for shift_model in list(current_loaded_models):
|
||||||
shift_model = current_loaded_models[i]
|
|
||||||
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))
|
can_unload.append((-shift_model.model_offloaded_memory(), sys.getrefcount(shift_model.model), shift_model.model_memory(), shift_model))
|
||||||
shift_model.currently_used = False
|
shift_model.currently_used = False
|
||||||
|
|
||||||
can_unload_sorted = sorted(can_unload)
|
can_unload_sorted = sorted(can_unload, key=lambda a: a[:3])
|
||||||
for x in can_unload_sorted:
|
for x in can_unload_sorted:
|
||||||
i = x[-1]
|
shift_model = x[-1]
|
||||||
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 shift_model.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 -= shift_model.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 {shift_model.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)
|
||||||
|
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)
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user