diff --git a/__init__.py b/__init__.py index f05e973..c1e1dd5 100644 --- a/__init__.py +++ b/__init__.py @@ -45,6 +45,7 @@ from .prompt_manager import ( # noqa: E402 PromptManagerFocused, PromptManagerSplit, ) +from .prompt_manager_pair import PromptManagerPair # noqa: E402 from .prompt_manager_text import PromptManagerText # noqa: E402 from .prompt_search_list import PromptSearchList # noqa: E402 @@ -52,6 +53,7 @@ NODE_CLASS_MAPPINGS = { "PromptManager": PromptManager, "PromptManagerFocused": PromptManagerFocused, "PromptManagerSplit": PromptManagerSplit, + "PromptManagerPair": PromptManagerPair, "PromptManagerText": PromptManagerText, "PromptSearchList": PromptSearchList, } @@ -60,6 +62,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "PromptManager": "Prompt Manager (Classic)", "PromptManagerFocused": "Prompt Manager (Focused)", "PromptManagerSplit": "Prompt Manager (Split)", + "PromptManagerPair": "Prompt Manager (+/-)", "PromptManagerText": "Prompt Manager Text", "PromptSearchList": "Prompt Search List", } diff --git a/prompt_manager_pair.py b/prompt_manager_pair.py new file mode 100644 index 0000000..0bea74c --- /dev/null +++ b/prompt_manager_pair.py @@ -0,0 +1,184 @@ +"""Paired positive and negative conditioning with reusable recipe storage.""" + +import hashlib +import json +import time +from typing import Any, Dict, Optional, Tuple + +try: + from .prompt_manager import IO, InputTypeDict, PromptManagerFocused + from .utils.hashing import generate_prompt_hash + from .utils.prompt_graph import intrinsic_pair_nodes, run_prompt_nodes + from .utils.prompt_pair_resolver import resolve_text_input + from .utils.usage_tracking import ( + PENDING_FIRST_USE, + PENDING_PAIR_FIRST_USE, + pair_run_key, + ) +except ImportError: + from prompt_manager import IO, InputTypeDict, PromptManagerFocused + from utils.hashing import generate_prompt_hash + from utils.prompt_graph import intrinsic_pair_nodes, run_prompt_nodes + from utils.prompt_pair_resolver import resolve_text_input + from utils.usage_tracking import ( + PENDING_FIRST_USE, + PENDING_PAIR_FIRST_USE, + pair_run_key, + ) + + +class PromptManagerPair(PromptManagerFocused): + """Keep raw editors as one recipe while encoding their effective texts.""" + + @classmethod + def INPUT_TYPES(cls) -> InputTypeDict: + inputs = super().INPUT_TYPES() + inputs["required"] = { + "text": inputs["required"]["text"], + "negative_text": ( + IO.STRING, + {"default": "", "multiline": True, "dynamicPrompts": True}, + ), + "clip": inputs["required"]["clip"], + } + return inputs + + RETURN_TYPES = (IO.CONDITIONING, IO.CONDITIONING, IO.STRING, IO.STRING) + RETURN_NAMES = ("positive", "negative", "positive_text", "negative_text") + OUTPUT_TOOLTIPS = ( + "Positive conditioning with composition and optional LoRA trigger words.", + "Negative conditioning from the negative editor, including empty text.", + "The effective positive text that was encoded.", + "The unchanged negative text that was encoded.", + ) + DESCRIPTION = "Encode and save a reusable positive/negative prompt pair." + + def encode_prompt( + self, + clip: Any, + text: str, + negative_text: str = "", + category: str = "", + tags: str = "", + search_text: str = "", + prepend_text: str = "", + append_text: str = "", + prompt: Optional[Dict[str, Any]] = None, + unique_id: Optional[str] = None, + ) -> Tuple[Any, Any, str, str]: + """Encode both sides once; database and tracking failures never block CLIP.""" + if clip is None: + raise RuntimeError("ERROR: clip input is invalid: None") + if not isinstance(text, str) or not isinstance(negative_text, str): + raise ValueError("Positive and negative prompt text must be strings") + positive = self._inject_lora_trigger_words( + " ".join( + part + for part in (prepend_text.strip(), text, append_text.strip()) + if part + ) + ) + positive_conditioning = clip.encode_from_tokens_scheduled( + clip.tokenize(positive) + ) + negative_conditioning = clip.encode_from_tokens_scheduled( + clip.tokenize(negative_text) + ) + pair_id = prompt_id = None + known = isinstance(prompt, dict) and unique_id is not None + node_id = str(unique_id) + is_positive = node_id in run_prompt_nodes(prompt) if known else True + pair_used = node_id in intrinsic_pair_nodes(prompt) if known else True + try: + pair_id = self.db.save_prompt_pair( + text, + negative_text, + category=category.strip() or None, + tags=self._parse_tags(tags), + ) + pair = self.db.get_prompt_pair(pair_id) + prompt_id = pair["positive_prompt_id"] + resolved_positive = ( + resolve_text_input(prompt, node_id, "text") if known else None + ) + resolved_negative = ( + resolve_text_input(prompt, node_id, "negative_text") if known else None + ) + positive_seen = resolved_positive is not None and generate_prompt_hash( + resolved_positive + ) == generate_prompt_hash(text) + pair_seen = ( + positive_seen + and resolved_negative is not None + and generate_prompt_hash(resolved_negative) + == generate_prompt_hash(negative_text) + ) + positive_pending = getattr(self, "pending_first_use", PENDING_FIRST_USE) + pair_pending = getattr( + self, "pending_pair_first_use", PENDING_PAIR_FIRST_USE + ) + if is_positive: + runs = positive_pending.consume(generate_prompt_hash(text)) + int( + not positive_seen + ) + if runs: + self.db.record_prompt_use(prompt_id, times=runs) + if pair_used: + runs = pair_pending.consume(pair_run_key(text, negative_text)) + int( + not pair_seen + ) + if runs: + self.db.record_prompt_pair_use(pair_id, times=runs) + if is_positive: + self.prompt_tracker.set_current_prompt( + prompt_text=positive, + additional_data={ + "prompt_id": prompt_id, + "pair_id": pair_id if pair_used else None, + "positive_text": positive, + "negative_text": negative_text, + "positive_storage_text": text.strip(), + "negative_storage_text": negative_text.strip(), + "prepend_text": prepend_text.strip(), + "append_text": append_text.strip(), + "final_text": positive, + "category": category.strip() or None, + "tags": self._parse_tags(tags) or [], + }, + push_to_queue=not positive_seen or (pair_used and not pair_seen), + ) + except Exception as error: + self.logger.warning(f"Could not save or track paired prompt: {error}") + try: + self.comfyui_integration.register_prompt( + f"promptmanager_pair_{time.time_ns()}", + positive, + { + "prompt_id": prompt_id, + "pair_id": pair_id if pair_used else None, + "negative_text": negative_text, + "category": category.strip() or None, + }, + ) + except Exception as error: + self.logger.warning(f"Could not register paired metadata: {error}") + return positive_conditioning, negative_conditioning, positive, negative_text + + @classmethod + def IS_CHANGED( + cls, + clip: Any, + text: str = "", + negative_text: str = "", + category: str = "", + tags: str = "", + search_text: str = "", + prepend_text: str = "", + append_text: str = "", + **kwargs: Any, + ) -> str: + """Either editor, composition or pair metadata invalidates the node cache.""" + serialized = json.dumps( + [text, negative_text, prepend_text, append_text, category, tags] + ) + return hashlib.sha256(serialized.encode("utf-8")).hexdigest() diff --git a/tests/test_init_package.py b/tests/test_init_package.py index 53ec2bb..39fc9bf 100644 --- a/tests/test_init_package.py +++ b/tests/test_init_package.py @@ -138,6 +138,7 @@ class TestPackageWiring(PackageTestCase): [ "PromptManager", "PromptManagerFocused", + "PromptManagerPair", "PromptManagerSplit", "PromptManagerText", "PromptSearchList", diff --git a/tests/test_pair_usage_tracking.py b/tests/test_pair_usage_tracking.py new file mode 100644 index 0000000..1010e13 --- /dev/null +++ b/tests/test_pair_usage_tracking.py @@ -0,0 +1,248 @@ +"""Paired usage is counted at dequeue, including cached node executions.""" + +import logging +import os +import tempfile +import unittest +from unittest.mock import Mock + +from database.operations import PromptDatabase +from utils import usage_tracking +from utils.prompt_graph import run_prompt_nodes +from utils.usage_tracking import ( + PendingFirstUse, + QueuedRuns, + handle_queued_prompt, + on_dequeue, +) + +try: + from tests.test_prompt_pair_node import pair_graph +except ImportError: + from test_prompt_pair_node import pair_graph + + +class TestPairUsageTracking(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.db = PromptDatabase(os.path.join(self.tmp.name, "pairs.db")) + self.addCleanup(self.db.close_all) + self.queued = QueuedRuns() + self.pair_queued = QueuedRuns() + self.pending = PendingFirstUse() + self.pair_pending = PendingFirstUse() + + def submit(self, graph): + return handle_queued_prompt( + {"prompt": graph}, self.queued, pair_queued=self.pair_queued + ) + + def dequeue(self, request): + on_dequeue( + request["prompt_id"], + self.queued, + lambda: self.db, + self.pending, + pair_queued=self.pair_queued, + pair_pending=self.pair_pending, + ) + + def node(self): + from prompt_manager_pair import PromptManagerPair + + node = PromptManagerPair.__new__(PromptManagerPair) + node.db = self.db + node.logger = logging.getLogger("test.pair_usage") + node.pending_first_use = self.pending + node.pending_pair_first_use = self.pair_pending + node._inject_lora_trigger_words = lambda text: text + node.prompt_tracker = Mock() + node.comfyui_integration = Mock() + return node + + def execute(self, graph, text="cat", negative="blur"): + return self.node().encode_prompt( + Mock(), text, negative_text=negative, prompt=graph, unique_id="1" + ) + + def test_existing_pair_and_positive_count_on_every_dequeue_without_encoding(self): + pid = self.db.save_prompt_pair("cat", "blur") + for _ in range(3): + self.dequeue(self.submit(pair_graph())) + pair = self.db.get_prompt_pair(pid) + self.assertEqual(pair["run_count"], 3) + self.assertEqual( + self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 3 + ) + self.assertEqual( + self.db.get_prompt_by_id(pair["negative_prompt_id"])["run_count"], 0 + ) + + def test_validation_rejection_never_counts_or_saves(self): + self.submit(pair_graph()) + self.assertEqual(self.db.list_prompt_pairs()["total"], 0) + self.assertEqual(self.db.search_prompts(), []) + + def test_new_pair_pending_uses_are_consumed_once_when_first_encoding_saves(self): + graph = pair_graph() + self.dequeue(self.submit(graph)) + self.dequeue(self.submit(graph)) + self.execute(graph) + pair = self.db.list_prompt_pairs()["pairs"][0] + self.assertEqual(pair["run_count"], 2) + self.assertEqual( + self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 2 + ) + self.execute(graph) + self.assertEqual(self.db.get_prompt_pair(pair["id"])["run_count"], 2) + + def test_negative_only_changes_create_distinct_pair_counts_and_positive_counts( + self, + ): + self.db.save_prompt_pair("cat", "blur") + self.db.save_prompt_pair("cat", "watermark") + self.dequeue(self.submit(pair_graph())) + self.dequeue(self.submit(pair_graph(negative="watermark"))) + self.assertEqual( + [p["run_count"] for p in self.db.list_prompt_pairs()["pairs"]], [1, 1] + ) + positive = self.db.get_prompt_by_hash( + usage_tracking.generate_prompt_hash("cat") + ) + self.assertEqual(positive["run_count"], 2) + + def test_negative_output_slots_are_not_positive_prompt_sources(self): + for slot in (1, 3): + with self.subTest(slot=slot): + self.assertEqual( + run_prompt_nodes(pair_graph(positive_slot=slot, negative_slot=0)), + [], + ) + + def test_mismatched_negative_does_not_count_unused_intrinsic_pair(self): + pid = self.db.save_prompt_pair("cat", "blur") + graph = pair_graph() + graph["2"] = { + "class_type": "CLIPTextEncode", + "inputs": {"text": "another negative"}, + } + graph["5"]["inputs"]["negative"] = ["2", 0] + self.dequeue(self.submit(graph)) + self.execute(graph) + self.assertEqual(self.db.get_prompt_pair(pid)["run_count"], 0) + self.assertEqual( + self.db.get_prompt_by_hash(usage_tracking.generate_prompt_hash("cat"))[ + "run_count" + ], + 1, + ) + + def test_unknown_batch_text_counts_each_execution_once_and_tracks_fifo(self): + graph = pair_graph(text=["2", 0]) + graph["2"] = {"class_type": "PromptSearchList", "inputs": {}} + self.submit(graph) + node = self.node() + node.encode_prompt( + Mock(), "batch cat", negative_text="blur", prompt=graph, unique_id="1" + ) + pair = self.db.list_prompt_pairs()["pairs"][0] + self.assertEqual(pair["run_count"], 1) + self.assertEqual( + self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 1 + ) + self.assertTrue( + node.prompt_tracker.set_current_prompt.call_args.kwargs["push_to_queue"] + ) + + def test_standalone_pair_counts_cached_dequeues(self): + pid = self.db.save_prompt_pair("cat", "") + graph = { + "1": { + "class_type": "PromptManagerPair", + "inputs": {"text": "cat", "negative_text": ""}, + } + } + self.dequeue(self.submit(graph)) + self.dequeue(self.submit(graph)) + self.assertEqual(self.db.get_prompt_pair(pid)["run_count"], 2) + + def test_paired_string_outputs_through_encoders_keep_roles(self): + from utils.prompt_graph import intrinsic_pair_nodes + + graph = pair_graph() + graph["2"] = {"class_type": "CLIPTextEncode", "inputs": {"text": ["1", 2]}} + graph["3"] = {"class_type": "CLIPTextEncode", "inputs": {"text": ["1", 3]}} + graph["4"] = {"class_type": "Reroute", "inputs": {"value": ["2", 0]}} + graph["5"]["inputs"].update(positive=["4", 0], negative=["3", 0]) + self.assertEqual(intrinsic_pair_nodes(graph), ["1"]) + pid = self.db.save_prompt_pair("cat", "blur") + self.dequeue(self.submit(graph)) + self.assertEqual(self.db.get_prompt_pair(pid)["run_count"], 1) + graph["3"]["inputs"]["text"] = ["1", 2] + self.assertEqual(intrinsic_pair_nodes(graph), []) + + def test_two_samplers_using_same_pair_count_one_run_and_dequeue_is_idempotent(self): + graph = pair_graph() + graph["6"] = {**graph["5"]} + pid = self.db.save_prompt_pair("cat", "blur") + request = self.submit(graph) + self.dequeue(request) + self.dequeue(request) + self.assertEqual(self.db.get_prompt_pair(pid)["run_count"], 1) + + def test_new_negative_with_existing_positive_counts_pending_pair_only_once(self): + self.db.save_prompt_pair("cat", "old negative") + graph = pair_graph(negative="new negative") + self.dequeue(self.submit(graph)) + self.execute(graph, negative="new negative") + pair = next( + p + for p in self.db.list_prompt_pairs()["pairs"] + if p["negative_text"] == "new negative" + ) + self.assertEqual(pair["run_count"], 1) + self.assertEqual( + self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 1 + ) + + def test_unknown_negative_counts_at_execution_without_double_positive_count(self): + graph = pair_graph(negative=["2", 0]) + graph["2"] = {"class_type": "PromptSearchList", "inputs": {}} + self.dequeue(self.submit(graph)) + self.execute(graph, negative="batch negative") + pair = self.db.list_prompt_pairs()["pairs"][0] + self.assertEqual(pair["run_count"], 1) + self.assertEqual( + self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 1 + ) + + def test_cyclic_unknown_and_oversized_pair_graphs_are_bounded(self): + from utils.prompt_graph import intrinsic_pair_nodes + from utils.prompt_pair_resolver import resolve_text_input + + self.assertEqual(intrinsic_pair_nodes(None), []) + self.assertEqual(intrinsic_pair_nodes({str(n): {} for n in range(501)}), []) + graph = pair_graph() + graph["2"] = {"class_type": "Reroute", "inputs": {"value": ["2", 0]}} + graph["5"]["inputs"]["positive"] = ["2", 0] + self.assertEqual(intrinsic_pair_nodes(graph), []) + self.assertIsNone(resolve_text_input(None, "1", "text")) + self.assertIsNone(resolve_text_input({1: {}, "1": {}}, "1", "text")) + self.assertIsNone(resolve_text_input(graph, "missing", "text")) + self.assertIsNone(resolve_text_input(pair_graph(), "1", "text", max_depth=0)) + graph["1"]["inputs"]["text"] = ["1", 3] + self.assertEqual(resolve_text_input(graph, "1", "text"), "blur") + self.assertIsNone(resolve_text_input(graph, "1", "text", max_text_length=2)) + + def test_replacing_client_prompt_id_clears_obsolete_pair_records(self): + pid = self.db.save_prompt_pair("cat", "blur") + request = {"prompt_id": "same-id", "prompt": pair_graph()} + handle_queued_prompt(request, self.queued, pair_queued=self.pair_queued) + replacement = pair_graph() + replacement["5"]["inputs"]["positive"] = ["1", 1] + replacement["5"]["inputs"]["negative"] = ["1", 0] + request = {"prompt_id": "same-id", "prompt": replacement} + handle_queued_prompt(request, self.queued, pair_queued=self.pair_queued) + self.dequeue(request) + self.assertEqual(self.db.get_prompt_pair(pid)["run_count"], 0) diff --git a/tests/test_prompt_pair_node.py b/tests/test_prompt_pair_node.py new file mode 100644 index 0000000..eff150c --- /dev/null +++ b/tests/test_prompt_pair_node.py @@ -0,0 +1,154 @@ +"""Paired node encoding preserves recipe storage and output roles.""" + +import logging +import os +import tempfile +import unittest +from unittest.mock import Mock + +from database.operations import PromptDatabase +from prompt_manager_pair import PromptManagerPair +from utils.usage_tracking import PendingFirstUse + + +def pair_graph(text="cat", negative="blur", positive_slot=0, negative_slot=1): + return { + "1": { + "class_type": "PromptManagerPair", + "inputs": {"text": text, "negative_text": negative}, + }, + "5": { + "class_type": "KSampler", + "inputs": { + "model": ["9", 0], + "positive": ["1", positive_slot], + "negative": ["1", negative_slot], + }, + }, + } + + +class TestPromptPairNode(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.db = PromptDatabase(os.path.join(self.tmp.name, "pairs.db")) + self.addCleanup(self.db.close_all) + self.node = PromptManagerPair.__new__(PromptManagerPair) + self.node.db = self.db + self.node.logger = logging.getLogger("test.pair_node") + self.node.prompt_tracker = Mock() + self.node.comfyui_integration = Mock() + self.node._inject_lora_trigger_words = Mock( + side_effect=lambda text: text + " LORA" + ) + self.node.pending_first_use = PendingFirstUse() + self.node.pending_pair_first_use = PendingFirstUse() + self.clip = Mock() + self.clip.encode_from_tokens_scheduled.side_effect = lambda tokens: ( + "conditioning", + tokens, + ) + self.clip.tokenize.side_effect = lambda text: text + + def test_schema_has_positive_composition_and_four_ordered_outputs(self): + schema = PromptManagerPair.INPUT_TYPES() + self.assertEqual(list(schema["required"]), ["text", "negative_text", "clip"]) + self.assertTrue(schema["required"]["negative_text"][1]["multiline"]) + self.assertEqual(schema["required"]["negative_text"][1]["default"], "") + for name in ("prepend_text", "append_text"): + self.assertTrue(schema["optional"][name][1]["forceInput"]) + self.assertEqual( + PromptManagerPair.RETURN_TYPES, + ("CONDITIONING", "CONDITIONING", "STRING", "STRING"), + ) + self.assertNotEqual( + PromptManagerPair.IS_CHANGED(None, "cat", negative_text="blur"), + PromptManagerPair.IS_CHANGED(None, "cat", negative_text="watermark"), + ) + + def test_encodes_each_side_once_positive_composition_and_lora_only(self): + graph = pair_graph() + result = self.node.encode_prompt( + self.clip, + "cat", + negative_text=" blur ", + prepend_text=" trigger ", + append_text=" detail ", + category="Animals", + tags="studio", + prompt=graph, + unique_id="1", + ) + self.assertEqual( + result, + ( + ("conditioning", "trigger cat detail LORA"), + ("conditioning", " blur "), + "trigger cat detail LORA", + " blur ", + ), + ) + self.node._inject_lora_trigger_words.assert_called_once_with( + "trigger cat detail" + ) + self.assertEqual(self.clip.tokenize.call_count, 2) + pair = self.db.list_prompt_pairs()["pairs"][0] + self.assertEqual( + ( + pair["positive_text"], + pair["negative_text"], + pair["category"], + pair["tags"], + ), + ("cat", "blur", "Animals", ["studio"]), + ) + self.assertIsNone( + self.db.get_prompt_by_id(pair["positive_prompt_id"])["category"] + ) + tracked = self.node.prompt_tracker.set_current_prompt.call_args.kwargs + self.assertEqual(tracked["additional_data"]["pair_id"], pair["id"]) + self.assertEqual(tracked["additional_data"]["negative_text"], " blur ") + self.assertEqual(tracked["additional_data"]["positive_storage_text"], "cat") + + def test_empty_negative_encodes_empty_text_and_stores_null_component(self): + result = self.node.encode_prompt(self.clip, "cat", negative_text="") + self.assertEqual(result[3], "") + self.assertIsNone(self.db.list_prompt_pairs()["pairs"][0]["negative_prompt_id"]) + + def test_invalid_clip_cannot_save_or_count(self): + with self.assertRaises(RuntimeError): + self.node.encode_prompt(None, "cat", negative_text="blur") + self.assertEqual(self.db.list_prompt_pairs()["total"], 0) + self.node.prompt_tracker.set_current_prompt.assert_not_called() + + def test_persistence_failure_does_not_block_conditioning(self): + self.node.db = Mock() + self.node.db.save_prompt_pair.side_effect = RuntimeError("unavailable") + result = self.node.encode_prompt(self.clip, "cat", negative_text="blur") + self.assertEqual(result[2:], ("cat LORA", "blur")) + self.node.prompt_tracker.set_current_prompt.assert_not_called() + + def test_negative_output_used_as_positive_does_not_claim_positive_usage(self): + self.node.encode_prompt( + self.clip, + "cat", + negative_text="blur", + prompt=pair_graph(positive_slot=1, negative_slot=0), + unique_id="1", + ) + pair = self.db.list_prompt_pairs()["pairs"][0] + self.assertEqual(pair["run_count"], 0) + self.assertEqual( + self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 0 + ) + self.node.prompt_tracker.set_current_prompt.assert_not_called() + + def test_invalid_text_does_not_write_and_integration_failures_do_not_block(self): + with self.assertRaises(ValueError): + self.node.encode_prompt(self.clip, [], negative_text="blur") + self.node.comfyui_integration.register_prompt.side_effect = RuntimeError( + "metadata unavailable" + ) + result = self.node.encode_prompt(self.clip, "cat", negative_text="blur") + self.assertEqual(result[2:], ("cat LORA", "blur")) diff --git a/utils/prompt_graph.py b/utils/prompt_graph.py index acc3b2e..c555775 100644 --- a/utils/prompt_graph.py +++ b/utils/prompt_graph.py @@ -12,7 +12,13 @@ import re from typing import Any, Dict, List, Optional PROMPT_MANAGER_TYPES = frozenset( - {"PromptManager", "PromptManagerText", "PromptManagerFocused", "PromptManagerSplit"} + { + "PromptManager", + "PromptManagerText", + "PromptManagerFocused", + "PromptManagerSplit", + "PromptManagerPair", + } ) # Upper bound on nodes expanded per graph; queued and embedded graphs are untrusted @@ -105,15 +111,19 @@ def positive_prompt_nodes(graph: Any) -> 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: + link = stack.pop() + node_id, slot = str(link[0]), link[1] + if (node_id, slot) in seen: continue - seen.add(node_id) + seen.add((node_id, slot)) node = graph.get(node_id) if not isinstance(node, dict): continue + if node.get("class_type") == "PromptManagerPair" and slot not in (0, 2): + continue if node.get("class_type") in PROMPT_MANAGER_TYPES: - found.append(node_id) + if node_id not in found: + found.append(node_id) continue stack.extend(reversed(_upstream_links(_inputs(node)))) return found @@ -228,3 +238,67 @@ def resolve_text(graph: Any, node_id: str) -> Optional[str]: text = resolve(_inputs(graph.get(str(node_id))).get("text")) return text.strip() if isinstance(text, str) and text.strip() else None + + +def intrinsic_pair_nodes(graph: Any) -> List[str]: + """Paired nodes whose matching outputs feed the same sampler, or standalone.""" + if not isinstance(graph, dict) or len(graph) > MAX_VISITS: + return [] + if not has_sampler(graph): + return [ + str(key) + for key, value in graph.items() + if isinstance(value, dict) + and value.get("class_type") == "PromptManagerPair" + ] + + def origin(link: Any, role: str, kind: str = "conditioning") -> Optional[str]: + seen = set() + while _is_link(link) and len(seen) < MAX_DEPTH: + node_id, slot = str(link[0]), link[1] + if (node_id, slot, kind) in seen: + return None + seen.add((node_id, slot, kind)) + node = graph.get(node_id) + if not isinstance(node, dict): + return None + name, inputs = node.get("class_type"), _inputs(node) + if name == "PromptManagerPair": + expected = { + ("positive", "conditioning"): 0, + ("negative", "conditioning"): 1, + ("positive", "string"): 2, + ("negative", "string"): 3, + } + return node_id if slot == expected[(role, kind)] else None + if name == "CLIPTextEncode" and slot == 0 and kind == "conditioning": + link, kind = inputs.get("text"), "string" + elif name == "Reroute" and slot == 0: + keys = [key for key in ("value", "input", "in") if key in inputs] + if len(keys) != 1: + return None + link = inputs[keys[0]] + else: + return None + return None + + found = [] + for node in graph.values(): + if not isinstance(node, dict): + continue + inputs = _inputs(node) + if not _is_link(inputs.get("positive")) or not _is_link(inputs.get("negative")): + continue + if ( + not _is_link(inputs.get("model")) + and "sampler" not in str(node.get("class_type", "")).lower() + ): + continue + positive = origin(inputs["positive"], "positive") + if ( + positive + and positive == origin(inputs["negative"], "negative") + and positive not in found + ): + found.append(positive) + return found diff --git a/utils/prompt_pair_resolver.py b/utils/prompt_pair_resolver.py index 9edc940..c9bd353 100644 --- a/utils/prompt_pair_resolver.py +++ b/utils/prompt_pair_resolver.py @@ -392,3 +392,26 @@ def resolve_prompt_pair( candidates=[resolver.pair(node_id) for node_id in sampler_ids], ) return resolver.pair(sampler_ids[0]) + + +def resolve_text_input( + graph: Any, + node_id: str, + input_name: str, + *, + max_visits: int = 500, + max_depth: int = 64, + max_text_length: int = 100_000, +) -> Optional[str]: + """Resolve one raw STRING input; empty text is distinct from unknown input.""" + if not isinstance(graph, dict) or len(graph) > max_visits: + return None + normalized = {str(key): value for key, value in graph.items()} + if len(normalized) != len(graph) or max_depth <= 0 or max_text_length <= 0: + return None + resolver = _GraphResolver(normalized, {}, max_visits, max_depth, max_text_length) + try: + _, inputs = resolver.node(str(node_id)) + return resolver.text(inputs.get(input_name), "string")[0] + except _Unresolved: + return None diff --git a/utils/usage_tracking.py b/utils/usage_tracking.py index 7dcaadf..217241d 100644 --- a/utils/usage_tracking.py +++ b/utils/usage_tracking.py @@ -29,11 +29,13 @@ from typing import Any, Callable, List, Optional, Tuple try: from .hashing import generate_prompt_hash from .logging_config import get_logger - from .prompt_graph import resolve_text, run_prompt_nodes + from .prompt_graph import intrinsic_pair_nodes, resolve_text, run_prompt_nodes + from .prompt_pair_resolver import resolve_text_input except ImportError: from utils.hashing import generate_prompt_hash from utils.logging_config import get_logger - from utils.prompt_graph import resolve_text, run_prompt_nodes + from utils.prompt_graph import intrinsic_pair_nodes, resolve_text, run_prompt_nodes + from utils.prompt_pair_resolver import resolve_text_input logger = get_logger("prompt_manager.usage_tracking") @@ -88,10 +90,10 @@ class QueuedRuns: self._lock = threading.Lock() def put(self, prompt_id: str, runs: List[Run]) -> None: - if not runs: - return with self._lock: self._entries.pop(prompt_id, None) + if not runs: + return self._entries[prompt_id] = list(runs) while len(self._entries) > self._max: self._entries.popitem(last=False) @@ -108,6 +110,8 @@ class QueuedRuns: # Shared between the queue hook and the nodes PENDING_FIRST_USE = PendingFirstUse() QUEUED_RUNS = QueuedRuns() +PENDING_PAIR_FIRST_USE = PendingFirstUse() +QUEUED_PAIR_RUNS = QueuedRuns() def queued_prompt_runs(graph: Any) -> List[Run]: @@ -118,7 +122,12 @@ def queued_prompt_runs(graph: Any) -> List[Run]: runs: List[Run] = [] seen = set() for node_id in run_prompt_nodes(graph): - text = resolve_text(graph, node_id) + node = graph.get(node_id, {}) if isinstance(graph, dict) else {} + text = ( + resolve_text_input(graph, node_id, "text") + if node.get("class_type") == "PromptManagerPair" + else resolve_text(graph, node_id) + ) if not text: continue prompt_hash = generate_prompt_hash(text) @@ -129,6 +138,62 @@ def queued_prompt_runs(graph: Any) -> List[Run]: return runs +def pair_run_key(positive: str, negative: str) -> str: + """Stable ordered identity; an empty negative remains an explicit component.""" + return generate_prompt_hash(positive) + ":" + generate_prompt_hash(negative) + + +def queued_pair_runs(graph: Any) -> List[Tuple[str, str, str]]: + """Resolve raw paired recipes only when both outputs form the used pair.""" + found, seen = [], set() + for node_id in intrinsic_pair_nodes(graph): + positive = resolve_text_input(graph, node_id, "text") + negative = resolve_text_input(graph, node_id, "negative_text") + if positive is None or not positive.strip() or negative is None: + continue + key = pair_run_key(positive, negative) + if key not in seen: + seen.add(key) + found.append((key, positive, negative)) + return found + + +def find_pair_id(db: Any, positive: str, negative: str) -> Optional[int]: + """Find a saved pair without creating or overwriting its metadata.""" + positive_row = db.get_prompt_by_hash(generate_prompt_hash(positive)) + negative_row = ( + db.get_prompt_by_hash(generate_prompt_hash(negative)) + if negative.strip() + else None + ) + if not positive_row or (negative.strip() and not negative_row): + return None + row = ( + db.model.get_connection() + .execute( + "SELECT id FROM prompt_pairs " + "WHERE positive_prompt_id=? AND negative_prompt_id IS ?", + (positive_row["id"], negative_row["id"] if negative_row else None), + ) + .fetchone() + ) + return row["id"] if row else None + + +def count_pair_runs( + runs: List[Tuple[str, str, str]], + db: Any, + pending: PendingFirstUse, +) -> None: + """Count existing recipes at dequeue and defer new recipes until first save.""" + for key, positive, negative in runs: + pair_id = find_pair_id(db, positive, negative) + if pair_id: + db.record_prompt_pair_use(pair_id) + else: + pending.add(key) + + def count_runs(runs: List[Run], db: Any, pending: PendingFirstUse) -> None: """Record one use per prompt in the database; remember the rest for first save.""" for prompt_hash, _text in runs: @@ -139,7 +204,12 @@ def count_runs(runs: List[Run], db: Any, pending: PendingFirstUse) -> None: pending.add(prompt_hash) -def handle_queued_prompt(json_data: Any, queued: QueuedRuns = QUEUED_RUNS) -> Any: +def handle_queued_prompt( + json_data: Any, + queued: QueuedRuns = QUEUED_RUNS, + *, + pair_queued: QueuedRuns = QUEUED_PAIR_RUNS, +) -> Any: """ComfyUI on-prompt handler body: stash the runs, always return the request. Adds a canonical UUID4 ``prompt_id`` when the client sent none, so the stash @@ -149,13 +219,18 @@ def handle_queued_prompt(json_data: Any, queued: QueuedRuns = QUEUED_RUNS) -> An try: graph = json_data.get("prompt") if isinstance(json_data, dict) else None runs = queued_prompt_runs(graph) - if not runs: + pair_runs = queued_pair_runs(graph) + if not runs and not pair_runs: + if isinstance(json_data, dict) and json_data.get("prompt_id") is not None: + queued.put(str(json_data["prompt_id"]), []) + pair_queued.put(str(json_data["prompt_id"]), []) return json_data prompt_id = json_data.get("prompt_id") if prompt_id is None: prompt_id = str(uuid.uuid4()) json_data["prompt_id"] = prompt_id queued.put(str(prompt_id), runs) + pair_queued.put(str(prompt_id), pair_runs) except Exception as e: logger.warning(f"Could not resolve queued prompt runs: {e}") return json_data @@ -166,12 +241,18 @@ def on_dequeue( queued: QueuedRuns, db_factory: Callable[[], Any], pending: PendingFirstUse, + *, + pair_queued: QueuedRuns = QUEUED_PAIR_RUNS, + pair_pending: PendingFirstUse = PENDING_PAIR_FIRST_USE, ) -> None: """Count the runs stashed for a prompt ComfyUI is about to execute.""" try: runs = queued.pop(str(prompt_id)) - if runs: - count_runs(runs, db_factory(), pending) + pair_runs = pair_queued.pop(str(prompt_id)) + if runs or pair_runs: + db = db_factory() + count_runs(runs, db, pending) + count_pair_runs(pair_runs, db, pair_pending) except Exception as e: logger.warning(f"Could not count runs of dequeued prompt {prompt_id}: {e}") @@ -183,8 +264,11 @@ def count_queued_prompt( try: graph = json_data.get("prompt") if isinstance(json_data, dict) else None runs = queued_prompt_runs(graph) - if runs: - count_runs(runs, db_factory(), pending) + pair_runs = queued_pair_runs(graph) + if runs or pair_runs: + db = db_factory() + count_runs(runs, db, pending) + count_pair_runs(pair_runs, db, PENDING_PAIR_FIRST_USE) except Exception as e: logger.warning(f"Could not count queued prompt runs: {e}") return json_data