Vito Sansevero
113d663fce
- Fix XSS in showFullPrompt() across admin.js, gallery.js, metadata.html by replacing HTML string interpolation with DOM manipulation (createElement + textContent), eliminating unsafe user content injection - Fix stack trace exposure in py/api error responses by replacing str(e) with generic messages and adding logger.exception() for server-side traceability
1179 lines
47 KiB
Python
1179 lines
47 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:
|
|
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)
|