From c01175530ed36fcb5961c2f2f2598e19b73287b9 Mon Sep 17 00:00:00 2001 From: rattus <46076784+rattus128@users.noreply.github.com> Date: Wed, 29 Jul 2026 07:05:57 +1000 Subject: [PATCH] Load weights to process RAM with MRU policy using pinning infrastructure (#15027) --- comfy/model_management.py | 87 +++++++++++++++++++++++-------------- comfy/model_patcher.py | 89 ++++++++++++++++++++++++++++++++------ comfy/ops.py | 25 ++++++++--- comfy/pinned_memory.py | 70 +++++++++++++++++------------- comfy_execution/caching.py | 4 +- comfy_execution/graph.py | 18 +++++--- execution.py | 6 ++- 7 files changed, 207 insertions(+), 92 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 766e9ea89..eb768d783 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -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() diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index d70b42bf8..39246b95c 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -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 diff --git a/comfy/ops.py b/comfy/ops.py index 1f7cc9575..5e1cce333 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -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 diff --git a/comfy/pinned_memory.py b/comfy/pinned_memory.py index cb77c517a..d78ab3c76 100644 --- a/comfy/pinned_memory.py +++ b/comfy/pinned_memory.py @@ -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 diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py index 6bd99b68f..d60aa1e50 100644 --- a/comfy_execution/caching.py +++ b/comfy_execution/caching.py @@ -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) diff --git a/comfy_execution/graph.py b/comfy_execution/graph.py index 479ee8a53..64dec2045 100644 --- a/comfy_execution/graph.py +++ b/comfy_execution/graph.py @@ -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 diff --git a/execution.py b/execution.py index 387772629..b17ace65a 100644 --- a/execution.py +++ b/execution.py @@ -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)