diff --git a/utils/comfyui_integration.py b/utils/comfyui_integration.py index 95a4e7d..cc2b4d2 100644 --- a/utils/comfyui_integration.py +++ b/utils/comfyui_integration.py @@ -185,7 +185,8 @@ class ComfyUIMetadataIntegration: Args: max_age_seconds: Maximum age in seconds before a prompt registration - is considered stale and removed (default: 600 seconds/10 minutes) + is considered stale and removed (default: 600 + seconds/10 minutes) """ current_time = time.time() diff --git a/utils/diagnostics.py b/utils/diagnostics.py index 10146ff..14f1581 100644 --- a/utils/diagnostics.py +++ b/utils/diagnostics.py @@ -33,7 +33,7 @@ The diagnostics provide: import os import sqlite3 from pathlib import Path -from typing import Dict, Any, List +from typing import Dict, Any from .logging_config import get_logger @@ -157,7 +157,7 @@ class GalleryDiagnostics: # Check if generated_images table exists cursor = conn.execute(""" - SELECT name FROM sqlite_master + SELECT name FROM sqlite_master WHERE type='table' AND name='generated_images' """) @@ -194,14 +194,17 @@ class GalleryDiagnostics: # Check if table exists cursor = conn.execute(""" - SELECT name FROM sqlite_master + SELECT name FROM sqlite_master WHERE type='table' AND name='generated_images' """) if not cursor.fetchone(): return { "status": "error", - "message": "generated_images table does not exist - run the updated code to create it", + "message": ( + "generated_images table does not exist - " + "run the updated code to create it" + ), } # Check image records @@ -211,10 +214,10 @@ class GalleryDiagnostics: # Get recent images cursor = conn.execute(""" - SELECT gi.*, p.text - FROM generated_images gi - LEFT JOIN prompts p ON gi.prompt_id = p.id - ORDER BY gi.generation_time DESC + SELECT gi.*, p.text + FROM generated_images gi + LEFT JOIN prompts p ON gi.prompt_id = p.id + ORDER BY gi.generation_time DESC LIMIT 5 """) recent_images = [dict(row) for row in cursor.fetchall()] @@ -368,10 +371,10 @@ class GalleryDiagnostics: # Check PIL try: - from PIL import Image + from PIL import Image # noqa: F401 - availability probe dependencies["PIL"] = True - self.logger.info(f" [PASS] PIL (Pillow): Available") + self.logger.info(" [PASS] PIL (Pillow): Available") except ImportError: self.logger.error(" [FAIL] PIL (Pillow): NOT AVAILABLE") diff --git a/utils/image_monitor.py b/utils/image_monitor.py index 51885ed..ae38baa 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -1,9 +1,9 @@ """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. +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 @@ -211,7 +211,8 @@ class ImageGenerationHandler(FileSystemEventHandler): 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 + # 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) @@ -363,9 +364,9 @@ class ImageGenerationHandler(FileSystemEventHandler): # 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). + # 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( @@ -485,13 +486,15 @@ class ImageGenerationHandler(FileSystemEventHandler): ): """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. + 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.) + prompt_context: Dictionary containing prompt information including ID + and text + metadata: Extracted metadata from the image file (workflow, parameters, + etc.) """ try: image_id = self.db_manager.link_image_to_prompt( @@ -499,7 +502,8 @@ class ImageGenerationHandler(FileSystemEventHandler): ) fallback_note = " (fallback)" if prompt_context.get("fallback") else "" self.logger.debug( - f"Successfully linked image {image_id} to prompt {prompt_context['id']}{fallback_note}" + 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}") @@ -543,7 +547,8 @@ class ImageMonitor: 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. + 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") @@ -641,7 +646,8 @@ class ImageMonitor: 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 + 2. Search common relative paths where ComfyUI output directories are + typically located 3. Verify that detected directories actually exist Returns: diff --git a/utils/metadata_extractor.py b/utils/metadata_extractor.py index d24f1f4..c6e8541 100644 --- a/utils/metadata_extractor.py +++ b/utils/metadata_extractor.py @@ -1,8 +1,9 @@ """ComfyUI metadata extraction utilities. -This module provides comprehensive metadata extraction capabilities for ComfyUI-generated -images. It can extract workflow information, prompt data, and generation parameters from -PNG images that contain embedded ComfyUI metadata in their text chunks. +This module provides comprehensive metadata extraction capabilities for +ComfyUI-generated images. It can extract workflow information, prompt data, and +generation parameters from PNG images that contain embedded ComfyUI metadata in +their text chunks. The extractor supports: - Complete workflow data extraction from PNG text chunks @@ -34,7 +35,6 @@ import os import json from typing import Optional, Dict, Any from PIL import Image -from PIL.PngImagePlugin import PngInfo from .logging_config import get_logger diff --git a/utils/prompt_graph.py b/utils/prompt_graph.py index 52423de..a22c4a7 100644 --- a/utils/prompt_graph.py +++ b/utils/prompt_graph.py @@ -134,7 +134,8 @@ def run_prompt_nodes(graph: Any) -> List[str]: def _join_text_inputs( inputs: Dict[str, Any], keys: List[str], resolve ) -> Optional[str]: - """WAS-style join: resolve each key, optionally strip, skip empty, join by delimiter.""" + """WAS-style join: resolve each key, optionally strip, skip empty, join by + delimiter.""" delimiter = inputs.get("delimiter", " ") if not isinstance(delimiter, str): return None diff --git a/utils/prompt_tracker.py b/utils/prompt_tracker.py index cd65109..bc3ab35 100644 --- a/utils/prompt_tracker.py +++ b/utils/prompt_tracker.py @@ -1,9 +1,9 @@ """Prompt tracking system for linking generated images to prompts. -This module provides thread-safe tracking of active prompt executions to enable automatic -association of generated images with their source prompts. The system maintains prompt -context across different execution threads and provides fallback mechanisms for reliable -image-prompt linking. +This module provides thread-safe tracking of active prompt executions to enable +automatic association of generated images with their source prompts. The system +maintains prompt context across different execution threads and provides fallback +mechanisms for reliable image-prompt linking. Key features: - Thread-safe prompt tracking using threading.local and locks @@ -115,7 +115,8 @@ class PromptTracker: Args: prompt_text: The prompt text being executed - additional_data: Additional prompt metadata, should include prompt_id from PromptManager + additional_data: Additional prompt metadata, should include prompt_id + from PromptManager 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. @@ -140,7 +141,8 @@ class PromptTracker: prompt_id = existing_prompt["id"] self.logger.debug(f"Found existing prompt ID: {prompt_id}") else: - # Generate a temporary ID for tracking (prompt should be saved by PromptManager) + # Generate a temporary ID for tracking (prompt should be saved + # by PromptManager) prompt_id = f"temp_{int(time.time())}" self.logger.debug(f"Using temporary ID for tracking: {prompt_id}") except Exception as e: @@ -169,11 +171,13 @@ class PromptTracker: with self._queue_lock: self._prompt_queue.append(execution_context) self.logger.debug( - f"Queue size: {len(self._prompt_queue)} after push for prompt {prompt_id}" + f"Queue size: {len(self._prompt_queue)} after push " + f"for prompt {prompt_id}" ) self.logger.debug( - f"Set current prompt: {execution_id} -> {prompt_text[:50]}... (thread: {threading.current_thread().ident})" + 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)}") return execution_id @@ -211,7 +215,8 @@ class PromptTracker: recent_prompt = self._find_recent_prompt() if recent_prompt: self.logger.debug( - f"Using recent prompt from global tracking: {recent_prompt['execution_id']}" + "Using recent prompt from global tracking: " + f"{recent_prompt['execution_id']}" ) return recent_prompt @@ -259,14 +264,16 @@ class PromptTracker: with self.lock: current_time = time.time() self.logger.debug( - f"Searching for recent prompt among {len(self.active_prompts)} active prompts" + f"Searching for recent prompt among {len(self.active_prompts)} " + "active prompts" ) recent_prompts = [] for exec_id, prompt in self.active_prompts.items(): age_seconds = current_time - prompt["timestamp"] self.logger.debug( - f"Prompt {exec_id}: age={age_seconds:.1f}s, timeout={self.prompt_timeout}s" + f"Prompt {exec_id}: age={age_seconds:.1f}s, " + f"timeout={self.prompt_timeout}s" ) if age_seconds < self.prompt_timeout: @@ -280,7 +287,7 @@ class PromptTracker: self.logger.debug(f"Found recent prompt: {most_recent['execution_id']}") return most_recent else: - self.logger.debug(f"No recent prompts found") + self.logger.debug("No recent prompts found") return None @@ -295,7 +302,8 @@ class PromptTracker: if current: execution_id = current["execution_id"] self.logger.debug( - f"Clearing current prompt: {execution_id} (thread: {threading.current_thread().ident})" + f"Clearing current prompt: {execution_id} " + f"(thread: {threading.current_thread().ident})" ) # Clear from thread-local storage @@ -317,7 +325,8 @@ class PromptTracker: ) else: self.logger.debug( - f"No current prompt to clear (thread: {threading.current_thread().ident})" + "No current prompt to clear " + f"(thread: {threading.current_thread().ident})" ) def extend_prompt_timeout(self, execution_id: str, additional_seconds: int = 60): @@ -431,7 +440,8 @@ class PromptTracker: Dictionary containing: - active_prompts_count: Number of currently active prompts - current_prompt_id: ID of current prompt in this thread (if any) - - current_execution_id: Execution ID of current prompt in this thread (if any) + - current_execution_id: Execution ID of current prompt in this thread + (if any) - thread_id: Current thread identifier - prompt_timeout: Configured prompt timeout in seconds - cleanup_interval: Configured cleanup interval in seconds