diff --git a/flashdreams/flashdreams/serving/output_targets.py b/flashdreams/flashdreams/serving/output_targets.py index b7e7be98..f5306dbb 100644 --- a/flashdreams/flashdreams/serving/output_targets.py +++ b/flashdreams/flashdreams/serving/output_targets.py @@ -154,7 +154,8 @@ def _lingbot_webrtc_spec( options: OutputLaunchOptions, ) -> OutputTargetSpec: argv = [ - "--config_name", + "webrtc", + "--preset-id", _pipeline_name(config), "--device", _device(config), @@ -166,15 +167,20 @@ def _lingbot_webrtc_spec( str(getattr(config, "pixel_width", 832)), ] if _compile_network(config) is False: - argv.append("--no_compile") + argv.append("--no-compile") example_idx = getattr(config, "example_idx", None) if example_idx is not None: argv.extend(("--example-idx", str(example_idx))) - _append_webrtc_bind_args(argv, options) + if options.host: + argv.extend(("--host", options.host)) + if options.port is not None: + argv.extend(("--port", str(options.port))) + if options.prefer_sw_encoder: + argv.append("--prefer-sw-encoder") return OutputTargetSpec( mode="webrtc", - label="LingBot WebRTC server", - module="lingbot.webrtc.server", + label="LingBot shared demo WebRTC server", + module="lingbot.demo.cli", argv=tuple(argv), ) diff --git a/flashdreams/tests/test_output_targets.py b/flashdreams/tests/test_output_targets.py index 171e0478..be4e11fb 100644 --- a/flashdreams/tests/test_output_targets.py +++ b/flashdreams/tests/test_output_targets.py @@ -71,9 +71,10 @@ def test_lingbot_webrtc_target_translates_runner_config() -> None: ), ) - assert spec.module == "lingbot.webrtc.server" + assert spec.module == "lingbot.demo.cli" assert spec.argv == ( - "--config_name", + "webrtc", + "--preset-id", "lingbot-world-fast", "--device", "cuda:1", @@ -83,14 +84,14 @@ def test_lingbot_webrtc_target_translates_runner_config() -> None: "480", "--video-width", "832", - "--no_compile", + "--no-compile", "--example-idx", "3", "--host", "127.0.0.1", "--port", "9010", - "--prefer_sw_encoder", + "--prefer-sw-encoder", ) diff --git a/integrations/lingbot/lingbot/demo/__init__.py b/integrations/lingbot/lingbot/demo/__init__.py new file mode 100644 index 00000000..a3da5783 --- /dev/null +++ b/integrations/lingbot/lingbot/demo/__init__.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Experimental Lingbot demo adapter built on ``flashdreams.runtime.demo``.""" + +from lingbot.demo.adapter import LingbotDemoAdapter +from lingbot.demo.spec import ( + DEFAULT_LINGBOT_PRESET, + LINGBOT_MODEL_ID, + LingbotReplayInputs, + LingbotWebRTCScenario, +) + +__all__ = [ + "DEFAULT_LINGBOT_PRESET", + "LINGBOT_MODEL_ID", + "LingbotDemoAdapter", + "LingbotReplayInputs", + "LingbotWebRTCScenario", +] diff --git a/integrations/lingbot/lingbot/demo/adapter.py b/integrations/lingbot/lingbot/demo/adapter.py new file mode 100644 index 00000000..354ea56b --- /dev/null +++ b/integrations/lingbot/lingbot/demo/adapter.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot adapter for the shared demo API.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Any + +from flashdreams.runtime import ( + InferenceConfig, + InputCanonicalizer, + UserInputCapability, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + PreparedScenario, + WebRTCOutputSpec, +) +from flashdreams.runtime.interfaces import InferenceRuntime +from lingbot.runtime import ( + LingbotModelAdapter, + LingbotReplayRuntime, + PipelineFactory, + build_lingbot_webrtc_runtime_config, + inference_input_from_replay_inputs, +) +from lingbot.input_mapping import ( + KeyboardToCameraCommand, + TextEventSelection, +) +from lingbot.webrtc.session import ( + LingbotInferenceRuntime, + LingbotRuntimeConfig, +) + +from .spec import ( + resolve_replay_inputs, + resolve_text_event_prompts, + resolve_user_input_events, + resolve_webrtc_scenario, +) +from .webrtc import ( + LingbotDemoWebRTCSessionManager, + create_lingbot_webrtc_app, +) + +ReplayRuntimeFactory = Callable[..., InferenceRuntime] +WebRTCRuntimeFactory = Callable[..., Any] + + +class LingbotDemoAdapter(LingbotModelAdapter): + """Model-owned Lingbot adapter consumed by shared demo launchers.""" + + def __init__( + self, + *, + replay_runtime_factory: ReplayRuntimeFactory = LingbotReplayRuntime, + webrtc_runtime_factory: WebRTCRuntimeFactory = LingbotInferenceRuntime, + pipeline_factory: PipelineFactory | None = None, + ) -> None: + super().__init__( + runtime_factory=replay_runtime_factory, + pipeline_factory=pipeline_factory, + ) + self._webrtc_runtime_factory = webrtc_runtime_factory + + def supported_input_modes(self) -> tuple[str, ...]: + return ("replay", "keyboard-driving") + + def supported_output_modes(self) -> tuple[str, ...]: + return ("mp4", "webrtc") + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + if spec.input_mode != "replay": + raise ValueError( + "Lingbot prepare_scenario currently supports only " + f"input_mode='replay', got {spec.input_mode!r}." + ) + if not isinstance(spec.output, Mp4OutputSpec): + raise ValueError("Lingbot replay demo currently requires MP4 output.") + + replay_inputs = resolve_replay_inputs( + spec.scenario, + default_prompt=self.default_replay_prompt(spec.config), + ) + text_event_prompts = resolve_text_event_prompts(spec.scenario) + user_inputs = resolve_user_input_events(spec.scenario) + if _camera_source(spec.scenario) == "events": + # Live control still needs the scenario's calibration, so the trace + # is loaded for its intrinsics and world scale and then discarded + # as a trajectory source. + trace = self.create_input_mapping(replay_inputs).camera_trace + mapping = self.create_live_input_mapping( + fps=replay_inputs.fps, + base_intrinsics=trace.intrinsics[0], + # A trace's world scale is derived from how far its poses + # travel, so a stationary example yields 0. Live control has no + # trajectory to normalize against, so it falls back to the same + # unit scale the WebRTC runtime uses. + world_scale=trace.world_scale or 1.0, + prompt=replay_inputs.prompt, + text_event_prompts=text_event_prompts, + ) + else: + mapping = self.create_input_mapping( + replay_inputs, + text_event_prompts=text_event_prompts, + ) + return PreparedScenario( + initial_inputs=inference_input_from_replay_inputs(replay_inputs), + user_inputs=user_inputs, + source_schema=_source_schema(user_inputs), + canonicalizer=_canonicalizer(text_event_prompts), + mapping=mapping, + metadata={ + "model_id": self.model_id, + "preset_id": self.preset_id(spec.config), + }, + ) + + def create_webrtc_runtime(self, spec: DemoSpec) -> Any: + runtime_config = self.create_webrtc_runtime_config(spec=spec, runtime=None) + return self._webrtc_runtime_factory(config=runtime_config) + + def create_webrtc_runtime_config( + self, + *, + spec: DemoSpec, + runtime: Any, + ) -> LingbotRuntimeConfig: + runtime_config = getattr(runtime, "config", None) + if isinstance(runtime_config, LingbotRuntimeConfig): + return runtime_config + if spec.input_mode != "keyboard-driving": + raise ValueError( + "Lingbot WebRTC requires input_mode='keyboard-driving', " + f"got {spec.input_mode!r}." + ) + if not isinstance(spec.output, WebRTCOutputSpec): + raise ValueError("Lingbot WebRTC requires WebRTC output.") + config = spec.config + if config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + self.validate_config(config) + scenario = resolve_webrtc_scenario(spec.scenario) + + compile_network = ( + bool(config.compile) + if config.compile is not None + else bool(_option(config, "compile_network", True)) + ) + return build_lingbot_webrtc_runtime_config( + preset_id=self.preset_id(config), + pipeline_config=self.pipeline_config(config), + seed=int(_option(config, "seed", 42)), + compile_network=compile_network, + context_parallel_size=int(_option(config, "context_parallel_size", 1)), + device=config.device or str(_option(config, "device", "cuda:0")), + video_height=spec.output.video_height, + video_width=spec.output.video_width, + fps=spec.output.fps, + warmup_chunks=spec.output.warmup_chunks, + warmup_timeout_s=spec.output.warmup_timeout_s, + example_idx=int(_option(config, "example_idx", scenario.example_idx)), + prefer_sw_encoder=scenario.prefer_sw_encoder, + runtime_options=config.runtime_options, + ) + + def create_webrtc_session_manager( + self, + *, + spec: DemoSpec, + runtime: Any, + runtime_config: LingbotRuntimeConfig, + fps: int, + client_liveness_timeout_s: float, + ) -> LingbotDemoWebRTCSessionManager: + del spec + return LingbotDemoWebRTCSessionManager( + runtime=runtime, + runtime_config=runtime_config, + fps=fps, + client_liveness_timeout_s=client_liveness_timeout_s, + ) + + def create_webrtc_app( + self, + *, + spec: DemoSpec, + session_manager: Any, + request_session_url: str, + ) -> Any: + return create_lingbot_webrtc_app( + spec=spec, + session_manager=session_manager, + request_session_url=request_session_url, + ) + + +def _option(config: InferenceConfig, name: str, default: Any) -> Any: + return config.runtime_options.get(name, default) + + +def _camera_source(scenario: Any) -> str: + if isinstance(scenario, Mapping): + return str(scenario.get("camera_source", "trace")) + return "trace" + + +_KEY_EVENT_TYPES = frozenset({"key_down", "key_up"}) + +_KEYBOARD_CAPABILITIES = ( + UserInputCapability(event_type="key_down", payload_fields=frozenset({"key"})), + UserInputCapability(event_type="key_up", payload_fields=frozenset({"key"})), +) + +_TEXT_EVENT_CAPABILITY = UserInputCapability( + event_type="text_event", + payload_fields=frozenset({"event_id"}), +) + + +def _source_schema(user_inputs: UserInputs) -> UserInputSchema: + """Declare what this scenario's event source can provide. + + Capabilities describe the source, not the particular trace. A keyboard + source is declared to provide both key edges even if one recording happens + to contain no ``key_up`` -- a key held for the whole run is a normal trace. + Declaring only the observed types would fail the keyboard converter's + consumed set, and ``converters_for`` would silently drop it, leaving the run + with no camera control. + """ + observed = {event.event_type for event in user_inputs.events} + capabilities: list[UserInputCapability] = [] + if observed & _KEY_EVENT_TYPES: + capabilities.extend(_KEYBOARD_CAPABILITIES) + if "text_event" in observed: + capabilities.append(_TEXT_EVENT_CAPABILITY) + for event_type in sorted(observed - _KEY_EVENT_TYPES - {"text_event"}): + payload_fields: frozenset[str] = frozenset() + for event in user_inputs.events: + if event.event_type == event_type: + payload_fields = frozenset(event.payload) + break + capabilities.append( + UserInputCapability( + event_type=event_type, + payload_fields=payload_fields, + ) + ) + return UserInputSchema( + capabilities=tuple(capabilities), + description=( + "Lingbot replay event trace" + if capabilities + else "fixed Lingbot replay input" + ), + ) + + +def _canonicalizer(text_event_prompts: Mapping[str, str] | None) -> InputCanonicalizer: + converters: list[Any] = [KeyboardToCameraCommand()] + if text_event_prompts: + converters.append(TextEventSelection()) + return InputCanonicalizer(converters) + + +__all__ = [ + "LingbotDemoAdapter", + "ReplayRuntimeFactory", + "WebRTCRuntimeFactory", +] diff --git a/integrations/lingbot/lingbot/demo/cli.py b/integrations/lingbot/lingbot/demo/cli.py new file mode 100644 index 00000000..9b08469f --- /dev/null +++ b/integrations/lingbot/lingbot/demo/cli.py @@ -0,0 +1,233 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CLI for the experimental shared Lingbot demo path.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch +import torch.distributed as dist + +from flashdreams.core.distributed import init as distributed_init +from flashdreams.runtime import InferenceConfig +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + WebRTCOutputSpec, + run_flashdreams_demo, + serve_flashdreams_demo, +) +from flashdreams.serving.webrtc.bootstrap import ( + configure_logging, + initialize_cuda_distributed, +) +from lingbot.example_data import ( + EXAMPLE_DATA_AVAILABLE_IDXS, + ensure_example_data_downloaded, +) +from lingbot.runtime import ( + FIELD_CAMERA_INTRINSICS_PATH, + FIELD_CAMERA_POSES_PATH, + FIELD_FIRST_FRAME_PATH, + FIELD_FPS, + FIELD_PIXEL_HEIGHT, + FIELD_PIXEL_WIDTH, + FIELD_PROMPT, + FIELD_TOTAL_BLOCKS, +) + +from .adapter import LingbotDemoAdapter +from .spec import ( + DEFAULT_FPS, + DEFAULT_LINGBOT_PRESET, + DEFAULT_PIXEL_HEIGHT, + DEFAULT_PIXEL_WIDTH, + LINGBOT_MODEL_ID, + LingbotWebRTCScenario, +) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Experimental Lingbot demo using flashdreams.runtime.demo." + ) + subparsers = parser.add_subparsers(dest="command", required=True) + + replay = subparsers.add_parser("replay", help="Run an MP4 replay demo.") + replay.add_argument("--preset-id", "--config-name", default=DEFAULT_LINGBOT_PRESET) + replay.add_argument("--device", default="cuda") + replay.add_argument("--prompt", default=None) + replay.add_argument("--prompt-path", type=Path, default=None) + replay.add_argument("--image-path", type=Path, default=None) + replay.add_argument("--pose-path", type=Path, default=None) + replay.add_argument( + "--intrinsic-path", + "--intrinsics-path", + type=Path, + default=None, + ) + replay.add_argument( + "--example-data", + action=argparse.BooleanOptionalAction, + default=None, + help=( + "Use the bundled Lingbot example when asset paths are omitted " + "(default: auto)." + ), + ) + replay.add_argument( + "--example-idx", + "--example_idx", + type=int, + default=0, + choices=EXAMPLE_DATA_AVAILABLE_IDXS, + ) + replay.add_argument("--total-blocks", type=int, default=20) + replay.add_argument("--pixel-height", type=int, default=DEFAULT_PIXEL_HEIGHT) + replay.add_argument("--pixel-width", type=int, default=DEFAULT_PIXEL_WIDTH) + replay.add_argument("--fps", type=int, default=DEFAULT_FPS) + replay.add_argument("--output", type=Path, required=True) + + webrtc = subparsers.add_parser("webrtc", help="Serve a WebRTC driving demo.") + webrtc.add_argument("--preset-id", "--config-name", default=DEFAULT_LINGBOT_PRESET) + webrtc.add_argument("--host", default="0.0.0.0") + webrtc.add_argument("--port", type=int, default=8080) + webrtc.add_argument("--device", default="cuda:0") + webrtc.add_argument("--seed", type=int, default=42) + webrtc.add_argument( + "--compile", + action=argparse.BooleanOptionalAction, + default=True, + help="Enable or disable torch.compile for the Lingbot transformer.", + ) + webrtc.add_argument("--fps", type=int, default=DEFAULT_FPS) + webrtc.add_argument("--video-height", type=int, default=DEFAULT_PIXEL_HEIGHT) + webrtc.add_argument("--video-width", type=int, default=DEFAULT_PIXEL_WIDTH) + webrtc.add_argument("--warmup-chunks", type=int, default=10) + webrtc.add_argument("--warmup-timeout-s", type=float, default=600.0) + webrtc.add_argument("--client-liveness-timeout-s", type=float, default=30.0) + webrtc.add_argument("--prefer-sw-encoder", action="store_true") + webrtc.add_argument( + "--example-idx", + "--example_idx", + type=int, + default=0, + choices=EXAMPLE_DATA_AVAILABLE_IDXS, + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> None: + configure_logging() + args = parse_args(argv) + adapter = LingbotDemoAdapter() + if args.command == "replay": + run_flashdreams_demo(spec=_replay_spec(args), adapter=adapter) + return + if args.command == "webrtc": + context = initialize_cuda_distributed( + default_device=args.device, + distributed_init_fn=distributed_init, + configure_logging_fn=configure_logging, + torch_module=torch, + dist_module=dist, + ) + ensure_example_data_downloaded( + is_rank_zero=(context.world_rank == 0), + example_idx=args.example_idx, + ) + serve_flashdreams_demo( + spec=_webrtc_spec( + args, + device=str(context.device), + context_parallel_size=context.world_size, + ), + adapter=adapter, + world_rank=context.world_rank, + ) + return + raise AssertionError(f"Unhandled command: {args.command}") + + +def _replay_spec(args: argparse.Namespace) -> DemoSpec: + scenario: dict[str, object] = { + "example_data": args.example_data, + "example_idx": args.example_idx, + FIELD_TOTAL_BLOCKS: args.total_blocks, + FIELD_PIXEL_HEIGHT: args.pixel_height, + FIELD_PIXEL_WIDTH: args.pixel_width, + FIELD_FPS: args.fps, + } + if args.prompt: + scenario[FIELD_PROMPT] = args.prompt + if args.prompt_path is not None: + scenario["prompt_path"] = args.prompt_path + if args.image_path is not None: + scenario[FIELD_FIRST_FRAME_PATH] = args.image_path + if args.pose_path is not None: + scenario[FIELD_CAMERA_POSES_PATH] = args.pose_path + if args.intrinsic_path is not None: + scenario[FIELD_CAMERA_INTRINSICS_PATH] = args.intrinsic_path + + return DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=args.preset_id, + input_mode="replay", + scenario=scenario, + output=Mp4OutputSpec( + path=args.output, + fps=args.fps, + output_layout="tchw", + ), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=args.preset_id, + device=args.device, + ), + ) + + +def _webrtc_spec( + args: argparse.Namespace, + *, + device: str, + context_parallel_size: int = 1, +) -> DemoSpec: + return DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=args.preset_id, + input_mode="keyboard-driving", + scenario=LingbotWebRTCScenario( + example_idx=args.example_idx, + prefer_sw_encoder=args.prefer_sw_encoder, + ), + output=WebRTCOutputSpec( + host=args.host, + port=args.port, + fps=args.fps, + video_width=args.video_width, + video_height=args.video_height, + warmup_chunks=args.warmup_chunks, + warmup_timeout_s=args.warmup_timeout_s, + client_liveness_timeout_s=args.client_liveness_timeout_s, + preload_name="Lingbot", + ), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=args.preset_id, + device=device, + compile=args.compile, + runtime_options={ + "seed": args.seed, + "context_parallel_size": context_parallel_size, + "example_idx": args.example_idx, + }, + ), + ) + + +if __name__ == "__main__": + main() diff --git a/integrations/lingbot/lingbot/demo/replay.py b/integrations/lingbot/lingbot/demo/replay.py new file mode 100644 index 00000000..668ba63a --- /dev/null +++ b/integrations/lingbot/lingbot/demo/replay.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot replay runtime re-export for shared demo entry points.""" + +from __future__ import annotations + +from lingbot.runtime import ( + LingbotReplayRuntime, + LingbotReplayRuntimeOptions, + LingbotReplaySession, + PipelineFactory, +) + +__all__ = [ + "LingbotReplayRuntime", + "LingbotReplayRuntimeOptions", + "LingbotReplaySession", + "PipelineFactory", +] diff --git a/integrations/lingbot/lingbot/demo/spec.py b/integrations/lingbot/lingbot/demo/spec.py new file mode 100644 index 00000000..ef153fdb --- /dev/null +++ b/integrations/lingbot/lingbot/demo/spec.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot demo-specific input shapes.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from flashdreams.runtime import UserInputEvent, UserInputs +from lingbot.example_data import ( + EXAMPLE_DATA_AVAILABLE_IDXS, + EXAMPLE_DATA_BASE_URL, + EXAMPLE_DATA_DIR_LOCAL, + EXAMPLE_DATA_FILENAMES, + EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS, + example_asset_urls, + example_data_dirname, +) +from lingbot.runtime import ( + DEFAULT_FPS, + DEFAULT_LINGBOT_PRESET, + DEFAULT_PIXEL_HEIGHT, + DEFAULT_PIXEL_WIDTH, + LINGBOT_MODEL_ID, + LingbotReplayInputs, + replay_inputs_from_mapping, +) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class LingbotWebRTCScenario: + """Example-data and serving options for the shared WebRTC demo path.""" + + example_idx: int = 0 + prefer_sw_encoder: bool = False + + def __post_init__(self) -> None: + if self.example_idx not in EXAMPLE_DATA_AVAILABLE_IDXS: + raise ValueError( + "LingbotWebRTCScenario.example_idx must be one of " + f"{EXAMPLE_DATA_AVAILABLE_IDXS}." + ) + + +def resolve_replay_inputs( + value: Any, + *, + default_prompt: str = "", + is_rank_zero: bool = True, +) -> LingbotReplayInputs: + """Normalize a user/demo scenario into direct Lingbot runtime inputs.""" + return replay_inputs_from_mapping( + value, + default_prompt=default_prompt, + is_rank_zero=is_rank_zero, + ) + + +def resolve_text_event_prompts(value: Any) -> dict[str, str]: + """Return the scenario's text-event catalog as ``{event_id: prompt}``.""" + if not isinstance(value, Mapping): + return {} + catalog = value.get("text_events") + if not catalog: + return {} + if isinstance(catalog, Mapping): + return {str(key): str(prompt) for key, prompt in catalog.items()} + prompts: dict[str, str] = {} + for entry in catalog: + event_id = getattr(entry, "event_id", None) + prompt = getattr(entry, "prompt", None) + if event_id is None and isinstance(entry, Mapping): + event_id = entry.get("event_id") + prompt = entry.get("prompt") + if event_id is None: + raise ValueError("Lingbot text events require an 'event_id'.") + prompts[str(event_id)] = "" if prompt is None else str(prompt) + return prompts + + +def resolve_user_input_events(value: Any) -> UserInputs: + """Normalize a scenario's recorded event trace into :class:`UserInputs`. + + Each record is ``{"t": seconds, "type": event_type, ...payload}``, which + maps one-to-one onto ``UserInputEvent``. Events are sorted by timestamp + because ``UserInputs`` requires non-decreasing order. + """ + if not isinstance(value, Mapping): + return UserInputs() + records = value.get("events") + if not records: + return UserInputs() + + events: list[UserInputEvent] = [] + for record in records: + if isinstance(record, UserInputEvent): + events.append(record) + continue + if not isinstance(record, Mapping): + raise TypeError( + "Lingbot scenario events must be UserInputEvent objects or " + "mappings." + ) + payload = { + key: item + for key, item in record.items() + if key not in {"t", "timestamp_s", "type", "event_type", "source"} + } + timestamp_s = record.get("t", record.get("timestamp_s")) + event_type = record.get("type", record.get("event_type")) + if timestamp_s is None or event_type is None: + raise ValueError( + "Lingbot scenario events require a timestamp ('t') and a " + "type ('type')." + ) + events.append( + UserInputEvent( + timestamp_s=float(timestamp_s), + event_type=str(event_type), + payload=payload, + source=record.get("source"), + ) + ) + events.sort(key=lambda event: event.timestamp_s) + return UserInputs(events=tuple(events)) + + +def resolve_webrtc_scenario(value: Any) -> LingbotWebRTCScenario: + """Normalize a user/demo scenario into a WebRTC scenario.""" + if value is None: + return LingbotWebRTCScenario() + if isinstance(value, LingbotWebRTCScenario): + return value + if not isinstance(value, Mapping): + raise TypeError( + "Lingbot WebRTC scenario must be a LingbotWebRTCScenario, " + "a mapping, or None." + ) + return LingbotWebRTCScenario( + example_idx=int(value.get("example_idx", 0)), + prefer_sw_encoder=bool(value.get("prefer_sw_encoder", False)), + ) + + +__all__ = [ + "DEFAULT_FPS", + "DEFAULT_LINGBOT_PRESET", + "DEFAULT_PIXEL_HEIGHT", + "DEFAULT_PIXEL_WIDTH", + "EXAMPLE_DATA_AVAILABLE_IDXS", + "EXAMPLE_DATA_BASE_URL", + "EXAMPLE_DATA_DIR_LOCAL", + "EXAMPLE_DATA_FILENAMES", + "EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS", + "LINGBOT_MODEL_ID", + "LingbotReplayInputs", + "LingbotWebRTCScenario", + "example_asset_urls", + "example_data_dirname", + "resolve_replay_inputs", + "resolve_text_event_prompts", + "resolve_user_input_events", + "resolve_webrtc_scenario", +] diff --git a/integrations/lingbot/lingbot/demo/webrtc.py b/integrations/lingbot/lingbot/demo/webrtc.py new file mode 100644 index 00000000..966a79fb --- /dev/null +++ b/integrations/lingbot/lingbot/demo/webrtc.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot WebRTC hooks for the shared demo API.""" + +from __future__ import annotations + +from typing import Any + +from aiohttp import web + +from flashdreams.runtime.demo import DemoSpec +from lingbot.webrtc.server import create_app +from lingbot.webrtc.session import ( + LingbotInferenceRuntime, + LingbotRuntimeConfig, + LingbotWebRTCSessionManager, +) + + +class LingbotDemoWebRTCSessionManager(LingbotWebRTCSessionManager): + """Shared demo session manager using Lingbot's existing WebRTC semantics.""" + + def __init__( + self, + *, + runtime: LingbotInferenceRuntime, + runtime_config: LingbotRuntimeConfig, + fps: int, + client_liveness_timeout_s: float, + ) -> None: + super().__init__( + runtime=runtime, + runtime_config=runtime_config, + fps=fps, + client_liveness_timeout_s=client_liveness_timeout_s, + ) + + +def create_lingbot_webrtc_app( + *, + spec: DemoSpec, + session_manager: Any, + request_session_url: str, +) -> web.Application: + """Create the packaged Lingbot browser app through existing serving glue.""" + del spec + return create_app( + session_manager=session_manager, + request_session_url=request_session_url, + ) + + +__all__ = [ + "LingbotDemoWebRTCSessionManager", + "create_lingbot_webrtc_app", +] diff --git a/integrations/lingbot/lingbot/example_data.py b/integrations/lingbot/lingbot/example_data.py new file mode 100644 index 00000000..af33628a --- /dev/null +++ b/integrations/lingbot/lingbot/example_data.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bundled LingBot-World example-data helpers.""" + +from __future__ import annotations + +from pathlib import Path + +import torch + +from flashdreams.core.io.disk import default_flashdreams_cache_dir +from flashdreams.core.io.download import download_to_cache + +EXAMPLE_DATA_BASE_URL = ( + "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples" +) +"""HTTP base URL for the canonical examples shared by all LingBot versions.""" + +EXAMPLE_DATA_DIR_LOCAL = default_flashdreams_cache_dir() / "example_data/lingbot_world" +"""Local cache root where downloaded example folders are stored.""" + +EXAMPLE_DATA_FILENAMES = ( + "image.jpg", + "poses.npy", + "intrinsics.npy", + "prompt.txt", +) +"""Example assets downloaded when each file is available upstream.""" + +EXAMPLE_DATA_AVAILABLE_IDXS = (0, 1, 2, 3, 4, 5) +"""Supported upstream example indices currently hosted under ``examples/``.""" + +EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS = (0, 1, 2, 5) +"""Example indices that provide their own upstream ``prompt.txt`` file.""" + + +def example_data_dirname(example_idx: int) -> str: + """Format ``example_idx`` into the upstream folder naming convention.""" + assert example_idx in EXAMPLE_DATA_AVAILABLE_IDXS, ( + f"--example_idx must be one of {EXAMPLE_DATA_AVAILABLE_IDXS}." + ) + return f"{example_idx:02d}" + + +def example_asset_urls(example_idx: int) -> dict[str, str]: + """Return canonical upstream URLs for a Lingbot example.""" + dirname = example_data_dirname(example_idx) + return { + "image": f"{EXAMPLE_DATA_BASE_URL}/{dirname}/image.jpg", + "intrinsics": f"{EXAMPLE_DATA_BASE_URL}/{dirname}/intrinsics.npy", + "poses": f"{EXAMPLE_DATA_BASE_URL}/{dirname}/poses.npy", + } + + +def ensure_example_data_downloaded(*, is_rank_zero: bool, example_idx: int) -> Path: + """Download bundled GitHub example files on rank 0; barrier other ranks.""" + example_dirname = example_data_dirname(example_idx) + cache_dir = EXAMPLE_DATA_DIR_LOCAL / example_dirname + if is_rank_zero: + for filename in EXAMPLE_DATA_FILENAMES: + if ( + filename == "prompt.txt" + and example_idx not in EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS + ): + continue + download_to_cache( + f"{EXAMPLE_DATA_BASE_URL}/{example_dirname}/{filename}", + cache_dir=cache_dir, + filename=filename, + ) + if torch.distributed.is_initialized(): + torch.distributed.barrier() + return cache_dir + + +__all__ = [ + "EXAMPLE_DATA_AVAILABLE_IDXS", + "EXAMPLE_DATA_BASE_URL", + "EXAMPLE_DATA_DIR_LOCAL", + "EXAMPLE_DATA_FILENAMES", + "EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS", + "ensure_example_data_downloaded", + "example_asset_urls", + "example_data_dirname", +] diff --git a/integrations/lingbot/lingbot/input_mapping.py b/integrations/lingbot/lingbot/input_mapping.py new file mode 100644 index 00000000..9a05c217 --- /dev/null +++ b/integrations/lingbot/lingbot/input_mapping.py @@ -0,0 +1,670 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot user-event canonicalization and canonical-to-model input mapping. + +Lingbot's two live controls are a free camera driven from the keyboard and a +catalog of server-owned text events. This module carries both across the +``UserInputs -> CanonicalInputs -> InferenceInput`` boundary: + +- :data:`CAMERA_COMMAND` and :class:`KeyboardToCameraCommand` turn raw key + edges into device-independent camera intent; +- :data:`TEXT_EVENT` and :class:`TextEventSelection` track which text event is + active; +- :class:`LingbotInputMapping` turns that canonical intent into the per-step + camera trajectory the session consumes, and requests a session-global prompt + update when the active text event changes. + +The modalities live here rather than in ``flashdreams.runtime.canonical`` +because Lingbot is currently their only consumer. Both are plain +``CanonicalModality`` values, so lifting them into the shared canonical layer +later is a move plus an export, with no change to this mapping. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from flashdreams.runtime.canonical import DeviceConverterSchema +from flashdreams.runtime.inputs import ( + CanonicalInputs, + CanonicalInputSchema, + CanonicalModality, + InferenceInput, + InferenceInputSchema, + InputField, + TimeWindow, + UserInputCapability, + UserInputs, +) +from flashdreams.runtime.mapping import InputMappingSchema +from flashdreams.runtime.types import StepRequest +from flashdreams.serving.webrtc.controls import ( + CameraPoseIntegrator, + KeyboardState, + PoseSegment, +) +from flashdreams.serving.realtime.input import DEFAULT_SUPPORTED_KEYS + +FIELD_CAMERA_TRAJECTORY = "camera_trajectory" +FIELD_CAMERA_INTRINSICS = "camera_intrinsics" +FIELD_PROMPT = "prompt" +FIELD_WORLD_SCALE = "world_scale" +FIELD_TOTAL_CAMERA_FRAMES = "total_camera_frames" + +_PASSTHROUGH_GLOBAL_FIELDS: tuple[InputField, ...] = ( + InputField(name="first_frame_path", input_modality="image/path"), + InputField(name="total_blocks", input_modality="count"), + InputField(name="pixel_height", input_modality="pixel-height"), + InputField(name="pixel_width", input_modality="pixel-width"), + InputField(name="fps", input_modality="fps"), +) +"""App-owned session inputs this mapping forwards without interpreting them.""" + +_CLEAR_STATES = frozenset({"clear", "release", "off", "none"}) +_TRIGGER_STATES = frozenset({"trigger", "hold", "on"}) + +_AXES: tuple[str, ...] = ("move_forward", "move_right", "yaw", "pitch") + +_AXIS_KEYS: Mapping[str, tuple[str, str]] = { + # Axis -> (positive key, negative key) in CameraPoseIntegrator's vocabulary. + # Both directions of the keyboard/axis conversion are derived from this one + # table so a rebind cannot make them disagree. + "move_forward": ("w", "s"), + "move_right": ("e", "q"), + "yaw": ("a", "d"), + "pitch": ("i", "k"), +} + +_KEY_ALIASES: Mapping[str, str] = {"j": "a", "l": "d"} +"""Alternate yaw keys accepted by ``KeyboardState``, folded onto ``a``/``d``.""" + + +CAMERA_COMMAND = CanonicalModality( + name="camera_command", + payload_fields=frozenset({*_AXES, "segments"}), + description=( + "Free-camera intent. move_forward, move_right, yaw, and pitch are in " + "[-1, 1] and hold the level state at the end of the window. segments " + "carries the piecewise-constant timeline inside the window as " + "((start_s, end_s, axes), ...), so a consumer can integrate sub-window " + "timing instead of quantizing control to the chunk boundary." + ), +) + +TEXT_EVENT = CanonicalModality( + name="text_event", + payload_fields=frozenset({"event_id"}), + description=( + "Identifier of the active server-owned text event, or None when no " + "event is active. Level-triggered: the value is held until cleared." + ), +) + + +def _axes_from_keys(pressed: Iterable[str]) -> dict[str, float]: + """Return camera axis values for a resolved set of pressed keys.""" + keys = {_KEY_ALIASES.get(key, key) for key in pressed} + axes: dict[str, float] = {} + for axis, (positive, negative) in _AXIS_KEYS.items(): + value = 0.0 + if positive in keys: + value += 1.0 + if negative in keys: + value -= 1.0 + axes[axis] = value + return axes + + +def _keys_from_axes(axes: Mapping[str, float]) -> frozenset[str]: + """Return the integrator key set equivalent to ``axes``. + + Pose integration stays the single implementation in + :class:`CameraPoseIntegrator`, which is expressed over key sets. Converting + back here keeps live Lingbot trajectories identical to the WebRTC path + instead of forking the integration math. + """ + keys: set[str] = set() + for axis, (positive, negative) in _AXIS_KEYS.items(): + value = float(axes.get(axis, 0.0)) + if value > 0: + keys.add(positive) + elif value < 0: + keys.add(negative) + return frozenset(keys) + + +class KeyboardToCameraCommand: + """Convert keyboard edges into :data:`CAMERA_COMMAND` level state.""" + + def __init__( + self, + *, + name: str = "keyboard-to-camera-command", + supported_keys: frozenset[str] = DEFAULT_SUPPORTED_KEYS, + priority: int = 0, + ) -> None: + self._supported_keys = supported_keys + self._state = KeyboardState(supported_keys=supported_keys) + self._schema = DeviceConverterSchema( + name=name, + produces=CAMERA_COMMAND, + device_kind="keyboard", + priority=priority, + consumes=( + UserInputCapability( + event_type="key_down", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_type="key_up", + payload_fields=frozenset({"key"}), + ), + ), + ) + + @property + def schema(self) -> DeviceConverterSchema: + return self._schema + + def reset(self) -> None: + self._state = KeyboardState(supported_keys=self._supported_keys) + + def convert( + self, + user_inputs: UserInputs, + window: TimeWindow, + ) -> Mapping[str, Any] | None: + segments: list[tuple[float, float, dict[str, float]]] = [] + segment_start = window.start_s + axes = _axes_from_keys(self._state.resolved_effective_keys()) + + for event in user_inputs.events: + if event.event_type not in {"key_down", "key_up"}: + continue + key = event.payload.get("key") + if not isinstance(key, str): + continue + edge_t = min(max(float(event.timestamp_s), window.start_s), window.end_s) + if edge_t > segment_start: + segments.append((segment_start, edge_t, axes)) + segment_start = edge_t + self._state.apply_event( + event="keydown" if event.event_type == "key_down" else "keyup", + key=key, + ) + axes = _axes_from_keys(self._state.resolved_effective_keys()) + + if window.end_s > segment_start or not segments: + segments.append((segment_start, window.end_s, axes)) + + return CAMERA_COMMAND.value({**axes, "segments": tuple(segments)}) + + +class TextEventSelection: + """Track the active :data:`TEXT_EVENT` id across windows.""" + + def __init__( + self, + *, + name: str = "text-event-selection", + priority: int = 0, + ) -> None: + self._active_event_id: str | None = None + self._schema = DeviceConverterSchema( + name=name, + produces=TEXT_EVENT, + device_kind="text-event", + priority=priority, + consumes=( + UserInputCapability( + event_type="text_event", + payload_fields=frozenset({"event_id"}), + ), + ), + ) + + @property + def schema(self) -> DeviceConverterSchema: + return self._schema + + def reset(self) -> None: + self._active_event_id = None + + def convert( + self, + user_inputs: UserInputs, + window: TimeWindow, + ) -> Mapping[str, Any] | None: + del window + for event in user_inputs.events: + if event.event_type != "text_event": + continue + event_id = event.payload.get("event_id") + state = str(event.payload.get("state", "trigger")).strip().lower() + if state and state not in _CLEAR_STATES and state not in _TRIGGER_STATES: + raise ValueError( + f"Unsupported text event state {state!r}. Supported states: " + f"{sorted(_CLEAR_STATES | _TRIGGER_STATES)}." + ) + if event_id is None or state in _CLEAR_STATES: + self._active_event_id = None + continue + self._active_event_id = str(event_id) + return TEXT_EVENT.value({"event_id": self._active_event_id}) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class LingbotCameraTrace: + """Fixed camera trajectory resolved from a replay scenario. + + Tensors are CPU float32; the session owns device placement. + """ + + __hash__ = None + + poses: torch.Tensor + """Camera-to-world poses, shape ``[T, 4, 4]``.""" + + intrinsics: torch.Tensor + """Per-frame intrinsics, shape ``[T, 4]``, already rescaled to output size.""" + + world_scale: float + + def __post_init__(self) -> None: + if self.poses.ndim != 3 or self.poses.shape[1:] != (4, 4): + raise ValueError( + f"LingbotCameraTrace.poses must be [T, 4, 4], got " + f"{tuple(self.poses.shape)}." + ) + if self.intrinsics.ndim != 2 or self.intrinsics.shape[1] != 4: + raise ValueError( + f"LingbotCameraTrace.intrinsics must be [T, 4], got " + f"{tuple(self.intrinsics.shape)}." + ) + if self.world_scale < 0: + # Zero is legal: preprocess_example_poses derives the scale from + # pose spread, so a stationary trace yields 0. That reached the + # model before this mapping existed, so it still does. + raise ValueError("LingbotCameraTrace.world_scale must be >= 0.") + + @property + def frame_count(self) -> int: + return int(self.poses.shape[0]) + + +def load_camera_trace( + *, + camera_poses_path: str | Path, + camera_intrinsics_path: str | Path, + pixel_height: int, + pixel_width: int, + intrinsics_reference_height: int, + intrinsics_reference_width: int, + world_scale: float | None = None, +) -> LingbotCameraTrace: + """Load and preprocess a fixed Lingbot camera trajectory from ``.npy`` files.""" + from lingbot.encoder.utils import ( # noqa: PLC0415 + get_Ks_transformed, + preprocess_example_poses, + ) + + intrinsics = torch.from_numpy( + np.asarray(np.load(camera_intrinsics_path), dtype=np.float32) + ) + intrinsics = get_Ks_transformed( + intrinsics, + height_org=intrinsics_reference_height, + width_org=intrinsics_reference_width, + height_resize=pixel_height, + width_resize=pixel_width, + height_final=pixel_height, + width_final=pixel_width, + ) + poses, inferred_world_scale = preprocess_example_poses( + np.asarray(np.load(camera_poses_path)) + ) + return LingbotCameraTrace( + poses=torch.from_numpy(np.ascontiguousarray(poses)).to(torch.float32), + intrinsics=intrinsics.to(torch.float32), + world_scale=float( + inferred_world_scale if world_scale is None else world_scale + ), + ) + + +class LingbotInputMapping: + """Build Lingbot per-step camera inputs from canonical user input. + + Two trajectory sources are supported through one mapping object, because + ``run_inference_session`` takes a single mapping: + + - a fixed :class:`LingbotCameraTrace`, sliced per step, which consumes no + canonical modality and keeps MP4/benchmark runs deterministic; + - live :data:`CAMERA_COMMAND` intent integrated into a trajectory, for + event-driven runs. + + Text events are mapped to a session-global prompt update rather than a + per-step field: swapping the rollout's text context is session-global model + state, so it travels in the ``global_conditioning`` slot of the step + payload. Whether the model can apply that update is session-owned. + """ + + def __init__( + self, + *, + fps: int, + trace: LingbotCameraTrace | None = None, + base_intrinsics: torch.Tensor | Sequence[float] | None = None, + world_scale: float | None = None, + text_event_prompts: Mapping[str, str] | None = None, + integrator: CameraPoseIntegrator | None = None, + ) -> None: + if fps <= 0: + raise ValueError("LingbotInputMapping.fps must be > 0.") + if trace is None and base_intrinsics is None: + raise ValueError( + "LingbotInputMapping requires either a fixed camera trace or " + "base_intrinsics for live camera control." + ) + self._fps = int(fps) + self._trace = trace + self._text_event_prompts = dict(text_event_prompts or {}) + self._applied_event_id: str | None = None + self._base_prompt: str | None = None + + if trace is None: + intrinsics = torch.as_tensor(base_intrinsics, dtype=torch.float32).reshape( + 4 + ) + if world_scale is None or world_scale <= 0: + raise ValueError( + "Live Lingbot camera control requires a positive world_scale." + ) + self._base_intrinsics = intrinsics + self._world_scale = float(world_scale) + self._integrator = integrator or CameraPoseIntegrator() + else: + self._base_intrinsics = None + self._world_scale = trace.world_scale + self._integrator = None + + consumes: list[CanonicalModality] = [] + if trace is None: + consumes.append(CAMERA_COMMAND) + if self._text_event_prompts: + consumes.append(TEXT_EVENT) + self._mapping_schema = InputMappingSchema( + name="lingbot-input-mapping", + consumes=tuple(consumes), + produces_global_conditioning=( + # map_global_conditioning_inputs returns the app-owned session + # payload augmented with the fields below, so the pass-through + # fields are part of what this mapping produces. Declaring them + # keeps undeclared_inference_inputs() quiet and lets the + # compatibility check see that required session inputs are + # reachable; omitting them makes the check reject every run. + *_PASSTHROUGH_GLOBAL_FIELDS, + InputField( + name=FIELD_WORLD_SCALE, + required=False, + input_modality="scale", + frequency_consumed="once", + ), + InputField( + name=FIELD_TOTAL_CAMERA_FRAMES, + required=False, + input_modality="count", + frequency_consumed="once", + ), + InputField( + name=FIELD_PROMPT, + required=False, + input_modality="text", + frequency_consumed="once", + description="Text-event prompt update for an active rollout.", + ), + ), + produces_step=( + InputField( + name=FIELD_CAMERA_TRAJECTORY, + input_modality="c2w_sequence", + frequency_consumed="per_step", + metadata={"shape": "[T,4,4]", "frame": "camera_to_world"}, + ), + InputField( + name=FIELD_CAMERA_INTRINSICS, + input_modality="intrinsics_vec4_sequence", + frequency_consumed="per_step", + metadata={"shape": "[T,4]"}, + ), + ), + ) + + @property + def mapping_schema(self) -> InputMappingSchema: + return self._mapping_schema + + @property + def camera_trace(self) -> LingbotCameraTrace: + """Return the fixed trace, for callers reusing its calibration.""" + if self._trace is None: + raise ValueError("This Lingbot mapping has no fixed camera trace.") + return self._trace + + @property + def canonical_input_schema(self) -> CanonicalInputSchema: + """Return the modalities this mapping consumes, for adapter reporting.""" + return CanonicalInputSchema( + modalities=self._mapping_schema.consumes, + description="Lingbot live camera and text-event control.", + ) + + def validate( + self, + *, + canonical_schema: CanonicalInputSchema | None = None, + inference_input_schema: InferenceInputSchema | None = None, + ) -> None: + if canonical_schema is not None: + for modality in self._mapping_schema.consumes: + if not canonical_schema.supports(modality): + raise ValueError( + f"Lingbot input mapping requires canonical modality " + f"{modality.name!r}, which the selected input source " + f"cannot supply." + ) + if inference_input_schema is not None: + for name in (FIELD_CAMERA_TRAJECTORY, FIELD_CAMERA_INTRINSICS): + if inference_input_schema.field_for(name=name, phase="step") is None: + raise ValueError( + f"Lingbot input mapping produces step input {name!r}, " + f"which this model does not declare." + ) + + def map_global_conditioning_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + ) -> InferenceInput: + del canonical_inputs + payload = dict(inference_input.global_conditioning) + payload[FIELD_WORLD_SCALE] = self._world_scale + if self._trace is not None: + payload[FIELD_TOTAL_CAMERA_FRAMES] = self._trace.frame_count + return InferenceInput( + global_conditioning=payload, + step=inference_input.step, + metadata=inference_input.metadata, + ) + + def map_step_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + request: StepRequest, + ) -> InferenceInput: + num_frames = _required_int(request.metadata, "num_frames") + frame_start = _required_int(request.metadata, "frame_start") + + if self._trace is not None: + poses, intrinsics = self._slice_trace( + frame_start=frame_start, + num_frames=num_frames, + ) + else: + poses, intrinsics = self._integrate( + canonical_inputs=canonical_inputs, + request=request, + frame_start=frame_start, + num_frames=num_frames, + ) + + step = dict(inference_input.step) + step[FIELD_CAMERA_TRAJECTORY] = poses + step[FIELD_CAMERA_INTRINSICS] = intrinsics + return InferenceInput( + global_conditioning=self._text_event_update(canonical_inputs), + step=step, + metadata=inference_input.metadata, + ) + + def _slice_trace( + self, + *, + frame_start: int, + num_frames: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert self._trace is not None + frame_end = frame_start + num_frames + if frame_end > self._trace.frame_count: + raise ValueError( + f"Lingbot camera trace has {self._trace.frame_count} frames, but " + f"step needs frames [{frame_start}, {frame_end})." + ) + return ( + self._trace.poses[frame_start:frame_end], + self._trace.intrinsics[frame_start:frame_end], + ) + + def _integrate( + self, + *, + canonical_inputs: CanonicalInputs, + request: StepRequest, + frame_start: int, + num_frames: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert self._integrator is not None + assert self._base_intrinsics is not None + command = canonical_inputs.values.get(CAMERA_COMMAND.name) + if command is None: + raise ValueError( + "Lingbot live camera control requires a 'camera_command' " + "canonical value for every step; the selected input source " + "produced none." + ) + + window = request.user_input_window + start_s = window.start_s if window is not None else frame_start / self._fps + end_s = window.end_s if window is not None else ( + frame_start + num_frames + ) / self._fps + segments = _pose_segments(command, start_s=start_s, end_s=end_s) + frame_times = [start_s + (index + 1) / self._fps for index in range(num_frames)] + # The integrator rejects frame times outside the segment span, and float + # accumulation can leave the last one a hair past the window end. + frame_times[-1] = min(frame_times[-1], end_s) + + poses = self._integrator.integrate_chunk( + segments=segments, + frame_times=frame_times, + ) + poses_t = torch.from_numpy(np.ascontiguousarray(poses)).to(torch.float32) + poses_t = poses_t.reshape(num_frames, 4, 4) + intrinsics_t = self._base_intrinsics.reshape(1, 4).repeat(num_frames, 1) + return poses_t, intrinsics_t + + def _text_event_update( + self, + canonical_inputs: CanonicalInputs, + ) -> Mapping[str, Any]: + if not self._text_event_prompts: + return {} + value = canonical_inputs.values.get(TEXT_EVENT.name) + if value is None: + return {} + event_id = value.get("event_id") + if event_id == self._applied_event_id: + return {} + if event_id is not None and event_id not in self._text_event_prompts: + supported = ", ".join(sorted(self._text_event_prompts)) + raise ValueError( + f"Unknown Lingbot text event_id={event_id!r}. Supported: {supported}" + ) + self._applied_event_id = event_id + prompt = ( + self._base_prompt + if event_id is None + else self._text_event_prompts[event_id] + ) + return {} if prompt is None else {FIELD_PROMPT: prompt} + + def set_base_prompt(self, prompt: str) -> None: + """Record the rollout prompt restored when a text event is cleared.""" + self._base_prompt = prompt + + +def _pose_segments( + command: Mapping[str, Any], + *, + start_s: float, + end_s: float, +) -> list[PoseSegment]: + """Return integrator-ready segments for one step window.""" + raw = command.get("segments") + if not raw: + # A source that supplies only level state still drives the step; the + # whole window then holds one constant command. + return [(start_s, end_s, _keys_from_axes(command))] + segments: list[PoseSegment] = [] + for segment_start, segment_end, axes in raw: + if float(segment_end) <= float(segment_start): + continue + segments.append( + (float(segment_start), float(segment_end), _keys_from_axes(axes)) + ) + if not segments: + return [(start_s, end_s, _keys_from_axes(command))] + return segments + + +def _required_int(metadata: Mapping[str, Any], name: str) -> int: + if name not in metadata: + raise ValueError( + f"Lingbot input mapping requires StepRequest.metadata[{name!r}]; the " + f"session did not provide it." + ) + return int(metadata[name]) + + +__all__ = [ + "CAMERA_COMMAND", + "FIELD_CAMERA_INTRINSICS", + "FIELD_CAMERA_TRAJECTORY", + "FIELD_TOTAL_CAMERA_FRAMES", + "KeyboardToCameraCommand", + "LingbotCameraTrace", + "LingbotInputMapping", + "TEXT_EVENT", + "TextEventSelection", + "load_camera_trace", +] diff --git a/integrations/lingbot/lingbot/runner.py b/integrations/lingbot/lingbot/runner.py index de269ae7..24016bb4 100644 --- a/integrations/lingbot/lingbot/runner.py +++ b/integrations/lingbot/lingbot/runner.py @@ -20,29 +20,32 @@ from dataclasses import dataclass, field from pathlib import Path -import numpy as np -import torch from loguru import logger -from flashdreams.core.io.disk import default_flashdreams_cache_dir -from flashdreams.core.io.download import download_to_cache from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig -from flashdreams.infra.runner_io import ( - ensure_output_dir, - load_first_frame_tensor, - runner_artifact_path, - write_runner_stats, - write_video_tensor, -) -from lingbot.encoder.camctrl import CamCtrlInput -from lingbot.encoder.utils import ( - get_Ks_transformed, - preprocess_example_poses, +from flashdreams.runtime import InputCanonicalizer, UserInputs, UserInputSchema +from flashdreams.runtime.metrics import NullMetricsRecorder +from flashdreams.runtime.runner import run_inference_session +from lingbot.example_data import ( + EXAMPLE_DATA_AVAILABLE_IDXS, + EXAMPLE_DATA_BASE_URL, + EXAMPLE_DATA_DIR_LOCAL, + EXAMPLE_DATA_FILENAMES, + EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS, + ensure_example_data_downloaded, + example_data_dirname, ) from lingbot.pipeline import ( LingbotWorldInferencePipeline, ) +from lingbot.runtime import ( + LingbotModelAdapter, + LingbotRunnerOutputTarget, + inference_config_from_runner_config, + inference_input_from_replay_inputs, + replay_inputs_from_runner_config, +) __all__ = [ "LingbotWorldRunnerConfig", @@ -58,68 +61,6 @@ _INTRINSICS_REFERENCE_WIDTH = 832 """Capture-resolution width matching :data:`_INTRINSICS_REFERENCE_HEIGHT`.""" -EXAMPLE_DATA_BASE_URL = ( - "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples" -) -"""HTTP base URL for the canonical examples shared by all LingBot versions.""" - -EXAMPLE_DATA_DIR_LOCAL = default_flashdreams_cache_dir() / "example_data/lingbot_world" -"""Local cache root where downloaded example folders are stored.""" - -EXAMPLE_DATA_FILENAMES = ( - "image.jpg", - "poses.npy", - "intrinsics.npy", - "prompt.txt", -) -"""Example assets downloaded when each file is available upstream.""" - -EXAMPLE_DATA_AVAILABLE_IDXS = (0, 1, 2, 3, 4, 5) -"""Supported upstream example indices currently hosted under ``examples/``.""" - -EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS = (0, 1, 2, 5) -"""Example indices that provide their own upstream ``prompt.txt`` file.""" - - -def example_data_dirname(example_idx: int) -> str: - """Format ``example_idx`` into the upstream folder naming convention.""" - assert example_idx in EXAMPLE_DATA_AVAILABLE_IDXS, ( - f"--example_idx must be one of {EXAMPLE_DATA_AVAILABLE_IDXS}." - ) - return f"{example_idx:02d}" - - -def ensure_example_data_downloaded(*, is_rank_zero: bool, example_idx: int) -> Path: - """Download bundled GitHub example files on rank 0; barrier other ranks. - - The runner calls this from :meth:`LingbotWorldRunner._fill_example_data_defaults`; - the WebRTC server calls it from its ``main()`` so the same files - land on disk before the server's - ``LingbotWebRTCSessionManager._initialize_sync`` checks for them. The - download itself is small (image + intrinsics + poses, plus a prompt - when available), uses the public LingBot-World GitHub raw URLs, and - is cached at :data:`EXAMPLE_DATA_DIR_LOCAL` so repeat calls are - no-ops. - """ - example_dirname = example_data_dirname(example_idx) - cache_dir = EXAMPLE_DATA_DIR_LOCAL / example_dirname - if is_rank_zero: - for filename in EXAMPLE_DATA_FILENAMES: - if ( - filename == "prompt.txt" - and example_idx not in EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS - ): - continue - download_to_cache( - f"{EXAMPLE_DATA_BASE_URL}/{example_dirname}/{filename}", - cache_dir=cache_dir, - filename=filename, - ) - if torch.distributed.is_initialized(): - torch.distributed.barrier() - return cache_dir - - @dataclass(kw_only=True) class LingbotWorldRunnerConfig(RunnerConfig): """Runner config for every shipped LingBot-World variant.""" @@ -222,118 +163,34 @@ def _fill_example_data_defaults(self) -> None: cfg.prompt_path = example_dir / "prompt.txt" def run(self) -> None: - """Drive an AR rollout until the camera stream is exhausted.""" + """Drive an AR rollout through the Lingbot runtime API path.""" cfg = self.config - if cfg.example_data: - self._fill_example_data_defaults() - assert cfg.image_path is not None, ( - "LingbotWorldRunner requires --image_path (first-frame RGB image)." - ) - assert cfg.pose_path is not None, ( - "LingbotWorldRunner requires --pose_path " - "(.npy of [T, 4, 4] camera-to-world matrices)." + adapter = LingbotModelAdapter() + inference_config = inference_config_from_runner_config( + cfg, + device=f"cuda:{self.local_rank}" if self.world_size > 1 else cfg.device, + pipeline=self.pipeline, ) - assert cfg.intrinsic_path is not None, ( - "LingbotWorldRunner requires --intrinsic_path " - "(.npy of [T, 4] camera intrinsics)." - ) - - prompt = self._resolve_prompt() - device = torch.device(f"cuda:{self.local_rank}") - - # Pipeline / encoder accept ``[*batch_shape, ...]`` shapes; the - # shipped configs pin ``batch_shape=()`` so a single-rollout layout - # is just ``[T, C, H, W]`` (image) / ``[T, 4, 4]`` (poses) / - # ``[T, 4]`` (intrinsics). - first_frames_t = load_first_frame_tensor( - cfg.image_path, - pixel_height=cfg.pixel_height, - pixel_width=cfg.pixel_width, - device=device, - dtype=torch.bfloat16, - interpolation="cubic", - install_hint="Install the lingbot plugin: pip install flashdreams-lingbot.", - ) - - Ks = np.load(cfg.intrinsic_path) - Ks_t = torch.from_numpy(Ks).to(device=device, dtype=torch.float32) - # Rescale capture-resolution intrinsics to the runner's frame size. - camera_intrinsics_t = get_Ks_transformed( - Ks_t, - height_org=_INTRINSICS_REFERENCE_HEIGHT, - width_org=_INTRINSICS_REFERENCE_WIDTH, - height_resize=cfg.pixel_height, - width_resize=cfg.pixel_width, - height_final=cfg.pixel_height, - width_final=cfg.pixel_width, + replay_inputs = replay_inputs_from_runner_config( + cfg, + is_rank_zero=self.is_rank_zero, ) - - c2ws = np.load(cfg.pose_path) - c2ws, trans_normalizer = preprocess_example_poses(c2ws) - camera_poses_t = torch.from_numpy(c2ws).to(device=device, dtype=torch.float32) - total_camera_frames = camera_poses_t.shape[0] - - if self.is_rank_zero: - logger.info( - f"[{cfg.runner_name}] loaded first_frame=" - f"{tuple(first_frames_t.shape)}, camera_poses=" - f"{tuple(camera_poses_t.shape)}" - ) - - cache = self.pipeline.initialize_cache(text=[prompt], image=first_frames_t) - - torch.cuda.synchronize() - if torch.distributed.is_initialized(): - torch.distributed.barrier() - - output_stream = self.create_video_output_stream(fps=cfg.fps) - start = 0 - for i in range(cfg.total_blocks): - num_frames = self.pipeline.get_num_output_frames(i) - end = start + num_frames - if end > total_camera_frames: - break - if self.is_rank_zero: - logger.info( - f"[{cfg.runner_name}] AR step {i}/{cfg.total_blocks}, " - f"num_frames={num_frames}, frames=[{start}, {end})" - ) - camctrl_input = CamCtrlInput( - intrinsics=camera_intrinsics_t[start:end], - poses=camera_poses_t[start:end], - world_scale=float(trans_normalizer), - ) - video_chunk = self.pipeline.generate( - autoregressive_index=i, - cache=cache, - input=camctrl_input, - ) - stats = self.pipeline.finalize(autoregressive_index=i, cache=cache) - output_stream.process(video_chunk, autoregressive_index=i, stats=stats) - start = end - - video = output_stream.finish() - if video is None: - return - - ensure_output_dir(cfg.output_dir) - video_path = runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4") - write_video_tensor( - video, - video_path, + initial_inputs = inference_input_from_replay_inputs(replay_inputs) + output_target = LingbotRunnerOutputTarget( + output_stream=self.create_video_output_stream(fps=cfg.fps), + output_dir=cfg.output_dir, + runner_name=cfg.runner_name, fps=cfg.fps, - layout="tchw", - install_hint="Install the lingbot plugin: pip install flashdreams-lingbot.", ) - logger.info( - f"[{cfg.runner_name}] wrote video {tuple(video.shape)} " - f"-> {video_path.resolve()}" + mapping = adapter.create_input_mapping(replay_inputs) + run_inference_session( + adapter=adapter, + config=inference_config, + mapping=mapping, + canonicalizer=InputCanonicalizer(), + source_schema=UserInputSchema(description="Lingbot runner fixed inputs"), + user_inputs=UserInputs(), + initial_inputs=initial_inputs, + output=output_target, + metrics=NullMetricsRecorder(), ) - - if output_stream.stats_history: - stats_path = write_runner_stats( - cfg.output_dir, cfg.runner_name, output_stream.stats_history - ) - logger.info( - f"[{cfg.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" - ) diff --git a/integrations/lingbot/lingbot/runtime.py b/integrations/lingbot/lingbot/runtime.py new file mode 100644 index 00000000..c15057aa --- /dev/null +++ b/integrations/lingbot/lingbot/runtime.py @@ -0,0 +1,1084 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot runtime API adapter and replay session implementation.""" + +from __future__ import annotations + +import os +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import torch.distributed as dist +from loguru import logger + +from flashdreams.core.distributed import init as init_distributed +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.infra.runner_io import ( + ensure_output_dir, + load_first_frame_tensor, + runner_artifact_path, + write_runner_stats, + write_video_tensor, +) +from flashdreams.infra.video_output import RunnerVideoOutputStream, VideoStepResult +from flashdreams.runtime import ( + CanonicalInputSchema, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InputField, + OutputArtifact, +) +from flashdreams.runtime.interfaces import InferenceRuntime, InferenceSession +from flashdreams.runtime.types import StepRequest, StepResult, TimeWindow +from lingbot.encoder.camctrl import CamCtrlInput +from lingbot.example_data import ( + EXAMPLE_DATA_AVAILABLE_IDXS, + EXAMPLE_DATA_DIR_LOCAL, + EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS, + ensure_example_data_downloaded, + example_asset_urls, + example_data_dirname, +) +from lingbot.input_mapping import ( + CAMERA_COMMAND, + FIELD_CAMERA_INTRINSICS, + FIELD_CAMERA_TRAJECTORY, + FIELD_TOTAL_CAMERA_FRAMES, + TEXT_EVENT, + LingbotCameraTrace, + LingbotInputMapping, + load_camera_trace, +) + +LINGBOT_MODEL_ID = "lingbot" +DEFAULT_LINGBOT_PRESET = "lingbot-world-fast-taehv-window15-sink3" +DEFAULT_PIXEL_HEIGHT = 464 +DEFAULT_PIXEL_WIDTH = 832 +DEFAULT_FPS = 16 + +_INTRINSICS_REFERENCE_HEIGHT = 480 +_INTRINSICS_REFERENCE_WIDTH = 832 +_INSTALL_HINT = "Install the lingbot plugin: pip install flashdreams-lingbot." + +FIELD_PROMPT = "prompt" +FIELD_FIRST_FRAME_PATH = "first_frame_path" +FIELD_CAMERA_POSES_PATH = "camera_poses_path" +FIELD_CAMERA_INTRINSICS_PATH = "camera_intrinsics_path" +FIELD_TOTAL_BLOCKS = "total_blocks" +FIELD_PIXEL_HEIGHT = "pixel_height" +FIELD_PIXEL_WIDTH = "pixel_width" +FIELD_FPS = "fps" +FIELD_WORLD_SCALE = "world_scale" + +PipelineFactory = Callable[[Any, str], Any] + + +@dataclass(frozen=True, kw_only=True, slots=True) +class LingbotReplayInputs: + """Resolved model-facing Lingbot replay inputs.""" + + prompt: str + first_frame_path: Path + camera_poses_path: Path + camera_intrinsics_path: Path + total_blocks: int = 20 + pixel_height: int = DEFAULT_PIXEL_HEIGHT + pixel_width: int = DEFAULT_PIXEL_WIDTH + fps: int = DEFAULT_FPS + world_scale: float | None = None + + def __post_init__(self) -> None: + if self.total_blocks <= 0: + raise ValueError("LingbotReplayInputs.total_blocks must be > 0.") + if self.pixel_height <= 0 or self.pixel_width <= 0: + raise ValueError("LingbotReplayInputs pixel dimensions must be > 0.") + if self.fps <= 0: + raise ValueError("LingbotReplayInputs.fps must be > 0.") + if self.world_scale is not None and self.world_scale <= 0: + raise ValueError("LingbotReplayInputs.world_scale must be > 0.") + object.__setattr__(self, "prompt", " ".join(self.prompt.split())) + object.__setattr__(self, "first_frame_path", Path(self.first_frame_path)) + object.__setattr__(self, "camera_poses_path", Path(self.camera_poses_path)) + object.__setattr__( + self, + "camera_intrinsics_path", + Path(self.camera_intrinsics_path), + ) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class LingbotSessionInputs: + """Session-global Lingbot state established at session start or reset. + + The camera trajectory is deliberately absent: it arrives per step through + ``InferenceInput.step``, built by the selected input mapping from either a + fixed trace or live user events. + """ + + prompt: str + first_frame_path: Path + total_blocks: int + pixel_height: int + pixel_width: int + fps: int + world_scale: float + total_camera_frames: int | None = None + + def __post_init__(self) -> None: + if self.total_blocks <= 0: + raise ValueError("LingbotSessionInputs.total_blocks must be > 0.") + if self.pixel_height <= 0 or self.pixel_width <= 0: + raise ValueError("LingbotSessionInputs pixel dimensions must be > 0.") + if self.fps <= 0: + raise ValueError("LingbotSessionInputs.fps must be > 0.") + if self.world_scale < 0: + raise ValueError("LingbotSessionInputs.world_scale must be >= 0.") + if self.total_camera_frames is not None and self.total_camera_frames <= 0: + raise ValueError( + "LingbotSessionInputs.total_camera_frames must be > 0 when set." + ) + object.__setattr__(self, "first_frame_path", Path(self.first_frame_path)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class LingbotReplayRuntimeOptions: + """Construction knobs for the Lingbot replay runtime.""" + + pipeline_config: Any + pipeline: Any | None = None + pipeline_factory: PipelineFactory | None = None + output_layout: VideoTensorLayout = "tchw" + + +class LingbotModelAdapter: + """Model adapter exposing Lingbot through ``flashdreams.runtime``.""" + + def __init__( + self, + *, + runtime_factory: Callable[..., InferenceRuntime] | None = None, + pipeline_factory: PipelineFactory | None = None, + ) -> None: + self._runtime_factory = runtime_factory or LingbotReplayRuntime + self._pipeline_factory = pipeline_factory + + @property + def model_id(self) -> str: + return LINGBOT_MODEL_ID + + @property + def inference_input_schema(self) -> InferenceInputSchema: + return InferenceInputSchema( + description="Lingbot camera-control model inputs.", + global_conditioning_fields=( + InputField( + name=FIELD_PROMPT, + input_modality="text", + frequency_consumed="once", + description=( + "Prompt text for the rollout. A non-empty value passed " + "to step() requests a text-event context swap." + ), + ), + InputField( + name=FIELD_FIRST_FRAME_PATH, + input_modality="image/path", + frequency_consumed="once", + description="First-frame RGB image path.", + ), + InputField(name=FIELD_TOTAL_BLOCKS, input_modality="count"), + InputField(name=FIELD_PIXEL_HEIGHT, input_modality="pixel-height"), + InputField(name=FIELD_PIXEL_WIDTH, input_modality="pixel-width"), + InputField(name=FIELD_FPS, input_modality="fps"), + InputField( + name=FIELD_WORLD_SCALE, + required=False, + input_modality="scale", + frequency_consumed="once", + description="Pose normalizer; supplied by the input mapping.", + ), + InputField( + name=FIELD_TOTAL_CAMERA_FRAMES, + required=False, + input_modality="count", + frequency_consumed="once", + description=( + "Frames the input source can supply. Absent means " + "unbounded, so only total_blocks ends the rollout." + ), + ), + ), + step_fields=( + InputField( + name=FIELD_CAMERA_TRAJECTORY, + input_modality="c2w_sequence", + frequency_consumed="per_step", + metadata={"shape": "[T,4,4]", "frame": "camera_to_world"}, + description="Camera-to-world poses for this chunk's frames.", + ), + InputField( + name=FIELD_CAMERA_INTRINSICS, + input_modality="intrinsics_vec4_sequence", + frequency_consumed="per_step", + metadata={"shape": "[T,4]"}, + description="Per-frame intrinsics for this chunk's frames.", + ), + ), + ) + + @property + def canonical_input_schema(self) -> CanonicalInputSchema | None: + return CanonicalInputSchema( + modalities=(CAMERA_COMMAND, TEXT_EVENT), + description="Lingbot live camera control and text events.", + ) + + def default_input_mapping(self) -> LingbotInputMapping | None: + """Return no default mapping; Lingbot mappings are scenario-bound. + + Both trajectory sources need scenario data the adapter does not have + here: a fixed trace needs its ``.npy`` files, and live control needs + base intrinsics and a world scale. Callers build one with + :meth:`create_input_mapping`. + """ + return None + + def create_input_mapping( + self, + replay_inputs: LingbotReplayInputs, + *, + text_event_prompts: Mapping[str, str] | None = None, + ) -> LingbotInputMapping: + """Build the fixed-trace mapping for a resolved replay scenario.""" + mapping = LingbotInputMapping( + fps=replay_inputs.fps, + trace=load_camera_trace( + camera_poses_path=replay_inputs.camera_poses_path, + camera_intrinsics_path=replay_inputs.camera_intrinsics_path, + pixel_height=replay_inputs.pixel_height, + pixel_width=replay_inputs.pixel_width, + intrinsics_reference_height=_INTRINSICS_REFERENCE_HEIGHT, + intrinsics_reference_width=_INTRINSICS_REFERENCE_WIDTH, + world_scale=replay_inputs.world_scale, + ), + text_event_prompts=text_event_prompts, + ) + mapping.set_base_prompt(replay_inputs.prompt) + return mapping + + def create_live_input_mapping( + self, + *, + fps: int, + base_intrinsics: Any, + world_scale: float, + prompt: str = "", + text_event_prompts: Mapping[str, str] | None = None, + trace: LingbotCameraTrace | None = None, + ) -> LingbotInputMapping: + """Build the event-driven mapping used by keyboard-driving scenarios.""" + mapping = LingbotInputMapping( + fps=fps, + trace=trace, + base_intrinsics=base_intrinsics, + world_scale=world_scale, + text_event_prompts=text_event_prompts, + ) + mapping.set_base_prompt(prompt) + return mapping + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.model_id: + raise ValueError( + f"Lingbot adapter requires model_id={self.model_id!r}, " + f"got {config.model_id!r}." + ) + self.pipeline_config(config) + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + return self._runtime_factory( + config=config, + options=LingbotReplayRuntimeOptions( + pipeline_config=self.pipeline_config(config), + pipeline=config.runtime_options.get("pipeline"), + pipeline_factory=self._pipeline_factory, + output_layout=str(config.runtime_options.get("output_layout", "tchw")), + ), + ) + + def preset_id(self, config: InferenceConfig | None) -> str: + return ( + DEFAULT_LINGBOT_PRESET + if config is None or config.preset_id is None + else config.preset_id + ) + + def pipeline_config(self, config: InferenceConfig) -> Any: + custom = config.runtime_options.get("pipeline_config") + if custom is not None: + return custom + preset_id = self.preset_id(config) + from lingbot.config import PIPELINE_CONFIGS # noqa: PLC0415 + + try: + return PIPELINE_CONFIGS[preset_id] + except KeyError as exc: + supported = ", ".join(sorted(PIPELINE_CONFIGS)) + raise ValueError( + f"Unsupported Lingbot preset_id={preset_id!r}. " + f"Supported presets: {supported}." + ) from exc + + def default_replay_prompt(self, config: InferenceConfig | None) -> str: + from lingbot.config import RUNNER_CONFIGS # noqa: PLC0415 + + runner = RUNNER_CONFIGS.get(self.preset_id(config)) + return "" if runner is None else str(getattr(runner, "prompt", "")) + + +class LingbotReplayRuntime: + """Heavyweight Lingbot runtime consumed by the standard loop.""" + + def __init__( + self, + *, + config: InferenceConfig, + options: LingbotReplayRuntimeOptions, + ) -> None: + self.config = config + self.options = options + if _is_torchrun_env() and not dist.is_initialized(): + init_distributed() + + if dist.is_initialized(): + self.local_rank = int(os.environ.get("LOCAL_RANK", "0")) + self.world_size = dist.get_world_size() + self.global_rank = dist.get_rank() + device = f"cuda:{self.local_rank}" + else: + self.local_rank = 0 + self.world_size = 1 + self.global_rank = 0 + device = config.device or "cuda" + + self.is_rank_zero = self.global_rank == 0 + if options.pipeline is not None: + self.pipeline = options.pipeline + self._owns_pipeline = False + else: + factory = options.pipeline_factory or _default_pipeline_factory + self.pipeline = factory(options.pipeline_config, device) + self._owns_pipeline = True + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + session_inputs = session_inputs_from_inference_input(inputs) + return LingbotReplaySession( + pipeline=self.pipeline, + session_inputs=session_inputs, + device=torch.device(f"cuda:{self.local_rank}") + if dist.is_initialized() + else torch.device(self.config.device or "cuda"), + is_rank_zero=self.is_rank_zero, + output_layout=self.options.output_layout, + ) + + def close(self) -> None: + pipeline = getattr(self, "pipeline", None) + if self._owns_pipeline and pipeline is not None: + close = getattr(pipeline, "close", None) + if callable(close): + close() + del self.pipeline + device = torch.device(self.config.device or "cuda") + if device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.empty_cache() + + +class LingbotReplaySession: + """One Lingbot rollout driven by per-step camera inputs.""" + + def __init__( + self, + *, + pipeline: Any, + session_inputs: LingbotSessionInputs, + device: torch.device, + is_rank_zero: bool, + output_layout: VideoTensorLayout, + ) -> None: + self.pipeline = pipeline + self.inputs = session_inputs + self.device = device + self.is_rank_zero = is_rank_zero + self.output_layout = output_layout + self.dtype = torch.bfloat16 + self._closed = False + self._step_index = 0 + self._frame_start = 0 + self._active_prompt = session_inputs.prompt + self._cache = self._initialize_cache() + if self.device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.synchronize(device=self.device) + if dist.is_initialized(): + dist.barrier() + + def next_step_request(self) -> StepRequest | None: + if self._closed: + return None + if self._step_index >= self.inputs.total_blocks: + return None + num_frames = int(self.pipeline.get_num_output_frames(self._step_index)) + frame_end = self._frame_start + num_frames + total_frames = self.inputs.total_camera_frames + if total_frames is not None and frame_end > total_frames: + return None + fps = self.inputs.fps + return StepRequest( + step_index=self._step_index, + # The window is what lets a mapping slice user events for exactly + # this chunk instead of replaying the whole session history. + user_input_window=TimeWindow( + start_s=self._frame_start / fps, + end_s=frame_end / fps, + ), + metadata={ + "num_frames": num_frames, + "frame_start": self._frame_start, + }, + ) + + def step(self, inputs: InferenceInput) -> StepResult: + if self._closed: + raise RuntimeError("Lingbot replay session is closed.") + + step_index = self._step_index + num_frames = int(self.pipeline.get_num_output_frames(step_index)) + self._apply_global_conditioning_update(inputs) + camera_poses = _require_step_tensor( + inputs, + FIELD_CAMERA_TRAJECTORY, + expected_shape=(num_frames, 4, 4), + ) + camera_intrinsics = _require_step_tensor( + inputs, + FIELD_CAMERA_INTRINSICS, + expected_shape=(num_frames, 4), + ) + frame_start = self._frame_start + frame_end = frame_start + num_frames + + if self.is_rank_zero: + logger.info( + "Lingbot runtime step {} frames=[{}, {})", + step_index, + frame_start, + frame_end, + ) + camctrl_input = CamCtrlInput( + intrinsics=camera_intrinsics.to(device=self.device, dtype=torch.float32), + poses=camera_poses.to(device=self.device, dtype=torch.float32), + world_scale=self.inputs.world_scale, + ) + start_t = time.perf_counter() + video_chunk = self.pipeline.generate( + autoregressive_index=step_index, + cache=self._cache, + input=camctrl_input, + ) + stats = self.pipeline.finalize( + autoregressive_index=step_index, + cache=self._cache, + ) + elapsed_s = time.perf_counter() - start_t + self._step_index += 1 + self._frame_start = frame_end + + metrics = _numeric_stats(stats) + metrics.setdefault("model_step_s", elapsed_s) + return StepResult( + step_index=step_index, + output=VideoStepResult.from_video_chunk( + chunk_index=step_index, + video_chunk=video_chunk, + layout=self.output_layout, + stats=metrics, + ), + frame_count=num_frames, + output_window=TimeWindow( + start_s=frame_start / self.inputs.fps, + end_s=frame_end / self.inputs.fps, + ), + metrics=metrics, + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + if inputs is not None: + session_inputs = session_inputs_from_inference_input(inputs) + if session_inputs != self.inputs: + raise ValueError("Lingbot replay reset cannot swap inputs.") + cache = getattr(self, "_cache", None) + if cache is not None: + del self._cache + self._active_prompt = self.inputs.prompt + self._cache = self._initialize_cache() + self._step_index = 0 + self._frame_start = 0 + + def _apply_global_conditioning_update(self, inputs: InferenceInput) -> None: + """Apply a mid-rollout text-event context swap, when one was requested. + + Text events reach the model as a session-global prompt update rather + than a per-step field, because they replace the rollout's whole + cross-attention text context. Not every pipeline can do this, so the + capability is probed the same way the WebRTC runtime probes it, and + only when a swap is actually requested. + """ + prompt = inputs.global_conditioning.get(FIELD_PROMPT) + if prompt is None or prompt == self._active_prompt: + return + transformer = self.pipeline.diffusion_model.transformer + replace_text_embeddings = getattr(transformer, "replace_text_embeddings", None) + if not callable(replace_text_embeddings): + raise RuntimeError( + "Lingbot text events need a pipeline whose transformer supports " + "replace_text_embeddings; this pipeline does not." + ) + self.pipeline._ensure_oneshot_encoders_loaded() + embeddings = self.pipeline.text_encoder([prompt]).to(device=self.device) + replace_text_embeddings(self._cache.transformer_cache, embeddings) + self._active_prompt = prompt + if self.is_rank_zero: + logger.info("Lingbot text context updated at step {}", self._step_index) + + def close(self) -> None: + self._closed = True + cache = getattr(self, "_cache", None) + if cache is not None: + del self._cache + + def _initialize_cache(self) -> Any: + first_frames = load_first_frame_tensor( + self.inputs.first_frame_path, + pixel_height=self.inputs.pixel_height, + pixel_width=self.inputs.pixel_width, + device=self.device, + dtype=self.dtype, + interpolation="cubic", + install_hint=_INSTALL_HINT, + ) + return self.pipeline.initialize_cache( + text=[self.inputs.prompt], + image=first_frames, + ) + + +def _require_step_tensor( + inputs: InferenceInput, + name: str, + *, + expected_shape: tuple[int, ...], +) -> torch.Tensor: + """Return one required per-step camera tensor, shape-checked.""" + if name not in inputs.step: + raise ValueError( + f"Lingbot step inputs are missing {name!r}. The selected input " + f"mapping must produce it for every step." + ) + value = inputs.step[name] + if not isinstance(value, torch.Tensor): + value = torch.as_tensor(np.asarray(value), dtype=torch.float32) + if tuple(value.shape) != expected_shape: + raise ValueError( + f"Lingbot step input {name!r} must have shape {expected_shape}, got " + f"{tuple(value.shape)}." + ) + return value + + +@dataclass(slots=True) +class LingbotRunnerOutputTarget: + """Runner-compatible MP4/stats output target for Lingbot replay results.""" + + output_stream: RunnerVideoOutputStream + output_dir: Path + runner_name: str + fps: int | float + install_hint: str = _INSTALL_HINT + _opened: bool = False + + def open(self) -> None: + self._opened = True + + def write(self, result: StepResult) -> None: + if not self._opened: + raise RuntimeError("Cannot write to a closed Lingbot output target.") + video_result = result.output + if not isinstance(video_result, VideoStepResult): + raise TypeError( + "LingbotRunnerOutputTarget requires VideoStepResult output, " + f"got {type(video_result).__name__}." + ) + self.output_stream.process( + video_result.video_chunk, + autoregressive_index=video_result.chunk_index, + stats=video_result.stats or dict(result.metrics), + ) + + def close(self) -> tuple[OutputArtifact, ...]: + self._opened = False + artifacts: list[OutputArtifact] = [] + video = self.output_stream.finish() + if video is None: + return () + + ensure_output_dir(self.output_dir) + video_path = runner_artifact_path(self.output_dir, self.runner_name, "mp4") + write_video_tensor( + video, + video_path, + fps=self.fps, + layout="tchw", + install_hint=self.install_hint, + ) + logger.info( + "[{}] wrote video {} -> {}", + self.runner_name, + tuple(video.shape), + video_path.resolve(), + ) + artifacts.append( + OutputArtifact(kind="video/mp4", uri=str(video_path.resolve())) + ) + if self.output_stream.stats_history: + stats_path = write_runner_stats( + self.output_dir, + self.runner_name, + self.output_stream.stats_history, + ) + logger.info( + "[{}] wrote per-AR-step stats -> {}", + self.runner_name, + stats_path.resolve(), + ) + artifacts.append( + OutputArtifact(kind="application/json", uri=str(stats_path.resolve())) + ) + return tuple(artifacts) + + +def inference_config_from_runner_config( + runner_config: Any, + *, + device: str, + pipeline: Any | None = None, +) -> InferenceConfig: + """Build the runtime config directly from a Lingbot runner config.""" + runtime_options: dict[str, Any] = { + "pipeline_config": runner_config.pipeline, + "output_layout": runner_config.postprocess_output_layout or "tchw", + } + if pipeline is not None: + runtime_options["pipeline"] = pipeline + compile_network = getattr( + runner_config.pipeline.diffusion_model.transformer, + "compile_network", + None, + ) + return InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=str(runner_config.pipeline.name), + device=device, + compile=None if compile_network is None else bool(compile_network), + runtime_options=runtime_options, + ) + + +def inference_input_from_runner_config( + runner_config: Any, + *, + is_rank_zero: bool, +) -> InferenceInput: + """Build session-global runtime inputs from a Lingbot runner config.""" + return inference_input_from_replay_inputs( + replay_inputs_from_runner_config(runner_config, is_rank_zero=is_rank_zero) + ) + + +def replay_inputs_from_runner_config( + runner_config: Any, + *, + is_rank_zero: bool, +) -> LingbotReplayInputs: + """Resolve a Lingbot runner config into scenario-level replay inputs.""" + return replay_inputs_from_mapping( + { + FIELD_PROMPT: getattr(runner_config, "prompt", ""), + "prompt_path": getattr(runner_config, "prompt_path", None), + FIELD_FIRST_FRAME_PATH: getattr(runner_config, "image_path", None), + FIELD_CAMERA_POSES_PATH: getattr(runner_config, "pose_path", None), + FIELD_CAMERA_INTRINSICS_PATH: getattr( + runner_config, + "intrinsic_path", + None, + ), + FIELD_TOTAL_BLOCKS: getattr(runner_config, "total_blocks", 20), + FIELD_PIXEL_HEIGHT: getattr( + runner_config, + "pixel_height", + DEFAULT_PIXEL_HEIGHT, + ), + FIELD_PIXEL_WIDTH: getattr( + runner_config, + "pixel_width", + DEFAULT_PIXEL_WIDTH, + ), + FIELD_FPS: getattr(runner_config, "fps", DEFAULT_FPS), + "example_data": getattr(runner_config, "example_data", False), + "example_idx": getattr(runner_config, "example_idx", 0), + }, + is_rank_zero=is_rank_zero, + ) + + +def inference_input_from_replay_inputs( + replay_inputs: LingbotReplayInputs, +) -> InferenceInput: + """Encode resolved Lingbot replay inputs into ``InferenceInput``.""" + payload: dict[str, Any] = { + FIELD_PROMPT: replay_inputs.prompt, + FIELD_FIRST_FRAME_PATH: replay_inputs.first_frame_path, + FIELD_TOTAL_BLOCKS: replay_inputs.total_blocks, + FIELD_PIXEL_HEIGHT: replay_inputs.pixel_height, + FIELD_PIXEL_WIDTH: replay_inputs.pixel_width, + FIELD_FPS: replay_inputs.fps, + } + if replay_inputs.world_scale is not None: + payload[FIELD_WORLD_SCALE] = replay_inputs.world_scale + return InferenceInput(global_conditioning=payload) + + +def session_inputs_from_inference_input( + inputs: InferenceInput, +) -> LingbotSessionInputs: + """Decode and validate session-global Lingbot inputs.""" + missing = LingbotModelAdapter().inference_input_schema.missing_global_conditioning( + inputs + ) + if missing: + raise ValueError(f"Lingbot session inputs missing required fields: {missing}.") + gc = inputs.global_conditioning + if gc.get(FIELD_WORLD_SCALE) is None: + raise ValueError( + "Lingbot session inputs require 'world_scale'; the selected input " + "mapping supplies it from the camera trace or live control setup." + ) + total_camera_frames = gc.get(FIELD_TOTAL_CAMERA_FRAMES) + return LingbotSessionInputs( + prompt=str(gc[FIELD_PROMPT]), + first_frame_path=Path(gc[FIELD_FIRST_FRAME_PATH]), + total_blocks=int(gc[FIELD_TOTAL_BLOCKS]), + pixel_height=int(gc[FIELD_PIXEL_HEIGHT]), + pixel_width=int(gc[FIELD_PIXEL_WIDTH]), + fps=int(gc[FIELD_FPS]), + world_scale=float(gc[FIELD_WORLD_SCALE]), + total_camera_frames=( + None if total_camera_frames is None else int(total_camera_frames) + ), + ) + + +def replay_inputs_from_mapping( + value: Any, + *, + default_prompt: str = "", + is_rank_zero: bool = True, +) -> LingbotReplayInputs: + """Resolve app/CLI replay values into direct Lingbot runtime inputs.""" + if isinstance(value, LingbotReplayInputs): + _require_existing_replay_paths(value) + return value + if value is None: + value = {} + if not isinstance(value, Mapping): + raise TypeError( + "Lingbot replay inputs must be a LingbotReplayInputs, mapping, or None." + ) + + example_idx = int(value.get("example_idx", 0)) + if example_idx not in EXAMPLE_DATA_AVAILABLE_IDXS: + raise ValueError( + f"Lingbot replay example_idx must be one of {EXAMPLE_DATA_AVAILABLE_IDXS}." + ) + + first_frame_path = _optional_path( + value.get(FIELD_FIRST_FRAME_PATH, value.get("image_path")) + ) + poses_path = _optional_path( + value.get(FIELD_CAMERA_POSES_PATH, value.get("pose_path")) + ) + intrinsics_path = _optional_path( + value.get(FIELD_CAMERA_INTRINSICS_PATH, value.get("intrinsic_path")) + ) + prompt_path = _optional_path(value.get("prompt_path")) + example_data = _resolve_example_data_default(value) + + if example_data and ( + first_frame_path is None + or poses_path is None + or intrinsics_path is None + or ( + prompt_path is None + and not _has_nonempty_value(value, FIELD_PROMPT) + and example_idx in EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS + ) + ): + example_dir = ensure_example_data_downloaded( + is_rank_zero=is_rank_zero, + example_idx=example_idx, + ) + first_frame_path = first_frame_path or example_dir / "image.jpg" + poses_path = poses_path or example_dir / "poses.npy" + intrinsics_path = intrinsics_path or example_dir / "intrinsics.npy" + if ( + prompt_path is None + and not _has_nonempty_value(value, FIELD_PROMPT) + and example_idx in EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS + ): + prompt_path = example_dir / "prompt.txt" + + replay_inputs = LingbotReplayInputs( + prompt=_resolve_prompt( + value, + prompt_path=prompt_path, + default_prompt=default_prompt, + ), + first_frame_path=_require_path_value( + first_frame_path, + label=FIELD_FIRST_FRAME_PATH, + ), + camera_poses_path=_require_path_value( + poses_path, + label=FIELD_CAMERA_POSES_PATH, + ), + camera_intrinsics_path=_require_path_value( + intrinsics_path, + label=FIELD_CAMERA_INTRINSICS_PATH, + ), + total_blocks=int(value.get(FIELD_TOTAL_BLOCKS, 20)), + pixel_height=int(value.get(FIELD_PIXEL_HEIGHT, DEFAULT_PIXEL_HEIGHT)), + pixel_width=int(value.get(FIELD_PIXEL_WIDTH, DEFAULT_PIXEL_WIDTH)), + fps=int(value.get(FIELD_FPS, DEFAULT_FPS)), + world_scale=( + None + if FIELD_WORLD_SCALE not in value or value[FIELD_WORLD_SCALE] is None + else float(value[FIELD_WORLD_SCALE]) + ), + ) + _require_existing_replay_paths(replay_inputs) + if prompt_path is not None: + _require_existing_path(prompt_path, label="prompt_path") + return replay_inputs + + +def build_lingbot_webrtc_runtime_config( + *, + preset_id: str, + pipeline_config: Any, + device: str, + seed: int, + compile_network: bool, + context_parallel_size: int, + video_height: int, + video_width: int, + fps: int, + warmup_chunks: int, + warmup_timeout_s: float, + example_idx: int, + prefer_sw_encoder: bool, + runtime_options: Mapping[str, Any] | None = None, +) -> Any: + """Build the Lingbot WebRTC runtime config from shared runtime inputs.""" + from lingbot.webrtc.session import LingbotRuntimeConfig # noqa: PLC0415 + + example_dirname = example_data_dirname(example_idx) + example_dir = EXAMPLE_DATA_DIR_LOCAL / example_dirname + if ( + example_idx == 0 + and not example_dir.exists() + and (EXAMPLE_DATA_DIR_LOCAL / "image.jpg").exists() + ): + example_dir = EXAMPLE_DATA_DIR_LOCAL + urls = example_asset_urls(example_idx) + runtime_config = LingbotRuntimeConfig( + config_name=preset_id, + pipeline_config=pipeline_config, + compile_network=compile_network, + seed=seed, + context_parallel_size=context_parallel_size, + device=device, + video_height=video_height, + video_width=video_width, + fps=fps, + warmup_chunks=warmup_chunks, + warmup_timeout_s=warmup_timeout_s, + encoder_backend="default" if prefer_sw_encoder else "auto", + example_data_dir=example_dir, + default_image_url=urls["image"], + default_intrinsics_url=urls["intrinsics"], + default_poses_url=urls["poses"], + ) + return _apply_webrtc_runtime_options(runtime_config, runtime_options or {}) + + +def _apply_webrtc_runtime_options(runtime_config: Any, options: Mapping[str, Any]) -> Any: + overrides: dict[str, Any] = {} + for name in ( + "world_scale", + "default_intrinsics", + "default_prompt", + "default_image_url", + "default_intrinsics_url", + "default_poses_url", + "encoder_bitrate_bps", + "encoder_gop", + "text_events", + ): + if name in options: + overrides[name] = options[name] + return replace(runtime_config, **overrides) if overrides else runtime_config + + +def _default_pipeline_factory(pipeline_config: Any, device: str) -> Any: + return pipeline_config.setup().to(device=device).eval() + + +def _numeric_stats(stats: Any) -> dict[str, float | int]: + if not isinstance(stats, Mapping): + return {} + return { + str(key): value + for key, value in stats.items() + if isinstance(value, (float, int)) and not isinstance(value, bool) + } + + +def _resolve_prompt( + value: Mapping[str, Any], + *, + prompt_path: Path | None, + default_prompt: str, +) -> str: + prompt = str(value.get(FIELD_PROMPT, value.get("prompt", ""))).strip() + if prompt: + return prompt + if prompt_path is not None: + lines = prompt_path.read_text(encoding="utf-8").splitlines() + if lines: + prompt = lines[0].strip() + if prompt: + return prompt + return default_prompt.strip() + + +def _resolve_example_data_default(value: Mapping[str, Any]) -> bool: + explicit = value.get("example_data") + if explicit is not None: + return _bool_value(explicit) + return not ( + _has_nonempty_value(value, FIELD_FIRST_FRAME_PATH) + or _has_nonempty_value(value, "image_path") + ) or not ( + _has_nonempty_value(value, FIELD_CAMERA_POSES_PATH) + or _has_nonempty_value(value, "pose_path") + ) or not ( + _has_nonempty_value(value, FIELD_CAMERA_INTRINSICS_PATH) + or _has_nonempty_value(value, "intrinsic_path") + ) + + +def _bool_value(value: Any) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + lowered = value.strip().lower() + if lowered in {"1", "true", "yes", "on"}: + return True + if lowered in {"0", "false", "no", "off"}: + return False + return bool(value) + + +def _has_nonempty_value(value: Mapping[str, Any], key: str) -> bool: + if key not in value: + return False + raw = value[key] + return raw is not None and raw != "" + + +def _optional_path(value: Any) -> Path | None: + if value is None or value == "": + return None + return Path(value) + + +def _require_path_value(value: Path | None, *, label: str) -> Path: + if value is None: + raise ValueError(f"Lingbot replay inputs require {label}.") + return value + + +def _require_existing_replay_paths(replay_inputs: LingbotReplayInputs) -> None: + _require_existing_path(replay_inputs.first_frame_path, label=FIELD_FIRST_FRAME_PATH) + _require_existing_path(replay_inputs.camera_poses_path, label=FIELD_CAMERA_POSES_PATH) + _require_existing_path( + replay_inputs.camera_intrinsics_path, + label=FIELD_CAMERA_INTRINSICS_PATH, + ) + + +def _require_existing_path(path: Path, *, label: str) -> None: + if not path.exists(): + raise FileNotFoundError(f"Lingbot replay inputs missing {label}: {path}") + + +def _is_torchrun_env() -> bool: + return "RANK" in os.environ and "WORLD_SIZE" in os.environ + + +__all__ = [ + "DEFAULT_FPS", + "DEFAULT_LINGBOT_PRESET", + "DEFAULT_PIXEL_HEIGHT", + "DEFAULT_PIXEL_WIDTH", + "FIELD_CAMERA_INTRINSICS_PATH", + "FIELD_CAMERA_POSES_PATH", + "FIELD_FIRST_FRAME_PATH", + "FIELD_FPS", + "FIELD_PIXEL_HEIGHT", + "FIELD_PIXEL_WIDTH", + "FIELD_PROMPT", + "FIELD_TOTAL_BLOCKS", + "FIELD_WORLD_SCALE", + "LINGBOT_MODEL_ID", + "LingbotModelAdapter", + "LingbotReplayInputs", + "LingbotReplayRuntime", + "LingbotReplayRuntimeOptions", + "LingbotReplaySession", + "LingbotRunnerOutputTarget", + "PipelineFactory", + "build_lingbot_webrtc_runtime_config", + "inference_config_from_runner_config", + "inference_input_from_replay_inputs", + "inference_input_from_runner_config", + "replay_inputs_from_inference_input", + "replay_inputs_from_mapping", +] diff --git a/integrations/lingbot/lingbot/webrtc/server.py b/integrations/lingbot/lingbot/webrtc/server.py index 4714ed85..3825a555 100644 --- a/integrations/lingbot/lingbot/webrtc/server.py +++ b/integrations/lingbot/lingbot/webrtc/server.py @@ -48,12 +48,15 @@ from flashdreams.serving.webrtc.server import ( close_package_resources as _close_package_resources, ) -from lingbot.runner import ( +from flashdreams.runtime import InferenceConfig +from lingbot.example_data import ( EXAMPLE_DATA_AVAILABLE_IDXS, - EXAMPLE_DATA_BASE_URL, - EXAMPLE_DATA_DIR_LOCAL, ensure_example_data_downloaded, - example_data_dirname, +) +from lingbot.runtime import ( + LINGBOT_MODEL_ID, + LingbotModelAdapter, + build_lingbot_webrtc_runtime_config, ) from lingbot.webrtc.session import ( LingbotImagePayload, @@ -105,6 +108,12 @@ def parse_args() -> argparse.Namespace: default="cuda:0", help="Torch device used for the Lingbot runtime.", ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Base random seed for the Lingbot rollout.", + ) parser.add_argument( "--warmup_chunks", type=int, @@ -335,34 +344,31 @@ def build_runtime_config( raise ValueError("--video-height and --video-width must be > 0") if args.video_height % 16 != 0 or args.video_width % 16 != 0: raise ValueError("--video-height and --video-width must be divisible by 16") - example_idx = getattr(args, "example_idx", 0) - example_dirname = example_data_dirname(example_idx) - example_dir = EXAMPLE_DATA_DIR_LOCAL / example_dirname - if ( - example_idx == 0 - and not example_dir.exists() - and (EXAMPLE_DATA_DIR_LOCAL / "image.jpg").exists() - ): - example_dir = EXAMPLE_DATA_DIR_LOCAL - return LingbotRuntimeConfig( - config_name=args.config_name, + + inference_config = InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=args.config_name, + device=device_override or args.device, + compile=not args.no_compile, + runtime_options={"context_parallel_size": context_parallel_size}, + ) + adapter = LingbotModelAdapter() + adapter.validate_config(inference_config) + return build_lingbot_webrtc_runtime_config( + preset_id=adapter.preset_id(inference_config), + pipeline_config=adapter.pipeline_config(inference_config), + device=inference_config.device or args.device, + seed=int(getattr(args, "seed", 42)), compile_network=not args.no_compile, context_parallel_size=context_parallel_size, - device=device_override or args.device, - warmup_chunks=args.warmup_chunks, - warmup_timeout_s=args.warmup_timeout_s, video_height=args.video_height, video_width=args.video_width, fps=args.fps, - encoder_backend=( - "default" if getattr(args, "prefer_sw_encoder", False) else "auto" - ), - example_data_dir=example_dir, - default_image_url=f"{EXAMPLE_DATA_BASE_URL}/{example_dirname}/image.jpg", - default_intrinsics_url=( - f"{EXAMPLE_DATA_BASE_URL}/{example_dirname}/intrinsics.npy" - ), - default_poses_url=f"{EXAMPLE_DATA_BASE_URL}/{example_dirname}/poses.npy", + warmup_chunks=args.warmup_chunks, + warmup_timeout_s=args.warmup_timeout_s, + example_idx=getattr(args, "example_idx", 0), + prefer_sw_encoder=getattr(args, "prefer_sw_encoder", False), + runtime_options=inference_config.runtime_options, ) diff --git a/integrations/lingbot/lingbot/webrtc/session.py b/integrations/lingbot/lingbot/webrtc/session.py index c0b59680..ee23bb96 100644 --- a/integrations/lingbot/lingbot/webrtc/session.py +++ b/integrations/lingbot/lingbot/webrtc/session.py @@ -476,6 +476,8 @@ class LingbotRuntimeConfig: text_events: tuple[TextEventSpec, ...] = field( default_factory=lambda: DEFAULT_TEXT_EVENTS ) + pipeline_config: Any | None = None + """Optional pre-resolved pipeline config used by shared demo adapters.""" @dataclass(frozen=True, slots=True) @@ -775,13 +777,16 @@ def _initialize_sync(self) -> None: if self._pipeline is not None: return - pipeline_configs = _pipeline_configs() - if self.config.config_name not in pipeline_configs: - supported = ", ".join(sorted(pipeline_configs)) - raise ValueError( - f"Unknown config_name={self.config.config_name!r}. " - f"Supported: {supported}" - ) + pipeline_config_base = self.config.pipeline_config + if pipeline_config_base is None: + pipeline_configs = _pipeline_configs() + if self.config.config_name not in pipeline_configs: + supported = ", ".join(sorted(pipeline_configs)) + raise ValueError( + f"Unknown config_name={self.config.config_name!r}. " + f"Supported: {supported}" + ) + pipeline_config_base = pipeline_configs[self.config.config_name] self._device = torch.device(self.config.device) if self._device.type == "cuda" and not torch.cuda.is_available(): @@ -796,7 +801,7 @@ def _initialize_sync(self) -> None: else self.config.seed ) pipeline_config = derive_config( - base_config=pipeline_configs[self.config.config_name], + base_config=pipeline_config_base, enable_sync_and_profile=True, diffusion_model=dict( seed=rollout_seed, @@ -1204,16 +1209,20 @@ class LingbotWebRTCSessionManager( def __init__( self, *, + runtime: LingbotInferenceRuntime | None = None, runtime_config: LingbotRuntimeConfig | None = None, fps: int | None = None, client_liveness_timeout_s: float = DEFAULT_CLIENT_LIVENESS_TIMEOUT_S, ) -> None: - runtime_config = runtime_config or LingbotRuntimeConfig() + runtime_config = runtime_config or getattr(runtime, "config", None) + if not isinstance(runtime_config, LingbotRuntimeConfig): + runtime_config = LingbotRuntimeConfig() fps = runtime_config.fps if fps is None else fps if fps <= 0: raise ValueError("fps must be > 0") + runtime = runtime or LingbotInferenceRuntime(config=runtime_config) super().__init__( - runtime=LingbotInferenceRuntime(config=runtime_config), + runtime=runtime, runtime_config=runtime_config, fps=fps, client_liveness_timeout_s=client_liveness_timeout_s, diff --git a/integrations/lingbot/pyproject.toml b/integrations/lingbot/pyproject.toml index f9e39cd9..b14cd2c9 100644 --- a/integrations/lingbot/pyproject.toml +++ b/integrations/lingbot/pyproject.toml @@ -41,6 +41,9 @@ dev = [ "pytest-asyncio>=0.23", ] +[project.scripts] +lingbot-demo = "lingbot.demo.cli:main" + # Each entry registers one ``runner_name`` slug with ``flashdreams-run``. # The discovery layer (``flashdreams.plugins.registry.discover_runners``) # scans this group at CLI startup; the entry-point name itself is purely diff --git a/integrations/lingbot/tests/test_demo_api.py b/integrations/lingbot/tests/test_demo_api.py new file mode 100644 index 00000000..95515e02 --- /dev/null +++ b/integrations/lingbot/tests/test_demo_api.py @@ -0,0 +1,625 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import Any, cast + +import numpy as np +import pytest +import torch +from aiohttp import web +from lingbot.demo import ( + DEFAULT_LINGBOT_PRESET, + LINGBOT_MODEL_ID, + LingbotDemoAdapter, + LingbotReplayInputs, + LingbotWebRTCScenario, +) +from lingbot.demo.cli import _replay_spec, _webrtc_spec, parse_args +from lingbot.demo.replay import ( + LingbotReplayRuntime, + LingbotReplayRuntimeOptions, +) +from lingbot.demo.webrtc import LingbotDemoWebRTCSessionManager +from lingbot.input_mapping import ( + FIELD_CAMERA_INTRINSICS, + FIELD_CAMERA_TRAJECTORY, +) +from lingbot.runtime import ( + FIELD_FIRST_FRAME_PATH, + FIELD_FPS, + FIELD_PIXEL_HEIGHT, + FIELD_PIXEL_WIDTH, + FIELD_PROMPT, + FIELD_TOTAL_BLOCKS, + inference_input_from_replay_inputs, +) +from lingbot.webrtc.session import LingbotRuntimeConfig + +from flashdreams.infra.video_output import VideoStepResult +from flashdreams.runtime import ( + CanonicalInputs, + InferenceConfig, + InferenceInput, + OutputArtifact, + OutputTarget, + StepResult, +) +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + WebRTCOutputSpec, + serve_flashdreams_demo, +) +from flashdreams.runtime.demo.replay import run_replay_demo +from flashdreams.runtime.demo.webrtc import WebRTCDemo, build_webrtc_demo +from flashdreams.serving.webrtc.server import SESSION_MANAGER_KEY + +pytestmark = pytest.mark.ci_cpu + + +def _write_camera_assets(poses: Path, intrinsics: Path, *, frames: int = 64) -> None: + """Write real .npy camera assets; the input mapping loads them for real.""" + trajectory = np.tile(np.eye(4, dtype=np.float32), (frames, 1, 1)) + trajectory[:, 2, 3] = np.linspace(0.0, 1.0, frames, dtype=np.float32) + np.save(poses, trajectory) + np.save( + intrinsics, + np.tile( + np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (frames, 1) + ), + ) + + +def test_lingbot_demo_defaults_to_interactive_preset() -> None: + args = parse_args(["replay", "--output", "demo.mp4"]) + + assert args.preset_id == "lingbot-world-fast-taehv-window15-sink3" + + +def test_lingbot_demo_adapter_declares_mp4_and_webrtc_modes() -> None: + adapter = LingbotDemoAdapter() + + assert adapter.model_id == LINGBOT_MODEL_ID + assert adapter.supported_input_modes() == ("replay", "keyboard-driving") + assert adapter.supported_output_modes() == ("mp4", "webrtc") + fields = { + field.name + for field in adapter.inference_input_schema.global_conditioning_fields + } + assert "scenario" not in fields + assert { + FIELD_PROMPT, + FIELD_FIRST_FRAME_PATH, + FIELD_TOTAL_BLOCKS, + FIELD_PIXEL_HEIGHT, + FIELD_PIXEL_WIDTH, + FIELD_FPS, + }.issubset(fields) + # Camera control is per-step model input, not session-global scenario data. + step_fields = { + field.name for field in adapter.inference_input_schema.step_fields + } + assert step_fields == {FIELD_CAMERA_TRAJECTORY, FIELD_CAMERA_INTRINSICS} + + +def test_lingbot_replay_demo_uses_shared_runner(tmp_path: Path) -> None: + image = tmp_path / "image.jpg" + poses = tmp_path / "poses.npy" + intrinsics = tmp_path / "intrinsics.npy" + image.write_bytes(b"fake") + _write_camera_assets(poses, intrinsics) + pipeline_config = object() + adapter = LingbotDemoAdapter() + output = _RecordingOutputTarget() + calls: list[dict[str, Any]] = [] + + def fake_runner(**kwargs: Any) -> Sequence[OutputArtifact]: + calls.append(kwargs) + return (OutputArtifact(kind="video/mp4", uri="memory://lingbot"),) + + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + input_mode="replay", + scenario={ + "prompt": "drive through a city", + "image_path": image, + "pose_path": poses, + "intrinsic_path": intrinsics, + "total_blocks": 1, + }, + output=Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=16, output_layout="tchw"), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + runtime_options={"pipeline_config": pipeline_config}, + ), + ) + + artifacts = run_replay_demo( + spec=spec, + adapter=adapter, + output_target_factory=lambda output_spec: output, + runner=fake_runner, + ) + + assert artifacts == (OutputArtifact(kind="video/mp4", uri="memory://lingbot"),) + assert len(calls) == 1 + assert calls[0]["adapter"] is adapter + assert calls[0]["config"] == spec.config + inputs = calls[0]["initial_inputs"].global_conditioning + assert inputs[FIELD_PROMPT] == "drive through a city" + assert inputs[FIELD_FIRST_FRAME_PATH] == image + assert inputs[FIELD_TOTAL_BLOCKS] == 1 + + +def test_lingbot_replay_invalid_scenario_fails_before_runtime_creation( + tmp_path: Path, +) -> None: + adapter = LingbotDemoAdapter( + replay_runtime_factory=lambda **kwargs: pytest.fail( + f"runtime should not be created: {kwargs}" + ) + ) + output_factory_calls = 0 + + def output_factory(output_spec: object) -> OutputTarget: + nonlocal output_factory_calls + del output_spec + output_factory_calls += 1 + return _RecordingOutputTarget() + + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + input_mode="replay", + scenario={ + "prompt": "drive", + "image_path": tmp_path / "missing.jpg", + "pose_path": tmp_path / "missing-poses.npy", + "intrinsic_path": tmp_path / "missing-intrinsics.npy", + }, + output=Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=16), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + runtime_options={"pipeline_config": object()}, + ), + ) + + with pytest.raises(FileNotFoundError, match=f"missing {FIELD_FIRST_FRAME_PATH}"): + run_replay_demo( + spec=spec, + adapter=adapter, + output_target_factory=output_factory, + ) + + assert output_factory_calls == 0 + + +def test_lingbot_replay_cli_defaults_to_example_data( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lingbot.runtime as runtime_module + + example_dir = tmp_path / "example" + example_dir.mkdir() + (example_dir / "image.jpg").write_bytes(b"fake") + _write_camera_assets( + example_dir / "poses.npy", example_dir / "intrinsics.npy" + ) + (example_dir / "prompt.txt").write_text("drive through a forest\n") + downloaded: list[int] = [] + + def fake_download(*, is_rank_zero: bool, example_idx: int) -> Path: + assert is_rank_zero is True + downloaded.append(example_idx) + return example_dir + + monkeypatch.setattr( + runtime_module, + "ensure_example_data_downloaded", + fake_download, + ) + args = parse_args(["replay", "--output", str(tmp_path / "demo.mp4")]) + spec = _replay_spec(args) + + prepared = LingbotDemoAdapter().prepare_scenario(spec) + + inputs = prepared.initial_inputs.global_conditioning + assert downloaded == [0] + assert inputs[FIELD_FIRST_FRAME_PATH] == example_dir / "image.jpg" + assert inputs[FIELD_PROMPT] == "drive through a forest" + + +def test_lingbot_replay_cli_can_disable_example_data(tmp_path: Path) -> None: + args = parse_args( + ["replay", "--no-example-data", "--output", str(tmp_path / "demo.mp4")] + ) + spec = _replay_spec(args) + + with pytest.raises(ValueError, match=f"require {FIELD_FIRST_FRAME_PATH}"): + LingbotDemoAdapter().prepare_scenario(spec) + + +def test_lingbot_replay_runtime_generates_video_step_result( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lingbot.runtime as runtime_module + + image = tmp_path / "image.jpg" + poses = tmp_path / "poses.npy" + intrinsics = tmp_path / "intrinsics.npy" + image.write_bytes(b"fake") + _write_camera_assets(poses, intrinsics, frames=16) + pipeline = _FakeLingbotPipeline() + monkeypatch.setattr( + runtime_module, + "load_first_frame_tensor", + lambda *args, **kwargs: torch.zeros(1, 3, 2, 2), + ) + + runtime = LingbotReplayRuntime( + config=InferenceConfig(model_id=LINGBOT_MODEL_ID, device="cpu"), + options=LingbotReplayRuntimeOptions( + pipeline_config=object(), + pipeline_factory=lambda pipeline_config, device: pipeline, + ), + ) + replay_inputs = LingbotReplayInputs( + prompt="drive", + first_frame_path=image, + camera_poses_path=poses, + camera_intrinsics_path=intrinsics, + total_blocks=1, + pixel_height=2, + pixel_width=2, + fps=16, + ) + adapter = LingbotDemoAdapter() + mapping = adapter.create_input_mapping(replay_inputs) + initial_inputs = mapping.map_global_conditioning_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=inference_input_from_replay_inputs(replay_inputs), + ) + session = runtime.start_session(initial_inputs) + + request = session.next_step_request() + assert request is not None + assert request.step_index == 0 + # The session asks for exactly this chunk's slice of the input timeline. + assert request.user_input_window is not None + assert request.user_input_window.start_s == 0.0 + assert request.user_input_window.end_s == 1 / 16 + assert request.metadata["num_frames"] == 1 + + step_inputs = mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput(), + request=request, + ) + assert step_inputs.step[FIELD_CAMERA_TRAJECTORY].shape == (1, 4, 4) + assert step_inputs.step[FIELD_CAMERA_INTRINSICS].shape == (1, 4) + result = session.step(step_inputs) + + assert result.step_index == 0 + assert result.frame_count == 1 + assert isinstance(result.output, VideoStepResult) + assert result.output.layout == "tchw" + assert result.output.video_chunk.shape == (1, 3, 2, 2) + assert result.output_window is not None + assert result.output_window.start_s == 0.0 + assert result.output_window.end_s == 1 / 16 + assert result.metrics["denoise_s"] == 0.25 + assert session.next_step_request() is None + assert pipeline.initialize_cache_calls == [ + {"text": ["drive"], "image_shape": (1, 3, 2, 2)} + ] + assert pipeline.generate_calls == [ + { + "autoregressive_index": 0, + "intrinsics_shape": (1, 4), + "poses_shape": (1, 4, 4), + "world_scale": pytest.approx(mapping.camera_trace.world_scale), + } + ] + runtime.close() + + +def test_lingbot_webrtc_cli_builds_keyboard_driving_spec() -> None: + args = parse_args( + [ + "webrtc", + "--host", + "127.0.0.1", + "--port", + "9090", + "--device", + "cuda:2", + "--seed", + "123", + "--no-compile", + "--fps", + "12", + "--video-height", + "32", + "--video-width", + "64", + "--warmup-chunks", + "0", + "--warmup-timeout-s", + "1.5", + "--client-liveness-timeout-s", + "2.5", + "--prefer-sw-encoder", + "--example-idx", + "2", + ] + ) + + spec = _webrtc_spec(args, device="cuda:3", context_parallel_size=4) + + assert spec.model_id == LINGBOT_MODEL_ID + assert spec.preset_id == DEFAULT_LINGBOT_PRESET + assert spec.input_mode == "keyboard-driving" + assert isinstance(spec.scenario, LingbotWebRTCScenario) + assert spec.scenario.example_idx == 2 + assert spec.scenario.prefer_sw_encoder is True + assert isinstance(spec.output, WebRTCOutputSpec) + assert spec.output.host == "127.0.0.1" + assert spec.output.port == 9090 + assert spec.output.fps == 12 + assert spec.output.video_width == 64 + assert spec.output.video_height == 32 + assert spec.output.warmup_chunks == 0 + assert spec.output.warmup_timeout_s == 1.5 + assert spec.output.client_liveness_timeout_s == 2.5 + assert spec.config is not None + assert spec.config.device == "cuda:3" + assert spec.config.compile is False + assert spec.config.runtime_options["seed"] == 123 + assert spec.config.runtime_options["context_parallel_size"] == 4 + assert spec.config.runtime_options["example_idx"] == 2 + + +def test_lingbot_webrtc_demo_uses_existing_manager_with_model_config() -> None: + pipeline_config = object() + adapter = LingbotDemoAdapter(webrtc_runtime_factory=_FakeWebRTCRuntime) + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + input_mode="keyboard-driving", + scenario=LingbotWebRTCScenario(example_idx=2, prefer_sw_encoder=True), + output=WebRTCOutputSpec( + host="0.0.0.0", + port=8080, + fps=24, + video_width=64, + video_height=32, + warmup_chunks=0, + warmup_timeout_s=1.0, + ), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + device="cuda:7", + runtime_options={"pipeline_config": pipeline_config, "seed": 123}, + ), + ) + + demo = build_webrtc_demo(spec=spec, adapter=adapter) + + assert isinstance(demo.runtime, _FakeWebRTCRuntime) + assert isinstance(demo.session_manager, LingbotDemoWebRTCSessionManager) + assert demo.session_manager._runtime is demo.runtime + assert demo.session_manager.runtime_config is demo.runtime.config + assert demo.runtime_config is demo.runtime.config + assert demo.runtime_config.pipeline_config is pipeline_config + assert demo.runtime_config.config_name == DEFAULT_LINGBOT_PRESET + assert demo.runtime_config.seed == 123 + assert demo.runtime_config.device == "cuda:7" + assert demo.runtime_config.video_width == 64 + assert demo.runtime_config.video_height == 32 + assert demo.runtime_config.fps == 24 + assert demo.runtime_config.encoder_backend == "default" + assert demo.runtime_config.example_data_dir.name == "02" + assert demo.session_manager._model_name() == DEFAULT_LINGBOT_PRESET + assert demo.host == "0.0.0.0" + assert demo.port == 8080 + + +def test_lingbot_webrtc_demo_installs_model_routes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lingbot.demo.webrtc as demo_webrtc_module + + app_calls: list[dict[str, Any]] = [] + + async def _ok(request: web.Request) -> web.Response: + del request + return web.Response(text="ok") + + def fake_create_app(**kwargs: Any) -> web.Application: + app_calls.append(kwargs) + app = web.Application() + app[SESSION_MANAGER_KEY] = kwargs["session_manager"] + app.router.add_get("/api/session/initial_scene", _ok) + app.router.add_get("/api/session/first_frame", _ok) + app.router.add_post("/api/session/input", _ok) + return app + + monkeypatch.setattr(demo_webrtc_module, "create_app", fake_create_app) + adapter = LingbotDemoAdapter(webrtc_runtime_factory=_FakeWebRTCRuntime) + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + input_mode="keyboard-driving", + scenario=LingbotWebRTCScenario(), + output=WebRTCOutputSpec( + host="0.0.0.0", + port=8080, + warmup_timeout_s=1.0, + preload_name="Test Lingbot", + ), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + runtime_options={"pipeline_config": object()}, + ), + ) + + demo = build_webrtc_demo(spec=spec, adapter=adapter, create_app=True) + + assert demo.app is not None + assert app_calls[0]["session_manager"] is demo.session_manager + assert app_calls[0]["request_session_url"] == ( + "http://127.0.0.1:8080/request_session" + ) + route_paths = {resource.canonical for resource in demo.app.router.resources()} + assert "/api/session/initial_scene" in route_paths + assert "/api/session/first_frame" in route_paths + assert "/api/session/input" in route_paths + + +def test_lingbot_webrtc_demo_serves_through_shared_runner( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lingbot.demo.webrtc as demo_webrtc_module + + server_calls: list[dict[str, Any]] = [] + + def fake_create_app(**kwargs: Any) -> web.Application: + app = web.Application() + app[SESSION_MANAGER_KEY] = kwargs["session_manager"] + return app + + def fake_server_runner(**kwargs: Any) -> None: + server_calls.append(kwargs) + + monkeypatch.setattr(demo_webrtc_module, "create_app", fake_create_app) + adapter = LingbotDemoAdapter(webrtc_runtime_factory=_FakeWebRTCRuntime) + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + input_mode="keyboard-driving", + scenario={"example_idx": 0}, + output=WebRTCOutputSpec( + host="0.0.0.0", + port=8080, + warmup_timeout_s=1.0, + ), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + runtime_options={"pipeline_config": object()}, + ), + ) + + demo = cast( + WebRTCDemo, + serve_flashdreams_demo( + spec=spec, + adapter=adapter, + world_rank=0, + server_runner=fake_server_runner, + ), + ) + + assert len(server_calls) == 1 + assert server_calls[0]["world_rank"] == 0 + assert server_calls[0]["session_manager"] is demo.session_manager + assert server_calls[0]["app"] is demo.app + assert server_calls[0]["host"] == "0.0.0.0" + assert server_calls[0]["port"] == 8080 + assert isinstance(demo.session_manager, LingbotDemoWebRTCSessionManager) + + +class _RecordingOutputTarget: + def open(self) -> None: + return None + + def write(self, result: StepResult) -> None: + del result + + def close(self) -> Sequence[OutputArtifact]: + return () + + +class _FakeLingbotPipeline: + def __init__(self) -> None: + self.initialize_cache_calls: list[dict[str, Any]] = [] + self.generate_calls: list[dict[str, Any]] = [] + + def initialize_cache(self, *, text: list[str], image: torch.Tensor) -> object: + self.initialize_cache_calls.append( + { + "text": text, + "image_shape": tuple(image.shape), + } + ) + return object() + + def get_num_output_frames(self, autoregressive_index: int) -> int: + del autoregressive_index + return 1 + + def generate( + self, + *, + autoregressive_index: int, + cache: object, + input: Any, + ) -> torch.Tensor: + del cache + self.generate_calls.append( + { + "autoregressive_index": autoregressive_index, + "intrinsics_shape": tuple(input.intrinsics.shape), + "poses_shape": tuple(input.poses.shape), + "world_scale": input.world_scale, + } + ) + return torch.full((1, 3, 2, 2), float(autoregressive_index)) + + def finalize(self, *, autoregressive_index: int, cache: object) -> dict[str, float]: + del autoregressive_index, cache + return {"denoise_s": 0.25} + + +class _FakeWebRTCRuntime: + def __init__(self, config: LingbotRuntimeConfig) -> None: + self.config = config + + async def initialize(self) -> None: + return None + + async def reset_for_new_session(self, *args: Any, **kwargs: Any) -> None: + return None + + def peek_steady_chunk_num_frames(self) -> int: + return 1 + + def peek_next_chunk_num_frames(self) -> int: + return 1 + + async def generate_chunk( + self, + *, + segments: list[Any], + frame_times: list[float], + ) -> Any: + del segments, frame_times + return None + + async def close(self) -> None: + return None + + def send_exit_signal(self) -> None: + return None + + def wait_for_termination(self) -> None: + return None diff --git a/integrations/lingbot/tests/test_input_mapping.py b/integrations/lingbot/tests/test_input_mapping.py new file mode 100644 index 00000000..1db76bf4 --- /dev/null +++ b/integrations/lingbot/tests/test_input_mapping.py @@ -0,0 +1,448 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import torch +from lingbot.demo.spec import resolve_text_event_prompts, resolve_user_input_events +from lingbot.input_mapping import ( + CAMERA_COMMAND, + FIELD_CAMERA_INTRINSICS, + FIELD_CAMERA_TRAJECTORY, + TEXT_EVENT, + KeyboardToCameraCommand, + LingbotInputMapping, + TextEventSelection, + load_camera_trace, +) + +from flashdreams.runtime import ( + CanonicalInputs, + InferenceInput, + InputCanonicalizer, + StepRequest, + TimeWindow, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, +) + +pytestmark = pytest.mark.ci_cpu + +_KEYBOARD_SOURCE = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="key_down", payload_fields=frozenset({"key"}) + ), + UserInputCapability(event_type="key_up", payload_fields=frozenset({"key"})), + UserInputCapability( + event_type="text_event", payload_fields=frozenset({"event_id"}) + ), + ) +) + + +def _step_request(*, step_index: int, frame_start: int, num_frames: int, fps: int = 16): + return StepRequest( + step_index=step_index, + user_input_window=TimeWindow( + start_s=frame_start / fps, + end_s=(frame_start + num_frames) / fps, + ), + metadata={"num_frames": num_frames, "frame_start": frame_start}, + ) + + +def _live_mapping(**kwargs) -> LingbotInputMapping: + return LingbotInputMapping( + fps=16, + base_intrinsics=torch.tensor([416.0, 416.0, 416.0, 240.0]), + world_scale=1.0, + **kwargs, + ) + + +def test_keyboard_events_become_camera_command_axes() -> None: + converter = KeyboardToCameraCommand() + inputs = UserInputs( + events=( + UserInputEvent(timestamp_s=0.0, event_type="key_down", payload={"key": "w"}), + ) + ) + window = TimeWindow(start_s=0.0, end_s=1.0) + + value = converter.convert(inputs.window(window), window) + + assert value is not None + assert value["move_forward"] == 1.0 + assert value["yaw"] == 0.0 + # Level-triggered: a key held across the next window still means forward. + next_window = TimeWindow(start_s=1.0, end_s=2.0) + held = converter.convert(UserInputs().window(next_window), next_window) + assert held is not None + assert held["move_forward"] == 1.0 + + +def test_camera_command_segments_preserve_sub_window_timing() -> None: + converter = KeyboardToCameraCommand() + window = TimeWindow(start_s=0.0, end_s=1.0) + inputs = UserInputs( + events=( + UserInputEvent(timestamp_s=0.5, event_type="key_down", payload={"key": "w"}), + ) + ) + + value = converter.convert(inputs.window(window), window) + + assert value is not None + segments = value["segments"] + assert [(start, end) for start, end, _ in segments] == [(0.0, 0.5), (0.5, 1.0)] + assert segments[0][2]["move_forward"] == 0.0 + assert segments[1][2]["move_forward"] == 1.0 + + +def test_key_events_drive_a_camera_trajectory() -> None: + canonicalizer = InputCanonicalizer([KeyboardToCameraCommand()]) + mapping = _live_mapping() + user_inputs = UserInputs( + events=( + UserInputEvent(timestamp_s=0.0, event_type="key_down", payload={"key": "w"}), + ) + ) + request = _step_request(step_index=0, frame_start=0, num_frames=4) + assert request.user_input_window is not None + + step_inputs = mapping.map_step_inputs( + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=request.user_input_window, + source_schema=_KEYBOARD_SOURCE, + ), + inference_input=InferenceInput(), + request=request, + ) + + poses = step_inputs.step[FIELD_CAMERA_TRAJECTORY] + assert poses.shape == (4, 4, 4) + assert step_inputs.step[FIELD_CAMERA_INTRINSICS].shape == (4, 4) + # Holding forward has to actually move the camera along the trajectory. + assert not torch.allclose(poses[0], poses[-1]) + assert poses[-1][:3, 3].abs().sum() > 0 + + +def test_idle_keyboard_leaves_the_camera_stationary() -> None: + canonicalizer = InputCanonicalizer([KeyboardToCameraCommand()]) + mapping = _live_mapping() + request = _step_request(step_index=0, frame_start=0, num_frames=4) + assert request.user_input_window is not None + + step_inputs = mapping.map_step_inputs( + canonical_inputs=canonicalizer.canonicalize( + UserInputs(), + window=request.user_input_window, + source_schema=_KEYBOARD_SOURCE, + ), + inference_input=InferenceInput(), + request=request, + ) + + poses = step_inputs.step[FIELD_CAMERA_TRAJECTORY] + assert torch.allclose(poses[0], poses[-1]) + + +def test_text_event_becomes_a_global_conditioning_prompt_update() -> None: + canonicalizer = InputCanonicalizer( + [KeyboardToCameraCommand(), TextEventSelection()] + ) + mapping = _live_mapping(text_event_prompts={"storm": "a violent storm"}) + mapping.set_base_prompt("a calm street") + user_inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, + event_type="text_event", + payload={"event_id": "storm"}, + ), + ) + ) + request = _step_request(step_index=0, frame_start=0, num_frames=4) + assert request.user_input_window is not None + + step_inputs = mapping.map_step_inputs( + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=request.user_input_window, + source_schema=_KEYBOARD_SOURCE, + ), + inference_input=InferenceInput(), + request=request, + ) + + assert step_inputs.global_conditioning["prompt"] == "a violent storm" + + # The swap is requested once, not re-sent on every later step. + next_request = _step_request(step_index=1, frame_start=4, num_frames=4) + assert next_request.user_input_window is not None + held = mapping.map_step_inputs( + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=next_request.user_input_window, + source_schema=_KEYBOARD_SOURCE, + ), + inference_input=InferenceInput(), + request=next_request, + ) + assert held.global_conditioning == {} + + +def test_clearing_a_text_event_restores_the_base_prompt() -> None: + converter = TextEventSelection() + window = TimeWindow(start_s=0.0, end_s=1.0) + triggered = converter.convert( + UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, + event_type="text_event", + payload={"event_id": "storm"}, + ), + ) + ), + window, + ) + assert triggered is not None and triggered["event_id"] == "storm" + + cleared = converter.convert( + UserInputs( + events=( + UserInputEvent( + timestamp_s=0.5, + event_type="text_event", + payload={"event_id": "storm", "state": "clear"}, + ), + ) + ), + window, + ) + assert cleared is not None and cleared["event_id"] is None + + +def test_unknown_text_event_is_rejected_by_the_mapping() -> None: + canonicalizer = InputCanonicalizer( + [KeyboardToCameraCommand(), TextEventSelection()] + ) + mapping = _live_mapping(text_event_prompts={"storm": "a violent storm"}) + request = _step_request(step_index=0, frame_start=0, num_frames=4) + assert request.user_input_window is not None + canonical = canonicalizer.canonicalize( + UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, + event_type="text_event", + payload={"event_id": "volcano"}, + ), + ) + ), + window=request.user_input_window, + source_schema=_KEYBOARD_SOURCE, + ) + + with pytest.raises(ValueError, match="Unknown Lingbot text event_id"): + mapping.map_step_inputs( + canonical_inputs=canonical, + inference_input=InferenceInput(), + request=request, + ) + + +def test_live_mapping_requires_camera_command_from_the_source() -> None: + mapping = _live_mapping() + + with pytest.raises(ValueError, match="requires a 'camera_command'"): + mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput(), + request=_step_request(step_index=0, frame_start=0, num_frames=4), + ) + + +def test_trace_mapping_slices_successive_chunks(tmp_path: Path) -> None: + poses_path = tmp_path / "poses.npy" + intrinsics_path = tmp_path / "intrinsics.npy" + trajectory = np.tile(np.eye(4, dtype=np.float32), (32, 1, 1)) + trajectory[:, 2, 3] = np.arange(32, dtype=np.float32) + np.save(poses_path, trajectory) + np.save( + intrinsics_path, + np.tile(np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (32, 1)), + ) + trace = load_camera_trace( + camera_poses_path=poses_path, + camera_intrinsics_path=intrinsics_path, + pixel_height=464, + pixel_width=832, + intrinsics_reference_height=480, + intrinsics_reference_width=832, + ) + mapping = LingbotInputMapping(fps=16, trace=trace) + + first = mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput(), + request=_step_request(step_index=0, frame_start=0, num_frames=4), + ) + second = mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput(), + request=_step_request(step_index=1, frame_start=4, num_frames=4), + ) + + assert first.step[FIELD_CAMERA_TRAJECTORY].shape == (4, 4, 4) + # Consecutive steps must advance through the trace, not restart it. + assert not torch.allclose( + first.step[FIELD_CAMERA_TRAJECTORY], second.step[FIELD_CAMERA_TRAJECTORY] + ) + assert mapping.mapping_schema.consumes == () + + +def test_trace_mapping_reports_running_past_the_end(tmp_path: Path) -> None: + poses_path = tmp_path / "poses.npy" + intrinsics_path = tmp_path / "intrinsics.npy" + np.save(poses_path, np.tile(np.eye(4, dtype=np.float32), (16, 1, 1))) + np.save( + intrinsics_path, + np.tile(np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (16, 1)), + ) + mapping = LingbotInputMapping( + fps=16, + trace=load_camera_trace( + camera_poses_path=poses_path, + camera_intrinsics_path=intrinsics_path, + pixel_height=464, + pixel_width=832, + intrinsics_reference_height=480, + intrinsics_reference_width=832, + ), + ) + + with pytest.raises(ValueError, match="camera trace has"): + mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput(), + request=_step_request(step_index=0, frame_start=0, num_frames=999), + ) + + +def test_scenario_event_trace_resolves_into_user_inputs() -> None: + user_inputs = resolve_user_input_events( + { + "events": [ + {"t": 1.5, "type": "key_down", "key": "a"}, + {"t": 0.0, "type": "key_down", "key": "w"}, + {"t": 2.0, "type": "text_event", "event_id": "storm"}, + ] + } + ) + + # UserInputs requires non-decreasing timestamps, so resolution must sort. + assert [event.timestamp_s for event in user_inputs.events] == [0.0, 1.5, 2.0] + assert user_inputs.events[0].payload == {"key": "w"} + assert user_inputs.events[2].event_type == "text_event" + + +def test_scenario_text_event_catalog_resolves() -> None: + assert resolve_text_event_prompts({"text_events": {"storm": "a storm"}}) == { + "storm": "a storm" + } + assert resolve_text_event_prompts( + {"text_events": [{"event_id": "portal", "prompt": "a glowing portal"}]} + ) == {"portal": "a glowing portal"} + assert resolve_text_event_prompts(None) == {} + + +def test_declared_modalities_match_what_the_converters_produce() -> None: + canonicalizer = InputCanonicalizer( + [KeyboardToCameraCommand(), TextEventSelection()] + ) + + schema = canonicalizer.canonical_schema(_KEYBOARD_SOURCE) + + assert schema.supports(CAMERA_COMMAND) + assert schema.supports(TEXT_EVENT) + # A source with no key events cannot feed the keyboard converter. + empty = canonicalizer.canonical_schema(UserInputSchema()) + assert not empty.supports(CAMERA_COMMAND) + + +def test_event_driven_scenario_builds_a_live_mapping(tmp_path: Path) -> None: + """A scenario can drive the camera from events instead of the pose trace.""" + from lingbot.demo.adapter import LingbotDemoAdapter + from lingbot.runtime import LINGBOT_MODEL_ID + + from flashdreams.runtime import InferenceConfig + from flashdreams.runtime.demo import DemoSpec, Mp4OutputSpec + + image = tmp_path / "image.jpg" + image.write_bytes(b"fake") + poses_path = tmp_path / "poses.npy" + intrinsics_path = tmp_path / "intrinsics.npy" + np.save(poses_path, np.tile(np.eye(4, dtype=np.float32), (32, 1, 1))) + np.save( + intrinsics_path, + np.tile(np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (32, 1)), + ) + + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + input_mode="replay", + scenario={ + "prompt": "a calm street", + "image_path": image, + "pose_path": poses_path, + "intrinsic_path": intrinsics_path, + "camera_source": "events", + "text_events": {"storm": "a violent storm"}, + "events": [ + {"t": 0.0, "type": "key_down", "key": "w"}, + {"t": 0.2, "type": "text_event", "event_id": "storm"}, + ], + }, + output=Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=16), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + runtime_options={"pipeline_config": object()}, + ), + ) + + prepared = LingbotDemoAdapter().prepare_scenario(spec) + + assert prepared.mapping is not None + assert prepared.mapping.mapping_schema.consumes == (CAMERA_COMMAND, TEXT_EVENT) + assert len(prepared.user_inputs.events) == 2 + # The declared source must actually cover the trace it carries, or the + # canonicalizer silently drops the keyboard converter. + canonical_schema = prepared.canonicalizer.canonical_schema(prepared.source_schema) + assert canonical_schema.supports(CAMERA_COMMAND) + assert canonical_schema.supports(TEXT_EVENT) + + request = _step_request(step_index=0, frame_start=0, num_frames=4) + assert request.user_input_window is not None + step_inputs = prepared.mapping.map_step_inputs( + canonical_inputs=prepared.canonicalizer.canonicalize( + prepared.user_inputs, + window=request.user_input_window, + source_schema=prepared.source_schema, + ), + inference_input=InferenceInput(), + request=request, + ) + poses = step_inputs.step[FIELD_CAMERA_TRAJECTORY] + assert poses.shape == (4, 4, 4) + assert not torch.allclose(poses[0], poses[-1]) + assert step_inputs.global_conditioning["prompt"] == "a violent storm" diff --git a/integrations/lingbot/tests/test_keyboard_parity.py b/integrations/lingbot/tests/test_keyboard_parity.py new file mode 100644 index 00000000..b233d726 --- /dev/null +++ b/integrations/lingbot/tests/test_keyboard_parity.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Live camera control must match the WebRTC path it will eventually replace. + +The WebRTC runtime drives the camera with ``KeyboardResampler.sample_chunk`` +feeding ``CameraPoseIntegrator``. The runtime-API path instead windows events +with ``StepRequest.user_input_window``, canonicalizes them into camera intent, +and integrates that. Both should produce the same trajectory for the same key +stream; these tests pin that, so a divergence shows up here rather than as +different handling in a live session. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from lingbot.input_mapping import ( + FIELD_CAMERA_TRAJECTORY, + KeyboardToCameraCommand, + LingbotInputMapping, +) + +from flashdreams.runtime import ( + InferenceInput, + InputCanonicalizer, + StepRequest, + TimeWindow, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, +) +from flashdreams.serving.webrtc.controls import CameraPoseIntegrator, KeyboardResampler + +pytestmark = pytest.mark.ci_cpu + +_FPS = 16 +_NUM_FRAMES = 4 + +_SOURCE = UserInputSchema( + capabilities=( + UserInputCapability(event_type="key_down", payload_fields=frozenset({"key"})), + UserInputCapability(event_type="key_up", payload_fields=frozenset({"key"})), + ) +) + +# (timestamp_s, event_type, key) +Edge = tuple[float, str, str] + +_STREAMS: dict[str, list[Edge]] = { + "hold_forward": [(0.0, "key_down", "w")], + "forward_then_release": [(0.0, "key_down", "w"), (0.1, "key_up", "w")], + "mid_chunk_turn": [(0.0, "key_down", "w"), (0.13, "key_down", "a")], + "strafe_and_pitch": [(0.02, "key_down", "e"), (0.09, "key_down", "i")], + "alternate_yaw_keys": [(0.0, "key_down", "w"), (0.05, "key_down", "j")], + "conflicting_yaw": [(0.0, "key_down", "a"), (0.07, "key_down", "d")], + "rapid_toggle": [ + (0.01, "key_down", "w"), + (0.04, "key_up", "w"), + (0.08, "key_down", "w"), + (0.2, "key_up", "w"), + ], + "idle": [], + # KeyboardResampler drains events with `event_t <= chunk_end` while + # TimeWindow is half-open, so an edge landing exactly on a chunk boundary + # is the most likely place for the two paths to disagree. + "exact_chunk_boundary": [(0.0, "key_down", "w"), (0.25, "key_down", "a")], + "boundary_release": [(0.0, "key_down", "w"), (0.25, "key_up", "w")], + "second_boundary": [(0.0, "key_down", "w"), (0.5, "key_down", "d")], +} + + +def _legacy_poses(edges: list[Edge], *, chunks: int) -> np.ndarray: + """Integrate a key stream the way the WebRTC session does.""" + resampler = KeyboardResampler(fps=_FPS, start_v=0.0) + integrator = CameraPoseIntegrator() + for timestamp_s, event_type, key in edges: + resampler.on_edge( + arrival_t=timestamp_s, + event="keydown" if event_type == "key_down" else "keyup", + key=key, + ) + poses = [] + for _ in range(chunks): + segments, frame_times = resampler.sample_chunk(_NUM_FRAMES) + poses.append( + integrator.integrate_chunk(segments=segments, frame_times=frame_times) + ) + return np.concatenate(poses) + + +def _runtime_api_poses(edges: list[Edge], *, chunks: int) -> np.ndarray: + """Integrate the same key stream through the runtime API input path.""" + canonicalizer = InputCanonicalizer([KeyboardToCameraCommand()]) + mapping = LingbotInputMapping( + fps=_FPS, + base_intrinsics=torch.tensor([416.0, 416.0, 416.0, 240.0]), + world_scale=1.0, + ) + user_inputs = UserInputs( + events=tuple( + UserInputEvent( + timestamp_s=timestamp_s, + event_type=event_type, + payload={"key": key}, + ) + for timestamp_s, event_type, key in edges + ) + ) + poses = [] + for chunk_index in range(chunks): + frame_start = chunk_index * _NUM_FRAMES + request = StepRequest( + step_index=chunk_index, + user_input_window=TimeWindow( + start_s=frame_start / _FPS, + end_s=(frame_start + _NUM_FRAMES) / _FPS, + ), + metadata={"num_frames": _NUM_FRAMES, "frame_start": frame_start}, + ) + assert request.user_input_window is not None + step_inputs = mapping.map_step_inputs( + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=request.user_input_window, + source_schema=_SOURCE, + ), + inference_input=InferenceInput(), + request=request, + ) + poses.append(step_inputs.step[FIELD_CAMERA_TRAJECTORY].numpy()) + return np.concatenate(poses) + + +@pytest.mark.parametrize("name", sorted(_STREAMS)) +def test_single_chunk_matches_the_webrtc_path(name: str) -> None: + edges = _STREAMS[name] + + legacy = _legacy_poses(edges, chunks=1) + runtime_api = _runtime_api_poses(edges, chunks=1) + + assert legacy.shape == runtime_api.shape + np.testing.assert_allclose(runtime_api, legacy, atol=1e-5) + + +@pytest.mark.parametrize("name", sorted(_STREAMS)) +def test_multi_chunk_matches_the_webrtc_path(name: str) -> None: + """Carried key state across chunk boundaries must agree too.""" + edges = _STREAMS[name] + + legacy = _legacy_poses(edges, chunks=3) + runtime_api = _runtime_api_poses(edges, chunks=3) + + assert legacy.shape == runtime_api.shape + np.testing.assert_allclose(runtime_api, legacy, atol=1e-5) diff --git a/integrations/lingbot/tests/test_runtime_gpu.py b/integrations/lingbot/tests/test_runtime_gpu.py new file mode 100644 index 00000000..db88ccd3 --- /dev/null +++ b/integrations/lingbot/tests/test_runtime_gpu.py @@ -0,0 +1,259 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +import torch +from lingbot import runtime as runtime_module +from lingbot.runtime import ( + LINGBOT_MODEL_ID, + LingbotModelAdapter, + LingbotReplayInputs, + LingbotReplayRuntime, + LingbotReplayRuntimeOptions, + inference_input_from_replay_inputs, +) + +from flashdreams.infra.video_output import VideoStepResult +from flashdreams.runtime import CanonicalInputs, InferenceConfig, InferenceInput + +pytestmark = pytest.mark.ci_gpu + + +def test_lingbot_replay_runtime_accepts_direct_inputs_on_cuda( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Exercise the migrated Lingbot runtime API path with CUDA tensors.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA required.") + + image = tmp_path / "image.jpg" + poses = tmp_path / "poses.npy" + intrinsics = tmp_path / "intrinsics.npy" + image.write_bytes(b"fake") + trajectory = np.tile(np.eye(4, dtype=np.float32), (32, 1, 1)) + trajectory[:, 2, 3] = np.arange(32, dtype=np.float32) + np.save(poses, trajectory) + np.save( + intrinsics, + np.tile(np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (32, 1)), + ) + pipeline = _FakeCudaLingbotPipeline() + + def _fake_load_first_frame_tensor( + path: Path, + **kwargs: Any, + ) -> torch.Tensor: + assert path == image + return torch.zeros( + (1, 3, 2, 2), + device=kwargs["device"], + dtype=kwargs["dtype"], + ) + + monkeypatch.setattr( + runtime_module, + "load_first_frame_tensor", + _fake_load_first_frame_tensor, + ) + runtime = LingbotReplayRuntime( + config=InferenceConfig(model_id=LINGBOT_MODEL_ID, device="cuda"), + options=LingbotReplayRuntimeOptions( + pipeline_config=object(), + pipeline_factory=lambda _pipeline_config, _device: pipeline, + ), + ) + replay_inputs = LingbotReplayInputs( + prompt="drive", + first_frame_path=image, + camera_poses_path=poses, + camera_intrinsics_path=intrinsics, + total_blocks=1, + pixel_height=2, + pixel_width=2, + fps=16, + ) + # Camera inputs now reach the session per step through the mapping, so the + # GPU path has to be driven the same way the standard loop drives it. + mapping = LingbotModelAdapter().create_input_mapping(replay_inputs) + session = runtime.start_session( + mapping.map_global_conditioning_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=inference_input_from_replay_inputs(replay_inputs), + ) + ) + try: + request = session.next_step_request() + assert request is not None + result = session.step( + mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput(), + request=request, + ) + ) + torch.cuda.synchronize() + finally: + session.close() + runtime.close() + + assert result.frame_count == 1 + assert isinstance(result.output, VideoStepResult) + assert result.output.video_chunk.is_cuda + assert result.output.video_chunk.shape == (1, 3, 2, 2) + assert pipeline.initialize_cache_devices == ["cuda"] + assert pipeline.generate_world_scales == [mapping.camera_trace.world_scale] + + +class _FakeCudaLingbotPipeline: + def __init__(self) -> None: + self.initialize_cache_devices: list[str] = [] + self.generate_world_scales: list[float] = [] + + def initialize_cache(self, *, text: list[str], image: torch.Tensor) -> object: + assert text == ["drive"] + assert image.is_cuda + self.initialize_cache_devices.append(image.device.type) + return object() + + def get_num_output_frames(self, autoregressive_index: int) -> int: + assert autoregressive_index == 0 + return 1 + + def generate( + self, + *, + autoregressive_index: int, + cache: object, + input: Any, + ) -> torch.Tensor: + del cache + assert autoregressive_index == 0 + assert input.intrinsics.is_cuda + assert input.poses.is_cuda + self.generate_world_scales.append(input.world_scale) + return torch.zeros( + (1, 3, 2, 2), + device=input.intrinsics.device, + dtype=torch.bfloat16, + ) + + def finalize(self, *, autoregressive_index: int, cache: object) -> dict[str, float]: + del autoregressive_index, cache + return {"denoise_s": 0.25} + + +def test_event_driven_camera_control_on_cuda( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Drive the CUDA session from key events instead of a fixed pose trace.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA required.") + + from lingbot.input_mapping import ( + FIELD_CAMERA_TRAJECTORY, + KeyboardToCameraCommand, + ) + + from flashdreams.runtime import ( + InputCanonicalizer, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, + ) + + image = tmp_path / "image.jpg" + image.write_bytes(b"fake") + pipeline = _FakeCudaLingbotPipeline() + + def _fake_load_first_frame_tensor(path: Path, **kwargs: Any) -> torch.Tensor: + del path + return torch.zeros((1, 3, 2, 2), device=kwargs["device"], dtype=kwargs["dtype"]) + + monkeypatch.setattr( + runtime_module, + "load_first_frame_tensor", + _fake_load_first_frame_tensor, + ) + + runtime = LingbotReplayRuntime( + config=InferenceConfig(model_id=LINGBOT_MODEL_ID, device="cuda"), + options=LingbotReplayRuntimeOptions( + pipeline_config=object(), + pipeline_factory=lambda _pipeline_config, _device: pipeline, + ), + ) + adapter = LingbotModelAdapter() + mapping = adapter.create_live_input_mapping( + fps=16, + base_intrinsics=torch.tensor([416.0, 416.0, 416.0, 240.0]), + world_scale=1.0, + prompt="drive", + ) + initial_inputs = mapping.map_global_conditioning_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput( + global_conditioning={ + "prompt": "drive", + "first_frame_path": image, + "total_blocks": 1, + "pixel_height": 2, + "pixel_width": 2, + "fps": 16, + } + ), + ) + canonicalizer = InputCanonicalizer([KeyboardToCameraCommand()]) + source = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="key_down", payload_fields=frozenset({"key"}) + ), + UserInputCapability( + event_type="key_up", payload_fields=frozenset({"key"}) + ), + ) + ) + user_inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, event_type="key_down", payload={"key": "w"} + ), + ) + ) + + session = runtime.start_session(initial_inputs) + try: + request = session.next_step_request() + assert request is not None + assert request.user_input_window is not None + step_inputs = mapping.map_step_inputs( + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=request.user_input_window, + source_schema=source, + ), + inference_input=InferenceInput(), + request=request, + ) + # Holding forward must produce real motion before it reaches the model. + # This chunk is one frame, so compare against the identity start pose + # rather than across frames: 0.8 m/s at 16fps advances 0.05 along +z. + poses = step_inputs.step[FIELD_CAMERA_TRAJECTORY] + assert poses[-1][2, 3].item() == pytest.approx(0.05, abs=1e-4) + result = session.step(step_inputs) + torch.cuda.synchronize() + finally: + session.close() + runtime.close() + + assert result.output.video_chunk.is_cuda + assert pipeline.generate_world_scales == [1.0] diff --git a/integrations/lingbot/tests/test_runtime_session_inputs.py b/integrations/lingbot/tests/test_runtime_session_inputs.py new file mode 100644 index 00000000..def9cfb9 --- /dev/null +++ b/integrations/lingbot/tests/test_runtime_session_inputs.py @@ -0,0 +1,395 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The replay session must consume its per-step inputs, not ignore them.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest +import torch +import lingbot.runtime as runtime_module +from lingbot.input_mapping import FIELD_CAMERA_INTRINSICS, FIELD_CAMERA_TRAJECTORY +from lingbot.runtime import ( + LINGBOT_MODEL_ID, + LingbotReplayRuntime, + LingbotReplayRuntimeOptions, + LingbotSessionInputs, +) + +from flashdreams.runtime import InferenceConfig, InferenceInput + +pytestmark = pytest.mark.ci_cpu + + +class _FakePipeline: + """Records what the session hands the model.""" + + def __init__(self, *, supports_text_swap: bool = True) -> None: + self.generate_calls: list[dict[str, Any]] = [] + self.text_encoder_calls: list[list[str]] = [] + self.encoders_loaded = 0 + self.diffusion_model = _FakeDiffusionModel( + supports_text_swap=supports_text_swap + ) + + def initialize_cache(self, *, text: list[str], image: torch.Tensor) -> Any: + del text, image + return _FakeCache() + + def get_num_output_frames(self, autoregressive_index: int) -> int: + del autoregressive_index + return 2 + + def generate( + self, + *, + autoregressive_index: int, + cache: object, + input: Any, + ) -> torch.Tensor: + del cache + self.generate_calls.append( + { + "autoregressive_index": autoregressive_index, + "poses": input.poses.clone(), + "intrinsics": input.intrinsics.clone(), + "world_scale": input.world_scale, + } + ) + return torch.zeros(2, 3, 2, 2) + + def finalize(self, *, autoregressive_index: int, cache: object) -> dict[str, float]: + del autoregressive_index, cache + return {"denoise_s": 0.1} + + def _ensure_oneshot_encoders_loaded(self) -> None: + self.encoders_loaded += 1 + + def text_encoder(self, texts: list[str]) -> torch.Tensor: + self.text_encoder_calls.append(list(texts)) + return torch.ones(1, 4) + + +class _FakeCache: + def __init__(self) -> None: + self.transformer_cache = object() + + +class _FakeTransformer: + def __init__(self) -> None: + self.replaced: list[torch.Tensor] = [] + + def replace_text_embeddings(self, cache: object, embeddings: torch.Tensor) -> None: + del cache + self.replaced.append(embeddings) + + +class _FakeDiffusionModel: + def __init__(self, *, supports_text_swap: bool) -> None: + self.transformer = _FakeTransformer() if supports_text_swap else object() + + +def _session( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + pipeline: _FakePipeline, + *, + total_blocks: int = 2, + total_camera_frames: int | None = None, +): + monkeypatch.setattr( + runtime_module, + "load_first_frame_tensor", + lambda *args, **kwargs: torch.zeros(1, 3, 2, 2), + ) + runtime = LingbotReplayRuntime( + config=InferenceConfig(model_id=LINGBOT_MODEL_ID, device="cpu"), + options=LingbotReplayRuntimeOptions( + pipeline_config=object(), + pipeline_factory=lambda pipeline_config, device: pipeline, + ), + ) + image = tmp_path / "image.jpg" + image.write_bytes(b"fake") + session = runtime_module.LingbotReplaySession( + pipeline=pipeline, + session_inputs=LingbotSessionInputs( + prompt="a calm street", + first_frame_path=image, + total_blocks=total_blocks, + pixel_height=2, + pixel_width=2, + fps=16, + world_scale=2.5, + total_camera_frames=total_camera_frames, + ), + device=torch.device("cpu"), + is_rank_zero=True, + output_layout="tchw", + ) + return runtime, session + + +def _step_payload(value: float) -> InferenceInput: + poses = torch.eye(4).repeat(2, 1, 1) + poses[:, 2, 3] = value + return InferenceInput( + step={ + FIELD_CAMERA_TRAJECTORY: poses, + FIELD_CAMERA_INTRINSICS: torch.full((2, 4), value), + } + ) + + +def test_step_forwards_its_camera_inputs_to_the_model( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline() + runtime, session = _session(tmp_path, monkeypatch, pipeline) + + session.step(_step_payload(1.0)) + session.step(_step_payload(2.0)) + + assert len(pipeline.generate_calls) == 2 + # Distinct per-step payloads must reach the model distinctly; a session that + # ignored its inputs would send the same slice twice. + assert pipeline.generate_calls[0]["poses"][0, 2, 3] == 1.0 + assert pipeline.generate_calls[1]["poses"][0, 2, 3] == 2.0 + assert pipeline.generate_calls[0]["intrinsics"][0, 0] == 1.0 + assert pipeline.generate_calls[1]["intrinsics"][0, 0] == 2.0 + assert pipeline.generate_calls[0]["world_scale"] == 2.5 + runtime.close() + + +def test_step_rejects_missing_camera_inputs( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline() + runtime, session = _session(tmp_path, monkeypatch, pipeline) + + with pytest.raises(ValueError, match="missing 'camera_trajectory'"): + session.step(InferenceInput()) + + assert pipeline.generate_calls == [] + runtime.close() + + +def test_step_rejects_wrongly_shaped_camera_inputs( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline() + runtime, session = _session(tmp_path, monkeypatch, pipeline) + + with pytest.raises(ValueError, match=r"must have shape \(2, 4, 4\)"): + session.step( + InferenceInput( + step={ + FIELD_CAMERA_TRAJECTORY: torch.eye(4).repeat(5, 1, 1), + FIELD_CAMERA_INTRINSICS: torch.zeros(5, 4), + } + ) + ) + + runtime.close() + + +def test_step_request_publishes_the_input_window( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline() + runtime, session = _session(tmp_path, monkeypatch, pipeline) + + first = session.next_step_request() + assert first is not None + assert first.user_input_window is not None + assert first.user_input_window.start_s == 0.0 + assert first.user_input_window.end_s == 2 / 16 + assert first.metadata == {"num_frames": 2, "frame_start": 0} + + session.step(_step_payload(1.0)) + second = session.next_step_request() + assert second is not None + assert second.user_input_window is not None + # Windows must advance with the rollout so each step maps its own events. + assert second.user_input_window.start_s == 2 / 16 + assert second.user_input_window.end_s == 4 / 16 + runtime.close() + + +def test_rollout_ends_when_the_camera_source_runs_out( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline() + runtime, session = _session( + tmp_path, monkeypatch, pipeline, total_blocks=10, total_camera_frames=3 + ) + + assert session.next_step_request() is not None + session.step(_step_payload(1.0)) + # Only 3 frames are available and each step needs 2, so the second step + # would overrun the source. + assert session.next_step_request() is None + runtime.close() + + +def test_unbounded_source_runs_until_total_blocks( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline() + runtime, session = _session( + tmp_path, monkeypatch, pipeline, total_blocks=2, total_camera_frames=None + ) + + steps = 0 + while session.next_step_request() is not None: + session.step(_step_payload(float(steps))) + steps += 1 + + assert steps == 2 + runtime.close() + + +def test_text_event_prompt_update_swaps_the_rollout_context( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline() + runtime, session = _session(tmp_path, monkeypatch, pipeline) + + payload = _step_payload(1.0) + session.step( + InferenceInput( + global_conditioning={"prompt": "a violent storm"}, + step=payload.step, + ) + ) + + assert pipeline.text_encoder_calls == [["a violent storm"]] + assert len(pipeline.diffusion_model.transformer.replaced) == 1 + + # Re-sending the same prompt must not re-encode or re-swap. + session.step( + InferenceInput( + global_conditioning={"prompt": "a violent storm"}, + step=payload.step, + ) + ) + assert pipeline.text_encoder_calls == [["a violent storm"]] + runtime.close() + + +def test_text_event_on_an_unsupported_pipeline_fails_clearly( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline(supports_text_swap=False) + runtime, session = _session(tmp_path, monkeypatch, pipeline) + + with pytest.raises(RuntimeError, match="replace_text_embeddings"): + session.step( + InferenceInput( + global_conditioning={"prompt": "a violent storm"}, + step=_step_payload(1.0).step, + ) + ) + + # A rollout with no text event must not need the capability at all. + pipeline.generate_calls.clear() + session.step(_step_payload(1.0)) + assert len(pipeline.generate_calls) == 1 + runtime.close() + + +def test_full_standard_loop_drives_the_session_through_the_mapping( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The real runner must accept the mapping and feed the session per step. + + A fake runner cannot catch a mapping the compatibility check rejects, so + this exercises flashdreams.runtime.run_inference_session itself. + """ + import numpy as np + from lingbot.runtime import ( + LingbotModelAdapter, + LingbotReplayInputs, + inference_input_from_replay_inputs, + ) + + from flashdreams.runtime import InputCanonicalizer, UserInputs, UserInputSchema + from flashdreams.runtime.metrics import NullMetricsRecorder + from flashdreams.runtime.output import OutputArtifact + from flashdreams.runtime.runner import run_inference_session + + image = tmp_path / "image.jpg" + image.write_bytes(b"fake") + poses_path = tmp_path / "poses.npy" + intrinsics_path = tmp_path / "intrinsics.npy" + trajectory = np.tile(np.eye(4, dtype=np.float32), (32, 1, 1)) + trajectory[:, 2, 3] = np.arange(32, dtype=np.float32) + np.save(poses_path, trajectory) + np.save( + intrinsics_path, + np.tile(np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (32, 1)), + ) + + pipeline = _FakePipeline() + monkeypatch.setattr( + runtime_module, + "load_first_frame_tensor", + lambda *args, **kwargs: torch.zeros(1, 3, 2, 2), + ) + adapter = LingbotModelAdapter( + pipeline_factory=lambda pipeline_config, device: pipeline, + ) + replay_inputs = LingbotReplayInputs( + prompt="a calm street", + first_frame_path=image, + camera_poses_path=poses_path, + camera_intrinsics_path=intrinsics_path, + total_blocks=3, + pixel_height=2, + pixel_width=2, + fps=16, + ) + + class _Collecting: + def __init__(self) -> None: + self.results: list[Any] = [] + + def open(self) -> None: + return None + + def write(self, result: Any) -> None: + self.results.append(result) + + def close(self) -> tuple[OutputArtifact, ...]: + return () + + output = _Collecting() + mapping = adapter.create_input_mapping(replay_inputs) + run_inference_session( + adapter=adapter, + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + device="cpu", + runtime_options={"pipeline_config": object()}, + ), + mapping=mapping, + canonicalizer=InputCanonicalizer(), + source_schema=UserInputSchema(), + user_inputs=UserInputs(), + initial_inputs=inference_input_from_replay_inputs(replay_inputs), + output=output, + metrics=NullMetricsRecorder(), + ) + + assert len(output.results) == 3 + assert len(pipeline.generate_calls) == 3 + # Each step must receive its own successive slice of the trace. Comparing + # against the trace directly is the real property; consecutive pose values + # can repeat because preprocess_example_poses re-expands encoded poses at + # stride-4 cadence. + trace_poses = mapping.camera_trace.poses + received = torch.cat([call["poses"] for call in pipeline.generate_calls]) + assert torch.allclose(received, trace_poses[:6]) diff --git a/integrations/lingbot/tests/test_smoke.py b/integrations/lingbot/tests/test_smoke.py index bf9c1ab1..3cd0a93a 100644 --- a/integrations/lingbot/tests/test_smoke.py +++ b/integrations/lingbot/tests/test_smoke.py @@ -21,9 +21,11 @@ from pathlib import Path from typing import cast +import numpy as np import pytest import tomli as tomllib from lingbot import config as config_mod +from lingbot import example_data as example_data_mod from lingbot import runner as runner_mod from lingbot.config import ( LINGBOT_WORLD_V2_CHECKPOINT_PATH, @@ -39,6 +41,14 @@ LingbotWorldRunnerConfig, example_data_dirname, ) +from lingbot.runtime import ( + FIELD_FIRST_FRAME_PATH, + FIELD_PROMPT, + FIELD_TOTAL_BLOCKS, + LINGBOT_MODEL_ID, + LingbotModelAdapter, + LingbotRunnerOutputTarget, +) from lingbot.transformer import ( LINGBOT_WORLD_MIN_CHECKPOINT_FREE_GB, LingbotWorldTransformer, @@ -50,6 +60,19 @@ pytestmark = pytest.mark.ci_cpu + +def _write_camera_assets(poses: Path, intrinsics: Path, *, frames: int = 64) -> None: + """Write real .npy camera assets; the input mapping loads them for real.""" + trajectory = np.tile(np.eye(4, dtype=np.float32), (frames, 1, 1)) + trajectory[:, 2, 3] = np.linspace(0.0, 1.0, frames, dtype=np.float32) + np.save(poses, trajectory) + np.save( + intrinsics, + np.tile( + np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (frames, 1) + ), + ) + ENTRY_POINT_GROUP = "flashdreams.runner_configs" @@ -77,10 +100,10 @@ def _record_download(url: str, *, cache_dir: Path, filename: str) -> None: del cache_dir, filename urls.append(url) - monkeypatch.setattr(runner_mod, "EXAMPLE_DATA_DIR_LOCAL", tmp_path) - monkeypatch.setattr(runner_mod, "download_to_cache", _record_download) + monkeypatch.setattr(example_data_mod, "EXAMPLE_DATA_DIR_LOCAL", tmp_path) + monkeypatch.setattr(example_data_mod, "download_to_cache", _record_download) - runner_mod.ensure_example_data_downloaded(is_rank_zero=True, example_idx=0) + example_data_mod.ensure_example_data_downloaded(is_rank_zero=True, example_idx=0) expected_base_url = ( "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00" @@ -105,10 +128,10 @@ def test_promptless_examples_skip_the_prompt_download( def _record_download(url: str, *, cache_dir: Path, filename: str) -> None: downloads.append((url, cache_dir, filename)) - monkeypatch.setattr(runner_mod, "EXAMPLE_DATA_DIR_LOCAL", tmp_path) - monkeypatch.setattr(runner_mod, "download_to_cache", _record_download) + monkeypatch.setattr(example_data_mod, "EXAMPLE_DATA_DIR_LOCAL", tmp_path) + monkeypatch.setattr(example_data_mod, "download_to_cache", _record_download) - cache_dir = runner_mod.ensure_example_data_downloaded( + cache_dir = example_data_mod.ensure_example_data_downloaded( is_rank_zero=True, example_idx=example_idx, ) @@ -162,6 +185,69 @@ def test_promptless_example_resolves_to_empty_string( ] +def test_runner_delegates_to_runtime_api_with_direct_inputs( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Keep the CLI runner on the new runtime path, not the old rollout loop.""" + image = tmp_path / "image.jpg" + poses = tmp_path / "poses.npy" + intrinsics = tmp_path / "intrinsics.npy" + image.write_bytes(b"fake") + _write_camera_assets(poses, intrinsics) + runner = object.__new__(LingbotWorldRunner) + runner_config = cast( + LingbotWorldRunnerConfig, + derive_config( + RUNNER_CONFIGS["lingbot-world-fast-taehv-window15-sink3"], + prompt="drive through a city", + image_path=image, + pose_path=poses, + intrinsic_path=intrinsics, + total_blocks=1, + device="cpu", + ), + ) + pipeline = object() + output_stream = object() + captured: dict[str, object] = {} + + def _fake_run_inference_session(**kwargs: object) -> tuple[object, ...]: + captured.update(kwargs) + return () + + monkeypatch.setattr( + runner, + "create_video_output_stream", + lambda **_kwargs: output_stream, + ) + monkeypatch.setattr( + runner_mod, + "run_inference_session", + _fake_run_inference_session, + ) + runner.config = runner_config + runner.pipeline = pipeline + runner.local_rank = 0 + runner.world_size = 1 + runner.is_rank_zero = True + + runner.run() + + assert isinstance(captured["adapter"], LingbotModelAdapter) + config = captured["config"] + assert getattr(config, "model_id") == LINGBOT_MODEL_ID + assert getattr(config, "device") == "cpu" + assert config.runtime_options["pipeline"] is pipeline + inputs = captured["initial_inputs"].global_conditioning + assert inputs[FIELD_PROMPT] == "drive through a city" + assert inputs[FIELD_FIRST_FRAME_PATH] == image + assert inputs[FIELD_TOTAL_BLOCKS] == 1 + output = captured["output"] + assert isinstance(output, LingbotRunnerOutputTarget) + assert output.output_stream is output_stream + + def test_runners_dict_is_non_empty() -> None: """Plugin must expose at least one runner.""" assert RUNNER_CONFIGS, "RUNNER_CONFIGS is empty"