Compare commits
10
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9ab4b95fb5 | ||
|
|
eb309c5737 | ||
|
|
42c96c9ff5 | ||
|
|
9a6a8a5312 | ||
|
|
16e88342ad | ||
|
|
417a05a2db | ||
|
|
4ab0119580 | ||
|
|
b05b5dbfd9 | ||
|
|
628e000765 | ||
|
|
144e6b1145 |
+107
-6
@@ -802,8 +802,8 @@ Recorded shell deviations:
|
||||
`file_size` equal to the file's size in this walk. Both paths go through
|
||||
`normalize_image_path` (absolute, normalized, `os.path.normcase`); a row without
|
||||
`file_path` matches by its `image_path`.
|
||||
- Read again: files linked without a pair (the monitor's fallback and FIFO links, the
|
||||
legacy positive-only import), linked to more than one prompt or pair, not linked, a
|
||||
- Read again: files linked without a pair (the monitor's queue and run-key links,
|
||||
fallback links made before #342, the legacy positive-only import), linked to more than one prompt or pair, not linked, a
|
||||
file whose size changed (one replaced at the same path, e.g. a reused ComfyUI
|
||||
counter name, is read and its conflict lands under Needs review), and a row with no
|
||||
stored size (unknown is not known). The scan now stores each file's size when it
|
||||
@@ -3352,14 +3352,14 @@ generation did not change meanwhile (`PromptModel.get_connection`).
|
||||
Writes racing the switch: a background write that was waiting on the old file's lock
|
||||
when the switch closed its connection fails, and run counting (`count_runs`,
|
||||
`count_pair_runs`) and the image monitor's link (`ImageGenerationHandler.
|
||||
link_image_to_prompt`: the resolved-pair ingest and the fallback or recent-prompt link)
|
||||
link_image_to_prompt`: the resolved-pair ingest and the plain prompt link)
|
||||
run it once more through `PromptDatabase.survive_library_switch` when the model's
|
||||
generation changed during the attempt; the retry opens the new library. The monitor
|
||||
links through `PromptDatabase.link_image_to_prompt_or_raise`, which raises database
|
||||
errors (`link_image_to_prompt` keeps its log-and-return-0 contract for other callers);
|
||||
each attempt is one transaction, so a failed attempt rolls back and a committed one
|
||||
does not retry, and the image is linked once. `TestBackgroundWritesSurviveTheSwitch`
|
||||
races a real UPDATE (run counting) and a real fallback image link against the lock and
|
||||
races a real UPDATE (run counting) and a real plain image link against the lock and
|
||||
asserts each lands in the new file only. Any other failure is not retried. Known limit: that
|
||||
blocked writer, and the mover with it, waits out the writer's busy timeout (30 s)
|
||||
before its connection closes: SQLite holds the connection's mutex during the busy
|
||||
@@ -3958,6 +3958,106 @@ Viewer, including nested modal focus, mobile layouts and per-role sending.
|
||||
real Classic and Nodes 2.0 widgets and verifies workflow serialization without
|
||||
loading a workflow or queueing generation.
|
||||
|
||||
## Image attribution: certain or unlinked (#342)
|
||||
|
||||
The image monitor and the output scan link an output only to a prompt that
|
||||
provably produced it; anything uncertain stays unlinked and the monitor logs it
|
||||
at info ("No certain prompt for <file>; left unlinked"). There is no recency
|
||||
fallback (`_get_fallback_prompt` is gone) and no unproven tracker entry.
|
||||
|
||||
- **Empty editor.** Prepend and append text are never stored as the prompt
|
||||
(maintainer's product decision; #343 tracks letting a link drive `text`). A
|
||||
PromptManager, Text, Focused, Split or Pair node whose `text` is empty saves,
|
||||
counts and tracks nothing for the positive side, logs at info (no "Could not
|
||||
save or track" warning), and marks its run key empty, so no earlier entry of
|
||||
that key can be used and a batch item keeps its queue place
|
||||
(`PromptTracker.mark_empty`). The queue hook counts nothing for it either. A
|
||||
save or tracking failure of a positive node marks its run key the same way, so
|
||||
an older entry of that key never stands in for the failed run.
|
||||
- **Run keys** (`utils/run_keys.py`): sha256 of a prompt node plus its whole
|
||||
upstream closure (`class_type` and `inputs`) in the API graph. Nodes register
|
||||
tracker entries under it; the monitor computes it from the image's embedded
|
||||
graph. The snapshot taken when the file appears names the key's latest run; a
|
||||
queued entry is taken only by key and from that run (batch items in order, a
|
||||
placeholder leaves its image unlinked, same-key entries of earlier runs ahead
|
||||
of it are stale and dropped, entries of later runs that executed during the
|
||||
monitor's delay wait for their own images). Without a queued entry the
|
||||
snapshot's entry itself is used only when its run id is known
|
||||
(`comfy_execution.utils.get_executing_context`) and that run registered one
|
||||
entry for the key, which covers a cached re-run (sampler seed changed) without
|
||||
applying a batch's last item to every image. When the key's latest entry is a
|
||||
placeholder the image stays unlinked without consulting history (the same
|
||||
rule as a popped placeholder). More than one positive PromptManager node in
|
||||
the image graph, or any positive node without a key: never linked from the
|
||||
tracker. Links are followed exactly as ComfyUI's `is_link` defines them (a
|
||||
`str` node id and an `int` or `float` slot in a two-item list). Equal keys do
|
||||
not always mean equal text: the key can't see `IS_CHANGED`, so an upstream
|
||||
node that forces re-runs gives two runs one key; that is why the run named by
|
||||
the snapshot is required as well. A file event delivered after the next run
|
||||
already re-registered the key is the one remaining exposure (see the ledger).
|
||||
- **Recorded history** (`database/run_attribution.py`): a file is attributed when
|
||||
exactly one recorded run reported its output location (root key, subfolder,
|
||||
name), its mtime lies within that run (dequeue to terminal or callback, +-2 s),
|
||||
and every positive context names the same stored prompt (the live
|
||||
`positive_prompt_id` hint). The monitor tries it last; the scan tries it when the
|
||||
#170 browser-parser floor finds nothing, and also for files with no prompt
|
||||
metadata at all (`--disable-metadata`, savers that embed none); it never
|
||||
creates a prompt. Lookups are batched (`stored_run_prompts`): one
|
||||
parameterized `filename IN (...)` query per chunk of at most 500 names
|
||||
(`LOOKUP_CHUNK`), matched in memory, so a scan batch asks history once per
|
||||
chunk, never once per file. The history schema is unchanged (no extra index),
|
||||
so older builds still open the library.
|
||||
- **After the run ends.** An image with no run keys (no embedded API graph) that
|
||||
nothing identifies waits for its run: when the run's outputs are recorded
|
||||
(`write_execution_outputs_for_finish` notifies `subscribe_outputs_sealed`
|
||||
listeners after commit), each waiting image the run reported gets exactly one
|
||||
history-only re-check on the monitor's worker. At most 1000 images wait, each
|
||||
for at most an hour; shutdown forgets them. The image is registered before the
|
||||
first history check, so a run sealed in between is not missed; a re-check that
|
||||
can't be scheduled yet (the image's own pass still pending, or the pending set
|
||||
full) is kept due and scheduled when a pending image finishes. A keyless image
|
||||
outside ComfyUI's output folder never waits (history can't attribute it).
|
||||
- **Spellings.** History keeps ComfyUI's own spelling of its output folder (the
|
||||
root key). A file is matched through that spelling first, then with both sides
|
||||
resolved (`realpath`), so an output folder reached through a symlink or a
|
||||
Windows junction still attributes the scan's resolved file paths; waiting
|
||||
files are keyed by their resolved path for the same reason. The scan makes
|
||||
its batch's history lookup only when the first file in the batch needs it.
|
||||
- **Blank positive.** A resolved pair whose stored positive is empty (the #342
|
||||
image with its Pair execution record) is `unresolved` for the scan and for
|
||||
`ingest_resolved_pair` (it used to raise in `_save_pair`).
|
||||
- **Parity floor.** Metadata linking, the scan's resolver and parser floor are
|
||||
unchanged for every non-blank result; only guesses are gone.
|
||||
|
||||
Tests: `tests/test_prompt_attribution.py` (nodes, tracker, monitor, run keys),
|
||||
`tests/test_run_attribution.py` (history attribution, monitor and scan),
|
||||
`tests/test_review_image_links_tool.py`, and the real-host
|
||||
`tests/comfyui/image-attribution.spec.mjs` (both renderers): a typed Pair image
|
||||
links, a Pair image whose positive is only linked prepend text from an unknown
|
||||
node stays unlinked with nothing stored, an unknown node feeding `text` links
|
||||
through its run key, and its cached re-run with a new sampler seed links again
|
||||
(the host's `execution_cached` event names the PromptManager node, under the
|
||||
fixture's Classic cache). A second case saves with metadata disabled and holds
|
||||
the run open behind the fixture gate after SaveImage: the output is waiting and
|
||||
unlinked while the run is open, and links to the stored prompt once the run's
|
||||
outputs are recorded. The managed host config disables monitoring, so these
|
||||
cases start and stop the plugin's monitor through the owned fixture route
|
||||
`/pm_execution_fixture/monitor` (its GET lists the waiting files).
|
||||
On origin/main the prepend image was linked to an unrelated prompt.
|
||||
|
||||
Repair: `tools/review_image_links.py <db> --output-dir <ComfyUI output> --ids
|
||||
<ids>` (or `--prompt <id>`) judges links by these rules (`keep`, `relink`,
|
||||
`unlink`, `conflict`, `missing`) on a backup copy and changes nothing.
|
||||
`--apply` writes `relink`/`unlink` only for ids named with `--ids`, after saving
|
||||
`<db>.pre-relink_<stamp>`, and only while each image is still linked to the
|
||||
prompt it was judged on (otherwise it is counted as skipped). `relink` points the
|
||||
link record at the other prompt and drops its pair links. `unlink` removes only
|
||||
the image's link record (its `generated_images` row, the metadata snapshot stored
|
||||
in it and its pair links): that is how the app represents an unlinked image, the
|
||||
same as a file the scan could not attribute, so the file stays on disk and in the
|
||||
gallery as unlinked, and its favorite and tags (kept by file) are untouched. Every
|
||||
run prints this. Queue-linked batch images can't be proven after the fact and read
|
||||
`unlink`, so review before applying. Stop ComfyUI first.
|
||||
## Linked prompt text (#343)
|
||||
|
||||
A link drives `text` on every PromptManager node, and `negative_text` on the Pair,
|
||||
@@ -4040,8 +4140,9 @@ PM_COMFY_ROOT=... PM_TEST_PYTHON=... npm run --prefix tools/vendor test:executio
|
||||
|
||||
## Durable execution ledger
|
||||
|
||||
`npm run --prefix tools/vendor test:executions` runs 46 real execution cases
|
||||
(23 each in Classic and Nodes 2.0, retaining all 42 prior execution cases),
|
||||
`npm run --prefix tools/vendor test:executions` runs 52 real execution cases
|
||||
(26 each in Classic and Nodes 2.0, retaining all 48 prior execution cases and adding
|
||||
the two #342 image-attribution cases),
|
||||
alongside the existing 80 stock host cases (the original 78 plus the native
|
||||
editor host case) and six pack compatibility cases. Run
|
||||
these suites serially because they share report cleanup and owned host resources. Use the same pinned disposable `PM_COMFY_ROOT`,
|
||||
|
||||
@@ -4,6 +4,8 @@ import sqlite3
|
||||
from contextlib import ExitStack
|
||||
from typing import Callable, ContextManager
|
||||
|
||||
from .run_attribution import notify_outputs_sealed
|
||||
|
||||
try:
|
||||
from ..utils.execution_history_events import RunFinish, encode_command
|
||||
from ..utils.execution_outputs import OutputBundle, validate_output_bundle
|
||||
@@ -113,6 +115,10 @@ class ExecutionOutputOperations:
|
||||
"VALUES(?,?,?,?,?,?,?,?)",
|
||||
[(run_id, *reference) for reference in prepared],
|
||||
)
|
||||
# Committed: images waiting for this run may now be attributed (#342)
|
||||
notify_outputs_sealed(
|
||||
run_id, [dict(zip(_REFERENCE_FIELDS, row)) for row in prepared]
|
||||
)
|
||||
return True
|
||||
|
||||
def get_execution_outputs(self, run_id: str) -> dict:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -0,0 +1,244 @@
|
||||
"""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.
|
||||
|
||||
Lookups are batched (``stored_run_prompts``): one parameterized query per
|
||||
chunk of at most ``LOOKUP_CHUNK`` distinct file names, matched in memory. The
|
||||
history schema is left exactly as it is (no extra index), so older builds can
|
||||
still open the library.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
import threading
|
||||
import weakref
|
||||
from typing import Any, Callable, Dict, Iterable, List, 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
|
||||
|
||||
# Distinct file names per lookup query: well under SQLite's bound-parameter
|
||||
# limit (999 on old builds), so a scan batch costs one query per chunk.
|
||||
LOOKUP_CHUNK = 500
|
||||
|
||||
_RUNS_FOR_NAMES = (
|
||||
"SELECT DISTINCT o.filename, o.subfolder, o.root_key, o.path_style, 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.filename IN ({marks})"
|
||||
)
|
||||
|
||||
|
||||
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 _inside(image_path: str, output_root: str) -> Optional[str]:
|
||||
"""``image_path`` relative to ``output_root``, or None when outside it."""
|
||||
try:
|
||||
relative = os.path.relpath(image_path, 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
|
||||
return relative
|
||||
|
||||
|
||||
def _resolved(path: str) -> str:
|
||||
try:
|
||||
return os.path.realpath(path)
|
||||
except (OSError, ValueError):
|
||||
return os.path.abspath(path)
|
||||
|
||||
|
||||
def _location(image_path: str, output_root: str, path_style: str):
|
||||
"""(root key, subfolder, filename) as history records them, or None.
|
||||
|
||||
The root key is always the host's own spelling of its output folder,
|
||||
because history recorded that one. The file may be spelled through the
|
||||
same folder or through its resolved path (a symlinked or junction output
|
||||
folder, as scan roots are): the lexical spelling is tried first, then
|
||||
both sides resolved.
|
||||
"""
|
||||
root_key = snapshot_output_roots(path_style, output_root, output_root).output_key
|
||||
if root_key is None:
|
||||
return None
|
||||
relative = _inside(os.path.abspath(image_path), os.path.abspath(output_root))
|
||||
if relative is None:
|
||||
relative = _inside(_resolved(image_path), _resolved(output_root))
|
||||
if relative is None:
|
||||
return None
|
||||
subfolder, filename = os.path.split(relative)
|
||||
if not filename:
|
||||
return None
|
||||
return root_key, subfolder, filename
|
||||
|
||||
|
||||
def in_output_folder(image_path: str, output_root: Optional[str]) -> bool:
|
||||
"""Whether history could ever attribute ``image_path`` (inside the host's
|
||||
output folder, in either spelling)."""
|
||||
if not image_path or not output_root:
|
||||
return False
|
||||
style = "windows" if os.name == "nt" else "posix"
|
||||
return _location(image_path, output_root, style) is not None
|
||||
|
||||
|
||||
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:
|
||||
return None
|
||||
return stored_run_prompts(db, [image_path], output_root, path_style).get(image_path)
|
||||
|
||||
|
||||
def _runs_by_location(db: Any, names: List[str]) -> Optional[Dict[tuple, dict]]:
|
||||
"""``{(filename, subfolder, root_key, style): {run_id: (start, end)}}``.
|
||||
|
||||
One query per ``LOOKUP_CHUNK`` names; None without history tables.
|
||||
"""
|
||||
found: Dict[tuple, dict] = {}
|
||||
conn = db.model.get_connection()
|
||||
for start in range(0, len(names), LOOKUP_CHUNK):
|
||||
chunk = names[start : start + LOOKUP_CHUNK]
|
||||
query = _RUNS_FOR_NAMES.format(marks=",".join("?" for _ in chunk))
|
||||
try:
|
||||
rows = conn.execute(query, chunk).fetchall() # nosec B608
|
||||
except sqlite3.Error:
|
||||
return None # no history tables (a library without history)
|
||||
for name, subfolder, root_key, style, run_id, began, ended in rows:
|
||||
runs = found.setdefault((name, subfolder, root_key, style), {})
|
||||
runs[run_id] = (began, ended)
|
||||
return found
|
||||
|
||||
|
||||
def _named_prompt(db: Any, run_id: str) -> Optional[int]:
|
||||
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
|
||||
|
||||
|
||||
def stored_run_prompts(
|
||||
db: Any,
|
||||
image_paths: Iterable[str],
|
||||
output_root: Optional[str],
|
||||
path_style: Optional[str] = None,
|
||||
) -> Dict[str, Optional[int]]:
|
||||
"""``{path: stored prompt id or None}`` for many files in few queries.
|
||||
|
||||
The output lookup is one query per ``LOOKUP_CHUNK`` distinct names; each
|
||||
run that alone reported a file is then read once for its contexts.
|
||||
"""
|
||||
paths = list(dict.fromkeys(path for path in image_paths if path))
|
||||
result: Dict[str, Optional[int]] = dict.fromkeys(paths)
|
||||
if not paths or not output_root:
|
||||
return result
|
||||
style = path_style or ("windows" if os.name == "nt" else "posix")
|
||||
located = {}
|
||||
for path in paths:
|
||||
location = _location(path, output_root, style)
|
||||
if location is not None:
|
||||
located[path] = location
|
||||
if not located:
|
||||
return result
|
||||
names = sorted({filename for _key, _sub, filename in located.values()})
|
||||
runs = _runs_by_location(db, names)
|
||||
if not runs:
|
||||
return result
|
||||
prompts: Dict[str, Optional[int]] = {}
|
||||
for path, (root_key, subfolder, filename) in located.items():
|
||||
reported = runs.get((filename, subfolder, root_key, style), {})
|
||||
if len(reported) != 1:
|
||||
continue
|
||||
((run_id, (began, ended)),) = reported.items()
|
||||
if not _written_during(path, began, ended):
|
||||
continue
|
||||
if run_id not in prompts:
|
||||
prompts[run_id] = _named_prompt(db, run_id)
|
||||
result[path] = prompts[run_id]
|
||||
return result
|
||||
|
||||
|
||||
_LOGGER = logging.getLogger("prompt_manager.run_attribution")
|
||||
_SEALED_LISTENERS: List[weakref.WeakMethod] = []
|
||||
_SEALED_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def subscribe_outputs_sealed(method: Callable[[str, List[dict]], None]) -> None:
|
||||
"""Call bound ``method(run_id, references)`` after a run's outputs are recorded.
|
||||
|
||||
Held weakly: a listener that is garbage collected is dropped.
|
||||
"""
|
||||
with _SEALED_LOCK:
|
||||
_SEALED_LISTENERS.append(weakref.WeakMethod(method))
|
||||
|
||||
|
||||
def notify_outputs_sealed(run_id: str, references: Iterable[dict]) -> None:
|
||||
"""Tell listeners a run's output references were just committed.
|
||||
|
||||
A listener failure is logged and never reaches the history writer.
|
||||
"""
|
||||
references = list(references)
|
||||
with _SEALED_LOCK:
|
||||
alive = [ref for ref in _SEALED_LISTENERS if ref() is not None]
|
||||
_SEALED_LISTENERS[:] = alive
|
||||
listeners = [ref() for ref in alive]
|
||||
for listener in listeners:
|
||||
if listener is None:
|
||||
continue
|
||||
try:
|
||||
listener(run_id, references)
|
||||
except Exception:
|
||||
_LOGGER.warning("Output seal listener failed", exc_info=True)
|
||||
@@ -238,6 +238,9 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
|
||||
except Exception as e:
|
||||
# Log error but don't fail the encoding
|
||||
self.logger.warning(f"Failed to save prompt to database: {e}")
|
||||
self._mark_untracked(prompt, unique_id)
|
||||
else:
|
||||
self._skip_empty_prompt(prompt, unique_id)
|
||||
|
||||
# Perform standard CLIP text encoding using the combined text
|
||||
self.logger.debug(
|
||||
|
||||
+56
-1
@@ -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
|
||||
|
||||
|
||||
@@ -168,14 +170,67 @@ class PromptManagerBase:
|
||||
prompt_id = self._save_prompt_to_database(
|
||||
text=text, category=category, tags=tags, count_run=batch_item
|
||||
)
|
||||
if prompt_id and role["is_run_prompt"]:
|
||||
if role["is_run_prompt"] and not prompt_id:
|
||||
# The save failed: an older entry of this key must not stand in (#342)
|
||||
self._mark_untracked(prompt_graph, unique_id)
|
||||
elif role["is_run_prompt"]:
|
||||
self.prompt_tracker.set_current_prompt(
|
||||
prompt_text=encoding_text,
|
||||
additional_data={**additional_data, "prompt_id": prompt_id},
|
||||
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 (``_mark_untracked``).
|
||||
"""
|
||||
self._mark_untracked(
|
||||
prompt_graph,
|
||||
unique_id,
|
||||
"Prompt text is empty (prepend and append text are not stored): "
|
||||
"nothing saved or tracked, and this run's images stay unlinked",
|
||||
)
|
||||
|
||||
def _mark_untracked(
|
||||
self, prompt_graph: Any, unique_id: Any, message: Optional[str] = None
|
||||
) -> None:
|
||||
"""This execution of a positive node has no prompt to link (#342).
|
||||
|
||||
An empty editor, or a save or tracking failure. The node's run key is
|
||||
marked with a placeholder, so no earlier prompt of that key can be
|
||||
linked to this run's images, and a batch item keeps its queue place.
|
||||
``message`` is logged at info for a positive node.
|
||||
"""
|
||||
try:
|
||||
role = self._prompt_role(prompt_graph, unique_id, "")
|
||||
if not role["is_run_prompt"]:
|
||||
if message:
|
||||
self.logger.debug("Empty text on a non-positive node; not saved")
|
||||
return
|
||||
if message:
|
||||
self.logger.info(message)
|
||||
self.prompt_tracker.mark_empty(
|
||||
push_to_queue=role["known"],
|
||||
**self._run_identity(prompt_graph, unique_id),
|
||||
)
|
||||
except Exception as error:
|
||||
# Tracking must never block generation
|
||||
self.logger.warning(f"Could not mark the untracked prompt: {error}")
|
||||
|
||||
def _record_use(self, prompt_id: int, times: int = 1) -> None:
|
||||
"""Count runs; a tracking failure must never block generation."""
|
||||
try:
|
||||
|
||||
+117
-77
@@ -117,85 +117,29 @@ 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}")
|
||||
if is_positive:
|
||||
self._mark_untracked(prompt, unique_id)
|
||||
try:
|
||||
self.comfyui_integration.register_prompt(
|
||||
f"promptmanager_pair_{time.time_ns()}",
|
||||
@@ -223,6 +167,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,
|
||||
|
||||
@@ -216,6 +216,9 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
|
||||
except Exception as e:
|
||||
# Log error but don't fail the processing
|
||||
self.logger.warning(f"Failed to save prompt to database: {e}")
|
||||
self._mark_untracked(prompt, unique_id)
|
||||
else:
|
||||
self._skip_empty_prompt(prompt, unique_id)
|
||||
|
||||
# Register with ComfyUI integration for standard metadata compatibility
|
||||
node_id = f"promptmanagertext_{int(time.time() * 1000)}"
|
||||
|
||||
+102
-4
@@ -35,7 +35,12 @@ 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,
|
||||
stored_run_prompts,
|
||||
)
|
||||
from ...utils.generation_params import (
|
||||
PLACEHOLDERS,
|
||||
GenerationParamsError,
|
||||
@@ -45,7 +50,12 @@ 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,
|
||||
stored_run_prompts,
|
||||
)
|
||||
from utils.generation_params import (
|
||||
PLACEHOLDERS,
|
||||
GenerationParamsError,
|
||||
@@ -548,6 +558,24 @@ def _current_image_monitor():
|
||||
return getattr(_image_monitor_module(), "_monitor_instance", None)
|
||||
|
||||
|
||||
class _BatchHistory:
|
||||
"""A scan batch's recorded-history lookup, made on the first file needing it.
|
||||
|
||||
Most files link through their metadata and never ask; the first one that
|
||||
does triggers one lookup for the whole batch (``_recorded_runs_for``).
|
||||
"""
|
||||
|
||||
def __init__(self, api, media_files):
|
||||
self._api = api
|
||||
self._files = media_files
|
||||
self._found = None
|
||||
|
||||
def get(self, path):
|
||||
if self._found is None:
|
||||
self._found = self._api._recorded_runs_for(self._files)
|
||||
return self._found.get(path)
|
||||
|
||||
|
||||
class AdminRoutesMixin:
|
||||
"""Mixin providing admin, diagnostics, and maintenance API endpoints."""
|
||||
|
||||
@@ -1564,6 +1592,8 @@ class AdminRoutesMixin:
|
||||
if outcomes is not None:
|
||||
outcomes.append(_scan_outcome(media_file, roots, status, title))
|
||||
|
||||
batch_results = list(batch_results)
|
||||
history = _BatchHistory(self, [media for media, _meta in batch_results])
|
||||
for media_file, metadata in batch_results:
|
||||
if self._scan_job and self._scan_job.cancelled:
|
||||
break
|
||||
@@ -1577,6 +1607,7 @@ class AdminRoutesMixin:
|
||||
ambiguous=ambiguous,
|
||||
untraced=untraced,
|
||||
records=records,
|
||||
history=history,
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error processing {media_file.name}: {e}")
|
||||
@@ -1635,6 +1666,7 @@ class AdminRoutesMixin:
|
||||
ambiguous=None,
|
||||
untraced=None,
|
||||
records=None,
|
||||
history=None,
|
||||
):
|
||||
"""Enrich exact pairs; retain legacy positive-only import for absent pairs.
|
||||
|
||||
@@ -1652,7 +1684,9 @@ class AdminRoutesMixin:
|
||||
a dict.
|
||||
"""
|
||||
if not metadata:
|
||||
return None
|
||||
# No prompt metadata (--disable-metadata, bare savers): only the
|
||||
# recorded run can say which prompt made it (#342)
|
||||
return self._link_recorded_run(media_file, {}, detailed, history)
|
||||
canonical = dict(metadata)
|
||||
for key in ("prompt", "workflow"):
|
||||
if isinstance(canonical.get(key), str):
|
||||
@@ -1662,9 +1696,12 @@ class AdminRoutesMixin:
|
||||
# Preserve malformed metadata for conservative resolver rejection.
|
||||
pass
|
||||
if not canonical.get("prompt") and not canonical.get("parameters"):
|
||||
return None
|
||||
return self._link_recorded_run(media_file, canonical, detailed, history)
|
||||
_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 +1730,11 @@ 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, detailed, history
|
||||
)
|
||||
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 +1758,62 @@ class AdminRoutesMixin:
|
||||
return status
|
||||
return {"status": status, "title": self._scan_title(media_file, prompt_text)}
|
||||
|
||||
def _recorded_runs_for(self, media_files):
|
||||
"""``{path: prompt id or None}`` from recorded history for a scan batch.
|
||||
|
||||
One output lookup per chunk of names (``stored_run_prompts``), so a
|
||||
batch never queries history once per file. An empty dict on failure.
|
||||
"""
|
||||
try:
|
||||
return stored_run_prompts(
|
||||
self.db, [str(path) for path in media_files], host_output_directory()
|
||||
)
|
||||
except Exception:
|
||||
self.logger.warning(
|
||||
"History attribution failed for a scan batch", exc_info=True
|
||||
)
|
||||
return {}
|
||||
|
||||
def _link_recorded_run(self, media_file, canonical, detailed=False, history=None):
|
||||
"""Link a file recorded history attributes to one stored prompt (#342).
|
||||
|
||||
Returns the scan status (with ``detailed``, a stored or linked file's
|
||||
``{"status", "title"}``), or None when history names no single stored
|
||||
prompt for the file (``database.run_attribution``). Never creates a
|
||||
prompt. Works without any prompt metadata in the file. ``history`` is
|
||||
the batch's precomputed lookup; without it the file is looked up alone.
|
||||
"""
|
||||
if history is not None:
|
||||
prompt_id = history.get(str(media_file))
|
||||
else:
|
||||
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"
|
||||
status = "linked" if image_id else "unchanged"
|
||||
if not detailed or status not in SCAN_TITLED_DETAILS:
|
||||
return status
|
||||
prompt = self.db.get_prompt_by_id(prompt_id) or {}
|
||||
return {
|
||||
"status": status,
|
||||
"title": self._scan_title(media_file, prompt.get("text") or ""),
|
||||
}
|
||||
|
||||
def _scan_title(self, media_file, text):
|
||||
"""The linked prompt's Library title (its stored title, else ``text``'s
|
||||
first line), as the Library and the inspector name it."""
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Deterministic nodes and queue controls installed only in an owned test host."""
|
||||
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
|
||||
@@ -255,3 +256,50 @@ async def cache_mode(request: web.Request) -> web.Response:
|
||||
raise web.HTTPBadRequest()
|
||||
_CACHE_MODE = mode
|
||||
return web.json_response({"configured": True})
|
||||
|
||||
|
||||
def _plugin_module(suffix: str):
|
||||
"""A loaded PromptManager module (ComfyUI names the package by its path)."""
|
||||
import sys
|
||||
|
||||
for name, module in list(sys.modules.items()):
|
||||
if name.endswith("promptmanager." + suffix):
|
||||
return module
|
||||
raise web.HTTPServiceUnavailable(text=f"PromptManager {suffix} not loaded")
|
||||
|
||||
|
||||
@PromptServer.instance.routes.post("/pm_execution_fixture/monitor")
|
||||
async def image_monitor(request: web.Request) -> web.Response:
|
||||
"""Start or stop PromptManager's image monitor on the host output folder.
|
||||
|
||||
The managed host's config disables monitoring; only the attribution cases
|
||||
(#342) turn it on, and they turn it off again.
|
||||
"""
|
||||
enabled = (await request.json())["enabled"]
|
||||
if type(enabled) is not bool:
|
||||
raise web.HTTPBadRequest()
|
||||
import folder_paths
|
||||
|
||||
monitor_module = _plugin_module("utils.image_monitor")
|
||||
config = _plugin_module("py.config").GalleryConfig
|
||||
monitor = monitor_module._monitor_instance
|
||||
if monitor is None:
|
||||
database = _plugin_module("database.operations").PromptDatabase()
|
||||
tracker = _plugin_module("utils.prompt_tracker").get_prompt_tracker(database)
|
||||
monitor = monitor_module.get_image_monitor(database, tracker)
|
||||
if enabled:
|
||||
config.MONITORING_ENABLED = True
|
||||
monitor.start_monitoring([folder_paths.get_output_directory()])
|
||||
else:
|
||||
monitor.stop_monitoring()
|
||||
config.MONITORING_ENABLED = False
|
||||
return web.json_response({"monitoring": monitor.is_monitoring})
|
||||
|
||||
|
||||
@PromptServer.instance.routes.get("/pm_execution_fixture/monitor")
|
||||
async def image_monitor_state(request: web.Request) -> web.Response:
|
||||
"""File names of keyless images waiting for their run's recorded outputs."""
|
||||
monitor = _plugin_module("utils.image_monitor")._monitor_instance
|
||||
handler = getattr(monitor, "handler", None)
|
||||
waiting = handler.awaiting_history() if handler is not None else []
|
||||
return web.json_response({"awaiting": [os.path.basename(p) for p in waiting]})
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
// #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, history = {}) {
|
||||
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');
|
||||
history.item = item;
|
||||
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);
|
||||
}
|
||||
|
||||
async function monitoring(execution, enabled) {
|
||||
const response = await execution.request.post('/pm_execution_fixture/monitor', { data: { enabled } });
|
||||
expect(response.ok()).toBe(true);
|
||||
expect((await response.json()).monitoring).toBe(enabled);
|
||||
}
|
||||
|
||||
async function cacheMode(execution, mode) {
|
||||
expect((await execution.request.post('/pm_execution_fixture/cache', { data: { mode } })).ok()).toBe(true);
|
||||
}
|
||||
|
||||
test.beforeEach(async ({ execution }) => {
|
||||
expect((await execution.request.post('/pm_execution_fixture/metadata', { data: { disabled: false } })).ok()).toBe(true);
|
||||
await cacheMode(execution, 'CLASSIC'); // keeps the re-run's PromptManager node cached
|
||||
await monitoring(execution, true);
|
||||
});
|
||||
test.afterEach(async ({ execution }) => {
|
||||
await monitoring(execution, false);
|
||||
await cacheMode(execution, 'RAM_PRESSURE');
|
||||
});
|
||||
|
||||
test('an output whose positive is only linked prepend text stays unlinked, never given an earlier prompt', async ({ execution }, testInfo) => {
|
||||
const tag = `${testInfo.project.name} ${randomUUID()}`;
|
||||
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 is served from the cache
|
||||
// without running (the host says so), so no queue entry exists for this run and
|
||||
// only the snapshot's run match can link it.
|
||||
const rerunHistory = {};
|
||||
const rerun = await saved(execution.request, { ...llmGraph, '40': { ...llmGraph['40'],
|
||||
inputs: { ...llmGraph['40'].inputs, seed: 2 } } }, rerunHistory);
|
||||
const cached = rerunHistory.item.status.messages.find(([kind]) => kind === 'execution_cached');
|
||||
expect(cached?.[1].nodes).toEqual(expect.arrayContaining(['11', '20']));
|
||||
await linkedTo(execution, rerun, llmText);
|
||||
});
|
||||
|
||||
test('an output saved without metadata links once its recorded run names a stored prompt', async ({ execution }, testInfo) => {
|
||||
const llmText = `bare output landscape ${testInfo.project.name} ${randomUUID()}`;
|
||||
const request = execution.request;
|
||||
expect((await request.post('/pm_execution_fixture/metadata', { data: { disabled: true } })).ok()).toBe(true);
|
||||
try {
|
||||
// No embedded graph: no run key can prove the tracker entry, so only the
|
||||
// run's recorded history can, and only once the run's outputs are sealed.
|
||||
// The gate after SaveImage holds the run open while the monitor looks.
|
||||
const prompt = graph({ kind: 'PromptManager', text: 'llm', llm: llmText, prefix: 'misslink/bare' });
|
||||
prompt['80'] = { class_type: 'PMHistoryImageOutput', inputs: { images: ['70', 0], mode: 'gate' } };
|
||||
expect((await request.post('/pm_execution_fixture/gate/hold')).ok()).toBe(true);
|
||||
const id = randomUUID();
|
||||
try {
|
||||
expect((await request.post('/prompt', { data: { prompt, prompt_id: id } })).ok()).toBe(true);
|
||||
let waiting = [];
|
||||
await expect.poll(async () => {
|
||||
waiting = (await (await request.get('/pm_execution_fixture/monitor')).json()).awaiting;
|
||||
return waiting.filter(name => name.startsWith('bare_')).length;
|
||||
}, { timeout: 20000, message: 'the monitor processed the bare output while its run was open' }).toBe(1);
|
||||
const [bare] = waiting.filter(name => name.startsWith('bare_'));
|
||||
expect((await readLinks(execution.databasePath, [bare])).links[bare]).toBeNull();
|
||||
expect((await request.post('/pm_execution_fixture/gate/release')).ok()).toBe(true);
|
||||
await linkedTo(execution, bare, llmText);
|
||||
expect((await (await request.get('/pm_execution_fixture/monitor')).json()).awaiting).not.toContain(bare);
|
||||
} finally {
|
||||
expect((await request.post('/pm_execution_fixture/gate/release')).ok()).toBe(true);
|
||||
}
|
||||
} finally {
|
||||
expect((await request.post('/pm_execution_fixture/metadata', { data: { disabled: false } })).ok()).toBe(true);
|
||||
}
|
||||
});
|
||||
@@ -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)
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
@@ -2466,6 +2466,7 @@ class TestActivityJobs(AdminAPITestCase):
|
||||
ambiguous=None,
|
||||
untraced=None,
|
||||
records=None,
|
||||
history=None,
|
||||
):
|
||||
key = Path(media_file).relative_to(self.output_dir).as_posix()
|
||||
outcome = names[key]
|
||||
|
||||
+17
-13
@@ -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")
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -0,0 +1,662 @@
|
||||
"""#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"))
|
||||
self.assertIsNone(self.tracker.pop_for_key(key, "another run"))
|
||||
popped = self.tracker.pop_for_key(key)
|
||||
self.assertEqual(popped["id"], older, "an earlier queued item stays first")
|
||||
self.assertTrue(self.tracker.pop_for_key(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_newer_run_drops_stale_entries_of_the_same_key(self):
|
||||
graph = llm_text_graph()
|
||||
# run-1 saved no image, so its entry was never taken
|
||||
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)
|
||||
self.assertEqual(self.tracker._prompt_queue, [])
|
||||
|
||||
def test_a_later_run_queued_before_processing_keeps_its_own_entry(self):
|
||||
# The image appears (snapshot), the next queued run executes the same
|
||||
# node again before the monitor's delay is over, then the image is
|
||||
# processed: it takes its own run's entry, the later run keeps its own.
|
||||
graph = llm_text_graph()
|
||||
own = self.track(graph, "first run text", run_id="run-1")
|
||||
snapshot = self.tracker.keyed_snapshot()
|
||||
later = self.track(graph, "second run text", run_id="run-2")
|
||||
self.assertEqual(self.process({"prompt": graph}, snapshot), own)
|
||||
self.assertEqual(self.process({"prompt": graph}), later)
|
||||
|
||||
def test_no_entry_when_the_image_appeared_means_no_queue_match(self):
|
||||
graph = llm_text_graph()
|
||||
empty_snapshot = self.tracker.keyed_snapshot()
|
||||
self.track(graph, "queued after the image appeared")
|
||||
self.assertIsNone(self.process({"prompt": graph}, empty_snapshot))
|
||||
self.assertEqual(len(self.tracker._prompt_queue), 1)
|
||||
|
||||
def test_hosts_without_run_ids_still_link_queued_entries_in_order(self):
|
||||
graph = llm_text_graph()
|
||||
first = self.track(graph, "item one", run_id=None)
|
||||
second = self.track(graph, "item two", run_id=None)
|
||||
self.assertEqual(self.process({"prompt": graph}), first)
|
||||
self.assertEqual(self.process({"prompt": graph}), second)
|
||||
self.assertIsNone(self.process({"prompt": graph}))
|
||||
|
||||
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"))
|
||||
|
||||
|
||||
class TestReviewFindings(MonitorCase, NodeCase):
|
||||
"""Final-review findings M1-M4 (#344)."""
|
||||
|
||||
def test_m1_a_positive_node_without_a_key_makes_the_image_unkeyable(self):
|
||||
graph = llm_text_graph()
|
||||
graph["30"] = {
|
||||
"class_type": "PromptManager",
|
||||
"inputs": {"text": ["missing", 0], "clip": ["4", 1]},
|
||||
}
|
||||
graph["31"] = {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {**graph["3"]["inputs"], "positive": ["30", 0]},
|
||||
}
|
||||
self.assertIsNone(prompt_node_key(graph, "30"))
|
||||
self.assertEqual(image_run_keys({"prompt": graph}), {})
|
||||
prompt_id = self.save(LLM_TEXT)
|
||||
self.tracker.set_current_prompt(
|
||||
LLM_TEXT,
|
||||
{"prompt_id": prompt_id},
|
||||
run_key=prompt_node_key(graph, "18:35"),
|
||||
run_id="run-1",
|
||||
)
|
||||
self.assertIsNone(self.process({"prompt": graph}))
|
||||
|
||||
def test_m2_links_are_exactly_comfyui_is_link(self):
|
||||
from utils.run_keys import _links
|
||||
|
||||
inputs = {
|
||||
"string_slot_int": ["7", 0],
|
||||
"string_slot_float": ["7", 0.0],
|
||||
"string_slot_bool": ["7", True], # bool is an int to ComfyUI too
|
||||
"int_source": [7, 0],
|
||||
"tuple": ("7", 0),
|
||||
"three": ["7", 0, 1],
|
||||
"slot_text": ["7", "0"],
|
||||
}
|
||||
self.assertEqual(sorted(_links(inputs)), ["7", "7", "7"])
|
||||
widget = issue_graph()
|
||||
widget["18:17"]["inputs"]["size"] = [512, 768] # a value, not a link
|
||||
self.assertIsNotNone(prompt_node_key(widget, "18:35"))
|
||||
floated = issue_graph()
|
||||
floated["18:35"]["inputs"]["prepend_text"] = ["18:13", 0.0]
|
||||
changed = copy.deepcopy(floated)
|
||||
changed["18:17"]["inputs"]["value"] = "another instruction"
|
||||
self.assertNotEqual(
|
||||
prompt_node_key(floated, "18:35"), prompt_node_key(changed, "18:35")
|
||||
)
|
||||
|
||||
def _failing_save(self, node):
|
||||
def explode(*args, **kwargs):
|
||||
raise RuntimeError("database is locked")
|
||||
|
||||
node.db = mock.Mock(wraps=self.db)
|
||||
for name in ("save_prompt", "save_prompt_pair", "get_prompt_by_hash"):
|
||||
setattr(node.db, name, mock.Mock(side_effect=explode))
|
||||
return node
|
||||
|
||||
def test_m3_a_failed_save_takes_over_its_run_key(self):
|
||||
graph = llm_text_graph()
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
older = self.save("previous run text")
|
||||
self.tracker.set_current_prompt(
|
||||
"previous run text", {"prompt_id": older}, run_key=key, run_id="run-1"
|
||||
)
|
||||
node = self._failing_save(self.node(PromptManager))
|
||||
with mock.patch("prompt_manager_base.executing_run_id", return_value="run-2"):
|
||||
node.encode_prompt(
|
||||
self.clip(), "new run text", prompt=graph, unique_id="18:35"
|
||||
)
|
||||
latest, run_id, _count = self.tracker.keyed_snapshot()[key]
|
||||
self.assertTrue(latest.get("placeholder"))
|
||||
self.assertEqual(run_id, "run-2")
|
||||
self.assertIsNone(self.process({"prompt": graph}))
|
||||
|
||||
def test_m3_a_failed_pair_save_takes_over_its_run_key(self):
|
||||
graph = issue_graph(text="typed")
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
node = self._failing_save(self.node(PromptManagerPair))
|
||||
with self.assertLogs("test.prompt_attribution.node", "WARNING"):
|
||||
node.encode_prompt(self.clip(), "typed", prompt=graph, unique_id="18:35")
|
||||
self.assertTrue(self.tracker.keyed_snapshot()[key][0].get("placeholder"))
|
||||
|
||||
def test_m3_a_failed_text_node_save_takes_over_its_run_key(self):
|
||||
graph = issue_graph()
|
||||
graph["18:35"] = {
|
||||
"class_type": "PromptManagerText",
|
||||
"inputs": {"text": "typed"},
|
||||
}
|
||||
graph["40"] = {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": ["18:35", 0], "clip": ["4", 1]},
|
||||
}
|
||||
graph["3"]["inputs"]["positive"] = ["40", 0]
|
||||
node = self._failing_save(self.node(PromptManagerText))
|
||||
with mock.patch.object(
|
||||
node, "_track_prompt_execution", side_effect=RuntimeError("boom")
|
||||
):
|
||||
node.process_text("typed", prompt=graph, unique_id="18:35")
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
self.assertTrue(self.tracker.keyed_snapshot()[key][0].get("placeholder"))
|
||||
|
||||
def test_m4_a_placeholder_never_falls_through_to_history(self):
|
||||
graph = llm_text_graph()
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
self.tracker.mark_empty(key, "run-1", push_to_queue=False)
|
||||
with mock.patch.object(
|
||||
self.handler, "_history_prompt", return_value=99
|
||||
) as history:
|
||||
self.assertIsNone(self.process({"prompt": graph})) # snapshot form
|
||||
self.tracker.mark_empty(key, "run-2", push_to_queue=True)
|
||||
self.assertIsNone(self.process({"prompt": graph})) # popped form
|
||||
history.assert_not_called()
|
||||
|
||||
|
||||
class TestRereviewFindings(MonitorCase, NodeCase):
|
||||
"""Re-review findings R4 and R5 (#344)."""
|
||||
|
||||
def test_r4_a_logging_failure_after_tracking_never_doubles_a_queue_slot(self):
|
||||
graph = llm_text_graph()
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
node = self.node(PromptManager)
|
||||
logger = mock.Mock(wraps=self.tracker.logger)
|
||||
logger.debug.side_effect = RuntimeError("log handler broke")
|
||||
self.tracker.logger = logger
|
||||
node.encode_prompt(self.clip(), LLM_TEXT, prompt=graph, unique_id="18:35")
|
||||
queued = [ctx for ctx in self.tracker._prompt_queue if ctx["run_key"] == key]
|
||||
self.assertEqual(len(queued), 1)
|
||||
|
||||
def test_r5_a_keyless_image_outside_the_output_folder_does_not_wait(self):
|
||||
outside = tempfile.mkdtemp()
|
||||
self.addCleanup(lambda: __import__("shutil").rmtree(outside, True))
|
||||
stray = os.path.join(outside, "copied.png")
|
||||
with open(stray, "wb") as handle:
|
||||
handle.write(b"png")
|
||||
with mock.patch(
|
||||
"utils.image_monitor.host_output_directory", return_value=self.tmp.name
|
||||
):
|
||||
with mock.patch.object(
|
||||
self.handler.metadata_extractor, "extract_metadata", return_value=None
|
||||
):
|
||||
self.handler.process_new_image(stray)
|
||||
self.handler.process_new_image(self.image) # inside: waits
|
||||
self.assertEqual(
|
||||
self.handler.awaiting_history(),
|
||||
[os.path.normcase(os.path.realpath(self.image))],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,255 @@
|
||||
"""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_api_images import ImageAPITestCase
|
||||
from tests.test_prompt_attribution import issue_graph, pair_evidence
|
||||
from tools import review_image_links as tool
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
|
||||
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 ToolFixture(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()
|
||||
|
||||
|
||||
class ReviewToolTests(ToolFixture):
|
||||
def test_dry_run_judges_and_changes_nothing(self):
|
||||
before = self.links()
|
||||
ids = f"{self.guessed[0]}-{self.gone}"
|
||||
code, out = self.run_tool("--ids", ids, "--json")
|
||||
self.assertEqual(code, 0)
|
||||
verdicts = {v["image_id"]: v for v in json.loads(out.split("--apply, for")[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, 0 skipped", 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")
|
||||
|
||||
|
||||
class ApplyGuardTests(ToolFixture):
|
||||
"""M6: unlink clears only the link, and only the link it was judged on."""
|
||||
|
||||
def test_unlink_skips_a_row_linked_elsewhere_since_the_review(self):
|
||||
db = PromptDatabase(self.db_path)
|
||||
try:
|
||||
verdicts = tool.review(db, [self.guessed[0]], None, self.output)
|
||||
self.assertEqual(verdicts[0]["verdict"], "unlink")
|
||||
db.model.get_connection().execute(
|
||||
"UPDATE generated_images SET prompt_id=? WHERE id=?",
|
||||
(self.typed, self.guessed[0]),
|
||||
)
|
||||
db.model.get_connection().commit()
|
||||
result = tool.apply(db, verdicts, [self.guessed[0]])
|
||||
finally:
|
||||
db.close_all()
|
||||
self.assertEqual(result["unlinked"], 0)
|
||||
self.assertEqual(self.links()[self.guessed[0]], self.typed)
|
||||
|
||||
def test_unlink_keeps_the_file_and_its_favorite(self):
|
||||
path = os.path.join(self.output, "moss_00009_.png")
|
||||
db = PromptDatabase(self.db_path)
|
||||
try:
|
||||
db.set_image_favorite("a" * 64, "moss_00009_.png")
|
||||
finally:
|
||||
db.close_all()
|
||||
self.run_tool("--ids", str(self.guessed[0]), "--apply")
|
||||
self.assertTrue(os.path.isfile(path))
|
||||
db = PromptDatabase(self.db_path)
|
||||
try:
|
||||
self.assertEqual(db.get_root_favorites("a" * 64), {"moss_00009_.png"})
|
||||
summary = db.get_image_link_summaries([path]).get(path) or {}
|
||||
finally:
|
||||
db.close_all()
|
||||
self.assertNotEqual(summary.get("status"), "linked")
|
||||
|
||||
def test_dry_run_states_what_apply_does(self):
|
||||
_code, out = self.run_tool("--ids", str(self.guessed[0]))
|
||||
self.assertIn("removes only the image's link record", out)
|
||||
self.assertIn("file stays on disk and in the gallery as unlinked", out)
|
||||
|
||||
|
||||
class GalleryAfterUnlinkTests(ImageAPITestCase):
|
||||
"""M6: after --apply the file is still listed, as an unlinked image."""
|
||||
|
||||
async def test_unlinked_file_stays_in_the_workspace_gallery(self):
|
||||
db = self.api.db
|
||||
prompt = db.save_prompt(
|
||||
text="a woman holding a coffee cup",
|
||||
prompt_hash=generate_prompt_hash("a woman holding a coffee cup"),
|
||||
)
|
||||
path = self.output_dir / "moss_00009_.png"
|
||||
png(str(path), issue_graph(), pair_evidence())
|
||||
image_id = db.link_image_to_prompt(prompt, str(path), metadata={})
|
||||
out = io.StringIO()
|
||||
tool.main(
|
||||
[
|
||||
db.model.db_path,
|
||||
"--output-dir",
|
||||
str(self.output_dir),
|
||||
"--ids",
|
||||
str(image_id),
|
||||
"--apply",
|
||||
],
|
||||
out=out,
|
||||
)
|
||||
self.assertIn("Applied: 0 relinked, 1 unlinked", out.getvalue())
|
||||
response = await self.client.get(
|
||||
"/prompt_manager/images/output", params={"mode": "workspace"}
|
||||
)
|
||||
self.assertEqual(response.status, 200)
|
||||
images = {i["filename"]: i for i in (await response.json())["images"]}
|
||||
self.assertIn("moss_00009_.png", images)
|
||||
link = images["moss_00009_.png"].get("linked") or {}
|
||||
self.assertNotEqual(link.get("status"), "linked")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,674 @@
|
||||
"""#342: recorded history attributes a file only to one certain, stored prompt."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import threading
|
||||
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_a_metadata_less_file_links_through_its_recorded_run(self):
|
||||
"""I1: --disable-metadata outputs carry no prompt; history still proves one."""
|
||||
db = self.api.db
|
||||
prompt_id = db.save_prompt(
|
||||
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
||||
)
|
||||
path = self.write("bare_00001_.png")
|
||||
self.record_run(db, str(self.output_dir), [("", path.name)], [LLM_TEXT])
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, {}), "linked")
|
||||
self.assertEqual(self.linked_prompt(path), prompt_id)
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, None), "unchanged")
|
||||
detailed = self.api._ingest_scanned_file(path, {}, detailed=True)
|
||||
self.assertEqual(detailed["status"], "unchanged")
|
||||
self.assertTrue(detailed["title"])
|
||||
|
||||
def test_a_metadata_less_file_without_a_recorded_run_is_still_skipped(self):
|
||||
path = self.write("bare_00002_.png")
|
||||
self.assertIsNone(self.api._ingest_scanned_file(path, {}))
|
||||
self.assertIsNone(self.api._ingest_scanned_file(path, {"workflow": {}}))
|
||||
self.assertIsNone(self.linked_prompt(path))
|
||||
|
||||
def test_the_issue_file_stays_unlinked_even_with_its_pair_record(self):
|
||||
"""Resolved with an empty stored positive: before #342 the scan raised."""
|
||||
db = self.api.db
|
||||
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"], [])
|
||||
|
||||
|
||||
class TestHostHelpers(unittest.TestCase):
|
||||
"""The host lookups degrade to "no proof", never to an error."""
|
||||
|
||||
def test_host_output_directory(self):
|
||||
import sys
|
||||
import types
|
||||
|
||||
from database import run_attribution
|
||||
|
||||
cases = (
|
||||
(types.SimpleNamespace(get_output_directory=lambda: "/out"), "/out"),
|
||||
(types.SimpleNamespace(get_output_directory=lambda: ""), None),
|
||||
(types.SimpleNamespace(get_output_directory=lambda: 1 / 0), None),
|
||||
)
|
||||
original = sys.modules.get("folder_paths")
|
||||
try:
|
||||
for module, expected in cases:
|
||||
sys.modules["folder_paths"] = module
|
||||
self.assertEqual(run_attribution.host_output_directory(), expected)
|
||||
sys.modules["folder_paths"] = None # import fails: outside ComfyUI
|
||||
self.assertIsNone(run_attribution.host_output_directory())
|
||||
finally:
|
||||
if original is None:
|
||||
sys.modules.pop("folder_paths", None)
|
||||
else:
|
||||
sys.modules["folder_paths"] = original
|
||||
|
||||
def test_run_window_needs_integer_times_and_a_readable_file(self):
|
||||
from database.run_attribution import _written_during
|
||||
|
||||
self.assertFalse(_written_during(__file__, None, 5))
|
||||
self.assertFalse(_written_during(os.path.join("no", "such.png"), 0, 5))
|
||||
|
||||
def test_executing_run_id(self):
|
||||
import sys
|
||||
import types
|
||||
|
||||
from utils.run_keys import executing_run_id
|
||||
|
||||
def host(context):
|
||||
def get_executing_context():
|
||||
if isinstance(context, Exception):
|
||||
raise context
|
||||
return context
|
||||
|
||||
return types.SimpleNamespace(get_executing_context=get_executing_context)
|
||||
|
||||
original = sys.modules.get("comfy_execution.utils")
|
||||
try:
|
||||
for context, expected in (
|
||||
(types.SimpleNamespace(prompt_id="run-1"), "run-1"),
|
||||
(types.SimpleNamespace(prompt_id=""), None),
|
||||
(None, None),
|
||||
(RuntimeError("no context"), None),
|
||||
):
|
||||
sys.modules["comfy_execution.utils"] = host(context)
|
||||
self.assertEqual(executing_run_id(), expected)
|
||||
sys.modules["comfy_execution.utils"] = None
|
||||
self.assertIsNone(executing_run_id())
|
||||
finally:
|
||||
if original is None:
|
||||
sys.modules.pop("comfy_execution.utils", None)
|
||||
else:
|
||||
sys.modules["comfy_execution.utils"] = original
|
||||
|
||||
|
||||
class TestMonitorHistoryRecheck(RecordedRuns, unittest.TestCase):
|
||||
"""I1: a keyless image gets one history-only re-check when its run is sealed."""
|
||||
|
||||
def setUp(self):
|
||||
history_tests.HistoryDatabaseTests.setUp(self)
|
||||
self.db = history_tests.HistoryDatabaseTests.database(self)
|
||||
self.output = str(Path(self.temp.name) / "output")
|
||||
os.makedirs(self.output)
|
||||
self.handler = ImageGenerationHandler(self.db, make_tracker(self.db))
|
||||
self.handler.processing_delay = 0
|
||||
self.handler.settle_interval = 0.01
|
||||
self.addCleanup(self.handler.shutdown)
|
||||
self.linked = []
|
||||
self.handler.link_image_to_prompt = lambda path, prompt, meta: (
|
||||
self.linked.append((os.path.basename(path), prompt["id"]))
|
||||
)
|
||||
patches = (
|
||||
mock.patch(
|
||||
"utils.image_monitor.host_output_directory", return_value=self.output
|
||||
),
|
||||
mock.patch.object(
|
||||
self.handler.metadata_extractor, "extract_metadata", return_value=None
|
||||
),
|
||||
)
|
||||
for patch in patches:
|
||||
patch.start()
|
||||
self.addCleanup(patch.stop)
|
||||
|
||||
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
||||
|
||||
def image(self, name):
|
||||
path = os.path.join(self.output, name)
|
||||
Path(path).write_bytes(b"png")
|
||||
return path
|
||||
|
||||
def save(self, text):
|
||||
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
||||
|
||||
def wait_for(self, predicate, timeout=5.0):
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
if predicate():
|
||||
return True
|
||||
time.sleep(0.01)
|
||||
return predicate()
|
||||
|
||||
def test_the_run_ending_after_processing_links_the_image_once(self):
|
||||
prompt_id = self.save(LLM_TEXT)
|
||||
path = self.image("late_00001_.png")
|
||||
self.handler.process_new_image(path) # the run is still executing
|
||||
self.assertEqual(self.linked, [])
|
||||
self.assertIn(
|
||||
os.path.normcase(os.path.realpath(path)), self.handler.awaiting_history()
|
||||
)
|
||||
self.record_run(self.db, self.output, [("", "late_00001_.png")], [LLM_TEXT])
|
||||
self.assertTrue(self.wait_for(lambda: self.linked))
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(self.linked, [("late_00001_.png", prompt_id)])
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
# Another run reporting the same name later is not a second re-check
|
||||
self.record_run(
|
||||
self.db, self.output, [("", "late_00001_.png")], [LLM_TEXT], item=1
|
||||
)
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(len(self.linked), 1)
|
||||
|
||||
def test_a_seal_during_the_first_pass_is_not_lost(self):
|
||||
"""R1: the run seals while the image's own pass is still pending."""
|
||||
prompt_id = self.save(LLM_TEXT)
|
||||
path = self.image("race_00001_.png")
|
||||
real = self.handler._history_prompt
|
||||
sealed = []
|
||||
|
||||
def history(image_path):
|
||||
result = real(image_path) # before the run's outputs commit: None
|
||||
if not sealed:
|
||||
sealed.append(True)
|
||||
# The writer commits and notifies now; the path is still pending
|
||||
self.record_run(
|
||||
self.db, self.output, [("", "race_00001_.png")], [LLM_TEXT]
|
||||
)
|
||||
return result
|
||||
|
||||
self.handler._history_prompt = history
|
||||
self.assertTrue(self.handler.schedule(path))
|
||||
self.assertTrue(self.wait_for(lambda: self.linked))
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(self.linked, [("race_00001_.png", prompt_id)])
|
||||
|
||||
def test_a_seal_while_the_pending_set_is_full_is_rechecked_later(self):
|
||||
from utils import image_monitor
|
||||
|
||||
prompt_id = self.save(LLM_TEXT)
|
||||
path = self.image("full_00001_.png")
|
||||
self.handler.process_new_image(path) # waiting, unlinked
|
||||
blocker = self.image("blocker.png")
|
||||
release = threading.Event()
|
||||
original = self.handler.process_new_image
|
||||
|
||||
def slow(image_path, prompt_snapshot=None):
|
||||
if image_path == blocker:
|
||||
release.wait(5)
|
||||
return original(image_path, prompt_snapshot=prompt_snapshot)
|
||||
|
||||
self.handler.process_new_image = slow
|
||||
with mock.patch.object(image_monitor, "MAX_PENDING_IMAGES", 1):
|
||||
self.assertTrue(self.handler.schedule(blocker))
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 1))
|
||||
self.record_run(self.db, self.output, [("", "full_00001_.png")], [LLM_TEXT])
|
||||
self.assertEqual(self.linked, [], "not scheduled while the set is full")
|
||||
release.set()
|
||||
self.assertTrue(self.wait_for(lambda: self.linked))
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(self.linked, [("full_00001_.png", prompt_id)])
|
||||
|
||||
def test_a_run_naming_nothing_stored_leaves_the_image_unlinked(self):
|
||||
path = self.image("empty_00001_.png")
|
||||
self.handler.process_new_image(path)
|
||||
self.record_run(self.db, self.output, [("", "empty_00001_.png")], [""])
|
||||
self.assertTrue(self.wait_for(lambda: not self.handler.awaiting_history()))
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(self.linked, [])
|
||||
|
||||
def test_history_found_at_once_leaves_nothing_awaiting(self):
|
||||
prompt_id = self.save(LLM_TEXT)
|
||||
self.record_run(self.db, self.output, [("", "done_00001_.png")], [LLM_TEXT])
|
||||
self.handler.process_new_image(self.image("done_00001_.png"))
|
||||
self.assertEqual(self.linked, [("done_00001_.png", prompt_id)])
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
|
||||
def test_keyed_images_do_not_wait_for_history(self):
|
||||
with mock.patch.object(
|
||||
self.handler.metadata_extractor,
|
||||
"extract_metadata",
|
||||
return_value={"prompt": llm_text_graph()},
|
||||
):
|
||||
self.handler.process_new_image(self.image("keyed_00001_.png"))
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
|
||||
def test_the_waiting_set_is_bounded_and_expires(self):
|
||||
from utils import image_monitor
|
||||
|
||||
with mock.patch.object(image_monitor, "MAX_AWAITING_HISTORY", 2):
|
||||
for n in range(3):
|
||||
self.handler.process_new_image(self.image(f"b_{n}.png"))
|
||||
self.assertEqual(
|
||||
[os.path.basename(p) for p in self.handler.awaiting_history()],
|
||||
["b_1.png", "b_2.png"],
|
||||
)
|
||||
later = time.monotonic() + image_monitor.AWAIT_HISTORY_SECONDS + 1
|
||||
with mock.patch.object(image_monitor.time, "monotonic", return_value=later):
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
|
||||
def test_shutdown_forgets_waiting_images_and_a_dead_handler_is_dropped(self):
|
||||
import gc
|
||||
|
||||
from database import run_attribution
|
||||
|
||||
self.handler.process_new_image(self.image("gone_00001_.png"))
|
||||
self.handler.shutdown()
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
other = ImageGenerationHandler(self.db, make_tracker(self.db))
|
||||
with mock.patch.object(
|
||||
other.metadata_extractor, "extract_metadata", return_value=None
|
||||
):
|
||||
other.process_new_image(self.image("gone_00002_.png"))
|
||||
del other
|
||||
gc.collect()
|
||||
run_attribution.notify_outputs_sealed("run", []) # no dead callback raises
|
||||
|
||||
|
||||
class TestBoundedLookup(RecordedRuns, unittest.TestCase):
|
||||
"""M5: history lookups are batched, and the history schema is unchanged."""
|
||||
|
||||
setUp = history_tests.HistoryDatabaseTests.setUp
|
||||
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
||||
database = history_tests.HistoryDatabaseTests.database
|
||||
|
||||
def output_lookups(self, db, action):
|
||||
statements = []
|
||||
conn = db.model.get_connection()
|
||||
conn.set_trace_callback(statements.append)
|
||||
try:
|
||||
result = action()
|
||||
finally:
|
||||
conn.set_trace_callback(None)
|
||||
return result, [s for s in statements if "FROM execution_run_outputs" in s]
|
||||
|
||||
def test_one_query_per_chunk_not_per_file_and_same_answers(self):
|
||||
from database import run_attribution
|
||||
|
||||
db = self.database()
|
||||
output = str(Path(self.temp.name) / "output")
|
||||
os.makedirs(output)
|
||||
prompt_id = db.save_prompt(
|
||||
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
||||
)
|
||||
names = [f"img_{n:05}_.png" for n in range(25)]
|
||||
paths = [os.path.join(output, name) for name in names]
|
||||
for path in paths:
|
||||
Path(path).write_bytes(b"png")
|
||||
self.record_run(db, output, [("", n) for n in names[:10]], [LLM_TEXT])
|
||||
self.record_run(db, output, [("", n) for n in names[5:10]], [""], item=1)
|
||||
with mock.patch.object(run_attribution, "LOOKUP_CHUNK", 10):
|
||||
result, lookups = self.output_lookups(
|
||||
db, lambda: run_attribution.stored_run_prompts(db, paths, output)
|
||||
)
|
||||
self.assertEqual(len(lookups), 3) # 25 names in chunks of 10
|
||||
expected = {path: None for path in paths}
|
||||
expected.update({path: prompt_id for path in paths[:5]})
|
||||
self.assertEqual(result, expected)
|
||||
for path in paths[:12]:
|
||||
self.assertEqual(
|
||||
run_attribution.stored_run_prompt(db, path, output), expected[path]
|
||||
)
|
||||
|
||||
def test_the_history_schema_gains_no_index(self):
|
||||
from database.execution_history_schema import history_schema_error
|
||||
|
||||
db = self.database()
|
||||
conn = db.model.get_connection()
|
||||
indexes = [
|
||||
row[1]
|
||||
for row in conn.execute("PRAGMA index_list(execution_run_outputs)")
|
||||
if row[3] != "pk"
|
||||
]
|
||||
self.assertEqual(indexes, [])
|
||||
self.assertIsNone(history_schema_error(conn, allow_missing=False))
|
||||
|
||||
|
||||
class TestScanBatchLookup(RecordedRuns, AdminAPITestCase):
|
||||
"""M5: a scan batch asks history once per chunk of names, not per file."""
|
||||
|
||||
def test_a_batch_its_metadata_links_never_asks_history(self):
|
||||
"""R3: the batch lookup is made only when a file needs it."""
|
||||
files = []
|
||||
for n in range(3):
|
||||
path = self.output_dir / f"typed_{n:05}_.png"
|
||||
path.write_bytes(b"png")
|
||||
files.append(path)
|
||||
meta = {"prompt": json.dumps(issue_graph(text="typed mountains"))}
|
||||
statements = []
|
||||
conn = self.api.db.model.get_connection()
|
||||
conn.set_trace_callback(statements.append)
|
||||
try:
|
||||
counts = self.api._ingest_scan_batch_sync([(f, meta) for f in files])
|
||||
finally:
|
||||
conn.set_trace_callback(None)
|
||||
self.assertEqual(counts["added"] + counts["linked"], 3, counts)
|
||||
lookups = [s for s in statements if "FROM execution_run_outputs" in s]
|
||||
self.assertEqual(lookups, [])
|
||||
|
||||
def test_metadata_less_batch_links_with_one_history_lookup(self):
|
||||
db = self.api.db
|
||||
prompt_id = db.save_prompt(
|
||||
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
||||
)
|
||||
files = []
|
||||
for n in range(6):
|
||||
path = self.output_dir / f"bare_{n:05}_.png"
|
||||
path.write_bytes(b"png")
|
||||
files.append(path)
|
||||
self.record_run(
|
||||
db, str(self.output_dir), [("", f.name) for f in files[:4]], [LLM_TEXT]
|
||||
)
|
||||
statements = []
|
||||
conn = db.model.get_connection()
|
||||
conn.set_trace_callback(statements.append)
|
||||
try:
|
||||
counts = self.api._ingest_scan_batch_sync([(f, {}) for f in files])
|
||||
finally:
|
||||
conn.set_trace_callback(None)
|
||||
lookups = [s for s in statements if "FROM execution_run_outputs" in s]
|
||||
self.assertEqual(len(lookups), 1)
|
||||
self.assertEqual(counts["linked"], 4)
|
||||
linked = {
|
||||
row[0]
|
||||
for row in conn.execute(
|
||||
"SELECT filename FROM generated_images WHERE prompt_id=?",
|
||||
(prompt_id,),
|
||||
)
|
||||
}
|
||||
self.assertEqual(linked, {f.name for f in files[:4]})
|
||||
|
||||
|
||||
class TestLinkedOutputRoot(RecordedRuns, unittest.TestCase):
|
||||
"""R2: ComfyUI's output folder reached through a symlink or junction."""
|
||||
|
||||
setUp = history_tests.HistoryDatabaseTests.setUp
|
||||
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
||||
database = history_tests.HistoryDatabaseTests.database
|
||||
|
||||
def linked_root(self):
|
||||
self.real = str(Path(self.temp.name) / "big_drive" / "output")
|
||||
os.makedirs(os.path.join(self.real, "sub"))
|
||||
link = str(Path(self.temp.name) / "output_link")
|
||||
try:
|
||||
os.symlink(self.real, link, target_is_directory=True)
|
||||
except (OSError, NotImplementedError) as error:
|
||||
self.skipTest(f"symlinks unavailable: {error}")
|
||||
return link
|
||||
|
||||
def test_either_spelling_of_the_file_is_attributed(self):
|
||||
db = self.database()
|
||||
link = self.linked_root()
|
||||
prompt_id = db.save_prompt(
|
||||
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
||||
)
|
||||
real_file = os.path.join(self.real, "sub", "img_00001_.png")
|
||||
Path(real_file).write_bytes(b"png")
|
||||
# History keeps the host's spelling of its output folder (the link)
|
||||
self.record_run(db, link, [("sub", "img_00001_.png")], [LLM_TEXT])
|
||||
link_file = os.path.join(link, "sub", "img_00001_.png")
|
||||
self.assertEqual(stored_run_prompt(db, real_file, link), prompt_id)
|
||||
self.assertEqual(stored_run_prompt(db, link_file, link), prompt_id)
|
||||
# The real folder named as the root is another root key: not attributed
|
||||
self.assertIsNone(stored_run_prompt(db, real_file, self.real))
|
||||
|
||||
def test_the_monitor_rechecks_a_file_seen_through_the_real_path(self):
|
||||
db = self.database()
|
||||
link = self.linked_root()
|
||||
prompt_id = db.save_prompt(
|
||||
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
||||
)
|
||||
real_file = os.path.join(self.real, "late_00001_.png")
|
||||
Path(real_file).write_bytes(b"png")
|
||||
handler = ImageGenerationHandler(db, make_tracker(db))
|
||||
handler.processing_delay = 0
|
||||
handler.settle_interval = 0.01
|
||||
self.addCleanup(handler.shutdown)
|
||||
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=link),
|
||||
mock.patch.object(
|
||||
handler.metadata_extractor, "extract_metadata", return_value=None
|
||||
),
|
||||
):
|
||||
handler.process_new_image(real_file)
|
||||
self.assertEqual(linked, [])
|
||||
self.record_run(db, link, [("", "late_00001_.png")], [LLM_TEXT])
|
||||
deadline = time.monotonic() + 5
|
||||
while not linked and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
self.assertEqual(linked, [prompt_id])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -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):
|
||||
|
||||
Vendored
+4
-1
@@ -29,15 +29,18 @@ test('execution script discovers only the new matrix and preserves all 78 origin
|
||||
assert.doesNotMatch(stock, /execution-history/);
|
||||
assert.doesNotMatch(stock, /execution-outputs/);
|
||||
assert.doesNotMatch(stock, /library-location/); // it starts its own fresh host (#249 S6b)
|
||||
assert.doesNotMatch(stock, /image-attribution/); // needs the owned fixture pack (#342)
|
||||
const executions = list('tools/vendor/playwright.execution.config.mjs');
|
||||
assert.match(executions, /\[classic-start\]/);
|
||||
assert.match(executions, /\[nodes2-start\]/);
|
||||
assert.match(executions, /Total: 86 tests in 4 files/); // + the fresh-host library location case, + #343 linked text
|
||||
// + the fresh-host library location case, + #342 image attribution, + #343 linked text
|
||||
assert.match(executions, /Total: 90 tests in 5 files/);
|
||||
for (const project of ['classic-start', 'nodes2-start']) {
|
||||
const cases = executions.split('\n').filter(line => line.includes(`[${project}]`));
|
||||
assert.equal(cases.filter(line => line.includes('execution-history.spec.mjs')).length, 15);
|
||||
assert.equal(cases.filter(line => line.includes('execution-outputs.spec.mjs')).length, 8);
|
||||
assert.equal(cases.filter(line => line.includes('library-location.spec.mjs')).length, 1);
|
||||
assert.equal(cases.filter(line => line.includes('image-attribution.spec.mjs')).length, 2);
|
||||
assert.equal(cases.filter(line => line.includes('execution-text-link.spec.mjs')).length, 19);
|
||||
}
|
||||
assert.doesNotMatch(executions, /prompt-manager.spec/);
|
||||
|
||||
@@ -0,0 +1,330 @@
|
||||
"""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, and only while each is still linked to the prompt it was judged on.
|
||||
|
||||
What ``--apply`` does (APPLY_EFFECTS below, printed by every run):
|
||||
|
||||
- ``relink`` points the image's link record at the other prompt and drops its
|
||||
pair links (they named the old positive).
|
||||
- ``unlink`` removes only the image's link record (its ``generated_images`` row
|
||||
with the metadata snapshot stored in it, and its pair links). That is how the
|
||||
app represents an unlinked image, exactly like a file the scan could not
|
||||
attribute: the file stays on disk and in the gallery as unlinked, and its
|
||||
favorite star and tags (kept by file, not by link) are untouched.
|
||||
|
||||
A batch item linked through the monitor's queue (PromptSearchList) can't be
|
||||
proven after the fact and reads ``unlink``: review before applying.
|
||||
|
||||
python3 tools/review_image_links.py <db> --output-dir <ComfyUI output> \\
|
||||
--ids 58005-58010
|
||||
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")
|
||||
APPLY_EFFECTS = (
|
||||
"--apply, for ids named with --ids only, while each is still linked to the "
|
||||
"prompt it was judged on:\n"
|
||||
" relink: points the image's link record at the other prompt and drops "
|
||||
"its pair links;\n"
|
||||
" unlink: removes only the image's link record (its stored metadata "
|
||||
"snapshot and pair links go with it); the file stays on disk and in the "
|
||||
"gallery as unlinked, with its favorite and tags.\n"
|
||||
)
|
||||
# Ids per query (well under SQLite's bound-parameter limit)
|
||||
ID_BATCH = 500
|
||||
|
||||
|
||||
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.
|
||||
|
||||
An image whose link changed since it was judged (now another prompt, or
|
||||
gone) is skipped.
|
||||
"""
|
||||
allowed = set(ids)
|
||||
chosen = [
|
||||
v
|
||||
for v in verdicts
|
||||
if v["verdict"] in ("relink", "unlink") and v["image_id"] in allowed
|
||||
]
|
||||
with db._pair_transaction() as conn:
|
||||
current = [
|
||||
item
|
||||
for item in chosen
|
||||
if (
|
||||
row := conn.execute(
|
||||
"SELECT prompt_id FROM generated_images WHERE id=?",
|
||||
(item["image_id"],),
|
||||
).fetchone()
|
||||
)
|
||||
and row[0] == item["linked_prompt"]
|
||||
]
|
||||
relink = [item for item in current if item["verdict"] == "relink"]
|
||||
for item in relink:
|
||||
# A pair link records the old positive: it goes with the old link
|
||||
conn.execute(
|
||||
"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",
|
||||
[item["image_id"] for item in current if item["verdict"] == "unlink"],
|
||||
)
|
||||
return {
|
||||
"relinked": len(relink),
|
||||
"unlinked": removed,
|
||||
"skipped": len(chosen) - len(current),
|
||||
}
|
||||
|
||||
|
||||
def backup_database(source, target):
|
||||
"""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(APPLY_EFFECTS)
|
||||
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, "
|
||||
f"{result['skipped']} skipped (link changed since review)\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())
|
||||
Vendored
+2
-1
@@ -19,7 +19,8 @@ const servers = [{
|
||||
|
||||
export default defineConfig({
|
||||
testDir: '../../tests/comfyui',
|
||||
testIgnore: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'execution-text-link.spec.mjs', 'library-location.spec.mjs'],
|
||||
testIgnore: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'execution-text-link.spec.mjs', 'library-location.spec.mjs',
|
||||
'image-attribution.spec.mjs'],
|
||||
outputDir: './test-results/comfyui',
|
||||
timeout: 60000,
|
||||
workers: 1,
|
||||
|
||||
+1
-1
@@ -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', 'execution-text-link.spec.mjs', 'library-location.spec.mjs'],
|
||||
testMatch: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'execution-text-link.spec.mjs', 'library-location.spec.mjs', 'image-attribution.spec.mjs'],
|
||||
outputDir: './test-results/executions/cases',
|
||||
timeout: 60000,
|
||||
workers: 1,
|
||||
|
||||
+254
-82
@@ -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
|
||||
"""
|
||||
|
||||
@@ -27,7 +28,8 @@ import os
|
||||
import queue
|
||||
import time
|
||||
import threading
|
||||
from typing import Callable, Optional, Dict, Any, Tuple
|
||||
from collections import OrderedDict
|
||||
from typing import Callable, Optional, Dict, Any, List, Tuple
|
||||
from watchdog.observers import Observer
|
||||
from watchdog.events import FileSystemEventHandler
|
||||
|
||||
@@ -36,18 +38,38 @@ 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,
|
||||
in_output_folder,
|
||||
stored_run_prompt,
|
||||
subscribe_outputs_sealed,
|
||||
)
|
||||
except ImportError:
|
||||
from database.pair_ingestion import ingest_resolved_pair
|
||||
from database.run_attribution import (
|
||||
host_output_directory,
|
||||
in_output_folder,
|
||||
stored_run_prompt,
|
||||
subscribe_outputs_sealed,
|
||||
)
|
||||
|
||||
# Upper bound on images waiting to be processed (a burst beyond this is skipped)
|
||||
MAX_PENDING_IMAGES = 1000
|
||||
# A file counts as fully written once its size is unchanged across two reads
|
||||
SETTLE_INTERVAL_SECONDS = 0.1
|
||||
SETTLE_ATTEMPTS = 20
|
||||
# Keyless images (no embedded API graph) left unlinked wait for their run's
|
||||
# recorded outputs for one history-only re-check (#342): at most this many,
|
||||
# each for at most this long.
|
||||
MAX_AWAITING_HISTORY = 1000
|
||||
AWAIT_HISTORY_SECONDS = 3600
|
||||
# ``schedule`` snapshot meaning "history-only re-check" (not a new image)
|
||||
HISTORY_ONLY = object()
|
||||
|
||||
|
||||
class ImageGenerationHandler(FileSystemEventHandler):
|
||||
@@ -109,6 +131,14 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
self._work: "queue.Queue[Optional[str]]" = queue.Queue()
|
||||
self._stop = threading.Event()
|
||||
self._worker: Optional[threading.Thread] = None
|
||||
# Keyless unlinked images awaiting their run's recorded outputs:
|
||||
# normalized path -> (path, monotonic deadline), oldest first.
|
||||
self._awaiting: "OrderedDict[str, Tuple[str, float]]" = OrderedDict()
|
||||
# Sealed re-checks ``schedule`` refused (the path still pending, or the
|
||||
# pending set full): retried each time a pending image finishes.
|
||||
self._history_due: "OrderedDict[str, None]" = OrderedDict()
|
||||
self._awaiting_lock = threading.Lock()
|
||||
self._sealed_subscribed = False
|
||||
|
||||
def on_created(self, event):
|
||||
"""Handle filesystem creation events.
|
||||
@@ -126,15 +156,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.
|
||||
@@ -174,6 +206,9 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
and exits; the next ``schedule`` starts a fresh worker with its own
|
||||
queue. Safe to call repeatedly and before anything was scheduled.
|
||||
"""
|
||||
with self._awaiting_lock:
|
||||
self._awaiting.clear()
|
||||
self._history_due.clear()
|
||||
with self._pending_lock:
|
||||
self._pending.clear()
|
||||
worker, self._worker = self._worker, None
|
||||
@@ -229,6 +264,7 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
finally:
|
||||
with self._pending_lock:
|
||||
self._pending.pop(image_path, None)
|
||||
self._drain_history_due()
|
||||
|
||||
def _queue_thumbnail(self, image_path: str) -> None:
|
||||
"""Queue a settled output's preview, best effort (linked or not).
|
||||
@@ -290,15 +326,20 @@ 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:
|
||||
if prompt_snapshot is HISTORY_ONLY:
|
||||
self._recheck_history(image_path)
|
||||
return
|
||||
self.logger.info(f"Processing image: {image_path}")
|
||||
|
||||
if not os.path.exists(image_path):
|
||||
@@ -329,43 +370,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 +408,193 @@ 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).
|
||||
|
||||
The snapshot taken when the image appeared names the key's latest run.
|
||||
|
||||
1. Batch queue: the oldest queued entry of that key from that run.
|
||||
Batch items (text from PromptSearchList) share one key and are
|
||||
linked in order; an empty item's placeholder leaves its image
|
||||
unlinked; entries queued by later runs stay for their own images.
|
||||
2. The snapshot's entry itself, when its run is known and registered
|
||||
exactly one entry for the key (a batch's latest item says nothing
|
||||
about the other images). This covers a cached node that did not
|
||||
run again.
|
||||
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 not keys:
|
||||
# No tracked run can be proven: history now, and once more after
|
||||
# the run's outputs are recorded (``_on_outputs_sealed``). A file
|
||||
# outside ComfyUI's output folder can never be attributed, so it
|
||||
# does not take a waiting slot.
|
||||
if not in_output_folder(image_path, host_output_directory()):
|
||||
return None
|
||||
self._await_history(image_path)
|
||||
prompt_id = self._history_prompt(image_path)
|
||||
if prompt_id is None:
|
||||
return None
|
||||
self._stop_awaiting(image_path)
|
||||
self.logger.info(f"History match: prompt {prompt_id} for {name}")
|
||||
return {"id": prompt_id, "from_history": True}
|
||||
if len(keys) == 1:
|
||||
(key,) = keys.values()
|
||||
if prompt_snapshot is None:
|
||||
prompt_snapshot = self._keyed_snapshot()
|
||||
# No entry of the key when the image appeared: no tracked run made it
|
||||
entry = (prompt_snapshot or {}).get(key)
|
||||
if entry is not None:
|
||||
context, run_id, count = entry
|
||||
queued = self.prompt_tracker.pop_for_key(key, run_id)
|
||||
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 context.get("placeholder"):
|
||||
# Same rule as a popped placeholder: that run had no prompt
|
||||
self.logger.info(f"Latest run for {name} had no prompt text")
|
||||
return None
|
||||
if run_id and count == 1:
|
||||
self.logger.info(f"Run match: prompt {context['id']} for {name}")
|
||||
return context
|
||||
else:
|
||||
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}
|
||||
|
||||
@staticmethod
|
||||
def _awaiting_key(image_path: str) -> str:
|
||||
"""One spelling for a waiting file, through symlinks and junctions."""
|
||||
try:
|
||||
resolved = os.path.realpath(image_path)
|
||||
except (OSError, ValueError):
|
||||
resolved = os.path.abspath(image_path)
|
||||
return os.path.normcase(resolved)
|
||||
|
||||
def _prune_awaiting(self) -> None:
|
||||
"""Drop expired and surplus entries (``_awaiting_lock`` held)."""
|
||||
now = time.monotonic()
|
||||
for key in [k for k, (_p, due) in self._awaiting.items() if due <= now]:
|
||||
del self._awaiting[key]
|
||||
while len(self._awaiting) > MAX_AWAITING_HISTORY:
|
||||
self._awaiting.popitem(last=False)
|
||||
|
||||
def _await_history(self, image_path: str) -> None:
|
||||
"""Re-check ``image_path`` once, history only, when its run is sealed."""
|
||||
if not self._sealed_subscribed:
|
||||
subscribe_outputs_sealed(self._on_outputs_sealed)
|
||||
self._sealed_subscribed = True
|
||||
key = self._awaiting_key(image_path)
|
||||
with self._awaiting_lock:
|
||||
self._awaiting.pop(key, None)
|
||||
self._awaiting[key] = (
|
||||
image_path,
|
||||
time.monotonic() + AWAIT_HISTORY_SECONDS,
|
||||
)
|
||||
self._prune_awaiting()
|
||||
|
||||
def _stop_awaiting(self, image_path: str) -> None:
|
||||
with self._awaiting_lock:
|
||||
self._awaiting.pop(self._awaiting_key(image_path), None)
|
||||
|
||||
def awaiting_history(self) -> List[str]:
|
||||
"""Normalized paths still waiting for their run's recorded outputs."""
|
||||
with self._awaiting_lock:
|
||||
self._prune_awaiting()
|
||||
return list(self._awaiting)
|
||||
|
||||
def _on_outputs_sealed(self, run_id: str, references: List[dict]) -> None:
|
||||
"""A run's outputs were recorded: re-check its waiting images once.
|
||||
|
||||
Called on the history writer's thread, so it only schedules the work.
|
||||
"""
|
||||
root = host_output_directory()
|
||||
if not root:
|
||||
return
|
||||
due = []
|
||||
with self._awaiting_lock:
|
||||
self._prune_awaiting()
|
||||
for reference in references:
|
||||
if reference.get("folder_type") != "output":
|
||||
continue
|
||||
key = self._awaiting_key(
|
||||
os.path.join(
|
||||
root, reference.get("subfolder", ""), reference["filename"]
|
||||
)
|
||||
)
|
||||
waiting = self._awaiting.pop(key, None)
|
||||
if waiting is not None:
|
||||
due.append(waiting[0])
|
||||
for image_path in due:
|
||||
self._schedule_recheck(image_path)
|
||||
|
||||
def _schedule_recheck(self, image_path: str) -> None:
|
||||
"""Schedule the history-only re-check, or keep it due for later."""
|
||||
if self.schedule(image_path, HISTORY_ONLY):
|
||||
return
|
||||
with self._awaiting_lock:
|
||||
self._history_due.pop(image_path, None)
|
||||
self._history_due[image_path] = None
|
||||
while len(self._history_due) > MAX_AWAITING_HISTORY:
|
||||
self._history_due.popitem(last=False)
|
||||
|
||||
def _drain_history_due(self) -> None:
|
||||
"""Schedule due re-checks now that a pending slot or path is free."""
|
||||
with self._awaiting_lock:
|
||||
if not self._history_due:
|
||||
return
|
||||
due = list(self._history_due)
|
||||
for image_path in due:
|
||||
if self.schedule(image_path, HISTORY_ONLY):
|
||||
with self._awaiting_lock:
|
||||
self._history_due.pop(image_path, None)
|
||||
|
||||
def _recheck_history(self, image_path: str) -> None:
|
||||
"""The one history-only re-check of a keyless image (#342)."""
|
||||
name = os.path.basename(image_path)
|
||||
if not os.path.exists(image_path):
|
||||
return
|
||||
prompt_id = self._history_prompt(image_path)
|
||||
if prompt_id is None:
|
||||
self.logger.info(
|
||||
f"No certain prompt for {name} after its run ended; left unlinked"
|
||||
)
|
||||
return
|
||||
self.logger.info(f"History match after its run ended: {prompt_id} for {name}")
|
||||
try:
|
||||
metadata = self.metadata_extractor.extract_metadata(image_path)
|
||||
except Exception:
|
||||
metadata = None
|
||||
self.link_image_to_prompt(
|
||||
image_path,
|
||||
{"id": prompt_id, "from_history": True},
|
||||
metadata or {"file_info": self.get_basic_file_info(image_path)},
|
||||
)
|
||||
|
||||
def _history_prompt(self, image_path: str) -> Optional[int]:
|
||||
"""The stored prompt recorded history attributes to this file, if one."""
|
||||
try:
|
||||
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 +744,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 +788,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}")
|
||||
|
||||
+151
-9
@@ -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, 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,20 +175,20 @@ class PromptTracker:
|
||||
with self.lock:
|
||||
self.active_prompts[execution_id] = execution_context
|
||||
|
||||
# Push to batch queue for ordered image linking
|
||||
if push_to_queue:
|
||||
with self._queue_lock:
|
||||
self._prompt_queue.append(execution_context)
|
||||
self.logger.debug(
|
||||
f"Queue size: {len(self._prompt_queue)} after push "
|
||||
f"for prompt {prompt_id}"
|
||||
)
|
||||
self._remember_keyed(execution_context)
|
||||
|
||||
self.logger.debug(
|
||||
f"Set current prompt: {execution_id} -> {prompt_text[:50]}... "
|
||||
f"(thread: {threading.current_thread().ident})"
|
||||
)
|
||||
self.logger.debug(f"Active prompts count: {len(self.active_prompts)}")
|
||||
|
||||
# Push to batch queue for ordered image linking. Last, with nothing
|
||||
# after it that can raise: a caller that marks a failed execution
|
||||
# (#342) must never add a second queue slot for the same item.
|
||||
if push_to_queue:
|
||||
with self._queue_lock:
|
||||
self._prompt_queue.append(execution_context)
|
||||
return execution_id
|
||||
|
||||
def get_current_prompt(self) -> Optional[Dict[str, Any]]:
|
||||
@@ -225,6 +234,99 @@ 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._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 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_key(
|
||||
self, run_key: Optional[str], run_id: Optional[str] = None
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Pop the oldest queued entry of ``run_key`` (FIFO), from ``run_id`` if set.
|
||||
|
||||
``run_id`` is the run the key's latest entry came from when the image
|
||||
appeared (``keyed_snapshot``): entries of that key queued by earlier
|
||||
runs ahead of it are stale and dropped, while entries of later runs
|
||||
(queued after the image appeared) stay for their own images. Entries
|
||||
of other keys stay queued; expired entries are discarded.
|
||||
"""
|
||||
if not run_key:
|
||||
return None
|
||||
now = time.time()
|
||||
with self._queue_lock:
|
||||
live = [
|
||||
ctx
|
||||
for ctx in self._prompt_queue
|
||||
if now - ctx.get("timestamp", 0) < self.prompt_timeout
|
||||
]
|
||||
index = next(
|
||||
(
|
||||
i
|
||||
for i, ctx in enumerate(live)
|
||||
if ctx.get("run_key") == run_key
|
||||
and (run_id is None or ctx.get("run_id") == run_id)
|
||||
),
|
||||
None,
|
||||
)
|
||||
if index is None:
|
||||
self._prompt_queue = live
|
||||
return None
|
||||
found = live[index]
|
||||
# Same-key entries ahead of it came from earlier runs: stale
|
||||
stale_or_found = {
|
||||
i for i in range(index + 1) if live[i].get("run_key") == run_key
|
||||
}
|
||||
self._prompt_queue = [
|
||||
ctx for i, ctx in enumerate(live) if i not in stale_or_found
|
||||
]
|
||||
return found
|
||||
|
||||
def pop_next_prompt(self) -> Optional[Dict[str, Any]]:
|
||||
"""Pop the next prompt from the batch queue (FIFO).
|
||||
|
||||
@@ -463,6 +565,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.
|
||||
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
"""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.
|
||||
|
||||
Equal keys do not always mean equal text. ComfyUI's cache signature also folds
|
||||
in each node's ``IS_CHANGED``, which the key cannot see: an upstream node that
|
||||
forces a re-run (a random line, an "always run" LLM) can give two runs one key
|
||||
and two texts. The key is therefore never used alone: the monitor also requires
|
||||
the entry to come from the run the snapshot taken when the image appeared names
|
||||
(``PromptTracker.pop_for_key``, ``ImageGenerationHandler._find_prompt_from_run``).
|
||||
|
||||
Pure functions; the only host dependency (the executing run id) is optional.
|
||||
"""
|
||||
|
||||
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 is_link(value: Any) -> bool:
|
||||
"""ComfyUI's own test (``comfy_execution.graph_utils.is_link``), exactly.
|
||||
|
||||
A two-item list of a ``str`` node id and an ``int`` or ``float`` slot (a
|
||||
``bool`` slot passes too, as it does in ComfyUI). Tuples are values.
|
||||
"""
|
||||
if not isinstance(value, list):
|
||||
return False
|
||||
if len(value) != 2:
|
||||
return False
|
||||
if not isinstance(value[0], str):
|
||||
return False
|
||||
if not isinstance(value[1], int) and not isinstance(value[1], float):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _links(inputs: Any) -> List[str]:
|
||||
if not isinstance(inputs, dict):
|
||||
return []
|
||||
return [value[0] for value in inputs.values() if is_link(value)]
|
||||
|
||||
|
||||
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, and empty when any positive node has
|
||||
no key: one unkeyable node must not leave another looking like the only
|
||||
positive prompt (a two-pass workflow would link the wrong pass).
|
||||
"""
|
||||
graph = metadata.get("prompt") if isinstance(metadata, dict) else None
|
||||
if isinstance(graph, str):
|
||||
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 None:
|
||||
return {}
|
||||
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
|
||||
Reference in New Issue
Block a user