OmniDreams WebRTC: Add support for GPU-accelerated H.264 (NVENC) encoding

[ksheth-dev]

Summary

Move the WebRTC video encode off the CPU and onto the GPU via NVIDIA NVENC
(H.264, PyNvVideoCodec 2.1) for integrations that opt in, behind a small
VideoEncoder abstraction the session picks at startup:

• PyNvHardwareEncoder (flashdreams/serving/webrtc/nvenc.py) —
  hardware H.264 via PyNvVideoCodec. Emits pre-encoded av.Packets that
  aiortc's H264Encoder.pack() forwards straight to the RTP packetizer
  through a new NVENCVideoTrack.
• DefaultRTCEncoder (flashdreams/serving/webrtc/encoders.py) —
  thin wrapper over aiortc's built-in software encoder: raw frames on a
  BufferedVideoTrack, aiortc negotiates the codec via SDP.

select_encoder(backend=...) picks between them:

• backend="auto" (default on omnidreams) — two-stage probe:
  Stage 1 (GetEncoderCaps + resolution bounds) is a silent fallback
  to DefaultRTCEncoder (with an INFO log); Stage 2 (CreateEncoder)
  raising after Stage 1 passed is a hard error — never a silent degrade,
  because it means the environment reports support but session
  allocation still fails.
• backend="nvenc" — any probe/init failure raises EncoderInitError
  at startup so misconfigured deployments surface loudly.
• backend="default" — bypasses the probe entirely; forced software path.

Availability is checked with importlib.util.find_spec — no
PyNvVideoCodec import, no CUDA-driver side effects on processes that
don't opt in. PyNvHardwareEncoder and its top-level
import PyNvVideoCodec live in a dedicated nvenc.py module that
select_encoder lazy-imports only when constructing the hardware
backend; integrations on the software path (Lingbot today) never load
the library.

New CLI flag --prefer_sw_encoder on omnidreams.webrtc.server maps
to backend="default" for A/B profiling or forced-software runs.

Also on this branch:

• BaseWebRTCSessionManager._resolve_video_encoder falls back to a
  session-scope DefaultRTCEncoder when the runtime does not expose
  video_encoder, fixing an AttributeError a real Lingbot WebRTC
  session would otherwise raise.
• Master-rank-only NVENC allocation on distributed launches:
  _initialize_video_encoder_sync early-returns on non-master ranks,
  so worker ranks (rank > 0) do not burn NVENC session slots they
  never use.
• H.264 SDP-fallback ownership fix: _enforce_h264_or_fallback awaits
  close on the orphaned session-scope track before installing the
  software fallback, and does not close the runtime-owned encoder
  (which is reused across sessions).

Measured impact with NVENC

At 1280x704, 30 fps, 10 Mb/s target on a single-GPU host:

• Actual encode wall-clock: ~5.7x reduction (40 ms -> 7 ms per chunk of 8 frames)
• Producer-side enqueue_ms: ~4x reduction (28 ms -> 7 ms)

Test plan

CI-safe unit tests (CPU)

• uv run pre-commit run -a — ruff format/fix, ty type check, uv-lockfile validation.
• uv run pytest -m ci_cpu — Full CPU-safe suite.
• uv run pytest flashdreams/tests/test_encoders.py -v — select_encoder branch coverage: backend="default" short-circuit, Stage-1 (find_spec / GetEncoderCaps) silent fallback vs. backend="nvenc" hard-raise, Stage-2 CreateEncoder hard-error guard, cross-chunk packet-ordering invariant (240-packet regression), aiortc/PyNvVideoCodec compatibility guards.
• uv run pytest flashdreams/tests/test_webrtc_manager.py -v — Base manager: _resolve_video_encoder runtime-encoder pass-through and software fallback (fix for the missing-video_encoder AttributeError), H.264 codec preference, SDP-fallback session-scope cleanup.
• uv run pytest integrations/omnidreams/tests/test_webrtc_runtime.py -v — Omnidreams session init: select_encoder wiring, --prefer_sw_encoder → encoder_backend mapping, master-rank-only NVENC allocation guard, session-scope orphan-track close on H.264 SDP fallback (runtime-owned encoder is not closed).
• uv run pytest integrations/lingbot/tests/test_webrtc_runtime.py -v — Lingbot session construction still passes with the base-manager software fallback (no video_encoder on the runtime).

