171 lines
7.0 KiB
Python
171 lines
7.0 KiB
Python
"""Generated-output gallery indexing and mutation routes."""
|
|
|
|
import asyncio
|
|
import os
|
|
import threading
|
|
import time
|
|
|
|
from aiohttp import web
|
|
import folder_paths
|
|
|
|
from .image_search import filter_images
|
|
from .media_routes import (
|
|
CARD_POSTER_VIDEO_EXTENSIONS,
|
|
CARD_THUMBNAIL_STATIC_EXTENSIONS,
|
|
_build_card_thumbnail,
|
|
_build_video_poster,
|
|
)
|
|
from .path_utils import require_filename, resolve_within
|
|
from .trash import move_to_trash, trash_failure
|
|
|
|
|
|
GALLERY_SNAPSHOT_SECONDS = 10
|
|
MAX_CACHED_GALLERY_IMAGES = 50_000
|
|
SUPPORTED_GALLERY_EXTENSIONS = {
|
|
'.png', '.jpg', '.jpeg', '.webp', '.gif', '.avif',
|
|
'.mp4', '.webm', '.mov', '.avi', '.mkv',
|
|
}
|
|
_gallery_snapshot_lock = threading.Lock()
|
|
_gallery_snapshot = {"root": None, "expires": 0, "images": []}
|
|
|
|
|
|
def _collect_gallery_images(output_dir):
|
|
images = []
|
|
for root, _, files in os.walk(output_dir):
|
|
for filename in files:
|
|
ext = os.path.splitext(filename)[1].lower()
|
|
if ext not in SUPPORTED_GALLERY_EXTENSIONS:
|
|
continue
|
|
try:
|
|
mtime = os.path.getmtime(os.path.join(root, filename))
|
|
except OSError:
|
|
continue
|
|
subfolder = os.path.relpath(root, output_dir)
|
|
images.append({
|
|
"filename": filename,
|
|
"subfolder": "" if subfolder == "." else subfolder.replace(os.sep, '/'),
|
|
"type": "output",
|
|
"mtime": mtime,
|
|
"ext": ext.lstrip('.'),
|
|
"media_type": "video" if ext in CARD_POSTER_VIDEO_EXTENSIONS else "image",
|
|
})
|
|
images.sort(key=lambda item: (item['mtime'], item['subfolder'], item['filename']), reverse=True)
|
|
return images
|
|
|
|
|
|
def _gallery_images(output_dir, refresh=False):
|
|
root = os.path.realpath(output_dir)
|
|
with _gallery_snapshot_lock:
|
|
if not refresh and _gallery_snapshot['root'] == root and time.monotonic() < _gallery_snapshot['expires']:
|
|
return _gallery_snapshot['images']
|
|
images = _collect_gallery_images(root)
|
|
_gallery_snapshot.update(root=root, expires=time.monotonic() + GALLERY_SNAPSHOT_SECONDS,
|
|
images=images if len(images) <= MAX_CACHED_GALLERY_IMAGES else [])
|
|
if len(images) > MAX_CACHED_GALLERY_IMAGES:
|
|
_gallery_snapshot['expires'] = 0
|
|
return images
|
|
|
|
|
|
def _invalidate_gallery_snapshot():
|
|
with _gallery_snapshot_lock:
|
|
_gallery_snapshot.update(root=None, expires=0, images=[])
|
|
|
|
|
|
def _gallery_listing(output_dir, refresh, query, allowed_exts=None):
|
|
images = _gallery_images(output_dir, refresh)
|
|
if allowed_exts is not None:
|
|
images = [img for img in images if os.path.splitext(img['filename'])[1].lower().lstrip('.') in allowed_exts]
|
|
return filter_images(os.path.realpath(output_dir), images, query) if query else images
|
|
|
|
|
|
async def api_get_gallery_images(request):
|
|
"""GET /anomalous/gallery_images?page&limit&refresh&term=..&ext=.. - Output images/videos, optionally filtered by extension."""
|
|
try:
|
|
page = max(1, int(request.query.get('page', 1)))
|
|
limit = min(200, max(1, int(request.query.get('limit', 50))))
|
|
# Repeated `term` params keep phrases whole; `q` is split on whitespace.
|
|
terms = [term for term in request.query.getall('term', []) if term.strip()]
|
|
query = terms or request.query.get('q', '').strip()
|
|
|
|
# Optional extension filter (?ext=png&ext=mp4 or ?ext=png,mp4)
|
|
ext_filter = None
|
|
if 'ext' in request.query:
|
|
allowed_exts = set()
|
|
for raw in request.query.getall('ext', []):
|
|
for part in raw.split(','):
|
|
cleaned = part.strip().lower().lstrip('.')
|
|
if cleaned:
|
|
allowed_exts.add(cleaned)
|
|
ext_filter = allowed_exts
|
|
|
|
images = await asyncio.to_thread(
|
|
_gallery_listing,
|
|
folder_paths.get_output_directory(),
|
|
request.query.get('refresh') == '1',
|
|
query,
|
|
ext_filter,
|
|
)
|
|
total = len(images)
|
|
start = (page - 1) * limit
|
|
return web.json_response({"images": images[start:start + limit], "total": total,
|
|
"page": page, "pages": (total + limit - 1) // limit})
|
|
except (TypeError, ValueError):
|
|
return web.json_response({"error": "Invalid pagination"}, status=400)
|
|
except OSError:
|
|
return web.json_response({"error": "Could not list output images"}, status=500)
|
|
|
|
|
|
async def api_output_thumbnail(request):
|
|
"""GET /anomalous/output_thumbnail?filename&subfolder - the gallery grid's 512 px WebP of an
|
|
output image/video (cached like model cards); the original stays at ComfyUI's /view."""
|
|
output_dir = folder_paths.get_output_directory()
|
|
try:
|
|
filename = require_filename(request.query.get('filename'))
|
|
file_path = resolve_within(resolve_within(output_dir, request.query.get('subfolder', '')), filename)
|
|
except ValueError:
|
|
return web.json_response({"status": "error", "message": "Invalid parameters"}, status=400)
|
|
if not os.path.isfile(file_path):
|
|
return web.json_response({"status": "error", "message": "File not found"}, status=404)
|
|
served = file_path
|
|
ext = os.path.splitext(file_path)[1].lower()
|
|
if ext in CARD_THUMBNAIL_STATIC_EXTENSIONS:
|
|
served = await asyncio.to_thread(_build_card_thumbnail, file_path)
|
|
elif ext in CARD_POSTER_VIDEO_EXTENSIONS:
|
|
poster = await asyncio.to_thread(_build_video_poster, file_path)
|
|
if poster:
|
|
served = poster
|
|
# Revalidate (cheap 304s): an output file can be replaced under the same name.
|
|
return web.FileResponse(served, headers={'Cache-Control': 'no-cache'})
|
|
|
|
|
|
async def api_delete_gallery_image(request):
|
|
try:
|
|
data = await request.json()
|
|
except ValueError:
|
|
data = None
|
|
if not isinstance(data, dict):
|
|
return web.json_response({"status": "error", "message": "Invalid request body"}, status=400)
|
|
try:
|
|
filename = data.get("filename")
|
|
subfolder = data.get("subfolder", "")
|
|
|
|
output_dir = folder_paths.get_output_directory()
|
|
try:
|
|
filename = require_filename(filename)
|
|
target_dir = resolve_within(output_dir, subfolder)
|
|
file_path = resolve_within(target_dir, filename)
|
|
except ValueError:
|
|
return web.json_response({"status": "error", "message": "Invalid parameters"}, status=400)
|
|
|
|
if os.path.exists(file_path):
|
|
try:
|
|
await asyncio.to_thread(move_to_trash, file_path)
|
|
except OSError as e:
|
|
return web.json_response({"status": "error", "message": trash_failure(e)}, status=409)
|
|
_invalidate_gallery_snapshot()
|
|
return web.json_response({"status": "success"})
|
|
else:
|
|
return web.json_response({"status": "error", "message": "File not found"})
|
|
except Exception as e:
|
|
return web.json_response({"status": "error", "message": str(e)}, status=500)
|