Update app.py

app.py

@@ -6,90 +6,151 @@ import subprocess
 import shutil
 from pathlib import Path
 
-# Root of the repo — wherever app.py lives
-REPO_ROOT = Path(__file__).parent.resolve()
-
 import gradio as gr
 import spaces
 import numpy as np
 import torch
 
+# ──────────────────────────────────────────────
+# Paths
+# ──────────────────────────────────────────────
+REPO_ROOT = Path(__file__).parent.resolve()
+CKPT_DIR  = REPO_ROOT / "checkpoints"
+
+# Sentinel files — if any are missing we trigger a fresh download
+_REQUIRED_FILES = [
+    CKPT_DIR / "text_encoder" / "negative_prompt.pt",
+    CKPT_DIR / "model",          # directory is enough as a sentinel
+]
+
+_ckpts_ready = False
+
 
 # ──────────────────────────────────────────────
-# Lazy model loading (done inside GPU decorator)
+# Checkpoint helpers
 # ──────────────────────────────────────────────
-_model_loaded = False
+
+def _checkpoints_present() -> bool:
+    for p in _REQUIRED_FILES:
+        if not p.exists():
+            print(f"[Lyra] Missing checkpoint path: {p}")
+            return False
+    return True
+
+
+def _download_checkpoints() -> str:
+    """
+    Download all checkpoints from nvidia/Lyra-2.0 on HuggingFace Hub.
+    Uses snapshot_download so the full checkpoints/ tree lands at
+    REPO_ROOT/checkpoints/ — exactly where the inference scripts expect them.
+    Returns a status string for display in the UI.
+    """
+    if _checkpoints_present():
+        return "✅ Checkpoints already present — skipping download."
+
+    print("[Lyra] Checkpoints not found — downloading from nvidia/Lyra-2.0 …")
+
+    try:
+        from huggingface_hub import snapshot_download
+    except ImportError:
+        subprocess.run(
+            [sys.executable, "-m", "pip", "install", "huggingface_hub", "-q"],
+            check=True,
+        )
+        from huggingface_hub import snapshot_download
+
+    snapshot_download(
+        repo_id="nvidia/Lyra-2.0",
+        allow_patterns=["checkpoints/**"],   # only grab the checkpoints subtree
+        local_dir=str(REPO_ROOT),            # → REPO_ROOT/checkpoints/...
+        local_dir_use_symlinks=False,
+    )
+
+    if _checkpoints_present():
+        msg = "✅ Checkpoints downloaded successfully."
+        print(f"[Lyra] {msg}")
+        return msg
+
+    msg = ("❌ Download finished but sentinel files are still missing. "
+           "Check Space logs and available disk space (~50 GB required).")
+    print(f"[Lyra] {msg}")
+    return msg
 
 
 def _ensure_models():
-    global _model_loaded
-    if _model_loaded:
+    """Called before every inference run — downloads checkpoints if needed."""
+    global _ckpts_ready
+    if _ckpts_ready:
         return
-    # Models are expected to be pre-downloaded to ./checkpoints/
-    # on the Space via the HF repo or a setup script.
-    _model_loaded = True
+    _download_checkpoints()
+    _ckpts_ready = _checkpoints_present()
+    if not _ckpts_ready:
+        raise RuntimeError(
+            "Checkpoints unavailable. Please click 'Download Checkpoints', "
+            "check the Space logs, and ensure ~50 GB of storage is available."
+        )
 
 
-# ──────────────────────────────────────────────
-# Core inference helpers
-# ──────────────────────────────────────────────
+# ── Kick off download at module load so it's done before the first request ──
+print("[Lyra] Startup checkpoint check …")
+_download_checkpoints()
+_ckpts_ready = _checkpoints_present()
 
-def _write_caption_txt(image_path: str, caption: str) -> str:
-    """Write a .txt caption file beside the image and return the directory."""
-    img_path = Path(image_path)
-    txt_path = img_path.with_suffix(".txt")
-    txt_path.write_text(caption)
-    return str(img_path.parent)
 
+# ──────────────────────────────────────────────
+# Subprocess helper
+# ──────────────────────────────────────────────
 
