mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-10-09 11:02:12 -03:00
Will Miao
406096c0e1
trace_model_path dead-ended on KJNodes GetNode virtual links and fell back to the first entry in MODELS, which could be a stale checkpoint name resurrected from the process-lifetime node_cache — images ended up with a wrong or missing Hashes.model. - SetNodeExtractor now records MODEL passthrough as a model_variable entry (variable name -> source node id), cached per node like other metadata - trace_model_path resolves GetNode variable references through those entries instead of giving up (last SetNode wins, per KJNodes semantics) - _fill_missing_metadata validates cached checkpoint names against the node's current widget values and rebuilds stale entries from the current inputs instead of resurrecting the previous model
408 lines
17 KiB
Python
408 lines
17 KiB
Python
import os
|
|
import re
|
|
import time
|
|
from typing import Any
|
|
from nodes import NODE_CLASS_MAPPINGS # pyright: ignore[reportMissingImports, reportAttributeAccessIssue]
|
|
from .node_extractors import NODE_EXTRACTORS, GenericNodeExtractor
|
|
from .constants import METADATA_CATEGORIES, IMAGES, MODELS, OVERWRITE
|
|
|
|
|
|
# Input fields that carry a model filename in loader-style nodes. Mirrors
|
|
# GenericNodeExtractor._MODEL_NAME_FIELDS, plus the TensorRT engine extension.
|
|
_MODEL_NAME_FIELDS = ("ckpt_name", "unet_name", "model_path", "model_name", "gguf_name")
|
|
_MODEL_FILE_EXTENSIONS = (
|
|
".ckpt", ".pt", ".pt2", ".bin", ".pth",
|
|
".safetensors", ".pkl", ".sft", ".gguf", ".engine",
|
|
)
|
|
|
|
|
|
def _checkpoint_name_candidates(inputs):
|
|
"""Derive the checkpoint names a loader node's current inputs could record.
|
|
|
|
Used to keep stale ``node_cache`` entries (from before the user switched
|
|
models) out of prompts they no longer belong to. Returns an empty set when
|
|
the inputs carry no recognizable model field.
|
|
"""
|
|
candidates = set()
|
|
for field in _MODEL_NAME_FIELDS:
|
|
value = inputs.get(field)
|
|
if not isinstance(value, str) or not value.strip():
|
|
continue
|
|
name = value.strip()
|
|
if not name.lower().endswith(_MODEL_FILE_EXTENSIONS):
|
|
continue
|
|
candidates.add(name)
|
|
base = os.path.splitext(os.path.basename(name))[0]
|
|
candidates.add(base)
|
|
# TensorRTLoaderExtractor derivation: drop the "_$profile" part and
|
|
# any trailing save counter (e.g. "_00001_").
|
|
derived = base
|
|
if "_$" in derived:
|
|
derived = derived[: derived.index("_$")]
|
|
derived = re.sub(r"_\d+_?$", "", derived)
|
|
candidates.add(derived)
|
|
return candidates
|
|
|
|
|
|
class MetadataRegistry:
|
|
"""A singleton registry to store and retrieve workflow metadata"""
|
|
|
|
_instance = None
|
|
|
|
current_prompt_id: Any = None
|
|
current_prompt: Any = None
|
|
metadata: dict[str, Any] = {}
|
|
prompt_metadata: dict[str, Any] = {}
|
|
executed_nodes: set[str] = set()
|
|
node_cache: dict[str, Any] = {}
|
|
max_prompt_history: int = 3
|
|
metadata_categories: list[str] = METADATA_CATEGORIES
|
|
|
|
def __new__(cls):
|
|
if cls._instance is None:
|
|
cls._instance = super().__new__(cls)
|
|
cls._instance._reset()
|
|
return cls._instance
|
|
|
|
def _reset(self):
|
|
self.current_prompt_id = None
|
|
self.current_prompt = None
|
|
self.metadata = {}
|
|
self.prompt_metadata = {}
|
|
self.executed_nodes = set()
|
|
|
|
# Node-level cache for metadata
|
|
self.node_cache = {}
|
|
|
|
# Limit the number of stored prompts
|
|
self.max_prompt_history = 3
|
|
|
|
# Categories we want to track and retrieve from cache
|
|
self.metadata_categories = METADATA_CATEGORIES
|
|
|
|
def _clean_old_prompts(self):
|
|
"""Clean up old prompt metadata, keeping only recent ones"""
|
|
if len(self.prompt_metadata) <= self.max_prompt_history:
|
|
return
|
|
|
|
# Sort all prompt_ids by timestamp
|
|
sorted_prompts = sorted(
|
|
self.prompt_metadata.keys(),
|
|
key=lambda pid: self.prompt_metadata[pid].get("timestamp", 0),
|
|
)
|
|
|
|
# Remove oldest records
|
|
prompts_to_remove = sorted_prompts[
|
|
: len(sorted_prompts) - self.max_prompt_history
|
|
]
|
|
for pid in prompts_to_remove:
|
|
del self.prompt_metadata[pid]
|
|
|
|
def start_collection(self, prompt_id):
|
|
"""Begin metadata collection for a new prompt"""
|
|
self.current_prompt_id = prompt_id
|
|
self.executed_nodes = set()
|
|
self.prompt_metadata[prompt_id] = {
|
|
category: {} for category in METADATA_CATEGORIES
|
|
}
|
|
# Add additional metadata fields
|
|
self.prompt_metadata[prompt_id].update(
|
|
{
|
|
"execution_order": [],
|
|
"current_prompt": None, # Will store the prompt object
|
|
"extra_data": None, # Will store the API extra_data for workflow metadata
|
|
"timestamp": time.time(),
|
|
}
|
|
)
|
|
|
|
# Clean up old prompt data
|
|
self._clean_old_prompts()
|
|
|
|
def set_current_prompt(self, prompt):
|
|
"""Set the current prompt object reference"""
|
|
self.current_prompt = prompt
|
|
if self.current_prompt_id and self.current_prompt_id in self.prompt_metadata:
|
|
# Store the prompt in the metadata for later relationship tracing
|
|
self.prompt_metadata[self.current_prompt_id]["current_prompt"] = prompt
|
|
|
|
def set_extra_data(self, extra_data):
|
|
"""Store the API extra_data (contains extra_pnginfo.workflow with node properties)"""
|
|
if self.current_prompt_id and self.current_prompt_id in self.prompt_metadata:
|
|
self.prompt_metadata[self.current_prompt_id]["extra_data"] = extra_data
|
|
|
|
def get_metadata(self, prompt_id=None):
|
|
"""Get collected metadata for a prompt"""
|
|
key = prompt_id if prompt_id is not None else self.current_prompt_id
|
|
if key not in self.prompt_metadata:
|
|
return {}
|
|
|
|
metadata = self.prompt_metadata[key]
|
|
|
|
# If we have a current prompt object, check for non-executed nodes
|
|
prompt_obj = metadata.get("current_prompt")
|
|
if prompt_obj and hasattr(prompt_obj, "original_prompt"):
|
|
original_prompt = prompt_obj.original_prompt
|
|
|
|
# Fill in missing metadata from cache for nodes that weren't executed
|
|
self._fill_missing_metadata(key, original_prompt)
|
|
|
|
return self.prompt_metadata.get(key, {})
|
|
|
|
def _fill_missing_metadata(self, prompt_id, original_prompt):
|
|
"""Fill missing metadata from cache for non-executed nodes"""
|
|
if not original_prompt:
|
|
return
|
|
|
|
executed_nodes = self.executed_nodes
|
|
metadata = self.prompt_metadata[prompt_id]
|
|
|
|
# Iterate through nodes in the original prompt
|
|
for node_id, node_data in original_prompt.items():
|
|
# Skip if already executed in this run
|
|
if node_id in executed_nodes:
|
|
continue
|
|
|
|
# Get the node type from the prompt (this is the key in NODE_CLASS_MAPPINGS)
|
|
prompt_class_type = node_data.get("class_type")
|
|
if not prompt_class_type:
|
|
continue
|
|
|
|
# Convert to actual class name (which is what we use in our cache)
|
|
class_type = prompt_class_type
|
|
if prompt_class_type in NODE_CLASS_MAPPINGS:
|
|
class_obj = NODE_CLASS_MAPPINGS[prompt_class_type]
|
|
class_type = class_obj.__name__
|
|
|
|
# Create cache key using the actual class name
|
|
cache_key = f"{node_id}:{class_type}"
|
|
|
|
# Check if this node type is relevant for metadata collection
|
|
if class_type in NODE_EXTRACTORS or cache_key in self.node_cache:
|
|
# Check if we have cached metadata for this node
|
|
if cache_key in self.node_cache:
|
|
cached_data = self.node_cache[cache_key]
|
|
|
|
# Detect bypass (mode=4) / mute (mode=2) — these nodes
|
|
# were intentionally disabled and should not contribute
|
|
# overwrite values from a previous execution's cache.
|
|
node_mode = node_data.get("mode", 0)
|
|
node_is_disabled = node_mode in (2, 4)
|
|
|
|
# Apply cached metadata to the current metadata
|
|
for category in self.metadata_categories:
|
|
if category == OVERWRITE and node_is_disabled:
|
|
continue
|
|
if category in cached_data and node_id in cached_data[category]:
|
|
if node_id not in metadata[category]:
|
|
cached_entry = cached_data[category][node_id]
|
|
if category == MODELS:
|
|
cached_entry = self._validated_checkpoint_fill(
|
|
cached_entry, node_data
|
|
)
|
|
if cached_entry is None:
|
|
continue
|
|
metadata[category][node_id] = cached_entry
|
|
|
|
@staticmethod
|
|
def _validated_checkpoint_fill(cached_entry, node_data):
|
|
"""Guard a MODELS cache fill against stale checkpoint names.
|
|
|
|
A loader that has not executed since the user switched models still
|
|
holds the previous model in its cache entry. When the node's current
|
|
inputs name a different file, the cached name is replaced with the one
|
|
derived from the current inputs; when no current name can be derived,
|
|
the stale entry is dropped entirely.
|
|
|
|
The cache is trusted as-is when:
|
|
* the entry is not a checkpoint record,
|
|
* the cached name carries no model-file extension (e.g. names set
|
|
from runtime attachments rather than a widget value), or
|
|
* the node inputs carry no recognizable model field (unknown loader).
|
|
"""
|
|
if not isinstance(cached_entry, dict):
|
|
return cached_entry
|
|
if cached_entry.get("type") != "checkpoint":
|
|
return cached_entry
|
|
cached_name = cached_entry.get("name")
|
|
if not cached_name or not cached_name.lower().endswith(_MODEL_FILE_EXTENSIONS):
|
|
return cached_entry
|
|
|
|
inputs = node_data.get("inputs", {})
|
|
candidates = _checkpoint_name_candidates(inputs)
|
|
if not candidates:
|
|
return cached_entry
|
|
if cached_name in candidates:
|
|
return cached_entry
|
|
|
|
# Stale — rebuild from the current inputs when possible.
|
|
for field in _MODEL_NAME_FIELDS:
|
|
value = inputs.get(field)
|
|
if isinstance(value, str) and value.strip().lower().endswith(
|
|
_MODEL_FILE_EXTENSIONS
|
|
):
|
|
rebuilt = dict(cached_entry)
|
|
rebuilt["name"] = value.strip()
|
|
return rebuilt
|
|
return None
|
|
|
|
def record_node_execution(self, node_id, class_type, inputs, outputs, return_types=None):
|
|
"""Record information about a node's execution"""
|
|
if not self.current_prompt_id:
|
|
return
|
|
|
|
# Add to execution order and mark as executed
|
|
if node_id not in self.executed_nodes:
|
|
self.executed_nodes.add(node_id)
|
|
self.prompt_metadata[self.current_prompt_id]["execution_order"].append(
|
|
node_id
|
|
)
|
|
|
|
# Process inputs to simplify working with them
|
|
processed_inputs = {}
|
|
for input_name, input_values in inputs.items():
|
|
if isinstance(input_values, list) and len(input_values) > 0:
|
|
# For single values, just use the first one (most common case)
|
|
processed_inputs[input_name] = input_values[0]
|
|
else:
|
|
processed_inputs[input_name] = input_values
|
|
|
|
# Extract node-specific metadata
|
|
extractor = NODE_EXTRACTORS.get(class_type, GenericNodeExtractor)
|
|
if extractor is GenericNodeExtractor:
|
|
extractor.extract(node_id, processed_inputs, outputs,
|
|
self.prompt_metadata[self.current_prompt_id],
|
|
return_types=return_types)
|
|
else:
|
|
extractor.extract(node_id, processed_inputs, outputs,
|
|
self.prompt_metadata[self.current_prompt_id])
|
|
|
|
# Cache this node's metadata
|
|
self._cache_node_metadata(node_id, class_type)
|
|
|
|
def update_node_execution(self, node_id, class_type, outputs, return_types=None):
|
|
"""Update node metadata with output information"""
|
|
if not self.current_prompt_id:
|
|
return
|
|
|
|
# Process outputs to make them more usable
|
|
processed_outputs = outputs
|
|
|
|
# Use the same extractor to update with outputs
|
|
extractor = NODE_EXTRACTORS.get(class_type, GenericNodeExtractor)
|
|
if hasattr(extractor, "update"):
|
|
if extractor is GenericNodeExtractor:
|
|
extractor.update(
|
|
node_id, processed_outputs,
|
|
self.prompt_metadata[self.current_prompt_id],
|
|
return_types=return_types,
|
|
)
|
|
else:
|
|
extractor.update(
|
|
node_id, processed_outputs,
|
|
self.prompt_metadata[self.current_prompt_id],
|
|
)
|
|
|
|
# Update the cached metadata for this node
|
|
self._cache_node_metadata(node_id, class_type)
|
|
|
|
def _cache_node_metadata(self, node_id, class_type):
|
|
"""Cache the metadata for a specific node"""
|
|
if not self.current_prompt_id or not node_id or not class_type:
|
|
return
|
|
|
|
# Create a cache key combining node_id and class_type
|
|
cache_key = f"{node_id}:{class_type}"
|
|
|
|
# Create a shallow copy of the node's metadata
|
|
node_metadata = {}
|
|
current_metadata = self.prompt_metadata[self.current_prompt_id]
|
|
|
|
for category in self.metadata_categories:
|
|
if category in current_metadata and node_id in current_metadata[category]:
|
|
if category not in node_metadata:
|
|
node_metadata[category] = {}
|
|
node_metadata[category][node_id] = current_metadata[category][node_id]
|
|
|
|
# Save new metadata or clear stale cache entries when metadata is empty
|
|
if any(node_metadata.values()):
|
|
self.node_cache[cache_key] = node_metadata
|
|
else:
|
|
self.node_cache.pop(cache_key, None)
|
|
|
|
def clear_unused_cache(self):
|
|
"""Clean up node_cache entries that are no longer in use"""
|
|
# Collect all node_ids currently in prompt_metadata
|
|
active_node_ids = set()
|
|
for prompt_data in self.prompt_metadata.values():
|
|
for category in self.metadata_categories:
|
|
if category in prompt_data:
|
|
active_node_ids.update(prompt_data[category].keys())
|
|
|
|
# Find cache keys that are no longer needed
|
|
keys_to_remove = []
|
|
for cache_key in self.node_cache:
|
|
node_id = cache_key.split(":")[0]
|
|
if node_id not in active_node_ids:
|
|
keys_to_remove.append(cache_key)
|
|
|
|
# Remove cache entries that are no longer needed
|
|
for key in keys_to_remove:
|
|
del self.node_cache[key]
|
|
|
|
def clear_metadata(self, prompt_id=None):
|
|
"""Clear metadata for a specific prompt or reset all data"""
|
|
if prompt_id is not None:
|
|
if prompt_id in self.prompt_metadata:
|
|
del self.prompt_metadata[prompt_id]
|
|
# Clean up cache after removing prompt
|
|
self.clear_unused_cache()
|
|
else:
|
|
# Reset all data
|
|
self._reset()
|
|
|
|
def get_first_decoded_image(self, prompt_id=None):
|
|
"""Get the first decoded image result"""
|
|
key = prompt_id if prompt_id is not None else self.current_prompt_id
|
|
if key not in self.prompt_metadata:
|
|
return None
|
|
|
|
metadata = self.prompt_metadata[key]
|
|
if IMAGES in metadata and "first_decode" in metadata[IMAGES]:
|
|
image_data = metadata[IMAGES]["first_decode"]["image"]
|
|
|
|
# If it's an image batch or tuple, handle various formats
|
|
if isinstance(image_data, (list, tuple)) and len(image_data) > 0:
|
|
# Return first element of list/tuple
|
|
return image_data[0]
|
|
|
|
# If it's a tensor, return as is for processing in the route handler
|
|
return image_data
|
|
|
|
# If no image is found in the current metadata, try to find it in the cache
|
|
# This handles the case where VAEDecode was cached by ComfyUI and not executed
|
|
prompt_obj = metadata.get("current_prompt")
|
|
if prompt_obj and hasattr(prompt_obj, "original_prompt"):
|
|
original_prompt = prompt_obj.original_prompt
|
|
for node_id, node_data in original_prompt.items():
|
|
class_type = node_data.get("class_type")
|
|
if class_type and class_type in NODE_CLASS_MAPPINGS:
|
|
class_obj = NODE_CLASS_MAPPINGS[class_type]
|
|
class_name = class_obj.__name__
|
|
# Check if this is a VAEDecode node
|
|
if class_name == "VAEDecode":
|
|
# Try to find this node in the cache
|
|
cache_key = f"{node_id}:{class_name}"
|
|
if cache_key in self.node_cache:
|
|
cached_data = self.node_cache[cache_key]
|
|
if IMAGES in cached_data and node_id in cached_data[IMAGES]:
|
|
image_data = cached_data[IMAGES][node_id]["image"]
|
|
# Handle different image formats
|
|
if (
|
|
isinstance(image_data, (list, tuple))
|
|
and len(image_data) > 0
|
|
):
|
|
return image_data[0]
|
|
return image_data
|
|
|
|
return None
|