"""Image and gallery API routes for PromptManager.""" import asyncio import json import os import re import time as _time import urllib.parse from pathlib import Path from aiohttp import web from PIL import Image class ImageRoutesMixin: """Mixin providing image and gallery-related API endpoints.""" def _register_image_routes(self, routes): @routes.get("/prompt_manager/prompts/{prompt_id}/images") async def get_prompt_images_route(request): return await self.get_prompt_images(request) @routes.get("/prompt_manager/images/recent") async def get_recent_images_route(request): return await self.get_recent_images(request) @routes.get("/prompt_manager/images/all") async def get_all_images_route(request): return await self.get_all_images(request) @routes.get("/prompt_manager/images/search") async def search_images_route(request): return await self.search_images(request) @routes.get("/prompt_manager/images/output") async def get_output_images_route(request): return await self.get_output_images(request) @routes.get("/prompt_manager/images/{image_id}/file") async def serve_image_route(request): return await self.serve_image(request) @routes.get("/prompt_manager/images/serve/{filepath:.*}") async def serve_output_image_route(request): return await self.serve_output_image(request) @routes.post("/prompt_manager/images/link") async def link_image_route(request): return await self.link_image_to_prompt(request) @routes.get("/prompt_manager/images/prompt/{image_path:.*}") async def get_image_prompt_route(request): return await self.get_image_prompt(request) @routes.delete("/prompt_manager/images/{image_id}") async def delete_image_route(request): return await self.delete_image(request) @routes.post("/prompt_manager/images/generate-thumbnails") async def generate_thumbnails_route(request): return await self.generate_thumbnails(request) @routes.get("/prompt_manager/images/generate-thumbnails/progress") async def generate_thumbnails_progress_route(request): return await self.generate_thumbnails_with_progress(request) @routes.post("/prompt_manager/images/clear-thumbnails") async def clear_thumbnails_route(request): return await self.clear_thumbnails(request) async def get_prompt_images(self, request): """Get all images for a specific prompt.""" try: prompt_id = request.match_info["prompt_id"] images = await self._run_in_executor(self.db.get_prompt_images, prompt_id) # Clean up any NaN values that cause JSON parsing errors (recursive) cleaned_images = [self._clean_nan_recursive(image) for image in images] # Additional fallback: convert to JSON string and clean NaN values manually try: response_data = {"success": True, "images": cleaned_images} # Convert to JSON string json_str = json.dumps(response_data, default=str) # Clean any remaining NaN values with regex json_str = re.sub(r":\s*NaN", ": null", json_str) json_str = re.sub(r"\[\s*NaN\s*\]", "[null]", json_str) json_str = re.sub(r",\s*NaN\s*,", ", null,", json_str) json_str = re.sub(r",\s*NaN\s*\]", ", null]", json_str) json_str = re.sub(r"\[\s*NaN\s*,", "[null,", json_str) # Parse back to verify it's valid JSON cleaned_data = json.loads(json_str) return web.json_response(cleaned_data) except Exception as json_error: self.logger.error(f"JSON cleaning error: {json_error}") # Fallback to original response return web.json_response({"success": True, "images": cleaned_images}) except Exception as e: self.logger.error(f"Get prompt images error: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) async def get_recent_images(self, request): """Get recently generated images.""" try: limit = int(request.query.get("limit", 50)) images = await self._run_in_executor(self.db.get_recent_images, limit) return web.json_response({"success": True, "images": images}) except Exception as e: self.logger.error(f"Get recent images error: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) async def get_all_images(self, request): """Get all generated images with linked prompts.""" try: images = await self._run_in_executor(self.db.get_all_images) return web.json_response( {"success": True, "images": images, "count": len(images)} ) except Exception as e: self.logger.error(f"Get all images error: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) async def search_images(self, request): """Search images by prompt text.""" try: query = request.query.get("q", "") if not query: return web.json_response( {"success": False, "error": "Search query required"}, status=400 ) images = await self._run_in_executor(self.db.search_images_by_prompt, query) return web.json_response( {"success": True, "images": images, "query": query} ) except Exception as e: self.logger.error(f"Search images error: {e}") 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).""" image_extensions = [".png", ".jpg", ".jpeg", ".webp", ".gif"] video_extensions = [".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"] media_extensions = image_extensions + video_extensions all_images = [] seen_paths = set() for ext in media_extensions: for pattern in [f"*{ext}", f"*{ext.upper()}"]: for media_path in output_path.rglob(pattern): if "thumbnails" not in media_path.parts: normalized_path = str(media_path).lower() if normalized_path not in seen_paths: seen_paths.add(normalized_path) all_images.append(media_path) # Sort by modification time (newest first) all_images.sort(key=lambda x: x.stat().st_mtime, reverse=True) return all_images async def _get_gallery_files(self, output_path): """Get gallery files with TTL cache. Invalidated by image monitor.""" 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 async def get_output_images(self, request): """Get all images from ComfyUI output folder.""" try: from urllib.parse import quote # Find ComfyUI output directory output_dir = self._find_comfyui_output_dir() if not output_dir: return web.json_response( { "success": False, "error": "ComfyUI output directory not found", "images": [], } ) # Get pagination parameters limit = int(request.query.get("limit", 100)) offset = int(request.query.get("offset", 0)) 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: try: stat = media_path.stat() rel_path = media_path.relative_to(output_path) extension = media_path.suffix.lower() is_video = extension in video_extensions media_type = "video" if is_video else "image" thumbnail_url = None if thumbnails_dir.exists(): thumbnail_ext = ".jpg" if is_video else extension rel_path_no_ext = rel_path.with_suffix("") thumbnail_rel_path = f"thumbnails/{rel_path_no_ext.as_posix()}_thumb{thumbnail_ext}" thumbnail_abs_path = output_path / thumbnail_rel_path if thumbnail_abs_path.exists(): thumbnail_url = f'/prompt_manager/images/serve/{quote(thumbnail_rel_path, safe="/")}' images.append( { "id": str(hash(str(media_path))), "filename": media_path.name, "path": str(media_path), "relative_path": str(rel_path), "url": f"/prompt_manager/images/serve/{rel_path.as_posix()}", "thumbnail_url": thumbnail_url, "size": stat.st_size, "modified_time": stat.st_mtime, "extension": extension, "media_type": media_type, "is_video": is_video, } ) except Exception as e: self.logger.error(f"Error processing media {media_path}: {e}") continue return images images = await self._run_in_executor(_format_page) return web.json_response( { "success": True, "images": images, "total": len(all_images), "offset": offset, "limit": limit, "has_more": offset + limit < len(all_images), } ) 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 ) async def serve_image(self, request): """Serve the actual image file using streamed FileResponse.""" try: image_id = int(request.match_info["image_id"]) image = await self._run_in_executor(self.db.get_image_by_id, image_id) if not image: return web.json_response( {"success": False, "error": "Image not found"}, status=404 ) 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): return web.json_response( {"success": False, "error": "Access denied"}, status=403 ) if not image_path.exists(): return web.json_response( {"success": False, "error": "Image file not found"}, status=404 ) response = web.FileResponse(image_path) response.headers["Cache-Control"] = "public, max-age=3600" return response except ValueError: return web.json_response( {"success": False, "error": "Invalid image ID"}, status=400 ) except Exception as e: self.logger.error(f"Serve image error: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) async def serve_output_image(self, request): """Serve image file directly from ComfyUI output folder using streamed FileResponse.""" try: filepath = request.match_info["filepath"] # Find ComfyUI output directory output_dir = self._find_comfyui_output_dir() if not output_dir: return web.json_response( {"success": False, "error": "ComfyUI output directory not found"}, status=404, ) # Construct full image path image_path = Path(output_dir) / filepath # 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 ) if not image_path.exists(): return web.json_response( {"success": False, "error": "Image file not found"}, status=404 ) response = web.FileResponse(image_path) response.headers["Cache-Control"] = "public, max-age=3600" return response 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 generate_thumbnails(self, request): """Generate thumbnails for all images and videos in the ComfyUI output directory.""" try: # Get request parameters data = await request.json() quality = data.get("quality", "medium") # Map quality to size size_map = {"low": (150, 150), "medium": (300, 300), "high": (600, 600)} thumbnail_size = size_map.get(quality, (300, 300)) # Find ComfyUI output directory output_dir = self._find_comfyui_output_dir() if not output_dir: return web.json_response( {"success": False, "error": "ComfyUI output directory not found"}, status=404, ) output_path = Path(output_dir) thumbnails_dir = output_path / "thumbnails" # Run the entire thumbnail generation in executor (heavy PIL I/O) result = await self._run_in_executor( self._generate_thumbnails_sync, output_path, thumbnails_dir, thumbnail_size, ) # Invalidate gallery cache since thumbnails changed self.invalidate_gallery_cache() return web.json_response(result) except ImportError: return web.json_response( { "success": False, "error": "PIL (Pillow) library not available. Install with: pip install Pillow", }, status=500, ) except Exception as e: self.logger.error(f"Generate thumbnails error: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) def _generate_thumbnails_sync(self, output_path, thumbnails_dir, thumbnail_size): """Blocking thumbnail generation loop (run in executor).""" import time thumbnails_dir.mkdir(exist_ok=True) self.logger.info("Scanning for media files to generate thumbnails...") image_extensions = [".png", ".jpg", ".jpeg", ".webp", ".gif"] video_extensions = [".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"] media_extensions = image_extensions + video_extensions media_files = [] for root, dirs, files in os.walk(output_path): if "thumbnails" in Path(root).parts: continue for file in files: if any(file.lower().endswith(ext) for ext in media_extensions): media_files.append(Path(root) / file) total_images = len(media_files) self.logger.info(f"Found {total_images} media files to process for thumbnails") if total_images == 0: return { "success": True, "count": 0, "total_images": 0, "message": "No media files found to process", "errors": [], } generated_count = 0 skipped_count = 0 errors = [] start_time = time.time() for i, media_file in enumerate(media_files): try: rel_path = media_file.relative_to(output_path) is_video = any( media_file.name.lower().endswith(ext) for ext in video_extensions ) rel_path_no_ext = rel_path.with_suffix("") if is_video: thumbnail_path = ( thumbnails_dir / f"{rel_path_no_ext.as_posix()}_thumb.jpg" ) else: thumbnail_path = ( thumbnails_dir / f"{rel_path_no_ext.as_posix()}_thumb{rel_path.suffix}" ) if ( thumbnail_path.exists() and thumbnail_path.stat().st_mtime > media_file.stat().st_mtime ): skipped_count += 1 continue thumbnail_path.parent.mkdir(parents=True, exist_ok=True) if is_video: if self._generate_video_thumbnail( media_file, thumbnail_path, thumbnail_size ): generated_count += 1 else: errors.append( f"Failed to generate video thumbnail for {media_file.name}" ) else: with Image.open(media_file) as img: if img.mode in ("RGBA", "LA", "P"): img = img.convert("RGB") img.thumbnail(thumbnail_size, Image.Resampling.LANCZOS) save_kwargs = {"quality": 85, "optimize": True} if thumbnail_path.suffix.lower() == ".png": save_kwargs = {"optimize": True} img.save(thumbnail_path, **save_kwargs) generated_count += 1 if ( generated_count % 100 == 0 or (i + 1) % max(1, total_images // 10) == 0 ): progress = ((i + 1) / total_images) * 100 elapsed = time.time() - start_time rate = (i + 1) / elapsed if elapsed > 0 else 0 eta = ((total_images - i - 1) / rate) if rate > 0 else 0 self.logger.info( f"Thumbnail progress: {i+1}/{total_images} ({progress:.1f}%) - " f"Generated: {generated_count}, Skipped: {skipped_count}, " f"Rate: {rate:.1f} img/s, ETA: {eta:.0f}s" ) except Exception as e: error_msg = ( f"Failed to generate thumbnail for {media_file.name}: {str(e)}" ) errors.append(error_msg) self.logger.warning(error_msg) elapsed_time = time.time() - start_time self.logger.info( f"Thumbnail generation completed: {generated_count} generated, " f"{skipped_count} skipped, {len(errors)} errors in {elapsed_time:.1f}s" ) return { "success": True, "count": generated_count, "skipped": skipped_count, "total_images": total_images, "errors": errors, "thumbnails_path": str(thumbnails_dir), "elapsed_time": round(elapsed_time, 2), "processing_rate": round( (total_images / elapsed_time) if elapsed_time > 0 else 0, 2 ), } async def generate_thumbnails_with_progress(self, request): """Generate thumbnails with Server-Sent Events progress updates.""" try: import time # Parse query parameters quality = request.query.get("quality", "medium") # Map quality to size size_map = {"low": (150, 150), "medium": (300, 300), "high": (600, 600)} thumbnail_size = size_map.get(quality, (300, 300)) # Set up SSE response response = web.StreamResponse( status=200, headers={ "Content-Type": "text/event-stream", "Cache-Control": "no-cache", "Connection": "keep-alive", "Access-Control-Allow-Origin": "*", }, ) await response.prepare(request) async def send_progress(event_type, data): """Send SSE event to client.""" try: message = f"event: {event_type}\ndata: {json.dumps(data)}\n\n" await response.write(message.encode("utf-8")) await asyncio.sleep(0.01) except Exception as e: self.logger.warning(f"Failed to send SSE message: {e}") try: # Find ComfyUI output directory output_dir = self._find_comfyui_output_dir() if not output_dir: await send_progress( "error", {"error": "ComfyUI output directory not found"} ) return response output_path = Path(output_dir) thumbnails_dir = output_path / "thumbnails" thumbnails_dir.mkdir(exist_ok=True) # Send scanning event await send_progress( "status", { "phase": "scanning", "message": f"Scanning {output_path} for images and videos to process...", }, ) self.logger.info( f"Starting thumbnail generation scan in: {output_path}" ) # Find all media files (images and videos) image_extensions = [".png", ".jpg", ".jpeg", ".webp", ".gif"] video_extensions = [ ".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv", ] media_extensions = image_extensions + video_extensions media_files = [] scanned_dirs = 0 for root, dirs, files in os.walk(output_path): if "thumbnails" in Path(root).parts: continue scanned_dirs += 1 if scanned_dirs % 5 == 0: await send_progress( "status", { "phase": "scanning", "message": f"Scanning directories... ({scanned_dirs} checked, {len(media_files)} files found)", }, ) for file in files: if any(file.lower().endswith(ext) for ext in media_extensions): media_files.append(Path(root) / file) self.logger.info( f"Scan complete: Found {len(media_files)} media files in {scanned_dirs} directories" ) total_images = len(media_files) # Count images vs videos for more detail image_count = sum( 1 for f in media_files if any(f.name.lower().endswith(ext) for ext in image_extensions) ) video_count = total_images - image_count await send_progress( "start", { "total_images": total_images, "phase": "processing", "message": f"Found {image_count} images and {video_count} videos to process", "image_count": image_count, "video_count": video_count, }, ) self.logger.info( f"Starting thumbnail generation for {image_count} images and {video_count} videos" ) if total_images == 0: await send_progress( "complete", { "count": 0, "skipped": 0, "total_images": 0, "elapsed_time": 0, "message": "No media files found to process", }, ) return response generated_count = 0 skipped_count = 0 errors = [] start_time = time.time() for i, media_file in enumerate(media_files): try: if not media_file.exists() or not media_file.is_file(): continue is_video = any( media_file.name.lower().endswith(ext) for ext in video_extensions ) rel_path = media_file.relative_to(output_path) rel_path_no_ext = rel_path.with_suffix("") if is_video: thumbnail_path = ( thumbnails_dir / f"{rel_path_no_ext.as_posix()}_thumb.jpg" ) else: thumbnail_path = ( thumbnails_dir / f"{rel_path_no_ext.as_posix()}_thumb{rel_path.suffix}" ) # Ensure thumbnail path is within our thumbnails directory try: thumbnail_path = thumbnail_path.resolve() thumbnails_dir_resolved = thumbnails_dir.resolve() if not str(thumbnail_path).startswith( str(thumbnails_dir_resolved) ): self.logger.warning( f"Skipping thumbnail outside safe directory: {thumbnail_path}" ) continue except Exception as e: self.logger.warning( f"Path validation failed for {rel_path}: {e}" ) continue # Skip if thumbnail already exists and is newer than original if ( thumbnail_path.exists() and thumbnail_path.stat().st_mtime > media_file.stat().st_mtime ): skipped_count += 1 else: thumbnail_path.parent.mkdir(parents=True, exist_ok=True) if is_video: if self._generate_video_thumbnail( media_file, thumbnail_path, thumbnail_size ): generated_count += 1 else: errors.append( f"Failed to generate video thumbnail for {media_file.name}" ) else: with Image.open(media_file) as img: if img.mode in ("RGBA", "LA", "P"): img = img.convert("RGB") img.thumbnail( thumbnail_size, Image.Resampling.LANCZOS ) save_kwargs = {"quality": 85, "optimize": True} if thumbnail_path.suffix.lower() == ".png": save_kwargs = {"optimize": True} img.save(thumbnail_path, **save_kwargs) generated_count += 1 # Send progress update if ( (i + 1) % 5 == 0 or (i + 1) % max(1, total_images // 100) == 0 or i == total_images - 1 ): elapsed = time.time() - start_time progress_percent = ((i + 1) / total_images) * 100 rate = (i + 1) / elapsed if elapsed > 0 else 0 eta = ((total_images - i - 1) / rate) if rate > 0 else 0 file_info = { "name": media_file.name, "dir": media_file.parent.name, "type": "video" if is_video else "image", "action": ( "skipped" if thumbnail_path.exists() else "generating" ), } await send_progress( "progress", { "processed": i + 1, "total_images": total_images, "generated": generated_count, "skipped": skipped_count, "percentage": round(progress_percent, 1), "rate": round(rate, 1), "eta": round(eta, 0), "elapsed": round(elapsed, 1), "current_file": f"{file_info['dir']}/{file_info['name']}", "file_type": file_info["type"], "action": file_info["action"], }, ) if (i + 1) % 50 == 0: self.logger.info( f"Thumbnail progress: {i+1}/{total_images} ({progress_percent:.1f}%) - Generated: {generated_count}, Skipped: {skipped_count}" ) except Exception as e: error_msg = f"Failed to generate thumbnail for {media_file.name}: {str(e)}" errors.append(error_msg) self.logger.warning(error_msg) if len(errors) <= 5: await send_progress( "status", { "phase": "processing", "message": f"Error processing {media_file.name}: {str(e)}", }, ) elapsed_time = time.time() - start_time completion_message = f"Successfully generated {generated_count} new thumbnails, skipped {skipped_count} existing" if errors: completion_message += f" ({len(errors)} errors occurred)" await send_progress( "complete", { "count": generated_count, "skipped": skipped_count, "total_images": total_images, "errors": errors[:10], "error_count": len(errors), "elapsed_time": round(elapsed_time, 2), "processing_rate": round( (total_images / elapsed_time) if elapsed_time > 0 else 0, 2 ), "message": completion_message, }, ) self.logger.info( f"Thumbnail generation completed: {generated_count} generated, {skipped_count} skipped, {len(errors)} errors in {elapsed_time:.2f}s" ) except Exception as e: self.logger.exception("Thumbnail generation failed") await send_progress( "error", { "error": "An internal error occurred", "message": "Thumbnail generation failed. Check server logs for details.", }, ) return response except ImportError: return web.json_response( { "success": False, "error": "PIL (Pillow) library not available. Install with: pip install Pillow", }, status=500, ) except Exception as e: self.logger.error(f"Generate thumbnails with progress error: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) def _generate_video_thumbnail(self, video_path, thumbnail_path, thumbnail_size): """Generate thumbnail from video file. Returns True if successful.""" try: # Try using OpenCV first (most reliable) try: import cv2 cap = cv2.VideoCapture(str(video_path)) if not cap.isOpened(): return False frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) target_frame = max(1, int(frame_count * 0.1)) cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame) ret, frame = cap.read() cap.release() if not ret or frame is None: return False frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) img = Image.fromarray(frame_rgb) img.thumbnail(thumbnail_size, Image.Resampling.LANCZOS) img.save(thumbnail_path, "JPEG", quality=85, optimize=True) self.logger.debug( f"Generated video thumbnail using OpenCV: {thumbnail_path}" ) return True except ImportError: pass # Fallback to ffmpeg try: import subprocess cmd = [ "ffmpeg", "-i", str(video_path), "-ss", "00:00:01", "-vframes", "1", "-s", f"{thumbnail_size[0]}x{thumbnail_size[1]}", "-y", str(thumbnail_path), ] result = subprocess.run(cmd, capture_output=True, text=True, timeout=30) if result.returncode == 0: self.logger.debug( f"Generated video thumbnail using ffmpeg: {thumbnail_path}" ) return True else: self.logger.warning( f"ffmpeg failed for {video_path}: {result.stderr}" ) except (ImportError, subprocess.TimeoutExpired, FileNotFoundError): pass # Last resort: create a placeholder thumbnail try: from PIL import ImageDraw, ImageFont img = Image.new("RGB", thumbnail_size, color=(50, 50, 50)) draw = ImageDraw.Draw(img) center_x, center_y = thumbnail_size[0] // 2, thumbnail_size[1] // 2 triangle_size = min(thumbnail_size) // 4 points = [ (center_x - triangle_size // 2, center_y - triangle_size // 2), (center_x - triangle_size // 2, center_y + triangle_size // 2), (center_x + triangle_size // 2, center_y), ] draw.polygon(points, fill=(255, 255, 255)) try: font = ImageFont.load_default() text = "VIDEO" bbox = draw.textbbox((0, 0), text, font=font) text_width = bbox[2] - bbox[0] draw.text( ( center_x - text_width // 2, center_y + triangle_size // 2 + 10, ), text, fill=(255, 255, 255), font=font, ) except (OSError, AttributeError): pass img.save(thumbnail_path, "JPEG", quality=85) self.logger.debug( f"Generated placeholder video thumbnail: {thumbnail_path}" ) return True except Exception as e: self.logger.warning( f"Failed to create placeholder thumbnail for {video_path}: {e}" ) return False except Exception as e: self.logger.error( f"Video thumbnail generation failed for {video_path}: {e}" ) return False async def clear_thumbnails(self, request): """Safely clear only our generated thumbnails, never touch original images.""" try: import shutil # Find ComfyUI output directory output_dir = self._find_comfyui_output_dir() if not output_dir: return web.json_response( {"success": False, "error": "ComfyUI output directory not found"}, status=404, ) output_path = Path(output_dir) thumbnails_dir = output_path / "thumbnails" if not thumbnails_dir.exists(): return web.json_response( { "success": True, "message": "No thumbnails directory found - nothing to clear", "cleared_files": 0, } ) # Verify this is actually our thumbnails directory try: thumbnails_dir_resolved = thumbnails_dir.resolve() output_path_resolved = output_path.resolve() if ( not str(thumbnails_dir_resolved).startswith( str(output_path_resolved) ) or thumbnails_dir.name != "thumbnails" ): self.logger.error( f"Safety check failed: thumbnails directory path invalid: {thumbnails_dir}" ) return web.json_response( { "success": False, "error": "Safety check failed: invalid thumbnails directory path", }, status=400, ) except Exception as e: self.logger.error(f"Path validation failed: {e}") return web.json_response( {"success": False, "error": "Path validation failed"}, status=500 ) # Count files before deletion cleared_count = 0 cleared_size = 0 for root, dirs, files in os.walk(thumbnails_dir): for file in files: file_path = Path(root) / file if "_thumb" in file.lower() and any( file.lower().endswith(ext) for ext in [".png", ".jpg", ".jpeg", ".webp", ".gif"] ): try: file_size = file_path.stat().st_size file_path.unlink() cleared_count += 1 cleared_size += file_size self.logger.debug(f"Cleared thumbnail: {file_path}") except Exception as e: self.logger.warning( f"Failed to delete thumbnail {file_path}: {e}" ) # Remove empty directories within thumbnails folder try: for root, dirs, files in os.walk(thumbnails_dir, topdown=False): if root != str(thumbnails_dir): try: Path(root).rmdir() except OSError: pass except Exception as e: self.logger.debug(f"Directory cleanup info: {e}") def format_size(bytes_size): for unit in ["B", "KB", "MB", "GB"]: if bytes_size < 1024.0: return f"{bytes_size:.1f} {unit}" bytes_size /= 1024.0 return f"{bytes_size:.1f} TB" self.logger.info( f"Thumbnail cleanup: cleared {cleared_count} files ({format_size(cleared_size)})" ) return web.json_response( { "success": True, "cleared_files": cleared_count, "cleared_size": cleared_size, "cleared_size_formatted": format_size(cleared_size), "message": f"Cleared {cleared_count} thumbnail files ({format_size(cleared_size)})", } ) except Exception as e: self.logger.error(f"Clear thumbnails error: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) async def link_image_to_prompt(self, request): """Link a generated image to a prompt.""" try: data = await request.json() prompt_id = data.get("prompt_id") image_path = data.get("image_path") metadata = data.get("metadata", {}) if not prompt_id or not image_path: return web.json_response( { "success": False, "error": "prompt_id and image_path are required", }, status=400, ) if not os.path.exists(image_path): return web.json_response( {"success": False, "error": "Image file not found"}, status=404 ) image_id = await self._run_in_executor( self.db.link_image_to_prompt, prompt_id, image_path, metadata ) return web.json_response( { "success": True, "image_id": image_id, "message": "Image linked successfully", } ) except Exception as e: self.logger.error(f"Link image error: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) async def get_image_prompt(self, request): """Get prompt information for a specific image path.""" try: # Get the image path from URL raw_image_path = request.match_info.get("image_path", "") image_path = urllib.parse.unquote(raw_image_path) if not image_path: return web.json_response( {"success": False, "error": "Image path is required"}, status=400 ) # Convert relative path to absolute if needed if not os.path.isabs(image_path): output_dir = self._find_comfyui_output_dir() if output_dir: image_path = str(Path(output_dir) / image_path) # Look up the image in generated_images table try: prompt_data = await self._run_in_executor( self.db.get_image_prompt_info, image_path ) if prompt_data: prompt_data["image_path"] = image_path if prompt_data: return web.json_response({"success": True, "prompt": prompt_data}) else: return web.json_response( { "success": False, "error": "No prompt found for this image", "image_path": image_path, } ) except Exception as db_error: self.logger.error(f"Database error in get_image_prompt: {db_error}") return web.json_response( {"success": False, "error": "Database error occurred"}, status=500 ) except Exception as e: self.logger.error(f"Get image prompt error: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) async def delete_image(self, request): """Delete an image record.""" try: image_id = int(request.match_info["image_id"]) success = await self._run_in_executor(self.db.delete_image, image_id) if success: return web.json_response( {"success": True, "message": "Image deleted successfully"} ) else: return web.json_response( {"success": False, "error": "Image not found"}, status=404 ) except ValueError: return web.json_response( {"success": False, "error": "Invalid image ID"}, status=400 ) except Exception as e: self.logger.error(f"Delete image error: {e}") return web.json_response({"success": False, "error": str(e)}, status=500)