Add RTX super resolution postprocessor

[gtong-nv]

Benchmark Report

Summary

Add optional NVIDIA RTX Video Super Resolution and deblur postprocessing to FlashDreams, then wire it through all OmniDreams output paths:

• offline flashdreams-run generation;
• the OmniDreams WebRTC server and request-session browser UI; and
• the local interactive-drive Vulkan/HUD window.

The feature remains opt-in. Existing runs produce the same native-resolution output unless a postprocess preset is selected. The proprietary NVIDIA VFX bindings live in the dedicated flashdreams[rtx-postprocess] optional extra, so they are not part of a default install. Video persistence requires a host-installed ffmpeg executable or fails with an actionable error; there is no imageio-ffmpeg dependency or fallback.

This branch is rebased onto origin/main at eadf2596 and adapts the feature to main's shared runner-I/O and reusable packaged-WebRTC infrastructure.

Motivation

OmniDreams normally generates 1280×704 RGB frames. Displaying those frames in a larger browser or local window only performs conventional presentation scaling, which makes the postprocess toggle appear ineffective and does not expose the higher-resolution VSR result.

This change makes postprocessing part of the decoded-frame pipeline and preserves its native output size through presentation. With RTX VSR enabled, clients receive and display a true 2560×1408 raster. With it disabled, the 1280×704 stream remains centered at native size while the surrounding HUD canvas stays visible.

What changed

RTX postprocessor

• Add a generic RTXVideoSuperResolutionPostProcessorConfig, processor factory, and per-stream session under flashdreams.infra.postprocess.
• Use DLPack/CUDA interop for decoded GPU tensors and preserve the FlashDreams [-1, 1] video contract.
• Validate device, layout, scaling, quality mode, and distributed execution before processing.
• Keep mutable VFX state inside the stream session and support explicit teardown/flush.
• Dynamically import nvvfx.VideoSuperRes only when an RTX preset is selected, with a clear runtime error when nvidia-vfx is unavailable.
• Register three presets:

Preset	VFX quality	Output size
rtx-super-resolution	HIGH	2× width and height
rtx-super-resolution-ultra	ULTRA	2× width and height
rtx-deblur-ultra	DEBLUR_ULTRA	unchanged

Deblur and super resolution are exposed as separate presets because the NVIDIA VFX API represents them as distinct quality modes; the same VideoSuperRes invocation cannot request DEBLUR_ULTRA and 2× scaling simultaneously.

Offline generation

• Route every decoded OmniDreams AR chunk through one VideoPostprocessStream.
• Preserve one postprocessor session per camera view.
• Flush the stream once at end of rollout and include postprocess CUDA timing in the per-AR-step stats JSON.
• Write postprocessed RGB directly instead of combining it with the native-size HDMap diagnostic pane, avoiding mixed-resolution canvases.
• Port video persistence to main's shared runner_io.write_video_tensor helper.
• Stream MP4 encoding through host ffmpeg, pad odd dimensions for YUV420p, reap failed encoder processes, and preserve FFmpeg diagnostics on broken pipes.
• Do not fall back to imageio-ffmpeg; missing host FFmpeg fails loudly.

Example:

uv run --with "nvidia-vfx==0.1.0.1" `
  --package flashdreams-omnidreams `
  flashdreams-run omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae `
  --example-data True `
  --postprocess.preset rtx-super-resolution

WebRTC demo

• Add --postprocess-preset as the server-side default.
• Advertise discovered presets through /api/postprocess/options.
• Allow the request-session page to select a preset, inherit the CLI default, or explicitly disable processing for the next connection.
• Reject session-option changes while a peer session is active, preventing state changes during an in-flight rollout.
• Reset and close the postprocess stream at session boundaries.
• Include the active preset in chunk metadata.
• Configure the OmniDreams-specific API routes through main's new create_packaged_webrtc_app(..., configure_app=...) hook.
• Fall back from symlink to hard link/copy when staging ClipGT parquet files, allowing the demo to run on Windows without Developer Mode or elevated symlink privileges.

