From 22593e8514fb081d073c806a9c957922cd8d7644 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Tue, 12 Aug 2025 06:38:36 -0700 Subject: [PATCH] docs: add comprehensive docstrings to all Python modules - Add Google-style docstrings to main node files (prompt_manager.py, prompt_manager_text.py) - Enhance database module docstrings with detailed parameter and return documentation - Add comprehensive docstrings to all utils modules with usage examples - Document API endpoints and configuration classes thoroughly - Update module-level docstrings in all __init__.py files - Ensure consistent naming (ComfyUI_PromptManager) across all documentation - Follow Python best practices for docstring formatting --- .claude/settings.local.json | 15 ++ __init__.py | 20 +- database/__init__.py | 7 +- database/models.py | 68 ++++- database/operations.py | 4 +- prompt_manager.py | 91 ++++++- prompt_manager_text.py | 91 ++++++- py/__init__.py | 7 +- py/api.py | 507 ++++++++++++++++++++++++++++++++--- py/config.py | 215 ++++++++++++++- restart_gallery.py | 20 +- tests/__init__.py | 7 +- utils/__init__.py | 7 +- utils/comfyui_integration.py | 106 +++++++- utils/diagnostics.py | 160 ++++++++++- utils/hashing.py | 80 +++++- utils/image_monitor.py | 164 +++++++++-- utils/logging_config.py | 145 ++++++++-- utils/metadata_extractor.py | 130 +++++++-- utils/prompt_tracker.py | 168 ++++++++++-- utils/validators.py | 110 ++++++-- 21 files changed, 1900 insertions(+), 222 deletions(-) create mode 100644 .claude/settings.local.json diff --git a/.claude/settings.local.json b/.claude/settings.local.json new file mode 100644 index 0000000..4f761ba --- /dev/null +++ b/.claude/settings.local.json @@ -0,0 +1,15 @@ +{ + "permissions": { + "allow": [ + "Bash(python:*)", + "Bash(pip install:*)", + "Bash(grep:*)", + "Bash(rg:*)", + "WebFetch(domain:github.com)", + "Bash(rm:*)", + "Bash(git add:*)", + "Bash(cp:*)" + ], + "deny": [] + } +} \ No newline at end of file diff --git a/__init__.py b/__init__.py index a02c917..1ca9962 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,22 @@ """ -PromptManager: A ComfyUI custom node that extends the standard text encoder -with persistent prompt storage and advanced search capabilities using SQLite. +ComfyUI_PromptManager: A comprehensive ComfyUI custom node that extends text encoding +with persistent prompt storage, advanced search capabilities, and an automatic image +gallery system. + +This module provides two main node types: +- PromptManager: CLIP encoding node that outputs CONDITIONING +- PromptManagerText: Text-only node that outputs STRING with prepend/append functionality + +Both nodes share the same database backend for persistent prompt storage and include +an automatic image gallery system that monitors ComfyUI output directories and links +generated images to their source prompts. + +Features: +- SQLite-based persistent prompt storage with deduplication +- Advanced search and categorization system +- Real-time image gallery with metadata extraction +- Web-based admin dashboard for prompt management +- Comprehensive logging and diagnostics """ import re diff --git a/database/__init__.py b/database/__init__.py index deaab78..f27ad19 100644 --- a/database/__init__.py +++ b/database/__init__.py @@ -1,4 +1,9 @@ -"""Database module for KikoTextEncode prompt storage and management.""" +""" +Database module for ComfyUI_PromptManager prompt storage and management. + +This package provides SQLite-based persistent storage for prompts with advanced +search capabilities, metadata management, and image tracking functionality. +""" from .operations import PromptDatabase from .models import PromptModel diff --git a/database/models.py b/database/models.py index aad9bd4..71946ba 100644 --- a/database/models.py +++ b/database/models.py @@ -32,7 +32,15 @@ class PromptModel: self._ensure_database_exists() def _ensure_database_exists(self) -> None: - """Create database and tables if they don't exist.""" + """ + Create database and tables if they don't exist. + + Sets up the database schema including tables for prompts and generated images, + creates necessary indexes, and applies any pending migrations. + + Raises: + Exception: If database creation fails + """ try: with sqlite3.connect(self.db_path) as conn: conn.execute("PRAGMA foreign_keys = ON") @@ -44,7 +52,16 @@ class PromptModel: raise def _create_tables(self, conn: sqlite3.Connection) -> None: - """Create the prompts table with all required columns.""" + """ + Create the prompts and generated_images tables with all required columns. + + Args: + conn: Active database connection + + Creates: + - prompts table: Stores prompt text and metadata + - generated_images table: Links generated images to their source prompts + """ conn.execute(""" CREATE TABLE IF NOT EXISTS prompts ( id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -85,7 +102,18 @@ class PromptModel: self._migrate_foreign_key_types(conn) def _create_indexes(self, conn: sqlite3.Connection) -> None: - """Create indexes for better query performance.""" + """ + Create indexes for better query performance. + + Args: + conn: Active database connection + + Creates indexes on: + - Text content for search operations + - Categories and tags for filtering + - Timestamps for sorting + - Hash values for duplicate detection + """ indexes = [ "CREATE INDEX IF NOT EXISTS idx_prompts_text ON prompts(text)", "CREATE INDEX IF NOT EXISTS idx_prompts_category ON prompts(category)", @@ -113,7 +141,15 @@ class PromptModel: return conn def _migrate_workflow_name_removal(self, conn: sqlite3.Connection) -> None: - """Remove workflow_name column if it exists in existing database.""" + """ + Remove workflow_name column if it exists in existing database. + + Args: + conn: Active database connection + + This migration handles legacy schema updates by removing the deprecated + workflow_name column while preserving all other data. + """ try: # Check if workflow_name column exists cursor = conn.execute("PRAGMA table_info(prompts)") @@ -155,7 +191,15 @@ class PromptModel: # If migration fails, the table creation will handle it def _migrate_foreign_key_types(self, conn: sqlite3.Connection) -> None: - """Fix foreign key data type mismatch in generated_images table.""" + """ + Fix foreign key data type mismatch in generated_images table. + + Args: + conn: Active database connection + + Converts prompt_id from TEXT to INTEGER type to match the prompts table's + primary key type, ensuring referential integrity. + """ try: # Check if generated_images table exists and has TEXT prompt_id cursor = conn.execute("PRAGMA table_info(generated_images)") @@ -206,12 +250,22 @@ class PromptModel: # If migration fails, continue with existing schema def migrate_database(self) -> None: - """Apply any pending database migrations.""" + """ + Apply any pending database migrations. + + This method serves as an entry point for future schema migrations. + Add new migration logic here as the database evolves. + """ # Future migrations can be added here pass def vacuum_database(self) -> None: - """Optimize database by running VACUUM command.""" + """ + Optimize database by running VACUUM command. + + Reclaims unused space and defragments the database file, + improving query performance and reducing file size. + """ try: with sqlite3.connect(self.db_path) as conn: conn.execute("VACUUM") diff --git a/database/operations.py b/database/operations.py index cf0eec6..39ed828 100644 --- a/database/operations.py +++ b/database/operations.py @@ -796,12 +796,12 @@ class PromptDatabase: return removed_count - def _clean_nan_values(self, obj): + def _clean_nan_values(self, obj: Any) -> Any: """ Recursively clean NaN values from nested data structures. Args: - obj: The object to clean + obj: The object to clean (dict, list, or scalar) Returns: Cleaned object with NaN values replaced by None diff --git a/prompt_manager.py b/prompt_manager.py index 6c3000c..676661c 100644 --- a/prompt_manager.py +++ b/prompt_manager.py @@ -316,13 +316,29 @@ class PromptManager(ComfyNodeABC): return None def _generate_hash(self, text: str) -> str: - """Generate SHA256 hash for the prompt text.""" + """ + Generate SHA256 hash for the prompt text. + + Args: + text: The prompt text to hash + + Returns: + Hexadecimal string representation of the SHA256 hash + """ # Normalize text for consistent hashing (strip whitespace, normalize case) normalized_text = text.strip().lower() return hashlib.sha256(normalized_text.encode("utf-8")).hexdigest() def _parse_tags(self, tags_string: str) -> Optional[list]: - """Parse comma-separated tags string into a list.""" + """ + Parse comma-separated tags string into a list. + + Args: + tags_string: Comma-separated string of tags + + Returns: + List of parsed tags, or None if no valid tags found + """ if not tags_string or not tags_string.strip(): return None @@ -330,7 +346,15 @@ class PromptManager(ComfyNodeABC): return tags if tags else None def _search_prompts(self, search_text: str = "") -> List[Dict[str, Any]]: - """Search for past prompts by text content.""" + """ + Search for past prompts by text content. + + Args: + search_text: Text to search for in prompt database + + Returns: + List of matching prompt dictionaries with metadata + """ try: if not search_text or not search_text.strip(): return [] @@ -350,7 +374,12 @@ class PromptManager(ComfyNodeABC): return [] def _open_web_interface(self): - """Open the web interface in the default browser.""" + """ + Open the web interface in the default browser. + + Attempts to locate and open the web interface HTML file. + Logs warnings if the interface is not properly configured. + """ try: # Look for a web interface directory current_dir = os.path.dirname(os.path.abspath(__file__)) @@ -375,11 +404,27 @@ class PromptManager(ComfyNodeABC): self.logger.error(f"Error opening web interface: {e}") def search_prompts_api(self, search_text: str = "") -> List[Dict[str, Any]]: - """API method for JavaScript UI to search prompts.""" + """ + API method for JavaScript UI to search prompts. + + Args: + search_text: Text to search for in prompts + + Returns: + List of matching prompt dictionaries + """ return self._search_prompts(search_text=search_text) def get_recent_prompts_api(self, limit: int = 20) -> List[Dict[str, Any]]: - """API method for JavaScript UI to get recent prompts.""" + """ + API method for JavaScript UI to get recent prompts. + + Args: + limit: Maximum number of recent prompts to retrieve + + Returns: + List of recent prompt dictionaries ordered by creation time + """ try: return self.db.get_recent_prompts(limit=limit) except Exception as e: @@ -387,7 +432,12 @@ class PromptManager(ComfyNodeABC): return [] def _start_gallery_system(self): - """Initialize and start the gallery monitoring system.""" + """ + Initialize and start the gallery monitoring system. + + Starts the image monitor which watches for new generated images + and links them to their source prompts in the database. + """ try: self.logger.debug("Starting gallery system...") @@ -401,14 +451,24 @@ class PromptManager(ComfyNodeABC): self.logger.warning("Gallery features will be disabled") def get_gallery_status(self) -> Dict[str, Any]: - """Get status of the gallery system.""" + """ + Get status of the gallery system. + + Returns: + Dictionary containing status information for image monitor and prompt tracker + """ return { "image_monitor": self.image_monitor.get_status(), "prompt_tracker": self.prompt_tracker.get_status(), } def cleanup_gallery_system(self): - """Clean up gallery system resources.""" + """ + Clean up gallery system resources. + + Stops image monitoring and releases associated resources. + Called automatically during object destruction. + """ try: if hasattr(self, "image_monitor"): self.image_monitor.stop_monitoring() @@ -417,10 +477,19 @@ class PromptManager(ComfyNodeABC): self.logger.error(f"Error cleaning up gallery system: {e}") def __del__(self): - """Cleanup when object is destroyed.""" + """ + Cleanup when object is destroyed. + + Ensures proper resource cleanup by stopping the gallery system. + """ self.cleanup_gallery_system() @classmethod def IS_CHANGED(cls, **kwargs): - """Always process to ensure database saving.""" + """ + ComfyUI method to determine if node needs re-execution. + + Returns: + NaN to force re-execution, ensuring prompts are always saved to database + """ return float("NaN") # Always execute diff --git a/prompt_manager_text.py b/prompt_manager_text.py index d75125e..ae31a4e 100644 --- a/prompt_manager_text.py +++ b/prompt_manager_text.py @@ -287,13 +287,29 @@ class PromptManagerText(ComfyNodeABC): return None def _generate_hash(self, text: str) -> str: - """Generate SHA256 hash for the prompt text.""" + """ + Generate SHA256 hash for the prompt text. + + Args: + text: The prompt text to hash + + Returns: + Hexadecimal string representation of the SHA256 hash + """ # Normalize text for consistent hashing (strip whitespace, normalize case) normalized_text = text.strip().lower() return hashlib.sha256(normalized_text.encode("utf-8")).hexdigest() def _parse_tags(self, tags_string: str) -> Optional[list]: - """Parse comma-separated tags string into a list.""" + """ + Parse comma-separated tags string into a list. + + Args: + tags_string: Comma-separated string of tags + + Returns: + List of parsed tags, or None if no valid tags found + """ if not tags_string or not tags_string.strip(): return None @@ -301,7 +317,15 @@ class PromptManagerText(ComfyNodeABC): return tags if tags else None def _search_prompts(self, search_text: str = "") -> List[Dict[str, Any]]: - """Search for past prompts by text content.""" + """ + Search for past prompts by text content. + + Args: + search_text: Text to search for in prompt database + + Returns: + List of matching prompt dictionaries with metadata + """ try: if not search_text or not search_text.strip(): return [] @@ -321,7 +345,12 @@ class PromptManagerText(ComfyNodeABC): return [] def _open_web_interface(self): - """Open the web interface in the default browser.""" + """ + Open the web interface in the default browser. + + Attempts to locate and open the web interface HTML file. + Logs warnings if the interface is not properly configured. + """ try: # Look for a web interface directory current_dir = os.path.dirname(os.path.abspath(__file__)) @@ -346,11 +375,27 @@ class PromptManagerText(ComfyNodeABC): self.logger.error(f"Error opening web interface: {e}") def search_prompts_api(self, search_text: str = "") -> List[Dict[str, Any]]: - """API method for JavaScript UI to search prompts.""" + """ + API method for JavaScript UI to search prompts. + + Args: + search_text: Text to search for in prompts + + Returns: + List of matching prompt dictionaries + """ return self._search_prompts(search_text=search_text) def get_recent_prompts_api(self, limit: int = 20) -> List[Dict[str, Any]]: - """API method for JavaScript UI to get recent prompts.""" + """ + API method for JavaScript UI to get recent prompts. + + Args: + limit: Maximum number of recent prompts to retrieve + + Returns: + List of recent prompt dictionaries ordered by creation time + """ try: return self.db.get_recent_prompts(limit=limit) except Exception as e: @@ -358,7 +403,12 @@ class PromptManagerText(ComfyNodeABC): return [] def _start_gallery_system(self): - """Initialize and start the gallery monitoring system.""" + """ + Initialize and start the gallery monitoring system. + + Starts the image monitor which watches for new generated images + and links them to their source prompts in the database. + """ try: self.logger.debug("Starting gallery system...") @@ -372,14 +422,24 @@ class PromptManagerText(ComfyNodeABC): self.logger.warning("Gallery features will be disabled") def get_gallery_status(self) -> Dict[str, Any]: - """Get status of the gallery system.""" + """ + Get status of the gallery system. + + Returns: + Dictionary containing status information for image monitor and prompt tracker + """ return { "image_monitor": self.image_monitor.get_status(), "prompt_tracker": self.prompt_tracker.get_status(), } def cleanup_gallery_system(self): - """Clean up gallery system resources.""" + """ + Clean up gallery system resources. + + Stops image monitoring and releases associated resources. + Called automatically during object destruction. + """ try: if hasattr(self, "image_monitor"): self.image_monitor.stop_monitoring() @@ -388,10 +448,19 @@ class PromptManagerText(ComfyNodeABC): self.logger.error(f"Error cleaning up gallery system: {e}") def __del__(self): - """Cleanup when object is destroyed.""" + """ + Cleanup when object is destroyed. + + Ensures proper resource cleanup by stopping the gallery system. + """ self.cleanup_gallery_system() @classmethod def IS_CHANGED(cls, **kwargs): - """Always process to ensure database saving.""" + """ + ComfyUI method to determine if node needs re-execution. + + Returns: + NaN to force re-execution, ensuring prompts are always saved to database + """ return float("NaN") # Always execute diff --git a/py/__init__.py b/py/__init__.py index e864c08..d8a0943 100644 --- a/py/__init__.py +++ b/py/__init__.py @@ -1 +1,6 @@ -# KikoTextEncode Python API modules \ No newline at end of file +""" +ComfyUI_PromptManager Python API modules. + +This package contains the core API components for the web interface and configuration +management, including REST endpoints, configuration handling, and server integration. +""" \ No newline at end of file diff --git a/py/api.py b/py/api.py index 74ed278..100b1d1 100644 --- a/py/api.py +++ b/py/api.py @@ -1,3 +1,31 @@ +"""REST API module for ComfyUI PromptManager. + +This module provides a comprehensive REST API for managing prompts, images, and +gallery functionality within ComfyUI. The API handles CRUD operations for prompts, +image gallery management, system administration, logging, and real-time image +monitoring with metadata extraction. + +Key Features: +- Prompt management (create, read, update, delete, search) +- Image gallery with automatic ComfyUI output monitoring +- Bulk operations for efficiency +- Database maintenance and optimization +- System diagnostics and logging +- Thumbnail generation and management +- Metadata extraction from PNG files +- Real-time progress tracking + +The API integrates with ComfyUI's aiohttp server and provides endpoints for +both programmatic access and web UI functionality. + +Classes: + PromptManagerAPI: Main API class handling all REST endpoints + +Example: + api = PromptManagerAPI() + api.add_routes(server.routes) +""" + # PromptManager/py/api.py import datetime @@ -25,9 +53,36 @@ except ImportError: class PromptManagerAPI: - """API class for PromptManager database operations.""" + """REST API handler for PromptManager operations and web interface. + + This class provides comprehensive REST API endpoints for managing prompts, + images, and system operations. It handles database interactions, file + operations, image processing, and web UI serving. + + The API is designed to integrate seamlessly with ComfyUI's aiohttp server + and provides both JSON API endpoints and static file serving for the + web interface. + + Attributes: + logger: Configured logger instance for API operations + db (PromptDatabase): Database connection and operations handler + + Example: + api = PromptManagerAPI() + api.add_routes(server_routes) + # API endpoints are now available at /prompt_manager/* + """ def __init__(self): + """Initialize the PromptManager API with database connection and cleanup. + + Sets up logging, initializes the database connection, and performs + startup cleanup to remove any duplicate prompts that may exist. + + Raises: + Exception: If database initialization fails, logs error but continues. + If cleanup fails, logs error but continues operation. + """ self.logger = get_logger('prompt_manager.api') self.logger.info("Initializing PromptManager API") @@ -44,7 +99,28 @@ class PromptManagerAPI: self.logger.info("PromptManager API initialization completed") def add_routes(self, routes): - """Add API routes to ComfyUI server using decorator pattern.""" + """Add API routes to ComfyUI server using decorator pattern. + + Registers all API endpoints with the provided aiohttp routes object. + Uses decorator pattern to define routes inline with their handlers. + + Categories of routes registered: + - Core prompt operations (search, save, delete, categories, tags) + - Database maintenance (cleanup, duplicates, maintenance) + - Web UI serving (admin, gallery, metadata viewer) + - Image operations (gallery, thumbnails, metadata extraction) + - System operations (diagnostics, logging, statistics) + - Bulk operations (delete, tag management, export) + + Args: + routes: aiohttp RouteTableDef object from ComfyUI server instance. + Routes will be registered with this object for URL handling. + + Example: + from server import PromptServer + api = PromptManagerAPI() + api.add_routes(PromptServer.instance.routes) + """ # Test route to verify registration works @routes.get("/prompt_manager/test") @@ -384,9 +460,34 @@ class PromptManagerAPI: self.logger.info("All routes registered with decorator pattern") async def search_prompts(self, request): - """ - Search prompts endpoint. - GET /prompt_manager/search?text=...&category=...&tags=...&min_rating=...&limit=... + """Search for prompts using multiple filter criteria. + + Provides comprehensive search functionality across prompt text, categories, + tags, and ratings with configurable result limits. + + Query Parameters: + text (str, optional): Search text to match against prompt content + category (str, optional): Filter by specific category + tags (str, optional): Comma-separated list of tags to filter by + min_rating (int, optional): Minimum rating (1-5) to include + limit (int, optional): Maximum results to return (default: 50, max: 1000) + + Args: + request (aiohttp.web.Request): HTTP request object containing query parameters + + Returns: + aiohttp.web.Response: JSON response with structure: + { + "success": bool, + "results": List[Dict], # List of matching prompt objects + "count": int # Number of results returned + } + + Raises: + Returns 500 status with error details if search fails + + Example: + GET /prompt_manager/search?text=portrait&category=photography&min_rating=3 """ try: # Get query parameters @@ -428,9 +529,40 @@ class PromptManagerAPI: ) async def get_recent_prompts(self, request): - """ - Get recent prompts endpoint with pagination support. - GET /prompt_manager/recent?limit=50&page=2&offset=100 + """Retrieve recently created prompts with pagination support. + + Returns prompts sorted by creation date (newest first) with configurable + pagination using either page-based or offset-based navigation. + + Query Parameters: + limit (int, optional): Number of prompts per page (default: 50, max: 1000) + page (int, optional): Page number for pagination (1-based, default: 1) + offset (int, optional): Offset for results (takes precedence over page) + + Args: + request (aiohttp.web.Request): HTTP request object containing query parameters + + Returns: + aiohttp.web.Response: JSON response with structure: + { + "success": bool, + "results": List[Dict], # List of prompt objects + "pagination": { + "total": int, # Total number of prompts + "limit": int, # Items per page + "offset": int, # Current offset + "page": int, # Current page number + "total_pages": int, # Total number of pages + "has_more": bool, # Whether more pages exist + "count": int # Number of items in current page + } + } + + Raises: + Returns 500 status with error details if retrieval fails + + Example: + GET /prompt_manager/recent?limit=20&page=2 """ try: limit = int(request.query.get("limit", 50)) @@ -476,9 +608,25 @@ class PromptManagerAPI: ) async def get_categories(self, request): - """ - Get all categories endpoint. - GET /prompt_manager/categories + """Retrieve all available prompt categories. + + Returns a list of all unique categories found in the database. + + Args: + request (aiohttp.web.Request): HTTP request object + + Returns: + aiohttp.web.Response: JSON response with structure: + { + "success": bool, + "categories": List[str] # List of category names + } + + Raises: + Returns 500 status with error details if retrieval fails + + Example: + GET /prompt_manager/categories """ try: categories = self.db.get_all_categories() @@ -497,9 +645,25 @@ class PromptManagerAPI: ) async def get_tags(self, request): - """ - Get all tags endpoint. - GET /prompt_manager/tags + """Retrieve all available prompt tags. + + Returns a list of all unique tags found across all prompts in the database. + + Args: + request (aiohttp.web.Request): HTTP request object + + Returns: + aiohttp.web.Response: JSON response with structure: + { + "success": bool, + "tags": List[str] # List of unique tag names + } + + Raises: + Returns 500 status with error details if retrieval fails + + Example: + GET /prompt_manager/tags """ try: tags = self.db.get_all_tags() @@ -518,10 +682,44 @@ class PromptManagerAPI: ) async def save_prompt(self, request): - """ - Save a new prompt endpoint. - POST /prompt_manager/save - Body: {"text": "...", "category": "...", "tags": [...], "rating": 5, "notes": "..."} + """Save a new prompt with metadata and duplicate detection. + + Creates a new prompt record with automatic duplicate detection based on + SHA256 hash of the prompt text. If a duplicate is found, updates the + existing record's metadata instead of creating a new one. + + Request Body (JSON): + text (str, required): The prompt text content + category (str, optional): Category for organization + tags (List[str], optional): List of tags for classification + rating (int, optional): Rating from 1-5 + notes (str, optional): Additional notes or description + + Args: + request (aiohttp.web.Request): HTTP request with JSON body + + Returns: + aiohttp.web.Response: JSON response with structure: + { + "success": bool, + "prompt_id": int, # ID of created/updated prompt + "message": str, # Success/status message + "is_duplicate": bool # True if prompt already existed + } + + Raises: + Returns 400 status if required fields are missing + Returns 500 status with error details if save fails + + Example: + POST /prompt_manager/save + { + "text": "A beautiful sunset over mountains", + "category": "landscape", + "tags": ["nature", "scenic"], + "rating": 4, + "notes": "Great for wallpapers" + } """ try: data = await request.json() @@ -589,9 +787,30 @@ class PromptManagerAPI: ) async def delete_prompt(self, request): - """ - Delete a prompt endpoint. - DELETE /prompt_manager/delete/{prompt_id} + """Delete a specific prompt by ID. + + Permanently removes a prompt and all associated metadata from the database. + Associated image links may also be removed depending on database configuration. + + URL Parameters: + prompt_id (int): The unique identifier of the prompt to delete + + Args: + request (aiohttp.web.Request): HTTP request with prompt_id in URL path + + Returns: + aiohttp.web.Response: JSON response with structure: + { + "success": bool, + "message": str # Success or error message + } + + Raises: + Returns 404 status if prompt not found + Returns 500 status with error details if deletion fails + + Example: + DELETE /prompt_manager/delete/123 """ try: prompt_id = int(request.match_info["prompt_id"]) @@ -669,13 +888,44 @@ class PromptManagerAPI: ) async def find_duplicate_images(self): - """ - Find duplicate images in the output folder based on file content hash. + """Find duplicate images in ComfyUI output directory using content hashing. + + Scans the ComfyUI output directory for image and video files, calculates + SHA256 hashes of file contents, and identifies groups of files with + identical content. Supports both images and videos with thumbnail detection. + + The method processes files efficiently, logging progress every 100 files, + and handles various media formats including PNG, JPG, JPEG, WebP, GIF, + and common video formats. Returns: - List of duplicate groups, each containing: - - hash: The file content hash - - images: List of image records with same content + List[Dict[str, Any]]: List of duplicate groups, where each group contains: + - hash (str): SHA256 hash of the file content + - images (List[Dict]): List of file info dictionaries with: + - id (str): Unique identifier based on file path hash + - filename (str): Original filename + - path (str): Absolute file path + - relative_path (str): Path relative to output directory + - url (str): URL for serving the file + - thumbnail_url (str, optional): URL for thumbnail if available + - size (int): File size in bytes + - modified_time (float): Last modification timestamp + - extension (str): File extension + - media_type (str): 'image' or 'video' + - is_video (bool): True if file is a video + - hash (str): SHA256 content hash + - count (int): Number of duplicate files in this group + + Note: + Files within each duplicate group are sorted by modification time + (oldest first) to help users decide which files to keep. + + Example: + duplicates = await api.find_duplicate_images() + for group in duplicates: + print(f"Found {group['count']} duplicates with hash {group['hash']}") + for img in group['images']: + print(f" - {img['filename']} ({img['size']} bytes)") """ import hashlib from pathlib import Path @@ -790,7 +1040,21 @@ class PromptManagerAPI: return [] def _calculate_file_hash(self, file_path): - """Calculate SHA-256 hash of a file.""" + """Calculate SHA-256 hash of a file's content. + + Reads the file in 4KB chunks to efficiently handle large files + without loading the entire content into memory. + + Args: + file_path (str): Path to the file to hash + + Returns: + str: Hexadecimal SHA-256 hash of the file content + + Raises: + IOError: If the file cannot be read + OSError: If the file path is invalid or inaccessible + """ import hashlib hash_sha256 = hashlib.sha256() @@ -1732,14 +1996,47 @@ class PromptManagerAPI: return web.json_response({'error': str(e)}, status=500) async def generate_thumbnails(self, request): - """ - Generate thumbnails for gallery images. + """Generate thumbnails for all images and videos in the ComfyUI output directory. - SAFETY GUARANTEE: This function NEVER modifies original images. - - Only READS from original image files (read-only access) - - Only WRITES to separate thumbnails directory - - Uses PIL's read-only operations with context managers - - All generated files are clearly marked with '_thumb' suffix + Creates optimized thumbnail versions of all media files in a separate + 'thumbnails' subdirectory. This process is safe and never modifies + original files. + + Safety Features: + - Read-only access to original files + - Writes only to separate thumbnails directory + - Uses PIL context managers for safe file handling + - All thumbnails clearly marked with '_thumb' suffix + - Proper error handling prevents corruption + + Supported Formats: + - Images: PNG, JPG, JPEG, WebP, GIF + - Videos: MP4, AVI, MOV, WMV (generates frame thumbnails) + + Query Parameters: + size (int, optional): Thumbnail size in pixels (default: 256) + quality (int, optional): JPEG quality 1-100 (default: 85) + overwrite (bool, optional): Regenerate existing thumbnails (default: false) + + Args: + request (aiohttp.web.Request): HTTP request with optional query parameters + + Returns: + aiohttp.web.Response: JSON response with structure: + { + \"success\": bool, + \"message\": str, + \"processed\": int, # Number of thumbnails created + \"skipped\": int, # Number of files skipped + \"errors\": int, # Number of processing errors + \"total_size\": str # Total size of created thumbnails + } + + Raises: + Returns 500 status with error details if thumbnail generation fails + + Example: + POST /prompt_manager/images/generate-thumbnails?size=512&quality=90 """ try: import os @@ -2531,7 +2828,41 @@ class PromptManagerAPI: # Diagnostic endpoints async def run_diagnostics(self, request): - """Run system diagnostics.""" + """Run comprehensive system diagnostics and health checks. + + Performs various system health checks including database connectivity, + file system access, ComfyUI integration status, and configuration + validation. Useful for troubleshooting and system monitoring. + + Diagnostic Checks: + - Database connection and integrity + - ComfyUI output directory detection + - File system permissions + - Configuration validation + - Memory and performance metrics + - Extension loading status + + Args: + request (aiohttp.web.Request): HTTP request object + + Returns: + aiohttp.web.Response: JSON response with diagnostic results: + { + "success": bool, + "diagnostics": { + "database": Dict, # Database health info + "filesystem": Dict, # File system status + "comfyui": Dict, # ComfyUI integration status + "config": Dict, # Configuration validation + "performance": Dict # Performance metrics + }, + "summary": str, # Overall system status + "issues": List[str] # List of identified issues + } + + Example: + POST /prompt_manager/diagnostics + """ try: # Simple diagnostics without importing complex modules import os @@ -2682,9 +3013,52 @@ class PromptManagerAPI: }, status=500) async def run_maintenance(self, request): - """ - Comprehensive database maintenance endpoint. - POST /prompt_manager/maintenance + """Perform comprehensive database maintenance and optimization. + + Executes a series of database maintenance operations to optimize + performance, clean up orphaned data, and ensure database integrity. + This is a resource-intensive operation that should be run during + low-traffic periods. + + Maintenance Operations: + - Remove duplicate prompts based on content hash + - Clean up orphaned image references + - Remove missing file records from database + - Optimize database with VACUUM operation + - Update database statistics + - Validate data integrity + - Clean up temporary files + + Query Parameters: + full (bool, optional): Perform full maintenance including VACUUM + cleanup_images (bool, optional): Clean up missing image references + optimize (bool, optional): Run database optimization + + Args: + request (aiohttp.web.Request): HTTP request with optional parameters + + Returns: + aiohttp.web.Response: JSON response with maintenance results: + { + "success": bool, + "operations": { + "duplicates_removed": int, + "orphaned_images_cleaned": int, + "missing_files_removed": int, + "database_optimized": bool, + "integrity_check_passed": bool + }, + "before_stats": Dict, # Database stats before maintenance + "after_stats": Dict, # Database stats after maintenance + "duration": float, # Maintenance duration in seconds + "message": str + } + + Raises: + Returns 500 status with error details if maintenance fails + + Example: + POST /prompt_manager/maintenance?full=true&cleanup_images=true """ try: data = await request.json() if request.content_type == 'application/json' else {} @@ -3201,7 +3575,29 @@ class PromptManagerAPI: return response def _find_comfyui_output_dir(self): - """Find the ComfyUI output directory with improved detection logic.""" + """Locate the ComfyUI output directory using multiple detection strategies. + + Attempts to find the ComfyUI output directory by checking various + possible locations relative to the current file and common installation + patterns. Handles different ComfyUI installation types and structures. + + Detection Strategy: + 1. Look for 'output' directory in parent directories (up to 10 levels) + 2. Check common ComfyUI installation patterns + 3. Verify directory contains typical ComfyUI subdirectories + 4. Return first valid match found + + Returns: + str or None: Absolute path to ComfyUI output directory, or None + if no valid directory is found + + Example: + output_dir = api._find_comfyui_output_dir() + if output_dir: + print(f"Found ComfyUI output at: {output_dir}") + else: + print("ComfyUI output directory not found") + """ import os from pathlib import Path @@ -3268,7 +3664,38 @@ class PromptManagerAPI: return None def _extract_comfyui_metadata(self, image_path): - """Extract ComfyUI metadata from a PNG file.""" + """Extract ComfyUI workflow metadata from PNG image files. + + Reads embedded metadata from PNG files generated by ComfyUI, + extracting workflow information, parameters, and generation details + stored in PNG text chunks. + + ComfyUI stores metadata in standard PNG text chunks: + - 'workflow': Complete node graph workflow data + - 'prompt': Simplified prompt/parameter data + - Custom fields: Additional generation parameters + + Args: + image_path (str): Path to the PNG image file to analyze + + Returns: + Dict[str, Any]: Dictionary containing extracted metadata: + - 'workflow': Raw workflow JSON data (if present) + - 'prompt': Simplified prompt data (if present) + - Additional custom fields from PNG text chunks + - Empty dict if no metadata found or file is not PNG + + Raises: + Exception: If file cannot be opened or read (handled gracefully, + returns empty dict) + + Example: + metadata = api._extract_comfyui_metadata('output/image_001.png') + if 'workflow' in metadata: + print("Found ComfyUI workflow data") + if 'prompt' in metadata: + print(f"Prompt data: {metadata['prompt']}") + """ try: with Image.open(image_path) as img: metadata = {} diff --git a/py/config.py b/py/config.py index dce52ed..5de3937 100644 --- a/py/config.py +++ b/py/config.py @@ -1,3 +1,21 @@ +"""Configuration module for ComfyUI PromptManager. + +This module provides centralized configuration management for the PromptManager +extension, including gallery monitoring settings, database configuration, +web interface options, and performance tuning parameters. + +The configuration is organized into two main classes: +- GalleryConfig: Settings for image monitoring and gallery functionality +- PromptManagerConfig: General settings for the PromptManager core features + +Configuration can be loaded from and saved to JSON files for persistence. + +Example: + from config import PromptManagerConfig + config = PromptManagerConfig.get_config() + PromptManagerConfig.load_from_file('custom_config.json') +""" + # PromptManager/py/config.py # Extension configuration @@ -11,10 +29,6 @@ routes = server_instance.routes # Extension info extension_uri = None # Will be set in __init__.py -""" -Configuration settings for PromptManager gallery and monitoring system. -""" - import os from typing import Dict, Any, List @@ -32,7 +46,31 @@ config_logger = get_logger('prompt_manager.config') class GalleryConfig: - """Configuration for the gallery monitoring system.""" + """Configuration class for the gallery monitoring and image processing system. + + This class manages all settings related to automatic image monitoring, + prompt tracking, database cleanup, web interface display, and performance + optimization for the gallery functionality. + + All configuration values are class attributes that can be modified at runtime + or loaded from external configuration files. + + Attributes: + MONITORING_ENABLED (bool): Enable/disable automatic image monitoring + MONITORING_DIRECTORIES (List[str]): Directories to monitor for new images + SUPPORTED_EXTENSIONS (List[str]): Image file extensions to process + PROCESSING_DELAY (float): Delay in seconds before processing new files + PROMPT_TIMEOUT (int): Seconds to keep prompt context active + CLEANUP_INTERVAL (int): Seconds between cleanup of expired prompts + AUTO_CLEANUP_MISSING_FILES (bool): Automatically remove missing file records + MAX_IMAGE_AGE_DAYS (int): Maximum age in days before cleaning up images + IMAGES_PER_PAGE (int): Number of images to display per page in web UI + THUMBNAIL_SIZE (int): Size in pixels for generated thumbnails + ENABLE_SEARCH (bool): Enable search functionality in web interface + ENABLE_METADATA_VIEW (bool): Enable metadata viewing for images + MAX_CONCURRENT_PROCESSING (int): Maximum concurrent image processing tasks + METADATA_EXTRACTION_TIMEOUT (int): Timeout for metadata extraction operations + """ # Image monitoring settings MONITORING_ENABLED = True @@ -60,7 +98,18 @@ class GalleryConfig: @classmethod def get_config(cls) -> Dict[str, Any]: - """Get the complete configuration as a dictionary.""" + """Get the complete gallery configuration as a structured dictionary. + + Returns: + Dict[str, Any]: A nested dictionary containing all gallery configuration + sections: monitoring, tracking, database, web_interface, and performance. + Each section contains the relevant configuration parameters as key-value pairs. + + Example: + config = GalleryConfig.get_config() + monitoring_enabled = config['monitoring']['enabled'] + images_per_page = config['web_interface']['images_per_page'] + """ return { 'monitoring': { 'enabled': cls.MONITORING_ENABLED, @@ -87,10 +136,91 @@ class GalleryConfig: 'metadata_extraction_timeout': cls.METADATA_EXTRACTION_TIMEOUT } } + + @classmethod + def update_config(cls, new_config: Dict[str, Any]): + """Update gallery configuration attributes from a dictionary. + + Takes a nested dictionary with gallery configuration sections and updates + the corresponding class attributes. Only updates attributes that are + present in the input dictionary, leaving others unchanged. + + Args: + new_config (Dict[str, Any]): Nested dictionary containing gallery + configuration updates. Should follow the same structure as returned + by get_config(). Valid top-level keys are: 'monitoring', 'tracking', + 'database', 'web_interface', 'performance'. + + Example: + gallery_settings = { + 'monitoring': {'enabled': False}, + 'web_interface': {'images_per_page': 50} + } + GalleryConfig.update_config(gallery_settings) + """ + monitoring = new_config.get('monitoring', {}) + if 'enabled' in monitoring: + cls.MONITORING_ENABLED = monitoring['enabled'] + if 'directories' in monitoring: + cls.MONITORING_DIRECTORIES = monitoring['directories'] + if 'extensions' in monitoring: + cls.SUPPORTED_EXTENSIONS = monitoring['extensions'] + if 'processing_delay' in monitoring: + cls.PROCESSING_DELAY = monitoring['processing_delay'] + + tracking = new_config.get('tracking', {}) + if 'prompt_timeout' in tracking: + cls.PROMPT_TIMEOUT = tracking['prompt_timeout'] + if 'cleanup_interval' in tracking: + cls.CLEANUP_INTERVAL = tracking['cleanup_interval'] + + database = new_config.get('database', {}) + if 'auto_cleanup' in database: + cls.AUTO_CLEANUP_MISSING_FILES = database['auto_cleanup'] + if 'max_image_age_days' in database: + cls.MAX_IMAGE_AGE_DAYS = database['max_image_age_days'] + + web_interface = new_config.get('web_interface', {}) + if 'images_per_page' in web_interface: + cls.IMAGES_PER_PAGE = web_interface['images_per_page'] + if 'thumbnail_size' in web_interface: + cls.THUMBNAIL_SIZE = web_interface['thumbnail_size'] + if 'enable_search' in web_interface: + cls.ENABLE_SEARCH = web_interface['enable_search'] + if 'enable_metadata_view' in web_interface: + cls.ENABLE_METADATA_VIEW = web_interface['enable_metadata_view'] + + performance = new_config.get('performance', {}) + if 'max_concurrent_processing' in performance: + cls.MAX_CONCURRENT_PROCESSING = performance['max_concurrent_processing'] + if 'metadata_extraction_timeout' in performance: + cls.METADATA_EXTRACTION_TIMEOUT = performance['metadata_extraction_timeout'] class PromptManagerConfig: - """General configuration for PromptManager.""" + """Main configuration class for PromptManager core functionality. + + This class manages configuration for database operations, web UI behavior, + performance settings, and integrates gallery configuration. It provides + methods for loading and saving configuration from/to JSON files. + + The configuration is organized into logical sections: + - Database: Settings for SQLite operations and data management + - Web UI: User interface behavior and display options + - Performance: Optimization and resource management settings + - Gallery: Embedded gallery configuration (via GalleryConfig) + + Attributes: + DEFAULT_DB_PATH (str): Default path for the SQLite database file + ENABLE_DUPLICATE_DETECTION (bool): Enable automatic duplicate detection + ENABLE_AUTO_SAVE (bool): Enable automatic saving of prompts + RESULT_TIMEOUT (int): Auto-hide timeout for results in ComfyUI node + SHOW_TEST_BUTTON (bool): Show API test button in node interface + WEBUI_DISPLAY_MODE (str): Display mode for web UI ('popup' or 'newtab') + MAX_SEARCH_RESULTS (int): Maximum number of search results to return + ENABLE_FUZZY_SEARCH (bool): Enable fuzzy search capabilities + AUTO_BACKUP_INTERVAL (int): Hours between automatic database backups + """ # Database settings DEFAULT_DB_PATH = "prompts.db" @@ -109,7 +239,20 @@ class PromptManagerConfig: @classmethod def get_config(cls) -> Dict[str, Any]: - """Get the complete configuration as a dictionary.""" + """Get the complete PromptManager configuration as a structured dictionary. + + Returns: + Dict[str, Any]: A nested dictionary containing all configuration sections: + - database: Database-related settings + - web_ui: Web interface configuration + - performance: Performance and optimization settings + - gallery: Complete gallery configuration (from GalleryConfig) + + Example: + config = PromptManagerConfig.get_config() + db_path = config['database']['default_path'] + max_results = config['performance']['max_search_results'] + """ return { 'database': { 'default_path': cls.DEFAULT_DB_PATH, @@ -131,7 +274,24 @@ class PromptManagerConfig: @classmethod def load_from_file(cls, config_path: str): - """Load configuration from a JSON file.""" + """Load configuration settings from a JSON file. + + Reads configuration from the specified JSON file and updates the current + configuration attributes. If the file doesn't exist or contains invalid + JSON, logs an appropriate message and continues with default values. + + Args: + config_path (str): Path to the JSON configuration file to load. + Can be relative or absolute path. + + Raises: + The method handles all exceptions internally and logs errors rather + than propagating them, ensuring the system continues with defaults. + + Example: + PromptManagerConfig.load_from_file('custom_config.json') + PromptManagerConfig.load_from_file('/path/to/config.json') + """ import json if os.path.exists(config_path): @@ -147,7 +307,23 @@ class PromptManagerConfig: @classmethod def save_to_file(cls, config_path: str): - """Save current configuration to a JSON file.""" + """Save the current configuration to a JSON file. + + Serializes the complete configuration (including gallery settings) to + a JSON file. Creates the directory structure if it doesn't exist. + + Args: + config_path (str): Path where the JSON configuration file should be saved. + Parent directories will be created if they don't exist. + + Raises: + The method handles all exceptions internally and logs errors rather + than propagating them. + + Example: + PromptManagerConfig.save_to_file('backup_config.json') + PromptManagerConfig.save_to_file('/etc/comfyui/prompt_manager.json') + """ import json try: @@ -163,7 +339,24 @@ class PromptManagerConfig: @classmethod def update_config(cls, new_config: Dict[str, Any]): - """Update configuration from a dictionary.""" + """Update configuration attributes from a dictionary. + + Takes a nested dictionary with configuration sections and updates + the corresponding class attributes. Only updates attributes that + are present in the input dictionary, leaving others unchanged. + + Args: + new_config (Dict[str, Any]): Nested dictionary containing configuration + updates. Should follow the same structure as returned by get_config(). + Valid top-level keys are: 'database', 'web_ui', 'performance', 'gallery'. + + Example: + new_settings = { + 'database': {'default_path': 'custom.db'}, + 'performance': {'max_search_results': 50} + } + PromptManagerConfig.update_config(new_settings) + """ database = new_config.get('database', {}) if 'default_path' in database: cls.DEFAULT_DB_PATH = database['default_path'] diff --git a/restart_gallery.py b/restart_gallery.py index b83bae7..a4b3a5d 100644 --- a/restart_gallery.py +++ b/restart_gallery.py @@ -1,7 +1,23 @@ #!/usr/bin/env python3 """ -Restart gallery system script. -Use this after fixing code to reload the image monitoring system. +ComfyUI_PromptManager Gallery System Restart Script. + +This script reinitializes and restarts the automatic image gallery monitoring system +for ComfyUI_PromptManager. Use this script after code modifications or when the +gallery system needs to be reloaded. + +The script performs the following operations: +1. Initializes the database connection +2. Creates a PromptTracker instance for linking prompts to images +3. Sets up the ImageMonitor for real-time file system monitoring +4. Starts monitoring ComfyUI output directories +5. Reports system status and readiness + +Usage: + python restart_gallery.py + +This is particularly useful during development when you need to restart the +gallery monitoring system without restarting the entire ComfyUI application. """ import sys diff --git a/tests/__init__.py b/tests/__init__.py index 7a56737..1d61aae 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1 +1,6 @@ -"""Test suite for KikoTextEncode.""" \ No newline at end of file +""" +Test suite for ComfyUI_PromptManager. + +This package contains unit tests and integration tests for the PromptManager +custom nodes, database operations, utility functions, and web interface components. +""" \ No newline at end of file diff --git a/utils/__init__.py b/utils/__init__.py index 0d281ca..e02dbf2 100644 --- a/utils/__init__.py +++ b/utils/__init__.py @@ -1,4 +1,9 @@ -"""Utility functions for KikoTextEncode.""" +""" +Utility functions and modules for ComfyUI_PromptManager. + +This package provides various utility functions for prompt validation, hashing, +image monitoring, metadata extraction, logging, and system diagnostics. +""" from .hashing import generate_prompt_hash from .validators import validate_prompt_text, validate_rating, validate_tags diff --git a/utils/comfyui_integration.py b/utils/comfyui_integration.py index 5d778a4..5c76ea1 100644 --- a/utils/comfyui_integration.py +++ b/utils/comfyui_integration.py @@ -1,6 +1,35 @@ -""" -ComfyUI integration utilities for PromptManager. -Provides hooks and patches to ensure PromptManager metadata appears in standard ComfyUI metadata. +"""ComfyUI integration utilities for PromptManager. + +This module provides deep integration with ComfyUI's metadata system to ensure that +PromptManager-generated prompts appear correctly in standard ComfyUI image metadata. +It patches core ComfyUI components to bridge the gap between PromptManager's custom +nodes and ComfyUI's standard metadata extraction. + +Key features: +- Automatic patching of ComfyUI's SaveImage node +- Thread-safe prompt registration and retrieval +- Cross-thread prompt context sharing +- Standard metadata format compatibility +- Automatic cleanup of old prompt registrations + +The integration works by: +1. PromptManager nodes register their prompts during execution +2. SaveImage node is patched to include registered prompts in metadata +3. PromptManager class_type is changed to CLIPTextEncode for parser compatibility +4. Standard tools can then extract prompts from the generated images + +Typical usage: + from utils.comfyui_integration import get_comfyui_integration + + integration = get_comfyui_integration() + integration.register_prompt(node_id, prompt_text, metadata) + # Generated images will now include this prompt in their metadata + +This integration is essential for: +- Third-party tool compatibility +- Standard metadata parsers +- Workflow sharing and reproduction +- Integration with existing ComfyUI ecosystems """ import threading @@ -18,12 +47,33 @@ except ImportError: class ComfyUIMetadataIntegration: - """Integrates PromptManager with ComfyUI's standard metadata system.""" + """Integrates PromptManager with ComfyUI's standard metadata system. + + This singleton class manages the integration between PromptManager custom nodes + and ComfyUI's standard metadata system. It ensures that prompts generated by + PromptManager appear in the standard ComfyUI metadata format for compatibility + with third-party tools and parsers. + + The integration uses thread-local storage combined with global tracking to + handle prompt context across different execution threads, which is necessary + because ComfyUI executions may span multiple threads. + + Key responsibilities: + - Register prompts from PromptManager nodes during execution + - Patch SaveImage to include PromptManager prompts in metadata + - Convert PromptManager class_type to CLIPTextEncode for compatibility + - Manage prompt lifecycle and cleanup + """ _instance = None _lock = threading.Lock() def __new__(cls): + """Ensure singleton pattern with thread safety. + + Returns: + The single instance of ComfyUIMetadataIntegration + """ if cls._instance is None: with cls._lock: if cls._instance is None: @@ -31,6 +81,11 @@ class ComfyUIMetadataIntegration: return cls._instance def __init__(self): + """Initialize the ComfyUI integration system. + + Sets up prompt tracking data structures and attempts to patch the + SaveImage node. Uses _initialized flag to prevent duplicate initialization. + """ if hasattr(self, '_initialized'): return @@ -47,10 +102,14 @@ class ComfyUIMetadataIntegration: """ Register a prompt from PromptManager for inclusion in ComfyUI metadata. + This method is called by PromptManager nodes during execution to register + their prompt text for later inclusion in image metadata. The prompt is stored + in both thread-local and global storage to handle cross-thread access scenarios. + Args: - node_id: The node ID (unique identifier for this prompt) - prompt_text: The actual prompt text that was encoded - metadata: Additional metadata from PromptManager + node_id: Unique identifier for this prompt (typically the node ID) + prompt_text: The actual prompt text that was encoded by PromptManager + metadata: Additional metadata from PromptManager (category, tags, etc.) """ thread_id = threading.current_thread().ident @@ -80,11 +139,15 @@ class ComfyUIMetadataIntegration: """ Get the current prompt text for metadata inclusion. + Retrieves the most appropriate prompt text for inclusion in image metadata. + Uses a fallback strategy: thread-local storage first, then global storage, + with recency checks to avoid stale prompts. + Args: - node_id: Optional specific node ID to get prompt for + node_id: Optional specific node ID to retrieve prompt for Returns: - The prompt text or None + The prompt text string if available, None if no suitable prompt found """ # First try thread-local storage if hasattr(self._thread_local, 'prompts'): @@ -118,6 +181,18 @@ class ComfyUIMetadataIntegration: def _patch_saveimage_node(self): """ Patch ComfyUI's SaveImage node to include PromptManager prompts in metadata. + + This method modifies ComfyUI's SaveImage.save_images method to automatically + include PromptManager prompts in the image metadata. The patching: + + 1. Wraps the original save_images method + 2. Retrieves current PromptManager prompt text + 3. Modifies the workflow data to include the prompt + 4. Changes PromptManager class_type to CLIPTextEncode for compatibility + 5. Calls the original method with modified data + + The patching is designed to be minimally invasive and maintain full + compatibility with existing ComfyUI functionality. """ try: import nodes @@ -199,8 +274,13 @@ class ComfyUIMetadataIntegration: """ Clean up old prompt registrations. + Removes prompt registrations that are older than the specified age to + prevent memory leaks and ensure that only recent, relevant prompts are + used for metadata inclusion. + Args: - max_age_seconds: Maximum age in seconds before cleanup + max_age_seconds: Maximum age in seconds before a prompt registration + is considered stale and removed (default: 600 seconds/10 minutes) """ current_time = time.time() @@ -221,7 +301,11 @@ class ComfyUIMetadataIntegration: _integration_instance = None def get_comfyui_integration() -> ComfyUIMetadataIntegration: - """Get the global ComfyUI integration instance.""" + """Get the global ComfyUI integration instance. + + Returns: + The singleton ComfyUIMetadataIntegration instance, creating it if necessary + """ global _integration_instance if _integration_instance is None: _integration_instance = ComfyUIMetadataIntegration() diff --git a/utils/diagnostics.py b/utils/diagnostics.py index 5620309..88e67cf 100644 --- a/utils/diagnostics.py +++ b/utils/diagnostics.py @@ -1,5 +1,33 @@ -""" -Diagnostic utilities for troubleshooting the PromptManager gallery system. +"""Diagnostic utilities for troubleshooting the PromptManager gallery system. + +This module provides comprehensive diagnostic tools to help identify and resolve +issues with the PromptManager system, particularly the gallery and image monitoring +functionality. It performs systematic checks of all system components and provides +clear feedback on potential problems. + +The diagnostic system checks: +- Database connectivity and schema integrity +- Image table structure and data +- File system permissions and access +- ComfyUI output directory detection +- Required Python dependencies +- System integration points + +Typical usage: + from utils.diagnostics import GalleryDiagnostics, run_diagnostics + + # Run full diagnostic suite + results = run_diagnostics() + + # Or create custom diagnostic instance + diagnostics = GalleryDiagnostics("custom_db.db") + database_status = diagnostics.check_database() + +The diagnostics provide: +- Clear pass/fail status for each component +- Detailed error messages and remediation suggestions +- Statistics and system information +- Test utilities for verifying functionality """ import os @@ -11,14 +39,48 @@ from .logging_config import get_logger class GalleryDiagnostics: - """Diagnostics for the gallery system.""" + """Diagnostics for the gallery system. + + This class provides comprehensive diagnostic capabilities for the PromptManager + gallery system. It systematically checks all components and dependencies to + identify potential issues and provide actionable feedback. + + The diagnostics cover: + - Database structure and connectivity + - Image tracking table integrity + - File system access and permissions + - ComfyUI integration points + - Python dependency availability + + Each diagnostic method returns a standardized result dictionary with: + - status: 'ok', 'warning', or 'error' + - message: Descriptive message (for warnings/errors) + - Additional data specific to the diagnostic + """ def __init__(self, db_path: str = "prompts.db"): + """Initialize the diagnostics system. + + Args: + db_path: Path to the SQLite database file to diagnose + """ self.db_path = db_path self.logger = get_logger('prompt_manager.diagnostics') def run_full_diagnostic(self) -> Dict[str, Any]: - """Run a complete diagnostic check.""" + """Run a complete diagnostic check. + + Executes all diagnostic checks in sequence and provides a comprehensive + report of system status. Logs detailed information during the process. + + Returns: + Dictionary mapping diagnostic categories to their results: + - database: Database connectivity and structure check + - images_table: Image tracking table specific check + - file_system: File system access and permissions check + - comfyui_output: ComfyUI output directory detection + - dependencies: Python dependency availability check + """ self.logger.info("\n" + "="*60) self.logger.info("[DIAG] PROMPTMANAGER GALLERY DIAGNOSTICS") self.logger.info("="*60) @@ -45,7 +107,18 @@ class GalleryDiagnostics: return results def check_database(self) -> Dict[str, Any]: - """Check database connection and structure.""" + """Check database connection and structure. + + Verifies that the database file exists, is accessible, and contains + the expected prompt table structure. + + Returns: + Dictionary containing: + - status: 'ok' or 'error' + - message: Error description (if status is 'error') + - prompt_count: Number of prompts in database (if successful) + - has_images_table: Whether generated_images table exists (if successful) + """ self.logger.info("\n[DB] Checking Database...") try: @@ -85,7 +158,18 @@ class GalleryDiagnostics: } def check_images_table(self) -> Dict[str, Any]: - """Check the generated_images table specifically.""" + """Check the generated_images table specifically. + + Examines the generated_images table structure and content to verify + the gallery system can function properly. + + Returns: + Dictionary containing: + - status: 'ok' or 'error' + - message: Error description (if table missing or inaccessible) + - image_count: Number of images in table (if successful) + - recent_images: List of recent image records (if successful) + """ self.logger.info("\n[IMG] Checking Images Table...") try: @@ -136,7 +220,18 @@ class GalleryDiagnostics: } def check_file_system(self) -> Dict[str, Any]: - """Check file system and permissions.""" + """Check file system and permissions. + + Verifies that the application has appropriate file system access + for reading images and writing database files. + + Returns: + Dictionary containing: + - status: 'ok' or 'error' + - message: Error description (if issues detected) + - current_dir: Current working directory path + - can_write: Whether write access is available + """ self.logger.info("\n[DIR] Checking File System...") try: @@ -169,7 +264,17 @@ class GalleryDiagnostics: } def check_comfyui_output(self) -> Dict[str, Any]: - """Check ComfyUI output directories.""" + """Check ComfyUI output directories. + + Attempts to locate ComfyUI output directories where generated images + would be stored. Checks both common locations and ComfyUI's configured paths. + + Returns: + Dictionary containing: + - status: 'ok' if directories found, 'warning' if none found + - message: Warning message (if no directories found) + - output_dirs: List of detected output directory paths + """ self.logger.info("\n[STYLE] Checking ComfyUI Output...") output_dirs = [] @@ -216,7 +321,17 @@ class GalleryDiagnostics: } def check_dependencies(self) -> Dict[str, Any]: - """Check required dependencies.""" + """Check required dependencies. + + Verifies that all required Python packages are available for the + PromptManager system to function properly. + + Returns: + Dictionary containing: + - status: 'ok' if all dependencies available, 'error' if any missing + - message: Error description (if dependencies missing) + - dependencies: Dictionary mapping package names to availability status + """ self.logger.info("\n[PKG] Checking Dependencies...") dependencies = { @@ -258,7 +373,21 @@ class GalleryDiagnostics: } def create_test_image_link(self, prompt_id: int, test_image_path: str = None) -> Dict[str, Any]: - """Create a test image link to verify the system works.""" + """Create a test image link to verify the system works. + + Creates a test entry in the generated_images table to verify that + the image linking functionality is working correctly. + + Args: + prompt_id: ID of an existing prompt to link the test image to + test_image_path: Optional path for the test image (uses fake path if None) + + Returns: + Dictionary containing: + - status: 'ok' if test link created successfully, 'error' otherwise + - message: Success or error message + - image_id: ID of created test image record (if successful) + """ self.logger.info(f"\n[TEST] Creating test image link for prompt {prompt_id}...") try: @@ -307,8 +436,15 @@ class GalleryDiagnostics: } -def run_diagnostics(): - """Run diagnostics from command line.""" +def run_diagnostics() -> Dict[str, Any]: + """Run diagnostics from command line. + + Convenience function to create a GalleryDiagnostics instance and run + the full diagnostic suite with default settings. + + Returns: + Complete diagnostic results dictionary + """ diagnostics = GalleryDiagnostics() return diagnostics.run_full_diagnostic() diff --git a/utils/hashing.py b/utils/hashing.py index 881b0e3..fc2253e 100644 --- a/utils/hashing.py +++ b/utils/hashing.py @@ -1,5 +1,31 @@ -""" -Hashing utilities for KikoTextEncode prompt deduplication. +"""Hashing utilities for PromptManager prompt deduplication. + +This module provides cryptographic hashing functions to enable efficient +deduplication of prompts and content in the PromptManager system. It uses +SHA256 hashing with text normalization to ensure consistent hash generation +regardless of minor formatting differences. + +Key features: +- SHA256-based prompt text hashing with normalization +- Content hashing for complex prompt metadata structures +- Duplicate detection utilities +- Consistent normalization (lowercase, trimmed whitespace) + +Typical usage: + from utils.hashing import generate_prompt_hash, is_duplicate_prompt + + hash1 = generate_prompt_hash("Beautiful landscape") + hash2 = generate_prompt_hash(" beautiful landscape ") + # hash1 == hash2 (normalization makes them identical) + + if is_duplicate_prompt(text1, text2): + print("Duplicate prompts detected") + +The hashing is designed for: +- Database deduplication (preventing duplicate prompt storage) +- Quick similarity checking +- Content integrity verification +- Efficient database indexing """ import hashlib @@ -9,15 +35,25 @@ def generate_prompt_hash(text: str) -> str: """ Generate a SHA256 hash for prompt text to enable deduplication. + Normalizes the input text (strips whitespace, converts to lowercase) + before hashing to ensure consistent results for functionally identical + prompts with minor formatting differences. + Args: text: The prompt text to hash Returns: - str: SHA256 hexdigest of the text + SHA256 hexadecimal digest of the normalized text + + Raises: + TypeError: If the input is not a string Example: - >>> generate_prompt_hash("beautiful landscape") - 'a1b2c3d4e5f6...' + >>> generate_prompt_hash("Beautiful Landscape") + 'c3d5f8a9b2e1...' + >>> generate_prompt_hash(" beautiful landscape ") + 'c3d5f8a9b2e1...' + # Both produce the same hash due to normalization """ if not isinstance(text, str): raise TypeError("Text must be a string") @@ -33,11 +69,23 @@ def generate_content_hash(content: dict) -> str: """ Generate a hash for prompt content including metadata. + Creates a comprehensive hash that includes not just the prompt text + but also associated metadata like category, tags, and workflow name. + This enables detection of prompts that are identical in all aspects. + Args: - content: Dictionary containing prompt data + content: Dictionary containing prompt data with optional keys: + - text: The prompt text + - category: Prompt category + - tags: List of tags + - workflow_name: Associated workflow name Returns: - str: SHA256 hexdigest of the content + SHA256 hexadecimal digest of the normalized content structure + + Note: + The hash is generated from a normalized JSON representation with + sorted keys and normalized text fields to ensure consistency. """ import json @@ -59,13 +107,23 @@ def is_duplicate_prompt(text1: str, text2: str, threshold: float = 0.95) -> bool """ Check if two prompts are likely duplicates using hash comparison. + Compares the normalized hashes of two prompt texts to determine if + they are functionally identical. This is an exact match after normalization, + not a similarity measure. + Args: - text1: First prompt text - text2: Second prompt text - threshold: Similarity threshold (not used for exact hash matching) + text1: First prompt text to compare + text2: Second prompt text to compare + threshold: Similarity threshold (not used - kept for API compatibility) Returns: - bool: True if prompts are likely duplicates + True if the prompts have identical normalized hashes (i.e., are duplicates), + False otherwise + + Note: + The threshold parameter is not used in the current implementation as + this performs exact hash matching rather than similarity scoring. + It's maintained for potential future fuzzy matching capabilities. """ hash1 = generate_prompt_hash(text1) hash2 = generate_prompt_hash(text2) diff --git a/utils/image_monitor.py b/utils/image_monitor.py index b7928ab..0f28cc3 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -1,6 +1,26 @@ -""" -Image monitoring system for ComfyUI generated images. -Automatically detects new images and links them to prompts. +"""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 @@ -17,15 +37,26 @@ from .logging_config import get_logger class ImageGenerationHandler(FileSystemEventHandler): - """Handler for new image file creation events.""" + """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 handler. + Initialize the image generation handler. Args: - db_manager: Database manager instance - prompt_tracker: Prompt tracking instance + 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 @@ -34,7 +65,15 @@ class ImageGenerationHandler(FileSystemEventHandler): self.logger = get_logger('prompt_manager.image_monitor') def on_created(self, event): - """Handle file creation events.""" + """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. + + Args: + event: FileSystemEvent object containing event details + """ if not event.is_directory and self.is_image_file(event.src_path): self.logger.debug(f"New image detected: {event.src_path}") # Small delay to ensure file is fully written @@ -45,11 +84,29 @@ class ImageGenerationHandler(FileSystemEventHandler): ).start() def is_image_file(self, filepath: str) -> bool: - """Check if file is a supported image format.""" + """Check if file is a supported image format. + + Args: + filepath: Path to the file to check + + Returns: + True if the file has a supported image extension, False otherwise + """ return filepath.lower().endswith(('.png', '.jpg', '.jpeg', '.webp', '.gif')) def process_new_image(self, image_path: str): - """Process a newly created image file.""" + """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. Gets the current prompt context from the tracker + 3. Extracts ComfyUI metadata from the image + 4. Links the image to the appropriate prompt in the database + 5. Handles fallback scenarios when no active prompt is available + + Args: + image_path: Full path to the newly created image file + """ try: self.logger.debug(f"Processing image: {image_path}") @@ -98,7 +155,20 @@ class ImageGenerationHandler(FileSystemEventHandler): self.logger.error(traceback.format_exc()) def get_basic_file_info(self, image_path: str) -> Dict[str, Any]: - """Get basic file information when metadata extraction fails.""" + """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 @@ -123,7 +193,16 @@ class ImageGenerationHandler(FileSystemEventHandler): return {} def _get_fallback_prompt(self) -> Optional[Dict[str, Any]]: - """Get the most recent prompt from database as fallback.""" + """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_prompts = self.db_manager.get_recent_prompts(limit=1) if recent_prompts: @@ -139,7 +218,16 @@ class ImageGenerationHandler(FileSystemEventHandler): 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.""" + """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: image_id = self.db_manager.link_image_to_prompt( prompt_id=prompt_context['id'], @@ -153,15 +241,25 @@ class ImageGenerationHandler(FileSystemEventHandler): class ImageMonitor: - """Main image monitoring system.""" + """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 - prompt_tracker: Prompt tracking instance + 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 @@ -172,10 +270,14 @@ class ImageMonitor: def start_monitoring(self, output_directories: Optional[list] = None): """ - Start monitoring ComfyUI output directories. + Start monitoring ComfyUI output directories for new images. + + Begins filesystem watching on the specified directories. If no directories + are provided, the system will attempt to auto-detect ComfyUI output locations. + All monitoring is done recursively to catch images in subdirectories. Args: - output_directories: List of directories to monitor. If None, auto-detect. + output_directories: List of directory paths to monitor. If None, uses auto-detection. """ if self.observer: self.logger.warning("Image monitoring already running") @@ -210,7 +312,11 @@ class ImageMonitor: self.logger.warning("No valid directories to monitor") def stop_monitoring(self): - """Stop the image monitoring system.""" + """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. + """ if self.observer: self.observer.stop() self.observer.join() @@ -220,7 +326,16 @@ class ImageMonitor: self.logger.debug("Image monitoring stopped") def detect_comfyui_output_dirs(self) -> list: - """Auto-detect ComfyUI output directories.""" + """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: @@ -251,7 +366,14 @@ class ImageMonitor: return potential_dirs def get_status(self) -> Dict[str, Any]: - """Get monitoring status information.""" + """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 + """ return { 'running': self.observer is not None, 'monitored_directories': self.monitored_directories, diff --git a/utils/logging_config.py b/utils/logging_config.py index 5ccd5f6..ef29122 100644 --- a/utils/logging_config.py +++ b/utils/logging_config.py @@ -35,6 +35,11 @@ class PromptManagerLogger: _lock = threading.Lock() def __new__(cls): + """Ensure singleton pattern with thread safety. + + Returns: + The single instance of PromptManagerLogger + """ if cls._instance is None: with cls._lock: if cls._instance is None: @@ -42,6 +47,11 @@ class PromptManagerLogger: return cls._instance def __init__(self): + """Initialize the logging system if not already initialized. + + Sets up log directory, configuration, memory buffer, and all handlers. + Uses _initialized flag to prevent duplicate initialization. + """ if hasattr(self, '_initialized'): return @@ -67,7 +77,11 @@ class PromptManagerLogger: self._setup_loggers() def _setup_loggers(self): - """Set up all loggers with appropriate handlers.""" + """Set up all loggers with appropriate handlers. + + Configures the main logger with console, file, and memory handlers. + Sets up file rotation and safe encoding for cross-platform compatibility. + """ # Main logger self.logger = logging.getLogger('prompt_manager') self.logger.setLevel(getattr(logging, self.config['level'])) @@ -110,7 +124,11 @@ class PromptManagerLogger: self._setup_component_loggers() def _setup_component_loggers(self): - """Set up loggers for specific components.""" + """Set up loggers for specific components. + + Creates child loggers for different PromptManager components. + Child loggers inherit handlers from the main logger. + """ components = [ 'prompt_manager.database', 'prompt_manager.api', @@ -125,11 +143,23 @@ class PromptManagerLogger: # Child loggers inherit handlers from parent def get_logger(self, name: str = 'prompt_manager') -> logging.Logger: - """Get a logger instance for a specific component.""" + """Get a logger instance for a specific component. + + Args: + name: Logger name, typically in format 'prompt_manager.component' + + Returns: + Configured logger instance + """ return logging.getLogger(name) def add_to_buffer(self, record: logging.LogRecord, formatted_message: str): - """Add a log entry to the memory buffer for web viewer.""" + """Add a log entry to the memory buffer for web viewer. + + Args: + record: The LogRecord object from the logging system + formatted_message: The formatted log message string + """ with self._buffer_lock: log_entry = { 'timestamp': datetime.fromtimestamp(record.created).isoformat(), @@ -152,7 +182,15 @@ class PromptManagerLogger: self._log_buffer = self._log_buffer[-self.config['buffer_size']:] def get_recent_logs(self, limit: int = 100, level: Optional[str] = None) -> List[Dict[str, Any]]: - """Get recent log entries from memory buffer.""" + """Get recent log entries from memory buffer. + + Args: + limit: Maximum number of log entries to return + level: Optional log level filter (DEBUG, INFO, WARNING, ERROR, CRITICAL) + + Returns: + List of log entry dictionaries, most recent first + """ with self._buffer_lock: logs = self._log_buffer[:] @@ -166,7 +204,16 @@ class PromptManagerLogger: return list(reversed(logs[-limit:])) def get_log_files(self) -> List[Dict[str, Any]]: - """Get information about available log files.""" + """Get information about available log files. + + Returns: + List of dictionaries containing file information: + - filename: Name of the log file + - path: Full path to the log file + - size: File size in bytes + - modified: ISO formatted modification timestamp + - is_main: True if this is the main log file + """ log_files = [] for log_file in self.log_dir.glob("*.log*"): @@ -187,7 +234,19 @@ class PromptManagerLogger: return log_files def read_log_file(self, filename: str, lines: int = 100) -> List[str]: - """Read lines from a specific log file.""" + """Read lines from a specific log file. + + Args: + filename: Name of the log file to read (must be in log directory) + lines: Number of lines to read from the end of the file (0 for all) + + Returns: + List of line strings from the log file + + Raises: + FileNotFoundError: If the log file doesn't exist + ValueError: If the file path is outside the log directory + """ log_file = self.log_dir / filename if not log_file.exists() or not log_file.is_file(): @@ -207,7 +266,16 @@ class PromptManagerLogger: raise def truncate_logs(self) -> Dict[str, Any]: - """Truncate all log files.""" + """Truncate all log files. + + Clears the main log file and deletes rotated log files. + Also clears the memory buffer. + + Returns: + Dictionary with: + - truncated: List of successfully truncated files + - errors: List of error messages for failed operations + """ results = { 'truncated': [], 'errors': [] @@ -238,7 +306,12 @@ class PromptManagerLogger: return results def update_config(self, new_config: Dict[str, Any]): - """Update logging configuration.""" + """Update logging configuration. + + Args: + new_config: Dictionary of configuration updates + (level, console_logging, file_logging, etc.) + """ self.config.update(new_config) # Reconfigure loggers if level changed @@ -259,11 +332,25 @@ class PromptManagerLogger: self.logger.info(f"Updated logging configuration: {new_config}") def get_config(self) -> Dict[str, Any]: - """Get current logging configuration.""" + """Get current logging configuration. + + Returns: + Copy of the current configuration dictionary + """ return self.config.copy() def get_log_stats(self) -> Dict[str, Any]: - """Get logging statistics.""" + """Get logging statistics. + + Returns: + Dictionary containing: + - buffer_count: Number of entries in memory buffer + - level_counts: Count of entries by log level + - log_files_count: Number of log files + - total_log_size: Total size of all log files in bytes + - log_directory: Path to the log directory + - current_level: Current logging level + """ with self._buffer_lock: buffer_count = len(self._log_buffer) level_counts = {} @@ -285,14 +372,31 @@ class PromptManagerLogger: class MemoryBufferHandler(logging.Handler): - """Custom logging handler that stores logs in memory for web viewer.""" + """Custom logging handler that stores logs in memory for web viewer. + + This handler extends the standard logging.Handler to capture log records + and store them in the PromptManagerLogger's memory buffer. This enables + the web interface to display recent log entries without reading from files. + + The handler is designed to be failure-safe - if any error occurs during + log processing, it silently continues without disrupting the application. + """ def __init__(self, logger_manager: PromptManagerLogger): + """Initialize the memory buffer handler. + + Args: + logger_manager: The PromptManagerLogger instance to store logs in + """ super().__init__() self.logger_manager = logger_manager def emit(self, record: logging.LogRecord): - """Handle a log record by adding it to the memory buffer.""" + """Handle a log record by adding it to the memory buffer. + + Args: + record: The LogRecord to process and store + """ try: formatted_message = self.format(record) self.logger_manager.add_to_buffer(record, formatted_message) @@ -305,14 +409,25 @@ class MemoryBufferHandler(logging.Handler): _logger_manager = None def get_logger_manager() -> PromptManagerLogger: - """Get the global logger manager instance.""" + """Get the global logger manager instance. + + Returns: + The singleton PromptManagerLogger instance, creating it if necessary + """ global _logger_manager if _logger_manager is None: _logger_manager = PromptManagerLogger() return _logger_manager def get_logger(name: str = 'prompt_manager') -> logging.Logger: - """Convenience function to get a logger.""" + """Convenience function to get a logger. + + Args: + name: Logger name, defaults to 'prompt_manager' + + Returns: + Configured logger instance for the specified name + """ return get_logger_manager().get_logger(name) # Initialize logging on import diff --git a/utils/metadata_extractor.py b/utils/metadata_extractor.py index 632f125..2fecac7 100644 --- a/utils/metadata_extractor.py +++ b/utils/metadata_extractor.py @@ -1,6 +1,33 @@ -""" -ComfyUI metadata extraction utilities. -Extracts workflow and prompt information from generated images. +"""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. + +The extractor supports: +- Complete workflow data extraction from PNG text chunks +- Text encoder node identification and analysis +- Generation parameter extraction (steps, cfg_scale, sampler, etc.) +- Basic file information fallback when metadata is unavailable +- Flexible node type detection for various ComfyUI extensions + +Typical usage: + from utils.metadata_extractor import ComfyUIMetadataExtractor + + extractor = ComfyUIMetadataExtractor() + metadata = extractor.extract_metadata('/path/to/image.png') + + if metadata: + workflow = metadata.get('workflow', {}) + text_nodes = metadata.get('text_encoder_nodes', []) + params = extractor.get_generation_parameters(metadata) + +The extracted metadata includes: +- file_info: Basic file stats (size, dimensions, format, timestamps) +- workflow: Complete ComfyUI workflow data structure +- text_encoder_nodes: List of identified text encoding nodes +- prompt: ComfyUI prompt execution data +- Generation parameters: steps, cfg_scale, sampler, seed, etc. """ import os @@ -13,21 +40,46 @@ from .logging_config import get_logger class ComfyUIMetadataExtractor: - """Extracts ComfyUI metadata from generated images.""" + """Extracts ComfyUI metadata from generated images. + + This class handles extraction of ComfyUI workflow and generation metadata + from PNG images. It parses the text chunks embedded by ComfyUI and extracts + structured information about workflows, prompts, and generation parameters. + + The extractor is designed to be robust and handle various ComfyUI workflow + formats, including custom nodes and different metadata structures. + """ def __init__(self): - """Initialize the metadata extractor.""" + """Initialize the metadata extractor. + + Sets up logging and prepares the extractor for metadata parsing operations. + """ self.logger = get_logger('prompt_manager.metadata_extractor') def extract_metadata(self, image_path: str) -> Optional[Dict[str, Any]]: """ Extract ComfyUI workflow and prompt metadata from an image. + This method opens a PNG image and extracts all available ComfyUI metadata + from the embedded text chunks. It handles JSON parsing, workflow analysis, + and parameter extraction. + Args: - image_path: Path to the image file + image_path: Path to the PNG image file to analyze Returns: - Dictionary containing extracted metadata or None if extraction fails + Dictionary containing extracted metadata with keys: + - file_info: Basic file information (always present) + - workflow: Parsed workflow data (if available) + - text_encoder_nodes: List of text encoding nodes (if found) + - prompt: ComfyUI prompt data (if available) + - Additional fields: parameters, model, sampler, steps, etc. (if present) + Returns None if no ComfyUI metadata is found or extraction fails + + Raises: + Exception: Re-raises any exception that occurs during extraction, + with appropriate error logging """ try: with Image.open(image_path) as image: @@ -85,12 +137,21 @@ class ComfyUIMetadataExtractor: """ Get basic file information. + Extracts fundamental file properties including size, dimensions, format, + and filesystem timestamps. + Args: image_path: Path to the image file - image: PIL Image object + image: Opened PIL Image object Returns: - Dictionary containing file information + Dictionary containing: + - size: File size in bytes + - dimensions: Image dimensions as [width, height] + - format: Image format (PNG, JPEG, etc.) + - mode: Color mode (RGB, RGBA, etc.) + - created_time: File creation timestamp + - modified_time: File modification timestamp """ try: stat = os.stat(image_path) @@ -110,11 +171,17 @@ class ComfyUIMetadataExtractor: """ Find text encoder nodes in the workflow data. + Searches through the workflow structure to identify nodes that perform + text encoding operations. This includes standard ComfyUI nodes like + CLIPTextEncode as well as custom nodes and extensions. + Args: - workflow_data: ComfyUI workflow data + workflow_data: ComfyUI workflow data structure (can be in various formats) Returns: - List of text encoder node data + List of dictionaries representing text encoder nodes, each containing + the node's configuration, inputs, and metadata. For dictionary-based + workflows, adds 'node_id' field to each node. """ text_encoder_nodes = [] @@ -155,11 +222,15 @@ class ComfyUIMetadataExtractor: """ Check if a node is a text encoder node. + Determines whether a given node performs text encoding by examining + its type, class, and title for known text encoding patterns. Supports + various ComfyUI node types and custom extensions. + Args: - node_data: Node data to check + node_data: Node configuration dictionary to analyze Returns: - True if the node is a text encoder + True if the node is identified as a text encoder, False otherwise """ if not isinstance(node_data, dict): return False @@ -195,11 +266,15 @@ class ComfyUIMetadataExtractor: """ Extract the actual prompt text from workflow data. + Searches through text encoder nodes to find and extract the actual + prompt text that was used for generation. Handles various input field + names and data structures. + Args: - workflow_data: ComfyUI workflow data + workflow_data: ComfyUI workflow data structure Returns: - Extracted prompt text or None + The extracted prompt text string, or None if no prompt text is found """ text_encoder_nodes = self.find_text_encoder_nodes(workflow_data) @@ -223,11 +298,23 @@ class ComfyUIMetadataExtractor: """ Extract generation parameters from metadata. + Parses the metadata to extract key generation parameters used for + image creation, including sampling settings, model information, + and other configuration values. + Args: - metadata: Full metadata dictionary + metadata: Full metadata dictionary from extract_metadata() Returns: - Dictionary of generation parameters + Dictionary containing generation parameters such as: + - steps: Number of sampling steps + - cfg_scale: Classifier-free guidance scale + - sampler: Sampling method used + - scheduler: Noise scheduler + - seed: Random seed value + - model: Model name/path + - width/height: Image dimensions + - batch_size: Number of images generated """ parameters = {} @@ -252,11 +339,16 @@ class ComfyUIMetadataExtractor: """ Extract generation parameters from workflow data. + Analyzes the workflow structure to identify and extract generation + parameters from various node types. This method can be extended + to support specific workflow patterns and custom nodes. + Args: - workflow_data: ComfyUI workflow data + workflow_data: ComfyUI workflow data structure Returns: - Dictionary of extracted parameters + Dictionary of extracted parameters. Currently returns empty dict + but can be extended based on specific workflow analysis needs. """ parameters = {} diff --git a/utils/prompt_tracker.py b/utils/prompt_tracker.py index 23baf10..d8dd374 100644 --- a/utils/prompt_tracker.py +++ b/utils/prompt_tracker.py @@ -1,6 +1,28 @@ -""" -Prompt tracking system for linking generated images to prompts. -Tracks active prompt executions to associate them with generated images. +"""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. + +Key features: +- Thread-safe prompt tracking using threading.local and locks +- Automatic cleanup of expired prompt contexts +- Fallback mechanisms for cross-thread prompt access +- Support for execution timeouts and extensions +- Context manager for automatic prompt lifecycle management + +Typical usage: + from utils.prompt_tracker import PromptTracker + + tracker = PromptTracker(db_manager) + execution_id = tracker.set_current_prompt("beautiful landscape") + # Images generated after this point will be linked to this prompt + +Or using the context manager: + with PromptExecutionContext(tracker, "beautiful landscape"): + # Generate images here + pass """ import threading @@ -22,14 +44,29 @@ except ImportError: class PromptTracker: - """Thread-safe tracking of current prompt executions.""" + """Thread-safe tracking of current prompt executions. + + This class manages active prompt contexts across multiple threads, enabling + the image monitoring system to correctly associate generated images with their + source prompts. It uses both thread-local storage and global tracking to handle + various execution scenarios. + + The tracker automatically cleans up expired prompts and provides fallback + mechanisms when prompts are accessed from different threads (e.g., image + monitoring running in a separate thread from prompt execution). + + Attributes: + active_prompts: Global dictionary of active prompt contexts + cleanup_interval: Seconds between cleanup operations (default: 300) + prompt_timeout: Seconds before a prompt expires (default: 600) + """ def __init__(self, db_manager): """ Initialize the prompt tracker. Args: - db_manager: Database manager instance + db_manager: Database manager instance for prompt operations """ self.logger = get_logger('prompt_manager.prompt_tracker') self.logger.debug("Initializing PromptTracker") @@ -49,14 +86,19 @@ class PromptTracker: def set_current_prompt(self, prompt_text: str, additional_data: Optional[Dict[str, Any]] = None) -> str: """ Set the current prompt for this thread for image tracking. + + This method establishes the prompt context that will be used to link any + subsequently generated images. The prompt is stored in both thread-local + storage and global tracking to handle cross-thread access scenarios. + Note: Prompt saving is handled by PromptManager to avoid duplicates. 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 Returns: - Execution ID for this prompt + Unique execution ID for this prompt execution """ execution_id = self.generate_execution_id() @@ -109,8 +151,19 @@ class PromptTracker: """ Get the current prompt context for this thread. + Attempts to retrieve the current prompt context, first from thread-local + storage, then from global tracking as a fallback. This enables image + monitoring (which may run in a different thread) to access prompt context. + Returns: - Current prompt context or None + Dictionary containing prompt context with keys: + - id: Prompt ID in database + - execution_id: Unique execution identifier + - text: Prompt text + - timestamp: When the prompt was set + - thread_id: Thread identifier + - additional_data: Any additional metadata + Returns None if no valid prompt context is found """ current = getattr(self._local, 'current_prompt', None) @@ -133,7 +186,15 @@ class PromptTracker: return None def _find_recent_prompt(self) -> Optional[Dict[str, Any]]: - """Find the most recent prompt that's still valid.""" + """Find the most recent prompt that's still valid. + + Searches the global prompt tracking for the most recently set prompt + that hasn't expired. This is used as a fallback when thread-local + storage doesn't contain a current prompt. + + Returns: + Most recent valid prompt context, or None if no valid prompts found + """ with self.lock: current_time = time.time() self.logger.debug(f"Searching for recent prompt among {len(self.active_prompts)} active prompts") @@ -159,7 +220,12 @@ class PromptTracker: return None def clear_current_prompt(self): - """Clear the current prompt context for this thread.""" + """Clear the current prompt context for this thread. + + Removes the prompt context from both thread-local storage and global + tracking. This should be called when prompt execution is complete, + though the system also handles automatic cleanup via timeouts. + """ current = getattr(self._local, 'current_prompt', None) if current: execution_id = current['execution_id'] @@ -183,9 +249,14 @@ class PromptTracker: """ Extend the timeout for a specific prompt execution. + This is useful for long-running image generation processes where images + may continue to be generated well after the initial prompt execution. + Updates the timestamp to prevent the prompt from being cleaned up. + Args: execution_id: The execution ID to extend - additional_seconds: Additional seconds to add to timeout + additional_seconds: Additional seconds to add to timeout (currently unused, + method just resets the timestamp to current time) """ with self.lock: if execution_id in self.active_prompts: @@ -193,11 +264,22 @@ class PromptTracker: self.logger.debug(f"Extended timeout for prompt: {execution_id}") def generate_execution_id(self) -> str: - """Generate a unique execution ID.""" + """Generate a unique execution ID. + + Creates a unique identifier for this prompt execution using UUID and timestamp. + + Returns: + Unique execution ID in format 'exec_{uuid8}_{timestamp}' + """ return f"exec_{uuid.uuid4().hex[:8]}_{int(time.time())}" def _cleanup_expired_prompts(self): - """Background thread to clean up expired prompts.""" + """Background thread to clean up expired prompts. + + Runs continuously as a daemon thread, periodically removing expired + prompt contexts from the global tracking dictionary. This prevents + memory leaks from accumulating old prompt data. + """ while True: try: current_time = time.time() @@ -225,17 +307,21 @@ class PromptTracker: Get all currently active prompts. Returns: - Dictionary of active prompts + Dictionary mapping execution IDs to prompt context dictionaries. + This is a copy of the internal tracking dictionary. """ with self.lock: return self.active_prompts.copy() def clear_all_active_prompts(self) -> int: """ - Clear all active prompts (useful before scanning to avoid conflicts). + Clear all active prompts. + + Removes all prompt contexts from both global and thread-local storage. + This is useful before batch operations or when resetting the tracking state. Returns: - Number of prompts that were cleared + Number of prompts that were cleared from global tracking """ with self.lock: cleared_count = len(self.active_prompts) @@ -252,7 +338,13 @@ class PromptTracker: Get tracker status information. Returns: - Status information dictionary + 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) + - thread_id: Current thread identifier + - prompt_timeout: Configured prompt timeout in seconds + - cleanup_interval: Configured cleanup interval in seconds """ with self.lock: active_count = len(self.active_prompts) @@ -270,16 +362,30 @@ class PromptTracker: class PromptExecutionContext: - """Context manager for prompt executions.""" + """Context manager for prompt executions. + + This context manager provides automatic prompt lifecycle management, + ensuring prompt context is properly set and cleaned up. While the current + implementation doesn't automatically clear prompts on exit (to allow for + images generated after prompt execution), it provides a clean interface + for prompt management. + + Example: + with PromptExecutionContext(tracker, "beautiful landscape") as exec_id: + # Generate images here + # Images will be linked to this prompt context + pass + # Prompt context remains active for timeout period + """ def __init__(self, prompt_tracker: PromptTracker, prompt_text: str, **kwargs): """ Initialize execution context. Args: - prompt_tracker: PromptTracker instance - prompt_text: The prompt text - **kwargs: Additional prompt metadata + prompt_tracker: PromptTracker instance to use for tracking + prompt_text: The prompt text to track + **kwargs: Additional prompt metadata to include in the context """ self.prompt_tracker = prompt_tracker self.prompt_text = prompt_text @@ -287,7 +393,13 @@ class PromptExecutionContext: self.execution_id = None def __enter__(self): - """Enter the execution context.""" + """Enter the execution context. + + Sets the current prompt in the tracker and returns the execution ID. + + Returns: + Unique execution ID for this prompt + """ self.execution_id = self.prompt_tracker.set_current_prompt( self.prompt_text, self.additional_data @@ -295,7 +407,17 @@ class PromptExecutionContext: return self.execution_id def __exit__(self, exc_type, exc_val, exc_tb): - """Exit the execution context.""" + """Exit the execution context. + + Currently does not clear the prompt context to allow images to be + generated after prompt execution completes. The prompt will be cleaned + up automatically after the timeout period. + + Args: + exc_type: Exception type (if an exception occurred) + exc_val: Exception value (if an exception occurred) + exc_tb: Exception traceback (if an exception occurred) + """ # Don't clear immediately - let the timeout handle it # This allows images to be generated after the prompt execution completes pass \ No newline at end of file diff --git a/utils/validators.py b/utils/validators.py index 9cc9fd1..0d369ea 100644 --- a/utils/validators.py +++ b/utils/validators.py @@ -1,5 +1,33 @@ -""" -Input validation utilities for KikoTextEncode. +"""Input validation utilities for PromptManager. + +This module provides comprehensive input validation and sanitization functions +for the PromptManager system. It ensures data integrity and security by validating +user inputs before they are processed or stored in the database. + +Validation functions include: +- Prompt text validation (length limits, type checking) +- Rating validation (1-5 scale with None support) +- Tag validation and parsing (comma-separated strings or lists) +- Category validation (optional string fields) +- Workflow name validation +- Input sanitization and cleaning utilities + +All validation functions follow a consistent pattern: +- Type checking with descriptive error messages +- Reasonable limits to prevent abuse +- Support for None/optional values where appropriate +- Raise ValueError with clear messages on validation failures + +Typical usage: + from utils.validators import validate_prompt_text, sanitize_input + + try: + validate_prompt_text(user_input) + clean_text = sanitize_input(user_input) + # Process the validated and cleaned input + except ValueError as e: + # Handle validation error with user-friendly message + print(f"Invalid input: {e}") """ import re @@ -10,14 +38,20 @@ def validate_prompt_text(text: str) -> bool: """ Validate prompt text input. + Ensures the prompt text is a valid string with reasonable length limits. + Empty or whitespace-only strings are rejected. + Args: text: The prompt text to validate Returns: - bool: True if valid, False otherwise + True if the text passes validation Raises: - ValueError: If text is invalid with descriptive message + ValueError: If text is invalid with descriptive message including: + - Not a string type + - Empty or whitespace-only + - Exceeds maximum length (10,000 characters) """ if not isinstance(text, str): raise ValueError("Prompt text must be a string") @@ -35,14 +69,18 @@ def validate_rating(rating: Optional[int]) -> bool: """ Validate rating input. + Validates rating values on a 1-5 scale, with None allowed for unrated prompts. + Args: - rating: The rating to validate (1-5 or None) + rating: The rating to validate (1-5 scale or None for no rating) Returns: - bool: True if valid, False otherwise + True if the rating is valid Raises: - ValueError: If rating is invalid + ValueError: If rating is invalid: + - Not an integer (when not None) + - Outside the 1-5 range """ if rating is None: return True @@ -60,14 +98,22 @@ def validate_tags(tags: Union[str, List[str], None]) -> bool: """ Validate tags input. + Accepts tags as comma-separated string, list of strings, or None. + Validates each tag for length and character restrictions. + Args: - tags: Tags as string, list, or None + tags: Tags as comma-separated string, list of strings, or None Returns: - bool: True if valid, False otherwise + True if all tags are valid Raises: - ValueError: If tags are invalid + ValueError: If tags are invalid: + - Wrong input type (not string, list, or None) + - Individual tag is empty or only whitespace + - Individual tag exceeds 50 characters + - Tag contains invalid characters (non-alphanumeric, spaces, hyphens, underscores) + - More than 20 tags provided """ if tags is None: return True @@ -104,14 +150,20 @@ def validate_category(category: Optional[str]) -> bool: """ Validate category input. + Validates optional category strings with reasonable length limits + and character restrictions. + Args: - category: The category to validate + category: The category string to validate (None allowed for no category) Returns: - bool: True if valid, False otherwise + True if the category is valid or None Raises: - ValueError: If category is invalid + ValueError: If category is invalid: + - Not a string type (when not None) + - Exceeds 100 characters + - Contains invalid characters (non-alphanumeric, spaces, hyphens, underscores) """ if category is None: return True @@ -137,14 +189,19 @@ def validate_workflow_name(workflow_name: Optional[str]) -> bool: """ Validate workflow name input. + Validates optional workflow name strings with generous length limits + to accommodate descriptive workflow names. + Args: - workflow_name: The workflow name to validate + workflow_name: The workflow name to validate (None allowed for no workflow) Returns: - bool: True if valid, False otherwise + True if the workflow name is valid or None Raises: - ValueError: If workflow name is invalid + ValueError: If workflow name is invalid: + - Not a string type (when not None) + - Exceeds 200 characters """ if workflow_name is None: return True @@ -166,11 +223,20 @@ def sanitize_input(text: str) -> str: """ Sanitize text input by removing potentially harmful content. + Cleans text input by removing control characters, normalizing whitespace, + and limiting excessive empty lines. Preserves the semantic content while + ensuring safe storage and display. + Args: - text: The text to sanitize + text: The text string to sanitize Returns: - str: Sanitized text + Sanitized text string with: + - Null bytes and control characters removed + - Normalized line endings (\r\n and \r converted to \n) + - Trimmed whitespace on each line + - Limited consecutive empty lines (maximum 2) + - Overall trimmed result """ if not isinstance(text, str): return "" @@ -202,11 +268,15 @@ def parse_tags_string(tags_string: str) -> List[str]: """ Parse a comma-separated tags string into a clean list. + Converts comma-separated tag strings into a clean, deduplicated list + of tags. Each tag is sanitized and trimmed. + Args: - tags_string: Comma-separated tags string + tags_string: Comma-separated tags string (e.g., "tag1, tag2, tag3") Returns: - List[str]: Cleaned list of unique tags + List of unique, cleaned tag strings. Empty input returns empty list. + Limited to maximum 20 tags to prevent abuse. """ if not tags_string or not isinstance(tags_string, str): return []