Vito
4c9cb7f5ea
* refactor: remove dead and test-only code and the write-only prompt registry (#620) 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. * fix: network-share readers, model unload off the loop, bad save requests (#620) Every read-only SQLite open now uses one URI helper, so a library on a Windows network share works in backups, restores, moves and history reads. Unloading a model no longer blocks the server. A malformed /save body gets a 400, not a 500. * feat(library): add Recently used, Most used and Name Z-A sorting (#620) The old Dashboard's sorts the community asked for, on the Prompt library, in saved views, and with paging through the whole list. * feat(workspace): stop linking the old pages (#620) They show a removal notice since 3.5.1. System no longer links them, the load failure offers a reload and ComfyUI instead, and the README says they are going away. The pages still work at their URLs.
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)
|