app.py

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  — must precede torch / CUDA-touching imports

import logging  # noqa: E402
import tempfile  # noqa: E402

import imageio  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
import gradio as gr  # noqa: E402
from omegaconf import OmegaConf  # noqa: E402
from PIL import Image  # noqa: E402

# torch>=2.6 defaults torch.load(weights_only=True). The Orbis-2 checkpoints are
# trusted upstream Lightning checkpoints; force weights_only=False so the main L1
# state-dict load succeeds (the sub-loaders already pass weights_only explicitly).
_orig_torch_load = torch.load
def _patched_torch_load(*args, **kwargs):  # noqa: ANN001
    kwargs.setdefault("weights_only", False)
    return _orig_torch_load(*args, **kwargs)
torch.load = _patched_torch_load

from huggingface_hub import snapshot_download  # noqa: E402

from data.l2_context import L2ContextMixin  # noqa: E402
from data.video_loaders import (  # noqa: E402
    ClipAugmenter,
    DecordFrameAdapter,
    ResizeCenterPolicy,
    TensorFrameAdapter,
)
from evaluate.utils import (  # noqa: E402
    compute_frame_interval,
    decode_video_frames,
    get_rollout_future_frame_count,
    get_video_fps_and_length,
    maybe_apply_condition_preprocessor_scales,
    overlay_trajectory_on_images,
    resolve_video_backend,
)
from util import instantiate_from_config  # noqa: E402

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("orbis2-demo")

STEERING_FORMAT = "speed_yawrate"
DEVICE = "cuda"

# ---------------------------------------------------------------------------
# Download checkpoints (tokenizer + L1 + distilled L2) at module scope.
# ---------------------------------------------------------------------------
MODELS_DIR = snapshot_download(
    repo_id="sud0301/orbis2",
    allow_patterns=[
        "tok/**",
        "L1/config*.yaml",
        "L1/checkpoints/last.ckpt",
        "L2_distilled/**",
    ],
)
os.environ["ORBIS2_MODELS_DIR"] = MODELS_DIR
logger.info("Checkpoints downloaded to %s", MODELS_DIR)

# Use the distilled L2 config: the distilled abstract predictor runs with very
# few sampler steps (l2_pred_NFE=4), which is what the authors recommend for
# fast inference.
CONFIG_PATH = os.path.join(MODELS_DIR, "L1", "config_distill.yaml")
CKPT_PATH = os.path.join(MODELS_DIR, "L1", "checkpoints", "last.ckpt")


def _build_model():
    config = OmegaConf.load(CONFIG_PATH)
    model = instantiate_from_config(config.model)
    ckpt_result = model.load_state_dict(torch.load(CKPT_PATH)["state_dict"], strict=False)
    exempt = tuple(getattr(model, "checkpoint_exempt_key_prefixes", ()))
    unexpected_missing = [k for k in ckpt_result.missing_keys if not k.startswith(exempt)]
    assert unexpected_missing == [], unexpected_missing
    model = model.to(DEVICE)
    model.eval()
    return model, config


logger.info("Building Orbis-2 hierarchical world model ...")
MODEL, CONFIG = _build_model()

# Resolve inference-time constants from the model / config.
SIZE = OmegaConf.select(CONFIG, "data.params.validation.params.size")
HEIGHT, WIDTH = (int(SIZE[0]), int(SIZE[1])) if not isinstance(SIZE, int) else (int(SIZE), int(SIZE))
L1_FRAME_RATE = float(OmegaConf.select(CONFIG, "data.params.validation.params.frame_rate"))
L2_FRAME_RATE = float(MODEL.condition_preprocessor.l2_predictor_frame_rate)
L1_CONTEXT_FRAMES = int(MODEL.vit.num_context_frames)
L2_CONTEXT_FRAMES = int(MODEL.condition_preprocessor.num_context_frames)
NUM_PRED_FRAMES = int(MODEL.num_pred_frames)
VIDEO_BACKEND = resolve_video_backend()

logger.info(
    "Model ready: size=%dx%d L1_rate=%g L2_rate=%g L1_ctx=%d L2_ctx=%d pred_frames=%d backend=%s",
    HEIGHT, WIDTH, L1_FRAME_RATE, L2_FRAME_RATE, L1_CONTEXT_FRAMES, L2_CONTEXT_FRAMES,
    NUM_PRED_FRAMES, VIDEO_BACKEND,
)


