Try to fix the model reloading issue some people have.

This commit is contained in:
comfyanonymous 2026-07-07 22:46:23 -04:00
parent 6880614319
commit 1ed48eba5a
3 changed files with 39 additions and 14 deletions

View File

@ -616,6 +616,8 @@ PIN_PRESSURE_HYSTERESIS = 256 * 1024 * 1024
#Freeing registerables on pressure does imply a GPU sync, so go big on #Freeing registerables on pressure does imply a GPU sync, so go big on
#the hysteresis so each expensive sync gives us back a good chunk. #the hysteresis so each expensive sync gives us back a good chunk.
REGISTERABLE_PIN_HYSTERESIS = 2048 * 1024 * 1024 REGISTERABLE_PIN_HYSTERESIS = 2048 * 1024 * 1024
WINDOWS_PIN_EVICTION_SWAP_PERCENT = 5.0
WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE = 512 * 1024 ** 2
def module_size(module): def module_size(module):
module_mem = 0 module_mem = 0
@ -642,6 +644,15 @@ def free_pins(size, evict_active=False):
size -= freed size -= freed
return freed_total return freed_total
def should_free_pins_for_ram_pressure(shortfall):
if shortfall <= 0:
return False
if not WINDOWS:
return True
if psutil.virtual_memory().available < WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE:
return True
return psutil.swap_memory().percent >= WINDOWS_PIN_EVICTION_SWAP_PERCENT
def ensure_pin_budget(size, evict_active=False): def ensure_pin_budget(size, evict_active=False):
if args.high_ram: if args.high_ram:
return True return True

View File

@ -503,6 +503,8 @@ RAM_CACHE_DEFAULT_RAM_USAGE = 0.05
RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3 RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3
RAM_CACHE_LARGE_INTERMEDIATE = 512 * 1024 ** 2
def all_outputs_dynamic(outputs): def all_outputs_dynamic(outputs):
if outputs is None: if outputs is None:
@ -517,7 +519,6 @@ def all_outputs_dynamic(outputs):
return True return True
class RAMPressureCache(LRUCache): class RAMPressureCache(LRUCache):
def __init__(self, key_class, enable_providers=False): 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() self.timestamps[self.cache_key_set.get_data_key(node_id)] = time.time()
super().set_local(node_id, value) 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: if psutil.virtual_memory().available >= target:
return return 0
clean_list = [] clean_list = []
@ -555,8 +556,9 @@ class RAMPressureCache(LRUCache):
oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key]) oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE
oom_ram_usage = ram_usage
def scan_list_for_ram_usage(outputs): def scan_list_for_ram_usage(outputs):
nonlocal ram_usage nonlocal ram_usage, oom_ram_usage
if outputs is None: if outputs is None:
return return
for output in outputs: for output in outputs:
@ -564,19 +566,26 @@ class RAMPressureCache(LRUCache):
scan_list_for_ram_usage(output) scan_list_for_ram_usage(output)
elif isinstance(output, torch.Tensor) and output.device.type == 'cpu': elif isinstance(output, torch.Tensor) and output.device.type == 'cpu':
ram_usage += output.numel() * output.element_size() 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: elif isinstance(output, ModelPatcher) and self.used_generation[key] != self.generation:
#old ModelPatchers are the first to go #old ModelPatchers are the first to go
ram_usage = 1e30 oom_ram_usage = 1e30
scan_list_for_ram_usage(cache_entry.outputs) 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, #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) #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: while psutil.virtual_memory().available < target and clean_list:
_, _, key = clean_list.pop() _, _, key, ram_usage = clean_list.pop()
del self.cache[key] del self.cache[key]
self.used_generation.pop(key, None) self.used_generation.pop(key, None)
self.timestamps.pop(key, None) self.timestamps.pop(key, None)
self.children.pop(key, None) self.children.pop(key, None)
freed += ram_usage
return freed

View File

@ -29,6 +29,7 @@ from comfy_execution.caching import (
HierarchicalCache, HierarchicalCache,
LRUCache, LRUCache,
RAMPressureCache, RAMPressureCache,
RAM_CACHE_LARGE_INTERMEDIATE,
) )
from comfy_execution.graph import ( from comfy_execution.graph import (
DynamicPrompt, DynamicPrompt,
@ -794,12 +795,16 @@ class PromptExecutor:
if self.cache_type == CacheType.RAM_PRESSURE: if self.cache_type == CacheType.RAM_PRESSURE:
ram_release_callback(ram_inactive_headroom) ram_release_callback(ram_inactive_headroom)
ram_shortfall = ram_headroom - psutil.virtual_memory().available ram_shortfall = ram_headroom - psutil.virtual_memory().available
freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2)) if ram_shortfall > 0:
if freed < ram_shortfall: freed = ram_release_callback(ram_headroom, free_active=True, min_entry_size=RAM_CACHE_LARGE_INTERMEDIATE)
if freed > 64 * (1024 ** 2): ram_shortfall -= freed
# AIMDO MEM_DECOMMIT can outrun psutil.available catching up. if comfy.model_management.should_free_pins_for_ram_pressure(ram_shortfall):
time.sleep(0.05) freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2))
ram_release_callback(ram_headroom, free_active=True) if freed < ram_shortfall:
if freed > 64 * (1024 ** 2):
# AIMDO MEM_DECOMMIT can outrun psutil.available catching up.
time.sleep(0.05)
ram_release_callback(ram_headroom, free_active=True)
else: else:
# Only execute when the while-loop ends without break # Only execute when the while-loop ends without break
# Send cached UI for intermediate output nodes that weren't executed # Send cached UI for intermediate output nodes that weren't executed