mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-06-09 13:02:08 +08:00
Optimize memory unload strategy for more optimized performance.
This commit is contained in:
parent
6a32c06f06
commit
c6de09b02e
@ -349,20 +349,27 @@ def unload_model_clones(model, unload_weights_only=True, force_unload=True):
|
|||||||
return unload_weight
|
return unload_weight
|
||||||
|
|
||||||
def free_memory(memory_required, device, keep_loaded=[]):
|
def free_memory(memory_required, device, keep_loaded=[]):
|
||||||
unloaded_model = False
|
unloaded_model = []
|
||||||
|
can_unload = []
|
||||||
|
|
||||||
for i in range(len(current_loaded_models) -1, -1, -1):
|
for i in range(len(current_loaded_models) -1, -1, -1):
|
||||||
if not DISABLE_SMART_MEMORY:
|
|
||||||
if get_free_memory(device) > memory_required:
|
|
||||||
break
|
|
||||||
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:
|
||||||
m = current_loaded_models.pop(i)
|
can_unload.append((sys.getrefcount(shift_model.model), shift_model.model_memory(), i))
|
||||||
m.model_unload()
|
|
||||||
del m
|
|
||||||
unloaded_model = True
|
|
||||||
|
|
||||||
if unloaded_model:
|
for x in sorted(can_unload):
|
||||||
|
i = x[-1]
|
||||||
|
if not DISABLE_SMART_MEMORY:
|
||||||
|
if get_free_memory(device) > memory_required:
|
||||||
|
break
|
||||||
|
current_loaded_models[i].model_unload()
|
||||||
|
unloaded_model.append(i)
|
||||||
|
|
||||||
|
for i in sorted(unloaded_model, reverse=True):
|
||||||
|
current_loaded_models.pop(i)
|
||||||
|
|
||||||
|
if len(unloaded_model) > 0:
|
||||||
soft_empty_cache()
|
soft_empty_cache()
|
||||||
else:
|
else:
|
||||||
if vram_state != VRAMState.HIGH_VRAM:
|
if vram_state != VRAMState.HIGH_VRAM:
|
||||||
|
Loading…
x
Reference in New Issue
Block a user