diff --git a/__init__.py b/__init__.py index ee90797..2e07092 100644 --- a/__init__.py +++ b/__init__.py @@ -124,7 +124,11 @@ except Exception as e: get_logger("prompt_manager.init").error(f"Failed to register usage hook: {e}") except Exception: - print(f"[ComfyUI-PromptManager] Warning: Failed to register usage hook: {e}") + import logging + + logging.getLogger("prompt_manager.init").error( + f"Failed to register usage hook: {e}" + ) __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/database/operations.py b/database/operations.py index 5a82cfd..f3c0765 100644 --- a/database/operations.py +++ b/database/operations.py @@ -163,21 +163,22 @@ class PromptDatabase: self.logger.debug(f"Successfully saved prompt with ID: {prompt_id}") return prompt_id - def record_prompt_use(self, prompt_id: int) -> bool: + def record_prompt_use(self, prompt_id: int, times: int = 1) -> bool: """ - Count one run of a prompt and mark it as most recently used. + Count runs of a prompt and mark it as most recently used. Args: prompt_id: The prompt ID + times: Number of runs to add (queued runs counted on first save) Returns: True if the prompt exists and was updated """ with self.model.get_connection() as conn: cursor = conn.execute( - "UPDATE prompts SET run_count = run_count + 1, last_used_at = ?" + "UPDATE prompts SET run_count = run_count + ?, last_used_at = ?" " WHERE id = ?", - (utc_now_iso(), prompt_id), + (max(1, int(times)), utc_now_iso(), prompt_id), ) conn.commit() return cursor.rowcount > 0 diff --git a/prompt_manager_base.py b/prompt_manager_base.py index 3d05ce0..ee190d5 100644 --- a/prompt_manager_base.py +++ b/prompt_manager_base.py @@ -24,7 +24,8 @@ try: from .utils.comfyui_integration import get_comfyui_integration from .utils.image_monitor import get_image_monitor from .utils.prompt_tracker import get_prompt_tracker - from .utils.prompt_graph import is_text_linked, run_prompt_nodes + from .utils.hashing import generate_prompt_hash + from .utils.prompt_graph import resolve_text, run_prompt_nodes from .utils.usage_tracking import PENDING_FIRST_USE except ImportError: import sys @@ -34,7 +35,8 @@ except ImportError: from utils.comfyui_integration import get_comfyui_integration from utils.image_monitor import get_image_monitor from utils.prompt_tracker import get_prompt_tracker - from utils.prompt_graph import is_text_linked, run_prompt_nodes + from utils.hashing import generate_prompt_hash + from utils.prompt_graph import resolve_text, run_prompt_nodes from utils.usage_tracking import PENDING_FIRST_USE @@ -109,8 +111,9 @@ class PromptManagerBase: if prompt_id: self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}") pending = getattr(self, "pending_first_use", PENDING_FIRST_USE) - if count_run or pending.consume(prompt_hash): - self._record_use(prompt_id) + runs = pending.consume(prompt_hash) + (1 if count_run else 0) + if runs: + self._record_use(prompt_id, times=runs) else: self.logger.warning("Failed to save prompt - no ID returned") @@ -120,20 +123,25 @@ class PromptManagerBase: self.logger.error(f"Error saving prompt to database: {e}") return None - def _prompt_role(self, prompt_graph: Any, unique_id: Any) -> Dict[str, bool]: + def _prompt_role( + self, prompt_graph: Any, unique_id: Any, text: str + ) -> Dict[str, bool]: """How this node's execution relates to the run, from the queued graph. Returns flags: is_run_prompt: feeds a sampler's positive input (images belong to it) - text_linked: text comes from another node (e.g. PromptSearchList) + seen_at_queue: the queue hook resolved exactly this text, so it already + counted the run and the image's metadata can identify the prompt known: the graph and node id were available """ if not isinstance(prompt_graph, dict) or unique_id is None: - return {"is_run_prompt": True, "text_linked": False, "known": False} + return {"is_run_prompt": True, "seen_at_queue": False, "known": False} node_id = str(unique_id) + resolved = resolve_text(prompt_graph, node_id) return { "is_run_prompt": node_id in run_prompt_nodes(prompt_graph), - "text_linked": is_text_linked(prompt_graph, node_id), + "seen_at_queue": resolved is not None + and generate_prompt_hash(resolved) == generate_prompt_hash(text), "known": True, } @@ -150,12 +158,15 @@ class PromptManagerBase: """Save the prompt and, for the run's positive prompt, register it for linking. Negative-prompt nodes are saved but never become the prompt their images are - linked to. Typed prompts are linked from the image's own metadata, so only - batch items (linked text) are pushed onto the FIFO queue. Without graph - context the pre-3.2.4 behaviour is kept. + linked to. Text the queue hook could resolve is counted there and linked from + the image's own metadata; anything else (PromptSearchList batches, unknown + text nodes, or a resolution that doesn't match) is counted here and pushed + onto the FIFO queue. Without graph context the pre-3.2.4 behaviour is kept. """ - role = self._prompt_role(prompt_graph, unique_id) - batch_item = role["is_run_prompt"] and role["text_linked"] + role = self._prompt_role(prompt_graph, unique_id, text) + batch_item = ( + role["known"] and role["is_run_prompt"] and not role["seen_at_queue"] + ) prompt_id = self._save_prompt_to_database( text=text, category=category, tags=tags, count_run=batch_item ) @@ -167,10 +178,10 @@ class PromptManagerBase: ) return prompt_id - def _record_use(self, prompt_id: int) -> None: - """Count this execution; a tracking failure must never block generation.""" + def _record_use(self, prompt_id: int, times: int = 1) -> None: + """Count runs; a tracking failure must never block generation.""" try: - self.db.record_prompt_use(prompt_id) + self.db.record_prompt_use(prompt_id, times=times) except Exception as e: self.logger.warning(f"Could not record use of prompt {prompt_id}: {e}") diff --git a/tests/test_image_linking.py b/tests/test_image_linking.py index 167b7a7..9c2827a 100644 --- a/tests/test_image_linking.py +++ b/tests/test_image_linking.py @@ -4,7 +4,7 @@ import os import sys import tempfile import unittest -from unittest import mock +import unittest.mock as mock sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -87,6 +87,16 @@ class TestMetadataLinking(LinkingTestCase): self.tracker.pop_next_prompt()["id"], stale, "queue left intact" ) + def test_text_from_a_known_string_node_links_from_metadata(self): + positive = self._save("a cat") + linked_graph = graph() + linked_graph["134"]["inputs"]["text"] = ["8", 0] + linked_graph["8"] = { + "class_type": "PrimitiveString", + "inputs": {"value": "a cat"}, + } + self.assertEqual(self._process({"prompt": linked_graph}), positive) + def test_batch_items_with_linked_text_use_the_queue(self): first, second = self._save("item one"), self._save("item two") batch_graph = graph() diff --git a/tests/test_node_tracking_wiring.py b/tests/test_node_tracking_wiring.py index fe3ae9a..e880cf5 100644 --- a/tests/test_node_tracking_wiring.py +++ b/tests/test_node_tracking_wiring.py @@ -2,8 +2,9 @@ import os import sys +import threading import unittest -from unittest import mock +import unittest.mock as mock sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -68,7 +69,7 @@ class TestQueueHookRegistration(unittest.TestCase): server = mock.Mock() db = mock.Mock() db.get_prompt_by_hash.return_value = None - with mock.patch.object(usage_tracking, "_registered", False): + with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()): self.assertTrue(usage_tracking.register_queue_hook(server, lambda: db)) self.assertFalse(usage_tracking.register_queue_hook(server, lambda: db)) diff --git a/tests/test_prompt_graph.py b/tests/test_prompt_graph.py index 7aa96c3..b0b6185 100644 --- a/tests/test_prompt_graph.py +++ b/tests/test_prompt_graph.py @@ -7,9 +7,8 @@ import unittest sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from utils.prompt_graph import ( - is_text_linked, - literal_text, positive_prompt_nodes, + resolve_text, run_prompt_nodes, ) @@ -169,22 +168,76 @@ class TestRunPromptNodes(unittest.TestCase): self.assertEqual(run_prompt_nodes(None), []) -class TestNodeText(unittest.TestCase): - def test_literal_and_linked_text(self): - graph = { - "1": pm(" a cat "), - "2": pm(["7", 0]), - "7": {"class_type": "PromptSearchList", "inputs": {}}, - } - self.assertEqual(literal_text(graph, "1"), "a cat") - self.assertIsNone(literal_text(graph, "2")) - self.assertFalse(is_text_linked(graph, "1")) - self.assertTrue(is_text_linked(graph, "2")) +class TestResolveText(unittest.TestCase): + """Text known at queue time: typed, or from known pure string nodes.""" - def test_missing_nodes(self): - self.assertIsNone(literal_text({}, "1")) - self.assertFalse(is_text_linked({}, "1")) - self.assertIsNone(literal_text({"1": pm(" ")}, "1")) + 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_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__": diff --git a/tests/test_prompt_usage.py b/tests/test_prompt_usage.py index 41936fe..5a0ce6c 100644 --- a/tests/test_prompt_usage.py +++ b/tests/test_prompt_usage.py @@ -77,6 +77,11 @@ class TestRecordPromptUse(UsageTestCase): result = self.db.get_recent_prompts(limit=2, sort="last_used_desc") self.assertEqual(self._ids(result["prompts"]), [recent, backfilled]) + def test_record_use_can_add_several_runs_at_once(self): + pid = self._save("several") + self.db.record_prompt_use(pid, times=3) + self.assertEqual(self.db.get_prompt_by_id(pid)["run_count"], 3) + def test_record_use_on_missing_prompt_returns_false(self): self.assertFalse(self.db.record_prompt_use(999999)) diff --git a/tests/test_usage_tracking.py b/tests/test_usage_tracking.py index 664b43a..d5740f2 100644 --- a/tests/test_usage_tracking.py +++ b/tests/test_usage_tracking.py @@ -8,7 +8,7 @@ import os import sys import tempfile import unittest -from unittest import mock +import unittest.mock as mock sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -99,6 +99,20 @@ class TestQueueHook(DbTestCase): self._queue(graph) self.assertEqual(self._runs(positive), 1) + def test_text_from_a_known_string_node_is_counted_at_queue_time(self): + positive = self._save("a cat") + graph = workflow() + graph["134"]["inputs"]["text"] = ["8", 0] + graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}} + self._queue(graph) + self._queue(graph) # cached re-run: the node won't execute + self.assertEqual(self._runs(positive), 2) + + def test_new_prompt_queued_twice_before_first_run_counts_twice(self): + self._queue(workflow(positive_text="queued twice")) + self._queue(workflow(positive_text="queued twice")) + self.assertEqual(self.pending.consume(generate_prompt_hash("queued twice")), 2) + def test_linked_text_is_left_to_node_execution(self): positive = self._save("from a batch") graph = workflow() @@ -125,6 +139,14 @@ class TestPendingFirstUse(unittest.TestCase): with mock.patch("utils.usage_tracking.time.monotonic", return_value=111.0): self.assertFalse(pending.consume("h")) + def test_consume_returns_how_many_runs_were_pending(self): + pending = PendingFirstUse() + pending.add("h") + pending.add("h") + pending.add("h") + self.assertEqual(pending.consume("h"), 3) + self.assertEqual(pending.consume("h"), 0) + def test_size_is_capped(self): pending = PendingFirstUse(max_entries=3) for h in ("a", "b", "c", "d"): @@ -186,6 +208,39 @@ class TestNodeRoles(DbTestCase): ) self.assertEqual(self._runs(prompt_id), 1) + def test_new_prompt_queued_twice_counts_both_runs_on_first_save(self): + graph = workflow(positive_text="double queued") + self._queue(graph) + self._queue(graph) + prompt_id = self._run(self._node(), graph, "134", "double queued") + self.assertEqual(self._runs(prompt_id), 2) + + def test_text_resolved_at_queue_time_is_not_counted_again_or_queued(self): + graph = workflow() + graph["134"]["inputs"]["text"] = ["8", 0] + graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}} + positive = self._save("a cat") + self._queue(graph) + node = self._node() + self._run(node, graph, "134", "a cat") + self.assertEqual(self._runs(positive), 1) + self.assertEqual( + node.prompt_tracker.calls, [{"text": "a cat", "push_to_queue": False}] + ) + + def test_resolution_mismatch_falls_back_to_counting_and_queueing(self): + # Hook guessed "a cat" but the node received different text: trust the node + graph = workflow() + graph["134"]["inputs"]["text"] = ["8", 0] + graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}} + node = self._node() + prompt_id = self._run(node, graph, "134", "actually a dog") + self.assertEqual(self._runs(prompt_id), 1) + self.assertEqual( + node.prompt_tracker.calls, + [{"text": "actually a dog", "push_to_queue": True}], + ) + def test_existing_prompt_is_not_counted_twice_by_hook_and_node(self): positive = self._save("a cat") self._queue(workflow()) diff --git a/utils/image_monitor.py b/utils/image_monitor.py index c011721..efcfb7f 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -34,7 +34,7 @@ from watchdog.events import FileSystemEventHandler from .metadata_extractor import ComfyUIMetadataExtractor from .logging_config import get_logger -from .prompt_graph import literal_text, run_prompt_nodes +from .prompt_graph import resolve_text, run_prompt_nodes class ImageGenerationHandler(FileSystemEventHandler): @@ -241,12 +241,12 @@ class ImageGenerationHandler(FileSystemEventHandler): prompt_data = metadata.get("prompt") if isinstance(prompt_data, dict): for node_id in run_prompt_nodes(prompt_data): - prompt_text = literal_text(prompt_data, node_id) + prompt_text = resolve_text(prompt_data, node_id) if prompt_text: break if not prompt_text: - # Positive text came from another node (batch item): leave it to the - # queue rather than guess from workflow widgets, which don't know roles + # Positive text can't be known from the graph (batch item or unknown + # node): leave it to the queue rather than guess from workflow widgets return None # Fallback (no executed graph): check text_encoder_nodes from workflow, but only if diff --git a/utils/prompt_graph.py b/utils/prompt_graph.py index 2efd874..fbb0df0 100644 --- a/utils/prompt_graph.py +++ b/utils/prompt_graph.py @@ -106,16 +106,73 @@ def run_prompt_nodes(graph: Any) -> List[str]: ] -def literal_text(graph: Any, node_id: str) -> Optional[str]: - """A node's `text` input when typed in (not linked), stripped; else None.""" - node = graph.get(str(node_id)) if isinstance(graph, dict) else None - text = _inputs(node).get("text") - if isinstance(text, str) and text.strip(): - return text.strip() +def _join_text_inputs( + inputs: Dict[str, Any], keys: List[str], resolve +) -> Optional[str]: + """WAS-style join: resolve each key, optionally strip, skip empty, join by delimiter.""" + delimiter = inputs.get("delimiter", " ") + if not isinstance(delimiter, str): + return None + clean = str(inputs.get("clean_whitespace", "true")).lower() == "true" + parts = [] + for key in keys: + value = resolve(inputs.get(key, "")) + if value is None: + return None + value = value.strip() if clean else value + if value: + parts.append(value) + return delimiter.join(parts) + + +def _resolve_string_node( + class_type: str, inputs: Dict[str, Any], resolve +) -> Optional[str]: + """Output of a known pure string node, or None when the type isn't known.""" + if class_type in ("PrimitiveString", "PrimitiveStringMultiline"): + return resolve(inputs.get("value")) + if class_type in ("ShowText|pysssss", "Text Multiline", "String Literal"): + return resolve(inputs.get("text", inputs.get("string"))) + if class_type == "StringConcatenate": + a, b = resolve(inputs.get("string_a", "")), resolve(inputs.get("string_b", "")) + delimiter = inputs.get("delimiter", "") + if a is None or b is None or not isinstance(delimiter, str): + return None + return f"{a}{delimiter}{b}" + if class_type == "Text Concatenate": + keys = sorted(k for k in inputs if re.fullmatch(r"text_[a-z]", k)) + return _join_text_inputs(inputs, keys, resolve) return None -def is_text_linked(graph: Any, node_id: str) -> bool: - """True when the node's `text` comes from another node (e.g. PromptSearchList).""" - node = graph.get(str(node_id)) if isinstance(graph, dict) else None - return _is_link(_inputs(node).get("text")) +def resolve_text(graph: Any, node_id: str) -> Optional[str]: + """A prompt node's `text` as it will be at run time, if knowable from the graph. + + Typed text, or text produced by known pure string nodes (core primitives and + StringConcatenate, WAS Text Concatenate/Multiline, pysssss ShowText). Anything + else, including PromptSearchList batches, returns None rather than a guess. + Callers compare the result with the text the node actually received. + """ + if not isinstance(graph, dict): + return None + seen = set() + + def resolve(value: Any) -> Optional[str]: + if isinstance(value, str): + return value + if not _is_link(value) or len(seen) >= MAX_VISITS: + return None + source_id = str(value[0]) + if source_id in seen: + return None + seen.add(source_id) + source = graph.get(source_id) + if not isinstance(source, dict): + return None + return _resolve_string_node( + str(source.get("class_type", "")), _inputs(source), resolve + ) + + seen.add(str(node_id)) + text = resolve(_inputs(graph.get(str(node_id))).get("text")) + return text.strip() if isinstance(text, str) and text.strip() else None diff --git a/utils/usage_tracking.py b/utils/usage_tracking.py index 034f812..2052848 100644 --- a/utils/usage_tracking.py +++ b/utils/usage_tracking.py @@ -6,23 +6,23 @@ workflow instead: it counts a run for each PromptManager node feeding a sampler' positive input. Prompts not yet in the database are remembered as pending and counted by the node when it first saves them. -Text supplied by another node (PromptSearchList batches) is unknown at queue time; -those runs are counted by the node as each item executes. +Text from known pure string nodes is resolved at queue time too. Anything else +(PromptSearchList batches, unknown nodes) is counted by the node as it executes. """ import threading import time from collections import OrderedDict -from typing import Any +from typing import Any, Tuple try: from .hashing import generate_prompt_hash from .logging_config import get_logger - from .prompt_graph import literal_text, run_prompt_nodes + from .prompt_graph import resolve_text, run_prompt_nodes except ImportError: from utils.hashing import generate_prompt_hash from utils.logging_config import get_logger - from utils.prompt_graph import literal_text, run_prompt_nodes + from utils.prompt_graph import resolve_text, run_prompt_nodes logger = get_logger("prompt_manager.usage_tracking") @@ -31,7 +31,11 @@ PENDING_MAX_ENTRIES = 1000 class PendingFirstUse: - """Hashes of queued positive prompts that were not in the database yet.""" + """Queued runs of positive prompts that were not in the database yet. + + Counts per hash: the same new prompt can be queued several times before its + first job runs, and later jobs may reuse ComfyUI's cached node. + """ def __init__( self, @@ -40,21 +44,21 @@ class PendingFirstUse: ): self._ttl = ttl_seconds self._max = max_entries - self._entries: "OrderedDict[str, float]" = OrderedDict() + self._entries: "OrderedDict[str, Tuple[int, float]]" = OrderedDict() self._lock = threading.Lock() def add(self, prompt_hash: str) -> None: with self._lock: - self._entries.pop(prompt_hash, None) - self._entries[prompt_hash] = time.monotonic() + count, _ = self._entries.pop(prompt_hash, (0, 0.0)) + self._entries[prompt_hash] = (count + 1, time.monotonic()) while len(self._entries) > self._max: self._entries.popitem(last=False) - def consume(self, prompt_hash: str) -> bool: - """True (once) if this hash was queued recently and not yet counted.""" + def consume(self, prompt_hash: str) -> int: + """Number of recent queued runs not yet counted (0 if none); clears them.""" with self._lock: - added = self._entries.pop(prompt_hash, None) - return added is not None and time.monotonic() - added <= self._ttl + count, added = self._entries.pop(prompt_hash, (0, 0.0)) + return count if count and time.monotonic() - added <= self._ttl else 0 # Shared between the queue hook and the nodes @@ -67,7 +71,7 @@ def handle_queued_prompt(json_data: Any, db: Any, pending: PendingFirstUse) -> A graph = json_data.get("prompt") if isinstance(json_data, dict) else None counted = set() for node_id in run_prompt_nodes(graph): - text = literal_text(graph, node_id) + text = resolve_text(graph, node_id) if not text: continue prompt_hash = generate_prompt_hash(text) @@ -84,7 +88,7 @@ def handle_queued_prompt(json_data: Any, db: Any, pending: PendingFirstUse) -> A return json_data -_registered = False +_hook_registered = threading.Event() _register_lock = threading.Lock() @@ -92,13 +96,12 @@ def register_queue_hook( server: Any, db_factory: Any, pending: PendingFirstUse = PENDING_FIRST_USE ) -> bool: """Register the on-prompt handler once; db_factory is called lazily per request.""" - global _registered with _register_lock: - if _registered: + if _hook_registered.is_set(): return False server.add_on_prompt_handler( lambda json_data: handle_queued_prompt(json_data, db_factory(), pending) ) - _registered = True + _hook_registered.set() logger.info("Registered prompt usage hook") return True