Blame for app.py

"""TTB-DVSM HF Spaces (ZeroGPU) demo.

Upload images / a short orbit video -> π³ pose estimation + TTB-DVSM
reconstruction (one GPU call) -> a *time-boxed* interactive viser viewer that
holds the GPU for the duration of one `@spaces.GPU` session.

Why this shape (ZeroGPU constraints, see README):
  - ZeroGPU forks a fresh worker process for each `@spaces.GPU` call and
    releases the GPU when the call returns; GPUs cannot be held permanently.
  - A generator decorated with `@spaces.GPU` keeps the GPU for as long as it is
    being iterated. We exploit that: `_live_session_worker` is a generator that
    starts a viser server *inside the GPU worker* and stays alive (yielding
    heartbeats) until its `duration` budget runs out. The main process reverse-
    proxies `/viser-proxy/*` to that worker-local viser port so the viewer is
    reachable from the single public Gradio port.
  - When the session expires the worker dies (GPU released). The user clicks
    "Start / Continue" to fork a new worker and rebuild the viewer from the
    already-computed artifact on disk (no re-reconstruction).

Models (nvidia/DVSM): uploads are reconstructed with DVSM-TTB-p8-general-960p; the
gallery examples were baked with DVSM-TTB-p8-DL3DV-960p. A live session renders an
artifact with the model that baked it.
"""
from __future__ import annotations

import os
import shutil
import socket
import sys
import tempfile
import threading
import time
import uuid
from pathlib import Path

import numpy as np

PROJECT_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(PROJECT_ROOT))

# ZeroGPU. `spaces` is a no-op-ish shim off-platform; provide a fallback so the
# module is importable for local testing without the package.
try:
    import spaces
except Exception:  # pragma: no cover - local dev without `spaces`
    class _SpacesShim:
        def GPU(self, *args, **kwargs):
            # Support both @spaces.GPU and @spaces.GPU(duration=...)
            if args and callable(args[0]):
                return args[0]
            def _wrap(fn):
                return fn
            return _wrap
    spaces = _SpacesShim()

import gradio as gr

from viser_app.scene_loaders import (
    IMAGE_EXTS, VIDEO_EXTS, extract_video_frames,
)
from viser_app.zerogpu_compat import repair_torch_cuda_state

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
# Uploaded images / videos are reconstructed with the general model. The gallery
# examples were baked with the DL3DV model. An artifact records the variant that
# baked it ("hf_variant") and has to be rendered with that same model.
UPLOAD_VARIANT = "DVSM-TTB-p8-general-960p"
EXAMPLE_VARIANT = "DVSM-TTB-p8-DL3DV-960p"
# How long one held-GPU interactive session lasts before it must be renewed.
# Bounded by the per-call duration cap and the user's daily quota
# (anon 2min / free 5min / PRO 40min). Override via env for PRO accounts.
SESSION_SECONDS = int(os.environ.get("SESSION_SECONDS", "120"))
RECON_DURATION = int(os.environ.get("RECON_DURATION", "120"))
# Reduced benchmark width set keeps per-session warmup short (every session
# re-warms from scratch — CUDA graphs can't persist across GPU workers).
BENCH_WIDTHS = [int(x) for x in os.environ.get("BENCH_WIDTHS", "512 768").split()]
# Upload budget: frames x tokens per frame may not exceed 64 views at 960 px
# (960x544 at patch 8 = 120x68 tokens per view). The two sliders trade one for
# the other: the one being moved wins and the other gives way (_FIT_*_JS), using
# the aspect-free form frames * width^2 <= 64 * 960^2. reconstruct_cb then checks
# the real token count, because an input taller than 16:9 has more tokens per
# frame at the same long axis.
PATCH = 8
TOKEN_BUDGET = 64 * (960 // PATCH) * (544 // PATCH)
FRAMES_MIN, FRAMES_MAX, FRAMES_STEP = 8, 128, 4
WIDTH_MIN, WIDTH_MAX, WIDTH_STEP = 384, 960, 32
_SLIDER_BUDGET = 64 * 960 ** 2

ARTIFACTS_DIR = Path(tempfile.gettempdir()) / "dvsm_artifacts"
ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)

# Pre-built example scenes are auto-discovered from examples/: every
# "<stem>.recon.pt" (baked with THIS model) that has a matching "<stem>.jpg|png"
# thumbnail becomes a gallery entry. Drop files in / remove them to change the
# gallery — no code edit needed. Clicking a thumbnail skips reconstruction and
# opens the live session directly on the artifact.
EXAMPLES_DIR = PROJECT_ROOT / "examples"


