olivv-cs-ComfyUI-FunPack/core/media.py

"""Media: files the user brought in, rather than something a run produced.

A reference image, an imported clip -- anything meant to feed INTO a
generation instead of coming out of one. Files live under config.MEDIA_DIR as
<id><ext>; a sibling index.json records the display name, kind, size and
timestamp. Pure file/JSON I/O, same shape as projects.py, deliberately not
importing it: a media id and a project id look alike but mean different
things, and a store that borrowed the other's helpers would make that
harder to see, not easier.
"""

from __future__ import annotations

import json
import mimetypes
import re
import threading
import time
import uuid
from pathlib import Path

from . import config

#: Guards the index's read-modify-write. Nothing in save_upload/delete awaits,
#: so on today's single request per event-loop-turn this is never contended --
#: but that safety is an accident of the call path, not something the index
#: itself enforces, and the first `asyncio.to_thread` wrapped around a big
#: upload (core/routes.py already does this for other slow I/O) turns two
#: concurrent writers into a lost update or a crash on the shared tmp file.
_LOCK = threading.Lock()

_ID = re.compile(r"\A[0-9a-f]{12}\Z")

#: What a name is allowed to become on disk. Not a sanitiser: a name that
#: merely survives cleaning can still be "..", so path safety comes from the
#: generated id, never from this.
MAX_NAME = 120

_IMAGE_EXT = frozenset({".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff"})
_VIDEO_EXT = frozenset({".mp4", ".webm", ".mov", ".mkv", ".avi"})
_AUDIO_EXT = frozenset({".mp3", ".wav", ".m4a", ".aac", ".ogg", ".flac", ".opus", ".weba"})

#: What core will accept onto disk at all. A media route is a place a browser
#: can send arbitrary bytes; this is the refusal, not a UI hint -- anything
#: not on this list is rejected before it reaches a filename.
ALLOWED_EXT = _IMAGE_EXT | _VIDEO_EXT | _AUDIO_EXT

#: Matches ComfyUI's own --max-upload-size default (100MB): FunUI's routes
#: ride on ComfyUI's Application in production, which already refuses a bigger
#: body before this ever runs. A larger number here would be a limit this
#: store claims to enforce but never gets the chance to.
MAX_BYTES = 100 * 1024 * 1024


def is_id(value) -> bool:
    return isinstance(value, str) and bool(_ID.match(value))


def _kind(ext: str) -> str:
    if ext in _IMAGE_EXT:
        return "image"
    if ext in _VIDEO_EXT:
        return "video"
    if ext in _AUDIO_EXT:
        return "audio"
    return "other"


def _dir():
    config.MEDIA_DIR.mkdir(parents=True, exist_ok=True)
    return config.MEDIA_DIR


def _index_path() -> Path:
    return _dir() / "index.json"


def _load_index() -> list[dict]:
    p = _index_path()
    if not p.exists():
        return []
    try:
        data = json.loads(p.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError, UnicodeDecodeError):
        return []  # an unreadable index is not a reason to have no list
    return data if isinstance(data, list) else []


def _save_index(items: list[dict]) -> None:
    path = _index_path()
    tmp = path.with_suffix(".tmp")
    tmp.write_text(json.dumps(items, indent=2), encoding="utf-8")
    tmp.replace(path)  # a concurrent read never sees a half-written index


def _has_sound(path: Path) -> bool:
    """Whether a video carries an audio track. Anything unreadable counts as silent."""
    try:
        import av
        with av.open(str(path)) as container:
            return bool(container.streams.audio)
    except Exception:  # noqa: BLE001 -- av missing or an unreadable file: no sound is the safe answer
        return False


def listing() -> list[dict]:
    """Every upload, newest first. An entry whose file vanished is dropped
    rather than shown -- there is nothing behind it to open. A video says whether it has sound
    (`sound`), looked up once and kept: a reference with a soundtrack feeds the model two inputs."""
    items = _load_index()
    fresh = False
    for it in items:
        path = config.MEDIA_DIR / it.get("filename", "")
        if it.get("kind") == "video" and "sound" not in it and path.is_file():
            it["sound"] = _has_sound(path)
            fresh = True
    if fresh:
        with _LOCK:
            _save_index(items)
    out = [it for it in items if not it.get("scratch") and (config.MEDIA_DIR / it.get("filename", "")).is_file()]      # scratch files are not the user's: never listed
    out.sort(key=lambda i: i.get("added", 0), reverse=True)
    return out


def get(mid: str) -> dict | None:
    if not is_id(mid):
        return None
    return next((it for it in _load_index() if it.get("id") == mid), None)


def save_upload(orig_name: str, data: bytes, limit=...) -> dict:
    """Write one uploaded file and record it. Raises ValueError for anything
    this store refuses -- an empty upload, an extension not on the allowlist,
    or a file over MAX_BYTES -- so the route can turn that into a 400 naming
    why, rather than a file silently written and then never explained."""
    if not data:
        raise ValueError("that file is empty")
    limit = MAX_BYTES if limit is ... else limit       # an upload's limit (read now); a file the server made itself has none
    if limit is not None and len(data) > limit:
        raise ValueError(f"that file is over the {limit // (1024 * 1024)}MB limit")
    ext = Path(orig_name or "").suffix.lower()
    if ext not in ALLOWED_EXT:
        raise ValueError(f"{ext or 'that file type'} is not something FunUI imports")

    mid = uuid.uuid4().hex[:12]
    filename = mid + ext
    _dir()
    (config.MEDIA_DIR / filename).write_bytes(data)
    entry = {
        "id": mid,
        "name": (Path(orig_name or "").name or filename)[:MAX_NAME],
        "filename": filename,
        "kind": _kind(ext),
        "size": len(data),
        "added": time.time(),
    }
    if entry["kind"] == "video":
        entry["sound"] = _has_sound(config.MEDIA_DIR / filename)
    # The write above needs no lock -- filename is this call's own uuid, so two
    # uploads never touch the same file. The index below is shared, and a
    # read here that raced another upload's write would silently drop it.
    with _LOCK:
        items = _load_index()
        items.append(entry)
        _save_index(items)
    return entry


