Update app.py

app.py

@@ -1,264 +1,524 @@
-import argparse
-import gc
 import os
-import sys
+import json
+import tempfile
+import subprocess
+import shutil
+from pathlib import Path
 
 import gradio as gr
-import torch
 import spaces
-from omegaconf import OmegaConf
-
-# Handle cubvh dynamic installation to avoid build isolation errors on HF Spaces
-try:
-    import cubvh
-except ImportError:
-    import subprocess
-    print("Installing cubvh...")
-    subprocess.check_call([sys.executable, "-m", "pip", "install", "git+https://github.com/ashawkey/cubvh.git", "--no-build-isolation"])
-
-# Add project root to path so we can run from anywhere
-sys.path.append(os.path.dirname(os.path.abspath(__file__)))
-
-from ultrashape.rembg import BackgroundRemover
-from ultrashape.utils.misc import instantiate_from_config
-from ultrashape.surface_loaders import SharpEdgeSurfaceLoader
-from ultrashape.utils import voxelize_from_point
-from ultrashape.pipelines import UltraShapePipeline
-
-# Global variables to cache the model
-MODEL_CACHE = {}
-
-
-def get_pipeline_cached(config_path, ckpt_path, device='cuda', low_vram=False):
-    # Check if we have a valid cached pipeline for this checkpoint
-    if "pipeline" in MODEL_CACHE and MODEL_CACHE.get("ckpt_path") == ckpt_path:
-        print("Using cached pipeline...")
-        return MODEL_CACHE["pipeline"], MODEL_CACHE["config"]
-
-    # Clear old cache if it exists (e.g. different checkpoint)
-    if MODEL_CACHE:
-        print("Clearing old model cache...")
-        MODEL_CACHE.clear()
-        gc.collect()
-        torch.cuda.empty_cache()
-
-    print(f"Loading config from {config_path}...")
-    config = OmegaConf.load(config_path)
-
-    print("Instantiating VAE...")
-    vae = instantiate_from_config(config.model.params.vae_config)
-
-    print("Instantiating DiT...")
-    dit = instantiate_from_config(config.model.params.dit_cfg)
-
-    print("Instantiating Conditioner...")
-    conditioner = instantiate_from_config(config.model.params.conditioner_config)
-
-    print("Instantiating Scheduler & Processor...")
-    scheduler = instantiate_from_config(config.model.params.scheduler_cfg)
-    image_processor = instantiate_from_config(config.model.params.image_processor_cfg)
-
-    print(f"Loading weights from {ckpt_path}...")
-    weights = torch.load(ckpt_path, map_location='cpu')
-
-    vae.load_state_dict(weights['vae'], strict=True)
-    dit.load_state_dict(weights['dit'], strict=True)
-    conditioner.load_state_dict(weights['conditioner'], strict=True)
-
-    vae.eval().to(device)
-    dit.eval().to(device)
-    conditioner.eval().to(device)
-
-    if hasattr(vae, 'enable_flashvdm_decoder'):
-        vae.enable_flashvdm_decoder()
-
-    print("Creating Pipeline...")
-    pipeline = UltraShapePipeline(
-        vae=vae,
-        model=dit,
-        scheduler=scheduler,
-        conditioner=conditioner,
-        image_processor=image_processor
-    )
+import numpy as np
+import torch
+
+
+# ──────────────────────────────────────────────
+# Lazy model loading (done inside GPU decorator)
+# ──────────────────────────────────────────────
+_model_loaded = False
+
+
+def _ensure_models():
+    global _model_loaded
+    if _model_loaded:
+        return
+    # Models are expected to be pre-downloaded to ./checkpoints/
+    # on the Space via the HF repo or a setup script.
+    _model_loaded = True
+
+
+# ──────────────────────────────────────────────
+# Core inference helpers
+# ──────────────────────────────────────────────
+
+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)
 