def _discover_examples():
    """Every "<stem>.recon.pt" with a matching "<stem>.jpg|jpeg|png|webp"
    thumbnail. The gallery shows thumbnails only (no labels)."""
    scenes = []
    # Gallery order: Mip-NeRF 360, then DL3DV, then our own captures.
    order = ("mipnerf360_", "dl3dv_")
    rank = lambda p: next((i for i, k in enumerate(order) if p.name.startswith(k)), len(order))
    for recon in sorted(EXAMPLES_DIR.glob("*.recon.pt"), key=lambda p: (rank(p), p.name)):
        stem = recon.name[: -len(".recon.pt")]
        thumb = next((EXAMPLES_DIR / f"{stem}{ext}"
                      for ext in (".jpg", ".jpeg", ".png", ".webp")
                      if (EXAMPLES_DIR / f"{stem}{ext}").exists()), None)
        if thumb is not None:
            scenes.append({"thumb": thumb, "recon": recon})
    return scenes


EXAMPLE_SCENES = _discover_examples()

# ---------------------------------------------------------------------------
# Module-level model load: download + keep on CPU (ZeroGPU convention).
# Real `.cuda()` happens inside @spaces.GPU workers.
# ---------------------------------------------------------------------------
print(f"[app] loading {UPLOAD_VARIANT}, {EXAMPLE_VARIANT} + π³ to CPU ...", flush=True)
from dvsm.models.DVSM import DecoderViewSynthesisModel
from viser_app.pi3_estimator import load_pi3_model

# APP_MODEL_DIR: load the DVSM models from a local directory instead of the Hub
# (faster / offline; used for local GPU-node testing). It holds one HF-format
# sub-directory per variant, i.e. the layout of the nvidia/DVSM repo.
_LOCAL_MODEL_DIR = os.environ.get("APP_MODEL_DIR")


def _load_dvsm(variant):
    if _LOCAL_MODEL_DIR:
        path = Path(_LOCAL_MODEL_DIR) / variant
        print(f"[app] loading {variant} from local dir {path}", flush=True)
        return DecoderViewSynthesisModel.from_pretrained(str(path)).eval()
    return DecoderViewSynthesisModel.from_hf_pretrained(variant).eval()


DVSM_MODELS_CPU = {v: _load_dvsm(v) for v in (UPLOAD_VARIANT, EXAMPLE_VARIANT)}
PI3_MODEL_CPU = load_pi3_model(flavor="pi3x", device="cpu").eval()
print("[app] models ready (CPU)", flush=True)

# ---------------------------------------------------------------------------
# Per-session viser registry (token -> worker viser port). The Gradio Space is
# a SINGLE shared process for ALL users — ZeroGPU only isolates GPU compute per
# @spaces.GPU call, NOT the app/server/state. So proxy routing MUST be keyed by
# a per-session token; a single global port would make every user's iframe
# point at whichever session set it last (they'd all collide on one viser).
# ---------------------------------------------------------------------------
_sessions: dict = {}              # token -> viser port (present == alive)
_sessions_lock = threading.Lock()


def _session_port(token: str):
    with _sessions_lock:
        return _sessions.get(token)


def _pick_free_port() -> int:
    s = socket.socket()
    s.bind(("", 0))
    port = s.getsockname()[1]
    s.close()
    return port


# ---------------------------------------------------------------------------
# GPU boundary (a): reconstruction
# ---------------------------------------------------------------------------
@spaces.GPU(duration=RECON_DURATION)
def reconstruct_gpu(img_dir: str, target_h: int, target_w: int, max_views: int) -> str:
    """π³ pose estimation + TTB-DVSM forward -> saved .recon.pt. Runs in a GPU
    worker; returns the artifact path (on shared /tmp) for the session call."""
    import torch
    from viser_app.pi3_estimator import _list_images, pi3_to_scene_data, run_pi3
    from viser_app.recon_scene import reconstruct, save_artifact

    repair_torch_cuda_state()  # torch.compile vs ZeroGPU, see zerogpu_compat
    image_paths = _list_images(Path(img_dir))
    if max_views and len(image_paths) > max_views:
        idxs = np.linspace(0, len(image_paths) - 1, max_views).round().astype(int)
        image_paths = [image_paths[i] for i in idxs]

    pi3 = PI3_MODEL_CPU.cuda()
    est = run_pi3(pi3, image_paths, image_size=518)
    del pi3
    torch.cuda.empty_cache()
    scene_data = pi3_to_scene_data(est, target_h, target_w)

    # Reuse the module-level CPU model (moved to cuda inside reconstruct()).
    renderer, batch, ishape, c2w_r, c2w_t, intr = reconstruct(
        ckpt_root=None, ckpt_name="", scene_data=scene_data, cmd_lst=[],
        hf_variant=UPLOAD_VARIANT, model=DVSM_MODELS_CPU[UPLOAD_VARIANT],
    )
    out_path = ARTIFACTS_DIR / f"{int(time.time() * 1000)}_{uuid.uuid4().hex[:8]}.recon.pt"
    save_artifact(renderer, batch, ishape, c2w_r, c2w_t, intr, out_path,
                  hf_variant=UPLOAD_VARIANT)
    return str(out_path)


# ---------------------------------------------------------------------------
# GPU boundary (b): time-boxed interactive viser session (generator)
# ---------------------------------------------------------------------------
class _SessionArgs:
    """Mimics the argparse namespace install_viser_handlers expects."""
    max_res = 768
    no_async_push = False
    bench_widths = BENCH_WIDTHS