-def _run(cmd: list[str], desc: str = "") -> tuple[bool, str]:
-    """Run a Python -m command using the current interpreter with PYTHONPATH set."""
-    # Replace bare "python" / "python3" with the real interpreter path
+def _run(cmd: list, desc: str = "") -> tuple:
+    """
+    Run a `python -m …` command with:
+      • sys.executable  → correct venv interpreter
+      • PYTHONPATH      → REPO_ROOT so lyra_2._src.* is importable
+      • cwd             → REPO_ROOT so 'checkpoints/model' resolves correctly
+    """
     if cmd[0] in ("python", "python3"):
         cmd = [sys.executable] + cmd[1:]
 
     run_env = {
         **os.environ,
-        "PYTHONPATH": str(REPO_ROOT),          # makes lyra_2._src.* importable
+        "PYTHONPATH": str(REPO_ROOT),
         "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
     }
-    print(f"[Lyra] {desc}: {' '.join(cmd)}")
+    print(f"[Lyra] {desc}: {' '.join(str(c) for c in cmd)}")
     result = subprocess.run(
         cmd,
         capture_output=True,
         text=True,
         env=run_env,
-        cwd=str(REPO_ROOT),                    # run from repo root
+        cwd=str(REPO_ROOT),
     )
     log = result.stdout + "\n" + result.stderr
     return result.returncode == 0, log
 
 
 # ──────────────────────────────────────────────
-# Zoom-in / Zoom-out trajectory (Option 1)
+# Zoom-in / Zoom-out trajectory  (Option 1)
 # ──────────────────────────────────────────────
 
-@spaces.GPU(duration=190)
+@spaces.GPU(duration=900)
 def run_zoomgs(
     image,
-    caption: str,
-    sample_id: int,
-    zoom_in_strength: float,
-    zoom_out_strength: float,
-    num_frames_in: int,
-    num_frames_out: int,
-    use_dmd: bool,
-    run_reconstruction: bool,
+    caption,
+    zoom_in_strength,
+    zoom_out_strength,
+    num_frames_in,
+    num_frames_out,
+    use_dmd,
+    run_reconstruction,
 ):
     _ensure_models()
 
     with tempfile.TemporaryDirectory() as tmp:
-        # Save uploaded image + caption
-        img_path = Path(tmp) / "input.png"
-        caption_path = Path(tmp) / "input.txt"
+        from PIL import Image as PILImage
 
-        from PIL import Image
-        Image.fromarray(image).save(img_path)
+        img_path     = Path(tmp) / "input.png"
+        caption_path = Path(tmp) / "input.txt"
+        PILImage.fromarray(image).save(img_path)
         caption_path.write_text(caption.strip() or "A scenic outdoor environment.")
 
         output_dir = Path(tmp) / "outputs" / "zoomgs"
