fix(ingestion): link an image only to a prompt that provably produced it (#344)
* fix(ingestion): link an image only to a prompt that provably produced it An LLM fed a PromptManagerPair's prepend_text while its text editor was empty. The node refused to save (an empty positive raised and logged a warning), and the image monitor then linked the image to a stale tracker entry or to the newest prompt in the database. - Nodes: an empty text editor saves and tracks nothing (prepend and append text are never the stored prompt), logs at info, and marks the node's run key empty so no earlier entry can be used for its images. - Tracker entries carry a run key: a digest of the prompt node and its whole upstream graph, the identity ComfyUI's cache uses. The monitor computes the same key from the image's embedded graph and uses a queued or snapshot entry only when the keys match (a snapshot also needs a known run with one entry, so a batch's last item is never applied to every image). - The recency fallback is gone; anything uncertain stays unlinked and is logged at info. - Recorded execution history attributes a file when exactly one run reported it, the file was written during that run, and the run's positive contexts name one stored prompt. The monitor and the output scan use it as their last certain source. - The scan and pair ingestion treat a resolved pair with an empty stored positive as unresolved instead of raising. - tools/review_image_links.py reviews selected links by these rules; dry run by default, --apply changes only named ids after a backup. Closes #342 * test(ingestion): real-host attribution case and the #342 testing contract The managed host disables the image monitor, so the owned fixture pack gains a route that starts and stops the plugin's monitor singleton for the attribution case only. The case runs in both renderers: a typed Pair image links, a Pair image whose positive is only linked prepend text stays unlinked with nothing stored, an unknown node feeding text links through its run key, and its cached re-run links again. * fix(ingestion): take a queued entry only from the run the image appeared in Dropping a key's older queue entries whenever the node ran again let a later queued run, executing during the monitor's processing delay, take the earlier image's entry. The snapshot taken when the image appears now names the key's latest run, and the queue is popped by key and run: same-key entries of earlier runs ahead of it are stale, entries of later runs wait for their own images. Also keeps the attribution case out of the stock real-host suite. * test(ingestion): host lookups degrade to no proof * fix(ingestion): address the #344 review (history re-check, keys, repair) - Images without an embedded API graph: the scan now asks recorded history before skipping a file with no prompt metadata, and the monitor gives a keyless unlinked image one history-only re-check when its run's outputs are recorded (a seal notification from the outputs write; at most 1000 waiting images, each for an hour). - Run keys: an image is unkeyable when any positive node has no key, so a two-pass graph is never linked through the other pass; links follow ComfyUI's is_link exactly. - A failed save marks the node's run key like an empty editor, so an older entry of that key can't stand in for this run. - A placeholder is the latest entry: the image stays unlinked without falling through to history, the same rule as a popped placeholder. - An index on execution_run_outputs(filename, subfolder, root_key, path_style) for the history lookup, added by migration and accepted by the output schema check (any other index is still refused). - Repair tool: unlink removes only the link record and only while the image is still linked to the prompt it was judged on; every run prints what --apply does. - Real-host case proves the re-run's PromptManager node was cached. * test(ingestion): real-host bare-output re-check and the review-round contract A metadata-disabled save, held open behind the fixture gate after SaveImage, waits unlinked while its run is open and links once the run's outputs are recorded. The fixture's monitor route lists the waiting files. TESTING.md covers the history re-check, the read index and what the repair tool's --apply does. * fix(ingestion): batch history lookups instead of adding an index An index on execution_run_outputs would make the library unopenable by older builds, whose schema check refuses any extra index there. The migration and the schema-check allowance are reverted, so the history schema is exactly main's. stored_run_prompts instead looks up a batch of files with one parameterized filename IN (...) query per chunk of at most 500 names and matches in memory; the scan computes it once per scan batch. * fix(ingestion): address the #344 re-review (seal race, linked roots) - A sealed history 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, instead of being dropped. - History matches a file through ComfyUI's spelling of its output folder first, then with both sides resolved, so a symlinked or junction output folder still attributes the scan's resolved paths; waiting files are keyed by their resolved path. - The scan makes its batch history lookup only when a file needs it. - set_current_prompt pushes to the queue last, so marking a failed execution never adds a second queue slot. - A keyless image outside the output folder doesn't take a waiting slot.
This commit is contained in:
Vito
+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