feat(nodes): encode paired prompts and count cached recipe runs (#193)
This commit is contained in:
Vito Sansevero
@@ -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",
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
@@ -138,6 +138,7 @@ class TestPackageWiring(PackageTestCase):
|
||||
[
|
||||
"PromptManager",
|
||||
"PromptManagerFocused",
|
||||
"PromptManagerPair",
|
||||
"PromptManagerSplit",
|
||||
"PromptManagerText",
|
||||
"PromptSearchList",
|
||||
|
||||
@@ -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)
|
||||
@@ -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"))
|
||||
+79
-5
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+95
-11
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user