Vito
a52ff4e052
* feat(ui): add Focused Prompt Manager node (#186) * feat(ui): autocomplete saved Focused metadata with shared controls (#186) * feat(ui): autocomplete Focused library filters (#186)
231 lines
8.8 KiB
Python
231 lines
8.8 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", "PromptManagerFocused"}
|
|
)
|
|
|
|
# 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 family 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
|