diff --git a/.gitignore b/.gitignore index eea9d96..9248013 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,9 @@ docs/ logs/ standalone_tagger.py prompts.db +prompts.db-shm +prompts.db-wal +prompts.db-journal # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/database/operations.py b/database/operations.py index 588f35b..ad9b6ae 100644 --- a/database/operations.py +++ b/database/operations.py @@ -187,6 +187,7 @@ def search_prompts( rating_max: Optional[int] = None, date_from: Optional[str] = None, date_to: Optional[str] = None, + folder: Optional[str] = None, limit: int = 100, offset: int = 0, ) -> List[Dict[str, Any]]: @@ -201,6 +202,7 @@ def search_prompts( rating_max: Maximum rating filter date_from: Start date filter (ISO format) date_to: End date filter (ISO format) + folder: Filter by subfolder name in generated image paths limit: Maximum number of results offset: Number of results to skip @@ -227,6 +229,26 @@ def search_prompts( ) params.append(tag) + if folder: + # Escape LIKE wildcards in the folder name + safe_folder = ( + folder.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + ) + query_parts.append( + "AND prompts.id IN (" + " SELECT DISTINCT prompt_id FROM generated_images" + " WHERE image_path LIKE ? ESCAPE '\\'" + " OR image_path LIKE ? ESCAPE '\\'" + " OR image_path LIKE ? ESCAPE '\\'" + ")" + ) + # Pattern 1: folder between forward slashes (Unix + normalized Windows) + params.append(f"%/{safe_folder}/%") + # Pattern 2: folder between backslashes (Windows native) + params.append(f"%\\\\{safe_folder}\\\\%") + # Pattern 3: folder at start of a relative path + params.append(f"{safe_folder}/%") + if rating_min is not None: query_parts.append("AND rating >= ?") params.append(rating_min) @@ -435,6 +457,46 @@ def get_all_categories(self) -> List[str]: ) return [row["category"] for row in cursor.fetchall()] + def get_prompt_subfolders(self, root_dirs: Optional[List[str]] = None) -> List[str]: + """ + Get distinct subfolder paths from generated_images. + + Extracts the directory portion of image_path, made relative to + root_dirs if provided. Returns sorted unique folder names. + """ + with self.model.get_connection() as conn: + cursor = conn.execute( + "SELECT DISTINCT image_path FROM generated_images " + "WHERE image_path IS NOT NULL AND image_path != ''" + ) + + folders = set() + for row in cursor.fetchall(): + image_path = row["image_path"] + parent = os.path.dirname(image_path) + if not parent: + continue + + made_relative = False + if root_dirs: + for root in root_dirs: + try: + rel = os.path.relpath(parent, root) + if not rel.startswith(".."): + if rel != ".": + folders.add(rel) + made_relative = True + break + except ValueError: + continue + + if not made_relative: + basename = os.path.basename(parent) + if basename: + folders.add(basename) + + return sorted(folders) + def get_all_tags(self) -> List[str]: """ Get all unique tags that are in use (linked to at least one prompt). diff --git a/py/api/__init__.py b/py/api/__init__.py index 723c78c..8d7af45 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -123,8 +123,7 @@ def __init__(self): self.db = PromptDatabase() self._cached_output_dir = None # Lazy-cached by _find_comfyui_output_dir() self._html_cache = {} # Cached HTML file contents keyed by path - self._gallery_cache = None # Cached gallery file listing (Fix 2.4) - self._gallery_cache_time = 0 # Timestamp of last cache fill + self._gallery_cache = {} # dict: path_str -> (files, timestamp) self._gallery_cache_ttl = 30 # Cache TTL in seconds # Run cleanup on initialization to remove any existing duplicates @@ -156,8 +155,7 @@ def invalidate_gallery_cache(self): Called by the image monitor when new files are detected. """ - self._gallery_cache = None - self._gallery_cache_time = 0 + self._gallery_cache = {} def add_routes(self, routes): """Register all API routes with the ComfyUI server. @@ -391,47 +389,52 @@ async def serve_js_static(request): # ── Shared utilities used by multiple mixins ────────────────────── - def _enrich_prompt_images(self, prompts): - """Add url and thumbnail_url to each image in prompt results.""" + def _enrich_images(self, images): + """Add url and thumbnail_url to a flat list of image dicts.""" from urllib.parse import quote as url_quote - output_dir = self._find_comfyui_output_dir() - output_path = Path(output_dir) if output_dir else None + output_dirs = self._get_all_output_dirs() - for prompt in prompts: - for image in prompt.get("images", []): - image_path_str = image.get("image_path", "") - if not image_path_str: - continue + for image in images: + image_path_str = image.get("image_path", "") + if not image_path_str: + continue - img_path = Path(image_path_str) + img_path = Path(image_path_str) - # Set fallback url via image ID - if image.get("id"): - image["url"] = f"/prompt_manager/images/{image['id']}/file" + # Set fallback url via image ID + if image.get("id"): + image["url"] = f"/prompt_manager/images/{image['id']}/file" - # Try to compute relative path and thumbnail URL - if output_path: - try: - rel_path = img_path.resolve().relative_to(output_path.resolve()) - image["relative_path"] = str(rel_path) - image["url"] = ( - f"/prompt_manager/images/serve/{url_quote(rel_path.as_posix(), safe='/')}" - ) + # Try each root to compute relative path and thumbnail URL + for output_path in output_dirs: + try: + rel_path = img_path.resolve().relative_to(output_path.resolve()) + image["relative_path"] = str(rel_path) + image["url"] = ( + f"/prompt_manager/images/serve/{url_quote(rel_path.as_posix(), safe='/')}" + ) - # Check for thumbnail - rel_no_ext = rel_path.with_suffix("") - thumb_rel = ( - f"thumbnails/{rel_no_ext.as_posix()}_thumb{rel_path.suffix}" + # Check for thumbnail + rel_no_ext = rel_path.with_suffix("") + thumb_rel = ( + f"thumbnails/{rel_no_ext.as_posix()}_thumb{rel_path.suffix}" + ) + thumb_abs = output_path / thumb_rel + if thumb_abs.exists(): + image["thumbnail_url"] = ( + f"/prompt_manager/images/serve/{url_quote(thumb_rel, safe='/')}" ) - thumb_abs = output_path / thumb_rel - if thumb_abs.exists(): - image["thumbnail_url"] = ( - f"/prompt_manager/images/serve/{url_quote(thumb_rel, safe='/')}" - ) - except (ValueError, RuntimeError): - pass + break # Found matching root, stop searching + except (ValueError, RuntimeError): + continue # Try next root + return images + + def _enrich_prompt_images(self, prompts): + """Add url and thumbnail_url to each image in prompt results.""" + for prompt in prompts: + self._enrich_images(prompt.get("images", [])) return prompts def _clean_nan_recursive(self, obj): @@ -550,6 +553,29 @@ def _find_comfyui_output_dir(self): return None + def _get_all_output_dirs(self): + """Get all configured output directories, falling back to auto-detect. + + Returns: + List[Path]: Valid output directory paths, possibly empty. + """ + from ..config import GalleryConfig + + output_dirs = [] + if GalleryConfig.MONITORING_DIRECTORIES: + for d in GalleryConfig.MONITORING_DIRECTORIES: + p = Path(d).resolve() + if p.is_dir(): + output_dirs.append(p) + + # Fallback to auto-detect if no configured dirs are valid + if not output_dirs: + fallback = self._find_comfyui_output_dir() + if fallback: + output_dirs.append(Path(fallback)) + + return output_dirs + def _extract_comfyui_metadata(self, image_path): """Extract ComfyUI workflow metadata from PNG image files.""" try: diff --git a/py/api/admin.py b/py/api/admin.py index 072deb7..7c0ddbd 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -371,6 +371,9 @@ async def get_settings(self, request): "settings": { "result_timeout": PromptManagerConfig.RESULT_TIMEOUT, "webui_display_mode": PromptManagerConfig.WEBUI_DISPLAY_MODE, + "gallery_root_paths": list( + GalleryConfig.MONITORING_DIRECTORIES + ), "gallery_root_path": ( GalleryConfig.MONITORING_DIRECTORIES[0] if GalleryConfig.MONITORING_DIRECTORIES @@ -400,8 +403,68 @@ async def save_settings(self, request): if "webui_display_mode" in data: PromptManagerConfig.WEBUI_DISPLAY_MODE = data["webui_display_mode"] - # Handle gallery root path - if "gallery_root_path" in data: + # Blocked system directories (shared by both path handlers) + blocked = [ + "/etc", + "/usr", + "/bin", + "/sbin", + "/boot", + "/proc", + "/sys", + "/dev", + "/var/log", + "/root", + "C:\\Windows", + "C:\\Program Files", + ] + + # Handle gallery root paths (array — preferred) + if "gallery_root_paths" in data: + new_paths = data["gallery_root_paths"] + if not isinstance(new_paths, list): + return web.json_response( + { + "success": False, + "error": "gallery_root_paths must be a list", + }, + status=400, + ) + + validated_paths = [] + for path_str in new_paths: + path_str = path_str.strip() + if not path_str: + continue + resolved = Path(path_str).resolve() + if not resolved.is_dir(): + return web.json_response( + { + "success": False, + "error": f"Path does not exist or is not a directory: {path_str}", + }, + status=400, + ) + for b in blocked: + if str(resolved).startswith(b): + return web.json_response( + { + "success": False, + "error": f"Cannot use system directory: {path_str}", + }, + status=400, + ) + validated_paths.append(path_str) + + old_paths = list(GalleryConfig.MONITORING_DIRECTORIES) + if validated_paths != old_paths: + GalleryConfig.MONITORING_DIRECTORIES = validated_paths + self._cached_output_dir = None + self._gallery_cache = {} + restart_required = True + + elif "gallery_root_path" in data: + # Backward compat: single path string new_path = data["gallery_root_path"].strip() old_path = ( GalleryConfig.MONITORING_DIRECTORIES[0] @@ -411,9 +474,7 @@ async def save_settings(self, request): if new_path != old_path: if new_path: - from pathlib import Path as _Path - - resolved = _Path(new_path).resolve() + resolved = Path(new_path).resolve() if not resolved.is_dir(): return web.json_response( { @@ -422,20 +483,6 @@ async def save_settings(self, request): }, status=400, ) - blocked = [ - "/etc", - "/usr", - "/bin", - "/sbin", - "/boot", - "/proc", - "/sys", - "/dev", - "/var/log", - "/root", - "C:\\Windows", - "C:\\Program Files", - ] for b in blocked: if str(resolved).startswith(b): return web.json_response( @@ -448,10 +495,8 @@ async def save_settings(self, request): GalleryConfig.MONITORING_DIRECTORIES = [new_path] else: GalleryConfig.MONITORING_DIRECTORIES = [] - # Invalidate caches so next lookup uses new config self._cached_output_dir = None - self._gallery_cache = None - self._gallery_cache_time = 0 + self._gallery_cache = {} restart_required = True # Save to config file for persistence @@ -599,10 +644,11 @@ async def run_diagnostics(self, request): # Check image monitor status try: - from ...utils.image_monitor import _monitor_instance + from ...utils import image_monitor as im_mod - if _monitor_instance is not None: - monitor_status = _monitor_instance.get_status() + monitor = im_mod._monitor_instance + if monitor is not None: + monitor_status = monitor.get_status() results["image_monitor"] = { "status": ( "ok" if monitor_status.get("observer_alive") else "error" @@ -978,29 +1024,60 @@ async def restore_database(self, request): async def scan_images(self, request): """Scan ComfyUI output images for prompt metadata and add them to the database.""" + BATCH_SIZE = 50 # Files per executor call for metadata extraction + + def _collect_media_files(output_dirs): + """Collect all media files from output directories (blocking I/O).""" + image_extensions = [".png", ".jpg", ".jpeg", ".webp", ".gif"] + video_extensions = [".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"] + media_extensions = image_extensions + video_extensions + all_files = [] + seen = set() + for output_dir in output_dirs: + for ext in media_extensions: + for pattern in [f"*{ext}", f"*{ext.upper()}"]: + for f in output_dir.rglob(pattern): + if "thumbnails" not in f.parts: + norm = str(f).lower() + if norm not in seen: + seen.add(norm) + all_files.append(f) + return all_files + + def _extract_batch_metadata(file_batch): + """Extract metadata from a batch of files in one executor call.""" + results = [] + for f in file_batch: + try: + meta = self._extract_comfyui_metadata(str(f)) + results.append((f, meta)) + except Exception: + results.append((f, {})) + return results + async def stream_response(): try: self.logger.info("Starting image scan operation") - self.logger.info("Starting scan (timer clearing not implemented yet)") - output_dir = self._find_comfyui_output_dir() - if not output_dir: - self.logger.error("ComfyUI output directory not found") - yield f"data: {json.dumps({'type': 'error', 'message': 'ComfyUI output directory not found'})}\n\n" + output_dirs = self._get_all_output_dirs() + if not output_dirs: + self.logger.error("No output directories found") + yield f"data: {json.dumps({'type': 'error', 'message': 'No output directories found. Configure scan directories in Settings.'})}\n\n" return - yield f"data: {json.dumps({'type': 'progress', 'progress': 0, 'status': 'Scanning for PNG files...', 'processed': 0, 'found': 0})}\n\n" + dir_names = [str(d) for d in output_dirs] + yield f"data: {json.dumps({'type': 'progress', 'progress': 0, 'status': f'Scanning {len(output_dirs)} directory(ies) for media files...', 'processed': 0, 'found': 0})}\n\n" - png_files = await self._run_in_executor( - lambda: list(Path(output_dir).rglob("*.png")) + media_files = await self._run_in_executor( + _collect_media_files, output_dirs ) - total_files = len(png_files) + total_files = len(media_files) if total_files == 0: - yield f"data: {json.dumps({'type': 'complete', 'processed': 0, 'found': 0, 'added': 0})}\n\n" + yield f"data: {json.dumps({'type': 'complete', 'processed': 0, 'found': 0, 'added': 0, 'linked': 0, 'directories': dir_names})}\n\n" return - yield f"data: {json.dumps({'type': 'progress', 'progress': 5, 'status': f'Found {total_files} PNG files to process...', 'processed': 0, 'found': 0})}\n\n" + yield f"data: {json.dumps({'type': 'progress', 'progress': 5, 'status': f'Found {total_files} media files to process...', 'processed': 0, 'found': 0})}\n\n" processed_count = 0 found_count = 0 @@ -1018,133 +1095,97 @@ async def stream_response(): sys.path.insert(0, current_dir) from utils.hashing import generate_prompt_hash - for i, png_file in enumerate(png_files): - try: - metadata = await self._run_in_executor( - self._extract_comfyui_metadata, str(png_file) - ) - processed_count += 1 + # Process files in batches for performance + for batch_start in range(0, total_files, BATCH_SIZE): + batch = media_files[batch_start : batch_start + BATCH_SIZE] - if metadata: - self.logger.debug( - f"Found metadata in {os.path.basename(png_file)}: {list(metadata.keys())}" - ) + # Extract metadata for entire batch in one executor call + batch_results = await self._run_in_executor( + _extract_batch_metadata, batch + ) + + for media_file, metadata in batch_results: + try: + processed_count += 1 + + if not metadata: + continue parsed_data = self._parse_comfyui_prompt(metadata) - self.logger.debug( - f"Parsed data keys: {list(parsed_data.keys())}, has prompt: {bool(parsed_data.get('prompt'))}, has parameters: {bool(parsed_data.get('parameters'))}" - ) - if parsed_data.get("prompt") or parsed_data.get( - "parameters" + if not ( + parsed_data.get("prompt") + or parsed_data.get("parameters") ): - found_count += 1 - prompt_text = self._extract_readable_prompt(parsed_data) + continue - if prompt_text: - self.logger.debug( - f"Found prompt in {os.path.basename(png_file)} (type: {type(prompt_text)}): {str(prompt_text)[:100]}..." - ) - else: - self.logger.debug( - f"No readable prompt found in {os.path.basename(png_file)}, parsed_data keys: {list(parsed_data.keys())}" - ) + found_count += 1 + prompt_text = self._extract_readable_prompt(parsed_data) - if prompt_text and not isinstance(prompt_text, str): - self.logger.debug( - f"Converting prompt_text from {type(prompt_text)} to string" - ) - prompt_text = str(prompt_text) + if prompt_text and not isinstance(prompt_text, str): + prompt_text = str(prompt_text) - if prompt_text and prompt_text.strip(): - try: - prompt_hash = generate_prompt_hash( - prompt_text.strip() - ) - self.logger.debug( - f"Generated hash for prompt: {prompt_hash[:16]}..." - ) + if not (prompt_text and prompt_text.strip()): + continue - existing = await self._run_in_executor( - self.db.get_prompt_by_hash, prompt_hash - ) - if existing: - self.logger.debug( - f"Found existing prompt ID {existing['id']} for image {os.path.basename(png_file)}" - ) - try: - await self._run_in_executor( - self.db.link_image_to_prompt, - existing["id"], - str(png_file), - ) - linked_count += 1 - self.logger.debug( - f"Linked image {os.path.basename(png_file)} to existing prompt {existing['id']}" - ) - except Exception as e: - self.logger.error( - f"Failed to link image {png_file} to existing prompt: {e}" - ) - else: - self.logger.debug( - f"Saving new prompt from {os.path.basename(png_file)}" - ) - prompt_id = await self._run_in_executor( - self.db.save_prompt, - prompt_text.strip(), - "scanned", - ["auto-scanned"], - None, - f"Auto-scanned from {os.path.basename(png_file)}", - prompt_hash, - ) - - if prompt_id: - added_count += 1 - self.logger.info( - f"Successfully saved new prompt with ID {prompt_id} from {os.path.basename(png_file)}" - ) - try: - await self._run_in_executor( - self.db.link_image_to_prompt, - prompt_id, - str(png_file), - ) - self.logger.debug( - f"Linked image {os.path.basename(png_file)} to new prompt {prompt_id}" - ) - except Exception as e: - self.logger.error( - f"Failed to link image {png_file} to new prompt: {e}" - ) - else: - self.logger.error( - f"Failed to save prompt from {os.path.basename(png_file)} - no ID returned" - ) + prompt_hash = generate_prompt_hash(prompt_text.strip()) + existing = await self._run_in_executor( + self.db.get_prompt_by_hash, prompt_hash + ) + if existing: + try: + await self._run_in_executor( + self.db.link_image_to_prompt, + existing["id"], + str(media_file), + ) + linked_count += 1 + except Exception as e: + self.logger.error( + f"Failed to link {media_file.name} to existing prompt: {e}" + ) + else: + prompt_id = await self._run_in_executor( + self.db.save_prompt, + prompt_text.strip(), + "scanned", + ["auto-scanned"], + None, + f"Auto-scanned from {media_file.name}", + prompt_hash, + ) + if prompt_id: + added_count += 1 + try: + await self._run_in_executor( + self.db.link_image_to_prompt, + prompt_id, + str(media_file), + ) except Exception as e: self.logger.error( - f"Failed to save prompt from {png_file}: {e}" + f"Failed to link {media_file.name} to new prompt: {e}" ) - # Update progress every 10 files - if i % 10 == 0 or i == total_files - 1: - progress = int((i + 1) / total_files * 100) - status = f"Processing file {i + 1}/{total_files}..." - - yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': status, 'processed': processed_count, 'found': found_count})}\n\n" - - await asyncio.sleep(0.01) + except Exception as e: + self.logger.error( + f"Error processing {media_file.name}: {e}" + ) + continue - except Exception as e: - self.logger.error(f"Error processing {png_file}: {e}") - continue + # Progress update after each batch + progress = int( + min(batch_start + len(batch), total_files) / total_files * 100 + ) + yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Processing file {min(batch_start + len(batch), total_files)}/{total_files}...', 'processed': processed_count, 'found': found_count})}\n\n" + await asyncio.sleep(0) self.logger.info( - f"Scan completed: processed={processed_count}, found={found_count}, new_prompts_added={added_count}, images_linked_to_existing={linked_count}" + f"Scan completed: processed={processed_count}, found={found_count}, " + f"new_prompts_added={added_count}, images_linked_to_existing={linked_count}" ) - yield f"data: {json.dumps({'type': 'complete', 'processed': processed_count, 'found': found_count, 'added': added_count, 'linked': linked_count})}\n\n" + yield f"data: {json.dumps({'type': 'complete', 'processed': processed_count, 'found': found_count, 'added': added_count, 'linked': linked_count, 'directories': dir_names})}\n\n" except Exception as e: self.logger.exception("Scan error") diff --git a/py/api/images.py b/py/api/images.py index 1dd7411..0acdb4e 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -68,6 +68,10 @@ async def generate_thumbnails_progress_route(request): async def clear_thumbnails_route(request): return await self.clear_thumbnails(request) + @routes.get("/prompt_manager/gallery/subfolders") + async def get_gallery_subfolders_route(request): + return await self.get_gallery_subfolders(request) + async def get_prompt_images(self, request): """Get all images for a specific prompt.""" try: @@ -77,6 +81,9 @@ async def get_prompt_images(self, request): # Clean up any NaN values that cause JSON parsing errors (recursive) cleaned_images = [self._clean_nan_recursive(image) for image in images] + # Add url and thumbnail_url so the frontend can serve them + self._enrich_images(cleaned_images) + # Additional fallback: convert to JSON string and clean NaN values manually try: response_data = {"success": True, "images": cleaned_images} @@ -145,7 +152,10 @@ async def search_images(self, request): return web.json_response({"success": False, "error": str(e)}, status=500) def _scan_gallery_files_sync(self, output_path): - """Scan output directory for media files (blocking I/O, run in executor).""" + """Scan output directory for media files (blocking I/O, run in executor). + + Returns list of (path, mtime) tuples sorted by mtime descending. + """ image_extensions = [".png", ".jpg", ".jpeg", ".webp", ".gif"] video_extensions = [".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"] media_extensions = image_extensions + video_extensions @@ -159,62 +169,73 @@ def _scan_gallery_files_sync(self, output_path): normalized_path = str(media_path).lower() if normalized_path not in seen_paths: seen_paths.add(normalized_path) - all_images.append(media_path) + try: + mtime = media_path.stat().st_mtime + all_images.append((media_path, mtime)) + except OSError: + continue - # Sort by modification time (newest first) - all_images.sort(key=lambda x: x.stat().st_mtime, reverse=True) + all_images.sort(key=lambda x: x[1], reverse=True) return all_images async def _get_gallery_files(self, output_path): - """Get gallery files with TTL cache. Invalidated by image monitor.""" + """Get gallery files with per-directory TTL cache.""" now = _time.monotonic() - if ( - self._gallery_cache is not None - and (now - self._gallery_cache_time) < self._gallery_cache_ttl - ): - return self._gallery_cache - - all_images = await self._run_in_executor( - self._scan_gallery_files_sync, output_path - ) - self._gallery_cache = all_images - self._gallery_cache_time = now - return all_images + key = str(output_path) + cached = self._gallery_cache.get(key) + if cached and (now - cached[1]) < self._gallery_cache_ttl: + return cached[0] + + files = await self._run_in_executor(self._scan_gallery_files_sync, output_path) + self._gallery_cache[key] = (files, now) + return files async def get_output_images(self, request): - """Get all images from ComfyUI output folder.""" + """Get all images from ComfyUI output folder(s).""" try: from urllib.parse import quote - # Find ComfyUI output directory - output_dir = self._find_comfyui_output_dir() - if not output_dir: + output_dirs = self._get_all_output_dirs() + if not output_dirs: return web.json_response( { "success": False, - "error": "ComfyUI output directory not found", + "error": "No output directories found", "images": [], - } + }, ) - # Get pagination parameters limit = int(request.query.get("limit", 100)) offset = int(request.query.get("offset", 0)) + subfolder = request.query.get("subfolder", "").strip() + + # Collect (path, mtime, root) from all directories + all_images = [] + for output_path in output_dirs: + dir_images = await self._get_gallery_files(output_path) + for img_path, mtime in dir_images: + all_images.append((img_path, mtime, output_path)) + + # Sort combined results by mtime (newest first) — no .stat() calls + all_images.sort(key=lambda x: x[1], reverse=True) + + # Apply subfolder filter if provided + if subfolder: + filtered = [] + for img_path, mtime, root in all_images: + rel_dir = str(img_path.relative_to(root).parent) + if rel_dir == subfolder or rel_dir.startswith(subfolder + os.sep): + filtered.append((img_path, mtime, root)) + all_images = filtered + + total = len(all_images) + paginated = all_images[offset : offset + limit] - output_path = Path(output_dir) - thumbnails_dir = output_path / "thumbnails" video_extensions = [".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"] - # Use cached file listing (Fix 2.4) - all_images = await self._get_gallery_files(output_path) - - # Apply pagination - paginated_images = all_images[offset : offset + limit] - - # Format media data in executor (stat calls are blocking) def _format_page(): images = [] - for media_path in paginated_images: + for media_path, mtime, output_path in paginated: try: stat = media_path.stat() rel_path = media_path.relative_to(output_path) @@ -223,6 +244,7 @@ def _format_page(): media_type = "video" if is_video else "image" thumbnail_url = None + thumbnails_dir = output_path / "thumbnails" if thumbnails_dir.exists(): thumbnail_ext = ".jpg" if is_video else extension rel_path_no_ext = rel_path.with_suffix("") @@ -237,6 +259,7 @@ def _format_page(): "filename": media_path.name, "path": str(media_path), "relative_path": str(rel_path), + "root_dir": str(output_path), "url": f"/prompt_manager/images/serve/{rel_path.as_posix()}", "thumbnail_url": thumbnail_url, "size": stat.st_size, @@ -257,18 +280,15 @@ def _format_page(): { "success": True, "images": images, - "total": len(all_images), + "total": total, "offset": offset, "limit": limit, - "has_more": offset + limit < len(all_images), + "has_more": offset + limit < total, } ) - except Exception as e: - self.logger.error(f"Get output images error: {e}") - return web.json_response( - {"success": False, "error": str(e), "images": []}, status=500 - ) + self.logger.error(f"Output images error: {e}") + return web.json_response({"success": False, "error": str(e)}, status=500) async def serve_image(self, request): """Serve the actual image file using streamed FileResponse.""" @@ -283,11 +303,13 @@ async def serve_image(self, request): image_path = Path(image["image_path"]).resolve() - # Validate path is within the ComfyUI output directory - output_dir = self._find_comfyui_output_dir() - if output_dir: - output_path = Path(output_dir).resolve() - if not image_path.is_relative_to(output_path): + # Validate path is within any configured output directory + output_dirs = self._get_all_output_dirs() + if output_dirs: + allowed = any( + image_path.is_relative_to(d.resolve()) for d in output_dirs + ) + if not allowed: return web.json_response( {"success": False, "error": "Access denied"}, status=403 ) @@ -314,43 +336,54 @@ async def serve_output_image(self, request): try: filepath = request.match_info["filepath"] - # Find ComfyUI output directory - output_dir = self._find_comfyui_output_dir() - if not output_dir: + # Search all configured output directories + output_dirs = self._get_all_output_dirs() + if not output_dirs: return web.json_response( {"success": False, "error": "ComfyUI output directory not found"}, status=404, ) - # Construct full image path - image_path = Path(output_dir) / filepath + # Try each root directory until the file is found + for output_path in output_dirs: + image_path = (output_path / filepath).resolve() - # Security check: make sure the path is within the output directory - try: - image_path = image_path.resolve() - output_path = Path(output_dir).resolve() - if not image_path.is_relative_to(output_path): - return web.json_response( - {"success": False, "error": "Access denied"}, status=403 - ) - except Exception: - return web.json_response( - {"success": False, "error": "Invalid file path"}, status=400 - ) + # Security check: must be within this output directory + if not image_path.is_relative_to(output_path.resolve()): + continue - if not image_path.exists(): - return web.json_response( - {"success": False, "error": "Image file not found"}, status=404 - ) + if image_path.exists(): + response = web.FileResponse(image_path) + response.headers["Cache-Control"] = "public, max-age=3600" + return response - response = web.FileResponse(image_path) - response.headers["Cache-Control"] = "public, max-age=3600" - return response + return web.json_response( + {"success": False, "error": "Image file not found"}, status=404 + ) except Exception as e: self.logger.error(f"Serve output image error: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) + async def get_gallery_subfolders(self, request): + """Get distinct subfolders from gallery output directories.""" + try: + output_dirs = self._get_all_output_dirs() + subfolders = set() + for output_path in output_dirs: + files = await self._get_gallery_files(output_path) + for f, _mtime in files: + rel_dir = str(f.relative_to(output_path).parent) + if rel_dir and rel_dir != ".": + subfolders.add(rel_dir) + + return web.json_response( + {"success": True, "subfolders": sorted(subfolders)} + ) + except Exception as e: + self.logger.error(f"Subfolders error: {e}") + return web.json_response({"success": False, "error": str(e)}, status=500) + async def generate_thumbnails(self, request): """Generate thumbnails for all images and videos in the ComfyUI output directory.""" try: diff --git a/py/api/prompts.py b/py/api/prompts.py index 4157e6d..1de4c32 100644 --- a/py/api/prompts.py +++ b/py/api/prompts.py @@ -113,6 +113,10 @@ async def bulk_add_tags_route(request): async def bulk_set_category_route(request): return await self.bulk_set_category(request) + @routes.get("/prompt_manager/subfolders") + async def get_subfolders_route(request): + return await self.get_subfolders(request) + # Export functionality @routes.get("/prompt_manager/export") async def export_prompts_route(request): @@ -127,6 +131,8 @@ async def search_prompts(self, request): min_rating = request.query.get("min_rating", 0) limit = int(request.query.get("limit", 50)) + folder = request.query.get("folder", "").strip() or None + tags = None if tags_str: tags = [tag.strip() for tag in tags_str.split(",") if tag.strip()] @@ -143,6 +149,7 @@ async def search_prompts(self, request): tags=tags, rating_min=min_rating, limit=limit, + folder=folder, ) self._enrich_prompt_images(results) @@ -157,6 +164,20 @@ async def search_prompts(self, request): status=500, ) + async def get_subfolders(self, request): + """Get distinct subfolder values derived from generated image paths.""" + try: + from ..config import GalleryConfig + + root_dirs = list(GalleryConfig.MONITORING_DIRECTORIES) or None + subfolders = await self._run_in_executor( + self.db.get_prompt_subfolders, root_dirs + ) + return web.json_response({"success": True, "subfolders": subfolders}) + except Exception as e: + self.logger.error(f"Subfolders error: {e}") + return web.json_response({"success": False, "error": str(e)}, status=500) + async def get_recent_prompts(self, request): """Retrieve recently created prompts with pagination support.""" try: diff --git a/pyproject.toml b/pyproject.toml index 7f234cb..d782192 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "promptmanager" description = "A powerful ComfyUI custom node that extends the standard text encoder with persistent prompt storage, advanced search capabilities, and an automatic image gallery system using SQLite." -version = "3.1.6" +version = "3.2.0" license = {file = "LICENSE"} dependencies = ["# Core dependencies for PromptManager", "# Note: Most dependencies are already included with ComfyUI", "# Already included with Python standard library:", "# - sqlite3", "# - hashlib", "# - json", "# - datetime", "# - os", "# - typing", "# - threading", "# - uuid", "# Required for gallery functionality:", "watchdog>=2.1.0 # For file system monitoring", "Pillow>=8.0.0 # For image metadata extraction (usually included with ComfyUI)", "# Optional dependencies for enhanced search functionality:", "# fuzzywuzzy[speedup]>=0.18.0 # For fuzzy string matching (optional)", "# sqlalchemy>=1.4.0 # For advanced ORM features (optional)", "# Development dependencies (optional):", "# pytest>=6.0.0 # For running tests", "# black>=22.0.0 # For code formatting", "# flake8>=4.0.0 # For linting", "# mypy>=0.910 # For type checking"] diff --git a/web/admin.html b/web/admin.html index ff17d2e..05282d4 100644 --- a/web/admin.html +++ b/web/admin.html @@ -78,7 +78,7 @@

-
+
+
+ + +
Settings

Gallery Settings

- - -

Custom path to scan for images. Leave empty to use ComfyUI's default output directory.

+ +
+ +
+ +

Directories to scan for images. Leave empty to use ComfyUI default output directory.

@@ -1224,6 +1239,52 @@

Image Already Tagged

+ + +