Vito
caa9649f32
* feat(db): add pair schema and stage restore migrations for #193 * feat(database): persist reusable ordered prompt pairs (#193) * feat(api): expose reusable pairs and atomic metadata enrichment (#193) * feat(metadata): resolve canonical prompt pairs with bounded provenance (#193) * feat(scan): enrich image pairs without guessing or duplicate links (#193) * fix(database): preserve pair tags during global maintenance (#193) * feat(tracking): enrich live image pairs and cancel scans safely (#193) * feat(nodes): encode paired prompts and count cached recipe runs (#193) * fix(database): validate restored pair ownership and provenance (#193) * feat(admin): report pair scan outcomes and cancellation (#193) * feat(ui): add paired prompt editor and explicit recipe library for #193 * fix(autocomplete): include reusable pair metadata vocabulary (#193) * fix(pairs): honor recent sorting and conservative scan outcomes (#193) * fix(ui): protect paired draft saves and clarify negative variants (#193) * fix(metadata): preserve bounded paired execution evidence (#193) * fix(tracking): follow actual pair output roles and prefer execution evidence (#193) * docs(pairs): document live validation and retain renderer screenshots (#193) * test(api): inject disposable databases before API startup (#194) * test(api): canonicalize portable temporary gallery roots (#194) * docs(ui): prototype branching negative variant history (#193) * docs(ui): show selected negative beneath positive in timeline concept (#193) * docs: retain approved timeline for later gallery redesign (#193) * test(db): close restore fixture connections before cleanup (#194) * test(db): reuse existing handle guards in all pair fixtures (#194) * chore: clarify rollback test and intentional exception handling (#193)
822 lines
34 KiB
Python
822 lines
34 KiB
Python
"""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 active prompts using the prompt tracker
|
|
- Handles fallback linking when no active prompt is available
|
|
- Provides status information and monitoring control
|
|
"""
|
|
|
|
import os
|
|
import queue
|
|
import time
|
|
import threading
|
|
from typing import Optional, Dict, Any, 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
|
|
|
|
try:
|
|
from ..database.pair_ingestion import ingest_resolved_pair
|
|
except ImportError:
|
|
from database.pair_ingestion import ingest_resolved_pair
|
|
|
|
# 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
|
|
|
|
|
|
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")
|
|
|
|
# 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
|
|
|
|
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 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)
|
|
|
|
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._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)
|
|
finally:
|
|
with self._pending_lock:
|
|
self._pending.pop(image_path, None)
|
|
|
|
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. 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
|
|
|
|
Args:
|
|
image_path: Full path to the newly created image file
|
|
prompt_snapshot: Prompt context captured at file-creation time (optional)
|
|
"""
|
|
try:
|
|
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)}"
|
|
)
|
|
|
|
# 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')}"
|
|
)
|
|
|
|
# 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
|
|
|
|
# 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_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 _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
|
|
):
|
|
"""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)
|
|
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
|
|
image_id = self.db_manager.link_image_to_prompt(
|
|
prompt_id=prompt_context["id"], image_path=image_path, metadata=metadata
|
|
)
|
|
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}"
|
|
)
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to link image to prompt: {e}")
|
|
|
|
|
|
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")
|
|
|
|
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)
|
|
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
|