app.py

"""JoyAI-Video-Edit — streaming v2v editing on HF ZeroGPU.

Port of the standalone uvicorn server (xvideo/serving/serve_joyomni_streaming.py)
to ZeroGPU's fork model, without modifying the code it vendors:
  * the main web process holds zero CUDA state — a thin byte pipe between the
    browser WebSocket and two multiprocessing.Queues;
  * each session runs ONE @spaces.GPU generator, acquired through the ladder
    in ws(); the entire server-side machine (gate, PE, streaming session,
    output pump) runs unchanged inside the fork.

Hardware: ZeroGPU, JOYOMNI_GPU_SIZE env picks the tier — xlarge (default, full
RTX Pro 6000 Blackwell, 96GB) or large (half-card MIG, 48GB + JOYOMNI_LOW_VRAM=1);
both are sm_120, matching the joyomni_ops wheel in wheels/.
"""
from __future__ import annotations

import os
import shutil
from pathlib import Path

# ZeroGPU's libstdc++ predates GCC 13; preload ours RTLD_GLOBAL before any native import.
def _preload_libstdcxx() -> None:
    import ctypes

    cand = Path(__file__).resolve().parent / "libs" / "libstdc++.so.6"
    if not cand.is_file():
        return
    try:
        ctypes.CDLL(str(cand), mode=ctypes.RTLD_GLOBAL)
        print(f"[boot] preloaded {cand}", flush=True)
    except OSError as e:
        print(f"[boot] WARNING: could not preload libstdc++: {e}", flush=True)


_preload_libstdcxx()

# Weights persist on /data; compile/JIT caches stay ephemeral (forks never compile).
_DATA = Path("/data") if Path("/data").is_dir() and os.access("/data", os.W_OK) else (Path(__file__).resolve().parent / "deps")
shutil.rmtree(_DATA / "cache", ignore_errors=True)  # legacy compile-cache dir
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("GRADIO_SSR_MODE", "false")

# Kernel path: FP8 + CUDA graph on plain cuDNN SDPA — on the Space GPU (RTX PRO 6000
# Blackwell) cuDNN is the fastest attention at the serving shapes.
os.environ.setdefault("JOYOMNI_FP8_IMG", "1")
os.environ.setdefault("JOYOMNI_FP8_TXT", "1")
os.environ.setdefault("JOYOMNI_CUDA_GRAPH", "1")

# Module-scope load runs under CUDA emulation (no real GPU) — skip the load-time
# warmup; each session fork warms up as it runs. The standalone server leaves this unset.
os.environ.setdefault("JOYOMNI_SKIP_LOAD_WARMUP", "1")
# Do NOT set JOYOMNI_SKIP_VAE_COMPILE_WARMUP here: skipping the load-time VAE
# warmups looks like a free ~6 min boot saving, but ZeroGPU's tensor packing hangs
# without them (the warmup forwards materialize state packing depends on).
# Empirically verified 2026-08-18 over five build cycles.

import sys

DEPLOY_ROOT = Path(__file__).resolve().parent
if str(DEPLOY_ROOT) not in sys.path:
    sys.path.insert(0, str(DEPLOY_ROOT))

import asyncio
import base64
import json
import queue
import re
import tempfile
import threading
import time
import traceback
from multiprocessing import Queue as MPQueue
from types import SimpleNamespace

import warnings

warnings.filterwarnings("ignore", message="ZeroGPU: Cannot get Gradio app Queue instance")

import spaces

CKPT_ROOT = Path(os.environ.get("JOYOMNI_CKPT_ROOT", str(_DATA / "checkpoints")))
# 60s = 120s billed on xlarge.
SESSION_DURATION = int(os.environ.get("JOYOMNI_SPACES_DURATION", "60"))
# xlarge = full RTX Pro 6000 (96GB); large = half-card MIG slice (48GB, needs
# JOYOMNI_LOW_VRAM=1) — the larger pool when the xlarge tier is saturated.
GPU_SIZE = os.environ.get("JOYOMNI_GPU_SIZE", "xlarge")


# Prebuilt sm_120 wheel: pip runs before repo files exist on Spaces — install here.
def _ensure_local_wheels() -> None:
    import importlib.util
    import subprocess

    wheel_dir = DEPLOY_ROOT / "wheels"
    specs = [
        ("joyomni_ops", "joyomni_ops-0.1.0-cp310-cp310-linux_x86_64.whl"),
    ]
    for mod, whl in specs:
        if importlib.util.find_spec(mod) is not None:
            continue  # already installed (warm container)
        path = wheel_dir / whl
        if not path.is_file():
            print(f"[boot] WARNING: wheel missing, cannot install {mod}: {path}", flush=True)
            continue
        print(f"[boot] installing {mod} from {whl} ...", flush=True)
        subprocess.check_call(
            [sys.executable, "-m", "pip", "install", "--no-deps", "--no-index", str(path)]
        )


_ensure_local_wheels()