-    if low_vram:
-        pipeline.enable_model_cpu_offload()
-
-    MODEL_CACHE["pipeline"] = pipeline
-    MODEL_CACHE["config"] = config
-    MODEL_CACHE["ckpt_path"] = ckpt_path
-
-    return pipeline, config
-
-
-@spaces.GPU
-def predict(
-        image_input,
-        mesh_input,
-        steps,
-        scale,
-        octree_res,
-        num_latents,
-        chunk_size,
-        seed,
-        remove_bg,
-        ckpt_path,
-        low_vram
+
+def _run(cmd: list[str], desc: str = "") -> tuple[bool, str]:
+    """Run a shell command and return (success, stderr/stdout)."""
+    print(f"[Lyra] {desc}: {' '.join(cmd)}")
+    result = subprocess.run(
+        cmd,
+        capture_output=True,
+        text=True,
+        env={**os.environ, "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"},
+    )
+    log = result.stdout + "\n" + result.stderr
+    return result.returncode == 0, log
+
+
+# ─────────────��────────────────────────────────
+# Zoom-in / Zoom-out trajectory (Option 1)
+# ──────────────────────────────────────────────
+
+@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,
 ):
-    # Aggressive memory cleanup at start
-    gc.collect()
-    torch.cuda.empty_cache()
-
-    try:
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        
-        # Resolve config path relative to this script
-        base_dir = os.path.dirname(os.path.abspath(__file__))
-        config_path = os.path.join(base_dir, "configs", "infer_dit_refine.yaml")
-
-        if not os.path.exists(config_path):
-            raise gr.Error(f"Config not found at {config_path}")
-            
-        if not os.path.exists(ckpt_path):
-            raise gr.Error(f"Checkpoint not found at '{ckpt_path}'. Please provide a valid checkpoint path or download the model weights.")
-
-        pipeline, config = get_pipeline_cached(config_path, ckpt_path, device, low_vram)
-
-        voxel_res = config.model.params.vae_config.params.voxel_query_res
-
-        print(f"Initializing Surface Loader (Token Num: {num_latents})...")
-        loader = SharpEdgeSurfaceLoader(
-            num_sharp_points=204800,
-            num_uniform_points=204800,
+    _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
+        Image.fromarray(image).save(img_path)
+        caption_path.write_text(caption.strip() or "A scenic outdoor environment.")
+
+        output_dir = Path(tmp) / "outputs" / "zoomgs"
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        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
+            "--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),
+        ]
+        if use_dmd:
+            cmd.append("--use_dmd")
+
+        ok, log = _run(["env", "PYTHONPATH=."] + cmd[1:] if cmd[0] == "python" else cmd,
+                       "ZoomGS video generation")
+
+        # Locate output video
+        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():
+            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])
+
+        return (
+            str(video_path) if video_path and video_path.exists() else None,
+            gs_video,
+            log[-4000:],
         )
 
-        print(f"Processing inputs...")
-        if image_input is None:
-            raise gr.Error("Image input is required")
-        if mesh_input is None:
-            raise gr.Error("Mesh input is required")
-
-        image = image_input.convert("RGBA")
-
-        if remove_bg or image.mode != 'RGBA':
-            rembg = BackgroundRemover()
-            image = rembg(image)
-
-        # Handle mesh input - Gradio Model3D returns path to file
-        surface = loader(mesh_input, normalize_scale=scale).to(device, dtype=torch.float16)
-        pc = surface[:, :, :3]  # [B, N, 3]
-
-        # Voxelize
-        _, voxel_idx = voxelize_from_point(pc, num_latents, resolution=voxel_res)
-
-        print("Running diffusion process...")
-        gen_device = "cpu" if low_vram else device
-        generator = torch.Generator(gen_device).manual_seed(int(seed))
-
-        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
-            mesh_out_list, _ = pipeline(
-                image=image,
-                voxel_cond=voxel_idx,
-                generator=generator,
-                box_v=1.0,
-                mc_level=0.0,
-                octree_resolution=int(octree_res),
-                num_chunks=int(chunk_size),
-                num_inference_steps=int(steps)
-            )
 