Structural / no-side-effect invariants

• Confirm encoders.py does not import PyNvVideoCodec or nvenc at module load:

uv run python -c "
import sys
from flashdreams.serving.webrtc import encoders
assert 'PyNvVideoCodec' not in sys.modules
assert 'flashdreams.serving.webrtc.nvenc' not in sys.modules
print('OK')"

• Confirm select_encoder(backend='default') does not touch NVENC even when the library is present:

uv run python -c "
import sys
from flashdreams.serving.webrtc.encoders import select_encoder
select_encoder(backend='default', width=1280, height=704, fps=30,
               bitrate=6_000_000, gpu_id=0)
assert 'PyNvVideoCodec' not in sys.modules
assert 'flashdreams.serving.webrtc.nvenc' not in sys.modules
print('OK')"

• Confirm the availability probe is side-effect-free:

uv run python -c "
import sys
from flashdreams.serving.webrtc.encoders import _pynvvideocodec_installed
present = _pynvvideocodec_installed()
assert 'PyNvVideoCodec' not in sys.modules
print(f'probe={present}, no import triggered')"

GPU tests (needs CUDA + NVENC-capable GPU)

• uv run pytest -m ci_gpu — Full GPU-safe suite.
• uv run pytest integrations/omnidreams/tests/test_nvenc_smoke.py -v — Hardware NVENC smoke: PyNvHardwareEncoder.is_supported probe, encoder init at the target resolution, emitted packets start with Annex-B start code and every IDR is preceded by SPS+PPS.

End-to-end (needs GPU + browser)

• Auto backend on an NVENC-capable host → NVENC selected:

  uv run --package flashdreams-omnidreams torchrun --nproc_per_node 1 \
    -m omnidreams.webrtc.server \
    --pipeline_config_name omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf \
    --scene-uuid 0d404ff7-2b66-498c-b047-1ed8cded60d4 \
    --port 8089

  • Expected: startup log Video encoder ready: backend=pynvvideocodec ..., browser connects and receives H.264 video.
• Force software encoder:

  uv run --package flashdreams-omnidreams torchrun --nproc_per_node 1 \
    -m omnidreams.webrtc.server \
    --pipeline_config_name omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf \
    --scene-uuid 0d404ff7-2b66-498c-b047-1ed8cded60d4 \
    --port 8089 \
    --prefer_sw_encoder

  • Expected: log Video encoder ready: backend=aiortc ..., behavior matches pre-branch software path.
• Multi-rank launch — only master rank allocates NVENC:

  uv run --package flashdreams-omnidreams torchrun --nproc_per_node 2 \
    -m omnidreams.webrtc.server \
    --pipeline_config_name omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf \
    --scene-uuid 0d404ff7-2b66-498c-b047-1ed8cded60d4 \
    --port 8089

  • Expected: only rank 0 emits Video encoder ready:; rank 1 does not construct a PyNvHardwareEncoder and does not consume an NVENC session slot.
• Lingbot software-path launch reaches the WebRTC listener without AttributeError and without silent exit-0 when PyNvVideoCodec is installed on the host:

  uv run --package flashdreams-lingbot python -m lingbot.webrtc.server \
    --host 0.0.0.0 --port 8089 \
    --config_name lingbot-world-fast-taehv-window15-sink3

  • Expected: server reaches Listening on ..., and inside the process PyNvVideoCodec never appears in sys.modules.