# Pull weights to /data on first boot; later boots skip. ONNX detectors optional.
def _resolve_weights() -> dict:
    CKPT_ROOT.mkdir(parents=True, exist_ok=True)
    jve_local = CKPT_ROOT / "JoyAI-Video-Edit"
    mimo_local = CKPT_ROOT / "MiMo-VL-7B-RL-2508"

    # Skip snapshot_download when the key files are already on /data (~30s less
    # cold-start). Either DIT (0811 preferred, 0804 fallback) satisfies readiness.
    _dit_0811 = jve_local / "dit" / "joyai_video_edit_dit_0811.pth"
    _dit_0804 = jve_local / "dit" / "joyai_video_edit_dit_0804.pth"
    jve_ready = (_dit_0811.is_file() or _dit_0804.is_file()) and \
                (jve_local / "vae" / "diffusion_pytorch_model.safetensors").is_file()
    mimo_ready = (mimo_local / "config.json").is_file() and \
                 any(mimo_local.glob("*.safetensors"))

    if jve_ready and mimo_ready:
        print("[boot] weights already in /data — skipping hub check entirely", flush=True)
    else:
        from huggingface_hub import snapshot_download
        snapshot_download("jdopensource/JoyAI-Video-Edit", local_dir=str(jve_local))
        snapshot_download("XiaomiMiMo/MiMo-VL-7B-RL-2508", local_dir=str(mimo_local))

    # Ensure 0811 specifically via a targeted single-file fetch; keep 0804 if absent.
    if not _dit_0811.is_file():
        try:
            from huggingface_hub import hf_hub_download
            hf_hub_download("jdopensource/JoyAI-Video-Edit",
                            "dit/joyai_video_edit_dit_0811.pth", local_dir=str(jve_local))
        except Exception as e:  # noqa: BLE001
            print(f"[boot] 0811 DIT not fetched ({type(e).__name__}: {e}); using 0804", flush=True)
    _dit = _dit_0811 if _dit_0811.is_file() else _dit_0804
    return {
        "dit": str(_dit),
        "vae": str(jve_local / "vae"),
        "text_encoder": str(mimo_local),
        "face_onnx": str(DEPLOY_ROOT / "detectors" / "face_detection_yunet_2023mar.onnx"),
        "person_onnx": str(DEPLOY_ROOT / "detectors" / "yolov8n.onnx"),
    }


# ZeroGPU requires models on cuda at module scope (packed once, mapped into forks).
_RUNTIME = None
_PATHS: dict | None = None
_BUILD_ERROR: str | None = None


def _build_runtime():
    from xvideo.inductor_autotune_fix import install as _install_autotune_fix
    from xvideo.serving.joyomni_streaming import JoyOmniRuntime

    _install_autotune_fix()
    paths = _resolve_weights()
    rt = JoyOmniRuntime.load(
        paths["dit"],
        vae_ckpt=paths["vae"],
        text_encoder_ckpt=paths["text_encoder"],
        device="cuda",
        vae_encode_device="cuda",
        vae_decode_device="cuda",
        vae_pseudo_device="cuda",
        postprocess_device="cuda",
        seed=42,
        warmup_height=480,
        warmup_width=840,
    )
    return rt, paths


try:
    print("[boot] building JoyAI-Video-Edit runtime (first boot downloads ~63GB to /data)...", flush=True)
    _t0 = time.perf_counter()
    _RUNTIME, _PATHS = _build_runtime()
    print(f"[boot] runtime ready in {time.perf_counter() - _t0:.1f}s", flush=True)
except Exception as e:  # noqa: BLE001
    _BUILD_ERROR = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
    print("[boot] runtime build FAILED:\n" + _BUILD_ERROR, flush=True)


# GPU session: full pipeline inside the fork; WS replaced by
#   in_q {"kind":"text"|"bytes"|"close","data":...} / out_q ("json",dict)|("bin",bytes)
def make_gpu_session(in_q, out_q, paths, duration, dead):
    """Build the @spaces.GPU generator for one session. in_q/out_q captured by
    CLOSURE — inherited through the fork (MPQueues are not picklable, so they
    cannot be decorated-function args)."""

    sched_t0 = time.monotonic()  # ZeroGPU's duration window opens here — queue wait included

    @spaces.GPU(duration=duration, size=GPU_SIZE)
    def gpu_session():
        # Queue slots can't be cancelled: if the client left, release without touching the engine.
        if dead.is_set():
            print("[ws] fork: client left during queue wait — GPU released", flush=True)
            return
        yield from _run_session_in_fork(in_q, out_q, paths, duration, sched_t0)

    return gpu_session


def _run_session_in_fork(in_q, out_q, paths, duration, sched_t0):
    """Inside the GPU fork: drive the vendored streaming engine on a worker thread
    and yield a sentinel per tick so the @spaces.GPU generator keeps the GPU lease
    alive for the session's lifetime."""
    from xvideo.serving.zerogpu_engine import run_session_blocking

    done = threading.Event()

    def _drive():
        try:
            run_session_blocking(_RUNTIME, paths, in_q, out_q)  # reports its own errors
        finally:
            done.set()

    threading.Thread(target=_drive, daemon=True).start()
    while not done.wait(timeout=1.0):
        left = duration - (time.monotonic() - sched_t0)
        try:
            out_q.put_nowait(("json", {"type": "lease_left", "left_s": max(0, round(left))}))
        except Exception:  # noqa: BLE001
            pass
        yield "tick"
    yield "done"


