Vito Sansevero
756509b7ea
- PromptTracker: the cleanup thread waits on a stop event instead of time.sleep; new shutdown(timeout=2) joins it and returns whether it ended. - ImageGenerationHandler.shutdown(timeout=5) now joins the worker and returns a bool; new close() is final: schedule() refuses and a sealed run's history re-check is dropped, so no worker restarts afterwards. - ImageMonitor.stop_monitoring(timeout=5) closes the handler first, then stops and joins the watchdog observer within one deadline. - Every test that builds a real tracker, handler or monitor ends with its final stop and an assertion, before its folder and databases go. - #370: the wrapper/session history test drains before dispose(); the held-write bound test uses 30 s ceilings instead of tight wall clocks.
255 lines
9.3 KiB
Python
255 lines
9.3 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 tests.background_stops import close_handler # noqa: E402
|
|
from database.operations import PromptDatabase
|
|
from utils.hashing import generate_prompt_hash
|
|
from utils.image_monitor import ImageGenerationHandler
|
|
from utils.prompt_tracker import PromptTracker
|
|
from utils.run_keys import prompt_node_key
|
|
|
|
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):
|
|
close_handler(self, self.handler) # and the tracker's cleanup thread
|
|
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": {}}
|
|
key = prompt_node_key(batch_graph, "134")
|
|
for text, prompt_id in (("item one", first), ("item two", second)):
|
|
self.tracker.set_current_prompt(
|
|
text, {"prompt_id": prompt_id}, run_key=key, run_id="run-1"
|
|
)
|
|
|
|
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)
|
|
)
|
|
self._save("most recent prompt") # never picked by recency (#342)
|
|
self.assertIsNone(self._process(metadata))
|
|
self.assertEqual(self.linked, [])
|
|
self.assertIsNotNone(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_left_unlinked(self):
|
|
first = self._queue_batch_item("item one")
|
|
self.assertEqual(
|
|
self.handler._find_prompt_from_metadata(None), ("unknown", None)
|
|
)
|
|
# Nothing proves which run made it (#342): unlinked, queue untouched
|
|
self.assertIsNone(self._process(None))
|
|
self.assertEqual(self.tracker.pop_next_prompt()["id"], 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_is_never_given_the_latest_prompt(self):
|
|
self._save("older")
|
|
self._save("latest")
|
|
self.assertIsNone(self._process(None))
|
|
self.assertIsNone(self._process({"prompt": graph("not saved")}))
|
|
|
|
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()
|