app.py

"""Lift4D — text-prompted 4D reconstruction from a monocular video.

Runs the first two stages of the Lift4D pipeline
(https://github.com/yehonathanlitman/Lift4D) through the *custom dataset*
workflow documented in the repo README:

    stage 1  segment_video.py        SAM 3 text prompt -> per-frame masks
    stage 2  sam3d/run_inference.py  Causal SAM 3D Objects -> per-frame 3DGS
"""
import os
import shutil
import subprocess
import sys
import time
import uuid
from pathlib import Path

os.environ.setdefault("CUDA_HOME", "/usr/local/cuda")
os.environ.setdefault("CONDA_PREFIX", "/usr/local")
os.environ["LIDRA_SKIP_INIT"] = "true"
os.environ["ATTN_BACKEND"] = "sdpa"
os.environ["SPARSE_ATTN_BACKEND"] = "sdpa"
os.environ["SPARSE_BACKEND"] = "spconv"
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import spaces  # noqa: E402  (must precede torch)
import gradio as gr  # noqa: E402
import numpy as np  # noqa: E402
from huggingface_hub import login, snapshot_download  # noqa: E402

if os.environ.get("HF_TOKEN"):
    login(token=os.environ["HF_TOKEN"])

APP_ROOT = Path(__file__).parent
sys.path.insert(0, str(APP_ROOT / "kaolin_stub"))

# --- released 4D reconstructions for the interactive 4D-visualization viewer ---
# The client-side WebGL 4D Gaussian-splat player (static/viewer4d.html +
# static/js/gs4d_player.js) renders the authors' released deformable-Gaussian
# scenes from per-scene params under static/params/<scene>/.  It is embedded in
# an iframe alongside the live per-frame splat / orbit viewers below.
STATIC_DIR = APP_ROOT / "static"
gr.set_static_paths(paths=[STATIC_DIR])
_STATIC_PREFIX = "/gradio_api/file="  # Gradio's static-serve endpoint
VIEWER4D_URL = _STATIC_PREFIX + "static/viewer4d.html"
PARAMS4D_URL = _STATIC_PREFIX + "static/params"
# Custom-run 4D param bundles are written here so they are served by the same
# static endpoint (STATIC_DIR is whitelisted above) and picked up by the viewer.
USER_RUNS_DIR = STATIC_DIR / "user_runs"
USER_RUNS_DIR.mkdir(parents=True, exist_ok=True)
SCENES_4D = ["goat", "rhino", "horsejump-low", "ball", "bicycle", "bird", "robot", "yarn"]


def _viewer4d_iframe_src(src: str) -> str:
    return (
        f'<iframe src="{src}" title="Lift4D interactive 4D Gaussian-splat viewer" '
        'style="width:100%;height:520px;border:0;border-radius:12px;'
        'background:#0d0f12;" allow="fullscreen"></iframe>'
    )


def viewer4d_iframe(scene: str) -> str:
    """HTML embed of the client-side WebGL 4D Gaussian-splat player for a scene."""
    if scene not in SCENES_4D:
        scene = SCENES_4D[0]
    src = f"{VIEWER4D_URL}?embed=1&v=3&scene={scene}&params={PARAMS4D_URL}"
    return _viewer4d_iframe_src(src)


def user_run_viewer4d_iframe(params_url: str) -> str:
    """HTML embed of the 4D viewer pointed at a *custom run's* exported params.

    ``params_url`` is the static-serve URL of the bundle directory written by
    viewer4d_export.export_run (meta.json + *.bin live directly there).  The
    ``custom=1`` flag tells the viewer this is a single user-run bundle (no
    per-scene sub-dir) rather than the released example scenes.
    """
    # cache-bust with the unique run dir so the iframe reloads on each new run
    src = f"{VIEWER4D_URL}?embed=1&custom=1&scene=custom&params={params_url}&v={uuid.uuid4().hex[:8]}"
    return _viewer4d_iframe_src(src)


def load_4d_scene(scene: str) -> str:
    """Load a released 4D reconstruction into the interactive 4D viewer.

    Args:
        scene: scene name, one of the released reconstructions
            (goat, rhino, horsejump-low, ball, bicycle, bird, robot, yarn).

    Returns:
        HTML embedding the WebGL 4D Gaussian-splat player for that scene.
    """
    return viewer4d_iframe(scene)

LIFT4D_COMMIT = "aa295d66d1baabc6f9954b619f8b1bdb557bd862"
WORK = Path("/tmp/lift4d")
DATA_ROOT = WORK / "data"
SAM3D_OUT = WORK / "sam3d_output"
for d in (DATA_ROOT / "custom", SAM3D_OUT):
    d.mkdir(parents=True, exist_ok=True)
os.environ["LIFT4D_DATA_ROOT"] = str(DATA_ROOT)
os.environ["LIFT4D_SAM3D_OUT"] = str(SAM3D_OUT)