@spaces.GPU(duration=SESSION_SECONDS)
def _live_session_worker(artifact_path: str):
    """Generator running in the GPU worker. Builds the renderer, starts a viser
    server on a self-chosen free port, then holds the GPU until the session
    budget expires.

    Yields tagged tuples so the UI can show a prep progress bar then a live
    countdown:
        ("prep",  frac, desc)     during setup
        ("ready", port, None)     viser is up and serving on `port`
        ("alive", remaining, None) per-second heartbeat while holding the GPU
    """
    import torch
    import viser
    from dvsm.models.pose_preprocessor import PosePreprocessor
    from viser_app.render import build_renderer
    from viser_app.viewer import install_viser_handlers
    from viser_app.viewer_optim import install_optimizations

    repair_torch_cuda_state()  # torch.compile vs ZeroGPU, see zerogpu_compat
    yield ("prep", 0.15, "Loading scene…")
    artifact = torch.load(artifact_path, weights_only=False, map_location="cpu")
    # Render with the model that baked this artifact (uploads vs gallery examples).
    variant = artifact.get("hf_variant") or EXAMPLE_VARIANT
    if variant not in DVSM_MODELS_CPU:
        raise RuntimeError(f"artifact was baked with {variant}, which this app does not load")
    print(f"[session] rendering with {variant}", flush=True)
    model = DVSM_MODELS_CPU[variant].cuda().eval()
    renderer = build_renderer(model, artifact)
    scene_meta = artifact["scene_meta"]

    pp_cpu = PosePreprocessor.__new__(PosePreprocessor)
    pp_cpu.r = artifact["pose_preproc"]["r"].cpu()
    pp_cpu.t = artifact["pose_preproc"]["t"].cpu()
    pp_cpu.s = artifact["pose_preproc"]["s"].cpu()
    renderer.pose_preprocessor = None
    pp_scale = float(pp_cpu.s.reshape(-1)[0])

    yield ("prep", 0.4, "Loading scene…")
    renderer = install_optimizations(
        renderer, use_fast_plucker=True, use_cuda_graph=True,
        n_warmup=5, max_graphs=len(BENCH_WIDTHS) + 2,
    )

    yield ("prep", 0.6, "Loading scene…")
    # Pick + bind the port HERE (inside the worker), with retry, so concurrent
    # sessions can't race on the same port (no cross-process TOCTOU).
    server = None
    for _ in range(10):
        port = _pick_free_port()
        try:
            server = viser.ViserServer(host="0.0.0.0", port=port)
            break
        except Exception:
            server = None
    if server is None:
        raise RuntimeError("could not bind a free viser port")
    server.scene.set_up_direction("+y")
    install_viser_handlers(server, renderer, scene_meta, pp_cpu, pp_scale,
                           model, _SessionArgs())

    deadline = time.monotonic() + SESSION_SECONDS - 5  # margin before hard kill
    yield ("ready", port, None)
    while True:
        remaining = int(deadline - time.monotonic())
        if remaining <= 0:
            break
        yield ("alive", remaining, None)
        time.sleep(1.0)
    try:
        server.stop()
    except Exception:
        pass


def _placeholder_html(msg: str = "Viser Viewer") -> str:
    """Same-size 16:9 placeholder shown before the viser iframe loads, so the
    layout doesn't jump and the user knows where the viewer will be."""
    # No own border/background — the gr.HTML(container=True) .block already
    # provides the single frame + background. Just a 16:9 centered hint, theme-
    # aware text; the ◧ accent stays NVIDIA green in light/dark.
    return (
        "<div style='width:100%;aspect-ratio:16/9;overflow:hidden;display:flex;"
        "flex-direction:column;align-items:center;justify-content:center;"
        "color:var(--body-text-color-subdued);text-align:center;gap:8px'>"
        "<div style='font-size:2.4em;line-height:1;color:#76B900'>◧</div>"
        f"<div style='font-size:1.05em;font-weight:600'>{msg}</div></div>"
    )


def _iframe_html(token: str) -> str:
    # Token in the PATH so every viser sub-request (assets + websocket) inherits
    # it and the proxy can route this user's traffic to their own worker.
    # No own border — the gr.HTML(container=True) .block frames it.
    return (
        f"<div style='width:100%;aspect-ratio:16/9;border-radius:6px;overflow:hidden'>"
        f"<iframe src='/viser-proxy/{token}/' style='width:100%;height:100%;border:0' "
        f"allow='clipboard-write'></iframe></div>"
    )


def _viewer_message(msg: str, color: str = "var(--body-text-color-subdued)") -> str:
    """A status message centered inside the same 16:9 viewer box, so recon
    progress / results overlay the placeholder area (aligned with the upload
    box on the left) instead of appearing below the viewer."""
    return (
        "<div style='width:100%;aspect-ratio:16/9;overflow:hidden;display:flex;"
        "align-items:center;justify-content:center;text-align:center;padding:0 1em;"
        f"color:{color};font-weight:600;font-size:1.05em'>{msg}</div>"
    )


