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)