diff --git a/database/pair_ingestion.py b/database/pair_ingestion.py index da1db08b..50011947 100644 --- a/database/pair_ingestion.py +++ b/database/pair_ingestion.py @@ -23,10 +23,26 @@ except ImportError: from utils.prompt_pair_resolver import RESOLVER_VERSION +def blank_positive(resolution: Dict[str, Any]) -> bool: + """A resolved pair whose stored positive text is empty (#342). + + A PromptManager node whose editor is empty encodes only its prepend and + append text, which is never stored as a prompt: such a result names no + prompt to link. + """ + if resolution.get("state") != "resolved": + return False + positive = resolution.get("positive") or {} + text = positive.get("storage_text", positive.get("text")) + return not (isinstance(text, str) and text.strip()) + + def ingest_resolved_pair( db: Any, image_path: str, metadata: Dict[str, Any], resolution: Dict[str, Any] ) -> Dict[str, Any]: state = resolution.get("state", "unresolved") + if state == "resolved" and blank_positive(resolution): + state = "unresolved" if state != "resolved": return {"status": state, "pairs_added": 0, "links_added": 0} with db._pair_transaction() as conn: diff --git a/database/run_attribution.py b/database/run_attribution.py new file mode 100644 index 00000000..f2a6c1ef --- /dev/null +++ b/database/run_attribution.py @@ -0,0 +1,122 @@ +"""Attribute an output file to a stored prompt through recorded history (#342). + +Execution history records, per run, the files ComfyUI reported +(``execution_run_outputs``) and the PromptManager texts the run used +(``execution_run_contexts``, whose ``positive_prompt_id`` hint is set only +while the raw text matches exactly one stored prompt). A file is attributed +only when all of this holds: + +- exactly one recorded run reported the file's output location (same output + root, subfolder and name); +- the file was last written while that run was executing, from its dequeue + to its terminal event (else the finish callback): a later run that history + did not record may have reused the name; +- every positive context of that run names the same stored prompt. + +Anything else returns None: the caller leaves the file unlinked. Read-only. +""" + +import os +import sqlite3 +from typing import Any, Optional + +try: + from ..utils.execution_outputs import snapshot_output_roots + from ..utils.library_location import import_folder_paths +except ImportError: + from utils.execution_outputs import snapshot_output_roots + from utils.library_location import import_folder_paths + +# File times may trail or lead the recorded run times a little (coarse file +# system clocks, the write finishing just after the host reported it). +MTIME_SLACK_MS = 2000 + +_RUNS_FOR_OUTPUT = ( + "SELECT DISTINCT o.run_id, r.observed_dequeue_at, " + "COALESCE(r.terminal_event_at, r.observed_callback_at) " + "FROM execution_run_outputs o JOIN execution_runs r ON r.run_id = o.run_id " + "WHERE o.folder_type = 'output' AND o.path_style = ? AND o.root_key = ? " + "AND o.subfolder = ? AND o.filename = ? LIMIT 2" +) + + +def host_output_directory() -> Optional[str]: + """ComfyUI's output directory, or None outside ComfyUI.""" + folder_paths = import_folder_paths() + if folder_paths is None: + return None + try: + value = folder_paths.get_output_directory() + except Exception: + return None + return value if isinstance(value, str) and value else None + + +def _location(image_path: str, output_root: str, path_style: str): + """(root key, subfolder, filename) as history records them, or None.""" + root_key = snapshot_output_roots(path_style, output_root, output_root).output_key + if root_key is None: + return None + try: + relative = os.path.relpath( + os.path.abspath(image_path), os.path.abspath(output_root) + ) + except ValueError: # another drive on Windows + return None + if os.path.isabs(relative) or relative == os.pardir: + return None + if relative.startswith(os.pardir + os.sep): + return None + subfolder, filename = os.path.split(relative) + if not filename: + return None + return root_key, subfolder, filename + + +def _written_during(image_path: str, started: Any, finished: Any) -> bool: + if not isinstance(started, int) or not isinstance(finished, int): + return False + try: + written = os.stat(image_path).st_mtime_ns // 1_000_000 + except OSError: + return False + return started - MTIME_SLACK_MS <= written <= finished + MTIME_SLACK_MS + + +def stored_run_prompt( + db: Any, + image_path: str, + output_root: Optional[str], + path_style: Optional[str] = None, +) -> Optional[int]: + """The stored prompt id history attributes to ``image_path``, else None.""" + if not image_path or not output_root: + return None + style = path_style or ("windows" if os.name == "nt" else "posix") + location = _location(image_path, output_root, style) + if location is None: + return None + root_key, subfolder, filename = location + try: + rows = ( + db.model.get_connection() + .execute(_RUNS_FOR_OUTPUT, (style, root_key, subfolder, filename)) + .fetchall() + ) + except sqlite3.Error: + return None # no history tables (a library without history) + if len(rows) != 1: + return None + run_id, started, finished = rows[0][0], rows[0][1], rows[0][2] + if not _written_during(image_path, started, finished): + return None + contexts = db.get_execution_contexts(run_id).get("contexts") or [] + named = { + context.get("positive_prompt_id") + for context in contexts + if context.get("positive") is not None + } + if len(named) != 1: + return None + (prompt_id,) = named + return prompt_id if isinstance(prompt_id, int) else None diff --git a/prompt_manager.py b/prompt_manager.py index 3f653e9a..4eb5d006 100644 --- a/prompt_manager.py +++ b/prompt_manager.py @@ -238,6 +238,8 @@ 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}") + else: + self._skip_empty_prompt(prompt, unique_id) # Perform standard CLIP text encoding using the combined text self.logger.debug( diff --git a/prompt_manager_base.py b/prompt_manager_base.py index 2c8550ff..74d95ddf 100644 --- a/prompt_manager_base.py +++ b/prompt_manager_base.py @@ -24,6 +24,7 @@ try: from .utils.prompt_tracker import get_prompt_tracker from .utils.hashing import generate_prompt_hash from .utils.prompt_graph import resolve_text, run_prompt_nodes + from .utils.run_keys import executing_run_id, prompt_node_key from .utils.usage_tracking import PENDING_FIRST_USE except ImportError: import sys @@ -35,6 +36,7 @@ except ImportError: from utils.prompt_tracker import get_prompt_tracker from utils.hashing import generate_prompt_hash from utils.prompt_graph import resolve_text, run_prompt_nodes + from utils.run_keys import executing_run_id, prompt_node_key from utils.usage_tracking import PENDING_FIRST_USE @@ -173,9 +175,44 @@ class PromptManagerBase: prompt_text=encoding_text, additional_data={**additional_data, "prompt_id": prompt_id}, push_to_queue=batch_item or not role["known"], + **self._run_identity(prompt_graph, unique_id), ) return prompt_id + @staticmethod + def _run_identity(prompt_graph: Any, unique_id: Any) -> Dict[str, Any]: + """The tracker's proof that an image belongs to this execution (#342).""" + return { + "run_key": prompt_node_key(prompt_graph, unique_id), + "run_id": executing_run_id(), + } + + def _skip_empty_prompt(self, prompt_graph: Any, unique_id: Any) -> None: + """The ``text`` editor is empty: save and track nothing (#342). + + 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). + """ + 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") + 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" + ) + 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}") + def _record_use(self, prompt_id: int, times: int = 1) -> None: """Count runs; a tracking failure must never block generation.""" try: diff --git a/prompt_manager_pair.py b/prompt_manager_pair.py index dcef20d4..1aeb2976 100644 --- a/prompt_manager_pair.py +++ b/prompt_manager_pair.py @@ -117,85 +117,27 @@ class PromptManagerPair(PromptManagerFocused): ) is_positive = "text" in used_fields pair_used = node_id in intrinsic_pair_nodes(prompt) if known else True - try: - pair_id = self.db.save_prompt_pair( - text, - negative_text, - category=category.strip() or None, - tags=self._parse_tags(tags), - ) - pair = self.db.get_prompt_pair(pair_id) - prompt_id = pair["positive_prompt_id"] - resolved_positive = ( - resolve_text_input(prompt, node_id, "text") if known else None - ) - resolved_negative = ( - resolve_text_input(prompt, node_id, "negative_text") if known else None - ) - positive_seen = resolved_positive is not None and generate_prompt_hash( - resolved_positive - ) == generate_prompt_hash(text) - pair_seen = ( - positive_seen - and resolved_negative is not None - and generate_prompt_hash(resolved_negative) - == generate_prompt_hash(negative_text) - ) - positive_pending = getattr(self, "pending_first_use", PENDING_FIRST_USE) - pair_pending = getattr( - self, "pending_pair_first_use", PENDING_PAIR_FIRST_USE - ) - counted = set() - for field, component_text, resolved_text, component_id in ( - ("text", text, resolved_positive, prompt_id), - ( - "negative_text", + if not text.strip(): + # Nothing to store: prepend/append are never the stored prompt (#342) + self._skip_empty_prompt(prompt, unique_id) + else: + try: + pair_id, prompt_id = self._save_and_track_pair( + text, negative_text, - resolved_negative, - pair["negative_prompt_id"], - ), - ): - component_hash = generate_prompt_hash(component_text) - if ( - field not in used_fields - or not component_id - or component_hash in counted - ): - continue - counted.add(component_hash) - seen_at_queue = ( - resolved_text is not None - and generate_prompt_hash(resolved_text) == component_hash + positive, + category=category, + tags=tags, + prepend_text=prepend_text, + append_text=append_text, + prompt=prompt, + unique_id=unique_id, + used_fields=used_fields, + is_positive=is_positive, + pair_used=pair_used, ) - runs = positive_pending.consume(component_hash) + int(not seen_at_queue) - if runs: - self.db.record_prompt_use(component_id, times=runs) - if pair_used: - runs = pair_pending.consume(pair_run_key(text, negative_text)) + int( - not pair_seen - ) - if runs: - self.db.record_prompt_pair_use(pair_id, times=runs) - if is_positive: - self.prompt_tracker.set_current_prompt( - prompt_text=positive, - additional_data={ - "prompt_id": prompt_id, - "pair_id": pair_id if pair_used else None, - "positive_text": positive, - "negative_text": negative_text, - "positive_storage_text": text.strip(), - "negative_storage_text": negative_text.strip(), - "prepend_text": prepend_text.strip(), - "append_text": append_text.strip(), - "final_text": positive, - "category": category.strip() or None, - "tags": self._parse_tags(tags) or [], - }, - push_to_queue=not positive_seen or (pair_used and not pair_seen), - ) - except Exception as error: - self.logger.warning(f"Could not save or track paired prompt: {error}") + except Exception as error: + self.logger.warning(f"Could not save or track paired prompt: {error}") try: self.comfyui_integration.register_prompt( f"promptmanager_pair_{time.time_ns()}", @@ -223,6 +165,102 @@ class PromptManagerPair(PromptManagerFocused): ) return result + def _save_and_track_pair( + self, + text: str, + negative_text: str, + positive: str, + *, + category: str, + tags: str, + prepend_text: str, + append_text: str, + prompt: Optional[Dict[str, Any]], + unique_id: Optional[str], + used_fields: set, + is_positive: bool, + pair_used: bool, + ) -> Tuple[Optional[int], Optional[int]]: + """Save the recipe, count its runs and register the positive for linking.""" + known = isinstance(prompt, dict) and unique_id is not None + node_id = str(unique_id) + pair_id = self.db.save_prompt_pair( + text, + negative_text, + category=category.strip() or None, + tags=self._parse_tags(tags), + ) + pair = self.db.get_prompt_pair(pair_id) + prompt_id = pair["positive_prompt_id"] + resolved_positive = ( + resolve_text_input(prompt, node_id, "text") if known else None + ) + resolved_negative = ( + resolve_text_input(prompt, node_id, "negative_text") if known else None + ) + positive_seen = resolved_positive is not None and generate_prompt_hash( + resolved_positive + ) == generate_prompt_hash(text) + pair_seen = ( + positive_seen + and resolved_negative is not None + and generate_prompt_hash(resolved_negative) + == generate_prompt_hash(negative_text) + ) + positive_pending = getattr(self, "pending_first_use", PENDING_FIRST_USE) + pair_pending = getattr(self, "pending_pair_first_use", PENDING_PAIR_FIRST_USE) + counted = set() + for field, component_text, resolved_text, component_id in ( + ("text", text, resolved_positive, prompt_id), + ( + "negative_text", + negative_text, + resolved_negative, + pair["negative_prompt_id"], + ), + ): + component_hash = generate_prompt_hash(component_text) + if ( + field not in used_fields + or not component_id + or component_hash in counted + ): + continue + counted.add(component_hash) + seen_at_queue = ( + resolved_text is not None + and generate_prompt_hash(resolved_text) == component_hash + ) + runs = positive_pending.consume(component_hash) + int(not seen_at_queue) + if runs: + self.db.record_prompt_use(component_id, times=runs) + if pair_used: + runs = pair_pending.consume(pair_run_key(text, negative_text)) + int( + not pair_seen + ) + if runs: + self.db.record_prompt_pair_use(pair_id, times=runs) + if is_positive: + self.prompt_tracker.set_current_prompt( + prompt_text=positive, + additional_data={ + "prompt_id": prompt_id, + "pair_id": pair_id if pair_used else None, + "positive_text": positive, + "negative_text": negative_text, + "positive_storage_text": text.strip(), + "negative_storage_text": negative_text.strip(), + "prepend_text": prepend_text.strip(), + "append_text": append_text.strip(), + "final_text": positive, + "category": category.strip() or None, + "tags": self._parse_tags(tags) or [], + }, + push_to_queue=not positive_seen or (pair_used and not pair_seen), + **self._run_identity(prompt, unique_id), + ) + return pair_id, prompt_id + @classmethod def IS_CHANGED( cls, diff --git a/prompt_manager_text.py b/prompt_manager_text.py index 92f8133d..272db9a8 100644 --- a/prompt_manager_text.py +++ b/prompt_manager_text.py @@ -216,6 +216,8 @@ 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}") + else: + self._skip_empty_prompt(prompt, unique_id) # Register with ComfyUI integration for standard metadata compatibility node_id = f"promptmanagertext_{int(time.time() * 1000)}" diff --git a/py/api/admin.py b/py/api/admin.py index d8630fd1..637baa65 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -35,7 +35,8 @@ except ImportError: from utils.parallel import map_parallel try: - from ...database.pair_ingestion import ingest_resolved_pair + from ...database.pair_ingestion import blank_positive, ingest_resolved_pair + from ...database.run_attribution import host_output_directory, stored_run_prompt from ...utils.generation_params import ( PLACEHOLDERS, GenerationParamsError, @@ -45,7 +46,8 @@ try: from ...utils.prompt_pair_resolver import resolve_prompt_pair from ...utils.review_capture import reviewable except ImportError: - from database.pair_ingestion import ingest_resolved_pair + from database.pair_ingestion import blank_positive, ingest_resolved_pair + from database.run_attribution import host_output_directory, stored_run_prompt from utils.generation_params import ( PLACEHOLDERS, GenerationParamsError, @@ -1665,6 +1667,9 @@ class AdminRoutesMixin: return None _with_file_size(canonical, media_file) resolution = resolve_prompt_pair(canonical) + if blank_positive(resolution): + # Only prepend/append text: no stored prompt can own it (#342) + resolution = {**resolution, "state": "unresolved"} # Ambiguous without a choice: unresolved unless something links it. unchoosable = resolution["state"] == "ambiguous" and not reviewable(resolution) if unchoosable: @@ -1693,6 +1698,9 @@ 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) + if linked is not None: + return linked return "unresolved" if unchoosable else resolution["state"] elif resolution["state"] == "ambiguous": return "unresolved" if unchoosable else "ambiguous" @@ -1716,6 +1724,33 @@ class AdminRoutesMixin: return status return {"status": status, "title": self._scan_title(media_file, prompt_text)} + def _link_recorded_run(self, media_file, canonical): + """Link a file recorded history attributes to one stored prompt (#342). + + Returns the scan status, or None when history names no single stored + prompt for the file (``database.run_attribution``). Never creates a + prompt. + """ + try: + prompt_id = stored_run_prompt( + self.db, str(media_file), host_output_directory() + ) + except Exception: + self.logger.warning( + f"History attribution failed for {media_file.name}", exc_info=True + ) + return None + if prompt_id is None: + return None + try: + image_id = self.db.link_image_to_prompt( + prompt_id, str(media_file), metadata=canonical + ) + 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" + def _scan_title(self, media_file, text): """The linked prompt's Library title (its stored title, else ``text``'s first line), as the Library and the inspector name it.""" diff --git a/tests/comfyui/image-attribution.spec.mjs b/tests/comfyui/image-attribution.spec.mjs new file mode 100644 index 00000000..addb4da8 --- /dev/null +++ b/tests/comfyui/image-attribution.spec.mjs @@ -0,0 +1,89 @@ +// #342 on the real pinned host: the image monitor links an output only to a prompt +// that provably produced it. An unknown node (PMHistoryCached) stands in for the +// maintainer's LLM; KSampler is the fixture's deterministic stand-in. +import { randomUUID } from 'node:crypto'; +import { execFile } from 'node:child_process'; +import { promisify } from 'node:util'; +import { fileURLToPath } from 'node:url'; +import { expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs'; +import { test } from './execution-fixture.mjs'; + +const run = promisify(execFile); +const reader = fileURLToPath(new URL('./read_execution_history.py', import.meta.url)); + +async function readLinks(databasePath, filenames, texts = []) { + const { stdout } = await run(process.env.PM_TEST_PYTHON || 'python', + [reader, databasePath, '--links', ...filenames, '--texts', ...texts], { timeout: 5000, maxBuffer: 1024 * 1024 }); + return JSON.parse(stdout); +} + +// kind: the positive node; text: a typed editor or 'llm' (linked from the unknown +// node); prepend: 'llm' links the unknown node into prepend_text. +function graph({ kind = 'PromptManagerPair', text = '', prepend = '', llm = '', seed = 1, prefix }) { + const linked = value => value === 'llm' ? ['11', 0] : value; + const nodes = { + '10': { class_type: 'PMHistoryInputs', inputs: {} }, + '20': { class_type: kind, inputs: { text: linked(text), clip: ['10', 0], append_text: '', + ...(prepend ? { prepend_text: linked(prepend) } : {}), + ...(kind === 'PromptManagerPair' ? { negative_text: '' } : {}) } }, + '30': { class_type: 'CLIPTextEncode', inputs: { text: '', clip: ['10', 0] } }, + '40': { class_type: 'KSampler', inputs: { model: ['10', 1], latent_image: ['10', 2], + positive: ['20', 0], negative: kind === 'PromptManagerPair' ? ['20', 1] : ['30', 0], + seed, steps: 1, cfg: 1, sampler_name: 'euler', scheduler: 'normal', denoise: 1 } }, + '60': { class_type: 'PMHistoryImages', inputs: { value: 64, batch: 1, latent: ['40', 0] } }, + '70': { class_type: 'SaveImage', inputs: { images: ['60', 0], filename_prefix: prefix } }, + }; + if (llm) nodes['11'] = { class_type: 'PMHistoryCached', inputs: { value: llm } }; + return nodes; +} + +async function saved(request, prompt) { + const id = randomUUID(); + const response = await request.post('/prompt', { data: { prompt, prompt_id: id } }); + expect(response.ok()).toBe(true); + let item; + await expect.poll(async () => { + item = (await (await request.get(`/history/${id}`)).json())[id]; + return item?.status?.messages?.some(([kind]) => kind === 'execution_success') || false; + }, { timeout: 15000 }).toBe(true); + const [image] = item.outputs['70'].images; + expect(image.type).toBe('output'); + return image.filename; +} + +async function linkedTo(execution, filename, expected) { + await expect.poll(async () => (await readLinks(execution.databasePath, [filename])).links[filename], + { timeout: 20000, message: `${filename} is linked by the image monitor` }).toBe(expected); +} + +test.beforeEach(async ({ execution }) => { + expect((await execution.request.post('/pm_execution_fixture/metadata', { data: { disabled: false } })).ok()).toBe(true); +}); + +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()}`; + const earlier = `a woman holding a coffee cup ${tag}`; + const prepended = `llm prepend landscape ${tag}`, llmText = `llm text landscape ${tag}`; + + // A typed prompt runs first: before #342 the next unidentifiable image took it. + const first = await saved(execution.request, graph({ text: earlier, prefix: 'misslink/earlier' })); + await linkedTo(execution, first, earlier); + + // The maintainer's workflow: empty editor, the generated text arrives in prepend. + const orphan = await saved(execution.request, + graph({ prepend: 'llm', llm: prepended, prefix: 'misslink/prepend' })); + + // An LLM feeding the editor itself is stored and links through its run key; the + // monitor works in creation order, so the prepend image was processed first. + const llmGraph = graph({ kind: 'PromptManager', text: 'llm', llm: llmText, prefix: 'misslink/text' }); + const proven = await saved(execution.request, llmGraph); + await linkedTo(execution, proven, llmText); + 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. + const rerun = await saved(execution.request, { ...llmGraph, '40': { ...llmGraph['40'], + inputs: { ...llmGraph['40'].inputs, seed: 2 } } }); + await linkedTo(execution, rerun, llmText); +}); diff --git a/tests/comfyui/read_execution_history.py b/tests/comfyui/read_execution_history.py index ccbc106a..ba209a0d 100644 --- a/tests/comfyui/read_execution_history.py +++ b/tests/comfyui/read_execution_history.py @@ -1,6 +1,7 @@ """Bounded read-only SQLite inspection for the owned execution fixture.""" import argparse +from contextlib import closing import json from pathlib import Path import sqlite3 @@ -89,14 +90,48 @@ def read_runs(database: str, job_id: str, session_id: str | None = None) -> dict } +def read_links(database: str, filenames: list, texts: list) -> dict: + """Each output file's linked prompt text (None: unlinked), and which texts + are stored as prompts (#342).""" + uri = Path(database).resolve().as_uri() + "?mode=ro" + with closing(sqlite3.connect(uri, uri=True, timeout=0.25)) as connection: + connection.execute("PRAGMA query_only=ON") + links = {} + for name in filenames[:64]: + rows = connection.execute( + "SELECT p.text FROM generated_images g JOIN prompts p " + "ON p.id = g.prompt_id WHERE g.filename = ? LIMIT 2", + (name,), + ).fetchall() + values = [row[0] for row in rows] + links[name] = values[0] if len(values) == 1 else None + if len(values) > 1: + links[name] = "" + stored = { + text: connection.execute( + "SELECT 1 FROM prompts WHERE text = ? LIMIT 1", (text,) + ).fetchone() + is not None + for text in texts[:64] + } + return {"links": links, "stored": stored} + + if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("database") - parser.add_argument("job_id") + parser.add_argument("job_id", nargs="?") parser.add_argument("--session-id") + parser.add_argument("--links", nargs="*", default=None) + parser.add_argument("--texts", nargs="*", default=[]) arguments = parser.parse_args() - print( - json.dumps( - read_runs(arguments.database, arguments.job_id, arguments.session_id) + if arguments.links is not None: + print( + json.dumps(read_links(arguments.database, arguments.links, arguments.texts)) + ) + else: + print( + json.dumps( + read_runs(arguments.database, arguments.job_id, arguments.session_id) + ) ) - ) diff --git a/tests/test_image_linking.py b/tests/test_image_linking.py index f75d18c8..65cc0152 100644 --- a/tests/test_image_linking.py +++ b/tests/test_image_linking.py @@ -12,6 +12,7 @@ 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", @@ -103,8 +104,11 @@ class TestMetadataLinking(LinkingTestCase): 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}) + 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) @@ -122,9 +126,10 @@ class TestMetadataLinking(LinkingTestCase): 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) + 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") @@ -200,12 +205,14 @@ class TestForeignWorkflows(LinkingTestCase): self.handler._find_prompt_from_metadata(metadata), ("foreign", None) ) - def test_missing_metadata_is_unknown_and_uses_the_queue(self): + 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) ) - self.assertEqual(self._process(None), first) + # 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") @@ -224,14 +231,11 @@ class TestForeignWorkflows(LinkingTestCase): ("unknown", None), ) - def test_unknown_image_without_any_context_falls_back_to_the_latest_prompt(self): + def test_unknown_image_is_never_given_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._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") diff --git a/tests/test_image_monitor_singleton.py b/tests/test_image_monitor_singleton.py index 58c34b09..b0659103 100644 --- a/tests/test_image_monitor_singleton.py +++ b/tests/test_image_monitor_singleton.py @@ -326,7 +326,7 @@ class TestHandlerScheduling(unittest.TestCase): self.tmp = tempfile.TemporaryDirectory() self.addCleanup(self.tmp.cleanup) self.tracker = MagicMock() - self.tracker.get_current_prompt.return_value = {"id": 7} + self.tracker.keyed_snapshot.return_value = {"key": ({"id": 7}, "run", 1)} self.handler = im_mod.ImageGenerationHandler(MagicMock(), self.tracker) self.handler.processing_delay = 0 self.handler.settle_interval = 0.01 @@ -354,7 +354,14 @@ class TestHandlerScheduling(unittest.TestCase): self.handler.on_created(FakeEvent(os.path.join(self.tmp.name, "c.txt"))) self.handler.on_created(FakeEvent(self.tmp.name, is_directory=True)) self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0)) - self.assertEqual(self.processed, [(path, {"id": 7})]) + self.assertEqual(self.processed, [(path, {"key": ({"id": 7}, "run", 1)})]) + + def test_a_failed_snapshot_still_schedules_the_image(self): + self.tracker.keyed_snapshot.side_effect = RuntimeError("tracker gone") + path = self._image("d.png") + self.handler.on_created(FakeEvent(path)) + self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0)) + self.assertEqual(self.processed, [(path, None)]) def test_images_are_processed_in_creation_order(self): paths = [self._image(f"{i}.png") for i in range(5)] diff --git a/tests/test_prompt_attribution.py b/tests/test_prompt_attribution.py new file mode 100644 index 00000000..66b67bca --- /dev/null +++ b/tests/test_prompt_attribution.py @@ -0,0 +1,485 @@ +"""#342: an image is linked only to a prompt that provably produced it. + +The maintainer's run: an LLM writes the prompt into a PromptManagerPair's +``prepend_text`` while its ``text`` editor is empty. Prepend and append text are +never stored as the prompt (product decision), so nothing is saved, and the +image must stay unlinked instead of taking the newest prompt in the database or +a stale tracker entry. +""" + +import copy +import logging +import os +import sys +import tempfile +import threading +import unittest +from unittest import mock + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +try: + from tests.open_handles import assert_closed +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed + +from database.operations import PromptDatabase +from prompt_manager import PromptManager +from prompt_manager_pair import PromptManagerPair +from prompt_manager_text import PromptManagerText +from utils.hashing import generate_prompt_hash +from utils.image_monitor import ImageGenerationHandler +from utils.prompt_tracker import PromptTracker +from utils.run_keys import image_run_keys, prompt_node_key +from utils.usage_tracking import ( + PendingFirstUse, + queued_pair_runs, + queued_prompt_runs, +) + +LLM_TEXT = "Photorealistic aerial landscape of misty mountains at dawn" +EVIDENCE_KEY = "comfyassets_prompt_pairs" + + +def issue_graph(text="", seed=1): + """The workflow from #342 (subgraph ids kept), PromptManagerPair positive.""" + return { + "18:35": { + "class_type": "PromptManagerPair", + "inputs": { + "text": text, + "negative_text": "", + "prepend_text": ["18:13", 0], + "append_text": "", + "category": "", + "tags": "", + "clip": ["4", 1], + }, + }, + "18:13": {"class_type": "PreviewAny", "inputs": {"source": ["18:16", 0]}}, + "18:16": { + "class_type": "ComfySwitchNode", + "inputs": { + "switch": True, + "on_false": ["18:17", 0], + "on_true": ["18:11", 0], + }, + }, + "18:11": { + "class_type": "TextGenerate", + "inputs": {"prompt": ["18:17", 0], "seed": 7}, + }, + "18:17": { + "class_type": "PrimitiveStringMultiline", + "inputs": {"value": "Describe an aerial mountain landscape"}, + }, + "18:5": { + "class_type": "CLIPTextEncode", + "inputs": {"text": "blurry", "clip": ["4", 1]}, + }, + "4": { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "moss.safetensors"}, + }, + "3": { + "class_type": "KSampler", + "inputs": { + "model": ["4", 0], + "positive": ["18:35", 0], + "negative": ["18:5", 0], + "latent_image": ["6", 0], + "seed": seed, + }, + }, + "6": {"class_type": "EmptyLatentImage", "inputs": {"batch_size": 1}}, + } + + +def llm_text_graph(seed=1): + """An LLM feeding a PromptManager's ``text``: stored, but not in the metadata.""" + graph = issue_graph(seed=seed) + graph["18:35"] = { + "class_type": "PromptManager", + "inputs": {"text": ["18:13", 0], "clip": ["4", 1]}, + } + return graph + + +def pair_evidence(text="", prepend=LLM_TEXT): + """The execution record PromptManagerPair embeds in the saved image.""" + return { + "version": 1, + "nodes": { + "18:35": { + "records": [ + { + "text": text, + "negative_text": "", + "positive_text": " ".join(p for p in (prepend, text) if p), + "negative_effective_text": "", + "prepend_text": prepend, + "append_text": "", + } + ] + } + }, + } + + +def make_tracker(db): + """A real tracker without its cleanup thread or config import.""" + with mock.patch.object(PromptTracker, "__init__", lambda self, db: None): + tracker = PromptTracker.__new__(PromptTracker) + tracker.logger = logging.getLogger("test.prompt_attribution.tracker") + tracker.db_manager = db + tracker._local = threading.local() + tracker.active_prompts = {} + tracker.lock = threading.Lock() + tracker._prompt_queue = [] + tracker._queue_lock = threading.Lock() + tracker.cleanup_interval = 300 + tracker.prompt_timeout = 600 + return tracker + + +class DatabaseCase(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.addCleanup(assert_closed, self, self.tmp.name) + self.db = PromptDatabase(os.path.join(self.tmp.name, "prompts.db")) + self.addCleanup(self.db.close_all) + self.tracker = make_tracker(self.db) + + def save(self, text): + return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text)) + + +class NodeCase(DatabaseCase): + def node(self, cls): + node = cls.__new__(cls) + node.db = self.db + node.logger = logging.getLogger("test.prompt_attribution.node") + node.prompt_tracker = self.tracker + node.comfyui_integration = mock.Mock() + node._inject_lora_trigger_words = lambda text: text + node.pending_first_use = PendingFirstUse() + node.pending_pair_first_use = PendingFirstUse() + return node + + def clip(self): + clip = mock.Mock() + clip.tokenize.side_effect = lambda text: text + clip.encode_from_tokens_scheduled.side_effect = lambda t: ("cond", t) + return clip + + def count_rows(self, table): + conn = self.db.model.get_connection() + return conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] + + +class TestEmptyEditorSavesNothing(NodeCase): + """Prepend/append are never the stored prompt; an empty editor saves nothing.""" + + def test_pair_with_only_prepend_text_saves_nothing_and_logs_info(self): + graph = issue_graph() + with self.assertLogs("test.prompt_attribution.node", "INFO") as logs: + result = self.node(PromptManagerPair).encode_prompt( + self.clip(), + "", + negative_text="", + prepend_text=LLM_TEXT, + prompt=graph, + unique_id="18:35", + ) + self.assertEqual(result[2], LLM_TEXT, "the composed text is still encoded") + self.assertEqual(self.count_rows("prompts"), 0) + self.assertEqual(self.count_rows("prompt_pairs"), 0) + self.assertFalse( + any(record.levelno >= logging.WARNING for record in logs.records), + [record.getMessage() for record in logs.records], + ) + self.assertTrue(any("Prompt text is empty" in line for line in logs.output)) + + def test_prompt_manager_and_text_nodes_with_only_prepend_save_nothing(self): + graph = issue_graph() + graph["18:35"] = { + "class_type": "PromptManager", + "inputs": {"text": "", "prepend_text": ["18:13", 0], "clip": ["4", 1]}, + } + self.node(PromptManager).encode_prompt( + self.clip(), "", prepend_text=LLM_TEXT, prompt=graph, unique_id="18:35" + ) + graph["18:35"]["class_type"] = "PromptManagerText" + self.node(PromptManagerText).process_text( + "", prepend_text=LLM_TEXT, prompt=graph, unique_id="18:35" + ) + self.assertEqual(self.count_rows("prompts"), 0) + + def test_empty_editor_marks_its_run_key_so_no_older_entry_is_used(self): + graph = issue_graph() + key = prompt_node_key(graph, "18:35") + older = self.save("a woman holding a coffee cup") + self.tracker.set_current_prompt( + "a woman holding a coffee cup", + {"prompt_id": older}, + push_to_queue=True, + run_key=key, + run_id="run-1", + ) + self.node(PromptManagerPair).encode_prompt( + self.clip(), "", prepend_text=LLM_TEXT, prompt=graph, unique_id="18:35" + ) + latest, _run, _count = self.tracker.keyed_snapshot()[key] + self.assertTrue(latest.get("placeholder")) + popped = self.tracker.pop_for_keys([key]) + self.assertEqual(popped["id"], older, "an earlier queued item stays first") + self.assertTrue(self.tracker.pop_for_keys([key]).get("placeholder")) + + def test_empty_negative_node_is_not_marked(self): + graph = issue_graph(text="a typed positive") + graph["18:5"] = { + "class_type": "PromptManager", + "inputs": {"text": "", "clip": ["4", 1]}, + } + self.node(PromptManager).encode_prompt( + self.clip(), "", prompt=graph, unique_id="18:5" + ) + self.assertEqual(self.tracker.keyed_snapshot(), {}) + + def test_queue_hook_counts_nothing_for_an_empty_editor(self): + graph = issue_graph() + self.assertEqual(queued_prompt_runs(graph), []) + self.assertEqual(queued_pair_runs(graph), []) + + def test_typed_text_with_prepend_still_stores_only_the_editor_text(self): + graph = issue_graph(text="mountains") + self.node(PromptManagerPair).encode_prompt( + self.clip(), + "mountains", + prepend_text=LLM_TEXT, + prompt=graph, + unique_id="18:35", + ) + self.assertIsNotNone( + self.db.get_prompt_by_hash(generate_prompt_hash("mountains")) + ) + self.assertIsNone(self.db.get_prompt_by_hash(generate_prompt_hash(LLM_TEXT))) + context, _run, _count = self.tracker.keyed_snapshot()[ + prompt_node_key(graph, "18:35") + ] + self.assertEqual(context["text"], f"{LLM_TEXT} mountains") + + +class MonitorCase(DatabaseCase): + def setUp(self): + super().setUp() + fd, self.image = tempfile.mkstemp(suffix=".png", dir=self.tmp.name) + os.close(fd) + 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 process(self, metadata, snapshot=None): + before = len(self.linked) + with mock.patch.object( + self.handler.metadata_extractor, "extract_metadata", return_value=metadata + ): + self.handler.process_new_image(self.image, prompt_snapshot=snapshot) + return self.linked[-1] if len(self.linked) > before else None + + +class TestIssueReproduction(MonitorCase): + def test_llm_prepend_image_is_left_unlinked_not_given_the_newest_prompt(self): + """Failing before #342: linked to the newest prompt in the database.""" + self.save("a woman holding a coffee cup") + metadata = {"prompt": issue_graph(), EVIDENCE_KEY: pair_evidence()} + with self.assertLogs("prompt_manager.image_monitor", "INFO") as logs: + self.assertIsNone(self.process(metadata)) + self.assertTrue(any("left unlinked" in line for line in logs.output)) + + def test_a_stale_entry_of_another_workflow_is_never_used(self): + """Failing before #342: the live tracker's entry was taken.""" + stale = self.save("a woman holding a coffee cup") + other = issue_graph(text="a woman holding a coffee cup") + self.tracker.set_current_prompt( + "a woman holding a coffee cup", + {"prompt_id": stale}, + push_to_queue=True, + run_key=prompt_node_key(other, "18:35"), + run_id="run-0", + ) + metadata = {"prompt": issue_graph(), EVIDENCE_KEY: pair_evidence()} + self.assertIsNone(self.process(metadata, self.tracker.keyed_snapshot())) + self.assertIsNone(self.process(metadata)) + self.assertEqual(len(self.tracker._prompt_queue), 1, "not consumed") + + def test_an_entry_without_a_run_key_is_never_used(self): + stale = self.save("previous run") + self.tracker.set_current_prompt("previous run", {"prompt_id": stale}) + self.assertIsNone(self.process({"prompt": llm_text_graph()})) + + def test_an_image_without_metadata_is_left_unlinked(self): + """Failing before #342: popped the queue, else took the newest prompt.""" + self.save("latest") + self.assertIsNone(self.process(None)) + + def test_the_recency_fallback_is_gone(self): + self.assertFalse(hasattr(self.handler, "_get_fallback_prompt")) + + +class TestProvenRuns(MonitorCase): + """Entries the image's own graph proves still link (LLM into ``text``).""" + + def track(self, graph, text, run_id="run-1", push=True): + prompt_id = ( + self.save(text) + if not self.db.get_prompt_by_hash(generate_prompt_hash(text)) + else self.db.get_prompt_by_hash(generate_prompt_hash(text))["id"] + ) + self.tracker.set_current_prompt( + text, + {"prompt_id": prompt_id}, + push_to_queue=push, + run_key=prompt_node_key(graph, "18:35"), + run_id=run_id, + ) + return prompt_id + + def test_queued_entry_of_the_same_run_key_links(self): + graph = llm_text_graph() + prompt_id = self.track(graph, LLM_TEXT) + self.assertEqual(self.process({"prompt": graph}), prompt_id) + + def test_cached_node_rerun_links_through_the_snapshot(self): + graph = llm_text_graph(seed=1) + prompt_id = self.track(graph, LLM_TEXT) + self.assertEqual(self.process({"prompt": graph}), prompt_id) + # Only the sampler seed changed: the node was cached and did not run, + # but its upstream (the run key) is unchanged. + rerun = llm_text_graph(seed=2) + self.assertEqual( + prompt_node_key(rerun, "18:35"), prompt_node_key(graph, "18:35") + ) + self.assertEqual( + self.process({"prompt": rerun}, self.tracker.keyed_snapshot()), prompt_id + ) + + def test_snapshot_needs_a_known_run(self): + graph = llm_text_graph() + self.track(graph, LLM_TEXT, run_id=None, push=False) + self.assertIsNone(self.process({"prompt": graph})) + + def test_batch_items_link_in_order_and_the_latest_alone_proves_nothing(self): + graph = llm_text_graph() + first = self.track(graph, "item one") + second = self.track(graph, "item two") + self.assertEqual(self.process({"prompt": graph}), first) + self.assertEqual(self.process({"prompt": graph}), second) + # Queue drained: two entries in one run, so the latest is not proof + self.assertIsNone(self.process({"prompt": graph})) + + def test_an_empty_batch_item_keeps_the_queue_aligned(self): + graph = llm_text_graph() + key = prompt_node_key(graph, "18:35") + first = self.track(graph, "item one") + self.tracker.mark_empty(key, "run-1", push_to_queue=True) + third = self.track(graph, "item three") + self.assertEqual(self.process({"prompt": graph}), first) + self.assertIsNone(self.process({"prompt": graph})) + self.assertEqual(self.process({"prompt": graph}), third) + + def test_a_new_run_supersedes_queued_entries_of_the_same_key(self): + graph = llm_text_graph() + self.track(graph, "earlier run text", run_id="run-1") + newer = self.track(graph, "newer run text", run_id="run-2") + self.assertEqual(self.process({"prompt": graph}), newer) + + def test_other_keys_stay_queued_for_their_own_images(self): + mine, other = llm_text_graph(), copy.deepcopy(llm_text_graph()) + other["18:17"]["inputs"]["value"] = "Describe a beach" + theirs = self.track(other, "a beach") + prompt_id = self.track(mine, LLM_TEXT) + self.assertEqual(self.process({"prompt": mine}), prompt_id) + self.assertEqual(self.process({"prompt": other}), theirs) + + def test_two_positive_prompt_nodes_are_not_guessed(self): + graph = llm_text_graph() + self.track(graph, LLM_TEXT) + graph["30"] = { + "class_type": "PromptManager", + "inputs": {"text": ["18:13", 0], "clip": ["4", 1]}, + } + graph["31"] = dict(graph["3"]) + graph["31"]["inputs"] = {**graph["3"]["inputs"], "positive": ["30", 0]} + self.assertEqual(len(image_run_keys({"prompt": graph})), 2) + self.assertIsNone(self.process({"prompt": graph})) + + def test_metadata_still_wins_and_leaves_the_queue_alone(self): + typed = self.save("a typed prompt") + graph = issue_graph(text="a typed prompt") + self.tracker.set_current_prompt( + "x", + {"prompt_id": typed}, + run_key=prompt_node_key(graph, "18:35"), + run_id="r", + ) + self.assertEqual(self.process({"prompt": graph}), typed) + self.assertEqual(len(self.tracker._prompt_queue), 1) + + +class TestRunKeys(unittest.TestCase): + def test_key_covers_the_upstream_and_ignores_downstream_and_titles(self): + graph = issue_graph() + key = prompt_node_key(graph, "18:35") + changed = copy.deepcopy(graph) + changed["3"]["inputs"]["seed"] = 99 + changed["18:35"]["_meta"] = {"title": "renamed"} + self.assertEqual(prompt_node_key(changed, "18:35"), key) + changed["18:11"]["inputs"]["seed"] = 8 + self.assertNotEqual(prompt_node_key(changed, "18:35"), key) + + def test_unknowable_keys_are_none(self): + graph = issue_graph() + self.assertIsNone(prompt_node_key(None, "1")) + self.assertIsNone(prompt_node_key(graph, None)) + self.assertIsNone(prompt_node_key(graph, "missing")) + broken = copy.deepcopy(graph) + broken["18:13"]["inputs"]["source"] = ["gone", 0] + self.assertIsNone(prompt_node_key(broken, "18:35")) + broken = copy.deepcopy(graph) + broken["18:17"]["inputs"]["value"] = object() + self.assertIsNone(prompt_node_key(broken, "18:35")) + + def test_image_keys_read_string_graphs_and_reject_garbage(self): + import json + + graph = issue_graph() + self.assertEqual( + image_run_keys({"prompt": json.dumps(graph)}), + {"18:35": prompt_node_key(graph, "18:35")}, + ) + self.assertEqual(image_run_keys({"prompt": "{not json"}), {}) + self.assertEqual(image_run_keys(None), {}) + self.assertEqual(image_run_keys({"prompt": []}), {}) + + def test_cyclic_and_huge_closures_have_no_key(self): + from utils.prompt_graph import MAX_VISITS + + cyclic = { + "1": {"class_type": "A", "inputs": {"x": ["2", 0]}}, + "2": {"class_type": "B", "inputs": {"x": ["1", 0]}}, + } + self.assertIsNotNone(prompt_node_key(cyclic, "1")) + chain = { + str(i): {"class_type": "N", "inputs": {"x": [str(i + 1), 0]}} + for i in range(MAX_VISITS + 1) + } + chain[str(MAX_VISITS + 1)] = {"class_type": "N", "inputs": {}} + self.assertIsNone(prompt_node_key(chain, "0")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_review_image_links_tool.py b/tests/test_review_image_links_tool.py new file mode 100644 index 00000000..98257760 --- /dev/null +++ b/tests/test_review_image_links_tool.py @@ -0,0 +1,175 @@ +"""tools/review_image_links.py: a dry run by default, changes only named images.""" + +import io +import json +import os +import shutil +import tempfile +import unittest +import unittest.mock + +from PIL import Image +from PIL.PngImagePlugin import PngInfo + +try: + from tests.open_handles import assert_closed +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed + +from database.operations import PromptDatabase +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 + + +def png(path, prompt=None, evidence=None): + info = PngInfo() + if prompt is not None: + info.add_text("prompt", json.dumps(prompt)) + if evidence is not None: + info.add_text("comfyassets_prompt_pairs", json.dumps(evidence)) + Image.new("RGB", (4, 4)).save(path, pnginfo=info) + + +class ReviewToolTests(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmp, True) + self.addCleanup(assert_closed, self, self.tmp) + self.output = os.path.join(self.tmp, "output") + os.makedirs(self.output) + self.db_path = os.path.join(self.tmp, "prompts.db") + db = PromptDatabase(self.db_path) + try: + save = lambda text: db.save_prompt( # noqa: E731 + text=text, prompt_hash=generate_prompt_hash(text) + ) + self.coffee = save("a woman holding a coffee cup") + self.typed = save("snowy peaks") + # The six #342 images: only prepend text, linked by the old fallback + self.guessed = [] + for n in range(9, 12): + path = os.path.join(self.output, f"moss_{n:05}_.png") + png(path, issue_graph(), pair_evidence()) + self.guessed.append( + db.link_image_to_prompt(self.coffee, path, metadata={}) + ) + # A typed image linked to the wrong prompt, and a correct one + wrong = os.path.join(self.output, "typed_wrong.png") + png(wrong, issue_graph(text="snowy peaks")) + self.wrong = db.link_image_to_prompt(self.coffee, wrong, metadata={}) + right = os.path.join(self.output, "typed_right.png") + png(right, issue_graph(text="snowy peaks")) + self.right = db.link_image_to_prompt(self.typed, right, metadata={}) + gone = os.path.join(self.output, "gone.png") + self.gone = db.link_image_to_prompt(self.coffee, gone, metadata={}) + finally: + db.close_all() + + def run_tool(self, *args): + out = io.StringIO() + code = tool.main([self.db_path, "--output-dir", self.output, *args], out=out) + return code, out.getvalue() + + def links(self): + db = PromptDatabase(self.db_path) + try: + rows = db.model.get_connection().execute( + "SELECT id, prompt_id FROM generated_images" + ) + return {row[0]: row[1] for row in rows} + finally: + db.close_all() + + 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])} + for image_id in self.guessed: + self.assertEqual(verdicts[image_id]["verdict"], "unlink") + self.assertEqual(verdicts[image_id]["sources"], {}) + self.assertEqual(verdicts[self.wrong]["verdict"], "relink") + self.assertEqual(verdicts[self.wrong]["target"], self.typed) + self.assertEqual(verdicts[self.right]["verdict"], "keep") + self.assertEqual(verdicts[self.gone]["verdict"], "missing") + self.assertIn("Dry run: nothing was changed.", out) + self.assertEqual(self.links(), before) + self.assertEqual( + sorted(os.listdir(self.tmp)), ["output", "prompts.db"], "no side files" + ) + + def test_prompt_selection_and_text_report(self): + code, out = self.run_tool("--prompt", str(self.coffee)) + self.assertEqual(code, 0) + self.assertIn("unlink: 3", out) + self.assertIn("relink: 1", out) + self.assertIn("missing: 1", out) + self.assertIn(f"-> {self.typed}", out) + + def test_apply_backs_up_then_changes_only_named_images(self): + named = f"{self.guessed[0]}-{self.guessed[1]},{self.wrong}" + code, out = self.run_tool( + "--ids", named, "--prompt", str(self.coffee), "--apply" + ) + self.assertEqual(code, 0) + self.assertIn("Applied: 1 relinked, 2 unlinked", out) + links = self.links() + self.assertNotIn(self.guessed[0], links) + self.assertNotIn(self.guessed[1], links) + self.assertEqual(links[self.guessed[2]], self.coffee, "not named: kept") + self.assertEqual(links[self.wrong], self.typed) + self.assertEqual(links[self.gone], self.coffee, "missing files are kept") + backups = [n for n in os.listdir(self.tmp) if ".pre-relink_" in n] + self.assertEqual(len(backups), 1) + restored = PromptDatabase(os.path.join(self.tmp, backups[0])) + try: + count = ( + restored.model.get_connection() + .execute("SELECT COUNT(*) FROM generated_images") + .fetchone()[0] + ) + finally: + restored.close_all() + self.assertEqual(count, 6) + + def test_arguments_are_validated(self): + for args in ( + [], + ["--ids", "5-2"], + ["--ids", "x"], + ["--prompt", "1", "--apply"], + ): + with self.subTest(args=args), self.assertRaises(SystemExit): + with unittest.mock.patch("sys.stderr", io.StringIO()): + self.run_tool(*args) + with ( + self.assertRaises(SystemExit), + unittest.mock.patch("sys.stderr", io.StringIO()), + ): + tool.main( + [os.path.join(self.tmp, "none.db"), "--output-dir", "x", "--ids", "1"] + ) + self.assertEqual(tool.parse_ids(" 3, 1-2 ,"), [1, 2, 3]) + + def test_conflicting_sources_are_left_to_the_maintainer(self): + db = PromptDatabase(self.db_path) + try: + handler = tool.ImageGenerationHandler(db, None) + row = { + "id": 1, + "prompt_id": self.coffee, + "image_path": os.path.join(self.output, "typed_right.png"), + } + with unittest.mock.patch.object( + tool, "certain_prompts", return_value={"scan": 1, "history": 2} + ): + verdict = tool.judge(db, handler, row, self.output) + finally: + db.close_all() + self.assertEqual(verdict["verdict"], "conflict") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_run_attribution.py b/tests/test_run_attribution.py new file mode 100644 index 00000000..585f619e --- /dev/null +++ b/tests/test_run_attribution.py @@ -0,0 +1,237 @@ +"""#342: recorded history attributes a file only to one certain, stored prompt.""" + +import json +import os +import sqlite3 +import time +import unittest +from pathlib import Path +from unittest import mock + +from database.run_attribution import stored_run_prompt +from tests import test_execution_history_database as history_tests +from tests.test_api_admin import AdminAPITestCase +from tests.test_execution_context_database import snapshot +from tests.test_prompt_attribution import issue_graph, llm_text_graph, make_tracker +from utils.execution_contexts import ContextBundle, RecipeContext +from utils.execution_outputs import ( + OutputBundle, + OutputReference, + snapshot_output_roots, +) +from utils.hashing import generate_prompt_hash +from utils.image_monitor import ImageGenerationHandler + +STYLE = "windows" if os.name == "nt" else "posix" +LLM_TEXT = "Photorealistic aerial landscape of misty mountains at dawn" + + +def now_ms(): + return time.time_ns() // 1_000_000 + + +class RecordedRuns: + """Write runs into execution history the way the capture pipeline does.""" + + start = history_tests.HistoryDatabaseTests.start + finish = history_tests.HistoryDatabaseTests.finish + + def record_run(self, db, output_root, files, positives, item=0, window=None): + started, finished = window or (now_ms() - 5000, now_ms() + 5000) + start = self.start(db, item, started) + db.write_execution_start(start) + finish = self.finish(start, time=finished) + db.write_execution_finish(finish) + db.write_execution_contexts( + start.run_id, + start.storage_epoch, + ContextBundle( + tuple( + RecipeContext(str(4 + i), snapshot(text), snapshot(""), "inferred") + for i, text in enumerate(positives) + ), + "complete", + "captured", + ), + ) + root = snapshot_output_roots(STYLE, output_root, output_root).output_key + references = tuple( + OutputReference("9", i, "output", sub, name, STYLE, root) + for i, (sub, name) in enumerate(files) + ) + db.write_execution_outputs_for_finish( + finish, OutputBundle(references, "complete", "reported") + ) + return start.run_id + + +class TestStoredRunPrompt(RecordedRuns, unittest.TestCase): + setUp = history_tests.HistoryDatabaseTests.setUp + close_databases = history_tests.HistoryDatabaseTests.close_databases + database = history_tests.HistoryDatabaseTests.database + + def setUp(self): # noqa: F811 - the shared temp dir and database list first + history_tests.HistoryDatabaseTests.setUp(self) + self.db = self.database() + self.output = str(Path(self.temp.name) / "output") + os.makedirs(os.path.join(self.output, "sub")) + self.image = os.path.join(self.output, "sub", "img_00001_.png") + Path(self.image).write_bytes(b"png") + + def save(self, text): + return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text)) + + def attribute(self, path=None): + return stored_run_prompt(self.db, path or self.image, self.output) + + def test_one_run_naming_one_stored_prompt_attributes_the_file(self): + prompt_id = self.save("an llm prompt") + self.record_run( + self.db, self.output, [("sub", "img_00001_.png")], ["an llm prompt"] + ) + self.assertEqual(self.attribute(), prompt_id) + + def test_text_that_is_not_stored_attributes_nothing(self): + """The #342 run: raw editor text was empty, nothing was stored.""" + self.save("a woman holding a coffee cup") + self.record_run(self.db, self.output, [("sub", "img_00001_.png")], [""]) + self.assertIsNone(self.attribute()) + + def test_two_runs_reporting_the_same_file_attribute_nothing(self): + self.save("first") + self.save("second") + self.record_run(self.db, self.output, [("sub", "img_00001_.png")], ["first"]) + self.record_run( + self.db, self.output, [("sub", "img_00001_.png")], ["second"], item=1 + ) + self.assertIsNone(self.attribute()) + + def test_a_file_written_outside_the_run_attributes_nothing(self): + self.save("an llm prompt") + self.record_run( + self.db, + self.output, + [("sub", "img_00001_.png")], + ["an llm prompt"], + window=(1_000_000, 2_000_000), + ) + self.assertIsNone(self.attribute()) + + def test_contexts_naming_different_prompts_attribute_nothing(self): + self.save("left") + self.save("right") + self.record_run( + self.db, self.output, [("sub", "img_00001_.png")], ["left", "right"] + ) + self.assertIsNone(self.attribute()) + + def test_files_outside_the_output_root_or_unrecorded_attribute_nothing(self): + self.save("an llm prompt") + self.record_run( + self.db, self.output, [("sub", "img_00001_.png")], ["an llm prompt"] + ) + outside = os.path.join(self.temp.name, "img_00001_.png") + Path(outside).write_bytes(b"png") + self.assertIsNone(self.attribute(outside)) + other = os.path.join(self.output, "sub", "other.png") + Path(other).write_bytes(b"png") + self.assertIsNone(self.attribute(other)) + self.assertIsNone(stored_run_prompt(self.db, self.image, None)) + self.assertIsNone(stored_run_prompt(self.db, self.image, "relative/root")) + + def test_a_library_without_history_tables_attributes_nothing(self): + broken = mock.Mock() + broken.model.get_connection.return_value.execute.side_effect = ( + sqlite3.OperationalError("no such table") + ) + self.assertIsNone(stored_run_prompt(broken, self.image, self.output)) + + def test_the_monitor_links_from_history_when_nothing_else_is_certain(self): + prompt_id = self.save(LLM_TEXT) + self.record_run(self.db, self.output, [("sub", "img_00001_.png")], [LLM_TEXT]) + handler = ImageGenerationHandler(self.db, make_tracker(self.db)) + linked = [] + handler.link_image_to_prompt = lambda path, prompt, meta: linked.append( + prompt["id"] + ) + with ( + mock.patch( + "utils.image_monitor.host_output_directory", return_value=self.output + ), + mock.patch.object( + handler.metadata_extractor, + "extract_metadata", + return_value={"prompt": llm_text_graph()}, + ), + ): + handler.process_new_image(self.image) + self.assertEqual(linked, [prompt_id]) + + +class TestScanUsesHistory(RecordedRuns, AdminAPITestCase): + """The output scan applies the same rules as the monitor.""" + + def linked_prompt(self, path): + row = ( + self.api.db.model.get_connection() + .execute( + "SELECT prompt_id FROM generated_images WHERE image_path=?", + (str(path),), + ) + .fetchone() + ) + return row[0] if row else None + + def write(self, name): + path = self.output_dir / name + path.write_bytes(b"png") + return path + + def test_an_llm_text_file_links_to_the_prompt_its_run_stored(self): + db = self.api.db + prompt_id = db.save_prompt( + text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT) + ) + path = self.write("llm_00001_.png") + self.record_run(db, str(self.output_dir), [("", path.name)], [LLM_TEXT]) + meta = {"prompt": json.dumps(llm_text_graph())} + self.assertEqual(self.api._ingest_scanned_file(path, meta), "linked") + self.assertEqual(self.linked_prompt(path), prompt_id) + + 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 + db.save_prompt( + text="a woman holding a coffee cup", + prompt_hash=generate_prompt_hash("a woman holding a coffee cup"), + ) + path = self.write("moss_00009_.png") + self.record_run(db, str(self.output_dir), [("", path.name)], [""]) + evidence = { + "version": 1, + "nodes": { + "18:35": { + "records": [ + { + "text": "", + "negative_text": "", + "positive_text": LLM_TEXT, + "negative_effective_text": "", + "prepend_text": LLM_TEXT, + "append_text": "", + } + ] + } + }, + } + meta = { + "prompt": json.dumps(issue_graph()), + "comfyassets_prompt_pairs": json.dumps(evidence), + } + self.assertEqual(self.api._ingest_scanned_file(path, meta), "unresolved") + self.assertIsNone(self.linked_prompt(path)) + self.assertEqual(db.list_prompt_pairs()["pairs"], []) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_usage_tracking.py b/tests/test_usage_tracking.py index eeabad5f..11c7282c 100644 --- a/tests/test_usage_tracking.py +++ b/tests/test_usage_tracking.py @@ -477,11 +477,20 @@ class TestPendingFirstUse(unittest.TestCase): class FakeTracker: def __init__(self): self.calls = [] + self.runs = [] - def set_current_prompt(self, prompt_text, additional_data=None, push_to_queue=True): + def set_current_prompt( + self, prompt_text, additional_data=None, push_to_queue=True, **run + ): self.calls.append({"text": prompt_text, "push_to_queue": push_to_queue}) + self.runs.append(run) return "exec-1" + def mark_empty(self, run_key=None, run_id=None, push_to_queue=False): + self.calls.append( + {"empty": True, "run_key": run_key, "push_to_queue": push_to_queue} + ) + class TestNodeRoles(DbTestCase): def _node(self): diff --git a/tools/review_image_links.py b/tools/review_image_links.py new file mode 100644 index 00000000..c07d1b97 --- /dev/null +++ b/tools/review_image_links.py @@ -0,0 +1,284 @@ +"""Review, and on request repair, image links that nothing certain supports (#342). + +Before #342 the image monitor linked an image it could not identify to the +newest prompt in the database or to a stale tracker entry. This tool judges the +selected images again by today's rules, using only what is certain: + +- the image's own metadata, as the monitor reads it; +- the output scan's resolver, and its legacy browser-parser floor (#170); +- recorded execution history (``database.run_attribution``). + +Each image gets one verdict: + +``keep`` a certain source names the prompt it is linked to +``relink`` every certain source names one other stored prompt +``unlink`` nothing certain names a prompt (the link was a guess) +``conflict`` certain sources disagree: decide by hand +``missing`` the file is gone: leave it to Maintenance's missing links + +Dry run by default: the database is copied (SQLite backup from a read-only +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 +(``.pre-relink_``). 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. + + python3 tools/review_image_links.py --output-dir \\ + --ids 58005-58010 + python3 tools/review_image_links.py --output-dir --prompt 5278 + python3 tools/review_image_links.py --output-dir \\ + --ids 58005-58010 --apply + +Stop ComfyUI before ``--apply``. Run from this repository with its venv; the +package's logger may write to its usual log file. +""" + +import argparse +import datetime +import json +import os +import sqlite3 +import sys +import tempfile +from contextlib import closing +from pathlib import Path + +_ROOT = Path(__file__).resolve().parents[1] +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +from database.maintenance import _delete_ids # noqa: E402 +from database.operations import PromptDatabase # noqa: E402 +from database.pair_ingestion import blank_positive # noqa: E402 +from database.run_attribution import stored_run_prompt # noqa: E402 +from utils.generation_params import ( # noqa: E402 + PLACEHOLDERS, + GenerationParamsError, + extract_generation_params, +) +from utils.hashing import generate_prompt_hash # noqa: E402 +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") +# Ids per query (well under SQLite's bound-parameter limit) +ID_BATCH = 500 + + +def parse_ids(text): + """``"1,4-6"`` -> ``[1, 4, 5, 6]``; raises ValueError on anything else.""" + ids = set() + for part in (text or "").split(","): + part = part.strip() + if not part: + continue + low, _, high = part.partition("-") + first, last = int(low), int(high or low) + if first < 1 or last < first or last - first > 100_000: + raise ValueError(f"Bad id range: {part}") + ids.update(range(first, last + 1)) + return sorted(ids) + + +def _stored_id(db, text): + if not isinstance(text, str) or not text.strip(): + return None + row = db.get_prompt_by_hash(generate_prompt_hash(text.strip())) + return row["id"] if row else None + + +def _parser_text(metadata): + """The legacy browser parser's positive prompt (the scan's floor), or None.""" + records = {key: metadata.get(key) for key in ("prompt", "workflow")} + try: + positive = extract_generation_params(records)["positivePrompt"] + except (GenerationParamsError, ValueError, TypeError, KeyError): + return None + if not isinstance(positive, str) or positive == PLACEHOLDERS["positivePrompt"]: + return None + return positive.strip() or None + + +def certain_prompts(db, handler, path, metadata, output_dir): + """``{source: prompt id}`` of every certain source that names a stored prompt.""" + found = {} + if metadata: + state, prompt = handler._find_prompt_from_metadata(metadata) + if state == "linked" and prompt: + found["metadata"] = prompt["id"] + resolution = resolve_prompt_pair(metadata) + if resolution["state"] == "resolved" and not blank_positive(resolution): + positive = resolution["positive"] + scanned = _stored_id(db, positive.get("storage_text", positive["text"])) + else: + scanned = _stored_id(db, _parser_text(metadata)) + if scanned is not None: + found["scan"] = scanned + recorded = stored_run_prompt(db, path, output_dir) + if recorded is not None: + found["history"] = recorded + return found + + +def judge(db, handler, row, output_dir): + """One image's verdict dict.""" + path = row["image_path"] + verdict = { + "image_id": row["id"], + "file": path, + "linked_prompt": row["prompt_id"], + "sources": {}, + "target": None, + } + if not os.path.isfile(path): + return {**verdict, "verdict": "missing"} + try: + metadata = handler.metadata_extractor.extract_metadata(path) + except Exception: + metadata = None + sources = certain_prompts(db, handler, path, metadata, output_dir) + named = set(sources.values()) + verdict["sources"] = sources + if not named: + return {**verdict, "verdict": "unlink"} + if len(named) > 1: + return {**verdict, "verdict": "conflict"} + (target,) = named + if target == row["prompt_id"]: + return {**verdict, "verdict": "keep", "target": target} + return {**verdict, "verdict": "relink", "target": target} + + +def selected_rows(db, ids, prompt_id): + """The named images and every image of ``prompt_id``, by id.""" + conn = db.model.get_connection() + columns = "SELECT id, prompt_id, image_path FROM generated_images " + rows = {} + for start in range(0, len(ids), ID_BATCH): + batch = ids[start : start + ID_BATCH] + marks = ",".join("?" for _ in batch) + query = columns + f"WHERE id IN ({marks})" # only ? marks + rows.update((row["id"], row) for row in conn.execute(query, batch)) + if prompt_id is not None: + found = conn.execute(columns + "WHERE prompt_id = ?", (prompt_id,)) + rows.update((row["id"], row) for row in found) + return [rows[key] for key in sorted(rows)] + + +def review(db, ids, prompt_id, output_dir): + handler = ImageGenerationHandler(db, None) + return [ + judge(db, handler, row, output_dir) for row in selected_rows(db, ids, prompt_id) + ] + + +def apply(db, verdicts, ids): + """Write the relink/unlink verdicts of explicitly named ids; returns counts.""" + 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 + ] + with db._pair_transaction() as conn: + for item in relink: + # A pair link records the old positive: it goes with the old link + conn.execute( + "DELETE FROM image_prompt_pairs WHERE image_id=?", (item["image_id"],) + ) + conn.execute( + "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} + + +def backup_database(source, target): + """SQLite online backup of ``source`` (opened read-only) into ``target``.""" + uri = Path(os.path.abspath(source)).as_uri() + "?mode=ro" + with ( + closing(sqlite3.connect(uri, uri=True)) as src, + closing(sqlite3.connect(target)) as dst, + ): + src.backup(dst) + + +def _print(verdicts, out): + for item in verdicts: + sources = ", ".join(f"{k}={v}" for k, v in sorted(item["sources"].items())) + target = f" -> {item['target']}" if item["verdict"] == "relink" else "" + linked = item["linked_prompt"] + out.write( + f"{item['image_id']:>8} {item['verdict']:<8} prompt {linked}{target}" + f" [{sources or 'nothing certain'}] {item['file']}\n" + ) + counts = {name: sum(v["verdict"] == name for v in verdicts) for name in VERDICTS} + out.write(" ".join(f"{name}: {count}" for name, count in counts.items()) + "\n") + + +def main(argv=None, out=sys.stdout): + parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0]) + parser.add_argument("database", help="the library database (prompts.db)") + parser.add_argument( + "--output-dir", required=True, help="ComfyUI's output directory (history)" + ) + parser.add_argument("--ids", default="", help="image ids, e.g. 58005-58010,58020") + parser.add_argument("--prompt", type=int, help="every image linked to this prompt") + parser.add_argument( + "--apply", action="store_true", help="write relink/unlink for --ids images" + ) + parser.add_argument("--json", action="store_true", help="print JSON verdicts") + args = parser.parse_args(argv) + try: + ids = parse_ids(args.ids) + except ValueError as error: + parser.error(str(error)) + if not ids and args.prompt is None: + parser.error("select images with --ids and/or --prompt") + if args.apply and not ids: + parser.error("--apply changes only images named with --ids") + if not os.path.isfile(args.database): + parser.error(f"no database at {args.database}") + + if not args.apply: + with tempfile.TemporaryDirectory() as scratch: + copy = os.path.join(scratch, "review.db") + backup_database(args.database, copy) + db = PromptDatabase(copy) + try: + verdicts = review(db, ids, args.prompt, args.output_dir) + finally: + db.close_all() + _report(verdicts, args.json, out) + out.write("Dry run: nothing was changed.\n") + return 0 + + stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + backup = f"{args.database}.pre-relink_{stamp}" + backup_database(args.database, backup) + out.write(f"Backup: {backup}\n") + db = PromptDatabase(args.database) + try: + verdicts = review(db, ids, args.prompt, args.output_dir) + _report(verdicts, args.json, out) + result = apply(db, verdicts, ids) + finally: + db.close_all() + out.write( + f"Applied: {result['relinked']} relinked, {result['unlinked']} unlinked\n" + ) + return 0 + + +def _report(verdicts, as_json, out): + if as_json: + out.write(json.dumps(verdicts, indent=2, sort_keys=True) + "\n") + else: + _print(verdicts, out) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/vendor/playwright.execution.config.mjs b/tools/vendor/playwright.execution.config.mjs index 82f1e054..911e1407 100644 --- a/tools/vendor/playwright.execution.config.mjs +++ b/tools/vendor/playwright.execution.config.mjs @@ -6,7 +6,7 @@ if (host.href !== 'http://127.0.0.1:8196/') throw new Error('Execution tests req export default defineConfig({ testDir: '../../tests/comfyui', - testMatch: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'library-location.spec.mjs'], + testMatch: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'library-location.spec.mjs', 'image-attribution.spec.mjs'], outputDir: './test-results/executions/cases', timeout: 60000, workers: 1, diff --git a/utils/image_monitor.py b/utils/image_monitor.py index b6fbc9bf..377951c6 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -18,8 +18,9 @@ Typical usage: The system automatically: - Detects new image files in monitored directories - Extracts ComfyUI workflow metadata from PNG chunks -- Links images to active prompts using the prompt tracker -- Handles fallback linking when no active prompt is available +- Links images to tracked prompts only when the image's own graph proves the + tracked entry belongs to its run (``utils.run_keys``) +- Leaves an image unlinked when nothing certain identifies its prompt (#342) - Provides status information and monitoring control """ @@ -36,12 +37,15 @@ from .metadata_extractor import ComfyUIMetadataExtractor from .logging_config import get_logger from .prompt_graph import resolve_text, run_prompt_nodes from .prompt_pair_resolver import resolve_prompt_pair +from .run_keys import image_run_keys 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 except ImportError: from database.pair_ingestion import ingest_resolved_pair + from database.run_attribution import host_output_directory, stored_run_prompt # Upper bound on images waiting to be processed (a burst beyond this is skipped) MAX_PENDING_IMAGES = 1000 @@ -126,15 +130,17 @@ class ImageGenerationHandler(FileSystemEventHandler): if event.is_directory or not self.is_image_file(event.src_path): return self.logger.info(f"New image detected: {event.src_path}") - # Snapshot prompt context NOW before the delay — in batch workflows - # the tracker advances to the next prompt before images are processed. - prompt_snapshot = self.prompt_tracker.get_current_prompt() - if prompt_snapshot: - self.logger.debug( - f"Snapshot prompt {prompt_snapshot.get('id', '?')} for " - f"{os.path.basename(event.src_path)}" - ) - self.schedule(event.src_path, prompt_snapshot) + # Snapshot the tracked runs NOW, before the delay: a later run may + # execute the same node again before this image is processed. + self.schedule(event.src_path, self._keyed_snapshot()) + + def _keyed_snapshot(self): + """The tracker's latest entry per run key, or None when unavailable.""" + try: + return self.prompt_tracker.keyed_snapshot() + except Exception as error: + self.logger.warning(f"Could not snapshot tracked prompts: {error}") + return None def schedule(self, image_path: str, prompt_snapshot=None) -> bool: """Queue an image for processing on the worker thread, once per path. @@ -290,13 +296,15 @@ class ImageGenerationHandler(FileSystemEventHandler): This method handles the complete processing pipeline for a new image: 1. Verifies the file still exists 2. Extracts ComfyUI metadata from the image - 3. Tries to identify the correct prompt via metadata-based lookup - 4. Falls back to the prompt snapshot captured at event time - 5. Links the image to the appropriate prompt in the database + 3. Identifies the prompt from the image's metadata, else from a tracked + run the image's graph proves (``_find_prompt_from_run``) + 4. Links the image to that prompt, or leaves it unlinked when nothing + certain identifies it (logged at info; #342) Args: image_path: Full path to the newly created image file - prompt_snapshot: Prompt context captured at file-creation time (optional) + prompt_snapshot: The tracker's keyed snapshot taken when the image + appeared (``PromptTracker.keyed_snapshot``); None reads it now """ try: self.logger.info(f"Processing image: {image_path}") @@ -329,43 +337,17 @@ class ImageGenerationHandler(FileSystemEventHandler): f"Metadata match: prompt {current_prompt['id']} for " f"{os.path.basename(image_path)}" ) - - # Strategy 2: Pop from batch queue. Batch items (text supplied by - # PromptSearchList) share one graph, so metadata can't tell them apart; - # they are queued in encoding order and images save in the same order. - if not current_prompt: - current_prompt = self.prompt_tracker.pop_next_prompt() - if current_prompt: - self.logger.info( - f"Queue match: prompt {current_prompt['id']} for " - f"{os.path.basename(image_path)}" - ) - - # Strategy 3: Use snapshot captured at file-creation time - if not current_prompt and prompt_snapshot: - current_prompt = prompt_snapshot - self.logger.info( - f"Snapshot match: prompt {current_prompt.get('id', 'unknown')}" + else: + current_prompt = self._find_prompt_from_run( + image_path, metadata, prompt_snapshot ) - - # Strategy 4: Check live prompt tracker if not current_prompt: - current_prompt = self.prompt_tracker.get_current_prompt() - if current_prompt: - self.logger.info( - f"Live tracker match: prompt {current_prompt['id']}" - ) - - # Strategy 5: Fallback to most recent prompt in DB - if not current_prompt: - current_prompt = self._get_fallback_prompt() - if current_prompt: - self.logger.debug(f"Fallback match: prompt {current_prompt['id']}") - else: - self.logger.warning( - f"No prompt context available, skipping image: {image_path}" - ) - return + # Nothing certain: never guess from recency or a stale entry (#342) + self.logger.info( + "No certain prompt for " + f"{os.path.basename(image_path)}; left unlinked" + ) + return # Extend timeout if we have an active execution if current_prompt.get("execution_id"): @@ -393,6 +375,61 @@ class ImageGenerationHandler(FileSystemEventHandler): self.logger.error(traceback.format_exc()) + def _find_prompt_from_run( + self, image_path: str, metadata, prompt_snapshot=None + ) -> Optional[Dict[str, Any]]: + """The prompt a tracked run proves for an image whose metadata can't say. + + Only an entry registered under the run key of the image's own positive + node counts (``utils.run_keys``): never the newest prompt in the + database and never an entry of another graph (#342). + + 1. Batch queue: the oldest queued entry of that key. Batch items (text + from PromptSearchList) share one key and are linked in order. An + empty item's placeholder leaves its image unlinked. + 2. Snapshot taken when the image appeared: the key's latest entry, when + its run registered exactly one entry for the key (a batch's latest + item says nothing about the other images) and the run is known. + 3. Recorded execution history: the run that wrote this file names one + stored prompt (``database.run_attribution``). + """ + keys = image_run_keys(metadata) + name = os.path.basename(image_path) + if len(keys) == 1: + (key,) = keys.values() + queued = self.prompt_tracker.pop_for_keys([key]) + if queued is not None: + if queued.get("placeholder"): + self.logger.info(f"Queue entry for {name} has no prompt text") + return None + self.logger.info(f"Queue match: prompt {queued['id']} for {name}") + return queued + if prompt_snapshot is None: + prompt_snapshot = self._keyed_snapshot() + entry = (prompt_snapshot or {}).get(key) + if entry is not None: + context, run_id, count = entry + if not context.get("placeholder") and run_id and count == 1: + self.logger.info(f"Run match: prompt {context['id']} for {name}") + return context + elif len(keys) > 1: + self.logger.debug(f"{len(keys)} positive prompt nodes for {name}") + prompt_id = self._history_prompt(image_path) + if prompt_id is None: + return None + self.logger.info(f"History match: prompt {prompt_id} for {name}") + return {"id": prompt_id, "from_history": True} + + def _history_prompt(self, image_path: str) -> Optional[int]: + """The stored prompt recorded history attributes to this file, if one.""" + try: + return stored_run_prompt( + self.db_manager, image_path, host_output_directory() + ) + except Exception as error: + self.logger.warning(f"History attribution failed: {error}") + return None + def _find_prompt_from_metadata(self, metadata) -> Tuple[str, Optional[Dict]]: """Identify the prompt that produced an image from its embedded metadata. @@ -542,35 +579,6 @@ class ImageGenerationHandler(FileSystemEventHandler): self.logger.error(f"Error getting file info: {e}") return {} - def _get_fallback_prompt(self) -> Optional[Dict[str, Any]]: - """Get the most recent prompt from database as fallback. - - When no active prompt is available from the tracker, this method attempts - to find the most recently created prompt in the database to use as a fallback - for image linking. - - Returns: - Dictionary containing prompt information with 'fallback' flag set to True, - or None if no recent prompt is available - """ - try: - recent = self.db_manager.get_recent_prompts(limit=1) - # Paginated dict ({"prompts": [...]}) or a plain list - recent_prompts = ( - recent.get("prompts") if isinstance(recent, dict) else recent - ) - if recent_prompts: - prompt = recent_prompts[0] - return { - "id": prompt["id"], - "text": prompt["text"], - "timestamp": prompt.get("created_at"), - "fallback": True, - } - except Exception as e: - self.logger.error(f"Error getting fallback prompt: {e}") - return None - def link_image_to_prompt( self, image_path: str, prompt_context: Dict, metadata: Dict ): @@ -615,10 +623,9 @@ class ImageGenerationHandler(FileSystemEventHandler): if image_id is None: return if image_id: # 0: nothing new was linked (logged by the database) - fallback_note = " (fallback)" if prompt_context.get("fallback") else "" self.logger.debug( f"Successfully linked image {image_id} to prompt " - f"{prompt_context['id']}{fallback_note}" + f"{prompt_context['id']}" ) except Exception as e: self.logger.error(f"Failed to link image to prompt: {e}") diff --git a/utils/prompt_tracker.py b/utils/prompt_tracker.py index bc3ab35e..82c75412 100644 --- a/utils/prompt_tracker.py +++ b/utils/prompt_tracker.py @@ -28,7 +28,8 @@ Or using the context manager: import threading import time import uuid -from typing import Optional, Dict, Any +from collections import OrderedDict +from typing import Any, Dict, Iterable, Optional, Tuple # Import logging system try: @@ -103,6 +104,8 @@ class PromptTracker: prompt_text: str, additional_data: Optional[Dict[str, Any]] = None, push_to_queue: bool = True, + run_key: Optional[str] = None, + run_id: Optional[str] = None, ) -> str: """ Set the current prompt for this thread for image tracking. @@ -120,6 +123,10 @@ class PromptTracker: push_to_queue: Also push onto the FIFO batch queue. Only batch items whose text comes from another node need it; typed prompts are linked from the image's own metadata instead. + run_key: The node's run key (``utils.run_keys.prompt_node_key``). The + image monitor links an image to a tracked prompt only when the + image's own graph gives the same key (#342). + run_id: The ComfyUI prompt id executing the node, when the host says. Returns: Unique execution ID for this prompt execution @@ -157,6 +164,8 @@ class PromptTracker: "timestamp": time.time(), "thread_id": threading.current_thread().ident, "additional_data": additional_data or {}, + "run_key": run_key, + "run_id": run_id, } # Store in thread-local storage @@ -166,9 +175,12 @@ class PromptTracker: with self.lock: self.active_prompts[execution_id] = execution_context + self._remember_keyed(execution_context) + # Push to batch queue for ordered image linking if push_to_queue: with self._queue_lock: + self._supersede_queued(run_key, run_id) self._prompt_queue.append(execution_context) self.logger.debug( f"Queue size: {len(self._prompt_queue)} after push " @@ -225,6 +237,95 @@ class PromptTracker: ) return None + def mark_empty( + self, + run_key: Optional[str], + run_id: Optional[str] = None, + push_to_queue: bool = False, + ) -> None: + """Record that a positive node ran with nothing to save (#342). + + The run key's latest entry becomes a placeholder, so no earlier prompt of + that key can be linked to this run's images; with ``push_to_queue`` the + placeholder also holds this item's place in the batch queue, so a later + item's image never takes the wrong queue entry. + """ + if not run_key: + return + placeholder = { + "id": None, + "placeholder": True, + "execution_id": self.generate_execution_id(), + "text": "", + "timestamp": time.time(), + "additional_data": {}, + "run_key": run_key, + "run_id": run_id, + } + self._remember_keyed(placeholder) + if push_to_queue: + with self._queue_lock: + self._supersede_queued(run_key, run_id) + self._prompt_queue.append(placeholder) + + def _keyed_runs(self) -> "KeyedRuns": + keyed = getattr(self, "_keyed", None) + if keyed is None: + with _KEYED_INIT_LOCK: + keyed = getattr(self, "_keyed", None) + if keyed is None: + keyed = self._keyed = KeyedRuns() + return keyed + + def _remember_keyed(self, context: Dict[str, Any]) -> None: + if context.get("run_key"): + self._keyed_runs().record(context) + + def _supersede_queued(self, run_key: Optional[str], run_id: Optional[str]): + """Drop queued entries of ``run_key`` from an earlier run. + + A node that executes again with the same key in a new run replaces what + the earlier run queued: images of the new run must not take them. + Called with ``_queue_lock`` held. + """ + if not run_key or run_id is None: + return + self._prompt_queue = [ + ctx + for ctx in self._prompt_queue + if ctx.get("run_key") != run_key or ctx.get("run_id") in (None, run_id) + ] + + def keyed_snapshot(self) -> Dict[str, Tuple[Dict[str, Any], Optional[str], int]]: + """``{run key: (latest entry, its run id, entries of that run)}`` now. + + The image monitor captures it when an image appears, before later runs + can move on. + """ + return self._keyed_runs().snapshot(self.prompt_timeout) + + def pop_for_keys(self, keys: Iterable[str]) -> Optional[Dict[str, Any]]: + """Pop the oldest queued entry whose run key is one of ``keys`` (FIFO). + + Entries of other keys stay queued; expired entries are discarded. + """ + wanted = {key for key in keys if key} + if not wanted: + return None + now = time.time() + with self._queue_lock: + kept = [] + found = None + for ctx in self._prompt_queue: + if now - ctx.get("timestamp", 0) >= self.prompt_timeout: + continue + if found is None and ctx.get("run_key") in wanted: + found = ctx + continue + kept.append(ctx) + self._prompt_queue = kept + return found + def pop_next_prompt(self) -> Optional[Dict[str, Any]]: """Pop the next prompt from the batch queue (FIFO). @@ -463,6 +564,46 @@ class PromptTracker: } +_KEYED_INIT_LOCK = threading.Lock() +# Run keys remembered for the image monitor (oldest forgotten first) +MAX_KEYED_RUNS = 256 + + +class KeyedRuns: + """The latest tracked entry per run key, and how many its run produced. + + More than one entry for a key in one run (a batch fed by a list) means the + latest one alone cannot say which item an image belongs to. + """ + + def __init__(self, max_keys: int = MAX_KEYED_RUNS): + self._max = max_keys + self._latest: "OrderedDict[str, Tuple[Dict[str, Any], Optional[str], int]]" + self._latest = OrderedDict() + self._lock = threading.Lock() + + def record(self, context: Dict[str, Any]) -> None: + key, run_id = context["run_key"], context.get("run_id") + with self._lock: + previous = self._latest.pop(key, None) + same_run = ( + previous is not None and run_id is not None and previous[1] == run_id + ) + count = previous[2] + 1 if same_run else 1 + self._latest[key] = (context, run_id, count) + while len(self._latest) > self._max: + self._latest.popitem(last=False) + + def snapshot(self, timeout: float) -> Dict[str, Tuple[Dict[str, Any], Any, int]]: + now = time.time() + with self._lock: + return { + key: entry + for key, entry in self._latest.items() + if now - entry[0].get("timestamp", 0) < timeout + } + + class PromptExecutionContext: """Context manager for prompt executions. diff --git a/utils/run_keys.py b/utils/run_keys.py new file mode 100644 index 00000000..fda46a87 --- /dev/null +++ b/utils/run_keys.py @@ -0,0 +1,115 @@ +"""Prove that a tracked prompt belongs to the run that saved an image (#342). + +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. + +Pure functions; the only host dependency (the executing run id) is optional. +""" + +import hashlib +import json +from typing import Any, Dict, List, Optional + +try: + from .prompt_graph import MAX_VISITS, run_prompt_nodes +except ImportError: + from utils.prompt_graph import MAX_VISITS, run_prompt_nodes + +RUN_KEY_VERSION = 1 + + +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 + + +def prompt_node_key(graph: Any, node_id: Any) -> Optional[str]: + """Digest of ``node_id`` and its whole upstream closure, or None. + + Only ``class_type`` and ``inputs`` count (titles and other UI fields do not + change what runs). None when the graph or node is missing, an upstream link + points at a missing node, the closure exceeds ``MAX_VISITS`` nodes, or a + value cannot be serialized: no key means no proof. + """ + if not isinstance(graph, dict) or node_id is None: + return None + start = str(node_id) + closure: Dict[str, Any] = {} + stack = [start] + while stack: + current = stack.pop() + if current in closure: + continue + node = graph.get(current) + if not isinstance(node, dict): + return None + if len(closure) >= MAX_VISITS: + return None + inputs = node.get("inputs") + closure[current] = [node.get("class_type"), inputs] + stack.extend(_links(inputs)) + try: + encoded = json.dumps( + [RUN_KEY_VERSION, start, closure], + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ) + except (TypeError, ValueError): + return None + return hashlib.sha256(encoded.encode("ascii")).hexdigest() + + +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. + """ + graph = metadata.get("prompt") if isinstance(metadata, dict) else None + if isinstance(graph, str): + try: + graph = json.loads(graph) + except (TypeError, ValueError, RecursionError): + return {} + if not isinstance(graph, dict): + return {} + 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 + return keys + + +def executing_run_id() -> Optional[str]: + """The ComfyUI prompt id currently executing on this thread, if the host says. + + None outside ComfyUI or on hosts without ``get_executing_context``. + """ + try: + from comfy_execution.utils import get_executing_context + except Exception: + return None + try: + context = get_executing_context() + except Exception: + return None + run_id = getattr(context, "prompt_id", None) + return run_id if isinstance(run_id, str) and run_id else None