diff --git a/__init__.py b/__init__.py index 86d9d6e..ee90797 100644 --- a/__init__.py +++ b/__init__.py @@ -111,6 +111,21 @@ except Exception as e: except Exception: print(f"[ComfyUI-PromptManager] Warning: Failed to start image monitoring: {e}") +# Count prompt runs when workflows are queued (ComfyUI skips unchanged nodes, so +# node execution alone misses re-runs). Separate from monitoring so either can fail. +try: + from .py.config import server_instance + from .utils.usage_tracking import register_queue_hook + + register_queue_hook(server_instance, lambda: _global_db) +except Exception as e: + try: + from .utils.logging_config import get_logger + + 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}") + __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] # Print startup message with loaded tools diff --git a/prompt_manager.py b/prompt_manager.py index ffbec23..9637c7c 100644 --- a/prompt_manager.py +++ b/prompt_manager.py @@ -4,7 +4,7 @@ with persistent prompt storage and search capabilities. """ import time -from typing import Any, Tuple +from typing import Any, Dict, Optional, Tuple try: from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict @@ -94,6 +94,12 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): }, ), }, + "hidden": { + # The queued graph and this node's id: used to tell the positive + # prompt of a run from negative-prompt PromptManager nodes + "prompt": "PROMPT", + "unique_id": "UNIQUE_ID", + }, } RETURN_TYPES = (IO.CONDITIONING, IO.STRING) @@ -119,6 +125,8 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): search_text: str = "", prepend_text: str = "", append_text: str = "", + prompt: Optional[Dict[str, Any]] = None, + unique_id: Optional[str] = None, ) -> Tuple[Any]: """ Encode the text prompt and save it to the database. @@ -181,30 +189,24 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): extended_tags.append(f"append:{append_text.strip()[:50]}") try: - prompt_id = self._save_prompt_to_database( + prompt_id = self._track_prompt_execution( text=storage_text.strip(), + encoding_text=encoding_text.strip(), category=category.strip() if category else None, tags=extended_tags if extended_tags else None, + additional_data={ + "category": category.strip() if category else None, + "tags": extended_tags, + "prepend_text": ( + prepend_text.strip() if prepend_text else None + ), + "append_text": append_text.strip() if append_text else None, + "final_text": encoding_text.strip(), + }, + prompt_graph=prompt, + unique_id=unique_id, ) - - # Set current prompt for image tracking - if prompt_id: - execution_id = self.prompt_tracker.set_current_prompt( - prompt_text=encoding_text.strip(), - additional_data={ - "category": category.strip() if category else None, - "tags": extended_tags, - "prompt_id": prompt_id, - "prepend_text": ( - prepend_text.strip() if prepend_text else None - ), - "append_text": append_text.strip() if append_text else None, - "final_text": encoding_text.strip(), - }, - ) - self.logger.debug( - f"Set execution context: {execution_id} for prompt ID: {prompt_id}" - ) + self.logger.debug(f"Tracked prompt ID: {prompt_id}") except Exception as e: # Log error but don't fail the encoding diff --git a/prompt_manager_base.py b/prompt_manager_base.py index 38e9dd0..3d05ce0 100644 --- a/prompt_manager_base.py +++ b/prompt_manager_base.py @@ -24,6 +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.usage_tracking import PENDING_FIRST_USE except ImportError: import sys @@ -32,6 +34,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.usage_tracking import PENDING_FIRST_USE class PromptManagerBase: @@ -56,14 +60,23 @@ class PromptManagerBase: self.logger.debug(f"{self.__class__.__name__} node initialization completed") def _save_prompt_to_database( - self, text: str, category: Optional[str] = None, tags: Optional[list] = None + self, + text: str, + category: Optional[str] = None, + tags: Optional[list] = None, + count_run: bool = False, ) -> Optional[int]: """Save the prompt to the SQLite database. + Runs are normally counted by the queue hook (utils/usage_tracking.py), so a + save only counts one when asked to (batch items whose text the hook could not + see) or when the hook queued this prompt before it existed. + Args: text: The prompt text category: Optional category tags: List of tags + count_run: Count this execution as a run of the prompt Returns: The prompt ID if saved successfully, None otherwise @@ -82,7 +95,8 @@ class PromptManagerBase: prompt_id=existing["id"], category=category, tags=tags ) self.logger.debug("Updated metadata for existing prompt") - self._record_use(existing["id"]) + if count_run: + self._record_use(existing["id"]) return existing["id"] self.logger.debug( @@ -94,7 +108,9 @@ class PromptManagerBase: if prompt_id: self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}") - self._record_use(prompt_id) + pending = getattr(self, "pending_first_use", PENDING_FIRST_USE) + if count_run or pending.consume(prompt_hash): + self._record_use(prompt_id) else: self.logger.warning("Failed to save prompt - no ID returned") @@ -104,6 +120,53 @@ 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]: + """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) + 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} + node_id = str(unique_id) + return { + "is_run_prompt": node_id in run_prompt_nodes(prompt_graph), + "text_linked": is_text_linked(prompt_graph, node_id), + "known": True, + } + + def _track_prompt_execution( + self, + text: str, + encoding_text: str, + category: Optional[str], + tags: Optional[list], + additional_data: Dict[str, Any], + prompt_graph: Any = None, + unique_id: Any = None, + ) -> Optional[int]: + """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. + """ + role = self._prompt_role(prompt_graph, unique_id) + batch_item = role["is_run_prompt"] and role["text_linked"] + prompt_id = self._save_prompt_to_database( + text=text, category=category, tags=tags, count_run=batch_item + ) + if prompt_id and role["is_run_prompt"]: + self.prompt_tracker.set_current_prompt( + prompt_text=encoding_text, + additional_data={**additional_data, "prompt_id": prompt_id}, + push_to_queue=batch_item or not role["known"], + ) + return prompt_id + def _record_use(self, prompt_id: int) -> None: """Count this execution; a tracking failure must never block generation.""" try: diff --git a/prompt_manager_text.py b/prompt_manager_text.py index 73bc061..37df856 100644 --- a/prompt_manager_text.py +++ b/prompt_manager_text.py @@ -4,7 +4,7 @@ without CLIP encoding, while maintaining all database and search features. """ import time -from typing import Tuple +from typing import Any, Dict, Optional, Tuple try: from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict @@ -94,6 +94,12 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): }, ), }, + "hidden": { + # The queued graph and this node's id: used to tell the positive + # prompt of a run from negative-prompt PromptManager nodes + "prompt": "PROMPT", + "unique_id": "UNIQUE_ID", + }, } RETURN_TYPES = (IO.STRING,) @@ -117,6 +123,8 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): search_text: str = "", prepend_text: str = "", append_text: str = "", + prompt: Optional[Dict[str, Any]] = None, + unique_id: Optional[str] = None, ) -> Tuple[str]: """ Process the text prompt and save it to the database. @@ -162,29 +170,23 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): extended_tags.append(f"append:{append_text.strip()[:50]}") try: - prompt_id = self._save_prompt_to_database( + prompt_id = self._track_prompt_execution( text=storage_text.strip(), + encoding_text=final_text.strip(), category=category.strip() if category else None, tags=extended_tags if extended_tags else None, + additional_data={ + "category": category.strip() if category else None, + "tags": extended_tags, + "prepend_text": ( + prepend_text.strip() if prepend_text else None + ), + "append_text": append_text.strip() if append_text else None, + }, + prompt_graph=prompt, + unique_id=unique_id, ) - - # Set current prompt for image tracking - if prompt_id: - execution_id = self.prompt_tracker.set_current_prompt( - prompt_text=final_text.strip(), - additional_data={ - "category": category.strip() if category else None, - "tags": extended_tags, - "prompt_id": prompt_id, - "prepend_text": ( - prepend_text.strip() if prepend_text else None - ), - "append_text": append_text.strip() if append_text else None, - }, - ) - self.logger.debug( - f"Set execution context: {execution_id} for prompt ID: {prompt_id}" - ) + self.logger.debug(f"Tracked prompt ID: {prompt_id}") except Exception as e: # Log error but don't fail the processing diff --git a/tests/test_image_linking.py b/tests/test_image_linking.py new file mode 100644 index 0000000..167b7a7 --- /dev/null +++ b/tests/test_image_linking.py @@ -0,0 +1,131 @@ +"""Image linking picks the positive prompt, preferring the image's own metadata.""" + +import os +import sys +import tempfile +import unittest +from unittest import mock + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from database.operations import PromptDatabase +from utils.hashing import generate_prompt_hash +from utils.image_monitor import ImageGenerationHandler +from utils.prompt_tracker import PromptTracker + +LOADER = { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "m.safetensors"}, +} + + +def graph(positive_text="a cat", negative_text="blurry"): + # Negative listed first: dict order must not decide which prompt an image gets + return { + "176": {"class_type": "PromptManager", "inputs": {"text": negative_text}}, + "134": {"class_type": "PromptManager", "inputs": {"text": positive_text}}, + "5": { + "class_type": "KSampler", + "inputs": { + "positive": ["134", 0], + "negative": ["176", 0], + "model": ["20", 0], + }, + }, + "20": LOADER, + } + + +class LinkingTestCase(unittest.TestCase): + def setUp(self): + fd, self.db_path = tempfile.mkstemp(suffix=".db") + os.close(fd) + fd, self.image = tempfile.mkstemp(suffix=".png") + os.close(fd) + self.db = PromptDatabase(self.db_path) + self.tracker = PromptTracker(self.db) + self.handler = ImageGenerationHandler(self.db, self.tracker) + self.linked = [] + self.handler.link_image_to_prompt = ( + lambda path, prompt, meta: self.linked.append(prompt["id"]) + ) + + def tearDown(self): + for path in ( + self.image, + self.db_path, + self.db_path + "-wal", + self.db_path + "-shm", + ): + if os.path.exists(path): + os.unlink(path) + + def _save(self, text): + return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text)) + + def _process(self, metadata): + with mock.patch.object( + self.handler.metadata_extractor, "extract_metadata", return_value=metadata + ): + self.handler.process_new_image(self.image) + return self.linked[-1] if self.linked else None + + +class TestMetadataLinking(LinkingTestCase): + def test_links_the_positive_prompt_not_the_first_prompt_manager_node(self): + positive = self._save("a cat") + self._save("blurry") + self.assertEqual(self._process({"prompt": graph()}), positive) + + def test_image_metadata_wins_over_a_stale_queue_entry(self): + positive = self._save("a cat") + stale = self._save("previous run") + self.tracker.set_current_prompt("previous run", {"prompt_id": stale}) + + self.assertEqual(self._process({"prompt": graph()}), positive) + self.assertEqual( + self.tracker.pop_next_prompt()["id"], stale, "queue left intact" + ) + + def test_batch_items_with_linked_text_use_the_queue(self): + first, second = self._save("item one"), self._save("item two") + batch_graph = graph() + batch_graph["134"]["inputs"]["text"] = ["7", 0] + batch_graph["7"] = {"class_type": "PromptSearchList", "inputs": {}} + self.tracker.set_current_prompt("item one", {"prompt_id": first}) + self.tracker.set_current_prompt("item two", {"prompt_id": second}) + + self.assertEqual(self._process({"prompt": batch_graph}), first) + self.assertEqual(self._process({"prompt": batch_graph}), second) + + def test_api_graph_present_skips_the_role_unaware_workflow_fallback(self): + negative = self._save("blurry") + batch_graph = graph() + batch_graph["134"]["inputs"]["text"] = ["7", 0] + metadata = { + "prompt": batch_graph, + "text_encoder_nodes": [ + {"type": "PromptManager", "inputs": [], "widgets_values": ["blurry"]} + ], + } + self.assertIsNone(self.handler._find_prompt_from_metadata(metadata)) + self.assertNotEqual(self._process(metadata), negative) + + def test_workflow_fallback_still_used_without_an_api_graph(self): + positive = self._save("from workflow only") + metadata = { + "text_encoder_nodes": [ + { + "type": "PromptManager", + "inputs": [], + "widgets_values": ["from workflow only"], + } + ] + } + self.assertEqual( + self.handler._find_prompt_from_metadata(metadata)["id"], positive + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_node_tracking_wiring.py b/tests/test_node_tracking_wiring.py new file mode 100644 index 0000000..fe3ae9a --- /dev/null +++ b/tests/test_node_tracking_wiring.py @@ -0,0 +1,82 @@ +"""Nodes pass ComfyUI's queued graph to role-aware tracking; the hook registers once.""" + +import os +import sys +import unittest +from unittest import mock + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from prompt_manager import PromptManager +from prompt_manager_text import PromptManagerText +from utils import usage_tracking + +GRAPH = {"134": {"class_type": "PromptManager", "inputs": {"text": "a cat"}}} + + +def bare(node_cls): + """A node without its database, tracker or gallery side effects.""" + node = node_cls.__new__(node_cls) + node.logger = __import__("logging").getLogger("test.wiring") + node.comfyui_integration = mock.Mock() + node._inject_lora_trigger_words = lambda text: text + node._track_prompt_execution = mock.Mock(return_value=42) + return node + + +class TestHiddenInputs(unittest.TestCase): + def test_both_nodes_request_the_queued_graph_and_their_id(self): + for node_cls in (PromptManager, PromptManagerText): + hidden = node_cls.INPUT_TYPES().get("hidden", {}) + self.assertEqual(hidden.get("prompt"), "PROMPT", node_cls.__name__) + self.assertEqual(hidden.get("unique_id"), "UNIQUE_ID", node_cls.__name__) + + +class TestNodesForwardGraph(unittest.TestCase): + def test_prompt_manager_forwards_graph_and_texts(self): + node = bare(PromptManager) + clip = mock.Mock() + node.encode_prompt( + clip, "a cat", prepend_text="best,", prompt=GRAPH, unique_id="134" + ) + + kwargs = node._track_prompt_execution.call_args.kwargs + self.assertEqual(kwargs["prompt_graph"], GRAPH) + self.assertEqual(kwargs["unique_id"], "134") + self.assertEqual(kwargs["text"], "a cat") + self.assertEqual(kwargs["encoding_text"], "best, a cat") + self.assertEqual(kwargs["additional_data"]["final_text"], "best, a cat") + + def test_prompt_manager_text_forwards_graph(self): + node = bare(PromptManagerText) + node.process_text("a castle", prompt=GRAPH, unique_id="7") + + kwargs = node._track_prompt_execution.call_args.kwargs + self.assertEqual((kwargs["prompt_graph"], kwargs["unique_id"]), (GRAPH, "7")) + self.assertEqual(kwargs["text"], "a castle") + + def test_nodes_still_work_without_graph_context(self): + node = bare(PromptManager) + node.encode_prompt(mock.Mock(), "a cat") + kwargs = node._track_prompt_execution.call_args.kwargs + self.assertIsNone(kwargs["prompt_graph"]) + self.assertIsNone(kwargs["unique_id"]) + + +class TestQueueHookRegistration(unittest.TestCase): + def test_registers_once_and_handler_passes_requests_through(self): + server = mock.Mock() + db = mock.Mock() + db.get_prompt_by_hash.return_value = None + with mock.patch.object(usage_tracking, "_registered", False): + self.assertTrue(usage_tracking.register_queue_hook(server, lambda: db)) + self.assertFalse(usage_tracking.register_queue_hook(server, lambda: db)) + + server.add_on_prompt_handler.assert_called_once() + handler = server.add_on_prompt_handler.call_args.args[0] + request = {"prompt": GRAPH} + self.assertIs(handler(request), request) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_prompt_graph.py b/tests/test_prompt_graph.py new file mode 100644 index 0000000..7aa96c3 --- /dev/null +++ b/tests/test_prompt_graph.py @@ -0,0 +1,191 @@ +"""Tests for utils.prompt_graph: which PromptManager nodes feed a positive input.""" + +import os +import sys +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, + 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), []) + + +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")) + + def test_missing_nodes(self): + self.assertIsNone(literal_text({}, "1")) + self.assertFalse(is_text_linked({}, "1")) + self.assertIsNone(literal_text({"1": pm(" ")}, "1")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_prompt_usage.py b/tests/test_prompt_usage.py index 2980f7d..41936fe 100644 --- a/tests/test_prompt_usage.py +++ b/tests/test_prompt_usage.py @@ -194,7 +194,11 @@ class TestUsageMigration(unittest.TestCase): class TestNodeRecordsUse(UsageTestCase): - """Every node execution counts, including re-runs of an existing prompt.""" + """Node saves count a run only when asked; re-runs are counted by the queue hook. + + ComfyUI skips unchanged nodes, so counting at execution misses re-runs; see + tests/test_usage_tracking.py for the hook that counts them. + """ def _node(self): node = PromptManagerBase.__new__(PromptManagerBase) @@ -202,13 +206,19 @@ class TestNodeRecordsUse(UsageTestCase): node.logger = __import__("logging").getLogger("test.prompt_usage") return node - def test_new_and_repeated_runs_are_counted(self): + def test_save_without_count_run_leaves_counting_to_the_queue_hook(self): node = self._node() first = node._save_prompt_to_database("a lighthouse at dusk") second = node._save_prompt_to_database("a lighthouse at dusk") self.assertEqual(first, second) - self.assertEqual(self.db.get_prompt_by_id(first)["run_count"], 2) + self.assertEqual(self.db.get_prompt_by_id(first)["run_count"], 0) + + def test_count_run_counts_new_and_existing_prompts(self): + node = self._node() + pid = node._save_prompt_to_database("batch item", count_run=True) + node._save_prompt_to_database("batch item", count_run=True) + self.assertEqual(self.db.get_prompt_by_id(pid)["run_count"], 2) if __name__ == "__main__": diff --git a/tests/test_usage_tracking.py b/tests/test_usage_tracking.py new file mode 100644 index 0000000..664b43a --- /dev/null +++ b/tests/test_usage_tracking.py @@ -0,0 +1,228 @@ +"""Run counting and node roles: only positive PromptManager nodes count and link. + +Runs are counted when a prompt is queued (ComfyUI's on-prompt hook), because +ComfyUI skips unchanged nodes, so node execution can't be relied on to count. +""" + +import os +import sys +import tempfile +import unittest +from unittest import mock + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from database.operations import PromptDatabase +from prompt_manager_base import PromptManagerBase +from utils.hashing import generate_prompt_hash +from utils.prompt_tracker import PromptTracker +from utils.usage_tracking import PendingFirstUse, handle_queued_prompt + +LOADER = { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "m.safetensors"}, +} + + +def workflow(positive_text="a cat", negative_text="blurry"): + """API graph with a positive (#134) and a negative (#176) PromptManager node.""" + return { + "176": {"class_type": "PromptManager", "inputs": {"text": negative_text}}, + "134": {"class_type": "PromptManager", "inputs": {"text": positive_text}}, + "5": { + "class_type": "KSampler", + "inputs": { + "positive": ["134", 0], + "negative": ["176", 0], + "model": ["20", 0], + }, + }, + "20": LOADER, + } + + +class DbTestCase(unittest.TestCase): + def setUp(self): + fd, self.path = tempfile.mkstemp(suffix=".db") + os.close(fd) + self.db = PromptDatabase(self.path) + self.pending = PendingFirstUse() + + def tearDown(self): + for suffix in ("", "-wal", "-shm"): + if os.path.exists(self.path + suffix): + os.unlink(self.path + suffix) + + def _save(self, text): + return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text)) + + def _runs(self, prompt_id): + return self.db.get_prompt_by_id(prompt_id)["run_count"] + + def _queue(self, graph): + return handle_queued_prompt( + {"prompt": graph, "client_id": "x"}, self.db, self.pending + ) + + +class TestQueueHook(DbTestCase): + def test_counts_the_positive_prompt_only(self): + positive, negative = self._save("a cat"), self._save("blurry") + self._queue(workflow()) + self.assertEqual(self._runs(positive), 1) + self.assertEqual(self._runs(negative), 0) + + def test_requeue_counts_again_even_when_comfyui_skips_the_cached_node(self): + positive = self._save("a cat") + self._queue(workflow()) + self._queue(workflow()) # no node execution in between + self.assertEqual(self._runs(positive), 2) + + def test_new_prompt_is_counted_when_the_node_first_saves_it(self): + self._queue(workflow(positive_text="a brand new prompt")) + self.assertTrue( + self.pending.consume(generate_prompt_hash("a brand new prompt")) + ) + self.assertFalse( + self.pending.consume(generate_prompt_hash("a brand new prompt")) + ) + + def test_same_text_in_two_positive_nodes_counts_once(self): + positive = self._save("a cat") + graph = workflow() + graph["135"] = {"class_type": "PromptManager", "inputs": {"text": "a cat"}} + graph["4"] = { + "class_type": "ConditioningCombine", + "inputs": {"conditioning_1": ["134", 0], "conditioning_2": ["135", 0]}, + } + graph["5"]["inputs"]["positive"] = ["4", 0] + self._queue(graph) + self.assertEqual(self._runs(positive), 1) + + def test_linked_text_is_left_to_node_execution(self): + positive = self._save("from a batch") + graph = workflow() + graph["134"]["inputs"]["text"] = ["7", 0] + graph["7"] = {"class_type": "PromptSearchList", "inputs": {}} + self._queue(graph) + self.assertEqual(self._runs(positive), 0) + + def test_returns_the_request_unchanged_and_never_raises(self): + request = {"prompt": workflow(), "client_id": "x", "extra_data": {}} + self.assertIs(handle_queued_prompt(request, self.db, self.pending), request) + for bad in (None, {}, {"prompt": "x"}, {"prompt": None}, []): + self.assertIs(handle_queued_prompt(bad, self.db, self.pending), bad) + broken_db = mock.Mock() + broken_db.get_prompt_by_hash.side_effect = RuntimeError("db locked") + self.assertIs(handle_queued_prompt(request, broken_db, self.pending), request) + + +class TestPendingFirstUse(unittest.TestCase): + def test_entries_expire(self): + pending = PendingFirstUse(ttl_seconds=10) + with mock.patch("utils.usage_tracking.time.monotonic", return_value=100.0): + pending.add("h") + with mock.patch("utils.usage_tracking.time.monotonic", return_value=111.0): + self.assertFalse(pending.consume("h")) + + def test_size_is_capped(self): + pending = PendingFirstUse(max_entries=3) + for h in ("a", "b", "c", "d"): + pending.add(h) + self.assertFalse(pending.consume("a")) # oldest evicted + self.assertTrue(pending.consume("d")) + + +class FakeTracker: + def __init__(self): + self.calls = [] + + def set_current_prompt(self, prompt_text, additional_data=None, push_to_queue=True): + self.calls.append({"text": prompt_text, "push_to_queue": push_to_queue}) + return "exec-1" + + +class TestNodeRoles(DbTestCase): + def _node(self): + node = PromptManagerBase.__new__(PromptManagerBase) + node.db = self.db + node.prompt_tracker = FakeTracker() + node.logger = __import__("logging").getLogger("test.usage_tracking") + node.pending_first_use = self.pending + return node + + def _run(self, node, graph, unique_id, text): + return node._track_prompt_execution( + text=text, + encoding_text=text, + category=None, + tags=None, + additional_data={}, + prompt_graph=graph, + unique_id=unique_id, + ) + + def test_negative_node_saves_but_never_becomes_the_current_prompt(self): + node = self._node() + prompt_id = self._run(node, workflow(), "176", "blurry") + self.assertIsNotNone(self.db.get_prompt_by_id(prompt_id)) + self.assertEqual(node.prompt_tracker.calls, []) + self.assertEqual(self._runs(prompt_id), 0) + + def test_positive_typed_text_is_current_without_queue_push(self): + node = self._node() + self._run(node, workflow(), "134", "a cat") + self.assertEqual( + node.prompt_tracker.calls, [{"text": "a cat", "push_to_queue": False}] + ) + + def test_first_save_after_queue_hook_counts_one_run(self): + self._queue(workflow(positive_text="never seen before")) + prompt_id = self._run( + self._node(), + workflow(positive_text="never seen before"), + "134", + "never seen before", + ) + self.assertEqual(self._runs(prompt_id), 1) + + def test_existing_prompt_is_not_counted_twice_by_hook_and_node(self): + positive = self._save("a cat") + self._queue(workflow()) + self._run(self._node(), workflow(), "134", "a cat") + self.assertEqual(self._runs(positive), 1) + + def test_positive_linked_text_pushes_queue_and_counts_each_item(self): + graph = workflow() + graph["134"]["inputs"]["text"] = ["7", 0] + graph["7"] = {"class_type": "PromptSearchList", "inputs": {}} + node = self._node() + first = self._run(node, graph, "134", "batch item one") + self._run(node, graph, "134", "batch item one") + self.assertEqual(self._runs(first), 2) + self.assertTrue(all(c["push_to_queue"] for c in node.prompt_tracker.calls)) + + def test_without_graph_context_keeps_legacy_behaviour(self): + node = self._node() + self._run(node, None, None, "standalone") + self.assertEqual( + node.prompt_tracker.calls, [{"text": "standalone", "push_to_queue": True}] + ) + + +class TestTrackerQueueFlag(unittest.TestCase): + def setUp(self): + self.tracker = PromptTracker(mock.Mock()) + + def test_push_to_queue_false_keeps_the_batch_queue_empty(self): + self.tracker.set_current_prompt("a cat", {"prompt_id": 1}, push_to_queue=False) + self.assertIsNone(self.tracker.pop_next_prompt()) + self.assertEqual(self.tracker.get_current_prompt()["id"], 1) + + def test_default_still_pushes(self): + self.tracker.set_current_prompt("a cat", {"prompt_id": 1}) + self.assertEqual(self.tracker.pop_next_prompt()["id"], 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/image_monitor.py b/utils/image_monitor.py index 892cec8..c011721 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -34,6 +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 class ImageGenerationHandler(FileSystemEventHandler): @@ -146,21 +147,26 @@ class ImageGenerationHandler(FileSystemEventHandler): except Exception as meta_error: self.logger.warning(f"Metadata extraction failed: {meta_error}") - # Strategy 1: Pop from batch queue (most reliable for batch workflows). - # Prompts are queued during CLIP encoding in order; images save in - # the same order, so FIFO pop gives the correct prompt per image. - current_prompt = self.prompt_tracker.pop_next_prompt() + # Strategy 1: The image's own metadata. It records the positive prompt + # that produced this image, independent of node caching and of any + # negative-prompt PromptManager nodes in the workflow. + current_prompt = self._find_prompt_from_metadata(metadata) if current_prompt: self.logger.info( - f"Queue match: prompt {current_prompt['id']} for " + f"Metadata match: prompt {current_prompt['id']} for " f"{os.path.basename(image_path)}" ) - # Strategy 2: Find prompt from image metadata + # Strategy 2: Pop from batch queue. Batch items (text supplied by + # PromptSearchList) share one graph, so metadata can't tell them apart; + # they are queued in encoding order and images save in the same order. if not current_prompt: - current_prompt = self._find_prompt_from_metadata(metadata) + current_prompt = self.prompt_tracker.pop_next_prompt() if current_prompt: - self.logger.info(f"Metadata match: prompt {current_prompt['id']}") + self.logger.info( + f"Queue match: prompt {current_prompt['id']} for " + f"{os.path.basename(image_path)}" + ) # Strategy 3: Use snapshot captured at file-creation time if not current_prompt and prompt_snapshot: @@ -231,21 +237,19 @@ class ImageGenerationHandler(FileSystemEventHandler): prompt_text = None - # Try to find PromptManager node in the prompt execution data + # The executed graph: use the PromptManager node feeding the positive input prompt_data = metadata.get("prompt") if isinstance(prompt_data, dict): - for node_id, node_info in prompt_data.items(): - if not isinstance(node_info, dict): - continue - class_type = node_info.get("class_type", "") - if class_type in ("PromptManager", "PromptManagerText"): - inputs = node_info.get("inputs", {}) - text = inputs.get("text", "") - if text and isinstance(text, str) and text.strip(): - prompt_text = text.strip() - break + for node_id in run_prompt_nodes(prompt_data): + prompt_text = literal_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 + return None - # Fallback: check text_encoder_nodes from workflow, but only if + # Fallback (no executed graph): check text_encoder_nodes from workflow, but only if # the text input is NOT connected (connected inputs override widget values, # so the widget value would be stale in batch workflows). if not prompt_text: diff --git a/utils/prompt_graph.py b/utils/prompt_graph.py new file mode 100644 index 0000000..2efd874 --- /dev/null +++ b/utils/prompt_graph.py @@ -0,0 +1,121 @@ +"""Find which PromptManager nodes feed a sampler's positive input in a ComfyUI graph. + +Works on the API-format prompt graph ({node_id: {"class_type", "inputs"}}) that +ComfyUI queues and embeds in saved images. Usage counting and image linking use it +so that negative-prompt PromptManager nodes never count as "the prompt" of a run. + +Pure functions with no ComfyUI imports, so they are unit-testable. Mirrors the +link-tracing rules of web/js/comfy-metadata.js. +""" + +import re +from typing import Any, Dict, List, Optional + +PROMPT_MANAGER_TYPES = frozenset({"PromptManager", "PromptManagerText"}) + +# Upper bound on nodes expanded per graph; queued and embedded graphs are untrusted +MAX_VISITS = 500 + +# Inputs that carry positive conditioning or prompt text towards a sampler +_CONDITIONING_KEY = re.compile(r"^conditioning(_\w+)?$") +_TEXT_KEY = re.compile(r"^(text|string|prompt)(_?[a-z0-9]+)?$", re.IGNORECASE) + + +def _is_link(value: Any) -> bool: + return ( + isinstance(value, (list, tuple)) + and len(value) == 2 + and isinstance(value[0], (str, int)) + and isinstance(value[1], int) + ) + + +def _inputs(node: Any) -> Dict[str, Any]: + inputs = node.get("inputs") if isinstance(node, dict) else None + return inputs if isinstance(inputs, dict) else {} + + +def _positive_roots(graph: Dict[str, Any]) -> List[Any]: + """Links into samplers' positive inputs (KSampler, CFGGuider, BasicGuider, ...).""" + roots = [] + for node in graph.values(): + inputs = _inputs(node) + if not _is_link(inputs.get("model")): + continue + 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"]) + return roots + + +def _upstream_links(inputs: Dict[str, Any]) -> List[Any]: + """Inputs to follow from a pass-through or encoder node, never `negative`.""" + links = [] + for key, value in inputs.items(): + if not _is_link(value): + continue + if key == "positive" or _CONDITIONING_KEY.match(key) or _TEXT_KEY.match(key): + links.append(value) + return links + + +def positive_prompt_nodes(graph: Any) -> List[str]: + """Ids of PromptManager/PromptManagerText nodes that feed a positive input. + + Ordered by discovery, each id once. Malformed graphs yield an empty list. + """ + if not isinstance(graph, dict): + return [] + + found: 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: + continue + seen.add(node_id) + node = graph.get(node_id) + if not isinstance(node, dict): + continue + if node.get("class_type") in PROMPT_MANAGER_TYPES: + found.append(node_id) + continue + stack.extend(reversed(_upstream_links(_inputs(node)))) + return found + + +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. + """ + positive = positive_prompt_nodes(graph) + if positive or not isinstance(graph, dict): + return positive + return [ + str(node_id) + for node_id, node in graph.items() + if isinstance(node, dict) and node.get("class_type") in PROMPT_MANAGER_TYPES + ] + + +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() + 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")) diff --git a/utils/prompt_tracker.py b/utils/prompt_tracker.py index 2c32298..91242dd 100644 --- a/utils/prompt_tracker.py +++ b/utils/prompt_tracker.py @@ -99,7 +99,10 @@ class PromptTracker: self.logger.debug("PromptTracker initialization completed") def set_current_prompt( - self, prompt_text: str, additional_data: Optional[Dict[str, Any]] = None + self, + prompt_text: str, + additional_data: Optional[Dict[str, Any]] = None, + push_to_queue: bool = True, ) -> str: """ Set the current prompt for this thread for image tracking. @@ -113,6 +116,9 @@ class PromptTracker: Args: prompt_text: The prompt text being executed additional_data: Additional prompt metadata, should include prompt_id from PromptManager + push_to_queue: Also push onto the FIFO batch queue. Only batch items whose + text comes from another node need it; typed prompts are linked from + the image's own metadata instead. Returns: Unique execution ID for this prompt execution @@ -164,11 +170,12 @@ class PromptTracker: self.active_prompts[execution_id] = execution_context # Push to batch queue for ordered image linking - with self._queue_lock: - self._prompt_queue.append(execution_context) - self.logger.debug( - f"Queue size: {len(self._prompt_queue)} after push for prompt {prompt_id}" - ) + if push_to_queue: + with self._queue_lock: + self._prompt_queue.append(execution_context) + self.logger.debug( + f"Queue size: {len(self._prompt_queue)} after push for prompt {prompt_id}" + ) self.logger.debug( f"Set current prompt: {execution_id} -> {prompt_text[:50]}... (thread: {threading.current_thread().ident})" diff --git a/utils/usage_tracking.py b/utils/usage_tracking.py new file mode 100644 index 0000000..034f812 --- /dev/null +++ b/utils/usage_tracking.py @@ -0,0 +1,104 @@ +"""Count prompt runs when a workflow is queued, not when nodes execute. + +ComfyUI skips nodes whose inputs are unchanged, so a PromptManager node does not +execute when the same prompt is re-run. The on-prompt hook sees every queued +workflow instead: it counts a run for each PromptManager node feeding a sampler's +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. +""" + +import threading +import time +from collections import OrderedDict +from typing import Any + +try: + from .hashing import generate_prompt_hash + from .logging_config import get_logger + from .prompt_graph import literal_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 + +logger = get_logger("prompt_manager.usage_tracking") + +PENDING_TTL_SECONDS = 3600 # a queued prompt that never executes is forgotten +PENDING_MAX_ENTRIES = 1000 + + +class PendingFirstUse: + """Hashes of queued positive prompts that were not in the database yet.""" + + def __init__( + self, + ttl_seconds: float = PENDING_TTL_SECONDS, + max_entries: int = PENDING_MAX_ENTRIES, + ): + self._ttl = ttl_seconds + self._max = max_entries + self._entries: "OrderedDict[str, 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() + 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.""" + with self._lock: + added = self._entries.pop(prompt_hash, None) + return added is not None and time.monotonic() - added <= self._ttl + + +# Shared between the queue hook and the nodes +PENDING_FIRST_USE = PendingFirstUse() + + +def handle_queued_prompt(json_data: Any, db: Any, pending: PendingFirstUse) -> Any: + """ComfyUI on-prompt handler body: count runs, always return the request as-is.""" + try: + 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) + if not text: + continue + prompt_hash = generate_prompt_hash(text) + if prompt_hash in counted: + continue + counted.add(prompt_hash) + existing = db.get_prompt_by_hash(prompt_hash) + if existing: + db.record_prompt_use(existing["id"]) + else: + pending.add(prompt_hash) + except Exception as e: + logger.warning(f"Could not count queued prompt runs: {e}") + return json_data + + +_registered = False +_register_lock = threading.Lock() + + +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: + return False + server.add_on_prompt_handler( + lambda json_data: handle_queued_prompt(json_data, db_factory(), pending) + ) + _registered = True + logger.info("Registered prompt usage hook") + return True