fix(ingestion): address the #344 review (history re-check, keys, repair)
- Images without an embedded API graph: the scan now asks recorded history before skipping a file with no prompt metadata, and the monitor gives a keyless unlinked image one history-only re-check when its run's outputs are recorded (a seal notification from the outputs write; at most 1000 waiting images, each for an hour). - Run keys: an image is unkeyable when any positive node has no key, so a two-pass graph is never linked through the other pass; links follow ComfyUI's is_link exactly. - A failed save marks the node's run key like an empty editor, so an older entry of that key can't stand in for this run. - A placeholder is the latest entry: the image stays unlinked without falling through to history, the same rule as a popped placeholder. - An index on execution_run_outputs(filename, subfolder, root_key, path_style) for the history lookup, added by migration and accepted by the output schema check (any other index is still refused). - Repair tool: unlink removes only the link record and only while the image is still linked to the prompt it was judged on; every run prints what --apply does. - Real-host case proves the re-run's PromptManager node was cached.
This commit is contained in:
@@ -314,12 +314,16 @@ def migrate_execution_history(conn: sqlite3.Connection) -> None:
|
||||
if version == 1:
|
||||
install_context_schema(conn)
|
||||
install_output_schema(conn)
|
||||
from .execution_output_schema import OUTPUT_READ_INDEX_COLUMNS
|
||||
|
||||
# Read projections need reverse content lookup and the complete run order.
|
||||
for table, columns in (
|
||||
("execution_prompt_snapshots", "raw_text_id,snapshot_id"),
|
||||
("execution_contexts", "positive_snapshot_id,context_id"),
|
||||
("execution_run_contexts", "context_id,run_id,anchor_id"),
|
||||
("execution_runs", "observed_dequeue_at,session_id,queue_item_id,run_id"),
|
||||
# Attribution of an output file to its run (#342)
|
||||
("execution_run_outputs", ",".join(OUTPUT_READ_INDEX_COLUMNS)),
|
||||
):
|
||||
conn.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_history_read_{table} "
|
||||
|
||||
@@ -119,6 +119,25 @@ def install_output_schema(conn: sqlite3.Connection) -> None:
|
||||
conn.execute(_table_sql(table))
|
||||
|
||||
|
||||
# The one optional, non-unique read index (#342): output file -> run. Older
|
||||
# libraries may lack it; migrate_execution_history adds it.
|
||||
OUTPUT_READ_INDEX = "idx_history_read_execution_run_outputs"
|
||||
OUTPUT_READ_INDEX_COLUMNS = ("filename", "subfolder", "root_key", "path_style")
|
||||
|
||||
|
||||
def _is_read_index(table: str, index: tuple, keys: tuple) -> bool:
|
||||
"""Exactly the attribution read index, as the migration creates it."""
|
||||
return (
|
||||
table == "execution_run_outputs"
|
||||
and index[1] == OUTPUT_READ_INDEX
|
||||
and not index[2]
|
||||
and index[3] == "c"
|
||||
and not index[4]
|
||||
and tuple(key[2] for key in keys) == OUTPUT_READ_INDEX_COLUMNS
|
||||
and not any(key[3] or key[4] != "BINARY" for key in keys)
|
||||
)
|
||||
|
||||
|
||||
def output_schema_error(conn: sqlite3.Connection) -> str | None:
|
||||
"""Reject altered schemas and invalid rows before returning bounded DTOs."""
|
||||
if conn.execute(
|
||||
@@ -157,6 +176,8 @@ def output_schema_error(conn: sqlite3.Connection) -> str | None:
|
||||
keys = tuple(
|
||||
row for row in conn.execute(f'PRAGMA index_xinfo("{name}")') if row[5]
|
||||
)
|
||||
if _is_read_index(table, index, keys):
|
||||
continue
|
||||
if (
|
||||
not index[2]
|
||||
or index[3] != "pk"
|
||||
|
||||
@@ -4,6 +4,8 @@ import sqlite3
|
||||
from contextlib import ExitStack
|
||||
from typing import Callable, ContextManager
|
||||
|
||||
from .run_attribution import notify_outputs_sealed
|
||||
|
||||
try:
|
||||
from ..utils.execution_history_events import RunFinish, encode_command
|
||||
from ..utils.execution_outputs import OutputBundle, validate_output_bundle
|
||||
@@ -113,6 +115,10 @@ class ExecutionOutputOperations:
|
||||
"VALUES(?,?,?,?,?,?,?,?)",
|
||||
[(run_id, *reference) for reference in prepared],
|
||||
)
|
||||
# Committed: images waiting for this run may now be attributed (#342)
|
||||
notify_outputs_sealed(
|
||||
run_id, [dict(zip(_REFERENCE_FIELDS, row)) for row in prepared]
|
||||
)
|
||||
return True
|
||||
|
||||
def get_execution_outputs(self, run_id: str) -> dict:
|
||||
|
||||
@@ -16,9 +16,12 @@ only when all of this holds:
|
||||
Anything else returns None: the caller leaves the file unlinked. Read-only.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
from typing import Any, Optional
|
||||
import threading
|
||||
import weakref
|
||||
from typing import Any, Callable, Iterable, List, Optional
|
||||
|
||||
try:
|
||||
from ..utils.execution_outputs import snapshot_output_roots
|
||||
@@ -120,3 +123,36 @@ def stored_run_prompt(
|
||||
return None
|
||||
(prompt_id,) = named
|
||||
return prompt_id if isinstance(prompt_id, int) else None
|
||||
|
||||
|
||||
_LOGGER = logging.getLogger("prompt_manager.run_attribution")
|
||||
_SEALED_LISTENERS: List[weakref.WeakMethod] = []
|
||||
_SEALED_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def subscribe_outputs_sealed(method: Callable[[str, List[dict]], None]) -> None:
|
||||
"""Call bound ``method(run_id, references)`` after a run's outputs are recorded.
|
||||
|
||||
Held weakly: a listener that is garbage collected is dropped.
|
||||
"""
|
||||
with _SEALED_LOCK:
|
||||
_SEALED_LISTENERS.append(weakref.WeakMethod(method))
|
||||
|
||||
|
||||
def notify_outputs_sealed(run_id: str, references: Iterable[dict]) -> None:
|
||||
"""Tell listeners a run's output references were just committed.
|
||||
|
||||
A listener failure is logged and never reaches the history writer.
|
||||
"""
|
||||
references = list(references)
|
||||
with _SEALED_LOCK:
|
||||
alive = [ref for ref in _SEALED_LISTENERS if ref() is not None]
|
||||
_SEALED_LISTENERS[:] = alive
|
||||
listeners = [ref() for ref in alive]
|
||||
for listener in listeners:
|
||||
if listener is None:
|
||||
continue
|
||||
try:
|
||||
listener(run_id, references)
|
||||
except Exception:
|
||||
_LOGGER.warning("Output seal listener failed", exc_info=True)
|
||||
|
||||
@@ -238,6 +238,7 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
|
||||
except Exception as e:
|
||||
# Log error but don't fail the encoding
|
||||
self.logger.warning(f"Failed to save prompt to database: {e}")
|
||||
self._mark_untracked(prompt, unique_id)
|
||||
else:
|
||||
self._skip_empty_prompt(prompt, unique_id)
|
||||
|
||||
|
||||
+27
-9
@@ -170,7 +170,10 @@ class PromptManagerBase:
|
||||
prompt_id = self._save_prompt_to_database(
|
||||
text=text, category=category, tags=tags, count_run=batch_item
|
||||
)
|
||||
if prompt_id and role["is_run_prompt"]:
|
||||
if role["is_run_prompt"] and not prompt_id:
|
||||
# The save failed: an older entry of this key must not stand in (#342)
|
||||
self._mark_untracked(prompt_graph, unique_id)
|
||||
elif role["is_run_prompt"]:
|
||||
self.prompt_tracker.set_current_prompt(
|
||||
prompt_text=encoding_text,
|
||||
additional_data={**additional_data, "prompt_id": prompt_id},
|
||||
@@ -193,25 +196,40 @@ class PromptManagerBase:
|
||||
Prepend and append text are deliberately never stored as the prompt, so
|
||||
an empty editor has no prompt to save, count or link. For the run's
|
||||
positive node this is logged at info and the node's run key is marked
|
||||
empty, so no earlier prompt can be linked to this run's images (and a
|
||||
batch item keeps its place in the queue).
|
||||
empty (``_mark_untracked``).
|
||||
"""
|
||||
self._mark_untracked(
|
||||
prompt_graph,
|
||||
unique_id,
|
||||
"Prompt text is empty (prepend and append text are not stored): "
|
||||
"nothing saved or tracked, and this run's images stay unlinked",
|
||||
)
|
||||
|
||||
def _mark_untracked(
|
||||
self, prompt_graph: Any, unique_id: Any, message: Optional[str] = None
|
||||
) -> None:
|
||||
"""This execution of a positive node has no prompt to link (#342).
|
||||
|
||||
An empty editor, or a save or tracking failure. The node's run key is
|
||||
marked with a placeholder, so no earlier prompt of that key can be
|
||||
linked to this run's images, and a batch item keeps its queue place.
|
||||
``message`` is logged at info for a positive node.
|
||||
"""
|
||||
try:
|
||||
role = self._prompt_role(prompt_graph, unique_id, "")
|
||||
if not role["is_run_prompt"]:
|
||||
self.logger.debug("Empty text on a non-positive node; nothing saved")
|
||||
if message:
|
||||
self.logger.debug("Empty text on a non-positive node; not saved")
|
||||
return
|
||||
self.logger.info(
|
||||
"Prompt text is empty (prepend and append text are not stored): "
|
||||
"nothing saved or tracked, and this run's images stay unlinked"
|
||||
)
|
||||
if message:
|
||||
self.logger.info(message)
|
||||
self.prompt_tracker.mark_empty(
|
||||
push_to_queue=role["known"],
|
||||
**self._run_identity(prompt_graph, unique_id),
|
||||
)
|
||||
except Exception as error:
|
||||
# Tracking must never block generation
|
||||
self.logger.warning(f"Could not mark the empty prompt: {error}")
|
||||
self.logger.warning(f"Could not mark the untracked prompt: {error}")
|
||||
|
||||
def _record_use(self, prompt_id: int, times: int = 1) -> None:
|
||||
"""Count runs; a tracking failure must never block generation."""
|
||||
|
||||
@@ -138,6 +138,8 @@ class PromptManagerPair(PromptManagerFocused):
|
||||
)
|
||||
except Exception as error:
|
||||
self.logger.warning(f"Could not save or track paired prompt: {error}")
|
||||
if is_positive:
|
||||
self._mark_untracked(prompt, unique_id)
|
||||
try:
|
||||
self.comfyui_integration.register_prompt(
|
||||
f"promptmanager_pair_{time.time_ns()}",
|
||||
|
||||
@@ -216,6 +216,7 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
|
||||
except Exception as e:
|
||||
# Log error but don't fail the processing
|
||||
self.logger.warning(f"Failed to save prompt to database: {e}")
|
||||
self._mark_untracked(prompt, unique_id)
|
||||
else:
|
||||
self._skip_empty_prompt(prompt, unique_id)
|
||||
|
||||
|
||||
+17
-7
@@ -1654,7 +1654,9 @@ class AdminRoutesMixin:
|
||||
a dict.
|
||||
"""
|
||||
if not metadata:
|
||||
return None
|
||||
# No prompt metadata (--disable-metadata, bare savers): only the
|
||||
# recorded run can say which prompt made it (#342)
|
||||
return self._link_recorded_run(media_file, {}, detailed)
|
||||
canonical = dict(metadata)
|
||||
for key in ("prompt", "workflow"):
|
||||
if isinstance(canonical.get(key), str):
|
||||
@@ -1664,7 +1666,7 @@ class AdminRoutesMixin:
|
||||
# Preserve malformed metadata for conservative resolver rejection.
|
||||
pass
|
||||
if not canonical.get("prompt") and not canonical.get("parameters"):
|
||||
return None
|
||||
return self._link_recorded_run(media_file, canonical, detailed)
|
||||
_with_file_size(canonical, media_file)
|
||||
resolution = resolve_prompt_pair(canonical)
|
||||
if blank_positive(resolution):
|
||||
@@ -1698,7 +1700,7 @@ class AdminRoutesMixin:
|
||||
# Graphs without a recognized sampler keep the legacy extractor below.
|
||||
prompt_text = self._graph_prompt_text(metadata)
|
||||
if not prompt_text:
|
||||
linked = self._link_recorded_run(media_file, canonical)
|
||||
linked = self._link_recorded_run(media_file, canonical, detailed)
|
||||
if linked is not None:
|
||||
return linked
|
||||
return "unresolved" if unchoosable else resolution["state"]
|
||||
@@ -1724,12 +1726,13 @@ class AdminRoutesMixin:
|
||||
return status
|
||||
return {"status": status, "title": self._scan_title(media_file, prompt_text)}
|
||||
|
||||
def _link_recorded_run(self, media_file, canonical):
|
||||
def _link_recorded_run(self, media_file, canonical, detailed=False):
|
||||
"""Link a file recorded history attributes to one stored prompt (#342).
|
||||
|
||||
Returns the scan status, or None when history names no single stored
|
||||
Returns the scan status (with ``detailed``, a stored or linked file's
|
||||
``{"status", "title"}``), or None when history names no single stored
|
||||
prompt for the file (``database.run_attribution``). Never creates a
|
||||
prompt.
|
||||
prompt. Works without any prompt metadata in the file.
|
||||
"""
|
||||
try:
|
||||
prompt_id = stored_run_prompt(
|
||||
@@ -1749,7 +1752,14 @@ class AdminRoutesMixin:
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to link {media_file.name} from history: {e}")
|
||||
return "found"
|
||||
return "linked" if image_id else "unchanged"
|
||||
status = "linked" if image_id else "unchanged"
|
||||
if not detailed or status not in SCAN_TITLED_DETAILS:
|
||||
return status
|
||||
prompt = self.db.get_prompt_by_id(prompt_id) or {}
|
||||
return {
|
||||
"status": status,
|
||||
"title": self._scan_title(media_file, prompt.get("text") or ""),
|
||||
}
|
||||
|
||||
def _scan_title(self, media_file, text):
|
||||
"""The linked prompt's Library title (its stored title, else ``text``'s
|
||||
|
||||
@@ -37,7 +37,7 @@ function graph({ kind = 'PromptManagerPair', text = '', prepend = '', llm = '',
|
||||
return nodes;
|
||||
}
|
||||
|
||||
async function saved(request, prompt) {
|
||||
async function saved(request, prompt, history = {}) {
|
||||
const id = randomUUID();
|
||||
const response = await request.post('/prompt', { data: { prompt, prompt_id: id } });
|
||||
expect(response.ok()).toBe(true);
|
||||
@@ -48,6 +48,7 @@ async function saved(request, prompt) {
|
||||
}, { timeout: 15000 }).toBe(true);
|
||||
const [image] = item.outputs['70'].images;
|
||||
expect(image.type).toBe('output');
|
||||
history.item = item;
|
||||
return image.filename;
|
||||
}
|
||||
|
||||
@@ -62,11 +63,19 @@ async function monitoring(execution, enabled) {
|
||||
expect((await response.json()).monitoring).toBe(enabled);
|
||||
}
|
||||
|
||||
async function cacheMode(execution, mode) {
|
||||
expect((await execution.request.post('/pm_execution_fixture/cache', { data: { mode } })).ok()).toBe(true);
|
||||
}
|
||||
|
||||
test.beforeEach(async ({ execution }) => {
|
||||
expect((await execution.request.post('/pm_execution_fixture/metadata', { data: { disabled: false } })).ok()).toBe(true);
|
||||
await cacheMode(execution, 'CLASSIC'); // keeps the re-run's PromptManager node cached
|
||||
await monitoring(execution, true);
|
||||
});
|
||||
test.afterEach(async ({ execution }) => { await monitoring(execution, false); });
|
||||
test.afterEach(async ({ execution }) => {
|
||||
await monitoring(execution, false);
|
||||
await cacheMode(execution, 'RAM_PRESSURE');
|
||||
});
|
||||
|
||||
test('an output whose positive is only linked prepend text stays unlinked, never given an earlier prompt', async ({ execution }, testInfo) => {
|
||||
const tag = `${testInfo.project.name} ${randomUUID()}`;
|
||||
@@ -89,9 +98,13 @@ test('an output whose positive is only linked prepend text stays unlinked, never
|
||||
const state = await readLinks(execution.databasePath, [orphan], [prepended]);
|
||||
expect(state).toEqual({ links: { [orphan]: null }, stored: { [prepended]: false } });
|
||||
|
||||
// Only the sampler seed changes: the PromptManager node may be served from the
|
||||
// cache without running, and its unchanged upstream still proves the prompt.
|
||||
// Only the sampler seed changes: the PromptManager node is served from the cache
|
||||
// without running (the host says so), so no queue entry exists for this run and
|
||||
// only the snapshot's run match can link it.
|
||||
const rerunHistory = {};
|
||||
const rerun = await saved(execution.request, { ...llmGraph, '40': { ...llmGraph['40'],
|
||||
inputs: { ...llmGraph['40'].inputs, seed: 2 } } });
|
||||
inputs: { ...llmGraph['40'].inputs, seed: 2 } } }, rerunHistory);
|
||||
const cached = rerunHistory.item.status.messages.find(([kind]) => kind === 'execution_cached');
|
||||
expect(cached?.[1].nodes).toEqual(expect.arrayContaining(['11', '20']));
|
||||
await linkedTo(execution, rerun, llmText);
|
||||
});
|
||||
|
||||
@@ -510,5 +510,119 @@ class TestRunKeys(unittest.TestCase):
|
||||
self.assertIsNone(prompt_node_key(chain, "0"))
|
||||
|
||||
|
||||
class TestReviewFindings(MonitorCase, NodeCase):
|
||||
"""Final-review findings M1-M4 (#344)."""
|
||||
|
||||
def test_m1_a_positive_node_without_a_key_makes_the_image_unkeyable(self):
|
||||
graph = llm_text_graph()
|
||||
graph["30"] = {
|
||||
"class_type": "PromptManager",
|
||||
"inputs": {"text": ["missing", 0], "clip": ["4", 1]},
|
||||
}
|
||||
graph["31"] = {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {**graph["3"]["inputs"], "positive": ["30", 0]},
|
||||
}
|
||||
self.assertIsNone(prompt_node_key(graph, "30"))
|
||||
self.assertEqual(image_run_keys({"prompt": graph}), {})
|
||||
prompt_id = self.save(LLM_TEXT)
|
||||
self.tracker.set_current_prompt(
|
||||
LLM_TEXT,
|
||||
{"prompt_id": prompt_id},
|
||||
run_key=prompt_node_key(graph, "18:35"),
|
||||
run_id="run-1",
|
||||
)
|
||||
self.assertIsNone(self.process({"prompt": graph}))
|
||||
|
||||
def test_m2_links_are_exactly_comfyui_is_link(self):
|
||||
from utils.run_keys import _links
|
||||
|
||||
inputs = {
|
||||
"string_slot_int": ["7", 0],
|
||||
"string_slot_float": ["7", 0.0],
|
||||
"string_slot_bool": ["7", True], # bool is an int to ComfyUI too
|
||||
"int_source": [7, 0],
|
||||
"tuple": ("7", 0),
|
||||
"three": ["7", 0, 1],
|
||||
"slot_text": ["7", "0"],
|
||||
}
|
||||
self.assertEqual(sorted(_links(inputs)), ["7", "7", "7"])
|
||||
widget = issue_graph()
|
||||
widget["18:17"]["inputs"]["size"] = [512, 768] # a value, not a link
|
||||
self.assertIsNotNone(prompt_node_key(widget, "18:35"))
|
||||
floated = issue_graph()
|
||||
floated["18:35"]["inputs"]["prepend_text"] = ["18:13", 0.0]
|
||||
changed = copy.deepcopy(floated)
|
||||
changed["18:17"]["inputs"]["value"] = "another instruction"
|
||||
self.assertNotEqual(
|
||||
prompt_node_key(floated, "18:35"), prompt_node_key(changed, "18:35")
|
||||
)
|
||||
|
||||
def _failing_save(self, node):
|
||||
def explode(*args, **kwargs):
|
||||
raise RuntimeError("database is locked")
|
||||
|
||||
node.db = mock.Mock(wraps=self.db)
|
||||
for name in ("save_prompt", "save_prompt_pair", "get_prompt_by_hash"):
|
||||
setattr(node.db, name, mock.Mock(side_effect=explode))
|
||||
return node
|
||||
|
||||
def test_m3_a_failed_save_takes_over_its_run_key(self):
|
||||
graph = llm_text_graph()
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
older = self.save("previous run text")
|
||||
self.tracker.set_current_prompt(
|
||||
"previous run text", {"prompt_id": older}, run_key=key, run_id="run-1"
|
||||
)
|
||||
node = self._failing_save(self.node(PromptManager))
|
||||
with mock.patch("prompt_manager_base.executing_run_id", return_value="run-2"):
|
||||
node.encode_prompt(
|
||||
self.clip(), "new run text", prompt=graph, unique_id="18:35"
|
||||
)
|
||||
latest, run_id, _count = self.tracker.keyed_snapshot()[key]
|
||||
self.assertTrue(latest.get("placeholder"))
|
||||
self.assertEqual(run_id, "run-2")
|
||||
self.assertIsNone(self.process({"prompt": graph}))
|
||||
|
||||
def test_m3_a_failed_pair_save_takes_over_its_run_key(self):
|
||||
graph = issue_graph(text="typed")
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
node = self._failing_save(self.node(PromptManagerPair))
|
||||
with self.assertLogs("test.prompt_attribution.node", "WARNING"):
|
||||
node.encode_prompt(self.clip(), "typed", prompt=graph, unique_id="18:35")
|
||||
self.assertTrue(self.tracker.keyed_snapshot()[key][0].get("placeholder"))
|
||||
|
||||
def test_m3_a_failed_text_node_save_takes_over_its_run_key(self):
|
||||
graph = issue_graph()
|
||||
graph["18:35"] = {
|
||||
"class_type": "PromptManagerText",
|
||||
"inputs": {"text": "typed"},
|
||||
}
|
||||
graph["40"] = {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": ["18:35", 0], "clip": ["4", 1]},
|
||||
}
|
||||
graph["3"]["inputs"]["positive"] = ["40", 0]
|
||||
node = self._failing_save(self.node(PromptManagerText))
|
||||
with mock.patch.object(
|
||||
node, "_track_prompt_execution", side_effect=RuntimeError("boom")
|
||||
):
|
||||
node.process_text("typed", prompt=graph, unique_id="18:35")
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
self.assertTrue(self.tracker.keyed_snapshot()[key][0].get("placeholder"))
|
||||
|
||||
def test_m4_a_placeholder_never_falls_through_to_history(self):
|
||||
graph = llm_text_graph()
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
self.tracker.mark_empty(key, "run-1", push_to_queue=False)
|
||||
with mock.patch.object(
|
||||
self.handler, "_history_prompt", return_value=99
|
||||
) as history:
|
||||
self.assertIsNone(self.process({"prompt": graph})) # snapshot form
|
||||
self.tracker.mark_empty(key, "run-2", push_to_queue=True)
|
||||
self.assertIsNone(self.process({"prompt": graph})) # popped form
|
||||
history.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -17,6 +17,7 @@ except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from tests.test_api_images import ImageAPITestCase
|
||||
from tests.test_prompt_attribution import issue_graph, pair_evidence
|
||||
from tools import review_image_links as tool
|
||||
from utils.hashing import generate_prompt_hash
|
||||
@@ -31,7 +32,7 @@ def png(path, prompt=None, evidence=None):
|
||||
Image.new("RGB", (4, 4)).save(path, pnginfo=info)
|
||||
|
||||
|
||||
class ReviewToolTests(unittest.TestCase):
|
||||
class ToolFixture(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.mkdtemp()
|
||||
self.addCleanup(shutil.rmtree, self.tmp, True)
|
||||
@@ -81,12 +82,14 @@ class ReviewToolTests(unittest.TestCase):
|
||||
finally:
|
||||
db.close_all()
|
||||
|
||||
|
||||
class ReviewToolTests(ToolFixture):
|
||||
def test_dry_run_judges_and_changes_nothing(self):
|
||||
before = self.links()
|
||||
ids = f"{self.guessed[0]}-{self.gone}"
|
||||
code, out = self.run_tool("--ids", ids, "--json")
|
||||
self.assertEqual(code, 0)
|
||||
verdicts = {v["image_id"]: v for v in json.loads(out.split("Dry run")[0])}
|
||||
verdicts = {v["image_id"]: v for v in json.loads(out.split("--apply, for")[0])}
|
||||
for image_id in self.guessed:
|
||||
self.assertEqual(verdicts[image_id]["verdict"], "unlink")
|
||||
self.assertEqual(verdicts[image_id]["sources"], {})
|
||||
@@ -114,7 +117,7 @@ class ReviewToolTests(unittest.TestCase):
|
||||
"--ids", named, "--prompt", str(self.coffee), "--apply"
|
||||
)
|
||||
self.assertEqual(code, 0)
|
||||
self.assertIn("Applied: 1 relinked, 2 unlinked", out)
|
||||
self.assertIn("Applied: 1 relinked, 2 unlinked, 0 skipped", out)
|
||||
links = self.links()
|
||||
self.assertNotIn(self.guessed[0], links)
|
||||
self.assertNotIn(self.guessed[1], links)
|
||||
@@ -171,5 +174,82 @@ class ReviewToolTests(unittest.TestCase):
|
||||
self.assertEqual(verdict["verdict"], "conflict")
|
||||
|
||||
|
||||
class ApplyGuardTests(ToolFixture):
|
||||
"""M6: unlink clears only the link, and only the link it was judged on."""
|
||||
|
||||
def test_unlink_skips_a_row_linked_elsewhere_since_the_review(self):
|
||||
db = PromptDatabase(self.db_path)
|
||||
try:
|
||||
verdicts = tool.review(db, [self.guessed[0]], None, self.output)
|
||||
self.assertEqual(verdicts[0]["verdict"], "unlink")
|
||||
db.model.get_connection().execute(
|
||||
"UPDATE generated_images SET prompt_id=? WHERE id=?",
|
||||
(self.typed, self.guessed[0]),
|
||||
)
|
||||
db.model.get_connection().commit()
|
||||
result = tool.apply(db, verdicts, [self.guessed[0]])
|
||||
finally:
|
||||
db.close_all()
|
||||
self.assertEqual(result["unlinked"], 0)
|
||||
self.assertEqual(self.links()[self.guessed[0]], self.typed)
|
||||
|
||||
def test_unlink_keeps_the_file_and_its_favorite(self):
|
||||
path = os.path.join(self.output, "moss_00009_.png")
|
||||
db = PromptDatabase(self.db_path)
|
||||
try:
|
||||
db.set_image_favorite("a" * 64, "moss_00009_.png")
|
||||
finally:
|
||||
db.close_all()
|
||||
self.run_tool("--ids", str(self.guessed[0]), "--apply")
|
||||
self.assertTrue(os.path.isfile(path))
|
||||
db = PromptDatabase(self.db_path)
|
||||
try:
|
||||
self.assertEqual(db.get_root_favorites("a" * 64), {"moss_00009_.png"})
|
||||
summary = db.get_image_link_summaries([path]).get(path) or {}
|
||||
finally:
|
||||
db.close_all()
|
||||
self.assertNotEqual(summary.get("status"), "linked")
|
||||
|
||||
def test_dry_run_states_what_apply_does(self):
|
||||
_code, out = self.run_tool("--ids", str(self.guessed[0]))
|
||||
self.assertIn("removes only the image's link record", out)
|
||||
self.assertIn("file stays on disk and in the gallery as unlinked", out)
|
||||
|
||||
|
||||
class GalleryAfterUnlinkTests(ImageAPITestCase):
|
||||
"""M6: after --apply the file is still listed, as an unlinked image."""
|
||||
|
||||
async def test_unlinked_file_stays_in_the_workspace_gallery(self):
|
||||
db = self.api.db
|
||||
prompt = db.save_prompt(
|
||||
text="a woman holding a coffee cup",
|
||||
prompt_hash=generate_prompt_hash("a woman holding a coffee cup"),
|
||||
)
|
||||
path = self.output_dir / "moss_00009_.png"
|
||||
png(str(path), issue_graph(), pair_evidence())
|
||||
image_id = db.link_image_to_prompt(prompt, str(path), metadata={})
|
||||
out = io.StringIO()
|
||||
tool.main(
|
||||
[
|
||||
db.model.db_path,
|
||||
"--output-dir",
|
||||
str(self.output_dir),
|
||||
"--ids",
|
||||
str(image_id),
|
||||
"--apply",
|
||||
],
|
||||
out=out,
|
||||
)
|
||||
self.assertIn("Applied: 0 relinked, 1 unlinked", out.getvalue())
|
||||
response = await self.client.get(
|
||||
"/prompt_manager/images/output", params={"mode": "workspace"}
|
||||
)
|
||||
self.assertEqual(response.status, 200)
|
||||
images = {i["filename"]: i for i in (await response.json())["images"]}
|
||||
self.assertIn("moss_00009_.png", images)
|
||||
link = images["moss_00009_.png"].get("linked") or {}
|
||||
self.assertNotEqual(link.get("status"), "linked")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -198,6 +198,27 @@ class TestScanUsesHistory(RecordedRuns, AdminAPITestCase):
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, meta), "linked")
|
||||
self.assertEqual(self.linked_prompt(path), prompt_id)
|
||||
|
||||
def test_a_metadata_less_file_links_through_its_recorded_run(self):
|
||||
"""I1: --disable-metadata outputs carry no prompt; history still proves one."""
|
||||
db = self.api.db
|
||||
prompt_id = db.save_prompt(
|
||||
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
||||
)
|
||||
path = self.write("bare_00001_.png")
|
||||
self.record_run(db, str(self.output_dir), [("", path.name)], [LLM_TEXT])
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, {}), "linked")
|
||||
self.assertEqual(self.linked_prompt(path), prompt_id)
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, None), "unchanged")
|
||||
detailed = self.api._ingest_scanned_file(path, {}, detailed=True)
|
||||
self.assertEqual(detailed["status"], "unchanged")
|
||||
self.assertTrue(detailed["title"])
|
||||
|
||||
def test_a_metadata_less_file_without_a_recorded_run_is_still_skipped(self):
|
||||
path = self.write("bare_00002_.png")
|
||||
self.assertIsNone(self.api._ingest_scanned_file(path, {}))
|
||||
self.assertIsNone(self.api._ingest_scanned_file(path, {"workflow": {}}))
|
||||
self.assertIsNone(self.linked_prompt(path))
|
||||
|
||||
def test_the_issue_file_stays_unlinked_even_with_its_pair_record(self):
|
||||
"""Resolved with an empty stored positive: before #342 the scan raised."""
|
||||
db = self.api.db
|
||||
@@ -233,10 +254,6 @@ class TestScanUsesHistory(RecordedRuns, AdminAPITestCase):
|
||||
self.assertEqual(db.list_prompt_pairs()["pairs"], [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
|
||||
class TestHostHelpers(unittest.TestCase):
|
||||
"""The host lookups degrade to "no proof", never to an error."""
|
||||
|
||||
@@ -301,3 +318,176 @@ class TestHostHelpers(unittest.TestCase):
|
||||
sys.modules.pop("comfy_execution.utils", None)
|
||||
else:
|
||||
sys.modules["comfy_execution.utils"] = original
|
||||
|
||||
|
||||
class TestMonitorHistoryRecheck(RecordedRuns, unittest.TestCase):
|
||||
"""I1: a keyless image gets one history-only re-check when its run is sealed."""
|
||||
|
||||
def setUp(self):
|
||||
history_tests.HistoryDatabaseTests.setUp(self)
|
||||
self.db = history_tests.HistoryDatabaseTests.database(self)
|
||||
self.output = str(Path(self.temp.name) / "output")
|
||||
os.makedirs(self.output)
|
||||
self.handler = ImageGenerationHandler(self.db, make_tracker(self.db))
|
||||
self.handler.processing_delay = 0
|
||||
self.handler.settle_interval = 0.01
|
||||
self.addCleanup(self.handler.shutdown)
|
||||
self.linked = []
|
||||
self.handler.link_image_to_prompt = lambda path, prompt, meta: (
|
||||
self.linked.append((os.path.basename(path), prompt["id"]))
|
||||
)
|
||||
patches = (
|
||||
mock.patch(
|
||||
"utils.image_monitor.host_output_directory", return_value=self.output
|
||||
),
|
||||
mock.patch.object(
|
||||
self.handler.metadata_extractor, "extract_metadata", return_value=None
|
||||
),
|
||||
)
|
||||
for patch in patches:
|
||||
patch.start()
|
||||
self.addCleanup(patch.stop)
|
||||
|
||||
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
||||
|
||||
def image(self, name):
|
||||
path = os.path.join(self.output, name)
|
||||
Path(path).write_bytes(b"png")
|
||||
return path
|
||||
|
||||
def save(self, text):
|
||||
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
||||
|
||||
def wait_for(self, predicate, timeout=5.0):
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
if predicate():
|
||||
return True
|
||||
time.sleep(0.01)
|
||||
return predicate()
|
||||
|
||||
def test_the_run_ending_after_processing_links_the_image_once(self):
|
||||
prompt_id = self.save(LLM_TEXT)
|
||||
path = self.image("late_00001_.png")
|
||||
self.handler.process_new_image(path) # the run is still executing
|
||||
self.assertEqual(self.linked, [])
|
||||
self.assertIn(os.path.normcase(path), self.handler.awaiting_history())
|
||||
self.record_run(self.db, self.output, [("", "late_00001_.png")], [LLM_TEXT])
|
||||
self.assertTrue(self.wait_for(lambda: self.linked))
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(self.linked, [("late_00001_.png", prompt_id)])
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
# Another run reporting the same name later is not a second re-check
|
||||
self.record_run(
|
||||
self.db, self.output, [("", "late_00001_.png")], [LLM_TEXT], item=1
|
||||
)
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(len(self.linked), 1)
|
||||
|
||||
def test_a_run_naming_nothing_stored_leaves_the_image_unlinked(self):
|
||||
path = self.image("empty_00001_.png")
|
||||
self.handler.process_new_image(path)
|
||||
self.record_run(self.db, self.output, [("", "empty_00001_.png")], [""])
|
||||
self.assertTrue(self.wait_for(lambda: not self.handler.awaiting_history()))
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(self.linked, [])
|
||||
|
||||
def test_history_found_at_once_leaves_nothing_awaiting(self):
|
||||
prompt_id = self.save(LLM_TEXT)
|
||||
self.record_run(self.db, self.output, [("", "done_00001_.png")], [LLM_TEXT])
|
||||
self.handler.process_new_image(self.image("done_00001_.png"))
|
||||
self.assertEqual(self.linked, [("done_00001_.png", prompt_id)])
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
|
||||
def test_keyed_images_do_not_wait_for_history(self):
|
||||
with mock.patch.object(
|
||||
self.handler.metadata_extractor,
|
||||
"extract_metadata",
|
||||
return_value={"prompt": llm_text_graph()},
|
||||
):
|
||||
self.handler.process_new_image(self.image("keyed_00001_.png"))
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
|
||||
def test_the_waiting_set_is_bounded_and_expires(self):
|
||||
from utils import image_monitor
|
||||
|
||||
with mock.patch.object(image_monitor, "MAX_AWAITING_HISTORY", 2):
|
||||
for n in range(3):
|
||||
self.handler.process_new_image(self.image(f"b_{n}.png"))
|
||||
self.assertEqual(
|
||||
[os.path.basename(p) for p in self.handler.awaiting_history()],
|
||||
["b_1.png", "b_2.png"],
|
||||
)
|
||||
later = time.monotonic() + image_monitor.AWAIT_HISTORY_SECONDS + 1
|
||||
with mock.patch.object(image_monitor.time, "monotonic", return_value=later):
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
|
||||
def test_shutdown_forgets_waiting_images_and_a_dead_handler_is_dropped(self):
|
||||
import gc
|
||||
|
||||
from database import run_attribution
|
||||
|
||||
self.handler.process_new_image(self.image("gone_00001_.png"))
|
||||
self.handler.shutdown()
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
other = ImageGenerationHandler(self.db, make_tracker(self.db))
|
||||
with mock.patch.object(
|
||||
other.metadata_extractor, "extract_metadata", return_value=None
|
||||
):
|
||||
other.process_new_image(self.image("gone_00002_.png"))
|
||||
del other
|
||||
gc.collect()
|
||||
run_attribution.notify_outputs_sealed("run", []) # no dead callback raises
|
||||
|
||||
|
||||
class TestOutputLookupIndex(unittest.TestCase):
|
||||
"""M5: the history lookup by file name uses an index, added by migration."""
|
||||
|
||||
setUp = history_tests.HistoryDatabaseTests.setUp
|
||||
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
||||
database = history_tests.HistoryDatabaseTests.database
|
||||
|
||||
def test_index_exists_is_used_validated_and_reinstalled(self):
|
||||
from database.execution_history_schema import (
|
||||
history_schema_error,
|
||||
migrate_execution_history,
|
||||
)
|
||||
from database.run_attribution import _RUNS_FOR_OUTPUT
|
||||
|
||||
db = self.database()
|
||||
conn = db.model.get_connection()
|
||||
name = "idx_history_read_execution_run_outputs"
|
||||
columns = [r[2] for r in conn.execute(f"PRAGMA index_info({name})")]
|
||||
self.assertEqual(columns, ["filename", "subfolder", "root_key", "path_style"])
|
||||
plan = " ".join(
|
||||
str(row[3])
|
||||
for row in conn.execute(
|
||||
"EXPLAIN QUERY PLAN " + _RUNS_FOR_OUTPUT, ("posix", "k", "", "f.png")
|
||||
)
|
||||
)
|
||||
self.assertIn(name, plan)
|
||||
self.assertIsNone(history_schema_error(conn, allow_missing=False))
|
||||
conn.execute(f"DROP INDEX {name}")
|
||||
self.assertIsNone(history_schema_error(conn, allow_missing=False))
|
||||
migrate_execution_history(conn)
|
||||
migrate_execution_history(conn)
|
||||
self.assertTrue(
|
||||
conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (name,)).fetchone()
|
||||
)
|
||||
|
||||
def test_any_other_index_on_the_outputs_table_is_still_refused(self):
|
||||
from database.execution_output_schema import output_schema_error
|
||||
|
||||
db = self.database()
|
||||
conn = db.model.get_connection()
|
||||
conn.execute("CREATE INDEX idx_rogue ON execution_run_outputs(filename)")
|
||||
self.assertIn("unexpected", output_schema_error(conn))
|
||||
conn.execute("DROP INDEX idx_rogue")
|
||||
name = "idx_history_read_execution_run_outputs"
|
||||
conn.execute(f"DROP INDEX {name}")
|
||||
conn.execute(f"CREATE INDEX {name} ON execution_run_outputs(filename)")
|
||||
self.assertIn("unexpected", output_schema_error(conn))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
+57
-11
@@ -21,8 +21,20 @@ connection) into a temporary folder and judged there, so the library is never
|
||||
written. ``--apply`` writes the ``relink`` and ``unlink`` verdicts of the
|
||||
selected images, after saving a backup beside the database
|
||||
(``<db>.pre-relink_<stamp>``). Only ids named on the command line are ever
|
||||
changed. A batch item linked through the monitor's queue (PromptSearchList)
|
||||
can't be proven after the fact and reads ``unlink``: review before applying.
|
||||
changed, and only while each is still linked to the prompt it was judged on.
|
||||
|
||||
What ``--apply`` does (APPLY_EFFECTS below, printed by every run):
|
||||
|
||||
- ``relink`` points the image's link record at the other prompt and drops its
|
||||
pair links (they named the old positive).
|
||||
- ``unlink`` removes only the image's link record (its ``generated_images`` row
|
||||
with the metadata snapshot stored in it, and its pair links). That is how the
|
||||
app represents an unlinked image, exactly like a file the scan could not
|
||||
attribute: the file stays on disk and in the gallery as unlinked, and its
|
||||
favorite star and tags (kept by file, not by link) are untouched.
|
||||
|
||||
A batch item linked through the monitor's queue (PromptSearchList) can't be
|
||||
proven after the fact and reads ``unlink``: review before applying.
|
||||
|
||||
python3 tools/review_image_links.py <db> --output-dir <ComfyUI output> \\
|
||||
--ids 58005-58010
|
||||
@@ -62,6 +74,15 @@ from utils.image_monitor import ImageGenerationHandler # noqa: E402
|
||||
from utils.prompt_pair_resolver import resolve_prompt_pair # noqa: E402
|
||||
|
||||
VERDICTS = ("keep", "relink", "unlink", "conflict", "missing")
|
||||
APPLY_EFFECTS = (
|
||||
"--apply, for ids named with --ids only, while each is still linked to the "
|
||||
"prompt it was judged on:\n"
|
||||
" relink: points the image's link record at the other prompt and drops "
|
||||
"its pair links;\n"
|
||||
" unlink: removes only the image's link record (its stored metadata "
|
||||
"snapshot and pair links go with it); the file stays on disk and in the "
|
||||
"gallery as unlinked, with its favorite and tags.\n"
|
||||
)
|
||||
# Ids per query (well under SQLite's bound-parameter limit)
|
||||
ID_BATCH = 500
|
||||
|
||||
@@ -174,15 +195,30 @@ def review(db, ids, prompt_id, output_dir):
|
||||
|
||||
|
||||
def apply(db, verdicts, ids):
|
||||
"""Write the relink/unlink verdicts of explicitly named ids; returns counts."""
|
||||
"""Write the relink/unlink verdicts of explicitly named ids; returns counts.
|
||||
|
||||
An image whose link changed since it was judged (now another prompt, or
|
||||
gone) is skipped.
|
||||
"""
|
||||
allowed = set(ids)
|
||||
relink = [
|
||||
v for v in verdicts if v["verdict"] == "relink" and v["image_id"] in allowed
|
||||
]
|
||||
unlink = [
|
||||
v for v in verdicts if v["verdict"] == "unlink" and v["image_id"] in allowed
|
||||
chosen = [
|
||||
v
|
||||
for v in verdicts
|
||||
if v["verdict"] in ("relink", "unlink") and v["image_id"] in allowed
|
||||
]
|
||||
with db._pair_transaction() as conn:
|
||||
current = [
|
||||
item
|
||||
for item in chosen
|
||||
if (
|
||||
row := conn.execute(
|
||||
"SELECT prompt_id FROM generated_images WHERE id=?",
|
||||
(item["image_id"],),
|
||||
).fetchone()
|
||||
)
|
||||
and row[0] == item["linked_prompt"]
|
||||
]
|
||||
relink = [item for item in current if item["verdict"] == "relink"]
|
||||
for item in relink:
|
||||
# A pair link records the old positive: it goes with the old link
|
||||
conn.execute(
|
||||
@@ -192,8 +228,16 @@ def apply(db, verdicts, ids):
|
||||
"UPDATE generated_images SET prompt_id=? WHERE id=? AND prompt_id=?",
|
||||
(item["target"], item["image_id"], item["linked_prompt"]),
|
||||
)
|
||||
removed = _delete_ids(conn, "generated_images", [v["image_id"] for v in unlink])
|
||||
return {"relinked": len(relink), "unlinked": removed}
|
||||
removed = _delete_ids(
|
||||
conn,
|
||||
"generated_images",
|
||||
[item["image_id"] for item in current if item["verdict"] == "unlink"],
|
||||
)
|
||||
return {
|
||||
"relinked": len(relink),
|
||||
"unlinked": removed,
|
||||
"skipped": len(chosen) - len(current),
|
||||
}
|
||||
|
||||
|
||||
def backup_database(source, target):
|
||||
@@ -253,6 +297,7 @@ def main(argv=None, out=sys.stdout):
|
||||
finally:
|
||||
db.close_all()
|
||||
_report(verdicts, args.json, out)
|
||||
out.write(APPLY_EFFECTS)
|
||||
out.write("Dry run: nothing was changed.\n")
|
||||
return 0
|
||||
|
||||
@@ -268,7 +313,8 @@ def main(argv=None, out=sys.stdout):
|
||||
finally:
|
||||
db.close_all()
|
||||
out.write(
|
||||
f"Applied: {result['relinked']} relinked, {result['unlinked']} unlinked\n"
|
||||
f"Applied: {result['relinked']} relinked, {result['unlinked']} unlinked, "
|
||||
f"{result['skipped']} skipped (link changed since review)\n"
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
+128
-5
@@ -28,7 +28,8 @@ import os
|
||||
import queue
|
||||
import time
|
||||
import threading
|
||||
from typing import Callable, Optional, Dict, Any, Tuple
|
||||
from collections import OrderedDict
|
||||
from typing import Callable, Optional, Dict, Any, List, Tuple
|
||||
from watchdog.observers import Observer
|
||||
from watchdog.events import FileSystemEventHandler
|
||||
|
||||
@@ -42,16 +43,31 @@ from .usage_tracking import survive_switch
|
||||
|
||||
try:
|
||||
from ..database.pair_ingestion import ingest_resolved_pair
|
||||
from ..database.run_attribution import host_output_directory, stored_run_prompt
|
||||
from ..database.run_attribution import (
|
||||
host_output_directory,
|
||||
stored_run_prompt,
|
||||
subscribe_outputs_sealed,
|
||||
)
|
||||
except ImportError:
|
||||
from database.pair_ingestion import ingest_resolved_pair
|
||||
from database.run_attribution import host_output_directory, stored_run_prompt
|
||||
from database.run_attribution import (
|
||||
host_output_directory,
|
||||
stored_run_prompt,
|
||||
subscribe_outputs_sealed,
|
||||
)
|
||||
|
||||
# Upper bound on images waiting to be processed (a burst beyond this is skipped)
|
||||
MAX_PENDING_IMAGES = 1000
|
||||
# A file counts as fully written once its size is unchanged across two reads
|
||||
SETTLE_INTERVAL_SECONDS = 0.1
|
||||
SETTLE_ATTEMPTS = 20
|
||||
# Keyless images (no embedded API graph) left unlinked wait for their run's
|
||||
# recorded outputs for one history-only re-check (#342): at most this many,
|
||||
# each for at most this long.
|
||||
MAX_AWAITING_HISTORY = 1000
|
||||
AWAIT_HISTORY_SECONDS = 3600
|
||||
# ``schedule`` snapshot meaning "history-only re-check" (not a new image)
|
||||
HISTORY_ONLY = object()
|
||||
|
||||
|
||||
class ImageGenerationHandler(FileSystemEventHandler):
|
||||
@@ -113,6 +129,11 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
self._work: "queue.Queue[Optional[str]]" = queue.Queue()
|
||||
self._stop = threading.Event()
|
||||
self._worker: Optional[threading.Thread] = None
|
||||
# Keyless unlinked images awaiting their run's recorded outputs:
|
||||
# normalized path -> (path, monotonic deadline), oldest first.
|
||||
self._awaiting: "OrderedDict[str, Tuple[str, float]]" = OrderedDict()
|
||||
self._awaiting_lock = threading.Lock()
|
||||
self._sealed_subscribed = False
|
||||
|
||||
def on_created(self, event):
|
||||
"""Handle filesystem creation events.
|
||||
@@ -180,6 +201,8 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
and exits; the next ``schedule`` starts a fresh worker with its own
|
||||
queue. Safe to call repeatedly and before anything was scheduled.
|
||||
"""
|
||||
with self._awaiting_lock:
|
||||
self._awaiting.clear()
|
||||
with self._pending_lock:
|
||||
self._pending.clear()
|
||||
worker, self._worker = self._worker, None
|
||||
@@ -307,6 +330,9 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
appeared (``PromptTracker.keyed_snapshot``); None reads it now
|
||||
"""
|
||||
try:
|
||||
if prompt_snapshot is HISTORY_ONLY:
|
||||
self._recheck_history(image_path)
|
||||
return
|
||||
self.logger.info(f"Processing image: {image_path}")
|
||||
|
||||
if not os.path.exists(image_path):
|
||||
@@ -399,6 +425,16 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
"""
|
||||
keys = image_run_keys(metadata)
|
||||
name = os.path.basename(image_path)
|
||||
if not keys:
|
||||
# No tracked run can be proven: history now, and once more after
|
||||
# the run's outputs are recorded (``_on_outputs_sealed``)
|
||||
self._await_history(image_path)
|
||||
prompt_id = self._history_prompt(image_path)
|
||||
if prompt_id is None:
|
||||
return None
|
||||
self._stop_awaiting(image_path)
|
||||
self.logger.info(f"History match: prompt {prompt_id} for {name}")
|
||||
return {"id": prompt_id, "from_history": True}
|
||||
if len(keys) == 1:
|
||||
(key,) = keys.values()
|
||||
if prompt_snapshot is None:
|
||||
@@ -414,10 +450,14 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
return None
|
||||
self.logger.info(f"Queue match: prompt {queued['id']} for {name}")
|
||||
return queued
|
||||
if not context.get("placeholder") and run_id and count == 1:
|
||||
if context.get("placeholder"):
|
||||
# Same rule as a popped placeholder: that run had no prompt
|
||||
self.logger.info(f"Latest run for {name} had no prompt text")
|
||||
return None
|
||||
if run_id and count == 1:
|
||||
self.logger.info(f"Run match: prompt {context['id']} for {name}")
|
||||
return context
|
||||
elif len(keys) > 1:
|
||||
else:
|
||||
self.logger.debug(f"{len(keys)} positive prompt nodes for {name}")
|
||||
prompt_id = self._history_prompt(image_path)
|
||||
if prompt_id is None:
|
||||
@@ -425,6 +465,89 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
self.logger.info(f"History match: prompt {prompt_id} for {name}")
|
||||
return {"id": prompt_id, "from_history": True}
|
||||
|
||||
@staticmethod
|
||||
def _awaiting_key(image_path: str) -> str:
|
||||
return os.path.normcase(os.path.abspath(image_path))
|
||||
|
||||
def _prune_awaiting(self) -> None:
|
||||
"""Drop expired and surplus entries (``_awaiting_lock`` held)."""
|
||||
now = time.monotonic()
|
||||
for key in [k for k, (_p, due) in self._awaiting.items() if due <= now]:
|
||||
del self._awaiting[key]
|
||||
while len(self._awaiting) > MAX_AWAITING_HISTORY:
|
||||
self._awaiting.popitem(last=False)
|
||||
|
||||
def _await_history(self, image_path: str) -> None:
|
||||
"""Re-check ``image_path`` once, history only, when its run is sealed."""
|
||||
if not self._sealed_subscribed:
|
||||
subscribe_outputs_sealed(self._on_outputs_sealed)
|
||||
self._sealed_subscribed = True
|
||||
key = self._awaiting_key(image_path)
|
||||
with self._awaiting_lock:
|
||||
self._awaiting.pop(key, None)
|
||||
self._awaiting[key] = (
|
||||
image_path,
|
||||
time.monotonic() + AWAIT_HISTORY_SECONDS,
|
||||
)
|
||||
self._prune_awaiting()
|
||||
|
||||
def _stop_awaiting(self, image_path: str) -> None:
|
||||
with self._awaiting_lock:
|
||||
self._awaiting.pop(self._awaiting_key(image_path), None)
|
||||
|
||||
def awaiting_history(self) -> List[str]:
|
||||
"""Normalized paths still waiting for their run's recorded outputs."""
|
||||
with self._awaiting_lock:
|
||||
self._prune_awaiting()
|
||||
return list(self._awaiting)
|
||||
|
||||
def _on_outputs_sealed(self, run_id: str, references: List[dict]) -> None:
|
||||
"""A run's outputs were recorded: re-check its waiting images once.
|
||||
|
||||
Called on the history writer's thread, so it only schedules the work.
|
||||
"""
|
||||
root = host_output_directory()
|
||||
if not root:
|
||||
return
|
||||
due = []
|
||||
with self._awaiting_lock:
|
||||
self._prune_awaiting()
|
||||
for reference in references:
|
||||
if reference.get("folder_type") != "output":
|
||||
continue
|
||||
key = self._awaiting_key(
|
||||
os.path.join(
|
||||
root, reference.get("subfolder", ""), reference["filename"]
|
||||
)
|
||||
)
|
||||
waiting = self._awaiting.pop(key, None)
|
||||
if waiting is not None:
|
||||
due.append(waiting[0])
|
||||
for image_path in due:
|
||||
self.schedule(image_path, HISTORY_ONLY)
|
||||
|
||||
def _recheck_history(self, image_path: str) -> None:
|
||||
"""The one history-only re-check of a keyless image (#342)."""
|
||||
name = os.path.basename(image_path)
|
||||
if not os.path.exists(image_path):
|
||||
return
|
||||
prompt_id = self._history_prompt(image_path)
|
||||
if prompt_id is None:
|
||||
self.logger.info(
|
||||
f"No certain prompt for {name} after its run ended; left unlinked"
|
||||
)
|
||||
return
|
||||
self.logger.info(f"History match after its run ended: {prompt_id} for {name}")
|
||||
try:
|
||||
metadata = self.metadata_extractor.extract_metadata(image_path)
|
||||
except Exception:
|
||||
metadata = None
|
||||
self.link_image_to_prompt(
|
||||
image_path,
|
||||
{"id": prompt_id, "from_history": True},
|
||||
metadata or {"file_info": self.get_basic_file_info(image_path)},
|
||||
)
|
||||
|
||||
def _history_prompt(self, image_path: str) -> Optional[int]:
|
||||
"""The stored prompt recorded history attributes to this file, if one."""
|
||||
try:
|
||||
|
||||
+32
-18
@@ -4,9 +4,14 @@ A PromptManager node registers the text it executed under a *run key*: a digest
|
||||
of the node and everything upstream of it in the queued API graph. The image
|
||||
embeds the same graph, so the image monitor can compute the same key for the
|
||||
image's positive prompt node. Equal keys mean the node ran with exactly the
|
||||
inputs (and upstream nodes) that produced this image. That is also the identity
|
||||
ComfyUI's own output cache uses: a node whose upstream is unchanged is served
|
||||
from the cache, so a re-run of an unchanged upstream reproduces the same text.
|
||||
inputs (and upstream nodes) that produced this image.
|
||||
|
||||
Equal keys do not always mean equal text. ComfyUI's cache signature also folds
|
||||
in each node's ``IS_CHANGED``, which the key cannot see: an upstream node that
|
||||
forces a re-run (a random line, an "always run" LLM) can give two runs one key
|
||||
and two texts. The key is therefore never used alone: the monitor also requires
|
||||
the entry to come from the run the snapshot taken when the image appeared names
|
||||
(``PromptTracker.pop_for_key``, ``ImageGenerationHandler._find_prompt_from_run``).
|
||||
|
||||
Pure functions; the only host dependency (the executing run id) is optional.
|
||||
"""
|
||||
@@ -23,20 +28,27 @@ except ImportError:
|
||||
RUN_KEY_VERSION = 1
|
||||
|
||||
|
||||
def is_link(value: Any) -> bool:
|
||||
"""ComfyUI's own test (``comfy_execution.graph_utils.is_link``), exactly.
|
||||
|
||||
A two-item list of a ``str`` node id and an ``int`` or ``float`` slot (a
|
||||
``bool`` slot passes too, as it does in ComfyUI). Tuples are values.
|
||||
"""
|
||||
if not isinstance(value, list):
|
||||
return False
|
||||
if len(value) != 2:
|
||||
return False
|
||||
if not isinstance(value[0], str):
|
||||
return False
|
||||
if not isinstance(value[1], int) and not isinstance(value[1], float):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _links(inputs: Any) -> List[str]:
|
||||
if not isinstance(inputs, dict):
|
||||
return []
|
||||
found = []
|
||||
for value in inputs.values():
|
||||
if (
|
||||
isinstance(value, (list, tuple))
|
||||
and len(value) == 2
|
||||
and isinstance(value[0], (str, int))
|
||||
and isinstance(value[1], int)
|
||||
and not isinstance(value[1], bool)
|
||||
):
|
||||
found.append(str(value[0]))
|
||||
return found
|
||||
return [value[0] for value in inputs.values() if is_link(value)]
|
||||
|
||||
|
||||
def prompt_node_key(graph: Any, node_id: Any) -> Optional[str]:
|
||||
@@ -79,8 +91,9 @@ def prompt_node_key(graph: Any, node_id: Any) -> Optional[str]:
|
||||
def image_run_keys(metadata: Any) -> Dict[str, str]:
|
||||
"""``{node_id: run key}`` for each positive PromptManager node of an image.
|
||||
|
||||
Empty without an embedded API graph: an image without one can never be
|
||||
proven to belong to a tracked run.
|
||||
Empty without an embedded API graph, and empty when any positive node has
|
||||
no key: one unkeyable node must not leave another looking like the only
|
||||
positive prompt (a two-pass workflow would link the wrong pass).
|
||||
"""
|
||||
graph = metadata.get("prompt") if isinstance(metadata, dict) else None
|
||||
if isinstance(graph, str):
|
||||
@@ -93,8 +106,9 @@ def image_run_keys(metadata: Any) -> Dict[str, str]:
|
||||
keys = {}
|
||||
for node_id in run_prompt_nodes(graph):
|
||||
key = prompt_node_key(graph, node_id)
|
||||
if key is not None:
|
||||
keys[node_id] = key
|
||||
if key is None:
|
||||
return {}
|
||||
keys[node_id] = key
|
||||
return keys
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user