Vito Sansevero
4548beb723
Phase 4 structural refactoring: - Split monolithic py/api.py (5.3k lines) into domain mixins: prompts.py, images.py, admin.py, logging_routes.py, autotag_routes.py - Extract inline JS from admin.html into web/js/admin.js - Extract inline JS from gallery.html into web/js/gallery.js - Add gzip compression middleware to API responses (Phase 5.4) - Standardize API error envelope with success: false (Phase 5.3) - Add filmstrip image enrichment to prompt list responses (Phase 5.2)
1046 lines
44 KiB
Python
1046 lines
44 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:
|
|
await send_progress('error', {
|
|
'error': str(e),
|
|
'message': f'Thumbnail generation failed: {str(e)}'
|
|
})
|
|
|
|
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)
|