diff --git a/flashdreams/flashdreams/serving/webrtc/encoders.py b/flashdreams/flashdreams/serving/webrtc/encoders.py new file mode 100644 index 000000000..b1bf0cae2 --- /dev/null +++ b/flashdreams/flashdreams/serving/webrtc/encoders.py @@ -0,0 +1,562 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Video encoder backends for the WebRTC serving path. + +Two implementations share a single ``VideoEncoder`` protocol: + +- :class:`DefaultRTCEncoder` — thin adapter over aiortc's built-in software + encoder path. Frames leave the model as raw RGB and aiortc's RTP sender + loop drives encoding. +- :class:`PyNvHardwareEncoder` — GPU-resident NVENC H.264 via + ``PyNvVideoCodec`` 2.1. Emits pre-encoded Annex-B packets that aiortc's + ``H264Encoder.pack()`` fragments into RTP payloads (no re-encoding). + +Selection is performed by :func:`select_encoder`, which layers a +capability query (``GetEncoderCaps``) with two failure semantics: an +environmental "no" from Stage 1 falls back silently to the software path, +while a Stage-2 ``CreateEncoder`` failure is logged with traceback and +re-raised so it cannot be masked by a silent fallback. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import time +from collections.abc import Callable +from dataclasses import dataclass +from fractions import Fraction +from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable + +import torch +from aiortc import MediaStreamTrack +from av.packet import Packet +from loguru import logger + +if TYPE_CHECKING: + from flashdreams.serving.webrtc.media import ( + BufferedVideoTrack, + NVENCVideoTrack, + ) + +# PyNvVideoCodec is a runtime dep. Members like ``nvc.GetEncoderCaps`` / +# ``nvc.FORCEIDR`` are not declared in any type stub, and the module may be +# absent at import time on non-NVENC hosts (or intentionally removed for +# testing the auto-fallback path). Hiding the import behind ``TYPE_CHECKING`` +# gives ty a single ``Any``-typed name to reason about while the runtime +# path keeps the guarded try / except. +if TYPE_CHECKING: + nvc: Any = None + _PYNVVIDEOCODEC_AVAILABLE: bool = False +else: + try: + import PyNvVideoCodec as nvc + + _PYNVVIDEOCODEC_AVAILABLE = True + except ImportError: + nvc = None + _PYNVVIDEOCODEC_AVAILABLE = False + + +EncoderBackend = Literal["auto", "nvenc", "default"] + + +# H.264 NAL type identifiers used when inspecting Annex-B bitstreams. +_H264_NAL_TYPE_IDR = 5 +_H264_NAL_TYPE_SPS = 7 +_H264_NAL_TYPE_PPS = 8 + +# RTP video clock, per RFC 6184. aiortc's ``H264Encoder.pack()`` rescales +# from whatever ``time_base`` the packet carries into this base, so setting +# ``time_base = 1 / _RTP_VIDEO_CLOCK`` on emitted packets is the +# lowest-conversion choice. +_RTP_VIDEO_CLOCK = 90_000 + + +class EncoderInitError(RuntimeError): + """Raised when a forced encoder backend cannot be initialized.""" + + +@dataclass(slots=True, frozen=True) +class ChunkDeliveryResult: + """Uniform result shape for :meth:`VideoEncoder.deliver_chunk`. + + Callers do not need to know which encoder produced the chunk; the + fields here cover both paths. + """ + + backend: str + num_frames: int + num_keyframes: int + encode_ms: float + + +@runtime_checkable +class VideoEncoder(Protocol): + """Encoder backend paired with a compatible :class:`MediaStreamTrack`. + + Each backend owns two responsibilities: creating a fresh media track + sized for one session (:meth:`create_track`), and encoding + enqueueing + one chunk of frames onto that track (:meth:`deliver_chunk`). Callers + pick one backend at startup and branch nowhere else. + """ + + fps: int + backend: str + prefers_codec: str | None + """SDP codec preference hint. ``"h264"`` for the NVENC path (must be + honoured or the pre-encoded bitstream will not decode); ``None`` to + let aiortc pick freely.""" + + def create_track(self, *, maxsize: int) -> BufferedVideoTrack | NVENCVideoTrack: + """Create a fresh session track compatible with this encoder.""" + ... + + async def deliver_chunk( + self, + chunk: torch.Tensor, + track: MediaStreamTrack, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + """Encode ``chunk`` and enqueue the resulting frames/packets.""" + ... + + def close(self) -> None: + """Release any encoder-owned resources (hardware handles, etc.).""" + ... + + +# --------------------------------------------------------------------------- +# Software path (aiortc's built-in encoder) +# --------------------------------------------------------------------------- + + +class DefaultRTCEncoder: + """Software path: hand raw RGB frames to aiortc and let it encode. + + This class does not encode itself; it wraps a + :class:`~flashdreams.serving.webrtc.media.BufferedVideoTrack` that + carries raw RGB frames, and aiortc's RTP sender loop drives encoding + frame-by-frame. The concrete codec (VP8, VP9, H.264) is chosen by SDP + negotiation. + """ + + prefers_codec: str | None = None + backend = "aiortc" + + def __init__(self, *, fps: int) -> None: + if fps <= 0: + raise ValueError(f"fps must be > 0, got {fps}") + self.fps = fps + logger.info( + "Video encoder ready: backend={} fps={}", + self.backend, + fps, + ) + + def create_track(self, *, maxsize: int) -> BufferedVideoTrack: + from flashdreams.serving.webrtc.media import BufferedVideoTrack + + return BufferedVideoTrack(fps=self.fps, maxsize=maxsize) + + async def deliver_chunk( + self, + chunk: torch.Tensor, + track: MediaStreamTrack, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + # aiortc's software encoder chooses keyframe cadence internally + # and responds to receiver PLI/FIR feedback, so a caller-supplied + # force_keyframe is neither honoured nor useful on this path. + del force_keyframe + from flashdreams.serving.webrtc.media import BufferedVideoTrack + + if not isinstance(track, BufferedVideoTrack): + raise TypeError( + "DefaultRTCEncoder requires a BufferedVideoTrack; got " + f"{type(track).__name__}. Create it via encoder.create_track()." + ) + enqueued = await track.enqueue_chunk(chunk) + return ChunkDeliveryResult( + backend=self.backend, + num_frames=enqueued, + num_keyframes=0, + encode_ms=0.0, + ) + + def close(self) -> None: + # No encoder-owned resources; the track is owned by aiortc's PC. + return + + +# --------------------------------------------------------------------------- +# Hardware path (NVENC via PyNvVideoCodec 2.1) +# --------------------------------------------------------------------------- + + +def _payload_contains_nal_type(payload: bytes, nal_type: int) -> bool: + """Scan an Annex-B H.264 payload for the presence of a specific NAL type.""" + i = 0 + while True: + idx = payload.find(b"\x00\x00\x01", i) + if idx < 0: + return False + nal_start = idx + 3 + if nal_start >= len(payload): + return False + if (payload[nal_start] & 0x1F) == nal_type: + return True + i = nal_start + 1 + + +def _chunk_to_abgr_cuda_frames(chunk: torch.Tensor) -> torch.Tensor: + """Convert a model-output chunk to NVENC-``ABGR``-formatted CUDA frames. + + Accepts ``[T, 3, H, W]`` or ``[1, 1, T, 3, H, W]`` (the shape produced + by the omnidreams runtime) in either ``uint8`` or float dtype + (float assumed to be in ``[-1, 1]``). Returns a contiguous + ``[T, H, W, 4]`` ``uint8`` CUDA tensor with alpha=255. + + **NVENC ``NV_ENC_BUFFER_FORMAT_ABGR`` is a word-ordered token, not + memory-ordered.** From ``nvEncodeAPI.h``: "a pixel is represented by + a 32-bit word with R in the lowest 8 bits, G in the next 8 bits, B + in the 8 bits after that and A in the highest 8 bits" (word + ``0xAABBGGRR``). In little-endian memory that is the byte sequence + ``[R, G, B, A]`` — so the channel-last tensor we hand to the encoder + must have channel 0 = R, 1 = G, 2 = B, 3 = A. Writing ``[A, B, G, R]`` + (the naive memory-order reading of the name) makes NVENC interpret + the alpha byte as R, producing a visible RGB↔BGR swap on the wire. + + ABGR (rather than NV12) is chosen so NVENC's driver-side RGB→YUV + conversion handles the colour transform, sparing us a bespoke NV12 + kernel. + """ + if not chunk.is_cuda: + raise ValueError("expected CUDA tensor for hardware encode path") + if chunk.ndim == 6: + if chunk.shape[0] != 1 or chunk.shape[1] != 1: + raise ValueError( + "expected single-batch, single-view chunk [1, 1, T, 3, H, W]; " + f"got {tuple(chunk.shape)}" + ) + chunk = chunk[0, 0] + if chunk.ndim != 4 or chunk.shape[1] != 3: + raise ValueError( + "expected chunk shape [T, 3, H, W] or [1, 1, T, 3, H, W]; " + f"got {tuple(chunk.shape)}" + ) + if chunk.dtype == torch.uint8: + rgb = chunk.permute(0, 2, 3, 1) + else: + rgb = ((chunk.float() + 1.0) / 2.0 * 255.0).clamp(0, 255).byte() + rgb = rgb.permute(0, 2, 3, 1) + t, h, w, _ = rgb.shape + a = torch.full((t, h, w, 1), 255, dtype=torch.uint8, device=rgb.device) + # Channel-last [R, G, B, A] → little-endian bytes [R, G, B, A] → + # NVENC word 0xAABBGGRR = NV_ENC_BUFFER_FORMAT_ABGR. + rgba = torch.cat([rgb, a], dim=-1) + return rgba.contiguous() + + +class PyNvHardwareEncoder: + """NVENC H.264 encoder backed by ``PyNvVideoCodec`` 2.1. + + Accepts CUDA tensors and emits Annex-B H.264 packets streamed onto an + :class:`~flashdreams.serving.webrtc.media.NVENCVideoTrack` as they are + encoded. Packets carry ``pts`` on the RTP 90 kHz video clock so + aiortc's ``H264Encoder.pack()`` can rescale without loss. + """ + + prefers_codec: str | None = "h264" + backend = "pynvvideocodec" + + @classmethod + def is_supported( + cls, + *, + gpu_id: int = 0, + width: int = 0, + height: int = 0, + ) -> tuple[bool, str]: + """Query the driver for NVENC H.264 support at the target resolution. + + Uses ``PyNvVideoCodec.GetEncoderCaps`` (public API) so that no + NVENC session is allocated for the probe itself. Returns + ``(True, "")`` when the environment supports NVENC H.264 at the + requested resolution, else ``(False, human_reason)`` where + ``human_reason`` is a diagnostic suitable for logging. + """ + if not _PYNVVIDEOCODEC_AVAILABLE: + return False, "PyNvVideoCodec library is not installed" + try: + # PyNvVideoCodec 2.1 signature: GetEncoderCaps(gpuid=, codec=). + # Keyword is ``gpuid`` (no underscore); positional order is + # (gpuid, codec). + caps = nvc.GetEncoderCaps(gpuid=gpu_id, codec="h264") + except Exception as exc: + return False, ( + f"GetEncoderCaps gpuid={gpu_id} raised {type(exc).__name__}: {exc}" + ) + if not caps: + return False, (f"GetEncoderCaps gpuid={gpu_id} returned no capabilities") + max_w = int(caps.get("width_max", 0) or 0) + max_h = int(caps.get("height_max", 0) or 0) + min_w = int(caps.get("width_min", 0) or 0) + min_h = int(caps.get("height_min", 0) or 0) + if width > 0 and height > 0: + if max_w > 0 and max_h > 0 and (width > max_w or height > max_h): + return False, ( + f"requested resolution {width}x{height} exceeds driver " + f"maximum {max_w}x{max_h} (gpuid={gpu_id})" + ) + if (min_w > 0 or min_h > 0) and (width < min_w or height < min_h): + return False, ( + f"requested resolution {width}x{height} below driver " + f"minimum {min_w}x{min_h} (gpuid={gpu_id})" + ) + return True, "" + + def __init__( + self, + *, + width: int, + height: int, + fps: int, + bitrate: int, + gpu_id: int, + gop: int = 30, + ) -> None: + if not _PYNVVIDEOCODEC_AVAILABLE: + raise RuntimeError( + "PyNvVideoCodec is not installed; PyNvHardwareEncoder cannot " + "be used. Install with `pip install PyNvVideoCodec` or " + "select DefaultRTCEncoder." + ) + if width <= 0 or height <= 0: + raise ValueError(f"width and height must be > 0, got {width}x{height}") + if fps <= 0: + raise ValueError(f"fps must be > 0, got {fps}") + if bitrate <= 0: + raise ValueError(f"bitrate must be > 0, got {bitrate}") + if gop <= 0: + raise ValueError(f"gop must be > 0, got {gop}") + + # Fail fast with a driver-specific diagnostic before allocating a + # session slot. When called via ``select_encoder`` this is + # redundant with the factory's own probe, but a direct instantiation + # (e.g. from a test) still gets a clear reason. + supported, reason = self.is_supported( + gpu_id=gpu_id, + width=width, + height=height, + ) + if not supported: + raise RuntimeError(f"NVENC H.264 not supported: {reason}") + + self.fps = fps + self._width = width + self._height = height + self._bitrate = bitrate + self._gpu_id = gpu_id + self._gop = gop + self._pts_counter = 0 + self._time_base = Fraction(1, _RTP_VIDEO_CLOCK) + # ``FORCEIDR`` is exposed by PyNvVideoCodec as an integer flag + # bitmask. Cache it so we don't hit ``getattr`` on every frame. + self._force_idr_flag = int(nvc.FORCEIDR) # type: ignore[attr-defined] + + # ``repeatspspps=1`` prepends SPS+PPS to every IDR: aiortc's + # ``H264Encoder.pack()`` does not synthesize parameter sets, so the + # RTP stream must carry them in-band or the receiver cannot lock on. + # ``bf=0`` and ``lookahead=0`` keep the output strictly 1:1 with + # input frames, which is what interactive streaming needs. + self._encoder = nvc.CreateEncoder( # type: ignore[attr-defined] + width=width, + height=height, + fmt="ABGR", + usecpuinputbuffer=False, + codec="h264", + preset="P4", + tuning_info="ultra_low_latency", + rc="cbr", + fps=fps, + bitrate=bitrate, + bf=0, + lookahead=0, + repeatspspps=1, + idrperiod=gop, + ) + logger.info( + "Video encoder ready: backend={} codec=h264 {}x{}@{}fps " + "bitrate={}bps gop={} gpu={}", + self.backend, + width, + height, + fps, + bitrate, + gop, + gpu_id, + ) + + def create_track(self, *, maxsize: int) -> NVENCVideoTrack: + # Lazy import to avoid an ``encoders`` ↔ ``media`` cycle. + from flashdreams.serving.webrtc.media import NVENCVideoTrack + + return NVENCVideoTrack(fps=self.fps, maxsize=maxsize) + + async def deliver_chunk( + self, + chunk: torch.Tensor, + track: MediaStreamTrack, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + from flashdreams.serving.webrtc.media import NVENCVideoTrack + + if not isinstance(track, NVENCVideoTrack): + raise TypeError( + "PyNvHardwareEncoder requires an NVENCVideoTrack; got " + f"{type(track).__name__}. Create it via encoder.create_track()." + ) + packets: list[Packet] = [] + + _num_frames, num_keyframes, encode_ms = await asyncio.to_thread( + self.encode_chunk_sync, + chunk, + force_keyframe=force_keyframe, + on_packet=packets.append, + ) + enqueued = await track.enqueue_encoded_packets(packets) + if enqueued < len(packets): + logger.debug( + "NVENC track closed while enqueueing encoded chunk; " + "enqueued {} of {} packet(s).", + enqueued, + len(packets), + ) + return ChunkDeliveryResult( + backend=self.backend, + num_frames=enqueued, + num_keyframes=num_keyframes, + encode_ms=encode_ms, + ) + + def encode_chunk_sync( + self, + chunk: torch.Tensor, + *, + force_keyframe: bool = False, + on_packet: Callable[[Packet], None] | None = None, + ) -> tuple[int, int, float]: + """Encode a chunk synchronously; returns ``(num_frames, num_keyframes, encode_ms)``. + + Kept public because callers (e.g. tests) that already run on a + worker thread should not have to route through :meth:`deliver_chunk` + just to get access to the emitted packets. + """ + frames = _chunk_to_abgr_cuda_frames(chunk) + num_frames = frames.shape[0] + num_keyframes = 0 + start_s = time.perf_counter() + for i in range(num_frames): + frame = frames[i].contiguous() + if force_keyframe and i == 0: + bs = self._encoder.Encode(frame, self._force_idr_flag) + else: + bs = self._encoder.Encode(frame) + if not bs: + continue + payload = bytes(bs) + packet = Packet(payload) + packet.pts = (self._pts_counter * _RTP_VIDEO_CLOCK) // self.fps + packet.time_base = self._time_base + self._pts_counter += 1 + if _payload_contains_nal_type(payload, _H264_NAL_TYPE_IDR): + num_keyframes += 1 + if on_packet is not None: + on_packet(packet) + encode_ms = (time.perf_counter() - start_s) * 1000.0 + return num_frames, num_keyframes, encode_ms + + def close(self) -> None: + with contextlib.suppress(Exception): + self._encoder.EndEncode() + + +# --------------------------------------------------------------------------- +# Factory +# --------------------------------------------------------------------------- + + +def select_encoder( + *, + backend: EncoderBackend, + width: int, + height: int, + fps: int, + bitrate: int, + gpu_id: int, + gop: int = 30, +) -> VideoEncoder: + """Choose an encoder implementation. + + ``backend == "default"`` returns a :class:`DefaultRTCEncoder` and does + not touch the NVENC probe path at all. + + ``backend == "nvenc"`` requires the hardware path; any probe failure + raises :class:`EncoderInitError` so misconfigured deployments surface + at startup rather than silently degrading. + + ``backend == "auto"`` prefers hardware when the environment supports + it. Selection uses two stages with two different failure semantics: + + * **Stage 1** (``GetEncoderCaps`` + resolution bounds): a "no" here + means the environment cannot do this, which is expected on hosts + without NVENC. Fall back silently to + :class:`DefaultRTCEncoder`, logging the reason at INFO. + * **Stage 2** (``CreateEncoder``): a raise here means Stage 1 said + the environment supports NVENC but session allocation failed anyway + — driver bug, session-pool exhaustion, hardware fault, or + misconfiguration. Log with traceback and re-raise; do not silently + degrade, because doing so would hide a real problem. + """ + if backend == "default": + return DefaultRTCEncoder(fps=fps) + + supported, reason = PyNvHardwareEncoder.is_supported( + gpu_id=gpu_id, + width=width, + height=height, + ) + if not supported: + if backend == "nvenc": + raise EncoderInitError( + f"encoder_backend='nvenc' requested but NVENC H.264 is not " + f"supported on this system: {reason}" + ) + logger.info( + "NVENC H.264 not supported on this system ({}); using aiortc " + "software encoder.", + reason, + ) + return DefaultRTCEncoder(fps=fps) + + try: + return PyNvHardwareEncoder( + width=width, + height=height, + fps=fps, + bitrate=bitrate, + gpu_id=gpu_id, + gop=gop, + ) + except Exception: + logger.exception( + "GetEncoderCaps reported NVENC H.264 supported, but CreateEncoder " + "failed. Not silently falling back to the software encoder — the " + "underlying failure needs to be diagnosed.", + ) + raise diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index 1a17c8b6b..b56d845e5 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -9,6 +9,7 @@ import contextlib import inspect import json +import math from collections import deque from collections.abc import Set as AbstractSet from dataclasses import dataclass, field @@ -16,11 +17,21 @@ from typing import Any import torch -from aiortc import RTCConfiguration, RTCPeerConnection, RTCSessionDescription +from aiortc import ( + RTCConfiguration, + RTCPeerConnection, + RTCRtpSender, + RTCSessionDescription, +) +from aiortc.sdp import SessionDescription from loguru import logger from flashdreams.serving.realtime.input import KeyboardResampler -from flashdreams.serving.webrtc.media import BufferedVideoTrack +from flashdreams.serving.webrtc.encoders import ( + DefaultRTCEncoder, + VideoEncoder, +) +from flashdreams.serving.webrtc.media import BufferedVideoTrack, NVENCVideoTrack from flashdreams.serving.webrtc.server import SessionBusyError from flashdreams.serving.webrtc.warmup import ( run_loopback_warmup_session, @@ -35,6 +46,45 @@ _CLIENT_LIVENESS_CHECK_INTERVAL_S = 1.0 +def _performance_stats_payload(stats: dict[str, float] | None) -> dict[str, float]: + if not stats: + return {} + payload: dict[str, float] = {} + for key, value in stats.items(): + if not isinstance(key, str): + continue + try: + number = float(value) + except (TypeError, ValueError): + continue + if math.isfinite(number): + payload[key] = round(number, 1) + return payload + + +def _format_performance_stats(stats: dict[str, float]) -> str: + return " ".join(f"{key}={value:.1f}" for key, value in sorted(stats.items())) + + +def _sdp_video_codecs(sdp: str) -> tuple[str, ...] | None: + """Return offered video codec MIME types, or ``None`` if parsing fails.""" + try: + description = SessionDescription.parse(sdp) + except Exception: + logger.debug("Could not parse remote SDP while checking video codecs.") + return None + + codecs: list[str] = [] + for media in description.media: + if getattr(media, "kind", None) != "video": + continue + for codec in media.rtp.codecs: + mime_type = str(getattr(codec, "mimeType", "")).strip() + if mime_type: + codecs.append(mime_type) + return tuple(dict.fromkeys(codecs)) + + class WebRTCControlSignal(IntEnum): """Rank-orchestration signals shared by the single-session runtimes.""" @@ -61,7 +111,8 @@ class ManagedWebRTCSession: """Per-session state for the single active WebRTC peer connection.""" runtime: Any - video_track: BufferedVideoTrack + video_track: BufferedVideoTrack | NVENCVideoTrack + video_encoder: VideoEncoder peer_connection: Any resampler: KeyboardResampler control_channel: Any | None = None @@ -158,6 +209,108 @@ def _make_resampler(self, *, start_v: float) -> KeyboardResampler: def _register_extra_peer_handlers(self, peer_connection: Any) -> None: """Register optional extra peer-connection event handlers.""" + def _prefer_h264_video_codec(self, *, transceiver: Any) -> None: + """Constrain the transceiver's codec preferences to H.264 variants. + + Required when the selected encoder emits pre-encoded H.264 packets + (``av.Packet`` route through ``H264Encoder.pack()``): if the SDP + negotiates VP8/VP9 instead, aiortc will pack the H.264 bitstream + under the wrong codec header and the receiver will fail to decode. + + If the local aiortc build does not advertise H.264, no preference + is set; the SDP-time fallback in ``_enforce_h264_or_fallback`` + will then swap the encoder to :class:`DefaultRTCEncoder`. + """ + caps = RTCRtpSender.getCapabilities("video") + h264_codecs = [c for c in caps.codecs if c.mimeType.lower() == "video/h264"] + if not h264_codecs: + return + transceiver.setCodecPreferences(h264_codecs) + + def _prepare_video_encoder_for_offer( + self, + *, + video_encoder: VideoEncoder, + offer_sdp: str, + ) -> VideoEncoder: + """Select a session encoder compatible with the browser offer.""" + if video_encoder.prefers_codec != "h264": + return video_encoder + + offered_video_codecs = _sdp_video_codecs(offer_sdp) + if offered_video_codecs is None: + return video_encoder + if any(codec.lower() == "video/h264" for codec in offered_video_codecs): + return video_encoder + + offered = ", ".join(offered_video_codecs) or "" + if getattr(self.runtime_config, "encoder_backend", None) == "nvenc": + video_encoder.close() + raise RuntimeError( + "encoder_backend='nvenc' requested but the browser offer does " + f"not advertise H.264; offered video codecs: {offered}." + ) + + logger.warning( + "Hardware encoder emits H.264, but the browser offer does not " + "advertise H.264 (offered: {}). Using aiortc software encoder " + "for this session.", + offered, + ) + video_encoder.close() + return DefaultRTCEncoder(fps=self.fps) + + def _enforce_h264_or_fallback( + self, + *, + transceiver: Any, + managed_session: ManagedWebRTCSession, + num_frames: int, + ) -> None: + """Verify H.264 was negotiated; swap to the software encoder if not. + + aiortc exposes the negotiated codec set on + ``RTCRtpTransceiver._codecs`` after ``setLocalDescription``. We + read it via that attribute (aiortc-internal, but stable in the + pinned version) and, if H.264 did not land, close the hardware + encoder and install a :class:`DefaultRTCEncoder` with a + :class:`BufferedVideoTrack` on the same sender before the first + RTP packet flies. ``replaceTrack`` does not renegotiate; aiortc's + RTP loop will encode raw ``av.VideoFrame`` output with whatever + codec (VP8/VP9/H.264) actually landed in the SDP. + """ + negotiated = getattr(transceiver, "_codecs", None) or [] + if negotiated and negotiated[0].mimeType.lower() == "video/h264": + logger.info( + "Video codec negotiated: {} (hardware encoder path active).", + negotiated[0].mimeType, + ) + return + + chosen = negotiated[0].mimeType if negotiated else "" + if getattr(self.runtime_config, "encoder_backend", None) == "nvenc": + managed_session.video_encoder.close() + raise RuntimeError( + "encoder_backend='nvenc' requested but SDP negotiation landed on " + f"{chosen!r}; cannot stream pre-encoded H.264 packets." + ) + logger.warning( + "H.264 preferred by hardware encoder but SDP negotiation " + "landed on {!r}; swapping to the software encoder before " + "streaming begins.", + chosen, + ) + # Close the hardware encoder so its NVENC session is released + # promptly; the software adapter has no hardware resources to + # release itself. + managed_session.video_encoder.close() + + fallback_encoder = DefaultRTCEncoder(fps=self.fps) + fallback_track = fallback_encoder.create_track(maxsize=num_frames) + transceiver.sender.replaceTrack(fallback_track) + managed_session.video_encoder = fallback_encoder + managed_session.video_track = fallback_track + def _on_offer_received(self, offer_sdp: str) -> None: """Hook invoked with the remote offer SDP before negotiation.""" @@ -283,8 +436,20 @@ async def _create_answer_with_runtime_ready_locked( # frames than steady state; sizing to it would force a per-chunk # stall, so we size to the steady-state count. num_frames = self._runtime.peek_steady_chunk_num_frames() - video_track = BufferedVideoTrack(fps=self.fps, maxsize=num_frames) - peer_connection.addTrack(video_track) + video_encoder: VideoEncoder = self._prepare_video_encoder_for_offer( + video_encoder=self._runtime.video_encoder, + offer_sdp=offer_sdp, + ) + video_track = video_encoder.create_track(maxsize=num_frames) + # Use ``addTransceiver`` (not ``addTrack``) so we can constrain the + # SDP m-line's codec list via ``setCodecPreferences`` when the + # encoder emits pre-encoded H.264 packets. + video_transceiver = peer_connection.addTransceiver( + video_track, + direction="sendonly", + ) + if video_encoder.prefers_codec == "h264": + self._prefer_h264_video_codec(transceiver=video_transceiver) # Start the resampler's virtual clock at 0; the real anchor is set # in the ``on_datachannel`` handler so chunk 0's window starts when # input can actually arrive. @@ -293,6 +458,7 @@ async def _create_answer_with_runtime_ready_locked( managed_session = ManagedWebRTCSession( runtime=self._runtime, video_track=video_track, + video_encoder=video_encoder, peer_connection=peer_connection, resampler=resampler, last_client_message_at=loop.time(), @@ -350,6 +516,12 @@ async def on_connectionstatechange() -> None: answer = await peer_connection.createAnswer() await peer_connection.setLocalDescription(answer) await wait_for_ice_gathering_complete(peer_connection) + if video_encoder.prefers_codec == "h264": + self._enforce_h264_or_fallback( + transceiver=video_transceiver, + managed_session=managed_session, + num_frames=num_frames, + ) local_description = peer_connection.localDescription if local_description is None: raise RuntimeError("Peer connection did not produce local description.") @@ -548,6 +720,7 @@ async def _generation_worker( runtime = managed_session.runtime resampler = managed_session.resampler video_track = managed_session.video_track + video_encoder = managed_session.video_encoder # Stay idle until the user interacts. Generating eagerly would burn # GPU cycles on a still scene the viewer never sees. Once an event @@ -603,6 +776,7 @@ async def _generation_worker( consumed_action_arrivals.append( managed_session.pending_action_arrivals.popleft() ) + t_after_sample = loop.time() try: result = await runtime.generate_chunk( segments=segments, frame_times=frame_times @@ -617,11 +791,24 @@ async def _generation_worker( return continue t_after_gen = loop.time() - enqueued = await video_track.enqueue_chunk(result.video_chunk) + delivery = await video_encoder.deliver_chunk( + result.video_chunk, + video_track, + force_keyframe=False, + ) + enqueued = delivery.num_frames t_after_enqueue = loop.time() gen_ms = (t_after_gen - t_before_gen) * 1e3 + sample_ms = (t_after_sample - t_before_gen) * 1e3 + runtime_call_ms = (t_after_gen - t_after_sample) * 1e3 enqueue_ms = (t_after_enqueue - t_after_gen) * 1e3 + chunk_total_ms = (t_after_enqueue - t_before_gen) * 1e3 + chunk_fps = ( + result.num_frames * 1000.0 / chunk_total_ms + if chunk_total_ms > 0 + else 0.0 + ) play_ms = result.num_frames * 1000.0 / video_track.fps lag_ms = (t_after_enqueue - resampler.next_chunk_start_v) * 1e3 control_latency_ms = ( @@ -629,19 +816,49 @@ async def _generation_worker( if consumed_action_arrivals else None ) - logger.debug( - "Chunk done chunk={} num_frames={} segments={} enqueued={} " - "gen_ms={:.1f} enqueue_ms={:.1f} play_ms={:.1f} queue_depth={} " - "lag_ms={:.1f}", + track_dropped_packets = getattr( + video_track, + "dropped_packets", + None, + ) + runtime_stats = _performance_stats_payload(result.stats) + log_stats = { + **runtime_stats, + "delivery_encode_ms": round(delivery.encode_ms, 1), + } + if isinstance(track_dropped_packets, int): + log_stats["track_dropped_packets"] = float(track_dropped_packets) + extra_stats = _format_performance_stats(log_stats) + log_level = ( + logger.info + if ( + result.chunk_index <= 2 + or result.chunk_index % 10 == 0 + or chunk_total_ms > play_ms + or lag_ms > play_ms + ) + else logger.debug + ) + log_level( + "WebRTC chunk done chunk={} num_frames={} segments={} " + "enqueued={} encoder={} gen_ms={:.1f} sample_ms={:.1f} " + "runtime_call_ms={:.1f} delivery_ms={:.1f} " + "play_ms={:.1f} chunk_fps={:.1f} queue_depth={} " + "lag_ms={:.1f} {}", result.chunk_index, result.num_frames, len(segments), enqueued, + delivery.backend, gen_ms, + sample_ms, + runtime_call_ms, enqueue_ms, play_ms, + chunk_fps, video_track.qsize(), lag_ms, + extra_stats, ) channel = managed_session.control_channel @@ -658,11 +875,22 @@ async def _generation_worker( }, "model": self._model_name(), "gen_ms": round(gen_ms, 1), + "sample_ms": round(sample_ms, 1), + "runtime_call_ms": round(runtime_call_ms, 1), "enqueue_ms": round(enqueue_ms, 1), + "delivery_ms": round(enqueue_ms, 1), + "delivery_encode_ms": round(delivery.encode_ms, 1), + "encoder_backend": delivery.backend, + "keyframes": delivery.num_keyframes, + "chunk_total_ms": round(chunk_total_ms, 1), + "chunk_fps": round(chunk_fps, 1), "play_ms": round(play_ms, 1), "queue_depth": video_track.qsize(), "lag_ms": round(lag_ms, 1), } + if isinstance(track_dropped_packets, int): + payload["track_dropped_packets"] = track_dropped_packets + payload.update(runtime_stats) payload.update(self._chunk_done_extra()) if control_latency_ms is not None: payload["latency_ms"] = round(control_latency_ms, 1) diff --git a/flashdreams/flashdreams/serving/webrtc/media.py b/flashdreams/flashdreams/serving/webrtc/media.py index e438c3b1d..2912a042f 100644 --- a/flashdreams/flashdreams/serving/webrtc/media.py +++ b/flashdreams/flashdreams/serving/webrtc/media.py @@ -4,7 +4,8 @@ from __future__ import annotations import asyncio -from collections.abc import Callable +import contextlib +from collections.abc import Callable, Sequence from fractions import Fraction from typing import TYPE_CHECKING @@ -12,6 +13,7 @@ from aiortc import MediaStreamTrack from aiortc.mediastreams import MediaStreamError from av import VideoFrame +from av.packet import Packet from loguru import logger from flashdreams.serving.realtime.media import tensor_chunk_to_rgb_frames @@ -133,3 +135,132 @@ async def close(self) -> None: break self._frames.put_nowait(None) self.stop() + + +class NVENCVideoTrack(MediaStreamTrack): + """WebRTC video track that delivers pre-encoded H.264 packets. + + Paired with ``PyNvHardwareEncoder``: :meth:`recv` returns + :class:`av.Packet` (not :class:`av.VideoFrame`), which aiortc's + ``RTCRtpSender`` routes through ``H264Encoder.pack()`` for RTP + fragmentation only. The encoder sets ``pts`` and ``time_base`` on + each packet before enqueueing; this track only paces delivery to + ``fps``. The async enqueue path applies backpressure so the browser + receives every frame in timestamp order instead of seeing silent + server-side drops. + """ + + kind = "video" + + def __init__(self, *, fps: int, maxsize: int) -> None: + super().__init__() + if fps <= 0: + raise ValueError("fps must be > 0") + if maxsize <= 0: + raise ValueError("maxsize must be > 0") + self._fps = fps + self._frame_interval_s = 1.0 / fps + self._next_deadline_s: float | None = None + self._maxsize = maxsize + self._packets: asyncio.Queue[Packet | None] = asyncio.Queue( + maxsize=maxsize, + ) + self._closed = False + self._dropped_packets = 0 + + @property + def fps(self) -> int: + return self._fps + + @property + def maxsize(self) -> int: + return self._maxsize + + @property + def dropped_packets(self) -> int: + return self._dropped_packets + + def qsize(self) -> int: + return self._packets.qsize() + + async def enqueue_encoded_packet(self, packet: Packet) -> bool: + """Enqueue one encoded packet, waiting for sender-side queue space.""" + if self._closed: + return False + await self._packets.put(packet) + return True + + async def enqueue_encoded_packets(self, packets: Sequence[Packet]) -> int: + """Enqueue encoded packets in order, applying backpressure if full.""" + for i, packet in enumerate(packets): + if not await self.enqueue_encoded_packet(packet): + return i + return len(packets) + + def enqueue_encoded_packet_nowait(self, packet: Packet) -> bool: + """Synchronously enqueue one packet on the loop thread. + + This compatibility helper is intentionally lossy on overflow. + The production NVENC path uses :meth:`enqueue_encoded_packets` + so playback preserves every frame; tests and diagnostic probes + can still use this helper when they explicitly want nonblocking + enqueue semantics. + """ + if self._closed: + return False + if self._maxsize > 0 and self._packets.full(): + with contextlib.suppress(asyncio.QueueEmpty): + self._packets.get_nowait() + self._dropped_packets += 1 + logger.debug( + "NVENCVideoTrack overflow: dropped oldest packet " + "(total dropped={})", + self._dropped_packets, + ) + self._packets.put_nowait(packet) + return True + + async def recv(self) -> Packet: + if self._closed: + raise MediaStreamError + + loop = asyncio.get_running_loop() + t_get_start = loop.time() + packet = await self._packets.get() + if packet is None: + raise MediaStreamError + get_wait_ms = (loop.time() - t_get_start) * 1000.0 + first_packet = self._next_deadline_s is None + just_stalled = (not first_packet) and get_wait_ms > _STALL_THRESHOLD_MS + + now_s = loop.time() + if first_packet or just_stalled: + self._next_deadline_s = now_s + else: + proposed = self._next_deadline_s + self._frame_interval_s + wait_s = proposed - now_s + if wait_s > 0: + await asyncio.sleep(wait_s) + self._next_deadline_s = proposed + else: + if -wait_s * 1000.0 > _PACING_LAG_LOG_MS: + logger.debug( + "NVENCVideoTrack pacing lag: deadline {:.1f}ms " + "behind walltime; re-anchoring (queue depth {}).", + -wait_s * 1000.0, + self._packets.qsize(), + ) + self._next_deadline_s = now_s + return packet + + async def close(self) -> None: + if self._closed: + return + self._closed = True + while True: + try: + self._packets.get_nowait() + except asyncio.QueueEmpty: + break + self._packets.put_nowait(None) + self.stop() diff --git a/flashdreams/tests/test_encoders.py b/flashdreams/tests/test_encoders.py new file mode 100644 index 000000000..0c7b574bd --- /dev/null +++ b/flashdreams/tests/test_encoders.py @@ -0,0 +1,508 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the ``VideoEncoder`` abstraction and ``select_encoder`` factory. + +Covers: + +- Protocol conformance for the software-path adapter. +- ``select_encoder`` branch coverage under ``backend={"default","auto","nvenc"}``. +- The two failure semantics of the capability probe: a Stage-1 no is a + silent fallback; a Stage-2 ``CreateEncoder`` raise is a hard error + (never a silent fallback). The latter is the regression guard that + keeps a future refactor from masking driver problems. +- The cross-chunk packet-ordering invariant: sequential ``await + deliver_chunk(...)`` calls must produce monotonically-ordered packets + on the paired track. A ``create_task``-style refactor would silently + break this — see :class:`TestDeliverChunkOrdering`. +- Compatibility guards for the two upstream libraries we couple to + (``aiortc`` and ``PyNvVideoCodec``). +""" + +from __future__ import annotations + +import asyncio +import threading +from fractions import Fraction +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from av.packet import Packet + +from flashdreams.serving.webrtc import encoders as enc_mod +from flashdreams.serving.webrtc.encoders import ( + ChunkDeliveryResult, + DefaultRTCEncoder, + EncoderInitError, + VideoEncoder, + select_encoder, +) +from flashdreams.serving.webrtc.media import NVENCVideoTrack + +pytestmark = pytest.mark.ci_cpu + +_SELECT_KW = dict( + width=1280, + height=704, + fps=30, + bitrate=6_000_000, + gpu_id=0, + gop=30, +) + + +# --------------------------------------------------------------------------- +# Protocol conformance +# --------------------------------------------------------------------------- + + +class TestVideoEncoderProtocol: + def test_default_encoder_satisfies_protocol(self) -> None: + assert isinstance(DefaultRTCEncoder(fps=30), VideoEncoder) + + def test_default_encoder_advertises_no_codec_preference(self) -> None: + assert DefaultRTCEncoder(fps=30).prefers_codec is None + + def test_default_encoder_rejects_bad_fps(self) -> None: + with pytest.raises(ValueError, match="fps"): + DefaultRTCEncoder(fps=0) + + +# --------------------------------------------------------------------------- +# select_encoder: "default" backend never touches NVENC +# --------------------------------------------------------------------------- + + +class TestSelectDefaultBackend: + def test_returns_default_encoder(self) -> None: + enc = select_encoder(backend="default", **_SELECT_KW) + assert isinstance(enc, DefaultRTCEncoder) + assert enc.fps == 30 + + def test_does_not_call_getencodercaps(self) -> None: + # Even if the library exists, "default" must not touch it — this + # keeps the software path deterministic and safe on hosts where + # GetEncoderCaps might have side effects. + fake_nvc = MagicMock() + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", True), + patch.object(enc_mod, "nvc", fake_nvc), + ): + select_encoder(backend="default", **_SELECT_KW) + fake_nvc.GetEncoderCaps.assert_not_called() + + def test_works_even_when_library_missing(self) -> None: + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", False), + patch.object(enc_mod, "nvc", None), + ): + enc = select_encoder(backend="default", **_SELECT_KW) + assert isinstance(enc, DefaultRTCEncoder) + + +# --------------------------------------------------------------------------- +# select_encoder: Stage 1 (GetEncoderCaps) failure paths +# --------------------------------------------------------------------------- + + +class TestSelectStage1Failure: + """Stage 1 says "environment can't do this" — auto silently falls back; + nvenc raises loudly.""" + + def test_library_missing_auto_falls_back(self, caplog) -> None: + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", False), + patch.object(enc_mod, "nvc", None), + ): + enc = select_encoder(backend="auto", **_SELECT_KW) + assert isinstance(enc, DefaultRTCEncoder) + + def test_library_missing_nvenc_raises(self) -> None: + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", False), + patch.object(enc_mod, "nvc", None), + ): + with pytest.raises(EncoderInitError, match="not installed"): + select_encoder(backend="nvenc", **_SELECT_KW) + + @pytest.mark.parametrize( + "caps_effect, expected_reason_frag", + [ + (RuntimeError("driver comms error"), "driver comms error"), + ({}, "no capabilities"), + ( + { + "width_max": 640, + "height_max": 480, + "width_min": 32, + "height_min": 32, + }, + "exceeds driver maximum", + ), + ( + { + "width_max": 8192, + "height_max": 8192, + "width_min": 1920, + "height_min": 1088, + }, + "below driver minimum", + ), + ], + ids=["caps_raise", "caps_empty", "caps_max_too_small", "caps_min_too_big"], + ) + def test_caps_failure_auto_falls_back( + self, + caps_effect, + expected_reason_frag, + ) -> None: + fake_nvc = MagicMock() + if isinstance(caps_effect, BaseException): + fake_nvc.GetEncoderCaps.side_effect = caps_effect + else: + fake_nvc.GetEncoderCaps.return_value = caps_effect + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", True), + patch.object(enc_mod, "nvc", fake_nvc), + ): + enc = select_encoder(backend="auto", **_SELECT_KW) + assert isinstance(enc, DefaultRTCEncoder) + # ``CreateEncoder`` must not be reached when Stage 1 fails. + fake_nvc.CreateEncoder.assert_not_called() + + def test_caps_failure_nvenc_raises_with_reason(self) -> None: + fake_nvc = MagicMock() + fake_nvc.GetEncoderCaps.side_effect = RuntimeError("driver comms error") + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", True), + patch.object(enc_mod, "nvc", fake_nvc), + ): + with pytest.raises(EncoderInitError, match="driver comms error"): + select_encoder(backend="nvenc", **_SELECT_KW) + + +# --------------------------------------------------------------------------- +# select_encoder: Stage 2 (CreateEncoder) failure — HARD ERROR +# --------------------------------------------------------------------------- + + +class TestSelectStage2HardError: + """The key regression guard: if Stage 1 says supported but Stage 2's + ``CreateEncoder`` raises, ``select_encoder`` must re-raise. Silently + falling back to the software encoder here would mask a real bug + (driver, session pool exhaustion, hardware fault). Any future refactor + that makes this test fail is almost certainly regressing that semantic. + """ + + def _fake_nvc_caps_ok(self) -> MagicMock: + fake = MagicMock() + fake.GetEncoderCaps.return_value = { + "width_max": 8192, + "height_max": 8192, + "width_min": 32, + "height_min": 32, + } + fake.FORCEIDR = 0x1 + return fake + + def test_construct_failure_reraises_under_auto(self) -> None: + fake_nvc = self._fake_nvc_caps_ok() + fake_nvc.CreateEncoder.side_effect = RuntimeError( + "NVENC session pool exhausted" + ) + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", True), + patch.object(enc_mod, "nvc", fake_nvc), + ): + with pytest.raises(RuntimeError, match="session pool exhausted"): + select_encoder(backend="auto", **_SELECT_KW) + + def test_construct_failure_reraises_under_nvenc(self) -> None: + fake_nvc = self._fake_nvc_caps_ok() + fake_nvc.CreateEncoder.side_effect = RuntimeError("hardware fault") + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", True), + patch.object(enc_mod, "nvc", fake_nvc), + ): + with pytest.raises(RuntimeError, match="hardware fault"): + select_encoder(backend="nvenc", **_SELECT_KW) + + +# --------------------------------------------------------------------------- +# ChunkDeliveryResult +# --------------------------------------------------------------------------- + + +class TestChunkDeliveryResult: + def test_is_frozen_dataclass(self) -> None: + result = ChunkDeliveryResult( + backend="fake", + num_frames=4, + num_keyframes=1, + encode_ms=1.5, + ) + with pytest.raises((AttributeError, Exception)): + result.backend = "other" # ty:ignore[invalid-assignment] + + +# --------------------------------------------------------------------------- +# DefaultRTCEncoder.deliver_chunk delegates to track.enqueue_chunk +# --------------------------------------------------------------------------- + + +class _FakeBufferedVideoTrack: + """Minimal stand-in that ``isinstance(track, BufferedVideoTrack)`` + treats as a real track (subclassing lets us bypass MediaStreamTrack's + aiortc runtime dependencies without breaking the isinstance check).""" + + def __init__(self) -> None: + self.enqueued_chunks: list = [] + + async def enqueue_chunk(self, chunk) -> int: + self.enqueued_chunks.append(chunk) + return 4 + + +class TestDefaultRTCEncoderDeliver: + @pytest.mark.asyncio + async def test_deliver_chunk_returns_frames_from_track(self) -> None: + from flashdreams.serving.webrtc import media as media_mod + + fake_track = _FakeBufferedVideoTrack() + # Patch the isinstance check inside deliver_chunk to accept our fake. + with patch.object(media_mod, "BufferedVideoTrack", _FakeBufferedVideoTrack): + enc = DefaultRTCEncoder(fps=30) + result = await enc.deliver_chunk( + SimpleNamespace(shape=(4, 3, 8, 8)), # ty:ignore[invalid-argument-type] + fake_track, # ty:ignore[invalid-argument-type] + ) + assert result.backend == "aiortc" + assert result.num_frames == 4 + assert result.num_keyframes == 0 + assert len(fake_track.enqueued_chunks) == 1 + + @pytest.mark.asyncio + async def test_deliver_chunk_rejects_wrong_track_type(self) -> None: + enc = DefaultRTCEncoder(fps=30) + with pytest.raises(TypeError, match="BufferedVideoTrack"): + await enc.deliver_chunk( + SimpleNamespace(), # ty:ignore[invalid-argument-type] + SimpleNamespace(), # ty:ignore[invalid-argument-type] + ) + + +# --------------------------------------------------------------------------- +# Compatibility guards for upstream libraries +# --------------------------------------------------------------------------- + + +class TestCompatGuards: + """These tests exist to fail loudly when an upstream dependency + changes its public surface underneath us. They are cheap and + catch dependency drift far earlier than a runtime failure would.""" + + def test_aiortc_h264_encoder_has_pack_and_encode(self) -> None: + from aiortc.codecs.h264 import H264Encoder + + assert callable(getattr(H264Encoder, "encode", None)), ( + "aiortc H264Encoder.encode disappeared — the software encoder " + "contract this design assumes has changed." + ) + assert callable(getattr(H264Encoder, "pack", None)), ( + "aiortc H264Encoder.pack disappeared — the pre-encoded packet " + "path this design relies on has changed." + ) + + def test_aiortc_sender_module_importable(self) -> None: + # Any structural change to rtcrtpsender that breaks import will + # break the runtime; catch it here before the first RTP packet. + import aiortc.rtcrtpsender # noqa: F401 + + def test_getencodercaps_callable_when_library_available(self) -> None: + if not enc_mod._PYNVVIDEOCODEC_AVAILABLE: + pytest.skip("PyNvVideoCodec is not installed on this host") + assert callable(getattr(enc_mod.nvc, "GetEncoderCaps", None)), ( + "PyNvVideoCodec.GetEncoderCaps renamed or removed — Stage 1 " + "capability probe cannot function." + ) + + +# --------------------------------------------------------------------------- +# Cross-chunk packet-ordering invariant +# --------------------------------------------------------------------------- + + +class _OrderingFakeEncoder: + """Minimal encoder that mimics :class:`PyNvHardwareEncoder`'s async + encode-then-enqueue contract without needing CUDA or PyNvVideoCodec. + + Encoding runs on an ``asyncio.to_thread`` worker, then the encoded + packets are enqueued through the track's backpressured async API. + This lets us assert the end-to-end ordering property without any GPU + dependency. + + The pts counter is guarded by a lock because the fire-and-forget + test dispatches multiple ``deliver_chunk`` calls concurrently, and + ``self._pts_counter += 1`` is not atomic across Python bytecodes. + """ + + backend = "fake" + prefers_codec: str | None = "h264" + + def __init__(self, *, fps: int, frames_per_chunk: int) -> None: + self.fps = fps + self._frames_per_chunk = frames_per_chunk + self._pts_counter = 0 + self._pts_counter_lock = threading.Lock() + self._time_base = Fraction(1, 90_000) + + def create_track(self, *, maxsize: int) -> NVENCVideoTrack: + return NVENCVideoTrack(fps=self.fps, maxsize=maxsize) + + async def deliver_chunk( + self, + chunk: object, + track: NVENCVideoTrack, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + del chunk, force_keyframe + frames = self._frames_per_chunk + + def _encode_worker() -> list[Packet]: + # One packet per frame, monotonically-increasing pts. The + # caller enqueues the collected packets after the worker + # returns, mirroring PyNvHardwareEncoder.deliver_chunk. + packets: list[Packet] = [] + for _ in range(frames): + packet = Packet(b"\x00\x00\x00\x01\x67") + with self._pts_counter_lock: + pts_frame_index = self._pts_counter + self._pts_counter += 1 + packet.pts = (pts_frame_index * 90_000) // self.fps + packet.time_base = self._time_base + packets.append(packet) + return packets + + packets = await asyncio.to_thread(_encode_worker) + enqueued = await track.enqueue_encoded_packets(packets) + return ChunkDeliveryResult( + backend=self.backend, + num_frames=enqueued, + num_keyframes=0, + encode_ms=0.0, + ) + + def close(self) -> None: + return + + +# 30 chunks × 8 frames = 240 packets — realistic scale for an interactive +# session (about 8 seconds of 30 fps video, or a few minutes of chunked +# generation at typical omnidreams cadence). At the encoder's real fps=30 +# the track's recv() pacing would make draining take ~8 s per test, so we +# use a much higher fps here to keep the pacing throttle negligible while +# preserving distinct, monotonic pts values (pts = i at fps=90_000 since +# ``(i * 90_000) // 90_000 == i``). +_ORDERING_NUM_CHUNKS = 30 +_ORDERING_FRAMES_PER_CHUNK = 8 +_ORDERING_TOTAL_FRAMES = _ORDERING_NUM_CHUNKS * _ORDERING_FRAMES_PER_CHUNK +_ORDERING_FPS = 90_000 + + +class TestDeliverChunkOrdering: + """Regression guard for the manager's sequential await pattern. + + The manager's ``_generation_worker`` does ``await deliver_chunk(...)`` + in a single loop, which forces chunk N's packets to land on the track + before chunk N+1 starts encoding. If a future refactor swaps that for + ``asyncio.create_task(deliver_chunk(...))`` — or introduces a + producer-consumer queue between generation and encoding without a + reorder buffer — chunks could complete out of order and packets + would land on the track with non-monotonic ``pts``. + + These tests replay the manager's sequential-await pattern against a + fake encoder that mirrors :class:`PyNvHardwareEncoder`'s async + marshaling. They pass under the current ``await``-based design and + would fail under a fire-and-forget rewrite. + """ + + @pytest.mark.asyncio + async def test_sequential_await_produces_monotonic_pts(self) -> None: + encoder = _OrderingFakeEncoder( + fps=_ORDERING_FPS, + frames_per_chunk=_ORDERING_FRAMES_PER_CHUNK, + ) + track = encoder.create_track(maxsize=_ORDERING_TOTAL_FRAMES) + + for _ in range(_ORDERING_NUM_CHUNKS): + await encoder.deliver_chunk(object(), track) + + seen_pts: list[int] = [] + for _ in range(_ORDERING_TOTAL_FRAMES): + packet = await asyncio.wait_for(track.recv(), timeout=1.0) + # ``_OrderingFakeEncoder`` always sets pts before enqueueing; the + # ``av.Packet.pts`` field is nullable at the type level, so narrow. + assert packet.pts is not None + seen_pts.append(int(packet.pts)) + + # Strictly monotonic and matches the exact expected pts sequence. + expected = [ + (i * 90_000) // _ORDERING_FPS for i in range(_ORDERING_TOTAL_FRAMES) + ] + assert seen_pts == expected, ( + "packets emitted out of order across chunks: " + f"first 16 got={seen_pts[:16]} expected={expected[:16]}" + ) + assert seen_pts == sorted(seen_pts) + + @pytest.mark.asyncio + async def test_fire_and_forget_pattern_would_break_ordering(self) -> None: + """Counter-check: if the manager ever spawns deliver_chunk via + ``asyncio.create_task`` (fire-and-forget), packets *can* interleave + across chunks. This test documents that failure mode so a future + refactor that reintroduces the anti-pattern can be caught by + comparing behaviour against the sequential-await test above. + + The test is not a bug — it is a demonstration that the sequential + await in the manager is *load-bearing* for ordering. + """ + encoder = _OrderingFakeEncoder( + fps=_ORDERING_FPS, + frames_per_chunk=_ORDERING_FRAMES_PER_CHUNK, + ) + track = encoder.create_track(maxsize=_ORDERING_TOTAL_FRAMES) + + # Bias interleaving: create_task, then gather. Individual chunks + # each finish quickly; scheduler order within the loop does not + # guarantee packets arrive in the same order the tasks were spawned. + tasks = [ + asyncio.create_task(encoder.deliver_chunk(object(), track)) + for _ in range(_ORDERING_NUM_CHUNKS) + ] + await asyncio.gather(*tasks) + + seen_pts: list[int] = [] + for _ in range(_ORDERING_TOTAL_FRAMES): + packet = await asyncio.wait_for(track.recv(), timeout=1.0) + # ``_OrderingFakeEncoder`` always sets pts before enqueueing; the + # ``av.Packet.pts`` field is nullable at the type level, so narrow. + assert packet.pts is not None + seen_pts.append(int(packet.pts)) + + # Every pts value must be present exactly once; that part is a + # correctness invariant of the fake encoder (guarded by the lock + # around ``_pts_counter``). Duplicate pts here would be a bug in + # the fake, not in the code under test. + expected_set = { + (i * 90_000) // _ORDERING_FPS for i in range(_ORDERING_TOTAL_FRAMES) + } + assert set(seen_pts) == expected_set + assert len(seen_pts) == _ORDERING_TOTAL_FRAMES + + # The full sequence, however, is not necessarily monotonic — the + # fire-and-forget pattern permits interleave. We do NOT assert + # monotonicity here; on the contrary, the moment this test starts + # asserting monotonic pts, the manager's sequential-await guard + # has become unnecessary (and the test above becomes redundant). diff --git a/flashdreams/tests/test_nvenc_track.py b/flashdreams/tests/test_nvenc_track.py new file mode 100644 index 000000000..09b8e6076 --- /dev/null +++ b/flashdreams/tests/test_nvenc_track.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for :class:`NVENCVideoTrack` — pre-encoded H.264 packet plumbing. + +The track is a thin queue in front of aiortc's ``RTCRtpSender``. What +matters is that: + +- ``recv()`` returns exactly the enqueued :class:`av.Packet` (aiortc's + ``H264Encoder.pack()`` reads ``payload``, ``pts`` and ``time_base`` + directly), and +- the async enqueue path applies backpressure instead of dropping packets, + preserving the 30 fps media timeline the browser receives. + +The legacy nonblocking helper is still tested separately because it is useful +for diagnostics, but the production NVENC encoder path does not use it. +""" + +from __future__ import annotations + +import asyncio +from fractions import Fraction + +import pytest +from aiortc.mediastreams import MediaStreamError +from av.packet import Packet + +from flashdreams.serving.webrtc.media import NVENCVideoTrack + +pytestmark = pytest.mark.ci_cpu + + +def _mk_packet(pts: int, payload: bytes = b"\x00\x00\x00\x01\x67") -> Packet: + pkt = Packet(payload) + pkt.pts = pts + pkt.time_base = Fraction(1, 90_000) + return pkt + + +class TestConstructorValidation: + def test_rejects_bad_fps(self) -> None: + with pytest.raises(ValueError, match="fps"): + NVENCVideoTrack(fps=0, maxsize=8) + + def test_rejects_bad_maxsize(self) -> None: + with pytest.raises(ValueError, match="maxsize"): + NVENCVideoTrack(fps=30, maxsize=0) + + def test_exposes_fps_and_maxsize(self) -> None: + track = NVENCVideoTrack(fps=30, maxsize=8) + assert track.fps == 30 + assert track.maxsize == 8 + assert track.qsize() == 0 + assert track.dropped_packets == 0 + + +class TestEnqueue: + def test_enqueue_returns_true_on_open_track(self) -> None: + track = NVENCVideoTrack(fps=30, maxsize=8) + assert track.enqueue_encoded_packet_nowait(_mk_packet(0)) is True + assert track.qsize() == 1 + + @pytest.mark.asyncio + async def test_async_enqueue_waits_for_queue_space(self) -> None: + track = NVENCVideoTrack(fps=30, maxsize=1) + assert await track.enqueue_encoded_packet(_mk_packet(0)) is True + + enqueue_task = asyncio.create_task(track.enqueue_encoded_packet(_mk_packet(1))) + await asyncio.sleep(0) + assert not enqueue_task.done() + + first = await asyncio.wait_for(track.recv(), timeout=1.0) + assert first.pts == 0 + assert await asyncio.wait_for(enqueue_task, timeout=1.0) is True + assert track.qsize() == 1 + assert track.dropped_packets == 0 + + second = await asyncio.wait_for(track.recv(), timeout=1.0) + assert second.pts == 1 + + @pytest.mark.asyncio + async def test_async_enqueue_many_preserves_order(self) -> None: + track = NVENCVideoTrack(fps=30, maxsize=2) + packets = [_mk_packet(pts=i) for i in range(4)] + + enqueue_task = asyncio.create_task(track.enqueue_encoded_packets(packets)) + seen_pts: list[int] = [] + for _ in packets: + packet = await asyncio.wait_for(track.recv(), timeout=1.0) + assert packet.pts is not None + seen_pts.append(packet.pts) + + assert await asyncio.wait_for(enqueue_task, timeout=1.0) == len(packets) + assert seen_pts == [0, 1, 2, 3] + assert track.dropped_packets == 0 + + @pytest.mark.asyncio + async def test_enqueue_returns_false_on_closed_track(self) -> None: + track = NVENCVideoTrack(fps=30, maxsize=8) + await track.close() + assert await track.enqueue_encoded_packet(_mk_packet(0)) is False + assert track.enqueue_encoded_packet_nowait(_mk_packet(0)) is False + + +class TestRecv: + @pytest.mark.asyncio + async def test_recv_returns_enqueued_packet_unchanged(self) -> None: + # aiortc's H264Encoder.pack() consumes bytes(packet), packet.pts + # and packet.time_base directly, so recv() must not mutate them. + track = NVENCVideoTrack(fps=30, maxsize=8) + original = _mk_packet(pts=1234, payload=b"\x00\x00\x00\x01\x25\xaa") + track.enqueue_encoded_packet_nowait(original) + got = await asyncio.wait_for(track.recv(), timeout=1.0) + assert bytes(got) == b"\x00\x00\x00\x01\x25\xaa" + assert got.pts == 1234 + assert got.time_base == Fraction(1, 90_000) + + @pytest.mark.asyncio + async def test_recv_on_closed_track_raises_mediastreamerror(self) -> None: + track = NVENCVideoTrack(fps=30, maxsize=8) + await track.close() + with pytest.raises(MediaStreamError): + await asyncio.wait_for(track.recv(), timeout=1.0) + + +class TestOverflow: + """The explicit nowait helper drops the oldest packet on overflow. + + The production async enqueue path above backpressures instead. This + branch exists for callers that deliberately choose nonblocking enqueue. + """ + + def test_full_queue_drops_oldest(self) -> None: + track = NVENCVideoTrack(fps=30, maxsize=2) + p0 = _mk_packet(pts=0) + p1 = _mk_packet(pts=1) + p2 = _mk_packet(pts=2) + assert track.enqueue_encoded_packet_nowait(p0) is True + assert track.enqueue_encoded_packet_nowait(p1) is True + # This must drop p0, not refuse p2. + assert track.enqueue_encoded_packet_nowait(p2) is True + assert track.qsize() == 2 + assert track.dropped_packets == 1 + + @pytest.mark.asyncio + async def test_overflow_preserves_newest_packet(self) -> None: + track = NVENCVideoTrack(fps=30, maxsize=2) + track.enqueue_encoded_packet_nowait(_mk_packet(pts=0)) + track.enqueue_encoded_packet_nowait(_mk_packet(pts=1)) + track.enqueue_encoded_packet_nowait(_mk_packet(pts=2)) + # Queue should now hold pts=1 and pts=2 (p0 dropped). + first = await asyncio.wait_for(track.recv(), timeout=1.0) + second = await asyncio.wait_for(track.recv(), timeout=1.0) + assert first.pts == 1 + assert second.pts == 2 diff --git a/flashdreams/tests/test_webrtc_manager.py b/flashdreams/tests/test_webrtc_manager.py index 69bb05ce5..d94180009 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -12,6 +12,7 @@ import torch from flashdreams.serving.webrtc.controls import WSAD_SUPPORTED_KEYS +from flashdreams.serving.webrtc.encoders import ChunkDeliveryResult from flashdreams.serving.webrtc.manager import ( BaseWebRTCSessionManager, ManagedWebRTCSession, @@ -48,6 +49,34 @@ async def close(self) -> None: self.closed = True +class _FakeVideoEncoder: + """Minimal ``VideoEncoder``-shaped stub for ``ManagedWebRTCSession`` + construction and generation-worker tests.""" + + fps = 30 + backend = "fake" + prefers_codec: str | None = None + + async def deliver_chunk( + self, + chunk: Any, + track: Any, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + del force_keyframe + num_frames = await track.enqueue_chunk(chunk) + return ChunkDeliveryResult( + backend=self.backend, + num_frames=num_frames, + num_keyframes=0, + encode_ms=0.0, + ) + + def close(self) -> None: + return + + class _FakePeerConnection: def __init__(self) -> None: self.closed = False @@ -101,6 +130,7 @@ def _managed_session( managed = ManagedWebRTCSession( runtime=runtime, video_track=video_track, # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), # ty:ignore[invalid-argument-type] peer_connection=peer, resampler=_FakeResampler(), # ty:ignore[invalid-argument-type] control_channel=channel, @@ -202,7 +232,10 @@ async def generate_chunk( chunk_index=0, num_frames=1, video_chunk=torch.zeros((1, 1, 1, 3, 2, 2), dtype=torch.uint8), - stats=None, + stats={ + "runtime_input_ms": 1.25, + "wrapper_generate_ms": 2.5, + }, ) class _ExtraManager(_BaseTestManager): @@ -229,6 +262,12 @@ def _chunk_done_extra(self) -> dict[str, Any]: assert payload["model"] == "fake-model" assert payload["stream"] == "rgb" assert payload["resolution"] == {"width": 8, "height": 4} + assert payload["encoder_backend"] == "fake" + assert payload["delivery_encode_ms"] == 0.0 + assert payload["runtime_input_ms"] == 1.2 + assert payload["wrapper_generate_ms"] == 2.5 + assert payload["chunk_total_ms"] >= 0.0 + assert payload["chunk_fps"] >= 0.0 @pytest.mark.asyncio diff --git a/integrations/lingbot/tests/test_webrtc_runtime.py b/integrations/lingbot/tests/test_webrtc_runtime.py index 5fc4f32ba..971e4612b 100644 --- a/integrations/lingbot/tests/test_webrtc_runtime.py +++ b/integrations/lingbot/tests/test_webrtc_runtime.py @@ -52,6 +52,19 @@ def send(self, payload: str) -> None: self.messages.append(decoded) +class _FakeVideoEncoder: + """Minimal ``VideoEncoder``-shaped stub for ``_ManagedLingbotSession`` + construction. Enough to satisfy the dataclass field; the tests here do + not exercise ``create_track`` / ``deliver_chunk`` on it.""" + + fps = 30 + backend = "fake" + prefers_codec: str | None = None + + def close(self) -> None: + return + + def _fake_runtime_factory(config: LingbotRuntimeConfig) -> object: del config return object() @@ -932,6 +945,7 @@ async def test_heartbeat_message_refreshes_client_liveness( managed_session = session._ManagedLingbotSession( runtime=object(), video_track=_FakeCloseable(), # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), # ty:ignore[invalid-argument-type] peer_connection=_FakeCloseable(), resampler=object(), # ty:ignore[invalid-argument-type] control_channel=object(), @@ -962,6 +976,7 @@ async def test_client_liveness_timeout_closes_active_session( managed_session = session._ManagedLingbotSession( runtime=object(), video_track=video_track, # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), # ty:ignore[invalid-argument-type] peer_connection=peer_connection, resampler=object(), # ty:ignore[invalid-argument-type] last_client_message_at=asyncio.get_running_loop().time() - 1.0, @@ -993,6 +1008,7 @@ async def test_disconnect_message_closes_active_session( managed_session = session._ManagedLingbotSession( runtime=object(), video_track=video_track, # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), # ty:ignore[invalid-argument-type] peer_connection=peer_connection, resampler=object(), # ty:ignore[invalid-argument-type] control_channel=object(), diff --git a/integrations/omnidreams/ludus-renderer/examples/render_mirror_augmented.py b/integrations/omnidreams/ludus-renderer/examples/render_mirror_augmented.py index 68253f1d1..b4e9e8c3c 100644 --- a/integrations/omnidreams/ludus-renderer/examples/render_mirror_augmented.py +++ b/integrations/omnidreams/ludus-renderer/examples/render_mirror_augmented.py @@ -123,7 +123,7 @@ def main(): import subprocess import tempfile - import PyNvVideoCodec as nvc # ty:ignore[unresolved-import] + import PyNvVideoCodec as nvc output_path = os.path.join( os.path.dirname(__file__), f"../_images/mirror_augmented_bev.mp4" diff --git a/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py b/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py index 2c7118dee..292b05e1f 100644 --- a/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py +++ b/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py @@ -23,7 +23,9 @@ from __future__ import annotations +import time from dataclasses import dataclass +from typing import Any import numpy as np import torch @@ -79,6 +81,7 @@ class GenerationOutput: None # Generated video frames [B, V, T, 3, H, W] (None if skip_video_generation) ) finalization_state: dict | None = None # Finalization state from the video model + stats: dict[str, float] | None = None class OmnidreamsConditioningWrapper(nn.Module): @@ -93,6 +96,7 @@ def __init__( self, *, pipeline_config_name: str, + pipeline_config: Any | None = None, resolution_wh: tuple[int, int], seed_for_every_rollout: int | None = None, device: torch.device = torch.device("cuda:0"), @@ -103,6 +107,9 @@ def __init__( pipeline_config_name: Key into :data:`omnidreams.config.OMNIDREAMS_CONFIGS` identifying the pipeline literal to instantiate. + pipeline_config: Optional already-derived pipeline config. When + provided, it is instantiated instead of looking up + ``pipeline_config_name``. resolution_wh: Decoded video ``(width, height)``. Not encoded in the pipeline config; supplied per server deployment. seed_for_every_rollout: Optional per-rollout RNG seed override. When @@ -110,20 +117,21 @@ def __init__( device: CUDA device the pipeline is moved to. Raises: - KeyError: ``pipeline_config_name`` is not a registered Omnidreams - integration. + KeyError: ``pipeline_config`` is omitted and ``pipeline_config_name`` + is not a registered Omnidreams integration. TypeError: The selected pipeline does not use :class:`CosmosTransformerConfig` (the wrapper relies on its ``num_views`` / ``len_t`` fields to size session state). """ super().__init__() - if pipeline_config_name not in OMNIDREAMS_CONFIGS: - raise KeyError( - f"Unknown Omnidreams pipeline config {pipeline_config_name!r}. " - f"Available: {sorted(OMNIDREAMS_CONFIGS)}" - ) - pipeline_config = OMNIDREAMS_CONFIGS[pipeline_config_name] + if pipeline_config is None: + if pipeline_config_name not in OMNIDREAMS_CONFIGS: + raise KeyError( + f"Unknown Omnidreams pipeline config {pipeline_config_name!r}. " + f"Available: {sorted(OMNIDREAMS_CONFIGS)}" + ) + pipeline_config = OMNIDREAMS_CONFIGS[pipeline_config_name] transformer_cfg = pipeline_config.diffusion_model.transformer if not isinstance(transformer_cfg, CosmosTransformerConfig): @@ -396,6 +404,7 @@ def start_generation( GenerationOutput with ``condition_frames`` ``[B, V, T, 3, H, W]`` and ``rgb_frames`` ``[B, V, T, 3, H, W]`` (or None). """ + t_start = time.perf_counter() assert len(text_prompts) == 1, ( "Only one text prompt (batch size == 1) is supported for now" @@ -407,6 +416,7 @@ def start_generation( frame_timestamps_us=frame_timestamps_us, expected_length=self.initial_frame_chunk_size, ) + t_validated = time.perf_counter() # condition_frames: [B, V, T, 3, H, W] condition_frames = self._render_condition_frames( @@ -416,13 +426,21 @@ def start_generation( frame_timestamps_us, dynamic_actor_pool, ) + t_rendered = time.perf_counter() if skip_video_generation: state = OmnidreamsConditioningState( renderer=renderer, ) return GenerationOutput( - state=state, condition_frames=condition_frames, rgb_frames=None + state=state, + condition_frames=condition_frames, + rgb_frames=None, + stats={ + "wrapper_validate_ms": (t_validated - t_start) * 1000.0, + "wrapper_render_ms": (t_rendered - t_validated) * 1000.0, + "wrapper_total_ms": (time.perf_counter() - t_start) * 1000.0, + }, ) initial_rgb_frames, condition_frames = self._normalize_start_inputs( @@ -433,16 +451,20 @@ def start_generation( first_frame = self._to_model_range(initial_rgb_frames).unsqueeze(2) condition = self._to_model_range(condition_frames) + t_prepared = time.perf_counter() pipeline_cache = self.pipeline.initialize_cache( text=text, image=first_frame, view_names=camera_names ) + t_cache_ready = time.perf_counter() rgb_frames = self.pipeline.generate( autoregressive_index=0, hdmap=condition, cache=pipeline_cache, ) + t_generated = time.perf_counter() rgb_frames = self._to_uint8(rgb_frames).contiguous() + t_uint8_ready = time.perf_counter() state = OmnidreamsConditioningState( renderer=renderer, @@ -453,6 +475,15 @@ def start_generation( condition_frames=condition_frames, rgb_frames=rgb_frames, finalization_state={"autoregressive_index": 0}, + stats={ + "wrapper_validate_ms": (t_validated - t_start) * 1000.0, + "wrapper_render_ms": (t_rendered - t_validated) * 1000.0, + "wrapper_prepare_ms": (t_prepared - t_rendered) * 1000.0, + "wrapper_cache_ms": (t_cache_ready - t_prepared) * 1000.0, + "wrapper_generate_ms": (t_generated - t_cache_ready) * 1000.0, + "wrapper_to_uint8_ms": (t_uint8_ready - t_generated) * 1000.0, + "wrapper_total_ms": (t_uint8_ready - t_start) * 1000.0, + }, ) def continue_generation( @@ -480,12 +511,14 @@ def continue_generation( GenerationOutput with ``condition_frames`` ``[B, V, T, 3, H, W]`` and ``rgb_frames`` ``[B, V, T, 3, H, W]`` (or None). """ + t_start = time.perf_counter() self._validate_camera_inputs( camera_names=camera_names, camera_poses_per_view=camera_poses_per_view, frame_timestamps_us=frame_timestamps_us, expected_length=self.frame_chunk_size, ) + t_validated = time.perf_counter() renderer = state.renderer profiler = get_profiler() @@ -502,10 +535,18 @@ def continue_generation( frame_timestamps_us, dynamic_actor_pool, ) + t_rendered = time.perf_counter() if skip_video_generation: return GenerationOutput( - state=state, condition_frames=condition_frames, rgb_frames=None + state=state, + condition_frames=condition_frames, + rgb_frames=None, + stats={ + "wrapper_validate_ms": (t_validated - t_start) * 1000.0, + "wrapper_render_ms": (t_rendered - t_validated) * 1000.0, + "wrapper_total_ms": (time.perf_counter() - t_start) * 1000.0, + }, ) if state.pipeline_cache is None: @@ -518,6 +559,7 @@ def continue_generation( condition = self._to_model_range(model_cond) prev_block_idx = state.pipeline_cache.autoregressive_index block_idx = 0 if prev_block_idx is None else prev_block_idx + 1 + t_prepared = time.perf_counter() with profiler.measure( "pipeline.continue_generation", @@ -530,7 +572,9 @@ def continue_generation( hdmap=condition, cache=state.pipeline_cache, ) + t_generated = time.perf_counter() rgb_frames = self._to_uint8(rgb_frames).contiguous() + t_uint8_ready = time.perf_counter() new_state = OmnidreamsConditioningState( renderer=renderer, @@ -541,6 +585,14 @@ def continue_generation( condition_frames=condition_frames, rgb_frames=rgb_frames, finalization_state={"autoregressive_index": block_idx}, + stats={ + "wrapper_validate_ms": (t_validated - t_start) * 1000.0, + "wrapper_render_ms": (t_rendered - t_validated) * 1000.0, + "wrapper_prepare_ms": (t_prepared - t_rendered) * 1000.0, + "wrapper_generate_ms": (t_generated - t_prepared) * 1000.0, + "wrapper_to_uint8_ms": (t_uint8_ready - t_generated) * 1000.0, + "wrapper_total_ms": (t_uint8_ready - t_start) * 1000.0, + }, ) def cleanup(self, state: OmnidreamsConditioningState) -> None: diff --git a/integrations/omnidreams/omnidreams/interactive_drive/cli.py b/integrations/omnidreams/omnidreams/interactive_drive/cli.py index 2313aabe4..6b552f787 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/cli.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/cli.py @@ -20,9 +20,13 @@ RasterConfig, WorldModelProfileConfig, ) +from omnidreams.interactive_drive.cli_args import ExplicitArgTrackingArgumentParser from omnidreams.interactive_drive.log import configure_logging from omnidreams.interactive_drive.synthetic_scene import build_synthetic_scene_to_temp -from omnidreams.interactive_drive.world_model.manifest import load_world_model_manifest +from omnidreams.interactive_drive.world_model.manifest import ( + load_world_model_manifest, + resolve_world_model_manifest_path, +) from omnidreams.scenes import local_scene_archive_path from flashdreams.serving.realtime.timing import TraceSink @@ -51,28 +55,11 @@ def resolve_manifest_path(path: str | Path) -> Path: config directory, so ``--manifest example_world_model_perf.yaml`` works from a workspace root. """ - raw_path = Path(path).expanduser() - if raw_path.is_absolute(): - return raw_path - - cwd_path = raw_path.resolve() - if cwd_path.exists(): - return cwd_path - - package_path = (_PACKAGE_ROOT / raw_path).resolve() - if package_path.exists(): - return package_path - - if len(raw_path.parts) == 1: - configs_path = (_CONFIGS_ROOT / raw_path).resolve() - if configs_path.exists(): - return configs_path - - return cwd_path + return resolve_world_model_manifest_path(path) def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( + parser = ExplicitArgTrackingArgumentParser( description="Single-process flashdreams driving demo" ) parser.add_argument( diff --git a/integrations/omnidreams/omnidreams/interactive_drive/cli_args.py b/integrations/omnidreams/omnidreams/interactive_drive/cli_args.py new file mode 100644 index 000000000..cec0946e3 --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/cli_args.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +import argparse +import sys +from collections.abc import Sequence + +_EXPLICIT_ARG_DESTS_ATTR = "_explicit_arg_dests" + + +def _resolve_option_action( + parser: argparse.ArgumentParser, token: str +) -> argparse.Action | None: + """Resolve a raw argv token to the argparse action that consumes it.""" + option = token.split("=", 1)[0] + action = parser._option_string_actions.get(option) + if action is not None: + return action + + if not getattr(parser, "allow_abbrev", True): + return None + if len(option) < 2 or option[0] not in parser.prefix_chars: + return None + + matches = [] + if option[1] in parser.prefix_chars: + for option_string, candidate in parser._option_string_actions.items(): + if option_string.startswith(option): + matches.append(candidate) + else: + short_option = option[:2] + for option_string, candidate in parser._option_string_actions.items(): + if option_string == short_option or option_string.startswith(option): + matches.append(candidate) + + if len(matches) == 1: + return matches[0] + return None + + +def explicit_arg_dests( + parser: argparse.ArgumentParser, argv: Sequence[str] +) -> frozenset[str]: + """Return parser destination names whose options appear in ``argv``.""" + explicit = set() + for token in argv: + action = _resolve_option_action(parser, token) + if action is not None: + explicit.add(action.dest) + return frozenset(explicit) + + +def arg_was_explicit(args: argparse.Namespace, dest: str) -> bool: + """Return whether a parsed namespace field came from an explicit CLI flag.""" + return dest in getattr(args, _EXPLICIT_ARG_DESTS_ATTR, frozenset()) + + +class ExplicitArgTrackingArgumentParser(argparse.ArgumentParser): + """ArgumentParser that records which optional arguments users supplied.""" + + def parse_args( + self, + args: Sequence[str] | None = None, + namespace: argparse.Namespace | None = None, + ) -> argparse.Namespace: + raw_args = sys.argv[1:] if args is None else list(args) + parsed = super().parse_args(raw_args, namespace) + setattr(parsed, _EXPLICIT_ARG_DESTS_ATTR, explicit_arg_dests(self, raw_args)) + return parsed diff --git a/integrations/omnidreams/omnidreams/interactive_drive/world_model/manifest.py b/integrations/omnidreams/omnidreams/interactive_drive/world_model/manifest.py index 9f353ffa4..6c0639acc 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/world_model/manifest.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/world_model/manifest.py @@ -26,6 +26,8 @@ _NATIVE_DIT_ACCELERATION_MODES = ("auto", "disabled", "required") _NATIVE_DIT_BACKENDS = ("fp8_kvcache_cudnn", "bf16") _NATIVE_VAE_ENCODERS = ("disabled", "fp8") +_INTERACTIVE_DRIVE_ROOT = Path(__file__).resolve().parents[1] +_CONFIGS_ROOT = _INTERACTIVE_DRIVE_ROOT / "configs" def _is_hf_url(raw: str) -> bool: @@ -99,6 +101,28 @@ def _resolve_manifest_path(raw_path: str | None, *, manifest_dir: Path) -> Path return path +def resolve_world_model_manifest_path(path: str | Path) -> Path: + """Resolve a CLI manifest value against cwd and bundled configs.""" + raw_path = Path(path).expanduser() + if raw_path.is_absolute(): + return raw_path + + cwd_path = raw_path.resolve() + if cwd_path.exists(): + return cwd_path + + package_path = (_INTERACTIVE_DRIVE_ROOT / raw_path).resolve() + if package_path.exists(): + return package_path + + if len(raw_path.parts) == 1: + configs_path = (_CONFIGS_ROOT / raw_path).resolve() + if configs_path.exists(): + return configs_path + + return cwd_path + + def _parse_resolution_wh(raw: object) -> tuple[int, int]: if raw is None: return _DEFAULT_RESOLUTION_WH diff --git a/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py b/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py index db394b749..bb7bc1b4b 100644 --- a/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py +++ b/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py @@ -26,10 +26,12 @@ import subprocess import sys import threading +import time from pathlib import Path from types import ModuleType from typing import Any, Iterator +from loguru import logger from omnidreams.native.acceleration import ( NativeAccelerationConfig, NativeAvailabilityCheck, @@ -528,6 +530,7 @@ def load_extension( if _extension is not None: return _extension _extension_load_error = None + build_started_s: float | None = None try: _ensure_windows_cuda13_toolkit() @@ -554,6 +557,19 @@ def load_extension( _add_windows_cuda_dll_directories(cudnn_package_dir) with _scoped_torch_max_jobs(max_jobs), _scoped_cuda_arch_list(): + build_started_s = time.perf_counter() + logger.info( + "Building/loading OmniDreams single-view native extension " + "{}. First-time builds can take several minutes; seeing " + "compiler processes such as cc1plus/nvcc is expected. " + "build_directory={} MAX_JOBS={} TORCH_CUDA_ARCH_LIST={} " + "verbose={}", + extension_name, + extension_build_dir, + os.environ.get(_PYTORCH_MAX_JOBS_ENV, ""), + os.environ.get(_PYTORCH_CUDA_ARCH_LIST_ENV, ""), + verbose, + ) _extension = load_torch_extension( name=extension_name, sources=[str(source) for source in _extension_sources()], @@ -653,8 +669,20 @@ def load_extension( with_cuda=True, verbose=verbose, ) + logger.info( + "OmniDreams single-view native extension {} ready in {:.1f}s.", + extension_name, + time.perf_counter() - build_started_s, + ) except Exception as exc: # pragma: no cover - environment-specific build path _extension_load_error = exc + if build_started_s is not None: + logger.warning( + "OmniDreams single-view native extension build/load failed " + "after {:.1f}s: {}", + time.perf_counter() - build_started_s, + exc, + ) return None return _extension diff --git a/integrations/omnidreams/omnidreams/webrtc/server.py b/integrations/omnidreams/omnidreams/webrtc/server.py index a36fda7a9..e75d2e36d 100644 --- a/integrations/omnidreams/omnidreams/webrtc/server.py +++ b/integrations/omnidreams/omnidreams/webrtc/server.py @@ -6,14 +6,28 @@ import argparse import os from contextlib import ExitStack +from dataclasses import replace from importlib.resources import as_file, files from pathlib import Path +from typing import Any import torch import torch.distributed as dist from aiohttp import web from loguru import logger from omnidreams.config import OMNIDREAMS_CONFIGS +from omnidreams.interactive_drive.cli_args import ( + ExplicitArgTrackingArgumentParser, + arg_was_explicit, +) +from omnidreams.interactive_drive.config import WorldModelProfileConfig +from omnidreams.interactive_drive.world_model.flashdreams_adapter import ( + _build_pipeline_config, +) +from omnidreams.interactive_drive.world_model.manifest import ( + load_world_model_manifest, + resolve_world_model_manifest_path, +) from omnidreams.transformer import CosmosTransformerConfig from omnidreams.webrtc.session import ( OmnidreamsRuntimeConfig, @@ -33,8 +47,8 @@ WEB_DIR_RESOURCE = files("omnidreams.webrtc").joinpath("web") -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = ExplicitArgTrackingArgumentParser( description=( "Omnidreams WebRTC server: serves /request_session and streams " "single-view WSAD-controlled video chunks over one peer connection." @@ -58,6 +72,17 @@ def parse_args() -> argparse.Namespace: "stages the selected Hugging Face scene." ), ) + parser.add_argument( + "--manifest", + type=Path, + default=None, + help=( + "Omnidreams world-model manifest (YAML). Accepts a path or a " + "bundled config filename such as example_world_model_perf.yaml. " + "When set, WebRTC uses the same pipeline perf toggles as the " + "interactive-drive world-model path." + ), + ) parser.add_argument( "--scene-uuid", type=str, @@ -106,7 +131,30 @@ def parse_args() -> argparse.Namespace: type=str, default="camera_front_wide_120fov", ) - return parser.parse_args() + encoder_group = parser.add_mutually_exclusive_group() + encoder_group.add_argument( + "--prefer_sw_encoder", + action="store_true", + help=( + "Prefer the FFmpeg software encoder (aiortc) over the " + "hardware encoder (PyNvVideoCodec/NVENC H.264). Useful on " + "hosts where NVENC is unavailable or misbehaving, and for " + "A/B profiling against the hardware path. Without this flag " + "the encoder is auto-selected at startup: NVENC when the " + "driver reports support at the target resolution, aiortc's " + "software encoder otherwise." + ), + ) + encoder_group.add_argument( + "--require_nvenc", + action="store_true", + help=( + "Require the PyNvVideoCodec/NVENC H.264 path. Startup fails " + "instead of falling back to aiortc if the library, driver, " + "codec negotiation, or target resolution cannot use NVENC." + ), + ) + return parser.parse_args(argv) async def _close_package_resources(app: web.Application) -> None: @@ -143,20 +191,69 @@ def build_runtime_config( *, device_override: str | None = None, ) -> OmnidreamsRuntimeConfig: + manifest_path = None + manifest = None + pipeline_config = None + pipeline_config_name = args.pipeline_config_name + device = args.device + seed = args.seed + fps = args.fps + video_width = args.video_width + video_height = args.video_height + + manifest_arg = getattr(args, "manifest", None) + if manifest_arg is not None: + manifest_path = resolve_world_model_manifest_path(manifest_arg) + manifest = load_world_model_manifest(manifest_path) + pipeline_config = _build_pipeline_config( + manifest, + profile=WorldModelProfileConfig(), + ) + pipeline_config_name = str(pipeline_config.name) + if ( + arg_was_explicit(args, "pipeline_config_name") + and args.pipeline_config_name != pipeline_config_name + ): + raise ValueError( + "--manifest selects pipeline config " + f"{pipeline_config_name!r}, but --pipeline_config_name was " + f"also set to {args.pipeline_config_name!r}." + ) + + if not arg_was_explicit(args, "device"): + device = manifest.device + if not arg_was_explicit(args, "seed"): + seed = manifest.seed_for_every_rollout + if not arg_was_explicit(args, "fps"): + fps = manifest.fps + if not arg_was_explicit(args, "video_width"): + video_width = manifest.resolution_wh[0] + if not arg_was_explicit(args, "video_height"): + video_height = manifest.resolution_wh[1] + return OmnidreamsRuntimeConfig( - pipeline_config_name=args.pipeline_config_name, + pipeline_config_name=pipeline_config_name, + pipeline_config=pipeline_config, + manifest_path=manifest_path, scene_dir=args.scene_dir, scene_uuid=args.scene_uuid, scene_variant=args.scene_variant, - seed=args.seed, - device=device_override or args.device, - video_height=args.video_height, - video_width=args.video_width, - fps=args.fps, + seed=seed, + device=device_override or device, + video_height=video_height, + video_width=video_width, + fps=fps, camera_name=args.camera_name, warmup_chunks=args.warmup_chunks, warmup_timeout_s=args.warmup_timeout_s, debug_serve_hdmaps=args.debug_serve_hdmaps, + encoder_backend=( + "nvenc" + if getattr(args, "require_nvenc", False) + else "default" + if getattr(args, "prefer_sw_encoder", False) + else "auto" + ), ) @@ -210,8 +307,10 @@ def initialize_distributed( return torch_device, world_rank, world_size -def _validate_single_view_config(config_name: str) -> None: - pipeline_cfg = OMNIDREAMS_CONFIGS[config_name] +def _validate_single_view_config( + config_name: str, pipeline_config: Any | None = None +) -> None: + pipeline_cfg = pipeline_config or OMNIDREAMS_CONFIGS[config_name] transformer_cfg = pipeline_cfg.diffusion_model.transformer if not isinstance(transformer_cfg, CosmosTransformerConfig): raise TypeError("Omnidreams WebRTC requires a CosmosTransformerConfig.") @@ -225,10 +324,16 @@ def _validate_single_view_config(config_name: str) -> None: def main() -> None: configure_logging() args = parse_args() - _validate_single_view_config(args.pipeline_config_name) + runtime_config = build_runtime_config(args) + _validate_single_view_config( + runtime_config.pipeline_config_name, + runtime_config.pipeline_config, + ) - runtime_device, world_rank, _ = initialize_distributed(default_device=args.device) - runtime_config = build_runtime_config(args, device_override=str(runtime_device)) + runtime_device, world_rank, _ = initialize_distributed( + default_device=runtime_config.device + ) + runtime_config = replace(runtime_config, device=str(runtime_device)) session_manager = OmnidreamsWebRTCSessionManager(runtime_config=runtime_config) app = None if world_rank == 0: diff --git a/integrations/omnidreams/omnidreams/webrtc/session.py b/integrations/omnidreams/omnidreams/webrtc/session.py index b307c7d21..91e2adeb4 100644 --- a/integrations/omnidreams/omnidreams/webrtc/session.py +++ b/integrations/omnidreams/omnidreams/webrtc/session.py @@ -54,6 +54,11 @@ CameraPoseIntegrator, PoseSegment, ) +from flashdreams.serving.webrtc.encoders import ( + EncoderBackend, + VideoEncoder, + select_encoder, +) from flashdreams.serving.webrtc.manager import ( DEFAULT_CLIENT_LIVENESS_TIMEOUT_S, BaseWebRTCSessionManager, @@ -405,6 +410,8 @@ class OmnidreamsRuntimeConfig: pipeline_config_name: str = ( "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf" ) + pipeline_config: Any | None = None + manifest_path: Path | None = None scene_dir: Path | None = None scene_uuid: str | None = None # Weather variant slug (default/rain/snow): picks the sibling USDZ + prompt. @@ -422,6 +429,14 @@ class OmnidreamsRuntimeConfig: warmup_chunks: int = 10 warmup_timeout_s: float = 600.0 debug_serve_hdmaps: bool = False + # Video encoder selection. ``"auto"`` prefers NVENC when the driver + # reports support at the target resolution (Stage-1 probe via + # ``PyNvVideoCodec.GetEncoderCaps``) and falls back to aiortc's + # software encoder otherwise. ``"nvenc"`` fails startup if NVENC + # cannot be initialized. ``"default"`` skips the probe entirely. + encoder_backend: EncoderBackend = "auto" + encoder_bitrate_bps: int = 6_000_000 + encoder_gop: int = 30 class OmnidreamsInferenceRuntime: @@ -450,9 +465,15 @@ def __init__(self, config: OmnidreamsRuntimeConfig | None = None) -> None: self._text_prompts: list[TextPrompt] | None = None self._camera_to_rig: torch.Tensor | None = None self._initial_ego_pose: np.ndarray | None = None + self._pending_finalization_state: dict | None = None self._next_timestamp_us: int = 0 self._closed = False self._clipgt_temp_dir: tempfile.TemporaryDirectory[str] | None = None + # Selected once at initialization; the concrete backend is chosen + # by ``select_encoder`` based on ``config.encoder_backend`` and + # the driver's ``GetEncoderCaps`` response at + # ``config.video_width`` / ``config.video_height``. + self._video_encoder: VideoEncoder | None = None # Pin every blocking runtime call to one OS thread: Omnidreams' CUDA # graph capture/replay state is thread-local, so spreading calls across # workers (e.g. asyncio.to_thread) crashes capture after a few chunks. @@ -474,6 +495,15 @@ def __init__(self, config: OmnidreamsRuntimeConfig | None = None) -> None: def is_master(self) -> bool: return self.rank == self.MASTER_RANK + @property + def video_encoder(self) -> VideoEncoder: + """Return the encoder selected at :meth:`initialize` time.""" + if self._video_encoder is None: + raise OmnidreamsRuntimeError( + "Video encoder is not initialized; call runtime.initialize() first." + ) + return self._video_encoder + def wait_for_termination(self) -> None: self.rank_coordinator.worker_loop(exit_signal=WebRTCControlSignal.EXIT) @@ -607,14 +637,19 @@ def _initialize_sync(self) -> None: camera_name=cfg.camera_name, variant=cfg.scene_variant, ) - if cfg.pipeline_config_name not in OMNIDREAMS_CONFIGS: + if ( + cfg.pipeline_config is None + and cfg.pipeline_config_name not in OMNIDREAMS_CONFIGS + ): supported = ", ".join(sorted(OMNIDREAMS_CONFIGS)) raise ValueError( f"Unknown pipeline_config_name={cfg.pipeline_config_name!r}. " f"Supported: {supported}" ) - pipeline_cfg = OMNIDREAMS_CONFIGS[cfg.pipeline_config_name] + pipeline_cfg = cfg.pipeline_config or OMNIDREAMS_CONFIGS[ + cfg.pipeline_config_name + ] transformer_cfg = pipeline_cfg.diffusion_model.transformer if not isinstance(transformer_cfg, CosmosTransformerConfig): raise TypeError( @@ -696,6 +731,7 @@ def _initialize_sync(self) -> None: pipeline_t0 = time.perf_counter() self._wrapper = OmnidreamsConditioningWrapper( pipeline_config_name=cfg.pipeline_config_name, + pipeline_config=cfg.pipeline_config, resolution_wh=(cfg.video_width, cfg.video_height), seed_for_every_rollout=cfg.seed, device=self._device, @@ -720,11 +756,39 @@ def _initialize_sync(self) -> None: self._initial_ego_pose = scene_data.ego_poses[0].transformation_matrix self._next_timestamp_us = int(scene_data.ego_poses[0].timestamp) self._reset_rollout_sync() + self._initialize_video_encoder_sync() logger.info( "Omnidreams runtime initialization complete in {:.1f}s.", time.perf_counter() - init_t0, ) + def _initialize_video_encoder_sync(self) -> None: + """Select the video encoder for this runtime. + + Runs on the runtime executor thread so any GPU-side probe + (``CreateEncoder``) sees the same CUDA context the model uses. + """ + if self._video_encoder is not None: + self._video_encoder.close() + self._video_encoder = None + device = ( + self._device + if self._device is not None + else _resolve_cuda_device( + self.config.device, + ) + ) + gpu_id = device.index if device.index is not None else 0 + self._video_encoder = select_encoder( + backend=self.config.encoder_backend, + width=self.config.video_width, + height=self.config.video_height, + fps=self.config.fps, + bitrate=self.config.encoder_bitrate_bps, + gpu_id=gpu_id, + gop=self.config.encoder_gop, + ) + def _prepare_clipgt_dir(self, clipgt_dir: Path) -> Path: def _has_prefixed_parquets(path: Path) -> bool: return any(path.glob("*.calibration_estimate.parquet")) @@ -768,6 +832,7 @@ def _reset_rollout_sync(self) -> None: if self._initial_ego_pose is None or self._scene_data is None: raise OmnidreamsRuntimeError("Scene state is not initialized.") + self._finalize_pending_generation_sync() if self._state is not None and self._state.pipeline_cache is not None: del self._state.pipeline_cache self._state = None @@ -782,6 +847,7 @@ def _reset_rollout_sync(self) -> None: self._wrapper.set_rollout_seed(self.config.seed) def _close_sync(self) -> None: + self._finalize_pending_generation_sync() state = self._state wrapper = self._wrapper self._state = None @@ -793,6 +859,10 @@ def _close_sync(self) -> None: self._camera_to_rig = None self._initial_ego_pose = None + if self._video_encoder is not None: + self._video_encoder.close() + self._video_encoder = None + if state is not None and wrapper is not None: wrapper.cleanup(state) if wrapper is not None: @@ -805,12 +875,32 @@ def _close_sync(self) -> None: torch.cuda.synchronize(device=self._device) torch.cuda.empty_cache() + def _finalize_pending_generation_sync(self) -> float: + finalization_state = self._pending_finalization_state + if finalization_state is None: + return 0.0 + self._pending_finalization_state = None + if ( + self._wrapper is None + or self._state is None + or self._state.pipeline_cache is None + ): + return 0.0 + + t_start = time.perf_counter() + self._wrapper.finalize_block_generation( + self._state.pipeline_cache, + finalization_state, + ) + return (time.perf_counter() - t_start) * 1000.0 + def _generate_one_chunk_sync( self, *, segments: list[PoseSegment], frame_times: list[float], ) -> WebRTCStepResult: + t_start = time.perf_counter() if ( self._wrapper is None or self._renderer is None @@ -822,6 +912,9 @@ def _generate_one_chunk_sync( if self._device is None: raise OmnidreamsRuntimeError("Runtime device is not initialized.") + prior_finalize_ms = self._finalize_pending_generation_sync() + t_after_prior_finalize = time.perf_counter() + num_frames = self.peek_next_chunk_num_frames() if len(frame_times) != num_frames: raise OmnidreamsRuntimeError( @@ -841,10 +934,12 @@ def _generate_one_chunk_sync( ) camera_poses = torch.einsum("nij,jk->nik", ego_poses_t, self._camera_to_rig) frame_timestamps_us = self._consume_timestamps(num_frames) + t_inputs_ready = time.perf_counter() camera_names = [self.config.camera_name] camera_poses_per_view = {self.config.camera_name: camera_poses} serve_hdmaps = self.config.debug_serve_hdmaps + t_wrapper_start = time.perf_counter() if self._state is None: output = self._wrapper.start_generation( text_prompts=self._text_prompts, @@ -865,12 +960,8 @@ def _generate_one_chunk_sync( skip_video_generation=serve_hdmaps, ) self._state = output.state - - if self._state.pipeline_cache is not None: - self._wrapper.finalize_block_generation( - self._state.pipeline_cache, - output.finalization_state, - ) + t_wrapper_done = time.perf_counter() + self._pending_finalization_state = output.finalization_state if serve_hdmaps: video_chunk = output.condition_frames @@ -878,12 +969,38 @@ def _generate_one_chunk_sync( raise OmnidreamsRuntimeError("Omnidreams WebRTC received no RGB frames.") else: video_chunk = output.rgb_frames + t_video_selected = time.perf_counter() + + # Preserve the compute-stream sync barrier that the previous + # ``.cpu()`` provided implicitly. The tensor stays on-device — the + # hardware encoder reads it via DLPack (zero-copy) while the + # software encoder path performs its own D2H copy inside + # ``BufferedVideoTrack``'s worker thread. Either way, the model's + # writes must be visible before those readers run. + t_sync_start = time.perf_counter() + if self._device is not None and self._device.type == "cuda": + torch.cuda.current_stream(self._device).synchronize() + t_synced = time.perf_counter() + + stats: dict[str, float] = dict(getattr(output, "stats", None) or {}) + stats.update( + { + "runtime_input_ms": (t_inputs_ready - t_after_prior_finalize) + * 1000.0, + "runtime_wrapper_ms": (t_wrapper_done - t_wrapper_start) * 1000.0, + "runtime_finalize_ms": prior_finalize_ms, + "runtime_prior_finalize_ms": prior_finalize_ms, + "runtime_select_ms": (t_video_selected - t_wrapper_done) * 1000.0, + "runtime_sync_ms": (t_synced - t_sync_start) * 1000.0, + "runtime_total_ms": (t_synced - t_start) * 1000.0, + } + ) result = WebRTCStepResult( chunk_index=self.autoregressive_index, num_frames=int(video_chunk.shape[2]), - video_chunk=video_chunk.detach().cpu(), - stats=None, + video_chunk=video_chunk.detach(), + stats=stats, ) self.autoregressive_index += 1 return result diff --git a/integrations/omnidreams/omnidreams/webrtc/web/request_session.css b/integrations/omnidreams/omnidreams/webrtc/web/request_session.css index d5c1998d5..b135c9469 100644 --- a/integrations/omnidreams/omnidreams/webrtc/web/request_session.css +++ b/integrations/omnidreams/omnidreams/webrtc/web/request_session.css @@ -34,6 +34,10 @@ body { color: var(--text); } +body.has-video { + overflow: hidden; +} + button { font: inherit; } @@ -81,16 +85,23 @@ button { object-fit: cover; opacity: 0; background: #050607; + transform: translateZ(0); transition: opacity 220ms ease; + will-change: opacity; } body.has-video .stageVideo { opacity: 1; } +body.has-video .stage { + min-height: 100vh; + min-height: 100svh; +} + body.has-video .idleCanvas, body.has-video .stageBackdrop { - opacity: 0; + display: none; } .stageVignette { @@ -118,6 +129,7 @@ body.has-video .stageBackdrop { background: var(--panel-bg); box-shadow: var(--panel-shadow); backdrop-filter: blur(18px) saturate(1.15); + contain: layout paint; } .brandOverlay { @@ -386,31 +398,135 @@ body[data-status="generating"] .connectButton { left: 50%; bottom: clamp(20px, 4vh, 40px); display: grid; - grid-template-columns: - minmax(max-content, 0.7fr) - minmax(max-content, 1fr) - minmax(max-content, 1.15fr) - minmax(max-content, 0.7fr) - minmax(max-content, 1.45fr); - width: min(960px, calc(100vw - 36px)); + grid-template-columns: repeat(auto-fit, minmax(108px, 1fr)); + width: min(1180px, calc(100vw - 36px)); overflow: hidden; border: 1px solid rgba(255, 255, 255, 0.14); border-radius: 8px; background: rgba(18, 20, 22, 0.74); box-shadow: 0 12px 38px rgba(0, 0, 0, 0.32); backdrop-filter: blur(16px) saturate(1.12); + contain: layout paint; transform: translateX(-50%); } +body.has-video .overlayPanel, +body.has-video .metricsBar { + background: rgba(5, 6, 7, 0.82); + box-shadow: none; + backdrop-filter: none; +} + +body.has-video .brandOverlay, +body.has-video .logCard { + display: none; +} + +body.has-video .statusCard { + top: 14px; + right: 14px; + width: auto; + min-width: 0; + padding: 7px 10px; + gap: 0; + border-color: rgba(255, 255, 255, 0.10); +} + +body.has-video .statusCard .panelLabel, +body.has-video .statusCard .flowLine { + display: none; +} + +body.has-video .statusLine { + min-height: 20px; + gap: 7px; + font-size: 0.82rem; +} + +body.has-video .statusDot { + width: 7px; + height: 7px; + box-shadow: none; +} + +body.has-video .controlCard { + left: 14px; + bottom: 70px; + width: auto; + padding: 8px; + border-color: rgba(255, 255, 255, 0.10); +} + +body.has-video .controlCard h2, +body.has-video .controlRow > span { + display: none; +} + +body.has-video .controlRow { + grid-template-columns: auto; + gap: 0; + min-height: 0; +} + +body.has-video .keyCluster { + grid-auto-columns: 32px; + gap: 6px; +} + +body.has-video .controlKey { + width: 32px; + height: 32px; + border-color: rgba(255, 255, 255, 0.16); + background: rgba(5, 6, 7, 0.82); + box-shadow: none; + font-size: 0.78rem; +} + +body.has-video .controlKey.is-active { + border-color: rgba(142, 240, 28, 0.68); + background: rgba(85, 156, 18, 0.42); + box-shadow: none; +} + +body.has-video .metricsBar { + bottom: 14px; + width: min(560px, calc(100vw - 28px)); + grid-template-columns: repeat(5, minmax(0, 1fr)); + border-color: rgba(255, 255, 255, 0.10); +} + +body.has-video .metricDetail { + display: none; +} + +body.has-video .metric { + min-height: 34px; + padding: 4px 8px; + border-bottom: 0; + font-size: 0.74rem; +} + +body.has-video .metricLiveEnd { + border-right: 0; +} + +body.has-video .metric strong { + font-size: 0.82rem; +} + +body.has-video[data-status="generating"] .statusLine strong { + animation: none; +} + .metric { display: grid; - grid-template-columns: max-content max-content; + grid-template-columns: 1fr; align-items: center; justify-content: center; - gap: 10px; - min-width: max-content; - min-height: 42px; - padding: 0 16px; + gap: 3px; + min-width: 0; + min-height: 48px; + padding: 7px 12px; border-right: 1px solid rgba(255, 255, 255, 0.11); font-size: 0.88rem; text-align: center; @@ -422,11 +538,16 @@ body[data-status="generating"] .connectButton { .metric span { color: var(--muted); + overflow: hidden; + text-overflow: ellipsis; white-space: nowrap; } .metric strong { - text-align: left; + min-width: 0; + overflow: hidden; + text-align: center; + text-overflow: ellipsis; white-space: nowrap; font-weight: 760; } @@ -451,7 +572,7 @@ body[data-status="generating"] .statusLine strong { } .stage { - min-height: max(100svh, 980px); + min-height: max(100svh, 1080px); } .brandOverlay { @@ -474,11 +595,11 @@ body[data-status="generating"] .statusLine strong { } .controlCard { - bottom: 440px; + bottom: 520px; } .logCard { - bottom: 164px; + bottom: 230px; } .metricsBar { @@ -499,7 +620,7 @@ body[data-status="generating"] .statusLine strong { } .metric:last-child { - grid-column: 1 / -1; + grid-column: auto; } } @@ -522,18 +643,20 @@ body[data-status="generating"] .statusLine strong { max-height: 150px; } - .metricsBar { - grid-template-columns: 1fr; - } - .metric { - min-height: 36px; + min-height: 38px; + padding: 6px 8px; border-right: 0; border-bottom: 1px solid rgba(255, 255, 255, 0.11); } + .metric:nth-child(2n - 1) { + border-right: 1px solid rgba(255, 255, 255, 0.11); + } + .metric:last-child { grid-column: auto; border-bottom: 0; + border-right: 0; } } diff --git a/integrations/omnidreams/omnidreams/webrtc/web/request_session.html b/integrations/omnidreams/omnidreams/webrtc/web/request_session.html index 4c0222ac9..ab58af66b 100644 --- a/integrations/omnidreams/omnidreams/webrtc/web/request_session.html +++ b/integrations/omnidreams/omnidreams/webrtc/web/request_session.html @@ -63,25 +63,45 @@

