zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-08-26 02:42:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

Load weights to process RAM with MRU policy using pinning infrastructure (#15027)

Browse Source
This commit is contained in:
rattus
2026-07-29 07:05:57 +10:00
committed by GitHub
parent 3d41e3ea4e
commit c01175530e
7 changed files with 207 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
comfy_execution/caching.py
View File

@@ -5,7 +5,7 @@ import psutil
import time
import torch
from typing import Sequence, Mapping, Dict
from comfy.model_patcher import ModelPatcher
from comfy.model_patcher import is_model_patcher_output
from comfy_execution.graph import DynamicPrompt
from abc import ABC, abstractmethod
@@ -567,7 +567,7 @@ class RAMPressureCache(LRUCache):
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:
elif is_model_patcher_output(output) and self.used_generation[key] != self.generation:
#old ModelPatchers are the first to go
oom_ram_usage = 1e30
scan_list_for_ram_usage(cache_entry.outputs)

18
comfy_execution/graph.py
View File

@@ -195,9 +195,10 @@ class ExecutionList(TopologicalSort):
ExecutionList implements a topological dissolve of the graph. After a node is staged for execution,
it can still be returned to the graph after having further dependencies added.
"""
def __init__(self, dynprompt, output_cache):
def __init__(self, dynprompt, output_cache, output_link_callback=None):
super().__init__(dynprompt)
self.output_cache = output_cache
self.output_link_callback = output_link_callback
self.staged_node_id = None
self.execution_cache = {}
self.execution_cache_listeners = {}
@@ -205,13 +206,16 @@ class ExecutionList(TopologicalSort):
def is_cached(self, node_id):
return self.output_cache.get_local(node_id) is not None
def cache_link(self, from_node_id, to_node_id):
def cache_link(self, from_node_id, to_node_id, from_socket=None):
if to_node_id not in self.execution_cache:
self.execution_cache[to_node_id] = {}
self.execution_cache[to_node_id][from_node_id] = self.output_cache.get_local(from_node_id)
value = self.output_cache.get_local(from_node_id)
self.execution_cache[to_node_id][from_node_id] = value
if from_node_id not in self.execution_cache_listeners:
self.execution_cache_listeners[from_node_id] = set()
self.execution_cache_listeners[from_node_id].add(to_node_id)
self.execution_cache_listeners[from_node_id].add((to_node_id, from_socket))
if value is not None and from_socket is not None and self.output_link_callback is not None:
self.output_link_callback(value.outputs[from_socket])
def get_cache(self, from_node_id, to_node_id):
if to_node_id not in self.execution_cache:
@@ -225,13 +229,15 @@ class ExecutionList(TopologicalSort):
def cache_update(self, node_id, value):
if node_id in self.execution_cache_listeners:
for to_node_id in self.execution_cache_listeners[node_id]:
for to_node_id, from_socket in self.execution_cache_listeners[node_id]:
if to_node_id in self.execution_cache:
self.execution_cache[to_node_id][node_id] = value
if from_socket is not None and self.output_link_callback is not None:
self.output_link_callback(value.outputs[from_socket])
def add_strong_link(self, from_node_id, from_socket, to_node_id):
super().add_strong_link(from_node_id, from_socket, to_node_id)
self.cache_link(from_node_id, to_node_id)
self.cache_link(from_node_id, to_node_id, from_socket)
async def stage_node_execution(self):
assert self.staged_node_id is None