Vito Sansevero
144e6b1145
An LLM fed a PromptManagerPair's prepend_text while its text editor was empty. The node refused to save (an empty positive raised and logged a warning), and the image monitor then linked the image to a stale tracker entry or to the newest prompt in the database. - Nodes: an empty text editor saves and tracks nothing (prepend and append text are never the stored prompt), logs at info, and marks the node's run key empty so no earlier entry can be used for its images. - Tracker entries carry a run key: a digest of the prompt node and its whole upstream graph, the identity ComfyUI's cache uses. The monitor computes the same key from the image's embedded graph and uses a queued or snapshot entry only when the keys match (a snapshot also needs a known run with one entry, so a batch's last item is never applied to every image). - The recency fallback is gone; anything uncertain stays unlinked and is logged at info. - Recorded execution history attributes a file when exactly one run reported it, the file was written during that run, and the run's positive contexts name one stored prompt. The monitor and the output scan use it as their last certain source. - The scan and pair ingestion treat a resolved pair with an empty stored positive as unresolved instead of raising. - tools/review_image_links.py reviews selected links by these rules; dry run by default, --apply changes only named ids after a backup. Closes #342
253 lines
9.2 KiB
Python
253 lines
9.2 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
|
|
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):
|
|
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()
|