def _pip(*args):
    r = subprocess.run([sys.executable, "-m", "pip", "install", "--no-cache-dir", *args],
                       capture_output=True, text=True, timeout=1800)
    print(f"  pip {'OK' if r.returncode == 0 else 'FAIL'}: {args[-1][:80]}", flush=True)
    if r.returncode != 0:
        print(r.stderr[-1500:], flush=True)
    return r.returncode == 0


print("=== runtime installs (need torch present) ===", flush=True)
# utils3d pinned to the commit MoGe expects; MoGe supplies the geometry helpers
# sam3d_objects.pipeline.utils.pointmap imports directly.
_pip("--no-deps", "git+https://github.com/EasternJournalist/utils3d.git@3913c65d81e05e47b9f367250cf8c0f7462a0900")
_pip("--no-deps", "git+https://github.com/microsoft/MoGe.git@a8c37341bc0325ca99b9d57981cc3bb2bd3e255b")
# gaussian_render.py imports gsplat at module scope; only its python side is
# touched (rendering goes through the Inria rasterizer / pytorch3d).
_pip("--no-deps", "gsplat")

LIFT4D_DIR = APP_ROOT / "Lift4D"
if not LIFT4D_DIR.exists():
    print("Cloning Lift4D...", flush=True)
    subprocess.run(["git", "clone", "https://github.com/yehonathanlitman/Lift4D.git",
                    str(LIFT4D_DIR)], check=True)
    subprocess.run(["git", "checkout", "-q", LIFT4D_COMMIT], cwd=str(LIFT4D_DIR), check=True)

SAM3D_DIR = LIFT4D_DIR / "sam3d"
sys.path.insert(0, str(LIFT4D_DIR))     # lift4d_datasets
sys.path.insert(0, str(SAM3D_DIR))      # sam3d_objects
sys.path.insert(0, str(APP_ROOT))

print("Downloading SAM 3D Objects checkpoints (gated: facebook/sam-3d-objects)...", flush=True)
CKPT = Path(snapshot_download("facebook/sam-3d-objects", token=os.environ.get("HF_TOKEN")))
hf_ckpt = CKPT / "checkpoints"
local_ckpt = SAM3D_DIR / "checkpoints" / "hf"
if hf_ckpt.exists() and not local_ckpt.exists():
    local_ckpt.parent.mkdir(parents=True, exist_ok=True)
    local_ckpt.symlink_to(hf_ckpt)
CONFIG_PATH = str(local_ckpt / "pipeline.yaml")

print("Downloading SAM 3 (gated: facebook/sam3)...", flush=True)
try:
    snapshot_download("facebook/sam3", token=os.environ.get("HF_TOKEN"),
                      allow_patterns=["*.json", "*.txt", "*.safetensors", "*.model", "*.bin"])
except Exception as exc:  # pragma: no cover
    print(f"  sam3 prefetch failed ({exc}); will download on first run", flush=True)

print("=== startup complete ===", flush=True)

# Model construction needs a real GPU (the SAM3D constructors run device-mixing
# tensor ops that ZeroGPU's startup CUDA shim rejects), so both models are
# built inside @spaces.GPU and cached in the worker.
_SAM3 = None
_PIPE = None


def _get_sam3():
    global _SAM3
    if _SAM3 is None:
        from transformers import Sam3Model, Sam3Processor
        model = Sam3Model.from_pretrained("facebook/sam3").to("cuda").eval()
        processor = Sam3Processor.from_pretrained("facebook/sam3")
        _SAM3 = (model, processor)
    return _SAM3


def _get_pipeline():
    global _PIPE
    if _PIPE is None:
        import lift4d_stage12
        _PIPE = lift4d_stage12.build_pipeline(CONFIG_PATH)
        if getattr(_PIPE, "rendering_engine", "pytorch3d") != "pytorch3d":
            _PIPE.rendering_engine = "pytorch3d"
    return _PIPE


MAX_FRAMES = 16


def _estimate(*args, **kwargs):
    """ZeroGPU budget, sized from measurements on the live Space.

    Timed on the RTX PRO 6000 worker: ~35 s one-off SAM-3D pipeline build,
    then ~18.5 s/frame at stage2_steps=25 (4 frames -> 110 s cold,
    8 frames -> 149 s warm).  Per-frame cost is roughly affine in the
    structured-latent step count, so scale with it and keep ~20% headroom.
    """
    def _get(i, default):
        try:
            return int(args[i])
        except (IndexError, TypeError, ValueError):
            return default

    n = _get(2, 8)
    s2 = _get(6, 25)
    return int(min(660, 45 + 1.2 * n * (6.0 + 0.5 * s2)))