class _L1L2FrameIndexer(L2ContextMixin):
    """Computes L1/L2 frame indices into a single source video."""

    def __init__(self, frame_interval, stored_data_frame_rate, num_l2_context, l2_frame_rate, l1_context_frames):
        self.frame_interval = frame_interval
        self.stored_data_frame_rate = stored_data_frame_rate
        self._init_l2_context(
            num_l2_context=num_l2_context,
            l2_frame_rate=l2_frame_rate,
            l1_context_frames=l1_context_frames,
            require_l2_context=True,
        )


def _load_context(video_path):
    """Sample L1 (high-rate) and L2 (low-rate, further back) context windows from one video."""
    native_fps, video_length = get_video_fps_and_length(video_path, VIDEO_BACKEND)
    frame_interval = compute_frame_interval(native_fps, L1_FRAME_RATE, "the L1 frame rate")
    compute_frame_interval(native_fps, L2_FRAME_RATE, "the L2 predictor's trained frame rate")

    indexer = _L1L2FrameIndexer(
        frame_interval=frame_interval,
        stored_data_frame_rate=native_fps,
        num_l2_context=L2_CONTEXT_FRAMES,
        l2_frame_rate=L2_FRAME_RATE,
        l1_context_frames=L1_CONTEXT_FRAMES,
    )

    l1_span = (L1_CONTEXT_FRAMES - 1) * frame_interval + 1
    start_frame = video_length - l1_span  # latest window that fits
    required_offset = indexer.get_required_l1_start_offset()
    if start_frame < required_offset:
        raise gr.Error(
            "Video is too short: the world model needs a longer history of driving footage "
            f"(at least ~{(required_offset + l1_span) / native_fps:.1f}s at {native_fps:g} fps) "
            "to build its long-range (L2) context. Use one of the example clips or a longer video."
        )

    l1_indices, l2_indices = indexer.get_l1_and_l2_indices(start_frame, L1_CONTEXT_FRAMES)
    all_frames = decode_video_frames(video_path, l1_indices + l2_indices, VIDEO_BACKEND)

    adapter = DecordFrameAdapter() if VIDEO_BACKEND == "decord" else TensorFrameAdapter()
    augmenter = ClipAugmenter(adapter, ResizeCenterPolicy((HEIGHT, WIDTH)))
    all_tensor = augmenter(all_frames)  # [F, C, H, W] in [-1, 1]

    l1_tensor = all_tensor[: len(l1_indices)].unsqueeze(0).to(DEVICE)
    l2_tensor = all_tensor[len(l1_indices):].unsqueeze(0).to(DEVICE)
    return l1_tensor, l2_tensor


def _frames_to_video(gen_frames, fps):
    """Save a [T, C, H, W] tensor in [-1, 1] to a temporary mp4 and return the path."""
    frames = ((gen_frames.clamp(-1, 1) + 1.0) / 2.0 * 255.0).round().to(torch.uint8)
    frames = frames.permute(0, 2, 3, 1).cpu().numpy()  # [T, H, W, C]
    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    tmp.close()
    imageio.mimsave(tmp.name, list(frames), fps=fps, codec="libx264", quality=8)
    return tmp.name


def _estimate_duration(video_path, num_rollout_steps, speed_scale, yaw_rate_scale, l1_nfe, *args, **kwargs):
    # Measured on ZeroGPU (RTX PRO 6000): ~18s for 10 steps, ~29s for 20 steps
    # (warm). Add headroom for a cold worker streaming ~6 GB of weights.
    steps = int(num_rollout_steps) if num_rollout_steps else 10
    return min(90, 20 + int(steps * 2))