def save_scratch(orig_name: str, data: bytes, folder) -> dict:
    """A picture the app made for one run (a carried last frame): written into `folder` (ComfyUI's temp folder), answerable by id
    like any media but never listed in the bin. It goes when the folder is cleared; entries whose file is gone are dropped here."""
    if not data:
        raise ValueError("that file is empty")
    ext = Path(orig_name or "").suffix.lower()
    if ext not in _IMAGE_EXT:
        raise ValueError(f"{ext or 'that file type'} is not a picture")
    mid = uuid.uuid4().hex[:12]
    target = Path(folder) / (mid + ext)
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_bytes(data)
    # `filename` is the absolute path: MEDIA_DIR / <absolute> is that path, so path_for, delete and the thumbnails need no second case.
    entry = {"id": mid, "name": (Path(orig_name or "").name or target.name)[:MAX_NAME], "filename": str(target), "kind": "image",
             "size": len(data), "added": time.time(), "scratch": True}
    with _LOCK:
        items = [it for it in _load_index() if not (it.get("scratch") and not Path(it.get("filename", "")).is_file())]
        items.append(entry)
        _save_index(items)
    return entry


def path_for(mid: str) -> Path | None:
    it = get(mid)
    if not it:
        return None
    p = config.MEDIA_DIR / it["filename"]
    return p if p.is_file() else None


def content_type(mid: str) -> str:
    it = get(mid)
    if not it:
        return "application/octet-stream"
    return mimetypes.guess_type(it["filename"])[0] or "application/octet-stream"


def delete(mid: str) -> bool:
    with _LOCK:
        items = _load_index()
        keep, removed = [], None
        for it in items:
            if it.get("id") == mid:
                removed = it
            else:
                keep.append(it)
        if removed is None:
            return False
        _save_index(keep)
    for gone in (config.MEDIA_DIR / removed["filename"], _thumb_file(mid)):
        try:
            gone.unlink(missing_ok=True)
        except OSError:
            pass
    return True


def rename(mid: str, name) -> dict | None:
    """Change the display name only: the file on disk is addressed by id and never moves."""
    name = " ".join(str(name or "").split())[:MAX_NAME]
    if not name:
        raise ValueError("a name is needed")
    with _LOCK:
        items = _load_index()
        for it in items:
            if it.get("id") == mid:
                it["name"] = name
                _save_index(items)
                return it
    return None


MAX_SUBJECT = 600


def set_subject(mid: str, text) -> dict | None:
    """What this picture MEANS when it is a reference ("<Subject 1> is the woman in <Picture 1>..."). Empty clears it."""
    text = " ".join(str(text or "").split())[:MAX_SUBJECT]
    with _LOCK:
        items = _load_index()
        for it in items:
            if it.get("id") == mid:
                if text:
                    it["subject"] = text
                else:
                    it.pop("subject", None)
                _save_index(items)
                return it
    return None


#: Long edge of a grid thumbnail in px: headroom over the largest cell (240) for retina, and
#: far below the 4K original the grid used to download and decode in full just to draw a square.
THUMB_MAX_DIM = 320


def _thumb_file(mid: str) -> Path:
    return config.MEDIA_DIR / "thumbnails" / f"{mid}.jpg"


def thumb_path_for(mid: str) -> Path | None:
    """A small cached JPEG for an image or video, made on first request. None for audio, an
    unreadable source, or no ffmpeg for a video: the grid then draws its placeholder. Blocking."""
    it = get(mid)
    src = path_for(mid) if it and it.get("kind") in ("image", "video") else None
    if src is None:
        return None
    out = _thumb_file(mid)
    try:
        if out.is_file() and out.stat().st_mtime >= src.stat().st_mtime:
            return out
        out.parent.mkdir(parents=True, exist_ok=True)
        # Unique per attempt: two requests for a not-yet-cached id must not interleave in one temp file.
        tmp = out.with_name(f"{out.stem}.{uuid.uuid4().hex[:8]}.tmp")
        try:
            if it["kind"] == "image":
                from PIL import Image
                from PIL import ImageOps
                img = ImageOps.exif_transpose(Image.open(src))      # the browser shows it turned upright
                img.thumbnail((THUMB_MAX_DIM, THUMB_MAX_DIM))      # in place, keeps aspect, never upscales
                (img if img.mode in ("RGB", "L") else img.convert("RGB")).save(tmp, "JPEG", quality=82)
            else:
                import shutil
                import subprocess
                ff = shutil.which("ffmpeg")
                if not ff:
                    return None
                box = f"scale='min({THUMB_MAX_DIM},iw)':'min({THUMB_MAX_DIM},ih)':force_original_aspect_ratio=decrease"
                # -f mjpeg: the temp name's extension gives ffmpeg nothing to infer a container from
                subprocess.run([ff, "-y", "-i", str(src), "-vframes", "1", "-vf", box, "-f", "mjpeg", str(tmp)],
                               capture_output=True, check=True)
            tmp.replace(out)
        finally:
            tmp.unlink(missing_ok=True)
        return out if out.is_file() else None
    except Exception:                                  # noqa: BLE001 -- a thumbnail is never worth an error
        return None