"""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 # ".recon.pt" (baked with THIS model) that has a matching ".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 ".recon.pt" with a matching ".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 ( "
" "
◧
" f"
{msg}
" ) 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"
" f"
" ) 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 ( "
{msg}
" ) 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 ( "
" f"
{desc}
" "
" f"
" ) # --------------------------------------------------------------------------- # 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 ( "
" "
" "
Session expired " "(GPU released). Click “Start / Continue Live Session” to resume.
" ) pct = max(0, min(100, round(100 * remaining / max(1, total)))) return ( "
" "
" f"
" f"
Live — ~{remaining}s " "of GPU time left. Drag to explore.
" ) 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 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 "
" + "/".join( f"{label}" for label, url in links if url) + "
" # 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( "

TTB-DVSM

" + _header_links() + "

" "TTB bakes the scene into fast weights in a single forward pass, " "and the result is browseable in the viser viewer." "

" "
" "
" "Getting Started
" "1. Drop 16–128 photos or a short orbit video into the input box.
" "2. Pick how many frames to use and the reconstruction width, " "then click Reconstruct.
" "3. Click Start / Continue Live Session to open the " "interactive viewer.
" "ZeroGPU holds the GPU only for a bounded session; when it expires, " "click Continue to resume." "
" "Viewer controls
" " • Left-drag orbit   " "Right-drag pan   " "Scroll zoom
" " • W A S D fly through the scene   " "Q / E down / up   " "Shift faster
" " • Right-side panel: Max W, " "Jump to input view, Reset View." "
" ) # 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()