Vito
5583159497
* feat(inspector): use recorded run history when the graph alone can't name the prompt (#386) When a file's own graph leaves the prompt open, both inspect routes ask recorded run history about the file (one run reported it and it was written during that run). The save node that run reported it from anchors the walk like a file-name match, so dataset branches that share one saver prefix show their own prompt, labelled "Recorded run". Still open at one sampler whose PromptManager text was recorded at execution (a language model's output), the recorded text shows as "Execution recorded". Results the file resolves alone never change; a failed lookup judges the file alone. * fix(inspector): recorded texts follow history's saver; conflicts anchor nothing (#386) Review follow-ups: the texts PromptManager recorded are looked up at the sampler history's saver leads to, never at a sampler the unanchored walk named; a file name naming a different saver than history anchors nothing; outputs and texts are read in one transaction; "Recorded run" ranks below every workflow-inferred source; unexpected lookup errors log at error level. * fix(history): look up recorded runs by history's own spelling of the file (#386) The history route inspects a case-normalized target on Windows, while history matches file names exactly, so the occurrence's reference names the file.
369 lines
14 KiB
Python
369 lines
14 KiB
Python
"""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.
|
|
|
|
``recorded_output`` applies the first two rules and returns, for the
|
|
inspector (#386), the one output node (save node) that run reported the file
|
|
from and the texts PromptManager recorded at execution per sampler.
|
|
|
|
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 dataclasses import dataclass
|
|
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), "
|
|
"o.output_node_id "
|
|
"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, nodes)}}``.
|
|
|
|
``nodes`` is the set of output node ids that reported the location in that
|
|
run. 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, node in rows:
|
|
runs = found.setdefault((name, subfolder, root_key, style), {})
|
|
runs.setdefault(run_id, (began, ended, set()))[2].add(node)
|
|
return found
|
|
|
|
|
|
def _recorded_runs(
|
|
db: Any,
|
|
image_paths: Iterable[str],
|
|
output_root: Optional[str],
|
|
path_style: Optional[str],
|
|
) -> Dict[str, tuple]:
|
|
"""``{path: (run_id, output node ids)}`` for each file history attributes.
|
|
|
|
A file is attributed when exactly one recorded run reported its location
|
|
and it was last written while that run executed.
|
|
"""
|
|
paths = list(dict.fromkeys(path for path in image_paths if path))
|
|
if not paths or not output_root:
|
|
return {}
|
|
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 {}
|
|
names = sorted({filename for _key, _sub, filename in located.values()})
|
|
runs = _runs_by_location(db, names)
|
|
if not runs:
|
|
return {}
|
|
attributed = {}
|
|
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, nodes)),) = reported.items()
|
|
if _written_during(path, began, ended):
|
|
attributed[path] = (run_id, nodes)
|
|
return attributed
|
|
|
|
|
|
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)
|
|
prompts: Dict[str, Optional[int]] = {}
|
|
for path, (run_id, _nodes) in _recorded_runs(
|
|
db, paths, output_root, path_style
|
|
).items():
|
|
if run_id not in prompts:
|
|
prompts[run_id] = _named_prompt(db, run_id)
|
|
result[path] = prompts[run_id]
|
|
return result
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RecordedOutput:
|
|
"""What recorded history knows about the run that wrote one file (#386).
|
|
|
|
``saver_id`` is the one output node the run reported the file from (None
|
|
when several did). ``texts`` maps each sampler anchor of the run to the
|
|
effective texts PromptManager recorded at execution for its roles
|
|
(``{"positive": text or None, "negative": text or None}``); a role whose
|
|
text was only inferred from the graph is None.
|
|
"""
|
|
|
|
saver_id: Optional[str]
|
|
texts: Dict[str, Dict[str, Optional[str]]]
|
|
|
|
|
|
def _recorded_texts(db: Any, run_id: str) -> Dict[str, Dict[str, Optional[str]]]:
|
|
texts = {}
|
|
for context in db.get_execution_contexts(run_id).get("contexts") or []:
|
|
evidence = context.get("evidence") or {}
|
|
roles = {}
|
|
for role in ("positive", "negative"):
|
|
values = context.get(role) or {}
|
|
proof = (evidence.get(role) or {}).get("effective") or {}
|
|
value = values.get("effective")
|
|
recorded = proof.get("provenance") == "recorded"
|
|
roles[role] = value if recorded and isinstance(value, str) else None
|
|
texts[context["anchor_id"]] = roles
|
|
return texts
|
|
|
|
|
|
def recorded_output(
|
|
db: Any,
|
|
image_path: str,
|
|
output_root: Optional[str],
|
|
path_style: Optional[str] = None,
|
|
) -> Optional[RecordedOutput]:
|
|
"""The recorded run facts for ``image_path``, else None (#386).
|
|
|
|
The file must be attributed to one run (``_recorded_runs``: one run
|
|
reported its location and it was written during that run). Read-only; it
|
|
never links anything.
|
|
"""
|
|
if not image_path:
|
|
return None
|
|
conn = db.model.get_connection()
|
|
own_transaction = not conn.in_transaction
|
|
if own_transaction: # outputs and texts from one state of history
|
|
conn.execute("BEGIN")
|
|
try:
|
|
found = _recorded_runs(db, [image_path], output_root, path_style)
|
|
if image_path not in found:
|
|
return None
|
|
run_id, nodes = found[image_path]
|
|
saver = next(iter(nodes)) if len(nodes) == 1 else None
|
|
return RecordedOutput(saver, _recorded_texts(db, run_id))
|
|
finally:
|
|
if own_transaction:
|
|
conn.rollback()
|
|
|
|
|
|
_LOGGER = logging.getLogger("prompt_manager.run_attribution")
|
|
|
|
|
|
def recorded_output_lookup(
|
|
db: Any, *, close: bool = False, reference: Optional[dict] = None
|
|
) -> Callable[[Any], Any]:
|
|
"""``lookup(path) -> RecordedOutput or None`` for the inspect routes (#386).
|
|
|
|
Uses ComfyUI's output folder. Any failure (no history, a locked library)
|
|
is logged and names no node, so the file is judged from its own metadata.
|
|
``close`` closes the calling thread's connection afterwards, for routes
|
|
that close theirs after reading. ``reference`` (a history occurrence's
|
|
output reference) spells the file as history recorded it: the route's
|
|
target is case-normalized on Windows, and history matches names exactly.
|
|
"""
|
|
|
|
def lookup(path: Any) -> Optional[RecordedOutput]:
|
|
root = host_output_directory()
|
|
if reference is not None:
|
|
if reference.get("folder_type") != "output" or not root:
|
|
return None
|
|
path = os.path.join(
|
|
root, reference.get("subfolder") or "", reference.get("filename") or ""
|
|
)
|
|
try:
|
|
return recorded_output(db, str(path), root)
|
|
except (sqlite3.Error, OSError, ValueError):
|
|
_LOGGER.warning("Recorded run lookup failed", exc_info=True)
|
|
return None
|
|
except Exception:
|
|
_LOGGER.exception("Recorded run lookup failed unexpectedly")
|
|
return None
|
|
finally:
|
|
if close:
|
|
db.model.close()
|
|
|
|
return lookup
|
|
|
|
|
|
_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)
|