Example:

$env:USE_LIBUV = "0"
uv run --with "nvidia-vfx==0.1.0.1" `
  --package flashdreams-omnidreams `
  python -m omnidreams.webrtc.server `
  --postprocess-preset rtx-super-resolution

USE_LIBUV=0 is required only for Windows PyTorch builds without libuv support; it is not part of the postprocessor itself.

Local interactive demo and HUD

• Add --postprocess-preset to the interactive-drive CLI.
• Keep the preset initialized for the model worker and expose a live enable/disable control.
• Label the HUD control Upsample 2x with ON, OFF, or N/A state.
• Disable the upsample control while either scene dropdown is open so dropdown clicks cannot toggle processing underneath the menu.
• Propagate toggles from the presenter through the chunk pipeline to the world-model postprocess stream.
• Recreate the postprocess stream at rollout/reset boundaries while preserving the selected state.
• Present model frames at native resolution instead of scaling them up to fill the camera pane.
• Grow the window when the processed source changes to a larger native resolution; do not shrink the window when processing is disabled.
• Center smaller native frames inside the existing camera canvas so the HUD remains visible.
• Preserve user window resizes until the model output resolution changes again.

Performance benchmark

The benchmark used matched offline OmniDreams generations with and without rtx-super-resolution.

Setup

Item	Value
GPU	NVIDIA RTX PRO 6000 Blackwell Workstation Edition, 97,887 MiB
Driver	591.86
PyTorch	2.12.1+cu130
nvidia-vfx	0.1.0.1
FFmpeg	8.1.2 host installation
Runner	omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae
Example scene UUID	239560dc-33d1-11ef-9720-00044bcbccac
Seed	42
Workload	6 AR chunks / 45 frames / 30 FPS
Native output	1280×704
RTX output	2560×1408
Steady-state sample	AR chunks 3–5; chunks 0–2 excluded for compile/CUDA-graph warmup

Both variants ran in fresh processes with the same dependency environment, input, prompt, seed, model configuration, and frame count. RTX effective latency is the pipeline's total_ms plus the separately measured postprocess event.

Results

Metric, steady state	Native	RTX VSR	Delta
Chunk latency, median	600.61 ms	610.69 ms	+10.08 ms / +1.68%
Chunk latency, mean	616.88 ms	629.63 ms	+12.75 ms / +2.07%
Chunk latency, p90	640.62 ms	656.86 ms	+16.24 ms / +2.54%
Effective throughput, median	13.32 FPS	13.10 FPS	-1.65%
RTX stage, median	—	8.94 ms/chunk	1.12 ms/frame
RTX stage, p90	—	9.33 ms/chunk	1.17 ms/frame
CUDA memory reserved, median	40.00 GiB	40.67 GiB	+0.67 GiB observed

Per measured chunk:

AR chunk	Native pipeline	RTX pipeline	RTX stage	RTX effective
3	650.62 ms	658.98 ms	9.42 ms	668.40 ms
4	599.41 ms	601.81 ms	8.89 ms	610.69 ms
5	600.61 ms	600.86 ms	8.94 ms	609.80 ms

The first RTX call took 1,599.98 ms for five frames because it created and initialized the VFX session. Subsequent calls stabilized at 8.59–9.72 ms for eight frames. Long-lived demos should therefore initialize/warm the selected processor before latency-sensitive interaction when possible.

Fresh-process wall times were 133.71 seconds native and 103.04 seconds RTX, but they are intentionally not used as a performance conclusion: the native run executed first and populated shared checkpoint/compile caches, and wall time also includes model loading, graph compilation, input decoding, and MP4 encoding.

The benchmark was captured on pre-rebase feature commit 34bbd2ae. The RTX execution path and measured preset configuration are unchanged by the rebase; post-rebase CPU contract tests and static checks are reported below. A repeated GPU sweep would still be appropriate before treating the figures as a release SLA.

Output quality analysis

There is no native high-resolution ground truth for this generated sequence. The quality comparison therefore uses a deterministic Lanczos 2× upscale of the exact 1280×704 generated output as the conventional-scaling reference and asks:

1. how much RTX differs from conventional scaling;
2. whether it increases measurable edge/high-frequency response; and
3. whether downsampled RTX preserves the native scene.

Metrics cover all 45 decoded H.264 CRF 18 frames:

Metric	Result	Meaning
RTX vs Lanczos PSNR, mean / median	37.75 / 37.95 dB	Close, non-identical outputs
RTX vs Lanczos luma SSIM, mean / median	0.9681 / 0.9708	Strong structural agreement
Downsampled RTX vs native PSNR, mean / median	38.48 / 38.58 dB	Native content mostly preserved
Downsampled RTX vs native luma SSIM, mean / median	0.9688 / 0.9734	Native structure mostly preserved
Laplacian variance, Lanczos / RTX mean	6.07 / 19.32	+218.4% high-frequency response
Sobel edge energy, Lanczos / RTX mean	9.39 / 10.69	+13.9% edge strength
Frame-to-frame MAE, Lanczos / RTX mean	2.190 / 2.307	+5.4% temporal pixel change

Interpretation

RTX VSR is working: it produces four times as many output pixels and materially strengthens fine contours. In a 1:1 crop, the clearest changes appear around vehicle bodywork and wheels, lane markings, vegetation, and utility-pole edges.

The difference remains subtle when the full 2560×1408 image is reduced to the same on-screen size as 1280×704. That is expected: display downscaling discards much of the additional edge information. This is why the HUD now preserves native output resolution and expands the window only when the processed frame is larger.

The sharpness metrics should not be read as proof that RTX recovered ground-truth detail. Increased Laplacian/Sobel response can include useful detail, stronger existing contours, ringing, compression residue, or noise. The 5.4% increase in temporal frame delta may reflect sharper moving detail, temporally varying enhancement, or both; this 1.5-second sample is too short to label it as flicker.

Both comparison videos were independently encoded as H.264 CRF 18/YUV420p, so the measurements include small codec differences. A reference-grade follow-up should start with known high-resolution footage, downsample it to the model resolution, apply VSR, and compare the result against the original high-resolution source.

The complete local benchmark report, logs, videos, full-frame comparison, close-up crop, and machine-readable metrics were generated separately in artifacts/omnidreams_rtx_sr_benchmark_20260722 and committed on dev/gtong/rtx-ss-report as fea5f1db.

Dependency and licensing behavior

• nvidia-vfx==0.1.0.1 is isolated in the flashdreams[rtx-postprocess] optional extra and recorded in the lockfile; it is not installed by default.
• Importing FlashDreams does not import or require nvidia-vfx; selecting an RTX preset performs the guarded import.
• Installed packages opt in with uv pip install 'flashdreams[rtx-postprocess]'; workspace one-off commands can continue to use uv run --with "nvidia-vfx==0.1.0.1" ....
• The host must provide ffmpeg on PATH. Missing FFmpeg raises a clear RuntimeError.
• There is no imageio-ffmpeg dependency or fallback, avoiding redistribution of an FFmpeg binary through this Python dependency path.

This keeps optional proprietary/runtime components outside the core Apache-2.0 package. Deployments remain responsible for reviewing and accepting the licenses of their separately installed NVIDIA VFX package and host FFmpeg build.

Validation

Post-rebase focused CPU suite:

130 passed in 85.48s

Command:

uv run --package flashdreams-omnidreams pytest -q `
  flashdreams/tests/test_runner_io.py `
  flashdreams/tests/test_postprocess_presets.py `
  flashdreams/tests/test_rtx_super_resolution_postprocess.py `
  flashdreams/tests/test_postprocess_stream.py `
  flashdreams/tests/test_runner_postprocess.py `
  integrations/omnidreams/tests/test_runner_video_output.py `
  integrations/omnidreams/tests/test_webrtc_runtime.py `
  integrations/omnidreams/tests/interactive_drive/test_chunk_pipeline.py `
  integrations/omnidreams/tests/interactive_drive/test_cli.py `
  integrations/omnidreams/tests/interactive_drive/test_presenter.py `
  integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py

