135 lines
5.1 KiB
Python
135 lines
5.1 KiB
Python
import logging
|
|
import os
|
|
import tempfile
|
|
import threading
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
|
|
from .history_db import resolve_image_path, get_data_dir
|
|
|
|
MAX_WIDTH = 600
|
|
QUALITY = 85
|
|
|
|
# Single-writer per hash within this process. The /generation POST generates a
|
|
# thumb on its own to_thread worker AFTER committing the DB row, and any
|
|
# concurrent GET /thumb (the sidebar workflow-cover especially) resolves that
|
|
# committed row and generates the SAME thumb on another worker — both miss the
|
|
# is_file() guard in the post-commit/pre-thumb window. The atomic publish below
|
|
# is the correctness guarantee against torn reads; this lock just collapses the
|
|
# redundant decode+resize. threading.Lock (not asyncio) — runs in worker threads.
|
|
_locks: dict[str, threading.Lock] = {}
|
|
_locks_guard = threading.Lock()
|
|
|
|
|
|
def _lock_for(image_hash: str) -> threading.Lock:
|
|
with _locks_guard:
|
|
return _locks.setdefault(image_hash, threading.Lock())
|
|
|
|
# Video outputs (Wan/AnimateDiff/…) get recorded like images. PIL can't open
|
|
# them, so thumbnail the first decoded frame instead — and never hand a video
|
|
# to Image.open (ultralytics' global PIL patch turns that miss into a noisy,
|
|
# offline-failing pi_heif auto-install attempt).
|
|
VIDEO_SUFFIXES = {".mp4", ".webm", ".mkv", ".mov", ".m4v"}
|
|
|
|
|
|
def _get_thumbs_dir() -> Path:
|
|
return get_data_dir() / "thumbs"
|
|
|
|
|
|
def _first_video_frame(source: Path) -> "Image.Image | None":
|
|
"""Decode a video's first frame as a poster image, or None if it can't."""
|
|
try:
|
|
import av
|
|
except ImportError:
|
|
return None
|
|
try:
|
|
with av.open(str(source)) as container:
|
|
stream = next((s for s in container.streams if s.type == "video"), None)
|
|
if stream is None:
|
|
return None
|
|
for frame in container.decode(stream):
|
|
return frame.to_image()
|
|
except Exception:
|
|
logging.debug("[PromptChain] could not decode video poster for %s",
|
|
source, exc_info=True)
|
|
return None
|
|
|
|
|
|
def _load_poster(source: Path) -> "Image.Image | None":
|
|
"""A still PIL image to thumbnail from: the first frame for video, the file
|
|
itself for images. Fully loaded and detached so the source handle closes."""
|
|
if source.suffix.lower() in VIDEO_SUFFIXES:
|
|
return _first_video_frame(source)
|
|
with Image.open(source) as im:
|
|
im.load()
|
|
return im.copy()
|
|
|
|
|
|
def get_or_create_thumbnail(image_hash: str) -> Path | None:
|
|
thumbs_dir = _get_thumbs_dir()
|
|
thumbs_dir.mkdir(parents=True, exist_ok=True)
|
|
thumb_path = thumbs_dir / f"{image_hash}.webp"
|
|
|
|
if thumb_path.is_file():
|
|
return thumb_path
|
|
|
|
with _lock_for(image_hash):
|
|
# another worker may have published while we waited on the lock
|
|
if thumb_path.is_file():
|
|
return thumb_path
|
|
|
|
source = resolve_image_path(image_hash)
|
|
if not source or not source.is_file():
|
|
return None
|
|
|
|
try:
|
|
img = _load_poster(source)
|
|
if img is None:
|
|
return None
|
|
|
|
if img.mode in ("RGBA", "LA", "P", "PA"):
|
|
background = Image.new("RGB", img.size, (255, 255, 255))
|
|
if img.mode == "P":
|
|
img = img.convert("RGBA")
|
|
background.paste(img, mask=img.split()[-1] if "A" in img.mode else None)
|
|
img = background
|
|
elif img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
|
|
if img.width > MAX_WIDTH:
|
|
ratio = MAX_WIDTH / img.width
|
|
img = img.resize((MAX_WIDTH, round(img.height * ratio)), Image.LANCZOS)
|
|
|
|
# Atomic publish: a concurrent GET /thumb must never FileResponse a
|
|
# half-written webp (PIL opens the path 'wb' = truncate-then-stream).
|
|
# Encode into a sibling temp file (same dir → same volume → os.replace
|
|
# is atomic) and swap it onto the final path in one indivisible step.
|
|
fd, tmp_name = tempfile.mkstemp(
|
|
dir=str(thumbs_dir), prefix=f".{image_hash}.", suffix=".webp.tmp"
|
|
)
|
|
try:
|
|
with os.fdopen(fd, "wb") as f:
|
|
img.save(f, format="WEBP", quality=QUALITY)
|
|
os.replace(tmp_name, thumb_path)
|
|
except PermissionError:
|
|
# Windows can refuse os.replace if the destination is open by a
|
|
# reader. The thumb is content-addressed, so if it now exists a
|
|
# peer already published identical bytes — accept that and drop
|
|
# our temp; otherwise report a clean miss.
|
|
try:
|
|
os.unlink(tmp_name)
|
|
except OSError:
|
|
pass
|
|
return thumb_path if thumb_path.is_file() else None
|
|
except BaseException:
|
|
try:
|
|
os.unlink(tmp_name)
|
|
except OSError:
|
|
pass
|
|
raise
|
|
return thumb_path
|
|
except Exception:
|
|
logging.exception("[PromptChain] thumbnail generation failed for %s", image_hash)
|
|
return None
|