def _prep_overlay(frac: float, desc: str) -> str:
    """Determinate green progress bar centered in the viewer box, shown while
    the viser session is being prepared (renderer build + warmup + bench)."""
    pct = max(0, min(100, round(100 * frac)))
    return (
        "<div style='width:100%;aspect-ratio:16/9;overflow:hidden;display:flex;"
        "flex-direction:column;align-items:center;justify-content:center;gap:10px;"
        "padding:0 12%;color:var(--body-text-color-subdued)'>"
        f"<div style='font-weight:600;font-size:1.05em'>{desc}</div>"
        "<div style='width:100%;height:8px;background:#333;border-radius:4px;overflow:hidden'>"
        f"<div style='height:100%;width:{pct}%;background:#76B900;"
        "transition:width .3s linear'></div></div></div>"
    )


# ---------------------------------------------------------------------------
# Gradio callbacks (main process)
# ---------------------------------------------------------------------------
def _delete_artifact(path):
    """Remove a user-generated recon artifact. Used both as the gr.State
    delete_callback (fires when the session ends / TTL expires) and to drop the
    previous artifact when a user reconstructs again. Guarded to ARTIFACTS_DIR so
    it never touches the permanent examples/ files."""
    try:
        if path and isinstance(path, str):
            p = Path(path)
            if p.parent == ARTIFACTS_DIR and p.exists():
                p.unlink()
    except Exception:
        pass


