Vito Sansevero
0e8108ae9b
With a negative-prompt PromptManager node in the workflow, every executed PromptManager node pushed onto the image-linking queue, so images could be linked to the negative prompt or to a previous run. ComfyUI also skips unchanged nodes, so re-runs of the same prompt were never counted. - utils/prompt_graph.py finds the PromptManager nodes feeding a sampler's positive input (KSampler, CFG/Basic guiders, ControlNet and combine pass-throughs, PromptManagerText into encoders), falling back to all PromptManager nodes when none can be traced. - Runs are counted by a ComfyUI on-prompt hook when a workflow is queued; prompts not yet saved are counted when the node first saves them. - Nodes receive the queued graph (hidden PROMPT/UNIQUE_ID). Only the run's positive prompt becomes the current prompt; only batch items with linked text are pushed onto the FIFO queue. - The image monitor links from the image's own metadata first, using the positive node, and falls back to the queue for batch items. Refs #90
83 lines
3.2 KiB
Python
83 lines
3.2 KiB
Python
"""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()
|