-        # Save output
-        output_dir = os.path.join(base_dir, "outputs_gradio")
-        os.makedirs(output_dir, exist_ok=True)
-        base_name = "output"
-        save_path = os.path.join(output_dir, f"{base_name}_refined.glb")
-
-        mesh_out = mesh_out_list[0]
-        mesh_out.export(save_path)
-        print(f"Successfully saved to {save_path}")
-
-        return save_path
-
-    except Exception as e:
-        import traceback
-        traceback.print_exc()
-        raise gr.Error(str(e))
-    finally:
-        # Aggressive memory cleanup at end
-        gc.collect()
-        torch.cuda.empty_cache()
-
-
-def main():
-    parser = argparse.ArgumentParser(description="UltraShape Gradio App")
-    parser.add_argument("--ckpt", type=str, default="checkpoints/ultrashape.pt", help="Path to split checkpoint (.pt)")
-    parser.add_argument("--share", action="store_true", help="Share the gradio app")
-    parser.add_argument("--low_vram", action="store_true", help="Optimize for low VRAM usage")
-
-    args = parser.parse_args()
-
-    # Define a clean & modern Gradio Theme
-    custom_theme = gr.themes.Soft(
-        primary_hue="blue",
-        secondary_hue="indigo",
-    ).set(
-        button_primary_background_fill="*primary_500",
-        button_primary_background_fill_hover="*primary_600",
-    )
+# ──────────────────────────────────────────────
+# Custom trajectory (Option 2)
+# ──────────────────────────────────────────────
 
-    # Define Gradio Interface
-    with gr.Blocks(title="UltraShape 1.0", theme=custom_theme) as demo:
-        gr.Markdown(
-            """
-            <div style="text-align: center;">
-                <h1>UltraShape 1.0</h1>
-                <h3>High-Fidelity 3D Shape Generation via Scalable Geometric Refinement</h3>
-                <p>Upload a reference image and a coarse mesh to generate a refined high-quality 3D shape.</p>
-            </div>
-            """
+@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,
+):
+    _ensure_models()
+
+    with tempfile.TemporaryDirectory() as tmp:
+        from PIL import Image
+        img_path = Path(tmp) / "first_frame.png"
+        Image.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
+            captions_path.write_text(captions_json)
+        except json.JSONDecodeError:
+            captions_path.write_text(json.dumps({"0": captions_json}))
+
+        output_dir = Path(tmp) / "outputs" / "custom"
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        cmd = [
+            "python", "-m", "lyra_2._src.inference.lyra2_custom_traj_inference",
+            "--input_image_path", str(img_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),
+        ]
+        if use_dmd:
+            cmd.append("--use_dmd")
+
+        ok, log = _run(cmd, "Custom trajectory video generation")
+
+        video_candidates = list(output_dir.rglob("*.mp4"))
+        video_path = video_candidates[0] if video_candidates else None
+
+        gs_video = None
+        if run_reconstruction and video_path:
+            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])
+
+        return (
+            str(video_path) if video_path else None,
+            gs_video,
+            log[-4000:],
         )
 
