diff --git a/docs/inference_runtime_serving_architecture_improvements.md b/docs/inference_runtime_serving_architecture_improvements.md new file mode 100644 index 000000000..9d71e5d6a --- /dev/null +++ b/docs/inference_runtime_serving_architecture_improvements.md @@ -0,0 +1,417 @@ +# Inference runtime and serving architecture improvements + +## Status + +Proposed. This document records the follow-up work needed to make the runtime, +output, WebRTC, and local-window architecture match the intended component +boundaries. It is an implementation checklist, not a compatibility promise. + +## Goal + +Use one model-session implementation and one generated-video result boundary +for runner CLI, WebRTC, and local-window execution: + +```mermaid +flowchart LR + INPUTS["CLI, WebRTC, and local input adapters"] --> WORKER["Model runtime worker"] + WORKER --> SESSION["Model session
pipeline, cache, AR state"] + SESSION --> STREAM["VideoOutputStream"] + STREAM --> RESULT["StepResult"] + RESULT --> MP4["MP4 collector"] + RESULT --> WEBRTC["WebRTC encoder"] + RESULT --> LOCAL["Local presenter"] +``` + +The model integration owns conditioning, pipeline/cache state, and generation. +Shared runtime code owns orchestration contracts. Output consumers own only +their transport or presentation behavior. + +## Non-goals + +- Do not change `StreamInferencePipeline.initialize_cache`, `generate`, or + `finalize`. +- Do not move Lingbot- or OmniDreams-specific conditioning into shared + `flashdreams` code. +- Do not force model output tensors to CPU before a consumer requires host + memory. +- Do not combine WebRTC encoding, MP4 writing, and local presentation into one + output class. +- Do not add a video-specific wrapper around `StepResult`. + +## Current problems + +### Parallel model runtimes + +The generic runtime API uses `InferenceRuntime` and `InferenceSession`, while +WebRTC uses a separate `WebRTCGenerationRuntime`. Lingbot and OmniDreams each +implement replay and WebRTC generation separately, and OmniDreams local-window +execution adds a third session implementation. + +This duplicates pipeline construction, cache lifecycle, AR indexing, +`generate`/`finalize`, reset behavior, and output packaging. + +### Duplicate result metadata + +The earlier design wrapped a video-specific result in `StepResult`, while both +carried equivalent step index, frame count, and metrics fields. Consumers need +one layout-aware `StepResult` boundary instead. + +### Mixed `VideoOutputStream` responsibilities + +`VideoOutputStream` currently performs post-processing, collection, statistics +collection, result construction, CUDA synchronization, MP4 conversion, and +writing. It also exposes both `process` and `make_step_result`, leaving callers +to choose between a tensor and a result object. + +### WebRTC discards the result abstraction + +The WebRTC manager receives `StepResult` but passes only +`result.video_chunk` to encoders. Encoder implementations then infer tensor +layout from rank and shape instead of consuming the declared layout. + +### Hidden WebRTC extension points + +The shared demo builder discovers undeclared adapter methods with `getattr`, +including runtime-config, session-manager, and app factories. Model-specific +manager subclasses also provide result metadata and session-reset behavior. +The effective server interface is therefore wider than the declared protocol. + +### Model behavior in the shared browser client + +The shared browser module owns peer connection, video, metrics, and data-channel +logic, but it also hardcodes driving controls and post-process REST behavior. +The model `adapter.js` contract is implicit and differs greatly between +integrations. + +### Hard-coded output launch routing + +`serving/output_targets.py` identifies integrations from runner-name prefixes +and launches different server families for Lingbot and OmniDreams. Adding an +integration or output mode requires editing shared routing code. + +## Target ownership + +### Shared runtime and infrastructure + +- Runtime/session protocols and orchestration. +- A thread-affine runtime worker for asynchronous serving. +- `StepResult` and layout-aware video conversion. +- Stateful output post-processing through `VideoOutputStream`. +- Generic output targets, WebRTC manager, encoders, and app construction. + +### Model integrations + +- Pipeline/config selection and checkpoint behavior. +- Model-specific global conditioning and per-step input mapping. +- One session core containing pipeline/cache/AR state. +- Model-specific session-input validation and optional browser routes/assets. +- Model-specific metadata placed on the generated result. + +### Output consumers + +- MP4: collect results and persist artifacts/statistics. +- WebRTC: encode and enqueue results, then report delivery metrics. +- Local window: convert results to lazy frames and present them. + +## Improvement workstreams + +### 1. Canonical step result + +Use one layout-aware `StepResult` as the direct boundary for generated video. +Do not introduce a separate video result type or nested result envelope. + +Target properties: + +- One step/chunk index. +- One frame count, derived from or validated against the declared layout. +- A required tensor layout. +- One metrics mapping. +- Optional output time window and model-specific metadata. +- No implicit CPU transfer. + +#### TODO + +- [x] Decide the final field names and update the runtime protocol. +- [x] Require `layout` on every video `StepResult`. +- [x] Derive or validate `num_frames` exactly once during construction. +- [x] Keep all step metrics in `StepResult.metrics`. +- [x] Move video output-window information onto the canonical result. +- [x] Change video sessions and output targets to pass `StepResult` directly. +- [x] Remove duplicate unwrap/type-check code from MP4 and runner output + targets. +- [x] Add CPU tests for layout validation, frame counts, metadata, and metrics. + +Acceptance criteria: + +- A generated video step crosses every model/output boundary as exactly one + layout-aware `StepResult`. +- No step index, frame count, or metrics mapping is duplicated in a second + envelope. + +### 2. Single `VideoOutputStream` operation + +Make the stream the only raw-tensor-to-generated-result stage: + +```python +result = output_stream.process( + video_chunk, + autoregressive_index=step_index, + metrics=metrics, +) +``` + +`process` should return `StepResult`. There should be no separate +`make_step_result` call. + +#### TODO + +- [x] Change `VideoOutputStream.process` to return `StepResult`. +- [x] Remove `VideoOutputStream.make_step_result`. +- [x] Keep streaming post-processing and result construction in the stream. +- [x] Move MP4 collection and writing into `Mp4VideoOutputTarget`. +- [x] Move runner statistics persistence into the runner/MP4 target. +- [x] Remove transport-specific CUDA synchronization from the stream. +- [x] Define how `finish` reports a buffered post-processor tail without + introducing a second result type. +- [x] Verify that a disabled postprocessor preserves tensor identity and device. +- [x] Verify that stateful postprocessors are reset between sessions. + +Acceptance criteria: + +- Every generated chunk makes one output-stream call. +- Post-processing occurs at most once per chunk. +- The stream does not know about WebRTC, local-window presentation, or MP4 + files. + +### 3. Result-aware WebRTC delivery + +The WebRTC manager and encoders should consume the complete generated result. + +#### TODO + +- [x] Change `VideoEncoder.deliver_chunk` to accept `StepResult`. +- [x] Pass the result directly from the session manager to the encoder. +- [x] Make software frame conversion use `result.layout`. +- [x] Make NVENC conversion use `result.layout` instead of tensor-rank + heuristics. +- [x] Move model-specific `chunk_done` fields into `result.metadata`. +- [x] Keep transport measurements such as enqueue time, queue depth, and + control latency in the WebRTC manager. +- [x] Test `tchw` and `bvtchw` delivery through both software and NVENC fakes. +- [x] Test that no host copy occurs before the software path requests one. + +Acceptance criteria: + +- The manager never unwraps `result.video_chunk` merely to cross the encoder + boundary. +- Encoder behavior is driven by the declared layout, not guessed shape. + +### 4. One model session core per integration + +Extract one synchronous model-session core for each integration. The core owns +pipeline/cache/AR state and returns `StepResult`. Input adapters prepare +the model-specific inputs for replay, WebRTC, or local use. + +#### Lingbot TODO + +- [x] Extract shared cache initialization, AR indexing, generation, finalize, + reset, and close logic from the replay and WebRTC sessions. +- [x] Reuse the core from the runner/replay path. +- [x] Map WebRTC keyboard actions and text events into the same per-step input + boundary. +- [x] Reuse the core from the WebRTC path. +- [x] Delete the duplicate Lingbot generation implementation. +- [x] Add parity tests comparing replay and live mappings for equivalent camera + inputs. + +#### OmniDreams TODO + +- [x] Extract shared pipeline/wrapper state, cache/finalization state, AR index, + post-processing, reset, and close logic. +- [x] Reuse the model-session boundary from replay and WebRTC. +- [x] Adapt interactive-drive trajectories to the same session-step input. +- [x] Carry `StepResult` to the local presentation boundary and use + `lazy_rgb_frames()` for presentation. +- [x] Preserve delayed-finalization behavior required by interactive drive. +- [x] Delete duplicate OmniDreams generation implementations after parity is + established. +- [x] Test RGB, debug-HDMap, post-process on/off, and scene-reset behavior. + +Acceptance criteria: + +- Each integration contains one implementation of cache initialization, + `generate`, `finalize`, reset, and AR-index advancement. +- Output mode changes input and presentation adapters, not model execution. + +### 5. Thread-affine runtime worker + +All asynchronous serving lifecycle calls must execute on one owned worker +thread so CUDA, Triton, and CUDA-graph state remain thread-affine. + +#### TODO + +- [x] Add a shared single-thread runtime worker under `flashdreams.runtime`. +- [x] Route runtime initialization, session creation/reset, step, and close + through that worker. +- [x] Set the CUDA device when the worker thread starts. +- [x] Keep distributed rank coordination inside model-owned operations. +- [x] Remove per-call `asyncio.to_thread` use from integration runtimes. +- [x] Make cancellation stop awaiting a call without abandoning runtime + cleanup. +- [x] Add CPU tests for call ordering, exception propagation, and shutdown. +- [x] Add a GPU regression test that runs enough chunks to exercise Triton and + CUDA-graph reuse on one thread. + +Acceptance criteria: + +- `initialize -> reset -> step* -> close` executes on the same OS thread for a + serving runtime. +- No integration independently invents its own thread-dispatch mechanism. + +### 6. Generic WebRTC session manager + +The shared manager should own only peer lifecycle, control-event timing, input +sampling, generation scheduling, encoding, and delivery. + +#### TODO + +- [x] Drive the canonical `StepRequest -> StepResult` runtime boundary instead + of a WebRTC-only `generate_chunk` method. +- [x] Use `StepRequest` metadata to determine the next input window and + frame count. +- [x] Replace model-specific reset hooks with mapped session inputs. +- [x] Replace `_model_name` with runtime/adapter identity. +- [x] Replace `_chunk_done_extra` with `StepResult.metadata`. +- [x] Replace integration-specific runtime-error tuples with shared runtime + errors. +- [x] Delete no-op manager wrappers. +- [x] Remove integration-specific manager subclasses; integration factories + configure the shared manager's control keys and generation-error policy. +- [x] Move model-specific HTTP input and preview behavior to app controllers. +- [x] Cover session negotiation, reset, reconnect, error, and warmup behavior in + shared CPU tests. + +Acceptance criteria: + +- Lingbot and OmniDreams use the same concrete manager unless a real transport + capability differs. +- The manager has no imports from integration packages. + +### 7. Explicit WebRTC app and browser adapter contracts + +Replace dynamic optional methods with explicit extension surfaces. + +#### Server TODO + +- [x] Declare a typed WebRTC demo-adapter protocol. +- [x] Replace `getattr` discovery of runtime-config, manager, and app factories. +- [x] Always construct the shared aiohttp/WebRTC app in shared code. +- [x] Let integrations provide model web resources and optional route + registration, not a complete replacement app factory. +- [x] Provide one generic session-input route that delegates parsing/validation + to the model adapter where practical. +- [x] Keep offer, health, static assets, preload, and shutdown routes shared. + +#### Browser TODO + +- [x] Document the `adapter.js` interface with a JSDoc typedef or equivalent. +- [x] Keep peer connection, video, heartbeat, metrics, and common control + rendering in the shared client. +- [x] Make control groups declarative instead of hardcoded as universal WSAD + controls. +- [x] Move model-specific session forms and control-message handling into the + model adapter. +- [x] Represent optional post-processing as an explicit capability. +- [x] Add shared adapter-contract tests for Lingbot and OmniDreams. + +Acceptance criteria: + +- The browser loads one shared client and one small model adapter. +- A model can add UI/session behavior without copying connection or playback + logic. +- The shared demo builder has no undeclared adapter calls. + +### 8. Capability-driven output launch + +Output discovery should come from registered model/demo adapters instead of +runner-name prefix checks. + +#### TODO + +- [x] Let adapters declare supported input and output modes. +- [x] Resolve `cli`, `webrtc`, and `local-window` through adapter capabilities. +- [x] Remove `_is_lingbot_runner` and `_is_omnidreams_runner` branches from + shared output routing. +- [x] Launch Lingbot and OmniDreams WebRTC through the same shared demo entry + point. +- [x] Keep local-window manifest selection inside the OmniDreams integration. +- [x] Add registry tests proving a new adapter can add an output without editing + shared routing code. + +Acceptance criteria: + +- Adding a model integration does not require a model-name branch under + `flashdreams/flashdreams`. +- All WebRTC-capable integrations use the same shared server construction. + +## Suggested pull-request sequence + +Keep each change behavior-preserving and independently testable: + +1. **Result contract:** canonicalize `StepResult` and remove duplicated + video fields from the outer result path. +2. **Output consumption:** simplify `VideoOutputStream`; make WebRTC encoders, + MP4, and local presentation consume the result directly. +3. **Runtime worker:** add thread-affine execution and migrate existing WebRTC + lifecycle calls without changing model behavior. +4. **Lingbot session:** unify replay/runner and WebRTC generation. +5. **OmniDreams session:** unify replay, WebRTC, and local-window generation. +6. **WebRTC manager:** remove model-specific manager hooks and wrappers. +7. **App/UI boundary:** formalize server and browser adapter contracts. +8. **Launch routing:** replace model-name branches with adapter capabilities. + +Do not combine the model-session migrations with the browser redesign. Keeping +those changes separate makes output parity and UI regressions easier to locate. + +## Verification checklist + +### Static and CPU checks + +- [x] `uv run --locked --group lint ty check` +- [x] `uv run --locked --group lint pre-commit run --all-files` +- [x] Runtime/result/output unit tests. +- [x] WebRTC manager, message, encoder, and server unit tests. +- [x] Lingbot and OmniDreams demo API CPU tests. +- [x] Local-window adapter and frame-conversion CPU tests. +- [x] Every new pytest test has exactly one CI marker. + +### GPU checks + +- [ ] Lingbot runner replay produces the expected chunk count and MP4. +- [ ] Lingbot WebRTC runs multiple chunks, resets, and reconnects. +- [ ] OmniDreams replay produces the expected chunk count and MP4. +- [ ] OmniDreams WebRTC runs multiple chunks with post-processing off and on. +- [ ] OmniDreams local window renders multiple chunks and resets scenes. +- [ ] Software and NVENC WebRTC delivery both work. +- [x] Compiled and CUDA-graph configurations run beyond capture/replay startup. +- [ ] Multi-GPU rank coordination still advances every AR step in order. + +### Parity checks + +- [x] Equivalent replay and live per-step inputs reach the same model session + shape and layout. +- [ ] Post-processing is applied once, with matching output across consumers. +- [ ] Frame count, step index, metrics, and metadata agree across CLI, WebRTC, + and local-window paths. +- [ ] No output consumer introduces an unexpected device transfer. + +## Definition of done + +- One model session implementation exists per integration. +- One `VideoOutputStream` call creates each generated `StepResult`. +- CLI, WebRTC, and local-window consumers accept that result directly. +- All WebRTC runtime lifecycle operations are thread-affine. +- Shared runtime and serving code contain no Lingbot/OmniDreams branches. +- The shared browser client owns connection/playback behavior; model adapters + own only model-specific UI and session behavior. +- CPU CI, lint/type checks, and targeted GPU serving tests pass. diff --git a/flashdreams/flashdreams/infra/postprocess/__init__.py b/flashdreams/flashdreams/infra/postprocess/__init__.py index c01c6ec5e..ded5fe12c 100644 --- a/flashdreams/flashdreams/infra/postprocess/__init__.py +++ b/flashdreams/flashdreams/infra/postprocess/__init__.py @@ -37,12 +37,14 @@ VideoPostprocessStepStats, VideoPostprocessStream, create_runner_postprocess_stream, + create_video_postprocess_stream, ) __all__ = [ "VideoPostprocessStream", "VideoPostprocessStepStats", "create_runner_postprocess_stream", + "create_video_postprocess_stream", "VideoChunk", "VideoPostProcessor", "VideoPostProcessorConfig", diff --git a/flashdreams/flashdreams/infra/postprocess/stream.py b/flashdreams/flashdreams/infra/postprocess/stream.py index 2dd90adf0..ca1eed49d 100644 --- a/flashdreams/flashdreams/infra/postprocess/stream.py +++ b/flashdreams/flashdreams/infra/postprocess/stream.py @@ -212,15 +212,17 @@ def _prepare(self, output: Tensor) -> None: self._prepared = True -def create_runner_postprocess_stream( - config: RunnerConfigT, +def create_video_postprocess_stream( *, + postprocess: VideoPostprocessChainConfig, + output_layout: VideoTensorLayout, + fps: float | None, + per_view: bool, world_size: int, is_rank_zero: bool = True, - fps: float | None = None, + profile: bool = False, ) -> VideoPostprocessStream | None: - """Create a runner post-processing stream, or ``None`` when skipped.""" - postprocess = getattr(config, "postprocess") + """Create a post-processing stream for one generated video rollout.""" if not postprocess.is_enabled(): return None postprocess.validate_execution(world_size=world_size) @@ -230,7 +232,27 @@ def create_runner_postprocess_stream( and not postprocess.requires_all_ranks(world_size=world_size) ): return None + return VideoPostprocessStream( + postprocess=postprocess, + output_layout=output_layout, + fps=fps, + per_view=per_view, + world_size=world_size, + profile=profile, + ) + +def create_runner_postprocess_stream( + config: RunnerConfigT, + *, + world_size: int, + is_rank_zero: bool = True, + fps: float | None = None, +) -> VideoPostprocessStream | None: + """Create a runner post-processing stream, or ``None`` when skipped.""" + postprocess = getattr(config, "postprocess") + if not postprocess.is_enabled(): + return None output_layout = getattr(config, "postprocess_output_layout") if output_layout is None: raise ValueError( @@ -242,12 +264,13 @@ def create_runner_postprocess_stream( if configured_fps is None: configured_fps = getattr(config, "fps", getattr(config, "output_fps", None)) - return VideoPostprocessStream( + return create_video_postprocess_stream( postprocess=postprocess, output_layout=output_layout, fps=configured_fps, per_view=getattr(config, "postprocess_per_view"), world_size=world_size, + is_rank_zero=is_rank_zero, profile=bool( getattr(getattr(config, "pipeline", None), "enable_sync_and_profile", False) ), diff --git a/flashdreams/flashdreams/infra/results.py b/flashdreams/flashdreams/infra/results.py new file mode 100644 index 000000000..d6d69488e --- /dev/null +++ b/flashdreams/flashdreams/infra/results.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generated inference result contracts shared by runtimes and consumers.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import TYPE_CHECKING, Any + +from torch import Tensor + +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.infra.time import TimeWindow + +if TYPE_CHECKING: + from flashdreams.infra.video_output import LazyRGBFrame + + +@dataclass(frozen=True, kw_only=True, slots=True) +class StepResult: + """Generated output and metadata returned by one inference step. + + Video results use :meth:`from_video_chunk`, which records a required tensor + layout and derives the frame count once. Non-video results may use the + regular constructor without a layout. + """ + + __hash__ = None + + step_index: int + output: Any = None + frame_count: int = 0 + layout: VideoTensorLayout | None = None + output_window: TimeWindow | None = None + metadata: Mapping[str, Any] = field(default_factory=dict) + metrics: Mapping[str, float | int] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.step_index < 0: + raise ValueError("StepResult.step_index must be >= 0.") + if self.frame_count < 0: + raise ValueError("StepResult.frame_count must be >= 0.") + if self.layout is not None: + from flashdreams.infra.video_output import infer_video_num_frames + + video_chunk = self.video_chunk + derived_frame_count = infer_video_num_frames( + video_chunk, + layout=self.layout, + ) + if self.frame_count not in (0, derived_frame_count): + raise ValueError( + "StepResult.frame_count does not match the declared video " + f"layout: expected {derived_frame_count}, got {self.frame_count}." + ) + object.__setattr__(self, "frame_count", derived_frame_count) + object.__setattr__(self, "metadata", MappingProxyType(dict(self.metadata))) + object.__setattr__(self, "metrics", MappingProxyType(dict(self.metrics))) + + @classmethod + def from_video_chunk( + cls, + *, + step_index: int, + video_chunk: Tensor, + layout: VideoTensorLayout, + output_window: TimeWindow | None = None, + metadata: Mapping[str, Any] | None = None, + metrics: Mapping[str, float | int] | None = None, + ) -> StepResult: + """Build one layout-aware generated-video result.""" + return cls( + step_index=step_index, + output=video_chunk, + layout=layout, + output_window=output_window, + metadata=dict(metadata or {}), + metrics=dict(metrics or {}), + ) + + @property + def video_chunk(self) -> Tensor: + """Return the video tensor or fail if this is not a video result.""" + if self.layout is None: + raise ValueError("StepResult.layout is required for video output.") + if not isinstance(self.output, Tensor): + raise TypeError( + "A video StepResult requires a torch.Tensor output, " + f"got {type(self.output).__name__}." + ) + return self.output + + def lazy_rgb_frames( + self, + *, + batch_index: int = 0, + view_index: int = 0, + record_cuda_event: bool = True, + ) -> list[LazyRGBFrame]: + """Expose this video result as lazy per-frame RGB handles.""" + from flashdreams.infra.video_output import lazy_rgb_frames_from_video_tensor + + return lazy_rgb_frames_from_video_tensor( + self.video_chunk, + layout=self._video_layout(), + batch_index=batch_index, + view_index=view_index, + record_cuda_event=record_cuda_event, + ) + + def video_hwc_uint8( + self, + *, + batch_index: int = 0, + view_index: int = 0, + ) -> Tensor: + """Return this video result as uint8 ``[T,H,W,C]`` on its device.""" + from flashdreams.infra.video_output import video_tensor_to_hwc_uint8 + + return video_tensor_to_hwc_uint8( + self.video_chunk, + layout=self._video_layout(), + batch_index=batch_index, + view_index=view_index, + ) + + def _video_layout(self) -> VideoTensorLayout: + if self.layout is None: + raise ValueError("StepResult.layout is required for video output.") + return self.layout + + +__all__ = ["StepResult"] diff --git a/flashdreams/flashdreams/infra/runner.py b/flashdreams/flashdreams/infra/runner.py index cf91c0960..060adbd31 100644 --- a/flashdreams/flashdreams/infra/runner.py +++ b/flashdreams/flashdreams/infra/runner.py @@ -63,6 +63,9 @@ class RunnerConfig(InstantiateConfig): per-runner ``--help`` (it's metadata, not a knob); a non-empty value is enforced for in-tree runners by the registry test.""" + output_adapter: Annotated[str | None, tyro.conf.Suppress] = None + """Optional ``module:attribute`` implementing non-CLI output capabilities.""" + pipeline: StreamInferencePipelineConfig """Wrapped pipeline config; the runner instantiates and drives it.""" @@ -183,9 +186,8 @@ def create_video_output_stream( self, *, fps: float | None = None, - move_to_cpu: bool = True, ) -> VideoOutputStream: - """Create the standard runner video output stream for one rollout.""" + """Create the standard post-processing stream for one rollout.""" layout = self.config.postprocess_output_layout if layout is None: raise ValueError( @@ -194,8 +196,6 @@ def create_video_output_stream( return VideoOutputStream( postprocess_stream=self.create_postprocess_stream(fps=fps), output_layout=layout, - collect_output=self.is_rank_zero, - move_to_cpu=move_to_cpu, ) @abstractmethod diff --git a/flashdreams/flashdreams/infra/time.py b/flashdreams/flashdreams/infra/time.py new file mode 100644 index 000000000..65629e2ec --- /dev/null +++ b/flashdreams/flashdreams/infra/time.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared time-domain value objects.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + + +@dataclass(frozen=True, kw_only=True, slots=True) +class TimeWindow: + """Half-open time window in seconds since session start.""" + + start_s: float + end_s: float + + def __post_init__(self) -> None: + if not math.isfinite(self.start_s) or not math.isfinite(self.end_s): + raise ValueError("TimeWindow bounds must be finite seconds.") + if self.start_s < 0 or self.end_s < 0: + raise ValueError("TimeWindow bounds must be non-negative.") + if self.end_s < self.start_s: + raise ValueError("TimeWindow.end_s must be >= start_s.") + + def contains(self, timestamp_s: float) -> bool: + """Return whether ``timestamp_s`` falls within this half-open window.""" + return self.start_s <= timestamp_s < self.end_s + + +__all__ = ["TimeWindow"] diff --git a/flashdreams/flashdreams/infra/video_output.py b/flashdreams/flashdreams/infra/video_output.py index 24f9ab70c..199c36b5c 100644 --- a/flashdreams/flashdreams/infra/video_output.py +++ b/flashdreams/flashdreams/infra/video_output.py @@ -17,16 +17,18 @@ from __future__ import annotations -from collections.abc import Callable, Mapping -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any +from collections.abc import Mapping +from typing import Any, Literal, TypeAlias, cast import torch from torch import Tensor from flashdreams.infra.acceleration.frame_prefetch import LazyCudaFrame from flashdreams.infra.postprocess import VideoPostprocessStream, VideoTensorLayout +from flashdreams.infra.results import StepResult +from flashdreams.infra.time import TimeWindow + +WritableVideoTensorLayout: TypeAlias = Literal["thwc", "tchw", "btchw", "bcthw"] def video_layout_time_dim(layout: VideoTensorLayout) -> int: @@ -42,6 +44,19 @@ def video_layout_time_dim(layout: VideoTensorLayout) -> int: def infer_video_num_frames(tensor: Tensor, *, layout: VideoTensorLayout) -> int: """Infer a video chunk's frame count from its declared layout.""" + expected_ndim = { + "tchw": 4, + "btchw": 5, + "bcthw": 5, + "bvtchw": 6, + }.get(layout) + if expected_ndim is None: + raise ValueError(f"unsupported video layout: {layout!r}") + if tensor.ndim != expected_ndim: + raise ValueError( + f"layout={layout!r} expects a {expected_ndim}D tensor, " + f"got shape {tuple(tensor.shape)}." + ) return int(tensor.shape[video_layout_time_dim(layout)]) @@ -133,249 +148,120 @@ def lazy_rgb_frames_from_video_tensor( ] -@dataclass(slots=True) -class VideoStepResult: - """One generated video chunk plus per-step metadata. - - The field names intentionally match the pre-existing WebRTC result shape - so serving runtimes and output helpers share layout-aware chunk metadata. - """ - - chunk_index: int - num_frames: int - video_chunk: Tensor - stats: dict[str, float] | None = None - layout: VideoTensorLayout | None = None - metadata: dict[str, Any] = field(default_factory=dict) - - @classmethod - def from_video_chunk( - cls, - *, - chunk_index: int, - video_chunk: Tensor, - layout: VideoTensorLayout, - stats: dict[str, float] | None = None, - metadata: Mapping[str, Any] | None = None, - ) -> VideoStepResult: - """Build a result and infer ``num_frames`` from ``layout``.""" - return cls( - chunk_index=chunk_index, - num_frames=infer_video_num_frames(video_chunk, layout=layout), - video_chunk=video_chunk, - stats=stats, - layout=layout, - metadata=dict(metadata or {}), - ) - - def lazy_rgb_frames( - self, - *, - batch_index: int = 0, - view_index: int = 0, - record_cuda_event: bool = True, - ) -> list[LazyRGBFrame]: - """Expose this chunk as lazy per-frame RGB handles.""" - if self.layout is None: - raise ValueError("VideoStepResult.layout is required for frame extraction") - return lazy_rgb_frames_from_video_tensor( - self.video_chunk, - layout=self.layout, - batch_index=batch_index, - view_index=view_index, - record_cuda_event=record_cuda_event, - ) - - def video_hwc_uint8( - self, - *, - batch_index: int = 0, - view_index: int = 0, - ) -> Tensor: - """Return this chunk as a uint8 ``[T,H,W,C]`` tensor on its source device.""" - if self.layout is None: - raise ValueError("VideoStepResult.layout is required for frame extraction") - return video_tensor_to_hwc_uint8( - self.video_chunk, - layout=self.layout, - batch_index=batch_index, - view_index=view_index, - ) - - class VideoOutputStream: - """Post-process and optionally collect generated video tensors. - - Runner CLI, realtime serving, and local presentation all use this same - tensor-in/tensor-out boundary. Transport-specific result envelopes and - frame conversions happen after :meth:`process`. - """ + """Turn generated tensors into post-processed step results.""" def __init__( self, *, postprocess_stream: VideoPostprocessStream | None, output_layout: VideoTensorLayout, - collect_output: bool = True, - move_to_cpu: bool = True, - empty_message: str = "runner emitted no video frames", ) -> None: self.postprocess_stream = postprocess_stream self.output_layout = output_layout - self._time_dim = video_layout_time_dim(output_layout) - self._collect_output = collect_output - self.move_to_cpu = move_to_cpu - self.empty_message = empty_message - self._chunks: list[Tensor] = [] self._closed = False - self.stats_history: list[dict[str, object]] = [] - - @property - def collect_output(self) -> bool: - """Return whether this stream collects chunks for rank-zero writing.""" - return self._collect_output + self._last_step_index: int | None = None def process( self, video_chunk: Tensor, *, autoregressive_index: int, - stats: dict[str, float] | None = None, - stats_extra: Mapping[str, object] | None = None, - ) -> Tensor: - """Process one chunk, optionally collect it, and return emitted frames.""" + metrics: Mapping[str, float | int] | None = None, + metadata: Mapping[str, Any] | None = None, + output_window: TimeWindow | None = None, + ) -> StepResult: + """Post-process one generated chunk into the shared result boundary.""" if self._closed: raise RuntimeError("cannot process video after finish()") processed = video_chunk + result_metadata = dict(metadata or {}) if self.postprocess_stream is not None: processed = self.postprocess_stream.process( video_chunk, autoregressive_index=autoregressive_index, ) - self._append_if_nonempty(processed) - if self.collect_output and stats is not None: - if self.postprocess_stream is None: - combined_stats: dict[str, object] = dict(stats) - else: - combined_stats = self.postprocess_stream.add_process_stats(stats) - entry: dict[str, object] = { - "autoregressive_index": autoregressive_index, - **combined_stats, - } - if stats_extra is not None: - entry.update(stats_extra) - self.stats_history.append(entry) - return processed - - def finish(self) -> Tensor | None: - """Flush post-processing and return the collected rank-zero video.""" - if self._closed: - return None - self._closed = True - if self.postprocess_stream is not None: - flushed = self.postprocess_stream.finish() - if flushed is not None: - self._append_if_nonempty(flushed) - return self._collected_output() - - def make_step_result( - self, - video_chunk: Tensor, - *, - autoregressive_index: int, - stats: dict[str, float] | None = None, - metadata: Mapping[str, Any] | None = None, - sync_device: torch.device | str | None = None, - ) -> VideoStepResult: - """Process a chunk and package the emitted frames for a live consumer. - - ``sync_device`` is useful for consumers such as WebRTC that hand a - GPU-resident result to another subsystem immediately after generation. - It synchronizes only when it names a CUDA device and never moves the - emitted tensor to the host. - """ - processed = self.process( - video_chunk, - autoregressive_index=autoregressive_index, - stats=stats, - ) - if sync_device is not None: - device = torch.device(sync_device) - if device.type == "cuda": - torch.cuda.current_stream(device).synchronize() - return VideoStepResult.from_video_chunk( - chunk_index=autoregressive_index, + postprocess_stats = self.postprocess_stream.last_process_stats + if postprocess_stats is not None: + result_metadata["postprocess"] = postprocess_stats.as_dict() + self._last_step_index = autoregressive_index + return StepResult.from_video_chunk( + step_index=autoregressive_index, video_chunk=processed.detach(), layout=self.output_layout, - stats=stats, - metadata=metadata, + output_window=output_window, + metrics=metrics, + metadata=result_metadata, ) - def finish_to_mp4( - self, - output_path: str | Path, - *, - fps: int | float, - writer: Callable[..., Path] | None = None, - install_hint: str | None = None, - ) -> Path | None: - """Finish this collecting stream and write its frames as one MP4. - - The stream converts its declared output layout to the runner I/O - layout, including tiling ``bvtchw`` views horizontally. - """ - video = self.finish() - if video is None: + def finish(self) -> StepResult | None: + """Close the stream and return a post-processing tail, when present.""" + if self._closed: return None - return self.write_mp4( - video, - output_path, - fps=fps, + self._closed = True + if self.postprocess_stream is None: + return None + flushed = self.postprocess_stream.finish() + if flushed is None: + return None + if self._last_step_index is None: + raise RuntimeError("post-processing emitted a tail before any video step") + return StepResult.from_video_chunk( + step_index=self._last_step_index, + video_chunk=flushed.detach(), layout=self.output_layout, - writer=writer, - install_hint=install_hint, + metadata={"postprocess_tail": True}, ) - def write_mp4( + +class VideoResultCollector: + """Collect video results for persistence or composed presentation.""" + + def __init__( self, - video: Tensor, - output_path: str | Path, *, - fps: int | float, - layout: VideoTensorLayout | str | None = None, - writer: Callable[..., Path] | None = None, - install_hint: str | None = None, - ) -> Path: - """Write video frames as MP4 using this stream's runner output path. - - ``layout`` defaults to :attr:`output_layout`; callers that compose a - presentation canvas can pass the runner-I/O ``thwc`` layout directly. - """ - from flashdreams.infra.runner_io import ( - DEFAULT_RUNNER_INSTALL_HINT, - write_video_tensor, - ) - - writable_video, writable_layout = prepare_video_for_mp4( - video, layout=layout or self.output_layout - ) - output_writer = writer or write_video_tensor - path = output_writer( - writable_video, - output_path, - fps=fps, - layout=writable_layout, - install_hint=install_hint or DEFAULT_RUNNER_INSTALL_HINT, - ) - return path + output_layout: VideoTensorLayout, + enabled: bool = True, + move_to_cpu: bool = True, + empty_message: str = "runner emitted no video frames", + ) -> None: + self.output_layout = output_layout + self.enabled = enabled + self.move_to_cpu = move_to_cpu + self.empty_message = empty_message + self._time_dim = video_layout_time_dim(output_layout) + self._chunks: list[Tensor] = [] + self.stats_history: list[dict[str, object]] = [] - def _append_if_nonempty(self, output: Tensor) -> None: - if not self.collect_output or output.shape[self._time_dim] == 0: + def add(self, result: StepResult) -> None: + """Collect one video result and its serializable statistics.""" + if result.layout != self.output_layout: + raise ValueError( + f"collector expected layout {self.output_layout!r}, " + f"got {result.layout!r}." + ) + if not self.enabled: return - self._chunks.append(output.cpu() if self.move_to_cpu else output) + if result.frame_count > 0: + chunk = result.video_chunk + self._chunks.append(chunk.cpu() if self.move_to_cpu else chunk) + entry: dict[str, object] = { + "step_index": result.step_index, + "frames": result.frame_count, + **result.metrics, + } + if result.output_window is not None: + entry["output_start_s"] = result.output_window.start_s + entry["output_end_s"] = result.output_window.end_s + if "postprocess" in result.metadata: + entry["postprocess"] = result.metadata["postprocess"] + if result.metadata.get("postprocess_tail"): + entry["postprocess_tail"] = True + self.stats_history.append(entry) - def _collected_output(self) -> Tensor | None: - if not self.collect_output: + def finish(self) -> Tensor | None: + """Concatenate and return all collected video chunks.""" + if not self.enabled: return None if not self._chunks: raise ValueError(self.empty_message) @@ -390,10 +276,10 @@ def prepare_video_for_mp4( video: Tensor, *, layout: VideoTensorLayout | str, -) -> tuple[Tensor, str]: +) -> tuple[Tensor, WritableVideoTensorLayout]: """Convert a stream output into a layout accepted by runner MP4 I/O.""" if layout in {"thwc", "tchw", "btchw", "bcthw"}: - return video, layout + return video, cast(WritableVideoTensorLayout, layout) if layout == "bvtchw": if video.ndim != 6: raise ValueError( @@ -419,7 +305,7 @@ def prepare_video_for_mp4( __all__ = [ "LazyRGBFrame", "VideoOutputStream", - "VideoStepResult", + "VideoResultCollector", "infer_video_num_frames", "lazy_rgb_frames_from_video_tensor", "prepare_video_for_mp4", diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py index 5f89196ba..fb6eb4b05 100644 --- a/flashdreams/flashdreams/runtime/__init__.py +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -59,6 +59,7 @@ from flashdreams.runtime.runner import run_inference_session from flashdreams.runtime.types import StepRequest, StepResult from flashdreams.runtime.video_output import Mp4VideoOutputTarget +from flashdreams.runtime.worker import ThreadAffineRuntimeWorker __all__ = [ "CanonicalInputs", @@ -101,6 +102,7 @@ "StepRequest", "StepResult", "TimeWindow", + "ThreadAffineRuntimeWorker", "run_inference_session", "undeclared_inference_inputs", "UserInputCapability", diff --git a/flashdreams/flashdreams/runtime/demo/__init__.py b/flashdreams/flashdreams/runtime/demo/__init__.py index 3d9d99919..7a0535556 100644 --- a/flashdreams/flashdreams/runtime/demo/__init__.py +++ b/flashdreams/flashdreams/runtime/demo/__init__.py @@ -3,7 +3,6 @@ """Experimental shared demo API above the inference runtime API.""" -from flashdreams.runtime.demo.app import run_flashdreams_demo, serve_flashdreams_demo from flashdreams.runtime.demo.outputs import build_output_target from flashdreams.runtime.demo.replay import run_replay_demo from flashdreams.runtime.demo.spec import ( @@ -13,6 +12,7 @@ NullOutputSpec, OutputSpec, PreparedScenario, + WebRTCAppResources, WebRTCOutputSpec, ) @@ -23,9 +23,8 @@ "NullOutputSpec", "OutputSpec", "PreparedScenario", + "WebRTCAppResources", "WebRTCOutputSpec", "build_output_target", - "run_flashdreams_demo", "run_replay_demo", - "serve_flashdreams_demo", ] diff --git a/flashdreams/flashdreams/runtime/demo/app.py b/flashdreams/flashdreams/runtime/demo/app.py index 7659859b4..b84e59163 100644 --- a/flashdreams/flashdreams/runtime/demo/app.py +++ b/flashdreams/flashdreams/runtime/demo/app.py @@ -1,36 +1,71 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Experimental shared demo entrypoints.""" +"""Shared command lifecycle for model demo applications.""" from __future__ import annotations +import argparse +from abc import ABC, abstractmethod from typing import Any -from .replay import run_replay_demo -from .spec import DemoAdapter, DemoSpec +import torch +import torch.distributed as dist +from flashdreams.core.distributed import init as distributed_init +from flashdreams.runtime.demo.replay import run_replay_demo +from flashdreams.runtime.demo.spec import DemoAdapter, DemoSpec +from flashdreams.serving.webrtc.bootstrap import ( + configure_logging, + initialize_cuda_distributed, +) -def run_flashdreams_demo( - *, - spec: DemoSpec, - adapter: DemoAdapter, - **kwargs: Any, -) -> object: - """Run a synchronous replay demo through the shared runtime runner.""" - return run_replay_demo(spec=spec, adapter=adapter, **kwargs) +class DemoApplication(ABC): + """Base command application shared by model replay and WebRTC demos.""" -def serve_flashdreams_demo( - *, - spec: DemoSpec, - adapter: DemoAdapter, - **kwargs: Any, -) -> object: - """Serve a WebRTC demo through the shared serving manager.""" - from .webrtc import serve_webrtc_demo + def main(self, argv: list[str] | None = None) -> None: + """Parse arguments and dispatch the selected demo mode.""" + configure_logging() + args = self.parse_args(argv) + if args.command == "replay": + run_replay_demo( + spec=self.replay_spec(args), + adapter=self.replay_adapter(), + ) + return + if args.command == "webrtc": + context = initialize_cuda_distributed( + default_device=args.device, + distributed_init_fn=distributed_init, + configure_logging_fn=configure_logging, + torch_module=torch, + dist_module=dist, + ) + self.prepare_webrtc(args, context=context) + self.serve_webrtc(args, context=context) + return + raise AssertionError(f"Unhandled command: {args.command}") - return serve_webrtc_demo(spec=spec, adapter=adapter, **kwargs) + @abstractmethod + def parse_args(self, argv: list[str] | None = None) -> argparse.Namespace: + """Parse this model's command-line arguments.""" + @abstractmethod + def replay_spec(self, args: argparse.Namespace) -> DemoSpec: + """Build the model-specific replay specification.""" -__all__ = ["run_flashdreams_demo", "serve_flashdreams_demo"] + @abstractmethod + def replay_adapter(self) -> DemoAdapter: + """Create the model-specific replay adapter.""" + + def prepare_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: + """Perform optional model-specific setup before serving WebRTC.""" + del args, context + + @abstractmethod + def serve_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: + """Build and serve the model-specific WebRTC demo.""" + + +__all__ = ["DemoApplication"] diff --git a/flashdreams/flashdreams/runtime/demo/spec.py b/flashdreams/flashdreams/runtime/demo/spec.py index bc2884ab3..6ba652f38 100644 --- a/flashdreams/flashdreams/runtime/demo/spec.py +++ b/flashdreams/flashdreams/runtime/demo/spec.py @@ -5,7 +5,7 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass, field, replace from pathlib import Path from typing import Any, Literal, Protocol, TypeAlias @@ -86,6 +86,15 @@ def __post_init__(self) -> None: OutputSpec: TypeAlias = NullOutputSpec | Mp4OutputSpec | WebRTCOutputSpec +@dataclass(frozen=True, kw_only=True, slots=True) +class WebRTCAppResources: + """Model-owned resources attached to the shared WebRTC application.""" + + model_web_resource: Any | None = None + configure_app: Callable[[Any], None] | None = None + preload_name: str | None = None + + @dataclass(frozen=True, kw_only=True, slots=True) class DemoSpec: """User-facing shared demo run description.""" @@ -146,7 +155,7 @@ def __post_init__(self) -> None: class DemoAdapter(ModelAdapter, Protocol): - """Model-owned adapter surface consumed by shared demo launchers.""" + """Transport-neutral model adapter consumed by demo runners.""" def supported_input_modes(self) -> tuple[str, ...]: """Return demo input modes this adapter can prepare.""" @@ -160,10 +169,6 @@ def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: """Validate and materialize scenario inputs before runtime creation.""" ... - def create_webrtc_runtime(self, spec: DemoSpec) -> Any: - """Create the model-owned runtime consumed by the shared WebRTC manager.""" - ... - __all__ = [ "DemoAdapter", @@ -173,4 +178,5 @@ def create_webrtc_runtime(self, spec: DemoSpec) -> Any: "OutputSpec", "PreparedScenario", "WebRTCOutputSpec", + "WebRTCAppResources", ] diff --git a/flashdreams/flashdreams/runtime/demo/webrtc.py b/flashdreams/flashdreams/runtime/demo/webrtc.py index f93a855db..b9444a197 100644 --- a/flashdreams/flashdreams/runtime/demo/webrtc.py +++ b/flashdreams/flashdreams/runtime/demo/webrtc.py @@ -6,7 +6,7 @@ from __future__ import annotations from collections.abc import Callable -from dataclasses import dataclass +from importlib.resources import files from pathlib import Path from typing import Any @@ -14,230 +14,74 @@ from flashdreams.serving.webrtc.bootstrap import run_webrtc_server from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager -from flashdreams.serving.webrtc.server import create_webrtc_app - -from .replay import _require_supported_mode -from .spec import DemoAdapter, DemoSpec, WebRTCOutputSpec - - -@dataclass(frozen=True, kw_only=True, slots=True) -class WebRTCDemoRuntimeConfig: - """Runtime config consumed by the shared WebRTC session manager.""" - - video_width: int - video_height: int - warmup_chunks: int - warmup_timeout_s: float - - -class SharedDemoWebRTCSessionManager(BaseWebRTCSessionManager[Any, Any]): - """Generic session manager wrapper for demo adapters.""" - - def __init__( - self, - *, - model_name: str, - runtime: Any, - runtime_config: Any, - fps: int, - client_liveness_timeout_s: float, - ) -> None: - self._demo_model_name = model_name - super().__init__( - runtime=runtime, - runtime_config=runtime_config, - fps=fps, - client_liveness_timeout_s=client_liveness_timeout_s, - ) - - def _model_name(self) -> str: - return self._demo_model_name - - -@dataclass(frozen=True, kw_only=True, slots=True) -class WebRTCDemo: - """Constructed WebRTC demo pieces, before or after serving.""" - - runtime: Any - runtime_config: Any - session_manager: BaseWebRTCSessionManager[Any, Any] - app: web.Application | None - host: str - port: int +from flashdreams.serving.webrtc.server import ( + close_package_resources, + create_packaged_webrtc_app, + create_webrtc_app, +) +from .spec import WebRTCAppResources, WebRTCOutputSpec CreateWebRTCApp = Callable[..., web.Application] RunWebRTCServer = Callable[..., None] -def build_webrtc_demo( +def serve_webrtc_demo( *, - spec: DemoSpec, - adapter: DemoAdapter, - create_app: bool = False, + output: WebRTCOutputSpec, + model_id: str, + session_manager: BaseWebRTCSessionManager[Any, Any], + app_resources: WebRTCAppResources, + world_rank: int = 0, create_app_fn: CreateWebRTCApp = create_webrtc_app, -) -> WebRTCDemo: - """Build shared WebRTC manager/app pieces for a demo adapter runtime.""" - if not isinstance(spec.output, WebRTCOutputSpec): - raise ValueError("build_webrtc_demo requires WebRTCOutputSpec output.") - _require_supported_mode( - mode=spec.input_mode, - supported=adapter.supported_input_modes(), - label="input_mode", - ) - _require_supported_mode( - mode=spec.output.mode, - supported=adapter.supported_output_modes(), - label="output.mode", - ) - - output = spec.output - runtime = adapter.create_webrtc_runtime(spec) - runtime_config = _create_runtime_config( - spec=spec, - adapter=adapter, - runtime=runtime, - ) - manager = _create_session_manager( - spec=spec, - adapter=adapter, - runtime=runtime, - runtime_config=runtime_config, - fps=output.fps, - client_liveness_timeout_s=output.client_liveness_timeout_s, - ) + server_runner: RunWebRTCServer = run_webrtc_server, +) -> web.Application | None: + """Serve a prepared model WebRTC runtime through the shared transport.""" app = ( _create_app( - spec=spec, - adapter=adapter, - session_manager=manager, + output=output, + model_id=model_id, + app_resources=app_resources, + session_manager=session_manager, create_app_fn=create_app_fn, ) - if create_app + if world_rank == 0 else None ) - return WebRTCDemo( - runtime=runtime, - runtime_config=runtime_config, - session_manager=manager, + server_runner( + world_rank=world_rank, + session_manager=session_manager, app=app, host=output.host, port=output.port, ) - - -def serve_webrtc_demo( - *, - spec: DemoSpec, - adapter: DemoAdapter, - world_rank: int = 0, - create_app_fn: CreateWebRTCApp = create_webrtc_app, - server_runner: RunWebRTCServer = run_webrtc_server, -) -> WebRTCDemo: - """Build and serve a shared WebRTC demo.""" - demo = build_webrtc_demo( - spec=spec, - adapter=adapter, - create_app=world_rank == 0, - create_app_fn=create_app_fn, - ) - server_runner( - world_rank=world_rank, - session_manager=demo.session_manager, - app=demo.app, - host=demo.host, - port=demo.port, - ) - return demo - - -def _create_runtime_config( - *, - spec: DemoSpec, - adapter: DemoAdapter, - runtime: Any, -) -> Any: - factory = getattr(adapter, "create_webrtc_runtime_config", None) - if callable(factory): - return factory(spec=spec, runtime=runtime) - - runtime_config = getattr(runtime, "config", None) - if _looks_like_webrtc_runtime_config(runtime_config): - return runtime_config - - output = spec.output - if not isinstance(output, WebRTCOutputSpec): - raise ValueError("WebRTC runtime config creation requires WebRTCOutputSpec.") - return WebRTCDemoRuntimeConfig( - video_width=output.video_width, - video_height=output.video_height, - warmup_chunks=output.warmup_chunks, - warmup_timeout_s=output.warmup_timeout_s, - ) - - -def _looks_like_webrtc_runtime_config(value: Any) -> bool: - return all( - hasattr(value, name) - for name in ( - "video_width", - "video_height", - "warmup_chunks", - "warmup_timeout_s", - ) - ) - - -def _create_session_manager( - *, - spec: DemoSpec, - adapter: DemoAdapter, - runtime: Any, - runtime_config: Any, - fps: int, - client_liveness_timeout_s: float, -) -> BaseWebRTCSessionManager[Any, Any]: - factory = getattr(adapter, "create_webrtc_session_manager", None) - if callable(factory): - return factory( - spec=spec, - runtime=runtime, - runtime_config=runtime_config, - fps=fps, - client_liveness_timeout_s=client_liveness_timeout_s, - ) - - return SharedDemoWebRTCSessionManager( - model_name=spec.model_id, - runtime=runtime, - runtime_config=runtime_config, - fps=fps, - client_liveness_timeout_s=client_liveness_timeout_s, - ) + return app def _create_app( *, - spec: DemoSpec, - adapter: DemoAdapter, + output: WebRTCOutputSpec, + model_id: str, + app_resources: WebRTCAppResources, session_manager: BaseWebRTCSessionManager[Any, Any], create_app_fn: CreateWebRTCApp, ) -> web.Application: - output = spec.output - if not isinstance(output, WebRTCOutputSpec): - raise ValueError("WebRTC app creation requires WebRTCOutputSpec output.") - factory = getattr(adapter, "create_webrtc_app", None) - if callable(factory): - return factory( - spec=spec, + if output.web_dir is not None: + return _build_webrtc_app( + output=output, session_manager=session_manager, - request_session_url=_request_session_url(output), + create_app_fn=create_app_fn, + preload_name=output.preload_name or app_resources.preload_name or model_id, ) - return _build_webrtc_app( - output=output, + return create_packaged_webrtc_app( + web_resource=files("flashdreams.serving.webrtc").joinpath("web"), + model_web_resource=app_resources.model_web_resource, session_manager=session_manager, + request_session_url=_request_session_url(output), + preload_name=output.preload_name or app_resources.preload_name or model_id, + configure_app=app_resources.configure_app, create_app_fn=create_app_fn, - preload_name=output.preload_name or spec.model_id, + cleanup_callback=close_package_resources, ) @@ -266,9 +110,5 @@ def _request_session_url(output: WebRTCOutputSpec) -> str: __all__ = [ "CreateWebRTCApp", "RunWebRTCServer", - "SharedDemoWebRTCSessionManager", - "WebRTCDemo", - "WebRTCDemoRuntimeConfig", - "build_webrtc_demo", "serve_webrtc_demo", ] diff --git a/flashdreams/flashdreams/runtime/inputs.py b/flashdreams/flashdreams/runtime/inputs.py index 9174b6a84..70260de2e 100644 --- a/flashdreams/flashdreams/runtime/inputs.py +++ b/flashdreams/flashdreams/runtime/inputs.py @@ -10,6 +10,7 @@ from dataclasses import dataclass, field from typing import Any, Literal, cast +from flashdreams.infra.time import TimeWindow from flashdreams.runtime._utils import freeze_mapping InputPhase = Literal["global_conditioning", "step"] @@ -26,26 +27,6 @@ def validate_phase(value: str) -> InputPhase: return cast(InputPhase, value) -@dataclass(frozen=True, kw_only=True, slots=True) -class TimeWindow: - """Half-open time window in seconds since session start.""" - - start_s: float - end_s: float - - def __post_init__(self) -> None: - if not math.isfinite(self.start_s) or not math.isfinite(self.end_s): - raise ValueError("TimeWindow bounds must be finite seconds.") - if self.start_s < 0 or self.end_s < 0: - raise ValueError("TimeWindow bounds must be non-negative.") - if self.end_s < self.start_s: - raise ValueError("TimeWindow.end_s must be >= start_s.") - - def contains(self, timestamp_s: float) -> bool: - """Return whether ``timestamp_s`` falls within this half-open window.""" - return self.start_s <= timestamp_s < self.end_s - - @dataclass(frozen=True, kw_only=True, slots=True) class InputField: """Lightweight schema field for user snapshots or model inputs. diff --git a/flashdreams/flashdreams/runtime/runner.py b/flashdreams/flashdreams/runtime/runner.py index 03d814472..c73a4f67c 100644 --- a/flashdreams/flashdreams/runtime/runner.py +++ b/flashdreams/flashdreams/runtime/runner.py @@ -143,7 +143,10 @@ def _all_user_inputs_window(user_inputs: UserInputs) -> TimeWindow: ) -def _record_timing_metrics(metrics: MetricsRecorder, result: StepResult) -> None: +def _record_timing_metrics( + metrics: MetricsRecorder, + result: StepResult, +) -> None: for name, value in result.metrics.items(): if not name.endswith("_s") or isinstance(value, bool): continue diff --git a/flashdreams/flashdreams/runtime/types.py b/flashdreams/flashdreams/runtime/types.py index 51d3846db..4130a925f 100644 --- a/flashdreams/flashdreams/runtime/types.py +++ b/flashdreams/flashdreams/runtime/types.py @@ -9,6 +9,7 @@ from dataclasses import dataclass, field from typing import Any +from flashdreams.infra.results import StepResult from flashdreams.runtime._utils import freeze_mapping from flashdreams.runtime.inputs import InferenceInputSchema, TimeWindow @@ -35,23 +36,4 @@ def __post_init__(self) -> None: object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) -@dataclass(frozen=True, kw_only=True, slots=True) -class StepResult: - """Generated output and metadata returned by one inference step.""" - - __hash__ = None - - step_index: int - output: Any = None - frame_count: int | None = None - output_window: TimeWindow | None = None - metadata: Mapping[str, Any] = field(default_factory=dict) - metrics: Mapping[str, float | int] = field(default_factory=dict) - - def __post_init__(self) -> None: - if self.step_index < 0: - raise ValueError("StepResult.step_index must be >= 0.") - if self.frame_count is not None and self.frame_count < 0: - raise ValueError("StepResult.frame_count must be >= 0.") - object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) - object.__setattr__(self, "metrics", freeze_mapping(self.metrics)) +__all__ = ["StepRequest", "StepResult"] diff --git a/flashdreams/flashdreams/runtime/video_output.py b/flashdreams/flashdreams/runtime/video_output.py index bd2e466e2..c92cf2032 100644 --- a/flashdreams/flashdreams/runtime/video_output.py +++ b/flashdreams/flashdreams/runtime/video_output.py @@ -14,7 +14,7 @@ DEFAULT_RUNNER_INSTALL_HINT, write_video_tensor, ) -from flashdreams.infra.video_output import VideoOutputStream, VideoStepResult +from flashdreams.infra.video_output import VideoResultCollector, prepare_video_for_mp4 from flashdreams.runtime.output import OutputArtifact from flashdreams.runtime.types import StepResult @@ -23,7 +23,7 @@ @dataclass(slots=True) class Mp4VideoOutputTarget: - """Write runtime ``VideoStepResult`` chunks to one MP4 artifact.""" + """Write layout-aware runtime step results to one MP4 artifact.""" output_path: Path fps: int | float @@ -31,8 +31,9 @@ class Mp4VideoOutputTarget: writer: VideoWriter = field(default=write_video_tensor, repr=False) install_hint: str = DEFAULT_RUNNER_INSTALL_HINT move_to_cpu: bool = True + enabled: bool = True _opened: bool = field(default=False, init=False, repr=False) - _stream: VideoOutputStream | None = field( + _collector: VideoResultCollector | None = field( default=None, init=False, repr=False, @@ -43,60 +44,46 @@ def closed(self) -> bool: return not self._opened def open(self) -> None: - self._stream = VideoOutputStream( - postprocess_stream=None, + self._collector = VideoResultCollector( output_layout=self.output_layout, - collect_output=True, + enabled=self.enabled, move_to_cpu=self.move_to_cpu, ) self._opened = True def write(self, result: StepResult) -> None: - if not self._opened or self._stream is None: + if not self._opened or self._collector is None: raise RuntimeError("Cannot write to a closed output target.") - video_result = result.output - if not isinstance(video_result, VideoStepResult): + if result.layout is None: raise TypeError( - "Mp4VideoOutputTarget requires StepResult.output to be " - f"VideoStepResult, got {type(video_result).__name__}." + "Mp4VideoOutputTarget requires a video StepResult with layout." ) - if video_result.layout != self.output_layout: + if result.layout != self.output_layout: raise ValueError( "Mp4VideoOutputTarget received layout " - f"{video_result.layout!r}; expected {self.output_layout!r}." + f"{result.layout!r}; expected {self.output_layout!r}." ) - stats = dict(video_result.stats or result.metrics) - stats_extra: dict[str, object] = { - "step_index": result.step_index, - "frames": video_result.num_frames, - } - if result.output_window is not None: - stats_extra["output_start_s"] = result.output_window.start_s - stats_extra["output_end_s"] = result.output_window.end_s - self._stream.process( - video_result.video_chunk, - autoregressive_index=video_result.chunk_index, - stats=stats if stats else None, - stats_extra=stats_extra, - ) + self._collector.add(result) def close(self) -> Sequence[OutputArtifact]: - if self._stream is None: + if self._collector is None: self._opened = False return () - stream = self._stream - self._stream = None + collector = self._collector + self._collector = None self._opened = False - video = stream.finish() + video = collector.finish() if video is None: return () - path = stream.write_mp4( - video, + writable_video, writable_layout = prepare_video_for_mp4( + video, layout=self.output_layout + ) + path = self.writer( + writable_video, self.output_path, fps=self.fps, - layout=self.output_layout, - writer=self.writer, + layout=writable_layout, install_hint=self.install_hint, ) return ( @@ -107,7 +94,7 @@ def close(self) -> Sequence[OutputArtifact]: "fps": self.fps, "source_layout": self.output_layout, "shape": tuple(int(dim) for dim in video.shape), - "stats_history": tuple(stream.stats_history), + "stats_history": tuple(collector.stats_history), }, ), ) diff --git a/flashdreams/flashdreams/runtime/worker.py b/flashdreams/flashdreams/runtime/worker.py new file mode 100644 index 000000000..6e0c17eb2 --- /dev/null +++ b/flashdreams/flashdreams/runtime/worker.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Thread-affine execution for stateful inference runtimes.""" + +from __future__ import annotations + +import asyncio +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Callable, TypeVar + +import torch + +_T = TypeVar("_T") + + +class ThreadAffineRuntimeWorker: + """Run ordered runtime lifecycle calls on one owned OS thread. + + CUDA graphs, Triton launchers, and some backend contexts are thread-local. + A runtime should therefore submit initialization, reset, generation, and + close operations through one worker instead of using ``asyncio.to_thread``. + + Cancelling an awaiting task does not cancel the submitted operation. The + operation remains ordered on the worker, and later calls run only after it + completes. + """ + + def __init__( + self, + *, + device: torch.device | str | None = None, + thread_name: str = "flashdreams-runtime", + ) -> None: + self._device = None if device is None else torch.device(device) + self._executor = ThreadPoolExecutor( + max_workers=1, + thread_name_prefix=thread_name, + initializer=self._initialize_thread, + ) + self._accepting = True + self._closed = False + self._close_lock = asyncio.Lock() + + @property + def closed(self) -> bool: + return self._closed + + async def call( + self, + func: Callable[..., _T], + /, + *args: Any, + **kwargs: Any, + ) -> _T: + """Run one callable after all previously submitted worker calls.""" + if not self._accepting: + raise RuntimeError("runtime worker is closed") + future = self._submit(func, args, kwargs) + try: + return await asyncio.shield(future) + except asyncio.CancelledError: + future.add_done_callback(_consume_exception) + raise + + async def close(self) -> None: + """Drain submitted work and stop accepting lifecycle calls.""" + async with self._close_lock: + if self._closed: + return + self._accepting = False + barrier = self._submit(_noop, (), {}) + await asyncio.shield(barrier) + self._executor.shutdown(wait=True, cancel_futures=False) + self._closed = True + + def _submit( + self, + func: Callable[..., _T], + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> asyncio.Future[_T]: + loop = asyncio.get_running_loop() + return loop.run_in_executor(self._executor, _invoke, func, args, kwargs) + + def _initialize_thread(self) -> None: + if self._device is not None and self._device.type == "cuda": + torch.cuda.set_device(self._device) + + +def _invoke( + func: Callable[..., _T], + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> _T: + return func(*args, **kwargs) + + +def _noop() -> None: + return + + +def _consume_exception(future: asyncio.Future[Any]) -> None: + if not future.cancelled(): + future.exception() + + +__all__ = ["ThreadAffineRuntimeWorker"] diff --git a/flashdreams/flashdreams/serving/output_targets.py b/flashdreams/flashdreams/serving/output_targets.py index f5306dbbf..a5e9f961b 100644 --- a/flashdreams/flashdreams/serving/output_targets.py +++ b/flashdreams/flashdreams/serving/output_targets.py @@ -5,27 +5,19 @@ from __future__ import annotations +import importlib import runpy import shlex import sys from dataclasses import dataclass +from functools import lru_cache from pathlib import Path -from typing import Any, Literal, TypeAlias +from typing import Literal, Protocol, TypeAlias, runtime_checkable from flashdreams.infra.runner import RunnerConfig OutputMode: TypeAlias = Literal["cli", "webrtc", "local-window"] -_OMNIDREAMS_LOCAL_WINDOW_MANIFESTS = { - "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae": ("example_world_model.yaml"), - "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf": ( - "example_world_model_perf.yaml" - ), - "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-native-perf": ( - "example_world_model_perf.yaml" - ), -} - class OutputTargetUnavailableError(ValueError): """Raised when a runner cannot be launched through a requested output.""" @@ -57,18 +49,39 @@ def command(self) -> str: return shlex.join(("python", "-m", self.module, *self.argv)) +@runtime_checkable +class OutputTargetAdapter(Protocol): + """Integration-owned non-CLI output capabilities for a runner config.""" + + def supported_modes( + self, + config: RunnerConfig, + options: OutputLaunchOptions, + ) -> tuple[OutputMode, ...]: ... + + def resolve( + self, + config: RunnerConfig, + *, + mode: OutputMode, + options: OutputLaunchOptions, + ) -> OutputTargetSpec | None: ... + + def available_output_modes( config: RunnerConfig, options: OutputLaunchOptions | None = None, ) -> tuple[OutputMode, ...]: """Return output modes known to support ``config``.""" options = options or OutputLaunchOptions() - modes: list[OutputMode] = ["cli"] - if _webrtc_spec(config, options) is not None: - modes.append("webrtc") - if _local_window_spec(config, options) is not None: - modes.append("local-window") - return tuple(modes) + adapter = _resolve_adapter(config) + if adapter is None: + return ("cli",) + modes = adapter.supported_modes(config, options) + invalid = [mode for mode in modes if mode == "cli"] + if invalid: + raise ValueError("Output adapters must not declare the built-in CLI mode.") + return ("cli", *dict.fromkeys(modes)) def resolve_output_target( @@ -81,10 +94,9 @@ def resolve_output_target( if mode == "cli": raise ValueError("CLI mode is run directly by the selected Runner.") options = options or OutputLaunchOptions() + adapter = _resolve_adapter(config) spec = ( - _webrtc_spec(config, options) - if mode == "webrtc" - else _local_window_spec(config, options) + None if adapter is None else adapter.resolve(config, mode=mode, options=options) ) if spec is None: supported = ", ".join(available_output_modes(config, options)) @@ -92,6 +104,10 @@ def resolve_output_target( f"Output mode {mode!r} is not available for runner " f"{config.runner_name!r}. Supported modes: {supported}." ) + if spec.mode != mode: + raise ValueError( + f"Output adapter returned mode {spec.mode!r} while resolving {mode!r}." + ) return spec @@ -105,185 +121,37 @@ def launch_output_target(spec: OutputTargetSpec) -> None: sys.argv = original_argv -def _webrtc_spec( - config: RunnerConfig, - options: OutputLaunchOptions, -) -> OutputTargetSpec | None: - name = _runner_name(config) - if _is_lingbot_runner(name): - return _lingbot_webrtc_spec(config, options) - if _is_omnidreams_runner(name) and _is_omnidreams_single_view(config): - return _omnidreams_webrtc_spec(config, options) - return None - - -def _local_window_spec( - config: RunnerConfig, - options: OutputLaunchOptions, -) -> OutputTargetSpec | None: - name = _runner_name(config) - if not _is_omnidreams_runner(name): +def _resolve_adapter(config: RunnerConfig) -> OutputTargetAdapter | None: + path = config.output_adapter + if not path: return None - manifest = options.local_window_manifest - if manifest is None: - manifest_name = _OMNIDREAMS_LOCAL_WINDOW_MANIFESTS.get(name) - if manifest_name is None: - return None - manifest_arg = manifest_name - else: - manifest_arg = str(manifest) - - argv = ["--manifest", manifest_arg] - _append_postprocess_preset(argv, config) - return OutputTargetSpec( - mode="local-window", - label="Omnidreams local interactive window", - module="omnidreams.interactive_drive", - argv=tuple(argv), - notes=( - ( - "Local-window uses the Omnidreams interactive-drive manifest for " - "scene, resolution, and runtime-specific controls." - ), - ), - ) - - -def _lingbot_webrtc_spec( - config: RunnerConfig, - options: OutputLaunchOptions, -) -> OutputTargetSpec: - argv = [ - "webrtc", - "--preset-id", - _pipeline_name(config), - "--device", - _device(config), - "--fps", - str(getattr(config, "fps", 16)), - "--video-height", - str(getattr(config, "pixel_height", 464)), - "--video-width", - str(getattr(config, "pixel_width", 832)), - ] - if _compile_network(config) is False: - argv.append("--no-compile") - example_idx = getattr(config, "example_idx", None) - if example_idx is not None: - argv.extend(("--example-idx", str(example_idx))) - if options.host: - argv.extend(("--host", options.host)) - if options.port is not None: - argv.extend(("--port", str(options.port))) - if options.prefer_sw_encoder: - argv.append("--prefer-sw-encoder") - return OutputTargetSpec( - mode="webrtc", - label="LingBot shared demo WebRTC server", - module="lingbot.demo.cli", - argv=tuple(argv), - ) - - -def _omnidreams_webrtc_spec( - config: RunnerConfig, - options: OutputLaunchOptions, -) -> OutputTargetSpec: - argv = [ - "--pipeline_config_name", - _pipeline_name(config), - "--device", - _device(config), - "--fps", - str(getattr(config, "output_fps", 30)), - "--video_height", - str(getattr(config, "pixel_height", 704)), - "--video_width", - str(getattr(config, "pixel_width", 1280)), - ] - seed = _diffusion_seed(config) - if seed is not None: - argv.extend(("--seed", str(seed))) - _append_postprocess_preset(argv, config) - _append_webrtc_bind_args(argv, options) - return OutputTargetSpec( - mode="webrtc", - label="Omnidreams WebRTC server", - module="omnidreams.webrtc.server", - argv=tuple(argv), - ) - + return _load_output_adapter(path) -def _append_webrtc_bind_args( - argv: list[str], - options: OutputLaunchOptions, -) -> None: - if options.host: - argv.extend(("--host", options.host)) - if options.port is not None: - argv.extend(("--port", str(options.port))) - if options.prefer_sw_encoder: - argv.append("--prefer_sw_encoder") - -def _append_postprocess_preset(argv: list[str], config: RunnerConfig) -> None: - preset = getattr(getattr(config, "postprocess", None), "preset", "") - if preset: - argv.extend(("--postprocess-preset", str(preset))) - - -def _runner_name(config: RunnerConfig) -> str: - return str(getattr(config, "runner_name", "")) - - -def _pipeline_name(config: RunnerConfig) -> str: - pipeline = getattr(config, "pipeline", None) - name = getattr(pipeline, "name", None) - return str(name or config.runner_name) - - -def _device(config: RunnerConfig) -> str: - return str(getattr(config, "device", "cuda")) - - -def _compile_network(config: RunnerConfig) -> bool | None: - transformer = _transformer_config(config) - value = getattr(transformer, "compile_network", None) - return None if value is None else bool(value) - - -def _diffusion_seed(config: RunnerConfig) -> int | None: - diffusion_model = getattr( - getattr(config, "pipeline", None), "diffusion_model", None - ) - seed = getattr(diffusion_model, "seed", None) - return None if seed is None else int(seed) - - -def _transformer_config(config: RunnerConfig) -> Any: - diffusion_model = getattr( - getattr(config, "pipeline", None), "diffusion_model", None - ) - return getattr(diffusion_model, "transformer", None) - - -def _is_lingbot_runner(name: str) -> bool: - return name.startswith("lingbot-world") - - -def _is_omnidreams_runner(name: str) -> bool: - return name.startswith("omnidreams-") - - -def _is_omnidreams_single_view(config: RunnerConfig) -> bool: - num_views = getattr(_transformer_config(config), "num_views", 1) - return int(num_views) == 1 +@lru_cache(maxsize=None) +def _load_output_adapter(path: str) -> OutputTargetAdapter: + try: + module_name, attribute = path.split(":", 1) + except ValueError as exc: + raise ValueError( + "RunnerConfig.output_adapter must use 'module:attribute' syntax; " + f"got {path!r}." + ) from exc + value = getattr(importlib.import_module(module_name), attribute) + if callable(value) and not isinstance(value, OutputTargetAdapter): + value = value() + if not isinstance(value, OutputTargetAdapter): + raise TypeError( + f"Output adapter {path!r} does not implement OutputTargetAdapter." + ) + return value __all__ = [ "OutputLaunchOptions", "OutputMode", "OutputTargetSpec", + "OutputTargetAdapter", "OutputTargetUnavailableError", "available_output_modes", "launch_output_target", diff --git a/flashdreams/flashdreams/serving/realtime/media.py b/flashdreams/flashdreams/serving/realtime/media.py index 9b93eaca1..ab63899e2 100644 --- a/flashdreams/flashdreams/serving/realtime/media.py +++ b/flashdreams/flashdreams/serving/realtime/media.py @@ -14,7 +14,15 @@ if TYPE_CHECKING: import torch -FrameLayout = Literal["hwc", "chw", "thwc", "tchw", "bvtchw"] +FrameLayout = Literal[ + "hwc", + "chw", + "thwc", + "tchw", + "btchw", + "bcthw", + "bvtchw", +] ValueRange = Literal["minus_one_one", "zero_one", "uint8"] @@ -125,6 +133,18 @@ def rgb_array_to_uint8_frames( f"[1, 1, T, 3, H, W], got {array.shape}" ) frames = np.transpose(array[0, 0], (0, 2, 3, 1)) + elif layout == "btchw": + if array.ndim != 5 or array.shape[0] != 1 or array.shape[2] != 3: + raise ValueError( + f"Expected single-batch video chunk [1, T, 3, H, W], got {array.shape}" + ) + frames = np.transpose(array[0], (0, 2, 3, 1)) + elif layout == "bcthw": + if array.ndim != 5 or array.shape[0] != 1 or array.shape[1] != 3: + raise ValueError( + f"Expected single-batch video chunk [1, 3, T, H, W], got {array.shape}" + ) + frames = np.transpose(array[0], (1, 2, 3, 0)) else: raise ValueError(f"Unsupported layout={layout!r}.") diff --git a/flashdreams/flashdreams/serving/webrtc/encoders.py b/flashdreams/flashdreams/serving/webrtc/encoders.py index bb27c7f6c..5eea69991 100644 --- a/flashdreams/flashdreams/serving/webrtc/encoders.py +++ b/flashdreams/flashdreams/serving/webrtc/encoders.py @@ -3,11 +3,9 @@ """Video encoder backends for the WebRTC serving path. -Integrations that opt in to hardware encoding call :func:`select_encoder` -from their own session init (omnidreams does this today via -``omnidreams.webrtc.session._initialize_video_encoder_sync``); those that -do not opt in pick up :class:`DefaultRTCEncoder` transparently through -:meth:`BaseWebRTCSessionManager._resolve_video_encoder`. +Thread-affine WebRTC runtimes call :func:`select_encoder` during shared runtime +initialization. Runtimes that do not opt in pick up :class:`DefaultRTCEncoder` +transparently through :meth:`BaseWebRTCSessionManager._resolve_video_encoder`. **This module deliberately does not import** ``PyNvVideoCodec``. The hardware encoder lives in a sibling module (:mod:`nvenc`) that @@ -28,6 +26,8 @@ from aiortc import MediaStreamTrack from loguru import logger +from flashdreams.runtime import StepResult + if TYPE_CHECKING: from flashdreams.serving.webrtc.media import BufferedVideoTrack, NVENCVideoTrack @@ -71,7 +71,7 @@ def create_track(self, *, maxsize: int) -> BufferedVideoTrack | NVENCVideoTrack: async def deliver_chunk( self, - chunk: torch.Tensor, + result: StepResult, track: MediaStreamTrack, *, force_keyframe: bool = False, @@ -112,7 +112,7 @@ def create_track(self, *, maxsize: int) -> BufferedVideoTrack: async def deliver_chunk( self, - chunk: torch.Tensor, + result: StepResult, track: MediaStreamTrack, *, force_keyframe: bool = False, @@ -128,7 +128,7 @@ async def deliver_chunk( "DefaultRTCEncoder requires a BufferedVideoTrack; got " f"{type(track).__name__}. Create it via encoder.create_track()." ) - enqueued = await track.enqueue_chunk(chunk) + enqueued = await track.enqueue_result(result) return ChunkDeliveryResult( backend=self.backend, num_frames=enqueued, diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index e35e70263..30fcbc56c 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -12,7 +12,6 @@ from collections import deque from collections.abc import Set as AbstractSet from dataclasses import dataclass, field, replace -from enum import IntEnum from typing import Any, Generic, TypeVar from aiortc import ( @@ -23,18 +22,9 @@ ) from loguru import logger -from flashdreams.infra.video_output import VideoStepResult -from flashdreams.runtime.inputs import ( - InferenceInput, - TimeWindow, - UserInputEvent, - UserInputs, -) -from flashdreams.serving.realtime.input import ( - DEFAULT_SUPPORTED_KEYS, - KeyboardResampler, - normalize_key, -) +from flashdreams.runtime.inputs import TimeWindow +from flashdreams.runtime.types import StepRequest, StepResult +from flashdreams.serving.realtime.input import KeyboardResampler from flashdreams.serving.webrtc.encoders import ( DefaultRTCEncoder, VideoEncoder, @@ -50,6 +40,7 @@ make_event_ack_payload, ) from flashdreams.serving.webrtc.runtime import ( + WebRTCControlSignal, WebRTCRuntimeConfig, WebRTCSessionRuntime, ) @@ -63,7 +54,7 @@ "BaseWebRTCSessionManager", "ManagedWebRTCSession", "WebRTCControlSignal", - "VideoStepResult", + "StepResult", ] # Close the active session if no client heartbeat/control message arrives @@ -82,6 +73,40 @@ _RuntimeConfigT = TypeVar("_RuntimeConfigT", bound=WebRTCRuntimeConfig) +def _summarize_sdp_candidates(sdp: str) -> str: + candidates = [ + line.removeprefix("a=candidate:") + for line in sdp.splitlines() + if line.startswith("a=candidate:") + ] + if not candidates: + return "0 candidates" + + protocols: dict[str, int] = {} + addresses: set[str] = set() + endpoints: list[str] = [] + for candidate in candidates: + parts = candidate.split() + if len(parts) >= 5: + protocols[parts[2].lower()] = protocols.get(parts[2].lower(), 0) + 1 + addresses.add(parts[4]) + if len(parts) >= 6: + endpoints.append(f"{parts[2].lower()}://{parts[4]}:{parts[5]}") + protocol_summary = ",".join( + f"{key}={value}" for key, value in sorted(protocols.items()) + ) + address_summary = ",".join(sorted(addresses)[:8]) + if len(addresses) > 8: + address_summary += f",+{len(addresses) - 8} more" + endpoint_summary = ",".join(endpoints[:12]) + if len(endpoints) > 12: + endpoint_summary += f",+{len(endpoints) - 12} more" + return ( + f"{len(candidates)} candidates protocols=[{protocol_summary}] " + f"addresses=[{address_summary}] endpoints=[{endpoint_summary}]" + ) + + def _stat_float(stats: dict[str, float], name: str, default: float = 0.0) -> float: value = stats.get(name) if value is None: @@ -97,19 +122,6 @@ def _stat_int(stats: dict[str, float], name: str) -> int: return int(round(_stat_float(stats, name))) -class WebRTCControlSignal(IntEnum): - """Rank-orchestration signals shared by the single-session runtimes.""" - - INITIALIZE = 0 - RESET_SESSION = 1 - ACTION_STEP = 2 - CLOSE = 3 - EVENT = 4 - SESSION_STEP = 5 - """One step driven by mapped ``InferenceInput`` rather than pose segments.""" - EXIT = 99 - - @dataclass(slots=True) class ManagedWebRTCSession: """Per-session state for the single active WebRTC peer connection.""" @@ -168,11 +180,6 @@ async def close(self) -> None: class BaseWebRTCSessionManager(Generic[_RuntimeT, _RuntimeConfigT]): """Owns one active WebRTC session and forwards actions into a model runtime.""" - _busy_message: str = "A WebRTC session is already active." - _warmup_label: str = "WebRTC" - _runtime_error_types: tuple[type[Exception], ...] = (RuntimeError,) - _close_session_on_generation_error: bool = False - _resampler_supported_keys: AbstractSet[str] | None = None _perf_log_interval_chunks: int = _DEFAULT_PERF_LOG_INTERVAL_CHUNKS def __init__( @@ -181,12 +188,26 @@ def __init__( runtime: _RuntimeT, runtime_config: _RuntimeConfigT, fps: int, + identity: str, + busy_message: str = "A WebRTC session is already active.", + warmup_label: str = "WebRTC", + supported_control_keys: AbstractSet[str] | None = None, + fatal_generation_errors: bool = False, client_liveness_timeout_s: float = DEFAULT_CLIENT_LIVENESS_TIMEOUT_S, ) -> None: if client_liveness_timeout_s <= 0: raise ValueError("client_liveness_timeout_s must be > 0") self.runtime_config = runtime_config self.fps = fps + self.identity = identity + self.busy_message = busy_message + self.warmup_label = warmup_label + self.supported_control_keys = ( + None + if supported_control_keys is None + else frozenset(supported_control_keys) + ) + self.fatal_generation_errors = fatal_generation_errors self.client_liveness_timeout_s = client_liveness_timeout_s self._runtime = runtime self._runtime_ready = False @@ -194,21 +215,23 @@ def __init__( self._active_session: ManagedWebRTCSession | None = None self._preload_lock = asyncio.Lock() self._session_lock = asyncio.Lock() + self._pending_session_input: Any = None - def _model_name(self) -> str: - """Human-readable model identifier reported in ``chunk_done``.""" - raise NotImplementedError + @property + def pending_session_input(self) -> Any: + """Input that will be applied to the next successfully negotiated session.""" + return self._pending_session_input - def _peek_pending_session_input(self) -> Any: - """Session input applied to the next ``create_answer`` (or ``None``).""" - return None + @property + def runtime(self) -> _RuntimeT: + """Model runtime driven by this transport manager.""" + return self._runtime - def _clear_pending_session_input(self) -> None: - """Clear the pending session input after a successful answer.""" - - async def _reset_runtime_for_session(self, session_input: Any) -> None: - """Reset the runtime for a new rollout, honoring ``session_input``.""" - await self._runtime.reset_for_new_session() + def set_pending_session_input(self, session_input: Any) -> None: + """Store validated model input for the next session.""" + if self.has_active_session(): + raise SessionBusyError(self.busy_message) + self._pending_session_input = session_input def _make_resampler(self, *, start_v: float) -> KeyboardResampler: return self._make_resampler_at_fps(start_v=start_v, fps=self.fps) @@ -216,12 +239,12 @@ def _make_resampler(self, *, start_v: float) -> KeyboardResampler: def _make_resampler_at_fps( self, *, start_v: float, fps: float ) -> KeyboardResampler: - if self._resampler_supported_keys is None: + if self.supported_control_keys is None: return KeyboardResampler(fps=fps, start_v=start_v) return KeyboardResampler( fps=fps, start_v=start_v, - supported_keys=frozenset(self._resampler_supported_keys), + supported_keys=self.supported_control_keys, ) @staticmethod @@ -245,46 +268,35 @@ def _positive_float_runtime_value(value: Any, *, label: str) -> float: return parsed def _runtime_input_fps(self, runtime: Any) -> float: - method = getattr(runtime, "peek_input_fps", None) - if callable(method): - return self._positive_float_runtime_value( - method(), - label="peek_input_fps", - ) - return float(self.fps) - - def _runtime_next_input_num_frames(self, runtime: Any) -> int: - method = getattr(runtime, "peek_next_input_num_frames", None) - if callable(method): - return self._positive_int_runtime_value( - method(), - label="peek_next_input_num_frames", + return self._positive_float_runtime_value( + runtime.peek_input_fps(), + label="peek_input_fps", + ) + + def _runtime_next_step_request(self, runtime: Any) -> tuple[StepRequest, int]: + request = runtime.next_step_request() + if not isinstance(request, StepRequest): + raise TypeError( + "next_step_request must return StepRequest, " + f"got {type(request).__name__}." ) - return self._positive_int_runtime_value( - runtime.peek_next_chunk_num_frames(), - label="peek_next_chunk_num_frames", + input_num_frames = self._positive_int_runtime_value( + request.metadata.get("input_frame_count"), + label="StepRequest.metadata['input_frame_count']", ) + return request, input_num_frames def _runtime_steady_output_num_frames(self, runtime: Any) -> int: - method = getattr(runtime, "peek_steady_output_num_frames", None) - if callable(method): - return self._positive_int_runtime_value( - method(), - label="peek_steady_output_num_frames", - ) return self._positive_int_runtime_value( - runtime.peek_steady_chunk_num_frames(), - label="peek_steady_chunk_num_frames", + runtime.peek_steady_output_num_frames(), + label="peek_steady_output_num_frames", ) - def _register_extra_peer_handlers(self, peer_connection: Any) -> None: - """Register optional extra peer-connection event handlers.""" - def _resolve_video_encoder(self) -> VideoEncoder: """Return the encoder to use for the next session. Default: read ``runtime.video_encoder`` if the runtime provides - one (omnidreams does, via ``_initialize_video_encoder_sync``); + one through the shared thread-affine runtime; otherwise construct a session-scope :class:`DefaultRTCEncoder`. Runtimes that do not participate in encoder selection transparently get the software path without having to opt in. @@ -351,7 +363,7 @@ async def _enforce_h264_or_fallback( # is drained. Otherwise ``ManagedWebRTCSession.close()`` would # only ever see the fallback track and never clean this one up. # The hardware encoder itself is owned by the runtime (created - # once in ``_initialize_video_encoder_sync`` and reused across + # once during runtime initialization and reused across # sessions), so it is intentionally NOT closed here — subsequent # sessions read the same object via ``runtime.video_encoder`` # and expect it live. Runtime shutdown releases it. @@ -363,297 +375,6 @@ async def _enforce_h264_or_fallback( 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.""" - - def _on_answer_created(self, answer_sdp: str) -> None: - """Hook invoked with the local answer SDP after negotiation.""" - - def _chunk_done_extra(self) -> dict[str, Any]: - """Extra fields merged into every ``chunk_done`` payload.""" - return {} - - @staticmethod - def _drives_inference_session(runtime: Any) -> bool: - """Return whether ``runtime`` should be driven through ``InferenceSession``.""" - return callable(getattr(runtime, "start_inference_session", None)) - - def _record_user_event( - self, - *, - managed_session: ManagedWebRTCSession, - timestamp_s: float, - event_type: str, - payload: dict[str, Any], - ) -> None: - """Buffer one raw user event for the session branch. - - Timestamps come from the same monotonic clock that anchors the - resampler, so a chunk's ``TimeWindow`` selects exactly the events that - arrived during that chunk's virtual window. - """ - if event_type in _KEY_USER_EVENT_TYPES and not self._supports_key_payload( - payload - ): - return - if len(managed_session.user_events) >= _MAX_SESSION_USER_EVENTS: - if event_type in _RELEASE_USER_EVENT_TYPES: - made_room = self._make_room_for_release_event( - managed_session=managed_session, - event_type=event_type, - payload=payload, - ) - if not made_room: - self._record_coalesced_release_event( - managed_session=managed_session, - timestamp_s=timestamp_s, - event_type=event_type, - payload=payload, - ) - return - else: - raise RuntimeError( - "Too many queued WebRTC user events; wait for inference to catch up." - ) - managed_session.user_events.append( - UserInputEvent( - timestamp_s=timestamp_s, - event_type=event_type, - payload=payload, - source="webrtc", - ) - ) - - def _make_room_for_release_event( - self, - *, - managed_session: ManagedWebRTCSession, - event_type: str, - payload: dict[str, Any], - ) -> bool: - events = managed_session.user_events - if not events: - return False - if event_type == "key_up": - released_key = payload.get("key") - normalized_released_key = ( - normalize_key(released_key) if isinstance(released_key, str) else None - ) - if normalized_released_key is not None: - for index, queued_event in enumerate(events): - queued_key = queued_event.payload.get("key") - if ( - queued_event.event_type == "key_down" - and isinstance(queued_key, str) - and normalize_key(queued_key) == normalized_released_key - ): - del events[index] - return True - for index, queued_event in enumerate(events): - queued_key = queued_event.payload.get("key") - if ( - queued_event.event_type == "key_up" - and isinstance(queued_key, str) - and normalize_key(queued_key) == normalized_released_key - ): - del events[index] - return True - return False - - def _record_coalesced_release_event( - self, - *, - managed_session: ManagedWebRTCSession, - timestamp_s: float, - event_type: str, - payload: dict[str, Any], - ) -> None: - if event_type != "key_up": - return - key = payload.get("key") - if not isinstance(key, str): - return - managed_session.coalesced_release_events[normalize_key(key)] = UserInputEvent( - timestamp_s=timestamp_s, - event_type=event_type, - payload=payload, - source="webrtc", - ) - - def _supported_key_names(self) -> frozenset[str]: - supported_keys = self._resampler_supported_keys - if supported_keys is None: - supported_keys = DEFAULT_SUPPORTED_KEYS - return frozenset(normalize_key(key) for key in supported_keys) - - def _supports_key_payload(self, payload: dict[str, Any]) -> bool: - key = payload.get("key") - return isinstance(key, str) and normalize_key(key) in self._supported_key_names() - - @staticmethod - def _pending_user_events( - managed_session: ManagedWebRTCSession, - ) -> tuple[UserInputEvent, ...]: - return tuple( - sorted( - ( - *managed_session.user_events, - *managed_session.coalesced_release_events.values(), - ), - key=lambda event: event.timestamp_s, - ) - ) - - def _catch_up_input_clock( - self, - *, - managed_session: ManagedWebRTCSession, - now: float, - chunk_duration: float, - ) -> None: - """Skip stale input windows without skipping session input state.""" - resampler = managed_session.resampler - lag = now - (resampler.next_chunk_start_v + chunk_duration) - if lag <= chunk_duration: - return - latest_chunk_start = now - chunk_duration - if managed_session.inference_session is not None: - catch_up_start = ( - 0.0 - if managed_session.session_steps_completed == 0 - else resampler.next_chunk_start_v - ) - if latest_chunk_start > catch_up_start: - self._advance_inference_input_state( - managed_session=managed_session, - window=TimeWindow( - start_s=catch_up_start, - end_s=latest_chunk_start, - ), - ) - resampler.next_chunk_start_v = latest_chunk_start - - def _advance_inference_input_state( - self, - *, - managed_session: ManagedWebRTCSession, - window: TimeWindow, - ) -> None: - """Advance session input converters over a skipped raw-event window.""" - if managed_session.inference_session is None or window.end_s <= window.start_s: - return - runtime = managed_session.runtime - runtime.input_canonicalizer.canonicalize( - UserInputs(events=self._pending_user_events(managed_session)), - window=window, - source_schema=runtime.input_source_schema, - ) - managed_session.session_input_state_advanced = True - self._prune_consumed_user_events( - managed_session, - before_s=window.end_s, - ) - - def _validate_user_event_payload( - self, - *, - managed_session: ManagedWebRTCSession, - event_type: str, - payload: dict[str, Any], - ) -> dict[str, Any]: - """Return a runtime-validated user-event payload.""" - validate = getattr(managed_session.runtime, "validate_user_event", None) - if not callable(validate): - return payload - result = validate(event_type=event_type, payload=dict(payload)) - if result is None: - return payload - if not isinstance(result, dict): - raise TypeError( - "validate_user_event must return a payload dict or None, got " - f"{type(result).__name__}." - ) - return result - - @staticmethod - def _prune_consumed_user_events( - managed_session: ManagedWebRTCSession, *, before_s: float - ) -> None: - """Drop events already folded into converter state. - - Converters are level-triggered and carry their own state across - windows, so an event older than the current window start cannot affect - any future window and would otherwise grow the buffer without bound. - """ - events = managed_session.user_events - while events and events[0].timestamp_s < before_s: - events.popleft() - for key, event in tuple(managed_session.coalesced_release_events.items()): - if event.timestamp_s < before_s: - del managed_session.coalesced_release_events[key] - - async def _step_inference_session( - self, - *, - managed_session: ManagedWebRTCSession, - window: TimeWindow, - ) -> VideoStepResult: - """Map this chunk's events into model inputs and run one session step.""" - session: Any = managed_session.inference_session - if session is None: - raise RuntimeError("Session branch invoked without an inference session.") - request = session.next_step_request() - if request is None: - raise RuntimeError("Inference session reported no further steps.") - # The transport owns input windowing. The session derives its own - # window from its frame counter, but live events are stamped on the - # manager's monotonic clock, so the manager's window wins. - if request.step_index == 0 and not managed_session.session_input_state_advanced: - # The resampler's clock is re-anchored to "now" at first - # interaction, but events that triggered it were stamped just - # before that anchor. Widening chunk 0 back to the session start - # keeps them in the first window; otherwise a text event that - # itself started generation would be dropped, since converters - # never see a window that has already passed. - window = TimeWindow(start_s=0.0, end_s=window.end_s) - request = replace(request, user_input_window=window) - step_inputs = self._build_step_inputs( - managed_session=managed_session, - request=request, - window=window, - ) - loop = asyncio.get_running_loop() - result = await loop.run_in_executor(None, session.step, step_inputs) - self._prune_consumed_user_events(managed_session, before_s=window.start_s) - output = result.output - if not isinstance(output, VideoStepResult): - raise TypeError( - "WebRTC session steps must produce VideoStepResult output, got " - f"{type(output).__name__}." - ) - managed_session.session_steps_completed += 1 - return output - - def _build_step_inputs( - self, - *, - managed_session: ManagedWebRTCSession, - request: Any, - window: TimeWindow, - ) -> InferenceInput: - """Canonicalize this chunk's events and map them into model inputs.""" - runtime = managed_session.runtime - canonical_inputs = runtime.input_canonicalizer.canonicalize( - UserInputs(events=self._pending_user_events(managed_session)), - window=window, - source_schema=runtime.input_source_schema, - ) - return runtime.input_mapping.map_step_inputs( - canonical_inputs=canonical_inputs, - inference_input=InferenceInput(), - request=request, - ) - async def _handle_event_message( self, *, @@ -769,15 +490,15 @@ async def create_answer(self, *, offer_sdp: str, offer_type: str) -> dict[str, s async with self._session_lock: if self._active_session is not None and not self._active_session.closed: - raise SessionBusyError(self._busy_message) + raise SessionBusyError(self.busy_message) - session_input = self._peek_pending_session_input() + session_input = self._pending_session_input answer = await self._create_answer_with_runtime_ready_locked( offer_sdp=offer_sdp, offer_type=offer_type, session_input=session_input, ) - self._clear_pending_session_input() + self._pending_session_input = None return answer async def _create_answer_with_runtime_ready_locked( @@ -790,11 +511,11 @@ async def _create_answer_with_runtime_ready_locked( enable_liveness_watchdog: bool = True, ) -> dict[str, str]: if self._active_session is not None and not self._active_session.closed: - raise SessionBusyError(self._busy_message) + raise SessionBusyError(self.busy_message) if not self._runtime_ready: - raise self._runtime_error_types[0]("Runtime is not initialized.") + raise RuntimeError("Runtime is not initialized.") - await self._reset_runtime_for_session(session_input) + await self._runtime.reset_for_new_session(session_input=session_input) peer_connection = RTCPeerConnection(rtc_configuration) # Bounded queue sized to one *steady-state* chunk so the producer @@ -878,11 +599,26 @@ async def on_connectionstatechange() -> None: }: await self.close_active_session() - self._register_extra_peer_handlers(peer_connection) + @peer_connection.on("iceconnectionstatechange") + def on_iceconnectionstatechange() -> None: + logger.info( + "Peer ICE connection state changed: {}", + peer_connection.iceConnectionState, + ) + + @peer_connection.on("icegatheringstatechange") + def on_icegatheringstatechange() -> None: + logger.debug( + "Peer ICE gathering state changed: {}", + peer_connection.iceGatheringState, + ) try: offer = RTCSessionDescription(sdp=offer_sdp, type=offer_type) - self._on_offer_received(offer_sdp) + logger.info( + "Received WebRTC offer with {}.", + _summarize_sdp_candidates(offer_sdp), + ) await peer_connection.setRemoteDescription(offer) answer = await peer_connection.createAnswer() await peer_connection.setLocalDescription(answer) @@ -896,7 +632,10 @@ async def on_connectionstatechange() -> None: local_description = peer_connection.localDescription if local_description is None: raise RuntimeError("Peer connection did not produce local description.") - self._on_answer_created(local_description.sdp) + logger.info( + "Created WebRTC answer with {}.", + _summarize_sdp_candidates(local_description.sdp), + ) return {"sdp": local_description.sdp, "type": local_description.type} except Exception: logger.exception("WebRTC negotiation failed while creating an answer.") @@ -906,13 +645,13 @@ async def on_connectionstatechange() -> None: async def _run_loopback_warmup_session(self, *, num_chunks: int) -> None: if not self._runtime_ready: - raise self._runtime_error_types[0]("Runtime is not initialized.") + raise RuntimeError("Runtime is not initialized.") await run_loopback_warmup_session( num_chunks=num_chunks, warmup_timeout_s=self.runtime_config.warmup_timeout_s, create_answer=self._create_loopback_warmup_answer, close_active_session=self.close_active_session, - label=self._warmup_label, + label=self.warmup_label, logger=logger, ) @@ -1085,7 +824,7 @@ async def _generation_worker( piecewise-constant timeline, hands segments and frame times to the runtime, and pushes the generated frames into the video track. The track's bounded queue then paces the loop to playback via - backpressure on ``BufferedVideoTrack.enqueue_chunk``. + backpressure on ``BufferedVideoTrack.enqueue_result``. """ loop = asyncio.get_running_loop() runtime = managed_session.runtime @@ -1117,8 +856,8 @@ async def _generation_worker( try: while not managed_session.closed: try: - input_num_frames = self._runtime_next_input_num_frames(runtime) - except self._runtime_error_types: + request, input_num_frames = self._runtime_next_step_request(runtime) + except RuntimeError: logger.exception("Runtime not ready; stopping generation worker.") return # Trigger when wallclock reaches the chunk's window end. @@ -1144,10 +883,15 @@ async def _generation_worker( t_before_gen = loop.time() chunk_start_v = resampler.next_chunk_start_v - # Sampled on both branches: the resampler owns the virtual - # clock, so it must advance even when its segments are unused. segments, frame_times = resampler.sample_chunk(input_num_frames) chunk_end_v = resampler.next_chunk_start_v + request = replace( + request, + user_input_window=TimeWindow( + start_s=chunk_start_v, + end_s=chunk_end_v, + ), + ) consumed_action_arrivals: list[float] = [] while ( managed_session.pending_action_arrivals @@ -1157,29 +901,27 @@ async def _generation_worker( managed_session.pending_action_arrivals.popleft() ) try: - if managed_session.inference_session is not None: - result = await self._step_inference_session( - managed_session=managed_session, - window=TimeWindow( - start_s=chunk_start_v, end_s=chunk_end_v - ), - ) - else: - result = await runtime.generate_chunk( - segments=segments, frame_times=frame_times + result = await runtime.step( + request=request, segments=segments, frame_times=frame_times + ) + if result.step_index != request.step_index: + raise RuntimeError( + "Runtime result step does not match its request: " + f"requested {request.step_index}, " + f"got {result.step_index}." ) except Exception as exc: logger.exception("Chunk generation failed.") channel = managed_session.control_channel if channel is not None: self._send_json(channel, make_error_payload(str(exc))) - if self._close_session_on_generation_error: + if self.fatal_generation_errors: await self.close_active_session() return continue t_after_gen = loop.time() delivery = await video_encoder.deliver_chunk( - result.video_chunk, + result, video_track, force_keyframe=False, ) @@ -1188,7 +930,7 @@ async def _generation_worker( gen_ms = (t_after_gen - t_before_gen) * 1e3 enqueue_ms = (t_after_enqueue - t_after_gen) * 1e3 - play_ms = result.num_frames * 1000.0 / video_track.fps + play_ms = result.frame_count * 1000.0 / video_track.fps lag_ms = (t_after_enqueue - resampler.next_chunk_start_v) * 1e3 control_latency_ms = ( (t_after_enqueue - consumed_action_arrivals[0]) * 1e3 @@ -1196,17 +938,16 @@ async def _generation_worker( else None ) perf_window_chunks += 1 - perf_window_frames += result.num_frames - if result.chunk_index == 0 or ( - perf_log_interval > 0 - and result.chunk_index % perf_log_interval == 0 + perf_window_frames += result.frame_count + if result.step_index == 0 or ( + perf_log_interval > 0 and result.step_index % perf_log_interval == 0 ): interval_s = max(t_after_enqueue - perf_window_start, 1.0e-6) interval_fps = perf_window_frames / interval_s - gen_fps = result.num_frames / max( + gen_fps = result.frame_count / max( t_after_gen - t_before_gen, 1.0e-6 ) - stats = result.stats or {} + stats = result.metrics logger.info( "WebRTC perf chunk={} interval_chunks={} frames={} " "gen_fps={:.1f} interval_fps={:.1f} playback_fps={} " @@ -1217,7 +958,7 @@ async def _generation_worker( "queue_depth={} lag_ms={:.0f} control_latency_ms={} " "compile_active={} compile_start_step={} cuda_graph={} " "cache_frames={} cache_tokens={}", - result.chunk_index, + result.step_index, perf_window_chunks, perf_window_frames, gen_fps, @@ -1252,9 +993,9 @@ async def _generation_worker( "segments={} enqueued={} " "gen_ms={:.1f} enqueue_ms={:.1f} play_ms={:.1f} queue_depth={} " "lag_ms={:.1f}", - result.chunk_index, + result.step_index, input_num_frames, - result.num_frames, + result.frame_count, len(segments), enqueued, gen_ms, @@ -1269,13 +1010,13 @@ async def _generation_worker( self._send_json( channel, make_chunk_done_payload( - chunk_index=result.chunk_index, - num_frames=result.num_frames, + chunk_index=result.step_index, + num_frames=result.frame_count, enqueued_frames=enqueued, fps=video_track.fps, width=self.runtime_config.video_width, height=self.runtime_config.video_height, - model=self._model_name(), + model=self.identity, gen_ms=gen_ms, enqueue_ms=enqueue_ms, play_ms=play_ms, @@ -1283,7 +1024,7 @@ async def _generation_worker( lag_ms=lag_ms, control_latency_ms=control_latency_ms, consumed_actions=len(consumed_action_arrivals), - extra=self._chunk_done_extra(), + extra=result.metadata, ), ) except asyncio.CancelledError: diff --git a/flashdreams/flashdreams/serving/webrtc/media.py b/flashdreams/flashdreams/serving/webrtc/media.py index 2912a042f..470569132 100644 --- a/flashdreams/flashdreams/serving/webrtc/media.py +++ b/flashdreams/flashdreams/serving/webrtc/media.py @@ -7,7 +7,7 @@ import contextlib from collections.abc import Callable, Sequence from fractions import Fraction -from typing import TYPE_CHECKING +from typing import cast import numpy as np from aiortc import MediaStreamTrack @@ -16,17 +16,31 @@ from av.packet import Packet from loguru import logger -from flashdreams.serving.realtime.media import tensor_chunk_to_rgb_frames - -if TYPE_CHECKING: - import torch +from flashdreams.runtime import StepResult +from flashdreams.serving.realtime.media import ( + FrameLayout, + ValueRange, + rgb_array_to_uint8_frames, +) +from flashdreams.serving.realtime.media import ( + tensor_chunk_to_rgb_frames as tensor_chunk_to_rgb_frames, +) _STALL_THRESHOLD_MS = 1.0 _PACING_LAG_LOG_MS = 5.0 -def _default_frame_converter(video_chunk: torch.Tensor) -> list[np.ndarray]: - return tensor_chunk_to_rgb_frames(video_chunk, sync_device=True) +def _default_frame_converter(result: StepResult) -> list[np.ndarray]: + video_chunk = result.video_chunk + value_range: ValueRange = ( + "minus_one_one" if video_chunk.is_floating_point() else "uint8" + ) + return rgb_array_to_uint8_frames( + video_chunk, + layout=cast(FrameLayout, result.layout), + value_range=value_range, + sync_device=True, + ) class BufferedVideoTrack(MediaStreamTrack): @@ -39,7 +53,7 @@ def __init__( *, fps: int, maxsize: int, - frame_converter: Callable[[torch.Tensor], list[np.ndarray]] | None = None, + frame_converter: Callable[[StepResult], list[np.ndarray]] | None = None, ) -> None: super().__init__() if fps <= 0: @@ -67,10 +81,10 @@ def maxsize(self) -> int: def qsize(self) -> int: return self._frames.qsize() - async def enqueue_chunk(self, video_chunk: torch.Tensor) -> int: + async def enqueue_result(self, result: StepResult) -> int: if self._closed: return 0 - frames = await asyncio.to_thread(self._frame_converter, video_chunk) + frames = await asyncio.to_thread(self._frame_converter, result) for i, frame in enumerate(frames): if self._closed: return i diff --git a/flashdreams/flashdreams/serving/webrtc/nvenc.py b/flashdreams/flashdreams/serving/webrtc/nvenc.py index ade33cb49..3a4cc93dc 100644 --- a/flashdreams/flashdreams/serving/webrtc/nvenc.py +++ b/flashdreams/flashdreams/serving/webrtc/nvenc.py @@ -31,6 +31,7 @@ from av.packet import Packet from loguru import logger +from flashdreams.runtime import StepResult from flashdreams.serving.webrtc.encoders import ChunkDeliveryResult # Runtime imports ``PyNvVideoCodec`` unconditionally (the isolation @@ -71,13 +72,12 @@ def _payload_contains_nal_type(payload: bytes, nal_type: int) -> bool: 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. +def _result_to_abgr_frames(result: StepResult) -> torch.Tensor: + """Convert a declared video result to NVENC-``ABGR``-formatted 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. + The result layout selects the time, channel, batch, and view axes; tensor + rank is never used to guess the model's output contract. The returned + contiguous ``[T, H, W, 4]`` uint8 tensor stays on the source device. **NVENC ``NV_ENC_BUFFER_FORMAT_ABGR`` is a word-ordered token, not memory-ordered.** From ``nvEncodeAPI.h``: "a pixel is represented by @@ -93,25 +93,7 @@ def _chunk_to_abgr_cuda_frames(chunk: torch.Tensor) -> torch.Tensor: 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) + rgb = result.video_hwc_uint8() 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] → @@ -260,7 +242,7 @@ def create_track(self, *, maxsize: int) -> NVENCVideoTrack: async def deliver_chunk( self, - chunk: torch.Tensor, + result: StepResult, track: MediaStreamTrack, *, force_keyframe: bool = False, @@ -297,7 +279,7 @@ def _stream(packet: Packet) -> None: _num_frames, num_keyframes, encode_ms = await asyncio.to_thread( self.encode_chunk_sync, - chunk, + result, force_keyframe=force_keyframe, on_packet=_stream, ) @@ -317,18 +299,20 @@ def _stream(packet: Packet) -> None: def encode_chunk_sync( self, - chunk: torch.Tensor, + result: StepResult, *, 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)``. + """Encode a result and return frame, keyframe, and timing counts. 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) + frames = _result_to_abgr_frames(result) + if not frames.is_cuda: + raise ValueError("expected CUDA tensor for hardware encode path") num_frames = frames.shape[0] num_keyframes = 0 start_s = time.perf_counter() diff --git a/flashdreams/flashdreams/serving/webrtc/runtime.py b/flashdreams/flashdreams/serving/webrtc/runtime.py index d402657e1..d80141f6b 100644 --- a/flashdreams/flashdreams/serving/webrtc/runtime.py +++ b/flashdreams/flashdreams/serving/webrtc/runtime.py @@ -1,19 +1,42 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Runtime contracts for shared WebRTC demo serving.""" +"""Runtime contracts and thread-affine execution for shared WebRTC serving.""" from __future__ import annotations +import asyncio +from abc import ABC, abstractmethod from collections.abc import Awaitable -from typing import Any, Protocol - -from flashdreams.infra.video_output import VideoStepResult -from flashdreams.runtime.canonical import InputCanonicalizer -from flashdreams.runtime.inputs import UserInputSchema -from flashdreams.runtime.interfaces import InferenceSession -from flashdreams.runtime.mapping import InputMapping +from enum import IntEnum +from typing import Any, Generic, Protocol, TypeVar + +import torch +import torch.distributed as dist + +from flashdreams.core.distributed.rank_orchestration import ( + RankCoordinator, + distributed_op, +) +from flashdreams.runtime.types import StepRequest, StepResult +from flashdreams.runtime.worker import ThreadAffineRuntimeWorker from flashdreams.serving.realtime.input import PoseSegment +from flashdreams.serving.webrtc.encoders import ( + EncoderBackend, + VideoEncoder, + select_encoder, +) + + +class WebRTCControlSignal(IntEnum): + """Rank-orchestration signals shared by WebRTC runtimes.""" + + INITIALIZE = 0 + RESET_SESSION = 1 + ACTION_STEP = 2 + CLOSE = 3 + EVENT = 4 + EXIT = 99 class WebRTCRuntimeConfig(Protocol): @@ -25,37 +48,49 @@ class WebRTCRuntimeConfig(Protocol): warmup_timeout_s: float -class WebRTCGenerationRuntime(Protocol): - """Generation lifecycle for one shared WebRTC session. +class ThreadAffineWebRTCRuntimeConfig(WebRTCRuntimeConfig, Protocol): + """Configuration consumed by the shared runtime execution layer.""" + + device: str + fps: int + encoder_backend: EncoderBackend + encoder_bitrate_bps: int + encoder_gop: int + + +class WebRTCServerLifecycle(Protocol): + """Distributed worker lifecycle used by the shared WebRTC serve loop.""" + + def send_exit_signal(self) -> None: ... + + def wait_for_termination(self) -> None: ... + + +class WebRTCSessionRuntime(WebRTCServerLifecycle, Protocol): + """Complete runtime contract consumed by the shared session manager. Integrations keep their model-specific state, checkpoints, conditioning, and cache logic inside their concrete runtime. The shared manager only needs this lifecycle and chunk-generation surface. - - By default, ``peek_next_chunk_num_frames`` and - ``peek_steady_chunk_num_frames`` are used for both input sampling and - output queue sizing. Runtimes whose model input clock differs from their - output video clock may also implement these optional methods: - - - ``peek_input_fps() -> float`` for the control/input sampling clock. - - ``peek_next_input_num_frames() -> int`` for the length of ``frame_times``. - - ``peek_steady_output_num_frames() -> int`` for video queue sizing. """ async def initialize(self) -> None: ... - async def reset_for_new_session(self) -> None: ... + async def reset_for_new_session(self, *, session_input: Any = None) -> None: ... + + def peek_input_fps(self) -> float: ... - def peek_steady_chunk_num_frames(self) -> int: ... + def next_step_request(self) -> StepRequest: ... - def peek_next_chunk_num_frames(self) -> int: ... + def peek_steady_output_num_frames(self) -> int: ... - async def generate_chunk( + async def step( self, *, + request: StepRequest, segments: list[PoseSegment], frame_times: list[float], - ) -> VideoStepResult: ... + ) -> StepResult: ... async def close(self) -> None: ... @@ -68,40 +103,229 @@ def trigger_event( ) -> dict[str, Any] | Awaitable[dict[str, Any]]: ... -class WebRTCInferenceSessionRuntime(Protocol): - """Optional runtime capability for driving an ``InferenceSession``. +_ConfigT = TypeVar("_ConfigT", bound=ThreadAffineWebRTCRuntimeConfig) +_SessionInputT = TypeVar("_SessionInputT") - A runtime implementing this opts into the manager's session branch, where - raw key and text events are canonicalized and mapped into per-step - ``InferenceInput`` instead of being handed to ``generate_chunk`` as - pre-integrated pose segments. The transport keeps owning event - timestamping and input-window selection; the model only declares its - mapping and consumes model-facing inputs. - Runtimes on this branch do not need ``generate_chunk`` or ``trigger_event``: - camera control arrives as mapped step inputs, and text events arrive as a - session-global conditioning update in the same payload. +class ThreadAffineDistributedWebRTCRuntime( + ABC, + Generic[_ConfigT, _SessionInputT], +): + """Coordinate one thread-affine, distributed WebRTC model runtime. + + Subclasses own model construction, rollout state, conditioning, and chunk + generation. This base owns the identical async-to-thread dispatch, rank + signaling, step ordering, and video-encoder lifecycle used by integrations. """ - async def start_inference_session(self) -> InferenceSession: ... + MASTER_RANK = 0 - @property - def input_mapping(self) -> InputMapping: ... + def __init__( + self, + *, + config: _ConfigT, + runtime_error_type: type[RuntimeError], + thread_name: str, + ) -> None: + self.config = config + self.rank = 0 if not dist.is_initialized() else dist.get_rank() + self._runtime_error_type = runtime_error_type + self._device = self._resolve_device(config.device) + self._closed = False + self._video_encoder: VideoEncoder | None = None + self._worker = ThreadAffineRuntimeWorker( + device=self._device, + thread_name=thread_name, + ) + self._step_lock = asyncio.Lock() + self.rank_coordinator = RankCoordinator( + device=self._device, + signal_type=WebRTCControlSignal, + is_master=self.is_master, + master_rank=self.MASTER_RANK, + ) + self.rank_coordinator.register_distributed_ops(self) + + @staticmethod + def _resolve_device(device_spec: str | torch.device) -> torch.device: + device = torch.device(device_spec) + if device.type == "cuda" and device.index is None: + device = torch.device( + f"cuda:{torch.cuda.current_device()}" + if torch.cuda.is_available() + else "cuda:0" + ) + return device @property - def input_canonicalizer(self) -> InputCanonicalizer: ... + def is_master(self) -> bool: + return self.rank == self.MASTER_RANK @property - def input_source_schema(self) -> UserInputSchema: ... + def video_encoder(self) -> VideoEncoder: + """Return the encoder selected during runtime initialization.""" + if self._video_encoder is None: + raise self._runtime_error( + "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) + + def send_exit_signal(self) -> None: + if self.is_master: + self.rank_coordinator.send_exit(exit_signal=WebRTCControlSignal.EXIT) + + async def initialize(self) -> None: + if self._is_runtime_initialized(): + return + await self._worker.call(self._initialize_sync_all_ranks) + + async def reset_for_new_session( + self, session_input: _SessionInputT | None = None + ) -> None: + self._require_open_and_initialized() + await self._worker.call(self._reset_rollout_sync_all_ranks, session_input) + + async def close(self) -> None: + if self._closed: + return + self._closed = True + try: + await self._worker.call(self._close_sync_all_ranks) + finally: + await self._worker.close() + + async def step( + self, + *, + request: StepRequest, + segments: list[PoseSegment], + frame_times: list[float], + ) -> StepResult: + self._require_open_and_initialized(session=True) + expected_step = self._runtime_step_index() + if request.step_index != expected_step: + raise self._runtime_error( + f"Expected request step {expected_step}, got {request.step_index}." + ) + + async with self._step_lock: + self._require_open_and_initialized(session=True) + return await self._worker.call( + self._generate_chunk_sync_all_ranks, + segments, + frame_times, + ) + + def peek_input_fps(self) -> float: + return float(self.config.fps) + + def next_step_request(self) -> StepRequest: + self._require_open_and_initialized() + return StepRequest( + step_index=self._runtime_step_index(), + metadata={"input_frame_count": self._next_input_frame_count()}, + ) + + def peek_steady_output_num_frames(self) -> int: + self._require_open_and_initialized() + return self._steady_output_frame_count() + + def _runtime_error(self, message: str) -> RuntimeError: + return self._runtime_error_type(message) + + def _require_open_and_initialized(self, *, session: bool = False) -> None: + if self._closed: + noun = "Session" if session else "Runtime" + raise self._runtime_error(f"{noun} is closed.") + if not self._is_runtime_initialized(): + raise self._runtime_error("Runtime is not initialized.") + + def _initialize_video_encoder_sync(self) -> None: + """Select the master rank's encoder on the model runtime thread.""" + if not self.is_master: + return + if self._video_encoder is not None: + self._video_encoder.close() + self._video_encoder = None + + backend = self.config.encoder_backend + if self._device.type != "cuda" and backend == "auto": + backend = "default" + if self._device.type != "cuda" and backend == "nvenc": + raise self._runtime_error( + "encoder_backend='nvenc' requires a CUDA runtime device." + ) + gpu_id = self._device.index if self._device.index is not None else 0 + self._video_encoder = select_encoder( + backend=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 _close_video_encoder_sync(self) -> None: + if self._video_encoder is not None: + self._video_encoder.close() + self._video_encoder = None + + @distributed_op(WebRTCControlSignal.INITIALIZE) + def _initialize_sync_all_ranks(self) -> None: + self._initialize_sync() + + @distributed_op(WebRTCControlSignal.RESET_SESSION) + def _reset_rollout_sync_all_ranks( + self, session_input: _SessionInputT | None = None + ) -> None: + self._reset_rollout_sync(session_input=session_input) + + @distributed_op(WebRTCControlSignal.ACTION_STEP) + def _generate_chunk_sync_all_ranks( + self, + segments: list[PoseSegment], + frame_times: list[float], + ) -> StepResult: + return self._generate_one_chunk_sync(segments=segments, frame_times=frame_times) + @distributed_op(WebRTCControlSignal.CLOSE) + def _close_sync_all_ranks(self) -> None: + try: + self._close_sync() + finally: + self._close_video_encoder_sync() -class WebRTCServerLifecycle(Protocol): - """Distributed worker lifecycle used by the shared WebRTC serve loop.""" + @abstractmethod + def _is_runtime_initialized(self) -> bool: ... - def send_exit_signal(self) -> None: ... + @abstractmethod + def _runtime_step_index(self) -> int: ... - def wait_for_termination(self) -> None: ... + @abstractmethod + def _next_input_frame_count(self) -> int: ... + + @abstractmethod + def _steady_output_frame_count(self) -> int: ... + + @abstractmethod + def _initialize_sync(self) -> None: ... + @abstractmethod + def _reset_rollout_sync( + self, session_input: _SessionInputT | None = None + ) -> None: ... + + @abstractmethod + def _generate_one_chunk_sync( + self, + *, + segments: list[PoseSegment], + frame_times: list[float], + ) -> StepResult: ... -class WebRTCSessionRuntime(WebRTCGenerationRuntime, WebRTCServerLifecycle, Protocol): - """Complete runtime contract consumed by the shared session manager.""" + @abstractmethod + def _close_sync(self) -> None: ... diff --git a/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py b/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py index c11c9cee3..735179a4a 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py +++ b/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py @@ -20,7 +20,10 @@ from functools import partial from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer from importlib.resources import as_file, files +from os import PathLike from pathlib import Path +from socket import socket +from socketserver import BaseServer from urllib.parse import urlsplit WEB_DIR_RESOURCE = files("flashdreams.serving.webrtc").joinpath("web") @@ -31,12 +34,20 @@ class MockUIRequestHandler(SimpleHTTPRequestHandler): def __init__( self, - *args: object, + request: socket | tuple[bytes, socket], + client_address: tuple[str, int], + server: BaseServer, + *, + directory: str | PathLike[str] | None = None, model_web_dir: Path | None = None, - **kwargs: object, ) -> None: self.model_web_dir = model_web_dir - super().__init__(*args, **kwargs) + super().__init__( + request, + client_address, + server, + directory=directory, + ) def _rewrite_path(self) -> bool: path = urlsplit(self.path).path diff --git a/flashdreams/flashdreams/serving/webrtc/web/request_session.js b/flashdreams/flashdreams/serving/webrtc/web/request_session.js index 8f708f0ba..95c340b2d 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/request_session.js +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.js @@ -3,6 +3,20 @@ const mockMode = new URLSearchParams(window.location.search).has("mock") +/** + * @typedef {Object} WebRTCModelAdapter + * @property {string=} modelName + * @property {string=} stylesheet + * @property {Array<{label: string, keys: Array}>=} controls + * @property {{postprocess?: boolean}=} capabilities + * @property {(context: Object) => (void|Promise)=} mount + * @property {(context: Object) => (void|Promise)=} beforeConnect + * @property {(action: Object, context: Object) => void=} onActionSent + * @property {(payload: Object, context: Object) => boolean=} onControlMessage + * @property {(visible: boolean, context: Object) => void=} onVideoVisibilityChanged + * @property {(context: Object) => void=} onDisconnect + */ + const connectButton = document.getElementById("connectButton") const statusText = document.getElementById("statusText") const flowText = document.getElementById("flowText") @@ -23,17 +37,6 @@ const modelPanelSlot = document.getElementById("modelPanelSlot") const modelControlSlot = document.getElementById("modelControlSlot") const controlRows = document.getElementById("controlRows") -const defaultControls = [ - { - label: "Drive / Turn", - keys: [ - { key: "w", label: "Forward" }, - { key: "a", label: "Turn left" }, - { key: "s", label: "Backward" }, - { key: "d", label: "Turn right" }, - ], - }, -] const keyAliases = new Map([ ["arrowup", "w"], ["arrowleft", "a"], @@ -50,6 +53,7 @@ const heartbeatIntervalMs = 2000 let allowedKeys = new Set() let controlButtons = [] +/** @type {WebRTCModelAdapter|null} */ let modelAdapter = null let peerConnection = null @@ -310,11 +314,11 @@ async function loadModelAdapter() { document.head.append(stylesheet) } const modelControls = Array.isArray(adapter.controls) ? adapter.controls : [] - renderControls([...defaultControls, ...modelControls]) + renderControls(modelControls) if (typeof adapter.modelName === "string") { modelContext.setModelName(adapter.modelName) } - if (adapter.enablePostprocess === true) { + if (adapter.capabilities?.postprocess === true) { try { await loadPostprocessOptions() } catch (error) { diff --git a/flashdreams/tests/test_encoders.py b/flashdreams/tests/test_encoders.py index 20d2183c6..610217807 100644 --- a/flashdreams/tests/test_encoders.py +++ b/flashdreams/tests/test_encoders.py @@ -35,11 +35,14 @@ from types import ModuleType, SimpleNamespace from unittest.mock import MagicMock, patch +import numpy as np import pytest +import torch from av.packet import Packet pytestmark = pytest.mark.ci_cpu +from flashdreams.runtime import StepResult from flashdreams.serving.webrtc import encoders as enc_mod from flashdreams.serving.webrtc.encoders import ( ChunkDeliveryResult, @@ -357,7 +360,7 @@ def test_is_frozen_dataclass(self) -> None: # --------------------------------------------------------------------------- -# DefaultRTCEncoder.deliver_chunk delegates to track.enqueue_chunk +# DefaultRTCEncoder.deliver_chunk delegates to track.enqueue_result # --------------------------------------------------------------------------- @@ -367,37 +370,100 @@ class _FakeBufferedVideoTrack: aiortc runtime dependencies without breaking the isinstance check).""" def __init__(self) -> None: - self.enqueued_chunks: list = [] + self.enqueued_results: list[StepResult] = [] - async def enqueue_chunk(self, chunk) -> int: - self.enqueued_chunks.append(chunk) - return 4 + async def enqueue_result(self, result: StepResult) -> int: + self.enqueued_results.append(result) + return result.frame_count class TestDefaultRTCEncoderDeliver: + @pytest.mark.parametrize( + ("layout", "shape"), + [("tchw", (4, 3, 8, 8)), ("bvtchw", (1, 1, 4, 3, 8, 8))], + ) @pytest.mark.asyncio - async def test_deliver_chunk_returns_frames_from_track(self) -> None: + async def test_deliver_chunk_returns_frames_from_track( + self, layout: str, shape: tuple[int, ...] + ) -> None: from flashdreams.serving.webrtc import media as media_mod fake_track = _FakeBufferedVideoTrack() + step_result = StepResult.from_video_chunk( + step_index=0, + video_chunk=torch.zeros(shape, dtype=torch.uint8), + layout=layout, # ty:ignore[invalid-argument-type] + ) # 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] + step_result, 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 + assert fake_track.enqueued_results == [step_result] + + @pytest.mark.parametrize( + ("layout", "shape"), + [("tchw", (3, 3, 2, 2)), ("bvtchw", (1, 1, 3, 3, 2, 2))], + ) + @pytest.mark.asyncio + async def test_software_conversion_uses_declared_layout( + self, layout: str, shape: tuple[int, ...] + ) -> None: + enc = DefaultRTCEncoder(fps=30) + track = enc.create_track(maxsize=3) + step_result = StepResult.from_video_chunk( + step_index=0, + video_chunk=torch.zeros(shape, dtype=torch.uint8), + layout=layout, # ty:ignore[invalid-argument-type] + ) + + delivery = await enc.deliver_chunk(step_result, track) + + assert delivery.num_frames == 3 + assert track.qsize() == 3 + await track.close() + + @pytest.mark.asyncio + async def test_software_path_defers_host_conversion_to_track(self) -> None: + from flashdreams.serving.webrtc.media import BufferedVideoTrack + + source = torch.zeros((2, 3, 2, 2), dtype=torch.uint8) + step_result = StepResult.from_video_chunk( + step_index=0, + video_chunk=source, + layout="tchw", + ) + seen: list[StepResult] = [] + + def _converter(delivered: StepResult) -> list[np.ndarray]: + seen.append(delivered) + assert delivered is step_result + assert delivered.video_chunk.data_ptr() == source.data_ptr() + return [np.zeros((2, 2, 3), dtype=np.uint8) for _ in range(2)] + + track = BufferedVideoTrack(fps=30, maxsize=2, frame_converter=_converter) + delivery = await DefaultRTCEncoder(fps=30).deliver_chunk(step_result, track) + + assert delivery.num_frames == 2 + assert seen == [step_result] + await track.close() @pytest.mark.asyncio async def test_deliver_chunk_rejects_wrong_track_type(self) -> None: enc = DefaultRTCEncoder(fps=30) + step_result = StepResult.from_video_chunk( + step_index=0, + video_chunk=torch.zeros((1, 3, 2, 2), dtype=torch.uint8), + layout="tchw", + ) with pytest.raises(TypeError, match="BufferedVideoTrack"): await enc.deliver_chunk( - SimpleNamespace(), # ty:ignore[invalid-argument-type] + step_result, SimpleNamespace(), # ty:ignore[invalid-argument-type] ) @@ -429,6 +495,35 @@ def test_aiortc_sender_module_importable(self) -> None: # break the runtime; catch it here before the first RTP packet. import aiortc.rtcrtpsender # noqa: F401 + +class TestNvencResultConversion: + @pytest.mark.parametrize( + ("layout", "shape"), + [("tchw", (2, 3, 2, 3)), ("bvtchw", (1, 1, 2, 3, 2, 3))], + ) + def test_conversion_uses_declared_layout( + self, + monkeypatch: pytest.MonkeyPatch, + layout: str, + shape: tuple[int, ...], + ) -> None: + nvenc_mod = _install_fake_nvc(monkeypatch, MagicMock()) + video = torch.empty(shape, dtype=torch.uint8) + channel_dim = 1 if layout == "tchw" else 3 + video.select(channel_dim, 0).fill_(10) + video.select(channel_dim, 1).fill_(20) + video.select(channel_dim, 2).fill_(30) + result = StepResult.from_video_chunk( + step_index=0, + video_chunk=video, + layout=layout, # ty:ignore[invalid-argument-type] + ) + + frames = nvenc_mod._result_to_abgr_frames(result) + + assert frames.shape == (2, 2, 3, 4) + assert torch.equal(frames[0, 0, 0], torch.tensor([10, 20, 30, 255])) + def test_getencodercaps_callable_when_library_available(self) -> None: # This guard exercises the *real* PyNvVideoCodec surface via # ``nvenc``. ``PyNvVideoCodec`` raises ``RuntimeError`` (not @@ -480,12 +575,12 @@ def create_track(self, *, maxsize: int) -> NVENCVideoTrack: async def deliver_chunk( self, - chunk: object, + result: StepResult, track: NVENCVideoTrack, *, force_keyframe: bool = False, ) -> ChunkDeliveryResult: - del chunk, force_keyframe + del result, force_keyframe loop = asyncio.get_running_loop() frames = self._frames_per_chunk @@ -573,12 +668,12 @@ def _packet(pts: int) -> Packet: return packet def _fake_encode_chunk_sync( - chunk: object, + result: StepResult, *, force_keyframe: bool = False, on_packet: Callable[[Packet], None] | None = None, ) -> tuple[int, int, float]: - del chunk, force_keyframe + del result, force_keyframe assert on_packet is not None on_packet(_packet(0)) loop.call_soon_threadsafe(first_packet_enqueued.set) @@ -591,7 +686,11 @@ def _fake_encode_chunk_sync( track = NVENCVideoTrack(fps=_ORDERING_FPS, maxsize=4) deliver_task = asyncio.create_task( encoder.deliver_chunk( - object(), + StepResult.from_video_chunk( + step_index=0, + video_chunk=torch.zeros((2, 3, 2, 2)), + layout="tchw", + ), track, ) ) @@ -622,8 +721,15 @@ async def test_sequential_await_produces_monotonic_pts(self) -> None: ) track = encoder.create_track(maxsize=_ORDERING_TOTAL_FRAMES) - for _ in range(_ORDERING_NUM_CHUNKS): - await encoder.deliver_chunk(object(), track) + for step_index in range(_ORDERING_NUM_CHUNKS): + await encoder.deliver_chunk( + StepResult.from_video_chunk( + step_index=step_index, + video_chunk=torch.zeros((_ORDERING_FRAMES_PER_CHUNK, 3, 1, 1)), + layout="tchw", + ), + track, + ) seen_pts: list[int] = [] for _ in range(_ORDERING_TOTAL_FRAMES): @@ -664,8 +770,17 @@ async def test_fire_and_forget_pattern_would_break_ordering(self) -> None: # 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) + asyncio.create_task( + encoder.deliver_chunk( + StepResult.from_video_chunk( + step_index=step_index, + video_chunk=torch.zeros((_ORDERING_FRAMES_PER_CHUNK, 3, 1, 1)), + layout="tchw", + ), + track, + ) + ) + for step_index in range(_ORDERING_NUM_CHUNKS) ] await asyncio.gather(*tasks) diff --git a/flashdreams/tests/test_output_targets.py b/flashdreams/tests/test_output_targets.py index be4e11fbc..1421ceb62 100644 --- a/flashdreams/tests/test_output_targets.py +++ b/flashdreams/tests/test_output_targets.py @@ -31,6 +31,11 @@ def _runner_config( num_views: int = 1, compile_network: bool = True, ) -> RunnerConfig: + output_adapter = None + if runner_name.startswith("lingbot-world"): + output_adapter = "lingbot.output_targets:OUTPUT_TARGET_ADAPTER" + elif runner_name.startswith("omnidreams-"): + output_adapter = "omnidreams.output_targets:OUTPUT_TARGET_ADAPTER" transformer = SimpleNamespace( num_views=num_views, compile_network=compile_network, @@ -43,6 +48,7 @@ def _runner_config( RunnerConfig, SimpleNamespace( runner_name=runner_name, + output_adapter=output_adapter, pipeline=pipeline, device="cuda:1", pixel_height=480, @@ -71,7 +77,7 @@ def test_lingbot_webrtc_target_translates_runner_config() -> None: ), ) - assert spec.module == "lingbot.demo.cli" + assert spec.module == "lingbot.demo.app" assert spec.argv == ( "webrtc", "--preset-id", @@ -109,6 +115,74 @@ def test_omnidreams_webrtc_target_rejects_multi_view() -> None: ) +def test_omnidreams_webrtc_target_uses_shared_demo_entry_point() -> None: + config = _runner_config( + runner_name="omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae", + ) + + spec = resolve_output_target( + config, + mode="webrtc", + options=OutputLaunchOptions( + host="127.0.0.1", + port=9011, + prefer_sw_encoder=True, + ), + ) + + assert spec.module == "omnidreams.demo.app" + assert spec.argv == ( + "webrtc", + "--preset-id", + "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae", + "--device", + "cuda:1", + "--fps", + "24", + "--video-height", + "480", + "--video-width", + "832", + "--seed", + "42", + "--host", + "127.0.0.1", + "--port", + "9011", + "--prefer-sw-encoder", + ) + + +def test_output_capabilities_can_be_added_without_shared_routing_change( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _FakeAdapter: + def supported_modes(self, config, options): + del config, options + return ("webrtc",) + + def resolve(self, config, *, mode, options): + del config, options + if mode != "webrtc": + return None + return OutputTargetSpec( + mode="webrtc", + label="plugin demo", + module="plugin.demo", + ) + + config = _runner_config(runner_name="third-party-model") + config.output_adapter = "plugin:adapter" + monkeypatch.setattr( + output_targets_module, + "_load_output_adapter", + lambda path: _FakeAdapter(), + ) + + assert available_output_modes(config) == ("cli", "webrtc") + assert resolve_output_target(config, mode="webrtc").module == "plugin.demo" + + def test_omnidreams_local_window_target_uses_manifest_override() -> None: config = _runner_config( runner_name="omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae" diff --git a/flashdreams/tests/test_rope_kernel.py b/flashdreams/tests/test_rope_kernel.py index 3aa76b8ec..fa5f8f71e 100644 --- a/flashdreams/tests/test_rope_kernel.py +++ b/flashdreams/tests/test_rope_kernel.py @@ -29,6 +29,8 @@ from __future__ import annotations +from collections.abc import Callable + import pytest import torch from torch import Tensor @@ -36,22 +38,27 @@ from flashdreams.core.attention.rope import apply_rope_freqs from flashdreams.core.attention.rope_kernel import apply_rotary_pos_emb -try: - from transformer_engine.pytorch.attention.rope import ( - apply_rotary_pos_emb as _te_apply_rotary_pos_emb, - ) - _TE_AVAILABLE = True -except (ImportError, OSError): +def _load_te_apply_rope() -> Callable[..., Tensor] | None: try: - from transformer_engine.pytorch.attention import ( - apply_rotary_pos_emb as _te_apply_rotary_pos_emb, - ) - - _TE_AVAILABLE = True + from transformer_engine.pytorch.attention.rope import apply_rotary_pos_emb except (ImportError, OSError): - _te_apply_rotary_pos_emb = None - _TE_AVAILABLE = False + try: + from transformer_engine.pytorch.attention import apply_rotary_pos_emb + except (ImportError, OSError): + return None + except RuntimeError as exc: + # CPU CI intentionally installs the TE meta-package without its + # framework extension. Treat only that known state as unavailable; + # unexpected TE initialization errors should still fail the test run. + if "empty `transformer-engine` meta package" not in str(exc): + raise + return None + return apply_rotary_pos_emb + + +_te_apply_rotary_pos_emb = _load_te_apply_rope() +_TE_AVAILABLE = _te_apply_rotary_pos_emb is not None _requires_te = pytest.mark.skipif( diff --git a/flashdreams/tests/test_runtime_demo_api.py b/flashdreams/tests/test_runtime_demo_api.py index 7719c7c49..6cdf772e5 100644 --- a/flashdreams/tests/test_runtime_demo_api.py +++ b/flashdreams/tests/test_runtime_demo_api.py @@ -5,12 +5,12 @@ from collections.abc import Sequence from pathlib import Path +from types import SimpleNamespace from typing import Any import pytest import torch -from flashdreams.infra.video_output import VideoStepResult from flashdreams.runtime import ( CanonicalInputs, CanonicalInputSchema, @@ -39,11 +39,14 @@ Mp4OutputSpec, NullOutputSpec, PreparedScenario, + WebRTCAppResources, WebRTCOutputSpec, build_output_target, run_replay_demo, ) -from flashdreams.runtime.demo.webrtc import build_webrtc_demo +from flashdreams.runtime.demo.webrtc import ( + serve_webrtc_demo, +) from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager pytestmark = pytest.mark.ci_cpu @@ -170,11 +173,11 @@ def output_factory(output_spec: object) -> OutputTarget: def test_demo_adapter_declares_supported_modes() -> None: adapter = _FakeDemoAdapter( input_modes=("replay",), - output_modes=("null", "mp4", "webrtc"), + output_modes=("null", "mp4"), ) assert adapter.supported_input_modes() == ("replay",) - assert adapter.supported_output_modes() == ("null", "mp4", "webrtc") + assert adapter.supported_output_modes() == ("null", "mp4") with pytest.raises(ValueError, match="input_mode='keyboard-driving'"): run_replay_demo( @@ -191,8 +194,7 @@ def test_demo_adapter_declares_supported_modes() -> None: assert not adapter.create_runtime_called -def test_webrtc_demo_uses_existing_session_manager_with_adapter_runtime() -> None: - adapter = _FakeDemoAdapter() +def test_webrtc_demo_serves_a_prepared_session_manager() -> None: spec = DemoSpec( model_id="fake-demo", scenario="valid-scenario", @@ -208,20 +210,46 @@ def test_webrtc_demo_uses_existing_session_manager_with_adapter_runtime() -> Non ), ) - demo = build_webrtc_demo(spec=spec, adapter=adapter) - - assert isinstance(demo.session_manager, BaseWebRTCSessionManager) - assert demo.runtime is adapter.webrtc_runtime - assert demo.session_manager._runtime is adapter.webrtc_runtime - assert demo.session_manager.runtime_config.video_width == 16 - assert demo.session_manager.runtime_config.video_height == 8 - assert demo.session_manager.fps == 24 - assert demo.session_manager._model_name() == "fake-demo" - assert demo.app is None - assert demo.host == "0.0.0.0" - assert demo.port == 8082 - assert adapter.create_webrtc_runtime_calls == [spec] - assert not adapter.create_runtime_called + assert isinstance(spec.output, WebRTCOutputSpec) + runtime = _FakeWebRTCRuntime( + SimpleNamespace( + video_width=16, + video_height=8, + warmup_chunks=0, + warmup_timeout_s=1.0, + ) + ) + manager = BaseWebRTCSessionManager( + runtime=runtime, + runtime_config=runtime.config, + fps=24, + identity="fake-demo", + client_liveness_timeout_s=spec.output.client_liveness_timeout_s, + ) + calls: list[dict[str, Any]] = [] + + def fake_server_runner(**kwargs: Any) -> None: + calls.append(kwargs) + + app = serve_webrtc_demo( + output=spec.output, + model_id=spec.model_id, + session_manager=manager, + app_resources=WebRTCAppResources(preload_name="Fake demo"), + world_rank=1, + server_runner=fake_server_runner, + ) + + assert app is None + assert calls == [ + { + "world_rank": 1, + "session_manager": manager, + "app": None, + "host": "0.0.0.0", + "port": 8082, + } + ] class _ChunkIndexMapping: @@ -276,8 +304,8 @@ def __init__( *, scenario_valid: bool = True, video_output: bool = False, - input_modes: tuple[str, ...] = ("replay", "keyboard-driving"), - output_modes: tuple[str, ...] = ("null", "mp4", "webrtc"), + input_modes: tuple[str, ...] = ("replay",), + output_modes: tuple[str, ...] = ("null", "mp4"), ) -> None: self._scenario_valid = scenario_valid self._video_output = video_output @@ -296,8 +324,6 @@ def __init__( self.prepare_scenario_calls: list[DemoSpec] = [] self.create_runtime_called = False self.runtime: _FakeRuntime | None = None - self.webrtc_runtime: _FakeWebRTCRuntime | None = None - self.create_webrtc_runtime_calls: list[DemoSpec] = [] def supported_input_modes(self) -> tuple[str, ...]: return self._input_modes @@ -327,11 +353,6 @@ def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: raise ValueError("invalid scenario") return self.prepared_scenario - def create_webrtc_runtime(self, spec: DemoSpec) -> "_FakeWebRTCRuntime": - self.create_webrtc_runtime_calls.append(spec) - self.webrtc_runtime = _FakeWebRTCRuntime() - return self.webrtc_runtime - class _FakeRuntime: def __init__( @@ -382,28 +403,30 @@ def next_step_request(self) -> StepRequest | None: def step(self, inputs: InferenceInput) -> StepResult: self._inference_input_schema.require_step(inputs) - output: object if self._video_output: - output = VideoStepResult.from_video_chunk( - chunk_index=self.step_index, + result = StepResult.from_video_chunk( + step_index=self.step_index, video_chunk=torch.full( (1, 1, 1, 3, 2, 2), self.step_index, dtype=torch.float32, ), layout="bvtchw", + output_window=TimeWindow( + start_s=0.5 * self.step_index, + end_s=0.5 * (self.step_index + 1), + ), ) else: - output = f"chunk-{self.step_index}" - result = StepResult( - step_index=self.step_index, - output=output, - frame_count=1, - output_window=TimeWindow( - start_s=0.5 * self.step_index, - end_s=0.5 * (self.step_index + 1), - ), - ) + result = StepResult( + step_index=self.step_index, + output=f"chunk-{self.step_index}", + frame_count=1, + output_window=TimeWindow( + start_s=0.5 * self.step_index, + end_s=0.5 * (self.step_index + 1), + ), + ) self.step_index += 1 return result @@ -427,25 +450,32 @@ def close(self) -> Sequence[OutputArtifact]: class _FakeWebRTCRuntime: + def __init__(self, config: Any) -> None: + self.config = config + async def initialize(self) -> None: return None - async def reset_for_new_session(self) -> None: - return None + async def reset_for_new_session(self, *, session_input: Any = None) -> None: + del session_input - def peek_steady_chunk_num_frames(self) -> int: - return 1 + def peek_input_fps(self) -> float: + return 24.0 - def peek_next_chunk_num_frames(self) -> int: + def peek_steady_output_num_frames(self) -> int: return 1 - async def generate_chunk( + def next_step_request(self) -> StepRequest: + return StepRequest(step_index=0, metadata={"input_frame_count": 1}) + + async def step( self, *, + request: StepRequest, segments: list[Any], frame_times: list[float], ) -> Any: - del segments, frame_times + del request, segments, frame_times return None async def close(self) -> None: diff --git a/flashdreams/tests/test_runtime_video_output.py b/flashdreams/tests/test_runtime_video_output.py index 898acf734..71a4271dd 100644 --- a/flashdreams/tests/test_runtime_video_output.py +++ b/flashdreams/tests/test_runtime_video_output.py @@ -9,7 +9,6 @@ import pytest import torch -from flashdreams.infra.video_output import VideoStepResult from flashdreams.runtime import Mp4VideoOutputTarget, StepResult, TimeWindow pytestmark = pytest.mark.ci_cpu @@ -19,7 +18,7 @@ def test_mp4_video_output_target_rejects_non_video_payload(tmp_path: Path) -> No target = Mp4VideoOutputTarget(output_path=tmp_path / "out.mp4", fps=30) target.open() - with pytest.raises(TypeError, match="VideoStepResult"): + with pytest.raises(TypeError, match="video StepResult"): target.write(StepResult(step_index=0, output="not-video")) @@ -53,15 +52,11 @@ def fake_writer( ) target.open() target.write( - StepResult( + StepResult.from_video_chunk( step_index=3, - output=VideoStepResult.from_video_chunk( - chunk_index=3, - video_chunk=torch.zeros((1, 2, 4, 3, 5, 6)), - layout="bvtchw", - stats={"model_step_s": 0.5}, - ), - frame_count=4, + video_chunk=torch.zeros((1, 2, 4, 3, 5, 6)), + layout="bvtchw", + metrics={"model_step_s": 0.5}, output_window=TimeWindow(start_s=1.0, end_s=2.0), ) ) @@ -81,10 +76,9 @@ def fake_writer( ] assert artifacts[0].metadata["stats_history"] == ( { - "autoregressive_index": 3, - "model_step_s": 0.5, "step_index": 3, "frames": 4, + "model_step_s": 0.5, "output_start_s": 1.0, "output_end_s": 2.0, }, diff --git a/flashdreams/tests/test_runtime_worker.py b/flashdreams/tests/test_runtime_worker.py new file mode 100644 index 000000000..f6fbf84aa --- /dev/null +++ b/flashdreams/tests/test_runtime_worker.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +import threading + +import pytest + +from flashdreams.runtime import ThreadAffineRuntimeWorker + +pytestmark = pytest.mark.ci_cpu + + +@pytest.mark.asyncio +async def test_worker_preserves_order_and_thread_affinity() -> None: + worker = ThreadAffineRuntimeWorker(thread_name="test-runtime") + calls: list[tuple[int, int]] = [] + + def _record(value: int) -> int: + calls.append((value, threading.get_ident())) + return value * 2 + + results = await asyncio.gather(*[worker.call(_record, value) for value in range(4)]) + await worker.close() + + assert results == [0, 2, 4, 6] + assert [value for value, _thread_id in calls] == [0, 1, 2, 3] + assert len({thread_id for _value, thread_id in calls}) == 1 + + +@pytest.mark.asyncio +async def test_worker_propagates_exceptions_and_remains_usable() -> None: + worker = ThreadAffineRuntimeWorker() + + def _raise() -> None: + raise ValueError("bad runtime call") + + with pytest.raises(ValueError, match="bad runtime call"): + await worker.call(_raise) + + assert await worker.call(lambda: 7) == 7 + await worker.close() + + +@pytest.mark.asyncio +async def test_cancelled_await_does_not_abandon_ordered_runtime_work() -> None: + worker = ThreadAffineRuntimeWorker() + started = threading.Event() + release = threading.Event() + completed: list[str] = [] + + def _blocking_call() -> None: + started.set() + assert release.wait(timeout=2.0) + completed.append("first") + + first = asyncio.create_task(worker.call(_blocking_call)) + assert await asyncio.to_thread(started.wait, 2.0) + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + + second = asyncio.create_task(worker.call(completed.append, "second")) + release.set() + await second + await worker.close() + + assert completed == ["first", "second"] + + +@pytest.mark.asyncio +async def test_close_drains_work_and_rejects_new_calls() -> None: + worker = ThreadAffineRuntimeWorker() + assert await worker.call(lambda: "done") == "done" + + await worker.close() + await worker.close() + + assert worker.closed + with pytest.raises(RuntimeError, match="closed"): + await worker.call(lambda: None) + + +@pytest.mark.asyncio +async def test_worker_sets_cuda_device_when_thread_starts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + seen: list[object] = [] + monkeypatch.setattr("torch.cuda.set_device", seen.append) + worker = ThreadAffineRuntimeWorker(device="cuda:3") + + await worker.call(lambda: None) + await worker.close() + + assert [str(device) for device in seen] == ["cuda:3"] diff --git a/flashdreams/tests/test_runtime_worker_gpu.py b/flashdreams/tests/test_runtime_worker_gpu.py new file mode 100644 index 000000000..d95e43bcf --- /dev/null +++ b/flashdreams/tests/test_runtime_worker_gpu.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +import torch + +from flashdreams.runtime import ThreadAffineRuntimeWorker + +pytestmark = pytest.mark.ci_gpu + + +@pytest.mark.asyncio +async def test_compiled_cuda_graph_replays_stay_on_runtime_thread() -> None: + """Exercise repeated Triton launches and CUDA-graph replay on one worker.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA required.") + + device = torch.device("cuda", torch.cuda.current_device()) + worker = ThreadAffineRuntimeWorker(device=device, thread_name="gpu-runtime-test") + state: dict[str, object] = {} + + def _initialize() -> None: + static_input = torch.ones(1024, device=device) + compiled = torch.compile(lambda value: torch.sin(value) + 1.0) + for _ in range(3): + compiled(static_input) + torch.cuda.synchronize(device) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_output = compiled(static_input) + state.update( + static_input=static_input, + static_output=static_output, + graph=graph, + ) + + def _step(value: float) -> float: + static_input = state["static_input"] + static_output = state["static_output"] + graph = state["graph"] + assert isinstance(static_input, torch.Tensor) + assert isinstance(static_output, torch.Tensor) + assert isinstance(graph, torch.cuda.CUDAGraph) + static_input.fill_(value) + graph.replay() + torch.cuda.synchronize(device) + return float(static_output[0].item()) + + try: + await worker.call(_initialize) + values = [await worker.call(_step, float(index)) for index in range(8)] + finally: + await worker.close() + + expected = [ + float(torch.sin(torch.tensor(float(index))) + 1.0) for index in range(8) + ] + assert values == pytest.approx(expected) diff --git a/flashdreams/tests/test_video_output.py b/flashdreams/tests/test_video_output.py index 51b375120..4cf70b72b 100644 --- a/flashdreams/tests/test_video_output.py +++ b/flashdreams/tests/test_video_output.py @@ -5,8 +5,7 @@ from __future__ import annotations -from pathlib import Path -from typing import Any +from typing import Any, cast import pytest import torch @@ -14,58 +13,94 @@ from flashdreams.infra.video_output import ( LazyRGBFrame, VideoOutputStream, - VideoStepResult, + VideoResultCollector, infer_video_num_frames, lazy_rgb_frames_from_video_tensor, + prepare_video_for_mp4, video_tensor_to_hwc_uint8, ) +from flashdreams.runtime import StepResult pytestmark = pytest.mark.ci_cpu -def test_video_step_result_infers_num_frames_from_layout() -> None: +def test_step_result_infers_video_frame_count_from_layout() -> None: video = torch.zeros((1, 2, 4, 3, 5, 6), dtype=torch.float32) - result = VideoStepResult.from_video_chunk( - chunk_index=7, + result = StepResult.from_video_chunk( + step_index=7, video_chunk=video, layout="bvtchw", - stats={"total_ms": 12.5}, + metrics={"total_ms": 12.5}, metadata={"stream": "rgb"}, ) - assert result.chunk_index == 7 - assert result.num_frames == 4 + assert result.step_index == 7 + assert result.frame_count == 4 assert result.video_chunk is video - assert result.stats == {"total_ms": 12.5} + assert result.metrics == {"total_ms": 12.5} assert result.layout == "bvtchw" assert result.metadata == {"stream": "rgb"} assert infer_video_num_frames(video, layout="bvtchw") == 4 -def test_video_output_stream_makes_step_result_without_host_copy() -> None: +def test_step_result_validates_video_step_and_layout_shape() -> None: + with pytest.raises(ValueError, match="step_index"): + StepResult.from_video_chunk( + step_index=-1, + video_chunk=torch.zeros((1, 3, 2, 4, 5)), + layout="bcthw", + ) + + with pytest.raises(ValueError, match="expects a 5D tensor"): + StepResult.from_video_chunk( + step_index=0, + video_chunk=torch.zeros((2, 3, 4, 5)), + layout="bcthw", + ) + + +def test_step_result_freezes_video_metadata_and_metrics() -> None: + metadata = {"stream": "rgb"} + metrics = {"model_step_s": 0.5} + result = StepResult.from_video_chunk( + step_index=0, + video_chunk=torch.zeros((2, 3, 4, 5)), + layout="tchw", + metadata=metadata, + metrics=metrics, + ) + + metadata["stream"] = "debug" + metrics["model_step_s"] = 1.0 + + assert result.metadata == {"stream": "rgb"} + assert result.metrics == {"model_step_s": 0.5} + with pytest.raises(TypeError): + cast(Any, result.metadata)["stream"] = "debug" + + +def test_video_output_stream_returns_step_result_without_host_copy() -> None: video = torch.zeros((3, 3, 4, 5), dtype=torch.float32, requires_grad=True) output_stream = VideoOutputStream( postprocess_stream=None, output_layout="tchw", - collect_output=False, - move_to_cpu=False, ) - result = output_stream.make_step_result( + result = output_stream.process( video, autoregressive_index=4, - stats={"decode_ms": 1.5}, + metrics={"decode_ms": 1.5}, ) - assert isinstance(result, VideoStepResult) - assert result.chunk_index == 4 - assert result.num_frames == 3 + assert isinstance(result, StepResult) + assert result.step_index == 4 + assert result.frame_count == 3 assert result.video_chunk.device == video.device assert result.video_chunk.data_ptr() == video.data_ptr() assert result.video_chunk.requires_grad is False assert result.layout == "tchw" - assert result.stats == {"decode_ms": 1.5} + assert result.metrics == {"decode_ms": 1.5} def test_video_tensor_to_hwc_uint8_preserves_device_layout_conversion() -> None: @@ -93,9 +128,9 @@ def test_lazy_rgb_frames_from_video_tensor_materializes_on_demand() -> None: assert frames[1].to_numpy()[2, 3].tolist() == [255, 255, 255] -def test_video_step_result_exposes_lazy_rgb_frames() -> None: - result = VideoStepResult.from_video_chunk( - chunk_index=0, +def test_step_result_exposes_lazy_rgb_frames() -> None: + result = StepResult.from_video_chunk( + step_index=0, video_chunk=torch.zeros((1, 2, 1, 3, 4, 5), dtype=torch.float32), layout="bvtchw", ) @@ -106,50 +141,52 @@ def test_video_step_result_exposes_lazy_rgb_frames() -> None: assert frames[0].to_numpy().shape == (4, 5, 3) -def test_video_output_stream_collects_chunks_and_stats() -> None: +def test_video_result_collector_collects_chunks_and_stats() -> None: output_stream = VideoOutputStream( postprocess_stream=None, output_layout="tchw", - move_to_cpu=False, ) + collector = VideoResultCollector(output_layout="tchw", move_to_cpu=False) chunk = torch.zeros((2, 3, 4, 5), dtype=torch.float32) - processed = output_stream.process( + result = output_stream.process( chunk, autoregressive_index=3, - stats={"total_ms": 8.0}, - stats_extra={"frames": 2, "fps": 250.0}, + metrics={"total_ms": 8.0, "pipeline_fps": 250.0}, ) - collected = output_stream.finish() + collector.add(result) + assert output_stream.finish() is None + collected = collector.finish() assert collected is not None assert collected.shape == chunk.shape assert collected.data_ptr() == chunk.data_ptr() - assert processed is chunk - assert output_stream.stats_history == [ + assert result.video_chunk.data_ptr() == chunk.data_ptr() + assert collector.stats_history == [ { - "autoregressive_index": 3, - "total_ms": 8.0, + "step_index": 3, "frames": 2, - "fps": 250.0, + "total_ms": 8.0, + "pipeline_fps": 250.0, } ] -def test_video_output_stream_collects_noop_chunks_without_postprocess() -> None: +def test_video_result_collector_skips_empty_chunks() -> None: output_stream = VideoOutputStream( postprocess_stream=None, output_layout="bcthw", - move_to_cpu=False, ) + collector = VideoResultCollector(output_layout="bcthw", move_to_cpu=False) first = torch.ones((1, 3, 2, 4, 5)) empty = torch.empty((1, 3, 0, 4, 5)) second = torch.full((1, 3, 1, 4, 5), 2.0) - output_stream.process(first, autoregressive_index=0) - output_stream.process(empty, autoregressive_index=1) - output_stream.process(second, autoregressive_index=2) - output = output_stream.finish() + collector.add(output_stream.process(first, autoregressive_index=0)) + collector.add(output_stream.process(empty, autoregressive_index=1)) + collector.add(output_stream.process(second, autoregressive_index=2)) + assert output_stream.finish() is None + output = collector.finish() assert output is not None assert output.shape == (1, 3, 3, 4, 5) @@ -157,41 +194,81 @@ def test_video_output_stream_collects_noop_chunks_without_postprocess() -> None: assert torch.equal(output[:, :, 2:], second) -def test_video_output_stream_finishes_to_mp4_with_multiview_tiling() -> None: - calls: list[dict[str, Any]] = [] - - def fake_writer( - video: torch.Tensor, - path: Path, - *, - fps: int | float, - layout: str, - install_hint: str, - ) -> Path: - calls.append( - { - "shape": tuple(video.shape), - "path": path, - "fps": fps, - "layout": layout, - "install_hint": install_hint, - } - ) - return path +def test_video_output_stream_returns_postprocess_tail_as_step_result() -> None: + class _TailPostprocess: + last_process_stats = None + + def process( + self, + output: torch.Tensor, + *, + autoregressive_index: int, + ) -> torch.Tensor: + del autoregressive_index + return output[:, :, :0] + + def finish(self) -> torch.Tensor: + return torch.ones((1, 3, 2, 4, 5)) output_stream = VideoOutputStream( - postprocess_stream=None, - output_layout="bvtchw", - move_to_cpu=False, + postprocess_stream=cast(Any, _TailPostprocess()), + output_layout="bcthw", ) - output_stream.process( - torch.zeros((1, 2, 3, 3, 4, 5)), autoregressive_index=0 + result = output_stream.process( + torch.zeros((1, 3, 2, 4, 5)), + autoregressive_index=6, ) - written = output_stream.finish_to_mp4( - Path("output.mp4"), fps=24, writer=fake_writer + tail = output_stream.finish() + + assert result.frame_count == 0 + assert tail is not None + assert tail.step_index == 6 + assert tail.frame_count == 2 + assert tail.metadata == {"postprocess_tail": True} + + +def test_video_output_stream_state_is_isolated_per_session() -> None: + class _StatefulPostprocess: + last_process_stats = None + + def __init__(self) -> None: + self.calls = 0 + + def process( + self, + output: torch.Tensor, + *, + autoregressive_index: int, + ) -> torch.Tensor: + del autoregressive_index + self.calls += 1 + return output + self.calls + + def finish(self) -> None: + return None + + first = VideoOutputStream( + postprocess_stream=cast(Any, _StatefulPostprocess()), + output_layout="tchw", + ) + second = VideoOutputStream( + postprocess_stream=cast(Any, _StatefulPostprocess()), + output_layout="tchw", ) + video = torch.zeros((1, 3, 2, 2)) + + first_result = first.process(video, autoregressive_index=0) + second_result = second.process(video, autoregressive_index=0) + + assert torch.equal(first_result.video_chunk, second_result.video_chunk) + assert first.postprocess_stream is not second.postprocess_stream + + +def test_prepare_video_for_mp4_tiles_multiview_video() -> None: + video = torch.zeros((1, 2, 3, 3, 4, 5)) + + writable, layout = prepare_video_for_mp4(video, layout="bvtchw") - assert written is not None - assert written == Path("output.mp4") - assert calls[0]["shape"] == (3, 4, 10, 3) + assert writable.shape == (3, 4, 10, 3) + assert layout == "thwc" diff --git a/flashdreams/tests/test_webrtc_manager.py b/flashdreams/tests/test_webrtc_manager.py index e4537edb0..e6089e869 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -11,6 +11,7 @@ import pytest import torch +from flashdreams.runtime import StepRequest, StepResult from flashdreams.serving.webrtc import manager as manager_module from flashdreams.serving.webrtc.controls import WSAD_SUPPORTED_KEYS from flashdreams.serving.webrtc.encoders import ChunkDeliveryResult @@ -18,7 +19,6 @@ BaseWebRTCSessionManager, ManagedWebRTCSession, ) -from flashdreams.infra.video_output import VideoStepResult from flashdreams.serving.webrtc.server import SessionBusyError pytestmark = pytest.mark.ci_cpu @@ -33,14 +33,21 @@ def _runtime_config() -> SimpleNamespace: ) +def _step_request(step_index: int = 0, input_frame_count: int = 1) -> StepRequest: + return StepRequest( + step_index=step_index, + metadata={"input_frame_count": input_frame_count}, + ) + + class _FakeVideoTrack: fps = 30 def __init__(self) -> None: self.closed = False - async def enqueue_chunk(self, chunk: Any) -> int: - del chunk + async def enqueue_result(self, result: StepResult) -> int: + del result return 1 def qsize(self) -> int: @@ -53,8 +60,8 @@ async def close(self) -> None: class _FakeVideoEncoder: """``VideoEncoder``-shaped stub for ``ManagedWebRTCSession`` construction and the base manager's generation-worker path. ``deliver_chunk`` - delegates to the paired track's ``enqueue_chunk`` so the manager - tests that drive one chunk end-to-end see the frames land.""" + delegates to the paired track's ``enqueue_result`` so the manager + tests that drive one result end-to-end see the frames land.""" fps = 30 backend = "fake" @@ -62,13 +69,13 @@ class _FakeVideoEncoder: async def deliver_chunk( self, - chunk: Any, + result: StepResult, track: Any, *, force_keyframe: bool = False, ) -> ChunkDeliveryResult: del force_keyframe - enqueued = await track.enqueue_chunk(chunk) + enqueued = await track.enqueue_result(result) return ChunkDeliveryResult( backend=self.backend, num_frames=enqueued, @@ -116,13 +123,12 @@ def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: class _CountingVideoTrack(_FakeVideoTrack): - async def enqueue_chunk(self, chunk: Any) -> int: - return int(chunk.shape[0]) + async def enqueue_result(self, result: StepResult) -> int: + return result.frame_count class _BaseTestManager(BaseWebRTCSessionManager): - def _model_name(self) -> str: - return "fake-model" + pass class _WOnlyTestManager(_BaseTestManager): @@ -130,28 +136,36 @@ class _WOnlyTestManager(_BaseTestManager): def _make_manager( - manager_cls: type[BaseWebRTCSessionManager], runtime: Any + manager_cls: type[BaseWebRTCSessionManager], runtime: Any, **kwargs: Any ) -> BaseWebRTCSessionManager: return manager_cls( runtime=runtime, runtime_config=_runtime_config(), fps=30, + identity="fake-model", + **kwargs, ) -def test_runtime_frame_timing_hooks_default_to_legacy_methods() -> None: - class _LegacyRuntime: - def peek_next_chunk_num_frames(self) -> int: - return 2 +def test_runtime_frame_timing_contract() -> None: + class _Runtime: + def peek_input_fps(self) -> float: + return 30.0 + + def next_step_request(self) -> StepRequest: + return _step_request(input_frame_count=2) - def peek_steady_chunk_num_frames(self) -> int: + def peek_steady_output_num_frames(self) -> int: return 3 - runtime = _LegacyRuntime() + runtime = _Runtime() manager = _make_manager(_BaseTestManager, runtime) assert manager._runtime_input_fps(runtime) == pytest.approx(30.0) - assert manager._runtime_next_input_num_frames(runtime) == 2 + assert manager._runtime_next_step_request(runtime) == ( + _step_request(input_frame_count=2), + 2, + ) assert manager._runtime_steady_output_num_frames(runtime) == 3 @@ -160,8 +174,8 @@ class _SplitRuntime: def peek_input_fps(self) -> float: return 6.0 - def peek_next_input_num_frames(self) -> int: - return 4 + def next_step_request(self) -> StepRequest: + return _step_request(input_frame_count=4) def peek_steady_output_num_frames(self) -> int: return 16 @@ -174,7 +188,10 @@ def peek_steady_output_num_frames(self) -> int: ) assert resampler.dt == pytest.approx(1.0 / 6.0) - assert manager._runtime_next_input_num_frames(runtime) == 4 + assert manager._runtime_next_step_request(runtime) == ( + _step_request(input_frame_count=4), + 4, + ) assert manager._runtime_steady_output_num_frames(runtime) == 16 @@ -474,21 +491,22 @@ class _ClosingRuntime: def __init__(self) -> None: self.generate_calls = 0 - def peek_next_chunk_num_frames(self) -> int: - return 1 + def next_step_request(self) -> StepRequest: + return _step_request(step_index=self.generate_calls) - async def generate_chunk( - self, *, segments: Any, frame_times: Any - ) -> VideoStepResult: - del segments, frame_times + async def step( + self, *, request: StepRequest, segments: Any, frame_times: Any + ) -> StepResult: + del request, segments, frame_times self.generate_calls += 1 raise RuntimeError("boom") - class _ClosingManager(_BaseTestManager): - _close_session_on_generation_error = True - runtime = _ClosingRuntime() - manager = _make_manager(_ClosingManager, runtime) + manager = _make_manager( + _BaseTestManager, + runtime, + fatal_generation_errors=True, + ) managed, video_track, peer, channel = _managed_session(runtime) manager._active_session = managed @@ -511,13 +529,13 @@ def __init__(self) -> None: self.generate_calls = 0 self.managed_session: ManagedWebRTCSession | None = None - def peek_next_chunk_num_frames(self) -> int: - return 1 + def next_step_request(self) -> StepRequest: + return _step_request(step_index=self.generate_calls) - async def generate_chunk( - self, *, segments: Any, frame_times: Any - ) -> VideoStepResult: - del segments, frame_times + async def step( + self, *, request: StepRequest, segments: Any, frame_times: Any + ) -> StepResult: + del request, segments, frame_times self.generate_calls += 1 # Stop the loop after the second attempt without tearing down. if self.generate_calls >= 2 and self.managed_session is not None: @@ -548,28 +566,24 @@ class _OneChunkRuntime: def __init__(self) -> None: self.managed_session: ManagedWebRTCSession | None = None - def peek_next_chunk_num_frames(self) -> int: - return 1 + def next_step_request(self) -> StepRequest: + return _step_request() - async def generate_chunk( - self, *, segments: Any, frame_times: Any - ) -> VideoStepResult: - del segments, frame_times + async def step( + self, *, request: StepRequest, segments: Any, frame_times: Any + ) -> StepResult: + del request, segments, frame_times if self.managed_session is not None: self.managed_session.closed = True - return VideoStepResult( - chunk_index=0, - num_frames=1, + return StepResult.from_video_chunk( + step_index=0, video_chunk=torch.zeros((1, 1, 1, 3, 2, 2), dtype=torch.uint8), - stats=None, + layout="bvtchw", + metadata={"stream": "rgb"}, ) - class _ExtraManager(_BaseTestManager): - def _chunk_done_extra(self) -> dict[str, Any]: - return {"stream": "rgb"} - runtime = _OneChunkRuntime() - manager = _make_manager(_ExtraManager, runtime) + manager = _make_manager(_BaseTestManager, runtime) managed, _video_track, _peer, channel = _managed_session(runtime) runtime.managed_session = managed manager._active_session = managed @@ -616,21 +630,24 @@ def __init__(self) -> None: def peek_input_fps(self) -> float: return 6.0 - def peek_next_input_num_frames(self) -> int: - return 2 + def next_step_request(self) -> StepRequest: + return _step_request(input_frame_count=2) - async def generate_chunk( - self, *, segments: Any, frame_times: list[float] - ) -> VideoStepResult: - del segments + async def step( + self, + *, + request: StepRequest, + segments: Any, + frame_times: list[float], + ) -> StepResult: + del request, segments self.frame_times = frame_times if self.managed_session is not None: self.managed_session.closed = True - return VideoStepResult( - chunk_index=0, - num_frames=5, - video_chunk=torch.zeros((5, 1, 1, 3, 2, 2), dtype=torch.uint8), - stats=None, + return StepResult.from_video_chunk( + step_index=0, + video_chunk=torch.zeros((5, 3, 2, 2), dtype=torch.uint8), + layout="tchw", ) runtime = _SplitRuntime() @@ -674,22 +691,22 @@ def __init__(self) -> None: self.managed_session: ManagedWebRTCSession | None = None self.chunk_index = 0 - def peek_next_chunk_num_frames(self) -> int: - return 1 + def next_step_request(self) -> StepRequest: + return _step_request(step_index=self.chunk_index) - async def generate_chunk( - self, *, segments: Any, frame_times: Any - ) -> VideoStepResult: - del segments, frame_times + async def step( + self, *, request: StepRequest, segments: Any, frame_times: Any + ) -> StepResult: + del request, segments, frame_times chunk_index = self.chunk_index self.chunk_index += 1 if chunk_index >= 2 and self.managed_session is not None: self.managed_session.closed = True - return VideoStepResult( - chunk_index=chunk_index, - num_frames=4, - video_chunk=torch.zeros((4, 1, 1, 3, 2, 2), dtype=torch.uint8), - stats={ + return StepResult.from_video_chunk( + step_index=chunk_index, + video_chunk=torch.zeros((4, 3, 2, 2), dtype=torch.uint8), + layout="tchw", + metrics={ "model_step_s": 0.02, "denoise_s": 0.01, "decode_s": 0.004, @@ -725,10 +742,11 @@ class _FrequentLogManager(_BaseTestManager): @pytest.mark.asyncio async def test_create_answer_raises_busy_with_subclass_message() -> None: - class _BusyManager(_BaseTestManager): - _busy_message = "custom busy message" - - manager = _make_manager(_BusyManager, runtime=SimpleNamespace()) + manager = _make_manager( + _BaseTestManager, + runtime=SimpleNamespace(), + busy_message="custom busy message", + ) manager._runtime_ready = True manager._warmup_complete = True existing, *_ = _managed_session(runtime=SimpleNamespace()) @@ -739,12 +757,11 @@ class _BusyManager(_BaseTestManager): def test_make_resampler_honors_supported_keys() -> None: - class _WsadManager(_BaseTestManager): - _resampler_supported_keys = WSAD_SUPPORTED_KEYS - - wsad = _make_manager(_WsadManager, runtime=SimpleNamespace())._make_resampler( - start_v=1.0 - ) + wsad = _make_manager( + _BaseTestManager, + runtime=SimpleNamespace(), + supported_control_keys=WSAD_SUPPORTED_KEYS, + )._make_resampler(start_v=1.0) wsad.on_edge(arrival_t=0.5, event="keydown", key="q") wsad_segments, _ = wsad.sample_chunk(num_frames=1) # 'q' is not a WSAD driving key, so it is rejected and never held. diff --git a/flashdreams/tests/test_webrtc_serving.py b/flashdreams/tests/test_webrtc_serving.py index c06d434aa..e376f7eb2 100644 --- a/flashdreams/tests/test_webrtc_serving.py +++ b/flashdreams/tests/test_webrtc_serving.py @@ -339,7 +339,9 @@ def test_shared_viewer_exposes_model_extension_slots() -> None: assert "sendCommand: sendModelCommand" in javascript assert 'id="postprocessField"' in html assert 'fetch("/api/postprocess/options")' in javascript - assert "adapter.enablePostprocess === true" in javascript + assert "@typedef {Object} WebRTCModelAdapter" in javascript + assert "adapter.capabilities?.postprocess === true" in javascript + assert "renderControls(modelControls)" in javascript assert "/api/session/initial_scene" not in javascript diff --git a/integrations/causal_forcing/causal_forcing/runner.py b/integrations/causal_forcing/causal_forcing/runner.py index 2d4634d75..574ec0561 100644 --- a/integrations/causal_forcing/causal_forcing/runner.py +++ b/integrations/causal_forcing/causal_forcing/runner.py @@ -39,6 +39,7 @@ WanInferencePipeline, WanInferencePipelineCache, ) +from flashdreams.runtime.video_output import Mp4VideoOutputTarget __all__ = [ "CausalForcingI2VRunnerConfig", @@ -166,25 +167,46 @@ def run(self) -> None: # Generate the autoregressive chunks. output_stream = self.create_video_output_stream(fps=config.fps) + video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") + output_target = Mp4VideoOutputTarget( + output_path=video_path, + fps=config.fps, + output_layout=output_stream.output_layout, + enabled=self.is_rank_zero, + ) + output_target.open() for i in range(config.total_blocks): video_chunk = self.pipeline.generate(autoregressive_index=i, cache=cache) stats = self.pipeline.finalize(autoregressive_index=i, cache=cache) - output_stream.process(video_chunk, autoregressive_index=i, stats=stats) + output_target.write( + output_stream.process( + video_chunk, + autoregressive_index=i, + metrics=stats, + ) + ) - video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - video_path = output_stream.finish_to_mp4(video_path, fps=config.fps) - if video_path is None: + tail = output_stream.finish() + if tail is not None: + output_target.write(tail) + artifacts = output_target.close() + if not artifacts: return + video_artifact = artifacts[0] + video_path = Path(video_artifact.uri) logger.info( - f"[{config.runner_name}] wrote video {tuple(generated.shape)} " + f"[{config.runner_name}] wrote video {video_artifact.metadata['shape']} " f"-> {video_path.resolve()}" ) # Write the perf stats. - if output_stream.stats_history: + stats_history = video_artifact.metadata["stats_history"] + if stats_history: stats_path = write_runner_stats( - config.output_dir, config.runner_name, output_stream.stats_history + config.output_dir, + config.runner_name, + list(stats_history), ) logger.info( f"[{config.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" diff --git a/integrations/cosmos_predict2/cosmos_predict2/runner.py b/integrations/cosmos_predict2/cosmos_predict2/runner.py index 057b1cbb7..d37bbad69 100644 --- a/integrations/cosmos_predict2/cosmos_predict2/runner.py +++ b/integrations/cosmos_predict2/cosmos_predict2/runner.py @@ -39,6 +39,7 @@ CosmosInferencePipeline, CosmosInferencePipelineCache, ) +from flashdreams.runtime.video_output import Mp4VideoOutputTarget __all__ = [ "Cosmos2I2VRunner", @@ -136,24 +137,39 @@ def run(self) -> None: cache = self._initialize_cache() output_stream = self.create_video_output_stream(fps=config.fps) + video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") + output_target = Mp4VideoOutputTarget( + output_path=video_path, + fps=config.fps, + output_layout=output_stream.output_layout, + enabled=self.is_rank_zero, + ) + output_target.open() generated = self.pipeline.generate(autoregressive_index=0, cache=cache) stats = self.pipeline.finalize(autoregressive_index=0, cache=cache) - output_stream.process(generated, autoregressive_index=0, stats=stats) - video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - video_path = output_stream.finish_to_mp4(video_path, fps=config.fps) - if video_path is None: + output_target.write( + output_stream.process(generated, autoregressive_index=0, metrics=stats) + ) + tail = output_stream.finish() + if tail is not None: + output_target.write(tail) + artifacts = output_target.close() + if not artifacts: return + video_artifact = artifacts[0] + video_path = Path(video_artifact.uri) logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " f"-> {video_path.resolve()}" ) - if output_stream.stats_history: + stats_history = video_artifact.metadata["stats_history"] + if stats_history: stats_path = write_runner_stats( config.output_dir, config.runner_name, - output_stream.stats_history, + list(stats_history), ) logger.info( f"[{config.runner_name}] wrote per-AR-step stats " diff --git a/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py b/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py index f5cf82e80..781ca58e5 100644 --- a/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py +++ b/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py @@ -34,6 +34,7 @@ WanInferencePipeline, WanInferencePipelineCache, ) +from flashdreams.runtime.video_output import Mp4VideoOutputTarget __all__ = [ "FastvideoCausalWan22T2VRunnerConfig", @@ -126,26 +127,47 @@ def run(self) -> None: cache = self._initialize_cache() output_stream = self.create_video_output_stream(fps=config.fps) + video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") + output_target = Mp4VideoOutputTarget( + output_path=video_path, + fps=config.fps, + output_layout=output_stream.output_layout, + enabled=self.is_rank_zero, + ) + output_target.open() for i in range(config.total_blocks): # Generate the autoregressive chunks. video_chunk = self.pipeline.generate(autoregressive_index=i, cache=cache) stats = self.pipeline.finalize(autoregressive_index=i, cache=cache) - output_stream.process(video_chunk, autoregressive_index=i, stats=stats) + output_target.write( + output_stream.process( + video_chunk, + autoregressive_index=i, + metrics=stats, + ) + ) - video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - video_path = output_stream.finish_to_mp4(video_path, fps=config.fps) - if video_path is None: + tail = output_stream.finish() + if tail is not None: + output_target.write(tail) + artifacts = output_target.close() + if not artifacts: return + video_artifact = artifacts[0] + video_path = Path(video_artifact.uri) logger.info( - f"[{config.runner_name}] wrote video {tuple(generated.shape)} " + f"[{config.runner_name}] wrote video {video_artifact.metadata['shape']} " f"-> {video_path.resolve()}" ) # Write the perf stats. - if output_stream.stats_history: + stats_history = video_artifact.metadata["stats_history"] + if stats_history: stats_path = write_runner_stats( - config.output_dir, config.runner_name, output_stream.stats_history + config.output_dir, + config.runner_name, + list(stats_history), ) logger.info( f"[{config.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" diff --git a/integrations/flashvsr/flashvsr/runner.py b/integrations/flashvsr/flashvsr/runner.py index 881943948..9f1f40cb0 100644 --- a/integrations/flashvsr/flashvsr/runner.py +++ b/integrations/flashvsr/flashvsr/runner.py @@ -37,6 +37,7 @@ runner_artifact_path, write_runner_stats, ) +from flashdreams.runtime.video_output import Mp4VideoOutputTarget from flashvsr.encoder import FlashVSREncoder from flashvsr.pipeline import ( FlashVSRPipeline, @@ -423,6 +424,14 @@ def run(self) -> None: cache = self._initialize_cache() output_stream = self.create_video_output_stream(fps=fps) + video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") + output_target = Mp4VideoOutputTarget( + output_path=video_path, + fps=fps, + output_layout=output_stream.output_layout, + enabled=self.is_rank_zero, + ) + output_target.open() for chunk_idx, (start, size) in enumerate(chunks): clip = video_t[:, :, start : start + size] video_chunk = self.pipeline.generate( @@ -432,7 +441,7 @@ def run(self) -> None: ) pipeline_frames = int(video_chunk.shape[2]) stats = self.pipeline.finalize(autoregressive_index=chunk_idx, cache=cache) - stats_extra: dict[str, float | int] | None = None + metrics = dict(stats or {}) if stats is not None: # Pipeline throughput is based on this AR step's direct output. # Postprocess emission/buffering is reported separately. @@ -442,27 +451,40 @@ def run(self) -> None: if chunk_total_ms > 0 else 0.0 ) - stats_extra = {"frames": pipeline_frames, "fps": chunk_fps} - output_stream.process( - video_chunk, - autoregressive_index=chunk_idx, - stats=stats, - stats_extra=stats_extra, + metrics.update( + { + "pipeline_frames": pipeline_frames, + "pipeline_fps": chunk_fps, + } + ) + output_target.write( + output_stream.process( + video_chunk, + autoregressive_index=chunk_idx, + metrics=metrics, + ) ) - video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - video_path = output_stream.finish_to_mp4(video_path, fps=fps) - if video_path is None: + tail = output_stream.finish() + if tail is not None: + output_target.write(tail) + artifacts = output_target.close() + if not artifacts: return + video_artifact = artifacts[0] + video_path = Path(video_artifact.uri) logger.info( - f"[{config.runner_name}] wrote video {tuple(generated.shape)} " + f"[{config.runner_name}] wrote video {video_artifact.metadata['shape']} " f"-> {video_path.resolve()}" ) - if output_stream.stats_history: + stats_history = video_artifact.metadata["stats_history"] + if stats_history: stats_path = write_runner_stats( - config.output_dir, config.runner_name, output_stream.stats_history + config.output_dir, + config.runner_name, + list(stats_history), ) logger.info( f"[{config.runner_name}] wrote per-AR-step stats -> " diff --git a/integrations/hy_worldplay/hy_worldplay/runner.py b/integrations/hy_worldplay/hy_worldplay/runner.py index c334809b9..d29a9030d 100644 --- a/integrations/hy_worldplay/hy_worldplay/runner.py +++ b/integrations/hy_worldplay/hy_worldplay/runner.py @@ -36,6 +36,7 @@ write_runner_stats, ) from flashdreams.recipes.wan.pipeline import WanInferencePipeline +from flashdreams.runtime.video_output import Mp4VideoOutputTarget __all__ = [ "DEFAULT_PROMPT", @@ -339,7 +340,16 @@ def run(self) -> None: device=device, dtype=first_param.dtype ) - output_stream = self.create_video_output_stream(fps=cfg.fps, move_to_cpu=False) + output_stream = self.create_video_output_stream(fps=cfg.fps) + out_path = runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4") + output_target = Mp4VideoOutputTarget( + output_path=out_path, + fps=cfg.fps, + output_layout=output_stream.output_layout, + move_to_cpu=False, + enabled=self.is_rank_zero, + ) + output_target.open() if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() start_time = time.time() @@ -350,21 +360,35 @@ def run(self) -> None: # advances the KV cache; called on every chunk # (including the last) for consistent stats. stats = self.pipeline.finalize(ar_idx, cache) - output_stream.process(chunk, autoregressive_index=ar_idx, stats=stats) + output_target.write( + output_stream.process( + chunk, + autoregressive_index=ar_idx, + metrics=stats, + ) + ) elapsed = time.time() - start_time - out_path = runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4") - out_path = output_stream.finish_to_mp4(out_path, fps=cfg.fps) - if out_path is None: + tail = output_stream.finish() + if tail is not None: + output_target.write(tail) + artifacts = output_target.close() + if not artifacts: return + video_artifact = artifacts[0] + out_path = Path(video_artifact.uri) logger.info( f"[{cfg.runner_name}] wrote video " - f"({tuple(video.shape)}) -> {out_path.resolve()} in {elapsed:.2f}s" + f"({video_artifact.metadata['shape']}) -> {out_path.resolve()} " + f"in {elapsed:.2f}s" ) - if output_stream.stats_history: + stats_history = video_artifact.metadata["stats_history"] + if stats_history: stats_path = write_runner_stats( - cfg.output_dir, cfg.runner_name, output_stream.stats_history + cfg.output_dir, + cfg.runner_name, + list(stats_history), ) logger.info( f"[{cfg.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" diff --git a/integrations/lingbot/lingbot/demo/adapter.py b/integrations/lingbot/lingbot/demo/adapter.py index 354ea56b4..f117cbfa7 100644 --- a/integrations/lingbot/lingbot/demo/adapter.py +++ b/integrations/lingbot/lingbot/demo/adapter.py @@ -9,7 +9,6 @@ from typing import Any from flashdreams.runtime import ( - InferenceConfig, InputCanonicalizer, UserInputCapability, UserInputs, @@ -19,38 +18,26 @@ DemoSpec, Mp4OutputSpec, PreparedScenario, - WebRTCOutputSpec, ) from flashdreams.runtime.interfaces import InferenceRuntime +from lingbot.input_mapping import ( + KeyboardToCameraCommand, + TextEventSelection, +) from lingbot.runtime import ( LingbotModelAdapter, LingbotReplayRuntime, PipelineFactory, - build_lingbot_webrtc_runtime_config, inference_input_from_replay_inputs, ) -from lingbot.input_mapping import ( - KeyboardToCameraCommand, - TextEventSelection, -) -from lingbot.webrtc.session import ( - LingbotInferenceRuntime, - LingbotRuntimeConfig, -) from .spec import ( resolve_replay_inputs, resolve_text_event_prompts, resolve_user_input_events, - resolve_webrtc_scenario, -) -from .webrtc import ( - LingbotDemoWebRTCSessionManager, - create_lingbot_webrtc_app, ) ReplayRuntimeFactory = Callable[..., InferenceRuntime] -WebRTCRuntimeFactory = Callable[..., Any] class LingbotDemoAdapter(LingbotModelAdapter): @@ -60,20 +47,18 @@ def __init__( self, *, replay_runtime_factory: ReplayRuntimeFactory = LingbotReplayRuntime, - webrtc_runtime_factory: WebRTCRuntimeFactory = LingbotInferenceRuntime, pipeline_factory: PipelineFactory | None = None, ) -> None: super().__init__( runtime_factory=replay_runtime_factory, pipeline_factory=pipeline_factory, ) - self._webrtc_runtime_factory = webrtc_runtime_factory def supported_input_modes(self) -> tuple[str, ...]: - return ("replay", "keyboard-driving") + return ("replay",) def supported_output_modes(self) -> tuple[str, ...]: - return ("mp4", "webrtc") + return ("mp4",) def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: if spec.input_mode != "replay": @@ -101,7 +86,7 @@ def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: # A trace's world scale is derived from how far its poses # travel, so a stationary example yields 0. Live control has no # trajectory to normalize against, so it falls back to the same - # unit scale the WebRTC runtime uses. + # unit scale the live runtime uses. world_scale=trace.world_scale or 1.0, prompt=replay_inputs.prompt, text_event_prompts=text_event_prompts, @@ -123,88 +108,6 @@ def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: }, ) - def create_webrtc_runtime(self, spec: DemoSpec) -> Any: - runtime_config = self.create_webrtc_runtime_config(spec=spec, runtime=None) - return self._webrtc_runtime_factory(config=runtime_config) - - def create_webrtc_runtime_config( - self, - *, - spec: DemoSpec, - runtime: Any, - ) -> LingbotRuntimeConfig: - runtime_config = getattr(runtime, "config", None) - if isinstance(runtime_config, LingbotRuntimeConfig): - return runtime_config - if spec.input_mode != "keyboard-driving": - raise ValueError( - "Lingbot WebRTC requires input_mode='keyboard-driving', " - f"got {spec.input_mode!r}." - ) - if not isinstance(spec.output, WebRTCOutputSpec): - raise ValueError("Lingbot WebRTC requires WebRTC output.") - config = spec.config - if config is None: - raise RuntimeError("DemoSpec.config was not initialized.") - self.validate_config(config) - scenario = resolve_webrtc_scenario(spec.scenario) - - compile_network = ( - bool(config.compile) - if config.compile is not None - else bool(_option(config, "compile_network", True)) - ) - return build_lingbot_webrtc_runtime_config( - preset_id=self.preset_id(config), - pipeline_config=self.pipeline_config(config), - seed=int(_option(config, "seed", 42)), - compile_network=compile_network, - context_parallel_size=int(_option(config, "context_parallel_size", 1)), - device=config.device or str(_option(config, "device", "cuda:0")), - video_height=spec.output.video_height, - video_width=spec.output.video_width, - fps=spec.output.fps, - warmup_chunks=spec.output.warmup_chunks, - warmup_timeout_s=spec.output.warmup_timeout_s, - example_idx=int(_option(config, "example_idx", scenario.example_idx)), - prefer_sw_encoder=scenario.prefer_sw_encoder, - runtime_options=config.runtime_options, - ) - - def create_webrtc_session_manager( - self, - *, - spec: DemoSpec, - runtime: Any, - runtime_config: LingbotRuntimeConfig, - fps: int, - client_liveness_timeout_s: float, - ) -> LingbotDemoWebRTCSessionManager: - del spec - return LingbotDemoWebRTCSessionManager( - runtime=runtime, - runtime_config=runtime_config, - fps=fps, - client_liveness_timeout_s=client_liveness_timeout_s, - ) - - def create_webrtc_app( - self, - *, - spec: DemoSpec, - session_manager: Any, - request_session_url: str, - ) -> Any: - return create_lingbot_webrtc_app( - spec=spec, - session_manager=session_manager, - request_session_url=request_session_url, - ) - - -def _option(config: InferenceConfig, name: str, default: Any) -> Any: - return config.runtime_options.get(name, default) - def _camera_source(scenario: Any) -> str: if isinstance(scenario, Mapping): @@ -273,5 +176,4 @@ def _canonicalizer(text_event_prompts: Mapping[str, str] | None) -> InputCanonic __all__ = [ "LingbotDemoAdapter", "ReplayRuntimeFactory", - "WebRTCRuntimeFactory", ] diff --git a/integrations/lingbot/lingbot/demo/cli.py b/integrations/lingbot/lingbot/demo/app.py similarity index 88% rename from integrations/lingbot/lingbot/demo/cli.py rename to integrations/lingbot/lingbot/demo/app.py index 9b08469f6..df6b9c703 100644 --- a/integrations/lingbot/lingbot/demo/cli.py +++ b/integrations/lingbot/lingbot/demo/app.py @@ -7,23 +7,15 @@ import argparse from pathlib import Path +from typing import Any -import torch -import torch.distributed as dist - -from flashdreams.core.distributed import init as distributed_init from flashdreams.runtime import InferenceConfig from flashdreams.runtime.demo import ( DemoSpec, Mp4OutputSpec, WebRTCOutputSpec, - run_flashdreams_demo, - serve_flashdreams_demo, -) -from flashdreams.serving.webrtc.bootstrap import ( - configure_logging, - initialize_cuda_distributed, ) +from flashdreams.runtime.demo.app import DemoApplication from lingbot.example_data import ( EXAMPLE_DATA_AVAILABLE_IDXS, ensure_example_data_downloaded, @@ -120,36 +112,43 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: return parser.parse_args(argv) -def main(argv: list[str] | None = None) -> None: - configure_logging() - args = parse_args(argv) - adapter = LingbotDemoAdapter() - if args.command == "replay": - run_flashdreams_demo(spec=_replay_spec(args), adapter=adapter) - return - if args.command == "webrtc": - context = initialize_cuda_distributed( - default_device=args.device, - distributed_init_fn=distributed_init, - configure_logging_fn=configure_logging, - torch_module=torch, - dist_module=dist, - ) +class LingbotDemoApplication(DemoApplication): + """Lingbot replay and WebRTC demo application.""" + + def parse_args(self, argv: list[str] | None = None) -> argparse.Namespace: + return parse_args(argv) + + def replay_spec(self, args: argparse.Namespace) -> DemoSpec: + return _replay_spec(args) + + def replay_adapter(self) -> LingbotDemoAdapter: + return LingbotDemoAdapter() + + def prepare_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: ensure_example_data_downloaded( is_rank_zero=(context.world_rank == 0), example_idx=args.example_idx, ) - serve_flashdreams_demo( + + def serve_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: + from .webrtc import serve_lingbot_webrtc_demo + + serve_lingbot_webrtc_demo( spec=_webrtc_spec( args, device=str(context.device), context_parallel_size=context.world_size, ), - adapter=adapter, world_rank=context.world_rank, ) - return - raise AssertionError(f"Unhandled command: {args.command}") + + +_APPLICATION = LingbotDemoApplication() + + +def main(argv: list[str] | None = None) -> None: + """Run the Lingbot demo application.""" + _APPLICATION.main(argv) def _replay_spec(args: argparse.Namespace) -> DemoSpec: diff --git a/integrations/lingbot/lingbot/demo/spec.py b/integrations/lingbot/lingbot/demo/spec.py index ef153fdb0..8fd15494e 100644 --- a/integrations/lingbot/lingbot/demo/spec.py +++ b/integrations/lingbot/lingbot/demo/spec.py @@ -101,8 +101,7 @@ def resolve_user_input_events(value: Any) -> UserInputs: continue if not isinstance(record, Mapping): raise TypeError( - "Lingbot scenario events must be UserInputEvent objects or " - "mappings." + "Lingbot scenario events must be UserInputEvent objects or mappings." ) payload = { key: item @@ -113,8 +112,7 @@ def resolve_user_input_events(value: Any) -> UserInputs: event_type = record.get("type", record.get("event_type")) if timestamp_s is None or event_type is None: raise ValueError( - "Lingbot scenario events require a timestamp ('t') and a " - "type ('type')." + "Lingbot scenario events require a timestamp ('t') and a type ('type')." ) events.append( UserInputEvent( diff --git a/integrations/lingbot/lingbot/demo/webrtc.py b/integrations/lingbot/lingbot/demo/webrtc.py index c9036670f..4efa4152e 100644 --- a/integrations/lingbot/lingbot/demo/webrtc.py +++ b/integrations/lingbot/lingbot/demo/webrtc.py @@ -5,64 +5,107 @@ from __future__ import annotations -from importlib.resources import as_file, files +from collections.abc import Callable +from importlib.resources import files from typing import Any -from aiohttp import web - -from flashdreams.runtime.demo import DemoSpec -from flashdreams.serving.webrtc.server import ( - close_package_resources, - create_packaged_webrtc_app, +from flashdreams.runtime import InferenceConfig +from flashdreams.runtime.demo import DemoSpec, WebRTCAppResources, WebRTCOutputSpec +from flashdreams.runtime.demo.webrtc import ( + CreateWebRTCApp, + RunWebRTCServer, + serve_webrtc_demo, +) +from flashdreams.serving.webrtc.bootstrap import run_webrtc_server +from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager +from flashdreams.serving.webrtc.server import create_webrtc_app +from lingbot.runtime import ( + LingbotModelAdapter, + build_lingbot_webrtc_runtime_config, ) +from lingbot.webrtc.server import configure_lingbot_webrtc_app from lingbot.webrtc.session import ( LingbotInferenceRuntime, LingbotRuntimeConfig, - LingbotWebRTCSessionManager, + create_lingbot_webrtc_session_manager, ) -from lingbot.webrtc.server import configure_lingbot_webrtc_app - -class LingbotDemoWebRTCSessionManager(LingbotWebRTCSessionManager): - """Shared demo session manager using Lingbot's existing WebRTC semantics.""" +from .spec import resolve_webrtc_scenario - def __init__( - self, - *, - runtime: LingbotInferenceRuntime, - runtime_config: LingbotRuntimeConfig, - fps: int, - client_liveness_timeout_s: float, - ) -> None: - super().__init__( - runtime=runtime, - runtime_config=runtime_config, - fps=fps, - client_liveness_timeout_s=client_liveness_timeout_s, - ) +WebRTCRuntimeFactory = Callable[..., Any] -def create_lingbot_webrtc_app( +def serve_lingbot_webrtc_demo( *, spec: DemoSpec, - session_manager: Any, - request_session_url: str, -) -> web.Application: - """Create Lingbot's shared browser app through generic serving glue.""" - del spec - return create_packaged_webrtc_app( - web_resource=files("flashdreams.serving.webrtc").joinpath("web"), - model_web_resource=files("lingbot.webrtc").joinpath("web"), - session_manager=session_manager, - preload_name="Lingbot", - request_session_url=request_session_url, - configure_app=configure_lingbot_webrtc_app, - as_file_fn=as_file, - cleanup_callback=close_package_resources, + world_rank: int = 0, + runtime_factory: WebRTCRuntimeFactory = LingbotInferenceRuntime, + model_adapter: LingbotModelAdapter | None = None, + create_app_fn: CreateWebRTCApp = create_webrtc_app, + server_runner: RunWebRTCServer = run_webrtc_server, +) -> object: + """Create Lingbot's runtime and serve it through the shared WebRTC transport.""" + if spec.input_mode != "keyboard-driving": + raise ValueError( + "Lingbot WebRTC requires input_mode='keyboard-driving', " + f"got {spec.input_mode!r}." + ) + if not isinstance(spec.output, WebRTCOutputSpec): + raise ValueError("Lingbot WebRTC requires WebRTC output.") + config = spec.config + if config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + model_adapter = model_adapter or LingbotModelAdapter() + model_adapter.validate_config(config) + scenario = resolve_webrtc_scenario(spec.scenario) + compile_network = ( + bool(config.compile) + if config.compile is not None + else bool(_option(config, "compile_network", True)) + ) + runtime_config = build_lingbot_webrtc_runtime_config( + preset_id=model_adapter.preset_id(config), + pipeline_config=model_adapter.pipeline_config(config), + seed=int(_option(config, "seed", 42)), + compile_network=compile_network, + context_parallel_size=int(_option(config, "context_parallel_size", 1)), + device=config.device or str(_option(config, "device", "cuda:0")), + video_height=spec.output.video_height, + video_width=spec.output.video_width, + fps=spec.output.fps, + warmup_chunks=spec.output.warmup_chunks, + warmup_timeout_s=spec.output.warmup_timeout_s, + example_idx=int(_option(config, "example_idx", scenario.example_idx)), + prefer_sw_encoder=scenario.prefer_sw_encoder, + runtime_options=config.runtime_options, + ) + runtime = runtime_factory(config=runtime_config) + manager = create_lingbot_webrtc_session_manager( + runtime=runtime, + runtime_config=runtime_config, + fps=spec.output.fps, + client_liveness_timeout_s=spec.output.client_liveness_timeout_s, ) + return serve_webrtc_demo( + output=spec.output, + model_id=spec.model_id, + session_manager=manager, + app_resources=WebRTCAppResources( + model_web_resource=files("lingbot.webrtc").joinpath("web"), + preload_name="Lingbot", + configure_app=configure_lingbot_webrtc_app, + ), + world_rank=world_rank, + create_app_fn=create_app_fn, + server_runner=server_runner, + ) + + +def _option(config: InferenceConfig, name: str, default: Any) -> Any: + return config.runtime_options.get(name, default) __all__ = [ - "LingbotDemoWebRTCSessionManager", - "create_lingbot_webrtc_app", + "WebRTCRuntimeFactory", + "serve_lingbot_webrtc_demo", ] diff --git a/integrations/lingbot/lingbot/input_mapping.py b/integrations/lingbot/lingbot/input_mapping.py index 9a05c2178..0814d9819 100644 --- a/integrations/lingbot/lingbot/input_mapping.py +++ b/integrations/lingbot/lingbot/input_mapping.py @@ -45,12 +45,12 @@ ) from flashdreams.runtime.mapping import InputMappingSchema from flashdreams.runtime.types import StepRequest +from flashdreams.serving.realtime.input import DEFAULT_SUPPORTED_KEYS from flashdreams.serving.webrtc.controls import ( CameraPoseIntegrator, KeyboardState, PoseSegment, ) -from flashdreams.serving.realtime.input import DEFAULT_SUPPORTED_KEYS FIELD_CAMERA_TRAJECTORY = "camera_trajectory" FIELD_CAMERA_INTRINSICS = "camera_intrinsics" @@ -333,9 +333,7 @@ def load_camera_trace( return LingbotCameraTrace( poses=torch.from_numpy(np.ascontiguousarray(poses)).to(torch.float32), intrinsics=intrinsics.to(torch.float32), - world_scale=float( - inferred_world_scale if world_scale is None else world_scale - ), + world_scale=float(inferred_world_scale if world_scale is None else world_scale), ) @@ -575,9 +573,11 @@ def _integrate( window = request.user_input_window start_s = window.start_s if window is not None else frame_start / self._fps - end_s = window.end_s if window is not None else ( - frame_start + num_frames - ) / self._fps + end_s = ( + window.end_s + if window is not None + else (frame_start + num_frames) / self._fps + ) segments = _pose_segments(command, start_s=start_s, end_s=end_s) frame_times = [start_s + (index + 1) / self._fps for index in range(num_frames)] # The integrator rejects frame times outside the segment span, and float diff --git a/integrations/lingbot/lingbot/model_session.py b/integrations/lingbot/lingbot/model_session.py new file mode 100644 index 000000000..751729ec9 --- /dev/null +++ b/integrations/lingbot/lingbot/model_session.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared synchronous Lingbot model-session state and execution.""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Mapping +from typing import Any + +import torch + +from flashdreams.infra.video_output import VideoOutputStream +from flashdreams.runtime import StepResult, TimeWindow + +OutputStreamFactory = Callable[[], VideoOutputStream] + + +class LingbotModelSessionCore: + """Own one Lingbot cache, AR index, and generated-output stream.""" + + def __init__( + self, + *, + pipeline: Any, + output_stream_factory: OutputStreamFactory, + ) -> None: + self.pipeline = pipeline + self._output_stream_factory = output_stream_factory + self._output_stream = output_stream_factory() + self._cache: Any | None = None + self._step_index = 0 + self._closed = False + + @property + def cache(self) -> Any: + if self._cache is None: + raise RuntimeError("Lingbot model session is not initialized.") + return self._cache + + @property + def step_index(self) -> int: + return self._step_index + + def next_num_frames(self) -> int: + self._require_open() + return int(self.pipeline.get_num_output_frames(self._step_index)) + + def reset(self, *, prompt: str, first_frames: torch.Tensor) -> None: + self._require_open() + self._cache = None + self._output_stream.finish() + self._output_stream = self._output_stream_factory() + self._cache = self.pipeline.initialize_cache( + text=[prompt], + image=first_frames, + ) + self._step_index = 0 + + def step( + self, + camctrl_input: Any, + *, + output_window: TimeWindow | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> StepResult: + self._require_open() + step_index = self._step_index + expected_frames = self.next_num_frames() + start_t = time.perf_counter() + video_chunk = self.pipeline.generate( + autoregressive_index=step_index, + cache=self.cache, + input=camctrl_input, + ) + stats = self.pipeline.finalize( + autoregressive_index=step_index, + cache=self.cache, + ) + metrics = _numeric_metrics(stats) + metrics.setdefault("model_step_s", time.perf_counter() - start_t) + result = self._output_stream.process( + video_chunk, + autoregressive_index=step_index, + metrics=metrics, + metadata=metadata, + output_window=output_window, + ) + if result.frame_count != expected_frames: + raise RuntimeError( + f"Expected generated chunk to contain {expected_frames} frames, " + f"got {result.frame_count}." + ) + self._step_index += 1 + return result + + def replace_text_embeddings(self, text_embeddings: torch.Tensor) -> None: + self._require_open() + transformer = self.pipeline.diffusion_model.transformer + replace = getattr(transformer, "replace_text_embeddings", None) + if not callable(replace): + raise RuntimeError( + "Current Lingbot pipeline does not support text-context swapping." + ) + replace(self.cache.transformer_cache, text_embeddings) + + def close(self) -> None: + if self._closed: + return + self._closed = True + self._cache = None + self._output_stream.finish() + + def _require_open(self) -> None: + if self._closed: + raise RuntimeError("Lingbot model session is closed.") + + +def _numeric_metrics(stats: object) -> dict[str, float | int]: + if not isinstance(stats, Mapping): + return {} + return { + str(name): value + for name, value in stats.items() + if isinstance(value, (int, float)) and not isinstance(value, bool) + } + + +__all__ = ["LingbotModelSessionCore"] diff --git a/integrations/lingbot/lingbot/output_targets.py b/integrations/lingbot/lingbot/output_targets.py new file mode 100644 index 000000000..86e77b69f --- /dev/null +++ b/integrations/lingbot/lingbot/output_targets.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot output capabilities for ``flashdreams-run``.""" + +from __future__ import annotations + +from typing import Any + +from flashdreams.infra.runner import RunnerConfig +from flashdreams.serving.output_targets import ( + OutputLaunchOptions, + OutputMode, + OutputTargetSpec, +) + + +class LingbotOutputTargetAdapter: + def supported_modes( + self, + config: RunnerConfig, + options: OutputLaunchOptions, + ) -> tuple[OutputMode, ...]: + del config, options + return ("webrtc",) + + def resolve( + self, + config: RunnerConfig, + *, + mode: OutputMode, + options: OutputLaunchOptions, + ) -> OutputTargetSpec | None: + if mode != "webrtc": + return None + argv = [ + "webrtc", + "--preset-id", + _pipeline_name(config), + "--device", + str(config.device), + "--fps", + str(getattr(config, "fps", 16)), + "--video-height", + str(getattr(config, "pixel_height", 464)), + "--video-width", + str(getattr(config, "pixel_width", 832)), + ] + if _compile_network(config) is False: + argv.append("--no-compile") + example_idx = getattr(config, "example_idx", None) + if example_idx is not None: + argv.extend(("--example-idx", str(example_idx))) + if options.host: + argv.extend(("--host", options.host)) + if options.port is not None: + argv.extend(("--port", str(options.port))) + if options.prefer_sw_encoder: + argv.append("--prefer-sw-encoder") + return OutputTargetSpec( + mode="webrtc", + label="LingBot shared demo WebRTC server", + module="lingbot.demo.app", + argv=tuple(argv), + ) + + +def _pipeline_name(config: RunnerConfig) -> str: + name = getattr(config.pipeline, "name", None) + return str(name or config.runner_name) + + +def _compile_network(config: RunnerConfig) -> bool | None: + diffusion_model = getattr(config.pipeline, "diffusion_model", None) + transformer: Any = getattr(diffusion_model, "transformer", None) + value = getattr(transformer, "compile_network", None) + return None if value is None else bool(value) + + +OUTPUT_TARGET_ADAPTER = LingbotOutputTargetAdapter() + +__all__ = ["OUTPUT_TARGET_ADAPTER", "LingbotOutputTargetAdapter"] diff --git a/integrations/lingbot/lingbot/runner.py b/integrations/lingbot/lingbot/runner.py index 24016bb46..7e4a139c4 100644 --- a/integrations/lingbot/lingbot/runner.py +++ b/integrations/lingbot/lingbot/runner.py @@ -29,9 +29,6 @@ from flashdreams.runtime.runner import run_inference_session from lingbot.example_data import ( EXAMPLE_DATA_AVAILABLE_IDXS, - EXAMPLE_DATA_BASE_URL, - EXAMPLE_DATA_DIR_LOCAL, - EXAMPLE_DATA_FILENAMES, EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS, ensure_example_data_downloaded, example_data_dirname, @@ -48,8 +45,10 @@ ) __all__ = [ + "EXAMPLE_DATA_AVAILABLE_IDXS", "LingbotWorldRunnerConfig", "LingbotWorldRunner", + "example_data_dirname", ] @@ -61,6 +60,7 @@ _INTRINSICS_REFERENCE_WIDTH = 832 """Capture-resolution width matching :data:`_INTRINSICS_REFERENCE_HEIGHT`.""" + @dataclass(kw_only=True) class LingbotWorldRunnerConfig(RunnerConfig): """Runner config for every shipped LingBot-World variant.""" @@ -68,6 +68,7 @@ class LingbotWorldRunnerConfig(RunnerConfig): _target: type["LingbotWorldRunner"] = field( default_factory=lambda: LingbotWorldRunner ) + output_adapter: str | None = "lingbot.output_targets:OUTPUT_TARGET_ADAPTER" prompt: str = "" """Text prompt. A non-empty value wins; otherwise the runner reads diff --git a/integrations/lingbot/lingbot/runtime.py b/integrations/lingbot/lingbot/runtime.py index 0bf12bddc..501376b7f 100644 --- a/integrations/lingbot/lingbot/runtime.py +++ b/integrations/lingbot/lingbot/runtime.py @@ -6,11 +6,10 @@ from __future__ import annotations import os -import time from collections.abc import Callable, Mapping from dataclasses import dataclass, replace from pathlib import Path -from typing import Any +from typing import Any, cast import numpy as np import torch @@ -24,13 +23,14 @@ runner_artifact_path, write_runner_stats, ) -from flashdreams.infra.video_output import VideoOutputStream, VideoStepResult +from flashdreams.infra.video_output import VideoOutputStream from flashdreams.runtime import ( CanonicalInputSchema, InferenceConfig, InferenceInput, InferenceInputSchema, InputField, + Mp4VideoOutputTarget, OutputArtifact, ) from flashdreams.runtime.interfaces import InferenceRuntime, InferenceSession @@ -54,6 +54,7 @@ LingbotInputMapping, load_camera_trace, ) +from lingbot.model_session import LingbotModelSessionCore LINGBOT_MODEL_ID = "lingbot" DEFAULT_LINGBOT_PRESET = "lingbot-world-fast-taehv-window15-sink3" @@ -302,13 +303,21 @@ def validate_config(self, config: InferenceConfig) -> None: def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: self.validate_config(config) + output_layout = config.runtime_options.get("output_layout", "tchw") + if not isinstance(output_layout, str) or output_layout not in { + "tchw", + "btchw", + "bcthw", + "bvtchw", + }: + raise ValueError(f"Unsupported Lingbot output layout: {output_layout!r}.") return self._runtime_factory( config=config, options=LingbotReplayRuntimeOptions( pipeline_config=self.pipeline_config(config), pipeline=config.runtime_options.get("pipeline"), pipeline_factory=self._pipeline_factory, - output_layout=str(config.runtime_options.get("output_layout", "tchw")), + output_layout=cast(VideoTensorLayout, output_layout), ), ) @@ -419,10 +428,16 @@ def __init__( self.output_layout = output_layout self.dtype = torch.bfloat16 self._closed = False - self._step_index = 0 self._frame_start = 0 self._active_prompt = session_inputs.prompt - self._cache = self._initialize_cache() + self._model_session = LingbotModelSessionCore( + pipeline=pipeline, + output_stream_factory=lambda: VideoOutputStream( + postprocess_stream=None, + output_layout=self.output_layout, + ), + ) + self._reset_model_session() if self.device.type == "cuda" and torch.cuda.is_available(): torch.cuda.synchronize(device=self.device) if dist.is_initialized(): @@ -431,16 +446,17 @@ def __init__( def next_step_request(self) -> StepRequest | None: if self._closed: return None - if self._step_index >= self.inputs.total_blocks: + step_index = self._model_session.step_index + if step_index >= self.inputs.total_blocks: return None - num_frames = int(self.pipeline.get_num_output_frames(self._step_index)) + num_frames = self._model_session.next_num_frames() frame_end = self._frame_start + num_frames total_frames = self.inputs.total_camera_frames if total_frames is not None and frame_end > total_frames: return None fps = self.inputs.fps return StepRequest( - step_index=self._step_index, + step_index=step_index, # The window is what lets a mapping slice user events for exactly # this chunk instead of replaying the whole session history. user_input_window=TimeWindow( @@ -457,8 +473,8 @@ def step(self, inputs: InferenceInput) -> StepResult: if self._closed: raise RuntimeError("Lingbot replay session is closed.") - step_index = self._step_index - num_frames = int(self.pipeline.get_num_output_frames(step_index)) + step_index = self._model_session.step_index + num_frames = self._model_session.next_num_frames() self._apply_global_conditioning_update(inputs) camera_poses = _require_step_tensor( inputs, @@ -485,49 +501,23 @@ def step(self, inputs: InferenceInput) -> StepResult: poses=camera_poses.to(device=self.device, dtype=torch.float32), world_scale=self.inputs.world_scale, ) - start_t = time.perf_counter() - video_chunk = self.pipeline.generate( - autoregressive_index=step_index, - cache=self._cache, - input=camctrl_input, - ) - stats = self.pipeline.finalize( - autoregressive_index=step_index, - cache=self._cache, - ) - elapsed_s = time.perf_counter() - start_t - self._step_index += 1 - self._frame_start = frame_end - - metrics = _numeric_stats(stats) - metrics.setdefault("model_step_s", elapsed_s) - return StepResult( - step_index=step_index, - output=VideoStepResult.from_video_chunk( - chunk_index=step_index, - video_chunk=video_chunk, - layout=self.output_layout, - stats=metrics, - ), - frame_count=num_frames, + result = self._model_session.step( + camctrl_input, output_window=TimeWindow( start_s=frame_start / self.inputs.fps, end_s=frame_end / self.inputs.fps, ), - metrics=metrics, ) + self._frame_start = frame_end + return result def reset(self, inputs: InferenceInput | None = None) -> None: if inputs is not None: session_inputs = session_inputs_from_inference_input(inputs) if session_inputs != self.inputs: raise ValueError("Lingbot replay reset cannot swap inputs.") - cache = getattr(self, "_cache", None) - if cache is not None: - del self._cache self._active_prompt = self.inputs.prompt - self._cache = self._initialize_cache() - self._step_index = 0 + self._reset_model_session() self._frame_start = 0 def _apply_global_conditioning_update(self, inputs: InferenceInput) -> None: @@ -551,18 +541,19 @@ def _apply_global_conditioning_update(self, inputs: InferenceInput) -> None: ) self.pipeline._ensure_oneshot_encoders_loaded() embeddings = self.pipeline.text_encoder([prompt]).to(device=self.device) - replace_text_embeddings(self._cache.transformer_cache, embeddings) + self._model_session.replace_text_embeddings(embeddings) self._active_prompt = prompt if self.is_rank_zero: - logger.info("Lingbot text context updated at step {}", self._step_index) + logger.info( + "Lingbot text context updated at step {}", + self._model_session.step_index, + ) def close(self) -> None: self._closed = True - cache = getattr(self, "_cache", None) - if cache is not None: - del self._cache + self._model_session.close() - def _initialize_cache(self) -> Any: + def _reset_model_session(self) -> None: first_frames = load_first_frame_tensor( self.inputs.first_frame_path, pixel_height=self.inputs.pixel_height, @@ -572,9 +563,9 @@ def _initialize_cache(self) -> Any: interpolation="cubic", install_hint=_INSTALL_HINT, ) - return self.pipeline.initialize_cache( - text=[self.inputs.prompt], - image=first_frames, + self._model_session.reset( + prompt=self.inputs.prompt, + first_frames=first_frames, ) @@ -611,49 +602,61 @@ class LingbotRunnerOutputTarget: fps: int | float install_hint: str = _INSTALL_HINT _opened: bool = False + _mp4_target: Mp4VideoOutputTarget | None = None def open(self) -> None: + video_path = runner_artifact_path(self.output_dir, self.runner_name, "mp4") + self._mp4_target = Mp4VideoOutputTarget( + output_path=video_path, + fps=self.fps, + output_layout=self.output_stream.output_layout, + install_hint=self.install_hint, + ) + self._mp4_target.open() self._opened = True def write(self, result: StepResult) -> None: if not self._opened: raise RuntimeError("Cannot write to a closed Lingbot output target.") - video_result = result.output - if not isinstance(video_result, VideoStepResult): + if result.layout is None: raise TypeError( - "LingbotRunnerOutputTarget requires VideoStepResult output, " - f"got {type(video_result).__name__}." + "LingbotRunnerOutputTarget requires a video StepResult with layout." ) - self.output_stream.process( - video_result.video_chunk, - autoregressive_index=video_result.chunk_index, - stats=video_result.stats or dict(result.metrics), + if self._mp4_target is None: + raise RuntimeError("Lingbot MP4 target is not open.") + processed = self.output_stream.process( + result.video_chunk, + autoregressive_index=result.step_index, + metrics=result.metrics, + metadata=result.metadata, + output_window=result.output_window, ) + self._mp4_target.write(processed) def close(self) -> tuple[OutputArtifact, ...]: self._opened = False - artifacts: list[OutputArtifact] = [] - video_path = runner_artifact_path(self.output_dir, self.runner_name, "mp4") - video_path = self.output_stream.finish_to_mp4( - video_path, - fps=self.fps, - install_hint=self.install_hint, - ) - if video_path is None: + target = self._mp4_target + self._mp4_target = None + if target is None: return () + tail = self.output_stream.finish() + if tail is not None: + target.write(tail) + artifacts = list(target.close()) + if not artifacts: + return () + video_path = Path(artifacts[0].uri) logger.info( "[{}] wrote video -> {}", self.runner_name, video_path.resolve(), ) - artifacts.append( - OutputArtifact(kind="video/mp4", uri=str(video_path.resolve())) - ) - if self.output_stream.stats_history: + stats_history = artifacts[0].metadata.get("stats_history", ()) + if stats_history: stats_path = write_runner_stats( self.output_dir, self.runner_name, - self.output_stream.stats_history, + list(stats_history), ) logger.info( "[{}] wrote per-AR-step stats -> {}", @@ -930,7 +933,9 @@ def build_lingbot_webrtc_runtime_config( return _apply_webrtc_runtime_options(runtime_config, runtime_options or {}) -def _apply_webrtc_runtime_options(runtime_config: Any, options: Mapping[str, Any]) -> Any: +def _apply_webrtc_runtime_options( + runtime_config: Any, options: Mapping[str, Any] +) -> Any: overrides: dict[str, Any] = {} for name in ( "world_scale", @@ -952,16 +957,6 @@ def _default_pipeline_factory(pipeline_config: Any, device: str) -> Any: return pipeline_config.setup().to(device=device).eval() -def _numeric_stats(stats: Any) -> dict[str, float | int]: - if not isinstance(stats, Mapping): - return {} - return { - str(key): value - for key, value in stats.items() - if isinstance(value, (float, int)) and not isinstance(value, bool) - } - - def _resolve_prompt( value: Mapping[str, Any], *, @@ -984,15 +979,19 @@ def _resolve_example_data_default(value: Mapping[str, Any]) -> bool: explicit = value.get("example_data") if explicit is not None: return _bool_value(explicit) - return not ( - _has_nonempty_value(value, FIELD_FIRST_FRAME_PATH) - or _has_nonempty_value(value, "image_path") - ) or not ( - _has_nonempty_value(value, FIELD_CAMERA_POSES_PATH) - or _has_nonempty_value(value, "pose_path") - ) or not ( - _has_nonempty_value(value, FIELD_CAMERA_INTRINSICS_PATH) - or _has_nonempty_value(value, "intrinsic_path") + return ( + not ( + _has_nonempty_value(value, FIELD_FIRST_FRAME_PATH) + or _has_nonempty_value(value, "image_path") + ) + or not ( + _has_nonempty_value(value, FIELD_CAMERA_POSES_PATH) + or _has_nonempty_value(value, "pose_path") + ) + or not ( + _has_nonempty_value(value, FIELD_CAMERA_INTRINSICS_PATH) + or _has_nonempty_value(value, "intrinsic_path") + ) ) @@ -1029,7 +1028,9 @@ def _require_path_value(value: Path | None, *, label: str) -> Path: def _require_existing_replay_paths(replay_inputs: LingbotReplayInputs) -> None: _require_existing_path(replay_inputs.first_frame_path, label=FIELD_FIRST_FRAME_PATH) - _require_existing_path(replay_inputs.camera_poses_path, label=FIELD_CAMERA_POSES_PATH) + _require_existing_path( + replay_inputs.camera_poses_path, label=FIELD_CAMERA_POSES_PATH + ) _require_existing_path( replay_inputs.camera_intrinsics_path, label=FIELD_CAMERA_INTRINSICS_PATH, diff --git a/integrations/lingbot/lingbot/webrtc/server.py b/integrations/lingbot/lingbot/webrtc/server.py index 1cd0bde1e..29ce67fc8 100644 --- a/integrations/lingbot/lingbot/webrtc/server.py +++ b/integrations/lingbot/lingbot/webrtc/server.py @@ -32,12 +32,14 @@ from flashdreams.core.distributed import ( init as distributed_init, ) +from flashdreams.runtime import InferenceConfig from flashdreams.serving.network import get_external_ip from flashdreams.serving.webrtc.bootstrap import ( configure_logging, initialize_cuda_distributed, run_webrtc_server, ) +from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager from flashdreams.serving.webrtc.server import ( SESSION_MANAGER_KEY, SessionBusyError, @@ -48,7 +50,6 @@ from flashdreams.serving.webrtc.server import ( close_package_resources as _close_package_resources, ) -from flashdreams.runtime import InferenceConfig from lingbot.example_data import ( EXAMPLE_DATA_AVAILABLE_IDXS, ensure_example_data_downloaded, @@ -60,9 +61,11 @@ ) from lingbot.webrtc.session import ( LingbotImagePayload, + LingbotInferenceRuntime, LingbotRuntimeConfig, LingbotSessionInput, - LingbotWebRTCSessionManager, + LingbotWebRTCSessionController, + create_lingbot_webrtc_session_manager, normalize_prompt_text, normalize_text_events, ) @@ -73,14 +76,20 @@ MAX_PROMPT_CHARS = 2_000 -class LingbotSessionManager(WebRTCSessionManager, Protocol): +class LingbotSessionController(Protocol): def get_initial_scene(self) -> dict[str, object]: ... def get_first_frame(self) -> LingbotImagePayload: ... def set_pending_session_input(self, session_input: LingbotSessionInput) -> None: ... -def _get_lingbot_manager(app: web.Application) -> LingbotSessionManager: - return cast(LingbotSessionManager, app[SESSION_MANAGER_KEY]) +LINGBOT_SESSION_CONTROLLER_KEY = web.AppKey( + "lingbot_session_controller", + LingbotSessionController, +) + + +def _get_lingbot_controller(app: web.Application) -> LingbotSessionController: + return app[LINGBOT_SESSION_CONTROLLER_KEY] def parse_args() -> argparse.Namespace: @@ -172,8 +181,18 @@ def create_app( *, request_session_url: str, session_manager: WebRTCSessionManager | None = None, + session_controller: LingbotSessionController | None = None, ) -> web.Application: - manager = session_manager or LingbotWebRTCSessionManager() + manager = session_manager or create_lingbot_webrtc_session_manager() + if session_controller is None and not isinstance(manager, BaseWebRTCSessionManager): + # Lightweight server tests may provide one object for both protocols. + session_controller = cast(LingbotSessionController, manager) + + def configure_app(app: web.Application) -> None: + configure_lingbot_webrtc_app( + app, + session_controller=session_controller, + ) return create_packaged_webrtc_app( web_resource=WEB_DIR_RESOURCE, @@ -181,28 +200,49 @@ def create_app( session_manager=manager, preload_name="Lingbot", request_session_url=request_session_url, - configure_app=configure_lingbot_webrtc_app, + configure_app=configure_app, as_file_fn=as_file, create_app_fn=create_webrtc_app, cleanup_callback=_close_package_resources, ) -def configure_lingbot_webrtc_app(app: web.Application) -> None: +def configure_lingbot_webrtc_app( + app: web.Application, + *, + session_controller: LingbotSessionController | None = None, +) -> None: """Register Lingbot-only initial-scene and session-input routes.""" + if session_controller is None: + manager = app[SESSION_MANAGER_KEY] + if not isinstance(manager, BaseWebRTCSessionManager): + raise TypeError( + "Lingbot routes require BaseWebRTCSessionManager or an " + "explicit session_controller." + ) + session_controller = LingbotWebRTCSessionController( + cast( + BaseWebRTCSessionManager[ + LingbotInferenceRuntime, + LingbotRuntimeConfig, + ], + manager, + ) + ) + app[LINGBOT_SESSION_CONTROLLER_KEY] = session_controller app.router.add_get("/api/session/initial_scene", _initial_scene) app.router.add_get("/api/session/first_frame", _first_frame) app.router.add_post("/api/session/input", _session_input) async def _initial_scene(request: web.Request) -> web.StreamResponse: - manager = _get_lingbot_manager(request.app) - return web.json_response(manager.get_initial_scene()) + controller = _get_lingbot_controller(request.app) + return web.json_response(controller.get_initial_scene()) async def _first_frame(request: web.Request) -> web.StreamResponse: - manager = _get_lingbot_manager(request.app) - payload = await asyncio.to_thread(manager.get_first_frame) + controller = _get_lingbot_controller(request.app) + payload = await asyncio.to_thread(controller.get_first_frame) if not isinstance(payload, LingbotImagePayload): raise web.HTTPInternalServerError(reason="Invalid Lingbot first-frame payload.") return web.Response(body=payload.data, content_type=payload.content_type) @@ -321,7 +361,7 @@ async def _session_input(request: web.Request) -> web.StreamResponse: ) ) - manager = _get_lingbot_manager(request.app) + controller = _get_lingbot_controller(request.app) session_input = LingbotSessionInput( prompt=prompt or None, first_frame_image_bytes=image_bytes, @@ -330,12 +370,12 @@ async def _session_input(request: web.Request) -> web.StreamResponse: text_events=normalized_text_events, ) try: - await asyncio.to_thread(manager.set_pending_session_input, session_input) + await asyncio.to_thread(controller.set_pending_session_input, session_input) except SessionBusyError as exc: raise web.HTTPConflict(reason=str(exc)) from exc except ValueError as exc: raise web.HTTPBadRequest(reason=str(exc)) from exc - return web.json_response(manager.get_initial_scene()) + return web.json_response(controller.get_initial_scene()) def build_runtime_config( @@ -421,7 +461,7 @@ def main() -> None: device_override=str(runtime_device), context_parallel_size=context_parallel_size, ) - session_manager = LingbotWebRTCSessionManager( + session_manager = create_lingbot_webrtc_session_manager( runtime_config=runtime_config, fps=args.fps, ) diff --git a/integrations/lingbot/lingbot/webrtc/session.py b/integrations/lingbot/lingbot/webrtc/session.py index e9b79c283..a45c3b2e1 100644 --- a/integrations/lingbot/lingbot/webrtc/session.py +++ b/integrations/lingbot/lingbot/webrtc/session.py @@ -17,7 +17,6 @@ from __future__ import annotations -import asyncio import http.client import io import ipaddress @@ -36,28 +35,24 @@ import torch.distributed as dist from loguru import logger -from flashdreams.core.distributed.rank_orchestration import ( - RankCoordinator, - distributed_op, -) +from flashdreams.core.distributed.rank_orchestration import distributed_op from flashdreams.core.io.disk import default_flashdreams_cache_dir from flashdreams.infra.config import derive_config -from flashdreams.infra.video_output import VideoOutputStream, VideoStepResult +from flashdreams.infra.video_output import VideoOutputStream +from flashdreams.runtime import StepResult from flashdreams.serving.webrtc.controls import ( CameraPoseIntegrator, PoseSegment, ) -from flashdreams.serving.webrtc.encoders import ( - EncoderBackend, - VideoEncoder, - select_encoder, -) +from flashdreams.serving.webrtc.encoders import EncoderBackend from flashdreams.serving.webrtc.manager import ( DEFAULT_CLIENT_LIVENESS_TIMEOUT_S, BaseWebRTCSessionManager, - ManagedWebRTCSession, WebRTCControlSignal, ) +from flashdreams.serving.webrtc.runtime import ( + ThreadAffineDistributedWebRTCRuntime, +) from flashdreams.serving.webrtc.server import SessionBusyError from flashdreams.runtime.canonical import InputCanonicalizer from flashdreams.runtime.inputs import ( @@ -74,6 +69,7 @@ TextEventSelection, ) from lingbot.encoder.utils import preprocess_example_poses +from lingbot.model_session import LingbotModelSessionCore _INTRINSICS_REFERENCE_HEIGHT = 480 _INTRINSICS_REFERENCE_WIDTH = 832 @@ -585,34 +581,25 @@ def normalize_text_events(raw_events: object) -> tuple[TextEventSpec, ...]: return tuple(text_events) -class LingbotInferenceRuntime: +class LingbotInferenceRuntime( + ThreadAffineDistributedWebRTCRuntime[ + LingbotRuntimeConfig, + LingbotSessionInput, + ] +): """Single-session Lingbot runtime with action-bound chunk generation.""" def __init__(self, config: LingbotRuntimeConfig | None = None) -> None: - self.config = config or LingbotRuntimeConfig() - self.MASTER_RANK = 0 - self.rank = 0 if not dist.is_initialized() else dist.get_rank() - - control_device = torch.device(self.config.device) - if control_device.type == "cuda" and control_device.index is None: - control_device = torch.device( - f"cuda:{torch.cuda.current_device()}" - if torch.cuda.is_available() - else "cuda:0" - ) + super().__init__( + config=config or LingbotRuntimeConfig(), + runtime_error_type=LingbotRuntimeError, + thread_name="lingbot-webrtc-runtime", + ) self.pose_integrator = CameraPoseIntegrator() - self.autoregressive_index = 0 - self._output_stream = VideoOutputStream( - postprocess_stream=None, - output_layout="tchw", - collect_output=False, - move_to_cpu=False, - ) - self._device: torch.device | None = None self._pipeline: Any | None = None - self._cache: Any | None = None + self._model_session: LingbotModelSessionCore | None = None self._base_intrinsics: torch.Tensor | None = None self._first_frames: torch.Tensor | None = None self._prompt: str | None = None @@ -624,55 +611,6 @@ def __init__(self, config: LingbotRuntimeConfig | None = None) -> None: self._input_canonicalizer: InputCanonicalizer | None = None self._sync_step_lock = threading.Lock() self._world_scale = 1.0 - self._video_encoder: VideoEncoder | None = None - self._closed = False - - self._step_lock = asyncio.Lock() - self.rank_coordinator = RankCoordinator( - device=control_device, - signal_type=WebRTCControlSignal, - is_master=self.is_master, - master_rank=self.MASTER_RANK, - ) - self.rank_coordinator.register_distributed_ops(self) - - @property - 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 LingbotRuntimeError( - "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) - - def send_exit_signal(self) -> None: - if self.is_master: - self.rank_coordinator.send_exit(exit_signal=WebRTCControlSignal.EXIT) - - async def initialize(self) -> None: - if self._pipeline is not None: - return - await asyncio.to_thread(self._initialize_sync_all_ranks) - - async def reset_for_new_session( - self, session_input: LingbotSessionInput | None = None - ) -> None: - if self._closed: - raise LingbotRuntimeError("Runtime is closed.") - if self._pipeline is None: - raise LingbotRuntimeError("Runtime is not initialized.") - await asyncio.to_thread(self._reset_rollout_sync_all_ranks, session_input) - - async def close(self) -> None: - self._closed = True - await asyncio.to_thread(self._close_sync_all_ranks) async def trigger_event( self, *, event_id: str, state: str = "trigger" @@ -680,158 +618,20 @@ async def trigger_event( """Activate or clear a precomputed text event for subsequent chunks.""" if self._closed: raise LingbotRuntimeError("Runtime is closed.") - if self._pipeline is None or self._cache is None: + if self._pipeline is None or self._model_session is None: raise LingbotRuntimeError("Runtime is not initialized.") event_id, state = self._validate_event_request(event_id=event_id, state=state) async with self._step_lock: if self._closed: raise LingbotRuntimeError("Runtime is closed.") - if self._pipeline is None or self._cache is None: + if self._pipeline is None or self._model_session is None: raise LingbotRuntimeError("Runtime is not initialized.") - return await asyncio.to_thread( + return await self._worker.call( self._trigger_event_sync_all_ranks, event_id, state, ) - async def start_inference_session(self) -> LingbotWebRTCInferenceSession: - """Return an ``InferenceSession`` view of the current rollout. - - The shared manager canonicalizes raw key and text events and maps them - into per-step model inputs before stepping the session. - """ - if self._closed: - raise LingbotRuntimeError("Runtime is closed.") - if self._input_mapping is None: - raise LingbotRuntimeError( - "Runtime input mapping is not initialized; reset the rollout first." - ) - return LingbotWebRTCInferenceSession(runtime=self) - - @property - def input_mapping(self) -> LingbotInputMapping: - if self._input_mapping is None: - raise LingbotRuntimeError("Runtime input mapping is not initialized.") - return self._input_mapping - - @property - def input_canonicalizer(self) -> InputCanonicalizer: - if self._input_canonicalizer is None: - raise LingbotRuntimeError("Runtime canonicalizer is not initialized.") - return self._input_canonicalizer - - @property - def input_source_schema(self) -> UserInputSchema: - return LINGBOT_WEBRTC_SOURCE_SCHEMA - - def validate_user_event( - self, *, event_type: str, payload: dict[str, Any] - ) -> dict[str, Any] | None: - """Validate one raw WebRTC user event before it is acknowledged.""" - if event_type != "text_event": - return payload - event_id_value = payload.get("event_id") - event_id = "" if event_id_value is None else str(event_id_value) - state = str(payload.get("state", "trigger")).strip().lower() or "trigger" - event_id, state = self._validate_event_request(event_id=event_id, state=state) - clears = state in {"clear", "release", "off", "none"} - return {"event_id": None if clears else event_id, "state": state} - - def _build_input_layers_sync( - self, text_events: tuple[TextEventSpec, ...] - ) -> None: - """Build the canonicalizer and mapping for the current rollout. - - A rollout can be reset before intrinsics are resolved; the mapping is - then left unbuilt and ``start_inference_session`` reports it. - """ - if self._base_intrinsics is None: - self._input_mapping = None - self._input_canonicalizer = None - return - self._input_canonicalizer = InputCanonicalizer( - [KeyboardToCameraCommand(), TextEventSelection()] - ) - # Mapping runs on the transport's event-loop thread, so hand it a CPU - # copy rather than the device tensor used inside generation. - self._input_mapping = LingbotInputMapping( - fps=int(self.config.fps), - base_intrinsics=self._base_intrinsics.detach().reshape(4).cpu(), - world_scale=self._world_scale or 1.0, - text_event_prompts={ - event.event_id: event.prompt for event in text_events - }, - ) - self._input_mapping.set_base_prompt(self._prompt or "") - - def _next_step_request_sync(self) -> StepRequest: - """Describe the next chunk for the mapping. - - The manager overrides ``user_input_window`` with its own clock; the - frame counter here only tells the mapping how much trajectory to build. - """ - num_frames = self.peek_next_chunk_num_frames() - return StepRequest( - step_index=self.autoregressive_index, - metadata={ - "num_frames": num_frames, - "frame_start": self.autoregressive_index * num_frames, - }, - ) - - def _step_blocking(self, inputs: InferenceInput) -> StepResult: - """Run one mapped step. Called from the manager's executor thread.""" - if self._closed: - raise LingbotRuntimeError("Session is closed.") - with self._sync_step_lock: - if self._closed: - raise LingbotRuntimeError("Session is closed.") - return self._step_sync_all_ranks(inputs) - - async def generate_chunk( - self, - *, - segments: list[PoseSegment], - frame_times: list[float], - ) -> VideoStepResult: - """Generate one autoregressive chunk from a piecewise-constant timeline. - - Args: - segments: Piecewise-constant keyboard-state segments covering the - chunk's virtual-time window. - frame_times: Virtual times at which to sample the camera pose; must - have length equal to :meth:`peek_next_chunk_num_frames` at call - time. - - Returns: - Video chunk and post-generation pipeline stats. - - Raises: - LingbotRuntimeError: Runtime is closed or not initialized. - """ - if self._closed: - raise LingbotRuntimeError("Session is closed.") - if self._pipeline is None or self._cache is None: - raise LingbotRuntimeError("Runtime is not initialized.") - - async with self._step_lock: - if self._closed: - raise LingbotRuntimeError("Session is closed.") - return await asyncio.to_thread( - self._generate_chunk_sync_all_ranks, segments, frame_times - ) - - def peek_next_chunk_num_frames(self) -> int: - """Return the number of frames the next chunk's pipeline call will emit. - - Master-only read with no distributed broadcast; safe to call from - the master rank's asyncio event loop to size the resampler's - per-chunk request. - """ - if self._pipeline is None: - raise LingbotRuntimeError("Runtime is not initialized.") - return int(self._pipeline.get_num_output_frames(self.autoregressive_index)) - # Arbitrary index well past the AR-step transient; for the Wan/lingbot # pipelines used here the per-step count is constant for any index # ``>= 1`` (only AR 0 emits fewer frames due to causal first-frame @@ -840,7 +640,20 @@ def peek_next_chunk_num_frames(self) -> int: # boundary of that transient. _STEADY_STATE_AR_PROBE_INDEX: int = 1000 - def peek_steady_chunk_num_frames(self) -> int: + def _is_runtime_initialized(self) -> bool: + return self._pipeline is not None and self._model_session is not None + + def _runtime_step_index(self) -> int: + if self._model_session is None: + raise LingbotRuntimeError("Runtime is not initialized.") + return self._model_session.step_index + + def _next_input_frame_count(self) -> int: + if self._model_session is None: + raise LingbotRuntimeError("Runtime is not initialized.") + return self._model_session.next_num_frames() + + def _steady_output_frame_count(self) -> int: """Return the steady-state per-chunk frame count. AR step 0 emits *fewer* frames than every subsequent step @@ -859,30 +672,6 @@ def peek_steady_chunk_num_frames(self) -> int: self._pipeline.get_num_output_frames(self._STEADY_STATE_AR_PROBE_INDEX) ) - @distributed_op(WebRTCControlSignal.INITIALIZE) - def _initialize_sync_all_ranks(self) -> None: - self._initialize_sync() - - @distributed_op(WebRTCControlSignal.RESET_SESSION) - def _reset_rollout_sync_all_ranks( - self, session_input: LingbotSessionInput | None = None - ) -> None: - self._reset_rollout_sync(session_input=session_input) - - @distributed_op(WebRTCControlSignal.ACTION_STEP) - def _generate_chunk_sync_all_ranks( - self, - segments: list[PoseSegment], - frame_times: list[float], - ) -> VideoStepResult: - return self._generate_one_chunk_sync(segments=segments, frame_times=frame_times) - - @distributed_op(WebRTCControlSignal.SESSION_STEP) - def _step_sync_all_ranks(self, inputs: InferenceInput) -> StepResult: - # distributed_op broadcasts rank-0 arguments, so worker ranks receive - # the mapped trajectory rather than recomputing it from raw events. - return self._step_sync(inputs) - @distributed_op(WebRTCControlSignal.EVENT) def _trigger_event_sync_all_ranks( self, @@ -891,10 +680,6 @@ def _trigger_event_sync_all_ranks( ) -> dict[str, str | None]: return self._trigger_event_sync(event_id=event_id, state=state) - @distributed_op(WebRTCControlSignal.CLOSE) - def _close_sync_all_ranks(self) -> None: - self._close_sync() - def _initialize_sync(self) -> None: if self._pipeline is not None: return @@ -910,7 +695,6 @@ def _initialize_sync(self) -> None: ) pipeline_config_base = pipeline_configs[self.config.config_name] - self._device = torch.device(self.config.device) if self._device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA is required for Lingbot runtime.") @@ -931,39 +715,16 @@ def _initialize_sync(self) -> None: ), ) self._pipeline = pipeline_config.setup().to(device=self._device) + self._model_session = LingbotModelSessionCore( + pipeline=self._pipeline, + output_stream_factory=lambda: VideoOutputStream( + postprocess_stream=None, + output_layout="tchw", + ), + ) self._reset_rollout_sync() self._initialize_video_encoder_sync() - def _initialize_video_encoder_sync(self) -> None: - """Select the video encoder for this runtime.""" - if not self.is_master: - return - 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 torch.device(self.config.device) - ) - backend: EncoderBackend = self.config.encoder_backend - if device.type != "cuda" and backend == "auto": - backend = "default" - if device.type != "cuda" and backend == "nvenc": - raise LingbotRuntimeError( - "encoder_backend='nvenc' requires a CUDA runtime device." - ) - gpu_id = device.index if device.index is not None else 0 - self._video_encoder = select_encoder( - backend=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 _encode_text_embeddings_sync(self, texts: list[str]) -> torch.Tensor: if self._pipeline is None: raise LingbotRuntimeError("Runtime pipeline is not initialized.") @@ -996,8 +757,6 @@ def _precompute_event_embeddings_sync( } def _build_base_intrinsics(self) -> torch.Tensor: - if self._device is None: - raise LingbotRuntimeError("Runtime device is not initialized.") intrinsics_path = self.config.example_data_dir / self.config.intrinsics_filename if self.config.default_intrinsics is not None: intrinsics = np.asarray(self.config.default_intrinsics, dtype=np.float32) @@ -1112,8 +871,6 @@ def _load_uploaded_first_frame_rgb(self, image_bytes: bytes) -> np.ndarray: ) def _first_frame_to_tensor(self, image_rgb: np.ndarray) -> torch.Tensor: - if self._device is None: - raise LingbotRuntimeError("Runtime device is not initialized.") # Bicubic to match the upstream Lingbot World demo / generate_fast.py # (which uses ``F.interpolate(mode='bicubic')`` over the ``[-1, 1]`` # tensor); bilinear here would give a different first-frame VAE latent. @@ -1169,13 +926,9 @@ def _prepare_session_input_state( def _reset_rollout_sync( self, session_input: LingbotSessionInput | None = None ) -> None: - if self._pipeline is None: + if self._pipeline is None or self._model_session is None: raise LingbotRuntimeError("Runtime pipeline is not initialized.") - if self._cache is not None: - del self._cache - self._cache = None - self._prepare_session_input_state(session_input) text_events = ( session_input.text_events @@ -1187,26 +940,19 @@ def _reset_rollout_sync( raise LingbotRuntimeError("Runtime input state is not initialized.") self.pose_integrator = CameraPoseIntegrator() - self.autoregressive_index = 0 self._active_event_id = None - self._cache = self._pipeline.initialize_cache( - text=[self._prompt], - image=self._first_frames, + self._model_session.reset( + prompt=self._prompt, + first_frames=self._first_frames, ) # Rebuilt per rollout: the mapping carries the rollout's text-event # catalog, base prompt, and pose integrator state. self._build_input_layers_sync(text_events) def _replace_rollout_text_embeddings(self, text_embeddings: torch.Tensor) -> None: - if self._pipeline is None or self._cache is None: + if self._pipeline is None or self._model_session is None: raise LingbotRuntimeError("Runtime is not initialized.") - transformer = self._pipeline.diffusion_model.transformer - replace_text_embeddings = getattr(transformer, "replace_text_embeddings", None) - if not callable(replace_text_embeddings): - raise LingbotRuntimeError( - "Current pipeline does not support runtime text-event swapping." - ) - replace_text_embeddings(self._cache.transformer_cache, text_embeddings) + self._model_session.replace_text_embeddings(text_embeddings) def _validate_event_request(self, *, event_id: str, state: str) -> tuple[str, str]: state = state.strip().lower() or "trigger" @@ -1240,9 +986,9 @@ def _trigger_event_sync( return {"active_event_id": event_id} def _close_sync(self) -> None: - cache = self._cache + model_session = self._model_session pipeline = self._pipeline - self._cache = None + self._model_session = None self._pipeline = None self._base_intrinsics = None self._first_frames = None @@ -1250,16 +996,12 @@ def _close_sync(self) -> None: self._base_text_embeddings = None self._event_embeddings = {} self._active_event_id = None - if self._video_encoder is not None: - self._video_encoder.close() - self._video_encoder = None - - if cache is not None: - del cache + if model_session is not None: + model_session.close() if pipeline is not None: del pipeline - if self._device is not None and self._device.type == "cuda": + if self._device.type == "cuda": torch.cuda.synchronize(device=self._device) torch.cuda.empty_cache() @@ -1268,28 +1010,22 @@ def _generate_one_chunk_sync( *, segments: list[PoseSegment], frame_times: list[float], - ) -> VideoStepResult: + ) -> StepResult: if ( self._pipeline is None - or self._cache is None + or self._model_session is None or self._base_intrinsics is None ): raise LingbotRuntimeError("Runtime is not initialized.") - if self._device is None: - raise LingbotRuntimeError("Runtime device is not initialized.") - - num_frames = int( - self._pipeline.get_num_output_frames(self.autoregressive_index) - ) + step_index = self._runtime_step_index() + num_frames = int(self._pipeline.get_num_output_frames(step_index)) if len(frame_times) != num_frames: raise LingbotRuntimeError( f"Expected {num_frames} frame_times for " - f"chunk={self.autoregressive_index}, got {len(frame_times)}." + f"chunk={step_index}, got {len(frame_times)}." ) if not segments: - raise LingbotRuntimeError( - f"Chunk={self.autoregressive_index} received empty segments." - ) + raise LingbotRuntimeError(f"Chunk={step_index} received empty segments.") poses = self.pose_integrator.integrate_chunk( segments=segments, frame_times=frame_times ) @@ -1326,24 +1062,13 @@ def _generate_from_camera_inputs( poses=poses.to(device=self._device, dtype=torch.float32), world_scale=self._world_scale, ) - video_chunk = self._pipeline.generate( - autoregressive_index=self.autoregressive_index, - cache=self._cache, - input=camctrl_input, - ) - stats = self._pipeline.finalize(self.autoregressive_index, self._cache) - result = self._output_stream.make_step_result( - video_chunk, - autoregressive_index=self.autoregressive_index, - stats=stats, - sync_device=self._device, - ) - if result.num_frames != num_frames: - raise LingbotRuntimeError( - f"Expected generated chunk to contain {num_frames} frames, " - f"got {result.num_frames}." + try: + result = self._model_session.step( + camctrl_input, + metadata={"active_event_id": self._active_event_id}, ) - self.autoregressive_index += 1 + except RuntimeError as exc: + raise LingbotRuntimeError(str(exc)) from exc return result def _step_sync(self, inputs: InferenceInput) -> StepResult: @@ -1456,68 +1181,57 @@ def close(self) -> None: return None -_ManagedLingbotSession = ManagedWebRTCSession - +def create_lingbot_webrtc_session_manager( + *, + runtime: LingbotInferenceRuntime | None = None, + runtime_config: LingbotRuntimeConfig | None = None, + fps: int | None = None, + client_liveness_timeout_s: float = DEFAULT_CLIENT_LIVENESS_TIMEOUT_S, +) -> BaseWebRTCSessionManager[LingbotInferenceRuntime, LingbotRuntimeConfig]: + """Configure the shared WebRTC manager for the Lingbot runtime.""" + runtime_config = runtime_config or getattr(runtime, "config", None) + if not isinstance(runtime_config, LingbotRuntimeConfig): + runtime_config = LingbotRuntimeConfig() + fps = runtime_config.fps if fps is None else fps + if fps <= 0: + raise ValueError("fps must be > 0") + runtime = runtime or LingbotInferenceRuntime(config=runtime_config) + return BaseWebRTCSessionManager( + runtime=runtime, + runtime_config=runtime_config, + fps=fps, + identity=runtime_config.config_name, + busy_message="A Lingbot session is already active.", + warmup_label="Lingbot WebRTC", + client_liveness_timeout_s=client_liveness_timeout_s, + ) -class LingbotWebRTCSessionManager( - BaseWebRTCSessionManager[LingbotInferenceRuntime, LingbotRuntimeConfig] -): - """Owns one active WebRTC session and forwards actions into Lingbot runtime.""" - _busy_message = "A Lingbot session is already active." - _warmup_label = "Lingbot WebRTC" - _runtime_error_types = (LingbotRuntimeError,) +class LingbotWebRTCSessionController: + """Own Lingbot browser inputs and preview data outside the transport manager.""" def __init__( self, - *, - runtime: LingbotInferenceRuntime | None = None, - runtime_config: LingbotRuntimeConfig | None = None, - fps: int | None = None, - client_liveness_timeout_s: float = DEFAULT_CLIENT_LIVENESS_TIMEOUT_S, + manager: BaseWebRTCSessionManager[ + LingbotInferenceRuntime, + LingbotRuntimeConfig, + ], ) -> None: - runtime_config = runtime_config or getattr(runtime, "config", None) - if not isinstance(runtime_config, LingbotRuntimeConfig): - runtime_config = LingbotRuntimeConfig() - fps = runtime_config.fps if fps is None else fps - if fps <= 0: - raise ValueError("fps must be > 0") - runtime = runtime or LingbotInferenceRuntime(config=runtime_config) - super().__init__( - runtime=runtime, - runtime_config=runtime_config, - fps=fps, - client_liveness_timeout_s=client_liveness_timeout_s, - ) - self._pending_session_input: LingbotSessionInput | None = None - - def _model_name(self) -> str: - return self.runtime_config.config_name - - def _chunk_done_extra(self) -> dict[str, object]: - return {"active_event_id": getattr(self._runtime, "_active_event_id", None)} - - def _peek_pending_session_input(self) -> LingbotSessionInput | None: - return self._pending_session_input - - def _clear_pending_session_input(self) -> None: - self._pending_session_input = None - - async def _reset_runtime_for_session( - self, session_input: LingbotSessionInput | None - ) -> None: - await self._runtime.reset_for_new_session(session_input=session_input) + self._manager = manager + self._runtime = manager.runtime + self._runtime_config = manager.runtime_config def _effective_text_events(self) -> tuple[TextEventSpec, ...]: + pending_session_input = self._manager.pending_session_input if ( - self._pending_session_input is not None - and self._pending_session_input.text_events is not None + pending_session_input is not None + and pending_session_input.text_events is not None ): - return self._pending_session_input.text_events - return self.runtime_config.text_events + return pending_session_input.text_events + return self._runtime_config.text_events def get_initial_scene(self) -> dict[str, object]: - pending_input = self._pending_session_input + pending_input = self._manager.pending_session_input text_events = self._effective_text_events() prompt = ( normalize_prompt_text(pending_input.prompt) @@ -1527,11 +1241,11 @@ def get_initial_scene(self) -> dict[str, object]: if pending_input is not None and pending_input.first_frame_image_url: image_url = pending_input.first_frame_image_url else: - image_url = self.runtime_config.default_image_url + image_url = self._runtime_config.default_image_url input_source = "uploaded" if pending_input is not None else "default" first_frame_path = ( - self.runtime_config.example_data_dir - / self.runtime_config.first_frame_filename + self._runtime_config.example_data_dir + / self._runtime_config.first_frame_filename ) has_first_frame = ( bool( @@ -1542,27 +1256,27 @@ def get_initial_scene(self) -> dict[str, object]: ) ) or first_frame_path.exists() - or bool(self.runtime_config.default_image_url) + or bool(self._runtime_config.default_image_url) ) return { "first_frame_url": "/api/session/first_frame", "image_url": image_url, - "default_image_url": self.runtime_config.default_image_url, + "default_image_url": self._runtime_config.default_image_url, "has_first_frame": has_first_frame, "prompt": prompt, "input_source": input_source, - "model": self.runtime_config.config_name, + "model": self._runtime_config.config_name, "capabilities": {"text_events": bool(text_events)}, "event_catalog": [event.as_public_dict() for event in text_events], "active_event_id": getattr(self._runtime, "_active_event_id", None), "resolution": { - "width": self.runtime_config.video_width, - "height": self.runtime_config.video_height, + "width": self._runtime_config.video_width, + "height": self._runtime_config.video_height, }, } def get_first_frame(self) -> LingbotImagePayload: - pending_input = self._pending_session_input + pending_input = self._manager.pending_session_input if pending_input is not None and pending_input.first_frame_image_bytes: return LingbotImagePayload( data=pending_input.first_frame_image_bytes, @@ -1582,8 +1296,8 @@ def get_first_frame(self) -> LingbotImagePayload: return LingbotImagePayload(data=image_bytes, content_type=content_type) first_frame_path = ( - self.runtime_config.example_data_dir - / self.runtime_config.first_frame_filename + self._runtime_config.example_data_dir + / self._runtime_config.first_frame_filename ) if first_frame_path.exists(): return LingbotImagePayload( @@ -1598,11 +1312,11 @@ def get_first_frame(self) -> LingbotImagePayload: return LingbotImagePayload(data=encoded.tobytes(), content_type="image/jpeg") def set_pending_session_input(self, session_input: LingbotSessionInput) -> None: - if self.has_active_session(): + if self._manager.has_active_session(): raise SessionBusyError( "Cannot update Lingbot input while a session is active." ) - current = self._pending_session_input + current = self._manager.pending_session_input first_frame_image_bytes = ( current.first_frame_image_bytes if current is not None else None @@ -1650,15 +1364,17 @@ def set_pending_session_input(self, session_input: LingbotSessionInput) -> None: if session_input.text_events is not None else (current.text_events if current is not None else None) ) - self._pending_session_input = LingbotSessionInput( - prompt=( - normalize_prompt_text(session_input.prompt) - if session_input.prompt is not None - else (current.prompt if current is not None else None) - ), - first_frame_image_bytes=first_frame_image_bytes, - first_frame_image_url=first_frame_image_url, - first_frame_content_type=first_frame_content_type, - first_frame_remote_payload=first_frame_remote_payload, - text_events=text_events, + self._manager.set_pending_session_input( + LingbotSessionInput( + prompt=( + normalize_prompt_text(session_input.prompt) + if session_input.prompt is not None + else (current.prompt if current is not None else None) + ), + first_frame_image_bytes=first_frame_image_bytes, + first_frame_image_url=first_frame_image_url, + first_frame_content_type=first_frame_content_type, + first_frame_remote_payload=first_frame_remote_payload, + text_events=text_events, + ) ) diff --git a/integrations/lingbot/lingbot/webrtc/web/adapter.js b/integrations/lingbot/lingbot/webrtc/web/adapter.js index d7bdb6ad2..b323c32be 100644 --- a/integrations/lingbot/lingbot/webrtc/web/adapter.js +++ b/integrations/lingbot/lingbot/webrtc/web/adapter.js @@ -4,6 +4,15 @@ const mockMode = new URLSearchParams(window.location.search).has("mock") const controls = [ + { + label: "Drive / Turn", + keys: [ + { key: "w", label: "Forward" }, + { key: "a", label: "Turn left" }, + { key: "s", label: "Backward" }, + { key: "d", label: "Turn right" }, + ], + }, { label: "Strafe", keys: [ diff --git a/integrations/lingbot/pyproject.toml b/integrations/lingbot/pyproject.toml index 211069042..f0175fb35 100644 --- a/integrations/lingbot/pyproject.toml +++ b/integrations/lingbot/pyproject.toml @@ -42,7 +42,7 @@ dev = [ ] [project.scripts] -lingbot-demo = "lingbot.demo.cli:main" +lingbot-demo = "lingbot.demo.app:main" # Each entry registers one ``runner_name`` slug with ``flashdreams-run``. # The discovery layer (``flashdreams.plugins.registry.discover_runners``) diff --git a/integrations/lingbot/tests/test_demo_api.py b/integrations/lingbot/tests/test_demo_api.py index b2c560545..2cd5edfac 100644 --- a/integrations/lingbot/tests/test_demo_api.py +++ b/integrations/lingbot/tests/test_demo_api.py @@ -5,7 +5,7 @@ from collections.abc import Sequence from pathlib import Path -from typing import Any, cast +from typing import Any import numpy as np import pytest @@ -18,12 +18,12 @@ LingbotReplayInputs, LingbotWebRTCScenario, ) -from lingbot.demo.cli import _replay_spec, _webrtc_spec, parse_args +from lingbot.demo.app import _replay_spec, _webrtc_spec, parse_args from lingbot.demo.replay import ( LingbotReplayRuntime, LingbotReplayRuntimeOptions, ) -from lingbot.demo.webrtc import LingbotDemoWebRTCSessionManager +from lingbot.demo.webrtc import serve_lingbot_webrtc_demo from lingbot.input_mapping import ( FIELD_CAMERA_INTRINSICS, FIELD_CAMERA_TRAJECTORY, @@ -39,23 +39,22 @@ ) from lingbot.webrtc.session import LingbotRuntimeConfig -from flashdreams.infra.video_output import VideoStepResult from flashdreams.runtime import ( CanonicalInputs, InferenceConfig, InferenceInput, OutputArtifact, OutputTarget, + StepRequest, StepResult, ) from flashdreams.runtime.demo import ( DemoSpec, Mp4OutputSpec, WebRTCOutputSpec, - serve_flashdreams_demo, ) from flashdreams.runtime.demo.replay import run_replay_demo -from flashdreams.runtime.demo.webrtc import WebRTCDemo, build_webrtc_demo +from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager from flashdreams.serving.webrtc.server import SESSION_MANAGER_KEY pytestmark = pytest.mark.ci_cpu @@ -68,9 +67,7 @@ def _write_camera_assets(poses: Path, intrinsics: Path, *, frames: int = 64) -> np.save(poses, trajectory) np.save( intrinsics, - np.tile( - np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (frames, 1) - ), + np.tile(np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (frames, 1)), ) @@ -80,12 +77,12 @@ def test_lingbot_demo_defaults_to_interactive_preset() -> None: assert args.preset_id == "lingbot-world-fast-taehv-window15-sink3" -def test_lingbot_demo_adapter_declares_mp4_and_webrtc_modes() -> None: +def test_lingbot_demo_adapter_declares_replay_modes_only() -> None: adapter = LingbotDemoAdapter() assert adapter.model_id == LINGBOT_MODEL_ID - assert adapter.supported_input_modes() == ("replay", "keyboard-driving") - assert adapter.supported_output_modes() == ("mp4", "webrtc") + assert adapter.supported_input_modes() == ("replay",) + assert adapter.supported_output_modes() == ("mp4",) fields = { field.name for field in adapter.inference_input_schema.global_conditioning_fields @@ -100,9 +97,7 @@ def test_lingbot_demo_adapter_declares_mp4_and_webrtc_modes() -> None: FIELD_FPS, }.issubset(fields) # Camera control is per-step model input, not session-global scenario data. - step_fields = { - field.name for field in adapter.inference_input_schema.step_fields - } + step_fields = {field.name for field in adapter.inference_input_schema.step_fields} assert step_fields == {FIELD_CAMERA_TRAJECTORY, FIELD_CAMERA_INTRINSICS} @@ -208,9 +203,7 @@ def test_lingbot_replay_cli_defaults_to_example_data( example_dir = tmp_path / "example" example_dir.mkdir() (example_dir / "image.jpg").write_bytes(b"fake") - _write_camera_assets( - example_dir / "poses.npy", example_dir / "intrinsics.npy" - ) + _write_camera_assets(example_dir / "poses.npy", example_dir / "intrinsics.npy") (example_dir / "prompt.txt").write_text("drive through a forest\n") downloaded: list[int] = [] @@ -308,9 +301,9 @@ def test_lingbot_replay_runtime_generates_video_step_result( assert result.step_index == 0 assert result.frame_count == 1 - assert isinstance(result.output, VideoStepResult) - assert result.output.layout == "tchw" - assert result.output.video_chunk.shape == (1, 3, 2, 2) + assert isinstance(result, StepResult) + assert result.layout == "tchw" + assert result.video_chunk.shape == (1, 3, 2, 2) assert result.output_window is not None assert result.output_window.start_s == 0.0 assert result.output_window.end_s == 1 / 16 @@ -388,7 +381,6 @@ def test_lingbot_webrtc_cli_builds_keyboard_driving_spec() -> None: def test_lingbot_webrtc_demo_uses_existing_manager_with_model_config() -> None: pipeline_config = object() - adapter = LingbotDemoAdapter(webrtc_runtime_factory=_FakeWebRTCRuntime) spec = DemoSpec( model_id=LINGBOT_MODEL_ID, preset_id=DEFAULT_LINGBOT_PRESET, @@ -411,31 +403,37 @@ def test_lingbot_webrtc_demo_uses_existing_manager_with_model_config() -> None: ), ) - demo = build_webrtc_demo(spec=spec, adapter=adapter) - - assert isinstance(demo.runtime, _FakeWebRTCRuntime) - assert isinstance(demo.session_manager, LingbotDemoWebRTCSessionManager) - assert demo.session_manager._runtime is demo.runtime - assert demo.session_manager.runtime_config is demo.runtime.config - assert demo.runtime_config is demo.runtime.config - assert demo.runtime_config.pipeline_config is pipeline_config - assert demo.runtime_config.config_name == DEFAULT_LINGBOT_PRESET - assert demo.runtime_config.seed == 123 - assert demo.runtime_config.device == "cuda:7" - assert demo.runtime_config.video_width == 64 - assert demo.runtime_config.video_height == 32 - assert demo.runtime_config.fps == 24 - assert demo.runtime_config.encoder_backend == "default" - assert demo.runtime_config.example_data_dir.name == "02" - assert demo.session_manager._model_name() == DEFAULT_LINGBOT_PRESET - assert demo.host == "0.0.0.0" - assert demo.port == 8080 + calls: list[dict[str, Any]] = [] + serve_lingbot_webrtc_demo( + spec=spec, + world_rank=1, + runtime_factory=_FakeWebRTCRuntime, + server_runner=lambda **kwargs: calls.append(kwargs), + ) + + manager = calls[0]["session_manager"] + runtime = manager._runtime + assert isinstance(runtime, _FakeWebRTCRuntime) + assert type(manager) is BaseWebRTCSessionManager + assert manager.runtime_config is runtime.config + assert runtime.config.pipeline_config is pipeline_config + assert runtime.config.config_name == DEFAULT_LINGBOT_PRESET + assert runtime.config.seed == 123 + assert runtime.config.device == "cuda:7" + assert runtime.config.video_width == 64 + assert runtime.config.video_height == 32 + assert runtime.config.fps == 24 + assert runtime.config.encoder_backend == "default" + assert runtime.config.example_data_dir.name == "02" + assert manager.identity == DEFAULT_LINGBOT_PRESET + assert calls[0]["host"] == "0.0.0.0" + assert calls[0]["port"] == 8080 def test_lingbot_webrtc_demo_uses_shared_viewer_shell( monkeypatch: pytest.MonkeyPatch, ) -> None: - import lingbot.demo.webrtc as demo_webrtc_module + import flashdreams.runtime.demo.webrtc as shared_webrtc_module app_calls: list[dict[str, Any]] = [] @@ -447,9 +445,8 @@ def fake_create_packaged_app(**kwargs: Any) -> web.Application: return app monkeypatch.setattr( - demo_webrtc_module, "create_packaged_webrtc_app", fake_create_packaged_app + shared_webrtc_module, "create_packaged_webrtc_app", fake_create_packaged_app ) - adapter = LingbotDemoAdapter(webrtc_runtime_factory=_FakeWebRTCRuntime) spec = DemoSpec( model_id=LINGBOT_MODEL_ID, preset_id=DEFAULT_LINGBOT_PRESET, @@ -468,18 +465,23 @@ def fake_create_packaged_app(**kwargs: Any) -> web.Application: ), ) - demo = build_webrtc_demo(spec=spec, adapter=adapter, create_app=True) + app = serve_lingbot_webrtc_demo( + spec=spec, + runtime_factory=_FakeWebRTCRuntime, + create_app_fn=lambda **kwargs: fake_create_packaged_app(**kwargs), + server_runner=lambda **kwargs: None, + ) - assert demo.app is not None - assert app_calls[0]["session_manager"] is demo.session_manager + assert isinstance(app, web.Application) + assert app_calls[0]["session_manager"] is app[SESSION_MANAGER_KEY] assert app_calls[0]["request_session_url"] == ( "http://127.0.0.1:8080/request_session" ) - assert app_calls[0]["preload_name"] == "Lingbot" + assert app_calls[0]["preload_name"] == "Test Lingbot" assert str(app_calls[0]["web_resource"]).endswith("serving/webrtc/web") assert str(app_calls[0]["model_web_resource"]).endswith("lingbot/webrtc/web") assert callable(app_calls[0]["configure_app"]) - route_paths = {resource.canonical for resource in demo.app.router.resources()} + route_paths = {resource.canonical for resource in app.router.resources()} assert "/api/session/initial_scene" in route_paths assert "/api/session/first_frame" in route_paths assert "/api/session/input" in route_paths @@ -488,7 +490,7 @@ def fake_create_packaged_app(**kwargs: Any) -> web.Application: def test_lingbot_webrtc_demo_serves_through_shared_runner( monkeypatch: pytest.MonkeyPatch, ) -> None: - import lingbot.demo.webrtc as demo_webrtc_module + import flashdreams.runtime.demo.webrtc as shared_webrtc_module server_calls: list[dict[str, Any]] = [] @@ -502,9 +504,8 @@ def fake_server_runner(**kwargs: Any) -> None: server_calls.append(kwargs) monkeypatch.setattr( - demo_webrtc_module, "create_packaged_webrtc_app", fake_create_packaged_app + shared_webrtc_module, "create_packaged_webrtc_app", fake_create_packaged_app ) - adapter = LingbotDemoAdapter(webrtc_runtime_factory=_FakeWebRTCRuntime) spec = DemoSpec( model_id=LINGBOT_MODEL_ID, preset_id=DEFAULT_LINGBOT_PRESET, @@ -522,23 +523,19 @@ def fake_server_runner(**kwargs: Any) -> None: ), ) - demo = cast( - WebRTCDemo, - serve_flashdreams_demo( - spec=spec, - adapter=adapter, - world_rank=0, - server_runner=fake_server_runner, - ), + app = serve_lingbot_webrtc_demo( + spec=spec, + world_rank=0, + runtime_factory=_FakeWebRTCRuntime, + server_runner=fake_server_runner, ) assert len(server_calls) == 1 assert server_calls[0]["world_rank"] == 0 - assert server_calls[0]["session_manager"] is demo.session_manager - assert server_calls[0]["app"] is demo.app + assert server_calls[0]["app"] is app assert server_calls[0]["host"] == "0.0.0.0" assert server_calls[0]["port"] == 8080 - assert isinstance(demo.session_manager, LingbotDemoWebRTCSessionManager) + assert type(server_calls[0]["session_manager"]) is BaseWebRTCSessionManager class _RecordingOutputTarget: @@ -603,19 +600,23 @@ async def initialize(self) -> None: async def reset_for_new_session(self, *args: Any, **kwargs: Any) -> None: return None - def peek_steady_chunk_num_frames(self) -> int: - return 1 + def peek_input_fps(self) -> float: + return 16.0 - def peek_next_chunk_num_frames(self) -> int: + def peek_steady_output_num_frames(self) -> int: return 1 - async def generate_chunk( + def next_step_request(self) -> StepRequest: + return StepRequest(step_index=0, metadata={"input_frame_count": 1}) + + async def step( self, *, + request: StepRequest, segments: list[Any], frame_times: list[float], ) -> Any: - del segments, frame_times + del request, segments, frame_times return None async def close(self) -> None: diff --git a/integrations/lingbot/tests/test_distributed_server_main.py b/integrations/lingbot/tests/test_distributed_server_main.py index 7f11a776b..dd57f1fdd 100644 --- a/integrations/lingbot/tests/test_distributed_server_main.py +++ b/integrations/lingbot/tests/test_distributed_server_main.py @@ -162,7 +162,11 @@ def _make_manager(runtime_config, fps): manager_fps.append(fps) return fake_manager - monkeypatch.setattr(server, "LingbotWebRTCSessionManager", _make_manager) + monkeypatch.setattr( + server, + "create_lingbot_webrtc_session_manager", + _make_manager, + ) monkeypatch.setattr(server, "get_external_ip", lambda: "203.0.113.10") def _create_app(*, session_manager, request_session_url=None): @@ -213,7 +217,11 @@ def _make_manager(runtime_config, fps): manager_fps.append(fps) return fake_manager - monkeypatch.setattr(server, "LingbotWebRTCSessionManager", _make_manager) + monkeypatch.setattr( + server, + "create_lingbot_webrtc_session_manager", + _make_manager, + ) server.main() diff --git a/integrations/lingbot/tests/test_input_mapping.py b/integrations/lingbot/tests/test_input_mapping.py index 1db76bf44..87ffb0b63 100644 --- a/integrations/lingbot/tests/test_input_mapping.py +++ b/integrations/lingbot/tests/test_input_mapping.py @@ -36,9 +36,7 @@ _KEYBOARD_SOURCE = UserInputSchema( capabilities=( - UserInputCapability( - event_type="key_down", payload_fields=frozenset({"key"}) - ), + UserInputCapability(event_type="key_down", payload_fields=frozenset({"key"})), UserInputCapability(event_type="key_up", payload_fields=frozenset({"key"})), UserInputCapability( event_type="text_event", payload_fields=frozenset({"event_id"}) @@ -71,7 +69,9 @@ def test_keyboard_events_become_camera_command_axes() -> None: converter = KeyboardToCameraCommand() inputs = UserInputs( events=( - UserInputEvent(timestamp_s=0.0, event_type="key_down", payload={"key": "w"}), + UserInputEvent( + timestamp_s=0.0, event_type="key_down", payload={"key": "w"} + ), ) ) window = TimeWindow(start_s=0.0, end_s=1.0) @@ -93,7 +93,9 @@ def test_camera_command_segments_preserve_sub_window_timing() -> None: window = TimeWindow(start_s=0.0, end_s=1.0) inputs = UserInputs( events=( - UserInputEvent(timestamp_s=0.5, event_type="key_down", payload={"key": "w"}), + UserInputEvent( + timestamp_s=0.5, event_type="key_down", payload={"key": "w"} + ), ) ) @@ -111,7 +113,9 @@ def test_key_events_drive_a_camera_trajectory() -> None: mapping = _live_mapping() user_inputs = UserInputs( events=( - UserInputEvent(timestamp_s=0.0, event_type="key_down", payload={"key": "w"}), + UserInputEvent( + timestamp_s=0.0, event_type="key_down", payload={"key": "w"} + ), ) ) request = _step_request(step_index=0, frame_start=0, num_frames=4) @@ -422,7 +426,7 @@ def test_event_driven_scenario_builds_a_live_mapping(tmp_path: Path) -> None: prepared = LingbotDemoAdapter().prepare_scenario(spec) - assert prepared.mapping is not None + assert isinstance(prepared.mapping, LingbotInputMapping) assert prepared.mapping.mapping_schema.consumes == (CAMERA_COMMAND, TEXT_EVENT) assert len(prepared.user_inputs.events) == 2 # The declared source must actually cover the trace it carries, or the diff --git a/integrations/lingbot/tests/test_runtime_gpu.py b/integrations/lingbot/tests/test_runtime_gpu.py index db88ccd3e..b5124c091 100644 --- a/integrations/lingbot/tests/test_runtime_gpu.py +++ b/integrations/lingbot/tests/test_runtime_gpu.py @@ -19,8 +19,12 @@ inference_input_from_replay_inputs, ) -from flashdreams.infra.video_output import VideoStepResult -from flashdreams.runtime import CanonicalInputs, InferenceConfig, InferenceInput +from flashdreams.runtime import ( + CanonicalInputs, + InferenceConfig, + InferenceInput, + StepResult, +) pytestmark = pytest.mark.ci_gpu @@ -104,9 +108,9 @@ def _fake_load_first_frame_tensor( runtime.close() assert result.frame_count == 1 - assert isinstance(result.output, VideoStepResult) - assert result.output.video_chunk.is_cuda - assert result.output.video_chunk.shape == (1, 3, 2, 2) + assert isinstance(result, StepResult) + assert result.video_chunk.is_cuda + assert result.video_chunk.shape == (1, 3, 2, 2) assert pipeline.initialize_cache_devices == ["cuda"] assert pipeline.generate_world_scales == [mapping.camera_trace.world_scale] @@ -217,9 +221,7 @@ def _fake_load_first_frame_tensor(path: Path, **kwargs: Any) -> torch.Tensor: UserInputCapability( event_type="key_down", payload_fields=frozenset({"key"}) ), - UserInputCapability( - event_type="key_up", payload_fields=frozenset({"key"}) - ), + UserInputCapability(event_type="key_up", payload_fields=frozenset({"key"})), ) ) user_inputs = UserInputs( @@ -255,5 +257,6 @@ def _fake_load_first_frame_tensor(path: Path, **kwargs: Any) -> torch.Tensor: session.close() runtime.close() - assert result.output.video_chunk.is_cuda + assert isinstance(result, StepResult) + assert result.video_chunk.is_cuda assert pipeline.generate_world_scales == [1.0] diff --git a/integrations/lingbot/tests/test_runtime_session_inputs.py b/integrations/lingbot/tests/test_runtime_session_inputs.py index def9cfb9e..27a9c4b93 100644 --- a/integrations/lingbot/tests/test_runtime_session_inputs.py +++ b/integrations/lingbot/tests/test_runtime_session_inputs.py @@ -8,9 +8,9 @@ from pathlib import Path from typing import Any +import lingbot.runtime as runtime_module import pytest import torch -import lingbot.runtime as runtime_module from lingbot.input_mapping import FIELD_CAMERA_INTRINSICS, FIELD_CAMERA_TRAJECTORY from lingbot.runtime import ( LINGBOT_MODEL_ID, @@ -267,7 +267,9 @@ def test_text_event_prompt_update_swaps_the_rollout_context( ) assert pipeline.text_encoder_calls == [["a violent storm"]] - assert len(pipeline.diffusion_model.transformer.replaced) == 1 + transformer = pipeline.diffusion_model.transformer + assert isinstance(transformer, _FakeTransformer) + assert len(transformer.replaced) == 1 # Re-sending the same prompt must not re-encode or re-swap. session.step( diff --git a/integrations/lingbot/tests/test_server_routes.py b/integrations/lingbot/tests/test_server_routes.py index 6c4acd49f..167f147a1 100644 --- a/integrations/lingbot/tests/test_server_routes.py +++ b/integrations/lingbot/tests/test_server_routes.py @@ -199,7 +199,7 @@ async def test_lingbot_model_adapter_is_served() -> None: assert response.status == 200 assert 'modelName: "Lingbot"' in body assert "/api/session/initial_scene" in body - assert '{ key: "w"' not in body + assert '{ key: "w"' in body assert '{ key: "q"' in body assert "enablePostprocess" not in body assert "RTCPeerConnection" not in body diff --git a/integrations/lingbot/tests/test_smoke.py b/integrations/lingbot/tests/test_smoke.py index 3cd0a93a6..a787431a3 100644 --- a/integrations/lingbot/tests/test_smoke.py +++ b/integrations/lingbot/tests/test_smoke.py @@ -57,6 +57,7 @@ from flashdreams.infra.config import derive_config from flashdreams.infra.runner import RunnerConfig +from flashdreams.runtime import InferenceConfig, InferenceInput pytestmark = pytest.mark.ci_cpu @@ -68,11 +69,10 @@ def _write_camera_assets(poses: Path, intrinsics: Path, *, frames: int = 64) -> np.save(poses, trajectory) np.save( intrinsics, - np.tile( - np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (frames, 1) - ), + np.tile(np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (frames, 1)), ) + ENTRY_POINT_GROUP = "flashdreams.runner_configs" @@ -236,10 +236,13 @@ def _fake_run_inference_session(**kwargs: object) -> tuple[object, ...]: assert isinstance(captured["adapter"], LingbotModelAdapter) config = captured["config"] - assert getattr(config, "model_id") == LINGBOT_MODEL_ID - assert getattr(config, "device") == "cpu" + assert isinstance(config, InferenceConfig) + assert config.model_id == LINGBOT_MODEL_ID + assert config.device == "cpu" assert config.runtime_options["pipeline"] is pipeline - inputs = captured["initial_inputs"].global_conditioning + initial_inputs = captured["initial_inputs"] + assert isinstance(initial_inputs, InferenceInput) + inputs = initial_inputs.global_conditioning assert inputs[FIELD_PROMPT] == "drive through a city" assert inputs[FIELD_FIRST_FRAME_PATH] == image assert inputs[FIELD_TOTAL_BLOCKS] == 1 diff --git a/integrations/lingbot/tests/test_webrtc_runtime.py b/integrations/lingbot/tests/test_webrtc_runtime.py index e38061a38..2c2f2ac7f 100644 --- a/integrations/lingbot/tests/test_webrtc_runtime.py +++ b/integrations/lingbot/tests/test_webrtc_runtime.py @@ -22,25 +22,42 @@ import pytest import torch -from lingbot.input_mapping import ( - KeyboardToCameraCommand, - LingbotInputMapping, - TextEventSelection, -) +from lingbot.model_session import LingbotModelSessionCore from lingbot.webrtc import session from lingbot.webrtc.session import ( LingbotRuntimeConfig, - LingbotWebRTCSessionManager, + create_lingbot_webrtc_session_manager, ) -from flashdreams.infra.video_output import VideoOutputStream, VideoStepResult -from flashdreams.runtime.canonical import InputCanonicalizer -from flashdreams.runtime.inputs import InferenceInput -from flashdreams.runtime.types import StepRequest, StepResult +from flashdreams.infra.video_output import VideoOutputStream +from flashdreams.runtime import StepRequest, StepResult +from flashdreams.serving.webrtc import runtime as webrtc_runtime +from flashdreams.serving.webrtc.manager import ( + BaseWebRTCSessionManager, + ManagedWebRTCSession, +) pytestmark = pytest.mark.ci_cpu +def _attach_model_session( + runtime: session.LingbotInferenceRuntime, + pipeline: object, + *, + cache: object | None = None, +) -> LingbotModelSessionCore: + core = LingbotModelSessionCore( + pipeline=pipeline, + output_stream_factory=lambda: VideoOutputStream( + postprocess_stream=None, + output_layout="tchw", + ), + ) + core._cache = cache # Test seam for already-initialized runtime state. + runtime._model_session = core + return core + + class _FakeCloseable: def __init__(self) -> None: self.closed = False @@ -50,7 +67,7 @@ async def close(self) -> None: class _FakeVideoEncoder: - """Minimal ``VideoEncoder``-shaped stub for ``_ManagedLingbotSession`` + """Minimal ``VideoEncoder``-shaped stub for ``ManagedWebRTCSession`` construction. Enough to satisfy the dataclass field; the tests here do not exercise ``create_track`` / ``deliver_chunk`` on it.""" @@ -68,19 +85,13 @@ def _fake_runtime_factory(config: LingbotRuntimeConfig) -> object: def test_session_manager_hooks_are_wired() -> None: - # Guards against the shared base-class attribute overrides being dropped - # (e.g. losing their leading underscore), which silently reverts behaviour - # to the base defaults. - assert ( - LingbotWebRTCSessionManager._busy_message - == "A Lingbot session is already active." + manager = create_lingbot_webrtc_session_manager( + runtime_config=LingbotRuntimeConfig(device="cpu") ) - assert LingbotWebRTCSessionManager._warmup_label == "Lingbot WebRTC" - assert LingbotWebRTCSessionManager._runtime_error_types == ( - session.LingbotRuntimeError, - ) - # Lingbot keeps streaming after a per-chunk failure rather than tearing down. - assert LingbotWebRTCSessionManager._close_session_on_generation_error is False + + assert manager.busy_message == "A Lingbot session is already active." + assert manager.warmup_label == "Lingbot WebRTC" + assert manager.fatal_generation_errors is False def test_runtime_defaults_use_canonical_v2_examples() -> None: @@ -102,7 +113,7 @@ def test_session_manager_uses_runtime_config_fps_by_default( ) -> None: monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) - manager = LingbotWebRTCSessionManager( + manager = create_lingbot_webrtc_session_manager( runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0, fps=12) ) @@ -115,7 +126,9 @@ def test_initialize_video_encoder_sync_skips_on_non_master( def _select_encoder_should_not_be_called(**_kw: object) -> object: raise AssertionError("worker ranks must not initialize WebRTC encoders") - monkeypatch.setattr(session, "select_encoder", _select_encoder_should_not_be_called) + monkeypatch.setattr( + webrtc_runtime, "select_encoder", _select_encoder_should_not_be_called + ) runtime = session.LingbotInferenceRuntime( config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) ) @@ -137,7 +150,7 @@ def _fake_select_encoder(**kwargs: object) -> _FakeVideoEncoder: calls.append(kwargs) return stub - monkeypatch.setattr(session, "select_encoder", _fake_select_encoder) + monkeypatch.setattr(webrtc_runtime, "select_encoder", _fake_select_encoder) runtime = session.LingbotInferenceRuntime( config=LingbotRuntimeConfig( device="cuda:2", @@ -201,17 +214,16 @@ def finalize(autoregressive_index: int, cache: object) -> dict[str, float]: captured: dict[str, object] = {} - def _fake_make_step_result( + def _fake_process( _stream: VideoOutputStream, video_chunk: object, **kwargs: object - ) -> VideoStepResult: + ) -> StepResult: captured["video_chunk"] = video_chunk captured.update(kwargs) - return VideoStepResult( - chunk_index=0, - num_frames=2, + return StepResult.from_video_chunk( + step_index=0, video_chunk=torch.zeros((2, 3, 4, 5)), - stats={"total_ms": 3.0}, layout="tchw", + metrics={"total_ms": 3.0}, ) runtime = session.LingbotInferenceRuntime( @@ -220,12 +232,12 @@ def _fake_make_step_result( pipeline = _FakePipeline() runtime._device = torch.device("cpu") runtime._pipeline = pipeline - runtime._cache = object() + _attach_model_session(runtime, pipeline, cache=object()) runtime._base_intrinsics = torch.ones(4) monkeypatch.setattr( VideoOutputStream, - "make_step_result", - _fake_make_step_result, + "process", + _fake_process, ) result = runtime._generate_one_chunk_sync( @@ -234,10 +246,13 @@ def _fake_make_step_result( ) assert captured["video_chunk"] is pipeline.output - assert captured["sync_device"] == torch.device("cpu") + metrics = cast(dict[str, float], captured["metrics"]) + assert metrics["total_ms"] == 3.0 + assert float(metrics["model_step_s"]) >= 0.0 assert pipeline.output.detach_calls == 0 - assert result.stats == {"total_ms": 3.0} - assert runtime.autoregressive_index == 1 + assert result.metrics == {"total_ms": 3.0} + assert runtime._model_session is not None + assert runtime._model_session.step_index == 1 def test_validate_remote_url_normalizes_github_blob_image_url( @@ -390,11 +405,12 @@ def _load_default_prompt(self) -> str: return "drive through a city" monkeypatch.setattr(session, "LingbotInferenceRuntime", _FakeRuntime) - manager = LingbotWebRTCSessionManager( + manager = create_lingbot_webrtc_session_manager( runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) ) + controller = session.LingbotWebRTCSessionController(manager) - scene = manager.get_initial_scene() + scene = controller.get_initial_scene() assert scene["capabilities"] == {"text_events": True} assert scene["active_event_id"] is None @@ -437,9 +453,10 @@ def _load_default_prompt(self) -> str: return "drive through a city" monkeypatch.setattr(session, "LingbotInferenceRuntime", _FakeRuntime) - manager = LingbotWebRTCSessionManager( + manager = create_lingbot_webrtc_session_manager( runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) ) + controller = session.LingbotWebRTCSessionController(manager) custom_events = ( session.TextEventSpec( event_id="rain", @@ -449,10 +466,10 @@ def _load_default_prompt(self) -> str: ), ) - manager.set_pending_session_input( + controller.set_pending_session_input( session.LingbotSessionInput(text_events=custom_events) ) - scene = manager.get_initial_scene() + scene = controller.get_initial_scene() assert scene["capabilities"] == {"text_events": True} assert scene["event_catalog"] == [custom_events[0].as_public_dict()] @@ -497,16 +514,17 @@ def _fake_read_remote_bytes( lambda hostname: (ipaddress.ip_address("93.184.216.34"),), ) monkeypatch.setattr(session, "_read_remote_bytes", _fake_read_remote_bytes) - manager = LingbotWebRTCSessionManager( + manager = create_lingbot_webrtc_session_manager( runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) ) + controller = session.LingbotWebRTCSessionController(manager) - manager.set_pending_session_input( + controller.set_pending_session_input( session.LingbotSessionInput( first_frame_image_url="https://example.test/scene.png" ) ) - payload = manager.get_first_frame() + payload = controller.get_first_frame() assert fake_runtime is not None assert fake_runtime.decoded_images == [b"remote-image"] @@ -515,8 +533,8 @@ def _fake_read_remote_bytes( data=b"remote-image", content_type="image/png", ) - assert manager._pending_session_input is not None - assert manager._pending_session_input.first_frame_remote_payload == payload + assert manager.pending_session_input is not None + assert manager.pending_session_input.first_frame_remote_payload == payload def test_prepare_session_input_state_uses_cached_remote_payload( @@ -577,36 +595,117 @@ def replace_text_embeddings( ) -> None: self.calls.append((cache, text_embeddings)) - class _FakeDiffusionModel: + monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) + manager = create_lingbot_webrtc_session_manager( + runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) + ) + runtime = _FakeRuntime() + channel = _FakeControlChannel() + managed_session = ManagedWebRTCSession( + runtime=runtime, + 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=channel, + ) + + await manager._handle_datachannel_message( + managed_session=managed_session, + raw_message='{"type":"event","event_id":"portal","state":"trigger"}', + ) + + assert runtime.calls == [("portal", "trigger")] + assert channel.messages == [ + { + "type": "event_ack", + "event_id": "portal", + "state": "trigger", + "active_event_id": "portal", + } + ] + assert managed_session.first_action_received.is_set() + + +@pytest.mark.asyncio +async def test_clear_event_message_does_not_require_event_id_and_preserves_ack_fields( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _FakeRuntime: def __init__(self) -> None: self.transformer = _FakeTransformer() - class _FakePipeline: + async def trigger_event( + self, *, event_id: str, state: str + ) -> dict[str, object]: + self.calls.append((event_id, state)) + return { + "type": "not_event_ack", + "event_id": "overwritten", + "state": "overwritten", + "active_event_id": None, + } + + monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) + manager = create_lingbot_webrtc_session_manager( + runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) + ) + runtime = _FakeRuntime() + channel = _FakeControlChannel() + managed_session = ManagedWebRTCSession( + runtime=runtime, + 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=channel, + ) + + await manager._handle_datachannel_message( + managed_session=managed_session, + raw_message='{"type":"event","state":"clear"}', + ) + + assert runtime.calls == [("", "clear")] + assert channel.messages == [ + { + "type": "event_ack", + "event_id": None, + "state": "clear", + "active_event_id": None, + } + ] + assert managed_session.first_action_received.is_set() + + +@pytest.mark.asyncio +async def test_event_message_without_id_is_rejected_for_trigger( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _FakeRuntime: def __init__(self) -> None: self.diffusion_model = _FakeDiffusionModel() - runtime = session.LingbotInferenceRuntime( - config=LingbotRuntimeConfig( - device="cpu", - warmup_chunks=0, - text_events=(), - ) - ) - transformer_cache = object() - cache = type("_FakeCache", (), {"transformer_cache": transformer_cache})() - base_text = torch.zeros((1, 2, 3)) - event_text = torch.ones((1, 2, 3)) - runtime._pipeline = _FakePipeline() - runtime._cache = cache - runtime._prompt = "base prompt" - runtime._event_embeddings = {"portal": event_text} - runtime._prompt_embeddings = { - "base prompt": base_text, - "a glowing portal opens": event_text, - } + async def trigger_event( + self, *, event_id: str, state: str + ) -> dict[str, object]: + del event_id, state + self.calls += 1 + return {} - runtime._apply_conditioning_update_sync( - InferenceInput(global_conditioning={"prompt": "a glowing portal opens"}) + monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) + manager = create_lingbot_webrtc_session_manager( + runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) + ) + runtime = _FakeRuntime() + channel = _FakeControlChannel() + managed_session = ManagedWebRTCSession( + runtime=runtime, + 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=channel, ) transformer = runtime._pipeline.diffusion_model.transformer @@ -653,7 +752,7 @@ def __init__(self) -> None: base_text = torch.zeros((1, 2, 3)) event_text = torch.ones((1, 2, 3)) runtime._pipeline = _FakePipeline() - runtime._cache = cache + _attach_model_session(runtime, runtime._pipeline, cache=cache) runtime._base_text_embeddings = base_text runtime._event_embeddings = {"portal": event_text} @@ -729,6 +828,7 @@ def initialize_cache(self, *, text: list[str], image: torch.Tensor) -> object: pipeline = _FakePipeline() runtime._device = torch.device("cpu") runtime._pipeline = pipeline + _attach_model_session(runtime, pipeline) def _fake_prepare_session_input_state( session_input: session.LingbotSessionInput | None, @@ -767,7 +867,7 @@ async def test_trigger_event_prevalidates_before_distributed_broadcast() -> None ) ) runtime._pipeline = object() - runtime._cache = object() + _attach_model_session(runtime, runtime._pipeline, cache=object()) runtime._event_embeddings = {"portal": torch.ones((1, 2, 3))} calls = 0 @@ -795,7 +895,7 @@ async def test_trigger_event_waits_for_generation_lock() -> None: ) ) runtime._pipeline = object() - runtime._cache = object() + _attach_model_session(runtime, runtime._pipeline, cache=object()) runtime._event_embeddings = {"portal": torch.ones((1, 2, 3))} calls: list[tuple[str, str]] = [] @@ -845,18 +945,18 @@ def _fake_runtime_factory(config: LingbotRuntimeConfig) -> _FakeRuntime: return fake_runtime async def _fake_loopback_warmup( - self: LingbotWebRTCSessionManager, *, num_chunks: int + self: BaseWebRTCSessionManager, *, num_chunks: int ) -> None: del self warmup_calls.append(num_chunks) monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) monkeypatch.setattr( - LingbotWebRTCSessionManager, + BaseWebRTCSessionManager, "_run_loopback_warmup_session", _fake_loopback_warmup, ) - manager = LingbotWebRTCSessionManager( + manager = create_lingbot_webrtc_session_manager( runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=2) ) @@ -929,14 +1029,32 @@ async def reset_for_new_session( del session_input self.reset_calls += 1 - def peek_steady_chunk_num_frames(self) -> int: - return 1 + def peek_input_fps(self) -> float: + return 30.0 - def peek_next_chunk_num_frames(self) -> int: + def peek_steady_output_num_frames(self) -> int: return 1 - async def start_inference_session(self) -> _FakeInferenceSession: - return _FakeInferenceSession(self) + def next_step_request(self) -> StepRequest: + return StepRequest( + step_index=len(self.generated_segments), + metadata={"input_frame_count": 1}, + ) + + async def step( + self, + *, + request: StepRequest, + segments: list[tuple[float, float, frozenset[str]]], + frame_times: list[float], + ) -> StepResult: + del frame_times + self.generated_segments.append(segments) + return StepResult.from_video_chunk( + step_index=request.step_index, + video_chunk=torch.zeros((1, 1, 1, 3, 2, 2), dtype=torch.uint8), + layout="bvtchw", + ) async def close(self) -> None: self.close_calls += 1 @@ -949,7 +1067,7 @@ def _fake_runtime_factory(config: LingbotRuntimeConfig) -> _FakeRuntime: return fake_runtime monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) - manager = LingbotWebRTCSessionManager( + manager = create_lingbot_webrtc_session_manager( runtime_config=LingbotRuntimeConfig( device="cpu", warmup_chunks=2, @@ -997,7 +1115,7 @@ def _fake_runtime_factory(config: LingbotRuntimeConfig) -> _FakeRuntime: return fake_runtime monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) - manager = LingbotWebRTCSessionManager( + manager = create_lingbot_webrtc_session_manager( runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) ) @@ -1014,7 +1132,7 @@ async def test_create_answer_passes_pending_session_input( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) - manager = LingbotWebRTCSessionManager( + manager = create_lingbot_webrtc_session_manager( runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) ) manager._runtime_ready = True @@ -1049,10 +1167,10 @@ async def test_heartbeat_message_refreshes_client_liveness( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) - manager = LingbotWebRTCSessionManager( + manager = create_lingbot_webrtc_session_manager( runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) ) - managed_session = session._ManagedLingbotSession( + managed_session = ManagedWebRTCSession( runtime=object(), video_track=_FakeCloseable(), # ty:ignore[invalid-argument-type] video_encoder=_FakeVideoEncoder(), # ty:ignore[invalid-argument-type] @@ -1077,13 +1195,13 @@ async def test_client_liveness_timeout_closes_active_session( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) - manager = LingbotWebRTCSessionManager( + manager = create_lingbot_webrtc_session_manager( runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0), client_liveness_timeout_s=0.01, ) video_track = _FakeCloseable() peer_connection = _FakeCloseable() - managed_session = session._ManagedLingbotSession( + managed_session = ManagedWebRTCSession( runtime=object(), video_track=video_track, # ty:ignore[invalid-argument-type] video_encoder=_FakeVideoEncoder(), # ty:ignore[invalid-argument-type] @@ -1110,12 +1228,12 @@ async def test_disconnect_message_closes_active_session( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) - manager = LingbotWebRTCSessionManager( + manager = create_lingbot_webrtc_session_manager( runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) ) video_track = _FakeCloseable() peer_connection = _FakeCloseable() - managed_session = session._ManagedLingbotSession( + managed_session = ManagedWebRTCSession( runtime=object(), video_track=video_track, # ty:ignore[invalid-argument-type] video_encoder=_FakeVideoEncoder(), # ty:ignore[invalid-argument-type] diff --git a/integrations/lingbot/tests/test_webrtc_runtime_distributed.py b/integrations/lingbot/tests/test_webrtc_runtime_distributed.py index 8dde9db82..691815bde 100644 --- a/integrations/lingbot/tests/test_webrtc_runtime_distributed.py +++ b/integrations/lingbot/tests/test_webrtc_runtime_distributed.py @@ -383,7 +383,7 @@ def test_runtime_distributed_ops_use_world_cp_and_rank_seed( if rank == 0: runtime._initialize_sync_all_ranks() runtime._reset_rollout_sync_all_ranks() - num_frames = runtime.peek_next_chunk_num_frames() + num_frames = runtime._next_input_frame_count() per_frame_keys = [frozenset() for _ in range(num_frames)] result = runtime._generate_chunk_sync_all_ranks(per_frame_keys) result_shape = tuple(result.video_chunk.shape) diff --git a/integrations/omnidreams/omnidreams/demo/README.md b/integrations/omnidreams/omnidreams/demo/README.md index d69c0170a..ce23b48cd 100644 --- a/integrations/omnidreams/omnidreams/demo/README.md +++ b/integrations/omnidreams/omnidreams/demo/README.md @@ -50,9 +50,9 @@ compile/cache behavior is reliable enough for the demo path. ## WebRTC -WebRTC uses the shared demo launcher around the existing Omnidreams live WebRTC -runtime. It is still scene-driven and uses Ludus to render HDMap conditioning -from a scene: +WebRTC uses the shared FlashDreams server, session manager, and runtime worker. +The small model adapter in this package loads one scene, renders HDMap +conditioning with Ludus, and runs OmniDreams from browser WASD controls: ```bash uv run --package flashdreams-omnidreams omnidreams-demo webrtc \ diff --git a/integrations/omnidreams/omnidreams/demo/adapter.py b/integrations/omnidreams/omnidreams/demo/adapter.py index e16c5c0a1..a1ca8e8f7 100644 --- a/integrations/omnidreams/omnidreams/demo/adapter.py +++ b/integrations/omnidreams/omnidreams/demo/adapter.py @@ -6,16 +6,10 @@ from __future__ import annotations from collections.abc import Callable -from dataclasses import replace from typing import Any from omnidreams.config import OMNIDREAMS_CONFIGS, OMNIDREAMS_RUNNERS -from omnidreams.webrtc.session import ( - OmnidreamsInferenceRuntime, - OmnidreamsRuntimeConfig, -) -from flashdreams.infra.postprocess import VideoPostprocessChainConfig from flashdreams.runtime import ( CanonicalInputSchema, IdentityInputMapping, @@ -30,7 +24,6 @@ DemoSpec, Mp4OutputSpec, PreparedScenario, - WebRTCOutputSpec, ) from flashdreams.runtime.interfaces import InferenceRuntime @@ -43,16 +36,9 @@ DEFAULT_OMNIDREAMS_PRESET, OMNIDREAMS_MODEL_ID, resolve_replay_scenario, - resolve_webrtc_scenario, -) -from .webrtc import ( - OmnidreamsDemoWebRTCSessionManager, - create_omnidreams_webrtc_app, - validate_postprocess_preset, ) ReplayRuntimeFactory = Callable[..., InferenceRuntime] -WebRTCRuntimeFactory = Callable[..., Any] class OmnidreamsDemoAdapter: @@ -62,11 +48,9 @@ def __init__( self, *, replay_runtime_factory: ReplayRuntimeFactory = OmnidreamsReplayRuntime, - webrtc_runtime_factory: WebRTCRuntimeFactory = OmnidreamsInferenceRuntime, pipeline_factory: PipelineFactory | None = None, ) -> None: self._replay_runtime_factory = replay_runtime_factory - self._webrtc_runtime_factory = webrtc_runtime_factory self._pipeline_factory = pipeline_factory self._mapping = IdentityInputMapping() @@ -94,10 +78,10 @@ def default_input_mapping(self) -> IdentityInputMapping: return self._mapping def supported_input_modes(self) -> tuple[str, ...]: - return ("replay", "keyboard-driving") + return ("replay",) def supported_output_modes(self) -> tuple[str, ...]: - return ("mp4", "webrtc") + return ("mp4",) def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: if spec.input_mode != "replay": @@ -143,87 +127,6 @@ def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: ), ) - def create_webrtc_runtime(self, spec: DemoSpec) -> Any: - runtime_config = self.create_webrtc_runtime_config(spec=spec, runtime=None) - return self._webrtc_runtime_factory(config=runtime_config) - - def create_webrtc_runtime_config( - self, - *, - spec: DemoSpec, - runtime: Any, - ) -> OmnidreamsRuntimeConfig: - runtime_config = getattr(runtime, "config", None) - if isinstance(runtime_config, OmnidreamsRuntimeConfig): - return runtime_config - if spec.input_mode != "keyboard-driving": - raise ValueError( - "OmniDreams WebRTC requires input_mode='keyboard-driving', " - f"got {spec.input_mode!r}." - ) - if not isinstance(spec.output, WebRTCOutputSpec): - raise ValueError("OmniDreams WebRTC requires WebRTC output.") - config = spec.config - if config is None: - raise RuntimeError("DemoSpec.config was not initialized.") - self.validate_config(config) - scenario = resolve_webrtc_scenario(spec.scenario) - validate_postprocess_preset(scenario.postprocess_preset) - - preset_id = self._preset_id(config) - pipeline_config = self._pipeline_config(config) - seed = _option(config, "seed", 42) - device = config.device or str(_option(config, "device", "cuda:0")) - runtime_config = OmnidreamsRuntimeConfig( - pipeline_config_name=preset_id, - pipeline_config=pipeline_config, - scene_dir=scenario.scene_dir, - scene_uuid=scenario.scene_uuid, - scene_variant=scenario.scene_variant, - seed=None if seed is None else int(seed), - device=device, - video_height=spec.output.video_height, - video_width=spec.output.video_width, - fps=spec.output.fps, - camera_name=scenario.camera_name, - warmup_chunks=spec.output.warmup_chunks, - warmup_timeout_s=spec.output.warmup_timeout_s, - debug_serve_hdmaps=scenario.debug_serve_hdmaps, - postprocess=VideoPostprocessChainConfig(preset=scenario.postprocess_preset), - encoder_backend="default" if scenario.prefer_sw_encoder else "auto", - ) - return _apply_webrtc_runtime_options(runtime_config, config.runtime_options) - - def create_webrtc_session_manager( - self, - *, - spec: DemoSpec, - runtime: Any, - runtime_config: OmnidreamsRuntimeConfig, - fps: int, - client_liveness_timeout_s: float, - ) -> OmnidreamsDemoWebRTCSessionManager: - del spec - return OmnidreamsDemoWebRTCSessionManager( - runtime=runtime, - runtime_config=runtime_config, - fps=fps, - client_liveness_timeout_s=client_liveness_timeout_s, - ) - - def create_webrtc_app( - self, - *, - spec: DemoSpec, - session_manager: Any, - request_session_url: str, - ) -> Any: - return create_omnidreams_webrtc_app( - spec=spec, - session_manager=session_manager, - request_session_url=request_session_url, - ) - def _preset_id(self, config: InferenceConfig | None) -> str: return ( DEFAULT_OMNIDREAMS_PRESET @@ -250,30 +153,7 @@ def _default_replay_prompt(self, config: InferenceConfig | None) -> str: return "" if runner is None else str(getattr(runner, "prompt", "")) -def _option(config: InferenceConfig, name: str, default: Any) -> Any: - return config.runtime_options.get(name, default) - - -def _apply_webrtc_runtime_options( - runtime_config: OmnidreamsRuntimeConfig, - options: Any, -) -> OmnidreamsRuntimeConfig: - if not isinstance(options, dict): - options = dict(options) - overrides: dict[str, Any] = {} - for name in ( - "move_speed_per_s", - "rotate_speed_rad_per_s", - "encoder_bitrate_bps", - "encoder_gop", - ): - if name in options: - overrides[name] = options[name] - return replace(runtime_config, **overrides) if overrides else runtime_config - - __all__ = [ "OmnidreamsDemoAdapter", "ReplayRuntimeFactory", - "WebRTCRuntimeFactory", ] diff --git a/integrations/omnidreams/omnidreams/demo/cli.py b/integrations/omnidreams/omnidreams/demo/app.py similarity index 85% rename from integrations/omnidreams/omnidreams/demo/cli.py rename to integrations/omnidreams/omnidreams/demo/app.py index d35a62b78..1366643f9 100644 --- a/integrations/omnidreams/omnidreams/demo/cli.py +++ b/integrations/omnidreams/omnidreams/demo/app.py @@ -7,24 +7,17 @@ import argparse from pathlib import Path +from typing import Any -import torch -import torch.distributed as dist from omnidreams.runner import DEFAULT_EXAMPLE_DATA_UUID_1V -from flashdreams.core.distributed import init as distributed_init from flashdreams.runtime import InferenceConfig from flashdreams.runtime.demo import ( DemoSpec, Mp4OutputSpec, WebRTCOutputSpec, - run_flashdreams_demo, - serve_flashdreams_demo, -) -from flashdreams.serving.webrtc.bootstrap import ( - configure_logging, - initialize_cuda_distributed, ) +from flashdreams.runtime.demo.app import DemoApplication from .adapter import OmnidreamsDemoAdapter from .spec import ( @@ -80,33 +73,37 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: webrtc.add_argument("--warmup-timeout-s", type=float, default=600.0) webrtc.add_argument("--client-liveness-timeout-s", type=float, default=10.0) webrtc.add_argument("--debug-serve-hdmaps", action="store_true") - webrtc.add_argument("--postprocess-preset", default="") webrtc.add_argument("--prefer-sw-encoder", action="store_true") return parser.parse_args(argv) -def main(argv: list[str] | None = None) -> None: - configure_logging() - args = parse_args(argv) - adapter = OmnidreamsDemoAdapter() - if args.command == "replay": - run_flashdreams_demo(spec=_replay_spec(args), adapter=adapter) - return - if args.command == "webrtc": - context = initialize_cuda_distributed( - default_device=args.device, - distributed_init_fn=distributed_init, - configure_logging_fn=configure_logging, - torch_module=torch, - dist_module=dist, - ) - serve_flashdreams_demo( +class OmnidreamsDemoApplication(DemoApplication): + """OmniDreams replay and WebRTC demo application.""" + + def parse_args(self, argv: list[str] | None = None) -> argparse.Namespace: + return parse_args(argv) + + def replay_spec(self, args: argparse.Namespace) -> DemoSpec: + return _replay_spec(args) + + def replay_adapter(self) -> OmnidreamsDemoAdapter: + return OmnidreamsDemoAdapter() + + def serve_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: + from .webrtc import serve_omnidreams_webrtc_demo + + serve_omnidreams_webrtc_demo( spec=_webrtc_spec(args, device=str(context.device)), - adapter=adapter, world_rank=context.world_rank, ) - return - raise AssertionError(f"Unhandled command: {args.command}") + + +_APPLICATION = OmnidreamsDemoApplication() + + +def main(argv: list[str] | None = None) -> None: + """Run the OmniDreams demo application.""" + _APPLICATION.main(argv) def _replay_spec(args: argparse.Namespace) -> DemoSpec: @@ -152,7 +149,6 @@ def _webrtc_spec(args: argparse.Namespace, *, device: str) -> DemoSpec: scene_variant=args.scene_variant, camera_name=args.camera_name, debug_serve_hdmaps=args.debug_serve_hdmaps, - postprocess_preset=args.postprocess_preset, prefer_sw_encoder=args.prefer_sw_encoder, ), output=WebRTCOutputSpec( diff --git a/integrations/omnidreams/omnidreams/demo/replay.py b/integrations/omnidreams/omnidreams/demo/replay.py index 908d84568..8ccb58650 100644 --- a/integrations/omnidreams/omnidreams/demo/replay.py +++ b/integrations/omnidreams/omnidreams/demo/replay.py @@ -6,14 +6,14 @@ from __future__ import annotations import os -import time -from collections.abc import Callable, Mapping +from collections.abc import Callable from dataclasses import dataclass from typing import Any import torch import torch.distributed as dist from loguru import logger +from omnidreams.model_session import OmnidreamsModelSessionCore from omnidreams.runner import _load_video from flashdreams.core.distributed import init as init_distributed @@ -22,7 +22,7 @@ DEFAULT_RUNNER_INSTALL_HINT, load_first_frame_tensor, ) -from flashdreams.infra.video_output import VideoStepResult +from flashdreams.infra.video_output import VideoOutputStream from flashdreams.runtime.config import InferenceConfig from flashdreams.runtime.inputs import InferenceInput from flashdreams.runtime.interfaces import InferenceSession @@ -114,9 +114,15 @@ def __init__( self.output_layout = output_layout self.dtype = torch.bfloat16 self._closed = False - self._step_index = 0 self._frame_start = 0 - self._cache = self._initialize_cache() + self._model_session = OmnidreamsModelSessionCore( + pipeline=pipeline, + output_stream_factory=lambda: VideoOutputStream( + postprocess_stream=None, + output_layout=self.output_layout, + ), + ) + self._model_session.reset(self._initialize_cache) self._hdmap_videos = self._load_hdmaps() if self.device.type == "cuda" and torch.cuda.is_available(): torch.cuda.synchronize(device=self.device) @@ -126,20 +132,21 @@ def __init__( def next_step_request(self) -> StepRequest | None: if self._closed: return None - if self._step_index >= self.scenario.total_blocks: + step_index = self._model_session.step_index + if step_index >= self.scenario.total_blocks: return None - num_frames = int(self.pipeline.get_num_frames(self._step_index)) + num_frames = self._model_session.next_num_frames() if self._frame_start + num_frames > self._hdmap_videos.shape[2]: return None - return StepRequest(step_index=self._step_index) + return StepRequest(step_index=step_index) def step(self, inputs: InferenceInput) -> StepResult: del inputs if self._closed: raise RuntimeError("OmniDreams replay session is closed.") - step_index = self._step_index - num_frames = int(self.pipeline.get_num_frames(step_index)) + step_index = self._model_session.step_index + num_frames = self._model_session.next_num_frames() frame_end = self._frame_start + num_frames logger.info( "OmniDreams demo replay step {} frames=[{}, {})", @@ -147,51 +154,23 @@ def step(self, inputs: InferenceInput) -> StepResult: self._frame_start, frame_end, ) - start_t = time.perf_counter() - video_chunk = self.pipeline.generate( - autoregressive_index=step_index, - cache=self._cache, - hdmap=self._hdmap_videos[:, :, self._frame_start : frame_end], + result = self._model_session.step( + self._hdmap_videos[:, :, self._frame_start : frame_end] ) - stats = self.pipeline.finalize( - autoregressive_index=step_index, - cache=self._cache, - ) - elapsed_s = time.perf_counter() - start_t - self._step_index += 1 self._frame_start = frame_end - - metrics = _numeric_stats(stats) - metrics.setdefault("model_step_s", elapsed_s) - return StepResult( - step_index=step_index, - output=VideoStepResult.from_video_chunk( - chunk_index=step_index, - video_chunk=video_chunk, - layout=self.output_layout, - stats=metrics, - ), - frame_count=num_frames, - metrics=metrics, - ) + return result def reset(self, inputs: InferenceInput | None = None) -> None: if inputs is not None: scenario = _scenario_from_inputs(inputs) if scenario != self.scenario: raise ValueError("OmniDreams replay reset cannot swap scenarios.") - cache = getattr(self, "_cache", None) - if cache is not None: - del self._cache - self._cache = self._initialize_cache() - self._step_index = 0 + self._model_session.reset(self._initialize_cache) self._frame_start = 0 def close(self) -> None: self._closed = True - cache = getattr(self, "_cache", None) - if cache is not None: - del self._cache + self._model_session.close() def _initialize_cache(self) -> Any: scenario = self.scenario @@ -255,16 +234,6 @@ def _scenario_from_inputs(inputs: InferenceInput) -> OmnidreamsReplayScenario: return scenario -def _numeric_stats(stats: Any) -> dict[str, float | int]: - if not isinstance(stats, Mapping): - return {} - return { - str(key): value - for key, value in stats.items() - if isinstance(value, (float, int)) and not isinstance(value, bool) - } - - def _is_torchrun_env() -> bool: return "RANK" in os.environ and "WORLD_SIZE" in os.environ diff --git a/integrations/omnidreams/omnidreams/demo/spec.py b/integrations/omnidreams/omnidreams/demo/spec.py index 0a5dcc062..6a4f147cf 100644 --- a/integrations/omnidreams/omnidreams/demo/spec.py +++ b/integrations/omnidreams/omnidreams/demo/spec.py @@ -18,10 +18,10 @@ _example_camera_names, ) from omnidreams.scenes import SCENE_VARIANT_DEFAULT -from omnidreams.webrtc.session import DEFAULT_WEBRTC_SCENE_UUID DEFAULT_OMNIDREAMS_PRESET = "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae" OMNIDREAMS_MODEL_ID = "omnidreams" +DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID = "0d404ff7-2b66-498c-b047-1ed8cded60d4" @dataclass(frozen=True, kw_only=True, slots=True) @@ -74,11 +74,10 @@ class OmnidreamsWebRTCScenario: """Scene/options for the shared WebRTC demo path.""" scene_dir: Path | None = None - scene_uuid: str | None = DEFAULT_WEBRTC_SCENE_UUID + scene_uuid: str | None = DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID scene_variant: str = SCENE_VARIANT_DEFAULT camera_name: str = "camera_front_wide_120fov" debug_serve_hdmaps: bool = False - postprocess_preset: str = "" prefer_sw_encoder: bool = False def __post_init__(self) -> None: @@ -166,11 +165,10 @@ def resolve_webrtc_scenario(value: Any) -> OmnidreamsWebRTCScenario: scene_dir = value.get("scene_dir") return OmnidreamsWebRTCScenario( scene_dir=Path(scene_dir) if scene_dir is not None else None, - scene_uuid=value.get("scene_uuid", DEFAULT_WEBRTC_SCENE_UUID), + scene_uuid=value.get("scene_uuid", DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID), scene_variant=str(value.get("scene_variant", SCENE_VARIANT_DEFAULT)), camera_name=str(value.get("camera_name", "camera_front_wide_120fov")), debug_serve_hdmaps=bool(value.get("debug_serve_hdmaps", False)), - postprocess_preset=str(value.get("postprocess_preset", "")), prefer_sw_encoder=bool(value.get("prefer_sw_encoder", False)), ) @@ -263,6 +261,7 @@ def _require_existing_paths(paths: tuple[Path, ...], *, label: str) -> None: __all__ = [ "DEFAULT_OMNIDREAMS_PRESET", + "DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID", "OMNIDREAMS_MODEL_ID", "OmnidreamsReplayScenario", "OmnidreamsWebRTCScenario", diff --git a/integrations/omnidreams/omnidreams/demo/web/adapter.js b/integrations/omnidreams/omnidreams/demo/web/adapter.js new file mode 100644 index 000000000..d07fb8cc1 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/web/adapter.js @@ -0,0 +1,17 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +export default { + modelName: "OmniDreams", + controls: [ + { + label: "Drive / Turn", + keys: [ + { key: "w", label: "Forward" }, + { key: "a", label: "Turn left" }, + { key: "s", label: "Backward" }, + { key: "d", label: "Turn right" }, + ], + }, + ], +} diff --git a/integrations/omnidreams/omnidreams/demo/webrtc.py b/integrations/omnidreams/omnidreams/demo/webrtc.py index 32c9e9a25..00857ea85 100644 --- a/integrations/omnidreams/omnidreams/demo/webrtc.py +++ b/integrations/omnidreams/omnidreams/demo/webrtc.py @@ -1,179 +1,584 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 - -"""OmniDreams WebRTC hooks for the shared demo API.""" +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""OmniDreams model runtime and browser hooks for the shared WebRTC demo.""" from __future__ import annotations -from typing import Any, cast - -from aiohttp import web -from omnidreams.webrtc.session import ( - OmnidreamsRuntimeConfig, - OmnidreamsRuntimeError, - OmnidreamsSessionInput, - _validate_requested_postprocess_preset, +import tempfile +import time +from collections.abc import Callable +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any + +import cv2 +import numpy as np +import torch +from loguru import logger +from omnidreams.conditioning.conditioning_wrapper import ( + AV_POSITIVE_PROMPT, + OmnidreamsConditioningState, + OmnidreamsConditioningWrapper, + TextPrompt, ) - -from flashdreams.plugins.registry import resolve_postprocess_preset -from flashdreams.runtime.demo import DemoSpec -from flashdreams.runtime.demo.webrtc import SharedDemoWebRTCSessionManager -from flashdreams.serving.webrtc.controls import WSAD_SUPPORTED_KEYS -from flashdreams.serving.webrtc.manager import DEFAULT_CLIENT_LIVENESS_TIMEOUT_S -from flashdreams.serving.webrtc.server import ( - SESSION_MANAGER_KEY, - SessionBusyError, - create_packaged_webrtc_app, +from omnidreams.conditioning.renderer import load_and_attach_ludus_scene +from omnidreams.conditioning.world_scenario.data_loaders import load_scene +from omnidreams.conditioning.world_scenario.settings import SETTINGS +from omnidreams.config import OMNIDREAMS_CONFIGS +from omnidreams.scenes import ( + SCENE_CLIPGT_DIRNAME, + SCENE_PROMPT_FILENAME, + SCENE_VARIANT_DEFAULT, + ensure_hf_scene_synced, + extract_local_scene, + prepare_clipgt_dir, + resolve_scene_assets, +) +from omnidreams.transformer import CosmosTransformerConfig + +from flashdreams.runtime import InferenceConfig, StepResult +from flashdreams.runtime.demo import DemoSpec, WebRTCAppResources, WebRTCOutputSpec +from flashdreams.runtime.demo.webrtc import ( + CreateWebRTCApp, + RunWebRTCServer, + serve_webrtc_demo, +) +from flashdreams.serving.webrtc.bootstrap import run_webrtc_server +from flashdreams.serving.webrtc.controls import ( + WSAD_SUPPORTED_KEYS, + CameraPoseIntegrator, + PoseSegment, ) -from flashdreams.serving.webrtc.server import ( - close_package_resources as _close_package_resources, +from flashdreams.serving.webrtc.encoders import EncoderBackend +from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager +from flashdreams.serving.webrtc.runtime import ThreadAffineDistributedWebRTCRuntime +from flashdreams.serving.webrtc.server import create_webrtc_app + +from .spec import ( + DEFAULT_OMNIDREAMS_PRESET, + DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID, + OMNIDREAMS_MODEL_ID, + resolve_webrtc_scenario, ) +WebRTCRuntimeFactory = Callable[..., Any] -class OmnidreamsDemoWebRTCSessionManager(SharedDemoWebRTCSessionManager): - """Shared WebRTC manager customized for OmniDreams session semantics.""" - _busy_message = "An Omnidreams session is already active." - _warmup_label = "Omnidreams WebRTC" - _runtime_error_types = (OmnidreamsRuntimeError,) - _close_session_on_generation_error = True - _resampler_supported_keys = WSAD_SUPPORTED_KEYS +class OmnidreamsWebRTCModelRuntimeError(RuntimeError): + """Raised when the OmniDreams demo runtime is used incorrectly.""" - runtime_config: OmnidreamsRuntimeConfig - _runtime: Any - def __init__( - self, - *, - runtime: Any, - runtime_config: OmnidreamsRuntimeConfig, - fps: int, - client_liveness_timeout_s: float = DEFAULT_CLIENT_LIVENESS_TIMEOUT_S, - ) -> None: +@dataclass(frozen=True, slots=True) +class OmnidreamsWebRTCModelRuntimeConfig: + """Configuration for one scene-driven OmniDreams WebRTC runtime.""" + + pipeline_config_name: str + """User-facing name of the selected OmniDreams pipeline.""" + + pipeline_config: Any + """Resolved single-view OmniDreams pipeline configuration.""" + + scene_dir: Path | None = None + """Local scene root; ``None`` downloads the selected Hugging Face scene.""" + + scene_uuid: str | None = DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID + """Scene UUID used for remote lookup or local archive selection.""" + + scene_variant: str = SCENE_VARIANT_DEFAULT + """Weather variant selected from the scene assets.""" + + seed: int | None = 42 + """Per-rollout seed; ``None`` selects fresh entropy for every session.""" + + device: str = "cuda:0" + """Device used for rendering and model inference.""" + + video_height: int = 704 + """Generated video height in pixels.""" + + video_width: int = 1280 + """Generated video width in pixels.""" + + fps: int = 30 + """Input sampling and output playback frame rate.""" + + camera_name: str = "camera_front_wide_120fov" + """Scene camera controlled by browser keyboard input.""" + + move_speed_per_s: float = 6.0 + """Forward and reverse translation speed in scene units per second.""" + + rotate_speed_rad_per_s: float = float(np.deg2rad(35.0)) + """Left and right rotation speed in radians per second.""" + + warmup_chunks: int = 10 + """Number of synthetic chunks generated before accepting sessions.""" + + warmup_timeout_s: float = 600.0 + """Maximum duration for WebRTC loopback warmup.""" + + debug_serve_hdmaps: bool = False + """Stream rendered conditioning frames without running video generation.""" + + encoder_backend: EncoderBackend = "auto" + """WebRTC video encoder selection policy.""" + + encoder_bitrate_bps: int = 6_000_000 + """Target WebRTC video bitrate in bits per second.""" + + encoder_gop: int = 30 + """WebRTC video encoder group-of-pictures length.""" + + +class OmnidreamsWebRTCModelRuntime( + ThreadAffineDistributedWebRTCRuntime[ + OmnidreamsWebRTCModelRuntimeConfig, + None, + ] +): + """Run one single-view OmniDreams scene with browser camera controls.""" + + def __init__(self, *, config: OmnidreamsWebRTCModelRuntimeConfig) -> None: super().__init__( - model_name=runtime_config.pipeline_config_name, - runtime=runtime, - runtime_config=runtime_config, - fps=fps, - client_liveness_timeout_s=client_liveness_timeout_s, - ) - self._pending_session_input: OmnidreamsSessionInput | None = None - - def _model_name(self) -> str: - return self.runtime_config.pipeline_config_name - - def _chunk_done_extra(self) -> dict[str, Any]: - return { - "stream": "hdmap" if self.runtime_config.debug_serve_hdmaps else "rgb", - "postprocess_preset": self._runtime.postprocess_preset, - } - - def _peek_pending_session_input(self) -> OmnidreamsSessionInput | None: - return self._pending_session_input - - def _clear_pending_session_input(self) -> None: - self._pending_session_input = None - - async def _reset_runtime_for_session( - self, session_input: OmnidreamsSessionInput | None - ) -> None: - await self._runtime.reset_for_new_session(session_input=session_input) - - def set_pending_session_input(self, session_input: OmnidreamsSessionInput) -> None: - if self.has_active_session(): - raise SessionBusyError(self._busy_message) - preset = session_input.postprocess_preset - if preset: - _validate_requested_postprocess_preset( - requested_preset=preset, - configured_preset=self.runtime_config.postprocess.preset, + config=config, + runtime_error_type=OmnidreamsWebRTCModelRuntimeError, + thread_name="omnidreams-demo-runtime", + ) + self.pose_integrator = self._new_pose_integrator() + self._wrapper: OmnidreamsConditioningWrapper | None = None + self._state: OmnidreamsConditioningState | None = None + self._renderer: Any | None = None + self._scene_data: Any | None = None + self._initial_rgb_frames: torch.Tensor | 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._step_index = 0 + self._next_timestamp_us = 0 + self._clipgt_temp_dir: tempfile.TemporaryDirectory[str] | None = None + + def _new_pose_integrator(self) -> CameraPoseIntegrator: + return CameraPoseIntegrator( + move_speed_per_s=self.config.move_speed_per_s, + rotate_speed_rad_per_s=self.config.rotate_speed_rad_per_s, + coordinate_system="FLU", + ) + + def _is_runtime_initialized(self) -> bool: + return self._wrapper is not None and self._renderer is not None + + def _runtime_step_index(self) -> int: + return self._step_index + + def _next_input_frame_count(self) -> int: + wrapper = self._require_wrapper() + if self._state is None: + return int(wrapper.initial_frame_chunk_size) + return int(wrapper.frame_chunk_size) + + def _steady_output_frame_count(self) -> int: + return int(self._require_wrapper().frame_chunk_size) + + def _initialize_sync(self) -> None: + if self._wrapper is not None: + return + + init_t0 = time.perf_counter() + cfg = self.config + transformer_cfg = cfg.pipeline_config.diffusion_model.transformer + if not isinstance(transformer_cfg, CosmosTransformerConfig): + raise TypeError( + "OmniDreams WebRTC requires a CosmosTransformerConfig pipeline." ) - self._pending_session_input = session_input - - -async def postprocess_options(request: web.Request) -> web.StreamResponse: - """Return the postprocess preset selected at server launch.""" - manager = _get_omnidreams_manager(request.app) - configured_preset = manager.runtime_config.postprocess.preset - presets = [configured_preset] if configured_preset else [] - return web.json_response( - { - "default_preset": configured_preset, - "presets": presets, - } - ) + if transformer_cfg.num_views != 1: + raise ValueError( + "OmniDreams WebRTC supports only single-view configs; " + f"{cfg.pipeline_config_name!r} has num_views=" + f"{transformer_cfg.num_views}." + ) + if self._device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for OmniDreams WebRTC inference.") + + scene_dir = self._prepare_scene() + clipgt_dir, first_frame_path, prompt_path = resolve_scene_assets( + scene_dir, + prompt_filename=SCENE_PROMPT_FILENAME, + clipgt_dirname=SCENE_CLIPGT_DIRNAME, + camera_name=cfg.camera_name, + variant=cfg.scene_variant, + ) + self._initial_rgb_frames = self._load_first_frame(first_frame_path) + prompt = prompt_path.read_text(encoding="utf-8").strip() or AV_POSITIVE_PROMPT + self._text_prompts = [TextPrompt(positive=prompt)] + + loadable_clipgt_dir, self._clipgt_temp_dir = prepare_clipgt_dir(clipgt_dir) + logger.info("Loading OmniDreams scene data from {}", loadable_clipgt_dir) + scene_data = load_scene( + loadable_clipgt_dir, + camera_names=[cfg.camera_name], + max_frames=-1, + input_pose_fps=SETTINGS["INPUT_POSE_FPS"], + resize_resolution_hw=(cfg.video_height, cfg.video_width), + ) + scene_data = load_and_attach_ludus_scene( + loadable_clipgt_dir, + scene_data, + device=self._device, + ) + self._validate_scene_data(scene_data, scene_dir=loadable_clipgt_dir) + logger.info( + "Setting up OmniDreams pipeline {} on {}.", + cfg.pipeline_config_name, + self._device, + ) + 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, + ) + renderer = wrapper.create_renderer(scene_data, [cfg.camera_name]) + + self._wrapper = wrapper + self._renderer = renderer + self._scene_data = scene_data + self._camera_to_rig = torch.as_tensor( + scene_data.camera_extrinsics[cfg.camera_name], + device=self._device, + dtype=torch.float32, + ) + 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, + ) -async def session_input(request: web.Request) -> web.StreamResponse: - """Apply browser-selected settings to the next WebRTC rollout.""" - try: - payload = await request.json() - except Exception as exc: - raise web.HTTPBadRequest(reason="Expected JSON session input.") from exc - if not isinstance(payload, dict): - raise web.HTTPBadRequest(reason="Session input must be a JSON object.") - preset = payload.get("postprocess_preset") - if not isinstance(preset, str): - raise web.HTTPBadRequest( - reason="Session input must include string 'postprocess_preset'." - ) - - manager = _get_omnidreams_manager(request.app) - try: - manager.set_pending_session_input( - OmnidreamsSessionInput(postprocess_preset=preset) + def _prepare_scene(self) -> Path: + cfg = self.config + if cfg.scene_dir is None: + return ensure_hf_scene_synced( + cfg.scene_uuid or DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID, + variant=cfg.scene_variant, + clipgt_dirname=SCENE_CLIPGT_DIRNAME, + ) + return extract_local_scene( + cfg.scene_dir, + scene_uuid=cfg.scene_uuid, + variant=cfg.scene_variant, + clipgt_dirname=SCENE_CLIPGT_DIRNAME, ) - except SessionBusyError as exc: - raise web.HTTPConflict(reason=str(exc)) from exc - except ValueError as exc: - raise web.HTTPBadRequest(reason=str(exc)) from exc - return web.json_response({"postprocess_preset": preset}) + def _load_first_frame(self, path: Path) -> torch.Tensor: + logger.info("Loading OmniDreams first frame from {}", path) + image_bgr = cv2.imread(str(path), cv2.IMREAD_COLOR) + if image_bgr is None: + raise RuntimeError(f"Failed to read first frame from {path}") + image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) + image_rgb = cv2.resize( + image_rgb, + (self.config.video_width, self.config.video_height), + interpolation=cv2.INTER_CUBIC, + ) + return ( + torch.from_numpy(image_rgb) + .permute(2, 0, 1) + .contiguous() + .unsqueeze(0) + .unsqueeze(0) + .to(device=self._device, dtype=torch.uint8) + ) -def configure_omnidreams_webrtc_app(app: web.Application) -> None: - """Register OmniDreams browser support routes on a shared WebRTC app.""" - app.router.add_get("/api/postprocess/options", postprocess_options) - app.router.add_post("/api/session/input", session_input) + def _validate_scene_data(self, scene_data: Any, *, scene_dir: Path) -> None: + camera_name = self.config.camera_name + if not scene_data.ego_poses: + raise ValueError(f"Scene {scene_dir} has no ego poses.") + if camera_name not in scene_data.camera_models: + raise ValueError(f"Camera {camera_name!r} was not loaded from {scene_dir}.") + if camera_name not in scene_data.camera_extrinsics: + raise ValueError( + f"Camera {camera_name!r} has no extrinsics in {scene_dir}." + ) + def _reset_rollout_sync(self, session_input: None = None) -> None: + del session_input + wrapper = self._require_wrapper() + if self._renderer is None or self._scene_data is None: + raise OmnidreamsWebRTCModelRuntimeError("Scene state is not initialized.") + if self._initial_ego_pose is None: + raise OmnidreamsWebRTCModelRuntimeError( + "Initial camera pose is unavailable." + ) + + if self._state is not None and self._state.pipeline_cache is not None: + del self._state.pipeline_cache + self._state = None + self._step_index = 0 + self.pose_integrator = self._new_pose_integrator() + self.pose_integrator.reset(self._initial_ego_pose) + self._next_timestamp_us = int(self._scene_data.ego_poses[0].timestamp) + wrapper.set_rollout_seed(self.config.seed) -def create_omnidreams_webrtc_app( + def _generate_one_chunk_sync( + self, + *, + segments: list[PoseSegment], + frame_times: list[float], + ) -> StepResult: + wrapper = self._require_wrapper() + if ( + self._renderer is None + or self._initial_rgb_frames is None + or self._text_prompts is None + or self._camera_to_rig is None + ): + raise OmnidreamsWebRTCModelRuntimeError("Runtime is not initialized.") + if len(frame_times) != self._next_input_frame_count(): + raise OmnidreamsWebRTCModelRuntimeError( + f"Expected {self._next_input_frame_count()} frame times for " + f"step {self._step_index}, got {len(frame_times)}." + ) + if not segments: + raise OmnidreamsWebRTCModelRuntimeError( + f"Step {self._step_index} received no control segments." + ) + + ego_poses = self.pose_integrator.integrate_chunk( + segments=segments, + frame_times=frame_times, + ) + ego_poses_t = torch.from_numpy(ego_poses).to( + device=self._device, + dtype=torch.float32, + ) + camera_poses = torch.einsum("nij,jk->nik", ego_poses_t, self._camera_to_rig) + frame_timestamps_us = self._consume_timestamps(len(frame_times)) + serve_hdmaps = self.config.debug_serve_hdmaps + + if self._state is None: + output = wrapper.start_generation( + text_prompts=self._text_prompts, + initial_rgb_frames=self._initial_rgb_frames, + renderer=self._renderer, + camera_names=[self.config.camera_name], + camera_poses_per_view={self.config.camera_name: camera_poses}, + frame_timestamps_us=frame_timestamps_us, + skip_video_generation=serve_hdmaps, + ) + else: + output = wrapper.continue_generation( + state=self._state, + camera_names=[self.config.camera_name], + camera_poses_per_view={self.config.camera_name: camera_poses}, + frame_timestamps_us=frame_timestamps_us, + skip_video_generation=serve_hdmaps, + ) + self._state = output.state + if self._state.pipeline_cache is not None: + wrapper.finalize_block_generation( + self._state.pipeline_cache, + output.finalization_state, + ) + + metadata = {"stream": "hdmap" if serve_hdmaps else "rgb"} + if serve_hdmaps: + video_chunk = output.condition_frames + else: + if output.rgb_frames is None: + raise OmnidreamsWebRTCModelRuntimeError( + "OmniDreams generation produced no RGB frames." + ) + video_chunk = output.rgb_frames + result = StepResult.from_video_chunk( + step_index=self._step_index, + video_chunk=video_chunk.detach(), + layout="bvtchw", + metadata=metadata, + ) + expected_frames = len(frame_times) + if result.frame_count != expected_frames: + raise OmnidreamsWebRTCModelRuntimeError( + f"Expected generated chunk to contain {expected_frames} frames, " + f"got {result.frame_count}." + ) + self._step_index += 1 + return result + + def _consume_timestamps(self, num_frames: int) -> list[int]: + step_us = int(round(1_000_000 / self.config.fps)) + timestamps = [ + self._next_timestamp_us + frame_index * step_us + for frame_index in range(num_frames) + ] + self._next_timestamp_us += num_frames * step_us + return timestamps + + def _close_sync(self) -> None: + if self._wrapper is not None and self._state is not None: + self._wrapper.cleanup(self._state) + elif self._renderer is not None: + self._renderer.cleanup() + self._state = None + self._wrapper = None + self._renderer = None + self._scene_data = None + self._initial_rgb_frames = None + self._text_prompts = None + self._camera_to_rig = None + self._initial_ego_pose = None + if self._clipgt_temp_dir is not None: + self._clipgt_temp_dir.cleanup() + self._clipgt_temp_dir = None + if self._device.type == "cuda": + torch.cuda.synchronize(device=self._device) + torch.cuda.empty_cache() + + def _require_wrapper(self) -> OmnidreamsConditioningWrapper: + if self._wrapper is None: + raise OmnidreamsWebRTCModelRuntimeError("Runtime is not initialized.") + return self._wrapper + + +def serve_omnidreams_webrtc_demo( *, spec: DemoSpec, - session_manager: Any, - request_session_url: str, -) -> web.Application: - """Create the packaged OmniDreams browser app through shared serving glue.""" - from importlib.resources import as_file, files - - output_preload_name = getattr(spec.output, "preload_name", None) - preload_name = output_preload_name if isinstance(output_preload_name, str) else "" - return create_packaged_webrtc_app( - web_resource=files("flashdreams.serving.webrtc").joinpath("web"), - model_web_resource=files("omnidreams.webrtc").joinpath("web"), - session_manager=session_manager, - preload_name=preload_name or "Omnidreams", - request_session_url=request_session_url, - configure_app=configure_omnidreams_webrtc_app, - as_file_fn=as_file, - cleanup_callback=_close_package_resources, + world_rank: int = 0, + runtime_factory: WebRTCRuntimeFactory = OmnidreamsWebRTCModelRuntime, + create_app_fn: CreateWebRTCApp = create_webrtc_app, + server_runner: RunWebRTCServer = run_webrtc_server, +) -> object: + """Create OmniDreams' runtime and serve it through the shared WebRTC transport.""" + if spec.input_mode != "keyboard-driving": + raise ValueError( + "OmniDreams WebRTC requires input_mode='keyboard-driving', " + f"got {spec.input_mode!r}." + ) + if not isinstance(spec.output, WebRTCOutputSpec): + raise ValueError("OmniDreams WebRTC requires WebRTC output.") + config = spec.config + if config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + if config.model_id != OMNIDREAMS_MODEL_ID: + raise ValueError( + f"OmniDreams WebRTC requires model_id={OMNIDREAMS_MODEL_ID!r}, " + f"got {config.model_id!r}." + ) + scenario = resolve_webrtc_scenario(spec.scenario) + preset_id = _preset_id(config) + seed = _option(config, "seed", 42) + runtime_config = OmnidreamsWebRTCModelRuntimeConfig( + pipeline_config_name=preset_id, + pipeline_config=_pipeline_config(config), + scene_dir=scenario.scene_dir, + scene_uuid=scenario.scene_uuid, + scene_variant=scenario.scene_variant, + seed=None if seed is None else int(seed), + device=config.device or str(_option(config, "device", "cuda:0")), + video_height=spec.output.video_height, + video_width=spec.output.video_width, + fps=spec.output.fps, + camera_name=scenario.camera_name, + warmup_chunks=spec.output.warmup_chunks, + warmup_timeout_s=spec.output.warmup_timeout_s, + debug_serve_hdmaps=scenario.debug_serve_hdmaps, + encoder_backend="default" if scenario.prefer_sw_encoder else "auto", + ) + runtime_config = _apply_runtime_options(runtime_config, config.runtime_options) + runtime = runtime_factory(config=runtime_config) + manager = BaseWebRTCSessionManager( + runtime=runtime, + runtime_config=runtime_config, + fps=runtime_config.fps, + identity=runtime_config.pipeline_config_name, + busy_message="An OmniDreams session is already active.", + warmup_label="OmniDreams WebRTC", + supported_control_keys=WSAD_SUPPORTED_KEYS, + fatal_generation_errors=True, + client_liveness_timeout_s=spec.output.client_liveness_timeout_s, + ) + from importlib.resources import files + + return serve_webrtc_demo( + output=spec.output, + model_id=spec.model_id, + session_manager=manager, + app_resources=WebRTCAppResources( + model_web_resource=files("omnidreams.demo").joinpath("web"), + preload_name="OmniDreams", + ), + world_rank=world_rank, + create_app_fn=create_app_fn, + server_runner=server_runner, ) -def validate_postprocess_preset(preset: str) -> None: - """Validate a configured preset without enabling the output system broadly.""" - if preset: - resolve_postprocess_preset(preset) +def _preset_id(config: InferenceConfig | None) -> str: + return ( + DEFAULT_OMNIDREAMS_PRESET + if config is None or config.preset_id is None + else config.preset_id + ) -def _get_omnidreams_manager(app: web.Application) -> OmnidreamsDemoWebRTCSessionManager: - return cast(OmnidreamsDemoWebRTCSessionManager, app[SESSION_MANAGER_KEY]) +def _pipeline_config(config: InferenceConfig) -> Any: + custom = config.runtime_options.get("pipeline_config") + if custom is not None: + return custom + preset_id = _preset_id(config) + try: + return OMNIDREAMS_CONFIGS[preset_id] + except KeyError as exc: + supported = ", ".join(sorted(OMNIDREAMS_CONFIGS)) + raise ValueError( + f"Unsupported OmniDreams preset_id={preset_id!r}. " + f"Supported presets: {supported}." + ) from exc + + +def _option(config: InferenceConfig, name: str, default: Any) -> Any: + return config.runtime_options.get(name, default) + + +def _apply_runtime_options( + runtime_config: OmnidreamsWebRTCModelRuntimeConfig, + options: Any, +) -> OmnidreamsWebRTCModelRuntimeConfig: + if not isinstance(options, dict): + options = dict(options) + overrides = { + name: options[name] + for name in ( + "move_speed_per_s", + "rotate_speed_rad_per_s", + "encoder_bitrate_bps", + "encoder_gop", + ) + if name in options + } + return replace(runtime_config, **overrides) if overrides else runtime_config __all__ = [ - "OmnidreamsDemoWebRTCSessionManager", - "configure_omnidreams_webrtc_app", - "create_omnidreams_webrtc_app", - "postprocess_options", - "session_input", - "validate_postprocess_preset", + "OmnidreamsWebRTCModelRuntime", + "OmnidreamsWebRTCModelRuntimeConfig", + "OmnidreamsWebRTCModelRuntimeError", + "WebRTCRuntimeFactory", + "serve_omnidreams_webrtc_demo", ] diff --git a/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py b/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py index c74c7fb3e..1dff41eac 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py @@ -18,6 +18,7 @@ build_synthetic_world_model_assets, default_synthetic_asset_dir, ) +from omnidreams.model_session import OmnidreamsModelSessionCore from flashdreams.infra.acceleration.encoder_lifecycle import ( collect_and_release_cuda_memory, @@ -31,10 +32,7 @@ VideoPostprocessChainConfig, VideoPostprocessStream, ) -from flashdreams.infra.video_output import ( - VideoOutputStream, - lazy_rgb_frames_from_video_tensor, -) +from flashdreams.infra.video_output import VideoOutputStream PipelineFactory = Callable[[WorldModelManifest, WorldModelProfileConfig], Any] _VIEW_NAMES = ["camera_front_wide_120fov"] @@ -496,13 +494,10 @@ def __init__( self._offload_text_encoder = bool(offload_text_encoder) self._pipeline_factory = pipeline_factory self._pipeline: Any | None = None - self._cache: Any | None = None + self._model_session: OmnidreamsModelSessionCore | None = None self._precomputed_embeddings: dict[str, torch.Tensor | None] | None = None - self._pending_finalization_index: int | None = None - self._next_block_index = 0 self._postprocess = postprocess or VideoPostprocessChainConfig() self._postprocess_enabled = self._postprocess.is_enabled() - self._output_stream: VideoOutputStream | None = None @property def pipeline(self) -> Any: @@ -610,6 +605,9 @@ def _validate_chunk_sizes(self) -> None: def _release_pipeline(self) -> None: if self._pipeline is None: return + if self._model_session is not None: + self._model_session.close() + self._model_session = None self._pipeline = None device = torch.device(self.manifest.device) collect_and_release_cuda_memory( @@ -624,7 +622,8 @@ def start( condition_frames: list[object], prompt: str, ) -> list[object]: - expected_frames = self.pipeline.get_num_frames(0) + model_session = self._ensure_model_session() + expected_frames = model_session.next_num_frames() if len(condition_frames) != expected_frames: raise ValueError( "First condition chunk length does not match flashdreams initial chunk size: " @@ -633,25 +632,22 @@ def start( start = time.perf_counter() with torch.no_grad(): - self._cache = self._initialize_cache(initial_rgb, prompt) - video = self.pipeline.generate( - autoregressive_index=0, - cache=self._cache, - hdmap=self._condition_tensor(condition_frames), + model_session.reset(lambda: self._initialize_cache(initial_rgb, prompt)) + result = model_session.step( + self._condition_tensor(condition_frames), + delay_finalization=True, ) - video = self._process_video(video, autoregressive_index=0) - model_frames = self._video_tensor_to_frames(video) + model_frames = list(result.lazy_rgb_frames()) _synchronize_cuda_frame_event(model_frames) - self._pending_finalization_index = 0 - self._next_block_index = 1 elapsed_ms = (time.perf_counter() - start) * 1000.0 logger.info(f"[flashdreams-session] start total_ms={elapsed_ms:.1f}") return model_frames def continue_generation(self, condition_frames: list[object]) -> list[object]: - if self._cache is None: + model_session = self._model_session + if model_session is None or not model_session.initialized: raise RuntimeError("start() must be called before continue_generation()") - expected_frames = self.pipeline.get_num_frames(self._next_block_index) + expected_frames = model_session.next_num_frames() if len(condition_frames) != expected_frames: raise ValueError( "Condition chunk length does not match flashdreams steady-state chunk size: " @@ -660,22 +656,13 @@ def continue_generation(self, condition_frames: list[object]) -> list[object]: start = time.perf_counter() with torch.no_grad(): - if self._pending_finalization_index is not None: - self.pipeline.finalize(self._pending_finalization_index, self._cache) - self._pending_finalization_index = None - video = self.pipeline.generate( - autoregressive_index=self._next_block_index, - cache=self._cache, - hdmap=self._condition_tensor(condition_frames), - ) - video = self._process_video( - video, autoregressive_index=self._next_block_index + result = model_session.step( + self._condition_tensor(condition_frames), + delay_finalization=True, ) - model_frames = self._video_tensor_to_frames(video) + model_frames = list(result.lazy_rgb_frames()) _synchronize_cuda_frame_event(model_frames) - block_index = self._next_block_index - self._pending_finalization_index = block_index - self._next_block_index += 1 + block_index = result.step_index elapsed_ms = (time.perf_counter() - start) * 1000.0 if block_index <= 3 or elapsed_ms > 500.0: logger.info( @@ -684,10 +671,8 @@ def continue_generation(self, condition_frames: list[object]) -> list[object]: return model_frames def reset(self, *, clear_precomputed_embeddings: bool = False) -> None: - self._close_postprocess_stream() - self._cache = None - self._pending_finalization_index = None - self._next_block_index = 0 + if self._model_session is not None: + self._model_session.clear(finalize_pending=False) if clear_precomputed_embeddings: self._precomputed_embeddings = None logger.info( @@ -696,11 +681,9 @@ def reset(self, *, clear_precomputed_embeddings: bool = False) -> None: ) def close(self) -> None: - self._close_postprocess_stream() - if self._cache is not None and self._pending_finalization_index is not None: - self.pipeline.finalize(self._pending_finalization_index, self._cache) - self._pending_finalization_index = None - self._cache = None + if self._model_session is not None: + self._model_session.close() + self._model_session = None self._pipeline = None def set_postprocess_enabled(self, enabled: bool) -> None: @@ -712,8 +695,9 @@ def set_postprocess_enabled(self, enabled: bool) -> None: ) if enabled == self._postprocess_enabled: return - self._close_postprocess_stream() self._postprocess_enabled = enabled + if self._model_session is not None: + self._model_session.replace_output_stream(self._new_output_stream) logger.info( "[flashdreams-session] post-processing {} preset={!r}", "enabled" if enabled else "disabled", @@ -733,32 +717,15 @@ def _new_output_stream(self) -> VideoOutputStream: return VideoOutputStream( postprocess_stream=postprocess_stream, output_layout="bvtchw", - collect_output=False, - move_to_cpu=False, ) - def _process_video( - self, video: torch.Tensor, *, autoregressive_index: int - ) -> torch.Tensor: - if self._output_stream is None: - self._output_stream = self._new_output_stream() - processed = self._output_stream.process( - video, - autoregressive_index=autoregressive_index, - ) - if processed.shape[2] != video.shape[2]: - raise RuntimeError( - "Interactive post-processing must emit one display frame for " - "each generated frame; got " - f"{processed.shape[2]} output frames for {video.shape[2]} inputs." + def _ensure_model_session(self) -> OmnidreamsModelSessionCore: + if self._model_session is None: + self._model_session = OmnidreamsModelSessionCore( + pipeline=self.pipeline, + output_stream_factory=self._new_output_stream, ) - return processed - - def _close_postprocess_stream(self) -> None: - if self._output_stream is None: - return - self._output_stream.finish() - self._output_stream = None + return self._model_session def _initialize_cache(self, initial_rgb: object, prompt: str) -> Any: if self.manifest.synthetic_model: @@ -856,21 +823,6 @@ def _condition_tensor(self, condition_frames: Sequence[object]) -> torch.Tensor: def _to_model_range(self, tensor: torch.Tensor) -> torch.Tensor: return _to_model_range(tensor, device=self.pipeline.device) - @staticmethod - def _video_tensor_to_frames(video: torch.Tensor) -> list[object]: - if video.ndim != 6: - raise ValueError( - f"Expected [B,V,T,3,H,W] video tensor, got shape {tuple(video.shape)}" - ) - return list( - lazy_rgb_frames_from_video_tensor( - video, - layout="bvtchw", - batch_index=0, - view_index=0, - ) - ) - def _rgb_hwc_uint8(frame: object) -> np.ndarray: return np.ascontiguousarray( diff --git a/integrations/omnidreams/omnidreams/model_session.py b/integrations/omnidreams/omnidreams/model_session.py new file mode 100644 index 000000000..96dbae65e --- /dev/null +++ b/integrations/omnidreams/omnidreams/model_session.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared synchronous OmniDreams pipeline-session execution.""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Mapping +from typing import Any + +import torch + +from flashdreams.infra.video_output import VideoOutputStream +from flashdreams.runtime import StepResult + +CacheFactory = Callable[[], Any] +OutputStreamFactory = Callable[[], VideoOutputStream] + + +class OmnidreamsModelSessionCore: + """Own one raw-pipeline cache, AR index, finalization, and output stream.""" + + def __init__( + self, + *, + pipeline: Any, + output_stream_factory: OutputStreamFactory, + ) -> None: + self.pipeline = pipeline + self._output_stream_factory = output_stream_factory + self._output_stream = output_stream_factory() + self._cache: Any | None = None + self._step_index = 0 + self._pending_finalization_index: int | None = None + self._closed = False + + @property + def step_index(self) -> int: + return self._step_index + + @property + def initialized(self) -> bool: + return self._cache is not None and not self._closed + + def next_num_frames(self) -> int: + self._require_open() + return int(self.pipeline.get_num_frames(self._step_index)) + + def reset(self, cache_factory: CacheFactory) -> None: + self._require_open() + if self._cache is not None or self._step_index != 0: + self._clear(finalize_pending=False, recreate_output_stream=True) + self._cache = cache_factory() + + def step( + self, + hdmap: torch.Tensor, + *, + delay_finalization: bool = False, + metadata: Mapping[str, Any] | None = None, + ) -> StepResult: + self._require_initialized() + self._finalize_pending() + step_index = self._step_index + expected_frames = self.next_num_frames() + start_t = time.perf_counter() + video_chunk = self.pipeline.generate( + autoregressive_index=step_index, + cache=self._cache, + hdmap=hdmap, + ) + metrics: dict[str, float | int] = {} + if delay_finalization: + self._pending_finalization_index = step_index + else: + metrics = _numeric_metrics( + self.pipeline.finalize( + autoregressive_index=step_index, + cache=self._cache, + ) + ) + metrics.setdefault("model_step_s", time.perf_counter() - start_t) + result = self._output_stream.process( + video_chunk, + autoregressive_index=step_index, + metrics=metrics, + metadata=metadata, + ) + if result.frame_count != expected_frames: + raise RuntimeError( + f"Expected generated chunk to contain {expected_frames} frames, " + f"got {result.frame_count}." + ) + self._step_index += 1 + return result + + def replace_output_stream(self, output_stream_factory: OutputStreamFactory) -> None: + self._require_open() + self._output_stream.finish() + self._output_stream_factory = output_stream_factory + self._output_stream = output_stream_factory() + + def finish_output(self) -> StepResult | None: + """Flush and return the output postprocessor tail, when present.""" + self._require_open() + return self._output_stream.finish() + + def clear(self, *, finalize_pending: bool = False) -> None: + self._require_open() + self._clear( + finalize_pending=finalize_pending, + recreate_output_stream=True, + ) + + def close(self) -> None: + if self._closed: + return + self._clear(finalize_pending=True, recreate_output_stream=False) + self._closed = True + + def _clear( + self, + *, + finalize_pending: bool, + recreate_output_stream: bool, + ) -> None: + if finalize_pending: + self._finalize_pending() + else: + self._pending_finalization_index = None + self._cache = None + self._step_index = 0 + self._output_stream.finish() + if recreate_output_stream: + self._output_stream = self._output_stream_factory() + + def _finalize_pending(self) -> None: + if self._cache is None or self._pending_finalization_index is None: + return + self.pipeline.finalize( + autoregressive_index=self._pending_finalization_index, + cache=self._cache, + ) + self._pending_finalization_index = None + + def _require_open(self) -> None: + if self._closed: + raise RuntimeError("OmniDreams model session is closed.") + + def _require_initialized(self) -> None: + self._require_open() + if self._cache is None: + raise RuntimeError("OmniDreams model session is not initialized.") + + +def _numeric_metrics(stats: object) -> dict[str, float | int]: + if not isinstance(stats, Mapping): + return {} + return { + str(name): value + for name, value in stats.items() + if isinstance(value, (int, float)) and not isinstance(value, bool) + } + + +__all__ = ["OmnidreamsModelSessionCore"] diff --git a/integrations/omnidreams/omnidreams/output_targets.py b/integrations/omnidreams/omnidreams/output_targets.py new file mode 100644 index 000000000..6bce82176 --- /dev/null +++ b/integrations/omnidreams/omnidreams/output_targets.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OmniDreams output capabilities for ``flashdreams-run``.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from flashdreams.infra.runner import RunnerConfig +from flashdreams.serving.output_targets import ( + OutputLaunchOptions, + OutputMode, + OutputTargetSpec, +) + +_LOCAL_WINDOW_MANIFESTS = { + "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae": "example_world_model.yaml", + "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf": ( + "example_world_model_perf.yaml" + ), + "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-native-perf": ( + "example_world_model_perf.yaml" + ), +} + + +class OmnidreamsOutputTargetAdapter: + def supported_modes( + self, + config: RunnerConfig, + options: OutputLaunchOptions, + ) -> tuple[OutputMode, ...]: + modes: list[OutputMode] = [] + if _is_single_view(config): + modes.append("webrtc") + if _local_window_manifest(config, options) is not None: + modes.append("local-window") + return tuple(modes) + + def resolve( + self, + config: RunnerConfig, + *, + mode: OutputMode, + options: OutputLaunchOptions, + ) -> OutputTargetSpec | None: + if mode == "webrtc" and _is_single_view(config): + return _webrtc_spec(config, options) + if mode == "local-window": + manifest = _local_window_manifest(config, options) + if manifest is not None: + return _local_window_spec(config, manifest) + return None + + +def _webrtc_spec( + config: RunnerConfig, + options: OutputLaunchOptions, +) -> OutputTargetSpec: + argv = [ + "webrtc", + "--preset-id", + _pipeline_name(config), + "--device", + str(config.device), + "--fps", + str(getattr(config, "output_fps", 30)), + "--video-height", + str(getattr(config, "pixel_height", 704)), + "--video-width", + str(getattr(config, "pixel_width", 1280)), + ] + seed = _diffusion_seed(config) + if seed is not None: + argv.extend(("--seed", str(seed))) + _append_postprocess_preset(argv, config) + if options.host: + argv.extend(("--host", options.host)) + if options.port is not None: + argv.extend(("--port", str(options.port))) + if options.prefer_sw_encoder: + argv.append("--prefer-sw-encoder") + return OutputTargetSpec( + mode="webrtc", + label="OmniDreams shared demo WebRTC server", + module="omnidreams.demo.app", + argv=tuple(argv), + ) + + +def _local_window_spec(config: RunnerConfig, manifest: Path) -> OutputTargetSpec: + argv = ["--manifest", str(manifest)] + _append_postprocess_preset(argv, config) + return OutputTargetSpec( + mode="local-window", + label="Omnidreams local interactive window", + module="omnidreams.interactive_drive", + argv=tuple(argv), + notes=( + "Local-window uses the OmniDreams interactive-drive manifest for " + "scene, resolution, and runtime-specific controls.", + ), + ) + + +def _local_window_manifest( + config: RunnerConfig, + options: OutputLaunchOptions, +) -> Path | None: + if options.local_window_manifest is not None: + return options.local_window_manifest + manifest = _LOCAL_WINDOW_MANIFESTS.get(config.runner_name) + return None if manifest is None else Path(manifest) + + +def _pipeline_name(config: RunnerConfig) -> str: + name = getattr(config.pipeline, "name", None) + return str(name or config.runner_name) + + +def _diffusion_seed(config: RunnerConfig) -> int | None: + diffusion_model = getattr(config.pipeline, "diffusion_model", None) + seed = getattr(diffusion_model, "seed", None) + return None if seed is None else int(seed) + + +def _is_single_view(config: RunnerConfig) -> bool: + diffusion_model = getattr(config.pipeline, "diffusion_model", None) + transformer: Any = getattr(diffusion_model, "transformer", None) + return int(getattr(transformer, "num_views", 1)) == 1 + + +def _append_postprocess_preset(argv: list[str], config: RunnerConfig) -> None: + preset = config.postprocess.preset + if preset: + argv.extend(("--postprocess-preset", str(preset))) + + +OUTPUT_TARGET_ADAPTER = OmnidreamsOutputTargetAdapter() + +__all__ = ["OUTPUT_TARGET_ADAPTER", "OmnidreamsOutputTargetAdapter"] diff --git a/integrations/omnidreams/omnidreams/runner.py b/integrations/omnidreams/omnidreams/runner.py index 177cb7f2d..589f2aa58 100644 --- a/integrations/omnidreams/omnidreams/runner.py +++ b/integrations/omnidreams/omnidreams/runner.py @@ -33,6 +33,7 @@ import torch from einops import rearrange from loguru import logger +from omnidreams.model_session import OmnidreamsModelSessionCore from omnidreams.pipeline import ( OmnidreamsPipeline, OmnidreamsPipelineCache, @@ -48,7 +49,9 @@ load_video_tensor, runner_artifact_path, write_runner_stats, + write_video_tensor, ) +from flashdreams.infra.video_output import VideoResultCollector DEFAULT_VIDEO_HEIGHT = 704 """Pixel-space rollout height (matches the trained 720p chassis).""" @@ -153,6 +156,7 @@ class OmnidreamsRunnerConfig(RunnerConfig): """ _target: type["OmnidreamsRunner"] = field(default_factory=lambda: OmnidreamsRunner) + output_adapter: str | None = "omnidreams.output_targets:OUTPUT_TARGET_ADAPTER" prompt: str = "" """Default text prompt applied to every camera. Override per-camera @@ -390,10 +394,20 @@ def _rollout_and_save( if torch.distributed.is_initialized(): torch.distributed.barrier() - output_stream = self.create_video_output_stream(fps=cfg.output_fps) + output_collector = VideoResultCollector( + output_layout=self.config.postprocess_output_layout or "bvtchw", + enabled=self.is_rank_zero, + ) + model_session = OmnidreamsModelSessionCore( + pipeline=self.pipeline, + output_stream_factory=lambda: self.create_video_output_stream( + fps=cfg.output_fps + ), + ) + model_session.reset(lambda: cache) start = 0 for i in range(cfg.total_blocks): - num_frames = self.pipeline.get_num_frames(i) + num_frames = model_session.next_num_frames() end = start + num_frames if end > hdmap_num_frames: break @@ -402,16 +416,14 @@ def _rollout_and_save( f"[{cfg.runner_name}] AR step {i}/{cfg.total_blocks}, " f"num_frames={num_frames}, frames=[{start}, {end})" ) - video_chunk = self.pipeline.generate( - autoregressive_index=i, - cache=cache, - hdmap=hdmap_videos_t[:, :, start:end], - ) - stats = self.pipeline.finalize(autoregressive_index=i, cache=cache) - output_stream.process(video_chunk, autoregressive_index=i, stats=stats) + output_collector.add(model_session.step(hdmap_videos_t[:, :, start:end])) start = end - video = output_stream.finish() + tail = model_session.finish_output() + if tail is not None: + output_collector.add(tail) + model_session.close() + video = output_collector.finish() if video is None: return generated_num_frames = video.shape[2] @@ -427,7 +439,7 @@ def _rollout_and_save( ) video_path = runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4") - video_path = output_stream.write_mp4( + video_path = write_video_tensor( canvas, video_path, fps=cfg.output_fps, @@ -440,9 +452,11 @@ def _rollout_and_save( f"-> {video_path.resolve()}" ) - if output_stream.stats_history: + if output_collector.stats_history: stats_path = write_runner_stats( - cfg.output_dir, cfg.runner_name, output_stream.stats_history + cfg.output_dir, + cfg.runner_name, + output_collector.stats_history, ) logger.info( f"[{cfg.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" diff --git a/integrations/omnidreams/omnidreams/scenes.py b/integrations/omnidreams/omnidreams/scenes.py index 49623d512..b4fc1c6b2 100644 --- a/integrations/omnidreams/omnidreams/scenes.py +++ b/integrations/omnidreams/omnidreams/scenes.py @@ -1,23 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Shared metadata + helpers for the ``omni-dreams-scenes`` HF dataset. - -Keeps the desktop ``interactive_drive`` demo (which uses the USDZ archive -intact) and ``webrtc.session`` (which extracts it) in lock-step on scene -naming, the HF org resolver, the variant-suffix parser, and the shared -on-disk cache layout under :data:`FLASHDREAMS_CACHE_DIR`/``omnidreams-scenes/``. -The archive (``/clipgt-.usdz``) and extracted dir -(``//``) coexist without name conflict. +"""Shared discovery and staging helpers for the ``omni-dreams-scenes`` dataset. + +The desktop demo consumes USDZ archives intact, while realtime demos extract +them into a normalized ClipGT layout. Both paths share scene naming, variant +selection, Hugging Face lookup, and the cache rooted at +``FLASHDREAMS_CACHE_DIR/omnidreams-scenes``. """ from __future__ import annotations import os import re -from pathlib import Path +import shutil +import tempfile +import zipfile +from collections.abc import Set as AbstractSet +from pathlib import Path, PurePosixPath from typing import Final +from filelock import FileLock +from loguru import logger from omnidreams.hf_org import hf_repo # First-frame image suffixes; both demo paths lowercase before comparison. @@ -25,11 +29,11 @@ {".bmp", ".jpeg", ".jpg", ".png", ".webp"} ) -# Per-scene prompt filename. interactive-drive also supports ``prompt_.txt`` -# variants (via ``variant_from_stem``); webrtc uses only this canonical name. +# Per-scene prompt filename. Interactive Drive also supports ``prompt_.txt`` +# variants through ``variant_from_stem``. SCENE_PROMPT_FILENAME: Final[str] = "prompt.txt" -# Subdir webrtc unpacks a USDZ payload into (``//clipgt/``). +# Subdirectory used for extracted USDZ payloads. SCENE_CLIPGT_DIRNAME: Final[str] = "clipgt" # Per-camera ground-truth frames live at ``frames//.jpeg``; @@ -94,7 +98,7 @@ def normalise_scene_uuid(scene_uuid: str) -> str: return parse_scene_stem(scene_uuid)[0] -def _variant_suffix(variant: str | None) -> str: +def scene_variant_suffix(variant: str | None) -> str: """Filename suffix for ``variant`` (``""`` for the default/base archive).""" slug = (variant or SCENE_VARIANT_DEFAULT).strip() return "" if slug in ("", SCENE_VARIANT_DEFAULT) else f"-{slug}" @@ -108,7 +112,10 @@ def scene_archive_filename( ``variant`` selects a weather sibling (``-rain`` / ``-snow``); the default maps to the base ``scenes/clipgt-.usdz``. """ - return f"scenes/clipgt-{normalise_scene_uuid(scene_uuid)}{_variant_suffix(variant)}.usdz" + return ( + f"scenes/clipgt-{normalise_scene_uuid(scene_uuid)}" + f"{scene_variant_suffix(variant)}.usdz" + ) def prompt_variant_for_scene_variant(variant: str) -> str: @@ -133,7 +140,9 @@ def resolve_variant_archive(scene_path: Path, variant: str) -> Path: """ scene_path = Path(scene_path) uuid, _current = parse_scene_stem(scene_path.stem) - candidate = scene_path.with_name(f"clipgt-{uuid}{_variant_suffix(variant)}.usdz") + candidate = scene_path.with_name( + f"clipgt-{uuid}{scene_variant_suffix(variant)}.usdz" + ) if candidate != scene_path and candidate.exists(): return candidate return scene_path @@ -161,7 +170,7 @@ def local_scene_archive_path( """ return ( scenes_cache_root() - / f"clipgt-{normalise_scene_uuid(scene_uuid)}{_variant_suffix(variant)}.usdz" + / f"clipgt-{normalise_scene_uuid(scene_uuid)}{scene_variant_suffix(variant)}.usdz" ) @@ -248,3 +257,308 @@ def hf_hub_download_scene( filename=scene_archive_filename(scene_uuid, variant), ) return Path(cached) + + +def _choose_existing_asset( + directory: Path, + *, + exact_name: str | None = None, + fallback_stems: tuple[str, ...] = (), + fallback_prefixes: tuple[str, ...] = (), + allowed_suffixes: AbstractSet[str] | None = None, + preferred_stems: tuple[str, ...] = (), +) -> Path | None: + if not directory.is_dir(): + return None + + if exact_name is not None: + exact_path = directory / exact_name + if exact_path.is_file() and ( + allowed_suffixes is None or exact_path.suffix.lower() in allowed_suffixes + ): + return exact_path + + candidates = [] + for path in directory.iterdir(): + if not path.is_file(): + continue + if allowed_suffixes is not None and path.suffix.lower() not in allowed_suffixes: + continue + if ( + path.stem in preferred_stems + or path.stem in fallback_stems + or any(path.stem.startswith(f"{prefix}-") for prefix in fallback_prefixes) + ): + candidates.append(path) + + if not candidates: + return None + + preferred_order = {stem: index for index, stem in enumerate(preferred_stems)} + return sorted( + candidates, + key=lambda path: ( + preferred_order.get(path.stem, len(preferred_order)), + path.name, + ), + )[0] + + +def _camera_name_candidates(camera_name: str) -> tuple[str, ...]: + underscore = camera_name.replace(":", "_") + colon = camera_name.replace("_", ":") + return tuple(dict.fromkeys((camera_name, underscore, colon))) + + +def _first_frame_sort_key(path: Path) -> tuple[int, str]: + stem = path.stem + return (int(stem), path.name) if stem.isdigit() else (2**63 - 1, path.name) + + +def _resolve_first_frame(clipgt_dir: Path, camera_name: str) -> Path | None: + frames_root = clipgt_dir / SCENE_FRAMES_DIRNAME + if not frames_root.is_dir(): + return None + candidate_dirs = [ + frames_root / name + for name in _camera_name_candidates(camera_name) + if (frames_root / name).is_dir() + ] + if not candidate_dirs: + candidate_dirs = [ + path for path in sorted(frames_root.iterdir()) if path.is_dir() + ] + for directory in candidate_dirs: + frames = [ + path + for path in directory.iterdir() + if path.is_file() and path.suffix.lower() in SCENE_FRAME_SUFFIXES + ] + if frames: + return sorted(frames, key=_first_frame_sort_key)[0] + return None + + +def resolve_scene_assets( + scene_dir: Path, + *, + prompt_filename: str, + clipgt_dirname: str, + camera_name: str = "camera_front_wide_120fov", + variant: str = SCENE_VARIANT_DEFAULT, +) -> tuple[Path, Path, Path]: + """Resolve the ClipGT root, first frame, and prompt for a scene.""" + missing_assets = [] + clipgt_dir = scene_dir / clipgt_dirname + if not clipgt_dir.is_dir(): + missing_assets.append(str(clipgt_dir)) + resolved_clipgt_dir = None + else: + resolved_clipgt_dir = clipgt_dir + + first_frame_path = ( + None + if resolved_clipgt_dir is None + else _resolve_first_frame(resolved_clipgt_dir, camera_name) + ) + if first_frame_path is None and resolved_clipgt_dir is not None: + first_frame_path = _choose_existing_asset( + resolved_clipgt_dir, + fallback_stems=("first_image_1",), + allowed_suffixes=SCENE_IMAGE_SUFFIXES, + preferred_stems=("first_image",), + ) + if first_frame_path is None: + missing_assets.append( + f"frames//*.jpeg or first_image.* under {resolved_clipgt_dir}/" + ) + + weather_prompt_stem = f"prompt{prompt_variant_for_scene_variant(variant)}" + prompt_path = ( + None + if resolved_clipgt_dir is None + else _choose_existing_asset( + resolved_clipgt_dir, + fallback_stems=("prompt1", "prompt2", "prompt3", "prompt"), + allowed_suffixes={".txt"}, + preferred_stems=(weather_prompt_stem, "prompt"), + ) + ) + if prompt_path is None: + missing_assets.append(f"{prompt_filename} under {resolved_clipgt_dir}/") + + if missing_assets: + raise FileNotFoundError( + "Missing Omnidreams scene assets: " + ", ".join(missing_assets) + ) + + assert resolved_clipgt_dir is not None + assert first_frame_path is not None + assert prompt_path is not None + return resolved_clipgt_dir, first_frame_path, prompt_path + + +def _safe_extract_zip(source: Path, destination: Path) -> None: + if destination.exists(): + if destination.is_file() or destination.is_symlink(): + destination.unlink() + else: + shutil.rmtree(destination) + destination.mkdir(parents=True, exist_ok=True) + destination_root = destination.resolve() + with zipfile.ZipFile(source) as zf: + for member in zf.infolist(): + member_path = PurePosixPath(member.filename) + if ( + member_path.is_absolute() + or not member_path.parts + or any(part in {"", ".", ".."} for part in member_path.parts) + ): + raise ValueError( + f"Unsafe archive member in {source}: {member.filename}" + ) + target = destination / Path(*member_path.parts) + target_resolved = target.resolve() + if destination_root != target_resolved and destination_root not in ( + target_resolved.parents + ): + raise ValueError( + f"Archive member escapes destination: {member.filename}" + ) + if member.is_dir(): + target.mkdir(parents=True, exist_ok=True) + continue + target.parent.mkdir(parents=True, exist_ok=True) + with zf.open(member) as src, target.open("wb") as dst: + shutil.copyfileobj(src, dst) + + +def extract_local_scene( + scene_dir: Path, + *, + scene_uuid: str | None, + variant: str = SCENE_VARIANT_DEFAULT, + clipgt_dirname: str, +) -> Path: + """Extract a local scene archive into the normalized scene layout.""" + if scene_uuid is None: + return scene_dir + + scene_uuid = scene_uuid.strip() + assert scene_uuid, "scene_uuid must be non-empty when provided." + if not scene_dir.is_dir(): + raise FileNotFoundError(f"scene_dir does not exist: {scene_dir}") + + suffix = scene_variant_suffix(variant) + expected_names = ( + f"clipgt-{scene_uuid}{suffix}.usdz", + f"{scene_uuid}{suffix}.usdz", + ) + archive_path = _choose_existing_asset(scene_dir, exact_name=expected_names[0]) or ( + _choose_existing_asset(scene_dir, exact_name=expected_names[1]) + ) + if archive_path is None: + archive_path = _choose_existing_asset( + scene_dir, + fallback_prefixes=( + f"clipgt-{scene_uuid}{suffix}", + f"{scene_uuid}{suffix}", + f"clipgt-{scene_uuid}", + scene_uuid, + ), + allowed_suffixes={".usdz"}, + preferred_stems=( + f"clipgt-{scene_uuid}{suffix}", + f"{scene_uuid}{suffix}", + f"clipgt-{scene_uuid}", + scene_uuid, + ), + ) + if archive_path is None: + raise FileNotFoundError( + "scene_uuid is set but no local USDZ archive was found in " + f"{scene_dir}. Expected one of: {', '.join(expected_names)}." + ) + + normalized_scene_dir = scene_dir / f"{scene_uuid}{suffix}" + _safe_extract_zip(archive_path, normalized_scene_dir / clipgt_dirname) + return normalized_scene_dir + + +def ensure_hf_scene_synced( + scene_uuid: str, + *, + variant: str = SCENE_VARIANT_DEFAULT, + clipgt_dirname: str = SCENE_CLIPGT_DIRNAME, +) -> Path: + """Download and extract a Hugging Face scene into the shared cache.""" + scene_uuid = scene_uuid.strip() + assert scene_uuid, "scene_uuid must be set." + suffix = scene_variant_suffix(variant) + cache_root = scenes_cache_root() + scene_dir = cache_root / f"{scene_uuid}{suffix}" + lock_path = cache_root / ".locks" / f"{scene_uuid}{suffix}.lock" + lock_path.parent.mkdir(parents=True, exist_ok=True) + + with FileLock(str(lock_path)): + archive_path = hf_hub_download_scene(scene_uuid, variant) + _safe_extract_zip(archive_path, scene_dir / clipgt_dirname) + + logger.info( + "Synced Omnidreams scene {} (variant {}) from Hugging Face ({}) to {}", + scene_uuid, + variant, + hf_scenes_repo_id(), + scene_dir, + ) + return scene_dir + + +def _link_or_copy_file(source: Path, target: Path) -> None: + try: + os.symlink(source, target) + return + except OSError: + pass + + try: + os.link(source, target) + return + except OSError: + shutil.copy2(source, target) + + +def prepare_clipgt_dir( + clipgt_dir: Path, +) -> tuple[Path, tempfile.TemporaryDirectory[str] | None]: + """Normalize supported ClipGT parquet layouts for the scene loader.""" + + def has_prefixed_parquets(path: Path) -> bool: + return any(path.glob("*.calibration_estimate.parquet")) + + def has_unprefixed_parquets(path: Path) -> bool: + return (path / "calibration_estimate.parquet").exists() + + if has_prefixed_parquets(clipgt_dir): + return clipgt_dir, None + + parquet_source_dir: Path | None = None + if has_unprefixed_parquets(clipgt_dir): + parquet_source_dir = clipgt_dir + else: + for candidate in (child for child in clipgt_dir.iterdir() if child.is_dir()): + if has_prefixed_parquets(candidate): + return candidate, None + if has_unprefixed_parquets(candidate): + parquet_source_dir = candidate + break + + if parquet_source_dir is None: + return clipgt_dir, None + + temp_dir = tempfile.TemporaryDirectory(prefix="omnidreams-clipgt-") + staged = Path(temp_dir.name) + for source in parquet_source_dir.glob("*.parquet"): + target = staged / f"clip.{source.name}" + _link_or_copy_file(source.resolve(), target) + return staged, temp_dir diff --git a/integrations/omnidreams/omnidreams/webrtc/__init__.py b/integrations/omnidreams/omnidreams/webrtc/__init__.py deleted file mode 100644 index b777e7889..000000000 --- a/integrations/omnidreams/omnidreams/webrtc/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Single-view Omnidreams WebRTC driving demo.""" diff --git a/integrations/omnidreams/omnidreams/webrtc/server.py b/integrations/omnidreams/omnidreams/webrtc/server.py deleted file mode 100644 index f577dc3df..000000000 --- a/integrations/omnidreams/omnidreams/webrtc/server.py +++ /dev/null @@ -1,377 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import argparse -from dataclasses import replace -from importlib.resources import as_file, files -from pathlib import Path -from typing import Any, Protocol, cast - -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, - OmnidreamsSessionInput, - OmnidreamsWebRTCSessionManager, -) - -from flashdreams.core.distributed import ( - init as distributed_init, -) -from flashdreams.infra.postprocess import VideoPostprocessChainConfig -from flashdreams.plugins.registry import discover_postprocess_presets -from flashdreams.serving.network import get_external_ip -from flashdreams.serving.webrtc.bootstrap import ( - configure_logging, - initialize_cuda_distributed, - run_webrtc_server, -) -from flashdreams.serving.webrtc.server import ( - SESSION_MANAGER_KEY, - SessionBusyError, - WebRTCSessionManager, - create_packaged_webrtc_app, - create_webrtc_app, -) -from flashdreams.serving.webrtc.server import ( - close_package_resources as _close_package_resources, -) - -WEB_DIR_RESOURCE = files("flashdreams.serving.webrtc").joinpath("web") -MODEL_WEB_DIR_RESOURCE = files("omnidreams.webrtc").joinpath("web") - - -class _OmnidreamsSessionManager(WebRTCSessionManager, Protocol): - runtime_config: OmnidreamsRuntimeConfig - - def set_pending_session_input( - self, session_input: OmnidreamsSessionInput - ) -> None: ... - - -def parse_args(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." - ) - ) - parser.add_argument("--host", type=str, default="0.0.0.0") - parser.add_argument("--port", type=int, default=8082) - parser.add_argument( - "--pipeline_config_name", - type=str, - default="omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf", - choices=sorted(OMNIDREAMS_CONFIGS), - ) - parser.add_argument( - "--scene_dir", - type=Path, - default=None, - help=( - "Local WebRTC scene directory containing clipgt/first_image.* " - "and clipgt/prompt.txt. If omitted, the server downloads and " - "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, - default=None, - help=( - "Scene UUID for nvidia/omni-dreams-scenes. Expected dataset asset: " - "scenes/clipgt-[-].usdz." - ), - ) - parser.add_argument( - "--scene-variant", - type=str, - default="default", - help=( - "Weather variant to serve: 'default' (clear), 'rain', or 'snow'. " - "Selects the matching sibling archive and weather prompt." - ), - ) - parser.add_argument("--device", type=str, default="cuda:0") - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--fps", type=int, default=30) - parser.add_argument("--video_height", type=int, default=704) - parser.add_argument("--video_width", type=int, default=1280) - parser.add_argument( - "--warmup_chunks", - type=int, - default=10, - help="Number of synthetic startup chunks to generate for kernel autotuning.", - ) - parser.add_argument( - "--warmup_timeout_s", - type=float, - default=600.0, - help="Maximum seconds to wait for synthetic startup warmup chunks.", - ) - parser.add_argument( - "--debug_serve_hdmaps", - action="store_true", - help=( - "Stream rendered HDMap conditioning frames instead of generated RGB " - "video. This skips video model generation after initialization." - ), - ) - parser.add_argument( - "--camera_name", - type=str, - default="camera_front_wide_120fov", - ) - parser.add_argument( - "--postprocess-preset", - "--postprocess_preset", - dest="postprocess_preset", - default="", - choices=sorted(discover_postprocess_presets()), - help=( - "Video post-process preset for WebRTC sessions. The browser can " - "only toggle this launched preset before connecting." - ), - ) - parser.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." - ), - ) - return parser.parse_args(argv) - - -def _get_omnidreams_manager(app: web.Application) -> _OmnidreamsSessionManager: - return cast(_OmnidreamsSessionManager, app[SESSION_MANAGER_KEY]) - - -async def _postprocess_options(request: web.Request) -> web.StreamResponse: - manager = _get_omnidreams_manager(request.app) - configured_preset = manager.runtime_config.postprocess.preset - presets = [configured_preset] if configured_preset else [] - return web.json_response( - { - "default_preset": configured_preset, - "presets": presets, - } - ) - - -async def _session_input(request: web.Request) -> web.StreamResponse: - try: - payload = await request.json() - except Exception as exc: - raise web.HTTPBadRequest(reason="Expected JSON session input.") from exc - if not isinstance(payload, dict): - raise web.HTTPBadRequest(reason="Session input must be a JSON object.") - preset = payload.get("postprocess_preset") - if not isinstance(preset, str): - raise web.HTTPBadRequest( - reason="Session input must include string 'postprocess_preset'." - ) - - manager = _get_omnidreams_manager(request.app) - try: - manager.set_pending_session_input( - OmnidreamsSessionInput(postprocess_preset=preset) - ) - except SessionBusyError as exc: - raise web.HTTPConflict(reason=str(exc)) from exc - except ValueError as exc: - raise web.HTTPBadRequest(reason=str(exc)) from exc - return web.json_response({"postprocess_preset": preset}) - - -def _configure_app(app: web.Application) -> None: - app.router.add_get("/api/postprocess/options", _postprocess_options) - app.router.add_post("/api/session/input", _session_input) - - -def create_app( - *, - request_session_url: str, - session_manager: WebRTCSessionManager | None = None, -) -> web.Application: - manager = session_manager or OmnidreamsWebRTCSessionManager() - return create_packaged_webrtc_app( - web_resource=WEB_DIR_RESOURCE, - model_web_resource=MODEL_WEB_DIR_RESOURCE, - session_manager=manager, - preload_name="Omnidreams", - request_session_url=request_session_url, - configure_app=_configure_app, - as_file_fn=as_file, - create_app_fn=create_webrtc_app, - cleanup_callback=_close_package_resources, - ) - - -def build_runtime_config( - args: argparse.Namespace, - *, - 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=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=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, - postprocess=VideoPostprocessChainConfig(preset=args.postprocess_preset), - encoder_backend="default" if args.prefer_sw_encoder else "auto", - ) - - -def initialize_distributed( - *, - default_device: str | torch.device = "cuda:0", -) -> tuple[torch.device, int, int]: - context = initialize_cuda_distributed( - default_device=default_device, - distributed_init_fn=distributed_init, - configure_logging_fn=configure_logging, - torch_module=torch, - dist_module=dist, - ) - logger.info( - "Rank {} initialized Omnidreams runtime with context_parallel_size {}", - context.world_rank, - context.world_size, - ) - return context.device, context.world_rank, context.world_size - - -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.") - if transformer_cfg.num_views != 1: - raise ValueError( - "Omnidreams WebRTC only serves single-view configs; " - f"{config_name!r} has num_views={transformer_cfg.num_views}." - ) - - -def main() -> None: - configure_logging() - args = parse_args() - 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=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: - external_ip = get_external_ip() - app = create_app( - session_manager=session_manager, - request_session_url=f"http://{external_ip}:{args.port}/request_session", - ) - logger.info("Starting on external IP: {}", external_ip) - run_webrtc_server( - world_rank=world_rank, - session_manager=session_manager, - app=app, - host=args.host, - port=args.port, - ) - - -if __name__ == "__main__": - main() diff --git a/integrations/omnidreams/omnidreams/webrtc/session.py b/integrations/omnidreams/omnidreams/webrtc/session.py deleted file mode 100644 index e45a07f80..000000000 --- a/integrations/omnidreams/omnidreams/webrtc/session.py +++ /dev/null @@ -1,1189 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import asyncio -import os -import shutil -import tempfile -import time -import zipfile -from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass, field -from pathlib import Path, PurePosixPath -from typing import AbstractSet, Any, Callable, TypeVar - -import cv2 -import numpy as np -import torch -import torch.distributed as dist -from filelock import FileLock -from loguru import logger -from omnidreams.conditioning.conditioning_wrapper import ( - AV_POSITIVE_PROMPT, - OmnidreamsConditioningState, - OmnidreamsConditioningWrapper, - TextPrompt, -) -from omnidreams.conditioning.renderer import load_and_attach_ludus_scene -from omnidreams.conditioning.world_scenario.data_loaders import load_scene -from omnidreams.conditioning.world_scenario.settings import SETTINGS -from omnidreams.config import OMNIDREAMS_CONFIGS -from omnidreams.scenes import ( - HF_DATASET_BROWSER_URL, - SCENE_CLIPGT_DIRNAME, - SCENE_FRAME_SUFFIXES, - SCENE_FRAMES_DIRNAME, - SCENE_IMAGE_SUFFIXES, - SCENE_PROMPT_FILENAME, - SCENE_VARIANT_DEFAULT, - hf_hub_download_scene, - hf_scenes_repo_id, - prompt_variant_for_scene_variant, - scenes_cache_root, -) -from omnidreams.transformer import CosmosTransformerConfig - -from flashdreams.core.distributed.rank_orchestration import ( - RankCoordinator, - distributed_op, -) -from flashdreams.infra.postprocess import ( - VideoPostprocessChainConfig, - VideoPostprocessStream, -) -from flashdreams.infra.video_output import VideoOutputStream, VideoStepResult -from flashdreams.plugins.registry import resolve_postprocess_preset -from flashdreams.serving.webrtc.controls import ( - WSAD_SUPPORTED_KEYS, - CameraPoseIntegrator, - PoseSegment, -) -from flashdreams.serving.webrtc.encoders import ( - EncoderBackend, - VideoEncoder, - select_encoder, -) -from flashdreams.serving.webrtc.manager import ( - DEFAULT_CLIENT_LIVENESS_TIMEOUT_S, - BaseWebRTCSessionManager, - ManagedWebRTCSession, - WebRTCControlSignal, -) -from flashdreams.serving.webrtc.server import SessionBusyError - -_T = TypeVar("_T") -# Default scene (clear-weather base archive). Weather siblings are selected -# via OmnidreamsRuntimeConfig.scene_variant / the server's --scene-variant. -DEFAULT_WEBRTC_SCENE_UUID = "0d404ff7-2b66-498c-b047-1ed8cded60d4" -# Back-compat aliases for ``omnidreams.scenes`` constants used by external imports. -WEBRTC_SCENES_HF_BROWSER_URL = HF_DATASET_BROWSER_URL -WEBRTC_SCENE_IMAGE_SUFFIXES = SCENE_IMAGE_SUFFIXES - - -def _resolve_cuda_device(device_spec: str | torch.device) -> torch.device: - """Resolve a device spec, filling in the active CUDA index when unspecified.""" - device = torch.device(device_spec) - if device.type == "cuda" and device.index is None: - device = torch.device( - f"cuda:{torch.cuda.current_device()}" - if torch.cuda.is_available() - else "cuda:0" - ) - return device - - -def _choose_existing_asset( - directory: Path, - *, - exact_name: str | None = None, - fallback_stems: tuple[str, ...] = (), - fallback_prefixes: tuple[str, ...] = (), - allowed_suffixes: AbstractSet[str] | None = None, - preferred_stems: tuple[str, ...] = (), -) -> Path | None: - if not directory.is_dir(): - return None - - if exact_name is not None: - exact_path = directory / exact_name - if exact_path.is_file() and ( - allowed_suffixes is None or exact_path.suffix.lower() in allowed_suffixes - ): - return exact_path - - candidates = [] - for path in directory.iterdir(): - if not path.is_file(): - continue - if allowed_suffixes is not None and path.suffix.lower() not in allowed_suffixes: - continue - if ( - path.stem in preferred_stems - or path.stem in fallback_stems - or any(path.stem.startswith(f"{prefix}-") for prefix in fallback_prefixes) - ): - candidates.append(path) - - if not candidates: - return None - - preferred_order = {stem: index for index, stem in enumerate(preferred_stems)} - return sorted( - candidates, - key=lambda path: ( - preferred_order.get(path.stem, len(preferred_order)), - path.name, - ), - )[0] - - -def _camera_name_candidates(camera_name: str) -> tuple[str, ...]: - """Colon/underscore spellings of ``camera_name`` (dataset uses underscores).""" - underscore = camera_name.replace(":", "_") - colon = camera_name.replace("_", ":") - return tuple(dict.fromkeys((camera_name, underscore, colon))) - - -def _first_frame_sort_key(path: Path) -> tuple[int, str]: - stem = path.stem - return (int(stem), path.name) if stem.isdigit() else (2**63 - 1, path.name) - - -def _resolve_webrtc_first_frame(clipgt_dir: Path, camera_name: str) -> Path | None: - """Earliest GT frame under ``clipgt/frames//``, else ``None``. - - ``None`` when the bundle ships no such frames, so the caller can fall back - to ``first_image.*``. - """ - frames_root = clipgt_dir / SCENE_FRAMES_DIRNAME - if not frames_root.is_dir(): - return None - candidate_dirs = [ - frames_root / name - for name in _camera_name_candidates(camera_name) - if (frames_root / name).is_dir() - ] - if not candidate_dirs: - # Fall back to any single camera directory present. - candidate_dirs = [ - path for path in sorted(frames_root.iterdir()) if path.is_dir() - ] - for directory in candidate_dirs: - frames = [ - path - for path in directory.iterdir() - if path.is_file() and path.suffix.lower() in SCENE_FRAME_SUFFIXES - ] - if frames: - return sorted(frames, key=_first_frame_sort_key)[0] - return None - - -def _resolve_webrtc_scene_assets( - scene_dir: Path, - *, - prompt_filename: str, - clipgt_dirname: str, - camera_name: str = "camera_front_wide_120fov", - variant: str = SCENE_VARIANT_DEFAULT, -) -> tuple[Path, Path, Path]: - missing_assets = [] - clipgt_dir = scene_dir / clipgt_dirname - if not clipgt_dir.is_dir(): - missing_assets.append(str(scene_dir / clipgt_dirname)) - clipgt_dir = None - - # Prefer the GT camera frame; fall back to ``first_image.*`` for bundles - # with no per-camera frames. - first_frame_path = ( - None - if clipgt_dir is None - else _resolve_webrtc_first_frame(clipgt_dir, camera_name) - ) - if first_frame_path is None and clipgt_dir is not None: - first_frame_path = _choose_existing_asset( - clipgt_dir, - fallback_stems=("first_image_1",), - allowed_suffixes=WEBRTC_SCENE_IMAGE_SUFFIXES, - preferred_stems=("first_image",), - ) - if first_frame_path is None: - missing_assets.append( - f"frames//*.jpeg or first_image.* under {clipgt_dir}/" - ) - - # Prompt matching the weather variant (``promptN.txt``); fall back to a - # bare ``prompt.txt`` for older bundles. - weather_prompt_stem = f"prompt{prompt_variant_for_scene_variant(variant)}" - prompt_path = ( - None - if clipgt_dir is None - else _choose_existing_asset( - clipgt_dir, - fallback_stems=("prompt1", "prompt2", "prompt3", "prompt"), - allowed_suffixes={".txt"}, - preferred_stems=(weather_prompt_stem, "prompt"), - ) - ) - if prompt_path is None: - missing_assets.append(f"{prompt_filename} under {clipgt_dir}/") - - if missing_assets: - raise FileNotFoundError( - "Missing Omnidreams WebRTC scene assets: " + ", ".join(missing_assets) - ) - - assert clipgt_dir is not None - assert first_frame_path is not None - assert prompt_path is not None - return clipgt_dir, first_frame_path, prompt_path - - -def _safe_extract_zip(source: Path, destination: Path) -> None: - if destination.exists(): - if destination.is_file() or destination.is_symlink(): - destination.unlink() - else: - shutil.rmtree(destination) - destination.mkdir(parents=True, exist_ok=True) - destination_root = destination.resolve() - with zipfile.ZipFile(source) as zf: - for member in zf.infolist(): - member_path = PurePosixPath(member.filename) - if ( - member_path.is_absolute() - or not member_path.parts - or any(part in {"", ".", ".."} for part in member_path.parts) - ): - raise ValueError( - f"Unsafe archive member in {source}: {member.filename}" - ) - target = destination / Path(*member_path.parts) - target_resolved = target.resolve() - if destination_root != target_resolved and destination_root not in ( - target_resolved.parents - ): - raise ValueError( - f"Archive member escapes destination: {member.filename}" - ) - if member.is_dir(): - target.mkdir(parents=True, exist_ok=True) - continue - target.parent.mkdir(parents=True, exist_ok=True) - with zf.open(member) as src, target.open("wb") as dst: - shutil.copyfileobj(src, dst) - - -def _variant_dir_suffix(variant: str | None) -> str: - """Cache subdir / filename suffix for ``variant`` (``""`` for default).""" - slug = (variant or SCENE_VARIANT_DEFAULT).strip() - return "" if slug in ("", SCENE_VARIANT_DEFAULT) else f"-{slug}" - - -def _extract_local_webrtc_scene_if_needed( - scene_dir: Path, - *, - scene_uuid: str | None, - variant: str = SCENE_VARIANT_DEFAULT, - clipgt_dirname: str, -) -> Path: - """Extract the ``scene_uuid`` (+ variant) archive into the local layout.""" - if scene_uuid is None: - return scene_dir - - scene_uuid = scene_uuid.strip() - assert scene_uuid, "scene_uuid must be non-empty when provided." - if not scene_dir.is_dir(): - raise FileNotFoundError(f"scene_dir does not exist: {scene_dir}") - - suffix = _variant_dir_suffix(variant) - expected_names = ( - f"clipgt-{scene_uuid}{suffix}.usdz", - f"{scene_uuid}{suffix}.usdz", - ) - archive_path = _choose_existing_asset(scene_dir, exact_name=expected_names[0]) or ( - _choose_existing_asset(scene_dir, exact_name=expected_names[1]) - ) - if archive_path is None: - # Prefer the variant suffix but accept the base archive too. - archive_path = _choose_existing_asset( - scene_dir, - fallback_prefixes=( - f"clipgt-{scene_uuid}{suffix}", - f"{scene_uuid}{suffix}", - f"clipgt-{scene_uuid}", - scene_uuid, - ), - allowed_suffixes={".usdz"}, - preferred_stems=( - f"clipgt-{scene_uuid}{suffix}", - f"{scene_uuid}{suffix}", - f"clipgt-{scene_uuid}", - scene_uuid, - ), - ) - if archive_path is None: - raise FileNotFoundError( - "scene_uuid is set but no local USDZ archive was found in " - f"{scene_dir}. Expected one of: {', '.join(expected_names)}." - ) - - normalized_scene_dir = scene_dir / f"{scene_uuid}{suffix}" - normalized_clipgt_root = normalized_scene_dir / clipgt_dirname - _safe_extract_zip(archive_path, normalized_clipgt_root) - return normalized_scene_dir - - -def _ensure_hf_webrtc_scene_synced( - scene_uuid: str, - *, - variant: str = SCENE_VARIANT_DEFAULT, - prompt_filename: str = SCENE_PROMPT_FILENAME, - clipgt_dirname: str = SCENE_CLIPGT_DIRNAME, -) -> Path: - """Stage an HF scene variant into the WebRTC cache layout. - - Downloads ``scenes/clipgt-[-].usdz`` and extracts it under - ``FLASHDREAMS_CACHE_DIR/omnidreams-scenes/[-]/clipgt/``. The - per-uuid+variant directory coexists with the desktop demo's archive files - in the same root. - """ - del prompt_filename # accepted for call-site symmetry; assets resolved later - scene_uuid = scene_uuid.strip() - assert scene_uuid, "scene_uuid must be set." - suffix = _variant_dir_suffix(variant) - cache_root = scenes_cache_root() - scene_dir = cache_root / f"{scene_uuid}{suffix}" - lock_path = cache_root / ".locks" / f"{scene_uuid}{suffix}.lock" - lock_path.parent.mkdir(parents=True, exist_ok=True) - - with FileLock(str(lock_path)): - archive_path = hf_hub_download_scene(scene_uuid, variant) - _safe_extract_zip(archive_path, scene_dir / clipgt_dirname) - - logger.info( - "Synced Omnidreams WebRTC scene {} (variant {}) from Hugging Face ({}) to {}", - scene_uuid, - variant, - hf_scenes_repo_id(), - scene_dir, - ) - return scene_dir - - -def _summarize_sdp_candidates(sdp: str) -> str: - candidates = [ - line.removeprefix("a=candidate:") - for line in sdp.splitlines() - if line.startswith("a=candidate:") - ] - if not candidates: - return "0 candidates" - - protocols: dict[str, int] = {} - addresses: set[str] = set() - endpoints: list[str] = [] - for candidate in candidates: - parts = candidate.split() - if len(parts) >= 5: - protocols[parts[2].lower()] = protocols.get(parts[2].lower(), 0) + 1 - addresses.add(parts[4]) - if len(parts) >= 6: - endpoints.append(f"{parts[2].lower()}://{parts[4]}:{parts[5]}") - protocol_summary = ",".join( - f"{key}={value}" for key, value in sorted(protocols.items()) - ) - address_summary = ",".join(sorted(addresses)[:8]) - if len(addresses) > 8: - address_summary += f",+{len(addresses) - 8} more" - endpoint_summary = ",".join(endpoints[:12]) - if len(endpoints) > 12: - endpoint_summary += f",+{len(endpoints) - 12} more" - return ( - f"{len(candidates)} candidates protocols=[{protocol_summary}] " - f"addresses=[{address_summary}] endpoints=[{endpoint_summary}]" - ) - - -def _link_or_copy_file(source: Path, target: Path) -> None: - """Stage a file efficiently without requiring Windows symlink privileges.""" - try: - os.symlink(source, target) - return - except OSError: - pass - - try: - os.link(source, target) - return - except OSError: - shutil.copy2(source, target) - - -class OmnidreamsRuntimeError(RuntimeError): - """Raised when the Omnidreams WebRTC runtime is used incorrectly.""" - - -@dataclass(slots=True) -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. - scene_variant: str = SCENE_VARIANT_DEFAULT - seed: int | None = 42 - device: str = "cuda:0" - video_height: int = 704 - video_width: int = 1280 - fps: int = 30 - camera_name: str = "camera_front_wide_120fov" - prompt_filename: str = SCENE_PROMPT_FILENAME - clipgt_dirname: str = SCENE_CLIPGT_DIRNAME - move_speed_per_s: float = 6.0 - rotate_speed_rad_per_s: float = float(np.deg2rad(35.0)) - warmup_chunks: int = 10 - warmup_timeout_s: float = 600.0 - debug_serve_hdmaps: bool = False - postprocess: VideoPostprocessChainConfig = field( - default_factory=VideoPostprocessChainConfig - ) - # 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 - - -@dataclass(frozen=True, slots=True) -class OmnidreamsSessionInput: - """Browser-selectable settings applied to the next WebRTC rollout.""" - - postprocess_preset: str | None = None - """Launched preset selection; ``None`` keeps the CLI default and ``""`` disables it.""" - - -def _validate_requested_postprocess_preset( - *, requested_preset: str, configured_preset: str -) -> None: - if not configured_preset: - raise ValueError( - "Post-processing is not enabled for this server; restart with " - "--postprocess-preset to make a preset available." - ) - if requested_preset != configured_preset: - raise ValueError( - "Post-processing preset must match the launched preset " - f"{configured_preset!r}; got {requested_preset!r}." - ) - resolve_postprocess_preset(requested_preset) - - -class OmnidreamsInferenceRuntime: - """Single-scene, single-view Omnidreams runtime for WebRTC control.""" - - def __init__(self, config: OmnidreamsRuntimeConfig | None = None) -> None: - self.config = config or OmnidreamsRuntimeConfig() - self.MASTER_RANK = 0 - self.rank = 0 if not dist.is_initialized() else dist.get_rank() - - control_device = _resolve_cuda_device(self.config.device) - - self.pose_integrator = CameraPoseIntegrator( - move_speed_per_s=self.config.move_speed_per_s, - rotate_speed_rad_per_s=self.config.rotate_speed_rad_per_s, - coordinate_system="FLU", - ) - self.autoregressive_index = 0 - - self._device: torch.device | None = None - self._wrapper: OmnidreamsConditioningWrapper | None = None - self._state: OmnidreamsConditioningState | None = None - self._renderer: Any | None = None - self._scene_data: Any | None = None - self._initial_rgb_frames: torch.Tensor | 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._next_timestamp_us: int = 0 - self._output_stream = self._new_output_stream(postprocess_stream=None) - self._postprocess_preset = self.config.postprocess.preset - 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. - self._executor = ThreadPoolExecutor( - max_workers=1, - thread_name_prefix="omnidreams-webrtc-runtime", - ) - - self._step_lock = asyncio.Lock() - self.rank_coordinator = RankCoordinator( - device=control_device, - signal_type=WebRTCControlSignal, - is_master=self.is_master, - master_rank=self.MASTER_RANK, - ) - self.rank_coordinator.register_distributed_ops(self) - - @property - def is_master(self) -> bool: - return self.rank == self.MASTER_RANK - - @property - def postprocess_preset(self) -> str: - """Preset active for the current rollout, or an empty string when off.""" - return self._postprocess_preset - - @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) - - def send_exit_signal(self) -> None: - if self.is_master: - self.rank_coordinator.send_exit(exit_signal=WebRTCControlSignal.EXIT) - - async def initialize(self) -> None: - if self._wrapper is not None: - return - await self._run_on_runtime_thread(self._initialize_sync_all_ranks) - - async def reset_for_new_session( - self, session_input: OmnidreamsSessionInput | None = None - ) -> None: - if self._closed: - raise OmnidreamsRuntimeError("Runtime is closed.") - if self._wrapper is None: - raise OmnidreamsRuntimeError("Runtime is not initialized.") - await self._run_on_runtime_thread( - self._reset_rollout_sync_all_ranks, - session_input, - ) - - async def close(self) -> None: - self._closed = True - try: - await self._run_on_runtime_thread(self._close_sync_all_ranks) - finally: - self._executor.shutdown(wait=False, cancel_futures=True) - - async def generate_chunk( - self, - *, - segments: list[PoseSegment], - frame_times: list[float], - ) -> VideoStepResult: - if self._closed: - raise OmnidreamsRuntimeError("Session is closed.") - if self._wrapper is None: - raise OmnidreamsRuntimeError("Runtime is not initialized.") - - async with self._step_lock: - if self._closed: - raise OmnidreamsRuntimeError("Session is closed.") - return await self._run_on_runtime_thread( - self._generate_chunk_sync_all_ranks, - segments, - frame_times, - ) - - async def _run_on_runtime_thread( - self, - func: Callable[..., _T], - *args: Any, - ) -> _T: - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - self._runtime_thread_entry, - func, - args, - ) - - def _runtime_thread_entry( - self, - func: Callable[..., _T], - args: tuple[Any, ...], - ) -> _T: - device = self._device - if device is None: - device = _resolve_cuda_device(self.config.device) - if device.type == "cuda": - torch.cuda.set_device(device) - return func(*args) - - def peek_next_chunk_num_frames(self) -> int: - if self._wrapper is None: - raise OmnidreamsRuntimeError("Runtime is not initialized.") - if self._state is None: - return int(self._wrapper.initial_frame_chunk_size) - return int(self._wrapper.frame_chunk_size) - - def peek_steady_chunk_num_frames(self) -> int: - if self._wrapper is None: - raise OmnidreamsRuntimeError("Runtime is not initialized.") - return int(self._wrapper.frame_chunk_size) - - @distributed_op(WebRTCControlSignal.INITIALIZE) - def _initialize_sync_all_ranks(self) -> None: - self._initialize_sync() - - @distributed_op(WebRTCControlSignal.RESET_SESSION) - def _reset_rollout_sync_all_ranks( - self, session_input: OmnidreamsSessionInput | None = None - ) -> None: - self._reset_rollout_sync(session_input=session_input) - - @distributed_op(WebRTCControlSignal.ACTION_STEP) - def _generate_chunk_sync_all_ranks( - self, - segments: list[PoseSegment], - frame_times: list[float], - ) -> VideoStepResult: - return self._generate_one_chunk_sync(segments=segments, frame_times=frame_times) - - @distributed_op(WebRTCControlSignal.CLOSE) - def _close_sync_all_ranks(self) -> None: - self._close_sync() - - def _initialize_sync(self) -> None: - if self._wrapper is not None: - return - - init_t0 = time.perf_counter() - cfg = self.config - if cfg.scene_dir is None: - scene_uuid = cfg.scene_uuid or DEFAULT_WEBRTC_SCENE_UUID - scene_dir = _ensure_hf_webrtc_scene_synced( - scene_uuid, - variant=cfg.scene_variant, - prompt_filename=cfg.prompt_filename, - clipgt_dirname=cfg.clipgt_dirname, - ) - else: - scene_dir = _extract_local_webrtc_scene_if_needed( - cfg.scene_dir, - scene_uuid=cfg.scene_uuid, - variant=cfg.scene_variant, - clipgt_dirname=cfg.clipgt_dirname, - ) - - cfg.scene_dir = scene_dir - clipgt_dir, first_frame_path, prompt_path = _resolve_webrtc_scene_assets( - scene_dir, - prompt_filename=cfg.prompt_filename, - clipgt_dirname=cfg.clipgt_dirname, - camera_name=cfg.camera_name, - variant=cfg.scene_variant, - ) - 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 = ( - 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( - "Omnidreams WebRTC requires a CosmosTransformerConfig pipeline." - ) - if transformer_cfg.num_views != 1: - raise ValueError( - "Omnidreams WebRTC v1 only supports single-view configs; " - f"{cfg.pipeline_config_name!r} has num_views={transformer_cfg.num_views}." - ) - - self._device = torch.device(cfg.device) - if self._device.type == "cuda" and not torch.cuda.is_available(): - raise RuntimeError("CUDA is required for Omnidreams WebRTC runtime.") - - logger.info("Loading Omnidreams first frame from {}", first_frame_path) - image_bgr = cv2.imread(str(first_frame_path), cv2.IMREAD_COLOR) - if image_bgr is None: - raise RuntimeError(f"Failed to read first frame from {first_frame_path}") - image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) - image_rgb = cv2.resize( - image_rgb, - (cfg.video_width, cfg.video_height), - interpolation=cv2.INTER_CUBIC, - ) - self._initial_rgb_frames = ( - torch.from_numpy(image_rgb) - .permute(2, 0, 1) - .contiguous() - .unsqueeze(0) - .unsqueeze(0) - .to(device=self._device, dtype=torch.uint8) - ) - - prompt = prompt_path.read_text(encoding="utf-8").strip() or AV_POSITIVE_PROMPT - self._text_prompts = [TextPrompt(positive=prompt)] - - loadable_clipgt_dir = self._prepare_clipgt_dir(clipgt_dir) - logger.info("Loading Omnidreams scene data from {}", loadable_clipgt_dir) - scene_t0 = time.perf_counter() - scene_data = load_scene( - loadable_clipgt_dir, - camera_names=[cfg.camera_name], - max_frames=-1, - input_pose_fps=SETTINGS["INPUT_POSE_FPS"], - resize_resolution_hw=(cfg.video_height, cfg.video_width), - ) - logger.info( - "Loaded Omnidreams scene data in {:.1f}s; attaching Ludus scene.", - time.perf_counter() - scene_t0, - ) - ludus_t0 = time.perf_counter() - scene_data = load_and_attach_ludus_scene( - loadable_clipgt_dir, - scene_data, - device=self._device, - ) - logger.info( - "Attached Omnidreams Ludus scene in {:.1f}s.", - time.perf_counter() - ludus_t0, - ) - if not scene_data.ego_poses: - raise ValueError(f"Scene {loadable_clipgt_dir} has no ego poses.") - if cfg.camera_name not in scene_data.camera_models: - raise ValueError( - f"Camera {cfg.camera_name!r} was not loaded from {loadable_clipgt_dir}." - ) - if cfg.camera_name not in scene_data.camera_extrinsics: - raise ValueError( - f"Camera {cfg.camera_name!r} has no extrinsics in {loadable_clipgt_dir}." - ) - - logger.info( - "Setting up Omnidreams pipeline {} on {}. This may load checkpoints, " - "compile modules, and initialize CUDA graphs.", - cfg.pipeline_config_name, - self._device, - ) - 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, - ) - logger.info( - "Omnidreams pipeline setup complete in {:.1f}s.", - time.perf_counter() - pipeline_t0, - ) - self._scene_data = scene_data - logger.info("Creating Omnidreams renderer for camera {}", cfg.camera_name) - renderer_t0 = time.perf_counter() - self._renderer = self._wrapper.create_renderer(scene_data, [cfg.camera_name]) - logger.info( - "Omnidreams renderer ready in {:.1f}s.", - time.perf_counter() - renderer_t0, - ) - self._camera_to_rig = torch.as_tensor( - scene_data.camera_extrinsics[cfg.camera_name], - device=self._device, - dtype=torch.float32, - ) - 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. - - Non-master ranks skip encoder initialization. WebRTC media is - served only by the master rank, so allocating an NVENC session - on a worker would consume one of the local GPU's concurrent - session slots without ever encoding a frame — and could fail - the worker's startup if the pool cannot accommodate one - allocation per rank. - """ - if not self.is_master: - return - 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")) - - def _has_unprefixed_parquets(path: Path) -> bool: - return (path / "calibration_estimate.parquet").exists() - - if _has_prefixed_parquets(clipgt_dir): - return clipgt_dir - - parquet_source_dir: Path | None = None - if _has_unprefixed_parquets(clipgt_dir): - parquet_source_dir = clipgt_dir - else: - # Some HF scenes extract into ``clipgt/clipgt`` (or another single - # nested directory) while first_image/prompt stay one level up. - # Discover that nested parquet root and normalize it for loader use. - nested_candidates = [ - child for child in clipgt_dir.iterdir() if child.is_dir() - ] - for candidate in nested_candidates: - if _has_prefixed_parquets(candidate): - return candidate - if _has_unprefixed_parquets(candidate): - parquet_source_dir = candidate - break - - if parquet_source_dir is None: - return clipgt_dir - - self._clipgt_temp_dir = tempfile.TemporaryDirectory(prefix="omnidreams-clipgt-") - staged = Path(self._clipgt_temp_dir.name) - for source in parquet_source_dir.glob("*.parquet"): - target = staged / f"clip.{source.name}" - _link_or_copy_file(source.resolve(), target) - return staged - - def _reset_rollout_sync( - self, session_input: OmnidreamsSessionInput | None = None - ) -> None: - if self._wrapper is None or self._renderer is None: - raise OmnidreamsRuntimeError("Runtime is not initialized.") - if self._initial_ego_pose is None or self._scene_data is None: - raise OmnidreamsRuntimeError("Scene state is not initialized.") - - self._reset_postprocess_stream(session_input) - if self._state is not None and self._state.pipeline_cache is not None: - del self._state.pipeline_cache - self._state = None - self.pose_integrator = CameraPoseIntegrator( - move_speed_per_s=self.config.move_speed_per_s, - rotate_speed_rad_per_s=self.config.rotate_speed_rad_per_s, - coordinate_system="FLU", - ) - self.pose_integrator.reset(self._initial_ego_pose) - self.autoregressive_index = 0 - self._next_timestamp_us = int(self._scene_data.ego_poses[0].timestamp) - self._wrapper.set_rollout_seed(self.config.seed) - - def _close_sync(self) -> None: - state = self._state - wrapper = self._wrapper - self._state = None - self._wrapper = None - self._renderer = None - self._scene_data = None - self._initial_rgb_frames = None - self._text_prompts = None - self._camera_to_rig = None - self._initial_ego_pose = None - self._close_postprocess_stream() - 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: - del wrapper - if self._clipgt_temp_dir is not None: - self._clipgt_temp_dir.cleanup() - self._clipgt_temp_dir = None - - if self._device is not None and self._device.type == "cuda": - torch.cuda.synchronize(device=self._device) - torch.cuda.empty_cache() - - def _reset_postprocess_stream( - self, session_input: OmnidreamsSessionInput | None - ) -> None: - self._close_postprocess_stream() - configured = self.config.postprocess - preset = ( - session_input.postprocess_preset - if session_input is not None - and session_input.postprocess_preset is not None - else configured.preset - ) - if preset: - _validate_requested_postprocess_preset( - requested_preset=preset, - configured_preset=configured.preset, - ) - postprocess = VideoPostprocessChainConfig( - processors=configured.processors, - preset=preset, - ) - world_size = dist.get_world_size() if dist.is_initialized() else 1 - postprocess.validate_execution(world_size=world_size) - self._postprocess_preset = preset - if not postprocess.is_enabled(): - return - if not self.is_master and not postprocess.requires_all_ranks( - world_size=world_size - ): - return - postprocess_stream = VideoPostprocessStream( - postprocess=postprocess, - output_layout="bvtchw", - fps=self.config.fps, - per_view=False, - world_size=world_size, - ) - self._output_stream = self._new_output_stream( - postprocess_stream=postprocess_stream, - ) - logger.info( - "Omnidreams WebRTC post-processing enabled with preset {!r}.", - preset, - ) - - def _close_postprocess_stream(self) -> None: - self._output_stream.finish() - self._output_stream = self._new_output_stream(postprocess_stream=None) - - @staticmethod - def _new_output_stream( - *, postprocess_stream: VideoPostprocessStream | None - ) -> VideoOutputStream: - return VideoOutputStream( - postprocess_stream=postprocess_stream, - output_layout="bvtchw", - collect_output=False, - move_to_cpu=False, - ) - - def _generate_one_chunk_sync( - self, - *, - segments: list[PoseSegment], - frame_times: list[float], - ) -> VideoStepResult: - if ( - self._wrapper is None - or self._renderer is None - or self._initial_rgb_frames is None - or self._text_prompts is None - or self._camera_to_rig is None - ): - raise OmnidreamsRuntimeError("Runtime is not initialized.") - if self._device is None: - raise OmnidreamsRuntimeError("Runtime device is not initialized.") - - num_frames = self.peek_next_chunk_num_frames() - if len(frame_times) != num_frames: - raise OmnidreamsRuntimeError( - f"Expected {num_frames} frame_times for chunk={self.autoregressive_index}, " - f"got {len(frame_times)}." - ) - if not segments: - raise OmnidreamsRuntimeError( - f"Chunk={self.autoregressive_index} received empty segments." - ) - - ego_poses = self.pose_integrator.integrate_chunk( - segments=segments, frame_times=frame_times - ) - ego_poses_t = torch.from_numpy(ego_poses).to( - device=self._device, dtype=torch.float32 - ) - camera_poses = torch.einsum("nij,jk->nik", ego_poses_t, self._camera_to_rig) - frame_timestamps_us = self._consume_timestamps(num_frames) - - camera_names = [self.config.camera_name] - camera_poses_per_view = {self.config.camera_name: camera_poses} - serve_hdmaps = self.config.debug_serve_hdmaps - if self._state is None: - output = self._wrapper.start_generation( - text_prompts=self._text_prompts, - initial_rgb_frames=self._initial_rgb_frames, - renderer=self._renderer, - camera_names=camera_names, - camera_poses_per_view=camera_poses_per_view, - frame_timestamps_us=frame_timestamps_us, - skip_video_generation=serve_hdmaps, - ) - self._state = output.state - else: - output = self._wrapper.continue_generation( - state=self._state, - camera_names=camera_names, - camera_poses_per_view=camera_poses_per_view, - frame_timestamps_us=frame_timestamps_us, - 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, - ) - - if serve_hdmaps: - video_chunk = output.condition_frames - elif output.rgb_frames is None: - raise OmnidreamsRuntimeError("Omnidreams WebRTC received no RGB frames.") - else: - video_chunk = output.rgb_frames - - if serve_hdmaps: - result = VideoStepResult.from_video_chunk( - chunk_index=self.autoregressive_index, - video_chunk=video_chunk.detach(), - layout="bvtchw", - ) - else: - result = self._output_stream.make_step_result( - video_chunk, - autoregressive_index=self.autoregressive_index, - sync_device=self._device, - ) - self.autoregressive_index += 1 - return result - - def _consume_timestamps(self, num_frames: int) -> list[int]: - step_us = int(round(1_000_000 / self.config.fps)) - timestamps = [self._next_timestamp_us + i * step_us for i in range(num_frames)] - self._next_timestamp_us += num_frames * step_us - return timestamps - - -_ManagedOmnidreamsSession = ManagedWebRTCSession - - -class OmnidreamsWebRTCSessionManager( - BaseWebRTCSessionManager[OmnidreamsInferenceRuntime, OmnidreamsRuntimeConfig] -): - """Owns one active WebRTC session and forwards WSAD actions.""" - - _busy_message = "An Omnidreams session is already active." - _warmup_label = "Omnidreams WebRTC" - _runtime_error_types = (OmnidreamsRuntimeError,) - # A chunk-generation failure here is fatal to the rollout, so tear the - # session down instead of retrying on the next tick. - _close_session_on_generation_error = True - _resampler_supported_keys = WSAD_SUPPORTED_KEYS - - def __init__( - self, - *, - runtime_config: OmnidreamsRuntimeConfig | None = None, - client_liveness_timeout_s: float = DEFAULT_CLIENT_LIVENESS_TIMEOUT_S, - ) -> None: - runtime_config = runtime_config or OmnidreamsRuntimeConfig() - super().__init__( - runtime=OmnidreamsInferenceRuntime(config=runtime_config), - runtime_config=runtime_config, - fps=runtime_config.fps, - client_liveness_timeout_s=client_liveness_timeout_s, - ) - self._pending_session_input: OmnidreamsSessionInput | None = None - - def _model_name(self) -> str: - return self.runtime_config.pipeline_config_name - - def _chunk_done_extra(self) -> dict[str, Any]: - return { - "stream": "hdmap" if self.runtime_config.debug_serve_hdmaps else "rgb", - "postprocess_preset": self._runtime.postprocess_preset, - } - - def _peek_pending_session_input(self) -> OmnidreamsSessionInput | None: - return self._pending_session_input - - def _clear_pending_session_input(self) -> None: - self._pending_session_input = None - - async def _reset_runtime_for_session( - self, session_input: OmnidreamsSessionInput | None - ) -> None: - await self._runtime.reset_for_new_session(session_input=session_input) - - def set_pending_session_input(self, session_input: OmnidreamsSessionInput) -> None: - if self.has_active_session(): - raise SessionBusyError(self._busy_message) - preset = session_input.postprocess_preset - if preset: - _validate_requested_postprocess_preset( - requested_preset=preset, - configured_preset=self.runtime_config.postprocess.preset, - ) - self._pending_session_input = session_input - - def _register_extra_peer_handlers(self, peer_connection: Any) -> None: - @peer_connection.on("iceconnectionstatechange") - def on_iceconnectionstatechange() -> None: - logger.info( - "Peer ICE connection state changed: {}", - peer_connection.iceConnectionState, - ) - - @peer_connection.on("icegatheringstatechange") - def on_icegatheringstatechange() -> None: - logger.debug( - "Peer ICE gathering state changed: {}", - peer_connection.iceGatheringState, - ) - - def _on_offer_received(self, offer_sdp: str) -> None: - logger.info( - "Received WebRTC offer with {}.", _summarize_sdp_candidates(offer_sdp) - ) - - def _on_answer_created(self, answer_sdp: str) -> None: - logger.info( - "Created WebRTC answer with {}.", _summarize_sdp_candidates(answer_sdp) - ) diff --git a/integrations/omnidreams/omnidreams/webrtc/web/adapter.js b/integrations/omnidreams/omnidreams/webrtc/web/adapter.js deleted file mode 100644 index 7ae561fe4..000000000 --- a/integrations/omnidreams/omnidreams/webrtc/web/adapter.js +++ /dev/null @@ -1,7 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -export default { - modelName: "OmniDreams", - enablePostprocess: true, -} diff --git a/integrations/omnidreams/pyproject.toml b/integrations/omnidreams/pyproject.toml index 4e05b745c..221cb042f 100644 --- a/integrations/omnidreams/pyproject.toml +++ b/integrations/omnidreams/pyproject.toml @@ -78,8 +78,8 @@ ludus-renderer = { workspace = true } interactive-drive = [ "slangpy==0.42.0", ] -# Optional NVIDIA VFX runtime for selecting RTX postprocess presets such as -# ``--postprocess-preset rtx-super-resolution`` in the WebRTC or local demo. +# Optional NVIDIA VFX runtime for selecting RTX postprocess presets in the +# local interactive demo. rtx-postprocess = [ "flashdreams[rtx-postprocess]", ] @@ -107,7 +107,7 @@ omnidreams-eval = "omnidreams.eval.cli:main" # Experimental shared demo API path. This coexists with the legacy # WebRTC/gRPC/interactive-drive demos until the new adapter is proven. -omnidreams-demo = "omnidreams.demo.cli:main" +omnidreams-demo = "omnidreams.demo.app:main" # Desktop interactive-drive demo entry point. Requires the # ``interactive-drive`` extra (it adds slangpy); without it the @@ -138,7 +138,7 @@ exclude = ["tests"] # workspace editable. Editable installs pick these up from the source # tree automatically. [tool.setuptools.package-data] -"omnidreams.webrtc.web" = ["adapter.js"] +"omnidreams.demo" = ["web/adapter.js"] "omnidreams.interactive_drive" = [ "configs/*.yaml", "configs/wheels/*.yaml", diff --git a/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py b/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py index 81d77a581..aef7aeb42 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py +++ b/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py @@ -288,6 +288,7 @@ def __init__(self, **kwargs: object) -> None: self.kwargs = kwargs self.calls: list[int] = [] self.finished = False + self.last_process_stats = None streams.append(self) def process( diff --git a/integrations/omnidreams/tests/test_demo_api.py b/integrations/omnidreams/tests/test_demo_api.py index 80928c5cf..1d2f81ef3 100644 --- a/integrations/omnidreams/tests/test_demo_api.py +++ b/integrations/omnidreams/tests/test_demo_api.py @@ -5,10 +5,11 @@ from collections.abc import Sequence from pathlib import Path -from typing import Any, cast +from types import SimpleNamespace +from typing import Any +import omnidreams.demo as demo_package import omnidreams.demo.spec as spec_module -import omnidreams.demo.webrtc as demo_webrtc_module import pytest import torch from aiohttp import web @@ -20,29 +21,32 @@ OmnidreamsReplayScenario, OmnidreamsWebRTCScenario, ) -from omnidreams.demo.cli import _replay_spec, _webrtc_spec, parse_args +from omnidreams.demo.app import _replay_spec, _webrtc_spec, parse_args from omnidreams.demo.replay import ( OmnidreamsReplayRuntime, OmnidreamsReplayRuntimeOptions, ) -from omnidreams.demo.webrtc import OmnidreamsDemoWebRTCSessionManager +from omnidreams.demo.webrtc import ( + OmnidreamsWebRTCModelRuntime, + OmnidreamsWebRTCModelRuntimeConfig, + serve_omnidreams_webrtc_demo, +) -from flashdreams.infra.video_output import VideoStepResult from flashdreams.runtime import ( InferenceConfig, InferenceInput, OutputArtifact, OutputTarget, + StepRequest, StepResult, ) from flashdreams.runtime.demo import ( DemoSpec, Mp4OutputSpec, WebRTCOutputSpec, - serve_flashdreams_demo, ) from flashdreams.runtime.demo.replay import run_replay_demo -from flashdreams.runtime.demo.webrtc import WebRTCDemo, build_webrtc_demo +from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager from flashdreams.serving.webrtc.server import SESSION_MANAGER_KEY pytestmark = pytest.mark.ci_cpu @@ -55,12 +59,19 @@ def test_omnidreams_demo_defaults_to_stable_non_perf_preset() -> None: assert not args.preset_id.endswith("-perf") -def test_omnidreams_demo_adapter_declares_mp4_and_webrtc_modes() -> None: +def test_omnidreams_demo_adapter_declares_replay_modes_only() -> None: adapter = OmnidreamsDemoAdapter() assert adapter.model_id == OMNIDREAMS_MODEL_ID - assert adapter.supported_input_modes() == ("replay", "keyboard-driving") - assert adapter.supported_output_modes() == ("mp4", "webrtc") + assert adapter.supported_input_modes() == ("replay",) + assert adapter.supported_output_modes() == ("mp4",) + + +def test_omnidreams_demo_does_not_import_legacy_webrtc_package() -> None: + demo_dir = Path(demo_package.__file__).parent + + for path in demo_dir.glob("*.py"): + assert "omnidreams.webrtc" not in path.read_text(encoding="utf-8"), path def test_omnidreams_replay_demo_uses_shared_runner(tmp_path: Path) -> None: @@ -251,9 +262,9 @@ def test_omnidreams_replay_runtime_generates_video_step_result( assert result.step_index == 0 assert result.frame_count == 1 - assert isinstance(result.output, VideoStepResult) - assert result.output.layout == "bvtchw" - assert result.output.video_chunk.shape == (1, 1, 1, 3, 2, 2) + assert isinstance(result, StepResult) + assert result.layout == "bvtchw" + assert result.video_chunk.shape == (1, 1, 1, 3, 2, 2) assert result.metrics["denoise_s"] == 0.25 assert session.next_step_request() is None assert pipeline.initialize_cache_calls == [ @@ -331,7 +342,6 @@ def test_omnidreams_webrtc_cli_builds_keyboard_driving_spec(tmp_path: Path) -> N def test_omnidreams_webrtc_demo_uses_shared_manager_with_model_config() -> None: pipeline_config = object() - adapter = OmnidreamsDemoAdapter(webrtc_runtime_factory=_FakeWebRTCRuntime) spec = DemoSpec( model_id=OMNIDREAMS_MODEL_ID, preset_id=DEFAULT_OMNIDREAMS_PRESET, @@ -360,47 +370,55 @@ def test_omnidreams_webrtc_demo_uses_shared_manager_with_model_config() -> None: ), ) - demo = build_webrtc_demo(spec=spec, adapter=adapter) - - assert isinstance(demo.runtime, _FakeWebRTCRuntime) - assert isinstance(demo.session_manager, OmnidreamsDemoWebRTCSessionManager) - assert demo.session_manager._runtime is demo.runtime - assert demo.session_manager.runtime_config is demo.runtime.config - assert demo.runtime_config is demo.runtime.config - assert demo.runtime_config.pipeline_config is pipeline_config - assert demo.runtime_config.pipeline_config_name == DEFAULT_OMNIDREAMS_PRESET - assert demo.runtime_config.scene_uuid == "scene-1" - assert demo.runtime_config.scene_variant == "rain" - assert demo.runtime_config.seed == 123 - assert demo.runtime_config.device == "cuda:7" - assert demo.runtime_config.video_width == 64 - assert demo.runtime_config.video_height == 32 - assert demo.runtime_config.fps == 24 - assert demo.runtime_config.debug_serve_hdmaps is True - assert demo.runtime_config.encoder_backend == "default" - assert demo.session_manager._model_name() == DEFAULT_OMNIDREAMS_PRESET - assert demo.host == "0.0.0.0" - assert demo.port == 8082 - - -def test_omnidreams_webrtc_demo_installs_model_routes( + calls: list[dict[str, Any]] = [] + serve_omnidreams_webrtc_demo( + spec=spec, + world_rank=1, + runtime_factory=_FakeWebRTCRuntime, + server_runner=lambda **kwargs: calls.append(kwargs), + ) + + manager = calls[0]["session_manager"] + runtime = manager._runtime + assert isinstance(runtime, _FakeWebRTCRuntime) + assert type(manager) is BaseWebRTCSessionManager + assert manager.runtime_config is runtime.config + assert runtime.config.pipeline_config is pipeline_config + assert runtime.config.pipeline_config_name == DEFAULT_OMNIDREAMS_PRESET + assert runtime.config.scene_uuid == "scene-1" + assert runtime.config.scene_variant == "rain" + assert runtime.config.seed == 123 + assert runtime.config.device == "cuda:7" + assert runtime.config.video_width == 64 + assert runtime.config.video_height == 32 + assert runtime.config.fps == 24 + assert runtime.config.debug_serve_hdmaps is True + assert runtime.config.encoder_backend == "default" + assert manager.identity == DEFAULT_OMNIDREAMS_PRESET + assert calls[0]["host"] == "0.0.0.0" + assert calls[0]["port"] == 8082 + + +def test_omnidreams_webrtc_demo_installs_model_assets_without_routes( monkeypatch: pytest.MonkeyPatch, ) -> None: + import flashdreams.runtime.demo.webrtc as shared_webrtc_module + app_calls: list[dict[str, Any]] = [] def fake_create_packaged_webrtc_app(**kwargs: Any) -> web.Application: app_calls.append(kwargs) app = web.Application() app[SESSION_MANAGER_KEY] = kwargs["session_manager"] - kwargs["configure_app"](app) + if configure_app := kwargs["configure_app"]: + configure_app(app) return app monkeypatch.setattr( - demo_webrtc_module, + shared_webrtc_module, "create_packaged_webrtc_app", fake_create_packaged_webrtc_app, ) - adapter = OmnidreamsDemoAdapter(webrtc_runtime_factory=_FakeWebRTCRuntime) spec = DemoSpec( model_id=OMNIDREAMS_MODEL_ID, preset_id=DEFAULT_OMNIDREAMS_PRESET, @@ -419,40 +437,44 @@ def fake_create_packaged_webrtc_app(**kwargs: Any) -> web.Application: ), ) - demo = build_webrtc_demo(spec=spec, adapter=adapter, create_app=True) + app = serve_omnidreams_webrtc_demo( + spec=spec, + runtime_factory=_FakeWebRTCRuntime, + server_runner=lambda **kwargs: None, + ) - assert demo.app is not None - assert app_calls[0]["session_manager"] is demo.session_manager + assert isinstance(app, web.Application) + assert app_calls[0]["session_manager"] is app[SESSION_MANAGER_KEY] assert app_calls[0]["request_session_url"] == ( "http://127.0.0.1:8082/request_session" ) assert app_calls[0]["preload_name"] == "Test Omnidreams" - assert str(app_calls[0]["model_web_resource"]).endswith("omnidreams/webrtc/web") - route_paths = {resource.canonical for resource in demo.app.router.resources()} - assert "/api/postprocess/options" in route_paths - assert "/api/session/input" in route_paths + assert str(app_calls[0]["model_web_resource"]).endswith("omnidreams/demo/web") + assert app_calls[0]["configure_app"] is None def test_omnidreams_webrtc_demo_serves_through_shared_runner( monkeypatch: pytest.MonkeyPatch, ) -> None: + import flashdreams.runtime.demo.webrtc as shared_webrtc_module + server_calls: list[dict[str, Any]] = [] def fake_create_packaged_webrtc_app(**kwargs: Any) -> web.Application: app = web.Application() app[SESSION_MANAGER_KEY] = kwargs["session_manager"] - kwargs["configure_app"](app) + if configure_app := kwargs["configure_app"]: + configure_app(app) return app def fake_server_runner(**kwargs: Any) -> None: server_calls.append(kwargs) monkeypatch.setattr( - demo_webrtc_module, + shared_webrtc_module, "create_packaged_webrtc_app", fake_create_packaged_webrtc_app, ) - adapter = OmnidreamsDemoAdapter(webrtc_runtime_factory=_FakeWebRTCRuntime) spec = DemoSpec( model_id=OMNIDREAMS_MODEL_ID, preset_id=DEFAULT_OMNIDREAMS_PRESET, @@ -470,23 +492,59 @@ def fake_server_runner(**kwargs: Any) -> None: ), ) - demo = cast( - WebRTCDemo, - serve_flashdreams_demo( - spec=spec, - adapter=adapter, - world_rank=0, - server_runner=fake_server_runner, - ), + app = serve_omnidreams_webrtc_demo( + spec=spec, + world_rank=0, + runtime_factory=_FakeWebRTCRuntime, + server_runner=fake_server_runner, ) assert len(server_calls) == 1 assert server_calls[0]["world_rank"] == 0 - assert server_calls[0]["session_manager"] is demo.session_manager - assert server_calls[0]["app"] is demo.app + assert server_calls[0]["app"] is app assert server_calls[0]["host"] == "0.0.0.0" assert server_calls[0]["port"] == 8082 - assert isinstance(demo.session_manager, OmnidreamsDemoWebRTCSessionManager) + assert type(server_calls[0]["session_manager"]) is BaseWebRTCSessionManager + + +@pytest.mark.asyncio +async def test_omnidreams_demo_runtime_generates_directly_from_controls() -> None: + config = OmnidreamsWebRTCModelRuntimeConfig( + pipeline_config_name="fake", + pipeline_config=object(), + device="cpu", + fps=30, + warmup_chunks=0, + ) + runtime = OmnidreamsWebRTCModelRuntime(config=config) + wrapper = _FakeConditioningWrapper() + runtime._wrapper = wrapper # ty:ignore[invalid-assignment] + runtime._renderer = _FakeRenderer() + runtime._scene_data = SimpleNamespace(ego_poses=[SimpleNamespace(timestamp=1_000)]) + runtime._initial_rgb_frames = torch.zeros((1, 1, 3, 4, 5), dtype=torch.uint8) + runtime._text_prompts = [] + runtime._camera_to_rig = torch.eye(4) + runtime._initial_ego_pose = torch.eye(4).numpy() + runtime.pose_integrator.reset() + runtime._next_timestamp_us = 1_000 + + first = runtime._generate_one_chunk_sync( + segments=[(0.0, 2 / 30, frozenset({"w"}))], + frame_times=[1 / 30, 2 / 30], + ) + second = runtime._generate_one_chunk_sync( + segments=[(2 / 30, 5 / 30, frozenset({"d"}))], + frame_times=[3 / 30, 4 / 30, 5 / 30], + ) + + assert (first.step_index, first.frame_count) == (0, 2) + assert (second.step_index, second.frame_count) == (1, 3) + assert wrapper.calls == [ + ("start", (2, 4, 4), [1_000, 34_333]), + ("continue", (3, 4, 4), [67_666, 100_999, 134_332]), + ] + assert wrapper.finalized == [0, 1] + await runtime.close() class _RecordingOutputTarget: @@ -543,6 +601,63 @@ def finalize(self, *, autoregressive_index: int, cache: object) -> dict[str, flo return {"denoise_s": 0.25} +class _FakeRenderer: + def __init__(self) -> None: + self.closed = False + + def cleanup(self) -> None: + self.closed = True + + +class _FakeConditioningWrapper: + initial_frame_chunk_size = 2 + frame_chunk_size = 3 + + def __init__(self) -> None: + self.calls: list[tuple[str, tuple[int, ...], list[int]]] = [] + self.finalized: list[int] = [] + self.cleaned = False + + def start_generation(self, **kwargs: Any) -> SimpleNamespace: + return self._output("start", kwargs=kwargs, frame_count=2, step_index=0) + + def continue_generation(self, **kwargs: Any) -> SimpleNamespace: + return self._output("continue", kwargs=kwargs, frame_count=3, step_index=1) + + def _output( + self, + operation: str, + *, + kwargs: dict[str, Any], + frame_count: int, + step_index: int, + ) -> SimpleNamespace: + poses = kwargs["camera_poses_per_view"]["camera_front_wide_120fov"] + timestamps = kwargs["frame_timestamps_us"] + self.calls.append((operation, tuple(poses.shape), timestamps)) + state = kwargs.get("state") or SimpleNamespace(pipeline_cache=object()) + return SimpleNamespace( + state=state, + condition_frames=torch.zeros( + (1, 1, frame_count, 3, 4, 5), dtype=torch.uint8 + ), + rgb_frames=torch.zeros((1, 1, frame_count, 3, 4, 5), dtype=torch.uint8), + finalization_state={"autoregressive_index": step_index}, + ) + + def finalize_block_generation( + self, + pipeline_cache: object, + finalization_state: dict[str, int], + ) -> None: + del pipeline_cache + self.finalized.append(finalization_state["autoregressive_index"]) + + def cleanup(self, state: object) -> None: + del state + self.cleaned = True + + class _FakeWebRTCRuntime: def __init__(self, config: Any) -> None: self.config = config @@ -553,19 +668,23 @@ async def initialize(self) -> None: async def reset_for_new_session(self, *args: Any, **kwargs: Any) -> None: return None - def peek_steady_chunk_num_frames(self) -> int: - return 1 + def peek_input_fps(self) -> float: + return 30.0 - def peek_next_chunk_num_frames(self) -> int: + def peek_steady_output_num_frames(self) -> int: return 1 - async def generate_chunk( + def next_step_request(self) -> StepRequest: + return StepRequest(step_index=0, metadata={"input_frame_count": 1}) + + async def step( self, *, + request: StepRequest, segments: list[Any], frame_times: list[float], ) -> Any: - del segments, frame_times + del request, segments, frame_times return None async def close(self) -> None: diff --git a/integrations/omnidreams/tests/test_nvenc_smoke.py b/integrations/omnidreams/tests/test_nvenc_smoke.py index b52ee8dfd..c0499a0ba 100644 --- a/integrations/omnidreams/tests/test_nvenc_smoke.py +++ b/integrations/omnidreams/tests/test_nvenc_smoke.py @@ -25,6 +25,8 @@ import pytest import torch +from flashdreams.runtime import StepResult + pytestmark = pytest.mark.ci_gpu # ``PyNvVideoCodec`` probes for the NVIDIA driver library at import time @@ -103,7 +105,11 @@ def test_encode_chunk_produces_annex_b_keyframe_with_sps_pps( ) packets: list = [] num_frames, num_keyframes, encode_ms = encoder.encode_chunk_sync( - chunk, + StepResult.from_video_chunk( + step_index=0, + video_chunk=chunk, + layout="tchw", + ), force_keyframe=True, on_packet=packets.append, ) diff --git a/integrations/omnidreams/tests/test_webrtc_runtime.py b/integrations/omnidreams/tests/test_webrtc_runtime.py deleted file mode 100644 index 0776203f1..000000000 --- a/integrations/omnidreams/tests/test_webrtc_runtime.py +++ /dev/null @@ -1,1482 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import argparse -import asyncio -import json -import sys -import zipfile -from dataclasses import dataclass -from importlib.resources import files -from pathlib import Path -from types import SimpleNamespace -from typing import Any - -import pytest -import torch -from aiohttp import web -from aiohttp.test_utils import make_mocked_request -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 ( - OmnidreamsInferenceRuntime, - OmnidreamsRuntimeConfig, - OmnidreamsWebRTCSessionManager, -) - -import flashdreams.plugins.registry as plugin_registry -from flashdreams.infra.postprocess import ( - VideoPostprocessChainConfig, - VideoPostProcessorConfig, -) -from flashdreams.infra.video_output import VideoStepResult -from flashdreams.serving.webrtc.controls import ( - WSAD_SUPPORTED_KEYS, - CameraPoseIntegrator, -) -from flashdreams.serving.webrtc.encoders import ( - ChunkDeliveryResult, - DefaultRTCEncoder, -) -from flashdreams.serving.webrtc.media import BufferedVideoTrack -from flashdreams.serving.webrtc.server import SESSION_MANAGER_KEY - -pytestmark = pytest.mark.ci_cpu - - -class _FakeCloseable: - def __init__(self) -> None: - self.closed = False - - 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 _json_response_payload(response: web.StreamResponse) -> dict[str, Any]: - assert isinstance(response, web.Response) - text = response.text - assert text is not None - payload = json.loads(text) - assert isinstance(payload, dict) - return payload - - -def _fake_runtime_factory(config: OmnidreamsRuntimeConfig) -> object: - del config - return object() - - -def test_session_manager_hooks_are_wired() -> None: - # Guards against the shared base-class attribute overrides being dropped - # (e.g. losing their leading underscore), which silently reverts behaviour - # to the base defaults. - assert ( - OmnidreamsWebRTCSessionManager._busy_message - == "An Omnidreams session is already active." - ) - assert OmnidreamsWebRTCSessionManager._warmup_label == "Omnidreams WebRTC" - assert OmnidreamsWebRTCSessionManager._runtime_error_types == ( - session.OmnidreamsRuntimeError, - ) - # A fatal chunk-generation error tears the omnidreams session down. - assert OmnidreamsWebRTCSessionManager._close_session_on_generation_error is True - # Only the WSAD driving keys are accepted by the resampler. - assert ( - OmnidreamsWebRTCSessionManager._resampler_supported_keys == WSAD_SUPPORTED_KEYS - ) - - -@dataclass -class _FakeOutput: - state: Any - condition_frames: torch.Tensor - rgb_frames: torch.Tensor | None - finalization_state: dict[str, int] - - -class _FakeWrapper: - initial_frame_chunk_size = 2 - frame_chunk_size = 3 - - def __init__(self) -> None: - self.calls: list[tuple[str, tuple[int, ...], list[int]]] = [] - self.finalized: list[dict[str, int]] = [] - self.skip_video_generation_flags: list[bool] = [] - - def start_generation(self, **kwargs: Any) -> _FakeOutput: - poses = kwargs["camera_poses_per_view"]["camera_front_wide_120fov"] - timestamps = kwargs["frame_timestamps_us"] - self.calls.append(("start", tuple(poses.shape), timestamps)) - skip_video_generation = bool(kwargs.get("skip_video_generation", False)) - self.skip_video_generation_flags.append(skip_video_generation) - return _FakeOutput( - state=SimpleNamespace( - pipeline_cache=None if skip_video_generation else object() - ), - condition_frames=torch.full((1, 1, 2, 3, 4, 5), 31, dtype=torch.uint8), - rgb_frames=( - None - if skip_video_generation - else torch.zeros((1, 1, 2, 3, 4, 5), dtype=torch.uint8) - ), - finalization_state={"autoregressive_index": 0}, - ) - - def continue_generation(self, **kwargs: Any) -> _FakeOutput: - poses = kwargs["camera_poses_per_view"]["camera_front_wide_120fov"] - timestamps = kwargs["frame_timestamps_us"] - self.calls.append(("continue", tuple(poses.shape), timestamps)) - skip_video_generation = bool(kwargs.get("skip_video_generation", False)) - self.skip_video_generation_flags.append(skip_video_generation) - return _FakeOutput( - state=kwargs["state"], - condition_frames=torch.full((1, 1, 3, 3, 4, 5), 47, dtype=torch.uint8), - rgb_frames=( - None - if skip_video_generation - else torch.zeros((1, 1, 3, 3, 4, 5), dtype=torch.uint8) - ), - finalization_state={"autoregressive_index": 1}, - ) - - def finalize_block_generation( - self, pipeline_cache: object, finalization_state: dict[str, int] - ) -> None: - del pipeline_cache - self.finalized.append(finalization_state) - - -def _build_fake_runtime() -> tuple[OmnidreamsInferenceRuntime, _FakeWrapper]: - runtime = OmnidreamsInferenceRuntime( - config=OmnidreamsRuntimeConfig(device="cpu", fps=30) - ) - wrapper = _FakeWrapper() - runtime._wrapper = wrapper # ty:ignore[invalid-assignment] - runtime._renderer = object() - runtime._initial_rgb_frames = torch.zeros((1, 1, 3, 4, 5), dtype=torch.uint8) - runtime._text_prompts = [] - runtime._camera_to_rig = torch.eye(4) - runtime._device = torch.device("cpu") - runtime._next_timestamp_us = 1000 - runtime.pose_integrator = CameraPoseIntegrator() - runtime.pose_integrator.reset() - return runtime, wrapper - - -def test_generate_chunk_dispatches_start_then_continue() -> None: - runtime, wrapper = _build_fake_runtime() - - result0 = runtime._generate_one_chunk_sync( - segments=[(0.0, 2 / 30, frozenset({"w"}))], - frame_times=[1 / 30, 2 / 30], - ) - result1 = runtime._generate_one_chunk_sync( - segments=[(2 / 30, 5 / 30, frozenset())], - frame_times=[3 / 30, 4 / 30, 5 / 30], - ) - - assert result0.chunk_index == 0 - assert result0.num_frames == 2 - assert result1.chunk_index == 1 - assert result1.num_frames == 3 - assert wrapper.calls[0][0] == "start" - assert wrapper.calls[0][1] == (2, 4, 4) - 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.skip_video_generation_flags == [False, False] - - -def test_generate_chunk_postprocesses_rgb_before_cpu_handoff() -> None: - class _FakePostprocessStream: - def __init__(self) -> None: - self.calls: list[int] = [] - - def process( - self, video_chunk: torch.Tensor, *, autoregressive_index: int - ) -> torch.Tensor: - self.calls.append(autoregressive_index) - return torch.full( - (1, 1, video_chunk.shape[2], 3, 8, 10), - 0.5, - device=video_chunk.device, - ) - - runtime, _wrapper = _build_fake_runtime() - postprocess_stream = _FakePostprocessStream() - runtime._output_stream.postprocess_stream = ( # ty:ignore[invalid-assignment] - postprocess_stream - ) - - result = runtime._generate_one_chunk_sync( - segments=[(0.0, 2 / 30, frozenset({"w"}))], - frame_times=[1 / 30, 2 / 30], - ) - - assert postprocess_stream.calls == [0] - assert result.video_chunk.device.type == "cpu" - assert result.video_chunk.shape == (1, 1, 2, 3, 8, 10) - assert result.video_chunk.unique().tolist() == [0.5] - - -def test_session_postprocess_override_replaces_the_rollout_stream( - monkeypatch: pytest.MonkeyPatch, -) -> None: - preset_config = VideoPostProcessorConfig() - monkeypatch.setattr( - session, - "resolve_postprocess_preset", - lambda name: preset_config, - ) - monkeypatch.setattr( - plugin_registry, - "resolve_postprocess_preset", - lambda name: preset_config, - ) - runtime = OmnidreamsInferenceRuntime( - config=OmnidreamsRuntimeConfig( - device="cpu", - fps=30, - postprocess=VideoPostprocessChainConfig(preset="fake-preset"), - ) - ) - - runtime._reset_postprocess_stream( - session.OmnidreamsSessionInput(postprocess_preset="fake-preset") - ) - first_stream = runtime._output_stream.postprocess_stream - - assert first_stream is not None - assert runtime.postprocess_preset == "fake-preset" - - runtime._reset_postprocess_stream( - session.OmnidreamsSessionInput(postprocess_preset="") - ) - - assert first_stream._closed is True - assert runtime._output_stream.postprocess_stream is None - assert runtime.postprocess_preset == "" - - -def test_generate_chunk_can_stream_debug_hdmaps_without_rgb_frames() -> None: - runtime, wrapper = _build_fake_runtime() - runtime.config.debug_serve_hdmaps = True - - result0 = runtime._generate_one_chunk_sync( - segments=[(0.0, 2 / 30, frozenset({"w"}))], - frame_times=[1 / 30, 2 / 30], - ) - result1 = runtime._generate_one_chunk_sync( - segments=[(2 / 30, 5 / 30, frozenset({"d"}))], - frame_times=[3 / 30, 4 / 30, 5 / 30], - ) - - assert result0.chunk_index == 0 - assert result0.num_frames == 2 - assert result0.video_chunk.shape == (1, 1, 2, 3, 4, 5) - assert result0.video_chunk.unique().tolist() == [31] - assert result1.chunk_index == 1 - assert result1.num_frames == 3 - assert result1.video_chunk.shape == (1, 1, 3, 3, 4, 5) - assert result1.video_chunk.unique().tolist() == [47] - assert wrapper.skip_video_generation_flags == [True, True] - assert wrapper.finalized == [] - - -def test_prepare_clipgt_dir_stages_unprefixed_parquets( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -) -> None: - clipgt = tmp_path / "clipgt" - clipgt.mkdir() - (clipgt / "calibration_estimate.parquet").touch() - (clipgt / "egomotion_estimate.parquet").touch() - (clipgt / "lane.parquet").touch() - runtime = OmnidreamsInferenceRuntime( - config=OmnidreamsRuntimeConfig(device="cpu", fps=30) - ) - - staged = runtime._prepare_clipgt_dir(clipgt) - - assert staged != clipgt - assert (staged / "clip.calibration_estimate.parquet").exists() - assert (staged / "clip.egomotion_estimate.parquet").exists() - assert (staged / "clip.lane.parquet").exists() - - monkeypatch.chdir(tmp_path) - staged_from_relative = runtime._prepare_clipgt_dir(Path("clipgt")) - assert (staged_from_relative / "clip.calibration_estimate.parquet").exists() - - -def test_prepare_clipgt_dir_stages_nested_unprefixed_parquets(tmp_path: Path) -> None: - clipgt = tmp_path / "clipgt" - clipgt.mkdir() - nested = clipgt / "clipgt" - nested.mkdir() - (clipgt / "first_image.png").touch() - (clipgt / "prompt.txt").touch() - (nested / "calibration_estimate.parquet").touch() - (nested / "egomotion_estimate.parquet").touch() - (nested / "lane.parquet").touch() - runtime = OmnidreamsInferenceRuntime( - config=OmnidreamsRuntimeConfig(device="cpu", fps=30) - ) - - staged = runtime._prepare_clipgt_dir(clipgt) - - assert staged != clipgt - assert (staged / "clip.calibration_estimate.parquet").exists() - assert (staged / "clip.egomotion_estimate.parquet").exists() - assert (staged / "clip.lane.parquet").exists() - - -def test_link_or_copy_file_falls_back_to_copy( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -) -> None: - source = tmp_path / "source.parquet" - target = tmp_path / "target.parquet" - source.write_bytes(b"parquet data") - - def _raise_link_error(*args: object, **kwargs: object) -> None: - del args, kwargs - raise OSError("links unavailable") - - monkeypatch.setattr(session.os, "symlink", _raise_link_error) - monkeypatch.setattr(session.os, "link", _raise_link_error) - - session._link_or_copy_file(source, target) - - assert target.read_bytes() == source.read_bytes() - assert not target.is_symlink() - - -def test_hf_webrtc_scene_sync_requires_usdz_first_frame( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -) -> None: - scene_uuid = "065dcac9-ee67-4434-a835-c6b816c88e48" - archive_repo_path = f"scenes/clipgt-{scene_uuid}.usdz" - archive_path = tmp_path / "clipgt.usdz" - with zipfile.ZipFile(archive_path, "w") as zf: - zf.writestr("calibration_estimate.parquet", "calibration") - zf.writestr("egomotion_estimate.parquet", "egomotion") - zf.writestr("prompt.txt", "archive prompt") - - def _fake_hf_hub_download(repo_id: str, repo_type: str, filename: str) -> str: - assert repo_id == session.hf_scenes_repo_id() - assert repo_type == "dataset" - assert filename == archive_repo_path - return str(archive_path) - - cache_dir = tmp_path / "flashdreams-cache" - stale_scene_dir = cache_dir / "omnidreams-scenes" / scene_uuid - stale_scene_dir.mkdir(parents=True) - (stale_scene_dir / "first_frame.jpeg").write_text( - "stale first frame", encoding="utf-8" - ) - (stale_scene_dir / "prompt.txt").write_text("stale prompt", encoding="utf-8") - - monkeypatch.setattr(scenes, "FLASHDREAMS_CACHE_DIR", cache_dir) - monkeypatch.setattr( - "huggingface_hub.hf_hub_download", - _fake_hf_hub_download, - ) - - scene_dir = session._ensure_hf_webrtc_scene_synced(scene_uuid) - - with pytest.raises(FileNotFoundError, match="first_image"): - session._resolve_webrtc_scene_assets( - scene_dir, - prompt_filename="prompt.txt", - clipgt_dirname="clipgt", - ) - - -def test_hf_webrtc_scene_sync_uses_extracted_first_image( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -) -> None: - scene_uuid = "065dcac9-ee67-4434-a835-c6b816c88e48" - archive_repo_path = f"scenes/clipgt-{scene_uuid}.usdz" - archive_path = tmp_path / "clipgt.usdz" - with zipfile.ZipFile(archive_path, "w") as zf: - zf.writestr("calibration_estimate.parquet", "calibration") - zf.writestr("egomotion_estimate.parquet", "egomotion") - zf.writestr("first_image.png", "first image") - zf.writestr("prompt.txt", "archive prompt") - - def _fake_hf_hub_download(repo_id: str, repo_type: str, filename: str) -> str: - assert repo_id == session.hf_scenes_repo_id() - assert repo_type == "dataset" - assert filename == archive_repo_path - return str(archive_path) - - cache_dir = tmp_path / "flashdreams-cache" - stale_scene_dir = cache_dir / "omnidreams-scenes" / scene_uuid - stale_scene_dir.mkdir(parents=True) - (stale_scene_dir / "first_frame.jpeg").write_text( - "stale first frame", encoding="utf-8" - ) - (stale_scene_dir / "prompt.txt").write_text("stale prompt", encoding="utf-8") - - monkeypatch.setattr(scenes, "FLASHDREAMS_CACHE_DIR", cache_dir) - monkeypatch.setattr( - "huggingface_hub.hf_hub_download", - _fake_hf_hub_download, - ) - - scene_dir = session._ensure_hf_webrtc_scene_synced(scene_uuid) - - assert (scene_dir / "clipgt" / "first_image.png").read_text( - encoding="utf-8" - ) == "first image" - assert (scene_dir / "clipgt" / "prompt.txt").read_text( - encoding="utf-8" - ) == "archive prompt" - - clipgt_dir, first_frame_path, prompt_path = session._resolve_webrtc_scene_assets( - scene_dir, - prompt_filename="prompt.txt", - clipgt_dirname="clipgt", - ) - assert clipgt_dir == scene_dir / "clipgt" - assert first_frame_path == scene_dir / "clipgt" / "first_image.png" - assert prompt_path == scene_dir / "clipgt" / "prompt.txt" - - -def test_hf_webrtc_scene_sync_requires_usdz_prompt( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -) -> None: - scene_uuid = "065dcac9-ee67-4434-a835-c6b816c88e48" - archive_repo_path = f"scenes/clipgt-{scene_uuid}.usdz" - archive_path = tmp_path / "clipgt.usdz" - with zipfile.ZipFile(archive_path, "w") as zf: - zf.writestr("calibration_estimate.parquet", "calibration") - zf.writestr("egomotion_estimate.parquet", "egomotion") - zf.writestr("first_image.png", "first image") - - def _fake_hf_hub_download(repo_id: str, repo_type: str, filename: str) -> str: - assert repo_id == session.hf_scenes_repo_id() - assert repo_type == "dataset" - assert filename == archive_repo_path - return str(archive_path) - - monkeypatch.setattr(scenes, "FLASHDREAMS_CACHE_DIR", tmp_path / "flashdreams-cache") - monkeypatch.setattr( - "huggingface_hub.hf_hub_download", - _fake_hf_hub_download, - ) - - scene_dir = session._ensure_hf_webrtc_scene_synced(scene_uuid) - - with pytest.raises(FileNotFoundError, match="prompt.txt"): - session._resolve_webrtc_scene_assets( - scene_dir, - prompt_filename="prompt.txt", - clipgt_dirname="clipgt", - ) - - -def test_resolved_empty_prompt_keeps_runtime_default_behavior(tmp_path: Path) -> None: - scene_dir = tmp_path / "scene" - clipgt_dir = scene_dir / "clipgt" - clipgt_dir.mkdir(parents=True) - (clipgt_dir / "first_image.png").write_text("first image", encoding="utf-8") - (clipgt_dir / "prompt.txt").write_text("", encoding="utf-8") - - _, _, prompt_path = session._resolve_webrtc_scene_assets( - scene_dir, - prompt_filename="prompt.txt", - clipgt_dirname="clipgt", - ) - - assert prompt_path == clipgt_dir / "prompt.txt" - assert ( - prompt_path.read_text(encoding="utf-8").strip() or session.AV_POSITIVE_PROMPT - ) == session.AV_POSITIVE_PROMPT - - -def test_build_runtime_config_threads_hf_scene_args(tmp_path: Path) -> None: - args = argparse.Namespace( - pipeline_config_name="omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf", - scene_dir=tmp_path / "local-scene", - scene_uuid="scene-123", - scene_variant="rain", - seed=123, - 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=True, - postprocess_preset="rtx-super-resolution", - prefer_sw_encoder=False, - ) - - cfg = webrtc_server.build_runtime_config(args, device_override="cuda:7") - - assert cfg.scene_dir == tmp_path / "local-scene" - assert cfg.scene_uuid == "scene-123" - assert cfg.scene_variant == "rain" - assert cfg.device == "cuda:7" - assert cfg.video_height == 360 - assert cfg.video_width == 640 - assert cfg.debug_serve_hdmaps is True - assert cfg.postprocess.preset == "rtx-super-resolution" - # ``--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, expected_backend", - [(False, "auto"), (True, "default")], -) -def test_build_runtime_config_maps_prefer_sw_encoder_to_backend( - tmp_path: Path, - prefer_sw_encoder: bool, - expected_backend: str, -) -> None: - """--prefer_sw_encoder is the single CLI switch that toggles between - the auto-probe path and the forced-software path. Any regression in - this mapping would silently disable the hardware encoder (or worse, - fail to disable it when explicitly asked).""" - 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, - postprocess_preset="", - prefer_sw_encoder=prefer_sw_encoder, - ) - cfg = webrtc_server.build_runtime_config(args) - assert cfg.encoder_backend == expected_backend - - -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_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_parse_args_omits_scene_dir_by_default( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr( - sys, - "argv", - [ - "omnidreams.webrtc.server", - "--debug_serve_hdmaps", - ], - ) - - args = webrtc_server.parse_args() - - assert args.scene_dir is None - assert args.scene_uuid is None - assert args.debug_serve_hdmaps is True - assert args.postprocess_preset == "" - - -def test_runtime_initialization_passes_manifest_pipeline_config( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -) -> None: - manifest_args = webrtc_server.parse_args( - [ - "--manifest", - "example_world_model_perf.yaml", - "--prefer_sw_encoder", - ] - ) - cfg = webrtc_server.build_runtime_config(manifest_args, device_override="cpu") - 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: - staged_scene_dir = tmp_path / "staged-scene" - calls: list[str] = [] - - def _fake_ensure_hf_webrtc_scene_synced( - scene_uuid: str, - *, - variant: str = "default", - prompt_filename: str, - clipgt_dirname: str, - ) -> Path: - del prompt_filename, clipgt_dirname, variant - calls.append(scene_uuid) - return staged_scene_dir - - def _fake_resolve_webrtc_scene_assets( - scene_dir: Path, - *, - prompt_filename: str, - clipgt_dirname: str, - camera_name: str = "camera_front_wide_120fov", - variant: str = "default", - ) -> tuple[Path, Path, Path]: - del prompt_filename, clipgt_dirname, camera_name, variant - clipgt_dir = scene_dir / "clipgt" - return clipgt_dir, clipgt_dir / "first_image.png", clipgt_dir / "prompt.txt" - - monkeypatch.setattr( - session, - "_ensure_hf_webrtc_scene_synced", - _fake_ensure_hf_webrtc_scene_synced, - ) - monkeypatch.setattr( - session, - "_resolve_webrtc_scene_assets", - _fake_resolve_webrtc_scene_assets, - ) - monkeypatch.setattr(session, "load_scene", lambda *args, **kwargs: None) - runtime = OmnidreamsInferenceRuntime( - config=OmnidreamsRuntimeConfig( - pipeline_config_name="missing-config", - device="cpu", - scene_dir=None, - scene_uuid=None, - ) - ) - - with pytest.raises(ValueError, match="Unknown pipeline_config_name"): - runtime._initialize_sync() - - assert calls == [session.DEFAULT_WEBRTC_SCENE_UUID] - - -def test_build_runtime_config_clears_scene_uuid_for_local_scene(tmp_path: Path) -> None: - 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=123, - 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=True, - postprocess_preset="", - prefer_sw_encoder=False, - ) - - cfg = webrtc_server.build_runtime_config(args) - - assert cfg.scene_dir == tmp_path / "local-scene" - assert cfg.scene_uuid is None - assert cfg.scene_variant == "default" - - -def test_session_manager_stores_postprocess_override_for_next_rollout() -> None: - manager = OmnidreamsWebRTCSessionManager( - runtime_config=OmnidreamsRuntimeConfig(device="cpu") - ) - session_input = session.OmnidreamsSessionInput(postprocess_preset="") - - manager.set_pending_session_input(session_input) - - assert manager._peek_pending_session_input() == session_input - manager._clear_pending_session_input() - assert manager._peek_pending_session_input() is None - - -def test_session_manager_rejects_unlaunched_postprocess_preset() -> None: - manager = OmnidreamsWebRTCSessionManager( - runtime_config=OmnidreamsRuntimeConfig(device="cpu") - ) - - with pytest.raises(ValueError, match="not enabled for this server"): - manager.set_pending_session_input( - session.OmnidreamsSessionInput(postprocess_preset="fake-preset") - ) - - -def test_session_manager_rejects_non_launched_postprocess_preset( - monkeypatch: pytest.MonkeyPatch, -) -> None: - preset_config = VideoPostProcessorConfig() - monkeypatch.setattr( - session, - "resolve_postprocess_preset", - lambda name: preset_config, - ) - manager = OmnidreamsWebRTCSessionManager( - runtime_config=OmnidreamsRuntimeConfig( - device="cpu", - postprocess=VideoPostprocessChainConfig(preset="launched-preset"), - ) - ) - - with pytest.raises(ValueError, match="must match the launched preset"): - manager.set_pending_session_input( - session.OmnidreamsSessionInput(postprocess_preset="other-preset") - ) - - -@pytest.mark.asyncio -async def test_postprocess_options_hide_unlaunched_presets() -> None: - manager = OmnidreamsWebRTCSessionManager( - runtime_config=OmnidreamsRuntimeConfig(device="cpu") - ) - app = web.Application() - app[SESSION_MANAGER_KEY] = manager - request = make_mocked_request("GET", "/api/postprocess/options", app=app) - - response = await webrtc_server._postprocess_options(request) - payload = _json_response_payload(response) - - assert payload == {"default_preset": "", "presets": []} - - -@pytest.mark.asyncio -async def test_postprocess_options_exposes_only_launch_preset() -> None: - manager = OmnidreamsWebRTCSessionManager( - runtime_config=OmnidreamsRuntimeConfig( - device="cpu", - postprocess=VideoPostprocessChainConfig(preset="launched-preset"), - ) - ) - app = web.Application() - app[SESSION_MANAGER_KEY] = manager - request = make_mocked_request("GET", "/api/postprocess/options", app=app) - - response = await webrtc_server._postprocess_options(request) - payload = _json_response_payload(response) - - assert payload == { - "default_preset": "launched-preset", - "presets": ["launched-preset"], - } - - -def test_webrtc_ui_posts_selected_postprocess_preset() -> None: - shared_web_dir = files("flashdreams.serving.webrtc").joinpath("web") - javascript = shared_web_dir.joinpath("request_session.js").read_text( - encoding="utf-8" - ) - adapter = ( - files("omnidreams.webrtc") - .joinpath("web", "adapter.js") - .read_text(encoding="utf-8") - ) - - assert 'fetch("/api/postprocess/options")' in javascript - assert 'fetch("/api/session/input"' in javascript - assert "postprocessAvailable" in javascript - assert "postprocessField.hidden = !postprocessAvailable" in javascript - assert "postprocess_preset: postprocessPreset" in javascript - assert "enablePostprocess: true" in adapter - assert "/api/postprocess/options" not in adapter - - -@pytest.mark.asyncio -async def test_session_manager_preload_runs_loopback_warmup_once( - monkeypatch: pytest.MonkeyPatch, -) -> None: - class _FakeRuntime: - def __init__(self, config: OmnidreamsRuntimeConfig) -> None: - self.config = config - self.initialize_calls = 0 - self.close_calls = 0 - - async def initialize(self) -> None: - self.initialize_calls += 1 - - async def close(self) -> None: - self.close_calls += 1 - - fake_runtime: _FakeRuntime | None = None - warmup_calls: list[int] = [] - - def _fake_runtime_factory(config: OmnidreamsRuntimeConfig) -> _FakeRuntime: - nonlocal fake_runtime - fake_runtime = _FakeRuntime(config) - return fake_runtime - - async def _fake_loopback_warmup( - self: OmnidreamsWebRTCSessionManager, *, num_chunks: int - ) -> None: - del self - warmup_calls.append(num_chunks) - - monkeypatch.setattr(session, "OmnidreamsInferenceRuntime", _fake_runtime_factory) - monkeypatch.setattr( - OmnidreamsWebRTCSessionManager, - "_run_loopback_warmup_session", - _fake_loopback_warmup, - ) - manager = OmnidreamsWebRTCSessionManager( - runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=2) - ) - - await manager.preload_runtime() - await manager.preload_runtime() - - assert fake_runtime is not None - assert fake_runtime.initialize_calls == 1 - assert warmup_calls == [2] - assert manager.is_runtime_ready() - - -@pytest.mark.asyncio -async def test_loopback_warmup_drives_session_generation( - monkeypatch: pytest.MonkeyPatch, -) -> None: - class _FakeRuntime: - def __init__(self, config: OmnidreamsRuntimeConfig) -> None: - self.config = config - self.initialize_calls = 0 - self.reset_calls = 0 - self.close_calls = 0 - self.postprocess_preset = config.postprocess.preset - self.generated_segments: list[ - list[tuple[float, float, frozenset[str]]] - ] = [] - # 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 - - async def reset_for_new_session( - self, session_input: session.OmnidreamsSessionInput | None = None - ) -> None: - del session_input - self.reset_calls += 1 - - def peek_steady_chunk_num_frames(self) -> int: - return 1 - - def peek_next_chunk_num_frames(self) -> int: - return 1 - - async def generate_chunk( - self, - *, - segments: list[tuple[float, float, frozenset[str]]], - frame_times: list[float], - ) -> VideoStepResult: - del frame_times - chunk_index = len(self.generated_segments) - self.generated_segments.append(segments) - return VideoStepResult( - chunk_index=chunk_index, - num_frames=1, - video_chunk=torch.zeros((1, 1, 1, 3, 2, 2), dtype=torch.uint8), - stats=None, - ) - - async def close(self) -> None: - self.close_calls += 1 - - fake_runtime: _FakeRuntime | None = None - - def _fake_runtime_factory(config: OmnidreamsRuntimeConfig) -> _FakeRuntime: - nonlocal fake_runtime - fake_runtime = _FakeRuntime(config) - return fake_runtime - - monkeypatch.setattr(session, "OmnidreamsInferenceRuntime", _fake_runtime_factory) - manager = OmnidreamsWebRTCSessionManager( - runtime_config=OmnidreamsRuntimeConfig( - device="cpu", - fps=30, - warmup_chunks=2, - ) - ) - - await asyncio.wait_for(manager.preload_runtime(), timeout=10.0) - - assert fake_runtime is not None - assert fake_runtime.initialize_calls == 1 - assert fake_runtime.reset_calls == 1 - # The close signal can race with the generation worker starting the next - # chunk; the warmup contract is that at least the requested chunks complete. - assert len(fake_runtime.generated_segments) >= 2 - assert not manager.has_active_session() - - -@pytest.mark.asyncio -async def test_heartbeat_message_refreshes_client_liveness( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(session, "OmnidreamsInferenceRuntime", _fake_runtime_factory) - manager = OmnidreamsWebRTCSessionManager( - runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=0) - ) - 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(), - last_client_message_at=0.0, - ) - manager._active_session = managed_session - - await manager._handle_datachannel_message( - managed_session=managed_session, - raw_message='{"type":"heartbeat"}', - ) - - assert managed_session.last_client_message_at > 0.0 - assert manager.has_active_session() - - -@pytest.mark.asyncio -async def test_client_liveness_timeout_closes_active_session( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(session, "OmnidreamsInferenceRuntime", _fake_runtime_factory) - manager = OmnidreamsWebRTCSessionManager( - runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=0), - client_liveness_timeout_s=0.01, - ) - video_track = _FakeCloseable() - peer_connection = _FakeCloseable() - 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, - ) - manager._active_session = managed_session - liveness_task = asyncio.create_task( - manager._client_liveness_watchdog(managed_session=managed_session) - ) - managed_session.liveness_task = liveness_task - - await asyncio.wait_for(liveness_task, timeout=1.0) - - assert not manager.has_active_session() - assert managed_session.closed - assert video_track.closed - assert peer_connection.closed - - -@pytest.mark.asyncio -async def test_disconnect_message_closes_active_session( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(session, "OmnidreamsInferenceRuntime", _fake_runtime_factory) - manager = OmnidreamsWebRTCSessionManager( - runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=0) - ) - video_track = _FakeCloseable() - peer_connection = _FakeCloseable() - 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(), - ) - manager._active_session = managed_session - - await manager._handle_datachannel_message( - managed_session=managed_session, - raw_message='{"type":"disconnect"}', - ) - - assert not manager.has_active_session() - assert managed_session.closed - assert video_track.closed - assert peer_connection.closed - - -@pytest.mark.asyncio -async def test_generation_worker_closes_session_after_generation_failure() -> None: - class _FailingRuntime: - def __init__(self) -> None: - self.generate_calls = 0 - - def peek_next_chunk_num_frames(self) -> int: - return 1 - - async def generate_chunk( - self, - *, - segments: list[tuple[float, float, frozenset[str]]], - frame_times: list[float], - ) -> VideoStepResult: - del segments, frame_times - self.generate_calls += 1 - raise RuntimeError("boom") - - class _FakeResampler: - dt = 0.0 - next_chunk_start_v = 0.0 - - def sample_chunk( - self, num_frames: int - ) -> tuple[list[tuple[float, float, frozenset[str]]], list[float]]: - assert num_frames == 1 - return [(0.0, 0.0, frozenset({"w"}))], [0.0] - - class _FakeVideoTrack: - fps = 30 - - def __init__(self) -> None: - self.closed = False - - async def close(self) -> None: - self.closed = True - - def qsize(self) -> int: - return 0 - - class _FakePeerConnection: - def __init__(self) -> None: - self.closed = False - - async def close(self) -> None: - self.closed = True - - class _FakeChannel: - def __init__(self) -> None: - self.messages: list[str] = [] - - def send(self, message: str) -> None: - self.messages.append(message) - - manager = OmnidreamsWebRTCSessionManager( - runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=0) - ) - runtime = _FailingRuntime() - video_track = _FakeVideoTrack() - peer_connection = _FakePeerConnection() - control_channel = _FakeChannel() - first_action_received = asyncio.Event() - first_action_received.set() - 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, - first_action_received=first_action_received, - ) - manager._active_session = managed_session - - task = asyncio.create_task( - manager._generation_worker(managed_session=managed_session) - ) - managed_session.generation_task = task - - await task - - assert runtime.generate_calls == 1 - assert not manager.has_active_session() - assert managed_session.closed - 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] - ) - - -@pytest.mark.asyncio -async 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")]) - - await manager._enforce_h264_or_fallback( - transceiver=transceiver, - managed_session=managed_session, - num_frames=4, - ) - - assert not hw_encoder.closed, ( - "runtime-owned hardware encoder must survive a session-scope fallback" - ) - assert original_track.closed, "orphaned hardware track was not 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 - - -@pytest.mark.asyncio -async 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")]) - - await manager._enforce_h264_or_fallback( - transceiver=transceiver, - managed_session=managed_session, - num_frames=4, - ) - - assert not hw_encoder.closed - assert not original_track.closed - assert managed_session.video_encoder is hw_encoder - assert managed_session.video_track is original_track - assert transceiver.sender.replaced_with is None - - -@pytest.mark.asyncio -async 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) - original_track = _FakeCloseable() - managed_session = _sdp_fallback_managed_session(hw_encoder) - managed_session.video_track = original_track # ty:ignore[invalid-assignment] - transceiver = _FakeTransceiver([]) - - await manager._enforce_h264_or_fallback( - transceiver=transceiver, - managed_session=managed_session, - num_frames=4, - ) - - assert not hw_encoder.closed, ( - "runtime-owned hardware encoder must survive a session-scope fallback" - ) - assert original_track.closed - assert isinstance(managed_session.video_encoder, DefaultRTCEncoder) - assert isinstance(managed_session.video_track, BufferedVideoTrack) - - -def test_initialize_video_encoder_sync_skips_on_non_master( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """WebRTC media is served only by the master rank, so worker ranks - must not reach ``select_encoder`` — allocating an NVENC session on a - worker would consume a local GPU concurrent-session slot without - ever encoding a frame, and could fail the worker's startup if the - pool cannot accommodate one allocation per rank.""" - - def _select_encoder_should_not_be_called(**_kw: Any) -> object: - raise AssertionError( - "_initialize_video_encoder_sync must not reach select_encoder " - "on non-master ranks" - ) - - monkeypatch.setattr( - session, - "select_encoder", - _select_encoder_should_not_be_called, - ) - - runtime = OmnidreamsInferenceRuntime( - config=OmnidreamsRuntimeConfig(device="cpu", fps=30) - ) - runtime.rank = 1 # simulate a worker rank - runtime._device = torch.device("cpu") - - runtime._initialize_video_encoder_sync() - - assert runtime._video_encoder is None - - -def test_initialize_video_encoder_sync_runs_on_master( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Master rank still initializes the encoder normally.""" - stub = _FakeVideoEncoder() - calls: list[dict[str, Any]] = [] - - def _fake_select_encoder(**kwargs: Any) -> _FakeVideoEncoder: - calls.append(kwargs) - return stub - - monkeypatch.setattr(session, "select_encoder", _fake_select_encoder) - - runtime = OmnidreamsInferenceRuntime( - config=OmnidreamsRuntimeConfig(device="cpu", fps=30) - ) - runtime.rank = 0 - runtime._device = torch.device("cpu") - - runtime._initialize_video_encoder_sync() - - assert len(calls) == 1 - assert runtime._video_encoder is stub diff --git a/integrations/omnidreams/tests/test_webrtc_server_routes.py b/integrations/omnidreams/tests/test_webrtc_server_routes.py deleted file mode 100644 index d1104057e..000000000 --- a/integrations/omnidreams/tests/test_webrtc_server_routes.py +++ /dev/null @@ -1,324 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import logging -from contextlib import ExitStack - -import pytest -from aiohttp.test_utils import TestClient, TestServer -from omnidreams.webrtc import server as webrtc_server -from omnidreams.webrtc.server import ( - _close_package_resources, - configure_logging, - create_app, -) - -from flashdreams.serving.webrtc.server import ( - PACKAGE_RESOURCE_STACK_KEY, - SessionBusyError, -) - -pytestmark = pytest.mark.ci_gpu - - -class FakeSessionManager: - def __init__(self) -> None: - self.answer_payload = {"sdp": "fake-answer-sdp", "type": "answer"} - self.raise_busy = False - self.preload_calls = 0 - self.offers: list[tuple[str, str]] = [] - self.active = False - self.runtime_ready = False - - def has_active_session(self) -> bool: - return self.active - - def is_runtime_ready(self) -> bool: - return self.runtime_ready - - async def preload_runtime(self) -> None: - self.preload_calls += 1 - self.runtime_ready = True - - async def create_answer(self, *, offer_sdp: str, offer_type: str) -> dict[str, str]: - self.offers.append((offer_sdp, offer_type)) - if self.raise_busy: - raise SessionBusyError("An Omnidreams session is already active.") - self.active = True - return self.answer_payload - - async def shutdown(self) -> None: - self.active = False - self.runtime_ready = False - - -async def _build_client(manager: FakeSessionManager) -> TestClient: - app = create_app( - session_manager=manager, - request_session_url="http://127.0.0.1:8080/request_session", - ) - server = TestServer(app) - client = TestClient(server) - await client.start_server() - return client - - -def test_create_app_keeps_package_web_resource_materialized() -> None: - app = create_app( - session_manager=FakeSessionManager(), - request_session_url="http://127.0.0.1:8080/request_session", - ) - try: - assert isinstance(app[PACKAGE_RESOURCE_STACK_KEY], ExitStack) - assert _close_package_resources in app.on_cleanup - - static_resources = [ - resource - for resource in app.router.resources() - if getattr(resource, "canonical", "") == "/static" - or resource.get_info().get("prefix") in {"/static", "/static/"} - ] - assert len(static_resources) == 1 - web_dir = static_resources[0].get_info()["directory"] - assert web_dir.is_dir() - assert ( - "FlashDreams WebRTC Drive" in (web_dir / "request_session.html").read_text() - ) - finally: - app[PACKAGE_RESOURCE_STACK_KEY].close() - - -def test_create_app_closes_package_resource_when_app_creation_fails( - monkeypatch, tmp_path -) -> None: - class TrackedResource: - closed = False - - def __enter__(self): - return tmp_path - - def __exit__(self, exc_type, exc_value, traceback): - self.closed = True - - tracked_resource = TrackedResource() - - def raise_app_creation_failure(**_kwargs): - raise RuntimeError("app creation failed") - - monkeypatch.setattr(webrtc_server, "as_file", lambda _resource: tracked_resource) - monkeypatch.setattr( - webrtc_server, - "create_webrtc_app", - raise_app_creation_failure, - ) - - with pytest.raises(RuntimeError, match="app creation failed"): - create_app( - session_manager=FakeSessionManager(), - request_session_url="http://127.0.0.1:8080/request_session", - ) - - assert tracked_resource.closed - - -@pytest.mark.asyncio -async def test_request_session_serves_html() -> None: - manager = FakeSessionManager() - client = await _build_client(manager) - try: - assert manager.preload_calls == 1 - response = await client.get("/request_session") - body = await response.text() - assert response.status == 200 - assert "FlashDreams WebRTC Drive" in body - finally: - await client.close() - - -@pytest.mark.asyncio -async def test_request_session_uses_lingbot_aligned_viewer_shell() -> None: - manager = FakeSessionManager() - client = await _build_client(manager) - try: - response = await client.get("/request_session") - body = await response.text() - assert response.status == 200 - assert 'class="brandOverlay"' in body - assert "FlashDreams" in body - assert "/static/assets/horizontal-dark.svg" in body - assert 'class="statusCard overlayPanel"' in body - assert 'class="controlCard overlayPanel"' in body - assert 'class="logCard overlayPanel"' in body - assert "Connect Session" in body - assert 'id="logState"' in body - assert "World Model" in body - assert 'id="controlRows"' in body - assert 'id="modelStatusSlot"' in body - finally: - await client.close() - - -@pytest.mark.asyncio -async def test_request_session_includes_idle_animation_canvas() -> None: - manager = FakeSessionManager() - client = await _build_client(manager) - try: - response = await client.get("/request_session") - body = await response.text() - assert response.status == 200 - assert ( - '' - in body - ) - finally: - await client.close() - - -@pytest.mark.asyncio -async def test_shared_flashdreams_brand_asset_is_served() -> None: - manager = FakeSessionManager() - client = await _build_client(manager) - try: - response = await client.get("/static/assets/horizontal-dark.svg") - assert response.status == 200 - assert response.content_type == "image/svg+xml" - finally: - await client.close() - - -@pytest.mark.asyncio -async def test_static_js_requests_recvonly_video_transceiver() -> None: - manager = FakeSessionManager() - client = await _build_client(manager) - try: - response = await client.get("/static/request_session.js") - body = await response.text() - assert response.status == 200 - assert 'addTransceiver("video", { direction: "recvonly" })' in body - finally: - await client.close() - - -@pytest.mark.asyncio -async def test_static_js_keeps_generic_controls_and_status_helpers() -> None: - manager = FakeSessionManager() - client = await _build_client(manager) - try: - response = await client.get("/static/request_session.js") - body = await response.text() - assert response.status == 200 - assert "const defaultControls = [" in body - assert "function renderControls(groups)" in body - assert 'const logState = document.getElementById("logState")' in body - assert 'logState.textContent = state === "idle" ? "Waiting" : message' in body - assert "eventLog.prepend(entry)" in body - finally: - await client.close() - - -@pytest.mark.asyncio -async def test_omnidreams_model_adapter_is_served() -> None: - client = await _build_client(FakeSessionManager()) - try: - config = await (await client.get("/api/ui/config")).json() - assert config["adapter_module"].startswith("/model-static/adapter.js") - response = await client.get("/model-static/adapter.js") - body = await response.text() - assert response.status == 200 - assert 'modelName: "OmniDreams"' in body - assert "enablePostprocess: true" in body - assert "/api/postprocess/options" not in body - assert "RTCPeerConnection" not in body - finally: - await client.close() - - -@pytest.mark.asyncio -async def test_static_js_draws_idle_animation_until_video_arrives() -> None: - manager = FakeSessionManager() - client = await _build_client(manager) - try: - response = await client.get("/static/request_session.js") - body = await response.text() - assert response.status == 200 - assert 'const idleCanvas = document.getElementById("idleCanvas")' in body - assert "function drawIdleScene(now)" in body - assert "window.requestAnimationFrame(drawIdleScene)" in body - finally: - await client.close() - - -@pytest.mark.asyncio -async def test_static_css_uses_lingbot_overlay_classes() -> None: - manager = FakeSessionManager() - client = await _build_client(manager) - try: - response = await client.get("/static/request_session.css") - body = await response.text() - assert response.status == 200 - for selector in ( - ".overlayPanel", - ".brandOverlay", - ".statusCard", - ".controlCard", - ".logCard", - ): - assert selector in body - finally: - await client.close() - - -@pytest.mark.asyncio -async def test_static_css_fades_idle_animation_after_video_arrives() -> None: - manager = FakeSessionManager() - client = await _build_client(manager) - try: - response = await client.get("/static/request_session.css") - body = await response.text() - assert response.status == 200 - assert ".idleCanvas" in body - assert "body.has-video .idleCanvas" in body - finally: - await client.close() - - -@pytest.mark.asyncio -async def test_offer_returns_answer_payload() -> None: - manager = FakeSessionManager() - client = await _build_client(manager) - try: - response = await client.post( - "/api/webrtc/offer", - json={"sdp": "offer-sdp", "type": "offer"}, - ) - payload = await response.json() - assert response.status == 200 - assert payload == manager.answer_payload - assert manager.offers == [("offer-sdp", "offer")] - finally: - await client.close() - - -@pytest.mark.asyncio -async def test_offer_busy_returns_conflict() -> None: - manager = FakeSessionManager() - manager.raise_busy = True - client = await _build_client(manager) - try: - response = await client.post( - "/api/webrtc/offer", - json={"sdp": "offer-sdp", "type": "offer"}, - ) - assert response.status == 409 - finally: - await client.close() - - -def test_configure_logging_suppresses_ice_info_spam() -> None: - configure_logging() - - assert logging.getLogger("aioice").getEffectiveLevel() == logging.WARNING - assert logging.getLogger("aioice.ice").getEffectiveLevel() == logging.WARNING - assert logging.getLogger("aiortc").getEffectiveLevel() == logging.WARNING diff --git a/integrations/self_forcing/self_forcing/runner.py b/integrations/self_forcing/self_forcing/runner.py index ba97c1440..3e4857c7f 100644 --- a/integrations/self_forcing/self_forcing/runner.py +++ b/integrations/self_forcing/self_forcing/runner.py @@ -34,6 +34,7 @@ WanInferencePipeline, WanInferencePipelineCache, ) +from flashdreams.runtime.video_output import Mp4VideoOutputTarget __all__ = [ "SelfForcingT2VRunnerConfig", @@ -125,25 +126,46 @@ def run(self) -> None: # Generate the autoregressive chunks. output_stream = self.create_video_output_stream(fps=config.fps) + video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") + output_target = Mp4VideoOutputTarget( + output_path=video_path, + fps=config.fps, + output_layout=output_stream.output_layout, + enabled=self.is_rank_zero, + ) + output_target.open() for i in range(config.total_blocks): video_chunk = self.pipeline.generate(autoregressive_index=i, cache=cache) stats = self.pipeline.finalize(autoregressive_index=i, cache=cache) - output_stream.process(video_chunk, autoregressive_index=i, stats=stats) + output_target.write( + output_stream.process( + video_chunk, + autoregressive_index=i, + metrics=stats, + ) + ) - video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - video_path = output_stream.finish_to_mp4(video_path, fps=config.fps) - if video_path is None: + tail = output_stream.finish() + if tail is not None: + output_target.write(tail) + artifacts = output_target.close() + if not artifacts: return + video_artifact = artifacts[0] + video_path = Path(video_artifact.uri) logger.info( - f"[{config.runner_name}] wrote video {tuple(generated.shape)} " + f"[{config.runner_name}] wrote video {video_artifact.metadata['shape']} " f"-> {video_path.resolve()}" ) # Write the perf stats. - if output_stream.stats_history: + stats_history = video_artifact.metadata["stats_history"] + if stats_history: stats_path = write_runner_stats( - config.output_dir, config.runner_name, output_stream.stats_history + config.output_dir, + config.runner_name, + list(stats_history), ) logger.info( f"[{config.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" diff --git a/integrations/wan21/wan21/runner.py b/integrations/wan21/wan21/runner.py index 33e2e4585..2098c746b 100644 --- a/integrations/wan21/wan21/runner.py +++ b/integrations/wan21/wan21/runner.py @@ -39,6 +39,7 @@ WanInferencePipeline, WanInferencePipelineCache, ) +from flashdreams.runtime.video_output import Mp4VideoOutputTarget __all__ = [ "Wan21I2VRunnerConfig", @@ -171,13 +172,27 @@ def run(self) -> None: # Generate the output in one AR step. output_stream = self.create_video_output_stream(fps=config.fps) + video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") + output_target = Mp4VideoOutputTarget( + output_path=video_path, + fps=config.fps, + output_layout=output_stream.output_layout, + enabled=self.is_rank_zero, + ) + output_target.open() generated = self.pipeline.generate(autoregressive_index=0, cache=cache) stats = self.pipeline.finalize(autoregressive_index=0, cache=cache) - output_stream.process(generated, autoregressive_index=0, stats=stats) - video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - video_path = output_stream.finish_to_mp4(video_path, fps=config.fps) - if video_path is None: + output_target.write( + output_stream.process(generated, autoregressive_index=0, metrics=stats) + ) + tail = output_stream.finish() + if tail is not None: + output_target.write(tail) + artifacts = output_target.close() + if not artifacts: return + video_artifact = artifacts[0] + video_path = Path(video_artifact.uri) logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " @@ -185,11 +200,12 @@ def run(self) -> None: ) # Write the perf stats. - if output_stream.stats_history: + stats_history = video_artifact.metadata["stats_history"] + if stats_history: stats_path = write_runner_stats( config.output_dir, config.runner_name, - output_stream.stats_history, + list(stats_history), ) logger.info( f"[{config.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}"