@spaces.GPU(duration=30, size=GPU_SIZE)
def gpu_probe():
    """Hardware introspection + micro-benchmarks inside the fork. Answers, for a
    few GPU-seconds: which physical GPU ZeroGPU attached, at what clocks/TGP, and
    the achievable bf16 / fp8 GEMM throughput + HBM bandwidth. Used to compare a
    dev box against this Space apples-to-apples."""
    import subprocess
    import time as _t

    import torch

    out = {"device": torch.cuda.get_device_name(0),
           "capability": list(torch.cuda.get_device_capability(0)),
           "torch": torch.__version__}
    try:
        out["cpu_count"] = os.cpu_count()
        out["cpu_affinity"] = len(os.sched_getaffinity(0))
        out["loadavg"] = [round(x, 1) for x in os.getloadavg()]
        model = ""
        for line in open("/proc/cpuinfo"):
            if line.startswith("model name"):
                model = line.split(":", 1)[1].strip()
                break
        out["cpu_model"] = model
        t0 = _t.perf_counter()
        x = 0
        for i in range(3_000_000):
            x += i * i
        out["cpu_1t_spin_ms"] = round((_t.perf_counter() - t0) * 1000, 1)
    except Exception as e:  # noqa: BLE001
        out["cpu_err"] = repr(e)
    try:
        free, total = torch.cuda.mem_get_info()
        out["mem_total_gb"] = round(total / 2**30, 1)
        out["mem_free_gb"] = round(free / 2**30, 1)
    except Exception as e:  # noqa: BLE001
        out["mem_err"] = repr(e)
    try:
        smi = subprocess.run(
            ["nvidia-smi", "--query-gpu=name,clocks.sm,clocks.max.sm,power.limit,power.draw,temperature.gpu",
             "--format=csv,noheader"], capture_output=True, text=True, timeout=10)
        out["smi"] = smi.stdout.strip() or smi.stderr.strip()
    except Exception as e:  # noqa: BLE001
        out["smi"] = repr(e)

    dev = torch.device("cuda:0")
    a = torch.randn(8192, 8192, device=dev, dtype=torch.bfloat16)
    b = torch.randn(8192, 8192, device=dev, dtype=torch.bfloat16)
    for _ in range(3):
        a @ b
    torch.cuda.synchronize()
    t0 = _t.perf_counter()
    n = 20
    for _ in range(n):
        a @ b
    torch.cuda.synchronize()
    dt = _t.perf_counter() - t0
    out["bf16_gemm_tflops"] = round(n * 2 * 8192**3 / dt / 1e12, 1)

    try:
        af = (a.float() / a.float().abs().amax()).to(torch.float8_e4m3fn)
        bf = (b.float() / b.float().abs().amax()).t().contiguous().t().to(torch.float8_e4m3fn)
        sa = torch.ones(1, device=dev)
        sb = torch.ones(1, device=dev)
        for _ in range(3):
            torch._scaled_mm(af, bf, scale_a=sa, scale_b=sb, out_dtype=torch.bfloat16)
        torch.cuda.synchronize()
        t0 = _t.perf_counter()
        for _ in range(n):
            torch._scaled_mm(af, bf, scale_a=sa, scale_b=sb, out_dtype=torch.bfloat16)
        torch.cuda.synchronize()
        dt = _t.perf_counter() - t0
        out["fp8_scaled_mm_tflops"] = round(n * 2 * 8192**3 / dt / 1e12, 1)
    except Exception as e:  # noqa: BLE001
        out["fp8_err"] = repr(e)

    try:
        big = torch.empty(2**30, device=dev, dtype=torch.uint8)
        dst = torch.empty_like(big)
        for _ in range(2):
            dst.copy_(big)
        torch.cuda.synchronize()
        t0 = _t.perf_counter()
        for _ in range(8):
            dst.copy_(big)
        torch.cuda.synchronize()
        dt = _t.perf_counter() - t0
        out["hbm_copy_gbps"] = round(8 * 2 * big.numel() / dt / 1e9, 0)
    except Exception as e:  # noqa: BLE001
        out["bw_err"] = repr(e)

    try:
        from xvideo.models.vae import vae_compile as _vc
        vae_d = _RUNTIME.decode_vae
        vae_e = _RUNTIME.pipeline.vae
        lat_c = int(getattr(vae_d, "latent_channels", 16) or 16)
        z = torch.zeros(1, lat_c, 2, 60, 105, device=dev, dtype=torch.bfloat16)
        z = _vc.prep_input(z)
        with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
            for _ in range(3):
                vae_d.decode(z, return_dict=False)[0]
            torch.cuda.synchronize()
            t0 = _t.perf_counter()
            for _ in range(15):
                vae_d.decode(z, return_dict=False)[0]
            torch.cuda.synchronize()
        out["vae_decode_direct_ms"] = round((_t.perf_counter() - t0) / 15 * 1000, 2)
        x = torch.zeros(1, 3, 9, 480, 840, device=dev, dtype=torch.bfloat16)
        x = _vc.prep_input(x)
        with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
            for _ in range(3):
                vae_e.encode(x)
            torch.cuda.synchronize()
            t0 = _t.perf_counter()
            for _ in range(10):
                vae_e.encode(x)
            torch.cuda.synchronize()
        out["vae_encode_direct_ms"] = round((_t.perf_counter() - t0) / 10 * 1000, 2)
    except Exception as e:  # noqa: BLE001
        out["vae_probe_err"] = repr(e)

    try:
        import torch.nn.functional as F
        torch.backends.cudnn.benchmark = True
        for tag, (shape, cout) in {"conv_a": ((1, 384, 2, 60, 105), 384),
                                   "conv_b": ((1, 128, 8, 120, 210), 128)}.items():
            x = torch.randn(*shape, device=dev, dtype=torch.bfloat16).to(memory_format=torch.channels_last_3d)
            w = torch.randn(cout, shape[1], 3, 3, 3, device=dev, dtype=torch.bfloat16).to(memory_format=torch.channels_last_3d)
            for _ in range(8):
                F.conv3d(x, w, padding=1)
            torch.cuda.synchronize()
            t0 = _t.perf_counter()
            for _ in range(30):
                F.conv3d(x, w, padding=1)
            torch.cuda.synchronize()
            out[tag + "_ms"] = round((_t.perf_counter() - t0) / 30 * 1000, 3)
    except Exception as e:  # noqa: BLE001
        out["conv_err"] = repr(e)

    try:
        tiny = torch.ones(16, device=dev)
        for _ in range(200):
            tiny.add_(1)
        torch.cuda.synchronize()
        n2 = 2000
        t0 = _t.perf_counter()
        for _ in range(n2):
            tiny.add_(1)
        torch.cuda.synchronize()
        out["launch_us_per_kernel"] = round((_t.perf_counter() - t0) / n2 * 1e6, 2)
        g = torch.cuda.CUDAGraph()
        s = torch.cuda.Stream()
        s.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(s):
            for _ in range(3):
                tiny.add_(1)
        torch.cuda.current_stream().wait_stream(s)
        with torch.cuda.graph(g):
            for _ in range(500):
                tiny.add_(1)
        for _ in range(3):
            g.replay()
        torch.cuda.synchronize()
        t0 = _t.perf_counter()
        for _ in range(20):
            g.replay()
        torch.cuda.synchronize()
        out["graph_replay_us_per_kernel"] = round((_t.perf_counter() - t0) / (20 * 500) * 1e6, 3)
    except Exception as e:  # noqa: BLE001
        out["launch_err"] = repr(e)
    return out