[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 moves WebRTC video encoding off the CPU and onto NVIDIA NVENC (H.264, PyNvVideoCodec 2.1) behind a VideoEncoder abstraction, delivering a measured ~5.7× encode wall-clock reduction at 1280×704/30fps. The select_encoder factory uses a two-stage probe — Stage 1 (capability/import) falls back silently to aiortc's software encoder; Stage 2 (CreateEncoder failure) is a hard error — with a lazy import that ensures processes on the software path never load PyNvVideoCodec.

• PyNvHardwareEncoder / NVENCVideoTrack (nvenc.py, media.py): hardware H.264 via PyNvVideoCodec; packets emitted per-frame via call_soon_threadsafe, drop-oldest overflow, and pacing logic mirroring BufferedVideoTrack.
• DefaultRTCEncoder (encoders.py): thin software-path wrapper; _resolve_video_encoder on the base manager falls back to this when the runtime doesn't expose video_encoder, fixing the Lingbot AttributeError.
• Session lifecycle changes (manager.py, session.py): master-rank-only NVENC allocation, H.264 SDP preference with _enforce_h264_or_fallback (closes the orphaned NVENCVideoTrack before installing the fallback), and an explicit CUDA stream sync replacing the implicit .cpu() barrier so chunks stay on-device for the hardware path.

Confidence Score: 5/5

Safe to merge; the encoder abstraction, two-stage probe, SDP fallback cleanup, and master-rank guard are all correctly implemented. The one latent defect (PTS counter skipping for empty NVENC output) cannot trigger in the configured zero-latency mode.

The encoder abstraction, two-stage probe, SDP fallback cleanup, and master-rank guard are all correctly implemented. The previously flagged import-isolation hole is now covered by try/except (ImportError, RuntimeError) around the deferred import. The PyNvVideoCodec version constraint is pinned to >=2.1,<3. The _enforce_h264_or_fallback path explicitly closes the orphaned NVENCVideoTrack before overwriting the reference. The single latent defect cannot fire under the configured bf=0/lookahead=0 zero-latency settings.

Files Needing Attention: flashdreams/flashdreams/serving/webrtc/nvenc.py — the encode_chunk_sync PTS counter logic.

Important Files Changed

Filename	Overview
flashdreams/flashdreams/serving/webrtc/encoders.py	New file: VideoEncoder Protocol + DefaultRTCEncoder + select_encoder factory; two-stage probe with correct ImportError/RuntimeError handling; lazy-import isolation for PyNvVideoCodec confirmed.
flashdreams/flashdreams/serving/webrtc/nvenc.py	New file: PyNvHardwareEncoder with NVENC H.264 via PyNvVideoCodec 2.1; PTS counter only advances on non-empty bitstream output — if NVENC returns empty bs despite bf=0/lookahead=0, subsequent packets receive stale timestamps.
flashdreams/flashdreams/serving/webrtc/media.py	Adds NVENCVideoTrack: pre-encoded Packet delivery track with drop-oldest overflow and correct pacing/stall re-anchoring; close() drains queue and signals None sentinel cleanly.
flashdreams/flashdreams/serving/webrtc/manager.py	Adds VideoEncoder plumbing to session lifecycle: _resolve_video_encoder fallback, _prefer_h264_video_codec codec preferences, _enforce_h264_or_fallback (closes orphaned NVENCVideoTrack before overwriting reference, intentionally skips runtime-owned encoder); generation worker now routes through deliver_chunk.
integrations/omnidreams/omnidreams/webrtc/session.py	Adds _initialize_video_encoder_sync (master-rank-only, select_encoder wiring) and replaces video_chunk.detach().cpu() with an explicit current_stream().synchronize() so the tensor stays on-device for the hardware path while preserving the GPU sync barrier.
integrations/omnidreams/pyproject.toml	Adds PyNvVideoCodec>=2.1,<3 dependency with correct version bound covering the 2.1 API surface (GetEncoderCaps gpuid=, FORCEIDR, CreateEncoder keyword signature).
integrations/omnidreams/omnidreams/webrtc/server.py	Adds --prefer_sw_encoder CLI flag that maps to encoder_backend='default'; wiring is clean and consistent with the OmnidreamsRuntimeConfig field.
flashdreams/tests/test_encoders.py	Comprehensive new test file covering Protocol conformance, all select_encoder branches, Stage-1 silent fallback vs Stage-2 hard-raise, cross-chunk packet-ordering invariant, and import-isolation invariants.
integrations/omnidreams/tests/test_webrtc_runtime.py	New integration tests: select_encoder wiring, --prefer_sw_encoder mapping, master-rank-only guard, and SDP fallback session-scope cleanup.
integrations/omnidreams/tests/test_nvenc_smoke.py	GPU smoke test for PyNvHardwareEncoder: verifies Annex-B start code, IDR/SPS/PPS presence, and PTS/time_base values; guards collection behind try/except to avoid CI collection failures on CPU hosts.

_{Reviews (7): Last reviewed commit: "webrtc: isolate PyNvVideoCodec in a dedi..." | Re-trigger Greptile}

[jarcherNV]

Thanks for putting this together. I think this is the right direction, but I would not merge it yet.

A few things need to be addressed first:

• The branch is currently conflicted with main, so it needs to be updated/rebased.
• Please resolve the software-encoder path issue Greptile flagged. generate_chunk() still routes through _encode_raw_output_sync(), which calls self._video_encoder.encode_chunk(), but DefaultRTCVideoEncoder does not implement encode_chunk(). That can fail at runtime when the software encoder path is selected.
• Please avoid relying on aiortc’s private _codecs attribute for codec negotiation. We should use a public API or parse the negotiated/local SDP instead.
• Since this depends on PyNvVideoCodec 2.1 APIs, please pin the dependency accordingly.
• The PR body still has CPU-only/non-NVIDIA fallback testing unchecked. Please either run that or add a focused test/mocking path for the fallback behavior.

After those are addressed and the branch has current CI results, I think this is worth a full review.

[jatentaki]

Agreeing with @jarcherNV request for rebase. I asked an agent for a rebase and I'm seeing issues (the stream freezes after a couple of frames, reporting increasing chunk id but remaining visually frozen). I would delay in-depth review until we have a concrete rebased version.

[gtong-nv]

There were some changes merged that unifies Lingbot and omnidreams WebRTC serving, which caused the conflict.
But the NVENC encoding should work any models uses WebRTC. @ksheth-dev can you address the comments above and also make it work for Lingbot? Thanks!

[gtong-nv]

It seems there are some ordering issues in the chunk generation and encoder

1. Generation pipeline produces raw chunks in order: chunk 0, chunk 1, chunk 2...
2. Encoding pipeline is spawned as fire-and-forget tasks, so multiple chunks can be encoding at the same time.

The generation order is sequential, but the packet enqueue order is whichever encode worker finishes/schedules first.
Codex suggests using one ordered encode consumer..

[jarcherNV]

Thanks for updating this. It looks like it is making good progress.

I still see a couple of issues that I think should be addressed before merge:

1. The shared BaseWebRTCSessionManager now unconditionally reads runtime.video_encoder, but the real LingbotInferenceRuntime does not appear to define that property. The Lingbot tests add fake encoders when constructing managed sessions, but I think the actual Lingbot WebRTC path would fail when creating an answer (I have not confirmed this myself by running it though). Perhaps the base manager should default to a software encoder when the runtime does not provide one. We can verify this on our end and confirm how Lingbot behaves.

2. In _enforce_h264_or_fallback(), the fallback path closes/replaces the session-local encoder, but the runtime still owns the original hardware encoder object. If that encoder is closed during fallback, a later session may still get the closed runtime encoder from runtime.video_encoder. I think the fallback should either update the runtime-owned encoder state or avoid closing/swapping a runtime-owned encoder at session scope.

3. Greptile’s remaining comment may be valid: the old NVENCVideoTrack should be closed when fallback replaces it with the software track, so stop() is called and its queue is drained.

4. I think the current testing may not be sufficient to fully verify this with CI. Maybe extra tests would be too expensive, I'm not sure. I'll leave this up to you.

[ksheth-dev]

The feature changes have been refactored in the latest version of dev/ksheth/pynvvideocodec_v2 branch which is rebased on latest from main branch.

With the latest refactored changes, the following above comments are addressed:

• Rebase done on main branch.

• SW encoder fallback path is added and verified.

• Removed usage of aiortc’s private _codecs attribute for codec negotiation.

• Pinned PyNvVideoCodec dependency to >=2.1.

• Verified the stream looks visibly fine.

• Ensured that the generation and encoding order are sequential.

• As per offline discussion, scoping this PR to only Omnidreams + WebRTC interactive drive. Extending to other integrations and/or moving NVENC support from serving/webrtc to a common core functionality will be part of a follow-up PR.

• Pending items:

• Verify Lingbot integration and ensure no impact to it.

• Measure NVENC vs SW encoder perf numbers and publish in this PR.

[gtong-nv]

Tested the nvenc code path on RTX6000 pro, and it did improve the omnidreams interactive demo demo in webRTC - from 12 FPS to 14 FPS

testing it with and without prefer_sw_encoder

uv run  --package flashdreams-omnidreams   python -m omnidreams.webrtc.server   --pipeline_config_name omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf   --scene-uuid 0d404ff7-2b66-498c-b047-1ed8cded60d4   --port 8091   --warmup_chunks 20   --warmup_timeout_s 900   --device cuda:0 

I think we can get this merged once @jarcherNV 's comments get addressed.

The PR is general enough and I think other demos that use webRTC can easily use this nvenc implementation.

[gtong-nv]

Another issue when testing in multi-gpu env -

The hardware encoder is initialized on every distributed rank even though only rank 0 serves WebRTC media. In torchrun --nproc_per_node=4, _initialize_sync_all_ranks() calls _initialize_video_encoder_sync() on all ranks, consuming NVENC sessions on workers and making startup fail if any non-serving rank cannot allocate NVENC. This should be master-only or lazy in create_answer()

[gtong-nv]

/ok to test 8e63d69

[ksheth-dev]

/ok to test 0ba2d2b

[ksheth-dev]

/ok to test f921e2c

[ksheth-dev]

/ok to test bec5ea2

[ksheth-dev]

Updated this PR to address all the comments and issues raised by the reviewers.

• Rebase artifacts folded into commit 2 (c00d773): session.py syntax after the postprocess field landed; Namespace fixture missing postprocess_preset.
• Runtime without video_encoder (Lingbot AttributeError) — commit 5 (63ff34c): base manager gains _resolve_video_encoder, falls back to a session-scope DefaultRTCEncoder when the runtime doesn't own one.
• Hardware encoder allocated on every distributed rank — commit 6 (b0878e3): _initialize_video_encoder_sync early-returns on non-master ranks, so worker ranks don't burn NVENC session slots.
• Session-scope cleanup on H.264 SDP fallback — commit 4 (defd093): _enforce_h264_or_fallback awaits close on the orphaned session-scope track but does not close the runtime-owned encoder.
• Lingbot silent exit-0 when PyNvVideoCodec is importable — commit 7 (6434238): PyNvHardwareEncoder split into a dedicated nvenc.py; select_encoder probes via importlib.util.find_spec and lazy-imports nvenc only on the hardware branch, so software-path processes never load PyNvVideoCodec.
• Pre-commit ty failures after the module split — folded into commit 7: narrowed VideoEncoder.create_track return to BufferedVideoTrack | NVENCVideoTrack, and put import PyNvVideoCodec as nvc behind a TYPE_CHECKING split so ty sees nvc: Any.
• CPU CI collection abort on test_nvenc_smoke.py — folded into commit 7: try around the nvenc imports with pytest.skip(allow_module_level=True) catching RuntimeError from PyNvVideoCodec's driver-load probe.
• CPU CI test failure via pytest.importorskip("PyNvVideoCodec") — folded into commit 7: replaced with try around the nvenc import catching both ImportError and RuntimeError, so the compat guard skips cleanly when the NVIDIA driver library is absent.

[ksheth-dev]

/ok to test 6434238

[jarcherNV]

I have not tested this myself, but it looks like my previous comments have been addressed. I am in favor of merging this, but will wait for @gtong-nv to confirm this works from his own testing.