From 0fa723c4a7315cf32c745a8efa56468fb83e481e Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Mon, 20 Jul 2026 16:21:36 +0000 Subject: [PATCH 1/5] Add RTX super resolution postprocessor Signed-off-by: Gangzheng Tong --- docs/source/api/cli.rst | 16 + .../flashdreams/infra/postprocess/__init__.py | 10 + .../flashdreams/infra/postprocess/rtx.py | 340 ++++++++++++++++++ .../flashdreams/infra/postprocess/stream.py | 4 +- flashdreams/pyproject.toml | 3 + flashdreams/tests/test_postprocess_presets.py | 11 + .../test_rtx_super_resolution_postprocess.py | 195 ++++++++++ 7 files changed, 578 insertions(+), 1 deletion(-) create mode 100644 flashdreams/flashdreams/infra/postprocess/rtx.py create mode 100644 flashdreams/tests/test_rtx_super_resolution_postprocess.py diff --git a/docs/source/api/cli.rst b/docs/source/api/cli.rst index e3f28302b..7dc6102fe 100644 --- a/docs/source/api/cli.rst +++ b/docs/source/api/cli.rst @@ -53,6 +53,22 @@ Resolve config only (no model instantiation): uv run flashdreams-run --no-instantiate self-forcing-wan2.1-t2v-1.3b-taehv +Post-processing presets +----------------------- + +Post-processing presets run on decoded RGB frames from a video runner. Select +one with ``--postprocess.preset``: + +.. code-block:: bash + + uv run flashdreams-run wan21-t2v-1.3b-480p \ + --postprocess.preset rtx-super-resolution + +The ``rtx-super-resolution`` preset wraps NVIDIA VFX Python bindings for RTX +Video Super Resolution. Install ``nvidia-vfx`` from NVIDIA's Python package +index in the active environment and run on a supported RTX GPU before selecting +this preset. + See also -------- diff --git a/flashdreams/flashdreams/infra/postprocess/__init__.py b/flashdreams/flashdreams/infra/postprocess/__init__.py index 814c6c934..126bed168 100644 --- a/flashdreams/flashdreams/infra/postprocess/__init__.py +++ b/flashdreams/flashdreams/infra/postprocess/__init__.py @@ -25,6 +25,12 @@ VideoTensorLayout, to_bvtchw, ) +from flashdreams.infra.postprocess.rtx import ( + POSTPROCESS_PRESET_RTX_SUPER_RESOLUTION, + RTXVideoSuperResolutionPostProcessor, + RTXVideoSuperResolutionPostProcessorConfig, + RTXVideoSuperResolutionQuality, +) from flashdreams.infra.postprocess.stream import ( VideoPostprocessStepStats, VideoPostprocessStream, @@ -43,4 +49,8 @@ "VideoSpec", "VideoTensorLayout", "to_bvtchw", + "RTXVideoSuperResolutionPostProcessor", + "RTXVideoSuperResolutionPostProcessorConfig", + "RTXVideoSuperResolutionQuality", + "POSTPROCESS_PRESET_RTX_SUPER_RESOLUTION", ] diff --git a/flashdreams/flashdreams/infra/postprocess/rtx.py b/flashdreams/flashdreams/infra/postprocess/rtx.py new file mode 100644 index 000000000..24d624472 --- /dev/null +++ b/flashdreams/flashdreams/infra/postprocess/rtx.py @@ -0,0 +1,340 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""NVIDIA RTX Video Super Resolution post-processor.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal, get_args + +import torch +from torch import Tensor + +from flashdreams.infra.postprocess.base import ( + VideoChunk, + VideoPostProcessor, + VideoPostProcessorConfig, + VideoPostProcessorSession, + VideoSpec, + to_bvtchw, +) + +RTXVideoSuperResolutionQuality = Literal[ + "BICUBIC", + "LOW", + "MEDIUM", + "HIGH", + "ULTRA", + "DENOISE_LOW", + "DENOISE_MEDIUM", + "DENOISE_HIGH", + "DENOISE_ULTRA", + "DEBLUR_LOW", + "DEBLUR_MEDIUM", + "DEBLUR_HIGH", + "DEBLUR_ULTRA", + "HIGHBITRATE_LOW", + "HIGHBITRATE_MEDIUM", + "HIGHBITRATE_HIGH", + "HIGHBITRATE_ULTRA", +] +"""Quality modes exposed by ``nvvfx.VideoSuperRes.QualityLevel``.""" + +_QUALITY_NAMES = get_args(RTXVideoSuperResolutionQuality) +_SAME_RESOLUTION_QUALITIES = { + "DENOISE_LOW", + "DENOISE_MEDIUM", + "DENOISE_HIGH", + "DENOISE_ULTRA", + "DEBLUR_LOW", + "DEBLUR_MEDIUM", + "DEBLUR_HIGH", + "DEBLUR_ULTRA", +} + + +@dataclass(kw_only=True) +class RTXVideoSuperResolutionPostProcessorConfig(VideoPostProcessorConfig): + """Post-process RGB video frames with NVIDIA RTX Video Super Resolution. + + The runtime is provided by the optional ``nvidia-vfx`` Python package, which + exposes ``nvvfx.VideoSuperRes``. The processor preserves FlashDreams' + ``[-1, 1]`` tensor range at its boundary and converts each frame to the + ``[0, 1]`` float32 CUDA input expected by the VFX binding. + """ + + _target: type["RTXVideoSuperResolutionPostProcessor"] = field( + default_factory=lambda: RTXVideoSuperResolutionPostProcessor + ) + + scale: float = 2.0 + """Spatial upsample factor used when explicit output dimensions are unset.""" + + output_width: int | None = None + """Optional explicit output width in pixels.""" + + output_height: int | None = None + """Optional explicit output height in pixels.""" + + quality: RTXVideoSuperResolutionQuality = "HIGH" + """RTX Video Super Resolution quality mode.""" + + device: int = 0 + """CUDA device index passed to ``nvvfx.VideoSuperRes``.""" + + clamp_input: bool = True + """Clamp incoming FlashDreams frames to ``[-1, 1]`` before VFX conversion.""" + + non_blocking: bool = False + """Forward ``non_blocking`` to ``VideoSuperRes.run``.""" + + use_current_stream: bool = True + """Pass the current PyTorch CUDA stream pointer to ``VideoSuperRes.run``.""" + + def output_spec(self, input_spec: VideoSpec) -> VideoSpec: + """Return the RGB stream specification produced by RTX VSR.""" + if input_spec.channels != 3: + raise ValueError( + f"RTX Video Super Resolution expects RGB chunks with 3 channels; " + f"got {input_spec.channels}." + ) + output_height, output_width = self._output_dimensions(input_spec) + return VideoSpec( + height=output_height, + width=output_width, + fps=input_spec.fps, + channels=3, + ) + + def _output_dimensions(self, input_spec: VideoSpec) -> tuple[int, int]: + if (self.output_width is None) != (self.output_height is None): + raise ValueError( + "RTX Video Super Resolution requires both output_width and " + "output_height when explicit output dimensions are configured." + ) + if self.output_width is not None and self.output_height is not None: + output_height = self.output_height + output_width = self.output_width + else: + if self.scale <= 0: + raise ValueError( + "RTX Video Super Resolution scale must be positive; " + f"got {self.scale}." + ) + output_height = int(round(input_spec.height * self.scale)) + output_width = int(round(input_spec.width * self.scale)) + if output_height <= 0 or output_width <= 0: + raise ValueError( + "RTX Video Super Resolution output dimensions must be positive; " + f"got {output_height}x{output_width}." + ) + if ( + self.quality in _SAME_RESOLUTION_QUALITIES + and (output_height, output_width) != (input_spec.height, input_spec.width) + ): + raise ValueError( + f"RTX Video Super Resolution quality {self.quality!r} is a " + "same-resolution mode; set scale=1.0 or explicit dimensions " + f"{input_spec.height}x{input_spec.width}." + ) + return output_height, output_width + + +class RTXVideoSuperResolutionPostProcessor( + VideoPostProcessor[RTXVideoSuperResolutionPostProcessorConfig] +): + """Factory for RTX Video Super Resolution sessions.""" + + def start(self, spec: VideoSpec) -> VideoPostProcessorSession: + """Start a lazy RTX VSR session for one generated stream.""" + return _RTXVideoSuperResolutionPostProcessorSession(self.config, spec) + + +class _RTXVideoSuperResolutionPostProcessorSession(VideoPostProcessorSession): + """Stateful RTX VSR processor that applies the VFX effect per frame.""" + + def __init__( + self, config: RTXVideoSuperResolutionPostProcessorConfig, spec: VideoSpec + ) -> None: + self._config = config + self._input_spec = spec + self._output_spec = config.output_spec(spec) + self._effect: Any | None = None + self._closed = False + + @torch.no_grad() + def process(self, chunk: VideoChunk) -> list[VideoChunk]: + """Upsample every frame in ``chunk`` and emit one output chunk.""" + if self._closed: + raise RuntimeError( + "cannot process RTX Video Super Resolution after flush()" + ) + canonical = self._chunk_to_bvtchw(chunk) + if canonical.shape[2] == 0: + return [ + VideoChunk( + tensor=_empty_output_like(canonical, spec=self._output_spec), + layout="bvtchw", + metadata={ + "source": "rtx_video_super_resolution", + "input_chunk": dict(chunk.metadata), + }, + ) + ] + effect = self._ensure_effect() + output = self._run_vsr(canonical, effect) + metadata = { + "source": "rtx_video_super_resolution", + "input_chunk": dict(chunk.metadata), + } + return [VideoChunk(tensor=output, layout="bvtchw", metadata=metadata)] + + def flush(self) -> list[VideoChunk]: + """Close the VFX effect and emit no tail frames.""" + if not self._closed: + self._closed = True + if self._effect is not None: + self._effect.close() + self._effect = None + return [] + + def _chunk_to_bvtchw(self, chunk: VideoChunk) -> Tensor: + canonical = to_bvtchw(chunk.tensor, layout=chunk.layout) + _, _, _, channels, height, width = canonical.shape + if channels != 3: + raise ValueError( + f"RTX Video Super Resolution expects RGB chunks with 3 channels; " + f"got {channels}." + ) + if height != self._input_spec.height or width != self._input_spec.width: + raise ValueError( + "RTX Video Super Resolution stream dimensions changed from " + f"{self._input_spec.height}x{self._input_spec.width} to " + f"{height}x{width}." + ) + return canonical + + def _ensure_effect(self) -> Any: + if self._effect is not None: + return self._effect + + video_super_res = _load_video_super_res_class() + effect = video_super_res( + quality=_resolve_quality(video_super_res, self._config.quality), + device=self._config.device, + ) + effect.output_width = self._output_spec.width + effect.output_height = self._output_spec.height + effect.load() + self._effect = effect + return effect + + def _run_vsr(self, canonical: Tensor, effect: Any) -> Tensor: + batch, views, frames, _, _, _ = canonical.shape + device = _torch_device_for_nvvfx_device(self._config.device) + stream_ptr = _current_cuda_stream_ptr( + device=device, + enabled=self._config.use_current_stream, + ) + outputs: list[Tensor] = [] + for batch_idx in range(batch): + for view_idx in range(views): + for frame_idx in range(frames): + outputs.append( + self._run_frame( + canonical[batch_idx, view_idx, frame_idx], + effect=effect, + device=device, + stream_ptr=stream_ptr, + ) + ) + stacked = torch.stack(outputs, dim=0) + return stacked.reshape( + batch, + views, + frames, + 3, + self._output_spec.height, + self._output_spec.width, + ) + + def _run_frame( + self, + frame: Tensor, + *, + effect: Any, + device: torch.device, + stream_ptr: int, + ) -> Tensor: + # The VFX binding expects contiguous channels-first float32 CUDA frames + # in [0, 1]. FlashDreams post-processing receives RGB in [-1, 1]. + frame = frame.to(device=device, dtype=torch.float32) + if self._config.clamp_input: + frame = frame.clamp(-1.0, 1.0) + frame = frame.add(1.0).mul(0.5).contiguous() + result = effect.run( + frame, + non_blocking=self._config.non_blocking, + stream_ptr=stream_ptr, + ) + output = torch.from_dlpack(result.image).clone() + return output.mul(2.0).sub(1.0) + + +def _empty_output_like(canonical: Tensor, *, spec: VideoSpec) -> Tensor: + batch, views, _, _, _, _ = canonical.shape + return canonical.new_empty( + (batch, views, 0, 3, spec.height, spec.width), + dtype=torch.float32, + ) + + +def _load_video_super_res_class() -> Any: + try: + from nvvfx import VideoSuperRes # noqa: PLC0415 + except ImportError as exc: + raise RuntimeError( + "RTX Video Super Resolution post-processing requires the optional " + "`nvidia-vfx` package, which provides `nvvfx.VideoSuperRes`. " + "Install it in the active environment before selecting the " + "`rtx-super-resolution` postprocess preset." + ) from exc + return VideoSuperRes + + +def _resolve_quality(video_super_res: Any, quality: str) -> Any: + try: + return getattr(video_super_res.QualityLevel, quality) + except AttributeError as exc: + available = ", ".join(_QUALITY_NAMES) + raise ValueError( + f"Unsupported RTX Video Super Resolution quality {quality!r}. " + f"Supported values: {available}." + ) from exc + + +def _torch_device_for_nvvfx_device(device: int) -> torch.device: + return torch.device(f"cuda:{device}") + + +def _current_cuda_stream_ptr(*, device: torch.device, enabled: bool) -> int: + if not enabled or device.type != "cuda": + return 0 + return torch.cuda.current_stream(device=device).cuda_stream + + +POSTPROCESS_PRESET_RTX_SUPER_RESOLUTION = RTXVideoSuperResolutionPostProcessorConfig() +"""Default RTX Video Super Resolution post-processing preset.""" diff --git a/flashdreams/flashdreams/infra/postprocess/stream.py b/flashdreams/flashdreams/infra/postprocess/stream.py index f67faef55..994bf6a4a 100644 --- a/flashdreams/flashdreams/infra/postprocess/stream.py +++ b/flashdreams/flashdreams/infra/postprocess/stream.py @@ -235,7 +235,9 @@ def _collected_output(self) -> Tensor | None: return None if not self._chunks: raise ValueError(self.empty_message) - return torch.cat(self._chunks, dim=self._time_dim) + output = torch.cat(self._chunks, dim=self._time_dim) + self._chunks.clear() + return output def create_runner_postprocess_stream( diff --git a/flashdreams/pyproject.toml b/flashdreams/pyproject.toml index c7553c5e9..9d4842150 100644 --- a/flashdreams/pyproject.toml +++ b/flashdreams/pyproject.toml @@ -70,6 +70,9 @@ dependencies = [ # binary, tyro union over runners as direct subcommands). flashdreams-run = "flashdreams.scripts.cli:entrypoint" +[project.entry-points."flashdreams.postprocess_presets"] +"rtx-super-resolution" = "flashdreams.infra.postprocess:POSTPROCESS_PRESET_RTX_SUPER_RESOLUTION" + [project.urls] Homepage = "https://github.com/NVIDIA/flashdreams" Repository = "https://github.com/NVIDIA/flashdreams" diff --git a/flashdreams/tests/test_postprocess_presets.py b/flashdreams/tests/test_postprocess_presets.py index a0a567c98..4fda167d3 100644 --- a/flashdreams/tests/test_postprocess_presets.py +++ b/flashdreams/tests/test_postprocess_presets.py @@ -23,6 +23,7 @@ from flashvsr.postprocess import FlashVSRPostProcessorConfig from flashdreams.infra.postprocess import ( + RTXVideoSuperResolutionPostProcessorConfig, VideoPostprocessChainConfig, VideoPostProcessorConfig, ) @@ -45,6 +46,7 @@ def test_discover_postprocess_presets_includes_flashvsr_entries() -> None: assert "flashvsr-v1.1-sparse-2.0" in presets assert "flashvsr-v1.1-sparse-1.5" in presets assert "flashvsr-v1.1-full-attn" in presets + assert "rtx-super-resolution" in presets def test_resolve_postprocess_preset_rejects_unknown_name() -> None: @@ -67,6 +69,15 @@ def test_chain_config_appends_preset_after_explicit_processors() -> None: assert preset.sparse_ratio == 2.0 +def test_chain_config_resolves_rtx_super_resolution_preset() -> None: + chain = VideoPostprocessChainConfig(preset="rtx-super-resolution") + + preset = chain.resolved_processors()[0] + + assert isinstance(preset, RTXVideoSuperResolutionPostProcessorConfig) + assert preset.scale == 2.0 + + def test_chain_config_is_enabled_for_preset_only() -> None: chain = VideoPostprocessChainConfig(preset="flashvsr-v1.1-sparse-2.0") diff --git a/flashdreams/tests/test_rtx_super_resolution_postprocess.py b/flashdreams/tests/test_rtx_super_resolution_postprocess.py new file mode 100644 index 000000000..9bba07b72 --- /dev/null +++ b/flashdreams/tests/test_rtx_super_resolution_postprocess.py @@ -0,0 +1,195 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU tests for the RTX Video Super Resolution post-processor.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +import flashdreams.infra.postprocess.rtx as rtx_module +from flashdreams.infra.postprocess import ( + RTXVideoSuperResolutionPostProcessorConfig, + VideoChunk, + VideoSpec, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_rtx_super_resolution_output_spec_uses_scale() -> None: + config = RTXVideoSuperResolutionPostProcessorConfig(scale=2.0) + + spec = config.output_spec(VideoSpec(height=540, width=960, fps=24)) + + assert spec == VideoSpec(height=1080, width=1920, fps=24, channels=3) + + +def test_rtx_super_resolution_output_spec_uses_explicit_dimensions() -> None: + config = RTXVideoSuperResolutionPostProcessorConfig( + output_height=720, + output_width=1280, + scale=4.0, + ) + + spec = config.output_spec(VideoSpec(height=360, width=640)) + + assert spec == VideoSpec(height=720, width=1280, channels=3) + + +def test_rtx_super_resolution_rejects_partial_explicit_dimensions() -> None: + config = RTXVideoSuperResolutionPostProcessorConfig(output_width=1280) + + with pytest.raises(ValueError, match="both output_width and output_height"): + config.output_spec(VideoSpec(height=360, width=640)) + + +def test_rtx_super_resolution_rejects_non_rgb_input() -> None: + config = RTXVideoSuperResolutionPostProcessorConfig() + + with pytest.raises(ValueError, match="expects RGB"): + config.output_spec(VideoSpec(height=8, width=8, channels=1)) + + +def test_rtx_super_resolution_rejects_scaled_same_resolution_quality() -> None: + config = RTXVideoSuperResolutionPostProcessorConfig( + quality="DENOISE_HIGH", + scale=2.0, + ) + + with pytest.raises(ValueError, match="same-resolution mode"): + config.output_spec(VideoSpec(height=8, width=8)) + + +def test_rtx_super_resolution_processes_frames_with_fake_backend( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _FakeVideoSuperRes: + class QualityLevel: + HIGH = "HIGH" + + def __init__(self, *, quality: str, device: int) -> None: + self.quality = quality + self.device = device + self.output_width: int | None = None + self.output_height: int | None = None + self.loaded = False + self.closed = False + self.calls: list[tuple[tuple[int, ...], bool, int]] = [] + created.append(self) + + def load(self) -> None: + self.loaded = True + + def run( + self, + input_array: torch.Tensor, + *, + non_blocking: bool, + stream_ptr: int, + ) -> SimpleNamespace: + assert self.loaded + assert self.output_height is not None + assert self.output_width is not None + assert input_array.dtype == torch.float32 + assert input_array.min() >= 0 + assert input_array.max() <= 1 + self.calls.append((tuple(input_array.shape), non_blocking, stream_ptr)) + image = F.interpolate( + input_array.unsqueeze(0), + size=(self.output_height, self.output_width), + mode="nearest", + )[0] + return SimpleNamespace(image=image) + + def close(self) -> None: + self.closed = True + + created: list[_FakeVideoSuperRes] = [] + + monkeypatch.setattr( + rtx_module, "_load_video_super_res_class", lambda: _FakeVideoSuperRes + ) + monkeypatch.setattr( + rtx_module, + "_torch_device_for_nvvfx_device", + lambda device: torch.device("cpu"), + ) + + config = RTXVideoSuperResolutionPostProcessorConfig(scale=2.0) + session = config.setup().start(VideoSpec(height=2, width=3, fps=12)) + video = torch.linspace(-1.0, 1.0, 2 * 3 * 2 * 3).reshape(2, 3, 2, 3) + + outputs = session.process( + VideoChunk( + tensor=video, + layout="tchw", + metadata={"autoregressive_index": 0}, + ) + ) + + assert len(outputs) == 1 + assert outputs[0].layout == "bvtchw" + assert outputs[0].tensor.shape == (1, 1, 2, 3, 4, 6) + assert outputs[0].metadata["source"] == "rtx_video_super_resolution" + assert len(created) == 1 + assert created[0].quality == "HIGH" + assert created[0].device == 0 + assert created[0].output_height == 4 + assert created[0].output_width == 6 + assert created[0].calls == [((3, 2, 3), False, 0), ((3, 2, 3), False, 0)] + expected_first = F.interpolate( + video[0].add(1.0).mul(0.5).unsqueeze(0), + size=(4, 6), + mode="nearest", + )[0].mul(2.0).sub(1.0) + assert torch.equal(outputs[0].tensor[0, 0, 0], expected_first) + + assert session.flush() == [] + assert created[0].closed + + +def test_rtx_super_resolution_empty_chunk_does_not_load_backend( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def _fail_load() -> object: + raise AssertionError("backend should not load for empty chunks") + + monkeypatch.setattr(rtx_module, "_load_video_super_res_class", _fail_load) + config = RTXVideoSuperResolutionPostProcessorConfig(scale=2.0) + session = config.setup().start(VideoSpec(height=2, width=3)) + + outputs = session.process(VideoChunk(tensor=torch.empty(0, 3, 2, 3))) + + assert len(outputs) == 1 + assert outputs[0].tensor.shape == (1, 1, 0, 3, 4, 6) + + +def test_rtx_super_resolution_missing_backend_error_is_actionable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def _missing_backend() -> object: + raise RuntimeError("install nvidia-vfx") + + monkeypatch.setattr(rtx_module, "_load_video_super_res_class", _missing_backend) + config = RTXVideoSuperResolutionPostProcessorConfig(scale=2.0) + session = config.setup().start(VideoSpec(height=2, width=3)) + + with pytest.raises(RuntimeError, match="nvidia-vfx"): + session.process(VideoChunk(tensor=torch.zeros(1, 3, 2, 3))) From e8eced21e95c5eace51f017f22ff8f10af475140 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Tue, 21 Jul 2026 13:01:59 -0700 Subject: [PATCH 2/5] Wire postprocessing into OmniDreams WebRTC demo Signed-off-by: Gangzheng Tong --- .../flashdreams/infra/postprocess/rtx.py | 39 +++-- .../test_rtx_super_resolution_postprocess.py | 115 +++++++++++++- integrations/omnidreams/README.md | 5 + .../omnidreams/omnidreams/webrtc/server.py | 72 +++++++++ .../omnidreams/omnidreams/webrtc/session.py | 141 +++++++++++++++++- .../omnidreams/webrtc/web/request_session.css | 30 ++++ .../webrtc/web/request_session.html | 6 + .../omnidreams/webrtc/web/request_session.js | 57 +++++++ .../tests/test_runner_video_output.py | 116 ++++++++++++++ .../omnidreams/tests/test_webrtc_runtime.py | 122 ++++++++++++++- 10 files changed, 678 insertions(+), 25 deletions(-) create mode 100644 integrations/omnidreams/tests/test_runner_video_output.py diff --git a/flashdreams/flashdreams/infra/postprocess/rtx.py b/flashdreams/flashdreams/infra/postprocess/rtx.py index 24d624472..b082ea04e 100644 --- a/flashdreams/flashdreams/infra/postprocess/rtx.py +++ b/flashdreams/flashdreams/infra/postprocess/rtx.py @@ -99,7 +99,7 @@ class RTXVideoSuperResolutionPostProcessorConfig(VideoPostProcessorConfig): """Clamp incoming FlashDreams frames to ``[-1, 1]`` before VFX conversion.""" non_blocking: bool = False - """Forward ``non_blocking`` to ``VideoSuperRes.run``.""" + """Request asynchronous VFX execution before synchronizing at the boundary.""" use_current_stream: bool = True """Pass the current PyTorch CUDA stream pointer to ``VideoSuperRes.run``.""" @@ -141,10 +141,10 @@ def _output_dimensions(self, input_spec: VideoSpec) -> tuple[int, int]: "RTX Video Super Resolution output dimensions must be positive; " f"got {output_height}x{output_width}." ) - if ( - self.quality in _SAME_RESOLUTION_QUALITIES - and (output_height, output_width) != (input_spec.height, input_spec.width) - ): + if self.quality in _SAME_RESOLUTION_QUALITIES and ( + output_height, + output_width, + ) != (input_spec.height, input_spec.width): raise ValueError( f"RTX Video Super Resolution quality {self.quality!r} is a " "same-resolution mode; set scale=1.0 or explicit dimensions " @@ -280,16 +280,25 @@ def _run_frame( stream_ptr: int, ) -> Tensor: # The VFX binding expects contiguous channels-first float32 CUDA frames - # in [0, 1]. FlashDreams post-processing receives RGB in [-1, 1]. - frame = frame.to(device=device, dtype=torch.float32) - if self._config.clamp_input: - frame = frame.clamp(-1.0, 1.0) - frame = frame.add(1.0).mul(0.5).contiguous() + # in [0, 1]. Most FlashDreams runners emit float RGB in [-1, 1], while + # serving integrations such as Omnidreams emit display-ready uint8 RGB. + if frame.dtype == torch.uint8: + frame = frame.to(device=device, dtype=torch.float32).mul(1.0 / 255.0) + else: + frame = frame.to(device=device, dtype=torch.float32) + if self._config.clamp_input: + frame = frame.clamp(-1.0, 1.0) + frame = frame.add(1.0).mul(0.5) + frame = frame.contiguous() result = effect.run( frame, non_blocking=self._config.non_blocking, stream_ptr=stream_ptr, ) + _synchronize_nonblocking_output( + device=device, + enabled=self._config.non_blocking, + ) output = torch.from_dlpack(result.image).clone() return output.mul(2.0).sub(1.0) @@ -304,7 +313,9 @@ def _empty_output_like(canonical: Tensor, *, spec: VideoSpec) -> Tensor: def _load_video_super_res_class() -> Any: try: - from nvvfx import VideoSuperRes # noqa: PLC0415 + from nvvfx import ( # ty: ignore[unresolved-import] # noqa: PLC0415 + VideoSuperRes, + ) except ImportError as exc: raise RuntimeError( "RTX Video Super Resolution post-processing requires the optional " @@ -336,5 +347,11 @@ def _current_cuda_stream_ptr(*, device: torch.device, enabled: bool) -> int: return torch.cuda.current_stream(device=device).cuda_stream +def _synchronize_nonblocking_output(*, device: torch.device, enabled: bool) -> None: + """Wait until an asynchronous VFX write is safe to import through DLPack.""" + if enabled: + torch.cuda.synchronize(device=device) + + POSTPROCESS_PRESET_RTX_SUPER_RESOLUTION = RTXVideoSuperResolutionPostProcessorConfig() """Default RTX Video Super Resolution post-processing preset.""" diff --git a/flashdreams/tests/test_rtx_super_resolution_postprocess.py b/flashdreams/tests/test_rtx_super_resolution_postprocess.py index 9bba07b72..e5a291466 100644 --- a/flashdreams/tests/test_rtx_super_resolution_postprocess.py +++ b/flashdreams/tests/test_rtx_super_resolution_postprocess.py @@ -154,17 +154,122 @@ def close(self) -> None: assert created[0].output_height == 4 assert created[0].output_width == 6 assert created[0].calls == [((3, 2, 3), False, 0), ((3, 2, 3), False, 0)] - expected_first = F.interpolate( - video[0].add(1.0).mul(0.5).unsqueeze(0), - size=(4, 6), - mode="nearest", - )[0].mul(2.0).sub(1.0) + expected_first = ( + F.interpolate( + video[0].add(1.0).mul(0.5).unsqueeze(0), + size=(4, 6), + mode="nearest", + )[0] + .mul(2.0) + .sub(1.0) + ) assert torch.equal(outputs[0].tensor[0, 0, 0], expected_first) assert session.flush() == [] assert created[0].closed +def test_rtx_super_resolution_normalizes_uint8_input( + monkeypatch: pytest.MonkeyPatch, +) -> None: + vfx_inputs: list[torch.Tensor] = [] + + class _FakeVideoSuperRes: + class QualityLevel: + HIGH = "HIGH" + + def __init__(self, **_kwargs: object) -> None: + self.output_width = 0 + self.output_height = 0 + + def load(self) -> None: + pass + + def run(self, input_array: torch.Tensor, **_kwargs: object) -> SimpleNamespace: + vfx_inputs.append(input_array.clone()) + image = F.interpolate( + input_array.unsqueeze(0), + size=(self.output_height, self.output_width), + mode="nearest", + )[0] + return SimpleNamespace(image=image) + + def close(self) -> None: + pass + + monkeypatch.setattr( + rtx_module, "_load_video_super_res_class", lambda: _FakeVideoSuperRes + ) + monkeypatch.setattr( + rtx_module, + "_torch_device_for_nvvfx_device", + lambda device: torch.device("cpu"), + ) + video = torch.tensor( + [[[[[[0, 64], [128, 255]], [[16, 80], [144, 240]], [[32, 96], [160, 224]]]]]], + dtype=torch.uint8, + ) + config = RTXVideoSuperResolutionPostProcessorConfig(scale=2.0) + session = config.setup().start(VideoSpec(height=2, width=2)) + + outputs = session.process(VideoChunk(tensor=video, layout="bvtchw")) + + expected_input = video[0, 0, 0].float().div(255.0) + assert len(vfx_inputs) == 1 + assert torch.allclose(vfx_inputs[0], expected_input) + expected_output = ( + F.interpolate(expected_input.unsqueeze(0), size=(4, 4), mode="nearest")[0] + .mul(2.0) + .sub(1.0) + ) + assert torch.allclose(outputs[0].tensor[0, 0, 0], expected_output) + + +def test_rtx_super_resolution_synchronizes_nonblocking_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _FakeVideoSuperRes: + class QualityLevel: + HIGH = "HIGH" + + def __init__(self, **_kwargs: object) -> None: + self.output_width = 0 + self.output_height = 0 + + def load(self) -> None: + pass + + def run(self, input_array: torch.Tensor, **_kwargs: object) -> SimpleNamespace: + return SimpleNamespace(image=input_array) + + def close(self) -> None: + pass + + synchronizations: list[tuple[torch.device, bool]] = [] + monkeypatch.setattr( + rtx_module, "_load_video_super_res_class", lambda: _FakeVideoSuperRes + ) + monkeypatch.setattr( + rtx_module, + "_torch_device_for_nvvfx_device", + lambda device: torch.device("cpu"), + ) + monkeypatch.setattr( + rtx_module, + "_synchronize_nonblocking_output", + lambda *, device, enabled: synchronizations.append((device, enabled)), + ) + config = RTXVideoSuperResolutionPostProcessorConfig( + scale=1.0, + non_blocking=True, + ) + session = config.setup().start(VideoSpec(height=2, width=3)) + + session.process(VideoChunk(tensor=torch.zeros(1, 3, 2, 3))) + + assert synchronizations == [(torch.device("cpu"), True)] + + def test_rtx_super_resolution_empty_chunk_does_not_load_backend( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/integrations/omnidreams/README.md b/integrations/omnidreams/README.md index 96a480db7..d486ee05b 100644 --- a/integrations/omnidreams/README.md +++ b/integrations/omnidreams/README.md @@ -174,6 +174,11 @@ from the scene's first ground-truth camera frame the weather-matched `clipgt/prompt.txt` (falling back to `clipgt/prompt.txt`). Pass `--scene_dir ` to use a pre-staged local scene instead. +To enable video post-processing by default, pass a registered preset such as +`--postprocess-preset rtx-super-resolution`. The request-session page also +offers a **Post-process** selector; its choice applies to the next connection +and can override the command-line default, including turning processing off. + ## Run gRPC server From the workspace root, run: diff --git a/integrations/omnidreams/omnidreams/webrtc/server.py b/integrations/omnidreams/omnidreams/webrtc/server.py index 4a2280f9a..58edde3c1 100644 --- a/integrations/omnidreams/omnidreams/webrtc/server.py +++ b/integrations/omnidreams/omnidreams/webrtc/server.py @@ -6,6 +6,7 @@ import argparse from importlib.resources import as_file, files from pathlib import Path +from typing import Protocol, cast import torch import torch.distributed as dist @@ -15,12 +16,15 @@ from omnidreams.transformer import CosmosTransformerConfig from omnidreams.webrtc.session import ( OmnidreamsRuntimeConfig, + OmnidreamsSessionInput, OmnidreamsWebRTCSessionManager, ) from flashdreams.core.distributed import ( init as distributed_init, ) +from flashdreams.infra.postprocess import VideoPostprocessChainConfig +from flashdreams.plugins.registry import discover_postprocess_presets from flashdreams.serving.network import get_external_ip from flashdreams.serving.webrtc.bootstrap import ( configure_logging, @@ -28,6 +32,8 @@ run_webrtc_server, ) from flashdreams.serving.webrtc.server import ( + SESSION_MANAGER_KEY, + SessionBusyError, WebRTCSessionManager, create_packaged_webrtc_app, create_webrtc_app, @@ -39,6 +45,14 @@ WEB_DIR_RESOURCE = files("omnidreams.webrtc").joinpath("web") +class _OmnidreamsSessionManager(WebRTCSessionManager, Protocol): + runtime_config: OmnidreamsRuntimeConfig + + def set_pending_session_input( + self, session_input: OmnidreamsSessionInput + ) -> None: ... + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=( @@ -112,9 +126,65 @@ def parse_args() -> argparse.Namespace: type=str, default="camera_front_wide_120fov", ) + parser.add_argument( + "--postprocess-preset", + "--postprocess_preset", + dest="postprocess_preset", + default="", + choices=sorted(discover_postprocess_presets()), + help=( + "Default video post-process preset for WebRTC sessions. The browser " + "can override this selection before connecting." + ), + ) return parser.parse_args() +def _get_omnidreams_manager(app: web.Application) -> _OmnidreamsSessionManager: + return cast(_OmnidreamsSessionManager, app[SESSION_MANAGER_KEY]) + + +async def _postprocess_options(request: web.Request) -> web.StreamResponse: + manager = _get_omnidreams_manager(request.app) + presets = sorted(discover_postprocess_presets()) + return web.json_response( + { + "default_preset": manager.runtime_config.postprocess.preset, + "presets": presets, + } + ) + + +async def _session_input(request: web.Request) -> web.StreamResponse: + try: + payload = await request.json() + except Exception as exc: + raise web.HTTPBadRequest(reason="Expected JSON session input.") from exc + if not isinstance(payload, dict): + raise web.HTTPBadRequest(reason="Session input must be a JSON object.") + preset = payload.get("postprocess_preset") + if not isinstance(preset, str): + raise web.HTTPBadRequest( + reason="Session input must include string 'postprocess_preset'." + ) + + manager = _get_omnidreams_manager(request.app) + try: + manager.set_pending_session_input( + OmnidreamsSessionInput(postprocess_preset=preset) + ) + except SessionBusyError as exc: + raise web.HTTPConflict(reason=str(exc)) from exc + except ValueError as exc: + raise web.HTTPBadRequest(reason=str(exc)) from exc + return web.json_response({"postprocess_preset": preset}) + + +def _configure_app(app: web.Application) -> None: + app.router.add_get("/api/postprocess/options", _postprocess_options) + app.router.add_post("/api/session/input", _session_input) + + def create_app( *, request_session_url: str, @@ -126,6 +196,7 @@ def create_app( session_manager=manager, preload_name="Omnidreams", request_session_url=request_session_url, + configure_app=_configure_app, as_file_fn=as_file, create_app_fn=create_webrtc_app, cleanup_callback=_close_package_resources, @@ -151,6 +222,7 @@ def build_runtime_config( warmup_chunks=args.warmup_chunks, warmup_timeout_s=args.warmup_timeout_s, debug_serve_hdmaps=args.debug_serve_hdmaps, + postprocess=VideoPostprocessChainConfig(preset=args.postprocess_preset), ) diff --git a/integrations/omnidreams/omnidreams/webrtc/session.py b/integrations/omnidreams/omnidreams/webrtc/session.py index 46d04d124..a12304f7f 100644 --- a/integrations/omnidreams/omnidreams/webrtc/session.py +++ b/integrations/omnidreams/omnidreams/webrtc/session.py @@ -10,7 +10,7 @@ import time import zipfile from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path, PurePosixPath from typing import AbstractSet, Any, Callable, TypeVar @@ -49,6 +49,11 @@ RankCoordinator, distributed_op, ) +from flashdreams.infra.postprocess import ( + VideoPostprocessChainConfig, + VideoPostprocessStream, +) +from flashdreams.plugins.registry import resolve_postprocess_preset from flashdreams.serving.webrtc.controls import ( WSAD_SUPPORTED_KEYS, CameraPoseIntegrator, @@ -61,6 +66,7 @@ WebRTCControlSignal, WebRTCStepResult, ) +from flashdreams.serving.webrtc.server import SessionBusyError _T = TypeVar("_T") # Default scene (clear-weather base archive). Weather siblings are selected @@ -396,6 +402,21 @@ def _summarize_sdp_candidates(sdp: str) -> str: ) +def _link_or_copy_file(source: Path, target: Path) -> None: + """Stage a file efficiently without requiring Windows symlink privileges.""" + try: + os.symlink(source, target) + return + except OSError: + pass + + try: + os.link(source, target) + return + except OSError: + shutil.copy2(source, target) + + class OmnidreamsRuntimeError(RuntimeError): """Raised when the Omnidreams WebRTC runtime is used incorrectly.""" @@ -422,6 +443,17 @@ class OmnidreamsRuntimeConfig: warmup_chunks: int = 10 warmup_timeout_s: float = 600.0 debug_serve_hdmaps: bool = False + postprocess: VideoPostprocessChainConfig = field( + default_factory=VideoPostprocessChainConfig + ) + + +@dataclass(frozen=True, slots=True) +class OmnidreamsSessionInput: + """Browser-selectable settings applied to the next WebRTC rollout.""" + + postprocess_preset: str | None = None + """Preset override; ``None`` keeps the CLI default and ``""`` disables it.""" class OmnidreamsInferenceRuntime: @@ -451,6 +483,8 @@ def __init__(self, config: OmnidreamsRuntimeConfig | None = None) -> None: self._camera_to_rig: torch.Tensor | None = None self._initial_ego_pose: np.ndarray | None = None self._next_timestamp_us: int = 0 + self._postprocess_stream: VideoPostprocessStream | None = None + self._postprocess_preset = self.config.postprocess.preset self._closed = False self._clipgt_temp_dir: tempfile.TemporaryDirectory[str] | None = None # Pin every blocking runtime call to one OS thread: Omnidreams' CUDA @@ -474,6 +508,11 @@ def __init__(self, config: OmnidreamsRuntimeConfig | None = None) -> None: def is_master(self) -> bool: return self.rank == self.MASTER_RANK + @property + def postprocess_preset(self) -> str: + """Preset active for the current rollout, or an empty string when off.""" + return self._postprocess_preset + def wait_for_termination(self) -> None: self.rank_coordinator.worker_loop(exit_signal=WebRTCControlSignal.EXIT) @@ -486,12 +525,17 @@ async def initialize(self) -> None: return await self._run_on_runtime_thread(self._initialize_sync_all_ranks) - async def reset_for_new_session(self) -> None: + async def reset_for_new_session( + self, session_input: OmnidreamsSessionInput | None = None + ) -> None: if self._closed: raise OmnidreamsRuntimeError("Runtime is closed.") if self._wrapper is None: raise OmnidreamsRuntimeError("Runtime is not initialized.") - await self._run_on_runtime_thread(self._reset_rollout_sync_all_ranks) + await self._run_on_runtime_thread( + self._reset_rollout_sync_all_ranks, + session_input, + ) async def close(self) -> None: self._closed = True @@ -562,8 +606,10 @@ def _initialize_sync_all_ranks(self) -> None: self._initialize_sync() @distributed_op(WebRTCControlSignal.RESET_SESSION) - def _reset_rollout_sync_all_ranks(self) -> None: - self._reset_rollout_sync() + def _reset_rollout_sync_all_ranks( + self, session_input: OmnidreamsSessionInput | None = None + ) -> None: + self._reset_rollout_sync(session_input=session_input) @distributed_op(WebRTCControlSignal.ACTION_STEP) def _generate_chunk_sync_all_ranks( @@ -759,15 +805,18 @@ def _has_unprefixed_parquets(path: Path) -> bool: staged = Path(self._clipgt_temp_dir.name) for source in parquet_source_dir.glob("*.parquet"): target = staged / f"clip.{source.name}" - os.symlink(source.resolve(), target) + _link_or_copy_file(source.resolve(), target) return staged - def _reset_rollout_sync(self) -> None: + def _reset_rollout_sync( + self, session_input: OmnidreamsSessionInput | None = None + ) -> None: if self._wrapper is None or self._renderer is None: raise OmnidreamsRuntimeError("Runtime is not initialized.") if self._initial_ego_pose is None or self._scene_data is None: raise OmnidreamsRuntimeError("Scene state is not initialized.") + self._reset_postprocess_stream(session_input) if self._state is not None and self._state.pipeline_cache is not None: del self._state.pipeline_cache self._state = None @@ -792,6 +841,7 @@ def _close_sync(self) -> None: self._text_prompts = None self._camera_to_rig = None self._initial_ego_pose = None + self._close_postprocess_stream() if state is not None and wrapper is not None: wrapper.cleanup(state) @@ -805,6 +855,52 @@ def _close_sync(self) -> None: torch.cuda.synchronize(device=self._device) torch.cuda.empty_cache() + def _reset_postprocess_stream( + self, session_input: OmnidreamsSessionInput | None + ) -> None: + self._close_postprocess_stream() + configured = self.config.postprocess + preset = ( + session_input.postprocess_preset + if session_input is not None + and session_input.postprocess_preset is not None + else configured.preset + ) + if preset: + resolve_postprocess_preset(preset) + postprocess = VideoPostprocessChainConfig( + processors=configured.processors, + preset=preset, + ) + world_size = dist.get_world_size() if dist.is_initialized() else 1 + postprocess.validate_execution(world_size=world_size) + self._postprocess_preset = preset + if not postprocess.is_enabled(): + return + if not self.is_master and not postprocess.requires_all_ranks( + world_size=world_size + ): + return + self._postprocess_stream = VideoPostprocessStream( + postprocess=postprocess, + output_layout="bvtchw", + fps=self.config.fps, + per_view=False, + world_size=world_size, + collect_output=False, + move_to_cpu=False, + ) + logger.info( + "Omnidreams WebRTC post-processing enabled with preset {!r}.", + preset, + ) + + def _close_postprocess_stream(self) -> None: + if self._postprocess_stream is None: + return + self._postprocess_stream.finish() + self._postprocess_stream = None + def _generate_one_chunk_sync( self, *, @@ -879,6 +975,12 @@ def _generate_one_chunk_sync( else: video_chunk = output.rgb_frames + if not serve_hdmaps and self._postprocess_stream is not None: + video_chunk = self._postprocess_stream.process( + video_chunk, + autoregressive_index=self.autoregressive_index, + ) + result = WebRTCStepResult( chunk_index=self.autoregressive_index, num_frames=int(video_chunk.shape[2]), @@ -924,12 +1026,35 @@ def __init__( fps=runtime_config.fps, client_liveness_timeout_s=client_liveness_timeout_s, ) + self._pending_session_input: OmnidreamsSessionInput | None = None def _model_name(self) -> str: return self.runtime_config.pipeline_config_name def _chunk_done_extra(self) -> dict[str, Any]: - return {"stream": "hdmap" if self.runtime_config.debug_serve_hdmaps else "rgb"} + return { + "stream": "hdmap" if self.runtime_config.debug_serve_hdmaps else "rgb", + "postprocess_preset": self._runtime.postprocess_preset, + } + + def _peek_pending_session_input(self) -> OmnidreamsSessionInput | None: + return self._pending_session_input + + def _clear_pending_session_input(self) -> None: + self._pending_session_input = None + + async def _reset_runtime_for_session( + self, session_input: OmnidreamsSessionInput | None + ) -> None: + await self._runtime.reset_for_new_session(session_input=session_input) + + def set_pending_session_input(self, session_input: OmnidreamsSessionInput) -> None: + if self.has_active_session(): + raise SessionBusyError(self._busy_message) + preset = session_input.postprocess_preset + if preset: + resolve_postprocess_preset(preset) + self._pending_session_input = session_input def _register_extra_peer_handlers(self, peer_connection: Any) -> None: @peer_connection.on("iceconnectionstatechange") diff --git a/integrations/omnidreams/omnidreams/webrtc/web/request_session.css b/integrations/omnidreams/omnidreams/webrtc/web/request_session.css index d5c1998d5..cbcc14538 100644 --- a/integrations/omnidreams/omnidreams/webrtc/web/request_session.css +++ b/integrations/omnidreams/omnidreams/webrtc/web/request_session.css @@ -38,6 +38,10 @@ button { font: inherit; } +select { + font: inherit; +} + .appShell { min-height: 100vh; } @@ -216,6 +220,32 @@ body[data-status="error"] .statusDot { opacity: 0.52; } +.postprocessField { + display: grid; + gap: 6px; + color: var(--muted); + font-size: 0.72rem; + font-weight: 700; + letter-spacing: 0.08em; + text-transform: uppercase; +} + +.postprocessField select { + width: 100%; + min-height: 34px; + border: 1px solid rgba(255, 255, 255, 0.22); + border-radius: 6px; + background: rgba(12, 14, 15, 0.86); + color: var(--text); + padding: 0 9px; + text-transform: none; +} + +.postprocessField select:disabled { + cursor: not-allowed; + opacity: 0.52; +} + body[data-status="connected"] .connectButton, body[data-status="waiting"] .connectButton, body[data-status="generating"] .connectButton { diff --git a/integrations/omnidreams/omnidreams/webrtc/web/request_session.html b/integrations/omnidreams/omnidreams/webrtc/web/request_session.html index 4c0222ac9..0b2723958 100644 --- a/integrations/omnidreams/omnidreams/webrtc/web/request_session.html +++ b/integrations/omnidreams/omnidreams/webrtc/web/request_session.html @@ -31,6 +31,12 @@

Omnidreams WebRTC Drive

Idle +
Flow waiting diff --git a/integrations/omnidreams/omnidreams/webrtc/web/request_session.js b/integrations/omnidreams/omnidreams/webrtc/web/request_session.js index 06fc7ea2e..cdbd9bbeb 100644 --- a/integrations/omnidreams/omnidreams/webrtc/web/request_session.js +++ b/integrations/omnidreams/omnidreams/webrtc/web/request_session.js @@ -13,6 +13,7 @@ const latencyValue = document.getElementById("latencyValue") const resolutionValue = document.getElementById("resolutionValue") const stepValue = document.getElementById("stepValue") const modelValue = document.getElementById("modelValue") +const postprocessSelect = document.getElementById("postprocessSelect") const controlButtons = Array.from(document.querySelectorAll("[data-control-key]")) const allowedKeys = new Set(["w", "a", "s", "d"]) @@ -119,6 +120,51 @@ function setVideoVisible(visible) { document.body.classList.toggle("has-video", visible) } +async function loadPostprocessOptions() { + const response = await fetch("/api/postprocess/options") + if (!response.ok) { + throw new Error(`post-process options failed (${response.status})`) + } + const payload = await response.json() + const presets = Array.isArray(payload.presets) ? payload.presets : [] + postprocessSelect.replaceChildren() + + const offOption = document.createElement("option") + offOption.value = "" + offOption.textContent = "Off" + postprocessSelect.append(offOption) + for (const preset of presets) { + if (typeof preset !== "string" || !preset) { + continue + } + const option = document.createElement("option") + option.value = preset + option.textContent = preset + postprocessSelect.append(option) + } + const defaultPreset = typeof payload.default_preset === "string" + ? payload.default_preset + : "" + postprocessSelect.value = presets.includes(defaultPreset) ? defaultPreset : "" +} + +async function configureSessionInput() { + const postprocessPreset = postprocessSelect.value + const response = await fetch("/api/session/input", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ postprocess_preset: postprocessPreset }), + }) + if (!response.ok) { + const text = await response.text() + throw new Error(`session configuration failed (${response.status}): ${text}`) + } + logEvent( + `post-process=${postprocessPreset || "off"}`, + { source: "client" } + ) +} + function renderMetrics() { const fps = firstFinite(metrics.fps, metrics.targetFps) const latency = firstFinite(metrics.latencyMs, metrics.rttMs) @@ -640,6 +686,7 @@ function disconnectSession({ notify = true } = {}) { stopStatsPolling() connected = false connectButton.disabled = false + postprocessSelect.disabled = false if (notify && controlChannel && controlChannel.readyState === "open") { try { controlChannel.send(JSON.stringify({ type: "disconnect" })) @@ -662,12 +709,14 @@ async function connectSession() { } connectButton.disabled = true + postprocessSelect.disabled = true setStatus("Connecting", "connecting") setFlow("creating peer connection") logEvent("connecting to server...", { source: "client" }) disconnecting = false try { + await configureSessionInput() const pc = new RTCPeerConnection() const channel = pc.createDataChannel("controls") peerConnection = pc @@ -726,6 +775,7 @@ async function connectSession() { if (["failed", "closed", "disconnected"].includes(state)) { connected = false connectButton.disabled = false + postprocessSelect.disabled = false stopHeartbeat() stopStatsPolling() setStatus(state === "failed" ? "Error" : "Idle", state === "failed" ? "error" : "idle") @@ -780,6 +830,7 @@ async function connectSession() { setFlow("failed") logEvent(`connect failed: ${error.message}`, { source: "client", level: "error" }) connectButton.disabled = false + postprocessSelect.disabled = false } } @@ -854,6 +905,12 @@ function initialize() { attachPointerControls() window.requestAnimationFrame(drawIdleScene) startVideoFrameMonitor() + void loadPostprocessOptions().catch((error) => { + logEvent(`post-process options unavailable: ${error.message}`, { + source: "client", + level: "error", + }) + }) } connectButton.addEventListener("click", () => { diff --git a/integrations/omnidreams/tests/test_runner_video_output.py b/integrations/omnidreams/tests/test_runner_video_output.py new file mode 100644 index 000000000..82ea653e1 --- /dev/null +++ b/integrations/omnidreams/tests/test_runner_video_output.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for the OmniDreams runner's streaming video writer.""" + +from __future__ import annotations + +import io +from pathlib import Path +from typing import Any + +import omnidreams.runner as runner_module +import pytest +import torch + +pytestmark = pytest.mark.ci_cpu + + +class _RecordingStdin: + def __init__(self, *, fail_write: bool = False, fail_close: bool = False) -> None: + self.fail_write = fail_write + self.fail_close = fail_close + self.writes: list[bytes] = [] + self.closed = False + + def write(self, data: bytes) -> None: + if self.fail_write: + raise BrokenPipeError("write failed") + self.writes.append(data) + + def close(self) -> None: + self.closed = True + if self.fail_close: + raise BrokenPipeError("close failed") + + +class _FakeProcess: + def __init__( + self, + *, + returncode: int = 0, + stderr: bytes = b"", + fail_write: bool = False, + fail_close: bool = False, + ) -> None: + self.stdin = _RecordingStdin( + fail_write=fail_write, + fail_close=fail_close, + ) + self.stderr = io.BytesIO(stderr) + self.returncode = returncode + self.wait_calls = 0 + + def wait(self) -> int: + self.wait_calls += 1 + return self.returncode + + +def _canvas(*, height: int = 3, width: int = 5) -> torch.Tensor: + return torch.zeros(1, height, width, 3) + + +def test_find_ffmpeg_binary_fails_loudly_when_host_binary_is_missing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(runner_module.shutil, "which", lambda _name: None) + + with pytest.raises(RuntimeError, match="installed on the host.*PATH"): + runner_module._find_ffmpeg_binary() + + +def test_write_video_pads_odd_dimensions_for_yuv420p( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + process = _FakeProcess() + commands: list[list[str]] = [] + + def _popen(cmd: list[str], **_kwargs: Any) -> _FakeProcess: + commands.append(cmd) + return process + + monkeypatch.setattr(runner_module, "_find_ffmpeg_binary", lambda: "ffmpeg") + monkeypatch.setattr(runner_module.subprocess, "Popen", _popen) + + runner_module._write_video(_canvas(), tmp_path / "odd.mp4", fps=24) + + assert commands[0][commands[0].index("-vf") + 1] == ( + "pad=ceil(iw/2)*2:ceil(ih/2)*2" + ) + assert commands[0][commands[0].index("-s") + 1] == "5x3" + assert process.wait_calls == 1 + + +def test_write_video_reaps_ffmpeg_and_preserves_broken_pipe_diagnostic( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + process = _FakeProcess( + returncode=7, + stderr=b"encoder rejected frame", + fail_write=True, + fail_close=True, + ) + monkeypatch.setattr(runner_module, "_find_ffmpeg_binary", lambda: "ffmpeg") + monkeypatch.setattr( + runner_module.subprocess, + "Popen", + lambda *_args, **_kwargs: process, + ) + + with pytest.raises(RuntimeError, match="exit code 7.*encoder rejected frame"): + runner_module._write_video(_canvas(), tmp_path / "broken.mp4", fps=24) + + assert process.stdin.closed + assert process.wait_calls == 1 diff --git a/integrations/omnidreams/tests/test_webrtc_runtime.py b/integrations/omnidreams/tests/test_webrtc_runtime.py index 39589ff4c..41b56aa20 100644 --- a/integrations/omnidreams/tests/test_webrtc_runtime.py +++ b/integrations/omnidreams/tests/test_webrtc_runtime.py @@ -8,6 +8,7 @@ import sys import zipfile from dataclasses import dataclass +from importlib.resources import files from pathlib import Path from types import SimpleNamespace from typing import Any @@ -23,6 +24,8 @@ OmnidreamsWebRTCSessionManager, ) +import flashdreams.plugins.registry as plugin_registry +from flashdreams.infra.postprocess import VideoPostProcessorConfig from flashdreams.serving.webrtc.controls import ( WSAD_SUPPORTED_KEYS, CameraPoseIntegrator, @@ -167,6 +170,71 @@ def test_generate_chunk_dispatches_start_then_continue() -> None: assert wrapper.skip_video_generation_flags == [False, False] +def test_generate_chunk_postprocesses_rgb_before_cpu_handoff() -> None: + class _FakePostprocessStream: + def __init__(self) -> None: + self.calls: list[int] = [] + + def process( + self, video_chunk: torch.Tensor, *, autoregressive_index: int + ) -> torch.Tensor: + self.calls.append(autoregressive_index) + return torch.full( + (1, 1, video_chunk.shape[2], 3, 8, 10), + 0.5, + device=video_chunk.device, + ) + + runtime, _wrapper = _build_fake_runtime() + postprocess_stream = _FakePostprocessStream() + runtime._postprocess_stream = postprocess_stream # ty:ignore[invalid-assignment] + + result = runtime._generate_one_chunk_sync( + segments=[(0.0, 2 / 30, frozenset({"w"}))], + frame_times=[1 / 30, 2 / 30], + ) + + assert postprocess_stream.calls == [0] + assert result.video_chunk.device.type == "cpu" + assert result.video_chunk.shape == (1, 1, 2, 3, 8, 10) + assert result.video_chunk.unique().tolist() == [0.5] + + +def test_session_postprocess_override_replaces_the_rollout_stream( + monkeypatch: pytest.MonkeyPatch, +) -> None: + preset_config = VideoPostProcessorConfig() + monkeypatch.setattr( + session, + "resolve_postprocess_preset", + lambda name: preset_config, + ) + monkeypatch.setattr( + plugin_registry, + "resolve_postprocess_preset", + lambda name: preset_config, + ) + runtime = OmnidreamsInferenceRuntime( + config=OmnidreamsRuntimeConfig(device="cpu", fps=30) + ) + + runtime._reset_postprocess_stream( + session.OmnidreamsSessionInput(postprocess_preset="fake-preset") + ) + first_stream = runtime._postprocess_stream + + assert first_stream is not None + assert runtime.postprocess_preset == "fake-preset" + + runtime._reset_postprocess_stream( + session.OmnidreamsSessionInput(postprocess_preset="") + ) + + assert first_stream._closed is True + assert runtime._postprocess_stream is None + assert runtime.postprocess_preset == "" + + def test_generate_chunk_can_stream_debug_hdmaps_without_rgb_frames() -> None: runtime, wrapper = _build_fake_runtime() runtime.config.debug_serve_hdmaps = True @@ -238,6 +306,26 @@ def test_prepare_clipgt_dir_stages_nested_unprefixed_parquets(tmp_path: Path) -> assert (staged / "clip.lane.parquet").exists() +def test_link_or_copy_file_falls_back_to_copy( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + source = tmp_path / "source.parquet" + target = tmp_path / "target.parquet" + source.write_bytes(b"parquet data") + + def _raise_link_error(*args: object, **kwargs: object) -> None: + del args, kwargs + raise OSError("links unavailable") + + monkeypatch.setattr(session.os, "symlink", _raise_link_error) + monkeypatch.setattr(session.os, "link", _raise_link_error) + + session._link_or_copy_file(source, target) + + assert target.read_bytes() == source.read_bytes() + assert not target.is_symlink() + + def test_hf_webrtc_scene_sync_requires_usdz_first_frame( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: @@ -397,6 +485,7 @@ def test_build_runtime_config_threads_hf_scene_args(tmp_path: Path) -> None: warmup_chunks=0, warmup_timeout_s=30.0, debug_serve_hdmaps=True, + postprocess_preset="rtx-super-resolution", ) cfg = webrtc_server.build_runtime_config(args, device_override="cuda:7") @@ -408,6 +497,7 @@ def test_build_runtime_config_threads_hf_scene_args(tmp_path: Path) -> None: assert cfg.video_height == 360 assert cfg.video_width == 640 assert cfg.debug_serve_hdmaps is True + assert cfg.postprocess.preset == "rtx-super-resolution" def test_parse_args_omits_scene_dir_by_default( @@ -427,6 +517,7 @@ def test_parse_args_omits_scene_dir_by_default( assert args.scene_dir is None assert args.scene_uuid is None assert args.debug_serve_hdmaps is True + assert args.postprocess_preset == "" def test_runtime_uses_default_scene_uuid_when_scene_is_unspecified( @@ -499,6 +590,7 @@ def test_build_runtime_config_clears_scene_uuid_for_local_scene(tmp_path: Path) warmup_chunks=0, warmup_timeout_s=30.0, debug_serve_hdmaps=True, + postprocess_preset="", ) cfg = webrtc_server.build_runtime_config(args) @@ -508,6 +600,30 @@ def test_build_runtime_config_clears_scene_uuid_for_local_scene(tmp_path: Path) assert cfg.scene_variant == "default" +def test_session_manager_stores_postprocess_override_for_next_rollout() -> None: + manager = OmnidreamsWebRTCSessionManager( + runtime_config=OmnidreamsRuntimeConfig(device="cpu") + ) + session_input = session.OmnidreamsSessionInput(postprocess_preset="") + + manager.set_pending_session_input(session_input) + + assert manager._peek_pending_session_input() == session_input + manager._clear_pending_session_input() + assert manager._peek_pending_session_input() is None + + +def test_webrtc_ui_posts_selected_postprocess_preset() -> None: + web_dir = files("omnidreams.webrtc").joinpath("web") + html = web_dir.joinpath("request_session.html").read_text(encoding="utf-8") + javascript = web_dir.joinpath("request_session.js").read_text(encoding="utf-8") + + assert 'id="postprocessSelect"' in html + assert 'fetch("/api/postprocess/options")' in javascript + assert 'fetch("/api/session/input"' in javascript + assert "postprocess_preset: postprocessPreset" in javascript + + @pytest.mark.asyncio async def test_session_manager_preload_runs_loopback_warmup_once( monkeypatch: pytest.MonkeyPatch, @@ -567,6 +683,7 @@ def __init__(self, config: OmnidreamsRuntimeConfig) -> None: self.initialize_calls = 0 self.reset_calls = 0 self.close_calls = 0 + self.postprocess_preset = config.postprocess.preset self.generated_segments: list[ list[tuple[float, float, frozenset[str]]] ] = [] @@ -574,7 +691,10 @@ def __init__(self, config: OmnidreamsRuntimeConfig) -> None: async def initialize(self) -> None: self.initialize_calls += 1 - async def reset_for_new_session(self) -> None: + async def reset_for_new_session( + self, session_input: session.OmnidreamsSessionInput | None = None + ) -> None: + del session_input self.reset_calls += 1 def peek_steady_chunk_num_frames(self) -> int: From e24d5ea481e540caaea558d9667c6244cc63b5e7 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Tue, 21 Jul 2026 16:41:37 -0700 Subject: [PATCH 3/5] Wire postprocessing into interactive demo Signed-off-by: Gangzheng Tong --- .../flashdreams/infra/postprocess/__init__.py | 4 + .../flashdreams/infra/postprocess/rtx.py | 13 +- flashdreams/pyproject.toml | 2 + flashdreams/tests/test_postprocess_presets.py | 23 +++ .../test_rtx_super_resolution_postprocess.py | 7 +- .../interactive_drive/_pipeline_fakes.py | 4 + .../omnidreams/interactive_drive/app.py | 4 + .../interactive_drive/backends/base.py | 8 + .../interactive_drive/backends/world_model.py | 6 + .../omnidreams/interactive_drive/cli.py | 15 ++ .../omnidreams/interactive_drive/config.py | 7 +- .../omnidreams/interactive_drive/demo.py | 5 + .../omnidreams/interactive_drive/presenter.py | 4 +- .../slangpy_hud_presenter.py | 173 +++++++++++++++++- .../video_model/chunk_pipeline.py | 12 ++ .../interactive_drive/video_model/local.py | 3 + .../world_model/flashdreams_adapter.py | 66 ++++++- .../interactive_drive/test_chunk_pipeline.py | 11 ++ .../tests/interactive_drive/test_cli.py | 12 ++ .../tests/interactive_drive/test_presenter.py | 89 +++++++++ .../test_world_model_adapter.py | 53 ++++++ 21 files changed, 510 insertions(+), 11 deletions(-) diff --git a/flashdreams/flashdreams/infra/postprocess/__init__.py b/flashdreams/flashdreams/infra/postprocess/__init__.py index 126bed168..c01c6ec5e 100644 --- a/flashdreams/flashdreams/infra/postprocess/__init__.py +++ b/flashdreams/flashdreams/infra/postprocess/__init__.py @@ -26,7 +26,9 @@ to_bvtchw, ) from flashdreams.infra.postprocess.rtx import ( + POSTPROCESS_PRESET_RTX_DEBLUR_ULTRA, POSTPROCESS_PRESET_RTX_SUPER_RESOLUTION, + POSTPROCESS_PRESET_RTX_SUPER_RESOLUTION_ULTRA, RTXVideoSuperResolutionPostProcessor, RTXVideoSuperResolutionPostProcessorConfig, RTXVideoSuperResolutionQuality, @@ -53,4 +55,6 @@ "RTXVideoSuperResolutionPostProcessorConfig", "RTXVideoSuperResolutionQuality", "POSTPROCESS_PRESET_RTX_SUPER_RESOLUTION", + "POSTPROCESS_PRESET_RTX_SUPER_RESOLUTION_ULTRA", + "POSTPROCESS_PRESET_RTX_DEBLUR_ULTRA", ] diff --git a/flashdreams/flashdreams/infra/postprocess/rtx.py b/flashdreams/flashdreams/infra/postprocess/rtx.py index b082ea04e..86fb7a5a8 100644 --- a/flashdreams/flashdreams/infra/postprocess/rtx.py +++ b/flashdreams/flashdreams/infra/postprocess/rtx.py @@ -354,4 +354,15 @@ def _synchronize_nonblocking_output(*, device: torch.device, enabled: bool) -> N POSTPROCESS_PRESET_RTX_SUPER_RESOLUTION = RTXVideoSuperResolutionPostProcessorConfig() -"""Default RTX Video Super Resolution post-processing preset.""" +"""Two-times RTX Video Super Resolution with ``HIGH`` quality.""" + +POSTPROCESS_PRESET_RTX_SUPER_RESOLUTION_ULTRA = ( + RTXVideoSuperResolutionPostProcessorConfig(quality="ULTRA") +) +"""Two-times RTX Video Super Resolution with ``ULTRA`` quality.""" + +POSTPROCESS_PRESET_RTX_DEBLUR_ULTRA = RTXVideoSuperResolutionPostProcessorConfig( + scale=1.0, + quality="DEBLUR_ULTRA", +) +"""Same-resolution RTX Video Deblur with ``ULTRA`` quality.""" diff --git a/flashdreams/pyproject.toml b/flashdreams/pyproject.toml index 9d4842150..8058985ae 100644 --- a/flashdreams/pyproject.toml +++ b/flashdreams/pyproject.toml @@ -72,6 +72,8 @@ flashdreams-run = "flashdreams.scripts.cli:entrypoint" [project.entry-points."flashdreams.postprocess_presets"] "rtx-super-resolution" = "flashdreams.infra.postprocess:POSTPROCESS_PRESET_RTX_SUPER_RESOLUTION" +"rtx-super-resolution-ultra" = "flashdreams.infra.postprocess:POSTPROCESS_PRESET_RTX_SUPER_RESOLUTION_ULTRA" +"rtx-deblur-ultra" = "flashdreams.infra.postprocess:POSTPROCESS_PRESET_RTX_DEBLUR_ULTRA" [project.urls] Homepage = "https://github.com/NVIDIA/flashdreams" diff --git a/flashdreams/tests/test_postprocess_presets.py b/flashdreams/tests/test_postprocess_presets.py index 4fda167d3..6549eec1d 100644 --- a/flashdreams/tests/test_postprocess_presets.py +++ b/flashdreams/tests/test_postprocess_presets.py @@ -47,6 +47,8 @@ def test_discover_postprocess_presets_includes_flashvsr_entries() -> None: assert "flashvsr-v1.1-sparse-1.5" in presets assert "flashvsr-v1.1-full-attn" in presets assert "rtx-super-resolution" in presets + assert "rtx-super-resolution-ultra" in presets + assert "rtx-deblur-ultra" in presets def test_resolve_postprocess_preset_rejects_unknown_name() -> None: @@ -78,6 +80,27 @@ def test_chain_config_resolves_rtx_super_resolution_preset() -> None: assert preset.scale == 2.0 +@pytest.mark.parametrize( + ("preset_name", "expected_quality", "expected_scale"), + [ + ("rtx-super-resolution-ultra", "ULTRA", 2.0), + ("rtx-deblur-ultra", "DEBLUR_ULTRA", 1.0), + ], +) +def test_chain_config_resolves_additional_rtx_presets( + preset_name: str, + expected_quality: str, + expected_scale: float, +) -> None: + chain = VideoPostprocessChainConfig(preset=preset_name) + + preset = chain.resolved_processors()[0] + + assert isinstance(preset, RTXVideoSuperResolutionPostProcessorConfig) + assert preset.quality == expected_quality + assert preset.scale == expected_scale + + def test_chain_config_is_enabled_for_preset_only() -> None: chain = VideoPostprocessChainConfig(preset="flashvsr-v1.1-sparse-2.0") diff --git a/flashdreams/tests/test_rtx_super_resolution_postprocess.py b/flashdreams/tests/test_rtx_super_resolution_postprocess.py index e5a291466..28356146f 100644 --- a/flashdreams/tests/test_rtx_super_resolution_postprocess.py +++ b/flashdreams/tests/test_rtx_super_resolution_postprocess.py @@ -67,9 +67,12 @@ def test_rtx_super_resolution_rejects_non_rgb_input() -> None: config.output_spec(VideoSpec(height=8, width=8, channels=1)) -def test_rtx_super_resolution_rejects_scaled_same_resolution_quality() -> None: +@pytest.mark.parametrize("quality", ["DENOISE_HIGH", "DEBLUR_ULTRA"]) +def test_rtx_super_resolution_rejects_scaled_same_resolution_quality( + quality: str, +) -> None: config = RTXVideoSuperResolutionPostProcessorConfig( - quality="DENOISE_HIGH", + quality=quality, # ty: ignore[invalid-argument-type] scale=2.0, ) diff --git a/integrations/omnidreams/omnidreams/interactive_drive/_pipeline_fakes.py b/integrations/omnidreams/omnidreams/interactive_drive/_pipeline_fakes.py index 96037e66f..918f1b509 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/_pipeline_fakes.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/_pipeline_fakes.py @@ -80,6 +80,7 @@ def __init__(self, frames_per_render: int, rgb_value: int = 0) -> None: self.warmup_model_calls = 0 self.load_scene_calls = 0 self.reset_calls = 0 + self.postprocess_enabled_calls: list[bool] = [] @property def can_prewarm(self) -> bool: @@ -95,6 +96,9 @@ def load_scene(self, scene: SceneBundle) -> None: def reset(self) -> None: self.reset_calls += 1 + def set_postprocess_enabled(self, enabled: bool) -> None: + self.postprocess_enabled_calls.append(enabled) + def render_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk: frames = tuple( PresentedFrame( diff --git a/integrations/omnidreams/omnidreams/interactive_drive/app.py b/integrations/omnidreams/omnidreams/interactive_drive/app.py index 0fccebabb..ac4cf9dd5 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/app.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/app.py @@ -168,6 +168,10 @@ def first_chunk_produced(self) -> bool: """``True`` once the model has produced its first generated chunk.""" return self._pipeline.first_chunk_produced.is_set() + def set_postprocess_enabled(self, enabled: bool) -> None: + """Queue a local-display post-process toggle on the model worker.""" + self._pipeline.set_postprocess_enabled(enabled) + def load_scene( self, scene_path: object, variant: str, prompt_override: str | None ) -> bool: diff --git a/integrations/omnidreams/omnidreams/interactive_drive/backends/base.py b/integrations/omnidreams/omnidreams/interactive_drive/backends/base.py index 1ad6e990d..42d90a9e6 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/backends/base.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/backends/base.py @@ -89,6 +89,14 @@ def reset_scene_conditioning(self) -> None: """ self.reset() + def set_postprocess_enabled(self, enabled: bool) -> None: + """Enable or disable generated-video post-processing. + + Pure raster backends have no generated-video post-process path, so the + default implementation intentionally does nothing. + """ + del enabled + @abstractmethod def render_first_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk: raise NotImplementedError diff --git a/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py b/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py index 33d8e64d8..73ae269c0 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py @@ -9,6 +9,7 @@ from pathlib import Path import numpy as np +from flashdreams.infra.postprocess import VideoPostprocessChainConfig from loguru import logger from omnidreams.interactive_drive.backends.base import RenderBackend from omnidreams.interactive_drive.config import ( @@ -45,6 +46,7 @@ def __init__( profile: WorldModelProfileConfig | None = None, bev: BevConfig | None = None, offload_text_encoder: bool = False, + postprocess: VideoPostprocessChainConfig | None = None, ) -> None: super().__init__(chunk=chunk, raster=raster) self._manifest = manifest @@ -53,6 +55,7 @@ def __init__( manifest, profile=profile, offload_text_encoder=offload_text_encoder, + postprocess=postprocess, ) self._scene: SceneBundle | None = None self._next_chunk_count = 0 @@ -237,6 +240,9 @@ def reset_scene_conditioning(self) -> None: self._session.reset(clear_precomputed_embeddings=True) self._next_chunk_count = 0 + def set_postprocess_enabled(self, enabled: bool) -> None: + self._session.set_postprocess_enabled(enabled) + def close(self) -> None: self._session.close() self._rasterizer.cleanup() diff --git a/integrations/omnidreams/omnidreams/interactive_drive/cli.py b/integrations/omnidreams/omnidreams/interactive_drive/cli.py index 2313aabe4..a6c0e2251 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/cli.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/cli.py @@ -7,6 +7,8 @@ from dataclasses import replace from pathlib import Path +from flashdreams.infra.postprocess import VideoPostprocessChainConfig +from flashdreams.plugins.registry import discover_postprocess_presets from loguru import logger from omnidreams.hf_org import DEFAULT_HF_ORG, apply_cli_to_env from omnidreams.hf_org import ENV_VAR as _HF_ORG_ENV_VAR @@ -205,6 +207,17 @@ def build_parser() -> argparse.ArgumentParser: "the cached embeddings across world-model resets." ), ) + parser.add_argument( + "--postprocess-preset", + "--postprocess_preset", + dest="postprocess_preset", + default="", + choices=sorted(discover_postprocess_presets()), + help=( + "Video post-process preset for generated frames. A configured " + "preset starts enabled and can be toggled in the local HUD." + ), + ) parser.add_argument( "--hf-org", default=None, @@ -472,6 +485,7 @@ def prepare_config_and_backend( enabled=bool(args.profile_world_model), ), world_model_offload_text_encoder=bool(args.offload_text_encoder), + postprocess=VideoPostprocessChainConfig(preset=args.postprocess_preset), bev=bev_config, stream_mjpeg_bind=args.stream_mjpeg, stop_after_consumed_chunks=args.stop_after_chunks, @@ -509,6 +523,7 @@ def prepare_config_and_backend( profile=config.world_model_profile, bev=config.bev, offload_text_encoder=config.world_model_offload_text_encoder, + postprocess=config.postprocess, ) return config, backend diff --git a/integrations/omnidreams/omnidreams/interactive_drive/config.py b/integrations/omnidreams/omnidreams/interactive_drive/config.py index de39f2a09..8e6c263c8 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/config.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/config.py @@ -3,10 +3,12 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Literal +from flashdreams.infra.postprocess import VideoPostprocessChainConfig + BackendName = Literal["raster", "omnidreams"] ViewMode = Literal["rgb", "model_rgb"] ComputeDeviceName = Literal["automatic", "cuda", "vulkan"] @@ -112,6 +114,9 @@ class AppConfig: vehicle: VehicleConfig = VehicleConfig() world_model_profile: WorldModelProfileConfig = WorldModelProfileConfig() world_model_offload_text_encoder: bool = False + postprocess: VideoPostprocessChainConfig = field( + default_factory=VideoPostprocessChainConfig + ) bev: BevConfig = BevConfig() # OOB thresholds plumbed to LoopConfig (overridable via CLI --oob-*). # Match alpasim's driver-side proximity: warn > 0.6, respawn >= 2.0 diff --git a/integrations/omnidreams/omnidreams/interactive_drive/demo.py b/integrations/omnidreams/omnidreams/interactive_drive/demo.py index 4e90af5be..c0af59990 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/demo.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/demo.py @@ -797,6 +797,11 @@ def _run_slangpy_hud(args: argparse.Namespace) -> None: close_presenter_on_exit=False, ) presenter.set_model_status(can_prewarm=app.can_prewarm, ready_probe=app.model_ready) + presenter.set_postprocess_control( + preset=config.postprocess.preset, + enabled=config.postprocess.is_enabled(), + callback=app.set_postprocess_enabled, + ) # Attach the wheel up front, bound to the app's long-lived keyboard, so # the HUD's steering / pedal chrome reacts to the physical device during diff --git a/integrations/omnidreams/omnidreams/interactive_drive/presenter.py b/integrations/omnidreams/omnidreams/interactive_drive/presenter.py index c1c44c5af..a9af81baa 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/presenter.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/presenter.py @@ -489,7 +489,9 @@ def enqueue_camera_to_shared_rgba( src_w = int(rgb_tensor.shape[1]) if src_h <= 0 or src_w <= 0: return False - scale = min(area_w / src_w, area_h / src_h) + # Preserve native pixels for smaller frames; only downscale when the + # source exceeds the available camera area. + scale = min(1.0, area_w / src_w, area_h / src_h) target_w = max(1, int(src_w * scale)) target_h = max(1, int(src_h * scale)) target_x = int(ax) + (area_w - target_w) // 2 diff --git a/integrations/omnidreams/omnidreams/interactive_drive/slangpy_hud_presenter.py b/integrations/omnidreams/omnidreams/interactive_drive/slangpy_hud_presenter.py index 0ed71b5ac..e59c14e89 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/slangpy_hud_presenter.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/slangpy_hud_presenter.py @@ -314,6 +314,11 @@ def __init__( # on the main thread, where it's safe to recreate Vulkan # resources. self._pending_resize: tuple[int, int] | None = None + # Last model-frame resolution that drove an automatic native-size + # window resize. User resizes at the same source resolution remain + # authoritative; a new source resolution (for example VSR on/off) + # triggers one fresh native-size request. + self._auto_sized_camera_src_size: tuple[int, int] | None = None self._window.on_resize = self._on_resize self._window.on_keyboard_event = self._on_keyboard_event self._window.on_mouse_event = self._on_mouse_event @@ -378,6 +383,7 @@ def __init__( self._variant_dropdown_open = False self._scene_header_rect: tuple[int, int, int, int] | None = None self._variant_header_rect: tuple[int, int, int, int] | None = None + self._postprocess_rect: tuple[int, int, int, int] | None = None self._scene_item_rects: list[tuple[tuple[int, int, int, int], Any]] = [] self._variant_item_rects: list[tuple[tuple[int, int, int, int], str]] = [] self._hovered_scene_label: str | None = None @@ -411,6 +417,9 @@ def __init__( # clicks and the placeholder shows a "Preloading scenes..." hint, so # the user can't pick a scene until every scene is cached. self._scene_selection_locked_probe: Callable[[], bool] = lambda: False + self._postprocess_preset = "" + self._postprocess_enabled = False + self._postprocess_callback: Callable[[bool], None] = lambda enabled: None # Scene-change request set by the dropdown click handlers. The # outer demo loop checks this after each ``app.run_scene`` returns: @@ -472,6 +481,15 @@ def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: self._apply_resize(new_size[0], new_size[1]) rgb = self._select_view_rgb(frame, view_mode) + if ( + view_mode == "model_rgb" + and frame.model_rgb_host_uint8 is not None + and self._resize_window_for_native_model_frame(rgb) + ): + # Window/swapchain resources are rebuilt at the start of the next + # presentation tick. Drop this transition frame instead of using + # old-size CUDA/Vulkan buffers against the newly resized window. + return try: if self._present_cuda_hud_frame(frame, rgb): return @@ -781,6 +799,52 @@ def _on_resize(self, width: int, height: int) -> None: # race with whatever frame is in flight. self._pending_resize = self._normalise_present_size(width, height) + def _resize_window_for_native_model_frame(self, rgb: object) -> bool: + """Grow the window when a model-frame resolution needs more room. + + The camera region never upscales smaller frames: they remain centered + at native resolution while the existing canvas and HUD stay visible. + Larger frames grow only the dimensions required to fit the source plus + the fixed-width HUD column. Window-manager clamping and later user + resizes remain authoritative until the source resolution changes. + + Returns: + ``True`` when a resize was requested and this frame should be + dropped while presentation resources are rebuilt. + """ + source_size = _rgb_source_size(rgb) + if source_size is None or source_size == getattr( + self, "_auto_sized_camera_src_size", None + ): + return False + source_width, source_height = source_size + current_width, current_height = self._current_window_size() + target_size = ( + max(current_width, source_width + HUD_PANEL_WIDTH), + max(current_height, source_height, MIN_WINDOW_H), + ) + if target_size == (current_width, current_height): + self._auto_sized_camera_src_size = source_size + return False + try: + self._window.resize(*target_size) + except Exception as exc: + logger.warning( + "[presenter] native model-frame window resize failed " + f"source={source_size} target={target_size} ({exc})", + ) + return False + # Some SDL/window-manager combinations deliver the resize callback + # asynchronously. Stash the request as well so the next frame always + # rebuilds resources before presenting at the new dimensions. + self._pending_resize = target_size + self._auto_sized_camera_src_size = source_size + logger.info( + "[presenter] native model-frame window resize " + f"source={source_size} target={target_size}", + ) + return True + def _submit_ready_cuda_hud(self) -> bool: interop = self._cuda_hud_interop if interop is None: @@ -935,7 +999,7 @@ def _compute_camera_fit(self) -> tuple[int, int, int, int] | None: cam_h = screen_h if src_w <= 0 or src_h <= 0: return None - scale = min(cam_w / src_w, cam_h / src_h) + scale = min(1.0, cam_w / src_w, cam_h / src_h) fit_w = max(1, int(src_w * scale)) fit_h = max(1, int(src_h * scale)) offset_x = (cam_w - fit_w) // 2 @@ -1164,7 +1228,7 @@ def _draw_camera( fw, fh = camera.size if fw <= 0 or fh <= 0 or aw <= 0 or ah <= 0: return - scale = min(aw / fw, ah / fh) + scale = min(1.0, aw / fw, ah / fh) target_w = max(1, int(fw * scale)) target_h = max(1, int(fh * scale)) cache_key = (id(camera), target_w, target_h) @@ -1298,6 +1362,7 @@ def _draw_panel( header_w = panel_size[0] - margin * 2 header_y = py + 8 variant_y = header_y + bar_h + 4 + postprocess_y = variant_y + bar_h + 4 self._scene_header_rect = ( header_x, header_y, @@ -1310,6 +1375,12 @@ def _draw_panel( header_x + header_w, variant_y + bar_h, ) + self._postprocess_rect = ( + header_x, + postprocess_y, + header_x + header_w, + postprocess_y + bar_h, + ) center_x = px + panel_size[0] // 2 # ``speed_y`` is the top of the speed-digit chip. PIL renders @@ -1318,7 +1389,7 @@ def _draw_panel( # bar would still land the visible glyph inside the bar. Add a # ~12 px clearance below ``variant_y + bar_h`` so the digit # never overlaps the headers. - speed_y = variant_y + bar_h + 12 + speed_y = postprocess_y + bar_h + 12 self._draw_speed(canvas, draw, center_x, speed_y, int(self._speed_mph)) # Light the reverse indicator red when reverse is engaged; the cached @@ -1382,6 +1453,8 @@ def _get_panel_chrome(self, panel_size: tuple[int, int]) -> Image.Image: # Scene header reads "Preloading scenes..." while locked, so the # lock state has to invalidate the cached chrome too. self._scene_selection_locked(), + self._postprocess_preset, + self._postprocess_enabled, ) if key == self._panel_chrome_cache_key and self._panel_chrome_cache is not None: return self._panel_chrome_cache @@ -1461,12 +1534,55 @@ def _get_panel_chrome(self, panel_size: tuple[int, int]) -> Image.Image: font=self._font_small, ) + # Post-processing is selected by CLI and can be switched live for the + # local window. An empty preset remains visible but disabled so users + # know which launch option unlocks the control. + postprocess_y = variant_y + bar_h + 4 + postprocess_rect = ( + margin, + postprocess_y, + margin + header_w, + postprocess_y + bar_h, + ) + postprocess_available = bool(self._postprocess_preset) + postprocess_clickable = postprocess_available and not ( + self._scene_dropdown_open or self._variant_dropdown_open + ) + d.rounded_rectangle(postprocess_rect, radius=6, fill=HEADER_BG + (255,)) + d.text( + (margin + 10, postprocess_y + 6), + "Upsample 2x", + fill=TEXT_COLOR if postprocess_clickable else LABEL_COLOR, + font=self._font_small, + ) + state_label = ( + ("ON" if self._postprocess_enabled else "OFF") + if postprocess_available + else "N/A" + ) + state_bbox = _measure_text(self._font_small, state_label) + state_w = state_bbox[2] - state_bbox[0] + state_fill = ( + NVIDIA_GREEN + if self._postprocess_enabled and postprocess_clickable + else LABEL_COLOR + ) + d.text( + ( + margin + header_w - state_w - 10 - state_bbox[0], + postprocess_y + 6, + ), + state_label, + fill=state_fill, + font=self._font_small, + ) + # ``mph`` label baseline + reverse-indicator box. Speed-y must # match the live ``_draw_panel`` calculation; both place the # speed-digit chip-top ~12 px below the variant bar so PIL's # tight-bbox glyph chip clears the headers. center_x = panel_w // 2 - speed_y = variant_y + bar_h + 12 + speed_y = postprocess_y + bar_h + 12 mbox = _measure_text(self._font_tiny, "mph") mw = mbox[2] - mbox[0] d.text( @@ -2171,6 +2287,23 @@ def _update_hover(self, pos: tuple[int, int]) -> None: break def _handle_click(self, pos: tuple[int, int]) -> None: + dropdown_open = self._scene_dropdown_open or self._variant_dropdown_open + if ( + not dropdown_open + and self._postprocess_rect + and _rect_contains(self._postprocess_rect, pos) + ): + if self._postprocess_preset: + self._postprocess_enabled = not self._postprocess_enabled + self._postprocess_callback(self._postprocess_enabled) + self._panel_chrome_cache_key = None + self._panel_chrome_cache = None + logger.info( + "[demo] post-processing {} preset={!r}", + "enabled" if self._postprocess_enabled else "disabled", + self._postprocess_preset, + ) + return # While scenes are still preloading, the scene/variant dropdowns are # locked (the only mouse-clickable HUD elements), so ignore clicks # until every scene is cached and selection is instant. @@ -2319,6 +2452,20 @@ def set_model_status( self._model_can_prewarm = bool(can_prewarm) self._model_ready_probe = ready_probe + def set_postprocess_control( + self, + *, + preset: str, + enabled: bool, + callback: Callable[[bool], None], + ) -> None: + """Bind the local HUD toggle to the model worker's post-process state.""" + self._postprocess_preset = preset + self._postprocess_enabled = bool(enabled and preset) + self._postprocess_callback = callback + self._panel_chrome_cache_key = None + self._panel_chrome_cache = None + def set_scene_selection_locked(self, probe: Callable[[], bool]) -> None: """Lock scene/variant selection while ``probe()`` returns True (--preload-scenes). @@ -2476,6 +2623,24 @@ def _has_cuda_tensor(frame: object) -> bool: return callable(getattr(frame, "to_cuda_tensor", None)) +def _rgb_source_size(frame: object) -> tuple[int, int] | None: + """Return an HWC RGB frame's ``(width, height)`` without a host copy.""" + source = frame + to_cuda_tensor = getattr(frame, "to_cuda_tensor", None) + if callable(to_cuda_tensor): + try: + source = to_cuda_tensor() + except RuntimeError: + return None + shape = getattr(source, "shape", None) + if shape is None or len(shape) != 3: + return None + height, width = int(shape[0]), int(shape[1]) + if width <= 0 or height <= 0: + return None + return (width, height) + + def _as_rgb_host_uint8(frame: object) -> np.ndarray: to_numpy = getattr(frame, "to_numpy", None) if callable(to_numpy): diff --git a/integrations/omnidreams/omnidreams/interactive_drive/video_model/chunk_pipeline.py b/integrations/omnidreams/omnidreams/interactive_drive/video_model/chunk_pipeline.py index 7c3091bd1..b8776e70d 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/video_model/chunk_pipeline.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/video_model/chunk_pipeline.py @@ -45,6 +45,8 @@ def render_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk: ... def reset(self) -> None: ... + def set_postprocess_enabled(self, enabled: bool) -> None: ... + @dataclass(frozen=True) class ChunkRequest: @@ -296,6 +298,16 @@ def reset_command(backend: VideoModelBackend) -> bool: self._command_queue.put(reset_command) + def set_postprocess_enabled(self, enabled: bool) -> None: + """Toggle post-processing on the model worker thread. Non-blocking.""" + self._raise_worker_error_if_any() + + def toggle_command(backend: VideoModelBackend) -> bool: + backend.set_postprocess_enabled(enabled) + return True + + self._command_queue.put(toggle_command) + def shutdown(self) -> None: self._command_queue.put(_shutdown_command) self._thread.join() diff --git a/integrations/omnidreams/omnidreams/interactive_drive/video_model/local.py b/integrations/omnidreams/omnidreams/interactive_drive/video_model/local.py index f7795be62..d270b2be3 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/video_model/local.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/video_model/local.py @@ -45,3 +45,6 @@ def render_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk: def reset(self) -> None: self._backend.reset() self._is_first_chunk = True + + def set_postprocess_enabled(self, enabled: bool) -> None: + self._backend.set_postprocess_enabled(enabled) diff --git a/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py b/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py index f2f9c54db..676bb155d 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py @@ -11,6 +11,10 @@ import numpy as np import torch +from flashdreams.infra.postprocess import ( + VideoPostprocessChainConfig, + VideoPostprocessStream, +) from loguru import logger from omnidreams.interactive_drive.config import WorldModelProfileConfig from omnidreams.interactive_drive.world_model.manifest import WorldModelManifest @@ -116,10 +120,9 @@ def _build_pipeline_config( manifest: WorldModelManifest, profile: WorldModelProfileConfig ) -> Any: try: - from omnidreams.config import OMNIDREAMS_CONFIGS - from flashdreams.infra.config import derive_config from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig + from omnidreams.config import OMNIDREAMS_CONFIGS except ModuleNotFoundError as exc: raise RuntimeError( "The flashdreams and flashdreams-omnidreams packages are required " @@ -482,6 +485,7 @@ def __init__( *, offload_text_encoder: bool = False, pipeline_factory: PipelineFactory | None = None, + postprocess: VideoPostprocessChainConfig | None = None, ) -> None: self.manifest = manifest self._profile_config = profile or WorldModelProfileConfig() @@ -492,6 +496,9 @@ def __init__( self._precomputed_embeddings: dict[str, torch.Tensor | None] | None = None self._pending_finalization_index: int | None = None self._next_block_index = 0 + self._postprocess = postprocess or VideoPostprocessChainConfig() + self._postprocess_enabled = self._postprocess.is_enabled() + self._postprocess_stream: VideoPostprocessStream | None = None @property def pipeline(self) -> Any: @@ -628,6 +635,7 @@ def start( cache=self._cache, hdmap=self._condition_tensor(condition_frames), ) + video = self._postprocess_video(video, autoregressive_index=0) model_frames = self._video_tensor_to_frames(video) _synchronize_cuda_frame_event(model_frames) self._pending_finalization_index = 0 @@ -656,6 +664,9 @@ def continue_generation(self, condition_frames: list[object]) -> list[object]: cache=self._cache, hdmap=self._condition_tensor(condition_frames), ) + video = self._postprocess_video( + video, autoregressive_index=self._next_block_index + ) model_frames = self._video_tensor_to_frames(video) _synchronize_cuda_frame_event(model_frames) block_index = self._next_block_index @@ -669,6 +680,7 @@ def continue_generation(self, condition_frames: list[object]) -> list[object]: return model_frames def reset(self, *, clear_precomputed_embeddings: bool = False) -> None: + self._close_postprocess_stream() self._cache = None self._pending_finalization_index = None self._next_block_index = 0 @@ -680,12 +692,62 @@ def reset(self, *, clear_precomputed_embeddings: bool = False) -> None: ) def close(self) -> None: + self._close_postprocess_stream() if self._cache is not None and self._pending_finalization_index is not None: self.pipeline.finalize(self._pending_finalization_index, self._cache) self._pending_finalization_index = None self._cache = None self._pipeline = None + def set_postprocess_enabled(self, enabled: bool) -> None: + """Toggle the configured post-process chain between generated chunks.""" + enabled = bool(enabled) + if enabled and not self._postprocess.is_enabled(): + raise RuntimeError( + "Cannot enable post-processing without --postprocess-preset." + ) + if enabled == self._postprocess_enabled: + return + self._close_postprocess_stream() + self._postprocess_enabled = enabled + logger.info( + "[flashdreams-session] post-processing {} preset={!r}", + "enabled" if enabled else "disabled", + self._postprocess.preset, + ) + + def _postprocess_video( + self, video: torch.Tensor, *, autoregressive_index: int + ) -> torch.Tensor: + if not self._postprocess_enabled: + return video + if self._postprocess_stream is None: + self._postprocess_stream = VideoPostprocessStream( + postprocess=self._postprocess, + output_layout="bvtchw", + fps=self.manifest.fps, + per_view=False, + world_size=1, + collect_output=False, + move_to_cpu=False, + ) + processed = self._postprocess_stream.process( + video, autoregressive_index=autoregressive_index + ) + if processed.shape[2] != video.shape[2]: + raise RuntimeError( + "Interactive post-processing must emit one display frame for " + "each generated frame; got " + f"{processed.shape[2]} output frames for {video.shape[2]} inputs." + ) + return processed + + def _close_postprocess_stream(self) -> None: + if self._postprocess_stream is None: + return + self._postprocess_stream.finish() + self._postprocess_stream = None + def _initialize_cache(self, initial_rgb: object, prompt: str) -> Any: if self.manifest.synthetic_model: return self._initialize_synthetic_cache() diff --git a/integrations/omnidreams/tests/interactive_drive/test_chunk_pipeline.py b/integrations/omnidreams/tests/interactive_drive/test_chunk_pipeline.py index 02c327a24..e83994b83 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_chunk_pipeline.py +++ b/integrations/omnidreams/tests/interactive_drive/test_chunk_pipeline.py @@ -217,6 +217,17 @@ def test_chunk_pipeline_reset_invokes_backend_reset() -> None: assert backend.reset_calls == 1 +def test_chunk_pipeline_toggles_postprocess_on_worker() -> None: + backend = FakeVideoModelBackend(frames_per_render=1) + pipeline = ChunkPipeline(backend) + + pipeline.set_postprocess_enabled(False) + pipeline.set_postprocess_enabled(True) + pipeline.shutdown() + + assert backend.postprocess_enabled_calls == [False, True] + + def test_chunk_pipeline_reuses_model_across_scene_changes() -> None: backend = FakeVideoModelBackend(frames_per_render=1) pipeline = ChunkPipeline(backend) diff --git a/integrations/omnidreams/tests/interactive_drive/test_cli.py b/integrations/omnidreams/tests/interactive_drive/test_cli.py index 47b37ef05..b0fa70baa 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_cli.py +++ b/integrations/omnidreams/tests/interactive_drive/test_cli.py @@ -14,3 +14,15 @@ def test_offload_text_encoder_flag_enables() -> None: args = build_parser().parse_args(["--offload-text-encoder"]) assert args.offload_text_encoder is True + + +def test_postprocess_preset_defaults_disabled() -> None: + args = build_parser().parse_args([]) + + assert args.postprocess_preset == "" + + +def test_postprocess_preset_accepts_rtx_super_resolution() -> None: + args = build_parser().parse_args(["--postprocess-preset", "rtx-super-resolution"]) + + assert args.postprocess_preset == "rtx-super-resolution" diff --git a/integrations/omnidreams/tests/interactive_drive/test_presenter.py b/integrations/omnidreams/tests/interactive_drive/test_presenter.py index 36f80b678..514a7dcc9 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_presenter.py +++ b/integrations/omnidreams/tests/interactive_drive/test_presenter.py @@ -452,6 +452,46 @@ def close(self) -> None: assert presenter._retired_cuda_hud_interops == [interop] +def test_hud_postprocess_control_toggles_configured_preset() -> None: + presenter = _hud_presenter_without_window() + calls: list[bool] = [] + presenter._postprocess_rect = (10, 20, 110, 52) + presenter._panel_chrome_cache_key = object() + presenter._panel_chrome_cache = object() + presenter._scene_dropdown_open = False + presenter._variant_dropdown_open = False + presenter.set_postprocess_control( + preset="rtx-super-resolution", + enabled=True, + callback=calls.append, + ) + + presenter._handle_click((20, 30)) + presenter._handle_click((20, 30)) + + assert calls == [False, True] + assert presenter._postprocess_enabled is True + + +def test_hud_scene_dropdown_blocks_underlying_upsample_toggle() -> None: + presenter = _hud_presenter_without_window() + calls: list[bool] = [] + presenter._postprocess_rect = (10, 20, 110, 52) + presenter._postprocess_preset = "rtx-super-resolution-ultra" + presenter._postprocess_enabled = True + presenter._postprocess_callback = calls.append + presenter._scene_dropdown_open = True + presenter._variant_dropdown_open = False + presenter._scene_item_rects = [] + presenter._scene_header_rect = None + presenter._scene_selection_locked_probe = lambda: False + + presenter._handle_click((20, 30)) + + assert calls == [] + assert presenter._postprocess_enabled is True + + def test_hud_resize_uses_actual_window_size_without_model_resolution_clamp() -> None: presenter = _hud_presenter_without_window() presenter._pending_resize = None @@ -461,6 +501,55 @@ def test_hud_resize_uses_actual_window_size_without_model_resolution_clamp() -> assert presenter._pending_resize == (320, 200) +def test_hud_auto_sizes_window_to_native_model_frame_resolution() -> None: + presenter = _hud_presenter_without_window() + resize_calls: list[tuple[int, int]] = [] + presenter._auto_sized_camera_src_size = None + presenter._pending_resize = None + presenter._window = SimpleNamespace( + size=SimpleNamespace(x=1000, y=600), + resize=lambda width, height: resize_calls.append((width, height)), + ) + + resized = presenter._resize_window_for_native_model_frame( + np.zeros((704, 1280, 3), dtype=np.uint8) + ) + + assert resized is True + assert resize_calls == [(1780, 704)] + assert presenter._pending_resize == (1780, 704) + + +def test_hud_keeps_larger_canvas_when_model_resolution_shrinks() -> None: + presenter = _hud_presenter_without_window() + resize_calls: list[tuple[int, int]] = [] + presenter._auto_sized_camera_src_size = (1280, 704) + presenter._pending_resize = None + presenter._window = SimpleNamespace( + size=SimpleNamespace(x=1500, y=800), + resize=lambda width, height: resize_calls.append((width, height)), + ) + + assert not presenter._resize_window_for_native_model_frame( + np.zeros((704, 1280, 3), dtype=np.uint8) + ) + assert not presenter._resize_window_for_native_model_frame( + np.zeros((352, 640, 3), dtype=np.uint8) + ) + + assert resize_calls == [] + + +def test_hud_centers_smaller_camera_at_native_resolution() -> None: + presenter = _hud_presenter_without_window() + presenter._latest_camera_src_size = (640, 352) + presenter._configured_size = (1920, 1080) + + fit = presenter._compute_camera_fit() + + assert fit == (640, 352, 390, 364) + + def test_hud_cuda_submit_abandons_ready_buffer_if_resize_retires_interop() -> None: presenter = _hud_presenter_without_window() mark_calls = 0 diff --git a/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py b/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py index 7d49b73cd..28424c024 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py +++ b/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py @@ -11,6 +11,7 @@ import omnidreams.interactive_drive.world_model.flashdreams_adapter as adapter_module import pytest import torch +from flashdreams.infra.postprocess import VideoPostprocessChainConfig from omnidreams.interactive_drive.config import WorldModelProfileConfig from omnidreams.interactive_drive.world_model.flashdreams_adapter import ( FlashdreamsWorldModelSession, @@ -276,6 +277,58 @@ def test_session_uses_flashdreams_pipeline_for_rollout() -> None: assert fake_pipeline.finalize_calls == [(0, "cache"), (1, "cache")] +def test_session_postprocesses_local_frames_and_supports_live_toggle( + monkeypatch, +) -> None: + streams: list[object] = [] + + class _FakePostprocessStream: + def __init__(self, **kwargs: object) -> None: + self.kwargs = kwargs + self.calls: list[int] = [] + self.finished = False + streams.append(self) + + def process( + self, output: torch.Tensor, *, autoregressive_index: int + ) -> torch.Tensor: + self.calls.append(autoregressive_index) + return output + + def finish(self) -> None: + self.finished = True + + monkeypatch.setattr( + adapter_module, "VideoPostprocessStream", _FakePostprocessStream + ) + fake_pipeline = _FakePipeline() + session = FlashdreamsWorldModelSession( + _manifest(), + pipeline_factory=lambda manifest, profile: fake_pipeline, + postprocess=VideoPostprocessChainConfig(preset="fake-preset"), + ) + session.warmup_model() + initial_rgb = np.zeros((2, 3, 3), dtype=np.uint8) + first_conditions = [np.zeros((2, 3, 3), dtype=np.uint8) for _ in range(5)] + next_conditions = [np.zeros((2, 3, 3), dtype=np.uint8) for _ in range(8)] + + session.start(initial_rgb, first_conditions, "demo prompt") + first_stream = streams[0] + assert first_stream.calls == [0] + assert first_stream.kwargs["output_layout"] == "bvtchw" + assert first_stream.kwargs["collect_output"] is False + + session.set_postprocess_enabled(False) + assert first_stream.finished is True + session.continue_generation(next_conditions) + assert len(streams) == 1 + + session.set_postprocess_enabled(True) + session.continue_generation(next_conditions) + assert len(streams) == 2 + assert streams[1].calls == [2] + + def test_session_synthetic_model_initializes_cache_from_synthetic_embeddings() -> None: fake_pipeline = _FakeSyntheticPipeline() manifest = replace(_manifest(), synthetic_model=True, resolution_wh=(64, 32)) From 1b1f9491a0604eded9520866b4a973e9a9513dfb Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Wed, 22 Jul 2026 12:19:44 -0700 Subject: [PATCH 4/5] Adapt RTX wiring to shared serving infrastructure Signed-off-by: Gangzheng Tong --- flashdreams/flashdreams/infra/runner_io.py | 83 ++++++++++++++++++- flashdreams/tests/test_runner_io.py | 27 +++--- .../interactive_drive/backends/world_model.py | 2 +- .../omnidreams/interactive_drive/cli.py | 4 +- .../world_model/flashdreams_adapter.py | 11 +-- .../test_world_model_adapter.py | 7 +- .../tests/test_runner_video_output.py | 13 ++- 7 files changed, 118 insertions(+), 29 deletions(-) diff --git a/flashdreams/flashdreams/infra/runner_io.py b/flashdreams/flashdreams/infra/runner_io.py index 465f98e88..aaed029a3 100644 --- a/flashdreams/flashdreams/infra/runner_io.py +++ b/flashdreams/flashdreams/infra/runner_io.py @@ -18,6 +18,8 @@ from __future__ import annotations import json +import shutil +import subprocess from collections.abc import Callable from pathlib import Path from typing import Any, Literal, TypeAlias @@ -314,13 +316,88 @@ def write_video_tensor( layout: VideoTensorLayout, install_hint: str = DEFAULT_RUNNER_INSTALL_HINT, ) -> Path: - """Write a ``[-1, 1]`` video tensor as an MP4 and return the path.""" - media = _import_mediapy("Writing videos", install_hint=install_hint) + """Write a ``[-1, 1]`` video tensor through the host FFmpeg executable.""" + del install_hint # Kept for API compatibility with existing runner call sites. path = Path(path) - media.write_video(str(path), video_tensor_to_uint8(video, layout=layout), fps=fps) + path.parent.mkdir(parents=True, exist_ok=True) + frames = video_tensor_to_uint8(video, layout=layout) + num_frames, height, width, channels = frames.shape + if channels != 3: + raise ValueError( + "expected video frames with shape [T, H, W, C=3], " + f"got {tuple(frames.shape)}" + ) + + cmd = [ + _find_ffmpeg_binary(), + "-y", + "-loglevel", + "error", + "-f", + "rawvideo", + "-vcodec", + "rawvideo", + "-pix_fmt", + "rgb24", + "-s", + f"{width}x{height}", + "-r", + str(fps), + "-i", + "-", + "-an", + "-vcodec", + "libx264", + "-vf", + "pad=ceil(iw/2)*2:ceil(ih/2)*2", + "-pix_fmt", + "yuv420p", + "-crf", + "18", + "-preset", + "medium", + "-movflags", + "+faststart", + str(path), + ] + process = subprocess.Popen(cmd, stdin=subprocess.PIPE, stderr=subprocess.PIPE) + assert process.stdin is not None + assert process.stderr is not None + write_error: BrokenPipeError | None = None + try: + for frame_index in range(num_frames): + process.stdin.write(frames[frame_index].tobytes()) + except BrokenPipeError as exc: + write_error = exc + finally: + try: + process.stdin.close() + except BrokenPipeError as exc: + if write_error is None: + write_error = exc + stderr = process.stderr.read().decode("utf-8", errors="replace") + returncode = process.wait() + + if write_error is not None: + raise RuntimeError( + f"ffmpeg closed while writing {path} (exit code {returncode}): {stderr}" + ) from write_error + if returncode != 0: + raise RuntimeError(f"ffmpeg failed while writing {path}: {stderr}") return path +def _find_ffmpeg_binary() -> str: + """Find the host-provided FFmpeg executable or fail with installation guidance.""" + ffmpeg_bin = shutil.which("ffmpeg") + if ffmpeg_bin is None: + raise RuntimeError( + "Writing the output video requires an ffmpeg executable installed " + "on the host and available on PATH." + ) + return ffmpeg_bin + + def _import_mediapy(action: str, *, install_hint: str) -> Any: """Import ``mediapy`` with a runner-specific dependency hint.""" try: diff --git a/flashdreams/tests/test_runner_io.py b/flashdreams/tests/test_runner_io.py index 30a61d4d0..57a6ce44e 100644 --- a/flashdreams/tests/test_runner_io.py +++ b/flashdreams/tests/test_runner_io.py @@ -17,6 +17,7 @@ from __future__ import annotations +import io import sys import types from pathlib import Path @@ -234,17 +235,22 @@ def from_path(cls, path: str) -> "VideoMetadata": assert calls == ["clip.mp4"] -def test_write_video_tensor_lazy_imports_mediapy( +def test_write_video_tensor_streams_to_host_ffmpeg( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: - calls: list[tuple[str, np.ndarray, int]] = [] - fake_media = types.ModuleType("mediapy") + commands: list[list[str]] = [] + process = types.SimpleNamespace( + stdin=io.BytesIO(), + stderr=io.BytesIO(), + wait=lambda: 0, + ) - def write_video(path: str, frames: np.ndarray, *, fps: int) -> None: - calls.append((path, frames, fps)) + def popen(cmd: list[str], **_kwargs: Any) -> Any: + commands.append(cmd) + return process - setattr(fake_media, "write_video", write_video) - monkeypatch.setitem(sys.modules, "mediapy", fake_media) + monkeypatch.setattr(runner_io, "_find_ffmpeg_binary", lambda: "ffmpeg") + monkeypatch.setattr(runner_io.subprocess, "Popen", popen) out_path = tmp_path / "out.mp4" returned = write_video_tensor( @@ -255,10 +261,9 @@ def write_video(path: str, frames: np.ndarray, *, fps: int) -> None: ) assert returned == out_path - assert calls[0][0] == str(out_path) - assert calls[0][1].shape == (1, 2, 2, 3) - assert calls[0][1].dtype == np.uint8 - assert calls[0][2] == 16 + assert commands[0][commands[0].index("-s") + 1] == "2x2" + assert commands[0][commands[0].index("-r") + 1] == "16" + assert commands[0][-1] == str(out_path) def test_load_first_frame_tensor_uses_requested_resize_interpolation( diff --git a/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py b/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py index 73ae269c0..5fecc0248 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py @@ -9,7 +9,6 @@ from pathlib import Path import numpy as np -from flashdreams.infra.postprocess import VideoPostprocessChainConfig from loguru import logger from omnidreams.interactive_drive.backends.base import RenderBackend from omnidreams.interactive_drive.config import ( @@ -33,6 +32,7 @@ from PIL import Image from flashdreams.infra.acceleration.prewarm import run_timed_prewarm +from flashdreams.infra.postprocess import VideoPostprocessChainConfig _FIRST_STEADY_STATE_WARMUP_MESSAGE = "Optimizing world model..." diff --git a/integrations/omnidreams/omnidreams/interactive_drive/cli.py b/integrations/omnidreams/omnidreams/interactive_drive/cli.py index a6c0e2251..a5b109418 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/cli.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/cli.py @@ -7,8 +7,6 @@ from dataclasses import replace from pathlib import Path -from flashdreams.infra.postprocess import VideoPostprocessChainConfig -from flashdreams.plugins.registry import discover_postprocess_presets from loguru import logger from omnidreams.hf_org import DEFAULT_HF_ORG, apply_cli_to_env from omnidreams.hf_org import ENV_VAR as _HF_ORG_ENV_VAR @@ -27,6 +25,8 @@ from omnidreams.interactive_drive.world_model.manifest import load_world_model_manifest from omnidreams.scenes import local_scene_archive_path +from flashdreams.infra.postprocess import VideoPostprocessChainConfig +from flashdreams.plugins.registry import discover_postprocess_presets from flashdreams.serving.realtime.timing import TraceSink # Package root (from this file's location) so packaged-asset defaults below diff --git a/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py b/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py index 676bb155d..838c19b90 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py @@ -11,10 +11,6 @@ import numpy as np import torch -from flashdreams.infra.postprocess import ( - VideoPostprocessChainConfig, - VideoPostprocessStream, -) from loguru import logger from omnidreams.interactive_drive.config import WorldModelProfileConfig from omnidreams.interactive_drive.world_model.manifest import WorldModelManifest @@ -32,6 +28,10 @@ ) from flashdreams.infra.acceleration.frame_prefetch import LazyCudaFrame from flashdreams.infra.acceleration.prewarm import run_timed_prewarm +from flashdreams.infra.postprocess import ( + VideoPostprocessChainConfig, + VideoPostprocessStream, +) PipelineFactory = Callable[[WorldModelManifest, WorldModelProfileConfig], Any] _VIEW_NAMES = ["camera_front_wide_120fov"] @@ -120,9 +120,10 @@ def _build_pipeline_config( manifest: WorldModelManifest, profile: WorldModelProfileConfig ) -> Any: try: + from omnidreams.config import OMNIDREAMS_CONFIGS + from flashdreams.infra.config import derive_config from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig - from omnidreams.config import OMNIDREAMS_CONFIGS except ModuleNotFoundError as exc: raise RuntimeError( "The flashdreams and flashdreams-omnidreams packages are required " diff --git a/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py b/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py index 28424c024..070bdac08 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py +++ b/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py @@ -11,7 +11,6 @@ import omnidreams.interactive_drive.world_model.flashdreams_adapter as adapter_module import pytest import torch -from flashdreams.infra.postprocess import VideoPostprocessChainConfig from omnidreams.interactive_drive.config import WorldModelProfileConfig from omnidreams.interactive_drive.world_model.flashdreams_adapter import ( FlashdreamsWorldModelSession, @@ -24,6 +23,8 @@ SyntheticWorldModelAssets, ) +from flashdreams.infra.postprocess import VideoPostprocessChainConfig + class _FakePipeline: def __init__(self) -> None: @@ -172,8 +173,8 @@ def test_build_pipeline_config_can_select_native_vae_encoder() -> None: ) assert config.image_encoder.native_vae_acceleration == "required" assert config.image_encoder.native_vae_backend == "fp8" - assert ( - config.image_encoder.native_vae_fp8_state_path == "/tmp/lightvae-fp8-state.pt" + assert Path(config.image_encoder.native_vae_fp8_state_path) == Path( + "/tmp/lightvae-fp8-state.pt" ) assert config.encoder.native_vae_acceleration == "required" assert config.encoder.native_vae_backend == "fp8" diff --git a/integrations/omnidreams/tests/test_runner_video_output.py b/integrations/omnidreams/tests/test_runner_video_output.py index 82ea653e1..8924798b7 100644 --- a/integrations/omnidreams/tests/test_runner_video_output.py +++ b/integrations/omnidreams/tests/test_runner_video_output.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""CPU tests for the OmniDreams runner's streaming video writer.""" +"""CPU tests for the shared streaming video writer used by OmniDreams.""" from __future__ import annotations @@ -9,10 +9,11 @@ from pathlib import Path from typing import Any -import omnidreams.runner as runner_module import pytest import torch +import flashdreams.infra.runner_io as runner_module + pytestmark = pytest.mark.ci_cpu @@ -83,7 +84,9 @@ def _popen(cmd: list[str], **_kwargs: Any) -> _FakeProcess: monkeypatch.setattr(runner_module, "_find_ffmpeg_binary", lambda: "ffmpeg") monkeypatch.setattr(runner_module.subprocess, "Popen", _popen) - runner_module._write_video(_canvas(), tmp_path / "odd.mp4", fps=24) + runner_module.write_video_tensor( + _canvas(), tmp_path / "odd.mp4", fps=24, layout="thwc" + ) assert commands[0][commands[0].index("-vf") + 1] == ( "pad=ceil(iw/2)*2:ceil(ih/2)*2" @@ -110,7 +113,9 @@ def test_write_video_reaps_ffmpeg_and_preserves_broken_pipe_diagnostic( ) with pytest.raises(RuntimeError, match="exit code 7.*encoder rejected frame"): - runner_module._write_video(_canvas(), tmp_path / "broken.mp4", fps=24) + runner_module.write_video_tensor( + _canvas(), tmp_path / "broken.mp4", fps=24, layout="thwc" + ) assert process.stdin.closed assert process.wait_calls == 1 From 1278c916552c36ef3fb966fb719b819687e9e931 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Wed, 22 Jul 2026 14:02:06 -0700 Subject: [PATCH 5/5] Add optional NVIDIA VFX dependency extra Signed-off-by: Gangzheng Tong --- THIRD-PARTY-NOTICES | 12 ++++++++++++ docs/source/api/cli.rst | 6 +++--- flashdreams/pyproject.toml | 6 ++++++ uv.lock | 36 +++++++++++++++--------------------- 4 files changed, 36 insertions(+), 24 deletions(-) diff --git a/THIRD-PARTY-NOTICES b/THIRD-PARTY-NOTICES index bc8f3ff47..21c389ba9 100644 --- a/THIRD-PARTY-NOTICES +++ b/THIRD-PARTY-NOTICES @@ -92,6 +92,18 @@ SpargeAttn Apache-2.0 https://github.com/thu-ml/SpargeAttn thirdparty_sources.json at bfd980b781784c04ad6a53e7ee657c0645d99171. +================================================================================ +Optional post-processing: RTX Video Super Resolution +================================================================================ + +nvidia-vfx LicenseRef-NvidiaProprietary + https://pypi.org/project/nvidia-vfx/ + +The ``rtx-postprocess`` extra installs NVIDIA's proprietary Python bindings +for RTX Video Super Resolution and deblur. It is optional, dynamically loaded +only when an RTX preset is selected, and is not redistributed as source in +this repository. + ================================================================================ Optional integration: integrations/lingbot ================================================================================ diff --git a/docs/source/api/cli.rst b/docs/source/api/cli.rst index 7dc6102fe..2e99ab447 100644 --- a/docs/source/api/cli.rst +++ b/docs/source/api/cli.rst @@ -65,9 +65,9 @@ one with ``--postprocess.preset``: --postprocess.preset rtx-super-resolution The ``rtx-super-resolution`` preset wraps NVIDIA VFX Python bindings for RTX -Video Super Resolution. Install ``nvidia-vfx`` from NVIDIA's Python package -index in the active environment and run on a supported RTX GPU before selecting -this preset. +Video Super Resolution. Install the optional dependency with +``uv pip install 'flashdreams[rtx-postprocess]'`` and run on a supported RTX GPU +before selecting this preset. See also -------- diff --git a/flashdreams/pyproject.toml b/flashdreams/pyproject.toml index 8058985ae..a27ea93a0 100644 --- a/flashdreams/pyproject.toml +++ b/flashdreams/pyproject.toml @@ -113,6 +113,12 @@ runners = [ "opencv-python-headless>=4.5", "scipy>=1.11", ] +# Optional NVIDIA VFX runtime used only by the RTX postprocess presets. +# Kept out of the default install because it is proprietary and requires +# supported NVIDIA RTX hardware. +rtx-postprocess = [ + "nvidia-vfx==0.1.0.1", +] serving = [ "aiohttp>=3.9", "aiortc>=1.9", diff --git a/uv.lock b/uv.lock index 2b0d03215..b78249b92 100644 --- a/uv.lock +++ b/uv.lock @@ -1247,6 +1247,9 @@ examples = [ { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, { name = "scipy", version = "1.18.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, ] +rtx-postprocess = [ + { name = "nvidia-vfx" }, +] runners = [ { name = "mediapy" }, { name = "opencv-python-headless" }, @@ -1289,6 +1292,7 @@ requires-dist = [ { name = "mediapy", marker = "extra == 'runners'", specifier = ">=1.1" }, { name = "numpy", specifier = ">=1.24,<2.5" }, { name = "nvidia-ml-py", specifier = ">=12.0" }, + { name = "nvidia-vfx", marker = "extra == 'rtx-postprocess'", specifier = "==0.1.0.1" }, { name = "opencv-python-headless", marker = "extra == 'examples'", specifier = ">=4.5" }, { name = "opencv-python-headless", marker = "extra == 'runners'", specifier = ">=4.5" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, @@ -1306,7 +1310,7 @@ requires-dist = [ { name = "tyro", specifier = ">=1.0" }, { name = "urllib3", specifier = ">=2.7.0" }, ] -provides-extras = ["dev", "examples", "runners", "serving"] +provides-extras = ["dev", "examples", "runners", "rtx-postprocess", "serving"] [package.metadata.requires-dev] cuda12 = [ @@ -3200,26 +3204,10 @@ name = "numpy" version = "2.4.6" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version == '3.13.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version == '3.12.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform == 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform != 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform != 'win32' and extra != 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform == 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version >= '3.14' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.13.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.12.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", - "python_full_version == '3.11.*' and sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12' and extra != 'group-11-flashdreams-cuda13'", + "python_full_version >= '3.12' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version >= '3.12' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform != 'win32'", ] sdist = { url = "https://files.pythonhosted.org/packages/d0/ad/fed0499ce6a338d2a03ebae59cd15093910c8875328855781952abf6c2fe/numpy-2.4.6.tar.gz", hash = "sha256:f3a3570c4a2a16746ac2c31a7c7c7b0c186b95ce902e33db6f28094ed7387dda", size = 20735807, upload-time = "2026-05-18T23:37:14.07Z" } wheels = [ @@ -3731,6 +3719,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9f/99/4c9c0c329bf9fc125008c3b54c7c94c0023518d06fc025ae36431375e1fe/nvidia_nvtx_cu12-12.8.90-py3-none-win_amd64.whl", hash = "sha256:619c8304aedc69f02ea82dd244541a83c3d9d40993381b3b590f1adaed3db41e", size = 56492, upload-time = "2025-03-07T01:52:24.69Z" }, ] +[[package]] +name = "nvidia-vfx" +version = "0.1.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/37/b4/58e1bbb8d6fc9ed786564d0878314ea2f2cd458c84861a5130927e431ff6/nvidia_vfx-0.1.0.1.tar.gz", hash = "sha256:8a26bae3a967a2ce29040f17ba9d75e106f3d0c68016d440a77ed9c7eb05daae", size = 2673, upload-time = "2026-03-09T19:29:40.556Z" } + [[package]] name = "onnx" version = "1.22.0"