nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
3be0ada6a129cbdd7e72193986b5ccdb2f86fc26
ComfyAssets-ComfyUI_PromptM…/utils/image_monitor.py
T
Vito 3be0ada6a1 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.
2026-10-25 19:03:35 -07:00

1082 lines
45 KiB
Python
Raw Blame History

"""Image monitoring system for ComfyUI generated images.
This module provides real-time monitoring of ComfyUI output directories to
automatically detect newly generated images and associate them with their
corresponding prompts. The system uses filesystem watchers to detect image creation
events and extract metadata from the images to maintain a gallery system.
The main components are:
- ImageGenerationHandler: Handles filesystem events for new image creation
- ImageMonitor: Main monitoring system that manages directory watching
Typical usage:
from utils.image_monitor import ImageMonitor
monitor = ImageMonitor(db_manager, prompt_tracker)
monitor.start_monitoring(['/path/to/comfyui/output'])
The system automatically:
- Detects new image files in monitored directories
- Extracts ComfyUI workflow metadata from PNG chunks
- 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
"""
import os
import queue
import time
import threading
from collections import OrderedDict
from typing import Callable, Optional, Dict, Any, List, Tuple
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
from .hashing import generate_prompt_hash
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):
"""Filesystem event handler for detecting new image generation.
This handler extends watchdog's FileSystemEventHandler to specifically handle
new image file creation events in ComfyUI output directories. When a new image
is detected, it attempts to:
1. Extract ComfyUI metadata from the image
2. Associate the image with the currently active prompt
3. Store the relationship in the database
The handler implements a small delay before processing to ensure files are
completely written before attempting to read them.
"""
def __init__(self, db_manager, prompt_tracker):
"""
Initialize the image generation handler.
Args:
db_manager: Database manager instance for storing image-prompt relationships
prompt_tracker: Prompt tracking instance for getting current active prompts
"""
self.db_manager = db_manager
self.prompt_tracker = prompt_tracker
self.metadata_extractor = ComfyUIMetadataExtractor()
self.logger = get_logger("prompt_manager.image_monitor")
# The Needs review pass for ambiguous outputs (the API's
# ``record_review_file``, wired at start-up): a linked output is
# never queued and its pending items leave the queue. None: no pass.
self.review_recorder: Optional[Callable[[str], str]] = None
# Queues a settled output's preview (the API's
# ``queue_capture_thumbnail``, wired at start-up; #249). It only
# queues: the preview is written on its own worker. None: no queue.
self.thumbnail_queue: Optional[Callable[[str], str]] = None
# Read from GalleryConfig if available, otherwise use defaults
try:
from ..py.config import GalleryConfig
self.processing_delay = GalleryConfig.PROCESSING_DELAY
self.supported_extensions = tuple(GalleryConfig.SUPPORTED_EXTENSIONS)
except Exception:
self.processing_delay = 2.0
self.supported_extensions = (".png", ".jpg", ".jpeg", ".webp", ".gif")
# Settle check: the file size must be unchanged across two reads
self.settle_interval = SETTLE_INTERVAL_SECONDS
self.settle_attempts = SETTLE_ATTEMPTS
# Images scheduled but not yet processed: path -> (prompt snapshot, due time).
# One daemon worker processes them in order, so batch images pop the
# prompt queue in creation order.
self._pending: Dict[str, Tuple[Any, float]] = {}
self._pending_lock = threading.Lock()
# Each worker owns its queue and stop event (replaced on shutdown),
# so a worker that is still draining never swallows a later image.
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.
This method is called by watchdog when a new file is created in a monitored
directory. It filters for image files and schedules them for processing after
a small delay to ensure the file is fully written.
The prompt context is captured immediately (before the delay) to handle
batch workflows where the prompt tracker advances before images are processed.
Args:
event: FileSystemEvent object containing event details
"""
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 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.
Images are processed in the order they were scheduled (batch images must
pop the prompt queue in creation order), each after ``processing_delay``
and once its size has stopped changing.
Returns:
False when the path is already pending or the pending set is full
"""
with self._pending_lock:
if image_path in self._pending:
self.logger.debug(f"Already pending, not rescheduled: {image_path}")
return False
if len(self._pending) >= MAX_PENDING_IMAGES:
self.logger.warning(
f"Too many images pending ({len(self._pending)}), "
f"skipping: {image_path}"
)
return False
due = time.monotonic() + self.processing_delay
self._pending[image_path] = (prompt_snapshot, due)
self._ensure_worker()
self._work.put(image_path)
return True
def pending_count(self) -> int:
"""Number of images scheduled but not yet processed."""
with self._pending_lock:
return len(self._pending)
def shutdown(self) -> None:
"""Stop the worker thread and drop every image still waiting.
The worker wakes from any delay or settle wait, finishes nothing more
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
work, self._work = self._work, queue.Queue()
stop, self._stop = self._stop, threading.Event()
stop.set()
if worker is not None and worker.is_alive():
work.put(None) # sentinel: wake a worker idle on an empty queue
def _ensure_worker(self) -> None:
"""Start the single daemon worker thread if it isn't running.
Called with ``_pending_lock`` held.
"""
if self._worker is None or not self._worker.is_alive():
self._worker = threading.Thread(
target=self._run_worker,
args=(self._work, self._stop),
name="PromptManagerImageWorker",
daemon=True,
)
self._worker.start()
def _run_worker(self, work: "queue.Queue[Optional[str]]", stop) -> None:
while not stop.is_set():
image_path = work.get()
try:
if image_path is None or stop.is_set():
break
self._process_pending(image_path, stop)
except Exception as e:
self.logger.error(f"Image worker failed on {image_path}: {e}")
finally:
work.task_done()
def _process_pending(self, image_path: str, stop=None) -> None:
"""Process one scheduled image; the pending entry is always cleared."""
stop = stop if stop is not None else self._stop
try:
with self._pending_lock:
prompt_snapshot, due = self._pending.get(image_path, (None, 0.0))
remaining = due - time.monotonic()
if remaining > 0 and stop.wait(remaining):
return
if not self.wait_until_settled(image_path, stop):
if not stop.is_set():
self.logger.warning(
f"Image never settled or vanished, skipping: {image_path}"
)
return
self.process_new_image(image_path, prompt_snapshot=prompt_snapshot)
self._queue_thumbnail(image_path)
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).
A failure is logged, never raised: the preview is then made on its
first view.
"""
queue_preview = getattr(self, "thumbnail_queue", None)
if queue_preview is None:
return
try:
outcome = queue_preview(image_path)
except Exception:
self.logger.warning(
f"Preview of {os.path.basename(image_path)} could not be queued",
exc_info=True,
)
return
self.logger.debug(f"Preview {outcome}: {os.path.basename(image_path)}")
def wait_until_settled(self, image_path: str, stop=None) -> bool:
"""True once the file's size is non-zero and unchanged across two reads.
Returns False as soon as ``stop`` (an Event) is set.
"""
stop = stop if stop is not None else self._stop
previous = None
for _ in range(self.settle_attempts):
try:
size = os.path.getsize(image_path)
except OSError:
return False
if previous is not None and size == previous and size > 0:
return True
previous = size
if stop.wait(self.settle_interval):
return False
return False
def is_image_file(self, filepath: str) -> bool:
"""Check if file is a supported image format and not a thumbnail.
Args:
filepath: Path to the file to check
Returns:
True if the file has a supported image extension and is not
inside a thumbnails directory, False otherwise
"""
# Skip files in thumbnails directory - those are derivatives, not
# generated images
if "/thumbnails/" in filepath or "\\thumbnails\\" in filepath:
return False
return filepath.lower().endswith(self.supported_extensions)
def process_new_image(self, image_path: str, prompt_snapshot=None):
"""Process a newly created image file for gallery integration.
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. 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: 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):
self.logger.warning(f"Image file no longer exists: {image_path}")
return
# Extract ComfyUI metadata first — needed both for linking and prompt lookup
metadata = None
try:
metadata = self.metadata_extractor.extract_metadata(image_path)
self.logger.debug(f"Extracted metadata: {bool(metadata)}")
except Exception as meta_error:
self.logger.warning(f"Metadata extraction failed: {meta_error}")
# Strategy 1: The image's own metadata. It records the positive prompt
# that produced this image, independent of node caching and of any
# negative-prompt PromptManager nodes in the workflow.
state, current_prompt = self._find_prompt_from_metadata(metadata)
if state == "foreign":
# Another workflow saved this image: it belongs to no prompt of
# ours, and must not consume a batch queue entry.
self.logger.info(
f"No PromptManager prompt in workflow, skipping image: {image_path}"
)
return
if current_prompt:
self.logger.info(
f"Metadata match: prompt {current_prompt['id']} for "
f"{os.path.basename(image_path)}"
)
else:
current_prompt = self._find_prompt_from_run(
image_path, metadata, prompt_snapshot
)
if not current_prompt:
# 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"):
self.prompt_tracker.extend_prompt_timeout(
current_prompt["execution_id"], 300
)
if metadata:
self.logger.debug(
f"Linking image with full metadata to prompt {current_prompt['id']}"
)
self.link_image_to_prompt(image_path, current_prompt, metadata)
else:
self.logger.debug(
f"Linking image with basic info to prompt {current_prompt['id']}"
)
basic_metadata = self.get_basic_file_info(image_path)
self.link_image_to_prompt(
image_path, current_prompt, {"file_info": basic_metadata}
)
except Exception as e:
self.logger.error(f"Error processing image {image_path}: {e}")
import traceback
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.
Parses the ComfyUI workflow/prompt data embedded in the image to find the
PromptManager node feeding the positive input, then matches its text
against the database by hash.
Args:
metadata: Extracted metadata dict from the image, or None
Returns:
``("linked", prompt)`` with a context dict holding 'id' and 'text';
``("foreign", None)`` when the workflow has no positive PromptManager
node, so the image belongs to no prompt of ours;
``("unknown", None)`` when metadata can't decide (missing, a batch
item whose text came from another node, or a lookup failure) and the
caller's queue and tracker fallbacks apply.
"""
if not metadata:
return "unknown", None
prompt_text = None
resolution = resolve_prompt_pair(metadata)
graph = metadata.get("prompt")
if resolution["state"] == "resolved" and isinstance(graph, dict):
positive = resolution["positive"]
path = positive["provenance"].get("path", [])
paired_source = any(
isinstance(graph.get(item.get("node_id")), dict)
and graph[item["node_id"]].get("class_type") == "PromptManagerPair"
for item in path
)
if paired_source:
try:
existing = self.db_manager.get_prompt_by_hash(
generate_prompt_hash(positive["storage_text"])
)
if existing:
return "linked", {
"id": existing["id"],
"text": existing["text"],
"from_metadata": True,
}
except Exception as error:
self.logger.warning(f"Paired metadata lookup failed: {error}")
# The executed graph: use the PromptManager node feeding the positive input
prompt_data = metadata.get("prompt")
if isinstance(prompt_data, dict):
run_nodes = run_prompt_nodes(prompt_data)
if not run_nodes:
return "foreign", None
for node_id in run_nodes:
prompt_text = resolve_text(prompt_data, node_id)
if prompt_text:
break
if not prompt_text:
# Positive text can't be known from the graph (batch item or unknown
# node): leave it to the queue rather than guess from workflow widgets
return "unknown", None
# Fallback (no executed graph): check text_encoder_nodes from workflow, but
# only if the text input is NOT connected (connected inputs override widget
# values, so the widget value would be stale in batch workflows).
if not prompt_text:
text_nodes = metadata.get("text_encoder_nodes", [])
if text_nodes and not any(
"PromptManager" in (n.get("type") or n.get("class_type") or "")
for n in text_nodes
if isinstance(n, dict)
):
return "foreign", None
for node in text_nodes:
node_type = node.get("type") or node.get("class_type") or ""
if "PromptManager" in node_type:
# Check if text input is connected — if so, widget value is stale
text_connected = False
for inp in node.get("inputs", []):
if isinstance(inp, dict) and inp.get("name") == "text":
if inp.get("link") is not None:
text_connected = True
break
if text_connected:
self.logger.debug(
"PromptManager text input is connected — "
"skipping stale widget value, deferring to snapshot"
)
continue
widgets = node.get("widgets_values", [])
if widgets and isinstance(widgets[0], str) and widgets[0].strip():
prompt_text = widgets[0].strip()
break
if not prompt_text:
return "unknown", None
# Look up by hash in database
try:
existing = self.db_manager.get_prompt_by_hash(
generate_prompt_hash(prompt_text)
)
if existing:
self.logger.debug(
f"Found DB prompt {existing['id']} from metadata text"
)
return "linked", {
"id": existing["id"],
"text": existing["text"],
"from_metadata": True,
}
except Exception as e:
self.logger.warning(f"Metadata-based prompt lookup failed: {e}")
return "unknown", None
def get_basic_file_info(self, image_path: str) -> Dict[str, Any]:
"""Get basic file information when metadata extraction fails.
Provides fallback file information when ComfyUI metadata cannot be extracted
from the image. Includes file size, format, and dimensions when possible.
Args:
image_path: Path to the image file
Returns:
Dictionary containing basic file information:
- size: File size in bytes
- format: Image format (PNG, JPEG, etc.)
- dimensions: Image width and height as list [width, height]
"""
try:
from PIL import Image
stat = os.stat(image_path)
file_info = {"size": stat.st_size, "format": None, "dimensions": None}
# Try to get image dimensions
try:
with Image.open(image_path) as img:
file_info["dimensions"] = list(img.size)
file_info["format"] = img.format
except Exception:
pass
return file_info
except Exception as e:
self.logger.error(f"Error getting file info: {e}")
return {}
def link_image_to_prompt(
self, image_path: str, prompt_context: Dict, metadata: Dict
):
"""Link an image to a prompt in the database.
Creates a database record associating the generated image with its source
prompt, including any extracted metadata from the image file.
Args:
image_path: Full path to the image file
prompt_context: Dictionary containing prompt information including ID
and text
metadata: Extracted metadata from the image file (workflow, parameters,
etc.)
"""
try:
resolution = resolve_prompt_pair(metadata)
def link():
# One write; run again on the moved library when a Move
# library switched while it waited (#249 S6b).
if resolution["state"] == "resolved":
positive = resolution["positive"]
expected = self.db_manager.get_prompt_by_hash(
generate_prompt_hash(
positive.get("storage_text", positive["text"])
)
)
if expected and expected["id"] == prompt_context["id"]:
ingest_resolved_pair(
self.db_manager, image_path, metadata, resolution
)
return None
# The raising form: a failed write reaches survive_switch.
return self.db_manager.link_image_to_prompt_or_raise(
prompt_id=prompt_context["id"],
image_path=image_path,
metadata=metadata,
)
image_id = survive_switch(self.db_manager, link)
if image_id is None:
return
if image_id: # 0: nothing new was linked (logged by the database)
self.logger.debug(
f"Successfully linked image {image_id} to prompt "
f"{prompt_context['id']}"
)
except Exception as e:
self.logger.error(f"Failed to link image to prompt: {e}")
return
if resolution["state"] == "ambiguous":
self._record_review(image_path)
def _record_review(self, image_path: str) -> None:
"""The Needs review pass for an ambiguous output, best effort.
The output is linked by now, so it is never queued; a pending item
an earlier scan left for its path leaves the queue. The link above
stays whatever happens here: a failure is logged, never raised.
"""
recorder = getattr(self, "review_recorder", None)
if recorder is None:
return
try:
outcome = recorder(image_path)
except Exception:
self.logger.warning(
f"{os.path.basename(image_path)} could not be added to Needs review",
exc_info=True,
)
return
self.logger.debug(f"Needs review {outcome}: {os.path.basename(image_path)}")
class ImageMonitor:
"""Main image monitoring system for ComfyUI gallery integration.
This class manages the overall image monitoring system, including:
- Setting up filesystem watchers for output directories
- Auto-detecting ComfyUI output locations
- Managing the lifecycle of monitoring operations
- Providing status information
The monitor uses watchdog to efficiently watch filesystem changes and can
monitor multiple directories simultaneously with recursive subdirectory support.
"""
def __init__(self, db_manager, prompt_tracker):
"""
Initialize the image monitor.
Args:
db_manager: Database manager instance for storing image relationships
prompt_tracker: Prompt tracking instance for getting active prompt context
"""
self.db_manager = db_manager
self.prompt_tracker = prompt_tracker
self.observer = None
self.handler = None
self.monitored_directories = []
self.logger = get_logger("prompt_manager.image_monitor")
self.review_recorder: Optional[Callable[[str], str]] = None
self.thumbnail_queue: Optional[Callable[[str], str]] = None
def set_thumbnail_queue(
self, queue_preview: Optional[Callable[[str], str]]
) -> None:
"""Queue every settled output's preview through ``queue_preview``
(None stops it), on the running handler and every later one."""
self.thumbnail_queue = queue_preview
if self.handler is not None:
self.handler.thumbnail_queue = queue_preview
def set_review_recorder(self, recorder: Optional[Callable[[str], str]]) -> None:
"""Record ambiguous outputs through ``recorder`` (None stops recording).
Applies to the running handler and to every handler started later.
"""
self.review_recorder = recorder
if self.handler is not None:
self.handler.review_recorder = recorder
def start_monitoring(self, output_directories: Optional[list] = None):
"""
Start monitoring ComfyUI output directories for new images.
Begins filesystem watching on the specified directories. If no directories
are provided, the system will first check GalleryConfig.MONITORING_DIRECTORIES,
then fall back to auto-detecting ComfyUI output locations.
All monitoring is done recursively to catch images in subdirectories.
Args:
output_directories: List of directory paths to monitor. If None, uses
config or auto-detection.
"""
if self.observer:
self.logger.warning("Image monitoring already running")
return
# Check config for monitoring settings
config = self._gallery_config()
if config is not None:
if not config.MONITORING_ENABLED:
self.logger.info("Image monitoring disabled in config")
return
# Use configured directories if set
if not output_directories and config.MONITORING_DIRECTORIES:
output_directories = config.MONITORING_DIRECTORIES
self.logger.info(
f"Using configured monitoring directories: {output_directories}"
)
# Auto-detect ComfyUI output directory if still none
if not output_directories:
output_directories = self.detect_comfyui_output_dirs()
valid_directories = []
for output_dir in output_directories or []:
if not os.path.isdir(output_dir):
self.logger.warning(f"Directory does not exist: {output_dir}")
continue
# Same containment rule as the settings endpoint, so a root
# hand-edited into config.json is never watched (and its files
# never linked) when it lies outside ComfyUI's directories.
validate = getattr(config, "validate_gallery_root", None)
if callable(validate):
ok, reason = validate(output_dir)
if not ok:
self.logger.warning(f"Not watching {output_dir}: {reason}")
continue
valid_directories.append(output_dir)
if not valid_directories:
self.logger.warning("No valid directories to monitor")
return
# Only a started observer becomes self.observer, so a failed start never
# leaves the monitor stuck "already running" with nothing to stop.
handler = ImageGenerationHandler(self.db_manager, self.prompt_tracker)
handler.review_recorder = self.review_recorder
handler.thumbnail_queue = self.thumbnail_queue
observer = Observer()
scheduled = []
for output_dir in valid_directories:
try:
observer.schedule(handler, output_dir, recursive=True)
scheduled.append(output_dir)
self.logger.info(f"Monitoring directory (recursive): {output_dir}")
except Exception as e:
self.logger.warning(f"Cannot watch {output_dir}: {e}")
if not scheduled:
self.logger.warning("No directory could be watched")
return
try:
observer.start()
except Exception as e:
self.logger.error(f"Image monitoring could not start: {e}")
return
self.handler = handler
self.observer = observer
self.monitored_directories = scheduled
self.logger.info(f"Image monitoring started for {len(scheduled)} directories")
@staticmethod
def _gallery_config():
"""GalleryConfig when the package's config is importable, else None."""
try:
from ..py.config import GalleryConfig
return GalleryConfig
except Exception:
return None
@property
def is_monitoring(self) -> bool:
"""Whether a filesystem observer is running."""
return self.observer is not None
def stop_monitoring(self):
"""Stop the image monitoring system.
Cleanly shuts down the filesystem watcher and clears all monitoring state.
This method should be called before program exit to ensure proper cleanup.
Safe to call when monitoring never started.
"""
observer, self.observer = self.observer, None
handler, self.handler = self.handler, None
self.monitored_directories = []
if observer is None:
return
try:
observer.stop()
observer.join()
except Exception as e:
self.logger.warning(f"Error while stopping image monitoring: {e}")
if handler is not None:
try:
handler.shutdown()
except Exception as e:
self.logger.warning(f"Error while stopping the image worker: {e}")
self.logger.debug("Image monitoring stopped")
def detect_comfyui_output_dirs(self) -> list:
"""Auto-detect ComfyUI output directories.
Attempts to locate ComfyUI output directories using multiple strategies:
1. Import ComfyUI's folder_paths module to get the configured output directory
2. Search common relative paths where ComfyUI output directories are
typically located
3. Verify that detected directories actually exist
Returns:
List of absolute paths to detected output directories
"""
potential_dirs = []
try:
# Try to import ComfyUI's folder_paths
import folder_paths
output_dir = folder_paths.get_output_directory()
if output_dir and os.path.exists(output_dir):
potential_dirs.append(output_dir)
self.logger.debug(f"Detected ComfyUI output directory: {output_dir}")
except ImportError:
self.logger.debug(
"ComfyUI folder_paths not available, using fallback detection"
)
# Fallback: Look for common ComfyUI directory structures
fallback_paths = [
"output",
"../output",
"../../output",
"ComfyUI/output",
"../ComfyUI/output",
]
for path in fallback_paths:
abs_path = os.path.abspath(path)
if os.path.exists(abs_path) and abs_path not in potential_dirs:
potential_dirs.append(abs_path)
self.logger.debug(f"Found output directory: {abs_path}")
return potential_dirs
def get_status(self) -> Dict[str, Any]:
"""Get monitoring status information.
Returns:
Dictionary containing:
- running: Boolean indicating if monitoring is active
- monitored_directories: List of currently monitored directory paths
- handler_active: Boolean indicating if the event handler is active
- observer_alive: Boolean indicating if observer thread is alive
"""
observer_alive = False
if self.observer is not None:
try:
observer_alive = self.observer.is_alive()
except Exception:
pass
return {
"running": self.observer is not None,
"monitored_directories": self.monitored_directories,
"handler_active": self.handler is not None,
"observer_alive": observer_alive,
}
# Singleton instance management
_monitor_instance: Optional[ImageMonitor] = None
_monitor_lock = threading.Lock()
def get_image_monitor(db_manager, prompt_tracker) -> ImageMonitor:
"""Get or create the singleton ImageMonitor instance.
This ensures only one ImageMonitor exists across all PromptManager nodes,
preventing duplicate image detection when multiple nodes are used.
Args:
db_manager: Database manager instance for storing image relationships
prompt_tracker: Prompt tracking instance for getting active prompt context
Returns:
The singleton ImageMonitor instance
"""
global _monitor_instance
if _monitor_instance is None:
with _monitor_lock:
if _monitor_instance is None:
_monitor_instance = ImageMonitor(db_manager, prompt_tracker)
return _monitor_instance
Reference in New Issue View Git Blame Copy Permalink