zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-08-17 14:59:12 +08:00
Code Issues Packages Projects Releases Wiki Activity

Try to fix the model reloading issue some people have. (#14822)

Browse Source
This commit is contained in:
comfyanonymous
2026-07-09 16:39:01 -07:00
committed by GitHub
parent 73e84d5ec8
commit b7a648ca20
3 changed files with 39 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
comfy_execution/caching.py
View File

@@ -503,6 +503,8 @@ RAM_CACHE_DEFAULT_RAM_USAGE = 0.05
RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3
RAM_CACHE_LARGE_INTERMEDIATE = 512 * 1024 ** 2
def all_outputs_dynamic(outputs):
if outputs is None:
@@ -517,7 +519,6 @@ def all_outputs_dynamic(outputs):
return True
class RAMPressureCache(LRUCache):
def __init__(self, key_class, enable_providers=False):
@@ -539,9 +540,9 @@ class RAMPressureCache(LRUCache):
self.timestamps[self.cache_key_set.get_data_key(node_id)] = time.time()
super().set_local(node_id, value)
def ram_release(self, target, free_active=False):
def ram_release(self, target, free_active=False, min_entry_size=0):
if psutil.virtual_memory().available >= target:
return
return 0
clean_list = []
@@ -555,8 +556,9 @@ class RAMPressureCache(LRUCache):
oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE
oom_ram_usage = ram_usage
def scan_list_for_ram_usage(outputs):
nonlocal ram_usage
nonlocal ram_usage, oom_ram_usage
if outputs is None:
return
for output in outputs:
@@ -564,19 +566,26 @@ class RAMPressureCache(LRUCache):
scan_list_for_ram_usage(output)
elif isinstance(output, torch.Tensor) and output.device.type == 'cpu':
ram_usage += output.numel() * output.element_size()
oom_ram_usage += output.numel() * output.element_size()
elif isinstance(output, ModelPatcher) and self.used_generation[key] != self.generation:
#old ModelPatchers are the first to go
ram_usage = 1e30
oom_ram_usage = 1e30
scan_list_for_ram_usage(cache_entry.outputs)
oom_score *= ram_usage
if ram_usage < min_entry_size:
continue
oom_score *= oom_ram_usage
#In the case where we have no information on the node ram usage at all,
#break OOM score ties on the last touch timestamp (pure LRU)
bisect.insort(clean_list, (oom_score, self.timestamps[key], key))
bisect.insort(clean_list, (oom_score, self.timestamps[key], key, ram_usage))
freed = 0
while psutil.virtual_memory().available < target and clean_list:
_, _, key = clean_list.pop()
_, _, key, ram_usage = clean_list.pop()
del self.cache[key]
self.used_generation.pop(key, None)
self.timestamps.pop(key, None)
self.children.pop(key, None)
freed += ram_usage
return freed