Try to fix the model reloading issue some people have.

Author: comfyanonymous

comfy/model_management.py

@@ -616,6 +616,8 @@ def get_torch_device_name(device):
 #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

comfy_execution/caching.py

@@ -503,6 +503,8 @@ async def ensure_subcache_for(self, node_id, children_ids):
 
 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 @@ def set_local(self, node_id, value):
         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,28 +556,36 @@ def ram_release(self, target, free_active=False):
             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:
                     if isinstance(output, (list, tuple)):
                         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

execution.py

@@ -29,6 +29,7 @@
     HierarchicalCache,
     LRUCache,
     RAMPressureCache,
+    RAM_CACHE_LARGE_INTERMEDIATE,
 )
 from comfy_execution.graph import (
     DynamicPrompt,
@@ -794,12 +795,16 @@ async def execute_async(self, prompt, prompt_id, extra_data={}, execute_outputs=
                     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