Vito Sansevero
f6432fec74
Removes helpers nothing calls, test-only wrappers and their tests, the unused config writer, three unused images, and the prompt registry the nodes wrote to but nothing read, which grew with every run.
290 lines
12 KiB
Python
290 lines
12 KiB
Python
"""Paired node encoding preserves recipe storage and output roles."""
|
|
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import Mock
|
|
|
|
try:
|
|
from tests.open_handles import assert_closed
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from open_handles import assert_closed
|
|
|
|
from database.operations import PromptDatabase
|
|
from prompt_manager_pair import PromptManagerPair
|
|
from utils.hashing import generate_prompt_hash
|
|
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.addCleanup(assert_closed, self, self.tmp.name)
|
|
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._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.assertEqual(
|
|
self.db.get_prompt_by_id(pair["positive_prompt_id"])["category"],
|
|
"Animals",
|
|
)
|
|
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 save(self, text="cat", negative="blur", **widgets):
|
|
self.node.encode_prompt(
|
|
self.clip,
|
|
text,
|
|
negative_text=negative,
|
|
prompt=pair_graph(text, negative),
|
|
unique_id="1",
|
|
**widgets,
|
|
)
|
|
hash_of = generate_prompt_hash
|
|
return (
|
|
self.db.get_prompt_by_hash(hash_of(text)),
|
|
self.db.get_prompt_by_hash(hash_of(negative)),
|
|
)
|
|
|
|
def test_first_save_gives_the_positive_its_category_and_tags(self):
|
|
positive, negative = self.save(category=" Animals ", tags="studio, macro")
|
|
self.assertEqual(positive["category"], "Animals")
|
|
self.assertCountEqual(positive["tags"], ["studio", "macro"])
|
|
self.assertIsNone(negative["category"])
|
|
self.assertEqual(negative["tags"], [])
|
|
|
|
def test_a_later_save_never_changes_a_category_or_tags_already_set(self):
|
|
self.save(category="Animals", tags="studio")
|
|
positive, _ = self.save(negative="noise", category="Pets", tags="other")
|
|
self.assertEqual(positive["category"], "Animals")
|
|
self.assertEqual(positive["tags"], ["studio"])
|
|
|
|
def test_a_save_keeps_metadata_a_person_set_on_the_prompt(self):
|
|
text = "cat"
|
|
self.db.save_prompt(
|
|
text,
|
|
category="Mine",
|
|
tags=["kept"],
|
|
prompt_hash=generate_prompt_hash(text),
|
|
)
|
|
positive, _ = self.save(category="Animals", tags="studio")
|
|
self.assertEqual(positive["category"], "Mine")
|
|
self.assertEqual(positive["tags"], ["kept"])
|
|
|
|
def test_a_run_on_an_existing_empty_prompt_changes_nothing(self):
|
|
self.db.save_prompt("cat", prompt_hash=generate_prompt_hash("cat"))
|
|
before = self.db.get_prompt_by_hash(generate_prompt_hash("cat"))["updated_at"]
|
|
positive, _ = self.save(category="Animals", tags="studio")
|
|
self.assertIsNone(positive["category"])
|
|
self.assertEqual(positive["tags"], [])
|
|
self.assertEqual(positive["updated_at"], before)
|
|
|
|
def test_a_pair_saved_without_metadata_stays_empty_on_a_later_run(self):
|
|
self.save()
|
|
positive, _ = self.save(negative="noise", category="Animals", tags="studio")
|
|
self.assertIsNone(positive["category"])
|
|
self.assertEqual(positive["tags"], [])
|
|
|
|
def test_a_cleared_category_and_tags_stay_cleared(self):
|
|
positive, _ = self.save(category="Animals", tags="studio")
|
|
self.db.update_prompt_metadata(positive["id"], category="", tags=[])
|
|
positive, _ = self.save(negative="noise", category="Animals", tags="studio")
|
|
self.assertEqual(positive["category"], "")
|
|
self.assertEqual(positive["tags"], [])
|
|
|
|
def test_a_prompt_first_seen_as_a_negative_stays_empty(self):
|
|
self.save(text="dog", negative="cat")
|
|
positive, _ = self.save(text="cat", negative="blur", category="Animals")
|
|
self.assertIsNone(positive["category"])
|
|
|
|
def test_blank_widgets_leave_the_prompt_empty(self):
|
|
positive, _ = self.save(category=" ", tags=" ")
|
|
self.assertIsNone(positive["category"])
|
|
self.assertEqual(positive["tags"], [])
|
|
|
|
def test_a_filled_prompt_does_not_change_the_saved_pair(self):
|
|
self.save(category="Animals", tags="studio")
|
|
positive, _ = self.save(category="Pets", tags="other")
|
|
pairs = self.db.list_prompt_pairs()["pairs"]
|
|
self.assertEqual(
|
|
[(pair["category"], pair["tags"]) for pair in pairs],
|
|
[("Animals", ["studio"])],
|
|
)
|
|
self.assertEqual(positive["category"], "Animals")
|
|
|
|
def test_library_shows_filled_metadata_without_repeating_a_tag(self):
|
|
self.save(category="Animals", tags="studio")
|
|
item = self.db.list_library()["items"][0]
|
|
self.assertEqual(item["category"], "Animals")
|
|
self.assertEqual(item["tags"], ["studio"])
|
|
tag = self.db.list_vocabulary_tags(limit=5)["tags"][0]
|
|
self.assertEqual(tag["prompt_count"], 1)
|
|
|
|
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(self):
|
|
with self.assertRaises(ValueError):
|
|
self.node.encode_prompt(self.clip, [], negative_text="blur")
|
|
result = self.node.encode_prompt(self.clip, "cat", negative_text="blur")
|
|
self.assertEqual(result[2:], ("cat LORA", "blur"))
|
|
|
|
def test_failed_negative_encoding_cannot_publish_positive_runtime_proof(self):
|
|
from contextlib import ExitStack
|
|
import sys
|
|
import types
|
|
from utils.execution_context_capture import ContextCapture
|
|
from utils.execution_context_budget import ContextBudget
|
|
from utils.execution_history_events import RunStart
|
|
from utils.execution_history_writer import HistoryWriterStatus
|
|
from unittest.mock import patch
|
|
|
|
identity = "11111111-1111-4111-8111-111111111111"
|
|
state = HistoryWriterStatus(identity, True, 0, 0, True, 1)
|
|
budget = ContextBudget()
|
|
manager = ContextCapture(1, lambda: state, budget)
|
|
self.addCleanup(manager.revoke_all)
|
|
cap = manager.begin(
|
|
RunStart(identity, identity, 1, identity, identity, 0),
|
|
1,
|
|
pair_graph(),
|
|
None,
|
|
)
|
|
host = types.ModuleType("comfy_execution.utils")
|
|
host.get_executing_context = lambda: types.SimpleNamespace(
|
|
prompt_id=identity, node_id="1", list_index=0
|
|
)
|
|
self.clip.encode_from_tokens_scheduled.side_effect = [
|
|
("conditioning", "cat LORA"),
|
|
RuntimeError("negative failed"),
|
|
]
|
|
with ExitStack() as stack:
|
|
stack.enter_context(
|
|
patch.dict(sys.modules, {"comfy_execution.utils": host})
|
|
)
|
|
with self.assertRaisesRegex(RuntimeError, "negative failed"):
|
|
self.node.encode_prompt(self.clip, "cat", negative_text="blur")
|
|
sealed = manager.seal(cap)
|
|
if sealed.lease:
|
|
self.addCleanup(sealed.lease.release)
|
|
self.assertIsNone(sealed.bundle.contexts[0].positive.effective.value)
|
|
self.assertIsNone(sealed.bundle.contexts[0].negative.effective.value)
|
|
self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
|