"""Find which PromptManager nodes feed a sampler's positive input in a ComfyUI graph. Works on the API-format prompt graph ({node_id: {"class_type", "inputs"}}) that ComfyUI queues and embeds in saved images. Usage counting and image linking use it so that negative-prompt PromptManager nodes never count as "the prompt" of a run. Pure functions with no ComfyUI imports, so they are unit-testable. Mirrors the link-tracing rules of web/js/comfy-metadata.js. """ import re from typing import Any, Dict, List, Optional PROMPT_MANAGER_TYPES = frozenset({"PromptManager", "PromptManagerText"}) # Upper bound on nodes expanded per graph; queued and embedded graphs are untrusted MAX_VISITS = 500 # Longest text resolve_text will build; StringConcatenate chains can double # the text at every level, so anything larger is treated as unknowable MAX_TEXT_LENGTH = 100_000 # Deepest chain of string nodes resolve_text follows (recursion stays bounded) MAX_DEPTH = 64 # Inputs that carry positive conditioning or prompt text towards a sampler _CONDITIONING_KEY = re.compile(r"^conditioning(_\w+)?$") _TEXT_KEY = re.compile(r"^(text|string|prompt)(_?[a-z0-9]+)?$", re.IGNORECASE) # Input names that carry a negative prompt (text_negative, neg_text, negative_prompt) _NEGATIVE_KEY = re.compile(r"(^|_)neg(ative)?($|_)", re.IGNORECASE) def follows_input(key: Any) -> bool: """Whether an input name carries positive conditioning or prompt text.""" if not isinstance(key, str) or _NEGATIVE_KEY.search(key): return False return key == "positive" or bool( _CONDITIONING_KEY.match(key) or _TEXT_KEY.match(key) ) def _is_link(value: Any) -> bool: return ( isinstance(value, (list, tuple)) and len(value) == 2 and isinstance(value[0], (str, int)) and isinstance(value[1], int) ) def _inputs(node: Any) -> Dict[str, Any]: inputs = node.get("inputs") if isinstance(node, dict) else None return inputs if isinstance(inputs, dict) else {} def _positive_roots(graph: Dict[str, Any]) -> List[Any]: """Links into samplers' positive inputs (KSampler, CFGGuider, BasicGuider, ...). A sampler is a node taking a model plus positive/negative conditioning, or a guider taking a model plus conditioning. Nodes named like a sampler also count without a model link, so a negative-only graph is never mistaken for one with no sampler at all. """ roots = [] for node in graph.values(): if not isinstance(node, dict): continue inputs = _inputs(node) class_name = str(node.get("class_type", "")).lower() has_model = _is_link(inputs.get("model")) if _is_link(inputs.get("positive")) and _is_link(inputs.get("negative")): if has_model or "sampler" in class_name: roots.append(inputs["positive"]) elif has_model and _is_link(inputs.get("conditioning")): if "guider" in class_name: roots.append(inputs["conditioning"]) return roots def _upstream_links(inputs: Dict[str, Any]) -> List[Any]: """Inputs to follow from a pass-through or encoder node, never `negative`.""" links = [] for key, value in inputs.items(): if not _is_link(value): continue if follows_input(key): links.append(value) return links def has_sampler(graph: Any) -> bool: """Whether the graph contains a node recognised as a sampler or guider.""" return isinstance(graph, dict) and bool(_positive_roots(graph)) def positive_prompt_nodes(graph: Any) -> List[str]: """Ids of PromptManager/PromptManagerText nodes that feed a positive input. Ordered by discovery, each id once. Malformed graphs yield an empty list. """ if not isinstance(graph, dict): return [] found: List[str] = [] seen = set() stack = list(reversed(_positive_roots(graph))) while stack and len(seen) < MAX_VISITS: node_id = str(stack.pop()[0]) if node_id in seen: continue seen.add(node_id) node = graph.get(node_id) if not isinstance(node, dict): continue if node.get("class_type") in PROMPT_MANAGER_TYPES: found.append(node_id) continue stack.extend(reversed(_upstream_links(_inputs(node)))) return found def run_prompt_nodes(graph: Any) -> List[str]: """Prompt nodes that represent a run: the positive ones, if any can be found. When the graph has no recognisable sampler (an unrecognised custom sampler, say), every PromptManager node counts, as before 3.2.4, so an exotic workflow never silently loses usage counting or image linking. When a sampler exists and no PromptManager node feeds its positive input, nothing counts: the only PromptManager nodes are negative prompts. """ positive = positive_prompt_nodes(graph) if positive or not isinstance(graph, dict) or has_sampler(graph): return positive return [ str(node_id) for node_id, node in graph.items() if isinstance(node, dict) and node.get("class_type") in PROMPT_MANAGER_TYPES ] def _join_text_inputs( inputs: Dict[str, Any], keys: List[str], resolve ) -> Optional[str]: """WAS-style join: resolve each key, optionally strip, skip empty, join by delimiter.""" delimiter = inputs.get("delimiter", " ") if not isinstance(delimiter, str): return None clean = str(inputs.get("clean_whitespace", "true")).lower() == "true" parts = [] for key in keys: value = resolve(inputs.get(key, "")) if value is None: return None value = value.strip() if clean else value if value: parts.append(value) return delimiter.join(parts) def _resolve_string_node( class_type: str, inputs: Dict[str, Any], resolve ) -> Optional[str]: """Output of a known pure string node, or None when the type isn't known.""" if class_type in ("PrimitiveString", "PrimitiveStringMultiline"): return resolve(inputs.get("value")) if class_type in ("ShowText|pysssss", "Text Multiline", "String Literal"): return resolve(inputs.get("text", inputs.get("string"))) if class_type == "StringConcatenate": a, b = resolve(inputs.get("string_a", "")), resolve(inputs.get("string_b", "")) delimiter = inputs.get("delimiter", "") if a is None or b is None or not isinstance(delimiter, str): return None return f"{a}{delimiter}{b}" if class_type == "Text Concatenate": keys = sorted(k for k in inputs if re.fullmatch(r"text_[a-z]", k)) return _join_text_inputs(inputs, keys, resolve) return None def resolve_text(graph: Any, node_id: str) -> Optional[str]: """A prompt node's `text` as it will be at run time, if knowable from the graph. Typed text, or text produced by known pure string nodes (core primitives and StringConcatenate, WAS Text Concatenate/Multiline, pysssss ShowText). Anything else, including PromptSearchList batches, returns None rather than a guess. Callers compare the result with the text the node actually received. """ if not isinstance(graph, dict): return None # Cycle detection is scoped to the current path so a node read by two # branches (a diamond, or StringConcatenate(string_a=X, string_b=X)) # resolves in both. Each node's text is computed once and memoised, so a # chain of diamonds costs one expansion per node rather than one per path, # and the shared visit budget bounds total work whatever the shape. A node # that fails because of a cycle fails from every path (the cycle is in the # graph, not in the walk), so memoising None is sound. path = {str(node_id)} memo: Dict[str, Optional[str]] = {} budget = [MAX_VISITS] def resolve(value: Any) -> Optional[str]: if isinstance(value, str): return value if len(value) <= MAX_TEXT_LENGTH else None if not _is_link(value): return None source_id = str(value[0]) if source_id in memo: return memo[source_id] if source_id in path or budget[0] <= 0 or len(path) > MAX_DEPTH: return None budget[0] -= 1 source = graph.get(source_id) if not isinstance(source, dict): memo[source_id] = None return None path.add(source_id) try: text = _resolve_string_node( str(source.get("class_type", "")), _inputs(source), resolve ) finally: path.discard(source_id) if isinstance(text, str) and len(text) > MAX_TEXT_LENGTH: text = None memo[source_id] = text return text text = resolve(_inputs(graph.get(str(node_id))).get("text")) return text.strip() if isinstance(text, str) and text.strip() else None