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.
436 lines
16 KiB
Python
436 lines
16 KiB
Python
"""Tests for utils.prompt_graph: which PromptManager nodes feed a positive input."""
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
import unittest
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from utils.prompt_graph import (
|
|
_upstream_links,
|
|
follows_input,
|
|
has_sampler,
|
|
positive_prompt_nodes,
|
|
resolve_text,
|
|
run_prompt_nodes,
|
|
)
|
|
|
|
|
|
def pm(text="a cat", **inputs):
|
|
return {"class_type": "PromptManager", "inputs": {"text": text, **inputs}}
|
|
|
|
|
|
def ksampler(positive, negative, model="20"):
|
|
return {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"positive": [positive, 0],
|
|
"negative": [negative, 0],
|
|
"model": [model, 0],
|
|
},
|
|
}
|
|
|
|
|
|
LOADER = {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": "m.safetensors"},
|
|
}
|
|
|
|
|
|
class TestPositivePromptNodes(unittest.TestCase):
|
|
def test_positive_and_negative_prompt_manager_nodes(self):
|
|
# Negative listed first, as in the #75 / pmtest workflow
|
|
graph = {
|
|
"176": pm("blurry"),
|
|
"134": pm("a cat"),
|
|
"5": ksampler("134", "176"),
|
|
"20": LOADER,
|
|
}
|
|
self.assertEqual(positive_prompt_nodes(graph), ["134"])
|
|
|
|
def test_negative_only_prompt_manager_is_never_positive(self):
|
|
graph = {
|
|
"176": pm("blurry"),
|
|
"6": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}},
|
|
"5": ksampler("6", "176"),
|
|
"20": LOADER,
|
|
}
|
|
self.assertEqual(positive_prompt_nodes(graph), [])
|
|
|
|
def test_through_controlnet_and_combine_pass_throughs(self):
|
|
graph = {
|
|
"1": pm("subject"),
|
|
"2": pm("style"),
|
|
"3": pm("blurry"),
|
|
"4": {
|
|
"class_type": "ConditioningCombine",
|
|
"inputs": {"conditioning_1": ["1", 0], "conditioning_2": ["2", 0]},
|
|
},
|
|
"5": {
|
|
"class_type": "ControlNetApplyAdvanced",
|
|
"inputs": {"positive": ["4", 0], "negative": ["3", 0], "strength": 1.0},
|
|
},
|
|
"6": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"positive": ["5", 0],
|
|
"negative": ["5", 1],
|
|
"model": ["20", 0],
|
|
},
|
|
},
|
|
"20": LOADER,
|
|
}
|
|
self.assertEqual(sorted(positive_prompt_nodes(graph)), ["1", "2"])
|
|
|
|
def test_cfg_guider_and_basic_guider(self):
|
|
graph = {
|
|
"1": pm("fox"),
|
|
"2": pm("watermark"),
|
|
"3": {
|
|
"class_type": "CFGGuider",
|
|
"inputs": {
|
|
"positive": ["1", 0],
|
|
"negative": ["2", 0],
|
|
"model": ["20", 0],
|
|
"cfg": 5,
|
|
},
|
|
},
|
|
"4": pm("owl"),
|
|
"5": {
|
|
"class_type": "BasicGuider",
|
|
"inputs": {"conditioning": ["4", 0], "model": ["20", 0]},
|
|
},
|
|
"20": LOADER,
|
|
}
|
|
self.assertEqual(sorted(positive_prompt_nodes(graph)), ["1", "4"])
|
|
|
|
def test_prompt_manager_text_wired_into_an_encoder(self):
|
|
graph = {
|
|
"1": {"class_type": "PromptManagerText", "inputs": {"text": "a castle"}},
|
|
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": ["1", 0]}},
|
|
"3": {"class_type": "PromptManagerText", "inputs": {"text": "lowres"}},
|
|
"4": {"class_type": "CLIPTextEncode", "inputs": {"text": ["3", 0]}},
|
|
"5": ksampler("2", "4"),
|
|
"20": LOADER,
|
|
}
|
|
self.assertEqual(positive_prompt_nodes(graph), ["1"])
|
|
|
|
def test_node_used_as_both_positive_and_negative_counts_as_positive(self):
|
|
graph = {"1": pm("same"), "5": ksampler("1", "1"), "20": LOADER}
|
|
self.assertEqual(positive_prompt_nodes(graph), ["1"])
|
|
|
|
def test_two_samplers_sharing_a_prompt_list_it_once(self):
|
|
graph = {
|
|
"1": pm("hero"),
|
|
"2": pm("blurry"),
|
|
"5": ksampler("1", "2"),
|
|
"14": ksampler("1", "2"),
|
|
"20": LOADER,
|
|
}
|
|
self.assertEqual(positive_prompt_nodes(graph), ["1"])
|
|
|
|
def test_graph_without_sampler_or_malformed_input(self):
|
|
self.assertEqual(positive_prompt_nodes({"1": pm()}), [])
|
|
for bad in (None, [], "x", {"1": None}, {"1": {"inputs": None}}):
|
|
self.assertEqual(positive_prompt_nodes(bad), [], repr(bad))
|
|
|
|
def test_cycles_terminate(self):
|
|
loop = {
|
|
"class_type": "Weird",
|
|
"inputs": {
|
|
"conditioning": ["9", 0],
|
|
"conditioning_1": ["9", 0],
|
|
"positive": ["9", 0],
|
|
},
|
|
}
|
|
graph = {"9": loop, "5": ksampler("9", "9"), "20": LOADER}
|
|
self.assertEqual(positive_prompt_nodes(graph), [])
|
|
|
|
|
|
class TestRunPromptNodes(unittest.TestCase):
|
|
"""The prompt nodes that represent a run, with a safe fallback."""
|
|
|
|
def test_uses_positive_nodes_when_found(self):
|
|
graph = {
|
|
"176": pm("blurry"),
|
|
"134": pm("a cat"),
|
|
"5": ksampler("134", "176"),
|
|
"20": LOADER,
|
|
}
|
|
self.assertEqual(run_prompt_nodes(graph), ["134"])
|
|
|
|
def test_unrecognised_sampler_falls_back_to_every_prompt_manager_node(self):
|
|
graph = {
|
|
"1": pm("a cat"),
|
|
"9": {"class_type": "ExoticSampler", "inputs": {"cond": ["1", 0]}},
|
|
"2": {"class_type": "PromptManagerText", "inputs": {"text": "x"}},
|
|
}
|
|
self.assertEqual(sorted(run_prompt_nodes(graph)), ["1", "2"])
|
|
|
|
def test_malformed_graph(self):
|
|
self.assertEqual(run_prompt_nodes(None), [])
|
|
|
|
def test_run_prompt_nodes_negative_only_graph_is_empty(self):
|
|
graph = {
|
|
"1": {"class_type": "PromptManager", "inputs": {"text": "bad hands"}},
|
|
"2": {
|
|
"class_type": "KSampler",
|
|
"inputs": {"negative": ["1", 0], "positive": ["9", 0]},
|
|
},
|
|
"9": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}},
|
|
}
|
|
self.assertEqual(run_prompt_nodes(graph), [])
|
|
|
|
def test_run_prompt_nodes_falls_back_when_no_sampler(self):
|
|
graph = {"1": {"class_type": "PromptManagerText", "inputs": {"text": "x"}}}
|
|
self.assertEqual(run_prompt_nodes(graph), ["1"])
|
|
|
|
|
|
class TestInputKeys(unittest.TestCase):
|
|
"""Which input names are followed when tracing the positive prompt upstream."""
|
|
|
|
def test_negative_flavoured_keys_are_skipped(self):
|
|
for key in ("negative_prompt", "text_negative", "neg_text", "neg", "negative"):
|
|
self.assertFalse(follows_input(key), key)
|
|
self.assertEqual(_upstream_links({key: ["1", 0]}), [], key)
|
|
|
|
def test_positive_and_plain_text_keys_are_followed(self):
|
|
for key in ("text_positive", "positive", "text", "prompt", "text_g"):
|
|
self.assertTrue(follows_input(key), key)
|
|
self.assertEqual(_upstream_links({key: ["1", 0]}), [["1", 0]], key)
|
|
|
|
def test_styler_node_only_follows_its_positive_text(self):
|
|
graph = {
|
|
"1": pm("a cat"),
|
|
"2": pm("bad hands"),
|
|
"3": {
|
|
"class_type": "SDXLPromptStyler",
|
|
"inputs": {"text_positive": ["1", 0], "text_negative": ["2", 0]},
|
|
},
|
|
"4": {"class_type": "CLIPTextEncode", "inputs": {"text": ["3", 0]}},
|
|
"5": ksampler("4", "2"),
|
|
"20": LOADER,
|
|
}
|
|
self.assertEqual(positive_prompt_nodes(graph), ["1"])
|
|
self.assertEqual(run_prompt_nodes(graph), ["1"])
|
|
self.assertEqual(resolve_text(graph, "1"), "a cat")
|
|
|
|
|
|
class TestHasSampler(unittest.TestCase):
|
|
def test_structural_sampler_and_guider_are_detected(self):
|
|
self.assertTrue(has_sampler({"5": ksampler("1", "2"), "20": LOADER}))
|
|
guider = {
|
|
"class_type": "BasicGuider",
|
|
"inputs": {"model": ["20", 0], "conditioning": ["1", 0]},
|
|
}
|
|
self.assertTrue(has_sampler({"6": guider, "20": LOADER}))
|
|
|
|
def test_sampler_without_model_link_is_detected_by_name(self):
|
|
sampler = {
|
|
"class_type": "KSamplerAdvanced",
|
|
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
|
|
}
|
|
self.assertTrue(has_sampler({"5": sampler}))
|
|
|
|
def test_non_samplers_and_malformed_graphs(self):
|
|
controlnet = {
|
|
"class_type": "ControlNetApplyAdvanced",
|
|
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
|
|
}
|
|
self.assertFalse(has_sampler({"7": controlnet}))
|
|
exotic = {"class_type": "ExoticSampler", "inputs": {"cond": ["1", 0]}}
|
|
self.assertFalse(has_sampler({"9": exotic}))
|
|
for bad in (None, [], "x", {"1": None}, {"1": {"inputs": None}}, {}):
|
|
self.assertFalse(has_sampler(bad), repr(bad))
|
|
|
|
|
|
class TestAdversarialGraphs(unittest.TestCase):
|
|
"""Queued and embedded graphs are untrusted: never hang, never raise."""
|
|
|
|
def _assert_fast(self, graph):
|
|
started = time.monotonic()
|
|
for fn in (positive_prompt_nodes, run_prompt_nodes, has_sampler):
|
|
fn(graph)
|
|
resolve_text(graph, "1")
|
|
self.assertLess(time.monotonic() - started, 1.0)
|
|
|
|
def test_cyclic_graph_terminates(self):
|
|
graph = {
|
|
"1": pm(["2", 0]),
|
|
"2": {"class_type": "StringConcatenate", "inputs": {"string_a": ["1", 0]}},
|
|
"3": {
|
|
"class_type": "ConditioningCombine",
|
|
"inputs": {"conditioning_1": ["3", 0], "conditioning_2": ["1", 0]},
|
|
},
|
|
"5": ksampler("3", "3"),
|
|
"20": LOADER,
|
|
}
|
|
self._assert_fast(graph)
|
|
self.assertEqual(run_prompt_nodes(graph), ["1"])
|
|
|
|
def test_ten_thousand_node_fan_out_terminates(self):
|
|
graph = {"20": LOADER}
|
|
combine_inputs = {}
|
|
for i in range(10_000):
|
|
graph[str(100 + i)] = pm(f"prompt {i}")
|
|
combine_inputs[f"conditioning_{i}"] = [str(100 + i), 0]
|
|
graph["3"] = {"class_type": "ConditioningCombine", "inputs": combine_inputs}
|
|
graph["5"] = ksampler("3", "3")
|
|
self._assert_fast(graph)
|
|
self.assertTrue(run_prompt_nodes(graph))
|
|
|
|
def test_non_dict_inputs_and_odd_values(self):
|
|
graph = {
|
|
"1": {"class_type": "PromptManager", "inputs": ["not", "a", "dict"]},
|
|
"2": {"class_type": "KSampler", "inputs": 42},
|
|
"3": {"class_type": None, "inputs": {"positive": "1", "negative": [1]}},
|
|
"4": "just a string",
|
|
"5": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"positive": [None, 0],
|
|
"negative": [{"x": 1}, 0],
|
|
"model": [[], 0],
|
|
},
|
|
},
|
|
6: {"class_type": "PromptManager", "inputs": {"text": 7}},
|
|
}
|
|
self._assert_fast(graph)
|
|
for fn in (positive_prompt_nodes, run_prompt_nodes):
|
|
self.assertIsInstance(fn(graph), list)
|
|
|
|
|
|
class TestResolveText(unittest.TestCase):
|
|
"""Text known at queue time: typed, or from known pure string nodes."""
|
|
|
|
def _linked(self, source):
|
|
return {"1": pm(["2", 0]), "2": source}
|
|
|
|
def test_typed_text(self):
|
|
self.assertEqual(resolve_text({"1": pm(" a cat ")}, "1"), "a cat")
|
|
|
|
def test_core_primitive_string(self):
|
|
for cls in ("PrimitiveString", "PrimitiveStringMultiline"):
|
|
graph = self._linked({"class_type": cls, "inputs": {"value": "a cat"}})
|
|
self.assertEqual(resolve_text(graph, "1"), "a cat", cls)
|
|
|
|
def test_core_string_concatenate(self):
|
|
graph = self._linked(
|
|
{
|
|
"class_type": "StringConcatenate",
|
|
"inputs": {
|
|
"string_a": "a cat",
|
|
"string_b": "at dusk",
|
|
"delimiter": ", ",
|
|
},
|
|
}
|
|
)
|
|
self.assertEqual(resolve_text(graph, "1"), "a cat, at dusk")
|
|
|
|
def test_was_text_concatenate_skips_empty_and_cleans_whitespace(self):
|
|
graph = self._linked(
|
|
{
|
|
"class_type": "Text Concatenate",
|
|
"inputs": {
|
|
"delimiter": " ",
|
|
"clean_whitespace": "true",
|
|
"text_a": " a cat ",
|
|
"text_b": "",
|
|
"text_c": ["3", 0],
|
|
},
|
|
}
|
|
)
|
|
graph["3"] = {"class_type": "PrimitiveString", "inputs": {"value": "at dusk"}}
|
|
self.assertEqual(resolve_text(graph, "1"), "a cat at dusk")
|
|
|
|
def test_pass_through_nodes(self):
|
|
for cls in ("ShowText|pysssss", "Text Multiline"):
|
|
graph = self._linked({"class_type": cls, "inputs": {"text": "a cat"}})
|
|
self.assertEqual(resolve_text(graph, "1"), "a cat", cls)
|
|
|
|
def test_unknown_and_batch_sources_are_never_guessed(self):
|
|
for source in (
|
|
{"class_type": "PromptSearchList", "inputs": {"search_text": "cat"}},
|
|
{
|
|
"class_type": "SomeTextReplace",
|
|
"inputs": {"text": "a cat", "find": "cat", "replace": "dog"},
|
|
},
|
|
):
|
|
self.assertIsNone(
|
|
resolve_text(self._linked(source), "1"), source["class_type"]
|
|
)
|
|
|
|
def test_same_source_used_twice_resolves_both_branches(self):
|
|
# StringConcatenate(string_a=X, string_b=X): X is not a cycle, it is
|
|
# simply shared, and must resolve in both branches.
|
|
graph = {
|
|
"1": pm(["2", 0]),
|
|
"2": {
|
|
"class_type": "StringConcatenate",
|
|
"inputs": {
|
|
"string_a": ["3", 0],
|
|
"string_b": ["3", 0],
|
|
"delimiter": "+",
|
|
},
|
|
},
|
|
"3": {"class_type": "PrimitiveString", "inputs": {"value": "x"}},
|
|
}
|
|
self.assertEqual(resolve_text(graph, "1"), "x+x")
|
|
|
|
def test_diamond_graph_resolves(self):
|
|
# 1 <- 2 <- (3, 4) and both 3 and 4 read from 5
|
|
graph = {
|
|
"1": pm(["2", 0]),
|
|
"2": {
|
|
"class_type": "StringConcatenate",
|
|
"inputs": {
|
|
"string_a": ["3", 0],
|
|
"string_b": ["4", 0],
|
|
"delimiter": " ",
|
|
},
|
|
},
|
|
"3": {"class_type": "ShowText|pysssss", "inputs": {"text": ["5", 0]}},
|
|
"4": {"class_type": "Text Multiline", "inputs": {"text": ["5", 0]}},
|
|
"5": {"class_type": "PrimitiveString", "inputs": {"value": "cat"}},
|
|
}
|
|
self.assertEqual(resolve_text(graph, "1"), "cat cat")
|
|
|
|
def test_indirect_cycle_returns_none(self):
|
|
graph = {
|
|
"1": pm(["2", 0]),
|
|
"2": {"class_type": "ShowText|pysssss", "inputs": {"text": ["3", 0]}},
|
|
"3": {"class_type": "Text Multiline", "inputs": {"text": ["2", 0]}},
|
|
}
|
|
self.assertIsNone(resolve_text(graph, "1"))
|
|
|
|
def test_visit_budget_bounds_shared_fan_out(self):
|
|
# A chain of concatenations that each read the previous node twice
|
|
# expands exponentially without memoisation; MAX_VISITS must cap it.
|
|
graph = {"0": {"class_type": "PrimitiveString", "inputs": {"value": "a"}}}
|
|
for i in range(1, 400):
|
|
prev = [str(i - 1), 0]
|
|
graph[str(i)] = {
|
|
"class_type": "StringConcatenate",
|
|
"inputs": {"string_a": prev, "string_b": prev, "delimiter": ""},
|
|
}
|
|
graph["pm"] = pm(["399", 0])
|
|
start = time.monotonic()
|
|
self.assertIsNone(resolve_text(graph, "pm"))
|
|
self.assertLess(time.monotonic() - start, 2.0)
|
|
|
|
def test_cycles_and_missing_links(self):
|
|
graph = {
|
|
"1": pm(["2", 0]),
|
|
"2": {
|
|
"class_type": "StringConcatenate",
|
|
"inputs": {"string_a": ["2", 0], "string_b": "x", "delimiter": ""},
|
|
},
|
|
}
|
|
self.assertIsNone(resolve_text(graph, "1"))
|
|
self.assertIsNone(resolve_text({"1": pm(["9", 0])}, "1"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|