# Main process: gr.Server (FastAPI) + thin byte-pipe WebSocket.
from gradio import Server  # noqa: E402
from fastapi import Request, WebSocket, WebSocketDisconnect  # noqa: E402
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, Response  # noqa: E402
from starlette.datastructures import Headers  # noqa: E402

# Mirrors zerogpu_engine._default_args: sessions record here (owner's decision).
_RECORD_DIR = os.environ.get("JOYOMNI_RECORD_DIR",
                             "/data/recordings" if Path("/data").is_dir() else "") or None

_INDEX_HTML = (DEPLOY_ROOT / "static" / "index.html").read_text(encoding="utf-8")

# JOYOMNI_MAINTENANCE=1 serves this instead of the live demo (delete the variable to restore).
_MAINTENANCE_HTML = """<!doctype html><html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>JoyAI Video Edit</title><style>
:root{--bg:#f0efe9;--panel:#fff;--line:#e6e4dd;--text:#1c1c1c;--muted:#8a8a86;--accent:#f2643f;}
*{box-sizing:border-box}body{margin:0;background:var(--bg);color:var(--text);
font-family:Inter,-apple-system,sans-serif}
.wrap{max-width:1080px;margin:0 auto;padding:48px 24px}
.notice{background:var(--panel);border:1px solid var(--line);border-radius:16px;padding:36px 28px;text-align:center}
.notice .emoji{font-size:44px}
.notice h1{font-size:20px;margin:14px 0 6px}
.notice p{color:var(--muted);font-size:14px;margin:4px 0}
.team{margin-top:34px}
.team .cap{font-size:13px;color:var(--muted);font-weight:600;margin-bottom:12px;text-align:center}
.team .cap b{color:var(--accent)}
.team img{width:100%;border-radius:12px;border:1px solid var(--line);cursor:zoom-in;background:#0d0d0d;display:block}
.team img:hover{border-color:var(--accent)}
.grid{display:grid;grid-template-columns:repeat(2,1fr);gap:14px;margin-top:14px}
@media(max-width:640px){.grid{grid-template-columns:1fr}}
.lb{display:none;position:fixed;inset:0;z-index:10;background:rgba(0,0,0,.85);align-items:center;justify-content:center;cursor:zoom-out;padding:20px}
.lb.show{display:flex}
.lb img{max-width:94vw;max-height:94vh;object-fit:contain;border-radius:10px}
</style></head><body><div class="wrap">
<div class="notice"><div class="emoji">\U0001f6e0️</div>
<h1>Down for maintenance — back soon</h1>
<p>The live demo is temporarily offline while we work on it.</p></div>
<div class="team"><div class="cap">Built by the JoyAI Video Edit team · runs natively at <b>1248×720</b> · prompt enhancement by <b>Claude Fable 5</b></div>
<div id="hero"></div>
<div class="grid" id="grid"></div></div></div>
<div class="lb" id="lb"><img id="lbImg" alt=""></div>
<script>
const lb=document.getElementById("lb"),lbImg=document.getElementById("lbImg");
lb.onclick=()=>lb.classList.remove("show");
document.addEventListener("keydown",e=>{if(e.key==="Escape")lb.classList.remove("show");});
fetch("team.json").then(r=>r.ok?r.json():[]).then(items=>{
  const hero=document.getElementById("hero"),g=document.getElementById("grid");
  items.forEach((it,i)=>{
    const im=document.createElement("img");
    im.src=i===0?it.full:it.thumb;im.alt=it.name;im.loading="lazy";
    im.onclick=()=>{lbImg.src=it.full;lb.classList.add("show");};
    (i===0?hero:g).appendChild(im);
  });
}).catch(()=>{});
</script></body></html>"""

