Jony
ba920f2b66
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
141 lines
4.9 KiB
Python
141 lines
4.9 KiB
Python
"""Bounded reference thumbnails for the unified image-generator panel."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from io import BytesIO
|
|
import os
|
|
from typing import Any
|
|
|
|
from PIL import Image, ImageOps
|
|
|
|
|
|
THUMBNAIL_MAX_SIZE = 256
|
|
THUMBNAIL_QUALITY = 82
|
|
THUMBNAIL_CACHE_SECONDS = 3600
|
|
_ALLOWED_FOLDER_TYPES = frozenset({"input", "output", "temp"})
|
|
|
|
|
|
def _is_within(root: str, candidate: str) -> bool:
|
|
try:
|
|
return os.path.commonpath(
|
|
[os.path.normcase(root), os.path.normcase(candidate)]
|
|
) == os.path.normcase(root)
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
def resolve_thumbnail_source(
|
|
folder_paths_module: Any,
|
|
filename: str,
|
|
subfolder: str = "",
|
|
folder_type: str = "input",
|
|
) -> str:
|
|
"""Resolve one public ComfyUI image descriptor without allowing traversal."""
|
|
|
|
folder_type = str(folder_type or "input").strip()
|
|
if folder_type not in _ALLOWED_FOLDER_TYPES:
|
|
raise ValueError("图片目录类型无效")
|
|
|
|
filename = str(filename or "").strip()
|
|
subfolder = str(subfolder or "").strip()
|
|
if not filename or os.path.basename(filename) != filename:
|
|
raise ValueError("图片文件名无效")
|
|
if os.path.isabs(subfolder) or ".." in subfolder.replace("\\", "/").split("/"):
|
|
raise ValueError("图片子目录无效")
|
|
|
|
root_value = folder_paths_module.get_directory_by_type(folder_type)
|
|
if not root_value:
|
|
raise ValueError("图片目录不存在")
|
|
root = os.path.realpath(os.path.abspath(root_value))
|
|
source = os.path.realpath(os.path.abspath(os.path.join(root, subfolder, filename)))
|
|
if not _is_within(root, source):
|
|
raise ValueError("图片路径不安全")
|
|
if not os.path.isfile(source):
|
|
raise FileNotFoundError(filename)
|
|
return source
|
|
|
|
|
|
def render_reference_thumbnail(
|
|
source: str,
|
|
max_size: int = THUMBNAIL_MAX_SIZE,
|
|
quality: int = THUMBNAIL_QUALITY,
|
|
) -> bytes:
|
|
"""Render a small WebP without modifying or replacing the source image."""
|
|
|
|
max_size = max(32, min(int(max_size), THUMBNAIL_MAX_SIZE))
|
|
quality = max(1, min(int(quality), 100))
|
|
with Image.open(source) as opened:
|
|
if getattr(opened, "is_animated", False):
|
|
opened.seek(0)
|
|
# JPEG decoders can use draft mode to avoid materialising every source
|
|
# pixel before the final thumbnail resize. Other formats still remain
|
|
# bounded by the route-level semaphore.
|
|
opened.draft("RGB", (max_size, max_size))
|
|
image = None
|
|
try:
|
|
image = ImageOps.exif_transpose(opened)
|
|
image.thumbnail(
|
|
(max_size, max_size),
|
|
Image.Resampling.LANCZOS,
|
|
reducing_gap=3.0,
|
|
)
|
|
has_alpha = "A" in image.getbands() or "transparency" in image.info
|
|
converted = image.convert("RGBA" if has_alpha else "RGB")
|
|
try:
|
|
output = BytesIO()
|
|
converted.save(output, format="WEBP", quality=quality, method=4)
|
|
return output.getvalue()
|
|
finally:
|
|
converted.close()
|
|
finally:
|
|
if image is not None and image is not opened:
|
|
image.close()
|
|
|
|
|
|
def _thumbnail_etag(source: str) -> str:
|
|
stat = os.stat(source)
|
|
return f'"o1key-thumb-{stat.st_mtime_ns:x}-{stat.st_size:x}"'
|
|
|
|
|
|
def register_o1key_image_thumbnail_route(PromptServer, web, folder_paths_module):
|
|
"""Register the thumbnail endpoint used only by lightweight panel previews."""
|
|
|
|
render_semaphore = asyncio.Semaphore(2)
|
|
|
|
@PromptServer.instance.routes.get("/o1key/image/thumbnail")
|
|
async def get_o1key_image_thumbnail(request):
|
|
try:
|
|
source = resolve_thumbnail_source(
|
|
folder_paths_module,
|
|
request.query.get("filename", ""),
|
|
request.query.get("subfolder", ""),
|
|
request.query.get("type", "input"),
|
|
)
|
|
etag = await asyncio.to_thread(_thumbnail_etag, source)
|
|
headers = {
|
|
"Cache-Control": f"private, max-age={THUMBNAIL_CACHE_SECONDS}",
|
|
"ETag": etag,
|
|
"X-Content-Type-Options": "nosniff",
|
|
}
|
|
if request.headers.get("If-None-Match") == etag:
|
|
return web.Response(status=304, headers=headers)
|
|
async with render_semaphore:
|
|
body = await asyncio.to_thread(render_reference_thumbnail, source)
|
|
return web.Response(body=body, content_type="image/webp", headers=headers)
|
|
except FileNotFoundError:
|
|
return web.Response(status=404)
|
|
except ValueError:
|
|
return web.Response(status=400)
|
|
except Exception:
|
|
# Do not expose local paths or decoder details to the browser.
|
|
return web.Response(status=415)
|
|
|
|
|
|
__all__ = [
|
|
"THUMBNAIL_MAX_SIZE",
|
|
"register_o1key_image_thumbnail_route",
|
|
"render_reference_thumbnail",
|
|
"resolve_thumbnail_source",
|
|
]
|