@spaces.GPU(duration=_estimate)
def reconstruct(video, prompt="", num_frames=8, consistency=0.2, seed=42,
                stage1_steps=50, stage2_steps=25, orbit_steps=48,
                progress=gr.Progress(track_tqdm=True)):
    import lift4d_datasets as ds_registry
    import lift4d_stage12 as stage12
    import orbit_render
    from PIL import Image

    if video is None:
        raise gr.Error("Please upload a video.")
    if not (prompt or "").strip():
        raise gr.Error("Please give a text prompt naming the object to reconstruct.")

    t0 = time.time()
    uid = f"clip_{uuid.uuid4().hex[:8]}"
    video_dir = DATA_ROOT / "custom" / uid
    frames_dir = video_dir / "frames"
    masks_dir = video_dir / "masks"

    # ---- stage 1: frames + SAM 3 masks in the README's custom layout --------
    progress(0.02, desc="Extracting frames")
    stems = stage12.extract_frames(video, frames_dir, int(num_frames))

    model, processor = _get_sam3()
    mask_name, masks = stage12.segment_frames(
        frames_dir, masks_dir, stems, prompt, model, processor,
        progress_cb=lambda f, d: progress(0.02 + 0.13 * f, desc=d))

    preview_dir = video_dir / "preview"
    preview_dir.mkdir(exist_ok=True)
    gallery = []
    for stem, m in zip(stems, masks):
        rgb = np.asarray(Image.open(frames_dir / f"{stem}.jpg").convert("RGB"))
        cut = (rgb * m[..., None] + 255 * (~m[..., None])).astype(np.uint8)
        p = preview_dir / f"{stem}.jpg"
        Image.fromarray(cut).save(p, quality=90)
        gallery.append(str(p))

    # ---- stage 2: causal SAM 3D Objects reconstruction ----------------------
    progress(0.18, desc="Loading SAM 3D Objects")
    pipe = _get_pipeline()
    spec = ds_registry.resolve("custom", uid, mask_name)

    progress(0.22, desc=f"Reconstructing {len(stems)} frames (causal SAM 3D)")
    saved, frame_images, out_base = stage12.run_stage2(
        pipe, spec, stems, mask_name,
        seed=int(seed), stage1_steps=int(stage1_steps), stage2_steps=int(stage2_steps),
        initial_frame_index=0, consistency_strength=float(consistency))
    t_recon = time.time() - t0

    # ---- placement, splat export, orbit render ------------------------------
    progress(0.85, desc="Exporting Gaussian splats")
    cam_frames, ply_paths, kept = [], [], []
    for rec in saved:
        gs = rec["gs"]
        if gs is None:
            continue
        packed = orbit_render.to_camera_space(
            gs, rec["scale"], rec["translation"], rec["rotation"], device="cuda")
        cam_frames.append(packed)
        kept.append(rec)

    if not cam_frames:
        raise gr.Error("SAM 3D Objects returned no Gaussians for this clip.")

    # Each frame is recentred on its own median so scrubbing the frame slider
    # keeps the subject framed instead of letting it walk out of the viewport.
    splat_dir = out_base / "splats"
    splat_dir.mkdir(exist_ok=True)
    for rec, packed in zip(kept, cam_frames):
        p = splat_dir / f"{rec['label']}.ply"
        orbit_render.write_viewer_ply(p, *packed, center=None)
        ply_paths.append(str(p))

    progress(0.92, desc="Rendering orbit video")
    imgs = orbit_render.render_orbit(cam_frames, size=512,
                                     orbit_steps=int(orbit_steps), device="cuda")
    orbit_path = out_base / "orbit.mp4"
    orbit_render.save_video(imgs, orbit_path, fps=12)

    # ---- populate the interactive 4D viewer with THIS run's output ----------
    # Export the per-frame reconstructed splats into the viewer's param format
    # (the same meta.json + *.bin bundle the released example scenes use) so the
    # 4D panel plays the user's actual reconstruction, not just the examples.
    progress(0.97, desc="Exporting 4D viewer params")
    import viewer4d_export
    run_params_dir = USER_RUNS_DIR / uid
    viewer4d_export.export_run(run_params_dir, cam_frames, fps=12)
    params_url = f"{_STATIC_PREFIX}static/user_runs/{uid}"
    viewer_html = user_run_viewer4d_iframe(params_url)

    total = time.time() - t0
    status = (
        f"**Done in {total:.0f}s** — {len(saved)} frames reconstructed "
        f"({t_recon:.0f}s in SAM 3D).\n\n"
        f"Custom dataset: `data/custom/{uid}/frames` + "
        f"`data/custom/{uid}/masks/<frame>_{mask_name}.png` → "
        f"stage-1 output tag `{spec.tag}`.\n\n"
        f"The **4D visualization** panel below now plays this run's reconstruction."
    )
    return str(orbit_path), ply_paths[0], gallery, status, str(splat_dir), viewer_html