Repository hooks over all PR files:

Trim Trailing Whitespace ........ passed
Fix End of Files ................ passed
ruff check/fix .................. passed
ruff format ..................... passed
validate uv lockfile ............ passed
sync package versions ........... passed
ty type check ................... passed

GPU/manual coverage completed before the rebase:

• matched offline native and RTX VSR generation;
• 2× output-dimension verification;
• per-chunk CUDA timing capture;
• full-frame and close-up visual inspection; and
• browser/local-window toggle behavior during iterative development.

Limitations and follow-ups

• RTX processing requires a supported NVIDIA GPU and a separately installed compatible nvidia-vfx package.
• The first VFX session initialization is expensive relative to steady state.
• rtx-deblur-ultra is same-resolution; it is not composited with the 2× presets in this PR.
• Quality measurements do not use high-resolution ground truth and should not be treated as reconstruction-accuracy claims.
• Multi-run, counterbalanced GPU measurements and a longer temporal-quality clip are recommended before publishing release-level performance/quality guarantees.

Reviewer guide

Suggested review order:

1. flashdreams/flashdreams/infra/postprocess/rtx.py — processor/session contract and VFX interop.
2. flashdreams/flashdreams/infra/runner_io.py — host-FFmpeg-only writer behavior after the main refactor.
3. integrations/omnidreams/omnidreams/webrtc/session.py and webrtc/server.py — per-session selection and lifecycle.
4. integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py — local worker postprocess state.
5. integrations/omnidreams/omnidreams/interactive_drive/slangpy_hud_presenter.py — native-resolution presentation and HUD interaction.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR adds an NVIDIA RTX Video Super Resolution post-processor (nvvfx.VideoSuperRes) to the FlashDreams postprocess chain and wires it end-to-end through the interactive-drive HUD, the WebRTC session manager, and the CLI. It also replaces the mediapy video writer with a direct FFmpeg subprocess, resolving previously reported issues around broken-pipe handling, odd-dimension encoding, and missing imageio-ffmpeg fallback.