-        with gr.Row():
-            with gr.Column(scale=1):
-                gr.Markdown("### Input Parameters")
-                image_input = gr.Image(type="pil", label="Reference Image", image_mode="RGBA")
-                mesh_input = gr.Model3D(label="Coarse Mesh (.glb, .obj)")
-                
-                with gr.Accordion("Advanced Options", open=False):
-                    steps = gr.Slider(minimum=1, maximum=200, value=50, step=1, label="Inference Steps (12-50)")
-                    scale = gr.Slider(minimum=0.1, maximum=2.0, value=0.99, label="Mesh Normalization Scale")
-                    octree_res = gr.Slider(minimum=64, maximum=2048, value=1024, step=64, label="Octree Resolution")
-                    num_latents = gr.Slider(minimum=1024, maximum=32768, value=32768, step=128,
-                                            label="Number of Latent Tokens (Decrease if OOM)")
-                    chunk_size = gr.Slider(minimum=512, maximum=10000, value=2048, step=512,
-                                           label="Chunk Size (Decrease if OOM)")
-                    seed = gr.Number(value=42, label="Random Seed")
-                    remove_bg = gr.Checkbox(label="Remove Background from Image", value=False)
-
-                run_btn = gr.Button("Generate Refined Shape", variant="primary", size="lg")
-
-            with gr.Column(scale=1):
-                gr.Markdown("### Refined Output")
-                output_model = gr.Model3D(label="High-Fidelity Mesh", interactive=False)
-                
-        gr.Markdown(
-            """
-            *Note: If you encounter Out-of-Memory (OOM) errors, try checking the 'Advanced Options' and lowering the `Number of Latent Tokens` (e.g., 8192) and `Chunk Size` (e.g., 2000), or run the app with the `--low_vram` flag.*
-            """
-        )
 
-        run_btn.click(
-            fn=lambda img, mesh, s, sc, oct, nml, chk, sd, rm: predict(
-                img, mesh, s, sc, oct, nml, chk, sd, rm, args.ckpt, args.low_vram
-            ),
-            inputs=[
-                image_input, mesh_input, steps, scale, octree_res, 
-                num_latents, chunk_size, seed, remove_bg
-            ],
-            outputs=[output_model]
-        )
+# ──────────────────────────────────────────────
+# 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;
+}
+"""
+
+
+def build_app():
+    with gr.Blocks(css=CSS, title="Lyra 2.0 — Explorable 3D Worlds") as demo:
+
+        # ── Header ──────────────────────────────────
+        gr.HTML("""
+        <div id="header">
+          <h1>✦ Lyra 2.0</h1>
+          <p>
+            <span class="badge">NVIDIA Research</span>
+            <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>
+        </div>
+        """)
+
+        # ── Tabs ─────────────────────────────────────
+        with gr.Tabs():
+
+            # ════════════════════════════════════════
+            # TAB 1 — Zoom-in / Zoom-out
+            # ════════════════════════════════════════
+            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_caption = gr.Textbox(
+                            label="Scene Caption",
+                            placeholder="A sunlit forest clearing with tall pine trees…",
+                            lines=2,
+                        )
+                        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")
+                            with gr.Row():
+                                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_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_btn.click(
+                    fn=run_zoomgs,
+                    inputs=[z_image, z_caption, gr.State(0),
+                            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
+            # ════════════════════════════════════════
+            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>')
+
+                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_captions = gr.Textbox(
+                            label='Per-chunk Captions (JSON or single string)',
+                            placeholder='{"0": "A grand hall interior", "81": "Corridor leading outside"}',
+                            lines=3,
+                        )
+                        with gr.Accordion("Advanced Options", open=False):
+                            with gr.Row():
+                                c_frames     = gr.Slider(81, 961, value=481, step=80, label="Num Frames (1+80k)")
+                                c_pose_scale = gr.Slider(0.1, 5.0, value=1.0, step=0.1, label="Pose Scale")
+                            with gr.Row():
+                                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_btn.click(
+                    fn=run_custom_traj,
+                    inputs=[c_image, c_traj, c_captions,
+                            c_frames, c_pose_scale, c_dmd, c_recon],
+                    outputs=[c_video, c_gs_vid, c_log],
+                )
+
+            # ════════════════════════════════════════
+            # TAB 3 — Model Info
+            # ════════════════════════════════════════
+            with gr.Tab("ℹ️ About"):
+                gr.Markdown("""
+## Lyra 2.0 — Explorable Generative 3D Worlds
+
+**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 |
+
+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 .
+```
+
+### 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}
+}
+```
+
+*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
 
-    demo.launch(share=args.share, server_name='0.0.0.0', server_port=7860)
 
 if __name__ == "__main__":
-    main()
+    demo = build_app()
+    demo.launch()