def reconstruct_cb(input_files, n_frames, target_w, prev_artifact,
                   progress=gr.Progress(track_tqdm=False)):
    """CPU pre-processing (decode video / copy images / size) then one GPU
    reconstruction call. Returns (artifact_path, viewer_html, start_btn_update).
    The status overlays the viewer box (see _viewer_message)."""
    # Drop this session's previous artifact (if any) before making a new one.
    _delete_artifact(prev_artifact)
    if not input_files:
        return "", _viewer_message("Upload a video or at least 3 images.", "#c66"), gr.update(interactive=False)

    work = Path(tempfile.mkdtemp(prefix="dvsm_in_"))
    img_dir = work / "images"
    img_dir.mkdir()
    # The uploaded/extracted input frames are only needed during reconstruction;
    # delete the whole work dir on exit so /tmp doesn't accumulate per request.
    try:
        paths = [Path(f.name if hasattr(f, "name") else f) for f in input_files]
        videos = [p for p in paths if p.suffix.lower() in VIDEO_EXTS]
        images = [p for p in paths if p.suffix.lower() in IMAGE_EXTS]
        n_frames = min(max(int(n_frames), 3), FRAMES_MAX)

        if videos:
            progress(0.05, desc=f"Extracting {n_frames} frames from video...")
            try:
                n = extract_video_frames(videos[0], img_dir, n_frames)
            except Exception as e:
                return "", _viewer_message(f"Video decode failed: {e}", "#c66"), gr.update(interactive=False)
            if n < 3:
                return "", _viewer_message(f"Only {n} frames extracted; need ≥3.", "#c66"), gr.update(interactive=False)
            n_inputs = n
        elif len(images) >= 3:
            if len(images) > n_frames:
                idxs = np.linspace(0, len(images) - 1, n_frames).round().astype(int)
                images = [images[i] for i in idxs]
            progress(0.05, desc=f"Copying {len(images)} files...")
            for src in images:
                shutil.copy(src, img_dir / src.name)
            n_inputs = len(images)
        else:
            return "", _viewer_message("Upload a video or at least 3 images.", "#c66"), gr.update(interactive=False)

        # Snap target dims to the input aspect; slider is the long axis; multiples of 8.
        try:
            from PIL import Image as _PILImage
            first_img = next(iter(sorted(img_dir.iterdir())))
            with _PILImage.open(first_img) as im:
                iw, ih = im.size
            input_aspect = iw / ih
        except Exception:
            input_aspect = 16 / 9
        long_axis = (min(int(target_w), WIDTH_MAX) // 8) * 8
        if input_aspect >= 1.0:
            target_w, target_h = long_axis, int(round(long_axis / input_aspect))
        else:
            target_h, target_w = long_axis, int(round(long_axis * input_aspect))
        target_h = max(8, (target_h // 8) * 8)
        target_w = max(8, (target_w // 8) * 8)

        # Token budget, on the real size. The sliders already keep a 16:9 input
        # inside it; a taller aspect gets fewer (evenly spaced) frames.
        tokens_per_view = (target_h // PATCH) * (target_w // PATCH)
        n_views = max(3, min(n_inputs, TOKEN_BUDGET // tokens_per_view))
        trimmed = (f" Used {n_views} of {n_inputs} frames to stay within the token budget."
                   if n_views < n_inputs else "")

        progress(0.2, desc=f"Running π³ + TTB-DVSM ({n_views} frames)...")
        try:
            artifact_path = reconstruct_gpu(str(img_dir), target_h, target_w, int(n_views))
        except Exception as e:
            return "", _viewer_message(f"Reconstruction failed: {e}", "#c66"), gr.update(interactive=False)

        return (
            artifact_path,
            _viewer_message(f"Reconstruction done ({n_views} views, {target_w}×{target_h}).{trimmed} "
                    f"Click “Start / Continue Live Session”."),
            gr.update(interactive=True),
        )
    finally:
        shutil.rmtree(work, ignore_errors=True)


def _session_bar(remaining: int, total: int, expired: bool = False) -> str:
    """Our own animated green countdown bar (one bar, rendered below the viewer).
    We draw it ourselves rather than relying on Gradio's auto progress, which
    would put a bar on every output component of the streaming event."""
    if expired:
        return (
            "<div style='margin-top:4px'>"
            "<div style='height:8px;background:#333;border-radius:4px'></div>"
            "<div style='font-weight:600;color:#c89b3c;margin-top:5px'>Session expired "
            "(GPU released). Click “Start / Continue Live Session” to resume.</div></div>"
        )
    pct = max(0, min(100, round(100 * remaining / max(1, total))))
    return (
        "<div style='margin-top:4px'>"
        "<div style='height:8px;background:#333;border-radius:4px;overflow:hidden'>"
        f"<div style='height:100%;width:{pct}%;background:#76B900;"
        "transition:width 1s linear'></div></div>"
        f"<div style='font-weight:600;color:#76B900;margin-top:5px'>Live — ~{remaining}s "
        "of GPU time left. Drag to explore.</div></div>"
    )


def live_session(artifact_path):
    """Gradio generator handler. Drives `_live_session_worker` (which holds the
    GPU) and streams the iframe + green countdown bar to the UI.

    Outputs: (viewer_html, session_status). The first yield swaps the recon
    message in the viewer box for the iframe; the countdown bar lives below."""
    if not artifact_path:
        yield _viewer_message("Reconstruct a scene first.", "#c66"), gr.update()
        return

    token = uuid.uuid4().hex
    try:
        for tag, a, _b in _live_session_worker(artifact_path):
            if tag == "prep":
                # a = fraction, _b = description -> prep bar overlaid in the viewer
                yield (_prep_overlay(a, _b), gr.update())
            elif tag == "ready":
                # viser is up on port `a` -> register this token's route, then
                # swap in the iframe (which carries the token).
                with _sessions_lock:
                    _sessions[token] = int(a)
                yield (_iframe_html(token), _session_bar(SESSION_SECONDS, SESSION_SECONDS))
            else:  # "alive": a = seconds remaining
                yield (gr.update(), _session_bar(int(a), SESSION_SECONDS))
    finally:
        with _sessions_lock:
            _sessions.pop(token, None)
    yield (gr.update(), _session_bar(0, SESSION_SECONDS, expired=True))


# ---------------------------------------------------------------------------
# Reverse proxy: /viser-proxy/* -> worker-local viser port
# ---------------------------------------------------------------------------
def _build_proxy_handlers():
    import asyncio
    import httpx
    import websockets as ws_pkg
    from starlette.requests import Request
    from starlette.responses import Response
    from starlette.websockets import WebSocket

    http_client = httpx.AsyncClient(timeout=60.0)

    async def proxy_http(request: Request) -> Response:
        token = request.path_params.get("token", "")
        path = request.path_params.get("path", "")
        port = _session_port(token)
        if port is None:
            return Response(b"No active viser session.", status_code=503)
        url = f"http://localhost:{port}/{path}"
        if request.url.query:
            url += "?" + request.url.query
        headers = {k: v for k, v in request.headers.items()
                   if k.lower() not in ("host", "content-length")}
        body = await request.body()
        try:
            r = await http_client.request(request.method, url, headers=headers,
                                          content=body, follow_redirects=False)
        except Exception as e:
            return Response(f"viser unreachable: {e}".encode(), status_code=502)
        resp_headers = {k: v for k, v in r.headers.items()
                        if k.lower() not in ("content-encoding", "transfer-encoding",
                                             "content-length", "connection")}
        return Response(content=r.content, status_code=r.status_code, headers=resp_headers)

    async def proxy_ws(websocket: WebSocket) -> None:
        token = websocket.path_params.get("token", "")
        path = websocket.path_params.get("path", "")
        port = _session_port(token)
        offered = list(websocket.scope.get("subprotocols") or [])
        if port is None:
            await websocket.close(code=1011)
            return
        try:
            upstream = await ws_pkg.connect(
                f"ws://localhost:{port}/{path}", subprotocols=offered or None,
                max_size=50 * 1024 * 1024, compression=None,
            )
        except Exception:
            await websocket.close(code=1011)
            return
        await websocket.accept(subprotocol=upstream.subprotocol)

        async def c2u():
            try:
                while True:
                    msg = await websocket.receive()
                    if msg.get("type") == "websocket.disconnect":
                        return
                    b, t = msg.get("bytes"), msg.get("text")
                    if b is not None:
                        await upstream.send(b)
                    elif t is not None:
                        await upstream.send(t)
            except Exception:
                pass
            finally:
                try:
                    await upstream.close()
                except Exception:
                    pass

        async def u2c():
            try:
                async for msg in upstream:
                    if isinstance(msg, (bytes, bytearray)):
                        await websocket.send_bytes(bytes(msg))
                    else:
                        await websocket.send_text(msg)
            except Exception:
                pass
            finally:
                try:
                    await websocket.close()
                except Exception:
                    pass

        await asyncio.gather(c2u(), u2c())

    return proxy_http, proxy_ws


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
NVIDIA_GREEN = "#76B900"
NVIDIA_GREEN_DARK = "#5C8F00"
_PALETTE = gr.themes.Color(
    name="nvidia_green", c50="#F0F8E0", c100="#E1F0C2", c200="#C7E690",
    c300="#A8D958", c400="#8BC53F", c500=NVIDIA_GREEN, c600=NVIDIA_GREEN_DARK,
    c700="#4A7300", c800="#395800", c900="#283D00", c950="#1A2900",
)
_theme = gr.themes.Default(primary_hue=_PALETTE).set(
    button_primary_background_fill=NVIDIA_GREEN,
    button_primary_background_fill_hover=NVIDIA_GREEN_DARK,
    button_primary_text_color="white",
    slider_color=NVIDIA_GREEN,
)
# Gradio 6 ignores Blocks(theme=...) when the app is served via
# mount_gradio_app (theme moved to launch()). Force NVIDIA green with CSS so the
# palette applies regardless of how the app is launched.
_custom_css = """
button.primary, button[variant='primary'], .gr-button-primary {
    background: #76B900 !important; background-image: none !important;
    border-color: #76B900 !important; box-shadow: none !important; color: #fff !important;
}
button.primary:hover, button[variant='primary']:hover, .gr-button-primary:hover {
    background: #5C8F00 !important; background-image: none !important; border-color: #5C8F00 !important;
}
input[type='range']::-webkit-slider-thumb { background: #76B900 !important; }
input[type='range']::-moz-range-thumb     { background: #76B900 !important; }
.gr-slider .filled, .slider .filled       { background: #76B900 !important; }
*:focus-visible, *:focus { outline-color: #76B900 !important; }
.gradio-container input:focus, .gradio-container textarea:focus, .gradio-container select:focus {
    border-color: #76B900 !important; box-shadow: 0 0 0 2px rgba(118, 185, 0, 0.4) !important;
}
.gallery .selected, .gr-gallery .selected,
.gallery .thumbnail:hover, .gr-gallery .thumbnail:hover { border-color: #76B900 !important; }
.progress, .gr-loader, .loader { background: #76B900 !important; }
.tab-nav .selected, .gr-tab-nav .selected { border-bottom-color: #76B900 !important; }
/* Flush blocks: zero the .block inner padding so the upload box and the viewer
   sit flush at the column top (kept aligned) and the viewer fills its frame
   edge-to-edge (no padding gap around the iframe). */
.flush-block { padding: 0 !important; }
/* Proportional columns: flex-basis 0 so the row width is always split 1:3 and
   BOTH columns scale with the window (default Gradio scale only distributes the
   leftover space, leaving the left column looking fixed). */
.col-left  { flex: 1 1 0 !important; min-width: 260px; }
.col-right { flex: 3 1 0 !important; min-width: 360px; }
/* Header text links under the title (see _header_links): a quiet grey line,
   brighter and underlined on hover. Overrides what Gradio gives every <a> in an
   HTML block (button-like padding, centred text, its own link colour). The
   light theme gets darker greys so the line stays readable on white. */
.ttb-links { display: flex; flex-wrap: wrap; gap: 0 0.6em; font-size: 1em; line-height: 1.6; margin-top: 0.15em; }
.ttb-links a {
    padding: 0 !important; text-align: left !important; font-weight: 400;
    color: var(--neutral-500) !important; text-decoration: none !important;
}
.dark .ttb-links a { color: var(--body-text-color-subdued) !important; }
.ttb-links a:hover { color: var(--body-text-color) !important; text-decoration: underline !important; }
.dark .ttb-links a:hover { color: var(--body-text-color) !important; }
.ttb-links span { color: var(--neutral-300) !important; }
.dark .ttb-links span { color: var(--border-color-primary) !important; }
"""

# Paper / project page / code links shown as a line of text links under the
# title. A link with an empty URL is left out (the arXiv link, until the paper
# has its arXiv page).
ARXIV_URL = "https://arxiv.org/abs/2611.17331"
PROJECT_URL = "https://test-time-baking.github.io/"
CODE_URL = "https://github.com/NVlabs/DVSM"


def _header_links() -> str:
    links = [("arXiv", ARXIV_URL), ("Project Page", PROJECT_URL), ("Code", CODE_URL)]
    # target=_blank: the Space runs in an iframe on huggingface.co.
    return "<div class='ttb-links'>" + "<span>/</span>".join(
        f"<a href='{url}' target='_blank' rel='noopener noreferrer'>{label}</a>"
        for label, url in links if url) + "</div>"


# Slider coupling (frames, width) -> the largest value of the OTHER slider that
# still fits frames * width^2 <= _SLIDER_BUDGET, never raising it.
_FIT_WIDTH_JS = f"""(n, w) => {{
    if (n * w * w <= {_SLIDER_BUDGET}) return [w];
    let m = {WIDTH_MAX};
    while (m > {WIDTH_MIN} && n * m * m > {_SLIDER_BUDGET}) m -= {WIDTH_STEP};
    return [Math.min(w, m)];
}}"""
_FIT_FRAMES_JS = f"""(n, w) => {{
    if (n * w * w <= {_SLIDER_BUDGET}) return [n];
    let m = {FRAMES_MAX};
    while (m > {FRAMES_MIN} && m * w * w > {_SLIDER_BUDGET}) m -= {FRAMES_STEP};
    return [Math.min(n, m)];
}}"""

# delete_cache=(freq, age): Gradio sweeps its OWN cache (uploaded files etc.)
# every hour, removing entries older than an hour — a safety net against disk
# exhaustion on HF Spaces.
with gr.Blocks(title="TTB-DVSM (ZeroGPU)", delete_cache=(3600, 3600)) as demo:
    # delete_callback ties each user's recon artifact lifetime to their session:
    # Gradio calls it when the State is cleaned up (tab closed +60min, or TTL),
    # so /tmp artifacts don't accumulate. time_to_live also bounds idle sessions.
    artifact_state = gr.State("", delete_callback=_delete_artifact, time_to_live=3600)
    # Full-width description on top.
    gr.HTML(
        "<h1 style='margin-bottom:0.1em;color:#76B900'>TTB-DVSM</h1>"
        + _header_links() +
        "<p style='color:var(--body-text-color);font-size:1.1em;margin-top:0.5em;line-height:1.6'>"
        "TTB bakes the scene into fast weights in a single forward pass, "
        "and the result is browseable in the viser viewer."
        "</p>"
        "<div style='color:var(--body-text-color);font-size:1.1em;margin-top:0.8em;line-height:1.7'>"
        "<div style='color:#76B900;font-weight:700;font-size:1.15em;margin-bottom:0.3em'>"
        "Getting Started</div>"
        "<b>1.</b> Drop 16–128 photos <i>or</i> a short orbit video into the input box.<br>"
        "<b>2.</b> Pick how many frames to use and the reconstruction width, "
        "then click <b>Reconstruct</b>.<br>"
        "<b>3.</b> Click <b>Start / Continue Live Session</b> to open the "
        "interactive viewer.<br>"
        "<i>ZeroGPU holds the GPU only for a bounded session; when it expires, "
        "click Continue to resume.</i>"
        "<div style='color:#76B900;font-weight:700;font-size:1.15em;margin:0.7em 0 0.3em'>"
        "Viewer controls</div>"
        "&nbsp;• <b>Left-drag</b> orbit &nbsp;&nbsp;"
        "<b>Right-drag</b> pan &nbsp;&nbsp;"
        "<b>Scroll</b> zoom<br>"
        "&nbsp;• <b>W A S D</b> fly through the scene &nbsp;&nbsp;"
        "<b>Q / E</b> down / up &nbsp;&nbsp;"
        "<b>Shift</b> faster<br>"
        "&nbsp;• Right-side panel: <b>Max W</b>, "
        "<b>Jump to input view</b>, <b>Reset View</b>."
        "</div>"
    )
    # Below the text: 1 (inputs + buttons) : 3 (viser viewer).
    with gr.Row():
        with gr.Column(scale=1, min_width=260, elem_classes=["col-left"]):
            inputs = gr.File(
                file_count="multiple",
                file_types=["image", "video", ".mp4", ".mov", ".mkv",
                            ".webm", ".avi", ".m4v"],
                label="Drop photos OR a short video (orbiting one scene)",
                show_label=False,
                elem_classes=["flush-block"],
            )
            n_frames = gr.Slider(FRAMES_MIN, FRAMES_MAX, 32, step=FRAMES_STEP,
                                 label="Frames to sample")
            tw = gr.Slider(WIDTH_MIN, WIDTH_MAX, 768, step=WIDTH_STEP,
                           label="Reconstruction width (long axis)",
                           info="Budget: 64 frames at 960 px. More frames need a lower width.")
            recon_btn = gr.Button("Reconstruct", variant="primary")
            session_btn = gr.Button("Start / Continue Live Session",
                                    variant="primary", interactive=False)
            if EXAMPLE_SCENES:
                gr.Markdown("### Example")
                example_gallery = gr.Gallery(
                    # No captions. Fixed box height -> vertical scrollbar when
                    # there are more thumbnails than fit; width still scales with
                    # the (resizable) column since it's 3 columns.
                    value=[str(s["thumb"]) for s in EXAMPLE_SCENES],
                    columns=3, height=200, object_fit="cover",
                    show_label=False, allow_preview=False,
                )
            else:
                example_gallery = None
        with gr.Column(scale=3, elem_classes=["col-right"]):
            # Viewer first so its top edge aligns with the upload box on the
            # left (both flush via .flush-block). Recon progress/status overlay
            # the viewer box itself (see reconstruct_cb); only the live-session
            # countdown bar lives BELOW the viewer.
            viewer_html = gr.HTML(_placeholder_html(), container=True,
                                  elem_classes=["flush-block"])
            session_status = gr.HTML()

    # The two sliders share one budget (see TOKEN_BUDGET). Client-side only: the
    # slider that changed keeps its value and the other is lowered just enough to
    # fit. `change` rather than `input` so the reset buttons are covered too; the
    # lowered slider's own change event is a no-op because the pair then fits.
    n_frames.change(None, [n_frames, tw], [tw], js=_FIT_WIDTH_JS)
    tw.change(None, [n_frames, tw], [n_frames], js=_FIT_FRAMES_JS)

    # Recon status overlays the viewer box; the gradio progress bar (the "20%")
    # therefore also shows over the viewer, aligned with the upload column.
    recon_btn.click(reconstruct_cb, [inputs, n_frames, tw, artifact_state],
                    [artifact_state, viewer_html, session_btn])
    # show_progress="hidden": the session generator streams for the whole GPU
    # hold; without this Gradio overlays an auto progress bar on EACH output
    # component. We draw our own green countdown bar instead.
    session_btn.click(live_session, [artifact_state],
                      [viewer_html, session_status],
                      show_progress="hidden")

    # Clicking an example thumbnail loads its prebuilt artifact and opens the
    # live session directly (no reconstruction).
    if example_gallery is not None:
        def _on_example_select(evt: gr.SelectData):
            idx = int(evt.index)
            if 0 <= idx < len(EXAMPLE_SCENES):
                return str(EXAMPLE_SCENES[idx]["recon"]), gr.update(interactive=True)
            return "", gr.update(interactive=False)

        example_gallery.select(
            _on_example_select, None, [artifact_state, session_btn],
        ).then(
            live_session, [artifact_state], [viewer_html, session_status],
            show_progress="hidden",
        )


# ---------------------------------------------------------------------------
# Launch Gradio (which starts uvicorn) and attach the viser reverse-proxy routes
# onto gradio's own FastAPI app. Using demo.launch() (not mount_gradio_app) lets
# us pass theme/css the Gradio-6 way and matches the proven viser_app pattern.
# ---------------------------------------------------------------------------
def _attach_proxy_routes(app):
    """PREPEND the /viser-proxy routes so they win over gradio's SPA catch-all.
    Appending (app.add_route) puts them AFTER gradio's catch-all route, which
    then serves index.html for /viser-proxy/* — the iframe ends up showing the
    gradio app again instead of viser. Inserting at the front fixes that."""
    from starlette.routing import Route, WebSocketRoute
    proxy_http, proxy_ws = _build_proxy_handlers()
    methods = ["GET", "POST", "PUT", "DELETE", "HEAD", "OPTIONS"]
    routes = [
        Route("/viser-proxy/{token}", proxy_http, methods=methods),
        Route("/viser-proxy/{token}/{path:path}", proxy_http, methods=methods),
        WebSocketRoute("/viser-proxy/{token}", proxy_ws),
        WebSocketRoute("/viser-proxy/{token}/{path:path}", proxy_ws),
    ]
    app.router.routes[0:0] = routes  # prepend → highest match priority
    print("[app] viser-proxy routes prepended; ready", flush=True)


def launch_app():
    """Launch gradio and attach the viser reverse-proxy. Done at MODULE level
    (below) so it runs no matter how HF Spaces starts the app — whether it
    `exec`s `python app.py` or `import`s the module and would otherwise call
    demo.launch() itself. prevent_thread_lock=True lets us attach routes after
    launch; we then block to keep the process alive."""
    # Unlimited Gradio-side concurrency: live_session is a long generator that
    # holds for the whole session, so Gradio's default limit of 1 would make a
    # second window queue behind the first and never show. We let ZeroGPU's own
    # scheduler handle GPU allocation/queueing across concurrent sessions.
    demo.queue(default_concurrency_limit=None)
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", "7860")),
        prevent_thread_lock=True,
        theme=_theme,
        css=_custom_css,
        # CRITICAL on HF Spaces: Gradio 6 defaults to SSR, which puts a Node
        # proxy on :7860 in front of the Python app on :7861. That Node proxy
        # serves the gradio SSR page for unknown paths, so /viser-proxy/* never
        # reaches our FastAPI routes (the iframe then shows gradio again).
        # Disabling SSR makes Python/FastAPI serve :7860 directly so the proxy
        # routes work — same topology as the local node run.
        ssr_mode=False,
    )
    _attach_proxy_routes(demo.app)
    threading.Event().wait()  # keep the process alive


# Launch at import/exec time (Spaces-robust). Set APP_NO_LAUNCH=1 to import this
# module without starting the server (e.g. for unit tests).
if os.environ.get("APP_NO_LAUNCH") != "1":
    launch_app()