Client Logs

- FPS - -- + Media FPS + --
+ Presented + -- +
+
+ Server FPS + -- +
+
+ Dropped + -- +
+
+ Jitter + -- +
+
+ Decode + -- +
+
+ Bitrate + -- +
+
Latency --
-
+
Resolution --
-
+
Step --
-
- World Model - Omnidreams -
diff --git a/integrations/omnidreams/omnidreams/webrtc/web/request_session.js b/integrations/omnidreams/omnidreams/webrtc/web/request_session.js index 06fc7ea2e..79bf532f4 100644 --- a/integrations/omnidreams/omnidreams/webrtc/web/request_session.js +++ b/integrations/omnidreams/omnidreams/webrtc/web/request_session.js @@ -8,11 +8,16 @@ const eventLog = document.getElementById("eventLog") const logState = document.getElementById("logState") const remoteVideo = document.getElementById("remoteVideo") const idleCanvas = document.getElementById("idleCanvas") -const fpsValue = document.getElementById("fpsValue") +const browserFpsValue = document.getElementById("browserFpsValue") +const presentedFpsValue = document.getElementById("presentedFpsValue") +const serverFpsValue = document.getElementById("serverFpsValue") +const dropValue = document.getElementById("dropValue") +const jitterValue = document.getElementById("jitterValue") +const decodeValue = document.getElementById("decodeValue") +const bitrateValue = document.getElementById("bitrateValue") const latencyValue = document.getElementById("latencyValue") const resolutionValue = document.getElementById("resolutionValue") const stepValue = document.getElementById("stepValue") -const modelValue = document.getElementById("modelValue") const controlButtons = Array.from(document.querySelectorAll("[data-control-key]")) const allowedKeys = new Set(["w", "a", "s", "d"]) @@ -25,26 +30,59 @@ const keyAliases = new Map([ const keySources = new Map() const heldKeyOrder = new Map() const activeKeys = new Set() -const frameTimes = [] +const videoFrameCallbackTimes = [] +const presentedFrameSamples = [] const pendingActions = [] const maxPendingActions = 32 const heartbeatIntervalMs = 2000 +const browserProfileLogIntervalMs = 5000 +const metricsRenderIntervalMs = 500 let peerConnection = null let controlChannel = null let statsTimer = null let videoMetricsTimer = null let heartbeatTimer = null +let metricsRenderTimer = null +let idleAnimationFrame = null +let previousInboundVideoStats = null +let previousBrowserProfileLogAt = 0 +let previousMetricsRenderAt = 0 let inferenceInFlight = false let connected = false let disconnecting = false let heldKeySequence = 0 const metrics = { - fps: null, + browserFps: null, + browserRtpFps: null, + receivedFps: null, + decodedFps: null, + presentedFps: null, + videoCallbackFps: null, + serverFps: null, targetFps: null, latencyMs: null, rttMs: null, + droppedFrames: null, + droppedFps: null, + dropPercent: null, + jitterMs: null, + jitterBufferMs: null, + decodeMs: null, + processingMs: null, + bitrateMbps: null, + packetsLost: null, + packetsLostPerSec: null, + freezeCount: null, + nackCount: null, + decoder: null, + presentationDelayMs: null, + serverQueueDepth: null, + serverLagMs: null, + serverDeliveryMs: null, + serverEncodeMs: null, + serverTrackDroppedPackets: null, resolution: null, step: null, model: "Omnidreams", @@ -61,6 +99,9 @@ function formatTime() { function firstFinite(...values) { for (const value of values) { + if (value === null || value === undefined || value === "") { + continue + } const number = Number(value) if (Number.isFinite(number)) { return number @@ -69,6 +110,15 @@ function firstFinite(...values) { return null } +function toFiniteNumber(value) { + return firstFinite(value) +} + +function formatFps(value) { + const number = toFiniteNumber(value) + return number === null ? "--" : number.toFixed(1) +} + function formatMs(value) { if (!Number.isFinite(value)) { return "--" @@ -79,12 +129,49 @@ function formatMs(value) { return `${Math.round(value)} ms` } -function logEvent(message, { source = "server", level = "info" } = {}) { - const consoleMessage = `[Omnidreams WebRTC][${source}] ${message}` - if (level === "error") { - console.error(consoleMessage) - } else { - console.info(consoleMessage) +function formatMbps(value) { + const number = toFiniteNumber(value) + if (number === null) { + return "--" + } + if (number >= 10) { + return `${number.toFixed(1)} Mb/s` + } + return `${number.toFixed(2)} Mb/s` +} + +function formatDropValue() { + const total = toFiniteNumber(metrics.droppedFrames) + if (total === null) { + return "--" + } + const droppedFps = toFiniteNumber(metrics.droppedFps) + if (droppedFps !== null && droppedFps > 0.05) { + return `${Math.round(total)} (${droppedFps.toFixed(1)}/s)` + } + return String(Math.round(total)) +} + +function setTextIfChanged(element, text) { + if (element.textContent !== text) { + element.textContent = text + } +} + +function logEvent( + message, + { source = "server", level = "info", dom = true, consoleOutput = true } = {} +) { + if (consoleOutput) { + const consoleMessage = `[Omnidreams WebRTC][${source}] ${message}` + if (level === "error") { + console.error(consoleMessage) + } else { + console.info(consoleMessage) + } + } + if (!dom) { + return } const entry = document.createElement("div") @@ -106,27 +193,65 @@ function logEvent(message, { source = "server", level = "info" } = {}) { } function setStatus(message, state = message.toLowerCase()) { - statusText.textContent = message - document.body.dataset.status = state - logState.textContent = state === "idle" ? "Waiting" : message + setTextIfChanged(statusText, message) + if (document.body.dataset.status !== state) { + document.body.dataset.status = state + } + setTextIfChanged(logState, state === "idle" ? "Waiting" : message) } function setFlow(message) { - flowText.textContent = message + setTextIfChanged(flowText, message) } function setVideoVisible(visible) { document.body.classList.toggle("has-video", visible) + if (visible) { + stopIdleAnimation() + } else { + startIdleAnimation() + } } -function renderMetrics() { - const fps = firstFinite(metrics.fps, metrics.targetFps) +function renderMetrics({ force = false } = {}) { + const now = performance.now() + if (!force && previousMetricsRenderAt > 0) { + const elapsed = now - previousMetricsRenderAt + if (elapsed < metricsRenderIntervalMs) { + if (metricsRenderTimer === null) { + metricsRenderTimer = window.setTimeout(() => { + metricsRenderTimer = null + renderMetrics({ force: true }) + }, metricsRenderIntervalMs - elapsed) + } + return + } + } + + previousMetricsRenderAt = now + const browserFps = firstFinite( + metrics.browserRtpFps, + metrics.decodedFps, + metrics.receivedFps, + metrics.presentedFps + ) + const presentedFps = firstFinite(metrics.presentedFps) + const serverFps = firstFinite(metrics.serverFps, metrics.targetFps) const latency = firstFinite(metrics.latencyMs, metrics.rttMs) - fpsValue.textContent = Number.isFinite(fps) ? String(Math.round(fps)) : "--" - latencyValue.textContent = formatMs(latency) - resolutionValue.textContent = metrics.resolution || "--" - stepValue.textContent = metrics.step === null ? "--" : String(metrics.step) - modelValue.textContent = metrics.model || "Omnidreams" + metrics.browserFps = browserFps + setTextIfChanged(browserFpsValue, formatFps(browserFps)) + setTextIfChanged(presentedFpsValue, formatFps(presentedFps)) + setTextIfChanged(serverFpsValue, formatFps(serverFps)) + setTextIfChanged(dropValue, formatDropValue()) + setTextIfChanged( + jitterValue, + formatMs(firstFinite(metrics.jitterBufferMs, metrics.jitterMs)) + ) + setTextIfChanged(decodeValue, formatMs(firstFinite(metrics.decodeMs, metrics.processingMs))) + setTextIfChanged(bitrateValue, formatMbps(metrics.bitrateMbps)) + setTextIfChanged(latencyValue, formatMs(latency)) + setTextIfChanged(resolutionValue, metrics.resolution || "--") + setTextIfChanged(stepValue, metrics.step === null ? "--" : String(metrics.step)) } function recordActionSent(action) { @@ -151,6 +276,15 @@ function takeObservedActionLatency(now = performance.now()) { function updateMetricsFromChunk(payload) { const observedLatencyMs = takeObservedActionLatency() metrics.targetFps = firstFinite(payload.fps, payload.target_fps, metrics.targetFps) + metrics.serverFps = firstFinite(payload.chunk_fps, metrics.serverFps) + metrics.serverQueueDepth = firstFinite(payload.queue_depth, metrics.serverQueueDepth) + metrics.serverLagMs = firstFinite(payload.lag_ms, metrics.serverLagMs) + metrics.serverDeliveryMs = firstFinite(payload.delivery_ms, payload.enqueue_ms, metrics.serverDeliveryMs) + metrics.serverEncodeMs = firstFinite(payload.delivery_encode_ms, metrics.serverEncodeMs) + metrics.serverTrackDroppedPackets = firstFinite( + payload.track_dropped_packets, + metrics.serverTrackDroppedPackets + ) metrics.latencyMs = firstFinite( payload.latency_ms, payload.control_latency_ms, @@ -178,8 +312,11 @@ function updateMetricsFromChunk(payload) { function updateMetricsFromVideo() { if (remoteVideo.videoWidth > 0 && remoteVideo.videoHeight > 0) { - metrics.resolution = `${remoteVideo.videoWidth}x${remoteVideo.videoHeight}` - renderMetrics() + const resolution = `${remoteVideo.videoWidth}x${remoteVideo.videoHeight}` + if (metrics.resolution !== resolution) { + metrics.resolution = resolution + renderMetrics({ force: true }) + } } } @@ -226,6 +363,11 @@ function drawRouteRibbon(ctx, width, height, t) { } function drawIdleScene(now) { + idleAnimationFrame = null + if (document.body.classList.contains("has-video")) { + return + } + const ctx = idleCanvas.getContext("2d") if (!ctx) { return @@ -319,23 +461,78 @@ function drawIdleScene(now) { ctx.fillStyle = `rgba(255, 255, 255, ${0.06 + Math.sin(t * 1.4) * 0.018})` ctx.fillRect(0, 0, width, height) - if (!document.body.classList.contains("has-video")) { - recordFrame(now) + idleAnimationFrame = window.requestAnimationFrame(drawIdleScene) +} + +function startIdleAnimation() { + if ( + idleAnimationFrame !== null || + document.body.classList.contains("has-video") + ) { + return } - window.requestAnimationFrame(drawIdleScene) + idleAnimationFrame = window.requestAnimationFrame(drawIdleScene) } -function recordFrame(timestamp) { +function stopIdleAnimation() { + if (idleAnimationFrame !== null) { + window.cancelAnimationFrame(idleAnimationFrame) + idleAnimationFrame = null + } +} + +function recordPresentedFrame(timestamp, metadata = {}) { const now = Number.isFinite(timestamp) ? timestamp : performance.now() - frameTimes.push(now) - while (frameTimes.length > 0 && now - frameTimes[0] > 1200) { - frameTimes.shift() + videoFrameCallbackTimes.push(now) + while ( + videoFrameCallbackTimes.length > 0 && + now - videoFrameCallbackTimes[0] > 1200 + ) { + videoFrameCallbackTimes.shift() } - if (frameTimes.length >= 2) { - const elapsed = frameTimes[frameTimes.length - 1] - frameTimes[0] - metrics.fps = elapsed > 0 ? ((frameTimes.length - 1) * 1000) / elapsed : metrics.fps - renderMetrics() + if (videoFrameCallbackTimes.length >= 2) { + const elapsed = + videoFrameCallbackTimes[videoFrameCallbackTimes.length - 1] - + videoFrameCallbackTimes[0] + metrics.videoCallbackFps = + elapsed > 0 + ? ((videoFrameCallbackTimes.length - 1) * 1000) / elapsed + : metrics.videoCallbackFps + } + + const presentedFrames = toFiniteNumber(metadata.presentedFrames) + if (presentedFrames !== null) { + presentedFrameSamples.push({ count: presentedFrames, timestamp: now }) + while ( + presentedFrameSamples.length > 0 && + now - presentedFrameSamples[0].timestamp > 1200 + ) { + presentedFrameSamples.shift() + } + if (presentedFrameSamples.length >= 2) { + const firstSample = presentedFrameSamples[0] + const lastSample = presentedFrameSamples[presentedFrameSamples.length - 1] + const elapsed = lastSample.timestamp - firstSample.timestamp + const frameDelta = lastSample.count - firstSample.count + if (elapsed > 0 && frameDelta >= 0) { + metrics.presentedFps = (frameDelta * 1000) / elapsed + } + } + } else { + metrics.presentedFps = metrics.videoCallbackFps + } + + const presentationTime = toFiniteNumber(metadata.presentationTime) + if (presentationTime !== null) { + metrics.presentationDelayMs = Math.max(0, now - presentationTime) } + + const processingDuration = toFiniteNumber(metadata.processingDuration) + if (processingDuration !== null) { + metrics.processingMs = processingDuration * 1000 + } + + renderMetrics() } function updateControlHighlights() { @@ -356,7 +553,7 @@ function actionLabel(action) { return `${action.event}${action.key ? `:${action.key}` : ""}` } -function sendControlAction(action) { +function sendControlAction(action, { log = true } = {}) { if (!connected || !controlChannel || controlChannel.readyState !== "open") { return false } @@ -371,12 +568,17 @@ function sendControlAction(action) { recordActionSent(action) setStatus("Generating", "generating") setFlow(`sent ${actionLabel(action)}, waiting=${inferenceInFlight}`) - logEvent(`control ${actionLabel(action)}`, { source: "client" }) + if (log) { + logEvent(`control ${actionLabel(action)}`, { + source: "client", + dom: !document.body.classList.contains("has-video"), + }) + } return true } -function enqueueAction(action) { - const sent = sendControlAction(action) +function enqueueAction(action, options = {}) { + const sent = sendControlAction(action, options) if (!sent) { setFlow(connected ? `not_sent ${actionLabel(action)}` : "connect session first") } @@ -387,7 +589,7 @@ function enqueueHeldKeyRepeats() { return (heldKeyOrder.get(a) || 0) - (heldKeyOrder.get(b) || 0) }) for (const key of heldKeys) { - enqueueAction({ event: "keydown", key }) + enqueueAction({ event: "keydown", key }, { log: false }) } } @@ -448,18 +650,50 @@ function handleControlMessage(rawMessage) { inferenceInFlight = false updateMetricsFromChunk(payload) const genMs = firstFinite(payload.gen_ms) + const runtimeCallMs = firstFinite(payload.runtime_call_ms) + const renderMs = firstFinite(payload.wrapper_render_ms) + const modelMs = firstFinite(payload.wrapper_generate_ms) + const syncMs = firstFinite(payload.runtime_sync_ms) + const deliveryMs = firstFinite(payload.delivery_ms, payload.enqueue_ms) + const encodeMs = firstFinite(payload.delivery_encode_ms) + const chunkFps = firstFinite(payload.chunk_fps) const lagMs = firstFinite(payload.lag_ms) const queueDepth = firstFinite(payload.queue_depth) + const trackDroppedPackets = firstFinite(payload.track_dropped_packets) const parts = [ `chunk_done index=${payload.chunk_index}`, `frames=${payload.num_frames}`, ] + if (typeof payload.encoder_backend === "string" && payload.encoder_backend) { + parts.push(`encoder=${payload.encoder_backend}`) + } if (Number.isFinite(Number(payload.enqueued_frames))) { parts.push(`enqueued=${payload.enqueued_frames}`) } + if (chunkFps !== null) { + parts.push(`chunk_fps=${Math.round(chunkFps)}`) + } if (genMs !== null) { parts.push(`gen=${Math.round(genMs)}ms`) } + if (runtimeCallMs !== null) { + parts.push(`runtime=${Math.round(runtimeCallMs)}ms`) + } + if (renderMs !== null) { + parts.push(`render=${Math.round(renderMs)}ms`) + } + if (modelMs !== null) { + parts.push(`model=${Math.round(modelMs)}ms`) + } + if (syncMs !== null) { + parts.push(`sync=${Math.round(syncMs)}ms`) + } + if (deliveryMs !== null) { + parts.push(`delivery=${Math.round(deliveryMs)}ms`) + } + if (encodeMs !== null) { + parts.push(`encode=${Math.round(encodeMs)}ms`) + } if (lagMs !== null) { parts.push(`lag=${Math.round(lagMs)}ms`) } @@ -469,9 +703,26 @@ function handleControlMessage(rawMessage) { if (queueDepth !== null) { parts.push(`queue=${queueDepth}`) } - logEvent(parts.join(", ")) + if (trackDroppedPackets !== null && trackDroppedPackets > 0) { + parts.push(`track_dropped=${trackDroppedPackets}`) + } + const chunkIndex = Number(payload.chunk_index) + const hasVideo = document.body.classList.contains("has-video") + const importantChunk = + !Number.isFinite(chunkIndex) || + chunkIndex <= 2 || + chunkIndex % 10 === 0 || + (trackDroppedPackets !== null && trackDroppedPackets > 0) + const showInPanel = !hasVideo && importantChunk + const showInConsole = importantChunk + logEvent(parts.join(", "), { + dom: showInPanel, + consoleOutput: showInConsole, + }) setStatus(activeKeys.size > 0 ? "Generating" : "Waiting", activeKeys.size > 0 ? "generating" : "waiting") - setFlow(`chunk ${payload.chunk_index} complete`) + if (showInPanel || !hasVideo) { + setFlow(`chunk ${payload.chunk_index} complete`) + } if (activeKeys.size > 0) { enqueueHeldKeyRepeats() } @@ -515,6 +766,273 @@ async function waitForIceGatheringComplete(pc) { }) } +function setLowLatencyReceiverHint(receiver) { + if (!receiver || !("playoutDelayHint" in receiver)) { + return + } + try { + receiver.playoutDelayHint = 0 + } catch (error) { + logEvent(`could not set low-latency playout hint: ${error.message}`, { + source: "client", + }) + } +} + +function applyLowLatencyReceiverHints(pc) { + for (const receiver of pc.getReceivers()) { + if (receiver.track && receiver.track.kind === "video") { + setLowLatencyReceiverHint(receiver) + } + } +} + +function preferH264VideoCodec(transceiver) { + if ( + !transceiver || + typeof transceiver.setCodecPreferences !== "function" || + !window.RTCRtpReceiver || + typeof RTCRtpReceiver.getCapabilities !== "function" + ) { + return + } + const capabilities = RTCRtpReceiver.getCapabilities("video") + const codecs = capabilities && Array.isArray(capabilities.codecs) + ? capabilities.codecs + : [] + const h264Codecs = codecs.filter((codec) => { + return String(codec.mimeType || "").toLowerCase() === "video/h264" + }) + if (h264Codecs.length === 0) { + return + } + try { + transceiver.setCodecPreferences(h264Codecs) + logEvent("preferred H.264 receive codec for NVENC compatibility", { + source: "client", + }) + } catch (error) { + logEvent(`could not prefer H.264 codec: ${error.message}`, { + source: "client", + }) + } +} + +function snapshotInboundVideoStats(report) { + const timestamp = toFiniteNumber(report.timestamp) + return { + timestamp: timestamp === null ? performance.now() : timestamp, + framesPerSecond: toFiniteNumber(report.framesPerSecond), + framesReceived: toFiniteNumber(report.framesReceived), + framesDecoded: toFiniteNumber(report.framesDecoded), + framesDropped: toFiniteNumber(report.framesDropped), + jitter: toFiniteNumber(report.jitter), + jitterBufferDelay: toFiniteNumber(report.jitterBufferDelay), + jitterBufferEmittedCount: toFiniteNumber(report.jitterBufferEmittedCount), + totalDecodeTime: toFiniteNumber(report.totalDecodeTime), + totalProcessingDelay: toFiniteNumber(report.totalProcessingDelay), + bytesReceived: toFiniteNumber(report.bytesReceived), + packetsLost: toFiniteNumber(report.packetsLost), + nackCount: toFiniteNumber(report.nackCount), + freezeCount: toFiniteNumber(report.freezeCount), + decoderImplementation: + typeof report.decoderImplementation === "string" ? report.decoderImplementation : null, + } +} + +function deltaValue(current, previous, key) { + if (!previous) { + return null + } + const currentValue = toFiniteNumber(current[key]) + const previousValue = toFiniteNumber(previous[key]) + if (currentValue === null || previousValue === null) { + return null + } + const delta = currentValue - previousValue + return delta >= 0 ? delta : null +} + +function deltaRate(current, previous, key, elapsedSeconds) { + const delta = deltaValue(current, previous, key) + if (delta === null || elapsedSeconds <= 0) { + return null + } + return delta / elapsedSeconds +} + +function deltaAverageMs(current, previous, totalKey, countKey) { + const totalDelta = deltaValue(current, previous, totalKey) + const countDelta = deltaValue(current, previous, countKey) + if (totalDelta === null || countDelta === null || countDelta <= 0) { + return null + } + return (totalDelta * 1000) / countDelta +} + +function updateInboundVideoMetrics(report) { + const current = snapshotInboundVideoStats(report) + const previous = previousInboundVideoStats + const elapsedMs = previous ? current.timestamp - previous.timestamp : null + const elapsedSeconds = elapsedMs !== null && elapsedMs > 0 ? elapsedMs / 1000 : null + + metrics.browserRtpFps = firstFinite(current.framesPerSecond, metrics.browserRtpFps) + metrics.jitterMs = current.jitter === null ? metrics.jitterMs : current.jitter * 1000 + metrics.droppedFrames = firstFinite(current.framesDropped, metrics.droppedFrames) + metrics.packetsLost = firstFinite(current.packetsLost, metrics.packetsLost) + metrics.freezeCount = firstFinite(current.freezeCount, metrics.freezeCount) + metrics.nackCount = firstFinite(current.nackCount, metrics.nackCount) + metrics.decoder = current.decoderImplementation || metrics.decoder + + if (elapsedSeconds !== null) { + const receivedFps = deltaRate(current, previous, "framesReceived", elapsedSeconds) + const decodedFps = deltaRate(current, previous, "framesDecoded", elapsedSeconds) + const droppedFps = deltaRate(current, previous, "framesDropped", elapsedSeconds) + const bytesPerSecond = deltaRate(current, previous, "bytesReceived", elapsedSeconds) + const packetsLostPerSecond = deltaRate(current, previous, "packetsLost", elapsedSeconds) + const jitterBufferMs = deltaAverageMs( + current, + previous, + "jitterBufferDelay", + "jitterBufferEmittedCount" + ) + const decodeMs = deltaAverageMs(current, previous, "totalDecodeTime", "framesDecoded") + const processingMs = deltaAverageMs( + current, + previous, + "totalProcessingDelay", + "framesDecoded" + ) + + metrics.receivedFps = firstFinite(receivedFps, metrics.receivedFps) + metrics.decodedFps = firstFinite(decodedFps, metrics.decodedFps) + metrics.droppedFps = firstFinite(droppedFps, metrics.droppedFps) + metrics.packetsLostPerSec = firstFinite(packetsLostPerSecond, metrics.packetsLostPerSec) + metrics.jitterBufferMs = firstFinite(jitterBufferMs, metrics.jitterBufferMs) + metrics.decodeMs = firstFinite(decodeMs, metrics.decodeMs) + metrics.processingMs = firstFinite(processingMs, metrics.processingMs) + if (bytesPerSecond !== null) { + metrics.bitrateMbps = (bytesPerSecond * 8) / 1_000_000 + } + + const dropped = deltaValue(current, previous, "framesDropped") + const decoded = deltaValue(current, previous, "framesDecoded") + if (dropped !== null && decoded !== null && dropped + decoded > 0) { + metrics.dropPercent = (dropped * 100) / (dropped + decoded) + } + } + + previousInboundVideoStats = current + renderMetrics() +} + +function diagnoseBrowserProfile() { + const serverFps = firstFinite(metrics.serverFps, metrics.targetFps) + const mediaFps = firstFinite( + metrics.browserFps, + metrics.browserRtpFps, + metrics.decodedFps, + metrics.receivedFps + ) + const presentedFps = firstFinite(metrics.presentedFps) + if (serverFps === null || mediaFps === null) { + return "collecting" + } + const frameBudgetMs = 1000 / Math.max(firstFinite(metrics.targetFps, 30), 1) + if (firstFinite(metrics.droppedFps, 0) > 1 || firstFinite(metrics.dropPercent, 0) > 5) { + return "browser_dropping_frames" + } + if ( + firstFinite(metrics.jitterBufferMs, 0) > frameBudgetMs * 2 || + firstFinite(metrics.jitterMs, 0) > 20 + ) { + return "network_or_jitter_buffer" + } + if (firstFinite(metrics.decodeMs, 0) > frameBudgetMs) { + return "decode_bound" + } + if (presentedFps !== null && mediaFps - presentedFps > 3) { + return "presentation_throttled" + } + if (serverFps - mediaFps <= 3) { + return "server_browser_matched" + } + if ( + metrics.receivedFps !== null && + metrics.decodedFps !== null && + metrics.receivedFps - metrics.decodedFps > 3 + ) { + return "decode_backlog" + } + if ( + firstFinite(metrics.serverQueueDepth, 0) > 2 || + firstFinite(metrics.serverLagMs, 0) > frameBudgetMs * 2 + ) { + return "server_queue_lag" + } + return "receiver_or_display_pacing" +} + +function maybeLogBrowserProfile(now = performance.now()) { + if (!connected || now - previousBrowserProfileLogAt < browserProfileLogIntervalMs) { + return + } + previousBrowserProfileLogAt = now + + const browserFps = firstFinite(metrics.browserFps, metrics.browserRtpFps) + const serverFps = firstFinite(metrics.serverFps, metrics.targetFps) + const parts = [`diagnosis=${diagnoseBrowserProfile()}`] + if (browserFps !== null) { + parts.push(`browser_fps=${browserFps.toFixed(1)}`) + } + if (serverFps !== null) { + parts.push(`server_fps=${serverFps.toFixed(1)}`) + } + if (metrics.receivedFps !== null) { + parts.push(`recv_fps=${metrics.receivedFps.toFixed(1)}`) + } + if (metrics.decodedFps !== null) { + parts.push(`decoded_fps=${metrics.decodedFps.toFixed(1)}`) + } + if (metrics.presentedFps !== null) { + parts.push(`presented_fps=${metrics.presentedFps.toFixed(1)}`) + } + if (metrics.videoCallbackFps !== null) { + parts.push(`callback_fps=${metrics.videoCallbackFps.toFixed(1)}`) + } + if (metrics.droppedFps !== null) { + parts.push(`dropped_fps=${metrics.droppedFps.toFixed(1)}`) + } + if (metrics.jitterBufferMs !== null) { + parts.push(`jitter_buffer=${Math.round(metrics.jitterBufferMs)}ms`) + } + if (metrics.decodeMs !== null) { + parts.push(`decode=${metrics.decodeMs.toFixed(1)}ms`) + } + if (metrics.bitrateMbps !== null) { + parts.push(`bitrate=${metrics.bitrateMbps.toFixed(2)}Mb/s`) + } + if (metrics.serverQueueDepth !== null) { + parts.push(`server_queue=${metrics.serverQueueDepth}`) + } + if (metrics.serverLagMs !== null) { + parts.push(`server_lag=${Math.round(metrics.serverLagMs)}ms`) + } + if ( + metrics.serverTrackDroppedPackets !== null && + metrics.serverTrackDroppedPackets > 0 + ) { + parts.push(`server_track_dropped=${metrics.serverTrackDroppedPackets}`) + } + if (metrics.decoder) { + parts.push(`decoder=${metrics.decoder}`) + } + logEvent(`browser_profile ${parts.join(", ")}`, { + source: "client", + dom: !document.body.classList.contains("has-video"), + }) +} + async function pollWebRtcStats() { if (!peerConnection) { return @@ -531,13 +1049,13 @@ async function pollWebRtcStats() { } if ( report.type === "inbound-rtp" && - (report.kind === "video" || report.mediaType === "video") && - Number.isFinite(report.framesPerSecond) + (report.kind === "video" || report.mediaType === "video") ) { - metrics.fps = report.framesPerSecond + updateInboundVideoMetrics(report) } } renderMetrics() + maybeLogBrowserProfile() } catch (error) { logEvent(`stats unavailable: ${error.message}`, { source: "client" }) } @@ -559,6 +1077,52 @@ function stopStatsPolling() { } } +function resetBrowserProfileMetrics() { + videoFrameCallbackTimes.length = 0 + presentedFrameSamples.length = 0 + previousInboundVideoStats = null + previousBrowserProfileLogAt = 0 + previousMetricsRenderAt = 0 + if (metricsRenderTimer !== null) { + window.clearTimeout(metricsRenderTimer) + metricsRenderTimer = null + } + Object.assign(metrics, { + browserFps: null, + browserRtpFps: null, + receivedFps: null, + decodedFps: null, + presentedFps: null, + videoCallbackFps: null, + serverFps: null, + targetFps: null, + latencyMs: null, + rttMs: null, + droppedFrames: null, + droppedFps: null, + dropPercent: null, + jitterMs: null, + jitterBufferMs: null, + decodeMs: null, + processingMs: null, + bitrateMbps: null, + packetsLost: null, + packetsLostPerSec: null, + freezeCount: null, + nackCount: null, + decoder: null, + presentationDelayMs: null, + serverQueueDepth: null, + serverLagMs: null, + serverDeliveryMs: null, + serverEncodeMs: null, + serverTrackDroppedPackets: null, + resolution: null, + step: null, + }) + renderMetrics({ force: true }) +} + function resetPeerHandles(pc = peerConnection, channel = controlChannel) { if (peerConnection === pc) { peerConnection = null @@ -665,6 +1229,7 @@ async function connectSession() { setStatus("Connecting", "connecting") setFlow("creating peer connection") logEvent("connecting to server...", { source: "client" }) + resetBrowserProfileMetrics() disconnecting = false try { @@ -672,7 +1237,9 @@ async function connectSession() { const channel = pc.createDataChannel("controls") peerConnection = pc controlChannel = channel - pc.addTransceiver("video", { direction: "recvonly" }) + const videoTransceiver = pc.addTransceiver("video", { direction: "recvonly" }) + preferH264VideoCodec(videoTransceiver) + setLowLatencyReceiverHint(videoTransceiver.receiver) channel.onopen = () => { connected = true @@ -705,6 +1272,8 @@ async function connectSession() { remoteVideo.srcObject = stream updateMetricsFromVideo() } + setLowLatencyReceiverHint(event.receiver) + applyLowLatencyReceiverHints(pc) setFlow("video track attached") logEvent("video track attached", { source: "client" }) } @@ -716,6 +1285,7 @@ async function connectSession() { connected = true setStatus("Waiting", "waiting") setFlow("connected; waiting for input") + applyLowLatencyReceiverHints(pc) startStatsPolling() return } @@ -836,9 +1406,9 @@ function startVideoFrameMonitor() { } return } - const onFrame = (now) => { + const onFrame = (now, metadata) => { if (document.body.classList.contains("has-video")) { - recordFrame(now) + recordPresentedFrame(now, metadata) updateMetricsFromVideo() } remoteVideo.requestVideoFrameCallback(onFrame) @@ -846,13 +1416,46 @@ function startVideoFrameMonitor() { remoteVideo.requestVideoFrameCallback(onFrame) } +function getBrowserProfileSnapshot() { + return { + ...metrics, + connected, + inferenceInFlight, + activeKeys: Array.from(activeKeys), + diagnosis: diagnoseBrowserProfile(), + } +} + +async function getPeerStatsSnapshot() { + if (!peerConnection) { + return [] + } + const stats = await peerConnection.getStats() + return Array.from(stats.values()).map((report) => ({ ...report })) +} + +function getRecentLogMessages(limit = 20) { + return Array.from(eventLog.querySelectorAll(".logEntry span")) + .slice(0, limit) + .map((entry) => entry.textContent) +} + +function exposeBrowserProfile() { + window.omnidreamsWebRTCProfile = { + snapshot: getBrowserProfileSnapshot, + peerStats: getPeerStatsSnapshot, + logs: getRecentLogMessages, + } +} + function initialize() { document.body.dataset.status = "idle" + exposeBrowserProfile() logEvent("viewer ready", { source: "client" }) setFlow("waiting") renderMetrics() attachPointerControls() - window.requestAnimationFrame(drawIdleScene) + startIdleAnimation() startVideoFrameMonitor() } diff --git a/integrations/omnidreams/pyproject.toml b/integrations/omnidreams/pyproject.toml index 5f796eb79..5b16daa97 100644 --- a/integrations/omnidreams/pyproject.toml +++ b/integrations/omnidreams/pyproject.toml @@ -52,6 +52,15 @@ dependencies = [ "tqdm>=4.67", "transformers>=5.0,<6", "huggingface_hub>=0.24", + # NVENC hardware H.264 encoding for the WebRTC path via + # ``flashdreams.serving.webrtc.encoders.PyNvHardwareEncoder``. Not + # optional: omnidreams already requires CUDA at runtime, so the + # PyNvVideoCodec install adds no extra platform surface. Hosts + # without an NVENC ASIC fall back to aiortc's software encoder + # automatically at session initialization (Stage-1 ``GetEncoderCaps`` + # probe returns unsupported, ``select_encoder`` logs the reason and + # constructs ``DefaultRTCEncoder``). + "PyNvVideoCodec>=2.1,<3", ] [tool.uv.sources] diff --git a/integrations/omnidreams/tests/interactive_drive/test_cli.py b/integrations/omnidreams/tests/interactive_drive/test_cli.py index 47b37ef05..e084eb2ee 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_cli.py +++ b/integrations/omnidreams/tests/interactive_drive/test_cli.py @@ -1,7 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import argparse + +import pytest + from omnidreams.interactive_drive.cli import build_parser +from omnidreams.interactive_drive.cli_args import ( + ExplicitArgTrackingArgumentParser, + arg_was_explicit, +) + +pytestmark = pytest.mark.ci_cpu def test_offload_text_encoder_flag_defaults_disabled() -> None: @@ -14,3 +24,51 @@ 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_parser_records_explicit_arg_destinations() -> None: + args = build_parser().parse_args( + [ + "--manifest", + "example_world_model_perf.yaml", + "--offload-text-encoder", + "--no-bev", + ] + ) + + assert arg_was_explicit(args, "manifest") + assert arg_was_explicit(args, "offload_text_encoder") + assert arg_was_explicit(args, "bev") + assert not arg_was_explicit(args, "camera") + + +def test_parser_records_abbreviated_arg_destinations() -> None: + parser = ExplicitArgTrackingArgumentParser() + parser.add_argument("--manifest") + parser.add_argument("--device") + parser.add_argument("--seed", type=int) + parser.add_argument( + "--feature-flag", + action=argparse.BooleanOptionalAction, + default=None, + ) + + args = parser.parse_args( + [ + "--man", + "example_world_model_perf.yaml", + "--dev=cuda:5", + "--see", + "123", + "--no-feat", + ] + ) + + assert args.manifest == "example_world_model_perf.yaml" + assert args.device == "cuda:5" + assert args.seed == 123 + assert args.feature_flag is False + assert arg_was_explicit(args, "manifest") + assert arg_was_explicit(args, "device") + assert arg_was_explicit(args, "seed") + assert arg_was_explicit(args, "feature_flag") diff --git a/integrations/omnidreams/tests/test_nvenc_smoke.py b/integrations/omnidreams/tests/test_nvenc_smoke.py new file mode 100644 index 000000000..b72742f94 --- /dev/null +++ b/integrations/omnidreams/tests/test_nvenc_smoke.py @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""GPU smoke test for :class:`PyNvHardwareEncoder`. + +Requires a host with NVENC-capable hardware and ``PyNvVideoCodec`` +installed. Skips cleanly on any other host so the ``ci_gpu`` job runs +elsewhere aren't confused by a hard failure. + +What we assert (kept small on purpose): + +- The capability probe reports supported. +- Encoding a modest CUDA chunk produces at least one packet. +- The first emitted packet is an Annex-B keyframe carrying SPS + PPS + + IDR NAL units. SPS/PPS presence verifies ``repeatspspps=1`` reached + the hardware — without it, receivers cannot lock on. +- Packet ``pts`` / ``time_base`` are set to the RTP 90 kHz clock, which + is what aiortc's ``H264Encoder.pack()`` expects. +""" + +from __future__ import annotations + +from fractions import Fraction + +import pytest +import torch + +pytestmark = pytest.mark.ci_gpu + +from flashdreams.serving.webrtc.encoders import ( # noqa: E402 + _PYNVVIDEOCODEC_AVAILABLE, + PyNvHardwareEncoder, + _payload_contains_nal_type, +) + +_H264_NAL_TYPE_IDR = 5 +_H264_NAL_TYPE_SPS = 7 +_H264_NAL_TYPE_PPS = 8 + + +def _has_annex_b_start_code(payload: bytes) -> bool: + return payload.startswith(b"\x00\x00\x00\x01") or payload.startswith( + b"\x00\x00\x01" + ) + + +@pytest.fixture(scope="module") +def nvenc_available() -> None: + if not _PYNVVIDEOCODEC_AVAILABLE: + pytest.skip("PyNvVideoCodec is not installed on this host") + if not torch.cuda.is_available(): + pytest.skip("CUDA is not available on this host") + supported, reason = PyNvHardwareEncoder.is_supported( + gpu_id=0, + width=512, + height=288, + ) + if not supported: + pytest.skip(f"NVENC H.264 not supported on this host: {reason}") + + +def test_probe_reports_supported(nvenc_available) -> None: + supported, reason = PyNvHardwareEncoder.is_supported( + gpu_id=0, + width=512, + height=288, + ) + assert supported, reason + + +def test_encode_chunk_produces_annex_b_keyframe_with_sps_pps( + nvenc_available, +) -> None: + encoder = PyNvHardwareEncoder( + width=512, + height=288, + fps=30, + bitrate=2_000_000, + gpu_id=0, + gop=15, + ) + try: + # 4-frame CUDA chunk in the omnidreams runtime's output layout. + chunk = torch.randint( + 0, + 255, + (4, 3, 288, 512), + dtype=torch.uint8, + device="cuda", + ) + packets: list = [] + num_frames, num_keyframes, encode_ms = encoder.encode_chunk_sync( + chunk, + force_keyframe=True, + on_packet=packets.append, + ) + assert num_frames == 4 + assert num_keyframes >= 1 + assert encode_ms > 0.0 + assert len(packets) >= 1 + + first = bytes(packets[0]) + assert _has_annex_b_start_code(first), ( + "first packet does not start with an Annex-B start code" + ) + assert _payload_contains_nal_type(first, _H264_NAL_TYPE_IDR), ( + "first packet does not carry an IDR NAL (nal_type=5) despite " + "force_keyframe=True" + ) + assert _payload_contains_nal_type(first, _H264_NAL_TYPE_SPS), ( + "SPS NAL missing from keyframe packet — repeatspspps=1 not " + "reaching the hardware; receivers will fail to lock on" + ) + assert _payload_contains_nal_type(first, _H264_NAL_TYPE_PPS), ( + "PPS NAL missing from keyframe packet — repeatspspps=1 not " + "reaching the hardware; receivers will fail to lock on" + ) + + # PTS uses the RTP 90 kHz clock so aiortc's H264Encoder.pack() + # rescales cleanly. First packet's pts must be 0. + assert packets[0].pts == 0 + assert packets[0].time_base == Fraction(1, 90_000) + # Second packet, if produced, must be exactly 90_000/fps ticks later. + if len(packets) >= 2: + assert packets[1].pts == 90_000 // 30 + finally: + encoder.close() diff --git a/integrations/omnidreams/tests/test_webrtc_runtime.py b/integrations/omnidreams/tests/test_webrtc_runtime.py index 39589ff4c..8a7404fac 100644 --- a/integrations/omnidreams/tests/test_webrtc_runtime.py +++ b/integrations/omnidreams/tests/test_webrtc_runtime.py @@ -15,6 +15,7 @@ import pytest import torch from omnidreams import scenes +from omnidreams.config import OMNIDREAMS_CONFIGS from omnidreams.webrtc import server as webrtc_server from omnidreams.webrtc import session from omnidreams.webrtc.session import ( @@ -27,7 +28,12 @@ WSAD_SUPPORTED_KEYS, CameraPoseIntegrator, ) -from flashdreams.serving.webrtc.manager import WebRTCStepResult +from flashdreams.serving.webrtc.encoders import ( + ChunkDeliveryResult, + DefaultRTCEncoder, +) +from flashdreams.serving.webrtc.manager import WebRTCStepResult, _sdp_video_codecs +from flashdreams.serving.webrtc.media import BufferedVideoTrack pytestmark = pytest.mark.ci_cpu @@ -40,6 +46,51 @@ async def close(self) -> None: self.closed = True +class _FakeVideoEncoder: + """Minimal :class:`VideoEncoder`-shaped stub for the manager tests. + + Wraps a real :class:`BufferedVideoTrack` because the manager attaches + the track to a real :class:`RTCPeerConnection` in the warmup path; + aiortc rejects anything that is not a genuine ``MediaStreamTrack``. + """ + + backend = "fake" + prefers_codec: str | None = None + + def __init__(self, *, fps: int = 30) -> None: + self.fps = fps + self.delivered_chunks: list[Any] = [] + self.closed = False + + def create_track(self, *, maxsize: int) -> BufferedVideoTrack: + return BufferedVideoTrack(fps=self.fps, maxsize=max(1, maxsize)) + + async def deliver_chunk( + self, + chunk: Any, + track: Any, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + del force_keyframe + self.delivered_chunks.append(chunk) + # If a real BufferedVideoTrack was provided, thread the chunk + # through its enqueue path so downstream consumers see frames. + if isinstance(track, BufferedVideoTrack): + enqueued = await track.enqueue_chunk(chunk) + else: + enqueued = int(chunk.shape[2]) if chunk.ndim == 6 else int(chunk.shape[0]) + return ChunkDeliveryResult( + backend=self.backend, + num_frames=enqueued, + num_keyframes=0, + encode_ms=0.1, + ) + + def close(self) -> None: + self.closed = True + + def _fake_runtime_factory(config: OmnidreamsRuntimeConfig) -> object: del config return object() @@ -163,7 +214,7 @@ def test_generate_chunk_dispatches_start_then_continue() -> None: assert wrapper.calls[0][2] == [1000, 34333] assert wrapper.calls[1][0] == "continue" assert wrapper.calls[1][1] == (3, 4, 4) - assert len(wrapper.finalized) == 2 + assert wrapper.finalized == [{"autoregressive_index": 0}] assert wrapper.skip_video_generation_flags == [False, False] @@ -397,6 +448,8 @@ 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, + prefer_sw_encoder=False, + require_nvenc=False, ) cfg = webrtc_server.build_runtime_config(args, device_override="cuda:7") @@ -408,6 +461,164 @@ 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 + # ``--prefer_sw_encoder`` unset maps to the ``auto`` backend, which + # still probes NVENC and only falls back to software when the driver + # reports it unsupported. + assert cfg.encoder_backend == "auto" + + +@pytest.mark.parametrize( + "prefer_sw_encoder, require_nvenc, expected_backend", + [(False, False, "auto"), (True, False, "default"), (False, True, "nvenc")], +) +def test_build_runtime_config_maps_prefer_sw_encoder_to_backend( + tmp_path: Path, + prefer_sw_encoder: bool, + require_nvenc: bool, + expected_backend: str, +) -> None: + """Encoder CLI switches map to the requested backend mode.""" + args = argparse.Namespace( + pipeline_config_name="omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf", + scene_dir=tmp_path / "local-scene", + scene_uuid=None, + scene_variant="default", + seed=1, + device="cuda:0", + video_height=360, + video_width=640, + fps=24, + camera_name="camera_front_wide_120fov", + warmup_chunks=0, + warmup_timeout_s=30.0, + debug_serve_hdmaps=False, + prefer_sw_encoder=prefer_sw_encoder, + require_nvenc=require_nvenc, + ) + cfg = webrtc_server.build_runtime_config(args) + assert cfg.encoder_backend == expected_backend + + +def test_parse_args_rejects_conflicting_encoder_modes() -> None: + with pytest.raises(SystemExit): + webrtc_server.parse_args(["--prefer_sw_encoder", "--require_nvenc"]) + + +def test_build_runtime_config_uses_manifest_perf_toggles() -> None: + args = webrtc_server.parse_args( + [ + "--manifest", + "example_world_model_perf.yaml", + "--warmup_chunks", + "0", + ] + ) + + cfg = webrtc_server.build_runtime_config(args) + + assert cfg.manifest_path is not None + assert cfg.manifest_path.name == "example_world_model_perf.yaml" + assert cfg.pipeline_config is not None + assert ( + cfg.pipeline_config_name + == "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf" + ) + assert cfg.video_width == 1168 + assert cfg.video_height == 640 + assert cfg.fps == 30 + assert cfg.seed is None + + transformer_cfg = cfg.pipeline_config.diffusion_model.transformer + scheduler_cfg = cfg.pipeline_config.diffusion_model.scheduler + assert transformer_cfg.skip_finalize_kv_cache is True + assert transformer_cfg.native_dit_acceleration == "required" + assert transformer_cfg.native_dit_backend == "fp8_kvcache_cudnn" + assert transformer_cfg.native_dit_attention_backend == "cudnn" + assert list(scheduler_cfg.denoising_timesteps) == [1000, 100] + assert scheduler_cfg.num_inference_steps == 2 + + +def test_build_runtime_config_manifest_allows_explicit_runtime_overrides() -> None: + args = webrtc_server.parse_args( + [ + "--manifest", + "example_world_model_perf.yaml", + "--device", + "cuda:5", + "--seed", + "123", + "--fps", + "24", + "--video_width", + "640", + "--video_height", + "352", + ] + ) + + cfg = webrtc_server.build_runtime_config(args) + + assert cfg.device == "cuda:5" + assert cfg.seed == 123 + assert cfg.fps == 24 + assert cfg.video_width == 640 + assert cfg.video_height == 352 + assert cfg.pipeline_config is not OMNIDREAMS_CONFIGS[cfg.pipeline_config_name] + + +def test_build_runtime_config_manifest_allows_abbreviated_runtime_overrides() -> None: + args = webrtc_server.parse_args( + [ + "--manif", + "example_world_model_perf.yaml", + "--dev", + "cuda:5", + "--see", + "123", + "--fp", + "24", + "--video_w", + "640", + "--video_h", + "352", + ] + ) + + cfg = webrtc_server.build_runtime_config(args) + + assert cfg.device == "cuda:5" + assert cfg.seed == 123 + assert cfg.fps == 24 + assert cfg.video_width == 640 + assert cfg.video_height == 352 + + +def test_build_runtime_config_rejects_manifest_config_name_conflict() -> None: + args = webrtc_server.parse_args( + [ + "--manifest", + "example_world_model_perf.yaml", + "--pipeline_config_name", + "omnidreams-sv-2steps-chunk3-loc6-vae-vae", + ] + ) + + with pytest.raises(ValueError, match="--manifest selects pipeline config"): + webrtc_server.build_runtime_config(args) + + +def test_build_runtime_config_rejects_manifest_abbreviated_config_name_conflict() -> None: + args = webrtc_server.parse_args( + [ + "--manif", + "example_world_model_perf.yaml", + "--pipeline", + "omnidreams-sv-2steps-chunk3-loc6-vae-vae", + ] + ) + + with pytest.raises(ValueError, match="--manifest selects pipeline config"): + webrtc_server.build_runtime_config(args) def test_parse_args_omits_scene_dir_by_default( @@ -429,6 +640,83 @@ def test_parse_args_omits_scene_dir_by_default( assert args.debug_serve_hdmaps is True +def test_runtime_initialization_passes_manifest_pipeline_config( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + manifest_args = webrtc_server.parse_args( + ["--manifest", "example_world_model_perf.yaml"] + ) + cfg = webrtc_server.build_runtime_config(manifest_args) + cfg.scene_dir = tmp_path / "scene" + clipgt_dir = cfg.scene_dir / "clipgt" + clipgt_dir.mkdir(parents=True) + first_frame_path = clipgt_dir / "first_image.png" + prompt_path = clipgt_dir / "prompt.txt" + first_frame_path.write_text("fake image", encoding="utf-8") + prompt_path.write_text("test prompt", encoding="utf-8") + captured: dict[str, object] = {} + + class _FakePose: + transformation_matrix = torch.eye(4).numpy() + timestamp = 123 + + class _FakeSceneData: + ego_poses = [_FakePose()] + camera_models = {cfg.camera_name: object()} + camera_extrinsics = {cfg.camera_name: torch.eye(4).numpy()} + + class _FakeConditioningWrapper: + initial_frame_chunk_size = 5 + frame_chunk_size = 8 + + def __init__(self, **kwargs: object) -> None: + captured.update(kwargs) + + def create_renderer(self, *_args: object) -> object: + return object() + + def set_rollout_seed(self, seed: int | None) -> None: + captured["rollout_seed"] = seed + + monkeypatch.setattr( + session, + "_extract_local_webrtc_scene_if_needed", + lambda scene_dir, **_kwargs: scene_dir, + ) + monkeypatch.setattr( + session, + "_resolve_webrtc_scene_assets", + lambda scene_dir, **_kwargs: (clipgt_dir, first_frame_path, prompt_path), + ) + monkeypatch.setattr( + session.cv2, + "imread", + lambda *_args, **_kwargs: torch.zeros((2, 2, 3), dtype=torch.uint8).numpy(), + ) + monkeypatch.setattr( + session, "load_scene", lambda *_args, **_kwargs: _FakeSceneData() + ) + monkeypatch.setattr( + session, + "load_and_attach_ludus_scene", + lambda _path, scene_data, **_kwargs: scene_data, + ) + monkeypatch.setattr( + session, + "OmnidreamsConditioningWrapper", + _FakeConditioningWrapper, + ) + runtime = OmnidreamsInferenceRuntime(config=cfg) + + runtime._initialize_sync() + + assert captured["pipeline_config_name"] == cfg.pipeline_config_name + assert captured["pipeline_config"] is cfg.pipeline_config + assert captured["resolution_wh"] == (cfg.video_width, cfg.video_height) + assert captured["seed_for_every_rollout"] is None + assert captured["rollout_seed"] is None + + def test_runtime_uses_default_scene_uuid_when_scene_is_unspecified( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: @@ -499,6 +787,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, + prefer_sw_encoder=False, ) cfg = webrtc_server.build_runtime_config(args) @@ -570,6 +859,9 @@ def __init__(self, config: OmnidreamsRuntimeConfig) -> None: self.generated_segments: list[ list[tuple[float, float, frozenset[str]]] ] = [] + # The manager reads ``runtime.video_encoder`` when it wires the + # peer connection during the warmup loopback session. + self.video_encoder = _FakeVideoEncoder(fps=config.fps) async def initialize(self) -> None: self.initialize_calls += 1 @@ -638,6 +930,7 @@ async def test_heartbeat_message_refreshes_client_liveness( managed_session = session._ManagedOmnidreamsSession( runtime=object(), video_track=_FakeCloseable(), # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), peer_connection=_FakeCloseable(), resampler=object(), # ty:ignore[invalid-argument-type] control_channel=object(), @@ -668,6 +961,7 @@ async def test_client_liveness_timeout_closes_active_session( managed_session = session._ManagedOmnidreamsSession( runtime=object(), video_track=video_track, # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), peer_connection=peer_connection, resampler=object(), # ty:ignore[invalid-argument-type] last_client_message_at=asyncio.get_running_loop().time() - 1.0, @@ -699,6 +993,7 @@ async def test_disconnect_message_closes_active_session( managed_session = session._ManagedOmnidreamsSession( runtime=object(), video_track=video_track, # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), peer_connection=peer_connection, resampler=object(), # ty:ignore[invalid-argument-type] control_channel=object(), @@ -783,6 +1078,7 @@ def send(self, message: str) -> None: managed_session = session._ManagedOmnidreamsSession( runtime=runtime, video_track=video_track, # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), peer_connection=peer_connection, resampler=_FakeResampler(), # ty:ignore[invalid-argument-type] control_channel=control_channel, @@ -803,3 +1099,237 @@ def send(self, message: str) -> None: assert video_track.closed assert peer_connection.closed assert len(control_channel.messages) == 1 + + +class _HardwareEncoderStub: + """A stand-in that ``_enforce_h264_or_fallback`` should recognize as a + hardware encoder (``prefers_codec == "h264"``) and, when H.264 fails to + negotiate, close and replace with :class:`DefaultRTCEncoder`.""" + + backend = "pynvvideocodec" + prefers_codec: str | None = "h264" + + def __init__(self, *, fps: int = 30) -> None: + self.fps = fps + self.closed = False + + def create_track(self, *, maxsize: int) -> Any: + del maxsize + return _FakeCloseable() + + async def deliver_chunk( + self, + chunk: Any, + track: Any, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + del chunk, track, force_keyframe + return ChunkDeliveryResult( + backend=self.backend, + num_frames=0, + num_keyframes=0, + encode_ms=0.0, + ) + + def close(self) -> None: + self.closed = True + + +@dataclass +class _FakeSdpCodec: + mimeType: str + + +class _FakeSender: + def __init__(self) -> None: + self.replaced_with: Any = None + + def replaceTrack(self, track: Any) -> None: + self.replaced_with = track + + +class _FakeTransceiver: + def __init__(self, negotiated: list[_FakeSdpCodec]) -> None: + self._codecs = negotiated + self.sender = _FakeSender() + + +def _sdp_fallback_managed_session( + hw_encoder: _HardwareEncoderStub, +) -> session._ManagedOmnidreamsSession: + return session._ManagedOmnidreamsSession( + runtime=object(), + video_track=_FakeCloseable(), # ty:ignore[invalid-argument-type] + video_encoder=hw_encoder, + peer_connection=_FakeCloseable(), + resampler=object(), # ty:ignore[invalid-argument-type] + ) + + +def _video_offer_sdp(*codecs: str) -> str: + payload_types = list(range(96, 96 + len(codecs))) + lines = [ + "v=0", + "o=- 1 2 IN IP4 127.0.0.1", + "s=-", + "t=0 0", + "m=video 9 UDP/TLS/RTP/SAVPF " + + " ".join(str(payload_type) for payload_type in payload_types), + "c=IN IP4 0.0.0.0", + "a=recvonly", + ] + for payload_type, codec in zip(payload_types, codecs, strict=True): + lines.append(f"a=rtpmap:{payload_type} {codec}/90000") + lines.extend( + [ + "m=application 9 UDP/DTLS/SCTP webrtc-datachannel", + "c=IN IP4 0.0.0.0", + "a=sctp-port:5000", + "", + ] + ) + return "\n".join(lines) + + +def test_sdp_video_codecs_extracts_video_mime_types() -> None: + assert _sdp_video_codecs(_video_offer_sdp("VP8", "H264")) == ( + "video/VP8", + "video/H264", + ) + + +def test_prepare_video_encoder_for_offer_keeps_hardware_when_h264_offered() -> None: + manager = OmnidreamsWebRTCSessionManager( + runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=0), + ) + hw_encoder = _HardwareEncoderStub(fps=30) + + selected = manager._prepare_video_encoder_for_offer( + video_encoder=hw_encoder, + offer_sdp=_video_offer_sdp("VP8", "H264"), + ) + + assert selected is hw_encoder + assert not hw_encoder.closed + + +def test_prepare_video_encoder_for_offer_falls_back_without_h264() -> None: + manager = OmnidreamsWebRTCSessionManager( + runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=0), + ) + hw_encoder = _HardwareEncoderStub(fps=30) + + selected = manager._prepare_video_encoder_for_offer( + video_encoder=hw_encoder, + offer_sdp=_video_offer_sdp("VP8", "VP9"), + ) + + assert hw_encoder.closed + assert isinstance(selected, DefaultRTCEncoder) + + +def test_prepare_video_encoder_raises_without_h264_when_nvenc_required() -> None: + manager = OmnidreamsWebRTCSessionManager( + runtime_config=OmnidreamsRuntimeConfig( + device="cpu", + warmup_chunks=0, + encoder_backend="nvenc", + ), + ) + hw_encoder = _HardwareEncoderStub(fps=30) + + with pytest.raises(RuntimeError, match="browser offer does not advertise H.264"): + manager._prepare_video_encoder_for_offer( + video_encoder=hw_encoder, + offer_sdp=_video_offer_sdp("VP8", "VP9"), + ) + + assert hw_encoder.closed + + +def test_enforce_h264_or_fallback_swaps_when_negotiation_lands_on_non_h264() -> None: + manager = OmnidreamsWebRTCSessionManager( + runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=0), + ) + hw_encoder = _HardwareEncoderStub(fps=30) + original_track = _FakeCloseable() + managed_session = _sdp_fallback_managed_session(hw_encoder) + managed_session.video_track = original_track # ty:ignore[invalid-assignment] + transceiver = _FakeTransceiver([_FakeSdpCodec(mimeType="video/VP8")]) + + manager._enforce_h264_or_fallback( + transceiver=transceiver, + managed_session=managed_session, + num_frames=4, + ) + + assert hw_encoder.closed, "hardware encoder should be closed on fallback" + assert isinstance(managed_session.video_encoder, DefaultRTCEncoder) + assert isinstance(managed_session.video_track, BufferedVideoTrack) + assert transceiver.sender.replaced_with is managed_session.video_track + + +def test_enforce_h264_or_fallback_raises_when_nvenc_required() -> None: + manager = OmnidreamsWebRTCSessionManager( + runtime_config=OmnidreamsRuntimeConfig( + device="cpu", + warmup_chunks=0, + encoder_backend="nvenc", + ), + ) + hw_encoder = _HardwareEncoderStub(fps=30) + managed_session = _sdp_fallback_managed_session(hw_encoder) + transceiver = _FakeTransceiver([_FakeSdpCodec(mimeType="video/VP8")]) + + with pytest.raises(RuntimeError, match="encoder_backend='nvenc' requested"): + manager._enforce_h264_or_fallback( + transceiver=transceiver, + managed_session=managed_session, + num_frames=4, + ) + + assert hw_encoder.closed + assert managed_session.video_encoder is hw_encoder + assert transceiver.sender.replaced_with is None + + +def test_enforce_h264_or_fallback_keeps_hardware_when_h264_negotiated() -> None: + manager = OmnidreamsWebRTCSessionManager( + runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=0), + ) + hw_encoder = _HardwareEncoderStub(fps=30) + original_track = _FakeCloseable() + managed_session = _sdp_fallback_managed_session(hw_encoder) + managed_session.video_track = original_track # ty:ignore[invalid-assignment] + transceiver = _FakeTransceiver([_FakeSdpCodec(mimeType="video/H264")]) + + manager._enforce_h264_or_fallback( + transceiver=transceiver, + managed_session=managed_session, + num_frames=4, + ) + + assert not hw_encoder.closed + assert managed_session.video_encoder is hw_encoder + assert managed_session.video_track is original_track + assert transceiver.sender.replaced_with is None + + +def test_enforce_h264_or_fallback_swaps_when_no_codecs_negotiated() -> None: + manager = OmnidreamsWebRTCSessionManager( + runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=0), + ) + hw_encoder = _HardwareEncoderStub(fps=30) + managed_session = _sdp_fallback_managed_session(hw_encoder) + transceiver = _FakeTransceiver([]) + + manager._enforce_h264_or_fallback( + transceiver=transceiver, + managed_session=managed_session, + num_frames=4, + ) + + assert hw_encoder.closed + assert isinstance(managed_session.video_encoder, DefaultRTCEncoder) + assert isinstance(managed_session.video_track, BufferedVideoTrack) diff --git a/uv.lock b/uv.lock index 2b0d03215..e34d45690 100644 --- a/uv.lock +++ b/uv.lock @@ -1553,6 +1553,7 @@ requires-dist = [ { name = "opencv-python-headless", specifier = ">=4.5" }, { name = "pillow", specifier = ">=10.0" }, { name = "pyarrow", specifier = ">=16.0" }, + { name = "pynvvideocodec", specifier = ">=2.1,<3" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, { name = "pytest-manual-marker", marker = "extra == 'dev'", specifier = ">=2.0" }, { name = "pyvirtualdisplay", marker = "extra == 'dev'", specifier = ">=3.0" },