RomanNum3ral
/
ComfyUI
mirror of https://github.com/Comfy-Org/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

This commit is contained in:
rattus 2026-07-29 07:05:57 +10:00 committed by GitHub
parent 3d41e3ea4e
commit c01175530e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 207 additions and 92 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

87
comfy/model_management.py
View File

@ -632,18 +632,50 @@ def mark_mmap_dirty(storage):
if mmap_refs is not None:
DIRTY_MMAPS.add(mmap_refs[0])
def free_pins(size, evict_active=False):
PIN_SUBSETS = [ "weights", "patches" ]
LOADED_PIN_SUBSETS = [ "weights-loaded", "patches-loaded" ]
def models_for_pin_eviction(active, current_prompt=None):
for loaded_model in current_loaded_models:
model = loaded_model.model
if model is None or not model.is_dynamic():
continue
pin_state = model.model.dynamic_pins[model.load_device]
if ((active is None or pin_state["active"] == active) and
(current_prompt is None or pin_state["current_prompt"] == current_prompt)):
yield model
def free_model_pins(size, subsets, current_prompt, active, registrations=False):
freed_total = 0
for loaded_model in reversed(current_loaded_models):
for model in models_for_pin_eviction(active, current_prompt=current_prompt):
if size <= 0:
return freed_total
model = loaded_model.model
if model is not None and model.is_dynamic() and (evict_active or not model.model.dynamic_pins[model.load_device]["active"]):
freed = model.partially_unload_ram(size)
freed_total += freed
size -= freed
if registrations:
freed = model.unregister_inactive_pins(size, subsets=subsets)
else:
freed = model.partially_unload_ram(size, subsets=subsets)
freed_total += freed
size -= freed
return freed_total
def pin_eviction_tiers(loaded, evict_active):
tiers = [
(PIN_SUBSETS, False, None),
(LOADED_PIN_SUBSETS, False, None),
(LOADED_PIN_SUBSETS, True, None),
]
if not loaded:
tiers.append((PIN_SUBSETS, True, False))
if evict_active:
tiers.append((PIN_SUBSETS, True, True))
return tiers
def free_pins(size, evict_active=False, loaded=False):
freed = 0
for subsets, current_prompt, active in pin_eviction_tiers(loaded, evict_active):
freed += free_model_pins(size - freed, subsets, current_prompt, active)
return freed
def should_free_pins_for_ram_pressure(shortfall):
if shortfall <= 0:
return False
@ -653,7 +685,7 @@ def should_free_pins_for_ram_pressure(shortfall):
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, loaded=False):
if args.high_ram:
return True
if args.fast_disk:
@ -664,32 +696,21 @@ def ensure_pin_budget(size, evict_active=False):
return True
to_free = shortfall + PIN_PRESSURE_HYSTERESIS
return free_pins(to_free, evict_active=evict_active) >= shortfall
return free_pins(to_free, evict_active=evict_active, loaded=loaded) >= shortfall
def free_registrations(shortfall, evict_active=True):
def free_registrations(shortfall, evict_active=True, loaded=False):
if MAX_PINNED_MEMORY <= 0:
return False
if shortfall <= 0:
return True
shortfall += REGISTERABLE_PIN_HYSTERESIS
for loaded_model in reversed(current_loaded_models):
model = loaded_model.model
if model is not None and model.is_dynamic() and not model.model.dynamic_pins[model.load_device]["active"]:
shortfall -= model.unregister_inactive_pins(shortfall)
if shortfall <= 0:
return True
if evict_active:
for loaded_model in current_loaded_models:
model = loaded_model.model
if model is not None and model.is_dynamic() and model.model.dynamic_pins[model.load_device]["active"]:
shortfall -= model.unregister_inactive_pins(shortfall)
if shortfall <= 0:
return True
for subsets, current_prompt, active in pin_eviction_tiers(loaded, evict_active):
shortfall -= free_model_pins(shortfall, subsets, current_prompt, active, registrations=True)
return shortfall <= REGISTERABLE_PIN_HYSTERESIS
def ensure_pin_registerable(size, evict_active=True):
return free_registrations(TOTAL_PINNED_MEMORY + size - MAX_PINNED_MEMORY, evict_active=evict_active)
def ensure_pin_registerable(size, evict_active=True, loaded=False):
return free_registrations(TOTAL_PINNED_MEMORY + size - MAX_PINNED_MEMORY, evict_active=evict_active, loaded=loaded)
class LoadedModel:
def __init__(self, model: ModelPatcher):
@ -1379,15 +1400,17 @@ def reset_cast_buffers():
pin_state = model.model.dynamic_pins[model.load_device]
if pin_state["active"]:
*_, buckets = pin_state["weights"]
for size, bucket in list(buckets.items()):
bucket[:] = [ entry for entry in bucket if entry[-1] is not None ]
if not bucket:
del buckets[size]
for subset in ("weights", "weights-loaded"):
*_, buckets = pin_state[subset]
for size, bucket in list(buckets.items()):
bucket[:] = [ entry for entry in bucket if entry[-1] is not None ]
if not bucket:
del buckets[size]
pin_state["active"] = False
model.partially_unload_ram(1e30, subsets=[ "patches" ])
model.model.dynamic_pins[model.load_device]["patches"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, pinned_hostbuf_size(model.model_size())), [], [-1], [0], [0], {})
model.partially_unload_ram(1e30, subsets=[ "patches", "patches-loaded" ])
for subset in ("patches", "patches-loaded"):
pin_state[subset] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, pinned_hostbuf_size(model.model_size())), [], [-1], [0], [0], {})
STREAM_CAST_BUFFERS.clear()
STREAM_AIMDO_CAST_BUFFERS.clear()