app = Server(title="JoyAI-Video-Edit (ZeroGPU)")


@app.get("/", response_class=HTMLResponse)
async def index():
    if os.environ.get("JOYOMNI_MAINTENANCE"):
        return HTMLResponse(_MAINTENANCE_HTML)
    # __SERVER_DEFAULTS__ filled from serve's argparse defaults.
    from xvideo.serving.serve_joyomni_streaming import build_parser
    _sd = build_parser().parse_args([])
    defaults = {k: getattr(_sd, k) for k in (
        "width", "height", "num_inference_steps", "output_quality", "seed", "fps",
        "online_gate", "kv_reset_frames", "static_diff_thresh", "freeze_kv_on_static",
        "profile_timings", "max_temporal_ids")}
    defaults.update({"use_pe": bool(os.environ.get("OPENAI_API_KEY")),
                     "pe_available": bool(os.environ.get("OPENAI_API_KEY")),
                     "record_enabled": _RECORD_DIR is not None})
    return _INDEX_HTML.replace("__SERVER_DEFAULTS__", json.dumps(defaults))


@app.get("/ref-images")
async def ref_images():
    # Reference presets: reuse serve's loader (reads rv2v_reference/, returns name->data-url).
    from xvideo.serving.serve_joyomni_streaming import _ref_images_cached
    return _ref_images_cached()


# Team gallery (Space-only): the client shows its strip when /team.json exists.
_TEAM_DIR = DEPLOY_ROOT / "teams"
_TEAM_ORDER = ("us", "wenxun", "wenxun-x", "xinran", "yicheng")


@app.get("/team.json")
async def team_manifest():
    items = [{"name": n, "thumb": f"/team/thumb/{n}", "full": f"/team/full/{n}"}
             for n in _TEAM_ORDER
             if (_TEAM_DIR / f"{n}.png").is_file() and (_TEAM_DIR / "thumb" / f"{n}.webp").is_file()]
    return JSONResponse(items, status_code=200 if items else 404)


@app.get("/team/{kind}/{name}")
async def team_image(kind: str, name: str):
    if name not in _TEAM_ORDER or kind not in ("thumb", "full"):
        return JSONResponse({"error": "not found"}, status_code=404)
    path = _TEAM_DIR / "thumb" / f"{name}.webp" if kind == "thumb" else _TEAM_DIR / f"{name}.png"
    if not path.is_file():
        return JSONResponse({"error": "not found"}, status_code=404)
    return FileResponse(str(path), media_type="image/webp" if kind == "thumb" else "image/png",
                        headers={"Cache-Control": "public, max-age=86400"})


@app.get("/health")
async def health():
    return {"ok": _BUILD_ERROR is None, "build_error": _BUILD_ERROR}


