fix(ingestion): link an image only to a prompt that provably produced it

An LLM fed a PromptManagerPair's prepend_text while its text editor was
empty. The node refused to save (an empty positive raised and logged a
warning), and the image monitor then linked the image to a stale tracker
entry or to the newest prompt in the database.

- Nodes: an empty text editor saves and tracks nothing (prepend and
  append text are never the stored prompt), logs at info, and marks the
  node's run key empty so no earlier entry can be used for its images.
- Tracker entries carry a run key: a digest of the prompt node and its
  whole upstream graph, the identity ComfyUI's cache uses. The monitor
  computes the same key from the image's embedded graph and uses a
  queued or snapshot entry only when the keys match (a snapshot also
  needs a known run with one entry, so a batch's last item is never
  applied to every image).
- The recency fallback is gone; anything uncertain stays unlinked and is
  logged at info.
- Recorded execution history attributes a file when exactly one run
  reported it, the file was written during that run, and the run's
  positive contexts name one stored prompt. The monitor and the output
  scan use it as their last certain source.
- The scan and pair ingestion treat a resolved pair with an empty stored
  positive as unresolved instead of raising.
- tools/review_image_links.py reviews selected links by these rules;
  dry run by default, --apply changes only named ids after a backup.

Closes #342
This commit is contained in:
Vito Sansevero
2026-10-25 12:58:16 -07:00
parent 863d717435
commit 144e6b1145
20 changed files with 2023 additions and 183 deletions
+16
View File
@@ -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:
+122
View File
@@ -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
+2
View File
@@ -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(
+37
View File
@@ -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:
+115 -77
View File
@@ -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,
+2
View File
@@ -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)}"
+37 -2
View File
@@ -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."""
+89
View File
@@ -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);
});
+40 -5
View File
@@ -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] = "<several links>"
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)
)
)
)
+17 -13
View File
@@ -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")
+9 -2
View File
@@ -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)]
+485
View File
@@ -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()
+175
View File
@@ -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()
+237
View File
@@ -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()
+10 -1
View File
@@ -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):
+284
View File
@@ -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
(``<db>.pre-relink_<stamp>``). Only ids named on the command line are ever
changed. A batch item linked through the monitor's queue (PromptSearchList)
can't be proven after the fact and reads ``unlink``: review before applying.
python3 tools/review_image_links.py <db> --output-dir <ComfyUI output> \\
--ids 58005-58010
python3 tools/review_image_links.py <db> --output-dir <dir> --prompt 5278
python3 tools/review_image_links.py <db> --output-dir <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())
+1 -1
View File
@@ -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,
+88 -81
View File
@@ -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}")
+142 -1
View File
@@ -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.
+115
View File
@@ -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