fix: never treat negative-only PromptManager nodes as the run prompt
run_prompt_nodes fell back to every PromptManager node whenever no positive node was traced, including graphs that do have a sampler and whose only PromptManager node feeds its negative input. Fall back only when the graph has no recognisable sampler (new has_sampler helper). Samplers named like one are recognised even without a model link. Adversarial graphs (cycles, 10k-node fan-out, non-dict inputs) are covered.
This commit is contained in:
Vito Sansevero
@@ -2,11 +2,13 @@
|
||||
|
||||
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 (
|
||||
has_sampler,
|
||||
positive_prompt_nodes,
|
||||
resolve_text,
|
||||
run_prompt_nodes,
|
||||
@@ -167,6 +169,105 @@ class TestRunPromptNodes(unittest.TestCase):
|
||||
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 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."""
|
||||
|
||||
+27
-13
@@ -36,19 +36,26 @@ def _inputs(node: Any) -> Dict[str, Any]:
|
||||
|
||||
|
||||
def _positive_roots(graph: Dict[str, Any]) -> List[Any]:
|
||||
"""Links into samplers' positive inputs (KSampler, CFGGuider, BasicGuider, ...)."""
|
||||
"""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():
|
||||
inputs = _inputs(node)
|
||||
if not _is_link(inputs.get("model")):
|
||||
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")):
|
||||
roots.append(inputs["positive"])
|
||||
elif (
|
||||
_is_link(inputs.get("conditioning"))
|
||||
and "guider" in str(node.get("class_type", "")).lower()
|
||||
):
|
||||
roots.append(inputs["conditioning"])
|
||||
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
|
||||
|
||||
|
||||
@@ -63,6 +70,11 @@ def _upstream_links(inputs: Dict[str, Any]) -> List[Any]:
|
||||
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.
|
||||
|
||||
@@ -92,12 +104,14 @@ def positive_prompt_nodes(graph: Any) -> List[str]:
|
||||
def run_prompt_nodes(graph: Any) -> List[str]:
|
||||
"""Prompt nodes that represent a run: the positive ones, if any can be found.
|
||||
|
||||
When no PromptManager node can be traced to a positive input (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 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):
|
||||
if positive or not isinstance(graph, dict) or has_sampler(graph):
|
||||
return positive
|
||||
return [
|
||||
str(node_id)
|
||||
|
||||
Reference in New Issue
Block a user