diff --git a/.github/workflows/omnidreams-demo-runtime.yml b/.github/workflows/omnidreams-demo-runtime.yml index 1f494fa83..361ad44a5 100644 --- a/.github/workflows/omnidreams-demo-runtime.yml +++ b/.github/workflows/omnidreams-demo-runtime.yml @@ -103,15 +103,62 @@ jobs: log_dir="${ARTIFACT_DIR}/logs" output_dir="${ARTIFACT_DIR}/outputs" + launch_dir="${ARTIFACT_DIR}/launch_manifests" summary="${ARTIFACT_DIR}/summary.md" status_file="${ARTIFACT_DIR}/command-status.env" - mkdir -p "${log_dir}" "${output_dir}" + mkdir -p "${log_dir}" "${output_dir}" "${launch_dir}" : > "${status_file}" - odemo() { - uv run --no-sync --package flashdreams-omnidreams omnidreams-demo "$@" + frun() { + uv run --no-sync --package flashdreams-omnidreams flashdreams-run "$@" } + runner="omnidreams" + null_manifest="${launch_dir}/null.yaml" + precomputed_manifest="${launch_dir}/precomputed-mp4.yaml" + ludus_manifest="${launch_dir}/ludus-mp4.yaml" + + cat > "${null_manifest}" < "${precomputed_manifest}" < "${ludus_manifest}" < "${summary}" run_demo null \ - odemo replay \ - --output-mode null \ - --device cuda:0 \ - --total-blocks "${NULL_BLOCKS}" + frun "${runner}" null --manifest "${null_manifest}" run_demo precomputed-mp4 \ - odemo replay \ - --device cuda:0 \ - --example-data \ - --example-data-uuid "${EXAMPLE_DATA_UUID}" \ - --total-blocks "${PRECOMPUTED_BLOCKS}" \ - --fps "${FPS}" \ - --output "${output_dir}/omnidreams-demo-precomputed-20s.mp4" + frun "${runner}" mp4 --manifest "${precomputed_manifest}" run_demo ludus-mp4 \ - odemo replay \ - --conditioning-mode ludus-scene-driving \ - --keyboard-trace "${LUDUS_TRACE}" \ - --device cuda:0 \ - --scene-uuid "${LUDUS_SCENE_UUID}" \ - --seed 42 \ - --total-blocks "${LUDUS_BLOCKS}" \ - --output "${output_dir}/omnidreams-demo-ludus-20s.mp4" + frun "${runner}" mp4 --manifest "${ludus_manifest}" - name: Validate OmniDreams demo artifacts run: | diff --git a/README.md b/README.md index fa18e70de..9ccf39758 100644 --- a/README.md +++ b/README.md @@ -102,7 +102,8 @@ pip install flashdreams ### Try the interactive driving demo -Drive a world model in real time with the OmniDreams `interactive-drive` demo. See the +Drive a world model in real time with the unified OmniDreams `local-window` or +`webrtc` launch mode. See the **[interactive demo guide](https://nvidia.github.io/flashdreams/main/models/omnidreams.html#launch-the-interactive-demo)**. ## Supported models diff --git a/configs/deterministic_quality_benchmarks.json b/configs/deterministic_quality_benchmarks.json index 29b732ff8..e2b4e7b1a 100644 --- a/configs/deterministic_quality_benchmarks.json +++ b/configs/deterministic_quality_benchmarks.json @@ -83,7 +83,7 @@ "-m", "tools.benchmarks.strict_run", "--", - "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae", + "omnidreams", "--example-data", "True", "--example-data-uuid", @@ -168,7 +168,7 @@ "-m", "tools.benchmarks.strict_run", "--", - "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae", + "omnidreams", "--example-data", "True", "--example-data-uuid", diff --git a/configs/launch_manifest/lingbot_mp4.yaml b/configs/launch_manifest/lingbot_mp4.yaml new file mode 100644 index 000000000..8d7c2aadb --- /dev/null +++ b/configs/launch_manifest/lingbot_mp4.yaml @@ -0,0 +1,14 @@ +schema_version: 1 +runner: lingbot-world-fast +mode: mp4 + +scenario: + example_data: true + example_idx: 0 + total_blocks: 20 + pixel_width: 832 + pixel_height: 464 + +output: + path: ../../outputs/lingbot-replay.mp4 + fps: 16 diff --git a/configs/launch_manifest/lingbot_webrtc.yaml b/configs/launch_manifest/lingbot_webrtc.yaml new file mode 100644 index 000000000..3e42b4b23 --- /dev/null +++ b/configs/launch_manifest/lingbot_webrtc.yaml @@ -0,0 +1,16 @@ +schema_version: 1 +runner: lingbot-world-fast +mode: webrtc + +scenario: + example_idx: 0 + +output: + host: 0.0.0.0 + port: 8080 + fps: 16 + video_width: 832 + video_height: 464 + warmup_chunks: 10 + warmup_timeout_s: 600 + client_liveness_timeout_s: 30 diff --git a/configs/launch_manifest/omnidreams_local_window.yaml b/configs/launch_manifest/omnidreams_local_window.yaml new file mode 100644 index 000000000..1d712ef96 --- /dev/null +++ b/configs/launch_manifest/omnidreams_local_window.yaml @@ -0,0 +1,10 @@ +schema_version: 1 +runner: omnidreams-perf +mode: local-window + +scenario: + auto_start: false + preload_scenes: false + +output: + world_model_manifest_path: ../../integrations/omnidreams/omnidreams/interactive_drive/configs/example_world_model_perf.yaml diff --git a/configs/launch_manifest/omnidreams_mp4.yaml b/configs/launch_manifest/omnidreams_mp4.yaml new file mode 100644 index 000000000..195e70608 --- /dev/null +++ b/configs/launch_manifest/omnidreams_mp4.yaml @@ -0,0 +1,13 @@ +schema_version: 1 +runner: omnidreams +mode: mp4 + +scenario: + example_data: true + total_blocks: 60 + pixel_width: 1280 + pixel_height: 704 + +output: + path: ../../outputs/omnidreams-replay.mp4 + fps: 30 diff --git a/configs/launch_manifest/omnidreams_null.yaml b/configs/launch_manifest/omnidreams_null.yaml new file mode 100644 index 000000000..334adf5c5 --- /dev/null +++ b/configs/launch_manifest/omnidreams_null.yaml @@ -0,0 +1,10 @@ +schema_version: 1 +runner: omnidreams +mode: "null" + +scenario: + example_data: true + total_blocks: 10 + pixel_width: 1280 + pixel_height: 704 + fps: 30 diff --git a/configs/launch_manifest/omnidreams_webrtc.yaml b/configs/launch_manifest/omnidreams_webrtc.yaml new file mode 100644 index 000000000..2a2ec9a72 --- /dev/null +++ b/configs/launch_manifest/omnidreams_webrtc.yaml @@ -0,0 +1,21 @@ +schema_version: 1 +runner: omnidreams +mode: webrtc + +runner_overrides: + device: cuda:0 + +scenario: + scene_uuid: 0d404ff7-2b66-498c-b047-1ed8cded60d4 + scene_variant: default + camera_name: camera_front_wide_120fov + +output: + host: 0.0.0.0 + port: 8089 + fps: 30 + video_width: 1280 + video_height: 704 + warmup_chunks: 10 + warmup_timeout_s: 600 + client_liveness_timeout_s: 10 diff --git a/configs/omnidreams_demo_replay_benchmarks.json b/configs/omnidreams_demo_replay_benchmarks.json index 3ac6e4019..09be6f95e 100644 --- a/configs/omnidreams_demo_replay_benchmarks.json +++ b/configs/omnidreams_demo_replay_benchmarks.json @@ -29,7 +29,7 @@ "--project", "integrations/omnidreams", "flashdreams-run", - "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae", + "omnidreams", "--example-data", "True", "--example-data-uuid", diff --git a/configs/one_minute_demo_benchmarks.json b/configs/one_minute_demo_benchmarks.json index 4dcb6f9f6..85112b1c5 100644 --- a/configs/one_minute_demo_benchmarks.json +++ b/configs/one_minute_demo_benchmarks.json @@ -68,7 +68,7 @@ "-m", "tools.benchmarks.strict_run", "--", - "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae", + "omnidreams", "--example-data", "True", "--example-data-uuid", diff --git a/docs/source/api/cli.rst b/docs/source/api/cli.rst index 2e99ab447..ff8594167 100644 --- a/docs/source/api/cli.rst +++ b/docs/source/api/cli.rst @@ -34,12 +34,24 @@ Inspect one runner's full options: uv run flashdreams-run self-forcing-wan2.1-t2v-1.3b-taehv --help -Run a single-GPU inference: +Run a single-GPU inference (``run`` is the default mode): .. code-block:: bash uv run flashdreams-run self-forcing-wan2.1-t2v-1.3b-taehv --total-blocks 7 +Launch a WebRTC demo from a versioned manifest: + +.. code-block:: bash + + uv run flashdreams-run lingbot-world-fast webrtc \ + --manifest configs/launch_manifest/lingbot_webrtc.yaml + +The common command shape is ``flashdreams-run [mode]``. A runner only +advertises modes it implements; unsupported pairs fail before CUDA +initialization. Shared modes are ``run``, ``mp4``, ``null``, ``webrtc``, and +``local-window``. + Run a multi-GPU inference: .. code-block:: bash @@ -51,7 +63,7 @@ Resolve config only (no model instantiation): .. code-block:: bash - uv run flashdreams-run --no-instantiate self-forcing-wan2.1-t2v-1.3b-taehv + uv run flashdreams-run self-forcing-wan2.1-t2v-1.3b-taehv --no-instantiate Post-processing presets ----------------------- @@ -73,5 +85,7 @@ See also -------- - :doc:`/quickstart/index` +- :doc:`/api/launch_manifests` - :doc:`/developer_guides/config_system` +- :doc:`/developer_guides/runner_slugs` - :doc:`/api/infra` diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst index c0694b322..f68d1df65 100644 --- a/docs/source/api/index.rst +++ b/docs/source/api/index.rst @@ -65,6 +65,7 @@ Reference pages for the FlashDreams command-line interface and Python APIs. :maxdepth: 1 cli + launch_manifests core infra integrations diff --git a/docs/source/api/launch_manifests.rst b/docs/source/api/launch_manifests.rst new file mode 100644 index 000000000..a7f4e8b65 --- /dev/null +++ b/docs/source/api/launch_manifests.rst @@ -0,0 +1,106 @@ +.. SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +.. SPDX-License-Identifier: Apache-2.0 + +Launch manifests +================ + +Every FlashDreams model launch starts with the same command shape: + +.. code-block:: bash + + uv run flashdreams-run [mode] [--manifest PATH] + +``run`` is the default mode. Integrations may additionally expose ``mp4``, +``null``, ``webrtc``, or ``local-window``. Inspect a resolved launch without +loading checkpoints or initializing CUDA with ``--no-instantiate``. + +Schema +------ + +Launch manifests are strict, versioned YAML documents: + +.. code-block:: yaml + + schema_version: 1 + runner: omnidreams + mode: webrtc + + runner_overrides: + device: cuda:0 + + scenario: + scene_uuid: 0d404ff7-2b66-498c-b047-1ed8cded60d4 + scene_variant: default + + output: + host: 0.0.0.0 + port: 8089 + +``schema_version``, ``runner``, and ``mode`` are required. The remaining +sections are optional mappings: + +.. note:: + + Quote the null-output mode as ``mode: "null"`` in YAML; an unquoted + ``null`` is YAML's null scalar rather than the FlashDreams mode name. + +``runner_overrides`` + Recursive overrides for the registered runner configuration. The same + runner fields remain available as explicit CLI flags. + +``scenario`` + Inputs and controls such as prompts, example data, scenes, traces, and + rollout length. The selected integration validates the accepted fields. + +``output`` + Transport or artifact settings such as output path, frame rate, WebRTC + bind address, warmup, and local-window presentation settings. + +Relative paths in keys named ``path`` or ending in ``_path``, ``_paths``, or +``_dir`` resolve relative to the manifest file, which makes checked-in launch +manifests reproducible from any working directory. Unknown top-level or +integration-specific fields fail before CUDA initialization. + +Precedence +---------- + +Settings resolve in this order, from lowest to highest precedence: + +.. code-block:: text + + registered runner preset + < manifest runner_overrides + < manifest scenario/output + < explicit CLI runner flags and --host/--port + +The runner and positional mode must agree with the manifest. For example, this +fails instead of silently launching a different preset: + +.. code-block:: bash + + uv run flashdreams-run lingbot-world-fast mp4 \ + --manifest configs/launch_manifest/lingbot_webrtc.yaml + +Examples +-------- + +.. code-block:: bash + + # WebRTC + uv run flashdreams-run lingbot-world-fast webrtc \ + --manifest configs/launch_manifest/lingbot_webrtc.yaml + + # MP4 replay + uv run flashdreams-run lingbot-world-fast mp4 \ + --manifest configs/launch_manifest/lingbot_mp4.yaml + + # Resolve an OmniDreams launch without loading the model + uv run flashdreams-run \ + omnidreams webrtc \ + --manifest configs/launch_manifest/omnidreams_webrtc.yaml \ + --no-instantiate + +OmniDreams local-window also accepts the existing +``example_world_model*.yaml`` format directly as a compatibility input. New +automation should use a versioned launch manifest whose +``output.world_model_manifest_path`` references that model-specific file. diff --git a/docs/source/developer_guides/index.rst b/docs/source/developer_guides/index.rst index 7cc2bac85..8a8874d8b 100644 --- a/docs/source/developer_guides/index.rst +++ b/docs/source/developer_guides/index.rst @@ -34,6 +34,13 @@ Developer Guides How every overridable field is surfaced as a CLI flag, how method defaults compose, and how to layer overrides on top. + .. grid-item-card:: Runner slugs and demo dispatch + :link: runner_slugs + :link-type: doc + + How public runner names are registered, parsed, matched to manifests, + and dispatched to integration-owned demo launch modes. + .. grid-item-card:: Add a new method :link: new_integration :link-type: doc @@ -62,6 +69,7 @@ generated clip, see :doc:`/quickstart/index`. inference_pipeline_overview config_system + runner_slugs new_integration local_benchmarks diff --git a/docs/source/developer_guides/latency_tuning.rst b/docs/source/developer_guides/latency_tuning.rst index 138762621..e038110b5 100644 --- a/docs/source/developer_guides/latency_tuning.rst +++ b/docs/source/developer_guides/latency_tuning.rst @@ -13,10 +13,10 @@ .. See the License for the specific language governing permissions and .. limitations under the License. -Interactive-drive latency tuning -================================ +OmniDreams interactive latency tuning +===================================== -``interactive-drive`` latency has two different components: +Interactive launch latency has two different components: - **Model / chunk latency** is the time spent preparing HDMap conditioning, running the OmniDreams DiT, decoding the generated chunk, and updating model @@ -50,8 +50,9 @@ Run the perf manifest only on hosts that can build and load the native extension .. code-block:: bash uv run --package flashdreams-omnidreams omnidreams-prepare --perf - uv run --package flashdreams-omnidreams interactive-drive \ - --manifest example_world_model_perf.yaml + uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams-perf local-window \ + --manifest configs/launch_manifest/omnidreams_local_window.yaml ``native_dit_acceleration: required`` is intentional. If the native extension is not available, startup fails instead of silently falling back to the slower @@ -142,10 +143,11 @@ Pick transport based on where the viewer runs: - Local Vulkan window: lowest-overhead local presentation when the host has a graphics-capable GPU and display stack. -- ``--stream-mjpeg [HOST:]PORT``: simple browser delivery from the same process. +- ``output.stream_mjpeg`` in a local-window manifest: simple browser delivery + from the same process. Use it on compute-only hosts such as GB300 systems without a graphics queue, or when a laptop browser views a remote model host. -- ``omnidreams.webrtc.server``: richer browser frontend with WebRTC's lower +- ``webrtc`` launch mode: richer browser frontend with WebRTC's lower video-delivery latency and streaming gRPC service support. Prefer this for product-style remote viewing or multi-client integration. @@ -156,8 +158,9 @@ resolution, and native-acceleration knobs first. Profiling and validated reference --------------------------------- -Use ``--profile-world-model`` to enable FlashDreams CUDA-event profiling for the -world-model runtime. Use ``--sync-gpu-timing`` only when you need raster compute +Set ``output.profile_world_model: true`` to enable FlashDreams CUDA-event +profiling for the world-model runtime. Set ``output.sync_gpu_timing: true`` only +when you need raster compute timings; it synchronizes GPU work and is not a throughput setting. The validated published reference for interactive-drive latency is the diff --git a/docs/source/developer_guides/runner_slugs.md b/docs/source/developer_guides/runner_slugs.md new file mode 100644 index 000000000..55f3a53aa --- /dev/null +++ b/docs/source/developer_guides/runner_slugs.md @@ -0,0 +1,148 @@ + + +# Runner slugs and demo launch dispatch + +A runner slug is the short public name after `flashdreams-run`. It selects a +registered `RunnerConfig`; it does not need to duplicate the detailed name of +the pipeline preset inside that config. + +## Quick start + +After installing the OmniDreams workspace package, the default MP4 demo is: + +```bash +uv sync --package flashdreams-omnidreams +uv run flashdreams-run omnidreams mp4 +``` + +This uses the bundled single-view example data and writes +`outputs/omnidreams.mp4`. Use a launch manifest when you need to change the +scenario or output: + +```bash +uv run flashdreams-run omnidreams mp4 \ + --manifest configs/launch_manifest/omnidreams_mp4.yaml +``` + +The shipped public OmniDreams runners are: + +| Runner slug | Registered config literal | Internal pipeline preset | +| --- | --- | --- | +| `omnidreams` | `RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE` | `omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae` | +| `omnidreams-perf` | `RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF` | `omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf` | + +The short slug is stable user interface. The internal pipeline name remains +specific enough for recipe selection, checkpoints, profiling, and direct demo +APIs. + +## How the command is dispatched + +For this command: + +```bash +uv run flashdreams-run omnidreams mp4 +``` + +the control flow is: + +```text +flashdreams-run console script + -> flashdreams.scripts.cli:entrypoint + -> _prepare_cli_args(["omnidreams", "mp4"]) + -> all_runners() + -> built-in runner registry + -> discover_runners() + -> importlib.metadata entry points + -> group: flashdreams.runner_configs + -> registry lookup by RunnerConfig.runner_name == "omnidreams" + -> remove the positional mode "mp4" before Tyro parsing + -> Tyro resolves runner flags over the registered RunnerConfig + -> main(..., mode="mp4") + -> resolve_launch() + -> omnidreams.launch:LAUNCH_CAPABILITY + -> OmnidreamsLaunchCapability.resolve() + -> ResolvedLaunch.launch() + -> omnidreams.demo.app.launch_from_runner() +``` + +Important details: + +1. `_prepare_cli_args` finds the runner token by checking the keys returned by + `all_runners()`. The next positional token is interpreted as a mode only + when it is one of `run`, `mp4`, `null`, `webrtc`, or + `local-window`. +2. A launch manifest must name the same runner and mode as the command. Its + `runner_overrides` are applied before Tyro parses explicit CLI overrides. +3. `run` calls `config.setup()` and the regular runner. Other modes are + delegated through the config's `launch_capability`. +4. The OmniDreams capability validates integration-specific scenario and + output fields, then calls the shared demo API directly. It does not invoke a + second CLI. +5. The demo derives its `preset_id` from `config.pipeline.name`. Therefore + `omnidreams` still selects the detailed stable non-performance pipeline + preset shown in the table. + +## How a slug is registered + +External integrations register runners with Python package entry points. The +OmniDreams package declares: + +```toml +[project.entry-points."flashdreams.runner_configs"] +"omnidreams" = "omnidreams.config:RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE" +"omnidreams-perf" = "omnidreams.config:RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF" +``` + +Each target resolves to an `OmnidreamsRunnerConfig` literal: + +```python +RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE = OmnidreamsRunnerConfig( + runner_name="omnidreams", + pipeline=SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE, + # ... +) +``` + +Discovery loads every entry point in the `flashdreams.runner_configs` group. +The registry key comes from `cfg.runner_name`, not from the entry-point name. +The entry-point name should still match `runner_name`: keeping both aligned +makes installed package metadata understandable, and the OmniDreams CPU tests +enforce that invariant. + +Slug collisions are deterministic. Built-in runners take precedence over +plugins, and the first discovered plugin with a given `runner_name` wins over +later plugins. The registry logs skipped collisions. + +## Adding or changing a public slug + +1. Choose a short, integration-level default such as `omnidreams`. Add a + concise suffix only when users need to select a materially different public + variant, such as `omnidreams-perf`. +2. Set `runner_name` on the exported runner config literal. +3. Add an entry with the same name under + `[project.entry-points."flashdreams.runner_configs"]`. +4. Keep `pipeline.name` unchanged unless the actual model recipe identity is + changing. +5. Update launch manifests, CI commands, benchmark scenarios, and user-facing + documentation. Do not rename checkpoint keys, asset paths, reference + artifacts, or internal preset IDs merely because the public slug changed. +6. Refresh the editable package metadata and run CPU-only checks: + +```bash +uv sync --package flashdreams-omnidreams --package flashdreams-lingbot \ + --group test +uv run flashdreams-run --help +uv run flashdreams-run omnidreams mp4 --no-instantiate +uv run pytest -m ci_cpu \ + integrations/omnidreams/tests/test_recipe_configs.py \ + integrations/omnidreams/tests/test_demo_api.py \ + flashdreams/tests/test_launch.py \ + flashdreams/tests/test_launch_manifest.py +``` + +The `--no-instantiate` check resolves registration, parsing, mode dispatch, +manifest validation, and default launch settings without loading checkpoints or +initializing the GPU. diff --git a/docs/source/models/lingbot_world.rst b/docs/source/models/lingbot_world.rst index 34e7fe7a0..c83097294 100644 --- a/docs/source/models/lingbot_world.rst +++ b/docs/source/models/lingbot_world.rst @@ -109,7 +109,7 @@ We provide the following variants: + ``sink_size_t=3`` streaming KV-cache. * - ``lingbot-world-v2-14b-causal-fast`` - LingBot-World V2 14B causal-fast on the shared LingBot pipeline - (Wan VAE decoder, 4-step). See `LingBot-World V2`_. + (Wan VAE decoder, 4-step). See :ref:`lingbot-world-v2`. * - ``lingbot-world-v2-14b-causal-fast-taehv-window15-sink3`` - LingBot-World V2 14B causal-fast with the TAEHV decoder, ``window_size_t=15`` + ``sink_size_t=3`` streaming KV-cache. @@ -123,6 +123,8 @@ To inspect all supported CLI arguments and their default values, run: lingbot-world-fast \ --help +.. _lingbot-world-v2: + LingBot-World V2 ---------------- @@ -227,14 +229,13 @@ Spin up the interactive LingBot-World server via WebRTC: .. code-block:: bash # from the repo root - uv run --package flashdreams-lingbot torchrun --nproc_per_node 4 \ - -m lingbot.webrtc.server \ - --host 0.0.0.0 --port 8089 \ - --config_name lingbot-world-fast-taehv-window15-sink3 \ - --example-idx 0 - -``--example-idx`` selects which example to download -(``0``, ``1``, ``2``, ``5``); assets auto-download on first launch. + uv run --package flashdreams-lingbot \ + torchrun --nproc_per_node 4 --no-python flashdreams-run \ + lingbot-world-fast-taehv-window15-sink3 webrtc \ + --host 0.0.0.0 --port 8089 + +``scenario.example_idx`` in a launch manifest selects which example to +download (``0``, ``1``, ``2``, ``5``); assets auto-download on first launch. The HTTP port opens only after model load + warmup — a few minutes on first launch, much faster afterwards. When ready the server prints ``Connect via http://:8089/request_session`` (use diff --git a/docs/source/models/omnidreams.rst b/docs/source/models/omnidreams.rst index 60db7e64f..1d888857f 100644 --- a/docs/source/models/omnidreams.rst +++ b/docs/source/models/omnidreams.rst @@ -69,6 +69,15 @@ Installation # from the repo root uv sync --project integrations/omnidreams +Generate the default MP4 demo from bundled example data: + +.. code-block:: bash + + uv run flashdreams-run omnidreams mp4 + +The command writes ``outputs/omnidreams.mp4``. Use a launch manifest to +override the scenario, rollout length, frame rate, or output path. + Running the method ------------------ @@ -79,7 +88,7 @@ example: uv run --project integrations/omnidreams \ flashdreams-run \ - omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf \ + omnidreams \ --example-data True \ --example_data_uuid "239560dc-33d1-11ef-9720-00044bcbccac" \ --total-blocks 20 @@ -95,8 +104,10 @@ We provide the following variants: * - Method - Description - * - ``omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf`` - - Single-view 2-step HDMap-conditioned I2V. + * - ``omnidreams`` + - Default single-view 2-step HDMap-conditioned I2V demo and runner. + * - ``omnidreams-perf`` + - Opt-in compile and CUDA-graph tuning across all pipeline stages. For multi-GPU inference, use: @@ -104,7 +115,7 @@ For multi-GPU inference, use: uv run --project integrations/omnidreams \ torchrun --nproc_per_node=4 --no-python flashdreams-run \ - omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf \ + omnidreams \ --example-data True \ --example_data_uuid "239560dc-33d1-11ef-9720-00044bcbccac" \ --total-blocks 20 @@ -115,7 +126,7 @@ To inspect all supported CLI arguments and their default values, run: uv run --project integrations/omnidreams \ flashdreams-run \ - omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf \ + omnidreams \ --help Some generated samples from the above commands: @@ -148,10 +159,9 @@ Some generated samples from the above commands: Launch the interactive demo --------------------------- -``interactive-drive`` runs the OmniDreams single-view pipeline in a -single process and streams the camera view to your browser. The demo -machine only needs a CUDA-capable GPU -- no graphics-capable GPU, -display server, or Vulkan support are required. +OmniDreams exposes ``webrtc`` and ``local-window`` through the shared +``flashdreams-run `` command. WebRTC only requires a +CUDA-capable GPU; local-window additionally requires a display and Vulkan. The demo requires access to `NVIDIA/flashdreams `_ and an ``HF_TOKEN`` with read access to @@ -176,14 +186,16 @@ isn't blocked on network I/O: uv run --package flashdreams-omnidreams omnidreams-prepare -Run the demo and stream to your browser: +Run the WebRTC demo: .. code-block:: bash - uv run --package flashdreams-omnidreams interactive-drive --stream-mjpeg :8080 + uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams webrtc \ + --manifest configs/launch_manifest/omnidreams_webrtc.yaml -Then open ``http://:8080/`` in any browser on the same -network and pick a scene from the picker in the bottom-right. +Then open ``http://:8089/request_session`` in any browser on the +same network. .. note:: @@ -198,13 +210,15 @@ network and pick a scene from the picker in the bottom-right. .. note:: - Add ``--offload-text-encoder`` to reduce peak VRAM usage by ~15 GB: + For local-window, set ``output.offload_text_encoder: true`` in a copy of + ``configs/launch_manifest/omnidreams_local_window.yaml`` to reduce peak VRAM + usage by ~15 GB, then launch it with the central command: .. code-block:: bash - uv run --package flashdreams-omnidreams interactive-drive \ - --stream-mjpeg :8080 \ - --offload-text-encoder + uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams-perf local-window \ + --manifest path/to/local-window.yaml The text and first-frame encoders are run once per scene and freed before the diffusion pipeline is built, and the resulting embeddings are cached and @@ -214,13 +228,13 @@ network and pick a scene from the picker in the bottom-right. resident, so the first load and scene/variant switches are slower. Prefer it when VRAM-constrained; otherwise leave it off for faster switching. -For execution using a consumer NVIDIA GPU that exposes a graphics stack, -omit the ``--stream-mjpeg`` flag to open the demo in a local Vulkan window -instead: +On a GPU with a graphics stack, launch the Vulkan window: .. code-block:: bash - uv run --package flashdreams-omnidreams interactive-drive + uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams-perf local-window \ + --manifest configs/launch_manifest/omnidreams_local_window.yaml The local window's HUD adds a weather-variant selector (clear, rain, snow) next to the scene picker, so the same scene can be switched between @@ -252,7 +266,7 @@ or joystick is viable. We provide a configuration tool to calibrate these: The demo auto-loads your default profile on subsequent launches. When you have more than one profile, the configuration tool's start screen lists them with **Make default** (plus Edit and Delete) buttons -- re-run the tool to -choose which profile ``interactive-drive`` loads by default, tweak a profile +choose which profile ``local-window`` loads by default, tweak a profile (steering sensitivity, deadzone, buttons, force feedback), or remove one. **Multiple devices.** A profile can bind controls across several devices -- @@ -307,38 +321,33 @@ point the demo at the perf manifest: .. code-block:: bash - uv run --package flashdreams-omnidreams interactive-drive \ - --manifest example_world_model_perf.yaml + uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams-perf local-window \ + --manifest configs/launch_manifest/omnidreams_local_window.yaml ``native_dit_acceleration: required`` makes the manifest fail loudly if the extension can't build or load, rather than silently falling back to PyTorch. -Alternative: WebRTC server --------------------------- +WebRTC server +------------- For deployments that require a richer browser frontend with WebRTC's lower video-delivery latency and a streaming gRPC service for -multi-client setups, the standalone server at -``omnidreams.webrtc.server`` ships a polished HTML5 client on top of -the same OmniDreams pipeline. The MJPEG path above is the -recommended starting point for most users; consider WebRTC if you -need bidirectional camera-control APIs or are already integrating -the gRPC service into a larger product. +multi-client setups, the ``webrtc`` launch mode ships a polished HTML5 client +on top of the same OmniDreams pipeline. .. code-block:: bash # from the repo root - uv run --package flashdreams-omnidreams torchrun --nproc_per_node 1 \ - -m omnidreams.webrtc.server \ - --host 0.0.0.0 --port 8089 \ - --pipeline_config_name omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf \ - --scene-uuid "0d404ff7-2b66-498c-b047-1ed8cded60d4" + uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams webrtc \ + --manifest configs/launch_manifest/omnidreams_webrtc.yaml Sample scene UUIDs for the interactive server are available in the `nvidia/omni-dreams-scenes Hugging Face dataset `_. Each scene ships clear, rain, and snow weather variants as sibling -archives; add ``--scene-variant rain`` (or ``snow``) to serve a specific -one (the default is the clear-weather scene). +archives; set ``scenario.scene_variant`` to ``rain`` or ``snow`` in the launch +manifest to serve a specific one (the default is clear weather). The server may take a few minutes to warm up. Once ready, it prints ``Connect via http://:8089/request_session``. diff --git a/docs/source/quickstart/index.rst b/docs/source/quickstart/index.rst index 2db544fba..4ae191a2d 100644 --- a/docs/source/quickstart/index.rst +++ b/docs/source/quickstart/index.rst @@ -54,28 +54,32 @@ Run your first model -------------------- Launch the OmniDreams interactive driving demo. It runs the world model -and streams the generated camera view to a browser: +and streams the generated camera view to a browser over WebRTC: .. code-block:: bash - uv run --package flashdreams-omnidreams interactive-drive --stream-mjpeg :8080 + uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams webrtc \ + --manifest configs/launch_manifest/omnidreams_webrtc.yaml -Then open ``http://:8080/`` in a browser on the same network +Then open ``http://:8089/request_session`` in a browser on the same network (use ``localhost`` on the same machine). The first launch spends several minutes loading checkpoints and compiling kernels; later launches reuse the cached assets. -On VRAM-constrained GPUs, add ``--offload-text-encoder`` to reduce peak -VRAM usage by about 15 GB: +Inspect the complete resolved runner, launch mode, scenario, and output without +loading checkpoints: .. code-block:: bash - uv run --package flashdreams-omnidreams interactive-drive \ - --stream-mjpeg :8080 \ - --offload-text-encoder + uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams webrtc \ + --manifest configs/launch_manifest/omnidreams_webrtc.yaml \ + --no-instantiate -See :doc:`/models/omnidreams` for the offload trade-offs, scripted -generation, scene variants, WebRTC serving, and multi-GPU options. +See :doc:`/models/omnidreams` for scripted generation, scene variants, +local-window serving, and multi-GPU options. See +:doc:`/api/launch_manifests` for the shared manifest schema. Where to next ------------- @@ -83,7 +87,7 @@ Where to next - :doc:`/models/index`: every shipped model with its CLI slug and the command to run it. - :doc:`/models/omnidreams`: drive a world model in real time with the - ``interactive-drive`` demo. + ``local-window`` or ``webrtc`` launch mode. - :doc:`/developer_guides/inference_pipeline_overview`: the generation loop end to end: KV cache, ring attention, CUDA-graph capture. - :doc:`/developer_guides/config_system`: the configuration layer diff --git a/docs/source/troubleshooting.rst b/docs/source/troubleshooting.rst index 53d039781..b4820c3a2 100644 --- a/docs/source/troubleshooting.rst +++ b/docs/source/troubleshooting.rst @@ -26,7 +26,7 @@ CUDA or PyTorch build mismatch **Symptoms:** - A CUDA extension fails to build or load. -- ``interactive-drive --manifest example_world_model_perf.yaml`` exits instead +- ``flashdreams-run local-window --manifest `` exits instead of falling back to the default PyTorch path. - Errors mention ``nvcc``, a CUDA version, a GPU architecture, or missing CUDA libraries. diff --git a/docs/unified_demo_launcher_migration_plan.md b/docs/unified_demo_launcher_migration_plan.md new file mode 100644 index 000000000..2236b58c2 --- /dev/null +++ b/docs/unified_demo_launcher_migration_plan.md @@ -0,0 +1,134 @@ +# Unified Demo Launcher Migration Plan + +## Goal + +Provide one memorable, extensible command for every FlashDreams launch: + +```bash +flashdreams-run [mode] [--manifest PATH] [overrides...] +``` + +`run` is the default mode. The initial shared modes are `run`, `mp4`, `null`, +`webrtc`, and `local-window`. A runner advertises only the modes it supports. + +Examples: + +```bash +flashdreams-run lingbot-world-fast webrtc --manifest lingbot-live.yaml +flashdreams-run lingbot-world-fast mp4 --manifest lingbot-replay.yaml +flashdreams-run omnidreams webrtc \ + --manifest omnidreams-drive.yaml +flashdreams-run omnidreams-perf local-window \ + --manifest example_world_model_perf.yaml +``` + +The runner slug remains the canonical model/preset selector. It is a stable, +fully specified public identity; a second model-name registry is not needed. + +## Architecture decisions + +- [x] Add a model-agnostic launch capability contract in `flashdreams`. + The contract must not import from `integrations/`. +- [x] Make integrations register their launch capability through a config slot + or plugin entry point, preserving `core -> infra -> integrations`. +- [x] Replace the current output-target `module + argv` translation and + `runpy` handoff with direct, typed launch construction. +- [x] Keep the generic runner path as mode `run`; do not force batch-only + integrations to implement demo modes. +- [x] Treat `local-window` as a first-class launch mode even when its backend + remains model-specific. + +## Manifest contract + +- [x] Add a versioned `FlashDreamsLaunchManifest` rather than extending + OmniDreams' `WorldModelManifest`, whose fields are intentionally + OmniDreams-specific. +- [x] The manifest must contain `schema_version`, `runner`, and `mode`, with + optional `runner_overrides`, `scenario`, and `output` mappings. +- [x] Resolve file paths relative to the manifest directory. +- [x] Reject unknown top-level fields and reject a manifest runner/mode that + conflicts with explicit command-line selection. +- [x] Make the selected integration validate typed `scenario` and `output` + sections; generic code must not learn model-specific fields. +- [x] Define and test precedence: + + ```text + registered runner preset + < manifest runner_overrides + < manifest scenario/output settings + < explicit CLI overrides + ``` + +- [x] Support `--no-instantiate` by printing the resolved runner, mode, + manifest path, scenario, and output without loading a model. + +## Central CLI + +- [x] Parse `flashdreams-run [mode]` before building the + mode-specific typed CLI. +- [x] Make `flashdreams-run --help` list supported modes and make + `flashdreams-run --help` list only valid options. +- [x] Fail unsupported runner/mode pairs before CUDA initialization. +- [x] Remove the temporary `flashdreams-run --output ` + compatibility alias after central mode-dispatch tests pass. +- [x] Retain `torchrun --nproc_per_node=N --no-python flashdreams-run ...` + as the multi-rank invocation form. + +## Integration migration + +### OmniDreams + +- [x] Move `omnidreams-demo replay`, `omnidreams-demo webrtc`, and + `interactive-drive` behind an OmniDreams launch capability. +- [x] Support `mp4`, `null`, `webrtc`, and `local-window`. +- [x] Accept existing `example_world_model*.yaml` files through either a + compatibility reader or converted manifests. +- [x] Preserve scene, weather, camera, conditioning, performance, native + acceleration, cache, VAE, seed, and post-process controls. +- [ ] Implement shared-runtime equivalents for every legacy WebRTC fallback + (including multi-rank serving and debug HDMap streaming) before deletion. + +### LingBot + +- [x] Move `lingbot-demo replay|webrtc` behind a LingBot launch capability. +- [x] Support `mp4` and `webrtc`. +- [x] Cover example assets, prompt, first frame, poses, intrinsics, compile, + warmup, WebRTC, and context-parallel settings in typed options/manifests. + +### Other integrations + +- [x] Expose `run` for every registered runner. +- [x] Add other modes only when an integration has a real implementation and + typed adapter for it. + +## Validation and documentation + +- [x] Add CPU tests for discovery, help, invalid pairs, manifest validation, + precedence, relative paths, and `--no-instantiate` output. +- [ ] Add GPU coverage for LingBot MP4/WebRTC and OmniDreams null/MP4/WebRTC/ + local-window, including multi-rank WebRTC where supported. +- [x] Rewrite the quickstart and model docs around the central command. +- [x] Publish one launch-manifest guide with schema, examples, precedence, + overrides, and reproducibility guidance. +- [x] Replace every demo-launch command in README, docs, CI, and package help. + +## Legacy removal gates + +- [x] Remove `OutputTargetAdapter`, `OutputTargetSpec`, and the `runpy` + output-launch bridge after direct launch parity is verified. +- [x] Remove the LingBot and OmniDreams output-target adapters. +- [x] Remove `lingbot-demo`, `omnidreams-demo`, direct WebRTC server launch + commands, and `interactive-drive` as model-launch entry points. +- [x] Retain preparation, evaluation, and controller-configuration utilities; + they are not model launch commands. +- [x] Remove compatibility aliases only after their central-command + replacements have test coverage and documentation. + +## Completion criteria + +- [x] Every supported demo launches through `flashdreams-run `. +- [x] Every supported demo configuration is expressible as a manifest plus + explicit overrides. +- [x] The central launcher does not invoke a second model CLI. +- [x] No legacy demo/server launch command remains documented or shipped. +- [ ] CPU and GPU tests cover every retained launch mode. diff --git a/flashdreams/flashdreams/configs/runner_configs.py b/flashdreams/flashdreams/configs/runner_configs.py index f017d118a..658b669fd 100644 --- a/flashdreams/flashdreams/configs/runner_configs.py +++ b/flashdreams/flashdreams/configs/runner_configs.py @@ -83,7 +83,9 @@ def all_runners() -> OrderedDict[str, RunnerConfig]: return _sort(runners) -def _annotated_base_runner_union(): +def _annotated_base_runner_union( + runners: Mapping[str, RunnerConfig] | None = None, +): """Build the tyro subcommand union over every discovered runner. Built lazily so importing this module never pays the entry-point @@ -97,7 +99,7 @@ def _annotated_base_runner_union(): * ``FlagConversionOff`` -- don't auto-flip booleans into ``--no-foo`` flags inside nested configs. """ - runners = all_runners() + runners = all_runners() if runners is None else _sort(runners) descriptions = {k: cfg.description for k, cfg in runners.items()} # ``Any`` because ty rejects the runtime tyro union as a type-form # arg to the ``SuppressFixed`` / ``FlagConversionOff`` markers below. diff --git a/flashdreams/flashdreams/infra/runner.py b/flashdreams/flashdreams/infra/runner.py index 060adbd31..f34723a0f 100644 --- a/flashdreams/flashdreams/infra/runner.py +++ b/flashdreams/flashdreams/infra/runner.py @@ -63,8 +63,8 @@ class RunnerConfig(InstantiateConfig): per-runner ``--help`` (it's metadata, not a knob); a non-empty value is enforced for in-tree runners by the registry test.""" - output_adapter: Annotated[str | None, tyro.conf.Suppress] = None - """Optional ``module:attribute`` implementing non-CLI output capabilities.""" + launch_capability: Annotated[str | None, tyro.conf.Suppress] = None + """Optional ``module:attribute`` implementing non-``run`` launch modes.""" pipeline: StreamInferencePipelineConfig """Wrapped pipeline config; the runner instantiates and drives it.""" diff --git a/flashdreams/flashdreams/scripts/cli.py b/flashdreams/flashdreams/scripts/cli.py index 55734f696..d9c847b53 100644 --- a/flashdreams/flashdreams/scripts/cli.py +++ b/flashdreams/flashdreams/scripts/cli.py @@ -28,8 +28,8 @@ flashdreams-run wan21-i2v-14b-480p --prompt "..." --image-path frame.png flashdreams-run --no-instantiate template-offline # resolve config only flashdreams-run wan21-t2v-1.3b-480p --postprocess.preset flashvsr-v1.1-sparse-2.0 - flashdreams-run --output webrtc lingbot-world-fast - flashdreams-run --output local-window omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae + flashdreams-run lingbot-world-fast webrtc --host 0.0.0.0 --port 8080 + flashdreams-run omnidreams local-window # Multi-GPU via context-parallelism (integration transformers auto-detect # CP size from the launcher's WORLD group). ``--no-python`` tells @@ -49,63 +49,77 @@ import tyro -from flashdreams.configs.runner_configs import _annotated_base_runner_union +from flashdreams.configs.runner_configs import _annotated_base_runner_union, all_runners from flashdreams.core.distributed import shutdown as shutdown_distributed from flashdreams.core.io.disk import disk_space_error_from_exception from flashdreams.infra.runner import RunnerConfig -from flashdreams.serving.output_targets import ( - OutputLaunchOptions, - OutputMode, - available_output_modes, - launch_output_target, - resolve_output_target, +from flashdreams.serving.launch import ( + LaunchMode, + LaunchOptions, + available_launch_modes, + resolve_launch, ) +from flashdreams.serving.launch_manifest import ( + FlashDreamsLaunchManifest, + load_launch_manifest, +) + +_POSITIONAL_MODES = frozenset({"run", "mp4", "null", "webrtc", "local-window"}) def main( config: RunnerConfig, no_instantiate: bool = False, *, - output: OutputMode = "cli", - output_host: str | None = None, - output_port: int | None = None, - output_manifest: Path | None = None, + mode: LaunchMode = "run", + host: str | None = None, + port: int | None = None, + legacy_world_manifest: Path | None = None, prefer_sw_encoder: bool = False, + launch_manifest: FlashDreamsLaunchManifest | None = None, ) -> None: """Print the resolved config and (by default) run the runner. Under ``torchrun`` only local-rank 0 prints; every rank holds the same resolved config. """ - output_spec = None - output_options = OutputLaunchOptions( - host=output_host, - port=output_port, + resolved_launch = None + launch_options = LaunchOptions( + host=host, + port=port, prefer_sw_encoder=prefer_sw_encoder, - local_window_manifest=output_manifest, + legacy_world_manifest=legacy_world_manifest, + launch_manifest=None if launch_manifest is None else launch_manifest.path, + scenario={} if launch_manifest is None else launch_manifest.scenario, + output={} if launch_manifest is None else launch_manifest.output, ) - if output != "cli": - output_spec = resolve_output_target( + if mode != "run": + resolved_launch = resolve_launch( config, - mode=output, - options=output_options, + mode=mode, + options=launch_options, ) if int(os.environ.get("LOCAL_RANK", "0")) == 0: print(f"Resolved config for {config.runner_name!r}:") print(config) print( - f"Available outputs: {', '.join(available_output_modes(config, output_options))}" + f"Available modes: {', '.join(available_launch_modes(config, launch_options))}" ) - if output_spec is not None: - print(f"Selected output: {output_spec.label}") - print(f"Launch command: {output_spec.command}") - for note in output_spec.notes: + if launch_manifest is not None: + print(f"Launch manifest: {launch_manifest.path}") + print(f"Launch mode: {launch_manifest.mode}") + print(f"Scenario: {dict(launch_manifest.scenario)}") + print(f"Output settings: {dict(launch_manifest.output)}") + if resolved_launch is not None: + print(f"Selected launch: {resolved_launch.label}") + print(f"Launch settings: {dict(resolved_launch.summary)}") + for note in resolved_launch.notes: print(f"Note: {note}") if no_instantiate: return - if output_spec is not None: - launch_output_target(output_spec) + if resolved_launch is not None: + resolved_launch.launch() return runner = config.setup() completed = False @@ -137,57 +151,89 @@ def _run_with_disk_error_handling(fn: Callable[[], None]) -> None: raise -def entrypoint() -> None: +def entrypoint(argv: list[str] | None = None) -> None: """``flashdreams-run`` console-script entry point. Plugin/entry-point discovery is deferred until call time so importing :mod:`flashdreams.scripts.cli` is cheap. """ tyro.extras.set_accent_color("bright_yellow") - union = _annotated_base_runner_union() + raw_args = list(sys.argv[1:] if argv is None else argv) + ( + normalized_args, + runners, + launch_manifest, + mode, + legacy_world_manifest, + ) = _prepare_cli_args(raw_args) + selected_runner_name = next( + (value for value in normalized_args if value in runners), + None, + ) + help_suffix = "" + if selected_runner_name is not None: + help_options = LaunchOptions( + legacy_world_manifest=legacy_world_manifest, + scenario={} if launch_manifest is None else launch_manifest.scenario, + output={} if launch_manifest is None else launch_manifest.output, + ) + supported = available_launch_modes( + runners[selected_runner_name], + help_options, + ) + help_suffix = ( + f" Selected mode: {mode}. Available modes: {', '.join(supported)}." + " Use --manifest PATH for scenario and output settings." + ) + if mode == "webrtc": + help_suffix += ( + " WebRTC CLI overrides: --host HOST, --port PORT, and" + " --prefer-sw-encoder." + ) + runners[selected_runner_name] = dataclasses.replace( + runners[selected_runner_name], + description=runners[selected_runner_name].description + help_suffix, + ) + union = _annotated_base_runner_union(runners) # ``name=""`` on the synthetic ``runner`` field suppresses its own # name from child prefixes, so ``--runner.prompt`` collapses to # ``--prompt`` and ``runner.pipeline.:`` # selectors collapse to ``pipeline.:``. Nested # struct fields keep their own names for disambiguation. + cli_fields: list[tuple] = [ + ("runner", Annotated[union, tyro.conf.arg(name="")]), + ( + "no_instantiate", + bool, + dataclasses.field(default=False), + ), + ] + if mode == "webrtc": + cli_fields.extend( + [ + ( + "host", + str | None, + dataclasses.field(default=None), + ), + ( + "port", + int | None, + dataclasses.field(default=None), + ), + ( + "prefer_sw_encoder", + bool, + dataclasses.field(default=False), + ), + ] + ) args_cls = dataclasses.make_dataclass( "FlashdreamsRunArgs", - [ - ("runner", Annotated[union, tyro.conf.arg(name="")]), - ( - "no_instantiate", - bool, - dataclasses.field(default=False), - ), - ( - "output", - OutputMode, - dataclasses.field(default="cli"), - ), - ( - "output_host", - str | None, - dataclasses.field(default=None), - ), - ( - "output_port", - int | None, - dataclasses.field(default=None), - ), - ( - "output_manifest", - Path | None, - dataclasses.field(default=None), - ), - ( - "prefer_sw_encoder", - bool, - dataclasses.field(default=False), - ), - ], + cli_fields, ) - args_cls.__doc__ = __doc__ + args_cls.__doc__ = (__doc__ or "") + help_suffix # Silence ``--help`` / parse-error banners on non-rank-0 ranks so # they print exactly once even though every rank parses argv. Every @@ -196,30 +242,150 @@ def entrypoint() -> None: args = tyro.cli( args_cls, prog="flashdreams-run", - description=__doc__, + description=args_cls.__doc__, console_outputs=_is_rank_zero(), + args=normalized_args, ) # ``args_cls`` is built dynamically; keep the untyped boundary explicit. parsed_args = cast(Any, args) runner_cfg: RunnerConfig = parsed_args.runner no_instantiate: bool = parsed_args.no_instantiate - output: OutputMode = parsed_args.output - output_host: str | None = parsed_args.output_host - output_port: int | None = parsed_args.output_port - output_manifest: Path | None = parsed_args.output_manifest - prefer_sw_encoder: bool = parsed_args.prefer_sw_encoder + host: str | None = getattr(parsed_args, "host", None) + port: int | None = getattr(parsed_args, "port", None) + prefer_sw_encoder: bool = getattr(parsed_args, "prefer_sw_encoder", False) _run_with_disk_error_handling( lambda: main( runner_cfg, no_instantiate, - output=output, - output_host=output_host, - output_port=output_port, - output_manifest=output_manifest, + mode=mode, + host=host, + port=port, + legacy_world_manifest=legacy_world_manifest, prefer_sw_encoder=prefer_sw_encoder, + launch_manifest=launch_manifest, ) ) +def _prepare_cli_args( + args: list[str], +) -> tuple[ + list[str], + dict[str, RunnerConfig], + FlashDreamsLaunchManifest | None, + LaunchMode, + Path | None, +]: + """Normalize positional launch modes and load an optional manifest.""" + normalized, manifest_path = _pop_option(args, "--manifest") + runners = dict(all_runners()) + runner_index = next( + (index for index, value in enumerate(normalized) if value in runners), + None, + ) + if runner_index is None: + if manifest_path is not None: + raise ValueError("--manifest requires an explicit runner slug.") + return normalized, runners, None, "run", None + + runner_name = normalized[runner_index] + positional_mode: LaunchMode | None = None + if runner_index + 1 < len(normalized): + candidate = normalized[runner_index + 1] + if candidate in _POSITIONAL_MODES: + positional_mode = cast(LaunchMode, candidate) + del normalized[runner_index + 1] + + launch_manifest: FlashDreamsLaunchManifest | None = None + legacy_world_manifest: Path | None = None + if manifest_path is not None: + try: + launch_manifest = load_launch_manifest(manifest_path) + except ValueError: + if positional_mode != "local-window": + raise + legacy_world_manifest = Path(manifest_path).expanduser().resolve() + else: + if launch_manifest.runner != runner_name: + raise ValueError( + f"Manifest runner {launch_manifest.runner!r} does not match " + f"selected runner {runner_name!r}." + ) + if positional_mode is not None and launch_manifest.mode != positional_mode: + raise ValueError( + f"Manifest mode {launch_manifest.mode!r} does not match " + f"selected mode {positional_mode!r}." + ) + runners[runner_name] = launch_manifest.apply_runner_overrides( + runners[runner_name] + ) + + raw_mode = positional_mode or ( + "run" if launch_manifest is None else launch_manifest.mode + ) + if raw_mode not in _POSITIONAL_MODES: + raise ValueError( + f"Unsupported launch mode {raw_mode!r}. Expected one of: " + f"{', '.join(sorted(_POSITIONAL_MODES))}." + ) + mode = cast(LaunchMode, raw_mode) + normalized = _hoist_global_options(normalized) + return normalized, runners, launch_manifest, mode, legacy_world_manifest + + +def _pop_option(args: list[str], name: str) -> tuple[list[str], str | None]: + remaining: list[str] = [] + value: str | None = None + index = 0 + while index < len(args): + item = args[index] + if item == name: + if value is not None: + raise ValueError(f"{name} may be specified only once.") + if index + 1 >= len(args): + raise ValueError(f"{name} requires a path.") + value = args[index + 1] + index += 2 + continue + prefix = name + "=" + if item.startswith(prefix): + if value is not None: + raise ValueError(f"{name} may be specified only once.") + value = item[len(prefix) :] + index += 1 + continue + remaining.append(item) + index += 1 + return remaining, value + + +def _hoist_global_options(args: list[str]) -> list[str]: + """Allow central launch flags before or after the runner subcommand.""" + value_options = {"--host", "--port"} + flag_options = {"--no-instantiate", "--prefer-sw-encoder"} + prefix: list[str] = [] + remaining: list[str] = [] + index = 0 + while index < len(args): + item = args[index] + if item in flag_options: + prefix.append(item) + index += 1 + continue + if item in value_options: + if index + 1 >= len(args): + raise ValueError(f"{item} requires a value.") + prefix.extend((item, args[index + 1])) + index += 2 + continue + if any(item.startswith(option + "=") for option in value_options): + prefix.append(item) + index += 1 + continue + remaining.append(item) + index += 1 + return [*prefix, *remaining] + + if __name__ == "__main__": entrypoint() diff --git a/flashdreams/flashdreams/serving/launch.py b/flashdreams/flashdreams/serving/launch.py new file mode 100644 index 000000000..0aaeb0861 --- /dev/null +++ b/flashdreams/flashdreams/serving/launch.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model-neutral launch capabilities for ``flashdreams-run``.""" + +from __future__ import annotations + +import importlib +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from functools import cache +from pathlib import Path +from typing import Literal, Protocol, TypeAlias, runtime_checkable + +from flashdreams.infra.runner import RunnerConfig + +LaunchMode: TypeAlias = Literal["run", "mp4", "null", "webrtc", "local-window"] + + +class LaunchModeUnavailableError(ValueError): + """Raised when a runner does not implement a selected launch mode.""" + + +@dataclass(frozen=True, slots=True) +class LaunchOptions: + """Model-neutral settings passed from the central CLI to an integration.""" + + host: str | None = None + port: int | None = None + prefer_sw_encoder: bool = False + legacy_world_manifest: Path | None = None + launch_manifest: Path | None = None + scenario: Mapping[str, object] = field(default_factory=dict) + output: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True) +class ResolvedLaunch: + """Validated launch ready to execute without invoking another CLI.""" + + mode: LaunchMode + label: str + launch: Callable[[], object] = field(repr=False) + summary: Mapping[str, object] = field(default_factory=dict) + notes: tuple[str, ...] = () + + +@runtime_checkable +class LaunchCapability(Protocol): + """Integration-owned modes and launch construction for one runner config.""" + + def supported_modes( + self, + config: RunnerConfig, + options: LaunchOptions, + ) -> tuple[LaunchMode, ...]: ... + + def resolve( + self, + config: RunnerConfig, + *, + mode: LaunchMode, + options: LaunchOptions, + ) -> ResolvedLaunch | None: ... + + +def available_launch_modes( + config: RunnerConfig, + options: LaunchOptions | None = None, +) -> tuple[LaunchMode, ...]: + """Return the modes implemented for ``config``.""" + options = options or LaunchOptions() + capability = _resolve_capability(config) + if capability is None: + return ("run",) + modes = capability.supported_modes(config, options) + if "run" in modes: + raise ValueError("Launch capabilities must not declare built-in mode 'run'.") + return ("run", *dict.fromkeys(modes)) + + +def resolve_launch( + config: RunnerConfig, + *, + mode: LaunchMode, + options: LaunchOptions | None = None, +) -> ResolvedLaunch: + """Validate and construct a non-``run`` launch.""" + if mode == "run": + raise ValueError("Mode 'run' is executed directly by the selected Runner.") + options = options or LaunchOptions() + capability = _resolve_capability(config) + resolved = ( + None + if capability is None + else capability.resolve(config, mode=mode, options=options) + ) + if resolved is None: + supported = ", ".join(available_launch_modes(config, options)) + raise LaunchModeUnavailableError( + f"Launch mode {mode!r} is not available for runner " + f"{config.runner_name!r}. Supported modes: {supported}." + ) + if resolved.mode != mode: + raise ValueError( + f"Launch capability returned mode {resolved.mode!r} while resolving " + f"{mode!r}." + ) + return resolved + + +def _resolve_capability(config: RunnerConfig) -> LaunchCapability | None: + path = config.launch_capability + if not path: + return None + return _load_launch_capability(path) + + +@cache +def _load_launch_capability(path: str) -> LaunchCapability: + try: + module_name, attribute = path.split(":", 1) + except ValueError as exc: + raise ValueError( + "RunnerConfig.launch_capability must use 'module:attribute' syntax; " + f"got {path!r}." + ) from exc + value = getattr(importlib.import_module(module_name), attribute) + if callable(value) and not isinstance(value, LaunchCapability): + value = value() + if not isinstance(value, LaunchCapability): + raise TypeError( + f"Launch capability {path!r} does not implement LaunchCapability." + ) + return value + + +__all__ = [ + "LaunchCapability", + "LaunchMode", + "LaunchModeUnavailableError", + "LaunchOptions", + "ResolvedLaunch", + "available_launch_modes", + "resolve_launch", +] diff --git a/flashdreams/flashdreams/serving/launch_manifest.py b/flashdreams/flashdreams/serving/launch_manifest.py new file mode 100644 index 000000000..4e6bf254b --- /dev/null +++ b/flashdreams/flashdreams/serving/launch_manifest.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Versioned manifests for ``flashdreams-run`` launch modes.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import yaml + +from flashdreams.infra.config import derive_config +from flashdreams.infra.runner import RunnerConfig + +SCHEMA_VERSION = 1 +_TOP_LEVEL_FIELDS = frozenset( + {"schema_version", "runner", "mode", "runner_overrides", "scenario", "output"} +) +_MODES = frozenset({"run", "mp4", "null", "webrtc", "local-window"}) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class FlashDreamsLaunchManifest: + """Resolved, model-neutral launch manifest.""" + + path: Path + schema_version: int + runner: str + mode: str + runner_overrides: Mapping[str, Any] = field(default_factory=dict) + scenario: Mapping[str, Any] = field(default_factory=dict) + output: Mapping[str, Any] = field(default_factory=dict) + + def apply_runner_overrides(self, config: RunnerConfig) -> RunnerConfig: + """Return ``config`` with manifest overrides applied recursively.""" + if config.runner_name != self.runner: + raise ValueError( + f"Manifest runner {self.runner!r} does not match " + f"selected runner {config.runner_name!r}." + ) + return derive_config(config, **dict(self.runner_overrides)) + + +def load_launch_manifest(path: str | Path) -> FlashDreamsLaunchManifest: + """Load and strictly validate one YAML launch manifest.""" + manifest_path = Path(path).expanduser().resolve() + if not manifest_path.is_file(): + raise FileNotFoundError( + f"Launch manifest path does not exist or is not a file: {manifest_path}. " + "Manifest paths are resolved relative to the current working directory " + f"({Path.cwd()})." + ) + raw = yaml.safe_load(manifest_path.read_text(encoding="utf-8")) or {} + if not isinstance(raw, dict): + raise TypeError(f"Launch manifest {manifest_path} must contain a YAML mapping.") + unknown = sorted(set(raw) - _TOP_LEVEL_FIELDS) + if unknown: + raise ValueError( + f"Unknown launch manifest fields in {manifest_path}: {', '.join(unknown)}." + ) + schema_version = raw.get("schema_version") + if schema_version != SCHEMA_VERSION: + raise ValueError( + f"Unsupported launch manifest schema_version={schema_version!r}; " + f"expected {SCHEMA_VERSION}." + ) + runner = _required_string(raw, "runner", manifest_path) + mode = _required_string(raw, "mode", manifest_path) + if mode not in _MODES: + raise ValueError( + f"Launch manifest {manifest_path} has unsupported mode {mode!r}; " + f"expected one of: {', '.join(sorted(_MODES))}." + ) + manifest_dir = manifest_path.parent + return FlashDreamsLaunchManifest( + path=manifest_path, + schema_version=SCHEMA_VERSION, + runner=runner, + mode=mode, + runner_overrides=_mapping(raw, "runner_overrides", manifest_path), + scenario=_resolve_paths( + _mapping(raw, "scenario", manifest_path), manifest_dir=manifest_dir + ), + output=_resolve_paths( + _mapping(raw, "output", manifest_path), manifest_dir=manifest_dir + ), + ) + + +def _required_string(raw: Mapping[str, Any], key: str, path: Path) -> str: + value = raw.get(key) + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Launch manifest {path} requires non-empty {key!r}.") + return value.strip() + + +def _mapping(raw: Mapping[str, Any], key: str, path: Path) -> Mapping[str, Any]: + value = raw.get(key, {}) + if not isinstance(value, dict): + raise TypeError(f"Launch manifest {path} field {key!r} must be a mapping.") + return value + + +def _resolve_paths(value: Any, *, manifest_dir: Path, key: str = "") -> Any: + if isinstance(value, dict): + return { + str(child_key): _resolve_paths( + child_value, + manifest_dir=manifest_dir, + key=str(child_key), + ) + for child_key, child_value in value.items() + } + if isinstance(value, list): + return [ + _resolve_paths(item, manifest_dir=manifest_dir, key=key) for item in value + ] + if isinstance(value, str) and _is_path_key(key): + path = Path(value).expanduser() + return path if path.is_absolute() else (manifest_dir / path).resolve() + return value + + +def _is_path_key(key: str) -> bool: + normalized = key.replace("-", "_") + return normalized in {"path", "output"} or normalized.endswith( + ("_path", "_paths", "_dir") + ) + + +__all__ = [ + "SCHEMA_VERSION", + "FlashDreamsLaunchManifest", + "load_launch_manifest", +] diff --git a/flashdreams/flashdreams/serving/output_targets.py b/flashdreams/flashdreams/serving/output_targets.py deleted file mode 100644 index a5e9f961b..000000000 --- a/flashdreams/flashdreams/serving/output_targets.py +++ /dev/null @@ -1,159 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Output target selection for ``flashdreams-run``.""" - -from __future__ import annotations - -import importlib -import runpy -import shlex -import sys -from dataclasses import dataclass -from functools import lru_cache -from pathlib import Path -from typing import Literal, Protocol, TypeAlias, runtime_checkable - -from flashdreams.infra.runner import RunnerConfig - -OutputMode: TypeAlias = Literal["cli", "webrtc", "local-window"] - - -class OutputTargetUnavailableError(ValueError): - """Raised when a runner cannot be launched through a requested output.""" - - -@dataclass(frozen=True, slots=True) -class OutputLaunchOptions: - """Common launch options shared by non-CLI output targets.""" - - host: str | None = None - port: int | None = None - prefer_sw_encoder: bool = False - local_window_manifest: Path | None = None - - -@dataclass(frozen=True, slots=True) -class OutputTargetSpec: - """A concrete output target module plus argv translated from a runner config.""" - - mode: OutputMode - label: str - module: str - argv: tuple[str, ...] = () - notes: tuple[str, ...] = () - - @property - def command(self) -> str: - """Return a copy-pasteable module command for diagnostics.""" - return shlex.join(("python", "-m", self.module, *self.argv)) - - -@runtime_checkable -class OutputTargetAdapter(Protocol): - """Integration-owned non-CLI output capabilities for a runner config.""" - - def supported_modes( - self, - config: RunnerConfig, - options: OutputLaunchOptions, - ) -> tuple[OutputMode, ...]: ... - - def resolve( - self, - config: RunnerConfig, - *, - mode: OutputMode, - options: OutputLaunchOptions, - ) -> OutputTargetSpec | None: ... - - -def available_output_modes( - config: RunnerConfig, - options: OutputLaunchOptions | None = None, -) -> tuple[OutputMode, ...]: - """Return output modes known to support ``config``.""" - options = options or OutputLaunchOptions() - adapter = _resolve_adapter(config) - if adapter is None: - return ("cli",) - modes = adapter.supported_modes(config, options) - invalid = [mode for mode in modes if mode == "cli"] - if invalid: - raise ValueError("Output adapters must not declare the built-in CLI mode.") - return ("cli", *dict.fromkeys(modes)) - - -def resolve_output_target( - config: RunnerConfig, - *, - mode: OutputMode, - options: OutputLaunchOptions | None = None, -) -> OutputTargetSpec: - """Resolve a non-CLI output target for a runner config.""" - if mode == "cli": - raise ValueError("CLI mode is run directly by the selected Runner.") - options = options or OutputLaunchOptions() - adapter = _resolve_adapter(config) - spec = ( - None if adapter is None else adapter.resolve(config, mode=mode, options=options) - ) - if spec is None: - supported = ", ".join(available_output_modes(config, options)) - raise OutputTargetUnavailableError( - f"Output mode {mode!r} is not available for runner " - f"{config.runner_name!r}. Supported modes: {supported}." - ) - if spec.mode != mode: - raise ValueError( - f"Output adapter returned mode {spec.mode!r} while resolving {mode!r}." - ) - return spec - - -def launch_output_target(spec: OutputTargetSpec) -> None: - """Execute an output target module as if launched with ``python -m``.""" - original_argv = sys.argv - sys.argv = [spec.module, *spec.argv] - try: - runpy.run_module(spec.module, run_name="__main__") - finally: - sys.argv = original_argv - - -def _resolve_adapter(config: RunnerConfig) -> OutputTargetAdapter | None: - path = config.output_adapter - if not path: - return None - return _load_output_adapter(path) - - -@lru_cache(maxsize=None) -def _load_output_adapter(path: str) -> OutputTargetAdapter: - try: - module_name, attribute = path.split(":", 1) - except ValueError as exc: - raise ValueError( - "RunnerConfig.output_adapter must use 'module:attribute' syntax; " - f"got {path!r}." - ) from exc - value = getattr(importlib.import_module(module_name), attribute) - if callable(value) and not isinstance(value, OutputTargetAdapter): - value = value() - if not isinstance(value, OutputTargetAdapter): - raise TypeError( - f"Output adapter {path!r} does not implement OutputTargetAdapter." - ) - return value - - -__all__ = [ - "OutputLaunchOptions", - "OutputMode", - "OutputTargetSpec", - "OutputTargetAdapter", - "OutputTargetUnavailableError", - "available_output_modes", - "launch_output_target", - "resolve_output_target", -] diff --git a/flashdreams/pyproject.toml b/flashdreams/pyproject.toml index fb5295781..c0286aa0c 100644 --- a/flashdreams/pyproject.toml +++ b/flashdreams/pyproject.toml @@ -36,6 +36,7 @@ dependencies = [ # non-subclassable TypeAliasType in NumPy 2.5 (notably on Python 3.13). "numpy>=1.24,<2.5", "nvidia-ml-py>=12.0", + "pyyaml>=6.0", "safetensors>=0.4", "tqdm>=4.60", "transformers>=5.0,<6", diff --git a/flashdreams/tests/test_benchmark_harness.py b/flashdreams/tests/test_benchmark_harness.py index 1de03b5db..1957baf0a 100644 --- a/flashdreams/tests/test_benchmark_harness.py +++ b/flashdreams/tests/test_benchmark_harness.py @@ -260,10 +260,8 @@ def test_shipped_omnidreams_demo_replay_scenarios_load() -> None: assert baseline.report_group is not None assert baseline.report_group.id == "omnidreams-demo" assert _command_value(baseline.command, "--total-blocks") == "226" - assert "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae" in baseline.command - assert "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf" not in ( - baseline.command - ) + assert "omnidreams" in baseline.command + assert "omnidreams-perf" not in baseline.command assert baseline.quality_baseline_compare is False demo = scenarios["omnidreams-sv-demo-replay"] @@ -311,10 +309,8 @@ def test_shipped_deterministic_quality_scenarios_load() -> None: "-m", "tools.benchmarks.strict_run", ) - assert "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae" in omnidreams.command - assert "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf" not in ( - omnidreams.command - ) + assert "omnidreams" in omnidreams.command + assert "omnidreams-perf" not in omnidreams.command assert omnidreams.report_group is not None assert omnidreams.report_group.id == "omnidreams" assert omnidreams.report_group.name == "Omnidreams" @@ -337,13 +333,8 @@ def test_shipped_deterministic_quality_scenarios_load() -> None: omnidreams_review = scenarios["omnidreams-sv-one-minute-review"] assert _command_value(omnidreams_review.command, "--total-blocks") == "226" - assert ( - "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae" - in omnidreams_review.command - ) - assert "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf" not in ( - omnidreams_review.command - ) + assert "omnidreams" in omnidreams_review.command + assert "omnidreams-perf" not in omnidreams_review.command assert "--pipeline.diffusion-model.seed" in omnidreams_review.command assert omnidreams_review.quality_baseline_compare is False diff --git a/flashdreams/tests/test_launch.py b/flashdreams/tests/test_launch.py new file mode 100644 index 000000000..e58528bd2 --- /dev/null +++ b/flashdreams/tests/test_launch.py @@ -0,0 +1,231 @@ +# 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 types import SimpleNamespace +from typing import cast + +import pytest + +from flashdreams.infra.runner import RunnerConfig +from flashdreams.scripts import cli +from flashdreams.serving import launch as launch_module +from flashdreams.serving.launch import ( + LaunchModeUnavailableError, + LaunchOptions, + ResolvedLaunch, + available_launch_modes, + resolve_launch, +) + +pytestmark = pytest.mark.ci_cpu + + +def _runner_config( + *, + runner_name: str, + num_views: int = 1, + pipeline_name: str | None = None, +) -> RunnerConfig: + launch_capability = None + if runner_name.startswith("lingbot-world"): + launch_capability = "lingbot.launch:LAUNCH_CAPABILITY" + elif runner_name == "omnidreams" or runner_name.startswith("omnidreams-"): + launch_capability = "omnidreams.launch:LAUNCH_CAPABILITY" + pipeline = SimpleNamespace( + name=pipeline_name or runner_name, + diffusion_model=SimpleNamespace( + seed=42, + transformer=SimpleNamespace(num_views=num_views, compile_network=True), + ), + ) + return cast( + RunnerConfig, + SimpleNamespace( + runner_name=runner_name, + launch_capability=launch_capability, + pipeline=pipeline, + device="cuda:1", + pixel_height=480, + pixel_width=832, + fps=20, + output_fps=24, + example_idx=3, + postprocess=SimpleNamespace(preset=""), + ), + ) + + +def test_lingbot_mp4_launch_validates_manifest_sections(tmp_path: Path) -> None: + resolved = resolve_launch( + _runner_config(runner_name="lingbot-world-fast"), + mode="mp4", + options=LaunchOptions( + scenario={"example_idx": 2, "total_blocks": 4}, + output={"path": tmp_path / "demo.mp4", "fps": 12}, + ), + ) + + assert resolved.mode == "mp4" + assert resolved.summary["output_path"] == tmp_path / "demo.mp4" + + +def test_lingbot_launch_rejects_unknown_integration_fields() -> None: + with pytest.raises(ValueError, match="Unsupported LingBot scenario fields: typo"): + resolve_launch( + _runner_config(runner_name="lingbot-world-fast"), + mode="webrtc", + options=LaunchOptions(scenario={"typo": True}), + ) + + +def test_omnidreams_webrtc_is_rejected_for_multi_view() -> None: + config = _runner_config( + runner_name="omnidreams-mv-2steps-chunk4-loc8-pshuffle-lighttae", + num_views=4, + ) + + assert available_launch_modes(config) == ("run", "mp4", "null") + with pytest.raises( + LaunchModeUnavailableError, + match="Supported modes: run, mp4, null", + ): + resolve_launch(config, mode="webrtc") + + +def test_omnidreams_webrtc_honors_explicit_network_precedence() -> None: + resolved = resolve_launch( + _runner_config( + runner_name="omnidreams", + ), + mode="webrtc", + options=LaunchOptions( + host="127.0.0.1", + port=9011, + output={"host": "0.0.0.0", "port": 8082}, + ), + ) + + assert resolved.summary["host"] == "127.0.0.1" + assert resolved.summary["port"] == 9011 + + +def test_omnidreams_mp4_short_slug_uses_default_output_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from omnidreams.demo import app + + calls: list[dict[str, object]] = [] + monkeypatch.setattr( + app, "launch_from_runner", lambda **kwargs: calls.append(kwargs) + ) + config = _runner_config( + runner_name="omnidreams", + pipeline_name="omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae", + ) + + resolved = resolve_launch(config, mode="mp4") + + assert resolved.summary == { + "runner": "omnidreams", + "mode": "mp4", + "device": "cuda:1", + "output_path": Path("outputs/omnidreams.mp4"), + } + resolved.launch() + assert calls[0]["config"] is config + assert calls[0]["mode"] == "mp4" + assert calls[0]["output"] == {"path": Path("outputs/omnidreams.mp4")} + + +def test_omnidreams_local_window_accepts_legacy_world_manifest() -> None: + config = _runner_config( + runner_name="omnidreams-perf", + pipeline_name="omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf", + ) + manifest = Path("custom.yaml") + options = LaunchOptions(legacy_world_manifest=manifest) + + assert available_launch_modes(config, options) == ( + "run", + "mp4", + "null", + "webrtc", + "local-window", + ) + resolved = resolve_launch(config, mode="local-window", options=options) + assert resolved.summary["world_model_manifest"] == manifest + + +def test_capabilities_extend_launch_without_shared_routing_changes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _FakeCapability: + def supported_modes(self, config, options): + del config, options + return ("webrtc",) + + def resolve(self, config, *, mode, options): + del config, options + if mode != "webrtc": + return None + return ResolvedLaunch( + mode="webrtc", + label="plugin launch", + launch=lambda: None, + ) + + config = _runner_config(runner_name="third-party-model") + config.launch_capability = "plugin:capability" + monkeypatch.setattr( + launch_module, + "_load_launch_capability", + lambda path: _FakeCapability(), + ) + + assert available_launch_modes(config) == ("run", "webrtc") + assert resolve_launch(config, mode="webrtc").label == "plugin launch" + + +def test_resolved_launch_calls_integration_directly( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from lingbot.demo import app + + calls: list[dict[str, object]] = [] + monkeypatch.setattr( + app, "launch_from_runner", lambda **kwargs: calls.append(kwargs) + ) + config = _runner_config(runner_name="lingbot-world-fast") + resolved = resolve_launch( + config, + mode="webrtc", + options=LaunchOptions(port=9000), + ) + + resolved.launch() + + assert calls[0]["config"] is config + assert calls[0]["mode"] == "webrtc" + assert calls[0]["port"] == 9000 + + +def test_no_instantiate_reports_launch_without_setting_up_model( + capsys: pytest.CaptureFixture[str], +) -> None: + config = _runner_config(runner_name="lingbot-world-fast") + + cli.main( + config, + no_instantiate=True, + mode="webrtc", + host="127.0.0.1", + port=9090, + ) + + output = capsys.readouterr().out + assert "Available modes: run, mp4, webrtc" in output + assert "Selected launch: LingBot WebRTC server" in output + assert "'host': '127.0.0.1'" in output diff --git a/flashdreams/tests/test_launch_manifest.py b/flashdreams/tests/test_launch_manifest.py new file mode 100644 index 000000000..08622bba2 --- /dev/null +++ b/flashdreams/tests/test_launch_manifest.py @@ -0,0 +1,312 @@ +# 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 types import SimpleNamespace +from typing import Any, cast + +import pytest + +from flashdreams.infra.runner import RunnerConfig +from flashdreams.scripts import cli +from flashdreams.serving.launch import resolve_launch +from flashdreams.serving.launch_manifest import load_launch_manifest + +pytestmark = pytest.mark.ci_cpu + + +def _config(name: str = "demo-runner") -> RunnerConfig: + return cast( + RunnerConfig, + SimpleNamespace( + runner_name=name, + launch_capability=None, + device="cuda:0", + pipeline=SimpleNamespace(diffusion_model=SimpleNamespace(seed=1)), + ), + ) + + +def test_launch_manifest_loads_strict_sections_and_relative_paths( + tmp_path: Path, +) -> None: + manifest_path = tmp_path / "launch.yaml" + manifest_path.write_text( + """\ +schema_version: 1 +runner: demo-runner +mode: mp4 +runner_overrides: + device: cuda:3 +scenario: + image_path: assets/frame.png +output: + path: results/demo.mp4 +""", + encoding="utf-8", + ) + + manifest = load_launch_manifest(manifest_path) + + assert manifest.runner == "demo-runner" + assert manifest.mode == "mp4" + assert manifest.scenario["image_path"] == tmp_path / "assets/frame.png" + assert manifest.output["path"] == tmp_path / "results/demo.mp4" + assert manifest.apply_runner_overrides(_config()).device == "cuda:3" + + +def test_launch_manifest_does_not_guess_configs_directory( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + manifest_path = tmp_path / "configs" / "launch_manifest" / "demo.yaml" + manifest_path.parent.mkdir(parents=True) + manifest_path.write_text( + 'schema_version: 1\nrunner: demo-runner\nmode: "null"\n', + encoding="utf-8", + ) + monkeypatch.chdir(tmp_path) + + requested_path = tmp_path / "launch_manifest" / "demo.yaml" + with pytest.raises(FileNotFoundError) as exc_info: + load_launch_manifest("launch_manifest/demo.yaml") + + message = str(exc_info.value) + assert str(requested_path) in message + assert "resolved relative to the current working directory" in message + + +@pytest.mark.parametrize( + "body, match", + [ + ("schema_version: 2\nrunner: demo\nmode: run\n", "schema_version"), + ("schema_version: 1\nrunner: demo\nmode: run\nextra: true\n", "extra"), + ("schema_version: 1\nrunner: ''\nmode: run\n", "runner"), + ("schema_version: 1\nrunner: demo\nmode: other\n", "unsupported mode"), + ], +) +def test_launch_manifest_rejects_invalid_documents( + tmp_path: Path, + body: str, + match: str, +) -> None: + path = tmp_path / "invalid.yaml" + path.write_text(body, encoding="utf-8") + + with pytest.raises(ValueError, match=match): + load_launch_manifest(path) + + +def test_positional_mode_and_manifest_are_normalized_before_tyro( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + manifest_path = tmp_path / "launch.yaml" + manifest_path.write_text( + """\ +schema_version: 1 +runner: demo-runner +mode: webrtc +runner_overrides: + device: cuda:2 +scenario: + scene_dir: scenes +output: + port: 9000 +""", + encoding="utf-8", + ) + monkeypatch.setattr(cli, "all_runners", lambda: {"demo-runner": _config()}) + + args, runners, manifest, mode, legacy_manifest = cli._prepare_cli_args( + ["demo-runner", "webrtc", "--manifest", str(manifest_path)] + ) + + assert args == ["demo-runner"] + assert runners["demo-runner"].device == "cuda:2" + assert manifest is not None + assert manifest.scenario["scene_dir"] == tmp_path / "scenes" + assert mode == "webrtc" + assert legacy_manifest is None + + +def test_positional_mode_must_match_manifest( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + path = tmp_path / "launch.yaml" + path.write_text( + 'schema_version: 1\nrunner: demo-runner\nmode: "null"\n', + encoding="utf-8", + ) + monkeypatch.setattr(cli, "all_runners", lambda: {"demo-runner": _config()}) + + with pytest.raises(ValueError, match="does not match selected mode"): + cli._prepare_cli_args(["demo-runner", "webrtc", "--manifest", str(path)]) + + +def test_run_mode_preserves_default_runner_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(cli, "all_runners", lambda: {"demo-runner": _config()}) + + args, _, manifest, mode, legacy_manifest = cli._prepare_cli_args( + ["demo-runner", "run"] + ) + + assert args == ["demo-runner"] + assert manifest is None + assert mode == "run" + assert legacy_manifest is None + + +def test_short_omnidreams_slug_and_mp4_mode_are_parsed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + cli, + "all_runners", + lambda: {"omnidreams": _config("omnidreams")}, + ) + + args, runners, manifest, mode, legacy_manifest = cli._prepare_cli_args( + ["omnidreams", "mp4"] + ) + + assert args == ["omnidreams"] + assert runners["omnidreams"].runner_name == "omnidreams" + assert manifest is None + assert mode == "mp4" + assert legacy_manifest is None + + +def test_central_options_are_allowed_after_runner( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(cli, "all_runners", lambda: {"demo-runner": _config()}) + + args, _, _, mode, _ = cli._prepare_cli_args( + [ + "demo-runner", + "webrtc", + "--host", + "127.0.0.1", + "--no-instantiate", + ] + ) + + assert args == [ + "--host", + "127.0.0.1", + "--no-instantiate", + "demo-runner", + ] + assert mode == "webrtc" + + +def test_legacy_local_window_manifest_is_routed_without_second_cli( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + path = tmp_path / "example_world_model.yaml" + path.write_text("resolution_wh: [1280, 704]\n", encoding="utf-8") + monkeypatch.setattr(cli, "all_runners", lambda: {"demo-runner": _config()}) + + args, _, launch_manifest, mode, legacy_manifest = cli._prepare_cli_args( + ["demo-runner", "local-window", "--manifest", str(path)] + ) + + assert args == ["demo-runner"] + assert launch_manifest is None + assert mode == "local-window" + assert legacy_manifest == path.resolve() + + +def test_mode_help_lists_only_mode_specific_central_overrides( + capsys: pytest.CaptureFixture[str], +) -> None: + with pytest.raises(SystemExit) as webrtc_exit: + cli.entrypoint(["lingbot-world-fast", "webrtc", "--help"]) + assert webrtc_exit.value.code == 0 + webrtc_help = capsys.readouterr().out + assert "Available modes: run, mp4, webrtc" in webrtc_help + assert "--host HOST" in webrtc_help + + with pytest.raises(SystemExit) as mp4_exit: + cli.entrypoint(["lingbot-world-fast", "mp4", "--help"]) + assert mp4_exit.value.code == 0 + mp4_help = capsys.readouterr().out + assert "Selected mode: mp4" in mp4_help + assert "--host HOST" not in mp4_help + + +def test_explicit_runner_cli_override_wins_over_manifest( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + path = tmp_path / "launch.yaml" + path.write_text( + """\ +schema_version: 1 +runner: lingbot-world-fast +mode: webrtc +runner_overrides: + device: cuda:2 +""", + encoding="utf-8", + ) + captured: list[tuple[RunnerConfig, dict[str, object]]] = [] + + def fake_main(config: RunnerConfig, no_instantiate: bool, **kwargs) -> None: + assert no_instantiate is True + captured.append((config, kwargs)) + + monkeypatch.setattr(cli, "main", fake_main) + + cli.entrypoint( + [ + "lingbot-world-fast", + "webrtc", + "--manifest", + str(path), + "--device", + "cuda:3", + "--no-instantiate", + ] + ) + + assert captured[0][0].device == "cuda:3" + assert captured[0][1]["mode"] == "webrtc" + + +@pytest.mark.parametrize( + "filename", + [ + "lingbot_mp4.yaml", + "lingbot_webrtc.yaml", + "omnidreams_local_window.yaml", + "omnidreams_mp4.yaml", + "omnidreams_null.yaml", + "omnidreams_webrtc.yaml", + ], +) +def test_documented_launch_manifests_resolve(filename: str) -> None: + repo_root = Path(__file__).resolve().parents[2] + manifest = load_launch_manifest( + repo_root / "configs" / "launch_manifest" / filename + ) + config = manifest.apply_runner_overrides(cli.all_runners()[manifest.runner]) + + if manifest.mode != "run": + resolved = resolve_launch( + config, + mode=cast(Any, manifest.mode), + options=cli.LaunchOptions( + launch_manifest=manifest.path, + scenario=manifest.scenario, + output=manifest.output, + ), + ) + assert resolved.mode == manifest.mode diff --git a/flashdreams/tests/test_output_targets.py b/flashdreams/tests/test_output_targets.py deleted file mode 100644 index 1421ceb62..000000000 --- a/flashdreams/tests/test_output_targets.py +++ /dev/null @@ -1,237 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import sys -from pathlib import Path -from types import SimpleNamespace -from typing import cast - -import pytest - -from flashdreams.infra.runner import RunnerConfig -from flashdreams.serving import output_targets as output_targets_module -from flashdreams.serving.output_targets import ( - OutputLaunchOptions, - OutputTargetSpec, - OutputTargetUnavailableError, - available_output_modes, - launch_output_target, - resolve_output_target, -) - -pytestmark = pytest.mark.ci_cpu - - -def _runner_config( - *, - runner_name: str, - pipeline_name: str | None = None, - num_views: int = 1, - compile_network: bool = True, -) -> RunnerConfig: - output_adapter = None - if runner_name.startswith("lingbot-world"): - output_adapter = "lingbot.output_targets:OUTPUT_TARGET_ADAPTER" - elif runner_name.startswith("omnidreams-"): - output_adapter = "omnidreams.output_targets:OUTPUT_TARGET_ADAPTER" - transformer = SimpleNamespace( - num_views=num_views, - compile_network=compile_network, - ) - pipeline = SimpleNamespace( - name=pipeline_name or runner_name, - diffusion_model=SimpleNamespace(seed=42, transformer=transformer), - ) - return cast( - RunnerConfig, - SimpleNamespace( - runner_name=runner_name, - output_adapter=output_adapter, - pipeline=pipeline, - device="cuda:1", - pixel_height=480, - pixel_width=832, - fps=20, - output_fps=24, - example_idx=3, - postprocess=SimpleNamespace(preset=""), - ), - ) - - -def test_lingbot_webrtc_target_translates_runner_config() -> None: - config = _runner_config( - runner_name="lingbot-world-fast", - compile_network=False, - ) - - spec = resolve_output_target( - config, - mode="webrtc", - options=OutputLaunchOptions( - host="127.0.0.1", - port=9010, - prefer_sw_encoder=True, - ), - ) - - assert spec.module == "lingbot.demo.app" - assert spec.argv == ( - "webrtc", - "--preset-id", - "lingbot-world-fast", - "--device", - "cuda:1", - "--fps", - "20", - "--video-height", - "480", - "--video-width", - "832", - "--no-compile", - "--example-idx", - "3", - "--host", - "127.0.0.1", - "--port", - "9010", - "--prefer-sw-encoder", - ) - - -def test_omnidreams_webrtc_target_rejects_multi_view() -> None: - config = _runner_config( - runner_name="omnidreams-mv-2steps-chunk4-loc8-pshuffle-lighttae", - num_views=4, - ) - - assert available_output_modes(config) == ("cli",) - with pytest.raises(OutputTargetUnavailableError, match="Supported modes: cli"): - resolve_output_target( - config, - mode="webrtc", - ) - - -def test_omnidreams_webrtc_target_uses_shared_demo_entry_point() -> None: - config = _runner_config( - runner_name="omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae", - ) - - spec = resolve_output_target( - config, - mode="webrtc", - options=OutputLaunchOptions( - host="127.0.0.1", - port=9011, - prefer_sw_encoder=True, - ), - ) - - assert spec.module == "omnidreams.demo.app" - assert spec.argv == ( - "webrtc", - "--preset-id", - "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae", - "--device", - "cuda:1", - "--fps", - "24", - "--video-height", - "480", - "--video-width", - "832", - "--seed", - "42", - "--host", - "127.0.0.1", - "--port", - "9011", - "--prefer-sw-encoder", - ) - - -def test_output_capabilities_can_be_added_without_shared_routing_change( - monkeypatch: pytest.MonkeyPatch, -) -> None: - class _FakeAdapter: - def supported_modes(self, config, options): - del config, options - return ("webrtc",) - - def resolve(self, config, *, mode, options): - del config, options - if mode != "webrtc": - return None - return OutputTargetSpec( - mode="webrtc", - label="plugin demo", - module="plugin.demo", - ) - - config = _runner_config(runner_name="third-party-model") - config.output_adapter = "plugin:adapter" - monkeypatch.setattr( - output_targets_module, - "_load_output_adapter", - lambda path: _FakeAdapter(), - ) - - assert available_output_modes(config) == ("cli", "webrtc") - assert resolve_output_target(config, mode="webrtc").module == "plugin.demo" - - -def test_omnidreams_local_window_target_uses_manifest_override() -> None: - config = _runner_config( - runner_name="omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae" - ) - config.postprocess.preset = "flashvsr-v1.1-sparse-2.0" - - spec = resolve_output_target( - config, - mode="local-window", - options=OutputLaunchOptions(local_window_manifest=Path("custom.yaml")), - ) - - assert spec.module == "omnidreams.interactive_drive" - assert spec.argv == ( - "--manifest", - "custom.yaml", - "--postprocess-preset", - "flashvsr-v1.1-sparse-2.0", - ) - assert spec.notes - - -def test_output_manifest_extends_local_window_availability() -> None: - config = _runner_config(runner_name="omnidreams-sv-2steps-chunk3-loc6-vae-vae") - options = OutputLaunchOptions(local_window_manifest=Path("custom.yaml")) - - assert available_output_modes( - config, - options, - ) == ("cli", "webrtc", "local-window") - - -def test_launch_output_target_runs_module_with_translated_argv(monkeypatch) -> None: - calls: list[tuple[str, str, tuple[str, ...]]] = [] - original_argv = list(sys.argv) - - def fake_run_module(module: str, *, run_name: str) -> None: - calls.append((module, run_name, tuple(sys.argv))) - - monkeypatch.setattr(output_targets_module.runpy, "run_module", fake_run_module) - - launch_output_target( - OutputTargetSpec( - mode="webrtc", - label="test", - module="demo.server", - argv=("--port", "9000"), - ) - ) - - assert calls == [("demo.server", "__main__", ("demo.server", "--port", "9000"))] - assert sys.argv == original_argv diff --git a/flashdreams/tools/benchmarks/scenarios.py b/flashdreams/tools/benchmarks/scenarios.py index 8e9b03a71..7b222beac 100644 --- a/flashdreams/tools/benchmarks/scenarios.py +++ b/flashdreams/tools/benchmarks/scenarios.py @@ -305,7 +305,7 @@ def built_in_scenarios() -> dict[str, BenchmarkScenario]: tags=("public", "runner", "world-model", "i2v"), command=( "flashdreams-run", - "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae", + "omnidreams", "--example-data", "--total-blocks", "8", diff --git a/integrations/lingbot/README.md b/integrations/lingbot/README.md index 71c71f91c..5693826ac 100644 --- a/integrations/lingbot/README.md +++ b/integrations/lingbot/README.md @@ -204,9 +204,8 @@ for i in range(total_blocks): ## Run (compatibility WebRTC server) -The `lingbot.webrtc` subpackage exposes a minimal WebRTC server that -binds the integration pipeline to keyboard input over a DataChannel and streams the -generated video back to the browser. +Use the central launcher to bind the LingBot pipeline to keyboard input over a +DataChannel and stream generated video back to the browser. - `GET /request_session` serves a standalone viewer page (HTML/CSS/JS files on disk, not inlined in Python). - `POST /api/webrtc/offer` performs SDP offer/answer signaling. @@ -220,17 +219,21 @@ generated video back to the browser. From repository root: ```bash -uv run --package flashdreams-lingbot python -m lingbot.webrtc.server \ - --host 0.0.0.0 --port 8089 --config_name lingbot-world-fast-taehv-window15-sink3 +uv run --package flashdreams-lingbot flashdreams-run \ + lingbot-world-fast-taehv-window15-sink3 webrtc \ + --host 0.0.0.0 --port 8089 # 4 GPUs uv run --package flashdreams-lingbot \ - python -m torch.distributed.run --standalone --nnodes=1 --nproc_per_node=4 \ - -m lingbot.webrtc.server \ - --host 0.0.0.0 --port 8089 \ - --config_name lingbot-world-fast-taehv-window15-sink3 + torchrun --standalone --nnodes=1 --nproc_per_node=4 --no-python \ + flashdreams-run lingbot-world-fast-taehv-window15-sink3 webrtc \ + --host 0.0.0.0 --port 8089 ``` +For a reproducible checked-in configuration, use +`--manifest configs/launch_manifest/lingbot_webrtc.yaml`; the manifest's runner +must match the selected runner slug. + Then open: - [http://localhost:8089/request_session](http://localhost:8089/request_session) diff --git a/integrations/lingbot/lingbot/demo/app.py b/integrations/lingbot/lingbot/demo/app.py index 8aa495dff..d83144974 100644 --- a/integrations/lingbot/lingbot/demo/app.py +++ b/integrations/lingbot/lingbot/demo/app.py @@ -7,8 +7,9 @@ import argparse from pathlib import Path -from typing import Any +from typing import Any, Literal, cast +from flashdreams.infra.runner import RunnerConfig from flashdreams.runtime import InferenceConfig from flashdreams.runtime.demo import ( DemoSpec, @@ -17,6 +18,11 @@ WebRTCOutputSpec, ) from flashdreams.runtime.demo.app import DemoApplication +from flashdreams.runtime.demo.replay import run_replay_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, @@ -159,6 +165,129 @@ def main(argv: list[str] | None = None) -> None: _APPLICATION.main(argv) +def launch_from_runner( + *, + config: RunnerConfig, + mode: Literal["mp4", "webrtc"], + scenario: dict[str, object], + output: dict[str, object], + host: str | None = None, + port: int | None = None, + prefer_sw_encoder: bool = False, +) -> object: + """Launch a LingBot demo directly from a resolved runner configuration.""" + configure_logging() + preset_id = str(getattr(config.pipeline, "name", config.runner_name)) + if mode == "mp4": + output_path = output.get("path") or output.get("output") + if output_path is None: + raise ValueError("LingBot mp4 mode requires output.path.") + args = argparse.Namespace( + preset_id=preset_id, + device=str(config.device), + prompt=scenario.get("prompt"), + prompt_path=_optional_path(scenario.get("prompt_path")), + image_path=_optional_path(scenario.get("image_path")), + pose_path=_optional_path(scenario.get("pose_path")), + intrinsic_path=_optional_path(scenario.get("intrinsic_path")), + example_data=scenario.get("example_data"), + example_idx=_as_int( + scenario.get("example_idx", getattr(config, "example_idx", 0)) + ), + total_blocks=_as_int( + scenario.get("total_blocks", getattr(config, "total_blocks", 20)) + ), + pixel_height=_as_int( + scenario.get( + "pixel_height", + getattr(config, "pixel_height", DEFAULT_PIXEL_HEIGHT), + ) + ), + pixel_width=_as_int( + scenario.get( + "pixel_width", getattr(config, "pixel_width", DEFAULT_PIXEL_WIDTH) + ) + ), + fps=_as_int( + output.get( + "fps", scenario.get("fps", getattr(config, "fps", DEFAULT_FPS)) + ) + ), + output_mode="mp4", + output=Path(cast(Any, output_path)), + ) + spec = _replay_spec(args) + return run_replay_demo(spec=spec, adapter=LingbotDemoAdapter()) + if mode != "webrtc": + raise ValueError(f"Unsupported LingBot launch mode: {mode!r}.") + + context = initialize_cuda_distributed(default_device=str(config.device)) + example_idx = _as_int( + scenario.get("example_idx", getattr(config, "example_idx", 0)) + ) + ensure_example_data_downloaded( + is_rank_zero=(context.world_rank == 0), + example_idx=example_idx, + ) + args = argparse.Namespace( + preset_id=preset_id, + device=str(context.device), + seed=_as_int(output.get("seed", 42)), + compile=_runner_compile(config), + fps=_as_int(output.get("fps", getattr(config, "fps", DEFAULT_FPS))), + video_height=_as_int( + output.get( + "video_height", getattr(config, "pixel_height", DEFAULT_PIXEL_HEIGHT) + ) + ), + video_width=_as_int( + output.get( + "video_width", getattr(config, "pixel_width", DEFAULT_PIXEL_WIDTH) + ) + ), + warmup_chunks=_as_int(output.get("warmup_chunks", 10)), + warmup_timeout_s=_as_float(output.get("warmup_timeout_s", 600.0)), + client_liveness_timeout_s=_as_float( + output.get("client_liveness_timeout_s", 30.0) + ), + prefer_sw_encoder=bool(output.get("prefer_sw_encoder", prefer_sw_encoder)), + example_idx=example_idx, + host=str(host or output.get("host", "0.0.0.0")), + port=_as_int(port if port is not None else output.get("port", 8080)), + ) + from .webrtc import serve_lingbot_webrtc_demo + + return serve_lingbot_webrtc_demo( + spec=_webrtc_spec( + args, + device=str(context.device), + context_parallel_size=context.world_size, + ), + world_rank=context.world_rank, + ) + + +def _optional_path(value: object) -> Path | None: + return None if value is None else Path(cast(Any, value)) + + +def _as_int(value: object) -> int: + return int(cast(Any, value)) + + +def _as_float(value: object) -> float: + return float(cast(Any, value)) + + +def _runner_compile(config: RunnerConfig) -> bool: + transformer = getattr( + getattr(config.pipeline, "diffusion_model", None), + "transformer", + None, + ) + return bool(getattr(transformer, "compile_network", True)) + + def _replay_spec(args: argparse.Namespace) -> DemoSpec: scenario: dict[str, object] = { "example_data": args.example_data, @@ -244,7 +373,3 @@ def _webrtc_spec( }, ), ) - - -if __name__ == "__main__": - main() diff --git a/integrations/lingbot/lingbot/launch.py b/integrations/lingbot/lingbot/launch.py new file mode 100644 index 000000000..7de74e0f0 --- /dev/null +++ b/integrations/lingbot/lingbot/launch.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""LingBot launch capability for ``flashdreams-run``.""" + +from __future__ import annotations + +from collections.abc import Mapping +from functools import partial +from typing import Literal, cast + +from flashdreams.infra.runner import RunnerConfig +from flashdreams.serving.launch import ( + LaunchMode, + LaunchOptions, + ResolvedLaunch, +) + +_REPLAY_SCENARIO_FIELDS = frozenset( + { + "prompt", + "prompt_path", + "image_path", + "pose_path", + "intrinsic_path", + "example_data", + "example_idx", + "total_blocks", + "pixel_height", + "pixel_width", + "fps", + } +) +_WEBRTC_SCENARIO_FIELDS = frozenset({"example_idx"}) +_WEBRTC_OUTPUT_FIELDS = frozenset( + { + "host", + "port", + "seed", + "fps", + "video_height", + "video_width", + "warmup_chunks", + "warmup_timeout_s", + "client_liveness_timeout_s", + "prefer_sw_encoder", + } +) + + +class LingbotLaunchCapability: + """Construct LingBot replay and WebRTC launches directly.""" + + def supported_modes( + self, + config: RunnerConfig, + options: LaunchOptions, + ) -> tuple[LaunchMode, ...]: + del config, options + return ("mp4", "webrtc") + + def resolve( + self, + config: RunnerConfig, + *, + mode: LaunchMode, + options: LaunchOptions, + ) -> ResolvedLaunch | None: + if mode == "mp4": + _validate_fields("scenario", options.scenario, _REPLAY_SCENARIO_FIELDS) + _validate_fields("output", options.output, {"path", "output", "fps"}) + output_path = options.output.get("path") or options.output.get("output") + if output_path is None: + raise ValueError( + "LingBot mp4 mode requires output.path in the manifest." + ) + return _resolved(config, mode, options, output_path=output_path) + if mode == "webrtc": + _validate_fields("scenario", options.scenario, _WEBRTC_SCENARIO_FIELDS) + _validate_fields("output", options.output, _WEBRTC_OUTPUT_FIELDS) + return _resolved(config, mode, options) + return None + + +def _resolved( + config: RunnerConfig, + mode: LaunchMode, + options: LaunchOptions, + *, + output_path: object | None = None, +) -> ResolvedLaunch: + summary: dict[str, object] = { + "runner": config.runner_name, + "mode": mode, + "device": config.device, + } + if output_path is not None: + summary["output_path"] = output_path + if mode == "webrtc": + summary["host"] = options.host or options.output.get("host", "0.0.0.0") + summary["port"] = ( + options.port + if options.port is not None + else options.output.get("port", 8080) + ) + return ResolvedLaunch( + mode=mode, + label=f"LingBot {'MP4 replay' if mode == 'mp4' else 'WebRTC server'}", + summary=summary, + launch=partial( + _launch, + config=config, + mode=mode, + options=options, + ), + ) + + +def _launch( + *, + config: RunnerConfig, + mode: LaunchMode, + options: LaunchOptions, +) -> object: + from lingbot.demo.app import launch_from_runner + + if mode not in {"mp4", "webrtc"}: + raise ValueError(f"Unsupported LingBot launch mode: {mode!r}.") + return launch_from_runner( + config=config, + mode=cast(Literal["mp4", "webrtc"], mode), + scenario=dict(options.scenario), + output=dict(options.output), + host=options.host, + port=options.port, + prefer_sw_encoder=options.prefer_sw_encoder, + ) + + +def _validate_fields( + section: str, + values: Mapping[str, object], + allowed: set[str] | frozenset[str], +) -> None: + unknown = sorted(set(values) - allowed) + if unknown: + raise ValueError(f"Unsupported LingBot {section} fields: {', '.join(unknown)}.") + + +LAUNCH_CAPABILITY = LingbotLaunchCapability() + +__all__ = ["LAUNCH_CAPABILITY", "LingbotLaunchCapability"] diff --git a/integrations/lingbot/lingbot/output_targets.py b/integrations/lingbot/lingbot/output_targets.py deleted file mode 100644 index 86e77b69f..000000000 --- a/integrations/lingbot/lingbot/output_targets.py +++ /dev/null @@ -1,82 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Lingbot output capabilities for ``flashdreams-run``.""" - -from __future__ import annotations - -from typing import Any - -from flashdreams.infra.runner import RunnerConfig -from flashdreams.serving.output_targets import ( - OutputLaunchOptions, - OutputMode, - OutputTargetSpec, -) - - -class LingbotOutputTargetAdapter: - def supported_modes( - self, - config: RunnerConfig, - options: OutputLaunchOptions, - ) -> tuple[OutputMode, ...]: - del config, options - return ("webrtc",) - - def resolve( - self, - config: RunnerConfig, - *, - mode: OutputMode, - options: OutputLaunchOptions, - ) -> OutputTargetSpec | None: - if mode != "webrtc": - return None - argv = [ - "webrtc", - "--preset-id", - _pipeline_name(config), - "--device", - str(config.device), - "--fps", - str(getattr(config, "fps", 16)), - "--video-height", - str(getattr(config, "pixel_height", 464)), - "--video-width", - str(getattr(config, "pixel_width", 832)), - ] - if _compile_network(config) is False: - argv.append("--no-compile") - example_idx = getattr(config, "example_idx", None) - if example_idx is not None: - argv.extend(("--example-idx", str(example_idx))) - 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 shared demo WebRTC server", - module="lingbot.demo.app", - argv=tuple(argv), - ) - - -def _pipeline_name(config: RunnerConfig) -> str: - name = getattr(config.pipeline, "name", None) - return str(name or config.runner_name) - - -def _compile_network(config: RunnerConfig) -> bool | None: - diffusion_model = getattr(config.pipeline, "diffusion_model", None) - transformer: Any = getattr(diffusion_model, "transformer", None) - value = getattr(transformer, "compile_network", None) - return None if value is None else bool(value) - - -OUTPUT_TARGET_ADAPTER = LingbotOutputTargetAdapter() - -__all__ = ["OUTPUT_TARGET_ADAPTER", "LingbotOutputTargetAdapter"] diff --git a/integrations/lingbot/lingbot/runner.py b/integrations/lingbot/lingbot/runner.py index 9a4ffe0f0..77462351d 100644 --- a/integrations/lingbot/lingbot/runner.py +++ b/integrations/lingbot/lingbot/runner.py @@ -19,7 +19,9 @@ from dataclasses import dataclass, field from pathlib import Path +from typing import Annotated +import tyro from loguru import logger from flashdreams.infra.postprocess import VideoTensorLayout @@ -68,7 +70,9 @@ class LingbotWorldRunnerConfig(RunnerConfig): _target: type["LingbotWorldRunner"] = field( default_factory=lambda: LingbotWorldRunner ) - output_adapter: str | None = "lingbot.output_targets:OUTPUT_TARGET_ADAPTER" + launch_capability: Annotated[str | None, tyro.conf.Suppress] = ( + "lingbot.launch:LAUNCH_CAPABILITY" + ) prompt: str = "" """Text prompt. A non-empty value wins; otherwise the runner reads diff --git a/integrations/lingbot/lingbot/webrtc/server.py b/integrations/lingbot/lingbot/webrtc/server.py index 29ce67fc8..a880b7801 100644 --- a/integrations/lingbot/lingbot/webrtc/server.py +++ b/integrations/lingbot/lingbot/webrtc/server.py @@ -480,7 +480,3 @@ def main() -> None: host=args.host, port=args.port, ) - - -if __name__ == "__main__": - main() diff --git a/integrations/lingbot/pyproject.toml b/integrations/lingbot/pyproject.toml index f0175fb35..db2475732 100644 --- a/integrations/lingbot/pyproject.toml +++ b/integrations/lingbot/pyproject.toml @@ -41,9 +41,6 @@ dev = [ "pytest-asyncio>=0.23", ] -[project.scripts] -lingbot-demo = "lingbot.demo.app: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/scripts/launch_webrtc.sh b/integrations/lingbot/scripts/launch_webrtc.sh deleted file mode 100755 index 391919c73..000000000 --- a/integrations/lingbot/scripts/launch_webrtc.sh +++ /dev/null @@ -1,28 +0,0 @@ -#!/usr/bin/env bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -set -euo pipefail - -HOST="${HOST:-0.0.0.0}" -PORT="${PORT:-8092}" -DEVICE="${DEVICE:-cuda:0}" -CONFIG_NAME="${CONFIG_NAME:-lingbot-world-v2-14b-causal-fast-taehv-window15-sink3}" -WARMUP_CHUNKS="${WARMUP_CHUNKS:-0}" -FPS="${FPS:-16}" -EXAMPLE_IDX="${EXAMPLE_IDX:-0}" -VIDEO_HEIGHT="${VIDEO_HEIGHT:-352}" -VIDEO_WIDTH="${VIDEO_WIDTH:-640}" - -cd "$(dirname "${BASH_SOURCE[0]}")/../../.." - -exec uv run --no-sync python -m lingbot.webrtc.server \ - --host "${HOST}" \ - --port "${PORT}" \ - --config_name "${CONFIG_NAME}" \ - --device "${DEVICE}" \ - --warmup_chunks "${WARMUP_CHUNKS}" \ - --fps "${FPS}" \ - --video-height "${VIDEO_HEIGHT}" \ - --video-width "${VIDEO_WIDTH}" \ - --example-idx "${EXAMPLE_IDX}" diff --git a/integrations/lingbot/tests/test_demo_api.py b/integrations/lingbot/tests/test_demo_api.py index 98444a7c5..d81b7a170 100644 --- a/integrations/lingbot/tests/test_demo_api.py +++ b/integrations/lingbot/tests/test_demo_api.py @@ -7,10 +7,12 @@ from pathlib import Path from typing import Any +import lingbot.demo.app as demo_app_module import numpy as np import pytest import torch from aiohttp import web +from lingbot.config import RUNNER_LINGBOT_WORLD_FAST from lingbot.demo import ( DEFAULT_LINGBOT_PRESET, LINGBOT_MODEL_ID, @@ -114,6 +116,32 @@ def test_lingbot_demo_defaults_to_interactive_preset() -> None: assert args.preset_id == "lingbot-world-fast-taehv-window15-sink3" +def test_lingbot_direct_runner_launch_builds_mp4_spec( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: list[DemoSpec] = [] + + def fake_run_replay_demo(*, spec: DemoSpec, adapter: object) -> str: + del adapter + captured.append(spec) + return "completed" + + monkeypatch.setattr(demo_app_module, "run_replay_demo", fake_run_replay_demo) + + result = demo_app_module.launch_from_runner( + config=RUNNER_LINGBOT_WORLD_FAST, + mode="mp4", + scenario={"example_idx": 2, "total_blocks": 3}, + output={"path": tmp_path / "demo.mp4", "fps": 12}, + ) + + assert result == "completed" + assert captured[0].preset_id == RUNNER_LINGBOT_WORLD_FAST.runner_name + assert isinstance(captured[0].output, Mp4OutputSpec) + assert captured[0].output.path == tmp_path / "demo.mp4" + + def test_lingbot_demo_adapter_declares_shared_demo_modes() -> None: adapter = LingbotDemoAdapter() diff --git a/integrations/omnidreams/README.md b/integrations/omnidreams/README.md index 47bc78f27..64f61319a 100644 --- a/integrations/omnidreams/README.md +++ b/integrations/omnidreams/README.md @@ -121,9 +121,11 @@ Interpret the report as follows: baseline JSON; the baseline should contain only expected metric values and tolerances. -## Run interactive-drive (desktop demo) -The `omnidreams.interactive_drive` subpackage ships a single-process -driving demo. Refer to [README for full guide](omnidreams/interactive_drive/README.md) +## Run the local-window desktop demo + +The `local-window` mode launches the single-process driving demo. Refer to the +[interactive guide](omnidreams/interactive_drive/README.md) for controls and +runtime requirements. Example execution below: ```bash @@ -138,11 +140,10 @@ uv run --package flashdreams-omnidreams python integrations/omnidreams/omnidream # Prepare to run tuned for performance uv run --package flashdreams-omnidreams omnidreams-prepare --perf -# Run demo -uv run --package flashdreams-omnidreams interactive-drive \ - --manifest example_world_model_perf.yaml - -# add `--stream-mjpeg :8080` to stream to your browser +# Run the centralized local-window launch +uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams-perf local-window \ + --manifest configs/launch_manifest/omnidreams_local_window.yaml ``` ## Native DiT defaults @@ -165,7 +166,9 @@ Sparge/SageAttention-3 hybrid schedule when the extension and GPU support it. From the workspace root, run: ```bash -uv run --package flashdreams-omnidreams torchrun --nproc_per_node 1 -m omnidreams.webrtc.server --pipeline_config_name omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf --scene-uuid 0d404ff7-2b66-498c-b047-1ed8cded60d4 --port 8089 +uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams webrtc \ + --manifest configs/launch_manifest/omnidreams_webrtc.yaml ``` When `--scene_dir` is omitted, the server downloads the selected scene from the @@ -177,10 +180,11 @@ Weather variants ship as sibling archives; pass `--scene-variant rain` (or from the scene's first ground-truth camera frame (`clipgt/frames//.jpeg`, falling back to `clipgt/first_image.*`) and the weather-matched `clipgt/prompt.txt` (falling back to `clipgt/prompt.txt`). -Pass `--scene_dir ` to use a pre-staged local scene instead. +Set `scenario.scene_dir` in the launch manifest to use a pre-staged local +scene instead. -To enable video post-processing by default, pass a registered preset such as -`--postprocess-preset rtx-super-resolution`. RTX postprocess presets require the +To enable video post-processing by default, override the runner's registered +post-process preset in the launch manifest. RTX postprocess presets require the optional NVIDIA VFX runtime: ```bash diff --git a/integrations/omnidreams/omnidreams/config.py b/integrations/omnidreams/omnidreams/config.py index 9cfac26ab..9031b4bca 100644 --- a/integrations/omnidreams/omnidreams/config.py +++ b/integrations/omnidreams/omnidreams/config.py @@ -430,7 +430,7 @@ def _lightvae_fp8_state_path() -> str | None: """All shipped Omnidreams variants, keyed by ``name``.""" -## Per-variant runner-config literals (slug == ``name``). +## Per-variant runner-config literals. _DEFAULT_PROMPT_1V = ( "Driving scene from a front-facing car camera. Urban environment with roads, " @@ -450,14 +450,14 @@ def _lightvae_fp8_state_path() -> str | None: ) RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE = OmnidreamsRunnerConfig( - runner_name=SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE.name, + runner_name="omnidreams", description="Single-view 2-step distilled chunk2 (LightVAE + LightTAE).", pipeline=SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE, prompt=_DEFAULT_PROMPT_1V, ) RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF = OmnidreamsRunnerConfig( - runner_name=SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF.name, + runner_name="omnidreams-perf", description=( "Single-view chunk2 perf preset (compile + CUDA graphs across all stages)." ), diff --git a/integrations/omnidreams/omnidreams/demo/README.md b/integrations/omnidreams/omnidreams/demo/README.md index 9dab46ed8..da59c7ce8 100644 --- a/integrations/omnidreams/omnidreams/demo/README.md +++ b/integrations/omnidreams/omnidreams/demo/README.md @@ -29,25 +29,27 @@ uv sync --python 3.12 --package flashdreams-omnidreams --no-dev Run a short replay without writing video output: ```bash -uv run --python 3.12 --package flashdreams-omnidreams omnidreams-demo replay \ - --output-mode null \ - --device cuda:0 \ - --total-blocks 10 +uv run --python 3.12 --package flashdreams-omnidreams flashdreams-run \ + omnidreams null \ + --manifest configs/launch_manifest/omnidreams_null.yaml ``` ## Precomputed MP4 Replay -Generate an MP4 from bundled single-view sample data and pre-rendered HDMaps: +Generate an MP4 from bundled single-view sample data and pre-rendered HDMaps +with no manifest: ```bash -mkdir -p outputs -uv run --python 3.12 --package flashdreams-omnidreams omnidreams-demo replay \ - --device cuda:0 \ - --example-data \ - --example-data-uuid 239560dc-33d1-11ef-9720-00044bcbccac \ - --total-blocks 225 \ - --fps 30 \ - --output outputs/omnidreams-demo-precomputed-1min.mp4 +uv run flashdreams-run omnidreams mp4 +``` + +The default output is `outputs/omnidreams.mp4`. To override the sample, +rollout length, frame rate, or output path, use the versioned manifest: + +```bash +uv run --python 3.12 --package flashdreams-omnidreams flashdreams-run \ + omnidreams mp4 \ + --manifest configs/launch_manifest/omnidreams_mp4.yaml ``` This replay path mirrors the benchmark runner path: it uses a prompt, first @@ -55,26 +57,21 @@ frame, and pre-rendered HDMap video. It does not load a Ludus scene or render HDMaps at runtime. The demo defaults to the stable non-perf OmniDreams preset used by the benchmark path. -Pass `--example-data-uuid ` to select another bundled single-view sample, -or `--no-example-data` to require explicit asset paths. +Set `scenario.example_data_uuid` to select another bundled single-view sample, +or set `scenario.example_data: false` and provide explicit asset paths. ## Ludus MP4 Replay Generate an MP4 by rendering HDMap conditioning from a recorded keyboard trace: ```bash -uv run --python 3.12 --package flashdreams-omnidreams omnidreams-demo replay \ - --conditioning-mode ludus-scene-driving \ - --keyboard-trace integrations/omnidreams/omnidreams/demo/traces/ludus_forward_sweep_60s.json \ - --device cuda:0 \ - --scene-uuid 0d404ff7-2b66-498c-b047-1ed8cded60d4 \ - --seed 42 \ - --total-blocks 226 \ - --output outputs/omnidreams-demo--ludus-1min.mp4 +uv run --python 3.12 --package flashdreams-omnidreams flashdreams-run \ + omnidreams mp4 \ + --manifest path/to/omnidreams-ludus-mp4.yaml ``` -The `omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf` preset remains an -explicit `--preset-id` opt-in. It should become the default only after the +The `omnidreams-perf` preset remains an +explicit runner-slug opt-in. It should become the default only after the compile/cache behavior is reliable enough for the demo path. ## WebRTC @@ -84,14 +81,12 @@ The small model adapter in this package loads one scene, renders HDMap conditioning with Ludus, and runs OmniDreams from browser WASD controls: ```bash -uv run --python 3.12 --package flashdreams-omnidreams omnidreams-demo webrtc \ - --host 0.0.0.0 \ - --port 8089 \ - --device cuda:0 \ - --scene-uuid 0d404ff7-2b66-498c-b047-1ed8cded60d4 +uv run --python 3.12 --package flashdreams-omnidreams flashdreams-run \ + omnidreams webrtc \ + --manifest configs/launch_manifest/omnidreams_webrtc.yaml ``` The scene UUID is optional; when omitted, the runtime uses the default -Hugging Face WebRTC scene. Override the scene with `--scene-uuid`, select a -weather variant with `--scene-variant default|rain|snow`, or use -`--scene-dir /path/to/local/scene` for a local staged scene. +Hugging Face WebRTC scene. Override ``scenario.scene_uuid``, select a weather +variant with ``scenario.scene_variant``, or set ``scenario.scene_dir`` for a +local staged scene. diff --git a/integrations/omnidreams/omnidreams/demo/adapter.py b/integrations/omnidreams/omnidreams/demo/adapter.py index 4dc7d1c87..e89ad5d4d 100644 --- a/integrations/omnidreams/omnidreams/demo/adapter.py +++ b/integrations/omnidreams/omnidreams/demo/adapter.py @@ -294,7 +294,15 @@ def _pipeline_config(self, config: InferenceConfig) -> Any: ) from exc def _default_replay_prompt(self, config: InferenceConfig | None) -> str: - runner = OMNIDREAMS_RUNNERS.get(self._preset_id(config)) + preset_id = self._preset_id(config) + runner = next( + ( + value + for value in OMNIDREAMS_RUNNERS.values() + if value.pipeline.name == preset_id + ), + None, + ) return "" if runner is None else str(getattr(runner, "prompt", "")) diff --git a/integrations/omnidreams/omnidreams/demo/app.py b/integrations/omnidreams/omnidreams/demo/app.py index ac1641917..5aa8bc1ea 100644 --- a/integrations/omnidreams/omnidreams/demo/app.py +++ b/integrations/omnidreams/omnidreams/demo/app.py @@ -7,11 +7,13 @@ import argparse import math +from collections.abc import Iterable from pathlib import Path -from typing import Any +from typing import Any, Literal, cast from omnidreams.runner import DEFAULT_EXAMPLE_DATA_UUID_1V +from flashdreams.infra.runner import RunnerConfig from flashdreams.runtime import InferenceConfig from flashdreams.runtime.demo import ( DemoSpec, @@ -20,6 +22,11 @@ WebRTCOutputSpec, ) from flashdreams.runtime.demo.app import DemoApplication +from flashdreams.runtime.demo.replay import run_replay_demo +from flashdreams.serving.webrtc.bootstrap import ( + configure_logging, + initialize_cuda_distributed, +) from .adapter import OmnidreamsDemoAdapter from .spec import ( @@ -146,6 +153,130 @@ def main(argv: list[str] | None = None) -> None: _APPLICATION.main(argv) +def launch_from_runner( + *, + config: RunnerConfig, + mode: Literal["mp4", "null", "webrtc"], + scenario: dict[str, object], + output: dict[str, object], + host: str | None = None, + port: int | None = None, + prefer_sw_encoder: bool = False, +) -> object: + """Launch an OmniDreams demo directly from a resolved runner config.""" + configure_logging() + preset_id = str(getattr(config.pipeline, "name", config.runner_name)) + seed = _runner_seed(config) + if mode in {"mp4", "null"}: + output_path = output.get("path") or output.get("output") + if mode == "mp4" and output_path is None: + raise ValueError("OmniDreams mp4 mode requires output.path.") + args = argparse.Namespace( + preset_id=preset_id, + device=str(config.device), + seed=seed, + conditioning_mode=scenario.get( + "conditioning_mode", OMNIDREAMS_CONDITIONING_PRECOMPUTED + ), + prompt=scenario.get("prompt"), + hdmap_video_paths=_as_tuple(scenario.get("hdmap_video_paths", ())), + first_frame_paths=_as_tuple(scenario.get("first_frame_paths", ())), + camera_names=_as_tuple(scenario.get("camera_names", ())), + keyboard_trace=_optional_path(scenario.get("keyboard_trace")), + scene_path=_optional_path(scenario.get("scene_path")), + scene_dir=_optional_path(scenario.get("scene_dir")), + scene_uuid=scenario.get("scene_uuid", DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID), + scene_variant=str(scenario.get("scene_variant", "default")), + camera_name=str(scenario.get("camera_name", "camera_front_wide_120fov")), + move_speed_per_s=_as_float(scenario.get("move_speed_per_s", 6.0)), + rotate_speed_rad_per_s=_as_float( + scenario.get("rotate_speed_rad_per_s", math.radians(35.0)) + ), + ludus_backend=str(scenario.get("ludus_backend", "cuda")), + example_data=scenario.get("example_data"), + example_data_uuid=scenario.get( + "example_data_uuid", DEFAULT_EXAMPLE_DATA_UUID_1V + ), + total_blocks=_as_int( + scenario.get("total_blocks", getattr(config, "total_blocks", 60)) + ), + pixel_height=_as_int( + scenario.get("pixel_height", getattr(config, "pixel_height", 704)) + ), + pixel_width=_as_int( + scenario.get("pixel_width", getattr(config, "pixel_width", 1280)) + ), + fps=_as_int( + output.get( + "fps", scenario.get("fps", getattr(config, "output_fps", 30)) + ) + ), + output_mode=mode, + output=None if output_path is None else Path(cast(Any, output_path)), + ) + spec = _replay_spec(args) + return run_replay_demo(spec=spec, adapter=OmnidreamsDemoAdapter()) + if mode != "webrtc": + raise ValueError(f"Unsupported OmniDreams launch mode: {mode!r}.") + + context = initialize_cuda_distributed(default_device=str(config.device)) + args = argparse.Namespace( + preset_id=preset_id, + device=str(context.device), + seed=seed, + scene_dir=_optional_path(scenario.get("scene_dir")), + scene_uuid=scenario.get("scene_uuid"), + scene_variant=str(scenario.get("scene_variant", "default")), + camera_name=str(scenario.get("camera_name", "camera_front_wide_120fov")), + fps=_as_int(output.get("fps", getattr(config, "output_fps", 30))), + video_height=_as_int( + output.get("video_height", getattr(config, "pixel_height", 704)) + ), + video_width=_as_int( + output.get("video_width", getattr(config, "pixel_width", 1280)) + ), + warmup_chunks=_as_int(output.get("warmup_chunks", 10)), + warmup_timeout_s=_as_float(output.get("warmup_timeout_s", 600.0)), + client_liveness_timeout_s=_as_float( + output.get("client_liveness_timeout_s", 10.0) + ), + debug_serve_hdmaps=bool(output.get("debug_serve_hdmaps", False)), + prefer_sw_encoder=bool(output.get("prefer_sw_encoder", prefer_sw_encoder)), + host=str(host or output.get("host", "0.0.0.0")), + port=_as_int(port if port is not None else output.get("port", 8082)), + ) + from .webrtc import serve_omnidreams_webrtc_demo + + return serve_omnidreams_webrtc_demo( + spec=_webrtc_spec(args, device=str(context.device)), + world_rank=context.world_rank, + ) + + +def _optional_path(value: object) -> Path | None: + return None if value is None else Path(cast(Any, value)) + + +def _as_int(value: object) -> int: + return int(cast(Any, value)) + + +def _as_float(value: object) -> float: + return float(cast(Any, value)) + + +def _as_tuple(value: object) -> tuple[object, ...]: + if isinstance(value, (str, bytes)) or not isinstance(value, Iterable): + raise TypeError("Expected a sequence value in the launch manifest.") + return tuple(value) + + +def _runner_seed(config: RunnerConfig) -> int: + diffusion_model = getattr(config.pipeline, "diffusion_model", None) + seed = getattr(diffusion_model, "seed", 42) + return 42 if seed is None else int(seed) + + def _replay_spec(args: argparse.Namespace) -> DemoSpec: scenario: dict[str, object] = { "conditioning_mode": args.conditioning_mode, @@ -247,7 +378,3 @@ def _split_paths(value: str) -> tuple[Path, ...]: def _split_strings(value: str) -> tuple[str, ...]: return tuple(part for part in value.split(",") if part) - - -if __name__ == "__main__": - main() diff --git a/integrations/omnidreams/omnidreams/interactive_drive/README.md b/integrations/omnidreams/omnidreams/interactive_drive/README.md index ff06081b8..5301030e0 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/README.md +++ b/integrations/omnidreams/omnidreams/interactive_drive/README.md @@ -182,7 +182,7 @@ Common `omnidreams-prepare` flags: - `--skip-text-encoder` — skip the ~14 GB text-encoder prewarm when you're using a precomputed prompt embedding or want a lighter first-time setup. - `--skip-scene` — don't stage any scene (for when you're supplying your - own USDZ via `interactive-drive --scene`). + own USDZ via ``scenario.scene`` in a local-window launch manifest). If `omnidreams-prepare` fails with `401`, `403`, or a gated-repo error, verify `HF_TOKEN` and confirm you have requested (and been granted) @@ -228,18 +228,19 @@ this subpackage (via `__file__`), so you don't have to pass long `integrations/omnidreams/omnidreams/interactive_drive/...` paths unless you want to override them. -There is one entry point — `interactive-drive` — and three modes selected by -flags: +There is one public entry point — `flashdreams-run` — with ``local-window`` +and ``webrtc`` modes. Local-window presentation variants are selected in the +launch manifest: | Mode | When to use | How | |---|---|---| -| **HUD (default)** | You have a graphical desktop session and want the full demo: scene/variant selector, steering wheel + pedals overlay, BEV minimap, keyboard *and* wheel input. | `interactive-drive ...` | -| **Bare backend, local window** | You want the lightweight setup: a single Vulkan window showing the world-model output, no HUD chrome. | `interactive-drive --no-hud ...` | -| **Bare backend, browser** | The demo machine has no graphics-capable GPU (e.g. compute-only GB300) or you want to view from a laptop browser while the model runs elsewhere. Implies `--no-hud`. | `interactive-drive --stream-mjpeg [HOST:]PORT ...` | +| **HUD (default)** | You have a graphical desktop session and want the full demo: scene/variant selector, steering wheel + pedals overlay, BEV minimap, keyboard *and* wheel input. | ``local-window`` with default output settings | +| **Bare backend, local window** | You want the lightweight setup: a single Vulkan window showing the world-model output, no HUD chrome. | ``local-window`` with ``output.no_hud: true`` | +| **Bare backend, browser** | The demo machine has no graphics-capable GPU (e.g. compute-only GB300) or you want to view from a laptop browser while the model runs elsewhere. Implies no HUD. | ``local-window`` with ``output.stream_mjpeg: :8080`` | For a richer remote-viewing experience with a polished frontend and lower -latency than an in-process MJPEG stream, prefer the separate -`omnidreams.webrtc.server` entry point (see +latency than an in-process MJPEG stream, prefer the centralized ``webrtc`` +mode (see [`integrations/omnidreams/README.md`](../../README.md)). The HUD itself uses pygame/SDL2 for rendering, which keeps the demo responsive @@ -251,7 +252,9 @@ process. ### HUD mode (default) ```bash -uv run --package flashdreams-omnidreams interactive-drive +uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams-perf local-window \ + --manifest configs/launch_manifest/omnidreams_local_window.yaml ``` The default `--scene` resolves to @@ -368,7 +371,9 @@ This is the lighter-weight path that matches the older standalone output, no HUD chrome, no scene selector. ```bash -uv run --package flashdreams-omnidreams interactive-drive --no-hud +uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams-perf local-window \ + --manifest path/to/local-window-with-no-hud.yaml ``` You should initially see the generated driving view. Press `2` to switch to the @@ -383,11 +388,12 @@ network, or when you want to demo from a laptop browser while the model runs elsewhere. Implies `--no-hud` because the user is then viewing through a browser, not a local Vulkan window — the slangpy HUD itself is a Vulkan presenter, so it can't run on the same hosts that need -`--stream-mjpeg`. +``output.stream_mjpeg`` in the launch manifest. ```bash -uv run --package flashdreams-omnidreams interactive-drive \ - --stream-mjpeg 8080 +uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams-perf local-window \ + --manifest path/to/local-window-mjpeg.yaml ``` Open `http://:8080/` in a browser on the same network; keyboard @@ -434,10 +440,10 @@ ssh -L 8080:localhost:8080 @ Then open `http://localhost:8080/`. -For a richer browser frontend with lower latency, prefer the separate -`omnidreams.webrtc.server` entry point. +For a richer browser frontend with lower latency, prefer the centralized +``webrtc`` launch mode. -#### Fully headless: `--stream-mjpeg` with `--auto-start` +#### Fully headless MJPEG with auto-start By default the streaming mode waits for the browser scene picker before it loads anything, so a freshly launched server idles on "Select a scene to @@ -447,8 +453,9 @@ add `--auto-start`. It skips the scene selection and immediately loads `--scene` or not yet staged): ```bash -uv run --package flashdreams-omnidreams interactive-drive \ - --stream-mjpeg 8080 --auto-start --scene +uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams-perf local-window \ + --manifest path/to/headless-auto-start.yaml ``` The interactive-drive CUDA fast path is enabled by default. HDMap raster frames @@ -457,7 +464,9 @@ SlangPy CUDA interop for generated RGB frames when the model output is still on CUDA: ```bash -uv run --no-sync --package flashdreams-omnidreams interactive-drive +uv run --no-sync --package flashdreams-omnidreams flashdreams-run \ + omnidreams-perf local-window \ + --manifest configs/launch_manifest/omnidreams_local_window.yaml ``` Set `INTERACTIVE_DRIVE_DISABLE_CUDA_INTEROP=1` to force the conservative host @@ -492,7 +501,9 @@ input-to-present timing while the demo runs: ```bash INTERACTIVE_DRIVE_PROFILE_INPUT_TO_PRESENT=1 \ - uv run --no-sync --package flashdreams-omnidreams interactive-drive --auto-start + uv run --no-sync --package flashdreams-omnidreams flashdreams-run \ + omnidreams-perf local-window \ + --manifest path/to/auto-start-local-window.yaml ``` The log line is `[profile] e2e ...`. `wall_present_fps` counts only frames @@ -578,9 +589,9 @@ feeds it to the same loader the regular flow uses: ```bash uv run --package flashdreams-omnidreams omnidreams-prepare --skip-scene -uv run --package flashdreams-omnidreams interactive-drive \ - --synthetic-scene \ - --synthetic-initial-rgb path/to/forward_facing_road_photo.jpg +uv run --package flashdreams-omnidreams flashdreams-run \ + omnidreams-perf local-window \ + --manifest path/to/synthetic-scene-local-window.yaml ``` The world model is trained on natural driving frames, so passing your own diff --git a/integrations/omnidreams/omnidreams/interactive_drive/__main__.py b/integrations/omnidreams/omnidreams/interactive_drive/__main__.py deleted file mode 100644 index eed736eb1..000000000 --- a/integrations/omnidreams/omnidreams/interactive_drive/__main__.py +++ /dev/null @@ -1,13 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - -# ``python -m omnidreams.interactive_drive`` and the ``interactive-drive`` console -# script both go through the demo wrapper so the same flags work in both -# the supervised HUD path and the bare backend path. The HUD is on by -# default; pass ``--no-hud`` to bypass it and fall through to the bare -# slangpy Vulkan window. Browser / remote streaming use cases live in -# the separate ``omnidreams.webrtc.server`` entry point. -from omnidreams.interactive_drive.demo import main - -if __name__ == "__main__": - main() diff --git a/integrations/omnidreams/omnidreams/interactive_drive/app.py b/integrations/omnidreams/omnidreams/interactive_drive/app.py index ac4cf9dd5..c8ca3fd0d 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/app.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/app.py @@ -81,7 +81,7 @@ def __init__( Vulkan window) or, when ``config.stream_mjpeg_bind`` is set, an :class:`MJPEGStreamingPresenter` that serves frames over HTTP with no GPU-graphics dependency. Browser viewers with a richer frontend - are served by ``omnidreams.webrtc.server`` instead. + are served by the centralized ``webrtc`` launch mode instead. """ self._config = config self._backend = backend @@ -575,7 +575,7 @@ def _build_presenter(config: AppConfig, keyboard: KeyboardState) -> PresenterBac :class:`SlangPyPresenter` -- a local Vulkan window. - For browser viewers with a richer frontend, ``omnidreams.webrtc.server`` + For browser viewers with a richer frontend, the centralized ``webrtc`` mode (a separate entry point) is the preferred path; this MJPEG fallback is the in-process, dependency-free alternative for headless boxes. """ diff --git a/integrations/omnidreams/omnidreams/interactive_drive/cli.py b/integrations/omnidreams/omnidreams/interactive_drive/cli.py index 5ff2456c2..e5ffcc03f 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/cli.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/cli.py @@ -230,8 +230,8 @@ def build_parser() -> argparse.ArgumentParser: "opens http://HOST:PORT/ in a browser to view the demo and " "send keyboard input. Useful on compute-only hosts (e.g. " "GB300-only DGX Station) where no Vulkan-capable GPU exists; " - "for a richer browser viewer prefer the separate " - "``omnidreams.webrtc.server`` entry point. Implies --no-hud " + "for a richer browser viewer prefer the centralized " + "``webrtc`` launch mode. Implies --no-hud " "when launched via the demo wrapper." ), ) diff --git a/integrations/omnidreams/omnidreams/interactive_drive/configs/example_world_model_synthetic.yaml b/integrations/omnidreams/omnidreams/interactive_drive/configs/example_world_model_synthetic.yaml index 56056e11c..f62b95bfc 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/configs/example_world_model_synthetic.yaml +++ b/integrations/omnidreams/omnidreams/interactive_drive/configs/example_world_model_synthetic.yaml @@ -5,7 +5,8 @@ # interactive-drive latency path runs on local random-initialized weights with # no checkpoint downloads. synthetic_model swaps only weight sources, not perf # knobs. To avoid drift, prefer running the real manifest with the flag: -# interactive-drive --manifest example_world_model_perf.yaml --synthetic-model +# Reference this file from output.world_model_manifest_path in a local-window +# FlashDreams launch manifest, and set output.synthetic_model: true. synthetic_model: true resolution_wh: [1168, 640] fps: 30 diff --git a/integrations/omnidreams/omnidreams/interactive_drive/demo.py b/integrations/omnidreams/omnidreams/interactive_drive/demo.py index c0af59990..783c53883 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/demo.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/demo.py @@ -508,7 +508,7 @@ def build_parser() -> argparse.ArgumentParser: " window entirely and serve frames to a browser as an MJPEG" " HTTP stream (useful on compute-only hosts without a Vulkan" " GPU). For a richer browser viewer use the separate" - " ``omnidreams.webrtc.server`` entry point." + " centralized ``webrtc`` launch mode." ) parser.add_argument( "--no-hud", @@ -526,7 +526,7 @@ def build_parser() -> argparse.ArgumentParser: "Directory of USDZ scenes shown in the HUD scene selector. " "Defaults to ``$FLASHDREAMS_CACHE_DIR/omnidreams-scenes/``, " "the shared cache root used by both this demo and the " - "``omnidreams.webrtc.server`` scene pipeline." + "centralized ``webrtc`` scene pipeline." ), ) parser.add_argument( @@ -692,8 +692,46 @@ def _maybe_autostage_scene(scene: Path, *, scene_dir: Path, allow_skip: bool) -> def main() -> None: + """Run the legacy parser entry point used by internal development tools.""" + _run_namespace(build_parser().parse_args()) + + +def launch_from_runner( + *, + config: object, + world_model_manifest: Path, + scenario: dict[str, object], + output: dict[str, object], +) -> None: + """Launch the local window directly from the central resolved launch.""" + args = build_parser().parse_args([]) + args.backend = "omnidreams" + args.manifest = world_model_manifest + preset = getattr(getattr(config, "postprocess", None), "preset", "") + args.postprocess_preset = output.get("postprocess_preset", preset) + for key, value in scenario.items(): + if value is not None: + setattr(args, key, _coerce_launch_path(key, value)) + for key, value in output.items(): + if ( + key not in {"world_model_manifest_path", "postprocess_preset"} + and value is not None + ): + setattr(args, key, _coerce_launch_path(key, value)) + _run_namespace(args) + + +def _coerce_launch_path(key: str, value: object) -> object: + if key.endswith(("_path", "_dir")) or key in {"scene", "wheel_device"}: + return Path(value) # type: ignore[arg-type] + if key.endswith("_axis") and isinstance(value, (list, tuple)): + return tuple(int(item) for item in value) + return value + + +def _run_namespace(args: argparse.Namespace) -> None: + """Execute one already-resolved local-window namespace.""" configure_logging() - args = build_parser().parse_args() if not args.synthetic_scene: # Only the bare ``--no-hud`` backend has no scene picker; the HUD # and MJPEG paths both let the user pick from ``--scene-dir``, so a diff --git a/integrations/omnidreams/omnidreams/interactive_drive/streaming_presenter.py b/integrations/omnidreams/omnidreams/interactive_drive/streaming_presenter.py index 41a53c8b0..60e89b7ca 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/streaming_presenter.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/streaming_presenter.py @@ -5,7 +5,7 @@ ``multipart/x-mixed-replace`` stream with keydown/keyup posted back. Dependency-free fallback for headless / compute-only hosts with no -graphics GPU; prefer ``omnidreams.webrtc.server`` for a richer viewer. +graphics GPU; prefer the centralized ``webrtc`` mode for a richer viewer. """ from __future__ import annotations diff --git a/integrations/omnidreams/omnidreams/launch.py b/integrations/omnidreams/omnidreams/launch.py new file mode 100644 index 000000000..7412e9119 --- /dev/null +++ b/integrations/omnidreams/omnidreams/launch.py @@ -0,0 +1,301 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OmniDreams launch capability for ``flashdreams-run``.""" + +from __future__ import annotations + +from collections.abc import Mapping +from functools import partial +from pathlib import Path +from typing import Any, Literal, cast + +from flashdreams.infra.runner import RunnerConfig +from flashdreams.serving.launch import ( + LaunchMode, + LaunchOptions, + ResolvedLaunch, +) + +_LOCAL_WINDOW_MANIFESTS = { + "omnidreams": "example_world_model.yaml", + "omnidreams-perf": "example_world_model_perf.yaml", +} +_DEFAULT_MP4_OUTPUT_PATH = Path("outputs/omnidreams.mp4") +_REPLAY_SCENARIO_FIELDS = frozenset( + { + "conditioning_mode", + "prompt", + "hdmap_video_paths", + "first_frame_paths", + "camera_names", + "keyboard_trace", + "scene_path", + "scene_dir", + "scene_uuid", + "scene_variant", + "camera_name", + "move_speed_per_s", + "rotate_speed_rad_per_s", + "ludus_backend", + "example_data", + "example_data_uuid", + "total_blocks", + "pixel_height", + "pixel_width", + "fps", + } +) +_WEBRTC_SCENARIO_FIELDS = frozenset( + {"scene_dir", "scene_uuid", "scene_variant", "camera_name"} +) +_WEBRTC_OUTPUT_FIELDS = frozenset( + { + "host", + "port", + "fps", + "video_height", + "video_width", + "warmup_chunks", + "warmup_timeout_s", + "client_liveness_timeout_s", + "debug_serve_hdmaps", + "prefer_sw_encoder", + } +) +_LOCAL_SCENARIO_FIELDS = frozenset( + { + "scene", + "scene_dir", + "camera", + "variant", + "prompt", + "synthetic_scene", + "synthetic_initial_rgb", + "synthetic_prompt", + "auto_start", + "preload_scenes", + "wheel_profile", + "wheel_profiles_dir", + "wheel_device", + "wheel_steering_axis", + "wheel_throttle_axis", + "wheel_brake_axis", + "wheel_pedals_inverted", + "no_wheel", + "control_assets_dir", + "official_hdmap_dir", + } +) +_LOCAL_OUTPUT_FIELDS = frozenset( + { + "world_model_manifest_path", + "no_hud", + "stream_mjpeg", + "cuda_visible_devices", + "compute_device", + "ludus_backend", + "sync_gpu_timing", + "profile_world_model", + "offload_text_encoder", + "postprocess_preset", + "hf_org", + "stop_after_chunks", + "synthetic_model", + "bev", + "bev_resolution", + "bev_height_m", + "bev_fov_deg", + "bev_tilt_deg", + "oob_warn_proximity", + "oob_respawn_proximity", + "oob_respawn_debounce_chunks", + "oob_margin_m", + "oob_warning_zone_m", + } +) + + +class OmnidreamsLaunchCapability: + """Construct OmniDreams replay, WebRTC, and local-window launches.""" + + def supported_modes( + self, + config: RunnerConfig, + options: LaunchOptions, + ) -> tuple[LaunchMode, ...]: + modes: list[LaunchMode] = ["mp4", "null"] + if _is_single_view(config): + modes.append("webrtc") + if _world_model_manifest(config, options) is not None: + modes.append("local-window") + return tuple(modes) + + def resolve( + self, + config: RunnerConfig, + *, + mode: LaunchMode, + options: LaunchOptions, + ) -> ResolvedLaunch | None: + if mode in {"mp4", "null"}: + _validate_fields("scenario", options.scenario, _REPLAY_SCENARIO_FIELDS) + _validate_fields("output", options.output, {"path", "output", "fps"}) + output_path = options.output.get("path") or options.output.get("output") + if mode == "mp4" and output_path is None: + output_path = _DEFAULT_MP4_OUTPUT_PATH + return _demo_launch(config, mode, options, output_path=output_path) + if mode == "webrtc" and _is_single_view(config): + _validate_fields("scenario", options.scenario, _WEBRTC_SCENARIO_FIELDS) + _validate_fields("output", options.output, _WEBRTC_OUTPUT_FIELDS) + return _demo_launch(config, mode, options) + if mode == "local-window": + world_manifest = _world_model_manifest(config, options) + if world_manifest is None: + return None + _validate_fields("scenario", options.scenario, _LOCAL_SCENARIO_FIELDS) + _validate_fields("output", options.output, _LOCAL_OUTPUT_FIELDS) + return _local_window_launch(config, options, world_manifest) + return None + + +def _demo_launch( + config: RunnerConfig, + mode: LaunchMode, + options: LaunchOptions, + *, + output_path: object | None = None, +) -> ResolvedLaunch: + summary: dict[str, object] = { + "runner": config.runner_name, + "mode": mode, + "device": config.device, + } + if output_path is not None: + summary["output_path"] = output_path + if mode == "webrtc": + summary["host"] = options.host or options.output.get("host", "0.0.0.0") + summary["port"] = ( + options.port + if options.port is not None + else options.output.get("port", 8082) + ) + return ResolvedLaunch( + mode=mode, + label=f"OmniDreams {mode} launch", + summary=summary, + launch=partial( + _launch_demo, + config=config, + mode=mode, + options=options, + output_path=output_path, + ), + ) + + +def _launch_demo( + *, + config: RunnerConfig, + mode: LaunchMode, + options: LaunchOptions, + output_path: object | None, +) -> object: + from omnidreams.demo.app import launch_from_runner + + if mode not in {"mp4", "null", "webrtc"}: + raise ValueError(f"Unsupported OmniDreams launch mode: {mode!r}.") + output = dict(options.output) + if output_path is not None: + output.setdefault("path", output_path) + return launch_from_runner( + config=config, + mode=cast(Literal["mp4", "null", "webrtc"], mode), + scenario=dict(options.scenario), + output=output, + host=options.host, + port=options.port, + prefer_sw_encoder=options.prefer_sw_encoder, + ) + + +def _local_window_launch( + config: RunnerConfig, + options: LaunchOptions, + world_manifest: Path, +) -> ResolvedLaunch: + return ResolvedLaunch( + mode="local-window", + label="OmniDreams local interactive window", + summary={ + "runner": config.runner_name, + "mode": "local-window", + "world_model_manifest": world_manifest, + }, + notes=( + ( + "The compatibility world-model manifest supplies interactive " + "runtime and native-acceleration settings." + ), + ), + launch=partial( + _launch_local_window, + config=config, + options=options, + world_manifest=world_manifest, + ), + ) + + +def _launch_local_window( + *, + config: RunnerConfig, + options: LaunchOptions, + world_manifest: Path, +) -> object: + from omnidreams.interactive_drive.demo import ( + launch_from_runner, + ) + + return launch_from_runner( + config=config, + world_model_manifest=world_manifest, + scenario=dict(options.scenario), + output=dict(options.output), + ) + + +def _world_model_manifest( + config: RunnerConfig, + options: LaunchOptions, +) -> Path | None: + configured = options.output.get("world_model_manifest_path") + if configured is not None: + return Path(cast(Any, configured)) + if options.legacy_world_manifest is not None: + return options.legacy_world_manifest + bundled = _LOCAL_WINDOW_MANIFESTS.get(config.runner_name) + return None if bundled is None else Path(bundled) + + +def _validate_fields( + section: str, + values: Mapping[str, object], + allowed: set[str] | frozenset[str], +) -> None: + unknown = sorted(set(values) - allowed) + if unknown: + raise ValueError( + f"Unsupported OmniDreams {section} fields: {', '.join(unknown)}." + ) + + +def _is_single_view(config: RunnerConfig) -> bool: + diffusion_model = getattr(config.pipeline, "diffusion_model", None) + transformer: Any = getattr(diffusion_model, "transformer", None) + return int(getattr(transformer, "num_views", 1)) == 1 + + +LAUNCH_CAPABILITY = OmnidreamsLaunchCapability() + +__all__ = ["LAUNCH_CAPABILITY", "OmnidreamsLaunchCapability"] diff --git a/integrations/omnidreams/omnidreams/output_targets.py b/integrations/omnidreams/omnidreams/output_targets.py deleted file mode 100644 index 6bce82176..000000000 --- a/integrations/omnidreams/omnidreams/output_targets.py +++ /dev/null @@ -1,143 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""OmniDreams output capabilities for ``flashdreams-run``.""" - -from __future__ import annotations - -from pathlib import Path -from typing import Any - -from flashdreams.infra.runner import RunnerConfig -from flashdreams.serving.output_targets import ( - OutputLaunchOptions, - OutputMode, - OutputTargetSpec, -) - -_LOCAL_WINDOW_MANIFESTS = { - "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae": "example_world_model.yaml", - "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf": ( - "example_world_model_perf.yaml" - ), - "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-native-perf": ( - "example_world_model_perf.yaml" - ), -} - - -class OmnidreamsOutputTargetAdapter: - def supported_modes( - self, - config: RunnerConfig, - options: OutputLaunchOptions, - ) -> tuple[OutputMode, ...]: - modes: list[OutputMode] = [] - if _is_single_view(config): - modes.append("webrtc") - if _local_window_manifest(config, options) is not None: - modes.append("local-window") - return tuple(modes) - - def resolve( - self, - config: RunnerConfig, - *, - mode: OutputMode, - options: OutputLaunchOptions, - ) -> OutputTargetSpec | None: - if mode == "webrtc" and _is_single_view(config): - return _webrtc_spec(config, options) - if mode == "local-window": - manifest = _local_window_manifest(config, options) - if manifest is not None: - return _local_window_spec(config, manifest) - return None - - -def _webrtc_spec( - config: RunnerConfig, - options: OutputLaunchOptions, -) -> OutputTargetSpec: - argv = [ - "webrtc", - "--preset-id", - _pipeline_name(config), - "--device", - str(config.device), - "--fps", - str(getattr(config, "output_fps", 30)), - "--video-height", - str(getattr(config, "pixel_height", 704)), - "--video-width", - str(getattr(config, "pixel_width", 1280)), - ] - seed = _diffusion_seed(config) - if seed is not None: - argv.extend(("--seed", str(seed))) - _append_postprocess_preset(argv, config) - 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="OmniDreams shared demo WebRTC server", - module="omnidreams.demo.app", - argv=tuple(argv), - ) - - -def _local_window_spec(config: RunnerConfig, manifest: Path) -> OutputTargetSpec: - argv = ["--manifest", str(manifest)] - _append_postprocess_preset(argv, config) - return OutputTargetSpec( - mode="local-window", - label="Omnidreams local interactive window", - module="omnidreams.interactive_drive", - argv=tuple(argv), - notes=( - "Local-window uses the OmniDreams interactive-drive manifest for " - "scene, resolution, and runtime-specific controls.", - ), - ) - - -def _local_window_manifest( - config: RunnerConfig, - options: OutputLaunchOptions, -) -> Path | None: - if options.local_window_manifest is not None: - return options.local_window_manifest - manifest = _LOCAL_WINDOW_MANIFESTS.get(config.runner_name) - return None if manifest is None else Path(manifest) - - -def _pipeline_name(config: RunnerConfig) -> str: - name = getattr(config.pipeline, "name", None) - return str(name or config.runner_name) - - -def _diffusion_seed(config: RunnerConfig) -> int | None: - diffusion_model = getattr(config.pipeline, "diffusion_model", None) - seed = getattr(diffusion_model, "seed", None) - return None if seed is None else int(seed) - - -def _is_single_view(config: RunnerConfig) -> bool: - diffusion_model = getattr(config.pipeline, "diffusion_model", None) - transformer: Any = getattr(diffusion_model, "transformer", None) - return int(getattr(transformer, "num_views", 1)) == 1 - - -def _append_postprocess_preset(argv: list[str], config: RunnerConfig) -> None: - preset = config.postprocess.preset - if preset: - argv.extend(("--postprocess-preset", str(preset))) - - -OUTPUT_TARGET_ADAPTER = OmnidreamsOutputTargetAdapter() - -__all__ = ["OUTPUT_TARGET_ADAPTER", "OmnidreamsOutputTargetAdapter"] diff --git a/integrations/omnidreams/omnidreams/prepare.py b/integrations/omnidreams/omnidreams/prepare.py index 2a8858dbb..651855402 100644 --- a/integrations/omnidreams/omnidreams/prepare.py +++ b/integrations/omnidreams/omnidreams/prepare.py @@ -8,7 +8,7 @@ * ``nvidia/omni-dreams-scenes`` USDZ archives -> consumed sealed by the desktop ``interactive-drive`` demo and unpacked on demand by - ``omnidreams.webrtc.server`` (both read from the shared cache at + the centralized ``webrtc`` launch (both read from the shared cache at ``$FLASHDREAMS_CACHE_DIR/omnidreams-scenes/``; see :mod:`omnidreams.scenes`). * The Cosmos-Reason1 text encoder used by the flashdreams world-model diff --git a/integrations/omnidreams/omnidreams/runner.py b/integrations/omnidreams/omnidreams/runner.py index 589f2aa58..643af0ce7 100644 --- a/integrations/omnidreams/omnidreams/runner.py +++ b/integrations/omnidreams/omnidreams/runner.py @@ -29,8 +29,10 @@ from dataclasses import dataclass, field from pathlib import Path +from typing import Annotated import torch +import tyro from einops import rearrange from loguru import logger from omnidreams.model_session import OmnidreamsModelSessionCore @@ -156,7 +158,9 @@ class OmnidreamsRunnerConfig(RunnerConfig): """ _target: type["OmnidreamsRunner"] = field(default_factory=lambda: OmnidreamsRunner) - output_adapter: str | None = "omnidreams.output_targets:OUTPUT_TARGET_ADAPTER" + launch_capability: Annotated[str | None, tyro.conf.Suppress] = ( + "omnidreams.launch:LAUNCH_CAPABILITY" + ) prompt: str = "" """Default text prompt applied to every camera. Override per-camera diff --git a/integrations/omnidreams/pyproject.toml b/integrations/omnidreams/pyproject.toml index fe48be188..558bd048a 100644 --- a/integrations/omnidreams/pyproject.toml +++ b/integrations/omnidreams/pyproject.toml @@ -20,7 +20,7 @@ build-backend = "setuptools.build_meta" [project] name = "flashdreams-omnidreams" version = "0.1.0" -description = "Omnidreams inference with flashdreams (webrtc / gRPC servers + the interactive-drive desktop demo)" +description = "Omnidreams inference with FlashDreams unified WebRTC and local-window launch modes" readme = "README.md" # PyNvVideoCodec 2.1 currently publishes wheels through CPython 3.12. requires-python = ">=3.10,<3.13" @@ -92,8 +92,8 @@ dev = [ ] [project.scripts] -# ``omnidreams-prepare`` stages resources used by *both* demo paths -# (desktop ``interactive-drive`` *and* ``omnidreams.webrtc.server``): +# ``omnidreams-prepare`` stages resources used by both ``local-window`` and +# ``webrtc`` launch modes: # the scene USDZs from ``nvidia/omni-dreams-scenes`` and the # Cosmos-Reason1 text encoder pinned to its runtime revision. Lives at # the package level (``omnidreams.prepare``) rather than under @@ -105,15 +105,6 @@ omnidreams-prepare = "omnidreams.prepare:main" # FlashDreams generation, and DrivingGen adapter setup. omnidreams-eval = "omnidreams.eval.cli:main" -# Experimental shared demo API path. This coexists with the legacy -# WebRTC/gRPC/interactive-drive demos until the new adapter is proven. -omnidreams-demo = "omnidreams.demo.app:main" - -# Desktop interactive-drive demo entry point. Requires the -# ``interactive-drive`` extra (it adds slangpy); without it the -# presenter import fails fast with a clear message. -interactive-drive = "omnidreams.interactive_drive.demo:main" - # Tkinter wizard that calibrates a connected steering wheel or game # controller and writes a local input profile the demo auto-discovers. # Stdlib-only (Tkinter); does not need the ``interactive-drive`` extra. @@ -124,8 +115,8 @@ interactive-drive-configuration = "omnidreams.interactive_drive.input_config.app # scans this group at CLI startup; the entry-point name itself is purely # informational, the registry key always comes from ``cfg.runner_name``. [project.entry-points."flashdreams.runner_configs"] -"omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae" = "omnidreams.config:RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE" -"omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf" = "omnidreams.config:RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF" +"omnidreams" = "omnidreams.config:RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE" +"omnidreams-perf" = "omnidreams.config:RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF" [tool.setuptools.packages.find] include = ["omnidreams*"] diff --git a/integrations/omnidreams/tests/interactive_drive/test_demo_scene_selection.py b/integrations/omnidreams/tests/interactive_drive/test_demo_scene_selection.py index 99297496e..b9535a722 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_demo_scene_selection.py +++ b/integrations/omnidreams/tests/interactive_drive/test_demo_scene_selection.py @@ -27,6 +27,35 @@ def test_auto_start_flag_and_deprecated_alias() -> None: assert parser.parse_args(["--no-autoload-scene"]).auto_start is False +def test_central_local_window_launch_builds_namespace_without_reparsing_argv( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + captured: list[argparse.Namespace] = [] + monkeypatch.setattr(demo_mod, "_run_namespace", captured.append) + world_manifest = tmp_path / "world.yaml" + scene = tmp_path / "scene.usdz" + config = types.SimpleNamespace( + postprocess=types.SimpleNamespace(preset="rtx-super-resolution") + ) + + demo_mod.launch_from_runner( + config=config, + world_model_manifest=world_manifest, + scenario={"scene": scene, "auto_start": True}, + output={"no_hud": True, "stream_mjpeg": ":8080"}, + ) + + args = captured[0] + assert args.backend == "omnidreams" + assert args.manifest == world_manifest + assert args.scene == scene + assert args.auto_start is True + assert args.no_hud is True + assert args.stream_mjpeg == ":8080" + assert args.postprocess_preset == "rtx-super-resolution" + + def test_resolve_scene_variant_prefers_weather_archive_path_for_default( tmp_path: Path, ) -> None: diff --git a/integrations/omnidreams/tests/test_demo_api.py b/integrations/omnidreams/tests/test_demo_api.py index 846e78f96..0761cd49b 100644 --- a/integrations/omnidreams/tests/test_demo_api.py +++ b/integrations/omnidreams/tests/test_demo_api.py @@ -13,6 +13,7 @@ import numpy as np import omnidreams.demo as demo_package +import omnidreams.demo.app as demo_app_module import omnidreams.demo.spec as spec_module import pytest import tomli as tomllib @@ -95,6 +96,31 @@ def test_omnidreams_demo_defaults_to_stable_non_perf_preset() -> None: assert not args.preset_id.endswith("-perf") +def test_omnidreams_direct_runner_launch_builds_null_spec( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: list[DemoSpec] = [] + + def fake_run_replay_demo(*, spec: DemoSpec, adapter: object) -> str: + del adapter + captured.append(spec) + return "completed" + + monkeypatch.setattr(demo_app_module, "run_replay_demo", fake_run_replay_demo) + config = OMNIDREAMS_RUNNERS["omnidreams"] + + result = demo_app_module.launch_from_runner( + config=config, + mode="null", + scenario={"example_data": True, "total_blocks": 2}, + output={}, + ) + + assert result == "completed" + assert captured[0].preset_id == config.pipeline.name + assert isinstance(captured[0].output, NullOutputSpec) + + def test_omnidreams_replay_cli_builds_null_output_spec() -> None: args = parse_args(["replay", "--output-mode", "null"]) @@ -364,7 +390,7 @@ def fake_sync(uuid: str) -> tuple[tuple[Path, ...], tuple[Path, ...]]: assert scenario.first_frame_paths == (first_frame,) assert scenario.camera_names == ("camera_front_wide_120fov",) assert scenario.prompts == ( - str(getattr(OMNIDREAMS_RUNNERS[DEFAULT_OMNIDREAMS_PRESET], "prompt")), + str(getattr(OMNIDREAMS_RUNNERS["omnidreams"], "prompt")), ) diff --git a/integrations/omnidreams/tests/test_quality_regression.py b/integrations/omnidreams/tests/test_quality_regression.py index f12dd8842..4cdd98897 100644 --- a/integrations/omnidreams/tests/test_quality_regression.py +++ b/integrations/omnidreams/tests/test_quality_regression.py @@ -32,7 +32,7 @@ pytestmark = pytest.mark.ci_gpu _ENV_PREFIX = "FLASHDREAMS_OMNIDREAMS_QUALITY_" -_DEFAULT_RUNNER = "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae" +_DEFAULT_RUNNER = "omnidreams" _DEFAULT_TOTAL_BLOCKS = 4 # ~1 second for chunk2 + 30fps output. diff --git a/integrations/omnidreams/tests/test_recipe_configs.py b/integrations/omnidreams/tests/test_recipe_configs.py index 4de5e2ac4..15e4004c6 100644 --- a/integrations/omnidreams/tests/test_recipe_configs.py +++ b/integrations/omnidreams/tests/test_recipe_configs.py @@ -17,8 +17,8 @@ The full numerics / GPU tests live alongside this file (they need GPU + checkpoints). These smoke tests just confirm the plugin is wired -correctly: importable, every ``runner_name`` mirrors its -``pipeline.name``, descriptions are non-empty, and the +correctly: importable, public slugs map to their intended internal +pipeline presets, descriptions are non-empty, and the entry-point declarations in ``pyproject.toml`` match the ``omnidreams.config`` ``RUNNER_*`` literals exactly. """ @@ -46,14 +46,15 @@ def test_runners_dict_is_non_empty() -> None: assert OMNIDREAMS_RUNNERS, "OMNIDREAMS_RUNNERS is empty" -def test_runner_name_mirrors_pipeline_name() -> None: - """``runner_name`` must equal ``pipeline.name`` per the CLI contract.""" - drifted = { - slug: (cfg.runner_name, cfg.pipeline.name) - for slug, cfg in OMNIDREAMS_RUNNERS.items() - if cfg.runner_name != cfg.pipeline.name +def test_public_runner_slugs_map_to_internal_pipeline_presets() -> None: + """Short public slugs must keep selecting the intended model presets.""" + expected = { + "omnidreams": "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae", + "omnidreams-perf": "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf", } - assert not drifted, f"runner_name != pipeline.name: {drifted}" + actual = {slug: cfg.pipeline.name for slug, cfg in OMNIDREAMS_RUNNERS.items()} + assert actual == expected + assert all(slug == cfg.runner_name for slug, cfg in OMNIDREAMS_RUNNERS.items()) def test_runners_have_descriptions() -> None: diff --git a/tests/README.md b/tests/README.md index 93f535458..fa0651a97 100644 --- a/tests/README.md +++ b/tests/README.md @@ -78,7 +78,7 @@ your environment needs Hugging Face authentication for model or dataset access. mkdir -p /tmp/omnidreams_quality_ref /tmp/omnidreams_quality_artifacts uv run --project integrations/omnidreams flashdreams-run \ - omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae \ + omnidreams \ --example-data True \ --example_data_uuid "239560dc-33d1-11ef-9720-00044bcbccac" \ --total-blocks 4 \ diff --git a/uv.lock b/uv.lock index 4d0a32ba5..ced04ff15 100644 --- a/uv.lock +++ b/uv.lock @@ -971,6 +971,7 @@ dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, { name = "nvidia-ml-py" }, + { name = "pyyaml" }, { name = "safetensors" }, { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(sys_platform != 'win32' and extra == 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, { name = "torch", version = "2.12.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'win32' and extra == 'extra-11-flashdreams-dev') or (sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, @@ -1047,6 +1048,7 @@ requires-dist = [ { name = "opencv-python-headless", marker = "extra == 'runners'", specifier = ">=4.5" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, { name = "pytest-manual-marker", marker = "extra == 'dev'", specifier = ">=2.0" }, + { name = "pyyaml", specifier = ">=6.0" }, { name = "safetensors", specifier = ">=0.4" }, { name = "scipy", marker = "extra == 'examples'", specifier = ">=1.11" }, { name = "scipy", marker = "extra == 'runners'", specifier = ">=1.11" },