@@ -98,77 +159,77 @@ def run_zoomgs(
         cmd = [
             "python", "-m", "lyra_2._src.inference.lyra2_zoomgs_inference",
             "--input_image_path", str(tmp),
-            "--sample_id", "0",           # we always name it input.png → id 0 equivalent
+            "--sample_id", "0",
             "--experiment", "lyra2",
-            "--checkpoint_dir", "checkpoints/model",
-            "--prompt_dir", str(tmp),
-            "--output_path", str(output_dir),
-            "--num_frames_zoom_in", str(num_frames_in),
-            "--num_frames_zoom_out", str(num_frames_out),
-            "--zoom_in_strength", str(zoom_in_strength),
-            "--zoom_out_strength", str(zoom_out_strength),
+            "--checkpoint_dir", str(CKPT_DIR / "model"),
+            "--prompt_dir",     str(tmp),
+            "--output_path",    str(output_dir),
+            "--num_frames_zoom_in",  str(int(num_frames_in)),
+            "--num_frames_zoom_out", str(int(num_frames_out)),
+            "--zoom_in_strength",    str(zoom_in_strength),
+            "--zoom_out_strength",   str(zoom_out_strength),
         ]
         if use_dmd:
             cmd.append("--use_dmd")
 
         ok, log = _run(cmd, "ZoomGS video generation")
 
-        # Locate output video
+        # The script writes to <output_dir>/<sample_id>/videos/<sample_id>.mp4
         video_path = output_dir / "0" / "videos" / "0.mp4"
         if not video_path.exists():
-            # Fallback: search recursively
             candidates = list(output_dir.rglob("*.mp4"))
             video_path = candidates[0] if candidates else None
 
         gs_video = None
-        if run_reconstruction and video_path and video_path.exists():
+        if run_reconstruction and video_path and Path(video_path).exists():
             ok2, log2 = _run(
                 ["python", "-m", "lyra_2._src.inference.vipe_da3_gs_recon",
                  "--input_video_path", str(video_path)],
                 "GS reconstruction",
             )
             log += "\n" + log2
-            ply_candidates = list(output_dir.rglob("gs_trajectory.mp4"))
-            if ply_candidates:
-                gs_video = str(ply_candidates[0])
+            gs_candidates = list(output_dir.rglob("gs_trajectory.mp4"))
+            if gs_candidates:
+                gs_video = str(gs_candidates[0])
 
         return (
-            str(video_path) if video_path and video_path.exists() else None,
+            str(video_path) if video_path and Path(video_path).exists() else None,
             gs_video,
             log[-4000:],
         )
 
 
 # ──────────────────────────────────────────────
-# Custom trajectory (Option 2)
+# Custom trajectory  (Option 2)
 # ──────────────────────────────────────────────
 
 @spaces.GPU(duration=900)
 def run_custom_traj(
     image,
     trajectory_file,
-    captions_json: str,
-    num_frames: int,
-    pose_scale: float,
-    use_dmd: bool,
-    run_reconstruction: bool,
+    captions_json,
+    num_frames,
+    pose_scale,
+    use_dmd,
+    run_reconstruction,
 ):
     _ensure_models()
 
     with tempfile.TemporaryDirectory() as tmp:
-        from PIL import Image
+        from PIL import Image as PILImage
+
         img_path = Path(tmp) / "first_frame.png"
-        Image.fromarray(image).save(img_path)
+        PILImage.fromarray(image).save(img_path)
 
         traj_path = Path(tmp) / "trajectory.npz"
         shutil.copy(trajectory_file.name, traj_path)
 
         captions_path = Path(tmp) / "captions.json"
         try:
-            json.loads(captions_json)   # validate
+            json.loads(captions_json)
             captions_path.write_text(captions_json)
-        except json.JSONDecodeError:
-            captions_path.write_text(json.dumps({"0": captions_json}))
+        except (json.JSONDecodeError, TypeError):
+            captions_path.write_text(json.dumps({"0": captions_json or ""}))
 
         output_dir = Path(tmp) / "outputs" / "custom"
         output_dir.mkdir(parents=True, exist_ok=True)
@@ -176,13 +237,13 @@ def run_custom_traj(
         cmd = [
             "python", "-m", "lyra_2._src.inference.lyra2_custom_traj_inference",
             "--input_image_path", str(img_path),
-            "--trajectory_path", str(traj_path),
+            "--trajectory_path",  str(traj_path),
             "--experiment", "lyra2",
-            "--checkpoint_dir", "checkpoints/model",
-            "--captions_path", str(captions_path),
-            "--num_frames", str(num_frames),
-            "--output_path", str(output_dir),
-            "--pose_scale", str(pose_scale),
+            "--checkpoint_dir", str(CKPT_DIR / "model"),
+            "--captions_path",  str(captions_path),
+            "--num_frames",     str(int(num_frames)),
+            "--output_path",    str(output_dir),
+            "--pose_scale",     str(pose_scale),
         ]
         if use_dmd:
             cmd.append("--use_dmd")
@@ -200,9 +261,9 @@ def run_custom_traj(
                 "GS reconstruction",
             )
             log += "\n" + log2
-            ply_candidates = list(output_dir.rglob("gs_trajectory.mp4"))
-            if ply_candidates:
-                gs_video = str(ply_candidates[0])
+            gs_candidates = list(output_dir.rglob("gs_trajectory.mp4"))
+            if gs_candidates:
+                gs_video = str(gs_candidates[0])
 
         return (
             str(video_path) if video_path else None,
@@ -211,180 +272,58 @@ def run_custom_traj(
         )
 
 
+# ──────────────────────────────────────────────
+# Manual download button handler
+# ──────────────────────────────────────────────
+
+def manual_download():
+    global _ckpts_ready
+    msg = _download_checkpoints()
+    _ckpts_ready = _checkpoints_present()
+    return msg
+
+
 # ──────────────────────────────────────────────
 # UI
 # ──────────────────────────────────────────────
 
 CSS = """
-/* ── Global reset & fonts ── */
 @import url('https://fonts.googleapis.com/css2?family=Syne:wght@400;600;700;800&family=DM+Mono:wght@300;400;500&display=swap');
-
 :root {
-    --bg:       #0a0c10;
-    --surface:  #111318;
-    --border:   #1e2230;
-    --accent:   #5affb0;
-    --accent2:  #a78bfa;
-    --text:     #e8eaf0;
-    --muted:    #5a5f72;
-    --radius:   12px;
-    --font-head: 'Syne', sans-serif;
-    --font-mono: 'DM Mono', monospace;
-}
-
-body, .gradio-container {
-    background: var(--bg) !important;
-    color: var(--text) !important;
-    font-family: var(--font-head) !important;
-}
-
-/* Header banner */
-#header {
-    background: linear-gradient(135deg, #0d1117 0%, #161b27 60%, #0f1520 100%);
-    border: 1px solid var(--border);
-    border-radius: var(--radius);
-    padding: 32px 40px 28px;
-    margin-bottom: 24px;
-    position: relative;
-    overflow: hidden;
-}
-#header::before {
-    content: '';
-    position: absolute;
-    inset: 0;
-    background: radial-gradient(ellipse 70% 60% at 80% 50%, rgba(94,255,176,0.06) 0%, transparent 70%),
-                radial-gradient(ellipse 50% 80% at 20% 80%, rgba(167,139,250,0.06) 0%, transparent 70%);
-    pointer-events: none;
-}
-#header h1 {
-    font-size: 2.4rem;
-    font-weight: 800;
-    letter-spacing: -0.02em;
-    margin: 0 0 8px;
-    background: linear-gradient(90deg, var(--accent) 0%, var(--accent2) 100%);
-    -webkit-background-clip: text;
-    -webkit-text-fill-color: transparent;
-}
-#header p {
-    color: var(--muted);
-    font-family: var(--font-mono);
-    font-size: 0.85rem;
-    margin: 0;
-    letter-spacing: 0.02em;
-}
-#header .badge {
-    display: inline-block;
-    margin-right: 8px;
-    padding: 3px 10px;
-    background: rgba(94,255,176,0.1);
-    border: 1px solid rgba(94,255,176,0.25);
-    border-radius: 20px;
-    color: var(--accent);
-    font-size: 0.75rem;
-    font-family: var(--font-mono);
-}
-
-/* Tabs */
-.tab-nav button {
-    background: transparent !important;
-    border: none !important;
-    border-bottom: 2px solid transparent !important;
-    color: var(--muted) !important;
-    font-family: var(--font-head) !important;
-    font-weight: 600 !important;
-    font-size: 0.95rem !important;
-    padding: 10px 20px !important;
-    transition: all .2s !important;
-}
-.tab-nav button.selected, .tab-nav button:hover {
-    color: var(--accent) !important;
-    border-bottom-color: var(--accent) !important;
-    background: transparent !important;
-}
-
-/* Panels / blocks */
-.gr-panel, .gr-box, .gradio-group {
-    background: var(--surface) !important;
-    border: 1px solid var(--border) !important;
-    border-radius: var(--radius) !important;
-}
-
-/* Inputs */
-input, textarea, .gr-input, .gr-textbox textarea {
-    background: #0d0f14 !important;
-    border: 1px solid var(--border) !important;
-    color: var(--text) !important;
-    font-family: var(--font-mono) !important;
-    border-radius: 8px !important;
-}
-input:focus, textarea:focus {
-    border-color: var(--accent) !important;
-    box-shadow: 0 0 0 2px rgba(94,255,176,0.12) !important;
-}
-
-/* Sliders */
-input[type=range] { accent-color: var(--accent) !important; }
-
-/* Buttons */
-button.primary, .gr-button-primary {
-    background: linear-gradient(135deg, var(--accent) 0%, #38d9a9 100%) !important;
-    color: #0a0c10 !important;
-    font-family: var(--font-head) !important;
-    font-weight: 700 !important;
-    border: none !important;
-    border-radius: 8px !important;
-    padding: 12px 28px !important;
-    font-size: 0.95rem !important;
-    letter-spacing: 0.01em !important;
-    transition: opacity .2s !important;
-}
-button.primary:hover { opacity: 0.85 !important; }
-
-button.secondary, .gr-button-secondary {
-    background: transparent !important;
-    border: 1px solid var(--border) !important;
-    color: var(--muted) !important;
-    font-family: var(--font-head) !important;
-    border-radius: 8px !important;
-}
-
-/* Labels */
-label, .gr-form > label, .block > label span {
-    color: var(--muted) !important;
-    font-family: var(--font-mono) !important;
-    font-size: 0.8rem !important;
-    letter-spacing: 0.04em !important;
-    text-transform: uppercase !important;
-}
-
-/* Log box */
-#log-box textarea {
-    font-size: 0.78rem !important;
-    color: #7af0b0 !important;
-    background: #060709 !important;
-}
-
-/* Accordion */
-.gr-accordion { border-color: var(--border) !important; }
-
-/* Info note */
-.info-note {
-    background: rgba(167,139,250,0.07);
-    border: 1px solid rgba(167,139,250,0.2);
-    border-radius: 8px;
-    padding: 12px 16px;
-    font-family: var(--font-mono);
-    font-size: 0.8rem;
-    color: #c4b5fd;
-    line-height: 1.6;
+    --bg:#0a0c10; --surface:#111318; --border:#1e2230;
+    --accent:#5affb0; --accent2:#a78bfa; --text:#e8eaf0; --muted:#5a5f72;
+    --radius:12px; --font-head:'Syne',sans-serif; --font-mono:'DM Mono',monospace;
 }
+body,.gradio-container{background:var(--bg)!important;color:var(--text)!important;font-family:var(--font-head)!important;}
+#header{background:linear-gradient(135deg,#0d1117 0%,#161b27 60%,#0f1520 100%);border:1px solid var(--border);border-radius:var(--radius);padding:32px 40px 28px;margin-bottom:24px;position:relative;overflow:hidden;}
+#header::before{content:'';position:absolute;inset:0;background:radial-gradient(ellipse 70% 60% at 80% 50%,rgba(94,255,176,.06) 0%,transparent 70%),radial-gradient(ellipse 50% 80% at 20% 80%,rgba(167,139,250,.06) 0%,transparent 70%);pointer-events:none;}
+#header h1{font-size:2.4rem;font-weight:800;letter-spacing:-.02em;margin:0 0 8px;background:linear-gradient(90deg,var(--accent) 0%,var(--accent2) 100%);-webkit-background-clip:text;-webkit-text-fill-color:transparent;}
+#header p{color:var(--muted);font-family:var(--font-mono);font-size:.85rem;margin:0;}
+#header .badge{display:inline-block;margin-right:8px;padding:3px 10px;background:rgba(94,255,176,.1);border:1px solid rgba(94,255,176,.25);border-radius:20px;color:var(--accent);font-size:.75rem;font-family:var(--font-mono);}
+.tab-nav button{background:transparent!important;border:none!important;border-bottom:2px solid transparent!important;color:var(--muted)!important;font-family:var(--font-head)!important;font-weight:600!important;font-size:.95rem!important;padding:10px 20px!important;transition:all .2s!important;}
+.tab-nav button.selected,.tab-nav button:hover{color:var(--accent)!important;border-bottom-color:var(--accent)!important;background:transparent!important;}
+.gr-panel,.gr-box,.gradio-group{background:var(--surface)!important;border:1px solid var(--border)!important;border-radius:var(--radius)!important;}
+input,textarea,.gr-input,.gr-textbox textarea{background:#0d0f14!important;border:1px solid var(--border)!important;color:var(--text)!important;font-family:var(--font-mono)!important;border-radius:8px!important;}
+input:focus,textarea:focus{border-color:var(--accent)!important;box-shadow:0 0 0 2px rgba(94,255,176,.12)!important;}
+input[type=range]{accent-color:var(--accent)!important;}
+button.primary,.gr-button-primary{background:linear-gradient(135deg,var(--accent) 0%,#38d9a9 100%)!important;color:#0a0c10!important;font-family:var(--font-head)!important;font-weight:700!important;border:none!important;border-radius:8px!important;padding:12px 28px!important;font-size:.95rem!important;transition:opacity .2s!important;}
+button.primary:hover{opacity:.85!important;}
+button.secondary,.gr-button-secondary{background:transparent!important;border:1px solid var(--border)!important;color:var(--muted)!important;font-family:var(--font-head)!important;border-radius:8px!important;}
+label,.gr-form>label,.block>label span{color:var(--muted)!important;font-family:var(--font-mono)!important;font-size:.8rem!important;letter-spacing:.04em!important;text-transform:uppercase!important;}
+#log-box textarea{font-size:.78rem!important;color:#7af0b0!important;background:#060709!important;}
+.info-note{background:rgba(167,139,250,.07);border:1px solid rgba(167,139,250,.2);border-radius:8px;padding:12px 16px;font-family:var(--font-mono);font-size:.8rem;color:#c4b5fd;line-height:1.6;}
 """
 
 
 def build_app():
+    initial_status = (
+        "✅ Checkpoints ready."
+        if _ckpts_ready else
+        "⚠️  Checkpoints not found locally. Click 'Download Checkpoints' or they will be fetched automatically on first inference."
+    )
+
     with gr.Blocks() as demo:
 
-        # ── Header ──────────────────────────────────
         gr.HTML("""
         <div id="header">
           <h1>✦ Lyra 2.0</h1>
@@ -393,25 +332,30 @@ def build_app():
             <span class="badge">3D Gaussian Splatting</span>
             <span class="badge">arXiv 2604.13036</span>
           </p>
-          <p style="margin-top:14px; color:#8892a4; font-size:0.9rem; font-family:'Syne',sans-serif;">
-            Generate persistent, explorable 3D worlds from a single image.
-            Walk through scenes, revisit areas — no spatial forgetting, no temporal drift.
+          <p style="margin-top:14px;color:#8892a4;font-size:.9rem;font-family:'Syne',sans-serif;">
+            Generate persistent, explorable 3D worlds from a single image —
+            no spatial forgetting, no temporal drift.
           </p>
         </div>
         """)
 
-        # ── Tabs ─────────────────────────────────────
+        # ── Checkpoint status bar ────────────────────
+        with gr.Row():
+            ckpt_status = gr.Textbox(
+                value=initial_status, label="Checkpoint Status",
+                interactive=False, scale=4,
+            )
+            ckpt_btn = gr.Button("⬇️ Download Checkpoints", variant="secondary", scale=1)
+        ckpt_btn.click(fn=manual_download, outputs=ckpt_status)
+
         with gr.Tabs():
 
-            # ════════════════════════════════════════
-            # TAB 1 — Zoom-in / Zoom-out
-            # ════════════════════════════════════════
+            # ── Tab 1: Zoom Trajectory ───────────────
             with gr.Tab("🔭 Zoom Trajectory"):
                 gr.HTML('<div class="info-note">Generate a zoom-in → zoom-out exploration video from a single image, then optionally lift it to a 3D Gaussian Splatting scene.</div>')
-
                 with gr.Row():
                     with gr.Column(scale=1):
-                        z_image = gr.Image(label="Input Image", type="numpy", height=280)
+                        z_image   = gr.Image(label="Input Image", type="numpy", height=280)
                         z_caption = gr.Textbox(
                             label="Scene Caption",
                             placeholder="A sunlit forest clearing with tall pine trees…",
@@ -419,39 +363,35 @@ def build_app():
                         )
                         with gr.Accordion("Advanced Options", open=False):
                             with gr.Row():
-                                z_in_str   = gr.Slider(0.1, 3.0, value=0.5,  step=0.1, label="Zoom-in Strength")
-                                z_out_str  = gr.Slider(0.1, 3.0, value=1.5,  step=0.1, label="Zoom-out Strength")
+                                z_in_str  = gr.Slider(0.1, 3.0, value=0.5,  step=0.1, label="Zoom-in Strength")
+                                z_out_str = gr.Slider(0.1, 3.0, value=1.5,  step=0.1, label="Zoom-out Strength")
                             with gr.Row():
-                                z_frames_in  = gr.Slider(81, 401, value=81,  step=80, label="Frames Zoom-in  (1+80k)")
+                                z_frames_in  = gr.Slider(81, 401, value=81,  step=80, label="Frames Zoom-in (1+80k)")
                                 z_frames_out = gr.Slider(81, 401, value=241, step=80, label="Frames Zoom-out (1+80k)")
                             with gr.Row():
-                                z_dmd    = gr.Checkbox(label="⚡ Fast Mode (DMD ×15 speedup, lower quality)", value=False)
-                                z_recon  = gr.Checkbox(label="🧊 Run 3DGS Reconstruction after video", value=True)
+                                z_dmd   = gr.Checkbox(label="⚡ Fast Mode (DMD ×15 speedup, lower quality)", value=False)
+                                z_recon = gr.Checkbox(label="🧊 Run 3DGS Reconstruction after video", value=True)
                         z_btn = gr.Button("Generate World", variant="primary")
-
                     with gr.Column(scale=1):
                         z_video  = gr.Video(label="Generated Exploration Video", height=280)
-                        z_gs_vid = gr.Video(label="3DGS Flythrough (if reconstruction enabled)", height=280)
-                        z_log    = gr.Textbox(label="Log", lines=6, interactive=False, elem_id="log-box")
+                        z_gs_vid = gr.Video(label="3DGS Flythrough", height=280)
+                        z_log    = gr.Textbox(label="Log", lines=8, interactive=False, elem_id="log-box")
 
                 z_btn.click(
                     fn=run_zoomgs,
-                    inputs=[z_image, z_caption, gr.State(0),
+                    inputs=[z_image, z_caption,
                             z_in_str, z_out_str, z_frames_in, z_frames_out,
                             z_dmd, z_recon],
                     outputs=[z_video, z_gs_vid, z_log],
                 )
 
-            # ════════════════════════════════════════
-            # TAB 2 — Custom Trajectory
-            # ════════════════════════════════════════
+            # ── Tab 2: Custom Trajectory ─────────────
             with gr.Tab("🎮 Custom Trajectory"):
-                gr.HTML('<div class="info-note">Provide your own camera trajectory (.npz with <code>w2c</code>, <code>intrinsics</code>, <code>image_height</code>, <code>image_width</code>) and per-chunk captions (JSON keyed by frame index, e.g. <code>{"0": "…", "81": "…"}</code>).</div>')
-
+                gr.HTML('<div class="info-note">Provide a custom camera trajectory (.npz with <code>w2c</code>, <code>intrinsics</code>, <code>image_height</code>, <code>image_width</code>) and per-chunk captions as JSON keyed by frame index.</div>')
                 with gr.Row():
                     with gr.Column(scale=1):
-                        c_image   = gr.Image(label="First Frame", type="numpy", height=240)
-                        c_traj    = gr.File(label="Trajectory (.npz)", file_types=[".npz"])
+                        c_image    = gr.Image(label="First Frame", type="numpy", height=240)
+                        c_traj     = gr.File(label="Trajectory (.npz)", file_types=[".npz"])
                         c_captions = gr.Textbox(
                             label='Per-chunk Captions (JSON or single string)',
                             placeholder='{"0": "A grand hall interior", "81": "Corridor leading outside"}',
@@ -465,11 +405,10 @@ def build_app():
                                 c_dmd   = gr.Checkbox(label="⚡ Fast Mode (DMD)", value=False)
                                 c_recon = gr.Checkbox(label="🧊 Run 3DGS Reconstruction", value=True)
                         c_btn = gr.Button("Generate World", variant="primary")
-
                     with gr.Column(scale=1):
                         c_video  = gr.Video(label="Generated Video", height=260)
                         c_gs_vid = gr.Video(label="3DGS Flythrough", height=260)
-                        c_log    = gr.Textbox(label="Log", lines=6, interactive=False, elem_id="log-box")
+                        c_log    = gr.Textbox(label="Log", lines=8, interactive=False, elem_id="log-box")
 
                 c_btn.click(
                     fn=run_custom_traj,
@@ -478,9 +417,7 @@ def build_app():
                     outputs=[c_video, c_gs_vid, c_log],
                 )
 
-            # ════════════════════════════════════════
-            # TAB 3 — Model Info
-            # ════════════════════════════════════════
+            # ── Tab 3: About ─────────────────────────
             with gr.Tab("ℹ️ About"):
                 gr.Markdown("""
 ## Lyra 2.0 — Explorable Generative 3D Worlds
@@ -488,45 +425,34 @@ def build_app():
 **NVIDIA Research** · [Paper](https://arxiv.org/abs/2604.13036) · [Project Page](https://research.nvidia.com/labs/sil/projects/lyra2/) · [HuggingFace](https://huggingface.co/nvidia/Lyra-2.0)
 
 ### How it works
-
-Lyra 2.0 solves two fundamental failure modes of long-horizon 3D world generation:
-
 | Problem | Solution |
 |---|---|
-| **Spatial Forgetting** — previously seen regions fall out of context and are hallucinated on revisit | Per-frame 3D geometry used for information routing — retrieve past frames and establish dense correspondences |
-| **Temporal Drifting** — autoregressive errors accumulate and distort appearance/geometry | Self-augmented training histories expose the model to its own degraded outputs, teaching correction not propagation |
+| **Spatial Forgetting** — revisited regions hallucinated when outside temporal context | Per-frame 3D geometry used for routing: retrieve past frames + dense correspondences |
+| **Temporal Drifting** — autoregressive errors accumulate over long trajectories | Self-augmented histories teach the model to correct drift rather than propagate it |
 
-The generated video is then lifted to a **3D Gaussian Splatting** scene via VIPE pose estimation + Depth Anything 3 depth.
-
-### GPU Requirements
-
-- Recommended: **H100 80 GB** (or A100 80 GB)  
-- ~9 min per 80 frames at full quality · ~35 s with `--use_dmd` (DMD fast mode)
-- GS reconstruction adds ~1 min on top
-
-### Checkpoint Setup
-
-Checkpoints are expected at `./checkpoints/model/`.  
-Download from HuggingFace:
-
-```bash
-huggingface-cli download nvidia/Lyra-2.0 \\
-  --include "checkpoints/*" \\
-  --local-dir .
+### Checkpoint layout
 ```
+checkpoints/
+├── model/                        ← main generation model weights
+└── text_encoder/
+    └── negative_prompt.pt        ← required for CFG negative guidance
+```
+Checkpoints are auto-downloaded (~50 GB) from `nvidia/Lyra-2.0` on HuggingFace at startup.
 
-### Citation
+### GPU requirements
+- Recommended: **H100 80 GB** or A100 80 GB
+- ~9 min / 80 frames full quality · ~35 s with DMD fast mode
+- GS reconstruction adds ~1 min
 
+### Citation
 ```bibtex
 @article{shen2026lyra2,
-    title={Lyra 2.0: Explorable Generative 3D Worlds},
-    author={Shen, Tianchang and Bahmani, Sherwin and He, Kai and ...},
-    journal={arXiv preprint arXiv:2604.13036},
-    year={2026}
+  title   = {Lyra 2.0: Explorable Generative 3D Worlds},
+  author  = {Shen, Tianchang and Bahmani, Sherwin and He, Kai and others},
+  journal = {arXiv preprint arXiv:2604.13036},
+  year    = {2026}
 }
 ```
-
-*Model weights released under [NVIDIA Internal Scientific Research and Development Model License](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-internal-scientific-research-and-development-model-license/).*
                 """)
 
     return demo
@@ -534,4 +460,4 @@ huggingface-cli download nvidia/Lyra-2.0 \\
 
 if __name__ == "__main__":
     demo = build_app()
-    demo.launch(css=CSS)
+    demo.launch(css=CSS, ssr_mode=False)