• rtx.py: New RTXVideoSuperResolutionPostProcessor with lazy VFX effect init, correct [-1,1]↔[0,1] normalization (including uint8 input), explicit torch.cuda.synchronize for the non_blocking path, and three registered entry-point presets (rtx-super-resolution, rtx-super-resolution-ultra, rtx-deblur-ultra).
• runner_io.py: Direct FFmpeg subprocess writer with BrokenPipeError recovery, pad=ceil(iw/2)*2:ceil(ih/2)*2 filter for yuv420p compatibility, and guaranteed process.wait().
• WebRTC / interactive-drive: Per-rollout VideoPostprocessStream lifecycle (create on reset, close on reset/close); browser UI gains a preset selector backed by /api/postprocess/options and /api/session/input; HUD gets a live toggle dispatched through the worker-thread command queue; camera scaling is capped at 1.0 so upscaled frames are not magnified beyond native pixels.

Confidence Score: 5/5

Safe to merge; new code paths are well-isolated behind optional nvvfx and the existing postprocess chain abstraction, and the FFmpeg rewrite addresses all previously flagged output-writing issues.

The RTX post-processor, WebRTC session input override, HUD toggle, and FFmpeg writer are all implemented correctly with proper resource lifecycle management and CUDA synchronization. No functional regressions were found across the changed paths.

No files require special attention beyond the cosmetic HUD label noted above.

