""" 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") # Disable the Xet downloader: on this Space it both fails to write in the GPU # worker AND crawls at ~1 kB/s for the 16 GB FLUX weights. Forcing the standard # HTTPS download path makes the main-process pre-fetch actually complete. os.environ.setdefault("HF_HUB_DISABLE_XET", "1") os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "0") 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 from fastapi.responses import FileResponse, HTMLResponse, Response, StreamingResponse 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", "") # Free Stadia Maps API key (https://client.stadiamaps.com) for the Stamen # Watercolor basemap on the Cesium globe. Optional — without it the globe keeps # its default imagery. STADIA_MAPS_KEY = os.environ.get("SOZAI_STADIA_KEY", "") @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), "stadiaMapsKey": STADIA_MAPS_KEY, "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, "snapshot_path": None, "snapshot_ready": False, "dl_started": False} def _img2img_start_bg_download(): """Download the FLUX weights in the MAIN process (a daemon thread), where cache writes succeed. The ZeroGPU worker can request a GPU but CANNOT write the HF cache during a download (it fails with 'Permission denied'), so the model must be fetched here first; the worker then only READS the cache (from_pretrained with local_files_only). Runs once, non-blocking.""" if _img2img_state.get("dl_started"): return _img2img_state["dl_started"] = True def _run(): import time from huggingface_hub import snapshot_download # FLUX.2-klein is ~24 GB across ~25 shards; HF's CDN intermittently drops # the connection mid-shard ("peer closed connection"). snapshot_download # RESUMES partial files from the cache, so we retry with backoff until the # whole repo is present rather than failing the first time a shard drops. attempt = 0 while not _img2img_state.get("snapshot_ready"): attempt += 1 try: print(f"[img2img] background download attempt {attempt} of {IMG2IMG_MODEL_ID}…") path = snapshot_download(repo_id=IMG2IMG_MODEL_ID, max_workers=2) _img2img_state["snapshot_path"] = path _img2img_state["snapshot_ready"] = True _img2img_state["build_error"] = None print(f"[img2img] model cached and ready at {path}") return except Exception as exc: # noqa: BLE001 — usually a dropped connection _img2img_state["build_error"] = f"download attempt {attempt} failed: {exc}" wait = min(60, 5 * attempt) print(f"[img2img] download attempt {attempt} failed ({exc}); resuming in {wait}s") time.sleep(wait) threading.Thread(target=_run, daemon=True).start() 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)'})") # FLUX is mounted as a read-only volume (HF Space "model volume") at # this path → the weights are already present, no download. The GPU # worker loads straight from the mount. Falls back to downloading # into the cache when the mount isn't there (local dev / no volume). mount = os.environ.get("SOZAI_IMG2IMG_MOUNT", "/models/flux") if os.path.isdir(mount) and os.listdir(mount): _img2img_state["snapshot_path"] = mount _img2img_state["snapshot_ready"] = True print(f"[img2img] using mounted model at {mount} (no download)") else: print("[img2img] no mount — pre-fetching weights to cache") _img2img_start_bg_download() 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 # The worker must NOT download (it can't write the cache). Require the # main-process background download to have finished, then load READ-ONLY # from the local cache snapshot. if not _img2img_state.get("snapshot_ready"): _img2img_start_bg_download() # ensure it's running raise RuntimeError("model still downloading in the background; try again shortly") model_src = _img2img_state.get("snapshot_path") or IMG2IMG_MODEL_ID lora_path = _img2img_state.get("lora_path") device = "cuda" if torch.cuda.is_available() else "cpu" print(f"[img2img] building from cache {model_src} on {device} " f"({'with LoRA' if lora_path else 'base, no LoRA'}) …") # low_cpu_mem_usage avoids staging the full ~14 GB in CPU RAM (which # can OOM-kill the ZeroGPU worker mid-load). local_files_only forbids # any network write from the worker. pipe = Flux2KleinPipeline.from_pretrained( model_src, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, local_files_only=True ).to(device) if lora_path and os.path.isfile(lora_path): # LoRA load is isolated: if it fails (e.g. peft missing, or an # incompatible adapter) we keep the BASE pipeline working rather # than failing the whole develop. try: 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)") except Exception as lexc: # noqa: BLE001 _img2img_state["lora_loaded"] = False print(f"[img2img] LoRA load failed ({lexc}); using base model") 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 import traceback _img2img_state["build_error"] = f"{type(exc).__name__}: {exc}" print(f"[img2img] in-worker build failed: {exc}") traceback.print_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=180) 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: " + str(_img2img_state.get("build_error") or "unknown")) 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} # --------------------------- Realtime sync (SSE) --------------------------- # _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", } # NOTE: This transport was a raw WebSocket (/ws/{room_id}). Hugging Face # Spaces' router silently drops custom (non-Gradio) WebSocket UPGRADE requests # at the proxy — the handshake never reaches this process, so every connection # loops in "reconnecting". Plain HTTP routes on this same app work fine, so we # mirror what Gradio itself does on Spaces: a one-way Server-Sent Events stream # for server→client messages, plus ordinary POST for client→server messages. # # Each client is identified by its session_id and holds one open SSE GET. The # Manager keys connections by an asyncio.Queue per client instead of a socket; # "sending" to a client just puts a message on its queue, which the SSE # generator drains to the browser. Presence, join-approval, the relay # allow-list, and the DB writes are all unchanged from the WebSocket version. # # SINGLE-REPLICA ASSUMPTION: a POST and the matching SSE stream must land on the # same process. Your Space runs one replica today, so this holds. If you ever # scale to multiple replicas, move the queues behind Redis pub/sub so a POST on # replica A reaches a stream on replica B. # The four fields that are safe to broadcast to other participants. `info` # never contains anything else now, but we filter through this everywhere it # crosses the wire so a stray non-JSON value (e.g. an asyncio.Event) can never # crash the SSE stream's json.dumps. _PUBLIC_FIELDS = ("session_id", "name", "participant_id", "color") def _public(info): """Project a participant's info down to only the JSON-safe public fields.""" return {k: info[k] for k in _PUBLIC_FIELDS} class _Client: """One connected participant: their info + the queue feeding their SSE.""" def __init__(self, info): self.info = info self.queue: asyncio.Queue = asyncio.Queue() self.alive = True async def put(self, message): if self.alive: await self.queue.put(message) class _PendingJoin: """A joiner awaiting host approval. Holds the decision Event SEPARATELY from `info` so `info` stays pure-JSON and can be broadcast safely.""" def __init__(self, info): self.info = info self.event = asyncio.Event() self.admitted = False class Manager: def __init__(self): self.rooms = {} # room_id -> {session_id: _Client} (admitted) self.pending = {} # room_id -> {session_id: info} (awaiting host) # ---- admitted participants -------------------------------------------- # def add_client(self, room_id, client): self.rooms.setdefault(room_id, {})[client.info["session_id"]] = client def get_client(self, room_id, session_id): return self.rooms.get(room_id, {}).get(session_id) def disconnect(self, room_id, session_id): peers = self.rooms.get(room_id) if peers and session_id in peers: peers[session_id].alive = False del peers[session_id] if not peers: self.rooms.pop(room_id, None) self.remove_pending(room_id, session_id) # ---- pending (awaiting host approval) --------------------------------- # def admit(self, room_id, info): """Promote a pending session to a full participant; returns its _Client.""" self.pending.get(room_id, {}).pop(info["session_id"], None) client = _Client(info) self.add_client(room_id, client) return client def add_pending(self, room_id, pending): self.pending.setdefault(room_id, {})[pending.info["session_id"]] = pending def remove_pending(self, room_id, session_id): p = self.pending.get(room_id) if p and session_id in p: del p[session_id] if not p: self.pending.pop(room_id, None) def find_pending(self, room_id, session_id): """Return the _PendingJoin for a session, or None.""" return self.pending.get(room_id, {}).get(session_id) # ---- host resolution (who approves joins) ----------------------------- # def host_client(self, room_id, owner_session_id): """The client 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 and owner_session_id in peers: return peers[owner_session_id] return next(iter(peers.values()), None) # ---- presence --------------------------------------------------------- # def presence(self, room_id): peers = self.rooms.get(room_id, {}) seen, out = set(), [] for client in peers.values(): info = client.info if info["session_id"] in seen: continue seen.add(info["session_id"]) out.append(_public(info)) return out # ---- delivery --------------------------------------------------------- # async def broadcast(self, room_id, message, exclude=None): for sid, client in list(self.rooms.get(room_id, {}).items()): if sid == exclude: continue try: await client.put(message) except Exception: pass async def send_to(self, room_id, session_id, message): client = self.get_client(room_id, session_id) if client is None: return False try: await client.put(message) return True except Exception: return False manager = Manager() # Pending joiners don't have a _Client/queue yet (they're not admitted), but # they still need to receive join_pending / join_declined over their own SSE # stream. We give every SSE connection a queue up front, keyed by session_id, # so the stream exists from the moment the browser connects — even while the # session is still waiting for host approval. _sse_queues = {} # (room_id, session_id) -> asyncio.Queue def _sse_queue(room_id, session_id): key = (room_id, session_id) q = _sse_queues.get(key) if q is None: q = asyncio.Queue() _sse_queues[key] = q return q async def _relay_incoming(room_id, session_id, data): """Apply one client→server event (formerly a WS receive_json). Shared by the POST /send endpoint. `data` already has a 'type'.""" client = manager.get_client(room_id, session_id) info = client.info if client else None t = data.get("type") if t == "rename" and info is not None: 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)}) return # ---- host approves / declines a pending joiner ---- if t in ("admit_join", "decline_join"): target_sid = data.get("session_id") pend = manager.find_pending(room_id, target_sid) if pend is not None: pend.admitted = (t == "admit_join") pend.event.set() return # 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 and info is not None: data["from"] = _public(info) await manager.broadcast(room_id, data, exclude=session_id) # 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) @app.get("/api/rooms/{room_id}/stream") async def room_stream(room_id: str, request: Request): """Server-Sent Events stream: server→client messages for one participant. Replaces the inbound half of the old WebSocket. The browser opens this with EventSource(...). Query params: session_id, name.""" if room_row(room_id) is None: # mimic the old 4404: emit a single error event then end the stream. async def _gone(): yield 'data: {"type":"error","message":"Room not found or expired."}\n\n' return StreamingResponse(_gone(), media_type="text/event-stream") session_id = (request.query_params.get("session_id") or secrets.token_hex(8))[:64] info = { "session_id": session_id, "name": sanitize_name(request.query_params.get("name")), "participant_id": secrets.token_hex(6), "color": color_for(session_id), } q = _sse_queue(room_id, session_id) # ---- join approval (same policy as the WS version) ---- 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 host = manager.host_client(room_id, owner_session) needs_approval = (not is_owner) and (host is not None) if needs_approval: pend = _PendingJoin(info) manager.add_pending(room_id, pend) await q.put({"type": "join_pending", "you": _public(info)}) await host.put({"type": "join_request", "participant": _public(info)}) try: await asyncio.wait_for(pend.event.wait(), timeout=120) except Exception: pass if not pend.admitted: manager.remove_pending(room_id, session_id) await q.put({"type": "join_declined"}) await q.put({"__close__": True}) # fall through to the generator, which will flush join_declined then end client = None else: client = manager.admit(room_id, info) client.queue = q # reuse the queue the SSE stream is already draining else: client = _Client(info) client.queue = q manager.add_client(room_id, client) if client is not None: 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 q.put({"type": "init", "you": _public(info), "participants": manager.presence(room_id)}) await manager.broadcast(room_id, {"type": "user_joined", "participant": _public(info)}, exclude=session_id) await manager.broadcast(room_id, {"type": "presence_update", "participants": manager.presence(room_id)}) async def event_gen(): try: while True: # heartbeat keeps the connection alive through idle proxies try: msg = await asyncio.wait_for(q.get(), timeout=20) except asyncio.TimeoutError: yield ": keep-alive\n\n" continue if isinstance(msg, dict) and msg.get("__close__"): break yield f"data: {json.dumps(msg, separators=(',', ':'))}\n\n" except asyncio.CancelledError: pass finally: _sse_queues.pop((room_id, session_id), None) if manager.get_client(room_id, session_id) is not None: manager.disconnect(room_id, session_id) await manager.broadcast(room_id, {"type": "user_left", "participant": _public(info)}) await manager.broadcast(room_id, {"type": "presence_update", "participants": manager.presence(room_id)}) return StreamingResponse(event_gen(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "X-Accel-Buffering": "no", # disable proxy buffering so events flush live "Connection": "keep-alive", }) @app.post("/api/rooms/{room_id}/send") async def room_send(room_id: str, request: Request): """Client→server messages: one event the client used to send over the socket (object_updated, chat_message, cursor_position, admit_join, ...).""" 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: raise HTTPException(status_code=400, detail="Invalid JSON.") session_id = (body or {}).get("session_id") if not session_id: raise HTTPException(status_code=400, detail="Missing session_id.") data = (body or {}).get("event") or {} if not isinstance(data, dict) or not data.get("type"): raise HTTPException(status_code=400, detail="Missing event.type.") await _relay_incoming(room_id, session_id, data) return {"ok": True} # ------------------------------ 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)