298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-06-15 07:45:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

Optimize memory unload strategy for more optimized performance.

Browse Source
This commit is contained in:
comfyanonymous 2024-03-24 02:36:30 -04:00
parent 6a32c06f06
commit c6de09b02e
1 changed files with 16 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
comfy/model_management.py
View File

@ -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: