249 lines
9.9 KiB
Python
249 lines
9.9 KiB
Python
"""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)
|