@spaces.GPU(duration=_estimate_duration)
@torch.no_grad()
def rollout(video_path, num_rollout_steps=10, speed_scale=1.0, yaw_rate_scale=1.0, l1_nfe=30,
            progress=gr.Progress(track_tqdm=True)):
    """Roll out the Orbis-2 driving world model from an input driving video.

    Given a short clip of front-camera driving footage, the model samples its
    high-rate (L1) and low-rate (L2) context windows from the video and then
    autoregressively predicts future frames.

    Args:
        video_path: path to an input driving video (front camera, >= ~4s history).
        num_rollout_steps: number of rollout steps; each step predicts several future frames.
        speed_scale: multiplicative factor on the ego speed conditioning (counterfactual "drive faster/slower").
        yaw_rate_scale: multiplicative factor on the ego yaw-rate conditioning (counterfactual "turn more/less").
        l1_nfe: number of sampler steps for the L1 detail predictor.

    Returns:
        A path to the predicted future-driving video (mp4).
    """
    if video_path is None:
        raise gr.Error("Please provide an input driving video.")

    num_rollout_steps = int(num_rollout_steps)
    l1_nfe = int(l1_nfe)

    maybe_apply_condition_preprocessor_scales(MODEL, float(speed_scale), float(yaw_rate_scale))

    l1_tensor, l2_tensor = _load_context(video_path)

    num_future_frames = get_rollout_future_frame_count(MODEL, num_rollout_steps)
    frame_rate = torch.tensor(L1_FRAME_RATE, device=DEVICE)
    data_batch = {"images": l1_tensor, "l2_context": l2_tensor, "frame_rate": frame_rate}

    # Unconditional steering: all-NaN placeholder = the framework's "no steering data" signal.
    get_required_steps = getattr(MODEL.condition_preprocessor, "get_required_rollout_odometry_steps", None)
    min_odo_steps = None
    if callable(get_required_steps):
        min_odo_steps = get_required_steps(
            validation_params=None,
            num_condition_frames=L1_CONTEXT_FRAMES,
            num_gen_frames=num_future_frames,
            rollout_steps=num_rollout_steps,
        )
    if min_odo_steps is None:
        raise gr.Error("Model did not report a required odometry length; cannot roll out.")
    data_batch["steering"] = torch.full(
        (1, int(min_odo_steps), 2), float("nan"), dtype=l1_tensor.dtype, device=DEVICE
    )
    data_batch["steering_format"] = STEERING_FORMAT

    condition_kwargs = MODEL.condition_preprocessor.get_condition_kwargs_from_batch(data_batch, split="rollout")

    with torch.autocast(dtype=torch.float16, device_type="cuda", enabled=True):
        _latents, gen_frames = MODEL.roll_out(
            x_0={"images": l1_tensor},
            num_gen_frames=num_rollout_steps,
            latent_input=False,
            NFE=l1_nfe,
            eta=0.0,
            sample_with_ema=True,
            num_samples=l1_tensor.size(0),
            frame_rate=frame_rate.reshape(1).repeat(l1_tensor.size(0)),
            condition_kwargs=condition_kwargs,
            decode_device="cpu",
            num_condition_frames=l1_tensor.size(1),
        )

    # gen_frames: [B, T, C, H, W] in [-1, 1]; take the single sample in the batch.
    out = _frames_to_video(gen_frames[0], fps=L1_FRAME_RATE)
    del _latents, condition_kwargs, l1_tensor, l2_tensor, gen_frames
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return out


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

INTRO = """
# 🚗 Orbis 2 — Hierarchical World Model for Driving

Orbis 2 is a **driving world model**: give it a short clip of front-camera driving
footage and it autoregressively **predicts the future**. A frozen low-frame-rate
**L2** predictor supplies abstract long-range context while the high-frame-rate
**L1** detail predictor generates the future frames.

Use the counterfactual **speed** / **yaw-rate** scales in *Advanced settings* to
nudge the imagined future (e.g. drive faster, or turn harder).

[Paper](https://arxiv.org/abs/2607.15898) · [Project page](https://lmb-freiburg.github.io/orbis2.github.io/) · [Code](https://github.com/lmb-freiburg/orbis2) · [Model](https://huggingface.co/sud0301/orbis2)
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(INTRO)
        with gr.Row():
            with gr.Column():
                video_in = gr.Video(label="Input driving video (front camera)", height=320)
                run = gr.Button("Predict the future", variant="primary")
            video_out = gr.Video(label="Predicted future rollout", height=320, autoplay=True)

        with gr.Accordion("Advanced settings", open=False):
            num_rollout_steps = gr.Slider(
                1, 20, value=10, step=1, label="Rollout steps",
                info="Each step predicts several future frames (~0.5s of video per step).",
            )
            speed_scale = gr.Slider(
                0.0, 2.0, value=1.0, step=0.1, label="Speed scale (counterfactual)",
                info="Multiplier on ego speed conditioning.",
            )
            yaw_rate_scale = gr.Slider(
                -2.0, 2.0, value=1.0, step=0.1, label="Yaw-rate scale (counterfactual)",
                info="Multiplier on ego yaw-rate (turning) conditioning.",
            )
            l1_nfe = gr.Slider(
                4, 50, value=30, step=1, label="L1 sampler steps (NFE)",
                info="More steps = higher detail, slower.",
            )

        gr.Examples(
            examples=[
                ["examples/highway_real.mp4"],
                ["examples/urban_real.mp4"],
                ["examples/urban_turn_real.mp4"],
            ],
            inputs=[video_in],
            outputs=video_out,
            fn=rollout,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        rollout,
        inputs=[video_in, num_rollout_steps, speed_scale, yaw_rate_scale, l1_nfe],
        outputs=video_out,
        api_name="rollout",
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)