277 lines
9.9 KiB
Python
277 lines
9.9 KiB
Python
"""Model cover, thumbnail-cache, and generated-image lookup routes."""
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import struct
|
|
import threading
|
|
import urllib.parse
|
|
|
|
from aiohttp import web
|
|
import folder_paths
|
|
|
|
from .path_utils import require_filename, resolve_folder_subdir, resolve_within
|
|
|
|
|
|
CARD_THUMBNAIL_EDGE = 512
|
|
CARD_THUMBNAIL_CACHE_LIMIT = 256 * 1024 * 1024
|
|
CARD_THUMBNAIL_STATIC_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.webp', '.avif'}
|
|
_thumbnail_cleanup_lock = threading.Lock()
|
|
_thumbnail_generation_slots = threading.BoundedSemaphore(2)
|
|
_thumbnail_last_cleanup = 0.0
|
|
_thumbnail_created_since_cleanup = 0
|
|
|
|
|
|
def _thumbnail_cache_directory():
|
|
try:
|
|
base_dir = folder_paths.get_temp_directory()
|
|
except Exception:
|
|
base_dir = None
|
|
if not isinstance(base_dir, (str, os.PathLike)):
|
|
base_dir = tempfile.gettempdir()
|
|
cache_dir = os.path.join(base_dir, "anomalous_model_browser", "card_thumbnails")
|
|
os.makedirs(cache_dir, exist_ok=True)
|
|
return cache_dir
|
|
|
|
|
|
def _prune_thumbnail_cache(cache_dir, created=False):
|
|
"""Occasionally cap the derived thumbnail cache without touching source covers."""
|
|
global _thumbnail_last_cleanup, _thumbnail_created_since_cleanup
|
|
now = time.monotonic()
|
|
with _thumbnail_cleanup_lock:
|
|
if created:
|
|
_thumbnail_created_since_cleanup += 1
|
|
now = time.monotonic()
|
|
if (
|
|
now - _thumbnail_last_cleanup < 3600
|
|
and _thumbnail_created_since_cleanup < 64
|
|
):
|
|
return
|
|
_thumbnail_last_cleanup = now
|
|
_thumbnail_created_since_cleanup = 0
|
|
entries = []
|
|
total_size = 0
|
|
try:
|
|
with os.scandir(cache_dir) as iterator:
|
|
for entry in iterator:
|
|
if not entry.is_file() or not entry.name.endswith('.webp'):
|
|
continue
|
|
try:
|
|
stat = entry.stat()
|
|
except OSError:
|
|
continue
|
|
total_size += stat.st_size
|
|
entries.append((stat.st_mtime_ns, stat.st_size, entry.path))
|
|
except OSError:
|
|
return
|
|
if total_size <= CARD_THUMBNAIL_CACHE_LIMIT:
|
|
return
|
|
target_size = int(CARD_THUMBNAIL_CACHE_LIMIT * 0.8)
|
|
for _, size, path in sorted(entries):
|
|
try:
|
|
os.remove(path)
|
|
total_size -= size
|
|
except OSError:
|
|
pass
|
|
if total_size <= target_size:
|
|
break
|
|
|
|
|
|
def _build_card_thumbnail_impl(source_path):
|
|
"""Return a cached 512px WebP card image, or the original on any safe fallback."""
|
|
try:
|
|
from PIL import Image, ImageOps
|
|
|
|
source_stat = os.stat(source_path)
|
|
cache_key = "\0".join((
|
|
os.path.realpath(source_path),
|
|
str(source_stat.st_size),
|
|
str(source_stat.st_mtime_ns),
|
|
str(getattr(source_stat, 'st_ctime_ns', 0)),
|
|
str(CARD_THUMBNAIL_EDGE),
|
|
))
|
|
digest = hashlib.sha256(cache_key.encode('utf-8', errors='surrogatepass')).hexdigest()
|
|
cache_dir = _thumbnail_cache_directory()
|
|
cached_path = os.path.join(cache_dir, f"{digest}.webp")
|
|
if os.path.isfile(cached_path):
|
|
try:
|
|
os.utime(cached_path, None)
|
|
except OSError:
|
|
pass
|
|
_prune_thumbnail_cache(cache_dir)
|
|
return cached_path
|
|
|
|
with Image.open(source_path) as image:
|
|
if getattr(image, 'is_animated', False):
|
|
return source_path
|
|
image = ImageOps.exif_transpose(image)
|
|
if max(image.size) <= CARD_THUMBNAIL_EDGE:
|
|
return source_path
|
|
resampling = getattr(Image, 'Resampling', Image).LANCZOS
|
|
image.thumbnail((CARD_THUMBNAIL_EDGE, CARD_THUMBNAIL_EDGE), resampling)
|
|
if image.mode not in ('RGB', 'RGBA'):
|
|
image = image.convert('RGBA' if 'transparency' in image.info else 'RGB')
|
|
temp_path = (
|
|
f"{cached_path}.{os.getpid()}.{threading.get_ident()}.tmp"
|
|
)
|
|
try:
|
|
image.save(temp_path, format='WEBP', quality=84, method=4)
|
|
os.replace(temp_path, cached_path)
|
|
finally:
|
|
if os.path.exists(temp_path):
|
|
try:
|
|
os.remove(temp_path)
|
|
except OSError:
|
|
pass
|
|
_prune_thumbnail_cache(cache_dir, created=True)
|
|
return cached_path
|
|
except Exception:
|
|
return source_path
|
|
|
|
|
|
def _build_card_thumbnail(source_path):
|
|
# Bound simultaneous decodes so opening a large folder cannot monopolize CPU/RAM.
|
|
with _thumbnail_generation_slots:
|
|
return _build_card_thumbnail_impl(source_path)
|
|
|
|
|
|
async def api_serve_image(request):
|
|
"""Dedicated image serving endpoint for model preview images."""
|
|
folder_type = request.query.get('type', 'checkpoints')
|
|
try:
|
|
path_idx = int(request.query.get('path_idx', 0))
|
|
except:
|
|
path_idx = 0
|
|
subfolder = request.query.get('subfolder', '')
|
|
filename = request.query.get('filename', '')
|
|
|
|
try:
|
|
filename = require_filename(filename)
|
|
_, target_dir = resolve_folder_subdir(folder_type, path_idx, subfolder)
|
|
file_path = resolve_within(target_dir, filename)
|
|
except (ValueError, KeyError):
|
|
return web.Response(status=400, text='Invalid request')
|
|
|
|
if not os.path.exists(file_path) or not os.path.isfile(file_path):
|
|
return web.Response(status=404, text='Image not found')
|
|
|
|
ext = os.path.splitext(filename)[1].lower()
|
|
content_types = {
|
|
'.png': 'image/png',
|
|
'.jpg': 'image/jpeg',
|
|
'.jpeg': 'image/jpeg',
|
|
'.webp': 'image/webp',
|
|
'.gif': 'image/gif',
|
|
'.avif': 'image/avif',
|
|
'.mp4': 'video/mp4',
|
|
'.webm': 'video/webm',
|
|
'.mov': 'video/quicktime',
|
|
'.avi': 'video/x-msvideo'
|
|
}
|
|
content_type = content_types.get(ext)
|
|
if content_type is None:
|
|
return web.Response(status=415, text='Unsupported media type')
|
|
|
|
served_path = file_path
|
|
if (
|
|
request.query.get('variant') == 'card'
|
|
and ext in CARD_THUMBNAIL_STATIC_EXTENSIONS
|
|
):
|
|
served_path = await asyncio.to_thread(_build_card_thumbnail, file_path)
|
|
if served_path != file_path:
|
|
content_type = 'image/webp'
|
|
|
|
headers = {'Content-Type': content_type}
|
|
if request.query.get('t'):
|
|
headers['Cache-Control'] = 'public, max-age=31536000, immutable'
|
|
return web.FileResponse(served_path, headers=headers)
|
|
|
|
|
|
async def api_clear_cache(request):
|
|
try:
|
|
from .metadata import clear_metadata_cache
|
|
|
|
if hasattr(folder_paths, "filename_list_cache"):
|
|
folder_paths.filename_list_cache.clear()
|
|
if hasattr(folder_paths, "cache_helper") and hasattr(folder_paths.cache_helper, "clear"):
|
|
folder_paths.cache_helper.clear()
|
|
clear_metadata_cache()
|
|
return web.json_response({"status": "success"})
|
|
except Exception as e:
|
|
return web.json_response({"status": "error", "message": str(e)})
|
|
|
|
|
|
def read_png_text_fast(path):
|
|
try:
|
|
with open(path, 'rb') as f:
|
|
signature = f.read(8)
|
|
if signature != b'\x89PNG\r\n\x1a\n':
|
|
return None
|
|
while True:
|
|
length_bytes = f.read(4)
|
|
if not length_bytes: break
|
|
length = struct.unpack('>I', length_bytes)[0]
|
|
chunk_type = f.read(4)
|
|
if chunk_type == b'tEXt':
|
|
data = f.read(length)
|
|
keyword, text = data.split(b'\0', 1)
|
|
if keyword == b'prompt':
|
|
return text.decode('utf-8', errors='ignore')
|
|
else:
|
|
f.seek(length, 1)
|
|
f.seek(4, 1)
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
async def api_get_model_images(request):
|
|
model_name = request.rel_url.query.get('model_name', '')
|
|
if not model_name:
|
|
return web.json_response({'images': []})
|
|
|
|
base_target = os.path.basename(model_name).lower()
|
|
|
|
# We will search the output directory
|
|
output_dir = folder_paths.get_output_directory()
|
|
if not os.path.exists(output_dir):
|
|
return web.json_response({'images': []})
|
|
|
|
def find_images():
|
|
matched_images = []
|
|
for root, _, files in os.walk(output_dir):
|
|
for file in files:
|
|
if not file.lower().endswith('.png'):
|
|
continue
|
|
full_path = os.path.join(root, file)
|
|
prompt_text = read_png_text_fast(full_path)
|
|
if not prompt_text:
|
|
continue
|
|
try:
|
|
prompt_data = json.loads(prompt_text)
|
|
matched = any(
|
|
os.path.basename(v).lower() == base_target
|
|
for node in prompt_data.values()
|
|
if isinstance(node, dict) and 'class_type' in node
|
|
for v in node.get('inputs', {}).values()
|
|
if isinstance(v, str)
|
|
)
|
|
if not matched:
|
|
continue
|
|
rel_path = os.path.relpath(root, output_dir).replace('\\', '/')
|
|
if rel_path == '.':
|
|
rel_path = ''
|
|
url = f'/view?filename={urllib.parse.quote(file)}&type=output'
|
|
if rel_path:
|
|
url += f'&subfolder={urllib.parse.quote(rel_path)}'
|
|
matched_images.append({'url': url, 'mtime': os.path.getmtime(full_path)})
|
|
except (OSError, ValueError, TypeError, json.JSONDecodeError):
|
|
pass
|
|
matched_images.sort(key=lambda x: x['mtime'], reverse=True)
|
|
return matched_images
|
|
|
|
matched_images = await asyncio.to_thread(find_images)
|
|
|
|
return web.json_response({'images': matched_images})
|