509 lines
20 KiB
Python
509 lines
20 KiB
Python
"""#90 item 3: a node's ``save_to_library`` switch.
|
|
|
|
Off: PromptManager ignores the run completely (no prompt saved, no run count,
|
|
no image link), while the node still returns its normal outputs. On (the
|
|
default, and every workflow or image made before the switch existed): as before.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
import unittest.mock as mock
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
try:
|
|
from tests.test_prompt_attribution import MonitorCase, NodeCase
|
|
from tests.test_execution_contexts import load_module
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from test_prompt_attribution import MonitorCase, NodeCase
|
|
from test_execution_contexts import load_module
|
|
|
|
from prompt_manager import PromptManager, PromptManagerFocused, PromptManagerSplit
|
|
from prompt_manager_pair import PromptManagerPair
|
|
from prompt_manager_text import PromptManagerText
|
|
from prompt_search_list import PromptSearchList
|
|
from utils.hashing import generate_prompt_hash
|
|
from utils.prompt_graph import (
|
|
intrinsic_pair_nodes,
|
|
positive_opt_outs,
|
|
run_prompt_components,
|
|
run_prompt_nodes,
|
|
save_setting,
|
|
)
|
|
from utils.run_keys import prompt_node_key
|
|
from utils.usage_tracking import (
|
|
PendingFirstUse,
|
|
QueuedRuns,
|
|
handle_queued_prompt,
|
|
on_dequeue,
|
|
queued_pair_runs,
|
|
queued_prompt_runs,
|
|
)
|
|
|
|
POSITIVE = "a lighthouse at dusk"
|
|
NEGATIVE = "blurry"
|
|
NODE_CLASSES = (
|
|
PromptManager,
|
|
PromptManagerText,
|
|
PromptManagerFocused,
|
|
PromptManagerSplit,
|
|
PromptManagerPair,
|
|
)
|
|
|
|
|
|
def run_graph(save=None, kind="PromptManager", negative_save=None):
|
|
"""A PromptManager positive feeding a KSampler; ``save`` None leaves it out."""
|
|
inputs = {"text": POSITIVE, "clip": ["4", 1]}
|
|
if save is not None:
|
|
inputs["save_to_library"] = save
|
|
negative_inputs = {"text": NEGATIVE, "clip": ["4", 1]}
|
|
if negative_save is not None:
|
|
negative_inputs["save_to_library"] = negative_save
|
|
return {
|
|
"1": {"class_type": kind, "inputs": inputs},
|
|
"2": {"class_type": "PromptManager", "inputs": negative_inputs},
|
|
"4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "m"}},
|
|
"3": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"model": ["4", 0],
|
|
"positive": ["1", 0],
|
|
"negative": ["2", 0],
|
|
"latent_image": ["6", 0],
|
|
},
|
|
},
|
|
"6": {"class_type": "EmptyLatentImage", "inputs": {"batch_size": 1}},
|
|
}
|
|
|
|
|
|
def pair_graph(save=None):
|
|
inputs = {"text": POSITIVE, "negative_text": NEGATIVE, "clip": ["4", 1]}
|
|
if save is not None:
|
|
inputs["save_to_library"] = save
|
|
return {
|
|
"1": {"class_type": "PromptManagerPair", "inputs": inputs},
|
|
"4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "m"}},
|
|
"3": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"model": ["4", 0],
|
|
"positive": ["1", 0],
|
|
"negative": ["1", 1],
|
|
"latent_image": ["6", 0],
|
|
},
|
|
},
|
|
"6": {"class_type": "EmptyLatentImage", "inputs": {"batch_size": 1}},
|
|
}
|
|
|
|
|
|
class TestInputDeclaration(unittest.TestCase):
|
|
def test_every_prompt_manager_node_declares_a_boolean_defaulting_on(self):
|
|
for cls in NODE_CLASSES:
|
|
optional = cls.INPUT_TYPES()["optional"]
|
|
kind, options = optional["save_to_library"]
|
|
self.assertEqual(kind, "BOOLEAN", cls.__name__)
|
|
self.assertIs(options["default"], True, cls.__name__)
|
|
# Declared last: no saved widget value of an older workflow shifts
|
|
self.assertEqual(list(optional)[-1], "save_to_library", cls.__name__)
|
|
|
|
def test_search_list_has_no_switch(self):
|
|
inputs = PromptSearchList.INPUT_TYPES()
|
|
self.assertNotIn("save_to_library", inputs.get("optional", {}))
|
|
self.assertNotIn("save_to_library", inputs.get("required", {}))
|
|
|
|
|
|
class TestSaveSetting(unittest.TestCase):
|
|
def test_missing_input_is_on(self):
|
|
self.assertIs(save_setting(run_graph(), "1"), True)
|
|
|
|
def test_literals_follow_comfyui_boolean_coercion(self):
|
|
for value, expected in (
|
|
(True, True),
|
|
(False, False),
|
|
("", False),
|
|
(0, False),
|
|
({"__value__": False}, False),
|
|
({"__value__": True}, True),
|
|
):
|
|
with self.subTest(value=value):
|
|
self.assertIs(save_setting(run_graph(save=value), "1"), expected)
|
|
|
|
def test_link_to_primitive_boolean_is_read(self):
|
|
for value in (True, False):
|
|
graph = run_graph(save=["9", 0])
|
|
graph["9"] = {"class_type": "PrimitiveBoolean", "inputs": {"value": value}}
|
|
self.assertIs(save_setting(graph, "1"), value)
|
|
|
|
def test_link_to_an_unknown_node_is_not_knowable(self):
|
|
graph = run_graph(save=["9", 0])
|
|
graph["9"] = {"class_type": "SomeCustomSwitch", "inputs": {}}
|
|
self.assertIsNone(save_setting(graph, "1"))
|
|
self.assertIsNone(save_setting(run_graph(save=["missing", 0]), "1"))
|
|
|
|
|
|
class TestGraphRoles(unittest.TestCase):
|
|
def test_opted_out_positive_is_not_a_run_prompt(self):
|
|
graph = run_graph(save=False)
|
|
self.assertEqual(run_prompt_nodes(graph), [])
|
|
self.assertEqual(run_prompt_components(graph), [])
|
|
self.assertEqual(positive_opt_outs(graph), ["1"])
|
|
|
|
def test_enabled_positive_is_unchanged(self):
|
|
for save in (None, True):
|
|
graph = run_graph(save=save)
|
|
self.assertEqual(run_prompt_nodes(graph), ["1"])
|
|
self.assertEqual(positive_opt_outs(graph), [])
|
|
|
|
def test_opted_out_negative_is_no_positive_opt_out(self):
|
|
graph = run_graph(negative_save=False)
|
|
self.assertEqual(run_prompt_nodes(graph), ["1"])
|
|
self.assertEqual(positive_opt_outs(graph), [])
|
|
|
|
def test_walk_stops_at_an_opted_out_node(self):
|
|
graph = run_graph(save=False)
|
|
graph["1"]["inputs"]["text"] = ["7", 0]
|
|
graph["7"] = {"class_type": "PromptManagerText", "inputs": {"text": "x"}}
|
|
self.assertEqual(run_prompt_nodes(graph), [])
|
|
|
|
def test_sampler_less_fallback_skips_opted_out_nodes(self):
|
|
graph = {
|
|
"1": {
|
|
"class_type": "PromptManagerText",
|
|
"inputs": {"text": "a", "save_to_library": False},
|
|
},
|
|
"2": {"class_type": "PromptManagerText", "inputs": {"text": "b"}},
|
|
}
|
|
self.assertEqual(run_prompt_nodes(graph), ["2"])
|
|
self.assertEqual(run_prompt_components(graph), [("2", "text")])
|
|
self.assertEqual(positive_opt_outs(graph), ["1"])
|
|
|
|
def test_opted_out_pair_is_not_an_intrinsic_pair(self):
|
|
self.assertEqual(intrinsic_pair_nodes(pair_graph()), ["1"])
|
|
self.assertEqual(intrinsic_pair_nodes(pair_graph(save=False)), [])
|
|
|
|
|
|
class TestQueueHook(unittest.TestCase):
|
|
def test_opted_out_node_is_never_stashed(self):
|
|
self.assertEqual(queued_prompt_runs(run_graph(save=False)), [])
|
|
self.assertEqual(queued_pair_runs(pair_graph(save=False)), [])
|
|
|
|
def test_unknown_setting_is_left_to_the_node(self):
|
|
graph = run_graph(save=["9", 0])
|
|
graph["9"] = {"class_type": "SomeCustomSwitch", "inputs": {}}
|
|
self.assertEqual(queued_prompt_runs(graph), [])
|
|
pair = pair_graph(save=["9", 0])
|
|
pair["9"] = {"class_type": "SomeCustomSwitch", "inputs": {}}
|
|
self.assertEqual(queued_pair_runs(pair), [])
|
|
|
|
def test_primitive_boolean_on_still_counts(self):
|
|
graph = run_graph(save=["9", 0])
|
|
graph["9"] = {"class_type": "PrimitiveBoolean", "inputs": {"value": True}}
|
|
self.assertEqual(
|
|
queued_prompt_runs(graph), [(generate_prompt_hash(POSITIVE), POSITIVE)]
|
|
)
|
|
|
|
def test_dequeue_counts_nothing_for_a_cached_opted_out_node(self):
|
|
"""A cached node never executes: only the queued graph can say it is off."""
|
|
db = mock.Mock()
|
|
db.get_prompt_by_hash.return_value = {"id": 5}
|
|
queued, pending = QueuedRuns(), PendingFirstUse()
|
|
request = {"prompt": run_graph(save=False), "prompt_id": "job-1"}
|
|
handle_queued_prompt(request, queued, pair_queued=QueuedRuns())
|
|
on_dequeue("job-1", queued, lambda: db, pending, pair_queued=QueuedRuns())
|
|
db.record_prompt_use.assert_not_called()
|
|
self.assertEqual(pending.consume(generate_prompt_hash(POSITIVE)), 0)
|
|
|
|
|
|
class TestNodesSaveNothing(NodeCase):
|
|
def run_node(self, cls, graph, save):
|
|
node = self.node(cls)
|
|
node.record_pair = mock.Mock()
|
|
kwargs = {"prompt": graph, "unique_id": "1", "save_to_library": save}
|
|
if cls is PromptManagerText:
|
|
return node, node.process_text(POSITIVE, **kwargs)
|
|
if cls is PromptManagerPair:
|
|
return node, node.encode_prompt(
|
|
self.clip(), POSITIVE, negative_text=NEGATIVE, **kwargs
|
|
)
|
|
return node, node.encode_prompt(self.clip(), POSITIVE, **kwargs)
|
|
|
|
def test_off_saves_counts_and_tracks_nothing_but_returns_outputs(self):
|
|
for cls in NODE_CLASSES:
|
|
with self.subTest(node=cls.__name__):
|
|
graph = (
|
|
pair_graph(save=False)
|
|
if cls is PromptManagerPair
|
|
else run_graph(save=False, kind=cls.__name__)
|
|
)
|
|
node, result = self.run_node(cls, graph, False)
|
|
self.assertEqual(self.count_rows("prompts"), 0)
|
|
self.assertEqual(self.count_rows("prompt_pairs"), 0)
|
|
self.assertEqual(self.tracker.keyed_snapshot(), {})
|
|
self.assertEqual(self.tracker._prompt_queue, [])
|
|
node.comfyui_integration.register_prompt.assert_not_called()
|
|
if cls is PromptManagerText:
|
|
self.assertEqual(result, (POSITIVE,))
|
|
elif cls is PromptManagerPair:
|
|
self.assertEqual(result[0], ("cond", POSITIVE))
|
|
self.assertEqual(result[1], ("cond", NEGATIVE))
|
|
self.assertEqual(result[2:], (POSITIVE, NEGATIVE))
|
|
else:
|
|
self.assertEqual(result, (("cond", POSITIVE), POSITIVE))
|
|
|
|
def test_off_without_graph_context_still_saves_nothing(self):
|
|
node = self.node(PromptManager)
|
|
node.encode_prompt(self.clip(), POSITIVE, save_to_library=False)
|
|
self.assertEqual(self.count_rows("prompts"), 0)
|
|
self.assertEqual(self.tracker._prompt_queue, [])
|
|
|
|
def test_off_leaves_an_existing_prompt_uncounted(self):
|
|
prompt_id = self.save(POSITIVE)
|
|
self.run_node(PromptManager, run_graph(save=False), False)
|
|
row = self.db.get_prompt_by_id(prompt_id)
|
|
self.assertFalse(row.get("run_count") or 0)
|
|
|
|
def test_on_is_unchanged(self):
|
|
self.run_node(PromptManager, run_graph(save=True), True)
|
|
self.assertIsNotNone(self.db.get_prompt_by_hash(generate_prompt_hash(POSITIVE)))
|
|
self.assertIn(
|
|
prompt_node_key(run_graph(save=True), "1"), self.tracker.keyed_snapshot()
|
|
)
|
|
|
|
def test_unknown_setting_on_at_run_time_counts_at_execution(self):
|
|
"""The hook could not read a wired switch, so the node counts its run."""
|
|
graph = run_graph(save=["9", 0])
|
|
graph["9"] = {"class_type": "SomeCustomSwitch", "inputs": {}}
|
|
prompt_id = self.save(POSITIVE)
|
|
self.run_node(PromptManager, graph, True)
|
|
self.assertEqual(self.db.get_prompt_by_id(prompt_id)["run_count"], 1)
|
|
|
|
|
|
class TestWiredSwitchOffMarksItsRun(NodeCase):
|
|
"""Review finding: a wired switch off at run time keeps the node a run prompt."""
|
|
|
|
def wired_graph(self):
|
|
graph = run_graph(save=["9", 0])
|
|
graph["9"] = {"class_type": "SomeCustomSwitch", "inputs": {}}
|
|
return graph
|
|
|
|
def test_an_earlier_entry_of_the_key_never_stands_in(self):
|
|
graph = self.wired_graph()
|
|
key = prompt_node_key(graph, "1")
|
|
node = self.node(PromptManager)
|
|
node.encode_prompt(
|
|
self.clip(), POSITIVE, prompt=graph, unique_id="1", save_to_library=True
|
|
)
|
|
node.encode_prompt(
|
|
self.clip(), POSITIVE, prompt=graph, unique_id="1", save_to_library=False
|
|
)
|
|
latest, _run, _count = self.tracker.keyed_snapshot()[key]
|
|
self.assertTrue(latest.get("placeholder"))
|
|
|
|
def test_literal_off_marks_nothing(self):
|
|
node = self.node(PromptManager)
|
|
node.encode_prompt(
|
|
self.clip(),
|
|
POSITIVE,
|
|
prompt=run_graph(save=False),
|
|
unique_id="1",
|
|
save_to_library=False,
|
|
)
|
|
self.assertEqual(self.tracker.keyed_snapshot(), {})
|
|
|
|
|
|
class TestMonitorLeavesOptedOutFilesUnlinked(MonitorCase):
|
|
def test_existing_prompt_is_not_linked_from_metadata(self):
|
|
self.save(POSITIVE)
|
|
self.assertIsNone(self.process({"prompt": run_graph(save=False)}))
|
|
|
|
def test_find_prompt_reports_opted_out(self):
|
|
self.save(POSITIVE)
|
|
state, prompt = self.handler._find_prompt_from_metadata(
|
|
{"prompt": run_graph(save=False)}
|
|
)
|
|
self.assertEqual((state, prompt), ("opted_out", None))
|
|
|
|
def test_mixed_positive_settings_leave_the_file_unlinked(self):
|
|
self.save(POSITIVE)
|
|
graph = run_graph(save=True)
|
|
graph["8"] = {
|
|
"class_type": "PromptManager",
|
|
"inputs": {"text": "other", "save_to_library": False},
|
|
}
|
|
graph["5"] = {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"model": ["4", 0],
|
|
"positive": ["8", 0],
|
|
"negative": ["2", 0],
|
|
"latent_image": ["6", 0],
|
|
},
|
|
}
|
|
self.assertIsNone(self.process({"prompt": graph}))
|
|
|
|
def test_workflow_only_metadata_with_switch_off_is_opted_out(self):
|
|
self.save(POSITIVE)
|
|
metadata = {
|
|
"text_encoder_nodes": [
|
|
{
|
|
"type": "PromptManager",
|
|
"inputs": [],
|
|
"widgets_values": [POSITIVE, "", "", "", "", "", False],
|
|
}
|
|
]
|
|
}
|
|
state, _prompt = self.handler._find_prompt_from_metadata(metadata)
|
|
self.assertEqual(state, "opted_out")
|
|
|
|
def test_workflow_only_ignores_muted_nodes_and_reads_subgraphs(self):
|
|
from utils.prompt_graph import metadata_opted_out
|
|
|
|
node = {
|
|
"type": "PromptManager",
|
|
"widgets_values": [POSITIVE, "", "", "", "", "", False],
|
|
}
|
|
self.assertFalse(
|
|
metadata_opted_out({"workflow": {"nodes": [{**node, "mode": 4}]}})
|
|
)
|
|
self.assertTrue(
|
|
metadata_opted_out(
|
|
{
|
|
"workflow": {
|
|
"nodes": [],
|
|
"definitions": {"subgraphs": [{"nodes": [node]}]},
|
|
}
|
|
}
|
|
)
|
|
)
|
|
|
|
def test_unknown_setting_is_left_to_the_run(self):
|
|
self.save(POSITIVE)
|
|
graph = run_graph(save=["9", 0])
|
|
graph["9"] = {"class_type": "SomeCustomSwitch", "inputs": {}}
|
|
state, _prompt = self.handler._find_prompt_from_metadata({"prompt": graph})
|
|
self.assertEqual(state, "unknown")
|
|
|
|
def test_enabled_positive_still_links(self):
|
|
prompt_id = self.save(POSITIVE)
|
|
self.assertEqual(self.process({"prompt": run_graph(save=True)}), prompt_id)
|
|
|
|
|
|
class TestLinkKeepsOptedOutNegativeOut(MonitorCase):
|
|
def setUp(self):
|
|
super().setUp()
|
|
# The real link, not the recording stub of MonitorCase
|
|
del self.handler.link_image_to_prompt
|
|
|
|
def test_opted_out_negative_is_never_stored_as_a_pair(self):
|
|
prompt_id = self.save(POSITIVE)
|
|
metadata = {"prompt": run_graph(save=True, negative_save=False)}
|
|
self.handler.link_image_to_prompt(self.image, {"id": prompt_id}, metadata)
|
|
self.assertIsNone(self.db.get_prompt_by_hash(generate_prompt_hash(NEGATIVE)))
|
|
self.assertEqual(self.count_pairs(), 0)
|
|
linked = (
|
|
self.db.model.get_connection()
|
|
.execute("SELECT prompt_id FROM generated_images")
|
|
.fetchall()
|
|
)
|
|
self.assertEqual([row[0] for row in linked], [prompt_id])
|
|
|
|
def test_opted_out_positive_links_nothing(self):
|
|
prompt_id = self.save(POSITIVE)
|
|
metadata = {"prompt": run_graph(save=False)}
|
|
self.handler.link_image_to_prompt(self.image, {"id": prompt_id}, metadata)
|
|
count = (
|
|
self.db.model.get_connection()
|
|
.execute("SELECT COUNT(*) FROM generated_images")
|
|
.fetchone()[0]
|
|
)
|
|
self.assertEqual(count, 0)
|
|
|
|
def count_pairs(self):
|
|
conn = self.db.model.get_connection()
|
|
return conn.execute("SELECT COUNT(*) FROM prompt_pairs").fetchone()[0]
|
|
|
|
|
|
class TestScanSkipsOptedOutFiles(NodeCase):
|
|
def setUp(self):
|
|
super().setUp()
|
|
from py.api import admin
|
|
|
|
self.mixin = admin.AdminRoutesMixin.__new__(admin.AdminRoutesMixin)
|
|
self.mixin.db = self.db
|
|
self.mixin.logger = logging.getLogger("test.save_to_library.scan")
|
|
self.file = Path(self.tmp.name) / "out.png"
|
|
|
|
def test_opted_out_file_stores_nothing(self):
|
|
status = self.mixin._ingest_scanned_file(
|
|
self.file, {"prompt": run_graph(save=False)}
|
|
)
|
|
self.assertEqual(status, "not_saved")
|
|
self.assertEqual(self.count_rows("prompts"), 0)
|
|
self.assertEqual(self.count_rows("generated_images"), 0)
|
|
|
|
def test_opted_out_file_is_never_queued_for_review(self):
|
|
ambiguous = []
|
|
graph = run_graph(save=False)
|
|
graph["5"] = dict(graph["3"]) # a second sampler: ambiguous otherwise
|
|
graph["5"]["inputs"] = {**graph["3"]["inputs"], "positive": ["2", 0]}
|
|
status = self.mixin._ingest_scanned_file(
|
|
self.file, {"prompt": graph}, ambiguous=ambiguous
|
|
)
|
|
self.assertEqual(status, "not_saved")
|
|
self.assertEqual(ambiguous, [])
|
|
|
|
def test_opted_out_negative_stores_the_positive_only(self):
|
|
status = self.mixin._ingest_scanned_file(
|
|
self.file, {"prompt": run_graph(save=True, negative_save=False)}
|
|
)
|
|
self.assertEqual(status, "added")
|
|
self.assertIsNotNone(self.db.get_prompt_by_hash(generate_prompt_hash(POSITIVE)))
|
|
self.assertIsNone(self.db.get_prompt_by_hash(generate_prompt_hash(NEGATIVE)))
|
|
self.assertEqual(self.count_rows("prompt_pairs"), 0)
|
|
|
|
def test_wired_switch_is_left_to_the_live_run(self):
|
|
graph = run_graph(save=["9", 0])
|
|
graph["9"] = {"class_type": "SomeCustomSwitch", "inputs": {}}
|
|
status = self.mixin._ingest_scanned_file(self.file, {"prompt": graph})
|
|
self.assertEqual(status, "not_saved_wired")
|
|
self.assertEqual(self.count_rows("prompts"), 0)
|
|
|
|
def test_status_has_a_scan_detail(self):
|
|
from py.api import admin
|
|
|
|
group, row, _line = admin._scan_outcome(self.file, [], "not_saved")
|
|
self.assertEqual(group, "no_metadata")
|
|
self.assertIn("does not save", row["detail"])
|
|
|
|
|
|
class TestHistoryRecordsNoText(unittest.TestCase):
|
|
def setUp(self):
|
|
self.m = load_module("execution_context_graph")
|
|
|
|
def test_opted_out_role_has_no_sources_and_no_text(self):
|
|
recipe = self.m.project_contexts(run_graph(save=False), None).recipes[0]
|
|
self.assertEqual(recipe.positive.sources, ())
|
|
self.assertIsNone(recipe.positive.inferred)
|
|
self.assertEqual(recipe.positive.reason, "not_saved")
|
|
# The enabled negative is recorded as before
|
|
self.assertEqual(recipe.negative.inferred.raw.value, NEGATIVE)
|
|
|
|
def test_unknown_setting_records_no_text_either(self):
|
|
graph = run_graph(save=["9", 0])
|
|
graph["9"] = {"class_type": "SomeCustomSwitch", "inputs": {}}
|
|
recipe = self.m.project_contexts(graph, None).recipes[0]
|
|
self.assertIsNone(recipe.positive.inferred)
|
|
self.assertEqual(recipe.positive.reason, "not_saved")
|
|
|
|
def test_enabled_role_is_unchanged(self):
|
|
recipe = self.m.project_contexts(run_graph(save=True), None).recipes[0]
|
|
self.assertEqual(recipe.positive.inferred.raw.value, POSITIVE)
|
|
self.assertEqual(len(recipe.positive.sources), 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|