"""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 ; 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 / 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 (" is the woman in ..."). 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