@app.get("/download_last")
async def download_last(rec: str = ""):
    if _RECORD_DIR is None:
        return JSONResponse({"error": "Recording is not enabled."}, status_code=404)
    if not re.fullmatch(r"\d+_\d+", rec):
        return JSONResponse({"error": "Missing or invalid rec id."}, status_code=400)
    base = Path(_RECORD_DIR) / rec
    if not base.is_dir():
        return JSONResponse({"error": "No such recording."}, status_code=404)
    segments = sorted(base.glob("output_*.mp4"))
    if not segments:
        return JSONResponse({"error": "Recording file has not been generated yet. Try again later."}, status_code=404)
    download_name = f"joyomni_{base.name}.mp4"
    if len(segments) == 1:
        return FileResponse(str(segments[0]), media_type="video/mp4", filename=download_name)
    try:
        import imageio_ffmpeg
        ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe()
    except Exception as exc:  # noqa: BLE001
        return JSONResponse({"error": f"ffmpeg unavailable: {exc!r}"}, status_code=500)
    list_path = out_path = None
    try:
        fd_list, list_path = tempfile.mkstemp(suffix=".txt", prefix="rv2v_cat_")
        with os.fdopen(fd_list, "w", encoding="utf-8") as f:
            for seg in segments:
                f.write(f"file '{seg.as_posix()}'\n")
        fd_out, out_path = tempfile.mkstemp(suffix=".mp4", prefix="rv2v_dl_")
        os.close(fd_out)
        proc = await asyncio.create_subprocess_exec(
            ffmpeg_exe, "-y", "-f", "concat", "-safe", "0", "-i", list_path,
            "-c", "copy", "-movflags", "+faststart", out_path,
            stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.PIPE,
        )
        _, stderr = await proc.communicate()
        if proc.returncode != 0:
            tail = (stderr or b"").decode("utf-8", "replace")
            return JSONResponse({"error": f"ffmpeg failed: {tail}"}, status_code=500)
        with open(out_path, "rb") as f:
            data = f.read()
        return Response(content=data, media_type="video/mp4",
                        headers={"Content-Disposition": f'attachment; filename="{download_name}"'})
    except Exception as exc:  # noqa: BLE001
        return JSONResponse({"error": f"download encode error: {exc!r}"}, status_code=500)
    finally:
        for p in (list_path, out_path):
            if p:
                try:
                    os.unlink(p)
                except OSError:
                    pass


# Debug endpoints burn quota: gated by the JOYOMNI_DEBUG_KEY secret (?key=...); no secret = disabled.
_DEBUG_KEY = os.environ.get("JOYOMNI_DEBUG_KEY", "")


def _debug_denied(request: Request):
    if not _DEBUG_KEY or request.query_params.get("key") != _DEBUG_KEY:
        return JSONResponse({"error": "forbidden"}, status_code=403)
    return None


@spaces.GPU(duration=120, size=GPU_SIZE)
def engine_bench(fresh_cache: bool = False, pace: float = 12.0, chunks: int = 16):
    """Engine-only throughput probe inside the fork: no session wrapper, no
    queues — the exact mirror of the dev-box bench. Separates 'fork environment'
    from 'session machinery' when chasing Space-vs-local speed gaps. Samples SM
    clocks while running to expose power-management effects."""
    import subprocess as _sp
    import threading as _th
    import time as _t

    import numpy as np
    from PIL import Image as _Image

    if fresh_cache:
        import torch._inductor.config  # noqa: F401
        os.environ["TORCHINDUCTOR_FORCE_DISABLE_CACHES"] = "1"

    from xvideo.serving.joyomni_streaming import StreamingSettings

    clocks = []
    stop_clk = _th.Event()

    def _clk():
        while not stop_clk.is_set():
            try:
                r = _sp.run(["nvidia-smi", "--query-gpu=clocks.sm,power.draw",
                             "--format=csv,noheader,nounits"],
                            capture_output=True, text=True, timeout=5)
                clocks.append(r.stdout.strip().split("\n")[0])
            except Exception:  # noqa: BLE001
                pass
            stop_clk.wait(0.5)

    _th.Thread(target=_clk, daemon=True).start()

    settings = StreamingSettings(height=480, width=840, num_inference_steps=2,
                                 seed=42, max_temporal_ids=8, profile_timings=True,
                                 output_codec="mjpeg")
    sess = _RUNTIME.create_v2v_session(
        "Turn the scene into an oil painting with warm colors.", settings=settings)
    rng = np.random.default_rng(7)
    base = rng.integers(0, 256, size=(480, 840, 3), dtype=np.uint8)
    recs = []
    t0 = _t.time()
    i = 0
    try:
        while len(recs) < chunks and i < 400 and _t.time() - t0 < 95:
            if pace > 0:
                target = t0 + i / pace
                now = _t.time()
                if target > now:
                    _t.sleep(target - now)
            frame = _Image.fromarray(np.roll(base, (i * 7) % 840, axis=1), mode="RGB")
            for r in sess.push_frame(frame, {"seq": i + 1, "t_capture_ms": 0.0}):
                p = r.profile or {}
                recs.append({k: round(float(p[k]), 4) for k in
                             ("dit_denoise_s", "vae_encode_s", "reference_prepare_s",
                              "vae_decode_s", "jpeg_encode_s", "frames_to_tensor_s",
                              "total_server_chunk_s") if p.get(k) is not None})
            i += 1
        deadline = _t.time() + 20.0
        while len(recs) < chunks and _t.time() < deadline:
            r = sess.wait_async_result(timeout=0.5)
            if r is not None:
                p = r.profile or {}
                recs.append({k: round(float(p[k]), 4) for k in
                             ("dit_denoise_s", "vae_encode_s", "reference_prepare_s",
                              "vae_decode_s", "jpeg_encode_s", "frames_to_tensor_s",
                              "total_server_chunk_s") if p.get(k) is not None})
    finally:
        stop_clk.set()
        try:
            sess.close()
        except Exception:  # noqa: BLE001
            pass
    st = recs[5:] or recs
    def _m(k):
        v = [c[k] for c in st if k in c]
        return round(sum(v) / len(v), 4) if v else None
    return {"chunks": len(recs), "fresh_cache": bool(fresh_cache), "pace": pace,
            "dit": _m("dit_denoise_s"), "enc": _m("vae_encode_s"),
            "ref": _m("reference_prepare_s"), "dec": _m("vae_decode_s"),
            "jpg": _m("jpeg_encode_s"), "f2t": _m("frames_to_tensor_s"),
            "total": _m("total_server_chunk_s"), "recs": recs,
            "clocks": clocks[-60:]}


