Vito Sansevero
5b7fa1476a
The cycle guard used one visited set for the whole walk, so a string node read twice (a diamond, or StringConcatenate with the same source on both inputs) made the second branch None and the prompt text unknowable. Cycle detection is now scoped to the current path; a separate visit counter keeps MAX_VISITS as the bound on total work.
214 lines
7.9 KiB
Python
214 lines
7.9 KiB
Python
"""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
|
|
|
|
# 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; the visit counter bounds total work regardless.
|
|
path = {str(node_id)}
|
|
visits = [0]
|
|
|
|
def resolve(value: Any) -> Optional[str]:
|
|
if isinstance(value, str):
|
|
return value
|
|
if not _is_link(value):
|
|
return None
|
|
visits[0] += 1
|
|
if visits[0] > MAX_VISITS:
|
|
return None
|
|
source_id = str(value[0])
|
|
if source_id in path:
|
|
return None
|
|
source = graph.get(source_id)
|
|
if not isinstance(source, dict):
|
|
return None
|
|
path.add(source_id)
|
|
try:
|
|
return _resolve_string_node(
|
|
str(source.get("class_type", "")), _inputs(source), resolve
|
|
)
|
|
finally:
|
|
path.discard(source_id)
|
|
|
|
text = resolve(_inputs(graph.get(str(node_id))).get("text"))
|
|
return text.strip() if isinstance(text, str) and text.strip() else None
|