def show_frame(splat_dir, idx):
    """Swap the Gaussian-splat viewer to another reconstructed frame (CPU)."""
    if not splat_dir:
        return gr.update()
    plys = sorted(Path(splat_dir).glob("*.ply"))
    if not plys:
        return gr.update()
    i = max(0, min(int(idx), len(plys) - 1))
    return gr.update(value=str(plys[i]))


CSS = """
.dark .gradio-container { background: #0b0b0d; }
#viewers .wrap { min-height: 420px; }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="Lift4D") as demo:
    gr.Markdown(
        "# Lift4D — 4D reconstruction from a casual video\n"
        "Name an object in your clip and Lift4D's front half runs: **SAM 3** "
        "segments it in every sampled frame, then **causal SAM 3D Objects** "
        "reconstructs a posed Gaussian splat per frame with SDEdit temporal "
        "warm-starting. Two viewers below: the per-frame **Gaussian splats** "
        "and a rendered **orbit video** of the reconstructed sequence.\n\n"
        "The **4D visualization** panel further down plays the authors' released "
        "deformable-Gaussian reconstructions in an interactive WebGL viewer "
        "(orbit / zoom / scrub time)."
    )

    with gr.Row():
        with gr.Column(scale=1):
            video = gr.Video(label="Input video", height=280)
            prompt = gr.Textbox(label="Object prompt", placeholder="e.g. goat",
                                value="")
            num_frames = gr.Slider(3, MAX_FRAMES, value=8, step=1,
                                   label="Frames to reconstruct",
                                   info="Sampled evenly across the clip. More frames = smoother 4D, longer run.")
            run = gr.Button("Reconstruct", variant="primary")
            with gr.Accordion("Advanced", open=False):
                consistency = gr.Slider(0.0, 1.0, value=0.2, step=0.05,
                                        label="Temporal consistency (SDEdit strength)",
                                        info="Higher preserves more structure between frames; lower allows more deformation.")
                seed = gr.Number(value=42, precision=0, label="Seed")
                stage1_steps = gr.Slider(10, 50, value=50, step=5,
                                         label="Sparse-structure steps")
                stage2_steps = gr.Slider(10, 50, value=25, step=5,
                                         label="Structured-latent steps")
                orbit_steps = gr.Slider(16, 96, value=48, step=8,
                                        label="Orbit video frames")

        with gr.Column(scale=2, elem_id="viewers"):
            with gr.Row():
                splat = gr.Model3D(label="Gaussian splats (per frame)",
                                   height=400, clear_color=[0.05, 0.05, 0.06, 1.0])
                orbit = gr.Video(label="Orbit video", height=400,
                                 autoplay=True, loop=True)
            frame_idx = gr.Slider(0, MAX_FRAMES - 1, value=0, step=1,
                                  label="Splat viewer: frame")
            status = gr.Markdown()
            gallery = gr.Gallery(label="SAM 3 segmentation (stage 1)", columns=8,
                                 height=130, object_fit="contain")
            splat_dir = gr.Textbox(visible=False)

    with gr.Accordion("4D visualization — interactive deformable-Gaussian viewer", open=True):
        gr.Markdown(
            "An interactive client-side WebGL Gaussian-splat viewer: drag to orbit, "
            "scroll to zoom, and scrub time to watch the Gaussians move. After you "
            "**Reconstruct** your own clip above, this panel plays *that* run's "
            "reconstruction. You can also browse the authors' released example "
            "**4D reconstructions** with the dropdown (**goat** is the repo README's "
            "in-the-wild custom-dataset clip)."
        )
        scene4d = gr.Dropdown(
            choices=SCENES_4D, value=SCENES_4D[0], label="Released example 4D scene",
            info="Released deformable-Gaussian reconstructions from the Lift4D authors.",
        )
        viewer4d = gr.HTML(viewer4d_iframe(SCENES_4D[0]))
        scene4d.change(load_4d_scene, inputs=[scene4d], outputs=[viewer4d],
                       api_name="load_4d_scene")

    inputs = [video, prompt, num_frames, consistency, seed,
              stage1_steps, stage2_steps, orbit_steps]
    outputs = [orbit, splat, gallery, status, splat_dir, viewer4d]

    run.click(reconstruct, inputs=inputs, outputs=outputs)
    frame_idx.change(show_frame, inputs=[splat_dir, frame_idx], outputs=[splat])

    gr.Examples(
        examples=[
            ["examples/horse_running.mp4", "horse"],
            ["examples/poodle_running.mp4", "dog"],
            ["examples/squirrel_eating.mp4", "squirrel"],
        ],
        inputs=[video, prompt],
        outputs=outputs,
        fn=reconstruct,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.queue(max_size=8).launch(mcp_server=True, show_error=True)