@app.get("/gpu-probe")
async def do_gpu_probe(request: Request):
    denied = _debug_denied(request)
    if denied is not None:
        return denied
    loop = asyncio.get_event_loop()
    try:
        return {"result": await loop.run_in_executor(None, gpu_probe)}
    except Exception as e:  # noqa: BLE001
        return JSONResponse({"error": str(e)}, status_code=500)


@app.get("/engine-bench")
async def do_engine_bench(request: Request, fresh_cache: int = 0, pace: float = 12.0, chunks: int = 16):
    denied = _debug_denied(request)
    if denied is not None:
        return denied
    loop = asyncio.get_event_loop()
    try:
        return {"result": await loop.run_in_executor(
            None, lambda: engine_bench(bool(fresh_cache), float(pace), int(chunks)))}
    except Exception as e:  # noqa: BLE001
        return JSONResponse({"error": str(e)}, status_code=500)


def _token_user(tok: str) -> bool:
    try:
        pl = tok.split(".")[1]
        claims = json.loads(base64.urlsafe_b64decode(pl + "=" * (-len(pl) % 4)))
        return claims.get("user") is not None
    except Exception:  # noqa: BLE001
        return False


# sid -> (headers, expiry). When embedded on huggingface.co, the parent page mints
# user-scoped quota headers via the "zerogpu-headers" postMessage handshake (same
# protocol gradio's frontend speaks). The frontend relays them here over a plain GET
# (WS upgrades can't carry custom headers) and /ws redeems the one-time ticket, so
# the visitor tier bills the logged-in user instead of the anonymous IP pool.
_MINTED: dict[str, tuple[dict, float]] = {}


@app.get("/mint/{sid}")
async def mint(sid: str, request: Request):
    picked = {k.lower(): v for k, v in request.headers.items()
              if k.lower() == "x-ip-token" or k.lower().startswith("x-zerogpu")}
    user = _token_user(picked.get("x-ip-token", ""))
    if picked:
        now = time.time()
        for k in [k for k, (_, exp) in _MINTED.items() if exp < now]:
            _MINTED.pop(k, None)
        if len(_MINTED) < 512:
            _MINTED[sid[:64]] = (picked, now + 120.0)
    print(f"[mint] stashed={sorted(picked)} user={'yes' if user else 'no'}", flush=True)
    return {"stashed": bool(picked), "user": user, "names": sorted(picked)}


