diff --git a/integrations_v2/omnidreams/README.md b/integrations_v2/omnidreams/README.md new file mode 100644 index 000000000..80776c677 --- /dev/null +++ b/integrations_v2/omnidreams/README.md @@ -0,0 +1,257 @@ + + +# Omnidreams driving video + +The first model on the v2 API that generates from something other than a +prompt. It continues from a first frame and follows an HDMap, the road layout, +for every frame it generates. That layout either comes from a recording of one +or is drawn as the run goes by the Ludus rasterizer. The model is the +single-view distilled checkpoint the `flashdreams-omnidreams` package already +configures for the v1 runner; this package is the application around it and +holds no model code of its own. + +The Python module here is `omnidreams_v2` rather than `omnidreams`, which the +v1 package owns and this one imports. + +## Set up + +From the workspace root: + +```bash +uv sync --package flashdreams-omnidreams-v2 --inexact +export HF_TOKEN= +``` + +The sync installs this application, the v1 `flashdreams-omnidreams` package +that holds the model and the Ludus rasterizer, and the `flashdreams-run-v2` +command itself. `--inexact` leaves anything else already in the environment +alone; drop it to have uv prune the environment down to what this needs. The +`interactive-drive` extra the v1 package documents is for its desktop +presenter, which nothing here uses. + +The token is for Hugging Face, where the checkpoint, the sample recordings and +the scenes all live. The checkpoint is fetched on the first run, tens of +gigabytes of it including the Cosmos-Reason1 text encoder, so expect to wait. + +## Generate a clip + +```bash +uv run --no-sync flashdreams-run-v2 omnidreams --output-path drive.mp4 +``` + +That downloads the default scene from `nvidia/omni-dreams-scenes` and has the +Ludus rasterizer draw the road layout while the run generates. Everything else +has a default too, including the frame the run continues from and the prompt, +both of which come out of the scene. + +Writing to a file, the runner says nothing until it is done, when it prints the +path. So a run is silent through the download, the checkpoint, the rasterizer +build and the generating, and only the model's own output breaks that up. Start +with `--max-blocks 8 --no-compile`, below, to keep the silence short. + +`uv run --no-sync` is what reaches the command, which the sync installed into +the workspace `.venv` rather than onto your `PATH`. Run +`source .venv/bin/activate` once instead and `flashdreams-run-v2` works on its +own. + +Arguments after `--` go to the application, and `-- --help` lists them, though +the runner insists on an output path before it will get that far: + +```bash +uv run --no-sync flashdreams-run-v2 omnidreams --output-path drive.mp4 -- --help +``` + +`--scene` drives a road other than the default, taking either an id to download +or a path to an archive of your own. This lists the ids the scenes dataset has: + +```bash +uv run --no-sync python -c 'from omnidreams.scenes import list_available_scene_uuids; print("\n".join(list_available_scene_uuids()))' +``` + +Any of them goes straight after `--scene`: + +```bash +uv run --no-sync flashdreams-run-v2 omnidreams --output-path drive.mp4 \ + -- --scene 0d404ff7-2b66-498c-b047-1ed8cded60d4 +``` + +That particular id is the default, named explicitly. + +## Where the road comes from + +Drawing is the default because it is what a run that eventually steers needs: a +recording cannot show a road nobody drove down. Nothing steers yet, so a drawn +run follows the drive its scene recorded, which makes it as repeatable as +replaying a video of one. Once input is wired up, the poses stop coming from the +recording and start coming from what the driver did, and nothing above the +renderer changes. + +A scene carries more than the layout. The frame the run continues from is the +one its front camera actually captured, and the drawing starts at that same +moment rather than at the top of the scene, so the model is shown the road it is +looking at. The scene's own description of the road becomes the prompt. The +recorded drive is sampled at 10Hz where the model generates at 30fps, so the +layout is resampled onto the generated rate. The default scene is 100 seconds of +road. + +## Something to point the examples at + +The examples below take a frame and a video of your own. If you have none to +hand, the bundled sample carries both, and this puts them in your shell: + +```bash +eval "$(uv run --no-sync python -c 'from omnidreams_v2.samples import DEFAULT_HDMAP_SAMPLE, fetch_hdmap_sample; v, f = fetch_hdmap_sample(DEFAULT_HDMAP_SAMPLE); print(f"export HDMAP_VIDEO={v}"); print(f"export HDMAP_FRAME={f}")')" +echo "$HDMAP_VIDEO" "$HDMAP_FRAME" +``` + +It prints nothing itself -- `eval` eats what it printed, which is the point -- +so the `echo` is how you see it worked. Give it a moment either way: importing +reaches the model package, and the files download the first time. + +## Drive the same road differently + +`--first-frame` and `--prompt` each replace what the scene supplied without +giving up the drawing, which is how one road is driven under weather it never +recorded: + +```bash +uv run --no-sync flashdreams-run-v2 omnidreams --output-path drive.mp4 \ + -- --first-frame "$HDMAP_FRAME" --prompt "The same street after dark." +``` + +The layout is still drawn from the scene, and still drawn from the moment the +scene recorded, since a frame of your own says nothing about where along the road +it was taken. Only the picture the model continues from changes. Naming a frame +is not asking to replay anything, so the default scene keeps being drawn when +`--scene` says nothing -- as above. + +The two are worth changing together. A first frame showing a different road than +the layout describes hands the model a contradiction, and a prompt still +describing the scene's own weather works against a frame that shows other +weather. The sample's frame above is a different road from the default scene, so +that command shows the path working rather than showing it working well; a frame +of the same road under other conditions is what this is for. + +The first drawn run pauses a few minutes to build the rasterizer's CUDA +extension, which needs `nvcc` new enough for your GPU on `PATH` or at +`CUDA_HOME` -- 12.8 or later for Blackwell. The build failing with `unsupported +gpu architecture` means the toolkit is older than the card. + +## Replay a recording instead + +`--hdmap` gives up the rasterizer and reads the layout from video someone +already rendered, which is what a benchmark wants and what runs on a machine +with no scene and no rasterizer: + +```bash +uv run --no-sync flashdreams-run-v2 omnidreams --output-path drive.mp4 \ + -- --hdmap +``` + +Bare, that fetches the default recording from the `nvidia/omni-dreams-samples` +dataset. A sample carries the frame to continue from as well as the layout, so +it needs nothing else said about it. Naming an id picks another, and they are +listed +[here](https://huggingface.co/datasets/nvidia/omni-dreams-samples/tree/main/data/single_view): + +```bash +uv run --no-sync flashdreams-run-v2 omnidreams --output-path drive.mp4 \ + -- --hdmap 239560dc-33d1-11ef-9720-00044bcbccac +``` + +That one is the default again, named rather than left unsaid. + +Point `--hdmap` at files instead to replay a recording of your own, one video +per camera, alongside the frame each continues from: + +```bash +uv run --no-sync flashdreams-run-v2 omnidreams --output-path drive.mp4 \ + -- --hdmap "$HDMAP_VIDEO" --first-frame "$HDMAP_FRAME" +``` + +Which of the two you meant is read off the file extension, the way the samples +are named: `--hdmap 239560dc-...` is an id to download and `--hdmap road.mp4` is +a file to read, so a misspelled path is reported as the missing file it is. + +Naming a recording and a scene at once is refused rather than resolved, since +they are two answers to the one question of where the layout comes from. +`--first-frame` is not one of those answers, which is why it works either way: +it says what the run continues from, and a drawn run answers that too. + +## How long a run is + +About a minute, which is what a run that was told nothing produces. Long enough +to see whether a drive holds together and short enough to wait for, since +generating is real time at best. + +`--max-blocks` says otherwise. `--max-blocks 8` is 61 frames, about two +seconds, which is what you want for a first run; pair it with `--no-compile`, +since compilation is on in the model's own config and costs minutes on the first +run to save milliseconds a block. + +```bash +uv run --no-sync flashdreams-run-v2 omnidreams --output-path drive.mp4 \ + -- --max-blocks 8 --no-compile +``` + +`--max-blocks 0` drives to the end of the road however long that takes, which is +what an interactive session wants and what the default minute would otherwise +cut off. On the default scene that is 100 seconds of video, so it is a longer +wait than anything else here: + +```bash +uv run --no-sync flashdreams-run-v2 omnidreams --output-path drive.mp4 \ + -- --max-blocks 0 +``` + +Either way the road can end first, and the run stops on a block boundary rather +than on a block it has only part of the layout for. The first block decodes 5 +frames and every block after it 8, at 30 frames per second, so a ten-second +recording is about 38 blocks and a minute is 226. + +## What it generates + +1280x704 frames at 30fps, laid out `bvtchw`, as `[-1, 1]` floats on the GPU. +Those numbers are the checkpoint's, read off the runner config the +`flashdreams-omnidreams` package ships rather than written down here. Something +else can be asked for with `--pixel-width` and `--pixel-height` before the `--`, +each a multiple of 8. + +One camera per run: an MP4 holds one sequence of frames, so the file window +rejects output with more than one view in it. The multi-view checkpoints need a +window that lays the cameras out, which is not built yet. + +## The seams underneath + +`HDMapSource` in `conditioning.py` is the seam a session reads its conditioning +through, and has two implementations. `RenderedHDMapSource` keeps a run's place +along a scene and draws each chunk as the run reaches it. +`PrecomputedHDMapSource` reads the same chunks out of recorded video instead. +A session cannot tell which it was given. + +Under the drawn one is a smaller seam, `SceneRenderer`, whose only +implementation is `LudusSceneRenderer` in `ludus.py`. That split is what keeps +the CUDA rasterizer out of the tests: everything around it -- working along a +scene consecutively, turning bytes into the pixels the model reads, starting the +layout at the moment the run continues from -- is covered on a CPU against a +stand-in renderer, and each of those would produce a plausible-looking wrong +drive rather than an error. + +Steering is what this shape is for. `ISession.step` already receives input +events and passes them to the source, where a drawn source is the one that could +act on them; today it ignores them and follows the recorded drive. + +## Tests + +```bash +uv sync --package flashdreams-omnidreams-v2 --group test --inexact +uv run --no-sync pytest integrations_v2/omnidreams -m ci_cpu -v +``` + +Those run the application against a stand-in model, a stand-in drive and a +stand-in renderer, which covers what is particular here: each block is +conditioned on exactly the frames it generates, a run ends when the road does, +and a drawn run is shown the road from the moment it continues from. diff --git a/integrations_v2/omnidreams/omnidreams_v2/__init__.py b/integrations_v2/omnidreams/omnidreams_v2/__init__.py new file mode 100644 index 000000000..3380fc50d --- /dev/null +++ b/integrations_v2/omnidreams/omnidreams_v2/__init__.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Omnidreams driving application for the FlashDreams v2 API.""" + +from .app import ( + OmnidreamsApplication, + OmnidreamsSessionConfig, + SceneDrive, + create_app, +) +from .conditioning import ( + HDMapSource, + PrecomputedHDMapSource, + RenderedHDMapSource, + SceneRenderer, +) +from .ludus import LudusSceneRenderer +from .session import OmnidreamsSession + +__all__ = [ + "HDMapSource", + "LudusSceneRenderer", + "OmnidreamsApplication", + "OmnidreamsSession", + "OmnidreamsSessionConfig", + "PrecomputedHDMapSource", + "RenderedHDMapSource", + "SceneDrive", + "SceneRenderer", + "create_app", +] diff --git a/integrations_v2/omnidreams/omnidreams_v2/app.py b/integrations_v2/omnidreams/omnidreams_v2/app.py new file mode 100644 index 000000000..b4c6bae58 --- /dev/null +++ b/integrations_v2/omnidreams/omnidreams_v2/app.py @@ -0,0 +1,606 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Omnidreams application, generating driving video from a road layout.""" + +import argparse +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from omnidreams.config import RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE +from omnidreams.interactive_drive.math3d import normalize_camera_name + +from flashdreams.api_v2.application import IApplication +from flashdreams.api_v2.session import ISession +from flashdreams.infra.config import derive_config +from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.video_tensor import VideoTensorLayout + +from .conditioning import HDMapSource, PrecomputedHDMapSource, RenderedHDMapSource +from .ludus import LudusSceneRenderer +from .samples import DEFAULT_HDMAP_SAMPLE, fetch_hdmap_sample +from .scenes import ( + DEFAULT_SCENE, + DEFAULT_SCENE_CAMERA, + fetch_scene, + read_prompt, + read_seed_frame, +) +from .session import OmnidreamsSession + +_RUNNER_CONFIG = RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE +"""Shipped single-view distilled model, and the defaults that come with it.""" + +_OUTPUT_LAYOUT = VideoTensorLayout.bvtchw +"""What this model emits: a batch of cameras, each a sequence of frames.""" + +_FRAMES_PER_SECOND_FOR_UI = 60 +"""Rate an interactive window would read input and present results at.""" + +_DEFAULT_RUN_SECONDS = 60 +"""How much video a run produces when it was not told how much to produce. + +Long enough to show whether a drive holds together, short enough to wait for. +A run generates in real time at best, so an hour of road is an hour of waiting, +which is not what someone who named no length was asking for. +""" + + +@dataclass(frozen=True, kw_only=True, slots=True) +class SceneDrive: + """A scene to draw the road from, and where along it a run starts.""" + + archive: Path + """Scene on disk, holding the road layout and the drive recorded along it.""" + + camera: str + """Camera to draw from, spelled the way the scene spells it.""" + + first_frame_path: Path + """Frame the run continues from: the one unpacked from the scene, or one the + command line named in its place.""" + + view_start_us: int + """When the scene's own frame was captured, which is where the drawing + starts whichever frame the run continues from.""" + + +@dataclass(frozen=True, kw_only=True, slots=True) +class OmnidreamsSessionConfig: + """What one command line resolved to, shared by every session it creates.""" + + prompt: str + """Text describing the drive, applied to every camera.""" + + device: str + """Device the pipeline is built on.""" + + max_blocks: int | None + """Blocks to stop after: a count, ``0`` for no limit, or ``None`` for the + default runtime, which takes the model's block sizes to work out.""" + + hdmap_video_paths: tuple[Path, ...] + """HDMap video per camera, in camera order. Empty for a drawn run.""" + + first_frame_paths: tuple[Path, ...] + """Frame to continue from per camera, in the same order. Empty for a drawn + run, which takes its frame from the scene instead.""" + + view_names: tuple[str, ...] + """Camera labels, in the same order.""" + + scene: SceneDrive | None + """Scene to draw the road from, or ``None`` to replay a recording of it.""" + + +class OmnidreamsApplication(IApplication): + """Omnidreams: driving video generated from a rendered road layout. + + The model continues from a first frame and follows the HDMap it is given, + so a run is described by where that layout comes from rather than by a + prompt alone. It comes either from a scene drawn by the Ludus rasterizer as + the run goes, which is what a run gets by default, or from recorded video, + which is what ``--hdmap`` asks for instead. + + Drawing is the default because it is what a run that eventually steers + needs: a recording cannot show a road nobody drove down. Nothing steers + yet, so a drawn run follows the drive its scene recorded, which is as + repeatable as replaying a video of it. + + A scene supplies the frame to continue from and the prompt as well as the + layout, and ``--first-frame`` and ``--prompt`` each replace what it supplied + without giving up the drawing. That is how the one road is driven under a + sky or a season the scene never recorded. + + The model is loaded once, on the first session, and shared by every session + after it, since loading reads a checkpoint of several gigabytes. + """ + + def __init__( + self, + *, + pipeline_config: Any | None = None, + source_factory: Callable[[OmnidreamsSessionConfig, SessionDesc], HDMapSource] + | None = None, + ) -> None: + """ + Args: + pipeline_config: Model to run, in place of the shipped distilled + checkpoint. A test passes a stand-in. + source_factory: Where conditioning comes from, in place of the + recording or the scene the command line named. A test passes a + stand-in for both. + """ + self._pipeline_config = ( + _RUNNER_CONFIG.pipeline if pipeline_config is None else pipeline_config + ) + self._source_factory = source_factory or _configured_source + self._config: OmnidreamsSessionConfig | None = None + self._pipeline: Any = None + + @property + def pipeline_config(self) -> Any: + """Model this will load, including whatever the command line changed.""" + return self._pipeline_config + + def init(self, commandline_args: Sequence[str]) -> None: + """Parse what road to drive, how far along it, and where the road is. + + Also where the road layout comes from, which is what a scene and a + recording are two answers to. Not what size or rate to generate at: that + describes the session, which the caller asks for. The model is not + loaded here either, though a scene is fetched so that a run that cannot + find its road says so before waiting on a checkpoint. + + Raises: + ValueError: The conditioning is described twice or by halves, or the + rollout length is not one this model can generate. + FileNotFoundError: The named scene or recording does not exist. + """ + parser = argparse.ArgumentParser( + prog="flashdreams-run-v2 omnidreams --", + description="Generate driving video from a road layout.", + ) + parser.add_argument( + "--scene", + default=None, + metavar="ID_OR_PATH", + help=( + "Scene to draw the road layout from, as an id to download or a " + "path to an archive. Drawing is what a run that names no " + f"source of its own does. Default: {DEFAULT_SCENE}." + ), + ) + parser.add_argument( + "--hdmap", + nargs="*", + default=None, + metavar="ID_OR_PATH", + help=( + "Replay an already-rendered HDMap rather than drawing one: a " + "sample id to download, or one video per camera. Bare, the " + f"default sample ({DEFAULT_HDMAP_SAMPLE})." + ), + ) + parser.add_argument( + "--first-frame", + type=Path, + nargs="+", + default=(), + metavar="PATH", + help=( + "Image or video to continue from, one per camera. Frame zero of " + "a video is taken. Default: the frame the scene recorded, or " + "the one a sample carries. Required alongside HDMap videos of " + "your own, which carry none." + ), + ) + parser.add_argument( + "--view-names", + nargs="+", + default=(), + metavar="NAME", + help=( + "Camera labels, one per camera. Default: indexed placeholders, " + "which is all a single-camera model needs." + ), + ) + parser.add_argument( + "--prompt", + default=None, + help=( + "Text describing the drive. Default: the scene's own " + "description of it, or failing that the model's." + ), + ) + parser.add_argument( + "--device", + default="cuda", + help="Device to load the model on. Default: %(default)s.", + ) + parser.add_argument( + "--max-blocks", + type=int, + default=None, + help=( + "Stop after this many blocks, or 0 to drive to the end of the " + f"road however long that takes. Default: about " + f"{_DEFAULT_RUN_SECONDS} seconds of video." + ), + ) + parser.add_argument( + "--compile", + action=argparse.BooleanOptionalAction, + default=None, + help=( + "Compile the network, costing minutes once and saving " + "milliseconds a step. Default: whatever the model's config says." + ), + ) + parser.add_argument( + "--seed", + type=int, + default=None, + help=( + "Seed the noise a run samples from, so the same command " + "generates the same clip. Default: whatever the model's config " + "says, which is usually a seed of its own." + ), + ) + args = parser.parse_args(list(commandline_args)) + + if args.max_blocks is not None and args.max_blocks < 0: + raise ValueError(f"--max-blocks cannot be negative, got {args.max_blocks}.") + scene = _resolve_scene(args) + if scene is None: + hdmap_videos, first_frames = _resolve_conditioning( + args.hdmap, tuple(args.first_frame) + ) + defaults = tuple(f"view_{index}" for index in range(len(hdmap_videos))) + else: + hdmap_videos, first_frames = (), () + # A scene names its cameras, so a drawn run can label its own view + # rather than leaving a placeholder where the label goes. + defaults = (normalize_camera_name(scene.camera)[1],) + view_names = _resolve_view_names(tuple(args.view_names), defaults) + + if args.compile is not None: + self._pipeline_config = derive_config( + self._pipeline_config, + diffusion_model={"transformer": {"compile_network": args.compile}}, + ) + if args.seed is not None: + self._pipeline_config = derive_config( + self._pipeline_config, diffusion_model={"seed": args.seed} + ) + self._config = OmnidreamsSessionConfig( + prompt=_resolve_prompt(args.prompt, scene), + device=args.device, + max_blocks=args.max_blocks, + hdmap_video_paths=hdmap_videos, + first_frame_paths=first_frames, + view_names=view_names, + scene=scene, + ) + + def session_desc(self) -> SessionDesc: + """Return the description of a session this application uses. + + The model generates its best video at the size and rate it was trained + at, so that is what a caller with none of its own in mind gets. + """ + return SessionDesc( + output_layout=_OUTPUT_LAYOUT, + frames_per_second_for_ui=_FRAMES_PER_SECOND_FOR_UI, + frames_per_second_for_step=_RUNNER_CONFIG.output_fps, + video_width=_RUNNER_CONFIG.pixel_width, + video_height=_RUNNER_CONFIG.pixel_height, + ) + + def create_session(self, session_desc: SessionDesc) -> ISession: + """Create one uninitialized session, loading the model if needed. + + Raises: + RuntimeError: :meth:`init` has not run yet. + ValueError: The description asks for output this cannot generate. + """ + config = self._config + if config is None: + raise RuntimeError( + f"{type(self).__name__}.init() must run before create_session()." + ) + # Before loading rather than after: a checkpoint of several gigabytes is + # a long wait for a layout this was never going to accept. + self._validate_layout(session_desc) + if self._pipeline is None: + self._pipeline = self._pipeline_config.setup().to(config.device).eval() + self._validate_frame_size(session_desc, self._pipeline) + return OmnidreamsSession( + self._pipeline, + config.prompt, + self._source_factory(config, session_desc), + session_desc, + self._resolve_max_blocks(config, session_desc), + ) + + def _resolve_max_blocks( + self, config: OmnidreamsSessionConfig, session_desc: SessionDesc + ) -> int | None: + """Return the block count to stop a run after, or ``None`` for no limit. + + A default has to be worked out here rather than while parsing, because + how many blocks a length of video takes is something only the loaded + model can say. + """ + if config.max_blocks == 0: + return None + if config.max_blocks is not None: + return config.max_blocks + return _blocks_for_seconds( + self._pipeline, + frames_per_second=session_desc.frames_per_second_for_step, + seconds=_DEFAULT_RUN_SECONDS, + ) + + def close(self) -> None: + """Release the model, and whatever memory it was holding.""" + pipeline = self._pipeline + self._pipeline = None + self._config = None + close = getattr(pipeline, "close", None) + if close is not None: + close() + + def _validate_layout(self, session_desc: SessionDesc) -> None: + """Reject a layout this model does not emit. + + Rejecting rather than resolving: a caller that asked for one video and + silently received another has no way to notice. + """ + if session_desc.output_layout is not _OUTPUT_LAYOUT: + raise ValueError( + f"This application only produces {_OUTPUT_LAYOUT.value} output, " + f"got {session_desc.output_layout.value}." + ) + + def _validate_frame_size(self, session_desc: SessionDesc, pipeline: Any) -> None: + """Reject a frame size that is not a whole number of latents across.""" + ratio = pipeline.decoder.spatial_compression_ratio + if session_desc.video_width % ratio or session_desc.video_height % ratio: + raise ValueError( + f"Frame dimensions must be multiples of {ratio}, got " + f"{session_desc.video_width}x{session_desc.video_height}." + ) + + +def create_app() -> IApplication: + """Return a new Omnidreams application.""" + return OmnidreamsApplication() + + +def _configured_source( + config: OmnidreamsSessionConfig, session_desc: SessionDesc +) -> HDMapSource: + """Return conditioning for a run, drawn or replayed as it asked.""" + if config.scene is not None: + return _rendered_source(config, session_desc) + return _recorded_source(config, session_desc) + + +def _recorded_source( + config: OmnidreamsSessionConfig, session_desc: SessionDesc +) -> HDMapSource: + """Return conditioning read from the recordings the command line named.""" + return PrecomputedHDMapSource( + hdmap_video_paths=config.hdmap_video_paths, + first_frame_paths=config.first_frame_paths, + view_names=config.view_names, + pixel_width=session_desc.video_width, + pixel_height=session_desc.video_height, + device=config.device, + ) + + +def _rendered_source( + config: OmnidreamsSessionConfig, session_desc: SessionDesc +) -> HDMapSource: + """Return conditioning drawn from the scene the command line named. + + Raises: + RuntimeError: The run has no scene, so there is nothing to draw. + """ + scene = config.scene + if scene is None: + raise RuntimeError("A drawn run needs a scene to draw.") + return RenderedHDMapSource( + renderer=LudusSceneRenderer( + scene_path=scene.archive, + camera=scene.camera, + view_start_us=scene.view_start_us, + # Drawn at the rate the model generates at, so one drawn frame is + # one generated frame and the drive runs at the speed it recorded. + frames_per_second=session_desc.frames_per_second_for_step, + pixel_width=session_desc.video_width, + pixel_height=session_desc.video_height, + device=config.device, + ), + first_frame_path=scene.first_frame_path, + view_name=config.view_names[0], + pixel_width=session_desc.video_width, + pixel_height=session_desc.video_height, + device=config.device, + ) + + +def _resolve_scene(args: argparse.Namespace) -> SceneDrive | None: + """Return the scene to draw the road from, or ``None`` to replay a recording. + + Drawing is what a run that named no source of its own does, a scene being + the source that has a road in it rather than a picture of one, and so the + only one a run could eventually be steered through. + + ``--first-frame`` is not what decides: it says what the run continues from, + which is a question a drawn run answers too, and answering it with a frame + of your own is how a scene's road is driven under a sky it never saw. + + Raises: + ValueError: A scene was named alongside a recording, which are two + answers to the one question of where the layout comes from. Or more + than one frame was named for the one camera drawn. + """ + if args.hdmap is not None: + if args.scene is not None: + raise ValueError( + "--scene draws the road layout as the run goes, so it cannot " + "be combined with --hdmap, which replays a recording of one." + ) + return None + first_frames = tuple(args.first_frame) + if len(first_frames) > 1: + raise ValueError( + f"Got {len(first_frames)} first frames for a drawn run, which draws " + "the one camera the scene is drawn from. Pass one." + ) + archive = fetch_scene(args.scene or DEFAULT_SCENE) + # Read even when overridden: the timestamp is where the drawing starts, and + # a frame of your own carries no such moment of its own. + recorded_frame, view_start_us = read_seed_frame( + archive, camera=DEFAULT_SCENE_CAMERA + ) + return SceneDrive( + archive=archive, + camera=DEFAULT_SCENE_CAMERA, + first_frame_path=first_frames[0] if first_frames else recorded_frame, + view_start_us=view_start_us, + ) + + +def _resolve_prompt(prompt: str | None, scene: SceneDrive | None) -> str: + """Return the text describing the drive. + + What the command line said, else what the scene says about itself, else the + model's own, which describes the sort of drive it was trained on rather than + this one in particular. + """ + if prompt is not None: + return prompt + if scene is not None: + described = read_prompt(scene.archive) + if described is not None: + return described + return _RUNNER_CONFIG.prompt + + +def _blocks_for_seconds(pipeline: Any, *, frames_per_second: int, seconds: int) -> int: + """Return the blocks whose frames add up to ``seconds`` of video. + + Counted off the model rather than divided out, because the blocks are not + all the same length: a causal decoder's first one is usually shorter, and + how it divides up a rollout is its own business. + + Raises: + ValueError: The model reports a block with no frames in it, so no + number of them would ever amount to a run. + """ + wanted = frames_per_second * seconds + frames = 0 + blocks = 0 + while frames < wanted: + length = pipeline.get_num_frames(blocks) + if length <= 0: + raise ValueError( + f"The model reports {length} frames in block {blocks}, so a " + "run length cannot be worked out from it. Pass --max-blocks." + ) + frames += length + blocks += 1 + return blocks + + +def _resolve_conditioning( + hdmap: list[str], + first_frames: tuple[Path, ...], +) -> tuple[tuple[Path, ...], tuple[Path, ...]]: + """Return the recording and the frame to continue from, per camera. + + Only reached by a run that named ``--hdmap``, which is what asking to replay + something means. What it was given decides which kind: nothing at all is the + default sample, a lone bare id is a sample to download, and anything else is + one video per camera. + + Raises: + ValueError: A sample was named alongside a frame it already carries, or + videos were named without one, or a different number of each was + named. The model rejects a camera count it was not trained for, but + only once it has loaded, so the two are counted against each other + here. + FileNotFoundError: A lone ``--hdmap`` word is neither a path nor a + sample the dataset has. + """ + sample = _named_sample(hdmap) + if sample is not None: + if first_frames: + raise ValueError( + f"--hdmap {sample} is a sample to download, which carries its " + "own frame to continue from, so it cannot be combined with " + "--first-frame." + ) + try: + recording, first_frame = fetch_hdmap_sample(sample) + except FileNotFoundError as exc: + raise FileNotFoundError( + f"No HDMap {sample!r}: there is no such path, and it is no " + f"sample either ({exc})." + ) from exc + return (recording,), (first_frame,) + hdmap_videos = tuple(Path(named) for named in hdmap) + if not first_frames: + raise ValueError("--first-frame is required alongside HDMap videos.") + if len(first_frames) != len(hdmap_videos): + raise ValueError( + f"Got {len(hdmap_videos)} HDMap video(s) and {len(first_frames)} " + "first frame(s): pass one of each per camera." + ) + return hdmap_videos, first_frames + + +def _named_sample(hdmap: list[str]) -> str | None: + """Return the sample id ``--hdmap`` named, or ``None`` if it named videos. + + A sample is named the way it is listed, by a bare id with no file extension + on it. Telling them apart by extension rather than by what exists means a + misspelled video is reported as the missing video it is, instead of being + looked for in a dataset it was never in. A file that is sitting right there + still wins, extension or not. + """ + if not hdmap: + return DEFAULT_HDMAP_SAMPLE + if len(hdmap) > 1: + return None + lone = Path(hdmap[0]) + return None if lone.suffix or lone.exists() else hdmap[0] + + +def _resolve_view_names( + view_names: tuple[str, ...], defaults: tuple[str, ...] +) -> tuple[str, ...]: + """Return the camera labels, which only a multi-camera model reads. + + Args: + view_names: What the command line asked for, if anything. + defaults: A label per camera to fall back on, which is also how many + cameras this run has. + + Raises: + ValueError: Some cameras were labelled and others were not. + """ + if not view_names: + return defaults + if len(view_names) != len(defaults): + raise ValueError( + f"Got {len(defaults)} camera(s) and {len(view_names)} view " + "name(s): pass one name per camera, or none at all." + ) + return view_names diff --git a/integrations_v2/omnidreams/omnidreams_v2/conditioning.py b/integrations_v2/omnidreams/omnidreams_v2/conditioning.py new file mode 100644 index 000000000..1d21ec0bf --- /dev/null +++ b/integrations_v2/omnidreams/omnidreams_v2/conditioning.py @@ -0,0 +1,351 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Where a run's HDMap conditioning comes from.""" + +from pathlib import Path +from typing import Protocol + +import torch +from torch import Tensor + +from flashdreams.infra.runner_io import load_first_frame_tensor, load_video_tensor +from flashdreams.runtime_v2.user_input_events import UserInputEvents + +_PIXEL_DTYPE = torch.bfloat16 +"""What the model reads its conditioning pixels as.""" + + +class HDMapSource(Protocol): + """One run's worth of HDMap conditioning, a chunk at a time. + + This model generates from rendered road layout rather than from a prompt + alone, so every step needs a chunk of HDMap pixels to condition on. A + recording has them already; a renderer draws them from wherever the driver + steered. Both answer the same questions, so a session does not know which + of them it was given. + + Not a client window and not an :class:`InputSource`: this produces model + input at the rate the model generates, and has to know how long the chunk + it is being asked for is. A window supplies what a person did, which is a + different thing arriving at a different rate. A source that renders is + given those events and turns them into geometry. + """ + + @property + def view_names(self) -> tuple[str, ...]: + """Cameras this supplies, in the order its chunks stack them.""" + ... + + def open(self) -> None: + """Load or start whatever produces the frames. + + Called once, from ``ISession.init``, so the cost of reading a video or + starting a renderer lands with the rest of a run's startup. + """ + ... + + def first_frame(self) -> Tensor: + """Return the frame the run continues from. + + Returns: + Pixels as ``[B, V, 1, 3, H, W]`` in ``[-1, 1]``. + """ + ... + + def has_frames(self, frame_count: int) -> bool: + """Report whether a chunk of ``frame_count`` more frames can be produced. + + Asked before each step, so a run ends on the boundary rather than on a + half-conditioned chunk. A renderer with no end to it always says yes. + """ + ... + + def next_chunk(self, frame_count: int, events: UserInputEvents) -> Tensor: + """Return the next chunk, advancing this source past it. + + Args: + frame_count: Frames the model is about to generate, which is how + many frames of conditioning it needs. + events: Input since the previous step. Ignored by a source that is + replaying something already recorded. + + Returns: + Pixels as ``[B, V, T, 3, H, W]`` in ``[-1, 1]``, with ``T`` equal to + ``frame_count``. + """ + ... + + def reset(self) -> None: + """Start this source over, so the run can generate from the top.""" + ... + + def close(self) -> None: + """Release whatever this holds.""" + ... + + +class SceneRenderer(Protocol): + """Frames of road layout, drawn on demand along a scene's timeline. + + The part of a rendered run that needs a GPU and a scene on disk, kept + behind a seam so the bookkeeping around it -- how far through the drive a + run is, and what the model wants its pixels to look like -- can be tested + without either. + """ + + @property + def frame_count(self) -> int: + """Frames this can draw, which is how long a run against it lasts.""" + ... + + def open(self) -> None: + """Load the scene and get the rasterizer ready to draw it.""" + ... + + def render(self, start: int, count: int) -> Tensor: + """Draw ``count`` frames of the timeline, beginning at ``start``. + + Returns: + Pixels as ``[T, H, W, 3]``, RGB bytes. + """ + ... + + def close(self) -> None: + """Release the scene and the rasterizer.""" + ... + + +class PrecomputedHDMapSource: + """HDMap conditioning read from recorded video, one file per camera. + + What a reproducible run uses: the same clip conditions every run, so two + runs of one command are comparable. The recording is also what says how + long the run is, since generating past the end of it would have nothing to + condition on. + """ + + def __init__( + self, + *, + hdmap_video_paths: tuple[Path, ...], + first_frame_paths: tuple[Path, ...], + view_names: tuple[str, ...], + pixel_width: int, + pixel_height: int, + device: str, + ) -> None: + """ + Args: + hdmap_video_paths: HDMap video per camera, in camera order. + first_frame_paths: Frame to continue from per camera, in the same + order. An image or a video, of which frame zero is taken. + view_names: Camera labels, in the same order. + pixel_width: Width to resize the conditioning to, which is the + width the run generates at. + pixel_height: Height to resize it to. + device: Device to load onto, alongside the model. + """ + self._hdmap_video_paths = hdmap_video_paths + self._first_frame_paths = first_frame_paths + self._view_names = view_names + self._pixel_width = pixel_width + self._pixel_height = pixel_height + self._device = torch.device(device) + self._hdmap: Tensor | None = None + self._cursor = 0 + + @property + def view_names(self) -> tuple[str, ...]: + return self._view_names + + def open(self) -> None: + """Decode every camera's HDMap video into one tensor. + + The whole clip is held, rather than decoded a chunk at a time, because + a reset has to replay it and decoding is the slow part. + """ + videos = [self._load_video(path) for path in self._hdmap_video_paths] + # [T, C, H, W] per camera, stacked into [B=1, V, T, C, H, W]. + self._hdmap = torch.stack(videos, dim=0).unsqueeze(0) + + def first_frame(self) -> Tensor: + frames = [ + load_first_frame_tensor( + path, + pixel_height=self._pixel_height, + pixel_width=self._pixel_width, + device=self._device, + dtype=_PIXEL_DTYPE, + allow_video=True, + ) + for path in self._first_frame_paths + ] + return torch.stack(frames, dim=0).unsqueeze(0) + + def has_frames(self, frame_count: int) -> bool: + return self._cursor + frame_count <= self._frames().shape[2] + + def next_chunk(self, frame_count: int, events: UserInputEvents) -> Tensor: + """Return the next ``frame_count`` frames of the recording. + + Args: + frame_count: Frames to return. + events: Ignored. A recording plays the drive it recorded, whatever + anyone does while watching it. + + Raises: + RuntimeError: The recording has fewer frames left than that, which + :meth:`has_frames` is asked in order to avoid. + """ + del events + hdmap = self._frames() + end = self._cursor + frame_count + if end > hdmap.shape[2]: + raise RuntimeError( + f"Asked for frames [{self._cursor}, {end}) of an HDMap " + f"recording {hdmap.shape[2]} frames long." + ) + chunk = hdmap[:, :, self._cursor : end] + self._cursor = end + return chunk + + def reset(self) -> None: + """Rewind to the start of the recording.""" + self._cursor = 0 + + def close(self) -> None: + self._hdmap = None + + def _frames(self) -> Tensor: + """Return the loaded recording. + + Raises: + RuntimeError: This was never opened, or has been closed. + """ + if self._hdmap is None: + raise RuntimeError( + f"{type(self).__name__}.open() must run before the recording " + "is read, and it cannot be read after close()." + ) + return self._hdmap + + def _load_video(self, path: Path) -> Tensor: + """Load and resize one camera's HDMap video to ``[T, C, H, W]``.""" + return load_video_tensor( + path, + pixel_height=self._pixel_height, + pixel_width=self._pixel_width, + device=self._device, + dtype=_PIXEL_DTYPE, + ) + + +class RenderedHDMapSource: + """HDMap conditioning drawn a chunk at a time by a renderer. + + The reason for drawing rather than replaying is that a drive can then go + somewhere the recording never went. Nothing steers yet, so what this + produces is the drive the scene recorded -- the same road as the matching + recording, which is what makes the two comparable while this is new. + + One camera, because the renderer draws one. A multi-camera run would stack + what several renderers drew, which is where this and the renderer seam grow + together rather than here alone. + """ + + def __init__( + self, + *, + renderer: SceneRenderer, + first_frame_path: Path, + view_name: str, + pixel_width: int, + pixel_height: int, + device: str, + ) -> None: + """ + Args: + renderer: Draws the road layout. Owned by this source, which opens + and closes it with the run. + first_frame_path: Frame the run continues from, which for a scene is + the recorded capture the drawn layout starts at. + view_name: Camera label, as the model spells it. + pixel_width: Width the run generates at, which the layout is drawn + at so the two line up. + pixel_height: Height the run generates at. + device: Device to hand the model its pixels on. + """ + self._renderer = renderer + self._first_frame_path = first_frame_path + self._view_name = view_name + self._pixel_width = pixel_width + self._pixel_height = pixel_height + self._device = torch.device(device) + self._cursor = 0 + + @property + def view_names(self) -> tuple[str, ...]: + return (self._view_name,) + + def open(self) -> None: + """Load the scene, which is where the run's startup cost mostly is.""" + self._renderer.open() + + def first_frame(self) -> Tensor: + frame = load_first_frame_tensor( + self._first_frame_path, + pixel_height=self._pixel_height, + pixel_width=self._pixel_width, + device=self._device, + dtype=_PIXEL_DTYPE, + allow_video=True, + ) + return frame.unsqueeze(0).unsqueeze(0) + + def has_frames(self, frame_count: int) -> bool: + return self._cursor + frame_count <= self._renderer.frame_count + + def next_chunk(self, frame_count: int, events: UserInputEvents) -> Tensor: + """Draw the next ``frame_count`` frames of road layout. + + Args: + frame_count: Frames to draw. + events: Ignored. This is where steering would enter, turning what a + driver did into where the next chunk is drawn from; until then a + rendered run follows the drive the scene recorded. + + Raises: + RuntimeError: The scene has fewer frames left than that, which + :meth:`has_frames` is asked in order to avoid. + """ + del events + available = self._renderer.frame_count + end = self._cursor + frame_count + if end > available: + raise RuntimeError( + f"Asked for frames [{self._cursor}, {end}) of a scene " + f"{available} frames long." + ) + frames = self._renderer.render(self._cursor, frame_count) + self._cursor = end + return _to_model_pixels(frames, device=self._device).unsqueeze(0).unsqueeze(0) + + def reset(self) -> None: + """Return to the start of the drive.""" + self._cursor = 0 + + def close(self) -> None: + self._renderer.close() + + +def _to_model_pixels(frames: Tensor, *, device: torch.device) -> Tensor: + """Convert drawn ``[T, H, W, 3]`` RGB bytes to model pixels ``[T, 3, H, W]``. + + A rasterizer hands back bytes in the layout an image viewer wants. The model + reads signed pixels in the layout a convolution wants, which is what the + recorded path's video loader already returns, so the conversion the two + paths share happens here. + """ + pixels = frames.to(device=device, dtype=_PIXEL_DTYPE) + return pixels.permute(0, 3, 1, 2) / 127.5 - 1.0 diff --git a/integrations_v2/omnidreams/omnidreams_v2/ludus.py b/integrations_v2/omnidreams/omnidreams_v2/ludus.py new file mode 100644 index 000000000..1db95047f --- /dev/null +++ b/integrations_v2/omnidreams/omnidreams_v2/ludus.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Road layout drawn by the Ludus rasterizer, from a scene on disk.""" + +from pathlib import Path +from typing import Any + +import torch +from torch import Tensor + + +class LudusSceneRenderer: + """Draws one camera's view of a scene's road layout, a chunk at a time. + + A recording of the layout fixes the drive before the run starts. Drawing it + does not, which is the point of this even while nothing steers: the geometry + is produced beside the frames it conditions, so it can eventually follow + wherever a driver goes. Until then the poses come from the drive the scene + recorded, which makes a drawn run and the recording of the same scene two + views of one drive. + + Everything Ludus is imported when the scene is loaded rather than when this + module is, so a run that never draws anything -- and a test on a machine + with no GPU -- does not pay for a rasterizer it will not use. + """ + + def __init__( + self, + *, + scene_path: Path, + camera: str, + view_start_us: int, + frames_per_second: int, + pixel_width: int, + pixel_height: int, + device: str, + ) -> None: + """ + Args: + scene_path: Scene archive to draw, holding the road layout and the + drive recorded along it. + camera: Camera to draw from, spelled the way the scene spells it. + view_start_us: Timestamp to start drawing at, which is when the + frame the run continues from was captured. Starting anywhere + else would show the model a road it is not looking at. + frames_per_second: Rate to draw at, matching the rate the model + generates at. The recorded drive is resampled onto it. + pixel_width: Width to draw at, which is the width the run + generates at. + pixel_height: Height to draw at. + device: Device holding the scene and the drawn frames. + """ + self._scene_path = scene_path + self._camera = camera + self._view_start_us = view_start_us + self._frames_per_second = frames_per_second + self._pixel_width = pixel_width + self._pixel_height = pixel_height + self._device = torch.device(device) + self._context: Any = None + self._scene_id: int | None = None + self._camera_id: int | None = None + self._camera_type: int | None = None + self._ego_tracks: Any = None + self._sensor_to_rig: Tensor | None = None + self._timestamps: Tensor | None = None + + @property + def frame_count(self) -> int: + """Frames of recorded drive there are to draw. + + Zero until the scene is loaded, since the length of the drive is + something the scene says. + """ + return 0 if self._timestamps is None else int(self._timestamps.shape[0]) + + def open(self) -> None: + """Load the scene, upload it, and lay out the timeline to draw along. + + Raises: + ValueError: The scene has no such camera, or no drive left after the + frame the run continues from. + """ + from ludus_renderer import LudusCudaTimestampedContext, load_scene + from ludus_renderer.render_utils import SceneAdapter + from ludus_renderer.torch.ops import CAMERA_TYPE_REGULAR + + scene = load_scene( + self._scene_path, + device=self._device, + # Drawn at the size the run generates at, which scales the camera's + # intrinsics with it. Drawing at a size the intrinsics disagree with + # puts the road in the wrong place rather than failing. + target_resolution=(self._pixel_width, self._pixel_height), + # The ego trajectory is drawn as a path along the road and the ego + # vehicle as a box in front of the camera. Neither is part of the + # layout this model was conditioned on. + include_ego_trajectory=False, + include_ego_obstacle=False, + ) + if self._camera not in scene.camera_name_to_id: + available = ", ".join(sorted(scene.camera_name_to_id)) + raise ValueError( + f"Scene {self._scene_path} has no camera {self._camera!r}. " + f"It has: {available}." + ) + + self._timestamps = self._timeline(scene.ego_track.timestamps) + self._ego_tracks = SceneAdapter(scene).ego_tracks + self._sensor_to_rig = scene.sensor_to_rig[self._camera].to(self._device) + self._camera_type = CAMERA_TYPE_REGULAR + + context = LudusCudaTimestampedContext(device=self._device) + # Every camera, so a camera's own index is what identifies it, which is + # what the scene's mapping already holds. + context.upload_cameras(list(scene.cameras)) + self._camera_id = scene.camera_name_to_id[self._camera] + self._scene_id = context.upload_scene(scene.timestamped_scene) + self._context = context + + def render(self, start: int, count: int) -> Tensor: + """Draw ``count`` frames of the drive, beginning at ``start``. + + Returns: + Pixels as ``[T, H, W, 3]``, RGB bytes on this renderer's device. + + Raises: + RuntimeError: :meth:`open` has not run yet. + """ + context = self._context + timestamps = self._timestamps + if context is None or timestamps is None: + raise RuntimeError( + f"{type(self).__name__}.open() must run before render()." + ) + chunk = timestamps[start : start + count] + images = context.render_uniform( + scene_id=self._scene_id, + camera_id=self._camera_id, + timestamps_us=chunk, + camera_type_id=self._camera_type, + camera_poses=self._camera_poses(chunk), + resolution=(self._pixel_height, self._pixel_width), + ) + # The rasterizer draws onto an opaque background, so the alpha it + # reports is of no use to a model reading three channels. + frames = images[:, :, :, :3] + if context.needs_vflip: + frames = frames.flip(1) + return frames.contiguous() + + def close(self) -> None: + """Drop the scene, and the device memory it was holding.""" + context = self._context + self._context = None + self._timestamps = None + self._ego_tracks = None + self._sensor_to_rig = None + if context is not None: + context.clear_scenes() + + def _timeline(self, recorded_us: Tensor) -> Tensor: + """Return evenly spaced timestamps to draw the recorded drive along. + + The drive is recorded at whatever rate its logger ran at, which is not + the rate the model generates at. Poses in between are interpolated, so + the timeline this returns is the generated rate and the recording is + read at whatever offsets that lands on. + + Raises: + ValueError: The drive ends at or before the frame the run continues + from, leaving nothing to draw. + """ + step_us = round(1_000_000 / self._frames_per_second) + # Clamped rather than trusted: the frame a run continues from and the + # recorded drive are two things the scene has to agree with itself + # about, and a scene that does not would otherwise draw an empty road. + start_us = max(int(recorded_us[0].item()), self._view_start_us) + end_us = int(recorded_us[-1].item()) + if end_us <= start_us: + raise ValueError( + f"Scene {self._scene_path} records a drive up to {end_us}us, " + f"which is not past {start_us}us, where the frame the run " + "continues from was captured. There is nothing to draw." + ) + return torch.arange(start_us, end_us + 1, step_us, dtype=torch.int64) + + def _camera_poses(self, timestamps: Tensor) -> Tensor: + """Return where the camera was at each timestamp, as Ludus wants it. + + Ludus draws from a world-to-camera transform, while the recorded drive + says where the vehicle was; the camera's mounting is what sits between + them. + """ + # [T, 1, 4, 4], one pose per timestamp, which Ludus wants as [T, 4, 4]. + rig_to_world = self._ego_tracks.get_transforms_at_timestamp(timestamps)[:, 0] + camera_to_world = torch.einsum("nij,jk->nik", rig_to_world, self._sensor_to_rig) + return torch.linalg.inv(camera_to_world) diff --git a/integrations_v2/omnidreams/omnidreams_v2/samples.py b/integrations_v2/omnidreams/omnidreams_v2/samples.py new file mode 100644 index 000000000..579e3335b --- /dev/null +++ b/integrations_v2/omnidreams/omnidreams_v2/samples.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bundled sample drives, fetched from Hugging Face when a run asks for one.""" + +from pathlib import Path + +from huggingface_hub import HfApi, hf_hub_download +from huggingface_hub.hf_api import RepoFile +from omnidreams.runner import ( + DEFAULT_EXAMPLE_DATA_UUID_1V, + EXAMPLE_DATA_HF_BROWSER_URL, + EXAMPLE_DATA_HF_REPO, +) + +DEFAULT_HDMAP_SAMPLE = DEFAULT_EXAMPLE_DATA_UUID_1V +"""Recording a run gets when it asks to replay one without naming it.""" + +_RECORDING_SUFFIX = "_hdmap.mp4" +"""What the HDMap recording in a sample directory ends with.""" + +_FIRST_FRAME_NAME = "first_frame.png" +"""What the frame to continue from is called in a sample directory.""" + + +def fetch_hdmap_sample(sample_id: str) -> tuple[Path, Path]: + """Download one bundled single-camera HDMap recording. + + Args: + sample_id: Directory under ``data/single_view`` in the samples dataset, + which is a clip UUID. + + Returns: + The HDMap recording, then the frame to continue from. A sample carries + both, which is why replaying one takes no other arguments. They land in + the Hugging Face cache, so asking again costs nothing. + + Raises: + FileNotFoundError: The sample has no recording in it, or more than one, + so which to drive through is not clear. + """ + directory = f"data/single_view/{sample_id}" + entries = HfApi().list_repo_tree( + repo_id=EXAMPLE_DATA_HF_REPO, + repo_type="dataset", + path_in_repo=directory, + recursive=False, + ) + # The recording is named after the clip rather than predictably, so the + # directory is listed rather than the filename built from the drive id. + recordings = [ + entry.path + for entry in entries + if isinstance(entry, RepoFile) and entry.path.endswith(_RECORDING_SUFFIX) + ] + if len(recordings) != 1: + found = ", ".join(recordings) if recordings else "none" + raise FileNotFoundError( + f"Expected one '*{_RECORDING_SUFFIX}' in {directory} of " + f"{EXAMPLE_DATA_HF_REPO}, found {found}. Samples are listed at " + f"{EXAMPLE_DATA_HF_BROWSER_URL}." + ) + return ( + _download(recordings[0]), + _download(f"{directory}/{_FIRST_FRAME_NAME}"), + ) + + +def _download(path_in_repo: str) -> Path: + """Return one file of the samples dataset, from the cache or from the hub.""" + return Path( + hf_hub_download( + repo_id=EXAMPLE_DATA_HF_REPO, + repo_type="dataset", + filename=path_in_repo, + ) + ) diff --git a/integrations_v2/omnidreams/omnidreams_v2/scenes.py b/integrations_v2/omnidreams/omnidreams_v2/scenes.py new file mode 100644 index 000000000..d48ebca3d --- /dev/null +++ b/integrations_v2/omnidreams/omnidreams_v2/scenes.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Scenes: the road a drawn run drives along, and what it starts from.""" + +import zipfile +from pathlib import Path + +from omnidreams.interactive_drive.math3d import normalize_camera_name +from omnidreams.scenes import ( + SCENE_FRAME_SUFFIXES, + SCENE_FRAMES_DIRNAME, + hf_hub_download_scene, + hf_scenes_repo_id, + scenes_cache_root, +) + +DEFAULT_SCENE = "0d404ff7-2b66-498c-b047-1ed8cded60d4" +"""Scene a run drives when it asks for one without naming it.""" + +DEFAULT_SCENE_CAMERA = "camera:front:wide:120fov" +"""Camera to draw from, being the one the shipped single-view model reads.""" + +_PROMPT_ENTRIES = ("prompt.txt", "prompt1.txt", "prompt_1.txt") +"""Where a scene keeps its own description of the drive, best first.""" + +_SEED_FRAME_DIRNAME = "v2-seed-frames" +"""Where frames unpacked from scene archives are kept, under the scene cache.""" + + +def fetch_scene(scene: str) -> Path: + """Return the local archive for a scene, downloading it if needed. + + Args: + scene: A path to a scene archive, or the id of one to download. + + Returns: + The archive on disk. A downloaded scene lands in the Hugging Face + cache, so asking for the same one again costs nothing. + + Raises: + FileNotFoundError: The scene is neither a path that exists nor an id + the scenes dataset has. + """ + named = Path(scene) + if named.exists(): + return named + try: + return hf_hub_download_scene(scene) + except Exception as exc: + raise FileNotFoundError( + f"No scene {scene!r}: there is no such path, and " + f"{hf_scenes_repo_id()} has no such scene ({exc})." + ) from exc + + +def read_seed_frame(scene_path: Path, *, camera: str) -> tuple[Path, int]: + """Unpack the earliest recorded frame for one camera of a scene. + + This is the frame a drawn run continues from. Its timestamp comes back with + it because that is also where the road layout has to start being drawn: a + run that continued from one moment while being shown the road at another + would be asked to drive a corner it cannot see. + + Args: + scene_path: Scene archive to read. + camera: Camera whose recording to take, in either spelling. + + Returns: + The unpacked image, and when it was captured, in microseconds. + + Raises: + FileNotFoundError: The scene has no timestamped frames for that camera, + so there is nothing to continue from. + """ + _, logical_name = normalize_camera_name(camera) + with zipfile.ZipFile(scene_path) as archive: + entry, captured_us = _earliest_frame(archive, scene_path, camera) + filename = f"{scene_path.stem}-{logical_name}-{Path(entry).name}" + unpacked = scenes_cache_root() / _SEED_FRAME_DIRNAME / filename + if not unpacked.exists(): + unpacked.parent.mkdir(parents=True, exist_ok=True) + # Written beside the cached scene rather than to a temporary + # directory, so a second run of the same command reuses it. + unpacked.write_bytes(archive.read(entry)) + return unpacked, captured_us + + +def read_prompt(scene_path: Path) -> str | None: + """Return a scene's own description of the drive, or ``None`` if it has none. + + Preferred over the model's generic prompt, since a description of this road + in this weather is what the drive being drawn actually looks like. + """ + with zipfile.ZipFile(scene_path) as archive: + entries = set(archive.namelist()) + for candidate in _PROMPT_ENTRIES: + if candidate in entries: + return archive.read(candidate).decode("utf-8").strip() or None + return None + + +def _earliest_frame( + archive: zipfile.ZipFile, scene_path: Path, camera: str +) -> tuple[str, int]: + """Return the first recorded frame for a camera, and its timestamp. + + Frames are kept as ``frames//.jpeg``, so the earliest + is the smallest of those names. Both spellings of a camera are looked for, + since scenes have been staged with each. + + Raises: + FileNotFoundError: No frame of that camera is named after a timestamp. + """ + clipgt_name, logical_name = normalize_camera_name(camera) + prefixes = tuple( + f"{SCENE_FRAMES_DIRNAME}/{name}/" + for name in {camera, clipgt_name, logical_name} + ) + captured: list[tuple[int, str]] = [] + for entry in archive.namelist(): + name = Path(entry) + if not entry.startswith(prefixes): + continue + if name.suffix.lower() not in SCENE_FRAME_SUFFIXES: + continue + if name.stem.isdigit(): + captured.append((int(name.stem), entry)) + if not captured: + raise FileNotFoundError( + f"Scene {scene_path} has no timestamped frames for camera " + f"{camera!r} under {SCENE_FRAMES_DIRNAME}/, so there is no " + "recorded frame for a run to continue from." + ) + captured_us, entry = min(captured) + return entry, captured_us diff --git a/integrations_v2/omnidreams/omnidreams_v2/session.py b/integrations_v2/omnidreams/omnidreams_v2/session.py new file mode 100644 index 000000000..5b059eab4 --- /dev/null +++ b/integrations_v2/omnidreams/omnidreams_v2/session.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""One Omnidreams rollout: road layout in, driving video out.""" + +from typing import Any + +from flashdreams.api_v2.session import ISession +from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.user_input_events import UserInputEvents + +from .conditioning import HDMapSource + + +class OmnidreamsSession(ISession): + """One HDMap-conditioned rollout, continuing from a first frame. + + A step is one autoregressive block, conditioned on the chunk of HDMap + pixels covering the frames it is about to generate. The source supplies + that chunk, and how it came by it is its own business: a run replaying a + recording and a run rendering what a driver steered differ only there. It + is also what says how long the run is, since generating past the end of + the road would have nothing to condition on. + + What belongs to a run is the cache this initializes, holding the encoded + prompt and first frame alongside the attention state, and the source, which + is positioned partway through a drive. The pipeline belongs to the + application. + """ + + def __init__( + self, + pipeline: Any, + prompt: str, + source: HDMapSource, + session_desc: SessionDesc, + max_blocks: int | None = None, + ) -> None: + """ + Args: + pipeline: Loaded pipeline, owned by the application. + prompt: Text describing the drive, applied to every camera. + source: Conditioning for this run, owned by this session. + session_desc: Session the application accepted, already checked + against what the model can produce. + max_blocks: Blocks to stop after, for a run shorter than the + conditioning. The default drives until the source runs out, + which for a recording is its end and for a renderer is never. + """ + self._pipeline = pipeline + self._prompt = prompt + self._source = source + self._session_desc = session_desc + self._max_blocks = max_blocks + self._blocks_generated = 0 + self._cache: Any = None + + def init(self) -> None: + """Open the conditioning, then encode the prompt and the first frame. + + Where the text encoder runs and the recording is decoded, so it is far + slower than a step. + """ + self._source.open() + self._cache = self._new_cache() + + @property + def session_desc(self) -> SessionDesc: + return self._session_desc + + def step(self, step_index: int, events: UserInputEvents) -> StepResult: + """Generate the next block of frames from the next chunk of HDMap. + + Args: + step_index: Also the autoregressive index the rollout is up to. The + pipeline rejects a step out of order. + events: Passed to the source, which is what turns them into + geometry. A recording ignores them; a renderer steers by them. + + Raises: + RuntimeError: :meth:`init` has not run yet. + """ + if self._cache is None: + raise RuntimeError(f"{type(self).__name__}.init() must run before step().") + # The model decides the chunk length, and the conditioning has to cover + # exactly it. A causal decoder's first block is usually shorter. + frame_count = self._pipeline.get_num_frames(step_index) + hdmap = self._source.next_chunk(frame_count, events) + self._blocks_generated += 1 + frames = self._pipeline.generate( + autoregressive_index=step_index, cache=self._cache, hdmap=hdmap + ) + # Advancing the attention state is what makes the next step continue + # this one, and it reports what the step cost. + metrics = self._pipeline.finalize( + autoregressive_index=step_index, cache=self._cache + ) + return StepResult( + step_index=step_index, + output=frames.detach(), + frame_count=int(frames.shape[2]), + output_layout=self._session_desc.output_layout, + metrics=dict(metrics or {}), + ) + + def is_finished(self) -> bool: + """Report whether the rollout has run out of road, or of blocks. + + The road is the length of a run: a recording ends when it ends, and a + renderer never does, so a run against one lasts until its client goes + away. ``max_blocks`` stops a run before either, and is how a smoke run + asks for a few seconds of a long drive. + """ + if self._max_blocks is not None and self._blocks_generated >= self._max_blocks: + return True + return not self._source.has_frames( + self._pipeline.get_num_frames(self._blocks_generated) + ) + + def reset(self) -> None: + """Drive the same route again from the first frame. + + The cache is replaced rather than cleared, so nothing of the abandoned + run reaches the new one, and the source rewinds with it. + """ + self._blocks_generated = 0 + self._source.reset() + self._cache = self._new_cache() + + def close(self) -> None: + """Release the rollout's cache and its conditioning, leaving the model.""" + self._cache = None + self._source.close() + + def _new_cache(self) -> Any: + """Encode the prompt and first frame into a cache for one rollout. + + The one-shot encoders are left loaded, unlike a batch run that releases + them to reclaim their memory: a reset comes back here, and an encoder + released is a session that cannot start over. + """ + return self._pipeline.initialize_cache( + text=[[self._prompt] * len(self._source.view_names)], + image=self._source.first_frame(), + view_names=list(self._source.view_names), + ) diff --git a/integrations_v2/omnidreams/omnidreams_v2/tests/test_stand_in_model.py b/integrations_v2/omnidreams/omnidreams_v2/tests/test_stand_in_model.py new file mode 100644 index 000000000..2658744b4 --- /dev/null +++ b/integrations_v2/omnidreams/omnidreams_v2/tests/test_stand_in_model.py @@ -0,0 +1,468 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for the Omnidreams application, against stand-ins. + +What is particular to a model conditioned on a road layout: every step is +conditioned on exactly the frames it generates, and a run ends when the layout +does. Neither needs a checkpoint to cover. A run against the real one, which +the other v2 integrations have as ``test_real_model.py``, is not written yet: +it needs a recorded drive to condition on as well as the checkpoint. +""" + +import shutil +from pathlib import Path + +import pytest +import torch +from omnidreams.config import RUNNER_SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE as _RUNNER +from omnidreams_v2 import OmnidreamsApplication, OmnidreamsSessionConfig +from omnidreams_v2 import app as app_module +from omnidreams_v2.samples import DEFAULT_HDMAP_SAMPLE + +from flashdreams.api_v2.client_window import IClientWindow +from flashdreams.runtime_v2.application_runner import ApplicationRunner +from flashdreams.runtime_v2.mp4_client_window import Mp4ClientWindow +from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.user_input_events import UserInputEvents +from flashdreams.runtime_v2.video_tensor import VideoTensorLayout + +pytestmark = pytest.mark.ci_cpu + +_NO_EVENTS = UserInputEvents([]) +"""What a window with nobody on it reports.""" + +_CONDITIONING_ARGS = [ + "--hdmap", + "drive_hdmap.mp4", + "--first-frame", + "drive_first_frame.png", + "--device", + "cpu", +] +"""Arguments naming conditioning that a stand-in source stands in for.""" + + +def test_the_model_says_what_it_generates_without_being_told() -> None: + """The numbers are the checkpoint's, read off the runner config this + integration already ships rather than written down again.""" + app = OmnidreamsApplication(pipeline_config=FakePipelineConfig()) + + desc = app.session_desc() + + assert (desc.video_width, desc.video_height) == ( + _RUNNER.pixel_width, + _RUNNER.pixel_height, + ) + assert desc.frames_per_second_for_step == _RUNNER.output_fps + assert desc.output_layout is VideoTensorLayout.bvtchw + + +def test_every_step_is_conditioned_on_the_frames_it_is_about_to_generate() -> None: + """The chunk handed to the model covers the block it generates and no more, + and consecutive blocks read consecutive stretches of the drive.""" + pipeline = FakePipeline() + source = FakeHDMapSource(total_frames=24, pipeline=pipeline) + window = RecordingClientWindow() + + _run(pipeline, source, window) + + assert source.chunks_read == [(0, 5), (5, 13), (13, 21)] + assert pipeline.conditioned_frames == [5, 8, 8] + assert [result.frame_count for result in window.results] == [5, 8, 8] + + +def test_a_run_ends_when_the_road_does() -> None: + """Three frames are left over, which is less than a block, so the run stops + rather than generating a partly conditioned one. Nothing asked it to stop: + the recording is what says how long a drive is.""" + pipeline = FakePipeline() + source = FakeHDMapSource(total_frames=24, pipeline=pipeline) + + _run(pipeline, source, RecordingClientWindow()) + + assert source.frames_left == 3 + + +def test_a_run_can_be_cut_short_of_the_road() -> None: + pipeline = FakePipeline() + source = FakeHDMapSource(total_frames=1000, pipeline=pipeline) + window = RecordingClientWindow() + + _run(pipeline, source, window, max_blocks=2) + + assert [result.frame_count for result in window.results] == [5, 8] + + +def test_a_run_told_no_length_produces_about_a_minute() -> None: + """Rather than driving to the end of the road, which on a long scene is an + hour of generating for someone who only asked to see it work. It stops on + the block that reaches a minute, a block being the smallest thing a run can + produce.""" + pipeline = FakePipeline() + source = FakeHDMapSource(total_frames=100_000, pipeline=pipeline) + window = RecordingClientWindow() + + _run(pipeline, source, window) + + a_minute = _RUNNER.output_fps * 60 + generated = sum(result.frame_count for result in window.results) + assert generated >= a_minute + assert generated - window.results[-1].frame_count < a_minute + + +def test_a_run_can_be_told_to_drive_to_the_end_of_the_road() -> None: + """Which is what a session someone is sitting in front of wants, and what + the default minute would otherwise cut off. This road is longer than a + minute, so the two answers are told apart.""" + pipeline = FakePipeline() + source = FakeHDMapSource(total_frames=2000, pipeline=pipeline) + + _run(pipeline, source, RecordingClientWindow(), max_blocks=0) + + assert source.frames_left == 3 + + +def test_a_reset_drives_the_same_route_again_from_the_start() -> None: + """Both halves of a run start over: the cache the model generates from, and + the source's place in the drive.""" + pipeline = FakePipeline() + source = FakeHDMapSource(total_frames=1000, pipeline=pipeline) + app = _application(pipeline, source) + app.init(_CONDITIONING_ARGS) + session = app.create_session(_stand_in_session_desc(pipeline)) + session.init() + session.step(0, _NO_EVENTS) + + session.reset() + session.step(0, _NO_EVENTS) + + assert source.chunks_read == [(0, 5), (0, 5)] + assert len(pipeline.caches) == 2 + + +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + (["--hdmap", "a.mp4"], "--first-frame is required"), + ( + ["--hdmap", "a.mp4", "b.mp4", "--first-frame", "a.png"], + "one of each per camera", + ), + ( + ["--hdmap", "5f0e1a2b-drive", "--first-frame", "a.png"], + "cannot be combined with --first-frame", + ), + ( + [*_CONDITIONING_ARGS, "--view-names", "left", "right"], + "one name per camera", + ), + ([*_CONDITIONING_ARGS, "--max-blocks", "-1"], "cannot be negative"), + ], +) +def test_a_run_that_describes_its_drive_incompletely_is_refused( + arguments: list[str], expected: str +) -> None: + """Before the model loads, since a checkpoint of several gigabytes is a long + wait for a typo.""" + app = OmnidreamsApplication(pipeline_config=FakePipelineConfig()) + + with pytest.raises(ValueError, match=expected): + app.init(arguments) + + +@pytest.mark.parametrize( + ("arguments", "expected_sample"), + [ + (["--hdmap"], DEFAULT_HDMAP_SAMPLE), + (["--hdmap", "5f0e1a2b-drive"], "5f0e1a2b-drive"), + ], +) +def test_a_run_that_asks_to_replay_without_naming_what_gets_a_sample( + monkeypatch: pytest.MonkeyPatch, arguments: list[str], expected_sample: str +) -> None: + """Bare ``--hdmap`` is how you replay something without having a recording + to hand, and a bare id is how you replay one of the listed samples: neither + has a file extension on it, which is what tells them from a video. The + download itself is stubbed, since a unit test that reaches Hugging Face is + not one.""" + recording, first_frame = Path("road_hdmap.mp4"), Path("road_first_frame.png") + requested: list[str] = [] + + def fetch(sample_id: str) -> tuple[Path, Path]: + requested.append(sample_id) + return recording, first_frame + + monkeypatch.setattr(app_module, "fetch_hdmap_sample", fetch) + pipeline = FakePipeline() + configs: list[OmnidreamsSessionConfig] = [] + + def factory( + config: OmnidreamsSessionConfig, session_desc: SessionDesc + ) -> FakeHDMapSource: + del session_desc + configs.append(config) + return FakeHDMapSource(total_frames=1000, pipeline=pipeline) + + app = OmnidreamsApplication( + pipeline_config=FakePipelineConfig(pipeline), source_factory=factory + ) + app.init([*arguments, "--device", "cpu"]) + app.create_session(_stand_in_session_desc(pipeline)) + + assert requested == [expected_sample] + assert configs[0].hdmap_video_paths == (recording,) + assert configs[0].first_frame_paths == (first_frame,) + + +def test_compilation_can_be_turned_off_for_a_run() -> None: + """Run against the real config rather than a stand-in, since what this + covers is the override landing where this model keeps the setting. No model + is loaded to answer it.""" + app = OmnidreamsApplication() + + app.init([*_CONDITIONING_ARGS, "--no-compile"]) + + assert app.pipeline_config.diffusion_model.transformer.compile_network is False + + +@pytest.mark.skipif( + shutil.which("ffmpeg") is None, reason="writing an MP4 needs ffmpeg on PATH" +) +def test_a_run_writes_every_generated_frame_to_an_mp4(tmp_path: Path) -> None: + """This model emits a camera dimension the others do not, so the run to a + file is worth covering rather than assuming.""" + pipeline = FakePipeline() + source = FakeHDMapSource(total_frames=1000, pipeline=pipeline) + path = tmp_path / "drive.mp4" + + _run(pipeline, source, Mp4ClientWindow(path), max_blocks=3) + + assert path.stat().st_size > 0 + + +## Stand-ins + + +class FakePipeline: + """A model's worth of behaviour, without a model. + + Generates frames of the shape and range the real pipeline does, including + the camera dimension, so the seam a checkpoint plugs into is covered on a + CPU. Every call is recorded, so a test can assert the rollout was driven in + order and conditioned on the right frames. + """ + + def __init__( + self, + *, + width: int = 128, + height: int = 64, + compression_ratio: int = 8, + first_block_frames: int = 5, + block_frames: int = 8, + ) -> None: + """ + Args: + width: Frame width to generate. Not square by default, so a + transposed frame cannot pass unnoticed. + height: Frame height to generate. + compression_ratio: Pixels one latent covers in each direction. + first_block_frames: Frames the first block decodes, which a causal + decoder has fewer of than the rest. + block_frames: Frames every block after the first decodes. + """ + self.decoder = FakeDecoder(compression_ratio) + self.width = width + self.height = height + self.first_block_frames = first_block_frames + self.block_frames = block_frames + self.device: str | None = None + self.caches: list[dict[str, object]] = [] + self.generated: list[int] = [] + self.conditioned_frames: list[int] = [] + self.closed = False + self._frames_generated = 0 + + def to(self, device: str) -> "FakePipeline": + self.device = device + return self + + def eval(self) -> "FakePipeline": + return self + + def initialize_cache(self, **kwargs: object) -> object: + self.caches.append(kwargs) + self._frames_generated = 0 + return object() + + def get_num_frames(self, autoregressive_index: int) -> int: + if autoregressive_index == 0: + return self.first_block_frames + return self.block_frames + + def generate( + self, *, autoregressive_index: int, cache: object, hdmap: torch.Tensor + ) -> torch.Tensor: + del cache + self.generated.append(autoregressive_index) + self.conditioned_frames.append(int(hdmap.shape[2])) + count = self.get_num_frames(autoregressive_index) + frames = torch.stack( + [self._frame(self._frames_generated + index) for index in range(count)] + ) + self._frames_generated += count + # [T, C, H, W] as one camera of one batch: [B, V, T, C, H, W]. + return frames.unsqueeze(0).unsqueeze(0) + + def finalize(self, *, autoregressive_index: int, cache: object) -> dict[str, float]: + del autoregressive_index, cache + return {"total_ms": 1.5} + + def close(self) -> None: + self.closed = True + + def _frame(self, frame_index: int) -> torch.Tensor: + """Return a grey frame whose shade moves with time, so a check made of a + real video is meaningful here too.""" + shade = -0.5 + (frame_index % 8) / 8.0 + return torch.full((3, self.height, self.width), shade, dtype=torch.float32) + + +class FakePipelineConfig: + """A pipeline config that builds a stand-in rather than loading a model.""" + + def __init__(self, pipeline: FakePipeline | None = None) -> None: + self.pipeline = pipeline if pipeline is not None else FakePipeline() + + def setup(self) -> FakePipeline: + return self.pipeline + + +class FakeDecoder: + """The one thing an application asks a decoder for.""" + + def __init__(self, spatial_compression_ratio: int) -> None: + self.spatial_compression_ratio = spatial_compression_ratio + + +class FakeHDMapSource: + """A drive of a known length, without a recording of one.""" + + def __init__( + self, + *, + total_frames: int, + pipeline: FakePipeline, + view_names: tuple[str, ...] = ("view_0",), + ) -> None: + """ + Args: + total_frames: Frames of conditioning this drive has in it. + pipeline: Stand-in whose frame size the chunks match. + view_names: Cameras this supplies. + """ + self._total_frames = total_frames + self._pipeline = pipeline + self._view_names = view_names + self._cursor = 0 + self.chunks_read: list[tuple[int, int]] = [] + self.opened = 0 + self.closed = 0 + + @property + def view_names(self) -> tuple[str, ...]: + return self._view_names + + @property + def frames_left(self) -> int: + """Conditioning the run never reached.""" + return self._total_frames - self._cursor + + def open(self) -> None: + self.opened += 1 + + def first_frame(self) -> torch.Tensor: + return torch.zeros( + 1, len(self._view_names), 1, 3, self._pipeline.height, self._pipeline.width + ) + + def has_frames(self, frame_count: int) -> bool: + return self._cursor + frame_count <= self._total_frames + + def next_chunk(self, frame_count: int, events: UserInputEvents) -> torch.Tensor: + del events + end = self._cursor + frame_count + self.chunks_read.append((self._cursor, end)) + self._cursor = end + return torch.zeros( + 1, + len(self._view_names), + frame_count, + 3, + self._pipeline.height, + self._pipeline.width, + ) + + def reset(self) -> None: + self._cursor = 0 + + def close(self) -> None: + self.closed += 1 + + +class RecordingClientWindow(IClientWindow): + """Keep what a run generated, and report no input.""" + + def __init__(self) -> None: + self.results: list[StepResult] = [] + + def get_user_input_events(self) -> UserInputEvents: + return _NO_EVENTS + + def open(self, session_desc: SessionDesc) -> None: + del session_desc + + def write(self, result: StepResult) -> None: + self.results.append(result) + + def close(self) -> None: + return + + +## Helpers + + +def _application( + pipeline: FakePipeline, source: FakeHDMapSource +) -> OmnidreamsApplication: + """Return an application over both stand-ins.""" + return OmnidreamsApplication( + pipeline_config=FakePipelineConfig(pipeline), + source_factory=lambda config, session_desc: source, + ) + + +def _stand_in_session_desc(pipeline: FakePipeline) -> SessionDesc: + """Return the session the stand-in generates, rather than the checkpoint's.""" + return SessionDesc( + output_layout=VideoTensorLayout.bvtchw, + frames_per_second_for_step=_RUNNER.output_fps, + video_width=pipeline.width, + video_height=pipeline.height, + ) + + +def _run( + pipeline: FakePipeline, + source: FakeHDMapSource, + window: IClientWindow, + *, + max_blocks: int | None = None, +) -> None: + """Run one application to completion against ``window``.""" + limit = [] if max_blocks is None else ["--max-blocks", str(max_blocks)] + ApplicationRunner(_application(pipeline, source), window).run( + _stand_in_session_desc(pipeline), [*_CONDITIONING_ARGS, *limit] + ) diff --git a/integrations_v2/omnidreams/omnidreams_v2/tests/test_stand_in_scene.py b/integrations_v2/omnidreams/omnidreams_v2/tests/test_stand_in_scene.py new file mode 100644 index 000000000..cc0fc1377 --- /dev/null +++ b/integrations_v2/omnidreams/omnidreams_v2/tests/test_stand_in_scene.py @@ -0,0 +1,424 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for a renderer rather than precomputed run. + +The rasterizer needs a GPU and a scene of real road, so what a stand-in covers +is everything around it. Three things in particular, each of which would produce +a plausible-looking but wrong drive rather than an error: that a run draws its +way along a scene consecutively, that what comes back reaches the model as +pixels of the range and layout it reads, and that the frame a run continues from +and the layout it is shown begin at the same moment. +""" + +import zipfile +from pathlib import Path + +import omnidreams.scenes +import pytest +import torch +from omnidreams_v2 import ( + LudusSceneRenderer, + OmnidreamsApplication, + OmnidreamsSessionConfig, + RenderedHDMapSource, +) +from omnidreams_v2 import app as app_module +from omnidreams_v2 import scenes as scenes_module +from omnidreams_v2.scenes import DEFAULT_SCENE +from PIL import Image + +from flashdreams.runtime_v2.user_input_events import UserInputEvents + +pytestmark = pytest.mark.ci_cpu + +_NO_EVENTS = UserInputEvents([]) +"""What a window with nobody on it reports.""" + +_SCENE_CAMERA = "camera_front_wide_120fov" +"""The camera a scene records, in the spelling the model uses for it.""" + +_CAPTURED_US = 1_700_000_000_000_000 +"""When the frame a run continues from was captured, as scenes count time.""" + + +def test_a_drawn_run_works_its_way_along_the_scene() -> None: + """Consecutively and without gaps, since a gap is a jump in the road the + model is shown while the video it generates stays continuous.""" + renderer = StubRenderer(frame_count=64) + source = _source(renderer) + source.open() + + source.next_chunk(5, _NO_EVENTS) + source.next_chunk(8, _NO_EVENTS) + source.next_chunk(8, _NO_EVENTS) + + assert renderer.drawn == [(0, 5), (5, 8), (13, 8)] + + +def test_what_is_drawn_reaches_the_model_as_pixels_it_reads() -> None: + """A rasterizer produces bytes laid out for a screen; the model reads signed + pixels laid out for a convolution. Getting this wrong leaves a run + conditioned on a washed-out or channel-swapped road.""" + renderer = StubRenderer(frame_count=8, height=4, width=6) + source = _source(renderer) + source.open() + + chunk = source.next_chunk(2, _NO_EVENTS) + + # [B, V, T, C, H, W], one camera of one batch. + assert chunk.shape == (1, 1, 2, 3, 4, 6) + assert chunk.dtype is torch.bfloat16 + # The stub draws the darkest and brightest byte it can, which are the ends + # of the range the model expects. + assert float(chunk.min()) == -1.0 + assert float(chunk.max()) == 1.0 + + +def test_a_drawn_run_ends_when_the_scene_does() -> None: + """The scene is what says how long a drawn run is, the same way a recording + does for a replayed one.""" + renderer = StubRenderer(frame_count=10) + source = _source(renderer) + source.open() + + source.next_chunk(5, _NO_EVENTS) + + assert source.has_frames(5) + assert not source.has_frames(6) + + +def test_a_drawn_run_refuses_to_draw_past_the_end_of_the_scene() -> None: + renderer = StubRenderer(frame_count=10) + source = _source(renderer) + source.open() + + with pytest.raises(RuntimeError, match="10 frames long"): + source.next_chunk(11, _NO_EVENTS) + + +def test_a_reset_returns_a_drawn_run_to_the_start_of_the_scene() -> None: + renderer = StubRenderer(frame_count=64) + source = _source(renderer) + source.open() + source.next_chunk(5, _NO_EVENTS) + + source.reset() + source.next_chunk(5, _NO_EVENTS) + + assert renderer.drawn == [(0, 5), (0, 5)] + + +def test_a_drawn_run_continues_from_the_scenes_own_recorded_frame( + tmp_path: Path, +) -> None: + """In the layout the model reads it in, which is the same one the replayed + path hands over, since it is the same model reading it.""" + source = _source(StubRenderer(frame_count=8), first_frame=_a_png(tmp_path, 8, 6)) + source.open() + + frame = source.first_frame() + + assert frame.shape == (1, 1, 1, 3, 4, 6) + + +def test_the_layout_is_drawn_from_the_moment_the_run_continues_from() -> None: + """The one alignment a drawn run cannot get wrong quietly: shown the road + from a different moment than the frame it continues from, the model would be + asked to drive a corner that is not in front of it. Reaching for the private + timeline because the alternative is covering it only on a GPU.""" + renderer = _renderer(view_start_us=100_000, frames_per_second=30) + recorded = torch.tensor([0, 500_000, 1_000_000], dtype=torch.int64) + + timeline = renderer._timeline(recorded) + + assert int(timeline[0]) == 100_000 + # Drawn at the rate the model generates at, not the rate the drive was + # recorded at, so one drawn frame is one generated frame. + assert int(timeline[1] - timeline[0]) == round(1_000_000 / 30) + assert int(timeline[-1]) <= 1_000_000 + + +def test_a_run_starting_before_its_scene_was_recorded_starts_where_it_was() -> None: + """A scene that disagrees with itself about when its drive began would + otherwise be drawn as empty road.""" + renderer = _renderer(view_start_us=0, frames_per_second=30) + recorded = torch.tensor([400_000, 900_000], dtype=torch.int64) + + timeline = renderer._timeline(recorded) + + assert int(timeline[0]) == 400_000 + + +def test_a_scene_whose_drive_ends_before_the_run_starts_is_refused() -> None: + renderer = _renderer(view_start_us=900_000, frames_per_second=30) + recorded = torch.tensor([0, 500_000], dtype=torch.int64) + + with pytest.raises(ValueError, match="nothing to draw"): + renderer._timeline(recorded) + + +def test_a_drawn_run_takes_its_prompt_and_first_frame_from_the_scene( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A scene describes the road it holds, which is a better description of the + drive than the model's own general one.""" + monkeypatch.setattr(omnidreams.scenes, "FLASHDREAMS_CACHE_DIR", tmp_path / "cache") + archive = _a_scene(tmp_path, prompt="A quiet suburban boulevard at dusk.") + app = OmnidreamsApplication(pipeline_config=object()) + + app.init(["--scene", str(archive), "--device", "cpu"]) + + config = _config_of(app) + assert config.prompt == "A quiet suburban boulevard at dusk." + assert config.scene is not None + assert config.scene.view_start_us == _CAPTURED_US + assert config.scene.first_frame_path.read_bytes() == b"a recorded frame" + # The scene names its camera, so a drawn run labels its view rather than + # leaving the placeholder a nameless recording gets. + assert config.view_names == (_SCENE_CAMERA,) + + +@pytest.mark.parametrize("arguments", [[], ["--scene", "a-named-scene"]]) +def test_a_run_that_names_no_source_draws_the_default_scene( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, arguments: list[str] +) -> None: + """Including a run given no arguments at all, so generating something takes + none. Drawing is the default because it is the source a run could one day + be steered through.""" + monkeypatch.setattr(omnidreams.scenes, "FLASHDREAMS_CACHE_DIR", tmp_path / "cache") + archive = _a_scene(tmp_path) + requested: list[str] = [] + + def fetch(scene: str) -> Path: + requested.append(scene) + return archive + + monkeypatch.setattr(app_module, "fetch_scene", fetch) + app = OmnidreamsApplication(pipeline_config=object()) + + app.init([*arguments, "--device", "cpu"]) + + expected = "a-named-scene" if arguments else DEFAULT_SCENE + assert requested == [expected] + assert _config_of(app).scene is not None + + +def test_a_prompt_on_the_command_line_beats_the_scenes_own( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(omnidreams.scenes, "FLASHDREAMS_CACHE_DIR", tmp_path / "cache") + archive = _a_scene(tmp_path, prompt="A quiet suburban boulevard at dusk.") + app = OmnidreamsApplication(pipeline_config=object()) + + app.init(["--scene", str(archive), "--prompt", "Heavy rain.", "--device", "cpu"]) + + assert _config_of(app).prompt == "Heavy rain." + + +def test_a_drawn_run_can_continue_from_a_frame_of_your_own( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Which is what drives the one road under a sky it never recorded: the + layout is still drawn, only the picture the model continues from changes.""" + monkeypatch.setattr(omnidreams.scenes, "FLASHDREAMS_CACHE_DIR", tmp_path / "cache") + archive = _a_scene(tmp_path) + mine = _a_png(tmp_path, 6, 4) + app = OmnidreamsApplication(pipeline_config=object()) + + app.init(["--scene", str(archive), "--first-frame", str(mine), "--device", "cpu"]) + + scene = _config_of(app).scene + assert scene is not None + assert scene.first_frame_path == mine + # Still drawn, and still drawn from the moment the scene recorded: a frame + # of your own says nothing about where along the road it was taken. + assert scene.view_start_us == _CAPTURED_US + + +def test_a_frame_of_your_own_is_enough_to_ask_for_a_drawn_run( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Naming a frame is not asking to replay anything, so it leaves the default + scene being drawn rather than demanding a recording to go with it.""" + monkeypatch.setattr(omnidreams.scenes, "FLASHDREAMS_CACHE_DIR", tmp_path / "cache") + archive = _a_scene(tmp_path) + monkeypatch.setattr(app_module, "fetch_scene", lambda scene: archive) + mine = _a_png(tmp_path, 6, 4) + app = OmnidreamsApplication(pipeline_config=object()) + + app.init(["--first-frame", str(mine), "--device", "cpu"]) + + scene = _config_of(app).scene + assert scene is not None + assert scene.first_frame_path == mine + + +def test_a_drawn_run_cannot_continue_from_more_frames_than_it_draws( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """One camera is drawn, so a second frame is a camera that never gets one.""" + monkeypatch.setattr(omnidreams.scenes, "FLASHDREAMS_CACHE_DIR", tmp_path / "cache") + archive = _a_scene(tmp_path) + app = OmnidreamsApplication(pipeline_config=object()) + + with pytest.raises(ValueError, match="Pass one"): + app.init( + [ + "--scene", + str(archive), + "--first-frame", + "a.png", + "b.png", + "--device", + "cpu", + ] + ) + + +@pytest.mark.parametrize( + "replaying", + [ + ["--hdmap"], + ["--hdmap", "road.mp4"], + ["--hdmap", "a-recorded-drive"], + ], +) +def test_a_scene_and_a_recording_cannot_both_say_where_the_layout_comes_from( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, replaying: list[str] +) -> None: + """They are two answers to one question, and quietly preferring either one + would leave a run doing something other than what it was asked to.""" + monkeypatch.setattr(omnidreams.scenes, "FLASHDREAMS_CACHE_DIR", tmp_path / "cache") + archive = _a_scene(tmp_path) + app = OmnidreamsApplication(pipeline_config=object()) + + with pytest.raises(ValueError, match="cannot be combined"): + app.init(["--scene", str(archive), *replaying, "--device", "cpu"]) + + +def test_a_scene_with_nothing_recorded_to_continue_from_is_refused( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Before the model loads, since a checkpoint of several gigabytes is a long + wait for a scene that was never going to work.""" + monkeypatch.setattr(omnidreams.scenes, "FLASHDREAMS_CACHE_DIR", tmp_path / "cache") + archive = tmp_path / "empty.usdz" + with zipfile.ZipFile(archive, "w") as scene: + scene.writestr("prompt.txt", "A road with no frames of it.") + app = OmnidreamsApplication(pipeline_config=object()) + + with pytest.raises(FileNotFoundError, match="no timestamped frames"): + app.init(["--scene", str(archive), "--device", "cpu"]) + + +def test_a_scene_that_is_neither_a_path_nor_a_known_id_is_refused( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Named as one message rather than two, since someone who mistyped a path + does not want to be told their typo is not a scene id either.""" + + def missing(scene_uuid: str) -> Path: + raise OSError(f"no scene {scene_uuid}") + + monkeypatch.setattr(scenes_module, "hf_hub_download_scene", missing) + app = OmnidreamsApplication(pipeline_config=object()) + + with pytest.raises(FileNotFoundError, match="no such path"): + app.init(["--scene", "not-a-scene", "--device", "cpu"]) + + +## Stand-ins + + +class StubRenderer: + """A scene of a known length, without a scene or a rasterizer. + + Draws the darkest and brightest bytes a rasterizer can, so a test can tell + whether the conversion to model pixels covered the range. + """ + + def __init__(self, *, frame_count: int, height: int = 4, width: int = 6) -> None: + self._frame_count = frame_count + self._height = height + self._width = width + self.drawn: list[tuple[int, int]] = [] + self.opened = 0 + self.closed = 0 + + @property + def frame_count(self) -> int: + return self._frame_count + + def open(self) -> None: + self.opened += 1 + + def render(self, start: int, count: int) -> torch.Tensor: + self.drawn.append((start, count)) + frames = torch.zeros(count, self._height, self._width, 3, dtype=torch.uint8) + frames[..., 0, :] = 255 + return frames + + def close(self) -> None: + self.closed += 1 + + +## Helpers + + +def _source( + renderer: StubRenderer, *, first_frame: Path | None = None +) -> RenderedHDMapSource: + """Return a drawn source over a stand-in renderer.""" + return RenderedHDMapSource( + renderer=renderer, + first_frame_path=first_frame or Path("unused.png"), + view_name=_SCENE_CAMERA, + pixel_width=6, + pixel_height=4, + device="cpu", + ) + + +def _renderer(*, view_start_us: int, frames_per_second: int) -> LudusSceneRenderer: + """Return a renderer that has not loaded anything, for its timeline alone.""" + return LudusSceneRenderer( + scene_path=Path("unused.usdz"), + camera=_SCENE_CAMERA, + view_start_us=view_start_us, + frames_per_second=frames_per_second, + pixel_width=6, + pixel_height=4, + device="cpu", + ) + + +def _a_scene(tmp_path: Path, *, prompt: str | None = None) -> Path: + """Write a scene archive holding one recorded frame, and maybe a prompt. + + The frame is not a real image: nothing here decodes one, and a test that + said otherwise would be claiming to cover more than it does. + """ + archive = tmp_path / "clipgt-a-scene.usdz" + with zipfile.ZipFile(archive, "w") as scene: + for captured_us in (_CAPTURED_US + 33_333, _CAPTURED_US): + scene.writestr( + f"frames/{_SCENE_CAMERA}/{captured_us}.jpeg", b"a recorded frame" + ) + if prompt is not None: + scene.writestr("prompt.txt", prompt) + return archive + + +def _a_png(tmp_path: Path, width: int, height: int) -> Path: + """Write an image for the one test that decodes what it continues from.""" + path = tmp_path / "first_frame.png" + Image.new("RGB", (width, height), color=(20, 40, 60)).save(path) + return path + + +def _config_of(app: OmnidreamsApplication) -> OmnidreamsSessionConfig: + """Return what an application's command line resolved to.""" + config = app._config + assert config is not None + return config diff --git a/integrations_v2/omnidreams/pyproject.toml b/integrations_v2/omnidreams/pyproject.toml new file mode 100644 index 000000000..20207a41b --- /dev/null +++ b/integrations_v2/omnidreams/pyproject.toml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "flashdreams-omnidreams-v2" +version = "0.1.0" +description = "Omnidreams driving application for the FlashDreams v2 API." +readme = "README.md" +requires-python = ">=3.10" +dependencies = [ + # Integration packages depend on the public framework, never the reverse. + "flashdreams", + # The model itself, whose pipeline config this application drives. It owns + # the module name ``omnidreams``, which is why this one is suffixed. + "flashdreams-omnidreams", +] + +[tool.uv.sources] +# Resolve both from this repository while developing the workspace. +flashdreams = { workspace = true } +flashdreams-omnidreams = { workspace = true } + +# What ``flashdreams-run-v2 omnidreams`` resolves to. A separate group from the +# v1 ``flashdreams.runner_configs``, which holds the older contract; the slug is +# the same because it is the same model. +[project.entry-points."flashdreams.applications_v2"] +"omnidreams" = "omnidreams_v2.app:create_app" + +[tool.setuptools.packages.find] +# Keep the distribution limited to the integration-owned Python package. +include = ["omnidreams_v2*"] + +[tool.uv] +managed = true diff --git a/pyproject.toml b/pyproject.toml index b6b8eb02a..ab914b2e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,7 @@ extraPaths = [ "integrations/wan22", "integrations_v2/color_fade", "integrations_v2/null_model", + "integrations_v2/omnidreams", "integrations_v2/red_screen", "integrations_v2/t2v_causal_forcing", "integrations_v2/t2v_cosmos_predict2", @@ -94,6 +95,7 @@ extra-paths = [ "integrations/wan22", "integrations_v2/color_fade", "integrations_v2/null_model", + "integrations_v2/omnidreams", "integrations_v2/red_screen", "integrations_v2/t2v_causal_forcing", "integrations_v2/t2v_cosmos_predict2", diff --git a/uv.lock b/uv.lock index edf6320b5..4c56e060f 100644 --- a/uv.lock +++ b/uv.lock @@ -29,6 +29,7 @@ members = [ "flashdreams-lingbot", "flashdreams-null-model", "flashdreams-omnidreams", + "flashdreams-omnidreams-v2", "flashdreams-red-screen", "flashdreams-sana-wm", "flashdreams-self-forcing", @@ -1400,6 +1401,21 @@ requires-dist = [ ] provides-extras = ["interactive-drive", "rtx-postprocess", "dev"] +[[package]] +name = "flashdreams-omnidreams-v2" +version = "0.1.0" +source = { editable = "integrations_v2/omnidreams" } +dependencies = [ + { name = "flashdreams" }, + { name = "flashdreams-omnidreams" }, +] + +[package.metadata] +requires-dist = [ + { name = "flashdreams", editable = "flashdreams" }, + { name = "flashdreams-omnidreams", editable = "integrations/omnidreams" }, +] + [[package]] name = "flashdreams-red-screen" version = "0.1.0"