Vito Sansevero
fea745deb9
- Prompt text now uses civitai example prompts (images[].meta.prompt) when available, falling back to model name instead of trigger words - Image serving now allows paths within LoRA directories when the integration is enabled, fixing 403 errors on preview images
1232 lines
50 KiB
Python
1232 lines
50 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)
|
|
|
|
@routes.get("/prompt_manager/gallery/subfolders")
|
|
async def get_gallery_subfolders_route(request):
|
|
return await self.get_gallery_subfolders(request)
|
|
|
|
async def get_prompt_images(self, request):
|
|
"""Get all images for a specific prompt."""
|
|
try:
|
|
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]
|
|
|
|
# Add url and thumbnail_url so the frontend can serve them
|
|
self._enrich_images(cleaned_images)
|
|
|
|
# Additional fallback: convert to JSON string and clean NaN values manually
|
|
try:
|
|
response_data = {"success": True, "images": cleaned_images}
|
|
# 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).
|
|
|
|
Returns list of (path, mtime) tuples sorted by mtime descending.
|
|
"""
|
|
image_extensions = [".png", ".jpg", ".jpeg", ".webp", ".gif"]
|
|
video_extensions = [".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"]
|
|
media_extensions = image_extensions + video_extensions
|
|
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)
|
|
try:
|
|
mtime = media_path.stat().st_mtime
|
|
all_images.append((media_path, mtime))
|
|
except OSError:
|
|
continue
|
|
|
|
all_images.sort(key=lambda x: x[1], reverse=True)
|
|
return all_images
|
|
|
|
async def _get_gallery_files(self, output_path):
|
|
"""Get gallery files with per-directory TTL cache."""
|
|
now = _time.monotonic()
|
|
key = str(output_path)
|
|
cached = self._gallery_cache.get(key)
|
|
if cached and (now - cached[1]) < self._gallery_cache_ttl:
|
|
return cached[0]
|
|
|
|
files = await self._run_in_executor(self._scan_gallery_files_sync, output_path)
|
|
self._gallery_cache[key] = (files, now)
|
|
return files
|
|
|
|
async def get_output_images(self, request):
|
|
"""Get all images from ComfyUI output folder(s)."""
|
|
try:
|
|
from urllib.parse import quote
|
|
|
|
output_dirs = self._get_all_output_dirs()
|
|
if not output_dirs:
|
|
return web.json_response(
|
|
{
|
|
"success": False,
|
|
"error": "No output directories found",
|
|
"images": [],
|
|
},
|
|
)
|
|
|
|
limit = int(request.query.get("limit", 100))
|
|
offset = int(request.query.get("offset", 0))
|
|
subfolder = request.query.get("subfolder", "").strip()
|
|
|
|
# Collect (path, mtime, root) from all directories
|
|
all_images = []
|
|
for output_path in output_dirs:
|
|
dir_images = await self._get_gallery_files(output_path)
|
|
for img_path, mtime in dir_images:
|
|
all_images.append((img_path, mtime, output_path))
|
|
|
|
# Sort combined results by mtime (newest first) — no .stat() calls
|
|
all_images.sort(key=lambda x: x[1], reverse=True)
|
|
|
|
# Apply subfolder filter if provided
|
|
if subfolder:
|
|
filtered = []
|
|
for img_path, mtime, root in all_images:
|
|
rel_dir = str(img_path.relative_to(root).parent)
|
|
if rel_dir == subfolder or rel_dir.startswith(subfolder + os.sep):
|
|
filtered.append((img_path, mtime, root))
|
|
all_images = filtered
|
|
|
|
total = len(all_images)
|
|
paginated = all_images[offset : offset + limit]
|
|
|
|
video_extensions = [".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"]
|
|
|
|
def _format_page():
|
|
images = []
|
|
for media_path, mtime, output_path in paginated:
|
|
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
|
|
thumbnails_dir = output_path / "thumbnails"
|
|
if thumbnails_dir.exists():
|
|
thumbnail_ext = ".jpg" if is_video else extension
|
|
rel_path_no_ext = rel_path.with_suffix("")
|
|
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),
|
|
"root_dir": str(output_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": total,
|
|
"offset": offset,
|
|
"limit": limit,
|
|
"has_more": offset + limit < total,
|
|
}
|
|
)
|
|
except Exception as e:
|
|
self.logger.error(f"Output images error: {e}")
|
|
return web.json_response({"success": False, "error": str(e)}, status=500)
|
|
|
|
async def serve_image(self, request):
|
|
"""Serve the actual image file using streamed FileResponse."""
|
|
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 any allowed directory
|
|
allowed_dirs = list(self._get_all_output_dirs())
|
|
|
|
# Also allow LoRA directories when integration is enabled
|
|
try:
|
|
from ..config import IntegrationConfig
|
|
|
|
if IntegrationConfig.LORA_MANAGER_ENABLED:
|
|
from ..lora_utils import find_lora_directories
|
|
|
|
lora_dirs = find_lora_directories(
|
|
IntegrationConfig.LORA_MANAGER_PATH
|
|
)
|
|
allowed_dirs.extend(Path(d) for d in lora_dirs)
|
|
except Exception:
|
|
pass
|
|
|
|
if allowed_dirs:
|
|
allowed = any(
|
|
image_path.is_relative_to(d.resolve()) for d in allowed_dirs
|
|
)
|
|
if not allowed:
|
|
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"]
|
|
|
|
# Search all configured output directories
|
|
output_dirs = self._get_all_output_dirs()
|
|
if not output_dirs:
|
|
return web.json_response(
|
|
{"success": False, "error": "ComfyUI output directory not found"},
|
|
status=404,
|
|
)
|
|
|
|
# Try each root directory until the file is found
|
|
for output_path in output_dirs:
|
|
image_path = (output_path / filepath).resolve()
|
|
|
|
# Security check: must be within this output directory
|
|
if not image_path.is_relative_to(output_path.resolve()):
|
|
continue
|
|
|
|
if image_path.exists():
|
|
response = web.FileResponse(image_path)
|
|
response.headers["Cache-Control"] = "public, max-age=3600"
|
|
return response
|
|
|
|
return web.json_response(
|
|
{"success": False, "error": "Image file not found"}, status=404
|
|
)
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Serve output image error: {e}")
|
|
return web.json_response({"success": False, "error": str(e)}, status=500)
|
|
|
|
async def get_gallery_subfolders(self, request):
|
|
"""Get distinct subfolders from gallery output directories."""
|
|
try:
|
|
output_dirs = self._get_all_output_dirs()
|
|
subfolders = set()
|
|
for output_path in output_dirs:
|
|
files = await self._get_gallery_files(output_path)
|
|
for f, _mtime in files:
|
|
rel_dir = str(f.relative_to(output_path).parent)
|
|
if rel_dir and rel_dir != ".":
|
|
subfolders.add(rel_dir)
|
|
|
|
return web.json_response(
|
|
{"success": True, "subfolders": sorted(subfolders)}
|
|
)
|
|
except Exception as e:
|
|
self.logger.error(f"Subfolders error: {e}")
|
|
return web.json_response({"success": False, "error": str(e)}, status=500)
|
|
|
|
async def generate_thumbnails(self, request):
|
|
"""Generate thumbnails for all images and videos in the ComfyUI output directory."""
|
|
try:
|
|
# 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}",
|
|
exc_info=True,
|
|
)
|
|
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)
|