Vito Sansevero
ed40be066c
Log exception details server-side only; send generic error messages to clients via SSE streams to resolve code scanning alerts.
1181 lines
48 KiB
Python
1181 lines
48 KiB
Python
"""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:
|
|
self.logger.warning(
|
|
f"Failed to generate thumbnail for {media_file.name}: {e}"
|
|
)
|
|
error_msg = f"Failed to generate thumbnail for {media_file.name}"
|
|
errors.append(error_msg)
|
|
|
|
if len(errors) <= 5:
|
|
await send_progress(
|
|
"status",
|
|
{
|
|
"phase": "processing",
|
|
"message": f"Error processing {media_file.name}",
|
|
},
|
|
)
|
|
|
|
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)
|