89
comfy/model_patcher.py
View File

@ -42,6 +42,52 @@ from comfy.patcher_extension import CallbacksMP, PatcherInjection, WrappersMP
import comfy_aimdo.model_vbar
def is_model_patcher_output(output):
return isinstance(output, ModelPatcher) or isinstance(getattr(output, "patcher", None), ModelPatcher)
class PromptModelTracker:
def __init__(self):
self.models = {}
def start(self):
self.end()
def add(self, outputs):
if isinstance(outputs, collections.abc.Mapping):
outputs = outputs.values()
elif not isinstance(outputs, (list, tuple)):
outputs = (outputs,)
for output in outputs:
if isinstance(output, (collections.abc.Mapping, list, tuple)):
self.add(output)
continue
models = []
if isinstance(output, ModelPatcher):
models.append(output)
models.extend(output.model_patches_models())
models.extend(output.get_nested_additional_models())
else:
patcher = getattr(output, "patcher", None)
if isinstance(patcher, ModelPatcher):
models.append(patcher)
get_models = getattr(output, "get_models", None)
if callable(get_models):
models.extend(get_models())
for model in models:
if not isinstance(model, ModelPatcher) or not model.is_dynamic():
continue
key = (id(model.model), model.load_device)
self.models[key] = model
model.set_in_use_by_current_prompt(True)
def end(self):
for model in self.models.values():
model.set_in_use_by_current_prompt(False)
self.models.clear()
def set_model_options_patch_replace(model_options, patch, name, block_name, number, transformer_index=None):
to = model_options["transformer_options"].copy()
@ -1724,14 +1770,20 @@ class ModelPatcherDynamic(ModelPatcher):
self.model.dynamic_pins[device] = {
"weights": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}),
"patches": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}),
"weights-loaded": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}),
"patches-loaded": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}),
"hostbufs_initialized": False,
"failed": False,
"active": False,
"current_prompt": False,
}
def is_dynamic(self):
return True
def set_in_use_by_current_prompt(self, in_use):
self.model.dynamic_pins[self.load_device]["current_prompt"] = in_use
def _vbar_get(self, create=False):
if self.load_device == torch.device("cpu"):
return None
@ -1802,6 +1854,8 @@ class ModelPatcherDynamic(ModelPatcher):
hostbuf_size = comfy.model_management.pinned_hostbuf_size(self.model_size())
pin_state["weights"] = (comfy_aimdo.host_buffer.HostBuffer(0, 64 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {})
pin_state["patches"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {})
pin_state["weights-loaded"] = (comfy_aimdo.host_buffer.HostBuffer(0, 64 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {})
pin_state["patches-loaded"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {})
pin_state["hostbufs_initialized"] = True
pin_state["failed"] = False
pin_state["active"] = True
@ -1943,12 +1997,14 @@ class ModelPatcherDynamic(ModelPatcher):
return freed
def loaded_ram_size(self):
return (self.model.dynamic_pins[self.load_device]["weights"][0].size)
pin_state = self.model.dynamic_pins[self.load_device]
return pin_state["weights"][0].size + pin_state["weights-loaded"][0].size
def pinned_memory_size(self):
return (self.model.dynamic_pins[self.load_device]["weights"][3][0])
pin_state = self.model.dynamic_pins[self.load_device]
return pin_state["weights"][3][0] + pin_state["weights-loaded"][3][0]
def unregister_inactive_pins(self, ram_to_unload, subsets=[ "weights", "patches" ]):
def unregister_inactive_pins(self, ram_to_unload, subsets=[ "weights-loaded", "patches-loaded", "weights", "patches" ]):
freed = 0
pin_state = self.model.dynamic_pins[self.load_device]
for subset in subsets:
@ -1956,15 +2012,17 @@ class ModelPatcherDynamic(ModelPatcher):
split = stack_split[0]
while split >= 0:
module, offset = stack[split]
module_pin = module._pins[subset]
split -= 1
stack_split[0] = split
if not module._pin_registered:
if not module_pin["registered"]:
continue
size = module._pin.numel() * module._pin.element_size()
if torch.cuda.cudart().cudaHostUnregister(module._pin.data_ptr()) != 0:
pin = module_pin["pin"]
size = pin.numel() * pin.element_size()
if torch.cuda.cudart().cudaHostUnregister(pin.data_ptr()) != 0:
comfy.model_management.discard_cuda_async_error()
continue
module._pin_registered = False
module_pin["registered"] = False
comfy.model_management.TOTAL_PINNED_MEMORY = max(0, comfy.model_management.TOTAL_PINNED_MEMORY - size)
pinned_size[0] = max(0, pinned_size[0] - size)
freed += size
@ -1973,20 +2031,23 @@ class ModelPatcherDynamic(ModelPatcher):
return freed
return freed
def partially_unload_ram(self, ram_to_unload, subsets=[ "weights", "patches" ]):
def partially_unload_ram(self, ram_to_unload, subsets=[ "weights-loaded", "patches-loaded", "weights", "patches" ]):
freed = 0
pin_state = self.model.dynamic_pins[self.load_device]
for subset in subsets:
hostbuf, stack, stack_split, pinned_size, *_ = pin_state[subset]
while len(stack) > 0:
module, offset = stack.pop()
size = module._pin.numel() * module._pin.element_size()
module._pin_balancer_entry[-1] = None
del module._pin_balancer_entry
del module._pin
hostbuf.truncate(offset, do_unregister=module._pin_registered)
module_pin = module._pins[subset]
pin = module_pin["pin"]
size = pin.numel() * pin.element_size()
module_pin["balancer_entry"][-1] = None
del module_pin["balancer_entry"]
del module_pin["pin"]
registered = module_pin["registered"]
hostbuf.truncate(offset, do_unregister=registered)
stack_split[0] = min(stack_split[0], len(stack) - 1)
if module._pin_registered:
if registered:
comfy.model_management.TOTAL_PINNED_MEMORY = max(0, comfy.model_management.TOTAL_PINNED_MEMORY - size)
pinned_size[0] = max(0, pinned_size[0] - size)
freed += size

25
comfy/ops.py
View File

@ -144,8 +144,13 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
needs_cast = False
xfer_source = [ s.weight, s.bias ]
pin = comfy.pinned_memory.get_pin(s)
subset = "weights"
pin = comfy.pinned_memory.get_pin(s, subset=subset)
if pin is None and not args.fast_disk:
loaded_pin = comfy.pinned_memory.get_pin(s, subset="weights-loaded")
if loaded_pin is not None or signature is not None:
subset = "weights-loaded"
pin = loaded_pin
if pin is not None:
xfer_source = [ pin ]
@ -182,12 +187,12 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
if pin is not None:
cast_maybe_lowvram_patch([pin], dest, offload_stream)
return
if signature is None or args.high_ram:
if signature is None or not args.fast_disk or args.high_ram:
comfy.pinned_memory.pin_memory(m, subset=subset, size=size)
pin = comfy.pinned_memory.get_pin(m, subset=subset)
cast_maybe_lowvram_patch(source, pin, offload_stream, xfer_dest2=dest)
handle_pin(s, pin, xfer_source, xfer_dest, size=dest_size)
handle_pin(s, pin, xfer_source, xfer_dest, subset=subset, size=dest_size)
for param_key in ("weight", "bias"):
lowvram_source = getattr(s, param_key + "_lowvram_function", None)
@ -197,8 +202,16 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
lowvram_dest = get_cast_buffer(lowvram_size)
lowvram_source.prepare(lowvram_dest, None, copy=False, commit=True)
pin = comfy.pinned_memory.get_pin(lowvram_source, subset="patches")
handle_pin(lowvram_source, pin, lowvram_source, lowvram_dest, subset="patches", size=lowvram_size)
subset = "patches"
pin = comfy.pinned_memory.get_pin(lowvram_source, subset=subset)
if pin is None:
loaded_pin = comfy.pinned_memory.get_pin(lowvram_source, subset="patches-loaded")
if loaded_pin is not None:
subset = "patches-loaded"
pin = loaded_pin
elif signature is not None and not args.fast_disk:
subset = "patches-loaded"
handle_pin(lowvram_source, pin, lowvram_source, lowvram_dest, subset=subset, size=lowvram_size)
prefetch["xfer_dest"] = xfer_dest

70
comfy/pinned_memory.py
View File

@ -9,14 +9,14 @@ import torch
from comfy.cli_args import args
def _add_to_bucket(module, buckets, size, priority):
def _add_to_bucket(module, module_pin, buckets, size, priority):
bucket = buckets.setdefault(size, [])
entry = [-priority, 0, module]
entry[1] = id(entry)
bisect.insort(bucket, entry)
module._pin_balancer_entry = entry
module_pin["balancer_entry"] = entry
def _steal_pin(module, stack, buckets, size, priority):
def _steal_pin(module, stack, buckets, size, priority, subset):
bucket = buckets.get(size)
if bucket is None:
return False
@ -31,34 +31,39 @@ def _steal_pin(module, stack, buckets, size, priority):
return False
*_, victim = bucket.pop()
module._pin = victim._pin
module._pin_registered = victim._pin_registered
module._pin_stack_index = victim._pin_stack_index
stack[module._pin_stack_index] = (module, stack[module._pin_stack_index][1])
module_pin = module._pins[subset]
victim_pin = victim._pins[subset]
module_pin["pin"] = victim_pin["pin"]
module_pin["registered"] = victim_pin["registered"]
module_pin["stack_index"] = victim_pin["stack_index"]
stack_index = module_pin["stack_index"]
stack[stack_index] = (module, stack[stack_index][1])
victim._pin_registered = False
del victim._pin
del victim._pin_stack_index
del victim._pin_balancer_entry
victim_pin["registered"] = False
del victim_pin["pin"]
del victim_pin["stack_index"]
del victim_pin["balancer_entry"]
_add_to_bucket(module, buckets, size, priority)
_add_to_bucket(module, module_pin, buckets, size, priority)
return True
def get_pin(module, subset="weights"):
pin = getattr(module, "_pin", None)
if pin is None or module._pin_registered or args.disable_pinned_memory:
pins = module.__dict__.get("_pins")
module_pin = None if pins is None else pins.get(subset)
pin = None if module_pin is None else module_pin.get("pin")
if pin is None or module_pin["registered"] or args.disable_pinned_memory:
return pin
_, _, stack_split, pinned_size, *_ = module._pin_state[subset]
size = pin.nbytes
comfy.model_management.ensure_pin_registerable(size)
comfy.model_management.ensure_pin_registerable(size, loaded=subset.endswith("-loaded"))
if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:
comfy.model_management.discard_cuda_async_error()
return pin
module._pin_registered = True
stack_split[0] = max(stack_split[0], module._pin_stack_index)
module_pin["registered"] = True
stack_split[0] = max(stack_split[0], module_pin["stack_index"])
comfy.model_management.TOTAL_PINNED_MEMORY += size
pinned_size[0] += size
return pin
@ -72,23 +77,26 @@ def pin_memory(module, subset="weights", size=None):
if pin is not None:
return
pins = module.__dict__.setdefault("_pins", {})
module_pin = pins.setdefault(subset, {})
hostbuf, stack, stack_split, pinned_size, counter, buckets = pin_state[subset]
if size is None:
size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ])
offset = hostbuf.size
registerable_size = size
priority = getattr(module, "_pin_balancer_priority", None)
loaded = subset.endswith("-loaded")
priority = module_pin.get("balancer_priority")
if priority is None:
priority = comfy.utils.bit_reverse_range(counter[0], 16)
counter[0] += 1
module._pin_balancer_priority = priority
module_pin["balancer_priority"] = priority
comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM)
if (not comfy.model_management.ensure_pin_budget(size) or
not comfy.model_management.ensure_pin_registerable(registerable_size)):
return _steal_pin(module, stack, buckets, size, priority)
if (not comfy.model_management.ensure_pin_budget(size, loaded=loaded) or
not comfy.model_management.ensure_pin_registerable(registerable_size, loaded=loaded)):
return _steal_pin(module, stack, buckets, size, priority, subset)
offset = hostbuf.size
extended = False
try:
hostbuf.extend(size=size, register=False)
@ -97,23 +105,23 @@ def pin_memory(module, subset="weights", size=None):
pin.untyped_storage()._comfy_hostbuf = hostbuf
if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:
comfy.model_management.discard_cuda_async_error()
comfy.model_management.free_registrations(size)
comfy.model_management.free_registrations(size, loaded=loaded)
if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:
comfy.model_management.discard_cuda_async_error()
del pin
hostbuf.truncate(offset, do_unregister=False)
return _steal_pin(module, stack, buckets, size, priority)
return _steal_pin(module, stack, buckets, size, priority, subset)
except RuntimeError:
if extended:
hostbuf.truncate(offset, do_unregister=False)
return _steal_pin(module, stack, buckets, size, priority)
return _steal_pin(module, stack, buckets, size, priority, subset)
module._pin = pin
module_pin["pin"] = pin
stack.append((module, offset))
module._pin_registered = True
module._pin_stack_index = len(stack) - 1
stack_split[0] = max(stack_split[0], module._pin_stack_index)
module_pin["registered"] = True
module_pin["stack_index"] = len(stack) - 1
stack_split[0] = max(stack_split[0], module_pin["stack_index"])
comfy.model_management.TOTAL_PINNED_MEMORY += size
pinned_size[0] += size
_add_to_bucket(module, buckets, size, priority)
_add_to_bucket(module, module_pin, buckets, size, priority)
return True

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

