191 lines
6.3 KiB
Python
191 lines
6.3 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Disk-backed persistence for generated images.
|
|
|
|
Each image is a PNG under ``studio_root()/images`` with its full generation
|
|
recipe embedded as PNG text chunks: a structured ``unsloth`` JSON blob (the
|
|
source of truth the gallery reads back) plus an Automatic1111-style
|
|
``parameters`` string for interop with other tools. Because the recipe lives
|
|
inside the file, a downloaded PNG carries its own settings.
|
|
|
|
The gallery is intentionally dumb storage: the route owns the metadata schema
|
|
and passes a plain dict; this module only writes/reads/sorts files.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
import re
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any, Optional
|
|
|
|
from loggers import get_logger
|
|
from utils.paths import ensure_dir, studio_root
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
# PNG text-chunk key holding our structured recipe JSON.
|
|
_META_KEY = "unsloth"
|
|
# Image ids are file stems; restrict to filename-safe chars so a crafted id
|
|
# can't escape the gallery directory.
|
|
_ID_RE = re.compile(r"^[A-Za-z0-9_-]{1,128}$")
|
|
|
|
|
|
def gallery_dir() -> Path:
|
|
return ensure_dir(studio_root() / "images")
|
|
|
|
|
|
def _params_text(meta: dict[str, Any]) -> str:
|
|
"""Automatic1111-style ``parameters`` string for cross-tool interop."""
|
|
lines = [str(meta.get("prompt", ""))]
|
|
negative = meta.get("negative_prompt")
|
|
if negative:
|
|
lines.append(f"Negative prompt: {negative}")
|
|
lines.append(
|
|
f"Steps: {meta.get('steps')}, CFG scale: {meta.get('guidance')}, "
|
|
f"Seed: {meta.get('seed')}, Size: {meta.get('width')}x{meta.get('height')}, "
|
|
f"Model: {meta.get('model', '')}"
|
|
)
|
|
return "\n".join(lines)
|
|
|
|
|
|
def _png_bytes(image: Any, meta: dict[str, Any]) -> bytes:
|
|
import io
|
|
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
info = PngInfo()
|
|
info.add_text(_META_KEY, json.dumps(meta))
|
|
info.add_text("parameters", _params_text(meta))
|
|
buf = io.BytesIO()
|
|
image.save(buf, format = "PNG", pnginfo = info)
|
|
return buf.getvalue()
|
|
|
|
|
|
def save(image: Any, meta: dict[str, Any]) -> dict[str, Any]:
|
|
"""Persist a PIL image with its recipe embedded; return the gallery record."""
|
|
image_id = uuid.uuid4().hex
|
|
(gallery_dir() / f"{image_id}.png").write_bytes(_png_bytes(image, meta))
|
|
return _record(image_id, meta)
|
|
|
|
|
|
def _record(image_id: str, meta: dict[str, Any]) -> dict[str, Any]:
|
|
return {
|
|
**meta,
|
|
"id": image_id,
|
|
"url": f"/api/inference/images/gallery/{image_id}/file",
|
|
}
|
|
|
|
|
|
def image_path(image_id: str) -> Optional[Path]:
|
|
"""Resolve an id to its on-disk PNG, or None if missing / unsafe."""
|
|
if not _ID_RE.match(image_id):
|
|
return None
|
|
path = gallery_dir() / f"{image_id}.png"
|
|
# Defence in depth: confirm the resolved path is still inside the gallery.
|
|
try:
|
|
path.resolve().relative_to(gallery_dir().resolve())
|
|
except ValueError:
|
|
return None
|
|
return path if path.is_file() else None
|
|
|
|
|
|
def image_b64(image_id: str) -> Optional[str]:
|
|
path = image_path(image_id)
|
|
if path is None:
|
|
return None
|
|
return base64.b64encode(path.read_bytes()).decode("ascii")
|
|
|
|
|
|
# Recipe keys a gallery record must carry (the required GalleryImage fields, minus
|
|
# id/url which _record adds). A PNG missing any is treated as foreign and skipped,
|
|
# so a hand-dropped or older-schema file can't 500 the whole listing.
|
|
_REQUIRED_META = ("prompt", "width", "height", "steps", "guidance", "seed", "created_at")
|
|
|
|
|
|
def _read_meta(path: Path) -> Optional[dict[str, Any]]:
|
|
from PIL import Image
|
|
|
|
try:
|
|
with Image.open(path) as im:
|
|
raw = im.text.get(_META_KEY) # type: ignore[attr-defined]
|
|
except Exception:
|
|
return None
|
|
if not raw:
|
|
return None
|
|
try:
|
|
meta = json.loads(raw)
|
|
except (ValueError, TypeError):
|
|
return None
|
|
if not isinstance(meta, dict) or any(k not in meta for k in _REQUIRED_META):
|
|
return None
|
|
return meta
|
|
|
|
|
|
def _mtime(path: Path) -> float:
|
|
try:
|
|
return path.stat().st_mtime
|
|
except OSError:
|
|
return 0.0
|
|
|
|
|
|
def list_images(limit: Optional[int] = None, offset: int = 0) -> list[dict[str, Any]]:
|
|
"""A newest-first window of images for infinite scroll.
|
|
|
|
Ordered by file mtime (a cheap stat, ~= generation order) so a months-old
|
|
gallery isn't opened in full just to sort it; only the window's recipes are
|
|
read. limit=None returns everything from ``offset`` on."""
|
|
try:
|
|
paths = list(gallery_dir().glob("*.png"))
|
|
except OSError:
|
|
return []
|
|
paths.sort(key = _mtime, reverse = True)
|
|
# Page over READABLE records, not raw files: filtering a foreign/corrupt PNG out of an
|
|
# already-sliced window would drop valid images that sort after it and make the route's
|
|
# has_more wrong. Read only as far as needed to fill the requested window.
|
|
# Known Phase-1 limit: this re-reads headers from the newest down to `offset+limit` on
|
|
# every page, so a deep infinite-scroll over a very large gallery (thousands of images,
|
|
# e.g. a long uncapped batch) is O(offset) header-opens per page. PIL opens are lazy
|
|
# (header only) and this runs off the event loop, so it's not a freeze; a later phase can
|
|
# switch to cursor-based paging (resume after the last-seen record) if it starts to bite.
|
|
want = None if limit is None else offset + limit
|
|
records = []
|
|
for path in paths:
|
|
meta = _read_meta(path)
|
|
if meta is None: # not one of ours (no recipe chunk) — skip
|
|
continue
|
|
records.append(_record(path.stem, meta))
|
|
if want is not None and len(records) >= want:
|
|
break
|
|
return records[offset:] if limit is None else records[offset : offset + limit]
|
|
|
|
|
|
def delete(image_id: str) -> bool:
|
|
path = image_path(image_id)
|
|
if path is None:
|
|
return False
|
|
try:
|
|
path.unlink()
|
|
return True
|
|
except OSError as exc:
|
|
logger.warning("image_gallery.delete_failed: %s", exc)
|
|
return False
|
|
|
|
|
|
def clear() -> int:
|
|
"""Delete every gallery PNG; return how many were removed."""
|
|
removed = 0
|
|
try:
|
|
paths = list(gallery_dir().glob("*.png"))
|
|
except OSError:
|
|
return 0
|
|
for path in paths:
|
|
try:
|
|
path.unlink()
|
|
removed += 1
|
|
except OSError:
|
|
continue
|
|
return removed
|