Vito Sansevero
205cf153ea
get_recent_prompts returns {"prompts": [...]} but _get_fallback_prompt
indexed it like a list, so the last-resort link to the most recent prompt
always failed silently. Accept both shapes.
248 lines
8.9 KiB
Python
248 lines
8.9 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):
|
|
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()
|