6
execution.py
View File

@ -16,6 +16,7 @@ import torch
from comfy.cli_args import args
import comfy.memory_management
import comfy.model_management
import comfy.model_patcher
import comfy.model_prefetch
import comfy_aimdo.model_vbar
@ -664,6 +665,7 @@ class PromptExecutor:
self.cache_args = cache_args
self.cache_type = cache_type
self.server = server
self.prompt_model_tracker = comfy.model_patcher.PromptModelTracker()
self.reset()
def reset(self):
@ -728,6 +730,7 @@ class PromptExecutor:
set_preview_method(extra_data.get("preview_method"))
nodes.interrupt_processing(False)
self.prompt_model_tracker.start()
if "client_id" in extra_data:
self.server.client_id = extra_data["client_id"]
@ -770,7 +773,7 @@ class PromptExecutor:
pending_async_nodes = {} # TODO - Unify this with pending_subgraph_results
ui_node_outputs = {}
executed = set()
execution_list = ExecutionList(dynamic_prompt, self.caches.outputs)
execution_list = ExecutionList(dynamic_prompt, self.caches.outputs, self.prompt_model_tracker.add)
current_outputs = self.caches.outputs.all_node_ids()
for node_id in list(execute_outputs):
execution_list.add_node(node_id)
@ -833,6 +836,7 @@ class PromptExecutor:
comfy.model_management.unload_all_models()
finally:
comfy.memory_management.set_ram_cache_release_state(None, 0)
self.prompt_model_tracker.end()
self._notify_prompt_lifecycle("end", prompt_id)