diff --git a/comfy/model_management.py b/comfy/model_management.py index b15d08ba1..222005b6f 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -616,6 +616,8 @@ PIN_PRESSURE_HYSTERESIS = 256 * 1024 * 1024 #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. 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): module_mem = 0 @@ -642,6 +644,15 @@ def free_pins(size, evict_active=False): size -= freed 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): if args.high_ram: return True diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py index ad75a0e50..6bd99b68f 100644 --- a/comfy_execution/caching.py +++ b/comfy_execution/caching.py @@ -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 diff --git a/execution.py b/execution.py index c45317593..19b8cdd68 100644 --- a/execution.py +++ b/execution.py @@ -29,6 +29,7 @@ from comfy_execution.caching import ( HierarchicalCache, LRUCache, RAMPressureCache, + RAM_CACHE_LARGE_INTERMEDIATE, ) from comfy_execution.graph import ( DynamicPrompt, @@ -794,12 +795,16 @@ class PromptExecutor: if self.cache_type == CacheType.RAM_PRESSURE: ram_release_callback(ram_inactive_headroom) ram_shortfall = ram_headroom - psutil.virtual_memory().available - freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2)) - 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) + if ram_shortfall > 0: + freed = ram_release_callback(ram_headroom, free_active=True, min_entry_size=RAM_CACHE_LARGE_INTERMEDIATE) + ram_shortfall -= freed + if comfy.model_management.should_free_pins_for_ram_pressure(ram_shortfall): + freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2)) + 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: # Only execute when the while-loop ends without break # Send cached UI for intermediate output nodes that weren't executed