@app.websocket("/ws")
async def ws(websocket: WebSocket):
    await websocket.accept()
    # Maintenance covers the GPU path too: with the app-pool rungs in the
    # ladder, a stale tab or a raw WS client could otherwise still bill the
    # owner while the front page says "closed".
    if os.environ.get("JOYOMNI_MAINTENANCE"):
        await websocket.send_json({"type": "error", "code": "maintenance",
                                   "message": "under maintenance"})
        await websocket.close()
        return
    if _BUILD_ERROR is not None:
        await websocket.send_json({"type": "error", "message": "model failed to load"})
        await websocket.close()
        return

    # 512 slots never fill (frame = 2 slots, frontend caps 32 unacked); a dropped
    # frame is never acked and would permanently strangle the credit gate.
    in_q: "MPQueue" = MPQueue(maxsize=512)
    out_q: "MPQueue" = MPQueue(maxsize=64)
    stop = threading.Event()
    gone = threading.Event()

    _hdrs = websocket.headers
    _src = "header"
    _sid = websocket.query_params.get("mint")
    if _sid and (_m := _MINTED.pop(_sid, None)) is not None and _m[1] > time.time():
        _hdrs = Headers({**dict(websocket.headers), **_m[0]})
        _src = "minted"
    _tok = _hdrs.get("x-ip-token")
    _user = bool(_tok) and _token_user(_tok)
    print(f"[ws] x-ip-token: {'yes' if _tok else 'no'} user={'yes' if _user else 'no'} src={_src}", flush=True)

    def _worker():
        from gradio.context import LocalContext

        # Acquisition ladder: visitor quota first (user token or anonymous IP
        # pool), then the shared app pool, each at SESSION_DURATION/30s.
        # Quota refusals are instant (checked at schedule time), so a step down
        # costs one round-trip; "No GPU was available" means the physical pool
        # is busy — no rung or pool fixes that, so the ladder aborts there.
        # Without an x-ip-token the "visitor" identity would silently bill the
        # app anyway (spaces falls back to the app token), so skip those rungs.
        _billx = 2 if GPU_SIZE == "xlarge" else 1
        tiers = [(pool, req, d)
                 for pool, req in (("visitor", SimpleNamespace(headers=_hdrs)),
                                   ("app", None))
                 if pool != "visitor" or _tok
                 for d in dict.fromkeys((SESSION_DURATION, 30))]

        got_gpu = False
        granted = None

        def _put(payload):
            try:
                out_q.put(("json", payload))
            except Exception:  # noqa: BLE001
                pass

        def _err(code, msg):
            payload = {"type": "error", "code": code, "message": msg}
            m = re.search(r"(?:re)?try\s+again\s+in\s+([0-9:]+)", msg, re.I)
            if m:
                payload["retry_in"] = m.group(1)
            _put(payload)

        last = ""
        try:
            for pool, req, dur in tiers:
                if stop.is_set() or gone.is_set():
                    return
                _put({"type": "gpu_wait", "pool": pool, "seconds": dur,
                      "gpu_size": GPU_SIZE, "billing_x": _billx})
                try:
                    LocalContext.request.set(req)
                    gen = make_gpu_session(in_q, out_q, _PATHS, dur, gone)
                    for _ in gen():
                        if not got_gpu:
                            got_gpu = True
                            granted = (pool, dur)
                            print(f"[ws] lease acquired: {pool}@{dur}s "
                                  f"(billed x{_billx})", flush=True)
                            _put({"type": "gpu_ready", "pool": pool, "seconds": dur,
                                  "gpu_size": GPU_SIZE, "billing_x": _billx})
                        if stop.is_set():
                            break
                    return  # session ended inside its lease (client stop/close)
                except Exception as exc:  # noqa: BLE001
                    msg = (str(exc) or exc.__class__.__name__).strip()
                    low = msg.lower()
                    if got_gpu:
                        # Ran, then died. Lease expiry surfaces as "GPU task
                        # aborted" — a normal end, not an error.
                        print("[ws] session ended:\n" + traceback.format_exc(), flush=True)
                        if "aborted" in low:
                            _put({"type": "lease_expired", "pool": granted[0],
                                  "seconds": granted[1], "billing_x": _billx})
                        else:
                            _err("gpu_error", msg)
                        return
                    print(f"[ws] {pool}@{dur}s failed: {msg!r}", flush=True)
                    last = msg
                    if "no gpu was available" in low:
                        _err("gpu_busy", msg)
                        return
                    if not any(k in low for k in ("quota", "exceeded", "limit", "credits")):
                        _err("gpu_error", msg)
                        return
                    # quota refusal -> next rung
            _err("gpu_quota", last or "quota exceeded")
        finally:
            stop.set()

    threading.Thread(target=_worker, daemon=True).start()
    loop = asyncio.get_event_loop()

    async def send_loop():
        while not stop.is_set() or not _q_empty(out_q):
            try:
                kind, payload = await loop.run_in_executor(None, lambda: out_q.get(timeout=0.1))
            except queue.Empty:
                continue
            except (OSError, ValueError):
                break
            try:
                if kind == "json":
                    await websocket.send_json(payload)
                else:
                    await websocket.send_bytes(payload)
            except (WebSocketDisconnect, RuntimeError):
                break
        try:
            await websocket.send_json({"type": "session_timeout",
                                       "message": "GPU lease ended; reconnect to re-queue"})
        except Exception:  # noqa: BLE001
            pass

    async def recv_loop():
        while not stop.is_set():
            try:
                msg = await websocket.receive()
            except (WebSocketDisconnect, RuntimeError):
                break
            if msg.get("type") == "websocket.disconnect":
                break
            if msg.get("text") is not None:
                # Answer pings here in the main process: during the GPU-queue wait
                # nothing drains in_q, and a queued ping would read as fake 10s+ RTT.
                if '"ping"' in msg["text"]:
                    try:
                        p = json.loads(msg["text"])
                        if p.get("type") == "ping":
                            await websocket.send_json({"type": "pong", "t": p.get("t")})
                            continue
                    except (ValueError, KeyError):
                        pass
                _q_put(in_q, {"kind": "text", "data": msg["text"]})
            elif msg.get("bytes") is not None:
                _q_put(in_q, {"kind": "bytes", "data": msg["bytes"]})
        gone.set()
        _q_put(in_q, {"kind": "close"})

    try:
        await asyncio.gather(send_loop(), recv_loop())
    except (WebSocketDisconnect, RuntimeError):
        pass
    finally:
        stop.set()
        gone.set()
        _q_put(in_q, {"kind": "close"})


def _q_put(q, item) -> None:
    try:
        q.put_nowait(item)
    except queue.Full:
        pass  # inbound frames are droppable under backpressure


def _q_empty(q) -> bool:
    try:
        return q.empty()
    except (OSError, ValueError, NotImplementedError):
        return True


# Satisfy ZeroGPU's startup scan (no Gradio Blocks events to walk with gr.Server).
# NOTE: do NOT create/mount any gr.Blocks here — its mere existence breaks the
# FastAPI-demo startup path (the spaces launch hook never fires -> the platform
# reports "No @spaces.GPU function detected"). Verified over four build cycles.
spaces.GPU(lambda: None)

demo = app

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)