feat(metadata): resolve canonical prompt pairs with bounded provenance (#193)
This commit is contained in:
Vito Sansevero
@@ -0,0 +1,427 @@
|
||||
"""Conservative prompt-pair resolution contracts, independent of ComfyUI."""
|
||||
|
||||
import unittest
|
||||
|
||||
from utils.prompt_pair_resolver import OutputAdapter, resolve_prompt_pair
|
||||
|
||||
|
||||
def graph(positive: str = "portrait", negative: str = "blurry") -> dict:
|
||||
return {
|
||||
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": positive}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": negative}},
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class TestPromptPairResolver(unittest.TestCase):
|
||||
def test_resolves_same_sampler_and_preserves_output_provenance(self):
|
||||
result = resolve_prompt_pair({"prompt": graph()})
|
||||
self.assertEqual(result["state"], "resolved")
|
||||
self.assertEqual(result["positive"]["text"], "portrait")
|
||||
self.assertEqual(result["negative"]["text"], "blurry")
|
||||
self.assertEqual(result["positive"]["provenance"]["node_id"], "1")
|
||||
self.assertEqual(result["positive"]["provenance"]["output_slot"], 0)
|
||||
self.assertEqual(result["sampler_id"], "3")
|
||||
|
||||
def test_explicit_empty_negative_is_resolved_but_missing_is_absent(self):
|
||||
self.assertEqual(resolve_prompt_pair(graph(negative=""))["state"], "resolved")
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(graph(negative=""))["negative"]["text"], ""
|
||||
)
|
||||
source = graph()
|
||||
del source["3"]["inputs"]["negative"]
|
||||
result = resolve_prompt_pair(source)
|
||||
self.assertEqual(result["state"], "absent")
|
||||
self.assertEqual(result["positive"]["text"], "portrait")
|
||||
self.assertIsNone(result["negative"])
|
||||
|
||||
def test_multiple_samplers_are_ambiguous_even_when_texts_match(self):
|
||||
source = graph()
|
||||
source["4"] = {
|
||||
"class_type": "KSamplerAdvanced",
|
||||
"inputs": source["3"]["inputs"],
|
||||
}
|
||||
result = resolve_prompt_pair(source)
|
||||
self.assertEqual(result["state"], "ambiguous")
|
||||
self.assertEqual(len(result["candidates"]), 2)
|
||||
self.assertIsNone(result["positive"])
|
||||
|
||||
def test_saver_anchor_selects_only_its_sampler(self):
|
||||
source = graph()
|
||||
source.update(
|
||||
{
|
||||
"4": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"positive": ["2", 0], "negative": ["1", 0]},
|
||||
},
|
||||
"5": {"class_type": "VAEDecode", "inputs": {"samples": ["4", 0]}},
|
||||
"6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
|
||||
}
|
||||
)
|
||||
result = resolve_prompt_pair(source, output_node_id="6")
|
||||
self.assertEqual(result["state"], "resolved")
|
||||
self.assertEqual(result["sampler_id"], "4")
|
||||
self.assertEqual(result["positive"]["text"], "blurry")
|
||||
|
||||
def test_unknown_anchor_transform_and_multistage_are_unresolved(self):
|
||||
source = graph()
|
||||
source["5"] = {
|
||||
"class_type": "MysteryImageTransform",
|
||||
"inputs": {"images": ["3", 0]},
|
||||
}
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source, output_node_id="5")["state"], "unresolved"
|
||||
)
|
||||
source["4"] = {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {**source["3"]["inputs"], "latent_image": ["3", 0]},
|
||||
}
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source, output_node_id="4")["state"], "unresolved"
|
||||
)
|
||||
|
||||
def test_combines_unknown_encoders_and_sdxl_channels_are_not_guessed(self):
|
||||
for name in ("ConditioningCombine", "CLIPTextEncodeSDXL", "UnknownEncoder"):
|
||||
with self.subTest(name=name):
|
||||
source = graph()
|
||||
source["1"] = {
|
||||
"class_type": name,
|
||||
"inputs": {
|
||||
"text": "guess",
|
||||
"text_g": "global",
|
||||
"text_l": "local",
|
||||
"conditioning_1": ["2", 0],
|
||||
},
|
||||
}
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_prompt_manager_composition_and_string_slot_are_resolved(self):
|
||||
source = graph()
|
||||
source["8"] = {
|
||||
"class_type": "PromptManagerSplit",
|
||||
"inputs": {
|
||||
"text": "portrait",
|
||||
"prepend_text": " trigger ",
|
||||
"append_text": " detail ",
|
||||
},
|
||||
}
|
||||
source["1"]["inputs"]["text"] = ["8", 1]
|
||||
result = resolve_prompt_pair(source)
|
||||
self.assertEqual(result["positive"]["text"], "trigger portrait detail")
|
||||
self.assertIn(
|
||||
{"node_id": "8", "output_slot": 1}, result["positive"]["provenance"]["path"]
|
||||
)
|
||||
source["3"]["inputs"]["positive"] = ["8", 0]
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source)["positive"]["text"], "trigger portrait detail"
|
||||
)
|
||||
source["1"]["inputs"]["text"] = ["8", 0]
|
||||
source["3"]["inputs"]["positive"] = ["1", 0]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_concat_and_reroute_resolve_safe_linked_strings(self):
|
||||
source = graph()
|
||||
source["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "portrait"}}
|
||||
source["9"] = {
|
||||
"class_type": "StringConcatenate",
|
||||
"inputs": {"string_a": ["8", 0], "string_b": ["8", 0], "delimiter": ", "},
|
||||
}
|
||||
source["10"] = {"class_type": "Reroute", "inputs": {"value": ["9", 0]}}
|
||||
source["1"]["inputs"]["text"] = ["10", 0]
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source)["positive"]["text"], "portrait, portrait"
|
||||
)
|
||||
|
||||
def test_slot_adapter_does_not_conflate_two_outputs_of_one_node(self):
|
||||
source = graph()
|
||||
source["8"] = {
|
||||
"class_type": "FuturePair",
|
||||
"inputs": {"first": "positive", "second": "negative"},
|
||||
}
|
||||
source["3"]["inputs"] = {"positive": ["8", 0], "negative": ["8", 1]}
|
||||
adapters = {
|
||||
("FuturePair", 0): OutputAdapter("first", "conditioning"),
|
||||
("FuturePair", 1): OutputAdapter("second", "conditioning"),
|
||||
}
|
||||
result = resolve_prompt_pair(source, adapters=adapters)
|
||||
self.assertEqual(result["positive"]["text"], "positive")
|
||||
self.assertEqual(result["negative"]["text"], "negative")
|
||||
self.assertNotEqual(
|
||||
result["positive"]["provenance"]["output_slot"],
|
||||
result["negative"]["provenance"]["output_slot"],
|
||||
)
|
||||
|
||||
def test_prompt_manager_text_has_only_its_real_string_output_zero(self):
|
||||
source = graph()
|
||||
source["8"] = {
|
||||
"class_type": "PromptManagerText",
|
||||
"inputs": {"text": "text node", "prepend_text": "before"},
|
||||
}
|
||||
source["1"]["inputs"]["text"] = ["8", 0]
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source)["positive"]["text"], "before text node"
|
||||
)
|
||||
source["1"]["inputs"]["text"] = ["8", 1]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
source["3"]["inputs"]["positive"] = ["8", 0]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_malformed_class_names_and_unknown_samplers_do_not_claim_absence(self):
|
||||
for name in ([], {}, None, "UnknownSampler"):
|
||||
source = graph()
|
||||
source["3"]["class_type"] = name
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_pair_adapter_composes_only_positive_and_retains_recipe_text(self):
|
||||
source = graph()
|
||||
source["8"] = {
|
||||
"class_type": "PromptManagerPair",
|
||||
"inputs": {
|
||||
"text": "positive",
|
||||
"negative_text": "negative",
|
||||
"prepend_text": "trigger",
|
||||
"append_text": "detail",
|
||||
},
|
||||
}
|
||||
source["3"]["inputs"] = {"positive": ["8", 0], "negative": ["8", 1]}
|
||||
result = resolve_prompt_pair(source)
|
||||
self.assertEqual(result["positive"]["text"], "trigger positive detail")
|
||||
self.assertEqual(result["positive"]["storage_text"], "positive")
|
||||
self.assertEqual(result["negative"]["text"], "negative")
|
||||
self.assertEqual(result["negative"]["storage_text"], "negative")
|
||||
source["1"]["inputs"]["text"] = ["8", 2]
|
||||
source["2"]["inputs"]["text"] = ["8", 3]
|
||||
source["3"]["inputs"] = {"positive": ["1", 0], "negative": ["2", 0]}
|
||||
result = resolve_prompt_pair(source)
|
||||
self.assertEqual(result["positive"]["storage_text"], "positive")
|
||||
self.assertEqual(result["negative"]["text"], "negative")
|
||||
source["3"]["inputs"]["negative"] = ["8", 3]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_linked_recipe_and_ordinary_clip_storage_values_are_explicit(self):
|
||||
source = graph()
|
||||
result = resolve_prompt_pair(source)
|
||||
self.assertEqual(result["positive"]["storage_text"], "portrait")
|
||||
source["8"] = {
|
||||
"class_type": "PrimitiveString",
|
||||
"inputs": {"value": "linked main"},
|
||||
}
|
||||
source["1"] = {
|
||||
"class_type": "PromptManager",
|
||||
"inputs": {"text": ["8", 0], "prepend_text": "before"},
|
||||
}
|
||||
result = resolve_prompt_pair(source)
|
||||
self.assertEqual(result["positive"]["text"], "before linked main")
|
||||
self.assertEqual(result["positive"]["storage_text"], "linked main")
|
||||
a1111 = resolve_prompt_pair(
|
||||
{
|
||||
"parameters": (
|
||||
"portrait\nNegative prompt: blurry\n" "Steps: 20, Sampler: Euler"
|
||||
)
|
||||
}
|
||||
)
|
||||
self.assertEqual(a1111["positive"]["storage_text"], "portrait")
|
||||
self.assertEqual(a1111["negative"]["storage_text"], "blurry")
|
||||
|
||||
def test_cycles_depth_visits_and_text_growth_are_bounded(self):
|
||||
source = graph()
|
||||
source["8"] = {"class_type": "Reroute", "inputs": {"value": ["8", 0]}}
|
||||
source["1"]["inputs"]["text"] = ["8", 0]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(graph(), max_visits=1)["state"], "unresolved"
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(graph(), max_depth=0)["state"], "unresolved"
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(graph(), max_text_length=3)["state"], "unresolved"
|
||||
)
|
||||
|
||||
def test_malformed_metadata_and_links_never_raise(self):
|
||||
for value in (
|
||||
None,
|
||||
[],
|
||||
"raw text",
|
||||
{},
|
||||
{"prompt": []},
|
||||
{"prompt": {"1": None}},
|
||||
):
|
||||
with self.subTest(value=value):
|
||||
self.assertIn(
|
||||
resolve_prompt_pair(value)["state"], ("absent", "unresolved")
|
||||
)
|
||||
for link in (["missing", 0], ["1", -1], ["1", True], ["1", 17], ["1"], 5):
|
||||
source = graph()
|
||||
source["3"]["inputs"]["positive"] = link
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_a1111_explicit_parameters_and_empty_negative(self):
|
||||
for negative in ("bad anatomy\nblurry", ""):
|
||||
result = resolve_prompt_pair(
|
||||
{
|
||||
"parameters": (
|
||||
f"portrait\nNegative prompt: {negative}\n"
|
||||
"Steps: 20, Sampler: Euler, Seed: 1"
|
||||
)
|
||||
}
|
||||
)
|
||||
self.assertEqual(result["state"], "resolved")
|
||||
self.assertEqual(result["negative"]["text"], negative)
|
||||
self.assertEqual(result["positive"]["provenance"]["format"], "a1111")
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair({"parameters": "portrait\nSteps: 20, Sampler: Euler"})[
|
||||
"state"
|
||||
],
|
||||
"absent",
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair({"parameters": "unstructured arbitrary text"})["state"],
|
||||
"unresolved",
|
||||
)
|
||||
|
||||
def test_was_concat_preserves_configured_whitespace_and_rejects_bad_delimiters(
|
||||
self,
|
||||
):
|
||||
source = graph()
|
||||
source["8"] = {
|
||||
"class_type": "Text Concatenate",
|
||||
"inputs": {
|
||||
"text_a": " portrait ",
|
||||
"text_b": "",
|
||||
"text_c": "studio",
|
||||
"delimiter": ", ",
|
||||
},
|
||||
}
|
||||
source["1"]["inputs"]["text"] = ["8", 0]
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source)["positive"]["text"], "portrait, studio"
|
||||
)
|
||||
source["8"]["inputs"]["clean_whitespace"] = "false"
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source)["positive"]["text"], " portrait , studio"
|
||||
)
|
||||
source["8"]["inputs"]["delimiter"] = []
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
source["8"]["inputs"] = {1: "bad key"}
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_valid_empty_latent_and_reroute_but_not_unknown_or_cyclic_latents(self):
|
||||
source = graph()
|
||||
source["8"] = {"class_type": "EmptyLatentImage", "inputs": {}}
|
||||
source["9"] = {"class_type": "Reroute", "inputs": {"value": ["8", 0]}}
|
||||
source["3"]["inputs"]["latent_image"] = ["9", 0]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "resolved")
|
||||
source["9"]["inputs"]["value"] = ["9", 0]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
source["3"]["inputs"]["latent_image"] = ["8", 0]
|
||||
source["8"]["class_type"] = "UnknownLatentTransform"
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
source["3"]["inputs"]["latent_image"] = None
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_cyclic_or_invalid_output_anchors_never_guess(self):
|
||||
source = graph()
|
||||
source["8"] = {"class_type": "Reroute", "inputs": {"value": ["8", 0]}}
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source, output_node_id="8")["state"], "unresolved"
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source, output_node_id="missing")["state"], "unresolved"
|
||||
)
|
||||
source["8"] = {"class_type": "SaveImage", "inputs": {"images": ["3", 1]}}
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source, output_node_id="8")["state"], "unresolved"
|
||||
)
|
||||
source["8"] = {"class_type": "SaveImage", "inputs": {"images": ["3", 0]}}
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source, output_node_id="8", max_depth=1)["state"],
|
||||
"unresolved",
|
||||
)
|
||||
|
||||
def test_malformed_reroute_empty_positive_and_duplicate_string_ids(self):
|
||||
source = graph(positive="")
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
source = graph()
|
||||
source["8"] = {
|
||||
"class_type": "Reroute",
|
||||
"inputs": {"value": ["1", 0], "input": ["2", 0]},
|
||||
}
|
||||
source["3"]["inputs"]["positive"] = ["8", 0]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
source[3] = source["3"]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_a1111_malformed_ambiguous_and_bounded_parameters(self):
|
||||
self.assertEqual(resolve_prompt_pair({"parameters": None})["state"], "absent")
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair({"parameters": "x" * 20}, max_text_length=10)["state"],
|
||||
"unresolved",
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(
|
||||
{"parameters": "\nNegative prompt: bad\nSteps: 20, Sampler: Euler"}
|
||||
)["state"],
|
||||
"unresolved",
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(
|
||||
{
|
||||
"parameters": (
|
||||
"portrait\nNegative prompt: bad\nNegative prompt: worse\n"
|
||||
"Steps: 20, Sampler: Euler"
|
||||
)
|
||||
}
|
||||
)["state"],
|
||||
"ambiguous",
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(
|
||||
{
|
||||
"parameters": (
|
||||
"portrait\nNegative prompt: bad\n"
|
||||
"Steps: 20, Sampler: Euler\nextra text"
|
||||
)
|
||||
}
|
||||
)["state"],
|
||||
"unresolved",
|
||||
)
|
||||
|
||||
def test_exponential_concat_text_is_bounded_before_building_a_large_string(self):
|
||||
source = graph()
|
||||
source["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "xx"}}
|
||||
for index in range(9, 19):
|
||||
source[str(index)] = {
|
||||
"class_type": "StringConcatenate",
|
||||
"inputs": {
|
||||
"string_a": [str(index - 1), 0],
|
||||
"string_b": [str(index - 1), 0],
|
||||
},
|
||||
}
|
||||
source["1"]["inputs"]["text"] = ["18", 0]
|
||||
result = resolve_prompt_pair(source, max_text_length=100)
|
||||
self.assertEqual(result["state"], "unresolved")
|
||||
self.assertIn("length limit", result["reason"])
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(graph(), max_depth=1)["state"], "unresolved"
|
||||
)
|
||||
|
||||
def test_api_graph_is_authoritative_over_unrelated_parameters(self):
|
||||
source = graph()
|
||||
source["1"]["class_type"] = "UnknownEncoder"
|
||||
result = resolve_prompt_pair(
|
||||
{
|
||||
"prompt": source,
|
||||
"parameters": (
|
||||
"fallback\nNegative prompt: wrong\n" "Steps: 20, Sampler: Euler"
|
||||
),
|
||||
}
|
||||
)
|
||||
self.assertEqual(result["state"], "unresolved")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,394 @@
|
||||
"""Conservatively recover a prompt pair from one sampler or A1111 parameters.
|
||||
|
||||
``resolve_prompt_pair`` returns a dictionary with ``state``, ``positive``,
|
||||
``negative``, ``sampler_id``, ``reason`` and ``candidates``. Each recovered side
|
||||
contains effective ``text``, deduplicated-recipe ``storage_text`` and
|
||||
``provenance``; ComfyUI provenance records the exact output
|
||||
slot and complete followed link path. States distinguish a resolved pair (an
|
||||
explicit empty negative is valid), absent data, unsupported/unresolved data and
|
||||
ambiguous sampler selection. Unknown conditioning semantics are never flattened.
|
||||
|
||||
Future nodes can register declarative ``OutputAdapter`` entries by
|
||||
``(class_type, output_slot)``. This does not prescribe a future paired-node schema.
|
||||
Effective text is reconstructed from recorded graph inputs; runtime-only trigger
|
||||
injection cannot be recovered unless it was recorded. The resolver has no
|
||||
ComfyUI, database, configuration or filesystem dependencies.
|
||||
"""
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Mapping, Optional, Tuple
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OutputAdapter:
|
||||
"""A known output's text input, output kind, and optional PM composition.
|
||||
|
||||
``kind`` is ``conditioning`` or ``string``. Composition is opt-in and follows
|
||||
the existing PromptManager prepend/main/append rule. Adapters describe pure
|
||||
text provenance, not arbitrary executable transforms. ``storage_input`` keeps
|
||||
the resolved input as the recipe even without composition (negative editor).
|
||||
"""
|
||||
|
||||
input_name: str
|
||||
kind: str
|
||||
compose: bool = False
|
||||
storage_input: bool = False
|
||||
|
||||
|
||||
_SAMPLERS = frozenset({"KSampler", "KSamplerAdvanced", "SamplerCustom"})
|
||||
_PROMPT_NODES = frozenset(
|
||||
{"PromptManager", "PromptManagerFocused", "PromptManagerSplit"}
|
||||
)
|
||||
_DEFAULT_ADAPTERS = {
|
||||
(name, slot): OutputAdapter("text", kind, True)
|
||||
for name in _PROMPT_NODES
|
||||
for slot, kind in ((0, "conditioning"), (1, "string"))
|
||||
}
|
||||
for _slot, _kind in ((0, "conditioning"), (2, "string")):
|
||||
_DEFAULT_ADAPTERS[("PromptManagerPair", _slot)] = OutputAdapter("text", _kind, True)
|
||||
for _slot, _kind in ((1, "conditioning"), (3, "string")):
|
||||
_DEFAULT_ADAPTERS[("PromptManagerPair", _slot)] = OutputAdapter(
|
||||
"negative_text", _kind, storage_input=True
|
||||
)
|
||||
_DEFAULT_ADAPTERS[("PromptManagerText", 0)] = OutputAdapter("text", "string", True)
|
||||
_DEFAULT_ADAPTERS[("CLIPTextEncode", 0)] = OutputAdapter("text", "conditioning")
|
||||
for _name in ("PrimitiveString", "PrimitiveStringMultiline"):
|
||||
_DEFAULT_ADAPTERS[(_name, 0)] = OutputAdapter("value", "string")
|
||||
for _name in ("Text Multiline", "String Literal", "ShowText|pysssss"):
|
||||
_DEFAULT_ADAPTERS[(_name, 0)] = OutputAdapter("text", "string")
|
||||
|
||||
|
||||
class _Unresolved(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def _link(value: Any) -> bool:
|
||||
return (
|
||||
isinstance(value, (list, tuple))
|
||||
and len(value) == 2
|
||||
and isinstance(value[0], (str, int))
|
||||
and not isinstance(value[0], bool)
|
||||
and isinstance(value[1], int)
|
||||
and not isinstance(value[1], bool)
|
||||
and value[1] >= 0
|
||||
)
|
||||
|
||||
|
||||
def _result(
|
||||
state: str,
|
||||
positive: Optional[dict] = None,
|
||||
negative: Optional[dict] = None,
|
||||
sampler_id: Optional[str] = None,
|
||||
reason: str = "",
|
||||
candidates: Optional[list] = None,
|
||||
) -> Dict[str, Any]:
|
||||
return {
|
||||
"state": state,
|
||||
"positive": positive,
|
||||
"negative": negative,
|
||||
"sampler_id": sampler_id,
|
||||
"reason": reason,
|
||||
"candidates": candidates or [],
|
||||
}
|
||||
|
||||
|
||||
class _GraphResolver:
|
||||
def __init__(
|
||||
self,
|
||||
graph: dict,
|
||||
adapters: Mapping[Tuple[str, int], OutputAdapter],
|
||||
max_visits: int,
|
||||
max_depth: int,
|
||||
max_text_length: int,
|
||||
) -> None:
|
||||
self.graph = graph
|
||||
self.adapters = {**_DEFAULT_ADAPTERS, **adapters}
|
||||
self.budget = max_visits
|
||||
self.max_depth = max_depth
|
||||
self.max_text_length = max_text_length
|
||||
self.memo: dict = {}
|
||||
self.active: set = set()
|
||||
|
||||
def node(self, node_id: str) -> Tuple[str, dict]:
|
||||
node = self.graph.get(node_id)
|
||||
if not isinstance(node, dict) or not isinstance(node.get("inputs"), dict):
|
||||
raise _Unresolved("Missing or malformed graph node")
|
||||
if any(not isinstance(key, str) for key in node["inputs"]):
|
||||
raise _Unresolved("Malformed node input names")
|
||||
return str(node.get("class_type", "")), node["inputs"]
|
||||
|
||||
def bounded(self, text: str) -> str:
|
||||
if len(text) > self.max_text_length:
|
||||
raise _Unresolved("Resolved text exceeds the length limit")
|
||||
return text
|
||||
|
||||
def text(self, value: Any, kind: str, depth: int = 0) -> Tuple[str, list, str]:
|
||||
if depth >= self.max_depth:
|
||||
raise _Unresolved("Resolution depth limit reached")
|
||||
if isinstance(value, str) and kind == "string":
|
||||
return self.bounded(value), [], value
|
||||
if not _link(value):
|
||||
raise _Unresolved("Expected a typed graph link or text input")
|
||||
node_id, slot = str(value[0]), value[1]
|
||||
key = (node_id, slot, kind)
|
||||
if key in self.active:
|
||||
raise _Unresolved("Cyclic prompt graph")
|
||||
if key in self.memo:
|
||||
return self.memo[key]
|
||||
self.budget -= 1
|
||||
if self.budget < 0:
|
||||
raise _Unresolved("Resolution visit limit reached")
|
||||
name, inputs = self.node(node_id)
|
||||
self.active.add(key)
|
||||
try:
|
||||
if name == "Reroute" and slot == 0:
|
||||
keys = [key for key in ("value", "input", "in") if key in inputs]
|
||||
if len(keys) != 1:
|
||||
raise _Unresolved("Reroute input is not uniquely identifiable")
|
||||
text, path, storage_text = self.text(inputs[keys[0]], kind, depth + 1)
|
||||
elif (name, slot) in self.adapters:
|
||||
adapter = self.adapters[(name, slot)]
|
||||
if not isinstance(adapter, OutputAdapter) or adapter.kind != kind:
|
||||
raise _Unresolved("Output slot has an incompatible type")
|
||||
text, path, storage_text = self.text(
|
||||
inputs.get(adapter.input_name), "string", depth + 1
|
||||
)
|
||||
if adapter.compose or adapter.storage_input:
|
||||
storage_text = text
|
||||
if adapter.compose:
|
||||
before, before_path, _ = self.text(
|
||||
inputs.get("prepend_text", ""), "string", depth + 1
|
||||
)
|
||||
after, after_path, _ = self.text(
|
||||
inputs.get("append_text", ""), "string", depth + 1
|
||||
)
|
||||
text = self.bounded(
|
||||
" ".join(
|
||||
part
|
||||
for part in (before.strip(), text, after.strip())
|
||||
if part
|
||||
)
|
||||
)
|
||||
path = before_path + path + after_path
|
||||
elif (
|
||||
name in ("StringConcatenate", "Text Concatenate")
|
||||
and slot == 0
|
||||
and kind == "string"
|
||||
):
|
||||
if name == "StringConcatenate":
|
||||
keys = ["string_a", "string_b"]
|
||||
delimiter = inputs.get("delimiter", "")
|
||||
clean = False
|
||||
else:
|
||||
keys = sorted(
|
||||
key for key in inputs if re.fullmatch(r"text_[a-z]", key)
|
||||
)
|
||||
delimiter = inputs.get("delimiter", " ")
|
||||
clean = (
|
||||
str(inputs.get("clean_whitespace", "true")).lower() == "true"
|
||||
)
|
||||
if not keys or not isinstance(delimiter, str):
|
||||
raise _Unresolved("Malformed concatenation")
|
||||
parts, path, size = [], [], 0
|
||||
for input_name in keys:
|
||||
part, followed, _ = self.text(
|
||||
inputs.get(input_name, ""), "string", depth + 1
|
||||
)
|
||||
if clean:
|
||||
part = part.strip()
|
||||
path.extend(followed)
|
||||
if part or name == "StringConcatenate":
|
||||
size += len(part) + (len(delimiter) if parts else 0)
|
||||
if size > self.max_text_length:
|
||||
raise _Unresolved("Resolved text exceeds the length limit")
|
||||
parts.append(part)
|
||||
text = delimiter.join(parts)
|
||||
storage_text = text
|
||||
else:
|
||||
raise _Unresolved("Unsupported node semantics or output slot")
|
||||
# Memoized paths stay bounded even for diamond-shaped string graphs.
|
||||
path = [{"node_id": node_id, "output_slot": slot}] + path
|
||||
if len(path) > self.max_depth * max(1, len(self.graph)):
|
||||
raise _Unresolved("Provenance path limit reached")
|
||||
self.memo[key] = (text, path, storage_text)
|
||||
return text, path, storage_text
|
||||
finally:
|
||||
self.active.discard(key)
|
||||
|
||||
def side(self, value: Any, sampler_id: str) -> dict:
|
||||
text, path, storage_text = self.text(value, "conditioning")
|
||||
origin = path[0]
|
||||
return {
|
||||
"text": text,
|
||||
"storage_text": storage_text,
|
||||
"provenance": {
|
||||
"format": "comfyui",
|
||||
"sampler_id": sampler_id,
|
||||
**origin,
|
||||
"path": path,
|
||||
},
|
||||
}
|
||||
|
||||
def pair(self, sampler_id: str) -> dict:
|
||||
positive = negative = None
|
||||
try:
|
||||
_, inputs = self.node(sampler_id)
|
||||
# A second sampler in the latent ancestry denotes a multistage run;
|
||||
# associating the final image with one prompt pair would be a guess.
|
||||
if "latent_image" in inputs:
|
||||
self.latent_source(inputs["latent_image"], set(), 0)
|
||||
positive = self.side(inputs.get("positive"), sampler_id)
|
||||
if not positive["text"].strip():
|
||||
raise _Unresolved("Positive prompt is empty")
|
||||
if "negative" not in inputs:
|
||||
return _result(
|
||||
"absent",
|
||||
positive,
|
||||
sampler_id=sampler_id,
|
||||
reason="Negative input is absent",
|
||||
)
|
||||
negative = self.side(inputs["negative"], sampler_id)
|
||||
return _result("resolved", positive, negative, sampler_id)
|
||||
except _Unresolved as error:
|
||||
return _result("unresolved", positive, negative, sampler_id, str(error))
|
||||
|
||||
def latent_source(self, value: Any, seen: set, depth: int) -> None:
|
||||
if not _link(value) or depth >= self.max_depth:
|
||||
raise _Unresolved("Unknown latent ancestry")
|
||||
node_id, slot = str(value[0]), value[1]
|
||||
if node_id in seen or len(seen) >= max(0, self.budget):
|
||||
raise _Unresolved("Cyclic or excessive latent ancestry")
|
||||
seen.add(node_id)
|
||||
name, inputs = self.node(node_id)
|
||||
if name in _SAMPLERS:
|
||||
raise _Unresolved("Multistage sampler ancestry is ambiguous")
|
||||
if name in ("EmptyLatentImage", "EmptySD3LatentImage") and slot == 0:
|
||||
return
|
||||
if name == "Reroute" and slot == 0 and "value" in inputs:
|
||||
self.latent_source(inputs["value"], seen, depth + 1)
|
||||
return
|
||||
raise _Unresolved("Unsupported latent ancestry")
|
||||
|
||||
def anchor(self, node_id: str) -> str:
|
||||
seen = set()
|
||||
slot = 0
|
||||
while len(seen) < self.max_depth and len(seen) < self.budget:
|
||||
if node_id in seen:
|
||||
raise _Unresolved("Cyclic output ancestry")
|
||||
seen.add(node_id)
|
||||
name, inputs = self.node(node_id)
|
||||
if name in _SAMPLERS and slot == 0:
|
||||
return node_id
|
||||
key = {
|
||||
"SaveImage": "images",
|
||||
"PreviewImage": "images",
|
||||
"VAEDecode": "samples",
|
||||
"VAEDecodeTiled": "samples",
|
||||
"Reroute": "value",
|
||||
}.get(name)
|
||||
value = inputs.get(key) if key and slot == 0 else None
|
||||
if not _link(value):
|
||||
raise _Unresolved("Unsupported output ancestry")
|
||||
node_id, slot = str(value[0]), value[1]
|
||||
raise _Unresolved("Output ancestry limit reached")
|
||||
|
||||
|
||||
def _parameters(value: Any, max_text_length: int) -> dict:
|
||||
if not isinstance(value, str) or not value:
|
||||
return _result("absent", reason="No A1111 parameters")
|
||||
if len(value) > max_text_length:
|
||||
return _result("unresolved", reason="Parameters exceed the length limit")
|
||||
footers = list(re.finditer(r"(?m)^Steps:\s*\d+\s*,[^\n]*$", value))
|
||||
if len(footers) != 1 or value[footers[0].end() :].strip():
|
||||
return _result("unresolved", reason="A1111 settings footer is not unique")
|
||||
body = value[: footers[0].start()].rstrip("\r\n")
|
||||
delimiter = "\nNegative prompt:"
|
||||
if body.count(delimiter) > 1:
|
||||
return _result("ambiguous", reason="Multiple negative prompt delimiters")
|
||||
positive, found, negative = body.partition(delimiter)
|
||||
positive = positive.strip()
|
||||
if not positive:
|
||||
return _result("unresolved", reason="Positive prompt is empty")
|
||||
provenance = {"format": "a1111", "field": "parameters"}
|
||||
positive_side = {
|
||||
"text": positive,
|
||||
"storage_text": positive,
|
||||
"provenance": {**provenance, "role": "positive"},
|
||||
}
|
||||
if not found:
|
||||
return _result(
|
||||
"absent", positive_side, reason="Negative prompt delimiter is absent"
|
||||
)
|
||||
negative_side = {
|
||||
"text": negative.strip(),
|
||||
"storage_text": negative.strip(),
|
||||
"provenance": {**provenance, "role": "negative"},
|
||||
}
|
||||
return _result("resolved", positive_side, negative_side)
|
||||
|
||||
|
||||
def resolve_prompt_pair(
|
||||
metadata: Any,
|
||||
*,
|
||||
output_node_id: Optional[str] = None,
|
||||
adapters: Optional[Mapping[Tuple[str, int], OutputAdapter]] = None,
|
||||
max_visits: int = 500,
|
||||
max_depth: int = 64,
|
||||
max_text_length: int = 100_000,
|
||||
) -> Dict[str, Any]:
|
||||
"""Resolve parsed metadata; a present API graph is authoritative over text.
|
||||
|
||||
Accepts ``{"prompt": api_graph, "parameters": optional_a1111_text}`` or a
|
||||
direct API graph. An optional saver/sampler node ID anchors branch selection.
|
||||
Unanchored multiple samplers are ambiguous, even when their texts are equal.
|
||||
"""
|
||||
if not isinstance(metadata, dict) or not metadata:
|
||||
return _result("absent", reason="No parsed prompt metadata")
|
||||
if "prompt" in metadata:
|
||||
graph = metadata["prompt"]
|
||||
elif "parameters" in metadata:
|
||||
return _parameters(metadata["parameters"], max_text_length)
|
||||
else:
|
||||
graph = metadata
|
||||
if not isinstance(graph, dict) or not graph:
|
||||
return _result("unresolved", reason="API graph is missing or malformed")
|
||||
if len(graph) > max_visits or max_depth <= 0 or max_text_length <= 0:
|
||||
return _result("unresolved", reason="Resolution limit reached")
|
||||
normalized = {str(node_id): node for node_id, node in graph.items()}
|
||||
if len(normalized) != len(graph):
|
||||
return _result("unresolved", reason="Graph node IDs are not unique")
|
||||
resolver = _GraphResolver(
|
||||
normalized, adapters or {}, max_visits, max_depth, max_text_length
|
||||
)
|
||||
if output_node_id is not None:
|
||||
try:
|
||||
sampler_ids = [resolver.anchor(str(output_node_id))]
|
||||
except _Unresolved as error:
|
||||
return _result("unresolved", reason=str(error))
|
||||
else:
|
||||
sampler_ids = [
|
||||
node_id
|
||||
for node_id, node in normalized.items()
|
||||
if isinstance(node, dict)
|
||||
and isinstance(node.get("class_type"), str)
|
||||
and node["class_type"] in _SAMPLERS
|
||||
]
|
||||
if not sampler_ids:
|
||||
has_pair_inputs = any(
|
||||
isinstance(node, dict)
|
||||
and isinstance(node.get("inputs"), dict)
|
||||
and ("positive" in node["inputs"] or "negative" in node["inputs"])
|
||||
for node in normalized.values()
|
||||
)
|
||||
return _result(
|
||||
"unresolved" if has_pair_inputs else "absent",
|
||||
reason="No supported sampler pair found",
|
||||
)
|
||||
if len(sampler_ids) > 1:
|
||||
return _result(
|
||||
"ambiguous",
|
||||
reason="Multiple sampler branches",
|
||||
candidates=[resolver.pair(node_id) for node_id in sampler_ids],
|
||||
)
|
||||
return resolver.pair(sampler_ids[0])
|
||||
Reference in New Issue
Block a user