# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Run LingBot-World 2.0 as an in-process realtime AR-Diffusion session. This example exercises the internal session API. It is not a public HTTP or WebSocket protocol. Each JSONL event advances the world by one three-latent frame AR block and writes the latent plus its identity metadata. Session-owned causal VAE encoder state produces each condition block without retaining or repeating a full image-condition horizon. """ from __future__ import annotations import argparse import asyncio import json import math import time from collections.abc import Sequence from pathlib import Path from typing import Any _MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers" _CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1" _FRAMES_PER_BLOCK = 3 def _camera_event_data(frames: list[list[str]]) -> dict[str, Any]: """Build the event-side script consumed by LingBotCameraControlReducer.""" return {"mode": "script", "frames": frames} def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description="Run an in-process realtime LingBot-World 2.0 session.") parser.add_argument("--model", default=_MODEL, help="Hugging Face model ID or local checkpoint path.") parser.add_argument("--image", required=True, help="Initial RGB image.") parser.add_argument("--prompt", required=True, help="Initial scene prompt.") parser.add_argument( "--events", required=True, help="JSONL file with one prompt/action event per AR block.", ) parser.add_argument("--output-dir", required=True, help="Directory for chunk latents and metadata.") parser.add_argument("--session-id", default="lingbot-world", help="Persistent world session identifier.") parser.add_argument("--height", type=int, default=480) parser.add_argument("--width", type=int, default=832) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--gpu-memory-fraction", type=float, default=0.1) parser.add_argument("--tensor-parallel-size", type=int, default=1) parser.add_argument("--enforce-eager", action="store_true") return parser.parse_args(argv) def _load_events(path: Path) -> list[dict[str, Any]]: events: list[dict[str, Any]] = [] for line_number, raw_line in enumerate(path.read_text().splitlines(), start=1): line = raw_line.strip() if not line: continue try: value = json.loads(line) except json.JSONDecodeError as exc: raise ValueError(f"events line {line_number} is not valid JSON: {exc.msg}.") from None if not isinstance(value, dict): raise ValueError(f"events line {line_number} must be a JSON object.") event_id = value.get("event_id") if isinstance(event_id, bool) or not isinstance(event_id, int) or event_id < 0: raise ValueError(f"events line {line_number} requires a non-negative integer event_id.") prompt = value.get("prompt") if prompt is not None and (not isinstance(prompt, str) or not prompt.strip()): raise ValueError(f"events line {line_number} prompt must be a non-empty string.") frames = value.get("frames") if frames is not None: if not isinstance(frames, list) or len(frames) != _FRAMES_PER_BLOCK: raise ValueError(f"events line {line_number} frames must contain exactly three lists.") for frame in frames: if not isinstance(frame, list) or any(not isinstance(action, str) for action in frame): raise ValueError(f"events line {line_number} frames must contain only action strings.") if prompt is None and frames is None: raise ValueError(f"events line {line_number} must update prompt and/or frames.") events.append({"event_id": event_id, "prompt": prompt, "frames": frames}) if not events: raise ValueError("events file must contain at least one event.") event_ids = [event["event_id"] for event in events] if event_ids != sorted(set(event_ids)): raise ValueError("event_id values must be unique and strictly increasing.") return events def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path]: image = Path(args.image).expanduser().resolve() events = Path(args.events).expanduser().resolve() output_dir = Path(args.output_dir).expanduser().resolve() if not image.is_file(): raise ValueError("--image must point to an existing file.") if not events.is_file(): raise ValueError("--events must point to an existing JSONL file.") if not args.prompt.strip(): raise ValueError("--prompt must contain non-whitespace text.") if args.height <= 0 or args.width <= 0 or args.height % 16 or args.width % 16: raise ValueError("--height and --width must be positive multiples of 16.") if args.tensor_parallel_size <= 0: raise ValueError("--tensor-parallel-size must be positive.") if not math.isfinite(args.gpu_memory_fraction) or not 0 < args.gpu_memory_fraction <= 1: raise ValueError("--gpu-memory-fraction must be in (0, 1].") return image, events, output_dir async def run(argv: Sequence[str] | None = None) -> Path: args = parse_args(argv) image, events_path, output_dir = _validate_args(args) events = _load_events(events_path) output_dir.mkdir(parents=True, exist_ok=True) # Imports stay below pure input validation so --help and helper tests do not # require a CUDA-enabled vLLM installation. import torch from vllm_omni.diffusion.models.lingbot_world.actions import LingBotCameraControlReducer from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer from vllm_omni.experimental.ar_diffusion.session import ( ARDiffusionSessionEvent, ARDiffusionSessionManager, ARDiffusionWorkerLifecycle, ) from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionControlInput from vllm_omni.inputs.data import OmniDiffusionSamplingParams engine = AsyncOmni( model=args.model, engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine", enforce_eager=args.enforce_eager, tensor_parallel_size=args.tensor_parallel_size, max_num_seqs=1, model_config={ "ar_diffusion_height": args.height, "ar_diffusion_width": args.width, "ar_diffusion_kv_config": { "gpu_memory_fraction": args.gpu_memory_fraction, "warmup_cudagraph": True, }, }, ) sampling = OmniDiffusionSamplingParams( height=args.height, width=args.width, num_frames=9, max_sequence_length=512, seed=args.seed, output_type="latent", extra_args={"flow_shift": 5.0}, ) consumer = ARDiffusionOmniTickConsumer( engine, prompt_provider=lambda tick: { "prompt": tick.prompt, "multi_modal_data": {"image": str(image)}, }, sampling_params_list=[sampling], diffusion_stage_id=0, ) manager = ARDiffusionSessionManager( tick_consumer=consumer, lifecycle=ARDiffusionWorkerLifecycle(engine, stage_ids=[0], timeout=180.0), max_pending_events=32, control_reducer_factory=LingBotCameraControlReducer, ) session = await manager.create_session(args.session_id) measurements: list[dict[str, Any]] = [] try: for chunk_index, event in enumerate(events): controls = () if event["frames"] is not None: controls = ( ARDiffusionControlInput( track="camera", schema=_CAMERA_ACTION_SCHEMA, data=_camera_event_data(event["frames"]), ), ) await session.accept_event( ARDiffusionSessionEvent( event_id=event["event_id"], prompt=event["prompt"] if event["prompt"] is not None else (args.prompt if chunk_index == 0 else None), controls=controls, ) ) started = time.perf_counter() output = await session.next_chunk() elapsed = time.perf_counter() - started if len(output.images) != 1 or not isinstance(output.images[0], torch.Tensor): raise RuntimeError("Expected one latent tensor from each realtime LingBot chunk.") latent = output.images[0].detach().float().cpu() metadata = consumer.chunk_metadata(output).to_dict() torch.save(latent, output_dir / f"chunk_{chunk_index:03d}.pt") (output_dir / f"chunk_{chunk_index:03d}.json").write_text( json.dumps(metadata, indent=2, sort_keys=True) + "\n" ) measurements.append( { "chunk_index": chunk_index, "latency_seconds": elapsed, "shape": list(latent.shape), "finite": bool(torch.isfinite(latent).all()), "metadata": metadata, } ) print(json.dumps(measurements[-1], sort_keys=True), flush=True) finally: try: await manager.close_session(args.session_id) finally: engine.shutdown() summary_path = output_dir / "summary.json" summary_path.write_text(json.dumps({"chunks": measurements}, indent=2, sort_keys=True) + "\n") return summary_path def main(argv: Sequence[str] | None = None) -> Path: return asyncio.run(run(argv)) if __name__ == "__main__": main()