Vito Sansevero
13d1115406
Windows refuses to unlink an open SQLite file and executor threads keep their own connections, so every owned test that creates a PromptDatabase calls close_all() before cleanup; the diagnostics temp dir ignores cleanup errors.
249 lines
9.0 KiB
Python
249 lines
9.0 KiB
Python
"""Image linking picks the positive prompt, preferring the image's own metadata."""
|
|
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
import unittest.mock as mock
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from database.operations import PromptDatabase
|
|
from utils.hashing import generate_prompt_hash
|
|
from utils.image_monitor import ImageGenerationHandler
|
|
from utils.prompt_tracker import PromptTracker
|
|
|
|
LOADER = {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": "m.safetensors"},
|
|
}
|
|
|
|
|
|
def graph(positive_text="a cat", negative_text="blurry"):
|
|
# Negative listed first: dict order must not decide which prompt an image gets
|
|
return {
|
|
"176": {"class_type": "PromptManager", "inputs": {"text": negative_text}},
|
|
"134": {"class_type": "PromptManager", "inputs": {"text": positive_text}},
|
|
"5": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"positive": ["134", 0],
|
|
"negative": ["176", 0],
|
|
"model": ["20", 0],
|
|
},
|
|
},
|
|
"20": LOADER,
|
|
}
|
|
|
|
|
|
class LinkingTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
fd, self.db_path = tempfile.mkstemp(suffix=".db")
|
|
os.close(fd)
|
|
fd, self.image = tempfile.mkstemp(suffix=".png")
|
|
os.close(fd)
|
|
self.db = PromptDatabase(self.db_path)
|
|
self.tracker = PromptTracker(self.db)
|
|
self.handler = ImageGenerationHandler(self.db, self.tracker)
|
|
self.linked = []
|
|
self.handler.link_image_to_prompt = (
|
|
lambda path, prompt, meta: self.linked.append(prompt["id"])
|
|
)
|
|
|
|
def tearDown(self):
|
|
self.db.close_all() # Windows cannot unlink an open database file
|
|
for path in (
|
|
self.image,
|
|
self.db_path,
|
|
self.db_path + "-wal",
|
|
self.db_path + "-shm",
|
|
):
|
|
if os.path.exists(path):
|
|
os.unlink(path)
|
|
|
|
def _save(self, text):
|
|
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
|
|
|
def _process(self, metadata):
|
|
with mock.patch.object(
|
|
self.handler.metadata_extractor, "extract_metadata", return_value=metadata
|
|
):
|
|
self.handler.process_new_image(self.image)
|
|
return self.linked[-1] if self.linked else None
|
|
|
|
|
|
class TestMetadataLinking(LinkingTestCase):
|
|
def test_links_the_positive_prompt_not_the_first_prompt_manager_node(self):
|
|
positive = self._save("a cat")
|
|
self._save("blurry")
|
|
self.assertEqual(self._process({"prompt": graph()}), positive)
|
|
|
|
def test_image_metadata_wins_over_a_stale_queue_entry(self):
|
|
positive = self._save("a cat")
|
|
stale = self._save("previous run")
|
|
self.tracker.set_current_prompt("previous run", {"prompt_id": stale})
|
|
|
|
self.assertEqual(self._process({"prompt": graph()}), positive)
|
|
self.assertEqual(
|
|
self.tracker.pop_next_prompt()["id"], stale, "queue left intact"
|
|
)
|
|
|
|
def test_text_from_a_known_string_node_links_from_metadata(self):
|
|
positive = self._save("a cat")
|
|
linked_graph = graph()
|
|
linked_graph["134"]["inputs"]["text"] = ["8", 0]
|
|
linked_graph["8"] = {
|
|
"class_type": "PrimitiveString",
|
|
"inputs": {"value": "a cat"},
|
|
}
|
|
self.assertEqual(self._process({"prompt": linked_graph}), positive)
|
|
|
|
def test_batch_items_with_linked_text_use_the_queue(self):
|
|
first, second = self._save("item one"), self._save("item two")
|
|
batch_graph = graph()
|
|
batch_graph["134"]["inputs"]["text"] = ["7", 0]
|
|
batch_graph["7"] = {"class_type": "PromptSearchList", "inputs": {}}
|
|
self.tracker.set_current_prompt("item one", {"prompt_id": first})
|
|
self.tracker.set_current_prompt("item two", {"prompt_id": second})
|
|
|
|
self.assertEqual(self._process({"prompt": batch_graph}), first)
|
|
self.assertEqual(self._process({"prompt": batch_graph}), second)
|
|
|
|
def test_api_graph_present_skips_the_role_unaware_workflow_fallback(self):
|
|
negative = self._save("blurry")
|
|
batch_graph = graph()
|
|
batch_graph["134"]["inputs"]["text"] = ["7", 0]
|
|
metadata = {
|
|
"prompt": batch_graph,
|
|
"text_encoder_nodes": [
|
|
{"type": "PromptManager", "inputs": [], "widgets_values": ["blurry"]}
|
|
],
|
|
}
|
|
self.assertEqual(
|
|
self.handler._find_prompt_from_metadata(metadata), ("unknown", None)
|
|
)
|
|
latest = self._save("most recent prompt") # recency fallback picks this
|
|
self.assertEqual(self._process(metadata), latest)
|
|
self.assertNotEqual(latest, negative)
|
|
|
|
def test_workflow_fallback_still_used_without_an_api_graph(self):
|
|
positive = self._save("from workflow only")
|
|
metadata = {
|
|
"text_encoder_nodes": [
|
|
{
|
|
"type": "PromptManager",
|
|
"inputs": [],
|
|
"widgets_values": ["from workflow only"],
|
|
}
|
|
]
|
|
}
|
|
state, prompt = self.handler._find_prompt_from_metadata(metadata)
|
|
self.assertEqual(state, "linked")
|
|
self.assertEqual(prompt["id"], positive)
|
|
|
|
|
|
def foreign_graph():
|
|
"""A workflow with no PromptManager node at all."""
|
|
return {
|
|
"9": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": "a cat", "clip": ["20", 1]},
|
|
},
|
|
"5": {
|
|
"class_type": "KSampler",
|
|
"inputs": {"positive": ["9", 0], "negative": ["9", 0], "model": ["20", 0]},
|
|
},
|
|
"20": LOADER,
|
|
}
|
|
|
|
|
|
class TestForeignWorkflows(LinkingTestCase):
|
|
"""Images from workflows without a positive PromptManager node are not ours."""
|
|
|
|
def _queue_batch_item(self, text):
|
|
prompt_id = self._save(text)
|
|
self.tracker.set_current_prompt(text, {"prompt_id": prompt_id})
|
|
return prompt_id
|
|
|
|
def test_foreign_image_leaves_the_batch_queue_aligned(self):
|
|
first = self._queue_batch_item("item one")
|
|
self.assertIsNone(self._process({"prompt": foreign_graph()}))
|
|
self.assertEqual(self.linked, [])
|
|
self.assertEqual(self.tracker.pop_next_prompt()["id"], first)
|
|
self.assertIsNone(self.tracker.pop_next_prompt())
|
|
|
|
def test_foreign_image_is_not_linked_by_the_later_fallbacks_either(self):
|
|
self._queue_batch_item("item one") # live tracker and recent prompt exist
|
|
self.assertIsNone(self._process({"prompt": foreign_graph()}))
|
|
self.assertEqual(self.linked, [])
|
|
|
|
def test_negative_only_workflow_is_foreign(self):
|
|
self._save("blurry")
|
|
negative_only = foreign_graph()
|
|
negative_only["176"] = {
|
|
"class_type": "PromptManager",
|
|
"inputs": {"text": "blurry"},
|
|
}
|
|
negative_only["5"]["inputs"]["negative"] = ["176", 0]
|
|
self.assertEqual(
|
|
self.handler._find_prompt_from_metadata({"prompt": negative_only}),
|
|
("foreign", None),
|
|
)
|
|
|
|
def test_workflow_only_metadata_without_prompt_manager_is_foreign(self):
|
|
metadata = {
|
|
"text_encoder_nodes": [
|
|
{"type": "CLIPTextEncode", "inputs": [], "widgets_values": ["a cat"]}
|
|
]
|
|
}
|
|
self.assertEqual(
|
|
self.handler._find_prompt_from_metadata(metadata), ("foreign", None)
|
|
)
|
|
|
|
def test_missing_metadata_is_unknown_and_uses_the_queue(self):
|
|
first = self._queue_batch_item("item one")
|
|
self.assertEqual(
|
|
self.handler._find_prompt_from_metadata(None), ("unknown", None)
|
|
)
|
|
self.assertEqual(self._process(None), first)
|
|
|
|
def test_tri_state_for_linked_unknown_and_foreign(self):
|
|
positive = self._save("a cat")
|
|
state, prompt = self.handler._find_prompt_from_metadata({"prompt": graph()})
|
|
self.assertEqual((state, prompt["id"]), ("linked", positive))
|
|
self.assertEqual(
|
|
self.handler._find_prompt_from_metadata({"prompt": foreign_graph()}),
|
|
("foreign", None),
|
|
)
|
|
self.assertEqual(
|
|
self.handler._find_prompt_from_metadata({"workflow": {}}),
|
|
("unknown", None),
|
|
)
|
|
self.assertEqual(
|
|
self.handler._find_prompt_from_metadata({"prompt": graph("not saved")}),
|
|
("unknown", None),
|
|
)
|
|
|
|
def test_unknown_image_without_any_context_falls_back_to_the_latest_prompt(self):
|
|
self._save("older")
|
|
latest = self._save("latest")
|
|
self.assertEqual(self._process(None), latest)
|
|
|
|
def test_recent_prompt_fallback_handles_an_empty_database(self):
|
|
self.assertIsNone(self.handler._get_fallback_prompt())
|
|
self.assertIsNone(self._process(None))
|
|
|
|
def test_metadata_lookup_errors_are_unknown(self):
|
|
self._save("a cat")
|
|
with mock.patch.object(
|
|
self.db, "get_prompt_by_hash", side_effect=RuntimeError("locked")
|
|
):
|
|
self.assertEqual(
|
|
self.handler._find_prompt_from_metadata({"prompt": graph()}),
|
|
("unknown", None),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|