diff --git "a/app.py" "b/app.py" new file mode 100644--- /dev/null +++ "b/app.py" @@ -0,0 +1,2840 @@ +""" +Sozai — hybrid room collaboration, served with gradio.Server (Server mode). + +Gradio is used ONLY as the backend (FastAPI + Gradio's API engine). The whole +UI is custom and comes from these files: + + landing.html the "Sozai" splash (make a room / join by code / skip) + app.html the collaborative album app (the React port from before) + +What this adds on top of the previous version: + * SQLite database (created locally on first run -> sozai_rooms.db) holding + `rooms` and `participants`, matching the project spec. + * Hybrid room model: secure room_id (uuid4) + short human room code whose + LENGTH THE USER CHOOSES (4-12) on the splash screen. + * REST endpoints: create room, get room state, resolve join-by-code. + * A WebSocket endpoint /ws/{room_id} for real-time sync (presence, chat, + shared selection, captions, live cursors, ...). + * A background cleanup thread that expires idle (1h) and old (24h) rooms. + * A "continue without a room" path: /app runs the app solo, no collaboration. + +Run it: + pip install gradio + python app.py # -> http://localhost:7860 +Drop app.py, landing.html, app.html next to storybook_render.zip (or an +extracted public/ folder). Assets are read straight out of the zip. +""" + +from __future__ import annotations + +import os + +# --------------------------------------------------------------------------- # +# ZeroGPU process-start setup. These MUST be set before NeMo / torch are +# imported, because on Hugging Face ZeroGPU the GPU is not attached at process +# start and CUDA must NOT be initialized in the main process. NeMo's RNNT path +# (and numba) can otherwise trigger a low-level CUDA init during import/load, +# which surfaces as: "Low-level CUDA init (torch._C._cuda_init) reached. This +# means ZeroGPU's PyTorch CUDA emulation mode did not intercept a CUDA +# operation". This block mirrors NVIDIA's reference Nemotron/Parakeet Spaces. +# --------------------------------------------------------------------------- # +_omp = os.environ.get("OMP_NUM_THREADS", "") +if _omp and not _omp.isdigit(): + os.environ["OMP_NUM_THREADS"] = str( + max(1, int("".join(c for c in _omp if c.isdigit()) or "1")) + ) +# Stop NeMo's RNNT/numba path from initializing CUDA in the main process. +os.environ.setdefault("NUMBA_DISABLE_CUDA", "1") +# Help torch's allocator play nicely with ZeroGPU's stateless GPU. +os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + +import asyncio +import functools +import hashlib +import io +import json +import mimetypes +import re +import secrets +import sqlite3 +import threading +import time +import uuid + +from fastapi import HTTPException, Request, WebSocket, WebSocketDisconnect +from fastapi.responses import FileResponse, HTMLResponse, Response + +from gradio import Server + +# --------------------------------------------------------------------------- # +# ZeroGPU support. +# +# On Hugging Face *ZeroGPU* Spaces, a GPU is attached only while a function +# decorated with `@spaces.GPU` is running, and the runtime REFUSES TO START if +# it can't find at least one such function during startup (that is the +# "No @spaces.GPU function detected during startup" error). The `spaces` +# package is preinstalled in the Space image. +# +# Locally (and on any non-ZeroGPU box) the `spaces` package usually isn't +# installed. The decorator is documented to be effect-free off ZeroGPU, so we +# provide a tiny no-op fallback that simply returns the function unchanged. +# This lets the SAME code run both locally and in a Space with no branching: +# * In a ZeroGPU Space -> real `spaces.GPU` requests/releases the GPU. +# * Anywhere else -> `spaces.GPU` is a transparent pass-through. +try: + import spaces # provided by the HF ZeroGPU Space image + _ZEROGPU = True +except Exception: # not on ZeroGPU (e.g. local dev, CPU Space, plain GPU Space) + _ZEROGPU = False + + class _SpacesShim: + """Minimal stand-in for the `spaces` module off ZeroGPU. + + Supports both `@spaces.GPU` and `@spaces.GPU(...)` call styles and just + returns the wrapped function untouched.""" + + @staticmethod + def GPU(fn=None, **_kwargs): + if fn is None: # used as @spaces.GPU(duration=...) + def _wrap(f): + return f + return _wrap + return fn # used as @spaces.GPU + + spaces = _SpacesShim() # type: ignore + +# Load environment variables from a .env file (if present) so SOZAI_*, +# PHOENIX_*, and the Cesium ion token can be set without exporting them. +# Done early so every os.environ.get(...) below sees the values. +try: + from dotenv import load_dotenv + load_dotenv() +except Exception: # python-dotenv not installed → rely on the real environment + pass + +BASE = os.path.dirname(os.path.abspath(__file__)) + +# --------------------------------------------------------------------------- # +# Frontend assets: read from a real public/ dir if present, else from the zip. +# --------------------------------------------------------------------------- # + +def _first_existing(paths): + return next((p for p in paths if p and os.path.exists(p)), None) + + +PUBLIC_DIR = _first_existing([ + os.environ.get("STORYBOOK_PUBLIC"), + os.path.join(BASE, "public"), + os.path.join(BASE, "storybook_render", "public"), + os.path.join(BASE, "src", "public"), +]) +ZIP_PATH = _first_existing([ + os.environ.get("STORYBOOK_ZIP"), + os.path.join(BASE, "storybook_render.zip"), +]) +_ZIP_PREFIX = None + + +@functools.lru_cache(maxsize=512) +def _read_from_zip(rel_path: str) -> bytes: + import zipfile + global _ZIP_PREFIX + with zipfile.ZipFile(ZIP_PATH) as z: + if _ZIP_PREFIX is None: + _ZIP_PREFIX = "" + for n in z.namelist(): + idx = n.find("public/") + if idx != -1: + _ZIP_PREFIX = n[:idx] + break + for name in (f"{_ZIP_PREFIX}public/{rel_path}", f"public/{rel_path}", rel_path): + try: + return z.read(name) + except KeyError: + continue + raise KeyError(rel_path) + + +if PUBLIC_DIR is None and ZIP_PATH is None: + raise SystemExit( + "Could not find frontend assets. Put storybook_render.zip next to " + "app.py, or extract its public/ folder alongside app.py." + ) + +# --------------------------------------------------------------------------- # +# Easter-egg pet sprites: the 2023 oneko-style icon library (76 animals, 32 +# sprites each). Served from a real folder if present, else from neko-main.zip. +# Each animal lives at /2023-icon-library//.png . +# --------------------------------------------------------------------------- # + +SPRITES_DIR = _first_existing([ + os.environ.get("SOZAI_SPRITES_DIR"), + os.path.join(BASE, "2023-icon-library"), + os.path.join(BASE, "public", "2023-icon-library"), + os.path.join(BASE, "sprites"), +]) +SPRITES_ZIP = _first_existing([ + os.environ.get("SOZAI_SPRITES_ZIP"), + os.path.join(BASE, "neko-main.zip"), +]) +_SPRITES_ZIP_PREFIX = None + +# The canonical sprite filenames shared by every animal in the library. +SPRITE_FILES = [ + "alert", "still", "sleep1", "sleep2", "wash", "yawn", "itch1", "itch2", + "nrun1", "nrun2", "nerun1", "nerun2", "erun1", "erun2", "serun1", "serun2", + "srun1", "srun2", "swrun1", "swrun2", "wrun1", "wrun2", "nwrun1", "nwrun2", + "nscratch1", "nscratch2", "escratch1", "escratch2", + "sscratch1", "sscratch2", "wscratch1", "wscratch2", +] + + +@functools.lru_cache(maxsize=1) +def list_sprite_animals(): + """Return the sorted list of available animal sprite-set names.""" + names = set() + if SPRITES_DIR and os.path.isdir(SPRITES_DIR): + for d in os.listdir(SPRITES_DIR): + if os.path.isdir(os.path.join(SPRITES_DIR, d)): + names.add(d) + elif SPRITES_ZIP: + import zipfile + try: + with zipfile.ZipFile(SPRITES_ZIP) as z: + for n in z.namelist(): + idx = n.find("2023-icon-library/") + if idx == -1: + continue + rest = n[idx + len("2023-icon-library/"):] + animal = rest.split("/", 1)[0] + if animal and "." not in animal: + names.add(animal) + except Exception: + pass + return sorted(names) + + +@functools.lru_cache(maxsize=4096) +def _read_sprite(animal: str, fname: str) -> bytes: + """Read a single sprite PNG for an animal from folder or zip.""" + safe_animal = "".join(ch for ch in animal if ch.isalnum() or ch in "-_") + safe_file = "".join(ch for ch in fname if ch.isalnum() or ch in "-_") + rel = f"2023-icon-library/{safe_animal}/{safe_file}.png" + if SPRITES_DIR and os.path.isdir(SPRITES_DIR): + fp = os.path.join(SPRITES_DIR, safe_animal, safe_file + ".png") + if os.path.isfile(fp): + with open(fp, "rb") as fh: + return fh.read() + if SPRITES_ZIP: + import zipfile + global _SPRITES_ZIP_PREFIX + with zipfile.ZipFile(SPRITES_ZIP) as z: + if _SPRITES_ZIP_PREFIX is None: + _SPRITES_ZIP_PREFIX = "" + for n in z.namelist(): + idx = n.find("2023-icon-library/") + if idx != -1: + _SPRITES_ZIP_PREFIX = n[:idx] + break + for name in (f"{_SPRITES_ZIP_PREFIX}{rel}", rel): + try: + return z.read(name) + except KeyError: + continue + raise KeyError(rel) + + +# --------------------------------------------------------------------------- # +# SQLite — created locally on first run. +# --------------------------------------------------------------------------- # + +DB_PATH = os.environ.get("SOZAI_DB", os.path.join(BASE, "sozai_rooms.db")) +DB_LOCK = threading.RLock() +_conn = sqlite3.connect(DB_PATH, check_same_thread=False) +_conn.row_factory = sqlite3.Row + + +def init_db(): + with DB_LOCK: + # auto_vacuum must be set BEFORE tables exist to take effect on a fresh + # DB; on an existing file it only applies after a full VACUUM (below). + # INCREMENTAL lets the cleanup loop reclaim freed pages to the OS so the + # file doesn't keep growing as rooms are created and deleted. + _conn.execute("PRAGMA auto_vacuum=INCREMENTAL;") + _conn.execute("PRAGMA journal_mode=WAL;") + _conn.executescript( + """ + CREATE TABLE IF NOT EXISTS rooms ( + room_id TEXT PRIMARY KEY, + share_code TEXT UNIQUE, + code_length INTEGER, + created_at REAL, + last_activity REAL, + owner_session_id TEXT, + is_active INTEGER DEFAULT 1, + metadata TEXT + ); + CREATE TABLE IF NOT EXISTS participants ( + participant_id TEXT PRIMARY KEY, + room_id TEXT, + session_id TEXT, + display_name TEXT, + joined_at REAL, + last_seen REAL + ); + CREATE TABLE IF NOT EXISTS room_state ( + room_id TEXT PRIMARY KEY, + state TEXT, + version INTEGER DEFAULT 0, + updated_at REAL + ); + CREATE INDEX IF NOT EXISTS idx_rooms_code ON rooms(share_code); + CREATE INDEX IF NOT EXISTS idx_rooms_activity ON rooms(last_activity); + CREATE INDEX IF NOT EXISTS idx_part_room ON participants(room_id); + """ + ) + _conn.commit() + # If this is an existing DB created without auto_vacuum, a one-time full + # VACUUM is needed for the INCREMENTAL setting to take hold. Cheap on a + # small/empty DB; safe to run on every startup. + try: + mode = _conn.execute("PRAGMA auto_vacuum;").fetchone()[0] + if mode != 2: # 2 == INCREMENTAL + _conn.execute("VACUUM;") + _conn.commit() + except Exception: + pass + + +init_db() + +# --------------------------------------------------------------------------- # +# Identifiers +# --------------------------------------------------------------------------- # + +# Unambiguous uppercase alphabet (no 0/O/1/I/L) -> ~4.9 bits/char. +CODE_ALPHABET = "23456789ABCDEFGHJKMNPQRSTVWXYZ" +MIN_CODE_LEN, MAX_CODE_LEN = 4, 12 +IDLE_TIMEOUT = int(os.environ.get("SOZAI_IDLE_TIMEOUT", str(60 * 60))) # 1 hour -> mark inactive +# Delete a room (and its photos/state) after this much inactivity. Keyed on +# last_activity, so the DB reflects live usage and doesn't bloat. Defaults to +# 1 hour + a 10-minute grace window; override with SOZAI_ROOM_DELETE_AFTER_IDLE. +ROOM_DELETE_AFTER_IDLE = int(os.environ.get("SOZAI_ROOM_DELETE_AFTER_IDLE", str(60 * 60 + 600))) +HARD_EXPIRATION = int(os.environ.get("SOZAI_HARD_EXPIRATION", str(24 * 60 * 60))) # absolute cap + + +def clamp_code_len(n) -> int: + try: + n = int(n) + except (TypeError, ValueError): + n = 6 + return max(MIN_CODE_LEN, min(MAX_CODE_LEN, n)) + + +def gen_share_code(length: int) -> str: + return "".join(secrets.choice(CODE_ALPHABET) for _ in range(length)) + + +def unique_share_code(length: int) -> str: + length = clamp_code_len(length) + with DB_LOCK: + for _ in range(40): + code = gen_share_code(length) + row = _conn.execute( + "SELECT 1 FROM rooms WHERE share_code=? AND is_active=1", (code,) + ).fetchone() + if row is None: + return code + # extremely unlikely fallback: widen the code + return gen_share_code(min(MAX_CODE_LEN, length + 2)) + + +def sanitize_name(name, default="Guest"): + name = (name or "").strip() + name = "".join(ch for ch in name if ch.isprintable()) + return name[:24] or default + + +def touch_room(room_id: str): + with DB_LOCK: + _conn.execute( + "UPDATE rooms SET last_activity=? WHERE room_id=?", (time.time(), room_id) + ) + _conn.commit() + + +def room_row(room_id: str): + with DB_LOCK: + return _conn.execute( + "SELECT * FROM rooms WHERE room_id=? AND is_active=1", (room_id,) + ).fetchone() + + +# --------------------------------------------------------------------------- # +# Gradio Server +# --------------------------------------------------------------------------- # + +app = Server() + + +# Kept from the previous version: a Gradio API endpoint the frontend can call +# through @gradio/client. Still works in collab mode. +@app.api(name="save_to_album") +def save_to_album(album: str = "", caption: str = "", tags: str = "") -> str: + """Save a photo + caption + tags to an album, returning a confirmation.""" + album = (album or "").strip() + if not album: + return "Select an album to save this photo." + return f'Added to "{album}".' + + +# --------------------------------------------------------------------------- # +# Auto-caption: a vision-language model that suggests a caption for a photo. +# The model (MiniCPM-V) is heavy, so it is loaded lazily on first use and +# cached. If transformers / llama.cpp / the model weights are unavailable, we +# fall back to a simple suggestion so the feature degrades gracefully instead +# of erroring. +# +# IMPORTANT: MiniCPM-V is a true VISION-LANGUAGE model. Both backends are wired +# to send the ACTUAL IMAGE to the model (not just alt-text), so the suggested +# title / caption / tags are grounded in what the photo actually shows. +# --------------------------------------------------------------------------- # + +AUTOCAPTION_MODEL_ID = os.environ.get("SOZAI_CAPTION_MODEL", "openbmb/MiniCPM-V-4") +# GGUF variant used when running inside a Hugging Face Space (where llama.cpp + +# a quantized GGUF is the practical way to run MiniCPM-V). Outside Spaces we use +# the normal transformers model above. Both are overridable via env. +# +# MiniCPM-V-4 is a vision-language model, so the GGUF build needs TWO files: the +# LM weights (CAPTION_GGUF_FILE) and the vision projector / CLIP encoder +# (CAPTION_GGUF_MMPROJ). Without the mmproj the model can't see images. Use the +# plain mmproj-model-f16.gguf — NOT the -IOS variant, which is iOS-only and +# fails to load in llama.cpp. +CAPTION_GGUF_REPO = os.environ.get("SOZAI_CAPTION_GGUF_REPO", "openbmb/MiniCPM-V-4-gguf") +CAPTION_GGUF_FILE = os.environ.get("SOZAI_CAPTION_GGUF_FILE", "ggml-model-Q4_K_M.gguf") +CAPTION_GGUF_MMPROJ = os.environ.get("SOZAI_CAPTION_GGUF_MMPROJ", "mmproj-model-f16.gguf") +_caption_lock = threading.Lock() +# backend: "" until loaded, then "gguf" or "transformers" +_caption_state = {"loaded": False, "ok": False, "model": None, "tokenizer": None, + "processor": None, "backend": "", + "gguf_model_path": None, "gguf_mmproj_path": None} + + +def _in_hf_space() -> bool: + """True when running inside a Hugging Face Space. HF sets SPACE_ID (and + related SPACE_* vars) automatically, so we use that to autoswitch to the + GGUF backend. Force on/off with SOZAI_USE_GGUF=1/0 if needed.""" + forced = os.environ.get("SOZAI_USE_GGUF", "").strip() + if forced in ("1", "true", "True", "yes"): + return True + if forced in ("0", "false", "False", "no"): + return False + return bool(os.environ.get("SPACE_ID") or os.environ.get("SPACE_HOST")) + + +def _load_caption_model_gguf() -> bool: + """Prepare the MiniCPM-V (vision) GGUF backend. Used inside HF Spaces. + + On ZeroGPU, CUDA must NOT be initialized in the main process — so this does + NOT build the llama.cpp model here. It only (a) checks the deps import, and + (b) downloads the LM + mmproj GGUF files (CPU/network only). The actual + `Llama(...)` construction happens lazily in `_ensure_gguf_model()`, which is + called from inside the @spaces.GPU worker where a GPU is attached. + """ + try: + from huggingface_hub import hf_hub_download + import llama_cpp # noqa: F401 (import check only) + from llama_cpp.llama_chat_format import MiniCPMv26ChatHandler # noqa: F401 + + model_path = hf_hub_download(repo_id=CAPTION_GGUF_REPO, filename=CAPTION_GGUF_FILE) + mmproj_path = hf_hub_download(repo_id=CAPTION_GGUF_REPO, filename=CAPTION_GGUF_MMPROJ) + # stash the resolved file paths; the model object is built later, lazily + _caption_state.update(model=None, tokenizer=None, processor=None, + gguf_model_path=model_path, gguf_mmproj_path=mmproj_path, + ok=True, backend="gguf") + print(f"[autocaption] GGUF vision backend ready " + f"({CAPTION_GGUF_REPO}: {CAPTION_GGUF_FILE} + {CAPTION_GGUF_MMPROJ}); " + f"model will load on first use inside the GPU worker") + return True + except Exception as exc: # llama_cpp/hub missing, download failed… + print(f"[autocaption] GGUF backend unavailable, will try transformers: {exc}") + return False + + +# Lock guarding the lazy in-worker GGUF build. On ZeroGPU each worker builds its +# own copy once; this prevents a double-build if two calls race. +_gguf_build_lock = threading.Lock() + + +def _ensure_gguf_model(): + """Build (once) and return the llama.cpp `Llama` model. MUST be called from + inside a @spaces.GPU context on ZeroGPU so the CUDA context is created while + a GPU is attached. Idempotent and thread-safe; caches into _caption_state. + Returns None on failure (caller falls back).""" + model = _caption_state.get("model") + if model is not None: + return model + with _gguf_build_lock: + model = _caption_state.get("model") + if model is not None: + return model + try: + from llama_cpp import Llama + from llama_cpp.llama_chat_format import MiniCPMv26ChatHandler + + model_path = _caption_state.get("gguf_model_path") + mmproj_path = _caption_state.get("gguf_mmproj_path") + if not model_path or not mmproj_path: + # files weren't resolved in the main-process prep step + from huggingface_hub import hf_hub_download + model_path = model_path or hf_hub_download( + repo_id=CAPTION_GGUF_REPO, filename=CAPTION_GGUF_FILE) + mmproj_path = mmproj_path or hf_hub_download( + repo_id=CAPTION_GGUF_REPO, filename=CAPTION_GGUF_MMPROJ) + chat_handler = MiniCPMv26ChatHandler(clip_model_path=mmproj_path, verbose=False) + llm = Llama( + model_path=model_path, + chat_handler=chat_handler, + n_ctx=4096, # vision tokens need a larger context window + n_gpu_layers=-1, # offload all layers (GPU is attached here) + verbose=False, + ) + _caption_state["model"] = llm + print(f"[autocaption] GGUF model built in worker ({CAPTION_GGUF_FILE})") + return llm + except Exception as exc: # noqa: BLE001 + print(f"[autocaption] GGUF in-worker build failed: {exc}") + return None + + +def _load_caption_model_transformers() -> bool: + """Load MiniCPM-V via transformers. Used outside HF Spaces. + + MiniCPM-V exposes a `.chat(image=..., msgs=...)` helper (trust_remote_code), + which we use so this path is ALSO image-grounded, not text-only. + """ + try: + import torch + from transformers import AutoModel, AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(AUTOCAPTION_MODEL_ID, trust_remote_code=True) + model = AutoModel.from_pretrained( + AUTOCAPTION_MODEL_ID, + trust_remote_code=True, + attn_implementation="sdpa", + torch_dtype=torch.bfloat16, + ) + model = model.eval() + if torch.cuda.is_available(): + model = model.cuda() + _caption_state.update(model=model, tokenizer=tokenizer, processor=None, + ok=True, backend="transformers") + print(f"[autocaption] transformers vision backend ready ({AUTOCAPTION_MODEL_ID})") + return True + except Exception as exc: # transformers/torch/model missing, no GPU, etc. + print(f"[autocaption] transformers backend unavailable: {exc}") + return False + + +def _load_caption_model(): + """Load and cache the captioning model once. Returns True on success. + + Autoswitches backend: inside a Hugging Face Space we use the GGUF + (llama.cpp) build; everywhere else we use the normal transformers model. + If the preferred backend fails to load we fall back to the other one before + giving up, so the feature still works wherever it can. + """ + if _caption_state["loaded"]: + return _caption_state["ok"] + with _caption_lock: + if _caption_state["loaded"]: + return _caption_state["ok"] + try: + if _in_hf_space(): + ok = _load_caption_model_gguf() or _load_caption_model_transformers() + else: + ok = _load_caption_model_transformers() or _load_caption_model_gguf() + _caption_state["ok"] = ok + finally: + _caption_state["loaded"] = True + return _caption_state["ok"] + + +def _resolve_image_for_model(image_path: str): + """Resolve whatever the frontend sends into a real local file path the + vision model can open. Handles three cases: + 1. a base64 data URL (data:image/...;base64,...) — what room uploads send + 2. a bundled asset URL (/assets/...) or zip entry — sample photos + 3. an already-local file path + Returns a filesystem path, or None if nothing usable was provided. + + This is the key to image-grounded captions: room photos arrive as data URLs, + NOT /assets/ paths, so without case (1) the model would never see the pixels + and would fall back to filename/alt-based guesses. + """ + if not image_path: + return None + s = image_path.strip() + # 1) data URL → write bytes to a temp file + if s.startswith("data:"): + try: + import base64 + import tempfile + header, b64 = (s.split(",", 1) if "," in s else ("", s)) + mime = "" + if header.startswith("data:") and ";" in header: + mime = header[len("data:"):header.index(";")] + ext = mimetypes.guess_extension(mime) or ".png" + raw = base64.b64decode(b64) + tmp = tempfile.NamedTemporaryFile(delete=False, suffix=ext) + tmp.write(raw) + tmp.close() + return tmp.name + except Exception as exc: # noqa: BLE001 + print(f"[autocaption] could not decode data URL: {exc}") + return None + # 2) asset / zip path + resolved = _resolve_asset_path(s) + if resolved: + return resolved + # 3) already a local file + if os.path.isfile(s): + return s + return None + + +def _resolve_asset_path(image_path: str): + """Map a frontend asset URL (/assets/...) to a real file or zip bytes.""" + if not image_path: + return None + rel = image_path.split("?")[0] + if rel.startswith("/assets/"): + rel = rel[len("/assets/"):] + rel = _safe_rel(rel) + if PUBLIC_DIR is not None: + fp = os.path.join(PUBLIC_DIR, *rel.split("/")) + if os.path.isfile(fp): + return fp + if ZIP_PATH is not None: + try: + import tempfile + + data = _read_from_zip(rel) + suffix = os.path.splitext(rel)[1] or ".png" + tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix) + tmp.write(data) + tmp.close() + return tmp.name + except KeyError: + return None + return None + + +def _load_pil_image(local_img: str): + """Open a resolved asset path as an RGB PIL image, or None on failure.""" + if not local_img or not os.path.isfile(local_img): + return None + try: + from PIL import Image + return Image.open(local_img).convert("RGB") + except Exception as exc: # noqa: BLE001 + print(f"[autocaption] could not open image {local_img}: {exc}") + return None + + +def _image_data_uri(local_img: str): + """Encode a local image file as a data: URI for the GGUF chat handler.""" + if not local_img or not os.path.isfile(local_img): + return None + try: + import base64 + mime = mimetypes.guess_type(local_img)[0] or "image/png" + with open(local_img, "rb") as fh: + b64 = base64.b64encode(fh.read()).decode() + return f"data:{mime};base64,{b64}" + except Exception as exc: # noqa: BLE001 + print(f"[autocaption] could not encode image {local_img}: {exc}") + return None + + +def _fallback_caption(alt: str) -> str: + base = (alt or "A moment worth remembering").strip().rstrip(".") + return base[:1].upper() + base[1:] + "." + + +def _fallback_tags(alt: str) -> str: + """Comma-separated tag fallback derived from the alt text.""" + words = re.findall(r"[A-Za-z0-9]+", (alt or "").lower()) + stop = {"the", "a", "an", "of", "and", "in", "on", "at", "to", "with", "photo", "image"} + tags = [w for w in words if w not in stop and len(w) > 2][:6] + if not tags: + tags = ["memory", "moment"] + return ", ".join(dict.fromkeys(tags)) # de-dupe, keep order + + +def _normalize_tags(text: str) -> str: + """Coerce model output into clean comma-separated tags. Splits on commas, + newlines, semicolons or bullet markers; trims; de-dupes; lowercases. Returns + "" if nothing usable is found (caller treats that as an error).""" + if not text: + return "" + # split on common separators the model might emit + parts = re.split(r"[,\n;•·\-\u2022]+", text) + out = [] + seen = set() + for p in parts: + tag = re.sub(r"^[\s\d\.\)\(]+", "", p).strip().strip('"').strip("'").lower() + tag = re.sub(r"\s+", " ", tag) + if tag and tag not in seen and len(tag) <= 40: + seen.add(tag) + out.append(tag) + return ", ".join(out) + + +def _is_comma_separated(text: str) -> bool: + """A valid tag string has at least one tag and, if multiple, comma sep.""" + t = (text or "").strip() + if not t: + return False + # single tag (no spaces-as-separator confusion) or multiple comma-separated + if "," in t: + return all(seg.strip() for seg in t.split(",")) + return bool(t) + + +@app.api(name="caption_model") +def caption_model() -> str: + """Return the id of the vision-language model used for auto-captioning, + so the UI can show the user which model is currently in use.""" + return AUTOCAPTION_MODEL_ID + + +# --------------------------------------------------------------------------- # +# Observability with Arize Phoenix + OpenTelemetry (no API keys). +# +# Phoenix runs IN-PROCESS — it launches automatically when app.py starts, so +# there is no separate `phoenix serve` to run. Every model call (caption / NSFW +# / ASR) is emitted as an OpenTelemetry span with OpenInference semantic +# attributes and shows up in the Phoenix UI, which is embedded directly in the +# app's "LLM Trace" section (via iframe). We also keep a small in-memory ring +# buffer so the in-app list works even if Phoenix can't start. +# +# Port can be overridden with PHOENIX_PORT (default 6006). +# --------------------------------------------------------------------------- # +_trace_lock = threading.Lock() +_trace_buffer: list = [] # most-recent-last ring buffer of spans +_TRACE_MAX = 200 +_otel_tracer = None +_otel_loaded = False +_otel_provider = None +_phoenix_session = None +_phoenix_launched = False +_phoenix_lock = threading.Lock() +PHOENIX_PORT = int(os.environ.get("PHOENIX_PORT", "6006")) +PHOENIX_HOST = os.environ.get("PHOENIX_HOST", "localhost") +PHOENIX_PROJECT = os.environ.get("PHOENIX_PROJECT_NAME", "sozai") +# the collector endpoint OTel sends spans to (same in-process server) +PHOENIX_ENDPOINT = os.environ.get( + "PHOENIX_COLLECTOR_ENDPOINT", f"http://{PHOENIX_HOST}:{PHOENIX_PORT}" +) + + +def _launch_phoenix(): + """Start the Phoenix app in-process (once). Returns the UI base URL, or + None if Phoenix isn't installed / failed to start. Idempotent + thread-safe. + This means the trace UI is available the moment app.py is running — no + separate `phoenix serve` needed. + + Host/port are configured via the PHOENIX_HOST / PHOENIX_PORT environment + variables (the launch_app host/port kwargs are deprecated).""" + global _phoenix_session, _phoenix_launched + if _phoenix_launched: + return _phoenix_session_url() + with _phoenix_lock: + if _phoenix_launched: + return _phoenix_session_url() + _phoenix_launched = True + # Phoenix reads these at import/launch time — set them before importing. + os.environ.setdefault("PHOENIX_HOST", PHOENIX_HOST) + os.environ.setdefault("PHOENIX_PORT", str(PHOENIX_PORT)) + try: + import phoenix as px + # launch_app() picks up PHOENIX_HOST / PHOENIX_PORT from the env. + # If something is already serving on the port, it returns a session + # pointed at it. + _phoenix_session = px.launch_app() + print(f"[phoenix] in-process UI at {_phoenix_session_url()}") + except Exception as exc: # noqa: BLE001 + print(f"[phoenix] could not launch in-process app: {exc}") + _phoenix_session = None + return _phoenix_session_url() + + +def _phoenix_session_url(): + if _phoenix_session is not None: + try: + return _phoenix_session.url + except Exception: # noqa: BLE001 + pass + return f"http://{PHOENIX_HOST}:{PHOENIX_PORT}" + + +def _get_tracer(): + """Lazily launch in-process Phoenix, register the OTel tracer provider, and + return a tracer. Falls back to a raw OTLP/HTTP exporter, then to the local + buffer only. Returns None if OpenTelemetry isn't installed.""" + global _otel_tracer, _otel_loaded + if _otel_loaded: + return _otel_tracer + _otel_loaded = True + # make sure the embedded Phoenix collector is up first + _launch_phoenix() + # Preferred: phoenix.otel.register — Phoenix-aware OTel defaults, no keys. + try: + from phoenix.otel import register + tracer_provider = register( + project_name=PHOENIX_PROJECT, + endpoint=PHOENIX_ENDPOINT.rstrip("/") + "/v1/traces", + batch=True, # use a BatchSpanProcessor (recommended) + set_global_tracer_provider=True, + ) + # batched spans may still be queued at exit — flush them on shutdown. + global _otel_provider + _otel_provider = tracer_provider + try: + import atexit + atexit.register(lambda: _otel_provider and _otel_provider.shutdown()) + except Exception: # noqa: BLE001 + pass + _otel_tracer = tracer_provider.get_tracer(__name__) + print(f"[otel] Phoenix tracing enabled → {PHOENIX_ENDPOINT} (project '{PHOENIX_PROJECT}')") + return _otel_tracer + except Exception as exc: # noqa: BLE001 + print(f"[otel] phoenix.otel unavailable ({exc}); trying raw OTLP exporter") + # Fallback: plain OpenTelemetry with an OTLP/HTTP exporter to Phoenix. + try: + from opentelemetry import trace as _t + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import BatchSpanProcessor + from opentelemetry.sdk.resources import Resource + from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter + provider = TracerProvider(resource=Resource.create({"service.name": PHOENIX_PROJECT})) + provider.add_span_processor(BatchSpanProcessor( + OTLPSpanExporter(endpoint=PHOENIX_ENDPOINT.rstrip("/") + "/v1/traces"))) + _t.set_tracer_provider(provider) + _otel_tracer = _t.get_tracer(__name__) + print(f"[otel] OTLP tracing enabled → {PHOENIX_ENDPOINT}") + return _otel_tracer + except Exception as exc: # noqa: BLE001 + print(f"[otel] OpenTelemetry unavailable, using local trace buffer only: {exc}") + _otel_tracer = None + return None + + +def _record_trace(event: dict): + """Append a trace event to the local ring buffer (newest last).""" + event = dict(event) + event.setdefault("ts", time.time()) + with _trace_lock: + _trace_buffer.append(event) + if len(_trace_buffer) > _TRACE_MAX: + del _trace_buffer[: len(_trace_buffer) - _TRACE_MAX] + + +# OpenInference semantic attribute keys (string fallbacks if the package that +# defines them isn't importable — keeps tracing working without it). These mirror +# the attributes Phoenix reads for its span detail + metrics. +_OI = { + "kind": "openinference.span.kind", + "input": "input.value", + "input_mime": "input.mime_type", + "output": "output.value", + "output_mime": "output.mime_type", + "model": "llm.model_name", + "provider": "llm.provider", + "system": "llm.system", + "invocation": "llm.invocation_parameters", + "finish": "llm.finish_reason", + "tok_prompt": "llm.token_count.prompt", + "tok_completion": "llm.token_count.completion", + "tok_total": "llm.token_count.total", + "metadata": "metadata", + "session": "session.id", +} +try: + from openinference.semconv.trace import SpanAttributes as _OISpan + _OI.update({ + "kind": _OISpan.OPENINFERENCE_SPAN_KIND, + "input": _OISpan.INPUT_VALUE, + "output": _OISpan.OUTPUT_VALUE, + "model": _OISpan.LLM_MODEL_NAME, + }) + for k, attr in (("input_mime", "INPUT_MIME_TYPE"), ("output_mime", "OUTPUT_MIME_TYPE"), + ("provider", "LLM_PROVIDER"), ("system", "LLM_SYSTEM"), + ("invocation", "LLM_INVOCATION_PARAMETERS"), ("finish", "LLM_FINISH_REASON"), + ("tok_prompt", "LLM_TOKEN_COUNT_PROMPT"), ("tok_completion", "LLM_TOKEN_COUNT_COMPLETION"), + ("tok_total", "LLM_TOKEN_COUNT_TOTAL"), ("metadata", "METADATA"), + ("session", "SESSION_ID")): + if hasattr(_OISpan, attr): + _OI[k] = getattr(_OISpan, attr) +except Exception: # noqa: BLE001 + pass +# back-compat aliases used elsewhere +_OI_KIND = _OI["kind"]; _OI_INPUT = _OI["input"]; _OI_OUTPUT = _OI["output"]; _OI_MODEL = _OI["model"] + + +def _provider_for(model_id: str) -> str: + """Best-effort provider label from a model id (for the trace 'provider').""" + m = (model_id or "").lower() + if m.startswith("nvidia/") or "nemotron" in m or "nemo" in m: + return "nvidia" + if "minicpm" in m or m.startswith("openbmb/"): + return "openbmb" + if "falconsai" in m or "nsfw" in m: + return "falconsai" + if "/" in m: + return m.split("/", 1)[0] + return "local" + + +def _count_tokens(text: str) -> int: + """Real token count using the caption model's tokenizer when available; + otherwise a clearly-approximate word-based estimate.""" + text = text or "" + if not text: + return 0 + tok = _caption_state.get("tokenizer") if isinstance(_caption_state, dict) else None + if tok is not None: + try: + return len(tok.encode(text)) + except Exception: # noqa: BLE001 + try: + return len(tok(text)["input_ids"]) + except Exception: # noqa: BLE001 + pass + return int(round(len(text.split()) / 0.75)) + + + +class traced: + """Context manager that times a model call, records it to the local buffer, + and emits an OpenTelemetry span (OpenInference 'LLM' kind) to Phoenix. + + Usage: + with traced("autocaption", model=MODEL_ID, file="x.jpg", + meta={"prompt": p}) as t: + ... + t.set_output(text) + """ + + def __init__(self, name, model="", file="", meta=None, invocation=None, system=""): + self.name = name + self.model = model + self.file = file + self.meta = meta or {} + self.invocation = invocation or {} # llm.invocation_parameters + self.system = system # llm.system (system prompt, if any) + self.output = None + self.finish_reason = "stop" + self.error = None + self._t0 = None + self._span_cm = None + self._span = None + self.trace_id = "" + self.span_id = "" + + def __enter__(self): + self._t0 = time.time() + tracer = _get_tracer() + if tracer is not None: + try: + self._span_cm = tracer.start_as_current_span(f"sozai.{self.name}") + self._span = self._span_cm.__enter__() + self._span.set_attribute(_OI["kind"], "LLM") + if self.model: + self._span.set_attribute(_OI["model"], self.model) + self._span.set_attribute(_OI["provider"], _provider_for(self.model)) + if self.system: + self._span.set_attribute(_OI["system"], str(self.system)[:1000]) + if self.invocation: + try: + self._span.set_attribute(_OI["invocation"], json.dumps(self.invocation)[:1000]) + except Exception: # noqa: BLE001 + pass + if self.file: + self._span.set_attribute("sozai.file", self.file) + # input value + mime for Phoenix display + try: + self._span.set_attribute(_OI["input"], json.dumps({"file": self.file, **self.meta})[:2000]) + self._span.set_attribute(_OI["input_mime"], "application/json") + except Exception: # noqa: BLE001 + pass + for k, v in self.meta.items(): + try: + self._span.set_attribute(f"sozai.{k}", str(v)[:500]) + except Exception: # noqa: BLE001 + pass + # capture the OTel trace/span ids so the in-app view can show them + try: + ctx = self._span.get_span_context() + self.trace_id = format(ctx.trace_id, "032x") + self.span_id = format(ctx.span_id, "016x") + except Exception: # noqa: BLE001 + pass + except Exception as exc: # noqa: BLE001 + print(f"[otel] span start failed: {exc}") + self._span = None + if not self.span_id: + self.span_id = uuid.uuid4().hex[:16] + if not self.trace_id: + self.trace_id = uuid.uuid4().hex + return self + + def set_output(self, output, finish_reason="stop"): + self.output = output + self.finish_reason = finish_reason + + def __exit__(self, exc_type, exc, tb): + t_end = time.time() + dur_ms = round((t_end - self._t0) * 1000.0, 1) if self._t0 else None + if exc is not None: + self.error = f"{exc_type.__name__}: {exc}" if exc_type else str(exc) + self.finish_reason = "error" + try: + input_payload = json.dumps({"file": self.file, **self.meta}) + except Exception: # noqa: BLE001 + input_payload = str({"file": self.file, **self.meta}) + out_str = str(self.output) if self.output is not None else "" + # real token counts via the model tokenizer when available, else estimate + tok_in = _count_tokens(input_payload) + tok_out = _count_tokens(out_str) + tok_total = tok_in + tok_out + tokens_per_sec = round(tok_out / (dur_ms / 1000.0), 1) if (dur_ms and tok_out) else 0 + provider = _provider_for(self.model) + _record_trace({ + "id": self.span_id, + "span_id": self.span_id, + "trace_id": self.trace_id, + "name": self.name, + "kind": "LLM", + "model": self.model, + "provider": provider, + "system": (str(self.system)[:400] if self.system else None), + "invocation": (self.invocation or None), + "finish_reason": self.finish_reason, + "file": self.file, + "meta": self.meta, + "input": input_payload[:2000], + "input_mime": "application/json", + "output": (out_str[:2000] if out_str else None), + "output_mime": "text/plain", + "error": self.error, + "duration_ms": dur_ms, + "started": self._t0, + "ended": t_end, + "tokens_in": tok_in, + "tokens_out": tok_out, + "tokens_total": tok_total, + "tokens_per_sec": tokens_per_sec, + "cost_usd": 0.0, # local open-source models — no API cost + "status": "error" if self.error else "ok", + }) + if self._span is not None: + try: + if self.output is not None: + self._span.set_attribute(_OI["output"], str(self.output)[:2000]) + self._span.set_attribute(_OI["output_mime"], "text/plain") + self._span.set_attribute(_OI["finish"], self.finish_reason) + if dur_ms is not None: + self._span.set_attribute("sozai.duration_ms", dur_ms) + self._span.set_attribute(_OI["tok_prompt"], tok_in) + self._span.set_attribute(_OI["tok_completion"], tok_out) + self._span.set_attribute(_OI["tok_total"], tok_total) + if self.error: + try: + from opentelemetry.trace import Status, StatusCode + self._span.set_status(Status(StatusCode.ERROR, self.error)) + except Exception: # noqa: BLE001 + pass + self._span.set_attribute("sozai.error", self.error) + except Exception: # noqa: BLE001 + pass + try: + self._span_cm.__exit__(exc_type, exc, tb) + except Exception: # noqa: BLE001 + pass + return False # never suppress exceptions + + +@app.get("/api/trace") +def get_trace(limit: int = 50): + """Return recent model-call traces plus summary stats for the in-app trace + view. Distributed traces also go to the in-process Phoenix UI.""" + tracer_ok = _get_tracer() is not None + with _trace_lock: + all_items = list(_trace_buffer) + items = _trace_buffer[-max(1, min(limit, _TRACE_MAX)):] + + # aggregate stats across everything in the buffer + total = len(all_items) + errors = sum(1 for e in all_items if e.get("status") == "error") + durations = [e["duration_ms"] for e in all_items if isinstance(e.get("duration_ms"), (int, float)) and e.get("duration_ms")] + avg_ms = round(sum(durations) / len(durations), 1) if durations else 0 + + def _pct(sorted_vals, q): + if not sorted_vals: + return 0 + return round(sorted_vals[min(len(sorted_vals) - 1, int(len(sorted_vals) * q))], 1) + s = sorted(durations) + p50_ms, p95_ms, p99_ms = _pct(s, 0.50), _pct(s, 0.95), _pct(s, 0.99) + + tokens_prompt = sum(int(e.get("tokens_in") or 0) for e in all_items) + tokens_completion = sum(int(e.get("tokens_out") or 0) for e in all_items) + tokens_total = tokens_prompt + tokens_completion + tps_vals = [e["tokens_per_sec"] for e in all_items if isinstance(e.get("tokens_per_sec"), (int, float)) and e.get("tokens_per_sec")] + avg_tps = round(sum(tps_vals) / len(tps_vals), 1) if tps_vals else 0 + cost_total = round(sum(float(e.get("cost_usd") or 0) for e in all_items), 4) + + by_type: dict = {} + by_kind: dict = {} + for e in all_items: + nm = e.get("name", "?") + b = by_type.setdefault(nm, {"count": 0, "errors": 0, "dur": [], "tok": 0}) + b["count"] += 1 + if e.get("status") == "error": + b["errors"] += 1 + if isinstance(e.get("duration_ms"), (int, float)) and e.get("duration_ms"): + b["dur"].append(e["duration_ms"]) + b["tok"] += int(e.get("tokens_total") or 0) + k = e.get("kind", "LLM") + by_kind[k] = by_kind.get(k, 0) + 1 + by_type_out = { + nm: { + "count": b["count"], + "errors": b["errors"], + "avg_ms": round(sum(b["dur"]) / len(b["dur"]), 1) if b["dur"] else 0, + "tokens": b["tok"], + } + for nm, b in by_type.items() + } + + return { + "phoenix": tracer_ok, + "phoenix_url": "/phoenix/" if _phoenix_session is not None else None, + "phoenix_direct": _phoenix_session_url() if _phoenix_session is not None else None, + "endpoint": PHOENIX_ENDPOINT if tracer_ok else None, + "project": PHOENIX_PROJECT, + "stats": { + "total": total, + "errors": errors, + "error_rate": round(errors / total, 3) if total else 0, + "avg_ms": avg_ms, + "p50_ms": p50_ms, + "p95_ms": p95_ms, + "p99_ms": p99_ms, + "tokens_total": tokens_total, + "tokens_prompt": tokens_prompt, + "tokens_completion": tokens_completion, + "avg_tokens_per_sec": avg_tps, + "cost_usd": cost_total, + "by_type": by_type_out, + "by_kind": by_kind, + }, + "events": list(reversed(items)), # newest first for display + } + + +# --- Same-origin reverse proxy for the embedded Phoenix UI ------------------ # +# Phoenix runs in-process on its own port (6006). To embed it in an iframe +# without cross-origin / X-Frame-Options problems, we proxy it under /phoenix/ +# on the app's own origin. Frame-blocking headers are stripped on the way back. +import urllib.request as _urlreq +import urllib.error as _urlerr + +_HOP_BY_HOP = { + "content-encoding", "content-length", "transfer-encoding", "connection", + "keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailers", + "upgrade", "x-frame-options", "content-security-policy", +} + + +_PHX_PREFIX = "/phoenix" + + +def _rewrite_phoenix_body(data: bytes, ctype: str) -> bytes: + """Phoenix's SPA references assets/APIs with root-absolute paths + (/assets/..., /graphql, /v1/...). Served under /phoenix/, those would + resolve against the app origin and 404 — leaving the iframe blank. Rewrite + HTML and JS so everything stays under the proxy prefix.""" + is_html = "text/html" in ctype + is_js = ("javascript" in ctype) or ("text/css" in ctype) + if not (is_html or is_js): + return data + try: + text = data.decode("utf-8") + except Exception: # noqa: BLE001 + return data # binary; leave untouched + + if is_html: + # 1) a so relative URLs resolve under /phoenix/ + if "]*>)", r'\1', + text, count=1, flags=re.IGNORECASE) + # 2) prefix root-absolute attribute URLs: href="/x" src="/x" + # (skip protocol-relative //, data:, and already-prefixed /phoenix) + text = re.sub(r'(\b(?:href|src)=")/(?!/|phoenix/)', + r"\1" + _PHX_PREFIX + "/", text) + + # 3) prefix root-absolute paths used in JS/CSS (fetch, websocket, url()). + # Phoenix hits "/graphql", "/v1/...", "/arize_phoenix_version", etc. + for p in ("/graphql", "/v1/", "/arize_phoenix_version", "/exports", + "/assets/", "/healthz", "/readyz"): + # in quotes: "/graphql" '/v1/...' + text = text.replace('"' + p, '"' + _PHX_PREFIX + p) + text = text.replace("'" + p, "'" + _PHX_PREFIX + p) + # css url(/assets/...) + text = text.replace("(" + p, "(" + _PHX_PREFIX + p) + return text.encode("utf-8") + + +async def _phoenix_proxy(request: Request, path: str): + base = _phoenix_session_url().rstrip("/") + target = base + "/" + path + if request.url.query: + target += "?" + request.url.query + body = await request.body() + method = request.method + headers = {k: v for k, v in request.headers.items() + if k.lower() not in ("host", "content-length", "accept-encoding")} + headers["accept-encoding"] = "identity" + req = _urlreq.Request(target, data=(body or None), method=method, headers=headers) + try: + with _urlreq.urlopen(req, timeout=30) as resp: + data = resp.read() + status = resp.status + ctype = resp.headers.get("Content-Type", "application/octet-stream") + out_headers = {} + for k, v in resp.headers.items(): + if k.lower() in _HOP_BY_HOP: + continue + out_headers[k] = v + data = _rewrite_phoenix_body(data, ctype) + return Response(content=data, status_code=status, media_type=ctype, headers=out_headers) + except _urlerr.HTTPError as e: + try: + err_body = e.read() + except Exception: # noqa: BLE001 + err_body = b"" + return Response(content=err_body, status_code=e.code, + media_type=e.headers.get("Content-Type", "text/plain")) + except Exception as exc: # noqa: BLE001 + return Response(content=f"Phoenix UI not reachable: {exc}".encode(), + status_code=502, media_type="text/plain") + + +@app.get("/phoenix") +@app.get("/phoenix/") +async def phoenix_root(request: Request): + return await _phoenix_proxy(request, "") + + +@app.get("/phoenix/{path:path}") +async def phoenix_proxy_get(request: Request, path: str): + return await _phoenix_proxy(request, path) + + +@app.post("/phoenix/{path:path}") +async def phoenix_proxy_post(request: Request, path: str): + return await _phoenix_proxy(request, path) + + +# Cesium ion access token. NEVER hardcode this in the HTML — it is read from +# the environment so it can be rotated without touching the frontend. Set it +# with: export SOZAI_CESIUM_ION_TOKEN="your-token" (or put it in a .env file). +from dotenv import load_dotenv +load_dotenv() +CESIUM_ION_TOKEN = os.environ.get("SOZAI_CESIUM_ION_TOKEN", "") + + +@app.get("/api/config") +def client_config(): + """Public, non-secret-ish runtime config for the frontend (e.g. whether a + Cesium ion token is configured, and the token itself if so). The token is + only as protected as the page that needs it; keeping it server-side env-only + means it can be rotated in one place.""" + import platform + return { + "cesiumIonToken": CESIUM_ION_TOKEN, + "cesiumEnabled": bool(CESIUM_ION_TOKEN), + "nsfwEnabled": True, # the gate always runs (open if model absent) + "asrEnabled": True, # the endpoint always exists (may be unavailable) + "asrModel": ASR_MODEL_ID, + "asrOsSupported": platform.system().lower() == "linux", + "asrLanguages": ASR_LANGUAGE_CHOICES, + "asrProfiles": list(ASR_CHUNK_PROFILES.keys()), + "asrDefaultProfile": ASR_DEFAULT_PROFILE, + "captionModel": AUTOCAPTION_MODEL_ID, + # "Proceed to Development" watercolour: the endpoint always exists; this + # only reports whether torch/diffusers are present (a CUDA-free check — + # the LoRA is optional, stage 1 runs base img2img). The ~14 GB pipeline + # still builds lazily on first use inside the GPU worker. + "img2imgEnabled": _img2img_prep(), + "img2imgModel": IMG2IMG_MODEL_ID, + "phoenix": _get_tracer() is not None, + } + + +# --------------------------------------------------------------------------- # +# NSFW image gate (item 1). Falconsai/nsfw_image_detection via transformers. +# If transformers/torch/model are missing, the gate is OPEN (score 0.0) so +# uploads are never blocked by a missing optional dependency. +# --------------------------------------------------------------------------- # +NSFW_THRESHOLD = float(os.environ.get("SOZAI_NSFW_THRESHOLD", "0.80")) +NSFW_MODEL_ID = "Falconsai/nsfw_image_detection" +_nsfw_lock = threading.Lock() +_nsfw_state = {"loaded": False, "pipe": None, "available": False} + + +def _get_nsfw_pipe(): + """Check that the NSFW classifier CAN be built, without initializing CUDA in + the main process (forbidden on ZeroGPU). This only verifies the import + + that transformers is present; the actual pipeline is constructed lazily + inside the @spaces.GPU worker by _ensure_nsfw_pipe(). Returns True when the + gate is available, None/False when it can't be (gate then fails OPEN).""" + if _nsfw_state["loaded"]: + return _nsfw_state["available"] + with _nsfw_lock: + if _nsfw_state["loaded"]: + return _nsfw_state["available"] + try: + import transformers # noqa: F401 (import check only — no CUDA here) + _nsfw_state["available"] = True + except Exception as exc: # noqa: BLE001 + print(f"[nsfw] transformers unavailable, gate open: {exc}") + _nsfw_state["available"] = False + finally: + _nsfw_state["loaded"] = True + return _nsfw_state["available"] + + +_nsfw_build_lock = threading.Lock() + + +def _ensure_nsfw_pipe(): + """Build (once) and return the transformers image-classification pipeline. + MUST run inside a @spaces.GPU context on ZeroGPU so CUDA initializes with a + GPU attached. Idempotent + thread-safe; caches into _nsfw_state. Returns the + pipeline, or None on failure (caller then treats the image as safe).""" + pipe = _nsfw_state.get("pipe") + if pipe is not None: + return pipe + with _nsfw_build_lock: + pipe = _nsfw_state.get("pipe") + if pipe is not None: + return pipe + try: + from transformers import pipeline + import torch + device = 0 if torch.cuda.is_available() else -1 + pipe = pipeline("image-classification", model=NSFW_MODEL_ID, device=device) + _nsfw_state["pipe"] = pipe + print(f"[nsfw] classifier built in worker (device={'cuda' if device == 0 else 'cpu'})") + return pipe + except Exception as exc: # noqa: BLE001 + print(f"[nsfw] classifier build failed, gate open: {exc}") + return None + + +@spaces.GPU(duration=30) +def _nsfw_infer(img): + """Score one PIL image for NSFW. Builds/uses the pipeline INSIDE this + @spaces.GPU context so CUDA is only touched with a GPU attached. The image + (picklable) crosses the boundary, never the model.""" + pipe = _ensure_nsfw_pipe() + if pipe is None: + return [] + return pipe(img.convert("RGB")) + + +def _nsfw_score_for(path_or_pil) -> float: + if not _get_nsfw_pipe(): + return 0.0 + try: + from PIL import Image + img = path_or_pil if hasattr(path_or_pil, "convert") else Image.open(path_or_pil) + for r in _nsfw_infer(img): + if str(r["label"]).lower() == "nsfw": + return float(r["score"]) + except Exception as exc: # noqa: BLE001 + print(f"[nsfw] scoring failed (treating as safe): {exc}") + return 0.0 + + +@app.api(name="nsfw_check") +def nsfw_check(image_path: str = "") -> dict: + """Score an uploaded image for NSFW content. Returns {score, blocked, + threshold, available}. Fails OPEN (never blocks) if the model is absent.""" + fp = _resolve_image_for_model(image_path) or image_path + available = bool(_get_nsfw_pipe()) + with traced("nsfw_check", model=NSFW_MODEL_ID, file=os.path.basename(image_path or "")) as t: + score = _nsfw_score_for(fp) if available else 0.0 + t.set_output({"score": round(score, 4)}) + return { + "score": round(score, 4), + "blocked": bool(score >= NSFW_THRESHOLD), + "threshold": NSFW_THRESHOLD, + "available": available, + } + + +@app.post("/api/nsfw") +async def nsfw_rest(request: Request): + """NSFW gate for client uploads sent as a base64 data URL. + Body: {data_url: "data:image/...;base64,...", name: "file.jpg"}. + Fails OPEN on any error so a missing model never blocks uploads.""" + import base64 + try: + body = await request.json() + except Exception: # noqa: BLE001 + return {"score": 0.0, "blocked": False, "available": bool(_get_nsfw_pipe()), "error": "bad request"} + data_url = (body or {}).get("data_url", "") + name = (body or {}).get("name", "upload") + available = bool(_get_nsfw_pipe()) + if not available or not data_url: + return {"score": 0.0, "blocked": False, "threshold": NSFW_THRESHOLD, "available": available} + try: + b64 = data_url.split(",", 1)[1] if "," in data_url else data_url + raw = base64.b64decode(b64) + from PIL import Image + img = Image.open(io.BytesIO(raw)) + with traced("nsfw_check", model=NSFW_MODEL_ID, file=str(name)[:60]) as t: + score = _nsfw_score_for(img) + t.set_output({"score": round(score, 4)}) + return {"score": round(score, 4), "blocked": bool(score >= NSFW_THRESHOLD), + "threshold": NSFW_THRESHOLD, "available": True} + except Exception as exc: # noqa: BLE001 + print(f"[nsfw] rest scoring failed (treating as safe): {exc}") + return {"score": 0.0, "blocked": False, "threshold": NSFW_THRESHOLD, "available": available, "error": str(exc)} + + +# --------------------------------------------------------------------------- # +# Develop → watercolour via Flux.2-klein + a scene LoRA (in-process img2img). +# This is what "Proceed to Development" triggers: each photo is sent here and +# painted into a watercolour by Flux2KleinPipeline (the source is VAE-encoded +# and used as reference conditioning while the LoRA steers the watercolour +# style). +# +# ZeroGPU port notes (mirrors the caption/NSFW/ASR helpers in this file): +# * The pipeline is ~14 GB and lives ONLY on the GPU, so it must be built +# INSIDE a @spaces.GPU context — never in the main process, where CUDA must +# stay un-initialized. `_img2img_prep()` does the CUDA-free part (import + +# LoRA file check) so /config and the endpoint can report availability +# cheaply; `_ensure_img2img_pipe()` does the actual GPU build, lazily, the +# first time `_img2img_infer` runs in a worker. +# * @spaces.GPU pickles arguments to/from the worker. The Flux pipe is NOT +# picklable, so it is fetched from module state inside the worker and never +# crosses the boundary. Only the PIL image in and the PIL image out cross +# it (both picklable), exactly like `_nsfw_infer`. +# * This file is self-contained on the Space (no backend/ package), so the +# squaring + load + stylize logic from backend/hf/3_infer_img2img.py is +# inlined here rather than imported. +# +# STAGED ROLLOUT: the scene LoRA is OPTIONAL. Stage 1 (now) is "pure img2img" — +# the base Flux2Klein model develops the photo with a generic watercolour prompt, +# no .safetensors required. Stage 2 is dropping the LoRA file in (or setting +# SOZAI_IMG2IMG_LORA_REPO): it's detected automatically and the trained +# watercolour style + trigger prompt switch on. So a missing LoRA degrades to +# base img2img rather than disabling the feature. +# +# Fails CLOSED (available=false) only when torch/diffusers are absent, so a CPU +# Space or local Mac keeps the original photo instead of erroring/OOMing. +# --------------------------------------------------------------------------- # +IMG2IMG_MODEL_ID = os.environ.get("SOZAI_IMG2IMG_MODEL", "black-forest-labs/FLUX.2-klein-4B") +# Path to the watercolour scene LoRA .safetensors. OPTIONAL: ship the file in the +# repo at this path, or point SOZAI_IMG2IMG_LORA_REPO/_FILE at a Hub repo to have +# it downloaded (CPU/network only) at prep time. Absent = base img2img (stage 1). +IMG2IMG_LORA_PATH = os.environ.get( + "SOZAI_IMG2IMG_LORA", + os.path.join(BASE, "data/places/checkpoints/hf-spaces/scene_lora.safetensors"), +) +IMG2IMG_LORA_REPO = os.environ.get("SOZAI_IMG2IMG_LORA_REPO", "") +IMG2IMG_LORA_FILE = os.environ.get("SOZAI_IMG2IMG_LORA_FILE", "scene_lora.safetensors") +IMG2IMG_LORA_SCALE = float(os.environ.get("SOZAI_IMG2IMG_LORA_SCALE", "1.0")) +# Generation defaults (from backend/hf/3_infer_img2img.py). +IMG2IMG_TRIGGER = "wtrclr8" +# Prompt used WITH the LoRA (its trigger token), vs. the base-model fallback used +# in stage 1 when no LoRA is loaded — the trigger token means nothing without it. +IMG2IMG_PROMPT = os.environ.get("SOZAI_IMG2IMG_PROMPT", "WTRCLR8 watercolour painting") +IMG2IMG_BASE_PROMPT = os.environ.get( + "SOZAI_IMG2IMG_BASE_PROMPT", + "a soft watercolour painting of this scene, painterly, textured paper, gentle washes", +) +IMG2IMG_SIZE = int(os.environ.get("SOZAI_IMG2IMG_SIZE", "1024")) +IMG2IMG_STEPS = int(os.environ.get("SOZAI_IMG2IMG_STEPS", "8")) +IMG2IMG_GUIDANCE = float(os.environ.get("SOZAI_IMG2IMG_GUIDANCE", "3.5")) +IMG2IMG_SEED = int(os.environ.get("SOZAI_IMG2IMG_SEED", "42")) +_img2img_lock = threading.Lock() +_img2img_state = {"loaded": False, "available": False, "pipe": None, + "lora_path": None, "lora_loaded": False} + + +def _img2img_prep() -> bool: + """CUDA-free availability check, run in the main process. Verifies diffusers + (with Flux2KleinPipeline) imports; the LoRA is OPTIONAL — if a file is present + (or downloadable via SOZAI_IMG2IMG_LORA_REPO) its path is cached for stage 2, + otherwise we fall back to base img2img. Does NOT touch CUDA or build the + pipeline. Returns True whenever the base feature can run in a GPU worker.""" + if _img2img_state["loaded"]: + return _img2img_state["available"] + with _img2img_lock: + if _img2img_state["loaded"]: + return _img2img_state["available"] + try: + import torch # noqa: F401 (import check only — no CUDA here) + from diffusers import Flux2KleinPipeline # noqa: F401 + # Resolve the LoRA if we can, but its absence is NOT fatal. + lora_path = IMG2IMG_LORA_PATH if (IMG2IMG_LORA_PATH and + os.path.isfile(IMG2IMG_LORA_PATH)) else None + if lora_path is None and IMG2IMG_LORA_REPO: + try: + from huggingface_hub import hf_hub_download + lora_path = hf_hub_download(repo_id=IMG2IMG_LORA_REPO, + filename=IMG2IMG_LORA_FILE) + except Exception as lexc: # noqa: BLE001 + print(f"[img2img] LoRA download skipped ({lexc}); using base model") + lora_path = None + _img2img_state["lora_path"] = lora_path + _img2img_state["available"] = True + print(f"[img2img] available ({IMG2IMG_MODEL_ID}; " + f"LoRA={lora_path or 'none — base img2img (stage 1)'}); " + f"pipeline will build on first use inside the GPU worker") + except Exception as exc: # noqa: BLE001 + print(f"[img2img] unavailable, endpoint will report disabled: {exc}") + _img2img_state["available"] = False + finally: + _img2img_state["loaded"] = True + return _img2img_state["available"] + + +_img2img_build_lock = threading.Lock() + + +def _ensure_img2img_pipe(): + """Build (once) and return the Flux2KleinPipeline on the GPU, applying the + watercolour LoRA only if one was resolved (stage 2); otherwise it's plain + base img2img (stage 1). MUST be called from inside a @spaces.GPU context on + ZeroGPU so the CUDA context is created with a GPU attached. Idempotent + + thread-safe; caches into _img2img_state. Returns the pipe, or None on failure + (caller falls back).""" + pipe = _img2img_state.get("pipe") + if pipe is not None: + return pipe + with _img2img_build_lock: + pipe = _img2img_state.get("pipe") + if pipe is not None: + return pipe + try: + import torch + from diffusers import Flux2KleinPipeline + try: + from pillow_heif import register_heif_opener + register_heif_opener() # so HEIC/HEIF phone photos open + except Exception: # noqa: BLE001 + pass + lora_path = _img2img_state.get("lora_path") + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"[img2img] building {IMG2IMG_MODEL_ID} on {device} " + f"({'with LoRA' if lora_path else 'base, no LoRA'}) …") + pipe = Flux2KleinPipeline.from_pretrained( + IMG2IMG_MODEL_ID, torch_dtype=torch.bfloat16 + ).to(device) + if lora_path and os.path.isfile(lora_path): + from pathlib import Path as _Path + lf = _Path(lora_path) + pipe.load_lora_weights(str(lf.parent), weight_name=lf.name, + adapter_name=IMG2IMG_TRIGGER) + pipe.set_adapters([IMG2IMG_TRIGGER], adapter_weights=[IMG2IMG_LORA_SCALE]) + _img2img_state["lora_loaded"] = True + print("[img2img] pipeline ready (watercolour LoRA applied)") + else: + _img2img_state["lora_loaded"] = False + print("[img2img] pipeline ready (base model — stage 1, no LoRA)") + _img2img_state["pipe"] = pipe + return pipe + except Exception as exc: # noqa: BLE001 + print(f"[img2img] in-worker build failed: {exc}") + return None + + +def _img2img_square(img, size: int = IMG2IMG_SIZE): + """EXIF-normalise, flatten to RGB, and centre-crop a PIL image to size×size — + the shape Klein expects (inlined from backend/hf/3_infer_img2img.py).""" + from PIL import Image as PILImage, ImageOps + src = ImageOps.exif_transpose(img).convert("RGB") + w, h = src.size + s = min(w, h) + src = src.crop(((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2)) + return src.resize((size, size), PILImage.LANCZOS) + + +@spaces.GPU(duration=120) +def _img2img_infer(src): + """Develop one PIL photo into a watercolour PIL image, inside a @spaces.GPU + context. Builds/reuses the pipeline here (never in the main process, never + pickled across the boundary). `src` and the returned image are PIL (both + picklable), the only things that cross the GPU boundary. Generous duration + because the FIRST call also pays the ~14 GB cold load; later calls reuse the + cached pipe and are fast.""" + import torch + pipe = _ensure_img2img_pipe() + if pipe is None: + raise RuntimeError("img2img pipeline could not be initialized") + sq = _img2img_square(src, size=IMG2IMG_SIZE) + # Trigger-token prompt only makes sense once the LoRA is loaded; otherwise + # use the generic base-model watercolour prompt (stage 1). + prompt = IMG2IMG_PROMPT if _img2img_state.get("lora_loaded") else IMG2IMG_BASE_PROMPT + device = getattr(pipe, "device", None) + generator = torch.Generator( + str(device) if device is not None else "cuda" + ).manual_seed(IMG2IMG_SEED) + result = pipe( + prompt=prompt, + image=sq, + height=IMG2IMG_SIZE, + width=IMG2IMG_SIZE, + num_inference_steps=IMG2IMG_STEPS, + guidance_scale=IMG2IMG_GUIDANCE, + generator=generator, + ) + return result.images[0] + + +def _pil_to_data_url(img, fmt: str = "PNG") -> str: + import base64 + buf = io.BytesIO() + img.save(buf, format=fmt) + b64 = base64.b64encode(buf.getvalue()).decode("ascii") + return f"data:image/{fmt.lower()};base64,{b64}" + + +@app.post("/api/img2img") +async def img2img_rest(request: Request): + """Develop one uploaded photo into a watercolour painting. + Body: {data_url: "data:image/...;base64,...", name?: str}. + Returns {available, image: , model}. When the pipeline can't + run on this host it returns {available: false, image: null} so the caller can + keep the original photo rather than treating it as an error.""" + import base64 + if not _img2img_prep(): + return {"available": False, "image": None, "model": IMG2IMG_MODEL_ID} + try: + body = await request.json() + except Exception: # noqa: BLE001 + raise HTTPException(status_code=400, detail="invalid JSON body") + data_url = (body or {}).get("data_url", "") + name = str((body or {}).get("name", "upload"))[:60] + if not data_url: + raise HTTPException(status_code=400, detail="missing data_url") + try: + b64 = data_url.split(",", 1)[1] if "," in data_url else data_url + raw = base64.b64decode(b64) + from PIL import Image + src = Image.open(io.BytesIO(raw)) + with traced("img2img", model=IMG2IMG_MODEL_ID, file=name) as t: + out = _img2img_infer(src) + t.set_output({"size": list(out.size)}) + return {"available": True, "image": _pil_to_data_url(out), "model": IMG2IMG_MODEL_ID} + except HTTPException: + raise + except Exception as exc: # noqa: BLE001 + print(f"[img2img] generation failed: {exc}") + raise HTTPException(status_code=500, detail=str(exc)) + + +# --------------------------------------------------------------------------- # +# Voice → text via NVIDIA NeMo ASR (item 5). Streaming Nemotron model; needs a +# GPU to be practical. Loaded lazily; if NeMo/torch/GPU are missing the endpoint +# returns {available: false} so the UI can show a graceful message. +# --------------------------------------------------------------------------- # +ASR_MODEL_ID = os.environ.get("SOZAI_ASR_MODEL", "nvidia/nemotron-3.5-asr-streaming-0.6b") +ASR_SAMPLE_RATE = 16_000 +_asr_lock = threading.Lock() +_asr_state = {"loaded": False, "model": None, "unsupported_os": False} + +# --- ported from the reference Nemotron Spaces app.py (working logic) ------- # +# Cache-aware streaming "lookahead" contexts. Larger = more accurate, higher +# latency. Keys are shown in the UI; values are [left, right] att_context_size. +ASR_CHUNK_PROFILES = { + "Lowest latency - 80 ms": [56, 0], + "Fast - 160 ms": [56, 1], + "Balanced - 320 ms": [56, 3], + "Accurate - 560 ms": [56, 6], + "Most accurate - 1.12 s": [56, 13], +} +ASR_DEFAULT_PROFILE = "Most accurate - 1.12 s" # default per spec + +# Ordered (label, value) choices for the UI language picker. +ASR_LANGUAGE_CHOICES = [ + ["Auto-detect", "auto"], ["Arabic", "ar-AR"], ["Bulgarian", "bg-BG"], + ["Croatian", "hr-HR"], ["Czech", "cs-CZ"], ["Danish", "da-DK"], + ["Dutch", "nl-NL"], ["English (UK)", "en-GB"], ["English (US)", "en-US"], + ["Estonian", "et-EE"], ["Finnish", "fi-FI"], ["French (Canada)", "fr-CA"], + ["French (France)", "fr-FR"], ["German", "de-DE"], ["Greek", "el-GR"], + ["Hebrew", "he-IL"], ["Hindi", "hi-IN"], ["Hungarian", "hu-HU"], + ["Italian", "it-IT"], ["Japanese", "ja-JP"], ["Korean", "ko-KR"], + ["Latvian", "lv-LV"], ["Lithuanian", "lt-LT"], ["Maltese", "mt-MT"], + ["Mandarin", "zh-CN"], ["Norwegian Bokmal", "nb-NO"], ["Norwegian Nynorsk", "nn-NO"], + ["Polish", "pl-PL"], ["Portuguese (Brazil)", "pt-BR"], ["Portuguese (Portugal)", "pt-PT"], + ["Romanian", "ro-RO"], ["Russian", "ru-RU"], ["Slovak", "sk-SK"], + ["Slovenian", "sl-SI"], ["Spanish (Spain)", "es-ES"], ["Spanish (US)", "es-US"], + ["Swedish", "sv-SE"], ["Thai", "th-TH"], ["Turkish", "tr-TR"], + ["Ukrainian", "uk-UA"], ["Vietnamese", "vi-VN"], +] + +# 40 supported language-locales (value sent to set_inference_prompt). "auto" +# lets the model auto-detect. +ASR_LANGUAGE_CODES = { + "auto", "ar-AR", "bg-BG", "hr-HR", "cs-CZ", "da-DK", "nl-NL", "en-GB", + "en-US", "et-EE", "fi-FI", "fr-CA", "fr-FR", "de-DE", "el-GR", "he-IL", + "hi-IN", "hu-HU", "it-IT", "ja-JP", "ko-KR", "lv-LV", "lt-LT", "mt-MT", + "zh-CN", "nb-NO", "nn-NO", "pl-PL", "pt-BR", "pt-PT", "ro-RO", "ru-RU", + "sk-SK", "sl-SI", "es-ES", "es-US", "sv-SE", "th-TH", "tr-TR", "uk-UA", + "vi-VN", +} + + +def _asr_os_supported() -> bool: + """NeMo / Nemotron streaming ASR is only supported on Linux. On macOS and + Windows the dependency stack (and ZeroGPU pathways) don't work, so we report + 'not supported' rather than attempting a load that will error out.""" + import platform + return platform.system().lower() == "linux" + + +def _get_asr_model(): + if _asr_state["loaded"]: + return _asr_state["model"] + with _asr_lock: + if _asr_state["loaded"]: + return _asr_state["model"] + try: + if not _asr_os_supported(): + import platform + print(f"[asr] Nemotron is not supported on {platform.system()} " + f"(Linux only); ASR disabled.") + _asr_state["model"] = None + _asr_state["unsupported_os"] = True + else: + import nemo.collections.asr as nemo_asr + import torch + torch.set_grad_enabled(False) + torch.set_float32_matmul_precision("high") + # Load on CPU at process start. On ZeroGPU the GPU isn't attached + # yet, so the model must be created on CPU here and only moved to + # cuda INSIDE the @spaces.GPU call (see _asr_run). map_location + + # explicit .to("cpu") keep all weight init on CPU. + model = nemo_asr.models.ASRModel.from_pretrained( + model_name=ASR_MODEL_ID, + map_location=torch.device("cpu"), + ) + # Configure greedy RNNT/CTC decoding for cache-aware streaming, + # mirroring the reference Spaces app so conformer_stream_step works. + try: + from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecodingConfig + from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTDecodingConfig + if hasattr(model, "change_decoding_strategy") and hasattr(model, "decoding"): + if hasattr(model, "joint"): + decoding_cfg = RNNTDecodingConfig(fused_batch_size=-1) + decoding_cfg.greedy.use_cuda_graph_decoder = False + model.change_decoding_strategy(decoding_cfg) + else: + model.change_decoding_strategy(CTCDecodingConfig()) + except Exception as dexc: # noqa: BLE001 + print(f"[asr] decoding strategy setup skipped: {dexc}") + model = model.to(device="cpu", dtype=torch.float32) + model.eval() + _asr_state["model"] = model + except Exception as exc: # noqa: BLE001 + print(f"[asr] NeMo model unavailable: {exc}") + _asr_state["model"] = None + finally: + _asr_state["loaded"] = True + return _asr_state["model"] + + +def _prep_wav_16k_mono(src_path: str) -> str: + """Decode any uploaded audio to 16kHz mono WAV (what the ASR model expects).""" + import tempfile + out = tempfile.NamedTemporaryFile(delete=False, suffix=".wav") + out.close() + # Prefer soundfile+librosa; fall back to ffmpeg if present. + try: + import soundfile as sf + import numpy as np # noqa: F401 + try: + import librosa + audio, _sr = librosa.load(src_path, sr=16000, mono=True) + except Exception: + data, sr = sf.read(src_path) + if getattr(data, "ndim", 1) > 1: + data = data.mean(axis=1) + audio = data + if sr != 16000: + import numpy as np + # naive resample as a last resort + idx = (np.arange(int(len(audio) * 16000 / sr)) * sr / 16000).astype(int) + idx = idx[idx < len(audio)] + audio = audio[idx] + sf.write(out.name, audio, 16000, subtype="PCM_16") + return out.name + except Exception as exc: # noqa: BLE001 + # ffmpeg fallback + try: + import subprocess + subprocess.run(["ffmpeg", "-y", "-i", src_path, "-ar", "16000", "-ac", "1", out.name], + check=True, capture_output=True) + return out.name + except Exception as exc2: # noqa: BLE001 + print(f"[asr] audio prep failed: {exc} / {exc2}") + return src_path + + +def _detect_language(text: str) -> str: + """Best-effort language id of a transcript (graceful if langdetect absent).""" + text = (text or "").strip() + if not text: + return "" + try: + from langdetect import detect + return detect(text) + except Exception: # noqa: BLE001 + return "" + + +def _asr_configure(model, language: str, chunk_profile: str) -> None: + """Apply the chosen lookahead profile + language prompt to the model. + Ported from the reference Nemotron Spaces app.""" + att = ASR_CHUNK_PROFILES.get(chunk_profile, ASR_CHUNK_PROFILES[ASR_DEFAULT_PROFILE]) + if hasattr(model.encoder, "set_default_att_context_size"): + model.encoder.set_default_att_context_size(att_context_size=att) + lang = language if language in ASR_LANGUAGE_CODES else "auto" + if hasattr(model, "set_inference_prompt"): + try: + model.set_inference_prompt(lang) + if hasattr(model, "decoding") and hasattr(model.decoding, "set_strip_lang_tags"): + # auto mode emits a tag; strip it for a clean transcript + model.decoding.set_strip_lang_tags(True, lang_tag_pattern=None) + except Exception as exc: # noqa: BLE001 + print(f"[asr] inference-prompt setup skipped: {exc}") + + +def _asr_transcribe_stream(model, wav_path: str, language: str, chunk_profile: str) -> str: + """Cache-aware streaming transcription, ported from NVIDIA NeMo's + Apache-licensed streaming example (via the reference Spaces app.py). Returns + the final transcript text. Falls back to model.transcribe() on any issue.""" + import torch + from nemo.collections.asr.parts.utils.streaming_utils import CacheAwareStreamingAudioBuffer + from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis + + def _extract(hyps): + hyps = list(hyps) + if not hyps: + return [""] + if isinstance(hyps[0], Hypothesis): + return [h.text for h in hyps] + return [str(h) for h in hyps] + + def _drop_extra(step_num, pad_and_drop): + if step_num == 0 and not pad_and_drop: + return 0 + return model.encoder.streaming_cfg.drop_extra_pre_encoded + + _asr_configure(model, language, chunk_profile) + model_device = next(model.parameters()).device + + buf = CacheAwareStreamingAudioBuffer(model=model, online_normalization=False, + pad_and_drop_preencoded=False) + buf.append_audio_file(wav_path, stream_id=-1) + + cache_last_channel, cache_last_time, cache_last_channel_len = \ + model.encoder.get_initial_cache_state(batch_size=1) + previous_hypotheses = None + previous_pred_out = None + text = "" + for step_num, (chunk_audio, chunk_lengths) in enumerate(buf, start=1): + with torch.inference_mode(): + chunk_audio = chunk_audio.to(device=model_device, dtype=torch.float32) + chunk_lengths = chunk_lengths.to(device=model_device) + (previous_pred_out, transcribed_texts, cache_last_channel, cache_last_time, + cache_last_channel_len, previous_hypotheses) = model.conformer_stream_step( + processed_signal=chunk_audio, + processed_signal_length=chunk_lengths, + cache_last_channel=cache_last_channel, + cache_last_time=cache_last_time, + cache_last_channel_len=cache_last_channel_len, + keep_all_outputs=buf.is_buffer_empty(), + previous_hypotheses=previous_hypotheses, + previous_pred_out=previous_pred_out, + drop_extra_pre_encoded=_drop_extra(step_num - 1, False), + return_transcription=True, + ) + text = _extract(transcribed_texts)[0] + return (text or "").strip() + + +@spaces.GPU(duration=120) +def _asr_run(wav_path: str, language: str = "auto", + chunk_profile: str = ASR_DEFAULT_PROFILE) -> str: + """Transcribe one prepared 16k-mono wav. Tries the cache-aware streaming + path first (honors language + latency/accuracy profile); falls back to the + plain transcribe() API if streaming isn't available for this checkpoint. + + ZeroGPU specifics, mirroring NVIDIA's reference Nemotron Space: + * The model is fetched from module state HERE, not passed as an argument — + @spaces.GPU pickles arguments to the GPU worker, and we don't want to + ship the NeMo model across that boundary. + * The model lives on CPU between calls (loaded on CPU at startup). Inside + this @spaces.GPU context the GPU is attached, so we move it to cuda for + the call and back to cpu in `finally`, exactly like the reference app. + This is what keeps CUDA out of the main process. + """ + import torch + model = _asr_state.get("model") + if model is None: + return "" + device = "cuda" if torch.cuda.is_available() else "cpu" + if device == "cuda": + torch.backends.cuda.matmul.allow_tf32 = True + try: + model.to(device=device, dtype=torch.float32) + try: + return _asr_transcribe_stream(model, wav_path, language, chunk_profile) + except Exception as exc: # noqa: BLE001 + print(f"[asr] streaming path failed ({exc}); falling back to transcribe()") + try: + outputs = model.transcribe([wav_path]) + first = outputs[0] if outputs else None + return (getattr(first, "text", None) or (first if isinstance(first, str) else "") or "").strip() + except Exception as exc2: # noqa: BLE001 + print(f"[asr] transcribe() fallback also failed: {exc2}") + return "" + finally: + # release the GPU: move weights back to CPU and clear the cache so the + # next ZeroGPU request starts clean. + try: + if device != "cpu": + model.to(device="cpu", dtype=torch.float32) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except Exception: # noqa: BLE001 + pass + + +@app.api(name="transcribe") +def transcribe(audio_path: str = "", field: str = "") -> dict: + """Transcribe a short audio clip to text with NeMo ASR and detect its + language. Returns {text, language, available, field}. If the model can't be + loaded (no NeMo/GPU), returns available=false with empty text.""" + model = _get_asr_model() + if model is None: + _record_trace({"name": "transcribe", "model": ASR_MODEL_ID, + "file": os.path.basename(audio_path or ""), "status": "unavailable", + "error": "NeMo ASR model not loaded (needs nemo_toolkit + GPU)", + "meta": {"field": field}, "output": None, "duration_ms": 0}) + return {"text": "", "language": "", "available": False, "field": field} + + wav = _prep_wav_16k_mono(audio_path) + text = "" + with traced("transcribe", model=ASR_MODEL_ID, + file=os.path.basename(audio_path or ""), meta={"field": field}) as t: + try: + # Go through _asr_run so the actual inference runs inside a + # @spaces.GPU context on ZeroGPU (and a no-op context elsewhere). + text = _asr_run(wav).strip() + except Exception as exc: # noqa: BLE001 + t.error = f"transcribe failed: {exc}" + t.set_output(text) + lang = _detect_language(text) + return {"text": text, "language": lang, "available": True, "field": field} + + +@app.post("/api/transcribe") +async def transcribe_rest(request: Request): + """Transcribe recorder audio sent as a base64 data URL. + Body: {data_url: "data:audio/webm;base64,...", field: "caption"}. + Returns {text, language, available, field}.""" + import base64 + import tempfile + try: + body = await request.json() + except Exception: # noqa: BLE001 + return {"text": "", "language": "", "available": False, "field": ""} + data_url = (body or {}).get("data_url", "") + field = (body or {}).get("field", "") + language = (body or {}).get("language", "auto") or "auto" + chunk_profile = (body or {}).get("chunk_profile", ASR_DEFAULT_PROFILE) or ASR_DEFAULT_PROFILE + model = _get_asr_model() + unsupported = bool(_asr_state.get("unsupported_os")) + if model is None or not data_url: + if not data_url: + return {"text": "", "language": "", "available": model is not None, + "unsupported_os": unsupported, "field": field} + _record_trace({"name": "transcribe", "model": ASR_MODEL_ID, "file": "recording", + "status": "unavailable", + "error": ("Nemotron not supported on this OS (Linux only)" if unsupported + else "NeMo ASR model not loaded (needs nemo_toolkit + GPU)"), + "meta": {"field": field}, "output": None, "duration_ms": 0}) + return {"text": "", "language": "", "available": False, + "unsupported_os": unsupported, "field": field} + try: + header, b64 = (data_url.split(",", 1) if "," in data_url else ("", data_url)) + ext = ".webm" + if "ogg" in header: + ext = ".ogg" + elif "wav" in header: + ext = ".wav" + elif "mp4" in header or "m4a" in header: + ext = ".mp4" + raw = base64.b64decode(b64) + src = tempfile.NamedTemporaryFile(delete=False, suffix=ext) + src.write(raw) + src.close() + wav = _prep_wav_16k_mono(src.name) + text = "" + with traced("transcribe", model=ASR_MODEL_ID, file="recording", + meta={"field": field, "language": language, "chunk_profile": chunk_profile}) as t: + try: + text = _asr_run(wav, language=language, chunk_profile=chunk_profile) + except Exception as exc: # noqa: BLE001 + t.error = f"transcribe failed: {exc}" + t.set_output(text) + return {"text": text, "language": _detect_language(text), "available": True, + "unsupported_os": False, "field": field} + except Exception as exc: # noqa: BLE001 + print(f"[asr] rest transcribe failed: {exc}") + return {"text": "", "language": "", "available": True, "field": field, "error": str(exc)} + + +# --------------------------------------------------------------------------- # +# Auto-caption inference. MiniCPM-V is a vision model, so the heavy generation +# is isolated in a @spaces.GPU-decorated helper: on ZeroGPU this requests a GPU +# for the duration of the call and releases it after; off ZeroGPU it's a no-op. +# Both backends (gguf via llama.cpp, transformers via .chat) are image-grounded. +# --------------------------------------------------------------------------- # + +@spaces.GPU(duration=60) +def _caption_infer(local_img, base_instruction, full_prompt, reply_suffix, max_new): + """Run one vision generation and return raw model text. + + `full_prompt` already bundles the per-kind instruction + reply directive, so + the same helper produces a title, a caption, or tags depending on what the + caller passed in. The IMAGE is sent to the model in both backends. + + IMPORTANT (ZeroGPU): two rules shape this function. + 1. @spaces.GPU ships ARGUMENTS to the GPU worker via pickle, and the + llama.cpp `Llama` object is NOT picklable ("cannot pickle 'module' + object"). So the model is never an argument — only plain str/int cross + the boundary; the model is obtained here. + 2. CUDA must NOT be initialized in the main process on ZeroGPU (GPUs only + exist inside this decorated call). So for the GGUF backend the model is + BUILT HERE on first use, while the GPU is attached, and cached in the + worker — never constructed in the main process. + """ + backend = _caption_state.get("backend") or "transformers" + if backend == "gguf": + # Build (once) and reuse the llama.cpp model inside the GPU worker so the + # CUDA context is created with a GPU attached. _ensure_gguf_model() + # caches it in module state; on ZeroGPU that state lives in the worker. + model = _ensure_gguf_model() + if model is None: + raise RuntimeError("GGUF caption model could not be initialized") + # llama.cpp vision chat completion: send the real image as a data: URI + # plus the instruction text; the mmproj-backed handler makes visual + # tokens. + user_content = [{"type": "text", "text": full_prompt}] + data_uri = _image_data_uri(local_img) + if data_uri: + user_content.insert(0, {"type": "image_url", "image_url": {"url": data_uri}}) + messages = [ + {"role": "system", "content": base_instruction}, + {"role": "user", "content": user_content}, + ] + out = model.create_chat_completion( + messages=messages, max_tokens=max_new, temperature=0.3, + ) + return (out["choices"][0]["message"]["content"] or "").strip() + + # transformers backend: MiniCPM-V's .chat() helper takes the PIL image and a + # msgs list. This keeps the off-Spaces path image-grounded too. + model = _caption_state["model"] + tokenizer = _caption_state["tokenizer"] + pil = _load_pil_image(local_img) + msgs = [{"role": "user", "content": [pil, full_prompt] if pil is not None else [full_prompt]}] + answer = model.chat( + msgs=msgs, + image=pil, + tokenizer=tokenizer, + sampling=True, + temperature=0.3, + max_new_tokens=max_new, + ) + if isinstance(answer, (list, tuple)): + answer = answer[0] if answer else "" + return (str(answer) or "").strip() + + +@app.api(name="autocaption") +def autocaption(image_path: str = "", alt: str = "", prompt: str = "", kind: str = "caption") -> str: + """Generate a caption, title, or tag suggestion for a photo by LOOKING AT + the photo with MiniCPM-V. + + `kind` is "caption", "title", or "tags" — it selects the instruction, the + reply directive, the generation length, and the trace span name, so the one + function serves all three call sites (caption field, title field, tags + field) with image-grounded output. + + `prompt` is an optional user-supplied instruction (from Settings) that lets + people customize the style. When empty, a sensible default is used. + + On any failure (model missing, generation error) it returns a lightweight + fallback derived from the photo's alt text so the UI button always works. + """ + kind_l = str(kind).lower() + is_title = kind_l == "title" + is_tags = kind_l == "tags" + span_name = "autotags" if is_tags else ("autotitle" if is_title else "autocaption") + default_instruction = ( + "Look at this personal photo and list 3 to 6 short, lowercase tags " + "describing what is in it, separated by commas." + if is_tags else + "Look at this personal photo and write a very short, evocative title " + "(2–5 words) for it." + if is_title else + "Look at this personal photo and write a short, warm one-sentence caption " + "for a personal photo album." + ) + reply_suffix = ( + " Reply with only the comma-separated tags." if is_tags else + " Reply with only the title." if is_title else + " Reply with only the caption." + ) + max_new = 32 if is_tags else 24 if is_title else 128 + + if not _load_caption_model(): + if is_tags: + cap = _fallback_tags(alt) + else: + cap = _fallback_caption(alt) + _record_trace({"name": span_name, "kind": "LLM", "model": AUTOCAPTION_MODEL_ID, + "file": os.path.basename(image_path or ""), "status": "fallback", + "error": "caption model unavailable", "output": cap[:300], + "meta": {"prompt": (prompt or "")[:120], "for": kind}, "duration_ms": 0}) + return cap + + with traced(span_name, model=AUTOCAPTION_MODEL_ID, + file=os.path.basename(image_path or ""), + meta={"prompt": (prompt or "")[:120], "alt": (alt or "")[:120], "for": kind}, + invocation={"max_new_tokens": max_new, "temperature": 0.3}, + system=default_instruction) as t: + try: + local_img = _resolve_image_for_model(image_path) + # Some frontends pass the photo as a data URL in `alt`; if image_path + # didn't resolve to a real file, try alt as an image source too. + if local_img is None and alt and alt.strip().startswith("data:"): + local_img = _resolve_image_for_model(alt) + # If we DID get an image, don't pollute the prompt with the filename; + # the pixels are the context. Only fall back to an alt text hint when + # no image could be resolved. + hint = "" if local_img else (f' (context hint: {alt})' if alt else "") + base_instruction = (prompt or "").strip() or default_instruction + full_prompt = base_instruction + hint + reply_suffix + + text = _caption_infer(local_img, base_instruction, + full_prompt, reply_suffix, max_new) + + if is_tags: + # tags must come out comma-separated; normalize whatever the + # model produced, then validate. If it can't be made into a + # comma list, fall back (which is always comma-separated). + norm = _normalize_tags(text) + if not _is_comma_separated(norm): + norm = _fallback_tags(alt) + result = norm or _fallback_tags(alt) + else: + # tidy: first line only, strip surrounding quotes + text = (text.splitlines()[0].strip().strip('"').strip()) if text else "" + result = text or _fallback_caption(alt) + t.set_output(result) + return result + except Exception as exc: + print(f"[{span_name}] generation failed, using fallback: {exc}") + t.error = f"generation failed: {exc}" + fb = _fallback_tags(alt) if is_tags else _fallback_caption(alt) + t.set_output(fb) + return fb + + +# ----------------------------- Room REST API ------------------------------- # + +_JOIN_HITS: dict = {} + + +def rate_ok(ip: str, limit=20, window=60) -> bool: + now = time.time() + hits = _JOIN_HITS.setdefault(ip, []) + hits[:] = [t for t in hits if now - t < window] + if len(hits) >= limit: + return False + hits.append(now) + return True + + +@app.post("/api/rooms") +async def create_room(request: Request): + """Create a room. Body: { code_length?, display_name?, session_id? }.""" + try: + data = await request.json() + except Exception: + data = {} + code_length = clamp_code_len(data.get("code_length", 6)) + name = sanitize_name(data.get("display_name")) + session_id = (data.get("session_id") or secrets.token_hex(8))[:64] + + room_id = str(uuid.uuid4()) # secure, never sequential + share_code = unique_share_code(code_length) + now = time.time() + with DB_LOCK: + _conn.execute( + "INSERT INTO rooms (room_id, share_code, code_length, created_at, " + "last_activity, owner_session_id, is_active, metadata) " + "VALUES (?,?,?,?,?,?,1,?)", + (room_id, share_code, code_length, now, now, session_id, json.dumps({})), + ) + _conn.commit() + return { + "room_id": room_id, + "share_code": share_code, + "code_length": code_length, + "room_path": f"/room/{room_id}", + } + + +@app.get("/api/rooms/{room_id}") +def get_room(room_id: str): + row = room_row(room_id) + if row is None: + raise HTTPException(status_code=404, detail="Room not found or expired.") + return { + "room_id": row["room_id"], + "share_code": row["share_code"], + "code_length": row["code_length"], + "is_active": bool(row["is_active"]), + "participants": manager.presence(room_id), + } + + +@app.post("/api/join-by-code") +async def join_by_code(request: Request): + """Resolve a human room code to a room_id. Body: { code }.""" + ip = request.client.host if request.client else "?" + if not rate_ok(ip): + raise HTTPException(status_code=429, detail="Too many attempts. Slow down.") + try: + data = await request.json() + except Exception: + data = {} + code = (data.get("code") or "").strip().upper() + code = "".join(ch for ch in code if ch in CODE_ALPHABET) + if not (MIN_CODE_LEN <= len(code) <= MAX_CODE_LEN): + raise HTTPException(status_code=400, detail="Invalid code format.") + with DB_LOCK: + row = _conn.execute( + "SELECT room_id FROM rooms WHERE share_code=? AND is_active=1", (code,) + ).fetchone() + if row is None: + raise HTTPException(status_code=404, detail="No active room with that code.") + return {"room_id": row["room_id"], "room_path": f"/room/{row['room_id']}"} + + +# --------------------------- Room content persistence ---------------------- # +# A room's scrapbook (photos, captions, pins, transforms) is stored as one JSON +# blob with a monotonic version. The frontend PUTs on change (debounced) and +# GETs on join, so a returning user sees the room as it was left. Live edits +# still flow over the WebSocket; this is the durable backstop + late-join seed. + +# cap the stored blob so a room can't be used as unbounded storage (data URLs +# for images can be large; 8 MB is generous for a handful of photos) +MAX_ROOM_STATE_BYTES = int(os.environ.get("SOZAI_MAX_ROOM_STATE_BYTES", str(8 * 1024 * 1024))) + + +@app.get("/api/rooms/{room_id}/state") +def get_room_state(room_id: str): + """Return the persisted scrapbook state for a room: {state, version}. + state is null if nothing has been saved yet.""" + if room_row(room_id) is None: + raise HTTPException(status_code=404, detail="Room not found or expired.") + with DB_LOCK: + row = _conn.execute( + "SELECT state, version, updated_at FROM room_state WHERE room_id=?", (room_id,) + ).fetchone() + if row is None: + return {"state": None, "version": 0, "updated_at": None} + try: + state = json.loads(row["state"]) if row["state"] else None + except Exception: # noqa: BLE001 + state = None + return {"state": state, "version": row["version"], "updated_at": row["updated_at"]} + + +@app.post("/api/rooms/{room_id}/state") +async def put_room_state(room_id: str, request: Request): + """Persist the room's scrapbook state. Body: {state, base_version?}. + + Uses optimistic concurrency: if base_version is supplied and no longer + matches the stored version, the write is rejected with the current state so + the client can reconcile (last-writer is NOT blindly allowed to clobber a + newer save). Returns {version, conflict?, state?}. + """ + if room_row(room_id) is None: + raise HTTPException(status_code=404, detail="Room not found or expired.") + try: + body = await request.json() + except Exception: # noqa: BLE001 + raise HTTPException(status_code=400, detail="Invalid JSON.") + state = (body or {}).get("state") + base_version = (body or {}).get("base_version") + + blob = json.dumps(state, separators=(",", ":")) + if len(blob.encode("utf-8")) > MAX_ROOM_STATE_BYTES: + raise HTTPException(status_code=413, detail="Scrapbook state too large to save.") + + now = time.time() + with DB_LOCK: + cur = _conn.execute( + "SELECT state, version FROM room_state WHERE room_id=?", (room_id,) + ).fetchone() + current_version = cur["version"] if cur else 0 + # optimistic concurrency check + if base_version is not None and cur is not None and int(base_version) != int(current_version): + try: + current_state = json.loads(cur["state"]) if cur["state"] else None + except Exception: # noqa: BLE001 + current_state = None + return {"version": current_version, "conflict": True, "state": current_state} + new_version = current_version + 1 + _conn.execute( + "INSERT INTO room_state (room_id, state, version, updated_at) VALUES (?,?,?,?) " + "ON CONFLICT(room_id) DO UPDATE SET state=excluded.state, " + "version=excluded.version, updated_at=excluded.updated_at", + (room_id, blob, new_version, now), + ) + _conn.commit() + touch_room(room_id) + return {"version": new_version, "conflict": False} + + +# --------------------------- WebSocket sync -------------------------------- # + +_PALETTE = [ + "#7c5cff", "#e9698f", "#2bb673", "#f2994a", "#3aa0ff", + "#9b51e0", "#eb5757", "#27ae60", "#f2c94c", "#56ccf2", +] + + +def color_for(session_id: str) -> str: + h = int(hashlib.sha1(session_id.encode()).hexdigest(), 16) + return _PALETTE[h % len(_PALETTE)] + + +# Events the server relays verbatim to the other participants. +# +# Every collaborative interaction the frontend emits must be listed here, or the +# server silently drops it and collaborators never see the action. Rather than a +# brittle allow-list that has to be hand-updated whenever the frontend grows a +# new event, we relay ANY event that is NOT a server-owned/internal type +# (see NON_RELAY_EVENTS below). This keeps cursors, selections, field typing, +# polaroid drags, card edits, screen-follow, etc. all in sync. +RELAY_EVENTS = { + # object lifecycle + text field edits (title / caption / date / time) + "object_created", "object_updated", "object_deleted", + # chat + presence-style ephemera + "chat_message", "cursor_position", "selection", "interaction", + # the photo the editor is currently viewing (drives optional follow) + "active_image", + # darkroom / scrapbook card edits (time / location / note) + "dv_edit", + # polaroid pins on the map: drop, move, style (rotate/scale), reset + "photo_pinned", "photo_moved", "photo_props", "photos_reset", + # navigation that collaborators can optionally follow + "screen_changed", "variant_changed", +} + +# Types the server originates or manages itself — these must NEVER be relayed +# back out from a client (a client cannot, e.g., forge presence for others). +NON_RELAY_EVENTS = { + "init", "me", "presence", "presence_update", "status", + "user_joined", "user_left", "rename", "error", + "join_request", "join_pending", "join_declined", "admit_join", "decline_join", "cancel_join", +} + + +class Manager: + def __init__(self): + self.rooms = {} # room_id -> {ws: info} (admitted participants) + self.pending = {} # room_id -> {ws: info} (awaiting host approval) + + async def connect(self, room_id, ws, info): + await ws.accept() + self.rooms.setdefault(room_id, {})[ws] = info + + def admit(self, room_id, ws, info): + """Promote a pending socket to a full participant.""" + self.pending.get(room_id, {}).pop(ws, None) + self.rooms.setdefault(room_id, {})[ws] = info + + def add_pending(self, room_id, ws, info): + self.pending.setdefault(room_id, {})[ws] = info + + def remove_pending(self, room_id, ws): + p = self.pending.get(room_id) + if p and ws in p: + del p[ws] + if not p: + self.pending.pop(room_id, None) + + def find_pending(self, room_id, session_id): + """Return (ws, info) for a pending session_id, or (None, None).""" + for ws, info in list(self.pending.get(room_id, {}).items()): + if info["session_id"] == session_id: + return ws, info + return None, None + + def host_ws(self, room_id, owner_session_id): + """The socket that should approve joins: the owner if connected, else + the earliest-connected participant (so approval never deadlocks).""" + peers = self.rooms.get(room_id, {}) + if not peers: + return None + if owner_session_id: + for ws, info in peers.items(): + if info["session_id"] == owner_session_id: + return ws + # fall back to the first admitted participant + return next(iter(peers), None) + + def disconnect(self, room_id, ws): + peers = self.rooms.get(room_id) + if peers and ws in peers: + del peers[ws] + if not peers: + self.rooms.pop(room_id, None) + self.remove_pending(room_id, ws) + + def presence(self, room_id): + peers = self.rooms.get(room_id, {}) + seen, out = set(), [] + for info in peers.values(): + if info["session_id"] in seen: + continue + seen.add(info["session_id"]) + out.append({k: info[k] for k in ("session_id", "name", "participant_id", "color")}) + return out + + async def broadcast(self, room_id, message, exclude=None): + for ws in list(self.rooms.get(room_id, {})): + if ws is exclude: + continue + try: + await ws.send_json(message) + except Exception: + pass + + async def send_to(self, ws, message): + try: + await ws.send_json(message) + return True + except Exception: + return False + + +manager = Manager() + + +@app.websocket("/ws/{room_id}") +async def ws_room(ws: WebSocket, room_id: str): + if room_row(room_id) is None: + await ws.accept() + await ws.send_json({"type": "error", "message": "Room not found or expired."}) + await ws.close(code=4404) + return + + session_id = (ws.query_params.get("session_id") or secrets.token_hex(8))[:64] + info = { + "session_id": session_id, + "name": sanitize_name(ws.query_params.get("name")), + "participant_id": secrets.token_hex(6), + "color": color_for(session_id), + } + + # ---- join approval ---- + # The owner is auto-admitted. Anyone else needs the host to admit them. + # If no host is currently connected, we auto-admit (so a room is never + # un-joinable / approval never deadlocks). Returning sessions that the host + # already admitted earlier in this room are also let straight back in. + row = room_row(room_id) + owner_session = row["owner_session_id"] if row else None + is_owner = bool(owner_session) and session_id == owner_session + + await ws.accept() + host = manager.host_ws(room_id, owner_session) + needs_approval = (not is_owner) and (host is not None) + + if needs_approval: + decision = asyncio.Event() + info["_decision_event"] = decision + info["_admitted"] = False + manager.add_pending(room_id, ws, info) + # tell the joiner they're waiting + await manager.send_to(ws, {"type": "join_pending", "you": {k: info[k] for k in ("session_id", "name", "participant_id", "color")}}) + # ask the host to admit / decline + await manager.send_to(host, { + "type": "join_request", + "participant": {k: info[k] for k in ("session_id", "name", "participant_id", "color")}, + }) + # wait for the host's decision (with a timeout → auto-decline) + try: + await asyncio.wait_for(decision.wait(), timeout=120) + except asyncio.TimeoutError: + pass + except Exception: + pass + if not info.get("_admitted"): + manager.remove_pending(room_id, ws) + await manager.send_to(ws, {"type": "join_declined"}) + try: + await ws.close(code=4403) + except Exception: + pass + return + manager.admit(room_id, ws, info) + else: + # owner, or no host present → straight in + manager.rooms.setdefault(room_id, {})[ws] = info + + now = time.time() + with DB_LOCK: + _conn.execute( + "INSERT OR REPLACE INTO participants (participant_id, room_id, " + "session_id, display_name, joined_at, last_seen) VALUES (?,?,?,?,?,?)", + (info["participant_id"], room_id, session_id, info["name"], now, now), + ) + _conn.commit() + touch_room(room_id) + + await ws.send_json({"type": "init", "you": info, "participants": manager.presence(room_id)}) + await manager.broadcast(room_id, {"type": "user_joined", "participant": info}, exclude=ws) + await manager.broadcast(room_id, {"type": "presence_update", "participants": manager.presence(room_id)}) + + try: + while True: + data = await ws.receive_json() + t = data.get("type") + if t == "rename": + info["name"] = sanitize_name(data.get("name")) + with DB_LOCK: + _conn.execute( + "UPDATE participants SET display_name=?, last_seen=? WHERE participant_id=?", + (info["name"], time.time(), info["participant_id"]), + ) + _conn.commit() + await manager.broadcast(room_id, {"type": "presence_update", "participants": manager.presence(room_id)}) + continue + # ---- host approves / declines a pending joiner ---- + if t in ("admit_join", "decline_join"): + target_sid = data.get("session_id") + pend_ws, pend_info = manager.find_pending(room_id, target_sid) + if pend_info is not None: + pend_info["_admitted"] = (t == "admit_join") + ev = pend_info.get("_decision_event") + if ev is not None: + ev.set() + continue + # explicitly known (RELAY_EVENTS) or simply not a server-owned type + # (NON_RELAY_EVENTS) — this way the frontend can introduce new + # interaction events without the server silently dropping them. + if t and t not in NON_RELAY_EVENTS: + data["from"] = {k: info[k] for k in ("session_id", "name", "participant_id", "color")} + await manager.broadcast(room_id, data, exclude=ws) + # cursors + transient interaction pings are too chatty to keep + # the room alive on; everything else counts as real activity. + if t not in ("cursor_position", "interaction"): + touch_room(room_id) + except WebSocketDisconnect: + pass + except Exception: + pass + finally: + manager.disconnect(room_id, ws) + await manager.broadcast(room_id, {"type": "user_left", "participant": info}) + await manager.broadcast(room_id, {"type": "presence_update", "participants": manager.presence(room_id)}) + + +# ------------------------------ Assets ------------------------------------- # + +def _safe_rel(path: str) -> str: + rel = os.path.normpath(path).lstrip("/\\") + if rel.startswith("..") or os.path.isabs(rel): + raise HTTPException(status_code=400, detail="bad path") + return rel.replace("\\", "/") + + +@app.get("/assets/{path:path}") +async def assets(request: Request, path: str): + rel = _safe_rel(path) + if PUBLIC_DIR is not None: + fp = os.path.join(PUBLIC_DIR, *rel.split("/")) + if os.path.isfile(fp): + return FileResponse(fp) + if ZIP_PATH is not None: + try: + data = _read_from_zip(rel) + media = mimetypes.guess_type(rel)[0] or "application/octet-stream" + return Response(content=data, media_type=media) + except KeyError: + pass + # Not a Sozai asset → it's almost certainly a Phoenix bundle asset (its JS + # requests /assets/vendor-*.js etc. with an absolute path). Proxy it to the + # in-process Phoenix server so the embedded trace UI loads correctly. + if _phoenix_session is not None: + return await _phoenix_proxy(request, "assets/" + rel) + raise HTTPException(status_code=404, detail=f"asset not found: {rel}") + + +@app.get("/favicon.ico") +def favicon(): + rel = "icon.svg" + if PUBLIC_DIR is not None: + fp = os.path.join(PUBLIC_DIR, rel) + if os.path.isfile(fp): + return FileResponse(fp) + if ZIP_PATH is not None: + try: + return Response(content=_read_from_zip(rel), media_type="image/svg+xml") + except KeyError: + pass + raise HTTPException(status_code=404, detail="favicon not found") + + +# ------------------------- Easter-egg sprite API --------------------------- # + +@app.get("/api/sprites") +def sprite_animals(): + """List the available pet sprite sets (animals) for the easter-egg picker.""" + return {"animals": list_sprite_animals(), "frames": SPRITE_FILES} + + +@app.get("/sprites/{animal}/{fname}") +def sprite_image(animal: str, fname: str): + """Serve one sprite PNG, e.g. /sprites/tabby/still or /sprites/dog/nrun1.""" + name = fname[:-4] if fname.lower().endswith(".png") else fname + try: + data = _read_sprite(animal, name) + except KeyError: + raise HTTPException(status_code=404, detail="sprite not found") + return Response(content=data, media_type="image/png") + + +# ------------------------- Vendored libraries ----------------------------- # +# Serve MapLibre GL JS (and CSS) from our OWN origin. On first request we fetch +# it once from a CDN and cache it to disk; thereafter it is served locally. +# This makes the map work even when the *browser* can't reach public CDNs +# (locked-down networks, offline kiosks, etc.) as long as the SERVER has +# outbound internet. Falls back to multiple CDNs in the page if this is empty. + +_VENDOR_DIR = os.path.join(BASE, ".vendor_cache") +_MAPLIBRE_VER = os.environ.get("SOZAI_MAPLIBRE_VER", "4.7.1") +_VENDOR_FILES = { + "maplibre-gl.js": [ + f"https://unpkg.com/maplibre-gl@{_MAPLIBRE_VER}/dist/maplibre-gl.js", + f"https://cdn.jsdelivr.net/npm/maplibre-gl@{_MAPLIBRE_VER}/dist/maplibre-gl.js", + f"https://cdnjs.cloudflare.com/ajax/libs/maplibre-gl/{_MAPLIBRE_VER}/maplibre-gl.js", + ], + "maplibre-gl.css": [ + f"https://unpkg.com/maplibre-gl@{_MAPLIBRE_VER}/dist/maplibre-gl.css", + f"https://cdn.jsdelivr.net/npm/maplibre-gl@{_MAPLIBRE_VER}/dist/maplibre-gl.css", + f"https://cdnjs.cloudflare.com/ajax/libs/maplibre-gl/{_MAPLIBRE_VER}/maplibre-gl.css", + ], +} +_vendor_lock = threading.Lock() + + +def _vendor_fetch(fname: str): + """Return cached bytes for a vendored file, downloading once if needed.""" + if fname not in _VENDOR_FILES: + return None + cache_path = os.path.join(_VENDOR_DIR, fname) + if os.path.isfile(cache_path) and os.path.getsize(cache_path) > 0: + with open(cache_path, "rb") as fh: + return fh.read() + with _vendor_lock: + if os.path.isfile(cache_path) and os.path.getsize(cache_path) > 0: + with open(cache_path, "rb") as fh: + return fh.read() + os.makedirs(_VENDOR_DIR, exist_ok=True) + import urllib.request + for url in _VENDOR_FILES[fname]: + try: + req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0 sozai"}) + with urllib.request.urlopen(req, timeout=20) as resp: + data = resp.read() + if data: + with open(cache_path, "wb") as fh: + fh.write(data) + print(f"[vendor] cached {fname} from {url} ({len(data)} bytes)") + return data + except Exception as exc: + print(f"[vendor] failed {url}: {exc}") + return None + + +@app.get("/vendor/{fname}") +def vendor_file(fname: str): + """Serve a vendored library (e.g. /vendor/maplibre-gl.js) from local cache.""" + data = _vendor_fetch(fname) + if data is None: + raise HTTPException(status_code=404, detail="vendor file unavailable") + media = "text/javascript" if fname.endswith(".js") else ("text/css" if fname.endswith(".css") else "application/octet-stream") + return Response(content=data, media_type=media) + + +# ------------------------------- Pages ------------------------------------- # + +def _read_html(name: str) -> str: + with open(os.path.join(BASE, name), "r", encoding="utf-8") as fh: + return fh.read() + + +def _inline_maplibre(html: str) -> str: + """Replace the marker in the page with the cached + MapLibre CSS + JS inlined directly. This is the most robust delivery: the + browser executes the library inline with no extra request, sidestepping any + routing/MIME/proxy issues. If the library can't be cached (server offline), + the marker is left as-is and the page's CDN fallback loader takes over. + """ + marker = "" + if marker not in html: + return html + js = _vendor_fetch("maplibre-gl.js") + css = _vendor_fetch("maplibre-gl.css") + if not js: + # leave the marker; the in-page CDN fallback will handle loading + return html + parts = [] + if css: + css_text = css.decode("utf-8", "replace").replace("\n" + css_text + "\n") + js_text = js.decode("utf-8", "replace").replace("\n" + js_text + "\n") + return html.replace(marker, "\n".join(parts), 1) + + +def _app_html() -> str: + return _inline_maplibre(_read_html("index.html")) + + +@app.get("/", response_class=HTMLResponse) +async def landing(): + return _read_html("landing.html") + + +@app.get("/app", response_class=HTMLResponse) +async def solo_app(): + return _app_html() + + +@app.get("/room/{room_id}", response_class=HTMLResponse) +async def room_app(room_id: str): + return _app_html() + + +# --------------------------- Cleanup daemon -------------------------------- # + +def _cleanup_loop(): + runs = 0 + while True: + time.sleep(300) # every 5 minutes + try: + now = time.time() + runs += 1 + with DB_LOCK: + # 1) flag rooms idle for > IDLE_TIMEOUT as inactive (so they + # can't be joined) — kept briefly for a grace window. + _conn.execute( + "UPDATE rooms SET is_active=0 WHERE is_active=1 AND ?-last_activity>?", + (now, IDLE_TIMEOUT), + ) + # 2) DELETE rooms (and their content) once they've been idle + # past the deletion window. This is keyed on last_activity, + # not creation time, so the DB tracks live usage and the + # heavy room_state blobs don't pile up. + _conn.execute( + "DELETE FROM rooms WHERE ?-last_activity>?", (now, ROOM_DELETE_AFTER_IDLE) + ) + # absolute safety cap: nothing survives past HARD_EXPIRATION + _conn.execute("DELETE FROM rooms WHERE ?-created_at>?", (now, HARD_EXPIRATION)) + # 3) cascade: drop orphaned participants + saved scrapbook state + _conn.execute( + "DELETE FROM participants WHERE room_id NOT IN (SELECT room_id FROM rooms)" + ) + _conn.execute( + "DELETE FROM room_state WHERE room_id NOT IN (SELECT room_id FROM rooms)" + ) + _conn.commit() + # 4) reclaim disk space. SQLite doesn't shrink the file on its + # own after deletes; incremental_vacuum returns freed pages to + # the OS. Run occasionally (hourly) to keep the file compact. + if runs % 12 == 0: + try: + _conn.execute("PRAGMA incremental_vacuum;") + _conn.commit() + except Exception: + pass + except Exception: + pass + + +threading.Thread(target=_cleanup_loop, daemon=True).start() + + +if __name__ == "__main__": + # Bring up the in-process Phoenix UI + OTel tracing immediately, so the + # embedded "LLM Trace" section is live as soon as the app starts — no + # separate `phoenix serve` required. + try: + _get_tracer() + print(f"[phoenix] embedded trace UI ready at {_phoenix_session_url()}") + except Exception as exc: # noqa: BLE001 + print(f"[phoenix] startup tracing init skipped: {exc}") + + # On Windows, asyncio's Proactor loop logs a noisy but harmless + # "ConnectionResetError: [WinError 10054]" whenever a client (browser tab, + # WebSocket, hot-reload) drops a connection abruptly. It's cosmetic — the + # connection simply closed — so we wrap the internal callback to ignore that + # one error. This affects nothing else and is a no-op off Windows. + import platform as _platform + if _platform.system() == "Windows": + try: + from asyncio.proactor_events import _ProactorBasePipeTransport + + _orig_ccl = _ProactorBasePipeTransport._call_connection_lost + + def _quiet_ccl(self, exc): + if isinstance(exc, ConnectionResetError): + return # ignore the benign 10054 reset + return _orig_ccl(self, exc) + + _ProactorBasePipeTransport._call_connection_lost = _quiet_ccl + except Exception: + pass # if internals change, just leave the default behavior + + app.launch(server_name="0.0.0.0", server_port=7860) \ No newline at end of file