Vito Sansevero
027fdc09bb
- /images/recent, /images/all and /images/output parse limit/offset through parse_page_params on every path (absent, 0, negative, float strings), and the clamped values are what reach the DB. /images/all gets its own MAX_BULK_LIMIT ceiling instead of an unbounded query; /images/recent caps its fetch window at MAX_RECENT_IMAGES_WINDOW. - parse_page_params also clamps offset to MAX_PAGE_OFFSET so values past int64 no longer surface as a 500 from SQLite. - Gallery and thumbnail scans share _iter_media_files: no symlink following, MAX_SCAN_DEPTH, and MAX_GALLERY_FILES / MAX_THUMBNAIL_FILES caps. Thumbnail size comes only from named presets. - JSON bodies on /images/link and /images/generate-thumbnails are read in bounded chunks; over MAX_JSON_BODY_BYTES answers 413, malformed or non-object JSON answers 400. - /images/link only accepts media paths inside the allowed directories (403 otherwise), closing a path-existence oracle.
1216 lines
46 KiB
Python
1216 lines
46 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, UnidentifiedImageError
|
|
|
|
from .prompts import bad_request, parse_page_params
|
|
from .prompts import safe_error_message as _safe_error
|
|
|
|
# Only these file types are ever served by the image routes, regardless of
|
|
# what sits inside an allowed directory (defence in depth for the gallery).
|
|
IMAGE_EXTENSIONS = frozenset(
|
|
{
|
|
".png",
|
|
".jpg",
|
|
".jpeg",
|
|
".webp",
|
|
".gif",
|
|
".bmp",
|
|
".tiff",
|
|
".tif",
|
|
".mp4",
|
|
".webm",
|
|
".mov",
|
|
}
|
|
)
|
|
|
|
|
|
def _canonical(path):
|
|
"""Resolve symlinks and normalise case so paths compare on every OS."""
|
|
return Path(os.path.normcase(os.path.realpath(str(path))))
|
|
|
|
|
|
def _is_within(path, root):
|
|
"""True when *path* (after realpath) lives under *root* (after realpath)."""
|
|
return _canonical(path).is_relative_to(_canonical(root))
|
|
|
|
|
|
def _has_traversal(relative):
|
|
"""True for absolute paths, drive-qualified paths or any ``..`` segment."""
|
|
if not relative or os.path.isabs(relative) or os.path.splitdrive(relative)[0]:
|
|
return True
|
|
if relative[0] in ("/", "\\"):
|
|
return True
|
|
return ".." in re.split(r"[\\/]+", relative)
|
|
|
|
|
|
def _forbidden(message="Access denied"):
|
|
return web.json_response({"success": False, "error": message}, status=403)
|
|
|
|
|
|
def _media_access_error(path, allowed_dirs):
|
|
"""Return a 403 response when *path* must not be served, else None.
|
|
|
|
Fails closed: an empty *allowed_dirs* denies everything.
|
|
"""
|
|
if Path(path).suffix.lower() not in IMAGE_EXTENSIONS:
|
|
return _forbidden("Only media files can be served")
|
|
if not allowed_dirs:
|
|
return _forbidden("No allowed image directories configured")
|
|
if not any(_is_within(path, root) for root in allowed_dirs):
|
|
return _forbidden()
|
|
return None
|
|
|
|
|
|
# Media types the gallery scans and thumbnails (a superset is never served).
|
|
GALLERY_IMAGE_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".webp", ".gif"})
|
|
VIDEO_EXTENSIONS = frozenset({".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"})
|
|
MEDIA_EXTENSIONS = GALLERY_IMAGE_EXTENSIONS | VIDEO_EXTENSIONS
|
|
|
|
# Caps on request-triggered filesystem work.
|
|
MAX_SCAN_DEPTH = 12
|
|
MAX_THUMBNAIL_FILES = 10_000
|
|
|
|
|
|
def _write_image_thumbnail(src, dst, thumbnail_size):
|
|
"""Resize *src* into *dst* (blocking PIL work)."""
|
|
with Image.open(src) 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 dst.suffix.lower() == ".png":
|
|
save_kwargs = {"optimize": True}
|
|
img.save(dst, **save_kwargs)
|
|
|
|
|
|
def _new_thumbnail_stats():
|
|
return {"generated": 0, "skipped": 0, "errors": []}
|
|
|
|
|
|
def _record_thumbnail_result(stats, result):
|
|
action = result["action"]
|
|
if action == "generated":
|
|
stats["generated"] += 1
|
|
elif action == "skipped":
|
|
stats["skipped"] += 1
|
|
else:
|
|
stats["errors"].append(f"{result['file']}: {result.get('error', 'failed')}")
|
|
|
|
|
|
def _thumbnail_progress_payload(index, total, stats, result, start):
|
|
elapsed = _time.monotonic() - start
|
|
rate = index / elapsed if elapsed > 0 else 0
|
|
eta = (total - index) / rate if rate > 0 else 0
|
|
return {
|
|
"processed": index,
|
|
"total_images": total,
|
|
"generated": stats["generated"],
|
|
"skipped": stats["skipped"],
|
|
"error_count": len(stats["errors"]),
|
|
"percentage": round(index / total * 100, 1) if total else 100.0,
|
|
"rate": round(rate, 1),
|
|
"eta": round(eta),
|
|
"elapsed": round(elapsed, 1),
|
|
"current_file": f"{result['dir']}/{result['file']}",
|
|
"file_type": result["type"],
|
|
"action": result["action"],
|
|
}
|
|
|
|
|
|
def _thumbnail_complete_payload(total, stats, start):
|
|
elapsed = _time.monotonic() - start
|
|
errors = stats["errors"]
|
|
message = (
|
|
f"Generated {stats['generated']} new thumbnails, "
|
|
f"skipped {stats['skipped']} existing"
|
|
)
|
|
if errors:
|
|
message += f" ({len(errors)} errors occurred)"
|
|
return {
|
|
"count": stats["generated"],
|
|
"skipped": stats["skipped"],
|
|
"total_images": total,
|
|
"errors": errors[:10],
|
|
"error_count": len(errors),
|
|
"elapsed_time": round(elapsed, 2),
|
|
"processing_rate": round(total / elapsed if elapsed > 0 else 0, 2),
|
|
"message": message,
|
|
}
|
|
|
|
|
|
# Caps on request-driven work (list sizes, scans, request bodies).
|
|
MAX_GALLERY_FILES = 50_000
|
|
MAX_RECENT_IMAGES_WINDOW = 5_000
|
|
MAX_BULK_LIMIT = 5_000
|
|
MAX_JSON_BODY_BYTES = 1024 * 1024
|
|
_READ_CHUNK = 64 * 1024
|
|
|
|
|
|
class _BodyTooLarge(Exception):
|
|
"""Raised when a request body exceeds MAX_JSON_BODY_BYTES."""
|
|
|
|
|
|
def _json_error(message, status):
|
|
return web.json_response({"success": False, "error": message}, status=status)
|
|
|
|
|
|
async def _read_json_body(request):
|
|
"""Read a JSON object body in bounded chunks.
|
|
|
|
Raises _BodyTooLarge past MAX_JSON_BODY_BYTES (declared or streamed) and
|
|
ValueError when the body is not a JSON object.
|
|
"""
|
|
limit = MAX_JSON_BODY_BYTES
|
|
if request.content_length is not None and request.content_length > limit:
|
|
raise _BodyTooLarge()
|
|
chunks, size = [], 0
|
|
while True:
|
|
chunk = await request.content.read(_READ_CHUNK)
|
|
if not chunk:
|
|
break
|
|
size += len(chunk)
|
|
if size > limit:
|
|
raise _BodyTooLarge()
|
|
chunks.append(chunk)
|
|
data = json.loads(b"".join(chunks).decode("utf-8"))
|
|
if not isinstance(data, dict):
|
|
raise ValueError("JSON body must be an object")
|
|
return data
|
|
|
|
|
|
async def _json_body_or_error(request):
|
|
"""Return ``(data, None)`` or ``(None, error_response)``."""
|
|
try:
|
|
return await _read_json_body(request), None
|
|
except _BodyTooLarge:
|
|
return None, _json_error("Request body too large", 413)
|
|
except (ValueError, UnicodeDecodeError):
|
|
return None, _json_error("Invalid JSON body", 400)
|
|
|
|
|
|
def _iter_media_files(output_path, max_files):
|
|
"""Yield media files under *output_path* with bounded effort.
|
|
|
|
Symlinks are never followed, ``thumbnails/`` is skipped, descent stops
|
|
at MAX_SCAN_DEPTH and at most *max_files* paths are produced.
|
|
"""
|
|
output_path = Path(output_path)
|
|
root_depth = len(output_path.parts)
|
|
produced = 0
|
|
for root, dirs, files in os.walk(output_path):
|
|
root_path = Path(root)
|
|
rel_parts = root_path.parts[root_depth:]
|
|
if "thumbnails" in rel_parts:
|
|
dirs[:] = []
|
|
continue
|
|
if len(rel_parts) >= MAX_SCAN_DEPTH:
|
|
dirs[:] = []
|
|
for name in sorted(files):
|
|
if Path(name).suffix.lower() not in MEDIA_EXTENSIONS:
|
|
continue
|
|
if produced >= max_files:
|
|
return
|
|
produced += 1
|
|
yield root_path / name
|
|
|
|
|
|
def _output_image_entry(media_path, output_path, root_index):
|
|
"""Build one gallery entry for *media_path* (blocking stat)."""
|
|
stat = media_path.stat()
|
|
rel_path = media_path.relative_to(output_path)
|
|
extension = media_path.suffix.lower()
|
|
is_video = extension in VIDEO_EXTENSIONS
|
|
|
|
thumbnail_url = None
|
|
thumb_rel = (
|
|
f"thumbnails/{rel_path.with_suffix('').as_posix()}_thumb"
|
|
f"{'.jpg' if is_video else extension}"
|
|
)
|
|
if (output_path / thumb_rel).exists():
|
|
thumbnail_url = (
|
|
f"/prompt_manager/images/serve/{urllib.parse.quote(thumb_rel, safe='/')}"
|
|
)
|
|
|
|
return {
|
|
"id": str(hash(str(media_path))),
|
|
"filename": media_path.name,
|
|
"path": str(media_path),
|
|
"relative_path": str(rel_path),
|
|
"root_dir": str(output_path),
|
|
"root_index": root_index,
|
|
"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": "video" if is_video else "image",
|
|
"is_video": is_video,
|
|
}
|
|
|
|
|
|
def _in_subfolder(rel_dir, subfolder):
|
|
return rel_dir == subfolder or rel_dir.startswith(subfolder + os.sep)
|
|
|
|
|
|
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 (bounded window, newest first)."""
|
|
try:
|
|
try:
|
|
limit, offset = parse_page_params(request.query)
|
|
except ValueError:
|
|
return bad_request("limit and offset must be integers")
|
|
|
|
# The DB query has no offset: fetch one bounded window and slice,
|
|
# so the offset is capped to keep the window small.
|
|
offset = min(offset, MAX_RECENT_IMAGES_WINDOW - limit)
|
|
rows = await self._run_in_executor(
|
|
self.db.get_recent_images, limit + offset
|
|
)
|
|
images = rows[offset:]
|
|
|
|
return web.json_response(
|
|
{
|
|
"success": True,
|
|
"images": images,
|
|
"pagination": {
|
|
"limit": limit,
|
|
"offset": offset,
|
|
"count": len(images),
|
|
},
|
|
}
|
|
)
|
|
except Exception as e:
|
|
self.logger.error(f"Get recent images error: {e}")
|
|
return web.json_response(
|
|
{"success": False, "error": _safe_error(e)}, status=500
|
|
)
|
|
|
|
async def get_all_images(self, request):
|
|
"""Get generated images with linked prompts (bulk-capped pages)."""
|
|
try:
|
|
try:
|
|
limit, offset = parse_page_params(
|
|
request.query,
|
|
default_limit=MAX_BULK_LIMIT,
|
|
max_limit=MAX_BULK_LIMIT,
|
|
)
|
|
except ValueError:
|
|
return bad_request("limit and offset must be integers")
|
|
|
|
images = await self._run_in_executor(
|
|
self.db.get_all_images, limit=limit, offset=offset
|
|
)
|
|
|
|
return web.json_response(
|
|
{
|
|
"success": True,
|
|
"images": images,
|
|
"count": len(images),
|
|
"pagination": {
|
|
"limit": limit,
|
|
"offset": offset,
|
|
"count": len(images),
|
|
"has_more": len(images) == limit,
|
|
},
|
|
}
|
|
)
|
|
except Exception as e:
|
|
self.logger.error(f"Get all images error: {e}")
|
|
return web.json_response(
|
|
{"success": False, "error": _safe_error(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; the
|
|
walk is bounded by MAX_SCAN_DEPTH and MAX_GALLERY_FILES.
|
|
"""
|
|
found = []
|
|
for media_path in _iter_media_files(output_path, MAX_GALLERY_FILES):
|
|
try:
|
|
found.append((media_path, media_path.stat().st_mtime))
|
|
except OSError:
|
|
continue
|
|
if len(found) >= MAX_GALLERY_FILES:
|
|
self.logger.warning(f"Gallery scan capped at {MAX_GALLERY_FILES} files")
|
|
found.sort(key=lambda x: x[1], reverse=True)
|
|
return found
|
|
|
|
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:
|
|
output_dirs = self._get_all_output_dirs()
|
|
if not output_dirs:
|
|
return web.json_response(
|
|
{
|
|
"success": False,
|
|
"error": "No output directories found",
|
|
"images": [],
|
|
},
|
|
)
|
|
try:
|
|
limit, offset = parse_page_params(request.query, default_limit=100)
|
|
except ValueError:
|
|
return bad_request("limit and offset must be integers")
|
|
subfolder = request.query.get("subfolder", "").strip()
|
|
|
|
all_images = await self._collect_output_images(output_dirs, subfolder)
|
|
total = len(all_images)
|
|
page = all_images[offset : offset + limit]
|
|
images = await self._run_in_executor(self._format_output_page, 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": _safe_error(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"])
|
|
allowed_dirs = list(self._get_all_output_dirs()) + self._lora_image_dirs()
|
|
|
|
denied = _media_access_error(image_path, allowed_dirs)
|
|
if denied is not None:
|
|
return denied
|
|
|
|
if not image_path.is_file():
|
|
return web.json_response(
|
|
{"success": False, "error": "Image file not found"}, status=404
|
|
)
|
|
|
|
return self._file_response(image_path)
|
|
|
|
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": _safe_error(e)}, status=500
|
|
)
|
|
|
|
def _lora_image_dirs(self):
|
|
"""Extra directories served when the LoRA integration is enabled."""
|
|
try:
|
|
from ..config import IntegrationConfig
|
|
|
|
if not IntegrationConfig.LORA_MANAGER_ENABLED:
|
|
return []
|
|
from ..lora_utils import find_lora_directories, get_lora_image_cache_dir
|
|
|
|
dirs = [
|
|
Path(d)
|
|
for d in find_lora_directories(IntegrationConfig.LORA_MANAGER_PATH)
|
|
]
|
|
dirs.append(Path(get_lora_image_cache_dir()))
|
|
return dirs
|
|
except Exception:
|
|
# LoRA integration is optional; never widen access on failure.
|
|
return []
|
|
|
|
@staticmethod
|
|
def _file_response(path):
|
|
response = web.FileResponse(path)
|
|
response.headers["Cache-Control"] = "public, max-age=3600"
|
|
return response
|
|
|
|
@staticmethod
|
|
def _select_root(output_dirs, root_param):
|
|
"""Narrow *output_dirs* to the ``?root=N`` entry when it is valid."""
|
|
if root_param is None:
|
|
return output_dirs
|
|
try:
|
|
idx = int(root_param)
|
|
except ValueError:
|
|
return output_dirs
|
|
if 0 <= idx < len(output_dirs):
|
|
return [output_dirs[idx]]
|
|
return output_dirs
|
|
|
|
async def serve_output_image(self, request):
|
|
"""Serve a media file from the ComfyUI output folder(s)."""
|
|
try:
|
|
filepath = request.match_info["filepath"]
|
|
if _has_traversal(filepath):
|
|
return _forbidden()
|
|
if Path(filepath).suffix.lower() not in IMAGE_EXTENSIONS:
|
|
return _forbidden("Only media files can be served")
|
|
|
|
output_dirs = self._select_root(
|
|
list(self._get_all_output_dirs()), request.query.get("root")
|
|
)
|
|
if not output_dirs:
|
|
return _forbidden("No allowed image directories configured")
|
|
|
|
for output_path in output_dirs:
|
|
image_path = output_path / filepath
|
|
if not image_path.exists():
|
|
continue
|
|
# A symlink that points outside the root is a hard deny, not
|
|
# a fall-through to the next root.
|
|
if not _is_within(image_path, output_path):
|
|
return _forbidden()
|
|
if image_path.is_file():
|
|
return self._file_response(image_path)
|
|
|
|
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": _safe_error(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)
|
|
|
|
include_ancestors = (
|
|
request.query.get("include_ancestors", "").lower() == "true"
|
|
)
|
|
if include_ancestors:
|
|
ancestors = set()
|
|
for folder in subfolders:
|
|
parts = folder.replace("\\", "/").split("/")
|
|
for i in range(1, len(parts)):
|
|
ancestors.add("/".join(parts[:i]))
|
|
subfolders.update(ancestors)
|
|
|
|
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)
|
|
|
|
# Fixed quality presets: clients choose a name, never a raw pixel size.
|
|
THUMBNAIL_SIZES = {"low": (150, 150), "medium": (300, 300), "high": (600, 600)}
|
|
|
|
@classmethod
|
|
def _thumbnail_size(cls, quality):
|
|
"""Map a quality preset name to a pixel size; anything else is medium."""
|
|
if not isinstance(quality, str):
|
|
return cls.THUMBNAIL_SIZES["medium"]
|
|
return cls.THUMBNAIL_SIZES.get(quality, cls.THUMBNAIL_SIZES["medium"])
|
|
|
|
async def generate_thumbnails(self, request):
|
|
"""Generate every thumbnail in one executor job (POST, blocking)."""
|
|
try:
|
|
data, error = await _json_body_or_error(request)
|
|
if error is not None:
|
|
return error
|
|
thumbnail_size = self._thumbnail_size(data.get("quality", "medium"))
|
|
|
|
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)
|
|
result = await self._run_in_executor(
|
|
self._generate_thumbnails_sync,
|
|
output_path,
|
|
output_path / "thumbnails",
|
|
thumbnail_size,
|
|
)
|
|
self.invalidate_gallery_cache()
|
|
return web.json_response(result)
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Generate thumbnails error: {e}")
|
|
return web.json_response(
|
|
{"success": False, "error": _safe_error(e)}, status=500
|
|
)
|
|
|
|
def _generate_thumbnails_sync(self, output_path, thumbnails_dir, thumbnail_size):
|
|
"""Blocking thumbnail generation loop (run in executor)."""
|
|
start = _time.monotonic()
|
|
targets = self._thumbnail_targets(output_path, thumbnails_dir)
|
|
stats = _new_thumbnail_stats()
|
|
for src, dst, is_video in targets:
|
|
_record_thumbnail_result(
|
|
stats, self._generate_one(src, dst, thumbnail_size, is_video)
|
|
)
|
|
payload = _thumbnail_complete_payload(len(targets), stats, start)
|
|
payload["success"] = True
|
|
payload["thumbnails_path"] = str(thumbnails_dir)
|
|
self.logger.info(payload["message"])
|
|
return payload
|
|
|
|
def _thumbnail_targets(self, output_path, thumbnails_dir):
|
|
"""Scan *output_path* for media (blocking); returns (src, dst, is_video).
|
|
|
|
Bounded by _iter_media_files (no symlinks, MAX_SCAN_DEPTH) and
|
|
MAX_THUMBNAIL_FILES per request.
|
|
"""
|
|
output_path = Path(output_path)
|
|
targets = []
|
|
for src in _iter_media_files(output_path, MAX_THUMBNAIL_FILES):
|
|
suffix = src.suffix.lower()
|
|
is_video = suffix in VIDEO_EXTENSIONS
|
|
rel_no_ext = src.relative_to(output_path).with_suffix("")
|
|
thumb_suffix = ".jpg" if is_video else suffix
|
|
dst = thumbnails_dir / f"{rel_no_ext.as_posix()}_thumb{thumb_suffix}"
|
|
if not _is_within(dst, thumbnails_dir):
|
|
self.logger.warning(f"Skipping thumbnail outside safe dir: {src.name}")
|
|
continue
|
|
targets.append((src, dst, is_video))
|
|
if len(targets) >= MAX_THUMBNAIL_FILES:
|
|
self.logger.warning(f"Thumbnail scan capped at {MAX_THUMBNAIL_FILES}")
|
|
return targets
|
|
|
|
def _generate_one(self, src, dst, thumbnail_size, is_video):
|
|
"""Create one thumbnail (blocking). Never raises; returns a result dict."""
|
|
result = {
|
|
"file": src.name,
|
|
"dir": src.parent.name,
|
|
"type": "video" if is_video else "image",
|
|
}
|
|
try:
|
|
if not src.is_file():
|
|
return {**result, "action": "error", "error": "File no longer exists"}
|
|
if dst.exists() and dst.stat().st_mtime > src.stat().st_mtime:
|
|
return {**result, "action": "skipped"}
|
|
dst.parent.mkdir(parents=True, exist_ok=True)
|
|
if is_video:
|
|
if not self._generate_video_thumbnail(src, dst, thumbnail_size):
|
|
return {**result, "action": "error", "error": "Video decode failed"}
|
|
else:
|
|
_write_image_thumbnail(src, dst, thumbnail_size)
|
|
return {**result, "action": "generated"}
|
|
except UnidentifiedImageError:
|
|
return {**result, "action": "error", "error": "Not a valid image file"}
|
|
except Exception as e:
|
|
self.logger.warning(f"Thumbnail failed for {src.name}", exc_info=True)
|
|
return {**result, "action": "error", "error": _safe_error(e)}
|
|
|
|
async def generate_thumbnails_with_progress(self, request):
|
|
"""Generate thumbnails with Server-Sent Events progress updates.
|
|
|
|
Every blocking step (scan, PIL/ffmpeg) runs in the executor; the
|
|
event loop only writes SSE frames between files.
|
|
"""
|
|
thumbnail_size = self._thumbnail_size(request.query.get("quality", "medium"))
|
|
response = web.StreamResponse(
|
|
status=200,
|
|
headers={
|
|
"Content-Type": "text/event-stream",
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
},
|
|
)
|
|
await response.prepare(request)
|
|
try:
|
|
output_dir = self._find_comfyui_output_dir()
|
|
if not output_dir:
|
|
await self._emit_sse(
|
|
response, "error", {"error": "ComfyUI output directory not found"}
|
|
)
|
|
return response
|
|
|
|
output_path = Path(output_dir)
|
|
await self._emit_sse(
|
|
response,
|
|
"status",
|
|
{"phase": "scanning", "message": "Scanning for images and videos..."},
|
|
)
|
|
targets = await self._run_in_executor(
|
|
self._thumbnail_targets, output_path, output_path / "thumbnails"
|
|
)
|
|
await self._stream_thumbnails(response, targets, thumbnail_size)
|
|
self.invalidate_gallery_cache()
|
|
except Exception:
|
|
self.logger.exception("Thumbnail generation failed")
|
|
await self._emit_sse(
|
|
response,
|
|
"error",
|
|
{
|
|
"error": "An internal error occurred",
|
|
"message": "Thumbnail generation failed. Check server logs.",
|
|
},
|
|
)
|
|
return response
|
|
|
|
async def _stream_thumbnails(self, response, targets, thumbnail_size):
|
|
"""Process *targets* one per executor job, emitting SSE after each."""
|
|
total = len(targets)
|
|
video_count = sum(1 for _, _, is_video in targets if is_video)
|
|
await self._emit_sse(
|
|
response,
|
|
"start",
|
|
{
|
|
"total_images": total,
|
|
"phase": "processing",
|
|
"image_count": total - video_count,
|
|
"video_count": video_count,
|
|
"message": (
|
|
f"Found {total - video_count} images and "
|
|
f"{video_count} videos to process"
|
|
),
|
|
},
|
|
)
|
|
stats = _new_thumbnail_stats()
|
|
start = _time.monotonic()
|
|
for index, (src, dst, is_video) in enumerate(targets, start=1):
|
|
result = await self._run_in_executor(
|
|
self._generate_one, src, dst, thumbnail_size, is_video
|
|
)
|
|
_record_thumbnail_result(stats, result)
|
|
if result["action"] == "error":
|
|
await self._emit_sse(
|
|
response,
|
|
"file_error",
|
|
{"file": result["file"], "error": result["error"]},
|
|
)
|
|
await self._emit_sse(
|
|
response,
|
|
"progress",
|
|
_thumbnail_progress_payload(index, total, stats, result, start),
|
|
)
|
|
await asyncio.sleep(0) # let other requests run between files
|
|
payload = _thumbnail_complete_payload(total, stats, start)
|
|
await self._emit_sse(response, "complete", payload)
|
|
self.logger.info(payload["message"])
|
|
|
|
async def _emit_sse(self, response, event, data):
|
|
"""Write one SSE frame; a closed client is logged, not raised."""
|
|
try:
|
|
frame = f"event: {event}\ndata: {json.dumps(data)}\n\n"
|
|
await response.write(frame.encode("utf-8"))
|
|
except Exception as e:
|
|
self.logger.warning(f"Failed to send SSE message: {e}")
|
|
|
|
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 (inside an allowed directory) to a prompt."""
|
|
try:
|
|
data, error = await _json_body_or_error(request)
|
|
if error is not None:
|
|
return error
|
|
prompt_id = data.get("prompt_id")
|
|
image_path = data.get("image_path")
|
|
metadata = data.get("metadata", {})
|
|
|
|
if not prompt_id or not isinstance(image_path, str) or not image_path:
|
|
return bad_request("prompt_id and image_path are required")
|
|
|
|
denied = _media_access_error(image_path, self._get_all_output_dirs())
|
|
if denied is not None:
|
|
return denied
|
|
|
|
if not os.path.isfile(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": _safe_error(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)
|
|
|
|
async def _collect_output_images(self, output_dirs, subfolder):
|
|
"""Merge cached per-root listings, newest first, optionally filtered."""
|
|
merged = []
|
|
for root_idx, output_path in enumerate(output_dirs):
|
|
for img_path, mtime in await self._get_gallery_files(output_path):
|
|
merged.append((img_path, mtime, output_path, root_idx))
|
|
merged.sort(key=lambda x: x[1], reverse=True)
|
|
if not subfolder:
|
|
return merged
|
|
return [
|
|
entry
|
|
for entry in merged
|
|
if _in_subfolder(str(entry[0].relative_to(entry[2]).parent), subfolder)
|
|
]
|
|
|
|
def _format_output_page(self, page):
|
|
entries = []
|
|
for media_path, _mtime, output_path, root_index in page:
|
|
try:
|
|
entries.append(_output_image_entry(media_path, output_path, root_index))
|
|
except Exception as e:
|
|
self.logger.error(f"Error processing media {media_path.name}: {e}")
|
|
return entries
|