Important Files Changed

Filename	Overview
flashdreams/flashdreams/infra/postprocess/rtx.py	New RTX VSR post-processor: well-structured lazy-init effect, correct [−1,1]↔[0,1] conversion, DLPack cloning, and explicit CUDA synchronization for the non-blocking path.
flashdreams/flashdreams/infra/runner_io.py	Replaces mediapy with a direct subprocess FFmpeg pipeline; handles BrokenPipeError, odd-dimension padding, process reaping, and stderr capture. Fixes multiple previously-reported issues.
flashdreams/flashdreams/infra/postprocess/stream.py	Adds self._chunks.clear() after torch.cat() in take_output(), releasing tensor references earlier. Safe defensive improvement.
integrations/omnidreams/omnidreams/webrtc/session.py	Adds per-rollout postprocess stream lifecycle (reset/close) and OmnidreamsSessionInput for browser-selectable preset overrides. Distributed-rank and master-only path logic is correct.
integrations/omnidreams/omnidreams/webrtc/server.py	Adds /api/postprocess/options and /api/session/input endpoints with proper validation, error codes, and SessionBusyError mapping.
integrations/omnidreams/omnidreams/interactive_drive/slangpy_hud_presenter.py	Adds postprocess HUD toggle with auto-resize for upscaled model frames; scale cap prevents blowing up smaller frames. "Upsample 2x" label is hardcoded regardless of active preset.
integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py	Integrates VideoPostprocessStream into FlashdreamsWorldModelSession with correct reset/close lifecycle and worker-thread-safe toggle dispatch via ChunkPipeline command queue.
integrations/omnidreams/omnidreams/webrtc/web/request_session.js	Adds loadPostprocessOptions() and configureSessionInput() with proper error handling and select disabling during connection; fires before RTCPeerConnection creation so race-free.
flashdreams/pyproject.toml	Registers three new entry points under flashdreams.postprocess_presets and adds an optional rtx-postprocess extra for nvidia-vfx==0.1.0.1.
integrations/omnidreams/omnidreams/interactive_drive/video_model/chunk_pipeline.py	Adds set_postprocess_enabled command dispatched through the existing worker-thread command queue, ensuring the toggle is serialized with generation.

Sequence Diagram

sequenceDiagram
    participant Browser
    participant Server as aiohttp Server
    participant Manager as OmnidreamsWebRTCSessionManager
    participant Runtime as OmnidreamsInferenceRuntime
    participant VFX as nvvfx.VideoSuperRes

    Browser->>Server: GET /api/postprocess/options
    Server-->>Browser: "{default_preset, presets[]}"

    Browser->>Server: "POST /api/session/input {postprocess_preset}"
    Server->>Manager: set_pending_session_input(OmnidreamsSessionInput)
    Manager->>Manager: resolve_postprocess_preset(preset) [validate]
    Server-->>Browser: "{postprocess_preset}"

    Browser->>Server: WebRTC connect
    Server->>Manager: _reset_runtime_for_session(session_input)
    Manager->>Runtime: reset_for_new_session(session_input)
    Runtime->>Runtime: _reset_postprocess_stream(session_input)
    Runtime->>Runtime: VideoPostprocessStream created

    loop Per generated chunk
        Runtime->>VFX: effect.run(frame, non_blocking, stream_ptr)
        VFX-->>Runtime: result.image (DLPack)
        Runtime->>Runtime: cuda.synchronize (if non_blocking)
        Runtime->>Runtime: from_dlpack(result.image).clone()
        Runtime-->>Manager: video_chunk (processed, on GPU)
        Manager-->>Browser: WebRTC video frame (CPU)
    end

Loading

_{Reviews (3): Last reviewed commit: "Add optional NVIDIA VFX dependency extra" | Re-trigger Greptile}

[jarcherNV]

LGTM

[gtong-nv]

/ok to test 1278c91