From eeda7e99f2a4e9a530cec3b33dd4f12f155d6f4d Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 30 Jul 2026 21:25:10 +0000 Subject: [PATCH 01/19] Add inference runtime API design proposal --- docs/inference_runtime_api_design.md | 640 +++++++++++++++++++++++++++ 1 file changed, 640 insertions(+) create mode 100644 docs/inference_runtime_api_design.md diff --git a/docs/inference_runtime_api_design.md b/docs/inference_runtime_api_design.md new file mode 100644 index 000000000..6a4eea9dc --- /dev/null +++ b/docs/inference_runtime_api_design.md @@ -0,0 +1,640 @@ + + +# FlashDreams Inference Runtime API Design Proposal + +Date: July 30, 2026 + +## Summary + +This proposal defines a standard inference runtime API for FlashDreams +integrations. The goal is to make world-model integrations easier to build, +benchmark, and run without forcing every model into the same input shape or +optimization stack. + +The proposed API separates the pieces that are currently mixed together in +integration-specific runner code: + +- `InferenceConfig`: how the model and inference stack should run; +- `UserInputs`: controls or events from an app, replay trace, or benchmark; +- `ModelInputs`: prompts, frames, videos, trajectories, maps, scene data, and + other values required by a specific model; +- input mapping: model/application-specific conversion from user-facing inputs + into model-facing inputs; +- runtime/session execution: model setup, warmup, per-rollout state, and + stepping; +- output targets: WebRTC, native display, MP4, benchmark artifacts, or headless + runs; +- metrics/profiling: timings, memory, traces, NVTX ranges, and benchmark + outputs. + +The API should standardize the envelope and lifecycle. It should not pretend +that all world models have the same inputs, that all models use the same +optimization stack, or that a raw checkpoint can fully describe how to run the +model. + +## Current Implementation Plan + +Implementation should happen on an experimental integration branch. PRs for this +work should target that branch until the API shape, LingBot migration, and +OmniDreams migration are all working well enough to merge to `main` together. + +The experimental branch can temporarily break or simplify command-line options +while the demos are being moved to the new API. The required outcome is that the +LingBot and OmniDreams demos still run through the new runtime path, and that +benchmark tooling can confirm they are at least broadly healthy before the +branch is merged back to `main`. + +Initial scope: + +- define the minimal runtime API envelope; +- migrate LingBot and OmniDreams to use it; +- support selectable output modes such as MP4, JPEG/MJPEG stream, WebRTC, and + headless/null where appropriate; +- use or update benchmark tooling to verify the migrated demos; +- defer broader model migrations, hosted execution, full autotune, and polished + metrics until the first branch proves the API shape. + +## Task Tracker + +| ID | Workstream | Can run in parallel? | Depends on | Done when | +| --- | --- | --- | --- | --- | +| T0 | Create experimental branch and contribution rules. | No, this starts the work. | None. | Branch exists, PR target is agreed, and main merge criteria are written down. | +| T1 | Minimal API envelope and naming. | Partly. | T0. | `InferenceConfig`, `UserInputs`, `ModelInputs`, runtime/session, output target, and mapping boundaries are defined well enough for demos to use. | +| T2 | Event-based `UserInputs`. | Yes, after T1 direction is agreed. | T1. | User inputs are primarily timestamped events; replay traces and derived snapshots are supported where needed. | +| T3 | `ModelInputs`, schemas, and mapping boundary. | Yes, after T1 direction is agreed. | T1. | Models can declare required initial/per-step inputs, and mappings can convert user events into model inputs. | +| T4 | `ModelRunner`, `InferenceRuntime`, and `InferenceSession` skeleton. | Partly. | T1. | A minimal standard loop can initialize a runtime, run at least one sequential session, and close cleanly. | +| T5 | Output mode selection. | Yes, after the result/output shape is agreed. | T1, T4. | A run can choose output behavior such as MP4, JPEG/MJPEG stream, WebRTC, benchmark artifact, or headless/null without changing model code. | +| T6 | LingBot migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | LingBot runs through the new API path with its event inputs mapped into model inputs. | +| T7 | OmniDreams migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | OmniDreams runs through the new API path with its model-specific inputs and mapping preserved. | +| T8 | Benchmark/smoke verification for LingBot and OmniDreams. | Preparation can run early; final gate is late. | T5, T6, T7. | Existing or updated benchmark tooling can run both migrated demos and produce enough evidence that they still work. | +| T9 | Metrics and profiling normalization for the branch. | Yes, but final integration is late. | T4, T5, T8. | Basic canonical metrics are emitted for migrated demos; deeper metrics can remain follow-up work. | +| T10 | CLI compatibility and migration cleanup. | Yes, after demo migrations start. | T6, T7. | Required demo commands are restored or replaced, temporary hacks are removed, and user-facing docs/notes match the branch behavior. | +| T11 | Stabilize and merge experimental branch to `main`. | No, final integration step. | T6-T10. | LingBot and OmniDreams pass agreed smoke/benchmark checks, review feedback is addressed, and the branch can merge as one API transition. | + +Suggested parallel split: + +- one person owns T1/T4, because the API envelope and standard loop are the + critical path; +- one person owns T2/T3, because event inputs, schemas, and mapping need to + stay coherent; +- one person owns T5/T8/T9, because outputs, benchmarks, and metrics are tightly + related; +- LingBot and OmniDreams can be assigned separately once the skeleton is usable; +- one person should track branch health, CLI compatibility, and merge readiness. + +## Architecture + +```text +Optional discovery for CLI, benchmark, hosted, or installed-package flows: + Model/preset registry + -> adapter/preset/default setup/scenario metadata + -> contributes defaults to the app-supplied run setup + +Main runtime flow: +App / integration / benchmark / transport + chooses how the run is driven and where output goes + supplies run setup: + InferenceConfig + UserInputs + ModelInputs + output/metrics options + | + v +ModelRunner / standard loop + orchestrates validation, lifecycle, stepping, output, and metrics + uses input mapping to: + validate that user/app inputs can drive the model + build initial and per-step ModelInputs during the run + | + v +InferenceRuntime + reusable heavyweight lifecycle: distributed init, model load, compile, warmup + load once; create sessions sequentially unless the backend supports concurrency + | + v +InferenceSession + one rollout/stream: prompt/initial inputs, cache/state, current step, reset + keeps per-run state from leaking across prompts, clients, or benchmark repeats + | + v +Model implementation / inference pipeline + hot path: encode -> model step -> decode -> cache/finalize + | + v +Output target + WebRTC | native window | MP4 | benchmark | headless/null + | + v +Metrics / artifacts / logs / reports / traces +``` + +## Example Sequential Session Flow + +The runtime/session split is primarily about reusing expensive model setup while +keeping each rollout's state isolated. The default mental model should be +sequential sessions, not required concurrent sessions. + +```text +ModelRunner / standard loop + | + v +Create InferenceRuntime from InferenceConfig + load checkpoint/model + initialize distributed/backend state + compile/capture/warm up if configured + | + v +Start InferenceSession A + initial ModelInputs: prompt/frame/scene/etc. + per-session state: cache, current step, reset state + step 0 -> step 1 -> ... -> done + outputs -> Output target + metrics -> Metrics recorder + close session A + | + v +Start InferenceSession B + new initial ModelInputs or replay scenario + independent cache/state + step 0 -> step 1 -> ... -> done + outputs -> Output target + metrics -> Metrics recorder + close session B + | + v +Close InferenceRuntime + release model/backend resources +``` + +For v0, an `InferenceRuntime` may support only one active session at a time. +Concurrent sessions should be treated as an optional backend/model capability, +not a baseline API requirement. + +`StreamInferencePipeline` should remain an important local implementation path +for models that already use it, but it should not be treated as the only +possible model boundary. A session may call `StreamInferencePipeline`, another +local model implementation, a Dynamo-like backend, or a hosted service. + +## System Components + +| Component | Role | Boundary | +| --- | --- | --- | +| Model/preset registry | Lists what can run: model/preset slugs, scenarios, capabilities, resource hints, and supported output modes. | Must remain cheap to query and must not load checkpoints. | +| App / integration / benchmark / transport | Owns the user-facing mode: CLI, native integration, WebRTC, benchmark, hosted request, or replay. | Supplies run setup, user inputs, model inputs, and output target selection. | +| User input library | Normalizes live or replayed controls into FlashDreams-supported user input events/windows. | Shared primitives for keyboard, reset, prompt/image updates, traces, and future scalar controls. | +| Input mapping | Converts user/app inputs plus initial model inputs into the model-specific inputs needed by the session. | Owned by model/application code; may be a no-op for simple runs. | +| ModelRunner / standard loop | Orchestrates one run from setup through runtime initialization, stepping, output, metrics, and teardown. | Shared orchestration layer used by CLIs, benchmarks, MP4 runs, and simple realtime flows. | +| InferenceRuntime | Owns heavyweight lifecycle: distributed init, model construction, checkpoint loading, compile/capture, warmup, hosted-service connection, and teardown. | Long-lived reusable runtime created from `InferenceConfig`; lets FlashDreams load/warm once and create sessions sequentially unless the backend supports concurrency. | +| InferenceSession | Owns one rollout or stream: initial inputs, cache state, current step, reset behavior, step requirements, and step execution. | Per-rollout interface consumed by the standard loop; keeps state isolated across prompts, browser clients, replay scenarios, or benchmark repeats. | +| Model implementation / inference pipeline | Implements encode, model step, decode, cache updates, and model-specific optimizations. | FlashDreams wraps this boundary; it should not replace every model implementation. | +| Output target | Consumes generated outputs and handles presentation or persistence. | Separate from model execution so the same session can feed WebRTC, MP4, benchmark, or headless output. | +| Metrics, artifacts, and profiling | Records timings, memory, quality data, logs, reports, traces, and optional NVTX ranges. | Shared observation layer for local runs, benchmarks, CI smoke, and hosted runs. | + +## API Layers + +FlashDreams should expose layered APIs rather than a single all-or-nothing +interface: + +```text +High-level runtime API + run setup -> standard loop -> output targets -> metrics/artifacts + +Adapter/runtime API + model adapter -> InferenceRuntime -> InferenceSession + +Low-level inference API + StreamInferencePipeline -> encoders/decoders -> cache/perf/profiling helpers +``` + +| Layer | Intended user | Provides | +| --- | --- | --- | +| High-level runtime API | Users who want FlashDreams to own the run loop. | Run setup, input mapping, runtime/session lifecycle, output targets, metrics, profiling, and benchmark artifacts. | +| Adapter/runtime API | Model owners who want their model to plug into the standard loop. | Model adapter, input requirements, runtime/session implementation, and model-specific mapping or validation. | +| Low-level inference API | Users who want to own their own loop while reusing FlashDreams building blocks. | `StreamInferencePipeline`, encoders, decoders, cache helpers, profiling tools, and optimization utilities. | + +These layers should remain compatible. The new runtime API sits above the +existing lower-level pieces; it does not replace them. + +## Goals + +- Make FlashDreams easier to use for new world-model integrations. +- Keep model-specific input semantics explicit instead of hiding them in runner + code. +- Avoid a single monolithic inference stack; different models should be able to + validate and use different optimization features. +- Separate model execution from presentation and persistence. +- Support both live input and deterministic replay through the same + runtime/session boundary. +- Make metrics, benchmark artifacts, and profiling first-class without forcing + profiling overhead into normal runs. +- Preserve room for local single-GPU, local distributed, Dynamo-like, and hosted + execution. + +## Non-Goals + +- Do not infer arbitrary model semantics from a raw checkpoint. +- Do not require every model to use the same encoder, decoder, scheduler, + control representation, transport, or optimization set. +- Do not make WebRTC or native display part of the model API. +- Do not make autotuning part of normal inference startup. +- Do not require users to use the high-level standard loop when they only need + lower-level inference building blocks. +- Do not require every existing integration to migrate in one large change. + +## API Placement + +The new API should sit above the existing `flashdreams.infra` layer. Existing +pipelines, encoders, decoders, runner configs, realtime input helpers, WebRTC +code, and quality/benchmark utilities should be reused where possible. + +The exact package layout and class definitions can be decided during +implementation. This document should define responsibilities and boundaries, not +the final Python shape. + +## InferenceConfig + +`InferenceConfig` describes how to run the model/runtime. It should cover: + +- model or preset identity; +- checkpoint or model asset selection; +- execution backend, such as local single GPU, local multi-GPU, Dynamo-like, or + hosted/external execution; +- device placement, precision, and resource hints; +- optimization choices such as compile, CUDA graph capture, attention backend, + cache policy, overlap, prefetch, and native extensions; +- runtime-affecting profiling or tracing options. + +It should not contain prompts, keyboard state, browser settings, MP4 paths, +benchmark output directories, or other app/output settings. Those belong in the +run setup around `InferenceConfig`. + +Existing `StreamInferencePipelineConfig` and `InstantiateConfig` style configs +can remain valid model references behind this layer. The model adapter should +validate which execution and optimization choices are supported. Unsupported +choices should fail clearly or be explicitly handled only when the user selected +an automatic mode. + +## UserInputs + +`UserInputs` describes user-facing controls produced by a live UI, browser, +native app, replay trace, synthetic benchmark driver, or no-op source. + +User inputs should primarily be represented as timestamped events. This gives +live apps, replay traces, and benchmarks the same basic shape, and lets +FlashDreams resample or window those events when a model session asks for the +next chunk of inputs. + +Initial supported user input types should stay close to what FlashDreams already +uses: + +- keyboard keydown/keyup events; +- reset requests; +- prompt update requests; +- image update requests; +- future scalar controls such as throttle, brake, steer, or camera axes once an + integration needs them. + +Snapshot-style inputs, such as current key state, can still be supported when +useful. They should be treated as a derived or compatibility form rather than +the primary user-input abstraction. + +User inputs are not model inputs. A keyboard event does not have one universal +meaning. One model may map it to pose segments, another to steering commands, +and another may ignore it. + +## ModelInputs + +`ModelInputs` describes the data the model or inference pipeline actually +requires. It should distinguish: + +- initial inputs: values needed to start or reset a rollout; +- per-step inputs: values needed for one generated chunk or frame window. + +Examples of initial model inputs include prompt, negative prompt, first frame, +input video, scene id, HD map asset, camera calibration, initial camera pose, +seed, or model-specific fields. + +Examples of per-step model inputs include frame timestamps, pose segments, +camera trajectory chunks, rendered HD map frames, conditioning video windows, +control tensors, event markers, or model-specific fields. + +Model input payloads should use semantic names, not only modality names. For +example, a first frame and an HD map frame should be distinct inputs even if +both are image-like values. + +For interactive runs, most `ModelInputs` will be initial values plus per-step +inputs produced by input mapping. For MP4 generation and benchmarking, the API +should also support fixed per-step model inputs so runs can be deterministic. + +## Schemas + +The API should support lightweight `UserInputSchema` and `ModelInputSchema` +metadata. + +These schemas are not meant to be a rich type system or a replacement for +model-specific validation. They should be just enough to answer: + +- what can this app, transport, trace, or benchmark source provide? +- what does this model require before startup and at each step? +- can this event source drive this model with the selected mapping? + +The purpose is to fail early before expensive model initialization, produce +clearer errors, make fixed scenarios easier to validate, and avoid ambiguous +dict payloads where keys only describe modality. + +For simple CLI text-to-video or image-to-video runs, `UserInputSchema` can be +trivial or omitted because there may be no live controls. `ModelInputSchema` is +more important because each supported model still needs to declare the +model-facing values it expects. + +## Model Requirements + +A raw checkpoint should not be treated as self-describing. It may imply tensor +shapes or architecture details, but it usually does not fully define: + +- required semantic inputs; +- initial versus per-step inputs; +- units for timestamps, poses, or calibration values; +- how user controls become model controls; +- preprocessing, encoder, decoder, mask, prompt, or cache rules. + +Therefore, a FlashDreams-supported model should have an adapter or integration +layer that declares its model input requirements and prepares inputs for the +underlying model implementation. + +Users running an existing FlashDreams-supported model should not need to write +that adapter. Developers bringing a new world model to FlashDreams should expect +to provide one. + +## External Model Usage + +Users should be able to run their own models without adding those models to the +FlashDreams repository. The flow depends on which API layer they use: + +```text +High-level runtime API + user supplies or installs model adapter + FlashDreams owns standard loop, outputs, metrics, benchmarks + +Adapter/runtime API + model owner implements adapter/runtime/session + adapter can be passed directly or registered by an installed package + +Low-level inference API + user owns loop and lifecycle + user reuses pipeline, encoder/decoder, cache, profiling, or optimization tools +``` + +| Flow | Registry needed? | Who provides model-specific code? | Result | +| --- | --- | --- | --- | +| Direct Python | No. | User or model owner passes an adapter/setup directly. | FlashDreams can run the standard loop without the model living in the repo. | +| Installed package | Yes, for discovery. | External or internal package registers adapters/presets. | CLIs, benchmarks, and hosted schedulers can discover the model cheaply. | +| Low-level only | No. | User owns the loop and calls lower-level FlashDreams pieces directly. | Useful when the user wants optimizations or pipeline helpers but not the standard loop. | + +The model adapter is a role/boundary, not necessarily a concrete class. It is +the model-specific code that declares input requirements, validates supported +configs, creates the runtime/session, and connects FlashDreams to the actual +model implementation. + +The registry should not be treated as a central FlashDreams-owned catalog of all +possible models. It is a discovery mechanism for installed adapters. Built-in +public integrations, internal GitLab-only integrations, and third-party packages +can all participate through the same mechanism. + +FlashDreams should not claim to run an arbitrary checkpoint with no adapter +unless the checkpoint already matches a supported generic adapter. + +## Input Mapping + +Input mapping is required whenever `UserInputs` need to become per-step +`ModelInputs`. The exact implementation does not need to be a required top-level +object. It could be: + +- a method on the model adapter; +- a method on an app/runtime adapter; +- a separate mapper object; +- a default no-op or identity mapping for simple T2V/I2V/fixed-input runs. + +There are two separate moments to keep clear: + +- before runtime initialization, FlashDreams should select the mapping and check + obvious compatibility between the app event source and the model; +- during the standard loop, the runner uses the mapping to build initial or + per-step `ModelInputs` from the relevant event window, often after the session + reports what it needs next. + +Examples: + +- T2V mapping validates a prompt and creates no per-step control inputs. +- I2V mapping validates a prompt plus first frame and creates no live controls. +- A keyboard-driven integration maps key events or event windows into pose + segments or steering controls. +- OmniDreams-like integrations may map driving commands into camera poses, HD + map frames, and dynamic actor state. +- Benchmark mapping can read fixed event traces and produce identical step + inputs each run. + +The compatibility check should be treated as early validation, not a guarantee +that the run will succeed. It can catch obvious mismatches, but the model +adapter/runtime still owns deep tensor validation and model semantics. + +## Runtime And Standard Loop + +The standard loop should be shared by CLI generation, headless playback, MP4 +generation, benchmarks, and simple realtime applications. + +A run should: + +1. Discover the model or preset without loading checkpoints. +2. Resolve inference config, user inputs, model inputs, output target, metrics, + profiling, and optional scenario setup. +3. Validate that the event source and mapping can drive the selected model. +4. Initialize the runtime. +5. Start a session from initial model inputs. +6. For each step, ask the session what it needs, gather live or fixed inputs, + build step model inputs, run the session step, route outputs, and record + metrics. +7. Finalize output artifacts, metrics, logs, reports, and traces. + +Realtime transports may need an async variant, backpressure, and explicit flow +control, but the conceptual boundary should remain the same: event/input source, +input mapping, session, output target, metrics. + +The session should expose what it needs for the next step rather than requiring +the app or output layer to guess. This matters because AR step 0 can differ from +steady-state steps, and encoder/decoder temporal compression can produce +different input and output frame windows. + +## Output Targets + +Output handling should be separate from model execution. The model session +returns generated outputs and metadata; the output target decides what to do +with them. + +Expected output targets include: + +- WebRTC streaming; +- native window display; +- MJPEG or lightweight remote preview; +- MP4 writing; +- benchmark artifact writing; +- headless playback; +- null output for pure throughput measurements. + +Display and transport can still affect measured performance through copies, +encoding, queueing, backpressure, and presentation timing. Those costs should be +measured as output-target or end-to-end metrics instead of being mixed into core +model-stage timings. + +## Fixed Inputs, Benchmarks + +The API should support fixed runs as a first-class case. This is needed for MP4 +generation, benchmarks, regression testing, and autotune. + +Two replay levels should be supported: + +- user-event replay: records timestamped key events, prompt updates, image + updates, reset events, and timing, then runs normal input mapping; +- model-input replay: records or defines already-mapped per-step model inputs + for stricter model-level regression tests. + +User-event replay tests more of the application stack. Model-input replay is +better for isolating model runtime performance and reproducibility. + +## Metrics And Profiling + +Metrics should have a small canonical baseline plus optional extras. + +The baseline should cover: + +- lifecycle timing: startup, load, warmup, first-step latency; +- model-stage timing: encode, model step, decode, finalize/cache update; +- memory: allocated, reserved, peak, and per-rank where applicable; +- throughput: frames per second, chunks per second, real-time factor. + +Realtime runs may add input-to-present latency, jitter, missed deadlines, queue +depth, dropped frames, WebRTC stats, encoder bitrate, and client stats. +Benchmark runs may add quality metrics, logs, MP4/image previews, and reports. + +Persisted timing metrics should use seconds as the canonical unit because +seconds compose cleanly across Python timers, traces, and long-running +durations. Reports and UIs can display milliseconds for short latencies. + +Profiling should be optional and controlled separately from normal metrics. +NVTX ranges should be supported for Nsight profiling, but profiling should not +be required for normal inference or benchmark runs. + +## Autotune + +Autotune should be a separate harness that evaluates candidate +`InferenceConfig` variants against fixed scenarios. It should not be part of +normal startup. + +Autotune may search over compile, CUDA graph capture, attention backend, +precision, cache policy, overlap, prefetch, native extensions, and chunk size +when the model supports those knobs. + +Results are only valid for a specific model, checkpoint, hardware, driver, +FlashDreams commit, and scenario. First-run compile/capture cost should be +separated from steady-state metrics. Agent assistance could help propose search +spaces or summarize results, but the measured selection process should be +deterministic code. + +## Distributed And Hosted Execution + +The API should leave room for local single-GPU, local multi-GPU, Dynamo-like +execution, and hosted execution such as a Reactor-style platform. + +At this stage, the proposal should not define Reactor- or Dynamo-specific +contracts in detail. It should preserve the right boundary: execution backend +selection belongs in `InferenceConfig`, while backend-specific scheduling, +authentication, asset access, output streaming, artifact handling, and failure +behavior belong behind the runtime/backend implementation. + +The practical order should be local first, then local distributed, then +hosted/distributed backends once concrete backend owners can validate the +requirements. + +## Existing Code And Migration + +The new API should reuse existing code instead of replacing everything: + +- keep `flashdreams.infra.pipeline` as the common local encode/model/decode + implementation path; +- keep existing encoder and decoder contracts and reuse temporal size helpers; +- keep existing runner configs and CLI compatibility during migration; +- reuse `KeyboardResampler` and realtime input helpers behind the new input + boundary; +- treat WebRTC as a transport/output adapter and bridge it gradually; +- reuse existing quality and benchmark utilities where applicable; +- keep internal-only integrations registered only in the GitLab/internal + workspace. + +The task tracker near the start of this document is the source of truth for the +first implementation branch. The first milestone is intentionally narrower than +the full design: prove the API with LingBot and OmniDreams, selectable output +modes, and enough benchmark/smoke coverage to merge the experimental branch +back to `main` safely. + +## Design Risks + +- `InferenceConfig` could become too broad if prompts, controls, output paths, + browser settings, and benchmark settings are added to it. Keep it focused on + model/runtime execution. +- Dict-like model inputs are flexible but can fail late. Keep dict payloads for + flexibility, but require lightweight schemas and adapter validation for + supported models. +- Schemas could become too heavy. Keep them minimal and role-oriented. +- User inputs are not model inputs. Keep input mapping explicit and + model/application-owned. +- Per-frame, per-chunk, and AR-step clocks are easy to confuse. The session + should expose step requirements instead of making app code guess. +- Output separation is necessary but not free. Measure output and transport + costs separately from core model timings. +- Hosted/distributed execution is still under-specified. Keep the API boundary + open until backend owners validate concrete requirements. +- Existing WebRTC behavior is nontrivial. Bridge it gradually to avoid + regressions. +- Public/internal boundaries must remain clean. Internal adapters, slugs, and + scenarios should not leak into the public repo. + +## Decisions To Make Before Implementation + +- What should the top-level package/API be called? +- Should the main registered object be called an adapter, integration, runtime + factory, or something else? +- What direct-Python API should let users pass an external adapter without + registering it? +- What package registration mechanism should third-party and internal adapters + use for CLI discovery and benchmarks? +- How lightweight should `UserInputSchema` and `ModelInputSchema` be? +- Where should input mapping live: model adapter, app adapter, separate object, + or a mix? +- What should the output abstraction be called? +- What is the minimum v0 set of supported user input events? +- What is the first public model to migrate? +- What metrics are required for every benchmark run? +- What metadata must be discoverable without loading checkpoints? +- What requirements do Dynamo/Reactor-style backends need before we commit to + hosted execution details? + +The document currently uses "integration" for model-specific packages and app +entrypoints. If the team prefers "model" as the public term, that can be changed +later without changing the architecture. + +## Recommendation + +Proceed with the proposed split: + +- `InferenceConfig` for model/runtime execution; +- `UserInputs` for app-facing controls and replay traces; +- `ModelInputs` for model-facing initial and per-step inputs; +- input mapping for model/application-specific conversion; +- runtime/session boundaries for lifecycle and stepping; +- output targets for display, streaming, files, and benchmarks; +- shared metrics and optional profiling. + +The main constraint is that arbitrary world-model inputs cannot be standardized +away. FlashDreams can provide the shared envelope, loop, metrics, replay, and +output tools, but each supported model still needs an adapter that declares and +validates its own input contract. From b6b7d557fb72be072bfb6f2f619cff316253d205 Mon Sep 17 00:00:00 2001 From: jarcherNV Date: Tue, 4 Aug 2026 02:11:54 -0700 Subject: [PATCH 02/19] Add experimental inference runtime API envelope (#403) Define the initial flashdreams.runtime package with minimal T1 boundaries for runtime config, user/model inputs, schemas, input mapping, model adapters, runtime/session protocols, output targets, and metrics. Add focused CPU tests for the new API surface without migrating existing runners. --- docs/inference_runtime_api_design.md | 112 ++-- flashdreams/flashdreams/runtime/__init__.py | 60 +++ flashdreams/flashdreams/runtime/_utils.py | 17 + flashdreams/flashdreams/runtime/config.py | 76 +++ flashdreams/flashdreams/runtime/inputs.py | 200 +++++++ flashdreams/flashdreams/runtime/interfaces.py | 90 ++++ flashdreams/flashdreams/runtime/mapping.py | 85 +++ flashdreams/flashdreams/runtime/metrics.py | 124 +++++ flashdreams/flashdreams/runtime/output.py | 78 +++ flashdreams/flashdreams/runtime/types.py | 56 ++ .../tests/test_inference_runtime_api.py | 492 ++++++++++++++++++ 11 files changed, 1350 insertions(+), 40 deletions(-) create mode 100644 flashdreams/flashdreams/runtime/__init__.py create mode 100644 flashdreams/flashdreams/runtime/_utils.py create mode 100644 flashdreams/flashdreams/runtime/config.py create mode 100644 flashdreams/flashdreams/runtime/inputs.py create mode 100644 flashdreams/flashdreams/runtime/interfaces.py create mode 100644 flashdreams/flashdreams/runtime/mapping.py create mode 100644 flashdreams/flashdreams/runtime/metrics.py create mode 100644 flashdreams/flashdreams/runtime/output.py create mode 100644 flashdreams/flashdreams/runtime/types.py create mode 100644 flashdreams/tests/test_inference_runtime_api.py diff --git a/docs/inference_runtime_api_design.md b/docs/inference_runtime_api_design.md index 6a4eea9dc..2f0ba19f8 100644 --- a/docs/inference_runtime_api_design.md +++ b/docs/inference_runtime_api_design.md @@ -59,25 +59,25 @@ Initial scope: ## Task Tracker -| ID | Workstream | Can run in parallel? | Depends on | Done when | -| --- | --- | --- | --- | --- | -| T0 | Create experimental branch and contribution rules. | No, this starts the work. | None. | Branch exists, PR target is agreed, and main merge criteria are written down. | -| T1 | Minimal API envelope and naming. | Partly. | T0. | `InferenceConfig`, `UserInputs`, `ModelInputs`, runtime/session, output target, and mapping boundaries are defined well enough for demos to use. | -| T2 | Event-based `UserInputs`. | Yes, after T1 direction is agreed. | T1. | User inputs are primarily timestamped events; replay traces and derived snapshots are supported where needed. | -| T3 | `ModelInputs`, schemas, and mapping boundary. | Yes, after T1 direction is agreed. | T1. | Models can declare required initial/per-step inputs, and mappings can convert user events into model inputs. | -| T4 | `ModelRunner`, `InferenceRuntime`, and `InferenceSession` skeleton. | Partly. | T1. | A minimal standard loop can initialize a runtime, run at least one sequential session, and close cleanly. | -| T5 | Output mode selection. | Yes, after the result/output shape is agreed. | T1, T4. | A run can choose output behavior such as MP4, JPEG/MJPEG stream, WebRTC, benchmark artifact, or headless/null without changing model code. | -| T6 | LingBot migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | LingBot runs through the new API path with its event inputs mapped into model inputs. | -| T7 | OmniDreams migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | OmniDreams runs through the new API path with its model-specific inputs and mapping preserved. | -| T8 | Benchmark/smoke verification for LingBot and OmniDreams. | Preparation can run early; final gate is late. | T5, T6, T7. | Existing or updated benchmark tooling can run both migrated demos and produce enough evidence that they still work. | -| T9 | Metrics and profiling normalization for the branch. | Yes, but final integration is late. | T4, T5, T8. | Basic canonical metrics are emitted for migrated demos; deeper metrics can remain follow-up work. | -| T10 | CLI compatibility and migration cleanup. | Yes, after demo migrations start. | T6, T7. | Required demo commands are restored or replaced, temporary hacks are removed, and user-facing docs/notes match the branch behavior. | -| T11 | Stabilize and merge experimental branch to `main`. | No, final integration step. | T6-T10. | LingBot and OmniDreams pass agreed smoke/benchmark checks, review feedback is addressed, and the branch can merge as one API transition. | +| ID | Status | Workstream | Can run in parallel? | Depends on | Done when | +| --- | --- | --- | --- | --- | --- | +| T0 | Complete | Create experimental branch and contribution rules. | No, this starts the work. | None. | Branch exists, PR target is agreed, and main merge criteria are written down. | +| T1 | Complete | Minimal API envelope and naming. | Partly. | T0. | `InferenceConfig`, `UserInputs`, `ModelInputs`, runtime/session, output target, and mapping boundaries are defined well enough for demos to use. | +| T2 | Planned | Event-based `UserInputs`. | Yes, after T1 direction is agreed. | T1. | User inputs are primarily timestamped events; replay traces and derived snapshots are supported where needed. | +| T3 | Planned | `ModelInputs`, schemas, and mapping boundary. | Yes, after T1 direction is agreed. | T1. | Models can declare required initial/per-step inputs, and mappings can convert user events into model inputs. | +| T4 | Planned | `ModelRunner`, `InferenceRuntime`, and `InferenceSession` skeleton. | Partly. | T1. | A minimal standard loop can initialize a runtime, run at least one sequential session, and close cleanly. | +| T5 | Planned | Output mode selection. | Yes, after the result/output shape is agreed. | T1, T4. | A run can choose output behavior such as MP4, JPEG/MJPEG stream, WebRTC, benchmark artifact, or headless/null without changing model code. | +| T6 | Planned | LingBot migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | LingBot runs through the new API path with its event inputs mapped into model inputs. | +| T7 | Planned | OmniDreams migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | OmniDreams runs through the new API path with its model-specific inputs and mapping preserved. | +| T8 | Planned | Benchmark/smoke verification for LingBot and OmniDreams. | Preparation can run early; final gate is late. | T5, T6, T7. | Existing or updated benchmark tooling can run both migrated demos and produce enough evidence that they still work. | +| T9 | Planned | Metrics and profiling normalization for the branch. | Yes, but final integration is late. | T4, T5, T8. | Basic canonical metrics are emitted for migrated demos; deeper metrics can remain follow-up work. | +| T10 | Planned | CLI compatibility and migration cleanup. | Yes, after demo migrations start. | T6, T7. | Required demo commands are restored or replaced, temporary hacks are removed, and user-facing docs/notes match the branch behavior. | +| T11 | Planned | Stabilize and merge experimental branch to `main`. | No, final integration step. | T6-T10. | LingBot and OmniDreams pass agreed smoke/benchmark checks, review feedback is addressed, and the branch can merge as one API transition. | Suggested parallel split: -- one person owns T1/T4, because the API envelope and standard loop are the - critical path; +- one person owns T4 and keeps it aligned with the completed T1 envelope, + because the standard loop is now the critical path; - one person owns T2/T3, because event inputs, schemas, and mapping need to stay coherent; - one person owns T5/T8/T9, because outputs, benchmarks, and metrics are tightly @@ -182,7 +182,7 @@ local model implementation, a Dynamo-like backend, or a hosted service. | Model/preset registry | Lists what can run: model/preset slugs, scenarios, capabilities, resource hints, and supported output modes. | Must remain cheap to query and must not load checkpoints. | | App / integration / benchmark / transport | Owns the user-facing mode: CLI, native integration, WebRTC, benchmark, hosted request, or replay. | Supplies run setup, user inputs, model inputs, and output target selection. | | User input library | Normalizes live or replayed controls into FlashDreams-supported user input events/windows. | Shared primitives for keyboard, reset, prompt/image updates, traces, and future scalar controls. | -| Input mapping | Converts user/app inputs plus initial model inputs into the model-specific inputs needed by the session. | Owned by model/application code; may be a no-op for simple runs. | +| Input mapping | Converts user/app inputs plus initial model inputs into the model-specific inputs needed by the session. | A model adapter may provide a default mapping; runtimes, applications, benchmarks, and replay tools may override it without changing the model step. | | ModelRunner / standard loop | Orchestrates one run from setup through runtime initialization, stepping, output, metrics, and teardown. | Shared orchestration layer used by CLIs, benchmarks, MP4 runs, and simple realtime flows. | | InferenceRuntime | Owns heavyweight lifecycle: distributed init, model construction, checkpoint loading, compile/capture, warmup, hosted-service connection, and teardown. | Long-lived reusable runtime created from `InferenceConfig`; lets FlashDreams load/warm once and create sessions sequentially unless the backend supports concurrency. | | InferenceSession | Owns one rollout or stream: initial inputs, cache state, current step, reset behavior, step requirements, and step execution. | Per-rollout interface consumed by the standard loop; keeps state isolated across prompts, browser clients, replay scenarios, or benchmark repeats. | @@ -281,8 +281,11 @@ native app, replay trace, synthetic benchmark driver, or no-op source. User inputs should primarily be represented as timestamped events. This gives live apps, replay traces, and benchmarks the same basic shape, and lets -FlashDreams resample or window those events when a model session asks for the -next chunk of inputs. +FlashDreams route, drain, or window those events when a model session asks for +the next chunk of inputs. Resampling and interpolation should remain +input-specific mapping or helper behavior, because controls such as rotations, +poses, or controller state may need semantics that generic runtime code cannot +infer safely. Initial supported user input types should stay close to what FlashDreams already uses: @@ -359,8 +362,8 @@ shapes or architecture details, but it usually does not fully define: - preprocessing, encoder, decoder, mask, prompt, or cache rules. Therefore, a FlashDreams-supported model should have an adapter or integration -layer that declares its model input requirements and prepares inputs for the -underlying model implementation. +layer that declares its model input requirements, declares any user inputs it can +map by default, and prepares inputs for the underlying model implementation. Users running an existing FlashDreams-supported model should not need to write that adapter. Developers bringing a new world model to FlashDreams should expect @@ -407,21 +410,25 @@ unless the checkpoint already matches a supported generic adapter. ## Input Mapping Input mapping is required whenever `UserInputs` need to become per-step -`ModelInputs`. The exact implementation does not need to be a required top-level -object. It could be: - -- a method on the model adapter; -- a method on an app/runtime adapter; -- a separate mapper object; -- a default no-op or identity mapping for simple T2V/I2V/fixed-input runs. +`ModelInputs`. In the T1 envelope this boundary is represented by a separate +`InputMapping` protocol. A model adapter may provide the default mapper because +it knows how its supported user controls affect model-facing inputs. Applications, +benchmarks, replay tools, or hosted runtimes may replace that mapper when they +need a different wire surface or aggregation policy. There are two separate moments to keep clear: - before runtime initialization, FlashDreams should select the mapping and check obvious compatibility between the app event source and the model; -- during the standard loop, the runner uses the mapping to build initial or - per-step `ModelInputs` from the relevant event window, often after the session - reports what it needs next. +- during the standard loop, the runtime or runner queues and timestamps user + events, then uses the selected mapping to build initial or per-step + `ModelInputs` from the relevant event window, often after the session reports + what it needs next. + +This keeps the Reactor-style contract intact: the model-side integration can +declare user inputs, declare model inputs, and provide a default mapping, while +the runtime owns transport, event validation, timestamping, input queue/window +selection, output delivery, and optional overrides. Examples: @@ -465,6 +472,11 @@ the app or output layer to guess. This matters because AR step 0 can differ from steady-state steps, and encoder/decoder temporal compression can produce different input and output frame windows. +Input and output timing should share a session timeline even when raw capture +rates and presentation rates differ. A session can request a user-input window +for mapping, then return an output window or equivalent metadata so an output +target can present the generated chunk at the intended cadence. + ## Output Targets Output handling should be separate from model execution. The model session @@ -598,20 +610,40 @@ back to `main` safely. - Public/internal boundaries must remain clean. Internal adapters, slugs, and scenarios should not leak into the public repo. -## Decisions To Make Before Implementation +## Decisions Made In T1 + +Task T1 settles the initial package and naming envelope without committing to a +registry, standard loop, concrete output modes, or model migrations: + +- The experimental API lives under `flashdreams.runtime`. +- The model-specific integration boundary is named `ModelAdapter`. +- Heavyweight lifecycle is split into `InferenceRuntime` and + `InferenceSession`. +- Step data carriers are named `StepRequest` and `StepResult`; a session returns + `None` from `next_step_request()` when the rollout is complete. +- User-facing inputs use `UserInputs`; model-facing inputs use `ModelInputs`. + Both remain lightweight payload envelopes with shallow read-only mappings. +- `UserInputSchema` and `ModelInputSchema` stay intentionally small: they + declare supported event types and required named fields for early validation, + not a full type system. +- Input mapping is represented by a separate `InputMapping` protocol. Model + adapters may provide a default mapping; runtimes and applications may override + it while preserving the `UserInputs` to `ModelInputs` boundary. Simple + fixed-input runs can use `IdentityInputMapping`. +- Output handling is represented by `OutputTarget`; `NullOutputTarget` is the + initial headless implementation. +- Metrics collection is represented by `MetricsRecorder`; timing samples use + seconds as the canonical unit. +- The minimum v0 user input shape is timestamped `UserInputEvent` records plus + optional snapshot data. Concrete event-type catalogs are left to T2 and demo + migrations. + +## Remaining Decisions -- What should the top-level package/API be called? -- Should the main registered object be called an adapter, integration, runtime - factory, or something else? - What direct-Python API should let users pass an external adapter without registering it? - What package registration mechanism should third-party and internal adapters use for CLI discovery and benchmarks? -- How lightweight should `UserInputSchema` and `ModelInputSchema` be? -- Where should input mapping live: model adapter, app adapter, separate object, - or a mix? -- What should the output abstraction be called? -- What is the minimum v0 set of supported user input events? - What is the first public model to migrate? - What metrics are required for every benchmark run? - What metadata must be discoverable without loading checkpoints? diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py new file mode 100644 index 000000000..03e6202b0 --- /dev/null +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Experimental inference runtime API envelope. + +This package defines the small v0 boundary above ``flashdreams.infra``. It is +intentionally additive while integrations migrate onto it. +""" + +from flashdreams.runtime.config import ExecutionBackend, InferenceConfig, Precision +from flashdreams.runtime.inputs import ( + InputField, + ModelInputs, + ModelInputSchema, + TimeWindow, + UserInputEvent, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.interfaces import ( + InferenceRuntime, + InferenceSession, + ModelAdapter, +) +from flashdreams.runtime.mapping import IdentityInputMapping, InputMapping +from flashdreams.runtime.metrics import ( + InMemoryMetricsRecorder, + MetricsRecorder, + NullMetricsRecorder, + RuntimeMetricSample, +) +from flashdreams.runtime.output import NullOutputTarget, OutputArtifact, OutputTarget +from flashdreams.runtime.types import StepRequest, StepResult + +__all__ = [ + "ExecutionBackend", + "IdentityInputMapping", + "InferenceConfig", + "InferenceRuntime", + "InferenceSession", + "InMemoryMetricsRecorder", + "InputField", + "InputMapping", + "MetricsRecorder", + "ModelAdapter", + "ModelInputs", + "ModelInputSchema", + "NullMetricsRecorder", + "NullOutputTarget", + "OutputArtifact", + "OutputTarget", + "Precision", + "RuntimeMetricSample", + "StepRequest", + "StepResult", + "TimeWindow", + "UserInputEvent", + "UserInputs", + "UserInputSchema", +] diff --git a/flashdreams/flashdreams/runtime/_utils.py b/flashdreams/flashdreams/runtime/_utils.py new file mode 100644 index 000000000..d8016c6b7 --- /dev/null +++ b/flashdreams/flashdreams/runtime/_utils.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Small helpers shared by the experimental runtime API.""" + +from __future__ import annotations + +from collections.abc import Mapping +from types import MappingProxyType +from typing import TypeVar + +ValueT = TypeVar("ValueT") + + +def freeze_mapping(value: Mapping[str, ValueT]) -> Mapping[str, ValueT]: + """Return a read-only shallow copy of ``value``.""" + return MappingProxyType(dict(value)) diff --git a/flashdreams/flashdreams/runtime/config.py b/flashdreams/flashdreams/runtime/config.py new file mode 100644 index 000000000..4b8752f13 --- /dev/null +++ b/flashdreams/flashdreams/runtime/config.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Runtime-facing configuration envelope.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Literal + +from flashdreams.runtime._utils import freeze_mapping + +ExecutionBackend = Literal["local", "local-distributed", "external", "hosted"] +"""Where and how inference compute is run.""" + +Precision = Literal["auto", "fp32", "fp16", "bf16"] +"""Coarse runtime precision choices.""" + + +@dataclass(frozen=True, kw_only=True, slots=True) +class InferenceConfig: + """Runtime settings that affect model execution. + + Prompts, user controls, browser settings, output paths, and benchmark + directories intentionally live outside this object. The typed optimization + fields cover common cross-backend knobs; open-ended adapter-specific choices + can use :attr:`runtime_options`. + """ + + __hash__ = None + + model_id: str + """Stable identity for the model adapter or runtime integration.""" + + preset_id: str | None = None + """Optional preset identity under :attr:`model_id`.""" + + checkpoint: str | Path | None = None + """Optional checkpoint or model-asset selector understood by the adapter.""" + + backend: ExecutionBackend = "local" + """Execution placement and backend family for inference compute.""" + + device: str | None = None + """Optional device selector such as ``cuda`` or ``cuda:0``; ``None`` leaves placement to the adapter/backend.""" + + precision: Precision = "auto" + """Preferred compute precision.""" + + compile: bool | None = None + """Optional - Whether model compilation is requested or disabled. `None` means left to the adapter to decide.""" + + cuda_graph: bool | None = None + """Optional - Whether CUDA graph capture is requested or disabled. `None` means left to the adapter to decide.""" + + attention_backend: str | None = None + """Optional attention implementation selector; ``None`` leaves the choice to the adapter.""" + + cache_policy: str | None = None + """Optional cache policy selector; ``None`` leaves the choice to the adapter.""" + + runtime_options: Mapping[str, Any] = field(default_factory=dict) + """Adapter/backend-specific runtime options.""" + + resource_hints: Mapping[str, Any] = field(default_factory=dict) + """Resource hints for launchers, schedulers, or hosted backends.""" + + def __post_init__(self) -> None: + if not self.model_id.strip(): + raise ValueError("InferenceConfig.model_id must be non-empty.") + object.__setattr__( + self, "runtime_options", freeze_mapping(self.runtime_options) + ) + object.__setattr__(self, "resource_hints", freeze_mapping(self.resource_hints)) diff --git a/flashdreams/flashdreams/runtime/inputs.py b/flashdreams/flashdreams/runtime/inputs.py new file mode 100644 index 000000000..e14b35722 --- /dev/null +++ b/flashdreams/flashdreams/runtime/inputs.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""User- and model-input envelopes for the experimental runtime API.""" + +from __future__ import annotations + +import math +from collections.abc import Iterable, Mapping +from dataclasses import dataclass, field +from typing import Any + +from flashdreams.runtime._utils import freeze_mapping + + +@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.""" + + name: str + required: bool = True + semantic_type: str | None = None + description: str = "" + + def __post_init__(self) -> None: + if not self.name.strip(): + raise ValueError("InputField.name must be non-empty.") + + +@dataclass(frozen=True, kw_only=True, slots=True) +class UserInputSchema: + """Minimal metadata for user events a source or mapping can provide.""" + + event_types: frozenset[str] = field(default_factory=frozenset) + snapshot_fields: tuple[InputField, ...] = () + description: str = "" + + def supports_event_types(self, event_types: Iterable[str]) -> bool: + """Return whether every requested event type is declared supported.""" + requested = frozenset(event_types) + if not requested: + return True + return requested.issubset(self.event_types) + + def missing_snapshot(self, inputs: "UserInputs") -> tuple[str, ...]: + """Return required snapshot fields absent from ``inputs``.""" + return _missing_required(self.snapshot_fields, inputs.snapshot) + + def require_snapshot(self, inputs: "UserInputs") -> None: + """Raise if required snapshot fields are absent.""" + missing = self.missing_snapshot(inputs) + if missing: + raise ValueError(f"Missing required user snapshot field(s): {missing}") + + +@dataclass(frozen=True, kw_only=True, slots=True) +class ModelInputSchema: + """Minimal metadata for model-facing initial and per-step inputs.""" + + initial_fields: tuple[InputField, ...] = () + """Model inputs required before starting the initial generation/session.""" + + step_fields: tuple[InputField, ...] = () + """Per-step model inputs required after the session starts.""" + + description: str = "" + + def missing_initial(self, inputs: "ModelInputs") -> tuple[str, ...]: + """Return required initial fields absent from ``inputs``.""" + return _missing_required(self.initial_fields, inputs.initial) + + def missing_step(self, inputs: "ModelInputs") -> tuple[str, ...]: + """Return required per-step fields absent from ``inputs``.""" + return _missing_required(self.step_fields, inputs.step) + + def require_initial(self, inputs: "ModelInputs") -> None: + """Raise if required initial fields are absent.""" + missing = self.missing_initial(inputs) + if missing: + raise ValueError(f"Missing required initial model input(s): {missing}") + + def require_step(self, inputs: "ModelInputs") -> None: + """Raise if required per-step fields are absent.""" + missing = self.missing_step(inputs) + if missing: + raise ValueError(f"Missing required step model input(s): {missing}") + + +@dataclass(frozen=True, kw_only=True, slots=True) +class UserInputEvent: + """User-facing input event timestamped in seconds since session start. + + Live runtimes, transports, replay loaders, or benchmark drivers stamp events + before queuing them for input mapping. Payload schema is intentionally minimal + in T1; concrete event catalogs belong to follow-up input-mapping work. + """ + + __hash__ = None + + timestamp_s: float + event_type: str + payload: Mapping[str, Any] = field(default_factory=dict) + source: str | None = None + source_event_id: str | None = None + + def __post_init__(self) -> None: + if not math.isfinite(self.timestamp_s) or self.timestamp_s < 0: + raise ValueError("UserInputEvent.timestamp_s must be finite and >= 0.") + if not self.event_type.strip(): + raise ValueError("UserInputEvent.event_type must be non-empty.") + object.__setattr__(self, "payload", freeze_mapping(self.payload)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class UserInputs: + """Transport-neutral user input batch or window. + + Events must be in non-decreasing timestamp order. Runtimes can pass the full + input history, a drained queue batch, or a session-requested time window to an + ``InputMapping``. + """ + + __hash__ = None + + events: tuple[UserInputEvent, ...] = () + snapshot: Mapping[str, Any] = field(default_factory=dict) + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + previous_timestamp_s = -math.inf + for event in self.events: + if event.timestamp_s < previous_timestamp_s: + raise ValueError( + "UserInputs.events must be sorted by non-decreasing timestamp_s." + ) + previous_timestamp_s = event.timestamp_s + object.__setattr__(self, "snapshot", freeze_mapping(self.snapshot)) + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + def window(self, time_window: TimeWindow) -> "UserInputs": + """Return inputs with events filtered to ``time_window``.""" + return UserInputs( + events=tuple( + event + for event in self.events + if time_window.contains(event.timestamp_s) + ), + snapshot=self.snapshot, + metadata=self.metadata, + ) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class ModelInputs: + """Model-facing payloads split by initial and per-step use.""" + + __hash__ = None + + initial: Mapping[str, Any] = field(default_factory=dict) + step: Mapping[str, Any] = field(default_factory=dict) + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__(self, "initial", freeze_mapping(self.initial)) + object.__setattr__(self, "step", freeze_mapping(self.step)) + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + def with_step(self, step: Mapping[str, Any]) -> "ModelInputs": + """Return a copy with replaced per-step payload.""" + return ModelInputs(initial=self.initial, step=step, metadata=self.metadata) + + +def _missing_required( + fields: tuple[InputField, ...], payload: Mapping[str, Any] +) -> tuple[str, ...]: + return tuple( + input_field.name + for input_field in fields + if input_field.required and input_field.name not in payload + ) diff --git a/flashdreams/flashdreams/runtime/interfaces.py b/flashdreams/flashdreams/runtime/interfaces.py new file mode 100644 index 000000000..9b6a064fd --- /dev/null +++ b/flashdreams/flashdreams/runtime/interfaces.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Protocols for model adapters, reusable runtimes, and sessions.""" + +from __future__ import annotations + +from typing import Protocol, runtime_checkable + +from flashdreams.runtime.config import InferenceConfig +from flashdreams.runtime.inputs import ( + ModelInputs, + ModelInputSchema, + UserInputSchema, +) +from flashdreams.runtime.mapping import InputMapping +from flashdreams.runtime.types import StepRequest, StepResult + + +@runtime_checkable +class InferenceSession(Protocol): + """One rollout or stream with isolated model/cache state.""" + + def next_step_request(self) -> StepRequest | None: + """Describe the next step's inputs, or return ``None`` when complete.""" + ... + + def step(self, inputs: ModelInputs) -> StepResult: + """Run one sequential inference step.""" + ... + + def reset(self, inputs: ModelInputs | None = None) -> None: + """Reset this session's rollout state when the backend supports it.""" + ... + + def close(self) -> None: + """Release per-session resources.""" + ... + + +@runtime_checkable +class InferenceRuntime(Protocol): + """Heavyweight reusable runtime created from :class:`InferenceConfig`.""" + + def start_session(self, inputs: ModelInputs) -> InferenceSession: + """Create an isolated session from initial model inputs.""" + ... + + def close(self) -> None: + """Release model/backend resources.""" + ... + + +# Do not mark ModelAdapter runtime-checkable: properties make issubclass() +# unreliable, and isinstance() would only verify attribute presence. +class ModelAdapter(Protocol): + """Model-specific boundary that declares defaults and creates runtimes. + + Adapters declare model-facing input requirements, optional user-input + capabilities, and an optional default mapping between the two. Runtime, + application, or benchmark code may override that mapping while preserving the + same ``UserInputs`` to ``ModelInputs`` boundary. + """ + + @property + def model_id(self) -> str: + """Stable identity for the model adapter or runtime integration.""" + ... + + @property + def model_input_schema(self) -> ModelInputSchema: + """Model-facing initial and per-step input requirements.""" + ... + + @property + def user_input_schema(self) -> UserInputSchema | None: + """User inputs supported by the adapter's default mapping, if any.""" + ... + + def default_input_mapping(self) -> InputMapping | None: + """Return the model-provided default user-to-model mapping, if any.""" + ... + + def validate_config(self, config: InferenceConfig) -> None: + """Fail early for unsupported runtime settings.""" + ... + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + """Initialize and return the heavyweight runtime.""" + ... diff --git a/flashdreams/flashdreams/runtime/mapping.py b/flashdreams/flashdreams/runtime/mapping.py new file mode 100644 index 000000000..756351081 --- /dev/null +++ b/flashdreams/flashdreams/runtime/mapping.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Input mapping boundary from user input windows to model inputs.""" + +from __future__ import annotations + +from typing import Protocol, runtime_checkable + +from flashdreams.runtime.inputs import ( + ModelInputs, + ModelInputSchema, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.types import StepRequest + + +@runtime_checkable +class InputMapping(Protocol): + """Convert user-facing inputs into model-facing inputs. + + A mapping may be supplied by the model adapter as a default or by an + application/runtime override. Step mappings usually receive a timestamped + event window selected by the runner for the current model step or chunk. + """ + + def validate( + self, + *, + user_schema: UserInputSchema | None = None, + model_schema: ModelInputSchema | None = None, + ) -> None: + """Fail early for obvious app, event-source, and model mismatches.""" + ... + + def map_initial_inputs( + self, + *, + user_inputs: UserInputs, + model_inputs: ModelInputs, + ) -> ModelInputs: + """Build initial model inputs before a session starts.""" + ... + + def map_step_inputs( + self, + *, + user_inputs: UserInputs, + model_inputs: ModelInputs, + request: StepRequest, + ) -> ModelInputs: + """Build model inputs for one session step from the current input window.""" + ... + + +class IdentityInputMapping: + """No-op mapper for fixed model-input or simple generation flows.""" + + def validate( + self, + *, + user_schema: UserInputSchema | None = None, + model_schema: ModelInputSchema | None = None, + ) -> None: + del user_schema, model_schema + + def map_initial_inputs( + self, + *, + user_inputs: UserInputs, + model_inputs: ModelInputs, + ) -> ModelInputs: + del user_inputs + return model_inputs + + def map_step_inputs( + self, + *, + user_inputs: UserInputs, + model_inputs: ModelInputs, + request: StepRequest, + ) -> ModelInputs: + del user_inputs, request + return model_inputs diff --git a/flashdreams/flashdreams/runtime/metrics.py b/flashdreams/flashdreams/runtime/metrics.py new file mode 100644 index 000000000..4286204f6 --- /dev/null +++ b/flashdreams/flashdreams/runtime/metrics.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Runtime metrics boundary for inference sessions.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any, Protocol, runtime_checkable + +from flashdreams.runtime._utils import freeze_mapping + + +@dataclass(frozen=True, kw_only=True, slots=True) +class RuntimeMetricSample: + """One runtime metric sample. + + Timing samples should use seconds as their canonical unit. + """ + + __hash__ = None + + name: str + value: float | int + unit: str = "s" + step_index: int | None = None + category: str = "runtime" + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.name.strip(): + raise ValueError("RuntimeMetricSample.name must be non-empty.") + if isinstance(self.value, bool) or not isinstance(self.value, (int, float)): + raise TypeError("RuntimeMetricSample.value must be numeric.") + if not math.isfinite(float(self.value)): + raise ValueError("RuntimeMetricSample.value must be finite.") + if self.step_index is not None and self.step_index < 0: + raise ValueError("RuntimeMetricSample.step_index must be >= 0.") + if not self.unit.strip(): + raise ValueError("RuntimeMetricSample.unit must be non-empty.") + if self.category == "timing" and self.unit != "s": + raise ValueError("Timing metric samples must use unit='s'.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@runtime_checkable +class MetricsRecorder(Protocol): + """Collector for runtime metrics.""" + + def record(self, sample: RuntimeMetricSample) -> None: + """Record one metric sample.""" + ... + + def record_timing( + self, + name: str, + duration_s: float, + *, + step_index: int | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> None: + """Record one timing sample in seconds.""" + ... + + def close(self) -> None: + """Finalize metric collection.""" + ... + + +@dataclass(slots=True) +class InMemoryMetricsRecorder: + """Simple metrics recorder useful for tests, smoke runs, and adapters.""" + + samples: list[RuntimeMetricSample] = field(default_factory=list) + closed: bool = False + + def record(self, sample: RuntimeMetricSample) -> None: + if self.closed: + raise RuntimeError("Cannot record metrics after close().") + self.samples.append(sample) + + def record_timing( + self, + name: str, + duration_s: float, + *, + step_index: int | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> None: + self.record( + RuntimeMetricSample( + name=name, + value=duration_s, + unit="s", + step_index=step_index, + category="timing", + metadata={} if metadata is None else metadata, + ) + ) + + def close(self) -> None: + self.closed = True + + +class NullMetricsRecorder: + """Metrics recorder that intentionally drops all samples.""" + + def record(self, sample: RuntimeMetricSample) -> None: + del sample + + def record_timing( + self, + name: str, + duration_s: float, + *, + step_index: int | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> None: + del name, duration_s, step_index, metadata + + def close(self) -> None: + return None diff --git a/flashdreams/flashdreams/runtime/output.py b/flashdreams/flashdreams/runtime/output.py new file mode 100644 index 000000000..aac341ee1 --- /dev/null +++ b/flashdreams/flashdreams/runtime/output.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Output target boundary for generated inference results.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, Protocol, runtime_checkable + +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.types import StepResult + + +@dataclass(frozen=True, kw_only=True, slots=True) +class OutputArtifact: + """Artifact produced by an output target.""" + + __hash__ = None + + kind: str + uri: str + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.kind.strip(): + raise ValueError("OutputArtifact.kind must be non-empty.") + if not self.uri.strip(): + raise ValueError("OutputArtifact.uri must be non-empty.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@runtime_checkable +class OutputTarget(Protocol): + """Consumes generated session outputs for presentation or persistence.""" + + def open(self) -> None: + """Prepare the target for a new run.""" + ... + + def write(self, result: StepResult) -> None: + """Consume one generated step result.""" + ... + + def close(self) -> Sequence[OutputArtifact]: + """Finalize and return any produced artifacts.""" + ... + + +@dataclass(slots=True) +class NullOutputTarget: + """Output target for headless runs and throughput measurements.""" + + store_results: bool = False + output_count: int = field(default=0, init=False) + results: list[StepResult] = field(default_factory=list, init=False) + _opened: bool = field(default=False, init=False, repr=False) + + @property + def closed(self) -> bool: + return not self._opened + + def open(self) -> None: + self._opened = True + self.output_count = 0 + self.results.clear() + + def write(self, result: StepResult) -> None: + if not self._opened: + raise RuntimeError("Cannot write to a closed output target.") + self.output_count += 1 + if self.store_results: + self.results.append(result) + + def close(self) -> Sequence[OutputArtifact]: + self._opened = False + return () diff --git a/flashdreams/flashdreams/runtime/types.py b/flashdreams/flashdreams/runtime/types.py new file mode 100644 index 000000000..52bf82166 --- /dev/null +++ b/flashdreams/flashdreams/runtime/types.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Plain data carriers shared by runtime protocols and adapters.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.inputs import ModelInputSchema, TimeWindow + + +@dataclass(frozen=True, kw_only=True, slots=True) +class StepRequest: + """Model-session request for the next step's inputs. + + ``user_input_window`` lets a runner drain or slice timestamped user events for + the current step before invoking the selected ``InputMapping``. + """ + + __hash__ = None + + step_index: int + model_input_schema: ModelInputSchema | None = None + user_input_window: TimeWindow | None = None + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.step_index < 0: + raise ValueError("StepRequest.step_index must be >= 0.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class StepResult: + """Generated output and metadata for 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)) diff --git a/flashdreams/tests/test_inference_runtime_api.py b/flashdreams/tests/test_inference_runtime_api.py new file mode 100644 index 000000000..1474383a0 --- /dev/null +++ b/flashdreams/tests/test_inference_runtime_api.py @@ -0,0 +1,492 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import fields +from typing import Any, cast + +import pytest + +from flashdreams.runtime import ( + IdentityInputMapping, + InferenceConfig, + InferenceRuntime, + InferenceSession, + InMemoryMetricsRecorder, + InputField, + InputMapping, + MetricsRecorder, + ModelAdapter, + ModelInputs, + ModelInputSchema, + NullOutputTarget, + OutputArtifact, + OutputTarget, + RuntimeMetricSample, + StepRequest, + StepResult, + TimeWindow, + UserInputEvent, + UserInputs, + UserInputSchema, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_inference_config_keeps_runtime_settings_separate() -> None: + denied_app_fields = {"prompt", "output_dir", "browser_settings"} + config = InferenceConfig( + model_id="lingbot-world", + preset_id="fast-taehv", + backend="local", + precision="bf16", + compile=False, + runtime_options={"chunk_size": 3}, + ) + + assert config.model_id == "lingbot-world" + assert config.preset_id == "fast-taehv" + assert config.runtime_options["chunk_size"] == 3 + assert denied_app_fields.isdisjoint(field.name for field in fields(InferenceConfig)) + with pytest.raises(TypeError): + cast(Any, config.runtime_options)["chunk_size"] = 4 + + +def test_inference_config_rejects_empty_model_id() -> None: + with pytest.raises(ValueError, match="model_id"): + InferenceConfig(model_id=" ") + + +@pytest.mark.parametrize( + ("factory", "match"), + [ + (lambda: InputField(name=" "), "InputField.name"), + (lambda: TimeWindow(start_s=1.0, end_s=0.0), "end_s"), + (lambda: TimeWindow(start_s=-1.0, end_s=0.0), "non-negative"), + (lambda: TimeWindow(start_s=0.0, end_s=float("nan")), "finite"), + ( + lambda: UserInputEvent(timestamp_s=-1.0, event_type="keydown"), + "timestamp_s", + ), + (lambda: UserInputEvent(timestamp_s=0.0, event_type=" "), "event_type"), + (lambda: StepRequest(step_index=-1), "step_index"), + (lambda: StepResult(step_index=-1), "step_index"), + (lambda: StepResult(step_index=0, frame_count=-1), "frame_count"), + (lambda: RuntimeMetricSample(name=" ", value=1.0), "name"), + (lambda: RuntimeMetricSample(name="sample", value=float("nan")), "finite"), + (lambda: OutputArtifact(kind=" ", uri="artifact://demo"), "kind"), + (lambda: OutputArtifact(kind="mp4", uri=" "), "uri"), + ], +) +def test_runtime_envelopes_reject_invalid_values(factory: object, match: str) -> None: + with pytest.raises(ValueError, match=match): + cast(Any, factory)() + + +def test_runtime_metric_sample_rejects_bool_values() -> None: + with pytest.raises(TypeError, match="numeric"): + RuntimeMetricSample(name="sample", value=True) + + +def test_model_input_schema_validates_initial_and_step_payloads() -> None: + schema = ModelInputSchema( + initial_fields=( + InputField(name="prompt"), + InputField(name="first_frame"), + ), + step_fields=(InputField(name="camera_poses"),), + ) + inputs = ModelInputs(initial={"prompt": "drive", "first_frame": object()}) + + schema.require_initial(inputs) + assert schema.missing_step(inputs) == ("camera_poses",) + + with pytest.raises(ValueError, match="camera_poses"): + schema.require_step(inputs) + + +def test_user_inputs_filter_timestamped_event_windows() -> None: + inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.1, + event_type="keyboard.keydown", + payload={"key": "w"}, + ), + UserInputEvent( + timestamp_s=0.4, + event_type="keyboard.keyup", + payload={"key": "w"}, + ), + UserInputEvent(timestamp_s=0.8, event_type="reset"), + ) + ) + + windowed = inputs.window(TimeWindow(start_s=0.25, end_s=0.75)) + + assert [event.event_type for event in windowed.events] == ["keyboard.keyup"] + + +def test_user_inputs_require_sorted_events() -> None: + with pytest.raises(ValueError, match="non-decreasing"): + UserInputs( + events=( + UserInputEvent(timestamp_s=1.0, event_type="late"), + UserInputEvent(timestamp_s=0.5, event_type="early"), + ) + ) + + +def test_user_input_schema_declares_event_capabilities() -> None: + schema = UserInputSchema( + event_types=frozenset({"keyboard.keydown", "keyboard.keyup", "reset"}) + ) + + assert schema.supports_event_types(["keyboard.keydown", "reset"]) + assert not schema.supports_event_types(["prompt.update"]) + + +def test_user_input_schema_validates_required_snapshot_fields() -> None: + schema = UserInputSchema( + snapshot_fields=( + InputField(name="pressed_keys"), + InputField(name="prompt", required=False), + ) + ) + inputs = UserInputs(snapshot={"pressed_keys": frozenset({"w"})}) + + schema.require_snapshot(inputs) + assert schema.missing_snapshot(UserInputs()) == ("pressed_keys",) + + with pytest.raises(ValueError, match="pressed_keys"): + schema.require_snapshot(UserInputs()) + + +def test_identity_input_mapping_leaves_model_inputs_unchanged() -> None: + mapping = IdentityInputMapping() + model_inputs = ModelInputs(initial={"prompt": "fixed"}, step={"hdmap": object()}) + request = StepRequest(step_index=0) + + assert ( + mapping.map_initial_inputs( + user_inputs=UserInputs(), + model_inputs=model_inputs, + ) + is model_inputs + ) + assert ( + mapping.map_step_inputs( + user_inputs=UserInputs(), + model_inputs=model_inputs, + request=request, + ) + is model_inputs + ) + + +def test_null_output_target_counts_and_optionally_stores_results() -> None: + target = NullOutputTarget(store_results=True) + result = StepResult(step_index=0, output=b"frame") + + assert target.closed + with pytest.raises(RuntimeError, match="closed output target"): + target.write(result) + + target.open() + assert not target.closed + target.write(result) + artifacts = target.close() + + assert target.closed + assert artifacts == () + assert target.output_count == 1 + assert target.results == [result] + with pytest.raises(RuntimeError, match="closed output target"): + target.write(StepResult(step_index=1)) + + +def test_null_output_target_open_resets_per_run_state() -> None: + target = NullOutputTarget(store_results=True) + + target.open() + target.write(StepResult(step_index=0, output=b"first")) + target.close() + target.open() + + assert target.output_count == 0 + assert target.results == [] + target.write(StepResult(step_index=0, output=b"second")) + assert target.output_count == 1 + assert target.results == [StepResult(step_index=0, output=b"second")] + + +def test_in_memory_metrics_recorder_uses_seconds_for_timing() -> None: + recorder = InMemoryMetricsRecorder() + + recorder.record_timing("model_step", 0.125, step_index=2) + + assert len(recorder.samples) == 1 + sample = recorder.samples[0] + assert sample.name == "model_step" + assert sample.value == pytest.approx(0.125) + assert sample.unit == "s" + assert sample.category == "timing" + assert sample.step_index == 2 + + +def test_timing_metric_samples_must_use_seconds() -> None: + with pytest.raises(ValueError, match="unit='s'"): + RuntimeMetricSample( + name="model_step", + value=12.5, + unit="ms", + category="timing", + ) + + +def test_runtime_api_components_compose_for_sequential_session() -> None: + adapter = _FakeAdapter() + config = InferenceConfig(model_id="fake-model") + user_inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.25, + event_type="keyboard.keydown", + payload={"key": "w"}, + ), + ) + ) + model_inputs = ModelInputs(initial={"prompt": "drive forward"}) + output = NullOutputTarget(store_results=True) + metrics = InMemoryMetricsRecorder() + + adapter.validate_config(config) + mapping = adapter.default_input_mapping() + assert mapping is not None + _drive_two_step_session( + adapter=adapter, + config=config, + mapping=mapping, + user_inputs=user_inputs, + model_inputs=model_inputs, + output=output, + metrics=metrics, + ) + + assert output.output_count == 2 + assert [result.output for result in output.results] == ["chunk-0", "chunk-1"] + assert [result.frame_count for result in output.results] == [3, 3] + assert output.results[0].output_window == TimeWindow(start_s=0.0, end_s=0.5) + assert [sample.step_index for sample in metrics.samples] == [0, 1] + assert metrics.closed + + +def test_reference_loop_validates_mapping_before_runtime_creation() -> None: + mapping = _OrderCheckingMapping() + adapter = _OrderCheckingAdapter(mapping=mapping) + + _drive_two_step_session( + adapter=adapter, + config=InferenceConfig(model_id="fake-model"), + mapping=mapping, + user_inputs=UserInputs(), + model_inputs=ModelInputs(initial={"prompt": "drive forward"}), + output=NullOutputTarget(), + metrics=InMemoryMetricsRecorder(), + ) + + assert mapping.validated + assert adapter.created_runtime_after_validate + + +def test_reference_loop_closes_runtime_when_session_start_fails() -> None: + adapter = _FailingStartAdapter() + output = NullOutputTarget() + metrics = InMemoryMetricsRecorder() + + with pytest.raises(RuntimeError, match="start failed"): + _drive_two_step_session( + adapter=adapter, + config=InferenceConfig(model_id="fake-model"), + mapping=IdentityInputMapping(), + user_inputs=UserInputs(), + model_inputs=ModelInputs(initial={"prompt": "drive forward"}), + output=output, + metrics=metrics, + ) + + assert adapter.runtime is not None + assert adapter.runtime.closed + assert output.closed + assert metrics.closed + + +def _drive_two_step_session( + *, + adapter: ModelAdapter, + config: InferenceConfig, + mapping: InputMapping, + user_inputs: UserInputs, + model_inputs: ModelInputs, + output: OutputTarget, + metrics: MetricsRecorder, +) -> None: + mapping.validate( + user_schema=adapter.user_input_schema, + model_schema=adapter.model_input_schema, + ) + initial_inputs = mapping.map_initial_inputs( + user_inputs=user_inputs, + model_inputs=model_inputs, + ) + runtime = adapter.create_runtime(config) + session: InferenceSession | None = None + output_opened = False + try: + session = runtime.start_session(initial_inputs) + output.open() + output_opened = True + while (request := session.next_step_request()) is not None: + step_inputs = mapping.map_step_inputs( + user_inputs=( + user_inputs.window(request.user_input_window) + if request.user_input_window is not None + else user_inputs + ), + model_inputs=ModelInputs( + initial=initial_inputs.initial, + step={"chunk_index": request.step_index}, + ), + request=request, + ) + result = session.step(step_inputs) + output.write(result) + metrics.record_timing( + "model_step", + float(result.metrics["model_step_s"]), + step_index=result.step_index, + ) + finally: + if output_opened: + output.close() + if session is not None: + session.close() + runtime.close() + metrics.close() + + +class _FakeAdapter: + model_id = "fake-model" + model_input_schema = ModelInputSchema( + initial_fields=(InputField(name="prompt"),), + step_fields=(InputField(name="chunk_index"),), + ) + user_input_schema = UserInputSchema(event_types=frozenset({"keyboard.keydown"})) + + def default_input_mapping(self) -> InputMapping: + return IdentityInputMapping() + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.model_id: + raise ValueError(f"Unsupported model_id={config.model_id!r}.") + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + return _FakeRuntime(model_input_schema=self.model_input_schema) + + +class _FakeRuntime: + def __init__(self, *, model_input_schema: ModelInputSchema) -> None: + self._model_input_schema = model_input_schema + self.closed = False + + def start_session(self, inputs: ModelInputs) -> InferenceSession: + self._model_input_schema.require_initial(inputs) + return _FakeSession(model_input_schema=self._model_input_schema) + + def close(self) -> None: + self.closed = True + + +class _FailingRuntime(_FakeRuntime): + def start_session(self, inputs: ModelInputs) -> InferenceSession: + del inputs + raise RuntimeError("start failed") + + +class _FakeSession: + def __init__(self, *, model_input_schema: ModelInputSchema) -> None: + self._model_input_schema = model_input_schema + self.step_index = 0 + self.closed = False + + def next_step_request(self) -> StepRequest | None: + if self.step_index >= 2: + return None + return StepRequest( + step_index=self.step_index, + model_input_schema=self._model_input_schema, + user_input_window=TimeWindow( + start_s=0.5 * self.step_index, + end_s=0.5 * (self.step_index + 1), + ), + ) + + def step(self, inputs: ModelInputs) -> StepResult: + self._model_input_schema.require_step(inputs) + result = StepResult( + step_index=self.step_index, + output=f"chunk-{self.step_index}", + frame_count=3, + output_window=TimeWindow( + start_s=0.5 * self.step_index, + end_s=0.5 * (self.step_index + 1), + ), + metrics={"model_step_s": 0.01}, + ) + self.step_index += 1 + return result + + def reset(self, inputs: ModelInputs | None = None) -> None: + del inputs + self.step_index = 0 + + def close(self) -> None: + self.closed = True + + +class _OrderCheckingMapping(IdentityInputMapping): + def __init__(self) -> None: + self.validated = False + + def validate( + self, + *, + user_schema: UserInputSchema | None = None, + model_schema: ModelInputSchema | None = None, + ) -> None: + super().validate(user_schema=user_schema, model_schema=model_schema) + self.validated = True + + +class _OrderCheckingAdapter(_FakeAdapter): + def __init__(self, *, mapping: _OrderCheckingMapping) -> None: + self._mapping = mapping + self.created_runtime_after_validate = False + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + self.created_runtime_after_validate = self._mapping.validated + return _FakeRuntime(model_input_schema=self.model_input_schema) + + +class _FailingStartAdapter(_FakeAdapter): + def __init__(self) -> None: + self.runtime: _FailingRuntime | None = None + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + self.runtime = _FailingRuntime(model_input_schema=self.model_input_schema) + return self.runtime From d107043aec5f8a2e4cc7a0ca76d2f2234a1017a8 Mon Sep 17 00:00:00 2001 From: aidanfnv Date: Wed, 5 Aug 2026 09:58:16 -0700 Subject: [PATCH 03/19] WIP Implement T2, T3, and part of T4 from API refactor plan (#413) * WIP implementation of T2, T3, partial T4 * Fix issues found by Claude * Rewrite based on discussion, port after merge * doc update * doc updates * Update based on new diagrams * Align closer to diagrams --- docs/inference_runtime_api_design.md | 111 +++- ...inference_runtime_inputs_implementation.md | 287 +++++++++ ...ence_runtime_supported_inputs_inventory.md | 321 ++++++++++ flashdreams/flashdreams/runtime/__init__.py | 59 +- flashdreams/flashdreams/runtime/canonical.py | 387 ++++++++++++ flashdreams/flashdreams/runtime/inputs.py | 359 ++++++++++- flashdreams/flashdreams/runtime/interfaces.py | 30 +- flashdreams/flashdreams/runtime/mapping.py | 356 ++++++++++- flashdreams/flashdreams/runtime/types.py | 4 +- .../tests/test_inference_runtime_api.py | 154 +++-- flashdreams/tests/test_runtime_canonical.py | 590 ++++++++++++++++++ .../tests/test_runtime_input_mapping.py | 573 +++++++++++++++++ 12 files changed, 3076 insertions(+), 155 deletions(-) create mode 100644 docs/inference_runtime_inputs_implementation.md create mode 100644 docs/inference_runtime_supported_inputs_inventory.md create mode 100644 flashdreams/flashdreams/runtime/canonical.py create mode 100644 flashdreams/tests/test_runtime_canonical.py create mode 100644 flashdreams/tests/test_runtime_input_mapping.py diff --git a/docs/inference_runtime_api_design.md b/docs/inference_runtime_api_design.md index 2f0ba19f8..f70fbd890 100644 --- a/docs/inference_runtime_api_design.md +++ b/docs/inference_runtime_api_design.md @@ -19,7 +19,7 @@ integration-specific runner code: - `InferenceConfig`: how the model and inference stack should run; - `UserInputs`: controls or events from an app, replay trace, or benchmark; -- `ModelInputs`: prompts, frames, videos, trajectories, maps, scene data, and +- `InferenceInput`: prompts, frames, videos, trajectories, maps, scene data, and other values required by a specific model; - input mapping: model/application-specific conversion from user-facing inputs into model-facing inputs; @@ -30,6 +30,12 @@ integration-specific runner code: - metrics/profiling: timings, memory, traces, NVTX ranges, and benchmark outputs. +Current T2/T3 implementation notes are in +`docs/inference_runtime_inputs_implementation.md`. + +The supported-model input inventory used to revisit T2/T3 is in +`docs/inference_runtime_supported_inputs_inventory.md`. + The API should standardize the envelope and lifecycle. It should not pretend that all world models have the same inputs, that all models use the same optimization stack, or that a raw checkpoint can fully describe how to run the @@ -62,9 +68,9 @@ Initial scope: | ID | Status | Workstream | Can run in parallel? | Depends on | Done when | | --- | --- | --- | --- | --- | --- | | T0 | Complete | Create experimental branch and contribution rules. | No, this starts the work. | None. | Branch exists, PR target is agreed, and main merge criteria are written down. | -| T1 | Complete | Minimal API envelope and naming. | Partly. | T0. | `InferenceConfig`, `UserInputs`, `ModelInputs`, runtime/session, output target, and mapping boundaries are defined well enough for demos to use. | -| T2 | Planned | Event-based `UserInputs`. | Yes, after T1 direction is agreed. | T1. | User inputs are primarily timestamped events; replay traces and derived snapshots are supported where needed. | -| T3 | Planned | `ModelInputs`, schemas, and mapping boundary. | Yes, after T1 direction is agreed. | T1. | Models can declare required initial/per-step inputs, and mappings can convert user events into model inputs. | +| T1 | Complete | Minimal API envelope and naming. | Partly. | T0. | `InferenceConfig`, `UserInputs`, `InferenceInput`, runtime/session, output target, and mapping boundaries are defined well enough for demos to use. | +| T2 | Complete | Event-based `UserInputs`. | Yes, after T1 direction is agreed. | T1. | User inputs are primarily timestamped events; replay traces and derived snapshots are supported where needed. | +| T3 | Complete | `CanonicalInputs`, `InferenceInput`, schemas, and mapping boundary. | Yes, after T1 direction is agreed. | T1. | Models can declare required global/per-step inputs, and mappings can convert canonical inputs into inference inputs. | | T4 | Planned | `ModelRunner`, `InferenceRuntime`, and `InferenceSession` skeleton. | Partly. | T1. | A minimal standard loop can initialize a runtime, run at least one sequential session, and close cleanly. | | T5 | Planned | Output mode selection. | Yes, after the result/output shape is agreed. | T1, T4. | A run can choose output behavior such as MP4, JPEG/MJPEG stream, WebRTC, benchmark artifact, or headless/null without changing model code. | | T6 | Planned | LingBot migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | LingBot runs through the new API path with its event inputs mapped into model inputs. | @@ -97,14 +103,14 @@ Main runtime flow: App / integration / benchmark / transport chooses how the run is driven and where output goes supplies run setup: - InferenceConfig + UserInputs + ModelInputs + output/metrics options + InferenceConfig + UserInputs + InferenceInput + output/metrics options | v ModelRunner / standard loop orchestrates validation, lifecycle, stepping, output, and metrics uses input mapping to: validate that user/app inputs can drive the model - build initial and per-step ModelInputs during the run + build global and per-step InferenceInput during the run | v InferenceRuntime @@ -145,7 +151,7 @@ Create InferenceRuntime from InferenceConfig | v Start InferenceSession A - initial ModelInputs: prompt/frame/scene/etc. + global conditioning: prompt/frame/scene/etc. per-session state: cache, current step, reset state step 0 -> step 1 -> ... -> done outputs -> Output target @@ -154,7 +160,7 @@ Start InferenceSession A | v Start InferenceSession B - new initial ModelInputs or replay scenario + new global conditioning or replay scenario independent cache/state step 0 -> step 1 -> ... -> done outputs -> Output target @@ -305,33 +311,63 @@ User inputs are not model inputs. A keyboard event does not have one universal meaning. One model may map it to pose segments, another to steering commands, and another may ignore it. -## ModelInputs +## CanonicalInputs And InferenceInput + +Inputs move through three layers: + +```text +UserInputs -> CanonicalInputs -> InferenceInput + raw canonicalized encoded +``` + +Raw device events are canonicalized into device-independent modalities before an +application sees them, so adding a keyboard, gamepad, or wheel is a converter +registration rather than an application change. `InferenceInput` is what an +`InferenceSession` actually receives. + +`InferenceInput` describes the data the model or inference pipeline actually +requires. Both it and `CanonicalInputs` distinguish two conditioning slots: -`ModelInputs` describes the data the model or inference pipeline actually -requires. It should distinguish: +- global conditioning: values that condition the whole rollout; +- per-step conditioning: values needed for one generated chunk or frame window. -- initial inputs: values needed to start or reset a rollout; -- per-step inputs: values needed for one generated chunk or frame window. +Examples of global conditioning include prompt, negative prompt, conditioning +frame, input video, scene id, HD map asset, camera calibration, initial camera +pose, seed, or model-specific fields. -Examples of initial model inputs include prompt, negative prompt, first frame, -input video, scene id, HD map asset, camera calibration, initial camera pose, -seed, or model-specific fields. +Global conditioning is normally supplied when a session starts, but a non-empty +global slot on a mid-rollout input is an update request rather than a reset; +resetting rollout state is a separate `InferenceSession.reset()` call. Whether a +given value can be swapped mid-rollout is declared per field by +`InputField.update_policy`. -Examples of per-step model inputs include frame timestamps, pose segments, +Examples of per-step conditioning include frame timestamps, pose segments, camera trajectory chunks, rendered HD map frames, conditioning video windows, control tensors, event markers, or model-specific fields. -Model input payloads should use semantic names, not only modality names. For +Inference input payloads should use semantic names, not only modality names. For example, a first frame and an HD map frame should be distinct inputs even if both are image-like values. -For interactive runs, most `ModelInputs` will be initial values plus per-step -inputs produced by input mapping. For MP4 generation and benchmarking, the API +Model input metadata may also include a lightweight lifecycle label, such as +runtime config, cache initialization, rollout binding, per-step input, or +session update. This should remain query metadata, not model-specific tensor +validation. + +Model input names, payload kinds, lifecycle labels, and schema metadata should +be open-ended. Supported integrations such as SANA-WM, LingBot, Omnidreams, and +future external adapters may need different semantic fields. Adding a new model +should usually mean adding adapter-owned schema declarations and mappings, not +changing a central FlashDreams enum. + +For interactive runs, most `InferenceInput` values will be global conditioning +plus per-step inputs produced by input mapping. For MP4 generation and benchmarking, the API should also support fixed per-step model inputs so runs can be deterministic. ## Schemas -The API should support lightweight `UserInputSchema` and `ModelInputSchema` +The API should support lightweight `UserInputSchema`, `CanonicalInputSchema`, +and `InferenceInputSchema` metadata. These schemas are not meant to be a rich type system or a replacement for @@ -345,8 +381,15 @@ The purpose is to fail early before expensive model initialization, produce clearer errors, make fixed scenarios easier to validate, and avoid ambiguous dict payloads where keys only describe modality. +Schema objects may carry open-ended metadata for query-time hints such as +coordinate frame, units, rough shape summary, accepted file suffixes, schema +URI, model family, or source/transport details. Metadata should help humans and +adapter selection code, but compatibility should still be based on the declared +event capabilities, semantic model fields, payload representation hints, and +lifecycle labels. + For simple CLI text-to-video or image-to-video runs, `UserInputSchema` can be -trivial or omitted because there may be no live controls. `ModelInputSchema` is +trivial or omitted because there may be no live controls. `InferenceInputSchema` is more important because each supported model still needs to declare the model-facing values it expects. @@ -410,19 +453,24 @@ unless the checkpoint already matches a supported generic adapter. ## Input Mapping Input mapping is required whenever `UserInputs` need to become per-step -`ModelInputs`. In the T1 envelope this boundary is represented by a separate +`InferenceInput`. In the T1 envelope this boundary is represented by a separate `InputMapping` protocol. A model adapter may provide the default mapper because it knows how its supported user controls affect model-facing inputs. Applications, benchmarks, replay tools, or hosted runtimes may replace that mapper when they need a different wire surface or aggregation policy. +The selected mapping may be a single mapper or a composed set of mappers, so one +run can combine separate prompt, first-frame, and live-control mappings instead +of routing everything through one object. + There are two separate moments to keep clear: -- before runtime initialization, FlashDreams should select the mapping and check - obvious compatibility between the app event source and the model; +- before runtime initialization, FlashDreams should select the mapping or mapper + set and check obvious compatibility between the app event source and the + model; - during the standard loop, the runtime or runner queues and timestamps user events, then uses the selected mapping to build initial or per-step - `ModelInputs` from the relevant event window, often after the session reports + `InferenceInput` from the relevant event window, often after the session reports what it needs next. This keeps the Reactor-style contract intact: the model-side integration can @@ -621,14 +669,16 @@ registry, standard loop, concrete output modes, or model migrations: `InferenceSession`. - Step data carriers are named `StepRequest` and `StepResult`; a session returns `None` from `next_step_request()` when the rollout is complete. -- User-facing inputs use `UserInputs`; model-facing inputs use `ModelInputs`. +- Raw inputs use `UserInputs`, canonicalized inputs use `CanonicalInputs`, and + model-facing inputs use `InferenceInput`. Both remain lightweight payload envelopes with shallow read-only mappings. -- `UserInputSchema` and `ModelInputSchema` stay intentionally small: they +- `UserInputSchema`, `CanonicalInputSchema`, and `InferenceInputSchema` stay + intentionally small: they declare supported event types and required named fields for early validation, not a full type system. - Input mapping is represented by a separate `InputMapping` protocol. Model adapters may provide a default mapping; runtimes and applications may override - it while preserving the `UserInputs` to `ModelInputs` boundary. Simple + it while preserving the `CanonicalInputs` to `InferenceInput` boundary. Simple fixed-input runs can use `IdentityInputMapping`. - Output handling is represented by `OutputTarget`; `NullOutputTarget` is the initial headless implementation. @@ -660,7 +710,8 @@ Proceed with the proposed split: - `InferenceConfig` for model/runtime execution; - `UserInputs` for app-facing controls and replay traces; -- `ModelInputs` for model-facing initial and per-step inputs; +- `CanonicalInputs` for device-independent application-facing inputs; +- `InferenceInput` for model-facing global and per-step conditioning; - input mapping for model/application-specific conversion; - runtime/session boundaries for lifecycle and stepping; - output targets for display, streaming, files, and benchmarks; diff --git a/docs/inference_runtime_inputs_implementation.md b/docs/inference_runtime_inputs_implementation.md new file mode 100644 index 000000000..8d485768a --- /dev/null +++ b/docs/inference_runtime_inputs_implementation.md @@ -0,0 +1,287 @@ + + +# Inference Runtime Inputs Implementation Notes + +This note documents the input layers of the experimental runtime API: what +exists, how the pieces fit together, what the compatibility query answers, and +what is intentionally still outside this layer. + +Implementation lives in `flashdreams.runtime`: + +- `flashdreams/flashdreams/runtime/inputs.py` — the input types and schemas +- `flashdreams/flashdreams/runtime/canonical.py` — raw device to canonical + modality conversion +- `flashdreams/flashdreams/runtime/mapping.py` — canonical to encoded mapping + and compatibility +- `flashdreams/tests/test_runtime_canonical.py` +- `flashdreams/tests/test_runtime_input_mapping.py` +- `flashdreams/tests/test_inference_runtime_api.py` — the T1 envelope tests, + including a reference loop that exercises all three layers + +The supported-model input inventory that informed this work is in +`docs/inference_runtime_supported_inputs_inventory.md`. + +## The Three Layers + +```text +UserInputs ──InputCanonicalizer──▶ CanonicalInputs ──InputMapping──▶ InferenceInput + raw canonicalized encoded +(device events) (device-independent) (what the session gets) +``` + +| Layer | Type | Owner | Example | +| --- | --- | --- | --- | +| raw | `UserInputs` / `UserInputEvent` | transport, replay loader, benchmark driver | `key_down {"key": "w"}`, wheel axis reading | +| canonicalized | `CanonicalInputs` | device converters registered on `InputCanonicalizer` | `driver_command {throttle, brake, steer, ...}` | +| encoded | `InferenceInput` | the selected `InputMapping` | whatever the model's session consumes | + +Applications and mappings consume `CanonicalInputs`. They never read raw device +events: `InputMapping.map_step_inputs` takes `canonical_inputs`, not +`user_inputs`, so this is enforced by the signature rather than by convention. +Adding a keyboard, gamepad, or wheel is an `InputCanonicalizer.register` call +that touches no application, mapping, or model code. + +This path covers **live user control only**. Global conditioning is +application-owned data and reaches `InferenceInput` directly, without passing +through canonicalization or a device converter. An application that wants a +trigger key to swap the prompt reads that as ordinary canonical control input +and updates its own global conditioning in response. + +## Conditioning Slots + +Both the canonical and encoded layers split into two slots, and the split means +the same thing at each: + +- **global conditioning** — conditions the whole rollout: prompt, conditioning + frame, scene. Normally supplied at session start. +- **per-step conditioning** — needed to generate the next chunk or frame: + steering, HD map frames, camera trajectory. + +`InputPhase` is `Literal["global", "step"]`. The axis names *which slot*, not +*when the value may arrive* — see the next section. + +## Global Conditioning Updates Are Not Resets + +A non-empty global slot on a mid-rollout `InferenceInput` is an **update +request**. The session should apply it when the model supports doing so. +Resetting rollout state is a separate, explicit `InferenceSession.reset()` call. +The motivating case is changing prompt and conditioning frame mid-run to change +the weather in an Omnidreams rollout. + +```python +from flashdreams.runtime import InferenceInput + +steady_state = InferenceInput(step={"steering": 0.25}) +assert not steady_state.requests_global_update + +changed_weather = steady_state.with_global_update({"prompt": "heavy rain"}) +assert changed_weather.requests_global_update +``` + +Because `with_step()` carries the global slot through unchanged, use +`without_global_update()` for the steady-state case; otherwise every step looks +like an update request. + +Whether a value can actually be swapped mid-rollout is declared per field: + +```python +from flashdreams.runtime import SESSION_START_ONLY, InferenceInputSchema, InputField + +schema = InferenceInputSchema( + global_fields=( + InputField(name="prompt", update_policy="step_boundary"), + InputField(name="scene_id", update_policy=SESSION_START_ONLY), + ) +) +schema.unsupported_global_updates( + InferenceInput(global_conditioning={"prompt": "heavy rain", "scene_id": "town_02"}) +) +# ("scene_id",) +``` + +`SESSION_START_ONLY` is the one reserved `update_policy` token. Everything else +in that vocabulary, and all of `lifecycle`, is open and adapter-owned; this layer +only carries it as queryable metadata. + +Steady-state steps must leave the global slot empty; otherwise every step reads +as an update request. Converters emit every window, because live control is +level-triggered: a key held across a step emits no events but still means full +throttle. + +## Raw Inputs + +`UserInputEvent` carries `timestamp_s`, `event_type`, `payload`, `source`, and +`source_event_id`. `UserInputs` holds an ordered batch plus a `snapshot` and +`metadata`, and slices to a half-open `TimeWindow`: + +```python +from flashdreams.runtime import TimeWindow, UserInputEvent, UserInputs + +inputs = UserInputs( + events=( + UserInputEvent(timestamp_s=0.0, event_type="prompt_set", + payload={"prompt": "drive forward"}), + UserInputEvent(timestamp_s=0.5, event_type="key_down", payload={"key": "w"}), + ) +) +step_window = inputs.window(TimeWindow(start_s=0.0, end_s=1.0)) +``` + +`UserInputSchema` describes what a transport, replay trace, or benchmark driver +can provide. `event_types` declares only that an event type exists; +`UserInputCapability` additionally pins the payload fields it carries, so a +converter can require `key_down` events that actually have a `key`. A bare +`event_types` entry still satisfies any consumer needing no specific payload +fields, so schemas written before capabilities existed keep working. + +## Canonical Modalities + +A `CanonicalModality` is a device-independent input: a name and the payload +fields it guarantees. Converters implement `DeviceConverter`, declaring +what raw capabilities they consume and which modality they produce. + +```python +from flashdreams.runtime import ( + DRIVER_COMMAND, InputCanonicalizer, KeyboardToDriverCommand, TimeWindow, +) + +canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) +canonicalizer.register(WheelToDriverCommand()) # a wheel is one call + +canonical = canonicalizer.canonicalize( + user_inputs, window=TimeWindow(start_s=0.0, end_s=1.0), source_schema=browser +) +canonical.values["driver_command"]["throttle"] +``` + +`DRIVER_COMMAND` is the one shipped modality. `KeyboardToDriverCommand` reuses +`KeyboardState`/`normalize_key` from `flashdreams.serving.realtime.input` and +mirrors the semantics the Omnidreams interactive-drive keyboard backend already +has. Its key bindings are data (`DEFAULT_DRIVING_BINDINGS`), and the set of +tracked keys is derived from them, so a rebound layout cannot leave an action +unreachable. + +`ScriptedModality` is the mock/replay converter. It consumes no raw +capabilities, so a benchmark or test can author a scenario at the canonical +level without knowing any device vocabulary: + +```python +canonicalizer = InputCanonicalizer([ + ScriptedModality(modality=DRIVER_COMMAND, timeline=[(0.0, full_throttle)]), +]) +canonicalizer.canonicalize( + UserInputs(), window=step_window, source_schema=UserInputSchema() +) +``` + +Application code is identical between a real run and a scripted one. + +Converters are stateful, so feed windows in session order and call +`InputCanonicalizer.reset()` at a rollout boundary. Replaying the same window +sequence reproduces the same `CanonicalInputs`. + +When several devices produce the same modality, the highest-priority one that +returned a value wins; `CanonicalInputs.metadata["canonical_sources"]` records +which device supplied each. Every feedable converter still sees each window, so +a preempted device's state stays current and unplugging the higher-priority +device does not resume from stale state. + +## Mapping And Compatibility + +`InputMapping` is the canonical-to-encoded boundary. `InputMappingSchema` is its +declarative surface: `consumes` names canonical modalities; `produces_global` +and `produces_step` name the `InferenceInput` fields it can build. + +`InputMapping.validate()` raises, which fails a run late and cannot say *which* +optional model input a source would enable or *which* missing modality makes a +required one unreachable. `check_mapping_compatibility` answers those before +expensive runtime initialization: + +```python +from flashdreams.runtime import check_mapping_set_compatibility + +compatibility = check_mapping_set_compatibility( + canonical_schema=canonicalizer.canonical_schema(browser), + inference_input_schema=adapter.inference_input_schema, + mapping_schemas=(prompt_mapping, frame_mapping, steering_mapping), +) +if not compatibility.can_drive: + compatibility.raise_if_incompatible() +``` + +`MappingCompatibility` reports `missing_modalities`, +`missing_required_model_fields`, `satisfied_required_model_fields`, +`available_optional_model_fields`, and `unavailable_mapping_schemas`. + +Compatibility is evaluated per mapping rather than over a flattened bag, so each +mapping keeps its own consumes/produces link. A mapping the source cannot feed +is dropped and reported, costing only the inputs it produced. So a dropped +mapping that fed only optional fields degrades the run instead of vetoing it, +and those fields are correctly absent from `available_optional_model_fields`; a +dropped mapping that was the only producer of a required field still blocks. + +Because a mapping consumes modalities rather than raw events, one mapping +written against `driver_command` works for a keyboard, a wheel, or any device +registered later, with no change to the mapping or the model schema. + +`undeclared_inference_inputs()` reports payload keys a mapping produced but did +not declare, which keeps hand-written schemas honest as the code drifts. + +## What This Does Not Validate + +The schemas intentionally avoid becoming a rich type system. These remain the +responsibility of the model adapter, runtime, session, or mapping: + +- tensor shape and dtype, image decode details; +- camera coordinate systems, pose and timestamp units; +- prompt-embedding swap mechanics; +- whether a model can actually apply a declared update policy at runtime; +- deep validation of scene, HD map, or actor-state data. + +The layer answers "can this source plausibly drive this model through this +mapping?" It does not replace model-owned validation. + +## Open Questions + +Tracked against the runtime API discussion, not yet settled: + +- **Alternative valid input combinations.** `InferenceInputSchema` has one flat + required set, so "accepts `{prompt}` OR `{prompt, conditioning_frame}`" cannot + be expressed. `MappingCompatibility.missing_required_model_fields` assumes a + single required set too. +- **`step()` returning a future**, for models with a dependency on their own + output. `InferenceSession.step()` is currently synchronous. +- **`Input System` ownership.** The diagrams show it pulling events, so the + Application owns an input system. `InputCanonicalizer` is currently a pure + function over a supplied window and owns no source. Whether it needs to grow + one depends on the loop-ownership decision. Mock input and key binding are + handled (`ScriptedModality`, `DEFAULT_DRIVING_BINDINGS`). + +## Owned Elsewhere + +Named here only so the boundary is explicit; these are not gaps in the input +layer: + +- **`FrameStream`**, which the architecture diagrams place between + `InferenceSession` and `Output Target`. The code writes `StepResult` straight + to `OutputTarget.write()`. Output shape is T5. +- **Declared output modalities**, so an output target or quality-eval can state + what it requires and be matched the way inputs now are. T5/T8. +- **`Application`**, the class that has-a input system, input map, global + conditioning, session, and output target. T4. +- **Loop ownership** — whether the application or the runtime/session drives the + main event loop, and whether inputs are queued and batched. + +## Validation + +```bash +.venv/bin/pytest flashdreams/tests/test_runtime_canonical.py \ + flashdreams/tests/test_runtime_input_mapping.py \ + flashdreams/tests/test_inference_runtime_api.py -q +.venv/bin/ty check flashdreams/flashdreams/runtime +``` + +At the time of writing these pass: 87 tests, and `ty` is clean. diff --git a/docs/inference_runtime_supported_inputs_inventory.md b/docs/inference_runtime_supported_inputs_inventory.md new file mode 100644 index 000000000..ebe9d853a --- /dev/null +++ b/docs/inference_runtime_supported_inputs_inventory.md @@ -0,0 +1,321 @@ + + +# Supported Model Input Inventory + +This note inventories the inputs used by the currently supported FlashDreams +runners and interactive runtimes, plus the SANA-WM input surface on `main`, then +records the T2/T3 API implications. It is intentionally about input contracts, +not tensor shape validation or model quality. + +## Inventory + +WAN 2.1 T2V, Self-Forcing WAN 2.1 T2V, Causal-Forcing T2V, +FastVideo Causal WAN 2.2 T2V, and Cosmos Predict2 T2V: + +- Source/app inputs: prompt text or prompt text file, pixel height/width, and + fps or block count depending on runner. +- Model-facing initial inputs: prompt text plus latent/output height and width + derived from run config. +- Model-facing step/update inputs: no live controls; AR loop steps with fixed + session state. + +WAN 2.1 I2V, Causal-Forcing I2V, and Cosmos Predict2 I2V: + +- Source/app inputs: prompt text or prompt file, first-frame image path or URL, + and pixel height/width. +- Model-facing initial inputs: prompt text and decoded first-frame tensor. +- Model-facing step/update inputs: no live controls. + +FlashVSR: + +- Source/app inputs: input video path or URL, chunk size, crop region, sparse + ratio, and optional output FPS. +- Model-facing initial inputs: no explicit prompt at runner time; the prompt + tensor is configured in the pipeline. Input video dimensions affect + per-video runtime/pipeline setup. +- Model-facing step/update inputs: video chunks passed to + `pipeline.generate(input=clip)`. + +LingBot CLI: + +- Source/app inputs: prompt or prompt path, first-frame image path, pose path, + intrinsics path, total blocks, dimensions, and fps. +- Model-facing initial inputs: prompt text and first-frame tensor. +- Model-facing step/update inputs: `CamCtrlInput` with intrinsics, camera poses, + and world scale. + +LingBot WebRTC: + +- Source/app inputs: session prompt, uploaded/remote/default first-frame image, + keyboard events, reset requests, text-event catalog, and trigger events. +- Model-facing initial inputs: prompt text, first-frame tensor, base text + embeddings, precomputed text-event embeddings, base intrinsics, and world + scale. +- Model-facing step/update inputs: keyboard event windows become pose segments + and camera trajectories. Text-event triggers can replace rollout text + embeddings when the model supports it. + +HY-WorldPlay WAN I2V: + +- Source/app inputs: prompt or prompt path, first-frame image path or example + image, pose string or pose JSON, memory-selection settings, dimensions, fps, + and seed. +- Model-facing initial inputs: prompt text and first-frame tensor for cache + initialization. +- Model-facing step/update inputs: pose data is bound for the rollout as action + labels, view matrices, intrinsics, and memory-selection state before AR steps. + +Omnidreams CLI: + +- Source/app inputs: shared prompt or per-camera prompts, HDMap video paths, + first-frame image/video paths, camera names, example-data UUID, and optional + embedding save/load paths. +- Model-facing initial inputs: prompt list, first-frame tensor, view names; or + precomputed text/image/negative-text embeddings. +- Model-facing step/update inputs: HDMap video chunks passed per AR step. + +Omnidreams WebRTC: + +- Source/app inputs: scene directory or scene UUID, scene variant, camera name, + prompt/first-frame assets resolved from the scene, keyboard events, reset + requests, and optional postprocess preset. +- Model-facing initial inputs: scene data, renderer, first-frame tensor, prompt, + camera calibration/extrinsics, initial ego pose, and initial timestamp. +- Model-facing step/update inputs: keyboard event windows become ego poses, + camera poses per view, and frame timestamps. The wrapper renders HDMap + conditioning internally for each step. + +Omnidreams interactive drive: + +- Source/app inputs: scene bundle, keyboard events or wheel/controller samples, + view-mode/reset/scene-exit controls, and vehicle/chunk config. +- Model-facing initial inputs: scene bundle, selected camera, prompt, initial + RGB frame, initial rig pose, and initial timestamp. +- Model-facing step/update inputs: `DriverCommand` samples become trajectory + chunks, rendered frames, and world-model conditioning. + +Template recipe: + +- Source/app inputs: synthetic runner config: batch size, height, width, context + tokens, AR steps, and seed. +- Model-facing initial inputs: synthetic transformer context, optional negative + context, height, and width. +- Model-facing step/update inputs: optional synthetic control tensor. + +WAN 2.2 TI2V pipeline config: + +- Source/app inputs: downstream runners use this rather than a standalone runner + in this tree. +- Model-facing initial inputs: prompt text and first-frame image for TI2V-style + cache initialization. +- Model-facing step/update inputs: downstream runners decide controls; + HY-WorldPlay currently binds action/camera state around it. + +SANA-WM bidirectional and streaming on `main`: + +- Source/app inputs: first-frame image path, prompt or prompt path, optional + negative prompt, camera trajectory path or action DSL, optional intrinsics + path or derived intrinsics, frame count, fps, Stage-1 sampling knobs, seed, + precision/refiner options, and streaming chunk/block settings. +- Model-facing initial inputs: decoder context such as prompt, fps, + `save_stage1`, refiner seed, sink size, and streaming refiner window/block + parameters. +- Model-facing step/update inputs: bidirectional passes one + `SanaWMI2VConditioningRequest` into the single generation step. Streaming + passes one `SanaWMStreamingI2VConditioningRequest` repeatedly; the + conditioning encoder caches rollout-wide prompt, first-frame, camera, latent + shape, and chunk-boundary state, then slices per AR chunk. +- Model-facing semantic fields include prompt, negative prompt, first frame, + camera-to-world trajectory, intrinsics vec4 sequence, frame count, fps, + sampling parameters, seed, and streaming chunking parameters. + +## API Implications + +The inventory changes the T2/T3 shape in four concrete ways. + +First, a selected mapping is often a composition. A LingBot-like run needs prompt +mapping, first-frame mapping, and keyboard-to-camera mapping. Omnidreams may add +scene selection, camera selection, and HDMap mapping. The implementation should +support checking a set of mapping schemas as one compatibility surface, while +still allowing a single mapping object when that is simpler. + +Second, `InferenceInputSchema` needs a lightweight lifecycle tag in addition to the +`initial` versus `step` phase. The phase answers when the value is needed at the +standard-loop level. The lifecycle tag distinguishes where the model adapter +uses it, such as: + +- `runtime_config`: values that affect setup before model/runtime construction, + such as FlashVSR input-video dimensions; +- `cache_init`: values passed when initializing or resetting a rollout cache, + such as prompts, first frames, view names, and precomputed embeddings; +- `rollout_binding`: values bound after cache initialization but before AR + steps, such as HY-WorldPlay action labels, camera tensors, and memory state; +- `step_input`: values consumed for one generated chunk, such as HDMap frames, + camera trajectories, driver commands, video chunks, and timestamps; +- `session_update`: values that can update an active session when supported, + such as LingBot text-event embedding swaps. + +The lifecycle tag is metadata, not a new deep type system. If both a model field +and mapping output specify lifecycle, compatibility should require them to agree. +If either side omits it, matching stays permissive for simple schemas. + +Third, `semantic_type` should be treated as a representation hint rather than a +universal semantic type. For example, `prompt` may arrive as inline text or a +path but become prompt text or text embeddings; the global conditioning frame +may arrive as a path, URL, bytes, or decoded tensor; camera motion may arrive as keys, pose JSON, +Numpy arrays, or integrated tensors. The semantic input name is still the main +contract. + +Fourth, schema objects need open-ended metadata for future adapters. This lets a +SANA-WM-like adapter advertise that `camera_trajectory_c2w` uses an +`[F,4,4]` OpenCV camera-to-world sequence, or lets another model advertise a +schema URI, units, coordinate frame, accepted file suffixes, cardinality hints, +or update notes. Metadata should remain query information and should not become +the compatibility type system. + +Fifth, `UserInputSchema` describes raw source capabilities, `CanonicalModality` +describes what an application consumes, and mapping schemas describe derived +model-facing semantics. A browser may provide `key_down`, `key_up`, +`prompt_set`, and `initial_frame_set` events. Those become canonical modalities +such as `driver_command` or `conditioning_prompt`; whether they can then drive +`steering`, `camera_trajectory`, or text embedding updates depends on the +selected mapping and model schema. + +## Implemented T2/T3 Shape + +The implementation that came out of this inventory is: + +1. Keep `UserInputEvent` and `UserInputs` as the raw event API, sliced by a + half-open `TimeWindow`. Static startup values remain timestamp-zero events. +2. Keep `UserInputSchema` lightweight and source-facing. `event_types` declares + that an event type exists; `UserInputCapability` additionally pins the + payload fields it carries. +3. Add a canonical layer between raw and encoded. `CanonicalModality` names a + device-independent input and its conditioning phase; `InputCanonicalizer` + registers per-device converters and produces `CanonicalInputs`. Applications + and mappings consume canonical inputs and never read raw device events. +4. Split `InferenceInput` (formerly `ModelInputs`) into `global_conditioning` + and `step`. A non-empty global slot mid-rollout is an update request, not a + reset; `InputField.update_policy` declares whether the model can apply it. +5. Extend `InputField` with `update_policy`, `lifecycle`, and `metadata` so + models can distinguish runtime config, cache initialization, rollout binding, + per-step inputs, and supported active-session updates. +6. Keep `InputMappingSchema` as the canonical-to-encoded boundary, with + mapping-set compatibility helpers for composed mappings. +7. Keep input names, semantic types, lifecycle labels, and metadata open-ended. + Adding a new model should usually mean adding adapter-owned schema + declarations and mappings, not changing the core input dataclasses. +8. Leave deep validation to model adapters, sessions, and mappings. The schema + layer catches obvious source/mapping/model mismatches before expensive + runtime initialization; it does not validate every tensor and coordinate + convention. + +See `docs/inference_runtime_inputs_implementation.md` for the resulting API. + +## Extensibility Contract + +The inventory above is not a vocabulary freeze. The core API does not contain a +closed enum of allowed input names. New adapters can introduce semantic field +names that match the model boundary they own. + +Use these conventions when adding future model schemas: + +- Prefer semantic names over modality names, such as `camera_trajectory_c2w` + instead of `array`, or `hdmap_frames` instead of `image`. +- Use `semantic_type` for a coarse representation hint, such as `path`, + `decoded_tensor`, `c2w_sequence`, `intrinsics_vec4_sequence`, or `embedding`. +- Use `lifecycle` to say where the adapter consumes the value, such as + `runtime_config`, `cache_init`, `rollout_binding`, `step_input`, or + `session_update`. +- Use `update_policy` to say when a value may change. `SESSION_START_ONLY` is + the one reserved token, meaning the value cannot be swapped mid-rollout. +- Use `metadata` for query hints: units, coordinate frame, shape summary, + accepted suffixes, schema URI, model family, value ranges, or cardinality. +- Keep deep validation in the adapter/mapping. The lightweight schemas answer + whether the selected source and mapping can plausibly drive the model before + expensive initialization. + +## Representative Schema Sketches + +These are not migration work for T4+, but they show that the current primitives +can describe the supported input surfaces. All use +`flashdreams.runtime.InferenceInputSchema` and `InputField`. + +```python +lingbot_model = InferenceInputSchema( + description="lingbot-world", + global_fields=( + InputField(name="prompt", lifecycle="cache_init"), + InputField(name="global_conditioning_frame", lifecycle="cache_init"), + ), + step_fields=( + InputField(name="camera_trajectory", lifecycle="step_input"), + InputField( + name="text_embeddings", + required=False, + update_policy="step_boundary", + lifecycle="session_update", + ), + ), +) +``` + +```python +omnidreams_model = InferenceInputSchema( + description="omnidreams", + global_fields=( + InputField(name="prompts", lifecycle="cache_init"), + InputField(name="global_conditioning_frames", lifecycle="cache_init"), + InputField(name="view_names", lifecycle="cache_init"), + InputField(name="text_embeddings", required=False, lifecycle="cache_init"), + InputField(name="image_embeddings", required=False, lifecycle="cache_init"), + ), + step_fields=(InputField(name="hdmap_frames", lifecycle="step_input"),), +) +``` + +```python +hy_worldplay_model = InferenceInputSchema( + description="hy-worldplay", + global_fields=( + InputField(name="prompt", lifecycle="cache_init"), + InputField(name="global_conditioning_frame", lifecycle="cache_init"), + InputField(name="action_labels", lifecycle="rollout_binding"), + InputField(name="camera_viewmats", lifecycle="rollout_binding"), + InputField(name="camera_intrinsics", lifecycle="rollout_binding"), + InputField(name="memory_config", lifecycle="rollout_binding"), + ), +) +``` + +```python +sana_wm_model = InferenceInputSchema( + description="sana-wm", + global_fields=( + InputField(name="prompt", lifecycle="cache_init"), + InputField(name="negative_prompt", required=False, lifecycle="cache_init"), + InputField(name="global_conditioning_frame", lifecycle="cache_init"), + InputField( + name="camera_trajectory_c2w", + semantic_type="c2w_sequence", + lifecycle="rollout_binding", + metadata={"shape": "[F,4,4]", "coordinates": "opencv_c2w"}, + ), + InputField( + name="camera_intrinsics_vec4", + required=False, + semantic_type="intrinsics_vec4_sequence", + lifecycle="rollout_binding", + metadata={"shape": "[F,4]"}, + ), + ), +) +``` + +SANA-WM's `stage1_sampling` and `streaming_chunking` are deliberately absent +above. They describe how to run the model rather than what conditions it, so +they belong in `InferenceConfig`, not in an input schema. Flagged here because +the runner currently threads them alongside the conditioning inputs. diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py index 03e6202b0..ab303c745 100644 --- a/flashdreams/flashdreams/runtime/__init__.py +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -7,22 +7,49 @@ intentionally additive while integrations migrate onto it. """ +from flashdreams.runtime.canonical import ( + DEFAULT_DRIVING_BINDINGS, + DRIVER_COMMAND, + DeviceConverter, + DeviceConverterSchema, + InputCanonicalizer, + KeyboardToDriverCommand, + ScriptedModality, +) from flashdreams.runtime.config import ExecutionBackend, InferenceConfig, Precision from flashdreams.runtime.inputs import ( + INPUT_PHASES, + SESSION_START_ONLY, + CanonicalInputs, + CanonicalInputSchema, + CanonicalModality, + InferenceInput, + InferenceInputSchema, InputField, - ModelInputs, - ModelInputSchema, + InputPhase, TimeWindow, + UserInputCapability, UserInputEvent, UserInputs, UserInputSchema, + validate_phase, ) from flashdreams.runtime.interfaces import ( InferenceRuntime, InferenceSession, ModelAdapter, ) -from flashdreams.runtime.mapping import IdentityInputMapping, InputMapping +from flashdreams.runtime.mapping import ( + DeclaresMappingSchema, + IdentityInputMapping, + InputMapping, + InputMappingSchema, + MappingCompatibility, + check_mapping_compatibility, + check_mapping_set_compatibility, + combine_mapping_schemas, + undeclared_inference_inputs, +) from flashdreams.runtime.metrics import ( InMemoryMetricsRecorder, MetricsRecorder, @@ -33,28 +60,50 @@ from flashdreams.runtime.types import StepRequest, StepResult __all__ = [ + "CanonicalInputs", + "CanonicalInputSchema", + "CanonicalModality", + "check_mapping_compatibility", + "check_mapping_set_compatibility", + "combine_mapping_schemas", + "DeclaresMappingSchema", + "DEFAULT_DRIVING_BINDINGS", + "DeviceConverter", + "DeviceConverterSchema", + "DRIVER_COMMAND", "ExecutionBackend", "IdentityInputMapping", "InferenceConfig", + "InferenceInput", + "InferenceInputSchema", "InferenceRuntime", "InferenceSession", "InMemoryMetricsRecorder", + "INPUT_PHASES", + "InputCanonicalizer", "InputField", "InputMapping", + "InputMappingSchema", + "InputPhase", + "KeyboardToDriverCommand", + "MappingCompatibility", "MetricsRecorder", "ModelAdapter", - "ModelInputs", - "ModelInputSchema", "NullMetricsRecorder", "NullOutputTarget", "OutputArtifact", "OutputTarget", "Precision", "RuntimeMetricSample", + "ScriptedModality", + "SESSION_START_ONLY", "StepRequest", "StepResult", "TimeWindow", + "undeclared_inference_inputs", + "UserInputCapability", "UserInputEvent", "UserInputs", "UserInputSchema", + "validate_phase", ] diff --git a/flashdreams/flashdreams/runtime/canonical.py b/flashdreams/flashdreams/runtime/canonical.py new file mode 100644 index 000000000..55f333ce7 --- /dev/null +++ b/flashdreams/flashdreams/runtime/canonical.py @@ -0,0 +1,387 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Raw device input to canonical modality conversion. + +This is the ``raw input -> canonicalized input`` leg. Applications consume +:class:`~flashdreams.runtime.inputs.CanonicalInputs`; they never read raw device +events. Adding a keyboard, gamepad, or force-feedback wheel is therefore a +:meth:`InputCanonicalizer.register` call that touches no application, mapping, +or model code. + +Converters are stateful, because HID input is edge-triggered while per-step +conditioning is level-triggered: a key held across a step emits no events yet +still means full throttle. Feed windows in session order and call +:meth:`InputCanonicalizer.reset` at a rollout boundary; replaying the same +window sequence then reproduces the same canonical inputs. + +This layer covers live user control only. Global conditioning such as a prompt +or conditioning frame is application-owned and reaches ``InferenceInput`` +directly, without passing through canonicalization. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any, Protocol, runtime_checkable + +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.inputs import ( + CanonicalInputs, + CanonicalInputSchema, + CanonicalModality, + TimeWindow, + UserInputCapability, + UserInputs, + UserInputSchema, +) +from flashdreams.serving.realtime.input import KeyboardState, normalize_key + +DriverBindings = Mapping[str, frozenset[str]] + +DEFAULT_DRIVING_BINDINGS: DriverBindings = MappingProxyType( + { + "throttle": frozenset({"w", "up"}), + "brake": frozenset({"s", "down"}), + "steer_left": frozenset({"a", "left"}), + "steer_right": frozenset({"d", "right"}), + "stop": frozenset({"space"}), + "reverse": frozenset(), + } +) +"""Default key bindings for :class:`KeyboardToDriverCommand`. + +Bindings are data so a layout can be rebound without editing the converter, and +so the set of tracked keys is derived from them rather than declared twice. +""" + +_DRIVER_ACTIONS = frozenset(DEFAULT_DRIVING_BINDINGS) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class DeviceConverterSchema: + """Metadata for one device-to-canonical-modality converter.""" + + name: str + produces: CanonicalModality + consumes: tuple[UserInputCapability, ...] = () + device_kind: str | None = None + priority: int = 0 + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) + + def __post_init__(self) -> None: + if not self.name.strip(): + raise ValueError("DeviceConverterSchema.name must be non-empty.") + if not isinstance(self.produces, CanonicalModality): + raise TypeError("produces must be a CanonicalModality object.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@runtime_checkable +class DeviceConverter(Protocol): + """Contract for turning one device's raw events into a canonical modality.""" + + @property + def schema(self) -> DeviceConverterSchema: + """Return converter metadata used for source selection.""" + ... + + def reset(self) -> None: + """Drop accumulated device state at a session or rollout boundary.""" + ... + + def convert( + self, + user_inputs: UserInputs, + window: TimeWindow, + ) -> Mapping[str, Any] | None: + """Return the modality value for ``window``, or ``None`` if inactive. + + ``user_inputs`` is already filtered to ``window``. Returning ``None`` + lets a present-but-idle device yield to a lower-priority one. + """ + ... + + +DRIVER_COMMAND = CanonicalModality( + name="driver_command", + payload_fields=frozenset({"throttle", "brake", "steer", "stop", "reverse"}), + description=( + "Normalized driving intent. throttle/brake are in [0, 1], steer is in " + "[-1, 1] with positive meaning left." + ), +) + + +class KeyboardToDriverCommand: + """Convert keyboard edges into :data:`DRIVER_COMMAND` level state. + + Mirrors the mapping the Omnidreams interactive-drive keyboard backend + already uses, so a keyboard reaches a model through the shared layer with + the same semantics it has today. + """ + + def __init__( + self, + *, + name: str = "keyboard-to-driver-command", + bindings: DriverBindings = DEFAULT_DRIVING_BINDINGS, + priority: int = 0, + ) -> None: + unknown = sorted(set(bindings) - _DRIVER_ACTIONS) + if unknown: + raise ValueError( + f"Unknown driver actions in bindings: {unknown}. " + f"Supported actions: {sorted(_DRIVER_ACTIONS)}." + ) + self._bindings = { + action: frozenset(normalize_key(key) for key in bindings.get(action, ())) + for action in _DRIVER_ACTIONS + } + # Tracked keys are derived, so they cannot drift from the bindings and + # silently make an action unreachable. + self._supported_keys = frozenset( + key for keys in self._bindings.values() for key in keys + ) + self._state = KeyboardState(supported_keys=self._supported_keys) + self._schema = DeviceConverterSchema( + name=name, + produces=DRIVER_COMMAND, + device_kind="keyboard", + priority=priority, + consumes=( + UserInputCapability( + event_type="key_down", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_type="key_up", + payload_fields=frozenset({"key"}), + ), + ), + ) + + @property + def schema(self) -> DeviceConverterSchema: + return self._schema + + def reset(self) -> None: + self._state = KeyboardState(supported_keys=self._supported_keys) + + def convert( + self, + user_inputs: UserInputs, + window: TimeWindow, + ) -> Mapping[str, Any] | None: + del window + for event in user_inputs.events: + if event.event_type not in {"key_down", "key_up"}: + continue + key = event.payload.get("key") + if not isinstance(key, str): + continue + self._state.apply_event( + event="keydown" if event.event_type == "key_down" else "keyup", + key=key, + ) + + pressed = {normalize_key(key) for key in self._state.snapshot()} + + def held(action: str) -> bool: + return bool(self._bindings[action] & pressed) + + steer = 0.0 + if held("steer_left"): + steer += 1.0 + if held("steer_right"): + steer -= 1.0 + return DRIVER_COMMAND.value( + { + "throttle": 1.0 if held("throttle") else 0.0, + "brake": 1.0 if held("brake") else 0.0, + "steer": steer, + "stop": held("stop"), + "reverse": held("reverse"), + } + ) + + +class ScriptedModality: + """Emit pre-authored canonical values, for benchmarks, replay, and tests. + + Mocking input should not require knowing the raw device vocabulary. This + converter consumes no raw capabilities, so it is feedable by any source + -- including an empty :class:`UserInputSchema` -- and application code is + identical between a real run and a scripted one. + + ``timeline`` is ``(start_s, value)`` pairs. Values are level-triggered and + held until the next entry begins, matching how live converters behave. An + entry applies to a window once it has begun by the window's end, and + ``None`` is returned for windows before the first entry. + """ + + def __init__( + self, + *, + modality: CanonicalModality, + timeline: Sequence[tuple[float, Mapping[str, Any]]], + name: str | None = None, + device_kind: str | None = "scripted", + priority: int = 0, + ) -> None: + entries = tuple(sorted(timeline, key=lambda entry: entry[0])) + for start_s, value in entries: + if start_s < 0: + raise ValueError("timeline start_s must be >= 0.") + modality.value(value) + self._entries = tuple( + (start_s, modality.value(value)) for start_s, value in entries + ) + self._modality = modality + self._schema = DeviceConverterSchema( + name=name or f"scripted-{modality.name}", + produces=modality, + device_kind=device_kind, + priority=priority, + ) + + @property + def schema(self) -> DeviceConverterSchema: + return self._schema + + def reset(self) -> None: + # The timeline is a pure function of the window, so replay is + # deterministic without any state to clear. + return None + + def convert( + self, + user_inputs: UserInputs, + window: TimeWindow, + ) -> Mapping[str, Any] | None: + del user_inputs + current: Mapping[str, Any] | None = None + for start_s, value in self._entries: + if start_s < window.end_s: + current = value + else: + break + return current + + +class InputCanonicalizer: + """Registry of device converters plus the raw-to-canonical rewrite. + + Registration is the whole extension point: a new device is a converter + registered against an existing modality, and a new modality is a converter + registered with a new :class:`CanonicalModality`. + """ + + def __init__(self, converters: Iterable[DeviceConverter] = ()) -> None: + self._converters: list[DeviceConverter] = [] + for converter in converters: + self.register(converter) + + def register(self, converter: DeviceConverter) -> None: + """Register one device converter.""" + if not isinstance(converter, DeviceConverter): + raise TypeError("converter must implement the DeviceConverter protocol.") + name = converter.schema.name + if any(existing.schema.name == name for existing in self._converters): + raise ValueError( + f"A device converter named {name!r} is already registered." + ) + self._converters.append(converter) + + @property + def converters(self) -> tuple[DeviceConverter, ...]: + """Return every registered converter.""" + return tuple(self._converters) + + def reset(self) -> None: + """Reset every registered converter's device state.""" + for converter in self._converters: + converter.reset() + + def converters_for( + self, + source_schema: UserInputSchema, + ) -> tuple[DeviceConverter, ...]: + """Return converters this source can feed, highest priority first.""" + feedable = [ + converter + for converter in self._converters + if all( + source_schema.supports(capability) + for capability in converter.schema.consumes + ) + ] + # Sort is stable, so equal-priority converters keep registration order. + return tuple(sorted(feedable, key=lambda each: -each.schema.priority)) + + def unavailable_converters( + self, + source_schema: UserInputSchema, + ) -> tuple[DeviceConverter, ...]: + """Return converters this source cannot feed, for diagnostics.""" + feedable = {id(converter) for converter in self.converters_for(source_schema)} + return tuple( + converter for converter in self._converters if id(converter) not in feedable + ) + + def canonical_schema( + self, + source_schema: UserInputSchema, + ) -> CanonicalInputSchema: + """Return the canonical modalities this raw source can supply. + + This is the boundary an application declares against. A mapping that + consumes ``driver_command`` then matches a keyboard source, a wheel + source, or any device registered later. + """ + modalities: list[CanonicalModality] = [] + for converter in self.converters_for(source_schema): + modality = converter.schema.produces + if modality not in modalities: + modalities.append(modality) + return CanonicalInputSchema( + modalities=tuple(modalities), + description=source_schema.description, + ) + + def canonicalize( + self, + user_inputs: UserInputs, + *, + window: TimeWindow, + source_schema: UserInputSchema, + ) -> CanonicalInputs: + """Convert one raw window into canonical inputs. + + Every feedable converter sees the window so its device state stays + current even while another device has precedence; that way unplugging + the higher-priority device does not resume from stale state. Among + converters producing the same modality, the highest-priority one that + returned a value wins. + """ + windowed = user_inputs.window(window) + values: dict[str, Any] = {} + sources: dict[str, str] = {} + for converter in self.converters_for(source_schema): + value = converter.convert(windowed, window) + modality = converter.schema.produces + if value is not None and modality.name not in values: + values[modality.name] = value + if converter.schema.device_kind is not None: + sources[modality.name] = converter.schema.device_kind + + metadata: dict[str, Any] = {} + if sources: + metadata["canonical_sources"] = freeze_mapping(sources) + return CanonicalInputs(values=values, metadata=metadata) diff --git a/flashdreams/flashdreams/runtime/inputs.py b/flashdreams/flashdreams/runtime/inputs.py index e14b35722..f0be31bea 100644 --- a/flashdreams/flashdreams/runtime/inputs.py +++ b/flashdreams/flashdreams/runtime/inputs.py @@ -8,10 +8,29 @@ import math from collections.abc import Iterable, Mapping from dataclasses import dataclass, field -from typing import Any +from typing import Any, Literal, cast from flashdreams.runtime._utils import freeze_mapping +InputPhase = Literal["global", "step"] + +INPUT_PHASES: tuple[InputPhase, ...] = ("global", "step") + +SESSION_START_ONLY = "session_start" +"""``InputField.update_policy`` value meaning "supply at session start only". + +``update_policy`` is otherwise an open, adapter-owned vocabulary. This is the +one reserved token, because the runtime needs to distinguish a conditioning +value that can be swapped mid-rollout from one that cannot. +""" + + +def validate_phase(value: str) -> InputPhase: + """Return ``value`` as a validated :data:`InputPhase`.""" + if value not in INPUT_PHASES: + raise ValueError(f"phase must be 'global' or 'step', got {value!r}.") + return cast(InputPhase, value) + @dataclass(frozen=True, kw_only=True, slots=True) class TimeWindow: @@ -35,16 +54,70 @@ def contains(self, timestamp_s: float) -> bool: @dataclass(frozen=True, kw_only=True, slots=True) class InputField: - """Lightweight schema field for user snapshots or model inputs.""" + """Lightweight schema field for user snapshots or model inputs. + + ``update_policy`` and ``lifecycle`` are plain query metadata. They let a + model advertise facts such as "prompt updates land at step boundaries" or + "this value is consumed at cache init" without making this layer + responsible for implementing or deeply validating that behavior. + """ name: str required: bool = True semantic_type: str | None = None + update_policy: str | None = None + lifecycle: str | None = None + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) description: str = "" def __post_init__(self) -> None: if not self.name.strip(): raise ValueError("InputField.name must be non-empty.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class UserInputCapability: + """One user event a source or mapping can provide, at payload granularity. + + ``UserInputSchema.event_types`` declares only that an event type exists. A + capability additionally pins the payload fields carried by that event, so a + mapping can state that it needs ``key_down`` events that actually carry a + ``key``. + """ + + event_type: str + semantic_type: str | None = None + payload_fields: frozenset[str] = field(default_factory=frozenset) + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) + description: str = "" + + def __post_init__(self) -> None: + if not self.event_type.strip(): + raise ValueError("UserInputCapability.event_type must be non-empty.") + for payload_field in self.payload_fields: + if not payload_field.strip(): + raise ValueError("payload field names must be non-empty.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + def is_satisfied_by(self, provider: "UserInputCapability") -> bool: + """Return whether ``provider`` can satisfy this consumed capability.""" + if self.event_type != provider.event_type: + return False + semantic_ok = ( + self.semantic_type is None + or provider.semantic_type is None + or self.semantic_type == provider.semantic_type + ) + return semantic_ok and self.payload_fields.issubset(provider.payload_fields) @dataclass(frozen=True, kw_only=True, slots=True) @@ -53,6 +126,7 @@ class UserInputSchema: event_types: frozenset[str] = field(default_factory=frozenset) snapshot_fields: tuple[InputField, ...] = () + capabilities: tuple[UserInputCapability, ...] = () description: str = "" def supports_event_types(self, event_types: Iterable[str]) -> bool: @@ -60,7 +134,63 @@ def supports_event_types(self, event_types: Iterable[str]) -> bool: requested = frozenset(event_types) if not requested: return True - return requested.issubset(self.event_types) + return requested.issubset(self.declared_event_types()) + + def declared_event_types(self) -> frozenset[str]: + """Return event types from ``event_types`` and from ``capabilities``.""" + return self.event_types | frozenset( + capability.event_type for capability in self.capabilities + ) + + def declared_capabilities(self) -> tuple[UserInputCapability, ...]: + """Return capabilities, widened with bare ``event_types`` entries. + + A plain ``event_types`` entry carries no payload promise, so it is + modeled as a capability with no payload fields. Coarse schemas written + before capabilities existed therefore still satisfy any consumer that + does not require specific payload fields. + """ + declared = list(self.capabilities) + covered = {capability.event_type for capability in declared} + declared.extend( + UserInputCapability(event_type=event_type) + for event_type in sorted(self.event_types - covered) + ) + return tuple(declared) + + def supports(self, capability: UserInputCapability) -> bool: + """Return whether this source can satisfy ``capability``.""" + return any( + capability.is_satisfied_by(provider) + for provider in self.declared_capabilities() + ) + + def validate_event(self, event: "UserInputEvent") -> None: + """Validate one event against the event types this source declares.""" + matching = [ + capability + for capability in self.declared_capabilities() + if capability.event_type == event.event_type + ] + if not matching: + raise ValueError( + f"User input source does not provide event type {event.event_type!r}." + ) + payload_keys = set(event.payload) + if not any( + capability.payload_fields.issubset(payload_keys) for capability in matching + ): + expected = sorted( + { + payload_field + for capability in matching + for payload_field in capability.payload_fields + } + ) + raise ValueError( + f"Event {event.event_type!r} payload is missing required " + f"fields: {expected}." + ) def missing_snapshot(self, inputs: "UserInputs") -> tuple[str, ...]: """Return required snapshot fields absent from ``inputs``.""" @@ -74,10 +204,10 @@ def require_snapshot(self, inputs: "UserInputs") -> None: @dataclass(frozen=True, kw_only=True, slots=True) -class ModelInputSchema: +class InferenceInputSchema: """Minimal metadata for model-facing initial and per-step inputs.""" - initial_fields: tuple[InputField, ...] = () + global_fields: tuple[InputField, ...] = () """Model inputs required before starting the initial generation/session.""" step_fields: tuple[InputField, ...] = () @@ -85,21 +215,81 @@ class ModelInputSchema: description: str = "" - def missing_initial(self, inputs: "ModelInputs") -> tuple[str, ...]: + def fields_for(self, phase: InputPhase) -> tuple[InputField, ...]: + """Return every declared field for ``phase``.""" + return ( + self.global_fields + if validate_phase(phase) == "global" + else self.step_fields + ) + + def required_fields( + self, + phase: InputPhase | None = None, + ) -> tuple[tuple[InputPhase, InputField], ...]: + """Return required fields as ``(phase, field)``, optionally filtered.""" + return self._select(phase, required=True) + + def optional_fields( + self, + phase: InputPhase | None = None, + ) -> tuple[tuple[InputPhase, InputField], ...]: + """Return optional fields as ``(phase, field)``, optionally filtered.""" + return self._select(phase, required=False) + + def field_for(self, *, name: str, phase: InputPhase) -> InputField | None: + """Return one declared field, if present.""" + for input_field in self.fields_for(phase): + if input_field.name == name: + return input_field + return None + + def _select( + self, + phase: InputPhase | None, + *, + required: bool, + ) -> tuple[tuple[InputPhase, InputField], ...]: + phases = INPUT_PHASES if phase is None else (validate_phase(phase),) + return tuple( + (each_phase, input_field) + for each_phase in phases + for input_field in self.fields_for(each_phase) + if input_field.required is required + ) + + def unsupported_global_updates(self, inputs: "InferenceInput") -> tuple[str, ...]: + """Return requested conditioning updates this model cannot apply. + + A field whose ``update_policy`` is :data:`SESSION_START_ONLY` can be + supplied when the session starts but not changed mid-rollout. Any other + policy, including ``None``, is treated as permissive here; the adapter + still owns whether the swap actually succeeds. + """ + return tuple( + name + for name in inputs.global_conditioning + if (declared := self.field_for(name=name, phase="global")) is not None + and declared.update_policy == SESSION_START_ONLY + ) + + def missing_global(self, inputs: "InferenceInput") -> tuple[str, ...]: """Return required initial fields absent from ``inputs``.""" - return _missing_required(self.initial_fields, inputs.initial) + return _missing_required(self.global_fields, inputs.global_conditioning) - def missing_step(self, inputs: "ModelInputs") -> tuple[str, ...]: + def missing_step(self, inputs: "InferenceInput") -> tuple[str, ...]: """Return required per-step fields absent from ``inputs``.""" return _missing_required(self.step_fields, inputs.step) - def require_initial(self, inputs: "ModelInputs") -> None: + def require_global(self, inputs: "InferenceInput") -> None: """Raise if required initial fields are absent.""" - missing = self.missing_initial(inputs) + missing = self.missing_global(inputs) if missing: - raise ValueError(f"Missing required initial model input(s): {missing}") + raise ValueError( + f"Missing required global conditioning input(s): {missing}" + ) - def require_step(self, inputs: "ModelInputs") -> None: + def require_step(self, inputs: "InferenceInput") -> None: """Raise if required per-step fields are absent.""" missing = self.missing_step(inputs) if missing: @@ -171,23 +361,154 @@ def window(self, time_window: TimeWindow) -> "UserInputs": @dataclass(frozen=True, kw_only=True, slots=True) -class ModelInputs: - """Model-facing payloads split by initial and per-step use.""" +class CanonicalModality: + """A device-independent user input an application consumes. + + This is the middle layer of ``raw input -> canonicalized input -> encoded + inference input``. Applications and benchmarks declare and consume + modalities; they never read raw device events, so adding a new device is a + converter registration rather than an application change. + + Modalities describe live user control only. Global conditioning such as a + prompt or conditioning frame is application-owned and reaches + :class:`InferenceInput` directly, without passing through this layer. + """ + + name: str + payload_fields: frozenset[str] = field(default_factory=frozenset) + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) + description: str = "" + + def __post_init__(self) -> None: + if not self.name.strip(): + raise ValueError("CanonicalModality.name must be non-empty.") + for payload_field in self.payload_fields: + if not payload_field.strip(): + raise ValueError("payload field names must be non-empty.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + def is_satisfied_by(self, provider: "CanonicalModality") -> bool: + """Return whether ``provider`` can satisfy this consumed modality.""" + return self.name == provider.name and self.payload_fields.issubset( + provider.payload_fields + ) + + def value(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: + """Return ``payload`` frozen, checking it covers this modality.""" + missing = sorted(self.payload_fields - set(payload)) + if missing: + raise ValueError( + f"Canonical modality {self.name!r} requires payload fields " + f"{missing}, which the converter did not produce." + ) + return freeze_mapping(payload) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class CanonicalInputSchema: + """Canonical modalities an application can be fed by a given source.""" + + modalities: tuple[CanonicalModality, ...] = () + description: str = "" + + def supports(self, modality: CanonicalModality) -> bool: + """Return whether this source can supply ``modality``.""" + return any(modality.is_satisfied_by(provided) for provided in self.modalities) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class CanonicalInputs: + """Canonicalized user input for one step, keyed by modality name. + + Values are level-triggered and normally present every step: a key held down + emits no events but still means full throttle. Global conditioning does not + appear here; it is application-owned and reaches :class:`InferenceInput` + directly. + """ __hash__ = None - initial: Mapping[str, Any] = field(default_factory=dict) + values: Mapping[str, Any] = field(default_factory=dict) + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__(self, "values", freeze_mapping(self.values)) + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class InferenceInput: + """Encoded inputs for one :class:`InferenceSession` call. + + Two conditioning slots: + + - ``global_conditioning``: values that condition the whole rollout, such as + the conditioning frame or prompt. Normally supplied when the session + starts. + - ``step``: values needed to generate the next chunk or frame. + + A non-empty ``global_conditioning`` on a mid-rollout input is an *update + request*, not a reset. The session should apply it when the model supports + that; resetting rollout state is a separate, explicit + :meth:`InferenceSession.reset` call. Whether a given value can be updated + mid-rollout is declared per field by ``InputField.update_policy``; see + :meth:`InferenceInputSchema.unsupported_global_updates`. + """ + + __hash__ = None + + global_conditioning: Mapping[str, Any] = field(default_factory=dict) step: Mapping[str, Any] = field(default_factory=dict) metadata: Mapping[str, Any] = field(default_factory=dict) def __post_init__(self) -> None: - object.__setattr__(self, "initial", freeze_mapping(self.initial)) + object.__setattr__( + self, "global_conditioning", freeze_mapping(self.global_conditioning) + ) object.__setattr__(self, "step", freeze_mapping(self.step)) object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) - def with_step(self, step: Mapping[str, Any]) -> "ModelInputs": - """Return a copy with replaced per-step payload.""" - return ModelInputs(initial=self.initial, step=step, metadata=self.metadata) + @property + def requests_global_update(self) -> bool: + """Return whether this input asks the session to update conditioning.""" + return bool(self.global_conditioning) + + def with_step(self, step: Mapping[str, Any]) -> "InferenceInput": + """Return a copy with replaced per-step payload. + + The global slot is carried through unchanged, so a mid-rollout input + built this way keeps whatever update request it already had. Use + :meth:`without_global_update` for the common steady-state case. + """ + return InferenceInput( + global_conditioning=self.global_conditioning, + step=step, + metadata=self.metadata, + ) + + def with_global_update( + self, global_conditioning: Mapping[str, Any] + ) -> "InferenceInput": + """Return a copy requesting a mid-rollout conditioning update.""" + return InferenceInput( + global_conditioning=global_conditioning, + step=self.step, + metadata=self.metadata, + ) + + def without_global_update(self) -> "InferenceInput": + """Return a copy that requests no conditioning update.""" + return InferenceInput(step=self.step, metadata=self.metadata) + + def for_phase(self, phase: InputPhase) -> Mapping[str, Any]: + """Return the payload mapping for ``phase``.""" + return ( + self.global_conditioning if validate_phase(phase) == "global" else self.step + ) def _missing_required( diff --git a/flashdreams/flashdreams/runtime/interfaces.py b/flashdreams/flashdreams/runtime/interfaces.py index 9b6a064fd..852a77f1c 100644 --- a/flashdreams/flashdreams/runtime/interfaces.py +++ b/flashdreams/flashdreams/runtime/interfaces.py @@ -9,9 +9,9 @@ from flashdreams.runtime.config import InferenceConfig from flashdreams.runtime.inputs import ( - ModelInputs, - ModelInputSchema, - UserInputSchema, + CanonicalInputSchema, + InferenceInput, + InferenceInputSchema, ) from flashdreams.runtime.mapping import InputMapping from flashdreams.runtime.types import StepRequest, StepResult @@ -25,11 +25,11 @@ def next_step_request(self) -> StepRequest | None: """Describe the next step's inputs, or return ``None`` when complete.""" ... - def step(self, inputs: ModelInputs) -> StepResult: + def step(self, inputs: InferenceInput) -> StepResult: """Run one sequential inference step.""" ... - def reset(self, inputs: ModelInputs | None = None) -> None: + def reset(self, inputs: InferenceInput | None = None) -> None: """Reset this session's rollout state when the backend supports it.""" ... @@ -42,8 +42,8 @@ def close(self) -> None: class InferenceRuntime(Protocol): """Heavyweight reusable runtime created from :class:`InferenceConfig`.""" - def start_session(self, inputs: ModelInputs) -> InferenceSession: - """Create an isolated session from initial model inputs.""" + def start_session(self, inputs: InferenceInput) -> InferenceSession: + """Create an isolated session from global conditioning inputs.""" ... def close(self) -> None: @@ -56,10 +56,10 @@ def close(self) -> None: class ModelAdapter(Protocol): """Model-specific boundary that declares defaults and creates runtimes. - Adapters declare model-facing input requirements, optional user-input - capabilities, and an optional default mapping between the two. Runtime, - application, or benchmark code may override that mapping while preserving the - same ``UserInputs`` to ``ModelInputs`` boundary. + Adapters declare model-facing input requirements, the canonical modalities + their default mapping consumes, and an optional default mapping between the + two. Runtime, application, or benchmark code may override that mapping while + preserving the same ``CanonicalInputs`` to ``InferenceInput`` boundary. """ @property @@ -68,17 +68,17 @@ def model_id(self) -> str: ... @property - def model_input_schema(self) -> ModelInputSchema: + def inference_input_schema(self) -> InferenceInputSchema: """Model-facing initial and per-step input requirements.""" ... @property - def user_input_schema(self) -> UserInputSchema | None: - """User inputs supported by the adapter's default mapping, if any.""" + def canonical_input_schema(self) -> CanonicalInputSchema | None: + """Canonical modalities the adapter's default mapping consumes.""" ... def default_input_mapping(self) -> InputMapping | None: - """Return the model-provided default user-to-model mapping, if any.""" + """Return the model-provided default canonical-to-model mapping.""" ... def validate_config(self, config: InferenceConfig) -> None: diff --git a/flashdreams/flashdreams/runtime/mapping.py b/flashdreams/flashdreams/runtime/mapping.py index 756351081..94f481406 100644 --- a/flashdreams/flashdreams/runtime/mapping.py +++ b/flashdreams/flashdreams/runtime/mapping.py @@ -1,17 +1,24 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Input mapping boundary from user input windows to model inputs.""" +"""Input mapping boundary from canonical inputs to encoded inference inputs.""" from __future__ import annotations -from typing import Protocol, runtime_checkable +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field, replace +from typing import Any, Protocol, runtime_checkable +from flashdreams.runtime._utils import freeze_mapping from flashdreams.runtime.inputs import ( - ModelInputs, - ModelInputSchema, - UserInputs, - UserInputSchema, + INPUT_PHASES, + CanonicalInputs, + CanonicalInputSchema, + CanonicalModality, + InferenceInput, + InferenceInputSchema, + InputField, + InputPhase, ) from flashdreams.runtime.types import StepRequest @@ -28,28 +35,28 @@ class InputMapping(Protocol): def validate( self, *, - user_schema: UserInputSchema | None = None, - model_schema: ModelInputSchema | None = None, + canonical_schema: CanonicalInputSchema | None = None, + inference_input_schema: InferenceInputSchema | None = None, ) -> None: """Fail early for obvious app, event-source, and model mismatches.""" ... - def map_initial_inputs( + def map_global_inputs( self, *, - user_inputs: UserInputs, - model_inputs: ModelInputs, - ) -> ModelInputs: - """Build initial model inputs before a session starts.""" + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + ) -> InferenceInput: + """Build global conditioning inputs before a session starts.""" ... def map_step_inputs( self, *, - user_inputs: UserInputs, - model_inputs: ModelInputs, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, request: StepRequest, - ) -> ModelInputs: + ) -> InferenceInput: """Build model inputs for one session step from the current input window.""" ... @@ -60,26 +67,315 @@ class IdentityInputMapping: def validate( self, *, - user_schema: UserInputSchema | None = None, - model_schema: ModelInputSchema | None = None, + canonical_schema: CanonicalInputSchema | None = None, + inference_input_schema: InferenceInputSchema | None = None, ) -> None: - del user_schema, model_schema + del canonical_schema, inference_input_schema - def map_initial_inputs( + def map_global_inputs( self, *, - user_inputs: UserInputs, - model_inputs: ModelInputs, - ) -> ModelInputs: - del user_inputs - return model_inputs + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + ) -> InferenceInput: + del canonical_inputs + return inference_input def map_step_inputs( self, *, - user_inputs: UserInputs, - model_inputs: ModelInputs, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, request: StepRequest, - ) -> ModelInputs: - del user_inputs, request - return model_inputs + ) -> InferenceInput: + del canonical_inputs, request + return inference_input + + +@dataclass(frozen=True, kw_only=True, slots=True) +class InputMappingSchema: + """Declarative compatibility surface for one mapping. + + ``InputMapping.validate`` fails a run late and opaquely: it raises, but it + cannot answer which optional model inputs a source would enable, or which + missing user capability is responsible for an unreachable model input. This + schema makes those questions answerable before runtime initialization. + """ + + name: str = "input-mapping" + consumes: tuple[CanonicalModality, ...] = () + produces_global: tuple[InputField, ...] = () + produces_step: tuple[InputField, ...] = () + metadata: Mapping[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) + + def __post_init__(self) -> None: + if not self.name.strip(): + raise ValueError("InputMappingSchema.name must be non-empty.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + def produces_for(self, phase: InputPhase) -> tuple[InputField, ...]: + """Return the fields this mapping produces for ``phase``.""" + return self.produces_global if phase == "global" else self.produces_step + + def can_produce(self, phase: InputPhase, required: InputField) -> bool: + """Return whether this mapping can produce ``required`` in ``phase``.""" + return any( + _field_matches(produced, required) for produced in self.produces_for(phase) + ) + + +def _field_matches(produced: InputField, required: InputField) -> bool: + if produced.name != required.name: + return False + semantic_ok = ( + produced.semantic_type is None + or required.semantic_type is None + or produced.semantic_type == required.semantic_type + ) + lifecycle_ok = ( + produced.lifecycle is None + or required.lifecycle is None + or produced.lifecycle == required.lifecycle + ) + return semantic_ok and lifecycle_ok + + +@dataclass(frozen=True, kw_only=True, slots=True) +class MappingCompatibility: + """Compatibility report for one source, model schema, and mapping set. + + Mappings whose consumed capabilities the source cannot provide are reported + in ``unavailable_mapping_schemas`` and excluded from the satisfied/available + reports, so those lists only name model inputs that can really be produced. + """ + + __hash__ = None + + canonical_schema: CanonicalInputSchema + inference_input_schema: InferenceInputSchema + mapping_schema: InputMappingSchema + missing_modalities: tuple[CanonicalModality, ...] = () + missing_required_model_fields: tuple[tuple[InputPhase, InputField], ...] = () + satisfied_required_model_fields: tuple[tuple[InputPhase, InputField], ...] = () + available_optional_model_fields: tuple[tuple[InputPhase, InputField], ...] = () + unavailable_mapping_schemas: tuple[InputMappingSchema, ...] = () + + @property + def can_drive(self) -> bool: + """Return whether this source can drive this model through the mapping. + + A mapping the source cannot feed does not block the run unless it was + the only way to produce a required model input. + """ + return not (self.missing_required_model_fields or self.missing_modalities) + + @property + def unavailable_mapping_names(self) -> tuple[str, ...]: + """Return names of mappings dropped because the source cannot feed them.""" + return tuple(schema.name for schema in self.unavailable_mapping_schemas) + + def raise_if_incompatible(self) -> None: + """Raise a compact error when this mapping cannot drive the model.""" + if self.can_drive: + return + problems: list[str] = [] + if self.missing_modalities: + missing = ", ".join(modality.name for modality in self.missing_modalities) + problems.append(f"missing canonical modalities: {missing}") + if self.missing_required_model_fields: + missing = ", ".join( + f"{phase}:{input_field.name}" + for phase, input_field in self.missing_required_model_fields + ) + problems.append(f"missing required model inputs: {missing}") + if self.unavailable_mapping_schemas: + problems.append( + "unavailable mappings: " + ", ".join(self.unavailable_mapping_names) + ) + raise ValueError( + f"Input mapping {self.mapping_schema.name!r} cannot drive this model " + f"from the selected source: " + "; ".join(problems) + ) + + +def _source_can_feed( + canonical_schema: CanonicalInputSchema, + mapping_schema: InputMappingSchema, +) -> bool: + return all( + canonical_schema.supports(modality) for modality in mapping_schema.consumes + ) + + +def combine_mapping_schemas( + mapping_schemas: Sequence[InputMappingSchema], + *, + name: str = "input-mapping-set", +) -> InputMappingSchema: + """Combine independently declared mappings into one compatibility surface. + + Duplicates are collapsed. Because ``metadata`` is excluded from equality, + the metadata of collapsed duplicates is merged rather than dropped, with the + first declaration winning on conflicting keys. + """ + consumes: list[CanonicalModality] = [] + produces: dict[InputPhase, list[InputField]] = {"global": [], "step": []} + + def _merge(target: list[Any], value: Any) -> None: + for index, existing in enumerate(target): + if existing == value: + if value.metadata: + target[index] = replace( + existing, + metadata={**dict(value.metadata), **dict(existing.metadata)}, + ) + return + target.append(value) + + for mapping_schema in mapping_schemas: + if not isinstance(mapping_schema, InputMappingSchema): + raise TypeError("mapping_schemas must contain InputMappingSchema objects.") + for modality in mapping_schema.consumes: + _merge(consumes, modality) + for phase in INPUT_PHASES: + for input_field in mapping_schema.produces_for(phase): + _merge(produces[phase], input_field) + + return InputMappingSchema( + name=name, + consumes=tuple(consumes), + produces_global=tuple(produces["global"]), + produces_step=tuple(produces["step"]), + ) + + +def _build_compatibility( + *, + canonical_schema: CanonicalInputSchema, + inference_input_schema: InferenceInputSchema, + mapping_schemas: Sequence[InputMappingSchema], + reported_schema: InputMappingSchema, +) -> MappingCompatibility: + feedable: list[InputMappingSchema] = [] + unavailable: list[InputMappingSchema] = [] + for mapping_schema in mapping_schemas: + if _source_can_feed(canonical_schema, mapping_schema): + feedable.append(mapping_schema) + else: + unavailable.append(mapping_schema) + + usable = combine_mapping_schemas(feedable, name=reported_schema.name) + required = inference_input_schema.required_fields() + missing_required = tuple( + (phase, input_field) + for phase, input_field in required + if not usable.can_produce(phase, input_field) + ) + satisfied_required = tuple( + (phase, input_field) + for phase, input_field in required + if usable.can_produce(phase, input_field) + ) + available_optional = tuple( + (phase, input_field) + for phase, input_field in inference_input_schema.optional_fields() + if usable.can_produce(phase, input_field) + ) + + # Only capabilities that block a required model input make the mapping + # unusable. A dropped mapping that fed nothing but optional fields degrades + # the run instead of vetoing it. + missing_modalities: list[CanonicalModality] = [] + for mapping_schema in unavailable: + if not any( + mapping_schema.can_produce(phase, input_field) + for phase, input_field in missing_required + ): + continue + for modality in mapping_schema.consumes: + if canonical_schema.supports(modality) or modality in missing_modalities: + continue + missing_modalities.append(modality) + + return MappingCompatibility( + canonical_schema=canonical_schema, + inference_input_schema=inference_input_schema, + mapping_schema=reported_schema, + missing_modalities=tuple(missing_modalities), + missing_required_model_fields=missing_required, + satisfied_required_model_fields=satisfied_required, + available_optional_model_fields=available_optional, + unavailable_mapping_schemas=tuple(unavailable), + ) + + +def check_mapping_compatibility( + *, + canonical_schema: CanonicalInputSchema, + inference_input_schema: InferenceInputSchema, + mapping_schema: InputMappingSchema, +) -> MappingCompatibility: + """Check whether a user-input source can drive a model through a mapping.""" + if not isinstance(mapping_schema, InputMappingSchema): + raise TypeError("mapping_schema must be an InputMappingSchema object.") + return _build_compatibility( + canonical_schema=canonical_schema, + inference_input_schema=inference_input_schema, + mapping_schemas=(mapping_schema,), + reported_schema=mapping_schema, + ) + + +def check_mapping_set_compatibility( + *, + canonical_schema: CanonicalInputSchema, + inference_input_schema: InferenceInputSchema, + mapping_schemas: Sequence[InputMappingSchema], + name: str = "input-mapping-set", +) -> MappingCompatibility: + """Check compatibility for a composed set of mappings. + + Each mapping keeps its own consumes/produces link, so a mapping the source + cannot feed only costs the model inputs that mapping produced. + """ + mapping_schemas = tuple(mapping_schemas) + return _build_compatibility( + canonical_schema=canonical_schema, + inference_input_schema=inference_input_schema, + mapping_schemas=mapping_schemas, + reported_schema=combine_mapping_schemas(mapping_schemas, name=name), + ) + + +def undeclared_inference_inputs( + inputs: InferenceInput, + mapping_schema: InputMappingSchema, +) -> tuple[tuple[InputPhase, str], ...]: + """Return payload keys a mapping produced but did not declare. + + Mapping schemas are hand-written, so they drift from what + ``map_global_inputs``/``map_step_inputs`` actually return. Mapping tests + can use this to keep the declared compatibility surface honest. + """ + return tuple( + (phase, key) + for phase in INPUT_PHASES + for key in inputs.for_phase(phase) + if not any( + declared.name == key for declared in mapping_schema.produces_for(phase) + ) + ) + + +@runtime_checkable +class DeclaresMappingSchema(Protocol): + """Optional refinement of :class:`InputMapping` that declares its surface.""" + + @property + def mapping_schema(self) -> InputMappingSchema: + """Return the declarative compatibility surface for this mapping.""" + ... diff --git a/flashdreams/flashdreams/runtime/types.py b/flashdreams/flashdreams/runtime/types.py index 52bf82166..467753026 100644 --- a/flashdreams/flashdreams/runtime/types.py +++ b/flashdreams/flashdreams/runtime/types.py @@ -10,7 +10,7 @@ from typing import Any from flashdreams.runtime._utils import freeze_mapping -from flashdreams.runtime.inputs import ModelInputSchema, TimeWindow +from flashdreams.runtime.inputs import InferenceInputSchema, TimeWindow @dataclass(frozen=True, kw_only=True, slots=True) @@ -24,7 +24,7 @@ class StepRequest: __hash__ = None step_index: int - model_input_schema: ModelInputSchema | None = None + inference_input_schema: InferenceInputSchema | None = None user_input_window: TimeWindow | None = None metadata: Mapping[str, Any] = field(default_factory=dict) diff --git a/flashdreams/tests/test_inference_runtime_api.py b/flashdreams/tests/test_inference_runtime_api.py index 1474383a0..edfafa634 100644 --- a/flashdreams/tests/test_inference_runtime_api.py +++ b/flashdreams/tests/test_inference_runtime_api.py @@ -9,17 +9,20 @@ import pytest from flashdreams.runtime import ( + CanonicalInputs, + CanonicalInputSchema, IdentityInputMapping, InferenceConfig, + InferenceInput, + InferenceInputSchema, InferenceRuntime, InferenceSession, InMemoryMetricsRecorder, + InputCanonicalizer, InputField, InputMapping, MetricsRecorder, ModelAdapter, - ModelInputs, - ModelInputSchema, NullOutputTarget, OutputArtifact, OutputTarget, @@ -27,6 +30,7 @@ StepRequest, StepResult, TimeWindow, + UserInputCapability, UserInputEvent, UserInputs, UserInputSchema, @@ -35,6 +39,18 @@ pytestmark = pytest.mark.ci_cpu +_SESSION_HORIZON_S = 3600.0 + +_KEYBOARD_SOURCE = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="keyboard.keydown", payload_fields=frozenset({"key"}) + ), + ) +) +_KEYBOARD_CANONICALIZER = InputCanonicalizer() + + def test_inference_config_keeps_runtime_settings_separate() -> None: denied_app_fields = {"prompt", "output_dir", "browser_settings"} config = InferenceConfig( @@ -90,17 +106,19 @@ def test_runtime_metric_sample_rejects_bool_values() -> None: RuntimeMetricSample(name="sample", value=True) -def test_model_input_schema_validates_initial_and_step_payloads() -> None: - schema = ModelInputSchema( - initial_fields=( +def test_inference_input_schema_validates_initial_and_step_payloads() -> None: + schema = InferenceInputSchema( + global_fields=( InputField(name="prompt"), - InputField(name="first_frame"), + InputField(name="global_conditioning_frame"), ), step_fields=(InputField(name="camera_poses"),), ) - inputs = ModelInputs(initial={"prompt": "drive", "first_frame": object()}) + inputs = InferenceInput( + global_conditioning={"prompt": "drive", "global_conditioning_frame": object()} + ) - schema.require_initial(inputs) + schema.require_global(inputs) assert schema.missing_step(inputs) == ("camera_poses",) with pytest.raises(ValueError, match="camera_poses"): @@ -164,25 +182,27 @@ def test_user_input_schema_validates_required_snapshot_fields() -> None: schema.require_snapshot(UserInputs()) -def test_identity_input_mapping_leaves_model_inputs_unchanged() -> None: +def test_identity_input_mapping_leaves_inference_input_unchanged() -> None: mapping = IdentityInputMapping() - model_inputs = ModelInputs(initial={"prompt": "fixed"}, step={"hdmap": object()}) + inference_input = InferenceInput( + global_conditioning={"prompt": "fixed"}, step={"hdmap": object()} + ) request = StepRequest(step_index=0) assert ( - mapping.map_initial_inputs( - user_inputs=UserInputs(), - model_inputs=model_inputs, + mapping.map_global_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=inference_input, ) - is model_inputs + is inference_input ) assert ( mapping.map_step_inputs( - user_inputs=UserInputs(), - model_inputs=model_inputs, + canonical_inputs=CanonicalInputs(), + inference_input=inference_input, request=request, ) - is model_inputs + is inference_input ) @@ -258,7 +278,7 @@ def test_runtime_api_components_compose_for_sequential_session() -> None: ), ) ) - model_inputs = ModelInputs(initial={"prompt": "drive forward"}) + inference_input = InferenceInput(global_conditioning={"prompt": "drive forward"}) output = NullOutputTarget(store_results=True) metrics = InMemoryMetricsRecorder() @@ -269,8 +289,10 @@ def test_runtime_api_components_compose_for_sequential_session() -> None: adapter=adapter, config=config, mapping=mapping, + canonicalizer=_KEYBOARD_CANONICALIZER, + source_schema=_KEYBOARD_SOURCE, user_inputs=user_inputs, - model_inputs=model_inputs, + inference_input=inference_input, output=output, metrics=metrics, ) @@ -291,8 +313,10 @@ def test_reference_loop_validates_mapping_before_runtime_creation() -> None: adapter=adapter, config=InferenceConfig(model_id="fake-model"), mapping=mapping, + canonicalizer=_KEYBOARD_CANONICALIZER, + source_schema=_KEYBOARD_SOURCE, user_inputs=UserInputs(), - model_inputs=ModelInputs(initial={"prompt": "drive forward"}), + inference_input=InferenceInput(global_conditioning={"prompt": "drive forward"}), output=NullOutputTarget(), metrics=InMemoryMetricsRecorder(), ) @@ -311,8 +335,12 @@ def test_reference_loop_closes_runtime_when_session_start_fails() -> None: adapter=adapter, config=InferenceConfig(model_id="fake-model"), mapping=IdentityInputMapping(), + canonicalizer=_KEYBOARD_CANONICALIZER, + source_schema=_KEYBOARD_SOURCE, user_inputs=UserInputs(), - model_inputs=ModelInputs(initial={"prompt": "drive forward"}), + inference_input=InferenceInput( + global_conditioning={"prompt": "drive forward"} + ), output=output, metrics=metrics, ) @@ -328,18 +356,24 @@ def _drive_two_step_session( adapter: ModelAdapter, config: InferenceConfig, mapping: InputMapping, + canonicalizer: InputCanonicalizer, + source_schema: UserInputSchema, user_inputs: UserInputs, - model_inputs: ModelInputs, + inference_input: InferenceInput, output: OutputTarget, metrics: MetricsRecorder, ) -> None: mapping.validate( - user_schema=adapter.user_input_schema, - model_schema=adapter.model_input_schema, + canonical_schema=adapter.canonical_input_schema, + inference_input_schema=adapter.inference_input_schema, ) - initial_inputs = mapping.map_initial_inputs( - user_inputs=user_inputs, - model_inputs=model_inputs, + initial_inputs = mapping.map_global_inputs( + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=TimeWindow(start_s=0.0, end_s=_SESSION_HORIZON_S), + source_schema=source_schema, + ), + inference_input=inference_input, ) runtime = adapter.create_runtime(config) session: InferenceSession | None = None @@ -350,13 +384,16 @@ def _drive_two_step_session( output_opened = True while (request := session.next_step_request()) is not None: step_inputs = mapping.map_step_inputs( - user_inputs=( - user_inputs.window(request.user_input_window) - if request.user_input_window is not None - else user_inputs + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=request.user_input_window + or TimeWindow(start_s=0.0, end_s=_SESSION_HORIZON_S), + source_schema=source_schema, ), - model_inputs=ModelInputs( - initial=initial_inputs.initial, + # The global slot stays empty in steady state. A mapping that + # sees ``canonical_inputs.has_global_change`` fills it via + # ``with_global_update`` to request a mid-rollout swap. + inference_input=InferenceInput( step={"chunk_index": request.step_index}, ), request=request, @@ -379,11 +416,11 @@ def _drive_two_step_session( class _FakeAdapter: model_id = "fake-model" - model_input_schema = ModelInputSchema( - initial_fields=(InputField(name="prompt"),), + inference_input_schema = InferenceInputSchema( + global_fields=(InputField(name="prompt"),), step_fields=(InputField(name="chunk_index"),), ) - user_input_schema = UserInputSchema(event_types=frozenset({"keyboard.keydown"})) + canonical_input_schema = CanonicalInputSchema() def default_input_mapping(self) -> InputMapping: return IdentityInputMapping() @@ -394,31 +431,31 @@ def validate_config(self, config: InferenceConfig) -> None: def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: self.validate_config(config) - return _FakeRuntime(model_input_schema=self.model_input_schema) + return _FakeRuntime(inference_input_schema=self.inference_input_schema) class _FakeRuntime: - def __init__(self, *, model_input_schema: ModelInputSchema) -> None: - self._model_input_schema = model_input_schema + def __init__(self, *, inference_input_schema: InferenceInputSchema) -> None: + self._inference_input_schema = inference_input_schema self.closed = False - def start_session(self, inputs: ModelInputs) -> InferenceSession: - self._model_input_schema.require_initial(inputs) - return _FakeSession(model_input_schema=self._model_input_schema) + def start_session(self, inputs: InferenceInput) -> InferenceSession: + self._inference_input_schema.require_global(inputs) + return _FakeSession(inference_input_schema=self._inference_input_schema) def close(self) -> None: self.closed = True class _FailingRuntime(_FakeRuntime): - def start_session(self, inputs: ModelInputs) -> InferenceSession: + def start_session(self, inputs: InferenceInput) -> InferenceSession: del inputs raise RuntimeError("start failed") class _FakeSession: - def __init__(self, *, model_input_schema: ModelInputSchema) -> None: - self._model_input_schema = model_input_schema + def __init__(self, *, inference_input_schema: InferenceInputSchema) -> None: + self._inference_input_schema = inference_input_schema self.step_index = 0 self.closed = False @@ -427,15 +464,15 @@ def next_step_request(self) -> StepRequest | None: return None return StepRequest( step_index=self.step_index, - model_input_schema=self._model_input_schema, + inference_input_schema=self._inference_input_schema, user_input_window=TimeWindow( start_s=0.5 * self.step_index, end_s=0.5 * (self.step_index + 1), ), ) - def step(self, inputs: ModelInputs) -> StepResult: - self._model_input_schema.require_step(inputs) + def step(self, inputs: InferenceInput) -> StepResult: + self._inference_input_schema.require_step(inputs) result = StepResult( step_index=self.step_index, output=f"chunk-{self.step_index}", @@ -449,7 +486,7 @@ def step(self, inputs: ModelInputs) -> StepResult: self.step_index += 1 return result - def reset(self, inputs: ModelInputs | None = None) -> None: + def reset(self, inputs: InferenceInput | None = None) -> None: del inputs self.step_index = 0 @@ -464,14 +501,19 @@ def __init__(self) -> None: def validate( self, *, - user_schema: UserInputSchema | None = None, - model_schema: ModelInputSchema | None = None, + canonical_schema: CanonicalInputSchema | None = None, + inference_input_schema: InferenceInputSchema | None = None, ) -> None: - super().validate(user_schema=user_schema, model_schema=model_schema) + super().validate( + canonical_schema=canonical_schema, + inference_input_schema=inference_input_schema, + ) self.validated = True class _OrderCheckingAdapter(_FakeAdapter): + canonical_input_schema = CanonicalInputSchema() + def __init__(self, *, mapping: _OrderCheckingMapping) -> None: self._mapping = mapping self.created_runtime_after_validate = False @@ -479,14 +521,18 @@ def __init__(self, *, mapping: _OrderCheckingMapping) -> None: def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: self.validate_config(config) self.created_runtime_after_validate = self._mapping.validated - return _FakeRuntime(model_input_schema=self.model_input_schema) + return _FakeRuntime(inference_input_schema=self.inference_input_schema) class _FailingStartAdapter(_FakeAdapter): + canonical_input_schema = CanonicalInputSchema() + def __init__(self) -> None: self.runtime: _FailingRuntime | None = None def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: self.validate_config(config) - self.runtime = _FailingRuntime(model_input_schema=self.model_input_schema) + self.runtime = _FailingRuntime( + inference_input_schema=self.inference_input_schema + ) return self.runtime diff --git a/flashdreams/tests/test_runtime_canonical.py b/flashdreams/tests/test_runtime_canonical.py new file mode 100644 index 000000000..1ad48d39e --- /dev/null +++ b/flashdreams/tests/test_runtime_canonical.py @@ -0,0 +1,590 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the raw-input to canonical-modality layer. + +These cover the middle leg of ``raw input -> canonicalized input -> encoded +inference input``: applications consume canonical modalities, never raw device +events, so adding a device is a registration rather than an application change. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import pytest + +from flashdreams.runtime import ( + DRIVER_COMMAND, + CanonicalInputs, + CanonicalModality, + DeviceConverterSchema, + InferenceInput, + InferenceInputSchema, + InputCanonicalizer, + InputField, + InputMappingSchema, + KeyboardToDriverCommand, + ScriptedModality, + TimeWindow, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, + check_mapping_compatibility, +) + +pytestmark = pytest.mark.ci_cpu + +KEYBOARD_SOURCE = UserInputSchema( + capabilities=( + UserInputCapability(event_type="key_down", payload_fields=frozenset({"key"})), + UserInputCapability(event_type="key_up", payload_fields=frozenset({"key"})), + ) +) +WHEEL_SOURCE = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="wheel_axis", payload_fields=frozenset({"axis", "value"}) + ), + ) +) +PROMPT_SOURCE = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="prompt_set", payload_fields=frozenset({"prompt"}) + ), + ) +) + +# Written once against the canonical modality. It names no key and no axis. +STEERING_MAPPING = InputMappingSchema( + name="driver-command-to-steering", + consumes=(DRIVER_COMMAND,), + produces_step=(InputField(name="steering"),), +) +STEERING_MODEL = InferenceInputSchema(step_fields=(InputField(name="steering"),)) + +WINDOW = TimeWindow(start_s=0.0, end_s=1.0) +NEXT_WINDOW = TimeWindow(start_s=1.0, end_s=2.0) + + +class WheelToDriverCommand: + """Minimal wheel converter standing in for a real evdev profile.""" + + def __init__(self, *, priority: int = 10) -> None: + self._steer = 0.0 + self._seen = False + self._schema = DeviceConverterSchema( + name="wheel-to-driver-command", + produces=DRIVER_COMMAND, + device_kind="wheel", + priority=priority, + consumes=( + UserInputCapability( + event_type="wheel_axis", + payload_fields=frozenset({"axis", "value"}), + ), + ), + ) + + @property + def schema(self) -> DeviceConverterSchema: + return self._schema + + def reset(self) -> None: + self._steer = 0.0 + self._seen = False + + def convert( + self, user_inputs: UserInputs, window: TimeWindow + ) -> Mapping[str, Any] | None: + del window + for event in user_inputs.events: + if event.event_type == "wheel_axis" and event.payload["axis"] == "steer": + self._seen = True + self._steer = float(event.payload["value"]) + if not self._seen: + return None + return DRIVER_COMMAND.value( + { + "throttle": 0.0, + "brake": 0.0, + "steer": self._steer, + "stop": False, + "reverse": False, + } + ) + + +def _key(event_type: str, key: str, timestamp_s: float) -> UserInputEvent: + return UserInputEvent( + timestamp_s=timestamp_s, event_type=event_type, payload={"key": key} + ) + + +def _command(canonical: CanonicalInputs) -> Mapping[str, Any]: + assert DRIVER_COMMAND.name in canonical.values + return canonical.values[DRIVER_COMMAND.name] + + +# --- per-step conditioning ---------------------------------------------- + + +def test_keyboard_edges_become_canonical_driver_command() -> None: + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + inputs = UserInputs(events=(_key("key_down", "w", 0.1),)) + + canonical = canonicalizer.canonicalize( + inputs, window=WINDOW, source_schema=KEYBOARD_SOURCE + ) + + assert _command(canonical)["throttle"] == 1.0 + assert _command(canonical)["steer"] == 0.0 + assert canonical.metadata["canonical_sources"]["driver_command"] == "keyboard" + + +def test_key_aliases_are_normalized() -> None: + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + inputs = UserInputs(events=(_key("key_down", "ArrowLeft", 0.1),)) + + canonical = canonicalizer.canonicalize( + inputs, window=WINDOW, source_schema=KEYBOARD_SOURCE + ) + + assert _command(canonical)["steer"] == 1.0 + + +def test_held_key_still_emits_in_a_window_with_no_events() -> None: + """Edge-triggered HID must become level-triggered per-step conditioning.""" + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + inputs = UserInputs(events=(_key("key_down", "w", 0.1),)) + canonicalizer.canonicalize(inputs, window=WINDOW, source_schema=KEYBOARD_SOURCE) + + quiet = canonicalizer.canonicalize( + inputs, window=NEXT_WINDOW, source_schema=KEYBOARD_SOURCE + ) + + assert _command(quiet)["throttle"] == 1.0 + + +def test_key_release_returns_to_neutral() -> None: + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + inputs = UserInputs(events=(_key("key_down", "a", 0.1), _key("key_up", "a", 1.5))) + canonicalizer.canonicalize(inputs, window=WINDOW, source_schema=KEYBOARD_SOURCE) + + released = canonicalizer.canonicalize( + inputs, window=NEXT_WINDOW, source_schema=KEYBOARD_SOURCE + ) + + assert _command(released)["steer"] == 0.0 + + +def test_reset_drops_device_state() -> None: + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + inputs = UserInputs(events=(_key("key_down", "w", 0.1),)) + canonicalizer.canonicalize(inputs, window=WINDOW, source_schema=KEYBOARD_SOURCE) + + canonicalizer.reset() + after = canonicalizer.canonicalize( + UserInputs(), window=NEXT_WINDOW, source_schema=KEYBOARD_SOURCE + ) + + assert _command(after)["throttle"] == 0.0 + + +# --- boundary: global conditioning is not canonicalized ----------------- + + +def test_canonical_inputs_carry_live_control_only() -> None: + """Global conditioning is application-owned and bypasses this layer.""" + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + inputs = UserInputs( + events=( + _key("key_down", "w", 0.1), + UserInputEvent( + timestamp_s=0.2, event_type="prompt_set", payload={"prompt": "rain"} + ), + ) + ) + + canonical = canonicalizer.canonicalize( + inputs, window=WINDOW, source_schema=KEYBOARD_SOURCE + ) + + assert set(canonical.values) == {"driver_command"} + + +def test_application_supplies_global_conditioning_directly() -> None: + """A prompt swap reaches the session without touching canonicalization.""" + update = InferenceInput(step={"steering": 0.0}).with_global_update( + {"prompt": "heavy rain"} + ) + + assert update.requests_global_update + assert update.global_conditioning["prompt"] == "heavy rain" + + +# --- device independence ------------------------------------------------ + + +def test_mapping_written_against_a_modality_accepts_a_keyboard() -> None: + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + + compatibility = check_mapping_compatibility( + canonical_schema=canonicalizer.canonical_schema(KEYBOARD_SOURCE), + inference_input_schema=STEERING_MODEL, + mapping_schema=STEERING_MAPPING, + ) + + assert compatibility.can_drive + + +def test_adding_a_device_needs_no_application_or_model_change() -> None: + """A wheel is one register() call; mapping and model schemas are untouched.""" + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + canonicalizer.register(WheelToDriverCommand()) + + compatibility = check_mapping_compatibility( + canonical_schema=canonicalizer.canonical_schema(WHEEL_SOURCE), + inference_input_schema=STEERING_MODEL, + mapping_schema=STEERING_MAPPING, + ) + assert compatibility.can_drive + + canonical = canonicalizer.canonicalize( + UserInputs( + events=( + UserInputEvent( + timestamp_s=0.5, + event_type="wheel_axis", + payload={"axis": "steer", "value": -0.4}, + ), + ) + ), + window=WINDOW, + source_schema=WHEEL_SOURCE, + ) + assert _command(canonical)["steer"] == pytest.approx(-0.4) + + +def test_source_with_no_feedable_converter_supplies_no_modalities() -> None: + canonicalizer = InputCanonicalizer([WheelToDriverCommand()]) + + schema = canonicalizer.canonical_schema(KEYBOARD_SOURCE) + + assert schema.modalities == () + assert not schema.supports(DRIVER_COMMAND) + assert canonicalizer.unavailable_converters(KEYBOARD_SOURCE) + + +def test_highest_priority_device_wins_when_both_are_present() -> None: + canonicalizer = InputCanonicalizer( + [KeyboardToDriverCommand(), WheelToDriverCommand()] + ) + both = UserInputSchema( + capabilities=KEYBOARD_SOURCE.capabilities + WHEEL_SOURCE.capabilities + ) + + canonical = canonicalizer.canonicalize( + UserInputs( + events=( + _key("key_down", "a", 0.2), + UserInputEvent( + timestamp_s=0.5, + event_type="wheel_axis", + payload={"axis": "steer", "value": -0.4}, + ), + ) + ), + window=WINDOW, + source_schema=both, + ) + + assert canonical.metadata["canonical_sources"]["driver_command"] == "wheel" + assert _command(canonical)["steer"] == pytest.approx(-0.4) + + +def test_preempted_device_keeps_its_state_current() -> None: + """Keyboard state must not be stale when the wheel disappears.""" + canonicalizer = InputCanonicalizer( + [KeyboardToDriverCommand(), WheelToDriverCommand()] + ) + both = UserInputSchema( + capabilities=KEYBOARD_SOURCE.capabilities + WHEEL_SOURCE.capabilities + ) + inputs = UserInputs( + events=( + _key("key_down", "w", 0.2), + UserInputEvent( + timestamp_s=0.5, + event_type="wheel_axis", + payload={"axis": "steer", "value": -0.4}, + ), + ) + ) + preempted = canonicalizer.canonicalize(inputs, window=WINDOW, source_schema=both) + assert preempted.metadata["canonical_sources"]["driver_command"] == "wheel" + + keyboard_only = canonicalizer.canonicalize( + inputs, window=NEXT_WINDOW, source_schema=KEYBOARD_SOURCE + ) + + assert keyboard_only.metadata["canonical_sources"]["driver_command"] == "keyboard" + assert _command(keyboard_only)["throttle"] == 1.0 + + +# --- registry ----------------------------------------------------------- + + +def test_duplicate_converter_names_are_rejected() -> None: + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + + with pytest.raises(ValueError, match="already registered"): + canonicalizer.register(KeyboardToDriverCommand()) + + +def test_converter_must_fill_the_declared_modality_payload() -> None: + modality = CanonicalModality( + name="steering_wheel", payload_fields=frozenset({"steer", "throttle"}) + ) + + with pytest.raises(ValueError, match="requires payload fields"): + modality.value({"steer": 0.0}) + + +def test_new_modality_is_a_registration_not_a_core_change() -> None: + pedals = CanonicalModality( + name="pedal_state", payload_fields=frozenset({"throttle"}) + ) + + class PedalsConverter: + schema = DeviceConverterSchema( + name="pedals", + produces=pedals, + device_kind="pedals", + consumes=( + UserInputCapability( + event_type="pedal_axis", + payload_fields=frozenset({"value"}), + ), + ), + ) + + def reset(self) -> None: + return None + + def convert( + self, user_inputs: UserInputs, window: TimeWindow + ) -> Mapping[str, Any] | None: + del window + if not user_inputs.events: + return None + return pedals.value( + {"throttle": float(user_inputs.events[-1].payload["value"])} + ) + + source = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="pedal_axis", payload_fields=frozenset({"value"}) + ), + ) + ) + canonicalizer = InputCanonicalizer([PedalsConverter()]) + + assert canonicalizer.canonical_schema(source).modalities == (pedals,) + canonical = canonicalizer.canonicalize( + UserInputs( + events=( + UserInputEvent( + timestamp_s=0.5, event_type="pedal_axis", payload={"value": 0.75} + ), + ) + ), + window=WINDOW, + source_schema=source, + ) + assert canonical.values["pedal_state"]["throttle"] == pytest.approx(0.75) + + +def test_replaying_the_same_windows_reproduces_the_same_canonical_inputs() -> None: + inputs = UserInputs(events=(_key("key_down", "w", 0.1), _key("key_down", "a", 1.2))) + + def run() -> list[dict[str, Any]]: + canonicalizer = InputCanonicalizer([KeyboardToDriverCommand()]) + return [ + dict( + _command( + canonicalizer.canonicalize( + inputs, window=window, source_schema=KEYBOARD_SOURCE + ) + ) + ) + for window in (WINDOW, NEXT_WINDOW) + ] + + assert run() == run() + + +# --- key bindings ------------------------------------------------------- + + +def test_bindings_are_data_and_can_be_rebound() -> None: + """A layout change must not require editing the converter.""" + azerty = InputCanonicalizer( + [ + KeyboardToDriverCommand( + bindings={ + "throttle": frozenset({"z"}), + "brake": frozenset({"s"}), + "steer_left": frozenset({"q"}), + "steer_right": frozenset({"d"}), + "stop": frozenset({"space"}), + } + ) + ] + ) + + canonical = azerty.canonicalize( + UserInputs(events=(_key("key_down", "z", 0.1),)), + window=WINDOW, + source_schema=KEYBOARD_SOURCE, + ) + + assert _command(canonical)["throttle"] == 1.0 + + +def test_tracked_keys_are_derived_so_an_action_cannot_go_unreachable() -> None: + """Declaring bindings and tracked keys separately used to disagree.""" + converter = KeyboardToDriverCommand( + bindings={"stop": frozenset({"escape"}), "throttle": frozenset({"w"})} + ) + canonicalizer = InputCanonicalizer([converter]) + + canonical = canonicalizer.canonicalize( + UserInputs(events=(_key("key_down", "escape", 0.1),)), + window=WINDOW, + source_schema=KEYBOARD_SOURCE, + ) + + assert _command(canonical)["stop"] is True + + +def test_unknown_driver_action_is_rejected() -> None: + with pytest.raises(ValueError, match="Unknown driver actions"): + KeyboardToDriverCommand(bindings={"turbo": frozenset({"t"})}) + + +def test_reverse_is_bindable() -> None: + canonicalizer = InputCanonicalizer( + [KeyboardToDriverCommand(bindings={"reverse": frozenset({"r"})})] + ) + + canonical = canonicalizer.canonicalize( + UserInputs(events=(_key("key_down", "r", 0.1),)), + window=WINDOW, + source_schema=KEYBOARD_SOURCE, + ) + + assert _command(canonical)["reverse"] is True + + +# --- scripted / mock input ---------------------------------------------- + + +def _scripted() -> InputCanonicalizer: + return InputCanonicalizer( + [ + ScriptedModality( + modality=DRIVER_COMMAND, + timeline=[ + ( + 0.0, + { + "throttle": 1.0, + "brake": 0.0, + "steer": 0.0, + "stop": False, + "reverse": False, + }, + ), + ( + 2.0, + { + "throttle": 0.0, + "brake": 0.0, + "steer": 1.0, + "stop": False, + "reverse": False, + }, + ), + ], + ) + ] + ) + + +def test_mock_input_needs_no_raw_events_or_source_schema() -> None: + """Authoring a benchmark scenario must not require raw device vocabulary.""" + canonical = _scripted().canonicalize( + UserInputs(), window=WINDOW, source_schema=UserInputSchema() + ) + + assert _command(canonical)["throttle"] == 1.0 + + +def test_scripted_values_hold_until_the_next_entry() -> None: + canonicalizer = _scripted() + windows = [TimeWindow(start_s=t, end_s=t + 1.0) for t in (0.0, 1.0, 2.0)] + + steer = [ + _command( + canonicalizer.canonicalize( + UserInputs(), window=w, source_schema=UserInputSchema() + ) + )["steer"] + for w in windows + ] + + assert steer == [0.0, 0.0, 1.0] + + +def test_scripted_converter_is_silent_before_its_first_entry() -> None: + canonicalizer = InputCanonicalizer( + [ + ScriptedModality( + modality=CanonicalModality(name="late", payload_fields=frozenset()), + timeline=[(5.0, {})], + ) + ] + ) + + canonical = canonicalizer.canonicalize( + UserInputs(), window=WINDOW, source_schema=UserInputSchema() + ) + + assert canonical.values == {} + + +def test_scripted_timeline_is_validated_against_the_modality() -> None: + with pytest.raises(ValueError, match="requires payload fields"): + ScriptedModality(modality=DRIVER_COMMAND, timeline=[(0.0, {"throttle": 1.0})]) + + +def test_scripted_replay_is_deterministic() -> None: + def run() -> list[float]: + canonicalizer = _scripted() + return [ + _command( + canonicalizer.canonicalize( + UserInputs(), + window=TimeWindow(start_s=t, end_s=t + 1.0), + source_schema=UserInputSchema(), + ) + )["steer"] + for t in (0.0, 1.0, 2.0) + ] + + assert run() == run() diff --git a/flashdreams/tests/test_runtime_input_mapping.py b/flashdreams/tests/test_runtime_input_mapping.py new file mode 100644 index 000000000..00cd9758f --- /dev/null +++ b/flashdreams/tests/test_runtime_input_mapping.py @@ -0,0 +1,573 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for declarative input-mapping compatibility in the runtime API. + +These cover the T2/T3 contract: sources declare what user events they can +provide at payload granularity, models declare required and optional +initial/per-step inputs, and a mapping declares what it consumes and produces so +compatibility can be answered before expensive runtime initialization. +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from flashdreams.runtime import ( + DRIVER_COMMAND, + SESSION_START_ONLY, + CanonicalInputs, + CanonicalInputSchema, + CanonicalModality, + IdentityInputMapping, + InferenceInput, + InferenceInputSchema, + InputField, + InputMappingSchema, + StepRequest, + TimeWindow, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, + check_mapping_compatibility, + check_mapping_set_compatibility, + combine_mapping_schemas, + undeclared_inference_inputs, +) + +pytestmark = pytest.mark.ci_cpu + +KEY_DOWN = UserInputCapability(event_type="key_down", payload_fields=frozenset({"key"})) +KEY_UP = UserInputCapability(event_type="key_up", payload_fields=frozenset({"key"})) +PROMPT_SET = UserInputCapability( + event_type="prompt_set", + semantic_type="text", + payload_fields=frozenset({"prompt"}), +) +FRAME_SET = UserInputCapability( + event_type="initial_frame_set", payload_fields=frozenset({"image"}) +) + +BROWSER_SOURCE = UserInputSchema( + capabilities=(KEY_DOWN, KEY_UP, PROMPT_SET, FRAME_SET), + description="browser webrtc client", +) + +CAMERA_LOOK = CanonicalModality( + name="camera_look", payload_fields=frozenset({"yaw", "pitch"}) +) + +CANONICAL_ALL = CanonicalInputSchema(modalities=(DRIVER_COMMAND, CAMERA_LOOK)) + +# Global conditioning is application-owned and does not come from a canonical +# modality, so this mapping consumes nothing and only declares what it produces. +PROMPT_MAPPING = InputMappingSchema( + name="prompt", + produces_global=(InputField(name="prompt", semantic_type="text"),), +) +FRAME_MAPPING = InputMappingSchema( + name="conditioning-frame", + produces_global=(InputField(name="global_conditioning_frame", required=False),), +) +STEERING_MAPPING = InputMappingSchema( + name="driver-command-to-steering", + consumes=(DRIVER_COMMAND,), + produces_step=(InputField(name="steering"),), +) +LOOK_MAPPING = InputMappingSchema( + name="camera-look", + consumes=(CAMERA_LOOK,), + produces_step=(InputField(name="camera_delta", required=False),), +) + +DRIVING_MODEL = InferenceInputSchema( + global_fields=( + InputField(name="prompt", semantic_type="text", lifecycle="cache_init"), + ), + step_fields=( + InputField(name="steering", lifecycle="step_input"), + InputField(name="camera_delta", required=False, lifecycle="step_input"), + ), +) + + +# --- user input events and windowing ------------------------------------ + + +def test_startup_values_are_represented_as_events() -> None: + inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, event_type="prompt_set", payload={"prompt": "drive"} + ), + UserInputEvent( + timestamp_s=0.5, event_type="key_down", payload={"key": "w"} + ), + ) + ) + + assert inputs.events[0].event_type == "prompt_set" + assert inputs.events[0].payload["prompt"] == "drive" + + +def test_windowing_is_half_open_and_deterministic() -> None: + inputs = UserInputs( + events=tuple( + UserInputEvent(timestamp_s=t, event_type="key_down", payload={"key": "w"}) + for t in (0.0, 0.5, 1.0, 1.5) + ) + ) + + windowed = inputs.window(TimeWindow(start_s=0.5, end_s=1.5)) + + assert [event.timestamp_s for event in windowed.events] == [0.5, 1.0] + + +def test_out_of_order_events_are_rejected() -> None: + with pytest.raises(ValueError, match="non-decreasing"): + UserInputs( + events=( + UserInputEvent(timestamp_s=1.0, event_type="key_down"), + UserInputEvent(timestamp_s=0.5, event_type="key_up"), + ) + ) + + +# --- user input schemas ------------------------------------------------- + + +def test_source_declares_capabilities_at_payload_granularity() -> None: + assert BROWSER_SOURCE.supports(KEY_DOWN) + assert not BROWSER_SOURCE.supports( + UserInputCapability( + event_type="key_down", payload_fields=frozenset({"key", "modifiers"}) + ) + ) + + +def test_bare_event_types_still_satisfy_payload_free_consumers() -> None: + """Coarse pre-capability schemas keep working against the finer query.""" + coarse = UserInputSchema(event_types=frozenset({"reset"})) + + assert coarse.supports(UserInputCapability(event_type="reset")) + assert not coarse.supports( + UserInputCapability(event_type="reset", payload_fields=frozenset({"reason"})) + ) + assert coarse.supports_event_types({"reset"}) + + +def test_capabilities_widen_declared_event_types() -> None: + assert "key_down" in BROWSER_SOURCE.declared_event_types() + assert BROWSER_SOURCE.supports_event_types({"key_down", "prompt_set"}) + + +def test_semantic_type_mismatch_blocks_capability_match() -> None: + source = UserInputSchema( + capabilities=( + UserInputCapability(event_type="prompt_set", semantic_type="embedding"), + ) + ) + + assert not source.supports( + UserInputCapability(event_type="prompt_set", semantic_type="text") + ) + + +def test_event_validation_reports_missing_payload_fields() -> None: + event = UserInputEvent(timestamp_s=0.0, event_type="key_down", payload={}) + + with pytest.raises(ValueError, match="missing required"): + BROWSER_SOURCE.validate_event(event) + + +def test_event_validation_rejects_undeclared_event_type() -> None: + event = UserInputEvent(timestamp_s=0.0, event_type="wheel_axis") + + with pytest.raises(ValueError, match="does not provide event type"): + BROWSER_SOURCE.validate_event(event) + + +# --- model input schemas ------------------------------------------------ + + +def test_model_declares_required_and_optional_fields_per_phase() -> None: + required = DRIVING_MODEL.required_fields() + optional = DRIVING_MODEL.optional_fields() + + assert {(phase, f.name) for phase, f in required} == { + ("global", "prompt"), + ("step", "steering"), + } + assert {(phase, f.name) for phase, f in optional} == {("step", "camera_delta")} + + +def test_required_fields_can_be_filtered_by_phase() -> None: + step_only = DRIVING_MODEL.required_fields("step") + + assert [f.name for _, f in step_only] == ["steering"] + + +def test_field_lookup_is_phase_scoped() -> None: + assert DRIVING_MODEL.field_for(name="prompt", phase="global") is not None + assert DRIVING_MODEL.field_for(name="prompt", phase="step") is None + + +def test_invalid_phase_is_rejected() -> None: + bad_phase: Any = "final" + + with pytest.raises(ValueError, match="phase must be"): + DRIVING_MODEL.fields_for(bad_phase) + + +def test_inference_input_expose_payload_per_phase() -> None: + inputs = InferenceInput( + global_conditioning={"prompt": "drive"}, step={"steering": 0.25} + ) + + assert inputs.for_phase("global")["prompt"] == "drive" + assert inputs.for_phase("step")["steering"] == 0.25 + + +def test_lifecycle_and_update_policy_are_queryable_metadata() -> None: + field = InputField( + name="prompt", + update_policy="step_boundary", + lifecycle="cache_init", + metadata={"coordinates": "opencv_c2w"}, + ) + + assert field.update_policy == "step_boundary" + assert field.lifecycle == "cache_init" + assert field.metadata["coordinates"] == "opencv_c2w" + + +def test_metadata_is_excluded_from_field_equality() -> None: + plain = InputField(name="prompt") + annotated = InputField(name="prompt", metadata={"note": "hint"}) + + assert plain == annotated + + +# --- mapping compatibility ---------------------------------------------- + + +def test_compatible_source_model_and_mapping_can_drive() -> None: + compatibility = check_mapping_set_compatibility( + canonical_schema=CANONICAL_ALL, + inference_input_schema=DRIVING_MODEL, + mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING, LOOK_MAPPING), + ) + + assert compatibility.can_drive + assert {(p, f.name) for p, f in compatibility.satisfied_required_model_fields} == { + ("global", "prompt"), + ("step", "steering"), + } + assert {(p, f.name) for p, f in compatibility.available_optional_model_fields} == { + ("step", "camera_delta") + } + + +def test_missing_required_model_field_blocks_the_run() -> None: + compatibility = check_mapping_set_compatibility( + canonical_schema=CANONICAL_ALL, + inference_input_schema=DRIVING_MODEL, + mapping_schemas=(PROMPT_MAPPING,), + ) + + assert not compatibility.can_drive + assert [f.name for _, f in compatibility.missing_required_model_fields] == [ + "steering" + ] + + +def test_missing_source_capability_is_reported_when_it_blocks() -> None: + no_wheel = CanonicalInputSchema(modalities=(CAMERA_LOOK,)) + + compatibility = check_mapping_set_compatibility( + canonical_schema=no_wheel, + inference_input_schema=DRIVING_MODEL, + mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING), + ) + + assert not compatibility.can_drive + assert compatibility.unavailable_mapping_names == ("driver-command-to-steering",) + assert {m.name for m in compatibility.missing_modalities} == {"driver_command"} + + +def test_unfeedable_optional_mapping_degrades_instead_of_vetoing() -> None: + """Losing a mapping that fed only optional fields must not block the run.""" + no_look = CanonicalInputSchema(modalities=(DRIVER_COMMAND,)) + + compatibility = check_mapping_set_compatibility( + canonical_schema=no_look, + inference_input_schema=DRIVING_MODEL, + mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING, LOOK_MAPPING), + ) + + assert compatibility.can_drive + assert compatibility.unavailable_mapping_names == ("camera-look",) + # The dropped mapping's field must not be advertised as available. + assert compatibility.available_optional_model_fields == () + + +def test_optional_field_needs_mapping_support_to_be_available() -> None: + compatibility = check_mapping_set_compatibility( + canonical_schema=CANONICAL_ALL, + inference_input_schema=DRIVING_MODEL, + mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING), + ) + + assert compatibility.can_drive + assert compatibility.available_optional_model_fields == () + + +def test_lifecycle_disagreement_blocks_a_field_match() -> None: + model = InferenceInputSchema( + global_fields=(InputField(name="prompt", lifecycle="rollout_binding"),) + ) + mapping = InputMappingSchema( + name="prompt", + produces_global=(InputField(name="prompt", lifecycle="cache_init"),), + ) + + compatibility = check_mapping_compatibility( + canonical_schema=CANONICAL_ALL, + inference_input_schema=model, + mapping_schema=mapping, + ) + + assert not compatibility.can_drive + + +def test_unspecified_lifecycle_stays_permissive() -> None: + model = InferenceInputSchema(global_fields=(InputField(name="prompt"),)) + + compatibility = check_mapping_compatibility( + canonical_schema=CANONICAL_ALL, + inference_input_schema=model, + mapping_schema=PROMPT_MAPPING, + ) + + assert compatibility.can_drive + + +def test_raise_if_incompatible_names_both_failure_kinds() -> None: + compatibility = check_mapping_set_compatibility( + canonical_schema=CanonicalInputSchema(modalities=(CAMERA_LOOK,)), + inference_input_schema=DRIVING_MODEL, + mapping_schemas=(PROMPT_MAPPING, STEERING_MAPPING), + ) + + with pytest.raises(ValueError) as excinfo: + compatibility.raise_if_incompatible() + + message = str(excinfo.value) + assert "missing canonical modalities" in message + assert "missing required model inputs" in message + + +def test_raise_if_incompatible_is_a_no_op_when_compatible() -> None: + compatibility = check_mapping_set_compatibility( + canonical_schema=CANONICAL_ALL, + inference_input_schema=DRIVING_MODEL, + mapping_schemas=(PROMPT_MAPPING, FRAME_MAPPING, STEERING_MAPPING), + ) + + compatibility.raise_if_incompatible() + + +def test_check_mapping_compatibility_rejects_a_non_schema() -> None: + not_a_schema: Any = object() + + with pytest.raises(TypeError, match="InputMappingSchema"): + check_mapping_compatibility( + canonical_schema=CANONICAL_ALL, + inference_input_schema=DRIVING_MODEL, + mapping_schema=not_a_schema, + ) + + +# --- mapping schema composition ----------------------------------------- + + +def test_combining_mappings_unions_their_surfaces() -> None: + combined = combine_mapping_schemas((PROMPT_MAPPING, STEERING_MAPPING)) + + assert {m.name for m in combined.consumes} == {"driver_command"} + assert [f.name for f in combined.produces_global] == ["prompt"] + assert [f.name for f in combined.produces_step] == ["steering"] + + +def test_duplicate_declarations_collapse_and_merge_metadata() -> None: + first = InputMappingSchema( + name="a", + produces_global=(InputField(name="prompt", metadata={"source": "a"}),), + ) + second = InputMappingSchema( + name="b", + produces_global=( + InputField(name="prompt", metadata={"source": "b", "extra": "kept"}), + ), + ) + + combined = combine_mapping_schemas((first, second)) + + assert len(combined.produces_global) == 1 + metadata = combined.produces_global[0].metadata + assert metadata["source"] == "a" + assert metadata["extra"] == "kept" + + +def test_combine_rejects_non_schema_entries() -> None: + not_a_schema: Any = object() + + with pytest.raises(TypeError, match="InputMappingSchema"): + combine_mapping_schemas((PROMPT_MAPPING, not_a_schema)) + + +# --- declaration drift -------------------------------------------------- + + +def test_undeclared_inference_input_catches_schema_drift() -> None: + produced = InferenceInput( + global_conditioning={"prompt": "drive"}, step={"steering": 0.0} + ) + + undeclared = undeclared_inference_inputs(produced, PROMPT_MAPPING) + + assert undeclared == (("step", "steering"),) + + +def test_declared_outputs_report_no_drift() -> None: + combined = combine_mapping_schemas((PROMPT_MAPPING, STEERING_MAPPING)) + produced = InferenceInput( + global_conditioning={"prompt": "drive"}, step={"steering": 0.0} + ) + + assert undeclared_inference_inputs(produced, combined) == () + + +# --- interoperability with the T1 envelope ------------------------------ + + +def test_identity_mapping_needs_no_declared_surface() -> None: + """Fixed-input runs stay possible without any schema declaration.""" + mapping = IdentityInputMapping() + fixed = InferenceInput( + global_conditioning={"prompt": "fixed"}, step={"steering": 0.0} + ) + + mapped = mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=fixed, + request=StepRequest(step_index=0), + ) + + assert mapped.step["steering"] == 0.0 + + +def test_empty_mapping_set_cannot_satisfy_a_required_field() -> None: + compatibility = check_mapping_set_compatibility( + canonical_schema=CANONICAL_ALL, + inference_input_schema=DRIVING_MODEL, + mapping_schemas=(), + ) + + assert not compatibility.can_drive + assert len(compatibility.missing_required_model_fields) == 2 + + +def test_model_with_no_requirements_is_always_drivable() -> None: + compatibility = check_mapping_set_compatibility( + canonical_schema=CanonicalInputSchema(), + inference_input_schema=InferenceInputSchema(), + mapping_schemas=(), + ) + + assert compatibility.can_drive + + +# --- global conditioning updates vs reset ------------------------------- + + +def test_empty_global_slot_requests_no_update() -> None: + steady_state = InferenceInput(step={"steering": 0.25}) + + assert not steady_state.requests_global_update + + +def test_non_empty_global_slot_mid_rollout_is_an_update_request() -> None: + """Changing weather mid-run updates conditioning; it is not a reset.""" + updated = InferenceInput(step={"steering": 0.0}).with_global_update( + {"prompt": "heavy rain"} + ) + + assert updated.requests_global_update + assert updated.global_conditioning["prompt"] == "heavy rain" + assert updated.step["steering"] == 0.0 + + +def test_with_step_carries_the_global_slot_through() -> None: + started = InferenceInput(global_conditioning={"prompt": "drive"}) + + stepped = started.with_step({"steering": 0.5}) + + assert stepped.global_conditioning["prompt"] == "drive" + + +def test_without_global_update_clears_the_request() -> None: + started = InferenceInput( + global_conditioning={"prompt": "drive"}, step={"steering": 0.5} + ) + + steady_state = started.without_global_update() + + assert not steady_state.requests_global_update + assert steady_state.step["steering"] == 0.5 + + +def test_model_can_declare_conditioning_it_cannot_swap_mid_rollout() -> None: + schema = InferenceInputSchema( + global_fields=( + InputField(name="prompt", update_policy="step_boundary"), + InputField(name="scene_id", update_policy=SESSION_START_ONLY), + ) + ) + update = InferenceInput( + global_conditioning={"prompt": "heavy rain", "scene_id": "town_02"} + ) + + assert schema.unsupported_global_updates(update) == ("scene_id",) + + +def test_permissive_when_no_update_policy_is_declared() -> None: + schema = InferenceInputSchema(global_fields=(InputField(name="prompt"),)) + update = InferenceInput(global_conditioning={"prompt": "heavy rain"}) + + assert schema.unsupported_global_updates(update) == () + + +def test_undeclared_global_values_are_left_to_the_adapter() -> None: + schema = InferenceInputSchema(global_fields=(InputField(name="prompt"),)) + update = InferenceInput(global_conditioning={"mystery": 1}) + + assert schema.unsupported_global_updates(update) == () + + +def test_steady_state_steps_do_not_request_a_global_update() -> None: + """Carrying session-start conditioning forward would look like an update.""" + started = InferenceInput(global_conditioning={"prompt": "drive"}) + + steady_state = InferenceInput(step={"chunk_index": 1}) + + assert started.requests_global_update + assert not steady_state.requests_global_update + assert ( + not started.with_step({"chunk_index": 1}) + .without_global_update() + .requests_global_update + ) From 2d2e5b3c956d7dc07ca3ec9377f1e2e99c5fcbec Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Wed, 5 Aug 2026 18:35:46 -0700 Subject: [PATCH 04/19] Clarify global conditioning input semantics Signed-off-by: Aidan Foster --- docs/inference_runtime_api_design.md | 65 ++++--- ...inference_runtime_inputs_implementation.md | 97 +++++----- ...ence_runtime_supported_inputs_inventory.md | 177 +++++++++--------- flashdreams/flashdreams/runtime/__init__.py | 2 - flashdreams/flashdreams/runtime/inputs.py | 114 +++-------- flashdreams/flashdreams/runtime/interfaces.py | 4 +- flashdreams/flashdreams/runtime/mapping.py | 33 ++-- flashdreams/flashdreams/runtime/types.py | 9 +- .../tests/test_inference_runtime_api.py | 19 +- flashdreams/tests/test_runtime_canonical.py | 11 +- .../tests/test_runtime_input_mapping.py | 167 +++++------------ 11 files changed, 293 insertions(+), 405 deletions(-) diff --git a/docs/inference_runtime_api_design.md b/docs/inference_runtime_api_design.md index f70fbd890..b6314a205 100644 --- a/docs/inference_runtime_api_design.md +++ b/docs/inference_runtime_api_design.md @@ -119,7 +119,7 @@ InferenceRuntime | v InferenceSession - one rollout/stream: prompt/initial inputs, cache/state, current step, reset + one rollout/stream: global conditioning, cache/state, current step, reset keeps per-run state from leaking across prompts, clients, or benchmark repeats | v @@ -187,11 +187,11 @@ local model implementation, a Dynamo-like backend, or a hosted service. | --- | --- | --- | | Model/preset registry | Lists what can run: model/preset slugs, scenarios, capabilities, resource hints, and supported output modes. | Must remain cheap to query and must not load checkpoints. | | App / integration / benchmark / transport | Owns the user-facing mode: CLI, native integration, WebRTC, benchmark, hosted request, or replay. | Supplies run setup, user inputs, model inputs, and output target selection. | -| User input library | Normalizes live or replayed controls into FlashDreams-supported user input events/windows. | Shared primitives for keyboard, reset, prompt/image updates, traces, and future scalar controls. | -| Input mapping | Converts user/app inputs plus initial model inputs into the model-specific inputs needed by the session. | A model adapter may provide a default mapping; runtimes, applications, benchmarks, and replay tools may override it without changing the model step. | +| User input library | Normalizes live or replayed controls into FlashDreams-supported user input events/windows. | Shared primitives for keyboard, reset, prompt/image selection, traces, and future scalar controls. | +| Input mapping | Converts user/app inputs plus global conditioning into the model-specific inputs needed by the session. | A model adapter may provide a default mapping; runtimes, applications, benchmarks, and replay tools may override it without changing the model step. | | ModelRunner / standard loop | Orchestrates one run from setup through runtime initialization, stepping, output, metrics, and teardown. | Shared orchestration layer used by CLIs, benchmarks, MP4 runs, and simple realtime flows. | | InferenceRuntime | Owns heavyweight lifecycle: distributed init, model construction, checkpoint loading, compile/capture, warmup, hosted-service connection, and teardown. | Long-lived reusable runtime created from `InferenceConfig`; lets FlashDreams load/warm once and create sessions sequentially unless the backend supports concurrency. | -| InferenceSession | Owns one rollout or stream: initial inputs, cache state, current step, reset behavior, step requirements, and step execution. | Per-rollout interface consumed by the standard loop; keeps state isolated across prompts, browser clients, replay scenarios, or benchmark repeats. | +| InferenceSession | Owns one rollout or stream: global conditioning, cache state, current step, reset behavior, step requirements, and step execution. | Per-rollout interface consumed by the standard loop; keeps state isolated across prompts, browser clients, replay scenarios, or benchmark repeats. | | Model implementation / inference pipeline | Implements encode, model step, decode, cache updates, and model-specific optimizations. | FlashDreams wraps this boundary; it should not replace every model implementation. | | Output target | Consumes generated outputs and handles presentation or persistence. | Separate from model execution so the same session can feed WebRTC, MP4, benchmark, or headless output. | | Metrics, artifacts, and profiling | Records timings, memory, quality data, logs, reports, traces, and optional NVTX ranges. | Shared observation layer for local runs, benchmarks, CI smoke, and hosted runs. | @@ -298,8 +298,7 @@ uses: - keyboard keydown/keyup events; - reset requests; -- prompt update requests; -- image update requests; +- prompt or image selection/update events; - future scalar controls such as throttle, brake, steer, or camera axes once an integration needs them. @@ -335,11 +334,11 @@ Examples of global conditioning include prompt, negative prompt, conditioning frame, input video, scene id, HD map asset, camera calibration, initial camera pose, seed, or model-specific fields. -Global conditioning is normally supplied when a session starts, but a non-empty -global slot on a mid-rollout input is an update request rather than a reset; -resetting rollout state is a separate `InferenceSession.reset()` call. Whether a -given value can be swapped mid-rollout is declared per field by -`InputField.update_policy`. +Global conditioning establishes session-global model state when a session +starts or resets. During an active rollout, a non-empty global-conditioning +payload passed to `InferenceSession.step()` asks the session to update that +state when the model supports it. Reset remains a separate explicit session +method. Examples of per-step conditioning include frame timestamps, pose segments, camera trajectory chunks, rendered HD map frames, conditioning video windows, @@ -349,16 +348,17 @@ Inference input payloads should use semantic names, not only modality names. For example, a first frame and an HD map frame should be distinct inputs even if both are image-like values. -Model input metadata may also include a lightweight lifecycle label, such as -runtime config, cache initialization, rollout binding, per-step input, or -session update. This should remain query metadata, not model-specific tensor -validation. +Model input names, payload kinds, semantic-type hints, and schema metadata +should be open-ended. Supported integrations such as SANA-WM, LingBot, +Omnidreams, and future external adapters may need different semantic fields. +Adding a new model should usually mean adding adapter-owned schema declarations +and mappings, not changing a central FlashDreams enum. -Model input names, payload kinds, lifecycle labels, and schema metadata should -be open-ended. Supported integrations such as SANA-WM, LingBot, Omnidreams, and -future external adapters may need different semantic fields. Adding a new model -should usually mean adding adapter-owned schema declarations and mappings, not -changing a central FlashDreams enum. +Consumption cadence is a separate hint from input scope. A field may be +provided through global conditioning because it is session-global state, while +the adapter consumes or slices it during every step. That can be recorded as +`frequency_consumed` metadata without changing whether the field belongs in +`global_conditioning_fields` or `step_fields`. For interactive runs, most `InferenceInput` values will be global conditioning plus per-step inputs produced by input mapping. For MP4 generation and benchmarking, the API @@ -374,7 +374,7 @@ These schemas are not meant to be a rich type system or a replacement for model-specific validation. They should be just enough to answer: - what can this app, transport, trace, or benchmark source provide? -- what does this model require before startup and at each step? +- what does this model require before session start and at each step? - can this event source drive this model with the selected mapping? The purpose is to fail early before expensive model initialization, produce @@ -386,7 +386,7 @@ coordinate frame, units, rough shape summary, accepted file suffixes, schema URI, model family, or source/transport details. Metadata should help humans and adapter selection code, but compatibility should still be based on the declared event capabilities, semantic model fields, payload representation hints, and -lifecycle labels. +schema phases. Consumption-cadence hints are descriptive and adapter-owned. For simple CLI text-to-video or image-to-video runs, `UserInputSchema` can be trivial or omitted because there may be no live controls. `InferenceInputSchema` is @@ -478,6 +478,14 @@ declare user inputs, declare model inputs, and provide a default mapping, while the runtime owns transport, event validation, timestamping, input queue/window selection, output delivery, and optional overrides. +`StepRequest` and `StepResult` are per-step runtime messages, not declarative +schemas. `InferenceSession.next_step_request()` returns a `StepRequest` to say +which step is next, which user-input time window to map, and whether this step +has any narrower `InferenceInputSchema` than the session default. The runner or +application then builds an `InferenceInput` and calls `InferenceSession.step()`, +which returns a `StepResult` carrying the generated output, output timing, +metrics, and step metadata. + Examples: - T2V mapping validates a prompt and creates no per-step control inputs. @@ -505,7 +513,7 @@ A run should: profiling, and optional scenario setup. 3. Validate that the event source and mapping can drive the selected model. 4. Initialize the runtime. -5. Start a session from initial model inputs. +5. Start a session from global conditioning inputs. 6. For each step, ask the session what it needs, gather live or fixed inputs, build step model inputs, run the session step, route outputs, and record metrics. @@ -553,8 +561,9 @@ generation, benchmarks, regression testing, and autotune. Two replay levels should be supported: -- user-event replay: records timestamped key events, prompt updates, image - updates, reset events, and timing, then runs normal input mapping; +- user-event replay: records timestamped key events, prompt or image + selection/update events, reset events, and timing, then runs normal input + mapping; - model-input replay: records or defines already-mapped per-step model inputs for stricter model-level regression tests. @@ -667,8 +676,10 @@ registry, standard loop, concrete output modes, or model migrations: - The model-specific integration boundary is named `ModelAdapter`. - Heavyweight lifecycle is split into `InferenceRuntime` and `InferenceSession`. -- Step data carriers are named `StepRequest` and `StepResult`; a session returns - `None` from `next_step_request()` when the rollout is complete. +- Step data carriers are named `StepRequest` and `StepResult`. They are runtime + messages around one call to `InferenceSession.step()`, not schema + declarations; a session returns `None` from `next_step_request()` when the + rollout is complete. - Raw inputs use `UserInputs`, canonicalized inputs use `CanonicalInputs`, and model-facing inputs use `InferenceInput`. Both remain lightweight payload envelopes with shallow read-only mappings. diff --git a/docs/inference_runtime_inputs_implementation.md b/docs/inference_runtime_inputs_implementation.md index 8d485768a..763b9810c 100644 --- a/docs/inference_runtime_inputs_implementation.md +++ b/docs/inference_runtime_inputs_implementation.md @@ -46,70 +46,60 @@ that touches no application, mapping, or model code. This path covers **live user control only**. Global conditioning is application-owned data and reaches `InferenceInput` directly, without passing -through canonicalization or a device converter. An application that wants a -trigger key to swap the prompt reads that as ordinary canonical control input -and updates its own global conditioning in response. +through canonicalization or a device converter. Session start/reset establishes +that global conditioning. During an active rollout, a non-empty +`global_conditioning` payload passed to `step()` requests an update of the +session-global state when the model supports it. ## Conditioning Slots -Both the canonical and encoded layers split into two slots, and the split means -the same thing at each: +The encoded layer splits model-facing inputs into two slots: -- **global conditioning** — conditions the whole rollout: prompt, conditioning - frame, scene. Normally supplied at session start. +- **global conditioning** — session-global model state: prompt, conditioning + frame, scene. - **per-step conditioning** — needed to generate the next chunk or frame: steering, HD map frames, camera trajectory. -`InputPhase` is `Literal["global", "step"]`. The axis names *which slot*, not -*when the value may arrive* — see the next section. +`InputPhase` is `Literal["global_conditioning", "step"]`. The phase names the +`InferenceInput` slot the caller provides. -## Global Conditioning Updates Are Not Resets +`InputField.frequency_consumed` is independent query metadata. It says how the +adapter consumes a field internally, such as `once` or `per_step`; it does not +decide whether the caller provides the field through `global_conditioning` or +`step`. -A non-empty global slot on a mid-rollout `InferenceInput` is an **update -request**. The session should apply it when the model supports doing so. -Resetting rollout state is a separate, explicit `InferenceSession.reset()` call. -The motivating case is changing prompt and conditioning frame mid-run to change -the weather in an Omnidreams rollout. +## Global Conditioning Is Session-Global State -```python -from flashdreams.runtime import InferenceInput - -steady_state = InferenceInput(step={"steering": 0.25}) -assert not steady_state.requests_global_update - -changed_weather = steady_state.with_global_update({"prompt": "heavy rain"}) -assert changed_weather.requests_global_update -``` - -Because `with_step()` carries the global slot through unchanged, use -`without_global_update()` for the steady-state case; otherwise every step looks -like an update request. - -Whether a value can actually be swapped mid-rollout is declared per field: +`InferenceInput.global_conditioning` carries session-scoped inputs. A runtime +passes those values to `InferenceRuntime.start_session()` or to +`InferenceSession.reset()` when the backend supports resetting a rollout. +During an active rollout, passing a non-empty `global_conditioning` payload to +`InferenceSession.step()` asks the session to update that session-global state. +The model/session owns whether that update is supported. ```python -from flashdreams.runtime import SESSION_START_ONLY, InferenceInputSchema, InputField +from flashdreams.runtime import InferenceInput, InferenceInputSchema, InputField schema = InferenceInputSchema( - global_fields=( - InputField(name="prompt", update_policy="step_boundary"), - InputField(name="scene_id", update_policy=SESSION_START_ONLY), + global_conditioning_fields=( + InputField(name="prompt"), + InputField(name="scene_id"), ) ) -schema.unsupported_global_updates( - InferenceInput(global_conditioning={"prompt": "heavy rain", "scene_id": "town_02"}) +schema.require_global_conditioning( + InferenceInput(global_conditioning={"prompt": "drive", "scene_id": "town_02"}) ) -# ("scene_id",) -``` -`SESSION_START_ONLY` is the one reserved `update_policy` token. Everything else -in that vocabulary, and all of `lifecycle`, is open and adapter-owned; this layer -only carries it as queryable metadata. +step_with_prompt_update = InferenceInput( + global_conditioning={"prompt": "heavy rain"}, + step={"steering": 0.0}, +) +``` -Steady-state steps must leave the global slot empty; otherwise every step reads -as an update request. Converters emit every window, because live control is -level-triggered: a key held across a step emits no events but still means full -throttle. +Per-step conditioning is different: those values are supplied through +`InferenceInput.step` for each generated chunk or frame. Converters still emit +every window, because live control is level-triggered: a key held across a step +emits no events but still means full throttle. ## Raw Inputs @@ -192,8 +182,9 @@ device does not resume from stale state. ## Mapping And Compatibility `InputMapping` is the canonical-to-encoded boundary. `InputMappingSchema` is its -declarative surface: `consumes` names canonical modalities; `produces_global` -and `produces_step` name the `InferenceInput` fields it can build. +declarative surface: `consumes` names canonical modalities; +`produces_global_conditioning` and `produces_step` name the `InferenceInput` +fields it can build. `InputMapping.validate()` raises, which fails a run late and cannot say *which* optional model input a source would enable or *which* missing modality makes a @@ -230,6 +221,13 @@ registered later, with no change to the mapping or the model schema. `undeclared_inference_inputs()` reports payload keys a mapping produced but did not declare, which keeps hand-written schemas honest as the code drifts. +`StepRequest` and `StepResult` sit around a single `InferenceSession.step()` +call. They are not schema declarations. A session returns `StepRequest` from +`next_step_request()` to name the next step, optionally provide a narrower +`InferenceInputSchema`, and request a `TimeWindow` of user inputs. The runner +then builds `InferenceInput` and calls `step()`, which returns a `StepResult` +for the output target and metrics recorder. + ## What This Does Not Validate The schemas intentionally avoid becoming a rich type system. These remain the @@ -237,8 +235,9 @@ responsibility of the model adapter, runtime, session, or mapping: - tensor shape and dtype, image decode details; - camera coordinate systems, pose and timestamp units; -- prompt-embedding swap mechanics; -- whether a model can actually apply a declared update policy at runtime; +- prompt-embedding mechanics; +- whether a model can actually apply a requested global-conditioning update; +- enforcing consumption-cadence metadata; - deep validation of scene, HD map, or actor-state data. The layer answers "can this source plausibly drive this model through this diff --git a/docs/inference_runtime_supported_inputs_inventory.md b/docs/inference_runtime_supported_inputs_inventory.md index ebe9d853a..d56cf1651 100644 --- a/docs/inference_runtime_supported_inputs_inventory.md +++ b/docs/inference_runtime_supported_inputs_inventory.md @@ -17,44 +17,44 @@ FastVideo Causal WAN 2.2 T2V, and Cosmos Predict2 T2V: - Source/app inputs: prompt text or prompt text file, pixel height/width, and fps or block count depending on runner. -- Model-facing initial inputs: prompt text plus latent/output height and width +- Model-facing global conditioning: prompt text plus latent/output height and width derived from run config. -- Model-facing step/update inputs: no live controls; AR loop steps with fixed +- Model-facing per-step inputs: no live controls; AR loop steps with fixed session state. WAN 2.1 I2V, Causal-Forcing I2V, and Cosmos Predict2 I2V: - Source/app inputs: prompt text or prompt file, first-frame image path or URL, and pixel height/width. -- Model-facing initial inputs: prompt text and decoded first-frame tensor. -- Model-facing step/update inputs: no live controls. +- Model-facing global conditioning: prompt text and decoded first-frame tensor. +- Model-facing per-step inputs: no live controls. FlashVSR: - Source/app inputs: input video path or URL, chunk size, crop region, sparse ratio, and optional output FPS. -- Model-facing initial inputs: no explicit prompt at runner time; the prompt +- Model-facing global conditioning: no explicit prompt at runner time; the prompt tensor is configured in the pipeline. Input video dimensions affect per-video runtime/pipeline setup. -- Model-facing step/update inputs: video chunks passed to +- Model-facing per-step inputs: video chunks passed to `pipeline.generate(input=clip)`. LingBot CLI: - Source/app inputs: prompt or prompt path, first-frame image path, pose path, intrinsics path, total blocks, dimensions, and fps. -- Model-facing initial inputs: prompt text and first-frame tensor. -- Model-facing step/update inputs: `CamCtrlInput` with intrinsics, camera poses, +- Model-facing global conditioning: prompt text and first-frame tensor. +- Model-facing per-step inputs: `CamCtrlInput` with intrinsics, camera poses, and world scale. LingBot WebRTC: - Source/app inputs: session prompt, uploaded/remote/default first-frame image, keyboard events, reset requests, text-event catalog, and trigger events. -- Model-facing initial inputs: prompt text, first-frame tensor, base text +- Model-facing global conditioning: prompt text, first-frame tensor, base text embeddings, precomputed text-event embeddings, base intrinsics, and world scale. -- Model-facing step/update inputs: keyboard event windows become pose segments +- Model-facing per-step inputs: keyboard event windows become pose segments and camera trajectories. Text-event triggers can replace rollout text embeddings when the model supports it. @@ -63,9 +63,9 @@ HY-WorldPlay WAN I2V: - Source/app inputs: prompt or prompt path, first-frame image path or example image, pose string or pose JSON, memory-selection settings, dimensions, fps, and seed. -- Model-facing initial inputs: prompt text and first-frame tensor for cache - initialization. -- Model-facing step/update inputs: pose data is bound for the rollout as action +- Model-facing global conditioning: prompt text and first-frame tensor for + session setup. +- Model-facing per-step inputs: pose data is bound for the rollout as action labels, view matrices, intrinsics, and memory-selection state before AR steps. Omnidreams CLI: @@ -73,18 +73,18 @@ Omnidreams CLI: - Source/app inputs: shared prompt or per-camera prompts, HDMap video paths, first-frame image/video paths, camera names, example-data UUID, and optional embedding save/load paths. -- Model-facing initial inputs: prompt list, first-frame tensor, view names; or +- Model-facing global conditioning: prompt list, first-frame tensor, view names; or precomputed text/image/negative-text embeddings. -- Model-facing step/update inputs: HDMap video chunks passed per AR step. +- Model-facing per-step inputs: HDMap video chunks passed per AR step. Omnidreams WebRTC: - Source/app inputs: scene directory or scene UUID, scene variant, camera name, prompt/first-frame assets resolved from the scene, keyboard events, reset requests, and optional postprocess preset. -- Model-facing initial inputs: scene data, renderer, first-frame tensor, prompt, +- Model-facing global conditioning: scene data, renderer, first-frame tensor, prompt, camera calibration/extrinsics, initial ego pose, and initial timestamp. -- Model-facing step/update inputs: keyboard event windows become ego poses, +- Model-facing per-step inputs: keyboard event windows become ego poses, camera poses per view, and frame timestamps. The wrapper renders HDMap conditioning internally for each step. @@ -92,26 +92,26 @@ Omnidreams interactive drive: - Source/app inputs: scene bundle, keyboard events or wheel/controller samples, view-mode/reset/scene-exit controls, and vehicle/chunk config. -- Model-facing initial inputs: scene bundle, selected camera, prompt, initial +- Model-facing global conditioning: scene bundle, selected camera, prompt, initial RGB frame, initial rig pose, and initial timestamp. -- Model-facing step/update inputs: `DriverCommand` samples become trajectory +- Model-facing per-step inputs: `DriverCommand` samples become trajectory chunks, rendered frames, and world-model conditioning. Template recipe: - Source/app inputs: synthetic runner config: batch size, height, width, context tokens, AR steps, and seed. -- Model-facing initial inputs: synthetic transformer context, optional negative +- Model-facing global conditioning: synthetic transformer context, optional negative context, height, and width. -- Model-facing step/update inputs: optional synthetic control tensor. +- Model-facing per-step inputs: optional synthetic control tensor. WAN 2.2 TI2V pipeline config: - Source/app inputs: downstream runners use this rather than a standalone runner in this tree. -- Model-facing initial inputs: prompt text and first-frame image for TI2V-style - cache initialization. -- Model-facing step/update inputs: downstream runners decide controls; +- Model-facing global conditioning: prompt text and first-frame image for + TI2V-style session setup. +- Model-facing per-step inputs: downstream runners decide controls; HY-WorldPlay currently binds action/camera state around it. SANA-WM bidirectional and streaming on `main`: @@ -120,10 +120,10 @@ SANA-WM bidirectional and streaming on `main`: negative prompt, camera trajectory path or action DSL, optional intrinsics path or derived intrinsics, frame count, fps, Stage-1 sampling knobs, seed, precision/refiner options, and streaming chunk/block settings. -- Model-facing initial inputs: decoder context such as prompt, fps, +- Model-facing global conditioning: decoder context such as prompt, fps, `save_stage1`, refiner seed, sink size, and streaming refiner window/block parameters. -- Model-facing step/update inputs: bidirectional passes one +- Model-facing per-step inputs: bidirectional passes one `SanaWMI2VConditioningRequest` into the single generation step. Streaming passes one `SanaWMStreamingI2VConditioningRequest` repeatedly; the conditioning encoder caches rollout-wide prompt, first-frame, camera, latent @@ -134,7 +134,7 @@ SANA-WM bidirectional and streaming on `main`: ## API Implications -The inventory changes the T2/T3 shape in four concrete ways. +The inventory changes the T2/T3 shape in five concrete ways. First, a selected mapping is often a composition. A LingBot-like run needs prompt mapping, first-frame mapping, and keyboard-to-camera mapping. Omnidreams may add @@ -142,25 +142,24 @@ scene selection, camera selection, and HDMap mapping. The implementation should support checking a set of mapping schemas as one compatibility surface, while still allowing a single mapping object when that is simpler. -Second, `InferenceInputSchema` needs a lightweight lifecycle tag in addition to the -`initial` versus `step` phase. The phase answers when the value is needed at the -standard-loop level. The lifecycle tag distinguishes where the model adapter -uses it, such as: - -- `runtime_config`: values that affect setup before model/runtime construction, - such as FlashVSR input-video dimensions; -- `cache_init`: values passed when initializing or resetting a rollout cache, - such as prompts, first frames, view names, and precomputed embeddings; -- `rollout_binding`: values bound after cache initialization but before AR - steps, such as HY-WorldPlay action labels, camera tensors, and memory state; -- `step_input`: values consumed for one generated chunk, such as HDMap frames, - camera trajectories, driver commands, video chunks, and timestamps; -- `session_update`: values that can update an active session when supported, - such as LingBot text-event embedding swaps. - -The lifecycle tag is metadata, not a new deep type system. If both a model field -and mapping output specify lifecycle, compatibility should require them to agree. -If either side omits it, matching stays permissive for simple schemas. +Second, `InferenceInputSchema` needs explicit global-conditioning and per-step +schema slots. `global_conditioning_fields` describe the session-global state +carried through `InferenceInput.global_conditioning`. Start/reset establishes +that state; a non-empty global-conditioning payload in a step context asks the +session to update it when the model supports that. `step_fields` arrive through +`InferenceInput.step` for one generated chunk or frame window. + +This distinction matters for rollout-wide values such as full camera +trajectories, action labels, intrinsics sequences, and memory-selection config. +Those can be supplied in the global-conditioning slot, even if the adapter later +slices them internally while executing steps. If the caller must supply a fresh +value for every generated chunk, that value belongs in `step_fields`. + +`frequency_consumed` is a separate optional hint for how the adapter uses a +field internally, such as `once` or `per_step`. It does not decide where the +caller provides the value. A field can live in `global_conditioning_fields` and +still have `frequency_consumed="per_step"` when the adapter slices or reads +rollout-wide state during step execution. Third, `semantic_type` should be treated as a representation hint rather than a universal semantic type. For example, `prompt` may arrive as inline text or a @@ -173,7 +172,7 @@ Fourth, schema objects need open-ended metadata for future adapters. This lets a SANA-WM-like adapter advertise that `camera_trajectory_c2w` uses an `[F,4,4]` OpenCV camera-to-world sequence, or lets another model advertise a schema URI, units, coordinate frame, accepted file suffixes, cardinality hints, -or update notes. Metadata should remain query information and should not become +or adapter notes. Metadata should remain query information and should not become the compatibility type system. Fifth, `UserInputSchema` describes raw source capabilities, `CanonicalModality` @@ -181,7 +180,7 @@ describes what an application consumes, and mapping schemas describe derived model-facing semantics. A browser may provide `key_down`, `key_up`, `prompt_set`, and `initial_frame_set` events. Those become canonical modalities such as `driver_command` or `conditioning_prompt`; whether they can then drive -`steering`, `camera_trajectory`, or text embedding updates depends on the +`steering`, `camera_trajectory`, or text embeddings depends on the selected mapping and model schema. ## Implemented T2/T3 Shape @@ -189,23 +188,24 @@ selected mapping and model schema. The implementation that came out of this inventory is: 1. Keep `UserInputEvent` and `UserInputs` as the raw event API, sliced by a - half-open `TimeWindow`. Static startup values remain timestamp-zero events. + half-open `TimeWindow`. Static session-start values remain timestamp-zero + events. 2. Keep `UserInputSchema` lightweight and source-facing. `event_types` declares that an event type exists; `UserInputCapability` additionally pins the payload fields it carries. 3. Add a canonical layer between raw and encoded. `CanonicalModality` names a - device-independent input and its conditioning phase; `InputCanonicalizer` + device-independent input and its payload fields; `InputCanonicalizer` registers per-device converters and produces `CanonicalInputs`. Applications and mappings consume canonical inputs and never read raw device events. 4. Split `InferenceInput` (formerly `ModelInputs`) into `global_conditioning` - and `step`. A non-empty global slot mid-rollout is an update request, not a - reset; `InputField.update_policy` declares whether the model can apply it. -5. Extend `InputField` with `update_policy`, `lifecycle`, and `metadata` so - models can distinguish runtime config, cache initialization, rollout binding, - per-step inputs, and supported active-session updates. + and `step`. Global conditioning is session-global state; `step` is the + payload for one generated chunk or frame window. +5. Keep `InputField.semantic_type` and `metadata` as lightweight query hints, + while leaving tensor shape, cadence, and model-specific validation to + adapters and sessions. 6. Keep `InputMappingSchema` as the canonical-to-encoded boundary, with mapping-set compatibility helpers for composed mappings. -7. Keep input names, semantic types, lifecycle labels, and metadata open-ended. +7. Keep input names, semantic types, and metadata open-ended. Adding a new model should usually mean adding adapter-owned schema declarations and mappings, not changing the core input dataclasses. 8. Leave deep validation to model adapters, sessions, and mappings. The schema @@ -227,11 +227,9 @@ Use these conventions when adding future model schemas: instead of `array`, or `hdmap_frames` instead of `image`. - Use `semantic_type` for a coarse representation hint, such as `path`, `decoded_tensor`, `c2w_sequence`, `intrinsics_vec4_sequence`, or `embedding`. -- Use `lifecycle` to say where the adapter consumes the value, such as - `runtime_config`, `cache_init`, `rollout_binding`, `step_input`, or - `session_update`. -- Use `update_policy` to say when a value may change. `SESSION_START_ONLY` is - the one reserved token, meaning the value cannot be swapped mid-rollout. +- Use `frequency_consumed` for adapter-consumption cadence, such as `once` or + `per_step`; keep it independent from whether the field is declared under + `global_conditioning_fields` or `step_fields`. - Use `metadata` for query hints: units, coordinate frame, shape summary, accepted suffixes, schema URI, model family, value ranges, or cardinality. - Keep deep validation in the adapter/mapping. The lightweight schemas answer @@ -247,18 +245,13 @@ can describe the supported input surfaces. All use ```python lingbot_model = InferenceInputSchema( description="lingbot-world", - global_fields=( - InputField(name="prompt", lifecycle="cache_init"), - InputField(name="global_conditioning_frame", lifecycle="cache_init"), + global_conditioning_fields=( + InputField(name="prompt", frequency_consumed="once"), + InputField(name="global_conditioning_frame", frequency_consumed="once"), + InputField(name="text_embeddings", required=False, frequency_consumed="once"), ), step_fields=( - InputField(name="camera_trajectory", lifecycle="step_input"), - InputField( - name="text_embeddings", - required=False, - update_policy="step_boundary", - lifecycle="session_update", - ), + InputField(name="camera_trajectory", frequency_consumed="per_step"), ), ) ``` @@ -266,27 +259,29 @@ lingbot_model = InferenceInputSchema( ```python omnidreams_model = InferenceInputSchema( description="omnidreams", - global_fields=( - InputField(name="prompts", lifecycle="cache_init"), - InputField(name="global_conditioning_frames", lifecycle="cache_init"), - InputField(name="view_names", lifecycle="cache_init"), - InputField(name="text_embeddings", required=False, lifecycle="cache_init"), - InputField(name="image_embeddings", required=False, lifecycle="cache_init"), + global_conditioning_fields=( + InputField(name="prompts", frequency_consumed="once"), + InputField(name="global_conditioning_frames", frequency_consumed="once"), + InputField(name="view_names", frequency_consumed="once"), + InputField(name="text_embeddings", required=False, frequency_consumed="once"), + InputField(name="image_embeddings", required=False, frequency_consumed="once"), + ), + step_fields=( + InputField(name="hdmap_frames", frequency_consumed="per_step"), ), - step_fields=(InputField(name="hdmap_frames", lifecycle="step_input"),), ) ``` ```python hy_worldplay_model = InferenceInputSchema( description="hy-worldplay", - global_fields=( - InputField(name="prompt", lifecycle="cache_init"), - InputField(name="global_conditioning_frame", lifecycle="cache_init"), - InputField(name="action_labels", lifecycle="rollout_binding"), - InputField(name="camera_viewmats", lifecycle="rollout_binding"), - InputField(name="camera_intrinsics", lifecycle="rollout_binding"), - InputField(name="memory_config", lifecycle="rollout_binding"), + global_conditioning_fields=( + InputField(name="prompt", frequency_consumed="once"), + InputField(name="global_conditioning_frame", frequency_consumed="once"), + InputField(name="action_labels", frequency_consumed="per_step"), + InputField(name="camera_viewmats", frequency_consumed="per_step"), + InputField(name="camera_intrinsics", frequency_consumed="per_step"), + InputField(name="memory_config", frequency_consumed="per_step"), ), ) ``` @@ -294,21 +289,21 @@ hy_worldplay_model = InferenceInputSchema( ```python sana_wm_model = InferenceInputSchema( description="sana-wm", - global_fields=( - InputField(name="prompt", lifecycle="cache_init"), - InputField(name="negative_prompt", required=False, lifecycle="cache_init"), - InputField(name="global_conditioning_frame", lifecycle="cache_init"), + global_conditioning_fields=( + InputField(name="prompt", frequency_consumed="once"), + InputField(name="negative_prompt", required=False, frequency_consumed="once"), + InputField(name="global_conditioning_frame", frequency_consumed="once"), InputField( name="camera_trajectory_c2w", semantic_type="c2w_sequence", - lifecycle="rollout_binding", + frequency_consumed="per_step", metadata={"shape": "[F,4,4]", "coordinates": "opencv_c2w"}, ), InputField( name="camera_intrinsics_vec4", required=False, semantic_type="intrinsics_vec4_sequence", - lifecycle="rollout_binding", + frequency_consumed="per_step", metadata={"shape": "[F,4]"}, ), ), diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py index ab303c745..04f84ae54 100644 --- a/flashdreams/flashdreams/runtime/__init__.py +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -19,7 +19,6 @@ from flashdreams.runtime.config import ExecutionBackend, InferenceConfig, Precision from flashdreams.runtime.inputs import ( INPUT_PHASES, - SESSION_START_ONLY, CanonicalInputs, CanonicalInputSchema, CanonicalModality, @@ -96,7 +95,6 @@ "Precision", "RuntimeMetricSample", "ScriptedModality", - "SESSION_START_ONLY", "StepRequest", "StepResult", "TimeWindow", diff --git a/flashdreams/flashdreams/runtime/inputs.py b/flashdreams/flashdreams/runtime/inputs.py index f0be31bea..d88fbdb7a 100644 --- a/flashdreams/flashdreams/runtime/inputs.py +++ b/flashdreams/flashdreams/runtime/inputs.py @@ -12,23 +12,17 @@ from flashdreams.runtime._utils import freeze_mapping -InputPhase = Literal["global", "step"] +InputPhase = Literal["global_conditioning", "step"] -INPUT_PHASES: tuple[InputPhase, ...] = ("global", "step") - -SESSION_START_ONLY = "session_start" -"""``InputField.update_policy`` value meaning "supply at session start only". - -``update_policy`` is otherwise an open, adapter-owned vocabulary. This is the -one reserved token, because the runtime needs to distinguish a conditioning -value that can be swapped mid-rollout from one that cannot. -""" +INPUT_PHASES: tuple[InputPhase, ...] = ("global_conditioning", "step") def validate_phase(value: str) -> InputPhase: """Return ``value`` as a validated :data:`InputPhase`.""" if value not in INPUT_PHASES: - raise ValueError(f"phase must be 'global' or 'step', got {value!r}.") + raise ValueError( + f"phase must be 'global_conditioning' or 'step', got {value!r}." + ) return cast(InputPhase, value) @@ -56,17 +50,15 @@ def contains(self, timestamp_s: float) -> bool: class InputField: """Lightweight schema field for user snapshots or model inputs. - ``update_policy`` and ``lifecycle`` are plain query metadata. They let a - model advertise facts such as "prompt updates land at step boundaries" or - "this value is consumed at cache init" without making this layer - responsible for implementing or deeply validating that behavior. + ``semantic_type``, ``frequency_consumed``, and ``metadata`` are query hints + only. Adapter-owned validation still decides concrete shape, dtype, units, + and tensor layout. """ name: str required: bool = True semantic_type: str | None = None - update_policy: str | None = None - lifecycle: str | None = None + frequency_consumed: str | None = None metadata: Mapping[str, Any] = field( default_factory=dict, compare=False, @@ -205,21 +197,21 @@ def require_snapshot(self, inputs: "UserInputs") -> None: @dataclass(frozen=True, kw_only=True, slots=True) class InferenceInputSchema: - """Minimal metadata for model-facing initial and per-step inputs.""" + """Minimal metadata for global conditioning and per-step inputs.""" - global_fields: tuple[InputField, ...] = () - """Model inputs required before starting the initial generation/session.""" + global_conditioning_fields: tuple[InputField, ...] = () + """Model inputs carried in the global conditioning slot.""" step_fields: tuple[InputField, ...] = () - """Per-step model inputs required after the session starts.""" + """Model inputs required for one session step.""" description: str = "" def fields_for(self, phase: InputPhase) -> tuple[InputField, ...]: """Return every declared field for ``phase``.""" return ( - self.global_fields - if validate_phase(phase) == "global" + self.global_conditioning_fields + if validate_phase(phase) == "global_conditioning" else self.step_fields ) @@ -258,32 +250,20 @@ def _select( if input_field.required is required ) - def unsupported_global_updates(self, inputs: "InferenceInput") -> tuple[str, ...]: - """Return requested conditioning updates this model cannot apply. - - A field whose ``update_policy`` is :data:`SESSION_START_ONLY` can be - supplied when the session starts but not changed mid-rollout. Any other - policy, including ``None``, is treated as permissive here; the adapter - still owns whether the swap actually succeeds. - """ - return tuple( - name - for name in inputs.global_conditioning - if (declared := self.field_for(name=name, phase="global")) is not None - and declared.update_policy == SESSION_START_ONLY + def missing_global_conditioning(self, inputs: "InferenceInput") -> tuple[str, ...]: + """Return required global conditioning fields absent from ``inputs``.""" + return _missing_required( + self.global_conditioning_fields, + inputs.global_conditioning, ) - def missing_global(self, inputs: "InferenceInput") -> tuple[str, ...]: - """Return required initial fields absent from ``inputs``.""" - return _missing_required(self.global_fields, inputs.global_conditioning) - def missing_step(self, inputs: "InferenceInput") -> tuple[str, ...]: """Return required per-step fields absent from ``inputs``.""" return _missing_required(self.step_fields, inputs.step) - def require_global(self, inputs: "InferenceInput") -> None: - """Raise if required initial fields are absent.""" - missing = self.missing_global(inputs) + def require_global_conditioning(self, inputs: "InferenceInput") -> None: + """Raise if required global conditioning fields are absent.""" + missing = self.missing_global_conditioning(inputs) if missing: raise ValueError( f"Missing required global conditioning input(s): {missing}" @@ -447,16 +427,10 @@ class InferenceInput: Two conditioning slots: - ``global_conditioning``: values that condition the whole rollout, such as - the conditioning frame or prompt. Normally supplied when the session - starts. + the conditioning frame or prompt. Session start/reset establishes this + state; a step call may carry a non-empty payload to request an update when + the model supports it. - ``step``: values needed to generate the next chunk or frame. - - A non-empty ``global_conditioning`` on a mid-rollout input is an *update - request*, not a reset. The session should apply it when the model supports - that; resetting rollout state is a separate, explicit - :meth:`InferenceSession.reset` call. Whether a given value can be updated - mid-rollout is declared per field by ``InputField.update_policy``; see - :meth:`InferenceInputSchema.unsupported_global_updates`. """ __hash__ = None @@ -472,42 +446,12 @@ def __post_init__(self) -> None: object.__setattr__(self, "step", freeze_mapping(self.step)) object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) - @property - def requests_global_update(self) -> bool: - """Return whether this input asks the session to update conditioning.""" - return bool(self.global_conditioning) - - def with_step(self, step: Mapping[str, Any]) -> "InferenceInput": - """Return a copy with replaced per-step payload. - - The global slot is carried through unchanged, so a mid-rollout input - built this way keeps whatever update request it already had. Use - :meth:`without_global_update` for the common steady-state case. - """ - return InferenceInput( - global_conditioning=self.global_conditioning, - step=step, - metadata=self.metadata, - ) - - def with_global_update( - self, global_conditioning: Mapping[str, Any] - ) -> "InferenceInput": - """Return a copy requesting a mid-rollout conditioning update.""" - return InferenceInput( - global_conditioning=global_conditioning, - step=self.step, - metadata=self.metadata, - ) - - def without_global_update(self) -> "InferenceInput": - """Return a copy that requests no conditioning update.""" - return InferenceInput(step=self.step, metadata=self.metadata) - def for_phase(self, phase: InputPhase) -> Mapping[str, Any]: """Return the payload mapping for ``phase``.""" return ( - self.global_conditioning if validate_phase(phase) == "global" else self.step + self.global_conditioning + if validate_phase(phase) == "global_conditioning" + else self.step ) diff --git a/flashdreams/flashdreams/runtime/interfaces.py b/flashdreams/flashdreams/runtime/interfaces.py index 852a77f1c..5c5054dc4 100644 --- a/flashdreams/flashdreams/runtime/interfaces.py +++ b/flashdreams/flashdreams/runtime/interfaces.py @@ -22,7 +22,7 @@ class InferenceSession(Protocol): """One rollout or stream with isolated model/cache state.""" def next_step_request(self) -> StepRequest | None: - """Describe the next step's inputs, or return ``None`` when complete.""" + """Return the next step's runtime request, or ``None`` when complete.""" ... def step(self, inputs: InferenceInput) -> StepResult: @@ -69,7 +69,7 @@ def model_id(self) -> str: @property def inference_input_schema(self) -> InferenceInputSchema: - """Model-facing initial and per-step input requirements.""" + """Model-facing global conditioning and per-step input requirements.""" ... @property diff --git a/flashdreams/flashdreams/runtime/mapping.py b/flashdreams/flashdreams/runtime/mapping.py index 94f481406..bbf11616e 100644 --- a/flashdreams/flashdreams/runtime/mapping.py +++ b/flashdreams/flashdreams/runtime/mapping.py @@ -41,13 +41,13 @@ def validate( """Fail early for obvious app, event-source, and model mismatches.""" ... - def map_global_inputs( + def map_global_conditioning_inputs( self, *, canonical_inputs: CanonicalInputs, inference_input: InferenceInput, ) -> InferenceInput: - """Build global conditioning inputs before a session starts.""" + """Build global conditioning inputs for session start or reset.""" ... def map_step_inputs( @@ -72,7 +72,7 @@ def validate( ) -> None: del canonical_schema, inference_input_schema - def map_global_inputs( + def map_global_conditioning_inputs( self, *, canonical_inputs: CanonicalInputs, @@ -104,7 +104,7 @@ class InputMappingSchema: name: str = "input-mapping" consumes: tuple[CanonicalModality, ...] = () - produces_global: tuple[InputField, ...] = () + produces_global_conditioning: tuple[InputField, ...] = () produces_step: tuple[InputField, ...] = () metadata: Mapping[str, Any] = field( default_factory=dict, @@ -119,7 +119,11 @@ def __post_init__(self) -> None: def produces_for(self, phase: InputPhase) -> tuple[InputField, ...]: """Return the fields this mapping produces for ``phase``.""" - return self.produces_global if phase == "global" else self.produces_step + return ( + self.produces_global_conditioning + if phase == "global_conditioning" + else self.produces_step + ) def can_produce(self, phase: InputPhase, required: InputField) -> bool: """Return whether this mapping can produce ``required`` in ``phase``.""" @@ -136,12 +140,7 @@ def _field_matches(produced: InputField, required: InputField) -> bool: or required.semantic_type is None or produced.semantic_type == required.semantic_type ) - lifecycle_ok = ( - produced.lifecycle is None - or required.lifecycle is None - or produced.lifecycle == required.lifecycle - ) - return semantic_ok and lifecycle_ok + return semantic_ok @dataclass(frozen=True, kw_only=True, slots=True) @@ -223,7 +222,10 @@ def combine_mapping_schemas( first declaration winning on conflicting keys. """ consumes: list[CanonicalModality] = [] - produces: dict[InputPhase, list[InputField]] = {"global": [], "step": []} + produces: dict[InputPhase, list[InputField]] = { + "global_conditioning": [], + "step": [], + } def _merge(target: list[Any], value: Any) -> None: for index, existing in enumerate(target): @@ -248,7 +250,7 @@ def _merge(target: list[Any], value: Any) -> None: return InputMappingSchema( name=name, consumes=tuple(consumes), - produces_global=tuple(produces["global"]), + produces_global_conditioning=tuple(produces["global_conditioning"]), produces_step=tuple(produces["step"]), ) @@ -358,8 +360,9 @@ def undeclared_inference_inputs( """Return payload keys a mapping produced but did not declare. Mapping schemas are hand-written, so they drift from what - ``map_global_inputs``/``map_step_inputs`` actually return. Mapping tests - can use this to keep the declared compatibility surface honest. + ``map_global_conditioning_inputs``/``map_step_inputs`` actually return. + Mapping tests can use this to keep the declared compatibility surface + honest. """ return tuple( (phase, key) diff --git a/flashdreams/flashdreams/runtime/types.py b/flashdreams/flashdreams/runtime/types.py index 467753026..51d3846db 100644 --- a/flashdreams/flashdreams/runtime/types.py +++ b/flashdreams/flashdreams/runtime/types.py @@ -15,10 +15,11 @@ @dataclass(frozen=True, kw_only=True, slots=True) class StepRequest: - """Model-session request for the next step's inputs. + """Per-step runtime request emitted by an inference session. - ``user_input_window`` lets a runner drain or slice timestamped user events for - the current step before invoking the selected ``InputMapping``. + This is not a schema declaration. ``user_input_window`` lets a runner drain + or slice timestamped user events for the current step before invoking the + selected ``InputMapping``. """ __hash__ = None @@ -36,7 +37,7 @@ def __post_init__(self) -> None: @dataclass(frozen=True, kw_only=True, slots=True) class StepResult: - """Generated output and metadata for one inference step.""" + """Generated output and metadata returned by one inference step.""" __hash__ = None diff --git a/flashdreams/tests/test_inference_runtime_api.py b/flashdreams/tests/test_inference_runtime_api.py index edfafa634..d454f6205 100644 --- a/flashdreams/tests/test_inference_runtime_api.py +++ b/flashdreams/tests/test_inference_runtime_api.py @@ -106,9 +106,9 @@ def test_runtime_metric_sample_rejects_bool_values() -> None: RuntimeMetricSample(name="sample", value=True) -def test_inference_input_schema_validates_initial_and_step_payloads() -> None: +def test_schema_validates_global_conditioning_and_step_payloads() -> None: schema = InferenceInputSchema( - global_fields=( + global_conditioning_fields=( InputField(name="prompt"), InputField(name="global_conditioning_frame"), ), @@ -118,7 +118,7 @@ def test_inference_input_schema_validates_initial_and_step_payloads() -> None: global_conditioning={"prompt": "drive", "global_conditioning_frame": object()} ) - schema.require_global(inputs) + schema.require_global_conditioning(inputs) assert schema.missing_step(inputs) == ("camera_poses",) with pytest.raises(ValueError, match="camera_poses"): @@ -190,7 +190,7 @@ def test_identity_input_mapping_leaves_inference_input_unchanged() -> None: request = StepRequest(step_index=0) assert ( - mapping.map_global_inputs( + mapping.map_global_conditioning_inputs( canonical_inputs=CanonicalInputs(), inference_input=inference_input, ) @@ -367,7 +367,7 @@ def _drive_two_step_session( canonical_schema=adapter.canonical_input_schema, inference_input_schema=adapter.inference_input_schema, ) - initial_inputs = mapping.map_global_inputs( + initial_inputs = mapping.map_global_conditioning_inputs( canonical_inputs=canonicalizer.canonicalize( user_inputs, window=TimeWindow(start_s=0.0, end_s=_SESSION_HORIZON_S), @@ -390,9 +390,8 @@ def _drive_two_step_session( or TimeWindow(start_s=0.0, end_s=_SESSION_HORIZON_S), source_schema=source_schema, ), - # The global slot stays empty in steady state. A mapping that - # sees ``canonical_inputs.has_global_change`` fills it via - # ``with_global_update`` to request a mid-rollout swap. + # Per-step calls carry only the step payload. A changed prompt + # or scene starts or resets a session outside this loop. inference_input=InferenceInput( step={"chunk_index": request.step_index}, ), @@ -417,7 +416,7 @@ def _drive_two_step_session( class _FakeAdapter: model_id = "fake-model" inference_input_schema = InferenceInputSchema( - global_fields=(InputField(name="prompt"),), + global_conditioning_fields=(InputField(name="prompt"),), step_fields=(InputField(name="chunk_index"),), ) canonical_input_schema = CanonicalInputSchema() @@ -440,7 +439,7 @@ def __init__(self, *, inference_input_schema: InferenceInputSchema) -> None: self.closed = False def start_session(self, inputs: InferenceInput) -> InferenceSession: - self._inference_input_schema.require_global(inputs) + self._inference_input_schema.require_global_conditioning(inputs) return _FakeSession(inference_input_schema=self._inference_input_schema) def close(self) -> None: diff --git a/flashdreams/tests/test_runtime_canonical.py b/flashdreams/tests/test_runtime_canonical.py index 1ad48d39e..cfe2e3646 100644 --- a/flashdreams/tests/test_runtime_canonical.py +++ b/flashdreams/tests/test_runtime_canonical.py @@ -217,13 +217,14 @@ def test_canonical_inputs_carry_live_control_only() -> None: def test_application_supplies_global_conditioning_directly() -> None: - """A prompt swap reaches the session without touching canonicalization.""" - update = InferenceInput(step={"steering": 0.0}).with_global_update( - {"prompt": "heavy rain"} + """A prompt reaches session start without touching canonicalization.""" + inputs = InferenceInput( + global_conditioning={"prompt": "heavy rain"}, + step={"steering": 0.0}, ) - assert update.requests_global_update - assert update.global_conditioning["prompt"] == "heavy rain" + assert inputs.global_conditioning["prompt"] == "heavy rain" + assert inputs.step["steering"] == 0.0 # --- device independence ------------------------------------------------ diff --git a/flashdreams/tests/test_runtime_input_mapping.py b/flashdreams/tests/test_runtime_input_mapping.py index 00cd9758f..0750e733e 100644 --- a/flashdreams/tests/test_runtime_input_mapping.py +++ b/flashdreams/tests/test_runtime_input_mapping.py @@ -4,9 +4,10 @@ """Tests for declarative input-mapping compatibility in the runtime API. These cover the T2/T3 contract: sources declare what user events they can -provide at payload granularity, models declare required and optional -initial/per-step inputs, and a mapping declares what it consumes and produces so -compatibility can be answered before expensive runtime initialization. +provide at payload granularity, models declare required and optional global +conditioning/per-step inputs, and a mapping declares what it consumes and +produces so compatibility can be answered before expensive runtime +initialization. """ from __future__ import annotations @@ -17,7 +18,6 @@ from flashdreams.runtime import ( DRIVER_COMMAND, - SESSION_START_ONLY, CanonicalInputs, CanonicalInputSchema, CanonicalModality, @@ -66,11 +66,13 @@ # modality, so this mapping consumes nothing and only declares what it produces. PROMPT_MAPPING = InputMappingSchema( name="prompt", - produces_global=(InputField(name="prompt", semantic_type="text"),), + produces_global_conditioning=(InputField(name="prompt", semantic_type="text"),), ) FRAME_MAPPING = InputMappingSchema( name="conditioning-frame", - produces_global=(InputField(name="global_conditioning_frame", required=False),), + produces_global_conditioning=( + InputField(name="global_conditioning_frame", required=False), + ), ) STEERING_MAPPING = InputMappingSchema( name="driver-command-to-steering", @@ -84,12 +86,10 @@ ) DRIVING_MODEL = InferenceInputSchema( - global_fields=( - InputField(name="prompt", semantic_type="text", lifecycle="cache_init"), - ), + global_conditioning_fields=(InputField(name="prompt", semantic_type="text"),), step_fields=( - InputField(name="steering", lifecycle="step_input"), - InputField(name="camera_delta", required=False, lifecycle="step_input"), + InputField(name="steering"), + InputField(name="camera_delta", required=False), ), ) @@ -97,7 +97,7 @@ # --- user input events and windowing ------------------------------------ -def test_startup_values_are_represented_as_events() -> None: +def test_session_start_values_are_represented_as_events() -> None: inputs = UserInputs( events=( UserInputEvent( @@ -198,7 +198,7 @@ def test_model_declares_required_and_optional_fields_per_phase() -> None: optional = DRIVING_MODEL.optional_fields() assert {(phase, f.name) for phase, f in required} == { - ("global", "prompt"), + ("global_conditioning", "prompt"), ("step", "steering"), } assert {(phase, f.name) for phase, f in optional} == {("step", "camera_delta")} @@ -211,7 +211,10 @@ def test_required_fields_can_be_filtered_by_phase() -> None: def test_field_lookup_is_phase_scoped() -> None: - assert DRIVING_MODEL.field_for(name="prompt", phase="global") is not None + assert ( + DRIVING_MODEL.field_for(name="prompt", phase="global_conditioning") + is not None + ) assert DRIVING_MODEL.field_for(name="prompt", phase="step") is None @@ -227,20 +230,28 @@ def test_inference_input_expose_payload_per_phase() -> None: global_conditioning={"prompt": "drive"}, step={"steering": 0.25} ) - assert inputs.for_phase("global")["prompt"] == "drive" + assert inputs.for_phase("global_conditioning")["prompt"] == "drive" assert inputs.for_phase("step")["steering"] == 0.25 -def test_lifecycle_and_update_policy_are_queryable_metadata() -> None: +def test_step_context_can_carry_global_conditioning_update_payload() -> None: + inputs = InferenceInput( + global_conditioning={"prompt": "heavy rain"}, + step={"steering": 0.25}, + ) + + assert inputs.global_conditioning["prompt"] == "heavy rain" + assert inputs.step["steering"] == 0.25 + + +def test_field_metadata_is_queryable() -> None: field = InputField( name="prompt", - update_policy="step_boundary", - lifecycle="cache_init", + frequency_consumed="once", metadata={"coordinates": "opencv_c2w"}, ) - assert field.update_policy == "step_boundary" - assert field.lifecycle == "cache_init" + assert field.frequency_consumed == "once" assert field.metadata["coordinates"] == "opencv_c2w" @@ -263,7 +274,7 @@ def test_compatible_source_model_and_mapping_can_drive() -> None: assert compatibility.can_drive assert {(p, f.name) for p, f in compatibility.satisfied_required_model_fields} == { - ("global", "prompt"), + ("global_conditioning", "prompt"), ("step", "steering"), } assert {(p, f.name) for p, f in compatibility.available_optional_model_fields} == { @@ -325,13 +336,17 @@ def test_optional_field_needs_mapping_support_to_be_available() -> None: assert compatibility.available_optional_model_fields == () -def test_lifecycle_disagreement_blocks_a_field_match() -> None: +def test_global_conditioning_mapping_matches_global_conditioning_field() -> None: model = InferenceInputSchema( - global_fields=(InputField(name="prompt", lifecycle="rollout_binding"),) + global_conditioning_fields=( + InputField(name="camera_trajectory", frequency_consumed="per_step"), + ) ) mapping = InputMappingSchema( - name="prompt", - produces_global=(InputField(name="prompt", lifecycle="cache_init"),), + name="trajectory", + produces_global_conditioning=( + InputField(name="camera_trajectory", frequency_consumed="once"), + ), ) compatibility = check_mapping_compatibility( @@ -340,11 +355,13 @@ def test_lifecycle_disagreement_blocks_a_field_match() -> None: mapping_schema=mapping, ) - assert not compatibility.can_drive + assert compatibility.can_drive -def test_unspecified_lifecycle_stays_permissive() -> None: - model = InferenceInputSchema(global_fields=(InputField(name="prompt"),)) +def test_unspecified_semantic_type_stays_permissive() -> None: + model = InferenceInputSchema( + global_conditioning_fields=(InputField(name="prompt"),) + ) compatibility = check_mapping_compatibility( canonical_schema=CANONICAL_ALL, @@ -398,26 +415,28 @@ def test_combining_mappings_unions_their_surfaces() -> None: combined = combine_mapping_schemas((PROMPT_MAPPING, STEERING_MAPPING)) assert {m.name for m in combined.consumes} == {"driver_command"} - assert [f.name for f in combined.produces_global] == ["prompt"] + assert [f.name for f in combined.produces_global_conditioning] == ["prompt"] assert [f.name for f in combined.produces_step] == ["steering"] def test_duplicate_declarations_collapse_and_merge_metadata() -> None: first = InputMappingSchema( name="a", - produces_global=(InputField(name="prompt", metadata={"source": "a"}),), + produces_global_conditioning=( + InputField(name="prompt", metadata={"source": "a"}), + ), ) second = InputMappingSchema( name="b", - produces_global=( + produces_global_conditioning=( InputField(name="prompt", metadata={"source": "b", "extra": "kept"}), ), ) combined = combine_mapping_schemas((first, second)) - assert len(combined.produces_global) == 1 - metadata = combined.produces_global[0].metadata + assert len(combined.produces_global_conditioning) == 1 + metadata = combined.produces_global_conditioning[0].metadata assert metadata["source"] == "a" assert metadata["extra"] == "kept" @@ -489,85 +508,3 @@ def test_model_with_no_requirements_is_always_drivable() -> None: ) assert compatibility.can_drive - - -# --- global conditioning updates vs reset ------------------------------- - - -def test_empty_global_slot_requests_no_update() -> None: - steady_state = InferenceInput(step={"steering": 0.25}) - - assert not steady_state.requests_global_update - - -def test_non_empty_global_slot_mid_rollout_is_an_update_request() -> None: - """Changing weather mid-run updates conditioning; it is not a reset.""" - updated = InferenceInput(step={"steering": 0.0}).with_global_update( - {"prompt": "heavy rain"} - ) - - assert updated.requests_global_update - assert updated.global_conditioning["prompt"] == "heavy rain" - assert updated.step["steering"] == 0.0 - - -def test_with_step_carries_the_global_slot_through() -> None: - started = InferenceInput(global_conditioning={"prompt": "drive"}) - - stepped = started.with_step({"steering": 0.5}) - - assert stepped.global_conditioning["prompt"] == "drive" - - -def test_without_global_update_clears_the_request() -> None: - started = InferenceInput( - global_conditioning={"prompt": "drive"}, step={"steering": 0.5} - ) - - steady_state = started.without_global_update() - - assert not steady_state.requests_global_update - assert steady_state.step["steering"] == 0.5 - - -def test_model_can_declare_conditioning_it_cannot_swap_mid_rollout() -> None: - schema = InferenceInputSchema( - global_fields=( - InputField(name="prompt", update_policy="step_boundary"), - InputField(name="scene_id", update_policy=SESSION_START_ONLY), - ) - ) - update = InferenceInput( - global_conditioning={"prompt": "heavy rain", "scene_id": "town_02"} - ) - - assert schema.unsupported_global_updates(update) == ("scene_id",) - - -def test_permissive_when_no_update_policy_is_declared() -> None: - schema = InferenceInputSchema(global_fields=(InputField(name="prompt"),)) - update = InferenceInput(global_conditioning={"prompt": "heavy rain"}) - - assert schema.unsupported_global_updates(update) == () - - -def test_undeclared_global_values_are_left_to_the_adapter() -> None: - schema = InferenceInputSchema(global_fields=(InputField(name="prompt"),)) - update = InferenceInput(global_conditioning={"mystery": 1}) - - assert schema.unsupported_global_updates(update) == () - - -def test_steady_state_steps_do_not_request_a_global_update() -> None: - """Carrying session-start conditioning forward would look like an update.""" - started = InferenceInput(global_conditioning={"prompt": "drive"}) - - steady_state = InferenceInput(step={"chunk_index": 1}) - - assert started.requests_global_update - assert not steady_state.requests_global_update - assert ( - not started.with_step({"chunk_index": 1}) - .without_global_update() - .requests_global_update - ) From c0ec4315c1a14d237dc03b64ea0e3b4d3966852f Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Wed, 5 Aug 2026 19:14:21 -0700 Subject: [PATCH 05/19] Rename misleading `semantic_type` to `input_modality` --- ...ence_runtime_supported_inputs_inventory.md | 94 +++++++++++++------ flashdreams/flashdreams/runtime/inputs.py | 23 +++-- flashdreams/flashdreams/runtime/mapping.py | 13 +-- .../tests/test_runtime_input_mapping.py | 41 ++++---- 4 files changed, 108 insertions(+), 63 deletions(-) diff --git a/docs/inference_runtime_supported_inputs_inventory.md b/docs/inference_runtime_supported_inputs_inventory.md index d56cf1651..f82f3192b 100644 --- a/docs/inference_runtime_supported_inputs_inventory.md +++ b/docs/inference_runtime_supported_inputs_inventory.md @@ -161,12 +161,12 @@ caller provides the value. A field can live in `global_conditioning_fields` and still have `frequency_consumed="per_step"` when the adapter slices or reads rollout-wide state during step execution. -Third, `semantic_type` should be treated as a representation hint rather than a -universal semantic type. For example, `prompt` may arrive as inline text or a -path but become prompt text or text embeddings; the global conditioning frame -may arrive as a path, URL, bytes, or decoded tensor; camera motion may arrive as keys, pose JSON, -Numpy arrays, or integrated tensors. The semantic input name is still the main -contract. +Third, `name` is the semantic model input role, while `input_modality` is only +a coarse value-kind hint. For example, `prompt` and `negative_prompt` are +different semantic names even though both usually have `input_modality="text"`. +The semantic input name is the main contract; source details such as path, URL, +bytes, decoded tensor layout, accepted suffixes, or file schema belong in +adapter validation or `metadata`. Fourth, schema objects need open-ended metadata for future adapters. This lets a SANA-WM-like adapter advertise that `camera_trajectory_c2w` uses an @@ -200,12 +200,13 @@ The implementation that came out of this inventory is: 4. Split `InferenceInput` (formerly `ModelInputs`) into `global_conditioning` and `step`. Global conditioning is session-global state; `step` is the payload for one generated chunk or frame window. -5. Keep `InputField.semantic_type` and `metadata` as lightweight query hints, - while leaving tensor shape, cadence, and model-specific validation to - adapters and sessions. +5. Keep `InputField.input_modality`, `frequency_consumed`, and `metadata` as + lightweight query hints, while leaving tensor shape and model-specific + validation to adapters and sessions. `InputField.name` remains the semantic + payload key. 6. Keep `InputMappingSchema` as the canonical-to-encoded boundary, with mapping-set compatibility helpers for composed mappings. -7. Keep input names, semantic types, and metadata open-ended. +7. Keep input names, input modalities, and metadata open-ended. Adding a new model should usually mean adding adapter-owned schema declarations and mappings, not changing the core input dataclasses. 8. Leave deep validation to model adapters, sessions, and mappings. The schema @@ -225,13 +226,14 @@ Use these conventions when adding future model schemas: - Prefer semantic names over modality names, such as `camera_trajectory_c2w` instead of `array`, or `hdmap_frames` instead of `image`. -- Use `semantic_type` for a coarse representation hint, such as `path`, - `decoded_tensor`, `c2w_sequence`, `intrinsics_vec4_sequence`, or `embedding`. +- Use `input_modality` for a coarse value-kind hint, such as `text`, `image`, + `embedding`, `c2w_sequence`, or `intrinsics_vec4_sequence`. +- Use `metadata` for representation details such as paths, decoded tensor + layout, units, coordinate frame, shape summary, accepted suffixes, schema URI, + model family, value ranges, or cardinality. - Use `frequency_consumed` for adapter-consumption cadence, such as `once` or `per_step`; keep it independent from whether the field is declared under `global_conditioning_fields` or `step_fields`. -- Use `metadata` for query hints: units, coordinate frame, shape summary, - accepted suffixes, schema URI, model family, value ranges, or cardinality. - Keep deep validation in the adapter/mapping. The lightweight schemas answer whether the selected source and mapping can plausibly drive the model before expensive initialization. @@ -246,9 +248,18 @@ can describe the supported input surfaces. All use lingbot_model = InferenceInputSchema( description="lingbot-world", global_conditioning_fields=( - InputField(name="prompt", frequency_consumed="once"), - InputField(name="global_conditioning_frame", frequency_consumed="once"), - InputField(name="text_embeddings", required=False, frequency_consumed="once"), + InputField(name="prompt", input_modality="text", frequency_consumed="once"), + InputField( + name="global_conditioning_frame", + input_modality="image", + frequency_consumed="once", + ), + InputField( + name="text_embeddings", + required=False, + input_modality="embedding", + frequency_consumed="once", + ), ), step_fields=( InputField(name="camera_trajectory", frequency_consumed="per_step"), @@ -260,11 +271,25 @@ lingbot_model = InferenceInputSchema( omnidreams_model = InferenceInputSchema( description="omnidreams", global_conditioning_fields=( - InputField(name="prompts", frequency_consumed="once"), - InputField(name="global_conditioning_frames", frequency_consumed="once"), + InputField(name="prompts", input_modality="text", frequency_consumed="once"), + InputField( + name="global_conditioning_frames", + input_modality="image", + frequency_consumed="once", + ), InputField(name="view_names", frequency_consumed="once"), - InputField(name="text_embeddings", required=False, frequency_consumed="once"), - InputField(name="image_embeddings", required=False, frequency_consumed="once"), + InputField( + name="text_embeddings", + required=False, + input_modality="embedding", + frequency_consumed="once", + ), + InputField( + name="image_embeddings", + required=False, + input_modality="embedding", + frequency_consumed="once", + ), ), step_fields=( InputField(name="hdmap_frames", frequency_consumed="per_step"), @@ -276,8 +301,12 @@ omnidreams_model = InferenceInputSchema( hy_worldplay_model = InferenceInputSchema( description="hy-worldplay", global_conditioning_fields=( - InputField(name="prompt", frequency_consumed="once"), - InputField(name="global_conditioning_frame", frequency_consumed="once"), + InputField(name="prompt", input_modality="text", frequency_consumed="once"), + InputField( + name="global_conditioning_frame", + input_modality="image", + frequency_consumed="once", + ), InputField(name="action_labels", frequency_consumed="per_step"), InputField(name="camera_viewmats", frequency_consumed="per_step"), InputField(name="camera_intrinsics", frequency_consumed="per_step"), @@ -290,19 +319,28 @@ hy_worldplay_model = InferenceInputSchema( sana_wm_model = InferenceInputSchema( description="sana-wm", global_conditioning_fields=( - InputField(name="prompt", frequency_consumed="once"), - InputField(name="negative_prompt", required=False, frequency_consumed="once"), - InputField(name="global_conditioning_frame", frequency_consumed="once"), + InputField(name="prompt", input_modality="text", frequency_consumed="once"), + InputField( + name="negative_prompt", + required=False, + input_modality="text", + frequency_consumed="once", + ), + InputField( + name="global_conditioning_frame", + input_modality="image", + frequency_consumed="once", + ), InputField( name="camera_trajectory_c2w", - semantic_type="c2w_sequence", + input_modality="c2w_sequence", frequency_consumed="per_step", metadata={"shape": "[F,4,4]", "coordinates": "opencv_c2w"}, ), InputField( name="camera_intrinsics_vec4", required=False, - semantic_type="intrinsics_vec4_sequence", + input_modality="intrinsics_vec4_sequence", frequency_consumed="per_step", metadata={"shape": "[F,4]"}, ), diff --git a/flashdreams/flashdreams/runtime/inputs.py b/flashdreams/flashdreams/runtime/inputs.py index d88fbdb7a..9174b6a84 100644 --- a/flashdreams/flashdreams/runtime/inputs.py +++ b/flashdreams/flashdreams/runtime/inputs.py @@ -50,14 +50,15 @@ def contains(self, timestamp_s: float) -> bool: class InputField: """Lightweight schema field for user snapshots or model inputs. - ``semantic_type``, ``frequency_consumed``, and ``metadata`` are query hints - only. Adapter-owned validation still decides concrete shape, dtype, units, - and tensor layout. + ``name`` is the model-facing input role and payload key, such as ``prompt`` + or ``negative_prompt``. ``input_modality``, ``frequency_consumed``, and + ``metadata`` are query hints only. Adapter-owned validation still decides + concrete shape, dtype, units, and tensor layout. """ name: str required: bool = True - semantic_type: str | None = None + input_modality: str | None = None frequency_consumed: str | None = None metadata: Mapping[str, Any] = field( default_factory=dict, @@ -83,7 +84,7 @@ class UserInputCapability: """ event_type: str - semantic_type: str | None = None + input_modality: str | None = None payload_fields: frozenset[str] = field(default_factory=frozenset) metadata: Mapping[str, Any] = field( default_factory=dict, @@ -104,12 +105,14 @@ def is_satisfied_by(self, provider: "UserInputCapability") -> bool: """Return whether ``provider`` can satisfy this consumed capability.""" if self.event_type != provider.event_type: return False - semantic_ok = ( - self.semantic_type is None - or provider.semantic_type is None - or self.semantic_type == provider.semantic_type + input_modality_ok = ( + self.input_modality is None + or provider.input_modality is None + or self.input_modality == provider.input_modality + ) + return input_modality_ok and self.payload_fields.issubset( + provider.payload_fields ) - return semantic_ok and self.payload_fields.issubset(provider.payload_fields) @dataclass(frozen=True, kw_only=True, slots=True) diff --git a/flashdreams/flashdreams/runtime/mapping.py b/flashdreams/flashdreams/runtime/mapping.py index bbf11616e..710ad8043 100644 --- a/flashdreams/flashdreams/runtime/mapping.py +++ b/flashdreams/flashdreams/runtime/mapping.py @@ -135,12 +135,12 @@ def can_produce(self, phase: InputPhase, required: InputField) -> bool: def _field_matches(produced: InputField, required: InputField) -> bool: if produced.name != required.name: return False - semantic_ok = ( - produced.semantic_type is None - or required.semantic_type is None - or produced.semantic_type == required.semantic_type + input_modality_ok = ( + produced.input_modality is None + or required.input_modality is None + or produced.input_modality == required.input_modality ) - return semantic_ok + return input_modality_ok @dataclass(frozen=True, kw_only=True, slots=True) @@ -369,7 +369,8 @@ def undeclared_inference_inputs( for phase in INPUT_PHASES for key in inputs.for_phase(phase) if not any( - declared.name == key for declared in mapping_schema.produces_for(phase) + declared.name == key + for declared in mapping_schema.produces_for(phase) ) ) diff --git a/flashdreams/tests/test_runtime_input_mapping.py b/flashdreams/tests/test_runtime_input_mapping.py index 0750e733e..b5f0744b9 100644 --- a/flashdreams/tests/test_runtime_input_mapping.py +++ b/flashdreams/tests/test_runtime_input_mapping.py @@ -44,7 +44,7 @@ KEY_UP = UserInputCapability(event_type="key_up", payload_fields=frozenset({"key"})) PROMPT_SET = UserInputCapability( event_type="prompt_set", - semantic_type="text", + input_modality="text", payload_fields=frozenset({"prompt"}), ) FRAME_SET = UserInputCapability( @@ -66,7 +66,7 @@ # modality, so this mapping consumes nothing and only declares what it produces. PROMPT_MAPPING = InputMappingSchema( name="prompt", - produces_global_conditioning=(InputField(name="prompt", semantic_type="text"),), + produces_global_conditioning=(InputField(name="prompt", input_modality="text"),), ) FRAME_MAPPING = InputMappingSchema( name="conditioning-frame", @@ -86,7 +86,7 @@ ) DRIVING_MODEL = InferenceInputSchema( - global_conditioning_fields=(InputField(name="prompt", semantic_type="text"),), + global_conditioning_fields=(InputField(name="prompt", input_modality="text"),), step_fields=( InputField(name="steering"), InputField(name="camera_delta", required=False), @@ -164,15 +164,15 @@ def test_capabilities_widen_declared_event_types() -> None: assert BROWSER_SOURCE.supports_event_types({"key_down", "prompt_set"}) -def test_semantic_type_mismatch_blocks_capability_match() -> None: +def test_input_modality_mismatch_blocks_capability_match() -> None: source = UserInputSchema( capabilities=( - UserInputCapability(event_type="prompt_set", semantic_type="embedding"), + UserInputCapability(event_type="prompt_set", input_modality="embedding"), ) ) assert not source.supports( - UserInputCapability(event_type="prompt_set", semantic_type="text") + UserInputCapability(event_type="prompt_set", input_modality="text") ) @@ -201,7 +201,9 @@ def test_model_declares_required_and_optional_fields_per_phase() -> None: ("global_conditioning", "prompt"), ("step", "steering"), } - assert {(phase, f.name) for phase, f in optional} == {("step", "camera_delta")} + assert {(phase, f.name) for phase, f in optional} == { + ("step", "camera_delta") + } def test_required_fields_can_be_filtered_by_phase() -> None: @@ -273,13 +275,12 @@ def test_compatible_source_model_and_mapping_can_drive() -> None: ) assert compatibility.can_drive - assert {(p, f.name) for p, f in compatibility.satisfied_required_model_fields} == { - ("global_conditioning", "prompt"), - ("step", "steering"), - } - assert {(p, f.name) for p, f in compatibility.available_optional_model_fields} == { - ("step", "camera_delta") - } + assert { + (p, f.name) for p, f in compatibility.satisfied_required_model_fields + } == {("global_conditioning", "prompt"), ("step", "steering")} + assert { + (p, f.name) for p, f in compatibility.available_optional_model_fields + } == {("step", "camera_delta")} def test_missing_required_model_field_blocks_the_run() -> None: @@ -290,9 +291,9 @@ def test_missing_required_model_field_blocks_the_run() -> None: ) assert not compatibility.can_drive - assert [f.name for _, f in compatibility.missing_required_model_fields] == [ - "steering" - ] + assert [ + f.name for _, f in compatibility.missing_required_model_fields + ] == ["steering"] def test_missing_source_capability_is_reported_when_it_blocks() -> None: @@ -358,7 +359,7 @@ def test_global_conditioning_mapping_matches_global_conditioning_field() -> None assert compatibility.can_drive -def test_unspecified_semantic_type_stays_permissive() -> None: +def test_unspecified_input_modality_stays_permissive() -> None: model = InferenceInputSchema( global_conditioning_fields=(InputField(name="prompt"),) ) @@ -415,7 +416,9 @@ def test_combining_mappings_unions_their_surfaces() -> None: combined = combine_mapping_schemas((PROMPT_MAPPING, STEERING_MAPPING)) assert {m.name for m in combined.consumes} == {"driver_command"} - assert [f.name for f in combined.produces_global_conditioning] == ["prompt"] + assert [f.name for f in combined.produces_global_conditioning] == [ + "prompt" + ] assert [f.name for f in combined.produces_step] == ["steering"] From 33a84fe43d6e12829ae0a4733aa07bf12d5d2ffe Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Wed, 5 Aug 2026 20:06:02 -0700 Subject: [PATCH 06/19] Correct disagreements in docs from PR #413 --- docs/inference_runtime_api_design.md | 44 +++++++++++++++------------- 1 file changed, 24 insertions(+), 20 deletions(-) diff --git a/docs/inference_runtime_api_design.md b/docs/inference_runtime_api_design.md index b6314a205..d1365d77e 100644 --- a/docs/inference_runtime_api_design.md +++ b/docs/inference_runtime_api_design.md @@ -319,13 +319,16 @@ UserInputs -> CanonicalInputs -> InferenceInput raw canonicalized encoded ``` -Raw device events are canonicalized into device-independent modalities before an -application sees them, so adding a keyboard, gamepad, or wheel is a converter -registration rather than an application change. `InferenceInput` is what an -`InferenceSession` actually receives. +Raw device events for live control are canonicalized into device-independent +modalities before application or mapping logic consumes them, so adding a +keyboard, gamepad, or wheel is a converter registration rather than an +application change. Global conditioning is application-owned and reaches +`InferenceInput` directly; it does not pass through live device canonicalization. +`InferenceInput` is what an `InferenceSession` actually receives. -`InferenceInput` describes the data the model or inference pipeline actually -requires. Both it and `CanonicalInputs` distinguish two conditioning slots: +`CanonicalInputs` describes device-independent live control for one requested +input window. `InferenceInput` describes the data the model or inference +pipeline actually requires, split into two conditioning slots: - global conditioning: values that condition the whole rollout; - per-step conditioning: values needed for one generated chunk or frame window. @@ -348,11 +351,11 @@ Inference input payloads should use semantic names, not only modality names. For example, a first frame and an HD map frame should be distinct inputs even if both are image-like values. -Model input names, payload kinds, semantic-type hints, and schema metadata -should be open-ended. Supported integrations such as SANA-WM, LingBot, -Omnidreams, and future external adapters may need different semantic fields. -Adding a new model should usually mean adding adapter-owned schema declarations -and mappings, not changing a central FlashDreams enum. +Model input names, input modalities, and schema metadata should be open-ended. +Supported integrations such as SANA-WM, LingBot, Omnidreams, and future +external adapters may need different semantic fields. Adding a new model should +usually mean adding adapter-owned schema declarations and mappings, not changing +a central FlashDreams enum. Consumption cadence is a separate hint from input scope. A field may be provided through global conditioning because it is session-global state, while @@ -360,9 +363,10 @@ the adapter consumes or slices it during every step. That can be recorded as `frequency_consumed` metadata without changing whether the field belongs in `global_conditioning_fields` or `step_fields`. -For interactive runs, most `InferenceInput` values will be global conditioning -plus per-step inputs produced by input mapping. For MP4 generation and benchmarking, the API -should also support fixed per-step model inputs so runs can be deterministic. +For interactive runs, most `InferenceInput` values will be app-owned global +conditioning plus per-step inputs produced by input mapping. For MP4 generation +and benchmarking, the API should also support fixed per-step model inputs so +runs can be deterministic. ## Schemas @@ -385,8 +389,8 @@ Schema objects may carry open-ended metadata for query-time hints such as coordinate frame, units, rough shape summary, accepted file suffixes, schema URI, model family, or source/transport details. Metadata should help humans and adapter selection code, but compatibility should still be based on the declared -event capabilities, semantic model fields, payload representation hints, and -schema phases. Consumption-cadence hints are descriptive and adapter-owned. +event capabilities, semantic model fields, input modalities, and schema phases. +Consumption-cadence hints are descriptive and adapter-owned. For simple CLI text-to-video or image-to-video runs, `UserInputSchema` can be trivial or omitted because there may be no live controls. `InferenceInputSchema` is @@ -468,10 +472,10 @@ There are two separate moments to keep clear: - before runtime initialization, FlashDreams should select the mapping or mapper set and check obvious compatibility between the app event source and the model; -- during the standard loop, the runtime or runner queues and timestamps user - events, then uses the selected mapping to build initial or per-step - `InferenceInput` from the relevant event window, often after the session reports - what it needs next. +- during the standard loop, the runtime or runner passes app-owned global + `InferenceInput` through the selected mapping before session start, then + queues and timestamps user events, canonicalizes the session-requested window, + and uses the selected mapping to build per-step `InferenceInput`. This keeps the Reactor-style contract intact: the model-side integration can declare user inputs, declare model inputs, and provide a default mapping, while From 1b7bac66120300efc8978992f717ced6edd39a1b Mon Sep 17 00:00:00 2001 From: Aidan Foster Date: Wed, 5 Aug 2026 21:32:39 -0700 Subject: [PATCH 07/19] Fix input canonicalization in tests --- .../tests/test_inference_runtime_api.py | 109 +++++++++++++++++- 1 file changed, 104 insertions(+), 5 deletions(-) diff --git a/flashdreams/tests/test_inference_runtime_api.py b/flashdreams/tests/test_inference_runtime_api.py index d454f6205..890d7efb9 100644 --- a/flashdreams/tests/test_inference_runtime_api.py +++ b/flashdreams/tests/test_inference_runtime_api.py @@ -3,6 +3,7 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import fields from typing import Any, cast @@ -11,6 +12,8 @@ from flashdreams.runtime import ( CanonicalInputs, CanonicalInputSchema, + CanonicalModality, + DeviceConverterSchema, IdentityInputMapping, InferenceConfig, InferenceInput, @@ -325,6 +328,40 @@ def test_reference_loop_validates_mapping_before_runtime_creation() -> None: assert adapter.created_runtime_after_validate +def test_reference_loop_does_not_canonicalize_global_conditioning() -> None: + mapping = _CanonicalRecordingMapping() + adapter = _FakeAdapter() + canonicalizer = InputCanonicalizer([_CountingDeviceConverter()]) + source_schema = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="stateful_event", + payload_fields=frozenset(), + ), + ) + ) + + _drive_two_step_session( + adapter=adapter, + config=InferenceConfig(model_id="fake-model"), + mapping=mapping, + canonicalizer=canonicalizer, + source_schema=source_schema, + user_inputs=UserInputs( + events=(UserInputEvent(timestamp_s=0.75, event_type="stateful_event"),) + ), + inference_input=InferenceInput(global_conditioning={"prompt": "drive forward"}), + output=NullOutputTarget(), + metrics=InMemoryMetricsRecorder(), + ) + + assert mapping.global_canonical_values == {} + assert mapping.step_canonical_values == ( + {"stateful_counter": {"count": 0}}, + {"stateful_counter": {"count": 1}}, + ) + + def test_reference_loop_closes_runtime_when_session_start_fails() -> None: adapter = _FailingStartAdapter() output = NullOutputTarget() @@ -367,12 +404,9 @@ def _drive_two_step_session( canonical_schema=adapter.canonical_input_schema, inference_input_schema=adapter.inference_input_schema, ) + canonicalizer.reset() initial_inputs = mapping.map_global_conditioning_inputs( - canonical_inputs=canonicalizer.canonicalize( - user_inputs, - window=TimeWindow(start_s=0.0, end_s=_SESSION_HORIZON_S), - source_schema=source_schema, - ), + canonical_inputs=CanonicalInputs(), inference_input=inference_input, ) runtime = adapter.create_runtime(config) @@ -510,6 +544,71 @@ def validate( self.validated = True +class _CanonicalRecordingMapping(IdentityInputMapping): + def __init__(self) -> None: + self.global_canonical_values: Mapping[str, Any] | None = None + self._step_canonical_values: list[Mapping[str, Any]] = [] + + @property + def step_canonical_values(self) -> tuple[Mapping[str, Any], ...]: + return tuple(self._step_canonical_values) + + def map_global_conditioning_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + ) -> InferenceInput: + self.global_canonical_values = canonical_inputs.values + return super().map_global_conditioning_inputs( + canonical_inputs=canonical_inputs, + inference_input=inference_input, + ) + + def map_step_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + request: StepRequest, + ) -> InferenceInput: + self._step_canonical_values.append(canonical_inputs.values) + return super().map_step_inputs( + canonical_inputs=canonical_inputs, + inference_input=inference_input, + request=request, + ) + + +_STATEFUL_COUNTER = CanonicalModality( + name="stateful_counter", + payload_fields=frozenset({"count"}), +) + + +class _CountingDeviceConverter: + schema = DeviceConverterSchema( + name="stateful-counter", + produces=_STATEFUL_COUNTER, + consumes=(UserInputCapability(event_type="stateful_event"),), + ) + + def __init__(self) -> None: + self.count = 0 + + def reset(self) -> None: + self.count = 0 + + def convert( + self, + user_inputs: UserInputs, + window: TimeWindow, + ) -> Mapping[str, Any] | None: + del window + self.count += len(user_inputs.events) + return _STATEFUL_COUNTER.value({"count": self.count}) + + class _OrderCheckingAdapter(_FakeAdapter): canonical_input_schema = CanonicalInputSchema() From c78920d4063e4abf505afbbdf650698063234402 Mon Sep 17 00:00:00 2001 From: jarcherNV Date: Thu, 6 Aug 2026 02:40:21 -0700 Subject: [PATCH 08/19] Add experimental inference runtime and shared demo API (#422) Introduce the experimental runtime/session/input envelopes and a shared demo-level API for replay and WebRTC flows. Add the shared runner, output target plumbing, fake-model coverage, and benchmark hooks. Port OmniDreams replay and WebRTC onto the shared demo path via a thin model-owned adapter, add local/remote validation docs, and update the migration plan to track remaining output/stat work and legacy demo cleanup. --- .../omnidreams_demo_replay_benchmarks.json | 88 +++ docs/inference_runtime_api_design.md | 64 +- ...inference_runtime_inputs_implementation.md | 13 +- .../developer_guides/local_benchmarks.rst | 21 + flashdreams/flashdreams/runtime/__init__.py | 4 + .../flashdreams/runtime/demo/__init__.py | 31 + flashdreams/flashdreams/runtime/demo/app.py | 36 + .../flashdreams/runtime/demo/outputs.py | 45 ++ .../flashdreams/runtime/demo/replay.py | 93 +++ flashdreams/flashdreams/runtime/demo/spec.py | 176 +++++ .../flashdreams/runtime/demo/webrtc.py | 274 +++++++ flashdreams/flashdreams/runtime/mapping.py | 3 +- flashdreams/flashdreams/runtime/runner.py | 199 +++++ .../flashdreams/runtime/video_output.py | 157 ++++ flashdreams/tests/test_benchmark_harness.py | 44 ++ .../tests/test_inference_runtime_api.py | 391 ---------- flashdreams/tests/test_runtime_demo_api.py | 458 +++++++++++ .../tests/test_runtime_input_mapping.py | 30 +- flashdreams/tests/test_runtime_runner.py | 660 ++++++++++++++++ .../tests/test_runtime_video_output.py | 91 +++ .../omnidreams/omnidreams/demo/README.md | 66 ++ .../omnidreams/omnidreams/demo/__init__.py | 20 + .../omnidreams/omnidreams/demo/adapter.py | 279 +++++++ .../omnidreams/omnidreams/demo/cli.py | 187 +++++ .../omnidreams/omnidreams/demo/replay.py | 277 +++++++ .../omnidreams/omnidreams/demo/spec.py | 271 +++++++ .../omnidreams/omnidreams/demo/webrtc.py | 178 +++++ integrations/omnidreams/pyproject.toml | 7 +- .../omnidreams/tests/test_demo_api.py | 577 ++++++++++++++ uv.lock | 720 +----------------- 30 files changed, 4313 insertions(+), 1147 deletions(-) create mode 100644 configs/omnidreams_demo_replay_benchmarks.json create mode 100644 flashdreams/flashdreams/runtime/demo/__init__.py create mode 100644 flashdreams/flashdreams/runtime/demo/app.py create mode 100644 flashdreams/flashdreams/runtime/demo/outputs.py create mode 100644 flashdreams/flashdreams/runtime/demo/replay.py create mode 100644 flashdreams/flashdreams/runtime/demo/spec.py create mode 100644 flashdreams/flashdreams/runtime/demo/webrtc.py create mode 100644 flashdreams/flashdreams/runtime/runner.py create mode 100644 flashdreams/flashdreams/runtime/video_output.py create mode 100644 flashdreams/tests/test_runtime_demo_api.py create mode 100644 flashdreams/tests/test_runtime_runner.py create mode 100644 flashdreams/tests/test_runtime_video_output.py create mode 100644 integrations/omnidreams/omnidreams/demo/README.md create mode 100644 integrations/omnidreams/omnidreams/demo/__init__.py create mode 100644 integrations/omnidreams/omnidreams/demo/adapter.py create mode 100644 integrations/omnidreams/omnidreams/demo/cli.py create mode 100644 integrations/omnidreams/omnidreams/demo/replay.py create mode 100644 integrations/omnidreams/omnidreams/demo/spec.py create mode 100644 integrations/omnidreams/omnidreams/demo/webrtc.py create mode 100644 integrations/omnidreams/tests/test_demo_api.py diff --git a/configs/omnidreams_demo_replay_benchmarks.json b/configs/omnidreams_demo_replay_benchmarks.json new file mode 100644 index 000000000..3ac6e4019 --- /dev/null +++ b/configs/omnidreams_demo_replay_benchmarks.json @@ -0,0 +1,88 @@ +{ + "schema_version": 1, + "description": "Manual one-minute local benchmark scenarios for comparing the legacy Omnidreams single-view runner against the experimental shared demo replay path. The runner writes the legacy stacked HDMap/RGB canvas while the shared demo writes generated RGB output, so use the report for manual MP4 comparison rather than automatic pixel quality scoring.", + "scenarios": [ + { + "id": "omnidreams-sv-runner-baseline", + "name": "Omnidreams single-view runner baseline", + "description": "Runs the stable legacy Omnidreams single-view runner with the bundled example data for the same one-minute block count used by the shipped Omnidreams baseline.", + "report_group": { + "id": "omnidreams-demo", + "name": "Omnidreams Demo Comparison" + }, + "tags": [ + "manual", + "gpu", + "real-demo", + "omnidreams", + "i2v", + "replay", + "baseline" + ], + "env": { + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True" + }, + "command": [ + "uv", + "run", + "--project", + "integrations/omnidreams", + "flashdreams-run", + "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae", + "--example-data", + "True", + "--example-data-uuid", + "239560dc-33d1-11ef-9720-00044bcbccac", + "--total-blocks", + "226" + ], + "warmup_steps": 1, + "quality_baseline_compare": false, + "timeout_s": 7200 + }, + { + "id": "omnidreams-sv-demo-replay", + "name": "Omnidreams shared demo replay", + "description": "Runs the experimental shared demo API replay path with the same stable non-perf preset, bundled example data, and one-minute block count as the legacy runner.", + "report_group": { + "id": "omnidreams-demo", + "name": "Omnidreams Demo Comparison" + }, + "tags": [ + "manual", + "gpu", + "real-demo", + "omnidreams", + "i2v", + "replay", + "shared-demo" + ], + "env": { + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True" + }, + "command": [ + "uv", + "run", + "--project", + "integrations/omnidreams", + "omnidreams-demo", + "replay", + "--preset-id", + "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae", + "--example-data", + "--example-data-uuid", + "239560dc-33d1-11ef-9720-00044bcbccac", + "--total-blocks", + "226", + "--output", + "{output_dir}/omnidreams-sv-demo-replay.mp4" + ], + "output_dir_arg": null, + "warmup_steps": 1, + "quality_baseline_compare": false, + "timeout_s": 7200 + } + ] +} diff --git a/docs/inference_runtime_api_design.md b/docs/inference_runtime_api_design.md index d1365d77e..6bd1018d9 100644 --- a/docs/inference_runtime_api_design.md +++ b/docs/inference_runtime_api_design.md @@ -44,22 +44,24 @@ model. ## Current Implementation Plan Implementation should happen on an experimental integration branch. PRs for this -work should target that branch until the API shape, LingBot migration, and -OmniDreams migration are all working well enough to merge to `main` together. +work should target that branch until the API shape and OmniDreams migration are +working well enough to merge to `main` together. LingBot migration is deferred +to a separate follow-up after the OmniDreams path has clarified the shared demo +API shape. The experimental branch can temporarily break or simplify command-line options -while the demos are being moved to the new API. The required outcome is that the -LingBot and OmniDreams demos still run through the new runtime path, and that -benchmark tooling can confirm they are at least broadly healthy before the -branch is merged back to `main`. +while the demos are being moved to the new API. The required outcome for this +branch is that the OmniDreams demo runs through the new shared demo/runtime path, +and that benchmark and manual WebRTC checks can confirm it is at least broadly +healthy before the branch is merged back to `main`. Initial scope: - define the minimal runtime API envelope; -- migrate LingBot and OmniDreams to use it; +- migrate OmniDreams to use it through a shared demo-level API; - support selectable output modes such as MP4, JPEG/MJPEG stream, WebRTC, and headless/null where appropriate; -- use or update benchmark tooling to verify the migrated demos; +- use or update benchmark tooling to verify the migrated OmniDreams demo; - defer broader model migrations, hosted execution, full autotune, and polished metrics until the first branch proves the API shape. @@ -71,14 +73,30 @@ Initial scope: | T1 | Complete | Minimal API envelope and naming. | Partly. | T0. | `InferenceConfig`, `UserInputs`, `InferenceInput`, runtime/session, output target, and mapping boundaries are defined well enough for demos to use. | | T2 | Complete | Event-based `UserInputs`. | Yes, after T1 direction is agreed. | T1. | User inputs are primarily timestamped events; replay traces and derived snapshots are supported where needed. | | T3 | Complete | `CanonicalInputs`, `InferenceInput`, schemas, and mapping boundary. | Yes, after T1 direction is agreed. | T1. | Models can declare required global/per-step inputs, and mappings can convert canonical inputs into inference inputs. | -| T4 | Planned | `ModelRunner`, `InferenceRuntime`, and `InferenceSession` skeleton. | Partly. | T1. | A minimal standard loop can initialize a runtime, run at least one sequential session, and close cleanly. | +| T4 | Complete | `ModelRunner`, `InferenceRuntime`, and `InferenceSession` skeleton. | Partly. | T1. | A minimal standard loop can initialize a runtime, run at least one sequential session, and close cleanly. | | T5 | Planned | Output mode selection. | Yes, after the result/output shape is agreed. | T1, T4. | A run can choose output behavior such as MP4, JPEG/MJPEG stream, WebRTC, benchmark artifact, or headless/null without changing model code. | -| T6 | Planned | LingBot migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | LingBot runs through the new API path with its event inputs mapped into model inputs. | -| T7 | Planned | OmniDreams migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | OmniDreams runs through the new API path with its model-specific inputs and mapping preserved. | -| T8 | Planned | Benchmark/smoke verification for LingBot and OmniDreams. | Preparation can run early; final gate is late. | T5, T6, T7. | Existing or updated benchmark tooling can run both migrated demos and produce enough evidence that they still work. | +| T6 | Deferred | LingBot migration. | Yes, but out of scope for this branch. | T2, T3, T4. | LingBot runs through the new API path with its event inputs mapped into model inputs. | +| T7 | Partially complete | OmniDreams migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | OmniDreams replay and WebRTC run through the shared demo API path; remaining work is output/stat integration, legacy demo retirement, and cleanup. | +| T8 | Partially complete | Benchmark/smoke verification for OmniDreams. | Preparation can run early; final gate is late. | T5, T7. | Existing or updated benchmark tooling can run the migrated OmniDreams demo and produce enough evidence that it still works. | | T9 | Planned | Metrics and profiling normalization for the branch. | Yes, but final integration is late. | T4, T5, T8. | Basic canonical metrics are emitted for migrated demos; deeper metrics can remain follow-up work. | -| T10 | Planned | CLI compatibility and migration cleanup. | Yes, after demo migrations start. | T6, T7. | Required demo commands are restored or replaced, temporary hacks are removed, and user-facing docs/notes match the branch behavior. | -| T11 | Planned | Stabilize and merge experimental branch to `main`. | No, final integration step. | T6-T10. | LingBot and OmniDreams pass agreed smoke/benchmark checks, review feedback is addressed, and the branch can merge as one API transition. | +| T10 | Planned | CLI compatibility, legacy retirement, and migration cleanup. | Yes, after demo migrations start. | T5, T7, T8. | Required demo commands are restored or replaced, old interactive-drive and old OmniDreams demo/server paths are removed or reduced to compatibility shims, code used only by retired demos is removed, and user-facing docs/notes match the branch behavior. | +| T11 | Planned | Stabilize and merge experimental branch to `main`. | No, final integration step. | T5, T7-T10. | OmniDreams passes agreed smoke/benchmark checks, review feedback is addressed, and the branch can merge as one API transition. | + +Current OmniDreams migration status: + +- The shared `flashdreams.runtime.demo` API and OmniDreams demo adapter exist. +- OmniDreams MP4 replay runs through the shared replay runner and MP4 output + target. +- The one-minute benchmark comparison can run the legacy replay path and the new + shared demo replay path side by side. +- OmniDreams WebRTC runs through `serve_flashdreams_demo(...)` and the shared + WebRTC manager path while still using the existing OmniDreams runtime and + packaged browser app. +- The migration is not complete until the new output target/stat artifact work + lands, the new OmniDreams path is updated to use it, the old interactive-drive + and old OmniDreams demo/server paths are removed or reduced to deliberate + compatibility shims, code used only by retired demos is deleted, and the + experimental demo/runtime/input code is cleaned up. Suggested parallel split: @@ -88,7 +106,8 @@ Suggested parallel split: stay coherent; - one person owns T5/T8/T9, because outputs, benchmarks, and metrics are tightly related; -- LingBot and OmniDreams can be assigned separately once the skeleton is usable; +- LingBot should be tracked as a separate follow-up once OmniDreams has settled + the shared demo API shape; - one person should track branch health, CLI compatibility, and merge readiness. ## Architecture @@ -510,6 +529,11 @@ adapter/runtime still owns deep tensor validation and model semantics. The standard loop should be shared by CLI generation, headless playback, MP4 generation, benchmarks, and simple realtime applications. +The current v0 production loop is `flashdreams.runtime.run_inference_session()`. +It is intentionally narrow: one adapter, one config, one canonicalizer/source, +one selected mapping, one initial input, one output target, one metrics +recorder, and one synchronous sequential session. + A run should: 1. Discover the model or preset without loading checkpoints. @@ -645,9 +669,11 @@ The new API should reuse existing code instead of replacing everything: The task tracker near the start of this document is the source of truth for the first implementation branch. The first milestone is intentionally narrower than -the full design: prove the API with LingBot and OmniDreams, selectable output -modes, and enough benchmark/smoke coverage to merge the experimental branch -back to `main` safely. +the full design: prove the API with OmniDreams, add shared output/stat artifact +selection, retire the old OmniDreams demo paths, clean up the experimental +runtime/demo code, and collect enough benchmark/smoke evidence to merge the +experimental branch back to `main` safely. LingBot should be handled in a +separate follow-up plan. ## Design Risks @@ -709,7 +735,7 @@ registry, standard loop, concrete output modes, or model migrations: registering it? - What package registration mechanism should third-party and internal adapters use for CLI discovery and benchmarks? -- What is the first public model to migrate? +- Which model should migrate after OmniDreams settles the shared demo API shape? - What metrics are required for every benchmark run? - What metadata must be discoverable without loading checkpoints? - What requirements do Dynamo/Reactor-style backends need before we commit to diff --git a/docs/inference_runtime_inputs_implementation.md b/docs/inference_runtime_inputs_implementation.md index 763b9810c..75460cced 100644 --- a/docs/inference_runtime_inputs_implementation.md +++ b/docs/inference_runtime_inputs_implementation.md @@ -18,8 +18,9 @@ Implementation lives in `flashdreams.runtime`: and compatibility - `flashdreams/tests/test_runtime_canonical.py` - `flashdreams/tests/test_runtime_input_mapping.py` -- `flashdreams/tests/test_inference_runtime_api.py` — the T1 envelope tests, - including a reference loop that exercises all three layers +- `flashdreams/tests/test_inference_runtime_api.py` — the T1 envelope tests +- `flashdreams/tests/test_runtime_runner.py` — the production standard loop + tests that exercise all three input layers with runtime/session cleanup The supported-model input inventory that informed this work is in `docs/inference_runtime_supported_inputs_inventory.md`. @@ -269,8 +270,9 @@ layer: to `OutputTarget.write()`. Output shape is T5. - **Declared output modalities**, so an output target or quality-eval can state what it requires and be matched the way inputs now are. T5/T8. -- **`Application`**, the class that has-a input system, input map, global - conditioning, session, and output target. T4. +- **Full `Application` ownership**, the class that has-a input system, input + map, global conditioning, session, and output target. T4 now provides the + narrow synchronous runner; richer application ownership remains outside T4. - **Loop ownership** — whether the application or the runtime/session drives the main event loop, and whether inputs are queued and batched. @@ -279,7 +281,8 @@ layer: ```bash .venv/bin/pytest flashdreams/tests/test_runtime_canonical.py \ flashdreams/tests/test_runtime_input_mapping.py \ - flashdreams/tests/test_inference_runtime_api.py -q + flashdreams/tests/test_inference_runtime_api.py \ + flashdreams/tests/test_runtime_runner.py -q .venv/bin/ty check flashdreams/flashdreams/runtime ``` diff --git a/docs/source/developer_guides/local_benchmarks.rst b/docs/source/developer_guides/local_benchmarks.rst index 687d95aa2..62842ba32 100644 --- a/docs/source/developer_guides/local_benchmarks.rst +++ b/docs/source/developer_guides/local_benchmarks.rst @@ -132,6 +132,27 @@ input stream is shorter than the requested duration. ``interactive-drive`` is left out of this shipped MP4 suite for now because its public CLI is a live presenter rather than a file-writing runner. +Omnidreams Shared Demo Comparison +--------------------------------- + +``configs/omnidreams_demo_replay_benchmarks.json`` contains a one-minute manual +comparison between the legacy Omnidreams single-view runner and the experimental +shared demo replay path: + +.. code-block:: bash + + uv run flashdreams-benchmark \ + --scenario-file configs/omnidreams_demo_replay_benchmarks.json \ + --scenario omnidreams-sv-runner-baseline \ + --scenario omnidreams-sv-demo-replay \ + --output-dir artifacts/benchmarks/omnidreams-demo-replay-compare + +Use the generated report's MP4 links for side-by-side manual review. The legacy +runner writes the stacked HDMap/RGB canvas while the shared demo writes generated +RGB output, so this comparison intentionally disables automatic baseline quality +scoring until those output layouts are aligned. Both scenarios use ``226`` +blocks, matching the shipped Omnidreams one-minute baseline. + Quality Hooks ------------- diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py index 04f84ae54..5f89196ba 100644 --- a/flashdreams/flashdreams/runtime/__init__.py +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -56,7 +56,9 @@ RuntimeMetricSample, ) from flashdreams.runtime.output import NullOutputTarget, OutputArtifact, OutputTarget +from flashdreams.runtime.runner import run_inference_session from flashdreams.runtime.types import StepRequest, StepResult +from flashdreams.runtime.video_output import Mp4VideoOutputTarget __all__ = [ "CanonicalInputs", @@ -88,6 +90,7 @@ "MappingCompatibility", "MetricsRecorder", "ModelAdapter", + "Mp4VideoOutputTarget", "NullMetricsRecorder", "NullOutputTarget", "OutputArtifact", @@ -98,6 +101,7 @@ "StepRequest", "StepResult", "TimeWindow", + "run_inference_session", "undeclared_inference_inputs", "UserInputCapability", "UserInputEvent", diff --git a/flashdreams/flashdreams/runtime/demo/__init__.py b/flashdreams/flashdreams/runtime/demo/__init__.py new file mode 100644 index 000000000..3d9d99919 --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/__init__.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""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 ( + DemoAdapter, + DemoSpec, + Mp4OutputSpec, + NullOutputSpec, + OutputSpec, + PreparedScenario, + WebRTCOutputSpec, +) + +__all__ = [ + "DemoAdapter", + "DemoSpec", + "Mp4OutputSpec", + "NullOutputSpec", + "OutputSpec", + "PreparedScenario", + "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 new file mode 100644 index 000000000..7659859b4 --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/app.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Experimental shared demo entrypoints.""" + +from __future__ import annotations + +from typing import Any + +from .replay import run_replay_demo +from .spec import DemoAdapter, DemoSpec + + +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) + + +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 + + return serve_webrtc_demo(spec=spec, adapter=adapter, **kwargs) + + +__all__ = ["run_flashdreams_demo", "serve_flashdreams_demo"] diff --git a/flashdreams/flashdreams/runtime/demo/outputs.py b/flashdreams/flashdreams/runtime/demo/outputs.py new file mode 100644 index 000000000..421ec3bb4 --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/outputs.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared demo output-target construction.""" + +from __future__ import annotations + +from pathlib import Path + +from flashdreams.runtime.output import NullOutputTarget, OutputTarget +from flashdreams.runtime.video_output import Mp4VideoOutputTarget, VideoWriter + +from .spec import Mp4OutputSpec, NullOutputSpec, OutputSpec, WebRTCOutputSpec + + +def build_output_target( + output: OutputSpec, + *, + mp4_writer: VideoWriter | None = None, +) -> OutputTarget: + """Build a replay output target from a demo output spec.""" + if isinstance(output, NullOutputSpec): + return NullOutputTarget(store_results=output.store_results) + if isinstance(output, Mp4OutputSpec): + output_path = Path(output.path) + if mp4_writer is not None: + return Mp4VideoOutputTarget( + output_path=output_path, + fps=output.fps, + output_layout=output.output_layout, + writer=mp4_writer, + move_to_cpu=output.move_to_cpu, + ) + return Mp4VideoOutputTarget( + output_path=output_path, + fps=output.fps, + output_layout=output.output_layout, + move_to_cpu=output.move_to_cpu, + ) + if isinstance(output, WebRTCOutputSpec): + raise ValueError("WebRTC output does not create a replay OutputTarget.") + raise TypeError(f"Unsupported demo output spec: {type(output).__name__}.") + + +__all__ = ["build_output_target"] diff --git a/flashdreams/flashdreams/runtime/demo/replay.py b/flashdreams/flashdreams/runtime/demo/replay.py new file mode 100644 index 000000000..18b873254 --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/replay.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared replay demo runner.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence + +from flashdreams.runtime.metrics import MetricsRecorder, NullMetricsRecorder +from flashdreams.runtime.output import OutputArtifact, OutputTarget +from flashdreams.runtime.runner import run_inference_session + +from .outputs import build_output_target +from .spec import DemoAdapter, DemoSpec, OutputSpec, WebRTCOutputSpec + +OutputTargetFactory = Callable[[OutputSpec], OutputTarget] +InferenceSessionRunner = Callable[..., Sequence[OutputArtifact]] + + +def run_replay_demo( + *, + spec: DemoSpec, + adapter: DemoAdapter, + output_target_factory: OutputTargetFactory = build_output_target, + metrics: MetricsRecorder | None = None, + runner: InferenceSessionRunner = run_inference_session, +) -> tuple[OutputArtifact, ...]: + """Run one prepared demo scenario through the shared runtime runner.""" + _require_supported_mode( + mode=spec.input_mode, + supported=adapter.supported_input_modes(), + label="input_mode", + ) + if spec.input_mode != "replay": + raise ValueError( + "run_replay_demo requires input_mode='replay', " + f"got input_mode={spec.input_mode!r}." + ) + _require_supported_mode( + mode=spec.output.mode, + supported=adapter.supported_output_modes(), + label="output.mode", + ) + if isinstance(spec.output, WebRTCOutputSpec): + raise ValueError("run_replay_demo does not support WebRTC output.") + + prepared = adapter.prepare_scenario(spec) + mapping = prepared.mapping or adapter.default_input_mapping() + if mapping is None: + raise ValueError( + "Demo scenario did not provide an input mapping, and the adapter " + "has no default input mapping." + ) + if spec.config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + + output = output_target_factory(spec.output) + metrics_recorder = metrics or NullMetricsRecorder() + return tuple( + runner( + adapter=adapter, + config=spec.config, + mapping=mapping, + canonicalizer=prepared.canonicalizer, + source_schema=prepared.source_schema, + user_inputs=prepared.user_inputs, + initial_inputs=prepared.initial_inputs, + output=output, + metrics=metrics_recorder, + ) + ) + + +def _require_supported_mode( + *, + mode: str, + supported: tuple[str, ...], + label: str, +) -> None: + if mode in supported: + return + supported_text = ", ".join(repr(each) for each in supported) or "" + raise ValueError( + f"Unsupported demo {label}={mode!r}; supported modes: {supported_text}." + ) + + +__all__ = [ + "InferenceSessionRunner", + "OutputTargetFactory", + "run_replay_demo", +] diff --git a/flashdreams/flashdreams/runtime/demo/spec.py b/flashdreams/flashdreams/runtime/demo/spec.py new file mode 100644 index 000000000..bc2884ab3 --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/spec.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Experimental shared demo API data shapes.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, Literal, Protocol, TypeAlias + +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.canonical import InputCanonicalizer +from flashdreams.runtime.config import InferenceConfig +from flashdreams.runtime.inputs import InferenceInput, UserInputs, UserInputSchema +from flashdreams.runtime.interfaces import ModelAdapter +from flashdreams.runtime.mapping import InputMapping + + +@dataclass(frozen=True, kw_only=True, slots=True) +class NullOutputSpec: + """Headless/null replay output.""" + + mode: Literal["null"] = "null" + store_results: bool = False + + +@dataclass(frozen=True, kw_only=True, slots=True) +class Mp4OutputSpec: + """MP4 replay output.""" + + path: str | Path + fps: int | float + mode: Literal["mp4"] = "mp4" + output_layout: VideoTensorLayout = "bvtchw" + move_to_cpu: bool = True + + def __post_init__(self) -> None: + if float(self.fps) <= 0: + raise ValueError("Mp4OutputSpec.fps must be > 0.") + object.__setattr__(self, "path", Path(self.path)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class WebRTCOutputSpec: + """Shared WebRTC serving output.""" + + mode: Literal["webrtc"] = "webrtc" + host: str = "127.0.0.1" + port: int = 8080 + fps: int = 30 + video_width: int = 1280 + video_height: int = 720 + warmup_chunks: int = 0 + warmup_timeout_s: float = 30.0 + client_liveness_timeout_s: float = 30.0 + web_dir: str | Path | None = None + request_session_path: str = "/request_session" + preload_name: str | None = None + + def __post_init__(self) -> None: + if not self.host.strip(): + raise ValueError("WebRTCOutputSpec.host must be non-empty.") + if not (0 < int(self.port) < 65536): + raise ValueError("WebRTCOutputSpec.port must be between 1 and 65535.") + if self.fps <= 0: + raise ValueError("WebRTCOutputSpec.fps must be > 0.") + if self.video_width <= 0 or self.video_height <= 0: + raise ValueError("WebRTCOutputSpec video dimensions must be > 0.") + if self.warmup_chunks < 0: + raise ValueError("WebRTCOutputSpec.warmup_chunks must be >= 0.") + if self.warmup_timeout_s <= 0: + raise ValueError("WebRTCOutputSpec.warmup_timeout_s must be > 0.") + if self.client_liveness_timeout_s <= 0: + raise ValueError("WebRTCOutputSpec.client_liveness_timeout_s must be > 0.") + if not self.request_session_path.startswith("/"): + raise ValueError( + "WebRTCOutputSpec.request_session_path must start with '/'." + ) + if self.web_dir is not None: + object.__setattr__(self, "web_dir", Path(self.web_dir)) + + +OutputSpec: TypeAlias = NullOutputSpec | Mp4OutputSpec | WebRTCOutputSpec + + +@dataclass(frozen=True, kw_only=True, slots=True) +class DemoSpec: + """User-facing shared demo run description.""" + + __hash__ = None + + model_id: str + input_mode: str + output: OutputSpec + preset_id: str | None = None + scenario: Any | None = None + config: InferenceConfig | None = None + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.model_id.strip(): + raise ValueError("DemoSpec.model_id must be non-empty.") + if not self.input_mode.strip(): + raise ValueError("DemoSpec.input_mode must be non-empty.") + config = self.config + if config is None: + config = InferenceConfig( + model_id=self.model_id, + preset_id=self.preset_id, + ) + else: + if config.model_id != self.model_id: + raise ValueError( + "DemoSpec.model_id must match InferenceConfig.model_id." + ) + if self.preset_id is None: + object.__setattr__(self, "preset_id", config.preset_id) + elif config.preset_id is None: + config = replace(config, preset_id=self.preset_id) + elif config.preset_id != self.preset_id: + raise ValueError( + "DemoSpec.preset_id must match InferenceConfig.preset_id." + ) + object.__setattr__(self, "config", config) + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class PreparedScenario: + """Runtime-ready scenario prepared by a model demo adapter.""" + + __hash__ = None + + initial_inputs: InferenceInput + user_inputs: UserInputs = field(default_factory=UserInputs) + source_schema: UserInputSchema = field(default_factory=UserInputSchema) + canonicalizer: InputCanonicalizer = field(default_factory=InputCanonicalizer) + mapping: InputMapping | None = None + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +class DemoAdapter(ModelAdapter, Protocol): + """Model-owned adapter surface consumed by shared demo launchers.""" + + def supported_input_modes(self) -> tuple[str, ...]: + """Return demo input modes this adapter can prepare.""" + ... + + def supported_output_modes(self) -> tuple[str, ...]: + """Return demo output modes this adapter can run.""" + ... + + 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", + "DemoSpec", + "Mp4OutputSpec", + "NullOutputSpec", + "OutputSpec", + "PreparedScenario", + "WebRTCOutputSpec", +] diff --git a/flashdreams/flashdreams/runtime/demo/webrtc.py b/flashdreams/flashdreams/runtime/demo/webrtc.py new file mode 100644 index 000000000..f93a855db --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/webrtc.py @@ -0,0 +1,274 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared WebRTC demo construction.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from aiohttp import web + +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 + + +CreateWebRTCApp = Callable[..., web.Application] +RunWebRTCServer = Callable[..., None] + + +def build_webrtc_demo( + *, + spec: DemoSpec, + adapter: DemoAdapter, + create_app: bool = False, + 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, + ) + app = ( + _create_app( + spec=spec, + adapter=adapter, + session_manager=manager, + create_app_fn=create_app_fn, + ) + if create_app + else None + ) + return WebRTCDemo( + runtime=runtime, + runtime_config=runtime_config, + session_manager=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, + ) + + +def _create_app( + *, + spec: DemoSpec, + adapter: DemoAdapter, + 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, + session_manager=session_manager, + request_session_url=_request_session_url(output), + ) + return _build_webrtc_app( + output=output, + session_manager=session_manager, + create_app_fn=create_app_fn, + preload_name=output.preload_name or spec.model_id, + ) + + +def _build_webrtc_app( + *, + output: WebRTCOutputSpec, + session_manager: BaseWebRTCSessionManager[Any, Any], + create_app_fn: CreateWebRTCApp, + preload_name: str, +) -> web.Application: + if output.web_dir is None: + raise ValueError("WebRTC app creation requires output.web_dir.") + return create_app_fn( + web_dir=Path(output.web_dir), + session_manager=session_manager, + request_session_url=_request_session_url(output), + preload_name=preload_name, + ) + + +def _request_session_url(output: WebRTCOutputSpec) -> str: + host = "127.0.0.1" if output.host in {"0.0.0.0", "::"} else output.host + return f"http://{host}:{output.port}{output.request_session_path}" + + +__all__ = [ + "CreateWebRTCApp", + "RunWebRTCServer", + "SharedDemoWebRTCSessionManager", + "WebRTCDemo", + "WebRTCDemoRuntimeConfig", + "build_webrtc_demo", + "serve_webrtc_demo", +] diff --git a/flashdreams/flashdreams/runtime/mapping.py b/flashdreams/flashdreams/runtime/mapping.py index 710ad8043..6dfb5cc45 100644 --- a/flashdreams/flashdreams/runtime/mapping.py +++ b/flashdreams/flashdreams/runtime/mapping.py @@ -369,8 +369,7 @@ def undeclared_inference_inputs( for phase in INPUT_PHASES for key in inputs.for_phase(phase) if not any( - declared.name == key - for declared in mapping_schema.produces_for(phase) + declared.name == key for declared in mapping_schema.produces_for(phase) ) ) diff --git a/flashdreams/flashdreams/runtime/runner.py b/flashdreams/flashdreams/runtime/runner.py new file mode 100644 index 000000000..03d814472 --- /dev/null +++ b/flashdreams/flashdreams/runtime/runner.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Minimal synchronous standard runner for the runtime API.""" + +from __future__ import annotations + +import math + +from flashdreams.runtime.canonical import InputCanonicalizer +from flashdreams.runtime.config import InferenceConfig +from flashdreams.runtime.inputs import ( + CanonicalInputs, + CanonicalInputSchema, + InferenceInput, + TimeWindow, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.interfaces import ( + InferenceRuntime, + InferenceSession, + ModelAdapter, +) +from flashdreams.runtime.mapping import ( + DeclaresMappingSchema, + InputMapping, + check_mapping_compatibility, +) +from flashdreams.runtime.metrics import MetricsRecorder +from flashdreams.runtime.output import OutputArtifact, OutputTarget +from flashdreams.runtime.types import StepResult + +_DEFAULT_SESSION_HORIZON_S = 3600.0 + + +def run_inference_session( + *, + adapter: ModelAdapter, + config: InferenceConfig, + mapping: InputMapping, + canonicalizer: InputCanonicalizer, + source_schema: UserInputSchema, + user_inputs: UserInputs, + initial_inputs: InferenceInput, + output: OutputTarget, + metrics: MetricsRecorder, +) -> tuple[OutputArtifact, ...]: + """Run one sequential inference session through the standard loop. + + This v0 loop intentionally handles one adapter/runtime/session, one selected + input mapping, one replay/live input batch, one output target, and one + metrics recorder. It is synchronous and owns only orchestration. + """ + + runtime: InferenceRuntime | None = None + session: InferenceSession | None = None + output_opened = False + output_artifacts: tuple[OutputArtifact, ...] = () + primary_error: BaseException | None = None + + try: + adapter.validate_config(config) + canonical_schema = canonicalizer.canonical_schema(source_schema) + _check_declared_mapping_compatibility( + mapping=mapping, + canonical_schema=canonical_schema, + adapter=adapter, + ) + mapping.validate( + canonical_schema=canonical_schema, + inference_input_schema=adapter.inference_input_schema, + ) + canonicalizer.reset() + mapped_initial_inputs = mapping.map_global_conditioning_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=initial_inputs, + ) + runtime = adapter.create_runtime(config) + session = runtime.start_session(mapped_initial_inputs) + output.open() + output_opened = True + step_base_inputs = InferenceInput( + step=initial_inputs.step, + metadata=initial_inputs.metadata, + ) + + while (request := session.next_step_request()) is not None: + step_inputs = mapping.map_step_inputs( + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=request.user_input_window + or _all_user_inputs_window(user_inputs), + source_schema=source_schema, + ), + inference_input=step_base_inputs, + request=request, + ) + result = session.step(step_inputs) + output.write(result) + _record_timing_metrics(metrics, result) + except BaseException as exc: + primary_error = exc + raise + finally: + cleanup_error, output_artifacts = _close_run_resources( + output=output if output_opened else None, + session=session, + runtime=runtime, + metrics=metrics, + ) + if cleanup_error is not None and primary_error is None: + raise cleanup_error + + return output_artifacts + + +def _check_declared_mapping_compatibility( + *, + mapping: InputMapping, + canonical_schema: CanonicalInputSchema, + adapter: ModelAdapter, +) -> None: + if not isinstance(mapping, DeclaresMappingSchema): + return + compatibility = check_mapping_compatibility( + canonical_schema=canonical_schema, + inference_input_schema=adapter.inference_input_schema, + mapping_schema=mapping.mapping_schema, + ) + compatibility.raise_if_incompatible() + + +def _all_user_inputs_window(user_inputs: UserInputs) -> TimeWindow: + if not user_inputs.events: + return TimeWindow(start_s=0.0, end_s=_DEFAULT_SESSION_HORIZON_S) + return TimeWindow( + start_s=0.0, + end_s=max( + _DEFAULT_SESSION_HORIZON_S, + math.nextafter(user_inputs.events[-1].timestamp_s, math.inf), + ), + ) + + +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 + sample_name = name[:-2] or name + metrics.record_timing( + sample_name, + float(value), + step_index=result.step_index, + ) + + +def _close_run_resources( + *, + output: OutputTarget | None, + session: InferenceSession | None, + runtime: InferenceRuntime | None, + metrics: MetricsRecorder, +) -> tuple[BaseException | None, tuple[OutputArtifact, ...]]: + cleanup_error: BaseException | None = None + artifacts: tuple[OutputArtifact, ...] = () + + def remember_error(exc: BaseException) -> None: + nonlocal cleanup_error + if cleanup_error is None: + cleanup_error = exc + + if output is not None: + try: + artifacts = tuple(output.close()) + except BaseException as exc: + remember_error(exc) + + if session is not None: + try: + session.close() + except BaseException as exc: + remember_error(exc) + + if runtime is not None: + try: + runtime.close() + except BaseException as exc: + remember_error(exc) + + try: + metrics.close() + except BaseException as exc: + remember_error(exc) + + return cleanup_error, artifacts + + +__all__ = ["run_inference_session"] diff --git a/flashdreams/flashdreams/runtime/video_output.py b/flashdreams/flashdreams/runtime/video_output.py new file mode 100644 index 000000000..b5c372125 --- /dev/null +++ b/flashdreams/flashdreams/runtime/video_output.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Video output targets for the runtime API.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import cast + +import torch + +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.infra.runner_io import ( + DEFAULT_RUNNER_INSTALL_HINT, + write_video_tensor, +) +from flashdreams.infra.runner_io import ( + VideoTensorLayout as WritableVideoTensorLayout, +) +from flashdreams.infra.video_output import RunnerVideoOutputStream, VideoStepResult +from flashdreams.runtime.output import OutputArtifact +from flashdreams.runtime.types import StepResult + +VideoWriter = Callable[..., Path] + + +@dataclass(slots=True) +class Mp4VideoOutputTarget: + """Write runtime ``VideoStepResult`` chunks to one MP4 artifact.""" + + output_path: Path + fps: int | float + output_layout: VideoTensorLayout = "bvtchw" + writer: VideoWriter = field(default=write_video_tensor, repr=False) + install_hint: str = DEFAULT_RUNNER_INSTALL_HINT + move_to_cpu: bool = True + _opened: bool = field(default=False, init=False, repr=False) + _stream: RunnerVideoOutputStream | None = field( + default=None, + init=False, + repr=False, + ) + + @property + def closed(self) -> bool: + return not self._opened + + def open(self) -> None: + self._stream = RunnerVideoOutputStream( + postprocess_stream=None, + output_layout=self.output_layout, + collect_output=True, + 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: + raise RuntimeError("Cannot write to a closed output target.") + video_result = result.output + if not isinstance(video_result, VideoStepResult): + raise TypeError( + "Mp4VideoOutputTarget requires StepResult.output to be " + f"VideoStepResult, got {type(video_result).__name__}." + ) + if video_result.layout != self.output_layout: + raise ValueError( + "Mp4VideoOutputTarget received layout " + f"{video_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, + ) + + def close(self) -> Sequence[OutputArtifact]: + if self._stream is None: + self._opened = False + return () + + stream = self._stream + self._stream = None + self._opened = False + video = stream.finish() + if video is None: + return () + + 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=writable_layout, + install_hint=self.install_hint, + ) + return ( + OutputArtifact( + kind="video/mp4", + uri=str(path), + metadata={ + "fps": self.fps, + "source_layout": self.output_layout, + "write_layout": writable_layout, + "shape": tuple(int(dim) for dim in writable_video.shape), + "stats_history": tuple(stream.stats_history), + }, + ), + ) + + +def _prepare_video_for_mp4( + video: torch.Tensor, + *, + layout: VideoTensorLayout, +) -> tuple[torch.Tensor, WritableVideoTensorLayout]: + """Convert runtime video layouts into layouts accepted by runner I/O.""" + if layout in {"tchw", "btchw", "bcthw"}: + return video, cast(WritableVideoTensorLayout, layout) + if layout == "bvtchw": + if video.ndim != 6: + raise ValueError( + "layout='bvtchw' expects a 6D [B,V,T,C,H,W] tensor, " + f"got {tuple(video.shape)}." + ) + if video.shape[0] != 1: + raise ValueError( + "layout='bvtchw' MP4 writing expects a single batch element, " + f"got {tuple(video.shape)}." + ) + _, views, frames, channels, height, width = video.shape + canvas = ( + video[0] + .permute(1, 3, 0, 4, 2) + .contiguous() + .reshape(frames, height, views * width, channels) + ) + return canvas, "thwc" + raise ValueError(f"unsupported runtime video layout for MP4: {layout!r}") + + +__all__ = ["Mp4VideoOutputTarget"] diff --git a/flashdreams/tests/test_benchmark_harness.py b/flashdreams/tests/test_benchmark_harness.py index d70f9eb5d..1de03b5db 100644 --- a/flashdreams/tests/test_benchmark_harness.py +++ b/flashdreams/tests/test_benchmark_harness.py @@ -245,6 +245,50 @@ def test_shipped_one_minute_demo_scenarios_load() -> None: } +def test_shipped_omnidreams_demo_replay_scenarios_load() -> None: + repo_root = Path(__file__).resolve().parents[2] + scenarios = load_scenario_file( + repo_root / "configs" / "omnidreams_demo_replay_benchmarks.json" + ) + + assert set(scenarios) == { + "omnidreams-sv-runner-baseline", + "omnidreams-sv-demo-replay", + } + + baseline = scenarios["omnidreams-sv-runner-baseline"] + assert baseline.report_group is not None + assert baseline.report_group.id == "omnidreams-demo" + assert _command_value(baseline.command, "--total-blocks") == "226" + assert "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae" in baseline.command + assert "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf" not in ( + baseline.command + ) + assert baseline.quality_baseline_compare is False + + demo = scenarios["omnidreams-sv-demo-replay"] + assert demo.output_dir_arg is None + assert demo.command[:5] == ( + "uv", + "run", + "--project", + "integrations/omnidreams", + "omnidreams-demo", + ) + assert demo.command[5] == "replay" + assert _command_value(demo.command, "--preset-id") == ( + "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae" + ) + assert _command_value(demo.command, "--total-blocks") == "226" + assert _command_value(demo.command, "--output") == ( + "{output_dir}/omnidreams-sv-demo-replay.mp4" + ) + assert "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf" not in ( + demo.command + ) + assert demo.quality_baseline_compare is False + + def test_shipped_deterministic_quality_scenarios_load() -> None: repo_root = Path(__file__).resolve().parents[2] scenarios = load_scenario_file( diff --git a/flashdreams/tests/test_inference_runtime_api.py b/flashdreams/tests/test_inference_runtime_api.py index 890d7efb9..42f75d688 100644 --- a/flashdreams/tests/test_inference_runtime_api.py +++ b/flashdreams/tests/test_inference_runtime_api.py @@ -3,7 +3,6 @@ from __future__ import annotations -from collections.abc import Mapping from dataclasses import fields from typing import Any, cast @@ -11,29 +10,18 @@ from flashdreams.runtime import ( CanonicalInputs, - CanonicalInputSchema, - CanonicalModality, - DeviceConverterSchema, IdentityInputMapping, InferenceConfig, InferenceInput, InferenceInputSchema, - InferenceRuntime, - InferenceSession, InMemoryMetricsRecorder, - InputCanonicalizer, InputField, - InputMapping, - MetricsRecorder, - ModelAdapter, NullOutputTarget, OutputArtifact, - OutputTarget, RuntimeMetricSample, StepRequest, StepResult, TimeWindow, - UserInputCapability, UserInputEvent, UserInputs, UserInputSchema, @@ -42,18 +30,6 @@ pytestmark = pytest.mark.ci_cpu -_SESSION_HORIZON_S = 3600.0 - -_KEYBOARD_SOURCE = UserInputSchema( - capabilities=( - UserInputCapability( - event_type="keyboard.keydown", payload_fields=frozenset({"key"}) - ), - ) -) -_KEYBOARD_CANONICALIZER = InputCanonicalizer() - - def test_inference_config_keeps_runtime_settings_separate() -> None: denied_app_fields = {"prompt", "output_dir", "browser_settings"} config = InferenceConfig( @@ -267,370 +243,3 @@ def test_timing_metric_samples_must_use_seconds() -> None: unit="ms", category="timing", ) - - -def test_runtime_api_components_compose_for_sequential_session() -> None: - adapter = _FakeAdapter() - config = InferenceConfig(model_id="fake-model") - user_inputs = UserInputs( - events=( - UserInputEvent( - timestamp_s=0.25, - event_type="keyboard.keydown", - payload={"key": "w"}, - ), - ) - ) - inference_input = InferenceInput(global_conditioning={"prompt": "drive forward"}) - output = NullOutputTarget(store_results=True) - metrics = InMemoryMetricsRecorder() - - adapter.validate_config(config) - mapping = adapter.default_input_mapping() - assert mapping is not None - _drive_two_step_session( - adapter=adapter, - config=config, - mapping=mapping, - canonicalizer=_KEYBOARD_CANONICALIZER, - source_schema=_KEYBOARD_SOURCE, - user_inputs=user_inputs, - inference_input=inference_input, - output=output, - metrics=metrics, - ) - - assert output.output_count == 2 - assert [result.output for result in output.results] == ["chunk-0", "chunk-1"] - assert [result.frame_count for result in output.results] == [3, 3] - assert output.results[0].output_window == TimeWindow(start_s=0.0, end_s=0.5) - assert [sample.step_index for sample in metrics.samples] == [0, 1] - assert metrics.closed - - -def test_reference_loop_validates_mapping_before_runtime_creation() -> None: - mapping = _OrderCheckingMapping() - adapter = _OrderCheckingAdapter(mapping=mapping) - - _drive_two_step_session( - adapter=adapter, - config=InferenceConfig(model_id="fake-model"), - mapping=mapping, - canonicalizer=_KEYBOARD_CANONICALIZER, - source_schema=_KEYBOARD_SOURCE, - user_inputs=UserInputs(), - inference_input=InferenceInput(global_conditioning={"prompt": "drive forward"}), - output=NullOutputTarget(), - metrics=InMemoryMetricsRecorder(), - ) - - assert mapping.validated - assert adapter.created_runtime_after_validate - - -def test_reference_loop_does_not_canonicalize_global_conditioning() -> None: - mapping = _CanonicalRecordingMapping() - adapter = _FakeAdapter() - canonicalizer = InputCanonicalizer([_CountingDeviceConverter()]) - source_schema = UserInputSchema( - capabilities=( - UserInputCapability( - event_type="stateful_event", - payload_fields=frozenset(), - ), - ) - ) - - _drive_two_step_session( - adapter=adapter, - config=InferenceConfig(model_id="fake-model"), - mapping=mapping, - canonicalizer=canonicalizer, - source_schema=source_schema, - user_inputs=UserInputs( - events=(UserInputEvent(timestamp_s=0.75, event_type="stateful_event"),) - ), - inference_input=InferenceInput(global_conditioning={"prompt": "drive forward"}), - output=NullOutputTarget(), - metrics=InMemoryMetricsRecorder(), - ) - - assert mapping.global_canonical_values == {} - assert mapping.step_canonical_values == ( - {"stateful_counter": {"count": 0}}, - {"stateful_counter": {"count": 1}}, - ) - - -def test_reference_loop_closes_runtime_when_session_start_fails() -> None: - adapter = _FailingStartAdapter() - output = NullOutputTarget() - metrics = InMemoryMetricsRecorder() - - with pytest.raises(RuntimeError, match="start failed"): - _drive_two_step_session( - adapter=adapter, - config=InferenceConfig(model_id="fake-model"), - mapping=IdentityInputMapping(), - canonicalizer=_KEYBOARD_CANONICALIZER, - source_schema=_KEYBOARD_SOURCE, - user_inputs=UserInputs(), - inference_input=InferenceInput( - global_conditioning={"prompt": "drive forward"} - ), - output=output, - metrics=metrics, - ) - - assert adapter.runtime is not None - assert adapter.runtime.closed - assert output.closed - assert metrics.closed - - -def _drive_two_step_session( - *, - adapter: ModelAdapter, - config: InferenceConfig, - mapping: InputMapping, - canonicalizer: InputCanonicalizer, - source_schema: UserInputSchema, - user_inputs: UserInputs, - inference_input: InferenceInput, - output: OutputTarget, - metrics: MetricsRecorder, -) -> None: - mapping.validate( - canonical_schema=adapter.canonical_input_schema, - inference_input_schema=adapter.inference_input_schema, - ) - canonicalizer.reset() - initial_inputs = mapping.map_global_conditioning_inputs( - canonical_inputs=CanonicalInputs(), - inference_input=inference_input, - ) - runtime = adapter.create_runtime(config) - session: InferenceSession | None = None - output_opened = False - try: - session = runtime.start_session(initial_inputs) - output.open() - output_opened = True - while (request := session.next_step_request()) is not None: - step_inputs = mapping.map_step_inputs( - canonical_inputs=canonicalizer.canonicalize( - user_inputs, - window=request.user_input_window - or TimeWindow(start_s=0.0, end_s=_SESSION_HORIZON_S), - source_schema=source_schema, - ), - # Per-step calls carry only the step payload. A changed prompt - # or scene starts or resets a session outside this loop. - inference_input=InferenceInput( - step={"chunk_index": request.step_index}, - ), - request=request, - ) - result = session.step(step_inputs) - output.write(result) - metrics.record_timing( - "model_step", - float(result.metrics["model_step_s"]), - step_index=result.step_index, - ) - finally: - if output_opened: - output.close() - if session is not None: - session.close() - runtime.close() - metrics.close() - - -class _FakeAdapter: - model_id = "fake-model" - inference_input_schema = InferenceInputSchema( - global_conditioning_fields=(InputField(name="prompt"),), - step_fields=(InputField(name="chunk_index"),), - ) - canonical_input_schema = CanonicalInputSchema() - - def default_input_mapping(self) -> InputMapping: - return IdentityInputMapping() - - def validate_config(self, config: InferenceConfig) -> None: - if config.model_id != self.model_id: - raise ValueError(f"Unsupported model_id={config.model_id!r}.") - - def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: - self.validate_config(config) - return _FakeRuntime(inference_input_schema=self.inference_input_schema) - - -class _FakeRuntime: - def __init__(self, *, inference_input_schema: InferenceInputSchema) -> None: - self._inference_input_schema = inference_input_schema - self.closed = False - - def start_session(self, inputs: InferenceInput) -> InferenceSession: - self._inference_input_schema.require_global_conditioning(inputs) - return _FakeSession(inference_input_schema=self._inference_input_schema) - - def close(self) -> None: - self.closed = True - - -class _FailingRuntime(_FakeRuntime): - def start_session(self, inputs: InferenceInput) -> InferenceSession: - del inputs - raise RuntimeError("start failed") - - -class _FakeSession: - def __init__(self, *, inference_input_schema: InferenceInputSchema) -> None: - self._inference_input_schema = inference_input_schema - self.step_index = 0 - self.closed = False - - def next_step_request(self) -> StepRequest | None: - if self.step_index >= 2: - return None - return StepRequest( - step_index=self.step_index, - inference_input_schema=self._inference_input_schema, - user_input_window=TimeWindow( - start_s=0.5 * self.step_index, - end_s=0.5 * (self.step_index + 1), - ), - ) - - def step(self, inputs: InferenceInput) -> StepResult: - self._inference_input_schema.require_step(inputs) - result = StepResult( - step_index=self.step_index, - output=f"chunk-{self.step_index}", - frame_count=3, - output_window=TimeWindow( - start_s=0.5 * self.step_index, - end_s=0.5 * (self.step_index + 1), - ), - metrics={"model_step_s": 0.01}, - ) - self.step_index += 1 - return result - - def reset(self, inputs: InferenceInput | None = None) -> None: - del inputs - self.step_index = 0 - - def close(self) -> None: - self.closed = True - - -class _OrderCheckingMapping(IdentityInputMapping): - def __init__(self) -> None: - self.validated = False - - def validate( - self, - *, - canonical_schema: CanonicalInputSchema | None = None, - inference_input_schema: InferenceInputSchema | None = None, - ) -> None: - super().validate( - canonical_schema=canonical_schema, - inference_input_schema=inference_input_schema, - ) - self.validated = True - - -class _CanonicalRecordingMapping(IdentityInputMapping): - def __init__(self) -> None: - self.global_canonical_values: Mapping[str, Any] | None = None - self._step_canonical_values: list[Mapping[str, Any]] = [] - - @property - def step_canonical_values(self) -> tuple[Mapping[str, Any], ...]: - return tuple(self._step_canonical_values) - - def map_global_conditioning_inputs( - self, - *, - canonical_inputs: CanonicalInputs, - inference_input: InferenceInput, - ) -> InferenceInput: - self.global_canonical_values = canonical_inputs.values - return super().map_global_conditioning_inputs( - canonical_inputs=canonical_inputs, - inference_input=inference_input, - ) - - def map_step_inputs( - self, - *, - canonical_inputs: CanonicalInputs, - inference_input: InferenceInput, - request: StepRequest, - ) -> InferenceInput: - self._step_canonical_values.append(canonical_inputs.values) - return super().map_step_inputs( - canonical_inputs=canonical_inputs, - inference_input=inference_input, - request=request, - ) - - -_STATEFUL_COUNTER = CanonicalModality( - name="stateful_counter", - payload_fields=frozenset({"count"}), -) - - -class _CountingDeviceConverter: - schema = DeviceConverterSchema( - name="stateful-counter", - produces=_STATEFUL_COUNTER, - consumes=(UserInputCapability(event_type="stateful_event"),), - ) - - def __init__(self) -> None: - self.count = 0 - - def reset(self) -> None: - self.count = 0 - - def convert( - self, - user_inputs: UserInputs, - window: TimeWindow, - ) -> Mapping[str, Any] | None: - del window - self.count += len(user_inputs.events) - return _STATEFUL_COUNTER.value({"count": self.count}) - - -class _OrderCheckingAdapter(_FakeAdapter): - canonical_input_schema = CanonicalInputSchema() - - def __init__(self, *, mapping: _OrderCheckingMapping) -> None: - self._mapping = mapping - self.created_runtime_after_validate = False - - def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: - self.validate_config(config) - self.created_runtime_after_validate = self._mapping.validated - return _FakeRuntime(inference_input_schema=self.inference_input_schema) - - -class _FailingStartAdapter(_FakeAdapter): - canonical_input_schema = CanonicalInputSchema() - - def __init__(self) -> None: - self.runtime: _FailingRuntime | None = None - - def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: - self.validate_config(config) - self.runtime = _FailingRuntime( - inference_input_schema=self.inference_input_schema - ) - return self.runtime diff --git a/flashdreams/tests/test_runtime_demo_api.py b/flashdreams/tests/test_runtime_demo_api.py new file mode 100644 index 000000000..7719c7c49 --- /dev/null +++ b/flashdreams/tests/test_runtime_demo_api.py @@ -0,0 +1,458 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +import pytest +import torch + +from flashdreams.infra.video_output import VideoStepResult +from flashdreams.runtime import ( + CanonicalInputs, + CanonicalInputSchema, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InferenceRuntime, + InferenceSession, + InputCanonicalizer, + InputField, + InputMapping, + InputMappingSchema, + NullMetricsRecorder, + NullOutputTarget, + OutputArtifact, + OutputTarget, + StepRequest, + StepResult, + TimeWindow, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + NullOutputSpec, + PreparedScenario, + WebRTCOutputSpec, + build_output_target, + run_replay_demo, +) +from flashdreams.runtime.demo.webrtc import build_webrtc_demo +from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager + +pytestmark = pytest.mark.ci_cpu + + +def test_replay_demo_uses_shared_runner() -> None: + adapter = _FakeDemoAdapter() + output = _RecordingOutputTarget() + calls: list[dict[str, Any]] = [] + + def fake_runner(**kwargs: Any) -> Sequence[OutputArtifact]: + calls.append(kwargs) + return (OutputArtifact(kind="test/artifact", uri="memory://artifact"),) + + spec = DemoSpec( + model_id="fake-demo", + scenario="valid-scenario", + input_mode="replay", + output=NullOutputSpec(), + ) + + artifacts = run_replay_demo( + spec=spec, + adapter=adapter, + output_target_factory=lambda output_spec: output, + metrics=NullMetricsRecorder(), + runner=fake_runner, + ) + + assert artifacts == (OutputArtifact(kind="test/artifact", uri="memory://artifact"),) + assert len(calls) == 1 + assert calls[0]["adapter"] is adapter + assert calls[0]["config"] == spec.config + assert calls[0]["mapping"] is adapter.prepared_scenario.mapping + assert calls[0]["canonicalizer"] is adapter.prepared_scenario.canonicalizer + assert calls[0]["source_schema"] is adapter.prepared_scenario.source_schema + assert calls[0]["user_inputs"] is adapter.prepared_scenario.user_inputs + assert calls[0]["initial_inputs"] is adapter.prepared_scenario.initial_inputs + assert calls[0]["output"] is output + assert adapter.prepare_scenario_calls == [spec] + assert not adapter.create_runtime_called + + +def test_replay_demo_builds_output_target_from_spec(tmp_path: Path) -> None: + writer_calls: list[dict[str, Any]] = [] + + def fake_writer( + video: torch.Tensor, + path: Path, + *, + fps: int | float, + layout: str, + install_hint: str, + ) -> Path: + del install_hint + writer_calls.append( + { + "shape": tuple(video.shape), + "path": path, + "fps": fps, + "layout": layout, + } + ) + return path + + spec = DemoSpec( + model_id="fake-demo", + scenario="valid-scenario", + input_mode="replay", + output=Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=12), + ) + + artifacts = run_replay_demo( + spec=spec, + adapter=_FakeDemoAdapter(video_output=True), + output_target_factory=lambda output_spec: build_output_target( + output_spec, + mp4_writer=fake_writer, + ), + ) + + assert len(artifacts) == 1 + assert artifacts[0].kind == "video/mp4" + assert artifacts[0].uri == str(tmp_path / "demo.mp4") + assert writer_calls == [ + { + "shape": (2, 2, 2, 3), + "path": tmp_path / "demo.mp4", + "fps": 12, + "layout": "thwc", + } + ] + + +def test_replay_demo_fails_before_runtime_creation_when_scenario_invalid() -> None: + adapter = _FakeDemoAdapter(scenario_valid=False) + output_factory_calls = 0 + + def output_factory(output_spec: object) -> OutputTarget: + nonlocal output_factory_calls + del output_spec + output_factory_calls += 1 + return NullOutputTarget() + + spec = DemoSpec( + model_id="fake-demo", + scenario="missing-scenario", + input_mode="replay", + output=NullOutputSpec(), + ) + + with pytest.raises(ValueError, match="invalid scenario"): + run_replay_demo( + spec=spec, + adapter=adapter, + output_target_factory=output_factory, + ) + + assert adapter.prepare_scenario_calls == [spec] + assert not adapter.create_runtime_called + assert output_factory_calls == 0 + + +def test_demo_adapter_declares_supported_modes() -> None: + adapter = _FakeDemoAdapter( + input_modes=("replay",), + output_modes=("null", "mp4", "webrtc"), + ) + + assert adapter.supported_input_modes() == ("replay",) + assert adapter.supported_output_modes() == ("null", "mp4", "webrtc") + + with pytest.raises(ValueError, match="input_mode='keyboard-driving'"): + run_replay_demo( + spec=DemoSpec( + model_id="fake-demo", + scenario="valid-scenario", + input_mode="keyboard-driving", + output=NullOutputSpec(), + ), + adapter=adapter, + ) + + assert adapter.prepare_scenario_calls == [] + assert not adapter.create_runtime_called + + +def test_webrtc_demo_uses_existing_session_manager_with_adapter_runtime() -> None: + adapter = _FakeDemoAdapter() + spec = DemoSpec( + model_id="fake-demo", + scenario="valid-scenario", + input_mode="keyboard-driving", + output=WebRTCOutputSpec( + host="0.0.0.0", + port=8082, + fps=24, + video_width=16, + video_height=8, + warmup_chunks=0, + warmup_timeout_s=1.0, + ), + ) + + 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 + + +class _ChunkIndexMapping: + mapping_schema = InputMappingSchema( + name="chunk-index", + produces_global_conditioning=(InputField(name="prompt"),), + produces_step=(InputField(name="chunk_index"),), + ) + + def validate( + self, + *, + canonical_schema: CanonicalInputSchema | None = None, + inference_input_schema: InferenceInputSchema | None = None, + ) -> None: + del canonical_schema, inference_input_schema + + def map_global_conditioning_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + ) -> InferenceInput: + del canonical_inputs + return inference_input + + def map_step_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + request: StepRequest, + ) -> InferenceInput: + del canonical_inputs + return InferenceInput( + global_conditioning=inference_input.global_conditioning, + step={"chunk_index": request.step_index}, + metadata=inference_input.metadata, + ) + + +class _FakeDemoAdapter: + model_id = "fake-demo" + inference_input_schema = InferenceInputSchema( + global_conditioning_fields=(InputField(name="prompt"),), + step_fields=(InputField(name="chunk_index"),), + ) + canonical_input_schema = CanonicalInputSchema() + + def __init__( + self, + *, + scenario_valid: bool = True, + video_output: bool = False, + input_modes: tuple[str, ...] = ("replay", "keyboard-driving"), + output_modes: tuple[str, ...] = ("null", "mp4", "webrtc"), + ) -> None: + self._scenario_valid = scenario_valid + self._video_output = video_output + self._input_modes = input_modes + self._output_modes = output_modes + self.mapping = _ChunkIndexMapping() + self.prepared_scenario = PreparedScenario( + initial_inputs=InferenceInput( + global_conditioning={"prompt": "drive forward"}, + ), + user_inputs=UserInputs(), + source_schema=UserInputSchema(), + canonicalizer=InputCanonicalizer(), + mapping=self.mapping, + ) + 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 + + def supported_output_modes(self) -> tuple[str, ...]: + return self._output_modes + + def default_input_mapping(self) -> InputMapping: + return IdentityInputMapping() + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.model_id: + raise ValueError(f"Unsupported model_id={config.model_id!r}.") + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + self.create_runtime_called = True + self.runtime = _FakeRuntime( + inference_input_schema=self.inference_input_schema, + video_output=self._video_output, + ) + return self.runtime + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + self.prepare_scenario_calls.append(spec) + if not self._scenario_valid: + 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__( + self, + *, + inference_input_schema: InferenceInputSchema, + video_output: bool, + ) -> None: + self._inference_input_schema = inference_input_schema + self._video_output = video_output + self.session: _FakeSession | None = None + self.closed = False + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + self._inference_input_schema.require_global_conditioning(inputs) + self.session = _FakeSession( + inference_input_schema=self._inference_input_schema, + video_output=self._video_output, + ) + return self.session + + def close(self) -> None: + self.closed = True + + +class _FakeSession: + def __init__( + self, + *, + inference_input_schema: InferenceInputSchema, + video_output: bool, + ) -> None: + self._inference_input_schema = inference_input_schema + self._video_output = video_output + self.step_index = 0 + self.closed = False + + def next_step_request(self) -> StepRequest | None: + if self.step_index >= 2: + return None + return StepRequest( + step_index=self.step_index, + user_input_window=TimeWindow( + start_s=0.5 * self.step_index, + end_s=0.5 * (self.step_index + 1), + ), + ) + + 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, + video_chunk=torch.full( + (1, 1, 1, 3, 2, 2), + self.step_index, + dtype=torch.float32, + ), + layout="bvtchw", + ) + 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), + ), + ) + self.step_index += 1 + return result + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self.step_index = 0 + + def close(self) -> None: + self.closed = True + + +class _RecordingOutputTarget: + def open(self) -> None: + return None + + def write(self, result: StepResult) -> None: + del result + + def close(self) -> Sequence[OutputArtifact]: + return () + + +class _FakeWebRTCRuntime: + async def initialize(self) -> None: + return None + + async def reset_for_new_session(self) -> None: + return None + + 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[Any], + frame_times: list[float], + ) -> Any: + del segments, frame_times + return None + + async def close(self) -> None: + return None + + def send_exit_signal(self) -> None: + return None + + def wait_for_termination(self) -> None: + return None diff --git a/flashdreams/tests/test_runtime_input_mapping.py b/flashdreams/tests/test_runtime_input_mapping.py index b5f0744b9..b774371f4 100644 --- a/flashdreams/tests/test_runtime_input_mapping.py +++ b/flashdreams/tests/test_runtime_input_mapping.py @@ -201,9 +201,7 @@ def test_model_declares_required_and_optional_fields_per_phase() -> None: ("global_conditioning", "prompt"), ("step", "steering"), } - assert {(phase, f.name) for phase, f in optional} == { - ("step", "camera_delta") - } + assert {(phase, f.name) for phase, f in optional} == {("step", "camera_delta")} def test_required_fields_can_be_filtered_by_phase() -> None: @@ -214,8 +212,7 @@ def test_required_fields_can_be_filtered_by_phase() -> None: def test_field_lookup_is_phase_scoped() -> None: assert ( - DRIVING_MODEL.field_for(name="prompt", phase="global_conditioning") - is not None + DRIVING_MODEL.field_for(name="prompt", phase="global_conditioning") is not None ) assert DRIVING_MODEL.field_for(name="prompt", phase="step") is None @@ -275,12 +272,13 @@ def test_compatible_source_model_and_mapping_can_drive() -> None: ) assert compatibility.can_drive - assert { - (p, f.name) for p, f in compatibility.satisfied_required_model_fields - } == {("global_conditioning", "prompt"), ("step", "steering")} - assert { - (p, f.name) for p, f in compatibility.available_optional_model_fields - } == {("step", "camera_delta")} + assert {(p, f.name) for p, f in compatibility.satisfied_required_model_fields} == { + ("global_conditioning", "prompt"), + ("step", "steering"), + } + assert {(p, f.name) for p, f in compatibility.available_optional_model_fields} == { + ("step", "camera_delta") + } def test_missing_required_model_field_blocks_the_run() -> None: @@ -291,9 +289,9 @@ def test_missing_required_model_field_blocks_the_run() -> None: ) assert not compatibility.can_drive - assert [ - f.name for _, f in compatibility.missing_required_model_fields - ] == ["steering"] + assert [f.name for _, f in compatibility.missing_required_model_fields] == [ + "steering" + ] def test_missing_source_capability_is_reported_when_it_blocks() -> None: @@ -416,9 +414,7 @@ def test_combining_mappings_unions_their_surfaces() -> None: combined = combine_mapping_schemas((PROMPT_MAPPING, STEERING_MAPPING)) assert {m.name for m in combined.consumes} == {"driver_command"} - assert [f.name for f in combined.produces_global_conditioning] == [ - "prompt" - ] + assert [f.name for f in combined.produces_global_conditioning] == ["prompt"] assert [f.name for f in combined.produces_step] == ["steering"] diff --git a/flashdreams/tests/test_runtime_runner.py b/flashdreams/tests/test_runtime_runner.py new file mode 100644 index 000000000..b755ff48a --- /dev/null +++ b/flashdreams/tests/test_runtime_runner.py @@ -0,0 +1,660 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +import pytest + +from flashdreams.runtime import ( + DRIVER_COMMAND, + CanonicalInputs, + CanonicalInputSchema, + CanonicalModality, + DeviceConverterSchema, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InferenceRuntime, + InferenceSession, + InMemoryMetricsRecorder, + InputCanonicalizer, + InputField, + InputMapping, + InputMappingSchema, + NullOutputTarget, + OutputArtifact, + RuntimeMetricSample, + StepRequest, + StepResult, + TimeWindow, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, + run_inference_session, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_run_inference_session_completes_two_step_run() -> None: + adapter = _FakeAdapter() + output = NullOutputTarget(store_results=True) + metrics = InMemoryMetricsRecorder() + + artifacts = run_inference_session( + adapter=adapter, + config=InferenceConfig(model_id="fake-model"), + mapping=_ChunkIndexMapping(), + canonicalizer=InputCanonicalizer(), + source_schema=UserInputSchema(), + user_inputs=UserInputs(), + initial_inputs=InferenceInput(global_conditioning={"prompt": "drive forward"}), + output=output, + metrics=metrics, + ) + + assert artifacts == () + assert output.closed + assert output.output_count == 2 + assert [result.output for result in output.results] == ["chunk-0", "chunk-1"] + assert [result.frame_count for result in output.results] == [3, 3] + assert output.results[0].output_window == TimeWindow(start_s=0.0, end_s=0.5) + assert adapter.runtime is not None + assert adapter.runtime.closed + assert adapter.runtime.session is not None + assert adapter.runtime.session.closed + assert [sample.name for sample in metrics.samples] == ["model_step", "model_step"] + assert [sample.step_index for sample in metrics.samples] == [0, 1] + assert metrics.closed + + +def test_runner_preserves_initial_step_inputs_for_identity_mapping() -> None: + adapter = _FakeAdapter() + + run_inference_session( + adapter=adapter, + config=InferenceConfig(model_id="fake-model"), + mapping=IdentityInputMapping(), + canonicalizer=InputCanonicalizer(), + source_schema=UserInputSchema(), + user_inputs=UserInputs(), + initial_inputs=InferenceInput( + global_conditioning={"prompt": "drive forward"}, + step={"chunk_index": 42}, + ), + output=NullOutputTarget(), + metrics=InMemoryMetricsRecorder(), + ) + + assert adapter.runtime is not None + assert adapter.runtime.session is not None + assert [dict(inputs.step) for inputs in adapter.runtime.session.step_inputs] == [ + {"chunk_index": 42}, + {"chunk_index": 42}, + ] + assert [ + dict(inputs.global_conditioning) + for inputs in adapter.runtime.session.step_inputs + ] == [{}, {}] + + +def test_runner_validates_mapping_before_runtime_creation() -> None: + mapping = _OrderCheckingMapping() + adapter = _OrderCheckingAdapter(mapping=mapping) + + run_inference_session( + adapter=adapter, + config=InferenceConfig(model_id="fake-model"), + mapping=mapping, + canonicalizer=InputCanonicalizer(), + source_schema=UserInputSchema(), + user_inputs=UserInputs(), + initial_inputs=InferenceInput(global_conditioning={"prompt": "drive forward"}), + output=NullOutputTarget(), + metrics=InMemoryMetricsRecorder(), + ) + + assert mapping.validated + assert adapter.created_runtime_after_validate + + +def test_runner_closes_runtime_when_session_start_fails() -> None: + adapter = _FailingStartAdapter() + output = _RecordingOutputTarget() + metrics = InMemoryMetricsRecorder() + + with pytest.raises(RuntimeError, match="start failed"): + run_inference_session( + adapter=adapter, + config=InferenceConfig(model_id="fake-model"), + mapping=_ChunkIndexMapping(), + canonicalizer=InputCanonicalizer(), + source_schema=UserInputSchema(), + user_inputs=UserInputs(), + initial_inputs=InferenceInput( + global_conditioning={"prompt": "drive forward"} + ), + output=output, + metrics=metrics, + ) + + assert adapter.runtime is not None + assert adapter.runtime.closed + assert output.events == () + assert metrics.closed + + +def test_runner_does_not_canonicalize_global_conditioning() -> None: + mapping = _CanonicalRecordingMapping() + + run_inference_session( + adapter=_FakeAdapter(), + config=InferenceConfig(model_id="fake-model"), + mapping=mapping, + canonicalizer=InputCanonicalizer([_CountingDeviceConverter()]), + source_schema=UserInputSchema( + capabilities=( + UserInputCapability( + event_type="stateful_event", + payload_fields=frozenset(), + ), + ) + ), + user_inputs=UserInputs( + events=(UserInputEvent(timestamp_s=0.75, event_type="stateful_event"),) + ), + initial_inputs=InferenceInput(global_conditioning={"prompt": "drive forward"}), + output=NullOutputTarget(), + metrics=InMemoryMetricsRecorder(), + ) + + assert mapping.global_canonical_values == {} + assert mapping.step_canonical_values == ( + {"stateful_counter": {"count": 0}}, + {"stateful_counter": {"count": 1}}, + ) + + +def test_runner_closes_opened_resources_after_output_failure() -> None: + events: list[str] = [] + adapter = _RecordingAdapter(events=events) + output = _FailingWriteOutputTarget(events=events) + metrics = _RecordingMetricsRecorder(events=events) + + with pytest.raises(RuntimeError, match="write failed"): + run_inference_session( + adapter=adapter, + config=InferenceConfig(model_id="fake-model"), + mapping=_ChunkIndexMapping(), + canonicalizer=InputCanonicalizer(), + source_schema=UserInputSchema(), + user_inputs=UserInputs(), + initial_inputs=InferenceInput( + global_conditioning={"prompt": "drive forward"} + ), + output=output, + metrics=metrics, + ) + + assert events == [ + "runtime.start_session", + "output.open", + "session.step:0", + "output.write:0", + "output.close", + "session.close", + "runtime.close", + "metrics.close", + ] + + +def test_runner_attempts_later_cleanup_when_output_close_fails() -> None: + events: list[str] = [] + adapter = _RecordingAdapter(events=events) + output = _FailingCloseOutputTarget(events=events) + metrics = _RecordingMetricsRecorder(events=events) + + with pytest.raises(RuntimeError, match="close failed"): + run_inference_session( + adapter=adapter, + config=InferenceConfig(model_id="fake-model"), + mapping=_ChunkIndexMapping(), + canonicalizer=InputCanonicalizer(), + source_schema=UserInputSchema(), + user_inputs=UserInputs(), + initial_inputs=InferenceInput( + global_conditioning={"prompt": "drive forward"} + ), + output=output, + metrics=metrics, + ) + + assert events == [ + "runtime.start_session", + "output.open", + "session.step:0", + "output.write:0", + "session.step:1", + "output.write:1", + "output.close", + "session.close", + "runtime.close", + "metrics.close", + ] + + +def test_runner_checks_declared_mapping_compatibility_before_runtime_creation() -> None: + adapter = _DrivingAdapter() + mapping = _UnfeedableDriverCommandMapping() + metrics = InMemoryMetricsRecorder() + + with pytest.raises(ValueError, match="cannot drive this model"): + run_inference_session( + adapter=adapter, + config=InferenceConfig(model_id="fake-model"), + mapping=mapping, + canonicalizer=InputCanonicalizer(), + source_schema=UserInputSchema(), + user_inputs=UserInputs(), + initial_inputs=InferenceInput( + global_conditioning={"prompt": "drive forward"} + ), + output=NullOutputTarget(), + metrics=metrics, + ) + + assert not adapter.create_runtime_called + assert not mapping.validated + assert metrics.closed + + +class _ChunkIndexMapping: + mapping_schema = InputMappingSchema( + name="chunk-index", + produces_global_conditioning=(InputField(name="prompt"),), + produces_step=(InputField(name="chunk_index"),), + ) + + def __init__(self) -> None: + self.validated = False + + def validate( + self, + *, + canonical_schema: CanonicalInputSchema | None = None, + inference_input_schema: InferenceInputSchema | None = None, + ) -> None: + del canonical_schema, inference_input_schema + self.validated = True + + def map_global_conditioning_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + ) -> InferenceInput: + del canonical_inputs + return inference_input + + def map_step_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + request: StepRequest, + ) -> InferenceInput: + del canonical_inputs + return InferenceInput( + global_conditioning=inference_input.global_conditioning, + step={"chunk_index": request.step_index}, + metadata=inference_input.metadata, + ) + + +class _UnfeedableDriverCommandMapping(_ChunkIndexMapping): + mapping_schema = InputMappingSchema( + name="driver-command", + consumes=(DRIVER_COMMAND,), + produces_global_conditioning=(InputField(name="prompt"),), + produces_step=(InputField(name="steering"),), + ) + + def map_step_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + request: StepRequest, + ) -> InferenceInput: + del request + return InferenceInput( + global_conditioning=inference_input.global_conditioning, + step={ + "steering": canonical_inputs.values[DRIVER_COMMAND.name]["steer"], + }, + metadata=inference_input.metadata, + ) + + +class _CanonicalRecordingMapping(_ChunkIndexMapping): + def __init__(self) -> None: + super().__init__() + self.global_canonical_values: Mapping[str, Any] | None = None + self._step_canonical_values: list[Mapping[str, Any]] = [] + + @property + def step_canonical_values(self) -> tuple[Mapping[str, Any], ...]: + return tuple(self._step_canonical_values) + + def map_global_conditioning_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + ) -> InferenceInput: + self.global_canonical_values = canonical_inputs.values + return super().map_global_conditioning_inputs( + canonical_inputs=canonical_inputs, + inference_input=inference_input, + ) + + def map_step_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + request: StepRequest, + ) -> InferenceInput: + self._step_canonical_values.append(canonical_inputs.values) + return super().map_step_inputs( + canonical_inputs=canonical_inputs, + inference_input=inference_input, + request=request, + ) + + +_STATEFUL_COUNTER = CanonicalModality( + name="stateful_counter", + payload_fields=frozenset({"count"}), +) + + +class _CountingDeviceConverter: + schema = DeviceConverterSchema( + name="stateful-counter", + produces=_STATEFUL_COUNTER, + consumes=(UserInputCapability(event_type="stateful_event"),), + ) + + def __init__(self) -> None: + self.count = 0 + + def reset(self) -> None: + self.count = 0 + + def convert( + self, + user_inputs: UserInputs, + window: TimeWindow, + ) -> Mapping[str, Any] | None: + del window + self.count += len(user_inputs.events) + return _STATEFUL_COUNTER.value({"count": self.count}) + + +class _FakeAdapter: + model_id = "fake-model" + inference_input_schema = InferenceInputSchema( + global_conditioning_fields=(InputField(name="prompt"),), + step_fields=(InputField(name="chunk_index"),), + ) + canonical_input_schema = CanonicalInputSchema() + + def __init__(self) -> None: + self.runtime: _FakeRuntime | None = None + self.create_runtime_called = False + + def default_input_mapping(self) -> InputMapping: + return _ChunkIndexMapping() + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.model_id: + raise ValueError(f"Unsupported model_id={config.model_id!r}.") + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + self.create_runtime_called = True + self.runtime = _FakeRuntime(inference_input_schema=self.inference_input_schema) + return self.runtime + + +class _DrivingAdapter(_FakeAdapter): + inference_input_schema = InferenceInputSchema( + global_conditioning_fields=(InputField(name="prompt"),), + step_fields=(InputField(name="steering"),), + ) + + +class _OrderCheckingMapping(_ChunkIndexMapping): + pass + + +class _OrderCheckingAdapter(_FakeAdapter): + def __init__(self, *, mapping: _OrderCheckingMapping) -> None: + super().__init__() + self._mapping = mapping + self.created_runtime_after_validate = False + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + self.created_runtime_after_validate = self._mapping.validated + self.create_runtime_called = True + self.runtime = _FakeRuntime(inference_input_schema=self.inference_input_schema) + return self.runtime + + +class _FailingStartAdapter(_FakeAdapter): + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + self.create_runtime_called = True + self.runtime = _FailingRuntime( + inference_input_schema=self.inference_input_schema + ) + return self.runtime + + +class _RecordingAdapter(_FakeAdapter): + def __init__(self, *, events: list[str]) -> None: + super().__init__() + self._events = events + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + self.create_runtime_called = True + self.runtime = _RecordingRuntime( + inference_input_schema=self.inference_input_schema, + events=self._events, + ) + return self.runtime + + +class _FakeRuntime: + def __init__(self, *, inference_input_schema: InferenceInputSchema) -> None: + self._inference_input_schema = inference_input_schema + self.session: _FakeSession | None = None + self.closed = False + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + self._inference_input_schema.require_global_conditioning(inputs) + self.session = _FakeSession(inference_input_schema=self._inference_input_schema) + return self.session + + def close(self) -> None: + self.closed = True + + +class _FailingRuntime(_FakeRuntime): + def start_session(self, inputs: InferenceInput) -> InferenceSession: + del inputs + raise RuntimeError("start failed") + + +class _RecordingRuntime(_FakeRuntime): + def __init__( + self, + *, + inference_input_schema: InferenceInputSchema, + events: list[str], + ) -> None: + super().__init__(inference_input_schema=inference_input_schema) + self._events = events + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + self._events.append("runtime.start_session") + self._inference_input_schema.require_global_conditioning(inputs) + self.session = _RecordingSession( + inference_input_schema=self._inference_input_schema, + events=self._events, + ) + return self.session + + def close(self) -> None: + self._events.append("runtime.close") + super().close() + + +class _FakeSession: + def __init__(self, *, inference_input_schema: InferenceInputSchema) -> None: + self._inference_input_schema = inference_input_schema + self.step_index = 0 + self.step_inputs: list[InferenceInput] = [] + self.closed = False + + def next_step_request(self) -> StepRequest | None: + if self.step_index >= 2: + return None + return StepRequest( + step_index=self.step_index, + inference_input_schema=self._inference_input_schema, + user_input_window=TimeWindow( + start_s=0.5 * self.step_index, + end_s=0.5 * (self.step_index + 1), + ), + ) + + def step(self, inputs: InferenceInput) -> StepResult: + self._inference_input_schema.require_step(inputs) + self.step_inputs.append(inputs) + result = StepResult( + step_index=self.step_index, + output=f"chunk-{self.step_index}", + frame_count=3, + output_window=TimeWindow( + start_s=0.5 * self.step_index, + end_s=0.5 * (self.step_index + 1), + ), + metrics={"model_step_s": 0.01, "frames": 3}, + ) + self.step_index += 1 + return result + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self.step_index = 0 + + def close(self) -> None: + self.closed = True + + +class _RecordingSession(_FakeSession): + def __init__( + self, + *, + inference_input_schema: InferenceInputSchema, + events: list[str], + ) -> None: + super().__init__(inference_input_schema=inference_input_schema) + self._events = events + + def step(self, inputs: InferenceInput) -> StepResult: + self._events.append(f"session.step:{self.step_index}") + return super().step(inputs) + + def close(self) -> None: + self._events.append("session.close") + super().close() + + +class _RecordingOutputTarget: + def __init__(self, *, events: list[str] | None = None) -> None: + self._events = events + self._opened = False + + @property + def events(self) -> tuple[str, ...]: + return () if self._events is None else tuple(self._events) + + def open(self) -> None: + self._opened = True + if self._events is not None: + self._events.append("output.open") + + def write(self, result: StepResult) -> None: + if not self._opened: + raise RuntimeError("Cannot write to a closed output target.") + if self._events is not None: + self._events.append(f"output.write:{result.step_index}") + + def close(self) -> Sequence[OutputArtifact]: + self._opened = False + if self._events is not None: + self._events.append("output.close") + return () + + +class _FailingWriteOutputTarget(_RecordingOutputTarget): + def write(self, result: StepResult) -> None: + super().write(result) + raise RuntimeError("write failed") + + +class _FailingCloseOutputTarget(_RecordingOutputTarget): + def close(self) -> Sequence[OutputArtifact]: + super().close() + raise RuntimeError("close failed") + + +class _RecordingMetricsRecorder: + def __init__(self, *, events: list[str]) -> None: + self._events = events + self.samples: list[RuntimeMetricSample] = [] + + def record(self, sample: RuntimeMetricSample) -> None: + self.samples.append(sample) + + def record_timing( + self, + name: str, + duration_s: float, + *, + step_index: int | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> None: + self.record( + RuntimeMetricSample( + name=name, + value=duration_s, + unit="s", + step_index=step_index, + category="timing", + metadata={} if metadata is None else metadata, + ) + ) + + def close(self) -> None: + self._events.append("metrics.close") diff --git a/flashdreams/tests/test_runtime_video_output.py b/flashdreams/tests/test_runtime_video_output.py new file mode 100644 index 000000000..898acf734 --- /dev/null +++ b/flashdreams/tests/test_runtime_video_output.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest +import torch + +from flashdreams.infra.video_output import VideoStepResult +from flashdreams.runtime import Mp4VideoOutputTarget, StepResult, TimeWindow + +pytestmark = pytest.mark.ci_cpu + + +def test_mp4_video_output_target_rejects_non_video_payload(tmp_path: Path) -> None: + target = Mp4VideoOutputTarget(output_path=tmp_path / "out.mp4", fps=30) + target.open() + + with pytest.raises(TypeError, match="VideoStepResult"): + target.write(StepResult(step_index=0, output="not-video")) + + +def test_mp4_video_output_target_writes_artifact_on_close(tmp_path: Path) -> None: + calls: list[dict[str, Any]] = [] + + def fake_writer( + video: torch.Tensor, + path: Path, + *, + fps: int | float, + layout: str, + install_hint: str, + ) -> Path: + del install_hint + calls.append( + { + "shape": tuple(video.shape), + "path": path, + "fps": fps, + "layout": layout, + } + ) + return path + + target = Mp4VideoOutputTarget( + output_path=tmp_path / "omnidreams.mp4", + fps=24, + writer=fake_writer, + move_to_cpu=False, + ) + target.open() + target.write( + StepResult( + 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, + output_window=TimeWindow(start_s=1.0, end_s=2.0), + ) + ) + + artifacts = target.close() + + assert len(artifacts) == 1 + assert artifacts[0].kind == "video/mp4" + assert artifacts[0].uri == str(tmp_path / "omnidreams.mp4") + assert calls == [ + { + "shape": (4, 5, 12, 3), + "path": tmp_path / "omnidreams.mp4", + "fps": 24, + "layout": "thwc", + } + ] + assert artifacts[0].metadata["stats_history"] == ( + { + "autoregressive_index": 3, + "model_step_s": 0.5, + "step_index": 3, + "frames": 4, + "output_start_s": 1.0, + "output_end_s": 2.0, + }, + ) diff --git a/integrations/omnidreams/omnidreams/demo/README.md b/integrations/omnidreams/omnidreams/demo/README.md new file mode 100644 index 000000000..d69c0170a --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/README.md @@ -0,0 +1,66 @@ + + +# OmniDreams Shared Demo API + +This folder contains the experimental OmniDreams demo built on +`flashdreams.runtime.demo`. + +Run commands from the FlashDreams workspace root: + +```bash +cd /path/to/flashdreams +export HF_TOKEN= +``` + +## MP4 Replay + +Generate an MP4 from the bundled single-view sample data: + +```bash +mkdir -p outputs +uv run --package flashdreams-omnidreams omnidreams-demo replay \ + --output outputs/omnidreams-demo.mp4 +``` + +This replay path mirrors the benchmark runner path: it uses a prompt, first +frame, and pre-rendered HDMap video. It does not load a Ludus scene or render +HDMaps at runtime. The demo defaults to the stable non-perf OmniDreams preset +used by the benchmark path. + +To provide benchmark-style assets explicitly: + +```bash +uv run --package flashdreams-omnidreams omnidreams-demo replay \ + --prompt "Driving scene from a front-facing car camera." \ + --hdmap-video-paths /path/to/camera_front_wide_120fov_hdmap.mp4 \ + --first-frame-paths /path/to/first_frame.png \ + --camera-names camera_front_wide_120fov \ + --output outputs/omnidreams-demo.mp4 +``` + +Pass `--example-data-uuid ` to select another bundled single-view sample, +or `--no-example-data` to require explicit asset paths. + +The `omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf` preset remains an +explicit `--preset-id` opt-in. It should become the default only after the +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: + +```bash +uv run --package flashdreams-omnidreams omnidreams-demo webrtc \ + --host 0.0.0.0 \ + --port 8082 +``` + +The scene UUID is optional; when omitted, the runtime uses the default +Hugging Face WebRTC scene. Override the scene with `--scene-uuid`, select a +weather variant with `--scene-variant default|rain|snow`, or use +`--scene-dir /path/to/local/scene` for a local staged scene. diff --git a/integrations/omnidreams/omnidreams/demo/__init__.py b/integrations/omnidreams/omnidreams/demo/__init__.py new file mode 100644 index 000000000..6fa3a9b21 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/__init__.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Experimental OmniDreams demo adapter built on ``flashdreams.runtime.demo``.""" + +from omnidreams.demo.adapter import OmnidreamsDemoAdapter +from omnidreams.demo.spec import ( + DEFAULT_OMNIDREAMS_PRESET, + OMNIDREAMS_MODEL_ID, + OmnidreamsReplayScenario, + OmnidreamsWebRTCScenario, +) + +__all__ = [ + "DEFAULT_OMNIDREAMS_PRESET", + "OMNIDREAMS_MODEL_ID", + "OmnidreamsDemoAdapter", + "OmnidreamsReplayScenario", + "OmnidreamsWebRTCScenario", +] diff --git a/integrations/omnidreams/omnidreams/demo/adapter.py b/integrations/omnidreams/omnidreams/demo/adapter.py new file mode 100644 index 000000000..e16c5c0a1 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/adapter.py @@ -0,0 +1,279 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OmniDreams adapter for the shared demo API.""" + +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, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InputCanonicalizer, + InputField, + UserInputSchema, +) +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + PreparedScenario, + WebRTCOutputSpec, +) +from flashdreams.runtime.interfaces import InferenceRuntime + +from .replay import ( + OmnidreamsReplayRuntime, + OmnidreamsReplayRuntimeOptions, + PipelineFactory, +) +from .spec import ( + 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: + """Model-owned OmniDreams adapter consumed by shared demo launchers.""" + + 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() + + @property + def model_id(self) -> str: + return OMNIDREAMS_MODEL_ID + + @property + def inference_input_schema(self) -> InferenceInputSchema: + return InferenceInputSchema( + global_conditioning_fields=( + InputField( + name="scenario", + input_modality="omnidreams/replay-scenario", + description="Resolved OmniDreams replay scenario.", + ), + ) + ) + + @property + def canonical_input_schema(self) -> CanonicalInputSchema | None: + return None + + def default_input_mapping(self) -> IdentityInputMapping: + return self._mapping + + def supported_input_modes(self) -> tuple[str, ...]: + return ("replay", "keyboard-driving") + + def supported_output_modes(self) -> tuple[str, ...]: + return ("mp4", "webrtc") + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + if spec.input_mode != "replay": + raise ValueError( + "OmniDreams prepare_scenario currently supports only " + f"input_mode='replay', got {spec.input_mode!r}." + ) + if not isinstance(spec.output, Mp4OutputSpec): + raise ValueError("OmniDreams replay demo currently requires MP4 output.") + scenario = resolve_replay_scenario( + spec.scenario, + default_prompt=self._default_replay_prompt(spec.config), + ) + return PreparedScenario( + initial_inputs=InferenceInput( + global_conditioning={"scenario": scenario}, + ), + source_schema=UserInputSchema(description="fixed OmniDreams replay input"), + canonicalizer=InputCanonicalizer(), + mapping=self._mapping, + metadata={ + "model_id": self.model_id, + "preset_id": self._preset_id(spec.config), + "num_views": len(scenario.camera_names), + }, + ) + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.model_id: + raise ValueError( + f"OmniDreams adapter requires model_id={self.model_id!r}, " + f"got {config.model_id!r}." + ) + self._pipeline_config(config) + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + return self._replay_runtime_factory( + config=config, + options=OmnidreamsReplayRuntimeOptions( + pipeline_config=self._pipeline_config(config), + pipeline_factory=self._pipeline_factory, + ), + ) + + 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 + if config is None or config.preset_id is None + else config.preset_id + ) + + def _pipeline_config(self, config: InferenceConfig) -> Any: + custom = config.runtime_options.get("pipeline_config") + if custom is not None: + return custom + preset_id = self._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 _default_replay_prompt(self, config: InferenceConfig | None) -> str: + runner = OMNIDREAMS_RUNNERS.get(self._preset_id(config)) + 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/cli.py new file mode 100644 index 000000000..d35a62b78 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/cli.py @@ -0,0 +1,187 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CLI for the experimental shared OmniDreams demo path.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +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 .adapter import OmnidreamsDemoAdapter +from .spec import ( + DEFAULT_OMNIDREAMS_PRESET, + OMNIDREAMS_MODEL_ID, + OmnidreamsWebRTCScenario, +) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Experimental OmniDreams demo using flashdreams.runtime.demo." + ) + subparsers = parser.add_subparsers(dest="command", required=True) + + replay = subparsers.add_parser("replay", help="Run an MP4 replay demo.") + replay.add_argument("--preset-id", default=DEFAULT_OMNIDREAMS_PRESET) + replay.add_argument("--device", default="cuda") + replay.add_argument("--prompt", default=None) + replay.add_argument("--hdmap-video-paths", type=_split_paths, default=()) + replay.add_argument("--first-frame-paths", type=_split_paths, default=()) + replay.add_argument("--camera-names", type=_split_strings, default=()) + replay.add_argument( + "--example-data", + action=argparse.BooleanOptionalAction, + default=None, + help=( + "Use the bundled single-view HF sample when asset paths are omitted " + "(default: auto)." + ), + ) + replay.add_argument("--example-data-uuid", default=DEFAULT_EXAMPLE_DATA_UUID_1V) + replay.add_argument("--total-blocks", type=int, default=60) + replay.add_argument("--pixel-height", type=int, default=704) + replay.add_argument("--pixel-width", type=int, default=1280) + replay.add_argument("--fps", type=int, default=30) + replay.add_argument("--output", type=Path, required=True) + + webrtc = subparsers.add_parser("webrtc", help="Serve a WebRTC driving demo.") + webrtc.add_argument("--preset-id", default=DEFAULT_OMNIDREAMS_PRESET) + webrtc.add_argument("--host", default="0.0.0.0") + webrtc.add_argument("--port", type=int, default=8082) + webrtc.add_argument("--device", default="cuda:0") + webrtc.add_argument("--seed", type=int, default=42) + webrtc.add_argument("--scene-dir", type=Path, default=None) + webrtc.add_argument("--scene-uuid", default=None) + webrtc.add_argument("--scene-variant", default="default") + webrtc.add_argument("--camera-name", default="camera_front_wide_120fov") + webrtc.add_argument("--fps", type=int, default=30) + webrtc.add_argument("--video-height", type=int, default=704) + webrtc.add_argument("--video-width", type=int, default=1280) + webrtc.add_argument("--warmup-chunks", type=int, default=10) + 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( + spec=_webrtc_spec(args, device=str(context.device)), + adapter=adapter, + world_rank=context.world_rank, + ) + return + raise AssertionError(f"Unhandled command: {args.command}") + + +def _replay_spec(args: argparse.Namespace) -> DemoSpec: + scenario: dict[str, object] = { + "example_data": args.example_data, + "example_data_uuid": args.example_data_uuid, + "total_blocks": args.total_blocks, + "pixel_height": args.pixel_height, + "pixel_width": args.pixel_width, + "fps": args.fps, + } + if args.prompt: + scenario["prompt"] = args.prompt + if args.hdmap_video_paths: + scenario["hdmap_video_paths"] = args.hdmap_video_paths + if args.first_frame_paths: + scenario["first_frame_paths"] = args.first_frame_paths + if args.camera_names: + scenario["camera_names"] = args.camera_names + + return DemoSpec( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=args.preset_id, + input_mode="replay", + scenario=scenario, + output=Mp4OutputSpec(path=args.output, fps=args.fps), + config=InferenceConfig( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=args.preset_id, + device=args.device, + ), + ) + + +def _webrtc_spec(args: argparse.Namespace, *, device: str) -> DemoSpec: + return DemoSpec( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=args.preset_id, + input_mode="keyboard-driving", + scenario=OmnidreamsWebRTCScenario( + scene_dir=args.scene_dir, + scene_uuid=args.scene_uuid, + 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( + host=args.host, + port=args.port, + fps=args.fps, + video_width=args.video_width, + video_height=args.video_height, + warmup_chunks=args.warmup_chunks, + warmup_timeout_s=args.warmup_timeout_s, + client_liveness_timeout_s=args.client_liveness_timeout_s, + preload_name="Omnidreams", + ), + config=InferenceConfig( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=args.preset_id, + device=device, + runtime_options={"seed": args.seed}, + ), + ) + + +def _split_paths(value: str) -> tuple[Path, ...]: + return tuple(Path(part) for part in value.split(",") if part) + + +def _split_strings(value: str) -> tuple[str, ...]: + return tuple(part for part in value.split(",") if part) + + +if __name__ == "__main__": + main() diff --git a/integrations/omnidreams/omnidreams/demo/replay.py b/integrations/omnidreams/omnidreams/demo/replay.py new file mode 100644 index 000000000..908d84568 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/replay.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OmniDreams replay runtime for the shared demo runner.""" + +from __future__ import annotations + +import os +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist +from loguru import logger +from omnidreams.runner import _load_video + +from flashdreams.core.distributed import init as init_distributed +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.infra.runner_io import ( + DEFAULT_RUNNER_INSTALL_HINT, + load_first_frame_tensor, +) +from flashdreams.infra.video_output import VideoStepResult +from flashdreams.runtime.config import InferenceConfig +from flashdreams.runtime.inputs import InferenceInput +from flashdreams.runtime.interfaces import InferenceSession +from flashdreams.runtime.types import StepRequest, StepResult + +from .spec import OmnidreamsReplayScenario + +PipelineFactory = Callable[[Any, str], Any] + + +@dataclass(frozen=True, kw_only=True, slots=True) +class OmnidreamsReplayRuntimeOptions: + """Construction knobs for the replay runtime.""" + + pipeline_config: Any + pipeline_factory: PipelineFactory | None = None + output_layout: VideoTensorLayout = "bvtchw" + + +class OmnidreamsReplayRuntime: + """Heavyweight OmniDreams runtime consumed by ``run_inference_session``.""" + + def __init__( + self, + *, + config: InferenceConfig, + options: OmnidreamsReplayRuntimeOptions, + ) -> None: + self.config = config + self.options = options + if _is_torchrun_env() and not dist.is_initialized(): + init_distributed() + + if dist.is_initialized(): + self.local_rank = int(os.environ.get("LOCAL_RANK", "0")) + self.world_size = dist.get_world_size() + self.global_rank = dist.get_rank() + device = f"cuda:{self.local_rank}" + else: + self.local_rank = 0 + self.world_size = 1 + self.global_rank = 0 + device = config.device or "cuda" + + self.is_rank_zero = self.global_rank == 0 + factory = options.pipeline_factory or _default_pipeline_factory + self.pipeline = factory(options.pipeline_config, device) + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + scenario = _scenario_from_inputs(inputs) + return OmnidreamsReplaySession( + pipeline=self.pipeline, + scenario=scenario, + device=torch.device(f"cuda:{self.local_rank}") + if dist.is_initialized() + else torch.device(self.config.device or "cuda"), + is_rank_zero=self.is_rank_zero, + output_layout=self.options.output_layout, + ) + + def close(self) -> None: + pipeline = getattr(self, "pipeline", None) + if pipeline is not None: + close = getattr(pipeline, "close", None) + if callable(close): + close() + del self.pipeline + device = torch.device(self.config.device or "cuda") + if device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.empty_cache() + + +class OmnidreamsReplaySession: + """One MP4 replay rollout over a prepared scenario.""" + + def __init__( + self, + *, + pipeline: Any, + scenario: OmnidreamsReplayScenario, + device: torch.device, + is_rank_zero: bool, + output_layout: VideoTensorLayout, + ) -> None: + self.pipeline = pipeline + self.scenario = scenario + self.device = device + self.is_rank_zero = is_rank_zero + 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._hdmap_videos = self._load_hdmaps() + if self.device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.synchronize(device=self.device) + if dist.is_initialized(): + dist.barrier() + + def next_step_request(self) -> StepRequest | None: + if self._closed: + return None + if self._step_index >= self.scenario.total_blocks: + return None + num_frames = int(self.pipeline.get_num_frames(self._step_index)) + if self._frame_start + num_frames > self._hdmap_videos.shape[2]: + return None + return StepRequest(step_index=self._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)) + frame_end = self._frame_start + num_frames + logger.info( + "OmniDreams demo replay step {} frames=[{}, {})", + step_index, + 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], + ) + 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, + ) + + 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._frame_start = 0 + + def close(self) -> None: + self._closed = True + cache = getattr(self, "_cache", None) + if cache is not None: + del self._cache + + def _initialize_cache(self) -> Any: + scenario = self.scenario + first_frames = [ + load_first_frame_tensor( + path, + pixel_height=scenario.pixel_height, + pixel_width=scenario.pixel_width, + device=self.device, + dtype=self.dtype, + allow_video=True, + install_hint=DEFAULT_RUNNER_INSTALL_HINT, + ) + for path in scenario.first_frame_paths + ] + first_frames_t = torch.stack(first_frames, dim=0).unsqueeze(0) + cache = self.pipeline.initialize_cache( + text=[list(scenario.prompts)], + image=first_frames_t, + view_names=list(scenario.camera_names), + ) + release = getattr(self.pipeline, "release_oneshot_encoders", None) + if callable(release): + release() + return cache + + def _load_hdmaps(self) -> torch.Tensor: + scenario = self.scenario + videos = [ + _load_video( + path, + pixel_height=scenario.pixel_height, + pixel_width=scenario.pixel_width, + device=self.device, + dtype=self.dtype, + ) + for path in scenario.hdmap_video_paths + ] + # [B=1, V, T, C, H, W] + hdmap_videos = torch.stack(videos, dim=0).unsqueeze(0) + if self.is_rank_zero: + logger.info( + "Loaded OmniDreams demo HDMaps shape={} views={}", + tuple(hdmap_videos.shape), + len(scenario.camera_names), + ) + return hdmap_videos + + +def _default_pipeline_factory(pipeline_config: Any, device: str) -> Any: + return pipeline_config.setup().to(device=device).eval() + + +def _scenario_from_inputs(inputs: InferenceInput) -> OmnidreamsReplayScenario: + scenario = inputs.global_conditioning.get("scenario") + if not isinstance(scenario, OmnidreamsReplayScenario): + raise TypeError( + "OmniDreams replay runtime requires global_conditioning['scenario'] " + "to be an 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 + + +__all__ = [ + "OmnidreamsReplayRuntime", + "OmnidreamsReplayRuntimeOptions", + "OmnidreamsReplaySession", + "PipelineFactory", +] diff --git a/integrations/omnidreams/omnidreams/demo/spec.py b/integrations/omnidreams/omnidreams/demo/spec.py new file mode 100644 index 000000000..0a5dcc062 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/spec.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OmniDreams demo-specific scenario shapes.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from omnidreams.runner import ( + DEFAULT_EXAMPLE_DATA_UUID_1V, + DEFAULT_VIDEO_HEIGHT, + DEFAULT_VIDEO_WIDTH, + _ensure_hf_single_view_example_data_synced, + _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" + + +@dataclass(frozen=True, kw_only=True, slots=True) +class OmnidreamsReplayScenario: + """Resolved replay assets for the shared MP4 demo path.""" + + prompts: tuple[str, ...] + hdmap_video_paths: tuple[Path, ...] + first_frame_paths: tuple[Path, ...] + camera_names: tuple[str, ...] + total_blocks: int = 60 + pixel_height: int = DEFAULT_VIDEO_HEIGHT + pixel_width: int = DEFAULT_VIDEO_WIDTH + fps: int = 30 + + def __post_init__(self) -> None: + if not self.prompts: + raise ValueError("OmnidreamsReplayScenario.prompts must be non-empty.") + num_views = len(self.prompts) + for name, values in ( + ("hdmap_video_paths", self.hdmap_video_paths), + ("first_frame_paths", self.first_frame_paths), + ("camera_names", self.camera_names), + ): + if len(values) != num_views: + raise ValueError( + f"OmnidreamsReplayScenario.{name} has {len(values)} " + f"entries but prompts has {num_views}." + ) + if self.total_blocks <= 0: + raise ValueError("OmnidreamsReplayScenario.total_blocks must be > 0.") + if self.pixel_height <= 0 or self.pixel_width <= 0: + raise ValueError("OmnidreamsReplayScenario pixel dimensions must be > 0.") + if self.fps <= 0: + raise ValueError("OmnidreamsReplayScenario.fps must be > 0.") + object.__setattr__( + self, + "hdmap_video_paths", + tuple(Path(path) for path in self.hdmap_video_paths), + ) + object.__setattr__( + self, + "first_frame_paths", + tuple(Path(path) for path in self.first_frame_paths), + ) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class OmnidreamsWebRTCScenario: + """Scene/options for the shared WebRTC demo path.""" + + scene_dir: Path | None = None + scene_uuid: str | None = DEFAULT_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: + if self.scene_dir is not None: + object.__setattr__(self, "scene_dir", Path(self.scene_dir)) + if not self.scene_variant.strip(): + raise ValueError("OmnidreamsWebRTCScenario.scene_variant is required.") + if not self.camera_name.strip(): + raise ValueError("OmnidreamsWebRTCScenario.camera_name is required.") + + +def resolve_replay_scenario( + value: Any, + *, + default_prompt: str = "", +) -> OmnidreamsReplayScenario: + """Normalize a user/demo scenario into a validated replay scenario.""" + if isinstance(value, OmnidreamsReplayScenario): + _require_existing_paths(value.hdmap_video_paths, label="hdmap_video_paths") + _require_existing_paths(value.first_frame_paths, label="first_frame_paths") + return value + if value is None: + value = {} + if not isinstance(value, Mapping): + raise TypeError( + "OmniDreams replay scenario must be an OmnidreamsReplayScenario " + "a mapping, or None." + ) + + hdmap_paths = _path_tuple(value.get("hdmap_video_paths", ())) + first_paths = _path_tuple(value.get("first_frame_paths", ())) + example_data = _resolve_example_data_default(value) + if example_data and (not hdmap_paths or not first_paths): + example_hdmaps, example_first_frames = ( + _ensure_hf_single_view_example_data_synced( + str(value.get("example_data_uuid", DEFAULT_EXAMPLE_DATA_UUID_1V)) + ) + ) + if not hdmap_paths: + hdmap_paths = example_hdmaps + if not first_paths: + first_paths = example_first_frames + + _require_existing_paths(hdmap_paths, label="hdmap_video_paths") + _require_existing_paths(first_paths, label="first_frame_paths") + if len(hdmap_paths) != len(first_paths): + raise ValueError( + "OmniDreams replay scenario requires one HDMap video and first " + "frame per view." + ) + + num_views = len(hdmap_paths) + prompts = _resolve_prompts(value, num_views, default_prompt=default_prompt) + camera_names = _string_tuple(value.get("camera_names", ())) + if not camera_names: + camera_names = ( + _example_camera_names(num_views) + if example_data + else tuple(f"view_{i}" for i in range(num_views)) + ) + + return OmnidreamsReplayScenario( + prompts=prompts, + hdmap_video_paths=hdmap_paths, + first_frame_paths=first_paths, + camera_names=camera_names, + total_blocks=int(value.get("total_blocks", 60)), + pixel_height=int(value.get("pixel_height", DEFAULT_VIDEO_HEIGHT)), + pixel_width=int(value.get("pixel_width", DEFAULT_VIDEO_WIDTH)), + fps=int(value.get("fps", 30)), + ) + + +def resolve_webrtc_scenario(value: Any) -> OmnidreamsWebRTCScenario: + """Normalize a user/demo scenario into a WebRTC scenario.""" + if value is None: + return OmnidreamsWebRTCScenario() + if isinstance(value, OmnidreamsWebRTCScenario): + return value + if not isinstance(value, Mapping): + raise TypeError( + "OmniDreams WebRTC scenario must be an OmnidreamsWebRTCScenario, " + "a mapping, or None." + ) + 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_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)), + ) + + +def _resolve_prompts( + value: Mapping[str, Any], + num_views: int, + *, + default_prompt: str, +) -> tuple[str, ...]: + prompts = _string_tuple(value.get("prompts", ())) + if prompts: + if len(prompts) != num_views: + raise ValueError( + f"OmniDreams replay prompts has {len(prompts)} entries but " + f"there are {num_views} views." + ) + return prompts + prompt = str(value.get("prompt", "")).strip() + if not prompt: + prompt = default_prompt.strip() + if not prompt: + raise ValueError("OmniDreams replay scenario requires prompt or prompts.") + return (prompt,) * num_views + + +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, "hdmap_video_paths") + or _has_nonempty_value(value, "first_frame_paths") + ) + + +def _bool_value(value: Any) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + lowered = value.strip().lower() + if lowered in {"1", "true", "yes", "on"}: + return True + if lowered in {"0", "false", "no", "off"}: + return False + return bool(value) + + +def _has_nonempty_value(value: Mapping[str, Any], key: str) -> bool: + if key not in value: + return False + raw = value[key] + if raw is None or raw == "": + return False + if isinstance(raw, Sequence) and not isinstance(raw, str): + return len(raw) > 0 + return True + + +def _path_tuple(value: Any) -> tuple[Path, ...]: + if value is None or value == "": + return () + if isinstance(value, (str, Path)): + return (Path(value),) + if isinstance(value, Sequence): + return tuple(Path(path) for path in value) + raise TypeError(f"Expected path or path sequence, got {type(value).__name__}.") + + +def _string_tuple(value: Any) -> tuple[str, ...]: + if value is None or value == "": + return () + if isinstance(value, str): + return (value,) + if isinstance(value, Sequence): + return tuple(str(item) for item in value) + raise TypeError(f"Expected string or string sequence, got {type(value).__name__}.") + + +def _require_existing_paths(paths: tuple[Path, ...], *, label: str) -> None: + if not paths: + raise ValueError(f"OmniDreams replay scenario requires {label}.") + missing = tuple(path for path in paths if not path.exists()) + if missing: + raise FileNotFoundError( + f"OmniDreams replay scenario missing {label}: " + + ", ".join(str(path) for path in missing) + ) + + +__all__ = [ + "DEFAULT_OMNIDREAMS_PRESET", + "OMNIDREAMS_MODEL_ID", + "OmnidreamsReplayScenario", + "OmnidreamsWebRTCScenario", + "resolve_replay_scenario", + "resolve_webrtc_scenario", +] diff --git a/integrations/omnidreams/omnidreams/demo/webrtc.py b/integrations/omnidreams/omnidreams/demo/webrtc.py new file mode 100644 index 000000000..699d47091 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/webrtc.py @@ -0,0 +1,178 @@ +# 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.""" + +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, +) + +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 flashdreams.serving.webrtc.server import ( + close_package_resources as _close_package_resources, +) + + +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 + + 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: + 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, + ) + 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, + } + ) + + +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) + ) + 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_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 create_omnidreams_webrtc_app( + *, + 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("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, + ) + + +def validate_postprocess_preset(preset: str) -> None: + """Validate a configured preset without enabling the output system broadly.""" + if preset: + resolve_postprocess_preset(preset) + + +def _get_omnidreams_manager(app: web.Application) -> OmnidreamsDemoWebRTCSessionManager: + return cast(OmnidreamsDemoWebRTCSessionManager, app[SESSION_MANAGER_KEY]) + + +__all__ = [ + "OmnidreamsDemoWebRTCSessionManager", + "configure_omnidreams_webrtc_app", + "create_omnidreams_webrtc_app", + "postprocess_options", + "session_input", + "validate_postprocess_preset", +] diff --git a/integrations/omnidreams/pyproject.toml b/integrations/omnidreams/pyproject.toml index 99d35d73e..eafd80a1c 100644 --- a/integrations/omnidreams/pyproject.toml +++ b/integrations/omnidreams/pyproject.toml @@ -22,7 +22,8 @@ name = "flashdreams-omnidreams" version = "0.1.0" description = "Omnidreams inference with flashdreams (webrtc / gRPC servers + the interactive-drive desktop demo)" readme = "README.md" -requires-python = ">=3.10,<3.14" +# PyNvVideoCodec 2.1 currently publishes wheels through CPython 3.12. +requires-python = ">=3.10,<3.13" dependencies = [ # Core inference / serving deps (consumed by ``omnidreams.webrtc``, # ``omnidreams.grpc``, and the ``omnidreams.interactive_drive`` desktop @@ -104,6 +105,10 @@ omnidreams-prepare = "omnidreams.prepare:main" # FlashDreams generation, and DrivingGen adapter setup. 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" + # Desktop interactive-drive demo entry point. Requires the # ``interactive-drive`` extra (it adds slangpy); without it the # presenter import fails fast with a clear message. diff --git a/integrations/omnidreams/tests/test_demo_api.py b/integrations/omnidreams/tests/test_demo_api.py new file mode 100644 index 000000000..d9411475a --- /dev/null +++ b/integrations/omnidreams/tests/test_demo_api.py @@ -0,0 +1,577 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import Any, cast + +import omnidreams.demo.spec as spec_module +import omnidreams.demo.webrtc as demo_webrtc_module +import pytest +import torch +from aiohttp import web +from omnidreams.config import OMNIDREAMS_RUNNERS +from omnidreams.demo import ( + DEFAULT_OMNIDREAMS_PRESET, + OMNIDREAMS_MODEL_ID, + OmnidreamsDemoAdapter, + OmnidreamsReplayScenario, + OmnidreamsWebRTCScenario, +) +from omnidreams.demo.cli import _replay_spec, _webrtc_spec, parse_args +from omnidreams.demo.replay import ( + OmnidreamsReplayRuntime, + OmnidreamsReplayRuntimeOptions, +) +from omnidreams.demo.webrtc import OmnidreamsDemoWebRTCSessionManager + +from flashdreams.infra.video_output import VideoStepResult +from flashdreams.runtime import ( + InferenceConfig, + InferenceInput, + OutputArtifact, + OutputTarget, + 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.server import SESSION_MANAGER_KEY + +pytestmark = pytest.mark.ci_cpu + + +def test_omnidreams_demo_defaults_to_stable_non_perf_preset() -> None: + args = parse_args(["replay", "--output", "demo.mp4"]) + + assert args.preset_id == "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae" + assert not args.preset_id.endswith("-perf") + + +def test_omnidreams_demo_adapter_declares_mp4_and_webrtc_modes() -> 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") + + +def test_omnidreams_replay_demo_uses_shared_runner(tmp_path: Path) -> None: + hdmap = tmp_path / "hdmap.mp4" + first_frame = tmp_path / "first.png" + hdmap.write_bytes(b"fake") + first_frame.write_bytes(b"fake") + pipeline_config = object() + adapter = OmnidreamsDemoAdapter() + output = _RecordingOutputTarget() + calls: list[dict[str, Any]] = [] + + def fake_runner(**kwargs: Any) -> Sequence[OutputArtifact]: + calls.append(kwargs) + return (OutputArtifact(kind="video/mp4", uri="memory://omnidreams"),) + + spec = DemoSpec( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + input_mode="replay", + scenario={ + "prompt": "drive through a city", + "hdmap_video_paths": (hdmap,), + "first_frame_paths": (first_frame,), + "camera_names": ("camera_front_wide_120fov",), + "total_blocks": 1, + }, + output=Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=30), + config=InferenceConfig( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + runtime_options={"pipeline_config": pipeline_config}, + ), + ) + + artifacts = run_replay_demo( + spec=spec, + adapter=adapter, + output_target_factory=lambda output_spec: output, + runner=fake_runner, + ) + + assert artifacts == (OutputArtifact(kind="video/mp4", uri="memory://omnidreams"),) + assert len(calls) == 1 + assert calls[0]["adapter"] is adapter + assert calls[0]["config"] == spec.config + scenario = calls[0]["initial_inputs"].global_conditioning["scenario"] + assert isinstance(scenario, OmnidreamsReplayScenario) + assert scenario.prompts == ("drive through a city",) + assert scenario.hdmap_video_paths == (hdmap,) + assert scenario.first_frame_paths == (first_frame,) + assert scenario.camera_names == ("camera_front_wide_120fov",) + + +def test_omnidreams_replay_invalid_scenario_fails_before_runtime_creation( + tmp_path: Path, +) -> None: + adapter = OmnidreamsDemoAdapter( + replay_runtime_factory=lambda **kwargs: pytest.fail( + f"runtime should not be created: {kwargs}" + ) + ) + output_factory_calls = 0 + + def output_factory(output_spec: object) -> OutputTarget: + nonlocal output_factory_calls + del output_spec + output_factory_calls += 1 + return _RecordingOutputTarget() + + spec = DemoSpec( + model_id=OMNIDREAMS_MODEL_ID, + input_mode="replay", + scenario={ + "prompt": "drive", + "hdmap_video_paths": (tmp_path / "missing-hdmap.mp4",), + "first_frame_paths": (tmp_path / "missing-first.png",), + }, + output=Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=30), + config=InferenceConfig( + model_id=OMNIDREAMS_MODEL_ID, + runtime_options={"pipeline_config": object()}, + ), + ) + + with pytest.raises(FileNotFoundError, match="missing hdmap_video_paths"): + run_replay_demo( + spec=spec, + adapter=adapter, + output_target_factory=output_factory, + ) + + assert output_factory_calls == 0 + + +def test_omnidreams_replay_cli_defaults_to_hf_example_data( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + hdmap = tmp_path / "hf-hdmap.mp4" + first_frame = tmp_path / "hf-first.png" + hdmap.write_bytes(b"fake") + first_frame.write_bytes(b"fake") + synced_uuids: list[str] = [] + + def fake_sync(uuid: str) -> tuple[tuple[Path, ...], tuple[Path, ...]]: + synced_uuids.append(uuid) + return (hdmap,), (first_frame,) + + monkeypatch.setattr( + spec_module, + "_ensure_hf_single_view_example_data_synced", + fake_sync, + ) + args = parse_args(["replay", "--output", str(tmp_path / "demo.mp4")]) + spec = _replay_spec(args) + + prepared = OmnidreamsDemoAdapter().prepare_scenario(spec) + + scenario = prepared.initial_inputs.global_conditioning["scenario"] + assert isinstance(scenario, OmnidreamsReplayScenario) + assert synced_uuids == ["239560dc-33d1-11ef-9720-00044bcbccac"] + assert scenario.hdmap_video_paths == (hdmap,) + assert scenario.first_frame_paths == (first_frame,) + assert scenario.camera_names == ("camera_front_wide_120fov",) + assert scenario.prompts == ( + str(getattr(OMNIDREAMS_RUNNERS[DEFAULT_OMNIDREAMS_PRESET], "prompt")), + ) + + +def test_omnidreams_replay_cli_can_disable_example_data(tmp_path: Path) -> None: + args = parse_args( + ["replay", "--no-example-data", "--output", str(tmp_path / "demo.mp4")] + ) + spec = _replay_spec(args) + + with pytest.raises(ValueError, match="requires hdmap_video_paths"): + OmnidreamsDemoAdapter().prepare_scenario(spec) + + +def test_omnidreams_replay_runtime_generates_video_step_result( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import omnidreams.demo.replay as replay_module + + hdmap = tmp_path / "hdmap.mp4" + first_frame = tmp_path / "first.png" + hdmap.write_bytes(b"fake") + first_frame.write_bytes(b"fake") + pipeline = _FakeOmnidreamsPipeline() + monkeypatch.setattr( + replay_module, + "load_first_frame_tensor", + lambda *args, **kwargs: torch.zeros(1, 3, 2, 2), + ) + monkeypatch.setattr( + replay_module, + "_load_video", + lambda *args, **kwargs: torch.zeros(2, 3, 2, 2), + ) + + runtime = OmnidreamsReplayRuntime( + config=InferenceConfig(model_id=OMNIDREAMS_MODEL_ID, device="cpu"), + options=OmnidreamsReplayRuntimeOptions( + pipeline_config=object(), + pipeline_factory=lambda pipeline_config, device: pipeline, + ), + ) + scenario = OmnidreamsReplayScenario( + prompts=("drive",), + hdmap_video_paths=(hdmap,), + first_frame_paths=(first_frame,), + camera_names=("camera_front_wide_120fov",), + total_blocks=1, + pixel_height=2, + pixel_width=2, + fps=30, + ) + session = runtime.start_session( + InferenceInput(global_conditioning={"scenario": scenario}) + ) + + request = session.next_step_request() + assert request is not None + assert request.step_index == 0 + result = session.step(InferenceInput()) + + 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 result.metrics["denoise_s"] == 0.25 + assert session.next_step_request() is None + assert pipeline.initialize_cache_calls == [ + { + "text": [["drive"]], + "image_shape": (1, 1, 1, 3, 2, 2), + "view_names": ["camera_front_wide_120fov"], + } + ] + runtime.close() + + +def test_omnidreams_webrtc_cli_builds_keyboard_driving_spec(tmp_path: Path) -> None: + args = parse_args( + [ + "webrtc", + "--host", + "127.0.0.1", + "--port", + "9090", + "--device", + "cuda:2", + "--seed", + "123", + "--scene-dir", + str(tmp_path / "scene"), + "--scene-uuid", + "scene-1", + "--scene-variant", + "rain", + "--camera-name", + "camera_front_wide_120fov", + "--fps", + "24", + "--video-height", + "32", + "--video-width", + "64", + "--warmup-chunks", + "0", + "--warmup-timeout-s", + "1.5", + "--client-liveness-timeout-s", + "2.5", + "--debug-serve-hdmaps", + "--prefer-sw-encoder", + ] + ) + + spec = _webrtc_spec(args, device="cuda:3") + + assert spec.model_id == OMNIDREAMS_MODEL_ID + assert spec.preset_id == DEFAULT_OMNIDREAMS_PRESET + assert spec.input_mode == "keyboard-driving" + assert isinstance(spec.scenario, OmnidreamsWebRTCScenario) + assert spec.scenario.scene_dir == tmp_path / "scene" + assert spec.scenario.scene_uuid == "scene-1" + assert spec.scenario.scene_variant == "rain" + assert spec.scenario.camera_name == "camera_front_wide_120fov" + assert spec.scenario.debug_serve_hdmaps is True + assert spec.scenario.prefer_sw_encoder is True + assert isinstance(spec.output, WebRTCOutputSpec) + assert spec.output.host == "127.0.0.1" + assert spec.output.port == 9090 + assert spec.output.fps == 24 + assert spec.output.video_width == 64 + assert spec.output.video_height == 32 + assert spec.output.warmup_chunks == 0 + assert spec.output.warmup_timeout_s == 1.5 + assert spec.output.client_liveness_timeout_s == 2.5 + assert spec.config is not None + assert spec.config.device == "cuda:3" + assert spec.config.runtime_options["seed"] == 123 + + +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, + input_mode="keyboard-driving", + scenario=OmnidreamsWebRTCScenario( + scene_uuid="scene-1", + scene_variant="rain", + camera_name="camera_front_wide_120fov", + debug_serve_hdmaps=True, + prefer_sw_encoder=True, + ), + output=WebRTCOutputSpec( + host="0.0.0.0", + port=8082, + fps=24, + video_width=64, + video_height=32, + warmup_chunks=0, + warmup_timeout_s=1.0, + ), + config=InferenceConfig( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + device="cuda:7", + runtime_options={"pipeline_config": pipeline_config, "seed": 123}, + ), + ) + + 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( + monkeypatch: pytest.MonkeyPatch, +) -> None: + 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) + return app + + monkeypatch.setattr( + demo_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, + input_mode="keyboard-driving", + scenario=OmnidreamsWebRTCScenario(), + output=WebRTCOutputSpec( + host="0.0.0.0", + port=8082, + warmup_timeout_s=1.0, + preload_name="Test Omnidreams", + ), + config=InferenceConfig( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + runtime_options={"pipeline_config": object()}, + ), + ) + + demo = build_webrtc_demo(spec=spec, adapter=adapter, create_app=True) + + assert demo.app is not None + assert app_calls[0]["session_manager"] is demo.session_manager + assert app_calls[0]["request_session_url"] == ( + "http://127.0.0.1:8082/request_session" + ) + assert app_calls[0]["preload_name"] == "Test Omnidreams" + 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 + + +def test_omnidreams_webrtc_demo_serves_through_shared_runner( + monkeypatch: pytest.MonkeyPatch, +) -> None: + 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) + return app + + def fake_server_runner(**kwargs: Any) -> None: + server_calls.append(kwargs) + + monkeypatch.setattr( + demo_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, + input_mode="keyboard-driving", + scenario={"scene_uuid": "scene-1"}, + output=WebRTCOutputSpec( + host="0.0.0.0", + port=8082, + warmup_timeout_s=1.0, + ), + config=InferenceConfig( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + runtime_options={"pipeline_config": object()}, + ), + ) + + demo = cast( + WebRTCDemo, + serve_flashdreams_demo( + spec=spec, + adapter=adapter, + world_rank=0, + 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]["host"] == "0.0.0.0" + assert server_calls[0]["port"] == 8082 + assert isinstance(demo.session_manager, OmnidreamsDemoWebRTCSessionManager) + + +class _RecordingOutputTarget: + def open(self) -> None: + return None + + def write(self, result: StepResult) -> None: + del result + + def close(self) -> Sequence[OutputArtifact]: + return () + + +class _FakeOmnidreamsPipeline: + def __init__(self) -> None: + self.initialize_cache_calls: list[dict[str, Any]] = [] + self.released_encoders = False + + def initialize_cache( + self, + *, + text: list[list[str]], + image: torch.Tensor, + view_names: list[str], + ) -> object: + self.initialize_cache_calls.append( + { + "text": text, + "image_shape": tuple(image.shape), + "view_names": view_names, + } + ) + return object() + + def release_oneshot_encoders(self) -> None: + self.released_encoders = True + + def get_num_frames(self, autoregressive_index: int) -> int: + del autoregressive_index + return 1 + + def generate( + self, + *, + autoregressive_index: int, + cache: object, + hdmap: torch.Tensor, + ) -> torch.Tensor: + del cache, hdmap + return torch.full((1, 1, 1, 3, 2, 2), float(autoregressive_index)) + + def finalize(self, *, autoregressive_index: int, cache: object) -> dict[str, float]: + del autoregressive_index, cache + return {"denoise_s": 0.25} + + +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, *args: Any, **kwargs: Any) -> None: + return None + + 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[Any], + frame_times: list[float], + ) -> Any: + del segments, frame_times + return None + + async def close(self) -> None: + return None + + def send_exit_signal(self) -> None: + return None + + def wait_for_termination(self) -> None: + return None diff --git a/uv.lock b/uv.lock index a077bb380..4d0a32ba5 100644 --- a/uv.lock +++ b/uv.lock @@ -1,6 +1,6 @@ version = 1 revision = 3 -requires-python = ">=3.10, <3.14" +requires-python = ">=3.10, <3.13" resolution-markers = [ "python_full_version >= '3.12' and sys_platform == 'win32'", "python_full_version == '3.11.*' and sys_platform == 'win32'", @@ -131,7 +131,7 @@ dependencies = [ { name = "frozenlist" }, { name = "multidict" }, { name = "propcache" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "typing-extensions" }, { name = "yarl" }, ] sdist = { url = "https://files.pythonhosted.org/packages/82/78/8ea7308cac6934de8c74a14f3d5f65d1c89287426688be79538d0e5c013d/aiohttp-3.14.1.tar.gz", hash = "sha256:307f2cff90a764d329e77040603fa032db89c5c24fdad50c4c15334cba744035", size = 7955794, upload-time = "2026-06-07T21:09:35.529Z" } @@ -190,29 +190,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3e/61/a0c0a8f327a9c52095cdd8e312391b00d3ed64ab6c72bb5c33d8ec251cf7/aiohttp-3.14.1-cp312-cp312-win32.whl", hash = "sha256:ec8dc383ee57ea3e883477dcca3f11b65d58199f1080acaf4cd6ad9a99698be4", size = 452771, upload-time = "2026-06-07T21:07:09.669Z" }, { url = "https://files.pythonhosted.org/packages/df/d9/ea367c75f16ac9c6cdc8febb25e8318fa21a2b1bc8d6514d4b2d890bface/aiohttp-3.14.1-cp312-cp312-win_amd64.whl", hash = "sha256:2aa92c87868cd13674989f9ee83e5f9f7ea4237589b728048e1f0c8f6caa3271", size = 479873, upload-time = "2026-06-07T21:07:11.538Z" }, { url = "https://files.pythonhosted.org/packages/03/64/8d96784a7851156db8a4c6c3f6f91042fdf39fb15a4cc38c8b3c14833c45/aiohttp-3.14.1-cp312-cp312-win_arm64.whl", hash = "sha256:2c840c90759922cb5e6dda94596e079a30fb5a5ba548e7e0dc00574703940847", size = 448073, upload-time = "2026-06-07T21:07:13.637Z" }, - { url = "https://files.pythonhosted.org/packages/bc/97/bd137012dd97e1649162b099135a80e1fd59aaa807b2430fc448d1029aff/aiohttp-3.14.1-cp313-cp313-android_21_arm64_v8a.whl", hash = "sha256:b3a03285a7f9c7b016324574a6d92a1c895da6b978cb8f1deee3ac72bc6da178", size = 506882, upload-time = "2026-06-07T21:07:15.501Z" }, - { url = "https://files.pythonhosted.org/packages/ef/79/e5cc690e9d922a66887ceeaca53a8ffd5a7b0be3816142b7abc433742d89/aiohttp-3.14.1-cp313-cp313-android_21_x86_64.whl", hash = "sha256:2a73f487ab8ef5abbb24b7aa9b73e98eaba9e9e031804ff2416f02eca315ccaf", size = 515270, upload-time = "2026-06-07T21:07:17.53Z" }, - { url = "https://files.pythonhosted.org/packages/fe/22/a73ccbf9dbd6e26dda0b24d5fd5db7da92ee3383a79f47677ffb834c5c5b/aiohttp-3.14.1-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:915fbb7b41b115192259f8c9ae58f3ddc444d2b5579917270211858e606a4afd", size = 485841, upload-time = "2026-06-07T21:07:19.555Z" }, - { url = "https://files.pythonhosted.org/packages/3b/b9/57ed8eaf596321c2ad747bd480fb1700dbd7177c60dfc9e4c187f629662e/aiohttp-3.14.1-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:7fb4bdf95b0561a79f259f9d28fbc109728c5ee7f27aff6391f0ca703a329abe", size = 492088, upload-time = "2026-06-07T21:07:21.581Z" }, - { url = "https://files.pythonhosted.org/packages/78/c0/5ebe5270a7c140d7c6f79dcb018640225f14d406c149e4eec04a7d82fe71/aiohttp-3.14.1-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:1b9748363260121d2927704f5d4fc498150669ca3ae93625986ee89c8f80dcd4", size = 501564, upload-time = "2026-06-07T21:07:23.388Z" }, - { url = "https://files.pythonhosted.org/packages/75/7f/8cdaa24fc7983865e0915153b96a9ac5bcdd3548d64c5a27d17cecccad2d/aiohttp-3.14.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:86a6dab78b0e43e2897a3bbe15745aa60dc5423ca437b7b0b164c069bf91b876", size = 751998, upload-time = "2026-06-07T21:07:25.046Z" }, - { url = "https://files.pythonhosted.org/packages/b2/f4/c4227aacfacc5cb0cc2d119b65301d177912a6842cd64e120c47af76064f/aiohttp-3.14.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:4dfd6e47d3c44c2279907607f73a4240b88c69eb8b90da7e2441a8045dfd21da", size = 510918, upload-time = "2026-06-07T21:07:27.28Z" }, - { url = "https://files.pythonhosted.org/packages/ab/01/a2d5f96cd4e74424864d30bc0a7e44d0a12dacdcfa91b5b2d1bd3dca6bf3/aiohttp-3.14.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:317acd9f8602858dc7d59679812c376c7f0b97bcbbf16e0d6237f54141d8a8a6", size = 508657, upload-time = "2026-06-07T21:07:29.252Z" }, - { url = "https://files.pythonhosted.org/packages/e8/ed/3c0fb5c500fdd8e7ebc10d1889c04384fffa1a9163eac1356088ca9da1b1/aiohttp-3.14.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bd869c427324e5cb15195793de951295710db28be7d818247f3097b4ab5d4b96", size = 1757907, upload-time = "2026-06-07T21:07:31.03Z" }, - { url = "https://files.pythonhosted.org/packages/0b/ab/d4c924d9bd5be3050c226612413ce68cb54c70d2c31b661bfc8d9a5b6a70/aiohttp-3.14.1-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:93b032b5ec3255473c143627d21a69ac74ae12f7f33974cb587c564d11b1066f", size = 1737565, upload-time = "2026-06-07T21:07:33.031Z" }, - { url = "https://files.pythonhosted.org/packages/19/2a/37326821ff779084020cdc33224d20b19f42f4183a500ff92022a739eda7/aiohttp-3.14.1-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f234b4deb12f3ad59127e037bc57c40c21e45b45282df7d3a55a0f409f595296", size = 1799018, upload-time = "2026-06-07T21:07:35.003Z" }, - { url = "https://files.pythonhosted.org/packages/b3/4f/6e947ba73e4ce09070761c05ed3a8ceb7c21f5e46798671d8b2aac0e4626/aiohttp-3.14.1-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:9af6779bfb46abf124068327abcdf9ce95c9ef8287a3e8da76ccf2d0f16c28fa", size = 1894416, upload-time = "2026-06-07T21:07:36.956Z" }, - { url = "https://files.pythonhosted.org/packages/9d/6e/dbf1d0625dc711fb2851f4f3c3055c39ed58bae92082d8c627dbe6013736/aiohttp-3.14.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:faccab372e66bc76d5731525e7f1143c922271725b9d38c9f97edcc66266b451", size = 1783881, upload-time = "2026-06-07T21:07:39.063Z" }, - { url = "https://files.pythonhosted.org/packages/44/c2/5e25098a67268ed369483ae7d1a58bd0a13d03aab860d2a0e4a6eb25b046/aiohttp-3.14.1-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f380468b09d2a81633ee863b0ec5648d364bd17bb8ecfb8c2f387f7ac1faf42c", size = 1587572, upload-time = "2026-06-07T21:07:41.058Z" }, - { url = "https://files.pythonhosted.org/packages/2a/bd/cf9cee17e140f942a3de73e658a543aa8fbf35a5fc67a9d2538d52d77f0b/aiohttp-3.14.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:97e704dcd26271f5bda3fa07c3ce0fb76d6d3f8659f4baa1a24442cc9ba177ca", size = 1722137, upload-time = "2026-06-07T21:07:43.014Z" }, - { url = "https://files.pythonhosted.org/packages/89/6d/5684f8c59045c96f81a18cefbc1fbbd79d25b88f1c622f2a5c5c08fcb632/aiohttp-3.14.1-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:269b76ac5394092b95bc4a098f4fc6c191c083c3bd12775d1e30e663132f6a09", size = 1755953, upload-time = "2026-06-07T21:07:45.933Z" }, - { url = "https://files.pythonhosted.org/packages/a8/40/35caf3170f8359760740a7d9aa0fff2e344bef98e1d1186f5a0f6dec17e6/aiohttp-3.14.1-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:5c0b3e614340c889d575451696374c9d17affd54cd607ca0babed8f8c37b9397", size = 1766479, upload-time = "2026-06-07T21:07:48.047Z" }, - { url = "https://files.pythonhosted.org/packages/6d/a1/b0c61e7a137f0d81de49a82023a6df73c3c16d6fefb0f8e4a93d21639002/aiohttp-3.14.1-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:5663ee9257cfa1add7253a7da3035a02f31b6600ec48261585e1800a81533080", size = 1580077, upload-time = "2026-06-07T21:07:50.069Z" }, - { url = "https://files.pythonhosted.org/packages/0b/41/194ea4623693009fcefebef7aef63c141754f153e9cd0d39d3b9e36c175c/aiohttp-3.14.1-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:603a2c834142172ffddc054067f5ec0ca65d57a0aa98a71bc81952573208e345", size = 1791688, upload-time = "2026-06-07T21:07:52.106Z" }, - { url = "https://files.pythonhosted.org/packages/ba/45/4de841f005cfe1fd63e2a2fe011262c515e2a62aa6994b15947e7d717ac9/aiohttp-3.14.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:cb21957bb8aca671c1765e32f58164cf0c50e6bf41c0bbbd16da20732ecaf588", size = 1761094, upload-time = "2026-06-07T21:07:54.113Z" }, - { url = "https://files.pythonhosted.org/packages/e4/ae/dbce10533d3896d544d5053939ed75b7dc31a1b0973d959b1b5ae21028d6/aiohttp-3.14.1-cp313-cp313-win32.whl", hash = "sha256:e509a55f681e6158c20f70f102f9cf61fb20fbc382272bc6d94b7343f2582780", size = 452662, upload-time = "2026-06-07T21:07:56.06Z" }, - { url = "https://files.pythonhosted.org/packages/7b/d9/0bf1a19362c32f06229da5e7ddfcec91f93474d6307f7a2d3135e9c674dc/aiohttp-3.14.1-cp313-cp313-win_amd64.whl", hash = "sha256:1ac8531b638959718e18c2207fbfe297819875da46a740b29dfa29beba64355a", size = 479748, upload-time = "2026-06-07T21:07:58.319Z" }, - { url = "https://files.pythonhosted.org/packages/22/0a/62e7232dc9484fbec112ceb32efb6a624cc7994ec6e2b019286f17c4e8f2/aiohttp-3.14.1-cp313-cp313-win_arm64.whl", hash = "sha256:250d14af67f6b6a1a4a811049b1afa69d61d617fca6bf33149b3ab1a6dbcf7b8", size = 447723, upload-time = "2026-06-07T21:08:00.154Z" }, ] [[package]] @@ -252,7 +229,7 @@ version = "1.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "frozenlist" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/61/62/06741b579156360248d1ec624842ad0edf697050bbaf7c3e46394e106ad1/aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7", size = 25007, upload-time = "2025-07-03T22:54:43.528Z" } wheels = [ @@ -293,7 +270,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "exceptiongroup", marker = "python_full_version < '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, { name = "idna" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/3b/72/5562aabb8dd7181e8e860622a38bea08d17842b99ecd4c91f84ac95251b0/anyio-4.14.1.tar.gz", hash = "sha256:8d648a3544c1a700e3ff78615cd679e4c5c3f149904287e73687b2596963629e", size = 254831, upload-time = "2026-06-24T20:56:06.017Z" } wheels = [ @@ -354,20 +331,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8e/ac/d90df7f1e3b97fc5554cf45076df5045f1e0a6adf13899e10121229b826c/av-16.1.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8cf065f9d438e1921dc31fc7aa045790b58aee71736897866420d80b5450f62a", size = 40817720, upload-time = "2026-01-11T09:57:39.039Z" }, { url = "https://files.pythonhosted.org/packages/80/6f/13c3a35f9dbcebafd03fe0c4cbd075d71ac8968ec849a3cfce406c35a9d2/av-16.1.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a345877a9d3cc0f08e2bc4ec163ee83176864b92587afb9d08dff50f37a9a829", size = 42267396, upload-time = "2026-01-11T09:57:42.115Z" }, { url = "https://files.pythonhosted.org/packages/c8/b9/275df9607f7fb44317ccb1d4be74827185c0d410f52b6e2cd770fe209118/av-16.1.0-cp312-cp312-win_amd64.whl", hash = "sha256:f49243b1d27c91cd8c66fdba90a674e344eb8eb917264f36117bf2b6879118fd", size = 31752045, upload-time = "2026-01-11T09:57:45.106Z" }, - { url = "https://files.pythonhosted.org/packages/75/2a/63797a4dde34283dd8054219fcb29294ba1c25d68ba8c8c8a6ae53c62c45/av-16.1.0-cp313-cp313-macosx_11_0_x86_64.whl", hash = "sha256:ce2a1b3d8bf619f6c47a9f28cfa7518ff75ddd516c234a4ee351037b05e6a587", size = 26916715, upload-time = "2026-01-11T09:57:47.682Z" }, - { url = "https://files.pythonhosted.org/packages/d2/c4/0b49cf730d0ae8cda925402f18ae814aef351f5772d14da72dd87ff66448/av-16.1.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:408dbe6a2573ca58a855eb8cd854112b33ea598651902c36709f5f84c991ed8e", size = 21452167, upload-time = "2026-01-11T09:57:50.606Z" }, - { url = "https://files.pythonhosted.org/packages/51/23/408806503e8d5d840975aad5699b153aaa21eb6de41ade75248a79b7a37f/av-16.1.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:57f657f86652a160a8a01887aaab82282f9e629abf94c780bbdbb01595d6f0f7", size = 39215659, upload-time = "2026-01-11T09:57:53.757Z" }, - { url = "https://files.pythonhosted.org/packages/c4/19/a8528d5bba592b3903f44c28dab9cc653c95fcf7393f382d2751a1d1523e/av-16.1.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:adbad2b355c2ee4552cac59762809d791bda90586d134a33c6f13727fb86cb3a", size = 40874970, upload-time = "2026-01-11T09:57:56.802Z" }, - { url = "https://files.pythonhosted.org/packages/e8/24/2dbcdf0e929ad56b7df078e514e7bd4ca0d45cba798aff3c8caac097d2f7/av-16.1.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f42e1a68ec2aebd21f7eb6895be69efa6aa27eec1670536876399725bbda4b99", size = 40530345, upload-time = "2026-01-11T09:58:00.421Z" }, - { url = "https://files.pythonhosted.org/packages/54/27/ae91b41207f34e99602d1c72ab6ffd9c51d7c67e3fbcd4e3a6c0e54f882c/av-16.1.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:58fe47aeaef0f100c40ec8a5de9abbd37f118d3ca03829a1009cf288e9aef67c", size = 41972163, upload-time = "2026-01-11T09:58:03.756Z" }, - { url = "https://files.pythonhosted.org/packages/fc/7a/22158fb923b2a9a00dfab0e96ef2e8a1763a94dd89e666a5858412383d46/av-16.1.0-cp313-cp313-win_amd64.whl", hash = "sha256:565093ebc93b2f4b76782589564869dadfa83af5b852edebedd8fee746457d06", size = 31729230, upload-time = "2026-01-11T09:58:07.254Z" }, - { url = "https://files.pythonhosted.org/packages/7f/f1/878f8687d801d6c4565d57ebec08449c46f75126ebca8e0fed6986599627/av-16.1.0-cp313-cp313t-macosx_11_0_x86_64.whl", hash = "sha256:574081a24edb98343fd9f473e21ae155bf61443d4ec9d7708987fa597d6b04b2", size = 27008769, upload-time = "2026-01-11T09:58:10.266Z" }, - { url = "https://files.pythonhosted.org/packages/30/f1/bd4ce8c8b5cbf1d43e27048e436cbc9de628d48ede088a1d0a993768eb86/av-16.1.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:9ab00ea29c25ebf2ea1d1e928d7babb3532d562481c5d96c0829212b70756ad0", size = 21590588, upload-time = "2026-01-11T09:58:12.629Z" }, - { url = "https://files.pythonhosted.org/packages/1d/dd/c81f6f9209201ff0b5d5bed6da6c6e641eef52d8fbc930d738c3f4f6f75d/av-16.1.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:a84a91188c1071f238a9523fd42dbe567fb2e2607b22b779851b2ce0eac1b560", size = 40638029, upload-time = "2026-01-11T09:58:15.399Z" }, - { url = "https://files.pythonhosted.org/packages/15/4d/07edff82b78d0459a6e807e01cd280d3180ce832efc1543de80d77676722/av-16.1.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:c2cd0de4dd022a7225ff224fde8e7971496d700be41c50adaaa26c07bb50bf97", size = 41970776, upload-time = "2026-01-11T09:58:19.075Z" }, - { url = "https://files.pythonhosted.org/packages/da/9d/1f48b354b82fa135d388477cd1b11b81bdd4384bd6a42a60808e2ec2d66b/av-16.1.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:0816143530624a5a93bc5494f8c6eeaf77549b9366709c2ac8566c1e9bff6df5", size = 41764751, upload-time = "2026-01-11T09:58:22.788Z" }, - { url = "https://files.pythonhosted.org/packages/2f/c7/a509801e98db35ec552dd79da7bdbcff7104044bfeb4c7d196c1ce121593/av-16.1.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:e3a28053af29644696d0c007e897d19b1197585834660a54773e12a40b16974c", size = 43034355, upload-time = "2026-01-11T09:58:26.125Z" }, - { url = "https://files.pythonhosted.org/packages/36/8b/e5f530d9e8f640da5f5c5f681a424c65f9dd171c871cd255d8a861785a6e/av-16.1.0-cp313-cp313t-win_amd64.whl", hash = "sha256:2e3e67144a202b95ed299d165232533989390a9ea3119d37eccec697dc6dbb0c", size = 31947047, upload-time = "2026-01-11T09:58:31.867Z" }, ] [[package]] @@ -484,18 +447,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7b/2b/2b6435f76bfeb6bbf055596976da087377ede68df465419d192acf00c437/cffi-2.0.0-cp312-cp312-win32.whl", hash = "sha256:da902562c3e9c550df360bfa53c035b2f241fed6d9aef119048073680ace4a18", size = 172932, upload-time = "2025-09-08T23:22:57.188Z" }, { url = "https://files.pythonhosted.org/packages/f8/ed/13bd4418627013bec4ed6e54283b1959cf6db888048c7cf4b4c3b5b36002/cffi-2.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:da68248800ad6320861f129cd9c1bf96ca849a2771a59e0344e88681905916f5", size = 183557, upload-time = "2025-09-08T23:22:58.351Z" }, { url = "https://files.pythonhosted.org/packages/95/31/9f7f93ad2f8eff1dbc1c3656d7ca5bfd8fb52c9d786b4dcf19b2d02217fa/cffi-2.0.0-cp312-cp312-win_arm64.whl", hash = "sha256:4671d9dd5ec934cb9a73e7ee9676f9362aba54f7f34910956b84d727b0d73fb6", size = 177762, upload-time = "2025-09-08T23:22:59.668Z" }, - { url = "https://files.pythonhosted.org/packages/4b/8d/a0a47a0c9e413a658623d014e91e74a50cdd2c423f7ccfd44086ef767f90/cffi-2.0.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:00bdf7acc5f795150faa6957054fbbca2439db2f775ce831222b66f192f03beb", size = 185230, upload-time = "2025-09-08T23:23:00.879Z" }, - { url = "https://files.pythonhosted.org/packages/4a/d2/a6c0296814556c68ee32009d9c2ad4f85f2707cdecfd7727951ec228005d/cffi-2.0.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:45d5e886156860dc35862657e1494b9bae8dfa63bf56796f2fb56e1679fc0bca", size = 181043, upload-time = "2025-09-08T23:23:02.231Z" }, - { url = "https://files.pythonhosted.org/packages/b0/1e/d22cc63332bd59b06481ceaac49d6c507598642e2230f201649058a7e704/cffi-2.0.0-cp313-cp313-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:07b271772c100085dd28b74fa0cd81c8fb1a3ba18b21e03d7c27f3436a10606b", size = 212446, upload-time = "2025-09-08T23:23:03.472Z" }, - { url = "https://files.pythonhosted.org/packages/a9/f5/a2c23eb03b61a0b8747f211eb716446c826ad66818ddc7810cc2cc19b3f2/cffi-2.0.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d48a880098c96020b02d5a1f7d9251308510ce8858940e6fa99ece33f610838b", size = 220101, upload-time = "2025-09-08T23:23:04.792Z" }, - { url = "https://files.pythonhosted.org/packages/f2/7f/e6647792fc5850d634695bc0e6ab4111ae88e89981d35ac269956605feba/cffi-2.0.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:f93fd8e5c8c0a4aa1f424d6173f14a892044054871c771f8566e4008eaa359d2", size = 207948, upload-time = "2025-09-08T23:23:06.127Z" }, - { url = "https://files.pythonhosted.org/packages/cb/1e/a5a1bd6f1fb30f22573f76533de12a00bf274abcdc55c8edab639078abb6/cffi-2.0.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:dd4f05f54a52fb558f1ba9f528228066954fee3ebe629fc1660d874d040ae5a3", size = 206422, upload-time = "2025-09-08T23:23:07.753Z" }, - { url = "https://files.pythonhosted.org/packages/98/df/0a1755e750013a2081e863e7cd37e0cdd02664372c754e5560099eb7aa44/cffi-2.0.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c8d3b5532fc71b7a77c09192b4a5a200ea992702734a2e9279a37f2478236f26", size = 219499, upload-time = "2025-09-08T23:23:09.648Z" }, - { url = "https://files.pythonhosted.org/packages/50/e1/a969e687fcf9ea58e6e2a928ad5e2dd88cc12f6f0ab477e9971f2309b57c/cffi-2.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d9b29c1f0ae438d5ee9acb31cadee00a58c46cc9c0b2f9038c6b0b3470877a8c", size = 222928, upload-time = "2025-09-08T23:23:10.928Z" }, - { url = "https://files.pythonhosted.org/packages/36/54/0362578dd2c9e557a28ac77698ed67323ed5b9775ca9d3fe73fe191bb5d8/cffi-2.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6d50360be4546678fc1b79ffe7a66265e28667840010348dd69a314145807a1b", size = 221302, upload-time = "2025-09-08T23:23:12.42Z" }, - { url = "https://files.pythonhosted.org/packages/eb/6d/bf9bda840d5f1dfdbf0feca87fbdb64a918a69bca42cfa0ba7b137c48cb8/cffi-2.0.0-cp313-cp313-win32.whl", hash = "sha256:74a03b9698e198d47562765773b4a8309919089150a0bb17d829ad7b44b60d27", size = 172909, upload-time = "2025-09-08T23:23:14.32Z" }, - { url = "https://files.pythonhosted.org/packages/37/18/6519e1ee6f5a1e579e04b9ddb6f1676c17368a7aba48299c3759bbc3c8b3/cffi-2.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:19f705ada2530c1167abacb171925dd886168931e0a7b78f5bffcae5c6b5be75", size = 183402, upload-time = "2025-09-08T23:23:15.535Z" }, - { url = "https://files.pythonhosted.org/packages/cb/0e/02ceeec9a7d6ee63bb596121c2c8e9b3a9e150936f4fbef6ca1943e6137c/cffi-2.0.0-cp313-cp313-win_arm64.whl", hash = "sha256:256f80b80ca3853f90c21b23ee78cd008713787b1b1e93eae9f3d6a7134abd91", size = 177780, upload-time = "2025-09-08T23:23:16.761Z" }, ] [[package]] @@ -561,22 +512,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/86/eb/890922a8b03a568ca2f336c36585a4713c55d4d67bf0f0c78924be6315ca/charset_normalizer-3.4.7-cp312-cp312-win32.whl", hash = "sha256:2257141f39fe65a3fdf38aeccae4b953e5f3b3324f4ff0daf9f15b8518666a2c", size = 148460, upload-time = "2026-04-02T09:26:41.416Z" }, { url = "https://files.pythonhosted.org/packages/35/d9/0e7dffa06c5ab081f75b1b786f0aefc88365825dfcd0ac544bdb7b2b6853/charset_normalizer-3.4.7-cp312-cp312-win_amd64.whl", hash = "sha256:5ed6ab538499c8644b8a3e18debabcd7ce684f3fa91cf867521a7a0279cab2d6", size = 159330, upload-time = "2026-04-02T09:26:42.554Z" }, { url = "https://files.pythonhosted.org/packages/9e/5d/481bcc2a7c88ea6b0878c299547843b2521ccbc40980cb406267088bc701/charset_normalizer-3.4.7-cp312-cp312-win_arm64.whl", hash = "sha256:56be790f86bfb2c98fb742ce566dfb4816e5a83384616ab59c49e0604d49c51d", size = 147828, upload-time = "2026-04-02T09:26:44.075Z" }, - { url = "https://files.pythonhosted.org/packages/c1/3b/66777e39d3ae1ddc77ee606be4ec6d8cbd4c801f65e5a1b6f2b11b8346dd/charset_normalizer-3.4.7-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:f496c9c3cc02230093d8330875c4c3cdfc3b73612a5fd921c65d39cbcef08063", size = 309627, upload-time = "2026-04-02T09:26:45.198Z" }, - { url = "https://files.pythonhosted.org/packages/2e/4e/b7f84e617b4854ade48a1b7915c8ccfadeba444d2a18c291f696e37f0d3b/charset_normalizer-3.4.7-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0ea948db76d31190bf08bd371623927ee1339d5f2a0b4b1b4a4439a65298703c", size = 207008, upload-time = "2026-04-02T09:26:46.824Z" }, - { url = "https://files.pythonhosted.org/packages/c4/bb/ec73c0257c9e11b268f018f068f5d00aa0ef8c8b09f7753ebd5f2880e248/charset_normalizer-3.4.7-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a277ab8928b9f299723bc1a2dabb1265911b1a76341f90a510368ca44ad9ab66", size = 228303, upload-time = "2026-04-02T09:26:48.397Z" }, - { url = "https://files.pythonhosted.org/packages/85/fb/32d1f5033484494619f701e719429c69b766bfc4dbc61aa9e9c8c166528b/charset_normalizer-3.4.7-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3bec022aec2c514d9cf199522a802bd007cd588ab17ab2525f20f9c34d067c18", size = 224282, upload-time = "2026-04-02T09:26:49.684Z" }, - { url = "https://files.pythonhosted.org/packages/fa/07/330e3a0dda4c404d6da83b327270906e9654a24f6c546dc886a0eb0ffb23/charset_normalizer-3.4.7-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e044c39e41b92c845bc815e5ae4230804e8e7bc29e399b0437d64222d92809dd", size = 215595, upload-time = "2026-04-02T09:26:50.915Z" }, - { url = "https://files.pythonhosted.org/packages/e3/7c/fc890655786e423f02556e0216d4b8c6bcb6bdfa890160dc66bf52dee468/charset_normalizer-3.4.7-cp313-cp313-manylinux_2_31_armv7l.whl", hash = "sha256:f495a1652cf3fbab2eb0639776dad966c2fb874d79d87ca07f9d5f059b8bd215", size = 201986, upload-time = "2026-04-02T09:26:52.197Z" }, - { url = "https://files.pythonhosted.org/packages/d8/97/bfb18b3db2aed3b90cf54dc292ad79fdd5ad65c4eae454099475cbeadd0d/charset_normalizer-3.4.7-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e712b419df8ba5e42b226c510472b37bd57b38e897d3eca5e8cfd410a29fa859", size = 211711, upload-time = "2026-04-02T09:26:53.49Z" }, - { url = "https://files.pythonhosted.org/packages/6f/a5/a581c13798546a7fd557c82614a5c65a13df2157e9ad6373166d2a3e645d/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:7804338df6fcc08105c7745f1502ba68d900f45fd770d5bdd5288ddccb8a42d8", size = 210036, upload-time = "2026-04-02T09:26:54.975Z" }, - { url = "https://files.pythonhosted.org/packages/8c/bf/b3ab5bcb478e4193d517644b0fb2bf5497fbceeaa7a1bc0f4d5b50953861/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:481551899c856c704d58119b5025793fa6730adda3571971af568f66d2424bb5", size = 202998, upload-time = "2026-04-02T09:26:56.303Z" }, - { url = "https://files.pythonhosted.org/packages/e7/4e/23efd79b65d314fa320ec6017b4b5834d5c12a58ba4610aa353af2e2f577/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:f59099f9b66f0d7145115e6f80dd8b1d847176df89b234a5a6b3f00437aa0832", size = 230056, upload-time = "2026-04-02T09:26:57.554Z" }, - { url = "https://files.pythonhosted.org/packages/b9/9f/1e1941bc3f0e01df116e68dc37a55c4d249df5e6fa77f008841aef68264f/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:f59ad4c0e8f6bba240a9bb85504faa1ab438237199d4cce5f622761507b8f6a6", size = 211537, upload-time = "2026-04-02T09:26:58.843Z" }, - { url = "https://files.pythonhosted.org/packages/80/0f/088cbb3020d44428964a6c97fe1edfb1b9550396bf6d278330281e8b709c/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:3dedcc22d73ec993f42055eff4fcfed9318d1eeb9a6606c55892a26964964e48", size = 226176, upload-time = "2026-04-02T09:27:00.437Z" }, - { url = "https://files.pythonhosted.org/packages/6a/9f/130394f9bbe06f4f63e22641d32fc9b202b7e251c9aef4db044324dac493/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:64f02c6841d7d83f832cd97ccf8eb8a906d06eb95d5276069175c696b024b60a", size = 217723, upload-time = "2026-04-02T09:27:02.021Z" }, - { url = "https://files.pythonhosted.org/packages/73/55/c469897448a06e49f8fa03f6caae97074fde823f432a98f979cc42b90e69/charset_normalizer-3.4.7-cp313-cp313-win32.whl", hash = "sha256:4042d5c8f957e15221d423ba781e85d553722fc4113f523f2feb7b188cc34c5e", size = 148085, upload-time = "2026-04-02T09:27:03.192Z" }, - { url = "https://files.pythonhosted.org/packages/5d/78/1b74c5bbb3f99b77a1715c91b3e0b5bdb6fe302d95ace4f5b1bec37b0167/charset_normalizer-3.4.7-cp313-cp313-win_amd64.whl", hash = "sha256:3946fa46a0cf3e4c8cb1cc52f56bb536310d34f25f01ca9b6c16afa767dab110", size = 158819, upload-time = "2026-04-02T09:27:04.454Z" }, - { url = "https://files.pythonhosted.org/packages/68/86/46bd42279d323deb8687c4a5a811fd548cb7d1de10cf6535d099877a9a9f/charset_normalizer-3.4.7-cp313-cp313-win_arm64.whl", hash = "sha256:80d04837f55fc81da168b98de4f4b797ef007fc8a79ab71c6ec9bc4dd662b15b", size = 147915, upload-time = "2026-04-02T09:27:05.971Z" }, { url = "https://files.pythonhosted.org/packages/db/8f/61959034484a4a7c527811f4721e75d02d653a35afb0b6054474d8185d4c/charset_normalizer-3.4.7-py3-none-any.whl", hash = "sha256:3dce51d0f5e7951f8bb4900c257dad282f49190fdbebecd4ba99bcc41fef404d", size = 61958, upload-time = "2026-04-02T09:28:37.794Z" }, ] @@ -644,26 +579,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/8a/bebe5a3f68b484d3a2b8ffaf84704b3e343ef1addea528132ef148e22b3b/contourpy-1.3.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:3d80b2c0300583228ac98d0a927a1ba6a2ba6b8a742463c564f1d419ee5b211e", size = 1380480, upload-time = "2025-04-15T17:38:06.7Z" }, { url = "https://files.pythonhosted.org/packages/34/db/fcd325f19b5978fb509a7d55e06d99f5f856294c1991097534360b307cf1/contourpy-1.3.2-cp312-cp312-win32.whl", hash = "sha256:90df94c89a91b7362e1142cbee7568f86514412ab8a2c0d0fca72d7e91b62912", size = 178489, upload-time = "2025-04-15T17:38:10.338Z" }, { url = "https://files.pythonhosted.org/packages/01/c8/fadd0b92ffa7b5eb5949bf340a63a4a496a6930a6c37a7ba0f12acb076d6/contourpy-1.3.2-cp312-cp312-win_amd64.whl", hash = "sha256:8c942a01d9163e2e5cfb05cb66110121b8d07ad438a17f9e766317bcb62abf73", size = 223042, upload-time = "2025-04-15T17:38:14.239Z" }, - { url = "https://files.pythonhosted.org/packages/2e/61/5673f7e364b31e4e7ef6f61a4b5121c5f170f941895912f773d95270f3a2/contourpy-1.3.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:de39db2604ae755316cb5967728f4bea92685884b1e767b7c24e983ef5f771cb", size = 271630, upload-time = "2025-04-15T17:38:19.142Z" }, - { url = "https://files.pythonhosted.org/packages/ff/66/a40badddd1223822c95798c55292844b7e871e50f6bfd9f158cb25e0bd39/contourpy-1.3.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:3f9e896f447c5c8618f1edb2bafa9a4030f22a575ec418ad70611450720b5b08", size = 255670, upload-time = "2025-04-15T17:38:23.688Z" }, - { url = "https://files.pythonhosted.org/packages/1e/c7/cf9fdee8200805c9bc3b148f49cb9482a4e3ea2719e772602a425c9b09f8/contourpy-1.3.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:71e2bd4a1c4188f5c2b8d274da78faab884b59df20df63c34f74aa1813c4427c", size = 306694, upload-time = "2025-04-15T17:38:28.238Z" }, - { url = "https://files.pythonhosted.org/packages/dd/e7/ccb9bec80e1ba121efbffad7f38021021cda5be87532ec16fd96533bb2e0/contourpy-1.3.2-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:de425af81b6cea33101ae95ece1f696af39446db9682a0b56daaa48cfc29f38f", size = 345986, upload-time = "2025-04-15T17:38:33.502Z" }, - { url = "https://files.pythonhosted.org/packages/dc/49/ca13bb2da90391fa4219fdb23b078d6065ada886658ac7818e5441448b78/contourpy-1.3.2-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:977e98a0e0480d3fe292246417239d2d45435904afd6d7332d8455981c408b85", size = 318060, upload-time = "2025-04-15T17:38:38.672Z" }, - { url = "https://files.pythonhosted.org/packages/c8/65/5245ce8c548a8422236c13ffcdcdada6a2a812c361e9e0c70548bb40b661/contourpy-1.3.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:434f0adf84911c924519d2b08fc10491dd282b20bdd3fa8f60fd816ea0b48841", size = 322747, upload-time = "2025-04-15T17:38:43.712Z" }, - { url = "https://files.pythonhosted.org/packages/72/30/669b8eb48e0a01c660ead3752a25b44fdb2e5ebc13a55782f639170772f9/contourpy-1.3.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:c66c4906cdbc50e9cba65978823e6e00b45682eb09adbb78c9775b74eb222422", size = 1308895, upload-time = "2025-04-15T17:39:00.224Z" }, - { url = "https://files.pythonhosted.org/packages/05/5a/b569f4250decee6e8d54498be7bdf29021a4c256e77fe8138c8319ef8eb3/contourpy-1.3.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8b7fc0cd78ba2f4695fd0a6ad81a19e7e3ab825c31b577f384aa9d7817dc3bef", size = 1379098, upload-time = "2025-04-15T17:43:29.649Z" }, - { url = "https://files.pythonhosted.org/packages/19/ba/b227c3886d120e60e41b28740ac3617b2f2b971b9f601c835661194579f1/contourpy-1.3.2-cp313-cp313-win32.whl", hash = "sha256:15ce6ab60957ca74cff444fe66d9045c1fd3e92c8936894ebd1f3eef2fff075f", size = 178535, upload-time = "2025-04-15T17:44:44.532Z" }, - { url = "https://files.pythonhosted.org/packages/12/6e/2fed56cd47ca739b43e892707ae9a13790a486a3173be063681ca67d2262/contourpy-1.3.2-cp313-cp313-win_amd64.whl", hash = "sha256:e1578f7eafce927b168752ed7e22646dad6cd9bca673c60bff55889fa236ebf9", size = 223096, upload-time = "2025-04-15T17:44:48.194Z" }, - { url = "https://files.pythonhosted.org/packages/54/4c/e76fe2a03014a7c767d79ea35c86a747e9325537a8b7627e0e5b3ba266b4/contourpy-1.3.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0475b1f6604896bc7c53bb070e355e9321e1bc0d381735421a2d2068ec56531f", size = 285090, upload-time = "2025-04-15T17:43:34.084Z" }, - { url = "https://files.pythonhosted.org/packages/7b/e2/5aba47debd55d668e00baf9651b721e7733975dc9fc27264a62b0dd26eb8/contourpy-1.3.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:c85bb486e9be652314bb5b9e2e3b0d1b2e643d5eec4992c0fbe8ac71775da739", size = 268643, upload-time = "2025-04-15T17:43:38.626Z" }, - { url = "https://files.pythonhosted.org/packages/a1/37/cd45f1f051fe6230f751cc5cdd2728bb3a203f5619510ef11e732109593c/contourpy-1.3.2-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:745b57db7758f3ffc05a10254edd3182a2a83402a89c00957a8e8a22f5582823", size = 310443, upload-time = "2025-04-15T17:43:44.522Z" }, - { url = "https://files.pythonhosted.org/packages/8b/a2/36ea6140c306c9ff6dd38e3bcec80b3b018474ef4d17eb68ceecd26675f4/contourpy-1.3.2-cp313-cp313t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:970e9173dbd7eba9b4e01aab19215a48ee5dd3f43cef736eebde064a171f89a5", size = 349865, upload-time = "2025-04-15T17:43:49.545Z" }, - { url = "https://files.pythonhosted.org/packages/95/b7/2fc76bc539693180488f7b6cc518da7acbbb9e3b931fd9280504128bf956/contourpy-1.3.2-cp313-cp313t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c6c4639a9c22230276b7bffb6a850dfc8258a2521305e1faefe804d006b2e532", size = 321162, upload-time = "2025-04-15T17:43:54.203Z" }, - { url = "https://files.pythonhosted.org/packages/f4/10/76d4f778458b0aa83f96e59d65ece72a060bacb20cfbee46cf6cd5ceba41/contourpy-1.3.2-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cc829960f34ba36aad4302e78eabf3ef16a3a100863f0d4eeddf30e8a485a03b", size = 327355, upload-time = "2025-04-15T17:44:01.025Z" }, - { url = "https://files.pythonhosted.org/packages/43/a3/10cf483ea683f9f8ab096c24bad3cce20e0d1dd9a4baa0e2093c1c962d9d/contourpy-1.3.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:d32530b534e986374fc19eaa77fcb87e8a99e5431499949b828312bdcd20ac52", size = 1307935, upload-time = "2025-04-15T17:44:17.322Z" }, - { url = "https://files.pythonhosted.org/packages/78/73/69dd9a024444489e22d86108e7b913f3528f56cfc312b5c5727a44188471/contourpy-1.3.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:e298e7e70cf4eb179cc1077be1c725b5fd131ebc81181bf0c03525c8abc297fd", size = 1372168, upload-time = "2025-04-15T17:44:33.43Z" }, - { url = "https://files.pythonhosted.org/packages/0f/1b/96d586ccf1b1a9d2004dd519b25fbf104a11589abfd05484ff12199cca21/contourpy-1.3.2-cp313-cp313t-win32.whl", hash = "sha256:d0e589ae0d55204991450bb5c23f571c64fe43adaa53f93fc902a84c96f52fe1", size = 189550, upload-time = "2025-04-15T17:44:37.092Z" }, - { url = "https://files.pythonhosted.org/packages/b0/e6/6000d0094e8a5e32ad62591c8609e269febb6e4db83a1c75ff8868b42731/contourpy-1.3.2-cp313-cp313t-win_amd64.whl", hash = "sha256:78e9253c3de756b3f6a5174d024c4835acd59eb3f8e2ca13e775dbffe1558f69", size = 238214, upload-time = "2025-04-15T17:44:40.827Z" }, { url = "https://files.pythonhosted.org/packages/33/05/b26e3c6ecc05f349ee0013f0bb850a761016d89cec528a98193a48c34033/contourpy-1.3.2-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:fd93cc7f3139b6dd7aab2f26a90dde0aa9fc264dbf70f6740d498a70b860b82c", size = 265681, upload-time = "2025-04-15T17:44:59.314Z" }, { url = "https://files.pythonhosted.org/packages/2b/25/ac07d6ad12affa7d1ffed11b77417d0a6308170f44ff20fa1d5aa6333f03/contourpy-1.3.2-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:107ba8a6a7eec58bb475329e6d3b95deba9440667c4d62b9b6063942b61d7f16", size = 315101, upload-time = "2025-04-15T17:45:04.165Z" }, { url = "https://files.pythonhosted.org/packages/8f/4d/5bb3192bbe9d3f27e3061a6a8e7733c9120e203cb8515767d30973f71030/contourpy-1.3.2-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:ded1706ed0c1049224531b81128efbd5084598f18d8a2d9efae833edbd2b40ad", size = 220599, upload-time = "2025-04-15T17:45:08.456Z" }, @@ -709,28 +624,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cf/8f/5847f44a7fddf859704217a99a23a4f6417b10e5ab1256a179264561540e/contourpy-1.3.3-cp312-cp312-win32.whl", hash = "sha256:023b44101dfe49d7d53932be418477dba359649246075c996866106da069af69", size = 185018, upload-time = "2025-07-26T12:01:35.64Z" }, { url = "https://files.pythonhosted.org/packages/19/e8/6026ed58a64563186a9ee3f29f41261fd1828f527dd93d33b60feca63352/contourpy-1.3.3-cp312-cp312-win_amd64.whl", hash = "sha256:8153b8bfc11e1e4d75bcb0bff1db232f9e10b274e0929de9d608027e0d34ff8b", size = 226567, upload-time = "2025-07-26T12:01:36.804Z" }, { url = "https://files.pythonhosted.org/packages/d1/e2/f05240d2c39a1ed228d8328a78b6f44cd695f7ef47beb3e684cf93604f86/contourpy-1.3.3-cp312-cp312-win_arm64.whl", hash = "sha256:07ce5ed73ecdc4a03ffe3e1b3e3c1166db35ae7584be76f65dbbe28a7791b0cc", size = 193655, upload-time = "2025-07-26T12:01:37.999Z" }, - { url = "https://files.pythonhosted.org/packages/68/35/0167aad910bbdb9599272bd96d01a9ec6852f36b9455cf2ca67bd4cc2d23/contourpy-1.3.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:177fb367556747a686509d6fef71d221a4b198a3905fe824430e5ea0fda54eb5", size = 293257, upload-time = "2025-07-26T12:01:39.367Z" }, - { url = "https://files.pythonhosted.org/packages/96/e4/7adcd9c8362745b2210728f209bfbcf7d91ba868a2c5f40d8b58f54c509b/contourpy-1.3.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d002b6f00d73d69333dac9d0b8d5e84d9724ff9ef044fd63c5986e62b7c9e1b1", size = 274034, upload-time = "2025-07-26T12:01:40.645Z" }, - { url = "https://files.pythonhosted.org/packages/73/23/90e31ceeed1de63058a02cb04b12f2de4b40e3bef5e082a7c18d9c8ae281/contourpy-1.3.3-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:348ac1f5d4f1d66d3322420f01d42e43122f43616e0f194fc1c9f5d830c5b286", size = 334672, upload-time = "2025-07-26T12:01:41.942Z" }, - { url = "https://files.pythonhosted.org/packages/ed/93/b43d8acbe67392e659e1d984700e79eb67e2acb2bd7f62012b583a7f1b55/contourpy-1.3.3-cp313-cp313-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:655456777ff65c2c548b7c454af9c6f33f16c8884f11083244b5819cc214f1b5", size = 381234, upload-time = "2025-07-26T12:01:43.499Z" }, - { url = "https://files.pythonhosted.org/packages/46/3b/bec82a3ea06f66711520f75a40c8fc0b113b2a75edb36aa633eb11c4f50f/contourpy-1.3.3-cp313-cp313-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:644a6853d15b2512d67881586bd03f462c7ab755db95f16f14d7e238f2852c67", size = 385169, upload-time = "2025-07-26T12:01:45.219Z" }, - { url = "https://files.pythonhosted.org/packages/4b/32/e0f13a1c5b0f8572d0ec6ae2f6c677b7991fafd95da523159c19eff0696a/contourpy-1.3.3-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4debd64f124ca62069f313a9cb86656ff087786016d76927ae2cf37846b006c9", size = 362859, upload-time = "2025-07-26T12:01:46.519Z" }, - { url = "https://files.pythonhosted.org/packages/33/71/e2a7945b7de4e58af42d708a219f3b2f4cff7386e6b6ab0a0fa0033c49a9/contourpy-1.3.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a15459b0f4615b00bbd1e91f1b9e19b7e63aea7483d03d804186f278c0af2659", size = 1332062, upload-time = "2025-07-26T12:01:48.964Z" }, - { url = "https://files.pythonhosted.org/packages/12/fc/4e87ac754220ccc0e807284f88e943d6d43b43843614f0a8afa469801db0/contourpy-1.3.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:ca0fdcd73925568ca027e0b17ab07aad764be4706d0a925b89227e447d9737b7", size = 1403932, upload-time = "2025-07-26T12:01:51.979Z" }, - { url = "https://files.pythonhosted.org/packages/a6/2e/adc197a37443f934594112222ac1aa7dc9a98faf9c3842884df9a9d8751d/contourpy-1.3.3-cp313-cp313-win32.whl", hash = "sha256:b20c7c9a3bf701366556e1b1984ed2d0cedf999903c51311417cf5f591d8c78d", size = 185024, upload-time = "2025-07-26T12:01:53.245Z" }, - { url = "https://files.pythonhosted.org/packages/18/0b/0098c214843213759692cc638fce7de5c289200a830e5035d1791d7a2338/contourpy-1.3.3-cp313-cp313-win_amd64.whl", hash = "sha256:1cadd8b8969f060ba45ed7c1b714fe69185812ab43bd6b86a9123fe8f99c3263", size = 226578, upload-time = "2025-07-26T12:01:54.422Z" }, - { url = "https://files.pythonhosted.org/packages/8a/9a/2f6024a0c5995243cd63afdeb3651c984f0d2bc727fd98066d40e141ad73/contourpy-1.3.3-cp313-cp313-win_arm64.whl", hash = "sha256:fd914713266421b7536de2bfa8181aa8c699432b6763a0ea64195ebe28bff6a9", size = 193524, upload-time = "2025-07-26T12:01:55.73Z" }, - { url = "https://files.pythonhosted.org/packages/c0/b3/f8a1a86bd3298513f500e5b1f5fd92b69896449f6cab6a146a5d52715479/contourpy-1.3.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:88df9880d507169449d434c293467418b9f6cbe82edd19284aa0409e7fdb933d", size = 306730, upload-time = "2025-07-26T12:01:57.051Z" }, - { url = "https://files.pythonhosted.org/packages/3f/11/4780db94ae62fc0c2053909b65dc3246bd7cecfc4f8a20d957ad43aa4ad8/contourpy-1.3.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:d06bb1f751ba5d417047db62bca3c8fde202b8c11fb50742ab3ab962c81e8216", size = 287897, upload-time = "2025-07-26T12:01:58.663Z" }, - { url = "https://files.pythonhosted.org/packages/ae/15/e59f5f3ffdd6f3d4daa3e47114c53daabcb18574a26c21f03dc9e4e42ff0/contourpy-1.3.3-cp313-cp313t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e4e6b05a45525357e382909a4c1600444e2a45b4795163d3b22669285591c1ae", size = 326751, upload-time = "2025-07-26T12:02:00.343Z" }, - { url = "https://files.pythonhosted.org/packages/0f/81/03b45cfad088e4770b1dcf72ea78d3802d04200009fb364d18a493857210/contourpy-1.3.3-cp313-cp313t-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ab3074b48c4e2cf1a960e6bbeb7f04566bf36b1861d5c9d4d8ac04b82e38ba20", size = 375486, upload-time = "2025-07-26T12:02:02.128Z" }, - { url = "https://files.pythonhosted.org/packages/0c/ba/49923366492ffbdd4486e970d421b289a670ae8cf539c1ea9a09822b371a/contourpy-1.3.3-cp313-cp313t-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:6c3d53c796f8647d6deb1abe867daeb66dcc8a97e8455efa729516b997b8ed99", size = 388106, upload-time = "2025-07-26T12:02:03.615Z" }, - { url = "https://files.pythonhosted.org/packages/9f/52/5b00ea89525f8f143651f9f03a0df371d3cbd2fccd21ca9b768c7a6500c2/contourpy-1.3.3-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:50ed930df7289ff2a8d7afeb9603f8289e5704755c7e5c3bbd929c90c817164b", size = 352548, upload-time = "2025-07-26T12:02:05.165Z" }, - { url = "https://files.pythonhosted.org/packages/32/1d/a209ec1a3a3452d490f6b14dd92e72280c99ae3d1e73da74f8277d4ee08f/contourpy-1.3.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:4feffb6537d64b84877da813a5c30f1422ea5739566abf0bd18065ac040e120a", size = 1322297, upload-time = "2025-07-26T12:02:07.379Z" }, - { url = "https://files.pythonhosted.org/packages/bc/9e/46f0e8ebdd884ca0e8877e46a3f4e633f6c9c8c4f3f6e72be3fe075994aa/contourpy-1.3.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:2b7e9480ffe2b0cd2e787e4df64270e3a0440d9db8dc823312e2c940c167df7e", size = 1391023, upload-time = "2025-07-26T12:02:10.171Z" }, - { url = "https://files.pythonhosted.org/packages/b9/70/f308384a3ae9cd2209e0849f33c913f658d3326900d0ff5d378d6a1422d2/contourpy-1.3.3-cp313-cp313t-win32.whl", hash = "sha256:283edd842a01e3dcd435b1c5116798d661378d83d36d337b8dde1d16a5fc9ba3", size = 196157, upload-time = "2025-07-26T12:02:11.488Z" }, - { url = "https://files.pythonhosted.org/packages/b2/dd/880f890a6663b84d9e34a6f88cded89d78f0091e0045a284427cb6b18521/contourpy-1.3.3-cp313-cp313t-win_amd64.whl", hash = "sha256:87acf5963fc2b34825e5b6b048f40e3635dd547f590b04d2ab317c2619ef7ae8", size = 240570, upload-time = "2025-07-26T12:02:12.754Z" }, - { url = "https://files.pythonhosted.org/packages/80/99/2adc7d8ffead633234817ef8e9a87115c8a11927a94478f6bb3d3f4d4f7d/contourpy-1.3.3-cp313-cp313t-win_arm64.whl", hash = "sha256:3c30273eb2a55024ff31ba7d052dde990d7d8e5450f4bbb6e913558b3d6c2301", size = 199713, upload-time = "2025-07-26T12:02:14.4Z" }, { url = "https://files.pythonhosted.org/packages/a5/29/8dcfe16f0107943fa92388c23f6e05cff0ba58058c4c95b00280d4c75a14/contourpy-1.3.3-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:cd5dfcaeb10f7b7f9dc8941717c6c2ade08f587be2226222c12b25f0483ed497", size = 278809, upload-time = "2025-07-26T12:02:52.74Z" }, { url = "https://files.pythonhosted.org/packages/85/a9/8b37ef4f7dafeb335daee3c8254645ef5725be4d9c6aa70b50ec46ef2f7e/contourpy-1.3.3-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:0c1fc238306b35f246d61a1d416a627348b5cf0648648a031e14bb8705fcdfe8", size = 261593, upload-time = "2025-07-26T12:02:54.037Z" }, { url = "https://files.pythonhosted.org/packages/0a/59/ebfb8c677c75605cc27f7122c90313fd2f375ff3c8d19a1694bda74aaa63/contourpy-1.3.3-pp311-pypy311_pp73-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:70f9aad7de812d6541d29d2bbf8feb22ff7e1c299523db288004e3157ff4674e", size = 302202, upload-time = "2025-07-26T12:02:55.947Z" }, @@ -804,12 +697,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/32/45/557d4ed1fa54f0c7db8aee083229f624990d69f7d00f55477eed5c7e169a/cuda_bindings-12.9.7-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0666d3c082ef8f4b2d670950589373550e9f3bf564d635dd883f24a0b40402ff", size = 7071026, upload-time = "2026-05-27T18:44:13.356Z" }, { url = "https://files.pythonhosted.org/packages/91/97/e3c6e58ece26a053419ba0a18444b5443cfc64451bbf37f84e8143b8bdca/cuda_bindings-12.9.7-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4c7ef48c5e13ae90f3b2ecfb72f8e99ac43c8f4c43e67e1325b8aae331453687", size = 7611059, upload-time = "2026-05-27T18:44:15.252Z" }, { url = "https://files.pythonhosted.org/packages/6d/39/afaa3de4d491a55af8961081e0b69c08d51bfbe471c359a7bddb4a28ca41/cuda_bindings-12.9.7-cp312-cp312-win_amd64.whl", hash = "sha256:3c089aaf4f5f570ec50244c68f5a2b00a2c9a8e01e04219fd2e36e340be0d88b", size = 7400841, upload-time = "2026-05-27T18:44:17.164Z" }, - { url = "https://files.pythonhosted.org/packages/eb/7b/f1575e41e1a17dc2f2a408b2e8e864c9324e41e3e23f6401e5efc54c152a/cuda_bindings-12.9.7-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:266379e4942051f544a8e7ea1a30ead8d7e8199b6b30fcdc8917cae2bf614e61", size = 6978549, upload-time = "2026-05-27T18:44:18.839Z" }, - { url = "https://files.pythonhosted.org/packages/9d/dc/62d62eb4f91eb721bcf46da51b13e9872ccd8fa7e60eb8ba7b7baeac72c6/cuda_bindings-12.9.7-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:59cf4a37b0d662ba15037c9ceebe1a306ebf2c01a8235a09be13cd07094fdb74", size = 7457675, upload-time = "2026-05-27T18:44:20.637Z" }, - { url = "https://files.pythonhosted.org/packages/43/b2/753fe88151001d0dc23f56a8e119fe06b991b0d1a885fa02f9852b12f523/cuda_bindings-12.9.7-cp313-cp313-win_amd64.whl", hash = "sha256:5bd89dcb78475a6d8a4620ea94b74edf0cbbeacee6d1622d8f94452c1e8d3f15", size = 7360097, upload-time = "2026-05-27T18:44:22.405Z" }, - { url = "https://files.pythonhosted.org/packages/f9/77/94d9b85f26add6fe9c9cb7c4ec3b96bc598f7ea5cfbd7490cc0a36adf5be/cuda_bindings-12.9.7-cp313-cp313t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2dbcd4801954eb3508f4dc2fa0d0c8eb93eb3f45326fd61be2731418c371e7a0", size = 6870886, upload-time = "2026-05-27T18:44:24.164Z" }, - { url = "https://files.pythonhosted.org/packages/04/dd/3ec34b569e1b990b11276feba306bf8f446656cc38e8ed0f49b5facfeffa/cuda_bindings-12.9.7-cp313-cp313t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3747ea132642416786a8e31bf229032df3a7856911ae5426a7be53d032df183d", size = 7345663, upload-time = "2026-05-27T18:44:26.333Z" }, - { url = "https://files.pythonhosted.org/packages/b8/c8/d79a20ba396e7ab2dfdd4b72b62356972b25b88aee2ded49a70c797ddea1/cuda_bindings-12.9.7-cp313-cp313t-win_amd64.whl", hash = "sha256:64f7ade7a7a3b69001489753acc21706d9dbda32db8deb68a767a0a0aab30b68", size = 7780136, upload-time = "2026-05-27T18:44:28.121Z" }, ] [[package]] @@ -834,9 +721,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ce/67/5e7dba1ba576dd73da5dee894ca076ca5e959450dfff66d6d510a255d1f7/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7855c4868aabc0cfae28abbe83d56734bdfbd08f08fc234ac1912a12858bf49", size = 6025351, upload-time = "2026-05-29T23:11:49.685Z" }, { url = "https://files.pythonhosted.org/packages/39/2a/6d2e9047d1fb243dbaa364b01e0297534b9ed7fd27dba1c9f361519cf69b/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e32d08f71ebcdf00f0f41eab2eb37e8da94c8ed411cc9f7f7a019ce6b34abe3a", size = 6657965, upload-time = "2026-05-29T23:11:52.227Z" }, { url = "https://files.pythonhosted.org/packages/7c/95/872a0392122f1fb43fcb06869790ef3171f37beee9f7db8f441739113570/cuda_bindings-13.3.1-cp312-cp312-win_amd64.whl", hash = "sha256:b134dd8c5c66ae4c4ad814f7aee88fd215353c077010cbc47e3b55ed35ec9eff", size = 5875099, upload-time = "2026-05-29T23:11:54.635Z" }, - { url = "https://files.pythonhosted.org/packages/cc/6e/2394f8163360f8391f8f1b7e72d300a82724edb81a7b7084c799fbd4c91f/cuda_bindings-13.3.1-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9efb21c1ee64981e184b9e0ba5eb3179e5ba3d4b51665a6cb52b8ef3d01a7cbf", size = 5920504, upload-time = "2026-05-29T23:11:56.883Z" }, - { url = "https://files.pythonhosted.org/packages/34/c2/ef9b6a63f7dc432712a462c816662e662e00d38caa9b861c8c2588195d03/cuda_bindings-13.3.1-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2732904099e0a4d4db774a5fc6d91ee95fae065b4d2ecabb4968c5fe2406c9d7", size = 6476660, upload-time = "2026-05-29T23:11:59.188Z" }, - { url = "https://files.pythonhosted.org/packages/0c/2f/6a0dd496550c6fafbf6aeb1bf40242eeabb2fd138a43892aabb4be8224c2/cuda_bindings-13.3.1-cp313-cp313-win_amd64.whl", hash = "sha256:18c8c167c8907b8f02531ca810534315c458dabef31f7965095619bf647b9202", size = 5830027, upload-time = "2026-05-29T23:12:01.205Z" }, ] [[package]] @@ -1574,13 +1458,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/de/72/8f5d083ef3ea86263a49296a4247343b111077b479d172b66f1d2971cd28/flip_evaluator-1.7-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:db5bd82d93c5be24e10134138cb54942fdce93fe771412a8b389e29e378cb59e", size = 906634, upload-time = "2025-11-07T15:33:16.866Z" }, { url = "https://files.pythonhosted.org/packages/53/bb/9a85a283efca7f57c6ba2c6457d299e4021984315e2ecf5a0d5fb13cc106/flip_evaluator-1.7-cp312-cp312-win32.whl", hash = "sha256:8a734f77b2f820110e67c78ad103cf62888cf26d9602e5fc0469551d9b81f0b8", size = 188882, upload-time = "2025-11-07T15:33:18.406Z" }, { url = "https://files.pythonhosted.org/packages/95/a6/fe3e220fc50783682662cc5fd9d1f86a4e9ee47d229c3079915bc226d0c9/flip_evaluator-1.7-cp312-cp312-win_amd64.whl", hash = "sha256:dc685bfc5acaab99adeb878b261c819637b2a0638bb7cbcffe9eb6a3d988c52b", size = 214363, upload-time = "2025-11-07T15:33:19.424Z" }, - { url = "https://files.pythonhosted.org/packages/25/86/e328522798cd53908ee875b69a3272306cc6229377fb2f70517fcc6820d9/flip_evaluator-1.7-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:76e689402877598b5f4a42d57d7109a4339381cd104317f361cefb103d0851b3", size = 189925, upload-time = "2025-11-07T15:33:20.51Z" }, - { url = "https://files.pythonhosted.org/packages/b2/bf/6d62779ed195adfda082f3cd7a91ebbe3a1cc9c1e936f6fd26283a248bdb/flip_evaluator-1.7-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3f37f778e02f15607450f11f6d1b21a9b8c6817415cdc3441c3a7c55c33b5664", size = 444745, upload-time = "2025-11-07T15:33:21.498Z" }, - { url = "https://files.pythonhosted.org/packages/a5/ed/fc88540c25b08aba2458835c26c8db54c6ea9c1dee058734492efe96eb1e/flip_evaluator-1.7-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d63eff9eaf9a68a6f9c30a3714c8a1a83fb6182f18df2fbc9bc361634a6691f7", size = 415164, upload-time = "2025-11-07T15:33:22.531Z" }, - { url = "https://files.pythonhosted.org/packages/3b/4d/18923dc2e5262d2e34cc09aa51d3d2b977855922afc8ce333df2a10e9fc8/flip_evaluator-1.7-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:1308484dfd89a90fef4d87db7ba26c18fbb4f7df36c7ed20c9369e1cb5b19689", size = 976486, upload-time = "2025-11-07T15:33:25.133Z" }, - { url = "https://files.pythonhosted.org/packages/9d/27/f552b286648c40bb5820b19c4a8200965abace4e17554ce098b33d7e4f28/flip_evaluator-1.7-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8fc68d919692ba147c6a3c6f8b2c41aa7d8aa578cad6ee86884019d90c99df66", size = 906636, upload-time = "2025-11-07T15:33:26.217Z" }, - { url = "https://files.pythonhosted.org/packages/da/e8/e8b84c26cabcfe939c9909b5414090506496015f3e5937fade4557467011/flip_evaluator-1.7-cp313-cp313-win32.whl", hash = "sha256:3735232d08f6128ff743e8405bf6c7666efbee6e0a6ee1b5e2f5aa0b64f9b090", size = 188882, upload-time = "2025-11-07T15:33:27.33Z" }, - { url = "https://files.pythonhosted.org/packages/b0/00/820a81e5d047298a622ca0538fede02b6ff09fbe85ad4bfa82649a92d919/flip_evaluator-1.7-cp313-cp313-win_amd64.whl", hash = "sha256:2081f715ea8190a5f58bc578b0cccd36a9a3f3cf78ad13fe6f86e5a75615bda7", size = 214374, upload-time = "2025-11-07T15:33:28.226Z" }, ] [[package]] @@ -1613,14 +1490,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f2/00/3bbab338c07c71fa56269953845e92c951a61457bbbb0f1022551ea266d9/fonttools-4.63.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:27fdc65af8da6f88b9c6121c47a464cbe359fcfff7ff6fc2d37a1f395d755b78", size = 5133598, upload-time = "2026-05-14T12:03:25.168Z" }, { url = "https://files.pythonhosted.org/packages/62/f2/aa27c7f98db5b064883dadcc5283947e81e034de42e22a33675878d98b54/fonttools-4.63.0-cp312-cp312-win32.whl", hash = "sha256:af2fd1664d00a397d75f806985ddb36282091c2131a73a6485c23b4a34722263", size = 2292575, upload-time = "2026-05-14T12:03:27.496Z" }, { url = "https://files.pythonhosted.org/packages/87/36/cccb9bc2a6ab63d1b2980374f0dca72ce95ae267c9b4cfe77455bb70d0d4/fonttools-4.63.0-cp312-cp312-win_amd64.whl", hash = "sha256:59ac449f8cca9b4ffa08d2e7bbadad87ce710d69d1eda5c3c1ce579baa987272", size = 2343211, upload-time = "2026-05-14T12:03:30.057Z" }, - { url = "https://files.pythonhosted.org/packages/0f/8d/d8fec3dcde2963f8c908fb315e5ff2cd0ac34f82394bbbf73a2aa5145ce3/fonttools-4.63.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:cd7e9857e5e63738b9d9fd707bc1f59c8b09e5177726d23664db393c59bb08bd", size = 2876062, upload-time = "2026-05-14T12:03:32.554Z" }, - { url = "https://files.pythonhosted.org/packages/ef/71/d935dc54e4ff121bfdd11e08702db63a7e6f25af21d8a3d7b7212df53641/fonttools-4.63.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c2a2a42198b696a6f48fad91709afb55176e66a5e566131219dba372fb7f8c59", size = 2424594, upload-time = "2026-05-14T12:03:34.86Z" }, - { url = "https://files.pythonhosted.org/packages/8e/40/e76320afa1df918e146155ef239b1719ee266092e96f5423bfd075affba1/fonttools-4.63.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1e874792a8212b44583ea02189d9e693906b2f78b261f372f95d6c563210ac1d", size = 5024840, upload-time = "2026-05-14T12:03:36.745Z" }, - { url = "https://files.pythonhosted.org/packages/ce/36/0b805d8c485f872f65a509cbe3b58a5d0d17bee855333b54a150c79d3061/fonttools-4.63.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:22135da48a348785c5e2d5d2d9d6bec5ed44adacbaeb9db12d9493bf6c6bfa68", size = 4975801, upload-time = "2026-05-14T12:03:38.833Z" }, - { url = "https://files.pythonhosted.org/packages/c8/26/2cee03d0aa083ab022da5c07aff9ed3f689da1defb81ad6917c9627896da/fonttools-4.63.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ccf41f2efdf56994d22d73bef4ced1052161958169428d06ba9724ea9e9a64be", size = 4965009, upload-time = "2026-05-14T12:03:41.494Z" }, - { url = "https://files.pythonhosted.org/packages/7e/48/cc4b66d9058c0d0982c833fad10127c4b0e9324606aafa41382295ca4102/fonttools-4.63.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:9ced0bd02ac751dd6319b0da88aaef24414e3b0dbc32bb4f24944821a3741a27", size = 5105892, upload-time = "2026-05-14T12:03:43.525Z" }, - { url = "https://files.pythonhosted.org/packages/d8/1f/a98a30a814b9ddef3a2e706025f90b9e0bc94890e6cb15254bc86547d11a/fonttools-4.63.0-cp313-cp313-win32.whl", hash = "sha256:85be818f5506e8a7753153def2c9550178f0ecae6a47b5e0e8dbb23f7cc90380", size = 2291313, upload-time = "2026-05-14T12:03:45.594Z" }, - { url = "https://files.pythonhosted.org/packages/92/46/5177b01f3b4abfdd4409f31cca4ab279c9343a26efbe9ec78c97fc612e02/fonttools-4.63.0-cp313-cp313-win_amd64.whl", hash = "sha256:ba04cb5891d4c0c21b6da95eda8d7b090021508a294fff33464fc7d241e0856b", size = 2342299, upload-time = "2026-05-14T12:03:47.414Z" }, { url = "https://files.pythonhosted.org/packages/2c/47/c99d5268f354002ce80f8d029cd9d7d872969da1de8b93d32de4dc56d6f4/fonttools-4.63.0-py3-none-any.whl", hash = "sha256:445af2eab030a16b9171ea8bdda7ebf7d96bda2df88ee182a464252f6e05e20d", size = 1164562, upload-time = "2026-05-14T12:04:29.092Z" }, ] @@ -1678,38 +1547,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/66/bb/852b9d6db2fa40be96f29c0d1205c306288f0684df8fd26ca1951d461a56/frozenlist-1.8.0-cp312-cp312-win32.whl", hash = "sha256:433403ae80709741ce34038da08511d4a77062aa924baf411ef73d1146e74faf", size = 39985, upload-time = "2025-10-06T05:36:23.661Z" }, { url = "https://files.pythonhosted.org/packages/b8/af/38e51a553dd66eb064cdf193841f16f077585d4d28394c2fa6235cb41765/frozenlist-1.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:34187385b08f866104f0c0617404c8eb08165ab1272e884abc89c112e9c00746", size = 44591, upload-time = "2025-10-06T05:36:24.958Z" }, { url = "https://files.pythonhosted.org/packages/a7/06/1dc65480ab147339fecc70797e9c2f69d9cea9cf38934ce08df070fdb9cb/frozenlist-1.8.0-cp312-cp312-win_arm64.whl", hash = "sha256:fe3c58d2f5db5fbd18c2987cba06d51b0529f52bc3a6cdc33d3f4eab725104bd", size = 40102, upload-time = "2025-10-06T05:36:26.333Z" }, - { url = "https://files.pythonhosted.org/packages/2d/40/0832c31a37d60f60ed79e9dfb5a92e1e2af4f40a16a29abcc7992af9edff/frozenlist-1.8.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8d92f1a84bb12d9e56f818b3a746f3efba93c1b63c8387a73dde655e1e42282a", size = 85717, upload-time = "2025-10-06T05:36:27.341Z" }, - { url = "https://files.pythonhosted.org/packages/30/ba/b0b3de23f40bc55a7057bd38434e25c34fa48e17f20ee273bbde5e0650f3/frozenlist-1.8.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:96153e77a591c8adc2ee805756c61f59fef4cf4073a9275ee86fe8cba41241f7", size = 49651, upload-time = "2025-10-06T05:36:28.855Z" }, - { url = "https://files.pythonhosted.org/packages/0c/ab/6e5080ee374f875296c4243c381bbdef97a9ac39c6e3ce1d5f7d42cb78d6/frozenlist-1.8.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f21f00a91358803399890ab167098c131ec2ddd5f8f5fd5fe9c9f2c6fcd91e40", size = 49417, upload-time = "2025-10-06T05:36:29.877Z" }, - { url = "https://files.pythonhosted.org/packages/d5/4e/e4691508f9477ce67da2015d8c00acd751e6287739123113a9fca6f1604e/frozenlist-1.8.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:fb30f9626572a76dfe4293c7194a09fb1fe93ba94c7d4f720dfae3b646b45027", size = 234391, upload-time = "2025-10-06T05:36:31.301Z" }, - { url = "https://files.pythonhosted.org/packages/40/76/c202df58e3acdf12969a7895fd6f3bc016c642e6726aa63bd3025e0fc71c/frozenlist-1.8.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:eaa352d7047a31d87dafcacbabe89df0aa506abb5b1b85a2fb91bc3faa02d822", size = 233048, upload-time = "2025-10-06T05:36:32.531Z" }, - { url = "https://files.pythonhosted.org/packages/f9/c0/8746afb90f17b73ca5979c7a3958116e105ff796e718575175319b5bb4ce/frozenlist-1.8.0-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:03ae967b4e297f58f8c774c7eabcce57fe3c2434817d4385c50661845a058121", size = 226549, upload-time = "2025-10-06T05:36:33.706Z" }, - { url = "https://files.pythonhosted.org/packages/7e/eb/4c7eefc718ff72f9b6c4893291abaae5fbc0c82226a32dcd8ef4f7a5dbef/frozenlist-1.8.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f6292f1de555ffcc675941d65fffffb0a5bcd992905015f85d0592201793e0e5", size = 239833, upload-time = "2025-10-06T05:36:34.947Z" }, - { url = "https://files.pythonhosted.org/packages/c2/4e/e5c02187cf704224f8b21bee886f3d713ca379535f16893233b9d672ea71/frozenlist-1.8.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:29548f9b5b5e3460ce7378144c3010363d8035cea44bc0bf02d57f5a685e084e", size = 245363, upload-time = "2025-10-06T05:36:36.534Z" }, - { url = "https://files.pythonhosted.org/packages/1f/96/cb85ec608464472e82ad37a17f844889c36100eed57bea094518bf270692/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ec3cc8c5d4084591b4237c0a272cc4f50a5b03396a47d9caaf76f5d7b38a4f11", size = 229314, upload-time = "2025-10-06T05:36:38.582Z" }, - { url = "https://files.pythonhosted.org/packages/5d/6f/4ae69c550e4cee66b57887daeebe006fe985917c01d0fff9caab9883f6d0/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:517279f58009d0b1f2e7c1b130b377a349405da3f7621ed6bfae50b10adf20c1", size = 243365, upload-time = "2025-10-06T05:36:40.152Z" }, - { url = "https://files.pythonhosted.org/packages/7a/58/afd56de246cf11780a40a2c28dc7cbabbf06337cc8ddb1c780a2d97e88d8/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:db1e72ede2d0d7ccb213f218df6a078a9c09a7de257c2fe8fcef16d5925230b1", size = 237763, upload-time = "2025-10-06T05:36:41.355Z" }, - { url = "https://files.pythonhosted.org/packages/cb/36/cdfaf6ed42e2644740d4a10452d8e97fa1c062e2a8006e4b09f1b5fd7d63/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:b4dec9482a65c54a5044486847b8a66bf10c9cb4926d42927ec4e8fd5db7fed8", size = 240110, upload-time = "2025-10-06T05:36:42.716Z" }, - { url = "https://files.pythonhosted.org/packages/03/a8/9ea226fbefad669f11b52e864c55f0bd57d3c8d7eb07e9f2e9a0b39502e1/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:21900c48ae04d13d416f0e1e0c4d81f7931f73a9dfa0b7a8746fb2fe7dd970ed", size = 233717, upload-time = "2025-10-06T05:36:44.251Z" }, - { url = "https://files.pythonhosted.org/packages/1e/0b/1b5531611e83ba7d13ccc9988967ea1b51186af64c42b7a7af465dcc9568/frozenlist-1.8.0-cp313-cp313-win32.whl", hash = "sha256:8b7b94a067d1c504ee0b16def57ad5738701e4ba10cec90529f13fa03c833496", size = 39628, upload-time = "2025-10-06T05:36:45.423Z" }, - { url = "https://files.pythonhosted.org/packages/d8/cf/174c91dbc9cc49bc7b7aab74d8b734e974d1faa8f191c74af9b7e80848e6/frozenlist-1.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:878be833caa6a3821caf85eb39c5ba92d28e85df26d57afb06b35b2efd937231", size = 43882, upload-time = "2025-10-06T05:36:46.796Z" }, - { url = "https://files.pythonhosted.org/packages/c1/17/502cd212cbfa96eb1388614fe39a3fc9ab87dbbe042b66f97acb57474834/frozenlist-1.8.0-cp313-cp313-win_arm64.whl", hash = "sha256:44389d135b3ff43ba8cc89ff7f51f5a0bb6b63d829c8300f79a2fe4fe61bcc62", size = 39676, upload-time = "2025-10-06T05:36:47.8Z" }, - { url = "https://files.pythonhosted.org/packages/d2/5c/3bbfaa920dfab09e76946a5d2833a7cbdf7b9b4a91c714666ac4855b88b4/frozenlist-1.8.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:e25ac20a2ef37e91c1b39938b591457666a0fa835c7783c3a8f33ea42870db94", size = 89235, upload-time = "2025-10-06T05:36:48.78Z" }, - { url = "https://files.pythonhosted.org/packages/d2/d6/f03961ef72166cec1687e84e8925838442b615bd0b8854b54923ce5b7b8a/frozenlist-1.8.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:07cdca25a91a4386d2e76ad992916a85038a9b97561bf7a3fd12d5d9ce31870c", size = 50742, upload-time = "2025-10-06T05:36:49.837Z" }, - { url = "https://files.pythonhosted.org/packages/1e/bb/a6d12b7ba4c3337667d0e421f7181c82dda448ce4e7ad7ecd249a16fa806/frozenlist-1.8.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:4e0c11f2cc6717e0a741f84a527c52616140741cd812a50422f83dc31749fb52", size = 51725, upload-time = "2025-10-06T05:36:50.851Z" }, - { url = "https://files.pythonhosted.org/packages/bc/71/d1fed0ffe2c2ccd70b43714c6cab0f4188f09f8a67a7914a6b46ee30f274/frozenlist-1.8.0-cp313-cp313t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:b3210649ee28062ea6099cfda39e147fa1bc039583c8ee4481cb7811e2448c51", size = 284533, upload-time = "2025-10-06T05:36:51.898Z" }, - { url = "https://files.pythonhosted.org/packages/c9/1f/fb1685a7b009d89f9bf78a42d94461bc06581f6e718c39344754a5d9bada/frozenlist-1.8.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:581ef5194c48035a7de2aefc72ac6539823bb71508189e5de01d60c9dcd5fa65", size = 292506, upload-time = "2025-10-06T05:36:53.101Z" }, - { url = "https://files.pythonhosted.org/packages/e6/3b/b991fe1612703f7e0d05c0cf734c1b77aaf7c7d321df4572e8d36e7048c8/frozenlist-1.8.0-cp313-cp313t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:3ef2d026f16a2b1866e1d86fc4e1291e1ed8a387b2c333809419a2f8b3a77b82", size = 274161, upload-time = "2025-10-06T05:36:54.309Z" }, - { url = "https://files.pythonhosted.org/packages/ca/ec/c5c618767bcdf66e88945ec0157d7f6c4a1322f1473392319b7a2501ded7/frozenlist-1.8.0-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5500ef82073f599ac84d888e3a8c1f77ac831183244bfd7f11eaa0289fb30714", size = 294676, upload-time = "2025-10-06T05:36:55.566Z" }, - { url = "https://files.pythonhosted.org/packages/7c/ce/3934758637d8f8a88d11f0585d6495ef54b2044ed6ec84492a91fa3b27aa/frozenlist-1.8.0-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:50066c3997d0091c411a66e710f4e11752251e6d2d73d70d8d5d4c76442a199d", size = 300638, upload-time = "2025-10-06T05:36:56.758Z" }, - { url = "https://files.pythonhosted.org/packages/fc/4f/a7e4d0d467298f42de4b41cbc7ddaf19d3cfeabaf9ff97c20c6c7ee409f9/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:5c1c8e78426e59b3f8005e9b19f6ff46e5845895adbde20ece9218319eca6506", size = 283067, upload-time = "2025-10-06T05:36:57.965Z" }, - { url = "https://files.pythonhosted.org/packages/dc/48/c7b163063d55a83772b268e6d1affb960771b0e203b632cfe09522d67ea5/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:eefdba20de0d938cec6a89bd4d70f346a03108a19b9df4248d3cf0d88f1b0f51", size = 292101, upload-time = "2025-10-06T05:36:59.237Z" }, - { url = "https://files.pythonhosted.org/packages/9f/d0/2366d3c4ecdc2fd391e0afa6e11500bfba0ea772764d631bbf82f0136c9d/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:cf253e0e1c3ceb4aaff6df637ce033ff6535fb8c70a764a8f46aafd3d6ab798e", size = 289901, upload-time = "2025-10-06T05:37:00.811Z" }, - { url = "https://files.pythonhosted.org/packages/b8/94/daff920e82c1b70e3618a2ac39fbc01ae3e2ff6124e80739ce5d71c9b920/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:032efa2674356903cd0261c4317a561a6850f3ac864a63fc1583147fb05a79b0", size = 289395, upload-time = "2025-10-06T05:37:02.115Z" }, - { url = "https://files.pythonhosted.org/packages/e3/20/bba307ab4235a09fdcd3cc5508dbabd17c4634a1af4b96e0f69bfe551ebd/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:6da155091429aeba16851ecb10a9104a108bcd32f6c1642867eadaee401c1c41", size = 283659, upload-time = "2025-10-06T05:37:03.711Z" }, - { url = "https://files.pythonhosted.org/packages/fd/00/04ca1c3a7a124b6de4f8a9a17cc2fcad138b4608e7a3fc5877804b8715d7/frozenlist-1.8.0-cp313-cp313t-win32.whl", hash = "sha256:0f96534f8bfebc1a394209427d0f8a63d343c9779cda6fc25e8e121b5fd8555b", size = 43492, upload-time = "2025-10-06T05:37:04.915Z" }, - { url = "https://files.pythonhosted.org/packages/59/5e/c69f733a86a94ab10f68e496dc6b7e8bc078ebb415281d5698313e3af3a1/frozenlist-1.8.0-cp313-cp313t-win_amd64.whl", hash = "sha256:5d63a068f978fc69421fb0e6eb91a9603187527c86b7cd3f534a5b77a592b888", size = 48034, upload-time = "2025-10-06T05:37:06.343Z" }, - { url = "https://files.pythonhosted.org/packages/16/6c/be9d79775d8abe79b05fa6d23da99ad6e7763a1d080fbae7290b286093fd/frozenlist-1.8.0-cp313-cp313t-win_arm64.whl", hash = "sha256:bf0a7e10b077bf5fb9380ad3ae8ce20ef919a6ad93b4552896419ac7e1d8e042", size = 41749, upload-time = "2025-10-06T05:37:07.431Z" }, { url = "https://files.pythonhosted.org/packages/9a/9a/e35b4a917281c0b8419d4207f4334c8e8c5dbf4f3f5f9ada73958d937dcc/frozenlist-1.8.0-py3-none-any.whl", hash = "sha256:0c18a16eab41e82c295618a77502e17b195883241c563b00f0aa5106fc4eaa0d", size = 13409, upload-time = "2025-10-06T05:38:16.721Z" }, ] @@ -1755,11 +1592,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ce/a9/a780cc66f86335a6019f557a8aaca8fbb970728f0efd2430d15ff1beae0e/google_crc32c-1.8.0-cp312-cp312-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:14f87e04d613dfa218d6135e81b78272c3b904e2a7053b841481b38a7d901411", size = 33364, upload-time = "2025-12-16T00:40:22.96Z" }, { url = "https://files.pythonhosted.org/packages/21/3f/3457ea803db0198c9aaca2dd373750972ce28a26f00544b6b85088811939/google_crc32c-1.8.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb5c869c2923d56cb0c8e6bcdd73c009c36ae39b652dbe46a05eb4ef0ad01454", size = 33740, upload-time = "2025-12-16T00:40:23.96Z" }, { url = "https://files.pythonhosted.org/packages/df/c0/87c2073e0c72515bb8733d4eef7b21548e8d189f094b5dad20b0ecaf64f6/google_crc32c-1.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:3cc0c8912038065eafa603b238abf252e204accab2a704c63b9e14837a854962", size = 34437, upload-time = "2025-12-16T00:35:21.395Z" }, - { url = "https://files.pythonhosted.org/packages/d1/db/000f15b41724589b0e7bc24bc7a8967898d8d3bc8caf64c513d91ef1f6c0/google_crc32c-1.8.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:3ebb04528e83b2634857f43f9bb8ef5b2bbe7f10f140daeb01b58f972d04736b", size = 31297, upload-time = "2025-12-16T00:23:20.709Z" }, - { url = "https://files.pythonhosted.org/packages/d7/0d/8ebed0c39c53a7e838e2a486da8abb0e52de135f1b376ae2f0b160eb4c1a/google_crc32c-1.8.0-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:450dc98429d3e33ed2926fc99ee81001928d63460f8538f21a5d6060912a8e27", size = 30867, upload-time = "2025-12-16T00:43:14.628Z" }, - { url = "https://files.pythonhosted.org/packages/ce/42/b468aec74a0354b34c8cbf748db20d6e350a68a2b0912e128cabee49806c/google_crc32c-1.8.0-cp313-cp313-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:3b9776774b24ba76831609ffbabce8cdf6fa2bd5e9df37b594221c7e333a81fa", size = 33344, upload-time = "2025-12-16T00:40:24.742Z" }, - { url = "https://files.pythonhosted.org/packages/1c/e8/b33784d6fc77fb5062a8a7854e43e1e618b87d5ddf610a88025e4de6226e/google_crc32c-1.8.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:89c17d53d75562edfff86679244830599ee0a48efc216200691de8b02ab6b2b8", size = 33694, upload-time = "2025-12-16T00:40:25.505Z" }, - { url = "https://files.pythonhosted.org/packages/92/b1/d3cbd4d988afb3d8e4db94ca953df429ed6db7282ed0e700d25e6c7bfc8d/google_crc32c-1.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:57a50a9035b75643996fbf224d6661e386c7162d1dfdab9bc4ca790947d1007f", size = 34435, upload-time = "2025-12-16T00:35:22.107Z" }, { url = "https://files.pythonhosted.org/packages/52/c5/c171e4d8c44fec1422d801a6d2e5d7ddabd733eeda505c79730ee9607f07/google_crc32c-1.8.0-pp311-pypy311_pp73-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:87fa445064e7db928226b2e6f0d5304ab4cd0339e664a4e9a25029f384d9bb93", size = 28615, upload-time = "2025-12-16T00:40:29.298Z" }, { url = "https://files.pythonhosted.org/packages/9c/97/7d75fe37a7a6ed171a2cf17117177e7aab7e6e0d115858741b41e9dd4254/google_crc32c-1.8.0-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f639065ea2042d5c034bf258a9f085eaa7af0cd250667c0635a3118e8f92c69c", size = 28800, upload-time = "2025-12-16T00:40:30.322Z" }, ] @@ -1803,16 +1635,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/25/7a/71437c7f3596e5246155c515852795a85a1a8d228190212432b13b97a95d/grpcio-1.81.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:8ea1936c26b99999b27479853039a7f34713f56c49375ad52b38535ec93a796c", size = 7849660, upload-time = "2026-06-11T12:45:40.627Z" }, { url = "https://files.pythonhosted.org/packages/65/40/7debc0da45d2efebafb82da75644be347497fe4ee250514b8cd3b86ae8bf/grpcio-1.81.1-cp312-cp312-win32.whl", hash = "sha256:a185a04039df6cae8648bc8ab6d6fde7bf94f7188ecf7828e76ac52eef1e41d6", size = 4185819, upload-time = "2026-06-11T12:45:43.027Z" }, { url = "https://files.pythonhosted.org/packages/2e/b9/8fe3ba5ed462067774ebc1f9c7f26aa7ebcc280ddd476be107153de1339e/grpcio-1.81.1-cp312-cp312-win_amd64.whl", hash = "sha256:3ad74f8bb1a18963914c5452d289422830b39459e8776ebbcd207be1fbfb1d94", size = 4930461, upload-time = "2026-06-11T12:45:45.775Z" }, - { url = "https://files.pythonhosted.org/packages/7a/42/dcc2e4b600538ef18327c0839d56b7d3c3812337c5d710df5877dbb39b1e/grpcio-1.81.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:b10e1ff4756ed27d5a29d7fc79cfce7ef1ff56ad20025b89bac7cf79e09abbbe", size = 6054466, upload-time = "2026-06-11T12:45:48.43Z" }, - { url = "https://files.pythonhosted.org/packages/7b/4a/a36e03210183a8a7d4c80c3936acee679f4bd77d5861f369db47b2cc5f05/grpcio-1.81.1-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:819edbdcb42ab8598b494bcf0222684bbb7a3c772bd1b1f0be7e029a6063c28e", size = 12048795, upload-time = "2026-06-11T12:45:54.011Z" }, - { url = "https://files.pythonhosted.org/packages/b0/d5/d68e30b29098f63beab6fe501100fe82674ff142b32c672532da86a99b3a/grpcio-1.81.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c5bf2dc311127d91230cc79b92188c082634a06cf66c5234db49a43b910183b0", size = 6599094, upload-time = "2026-06-11T12:45:57.799Z" }, - { url = "https://files.pythonhosted.org/packages/3d/b3/e837954d279754f638a11cca5dcf6b24a005efb398984cefaf7735945a54/grpcio-1.81.1-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:e8ca6a1fcdb2943c9cbc1804a1baf3acb6071d72a471591678ded84218006e14", size = 7307182, upload-time = "2026-06-11T12:46:00.568Z" }, - { url = "https://files.pythonhosted.org/packages/0d/1e/b47957057e729adc6cdf519a47f8be2562b7140e280f1418443eb4022192/grpcio-1.81.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e64dd101d380a115cc5a0c7856788adb535f1a4e21fc543775602f8be95180ae", size = 6810962, upload-time = "2026-06-11T12:46:03.312Z" }, - { url = "https://files.pythonhosted.org/packages/40/26/569868e364e05b19ec8f969da53d230bcd89c962cd198f7c29943155c4d3/grpcio-1.81.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:98a07f9bf591e3a8919797bee1c53f026ba4acd587e5a4404c8e57c9ec36b2a5", size = 7415698, upload-time = "2026-06-11T12:46:06.005Z" }, - { url = "https://files.pythonhosted.org/packages/36/0c/5440a0582cb5653fc42a6e262eeb22700943313f8076f9dc927491b20a59/grpcio-1.81.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:c261d74b1a945cf895a9d6eccd1685a8e837531beaab782da4d630a8d12deffb", size = 8407779, upload-time = "2026-06-11T12:46:08.84Z" }, - { url = "https://files.pythonhosted.org/packages/ff/aa/66fe9f39871d766987d869a03ee0842a026f499c7b1e62decb9e78a8088e/grpcio-1.81.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:58ad1131c300d3c9b933802b3cc4dc69d380822935ba50b28703156ea826fbf7", size = 7844521, upload-time = "2026-06-11T12:46:12.171Z" }, - { url = "https://files.pythonhosted.org/packages/f0/9e/69bb7194861bcd28fb3193261d4f9c3831b4446993f002cf59068943e7ab/grpcio-1.81.1-cp313-cp313-win32.whl", hash = "sha256:78e29211f26da2fdd0e9c6d2b79f489476140cf7029b6a64808ade7ca4156a42", size = 4182786, upload-time = "2026-06-11T12:46:15.192Z" }, - { url = "https://files.pythonhosted.org/packages/0d/20/3da8bb0d637feccdc3e1e419bb511ce93651ce7d54164f95de22cc0b8b34/grpcio-1.81.1-cp313-cp313-win_amd64.whl", hash = "sha256:edb59506291b647a30884b1d51a599d605f40b20af4a7dc3d33786a47a31de60", size = 4928648, upload-time = "2026-06-11T12:46:17.823Z" }, ] [[package]] @@ -1856,16 +1678,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/b5/67baeba7366162652cdc1dbd962289accde07241bc8f42f6f02b305efcc6/grpcio_tools-1.81.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:724ecb69af63d2f6d4ccea3e6fa0ca110ed9c5824d48c2f887c631bbb03c1c3c", size = 3370797, upload-time = "2026-06-11T12:50:21.501Z" }, { url = "https://files.pythonhosted.org/packages/f2/5d/34f2dce2125ccb107e32b57f5a9c1257edcc0793b0d2fef1e8b13a6bac3c/grpcio_tools-1.81.1-cp312-cp312-win32.whl", hash = "sha256:895a6782cec86beac71ccebb4b9848259c6f04a3028b8e42fa8d40cfe5146593", size = 1008453, upload-time = "2026-06-11T12:50:23.358Z" }, { url = "https://files.pythonhosted.org/packages/8a/be/09da8256ec8d2a5ce8a1acc51cbbc4ca52a462d78ed3412778440a56502e/grpcio_tools-1.81.1-cp312-cp312-win_amd64.whl", hash = "sha256:0265fd1386b7458302f79542558345880d484f8fa92ae196c0c0268242c5f23a", size = 1174857, upload-time = "2026-06-11T12:50:25.685Z" }, - { url = "https://files.pythonhosted.org/packages/76/90/5faa8b26e03495e5117f93bef8293cbada4af136362745dad7d1813ef0b0/grpcio_tools-1.81.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:3d604b4fd114b79ebb9f865bf3e04fd3ae93c704e1fad96f7fd03b0865c263b7", size = 2586071, upload-time = "2026-06-11T12:50:28.4Z" }, - { url = "https://files.pythonhosted.org/packages/e8/9a/85dc589fa6ae2439451eaa81a1578de31e29c676980d38bef7549b8a1f45/grpcio_tools-1.81.1-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:3389e705460efa3f3758141ba5520e6743b131c9576197c944fb9cbe49048126", size = 5813299, upload-time = "2026-06-11T12:50:31.295Z" }, - { url = "https://files.pythonhosted.org/packages/77/fd/c53994e58a837e6eefe48f53eb3492afc04f2b8af255df4adb37d14378f8/grpcio_tools-1.81.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8a17d8ceeb6a855fadf39f5171c80a382d97c4db98d5943eca553497fdebf84b", size = 2634668, upload-time = "2026-06-11T12:50:33.938Z" }, - { url = "https://files.pythonhosted.org/packages/34/32/de988e86688686a2117e7ce6ce9eff4f638c929bb55b0afe60d6fbd2e45c/grpcio_tools-1.81.1-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:43baf71dc60fd653062da2e95e95c73b35dd130be8f9fa3d544c3af3f808a290", size = 2957930, upload-time = "2026-06-11T12:50:36.726Z" }, - { url = "https://files.pythonhosted.org/packages/72/97/3f18a0ea32b5f809d21961dbd0bc382b589a4c3d501e3d67c345d5456ed3/grpcio_tools-1.81.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:136e90906af0df51ad929713244ba812d0dbb1844b4f467d5d86bdb054698f90", size = 2697760, upload-time = "2026-06-11T12:50:39.108Z" }, - { url = "https://files.pythonhosted.org/packages/49/c0/dbf5cbc877290ff7504a59959a8af4fdcfdaa1e84237948405ccf1aa82a6/grpcio_tools-1.81.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:bd6c3bf3ea6a61eb58c54368d72ada591f2a270f3a31a32e8536e773337e76d9", size = 3151456, upload-time = "2026-06-11T12:50:41.983Z" }, - { url = "https://files.pythonhosted.org/packages/de/ea/16fe2dc83140a59e5c0a0b9dc2693dd36bfaa6bd835724b4ec66a68eab7b/grpcio_tools-1.81.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:2c306c307f8f74cddc4056fdbb6f1da55de087a21120efbd02bd915daa5a52fd", size = 3710469, upload-time = "2026-06-11T12:50:44.596Z" }, - { url = "https://files.pythonhosted.org/packages/22/7d/df987d7d81e7ad2f7516d9e9d56ff29c54dbc6d8587e425688dca9a28e49/grpcio_tools-1.81.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:bdbdc927be2e0ea13c32564a72ee31d712a716fb6f8c0d53d37a77d8277c272c", size = 3370488, upload-time = "2026-06-11T12:50:47.199Z" }, - { url = "https://files.pythonhosted.org/packages/ba/c5/5a63444d694ea47bf670138208f71830cc1759c402c8818092b28ab2dc5f/grpcio_tools-1.81.1-cp313-cp313-win32.whl", hash = "sha256:9d383724bcd67244b6def9e9164c640ee9380c0b7534ee7545a6fb0022a59afe", size = 1008229, upload-time = "2026-06-11T12:50:49.527Z" }, - { url = "https://files.pythonhosted.org/packages/00/75/3945e26d5c94ae6ed9be5caef73d4d66c47dc8cfdd7b4995efaf942754e0/grpcio_tools-1.81.1-cp313-cp313-win_amd64.whl", hash = "sha256:f3eb15849979ca7bb864ce81a74d68b0f225a7f111ed3fe212bfc08cf9812b10", size = 1174523, upload-time = "2026-06-11T12:50:51.755Z" }, ] [[package]] @@ -1883,14 +1695,6 @@ version = "1.5.1" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/4b/2d/57fd21d84d93efb4bd0b962383790e19dd1bc053501b4264c97903b4e83e/hf_xet-1.5.1.tar.gz", hash = "sha256:51ef4500dab3764b41135ee1381a4b62ce56fc54d4c92b719b59e597d6df5bf6", size = 876636, upload-time = "2026-06-08T23:02:53.897Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/64/ee/dd9ba7beae1005e54131b7d45263cc74c8a066d47d354e6d58ae9445a388/hf_xet-1.5.1-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:dbf48c0d02cf0b2e568944330c60d9120c272dabe013bd892d48e25bc6797577", size = 4069485, upload-time = "2026-06-08T23:02:13.193Z" }, - { url = "https://files.pythonhosted.org/packages/b6/bc/9cae6cfeb4e03070874e73e5c97c66eb90369d3206b6a2b1ef5f96520888/hf_xet-1.5.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e78e4e5192ad2b674c2e1160b651cb9134db974f8ae1835bdfbfb0166b894a43", size = 3838493, upload-time = "2026-06-08T23:02:15.282Z" }, - { url = "https://files.pythonhosted.org/packages/ba/b4/d5c01e0eb6d9f2ca2dacd84d0d1b71e6cfbb2ef3208c968528e010e9b3d7/hf_xet-1.5.1-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6f7a04a8ad962422e225bc49fbbac99dc1806764b1f3e54dbd154bffa7593947", size = 4505658, upload-time = "2026-06-08T23:02:17.196Z" }, - { url = "https://files.pythonhosted.org/packages/76/c5/29a7598c0c6383c523dc22186d577f4e04267a626cd95ae60f67c00bfe66/hf_xet-1.5.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:d48199c2bf4f8df0adc55d31d1368b6ec0e4d4f45bc86b08038089c23db0bed8", size = 4292822, upload-time = "2026-06-08T23:02:18.608Z" }, - { url = "https://files.pythonhosted.org/packages/04/9a/dceaf6ca69390126b86ea825fb354b93d01163199070b7bd849225de9468/hf_xet-1.5.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:97f212a88d14bbf573619a74b7fecb238de77d08fc702e54dec6f78276ca3283", size = 4491255, upload-time = "2026-06-08T23:02:20.124Z" }, - { url = "https://files.pythonhosted.org/packages/48/a7/e5a7afaacf6c1791fdbeeac42951fb81c3d2bc482992b115dedcc86d963e/hf_xet-1.5.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:f61e3665892a6c8c5e765395838b8ddf36185da835253d4bc4509a81e49fb342", size = 4711062, upload-time = "2026-06-08T23:02:21.863Z" }, - { url = "https://files.pythonhosted.org/packages/53/49/2802f8433c9742ce281bddc1e65c02c32268ca3098d66828b05e12e45ee2/hf_xet-1.5.1-cp313-cp313t-win_amd64.whl", hash = "sha256:f4ad3ebd4c32dd2b27099d69dc7b2df821e30767e46fb6ee6a0713778243b8ff", size = 4017205, upload-time = "2026-06-08T23:02:23.495Z" }, - { url = "https://files.pythonhosted.org/packages/9e/5a/50c71195b9fb883659f596e7252faf4c18c58e753a9013bdbf9bac5d2250/hf_xet-1.5.1-cp313-cp313t-win_arm64.whl", hash = "sha256:8298485c1e36e7e67cbd01eeb1376619b7af43d4f1ec245caae306f890a8a32d", size = 3845426, upload-time = "2026-06-08T23:02:25.124Z" }, { url = "https://files.pythonhosted.org/packages/7a/d8/5e54cf37434759d1f4f2ba9b66077ff9d4c4e1f37b6bd7975da5c40d94ab/hf_xet-1.5.1-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:6abd35c3221eff63836618ddfb954dcf84798603f71d8e33e3ed7b04acfdbe6e", size = 4077794, upload-time = "2026-06-08T23:02:40.656Z" }, { url = "https://files.pythonhosted.org/packages/35/94/4b2ecfbad8f8b04701a23aefb62f540b9137d058b7e1dbef16a32676f0e9/hf_xet-1.5.1-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:94e761bbd266bf4c03cee73753916062665ce8365aa40ed321f45afcb934b41e", size = 3845354, upload-time = "2026-06-08T23:02:42.702Z" }, { url = "https://files.pythonhosted.org/packages/de/cc/f99f4bc7295023d7bd9ebbfd51f75cc530ca262c1227666268b8208f4b77/hf_xet-1.5.1-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:892e3a3a3aecc12aded8b93cf4f9cd059282c7de0732f7d55026f3abdf474350", size = 4514864, upload-time = "2026-06-08T23:02:44.497Z" }, @@ -2189,35 +1993,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e4/34/8aefdd0be9cfd00a44509251ba864f5caf2991e36772e61c408007e7f417/kiwisolver-1.5.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:1d9daea4ea6b9be74fe2f01f7fbade8d6ffab263e781274cffca0dba9be9eec9", size = 2294947, upload-time = "2026-03-09T13:13:43.343Z" }, { url = "https://files.pythonhosted.org/packages/ad/cf/0348374369ca588f8fe9c338fae49fa4e16eeb10ffb3d012f23a54578a9e/kiwisolver-1.5.0-cp312-cp312-win_amd64.whl", hash = "sha256:f18c2d9782259a6dc132fdc7a63c168cbc74b35284b6d75c673958982a378384", size = 73569, upload-time = "2026-03-09T13:13:45.792Z" }, { url = "https://files.pythonhosted.org/packages/28/26/192b26196e2316e2bd29deef67e37cdf9870d9af8e085e521afff0fed526/kiwisolver-1.5.0-cp312-cp312-win_arm64.whl", hash = "sha256:f7c7553b13f69c1b29a5bde08ddc6d9d0c8bfb84f9ed01c30db25944aeb852a7", size = 64997, upload-time = "2026-03-09T13:13:46.878Z" }, - { url = "https://files.pythonhosted.org/packages/9d/69/024d6711d5ba575aa65d5538042e99964104e97fa153a9f10bc369182bc2/kiwisolver-1.5.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:fd40bb9cd0891c4c3cb1ddf83f8bbfa15731a248fdc8162669405451e2724b09", size = 123166, upload-time = "2026-03-09T13:13:48.032Z" }, - { url = "https://files.pythonhosted.org/packages/ce/48/adbb40df306f587054a348831220812b9b1d787aff714cfbc8556e38fccd/kiwisolver-1.5.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c0e1403fd7c26d77c1f03e096dc58a5c726503fa0db0456678b8668f76f521e3", size = 66395, upload-time = "2026-03-09T13:13:49.365Z" }, - { url = "https://files.pythonhosted.org/packages/a8/3a/d0a972b34e1c63e2409413104216cd1caa02c5a37cb668d1687d466c1c45/kiwisolver-1.5.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:dda366d548e89a90d88a86c692377d18d8bd64b39c1fb2b92cb31370e2896bbd", size = 64065, upload-time = "2026-03-09T13:13:50.562Z" }, - { url = "https://files.pythonhosted.org/packages/2b/0a/7b98e1e119878a27ba8618ca1e18b14f992ff1eda40f47bccccf4de44121/kiwisolver-1.5.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:332b4f0145c30b5f5ad9374881133e5aa64320428a57c2c2b61e9d891a51c2f3", size = 1477903, upload-time = "2026-03-09T13:13:52.084Z" }, - { url = "https://files.pythonhosted.org/packages/18/d8/55638d89ffd27799d5cc3d8aa28e12f4ce7a64d67b285114dbedc8ea4136/kiwisolver-1.5.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0c50b89ffd3e1a911c69a1dd3de7173c0cd10b130f56222e57898683841e4f96", size = 1278751, upload-time = "2026-03-09T13:13:54.673Z" }, - { url = "https://files.pythonhosted.org/packages/b8/97/b4c8d0d18421ecceba20ad8701358453b88e32414e6f6950b5a4bad54e65/kiwisolver-1.5.0-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:4db576bb8c3ef9365f8b40fe0f671644de6736ae2c27a2c62d7d8a1b4329f099", size = 1296793, upload-time = "2026-03-09T13:13:56.287Z" }, - { url = "https://files.pythonhosted.org/packages/c4/10/f862f94b6389d8957448ec9df59450b81bec4abb318805375c401a1e6892/kiwisolver-1.5.0-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:0b85aad90cea8ac6797a53b5d5f2e967334fa4d1149f031c4537569972596cb8", size = 1346041, upload-time = "2026-03-09T13:13:58.269Z" }, - { url = "https://files.pythonhosted.org/packages/a3/6a/f1650af35821eaf09de398ec0bc2aefc8f211f0cda50204c9f1673741ba9/kiwisolver-1.5.0-cp313-cp313-manylinux_2_39_riscv64.whl", hash = "sha256:d36ca54cb4c6c4686f7cbb7b817f66f5911c12ddb519450bbe86707155028f87", size = 987292, upload-time = "2026-03-09T13:13:59.871Z" }, - { url = "https://files.pythonhosted.org/packages/de/19/d7fb82984b9238115fe629c915007be608ebd23dc8629703d917dbfaffd4/kiwisolver-1.5.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:38f4a703656f493b0ad185211ccfca7f0386120f022066b018eb5296d8613e23", size = 2227865, upload-time = "2026-03-09T13:14:01.401Z" }, - { url = "https://files.pythonhosted.org/packages/7f/b9/46b7f386589fd222dac9e9de9c956ce5bcefe2ee73b4e79891381dda8654/kiwisolver-1.5.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:3ac2360e93cb41be81121755c6462cff3beaa9967188c866e5fce5cf13170859", size = 2324369, upload-time = "2026-03-09T13:14:02.972Z" }, - { url = "https://files.pythonhosted.org/packages/92/8b/95e237cf3d9c642960153c769ddcbe278f182c8affb20cecc1cc983e7cc5/kiwisolver-1.5.0-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:c95cab08d1965db3d84a121f1c7ce7479bdd4072c9b3dafd8fecce48a2e6b902", size = 1977989, upload-time = "2026-03-09T13:14:04.503Z" }, - { url = "https://files.pythonhosted.org/packages/1b/95/980c9df53501892784997820136c01f62bc1865e31b82b9560f980c0e649/kiwisolver-1.5.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:fc20894c3d21194d8041a28b65622d5b86db786da6e3cfe73f0c762951a61167", size = 2491645, upload-time = "2026-03-09T13:14:06.106Z" }, - { url = "https://files.pythonhosted.org/packages/cb/32/900647fd0840abebe1561792c6b31e6a7c0e278fc3973d30572a965ca14c/kiwisolver-1.5.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:7a32f72973f0f950c1920475d5c5ea3d971b81b6f0ec53b8d0a956cc965f22e0", size = 2295237, upload-time = "2026-03-09T13:14:08.891Z" }, - { url = "https://files.pythonhosted.org/packages/be/8a/be60e3bbcf513cc5a50f4a3e88e1dcecebb79c1ad607a7222877becaa101/kiwisolver-1.5.0-cp313-cp313-win_amd64.whl", hash = "sha256:0bf3acf1419fa93064a4c2189ac0b58e3be7872bf6ee6177b0d4c63dc4cea276", size = 73573, upload-time = "2026-03-09T13:14:12.327Z" }, - { url = "https://files.pythonhosted.org/packages/4d/d2/64be2e429eb4fca7f7e1c52a91b12663aeaf25de3895e5cca0f47ef2a8d0/kiwisolver-1.5.0-cp313-cp313-win_arm64.whl", hash = "sha256:fa8eb9ecdb7efb0b226acec134e0d709e87a909fa4971a54c0c4f6e88635484c", size = 64998, upload-time = "2026-03-09T13:14:13.469Z" }, - { url = "https://files.pythonhosted.org/packages/b0/69/ce68dd0c85755ae2de490bf015b62f2cea5f6b14ff00a463f9d0774449ff/kiwisolver-1.5.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:db485b3847d182b908b483b2ed133c66d88d49cacf98fd278fadafe11b4478d1", size = 125700, upload-time = "2026-03-09T13:14:14.636Z" }, - { url = "https://files.pythonhosted.org/packages/74/aa/937aac021cf9d4349990d47eb319309a51355ed1dbdc9c077cdc9224cb11/kiwisolver-1.5.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:be12f931839a3bdfe28b584db0e640a65a8bcbc24560ae3fdb025a449b3d754e", size = 67537, upload-time = "2026-03-09T13:14:15.808Z" }, - { url = "https://files.pythonhosted.org/packages/ee/20/3a87fbece2c40ad0f6f0aefa93542559159c5f99831d596050e8afae7a9f/kiwisolver-1.5.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:16b85d37c2cbb3253226d26e64663f755d88a03439a9c47df6246b35defbdfb7", size = 65514, upload-time = "2026-03-09T13:14:18.035Z" }, - { url = "https://files.pythonhosted.org/packages/f0/7f/f943879cda9007c45e1f7dba216d705c3a18d6b35830e488b6c6a4e7cdf0/kiwisolver-1.5.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4432b835675f0ea7414aab3d37d119f7226d24869b7a829caeab49ebda407b0c", size = 1584848, upload-time = "2026-03-09T13:14:19.745Z" }, - { url = "https://files.pythonhosted.org/packages/37/f8/4d4f85cc1870c127c88d950913370dd76138482161cd07eabbc450deff01/kiwisolver-1.5.0-cp313-cp313t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1b0feb50971481a2cc44d94e88bdb02cdd497618252ae226b8eb1201b957e368", size = 1391542, upload-time = "2026-03-09T13:14:21.54Z" }, - { url = "https://files.pythonhosted.org/packages/04/0b/65dd2916c84d252b244bd405303220f729e7c17c9d7d33dca6feeff9ffc4/kiwisolver-1.5.0-cp313-cp313t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:56fa888f10d0f367155e76ce849fa1166fc9730d13bd2d65a2aa13b6f5424489", size = 1404447, upload-time = "2026-03-09T13:14:23.205Z" }, - { url = "https://files.pythonhosted.org/packages/39/5c/2606a373247babce9b1d056c03a04b65f3cf5290a8eac5d7bdead0a17e21/kiwisolver-1.5.0-cp313-cp313t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:940dda65d5e764406b9fb92761cbf462e4e63f712ab60ed98f70552e496f3bf1", size = 1455918, upload-time = "2026-03-09T13:14:24.74Z" }, - { url = "https://files.pythonhosted.org/packages/d5/d1/c6078b5756670658e9192a2ef11e939c92918833d2745f85cd14a6004bdf/kiwisolver-1.5.0-cp313-cp313t-manylinux_2_39_riscv64.whl", hash = "sha256:89fc958c702ee9a745e4700378f5d23fddbc46ff89e8fdbf5395c24d5c1452a3", size = 1072856, upload-time = "2026-03-09T13:14:26.597Z" }, - { url = "https://files.pythonhosted.org/packages/cb/c8/7def6ddf16eb2b3741d8b172bdaa9af882b03c78e9b0772975408801fa63/kiwisolver-1.5.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9027d773c4ff81487181a925945743413f6069634d0b122d0b37684ccf4f1e18", size = 2333580, upload-time = "2026-03-09T13:14:28.237Z" }, - { url = "https://files.pythonhosted.org/packages/9e/87/2ac1fce0eb1e616fcd3c35caa23e665e9b1948bb984f4764790924594128/kiwisolver-1.5.0-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:5b233ea3e165e43e35dba1d2b8ecc21cf070b45b65ae17dd2747d2713d942021", size = 2423018, upload-time = "2026-03-09T13:14:30.018Z" }, - { url = "https://files.pythonhosted.org/packages/67/13/c6700ccc6cc218716bfcda4935e4b2997039869b4ad8a94f364c5a3b8e63/kiwisolver-1.5.0-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:ce9bf03dad3b46408c08649c6fbd6ca28a9fce0eb32fdfffa6775a13103b5310", size = 2062804, upload-time = "2026-03-09T13:14:32.888Z" }, - { url = "https://files.pythonhosted.org/packages/1b/bd/877056304626943ff0f1f44c08f584300c199b887cb3176cd7e34f1515f1/kiwisolver-1.5.0-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:fc4d3f1fb9ca0ae9f97b095963bc6326f1dbfd3779d6679a1e016b9baaa153d3", size = 2597482, upload-time = "2026-03-09T13:14:34.971Z" }, - { url = "https://files.pythonhosted.org/packages/75/19/c60626c47bf0f8ac5dcf72c6c98e266d714f2fbbfd50cf6dab5ede3aaa50/kiwisolver-1.5.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:f443b4825c50a51ee68585522ab4a1d1257fac65896f282b4c6763337ac9f5d2", size = 2394328, upload-time = "2026-03-09T13:14:36.816Z" }, - { url = "https://files.pythonhosted.org/packages/47/84/6a6d5e5bb8273756c27b7d810d47f7ef2f1f9b9fd23c9ee9a3f8c75c9cef/kiwisolver-1.5.0-cp313-cp313t-win_arm64.whl", hash = "sha256:893ff3a711d1b515ba9da14ee090519bad4610ed1962fbe298a434e8c5f8db53", size = 68410, upload-time = "2026-03-09T13:14:38.695Z" }, { url = "https://files.pythonhosted.org/packages/1c/fa/2910df836372d8761bb6eff7d8bdcb1613b5c2e03f260efe7abe34d388a7/kiwisolver-1.5.0-graalpy312-graalpy250_312_native-macosx_10_13_x86_64.whl", hash = "sha256:5ae8e62c147495b01a0f4765c878e9bfdf843412446a247e28df59936e99e797", size = 130262, upload-time = "2026-03-09T13:15:35.629Z" }, { url = "https://files.pythonhosted.org/packages/0f/41/c5f71f9f00aabcc71fee8b7475e3f64747282580c2fe748961ba29b18385/kiwisolver-1.5.0-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:f6764a4ccab3078db14a632420930f6186058750df066b8ea2a7106df91d3203", size = 138036, upload-time = "2026-03-09T13:15:36.894Z" }, { url = "https://files.pythonhosted.org/packages/fa/06/7399a607f434119c6e1fdc8ec89a8d51ccccadf3341dee4ead6bd14caaf5/kiwisolver-1.5.0-graalpy312-graalpy250_312_native-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c31c13da98624f957b0fb1b5bae5383b2333c2c3f6793d9825dd5ce79b525cb7", size = 194295, upload-time = "2026-03-09T13:15:38.22Z" }, @@ -2276,7 +2051,7 @@ dev = [ { name = "pytest" }, ] video-codec = [ - { name = "pynvvideocodec", marker = "python_full_version < '3.13' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "pynvvideocodec" }, ] [package.metadata] @@ -2370,28 +2145,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, { url = "https://files.pythonhosted.org/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, - { url = "https://files.pythonhosted.org/packages/38/2f/907b9c7bbba283e68f20259574b13d005c121a0fa4c175f9bed27c4597ff/markupsafe-3.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e1cf1972137e83c5d4c136c43ced9ac51d0e124706ee1c8aa8532c1287fa8795", size = 11622, upload-time = "2025-09-27T18:36:41.777Z" }, - { url = "https://files.pythonhosted.org/packages/9c/d9/5f7756922cdd676869eca1c4e3c0cd0df60ed30199ffd775e319089cb3ed/markupsafe-3.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:116bb52f642a37c115f517494ea5feb03889e04df47eeff5b130b1808ce7c219", size = 12029, upload-time = "2025-09-27T18:36:43.257Z" }, - { url = "https://files.pythonhosted.org/packages/00/07/575a68c754943058c78f30db02ee03a64b3c638586fba6a6dd56830b30a3/markupsafe-3.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:133a43e73a802c5562be9bbcd03d090aa5a1fe899db609c29e8c8d815c5f6de6", size = 24374, upload-time = "2025-09-27T18:36:44.508Z" }, - { url = "https://files.pythonhosted.org/packages/a9/21/9b05698b46f218fc0e118e1f8168395c65c8a2c750ae2bab54fc4bd4e0e8/markupsafe-3.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ccfcd093f13f0f0b7fdd0f198b90053bf7b2f02a3927a30e63f3ccc9df56b676", size = 22980, upload-time = "2025-09-27T18:36:45.385Z" }, - { url = "https://files.pythonhosted.org/packages/7f/71/544260864f893f18b6827315b988c146b559391e6e7e8f7252839b1b846a/markupsafe-3.0.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:509fa21c6deb7a7a273d629cf5ec029bc209d1a51178615ddf718f5918992ab9", size = 21990, upload-time = "2025-09-27T18:36:46.916Z" }, - { url = "https://files.pythonhosted.org/packages/c2/28/b50fc2f74d1ad761af2f5dcce7492648b983d00a65b8c0e0cb457c82ebbe/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a4afe79fb3de0b7097d81da19090f4df4f8d3a2b3adaa8764138aac2e44f3af1", size = 23784, upload-time = "2025-09-27T18:36:47.884Z" }, - { url = "https://files.pythonhosted.org/packages/ed/76/104b2aa106a208da8b17a2fb72e033a5a9d7073c68f7e508b94916ed47a9/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:795e7751525cae078558e679d646ae45574b47ed6e7771863fcc079a6171a0fc", size = 21588, upload-time = "2025-09-27T18:36:48.82Z" }, - { url = "https://files.pythonhosted.org/packages/b5/99/16a5eb2d140087ebd97180d95249b00a03aa87e29cc224056274f2e45fd6/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8485f406a96febb5140bfeca44a73e3ce5116b2501ac54fe953e488fb1d03b12", size = 23041, upload-time = "2025-09-27T18:36:49.797Z" }, - { url = "https://files.pythonhosted.org/packages/19/bc/e7140ed90c5d61d77cea142eed9f9c303f4c4806f60a1044c13e3f1471d0/markupsafe-3.0.3-cp313-cp313-win32.whl", hash = "sha256:bdd37121970bfd8be76c5fb069c7751683bdf373db1ed6c010162b2a130248ed", size = 14543, upload-time = "2025-09-27T18:36:51.584Z" }, - { url = "https://files.pythonhosted.org/packages/05/73/c4abe620b841b6b791f2edc248f556900667a5a1cf023a6646967ae98335/markupsafe-3.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:9a1abfdc021a164803f4d485104931fb8f8c1efd55bc6b748d2f5774e78b62c5", size = 15113, upload-time = "2025-09-27T18:36:52.537Z" }, - { url = "https://files.pythonhosted.org/packages/f0/3a/fa34a0f7cfef23cf9500d68cb7c32dd64ffd58a12b09225fb03dd37d5b80/markupsafe-3.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:7e68f88e5b8799aa49c85cd116c932a1ac15caaa3f5db09087854d218359e485", size = 13911, upload-time = "2025-09-27T18:36:53.513Z" }, - { url = "https://files.pythonhosted.org/packages/e4/d7/e05cd7efe43a88a17a37b3ae96e79a19e846f3f456fe79c57ca61356ef01/markupsafe-3.0.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:218551f6df4868a8d527e3062d0fb968682fe92054e89978594c28e642c43a73", size = 11658, upload-time = "2025-09-27T18:36:54.819Z" }, - { url = "https://files.pythonhosted.org/packages/99/9e/e412117548182ce2148bdeacdda3bb494260c0b0184360fe0d56389b523b/markupsafe-3.0.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3524b778fe5cfb3452a09d31e7b5adefeea8c5be1d43c4f810ba09f2ceb29d37", size = 12066, upload-time = "2025-09-27T18:36:55.714Z" }, - { url = "https://files.pythonhosted.org/packages/bc/e6/fa0ffcda717ef64a5108eaa7b4f5ed28d56122c9a6d70ab8b72f9f715c80/markupsafe-3.0.3-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e885a3d1efa2eadc93c894a21770e4bc67899e3543680313b09f139e149ab19", size = 25639, upload-time = "2025-09-27T18:36:56.908Z" }, - { url = "https://files.pythonhosted.org/packages/96/ec/2102e881fe9d25fc16cb4b25d5f5cde50970967ffa5dddafdb771237062d/markupsafe-3.0.3-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8709b08f4a89aa7586de0aadc8da56180242ee0ada3999749b183aa23df95025", size = 23569, upload-time = "2025-09-27T18:36:57.913Z" }, - { url = "https://files.pythonhosted.org/packages/4b/30/6f2fce1f1f205fc9323255b216ca8a235b15860c34b6798f810f05828e32/markupsafe-3.0.3-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b8512a91625c9b3da6f127803b166b629725e68af71f8184ae7e7d54686a56d6", size = 23284, upload-time = "2025-09-27T18:36:58.833Z" }, - { url = "https://files.pythonhosted.org/packages/58/47/4a0ccea4ab9f5dcb6f79c0236d954acb382202721e704223a8aafa38b5c8/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9b79b7a16f7fedff2495d684f2b59b0457c3b493778c9eed31111be64d58279f", size = 24801, upload-time = "2025-09-27T18:36:59.739Z" }, - { url = "https://files.pythonhosted.org/packages/6a/70/3780e9b72180b6fecb83a4814d84c3bf4b4ae4bf0b19c27196104149734c/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:12c63dfb4a98206f045aa9563db46507995f7ef6d83b2f68eda65c307c6829eb", size = 22769, upload-time = "2025-09-27T18:37:00.719Z" }, - { url = "https://files.pythonhosted.org/packages/98/c5/c03c7f4125180fc215220c035beac6b9cb684bc7a067c84fc69414d315f5/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8f71bc33915be5186016f675cd83a1e08523649b0e33efdb898db577ef5bb009", size = 23642, upload-time = "2025-09-27T18:37:01.673Z" }, - { url = "https://files.pythonhosted.org/packages/80/d6/2d1b89f6ca4bff1036499b1e29a1d02d282259f3681540e16563f27ebc23/markupsafe-3.0.3-cp313-cp313t-win32.whl", hash = "sha256:69c0b73548bc525c8cb9a251cddf1931d1db4d2258e9599c28c07ef3580ef354", size = 14612, upload-time = "2025-09-27T18:37:02.639Z" }, - { url = "https://files.pythonhosted.org/packages/2b/98/e48a4bfba0a0ffcf9925fe2d69240bfaa19c6f7507b8cd09c70684a53c1e/markupsafe-3.0.3-cp313-cp313t-win_amd64.whl", hash = "sha256:1b4b79e8ebf6b55351f0d91fe80f893b4743f104bff22e90697db1590e47a218", size = 15200, upload-time = "2025-09-27T18:37:03.582Z" }, - { url = "https://files.pythonhosted.org/packages/0e/72/e3cc540f351f316e9ed0f092757459afbc595824ca724cbc5a5d4263713f/markupsafe-3.0.3-cp313-cp313t-win_arm64.whl", hash = "sha256:ad2cf8aa28b8c020ab2fc8287b0f823d0a7d8630784c31e9ee5edea20f406287", size = 13973, upload-time = "2025-09-27T18:37:04.929Z" }, ] [[package]] @@ -2435,20 +2188,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fc/c2/541e4d09d87bb6b5830fc28b4c887a9a8cf4e1c6cee698a8c05552ae2003/matplotlib-3.10.9-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:d75d11c949914165976c621b2324f9ef162af7ebf4b057ddf95dd1dba7e5edcf", size = 9670966, upload-time = "2026-04-24T00:12:32.131Z" }, { url = "https://files.pythonhosted.org/packages/04/a1/4571fc46e7702de8d0c2dc54ad1b2f8e29328dea3ee90831181f7353d93c/matplotlib-3.10.9-cp312-cp312-win_amd64.whl", hash = "sha256:d091f9d758b34aaaaa6331d13574bf01891d903b3dec59bfff458ef7551de5d6", size = 8217462, upload-time = "2026-04-24T00:12:35.226Z" }, { url = "https://files.pythonhosted.org/packages/4b/d0/2269edb12aa30c13c8bcc9382892e39943ce1d28aab4ec296e0381798e81/matplotlib-3.10.9-cp312-cp312-win_arm64.whl", hash = "sha256:10cc5ce06d10231c36f40e875f3c7e8050362a4ee8f0ee5d29a6b3277d57bb42", size = 8136688, upload-time = "2026-04-24T00:12:37.442Z" }, - { url = "https://files.pythonhosted.org/packages/aa/d3/8d4f6afbecb49fc04e060a57c0fce39ea51cc163a6bd87303ccd698e4fa6/matplotlib-3.10.9-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:b580440f1ff81a0e34122051a3dfabb7e4b7f9e380629929bde0eff9af72165f", size = 8320331, upload-time = "2026-04-24T00:12:39.688Z" }, - { url = "https://files.pythonhosted.org/packages/63/d9/9e14bc7564bf92d5ffa801ae5fac819ce74b925dfb55e3ebde61a3bbad3e/matplotlib-3.10.9-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:b1b745c489cd1a77a0dc1120a05dc87af9798faebc913601feb8c73d89bf2d1e", size = 8216461, upload-time = "2026-04-24T00:12:42.494Z" }, - { url = "https://files.pythonhosted.org/packages/8a/17/4402d0d14ccf1dfc70932600b68097fbbf9c898a4871d2cbbe79c7801a32/matplotlib-3.10.9-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8f3bcac1ca5ed000a6f4337d47ba67dfddf37ed6a46c15fd7f014997f7bf865f", size = 8790091, upload-time = "2026-04-24T00:12:44.789Z" }, - { url = "https://files.pythonhosted.org/packages/3e/0b/322aeec06dd9b91411f92028b37d447342770a24392aa4813e317064dad5/matplotlib-3.10.9-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7a8d66a55def891c33147ba3ba9bfcabf0b526a43764c818acbb4525e5ed0838", size = 9605027, upload-time = "2026-04-24T00:12:47.583Z" }, - { url = "https://files.pythonhosted.org/packages/74/88/5f13482f55e7b00bcfc09838b093c2456e1379978d2a146844aae05350ad/matplotlib-3.10.9-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:d843374407c4017a6403b59c6c81606773d136f3259d5b6da3131bc814542cc2", size = 9671269, upload-time = "2026-04-24T00:12:50.878Z" }, - { url = "https://files.pythonhosted.org/packages/c5/e0/0840fd2f93da988ec660b8ad1984abe9f25d2aed22a5e394ff1c68c88307/matplotlib-3.10.9-cp313-cp313-win_amd64.whl", hash = "sha256:f4399f64b3e94cd500195490972ae1ee81170df1636fa15364d157d5bdd7b921", size = 8217588, upload-time = "2026-04-24T00:12:53.784Z" }, - { url = "https://files.pythonhosted.org/packages/47/b9/d706d06dd605c49b9f83a2aed8c13e3e5db70697d7a80b7e3d7915de6b17/matplotlib-3.10.9-cp313-cp313-win_arm64.whl", hash = "sha256:ba7b3b8ef09eab7df0e86e9ae086faa433efbfbdb46afcb3aa16aabf779469a8", size = 8136913, upload-time = "2026-04-24T00:12:56.501Z" }, - { url = "https://files.pythonhosted.org/packages/9b/45/6e32d96978264c8ca8c4b1010adb955a1a49cfaf314e212bbc8908f04a61/matplotlib-3.10.9-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:09218df8a93712bd6ea133e83a153c755448cf7868316c531cffcc43f69d1cc9", size = 8368019, upload-time = "2026-04-24T00:12:58.896Z" }, - { url = "https://files.pythonhosted.org/packages/86/0a/c8e3d3bba245f0f7fc424937f8ff7ef77291a36af3edb97ccd78aa93d84f/matplotlib-3.10.9-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:82368699727bfb7b0182e1aa13082e3c08e092fa1a25d3e1fd92405bff96f6d4", size = 8264645, upload-time = "2026-04-24T00:13:01.406Z" }, - { url = "https://files.pythonhosted.org/packages/3d/aa/5bf5a14fe4fed73a4209a155606f8096ff797aad89c6c35179026571133e/matplotlib-3.10.9-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3225f4e1edcb8c86c884ddf79ebe20ecd0a67d30188f279897554ccd8fded4dc", size = 8802194, upload-time = "2026-04-24T00:13:03.702Z" }, - { url = "https://files.pythonhosted.org/packages/dd/5e/b4be852d6bba6fd15893fadf91ff26ae49cb91aac789e95dde9d342e664f/matplotlib-3.10.9-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:de2445a0c6690d21b7eb6ce071cebad6d40a2e9bdf10d039074a96ba19797b99", size = 9622684, upload-time = "2026-04-24T00:13:06.647Z" }, - { url = "https://files.pythonhosted.org/packages/4c/3d/ed428c971139112ef730f62770654d609467346d09d4b62617e1afd68a5a/matplotlib-3.10.9-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:b2b9516251cb89ff618d757daec0e2ed1bf21248013844a853d87ef85ab3081d", size = 9680790, upload-time = "2026-04-24T00:13:10.009Z" }, - { url = "https://files.pythonhosted.org/packages/e7/09/052e884aaf2b985c63cb79f715f1d5b6a3eaa7de78f6a52b9dbc077d5b53/matplotlib-3.10.9-cp313-cp313t-win_amd64.whl", hash = "sha256:e9fae004b941b23ff2edcf1567a857ed77bafc8086ffa258190462328434faf8", size = 8287571, upload-time = "2026-04-24T00:13:13.087Z" }, - { url = "https://files.pythonhosted.org/packages/f4/38/ae27288e788c35a4250491422f3db7750366fc8c97d6f36fbdecfc1f5518/matplotlib-3.10.9-cp313-cp313t-win_arm64.whl", hash = "sha256:6b63d9c7c769b88ab81e10dc86e4e0607cf56817b9f9e6cf24b2a5f1693b8e38", size = 8188292, upload-time = "2026-04-24T00:13:15.546Z" }, { url = "https://files.pythonhosted.org/packages/2c/2b/0e92ad0ac446633f928a1563db4aa8add407e1924faf0ded5b95b35afb27/matplotlib-3.10.9-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:1872fb212a05b729e649754a72d5da61d03e0554d76e80303b6f83d1d2c0552b", size = 8293058, upload-time = "2026-04-24T00:13:56.339Z" }, { url = "https://files.pythonhosted.org/packages/4b/23/74682fd369f5299ceda438fea2a0662e6383b85c9383fb9cdfcf04713e07/matplotlib-3.10.9-pp310-pypy310_pp73-macosx_11_0_arm64.whl", hash = "sha256:985f2238880e2e69093f588f5fe2e46771747febf0649f3cf7f7b7480875317f", size = 8186627, upload-time = "2026-04-24T00:13:58.623Z" }, { url = "https://files.pythonhosted.org/packages/ca/e8/368aab88f3c4cd8992800f31abfe0670c3e47540ba20a97e9fdbcde594b3/matplotlib-3.10.9-pp310-pypy310_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6640f75af2c6148293caa0a2b39dd806a492dd66c8a8b04035813e33d0fd2585", size = 8764117, upload-time = "2026-04-24T00:14:01.684Z" }, @@ -2494,20 +2233,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/85/c2/db15da2bbdf9e3ca66df7db8e2c33a1dfed67be24a24d2c878efaaff01d6/matplotlib-3.11.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:94f5000f67ca9faa300863ea17f8bce9175cb67b88bec4bc7780502d53dd7c9e", size = 10923899, upload-time = "2026-06-12T02:28:00.223Z" }, { url = "https://files.pythonhosted.org/packages/e5/2f/a58a4443a4d052a4ea77557478336aefc26c7981f6408d37adba763aa758/matplotlib-3.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:ac6f1ef39f3d0f9e2463303013094992cdbe0f85f43bc54155bc472b2042768e", size = 9329528, upload-time = "2026-06-12T02:28:02.27Z" }, { url = "https://files.pythonhosted.org/packages/61/0f/4b669589d47733b97ab9df4b58d6fc1e68acb5ea42a928dc7cbdd6bf5871/matplotlib-3.11.0-cp312-cp312-win_arm64.whl", hash = "sha256:9dd11fb612ce7bc60b1de5b4fc87ff959d22317b5de42aabf392f66f97af22eb", size = 9003413, upload-time = "2026-06-12T02:28:04.49Z" }, - { url = "https://files.pythonhosted.org/packages/55/41/aa47f156b061d14c98b906f76c428507397708ec63ff94f410ae1752b426/matplotlib-3.11.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:6ce3b839b34ae1f430b4616893a2945a2999debaa7e94e7e29a2a8bbf286f7b5", size = 9450532, upload-time = "2026-06-12T02:28:06.769Z" }, - { url = "https://files.pythonhosted.org/packages/8c/4f/5a9eb0375e81413953febf8af7b012a6b6357f53438a15c4f5ad86c6bbb5/matplotlib-3.11.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:373db8f91214e8ccaf35ac833cc1dd59dd961e148bbd55dd027141591dde1313", size = 9279760, upload-time = "2026-06-12T02:28:09.152Z" }, - { url = "https://files.pythonhosted.org/packages/a4/c0/1117d53077e3ac3152503a84e9cf7a5c239576805ee71276e80c2aaa7471/matplotlib-3.11.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:be152b7570324dc8d01574cc9474dd2d803237acf528bcbb5b211fa347461a09", size = 10031623, upload-time = "2026-06-12T02:28:11.26Z" }, - { url = "https://files.pythonhosted.org/packages/92/7e/e937138daffad65b71bf831a377809dcbc830fb4f31a31e067dc1faa2575/matplotlib-3.11.0-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:126f256df600652d7e4b394cf3164ff75210a00038f287c95a012a6f58d0e83f", size = 10839372, upload-time = "2026-06-12T02:28:14.102Z" }, - { url = "https://files.pythonhosted.org/packages/1d/c2/438ecc197ffb8023b6b9922915542f2172f5fd45b76703b0b4fc47322243/matplotlib-3.11.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:03acfeddf87b0dddb11b081ef7740ad445a3ca8bcb6b8e3011b08f2cf802b75c", size = 10924099, upload-time = "2026-06-12T02:28:16.383Z" }, - { url = "https://files.pythonhosted.org/packages/40/2e/395883da416f378b3ed2c9f3e843ac477eae1ce731b671b79adaa6f0bacd/matplotlib-3.11.0-cp313-cp313-win_amd64.whl", hash = "sha256:ab3722f04f3ff34c23b5012c5873d2894174e06c3822fcdac3610965a5ac7d06", size = 9329727, upload-time = "2026-06-12T02:28:18.581Z" }, - { url = "https://files.pythonhosted.org/packages/61/82/2c388956abf8bf392dfb5b8917c502f1082df6a941b781ab8c8e5ba2474b/matplotlib-3.11.0-cp313-cp313-win_arm64.whl", hash = "sha256:c945824670fb8915b4ac879e5e61f3c58e0913022f70a0de4c082b17372f8771", size = 9003506, upload-time = "2026-06-12T02:28:20.474Z" }, - { url = "https://files.pythonhosted.org/packages/c8/c1/34454baa44da7975ada82e9aea37105ec47059514dc967d3be14426ba8dc/matplotlib-3.11.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:3489c3dc487669b4a980bc3068f87856de7a1564248d3f6c629efb2a58b03f24", size = 9499838, upload-time = "2026-06-12T02:28:22.713Z" }, - { url = "https://files.pythonhosted.org/packages/b1/c3/98fe79a398cf232219f090163a7fa7e6766e9f2e0ad26df54d6f8934d8ee/matplotlib-3.11.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:6a98f5476ce784a50ce09998f4ae1e6a9f25043cef8a480c98949902eda74620", size = 9332298, upload-time = "2026-06-12T02:28:24.796Z" }, - { url = "https://files.pythonhosted.org/packages/95/e4/b4b7c33151e74e5c802f3cde1ba807ebfc38401e329b44e215a5888dd76d/matplotlib-3.11.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:565af866fd63e4bd3f987d580afe27c44c2552a3b3305f4ecbb85133601ea6f3", size = 10045491, upload-time = "2026-06-12T02:28:27.141Z" }, - { url = "https://files.pythonhosted.org/packages/71/28/394548efd68354110c1a1be11fe6b6e559e06d1a23da35908a0e316c55a9/matplotlib-3.11.0-cp313-cp313t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e6b3e64dea5062c570f04358e2711859f3531b459f29516274fbad889079e4f3", size = 10857059, upload-time = "2026-06-12T02:28:29.222Z" }, - { url = "https://files.pythonhosted.org/packages/c8/44/e7922e6e2a4d63bdfbc9dc4a53e3850ab438d46cf42e6779bb15ec92c948/matplotlib-3.11.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:942b37c5db1899610bd1543ce8e13e4ecff9a4633e7f63bb6aa9205d2644ebd1", size = 10939576, upload-time = "2026-06-12T02:28:31.66Z" }, - { url = "https://files.pythonhosted.org/packages/3d/be/b1ca96003a441d619b727fee21d671fdff7a5ce2f1bb797b2521aa2f679a/matplotlib-3.11.0-cp313-cp313t-win_amd64.whl", hash = "sha256:c08e649a6313e1291e713623b97a38e5bb4aa580b2a100a94a3309bc6b9c8eb3", size = 9379519, upload-time = "2026-06-12T02:28:33.888Z" }, - { url = "https://files.pythonhosted.org/packages/e3/72/4bf3b91821c34596dd6a7bdac5836d94f744144c8208939ef49d8ec43f7e/matplotlib-3.11.0-cp313-cp313t-win_arm64.whl", hash = "sha256:2746cd2c113742ff6ce37a864c5ac5fd7aa644568f445e66166e457ac78e40e0", size = 9055456, upload-time = "2026-06-12T02:28:35.878Z" }, { url = "https://files.pythonhosted.org/packages/0f/c2/f5da6cd37ed6871f5c9b3c0507ddb69f14d6c36fac4541e4e0c60cb8cdfc/matplotlib-3.11.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:81ae77077a1e16d37a5b61096ccb07c8d90a99b518fa8256b8f21578932f2f62", size = 9434094, upload-time = "2026-06-12T02:29:09.135Z" }, { url = "https://files.pythonhosted.org/packages/f8/07/56f66906e0f87a0c6d0d0acbd34dbc9432b1931d8f26ef618bd6f92932a9/matplotlib-3.11.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:ddef37840695f5eef65f9f070fe2d2f510f584c2156203f9f622a5b0584efffd", size = 9262183, upload-time = "2026-06-12T02:29:11.283Z" }, { url = "https://files.pythonhosted.org/packages/0c/d8/c4ecab06b7ea36a570c4f3bd2d48d1799fd5d9174470e45c2194199431e7/matplotlib-3.11.0-pp311-pypy311_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:cf662e5ac5707658cb931e19972c4bd99f7b4f8b7bf79d3c821d239fa6b71e64", size = 10015653, upload-time = "2026-06-12T02:29:13.251Z" }, @@ -2589,16 +2314,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3a/cb/28ce52eb94390dda42599c98ea0204d74799e4d8047a0eb559b6fd648056/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff", size = 5009002, upload-time = "2025-11-17T22:31:52.001Z" }, { url = "https://files.pythonhosted.org/packages/f5/f0/0cfadd537c5470378b1b32bd859cf2824972174b51b873c9d95cfd7475a5/ml_dtypes-0.5.4-cp312-cp312-win_amd64.whl", hash = "sha256:c1a953995cccb9e25a4ae19e34316671e4e2edaebe4cf538229b1fc7109087b7", size = 212222, upload-time = "2025-11-17T22:31:53.742Z" }, { url = "https://files.pythonhosted.org/packages/16/2e/9acc86985bfad8f2c2d30291b27cd2bb4c74cea08695bd540906ed744249/ml_dtypes-0.5.4-cp312-cp312-win_arm64.whl", hash = "sha256:9bad06436568442575beb2d03389aa7456c690a5b05892c471215bfd8cf39460", size = 160793, upload-time = "2025-11-17T22:31:55.358Z" }, - { url = "https://files.pythonhosted.org/packages/d9/a1/4008f14bbc616cfb1ac5b39ea485f9c63031c4634ab3f4cf72e7541f816a/ml_dtypes-0.5.4-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8c760d85a2f82e2bed75867079188c9d18dae2ee77c25a54d60e9cc79be1bc48", size = 676888, upload-time = "2025-11-17T22:31:56.907Z" }, - { url = "https://files.pythonhosted.org/packages/d3/b7/dff378afc2b0d5a7d6cd9d3209b60474d9819d1189d347521e1688a60a53/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ce756d3a10d0c4067172804c9cc276ba9cc0ff47af9078ad439b075d1abdc29b", size = 5036993, upload-time = "2025-11-17T22:31:58.497Z" }, - { url = "https://files.pythonhosted.org/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:533ce891ba774eabf607172254f2e7260ba5f57bdd64030c9a4fcfbd99815d0d", size = 5010956, upload-time = "2025-11-17T22:31:59.931Z" }, - { url = "https://files.pythonhosted.org/packages/e1/8b/200088c6859d8221454825959df35b5244fa9bdf263fd0249ac5fb75e281/ml_dtypes-0.5.4-cp313-cp313-win_amd64.whl", hash = "sha256:f21c9219ef48ca5ee78402d5cc831bd58ea27ce89beda894428bc67a52da5328", size = 212224, upload-time = "2025-11-17T22:32:01.349Z" }, - { url = "https://files.pythonhosted.org/packages/8f/75/dfc3775cb36367816e678f69a7843f6f03bd4e2bcd79941e01ea960a068e/ml_dtypes-0.5.4-cp313-cp313-win_arm64.whl", hash = "sha256:35f29491a3e478407f7047b8a4834e4640a77d2737e0b294d049746507af5175", size = 160798, upload-time = "2025-11-17T22:32:02.864Z" }, - { url = "https://files.pythonhosted.org/packages/4f/74/e9ddb35fd1dd43b1106c20ced3f53c2e8e7fc7598c15638e9f80677f81d4/ml_dtypes-0.5.4-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:304ad47faa395415b9ccbcc06a0350800bc50eda70f0e45326796e27c62f18b6", size = 702083, upload-time = "2025-11-17T22:32:04.08Z" }, - { url = "https://files.pythonhosted.org/packages/74/f5/667060b0aed1aa63166b22897fdf16dca9eb704e6b4bbf86848d5a181aa7/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6a0df4223b514d799b8a1629c65ddc351b3efa833ccf7f8ea0cf654a61d1e35d", size = 5354111, upload-time = "2025-11-17T22:32:05.546Z" }, - { url = "https://files.pythonhosted.org/packages/40/49/0f8c498a28c0efa5f5c95a9e374c83ec1385ca41d0e85e7cf40e5d519a21/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:531eff30e4d368cb6255bc2328d070e35836aa4f282a0fb5f3a0cd7260257298", size = 5366453, upload-time = "2025-11-17T22:32:07.115Z" }, - { url = "https://files.pythonhosted.org/packages/8c/27/12607423d0a9c6bbbcc780ad19f1f6baa2b68b18ce4bddcdc122c4c68dc9/ml_dtypes-0.5.4-cp313-cp313t-win_amd64.whl", hash = "sha256:cb73dccfc991691c444acc8c0012bee8f2470da826a92e3a20bb333b1a7894e6", size = 225612, upload-time = "2025-11-17T22:32:08.615Z" }, - { url = "https://files.pythonhosted.org/packages/e5/80/5a5929e92c72936d5b19872c5fb8fc09327c1da67b3b68c6a13139e77e20/ml_dtypes-0.5.4-cp313-cp313t-win_arm64.whl", hash = "sha256:3bbbe120b915090d9dd1375e4684dd17a20a2491ef25d640a908281da85e73f1", size = 164145, upload-time = "2025-11-17T22:32:09.782Z" }, ] [[package]] @@ -2689,42 +2404,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ca/a4/840f5b97339e27846c46307f2530a2805d9d537d8b8bd416af031cad7fa0/multidict-6.7.1-cp312-cp312-win32.whl", hash = "sha256:28ca5ce2fd9716631133d0e9a9b9a745ad7f60bac2bccafb56aa380fc0b6c511", size = 41887, upload-time = "2026-01-26T02:44:14.245Z" }, { url = "https://files.pythonhosted.org/packages/80/31/0b2517913687895f5904325c2069d6a3b78f66cc641a86a2baf75a05dcbb/multidict-6.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:fcee94dfbd638784645b066074b338bc9cc155d4b4bffa4adce1615c5a426c19", size = 46053, upload-time = "2026-01-26T02:44:15.371Z" }, { url = "https://files.pythonhosted.org/packages/0c/5b/aba28e4ee4006ae4c7df8d327d31025d760ffa992ea23812a601d226e682/multidict-6.7.1-cp312-cp312-win_arm64.whl", hash = "sha256:ba0a9fb644d0c1a2194cf7ffb043bd852cea63a57f66fbd33959f7dae18517bf", size = 43307, upload-time = "2026-01-26T02:44:16.852Z" }, - { url = "https://files.pythonhosted.org/packages/f2/22/929c141d6c0dba87d3e1d38fbdf1ba8baba86b7776469f2bc2d3227a1e67/multidict-6.7.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:2b41f5fed0ed563624f1c17630cb9941cf2309d4df00e494b551b5f3e3d67a23", size = 76174, upload-time = "2026-01-26T02:44:18.509Z" }, - { url = "https://files.pythonhosted.org/packages/c7/75/bc704ae15fee974f8fccd871305e254754167dce5f9e42d88a2def741a1d/multidict-6.7.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:84e61e3af5463c19b67ced91f6c634effb89ef8bfc5ca0267f954451ed4bb6a2", size = 45116, upload-time = "2026-01-26T02:44:19.745Z" }, - { url = "https://files.pythonhosted.org/packages/79/76/55cd7186f498ed080a18440c9013011eb548f77ae1b297206d030eb1180a/multidict-6.7.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:935434b9853c7c112eee7ac891bc4cb86455aa631269ae35442cb316790c1445", size = 43524, upload-time = "2026-01-26T02:44:21.571Z" }, - { url = "https://files.pythonhosted.org/packages/e9/3c/414842ef8d5a1628d68edee29ba0e5bcf235dbfb3ccd3ea303a7fe8c72ff/multidict-6.7.1-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:432feb25a1cb67fe82a9680b4d65fb542e4635cb3166cd9c01560651ad60f177", size = 249368, upload-time = "2026-01-26T02:44:22.803Z" }, - { url = "https://files.pythonhosted.org/packages/f6/32/befed7f74c458b4a525e60519fe8d87eef72bb1e99924fa2b0f9d97a221e/multidict-6.7.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e82d14e3c948952a1a85503817e038cba5905a3352de76b9a465075d072fba23", size = 256952, upload-time = "2026-01-26T02:44:24.306Z" }, - { url = "https://files.pythonhosted.org/packages/03/d6/c878a44ba877f366630c860fdf74bfb203c33778f12b6ac274936853c451/multidict-6.7.1-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:4cfb48c6ea66c83bcaaf7e4dfa7ec1b6bbcf751b7db85a328902796dfde4c060", size = 240317, upload-time = "2026-01-26T02:44:25.772Z" }, - { url = "https://files.pythonhosted.org/packages/68/49/57421b4d7ad2e9e60e25922b08ceb37e077b90444bde6ead629095327a6f/multidict-6.7.1-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1d540e51b7e8e170174555edecddbd5538105443754539193e3e1061864d444d", size = 267132, upload-time = "2026-01-26T02:44:27.648Z" }, - { url = "https://files.pythonhosted.org/packages/b7/fe/ec0edd52ddbcea2a2e89e174f0206444a61440b40f39704e64dc807a70bd/multidict-6.7.1-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:273d23f4b40f3dce4d6c8a821c741a86dec62cded82e1175ba3d99be128147ed", size = 268140, upload-time = "2026-01-26T02:44:29.588Z" }, - { url = "https://files.pythonhosted.org/packages/b0/73/6e1b01cbeb458807aa0831742232dbdd1fa92bfa33f52a3f176b4ff3dc11/multidict-6.7.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9d624335fd4fa1c08a53f8b4be7676ebde19cd092b3895c421045ca87895b429", size = 254277, upload-time = "2026-01-26T02:44:30.902Z" }, - { url = "https://files.pythonhosted.org/packages/6a/b2/5fb8c124d7561a4974c342bc8c778b471ebbeb3cc17df696f034a7e9afe7/multidict-6.7.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:12fad252f8b267cc75b66e8fc51b3079604e8d43a75428ffe193cd9e2195dfd6", size = 252291, upload-time = "2026-01-26T02:44:32.31Z" }, - { url = "https://files.pythonhosted.org/packages/5a/96/51d4e4e06bcce92577fcd488e22600bd38e4fd59c20cb49434d054903bd2/multidict-6.7.1-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:03ede2a6ffbe8ef936b92cb4529f27f42be7f56afcdab5ab739cd5f27fb1cbf9", size = 250156, upload-time = "2026-01-26T02:44:33.734Z" }, - { url = "https://files.pythonhosted.org/packages/db/6b/420e173eec5fba721a50e2a9f89eda89d9c98fded1124f8d5c675f7a0c0f/multidict-6.7.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:90efbcf47dbe33dcf643a1e400d67d59abeac5db07dc3f27d6bdeae497a2198c", size = 249742, upload-time = "2026-01-26T02:44:35.222Z" }, - { url = "https://files.pythonhosted.org/packages/44/a3/ec5b5bd98f306bc2aa297b8c6f11a46714a56b1e6ef5ebda50a4f5d7c5fb/multidict-6.7.1-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:5c4b9bfc148f5a91be9244d6264c53035c8a0dcd2f51f1c3c6e30e30ebaa1c84", size = 262221, upload-time = "2026-01-26T02:44:36.604Z" }, - { url = "https://files.pythonhosted.org/packages/cd/f7/e8c0d0da0cd1e28d10e624604e1a36bcc3353aaebdfdc3a43c72bc683a12/multidict-6.7.1-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:401c5a650f3add2472d1d288c26deebc540f99e2fb83e9525007a74cd2116f1d", size = 258664, upload-time = "2026-01-26T02:44:38.008Z" }, - { url = "https://files.pythonhosted.org/packages/52/da/151a44e8016dd33feed44f730bd856a66257c1ee7aed4f44b649fb7edeb3/multidict-6.7.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:97891f3b1b3ffbded884e2916cacf3c6fc87b66bb0dde46f7357404750559f33", size = 249490, upload-time = "2026-01-26T02:44:39.386Z" }, - { url = "https://files.pythonhosted.org/packages/87/af/a3b86bf9630b732897f6fc3f4c4714b90aa4361983ccbdcd6c0339b21b0c/multidict-6.7.1-cp313-cp313-win32.whl", hash = "sha256:e1c5988359516095535c4301af38d8a8838534158f649c05dd1050222321bcb3", size = 41695, upload-time = "2026-01-26T02:44:41.318Z" }, - { url = "https://files.pythonhosted.org/packages/b2/35/e994121b0e90e46134673422dd564623f93304614f5d11886b1b3e06f503/multidict-6.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:960c83bf01a95b12b08fd54324a4eb1d5b52c88932b5cba5d6e712bb3ed12eb5", size = 45884, upload-time = "2026-01-26T02:44:42.488Z" }, - { url = "https://files.pythonhosted.org/packages/ca/61/42d3e5dbf661242a69c97ea363f2d7b46c567da8eadef8890022be6e2ab0/multidict-6.7.1-cp313-cp313-win_arm64.whl", hash = "sha256:563fe25c678aaba333d5399408f5ec3c383ca5b663e7f774dd179a520b8144df", size = 43122, upload-time = "2026-01-26T02:44:43.664Z" }, - { url = "https://files.pythonhosted.org/packages/6d/b3/e6b21c6c4f314bb956016b0b3ef2162590a529b84cb831c257519e7fde44/multidict-6.7.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:c76c4bec1538375dad9d452d246ca5368ad6e1c9039dadcf007ae59c70619ea1", size = 83175, upload-time = "2026-01-26T02:44:44.894Z" }, - { url = "https://files.pythonhosted.org/packages/fb/76/23ecd2abfe0957b234f6c960f4ade497f55f2c16aeb684d4ecdbf1c95791/multidict-6.7.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:57b46b24b5d5ebcc978da4ec23a819a9402b4228b8a90d9c656422b4bdd8a963", size = 48460, upload-time = "2026-01-26T02:44:46.106Z" }, - { url = "https://files.pythonhosted.org/packages/c4/57/a0ed92b23f3a042c36bc4227b72b97eca803f5f1801c1ab77c8a212d455e/multidict-6.7.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e954b24433c768ce78ab7929e84ccf3422e46deb45a4dc9f93438f8217fa2d34", size = 46930, upload-time = "2026-01-26T02:44:47.278Z" }, - { url = "https://files.pythonhosted.org/packages/b5/66/02ec7ace29162e447f6382c495dc95826bf931d3818799bbef11e8f7df1a/multidict-6.7.1-cp313-cp313t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:3bd231490fa7217cc832528e1cd8752a96f0125ddd2b5749390f7c3ec8721b65", size = 242582, upload-time = "2026-01-26T02:44:48.604Z" }, - { url = "https://files.pythonhosted.org/packages/58/18/64f5a795e7677670e872673aca234162514696274597b3708b2c0d276cce/multidict-6.7.1-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:253282d70d67885a15c8a7716f3a73edf2d635793ceda8173b9ecc21f2fb8292", size = 250031, upload-time = "2026-01-26T02:44:50.544Z" }, - { url = "https://files.pythonhosted.org/packages/c8/ed/e192291dbbe51a8290c5686f482084d31bcd9d09af24f63358c3d42fd284/multidict-6.7.1-cp313-cp313t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:0b4c48648d7649c9335cf1927a8b87fa692de3dcb15faa676c6a6f1f1aabda43", size = 228596, upload-time = "2026-01-26T02:44:51.951Z" }, - { url = "https://files.pythonhosted.org/packages/1e/7e/3562a15a60cf747397e7f2180b0a11dc0c38d9175a650e75fa1b4d325e15/multidict-6.7.1-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:98bc624954ec4d2c7cb074b8eefc2b5d0ce7d482e410df446414355d158fe4ca", size = 257492, upload-time = "2026-01-26T02:44:53.902Z" }, - { url = "https://files.pythonhosted.org/packages/24/02/7d0f9eae92b5249bb50ac1595b295f10e263dd0078ebb55115c31e0eaccd/multidict-6.7.1-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:1b99af4d9eec0b49927b4402bcbb58dea89d3e0db8806a4086117019939ad3dd", size = 255899, upload-time = "2026-01-26T02:44:55.316Z" }, - { url = "https://files.pythonhosted.org/packages/00/e3/9b60ed9e23e64c73a5cde95269ef1330678e9c6e34dd4eb6b431b85b5a10/multidict-6.7.1-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6aac4f16b472d5b7dc6f66a0d49dd57b0e0902090be16594dc9ebfd3d17c47e7", size = 247970, upload-time = "2026-01-26T02:44:56.783Z" }, - { url = "https://files.pythonhosted.org/packages/3e/06/538e58a63ed5cfb0bd4517e346b91da32fde409d839720f664e9a4ae4f9d/multidict-6.7.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:21f830fe223215dffd51f538e78c172ed7c7f60c9b96a2bf05c4848ad49921c3", size = 245060, upload-time = "2026-01-26T02:44:58.195Z" }, - { url = "https://files.pythonhosted.org/packages/b2/2f/d743a3045a97c895d401e9bd29aaa09b94f5cbdf1bd561609e5a6c431c70/multidict-6.7.1-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:f5dd81c45b05518b9aa4da4aa74e1c93d715efa234fd3e8a179df611cc85e5f4", size = 235888, upload-time = "2026-01-26T02:44:59.57Z" }, - { url = "https://files.pythonhosted.org/packages/38/83/5a325cac191ab28b63c52f14f1131f3b0a55ba3b9aa65a6d0bf2a9b921a0/multidict-6.7.1-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:eb304767bca2bb92fb9c5bd33cedc95baee5bb5f6c88e63706533a1c06ad08c8", size = 243554, upload-time = "2026-01-26T02:45:01.054Z" }, - { url = "https://files.pythonhosted.org/packages/20/1f/9d2327086bd15da2725ef6aae624208e2ef828ed99892b17f60c344e57ed/multidict-6.7.1-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:c9035dde0f916702850ef66460bc4239d89d08df4d02023a5926e7446724212c", size = 252341, upload-time = "2026-01-26T02:45:02.484Z" }, - { url = "https://files.pythonhosted.org/packages/e8/2c/2a1aa0280cf579d0f6eed8ee5211c4f1730bd7e06c636ba2ee6aafda302e/multidict-6.7.1-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:af959b9beeb66c822380f222f0e0a1889331597e81f1ded7f374f3ecb0fd6c52", size = 246391, upload-time = "2026-01-26T02:45:03.862Z" }, - { url = "https://files.pythonhosted.org/packages/e5/03/7ca022ffc36c5a3f6e03b179a5ceb829be9da5783e6fe395f347c0794680/multidict-6.7.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:41f2952231456154ee479651491e94118229844dd7226541788be783be2b5108", size = 243422, upload-time = "2026-01-26T02:45:05.296Z" }, - { url = "https://files.pythonhosted.org/packages/dc/1d/b31650eab6c5778aceed46ba735bd97f7c7d2f54b319fa916c0f96e7805b/multidict-6.7.1-cp313-cp313t-win32.whl", hash = "sha256:df9f19c28adcb40b6aae30bbaa1478c389efd50c28d541d76760199fc1037c32", size = 47770, upload-time = "2026-01-26T02:45:06.754Z" }, - { url = "https://files.pythonhosted.org/packages/ac/5b/2d2d1d522e51285bd61b1e20df8f47ae1a9d80839db0b24ea783b3832832/multidict-6.7.1-cp313-cp313t-win_amd64.whl", hash = "sha256:d54ecf9f301853f2c5e802da559604b3e95bb7a3b01a9c295c6ee591b9882de8", size = 53109, upload-time = "2026-01-26T02:45:08.044Z" }, - { url = "https://files.pythonhosted.org/packages/3d/a3/cc409ba012c83ca024a308516703cf339bdc4b696195644a7215a5164a24/multidict-6.7.1-cp313-cp313t-win_arm64.whl", hash = "sha256:5a37ca18e360377cfda1d62f5f382ff41f2b8c4ccb329ed974cc2e1643440118", size = 45573, upload-time = "2026-01-26T02:45:09.349Z" }, { url = "https://files.pythonhosted.org/packages/81/08/7036c080d7117f28a4af526d794aab6a84463126db031b007717c1a6676e/multidict-6.7.1-py3-none-any.whl", hash = "sha256:55d97cc6dae627efa6a6e548885712d4864b81110ac76fa4e534c03819fa4a56", size = 12319, upload-time = "2026-01-26T02:46:44.004Z" }, ] @@ -2876,26 +2555,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/25/5761d832a81df431e260719ec45de696414266613c9ee268394dd5ad8236/numpy-2.2.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fe27749d33bb772c80dcd84ae7e8df2adc920ae8297400dabec45f0dedb3f6de", size = 18313783, upload-time = "2025-05-17T21:36:56.883Z" }, { url = "https://files.pythonhosted.org/packages/57/0a/72d5a3527c5ebffcd47bde9162c39fae1f90138c961e5296491ce778e682/numpy-2.2.6-cp312-cp312-win32.whl", hash = "sha256:4eeaae00d789f66c7a25ac5f34b71a7035bb474e679f410e5e1a94deb24cf2d4", size = 6246506, upload-time = "2025-05-17T21:37:07.368Z" }, { url = "https://files.pythonhosted.org/packages/36/fa/8c9210162ca1b88529ab76b41ba02d433fd54fecaf6feb70ef9f124683f1/numpy-2.2.6-cp312-cp312-win_amd64.whl", hash = "sha256:c1f9540be57940698ed329904db803cf7a402f3fc200bfe599334c9bd84a40b2", size = 12614190, upload-time = "2025-05-17T21:37:26.213Z" }, - { url = "https://files.pythonhosted.org/packages/f9/5c/6657823f4f594f72b5471f1db1ab12e26e890bb2e41897522d134d2a3e81/numpy-2.2.6-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0811bb762109d9708cca4d0b13c4f67146e3c3b7cf8d34018c722adb2d957c84", size = 20867828, upload-time = "2025-05-17T21:37:56.699Z" }, - { url = "https://files.pythonhosted.org/packages/dc/9e/14520dc3dadf3c803473bd07e9b2bd1b69bc583cb2497b47000fed2fa92f/numpy-2.2.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:287cc3162b6f01463ccd86be154f284d0893d2b3ed7292439ea97eafa8170e0b", size = 14143006, upload-time = "2025-05-17T21:38:18.291Z" }, - { url = "https://files.pythonhosted.org/packages/4f/06/7e96c57d90bebdce9918412087fc22ca9851cceaf5567a45c1f404480e9e/numpy-2.2.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:f1372f041402e37e5e633e586f62aa53de2eac8d98cbfb822806ce4bbefcb74d", size = 5076765, upload-time = "2025-05-17T21:38:27.319Z" }, - { url = "https://files.pythonhosted.org/packages/73/ed/63d920c23b4289fdac96ddbdd6132e9427790977d5457cd132f18e76eae0/numpy-2.2.6-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:55a4d33fa519660d69614a9fad433be87e5252f4b03850642f88993f7b2ca566", size = 6617736, upload-time = "2025-05-17T21:38:38.141Z" }, - { url = "https://files.pythonhosted.org/packages/85/c5/e19c8f99d83fd377ec8c7e0cf627a8049746da54afc24ef0a0cb73d5dfb5/numpy-2.2.6-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f92729c95468a2f4f15e9bb94c432a9229d0d50de67304399627a943201baa2f", size = 14010719, upload-time = "2025-05-17T21:38:58.433Z" }, - { url = "https://files.pythonhosted.org/packages/19/49/4df9123aafa7b539317bf6d342cb6d227e49f7a35b99c287a6109b13dd93/numpy-2.2.6-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1bc23a79bfabc5d056d106f9befb8d50c31ced2fbc70eedb8155aec74a45798f", size = 16526072, upload-time = "2025-05-17T21:39:22.638Z" }, - { url = "https://files.pythonhosted.org/packages/b2/6c/04b5f47f4f32f7c2b0e7260442a8cbcf8168b0e1a41ff1495da42f42a14f/numpy-2.2.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:e3143e4451880bed956e706a3220b4e5cf6172ef05fcc397f6f36a550b1dd868", size = 15503213, upload-time = "2025-05-17T21:39:45.865Z" }, - { url = "https://files.pythonhosted.org/packages/17/0a/5cd92e352c1307640d5b6fec1b2ffb06cd0dabe7d7b8227f97933d378422/numpy-2.2.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b4f13750ce79751586ae2eb824ba7e1e8dba64784086c98cdbbcc6a42112ce0d", size = 18316632, upload-time = "2025-05-17T21:40:13.331Z" }, - { url = "https://files.pythonhosted.org/packages/f0/3b/5cba2b1d88760ef86596ad0f3d484b1cbff7c115ae2429678465057c5155/numpy-2.2.6-cp313-cp313-win32.whl", hash = "sha256:5beb72339d9d4fa36522fc63802f469b13cdbe4fdab4a288f0c441b74272ebfd", size = 6244532, upload-time = "2025-05-17T21:43:46.099Z" }, - { url = "https://files.pythonhosted.org/packages/cb/3b/d58c12eafcb298d4e6d0d40216866ab15f59e55d148a5658bb3132311fcf/numpy-2.2.6-cp313-cp313-win_amd64.whl", hash = "sha256:b0544343a702fa80c95ad5d3d608ea3599dd54d4632df855e4c8d24eb6ecfa1c", size = 12610885, upload-time = "2025-05-17T21:44:05.145Z" }, - { url = "https://files.pythonhosted.org/packages/6b/9e/4bf918b818e516322db999ac25d00c75788ddfd2d2ade4fa66f1f38097e1/numpy-2.2.6-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0bca768cd85ae743b2affdc762d617eddf3bcf8724435498a1e80132d04879e6", size = 20963467, upload-time = "2025-05-17T21:40:44Z" }, - { url = "https://files.pythonhosted.org/packages/61/66/d2de6b291507517ff2e438e13ff7b1e2cdbdb7cb40b3ed475377aece69f9/numpy-2.2.6-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:fc0c5673685c508a142ca65209b4e79ed6740a4ed6b2267dbba90f34b0b3cfda", size = 14225144, upload-time = "2025-05-17T21:41:05.695Z" }, - { url = "https://files.pythonhosted.org/packages/e4/25/480387655407ead912e28ba3a820bc69af9adf13bcbe40b299d454ec011f/numpy-2.2.6-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:5bd4fc3ac8926b3819797a7c0e2631eb889b4118a9898c84f585a54d475b7e40", size = 5200217, upload-time = "2025-05-17T21:41:15.903Z" }, - { url = "https://files.pythonhosted.org/packages/aa/4a/6e313b5108f53dcbf3aca0c0f3e9c92f4c10ce57a0a721851f9785872895/numpy-2.2.6-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:fee4236c876c4e8369388054d02d0e9bb84821feb1a64dd59e137e6511a551f8", size = 6712014, upload-time = "2025-05-17T21:41:27.321Z" }, - { url = "https://files.pythonhosted.org/packages/b7/30/172c2d5c4be71fdf476e9de553443cf8e25feddbe185e0bd88b096915bcc/numpy-2.2.6-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e1dda9c7e08dc141e0247a5b8f49cf05984955246a327d4c48bda16821947b2f", size = 14077935, upload-time = "2025-05-17T21:41:49.738Z" }, - { url = "https://files.pythonhosted.org/packages/12/fb/9e743f8d4e4d3c710902cf87af3512082ae3d43b945d5d16563f26ec251d/numpy-2.2.6-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f447e6acb680fd307f40d3da4852208af94afdfab89cf850986c3ca00562f4fa", size = 16600122, upload-time = "2025-05-17T21:42:14.046Z" }, - { url = "https://files.pythonhosted.org/packages/12/75/ee20da0e58d3a66f204f38916757e01e33a9737d0b22373b3eb5a27358f9/numpy-2.2.6-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:389d771b1623ec92636b0786bc4ae56abafad4a4c513d36a55dce14bd9ce8571", size = 15586143, upload-time = "2025-05-17T21:42:37.464Z" }, - { url = "https://files.pythonhosted.org/packages/76/95/bef5b37f29fc5e739947e9ce5179ad402875633308504a52d188302319c8/numpy-2.2.6-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8e9ace4a37db23421249ed236fdcdd457d671e25146786dfc96835cd951aa7c1", size = 18385260, upload-time = "2025-05-17T21:43:05.189Z" }, - { url = "https://files.pythonhosted.org/packages/09/04/f2f83279d287407cf36a7a8053a5abe7be3622a4363337338f2585e4afda/numpy-2.2.6-cp313-cp313t-win32.whl", hash = "sha256:038613e9fb8c72b0a41f025a7e4c3f0b7a1b5d768ece4796b674c8f3fe13efff", size = 6377225, upload-time = "2025-05-17T21:43:16.254Z" }, - { url = "https://files.pythonhosted.org/packages/67/0e/35082d13c09c02c011cf21570543d202ad929d961c02a147493cb0c2bdf5/numpy-2.2.6-cp313-cp313t-win_amd64.whl", hash = "sha256:6031dd6dfecc0cf9f668681a37648373bddd6421fff6c66ec1624eed0180ee06", size = 12771374, upload-time = "2025-05-17T21:43:35.479Z" }, { url = "https://files.pythonhosted.org/packages/9e/3b/d94a75f4dbf1ef5d321523ecac21ef23a3cd2ac8b78ae2aac40873590229/numpy-2.2.6-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:0b605b275d7bd0c640cad4e5d30fa701a8d59302e127e5f79138ad62762c3e3d", size = 21040391, upload-time = "2025-05-17T21:44:35.948Z" }, { url = "https://files.pythonhosted.org/packages/17/f4/09b2fa1b58f0fb4f7c7963a1649c64c4d315752240377ed74d9cd878f7b5/numpy-2.2.6-pp310-pypy310_pp73-macosx_14_0_x86_64.whl", hash = "sha256:7befc596a7dc9da8a337f79802ee8adb30a552a94f792b9c9d18c840055907db", size = 6786754, upload-time = "2025-05-17T21:44:47.446Z" }, { url = "https://files.pythonhosted.org/packages/af/30/feba75f143bdc868a1cc3f44ccfa6c4b9ec522b36458e738cd00f67b573f/numpy-2.2.6-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ce47521a4754c8f4593837384bd3424880629f718d87c5d44f8ed763edd63543", size = 16643476, upload-time = "2025-05-17T21:45:11.871Z" }, @@ -2936,27 +2595,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f6/81/e1b27545deedce7f4a0b348618c6b62d74e36a4dc9ccd42f3eb2f85eee32/numpy-2.4.6-cp312-cp312-win32.whl", hash = "sha256:e3eeb0aabd6bd5ce64faae67e9935203a6991b4bc2a485a767fbafb2c5125f45", size = 5962825, upload-time = "2026-05-18T23:34:20.3Z" }, { url = "https://files.pythonhosted.org/packages/ab/ca/feab00bd44aa5fe1ad2c18f08b4d3bb92e26484b0b1d1443897809ed528c/numpy-2.4.6-cp312-cp312-win_amd64.whl", hash = "sha256:d8e8286dd7cea7895157318d1b91cdacac64c479f3cbc8dce548331728484751", size = 12321687, upload-time = "2026-05-18T23:34:23.095Z" }, { url = "https://files.pythonhosted.org/packages/63/cf/5a6d34850a39d1093558564f77ee8e8e0bee5061151b8f05a55711001ec7/numpy-2.4.6-cp312-cp312-win_arm64.whl", hash = "sha256:4081eb135ac24158bd51cdfbef16f1c64df7063b1143f24731387137c092bec8", size = 10221482, upload-time = "2026-05-18T23:34:25.876Z" }, - { url = "https://files.pythonhosted.org/packages/fb/82/bdab26d7438c6791ca31b7c024ca37c1eab8b726ba236129005cd4a06e45/numpy-2.4.6-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:511dbaf848decaaaf4b4ca48032619fb3138710c4bf7da7617765edad1ef96b0", size = 16684648, upload-time = "2026-05-18T23:34:29.41Z" }, - { url = "https://files.pythonhosted.org/packages/1b/30/a80189bcc7f5e4258b3fbc3968d909d1756f54d023299ecc39ad6fdb9ef8/numpy-2.4.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:bf162abab1c1a736333192707cef898e735a5ca00f38f27eeedf44b39d9e85eb", size = 14693902, upload-time = "2026-05-18T23:34:33.013Z" }, - { url = "https://files.pythonhosted.org/packages/97/12/70b5d0d7c15e1ebb8a6a84a8caa1d19e181d84fb58bb6d70aca29099dec1/numpy-2.4.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:043191bfa8eab18c776647b62723ac9dddece59743b13f49b2016094129c2b3f", size = 5198992, upload-time = "2026-05-18T23:34:36.132Z" }, - { url = "https://files.pythonhosted.org/packages/ba/8c/ebd2a8f8a83541f8d38cc5667e8c2b69cecfd30da6e45693e8158857d44b/numpy-2.4.6-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:6180d8b35af935aed8ece3a85e0a43f87393ae0ac87c8d2c8bd2c993f7270ef3", size = 6546944, upload-time = "2026-05-18T23:34:38.484Z" }, - { url = "https://files.pythonhosted.org/packages/bb/c5/7b863a97a91671a0338f4253bd3b5a3d3852f0692dae91711c9f4a10e787/numpy-2.4.6-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:72fbe16c6fac95aedf5937fa873445cec2110be35d8a4e9433d7501fd98dae6b", size = 15669392, upload-time = "2026-05-18T23:34:41.257Z" }, - { url = "https://files.pythonhosted.org/packages/a5/9d/3584b9984ca4c047aea75214ce1a4c4c73d849bd71b604264b7f5653f8a8/numpy-2.4.6-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a7830bab239b79cda9c08c2da014761cafb48da6150e1da17ac06283f43b6089", size = 16633220, upload-time = "2026-05-18T23:34:45.075Z" }, - { url = "https://files.pythonhosted.org/packages/05/ae/7c67fba23bd98caec7c99261f3a16072ade14813486b0282cb29846de832/numpy-2.4.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ef4aea96ce4d3b074422cb4f2f64e216bf9e213004bb58ecfdf50ea02ea8eb9a", size = 17020800, upload-time = "2026-05-18T23:34:49.065Z" }, - { url = "https://files.pythonhosted.org/packages/d9/5d/3b6725cb31d983c5e66916f5d36f6d7e5521129e4c4404d64f918292a5b6/numpy-2.4.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:dfa20cc6ca228e6b155b11da03825975ce66aea520985dbbddf0f2a5a495c605", size = 18357600, upload-time = "2026-05-18T23:34:52.709Z" }, - { url = "https://files.pythonhosted.org/packages/f7/da/2ccc6c2fe8898dee01d90c75c5f5f914a23daf99e3e0f59516a08760c8b5/numpy-2.4.6-cp313-cp313-win32.whl", hash = "sha256:56b39e5e0622a09a25bf5baf62f4bcf0cb8a41ae6e2819cf49bbc5a74c083f91", size = 5961134, upload-time = "2026-05-18T23:34:55.618Z" }, - { url = "https://files.pythonhosted.org/packages/b5/cd/9cc4dc876fb065d5c220aae4d5e14826b2715331bb7618ce1fb07a679d99/numpy-2.4.6-cp313-cp313-win_amd64.whl", hash = "sha256:c4fc99836233ea196540b17ab0983aff60ed07941751930f5f4d05bc3b3b7359", size = 12318598, upload-time = "2026-05-18T23:34:58.928Z" }, - { url = "https://files.pythonhosted.org/packages/39/1e/c0bcba1f8694116485fe28fd1be698c278fcda4141c5b0e53a2aed8b12a8/numpy-2.4.6-cp313-cp313-win_arm64.whl", hash = "sha256:a7c711e21628b52034bb5ab8d1bce291f752fcc5e92accc615778acee1ff4778", size = 10222272, upload-time = "2026-05-18T23:35:02.167Z" }, - { url = "https://files.pythonhosted.org/packages/63/6d/cc5619247c8f4204e507f5883528372e4ac4bb189e579fb859a12e480b1f/numpy-2.4.6-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:112b06a867b235ef466ed3508ddf0238050df9c727cafb5301ac385b899189a1", size = 14821197, upload-time = "2026-05-18T23:35:05.468Z" }, - { url = "https://files.pythonhosted.org/packages/00/58/f1c39161c87d9e9bed660f1ed4bafc0e403d5ec9650b6dd77aead07d489b/numpy-2.4.6-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:eaf7fa2de5c0be8ae6ff8e9bea2ccd725e980541244521d8d4b5f3354a27babe", size = 5326287, upload-time = "2026-05-18T23:35:08.693Z" }, - { url = "https://files.pythonhosted.org/packages/af/57/3917ab0fd97f271a8694513581b8a36c655f111c446852c302f04ccdb6fc/numpy-2.4.6-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:7265a2f3d436e54ef9f2b52b5c937e6be778781bd97a590319d7348f1c1ca997", size = 6646763, upload-time = "2026-05-18T23:35:11.459Z" }, - { url = "https://files.pythonhosted.org/packages/eb/0f/037e64c494b67581ae18193d770adef354c41f3f2c8ebf865602d949bf8f/numpy-2.4.6-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f74a575920ab21fe304421a3fc28793d82e299cae9eccb37084e9fc7f3617c20", size = 15728070, upload-time = "2026-05-18T23:35:14.79Z" }, - { url = "https://files.pythonhosted.org/packages/21/a6/5d2bae9c9542eb4df16dc9c46dc79c186e9bad53805dfa5399a6023c6db0/numpy-2.4.6-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ede83e07a75dd06bc501566c1eca2afc0d61677c1472ac9ad93fdee6e638a48d", size = 16681752, upload-time = "2026-05-18T23:35:18.836Z" }, - { url = "https://files.pythonhosted.org/packages/92/14/23d1dfb410ae362cd59ce53e936b1513d545eb40db3949ced632e19a459e/numpy-2.4.6-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:68bb27509ac1b9a3443094260f6326150663b06abe40b73a2f81160623da5b67", size = 17086024, upload-time = "2026-05-18T23:35:22.52Z" }, - { url = "https://files.pythonhosted.org/packages/4b/6e/23595a2c642cdf3bc567877064bdd7f91c8b0038a4453cf2daf7248eafe9/numpy-2.4.6-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:a0df0043bdb289bde1f62da130d20df23d58b45429f752bc7a8fc5325a225ecd", size = 18403398, upload-time = "2026-05-18T23:35:26.398Z" }, - { url = "https://files.pythonhosted.org/packages/8a/90/0ac3bc947217e66dec77e7cbc6a1979d1af70b6461b82f620d3bccd5e4c8/numpy-2.4.6-cp313-cp313t-win32.whl", hash = "sha256:29a287e0cf63ff528da061de6b9f64a4618da591ca1046aafc54062e40ca7eab", size = 6084971, upload-time = "2026-05-18T23:35:29.387Z" }, - { url = "https://files.pythonhosted.org/packages/77/71/5673e351671a1d2bd6063b91b44f70c0affea7d1516fa7a6572941ba4aa1/numpy-2.4.6-cp313-cp313t-win_amd64.whl", hash = "sha256:25c692919ac5a01f170a3bfcd62d745b24fd095c353d50812637d6fcab442e75", size = 12458532, upload-time = "2026-05-18T23:35:32.175Z" }, - { url = "https://files.pythonhosted.org/packages/3f/88/19d3503c5046e688f049274b27a3ef3d771152fa80d3ba3d01a3dff61abe/numpy-2.4.6-cp313-cp313t-win_arm64.whl", hash = "sha256:1e978ec1e8bd0e0e4de6bb75de9d30cbb74db6b6a2bb727618613703ca0167dd", size = 10291881, upload-time = "2026-05-18T23:35:35.465Z" }, { url = "https://files.pythonhosted.org/packages/de/12/b422cc84439adc0d00de605bf4a308890ae5c26f2c71fbd73e5d08fbb0dd/numpy-2.4.6-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:55cced7c52e981362f708ad635198e97a752dfba412cc03c23bbf3bd8d5cd662", size = 16847511, upload-time = "2026-05-18T23:36:50.673Z" }, { url = "https://files.pythonhosted.org/packages/44/53/f481bef68011740f8849418d82db07230e825013f31f4eef5ba5b805316a/numpy-2.4.6-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:d6da64deb6b8ed903e7560180a92f2d804ee1ba5eeb849ac2748b8c1aba1f6d7", size = 14889064, upload-time = "2026-05-18T23:36:53.879Z" }, { url = "https://files.pythonhosted.org/packages/7f/57/42ed575c10ced8af951d426bc4e1f8aff16fd851db33f067036215a7f860/numpy-2.4.6-pp311-pypy311_pp73-macosx_14_0_arm64.whl", hash = "sha256:68a5124b13fa6cc2086764a20005d30bc0548146f7f5322f02fce212ca14317f", size = 5394157, upload-time = "2026-05-18T23:36:57.194Z" }, @@ -3145,9 +2783,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1a/3f/523fb08d9b7be15242ade6e2a641900d05c0e9cfffab8260de37a04ac0d2/nvidia_cudnn_frontend-1.22.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f64fb4e0a45b7a8bb126f91a71d8afc03facf14b82dade51744ca48cf20d2974", size = 2722597, upload-time = "2026-04-10T17:33:54.366Z" }, { url = "https://files.pythonhosted.org/packages/34/b7/35c87c334d553bd45809ec957b53f3d7dd13c5a407e853c9eea29fcc5b3c/nvidia_cudnn_frontend-1.22.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:933275df405053001888875ee75d2138b20dc4e8bf4057461b1c74ca68b0e270", size = 2863367, upload-time = "2026-04-10T17:29:22.838Z" }, { url = "https://files.pythonhosted.org/packages/4f/42/af975c8937a4c331b1215a0b2bdd2a742d792c6f777f919fd70480d63762/nvidia_cudnn_frontend-1.22.1-cp312-cp312-win_amd64.whl", hash = "sha256:2da1c277f008ee64273a48a5cb8d07efbb6d6774fdc08bd889476cce93b2f69a", size = 2310595, upload-time = "2026-04-10T17:37:24.776Z" }, - { url = "https://files.pythonhosted.org/packages/29/d3/d698b020ced27b75f1e29862f0bc26759da96fc743570a094632c0dd14a9/nvidia_cudnn_frontend-1.22.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1bc0a0ec8004998a56f222cef618243bbee779930cdf3fe1f4a7604b2b412388", size = 2722225, upload-time = "2026-04-10T17:34:42.315Z" }, - { url = "https://files.pythonhosted.org/packages/2b/04/b7b66e3a0a7b036aca0f9704b335e663609359d0e3bdd7097f6d5ccdb40a/nvidia_cudnn_frontend-1.22.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b5295f8018cd92119968d948d25b0d2d834afd552627b47450759880dfe32110", size = 2863434, upload-time = "2026-04-10T17:29:55.721Z" }, - { url = "https://files.pythonhosted.org/packages/54/8c/e9da7bbdf197397d13bb418027951e6181d0bb74c70c648fd97376bc2ed7/nvidia_cudnn_frontend-1.22.1-cp313-cp313-win_amd64.whl", hash = "sha256:7ea7887facf23d5363159073b0080cc09185e73be16ae797831d89f09b96b0f4", size = 2310490, upload-time = "2026-04-10T17:37:47.625Z" }, ] [[package]] @@ -3537,19 +3172,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a6/de/8b1895b107277d52f2b42d3a6806e69cfef0d5cf1d0ba343470b9d8e0a04/pandas-2.3.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a68e15f780eddf2b07d242e17a04aa187a7ee12b40b930bfdd78070556550e98", size = 12771002, upload-time = "2025-09-29T23:20:26.76Z" }, { url = "https://files.pythonhosted.org/packages/87/21/84072af3187a677c5893b170ba2c8fbe450a6ff911234916da889b698220/pandas-2.3.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:371a4ab48e950033bcf52b6527eccb564f52dc826c02afd9a1bc0ab731bba084", size = 13450971, upload-time = "2025-09-29T23:20:41.344Z" }, { url = "https://files.pythonhosted.org/packages/86/41/585a168330ff063014880a80d744219dbf1dd7a1c706e75ab3425a987384/pandas-2.3.3-cp312-cp312-win_amd64.whl", hash = "sha256:a16dcec078a01eeef8ee61bf64074b4e524a2a3f4b3be9326420cabe59c4778b", size = 10992722, upload-time = "2025-09-29T23:20:54.139Z" }, - { url = "https://files.pythonhosted.org/packages/cd/4b/18b035ee18f97c1040d94debd8f2e737000ad70ccc8f5513f4eefad75f4b/pandas-2.3.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:56851a737e3470de7fa88e6131f41281ed440d29a9268dcbf0002da5ac366713", size = 11544671, upload-time = "2025-09-29T23:21:05.024Z" }, - { url = "https://files.pythonhosted.org/packages/31/94/72fac03573102779920099bcac1c3b05975c2cb5f01eac609faf34bed1ca/pandas-2.3.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:bdcd9d1167f4885211e401b3036c0c8d9e274eee67ea8d0758a256d60704cfe8", size = 10680807, upload-time = "2025-09-29T23:21:15.979Z" }, - { url = "https://files.pythonhosted.org/packages/16/87/9472cf4a487d848476865321de18cc8c920b8cab98453ab79dbbc98db63a/pandas-2.3.3-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e32e7cc9af0f1cc15548288a51a3b681cc2a219faa838e995f7dc53dbab1062d", size = 11709872, upload-time = "2025-09-29T23:21:27.165Z" }, - { url = "https://files.pythonhosted.org/packages/15/07/284f757f63f8a8d69ed4472bfd85122bd086e637bf4ed09de572d575a693/pandas-2.3.3-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:318d77e0e42a628c04dc56bcef4b40de67918f7041c2b061af1da41dcff670ac", size = 12306371, upload-time = "2025-09-29T23:21:40.532Z" }, - { url = "https://files.pythonhosted.org/packages/33/81/a3afc88fca4aa925804a27d2676d22dcd2031c2ebe08aabd0ae55b9ff282/pandas-2.3.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:4e0a175408804d566144e170d0476b15d78458795bb18f1304fb94160cabf40c", size = 12765333, upload-time = "2025-09-29T23:21:55.77Z" }, - { url = "https://files.pythonhosted.org/packages/8d/0f/b4d4ae743a83742f1153464cf1a8ecfafc3ac59722a0b5c8602310cb7158/pandas-2.3.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:93c2d9ab0fc11822b5eece72ec9587e172f63cff87c00b062f6e37448ced4493", size = 13418120, upload-time = "2025-09-29T23:22:10.109Z" }, - { url = "https://files.pythonhosted.org/packages/4f/c7/e54682c96a895d0c808453269e0b5928a07a127a15704fedb643e9b0a4c8/pandas-2.3.3-cp313-cp313-win_amd64.whl", hash = "sha256:f8bfc0e12dc78f777f323f55c58649591b2cd0c43534e8355c51d3fede5f4dee", size = 10993991, upload-time = "2025-09-29T23:25:04.889Z" }, - { url = "https://files.pythonhosted.org/packages/f9/ca/3f8d4f49740799189e1395812f3bf23b5e8fc7c190827d55a610da72ce55/pandas-2.3.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:75ea25f9529fdec2d2e93a42c523962261e567d250b0013b16210e1d40d7c2e5", size = 12048227, upload-time = "2025-09-29T23:22:24.343Z" }, - { url = "https://files.pythonhosted.org/packages/0e/5a/f43efec3e8c0cc92c4663ccad372dbdff72b60bdb56b2749f04aa1d07d7e/pandas-2.3.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:74ecdf1d301e812db96a465a525952f4dde225fdb6d8e5a521d47e1f42041e21", size = 11411056, upload-time = "2025-09-29T23:22:37.762Z" }, - { url = "https://files.pythonhosted.org/packages/46/b1/85331edfc591208c9d1a63a06baa67b21d332e63b7a591a5ba42a10bb507/pandas-2.3.3-cp313-cp313t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6435cb949cb34ec11cc9860246ccb2fdc9ecd742c12d3304989017d53f039a78", size = 11645189, upload-time = "2025-09-29T23:22:51.688Z" }, - { url = "https://files.pythonhosted.org/packages/44/23/78d645adc35d94d1ac4f2a3c4112ab6f5b8999f4898b8cdf01252f8df4a9/pandas-2.3.3-cp313-cp313t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:900f47d8f20860de523a1ac881c4c36d65efcb2eb850e6948140fa781736e110", size = 12121912, upload-time = "2025-09-29T23:23:05.042Z" }, - { url = "https://files.pythonhosted.org/packages/53/da/d10013df5e6aaef6b425aa0c32e1fc1f3e431e4bcabd420517dceadce354/pandas-2.3.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:a45c765238e2ed7d7c608fc5bc4a6f88b642f2f01e70c0c23d2224dd21829d86", size = 12712160, upload-time = "2025-09-29T23:23:28.57Z" }, - { url = "https://files.pythonhosted.org/packages/bd/17/e756653095a083d8a37cbd816cb87148debcfcd920129b25f99dd8d04271/pandas-2.3.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:c4fc4c21971a1a9f4bdb4c73978c7f7256caa3e62b323f70d6cb80db583350bc", size = 13199233, upload-time = "2025-09-29T23:24:24.876Z" }, ] [[package]] @@ -3585,21 +3207,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f8/45/830bb57f533a4604b355e07edcb8ea18cf88b5f94e5fca92f27052d7c597/pandas-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f8894dc474d648fe7b6ff0ca9b0bd73950d19952bc1a6534540762c5d79d305c", size = 11950889, upload-time = "2026-05-11T18:52:50.905Z" }, { url = "https://files.pythonhosted.org/packages/b9/c5/fc1b368f303087d20e8c9bf3d6ceb186263cfac0ade735cd938538bea839/pandas-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:c7be265b62cef88e253a941e4698604973736dcfe242fdb5198f0f7bc473cdcc", size = 9755463, upload-time = "2026-05-11T18:52:53.386Z" }, { url = "https://files.pythonhosted.org/packages/86/bd/fda8f9705b1b09c6ebe14bfc0fa0e4ec8584d54ea673628f157ff55131af/pandas-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:557409bc4178e70ee8d9ddb494798e51ebf6ea59330f6be22c51bab2a7db6c49", size = 9066158, upload-time = "2026-05-11T18:52:56.038Z" }, - { url = "https://files.pythonhosted.org/packages/c5/90/62d8302883c44308c477e222c3daf7c813a34c8e96985882fbd53d964352/pandas-3.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:67b3b64c11910cfa29f4e94a14d3bff9ee693b6fc76055e7cad549cee0aec5fa", size = 10331071, upload-time = "2026-05-11T18:52:58.838Z" }, - { url = "https://files.pythonhosted.org/packages/7f/ae/6a6493c783a101f165e4356953ba3c74d6f77f0042fa7d753da9dfbb640c/pandas-3.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:39436b377d56d2a2e52d0395bdbee171f01068e99af5250509aceeb929f765c7", size = 9875690, upload-time = "2026-05-11T18:53:01.431Z" }, - { url = "https://files.pythonhosted.org/packages/62/7c/5df8e9f56c69a2769fbe9382a5ef8f2658c007e376434e1e2cbb57ad895f/pandas-3.0.3-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d4be06d68f9ddcfc645b87534911da79a8fbffc7573c80e0edcf42a5020624d8", size = 10381634, upload-time = "2026-05-11T18:53:04.393Z" }, - { url = "https://files.pythonhosted.org/packages/99/68/1237369725aa617bb358263d535803e3053fdbc593513ec5ed9c9896b5b6/pandas-3.0.3-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a4eeb6830daf35a71cc09649bd823e2b542dac246cdee9614c6e4bd65028cd6a", size = 10891243, upload-time = "2026-05-11T18:53:07.643Z" }, - { url = "https://files.pythonhosted.org/packages/25/93/77d108e8af7222b4a503ebde0e30215b1c2e4f8e53a526431890f22d5586/pandas-3.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:1928e07221f82db493cd4af1e23c1bfca524a19a4699887975bff68f49a72bfb", size = 11388659, upload-time = "2026-05-11T18:53:10.634Z" }, - { url = "https://files.pythonhosted.org/packages/d0/bd/eff5b4399f332ac386c853f6cd2bd3fa2ca0061b9f36ecd9c4d7c4265649/pandas-3.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:51b1fe551acb77dac643c6fda86084d8d446c10fe64b06a9cc29c4cc8540e7f2", size = 11942880, upload-time = "2026-05-11T18:53:13.536Z" }, - { url = "https://files.pythonhosted.org/packages/2c/20/559ace4200982c3887d0b86bfd0d856a2143ef8ddab63cc07934951a964c/pandas-3.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:a82d532a3351d435432cd913edbccaf8b8e01d4dd0e5ced5a8d2e8ecd94c7e44", size = 9757091, upload-time = "2026-05-11T18:53:16.306Z" }, - { url = "https://files.pythonhosted.org/packages/3a/66/69055a09fe200f29f922a3eeec4804611900b95f52d932ece3393c3c0c19/pandas-3.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:275c14e0fce14a2ec20eee474aecd305478ea3c1e6f6a9d8fe219a165542717e", size = 9057282, upload-time = "2026-05-11T18:53:18.768Z" }, - { url = "https://files.pythonhosted.org/packages/57/0e/efe801b0e6811e8e650cd21b7f2608e30f08a7067e2bf6e8752b0d56ee3c/pandas-3.0.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:46997386d528eb40376ecd6b033cf4a8a1e5282580f68f43de875b78cba2199d", size = 10767016, upload-time = "2026-05-11T18:53:21.227Z" }, - { url = "https://files.pythonhosted.org/packages/ea/dc/eb55135a1d5f0f0519f28da1f609a206d2cad1f9c35c32d51e38dd7261ae/pandas-3.0.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:261e308dfb22448384b7580cf719d2f998fe2966c92893c3e77d14008af1f066", size = 10420210, upload-time = "2026-05-11T18:53:23.982Z" }, - { url = "https://files.pythonhosted.org/packages/c6/3e/b1d5d955ce33ffecb407465a60bc32769d74fcf68224b7ae67ae11d4dea4/pandas-3.0.3-cp313-cp313t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:dd1a5d1def6a46002e964510bdc67c368aa0951df5d1d9f8365336f5a1f490cd", size = 10336126, upload-time = "2026-05-11T18:53:26.731Z" }, - { url = "https://files.pythonhosted.org/packages/f5/76/a01261711ab60a22d71b862f0de20e4c504bf80457270ad8cb42110f6abc/pandas-3.0.3-cp313-cp313t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d72828c20c6d6e83e1e22a6a3b47b326b71664112fa9705dcbccfd7a39b62085", size = 10728051, upload-time = "2026-05-11T18:53:29.125Z" }, - { url = "https://files.pythonhosted.org/packages/e9/21/ea191195e587b18cf682e97f433f81b2d0fbe341380e80a3e0d6e4403c8e/pandas-3.0.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:d26cbe1fcfc12e8fd900e2454163e466b2d3af84f7c75481df7683ffc073d870", size = 11350796, upload-time = "2026-05-11T18:53:32.056Z" }, - { url = "https://files.pythonhosted.org/packages/64/69/f0eaaf54939f0e8c6768fd06be9af2cef9b36048b96dfb9e1b2c685a807e/pandas-3.0.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:3e91cec1879ada0624fc3dc9953c5cbd60208e59c0db28f540c5d6d47502422f", size = 11799741, upload-time = "2026-05-11T18:53:34.985Z" }, - { url = "https://files.pythonhosted.org/packages/45/a4/865e0e510cae5fc2194de4db28be638952de942571ba9125934fd9c01d47/pandas-3.0.3-cp313-cp313t-win_amd64.whl", hash = "sha256:08d789b41f87e0905880e293cedf6197ce71fe67cc081358b1e148a491b9bd13", size = 10499958, upload-time = "2026-05-11T18:53:37.857Z" }, ] [[package]] @@ -3662,31 +3269,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/be/42/025cfe05d1be22dbfdb4f264fe9de1ccda83f66e4fc3aac94748e784af04/pillow-12.2.0-cp312-cp312-win32.whl", hash = "sha256:58f62cc0f00fd29e64b29f4fd923ffdb3859c9f9e6105bfc37ba1d08994e8940", size = 6378489, upload-time = "2026-04-01T14:43:34.601Z" }, { url = "https://files.pythonhosted.org/packages/5d/7b/25a221d2c761c6a8ae21bfa3874988ff2583e19cf8a27bf2fee358df7942/pillow-12.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:7f84204dee22a783350679a0333981df803dac21a0190d706a50475e361c93f5", size = 7084129, upload-time = "2026-04-01T14:43:37.213Z" }, { url = "https://files.pythonhosted.org/packages/10/e1/542a474affab20fd4a0f1836cb234e8493519da6b76899e30bcc5d990b8b/pillow-12.2.0-cp312-cp312-win_arm64.whl", hash = "sha256:af73337013e0b3b46f175e79492d96845b16126ddf79c438d7ea7ff27783a414", size = 2463612, upload-time = "2026-04-01T14:43:39.421Z" }, - { url = "https://files.pythonhosted.org/packages/4a/01/53d10cf0dbad820a8db274d259a37ba50b88b24768ddccec07355382d5ad/pillow-12.2.0-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:8297651f5b5679c19968abefd6bb84d95fe30ef712eb1b2d9b2d31ca61267f4c", size = 4100837, upload-time = "2026-04-01T14:43:41.506Z" }, - { url = "https://files.pythonhosted.org/packages/0f/98/f3a6657ecb698c937f6c76ee564882945f29b79bad496abcba0e84659ec5/pillow-12.2.0-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:50d8520da2a6ce0af445fa6d648c4273c3eeefbc32d7ce049f22e8b5c3daecc2", size = 4176528, upload-time = "2026-04-01T14:43:43.773Z" }, - { url = "https://files.pythonhosted.org/packages/69/bc/8986948f05e3ea490b8442ea1c1d4d990b24a7e43d8a51b2c7d8b1dced36/pillow-12.2.0-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:766cef22385fa1091258ad7e6216792b156dc16d8d3fa607e7545b2b72061f1c", size = 3640401, upload-time = "2026-04-01T14:43:45.87Z" }, - { url = "https://files.pythonhosted.org/packages/34/46/6c717baadcd62bc8ed51d238d521ab651eaa74838291bda1f86fe1f864c9/pillow-12.2.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:5d2fd0fa6b5d9d1de415060363433f28da8b1526c1c129020435e186794b3795", size = 5308094, upload-time = "2026-04-01T14:43:48.438Z" }, - { url = "https://files.pythonhosted.org/packages/71/43/905a14a8b17fdb1ccb58d282454490662d2cb89a6bfec26af6d3520da5ec/pillow-12.2.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:56b25336f502b6ed02e889f4ece894a72612fe885889a6e8c4c80239ff6e5f5f", size = 4695402, upload-time = "2026-04-01T14:43:51.292Z" }, - { url = "https://files.pythonhosted.org/packages/73/dd/42107efcb777b16fa0393317eac58f5b5cf30e8392e266e76e51cff28c3d/pillow-12.2.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f1c943e96e85df3d3478f7b691f229887e143f81fedab9b20205349ab04d73ed", size = 6280005, upload-time = "2026-04-01T14:43:54.242Z" }, - { url = "https://files.pythonhosted.org/packages/a8/68/b93e09e5e8549019e61acf49f65b1a8530765a7f812c77a7461bca7e4494/pillow-12.2.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:03f6fab9219220f041c74aeaa2939ff0062bd5c364ba9ce037197f4c6d498cd9", size = 8090669, upload-time = "2026-04-01T14:43:57.335Z" }, - { url = "https://files.pythonhosted.org/packages/4b/6e/3ccb54ce8ec4ddd1accd2d89004308b7b0b21c4ac3d20fa70af4760a4330/pillow-12.2.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5cdfebd752ec52bf5bb4e35d9c64b40826bc5b40a13df7c3cda20a2c03a0f5ed", size = 6395194, upload-time = "2026-04-01T14:43:59.864Z" }, - { url = "https://files.pythonhosted.org/packages/67/ee/21d4e8536afd1a328f01b359b4d3997b291ffd35a237c877b331c1c3b71c/pillow-12.2.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:eedf4b74eda2b5a4b2b2fb4c006d6295df3bf29e459e198c90ea48e130dc75c3", size = 7082423, upload-time = "2026-04-01T14:44:02.74Z" }, - { url = "https://files.pythonhosted.org/packages/78/5f/e9f86ab0146464e8c133fe85df987ed9e77e08b29d8d35f9f9f4d6f917ba/pillow-12.2.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:00a2865911330191c0b818c59103b58a5e697cae67042366970a6b6f1b20b7f9", size = 6505667, upload-time = "2026-04-01T14:44:05.381Z" }, - { url = "https://files.pythonhosted.org/packages/ed/1e/409007f56a2fdce61584fd3acbc2bbc259857d555196cedcadc68c015c82/pillow-12.2.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:1e1757442ed87f4912397c6d35a0db6a7b52592156014706f17658ff58bbf795", size = 7208580, upload-time = "2026-04-01T14:44:08.39Z" }, - { url = "https://files.pythonhosted.org/packages/23/c4/7349421080b12fb35414607b8871e9534546c128a11965fd4a7002ccfbee/pillow-12.2.0-cp313-cp313-win32.whl", hash = "sha256:144748b3af2d1b358d41286056d0003f47cb339b8c43a9ea42f5fea4d8c66b6e", size = 6375896, upload-time = "2026-04-01T14:44:11.197Z" }, - { url = "https://files.pythonhosted.org/packages/3f/82/8a3739a5e470b3c6cbb1d21d315800d8e16bff503d1f16b03a4ec3212786/pillow-12.2.0-cp313-cp313-win_amd64.whl", hash = "sha256:390ede346628ccc626e5730107cde16c42d3836b89662a115a921f28440e6a3b", size = 7081266, upload-time = "2026-04-01T14:44:13.947Z" }, - { url = "https://files.pythonhosted.org/packages/c3/25/f968f618a062574294592f668218f8af564830ccebdd1fa6200f598e65c5/pillow-12.2.0-cp313-cp313-win_arm64.whl", hash = "sha256:8023abc91fba39036dbce14a7d6535632f99c0b857807cbbbf21ecc9f4717f06", size = 2463508, upload-time = "2026-04-01T14:44:16.312Z" }, - { url = "https://files.pythonhosted.org/packages/4d/a4/b342930964e3cb4dce5038ae34b0eab4653334995336cd486c5a8c25a00c/pillow-12.2.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:042db20a421b9bafecc4b84a8b6e444686bd9d836c7fd24542db3e7df7baad9b", size = 5309927, upload-time = "2026-04-01T14:44:18.89Z" }, - { url = "https://files.pythonhosted.org/packages/9f/de/23198e0a65a9cf06123f5435a5d95cea62a635697f8f03d134d3f3a96151/pillow-12.2.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:dd025009355c926a84a612fecf58bb315a3f6814b17ead51a8e48d3823d9087f", size = 4698624, upload-time = "2026-04-01T14:44:21.115Z" }, - { url = "https://files.pythonhosted.org/packages/01/a6/1265e977f17d93ea37aa28aa81bad4fa597933879fac2520d24e021c8da3/pillow-12.2.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:88ddbc66737e277852913bd1e07c150cc7bb124539f94c4e2df5344494e0a612", size = 6321252, upload-time = "2026-04-01T14:44:23.663Z" }, - { url = "https://files.pythonhosted.org/packages/3c/83/5982eb4a285967baa70340320be9f88e57665a387e3a53a7f0db8231a0cd/pillow-12.2.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d362d1878f00c142b7e1a16e6e5e780f02be8195123f164edf7eddd911eefe7c", size = 8126550, upload-time = "2026-04-01T14:44:26.772Z" }, - { url = "https://files.pythonhosted.org/packages/4e/48/6ffc514adce69f6050d0753b1a18fd920fce8cac87620d5a31231b04bfc5/pillow-12.2.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2c727a6d53cb0018aadd8018c2b938376af27914a68a492f59dfcaca650d5eea", size = 6433114, upload-time = "2026-04-01T14:44:29.615Z" }, - { url = "https://files.pythonhosted.org/packages/36/a3/f9a77144231fb8d40ee27107b4463e205fa4677e2ca2548e14da5cf18dce/pillow-12.2.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:efd8c21c98c5cc60653bcb311bef2ce0401642b7ce9d09e03a7da87c878289d4", size = 7115667, upload-time = "2026-04-01T14:44:32.773Z" }, - { url = "https://files.pythonhosted.org/packages/c1/fc/ac4ee3041e7d5a565e1c4fd72a113f03b6394cc72ab7089d27608f8aaccb/pillow-12.2.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9f08483a632889536b8139663db60f6724bfcb443c96f1b18855860d7d5c0fd4", size = 6538966, upload-time = "2026-04-01T14:44:35.252Z" }, - { url = "https://files.pythonhosted.org/packages/c0/a8/27fb307055087f3668f6d0a8ccb636e7431d56ed0750e07a60547b1e083e/pillow-12.2.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:dac8d77255a37e81a2efcbd1fc05f1c15ee82200e6c240d7e127e25e365c39ea", size = 7238241, upload-time = "2026-04-01T14:44:37.875Z" }, - { url = "https://files.pythonhosted.org/packages/ad/4b/926ab182c07fccae9fcb120043464e1ff1564775ec8864f21a0ebce6ac25/pillow-12.2.0-cp313-cp313t-win32.whl", hash = "sha256:ee3120ae9dff32f121610bb08e4313be87e03efeadfc6c0d18f89127e24d0c24", size = 6379592, upload-time = "2026-04-01T14:44:40.336Z" }, - { url = "https://files.pythonhosted.org/packages/c2/c4/f9e476451a098181b30050cc4c9a3556b64c02cf6497ea421ac047e89e4b/pillow-12.2.0-cp313-cp313t-win_amd64.whl", hash = "sha256:325ca0528c6788d2a6c3d40e3568639398137346c3d6e66bb61db96b96511c98", size = 7085542, upload-time = "2026-04-01T14:44:43.251Z" }, - { url = "https://files.pythonhosted.org/packages/00/a4/285f12aeacbe2d6dc36c407dfbbe9e96d4a80b0fb710a337f6d2ad978c75/pillow-12.2.0-cp313-cp313t-win_arm64.whl", hash = "sha256:2e5a76d03a6c6dcef67edabda7a52494afa4035021a79c8558e14af25313d453", size = 2465765, upload-time = "2026-04-01T14:44:45.996Z" }, { url = "https://files.pythonhosted.org/packages/4e/b7/2437044fb910f499610356d1352e3423753c98e34f915252aafecc64889f/pillow-12.2.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:0538bd5e05efec03ae613fd89c4ce0368ecd2ba239cc25b9f9be7ed426b0af1f", size = 5273969, upload-time = "2026-04-01T14:45:55.538Z" }, { url = "https://files.pythonhosted.org/packages/f6/f4/8316e31de11b780f4ac08ef3654a75555e624a98db1056ecb2122d008d5a/pillow-12.2.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:394167b21da716608eac917c60aa9b969421b5dcbbe02ae7f013e7b85811c69d", size = 4659674, upload-time = "2026-04-01T14:45:58.093Z" }, { url = "https://files.pythonhosted.org/packages/d4/37/664fca7201f8bb2aa1d20e2c3d5564a62e6ae5111741966c8319ca802361/pillow-12.2.0-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5d04bfa02cc2d23b497d1e90a0f927070043f6cbf303e738300532379a4b4e0f", size = 5288479, upload-time = "2026-04-01T14:46:01.141Z" }, @@ -3811,40 +3393,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cf/74/66bd798b5b3be70aa1b391f5cc9d6a0a5532d7fd3b19ec0b213e72e6ad9d/propcache-0.5.2-cp312-cp312-win32.whl", hash = "sha256:8c7972d8f193740d9175f0998ab38717e6cd322d5935c5b0fef8c0d323fd9031", size = 39018, upload-time = "2026-05-08T21:00:31.622Z" }, { url = "https://files.pythonhosted.org/packages/61/7c/5c0d34aa3024694d6dcb9271cdbdd08c4e47c1c0ad95ec7e7bc74cdea145/propcache-0.5.2-cp312-cp312-win_amd64.whl", hash = "sha256:d9ee8826a7d47863a08ac44e1a5f611a462eefc3a194b492da242128bec75b42", size = 42322, upload-time = "2026-05-08T21:00:32.918Z" }, { url = "https://files.pythonhosted.org/packages/4d/91/875812f1a3feb20ceba818ef39fbe4d92f1081e04ac815c822496d0d038b/propcache-0.5.2-cp312-cp312-win_arm64.whl", hash = "sha256:2800a4a8ead6b28cccd1ec54b59346f0def7922ee1c7598e8499c733cfbb7c84", size = 38172, upload-time = "2026-05-08T21:00:35.124Z" }, - { url = "https://files.pythonhosted.org/packages/c5/09/f049e45385503fe67db75a6b6186a7b9f0c3930366dc960522c312a825b1/propcache-0.5.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:099aaf4b4d1a02265b92a977edf00b5c4f63b3b17ac6de39b0d637c9cac0188a", size = 94457, upload-time = "2026-05-08T21:00:36.355Z" }, - { url = "https://files.pythonhosted.org/packages/6b/65/83d1d05655baf63113731bd5a1008435e14f8d1e5a06cbe4ec5b23ad7a31/propcache-0.5.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:68ce1c44c7a813a7f71ea04315a8c7b330b63db99d059a797a4651bb6f69f117", size = 53835, upload-time = "2026-05-08T21:00:38.072Z" }, - { url = "https://files.pythonhosted.org/packages/a9/12/a6ba6482bb5ea3260c000c9b20881c95fa11c6b30173715668259f844ed7/propcache-0.5.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:fc299c129490f55f254cd90be0deca4764e36e9a7c08b4aa588479a3bbed3098", size = 54545, upload-time = "2026-05-08T21:00:39.319Z" }, - { url = "https://files.pythonhosted.org/packages/a9/19/7fa086f5764c59ec8a8e157cd93aa8497acc00aba9dcdec56bfffb32602d/propcache-0.5.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a6ae2198be502c10f09b2516e7b5d019816924bc3183a43ce792a7bd6625e6f4", size = 59886, upload-time = "2026-05-08T21:00:40.621Z" }, - { url = "https://files.pythonhosted.org/packages/a1/e4/5d7663dc8235956c8f5281698a3af1d351d8820341ddd890f59d9a9127f2/propcache-0.5.2-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6041d31504dc1779d700e1edcfb08eea334b357620b06681a4eabb57a74e574e", size = 63261, upload-time = "2026-05-08T21:00:41.775Z" }, - { url = "https://files.pythonhosted.org/packages/4a/4a/15a03adee24d6350da4292caeac44c34c033d2afe5e87eb370f38854560f/propcache-0.5.2-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f7eabc04151c78a9f4d5bbb5f1faf571e4defeb4b585e0fe95b60ff2dbe4d3d7", size = 64184, upload-time = "2026-05-08T21:00:43.018Z" }, - { url = "https://files.pythonhosted.org/packages/8b/c6/979176efdaa3d239e36d503d5af63a0a773b36662ed8f52e5b6a6d9fd40e/propcache-0.5.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4db0ba63d693afd40d249bd93f842b5f144f8fcbb83de05660373bcf30517b1d", size = 61534, upload-time = "2026-05-08T21:00:44.507Z" }, - { url = "https://files.pythonhosted.org/packages/c8/22/63e8cd1bae4c2d2be6493b6b7d10566ddafad88137cfbc99964a1119853c/propcache-0.5.2-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:1dbcf7675229b35d31abb6547d8ebc8c27a830ac3f9a794edff6254873ec7c0a", size = 61500, upload-time = "2026-05-08T21:00:45.796Z" }, - { url = "https://files.pythonhosted.org/packages/60/5a/28e5d9acbac1cc9ccb67045e8c1b943aa8d79fdf39c93bd73cacd68008ea/propcache-0.5.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d310c013aad2c72f1c3f2f8dd3279d460a858c551f97aeb8c63e4693cca7b4d2", size = 59994, upload-time = "2026-05-08T21:00:47.093Z" }, - { url = "https://files.pythonhosted.org/packages/f3/40/db650677f554a95b9c01a7c9d93d629e93a15562f5deb4573c9ee136fed2/propcache-0.5.2-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:06187263ddad280d05b4d8a8b3bb7d164cbebd469236544a42e6d9b28ac6a4fa", size = 56884, upload-time = "2026-05-08T21:00:48.376Z" }, - { url = "https://files.pythonhosted.org/packages/80/45/70b39b89516ff8b96bf732fa6fded8cef20f293cb1508690101c3c07ec51/propcache-0.5.2-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:3115559b8effafd63b142ea5ed53d63a16ea6469cbc63dce4ee194b42db5d853", size = 63464, upload-time = "2026-05-08T21:00:49.954Z" }, - { url = "https://files.pythonhosted.org/packages/f9/e2/fa59d3a89eac5534293124af4f1d0d0ada091ce4a0ab4610ce03fd2bdd8d/propcache-0.5.2-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:c60462af8e6dc30c35407c7237ea908d777b22862bbee27bc4699c0d8bcdc45a", size = 61588, upload-time = "2026-05-08T21:00:51.281Z" }, - { url = "https://files.pythonhosted.org/packages/0b/97/efb547a55c4bc7381cfb202d6a2239ac621045277bc1ea5dfd3a7f0516c0/propcache-0.5.2-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:40314bca9ac559716fe374094fc81c11dcc34b64fd6c585360f5775690505704", size = 64667, upload-time = "2026-05-08T21:00:52.602Z" }, - { url = "https://files.pythonhosted.org/packages/92/56/f5c7d9b4b7595d5127da38974d791b2153f3d1eae6c674af3583ace92ad3/propcache-0.5.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:cfa21e036ce1e1db2be04ba3b85d2df1bb1702fa01932d984c5464c665228ff4", size = 62463, upload-time = "2026-05-08T21:00:54.303Z" }, - { url = "https://files.pythonhosted.org/packages/bd/3b/484a3a65fc9f9f60c41dcd17b428bace5389544e2c680994534a20755066/propcache-0.5.2-cp313-cp313-win32.whl", hash = "sha256:f156a3529f38063b6dbaf356e15602a7f95f8055b1295a438433a6386f10463d", size = 38621, upload-time = "2026-05-08T21:00:55.808Z" }, - { url = "https://files.pythonhosted.org/packages/1c/fd/3f0f10dba4dabad3bf53102be007abf55481067952bde0fdddff439e7c61/propcache-0.5.2-cp313-cp313-win_amd64.whl", hash = "sha256:dfed59d0a5aeb01e242e66ff0300bc4a265a7c05f612d30016f0b60b1017d757", size = 41649, upload-time = "2026-05-08T21:00:57.061Z" }, - { url = "https://files.pythonhosted.org/packages/90/ec/6ce619cc32bb500a482f811f9cd509368b4e58e638d13f2c68f370d6b475/propcache-0.5.2-cp313-cp313-win_arm64.whl", hash = "sha256:ba338430e87ceb9c8f0cf754de38a9860560261e56c00376debd628698a7364f", size = 37636, upload-time = "2026-05-08T21:00:58.646Z" }, - { url = "https://files.pythonhosted.org/packages/1b/82/c1d268bbbf2ef981c5bf0fbbe746db617c66e3bcefe431a1aa8943fbe23a/propcache-0.5.2-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:a592f5f3da71c8691c788c13cb6734b6d17663d2e1cb8caddf0673d01ef8847d", size = 98872, upload-time = "2026-05-08T21:00:59.889Z" }, - { url = "https://files.pythonhosted.org/packages/f4/d4/52c871e73e864e6b34c0e2d58ac1ec5ccd149497ddc7ad2137ae98323a35/propcache-0.5.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:6a997d0489e9668a384fcfd5061b857aa5361de73191cac204d04b889cfbbafa", size = 56257, upload-time = "2026-05-08T21:01:01.195Z" }, - { url = "https://files.pythonhosted.org/packages/67/f0/9b90ca2a210b3d09bcfcd96ecd0f55545c091535abce2a45de2775cfd357/propcache-0.5.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:10734b5484ea113152ee25a91dccedf81631791805d2c9ccb054958e51842c94", size = 56696, upload-time = "2026-05-08T21:01:02.941Z" }, - { url = "https://files.pythonhosted.org/packages/9d/0e/6e9d4ba07c8e56e21ddec1e75f12148142b21ca83a51871babce095334f4/propcache-0.5.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cafca7e56c12bb02ae16d283742bef25a61122e9dab2b5b3f2ccbe589ce32164", size = 62378, upload-time = "2026-05-08T21:01:04.475Z" }, - { url = "https://files.pythonhosted.org/packages/65/19/c10badaa463dde8a27ce884f8ee2ec37e6035b7c9f5ff0c8f74f06f08dac/propcache-0.5.2-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f064f8d2b59177878b7615df1735cd8fe3462ed6be8c7b217d17a276489c2b7f", size = 65283, upload-time = "2026-05-08T21:01:05.959Z" }, - { url = "https://files.pythonhosted.org/packages/b0/b6/93bea99ca80e19cef6512a8580e5b7857bbe09422d9daa7fd4ef5723306c/propcache-0.5.2-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f78abfa8dfc32376fd1aacf597b2f2fbbe0ea751419aee718af5d4f82537ef8c", size = 66616, upload-time = "2026-05-08T21:01:07.228Z" }, - { url = "https://files.pythonhosted.org/packages/83/e4/5c7462e50625f051f37fb38b8224f7639f667184bbd34424ec83819bb1b7/propcache-0.5.2-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f7467da8a9822bf1a55336f877340c5bcbd3c482afc43a99771169f74a26dedc", size = 63773, upload-time = "2026-05-08T21:01:08.514Z" }, - { url = "https://files.pythonhosted.org/packages/ca/b6/99238894047b13c823be25027e736626cd414a52a5e30d2c3347c2733529/propcache-0.5.2-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a6ddc6ac9e25de626c1f129c1b467d7ecd33ce2237d3fd0c4e429feef0a7ee1f", size = 63664, upload-time = "2026-05-08T21:01:09.874Z" }, - { url = "https://files.pythonhosted.org/packages/85/1e/a3a1a63116a2b8edb415a8bb9a6f0c34bd03830b1e18e8ce2904e1dc1cf4/propcache-0.5.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:2f22cbbac9e26a8e864c0985ff1268d5d939d53d9d9411a9824279097e03a2cb", size = 62643, upload-time = "2026-05-08T21:01:11.132Z" }, - { url = "https://files.pythonhosted.org/packages/e4/03/893cf147de2fc6543c5eaa07ad833170e7e2a2385725bbebe8c0503723bb/propcache-0.5.2-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:fc76378c62a0f04d0cd82fbb1a2cd2d7e28fcb40d5873f28a6c44e388aaa2751", size = 59595, upload-time = "2026-05-08T21:01:12.387Z" }, - { url = "https://files.pythonhosted.org/packages/86/3b/04c1a2e12c57766568ba75ba72b3bf2042818d4c1425fab6fc07155c7cff/propcache-0.5.2-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:acd2c8edba48e31e58a363b8cf4e5c7db3b04b3f9e371f601df30d9b0d244836", size = 65711, upload-time = "2026-05-08T21:01:13.676Z" }, - { url = "https://files.pythonhosted.org/packages/1c/34/80f8d0099f8d6bacc4de1624c85672681c8cd1149ca2da0e38fd120b817f/propcache-0.5.2-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:452b5065457eb9991ec5eb38ff41d6cd4c991c9ac7c531c4d5849ae473a9a13f", size = 64247, upload-time = "2026-05-08T21:01:14.936Z" }, - { url = "https://files.pythonhosted.org/packages/f3/1a/8b08f3a5f1037e9e370c55883ceeeee0f6dd0416fb2d2d67b8bfc91f2a79/propcache-0.5.2-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:3430bb2bfe1331885c427745a751e774ee679fd4344f80b97bf879815fe8fa55", size = 67102, upload-time = "2026-05-08T21:01:16.281Z" }, - { url = "https://files.pythonhosted.org/packages/34/68/8bdb7bb7756d76e005490649d10e4a8369e610c74d619f71e1aedf889e9c/propcache-0.5.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:cef6cea3922890dd6c9654971001fa797b526c16ab5e1e46c05fd6f877be7568", size = 64964, upload-time = "2026-05-08T21:01:17.57Z" }, - { url = "https://files.pythonhosted.org/packages/0a/aa/50fb0b5d3968b61a510926ff8b8465f1d6e976b3ab74496d7a4b9fc42515/propcache-0.5.2-cp313-cp313t-win32.whl", hash = "sha256:72d61e16dd78228b58c5d47be830ff3da7e5f139abdf0aef9d86cde1c5cf2191", size = 42546, upload-time = "2026-05-08T21:01:18.946Z" }, - { url = "https://files.pythonhosted.org/packages/ae/4c/0ddbae64321bd4a95bcbfc19307238016b5b1fee645c84626c8d539e5b74/propcache-0.5.2-cp313-cp313t-win_amd64.whl", hash = "sha256:0958834041a0166d343b8d2cedcd8bcbaeb4fdbe0cf08320c5379f143c3be6e7", size = 46330, upload-time = "2026-05-08T21:01:20.162Z" }, - { url = "https://files.pythonhosted.org/packages/00/d9/9cddc8efb78d8af264c5ec9f6d10b62f57c515feda8d321595f56010fb23/propcache-0.5.2-cp313-cp313t-win_arm64.whl", hash = "sha256:6de8bd93ddde9b992cf2b2e0d796d501a19026b5b9fd87356d7d0779531a8d96", size = 40521, upload-time = "2026-05-08T21:01:21.399Z" }, { url = "https://files.pythonhosted.org/packages/3a/ed/1cdcab6ba3d6ab7feca11fc14f0eeea80755bb53ef4e892079f31b10a25f/propcache-0.5.2-py3-none-any.whl", hash = "sha256:be1ddfcbb376e3de5d2e2db1d58d6d67463e6b4f9f040c000de8e300295465fe", size = 14036, upload-time = "2026-05-08T21:02:10.673Z" }, ] @@ -3869,12 +3417,6 @@ version = "7.2.2" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/51/08/510cbdb69c25a96f4ae523f733cdc963ae654904e8db864c07585ef99875/psutil-7.2.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2edccc433cbfa046b980b0df0171cd25bcaeb3a68fe9022db0979e7aa74a826b", size = 130595, upload-time = "2026-01-28T18:14:57.293Z" }, - { url = "https://files.pythonhosted.org/packages/d6/f5/97baea3fe7a5a9af7436301f85490905379b1c6f2dd51fe3ecf24b4c5fbf/psutil-7.2.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e78c8603dcd9a04c7364f1a3e670cea95d51ee865e4efb3556a3a63adef958ea", size = 131082, upload-time = "2026-01-28T18:14:59.732Z" }, - { url = "https://files.pythonhosted.org/packages/37/d6/246513fbf9fa174af531f28412297dd05241d97a75911ac8febefa1a53c6/psutil-7.2.2-cp313-cp313t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1a571f2330c966c62aeda00dd24620425d4b0cc86881c89861fbc04549e5dc63", size = 181476, upload-time = "2026-01-28T18:15:01.884Z" }, - { url = "https://files.pythonhosted.org/packages/b8/b5/9182c9af3836cca61696dabe4fd1304e17bc56cb62f17439e1154f225dd3/psutil-7.2.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:917e891983ca3c1887b4ef36447b1e0873e70c933afc831c6b6da078ba474312", size = 184062, upload-time = "2026-01-28T18:15:04.436Z" }, - { url = "https://files.pythonhosted.org/packages/16/ba/0756dca669f5a9300d0cbcbfae9a4c30e446dfc7440ffe43ded5724bfd93/psutil-7.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:ab486563df44c17f5173621c7b198955bd6b613fb87c71c161f827d3fb149a9b", size = 139893, upload-time = "2026-01-28T18:15:06.378Z" }, - { url = "https://files.pythonhosted.org/packages/1c/61/8fa0e26f33623b49949346de05ec1ddaad02ed8ba64af45f40a147dbfa97/psutil-7.2.2-cp313-cp313t-win_arm64.whl", hash = "sha256:ae0aefdd8796a7737eccea863f80f81e468a1e4cf14d926bd9b6f5f2d5f90ca9", size = 135589, upload-time = "2026-01-28T18:15:08.03Z" }, { url = "https://files.pythonhosted.org/packages/e7/36/5ee6e05c9bd427237b11b3937ad82bb8ad2752d72c6969314590dd0c2f6e/psutil-7.2.2-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486", size = 129090, upload-time = "2026-01-28T18:15:22.168Z" }, { url = "https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979", size = 129859, upload-time = "2026-01-28T18:15:23.795Z" }, { url = "https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9", size = 155560, upload-time = "2026-01-28T18:15:25.976Z" }, @@ -3930,20 +3472,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5f/e8/f88ce625fe8babaae64e8db2d417c7653adb3019b08aae85c5ed787dc816/pyarrow-24.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3b13dedfe76a0ad2d1d859b0811b53827a4e9d93a0bcb05cf59333ab4980cc7e", size = 49376032, upload-time = "2026-04-21T10:47:48.967Z" }, { url = "https://files.pythonhosted.org/packages/36/7a/82c363caa145fff88fb475da50d3bf52bb024f61917be5424c3392eaf878/pyarrow-24.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:25ea65d868eb04015cd18e6df2fbe98f07e5bda2abefabcb88fce39a947716f6", size = 51929490, upload-time = "2026-04-21T10:47:55.981Z" }, { url = "https://files.pythonhosted.org/packages/66/1c/e3e72c8014ad2743ca64a701652c733cc5cbcee15c0463a32a8c55518d9e/pyarrow-24.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:295f0a7f2e242dabd513737cf076007dc5b2d59237e3eca37b05c0c6446f3826", size = 27355660, upload-time = "2026-04-21T10:48:01.718Z" }, - { url = "https://files.pythonhosted.org/packages/6f/d3/a1abf004482026ddc17f4503db227787fa3cfe41ec5091ff20e4fea55e57/pyarrow-24.0.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:02b001b3ed4723caa44f6cd1af2d5c86aa2cf9971dacc2ffa55b21237713dfba", size = 34976759, upload-time = "2026-04-21T10:48:07.258Z" }, - { url = "https://files.pythonhosted.org/packages/4f/4a/34f0a36d28a2dd32225301b79daad44e243dc1a2bb77d43b60749be255c4/pyarrow-24.0.0-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:04920d6a71aabd08a0417709efce97d45ea8e6fb733d9ca9ecffb13c67839f68", size = 36658471, upload-time = "2026-04-21T10:48:13.347Z" }, - { url = "https://files.pythonhosted.org/packages/1f/78/543b94712ae8bb1a6023bcc1acf1a740fbff8286747c289cd9468fced2a5/pyarrow-24.0.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:a964266397740257f16f7bb2e4f08a0c81454004beab8ff59dd531b73610e9f2", size = 45675981, upload-time = "2026-04-21T10:48:20.201Z" }, - { url = "https://files.pythonhosted.org/packages/84/9f/8fb7c222b100d314137fa40ec050de56cd8c6d957d1cfff685ce72f15b17/pyarrow-24.0.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6f066b179d68c413374294bc1735f68475457c933258df594443bb9d88ddc2a0", size = 48859172, upload-time = "2026-04-21T10:48:27.541Z" }, - { url = "https://files.pythonhosted.org/packages/a7/d3/1ea72538e6c8b3b475ed78d1049a2c518e655761ea50fe1171fc855fcab7/pyarrow-24.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:1183baeb14c5f587b1ec52831e665718ce632caab84b7cd6b85fd44f96114495", size = 49385733, upload-time = "2026-04-21T10:48:34.7Z" }, - { url = "https://files.pythonhosted.org/packages/c3/be/c3d8b06a1ba35f2260f8e1f771abbee7d5e345c0937aab90675706b1690a/pyarrow-24.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:806f24b4085453c197a5078218d1ee08783ebbba271badd153d1ae22a3ee804f", size = 51934335, upload-time = "2026-04-21T10:48:42.099Z" }, - { url = "https://files.pythonhosted.org/packages/9c/62/89e07a1e7329d2cde3e3c6994ba0839a24977a2beda8be6005ea3d860b99/pyarrow-24.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:e4505fc6583f7b05ab854934896bcac8253b04ac1171a77dfb73efef92076d91", size = 27271748, upload-time = "2026-04-21T10:49:42.532Z" }, - { url = "https://files.pythonhosted.org/packages/17/1a/cff3a59f80b5b1658549d46611b67163f65e0664431c076ad728bf9d5af4/pyarrow-24.0.0-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:1a4e45017efbf115032e4475ee876d525e0e36c742214fbe405332480ecd6275", size = 35238554, upload-time = "2026-04-21T10:48:48.526Z" }, - { url = "https://files.pythonhosted.org/packages/a8/99/cce0f42a327bfef2c420fb6078a3eb834826e5d6697bf3009fe11d2ad051/pyarrow-24.0.0-cp313-cp313t-macosx_12_0_x86_64.whl", hash = "sha256:7986f1fa71cee060ad00758bcc79d3a93bab8559bf978fab9e53472a2e25a17b", size = 36782301, upload-time = "2026-04-21T10:48:55.181Z" }, - { url = "https://files.pythonhosted.org/packages/2a/66/8e560d5ff6793ca29aca213c53eec0dd482dd46cb93b2819e5aab52e4252/pyarrow-24.0.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:d3e0b61e8efb24ed38898e5cdc5fffa9124be480008d401a1f8071500494ae42", size = 45721929, upload-time = "2026-04-21T10:49:03.676Z" }, - { url = "https://files.pythonhosted.org/packages/27/0c/a26e25505d030716e078d9f16eb74973cbf0b33b672884e9f9da1c83b871/pyarrow-24.0.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:55a3bc1e3df3b5567b7d27ef551b2283f0c68a5e86f1cd56abc569da4f31335b", size = 48825365, upload-time = "2026-04-21T10:49:11.714Z" }, - { url = "https://files.pythonhosted.org/packages/5f/eb/771f9ecb0c65e73fe9dccdd1717901b9594f08c4515d000c7c62df573811/pyarrow-24.0.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:641f795b361874ac9da5294f8f443dfdbee355cf2bd9e3b8d97aaac2306b9b37", size = 49451819, upload-time = "2026-04-21T10:49:21.474Z" }, - { url = "https://files.pythonhosted.org/packages/48/da/61ae89a88732f5a785646f3ec6125dbb640fa98a540eb2b9889caa561403/pyarrow-24.0.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8adc8e6ce5fccf5dc707046ae4914fd537def529709cc0d285d37a7f9cd442ca", size = 51909252, upload-time = "2026-04-21T10:49:31.164Z" }, - { url = "https://files.pythonhosted.org/packages/cb/1a/8dd5cafab7b66573fa91c03d06d213356ad4edd71813aa75e08ce2b3a844/pyarrow-24.0.0-cp313-cp313t-win_amd64.whl", hash = "sha256:9b18371ad2f44044b81a8d23bc2d8a9b6a6226dca775e8e16cfee640473d6c5d", size = 27388127, upload-time = "2026-04-21T10:49:37.334Z" }, ] [[package]] @@ -4023,21 +3551,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/47/bd/6f2fc8188f31bf10590f1e98e7b306336161fac930a8c514cd7bd828c7dc/pydantic_core-2.46.4-cp312-cp312-win32.whl", hash = "sha256:9aa768456404a8bf48a4406685ac2bec8e72b62c69313734fa3b73cf33b3a894", size = 1974823, upload-time = "2026-05-06T13:40:47.985Z" }, { url = "https://files.pythonhosted.org/packages/40/8c/985c1d41ea1107c2534abd9870e4ed5c8e7669b5c308297835c001e7a1c4/pydantic_core-2.46.4-cp312-cp312-win_amd64.whl", hash = "sha256:e9c26f834c65f5752f3f06cb08cb86a913ceb7274d0db6e267808a708b46bc89", size = 2072919, upload-time = "2026-05-06T13:39:21.153Z" }, { url = "https://files.pythonhosted.org/packages/c4/ba/f463d006e0c47373ca7ec5e1a261c59dc01ef4d62b2657af925fb0deee3a/pydantic_core-2.46.4-cp312-cp312-win_arm64.whl", hash = "sha256:4fc73cb559bdb54b1134a706a2802a4cddd27a0633f5abb7e53056268751ac6a", size = 2027604, upload-time = "2026-05-06T13:39:03.753Z" }, - { url = "https://files.pythonhosted.org/packages/51/a2/5d30b469c5267a17b39dec53208222f76a8d351dfac4af661888c5aee77d/pydantic_core-2.46.4-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:5d5902252db0d3cedf8d4a1bc68f70eeb430f7e4c7104c8c476753519b423008", size = 2106306, upload-time = "2026-05-06T13:37:48.029Z" }, - { url = "https://files.pythonhosted.org/packages/c1/81/4fa520eaffa8bd7d1525e644cd6d39e7d60b1592bc5b516693c7340b50f1/pydantic_core-2.46.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:c94f0688e7b8d0a67abf40e57a7eaaecd17cc9586706a31b76c031f63df052b4", size = 1951906, upload-time = "2026-05-06T13:37:17.012Z" }, - { url = "https://files.pythonhosted.org/packages/03/d5/fd02da45b659668b05923b17ba3a0100a0a3d5541e3bd8fcc4ecb711309e/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f027324c56cd5406ca49c124b0db10e56c69064fec039acc571c29020cc87c76", size = 1976802, upload-time = "2026-05-06T13:37:35.113Z" }, - { url = "https://files.pythonhosted.org/packages/21/f2/95727e1368be3d3ed485eaab7adbd7dda408f33f7a36e8b48e0144002b91/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:e739fee756ba1010f8bcccb534252e85a35fe45ae92c295a06059ce58b74ccd3", size = 2052446, upload-time = "2026-05-06T13:37:12.313Z" }, - { url = "https://files.pythonhosted.org/packages/9c/86/5d99feea3f77c7234b8718075b23db11532773c1a0dbd9b9490215dc2eeb/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9d56801be94b86a9da183e5f3766e6310752b99ff647e38b09a9500d88e46e76", size = 2232757, upload-time = "2026-05-06T13:39:01.149Z" }, - { url = "https://files.pythonhosted.org/packages/d2/3a/508ac615935ef7588cf6d9e9b91309fdc2da751af865e02a9098de88258c/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2412e734dcb48da14d4e4006b82b46b74f2518b8a26ee7e58c6844a6cd6d03c4", size = 2309275, upload-time = "2026-05-06T13:37:41.406Z" }, - { url = "https://files.pythonhosted.org/packages/07/f8/41db9de19d7987d6b04715a02b3b40aea467000275d9d758ffaa31af7d50/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9551187363ffc0de2a00b2e47c25aeaeb1020b69b668762966df15fc5659dd5a", size = 2094467, upload-time = "2026-05-06T13:39:18.847Z" }, - { url = "https://files.pythonhosted.org/packages/2c/e2/f35033184cb11d0052daf4416e8e10a502ea2ac006fc4f459aee872727d1/pydantic_core-2.46.4-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:0186750b482eefa11d7f435892b09c5c606193ef3375bcf94aa00ae6bfb66262", size = 2134417, upload-time = "2026-05-06T13:40:17.944Z" }, - { url = "https://files.pythonhosted.org/packages/7e/7b/6ceeb1cc90e193862f444ebe373d8fdf613f0a82572dde03fb10734c6c71/pydantic_core-2.46.4-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:5855698a4856556d86e8e6cd8434bc3ac0314ee8e12089ae0e143f64c6256e4e", size = 2179782, upload-time = "2026-05-06T13:40:32.618Z" }, - { url = "https://files.pythonhosted.org/packages/5a/f2/c8d7773ede6af08036423a00ae0ceffce266c3c52a096c435d68c896083f/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:cbaf13819775b7f769bf4a1f066cb6df7a28d4480081a589828ef190226881cd", size = 2188782, upload-time = "2026-05-06T13:36:51.018Z" }, - { url = "https://files.pythonhosted.org/packages/59/31/0c864784e31f09f05cdd87606f08923b9c9e7f6e51dd27f20f62f975ce9f/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:633147d34cf4550417f12e2b1a0383973bdf5cdfde212cb09e9a581cf10820be", size = 2328334, upload-time = "2026-05-06T13:40:37.764Z" }, - { url = "https://files.pythonhosted.org/packages/c2/eb/4f6c8a41efa30baa755590f4141abf3a8c370fab610915733e74134a7270/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:82cf5301172168103724d49a1444d3378cb20cdee30b116a1bd6031236298a5d", size = 2372986, upload-time = "2026-05-06T13:39:34.152Z" }, - { url = "https://files.pythonhosted.org/packages/5b/24/b375a480d53113860c299764bfe9f349a3dc9108b3adc0d7f0d786492ebf/pydantic_core-2.46.4-cp313-cp313-win32.whl", hash = "sha256:9fa8ae11da9e2b3126c6426f147e0fba88d96d65921799bb30c6abd1cb2c97fb", size = 1973693, upload-time = "2026-05-06T13:37:55.072Z" }, - { url = "https://files.pythonhosted.org/packages/7e/e8/cff247591966f2d22ec8c003cd7587e27b7ba7b81ab2fb888e3ab75dc285/pydantic_core-2.46.4-cp313-cp313-win_amd64.whl", hash = "sha256:6b3ace8194b0e5204818c92802dcdca7fc6d88aabbb799d7c795540d9cd6d292", size = 2071819, upload-time = "2026-05-06T13:38:49.139Z" }, - { url = "https://files.pythonhosted.org/packages/c6/1a/f4aee670d5670e9e148e0c82c7db98d780be566c6e6a97ee8035528ca0b3/pydantic_core-2.46.4-cp313-cp313-win_arm64.whl", hash = "sha256:184c081504d17f1c1066e430e117142b2c77d9448a97f7b65c6ac9fd9aee238d", size = 2027411, upload-time = "2026-05-06T13:40:45.796Z" }, { url = "https://files.pythonhosted.org/packages/ee/a4/73995fd4ebbb46ba0ee51e6fa049b8f02c40daebb762208feda8a6b7894d/pydantic_core-2.46.4-graalpy311-graalpy242_311_native-macosx_10_12_x86_64.whl", hash = "sha256:14d4edf427bdcf950a8a02d7cb44a08614388dd6e1bdcbf4f67504fa7887da9c", size = 2111589, upload-time = "2026-05-06T13:37:10.817Z" }, { url = "https://files.pythonhosted.org/packages/fb/7f/f37d3a5e8bfcc2e403f5c57a730f2d815693fb42119e8ea48b3789335af1/pydantic_core-2.46.4-graalpy311-graalpy242_311_native-macosx_11_0_arm64.whl", hash = "sha256:0ce40cd7b21210e99342afafbd4d0f76d784eb5b1d60f3bdc566be4983c6c73b", size = 1944552, upload-time = "2026-05-06T13:36:56.717Z" }, { url = "https://files.pythonhosted.org/packages/15/3c/d7eb777b3ff43e8433a4efb39a17aa8fd98a4ee8561a24a67ef5db07b2d6/pydantic_core-2.46.4-graalpy311-graalpy242_311_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:90884113d8b48f760e9587002789ddd741e76ab9f89518cd1e43b1f1a52ec44b", size = 1982984, upload-time = "2026-05-06T13:39:06.207Z" }, @@ -4142,7 +3655,7 @@ version = "26.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cryptography" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/74/b7/da07bae88f5a9506b4def6f2f4903cf4c3b8831e560dba8fa18ca08f758f/pyopenssl-26.3.0.tar.gz", hash = "sha256:589de7fae1c9ea670d18422ed00fc04da787bbde8e1454aea872aa57b49ad341", size = 182024, upload-time = "2026-06-12T20:28:07.458Z" } wheels = [ @@ -4183,7 +3696,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "backports-asyncio-runner", marker = "python_full_version < '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, { name = "pytest" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/43/7c/d36d04db312ecf4298932ef77e6e4a9e8ad017906e24e34f0b0c361a2473/pytest_asyncio-1.4.0.tar.gz", hash = "sha256:c6c0d2259945122819f171a32ecea2c349ead889ee28176caaf492143424be42", size = 58514, upload-time = "2026-05-26T09:56:04.083Z" } wheels = [ @@ -4280,16 +3793,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/74/93/7baea19427dcfbe1e5a372d81473250b379f04b1bd3c4c5ff825e2327202/pyyaml-6.0.3-cp312-cp312-win32.whl", hash = "sha256:96b533f0e99f6579b3d4d4995707cf36df9100d67e0c8303a0c55b27b5f99bc5", size = 137658, upload-time = "2025-09-25T21:32:20.209Z" }, { url = "https://files.pythonhosted.org/packages/86/bf/899e81e4cce32febab4fb42bb97dcdf66bc135272882d1987881a4b519e9/pyyaml-6.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:5fcd34e47f6e0b794d17de1b4ff496c00986e1c83f7ab2fb8fcfe9616ff7477b", size = 154003, upload-time = "2025-09-25T21:32:21.167Z" }, { url = "https://files.pythonhosted.org/packages/1a/08/67bd04656199bbb51dbed1439b7f27601dfb576fb864099c7ef0c3e55531/pyyaml-6.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:64386e5e707d03a7e172c0701abfb7e10f0fb753ee1d773128192742712a98fd", size = 140344, upload-time = "2025-09-25T21:32:22.617Z" }, - { url = "https://files.pythonhosted.org/packages/d1/11/0fd08f8192109f7169db964b5707a2f1e8b745d4e239b784a5a1dd80d1db/pyyaml-6.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:8da9669d359f02c0b91ccc01cac4a67f16afec0dac22c2ad09f46bee0697eba8", size = 181669, upload-time = "2025-09-25T21:32:23.673Z" }, - { url = "https://files.pythonhosted.org/packages/b1/16/95309993f1d3748cd644e02e38b75d50cbc0d9561d21f390a76242ce073f/pyyaml-6.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2283a07e2c21a2aa78d9c4442724ec1eb15f5e42a723b99cb3d822d48f5f7ad1", size = 173252, upload-time = "2025-09-25T21:32:25.149Z" }, - { url = "https://files.pythonhosted.org/packages/50/31/b20f376d3f810b9b2371e72ef5adb33879b25edb7a6d072cb7ca0c486398/pyyaml-6.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ee2922902c45ae8ccada2c5b501ab86c36525b883eff4255313a253a3160861c", size = 767081, upload-time = "2025-09-25T21:32:26.575Z" }, - { url = "https://files.pythonhosted.org/packages/49/1e/a55ca81e949270d5d4432fbbd19dfea5321eda7c41a849d443dc92fd1ff7/pyyaml-6.0.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a33284e20b78bd4a18c8c2282d549d10bc8408a2a7ff57653c0cf0b9be0afce5", size = 841159, upload-time = "2025-09-25T21:32:27.727Z" }, - { url = "https://files.pythonhosted.org/packages/74/27/e5b8f34d02d9995b80abcef563ea1f8b56d20134d8f4e5e81733b1feceb2/pyyaml-6.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f29edc409a6392443abf94b9cf89ce99889a1dd5376d94316ae5145dfedd5d6", size = 801626, upload-time = "2025-09-25T21:32:28.878Z" }, - { url = "https://files.pythonhosted.org/packages/f9/11/ba845c23988798f40e52ba45f34849aa8a1f2d4af4b798588010792ebad6/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f7057c9a337546edc7973c0d3ba84ddcdf0daa14533c2065749c9075001090e6", size = 753613, upload-time = "2025-09-25T21:32:30.178Z" }, - { url = "https://files.pythonhosted.org/packages/3d/e0/7966e1a7bfc0a45bf0a7fb6b98ea03fc9b8d84fa7f2229e9659680b69ee3/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:eda16858a3cab07b80edaf74336ece1f986ba330fdb8ee0d6c0d68fe82bc96be", size = 794115, upload-time = "2025-09-25T21:32:31.353Z" }, - { url = "https://files.pythonhosted.org/packages/de/94/980b50a6531b3019e45ddeada0626d45fa85cbe22300844a7983285bed3b/pyyaml-6.0.3-cp313-cp313-win32.whl", hash = "sha256:d0eae10f8159e8fdad514efdc92d74fd8d682c933a6dd088030f3834bc8e6b26", size = 137427, upload-time = "2025-09-25T21:32:32.58Z" }, - { url = "https://files.pythonhosted.org/packages/97/c9/39d5b874e8b28845e4ec2202b5da735d0199dbe5b8fb85f91398814a9a46/pyyaml-6.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:79005a0d97d5ddabfeeea4cf676af11e647e41d81c9a7722a193022accdb6b7c", size = 154090, upload-time = "2025-09-25T21:32:33.659Z" }, - { url = "https://files.pythonhosted.org/packages/73/e8/2bdf3ca2090f68bb3d75b44da7bbc71843b19c9f2b9cb9b0f4ab7a5a4329/pyyaml-6.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:5498cd1645aa724a7c71c8f378eb29ebe23da2fc0d7a08071d89469bf1d2defb", size = 140246, upload-time = "2025-09-25T21:32:34.663Z" }, ] [[package]] @@ -4360,38 +3863,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c3/1c/bdcc98f9a4af4fdd166c74941174619ccff4726d3ce32faa8e9a2ecd38dd/regex-2026.5.9-cp312-cp312-win32.whl", hash = "sha256:164eba9b755ea6f244b0d881196fbc1fac09714e9782c9e2732b813142033c8e", size = 266699, upload-time = "2026-05-09T23:12:59.14Z" }, { url = "https://files.pythonhosted.org/packages/78/87/240d36864f9e48ace85f72e79ced97ceb7f27ce87739a947dcb834b4e6bc/regex-2026.5.9-cp312-cp312-win_amd64.whl", hash = "sha256:86f40a5d6444db30a125c9c9177e6b25dad981cbc37451fd838f145e6edac92e", size = 277783, upload-time = "2026-05-09T23:13:00.789Z" }, { url = "https://files.pythonhosted.org/packages/4f/b5/7b30f312b0669dff5beebe5b0989dc2d1a312b1a44fab852199c387a5b96/regex-2026.5.9-cp312-cp312-win_arm64.whl", hash = "sha256:96f5f58b54a063d7ea9dca08e1cf57bfe10499c4d579ee672da284f57f5f0070", size = 270513, upload-time = "2026-05-09T23:13:02.426Z" }, - { url = "https://files.pythonhosted.org/packages/aa/da/797e91ecec6f84135da778ddce78c20e0af5d2a15c26f87a81bc3eadb6db/regex-2026.5.9-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:d626b84406444b165fc0ba981604edea39f0588ff1f92baa23fe50799ea9afdb", size = 490303, upload-time = "2026-05-09T23:13:04.382Z" }, - { url = "https://files.pythonhosted.org/packages/44/da/bf30abaaa737b58f4a4b8c4a03659e02fd92092c822e0197ed9e0daab917/regex-2026.5.9-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:d7bdc0ab8f3dd7e1b4f9ab88634e13374669db86bb3c72e8292f07ae313f539f", size = 292019, upload-time = "2026-05-09T23:13:06.022Z" }, - { url = "https://files.pythonhosted.org/packages/2d/e7/d0eaf5713828417b9e5648cf81fa9bacd4961f6ab98c380c2034f8716e35/regex-2026.5.9-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a8820737949116ffff55fe18f9fc644530063ba6ebfcb8314239416e78f1347c", size = 289468, upload-time = "2026-05-09T23:13:08.214Z" }, - { url = "https://files.pythonhosted.org/packages/d3/9b/b3fdd62b003baa1a9b593cd8c8699c9651c2e80cc21a5c715707983c42d7/regex-2026.5.9-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aa0fbdbac82cb3e4450d0ccde7d7a35607f4cb2dd9fba4b8b69bfaf8c9fa6aed", size = 796749, upload-time = "2026-05-09T23:13:10.573Z" }, - { url = "https://files.pythonhosted.org/packages/d4/30/66ab84588765f5b4b271a9ca09ef7ce2b87caa95176ec3d2ad65d7bc4902/regex-2026.5.9-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:57e8915c7986aa33d25e4d3629cef711cd2863f2961b10409f0c04cb8b7d9020", size = 865445, upload-time = "2026-05-09T23:13:12.523Z" }, - { url = "https://files.pythonhosted.org/packages/1a/89/f05169e8588aac365f35ffc7f3bc3184f095ef4cfded7cfaa3c7fd5dbd89/regex-2026.5.9-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:508f56a89ba9cb26e4168cbc37dbd60a28d82430a9e18ad1d25fe0883c314ca2", size = 912322, upload-time = "2026-05-09T23:13:14.281Z" }, - { url = "https://files.pythonhosted.org/packages/30/e1/c93444052cf41581f3c884ab3fb5823daf0992f11cd4388d4275ca610558/regex-2026.5.9-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b6d189041f15691cfa2b6c4290448ec221244d225b3f5fe9e7771b34ffcdf6e2", size = 801269, upload-time = "2026-05-09T23:13:16.569Z" }, - { url = "https://files.pythonhosted.org/packages/50/fe/0cf96b882f540e62e8b9956599798203d599c44cf4c77917ca27400ff69b/regex-2026.5.9-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e82db382b44d0111b22601c509c89f64434816c9e0eef9d1989cda8cc6ff1c04", size = 777085, upload-time = "2026-05-09T23:13:18.675Z" }, - { url = "https://files.pythonhosted.org/packages/23/5c/d78d4924e7fc875557b9e9b768423925fdfaac5549d06da7810019a9bd26/regex-2026.5.9-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:2acfb48634f64996b57f90f39afa692ff362162722581921fe92239a59960f3c", size = 785153, upload-time = "2026-05-09T23:13:20.525Z" }, - { url = "https://files.pythonhosted.org/packages/bf/e0/5214774090e7b4524dcea3e3c4aa74141d43043f8beb49c1599db1c8b53a/regex-2026.5.9-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:d29eebfc9525db68cad3c97eedd7f754fa265aa5cd0cf4f863b2421e1b48fc9f", size = 860164, upload-time = "2026-05-09T23:13:22.263Z" }, - { url = "https://files.pythonhosted.org/packages/6e/e1/4a57a83350319b1271f0d7a249b8672513ed928b237a741631270de6caea/regex-2026.5.9-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:debb893095e944091c16e641a6e33c1b0f4cb61ab945ec5afbf53ce7068834d8", size = 765731, upload-time = "2026-05-09T23:13:24.277Z" }, - { url = "https://files.pythonhosted.org/packages/12/f4/499e74a20c156fc75836ee04a72a38d1a063978f600937f9760467beb1b0/regex-2026.5.9-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:d659eee77986549c9ea45b861c7567e44d6287c3dc9a4565478853f7b9fe2ff6", size = 852062, upload-time = "2026-05-09T23:13:26.125Z" }, - { url = "https://files.pythonhosted.org/packages/5b/92/7eebc0d0a01e78629695f342ba17e0deaff8fb45e79cc0d7b98287da6e3e/regex-2026.5.9-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:2efa205e6d98b24d1f3ab395c11aa15cdf10935bca283d0285e0499c284fba21", size = 789577, upload-time = "2026-05-09T23:13:27.814Z" }, - { url = "https://files.pythonhosted.org/packages/05/a4/018e71f7d2ad48c1ebe6d3ae0026f9b7cb4802fd15c7cc02fdf724355102/regex-2026.5.9-cp313-cp313-win32.whl", hash = "sha256:f3844f134e834076677dd369976e9f5068679fcb8e50102fdf6b7ac96a3ec127", size = 266691, upload-time = "2026-05-09T23:13:29.549Z" }, - { url = "https://files.pythonhosted.org/packages/e6/1d/861a93719fb9ee7dbfc3761b3797b7a3e112a5d42c6129459d2d741be9b5/regex-2026.5.9-cp313-cp313-win_amd64.whl", hash = "sha256:3527bb4942d2c14552155406cdedd906567456821848aed1cb4933a391bf5eca", size = 277747, upload-time = "2026-05-09T23:13:31.859Z" }, - { url = "https://files.pythonhosted.org/packages/d9/c6/0a2436ae4da1ba76e51cb98943c6838a9a721faa40ebe2dce07694ae34e3/regex-2026.5.9-cp313-cp313-win_arm64.whl", hash = "sha256:56a33f191f17d8c417f99945ebdc1e691d3af9605d86ec68c7e54a57e3e17af6", size = 270500, upload-time = "2026-05-09T23:13:33.525Z" }, - { url = "https://files.pythonhosted.org/packages/e8/e9/d21346f7b60ed58789371358ed66b09d00f832e1bd7c06e55d9da5679882/regex-2026.5.9-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:01f28d868834624c934b8d2e0aa1c8341337e37831f4a012f18a5afcba4cbaf3", size = 494172, upload-time = "2026-05-09T23:13:35.935Z" }, - { url = "https://files.pythonhosted.org/packages/c4/43/fd1177a2032037c681baecdb3422ee4e1424aec4e4f470ef47793d325274/regex-2026.5.9-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:48036f6374aaa79eb3b754ec29c61d1c6b1606749d705a13f8854fa2539671f6", size = 293952, upload-time = "2026-05-09T23:13:38.307Z" }, - { url = "https://files.pythonhosted.org/packages/f2/7d/9fbf919768368d3f8a4f6c692cf2aa61e482b2b81ec6a298ace4cbf02480/regex-2026.5.9-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:b96350aa424e79d4fd6b567b344dcbe2b2d6bfc48dfe7717587e1fa6d43da6ff", size = 292314, upload-time = "2026-05-09T23:13:40.353Z" }, - { url = "https://files.pythonhosted.org/packages/e2/6c/e41bfeecb589716843e7c4df09ba46ff2a42961457afece19059d85caeef/regex-2026.5.9-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8f3af7a4903c5c04a11a196a5aa75cdd7dd3f8508132f9fb3259d9f5908e3b88", size = 811681, upload-time = "2026-05-09T23:13:42.543Z" }, - { url = "https://files.pythonhosted.org/packages/87/83/a5c1c525fba0aa656e88ad0face0b1829788ef4c2fb6b26df58aa1151b84/regex-2026.5.9-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7e87577720152d2caae19fe2baaf1f8d5ca12091e9e229f03915c37d1e4b9178", size = 871135, upload-time = "2026-05-09T23:13:44.326Z" }, - { url = "https://files.pythonhosted.org/packages/18/d4/80882e799e440dd878b0979cbebf8fa4d54624a332c83037c7a701649e3f/regex-2026.5.9-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:c8b9b9d294cfea3cd19c718ade7cc93492b2c4991abd9a68d0b3477ae6d8e100", size = 917265, upload-time = "2026-05-09T23:13:47.295Z" }, - { url = "https://files.pythonhosted.org/packages/ae/ff/8db60211e2286e396aad7dc7725356c502bff0901ea05bd6cdc2e1a042b9/regex-2026.5.9-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:728d8bfd28a8845c8b6bc5dc7ce010453d206396786c0765c2740cb65f37791e", size = 816311, upload-time = "2026-05-09T23:13:49.885Z" }, - { url = "https://files.pythonhosted.org/packages/4c/47/742ef579c61730f8d268e5cf1f9ce0e37e2ea041ad0f5644724f2378e463/regex-2026.5.9-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7e30b874d341fac767d7df5a0870540541c2c054b80cfaac116e8d367a8a7ff2", size = 785498, upload-time = "2026-05-09T23:13:52.25Z" }, - { url = "https://files.pythonhosted.org/packages/7f/ab/cb0999802dcb0fb95b1ab005e8d4163d8afdd67efc2cb6b6630ac13f8cb1/regex-2026.5.9-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:fd190e88a895a8901325fad284a3f74ea52b1da8525b76cc811fa9b1edf0ce2b", size = 801348, upload-time = "2026-05-09T23:13:54.127Z" }, - { url = "https://files.pythonhosted.org/packages/7d/62/8ca59a24c55bc34d166eefaf3717bd77772f329fdbf984d86581e0a3571c/regex-2026.5.9-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:8e76e8161ad00694cfce6767d5dea860c6391ac5b83e5c3a39661e696f11fc7e", size = 866493, upload-time = "2026-05-09T23:13:56.067Z" }, - { url = "https://files.pythonhosted.org/packages/8d/3d/30f2ae62cef3278bb5bb821f467277a55fb73f01032cf85997e15e8289a8/regex-2026.5.9-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:ddda5340e6c01a293027dd46232fa79eaff1b48058ce7a98f572b6445b088041", size = 772811, upload-time = "2026-05-09T23:13:57.867Z" }, - { url = "https://files.pythonhosted.org/packages/d8/ae/7d2089bcd78ad0c0161bc684339df50032acb438a7bd3305e7ddb1193cec/regex-2026.5.9-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:205109e96b3cf5adf8f4cd62bedde9487feb282b9497a3535451e5a24cd706a0", size = 856584, upload-time = "2026-05-09T23:13:59.679Z" }, - { url = "https://files.pythonhosted.org/packages/a9/29/92ff47f75990131ea4f24ba17819e5a9d141e10819807e09addd73409af6/regex-2026.5.9-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:dfbe4579b9f08036aa7d101d1835437a20783574ac66327e6b29b4018a138081", size = 803453, upload-time = "2026-05-09T23:14:01.978Z" }, - { url = "https://files.pythonhosted.org/packages/04/99/eff29f1037dcab36702c9ee5d6858cf1ce2336ea8ea2987f64245b99ea5e/regex-2026.5.9-cp313-cp313t-win32.whl", hash = "sha256:ed2c9e8068b614c574d8d30e543d617cf5379b0535d46f97ef00e904745a08b5", size = 269951, upload-time = "2026-05-09T23:14:03.661Z" }, - { url = "https://files.pythonhosted.org/packages/0e/9d/8870b8981d27b22cda77bb26a5ac7ebfa9c7d9e0dea195a834a82380e748/regex-2026.5.9-cp313-cp313t-win_amd64.whl", hash = "sha256:b46b0f094dc1d3b90356c85a0bd2c9bafc4a6a190b9d6f8ddd5a033b6e088ed4", size = 281240, upload-time = "2026-05-09T23:14:05.56Z" }, - { url = "https://files.pythonhosted.org/packages/72/b1/3379415e8f135c13ac551353397cc4fe97b4978f3cac73c5fcbcded548b8/regex-2026.5.9-cp313-cp313t-win_arm64.whl", hash = "sha256:872acc074bd29ffc9913ecdfedf6ea77502312ca44a4aa0d3779089c6069d8de", size = 272383, upload-time = "2026-05-09T23:14:07.843Z" }, ] [[package]] @@ -4508,24 +3979,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/89/b1/fbb53137f42c4bf630b1ffdfc2151a62d1d1b903b249f030d2b1c0280af8/scipy-1.15.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:6cfd56fc1a8e53f6e89ba3a7a7251f7396412d655bca2aa5611c8ec9a6784a1e", size = 36885140, upload-time = "2025-05-08T16:06:39.249Z" }, { url = "https://files.pythonhosted.org/packages/2e/2e/025e39e339f5090df1ff266d021892694dbb7e63568edcfe43f892fa381d/scipy-1.15.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:0ff17c0bb1cb32952c09217d8d1eed9b53d1463e5f1dd6052c7857f83127d539", size = 39710549, upload-time = "2025-05-08T16:06:45.729Z" }, { url = "https://files.pythonhosted.org/packages/e6/eb/3bf6ea8ab7f1503dca3a10df2e4b9c3f6b3316df07f6c0ded94b281c7101/scipy-1.15.3-cp312-cp312-win_amd64.whl", hash = "sha256:52092bc0472cfd17df49ff17e70624345efece4e1a12b23783a1ac59a1b728ed", size = 40966184, upload-time = "2025-05-08T16:06:52.623Z" }, - { url = "https://files.pythonhosted.org/packages/73/18/ec27848c9baae6e0d6573eda6e01a602e5649ee72c27c3a8aad673ebecfd/scipy-1.15.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:2c620736bcc334782e24d173c0fdbb7590a0a436d2fdf39310a8902505008759", size = 38728256, upload-time = "2025-05-08T16:06:58.696Z" }, - { url = "https://files.pythonhosted.org/packages/74/cd/1aef2184948728b4b6e21267d53b3339762c285a46a274ebb7863c9e4742/scipy-1.15.3-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:7e11270a000969409d37ed399585ee530b9ef6aa99d50c019de4cb01e8e54e62", size = 30109540, upload-time = "2025-05-08T16:07:04.209Z" }, - { url = "https://files.pythonhosted.org/packages/5b/d8/59e452c0a255ec352bd0a833537a3bc1bfb679944c4938ab375b0a6b3a3e/scipy-1.15.3-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:8c9ed3ba2c8a2ce098163a9bdb26f891746d02136995df25227a20e71c396ebb", size = 22383115, upload-time = "2025-05-08T16:07:08.998Z" }, - { url = "https://files.pythonhosted.org/packages/08/f5/456f56bbbfccf696263b47095291040655e3cbaf05d063bdc7c7517f32ac/scipy-1.15.3-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:0bdd905264c0c9cfa74a4772cdb2070171790381a5c4d312c973382fc6eaf730", size = 25163884, upload-time = "2025-05-08T16:07:14.091Z" }, - { url = "https://files.pythonhosted.org/packages/a2/66/a9618b6a435a0f0c0b8a6d0a2efb32d4ec5a85f023c2b79d39512040355b/scipy-1.15.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:79167bba085c31f38603e11a267d862957cbb3ce018d8b38f79ac043bc92d825", size = 35174018, upload-time = "2025-05-08T16:07:19.427Z" }, - { url = "https://files.pythonhosted.org/packages/b5/09/c5b6734a50ad4882432b6bb7c02baf757f5b2f256041da5df242e2d7e6b6/scipy-1.15.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c9deabd6d547aee2c9a81dee6cc96c6d7e9a9b1953f74850c179f91fdc729cb7", size = 37269716, upload-time = "2025-05-08T16:07:25.712Z" }, - { url = "https://files.pythonhosted.org/packages/77/0a/eac00ff741f23bcabd352731ed9b8995a0a60ef57f5fd788d611d43d69a1/scipy-1.15.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:dde4fc32993071ac0c7dd2d82569e544f0bdaff66269cb475e0f369adad13f11", size = 36872342, upload-time = "2025-05-08T16:07:31.468Z" }, - { url = "https://files.pythonhosted.org/packages/fe/54/4379be86dd74b6ad81551689107360d9a3e18f24d20767a2d5b9253a3f0a/scipy-1.15.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f77f853d584e72e874d87357ad70f44b437331507d1c311457bed8ed2b956126", size = 39670869, upload-time = "2025-05-08T16:07:38.002Z" }, - { url = "https://files.pythonhosted.org/packages/87/2e/892ad2862ba54f084ffe8cc4a22667eaf9c2bcec6d2bff1d15713c6c0703/scipy-1.15.3-cp313-cp313-win_amd64.whl", hash = "sha256:b90ab29d0c37ec9bf55424c064312930ca5f4bde15ee8619ee44e69319aab163", size = 40988851, upload-time = "2025-05-08T16:08:33.671Z" }, - { url = "https://files.pythonhosted.org/packages/1b/e9/7a879c137f7e55b30d75d90ce3eb468197646bc7b443ac036ae3fe109055/scipy-1.15.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:3ac07623267feb3ae308487c260ac684b32ea35fd81e12845039952f558047b8", size = 38863011, upload-time = "2025-05-08T16:07:44.039Z" }, - { url = "https://files.pythonhosted.org/packages/51/d1/226a806bbd69f62ce5ef5f3ffadc35286e9fbc802f606a07eb83bf2359de/scipy-1.15.3-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:6487aa99c2a3d509a5227d9a5e889ff05830a06b2ce08ec30df6d79db5fcd5c5", size = 30266407, upload-time = "2025-05-08T16:07:49.891Z" }, - { url = "https://files.pythonhosted.org/packages/e5/9b/f32d1d6093ab9eeabbd839b0f7619c62e46cc4b7b6dbf05b6e615bbd4400/scipy-1.15.3-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:50f9e62461c95d933d5c5ef4a1f2ebf9a2b4e83b0db374cb3f1de104d935922e", size = 22540030, upload-time = "2025-05-08T16:07:54.121Z" }, - { url = "https://files.pythonhosted.org/packages/e7/29/c278f699b095c1a884f29fda126340fcc201461ee8bfea5c8bdb1c7c958b/scipy-1.15.3-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:14ed70039d182f411ffc74789a16df3835e05dc469b898233a245cdfd7f162cb", size = 25218709, upload-time = "2025-05-08T16:07:58.506Z" }, - { url = "https://files.pythonhosted.org/packages/24/18/9e5374b617aba742a990581373cd6b68a2945d65cc588482749ef2e64467/scipy-1.15.3-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0a769105537aa07a69468a0eefcd121be52006db61cdd8cac8a0e68980bbb723", size = 34809045, upload-time = "2025-05-08T16:08:03.929Z" }, - { url = "https://files.pythonhosted.org/packages/e1/fe/9c4361e7ba2927074360856db6135ef4904d505e9b3afbbcb073c4008328/scipy-1.15.3-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9db984639887e3dffb3928d118145ffe40eff2fa40cb241a306ec57c219ebbbb", size = 36703062, upload-time = "2025-05-08T16:08:09.558Z" }, - { url = "https://files.pythonhosted.org/packages/b7/8e/038ccfe29d272b30086b25a4960f757f97122cb2ec42e62b460d02fe98e9/scipy-1.15.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:40e54d5c7e7ebf1aa596c374c49fa3135f04648a0caabcb66c52884b943f02b4", size = 36393132, upload-time = "2025-05-08T16:08:15.34Z" }, - { url = "https://files.pythonhosted.org/packages/10/7e/5c12285452970be5bdbe8352c619250b97ebf7917d7a9a9e96b8a8140f17/scipy-1.15.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:5e721fed53187e71d0ccf382b6bf977644c533e506c4d33c3fb24de89f5c3ed5", size = 38979503, upload-time = "2025-05-08T16:08:21.513Z" }, - { url = "https://files.pythonhosted.org/packages/81/06/0a5e5349474e1cbc5757975b21bd4fad0e72ebf138c5592f191646154e06/scipy-1.15.3-cp313-cp313t-win_amd64.whl", hash = "sha256:76ad1fb5f8752eabf0fa02e4cc0336b4e8f021e2d5f061ed37d6d264db35e3ca", size = 40308097, upload-time = "2025-05-08T16:08:27.627Z" }, ] [[package]] @@ -4561,26 +4014,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/65/94/7698add8f276dbab7a9de9fb6b0e02fc13ee61d51c7c3f85ac28b65e1239/scipy-1.17.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f590cd684941912d10becc07325a3eeb77886fe981415660d9265c4c418d0bea", size = 37625856, upload-time = "2026-02-23T00:19:00.307Z" }, { url = "https://files.pythonhosted.org/packages/a2/84/dc08d77fbf3d87d3ee27f6a0c6dcce1de5829a64f2eae85a0ecc1f0daa73/scipy-1.17.1-cp312-cp312-win_amd64.whl", hash = "sha256:41b71f4a3a4cab9d366cd9065b288efc4d4f3c0b37a91a8e0947fb5bd7f31d87", size = 36549682, upload-time = "2026-02-23T00:19:07.67Z" }, { url = "https://files.pythonhosted.org/packages/bc/98/fe9ae9ffb3b54b62559f52dedaebe204b408db8109a8c66fdd04869e6424/scipy-1.17.1-cp312-cp312-win_arm64.whl", hash = "sha256:f4115102802df98b2b0db3cce5cb9b92572633a1197c77b7553e5203f284a5b3", size = 24547340, upload-time = "2026-02-23T00:19:12.024Z" }, - { url = "https://files.pythonhosted.org/packages/76/27/07ee1b57b65e92645f219b37148a7e7928b82e2b5dbeccecb4dff7c64f0b/scipy-1.17.1-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:5e3c5c011904115f88a39308379c17f91546f77c1667cea98739fe0fccea804c", size = 31590199, upload-time = "2026-02-23T00:19:17.192Z" }, - { url = "https://files.pythonhosted.org/packages/ec/ae/db19f8ab842e9b724bf5dbb7db29302a91f1e55bc4d04b1025d6d605a2c5/scipy-1.17.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:6fac755ca3d2c3edcb22f479fceaa241704111414831ddd3bc6056e18516892f", size = 28154001, upload-time = "2026-02-23T00:19:22.241Z" }, - { url = "https://files.pythonhosted.org/packages/5b/58/3ce96251560107b381cbd6e8413c483bbb1228a6b919fa8652b0d4090e7f/scipy-1.17.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:7ff200bf9d24f2e4d5dc6ee8c3ac64d739d3a89e2326ba68aaf6c4a2b838fd7d", size = 20325719, upload-time = "2026-02-23T00:19:26.329Z" }, - { url = "https://files.pythonhosted.org/packages/b2/83/15087d945e0e4d48ce2377498abf5ad171ae013232ae31d06f336e64c999/scipy-1.17.1-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:4b400bdc6f79fa02a4d86640310dde87a21fba0c979efff5248908c6f15fad1b", size = 22683595, upload-time = "2026-02-23T00:19:30.304Z" }, - { url = "https://files.pythonhosted.org/packages/b4/e0/e58fbde4a1a594c8be8114eb4aac1a55bcd6587047efc18a61eb1f5c0d30/scipy-1.17.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2b64ca7d4aee0102a97f3ba22124052b4bd2152522355073580bf4845e2550b6", size = 32896429, upload-time = "2026-02-23T00:19:35.536Z" }, - { url = "https://files.pythonhosted.org/packages/f5/5f/f17563f28ff03c7b6799c50d01d5d856a1d55f2676f537ca8d28c7f627cd/scipy-1.17.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:581b2264fc0aa555f3f435a5944da7504ea3a065d7029ad60e7c3d1ae09c5464", size = 35203952, upload-time = "2026-02-23T00:19:42.259Z" }, - { url = "https://files.pythonhosted.org/packages/8d/a5/9afd17de24f657fdfe4df9a3f1ea049b39aef7c06000c13db1530d81ccca/scipy-1.17.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:beeda3d4ae615106d7094f7e7cef6218392e4465cc95d25f900bebabfded0950", size = 34979063, upload-time = "2026-02-23T00:19:47.547Z" }, - { url = "https://files.pythonhosted.org/packages/8b/13/88b1d2384b424bf7c924f2038c1c409f8d88bb2a8d49d097861dd64a57b2/scipy-1.17.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6609bc224e9568f65064cfa72edc0f24ee6655b47575954ec6339534b2798369", size = 37598449, upload-time = "2026-02-23T00:19:53.238Z" }, - { url = "https://files.pythonhosted.org/packages/35/e5/d6d0e51fc888f692a35134336866341c08655d92614f492c6860dc45bb2c/scipy-1.17.1-cp313-cp313-win_amd64.whl", hash = "sha256:37425bc9175607b0268f493d79a292c39f9d001a357bebb6b88fdfaff13f6448", size = 36510943, upload-time = "2026-02-23T00:20:50.89Z" }, - { url = "https://files.pythonhosted.org/packages/2a/fd/3be73c564e2a01e690e19cc618811540ba5354c67c8680dce3281123fb79/scipy-1.17.1-cp313-cp313-win_arm64.whl", hash = "sha256:5cf36e801231b6a2059bf354720274b7558746f3b1a4efb43fcf557ccd484a87", size = 24545621, upload-time = "2026-02-23T00:20:55.871Z" }, - { url = "https://files.pythonhosted.org/packages/6f/6b/17787db8b8114933a66f9dcc479a8272e4b4da75fe03b0c282f7b0ade8cd/scipy-1.17.1-cp313-cp313t-macosx_10_14_x86_64.whl", hash = "sha256:d59c30000a16d8edc7e64152e30220bfbd724c9bbb08368c054e24c651314f0a", size = 31936708, upload-time = "2026-02-23T00:19:58.694Z" }, - { url = "https://files.pythonhosted.org/packages/38/2e/524405c2b6392765ab1e2b722a41d5da33dc5c7b7278184a8ad29b6cb206/scipy-1.17.1-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:010f4333c96c9bb1a4516269e33cb5917b08ef2166d5556ca2fd9f082a9e6ea0", size = 28570135, upload-time = "2026-02-23T00:20:03.934Z" }, - { url = "https://files.pythonhosted.org/packages/fd/c3/5bd7199f4ea8556c0c8e39f04ccb014ac37d1468e6cfa6a95c6b3562b76e/scipy-1.17.1-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:2ceb2d3e01c5f1d83c4189737a42d9cb2fc38a6eeed225e7515eef71ad301dce", size = 20741977, upload-time = "2026-02-23T00:20:07.935Z" }, - { url = "https://files.pythonhosted.org/packages/d9/b8/8ccd9b766ad14c78386599708eb745f6b44f08400a5fd0ade7cf89b6fc93/scipy-1.17.1-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:844e165636711ef41f80b4103ed234181646b98a53c8f05da12ca5ca289134f6", size = 23029601, upload-time = "2026-02-23T00:20:12.161Z" }, - { url = "https://files.pythonhosted.org/packages/6d/a0/3cb6f4d2fb3e17428ad2880333cac878909ad1a89f678527b5328b93c1d4/scipy-1.17.1-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:158dd96d2207e21c966063e1635b1063cd7787b627b6f07305315dd73d9c679e", size = 33019667, upload-time = "2026-02-23T00:20:17.208Z" }, - { url = "https://files.pythonhosted.org/packages/f3/c3/2d834a5ac7bf3a0c806ad1508efc02dda3c8c61472a56132d7894c312dea/scipy-1.17.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74cbb80d93260fe2ffa334efa24cb8f2f0f622a9b9febf8b483c0b865bfb3475", size = 35264159, upload-time = "2026-02-23T00:20:23.087Z" }, - { url = "https://files.pythonhosted.org/packages/4d/77/d3ed4becfdbd217c52062fafe35a72388d1bd82c2d0ba5ca19d6fcc93e11/scipy-1.17.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:dbc12c9f3d185f5c737d801da555fb74b3dcfa1a50b66a1a93e09190f41fab50", size = 35102771, upload-time = "2026-02-23T00:20:28.636Z" }, - { url = "https://files.pythonhosted.org/packages/bd/12/d19da97efde68ca1ee5538bb261d5d2c062f0c055575128f11a2730e3ac1/scipy-1.17.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:94055a11dfebe37c656e70317e1996dc197e1a15bbcc351bcdd4610e128fe1ca", size = 37665910, upload-time = "2026-02-23T00:20:34.743Z" }, - { url = "https://files.pythonhosted.org/packages/06/1c/1172a88d507a4baaf72c5a09bb6c018fe2ae0ab622e5830b703a46cc9e44/scipy-1.17.1-cp313-cp313t-win_amd64.whl", hash = "sha256:e30bdeaa5deed6bc27b4cc490823cd0347d7dae09119b8803ae576ea0ce52e4c", size = 36562980, upload-time = "2026-02-23T00:20:40.575Z" }, - { url = "https://files.pythonhosted.org/packages/70/b0/eb757336e5a76dfa7911f63252e3b7d1de00935d7705cf772db5b45ec238/scipy-1.17.1-cp313-cp313t-win_arm64.whl", hash = "sha256:a720477885a9d2411f94a93d16f9d89bad0f28ca23c3f8daa521e2dcc3f44d49", size = 24856543, upload-time = "2026-02-23T00:20:45.313Z" }, ] [[package]] @@ -4606,16 +4039,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/84/ca/210d4759c7210bb7d269437421959b39a33434e2776b60c5cb8a763bb30a/scipy-1.18.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:4a55985d54c769c872e64b7f4c8a81cc30ef700cc04296abbbf3705439c126de", size = 37421717, upload-time = "2026-06-19T15:00:05.102Z" }, { url = "https://files.pythonhosted.org/packages/2b/54/9a9edb45345bd6744da5ddfb6628e5d5185920494c6a67ec45b6381004cb/scipy-1.18.0-cp312-cp312-win_amd64.whl", hash = "sha256:71ccc8faa2dd16ac310233203474a8b5cb67f10dedd54a3116d34943f4b19132", size = 36597428, upload-time = "2026-06-19T15:00:08.112Z" }, { url = "https://files.pythonhosted.org/packages/99/0e/33f32a2a58987e26aec0f7df252cbbad1e90ae77bdbc76f40dd4ed0cf0ea/scipy-1.18.0-cp312-cp312-win_arm64.whl", hash = "sha256:d88363fd9d8fbd3511bd273f1a49efb2a540773ddf92a91d57498ce7dd7f3e76", size = 24351481, upload-time = "2026-06-19T15:00:11.103Z" }, - { url = "https://files.pythonhosted.org/packages/05/52/9c0136c2de7ae0779b7b366447766cec6d9f0702c56bb8ffeb04c8fd3af4/scipy-1.18.0-cp313-cp313-macosx_10_15_x86_64.whl", hash = "sha256:09143f676d157d9f546d663504ef9c1becb819824f1afc018814176411942446", size = 31036107, upload-time = "2026-06-19T15:00:14.03Z" }, - { url = "https://files.pythonhosted.org/packages/02/73/0291a64843270f4efb86cdcf2ee0f2048631b65ec6b405398b2b4dbf11bf/scipy-1.18.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:5efe260f69417b97ddae455bfb5a95e8359f7f66ad7fa9522a60feb66f169520", size = 28663303, upload-time = "2026-06-19T15:00:16.819Z" }, - { url = "https://files.pythonhosted.org/packages/d3/0f/10ffa0b697a572f4e0d48b92a88895d366422f019f723e7e14a84c050dac/scipy-1.18.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:68363b7eaacd8b5dd426df56d782cc156468ac79a127a1b87ca597d6e2e82197", size = 20404960, upload-time = "2026-06-19T15:00:19.635Z" }, - { url = "https://files.pythonhosted.org/packages/7e/d2/e896cea21ba8edd6c81d4c55b1ffcc717e79698dcbebf9641b4cfb4c6622/scipy-1.18.0-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:c5557d8be5da8e41353fcd4d21491fdbab83b062fc579e94dc09a7c8ab4f669b", size = 23034074, upload-time = "2026-06-19T15:00:22.107Z" }, - { url = "https://files.pythonhosted.org/packages/ea/b2/e83ea34279a52c03374477c74006256ec78df65fc877baa4617d6de1d202/scipy-1.18.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0d13bca67c096d89fb95ced0d8921807300fce0275643aef9533cc63a0773468", size = 33942038, upload-time = "2026-06-19T15:00:24.964Z" }, - { url = "https://files.pythonhosted.org/packages/f6/af/e8fe5fb136f51e2b01678b92cb4106d10d8cd68ec147ead2e7cb0ac75398/scipy-1.18.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a46f9273dbd0eb1cefba61c9b8648b4dfe3cbc14a080176f9a73e44b8336dc7f", size = 35266390, upload-time = "2026-06-19T15:00:28.059Z" }, - { url = "https://files.pythonhosted.org/packages/3a/49/2c5cbb907b56695fc67517811d1db234dfd83381a84814ec220aded2794d/scipy-1.18.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:5aba46108853ddfc77906b6557aac839d2b52e900c1d72a1180adaaab58d265f", size = 35551324, upload-time = "2026-06-19T15:00:31.014Z" }, - { url = "https://files.pythonhosted.org/packages/bb/73/eda39f7a2d306ff0ffc574afd13c0bbb6d10a603d9a413998ee269487a80/scipy-1.18.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b6f758e35f12757b5d95c00bc6de2438e229c2664b7a92e96f205959d9f2dfa4", size = 37404785, upload-time = "2026-06-19T15:00:34.072Z" }, - { url = "https://files.pythonhosted.org/packages/b7/d2/ae881ee28d014f38e0ccbfd974a06a919ba9af34f1f74bf42b5301891d63/scipy-1.18.0-cp313-cp313-win_amd64.whl", hash = "sha256:1afac4a847207c7ff8efd321734a50b06d0280b3b2a2c0fc2f413101747ad7c7", size = 36554943, upload-time = "2026-06-19T15:00:36.903Z" }, - { url = "https://files.pythonhosted.org/packages/70/3a/21154e2d54eb3639c6bf4dbae2e531c68356bfe95990daa30df33b30d556/scipy-1.18.0-cp313-cp313-win_arm64.whl", hash = "sha256:c5dbddf60e58c2312316d097271a8e73d40eaf2eabfa4d95ed7d3695bbf2ce7b", size = 24350911, upload-time = "2026-06-19T15:00:40.062Z" }, ] [[package]] @@ -4645,20 +4068,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b6/2d/37e3da037318a70066ded0d51bc2a7f35491ae6338dd993d5eb1503fc3b5/sentencepiece-0.2.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c8a168b040bc61681293f79a949b5d911c8e25086f4260285b8d97ab5f1195da", size = 1397736, upload-time = "2026-07-12T08:38:35.771Z" }, { url = "https://files.pythonhosted.org/packages/8d/11/753fca2e6b109be3ab7867abf357dfe48677fe726ae5a5363d0b54ca9450/sentencepiece-0.2.2-cp312-cp312-win_amd64.whl", hash = "sha256:7c6e7bf684dc12145bfa685d3060beaea55139134ba848289bee514ed42e7383", size = 1248030, upload-time = "2026-07-12T08:38:37.604Z" }, { url = "https://files.pythonhosted.org/packages/e2/0a/70efbe861ca182d7d4b6e1a20f58e043400848fa9f2915229f082e221648/sentencepiece-0.2.2-cp312-cp312-win_arm64.whl", hash = "sha256:76ff5814db72e7462dece042d7593cdf102b8ec82c2b1cc201a2add34ee3050d", size = 1187325, upload-time = "2026-07-12T08:38:39.348Z" }, - { url = "https://files.pythonhosted.org/packages/b9/a3/b3b05095c174d6e80d37d5ddc2f57c2c56237333e7bbd6079cf3243c2a8a/sentencepiece-0.2.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:77c3ce990b23441e5ecfa5bce181fd6f408b564aeb6d7e1d1e7de9c5612501c8", size = 2188346, upload-time = "2026-07-12T08:38:41.089Z" }, - { url = "https://files.pythonhosted.org/packages/ca/f3/72ebc4acb10a06bcf7503fbc6091c8f5db68300f6aac4356c09e6c76e0e1/sentencepiece-0.2.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:fd523c4992041faa5c2b3cde62253d11a96c30d73a34afe48a486e8e2254cd1c", size = 1441434, upload-time = "2026-07-12T08:38:42.56Z" }, - { url = "https://files.pythonhosted.org/packages/34/db/f9ea1a6844b4fa5dfe2312095cd866a1f724cd0905054ab9d5991778ba50/sentencepiece-0.2.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:201a8e0f55501a76e08dbf2c54bc45f4642b379271e89c667d517bfbc2191f2a", size = 1347267, upload-time = "2026-07-12T08:38:44.389Z" }, - { url = "https://files.pythonhosted.org/packages/32/4f/31c1073314ad94466bca37d29581761d70110237ee3d46b0efece59a8c1e/sentencepiece-0.2.2-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8eed98514bffe5ecac37f493f91869c351fbb05629328bfdbc08502c6c094dc0", size = 1324980, upload-time = "2026-07-12T08:38:46.304Z" }, - { url = "https://files.pythonhosted.org/packages/59/b4/a0356fa04d6a14337a6e0e443556785a0422c53ec58baae6b9568120eb0f/sentencepiece-0.2.2-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:64b656f025355cf8c51abe9fbe3848540756c6d7ca5e6791b1afa664bc24c7cb", size = 1397593, upload-time = "2026-07-12T08:38:48.302Z" }, - { url = "https://files.pythonhosted.org/packages/09/fa/d2d6369257fd2f0de616b1c7110b73fab409ef61b14f1b9e0010ed325914/sentencepiece-0.2.2-cp313-cp313-win_amd64.whl", hash = "sha256:74f0ee601047c0c12a783088b51be4e6214a62ecd9e02278c477433cd16e0ed9", size = 1247987, upload-time = "2026-07-12T08:38:50.15Z" }, - { url = "https://files.pythonhosted.org/packages/17/ee/2bb594da6fd95e32f29057f1aa7fa996701b8980090923c2d8711fdc0a24/sentencepiece-0.2.2-cp313-cp313-win_arm64.whl", hash = "sha256:b23fe17779834d3c27aaf2edac9486d04cca1a7deb8f5facda35150ac6263a91", size = 1187250, upload-time = "2026-07-12T08:38:52.246Z" }, - { url = "https://files.pythonhosted.org/packages/58/9c/dfc82846460e7a712310f5613f23d8b553cabb4e2e648663c11d8382af56/sentencepiece-0.2.2-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:72b7825b331b1b7e7c45be2e674b3e3c65af608fa376bad2d851b20aaf0cdc78", size = 2223080, upload-time = "2026-07-12T08:38:54.391Z" }, - { url = "https://files.pythonhosted.org/packages/8d/4e/3ff12cebe6d31662d9ceeabfb282de20bd0d6098fa282b4a3b8305abc7e8/sentencepiece-0.2.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:d795c4ac689a57f9d4ba2288126ec7901d389ad5827d2f8b8533c883974fe563", size = 1458511, upload-time = "2026-07-12T08:38:56.811Z" }, - { url = "https://files.pythonhosted.org/packages/59/5a/16d51d05360be4cee3ebfe4837c184054c4eed16cabaeb3b039524e9a000/sentencepiece-0.2.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3ab3f1ae98970b5590e2209341522718900ba19bcc2c207ffaa6bd417ad960c5", size = 1361138, upload-time = "2026-07-12T08:38:58.808Z" }, - { url = "https://files.pythonhosted.org/packages/0f/af/c30ee2a9f99d51db9844acaa8fa0b611a97c2fa7116646fa43db3300b187/sentencepiece-0.2.2-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3ec27c152a1f1b24bc9168b55a5880f3c16e2334e697da6f55a1046a22405a3d", size = 1328625, upload-time = "2026-07-12T08:39:00.849Z" }, - { url = "https://files.pythonhosted.org/packages/3e/1a/4c6b39d03f5ba8439509adbd5a23c9538088a3cb679e7a47b911e8442bc6/sentencepiece-0.2.2-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:59d6588712101ccfcae9b03692be3aaae1514c2078666d7b05f15ba3a702e41b", size = 1398595, upload-time = "2026-07-12T08:39:02.86Z" }, - { url = "https://files.pythonhosted.org/packages/0f/bc/9eedddcec1fd57bc70200fa3ebf792d18fa63527a5369581cd416c81f97f/sentencepiece-0.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:89625fb43765cccaa1443b9adb61f283e5fe4cb1536728205d06bada730caa53", size = 1259346, upload-time = "2026-07-12T08:39:04.559Z" }, - { url = "https://files.pythonhosted.org/packages/41/15/7e74c8533848866ff560b29f7d8719921b76c4ec7149592d6d28e0deee75/sentencepiece-0.2.2-cp313-cp313t-win_arm64.whl", hash = "sha256:4f0603267cd15b92b68c2c0e852a441507614b70dc7773659baa6b8c214a91fd", size = 1196596, upload-time = "2026-07-12T08:39:06.454Z" }, ] [[package]] @@ -4704,22 +4113,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/a0/704c7292f7014c7e74ec84eddb7b109e1fbae74a16deae9c1504b1d15565/shapely-2.1.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:980c777c612514c0cf99bc8a9de6d286f5e186dcaf9091252fcd444e5638193d", size = 4152714, upload-time = "2025-09-24T13:50:39.9Z" }, { url = "https://files.pythonhosted.org/packages/53/46/319c9dc788884ad0785242543cdffac0e6530e4d0deb6c4862bc4143dcf3/shapely-2.1.2-cp312-cp312-win32.whl", hash = "sha256:9111274b88e4d7b54a95218e243282709b330ef52b7b86bc6aaf4f805306f454", size = 1542745, upload-time = "2025-09-24T13:50:41.414Z" }, { url = "https://files.pythonhosted.org/packages/ec/bf/cb6c1c505cb31e818e900b9312d514f381fbfa5c4363edfce0fcc4f8c1a4/shapely-2.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:743044b4cfb34f9a67205cee9279feaf60ba7d02e69febc2afc609047cb49179", size = 1722861, upload-time = "2025-09-24T13:50:43.35Z" }, - { url = "https://files.pythonhosted.org/packages/c3/90/98ef257c23c46425dc4d1d31005ad7c8d649fe423a38b917db02c30f1f5a/shapely-2.1.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:b510dda1a3672d6879beb319bc7c5fd302c6c354584690973c838f46ec3e0fa8", size = 1832644, upload-time = "2025-09-24T13:50:44.886Z" }, - { url = "https://files.pythonhosted.org/packages/6d/ab/0bee5a830d209adcd3a01f2d4b70e587cdd9fd7380d5198c064091005af8/shapely-2.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8cff473e81017594d20ec55d86b54bc635544897e13a7cfc12e36909c5309a2a", size = 1642887, upload-time = "2025-09-24T13:50:46.735Z" }, - { url = "https://files.pythonhosted.org/packages/2d/5e/7d7f54ba960c13302584c73704d8c4d15404a51024631adb60b126a4ae88/shapely-2.1.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fe7b77dc63d707c09726b7908f575fc04ff1d1ad0f3fb92aec212396bc6cfe5e", size = 2970931, upload-time = "2025-09-24T13:50:48.374Z" }, - { url = "https://files.pythonhosted.org/packages/f2/a2/83fc37e2a58090e3d2ff79175a95493c664bcd0b653dd75cb9134645a4e5/shapely-2.1.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7ed1a5bbfb386ee8332713bf7508bc24e32d24b74fc9a7b9f8529a55db9f4ee6", size = 3082855, upload-time = "2025-09-24T13:50:50.037Z" }, - { url = "https://files.pythonhosted.org/packages/44/2b/578faf235a5b09f16b5f02833c53822294d7f21b242f8e2d0cf03fb64321/shapely-2.1.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a84e0582858d841d54355246ddfcbd1fce3179f185da7470f41ce39d001ee1af", size = 3979960, upload-time = "2025-09-24T13:50:51.74Z" }, - { url = "https://files.pythonhosted.org/packages/4d/04/167f096386120f692cc4ca02f75a17b961858997a95e67a3cb6a7bbd6b53/shapely-2.1.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:dc3487447a43d42adcdf52d7ac73804f2312cbfa5d433a7d2c506dcab0033dfd", size = 4142851, upload-time = "2025-09-24T13:50:53.49Z" }, - { url = "https://files.pythonhosted.org/packages/48/74/fb402c5a6235d1c65a97348b48cdedb75fb19eca2b1d66d04969fc1c6091/shapely-2.1.2-cp313-cp313-win32.whl", hash = "sha256:9c3a3c648aedc9f99c09263b39f2d8252f199cb3ac154fadc173283d7d111350", size = 1541890, upload-time = "2025-09-24T13:50:55.337Z" }, - { url = "https://files.pythonhosted.org/packages/41/47/3647fe7ad990af60ad98b889657a976042c9988c2807cf322a9d6685f462/shapely-2.1.2-cp313-cp313-win_amd64.whl", hash = "sha256:ca2591bff6645c216695bdf1614fca9c82ea1144d4a7591a466fef64f28f0715", size = 1722151, upload-time = "2025-09-24T13:50:57.153Z" }, - { url = "https://files.pythonhosted.org/packages/3c/49/63953754faa51ffe7d8189bfbe9ca34def29f8c0e34c67cbe2a2795f269d/shapely-2.1.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2d93d23bdd2ed9dc157b46bc2f19b7da143ca8714464249bef6771c679d5ff40", size = 1834130, upload-time = "2025-09-24T13:50:58.49Z" }, - { url = "https://files.pythonhosted.org/packages/7f/ee/dce001c1984052970ff60eb4727164892fb2d08052c575042a47f5a9e88f/shapely-2.1.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:01d0d304b25634d60bd7cf291828119ab55a3bab87dc4af1e44b07fb225f188b", size = 1642802, upload-time = "2025-09-24T13:50:59.871Z" }, - { url = "https://files.pythonhosted.org/packages/da/e7/fc4e9a19929522877fa602f705706b96e78376afb7fad09cad5b9af1553c/shapely-2.1.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8d8382dd120d64b03698b7298b89611a6ea6f55ada9d39942838b79c9bc89801", size = 3018460, upload-time = "2025-09-24T13:51:02.08Z" }, - { url = "https://files.pythonhosted.org/packages/a1/18/7519a25db21847b525696883ddc8e6a0ecaa36159ea88e0fef11466384d0/shapely-2.1.2-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:19efa3611eef966e776183e338b2d7ea43569ae99ab34f8d17c2c054d3205cc0", size = 3095223, upload-time = "2025-09-24T13:51:04.472Z" }, - { url = "https://files.pythonhosted.org/packages/48/de/b59a620b1f3a129c3fecc2737104a0a7e04e79335bd3b0a1f1609744cf17/shapely-2.1.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:346ec0c1a0fcd32f57f00e4134d1200e14bf3f5ae12af87ba83ca275c502498c", size = 4030760, upload-time = "2025-09-24T13:51:06.455Z" }, - { url = "https://files.pythonhosted.org/packages/96/b3/c6655ee7232b417562bae192ae0d3ceaadb1cc0ffc2088a2ddf415456cc2/shapely-2.1.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:6305993a35989391bd3476ee538a5c9a845861462327efe00dd11a5c8c709a99", size = 4170078, upload-time = "2025-09-24T13:51:08.584Z" }, - { url = "https://files.pythonhosted.org/packages/a0/8e/605c76808d73503c9333af8f6cbe7e1354d2d238bda5f88eea36bfe0f42a/shapely-2.1.2-cp313-cp313t-win32.whl", hash = "sha256:c8876673449f3401f278c86eb33224c5764582f72b653a415d0e6672fde887bf", size = 1559178, upload-time = "2025-09-24T13:51:10.73Z" }, - { url = "https://files.pythonhosted.org/packages/36/f7/d317eb232352a1f1444d11002d477e54514a4a6045536d49d0c59783c0da/shapely-2.1.2-cp313-cp313t-win_amd64.whl", hash = "sha256:4a44bc62a10d84c11a7a3d7c1c4fe857f7477c3506e24c9062da0db0ae0c449c", size = 1739756, upload-time = "2025-09-24T13:51:12.105Z" }, ] [[package]] @@ -4762,10 +4155,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/20/bc/446357c229692f51885f4c5f3894af3aff37ccaafebc4f24066c2b9b5b80/slangpy-0.42.0-cp312-cp312-manylinux_2_34_aarch64.whl", hash = "sha256:4e2f76d3709ee2152f54620a27ab45fbb459db392c03b2d492382f675bfccc2a", size = 82123007, upload-time = "2026-05-28T22:40:30.815Z" }, { url = "https://files.pythonhosted.org/packages/cd/d1/d9822a2c38dc583850608650e6d3304f6cf6e03a1335ba21078440790a68/slangpy-0.42.0-cp312-cp312-manylinux_2_34_x86_64.whl", hash = "sha256:8dbc36b32c1c1fc8dffa70d63ddf0ab538c473ba072f802c8d7730530bea7a7c", size = 83670144, upload-time = "2026-05-28T22:40:37.618Z" }, { url = "https://files.pythonhosted.org/packages/b7/b8/94d067236898b5bec62a7ae682c296ba2310ebf9aa894bb32de9a4cfd478/slangpy-0.42.0-cp312-cp312-win_amd64.whl", hash = "sha256:82e212b7b195aeafb23ab43a43069980e3be0952f91869b65b5b4288bad9129d", size = 78280523, upload-time = "2026-05-28T22:40:44.154Z" }, - { url = "https://files.pythonhosted.org/packages/15/3c/44f8d8c20b83e10dc2dffc45d9a93adc3e6d037ec1b72756e8dfeec9cd7d/slangpy-0.42.0-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:c070359285acf177fb7f4ca6cbd29c6c7f37a9b036f42f436353aa773c08e665", size = 37738648, upload-time = "2026-05-28T22:40:48.454Z" }, - { url = "https://files.pythonhosted.org/packages/da/23/6427597cd186477c6020124883eb1f5d7e15ae4e46bb7ee60125fbc30979/slangpy-0.42.0-cp313-cp313-manylinux_2_34_aarch64.whl", hash = "sha256:07a8a5bb32e5644ea0226e40db5c773f1218c901883e70d2f7b20ceef433f724", size = 82125926, upload-time = "2026-05-28T22:40:54.468Z" }, - { url = "https://files.pythonhosted.org/packages/9f/b1/cbde95d04d94d2c3f0241b54354f5e792e64a1c1643b4424a2ccb80f2788/slangpy-0.42.0-cp313-cp313-manylinux_2_34_x86_64.whl", hash = "sha256:4c26022765a414925efde972e26ca3f99458435879c1133c6bbe2e962a6c2c16", size = 83670289, upload-time = "2026-05-28T22:41:00.771Z" }, - { url = "https://files.pythonhosted.org/packages/83/9e/e57c578d8a576ddfea74f2f11d561d2ceca6b4fd92fa1341dc3d93a6b7f8/slangpy-0.42.0-cp313-cp313-win_amd64.whl", hash = "sha256:dd0a32184bce8a8c507adbfedf5b49694f55543223c019a3c540ffe4a7df1730", size = 78280954, upload-time = "2026-05-28T22:41:06.903Z" }, ] [[package]] @@ -5050,7 +4439,7 @@ version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/eb/e3/7c1dc7381d9f8ab7d854328ebfa884e62cb3f3d8549ddfd37c7814f42afa/starlette-1.3.1.tar.gz", hash = "sha256:05d0213193f2fbaae60e2ecb593b4add4262ad4e46536b54abe36f11a71724e0", size = 2703240, upload-time = "2026-06-12T09:23:11.602Z" } wheels = [ @@ -5123,15 +4512,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d6/2f/4a3c322f22c5c66c4b836ec58211641a4067364f5dcdd7b974b4c5da300c/tomli-2.4.1-cp312-cp312-win32.whl", hash = "sha256:da25dc3563bff5965356133435b757a795a17b17d01dbc0f42fb32447ddfd917", size = 98141, upload-time = "2026-03-25T20:21:28.492Z" }, { url = "https://files.pythonhosted.org/packages/24/22/4daacd05391b92c55759d55eaee21e1dfaea86ce5c571f10083360adf534/tomli-2.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:52c8ef851d9a240f11a88c003eacb03c31fc1c9c4ec64a99a0f922b93874fda9", size = 108847, upload-time = "2026-03-25T20:21:29.386Z" }, { url = "https://files.pythonhosted.org/packages/68/fd/70e768887666ddd9e9f5d85129e84910f2db2796f9096aa02b721a53098d/tomli-2.4.1-cp312-cp312-win_arm64.whl", hash = "sha256:f758f1b9299d059cc3f6546ae2af89670cb1c4d48ea29c3cacc4fe7de3058257", size = 95088, upload-time = "2026-03-25T20:21:30.677Z" }, - { url = "https://files.pythonhosted.org/packages/07/06/b823a7e818c756d9a7123ba2cda7d07bc2dd32835648d1a7b7b7a05d848d/tomli-2.4.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:36d2bd2ad5fb9eaddba5226aa02c8ec3fa4f192631e347b3ed28186d43be6b54", size = 155866, upload-time = "2026-03-25T20:21:31.65Z" }, - { url = "https://files.pythonhosted.org/packages/14/6f/12645cf7f08e1a20c7eb8c297c6f11d31c1b50f316a7e7e1e1de6e2e7b7e/tomli-2.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:eb0dc4e38e6a1fd579e5d50369aa2e10acfc9cace504579b2faabb478e76941a", size = 149887, upload-time = "2026-03-25T20:21:33.028Z" }, - { url = "https://files.pythonhosted.org/packages/5c/e0/90637574e5e7212c09099c67ad349b04ec4d6020324539297b634a0192b0/tomli-2.4.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7f2c7f2b9ca6bdeef8f0fa897f8e05085923eb091721675170254cbc5b02897", size = 243704, upload-time = "2026-03-25T20:21:34.51Z" }, - { url = "https://files.pythonhosted.org/packages/10/8f/d3ddb16c5a4befdf31a23307f72828686ab2096f068eaf56631e136c1fdd/tomli-2.4.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f3c6818a1a86dd6dca7ddcaaf76947d5ba31aecc28cb1b67009a5877c9a64f3f", size = 251628, upload-time = "2026-03-25T20:21:36.012Z" }, - { url = "https://files.pythonhosted.org/packages/e3/f1/dbeeb9116715abee2485bf0a12d07a8f31af94d71608c171c45f64c0469d/tomli-2.4.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d312ef37c91508b0ab2cee7da26ec0b3ed2f03ce12bd87a588d771ae15dcf82d", size = 247180, upload-time = "2026-03-25T20:21:37.136Z" }, - { url = "https://files.pythonhosted.org/packages/d3/74/16336ffd19ed4da28a70959f92f506233bd7cfc2332b20bdb01591e8b1d1/tomli-2.4.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:51529d40e3ca50046d7606fa99ce3956a617f9b36380da3b7f0dd3dd28e68cb5", size = 251674, upload-time = "2026-03-25T20:21:38.298Z" }, - { url = "https://files.pythonhosted.org/packages/16/f9/229fa3434c590ddf6c0aa9af64d3af4b752540686cace29e6281e3458469/tomli-2.4.1-cp313-cp313-win32.whl", hash = "sha256:2190f2e9dd7508d2a90ded5ed369255980a1bcdd58e52f7fe24b8162bf9fedbd", size = 97976, upload-time = "2026-03-25T20:21:39.316Z" }, - { url = "https://files.pythonhosted.org/packages/6a/1e/71dfd96bcc1c775420cb8befe7a9d35f2e5b1309798f009dca17b7708c1e/tomli-2.4.1-cp313-cp313-win_amd64.whl", hash = "sha256:8d65a2fbf9d2f8352685bc1364177ee3923d6baf5e7f43ea4959d7d8bc326a36", size = 108755, upload-time = "2026-03-25T20:21:40.248Z" }, - { url = "https://files.pythonhosted.org/packages/83/7a/d34f422a021d62420b78f5c538e5b102f62bea616d1d75a13f0a88acb04a/tomli-2.4.1-cp313-cp313-win_arm64.whl", hash = "sha256:4b605484e43cdc43f0954ddae319fb75f04cc10dd80d830540060ee7cd0243cd", size = 95265, upload-time = "2026-03-25T20:21:41.219Z" }, { url = "https://files.pythonhosted.org/packages/7b/61/cceae43728b7de99d9b847560c262873a1f6c98202171fd5ed62640b494b/tomli-2.4.1-py3-none-any.whl", hash = "sha256:0d85819802132122da43cb86656f8d1f8c6587d54ae7dcaf30e90533028b49fe", size = 14583, upload-time = "2026-03-25T20:22:03.012Z" }, ] @@ -5171,12 +4551,6 @@ wheels = [ { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:9c8f38efee365cb9d334de8a83ce52fc7e5fc9e5a7b0853285efa1b69e00b0f2", upload-time = "2026-04-27T17:41:30Z" }, { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:d252cf975fb18c94a85336323ad425f473df56dab35a44b00399bd70c7a3b997", upload-time = "2026-04-27T17:42:06Z" }, { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp312-cp312-win_amd64.whl", hash = "sha256:7c78215c3af4f62e63f2b2e360f1722fc719b0853c7ac22666483d9810613a4c", upload-time = "2026-04-27T17:43:49Z" }, - { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:7db3580106bba044da5b8950f3fb8fe5f31999eaab3f6a3aa2ac5d202c3684d2", upload-time = "2026-04-27T17:45:35Z" }, - { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:db964b33c55035a72ab3e2162287af8f1cc276039c65d015740cc88c26dcedf7", upload-time = "2026-04-27T17:46:18Z" }, - { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp313-cp313-win_amd64.whl", hash = "sha256:6f367e62fd81b75cdf23ca4b75ced834d2db2cf98d1588ac935bde345de9de23", upload-time = "2026-04-27T17:48:09Z" }, - { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:cd1cf1005c5fe419194ee294b7b584ba5ad0f2fb1778b3fe5a7b9c3f4617ddbc", upload-time = "2026-04-27T17:50:01Z" }, - { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:74b628dbc71603977b09f4e140792c6e997081a35ef3421555f3f6e201b81210", upload-time = "2026-04-27T17:50:42Z" }, - { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp313-cp313t-win_amd64.whl", hash = "sha256:c2a5984deba8e001d166bf9cb83b8351f63a28b009e1a2fa0e4bbf08c90b259b", upload-time = "2026-04-27T17:52:32Z" }, ] [[package]] @@ -5219,10 +4593,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1e/00/4210d76ca7424981f04033ebe7e48816ab83287a62538747a58825db770c/torch-2.12.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:2de4e19b88a481482c6c75291f2d6a52eda3ce51f311b29aa9b68499c830c07c", size = 426382721, upload-time = "2026-06-17T21:06:41.842Z" }, { url = "https://files.pythonhosted.org/packages/76/1f/bc9f5a5aa569307076365f25afcebacb22e9c754b1bcfbaaa146627c7fda/torch-2.12.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:649e4ced014ba646f76f8cb9c9726735a6323eb321b7919f942790a923f90921", size = 532261322, upload-time = "2026-06-17T21:06:06.673Z" }, { url = "https://files.pythonhosted.org/packages/9e/49/c549461daa008159d006a76a991fbc2f26fa8bac27a4030c858463dcb20f/torch-2.12.1-cp312-cp312-win_amd64.whl", hash = "sha256:e86550597877fb272ddc52db2f85b82cb601ea7bd932576a0340152cae2200b3", size = 122988095, upload-time = "2026-06-17T21:07:44.9Z" }, - { url = "https://files.pythonhosted.org/packages/ff/4a/0300261818e1560d72cc160ac826005507e8b7ca0a35788b591436d05b4a/torch-2.12.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:c75e93173c700bccd6bfcc4a9d19ce242ab6dacd1f1781483027a16239b9e650", size = 87992358, upload-time = "2026-06-17T21:07:40.299Z" }, - { url = "https://files.pythonhosted.org/packages/30/a7/874a5ca05e8f159211dca7921060f7057acc1adb26431e119fd150623efc/torch-2.12.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:fcb61ccd20784b62bdd78ec84238a5cfb383b4994902e03bac95505ab360884c", size = 426386134, upload-time = "2026-06-17T21:07:31.481Z" }, - { url = "https://files.pythonhosted.org/packages/e1/75/20bb8fe9c1ad6538cce8cd0391b51927ae5af0b17ed1eab44b8824465dc1/torch-2.12.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:f4afc8083dff08719edbea346644476e3cec0cf40ebe256be0ee5d5b7c7e8c0d", size = 532268019, upload-time = "2026-06-17T21:05:37.925Z" }, - { url = "https://files.pythonhosted.org/packages/d1/fa/824ddb662af55b2eabc0dbb7b57c7c0b1bcd93693754a2b8509ec4d16490/torch-2.12.1-cp313-cp313-win_amd64.whl", hash = "sha256:f92609e3b3ce72f25e2eb780d043ced2480c1a86c47c852604fc7a9108648386", size = 122987777, upload-time = "2026-06-17T21:07:09.49Z" }, ] [[package]] @@ -5255,9 +4625,6 @@ wheels = [ { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.12.1%2Bcu130-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:2d3e87d41ffb340ddf8c99e2a690a29feea9f5271459dd57621cd11317a434f2", upload-time = "2026-06-18T02:38:10Z" }, { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.12.1%2Bcu130-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:4bafc356fbb622e2756179406825c3a56c17b401196435a1487c5b40c657706c", upload-time = "2026-06-18T02:38:36Z" }, { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.12.1%2Bcu130-cp312-cp312-win_amd64.whl", hash = "sha256:52c5da6a0898d5d3473c02bd304b7a3bc0b72e351c6f3bfa0783e45ef9f4cd61", upload-time = "2026-06-18T02:39:58Z" }, - { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.12.1%2Bcu130-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:2a2bb858316615b90b14ff27d0c732d5af85d066f5ee7bf81fad2c9215839be7", upload-time = "2026-06-18T02:41:12Z" }, - { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.12.1%2Bcu130-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:d5e1840442d2182957b3d2f778cc325c90fa5cb42aa8b1ac949f029e9bdd7f06", upload-time = "2026-06-18T02:41:44Z" }, - { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.12.1%2Bcu130-cp313-cp313-win_amd64.whl", hash = "sha256:76dd848312a40d29499614b714a4318841734ff309ad922f2365868395b6b054", upload-time = "2026-06-18T02:43:02Z" }, ] [[package]] @@ -5285,12 +4652,6 @@ wheels = [ { url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.26.0%2Bcu128-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:63e35234aed13b6edda37056f417b5c281249669db631e706811917af36b21d7", upload-time = "2026-04-09T23:21:35Z" }, { url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.26.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:ccf26b4b659cfce6f2208cb8326071d51c70219a34856dfdf468d1e19af52c0d", upload-time = "2026-03-23T15:36:22Z" }, { url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.26.0%2Bcu128-cp312-cp312-win_amd64.whl", hash = "sha256:8c0d1c4fbb2c9a4d5d41d0aaa87da20e525bcb2a154ce405725b0be59456804b", upload-time = "2026-04-09T23:21:36Z" }, - { url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.26.0%2Bcu128-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:c4a9cacd521f2a4df0bcd9d8e96704771b928f478f1f3067e4085bb53a1da298", upload-time = "2026-04-09T23:21:37Z" }, - { url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.26.0%2Bcu128-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:cb1f6184a7ba30fba40580e1a01a6604a86c55e79fdda187f40116ee680441ec", upload-time = "2026-03-23T15:36:22Z" }, - { url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.26.0%2Bcu128-cp313-cp313-win_amd64.whl", hash = "sha256:0232cb219927a52d6c98ff202f32d1cdf4802c2195a85fc1f1a0c1b0b4983a4d", upload-time = "2026-04-09T23:21:38Z" }, - { url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.26.0%2Bcu128-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:e594732552a8c2fee2ace9c6475c6c6904fc44ccca622ee6765a89a045416a44", upload-time = "2026-04-09T23:21:38Z" }, - { url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.26.0%2Bcu128-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:6168abc019803ac9e97efce27eafd2fdb33db04dcc54a86039537729e5047b29", upload-time = "2026-03-23T15:36:23Z" }, - { url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.26.0%2Bcu128-cp313-cp313t-win_amd64.whl", hash = "sha256:367d42ea703844ecdb516e9d5eb09929012a58705d2622cf4e9e3c37f278cb85", upload-time = "2026-04-09T23:21:39Z" }, ] [[package]] @@ -5321,10 +4682,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/43/42/103fa8f9366cfd1329fe449d6b1a25a640c0c17862ed48f21c4af94af322/torchvision-0.27.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:9edfb5a549fc2f30ccadb24eca907901e92e426c91a59316be6703a9360e5098", size = 7830902, upload-time = "2026-06-17T21:09:29.739Z" }, { url = "https://files.pythonhosted.org/packages/97/70/fa6052a42110a3657fc94073648da6171220469f4bf9f27e6a0b9378075c/torchvision-0.27.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:ae3d49e57c4abc8eafc1a1971f80fc4948a6268fa69340737ca4466936def080", size = 7664211, upload-time = "2026-06-17T21:09:17.206Z" }, { url = "https://files.pythonhosted.org/packages/d0/95/27aca854da7e536a339f46bab1ef67823ac2ac97c59ab2b3203b373d46cf/torchvision-0.27.1-cp312-cp312-win_amd64.whl", hash = "sha256:0b6e3aa98b7433506bbce1d0d05cb13ec787fc6eb8c5fbd998b26ce05f047543", size = 4079076, upload-time = "2026-06-17T21:09:15.907Z" }, - { url = "https://files.pythonhosted.org/packages/32/bb/b21e0f598ca191bb2a9e9fda2fee37c06ad113313b43c6769dbefa0e921d/torchvision-0.27.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:d60311a6d08df905f9656a3a312f0a8f55f0d46321bc737bad30a8dec9644309", size = 1852110, upload-time = "2026-06-17T21:09:22.577Z" }, - { url = "https://files.pythonhosted.org/packages/2f/90/d61171daa5d6cd5f9315f84f9ef947b047a9fdf283d53241327045a8dd6d/torchvision-0.27.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:08aa33bc8e062cca32aefa90ac714916c5a855cbe1ab4c6148fc0453eb40ca5a", size = 7789476, upload-time = "2026-06-17T21:09:13.105Z" }, - { url = "https://files.pythonhosted.org/packages/b8/dc/b21d7801562c23a770e7037989814582f22ca4db479204293561de4b62e8/torchvision-0.27.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:916448be4b19676677b0dbf47d08f68b7955ea0abec7fc79340c31e217a824ba", size = 7664256, upload-time = "2026-06-17T21:09:07.549Z" }, - { url = "https://files.pythonhosted.org/packages/b9/b3/4386976ff77eda55f0aed504a288564f3ff8d170b6db49ee22e172eddfac/torchvision-0.27.1-cp313-cp313-win_amd64.whl", hash = "sha256:18bc906235bfa901c135acd239f05b8c8ab90d502830cf1ef2cba3301e1f8a23", size = 4150710, upload-time = "2026-06-17T21:09:14.457Z" }, ] [[package]] @@ -5352,9 +4709,6 @@ wheels = [ { url = "https://download-r2.pytorch.org/whl/cu130/torchvision-0.27.1%2Bcu130-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:d2b15508b03a8949d8ff1e67a61342e9191c3fde2bf750996b24b3d472bcd7cb", upload-time = "2026-06-17T15:44:44Z" }, { url = "https://download-r2.pytorch.org/whl/cu130/torchvision-0.27.1%2Bcu130-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:abbfc724597c16da177002a16979aa8c44c4898c97bcb731b647cc57507f5772", upload-time = "2026-06-17T15:44:44Z" }, { url = "https://download-r2.pytorch.org/whl/cu130/torchvision-0.27.1%2Bcu130-cp312-cp312-win_amd64.whl", hash = "sha256:1bf254c102bfaf97d3e7878b76b68999bcd4dcd4303c109e76b4fbf9b15265c5", upload-time = "2026-06-18T04:00:18Z" }, - { url = "https://download-r2.pytorch.org/whl/cu130/torchvision-0.27.1%2Bcu130-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:e1d81f9e5f99a73e239e143b73d143181ab2e71a8c8ef79fa90e908cc356218c", upload-time = "2026-06-17T15:44:44Z" }, - { url = "https://download-r2.pytorch.org/whl/cu130/torchvision-0.27.1%2Bcu130-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:f9f008d4b1b2f013eaf7bec1a9ff221263581d77668d4e1e0e9c3ed351d56465", upload-time = "2026-06-17T15:44:44Z" }, - { url = "https://download-r2.pytorch.org/whl/cu130/torchvision-0.27.1%2Bcu130-cp313-cp313-win_amd64.whl", hash = "sha256:04f6bda2ab9589ad0a63d44fe6e1dc3ff86f379c100a491b6daba9ec9ec72499", upload-time = "2026-06-18T04:00:19Z" }, ] [[package]] @@ -5460,10 +4814,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e0/12/b05ba554d2c623bffa59922b94b0775673de251f468a9609bc9e45de95e9/triton-3.6.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e8e323d608e3a9bfcc2d9efcc90ceefb764a82b99dea12a86d643c72539ad5d3", size = 188214640, upload-time = "2026-01-20T16:00:35.869Z" }, { url = "https://files.pythonhosted.org/packages/17/5d/08201db32823bdf77a0e2b9039540080b2e5c23a20706ddba942924ebcd6/triton-3.6.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:374f52c11a711fd062b4bfbb201fd9ac0a5febd28a96fb41b4a0f51dde3157f4", size = 176128243, upload-time = "2026-01-20T16:16:07.857Z" }, { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, - { url = "https://files.pythonhosted.org/packages/3c/12/34d71b350e89a204c2c7777a9bba0dcf2f19a5bfdd70b57c4dbc5ffd7154/triton-3.6.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:448e02fe6dc898e9e5aa89cf0ee5c371e99df5aa5e8ad976a80b93334f3494fd", size = 176133521, upload-time = "2026-01-20T16:16:13.321Z" }, - { url = "https://files.pythonhosted.org/packages/f9/0b/37d991d8c130ce81a8728ae3c25b6e60935838e9be1b58791f5997b24a54/triton-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10c7f76c6e72d2ef08df639e3d0d30729112f47a56b0c81672edc05ee5116ac9", size = 188289450, upload-time = "2026-01-20T16:00:49.136Z" }, - { url = "https://files.pythonhosted.org/packages/ce/4e/41b0c8033b503fd3cfcd12392cdd256945026a91ff02452bef40ec34bee7/triton-3.6.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1722e172d34e32abc3eb7711d0025bb69d7959ebea84e3b7f7a341cd7ed694d6", size = 176276087, upload-time = "2026-01-20T16:16:18.989Z" }, - { url = "https://files.pythonhosted.org/packages/35/f8/9c66bfc55361ec6d0e4040a0337fb5924ceb23de4648b8a81ae9d33b2b38/triton-3.6.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d002e07d7180fd65e622134fbd980c9a3d4211fb85224b56a0a0efbd422ab72f", size = 188400296, upload-time = "2026-01-20T16:00:56.042Z" }, ] [[package]] @@ -5482,8 +4832,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cd/5e/fce69606f7f240297f163e25539906732b199530d486ce67ae319877e821/triton-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6744957e9fd610a29680ec2346057d0c86948ed3812468670719f391e94b44a5", size = 197701306, upload-time = "2026-06-17T19:53:13.673Z" }, { url = "https://files.pythonhosted.org/packages/94/fa/f856e24deb462d5f18bd4b5a746957862ab9b6ee5834bda60605ec348366/triton-3.7.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9497f2e696ee368862a181a90b2dcc03ca978cc4f602abd67c7d81022a6988e1", size = 184692359, upload-time = "2026-06-17T20:03:48.288Z" }, { url = "https://files.pythonhosted.org/packages/c4/6f/fb96d15db6f36d6eae4cafb998c2e0353bf59d7c4ea1662d7497f269134a/triton-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7e40869937a68206ec70d7f25bb7ec6433cb083f9135e1f36dbd318dc449a728", size = 197719725, upload-time = "2026-06-17T19:53:20.419Z" }, - { url = "https://files.pythonhosted.org/packages/00/42/c5089d4d9327fcd1e862c599cc2927f39418f84dd11a84cb2ccff9d4787a/triton-3.7.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cdbfc09d9ec58bc5e68321525653220de7515c199e7a8097a97c85e62b52cd0a", size = 184694629, upload-time = "2026-06-17T20:03:53.444Z" }, - { url = "https://files.pythonhosted.org/packages/07/42/2c3ac59253ae8892b6f307875263dd23dc875cdf732d3aea40d6d41fb7cb/triton-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:58c0e131da05134a2a4788ccbcc0c1105cf0f54c8e98f19e34cd465396dc15eb", size = 197729241, upload-time = "2026-06-17T19:53:27.801Z" }, ] [[package]] @@ -5494,7 +4842,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4d/b4/c50a22dd2d493a3e8e78744cbdcb29932f367b68999c8e08874b94aebd3a/triton_windows-3.7.1.post27-cp310-cp310-win_amd64.whl", hash = "sha256:01ede775b102c91acc89fc6946b946451f966cb64856c650d244d37ea3bbdee5", size = 49678783, upload-time = "2026-06-21T16:47:51.235Z" }, { url = "https://files.pythonhosted.org/packages/26/f5/0f5eaf48abc0c9900f600dbdfa8139e678aa7d47dc1da51b0541979e96df/triton_windows-3.7.1.post27-cp311-cp311-win_amd64.whl", hash = "sha256:b739bd7d39f919280294d8af172a90aa2f17a4377bfbca2ea30a8afae61d5eaa", size = 49679173, upload-time = "2026-06-21T16:48:02.395Z" }, { url = "https://files.pythonhosted.org/packages/76/30/325b420efd0047e119679c646a9a410db216069800ec009fae3da26c69a3/triton_windows-3.7.1.post27-cp312-cp312-win_amd64.whl", hash = "sha256:f5406230d7dbf6965bc4051fcad27b81c39ba4a4bfde06f494dd7ff4eb325a9e", size = 49683004, upload-time = "2026-06-21T16:48:14.02Z" }, - { url = "https://files.pythonhosted.org/packages/ec/28/f0b2801c2cfd79be5878bf429e82b5da00fb78f046a9c21ba946681cc467/triton_windows-3.7.1.post27-cp313-cp313-win_amd64.whl", hash = "sha256:e8ed215c02afc85a81f0097196f8bebdf6f68085f1bc0fe8771b9a74783f15ab", size = 49684257, upload-time = "2026-06-21T16:48:24.901Z" }, ] [[package]] @@ -5682,31 +5029,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fa/08/d9e2e0f9e8e6791d33aefc694ad7eefa7f901f63caff84a81ded38692f9c/watchfiles-1.2.0-cp312-cp312-win32.whl", hash = "sha256:7a2cffd17d27d2ecbb310c2b1d8174f222a5495b1a721894afa88ec11e25b898", size = 275480, upload-time = "2026-05-18T04:30:31.307Z" }, { url = "https://files.pythonhosted.org/packages/1c/e6/9d42569c0102645cc8cea5d8c7d8a1e9d4ada2cb7f05f75e554b8aa2202a/watchfiles-1.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:f155b3a1b2a5fc89cdc70d47ee5d54e3b75e88efa34982028a35daef9ba00379", size = 288718, upload-time = "2026-05-18T04:32:10.745Z" }, { url = "https://files.pythonhosted.org/packages/0a/26/88e0dc6ee3898169d7fa22bb6a69cabf2502d2ee25cb8c876d1262d204f8/watchfiles-1.2.0-cp312-cp312-win_arm64.whl", hash = "sha256:8fa585ede612ee9f9e91b18bebf9ba11b9ae29a4e3a0d0cf6fca3e382133f0d5", size = 281026, upload-time = "2026-05-18T04:30:22.23Z" }, - { url = "https://files.pythonhosted.org/packages/d1/4d/70a7feced9f87e2ff26dba42667290f41694fc64646c67261fbb8cab5d5c/watchfiles-1.2.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:01ea8d66f0693b9b60a6541c8d10263091ca9a9060d242f3c1f3143f9aad2c98", size = 399730, upload-time = "2026-05-18T04:31:38.162Z" }, - { url = "https://files.pythonhosted.org/packages/31/3a/0da302f2307aee316922806ebd5726c542cbd787c938271cf14a074c7daf/watchfiles-1.2.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7ba0480b9a74af058f43b337e937a451e109295c420916d68ad24e3dc02f5e44", size = 392842, upload-time = "2026-05-18T04:30:27.051Z" }, - { url = "https://files.pythonhosted.org/packages/db/ef/d5bdb705c224dbc256aa0c1ec47bf4e61ec52558f2afb44a71a1fe4d7015/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4f34e26a19f91f710c08e0183429f0d1d15df734e6bc78c31e77b9ea9c433658", size = 452989, upload-time = "2026-05-18T04:31:11.945Z" }, - { url = "https://files.pythonhosted.org/packages/71/29/5495f2c1661949ef7a35e4d71111d129cfe7606414a26887a919d0a55406/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:b4e77f6a55f858504069abd35d336a637555c09bca453dde1ee1e5ada8a6a1fb", size = 458978, upload-time = "2026-05-18T04:30:52.606Z" }, - { url = "https://files.pythonhosted.org/packages/d5/8c/7f9c07c433811c2fffd93e13fdfb7135de9aab5f2ae41be08960fa0047dc/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0cb4d80e212f116474a545c21c912b445f16bb0cef9e6a73a498164223e14e2f", size = 490248, upload-time = "2026-05-18T04:31:36.003Z" }, - { url = "https://files.pythonhosted.org/packages/3c/11/d93632febc52fbc21be90231bb7c17fd5387f46c9076fd40a5f9c2ae6910/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b974946a10af379d425e2eef5b62f5c6ebeaccf91d45eaad6f5b27ecd4f91aa0", size = 571847, upload-time = "2026-05-18T04:31:10.862Z" }, - { url = "https://files.pythonhosted.org/packages/55/b4/383173e73aabb07ad1d9c7aa859d95437ac46a6d6a1e11005facda0c9d19/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:86bc13c25a8d1fcd70b51d0ce7c9b65e90de5666fcbfd3e34957cc73ee19aeb5", size = 465974, upload-time = "2026-05-18T04:30:17.006Z" }, - { url = "https://files.pythonhosted.org/packages/a7/6c/89b1a230a78f57c52dd8893adb1f92f94411721b6ec12596c56d98c74356/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ca148d73dea36c9763aaa351e4d7a51780ec1584217c45276f4fe8239c768b71", size = 454782, upload-time = "2026-05-18T04:30:35.656Z" }, - { url = "https://files.pythonhosted.org/packages/24/62/1732118367cfff0a9fce3bf62ff4bfded09ef5df21d9d446b858b3f70a96/watchfiles-1.2.0-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:c525543d91961c6955b2636b308569e84a1d1c5f5f2932041ab9ef46422f43e3", size = 465182, upload-time = "2026-05-18T04:30:20.846Z" }, - { url = "https://files.pythonhosted.org/packages/28/96/716f7e5f51339bf22963f3345f9f27d7f3b30e2eadc597e257c881dd3c53/watchfiles-1.2.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:a204794696ffb8f9b10fba6f7cb5216d42f3b2b71860ccac6b6e42f5f10973b0", size = 629841, upload-time = "2026-05-18T04:31:05.397Z" }, - { url = "https://files.pythonhosted.org/packages/4c/fe/c40783950fd771ccf66ab3ec2722d188a9af1c7f96c6e811f36e40c6e03f/watchfiles-1.2.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:10d86db20695afe7997ac9e1717637d6714a8d0220458c33f3d2061f54cec427", size = 658028, upload-time = "2026-05-18T04:31:48.22Z" }, - { url = "https://files.pythonhosted.org/packages/71/72/4508db1856d1d87fcbb3b63f4839bab1b5682cb0e8d224d122263c09654a/watchfiles-1.2.0-cp313-cp313-win32.whl", hash = "sha256:eb283ee99e21ad6443c8cdb06ac5b34b1308c329cbdf03fa02b445363714c799", size = 275183, upload-time = "2026-05-18T04:30:59.57Z" }, - { url = "https://files.pythonhosted.org/packages/f9/36/14b76ca57652e5cc5fd1c11f32a261292c08a0d19a00351013c2549cbfb2/watchfiles-1.2.0-cp313-cp313-win_amd64.whl", hash = "sha256:a0f27f01bee51861392bb6b7c4fdb290b27d1eb194e9e28788d68102a0e898d9", size = 288059, upload-time = "2026-05-18T04:32:07.937Z" }, - { url = "https://files.pythonhosted.org/packages/1b/8d/0a85e395398d8d20fadfe5c5d32c726eee17a519e78fb356f2cf7531bffe/watchfiles-1.2.0-cp313-cp313-win_arm64.whl", hash = "sha256:3651aa7058595e9cfb75d35dd5ada2bf9f48a5b8a0f3562821d3e210c507e077", size = 280186, upload-time = "2026-05-18T04:31:54.484Z" }, - { url = "https://files.pythonhosted.org/packages/37/68/36db056f1fdcc5f07302f56e631774d6835bcd6fa3ace402304621d5f9e5/watchfiles-1.2.0-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:faea288b6f0ab1902ef08f4ca6de005dccf856c4e0c4f21b8c5fce02d90a1b08", size = 399031, upload-time = "2026-05-18T04:30:44.576Z" }, - { url = "https://files.pythonhosted.org/packages/c1/64/01a9d6f66a82a5c101ce939274106cc72759d62427e153f01edd2b9f87c2/watchfiles-1.2.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:01859b11fd9fbca670f4d5da00fbac282cfea9bd67a2125d8b2833a3b5617ea9", size = 391205, upload-time = "2026-05-18T04:30:25.413Z" }, - { url = "https://files.pythonhosted.org/packages/84/2c/0a44fe058cb4bb7b8ede6b6670698bbb7c0400740e378d00022189b7b31d/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:fff610d7bb2256a317bb1e96f0d7862c7aa8076733ee5df0fd41bbe76a24a4f4", size = 451892, upload-time = "2026-05-18T04:32:14.005Z" }, - { url = "https://files.pythonhosted.org/packages/67/a1/351e0d56cd35e6488b5c8b4fb11a809a5bc923e8fe8fed9faf8920be0c89/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:b141a4891c995a039cd89e9a49e62df1dc8a559a5d1a6e4c7106d16c12777a55", size = 458867, upload-time = "2026-05-18T04:31:22.279Z" }, - { url = "https://files.pythonhosted.org/packages/d5/7d/9d09605187f1b838998624049fcf8bf47b73c1a3b76901fcac1782f62277/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f22943b7770483f6ea0721c6b11d022947a98eb0acae14694de034f4d0d38925", size = 490217, upload-time = "2026-05-18T04:31:43.657Z" }, - { url = "https://files.pythonhosted.org/packages/60/5d/a17a16eccb182f04188cd308ec24b1a71a9b5c4e7098269cf35d9fa56d02/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1bc6195825b7dcd217968bb1f801a60fd4c16e8eeab5bedc7fe917d7d5995ab4", size = 571458, upload-time = "2026-05-18T04:32:11.875Z" }, - { url = "https://files.pythonhosted.org/packages/d3/3d/4dd457062083ab1938e5dfd45032eb425cee2ac817287ca8ff4356183e5d/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d4a4b147f5dca2a5d325a06a832fb43f345751adfbc63204aec30e0d9ca965a2", size = 464707, upload-time = "2026-05-18T04:30:43.492Z" }, - { url = "https://files.pythonhosted.org/packages/c6/71/ea8c57b128f5383de74d0c7d2d9c57ad7c9a65a930c451bd25d524b295b7/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4543579a9bdb0c9560039b4ffddbdb39545707659fbc430ce4c10f3f68d557f9", size = 454663, upload-time = "2026-05-18T04:30:16.061Z" }, - { url = "https://files.pythonhosted.org/packages/53/fd/2e812bf938406d7db351f0703ddd3fc6c061cf30d96153a77bc79a943a44/watchfiles-1.2.0-cp313-cp313t-manylinux_2_31_riscv64.whl", hash = "sha256:20aa0e708b920bde876a4aa82dc7dd6ebea228a63a67cda6632c2fc87b787efa", size = 463537, upload-time = "2026-05-18T04:31:44.9Z" }, - { url = "https://files.pythonhosted.org/packages/86/56/d17a7f1dd1bc3035f1072694a551301272f1739c2d8e319c927cb9e29b38/watchfiles-1.2.0-cp313-cp313t-musllinux_1_1_aarch64.whl", hash = "sha256:d413349d565dab74297f2a63e84a097936be69bf8f3b3801f27f380e32040f44", size = 629194, upload-time = "2026-05-18T04:31:14.141Z" }, - { url = "https://files.pythonhosted.org/packages/be/06/f1ff66bf5cae50aa4062779a0ecd0bbaf15e466195719074078947d9a17d/watchfiles-1.2.0-cp313-cp313t-musllinux_1_1_x86_64.whl", hash = "sha256:f28b2725eb8cce327b9b3ab02415c853011dc55c95832fe90de6bc56f5315f72", size = 656194, upload-time = "2026-05-18T04:31:47.14Z" }, { url = "https://files.pythonhosted.org/packages/23/f4/7513ef1e85fc4c6331b59479d6d72661fc391fbe543678052ac72c8b6c19/watchfiles-1.2.0-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:4674d49eb94706dfe666c069fc0a1b646ffcf920473492e209f6d5f60d3f0cc2", size = 403050, upload-time = "2026-05-18T04:30:36.753Z" }, { url = "https://files.pythonhosted.org/packages/27/0b/a54103cfd732bb703c7a749222011a0483ef3705948dae3b203158601119/watchfiles-1.2.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:094b9b70103d4e963499bdea001ee3c2697b144cd9ae6218a62c0f89ec9e31db", size = 396629, upload-time = "2026-05-18T04:32:03.268Z" }, { url = "https://files.pythonhosted.org/packages/5e/2c/73f31a3b893886206c3f54d73e8ad8dee58cdb2f69ad2622e0a8a9e07f4e/watchfiles-1.2.0-pp311-pypy311_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b0ef001f8c25ad0fa9529f914c1600647ecd0f542d11c19b7894768c67b6acb7", size = 457318, upload-time = "2026-05-18T04:31:01.932Z" }, @@ -5755,15 +5077,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/47/88/4dd516068e1a3d6ab3c7c183288404cd424a9a02d585efbac226cb61ff2d/websockets-16.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:485c49116d0af10ac698623c513c1cc01c9446c058a4e61e3bf6c19dff7335a2", size = 184880, upload-time = "2026-01-10T09:22:55.033Z" }, { url = "https://files.pythonhosted.org/packages/91/d6/7d4553ad4bf1c0421e1ebd4b18de5d9098383b5caa1d937b63df8d04b565/websockets-16.0-cp312-cp312-win32.whl", hash = "sha256:eaded469f5e5b7294e2bdca0ab06becb6756ea86894a47806456089298813c89", size = 178261, upload-time = "2026-01-10T09:22:56.251Z" }, { url = "https://files.pythonhosted.org/packages/c3/f0/f3a17365441ed1c27f850a80b2bc680a0fa9505d733fe152fdf5e98c1c0b/websockets-16.0-cp312-cp312-win_amd64.whl", hash = "sha256:5569417dc80977fc8c2d43a86f78e0a5a22fee17565d78621b6bb264a115d4ea", size = 178693, upload-time = "2026-01-10T09:22:57.478Z" }, - { url = "https://files.pythonhosted.org/packages/cc/9c/baa8456050d1c1b08dd0ec7346026668cbc6f145ab4e314d707bb845bf0d/websockets-16.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:878b336ac47938b474c8f982ac2f7266a540adc3fa4ad74ae96fea9823a02cc9", size = 177364, upload-time = "2026-01-10T09:22:59.333Z" }, - { url = "https://files.pythonhosted.org/packages/7e/0c/8811fc53e9bcff68fe7de2bcbe75116a8d959ac699a3200f4847a8925210/websockets-16.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:52a0fec0e6c8d9a784c2c78276a48a2bdf099e4ccc2a4cad53b27718dbfd0230", size = 175039, upload-time = "2026-01-10T09:23:01.171Z" }, - { url = "https://files.pythonhosted.org/packages/aa/82/39a5f910cb99ec0b59e482971238c845af9220d3ab9fa76dd9162cda9d62/websockets-16.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:e6578ed5b6981005df1860a56e3617f14a6c307e6a71b4fff8c48fdc50f3ed2c", size = 175323, upload-time = "2026-01-10T09:23:02.341Z" }, - { url = "https://files.pythonhosted.org/packages/bd/28/0a25ee5342eb5d5f297d992a77e56892ecb65e7854c7898fb7d35e9b33bd/websockets-16.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:95724e638f0f9c350bb1c2b0a7ad0e83d9cc0c9259f3ea94e40d7b02a2179ae5", size = 184975, upload-time = "2026-01-10T09:23:03.756Z" }, - { url = "https://files.pythonhosted.org/packages/f9/66/27ea52741752f5107c2e41fda05e8395a682a1e11c4e592a809a90c6a506/websockets-16.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c0204dc62a89dc9d50d682412c10b3542d748260d743500a85c13cd1ee4bde82", size = 186203, upload-time = "2026-01-10T09:23:05.01Z" }, - { url = "https://files.pythonhosted.org/packages/37/e5/8e32857371406a757816a2b471939d51c463509be73fa538216ea52b792a/websockets-16.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:52ac480f44d32970d66763115edea932f1c5b1312de36df06d6b219f6741eed8", size = 185653, upload-time = "2026-01-10T09:23:06.301Z" }, - { url = "https://files.pythonhosted.org/packages/9b/67/f926bac29882894669368dc73f4da900fcdf47955d0a0185d60103df5737/websockets-16.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6e5a82b677f8f6f59e8dfc34ec06ca6b5b48bc4fcda346acd093694cc2c24d8f", size = 184920, upload-time = "2026-01-10T09:23:07.492Z" }, - { url = "https://files.pythonhosted.org/packages/3c/a1/3d6ccdcd125b0a42a311bcd15a7f705d688f73b2a22d8cf1c0875d35d34a/websockets-16.0-cp313-cp313-win32.whl", hash = "sha256:abf050a199613f64c886ea10f38b47770a65154dc37181bfaff70c160f45315a", size = 178255, upload-time = "2026-01-10T09:23:09.245Z" }, - { url = "https://files.pythonhosted.org/packages/6b/ae/90366304d7c2ce80f9b826096a9e9048b4bb760e44d3b873bb272cba696b/websockets-16.0-cp313-cp313-win_amd64.whl", hash = "sha256:3425ac5cf448801335d6fdc7ae1eb22072055417a96cc6b31b3861f455fbc156", size = 178689, upload-time = "2026-01-10T09:23:10.483Z" }, { url = "https://files.pythonhosted.org/packages/72/07/c98a68571dcf256e74f1f816b8cc5eae6eb2d3d5cfa44d37f801619d9166/websockets-16.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:349f83cd6c9a415428ee1005cadb5c2c56f4389bc06a9af16103c3bc3dcc8b7d", size = 174947, upload-time = "2026-01-10T09:23:36.166Z" }, { url = "https://files.pythonhosted.org/packages/7e/52/93e166a81e0305b33fe416338be92ae863563fe7bce446b0f687b9df5aea/websockets-16.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:4a1aba3340a8dca8db6eb5a7986157f52eb9e436b74813764241981ca4888f03", size = 175260, upload-time = "2026-01-10T09:23:37.409Z" }, { url = "https://files.pythonhosted.org/packages/56/0c/2dbf513bafd24889d33de2ff0368190a0e69f37bcfa19009ef819fe4d507/websockets-16.0-pp311-pypy311_pp73-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:f4a32d1bd841d4bcbffdcb3d2ce50c09c3909fbead375ab28d0181af89fd04da", size = 176071, upload-time = "2026-01-10T09:23:39.158Z" }, @@ -5843,23 +5156,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/48/2d/1c8d89c7c5f9cad9fb2902445d94e2ab1d7aa35de029afbb8ae95c42d00f/yarl-1.24.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e30dd55825dc554ec5b66a94953b8eda8745926514c5089dfcacecb9c99b5bd1", size = 105719, upload-time = "2026-05-19T21:29:18.367Z" }, { url = "https://files.pythonhosted.org/packages/a7/25/722e3b93bd687009afb2d59a35e13d30ddd8f80571445bb0c4e4ce26ec66/yarl-1.24.2-cp312-cp312-win_amd64.whl", hash = "sha256:7dafe10c12ddd4d120d528c4b5599c953bd7b12845347d507b95451195bb6cad", size = 92901, upload-time = "2026-05-19T21:29:20.014Z" }, { url = "https://files.pythonhosted.org/packages/39/47/4486ccfb674c04854a1ef8aa77868b6a6f765feaf69633409d7ca4f02cb8/yarl-1.24.2-cp312-cp312-win_arm64.whl", hash = "sha256:044a09d8401fcf8681977faef6d286b8ade1e2d2e9dceda175d1cfa5ca496f30", size = 87229, upload-time = "2026-05-19T21:29:22.1Z" }, - { url = "https://files.pythonhosted.org/packages/82/62/fcf0ce677f17e5c471c06311dd25964be38a4c586993632910d2e75278bc/yarl-1.24.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:491ac9141decf49ee8030199e1ee251cdff0e131f25678817ff6aa5f837a3536", size = 128978, upload-time = "2026-05-19T21:29:23.83Z" }, - { url = "https://files.pythonhosted.org/packages/d3/58/8e63299bb71ed61a834121d9d3fe6c9fcf2a6a5d09754ff4f20f2d20baf5/yarl-1.24.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e89418f65eda18f99030386305bd44d7d504e328a7945db1ead514fbe03a0607", size = 91733, upload-time = "2026-05-19T21:29:25.375Z" }, - { url = "https://files.pythonhosted.org/packages/c1/24/16748d5dab6daec8b0ed81ccec639a1cded0f18dcc62a4f696b4fe366c37/yarl-1.24.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:cdfcce633b4a4bb8281913c57fcafd4b5933fbc19111a5e3930bbd299d6102f1", size = 91113, upload-time = "2026-05-19T21:29:26.928Z" }, - { url = "https://files.pythonhosted.org/packages/1b/66/b63fff7b71211e866624b21432d5943cbb633eb0c2872d9ee3070648f22c/yarl-1.24.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:863297ddede92ee49024e9a9b11ecb59f310ca85b60d8537f56bed9bbb5b1986", size = 103899, upload-time = "2026-05-19T21:29:28.842Z" }, - { url = "https://files.pythonhosted.org/packages/9d/ac/ba1974b8533909636f7733fe86cf677e3619527c3c2fa913e0ea89c48757/yarl-1.24.2-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:374423f70754a2c96942ede36a29d37dc6b0cb8f92f8d009ddf3ed78d3da5488", size = 97862, upload-time = "2026-05-19T21:29:31.086Z" }, - { url = "https://files.pythonhosted.org/packages/1b/a5/123ac993b5c2ba6f554a140305620cb8f150fa543711bbc49be3ec0a65a4/yarl-1.24.2-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:33a29b5d00ccbf3219bb3e351d7875739c19481e030779f48cc46a7a71681a9b", size = 111060, upload-time = "2026-05-19T21:29:32.657Z" }, - { url = "https://files.pythonhosted.org/packages/23/37/c472d3af3509688392134a88a825276770a187f1daa4de3f6dc0a327a751/yarl-1.24.2-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a9532c57211730c515341af11fef6e9b61d157487272a096d0c04da445642592", size = 110613, upload-time = "2026-05-19T21:29:34.379Z" }, - { url = "https://files.pythonhosted.org/packages/df/88/09c28dad91e662ccfaa1b78f1c57badde74fc9d0b23e74aef644750ecd73/yarl-1.24.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:91e72cf093fd833483a97ee648e0c053c7c629f51ff4a0e7edd84f806b0c5617", size = 107012, upload-time = "2026-05-19T21:29:36.216Z" }, - { url = "https://files.pythonhosted.org/packages/07/ab/9d4f69d571a94f4d112fa7e2e007200f5a54d319f58c82ac7b7baa61f5c6/yarl-1.24.2-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b3177bc0a768ef3bacceb4f272632990b7bea352f1b2f1eee9d6d6ff16516f92", size = 105887, upload-time = "2026-05-19T21:29:38.746Z" }, - { url = "https://files.pythonhosted.org/packages/8e/9a/000b2b66c0d772a499fc531d21dab92dfeb73b640a12eed6ba89f49bb2d0/yarl-1.24.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:e196952aacaf3b232e265ff02980b64d483dc0972bd49bcb061171ff22ac203a", size = 103620, upload-time = "2026-05-19T21:29:40.368Z" }, - { url = "https://files.pythonhosted.org/packages/41/7c/7c1050f73450fbdaa3f0c72017059f00ce5e13366692f3dba25275a1083d/yarl-1.24.2-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:204e7a61ce99919c0de1bf904ab5d7aa188a129ea8f690a8f76cfb6e2844dc44", size = 100599, upload-time = "2026-05-19T21:29:42.66Z" }, - { url = "https://files.pythonhosted.org/packages/ec/b1/29e5756b3926705f5f6089bd5b9f50a56eaac550da6e260bf713ead44d04/yarl-1.24.2-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:4b156914620f0b9d78dc1adb3751141daee561cfec796088abb89ed49d220f1a", size = 110604, upload-time = "2026-05-19T21:29:44.632Z" }, - { url = "https://files.pythonhosted.org/packages/a3/4b/8415bc96e9b150cde942fbac9a8182985e58f40ce5c54c34ed015407d3ee/yarl-1.24.2-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:8372a2b976cf70654b2be6619ab6068acabb35f724c0fda7b277fbf53d66a5cf", size = 105161, upload-time = "2026-05-19T21:29:46.755Z" }, - { url = "https://files.pythonhosted.org/packages/8b/d4/cde059abfa229553b7298a2eadde2752e723d50aeedaef86ce59da2718ee/yarl-1.24.2-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:f9a1e9b622ca284143aab5d885848686dcd85453bb1ca9abcdb7503e64dc0056", size = 110619, upload-time = "2026-05-19T21:29:48.972Z" }, - { url = "https://files.pythonhosted.org/packages/e7/2c/d6a6c9a61549f7b6c7e6dc6937d195bcf069582b47b7200dcd0e7b256acf/yarl-1.24.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:810e19b685c8c3c5862f6a38160a1f4e4c0916c9390024ec347b6157a45a0992", size = 107362, upload-time = "2026-05-19T21:29:51Z" }, - { url = "https://files.pythonhosted.org/packages/92/dd/3ae5fe417e9d1c353a548553326eb9935e76b6b727161563b424cc296df3/yarl-1.24.2-cp313-cp313-win_amd64.whl", hash = "sha256:7d37fb7c38f2b6edab0f845c4f85148d4c44204f52bc127021bd2bc9fdbf1656", size = 92667, upload-time = "2026-05-19T21:29:52.743Z" }, - { url = "https://files.pythonhosted.org/packages/10/cc/a7beb239f78f27fca1b053c8e8595e4179c02e62249b4687ec218c370c50/yarl-1.24.2-cp313-cp313-win_arm64.whl", hash = "sha256:1e831894be7c2954240e49791fa4b50c05a0dc881de2552cfe3ffd8631c7f461", size = 87069, upload-time = "2026-05-19T21:29:54.442Z" }, { url = "https://files.pythonhosted.org/packages/fd/4d/4b880086bd0d3e034d25647be1d830afc3e3f610e98c4ab3490af6b1b6d5/yarl-1.24.2-py3-none-any.whl", hash = "sha256:2783d9226db8797636cd6896e4de81feed252d1db72265686c9558d97a4d94b9", size = 53576, upload-time = "2026-05-19T21:31:03.909Z" }, ] From 8068c8e33176ba1576c6086382b749798221beb8 Mon Sep 17 00:00:00 2001 From: aidanfnv Date: Thu, 6 Aug 2026 18:57:21 -0700 Subject: [PATCH 09/19] Port Lingbot to new API (#425) * initial commit * more porting * Claude reviewed fixes * Add keyboard parity and event-driven GPU tests for the Lingbot port Pin the runtime-API camera path against the WebRTC path it will replace. KeyboardResampler + CameraPoseIntegrator and the canonicalizer + mapping path are compared over 11 key streams, single- and multi-chunk, including edges landing exactly on a chunk boundary where KeyboardResampler's inclusive drain meets TimeWindow's half-open slice. They agree. Also cover event-driven camera control on CUDA, and repair the existing CUDA test, which monkeypatched trace-loading helpers that moved out of the session and called step() with an empty InferenceInput. Co-Authored-By: Claude Opus 5 (1M context) * Update integrations/lingbot/lingbot/runtime.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: aidanfnv --------- Signed-off-by: aidanfnv Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../flashdreams/serving/output_targets.py | 16 +- flashdreams/tests/test_output_targets.py | 9 +- integrations/lingbot/lingbot/demo/__init__.py | 20 + integrations/lingbot/lingbot/demo/adapter.py | 277 +++++ integrations/lingbot/lingbot/demo/cli.py | 233 ++++ integrations/lingbot/lingbot/demo/replay.py | 20 + integrations/lingbot/lingbot/demo/spec.py | 167 +++ integrations/lingbot/lingbot/demo/webrtc.py | 57 + integrations/lingbot/lingbot/example_data.py | 86 ++ integrations/lingbot/lingbot/input_mapping.py | 670 ++++++++++ integrations/lingbot/lingbot/runner.py | 229 +--- integrations/lingbot/lingbot/runtime.py | 1084 +++++++++++++++++ integrations/lingbot/lingbot/webrtc/server.py | 60 +- .../lingbot/lingbot/webrtc/session.py | 29 +- integrations/lingbot/pyproject.toml | 3 + integrations/lingbot/tests/test_demo_api.py | 625 ++++++++++ .../lingbot/tests/test_input_mapping.py | 448 +++++++ .../lingbot/tests/test_keyboard_parity.py | 157 +++ .../lingbot/tests/test_runtime_gpu.py | 259 ++++ .../tests/test_runtime_session_inputs.py | 395 ++++++ integrations/lingbot/tests/test_smoke.py | 98 +- 21 files changed, 4704 insertions(+), 238 deletions(-) create mode 100644 integrations/lingbot/lingbot/demo/__init__.py create mode 100644 integrations/lingbot/lingbot/demo/adapter.py create mode 100644 integrations/lingbot/lingbot/demo/cli.py create mode 100644 integrations/lingbot/lingbot/demo/replay.py create mode 100644 integrations/lingbot/lingbot/demo/spec.py create mode 100644 integrations/lingbot/lingbot/demo/webrtc.py create mode 100644 integrations/lingbot/lingbot/example_data.py create mode 100644 integrations/lingbot/lingbot/input_mapping.py create mode 100644 integrations/lingbot/lingbot/runtime.py create mode 100644 integrations/lingbot/tests/test_demo_api.py create mode 100644 integrations/lingbot/tests/test_input_mapping.py create mode 100644 integrations/lingbot/tests/test_keyboard_parity.py create mode 100644 integrations/lingbot/tests/test_runtime_gpu.py create mode 100644 integrations/lingbot/tests/test_runtime_session_inputs.py diff --git a/flashdreams/flashdreams/serving/output_targets.py b/flashdreams/flashdreams/serving/output_targets.py index b7e7be981..f5306dbbf 100644 --- a/flashdreams/flashdreams/serving/output_targets.py +++ b/flashdreams/flashdreams/serving/output_targets.py @@ -154,7 +154,8 @@ def _lingbot_webrtc_spec( options: OutputLaunchOptions, ) -> OutputTargetSpec: argv = [ - "--config_name", + "webrtc", + "--preset-id", _pipeline_name(config), "--device", _device(config), @@ -166,15 +167,20 @@ def _lingbot_webrtc_spec( str(getattr(config, "pixel_width", 832)), ] if _compile_network(config) is False: - argv.append("--no_compile") + argv.append("--no-compile") example_idx = getattr(config, "example_idx", None) if example_idx is not None: argv.extend(("--example-idx", str(example_idx))) - _append_webrtc_bind_args(argv, options) + 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 WebRTC server", - module="lingbot.webrtc.server", + label="LingBot shared demo WebRTC server", + module="lingbot.demo.cli", argv=tuple(argv), ) diff --git a/flashdreams/tests/test_output_targets.py b/flashdreams/tests/test_output_targets.py index 171e04786..be4e11fbc 100644 --- a/flashdreams/tests/test_output_targets.py +++ b/flashdreams/tests/test_output_targets.py @@ -71,9 +71,10 @@ def test_lingbot_webrtc_target_translates_runner_config() -> None: ), ) - assert spec.module == "lingbot.webrtc.server" + assert spec.module == "lingbot.demo.cli" assert spec.argv == ( - "--config_name", + "webrtc", + "--preset-id", "lingbot-world-fast", "--device", "cuda:1", @@ -83,14 +84,14 @@ def test_lingbot_webrtc_target_translates_runner_config() -> None: "480", "--video-width", "832", - "--no_compile", + "--no-compile", "--example-idx", "3", "--host", "127.0.0.1", "--port", "9010", - "--prefer_sw_encoder", + "--prefer-sw-encoder", ) diff --git a/integrations/lingbot/lingbot/demo/__init__.py b/integrations/lingbot/lingbot/demo/__init__.py new file mode 100644 index 000000000..a3da57837 --- /dev/null +++ b/integrations/lingbot/lingbot/demo/__init__.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Experimental Lingbot demo adapter built on ``flashdreams.runtime.demo``.""" + +from lingbot.demo.adapter import LingbotDemoAdapter +from lingbot.demo.spec import ( + DEFAULT_LINGBOT_PRESET, + LINGBOT_MODEL_ID, + LingbotReplayInputs, + LingbotWebRTCScenario, +) + +__all__ = [ + "DEFAULT_LINGBOT_PRESET", + "LINGBOT_MODEL_ID", + "LingbotDemoAdapter", + "LingbotReplayInputs", + "LingbotWebRTCScenario", +] diff --git a/integrations/lingbot/lingbot/demo/adapter.py b/integrations/lingbot/lingbot/demo/adapter.py new file mode 100644 index 000000000..354ea56b4 --- /dev/null +++ b/integrations/lingbot/lingbot/demo/adapter.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot adapter for the shared demo API.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Any + +from flashdreams.runtime import ( + InferenceConfig, + InputCanonicalizer, + UserInputCapability, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + PreparedScenario, + WebRTCOutputSpec, +) +from flashdreams.runtime.interfaces import InferenceRuntime +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): + """Model-owned Lingbot adapter consumed by shared demo launchers.""" + + 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") + + def supported_output_modes(self) -> tuple[str, ...]: + return ("mp4", "webrtc") + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + if spec.input_mode != "replay": + raise ValueError( + "Lingbot prepare_scenario currently supports only " + f"input_mode='replay', got {spec.input_mode!r}." + ) + if not isinstance(spec.output, Mp4OutputSpec): + raise ValueError("Lingbot replay demo currently requires MP4 output.") + + replay_inputs = resolve_replay_inputs( + spec.scenario, + default_prompt=self.default_replay_prompt(spec.config), + ) + text_event_prompts = resolve_text_event_prompts(spec.scenario) + user_inputs = resolve_user_input_events(spec.scenario) + if _camera_source(spec.scenario) == "events": + # Live control still needs the scenario's calibration, so the trace + # is loaded for its intrinsics and world scale and then discarded + # as a trajectory source. + trace = self.create_input_mapping(replay_inputs).camera_trace + mapping = self.create_live_input_mapping( + fps=replay_inputs.fps, + base_intrinsics=trace.intrinsics[0], + # 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. + world_scale=trace.world_scale or 1.0, + prompt=replay_inputs.prompt, + text_event_prompts=text_event_prompts, + ) + else: + mapping = self.create_input_mapping( + replay_inputs, + text_event_prompts=text_event_prompts, + ) + return PreparedScenario( + initial_inputs=inference_input_from_replay_inputs(replay_inputs), + user_inputs=user_inputs, + source_schema=_source_schema(user_inputs), + canonicalizer=_canonicalizer(text_event_prompts), + mapping=mapping, + metadata={ + "model_id": self.model_id, + "preset_id": self.preset_id(spec.config), + }, + ) + + 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): + return str(scenario.get("camera_source", "trace")) + return "trace" + + +_KEY_EVENT_TYPES = frozenset({"key_down", "key_up"}) + +_KEYBOARD_CAPABILITIES = ( + UserInputCapability(event_type="key_down", payload_fields=frozenset({"key"})), + UserInputCapability(event_type="key_up", payload_fields=frozenset({"key"})), +) + +_TEXT_EVENT_CAPABILITY = UserInputCapability( + event_type="text_event", + payload_fields=frozenset({"event_id"}), +) + + +def _source_schema(user_inputs: UserInputs) -> UserInputSchema: + """Declare what this scenario's event source can provide. + + Capabilities describe the source, not the particular trace. A keyboard + source is declared to provide both key edges even if one recording happens + to contain no ``key_up`` -- a key held for the whole run is a normal trace. + Declaring only the observed types would fail the keyboard converter's + consumed set, and ``converters_for`` would silently drop it, leaving the run + with no camera control. + """ + observed = {event.event_type for event in user_inputs.events} + capabilities: list[UserInputCapability] = [] + if observed & _KEY_EVENT_TYPES: + capabilities.extend(_KEYBOARD_CAPABILITIES) + if "text_event" in observed: + capabilities.append(_TEXT_EVENT_CAPABILITY) + for event_type in sorted(observed - _KEY_EVENT_TYPES - {"text_event"}): + payload_fields: frozenset[str] = frozenset() + for event in user_inputs.events: + if event.event_type == event_type: + payload_fields = frozenset(event.payload) + break + capabilities.append( + UserInputCapability( + event_type=event_type, + payload_fields=payload_fields, + ) + ) + return UserInputSchema( + capabilities=tuple(capabilities), + description=( + "Lingbot replay event trace" + if capabilities + else "fixed Lingbot replay input" + ), + ) + + +def _canonicalizer(text_event_prompts: Mapping[str, str] | None) -> InputCanonicalizer: + converters: list[Any] = [KeyboardToCameraCommand()] + if text_event_prompts: + converters.append(TextEventSelection()) + return InputCanonicalizer(converters) + + +__all__ = [ + "LingbotDemoAdapter", + "ReplayRuntimeFactory", + "WebRTCRuntimeFactory", +] diff --git a/integrations/lingbot/lingbot/demo/cli.py b/integrations/lingbot/lingbot/demo/cli.py new file mode 100644 index 000000000..9b08469f6 --- /dev/null +++ b/integrations/lingbot/lingbot/demo/cli.py @@ -0,0 +1,233 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CLI for the experimental shared Lingbot demo path.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +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 lingbot.example_data import ( + EXAMPLE_DATA_AVAILABLE_IDXS, + ensure_example_data_downloaded, +) +from lingbot.runtime import ( + FIELD_CAMERA_INTRINSICS_PATH, + FIELD_CAMERA_POSES_PATH, + FIELD_FIRST_FRAME_PATH, + FIELD_FPS, + FIELD_PIXEL_HEIGHT, + FIELD_PIXEL_WIDTH, + FIELD_PROMPT, + FIELD_TOTAL_BLOCKS, +) + +from .adapter import LingbotDemoAdapter +from .spec import ( + DEFAULT_FPS, + DEFAULT_LINGBOT_PRESET, + DEFAULT_PIXEL_HEIGHT, + DEFAULT_PIXEL_WIDTH, + LINGBOT_MODEL_ID, + LingbotWebRTCScenario, +) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Experimental Lingbot demo using flashdreams.runtime.demo." + ) + subparsers = parser.add_subparsers(dest="command", required=True) + + replay = subparsers.add_parser("replay", help="Run an MP4 replay demo.") + replay.add_argument("--preset-id", "--config-name", default=DEFAULT_LINGBOT_PRESET) + replay.add_argument("--device", default="cuda") + replay.add_argument("--prompt", default=None) + replay.add_argument("--prompt-path", type=Path, default=None) + replay.add_argument("--image-path", type=Path, default=None) + replay.add_argument("--pose-path", type=Path, default=None) + replay.add_argument( + "--intrinsic-path", + "--intrinsics-path", + type=Path, + default=None, + ) + replay.add_argument( + "--example-data", + action=argparse.BooleanOptionalAction, + default=None, + help=( + "Use the bundled Lingbot example when asset paths are omitted " + "(default: auto)." + ), + ) + replay.add_argument( + "--example-idx", + "--example_idx", + type=int, + default=0, + choices=EXAMPLE_DATA_AVAILABLE_IDXS, + ) + replay.add_argument("--total-blocks", type=int, default=20) + replay.add_argument("--pixel-height", type=int, default=DEFAULT_PIXEL_HEIGHT) + replay.add_argument("--pixel-width", type=int, default=DEFAULT_PIXEL_WIDTH) + replay.add_argument("--fps", type=int, default=DEFAULT_FPS) + replay.add_argument("--output", type=Path, required=True) + + webrtc = subparsers.add_parser("webrtc", help="Serve a WebRTC driving demo.") + webrtc.add_argument("--preset-id", "--config-name", default=DEFAULT_LINGBOT_PRESET) + webrtc.add_argument("--host", default="0.0.0.0") + webrtc.add_argument("--port", type=int, default=8080) + webrtc.add_argument("--device", default="cuda:0") + webrtc.add_argument("--seed", type=int, default=42) + webrtc.add_argument( + "--compile", + action=argparse.BooleanOptionalAction, + default=True, + help="Enable or disable torch.compile for the Lingbot transformer.", + ) + webrtc.add_argument("--fps", type=int, default=DEFAULT_FPS) + webrtc.add_argument("--video-height", type=int, default=DEFAULT_PIXEL_HEIGHT) + webrtc.add_argument("--video-width", type=int, default=DEFAULT_PIXEL_WIDTH) + webrtc.add_argument("--warmup-chunks", type=int, default=10) + webrtc.add_argument("--warmup-timeout-s", type=float, default=600.0) + webrtc.add_argument("--client-liveness-timeout-s", type=float, default=30.0) + webrtc.add_argument("--prefer-sw-encoder", action="store_true") + webrtc.add_argument( + "--example-idx", + "--example_idx", + type=int, + default=0, + choices=EXAMPLE_DATA_AVAILABLE_IDXS, + ) + 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, + ) + ensure_example_data_downloaded( + is_rank_zero=(context.world_rank == 0), + example_idx=args.example_idx, + ) + serve_flashdreams_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}") + + +def _replay_spec(args: argparse.Namespace) -> DemoSpec: + scenario: dict[str, object] = { + "example_data": args.example_data, + "example_idx": args.example_idx, + FIELD_TOTAL_BLOCKS: args.total_blocks, + FIELD_PIXEL_HEIGHT: args.pixel_height, + FIELD_PIXEL_WIDTH: args.pixel_width, + FIELD_FPS: args.fps, + } + if args.prompt: + scenario[FIELD_PROMPT] = args.prompt + if args.prompt_path is not None: + scenario["prompt_path"] = args.prompt_path + if args.image_path is not None: + scenario[FIELD_FIRST_FRAME_PATH] = args.image_path + if args.pose_path is not None: + scenario[FIELD_CAMERA_POSES_PATH] = args.pose_path + if args.intrinsic_path is not None: + scenario[FIELD_CAMERA_INTRINSICS_PATH] = args.intrinsic_path + + return DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=args.preset_id, + input_mode="replay", + scenario=scenario, + output=Mp4OutputSpec( + path=args.output, + fps=args.fps, + output_layout="tchw", + ), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=args.preset_id, + device=args.device, + ), + ) + + +def _webrtc_spec( + args: argparse.Namespace, + *, + device: str, + context_parallel_size: int = 1, +) -> DemoSpec: + return DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=args.preset_id, + input_mode="keyboard-driving", + scenario=LingbotWebRTCScenario( + example_idx=args.example_idx, + prefer_sw_encoder=args.prefer_sw_encoder, + ), + output=WebRTCOutputSpec( + host=args.host, + port=args.port, + fps=args.fps, + video_width=args.video_width, + video_height=args.video_height, + warmup_chunks=args.warmup_chunks, + warmup_timeout_s=args.warmup_timeout_s, + client_liveness_timeout_s=args.client_liveness_timeout_s, + preload_name="Lingbot", + ), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=args.preset_id, + device=device, + compile=args.compile, + runtime_options={ + "seed": args.seed, + "context_parallel_size": context_parallel_size, + "example_idx": args.example_idx, + }, + ), + ) + + +if __name__ == "__main__": + main() diff --git a/integrations/lingbot/lingbot/demo/replay.py b/integrations/lingbot/lingbot/demo/replay.py new file mode 100644 index 000000000..668ba63aa --- /dev/null +++ b/integrations/lingbot/lingbot/demo/replay.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot replay runtime re-export for shared demo entry points.""" + +from __future__ import annotations + +from lingbot.runtime import ( + LingbotReplayRuntime, + LingbotReplayRuntimeOptions, + LingbotReplaySession, + PipelineFactory, +) + +__all__ = [ + "LingbotReplayRuntime", + "LingbotReplayRuntimeOptions", + "LingbotReplaySession", + "PipelineFactory", +] diff --git a/integrations/lingbot/lingbot/demo/spec.py b/integrations/lingbot/lingbot/demo/spec.py new file mode 100644 index 000000000..ef153fdb0 --- /dev/null +++ b/integrations/lingbot/lingbot/demo/spec.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot demo-specific input shapes.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from flashdreams.runtime import UserInputEvent, UserInputs +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, + example_asset_urls, + example_data_dirname, +) +from lingbot.runtime import ( + DEFAULT_FPS, + DEFAULT_LINGBOT_PRESET, + DEFAULT_PIXEL_HEIGHT, + DEFAULT_PIXEL_WIDTH, + LINGBOT_MODEL_ID, + LingbotReplayInputs, + replay_inputs_from_mapping, +) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class LingbotWebRTCScenario: + """Example-data and serving options for the shared WebRTC demo path.""" + + example_idx: int = 0 + prefer_sw_encoder: bool = False + + def __post_init__(self) -> None: + if self.example_idx not in EXAMPLE_DATA_AVAILABLE_IDXS: + raise ValueError( + "LingbotWebRTCScenario.example_idx must be one of " + f"{EXAMPLE_DATA_AVAILABLE_IDXS}." + ) + + +def resolve_replay_inputs( + value: Any, + *, + default_prompt: str = "", + is_rank_zero: bool = True, +) -> LingbotReplayInputs: + """Normalize a user/demo scenario into direct Lingbot runtime inputs.""" + return replay_inputs_from_mapping( + value, + default_prompt=default_prompt, + is_rank_zero=is_rank_zero, + ) + + +def resolve_text_event_prompts(value: Any) -> dict[str, str]: + """Return the scenario's text-event catalog as ``{event_id: prompt}``.""" + if not isinstance(value, Mapping): + return {} + catalog = value.get("text_events") + if not catalog: + return {} + if isinstance(catalog, Mapping): + return {str(key): str(prompt) for key, prompt in catalog.items()} + prompts: dict[str, str] = {} + for entry in catalog: + event_id = getattr(entry, "event_id", None) + prompt = getattr(entry, "prompt", None) + if event_id is None and isinstance(entry, Mapping): + event_id = entry.get("event_id") + prompt = entry.get("prompt") + if event_id is None: + raise ValueError("Lingbot text events require an 'event_id'.") + prompts[str(event_id)] = "" if prompt is None else str(prompt) + return prompts + + +def resolve_user_input_events(value: Any) -> UserInputs: + """Normalize a scenario's recorded event trace into :class:`UserInputs`. + + Each record is ``{"t": seconds, "type": event_type, ...payload}``, which + maps one-to-one onto ``UserInputEvent``. Events are sorted by timestamp + because ``UserInputs`` requires non-decreasing order. + """ + if not isinstance(value, Mapping): + return UserInputs() + records = value.get("events") + if not records: + return UserInputs() + + events: list[UserInputEvent] = [] + for record in records: + if isinstance(record, UserInputEvent): + events.append(record) + continue + if not isinstance(record, Mapping): + raise TypeError( + "Lingbot scenario events must be UserInputEvent objects or " + "mappings." + ) + payload = { + key: item + for key, item in record.items() + if key not in {"t", "timestamp_s", "type", "event_type", "source"} + } + timestamp_s = record.get("t", record.get("timestamp_s")) + 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')." + ) + events.append( + UserInputEvent( + timestamp_s=float(timestamp_s), + event_type=str(event_type), + payload=payload, + source=record.get("source"), + ) + ) + events.sort(key=lambda event: event.timestamp_s) + return UserInputs(events=tuple(events)) + + +def resolve_webrtc_scenario(value: Any) -> LingbotWebRTCScenario: + """Normalize a user/demo scenario into a WebRTC scenario.""" + if value is None: + return LingbotWebRTCScenario() + if isinstance(value, LingbotWebRTCScenario): + return value + if not isinstance(value, Mapping): + raise TypeError( + "Lingbot WebRTC scenario must be a LingbotWebRTCScenario, " + "a mapping, or None." + ) + return LingbotWebRTCScenario( + example_idx=int(value.get("example_idx", 0)), + prefer_sw_encoder=bool(value.get("prefer_sw_encoder", False)), + ) + + +__all__ = [ + "DEFAULT_FPS", + "DEFAULT_LINGBOT_PRESET", + "DEFAULT_PIXEL_HEIGHT", + "DEFAULT_PIXEL_WIDTH", + "EXAMPLE_DATA_AVAILABLE_IDXS", + "EXAMPLE_DATA_BASE_URL", + "EXAMPLE_DATA_DIR_LOCAL", + "EXAMPLE_DATA_FILENAMES", + "EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS", + "LINGBOT_MODEL_ID", + "LingbotReplayInputs", + "LingbotWebRTCScenario", + "example_asset_urls", + "example_data_dirname", + "resolve_replay_inputs", + "resolve_text_event_prompts", + "resolve_user_input_events", + "resolve_webrtc_scenario", +] diff --git a/integrations/lingbot/lingbot/demo/webrtc.py b/integrations/lingbot/lingbot/demo/webrtc.py new file mode 100644 index 000000000..966a79fb7 --- /dev/null +++ b/integrations/lingbot/lingbot/demo/webrtc.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot WebRTC hooks for the shared demo API.""" + +from __future__ import annotations + +from typing import Any + +from aiohttp import web + +from flashdreams.runtime.demo import DemoSpec +from lingbot.webrtc.server import create_app +from lingbot.webrtc.session import ( + LingbotInferenceRuntime, + LingbotRuntimeConfig, + LingbotWebRTCSessionManager, +) + + +class LingbotDemoWebRTCSessionManager(LingbotWebRTCSessionManager): + """Shared demo session manager using Lingbot's existing WebRTC semantics.""" + + 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, + ) + + +def create_lingbot_webrtc_app( + *, + spec: DemoSpec, + session_manager: Any, + request_session_url: str, +) -> web.Application: + """Create the packaged Lingbot browser app through existing serving glue.""" + del spec + return create_app( + session_manager=session_manager, + request_session_url=request_session_url, + ) + + +__all__ = [ + "LingbotDemoWebRTCSessionManager", + "create_lingbot_webrtc_app", +] diff --git a/integrations/lingbot/lingbot/example_data.py b/integrations/lingbot/lingbot/example_data.py new file mode 100644 index 000000000..af33628a5 --- /dev/null +++ b/integrations/lingbot/lingbot/example_data.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bundled LingBot-World example-data helpers.""" + +from __future__ import annotations + +from pathlib import Path + +import torch + +from flashdreams.core.io.disk import default_flashdreams_cache_dir +from flashdreams.core.io.download import download_to_cache + +EXAMPLE_DATA_BASE_URL = ( + "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples" +) +"""HTTP base URL for the canonical examples shared by all LingBot versions.""" + +EXAMPLE_DATA_DIR_LOCAL = default_flashdreams_cache_dir() / "example_data/lingbot_world" +"""Local cache root where downloaded example folders are stored.""" + +EXAMPLE_DATA_FILENAMES = ( + "image.jpg", + "poses.npy", + "intrinsics.npy", + "prompt.txt", +) +"""Example assets downloaded when each file is available upstream.""" + +EXAMPLE_DATA_AVAILABLE_IDXS = (0, 1, 2, 3, 4, 5) +"""Supported upstream example indices currently hosted under ``examples/``.""" + +EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS = (0, 1, 2, 5) +"""Example indices that provide their own upstream ``prompt.txt`` file.""" + + +def example_data_dirname(example_idx: int) -> str: + """Format ``example_idx`` into the upstream folder naming convention.""" + assert example_idx in EXAMPLE_DATA_AVAILABLE_IDXS, ( + f"--example_idx must be one of {EXAMPLE_DATA_AVAILABLE_IDXS}." + ) + return f"{example_idx:02d}" + + +def example_asset_urls(example_idx: int) -> dict[str, str]: + """Return canonical upstream URLs for a Lingbot example.""" + dirname = example_data_dirname(example_idx) + return { + "image": f"{EXAMPLE_DATA_BASE_URL}/{dirname}/image.jpg", + "intrinsics": f"{EXAMPLE_DATA_BASE_URL}/{dirname}/intrinsics.npy", + "poses": f"{EXAMPLE_DATA_BASE_URL}/{dirname}/poses.npy", + } + + +def ensure_example_data_downloaded(*, is_rank_zero: bool, example_idx: int) -> Path: + """Download bundled GitHub example files on rank 0; barrier other ranks.""" + example_dirname = example_data_dirname(example_idx) + cache_dir = EXAMPLE_DATA_DIR_LOCAL / example_dirname + if is_rank_zero: + for filename in EXAMPLE_DATA_FILENAMES: + if ( + filename == "prompt.txt" + and example_idx not in EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS + ): + continue + download_to_cache( + f"{EXAMPLE_DATA_BASE_URL}/{example_dirname}/{filename}", + cache_dir=cache_dir, + filename=filename, + ) + if torch.distributed.is_initialized(): + torch.distributed.barrier() + return cache_dir + + +__all__ = [ + "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_asset_urls", + "example_data_dirname", +] diff --git a/integrations/lingbot/lingbot/input_mapping.py b/integrations/lingbot/lingbot/input_mapping.py new file mode 100644 index 000000000..9a05c2178 --- /dev/null +++ b/integrations/lingbot/lingbot/input_mapping.py @@ -0,0 +1,670 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot user-event canonicalization and canonical-to-model input mapping. + +Lingbot's two live controls are a free camera driven from the keyboard and a +catalog of server-owned text events. This module carries both across the +``UserInputs -> CanonicalInputs -> InferenceInput`` boundary: + +- :data:`CAMERA_COMMAND` and :class:`KeyboardToCameraCommand` turn raw key + edges into device-independent camera intent; +- :data:`TEXT_EVENT` and :class:`TextEventSelection` track which text event is + active; +- :class:`LingbotInputMapping` turns that canonical intent into the per-step + camera trajectory the session consumes, and requests a session-global prompt + update when the active text event changes. + +The modalities live here rather than in ``flashdreams.runtime.canonical`` +because Lingbot is currently their only consumer. Both are plain +``CanonicalModality`` values, so lifting them into the shared canonical layer +later is a move plus an export, with no change to this mapping. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from flashdreams.runtime.canonical import DeviceConverterSchema +from flashdreams.runtime.inputs import ( + CanonicalInputs, + CanonicalInputSchema, + CanonicalModality, + InferenceInput, + InferenceInputSchema, + InputField, + TimeWindow, + UserInputCapability, + UserInputs, +) +from flashdreams.runtime.mapping import InputMappingSchema +from flashdreams.runtime.types import StepRequest +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" +FIELD_PROMPT = "prompt" +FIELD_WORLD_SCALE = "world_scale" +FIELD_TOTAL_CAMERA_FRAMES = "total_camera_frames" + +_PASSTHROUGH_GLOBAL_FIELDS: tuple[InputField, ...] = ( + InputField(name="first_frame_path", input_modality="image/path"), + InputField(name="total_blocks", input_modality="count"), + InputField(name="pixel_height", input_modality="pixel-height"), + InputField(name="pixel_width", input_modality="pixel-width"), + InputField(name="fps", input_modality="fps"), +) +"""App-owned session inputs this mapping forwards without interpreting them.""" + +_CLEAR_STATES = frozenset({"clear", "release", "off", "none"}) +_TRIGGER_STATES = frozenset({"trigger", "hold", "on"}) + +_AXES: tuple[str, ...] = ("move_forward", "move_right", "yaw", "pitch") + +_AXIS_KEYS: Mapping[str, tuple[str, str]] = { + # Axis -> (positive key, negative key) in CameraPoseIntegrator's vocabulary. + # Both directions of the keyboard/axis conversion are derived from this one + # table so a rebind cannot make them disagree. + "move_forward": ("w", "s"), + "move_right": ("e", "q"), + "yaw": ("a", "d"), + "pitch": ("i", "k"), +} + +_KEY_ALIASES: Mapping[str, str] = {"j": "a", "l": "d"} +"""Alternate yaw keys accepted by ``KeyboardState``, folded onto ``a``/``d``.""" + + +CAMERA_COMMAND = CanonicalModality( + name="camera_command", + payload_fields=frozenset({*_AXES, "segments"}), + description=( + "Free-camera intent. move_forward, move_right, yaw, and pitch are in " + "[-1, 1] and hold the level state at the end of the window. segments " + "carries the piecewise-constant timeline inside the window as " + "((start_s, end_s, axes), ...), so a consumer can integrate sub-window " + "timing instead of quantizing control to the chunk boundary." + ), +) + +TEXT_EVENT = CanonicalModality( + name="text_event", + payload_fields=frozenset({"event_id"}), + description=( + "Identifier of the active server-owned text event, or None when no " + "event is active. Level-triggered: the value is held until cleared." + ), +) + + +def _axes_from_keys(pressed: Iterable[str]) -> dict[str, float]: + """Return camera axis values for a resolved set of pressed keys.""" + keys = {_KEY_ALIASES.get(key, key) for key in pressed} + axes: dict[str, float] = {} + for axis, (positive, negative) in _AXIS_KEYS.items(): + value = 0.0 + if positive in keys: + value += 1.0 + if negative in keys: + value -= 1.0 + axes[axis] = value + return axes + + +def _keys_from_axes(axes: Mapping[str, float]) -> frozenset[str]: + """Return the integrator key set equivalent to ``axes``. + + Pose integration stays the single implementation in + :class:`CameraPoseIntegrator`, which is expressed over key sets. Converting + back here keeps live Lingbot trajectories identical to the WebRTC path + instead of forking the integration math. + """ + keys: set[str] = set() + for axis, (positive, negative) in _AXIS_KEYS.items(): + value = float(axes.get(axis, 0.0)) + if value > 0: + keys.add(positive) + elif value < 0: + keys.add(negative) + return frozenset(keys) + + +class KeyboardToCameraCommand: + """Convert keyboard edges into :data:`CAMERA_COMMAND` level state.""" + + def __init__( + self, + *, + name: str = "keyboard-to-camera-command", + supported_keys: frozenset[str] = DEFAULT_SUPPORTED_KEYS, + priority: int = 0, + ) -> None: + self._supported_keys = supported_keys + self._state = KeyboardState(supported_keys=supported_keys) + self._schema = DeviceConverterSchema( + name=name, + produces=CAMERA_COMMAND, + device_kind="keyboard", + priority=priority, + consumes=( + UserInputCapability( + event_type="key_down", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_type="key_up", + payload_fields=frozenset({"key"}), + ), + ), + ) + + @property + def schema(self) -> DeviceConverterSchema: + return self._schema + + def reset(self) -> None: + self._state = KeyboardState(supported_keys=self._supported_keys) + + def convert( + self, + user_inputs: UserInputs, + window: TimeWindow, + ) -> Mapping[str, Any] | None: + segments: list[tuple[float, float, dict[str, float]]] = [] + segment_start = window.start_s + axes = _axes_from_keys(self._state.resolved_effective_keys()) + + for event in user_inputs.events: + if event.event_type not in {"key_down", "key_up"}: + continue + key = event.payload.get("key") + if not isinstance(key, str): + continue + edge_t = min(max(float(event.timestamp_s), window.start_s), window.end_s) + if edge_t > segment_start: + segments.append((segment_start, edge_t, axes)) + segment_start = edge_t + self._state.apply_event( + event="keydown" if event.event_type == "key_down" else "keyup", + key=key, + ) + axes = _axes_from_keys(self._state.resolved_effective_keys()) + + if window.end_s > segment_start or not segments: + segments.append((segment_start, window.end_s, axes)) + + return CAMERA_COMMAND.value({**axes, "segments": tuple(segments)}) + + +class TextEventSelection: + """Track the active :data:`TEXT_EVENT` id across windows.""" + + def __init__( + self, + *, + name: str = "text-event-selection", + priority: int = 0, + ) -> None: + self._active_event_id: str | None = None + self._schema = DeviceConverterSchema( + name=name, + produces=TEXT_EVENT, + device_kind="text-event", + priority=priority, + consumes=( + UserInputCapability( + event_type="text_event", + payload_fields=frozenset({"event_id"}), + ), + ), + ) + + @property + def schema(self) -> DeviceConverterSchema: + return self._schema + + def reset(self) -> None: + self._active_event_id = None + + def convert( + self, + user_inputs: UserInputs, + window: TimeWindow, + ) -> Mapping[str, Any] | None: + del window + for event in user_inputs.events: + if event.event_type != "text_event": + continue + event_id = event.payload.get("event_id") + state = str(event.payload.get("state", "trigger")).strip().lower() + if state and state not in _CLEAR_STATES and state not in _TRIGGER_STATES: + raise ValueError( + f"Unsupported text event state {state!r}. Supported states: " + f"{sorted(_CLEAR_STATES | _TRIGGER_STATES)}." + ) + if event_id is None or state in _CLEAR_STATES: + self._active_event_id = None + continue + self._active_event_id = str(event_id) + return TEXT_EVENT.value({"event_id": self._active_event_id}) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class LingbotCameraTrace: + """Fixed camera trajectory resolved from a replay scenario. + + Tensors are CPU float32; the session owns device placement. + """ + + __hash__ = None + + poses: torch.Tensor + """Camera-to-world poses, shape ``[T, 4, 4]``.""" + + intrinsics: torch.Tensor + """Per-frame intrinsics, shape ``[T, 4]``, already rescaled to output size.""" + + world_scale: float + + def __post_init__(self) -> None: + if self.poses.ndim != 3 or self.poses.shape[1:] != (4, 4): + raise ValueError( + f"LingbotCameraTrace.poses must be [T, 4, 4], got " + f"{tuple(self.poses.shape)}." + ) + if self.intrinsics.ndim != 2 or self.intrinsics.shape[1] != 4: + raise ValueError( + f"LingbotCameraTrace.intrinsics must be [T, 4], got " + f"{tuple(self.intrinsics.shape)}." + ) + if self.world_scale < 0: + # Zero is legal: preprocess_example_poses derives the scale from + # pose spread, so a stationary trace yields 0. That reached the + # model before this mapping existed, so it still does. + raise ValueError("LingbotCameraTrace.world_scale must be >= 0.") + + @property + def frame_count(self) -> int: + return int(self.poses.shape[0]) + + +def load_camera_trace( + *, + camera_poses_path: str | Path, + camera_intrinsics_path: str | Path, + pixel_height: int, + pixel_width: int, + intrinsics_reference_height: int, + intrinsics_reference_width: int, + world_scale: float | None = None, +) -> LingbotCameraTrace: + """Load and preprocess a fixed Lingbot camera trajectory from ``.npy`` files.""" + from lingbot.encoder.utils import ( # noqa: PLC0415 + get_Ks_transformed, + preprocess_example_poses, + ) + + intrinsics = torch.from_numpy( + np.asarray(np.load(camera_intrinsics_path), dtype=np.float32) + ) + intrinsics = get_Ks_transformed( + intrinsics, + height_org=intrinsics_reference_height, + width_org=intrinsics_reference_width, + height_resize=pixel_height, + width_resize=pixel_width, + height_final=pixel_height, + width_final=pixel_width, + ) + poses, inferred_world_scale = preprocess_example_poses( + np.asarray(np.load(camera_poses_path)) + ) + 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 + ), + ) + + +class LingbotInputMapping: + """Build Lingbot per-step camera inputs from canonical user input. + + Two trajectory sources are supported through one mapping object, because + ``run_inference_session`` takes a single mapping: + + - a fixed :class:`LingbotCameraTrace`, sliced per step, which consumes no + canonical modality and keeps MP4/benchmark runs deterministic; + - live :data:`CAMERA_COMMAND` intent integrated into a trajectory, for + event-driven runs. + + Text events are mapped to a session-global prompt update rather than a + per-step field: swapping the rollout's text context is session-global model + state, so it travels in the ``global_conditioning`` slot of the step + payload. Whether the model can apply that update is session-owned. + """ + + def __init__( + self, + *, + fps: int, + trace: LingbotCameraTrace | None = None, + base_intrinsics: torch.Tensor | Sequence[float] | None = None, + world_scale: float | None = None, + text_event_prompts: Mapping[str, str] | None = None, + integrator: CameraPoseIntegrator | None = None, + ) -> None: + if fps <= 0: + raise ValueError("LingbotInputMapping.fps must be > 0.") + if trace is None and base_intrinsics is None: + raise ValueError( + "LingbotInputMapping requires either a fixed camera trace or " + "base_intrinsics for live camera control." + ) + self._fps = int(fps) + self._trace = trace + self._text_event_prompts = dict(text_event_prompts or {}) + self._applied_event_id: str | None = None + self._base_prompt: str | None = None + + if trace is None: + intrinsics = torch.as_tensor(base_intrinsics, dtype=torch.float32).reshape( + 4 + ) + if world_scale is None or world_scale <= 0: + raise ValueError( + "Live Lingbot camera control requires a positive world_scale." + ) + self._base_intrinsics = intrinsics + self._world_scale = float(world_scale) + self._integrator = integrator or CameraPoseIntegrator() + else: + self._base_intrinsics = None + self._world_scale = trace.world_scale + self._integrator = None + + consumes: list[CanonicalModality] = [] + if trace is None: + consumes.append(CAMERA_COMMAND) + if self._text_event_prompts: + consumes.append(TEXT_EVENT) + self._mapping_schema = InputMappingSchema( + name="lingbot-input-mapping", + consumes=tuple(consumes), + produces_global_conditioning=( + # map_global_conditioning_inputs returns the app-owned session + # payload augmented with the fields below, so the pass-through + # fields are part of what this mapping produces. Declaring them + # keeps undeclared_inference_inputs() quiet and lets the + # compatibility check see that required session inputs are + # reachable; omitting them makes the check reject every run. + *_PASSTHROUGH_GLOBAL_FIELDS, + InputField( + name=FIELD_WORLD_SCALE, + required=False, + input_modality="scale", + frequency_consumed="once", + ), + InputField( + name=FIELD_TOTAL_CAMERA_FRAMES, + required=False, + input_modality="count", + frequency_consumed="once", + ), + InputField( + name=FIELD_PROMPT, + required=False, + input_modality="text", + frequency_consumed="once", + description="Text-event prompt update for an active rollout.", + ), + ), + produces_step=( + InputField( + name=FIELD_CAMERA_TRAJECTORY, + input_modality="c2w_sequence", + frequency_consumed="per_step", + metadata={"shape": "[T,4,4]", "frame": "camera_to_world"}, + ), + InputField( + name=FIELD_CAMERA_INTRINSICS, + input_modality="intrinsics_vec4_sequence", + frequency_consumed="per_step", + metadata={"shape": "[T,4]"}, + ), + ), + ) + + @property + def mapping_schema(self) -> InputMappingSchema: + return self._mapping_schema + + @property + def camera_trace(self) -> LingbotCameraTrace: + """Return the fixed trace, for callers reusing its calibration.""" + if self._trace is None: + raise ValueError("This Lingbot mapping has no fixed camera trace.") + return self._trace + + @property + def canonical_input_schema(self) -> CanonicalInputSchema: + """Return the modalities this mapping consumes, for adapter reporting.""" + return CanonicalInputSchema( + modalities=self._mapping_schema.consumes, + description="Lingbot live camera and text-event control.", + ) + + def validate( + self, + *, + canonical_schema: CanonicalInputSchema | None = None, + inference_input_schema: InferenceInputSchema | None = None, + ) -> None: + if canonical_schema is not None: + for modality in self._mapping_schema.consumes: + if not canonical_schema.supports(modality): + raise ValueError( + f"Lingbot input mapping requires canonical modality " + f"{modality.name!r}, which the selected input source " + f"cannot supply." + ) + if inference_input_schema is not None: + for name in (FIELD_CAMERA_TRAJECTORY, FIELD_CAMERA_INTRINSICS): + if inference_input_schema.field_for(name=name, phase="step") is None: + raise ValueError( + f"Lingbot input mapping produces step input {name!r}, " + f"which this model does not declare." + ) + + def map_global_conditioning_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + ) -> InferenceInput: + del canonical_inputs + payload = dict(inference_input.global_conditioning) + payload[FIELD_WORLD_SCALE] = self._world_scale + if self._trace is not None: + payload[FIELD_TOTAL_CAMERA_FRAMES] = self._trace.frame_count + return InferenceInput( + global_conditioning=payload, + step=inference_input.step, + metadata=inference_input.metadata, + ) + + def map_step_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + request: StepRequest, + ) -> InferenceInput: + num_frames = _required_int(request.metadata, "num_frames") + frame_start = _required_int(request.metadata, "frame_start") + + if self._trace is not None: + poses, intrinsics = self._slice_trace( + frame_start=frame_start, + num_frames=num_frames, + ) + else: + poses, intrinsics = self._integrate( + canonical_inputs=canonical_inputs, + request=request, + frame_start=frame_start, + num_frames=num_frames, + ) + + step = dict(inference_input.step) + step[FIELD_CAMERA_TRAJECTORY] = poses + step[FIELD_CAMERA_INTRINSICS] = intrinsics + return InferenceInput( + global_conditioning=self._text_event_update(canonical_inputs), + step=step, + metadata=inference_input.metadata, + ) + + def _slice_trace( + self, + *, + frame_start: int, + num_frames: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert self._trace is not None + frame_end = frame_start + num_frames + if frame_end > self._trace.frame_count: + raise ValueError( + f"Lingbot camera trace has {self._trace.frame_count} frames, but " + f"step needs frames [{frame_start}, {frame_end})." + ) + return ( + self._trace.poses[frame_start:frame_end], + self._trace.intrinsics[frame_start:frame_end], + ) + + def _integrate( + self, + *, + canonical_inputs: CanonicalInputs, + request: StepRequest, + frame_start: int, + num_frames: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert self._integrator is not None + assert self._base_intrinsics is not None + command = canonical_inputs.values.get(CAMERA_COMMAND.name) + if command is None: + raise ValueError( + "Lingbot live camera control requires a 'camera_command' " + "canonical value for every step; the selected input source " + "produced none." + ) + + 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 + 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 + # accumulation can leave the last one a hair past the window end. + frame_times[-1] = min(frame_times[-1], end_s) + + poses = self._integrator.integrate_chunk( + segments=segments, + frame_times=frame_times, + ) + poses_t = torch.from_numpy(np.ascontiguousarray(poses)).to(torch.float32) + poses_t = poses_t.reshape(num_frames, 4, 4) + intrinsics_t = self._base_intrinsics.reshape(1, 4).repeat(num_frames, 1) + return poses_t, intrinsics_t + + def _text_event_update( + self, + canonical_inputs: CanonicalInputs, + ) -> Mapping[str, Any]: + if not self._text_event_prompts: + return {} + value = canonical_inputs.values.get(TEXT_EVENT.name) + if value is None: + return {} + event_id = value.get("event_id") + if event_id == self._applied_event_id: + return {} + if event_id is not None and event_id not in self._text_event_prompts: + supported = ", ".join(sorted(self._text_event_prompts)) + raise ValueError( + f"Unknown Lingbot text event_id={event_id!r}. Supported: {supported}" + ) + self._applied_event_id = event_id + prompt = ( + self._base_prompt + if event_id is None + else self._text_event_prompts[event_id] + ) + return {} if prompt is None else {FIELD_PROMPT: prompt} + + def set_base_prompt(self, prompt: str) -> None: + """Record the rollout prompt restored when a text event is cleared.""" + self._base_prompt = prompt + + +def _pose_segments( + command: Mapping[str, Any], + *, + start_s: float, + end_s: float, +) -> list[PoseSegment]: + """Return integrator-ready segments for one step window.""" + raw = command.get("segments") + if not raw: + # A source that supplies only level state still drives the step; the + # whole window then holds one constant command. + return [(start_s, end_s, _keys_from_axes(command))] + segments: list[PoseSegment] = [] + for segment_start, segment_end, axes in raw: + if float(segment_end) <= float(segment_start): + continue + segments.append( + (float(segment_start), float(segment_end), _keys_from_axes(axes)) + ) + if not segments: + return [(start_s, end_s, _keys_from_axes(command))] + return segments + + +def _required_int(metadata: Mapping[str, Any], name: str) -> int: + if name not in metadata: + raise ValueError( + f"Lingbot input mapping requires StepRequest.metadata[{name!r}]; the " + f"session did not provide it." + ) + return int(metadata[name]) + + +__all__ = [ + "CAMERA_COMMAND", + "FIELD_CAMERA_INTRINSICS", + "FIELD_CAMERA_TRAJECTORY", + "FIELD_TOTAL_CAMERA_FRAMES", + "KeyboardToCameraCommand", + "LingbotCameraTrace", + "LingbotInputMapping", + "TEXT_EVENT", + "TextEventSelection", + "load_camera_trace", +] diff --git a/integrations/lingbot/lingbot/runner.py b/integrations/lingbot/lingbot/runner.py index de269ae78..24016bb46 100644 --- a/integrations/lingbot/lingbot/runner.py +++ b/integrations/lingbot/lingbot/runner.py @@ -20,29 +20,32 @@ from dataclasses import dataclass, field from pathlib import Path -import numpy as np -import torch from loguru import logger -from flashdreams.core.io.disk import default_flashdreams_cache_dir -from flashdreams.core.io.download import download_to_cache from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig -from flashdreams.infra.runner_io import ( - ensure_output_dir, - load_first_frame_tensor, - runner_artifact_path, - write_runner_stats, - write_video_tensor, -) -from lingbot.encoder.camctrl import CamCtrlInput -from lingbot.encoder.utils import ( - get_Ks_transformed, - preprocess_example_poses, +from flashdreams.runtime import InputCanonicalizer, UserInputs, UserInputSchema +from flashdreams.runtime.metrics import NullMetricsRecorder +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, ) from lingbot.pipeline import ( LingbotWorldInferencePipeline, ) +from lingbot.runtime import ( + LingbotModelAdapter, + LingbotRunnerOutputTarget, + inference_config_from_runner_config, + inference_input_from_replay_inputs, + replay_inputs_from_runner_config, +) __all__ = [ "LingbotWorldRunnerConfig", @@ -58,68 +61,6 @@ _INTRINSICS_REFERENCE_WIDTH = 832 """Capture-resolution width matching :data:`_INTRINSICS_REFERENCE_HEIGHT`.""" -EXAMPLE_DATA_BASE_URL = ( - "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples" -) -"""HTTP base URL for the canonical examples shared by all LingBot versions.""" - -EXAMPLE_DATA_DIR_LOCAL = default_flashdreams_cache_dir() / "example_data/lingbot_world" -"""Local cache root where downloaded example folders are stored.""" - -EXAMPLE_DATA_FILENAMES = ( - "image.jpg", - "poses.npy", - "intrinsics.npy", - "prompt.txt", -) -"""Example assets downloaded when each file is available upstream.""" - -EXAMPLE_DATA_AVAILABLE_IDXS = (0, 1, 2, 3, 4, 5) -"""Supported upstream example indices currently hosted under ``examples/``.""" - -EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS = (0, 1, 2, 5) -"""Example indices that provide their own upstream ``prompt.txt`` file.""" - - -def example_data_dirname(example_idx: int) -> str: - """Format ``example_idx`` into the upstream folder naming convention.""" - assert example_idx in EXAMPLE_DATA_AVAILABLE_IDXS, ( - f"--example_idx must be one of {EXAMPLE_DATA_AVAILABLE_IDXS}." - ) - return f"{example_idx:02d}" - - -def ensure_example_data_downloaded(*, is_rank_zero: bool, example_idx: int) -> Path: - """Download bundled GitHub example files on rank 0; barrier other ranks. - - The runner calls this from :meth:`LingbotWorldRunner._fill_example_data_defaults`; - the WebRTC server calls it from its ``main()`` so the same files - land on disk before the server's - ``LingbotWebRTCSessionManager._initialize_sync`` checks for them. The - download itself is small (image + intrinsics + poses, plus a prompt - when available), uses the public LingBot-World GitHub raw URLs, and - is cached at :data:`EXAMPLE_DATA_DIR_LOCAL` so repeat calls are - no-ops. - """ - example_dirname = example_data_dirname(example_idx) - cache_dir = EXAMPLE_DATA_DIR_LOCAL / example_dirname - if is_rank_zero: - for filename in EXAMPLE_DATA_FILENAMES: - if ( - filename == "prompt.txt" - and example_idx not in EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS - ): - continue - download_to_cache( - f"{EXAMPLE_DATA_BASE_URL}/{example_dirname}/{filename}", - cache_dir=cache_dir, - filename=filename, - ) - if torch.distributed.is_initialized(): - torch.distributed.barrier() - return cache_dir - - @dataclass(kw_only=True) class LingbotWorldRunnerConfig(RunnerConfig): """Runner config for every shipped LingBot-World variant.""" @@ -222,118 +163,34 @@ def _fill_example_data_defaults(self) -> None: cfg.prompt_path = example_dir / "prompt.txt" def run(self) -> None: - """Drive an AR rollout until the camera stream is exhausted.""" + """Drive an AR rollout through the Lingbot runtime API path.""" cfg = self.config - if cfg.example_data: - self._fill_example_data_defaults() - assert cfg.image_path is not None, ( - "LingbotWorldRunner requires --image_path (first-frame RGB image)." - ) - assert cfg.pose_path is not None, ( - "LingbotWorldRunner requires --pose_path " - "(.npy of [T, 4, 4] camera-to-world matrices)." + adapter = LingbotModelAdapter() + inference_config = inference_config_from_runner_config( + cfg, + device=f"cuda:{self.local_rank}" if self.world_size > 1 else cfg.device, + pipeline=self.pipeline, ) - assert cfg.intrinsic_path is not None, ( - "LingbotWorldRunner requires --intrinsic_path " - "(.npy of [T, 4] camera intrinsics)." - ) - - prompt = self._resolve_prompt() - device = torch.device(f"cuda:{self.local_rank}") - - # Pipeline / encoder accept ``[*batch_shape, ...]`` shapes; the - # shipped configs pin ``batch_shape=()`` so a single-rollout layout - # is just ``[T, C, H, W]`` (image) / ``[T, 4, 4]`` (poses) / - # ``[T, 4]`` (intrinsics). - first_frames_t = load_first_frame_tensor( - cfg.image_path, - pixel_height=cfg.pixel_height, - pixel_width=cfg.pixel_width, - device=device, - dtype=torch.bfloat16, - interpolation="cubic", - install_hint="Install the lingbot plugin: pip install flashdreams-lingbot.", - ) - - Ks = np.load(cfg.intrinsic_path) - Ks_t = torch.from_numpy(Ks).to(device=device, dtype=torch.float32) - # Rescale capture-resolution intrinsics to the runner's frame size. - camera_intrinsics_t = get_Ks_transformed( - Ks_t, - height_org=_INTRINSICS_REFERENCE_HEIGHT, - width_org=_INTRINSICS_REFERENCE_WIDTH, - height_resize=cfg.pixel_height, - width_resize=cfg.pixel_width, - height_final=cfg.pixel_height, - width_final=cfg.pixel_width, + replay_inputs = replay_inputs_from_runner_config( + cfg, + is_rank_zero=self.is_rank_zero, ) - - c2ws = np.load(cfg.pose_path) - c2ws, trans_normalizer = preprocess_example_poses(c2ws) - camera_poses_t = torch.from_numpy(c2ws).to(device=device, dtype=torch.float32) - total_camera_frames = camera_poses_t.shape[0] - - if self.is_rank_zero: - logger.info( - f"[{cfg.runner_name}] loaded first_frame=" - f"{tuple(first_frames_t.shape)}, camera_poses=" - f"{tuple(camera_poses_t.shape)}" - ) - - cache = self.pipeline.initialize_cache(text=[prompt], image=first_frames_t) - - torch.cuda.synchronize() - if torch.distributed.is_initialized(): - torch.distributed.barrier() - - output_stream = self.create_video_output_stream(fps=cfg.fps) - start = 0 - for i in range(cfg.total_blocks): - num_frames = self.pipeline.get_num_output_frames(i) - end = start + num_frames - if end > total_camera_frames: - break - if self.is_rank_zero: - logger.info( - f"[{cfg.runner_name}] AR step {i}/{cfg.total_blocks}, " - f"num_frames={num_frames}, frames=[{start}, {end})" - ) - camctrl_input = CamCtrlInput( - intrinsics=camera_intrinsics_t[start:end], - poses=camera_poses_t[start:end], - world_scale=float(trans_normalizer), - ) - video_chunk = self.pipeline.generate( - autoregressive_index=i, - cache=cache, - input=camctrl_input, - ) - stats = self.pipeline.finalize(autoregressive_index=i, cache=cache) - output_stream.process(video_chunk, autoregressive_index=i, stats=stats) - start = end - - video = output_stream.finish() - if video is None: - return - - ensure_output_dir(cfg.output_dir) - video_path = runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4") - write_video_tensor( - video, - video_path, + initial_inputs = inference_input_from_replay_inputs(replay_inputs) + output_target = LingbotRunnerOutputTarget( + output_stream=self.create_video_output_stream(fps=cfg.fps), + output_dir=cfg.output_dir, + runner_name=cfg.runner_name, fps=cfg.fps, - layout="tchw", - install_hint="Install the lingbot plugin: pip install flashdreams-lingbot.", ) - logger.info( - f"[{cfg.runner_name}] wrote video {tuple(video.shape)} " - f"-> {video_path.resolve()}" + mapping = adapter.create_input_mapping(replay_inputs) + run_inference_session( + adapter=adapter, + config=inference_config, + mapping=mapping, + canonicalizer=InputCanonicalizer(), + source_schema=UserInputSchema(description="Lingbot runner fixed inputs"), + user_inputs=UserInputs(), + initial_inputs=initial_inputs, + output=output_target, + metrics=NullMetricsRecorder(), ) - - if output_stream.stats_history: - stats_path = write_runner_stats( - cfg.output_dir, cfg.runner_name, output_stream.stats_history - ) - logger.info( - f"[{cfg.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" - ) diff --git a/integrations/lingbot/lingbot/runtime.py b/integrations/lingbot/lingbot/runtime.py new file mode 100644 index 000000000..c15057aa5 --- /dev/null +++ b/integrations/lingbot/lingbot/runtime.py @@ -0,0 +1,1084 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot runtime API adapter and replay session implementation.""" + +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 + +import numpy as np +import torch +import torch.distributed as dist +from loguru import logger + +from flashdreams.core.distributed import init as init_distributed +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.infra.runner_io import ( + ensure_output_dir, + load_first_frame_tensor, + runner_artifact_path, + write_runner_stats, + write_video_tensor, +) +from flashdreams.infra.video_output import RunnerVideoOutputStream, VideoStepResult +from flashdreams.runtime import ( + CanonicalInputSchema, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InputField, + OutputArtifact, +) +from flashdreams.runtime.interfaces import InferenceRuntime, InferenceSession +from flashdreams.runtime.types import StepRequest, StepResult, TimeWindow +from lingbot.encoder.camctrl import CamCtrlInput +from lingbot.example_data import ( + EXAMPLE_DATA_AVAILABLE_IDXS, + EXAMPLE_DATA_DIR_LOCAL, + EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS, + ensure_example_data_downloaded, + example_asset_urls, + example_data_dirname, +) +from lingbot.input_mapping import ( + CAMERA_COMMAND, + FIELD_CAMERA_INTRINSICS, + FIELD_CAMERA_TRAJECTORY, + FIELD_TOTAL_CAMERA_FRAMES, + TEXT_EVENT, + LingbotCameraTrace, + LingbotInputMapping, + load_camera_trace, +) + +LINGBOT_MODEL_ID = "lingbot" +DEFAULT_LINGBOT_PRESET = "lingbot-world-fast-taehv-window15-sink3" +DEFAULT_PIXEL_HEIGHT = 464 +DEFAULT_PIXEL_WIDTH = 832 +DEFAULT_FPS = 16 + +_INTRINSICS_REFERENCE_HEIGHT = 480 +_INTRINSICS_REFERENCE_WIDTH = 832 +_INSTALL_HINT = "Install the lingbot plugin: pip install flashdreams-lingbot." + +FIELD_PROMPT = "prompt" +FIELD_FIRST_FRAME_PATH = "first_frame_path" +FIELD_CAMERA_POSES_PATH = "camera_poses_path" +FIELD_CAMERA_INTRINSICS_PATH = "camera_intrinsics_path" +FIELD_TOTAL_BLOCKS = "total_blocks" +FIELD_PIXEL_HEIGHT = "pixel_height" +FIELD_PIXEL_WIDTH = "pixel_width" +FIELD_FPS = "fps" +FIELD_WORLD_SCALE = "world_scale" + +PipelineFactory = Callable[[Any, str], Any] + + +@dataclass(frozen=True, kw_only=True, slots=True) +class LingbotReplayInputs: + """Resolved model-facing Lingbot replay inputs.""" + + prompt: str + first_frame_path: Path + camera_poses_path: Path + camera_intrinsics_path: Path + total_blocks: int = 20 + pixel_height: int = DEFAULT_PIXEL_HEIGHT + pixel_width: int = DEFAULT_PIXEL_WIDTH + fps: int = DEFAULT_FPS + world_scale: float | None = None + + def __post_init__(self) -> None: + if self.total_blocks <= 0: + raise ValueError("LingbotReplayInputs.total_blocks must be > 0.") + if self.pixel_height <= 0 or self.pixel_width <= 0: + raise ValueError("LingbotReplayInputs pixel dimensions must be > 0.") + if self.fps <= 0: + raise ValueError("LingbotReplayInputs.fps must be > 0.") + if self.world_scale is not None and self.world_scale <= 0: + raise ValueError("LingbotReplayInputs.world_scale must be > 0.") + object.__setattr__(self, "prompt", " ".join(self.prompt.split())) + object.__setattr__(self, "first_frame_path", Path(self.first_frame_path)) + object.__setattr__(self, "camera_poses_path", Path(self.camera_poses_path)) + object.__setattr__( + self, + "camera_intrinsics_path", + Path(self.camera_intrinsics_path), + ) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class LingbotSessionInputs: + """Session-global Lingbot state established at session start or reset. + + The camera trajectory is deliberately absent: it arrives per step through + ``InferenceInput.step``, built by the selected input mapping from either a + fixed trace or live user events. + """ + + prompt: str + first_frame_path: Path + total_blocks: int + pixel_height: int + pixel_width: int + fps: int + world_scale: float + total_camera_frames: int | None = None + + def __post_init__(self) -> None: + if self.total_blocks <= 0: + raise ValueError("LingbotSessionInputs.total_blocks must be > 0.") + if self.pixel_height <= 0 or self.pixel_width <= 0: + raise ValueError("LingbotSessionInputs pixel dimensions must be > 0.") + if self.fps <= 0: + raise ValueError("LingbotSessionInputs.fps must be > 0.") + if self.world_scale < 0: + raise ValueError("LingbotSessionInputs.world_scale must be >= 0.") + if self.total_camera_frames is not None and self.total_camera_frames <= 0: + raise ValueError( + "LingbotSessionInputs.total_camera_frames must be > 0 when set." + ) + object.__setattr__(self, "first_frame_path", Path(self.first_frame_path)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class LingbotReplayRuntimeOptions: + """Construction knobs for the Lingbot replay runtime.""" + + pipeline_config: Any + pipeline: Any | None = None + pipeline_factory: PipelineFactory | None = None + output_layout: VideoTensorLayout = "tchw" + + +class LingbotModelAdapter: + """Model adapter exposing Lingbot through ``flashdreams.runtime``.""" + + def __init__( + self, + *, + runtime_factory: Callable[..., InferenceRuntime] | None = None, + pipeline_factory: PipelineFactory | None = None, + ) -> None: + self._runtime_factory = runtime_factory or LingbotReplayRuntime + self._pipeline_factory = pipeline_factory + + @property + def model_id(self) -> str: + return LINGBOT_MODEL_ID + + @property + def inference_input_schema(self) -> InferenceInputSchema: + return InferenceInputSchema( + description="Lingbot camera-control model inputs.", + global_conditioning_fields=( + InputField( + name=FIELD_PROMPT, + input_modality="text", + frequency_consumed="once", + description=( + "Prompt text for the rollout. A non-empty value passed " + "to step() requests a text-event context swap." + ), + ), + InputField( + name=FIELD_FIRST_FRAME_PATH, + input_modality="image/path", + frequency_consumed="once", + description="First-frame RGB image path.", + ), + InputField(name=FIELD_TOTAL_BLOCKS, input_modality="count"), + InputField(name=FIELD_PIXEL_HEIGHT, input_modality="pixel-height"), + InputField(name=FIELD_PIXEL_WIDTH, input_modality="pixel-width"), + InputField(name=FIELD_FPS, input_modality="fps"), + InputField( + name=FIELD_WORLD_SCALE, + required=False, + input_modality="scale", + frequency_consumed="once", + description="Pose normalizer; supplied by the input mapping.", + ), + InputField( + name=FIELD_TOTAL_CAMERA_FRAMES, + required=False, + input_modality="count", + frequency_consumed="once", + description=( + "Frames the input source can supply. Absent means " + "unbounded, so only total_blocks ends the rollout." + ), + ), + ), + step_fields=( + InputField( + name=FIELD_CAMERA_TRAJECTORY, + input_modality="c2w_sequence", + frequency_consumed="per_step", + metadata={"shape": "[T,4,4]", "frame": "camera_to_world"}, + description="Camera-to-world poses for this chunk's frames.", + ), + InputField( + name=FIELD_CAMERA_INTRINSICS, + input_modality="intrinsics_vec4_sequence", + frequency_consumed="per_step", + metadata={"shape": "[T,4]"}, + description="Per-frame intrinsics for this chunk's frames.", + ), + ), + ) + + @property + def canonical_input_schema(self) -> CanonicalInputSchema | None: + return CanonicalInputSchema( + modalities=(CAMERA_COMMAND, TEXT_EVENT), + description="Lingbot live camera control and text events.", + ) + + def default_input_mapping(self) -> LingbotInputMapping | None: + """Return no default mapping; Lingbot mappings are scenario-bound. + + Both trajectory sources need scenario data the adapter does not have + here: a fixed trace needs its ``.npy`` files, and live control needs + base intrinsics and a world scale. Callers build one with + :meth:`create_input_mapping`. + """ + return None + + def create_input_mapping( + self, + replay_inputs: LingbotReplayInputs, + *, + text_event_prompts: Mapping[str, str] | None = None, + ) -> LingbotInputMapping: + """Build the fixed-trace mapping for a resolved replay scenario.""" + mapping = LingbotInputMapping( + fps=replay_inputs.fps, + trace=load_camera_trace( + camera_poses_path=replay_inputs.camera_poses_path, + camera_intrinsics_path=replay_inputs.camera_intrinsics_path, + pixel_height=replay_inputs.pixel_height, + pixel_width=replay_inputs.pixel_width, + intrinsics_reference_height=_INTRINSICS_REFERENCE_HEIGHT, + intrinsics_reference_width=_INTRINSICS_REFERENCE_WIDTH, + world_scale=replay_inputs.world_scale, + ), + text_event_prompts=text_event_prompts, + ) + mapping.set_base_prompt(replay_inputs.prompt) + return mapping + + def create_live_input_mapping( + self, + *, + fps: int, + base_intrinsics: Any, + world_scale: float, + prompt: str = "", + text_event_prompts: Mapping[str, str] | None = None, + trace: LingbotCameraTrace | None = None, + ) -> LingbotInputMapping: + """Build the event-driven mapping used by keyboard-driving scenarios.""" + mapping = LingbotInputMapping( + fps=fps, + trace=trace, + base_intrinsics=base_intrinsics, + world_scale=world_scale, + text_event_prompts=text_event_prompts, + ) + mapping.set_base_prompt(prompt) + return mapping + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.model_id: + raise ValueError( + f"Lingbot adapter requires model_id={self.model_id!r}, " + f"got {config.model_id!r}." + ) + self.pipeline_config(config) + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + 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")), + ), + ) + + def preset_id(self, config: InferenceConfig | None) -> str: + return ( + DEFAULT_LINGBOT_PRESET + if config is None or config.preset_id is None + else config.preset_id + ) + + def pipeline_config(self, config: InferenceConfig) -> Any: + custom = config.runtime_options.get("pipeline_config") + if custom is not None: + return custom + preset_id = self.preset_id(config) + from lingbot.config import PIPELINE_CONFIGS # noqa: PLC0415 + + try: + return PIPELINE_CONFIGS[preset_id] + except KeyError as exc: + supported = ", ".join(sorted(PIPELINE_CONFIGS)) + raise ValueError( + f"Unsupported Lingbot preset_id={preset_id!r}. " + f"Supported presets: {supported}." + ) from exc + + def default_replay_prompt(self, config: InferenceConfig | None) -> str: + from lingbot.config import RUNNER_CONFIGS # noqa: PLC0415 + + runner = RUNNER_CONFIGS.get(self.preset_id(config)) + return "" if runner is None else str(getattr(runner, "prompt", "")) + + +class LingbotReplayRuntime: + """Heavyweight Lingbot runtime consumed by the standard loop.""" + + def __init__( + self, + *, + config: InferenceConfig, + options: LingbotReplayRuntimeOptions, + ) -> None: + self.config = config + self.options = options + if _is_torchrun_env() and not dist.is_initialized(): + init_distributed() + + if dist.is_initialized(): + self.local_rank = int(os.environ.get("LOCAL_RANK", "0")) + self.world_size = dist.get_world_size() + self.global_rank = dist.get_rank() + device = f"cuda:{self.local_rank}" + else: + self.local_rank = 0 + self.world_size = 1 + self.global_rank = 0 + device = config.device or "cuda" + + self.is_rank_zero = self.global_rank == 0 + if options.pipeline is not None: + self.pipeline = options.pipeline + self._owns_pipeline = False + else: + factory = options.pipeline_factory or _default_pipeline_factory + self.pipeline = factory(options.pipeline_config, device) + self._owns_pipeline = True + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + session_inputs = session_inputs_from_inference_input(inputs) + return LingbotReplaySession( + pipeline=self.pipeline, + session_inputs=session_inputs, + device=torch.device(f"cuda:{self.local_rank}") + if dist.is_initialized() + else torch.device(self.config.device or "cuda"), + is_rank_zero=self.is_rank_zero, + output_layout=self.options.output_layout, + ) + + def close(self) -> None: + pipeline = getattr(self, "pipeline", None) + if self._owns_pipeline and pipeline is not None: + close = getattr(pipeline, "close", None) + if callable(close): + close() + del self.pipeline + device = torch.device(self.config.device or "cuda") + if device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.empty_cache() + + +class LingbotReplaySession: + """One Lingbot rollout driven by per-step camera inputs.""" + + def __init__( + self, + *, + pipeline: Any, + session_inputs: LingbotSessionInputs, + device: torch.device, + is_rank_zero: bool, + output_layout: VideoTensorLayout, + ) -> None: + self.pipeline = pipeline + self.inputs = session_inputs + self.device = device + self.is_rank_zero = is_rank_zero + 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() + if self.device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.synchronize(device=self.device) + if dist.is_initialized(): + dist.barrier() + + def next_step_request(self) -> StepRequest | None: + if self._closed: + return None + if self._step_index >= self.inputs.total_blocks: + return None + num_frames = int(self.pipeline.get_num_output_frames(self._step_index)) + 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, + # 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( + start_s=self._frame_start / fps, + end_s=frame_end / fps, + ), + metadata={ + "num_frames": num_frames, + "frame_start": self._frame_start, + }, + ) + + 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)) + self._apply_global_conditioning_update(inputs) + camera_poses = _require_step_tensor( + inputs, + FIELD_CAMERA_TRAJECTORY, + expected_shape=(num_frames, 4, 4), + ) + camera_intrinsics = _require_step_tensor( + inputs, + FIELD_CAMERA_INTRINSICS, + expected_shape=(num_frames, 4), + ) + frame_start = self._frame_start + frame_end = frame_start + num_frames + + if self.is_rank_zero: + logger.info( + "Lingbot runtime step {} frames=[{}, {})", + step_index, + frame_start, + frame_end, + ) + camctrl_input = CamCtrlInput( + intrinsics=camera_intrinsics.to(device=self.device, dtype=torch.float32), + 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, + output_window=TimeWindow( + start_s=frame_start / self.inputs.fps, + end_s=frame_end / self.inputs.fps, + ), + metrics=metrics, + ) + + 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._frame_start = 0 + + def _apply_global_conditioning_update(self, inputs: InferenceInput) -> None: + """Apply a mid-rollout text-event context swap, when one was requested. + + Text events reach the model as a session-global prompt update rather + than a per-step field, because they replace the rollout's whole + cross-attention text context. Not every pipeline can do this, so the + capability is probed the same way the WebRTC runtime probes it, and + only when a swap is actually requested. + """ + prompt = inputs.global_conditioning.get(FIELD_PROMPT) + if prompt is None or prompt == self._active_prompt: + return + transformer = self.pipeline.diffusion_model.transformer + replace_text_embeddings = getattr(transformer, "replace_text_embeddings", None) + if not callable(replace_text_embeddings): + raise RuntimeError( + "Lingbot text events need a pipeline whose transformer supports " + "replace_text_embeddings; this pipeline does not." + ) + 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._active_prompt = prompt + if self.is_rank_zero: + logger.info("Lingbot text context updated at step {}", self._step_index) + + def close(self) -> None: + self._closed = True + cache = getattr(self, "_cache", None) + if cache is not None: + del self._cache + + def _initialize_cache(self) -> Any: + first_frames = load_first_frame_tensor( + self.inputs.first_frame_path, + pixel_height=self.inputs.pixel_height, + pixel_width=self.inputs.pixel_width, + device=self.device, + dtype=self.dtype, + interpolation="cubic", + install_hint=_INSTALL_HINT, + ) + return self.pipeline.initialize_cache( + text=[self.inputs.prompt], + image=first_frames, + ) + + +def _require_step_tensor( + inputs: InferenceInput, + name: str, + *, + expected_shape: tuple[int, ...], +) -> torch.Tensor: + """Return one required per-step camera tensor, shape-checked.""" + if name not in inputs.step: + raise ValueError( + f"Lingbot step inputs are missing {name!r}. The selected input " + f"mapping must produce it for every step." + ) + value = inputs.step[name] + if not isinstance(value, torch.Tensor): + value = torch.as_tensor(np.asarray(value), dtype=torch.float32) + if tuple(value.shape) != expected_shape: + raise ValueError( + f"Lingbot step input {name!r} must have shape {expected_shape}, got " + f"{tuple(value.shape)}." + ) + return value + + +@dataclass(slots=True) +class LingbotRunnerOutputTarget: + """Runner-compatible MP4/stats output target for Lingbot replay results.""" + + output_stream: RunnerVideoOutputStream + output_dir: Path + runner_name: str + fps: int | float + install_hint: str = _INSTALL_HINT + _opened: bool = False + + def open(self) -> None: + 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): + raise TypeError( + "LingbotRunnerOutputTarget requires VideoStepResult output, " + f"got {type(video_result).__name__}." + ) + self.output_stream.process( + video_result.video_chunk, + autoregressive_index=video_result.chunk_index, + stats=video_result.stats or dict(result.metrics), + ) + + def close(self) -> tuple[OutputArtifact, ...]: + self._opened = False + artifacts: list[OutputArtifact] = [] + video = self.output_stream.finish() + if video is None: + return () + + ensure_output_dir(self.output_dir) + video_path = runner_artifact_path(self.output_dir, self.runner_name, "mp4") + write_video_tensor( + video, + video_path, + fps=self.fps, + layout="tchw", + install_hint=self.install_hint, + ) + logger.info( + "[{}] wrote video {} -> {}", + self.runner_name, + tuple(video.shape), + video_path.resolve(), + ) + artifacts.append( + OutputArtifact(kind="video/mp4", uri=str(video_path.resolve())) + ) + if self.output_stream.stats_history: + stats_path = write_runner_stats( + self.output_dir, + self.runner_name, + self.output_stream.stats_history, + ) + logger.info( + "[{}] wrote per-AR-step stats -> {}", + self.runner_name, + stats_path.resolve(), + ) + artifacts.append( + OutputArtifact(kind="application/json", uri=str(stats_path.resolve())) + ) + return tuple(artifacts) + + +def inference_config_from_runner_config( + runner_config: Any, + *, + device: str, + pipeline: Any | None = None, +) -> InferenceConfig: + """Build the runtime config directly from a Lingbot runner config.""" + runtime_options: dict[str, Any] = { + "pipeline_config": runner_config.pipeline, + "output_layout": runner_config.postprocess_output_layout or "tchw", + } + if pipeline is not None: + runtime_options["pipeline"] = pipeline + compile_network = getattr( + runner_config.pipeline.diffusion_model.transformer, + "compile_network", + None, + ) + return InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=str(runner_config.pipeline.name), + device=device, + compile=None if compile_network is None else bool(compile_network), + runtime_options=runtime_options, + ) + + +def inference_input_from_runner_config( + runner_config: Any, + *, + is_rank_zero: bool, +) -> InferenceInput: + """Build session-global runtime inputs from a Lingbot runner config.""" + return inference_input_from_replay_inputs( + replay_inputs_from_runner_config(runner_config, is_rank_zero=is_rank_zero) + ) + + +def replay_inputs_from_runner_config( + runner_config: Any, + *, + is_rank_zero: bool, +) -> LingbotReplayInputs: + """Resolve a Lingbot runner config into scenario-level replay inputs.""" + return replay_inputs_from_mapping( + { + FIELD_PROMPT: getattr(runner_config, "prompt", ""), + "prompt_path": getattr(runner_config, "prompt_path", None), + FIELD_FIRST_FRAME_PATH: getattr(runner_config, "image_path", None), + FIELD_CAMERA_POSES_PATH: getattr(runner_config, "pose_path", None), + FIELD_CAMERA_INTRINSICS_PATH: getattr( + runner_config, + "intrinsic_path", + None, + ), + FIELD_TOTAL_BLOCKS: getattr(runner_config, "total_blocks", 20), + FIELD_PIXEL_HEIGHT: getattr( + runner_config, + "pixel_height", + DEFAULT_PIXEL_HEIGHT, + ), + FIELD_PIXEL_WIDTH: getattr( + runner_config, + "pixel_width", + DEFAULT_PIXEL_WIDTH, + ), + FIELD_FPS: getattr(runner_config, "fps", DEFAULT_FPS), + "example_data": getattr(runner_config, "example_data", False), + "example_idx": getattr(runner_config, "example_idx", 0), + }, + is_rank_zero=is_rank_zero, + ) + + +def inference_input_from_replay_inputs( + replay_inputs: LingbotReplayInputs, +) -> InferenceInput: + """Encode resolved Lingbot replay inputs into ``InferenceInput``.""" + payload: dict[str, Any] = { + FIELD_PROMPT: replay_inputs.prompt, + FIELD_FIRST_FRAME_PATH: replay_inputs.first_frame_path, + FIELD_TOTAL_BLOCKS: replay_inputs.total_blocks, + FIELD_PIXEL_HEIGHT: replay_inputs.pixel_height, + FIELD_PIXEL_WIDTH: replay_inputs.pixel_width, + FIELD_FPS: replay_inputs.fps, + } + if replay_inputs.world_scale is not None: + payload[FIELD_WORLD_SCALE] = replay_inputs.world_scale + return InferenceInput(global_conditioning=payload) + + +def session_inputs_from_inference_input( + inputs: InferenceInput, +) -> LingbotSessionInputs: + """Decode and validate session-global Lingbot inputs.""" + missing = LingbotModelAdapter().inference_input_schema.missing_global_conditioning( + inputs + ) + if missing: + raise ValueError(f"Lingbot session inputs missing required fields: {missing}.") + gc = inputs.global_conditioning + if gc.get(FIELD_WORLD_SCALE) is None: + raise ValueError( + "Lingbot session inputs require 'world_scale'; the selected input " + "mapping supplies it from the camera trace or live control setup." + ) + total_camera_frames = gc.get(FIELD_TOTAL_CAMERA_FRAMES) + return LingbotSessionInputs( + prompt=str(gc[FIELD_PROMPT]), + first_frame_path=Path(gc[FIELD_FIRST_FRAME_PATH]), + total_blocks=int(gc[FIELD_TOTAL_BLOCKS]), + pixel_height=int(gc[FIELD_PIXEL_HEIGHT]), + pixel_width=int(gc[FIELD_PIXEL_WIDTH]), + fps=int(gc[FIELD_FPS]), + world_scale=float(gc[FIELD_WORLD_SCALE]), + total_camera_frames=( + None if total_camera_frames is None else int(total_camera_frames) + ), + ) + + +def replay_inputs_from_mapping( + value: Any, + *, + default_prompt: str = "", + is_rank_zero: bool = True, +) -> LingbotReplayInputs: + """Resolve app/CLI replay values into direct Lingbot runtime inputs.""" + if isinstance(value, LingbotReplayInputs): + _require_existing_replay_paths(value) + return value + if value is None: + value = {} + if not isinstance(value, Mapping): + raise TypeError( + "Lingbot replay inputs must be a LingbotReplayInputs, mapping, or None." + ) + + example_idx = int(value.get("example_idx", 0)) + if example_idx not in EXAMPLE_DATA_AVAILABLE_IDXS: + raise ValueError( + f"Lingbot replay example_idx must be one of {EXAMPLE_DATA_AVAILABLE_IDXS}." + ) + + first_frame_path = _optional_path( + value.get(FIELD_FIRST_FRAME_PATH, value.get("image_path")) + ) + poses_path = _optional_path( + value.get(FIELD_CAMERA_POSES_PATH, value.get("pose_path")) + ) + intrinsics_path = _optional_path( + value.get(FIELD_CAMERA_INTRINSICS_PATH, value.get("intrinsic_path")) + ) + prompt_path = _optional_path(value.get("prompt_path")) + example_data = _resolve_example_data_default(value) + + if example_data and ( + first_frame_path is None + or poses_path is None + or intrinsics_path is None + or ( + prompt_path is None + and not _has_nonempty_value(value, FIELD_PROMPT) + and example_idx in EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS + ) + ): + example_dir = ensure_example_data_downloaded( + is_rank_zero=is_rank_zero, + example_idx=example_idx, + ) + first_frame_path = first_frame_path or example_dir / "image.jpg" + poses_path = poses_path or example_dir / "poses.npy" + intrinsics_path = intrinsics_path or example_dir / "intrinsics.npy" + if ( + prompt_path is None + and not _has_nonempty_value(value, FIELD_PROMPT) + and example_idx in EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS + ): + prompt_path = example_dir / "prompt.txt" + + replay_inputs = LingbotReplayInputs( + prompt=_resolve_prompt( + value, + prompt_path=prompt_path, + default_prompt=default_prompt, + ), + first_frame_path=_require_path_value( + first_frame_path, + label=FIELD_FIRST_FRAME_PATH, + ), + camera_poses_path=_require_path_value( + poses_path, + label=FIELD_CAMERA_POSES_PATH, + ), + camera_intrinsics_path=_require_path_value( + intrinsics_path, + label=FIELD_CAMERA_INTRINSICS_PATH, + ), + total_blocks=int(value.get(FIELD_TOTAL_BLOCKS, 20)), + pixel_height=int(value.get(FIELD_PIXEL_HEIGHT, DEFAULT_PIXEL_HEIGHT)), + pixel_width=int(value.get(FIELD_PIXEL_WIDTH, DEFAULT_PIXEL_WIDTH)), + fps=int(value.get(FIELD_FPS, DEFAULT_FPS)), + world_scale=( + None + if FIELD_WORLD_SCALE not in value or value[FIELD_WORLD_SCALE] is None + else float(value[FIELD_WORLD_SCALE]) + ), + ) + _require_existing_replay_paths(replay_inputs) + if prompt_path is not None: + _require_existing_path(prompt_path, label="prompt_path") + return replay_inputs + + +def build_lingbot_webrtc_runtime_config( + *, + preset_id: str, + pipeline_config: Any, + device: str, + seed: int, + compile_network: bool, + context_parallel_size: int, + video_height: int, + video_width: int, + fps: int, + warmup_chunks: int, + warmup_timeout_s: float, + example_idx: int, + prefer_sw_encoder: bool, + runtime_options: Mapping[str, Any] | None = None, +) -> Any: + """Build the Lingbot WebRTC runtime config from shared runtime inputs.""" + from lingbot.webrtc.session import LingbotRuntimeConfig # noqa: PLC0415 + + example_dirname = example_data_dirname(example_idx) + example_dir = EXAMPLE_DATA_DIR_LOCAL / example_dirname + if ( + example_idx == 0 + and not example_dir.exists() + and (EXAMPLE_DATA_DIR_LOCAL / "image.jpg").exists() + ): + example_dir = EXAMPLE_DATA_DIR_LOCAL + urls = example_asset_urls(example_idx) + runtime_config = LingbotRuntimeConfig( + config_name=preset_id, + pipeline_config=pipeline_config, + compile_network=compile_network, + seed=seed, + context_parallel_size=context_parallel_size, + device=device, + video_height=video_height, + video_width=video_width, + fps=fps, + warmup_chunks=warmup_chunks, + warmup_timeout_s=warmup_timeout_s, + encoder_backend="default" if prefer_sw_encoder else "auto", + example_data_dir=example_dir, + default_image_url=urls["image"], + default_intrinsics_url=urls["intrinsics"], + default_poses_url=urls["poses"], + ) + return _apply_webrtc_runtime_options(runtime_config, runtime_options or {}) + + +def _apply_webrtc_runtime_options(runtime_config: Any, options: Mapping[str, Any]) -> Any: + overrides: dict[str, Any] = {} + for name in ( + "world_scale", + "default_intrinsics", + "default_prompt", + "default_image_url", + "default_intrinsics_url", + "default_poses_url", + "encoder_bitrate_bps", + "encoder_gop", + "text_events", + ): + if name in options: + overrides[name] = options[name] + return replace(runtime_config, **overrides) if overrides else runtime_config + + +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], + *, + prompt_path: Path | None, + default_prompt: str, +) -> str: + prompt = str(value.get(FIELD_PROMPT, value.get("prompt", ""))).strip() + if prompt: + return prompt + if prompt_path is not None: + lines = prompt_path.read_text(encoding="utf-8").splitlines() + if lines: + prompt = lines[0].strip() + if prompt: + return prompt + return default_prompt.strip() + + +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") + ) + + +def _bool_value(value: Any) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + lowered = value.strip().lower() + if lowered in {"1", "true", "yes", "on"}: + return True + if lowered in {"0", "false", "no", "off"}: + return False + return bool(value) + + +def _has_nonempty_value(value: Mapping[str, Any], key: str) -> bool: + if key not in value: + return False + raw = value[key] + return raw is not None and raw != "" + + +def _optional_path(value: Any) -> Path | None: + if value is None or value == "": + return None + return Path(value) + + +def _require_path_value(value: Path | None, *, label: str) -> Path: + if value is None: + raise ValueError(f"Lingbot replay inputs require {label}.") + return value + + +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_intrinsics_path, + label=FIELD_CAMERA_INTRINSICS_PATH, + ) + + +def _require_existing_path(path: Path, *, label: str) -> None: + if not path.exists(): + raise FileNotFoundError(f"Lingbot replay inputs missing {label}: {path}") + + +def _is_torchrun_env() -> bool: + return "RANK" in os.environ and "WORLD_SIZE" in os.environ + + +__all__ = [ + "DEFAULT_FPS", + "DEFAULT_LINGBOT_PRESET", + "DEFAULT_PIXEL_HEIGHT", + "DEFAULT_PIXEL_WIDTH", + "FIELD_CAMERA_INTRINSICS_PATH", + "FIELD_CAMERA_POSES_PATH", + "FIELD_FIRST_FRAME_PATH", + "FIELD_FPS", + "FIELD_PIXEL_HEIGHT", + "FIELD_PIXEL_WIDTH", + "FIELD_PROMPT", + "FIELD_TOTAL_BLOCKS", + "FIELD_WORLD_SCALE", + "LINGBOT_MODEL_ID", + "LingbotModelAdapter", + "LingbotReplayInputs", + "LingbotReplayRuntime", + "LingbotReplayRuntimeOptions", + "LingbotReplaySession", + "LingbotRunnerOutputTarget", + "PipelineFactory", + "build_lingbot_webrtc_runtime_config", + "inference_config_from_runner_config", + "inference_input_from_replay_inputs", + "inference_input_from_runner_config", + "replay_inputs_from_inference_input", + "replay_inputs_from_mapping", +] diff --git a/integrations/lingbot/lingbot/webrtc/server.py b/integrations/lingbot/lingbot/webrtc/server.py index 4714ed85b..3825a5559 100644 --- a/integrations/lingbot/lingbot/webrtc/server.py +++ b/integrations/lingbot/lingbot/webrtc/server.py @@ -48,12 +48,15 @@ from flashdreams.serving.webrtc.server import ( close_package_resources as _close_package_resources, ) -from lingbot.runner import ( +from flashdreams.runtime import InferenceConfig +from lingbot.example_data import ( EXAMPLE_DATA_AVAILABLE_IDXS, - EXAMPLE_DATA_BASE_URL, - EXAMPLE_DATA_DIR_LOCAL, ensure_example_data_downloaded, - example_data_dirname, +) +from lingbot.runtime import ( + LINGBOT_MODEL_ID, + LingbotModelAdapter, + build_lingbot_webrtc_runtime_config, ) from lingbot.webrtc.session import ( LingbotImagePayload, @@ -105,6 +108,12 @@ def parse_args() -> argparse.Namespace: default="cuda:0", help="Torch device used for the Lingbot runtime.", ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Base random seed for the Lingbot rollout.", + ) parser.add_argument( "--warmup_chunks", type=int, @@ -335,34 +344,31 @@ def build_runtime_config( raise ValueError("--video-height and --video-width must be > 0") if args.video_height % 16 != 0 or args.video_width % 16 != 0: raise ValueError("--video-height and --video-width must be divisible by 16") - example_idx = getattr(args, "example_idx", 0) - example_dirname = example_data_dirname(example_idx) - example_dir = EXAMPLE_DATA_DIR_LOCAL / example_dirname - if ( - example_idx == 0 - and not example_dir.exists() - and (EXAMPLE_DATA_DIR_LOCAL / "image.jpg").exists() - ): - example_dir = EXAMPLE_DATA_DIR_LOCAL - return LingbotRuntimeConfig( - config_name=args.config_name, + + inference_config = InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=args.config_name, + device=device_override or args.device, + compile=not args.no_compile, + runtime_options={"context_parallel_size": context_parallel_size}, + ) + adapter = LingbotModelAdapter() + adapter.validate_config(inference_config) + return build_lingbot_webrtc_runtime_config( + preset_id=adapter.preset_id(inference_config), + pipeline_config=adapter.pipeline_config(inference_config), + device=inference_config.device or args.device, + seed=int(getattr(args, "seed", 42)), compile_network=not args.no_compile, context_parallel_size=context_parallel_size, - device=device_override or args.device, - warmup_chunks=args.warmup_chunks, - warmup_timeout_s=args.warmup_timeout_s, video_height=args.video_height, video_width=args.video_width, fps=args.fps, - encoder_backend=( - "default" if getattr(args, "prefer_sw_encoder", False) else "auto" - ), - example_data_dir=example_dir, - default_image_url=f"{EXAMPLE_DATA_BASE_URL}/{example_dirname}/image.jpg", - default_intrinsics_url=( - f"{EXAMPLE_DATA_BASE_URL}/{example_dirname}/intrinsics.npy" - ), - default_poses_url=f"{EXAMPLE_DATA_BASE_URL}/{example_dirname}/poses.npy", + warmup_chunks=args.warmup_chunks, + warmup_timeout_s=args.warmup_timeout_s, + example_idx=getattr(args, "example_idx", 0), + prefer_sw_encoder=getattr(args, "prefer_sw_encoder", False), + runtime_options=inference_config.runtime_options, ) diff --git a/integrations/lingbot/lingbot/webrtc/session.py b/integrations/lingbot/lingbot/webrtc/session.py index c0b59680b..ee23bb967 100644 --- a/integrations/lingbot/lingbot/webrtc/session.py +++ b/integrations/lingbot/lingbot/webrtc/session.py @@ -476,6 +476,8 @@ class LingbotRuntimeConfig: text_events: tuple[TextEventSpec, ...] = field( default_factory=lambda: DEFAULT_TEXT_EVENTS ) + pipeline_config: Any | None = None + """Optional pre-resolved pipeline config used by shared demo adapters.""" @dataclass(frozen=True, slots=True) @@ -775,13 +777,16 @@ def _initialize_sync(self) -> None: if self._pipeline is not None: return - pipeline_configs = _pipeline_configs() - if self.config.config_name not in pipeline_configs: - supported = ", ".join(sorted(pipeline_configs)) - raise ValueError( - f"Unknown config_name={self.config.config_name!r}. " - f"Supported: {supported}" - ) + pipeline_config_base = self.config.pipeline_config + if pipeline_config_base is None: + pipeline_configs = _pipeline_configs() + if self.config.config_name not in pipeline_configs: + supported = ", ".join(sorted(pipeline_configs)) + raise ValueError( + f"Unknown config_name={self.config.config_name!r}. " + f"Supported: {supported}" + ) + 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(): @@ -796,7 +801,7 @@ def _initialize_sync(self) -> None: else self.config.seed ) pipeline_config = derive_config( - base_config=pipeline_configs[self.config.config_name], + base_config=pipeline_config_base, enable_sync_and_profile=True, diffusion_model=dict( seed=rollout_seed, @@ -1204,16 +1209,20 @@ class LingbotWebRTCSessionManager( 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, ) -> None: - runtime_config = runtime_config or LingbotRuntimeConfig() + 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=LingbotInferenceRuntime(config=runtime_config), + runtime=runtime, runtime_config=runtime_config, fps=fps, client_liveness_timeout_s=client_liveness_timeout_s, diff --git a/integrations/lingbot/pyproject.toml b/integrations/lingbot/pyproject.toml index f9e39cd92..b14cd2c91 100644 --- a/integrations/lingbot/pyproject.toml +++ b/integrations/lingbot/pyproject.toml @@ -41,6 +41,9 @@ dev = [ "pytest-asyncio>=0.23", ] +[project.scripts] +lingbot-demo = "lingbot.demo.cli:main" + # Each entry registers one ``runner_name`` slug with ``flashdreams-run``. # The discovery layer (``flashdreams.plugins.registry.discover_runners``) # scans this group at CLI startup; the entry-point name itself is purely diff --git a/integrations/lingbot/tests/test_demo_api.py b/integrations/lingbot/tests/test_demo_api.py new file mode 100644 index 000000000..95515e027 --- /dev/null +++ b/integrations/lingbot/tests/test_demo_api.py @@ -0,0 +1,625 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import Any, cast + +import numpy as np +import pytest +import torch +from aiohttp import web +from lingbot.demo import ( + DEFAULT_LINGBOT_PRESET, + LINGBOT_MODEL_ID, + LingbotDemoAdapter, + LingbotReplayInputs, + LingbotWebRTCScenario, +) +from lingbot.demo.cli import _replay_spec, _webrtc_spec, parse_args +from lingbot.demo.replay import ( + LingbotReplayRuntime, + LingbotReplayRuntimeOptions, +) +from lingbot.demo.webrtc import LingbotDemoWebRTCSessionManager +from lingbot.input_mapping import ( + FIELD_CAMERA_INTRINSICS, + FIELD_CAMERA_TRAJECTORY, +) +from lingbot.runtime import ( + FIELD_FIRST_FRAME_PATH, + FIELD_FPS, + FIELD_PIXEL_HEIGHT, + FIELD_PIXEL_WIDTH, + FIELD_PROMPT, + FIELD_TOTAL_BLOCKS, + inference_input_from_replay_inputs, +) +from lingbot.webrtc.session import LingbotRuntimeConfig + +from flashdreams.infra.video_output import VideoStepResult +from flashdreams.runtime import ( + CanonicalInputs, + InferenceConfig, + InferenceInput, + OutputArtifact, + OutputTarget, + 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.server import SESSION_MANAGER_KEY + +pytestmark = pytest.mark.ci_cpu + + +def _write_camera_assets(poses: Path, intrinsics: Path, *, frames: int = 64) -> None: + """Write real .npy camera assets; the input mapping loads them for real.""" + trajectory = np.tile(np.eye(4, dtype=np.float32), (frames, 1, 1)) + trajectory[:, 2, 3] = np.linspace(0.0, 1.0, frames, dtype=np.float32) + np.save(poses, trajectory) + np.save( + intrinsics, + np.tile( + np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (frames, 1) + ), + ) + + +def test_lingbot_demo_defaults_to_interactive_preset() -> None: + args = parse_args(["replay", "--output", "demo.mp4"]) + + assert args.preset_id == "lingbot-world-fast-taehv-window15-sink3" + + +def test_lingbot_demo_adapter_declares_mp4_and_webrtc_modes() -> 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") + fields = { + field.name + for field in adapter.inference_input_schema.global_conditioning_fields + } + assert "scenario" not in fields + assert { + FIELD_PROMPT, + FIELD_FIRST_FRAME_PATH, + FIELD_TOTAL_BLOCKS, + FIELD_PIXEL_HEIGHT, + FIELD_PIXEL_WIDTH, + 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 + } + assert step_fields == {FIELD_CAMERA_TRAJECTORY, FIELD_CAMERA_INTRINSICS} + + +def test_lingbot_replay_demo_uses_shared_runner(tmp_path: Path) -> None: + image = tmp_path / "image.jpg" + poses = tmp_path / "poses.npy" + intrinsics = tmp_path / "intrinsics.npy" + image.write_bytes(b"fake") + _write_camera_assets(poses, intrinsics) + pipeline_config = object() + adapter = LingbotDemoAdapter() + output = _RecordingOutputTarget() + calls: list[dict[str, Any]] = [] + + def fake_runner(**kwargs: Any) -> Sequence[OutputArtifact]: + calls.append(kwargs) + return (OutputArtifact(kind="video/mp4", uri="memory://lingbot"),) + + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + input_mode="replay", + scenario={ + "prompt": "drive through a city", + "image_path": image, + "pose_path": poses, + "intrinsic_path": intrinsics, + "total_blocks": 1, + }, + output=Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=16, output_layout="tchw"), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + runtime_options={"pipeline_config": pipeline_config}, + ), + ) + + artifacts = run_replay_demo( + spec=spec, + adapter=adapter, + output_target_factory=lambda output_spec: output, + runner=fake_runner, + ) + + assert artifacts == (OutputArtifact(kind="video/mp4", uri="memory://lingbot"),) + assert len(calls) == 1 + assert calls[0]["adapter"] is adapter + assert calls[0]["config"] == spec.config + inputs = calls[0]["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 + + +def test_lingbot_replay_invalid_scenario_fails_before_runtime_creation( + tmp_path: Path, +) -> None: + adapter = LingbotDemoAdapter( + replay_runtime_factory=lambda **kwargs: pytest.fail( + f"runtime should not be created: {kwargs}" + ) + ) + output_factory_calls = 0 + + def output_factory(output_spec: object) -> OutputTarget: + nonlocal output_factory_calls + del output_spec + output_factory_calls += 1 + return _RecordingOutputTarget() + + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + input_mode="replay", + scenario={ + "prompt": "drive", + "image_path": tmp_path / "missing.jpg", + "pose_path": tmp_path / "missing-poses.npy", + "intrinsic_path": tmp_path / "missing-intrinsics.npy", + }, + output=Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=16), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + runtime_options={"pipeline_config": object()}, + ), + ) + + with pytest.raises(FileNotFoundError, match=f"missing {FIELD_FIRST_FRAME_PATH}"): + run_replay_demo( + spec=spec, + adapter=adapter, + output_target_factory=output_factory, + ) + + assert output_factory_calls == 0 + + +def test_lingbot_replay_cli_defaults_to_example_data( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lingbot.runtime as runtime_module + + 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" + ) + (example_dir / "prompt.txt").write_text("drive through a forest\n") + downloaded: list[int] = [] + + def fake_download(*, is_rank_zero: bool, example_idx: int) -> Path: + assert is_rank_zero is True + downloaded.append(example_idx) + return example_dir + + monkeypatch.setattr( + runtime_module, + "ensure_example_data_downloaded", + fake_download, + ) + args = parse_args(["replay", "--output", str(tmp_path / "demo.mp4")]) + spec = _replay_spec(args) + + prepared = LingbotDemoAdapter().prepare_scenario(spec) + + inputs = prepared.initial_inputs.global_conditioning + assert downloaded == [0] + assert inputs[FIELD_FIRST_FRAME_PATH] == example_dir / "image.jpg" + assert inputs[FIELD_PROMPT] == "drive through a forest" + + +def test_lingbot_replay_cli_can_disable_example_data(tmp_path: Path) -> None: + args = parse_args( + ["replay", "--no-example-data", "--output", str(tmp_path / "demo.mp4")] + ) + spec = _replay_spec(args) + + with pytest.raises(ValueError, match=f"require {FIELD_FIRST_FRAME_PATH}"): + LingbotDemoAdapter().prepare_scenario(spec) + + +def test_lingbot_replay_runtime_generates_video_step_result( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lingbot.runtime as runtime_module + + image = tmp_path / "image.jpg" + poses = tmp_path / "poses.npy" + intrinsics = tmp_path / "intrinsics.npy" + image.write_bytes(b"fake") + _write_camera_assets(poses, intrinsics, frames=16) + pipeline = _FakeLingbotPipeline() + monkeypatch.setattr( + runtime_module, + "load_first_frame_tensor", + lambda *args, **kwargs: torch.zeros(1, 3, 2, 2), + ) + + runtime = LingbotReplayRuntime( + config=InferenceConfig(model_id=LINGBOT_MODEL_ID, device="cpu"), + options=LingbotReplayRuntimeOptions( + pipeline_config=object(), + pipeline_factory=lambda pipeline_config, device: pipeline, + ), + ) + replay_inputs = LingbotReplayInputs( + prompt="drive", + first_frame_path=image, + camera_poses_path=poses, + camera_intrinsics_path=intrinsics, + total_blocks=1, + pixel_height=2, + pixel_width=2, + fps=16, + ) + adapter = LingbotDemoAdapter() + mapping = adapter.create_input_mapping(replay_inputs) + initial_inputs = mapping.map_global_conditioning_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=inference_input_from_replay_inputs(replay_inputs), + ) + session = runtime.start_session(initial_inputs) + + request = session.next_step_request() + assert request is not None + assert request.step_index == 0 + # The session asks for exactly this chunk's slice of the input timeline. + assert request.user_input_window is not None + assert request.user_input_window.start_s == 0.0 + assert request.user_input_window.end_s == 1 / 16 + assert request.metadata["num_frames"] == 1 + + step_inputs = mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput(), + request=request, + ) + assert step_inputs.step[FIELD_CAMERA_TRAJECTORY].shape == (1, 4, 4) + assert step_inputs.step[FIELD_CAMERA_INTRINSICS].shape == (1, 4) + result = session.step(step_inputs) + + 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 result.output_window is not None + assert result.output_window.start_s == 0.0 + assert result.output_window.end_s == 1 / 16 + assert result.metrics["denoise_s"] == 0.25 + assert session.next_step_request() is None + assert pipeline.initialize_cache_calls == [ + {"text": ["drive"], "image_shape": (1, 3, 2, 2)} + ] + assert pipeline.generate_calls == [ + { + "autoregressive_index": 0, + "intrinsics_shape": (1, 4), + "poses_shape": (1, 4, 4), + "world_scale": pytest.approx(mapping.camera_trace.world_scale), + } + ] + runtime.close() + + +def test_lingbot_webrtc_cli_builds_keyboard_driving_spec() -> None: + args = parse_args( + [ + "webrtc", + "--host", + "127.0.0.1", + "--port", + "9090", + "--device", + "cuda:2", + "--seed", + "123", + "--no-compile", + "--fps", + "12", + "--video-height", + "32", + "--video-width", + "64", + "--warmup-chunks", + "0", + "--warmup-timeout-s", + "1.5", + "--client-liveness-timeout-s", + "2.5", + "--prefer-sw-encoder", + "--example-idx", + "2", + ] + ) + + spec = _webrtc_spec(args, device="cuda:3", context_parallel_size=4) + + assert spec.model_id == LINGBOT_MODEL_ID + assert spec.preset_id == DEFAULT_LINGBOT_PRESET + assert spec.input_mode == "keyboard-driving" + assert isinstance(spec.scenario, LingbotWebRTCScenario) + assert spec.scenario.example_idx == 2 + assert spec.scenario.prefer_sw_encoder is True + assert isinstance(spec.output, WebRTCOutputSpec) + assert spec.output.host == "127.0.0.1" + assert spec.output.port == 9090 + assert spec.output.fps == 12 + assert spec.output.video_width == 64 + assert spec.output.video_height == 32 + assert spec.output.warmup_chunks == 0 + assert spec.output.warmup_timeout_s == 1.5 + assert spec.output.client_liveness_timeout_s == 2.5 + assert spec.config is not None + assert spec.config.device == "cuda:3" + assert spec.config.compile is False + assert spec.config.runtime_options["seed"] == 123 + assert spec.config.runtime_options["context_parallel_size"] == 4 + assert spec.config.runtime_options["example_idx"] == 2 + + +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, + input_mode="keyboard-driving", + scenario=LingbotWebRTCScenario(example_idx=2, prefer_sw_encoder=True), + output=WebRTCOutputSpec( + host="0.0.0.0", + port=8080, + fps=24, + video_width=64, + video_height=32, + warmup_chunks=0, + warmup_timeout_s=1.0, + ), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + device="cuda:7", + runtime_options={"pipeline_config": pipeline_config, "seed": 123}, + ), + ) + + 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 + + +def test_lingbot_webrtc_demo_installs_model_routes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lingbot.demo.webrtc as demo_webrtc_module + + app_calls: list[dict[str, Any]] = [] + + async def _ok(request: web.Request) -> web.Response: + del request + return web.Response(text="ok") + + def fake_create_app(**kwargs: Any) -> web.Application: + app_calls.append(kwargs) + app = web.Application() + app[SESSION_MANAGER_KEY] = kwargs["session_manager"] + app.router.add_get("/api/session/initial_scene", _ok) + app.router.add_get("/api/session/first_frame", _ok) + app.router.add_post("/api/session/input", _ok) + return app + + monkeypatch.setattr(demo_webrtc_module, "create_app", fake_create_app) + adapter = LingbotDemoAdapter(webrtc_runtime_factory=_FakeWebRTCRuntime) + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + input_mode="keyboard-driving", + scenario=LingbotWebRTCScenario(), + output=WebRTCOutputSpec( + host="0.0.0.0", + port=8080, + warmup_timeout_s=1.0, + preload_name="Test Lingbot", + ), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + runtime_options={"pipeline_config": object()}, + ), + ) + + demo = build_webrtc_demo(spec=spec, adapter=adapter, create_app=True) + + assert demo.app is not None + assert app_calls[0]["session_manager"] is demo.session_manager + assert app_calls[0]["request_session_url"] == ( + "http://127.0.0.1:8080/request_session" + ) + route_paths = {resource.canonical for resource in demo.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 + + +def test_lingbot_webrtc_demo_serves_through_shared_runner( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lingbot.demo.webrtc as demo_webrtc_module + + server_calls: list[dict[str, Any]] = [] + + def fake_create_app(**kwargs: Any) -> web.Application: + app = web.Application() + app[SESSION_MANAGER_KEY] = kwargs["session_manager"] + return app + + def fake_server_runner(**kwargs: Any) -> None: + server_calls.append(kwargs) + + monkeypatch.setattr(demo_webrtc_module, "create_app", fake_create_app) + adapter = LingbotDemoAdapter(webrtc_runtime_factory=_FakeWebRTCRuntime) + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + input_mode="keyboard-driving", + scenario={"example_idx": 0}, + output=WebRTCOutputSpec( + host="0.0.0.0", + port=8080, + warmup_timeout_s=1.0, + ), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + runtime_options={"pipeline_config": object()}, + ), + ) + + demo = cast( + WebRTCDemo, + serve_flashdreams_demo( + spec=spec, + adapter=adapter, + world_rank=0, + 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]["host"] == "0.0.0.0" + assert server_calls[0]["port"] == 8080 + assert isinstance(demo.session_manager, LingbotDemoWebRTCSessionManager) + + +class _RecordingOutputTarget: + def open(self) -> None: + return None + + def write(self, result: StepResult) -> None: + del result + + def close(self) -> Sequence[OutputArtifact]: + return () + + +class _FakeLingbotPipeline: + def __init__(self) -> None: + self.initialize_cache_calls: list[dict[str, Any]] = [] + self.generate_calls: list[dict[str, Any]] = [] + + def initialize_cache(self, *, text: list[str], image: torch.Tensor) -> object: + self.initialize_cache_calls.append( + { + "text": text, + "image_shape": tuple(image.shape), + } + ) + return object() + + def get_num_output_frames(self, autoregressive_index: int) -> int: + del autoregressive_index + return 1 + + def generate( + self, + *, + autoregressive_index: int, + cache: object, + input: Any, + ) -> torch.Tensor: + del cache + self.generate_calls.append( + { + "autoregressive_index": autoregressive_index, + "intrinsics_shape": tuple(input.intrinsics.shape), + "poses_shape": tuple(input.poses.shape), + "world_scale": input.world_scale, + } + ) + return torch.full((1, 3, 2, 2), float(autoregressive_index)) + + def finalize(self, *, autoregressive_index: int, cache: object) -> dict[str, float]: + del autoregressive_index, cache + return {"denoise_s": 0.25} + + +class _FakeWebRTCRuntime: + def __init__(self, config: LingbotRuntimeConfig) -> None: + self.config = config + + async def initialize(self) -> None: + return 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_next_chunk_num_frames(self) -> int: + return 1 + + async def generate_chunk( + self, + *, + segments: list[Any], + frame_times: list[float], + ) -> Any: + del segments, frame_times + return None + + async def close(self) -> None: + return None + + def send_exit_signal(self) -> None: + return None + + def wait_for_termination(self) -> None: + return None diff --git a/integrations/lingbot/tests/test_input_mapping.py b/integrations/lingbot/tests/test_input_mapping.py new file mode 100644 index 000000000..1db76bf44 --- /dev/null +++ b/integrations/lingbot/tests/test_input_mapping.py @@ -0,0 +1,448 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import torch +from lingbot.demo.spec import resolve_text_event_prompts, resolve_user_input_events +from lingbot.input_mapping import ( + CAMERA_COMMAND, + FIELD_CAMERA_INTRINSICS, + FIELD_CAMERA_TRAJECTORY, + TEXT_EVENT, + KeyboardToCameraCommand, + LingbotInputMapping, + TextEventSelection, + load_camera_trace, +) + +from flashdreams.runtime import ( + CanonicalInputs, + InferenceInput, + InputCanonicalizer, + StepRequest, + TimeWindow, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, +) + +pytestmark = pytest.mark.ci_cpu + +_KEYBOARD_SOURCE = UserInputSchema( + capabilities=( + 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"}) + ), + ) +) + + +def _step_request(*, step_index: int, frame_start: int, num_frames: int, fps: int = 16): + return StepRequest( + step_index=step_index, + user_input_window=TimeWindow( + start_s=frame_start / fps, + end_s=(frame_start + num_frames) / fps, + ), + metadata={"num_frames": num_frames, "frame_start": frame_start}, + ) + + +def _live_mapping(**kwargs) -> LingbotInputMapping: + return LingbotInputMapping( + fps=16, + base_intrinsics=torch.tensor([416.0, 416.0, 416.0, 240.0]), + world_scale=1.0, + **kwargs, + ) + + +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"}), + ) + ) + window = TimeWindow(start_s=0.0, end_s=1.0) + + value = converter.convert(inputs.window(window), window) + + assert value is not None + assert value["move_forward"] == 1.0 + assert value["yaw"] == 0.0 + # Level-triggered: a key held across the next window still means forward. + next_window = TimeWindow(start_s=1.0, end_s=2.0) + held = converter.convert(UserInputs().window(next_window), next_window) + assert held is not None + assert held["move_forward"] == 1.0 + + +def test_camera_command_segments_preserve_sub_window_timing() -> None: + converter = KeyboardToCameraCommand() + 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"}), + ) + ) + + value = converter.convert(inputs.window(window), window) + + assert value is not None + segments = value["segments"] + assert [(start, end) for start, end, _ in segments] == [(0.0, 0.5), (0.5, 1.0)] + assert segments[0][2]["move_forward"] == 0.0 + assert segments[1][2]["move_forward"] == 1.0 + + +def test_key_events_drive_a_camera_trajectory() -> None: + canonicalizer = InputCanonicalizer([KeyboardToCameraCommand()]) + mapping = _live_mapping() + user_inputs = UserInputs( + events=( + UserInputEvent(timestamp_s=0.0, event_type="key_down", payload={"key": "w"}), + ) + ) + request = _step_request(step_index=0, frame_start=0, num_frames=4) + assert request.user_input_window is not None + + step_inputs = mapping.map_step_inputs( + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=request.user_input_window, + source_schema=_KEYBOARD_SOURCE, + ), + inference_input=InferenceInput(), + request=request, + ) + + poses = step_inputs.step[FIELD_CAMERA_TRAJECTORY] + assert poses.shape == (4, 4, 4) + assert step_inputs.step[FIELD_CAMERA_INTRINSICS].shape == (4, 4) + # Holding forward has to actually move the camera along the trajectory. + assert not torch.allclose(poses[0], poses[-1]) + assert poses[-1][:3, 3].abs().sum() > 0 + + +def test_idle_keyboard_leaves_the_camera_stationary() -> None: + canonicalizer = InputCanonicalizer([KeyboardToCameraCommand()]) + mapping = _live_mapping() + request = _step_request(step_index=0, frame_start=0, num_frames=4) + assert request.user_input_window is not None + + step_inputs = mapping.map_step_inputs( + canonical_inputs=canonicalizer.canonicalize( + UserInputs(), + window=request.user_input_window, + source_schema=_KEYBOARD_SOURCE, + ), + inference_input=InferenceInput(), + request=request, + ) + + poses = step_inputs.step[FIELD_CAMERA_TRAJECTORY] + assert torch.allclose(poses[0], poses[-1]) + + +def test_text_event_becomes_a_global_conditioning_prompt_update() -> None: + canonicalizer = InputCanonicalizer( + [KeyboardToCameraCommand(), TextEventSelection()] + ) + mapping = _live_mapping(text_event_prompts={"storm": "a violent storm"}) + mapping.set_base_prompt("a calm street") + user_inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, + event_type="text_event", + payload={"event_id": "storm"}, + ), + ) + ) + request = _step_request(step_index=0, frame_start=0, num_frames=4) + assert request.user_input_window is not None + + step_inputs = mapping.map_step_inputs( + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=request.user_input_window, + source_schema=_KEYBOARD_SOURCE, + ), + inference_input=InferenceInput(), + request=request, + ) + + assert step_inputs.global_conditioning["prompt"] == "a violent storm" + + # The swap is requested once, not re-sent on every later step. + next_request = _step_request(step_index=1, frame_start=4, num_frames=4) + assert next_request.user_input_window is not None + held = mapping.map_step_inputs( + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=next_request.user_input_window, + source_schema=_KEYBOARD_SOURCE, + ), + inference_input=InferenceInput(), + request=next_request, + ) + assert held.global_conditioning == {} + + +def test_clearing_a_text_event_restores_the_base_prompt() -> None: + converter = TextEventSelection() + window = TimeWindow(start_s=0.0, end_s=1.0) + triggered = converter.convert( + UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, + event_type="text_event", + payload={"event_id": "storm"}, + ), + ) + ), + window, + ) + assert triggered is not None and triggered["event_id"] == "storm" + + cleared = converter.convert( + UserInputs( + events=( + UserInputEvent( + timestamp_s=0.5, + event_type="text_event", + payload={"event_id": "storm", "state": "clear"}, + ), + ) + ), + window, + ) + assert cleared is not None and cleared["event_id"] is None + + +def test_unknown_text_event_is_rejected_by_the_mapping() -> None: + canonicalizer = InputCanonicalizer( + [KeyboardToCameraCommand(), TextEventSelection()] + ) + mapping = _live_mapping(text_event_prompts={"storm": "a violent storm"}) + request = _step_request(step_index=0, frame_start=0, num_frames=4) + assert request.user_input_window is not None + canonical = canonicalizer.canonicalize( + UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, + event_type="text_event", + payload={"event_id": "volcano"}, + ), + ) + ), + window=request.user_input_window, + source_schema=_KEYBOARD_SOURCE, + ) + + with pytest.raises(ValueError, match="Unknown Lingbot text event_id"): + mapping.map_step_inputs( + canonical_inputs=canonical, + inference_input=InferenceInput(), + request=request, + ) + + +def test_live_mapping_requires_camera_command_from_the_source() -> None: + mapping = _live_mapping() + + with pytest.raises(ValueError, match="requires a 'camera_command'"): + mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput(), + request=_step_request(step_index=0, frame_start=0, num_frames=4), + ) + + +def test_trace_mapping_slices_successive_chunks(tmp_path: Path) -> None: + poses_path = tmp_path / "poses.npy" + intrinsics_path = tmp_path / "intrinsics.npy" + trajectory = np.tile(np.eye(4, dtype=np.float32), (32, 1, 1)) + trajectory[:, 2, 3] = np.arange(32, dtype=np.float32) + np.save(poses_path, trajectory) + np.save( + intrinsics_path, + np.tile(np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (32, 1)), + ) + trace = load_camera_trace( + camera_poses_path=poses_path, + camera_intrinsics_path=intrinsics_path, + pixel_height=464, + pixel_width=832, + intrinsics_reference_height=480, + intrinsics_reference_width=832, + ) + mapping = LingbotInputMapping(fps=16, trace=trace) + + first = mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput(), + request=_step_request(step_index=0, frame_start=0, num_frames=4), + ) + second = mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput(), + request=_step_request(step_index=1, frame_start=4, num_frames=4), + ) + + assert first.step[FIELD_CAMERA_TRAJECTORY].shape == (4, 4, 4) + # Consecutive steps must advance through the trace, not restart it. + assert not torch.allclose( + first.step[FIELD_CAMERA_TRAJECTORY], second.step[FIELD_CAMERA_TRAJECTORY] + ) + assert mapping.mapping_schema.consumes == () + + +def test_trace_mapping_reports_running_past_the_end(tmp_path: Path) -> None: + poses_path = tmp_path / "poses.npy" + intrinsics_path = tmp_path / "intrinsics.npy" + np.save(poses_path, np.tile(np.eye(4, dtype=np.float32), (16, 1, 1))) + np.save( + intrinsics_path, + np.tile(np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (16, 1)), + ) + mapping = LingbotInputMapping( + fps=16, + trace=load_camera_trace( + camera_poses_path=poses_path, + camera_intrinsics_path=intrinsics_path, + pixel_height=464, + pixel_width=832, + intrinsics_reference_height=480, + intrinsics_reference_width=832, + ), + ) + + with pytest.raises(ValueError, match="camera trace has"): + mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput(), + request=_step_request(step_index=0, frame_start=0, num_frames=999), + ) + + +def test_scenario_event_trace_resolves_into_user_inputs() -> None: + user_inputs = resolve_user_input_events( + { + "events": [ + {"t": 1.5, "type": "key_down", "key": "a"}, + {"t": 0.0, "type": "key_down", "key": "w"}, + {"t": 2.0, "type": "text_event", "event_id": "storm"}, + ] + } + ) + + # UserInputs requires non-decreasing timestamps, so resolution must sort. + assert [event.timestamp_s for event in user_inputs.events] == [0.0, 1.5, 2.0] + assert user_inputs.events[0].payload == {"key": "w"} + assert user_inputs.events[2].event_type == "text_event" + + +def test_scenario_text_event_catalog_resolves() -> None: + assert resolve_text_event_prompts({"text_events": {"storm": "a storm"}}) == { + "storm": "a storm" + } + assert resolve_text_event_prompts( + {"text_events": [{"event_id": "portal", "prompt": "a glowing portal"}]} + ) == {"portal": "a glowing portal"} + assert resolve_text_event_prompts(None) == {} + + +def test_declared_modalities_match_what_the_converters_produce() -> None: + canonicalizer = InputCanonicalizer( + [KeyboardToCameraCommand(), TextEventSelection()] + ) + + schema = canonicalizer.canonical_schema(_KEYBOARD_SOURCE) + + assert schema.supports(CAMERA_COMMAND) + assert schema.supports(TEXT_EVENT) + # A source with no key events cannot feed the keyboard converter. + empty = canonicalizer.canonical_schema(UserInputSchema()) + assert not empty.supports(CAMERA_COMMAND) + + +def test_event_driven_scenario_builds_a_live_mapping(tmp_path: Path) -> None: + """A scenario can drive the camera from events instead of the pose trace.""" + from lingbot.demo.adapter import LingbotDemoAdapter + from lingbot.runtime import LINGBOT_MODEL_ID + + from flashdreams.runtime import InferenceConfig + from flashdreams.runtime.demo import DemoSpec, Mp4OutputSpec + + image = tmp_path / "image.jpg" + image.write_bytes(b"fake") + poses_path = tmp_path / "poses.npy" + intrinsics_path = tmp_path / "intrinsics.npy" + np.save(poses_path, np.tile(np.eye(4, dtype=np.float32), (32, 1, 1))) + np.save( + intrinsics_path, + np.tile(np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (32, 1)), + ) + + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + input_mode="replay", + scenario={ + "prompt": "a calm street", + "image_path": image, + "pose_path": poses_path, + "intrinsic_path": intrinsics_path, + "camera_source": "events", + "text_events": {"storm": "a violent storm"}, + "events": [ + {"t": 0.0, "type": "key_down", "key": "w"}, + {"t": 0.2, "type": "text_event", "event_id": "storm"}, + ], + }, + output=Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=16), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + runtime_options={"pipeline_config": object()}, + ), + ) + + prepared = LingbotDemoAdapter().prepare_scenario(spec) + + assert prepared.mapping is not None + 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 + # canonicalizer silently drops the keyboard converter. + canonical_schema = prepared.canonicalizer.canonical_schema(prepared.source_schema) + assert canonical_schema.supports(CAMERA_COMMAND) + assert canonical_schema.supports(TEXT_EVENT) + + request = _step_request(step_index=0, frame_start=0, num_frames=4) + assert request.user_input_window is not None + step_inputs = prepared.mapping.map_step_inputs( + canonical_inputs=prepared.canonicalizer.canonicalize( + prepared.user_inputs, + window=request.user_input_window, + source_schema=prepared.source_schema, + ), + inference_input=InferenceInput(), + request=request, + ) + poses = step_inputs.step[FIELD_CAMERA_TRAJECTORY] + assert poses.shape == (4, 4, 4) + assert not torch.allclose(poses[0], poses[-1]) + assert step_inputs.global_conditioning["prompt"] == "a violent storm" diff --git a/integrations/lingbot/tests/test_keyboard_parity.py b/integrations/lingbot/tests/test_keyboard_parity.py new file mode 100644 index 000000000..b233d7266 --- /dev/null +++ b/integrations/lingbot/tests/test_keyboard_parity.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Live camera control must match the WebRTC path it will eventually replace. + +The WebRTC runtime drives the camera with ``KeyboardResampler.sample_chunk`` +feeding ``CameraPoseIntegrator``. The runtime-API path instead windows events +with ``StepRequest.user_input_window``, canonicalizes them into camera intent, +and integrates that. Both should produce the same trajectory for the same key +stream; these tests pin that, so a divergence shows up here rather than as +different handling in a live session. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from lingbot.input_mapping import ( + FIELD_CAMERA_TRAJECTORY, + KeyboardToCameraCommand, + LingbotInputMapping, +) + +from flashdreams.runtime import ( + InferenceInput, + InputCanonicalizer, + StepRequest, + TimeWindow, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, +) +from flashdreams.serving.webrtc.controls import CameraPoseIntegrator, KeyboardResampler + +pytestmark = pytest.mark.ci_cpu + +_FPS = 16 +_NUM_FRAMES = 4 + +_SOURCE = UserInputSchema( + capabilities=( + UserInputCapability(event_type="key_down", payload_fields=frozenset({"key"})), + UserInputCapability(event_type="key_up", payload_fields=frozenset({"key"})), + ) +) + +# (timestamp_s, event_type, key) +Edge = tuple[float, str, str] + +_STREAMS: dict[str, list[Edge]] = { + "hold_forward": [(0.0, "key_down", "w")], + "forward_then_release": [(0.0, "key_down", "w"), (0.1, "key_up", "w")], + "mid_chunk_turn": [(0.0, "key_down", "w"), (0.13, "key_down", "a")], + "strafe_and_pitch": [(0.02, "key_down", "e"), (0.09, "key_down", "i")], + "alternate_yaw_keys": [(0.0, "key_down", "w"), (0.05, "key_down", "j")], + "conflicting_yaw": [(0.0, "key_down", "a"), (0.07, "key_down", "d")], + "rapid_toggle": [ + (0.01, "key_down", "w"), + (0.04, "key_up", "w"), + (0.08, "key_down", "w"), + (0.2, "key_up", "w"), + ], + "idle": [], + # KeyboardResampler drains events with `event_t <= chunk_end` while + # TimeWindow is half-open, so an edge landing exactly on a chunk boundary + # is the most likely place for the two paths to disagree. + "exact_chunk_boundary": [(0.0, "key_down", "w"), (0.25, "key_down", "a")], + "boundary_release": [(0.0, "key_down", "w"), (0.25, "key_up", "w")], + "second_boundary": [(0.0, "key_down", "w"), (0.5, "key_down", "d")], +} + + +def _legacy_poses(edges: list[Edge], *, chunks: int) -> np.ndarray: + """Integrate a key stream the way the WebRTC session does.""" + resampler = KeyboardResampler(fps=_FPS, start_v=0.0) + integrator = CameraPoseIntegrator() + for timestamp_s, event_type, key in edges: + resampler.on_edge( + arrival_t=timestamp_s, + event="keydown" if event_type == "key_down" else "keyup", + key=key, + ) + poses = [] + for _ in range(chunks): + segments, frame_times = resampler.sample_chunk(_NUM_FRAMES) + poses.append( + integrator.integrate_chunk(segments=segments, frame_times=frame_times) + ) + return np.concatenate(poses) + + +def _runtime_api_poses(edges: list[Edge], *, chunks: int) -> np.ndarray: + """Integrate the same key stream through the runtime API input path.""" + canonicalizer = InputCanonicalizer([KeyboardToCameraCommand()]) + mapping = LingbotInputMapping( + fps=_FPS, + base_intrinsics=torch.tensor([416.0, 416.0, 416.0, 240.0]), + world_scale=1.0, + ) + user_inputs = UserInputs( + events=tuple( + UserInputEvent( + timestamp_s=timestamp_s, + event_type=event_type, + payload={"key": key}, + ) + for timestamp_s, event_type, key in edges + ) + ) + poses = [] + for chunk_index in range(chunks): + frame_start = chunk_index * _NUM_FRAMES + request = StepRequest( + step_index=chunk_index, + user_input_window=TimeWindow( + start_s=frame_start / _FPS, + end_s=(frame_start + _NUM_FRAMES) / _FPS, + ), + metadata={"num_frames": _NUM_FRAMES, "frame_start": frame_start}, + ) + assert request.user_input_window is not None + step_inputs = mapping.map_step_inputs( + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=request.user_input_window, + source_schema=_SOURCE, + ), + inference_input=InferenceInput(), + request=request, + ) + poses.append(step_inputs.step[FIELD_CAMERA_TRAJECTORY].numpy()) + return np.concatenate(poses) + + +@pytest.mark.parametrize("name", sorted(_STREAMS)) +def test_single_chunk_matches_the_webrtc_path(name: str) -> None: + edges = _STREAMS[name] + + legacy = _legacy_poses(edges, chunks=1) + runtime_api = _runtime_api_poses(edges, chunks=1) + + assert legacy.shape == runtime_api.shape + np.testing.assert_allclose(runtime_api, legacy, atol=1e-5) + + +@pytest.mark.parametrize("name", sorted(_STREAMS)) +def test_multi_chunk_matches_the_webrtc_path(name: str) -> None: + """Carried key state across chunk boundaries must agree too.""" + edges = _STREAMS[name] + + legacy = _legacy_poses(edges, chunks=3) + runtime_api = _runtime_api_poses(edges, chunks=3) + + assert legacy.shape == runtime_api.shape + np.testing.assert_allclose(runtime_api, legacy, atol=1e-5) diff --git a/integrations/lingbot/tests/test_runtime_gpu.py b/integrations/lingbot/tests/test_runtime_gpu.py new file mode 100644 index 000000000..db88ccd3e --- /dev/null +++ b/integrations/lingbot/tests/test_runtime_gpu.py @@ -0,0 +1,259 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +import torch +from lingbot import runtime as runtime_module +from lingbot.runtime import ( + LINGBOT_MODEL_ID, + LingbotModelAdapter, + LingbotReplayInputs, + LingbotReplayRuntime, + LingbotReplayRuntimeOptions, + inference_input_from_replay_inputs, +) + +from flashdreams.infra.video_output import VideoStepResult +from flashdreams.runtime import CanonicalInputs, InferenceConfig, InferenceInput + +pytestmark = pytest.mark.ci_gpu + + +def test_lingbot_replay_runtime_accepts_direct_inputs_on_cuda( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Exercise the migrated Lingbot runtime API path with CUDA tensors.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA required.") + + image = tmp_path / "image.jpg" + poses = tmp_path / "poses.npy" + intrinsics = tmp_path / "intrinsics.npy" + image.write_bytes(b"fake") + trajectory = np.tile(np.eye(4, dtype=np.float32), (32, 1, 1)) + trajectory[:, 2, 3] = np.arange(32, dtype=np.float32) + np.save(poses, trajectory) + np.save( + intrinsics, + np.tile(np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (32, 1)), + ) + pipeline = _FakeCudaLingbotPipeline() + + def _fake_load_first_frame_tensor( + path: Path, + **kwargs: Any, + ) -> torch.Tensor: + assert path == image + return torch.zeros( + (1, 3, 2, 2), + device=kwargs["device"], + dtype=kwargs["dtype"], + ) + + monkeypatch.setattr( + runtime_module, + "load_first_frame_tensor", + _fake_load_first_frame_tensor, + ) + runtime = LingbotReplayRuntime( + config=InferenceConfig(model_id=LINGBOT_MODEL_ID, device="cuda"), + options=LingbotReplayRuntimeOptions( + pipeline_config=object(), + pipeline_factory=lambda _pipeline_config, _device: pipeline, + ), + ) + replay_inputs = LingbotReplayInputs( + prompt="drive", + first_frame_path=image, + camera_poses_path=poses, + camera_intrinsics_path=intrinsics, + total_blocks=1, + pixel_height=2, + pixel_width=2, + fps=16, + ) + # Camera inputs now reach the session per step through the mapping, so the + # GPU path has to be driven the same way the standard loop drives it. + mapping = LingbotModelAdapter().create_input_mapping(replay_inputs) + session = runtime.start_session( + mapping.map_global_conditioning_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=inference_input_from_replay_inputs(replay_inputs), + ) + ) + try: + request = session.next_step_request() + assert request is not None + result = session.step( + mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput(), + request=request, + ) + ) + torch.cuda.synchronize() + finally: + session.close() + 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 pipeline.initialize_cache_devices == ["cuda"] + assert pipeline.generate_world_scales == [mapping.camera_trace.world_scale] + + +class _FakeCudaLingbotPipeline: + def __init__(self) -> None: + self.initialize_cache_devices: list[str] = [] + self.generate_world_scales: list[float] = [] + + def initialize_cache(self, *, text: list[str], image: torch.Tensor) -> object: + assert text == ["drive"] + assert image.is_cuda + self.initialize_cache_devices.append(image.device.type) + return object() + + def get_num_output_frames(self, autoregressive_index: int) -> int: + assert autoregressive_index == 0 + return 1 + + def generate( + self, + *, + autoregressive_index: int, + cache: object, + input: Any, + ) -> torch.Tensor: + del cache + assert autoregressive_index == 0 + assert input.intrinsics.is_cuda + assert input.poses.is_cuda + self.generate_world_scales.append(input.world_scale) + return torch.zeros( + (1, 3, 2, 2), + device=input.intrinsics.device, + dtype=torch.bfloat16, + ) + + def finalize(self, *, autoregressive_index: int, cache: object) -> dict[str, float]: + del autoregressive_index, cache + return {"denoise_s": 0.25} + + +def test_event_driven_camera_control_on_cuda( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Drive the CUDA session from key events instead of a fixed pose trace.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA required.") + + from lingbot.input_mapping import ( + FIELD_CAMERA_TRAJECTORY, + KeyboardToCameraCommand, + ) + + from flashdreams.runtime import ( + InputCanonicalizer, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, + ) + + image = tmp_path / "image.jpg" + image.write_bytes(b"fake") + pipeline = _FakeCudaLingbotPipeline() + + def _fake_load_first_frame_tensor(path: Path, **kwargs: Any) -> torch.Tensor: + del path + return torch.zeros((1, 3, 2, 2), device=kwargs["device"], dtype=kwargs["dtype"]) + + monkeypatch.setattr( + runtime_module, + "load_first_frame_tensor", + _fake_load_first_frame_tensor, + ) + + runtime = LingbotReplayRuntime( + config=InferenceConfig(model_id=LINGBOT_MODEL_ID, device="cuda"), + options=LingbotReplayRuntimeOptions( + pipeline_config=object(), + pipeline_factory=lambda _pipeline_config, _device: pipeline, + ), + ) + adapter = LingbotModelAdapter() + mapping = adapter.create_live_input_mapping( + fps=16, + base_intrinsics=torch.tensor([416.0, 416.0, 416.0, 240.0]), + world_scale=1.0, + prompt="drive", + ) + initial_inputs = mapping.map_global_conditioning_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput( + global_conditioning={ + "prompt": "drive", + "first_frame_path": image, + "total_blocks": 1, + "pixel_height": 2, + "pixel_width": 2, + "fps": 16, + } + ), + ) + canonicalizer = InputCanonicalizer([KeyboardToCameraCommand()]) + source = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="key_down", payload_fields=frozenset({"key"}) + ), + UserInputCapability( + event_type="key_up", payload_fields=frozenset({"key"}) + ), + ) + ) + user_inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, event_type="key_down", payload={"key": "w"} + ), + ) + ) + + session = runtime.start_session(initial_inputs) + try: + request = session.next_step_request() + assert request is not None + assert request.user_input_window is not None + step_inputs = mapping.map_step_inputs( + canonical_inputs=canonicalizer.canonicalize( + user_inputs, + window=request.user_input_window, + source_schema=source, + ), + inference_input=InferenceInput(), + request=request, + ) + # Holding forward must produce real motion before it reaches the model. + # This chunk is one frame, so compare against the identity start pose + # rather than across frames: 0.8 m/s at 16fps advances 0.05 along +z. + poses = step_inputs.step[FIELD_CAMERA_TRAJECTORY] + assert poses[-1][2, 3].item() == pytest.approx(0.05, abs=1e-4) + result = session.step(step_inputs) + torch.cuda.synchronize() + finally: + session.close() + runtime.close() + + assert result.output.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 new file mode 100644 index 000000000..def9cfb9e --- /dev/null +++ b/integrations/lingbot/tests/test_runtime_session_inputs.py @@ -0,0 +1,395 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The replay session must consume its per-step inputs, not ignore them.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +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, + LingbotReplayRuntime, + LingbotReplayRuntimeOptions, + LingbotSessionInputs, +) + +from flashdreams.runtime import InferenceConfig, InferenceInput + +pytestmark = pytest.mark.ci_cpu + + +class _FakePipeline: + """Records what the session hands the model.""" + + def __init__(self, *, supports_text_swap: bool = True) -> None: + self.generate_calls: list[dict[str, Any]] = [] + self.text_encoder_calls: list[list[str]] = [] + self.encoders_loaded = 0 + self.diffusion_model = _FakeDiffusionModel( + supports_text_swap=supports_text_swap + ) + + def initialize_cache(self, *, text: list[str], image: torch.Tensor) -> Any: + del text, image + return _FakeCache() + + def get_num_output_frames(self, autoregressive_index: int) -> int: + del autoregressive_index + return 2 + + def generate( + self, + *, + autoregressive_index: int, + cache: object, + input: Any, + ) -> torch.Tensor: + del cache + self.generate_calls.append( + { + "autoregressive_index": autoregressive_index, + "poses": input.poses.clone(), + "intrinsics": input.intrinsics.clone(), + "world_scale": input.world_scale, + } + ) + return torch.zeros(2, 3, 2, 2) + + def finalize(self, *, autoregressive_index: int, cache: object) -> dict[str, float]: + del autoregressive_index, cache + return {"denoise_s": 0.1} + + def _ensure_oneshot_encoders_loaded(self) -> None: + self.encoders_loaded += 1 + + def text_encoder(self, texts: list[str]) -> torch.Tensor: + self.text_encoder_calls.append(list(texts)) + return torch.ones(1, 4) + + +class _FakeCache: + def __init__(self) -> None: + self.transformer_cache = object() + + +class _FakeTransformer: + def __init__(self) -> None: + self.replaced: list[torch.Tensor] = [] + + def replace_text_embeddings(self, cache: object, embeddings: torch.Tensor) -> None: + del cache + self.replaced.append(embeddings) + + +class _FakeDiffusionModel: + def __init__(self, *, supports_text_swap: bool) -> None: + self.transformer = _FakeTransformer() if supports_text_swap else object() + + +def _session( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + pipeline: _FakePipeline, + *, + total_blocks: int = 2, + total_camera_frames: int | None = None, +): + monkeypatch.setattr( + runtime_module, + "load_first_frame_tensor", + lambda *args, **kwargs: torch.zeros(1, 3, 2, 2), + ) + runtime = LingbotReplayRuntime( + config=InferenceConfig(model_id=LINGBOT_MODEL_ID, device="cpu"), + options=LingbotReplayRuntimeOptions( + pipeline_config=object(), + pipeline_factory=lambda pipeline_config, device: pipeline, + ), + ) + image = tmp_path / "image.jpg" + image.write_bytes(b"fake") + session = runtime_module.LingbotReplaySession( + pipeline=pipeline, + session_inputs=LingbotSessionInputs( + prompt="a calm street", + first_frame_path=image, + total_blocks=total_blocks, + pixel_height=2, + pixel_width=2, + fps=16, + world_scale=2.5, + total_camera_frames=total_camera_frames, + ), + device=torch.device("cpu"), + is_rank_zero=True, + output_layout="tchw", + ) + return runtime, session + + +def _step_payload(value: float) -> InferenceInput: + poses = torch.eye(4).repeat(2, 1, 1) + poses[:, 2, 3] = value + return InferenceInput( + step={ + FIELD_CAMERA_TRAJECTORY: poses, + FIELD_CAMERA_INTRINSICS: torch.full((2, 4), value), + } + ) + + +def test_step_forwards_its_camera_inputs_to_the_model( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline() + runtime, session = _session(tmp_path, monkeypatch, pipeline) + + session.step(_step_payload(1.0)) + session.step(_step_payload(2.0)) + + assert len(pipeline.generate_calls) == 2 + # Distinct per-step payloads must reach the model distinctly; a session that + # ignored its inputs would send the same slice twice. + assert pipeline.generate_calls[0]["poses"][0, 2, 3] == 1.0 + assert pipeline.generate_calls[1]["poses"][0, 2, 3] == 2.0 + assert pipeline.generate_calls[0]["intrinsics"][0, 0] == 1.0 + assert pipeline.generate_calls[1]["intrinsics"][0, 0] == 2.0 + assert pipeline.generate_calls[0]["world_scale"] == 2.5 + runtime.close() + + +def test_step_rejects_missing_camera_inputs( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline() + runtime, session = _session(tmp_path, monkeypatch, pipeline) + + with pytest.raises(ValueError, match="missing 'camera_trajectory'"): + session.step(InferenceInput()) + + assert pipeline.generate_calls == [] + runtime.close() + + +def test_step_rejects_wrongly_shaped_camera_inputs( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline() + runtime, session = _session(tmp_path, monkeypatch, pipeline) + + with pytest.raises(ValueError, match=r"must have shape \(2, 4, 4\)"): + session.step( + InferenceInput( + step={ + FIELD_CAMERA_TRAJECTORY: torch.eye(4).repeat(5, 1, 1), + FIELD_CAMERA_INTRINSICS: torch.zeros(5, 4), + } + ) + ) + + runtime.close() + + +def test_step_request_publishes_the_input_window( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline() + runtime, session = _session(tmp_path, monkeypatch, pipeline) + + first = session.next_step_request() + assert first is not None + assert first.user_input_window is not None + assert first.user_input_window.start_s == 0.0 + assert first.user_input_window.end_s == 2 / 16 + assert first.metadata == {"num_frames": 2, "frame_start": 0} + + session.step(_step_payload(1.0)) + second = session.next_step_request() + assert second is not None + assert second.user_input_window is not None + # Windows must advance with the rollout so each step maps its own events. + assert second.user_input_window.start_s == 2 / 16 + assert second.user_input_window.end_s == 4 / 16 + runtime.close() + + +def test_rollout_ends_when_the_camera_source_runs_out( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline() + runtime, session = _session( + tmp_path, monkeypatch, pipeline, total_blocks=10, total_camera_frames=3 + ) + + assert session.next_step_request() is not None + session.step(_step_payload(1.0)) + # Only 3 frames are available and each step needs 2, so the second step + # would overrun the source. + assert session.next_step_request() is None + runtime.close() + + +def test_unbounded_source_runs_until_total_blocks( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline() + runtime, session = _session( + tmp_path, monkeypatch, pipeline, total_blocks=2, total_camera_frames=None + ) + + steps = 0 + while session.next_step_request() is not None: + session.step(_step_payload(float(steps))) + steps += 1 + + assert steps == 2 + runtime.close() + + +def test_text_event_prompt_update_swaps_the_rollout_context( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline() + runtime, session = _session(tmp_path, monkeypatch, pipeline) + + payload = _step_payload(1.0) + session.step( + InferenceInput( + global_conditioning={"prompt": "a violent storm"}, + step=payload.step, + ) + ) + + assert pipeline.text_encoder_calls == [["a violent storm"]] + assert len(pipeline.diffusion_model.transformer.replaced) == 1 + + # Re-sending the same prompt must not re-encode or re-swap. + session.step( + InferenceInput( + global_conditioning={"prompt": "a violent storm"}, + step=payload.step, + ) + ) + assert pipeline.text_encoder_calls == [["a violent storm"]] + runtime.close() + + +def test_text_event_on_an_unsupported_pipeline_fails_clearly( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pipeline = _FakePipeline(supports_text_swap=False) + runtime, session = _session(tmp_path, monkeypatch, pipeline) + + with pytest.raises(RuntimeError, match="replace_text_embeddings"): + session.step( + InferenceInput( + global_conditioning={"prompt": "a violent storm"}, + step=_step_payload(1.0).step, + ) + ) + + # A rollout with no text event must not need the capability at all. + pipeline.generate_calls.clear() + session.step(_step_payload(1.0)) + assert len(pipeline.generate_calls) == 1 + runtime.close() + + +def test_full_standard_loop_drives_the_session_through_the_mapping( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The real runner must accept the mapping and feed the session per step. + + A fake runner cannot catch a mapping the compatibility check rejects, so + this exercises flashdreams.runtime.run_inference_session itself. + """ + import numpy as np + from lingbot.runtime import ( + LingbotModelAdapter, + LingbotReplayInputs, + inference_input_from_replay_inputs, + ) + + from flashdreams.runtime import InputCanonicalizer, UserInputs, UserInputSchema + from flashdreams.runtime.metrics import NullMetricsRecorder + from flashdreams.runtime.output import OutputArtifact + from flashdreams.runtime.runner import run_inference_session + + image = tmp_path / "image.jpg" + image.write_bytes(b"fake") + poses_path = tmp_path / "poses.npy" + intrinsics_path = tmp_path / "intrinsics.npy" + trajectory = np.tile(np.eye(4, dtype=np.float32), (32, 1, 1)) + trajectory[:, 2, 3] = np.arange(32, dtype=np.float32) + np.save(poses_path, trajectory) + np.save( + intrinsics_path, + np.tile(np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (32, 1)), + ) + + pipeline = _FakePipeline() + monkeypatch.setattr( + runtime_module, + "load_first_frame_tensor", + lambda *args, **kwargs: torch.zeros(1, 3, 2, 2), + ) + adapter = LingbotModelAdapter( + pipeline_factory=lambda pipeline_config, device: pipeline, + ) + replay_inputs = LingbotReplayInputs( + prompt="a calm street", + first_frame_path=image, + camera_poses_path=poses_path, + camera_intrinsics_path=intrinsics_path, + total_blocks=3, + pixel_height=2, + pixel_width=2, + fps=16, + ) + + class _Collecting: + def __init__(self) -> None: + self.results: list[Any] = [] + + def open(self) -> None: + return None + + def write(self, result: Any) -> None: + self.results.append(result) + + def close(self) -> tuple[OutputArtifact, ...]: + return () + + output = _Collecting() + mapping = adapter.create_input_mapping(replay_inputs) + run_inference_session( + adapter=adapter, + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + device="cpu", + runtime_options={"pipeline_config": object()}, + ), + mapping=mapping, + canonicalizer=InputCanonicalizer(), + source_schema=UserInputSchema(), + user_inputs=UserInputs(), + initial_inputs=inference_input_from_replay_inputs(replay_inputs), + output=output, + metrics=NullMetricsRecorder(), + ) + + assert len(output.results) == 3 + assert len(pipeline.generate_calls) == 3 + # Each step must receive its own successive slice of the trace. Comparing + # against the trace directly is the real property; consecutive pose values + # can repeat because preprocess_example_poses re-expands encoded poses at + # stride-4 cadence. + trace_poses = mapping.camera_trace.poses + received = torch.cat([call["poses"] for call in pipeline.generate_calls]) + assert torch.allclose(received, trace_poses[:6]) diff --git a/integrations/lingbot/tests/test_smoke.py b/integrations/lingbot/tests/test_smoke.py index bf9c1ab1a..3cd0a93a6 100644 --- a/integrations/lingbot/tests/test_smoke.py +++ b/integrations/lingbot/tests/test_smoke.py @@ -21,9 +21,11 @@ from pathlib import Path from typing import cast +import numpy as np import pytest import tomli as tomllib from lingbot import config as config_mod +from lingbot import example_data as example_data_mod from lingbot import runner as runner_mod from lingbot.config import ( LINGBOT_WORLD_V2_CHECKPOINT_PATH, @@ -39,6 +41,14 @@ LingbotWorldRunnerConfig, example_data_dirname, ) +from lingbot.runtime import ( + FIELD_FIRST_FRAME_PATH, + FIELD_PROMPT, + FIELD_TOTAL_BLOCKS, + LINGBOT_MODEL_ID, + LingbotModelAdapter, + LingbotRunnerOutputTarget, +) from lingbot.transformer import ( LINGBOT_WORLD_MIN_CHECKPOINT_FREE_GB, LingbotWorldTransformer, @@ -50,6 +60,19 @@ pytestmark = pytest.mark.ci_cpu + +def _write_camera_assets(poses: Path, intrinsics: Path, *, frames: int = 64) -> None: + """Write real .npy camera assets; the input mapping loads them for real.""" + trajectory = np.tile(np.eye(4, dtype=np.float32), (frames, 1, 1)) + trajectory[:, 2, 3] = np.linspace(0.0, 1.0, frames, dtype=np.float32) + np.save(poses, trajectory) + np.save( + intrinsics, + np.tile( + np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (frames, 1) + ), + ) + ENTRY_POINT_GROUP = "flashdreams.runner_configs" @@ -77,10 +100,10 @@ def _record_download(url: str, *, cache_dir: Path, filename: str) -> None: del cache_dir, filename urls.append(url) - monkeypatch.setattr(runner_mod, "EXAMPLE_DATA_DIR_LOCAL", tmp_path) - monkeypatch.setattr(runner_mod, "download_to_cache", _record_download) + monkeypatch.setattr(example_data_mod, "EXAMPLE_DATA_DIR_LOCAL", tmp_path) + monkeypatch.setattr(example_data_mod, "download_to_cache", _record_download) - runner_mod.ensure_example_data_downloaded(is_rank_zero=True, example_idx=0) + example_data_mod.ensure_example_data_downloaded(is_rank_zero=True, example_idx=0) expected_base_url = ( "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00" @@ -105,10 +128,10 @@ def test_promptless_examples_skip_the_prompt_download( def _record_download(url: str, *, cache_dir: Path, filename: str) -> None: downloads.append((url, cache_dir, filename)) - monkeypatch.setattr(runner_mod, "EXAMPLE_DATA_DIR_LOCAL", tmp_path) - monkeypatch.setattr(runner_mod, "download_to_cache", _record_download) + monkeypatch.setattr(example_data_mod, "EXAMPLE_DATA_DIR_LOCAL", tmp_path) + monkeypatch.setattr(example_data_mod, "download_to_cache", _record_download) - cache_dir = runner_mod.ensure_example_data_downloaded( + cache_dir = example_data_mod.ensure_example_data_downloaded( is_rank_zero=True, example_idx=example_idx, ) @@ -162,6 +185,69 @@ def test_promptless_example_resolves_to_empty_string( ] +def test_runner_delegates_to_runtime_api_with_direct_inputs( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Keep the CLI runner on the new runtime path, not the old rollout loop.""" + image = tmp_path / "image.jpg" + poses = tmp_path / "poses.npy" + intrinsics = tmp_path / "intrinsics.npy" + image.write_bytes(b"fake") + _write_camera_assets(poses, intrinsics) + runner = object.__new__(LingbotWorldRunner) + runner_config = cast( + LingbotWorldRunnerConfig, + derive_config( + RUNNER_CONFIGS["lingbot-world-fast-taehv-window15-sink3"], + prompt="drive through a city", + image_path=image, + pose_path=poses, + intrinsic_path=intrinsics, + total_blocks=1, + device="cpu", + ), + ) + pipeline = object() + output_stream = object() + captured: dict[str, object] = {} + + def _fake_run_inference_session(**kwargs: object) -> tuple[object, ...]: + captured.update(kwargs) + return () + + monkeypatch.setattr( + runner, + "create_video_output_stream", + lambda **_kwargs: output_stream, + ) + monkeypatch.setattr( + runner_mod, + "run_inference_session", + _fake_run_inference_session, + ) + runner.config = runner_config + runner.pipeline = pipeline + runner.local_rank = 0 + runner.world_size = 1 + runner.is_rank_zero = True + + runner.run() + + assert isinstance(captured["adapter"], LingbotModelAdapter) + config = captured["config"] + assert getattr(config, "model_id") == LINGBOT_MODEL_ID + assert getattr(config, "device") == "cpu" + assert config.runtime_options["pipeline"] is pipeline + inputs = captured["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 + output = captured["output"] + assert isinstance(output, LingbotRunnerOutputTarget) + assert output.output_stream is output_stream + + def test_runners_dict_is_non_empty() -> None: """Plugin must expose at least one runner.""" assert RUNNER_CONFIGS, "RUNNER_CONFIGS is empty" From c9705f49bb52467fad7afe6a8573495eb841d90f Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Fri, 7 Aug 2026 03:12:45 +0000 Subject: [PATCH 10/19] Unify video output stream handling Signed-off-by: Gangzheng Tong --- flashdreams/flashdreams/infra/runner.py | 6 +- flashdreams/flashdreams/infra/video_output.py | 141 +++++++++++++++++- .../flashdreams/runtime/video_output.py | 55 +------ .../flashdreams/serving/webrtc/manager.py | 6 +- .../flashdreams/serving/webrtc/runtime.py | 38 +---- flashdreams/tests/test_video_output.py | 86 +++++++---- flashdreams/tests/test_webrtc_manager.py | 18 +-- .../causal_forcing/causal_forcing/runner.py | 12 +- .../cosmos_predict2/cosmos_predict2/runner.py | 11 +- .../fastvideo_causal_wan22/runner.py | 12 +- integrations/flashvsr/flashvsr/runner.py | 11 +- .../hy_worldplay/hy_worldplay/runner.py | 11 +- integrations/lingbot/lingbot/runtime.py | 20 +-- .../lingbot/lingbot/webrtc/session.py | 25 ++-- .../lingbot/tests/test_webrtc_runtime.py | 29 ++-- .../world_model/flashdreams_adapter.py | 45 ++++-- integrations/omnidreams/omnidreams/runner.py | 4 +- .../omnidreams/omnidreams/webrtc/session.py | 52 ++++--- .../omnidreams/tests/test_webrtc_runtime.py | 16 +- .../self_forcing/self_forcing/runner.py | 12 +- integrations/wan21/wan21/runner.py | 12 +- 21 files changed, 345 insertions(+), 277 deletions(-) diff --git a/flashdreams/flashdreams/infra/runner.py b/flashdreams/flashdreams/infra/runner.py index 2ef2fc296..cf91c0960 100644 --- a/flashdreams/flashdreams/infra/runner.py +++ b/flashdreams/flashdreams/infra/runner.py @@ -39,7 +39,7 @@ VideoTensorLayout, create_runner_postprocess_stream, ) -from flashdreams.infra.video_output import RunnerVideoOutputStream +from flashdreams.infra.video_output import VideoOutputStream def _is_torchrun_env() -> bool: @@ -184,14 +184,14 @@ def create_video_output_stream( *, fps: float | None = None, move_to_cpu: bool = True, - ) -> RunnerVideoOutputStream: + ) -> VideoOutputStream: """Create the standard runner video output stream for one rollout.""" layout = self.config.postprocess_output_layout if layout is None: raise ValueError( "Runner video output collection requires an output layout." ) - return RunnerVideoOutputStream( + return VideoOutputStream( postprocess_stream=self.create_postprocess_stream(fps=fps), output_layout=layout, collect_output=self.is_rank_zero, diff --git a/flashdreams/flashdreams/infra/video_output.py b/flashdreams/flashdreams/infra/video_output.py index 930bcf572..24f9ab70c 100644 --- a/flashdreams/flashdreams/infra/video_output.py +++ b/flashdreams/flashdreams/infra/video_output.py @@ -17,8 +17,9 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass, field +from pathlib import Path from typing import Any import torch @@ -202,8 +203,13 @@ def video_hwc_uint8( ) -class RunnerVideoOutputStream: - """Post-process, collect, and summarize runner video chunks.""" +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`. + """ def __init__( self, @@ -236,8 +242,8 @@ def process( autoregressive_index: int, stats: dict[str, float] | None = None, stats_extra: Mapping[str, object] | None = None, - ) -> None: - """Process one generated chunk and collect it when this rank writes output.""" + ) -> Tensor: + """Process one chunk, optionally collect it, and return emitted frames.""" if self._closed: raise RuntimeError("cannot process video after finish()") processed = video_chunk @@ -259,6 +265,7 @@ def process( 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.""" @@ -271,6 +278,97 @@ def finish(self) -> Tensor | 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, + video_chunk=processed.detach(), + layout=self.output_layout, + stats=stats, + metadata=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: + return None + return self.write_mp4( + video, + output_path, + fps=fps, + layout=self.output_layout, + writer=writer, + install_hint=install_hint, + ) + + def write_mp4( + 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 + def _append_if_nonempty(self, output: Tensor) -> None: if not self.collect_output or output.shape[self._time_dim] == 0: return @@ -288,12 +386,43 @@ def _collected_output(self) -> Tensor | None: return output +def prepare_video_for_mp4( + video: Tensor, + *, + layout: VideoTensorLayout | str, +) -> tuple[Tensor, str]: + """Convert a stream output into a layout accepted by runner MP4 I/O.""" + if layout in {"thwc", "tchw", "btchw", "bcthw"}: + return video, layout + if layout == "bvtchw": + if video.ndim != 6: + raise ValueError( + "layout='bvtchw' expects a 6D [B,V,T,C,H,W] tensor, " + f"got {tuple(video.shape)}." + ) + if video.shape[0] != 1: + raise ValueError( + "layout='bvtchw' MP4 writing expects a single batch element, " + f"got {tuple(video.shape)}." + ) + _, views, frames, channels, height, width = video.shape + canvas = ( + video[0] + .permute(1, 3, 0, 4, 2) + .contiguous() + .reshape(frames, height, views * width, channels) + ) + return canvas, "thwc" + raise ValueError(f"unsupported video layout for MP4: {layout!r}") + + __all__ = [ "LazyRGBFrame", - "RunnerVideoOutputStream", + "VideoOutputStream", "VideoStepResult", "infer_video_num_frames", "lazy_rgb_frames_from_video_tensor", + "prepare_video_for_mp4", "video_layout_time_dim", "video_tensor_to_hwc_uint8", ] diff --git a/flashdreams/flashdreams/runtime/video_output.py b/flashdreams/flashdreams/runtime/video_output.py index b5c372125..bd2e466e2 100644 --- a/flashdreams/flashdreams/runtime/video_output.py +++ b/flashdreams/flashdreams/runtime/video_output.py @@ -8,19 +8,13 @@ from collections.abc import Callable, Sequence from dataclasses import dataclass, field from pathlib import Path -from typing import cast - -import torch from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner_io import ( DEFAULT_RUNNER_INSTALL_HINT, write_video_tensor, ) -from flashdreams.infra.runner_io import ( - VideoTensorLayout as WritableVideoTensorLayout, -) -from flashdreams.infra.video_output import RunnerVideoOutputStream, VideoStepResult +from flashdreams.infra.video_output import VideoOutputStream, VideoStepResult from flashdreams.runtime.output import OutputArtifact from flashdreams.runtime.types import StepResult @@ -38,7 +32,7 @@ class Mp4VideoOutputTarget: install_hint: str = DEFAULT_RUNNER_INSTALL_HINT move_to_cpu: bool = True _opened: bool = field(default=False, init=False, repr=False) - _stream: RunnerVideoOutputStream | None = field( + _stream: VideoOutputStream | None = field( default=None, init=False, repr=False, @@ -49,7 +43,7 @@ def closed(self) -> bool: return not self._opened def open(self) -> None: - self._stream = RunnerVideoOutputStream( + self._stream = VideoOutputStream( postprocess_stream=None, output_layout=self.output_layout, collect_output=True, @@ -97,16 +91,12 @@ def close(self) -> Sequence[OutputArtifact]: video = stream.finish() if video is None: return () - - writable_video, writable_layout = _prepare_video_for_mp4( + path = stream.write_mp4( video, - layout=self.output_layout, - ) - path = self.writer( - writable_video, self.output_path, fps=self.fps, - layout=writable_layout, + layout=self.output_layout, + writer=self.writer, install_hint=self.install_hint, ) return ( @@ -116,42 +106,11 @@ def close(self) -> Sequence[OutputArtifact]: metadata={ "fps": self.fps, "source_layout": self.output_layout, - "write_layout": writable_layout, - "shape": tuple(int(dim) for dim in writable_video.shape), + "shape": tuple(int(dim) for dim in video.shape), "stats_history": tuple(stream.stats_history), }, ), ) -def _prepare_video_for_mp4( - video: torch.Tensor, - *, - layout: VideoTensorLayout, -) -> tuple[torch.Tensor, WritableVideoTensorLayout]: - """Convert runtime video layouts into layouts accepted by runner I/O.""" - if layout in {"tchw", "btchw", "bcthw"}: - return video, cast(WritableVideoTensorLayout, layout) - if layout == "bvtchw": - if video.ndim != 6: - raise ValueError( - "layout='bvtchw' expects a 6D [B,V,T,C,H,W] tensor, " - f"got {tuple(video.shape)}." - ) - if video.shape[0] != 1: - raise ValueError( - "layout='bvtchw' MP4 writing expects a single batch element, " - f"got {tuple(video.shape)}." - ) - _, views, frames, channels, height, width = video.shape - canvas = ( - video[0] - .permute(1, 3, 0, 4, 2) - .contiguous() - .reshape(frames, height, views * width, channels) - ) - return canvas, "thwc" - raise ValueError(f"unsupported runtime video layout for MP4: {layout!r}") - - __all__ = ["Mp4VideoOutputTarget"] diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index e752be2a4..97c0c578e 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -23,6 +23,7 @@ ) from loguru import logger +from flashdreams.infra.video_output import VideoStepResult from flashdreams.serving.realtime.input import KeyboardResampler from flashdreams.serving.webrtc.encoders import ( DefaultRTCEncoder, @@ -41,8 +42,6 @@ from flashdreams.serving.webrtc.runtime import ( WebRTCRuntimeConfig, WebRTCSessionRuntime, - WebRTCStepResult, - make_webrtc_step_result, ) from flashdreams.serving.webrtc.server import SessionBusyError from flashdreams.serving.webrtc.warmup import ( @@ -54,8 +53,7 @@ "BaseWebRTCSessionManager", "ManagedWebRTCSession", "WebRTCControlSignal", - "WebRTCStepResult", - "make_webrtc_step_result", + "VideoStepResult", ] # Close the active session if no client heartbeat/control message arrives diff --git a/flashdreams/flashdreams/serving/webrtc/runtime.py b/flashdreams/flashdreams/serving/webrtc/runtime.py index 79fa9e0e2..049574a13 100644 --- a/flashdreams/flashdreams/serving/webrtc/runtime.py +++ b/flashdreams/flashdreams/serving/webrtc/runtime.py @@ -5,45 +5,13 @@ from __future__ import annotations -from collections.abc import Awaitable, Mapping +from collections.abc import Awaitable from typing import Any, Protocol -import torch - -from flashdreams.infra.postprocess import VideoTensorLayout -from flashdreams.infra.video_output import VideoStepResult, infer_video_num_frames +from flashdreams.infra.video_output import VideoStepResult from flashdreams.serving.realtime.input import PoseSegment -class WebRTCStepResult(VideoStepResult): - """One generated chunk handed back by a WebRTC model runtime.""" - - -def make_webrtc_step_result( - *, - chunk_index: int, - video_chunk: torch.Tensor, - layout: VideoTensorLayout, - stats: dict[str, float] | None = None, - sync_device: torch.device | str | None = None, - metadata: Mapping[str, Any] | None = None, -) -> WebRTCStepResult: - """Package a generated chunk for WebRTC without forcing a host copy.""" - if sync_device is not None: - device = torch.device(sync_device) - if device.type == "cuda": - torch.cuda.current_stream(device).synchronize() - - return WebRTCStepResult( - chunk_index=chunk_index, - num_frames=infer_video_num_frames(video_chunk, layout=layout), - video_chunk=video_chunk.detach(), - stats=stats, - layout=layout, - metadata=dict(metadata or {}), - ) - - class WebRTCRuntimeConfig(Protocol): """Config fields consumed by the shared WebRTC session manager.""" @@ -83,7 +51,7 @@ async def generate_chunk( *, segments: list[PoseSegment], frame_times: list[float], - ) -> WebRTCStepResult: ... + ) -> VideoStepResult: ... async def close(self) -> None: ... diff --git a/flashdreams/tests/test_video_output.py b/flashdreams/tests/test_video_output.py index 4b598cb63..51b375120 100644 --- a/flashdreams/tests/test_video_output.py +++ b/flashdreams/tests/test_video_output.py @@ -5,18 +5,20 @@ from __future__ import annotations +from pathlib import Path +from typing import Any + import pytest import torch from flashdreams.infra.video_output import ( LazyRGBFrame, - RunnerVideoOutputStream, + VideoOutputStream, VideoStepResult, infer_video_num_frames, lazy_rgb_frames_from_video_tensor, video_tensor_to_hwc_uint8, ) -from flashdreams.serving.webrtc.manager import WebRTCStepResult, make_webrtc_step_result pytestmark = pytest.mark.ci_cpu @@ -41,29 +43,22 @@ def test_video_step_result_infers_num_frames_from_layout() -> None: assert infer_video_num_frames(video, layout="bvtchw") == 4 -def test_webrtc_step_result_uses_shared_video_result_contract() -> None: - result = WebRTCStepResult( - chunk_index=1, - num_frames=2, - video_chunk=torch.zeros((2, 3, 4, 5)), - stats=None, - ) - - assert isinstance(result, VideoStepResult) - assert result.num_frames == 2 - - -def test_make_webrtc_step_result_preserves_tensor_and_infers_frames() -> None: +def test_video_output_stream_makes_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 = make_webrtc_step_result( - chunk_index=4, - video_chunk=video, - layout="tchw", + result = output_stream.make_step_result( + video, + autoregressive_index=4, stats={"decode_ms": 1.5}, ) - assert isinstance(result, WebRTCStepResult) + assert isinstance(result, VideoStepResult) assert result.chunk_index == 4 assert result.num_frames == 3 assert result.video_chunk.device == video.device @@ -111,15 +106,15 @@ def test_video_step_result_exposes_lazy_rgb_frames() -> None: assert frames[0].to_numpy().shape == (4, 5, 3) -def test_runner_video_output_stream_collects_chunks_and_stats() -> None: - output_stream = RunnerVideoOutputStream( +def test_video_output_stream_collects_chunks_and_stats() -> None: + output_stream = VideoOutputStream( postprocess_stream=None, output_layout="tchw", move_to_cpu=False, ) chunk = torch.zeros((2, 3, 4, 5), dtype=torch.float32) - output_stream.process( + processed = output_stream.process( chunk, autoregressive_index=3, stats={"total_ms": 8.0}, @@ -130,6 +125,7 @@ def test_runner_video_output_stream_collects_chunks_and_stats() -> None: 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 == [ { "autoregressive_index": 3, @@ -140,8 +136,8 @@ def test_runner_video_output_stream_collects_chunks_and_stats() -> None: ] -def test_runner_video_output_stream_collects_noop_chunks_without_postprocess() -> None: - output_stream = RunnerVideoOutputStream( +def test_video_output_stream_collects_noop_chunks_without_postprocess() -> None: + output_stream = VideoOutputStream( postprocess_stream=None, output_layout="bcthw", move_to_cpu=False, @@ -159,3 +155,43 @@ def test_runner_video_output_stream_collects_noop_chunks_without_postprocess() - assert output.shape == (1, 3, 3, 4, 5) assert torch.equal(output[:, :, :2], first) 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 + + output_stream = VideoOutputStream( + postprocess_stream=None, + output_layout="bvtchw", + move_to_cpu=False, + ) + output_stream.process( + torch.zeros((1, 2, 3, 3, 4, 5)), autoregressive_index=0 + ) + + written = output_stream.finish_to_mp4( + Path("output.mp4"), fps=24, writer=fake_writer + ) + + assert written is not None + assert written == Path("output.mp4") + assert calls[0]["shape"] == (3, 4, 10, 3) diff --git a/flashdreams/tests/test_webrtc_manager.py b/flashdreams/tests/test_webrtc_manager.py index e213d0760..b7dca0582 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -17,8 +17,8 @@ from flashdreams.serving.webrtc.manager import ( BaseWebRTCSessionManager, ManagedWebRTCSession, - WebRTCStepResult, ) +from flashdreams.infra.video_output import VideoStepResult from flashdreams.serving.webrtc.server import SessionBusyError pytestmark = pytest.mark.ci_cpu @@ -197,7 +197,7 @@ def peek_next_chunk_num_frames(self) -> int: async def generate_chunk( self, *, segments: Any, frame_times: Any - ) -> WebRTCStepResult: + ) -> VideoStepResult: del segments, frame_times self.generate_calls += 1 raise RuntimeError("boom") @@ -234,7 +234,7 @@ def peek_next_chunk_num_frames(self) -> int: async def generate_chunk( self, *, segments: Any, frame_times: Any - ) -> WebRTCStepResult: + ) -> VideoStepResult: del segments, frame_times self.generate_calls += 1 # Stop the loop after the second attempt without tearing down. @@ -271,11 +271,11 @@ def peek_next_chunk_num_frames(self) -> int: async def generate_chunk( self, *, segments: Any, frame_times: Any - ) -> WebRTCStepResult: + ) -> VideoStepResult: del segments, frame_times if self.managed_session is not None: self.managed_session.closed = True - return WebRTCStepResult( + return VideoStepResult( chunk_index=0, num_frames=1, video_chunk=torch.zeros((1, 1, 1, 3, 2, 2), dtype=torch.uint8), @@ -339,12 +339,12 @@ def peek_next_input_num_frames(self) -> int: async def generate_chunk( self, *, segments: Any, frame_times: list[float] - ) -> WebRTCStepResult: + ) -> VideoStepResult: del segments self.frame_times = frame_times if self.managed_session is not None: self.managed_session.closed = True - return WebRTCStepResult( + return VideoStepResult( chunk_index=0, num_frames=5, video_chunk=torch.zeros((5, 1, 1, 3, 2, 2), dtype=torch.uint8), @@ -397,13 +397,13 @@ def peek_next_chunk_num_frames(self) -> int: async def generate_chunk( self, *, segments: Any, frame_times: Any - ) -> WebRTCStepResult: + ) -> VideoStepResult: del 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 WebRTCStepResult( + return VideoStepResult( chunk_index=chunk_index, num_frames=4, video_chunk=torch.zeros((4, 1, 1, 3, 2, 2), dtype=torch.uint8), diff --git a/integrations/causal_forcing/causal_forcing/runner.py b/integrations/causal_forcing/causal_forcing/runner.py index 1a362bec9..2d4634d75 100644 --- a/integrations/causal_forcing/causal_forcing/runner.py +++ b/integrations/causal_forcing/causal_forcing/runner.py @@ -28,14 +28,12 @@ from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig from flashdreams.infra.runner_io import ( - ensure_output_dir, load_first_frame_tensor, read_image_rgb, resolve_input_path, resolve_prompt_value, runner_artifact_path, write_runner_stats, - write_video_tensor, ) from flashdreams.recipes.wan import ( WanInferencePipeline, @@ -173,14 +171,10 @@ def run(self) -> None: stats = self.pipeline.finalize(autoregressive_index=i, cache=cache) output_stream.process(video_chunk, autoregressive_index=i, stats=stats) - generated = output_stream.finish() - if generated is None: - return - - # Write the video. - ensure_output_dir(config.output_dir) video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - write_video_tensor(generated, video_path, fps=config.fps, layout="tchw") + video_path = output_stream.finish_to_mp4(video_path, fps=config.fps) + if video_path is None: + return logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " diff --git a/integrations/cosmos_predict2/cosmos_predict2/runner.py b/integrations/cosmos_predict2/cosmos_predict2/runner.py index 8d4733601..057b1cbb7 100644 --- a/integrations/cosmos_predict2/cosmos_predict2/runner.py +++ b/integrations/cosmos_predict2/cosmos_predict2/runner.py @@ -28,14 +28,12 @@ from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig from flashdreams.infra.runner_io import ( - ensure_output_dir, load_first_frame_tensor, read_image_rgb, resolve_input_path, resolve_prompt_value, runner_artifact_path, write_runner_stats, - write_video_tensor, ) from flashdreams.recipes.cosmos.pipeline import ( CosmosInferencePipeline, @@ -141,13 +139,10 @@ def run(self) -> None: 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) - generated = output_stream.finish() - if generated is None: - return - - ensure_output_dir(config.output_dir) video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - write_video_tensor(generated, video_path, fps=config.fps, layout="tchw") + video_path = output_stream.finish_to_mp4(video_path, fps=config.fps) + if video_path is None: + return logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " diff --git a/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py b/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py index a1aa9f560..f5cf82e80 100644 --- a/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py +++ b/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py @@ -26,11 +26,9 @@ from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig from flashdreams.infra.runner_io import ( - ensure_output_dir, resolve_prompt_value, runner_artifact_path, write_runner_stats, - write_video_tensor, ) from flashdreams.recipes.wan import ( WanInferencePipeline, @@ -134,14 +132,10 @@ def run(self) -> None: stats = self.pipeline.finalize(autoregressive_index=i, cache=cache) output_stream.process(video_chunk, autoregressive_index=i, stats=stats) - generated = output_stream.finish() - if generated is None: - return - - # Write the video. - ensure_output_dir(config.output_dir) video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - write_video_tensor(generated, video_path, fps=config.fps, layout="tchw") + video_path = output_stream.finish_to_mp4(video_path, fps=config.fps) + if video_path is None: + return logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " diff --git a/integrations/flashvsr/flashvsr/runner.py b/integrations/flashvsr/flashvsr/runner.py index e1707dd14..881943948 100644 --- a/integrations/flashvsr/flashvsr/runner.py +++ b/integrations/flashvsr/flashvsr/runner.py @@ -30,14 +30,12 @@ from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig, _is_torchrun_env from flashdreams.infra.runner_io import ( - ensure_output_dir, read_video_fps, read_video_rgb, resolve_input_path, rgb_video_to_normalized_tensor, runner_artifact_path, write_runner_stats, - write_video_tensor, ) from flashvsr.encoder import FlashVSREncoder from flashvsr.pipeline import ( @@ -452,13 +450,10 @@ def run(self) -> None: stats_extra=stats_extra, ) - generated = output_stream.finish() - if generated is None: - return - - ensure_output_dir(config.output_dir) video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - write_video_tensor(generated, video_path, fps=fps, layout="bcthw") + video_path = output_stream.finish_to_mp4(video_path, fps=fps) + if video_path is None: + return logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " diff --git a/integrations/hy_worldplay/hy_worldplay/runner.py b/integrations/hy_worldplay/hy_worldplay/runner.py index a97faee4e..c334809b9 100644 --- a/integrations/hy_worldplay/hy_worldplay/runner.py +++ b/integrations/hy_worldplay/hy_worldplay/runner.py @@ -31,11 +31,9 @@ from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig from flashdreams.infra.runner_io import ( - ensure_output_dir, resolve_prompt_value, runner_artifact_path, write_runner_stats, - write_video_tensor, ) from flashdreams.recipes.wan.pipeline import WanInferencePipeline @@ -353,15 +351,12 @@ def run(self) -> None: # (including the last) for consistent stats. stats = self.pipeline.finalize(ar_idx, cache) output_stream.process(chunk, autoregressive_index=ar_idx, stats=stats) - video = output_stream.finish() elapsed = time.time() - start_time - if video is None: - return - - ensure_output_dir(cfg.output_dir) out_path = runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4") - write_video_tensor(video, out_path, fps=cfg.fps, layout="btchw") + out_path = output_stream.finish_to_mp4(out_path, fps=cfg.fps) + if out_path is None: + return logger.info( f"[{cfg.runner_name}] wrote video " f"({tuple(video.shape)}) -> {out_path.resolve()} in {elapsed:.2f}s" diff --git a/integrations/lingbot/lingbot/runtime.py b/integrations/lingbot/lingbot/runtime.py index c15057aa5..0bf12bddc 100644 --- a/integrations/lingbot/lingbot/runtime.py +++ b/integrations/lingbot/lingbot/runtime.py @@ -20,13 +20,11 @@ from flashdreams.core.distributed import init as init_distributed from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner_io import ( - ensure_output_dir, load_first_frame_tensor, runner_artifact_path, write_runner_stats, - write_video_tensor, ) -from flashdreams.infra.video_output import RunnerVideoOutputStream, VideoStepResult +from flashdreams.infra.video_output import VideoOutputStream, VideoStepResult from flashdreams.runtime import ( CanonicalInputSchema, InferenceConfig, @@ -607,7 +605,7 @@ def _require_step_tensor( class LingbotRunnerOutputTarget: """Runner-compatible MP4/stats output target for Lingbot replay results.""" - output_stream: RunnerVideoOutputStream + output_stream: VideoOutputStream output_dir: Path runner_name: str fps: int | float @@ -635,23 +633,17 @@ def write(self, result: StepResult) -> None: def close(self) -> tuple[OutputArtifact, ...]: self._opened = False artifacts: list[OutputArtifact] = [] - video = self.output_stream.finish() - if video is None: - return () - - ensure_output_dir(self.output_dir) video_path = runner_artifact_path(self.output_dir, self.runner_name, "mp4") - write_video_tensor( - video, + video_path = self.output_stream.finish_to_mp4( video_path, fps=self.fps, - layout="tchw", install_hint=self.install_hint, ) + if video_path is None: + return () logger.info( - "[{}] wrote video {} -> {}", + "[{}] wrote video -> {}", self.runner_name, - tuple(video.shape), video_path.resolve(), ) artifacts.append( diff --git a/integrations/lingbot/lingbot/webrtc/session.py b/integrations/lingbot/lingbot/webrtc/session.py index ee23bb967..62dbb8442 100644 --- a/integrations/lingbot/lingbot/webrtc/session.py +++ b/integrations/lingbot/lingbot/webrtc/session.py @@ -41,6 +41,7 @@ ) 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.serving.webrtc.controls import ( CameraPoseIntegrator, PoseSegment, @@ -55,8 +56,6 @@ BaseWebRTCSessionManager, ManagedWebRTCSession, WebRTCControlSignal, - WebRTCStepResult, - make_webrtc_step_result, ) from flashdreams.serving.webrtc.server import SessionBusyError from lingbot.encoder.utils import preprocess_example_poses @@ -589,6 +588,12 @@ def __init__(self, config: LingbotRuntimeConfig | None = None) -> None: 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 @@ -675,7 +680,7 @@ async def generate_chunk( *, segments: list[PoseSegment], frame_times: list[float], - ) -> WebRTCStepResult: + ) -> VideoStepResult: """Generate one autoregressive chunk from a piecewise-constant timeline. Args: @@ -687,7 +692,7 @@ async def generate_chunk( :meth:`peek_next_chunk_num_frames` at call time. Returns: - :class:`WebRTCStepResult` carrying the produced video chunk + :class:`VideoStepResult` carrying the produced video chunk and the post-generation pipeline stats. Raises: @@ -758,7 +763,7 @@ def _generate_chunk_sync_all_ranks( self, segments: list[PoseSegment], frame_times: list[float], - ) -> WebRTCStepResult: + ) -> VideoStepResult: return self._generate_one_chunk_sync(segments=segments, frame_times=frame_times) @distributed_op(WebRTCControlSignal.EVENT) @@ -1135,7 +1140,7 @@ def _generate_one_chunk_sync( *, segments: list[PoseSegment], frame_times: list[float], - ) -> WebRTCStepResult: + ) -> VideoStepResult: if ( self._pipeline is None or self._cache is None @@ -1177,11 +1182,9 @@ def _generate_one_chunk_sync( input=camctrl_input, ) stats = self._pipeline.finalize(self.autoregressive_index, self._cache) - - result = make_webrtc_step_result( - chunk_index=self.autoregressive_index, - video_chunk=video_chunk, - layout="tchw", + result = self._output_stream.make_step_result( + video_chunk, + autoregressive_index=self.autoregressive_index, stats=stats, sync_device=self._device, ) diff --git a/integrations/lingbot/tests/test_webrtc_runtime.py b/integrations/lingbot/tests/test_webrtc_runtime.py index 97a6212a9..4153efb1b 100644 --- a/integrations/lingbot/tests/test_webrtc_runtime.py +++ b/integrations/lingbot/tests/test_webrtc_runtime.py @@ -29,8 +29,7 @@ LingbotWebRTCSessionManager, ) -from flashdreams.infra.postprocess import VideoTensorLayout -from flashdreams.serving.webrtc.manager import WebRTCStepResult +from flashdreams.infra.video_output import VideoOutputStream, VideoStepResult pytestmark = pytest.mark.ci_cpu @@ -173,7 +172,7 @@ def _fake_select_encoder(**kwargs: object) -> _FakeVideoEncoder: ] -def test_generate_one_chunk_sync_hands_gpu_resident_output_to_webrtc_helper( +def test_generate_one_chunk_sync_hands_gpu_resident_output_to_output_stream( monkeypatch: pytest.MonkeyPatch, ) -> None: class _NoCpuChunk: @@ -205,16 +204,17 @@ def finalize(autoregressive_index: int, cache: object) -> dict[str, float]: captured: dict[str, object] = {} - def _fake_make_webrtc_step_result(**kwargs: object) -> WebRTCStepResult: + def _fake_make_step_result( + _stream: VideoOutputStream, video_chunk: object, **kwargs: object + ) -> VideoStepResult: + captured["video_chunk"] = video_chunk captured.update(kwargs) - stats = cast(dict[str, float] | None, kwargs["stats"]) - layout = cast(VideoTensorLayout | None, kwargs["layout"]) - return WebRTCStepResult( + return VideoStepResult( chunk_index=0, num_frames=2, video_chunk=torch.zeros((2, 3, 4, 5)), - stats=stats, - layout=layout, + stats={"total_ms": 3.0}, + layout="tchw", ) runtime = session.LingbotInferenceRuntime( @@ -226,9 +226,9 @@ def _fake_make_webrtc_step_result(**kwargs: object) -> WebRTCStepResult: runtime._cache = object() runtime._base_intrinsics = torch.ones(4) monkeypatch.setattr( - session, - "make_webrtc_step_result", - _fake_make_webrtc_step_result, + VideoOutputStream, + "make_step_result", + _fake_make_step_result, ) result = runtime._generate_one_chunk_sync( @@ -237,7 +237,6 @@ def _fake_make_webrtc_step_result(**kwargs: object) -> WebRTCStepResult: ) assert captured["video_chunk"] is pipeline.output - assert captured["layout"] == "tchw" assert captured["sync_device"] == torch.device("cpu") assert pipeline.output.detach_calls == 0 assert result.stats == {"total_ms": 3.0} @@ -966,11 +965,11 @@ async def generate_chunk( *, segments: list[tuple[float, float, frozenset[str]]], frame_times: list[float], - ) -> WebRTCStepResult: + ) -> VideoStepResult: del frame_times chunk_index = len(self.generated_segments) self.generated_segments.append(segments) - return WebRTCStepResult( + return VideoStepResult( chunk_index=chunk_index, num_frames=1, video_chunk=torch.zeros((1, 1, 1, 3, 2, 2), dtype=torch.uint8), 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 e61b4904c..c74c7fb3e 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py @@ -31,7 +31,10 @@ VideoPostprocessChainConfig, VideoPostprocessStream, ) -from flashdreams.infra.video_output import lazy_rgb_frames_from_video_tensor +from flashdreams.infra.video_output import ( + VideoOutputStream, + lazy_rgb_frames_from_video_tensor, +) PipelineFactory = Callable[[WorldModelManifest, WorldModelProfileConfig], Any] _VIEW_NAMES = ["camera_front_wide_120fov"] @@ -499,7 +502,7 @@ def __init__( self._next_block_index = 0 self._postprocess = postprocess or VideoPostprocessChainConfig() self._postprocess_enabled = self._postprocess.is_enabled() - self._postprocess_stream: VideoPostprocessStream | None = None + self._output_stream: VideoOutputStream | None = None @property def pipeline(self) -> Any: @@ -636,7 +639,7 @@ def start( cache=self._cache, hdmap=self._condition_tensor(condition_frames), ) - video = self._postprocess_video(video, autoregressive_index=0) + video = self._process_video(video, autoregressive_index=0) model_frames = self._video_tensor_to_frames(video) _synchronize_cuda_frame_event(model_frames) self._pending_finalization_index = 0 @@ -665,7 +668,7 @@ def continue_generation(self, condition_frames: list[object]) -> list[object]: cache=self._cache, hdmap=self._condition_tensor(condition_frames), ) - video = self._postprocess_video( + video = self._process_video( video, autoregressive_index=self._next_block_index ) model_frames = self._video_tensor_to_frames(video) @@ -717,21 +720,31 @@ def set_postprocess_enabled(self, enabled: bool) -> None: self._postprocess.preset, ) - def _postprocess_video( - self, video: torch.Tensor, *, autoregressive_index: int - ) -> torch.Tensor: - if not self._postprocess_enabled: - return video - if self._postprocess_stream is None: - self._postprocess_stream = VideoPostprocessStream( + def _new_output_stream(self) -> VideoOutputStream: + postprocess_stream = None + if self._postprocess_enabled: + postprocess_stream = VideoPostprocessStream( postprocess=self._postprocess, output_layout="bvtchw", fps=self.manifest.fps, per_view=False, world_size=1, ) - processed = self._postprocess_stream.process( - video, autoregressive_index=autoregressive_index + 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( @@ -742,10 +755,10 @@ def _postprocess_video( return processed def _close_postprocess_stream(self) -> None: - if self._postprocess_stream is None: + if self._output_stream is None: return - self._postprocess_stream.finish() - self._postprocess_stream = None + self._output_stream.finish() + self._output_stream = None def _initialize_cache(self, initial_rgb: object, prompt: str) -> Any: if self.manifest.synthetic_model: diff --git a/integrations/omnidreams/omnidreams/runner.py b/integrations/omnidreams/omnidreams/runner.py index 2332d1c65..177cb7f2d 100644 --- a/integrations/omnidreams/omnidreams/runner.py +++ b/integrations/omnidreams/omnidreams/runner.py @@ -48,7 +48,6 @@ load_video_tensor, runner_artifact_path, write_runner_stats, - write_video_tensor, ) DEFAULT_VIDEO_HEIGHT = 704 @@ -427,9 +426,8 @@ def _rollout_and_save( video=video, ) - ensure_output_dir(cfg.output_dir) video_path = runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4") - write_video_tensor( + video_path = output_stream.write_mp4( canvas, video_path, fps=cfg.output_fps, diff --git a/integrations/omnidreams/omnidreams/webrtc/session.py b/integrations/omnidreams/omnidreams/webrtc/session.py index 86e4d020b..e45a07f80 100644 --- a/integrations/omnidreams/omnidreams/webrtc/session.py +++ b/integrations/omnidreams/omnidreams/webrtc/session.py @@ -53,6 +53,7 @@ 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, @@ -69,8 +70,6 @@ BaseWebRTCSessionManager, ManagedWebRTCSession, WebRTCControlSignal, - WebRTCStepResult, - make_webrtc_step_result, ) from flashdreams.serving.webrtc.server import SessionBusyError @@ -515,7 +514,7 @@ def __init__(self, config: OmnidreamsRuntimeConfig | None = None) -> None: self._camera_to_rig: torch.Tensor | None = None self._initial_ego_pose: np.ndarray | None = None self._next_timestamp_us: int = 0 - self._postprocess_stream: VideoPostprocessStream | None = None + self._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 @@ -595,7 +594,7 @@ async def generate_chunk( *, segments: list[PoseSegment], frame_times: list[float], - ) -> WebRTCStepResult: + ) -> VideoStepResult: if self._closed: raise OmnidreamsRuntimeError("Session is closed.") if self._wrapper is None: @@ -662,7 +661,7 @@ def _generate_chunk_sync_all_ranks( self, segments: list[PoseSegment], frame_times: list[float], - ) -> WebRTCStepResult: + ) -> VideoStepResult: return self._generate_one_chunk_sync(segments=segments, frame_times=frame_times) @distributed_op(WebRTCControlSignal.CLOSE) @@ -976,30 +975,42 @@ def _reset_postprocess_stream( world_size=world_size ): return - self._postprocess_stream = VideoPostprocessStream( + 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: - if self._postprocess_stream is None: - return - self._postprocess_stream.finish() - self._postprocess_stream = 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], - ) -> WebRTCStepResult: + ) -> VideoStepResult: if ( self._wrapper is None or self._renderer is None @@ -1068,19 +1079,18 @@ def _generate_one_chunk_sync( else: video_chunk = output.rgb_frames - if not serve_hdmaps and self._postprocess_stream is not None: - video_chunk = self._postprocess_stream.process( + 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, ) - - result = make_webrtc_step_result( - chunk_index=self.autoregressive_index, - video_chunk=video_chunk, - layout="bvtchw", - stats=None, - sync_device=self._device, - ) self.autoregressive_index += 1 return result diff --git a/integrations/omnidreams/tests/test_webrtc_runtime.py b/integrations/omnidreams/tests/test_webrtc_runtime.py index 9f46e6be7..0c6055f3f 100644 --- a/integrations/omnidreams/tests/test_webrtc_runtime.py +++ b/integrations/omnidreams/tests/test_webrtc_runtime.py @@ -33,6 +33,7 @@ VideoPostprocessChainConfig, VideoPostProcessorConfig, ) +from flashdreams.infra.video_output import VideoStepResult from flashdreams.serving.webrtc.controls import ( WSAD_SUPPORTED_KEYS, CameraPoseIntegrator, @@ -41,7 +42,6 @@ ChunkDeliveryResult, DefaultRTCEncoder, ) -from flashdreams.serving.webrtc.manager import WebRTCStepResult from flashdreams.serving.webrtc.media import BufferedVideoTrack from flashdreams.serving.webrtc.server import SESSION_MANAGER_KEY @@ -254,7 +254,9 @@ def process( runtime, _wrapper = _build_fake_runtime() postprocess_stream = _FakePostprocessStream() - runtime._postprocess_stream = postprocess_stream # ty:ignore[invalid-assignment] + runtime._output_stream.postprocess_stream = ( # ty:ignore[invalid-assignment] + postprocess_stream + ) result = runtime._generate_one_chunk_sync( segments=[(0.0, 2 / 30, frozenset({"w"}))], @@ -292,7 +294,7 @@ def test_session_postprocess_override_replaces_the_rollout_stream( runtime._reset_postprocess_stream( session.OmnidreamsSessionInput(postprocess_preset="fake-preset") ) - first_stream = runtime._postprocess_stream + first_stream = runtime._output_stream.postprocess_stream assert first_stream is not None assert runtime.postprocess_preset == "fake-preset" @@ -302,7 +304,7 @@ def test_session_postprocess_override_replaces_the_rollout_stream( ) assert first_stream._closed is True - assert runtime._postprocess_stream is None + assert runtime._output_stream.postprocess_stream is None assert runtime.postprocess_preset == "" @@ -1052,11 +1054,11 @@ async def generate_chunk( *, segments: list[tuple[float, float, frozenset[str]]], frame_times: list[float], - ) -> WebRTCStepResult: + ) -> VideoStepResult: del frame_times chunk_index = len(self.generated_segments) self.generated_segments.append(segments) - return WebRTCStepResult( + return VideoStepResult( chunk_index=chunk_index, num_frames=1, video_chunk=torch.zeros((1, 1, 1, 3, 2, 2), dtype=torch.uint8), @@ -1199,7 +1201,7 @@ async def generate_chunk( *, segments: list[tuple[float, float, frozenset[str]]], frame_times: list[float], - ) -> WebRTCStepResult: + ) -> VideoStepResult: del segments, frame_times self.generate_calls += 1 raise RuntimeError("boom") diff --git a/integrations/self_forcing/self_forcing/runner.py b/integrations/self_forcing/self_forcing/runner.py index da01d9787..ba97c1440 100644 --- a/integrations/self_forcing/self_forcing/runner.py +++ b/integrations/self_forcing/self_forcing/runner.py @@ -26,11 +26,9 @@ from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig from flashdreams.infra.runner_io import ( - ensure_output_dir, resolve_prompt_value, runner_artifact_path, write_runner_stats, - write_video_tensor, ) from flashdreams.recipes.wan import ( WanInferencePipeline, @@ -132,14 +130,10 @@ def run(self) -> None: stats = self.pipeline.finalize(autoregressive_index=i, cache=cache) output_stream.process(video_chunk, autoregressive_index=i, stats=stats) - generated = output_stream.finish() - if generated is None: - return - - # Write the video. - ensure_output_dir(config.output_dir) video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - write_video_tensor(generated, video_path, fps=config.fps, layout="tchw") + video_path = output_stream.finish_to_mp4(video_path, fps=config.fps) + if video_path is None: + return logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " diff --git a/integrations/wan21/wan21/runner.py b/integrations/wan21/wan21/runner.py index 7e0c1fe5d..33e2e4585 100644 --- a/integrations/wan21/wan21/runner.py +++ b/integrations/wan21/wan21/runner.py @@ -28,14 +28,12 @@ from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig from flashdreams.infra.runner_io import ( - ensure_output_dir, load_first_frame_tensor, read_image_rgb, resolve_input_path, resolve_prompt_value, runner_artifact_path, write_runner_stats, - write_video_tensor, ) from flashdreams.recipes.wan import ( WanInferencePipeline, @@ -176,14 +174,10 @@ def run(self) -> None: 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) - generated = output_stream.finish() - if generated is None: - return - - # Write the video. - ensure_output_dir(config.output_dir) video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - write_video_tensor(generated, video_path, fps=config.fps, layout="tchw") + video_path = output_stream.finish_to_mp4(video_path, fps=config.fps) + if video_path is None: + return logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " From 31fc0c9c809c88168e6089f636825420c63241eb Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Fri, 7 Aug 2026 04:46:59 +0000 Subject: [PATCH 11/19] Unify WebRTC demo viewer Signed-off-by: Gangzheng Tong --- .../serving}/webrtc/web/__init__.py | 0 .../webrtc/web/assets/horizontal-dark.svg | 192 ++++++++++++++++++ .../webrtc/web/assets/horizontal-light.svg | 182 +++++++++++++++++ .../serving}/webrtc/web/mock_ui_server.py | 6 +- .../serving}/webrtc/web/request_session.css | 0 .../serving}/webrtc/web/request_session.html | 12 +- .../serving}/webrtc/web/request_session.js | 24 ++- flashdreams/pyproject.toml | 3 + flashdreams/tests/test_webrtc_serving.py | 12 ++ integrations/lingbot/lingbot/demo/webrtc.py | 14 +- integrations/lingbot/tests/test_demo_api.py | 27 +-- .../omnidreams/omnidreams/demo/webrtc.py | 2 +- .../omnidreams/omnidreams/webrtc/server.py | 2 +- .../omnidreams/tests/test_webrtc_runtime.py | 2 +- .../tests/test_webrtc_server_routes.py | 4 +- 15 files changed, 444 insertions(+), 38 deletions(-) rename {integrations/omnidreams/omnidreams => flashdreams/flashdreams/serving}/webrtc/web/__init__.py (100%) create mode 100644 flashdreams/flashdreams/serving/webrtc/web/assets/horizontal-dark.svg create mode 100644 flashdreams/flashdreams/serving/webrtc/web/assets/horizontal-light.svg rename {integrations/lingbot/lingbot => flashdreams/flashdreams/serving}/webrtc/web/mock_ui_server.py (90%) rename {integrations/omnidreams/omnidreams => flashdreams/flashdreams/serving}/webrtc/web/request_session.css (100%) rename {integrations/omnidreams/omnidreams => flashdreams/flashdreams/serving}/webrtc/web/request_session.html (89%) rename {integrations/omnidreams/omnidreams => flashdreams/flashdreams/serving}/webrtc/web/request_session.js (97%) diff --git a/integrations/omnidreams/omnidreams/webrtc/web/__init__.py b/flashdreams/flashdreams/serving/webrtc/web/__init__.py similarity index 100% rename from integrations/omnidreams/omnidreams/webrtc/web/__init__.py rename to flashdreams/flashdreams/serving/webrtc/web/__init__.py diff --git a/flashdreams/flashdreams/serving/webrtc/web/assets/horizontal-dark.svg b/flashdreams/flashdreams/serving/webrtc/web/assets/horizontal-dark.svg new file mode 100644 index 000000000..89b68f9d1 --- /dev/null +++ b/flashdreams/flashdreams/serving/webrtc/web/assets/horizontal-dark.svg @@ -0,0 +1,192 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/flashdreams/flashdreams/serving/webrtc/web/assets/horizontal-light.svg b/flashdreams/flashdreams/serving/webrtc/web/assets/horizontal-light.svg new file mode 100644 index 000000000..a491910db --- /dev/null +++ b/flashdreams/flashdreams/serving/webrtc/web/assets/horizontal-light.svg @@ -0,0 +1,182 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/integrations/lingbot/lingbot/webrtc/web/mock_ui_server.py b/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py similarity index 90% rename from integrations/lingbot/lingbot/webrtc/web/mock_ui_server.py rename to flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py index fecdadd6c..5ad682b82 100644 --- a/integrations/lingbot/lingbot/webrtc/web/mock_ui_server.py +++ b/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py @@ -21,7 +21,7 @@ from importlib.resources import as_file, files from urllib.parse import urlsplit -WEB_DIR_RESOURCE = files("lingbot.webrtc").joinpath("web") +WEB_DIR_RESOURCE = files("flashdreams.serving.webrtc").joinpath("web") class MockUIRequestHandler(SimpleHTTPRequestHandler): @@ -52,7 +52,7 @@ def do_HEAD(self) -> None: def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Serve the Lingbot mock UI.") + parser = argparse.ArgumentParser(description="Serve the shared WebRTC mock UI.") parser.add_argument("--host", type=str, default="127.0.0.1") parser.add_argument("--port", type=int, default=8090) return parser.parse_args() @@ -64,7 +64,7 @@ def main() -> None: handler = partial(MockUIRequestHandler, directory=str(web_dir)) server = ThreadingHTTPServer((args.host, args.port), handler) print( - f"Serving mock UI at http://{args.host}:{args.port}/request_session?mock=1" + f"Serving shared mock UI at http://{args.host}:{args.port}/request_session?mock=1" ) try: server.serve_forever() diff --git a/integrations/omnidreams/omnidreams/webrtc/web/request_session.css b/flashdreams/flashdreams/serving/webrtc/web/request_session.css similarity index 100% rename from integrations/omnidreams/omnidreams/webrtc/web/request_session.css rename to flashdreams/flashdreams/serving/webrtc/web/request_session.css diff --git a/integrations/omnidreams/omnidreams/webrtc/web/request_session.html b/flashdreams/flashdreams/serving/webrtc/web/request_session.html similarity index 89% rename from integrations/omnidreams/omnidreams/webrtc/web/request_session.html rename to flashdreams/flashdreams/serving/webrtc/web/request_session.html index 263a1b299..321532e5a 100644 --- a/integrations/omnidreams/omnidreams/webrtc/web/request_session.html +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.html @@ -8,13 +8,13 @@ - Omnidreams WebRTC Drive - + FlashDreams WebRTC Drive +
-
-

Omnidreams WebRTC Drive

+
+

FlashDreams WebRTC Drive

@@ -86,12 +86,12 @@

Client Logs

World Model - Omnidreams + World Model
- + diff --git a/integrations/omnidreams/omnidreams/webrtc/web/request_session.js b/flashdreams/flashdreams/serving/webrtc/web/request_session.js similarity index 97% rename from integrations/omnidreams/omnidreams/webrtc/web/request_session.js rename to flashdreams/flashdreams/serving/webrtc/web/request_session.js index 11b7302e8..a9ae9e74a 100644 --- a/integrations/omnidreams/omnidreams/webrtc/web/request_session.js +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.js @@ -50,7 +50,7 @@ const metrics = { rttMs: null, resolution: null, step: null, - model: "Omnidreams", + model: "World Model", } function normalizeKey(rawKey) { @@ -83,7 +83,7 @@ function formatMs(value) { } function logEvent(message, { source = "server", level = "info" } = {}) { - const consoleMessage = `[Omnidreams WebRTC][${source}] ${message}` + const consoleMessage = `[FlashDreams WebRTC][${source}] ${message}` if (level === "error") { console.error(consoleMessage) } else { @@ -123,11 +123,22 @@ function setVideoVisible(visible) { } function setPostprocessDisabled(disabled) { + if (!postprocessSelect) { + return + } postprocessSelect.disabled = disabled || !postprocessControlAvailable } async function loadPostprocessOptions() { + if (!postprocessField || !postprocessSelect) { + return + } const response = await fetch("/api/postprocess/options") + if (response.status === 404) { + postprocessControlAvailable = false + postprocessField.hidden = true + return + } if (!response.ok) { throw new Error(`post-process options failed (${response.status})`) } @@ -158,6 +169,9 @@ async function loadPostprocessOptions() { } async function configureSessionInput() { + if (!postprocessControlAvailable || !postprocessSelect) { + return + } const postprocessPreset = postprocessSelect.value const response = await fetch("/api/session/input", { method: "POST", @@ -181,7 +195,7 @@ function renderMetrics() { latencyValue.textContent = formatMs(latency) resolutionValue.textContent = metrics.resolution || "--" stepValue.textContent = metrics.step === null ? "--" : String(metrics.step) - modelValue.textContent = metrics.model || "Omnidreams" + modelValue.textContent = metrics.model || "World Model" } function recordActionSent(action) { @@ -633,7 +647,7 @@ async function dumpPeerStats(reason) { for (const report of stats.values()) { reports.set(report.id, report) } - console.group(`[Omnidreams WebRTC] peer stats: ${reason}`) + console.group(`[FlashDreams WebRTC] peer stats: ${reason}`) for (const report of stats.values()) { if (report.type !== "candidate-pair") { continue @@ -655,7 +669,7 @@ async function dumpPeerStats(reason) { } console.groupEnd() } catch (error) { - console.warn("[Omnidreams WebRTC] getStats failed", error) + console.warn("[FlashDreams WebRTC] getStats failed", error) } } diff --git a/flashdreams/pyproject.toml b/flashdreams/pyproject.toml index 06278038c..fb5295781 100644 --- a/flashdreams/pyproject.toml +++ b/flashdreams/pyproject.toml @@ -135,6 +135,9 @@ version = {attr = "flashdreams._version.__version__"} include = ["flashdreams*", "tools*"] exclude = ["tests", "flashdreams._pytest_plugins*"] +[tool.setuptools.package-data] +"flashdreams.serving.webrtc" = ["web/*.html", "web/*.css", "web/*.js", "web/assets/*.svg"] + [dependency-groups] # Default CUDA 13 profile. No source binding -- Linux falls through to # PyPI's CUDA-13 manylinux wheel, Windows is captured by the unconditional diff --git a/flashdreams/tests/test_webrtc_serving.py b/flashdreams/tests/test_webrtc_serving.py index e1d1c9fc1..43a2c066a 100644 --- a/flashdreams/tests/test_webrtc_serving.py +++ b/flashdreams/tests/test_webrtc_serving.py @@ -4,6 +4,7 @@ from __future__ import annotations from contextlib import nullcontext +from importlib.resources import files import numpy as np import pytest @@ -173,6 +174,17 @@ def _raise_creation_failure(**_kwargs) -> web.Application: assert closed +def test_shared_viewer_treats_postprocess_routes_as_optional() -> None: + web_dir = files("flashdreams.serving.webrtc").joinpath("web") + html = web_dir.joinpath("request_session.html").read_text(encoding="utf-8") + javascript = web_dir.joinpath("request_session.js").read_text(encoding="utf-8") + + assert "/static/request_session.js?v=shared-webrtc-v1" in html + assert "if (!postprocessField || !postprocessSelect)" in javascript + assert "if (response.status === 404)" in javascript + assert "if (!postprocessControlAvailable || !postprocessSelect)" in javascript + + @pytest.mark.asyncio async def test_packaged_webrtc_app_serves_common_routes(tmp_path) -> None: (tmp_path / "request_session.html").write_text( diff --git a/integrations/lingbot/lingbot/demo/webrtc.py b/integrations/lingbot/lingbot/demo/webrtc.py index 966a79fb7..0b27c1f74 100644 --- a/integrations/lingbot/lingbot/demo/webrtc.py +++ b/integrations/lingbot/lingbot/demo/webrtc.py @@ -5,12 +5,16 @@ from __future__ import annotations +from importlib.resources import as_file, files from typing import Any from aiohttp import web from flashdreams.runtime.demo import DemoSpec -from lingbot.webrtc.server import create_app +from flashdreams.serving.webrtc.server import ( + close_package_resources, + create_packaged_webrtc_app, +) from lingbot.webrtc.session import ( LingbotInferenceRuntime, LingbotRuntimeConfig, @@ -43,11 +47,15 @@ def create_lingbot_webrtc_app( session_manager: Any, request_session_url: str, ) -> web.Application: - """Create the packaged Lingbot browser app through existing serving glue.""" + """Create Lingbot's shared browser app through generic serving glue.""" del spec - return create_app( + return create_packaged_webrtc_app( + web_resource=files("flashdreams.serving.webrtc").joinpath("web"), session_manager=session_manager, + preload_name="Lingbot", request_session_url=request_session_url, + as_file_fn=as_file, + cleanup_callback=close_package_resources, ) diff --git a/integrations/lingbot/tests/test_demo_api.py b/integrations/lingbot/tests/test_demo_api.py index 95515e027..c9446b414 100644 --- a/integrations/lingbot/tests/test_demo_api.py +++ b/integrations/lingbot/tests/test_demo_api.py @@ -432,27 +432,22 @@ def test_lingbot_webrtc_demo_uses_existing_manager_with_model_config() -> None: assert demo.port == 8080 -def test_lingbot_webrtc_demo_installs_model_routes( +def test_lingbot_webrtc_demo_uses_shared_viewer_shell( monkeypatch: pytest.MonkeyPatch, ) -> None: import lingbot.demo.webrtc as demo_webrtc_module app_calls: list[dict[str, Any]] = [] - async def _ok(request: web.Request) -> web.Response: - del request - return web.Response(text="ok") - - def fake_create_app(**kwargs: Any) -> web.Application: + def fake_create_packaged_app(**kwargs: Any) -> web.Application: app_calls.append(kwargs) app = web.Application() app[SESSION_MANAGER_KEY] = kwargs["session_manager"] - app.router.add_get("/api/session/initial_scene", _ok) - app.router.add_get("/api/session/first_frame", _ok) - app.router.add_post("/api/session/input", _ok) return app - monkeypatch.setattr(demo_webrtc_module, "create_app", fake_create_app) + monkeypatch.setattr( + demo_webrtc_module, "create_packaged_webrtc_app", fake_create_packaged_app + ) adapter = LingbotDemoAdapter(webrtc_runtime_factory=_FakeWebRTCRuntime) spec = DemoSpec( model_id=LINGBOT_MODEL_ID, @@ -479,10 +474,8 @@ def fake_create_app(**kwargs: Any) -> web.Application: assert app_calls[0]["request_session_url"] == ( "http://127.0.0.1:8080/request_session" ) - route_paths = {resource.canonical for resource in demo.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 + assert app_calls[0]["preload_name"] == "Lingbot" + assert str(app_calls[0]["web_resource"]).endswith("serving/webrtc/web") def test_lingbot_webrtc_demo_serves_through_shared_runner( @@ -492,7 +485,7 @@ def test_lingbot_webrtc_demo_serves_through_shared_runner( server_calls: list[dict[str, Any]] = [] - def fake_create_app(**kwargs: Any) -> web.Application: + def fake_create_packaged_app(**kwargs: Any) -> web.Application: app = web.Application() app[SESSION_MANAGER_KEY] = kwargs["session_manager"] return app @@ -500,7 +493,9 @@ def fake_create_app(**kwargs: Any) -> web.Application: def fake_server_runner(**kwargs: Any) -> None: server_calls.append(kwargs) - monkeypatch.setattr(demo_webrtc_module, "create_app", fake_create_app) + monkeypatch.setattr( + demo_webrtc_module, "create_packaged_webrtc_app", fake_create_packaged_app + ) adapter = LingbotDemoAdapter(webrtc_runtime_factory=_FakeWebRTCRuntime) spec = DemoSpec( model_id=LINGBOT_MODEL_ID, diff --git a/integrations/omnidreams/omnidreams/demo/webrtc.py b/integrations/omnidreams/omnidreams/demo/webrtc.py index 699d47091..911c38c72 100644 --- a/integrations/omnidreams/omnidreams/demo/webrtc.py +++ b/integrations/omnidreams/omnidreams/demo/webrtc.py @@ -148,7 +148,7 @@ def create_omnidreams_webrtc_app( 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("omnidreams.webrtc").joinpath("web"), + web_resource=files("flashdreams.serving.webrtc").joinpath("web"), session_manager=session_manager, preload_name=preload_name or "Omnidreams", request_session_url=request_session_url, diff --git a/integrations/omnidreams/omnidreams/webrtc/server.py b/integrations/omnidreams/omnidreams/webrtc/server.py index c62724422..bc50bf49c 100644 --- a/integrations/omnidreams/omnidreams/webrtc/server.py +++ b/integrations/omnidreams/omnidreams/webrtc/server.py @@ -55,7 +55,7 @@ close_package_resources as _close_package_resources, ) -WEB_DIR_RESOURCE = files("omnidreams.webrtc").joinpath("web") +WEB_DIR_RESOURCE = files("flashdreams.serving.webrtc").joinpath("web") class _OmnidreamsSessionManager(WebRTCSessionManager, Protocol): diff --git a/integrations/omnidreams/tests/test_webrtc_runtime.py b/integrations/omnidreams/tests/test_webrtc_runtime.py index 0c6055f3f..1a11f18f8 100644 --- a/integrations/omnidreams/tests/test_webrtc_runtime.py +++ b/integrations/omnidreams/tests/test_webrtc_runtime.py @@ -953,7 +953,7 @@ async def test_postprocess_options_exposes_only_launch_preset() -> None: def test_webrtc_ui_posts_selected_postprocess_preset() -> None: - web_dir = files("omnidreams.webrtc").joinpath("web") + web_dir = files("flashdreams.serving.webrtc").joinpath("web") html = web_dir.joinpath("request_session.html").read_text(encoding="utf-8") javascript = web_dir.joinpath("request_session.js").read_text(encoding="utf-8") diff --git a/integrations/omnidreams/tests/test_webrtc_server_routes.py b/integrations/omnidreams/tests/test_webrtc_server_routes.py index f052aac10..bddb824bd 100644 --- a/integrations/omnidreams/tests/test_webrtc_server_routes.py +++ b/integrations/omnidreams/tests/test_webrtc_server_routes.py @@ -84,7 +84,7 @@ def test_create_app_keeps_package_web_resource_materialized() -> None: web_dir = static_resources[0].get_info()["directory"] assert web_dir.is_dir() assert ( - "Omnidreams WebRTC Drive" in (web_dir / "request_session.html").read_text() + "FlashDreams WebRTC Drive" in (web_dir / "request_session.html").read_text() ) finally: app[PACKAGE_RESOURCE_STACK_KEY].close() @@ -132,7 +132,7 @@ async def test_request_session_serves_html() -> None: response = await client.get("/request_session") body = await response.text() assert response.status == 200 - assert "Omnidreams WebRTC Drive" in body + assert "FlashDreams WebRTC Drive" in body finally: await client.close() From 20f92100553adc30dc4801bf25aba3b7b1904ec4 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Fri, 7 Aug 2026 06:24:59 +0000 Subject: [PATCH 12/19] Unify WebRTC UI --- .../flashdreams/serving/webrtc/server.py | 30 +- .../serving/webrtc/web/mock_ui_server.py | 73 +- .../serving/webrtc/web/request_session.css | 16 + .../serving/webrtc/web/request_session.html | 23 +- .../serving/webrtc/web/request_session.js | 258 ++- flashdreams/tests/test_webrtc_serving.py | 51 +- integrations/lingbot/lingbot/demo/webrtc.py | 3 + integrations/lingbot/lingbot/webrtc/server.py | 18 +- .../lingbot/lingbot/webrtc/web/adapter.css | 312 ++++ .../lingbot/lingbot/webrtc/web/adapter.js | 526 ++++++ .../webrtc/web/assets/horizontal-dark.svg | 192 -- .../webrtc/web/assets/horizontal-light.svg | 182 -- .../lingbot/webrtc/web/request_session.css | 960 ---------- .../lingbot/webrtc/web/request_session.html | 163 -- .../lingbot/webrtc/web/request_session.js | 1622 ----------------- integrations/lingbot/pyproject.toml | 6 +- integrations/lingbot/tests/test_demo_api.py | 8 + .../lingbot/tests/test_server_routes.py | 25 +- .../omnidreams/omnidreams/demo/webrtc.py | 1 + .../omnidreams/omnidreams/webrtc/server.py | 2 + .../omnidreams/webrtc/web/adapter.js | 7 + .../webrtc/web/assets/horizontal-dark.svg | 192 -- .../webrtc/web/assets/horizontal-light.svg | 182 -- integrations/omnidreams/pyproject.toml | 2 +- .../omnidreams/tests/test_demo_api.py | 1 + .../omnidreams/tests/test_webrtc_runtime.py | 21 +- .../tests/test_webrtc_server_routes.py | 28 +- 27 files changed, 1304 insertions(+), 3600 deletions(-) create mode 100644 integrations/lingbot/lingbot/webrtc/web/adapter.css create mode 100644 integrations/lingbot/lingbot/webrtc/web/adapter.js delete mode 100644 integrations/lingbot/lingbot/webrtc/web/assets/horizontal-dark.svg delete mode 100644 integrations/lingbot/lingbot/webrtc/web/assets/horizontal-light.svg delete mode 100644 integrations/lingbot/lingbot/webrtc/web/request_session.css delete mode 100644 integrations/lingbot/lingbot/webrtc/web/request_session.html delete mode 100644 integrations/lingbot/lingbot/webrtc/web/request_session.js create mode 100644 integrations/omnidreams/omnidreams/webrtc/web/adapter.js delete mode 100644 integrations/omnidreams/omnidreams/webrtc/web/assets/horizontal-dark.svg delete mode 100644 integrations/omnidreams/omnidreams/webrtc/web/assets/horizontal-light.svg diff --git a/flashdreams/flashdreams/serving/webrtc/server.py b/flashdreams/flashdreams/serving/webrtc/server.py index a9e386533..bf257701d 100644 --- a/flashdreams/flashdreams/serving/webrtc/server.py +++ b/flashdreams/flashdreams/serving/webrtc/server.py @@ -34,6 +34,7 @@ async def shutdown(self) -> None: ... def create_webrtc_app( *, web_dir: Path, + model_web_dir: Path | None = None, session_manager: WebRTCSessionManager, request_session_url: str, index_filename: str = "request_session.html", @@ -89,6 +90,12 @@ async def healthz(request: web.Request) -> web.StreamResponse: } ) + async def ui_config(_: web.Request) -> web.StreamResponse: + adapter_module = None + if model_web_dir is not None and (model_web_dir / "adapter.js").is_file(): + adapter_module = "/model-static/adapter.js?v=model-ui-v1" + return web.json_response({"adapter_module": adapter_module}) + async def on_startup(app: web.Application) -> None: manager = app[SESSION_MANAGER_KEY] logger.info("Preloading {} runtime on startup.", preload_name) @@ -104,7 +111,10 @@ async def on_shutdown(app: web.Application) -> None: app.router.add_get("/request_session", request_session_page) app.router.add_post("/api/webrtc/offer", offer) app.router.add_get("/healthz", healthz) + app.router.add_get("/api/ui/config", ui_config) app.router.add_static("/static/", web_dir, show_index=False) + if model_web_dir is not None: + app.router.add_static("/model-static/", model_web_dir, show_index=False) app.on_startup.append(on_startup) app.on_shutdown.append(on_shutdown) return app @@ -117,6 +127,7 @@ async def close_package_resources(app: web.Application) -> None: def create_packaged_webrtc_app( *, web_resource: Any, + model_web_resource: Any | None = None, session_manager: WebRTCSessionManager, request_session_url: str, preload_name: str, @@ -136,13 +147,18 @@ def create_packaged_webrtc_app( resource_stack = ExitStack() try: web_dir = resource_stack.enter_context(as_file_fn(web_resource)) - app = create_app_fn( - web_dir=web_dir, - session_manager=session_manager, - preload_name=preload_name, - request_session_url=request_session_url, - index_filename=index_filename, - ) + create_kwargs: dict[str, Any] = { + "web_dir": web_dir, + "session_manager": session_manager, + "preload_name": preload_name, + "request_session_url": request_session_url, + "index_filename": index_filename, + } + if model_web_resource is not None: + create_kwargs["model_web_dir"] = resource_stack.enter_context( + as_file_fn(model_web_resource) + ) + app = create_app_fn(**create_kwargs) if configure_app is not None: configure_app(app) app[PACKAGE_RESOURCE_STACK_KEY] = resource_stack diff --git a/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py b/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py index 5ad682b82..c11c9cee3 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py +++ b/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py @@ -16,16 +16,27 @@ from __future__ import annotations import argparse +import json from functools import partial from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer from importlib.resources import as_file, files +from pathlib import Path from urllib.parse import urlsplit WEB_DIR_RESOURCE = files("flashdreams.serving.webrtc").joinpath("web") class MockUIRequestHandler(SimpleHTTPRequestHandler): - """Serve the static viewer without preloading the Lingbot runtime.""" + """Serve the static viewer without preloading a model runtime.""" + + def __init__( + self, + *args: object, + model_web_dir: Path | None = None, + **kwargs: object, + ) -> None: + self.model_web_dir = model_web_dir + super().__init__(*args, **kwargs) def _rewrite_path(self) -> bool: path = urlsplit(self.path).path @@ -41,27 +52,85 @@ def _rewrite_path(self) -> bool: return False def do_GET(self) -> None: + if self._serve_ui_config(): + return + if self._serve_model_asset(head_only=False): + return if self._rewrite_path(): return super().do_GET() def do_HEAD(self) -> None: + if self._serve_ui_config(): + return + if self._serve_model_asset(head_only=True): + return if self._rewrite_path(): return super().do_HEAD() + def _serve_ui_config(self) -> bool: + if urlsplit(self.path).path != "/api/ui/config": + return False + adapter_module = ( + "/model-static/adapter.js?v=model-ui-v1" + if self.model_web_dir is not None + and (self.model_web_dir / "adapter.js").is_file() + else None + ) + payload = json.dumps({"adapter_module": adapter_module}).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + if self.command != "HEAD": + self.wfile.write(payload) + return True + + def _serve_model_asset(self, *, head_only: bool) -> bool: + path = urlsplit(self.path).path + if not path.startswith("/model-static/") or self.model_web_dir is None: + return False + relative = Path(path.removeprefix("/model-static/")) + if relative.is_absolute() or ".." in relative.parts: + self.send_error(404) + return True + original_directory = self.directory + original_path = self.path + try: + self.directory = str(self.model_web_dir) + self.path = "/" + relative.as_posix() + if head_only: + super().do_HEAD() + else: + super().do_GET() + finally: + self.directory = original_directory + self.path = original_path + return True + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Serve the shared WebRTC mock UI.") parser.add_argument("--host", type=str, default="127.0.0.1") parser.add_argument("--port", type=int, default=8090) + parser.add_argument( + "--model-web-dir", + type=Path, + default=None, + help="Optional integration web directory containing adapter.js.", + ) return parser.parse_args() def main() -> None: args = parse_args() with as_file(WEB_DIR_RESOURCE) as web_dir: - handler = partial(MockUIRequestHandler, directory=str(web_dir)) + handler = partial( + MockUIRequestHandler, + directory=str(web_dir), + model_web_dir=args.model_web_dir, + ) server = ThreadingHTTPServer((args.host, args.port), handler) print( f"Serving shared mock UI at http://{args.host}:{args.port}/request_session?mock=1" diff --git a/flashdreams/flashdreams/serving/webrtc/web/request_session.css b/flashdreams/flashdreams/serving/webrtc/web/request_session.css index 890c36147..237804419 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/request_session.css +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.css @@ -88,6 +88,22 @@ select { transition: opacity 220ms ease; } +.modelStageSlot { + position: absolute; + z-index: 2; + inset: 0; + pointer-events: none; +} + +.modelPanelSlot { + display: contents; +} + +.modelStatusSlot:empty, +.modelControlSlot:empty { + display: none; +} + body.has-video .stageVideo { opacity: 1; } diff --git a/flashdreams/flashdreams/serving/webrtc/web/request_session.html b/flashdreams/flashdreams/serving/webrtc/web/request_session.html index 321532e5a..ee42c82a0 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/request_session.html +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.html @@ -9,7 +9,7 @@ FlashDreams WebRTC Drive - +
@@ -18,6 +18,7 @@

FlashDreams WebRTC Drive

+
@@ -30,32 +31,26 @@

FlashDreams WebRTC Drive

Idle - + +
Flow waiting
+
+

Controls

-
-
-
- - - - -
- Drive / Turn -
-
+
+
@@ -92,6 +87,6 @@

Client Logs

- + diff --git a/flashdreams/flashdreams/serving/webrtc/web/request_session.js b/flashdreams/flashdreams/serving/webrtc/web/request_session.js index a9ae9e74a..8f708f0ba 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/request_session.js +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.js @@ -1,6 +1,8 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +const mockMode = new URLSearchParams(window.location.search).has("mock") + const connectButton = document.getElementById("connectButton") const statusText = document.getElementById("statusText") const flowText = document.getElementById("flowText") @@ -15,9 +17,23 @@ const stepValue = document.getElementById("stepValue") const modelValue = document.getElementById("modelValue") const postprocessField = document.getElementById("postprocessField") const postprocessSelect = document.getElementById("postprocessSelect") -const controlButtons = Array.from(document.querySelectorAll("[data-control-key]")) - -const allowedKeys = new Set(["w", "a", "s", "d"]) +const modelStageSlot = document.getElementById("modelStageSlot") +const modelStatusSlot = document.getElementById("modelStatusSlot") +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"], @@ -32,6 +48,10 @@ const pendingActions = [] const maxPendingActions = 32 const heartbeatIntervalMs = 2000 +let allowedKeys = new Set() +let controlButtons = [] +let modelAdapter = null + let peerConnection = null let controlChannel = null let statsTimer = null @@ -41,7 +61,7 @@ let inferenceInFlight = false let connected = false let disconnecting = false let heldKeySequence = 0 -let postprocessControlAvailable = false +let postprocessAvailable = false const metrics = { fps: null, @@ -58,6 +78,21 @@ function normalizeKey(rawKey) { return keyAliases.get(key) || key } +function isEditableControlTarget(target) { + if (!target || typeof target !== "object") { + return false + } + if (target.isContentEditable === true) { + return true + } + const tagName = typeof target.tagName === "string" ? target.tagName.toLowerCase() : "" + if (["input", "textarea", "select"].includes(tagName)) { + return true + } + return typeof target.closest === "function" + && target.closest("input, textarea, select, [contenteditable]") !== null +} + function formatTime() { return new Date().toLocaleTimeString([], { hour12: false }) } @@ -120,56 +155,76 @@ function setFlow(message) { function setVideoVisible(visible) { document.body.classList.toggle("has-video", visible) + modelAdapter?.onVideoVisibilityChanged?.(visible, modelContext) } -function setPostprocessDisabled(disabled) { - if (!postprocessSelect) { - return +function renderControls(groups) { + controlRows.replaceChildren() + allowedKeys = new Set() + for (const group of groups) { + if (!group || !Array.isArray(group.keys) || group.keys.length === 0) { + continue + } + const row = document.createElement("div") + row.className = "controlRow" + const cluster = document.createElement("div") + cluster.className = group.keys.length > 2 ? "keyCluster keyClusterWide" : "keyCluster" + for (const item of group.keys) { + const key = normalizeKey(typeof item === "string" ? item : item?.key) + if (!key) { + continue + } + allowedKeys.add(key) + const button = document.createElement("button") + button.className = "controlKey" + button.type = "button" + button.dataset.controlKey = key + button.textContent = key.toUpperCase() + button.setAttribute("aria-label", typeof item === "string" ? key : (item.label || key)) + cluster.append(button) + } + const label = document.createElement("span") + label.textContent = String(group.label || "Controls") + row.append(cluster, label) + controlRows.append(row) } - postprocessSelect.disabled = disabled || !postprocessControlAvailable + controlButtons = Array.from(controlRows.querySelectorAll("[data-control-key]")) +} + +function setPostprocessDisabled(disabled) { + postprocessSelect.disabled = disabled || !postprocessAvailable } async function loadPostprocessOptions() { - if (!postprocessField || !postprocessSelect) { - return - } - const response = await fetch("/api/postprocess/options") - if (response.status === 404) { - postprocessControlAvailable = false - postprocessField.hidden = true - return - } - if (!response.ok) { - throw new Error(`post-process options failed (${response.status})`) - } - const payload = await response.json() + const payload = mockMode + ? { default_preset: "rtx-super-resolution", presets: ["rtx-super-resolution"] } + : await fetch("/api/postprocess/options").then(async (response) => { + if (!response.ok) { + throw new Error(`post-process options failed (${response.status})`) + } + return response.json() + }) const presets = Array.isArray(payload.presets) ? payload.presets : [] const defaultPreset = typeof payload.default_preset === "string" ? payload.default_preset : "" - postprocessControlAvailable = Boolean(defaultPreset && presets.includes(defaultPreset)) - postprocessField.hidden = !postprocessControlAvailable - postprocessSelect.replaceChildren() - - const offOption = document.createElement("option") - offOption.value = "" - offOption.textContent = "Off" - postprocessSelect.append(offOption) + postprocessAvailable = Boolean(defaultPreset && presets.includes(defaultPreset)) + postprocessField.hidden = !postprocessAvailable + postprocessSelect.replaceChildren(new Option("Off", "")) for (const preset of presets) { - if (typeof preset !== "string" || !preset) { - continue + if (typeof preset === "string" && preset) { + postprocessSelect.append(new Option(preset, preset)) } - const option = document.createElement("option") - option.value = preset - option.textContent = preset - postprocessSelect.append(option) } - postprocessSelect.value = postprocessControlAvailable ? defaultPreset : "" + postprocessSelect.value = postprocessAvailable ? defaultPreset : "" setPostprocessDisabled(false) + if (postprocessAvailable) { + logEvent(`post-process=${postprocessSelect.value}`, { source: "client" }) + } } -async function configureSessionInput() { - if (!postprocessControlAvailable || !postprocessSelect) { +async function configurePostprocessSession() { + if (!postprocessAvailable || mockMode) { return } const postprocessPreset = postprocessSelect.value @@ -182,10 +237,97 @@ async function configureSessionInput() { const text = await response.text() throw new Error(`session configuration failed (${response.status}): ${text}`) } - logEvent( - `post-process=${postprocessPreset || "off"}`, - { source: "client" } - ) + logEvent(`post-process=${postprocessPreset || "off"}`, { source: "client" }) +} + +function sendModelMessage(payload) { + if (!connected || !controlChannel || controlChannel.readyState !== "open") { + return false + } + controlChannel.send(JSON.stringify(payload)) + return true +} + +function sendModelCommand(payload, label = "model command") { + if (!sendModelMessage(payload)) { + setFlow("connect session first") + return false + } + inferenceInFlight = true + setStatus("Generating", "generating") + setFlow(`sent ${label}`) + logEvent(label, { source: "client" }) + return true +} + +const modelContext = { + slots: { + stage: modelStageSlot, + status: modelStatusSlot, + panel: modelPanelSlot, + controls: modelControlSlot, + }, + isVideoVisible: () => document.body.classList.contains("has-video"), + logEvent, + releaseControls: releaseAllKeys, + sendCommand: sendModelCommand, + setModelName(name) { + if (typeof name === "string" && name) { + metrics.model = name + renderMetrics() + } + }, + setResolution(width, height) { + if (Number.isFinite(Number(width)) && Number.isFinite(Number(height))) { + metrics.resolution = `${Number(width)}x${Number(height)}` + renderMetrics() + } + }, +} + +async function loadModelAdapter() { + let adapter = {} + try { + const response = await fetch("/api/ui/config") + if (response.ok) { + const config = await response.json() + if (typeof config.adapter_module === "string" && config.adapter_module) { + const module = await import(config.adapter_module) + if (module.default && typeof module.default === "object") { + adapter = module.default + } + } + } + } catch (error) { + logEvent(`model UI unavailable: ${error.message}`, { source: "client", level: "error" }) + } + + modelAdapter = adapter + if (typeof adapter.stylesheet === "string" && adapter.stylesheet) { + const stylesheet = document.createElement("link") + stylesheet.rel = "stylesheet" + stylesheet.href = adapter.stylesheet + document.head.append(stylesheet) + } + const modelControls = Array.isArray(adapter.controls) ? adapter.controls : [] + renderControls([...defaultControls, ...modelControls]) + if (typeof adapter.modelName === "string") { + modelContext.setModelName(adapter.modelName) + } + if (adapter.enablePostprocess === true) { + try { + await loadPostprocessOptions() + } catch (error) { + postprocessAvailable = false + postprocessField.hidden = true + setPostprocessDisabled(false) + logEvent(`post-process unavailable: ${error.message}`, { + source: "client", + level: "error", + }) + } + } + await adapter.mount?.(modelContext) } function renderMetrics() { @@ -437,6 +579,7 @@ function sendControlAction(action) { action, }) ) + modelAdapter?.onActionSent?.(action, modelContext) recordActionSent(action) setStatus("Generating", "generating") setFlow(`sent ${actionLabel(action)}, waiting=${inferenceInFlight}`) @@ -544,6 +687,11 @@ function handleControlMessage(rawMessage) { if (activeKeys.size > 0) { enqueueHeldKeyRepeats() } + modelAdapter?.onControlMessage?.(payload, modelContext) + return + } + + if (modelAdapter?.onControlMessage?.(payload, modelContext)) { return } @@ -710,6 +858,7 @@ function disconnectSession({ notify = true } = {}) { connected = false connectButton.disabled = false setPostprocessDisabled(false) + modelAdapter?.onDisconnect?.(modelContext) if (notify && controlChannel && controlChannel.readyState === "open") { try { controlChannel.send(JSON.stringify({ type: "disconnect" })) @@ -739,7 +888,8 @@ async function connectSession() { disconnecting = false try { - await configureSessionInput() + await configurePostprocessSession() + await modelAdapter?.beforeConnect?.(modelContext) const pc = new RTCPeerConnection() const channel = pc.createDataChannel("controls") peerConnection = pc @@ -755,6 +905,7 @@ async function connectSession() { } channel.onclose = () => { connected = false + setPostprocessDisabled(false) if (document.body.dataset.status !== "error") { setStatus("Closed", "idle") } @@ -762,6 +913,7 @@ async function connectSession() { logEvent("control data channel closed", { source: "client" }) stopHeartbeat() stopStatsPolling() + modelAdapter?.onDisconnect?.(modelContext) if (!disconnecting && pc.connectionState !== "closed") { pc.close() } @@ -799,6 +951,7 @@ async function connectSession() { connected = false connectButton.disabled = false setPostprocessDisabled(false) + modelAdapter?.onDisconnect?.(modelContext) stopHeartbeat() stopStatsPolling() setStatus(state === "failed" ? "Error" : "Idle", state === "failed" ? "error" : "idle") @@ -854,10 +1007,14 @@ async function connectSession() { logEvent(`connect failed: ${error.message}`, { source: "client", level: "error" }) connectButton.disabled = false setPostprocessDisabled(false) + modelAdapter?.onDisconnect?.(modelContext) } } function handleKeyDown(event) { + if (isEditableControlTarget(event.target)) { + return + } const key = normalizeKey(event.key) if (!allowedKeys.has(key)) { return @@ -871,6 +1028,9 @@ function handleKeyDown(event) { } function handleKeyUp(event) { + if (isEditableControlTarget(event.target)) { + return + } const key = normalizeKey(event.key) if (!allowedKeys.has(key)) { return @@ -920,20 +1080,16 @@ function startVideoFrameMonitor() { remoteVideo.requestVideoFrameCallback(onFrame) } -function initialize() { +async function initialize() { document.body.dataset.status = "idle" logEvent("viewer ready", { source: "client" }) setFlow("waiting") renderMetrics() + await loadModelAdapter() attachPointerControls() window.requestAnimationFrame(drawIdleScene) startVideoFrameMonitor() - void loadPostprocessOptions().catch((error) => { - logEvent(`post-process options unavailable: ${error.message}`, { - source: "client", - level: "error", - }) - }) + connectButton.disabled = false } connectButton.addEventListener("click", () => { @@ -957,4 +1113,4 @@ window.addEventListener("beforeunload", () => { disconnectSession() }) -initialize() +void initialize() diff --git a/flashdreams/tests/test_webrtc_serving.py b/flashdreams/tests/test_webrtc_serving.py index 43a2c066a..c7b8a61b7 100644 --- a/flashdreams/tests/test_webrtc_serving.py +++ b/flashdreams/tests/test_webrtc_serving.py @@ -174,15 +174,26 @@ def _raise_creation_failure(**_kwargs) -> web.Application: assert closed -def test_shared_viewer_treats_postprocess_routes_as_optional() -> None: +def test_shared_viewer_exposes_model_extension_slots() -> None: web_dir = files("flashdreams.serving.webrtc").joinpath("web") html = web_dir.joinpath("request_session.html").read_text(encoding="utf-8") javascript = web_dir.joinpath("request_session.js").read_text(encoding="utf-8") - assert "/static/request_session.js?v=shared-webrtc-v1" in html - assert "if (!postprocessField || !postprocessSelect)" in javascript - assert "if (response.status === 404)" in javascript - assert "if (!postprocessControlAvailable || !postprocessSelect)" in javascript + assert "/static/request_session.js?v=shared-webrtc-v3" in html + for slot in ( + "modelStageSlot", + "modelStatusSlot", + "modelPanelSlot", + "modelControlSlot", + ): + assert f'id="{slot}"' in html + assert 'fetch("/api/ui/config")' in javascript + assert "await modelAdapter?.beforeConnect?.(modelContext)" in javascript + assert "sendCommand: sendModelCommand" in javascript + assert 'id="postprocessField"' in html + assert 'fetch("/api/postprocess/options")' in javascript + assert "adapter.enablePostprocess === true" in javascript + assert "/api/session/initial_scene" not in javascript @pytest.mark.asyncio @@ -212,6 +223,36 @@ async def test_packaged_webrtc_app_serves_common_routes(tmp_path) -> None: await client.close() +@pytest.mark.asyncio +async def test_packaged_webrtc_app_serves_model_adapter(tmp_path) -> None: + shared_dir = tmp_path / "shared" + model_dir = tmp_path / "model" + shared_dir.mkdir() + model_dir.mkdir() + (shared_dir / "request_session.html").write_text("session") + (model_dir / "adapter.js").write_text("export default {}") + app = create_packaged_webrtc_app( + web_resource=shared_dir, + model_web_resource=model_dir, + session_manager=_FakeSessionManager(), + request_session_url="http://127.0.0.1:8080/request_session", + preload_name="Test", + as_file_fn=lambda resource: nullcontext(resource), + ) + client = TestClient(TestServer(app)) + await client.start_server() + try: + config_response = await client.get("/api/ui/config") + assert await config_response.json() == { + "adapter_module": "/model-static/adapter.js?v=model-ui-v1" + } + adapter_response = await client.get("/model-static/adapter.js") + assert adapter_response.status == 200 + assert await adapter_response.text() == "export default {}" + finally: + await client.close() + + def test_webrtc_message_helpers_preserve_public_payload_shape() -> None: assert make_error_payload("boom") == {"type": "error", "message": "boom"} assert make_event_ack_payload( diff --git a/integrations/lingbot/lingbot/demo/webrtc.py b/integrations/lingbot/lingbot/demo/webrtc.py index 0b27c1f74..c9036670f 100644 --- a/integrations/lingbot/lingbot/demo/webrtc.py +++ b/integrations/lingbot/lingbot/demo/webrtc.py @@ -20,6 +20,7 @@ LingbotRuntimeConfig, LingbotWebRTCSessionManager, ) +from lingbot.webrtc.server import configure_lingbot_webrtc_app class LingbotDemoWebRTCSessionManager(LingbotWebRTCSessionManager): @@ -51,9 +52,11 @@ def create_lingbot_webrtc_app( 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, ) diff --git a/integrations/lingbot/lingbot/webrtc/server.py b/integrations/lingbot/lingbot/webrtc/server.py index 3825a5559..1cd0bde1e 100644 --- a/integrations/lingbot/lingbot/webrtc/server.py +++ b/integrations/lingbot/lingbot/webrtc/server.py @@ -67,7 +67,8 @@ normalize_text_events, ) -WEB_DIR_RESOURCE = files("lingbot.webrtc").joinpath("web") +WEB_DIR_RESOURCE = files("flashdreams.serving.webrtc").joinpath("web") +MODEL_WEB_DIR_RESOURCE = files("lingbot.webrtc").joinpath("web") MAX_UPLOAD_IMAGE_BYTES = 15 * 1024 * 1024 MAX_PROMPT_CHARS = 2_000 @@ -174,23 +175,26 @@ def create_app( ) -> web.Application: manager = session_manager or LingbotWebRTCSessionManager() - def _configure_app(app: web.Application) -> None: - 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) - return create_packaged_webrtc_app( web_resource=WEB_DIR_RESOURCE, + model_web_resource=MODEL_WEB_DIR_RESOURCE, session_manager=manager, preload_name="Lingbot", request_session_url=request_session_url, - configure_app=_configure_app, + configure_app=configure_lingbot_webrtc_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: + """Register Lingbot-only initial-scene and session-input routes.""" + 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()) diff --git a/integrations/lingbot/lingbot/webrtc/web/adapter.css b/integrations/lingbot/lingbot/webrtc/web/adapter.css new file mode 100644 index 000000000..61e8116ee --- /dev/null +++ b/integrations/lingbot/lingbot/webrtc/web/adapter.css @@ -0,0 +1,312 @@ +/* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. */ +/* SPDX-License-Identifier: Apache-2.0 */ + +.firstFramePreview { + position: absolute; + inset: 0; + width: 100%; + height: 100%; + object-fit: cover; + opacity: 0; + transition: opacity 220ms ease; +} + +body.is-ready-preview .firstFramePreview { + opacity: 1; +} + +.sceneCard { + position: absolute; + top: clamp(120px, 15vh, 168px); + left: clamp(18px, 3vw, 52px); + display: grid; + gap: 12px; + width: min(560px, calc(100vw - 36px)); + max-height: min(680px, calc(100vh - 224px)); + padding: 16px 18px 18px; + overflow: auto; +} + +.sceneCard[hidden], +.eventControls[hidden] { + display: none; +} + +.firstFrameSourceRow { + display: grid; + grid-template-columns: 86px minmax(0, 1fr) 86px; + gap: 8px; + min-height: 68px; +} + +.firstFrameSourceRow[data-mode="upload"] { + grid-template-columns: minmax(0, 1fr) 86px 86px; +} + +.sourcePane { + min-width: 0; + overflow: hidden; + border: 1px solid rgba(255, 255, 255, 0.13); + border-radius: 8px; + background: rgba(255, 255, 255, 0.045); +} + +.firstFrameSourceRow[data-mode="upload"] .sourcePaneUpload, +.firstFrameSourceRow[data-mode="url"] .sourcePaneUrl { + border-color: rgba(142, 240, 28, 0.38); + background: rgba(255, 255, 255, 0.075); +} + +.sourceModeButton, +.uploadControl { + width: 100%; + min-height: 68px; + border: 0; + background: transparent; + color: var(--text); + cursor: pointer; + font-size: 0.78rem; + font-weight: 800; +} + +.uploadControl { + display: none; + align-items: center; + justify-content: center; + padding: 0 14px; +} + +.uploadControl input { + position: absolute; + width: 1px; + height: 1px; + opacity: 0; +} + +.firstFrameSourceRow[data-mode="upload"] .sourcePaneUpload .uploadControl { + display: flex; +} + +.firstFrameSourceRow[data-mode="upload"] .sourcePaneUpload .sourceModeButton, +.firstFrameSourceRow[data-mode="url"] .sourcePaneUrl .sourceModeButton { + display: none; +} + +.promptControl, +.textEventEditor, +.urlControl { + display: grid; + gap: 6px; + color: var(--muted); + font-size: 0.76rem; + font-weight: 700; +} + +.urlControl { + display: none; + padding: 8px; +} + +.firstFrameSourceRow[data-mode="url"] .sourcePaneUrl .urlControl { + display: grid; +} + +.promptControl textarea, +.textEventPrompt, +.textEventLabel, +.urlControl input { + width: 100%; + padding: 10px 12px; + border: 1px solid rgba(255, 255, 255, 0.18); + border-radius: 6px; + outline: 0; + background: rgba(4, 6, 7, 0.62); + color: var(--text); + font: 0.84rem/1.3 inherit; +} + +.promptControl textarea { + min-height: 94px; + resize: vertical; +} + +.promptControl textarea:focus, +.textEventPrompt:focus, +.textEventLabel:focus, +.urlControl input:focus { + outline: 2px solid rgba(142, 240, 28, 0.38); + outline-offset: 2px; +} + +.urlUpdateButton, +.textEventAddButton, +.textEventRemoveButton { + border: 1px solid rgba(142, 240, 28, 0.42); + border-radius: 7px; + background: rgba(142, 240, 28, 0.14); + color: var(--text); + cursor: pointer; + font-size: 0.76rem; + font-weight: 800; +} + +.urlUpdateButton { + min-height: 68px; +} + +.firstFrameUpdateRow { + min-height: 0; +} + +.fieldStatus { + color: var(--muted); + font-size: 0.76rem; + font-weight: 700; +} + +.fieldStatus[data-state="error"] { + color: var(--danger); +} + +.fieldStatus[data-state="success"] { + color: var(--accent); +} + +.fieldStatus[data-state="pending"] { + color: var(--warning); +} + +.textEventHeader { + display: flex; + align-items: center; + justify-content: space-between; + min-height: 32px; +} + +.textEventAddButton, +.textEventRemoveButton { + min-width: 34px; + min-height: 30px; +} + +.textEventRemoveButton { + border-color: rgba(255, 255, 255, 0.18); + background: rgba(255, 255, 255, 0.06); + color: var(--muted); +} + +.textEventList, +.textEventFields { + display: grid; + gap: 8px; +} + +.textEventRow { + display: grid; + grid-template-columns: minmax(0, 1fr) 34px; + gap: 8px; + align-items: start; +} + +.textEventPrompt { + min-height: 62px; + resize: vertical; +} + +.eventControls { + display: grid; + grid-template-columns: minmax(0, 1fr) auto; + gap: 8px; + margin-top: 16px; + padding-top: 14px; + border-top: 1px solid rgba(255, 255, 255, 0.14); +} + +.eventButtons { + display: flex; + flex-wrap: wrap; + gap: 8px; +} + +.eventButton { + min-height: 34px; + padding: 0 12px; + border: 1px solid rgba(255, 255, 255, 0.20); + border-radius: 6px; + background: rgba(12, 14, 15, 0.62); + color: var(--text); + cursor: pointer; + font-size: 0.78rem; + font-weight: 800; +} + +.eventButton.is-active { + border-color: rgba(99, 216, 255, 0.72); + background: rgba(99, 216, 255, 0.16); + color: #dff7ff; +} + +@media (min-width: 901px) { + .sceneCard, + .controlCard { + width: min(560px, calc(100vw - 36px)); + } + + .sceneCard { + max-height: clamp(280px, calc(100vh - 500px), 430px); + } + + .controlRows { + grid-template-columns: repeat(2, minmax(0, 1fr)); + column-gap: 24px; + } + + .controlRow { + grid-template-columns: minmax(96px, 176px) 1fr; + gap: 12px; + } +} + +@media (max-width: 900px) { + .stage { + min-height: max(100svh, 1660px); + } + + .sceneCard { + top: 270px; + right: 18px; + left: 18px; + width: auto; + max-height: 570px; + } + + .controlCard { + top: 860px; + bottom: auto; + } + + .logCard { + top: 1216px; + bottom: auto; + } +} + +@media (max-width: 560px) { + .stage { + min-height: max(100svh, 1900px); + } + + .logCard { + top: 1395px; + } + + .firstFrameSourceRow, + .firstFrameSourceRow[data-mode="upload"] { + grid-template-columns: 74px minmax(0, 1fr); + } + + .urlUpdateButton { + grid-column: 1 / -1; + min-height: 44px; + } +} diff --git a/integrations/lingbot/lingbot/webrtc/web/adapter.js b/integrations/lingbot/lingbot/webrtc/web/adapter.js new file mode 100644 index 000000000..d7bdb6ad2 --- /dev/null +++ b/integrations/lingbot/lingbot/webrtc/web/adapter.js @@ -0,0 +1,526 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +const mockMode = new URLSearchParams(window.location.search).has("mock") + +const controls = [ + { + label: "Strafe", + keys: [ + { key: "q", label: "Strafe left" }, + { key: "e", label: "Strafe right" }, + ], + }, + { + label: "Pitch", + keys: [ + { key: "i", label: "Pitch up" }, + { key: "k", label: "Pitch down" }, + ], + }, + { + label: "Look", + keys: [ + { key: "j", label: "Look left" }, + { key: "l", label: "Look right" }, + ], + }, +] + +let context = null +let initialScene = null +let initialSceneLocked = false +let promptEdited = false +let textEventsEdited = false +let firstFrameUrlEdited = false +let firstFrameInputMode = "url" +let selectedFirstFrameFile = null +let selectedFirstFrameUrl = null +let firstFrameSelectionCommitted = false +let activeEventId = null +let textEventDrafts = [] +let textEventSequence = 0 + +let preview = null +let sceneCard = null +let firstFrameSourceRow = null +let uploadModeButton = null +let urlModeButton = null +let firstFrameInput = null +let firstFrameUrlInput = null +let firstFrameUrlUpdateButton = null +let firstFrameUrlStatus = null +let firstFrameName = null +let promptInput = null +let textEventList = null +let addTextEventButton = null +let eventControls = null +let eventButtons = null +let clearEventButton = null + +function makeSceneCard() { + const panel = document.createElement("section") + panel.className = "sceneCard overlayPanel" + panel.setAttribute("aria-label", "Initial Scene") + panel.innerHTML = ` + Initial Scene +
+
+ + +
+
+ +
+ + +
+
+ +
+
+ +
+ +
+
+ Text Events + +
+
+
+ ` + return panel +} + +function makeEventControls() { + const root = document.createElement("div") + root.className = "eventControls" + root.hidden = true + root.innerHTML = ` +
+ + ` + return root +} + +function bindElements() { + firstFrameSourceRow = sceneCard.querySelector(".firstFrameSourceRow") + uploadModeButton = sceneCard.querySelector(".uploadModeButton") + urlModeButton = sceneCard.querySelector(".urlModeButton") + firstFrameInput = sceneCard.querySelector(".firstFrameInput") + firstFrameUrlInput = sceneCard.querySelector(".firstFrameUrlInput") + firstFrameUrlUpdateButton = sceneCard.querySelector(".urlUpdateButton") + firstFrameUrlStatus = sceneCard.querySelector(".fieldStatus") + firstFrameName = sceneCard.querySelector(".firstFrameName") + promptInput = sceneCard.querySelector(".promptControl textarea") + textEventList = sceneCard.querySelector(".textEventList") + addTextEventButton = sceneCard.querySelector(".textEventAddButton") + eventButtons = eventControls.querySelector(".eventButtons") + clearEventButton = eventControls.querySelector(".eventButtonClear") +} + +function makeTextEventId(label = "") { + const slug = String(label) + .toLowerCase() + .replace(/[^a-z0-9]+/g, "-") + .replace(/^-+|-+$/g, "") + .slice(0, 48) + textEventSequence += 1 + return `${slug || "event"}-${textEventSequence}` +} + +function makeTextEventDraft(item = {}) { + const label = String(item.label || "").trim() + return { + event_id: String(item.event_id || item.id || "").trim() || makeTextEventId(label), + label, + prompt: String(item.prompt || "").trim(), + } +} + +function setFirstFrameInputMode(mode) { + if (mode !== "upload" && mode !== "url") { + return + } + firstFrameInputMode = mode + firstFrameSourceRow.dataset.mode = mode + uploadModeButton.setAttribute("aria-pressed", mode === "upload" ? "true" : "false") + urlModeButton.setAttribute("aria-pressed", mode === "url" ? "true" : "false") +} + +function setFirstFrameStatus(message = "", state = "idle") { + firstFrameUrlStatus.textContent = message + firstFrameUrlStatus.hidden = !message + firstFrameUrlStatus.dataset.state = state +} + +function defaultFirstFrameName() { + return initialScene?.has_first_frame ? "Example Image" : "Choose Image" +} + +function clearSelectedFile() { + selectedFirstFrameFile = null + firstFrameSelectionCommitted = false + firstFrameInput.value = "" + if (selectedFirstFrameUrl) { + URL.revokeObjectURL(selectedFirstFrameUrl) + selectedFirstFrameUrl = null + } +} + +function updatePreview() { + const selected = selectedFirstFrameUrl && firstFrameSelectionCommitted + const initial = initialScene?.has_first_frame && initialScene?.first_frame_url + if (selected) { + preview.src = selectedFirstFrameUrl + } else if (initial) { + const separator = initialScene.first_frame_url.includes("?") ? "&" : "?" + preview.src = `${initialScene.first_frame_url}${separator}t=${Date.now()}` + } + document.body.classList.toggle( + "is-ready-preview", + !context.isVideoVisible() && Boolean(selected || initial), + ) +} + +function setSessionLocked(locked) { + initialSceneLocked = locked + sceneCard.hidden = locked + for (const input of sceneCard.querySelectorAll("input, textarea, button")) { + input.disabled = locked + } +} + +function renderTextEventEditor() { + textEventList.replaceChildren() + for (const [index, draft] of textEventDrafts.entries()) { + const row = document.createElement("div") + row.className = "textEventRow" + const fields = document.createElement("div") + fields.className = "textEventFields" + const label = document.createElement("input") + label.className = "textEventLabel" + label.maxLength = 64 + label.placeholder = "Label" + label.value = draft.label + const prompt = document.createElement("textarea") + prompt.className = "textEventPrompt" + prompt.rows = 2 + prompt.maxLength = 1000 + prompt.placeholder = "Event prompt" + prompt.value = draft.prompt + const remove = document.createElement("button") + remove.className = "textEventRemoveButton" + remove.type = "button" + remove.textContent = "X" + remove.setAttribute("aria-label", `Remove text event ${index + 1}`) + for (const input of [label, prompt]) { + input.disabled = initialSceneLocked + input.addEventListener("focus", context.releaseControls) + } + label.addEventListener("input", () => { + draft.label = label.value + textEventsEdited = true + }) + prompt.addEventListener("input", () => { + draft.prompt = prompt.value + textEventsEdited = true + }) + remove.disabled = initialSceneLocked + remove.addEventListener("click", () => { + textEventDrafts.splice(index, 1) + textEventsEdited = true + renderTextEventEditor() + }) + fields.append(label, prompt) + row.append(fields, remove) + textEventList.append(row) + } +} + +function collectTextEvents() { + const events = [] + const usedIds = new Set() + for (const draft of textEventDrafts) { + const label = draft.label.trim() + const prompt = draft.prompt.trim() + if (!label && !prompt) { + continue + } + if (!prompt) { + throw new Error("Each text event needs a prompt.") + } + let eventId = String(draft.event_id || "").trim() || makeTextEventId(label) + while (usedIds.has(eventId)) { + eventId = makeTextEventId(label) + } + draft.event_id = eventId + usedIds.add(eventId) + events.push({ event_id: eventId, label: label || eventId, prompt, category: "custom" }) + } + return events +} + +function renderEventControls() { + const catalog = Array.isArray(initialScene?.event_catalog) ? initialScene.event_catalog : [] + eventControls.hidden = catalog.length === 0 + eventButtons.replaceChildren() + for (const item of catalog) { + const eventId = String(item.event_id || "").trim() + if (!eventId) { + continue + } + const button = document.createElement("button") + button.className = "eventButton" + button.type = "button" + button.textContent = String(item.label || eventId) + button.classList.toggle("is-active", activeEventId === eventId) + button.addEventListener("click", () => sendTextEvent(eventId, "trigger")) + eventButtons.append(button) + } + clearEventButton.classList.toggle("is-active", activeEventId === null) +} + +function applyInitialScene(scene) { + initialScene = scene + if (!promptEdited && typeof scene.prompt === "string") { + promptInput.value = scene.prompt + } + const imageUrl = typeof scene.image_url === "string" + ? scene.image_url + : (typeof scene.default_image_url === "string" ? scene.default_image_url : "") + if (!selectedFirstFrameFile && !firstFrameUrlEdited && imageUrl) { + firstFrameUrlInput.value = imageUrl + setFirstFrameInputMode("url") + } + firstFrameName.textContent = firstFrameUrlInput.value.trim() ? "Upload Image" : defaultFirstFrameName() + activeEventId = scene.active_event_id || null + if (!textEventsEdited) { + textEventDrafts = Array.isArray(scene.event_catalog) + ? scene.event_catalog.map((item) => makeTextEventDraft(item)) + : [] + renderTextEventEditor() + } + renderEventControls() + context.setModelName(scene.model || "Lingbot") + context.setResolution(scene.resolution?.width, scene.resolution?.height) + updatePreview() +} + +function mockInitialScene() { + return { + prompt: "Drive through a cinematic city street at sunset.", + has_first_frame: false, + model: "Lingbot", + resolution: { width: 832, height: 464 }, + event_catalog: [ + { event_id: "portal", label: "Portal", prompt: "A luminous portal opens." }, + { event_id: "storm", label: "Storm", prompt: "A dramatic storm rolls in." }, + ], + } +} + +async function loadInitialScene() { + if (mockMode) { + applyInitialScene(mockInitialScene()) + return + } + const response = await fetch("/api/session/initial_scene") + if (!response.ok) { + throw new Error(`initial scene failed (${response.status})`) + } + applyInitialScene(await response.json()) +} + +function validateImageUrl(value) { + const imageUrl = value.trim() + let parsed = null + try { + parsed = new URL(imageUrl) + } catch { + throw new Error("Enter a valid http(s) image URL.") + } + if (!["http:", "https:"].includes(parsed.protocol)) { + throw new Error("Enter a valid http(s) image URL.") + } + return imageUrl +} + +async function uploadSessionInput({ includeFirstFrame = false } = {}) { + const prompt = promptInput.value.trim() + const hasPrompt = promptEdited && Boolean(prompt) + const hasFile = includeFirstFrame && firstFrameInputMode === "upload" && selectedFirstFrameFile + let imageUrl = firstFrameUrlInput.value.trim() + const hasUrl = includeFirstFrame && firstFrameInputMode === "url" && Boolean(imageUrl) + const textEvents = textEventsEdited ? collectTextEvents() : null + if (!hasPrompt && !hasFile && !hasUrl && textEvents === null) { + return + } + if (hasUrl) { + imageUrl = validateImageUrl(imageUrl) + } + if (mockMode) { + applyInitialScene({ + ...mockInitialScene(), + prompt: hasPrompt ? prompt : initialScene.prompt, + event_catalog: textEvents ?? initialScene.event_catalog, + active_event_id: activeEventId, + }) + } else { + const form = new FormData() + if (hasPrompt) form.append("prompt", prompt) + if (hasFile) form.append("image", selectedFirstFrameFile, selectedFirstFrameFile.name) + if (hasUrl) form.append("image_url", imageUrl) + if (textEvents !== null) form.append("text_events", JSON.stringify(textEvents)) + const response = await fetch("/api/session/input", { method: "POST", body: form }) + if (!response.ok) { + const text = (await response.text()).trim().replace(/^\d+:\s*/, "") + throw new Error(text || `input upload failed (${response.status})`) + } + applyInitialScene(await response.json()) + } + promptEdited = false + textEventsEdited = false + firstFrameUrlEdited = false +} + +async function updateFirstFrame() { + if (initialSceneLocked) return + try { + if (firstFrameInputMode === "upload" && !selectedFirstFrameFile) { + throw new Error("Choose an image file.") + } + if (firstFrameInputMode === "url") { + firstFrameUrlInput.value = validateImageUrl(firstFrameUrlInput.value) + clearSelectedFile() + } + setFirstFrameStatus("Updating...", "pending") + firstFrameUrlUpdateButton.disabled = true + await uploadSessionInput({ includeFirstFrame: true }) + firstFrameSelectionCommitted = true + setFirstFrameStatus("Updated", "success") + updatePreview() + } catch (error) { + setFirstFrameStatus(error.message, "error") + context.logEvent(`first frame update failed: ${error.message}`, { source: "client", level: "error" }) + } finally { + firstFrameUrlUpdateButton.disabled = initialSceneLocked + } +} + +function sendTextEvent(eventId, state) { + const label = state === "clear" ? "clear event" : `event:${eventId}` + if (!context.sendCommand({ type: "event", event_id: eventId, state }, label)) { + return + } + setSessionLocked(true) +} + +function attachListeners() { + uploadModeButton.addEventListener("click", () => { + setFirstFrameInputMode("upload") + context.releaseControls() + }) + urlModeButton.addEventListener("click", () => { + setFirstFrameInputMode("url") + context.releaseControls() + }) + firstFrameInput.addEventListener("change", () => { + setFirstFrameInputMode("upload") + const [file] = firstFrameInput.files + selectedFirstFrameFile = file || null + firstFrameSelectionCommitted = false + if (selectedFirstFrameUrl) URL.revokeObjectURL(selectedFirstFrameUrl) + selectedFirstFrameUrl = selectedFirstFrameFile ? URL.createObjectURL(selectedFirstFrameFile) : null + firstFrameName.textContent = selectedFirstFrameFile?.name || defaultFirstFrameName() + firstFrameUrlInput.value = "" + firstFrameUrlEdited = false + setFirstFrameStatus(selectedFirstFrameFile ? "Image not updated" : "", "pending") + }) + firstFrameUrlInput.addEventListener("input", () => { + setFirstFrameInputMode("url") + if (selectedFirstFrameFile) clearSelectedFile() + firstFrameUrlEdited = true + firstFrameName.textContent = firstFrameUrlInput.value.trim() ? "Upload Image" : defaultFirstFrameName() + setFirstFrameStatus(firstFrameUrlInput.value.trim() ? "URL not updated" : "", "pending") + }) + firstFrameUrlUpdateButton.addEventListener("click", () => void updateFirstFrame()) + promptInput.addEventListener("input", () => { promptEdited = true }) + addTextEventButton.addEventListener("click", () => { + textEventDrafts.push(makeTextEventDraft()) + textEventsEdited = true + renderTextEventEditor() + context.releaseControls() + }) + clearEventButton.addEventListener("click", () => sendTextEvent(activeEventId || "clear", "clear")) + for (const input of [firstFrameUrlInput, promptInput, addTextEventButton]) { + input.addEventListener("focus", context.releaseControls) + } +} + +export default { + modelName: "Lingbot", + stylesheet: new URL("./adapter.css", import.meta.url).href, + controls, + + async mount(sharedContext) { + context = sharedContext + preview = document.createElement("img") + preview.className = "firstFramePreview" + preview.alt = "" + preview.setAttribute("aria-hidden", "true") + sceneCard = makeSceneCard() + eventControls = makeEventControls() + context.slots.stage.append(preview) + context.slots.panel.append(sceneCard) + context.slots.controls.append(eventControls) + bindElements() + setFirstFrameInputMode("url") + attachListeners() + try { + await loadInitialScene() + } catch (error) { + context.logEvent(`initial scene unavailable: ${error.message}`, { source: "client", level: "error" }) + } + }, + + async beforeConnect() { + await uploadSessionInput() + }, + + onActionSent() { + setSessionLocked(true) + updatePreview() + }, + + onControlMessage(payload) { + if (payload.type === "chunk_done" && Object.prototype.hasOwnProperty.call(payload, "active_event_id")) { + activeEventId = payload.active_event_id || null + renderEventControls() + return false + } + if (payload.type === "event_ack") { + activeEventId = payload.active_event_id || null + renderEventControls() + context.logEvent(`event ${payload.event_id} ${payload.state}`) + return true + } + return false + }, + + onVideoVisibilityChanged() { + updatePreview() + }, + + onDisconnect() { + setSessionLocked(false) + updatePreview() + }, +} diff --git a/integrations/lingbot/lingbot/webrtc/web/assets/horizontal-dark.svg b/integrations/lingbot/lingbot/webrtc/web/assets/horizontal-dark.svg deleted file mode 100644 index 89b68f9d1..000000000 --- a/integrations/lingbot/lingbot/webrtc/web/assets/horizontal-dark.svg +++ /dev/null @@ -1,192 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/integrations/lingbot/lingbot/webrtc/web/assets/horizontal-light.svg b/integrations/lingbot/lingbot/webrtc/web/assets/horizontal-light.svg deleted file mode 100644 index a491910db..000000000 --- a/integrations/lingbot/lingbot/webrtc/web/assets/horizontal-light.svg +++ /dev/null @@ -1,182 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/integrations/lingbot/lingbot/webrtc/web/request_session.css b/integrations/lingbot/lingbot/webrtc/web/request_session.css deleted file mode 100644 index f11e43090..000000000 --- a/integrations/lingbot/lingbot/webrtc/web/request_session.css +++ /dev/null @@ -1,960 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -:root { - color-scheme: dark; - font-family: Inter, ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, sans-serif; - --panel-bg: rgba(16, 17, 18, 0.70); - --panel-border: rgba(255, 255, 255, 0.15); - --panel-shadow: 0 20px 60px rgba(0, 0, 0, 0.35); - --text: #f3f7f0; - --muted: #b9c0bf; - --accent: #8ef01c; - --accent-strong: #a9ff2f; - --danger: #ff685f; - --warning: #ffbf4f; - --cyan: #63d8ff; -} - -* { - box-sizing: border-box; -} - -html, -body { - min-height: 100%; -} - -body { - margin: 0; - overflow: hidden; - background: #050607; - color: var(--text); -} - -button { - font: inherit; -} - -.appShell { - min-height: 100vh; -} - -.stage { - position: relative; - min-height: 100vh; - overflow: hidden; - background: #050607; - isolation: isolate; -} - -.mockCanvas, -.stageVideo, -.firstFramePreview, -.stageVignette { - position: absolute; - inset: 0; - width: 100%; - height: 100%; -} - -.mockCanvas { - z-index: 0; - background: #0b1013; -} - -.stageVideo { - z-index: 1; - display: block; - object-fit: cover; - opacity: 0; - transition: opacity 220ms ease; -} - -.firstFramePreview { - z-index: 0; - display: block; - object-fit: cover; - opacity: 0; - transition: opacity 220ms ease; -} - -body.has-video .stageVideo { - opacity: 1; -} - -body.has-video .mockCanvas { - opacity: 0; -} - -body.is-ready-preview .firstFramePreview { - opacity: 1; -} - -body.is-ready-preview .mockCanvas { - opacity: 0; -} - -.stageVignette { - z-index: 2; - pointer-events: none; - background: none; -} - -.srOnly { - position: absolute; - width: 1px; - height: 1px; - padding: 0; - margin: -1px; - overflow: hidden; - clip: rect(0, 0, 0, 0); - white-space: nowrap; - border: 0; -} - -.overlayPanel { - z-index: 3; - border: 1px solid var(--panel-border); - border-radius: 8px; - background: var(--panel-bg); - box-shadow: var(--panel-shadow); - backdrop-filter: blur(18px) saturate(1.15); -} - -.brandOverlay { - position: absolute; - z-index: 3; - top: clamp(18px, 3vw, 42px); - left: clamp(18px, 3vw, 52px); - width: clamp(260px, 25vw, 390px); - filter: drop-shadow(0 4px 18px rgba(0, 0, 0, 0.55)); -} - -.statusLine strong, -.metric strong, -.logEntry time { - color: var(--accent-strong); -} - -.brandLogo { - display: block; - width: 100%; - height: auto; -} - -.statusCard { - position: absolute; - top: clamp(18px, 3vw, 42px); - right: clamp(18px, 3vw, 42px); - width: min(240px, calc(100vw - 36px)); - padding: 16px 18px; - display: grid; - gap: 10px; -} - -.panelLabel { - color: var(--muted); - font-size: 0.72rem; - font-weight: 700; - letter-spacing: 0.12em; - text-transform: uppercase; -} - -.statusLine { - display: flex; - align-items: center; - gap: 10px; - min-height: 30px; - font-size: clamp(1.1rem, 2vw, 1.35rem); -} - -.statusDot, -.liveDot { - display: inline-block; - width: 10px; - height: 10px; - flex: 0 0 auto; - border-radius: 999px; - background: var(--muted); - box-shadow: 0 0 0 rgba(255, 255, 255, 0); -} - -body[data-status="connected"] .statusDot, -body[data-status="waiting"] .statusDot, -body[data-status="generating"] .statusDot, -.liveDot { - background: #47d65a; - box-shadow: 0 0 16px rgba(71, 214, 90, 0.62); -} - -body[data-status="connecting"] .statusDot { - background: var(--warning); - box-shadow: 0 0 16px rgba(255, 191, 79, 0.54); -} - -body[data-status="error"] .statusDot { - background: var(--danger); - box-shadow: 0 0 16px rgba(255, 104, 95, 0.54); -} - -.connectButton { - width: 100%; - min-height: 36px; - border: 1px solid rgba(142, 240, 28, 0.45); - border-radius: 6px; - background: rgba(142, 240, 28, 0.12); - color: var(--text); - cursor: pointer; - font-weight: 700; -} - -.connectButton:hover { - background: rgba(142, 240, 28, 0.20); -} - -.connectButton:disabled { - cursor: not-allowed; - opacity: 0.52; -} - -body[data-status="connected"] .connectButton, -body[data-status="waiting"] .connectButton, -body[data-status="generating"] .connectButton { - display: none; -} - -.flowLine { - display: grid; - grid-template-columns: auto 1fr; - gap: 8px; - align-items: center; - min-width: 0; - color: var(--muted); - font-size: 0.78rem; -} - -.flowLine strong { - min-width: 0; - overflow: hidden; - color: var(--text); - text-overflow: ellipsis; - white-space: nowrap; -} - -.sceneCard { - position: absolute; - top: clamp(120px, 15vh, 168px); - left: clamp(18px, 3vw, 52px); - width: min(560px, calc(100vw - 36px)); - max-height: min(680px, calc(100vh - 224px)); - overflow: auto; - padding: 16px 18px 18px; - display: grid; - gap: 12px; -} - -.sceneCard[hidden] { - display: none; -} - -.firstFrameSourceRow { - display: grid; - grid-template-columns: 86px minmax(0, 1fr) 86px; - gap: 8px; - align-items: stretch; - min-height: 68px; - transition: grid-template-columns 160ms ease; -} - -.firstFrameSourceRow[data-mode="upload"] { - grid-template-columns: minmax(0, 1fr) 86px 86px; -} - -.sourcePane { - min-width: 0; - overflow: hidden; - border: 1px solid rgba(255, 255, 255, 0.13); - border-radius: 8px; - background: rgba(255, 255, 255, 0.045); -} - -.firstFrameSourceRow[data-mode="upload"] .sourcePaneUpload, -.firstFrameSourceRow[data-mode="url"] .sourcePaneUrl { - border-color: rgba(142, 240, 28, 0.38); - background: rgba(255, 255, 255, 0.075); -} - -.sourceModeButton { - width: 100%; - height: 100%; - min-height: 68px; - border: 0; - background: transparent; - color: rgba(243, 247, 240, 0.60); - cursor: pointer; - font-size: 0.78rem; - font-weight: 800; -} - -.sourceModeButton:hover, -.sourceModeButton:focus-visible { - color: var(--text); - background: rgba(255, 255, 255, 0.06); - outline: 0; -} - -.uploadControl { - display: none; - align-items: center; - justify-content: center; - min-height: 68px; - height: 100%; - padding: 0 14px; - overflow: hidden; - border: 0; - background: transparent; - color: var(--text); - cursor: pointer; - font-size: 0.84rem; - font-weight: 800; -} - -.uploadControl:hover { - background: rgba(142, 240, 28, 0.14); -} - -.firstFrameSourceRow[data-mode="upload"] .sourcePaneUpload .uploadControl { - display: flex; -} - -.firstFrameSourceRow[data-mode="upload"] .sourcePaneUpload .sourceModeButton, -.firstFrameSourceRow[data-mode="url"] .sourcePaneUrl .sourceModeButton { - display: none; -} - -.uploadControl input { - position: absolute; - width: 1px; - height: 1px; - opacity: 0; - pointer-events: none; -} - -.uploadControl span { - max-width: 100%; - overflow: hidden; - text-overflow: ellipsis; - white-space: nowrap; -} - -.promptControl, -.textEventEditor, -.urlControl { - display: none; - gap: 6px; - color: var(--muted); - font-size: 0.76rem; - font-weight: 700; -} - -.promptControl { - display: grid; -} - -.textEventEditor { - display: grid; -} - -.firstFrameSourceRow[data-mode="url"] .sourcePaneUrl .urlControl { - display: grid; - padding: 8px; -} - -.promptControl textarea, -.textEventPrompt, -.textEventLabel, -.urlControl input { - width: 100%; - border: 1px solid rgba(255, 255, 255, 0.18); - border-radius: 6px; - background: rgba(4, 6, 7, 0.62); - color: var(--text); - font: 0.84rem/1.3 inherit; - padding: 10px 12px; -} - -.urlControl input { - min-height: 38px; -} - -.firstFrameUpdateRow { - display: flex; - align-items: center; - min-height: 0; -} - -.firstFrameUpdateRow:empty { - display: none; -} - -.urlUpdateButton { - height: 100%; - min-height: 68px; - padding: 0 10px; - border: 1px solid rgba(142, 240, 28, 0.42); - border-radius: 8px; - background: rgba(142, 240, 28, 0.14); - color: var(--text); - cursor: pointer; - font-size: 0.78rem; - font-weight: 800; -} - -.urlUpdateButton:hover:not(:disabled), -.urlUpdateButton:focus-visible { - border-color: rgba(142, 240, 28, 0.70); - background: rgba(142, 240, 28, 0.24); -} - -.urlUpdateButton:disabled { - cursor: not-allowed; - opacity: 0.54; -} - -.fieldStatus { - color: var(--muted); - font-size: 0.76rem; - font-weight: 700; - line-height: 1.25; -} - -.fieldStatus[data-state="error"] { - color: var(--danger); -} - -.fieldStatus[data-state="success"] { - color: var(--accent); -} - -.fieldStatus[data-state="pending"] { - color: var(--warning); -} - -@media (max-width: 560px) { - .firstFrameSourceRow { - grid-template-columns: 74px minmax(0, 1fr); - } - - .firstFrameSourceRow[data-mode="upload"] { - grid-template-columns: minmax(0, 1fr) 74px; - } - - .sourceModeButton { - font-size: 0.72rem; - } - - .urlUpdateButton { - grid-column: 1 / -1; - min-height: 44px; - } -} - -.promptControl textarea { - min-height: 94px; - resize: vertical; -} - -.promptControl textarea:focus, -.textEventPrompt:focus, -.textEventLabel:focus, -.urlControl input:focus { - outline: 2px solid rgba(142, 240, 28, 0.38); - outline-offset: 2px; -} - -.textEventHeader { - display: flex; - align-items: center; - justify-content: space-between; - gap: 10px; - min-height: 32px; -} - -.textEventHeader span { - font-weight: 700; -} - -.textEventAddButton, -.textEventRemoveButton { - min-width: 34px; - min-height: 30px; - border: 1px solid rgba(142, 240, 28, 0.42); - border-radius: 6px; - background: rgba(142, 240, 28, 0.12); - color: var(--text); - cursor: pointer; - font-size: 0.76rem; - font-weight: 800; -} - -.textEventRemoveButton { - border-color: rgba(255, 255, 255, 0.18); - background: rgba(255, 255, 255, 0.06); - color: var(--muted); -} - -.textEventAddButton:hover:not(:disabled), -.textEventRemoveButton:hover:not(:disabled) { - border-color: rgba(142, 240, 28, 0.62); - background: rgba(142, 240, 28, 0.20); -} - -.textEventAddButton:disabled, -.textEventRemoveButton:disabled { - cursor: not-allowed; - opacity: 0.54; -} - -.textEventList { - display: grid; - gap: 10px; -} - -.textEventRow { - display: grid; - grid-template-columns: minmax(0, 1fr) 34px; - gap: 8px; - align-items: start; -} - -.textEventFields { - display: grid; - gap: 8px; - min-width: 0; -} - -.textEventLabel { - min-height: 34px; -} - -.textEventPrompt { - min-height: 62px; - resize: vertical; -} - -.controlCard { - position: absolute; - left: clamp(18px, 3vw, 48px); - bottom: clamp(96px, 12vh, 132px); - width: min(380px, calc(100vw - 36px)); - padding: 18px 20px 20px; -} - -.controlCard h2, -.logCard h2 { - display: flex; - align-items: center; - gap: 10px; - margin: 0 0 18px; - font-size: 1.08rem; - font-weight: 740; - letter-spacing: 0; -} - -.controlCard h2 span, -.logCard h2 span { - width: 3px; - height: 22px; - border-radius: 999px; - background: var(--accent); - box-shadow: 0 0 14px rgba(142, 240, 28, 0.42); -} - -.controlRows { - display: grid; - gap: 12px; -} - -.controlRow { - display: grid; - grid-template-columns: 176px 1fr; - gap: 18px; - align-items: center; - min-height: 38px; - color: var(--text); - font-size: 0.95rem; -} - -.keyCluster { - display: grid; - grid-auto-flow: column; - grid-auto-columns: 38px; - gap: 8px; - justify-content: start; -} - -.keyClusterWide { - grid-auto-columns: 38px; -} - -.controlKey { - width: 38px; - height: 38px; - border: 1px solid rgba(255, 255, 255, 0.22); - border-bottom-color: rgba(255, 255, 255, 0.34); - border-radius: 6px; - background: rgba(12, 14, 15, 0.62); - color: #f9fbff; - cursor: pointer; - font-weight: 750; - line-height: 1; - box-shadow: inset 0 -2px 0 rgba(0, 0, 0, 0.30); - touch-action: none; - user-select: none; -} - -.controlKey:hover { - border-color: rgba(255, 255, 255, 0.42); - background: rgba(255, 255, 255, 0.09); -} - -.controlKey.is-active { - border-color: rgba(142, 240, 28, 0.78); - background: rgba(142, 240, 28, 0.26); - color: var(--accent-strong); - box-shadow: - 0 0 18px rgba(142, 240, 28, 0.32), - inset 0 0 0 1px rgba(142, 240, 28, 0.18); - transform: translateY(1px); -} - -.eventControls { - display: grid; - grid-template-columns: minmax(0, 1fr) auto; - gap: 8px; - margin-top: 16px; - padding-top: 14px; - border-top: 1px solid rgba(255, 255, 255, 0.14); -} - -.eventButtons { - display: flex; - flex-wrap: wrap; - gap: 8px; - min-width: 0; -} - -.eventButton { - min-height: 34px; - padding: 0 12px; - border: 1px solid rgba(255, 255, 255, 0.20); - border-radius: 6px; - background: rgba(12, 14, 15, 0.62); - color: var(--text); - cursor: pointer; - font-size: 0.78rem; - font-weight: 800; -} - -.eventButton:hover { - border-color: rgba(255, 255, 255, 0.42); - background: rgba(255, 255, 255, 0.09); -} - -.eventButton.is-active { - border-color: rgba(99, 216, 255, 0.72); - background: rgba(99, 216, 255, 0.16); - color: #dff7ff; -} - -.eventButtonClear { - color: var(--muted); -} - -.logCard { - position: absolute; - right: clamp(18px, 3vw, 42px); - bottom: clamp(78px, 9vh, 98px); - width: min(400px, calc(100vw - 36px)); - padding: 18px 20px 16px; -} - -.logList { - display: grid; - align-content: start; - gap: 6px; - min-height: 150px; - max-height: 300px; - overflow: auto; - padding-bottom: 12px; - border-bottom: 1px solid rgba(255, 255, 255, 0.16); - font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, monospace; - font-size: 0.78rem; - line-height: 1.4; -} - -.logEntry { - display: grid; - grid-template-columns: 76px 1fr; - gap: 8px; - min-width: 0; - color: #f3f6f7; -} - -.logEntry span { - min-width: 0; - overflow-wrap: anywhere; - white-space: pre-wrap; -} - -.logEntry.is-error span { - color: #ffb0aa; -} - -.logEntry.is-client time { - color: var(--cyan); -} - -.logFooter { - display: flex; - align-items: center; - gap: 8px; - padding-top: 12px; - color: var(--muted); - font-size: 0.86rem; -} - -.metricsBar { - position: absolute; - z-index: 3; - left: 50%; - bottom: clamp(20px, 4vh, 40px); - display: grid; - grid-template-columns: - minmax(max-content, 0.7fr) - minmax(max-content, 1fr) - minmax(max-content, 1.15fr) - minmax(max-content, 0.7fr) - minmax(max-content, 1.45fr); - width: min(960px, calc(100vw - 36px)); - overflow: hidden; - border: 1px solid rgba(255, 255, 255, 0.14); - border-radius: 8px; - background: rgba(18, 20, 22, 0.74); - box-shadow: 0 12px 38px rgba(0, 0, 0, 0.32); - backdrop-filter: blur(16px) saturate(1.12); - transform: translateX(-50%); -} - -.metric { - display: grid; - grid-template-columns: max-content max-content; - align-items: center; - justify-content: center; - gap: 10px; - min-width: max-content; - min-height: 42px; - padding: 0 16px; - border-right: 1px solid rgba(255, 255, 255, 0.11); - font-size: 0.88rem; - text-align: center; -} - -.metric:last-child { - border-right: 0; -} - -.metric span { - color: var(--muted); - white-space: nowrap; -} - -.metric strong { - text-align: left; - white-space: nowrap; - font-weight: 760; -} - -body[data-status="generating"] .statusLine strong { - animation: statusPulse 1.2s ease-in-out infinite; -} - -@keyframes statusPulse { - 0%, - 100% { - opacity: 1; - } - 50% { - opacity: 0.62; - } -} - -@media (max-width: 900px) { - body { - overflow: auto; - } - - .stage { - min-height: max(100svh, 1120px); - } - - .brandOverlay { - top: 18px; - left: 18px; - width: min(310px, calc(100vw - 36px)); - } - - .statusCard { - top: 88px; - left: 18px; - right: auto; - } - - .sceneCard { - top: 242px; - left: 18px; - right: 18px; - width: auto; - } - - .controlCard, - .logCard { - left: 18px; - right: 18px; - width: auto; - } - - .controlCard { - bottom: 360px; - } - - .logCard { - bottom: 164px; - } - - .metricsBar { - left: 18px; - right: 18px; - bottom: 18px; - grid-template-columns: repeat(2, minmax(0, 1fr)); - width: auto; - transform: none; - } - - .metric { - padding: 0 14px; - } - - .metric:nth-child(2n) { - border-right: 0; - } - - .metric:last-child { - grid-column: 1 / -1; - } -} - -/* Collapsible + movable overlay panels */ -.panelDragHandle { - display: flex; - align-items: center; - gap: 10px; - cursor: move; - touch-action: none; - user-select: none; -} - -.panelDragHandle .panelCollapseButton { - margin-left: auto; -} - -.panelCollapseButton { - display: inline-flex; - align-items: center; - justify-content: center; - width: 24px; - height: 24px; - padding: 0; - border: 1px solid rgba(255, 255, 255, 0.22); - border-radius: 6px; - background: rgba(255, 255, 255, 0.06); - color: var(--text); - cursor: pointer; - font-size: 1rem; - font-weight: 800; - line-height: 1; - flex: 0 0 auto; -} - -.panelCollapseButton:hover { - border-color: rgba(255, 255, 255, 0.42); - background: rgba(255, 255, 255, 0.12); -} - -.overlayPanel.is-floating { - right: auto !important; - bottom: auto !important; -} - -.overlayPanel.is-collapsed { - max-height: none !important; - overflow: visible !important; -} - -.overlayPanel.is-collapsed > :not(.panelDragHandle) { - display: none !important; -} - -.overlayPanel.is-collapsed .panelDragHandle { - margin-bottom: 0; -} - -@media (max-width: 520px) { - .controlRow { - grid-template-columns: 1fr; - gap: 8px; - } - - .keyCluster { - grid-auto-columns: minmax(34px, 1fr); - } - - .controlKey { - width: 100%; - } - - .logList { - min-height: 112px; - max-height: 150px; - } - - .metricsBar { - grid-template-columns: 1fr; - } - - .metric { - min-height: 36px; - border-right: 0; - border-bottom: 1px solid rgba(255, 255, 255, 0.11); - } - - .metric:last-child { - grid-column: auto; - border-bottom: 0; - } -} diff --git a/integrations/lingbot/lingbot/webrtc/web/request_session.html b/integrations/lingbot/lingbot/webrtc/web/request_session.html deleted file mode 100644 index 08806494f..000000000 --- a/integrations/lingbot/lingbot/webrtc/web/request_session.html +++ /dev/null @@ -1,163 +0,0 @@ - - - - - - - - Lingbot WebRTC Viewer - - - -
-
-

Lingbot WebRTC Viewer

- - - - - -
- -
- -
- Status -
- - Idle -
- -
- Flow - waiting -
-
- -
- Initial Scene -
-
- - -
-
- -
- - -
-
- -
-
- -
- -
-
- Text Events - -
-
-
-
- -
-

Controls

-
-
-
- - - - -
- Drive / Turn -
-
-
- - -
- Strafe -
-
-
- - -
- Pitch -
-
-
- - -
- Look -
-
- -
- -
-

Client Logs

-
-
- - Waiting -
-
- -
-
- FPS - -- -
-
- Latency - -- -
-
- Resolution - -- -
-
- Step - -- -
-
- World Model - Lingbot -
-
-
-
- - - - diff --git a/integrations/lingbot/lingbot/webrtc/web/request_session.js b/integrations/lingbot/lingbot/webrtc/web/request_session.js deleted file mode 100644 index 4178a73db..000000000 --- a/integrations/lingbot/lingbot/webrtc/web/request_session.js +++ /dev/null @@ -1,1622 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -const connectButton = document.getElementById("connectButton") -const statusText = document.getElementById("statusText") -const flowText = document.getElementById("flowText") -const eventLog = document.getElementById("eventLog") -const logState = document.getElementById("logState") -const remoteVideo = document.getElementById("remoteVideo") -const mockCanvas = document.getElementById("mockCanvas") -const firstFramePreview = document.getElementById("firstFramePreview") -const sceneCard = document.getElementById("sceneCard") -const firstFrameSourceRow = document.getElementById("firstFrameSourceRow") -const uploadModeButton = document.getElementById("uploadModeButton") -const urlModeButton = document.getElementById("urlModeButton") -const firstFrameInput = document.getElementById("firstFrameInput") -const firstFrameUrlInput = document.getElementById("firstFrameUrlInput") -const firstFrameUrlUpdateButton = document.getElementById("firstFrameUrlUpdateButton") -const firstFrameUrlStatus = document.getElementById("firstFrameUrlStatus") -const firstFrameName = document.getElementById("firstFrameName") -const promptInput = document.getElementById("promptInput") -const textEventList = document.getElementById("textEventList") -const addTextEventButton = document.getElementById("addTextEventButton") -const fpsValue = document.getElementById("fpsValue") -const latencyValue = document.getElementById("latencyValue") -const resolutionValue = document.getElementById("resolutionValue") -const stepValue = document.getElementById("stepValue") -const modelValue = document.getElementById("modelValue") -const controlButtons = Array.from(document.querySelectorAll("[data-control-key]")) -const eventControls = document.getElementById("eventControls") -const eventButtons = document.getElementById("eventButtons") -const clearEventButton = document.getElementById("clearEventButton") - -const params = new URLSearchParams(window.location.search) -const mockMode = params.has("mock") && params.get("mock") !== "0" -const allowedKeys = new Set(["w", "a", "s", "d", "q", "e", "i", "j", "k", "l"]) -const keySources = new Map() -const heldKeyOrder = new Map() -const activeKeys = new Set() -const frameTimes = [] -const pendingActions = [] -const maxPendingActions = 32 -const heartbeatIntervalMs = 2000 - -let peerConnection = null -let controlChannel = null -let statsTimer = null -let heartbeatTimer = null -let inferenceInFlight = false -let connected = false -let disconnecting = false -let heldKeySequence = 0 -let mockChunkIndex = 0 -let mockGenerationStarted = false -let mockChunkTimer = null -let actionStarted = false -let initialSceneLocked = false -let promptEdited = false -let textEventsEdited = false -let firstFrameUrlEdited = false -let firstFrameInputMode = "url" -let initialScene = null -let selectedFirstFrameUrl = null -let selectedFirstFrameFile = null -let firstFrameSelectionCommitted = false -let firstFramePreviewRefreshToken = 0 -let activeEventId = null -let textEventDrafts = [] -let textEventSequence = 0 - -const metrics = { - fps: null, - targetFps: null, - latencyMs: null, - rttMs: null, - resolution: null, - step: null, - model: "Lingbot", -} - -function normalizeKey(rawKey) { - return String(rawKey || "").toLowerCase() -} - -function makeTextEventId(label = "") { - const slug = String(label || "") - .toLowerCase() - .replace(/[^a-z0-9]+/g, "-") - .replace(/^-+|-+$/g, "") - .slice(0, 48) - textEventSequence += 1 - return `${slug || "event"}-${textEventSequence}` -} - -function isEditableControlTarget(target) { - if (!target || typeof target !== "object") { - return false - } - if (target.isContentEditable === true) { - return true - } - - const tagName = typeof target.tagName === "string" ? target.tagName.toLowerCase() : "" - if (tagName === "input" || tagName === "textarea" || tagName === "select") { - return true - } - if (typeof target.closest === "function") { - return target.closest("input, textarea, select, [contenteditable]") !== null - } - return false -} - -function formatTime() { - return new Date().toLocaleTimeString([], { hour12: false }) -} - -function firstFinite(...values) { - for (const value of values) { - if (value === null || value === undefined || value === "") { - continue - } - const number = Number(value) - if (Number.isFinite(number)) { - return number - } - } - return null -} - -function formatMs(value) { - if (!Number.isFinite(value)) { - return "--" - } - if (value >= 1000) { - return `${(value / 1000).toFixed(1)} s` - } - return `${Math.round(value)} ms` -} - -function logEvent(message, { source = "server", level = "info" } = {}) { - const entry = document.createElement("div") - entry.className = `logEntry is-${source}` - if (level === "error") { - entry.classList.add("is-error") - } - - const time = document.createElement("time") - time.textContent = `[${formatTime()}]` - const body = document.createElement("span") - body.textContent = message - entry.append(time, body) - eventLog.prepend(entry) - - while (eventLog.children.length > 36) { - eventLog.lastElementChild.remove() - } -} - -function setStatus(message, state = message.toLowerCase()) { - statusText.textContent = message - document.body.dataset.status = state - logState.textContent = state === "idle" ? "Waiting" : message -} - -function setFlow(message) { - flowText.textContent = message -} - -function setVideoVisible(visible) { - document.body.classList.toggle("has-video", visible) - updateReadyPreview() -} - -function setInitialSceneLocked(locked) { - initialSceneLocked = locked - sceneCard.hidden = locked - uploadModeButton.disabled = locked - urlModeButton.disabled = locked - firstFrameInput.disabled = locked - firstFrameUrlInput.disabled = locked - firstFrameUrlUpdateButton.disabled = locked - promptInput.disabled = locked - addTextEventButton.disabled = locked - for (const input of textEventList.querySelectorAll("input, textarea, button")) { - input.disabled = locked - } -} - -function setFirstFrameInputMode(mode) { - if (mode !== "upload" && mode !== "url") { - return - } - firstFrameInputMode = mode - firstFrameSourceRow.dataset.mode = mode - uploadModeButton.setAttribute("aria-pressed", mode === "upload" ? "true" : "false") - urlModeButton.setAttribute("aria-pressed", mode === "url" ? "true" : "false") -} - -function defaultFirstFrameName() { - return initialScene && initialScene.has_first_frame ? "Example Image" : "Choose Image" -} - -function setFirstFrameUrlStatus(message = "", state = "idle") { - firstFrameUrlStatus.textContent = message - firstFrameUrlStatus.hidden = message.length === 0 - firstFrameUrlStatus.dataset.state = state -} - -function validateFirstFrameUrl(value) { - const imageUrl = value.trim() - if (!imageUrl) { - throw new Error("Enter an image URL.") - } - let parsed - try { - parsed = new URL(imageUrl) - } catch { - throw new Error("Enter a valid image URL.") - } - if (parsed.protocol !== "http:" && parsed.protocol !== "https:") { - throw new Error("Enter an http(s) image URL.") - } - return imageUrl -} - -function clearSelectedFirstFrameFile() { - selectedFirstFrameFile = null - firstFrameSelectionCommitted = false - firstFrameInput.value = "" - if (selectedFirstFrameUrl) { - URL.revokeObjectURL(selectedFirstFrameUrl) - selectedFirstFrameUrl = null - } -} - -function clearFirstFrameUrlInput() { - firstFrameUrlInput.value = "" - firstFrameUrlEdited = false - setFirstFrameUrlStatus() -} - -function refreshedPreviewUrl(url) { - const separator = url.includes("?") ? "&" : "?" - return `${url}${separator}t=${firstFramePreviewRefreshToken}` -} - -function updateReadyPreview() { - const canPreview = !document.body.classList.contains("has-video") - const hasSelectedImage = selectedFirstFrameUrl !== null && firstFrameSelectionCommitted - const hasInitialImage = Boolean( - initialScene && initialScene.has_first_frame && initialScene.first_frame_url - ) - - if (hasSelectedImage) { - firstFramePreview.src = selectedFirstFrameUrl - } else if (hasInitialImage && initialScene.first_frame_url) { - firstFramePreview.src = refreshedPreviewUrl(initialScene.first_frame_url) - } - - document.body.classList.toggle( - "is-ready-preview", - canPreview && (hasSelectedImage || hasInitialImage) - ) -} - -function applyInitialScene(scene) { - initialScene = scene - firstFramePreviewRefreshToken = Date.now() - if (!promptEdited && typeof scene.prompt === "string") { - promptInput.value = scene.prompt - } - const sceneImageUrl = typeof scene.image_url === "string" - ? scene.image_url - : (typeof scene.default_image_url === "string" ? scene.default_image_url : "") - if (!selectedFirstFrameFile && !firstFrameUrlEdited && sceneImageUrl) { - firstFrameUrlInput.value = sceneImageUrl - setFirstFrameInputMode("url") - } - if (!selectedFirstFrameFile) { - firstFrameName.textContent = firstFrameUrlInput.value.trim() - ? "Upload Image" - : defaultFirstFrameName() - } - if (scene.model) { - metrics.model = scene.model - } - if (scene.resolution && typeof scene.resolution === "object") { - const width = Number(scene.resolution.width) - const height = Number(scene.resolution.height) - if (Number.isFinite(width) && Number.isFinite(height)) { - metrics.resolution = `${width}x${height}` - } - } - activeEventId = scene.active_event_id || null - if (!textEventsEdited) { - setTextEventDraftsFromCatalog(scene.event_catalog) - } - renderEventControls() - renderMetrics() - updateReadyPreview() -} - -function makeTextEventDraft(item = {}) { - const label = String(item.label || "").trim() - return { - event_id: String(item.event_id || item.id || "").trim() || makeTextEventId(label), - label, - prompt: String(item.prompt || "").trim(), - } -} - -function setTextEventDraftsFromCatalog(catalog) { - textEventDrafts = Array.isArray(catalog) - ? catalog.map((item) => makeTextEventDraft(item)) - : [] - renderTextEventEditor() -} - -function markTextEventsEdited() { - textEventsEdited = true -} - -function renderTextEventEditor() { - textEventList.replaceChildren() - for (const [index, draft] of textEventDrafts.entries()) { - const row = document.createElement("div") - row.className = "textEventRow" - - const fields = document.createElement("div") - fields.className = "textEventFields" - - const labelInput = document.createElement("input") - labelInput.className = "textEventLabel" - labelInput.type = "text" - labelInput.maxLength = 64 - labelInput.placeholder = "Label" - labelInput.value = draft.label - labelInput.disabled = initialSceneLocked - labelInput.addEventListener("input", () => { - draft.label = labelInput.value - markTextEventsEdited() - }) - labelInput.addEventListener("focus", releaseAllKeys) - - const promptTextarea = document.createElement("textarea") - promptTextarea.className = "textEventPrompt" - promptTextarea.rows = 2 - promptTextarea.maxLength = 1000 - promptTextarea.placeholder = "Event Prompt" - promptTextarea.value = draft.prompt - promptTextarea.disabled = initialSceneLocked - promptTextarea.addEventListener("input", () => { - draft.prompt = promptTextarea.value - markTextEventsEdited() - }) - promptTextarea.addEventListener("focus", releaseAllKeys) - - const removeButton = document.createElement("button") - removeButton.className = "textEventRemoveButton" - removeButton.type = "button" - removeButton.textContent = "X" - removeButton.setAttribute("aria-label", `Remove text event ${index + 1}`) - removeButton.disabled = initialSceneLocked - removeButton.addEventListener("click", () => { - textEventDrafts.splice(index, 1) - markTextEventsEdited() - renderTextEventEditor() - renderEventControls() - }) - - fields.append(labelInput, promptTextarea) - row.append(fields, removeButton) - textEventList.append(row) - } -} - -function collectTextEvents() { - const events = [] - const usedIds = new Set() - for (const draft of textEventDrafts) { - const label = draft.label.trim() - const prompt = draft.prompt.trim() - if (!label && !prompt) { - continue - } - if (!prompt) { - throw new Error("Each text event needs a prompt.") - } - let eventId = String(draft.event_id || "").trim() - if (!eventId) { - eventId = makeTextEventId(label) - draft.event_id = eventId - } - while (usedIds.has(eventId)) { - eventId = makeTextEventId(label) - draft.event_id = eventId - } - usedIds.add(eventId) - events.push({ - event_id: eventId, - label: label || eventId, - prompt, - category: "custom", - }) - } - return events -} - -function renderEventControls() { - const catalog = Array.isArray(initialScene && initialScene.event_catalog) - ? initialScene.event_catalog - : [] - eventControls.hidden = catalog.length === 0 - eventButtons.replaceChildren() - for (const item of catalog) { - const eventId = String(item.event_id || "").trim() - if (!eventId) { - continue - } - const button = document.createElement("button") - button.className = "eventButton" - button.type = "button" - button.textContent = String(item.label || eventId) - button.dataset.eventId = eventId - button.classList.toggle("is-active", activeEventId === eventId) - button.addEventListener("click", () => { - sendTextEvent(eventId, "trigger") - }) - eventButtons.append(button) - } - clearEventButton.hidden = catalog.length === 0 - clearEventButton.classList.toggle("is-active", activeEventId === null) -} - -async function loadInitialScene() { - if (mockMode) { - applyInitialScene({ - prompt: promptInput.value, - has_first_frame: selectedFirstFrameFile !== null || firstFrameUrlInput.value.trim().length > 0, - first_frame_url: firstFrameUrlInput.value.trim(), - image_url: firstFrameUrlInput.value.trim(), - model: metrics.model, - resolution: { width: 832, height: 464 }, - event_catalog: [ - { - event_id: "portal", - label: "Portal", - prompt: "A luminous magical portal opens in the scene.", - category: "environment", - }, - { - event_id: "storm", - label: "Storm", - prompt: "A dramatic storm rolls in with rain and lightning.", - category: "environment", - }, - { - event_id: "fireworks", - label: "Fireworks", - prompt: "Bright fireworks burst overhead.", - category: "environment", - }, - ], - input_source: selectedFirstFrameFile ? "uploaded" : "default", - }) - return - } - try { - const response = await fetch("/api/session/initial_scene") - if (!response.ok) { - throw new Error(`initial scene failed (${response.status})`) - } - applyInitialScene(await response.json()) - } catch (error) { - logEvent(`initial scene unavailable: ${error.message}`, { source: "client" }) - } -} - -async function uploadSessionInputIfNeeded({ includeFirstFrame = false } = {}) { - const prompt = promptInput.value.trim() - let imageUrl = firstFrameUrlInput.value.trim() - const hasPrompt = promptEdited && prompt.length > 0 - const hasImage = - includeFirstFrame && firstFrameInputMode === "upload" && selectedFirstFrameFile !== null - const hasImageUrl = - includeFirstFrame && firstFrameInputMode === "url" && imageUrl.length > 0 - let textEvents = null - if (textEventsEdited) { - textEvents = collectTextEvents() - } - const hasTextEvents = textEvents !== null - if (!hasPrompt && !hasImage && !hasImageUrl && !hasTextEvents) { - return - } - if (hasImageUrl) { - try { - imageUrl = validateFirstFrameUrl(imageUrl) - firstFrameUrlInput.value = imageUrl - } catch (error) { - setFirstFrameUrlStatus(error.message, "error") - throw error - } - } - - if (mockMode) { - applyInitialScene({ - prompt: hasPrompt ? prompt : promptInput.value, - has_first_frame: hasImage || hasImageUrl, - first_frame_url: hasImageUrl ? imageUrl : firstFrameUrlInput.value.trim(), - image_url: hasImageUrl ? imageUrl : firstFrameUrlInput.value.trim(), - model: metrics.model, - resolution: { width: 832, height: 464 }, - event_catalog: hasTextEvents - ? textEvents - : (initialScene ? initialScene.event_catalog : []), - active_event_id: activeEventId, - input_source: "uploaded", - }) - promptEdited = false - textEventsEdited = false - firstFrameUrlEdited = false - if (hasImage || hasImageUrl) { - setFirstFrameUrlStatus("Updated", "success") - } - return - } - - const form = new FormData() - if (hasPrompt) { - form.append("prompt", prompt) - } - if (hasImage) { - form.append("image", selectedFirstFrameFile, selectedFirstFrameFile.name) - } else if (hasImageUrl) { - form.append("image_url", imageUrl) - } - if (hasTextEvents) { - form.append("text_events", JSON.stringify(textEvents)) - } - - const response = await fetch("/api/session/input", { - method: "POST", - body: form, - }) - if (!response.ok) { - const text = (await response.text()).trim().replace(/^\d+:\s*/, "") - throw new Error(text || `input upload failed (${response.status})`) - } - promptEdited = false - textEventsEdited = false - firstFrameUrlEdited = false - applyInitialScene(await response.json()) - if (hasImage || hasImageUrl) { - setFirstFrameUrlStatus("Updated", "success") - } -} - -async function updateFirstFrameInput() { - if (initialSceneLocked) { - return - } - - if (firstFrameInputMode === "upload") { - if (!selectedFirstFrameFile) { - setFirstFrameUrlStatus("Choose an image file.", "error") - return - } - } else { - let imageUrl - try { - imageUrl = validateFirstFrameUrl(firstFrameUrlInput.value) - } catch (error) { - setFirstFrameUrlStatus(error.message, "error") - return - } - firstFrameUrlInput.value = imageUrl - clearSelectedFirstFrameFile() - } - - setFirstFrameUrlStatus("Updating...", "pending") - firstFrameUrlUpdateButton.disabled = true - - try { - await uploadSessionInputIfNeeded({ includeFirstFrame: true }) - firstFrameSelectionCommitted = true - updateReadyPreview() - setFirstFrameUrlStatus("Updated", "success") - logEvent("first frame updated", { source: "client" }) - } catch (error) { - setFirstFrameUrlStatus(error.message, "error") - logEvent(`first frame update failed: ${error.message}`, { - source: "client", - level: "error", - }) - } finally { - firstFrameUrlUpdateButton.disabled = initialSceneLocked - } -} - -function renderMetrics() { - const fps = firstFinite(metrics.fps, metrics.targetFps) - fpsValue.textContent = Number.isFinite(fps) ? String(Math.round(fps)) : "--" - latencyValue.textContent = formatMs(metrics.latencyMs) - resolutionValue.textContent = metrics.resolution || "--" - stepValue.textContent = metrics.step === null ? "--" : String(metrics.step) - modelValue.textContent = metrics.model || "Lingbot" -} - -function updateMetricsFromChunk(payload) { - const observedLatencyMs = takeObservedActionLatency() - metrics.targetFps = firstFinite(payload.fps, payload.target_fps, metrics.targetFps) - metrics.latencyMs = firstFinite( - payload.latency_ms, - payload.control_latency_ms, - observedLatencyMs, - payload.lag_ms, - payload.gen_ms, - metrics.latencyMs - ) - metrics.step = Number.isFinite(Number(payload.chunk_index)) - ? Number(payload.chunk_index) - : metrics.step - metrics.model = typeof payload.model === "string" && payload.model ? payload.model : metrics.model - - if (typeof payload.resolution === "string") { - metrics.resolution = payload.resolution - } else if (payload.resolution && typeof payload.resolution === "object") { - const width = Number(payload.resolution.width) - const height = Number(payload.resolution.height) - if (Number.isFinite(width) && Number.isFinite(height)) { - metrics.resolution = `${width}x${height}` - } - } - renderMetrics() -} - -function updateMetricsFromVideo() { - if (remoteVideo.videoWidth > 0 && remoteVideo.videoHeight > 0) { - metrics.resolution = `${remoteVideo.videoWidth}x${remoteVideo.videoHeight}` - renderMetrics() - } -} - -function recordFrame(timestamp) { - const now = Number.isFinite(timestamp) ? timestamp : performance.now() - frameTimes.push(now) - while (frameTimes.length > 0 && now - frameTimes[0] > 1200) { - frameTimes.shift() - } - if (frameTimes.length >= 2) { - const elapsed = frameTimes[frameTimes.length - 1] - frameTimes[0] - metrics.fps = elapsed > 0 ? ((frameTimes.length - 1) * 1000) / elapsed : metrics.fps - renderMetrics() - } -} - -function updateControlHighlights() { - activeKeys.clear() - for (const [key, sources] of keySources.entries()) { - if (sources.size > 0) { - activeKeys.add(key) - } - } - for (const button of controlButtons) { - const key = button.dataset.controlKey - button.classList.toggle("is-active", activeKeys.has(key)) - button.setAttribute("aria-pressed", activeKeys.has(key) ? "true" : "false") - } -} - -function actionLabel(action) { - return `${action.event}${action.key ? `:${action.key}` : ""}` -} - -function recordActionSent(action) { - pendingActions.push({ - sentAt: performance.now(), - label: actionLabel(action), - }) - while (pendingActions.length > maxPendingActions) { - pendingActions.shift() - } -} - -function takeObservedActionLatency(now = performance.now()) { - if (pendingActions.length === 0) { - return null - } - const oldest = pendingActions[0] - pendingActions.length = 0 - return Math.max(0, now - oldest.sentAt) -} - -function sendControlAction(action) { - if (mockMode && connected && !controlChannel) { - actionStarted = true - setInitialSceneLocked(true) - updateReadyPreview() - inferenceInFlight = true - mockGenerationStarted = true - recordActionSent(action) - setStatus("Generating", "generating") - setFlow(`sent ${actionLabel(action)}, waiting=true`) - logEvent(`control ${actionLabel(action)}`, { source: "client" }) - return true - } - - if (!connected || !controlChannel || controlChannel.readyState !== "open") { - return false - } - - actionStarted = true - setInitialSceneLocked(true) - updateReadyPreview() - inferenceInFlight = true - controlChannel.send( - JSON.stringify({ - type: "action", - action, - }) - ) - recordActionSent(action) - setStatus("Generating", "generating") - setFlow(`sent ${actionLabel(action)}, waiting=${inferenceInFlight}`) - logEvent(`control ${actionLabel(action)}`, { source: "client" }) - return true -} - -function sendTextEvent(eventId, state = "trigger") { - const label = state === "clear" ? "clear event" : `event:${eventId}` - if (mockMode && connected && !controlChannel) { - activeEventId = state === "clear" ? null : eventId - renderEventControls() - actionStarted = true - mockGenerationStarted = true - setStatus("Generating", "generating") - setFlow(`sent ${label}`) - logEvent(label, { source: "client" }) - return true - } - - if (!connected || !controlChannel || controlChannel.readyState !== "open") { - setFlow("connect session first") - return false - } - - actionStarted = true - setInitialSceneLocked(true) - updateReadyPreview() - inferenceInFlight = true - controlChannel.send( - JSON.stringify({ - type: "event", - event_id: eventId, - state, - }) - ) - setStatus("Generating", "generating") - setFlow(`sent ${label}`) - logEvent(label, { source: "client" }) - return true -} - -function enqueueAction(action) { - const sent = sendControlAction(action) - if (!sent) { - setFlow(connected ? `not_sent ${actionLabel(action)}` : "connect session first") - } -} - -function enqueueHeldKeyRepeats() { - const heldKeys = Array.from(activeKeys).sort((a, b) => { - return (heldKeyOrder.get(a) || 0) - (heldKeyOrder.get(b) || 0) - }) - for (const key of heldKeys) { - enqueueAction({ event: "keydown", key }) - } -} - -function setKeyHeld(key, source, held) { - const normalized = normalizeKey(key) - if (!allowedKeys.has(normalized)) { - return - } - - let sources = keySources.get(normalized) - if (!sources) { - sources = new Set() - keySources.set(normalized, sources) - } - - const wasActive = sources.size > 0 - if (held) { - sources.add(source) - } else { - sources.delete(source) - } - const isActive = sources.size > 0 - updateControlHighlights() - - if (held && !wasActive && isActive) { - heldKeySequence += 1 - heldKeyOrder.set(normalized, heldKeySequence) - enqueueAction({ event: "keydown", key: normalized }) - } - if (!held && wasActive && !isActive) { - heldKeyOrder.delete(normalized) - enqueueAction({ event: "keyup", key: normalized }) - } -} - -function releaseAllKeys() { - for (const key of Array.from(keySources.keys())) { - const sources = keySources.get(key) - if (sources && sources.size > 0) { - sources.clear() - heldKeyOrder.delete(key) - updateControlHighlights() - enqueueAction({ event: "keyup", key }) - } - } -} - -function handleControlMessage(rawMessage) { - let payload - try { - payload = JSON.parse(rawMessage) - } catch (error) { - logEvent(`invalid control payload: ${rawMessage}`, { level: "error" }) - return - } - - if (payload.type === "chunk_done") { - inferenceInFlight = false - if (Object.prototype.hasOwnProperty.call(payload, "active_event_id")) { - activeEventId = payload.active_event_id || null - } - renderEventControls() - updateMetricsFromChunk(payload) - const genMs = firstFinite(payload.gen_ms) - const lagMs = firstFinite(payload.lag_ms) - const queueDepth = firstFinite(payload.queue_depth) - const parts = [ - `chunk_done index=${payload.chunk_index}`, - `frames=${payload.num_frames}`, - `enqueued=${payload.enqueued_frames}`, - ] - if (genMs !== null) { - parts.push(`gen=${Math.round(genMs)}ms`) - } - if (lagMs !== null) { - parts.push(`lag=${Math.round(lagMs)}ms`) - } - if (metrics.latencyMs !== null) { - parts.push(`latency=${Math.round(metrics.latencyMs)}ms`) - } - if (queueDepth !== null) { - parts.push(`queue=${queueDepth}`) - } - logEvent(parts.join(", ")) - setStatus(activeKeys.size > 0 ? "Generating" : "Waiting", activeKeys.size > 0 ? "generating" : "waiting") - setFlow(`chunk ${payload.chunk_index} complete`) - if (activeKeys.size > 0) { - enqueueHeldKeyRepeats() - } - return - } - - if (payload.type === "event_ack") { - activeEventId = payload.active_event_id || null - renderEventControls() - logEvent(`event ${payload.event_id} ${payload.state}`, { source: "server" }) - return - } - - if (payload.type === "server_log") { - logEvent(payload.message || "server log") - return - } - - if (payload.type === "busy") { - logEvent(`server busy: ${payload.message}`, { level: "error" }) - setStatus("Waiting", "waiting") - return - } - - if (payload.type === "error") { - inferenceInFlight = false - logEvent(`server error: ${payload.message}`, { level: "error" }) - setStatus("Error", "error") - setFlow("server error") - return - } - - logEvent(`server message: ${rawMessage}`) -} - -async function waitForIceGatheringComplete(pc) { - if (pc.iceGatheringState === "complete") { - return - } - await new Promise((resolve) => { - const onStateChange = () => { - if (pc.iceGatheringState === "complete") { - pc.removeEventListener("icegatheringstatechange", onStateChange) - resolve() - } - } - pc.addEventListener("icegatheringstatechange", onStateChange) - }) -} - -async function pollWebRtcStats() { - if (!peerConnection) { - return - } - try { - const stats = await peerConnection.getStats() - for (const report of stats.values()) { - if ( - report.type === "candidate-pair" && - report.state === "succeeded" && - Number.isFinite(report.currentRoundTripTime) - ) { - metrics.rttMs = report.currentRoundTripTime * 1000 - } - if ( - report.type === "inbound-rtp" && - (report.kind === "video" || report.mediaType === "video") && - Number.isFinite(report.framesPerSecond) - ) { - metrics.fps = report.framesPerSecond - } - } - renderMetrics() - } catch (error) { - logEvent(`stats unavailable: ${error.message}`, { source: "client" }) - } -} - -function startStatsPolling() { - if (statsTimer !== null) { - return - } - statsTimer = window.setInterval(() => { - void pollWebRtcStats() - }, 1000) -} - -function stopStatsPolling() { - if (statsTimer !== null) { - window.clearInterval(statsTimer) - statsTimer = null - } -} - -function sendHeartbeat() { - if (!controlChannel || controlChannel.readyState !== "open") { - return - } - try { - controlChannel.send(JSON.stringify({ type: "heartbeat", t: Date.now() })) - } catch (error) { - logEvent(`heartbeat failed: ${error.message}`, { source: "client" }) - } -} - -function startHeartbeat() { - if (heartbeatTimer !== null) { - return - } - sendHeartbeat() - heartbeatTimer = window.setInterval(sendHeartbeat, heartbeatIntervalMs) -} - -function stopHeartbeat() { - if (heartbeatTimer !== null) { - window.clearInterval(heartbeatTimer) - heartbeatTimer = null - } -} - -function disconnectSession({ notify = true } = {}) { - if (disconnecting) { - return - } - disconnecting = true - stopHeartbeat() - stopStatsPolling() - connected = false - actionStarted = false - updateReadyPreview() - connectButton.disabled = false - if (notify && controlChannel && controlChannel.readyState === "open") { - try { - controlChannel.send(JSON.stringify({ type: "disconnect" })) - } catch { - // The browser may already be tearing the page down. - } - } - if (controlChannel && controlChannel.readyState !== "closed") { - controlChannel.close() - } - if (peerConnection) { - peerConnection.close() - } -} - -async function connectSession() { - if (mockMode) { - await startMockSession() - return - } - - connectButton.disabled = true - setStatus("Connecting", "connecting") - setFlow("preparing input") - logEvent("connecting to server...", { source: "client" }) - disconnecting = false - actionStarted = false - updateReadyPreview() - - try { - await uploadSessionInputIfNeeded() - setFlow("creating peer connection") - - peerConnection = new RTCPeerConnection() - controlChannel = peerConnection.createDataChannel("controls") - peerConnection.addTransceiver("video", { direction: "recvonly" }) - - controlChannel.onopen = () => { - logEvent("control data channel open") - setFlow("ready for action") - startHeartbeat() - } - controlChannel.onclose = () => { - logEvent("control data channel closed") - setFlow("channel closed") - stopHeartbeat() - } - controlChannel.onmessage = (event) => { - handleControlMessage(event.data) - } - - peerConnection.ontrack = (event) => { - const [stream] = event.streams - if (stream) { - remoteVideo.srcObject = stream - updateMetricsFromVideo() - } - } - - peerConnection.onconnectionstatechange = () => { - const state = peerConnection.connectionState - logEvent(`connection_state=${state}`, { source: "client" }) - if (state === "connected") { - connected = true - setStatus("Waiting", "waiting") - setFlow("connected; waiting for input") - startStatsPolling() - return - } - if (state === "connecting") { - setStatus("Connecting", "connecting") - return - } - if (["failed", "closed", "disconnected"].includes(state)) { - connected = false - actionStarted = false - updateReadyPreview() - connectButton.disabled = false - stopHeartbeat() - stopStatsPolling() - setStatus(state === "failed" ? "Error" : "Idle", state === "failed" ? "error" : "idle") - } - } - - const offer = await peerConnection.createOffer() - await peerConnection.setLocalDescription(offer) - await waitForIceGatheringComplete(peerConnection) - - const response = await fetch("/api/webrtc/offer", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(peerConnection.localDescription), - }) - if (!response.ok) { - const text = await response.text() - throw new Error(`offer failed (${response.status}): ${text}`) - } - const answer = await response.json() - await peerConnection.setRemoteDescription(answer) - logEvent("offer/answer completed") - } catch (error) { - stopHeartbeat() - if (peerConnection) { - peerConnection.close() - } - connected = false - setStatus("Error", "error") - setFlow("failed") - logEvent(`connect failed: ${error.message}`, { source: "client", level: "error" }) - connectButton.disabled = false - } -} - -function handleKeyDown(event) { - if (isEditableControlTarget(event.target)) { - return - } - - const key = normalizeKey(event.key) - if (!allowedKeys.has(key)) { - return - } - event.preventDefault() - - if (event.repeat) { - return - } - setKeyHeld(key, `keyboard:${key}`, true) -} - -function handleKeyUp(event) { - if (isEditableControlTarget(event.target)) { - return - } - - const key = normalizeKey(event.key) - if (!allowedKeys.has(key)) { - return - } - event.preventDefault() - setKeyHeld(key, `keyboard:${key}`, false) -} - -function attachPointerControls() { - for (const button of controlButtons) { - const key = button.dataset.controlKey - button.addEventListener("pointerdown", (event) => { - if (event.button !== 0) { - return - } - event.preventDefault() - button.setPointerCapture(event.pointerId) - setKeyHeld(key, `pointer:${event.pointerId}`, true) - }) - button.addEventListener("pointerup", (event) => { - event.preventDefault() - setKeyHeld(key, `pointer:${event.pointerId}`, false) - }) - button.addEventListener("pointercancel", (event) => { - setKeyHeld(key, `pointer:${event.pointerId}`, false) - }) - button.addEventListener("lostpointercapture", (event) => { - setKeyHeld(key, `pointer:${event.pointerId}`, false) - }) - } -} - -function resizeCanvas(ctx) { - const rect = mockCanvas.getBoundingClientRect() - const dpr = Math.min(window.devicePixelRatio || 1, 2) - const width = Math.max(1, Math.floor(rect.width * dpr)) - const height = Math.max(1, Math.floor(rect.height * dpr)) - if (mockCanvas.width !== width || mockCanvas.height !== height) { - mockCanvas.width = width - mockCanvas.height = height - } - ctx.setTransform(dpr, 0, 0, dpr, 0, 0) - return { width: rect.width, height: rect.height } -} - -function drawMountain(ctx, points, fill) { - ctx.beginPath() - ctx.moveTo(points[0][0], points[0][1]) - for (const point of points.slice(1)) { - ctx.lineTo(point[0], point[1]) - } - ctx.closePath() - ctx.fillStyle = fill - ctx.fill() -} - -function drawMockScene(now) { - const ctx = mockCanvas.getContext("2d") - const { width, height } = resizeCanvas(ctx) - const t = now * 0.001 - const horizon = height * 0.44 - - const sky = ctx.createLinearGradient(0, 0, width, height) - sky.addColorStop(0, "#718697") - sky.addColorStop(0.48, "#c8d8da") - sky.addColorStop(1, "#f4bf77") - ctx.fillStyle = sky - ctx.fillRect(0, 0, width, height) - - drawMountain( - ctx, - [ - [0, horizon + 40], - [width * 0.18, height * 0.15], - [width * 0.35, horizon + 18], - [width * 0.52, height * 0.22], - [width * 0.72, horizon + 30], - [width, height * 0.30], - [width, height], - [0, height], - ], - "rgba(36, 54, 56, 0.90)" - ) - drawMountain( - ctx, - [ - [width * 0.18, horizon + 36], - [width * 0.34, height * 0.25], - [width * 0.50, horizon + 12], - [width * 0.67, height * 0.31], - [width, horizon + 26], - [width, height], - [width * 0.18, height], - ], - "rgba(72, 93, 88, 0.72)" - ) - - const water = ctx.createLinearGradient(width * 0.58, horizon, width, height) - water.addColorStop(0, "rgba(166, 196, 199, 0.78)") - water.addColorStop(1, "rgba(61, 85, 93, 0.92)") - ctx.fillStyle = water - ctx.beginPath() - ctx.moveTo(width * 0.54, horizon + 36) - ctx.lineTo(width, horizon + 8) - ctx.lineTo(width, height) - ctx.lineTo(width * 0.64, height) - ctx.closePath() - ctx.fill() - - const road = ctx.createLinearGradient(width * 0.35, horizon, width * 0.45, height) - road.addColorStop(0, "#424a4b") - road.addColorStop(1, "#17191a") - ctx.fillStyle = road - ctx.beginPath() - ctx.moveTo(width * 0.36, horizon + 30) - ctx.lineTo(width * 0.60, horizon + 26) - ctx.lineTo(width * 0.70, height) - ctx.lineTo(width * 0.18, height) - ctx.closePath() - ctx.fill() - - ctx.strokeStyle = "rgba(255, 220, 105, 0.72)" - ctx.lineWidth = 3 - ctx.beginPath() - ctx.moveTo(width * 0.50, horizon + 32) - ctx.lineTo(width * 0.56, height) - ctx.stroke() - - ctx.strokeStyle = "rgba(240, 246, 242, 0.58)" - ctx.lineWidth = 2 - for (let i = 0; i < 12; i += 1) { - const y = horizon + 50 + ((i * 52 + t * 80) % (height - horizon + 90)) - const scale = (y - horizon) / (height - horizon) - ctx.beginPath() - ctx.moveTo(width * (0.42 + scale * 0.02), y) - ctx.lineTo(width * (0.46 + scale * 0.04), y + 18 + scale * 20) - ctx.stroke() - } - - ctx.fillStyle = "rgba(36, 38, 35, 0.88)" - for (let i = 0; i < 7; i += 1) { - const x = width * (0.02 + i * 0.055) - const y = horizon + 18 - i * 3 - const buildingWidth = width * 0.046 - const buildingHeight = height * (0.16 + (i % 3) * 0.035) - ctx.fillRect(x, y - buildingHeight, buildingWidth, buildingHeight) - ctx.fillStyle = "rgba(255, 214, 142, 0.58)" - ctx.fillRect(x + 8, y - buildingHeight + 18, 7, 11) - ctx.fillRect(x + buildingWidth - 15, y - buildingHeight + 42, 7, 11) - ctx.fillStyle = "rgba(36, 38, 35, 0.88)" - } - - ctx.strokeStyle = "rgba(19, 42, 30, 0.94)" - ctx.lineWidth = 5 - for (let i = 0; i < 6; i += 1) { - const x = width * (0.19 + i * 0.035) - const treeBase = horizon + 55 + i * 12 - ctx.beginPath() - ctx.moveTo(x, treeBase) - ctx.lineTo(x, treeBase - height * 0.18) - ctx.stroke() - ctx.fillStyle = "rgba(33, 77, 46, 0.85)" - ctx.beginPath() - ctx.ellipse(x, treeBase - height * 0.12, 8, 44, 0, 0, Math.PI * 2) - ctx.fill() - } - - ctx.fillStyle = `rgba(255, 255, 255, ${0.10 + Math.sin(t) * 0.025})` - ctx.fillRect(0, 0, width, height) - - if (!document.body.classList.contains("has-video")) { - recordFrame(now) - } - window.requestAnimationFrame(drawMockScene) -} - -function startVideoFrameMonitor() { - if (typeof remoteVideo.requestVideoFrameCallback !== "function") { - return - } - const onFrame = (now) => { - if (document.body.classList.contains("has-video")) { - recordFrame(now) - updateMetricsFromVideo() - } - remoteVideo.requestVideoFrameCallback(onFrame) - } - remoteVideo.requestVideoFrameCallback(onFrame) -} - -function mockChunkPayload() { - const numFrames = 12 - const targetFps = 16 - const genMs = 360 + Math.random() * 120 - const lagMs = 54 + Math.random() * 36 - return { - type: "chunk_done", - chunk_index: mockChunkIndex++, - num_frames: numFrames, - enqueued_frames: numFrames, - fps: targetFps, - resolution: { width: 1280, height: 720 }, - model: "lingbot-world-v2-14b-causal-fast-taehv-window15-sink3", - active_event_id: activeEventId, - latency_ms: 118 + Math.random() * 48, - consumed_actions: 1, - gen_ms: genMs, - enqueue_ms: 8 + Math.random() * 4, - play_ms: (numFrames * 1000) / targetFps, - lag_ms: lagMs, - queue_depth: Math.floor(3 + Math.random() * 7), - } -} - -function ensureMockChunks() { - if (mockChunkTimer !== null) { - return - } - mockChunkTimer = window.setInterval(() => { - if (!connected || !mockGenerationStarted) { - return - } - handleControlMessage(JSON.stringify(mockChunkPayload())) - }, 760) -} - -async function startMockSession() { - connectButton.disabled = true - setStatus("Connecting", "connecting") - setFlow("mock warmup") - logEvent("connecting to mock server...", { source: "client" }) - actionStarted = false - await uploadSessionInputIfNeeded() - await new Promise((resolve) => { - window.setTimeout(resolve, 260) - }) - connected = true - metrics.targetFps = 16 - metrics.resolution = "1280x720" - metrics.model = "lingbot-world-v2-14b-causal-fast-taehv-window15-sink3" - renderMetrics() - setStatus("Waiting", "waiting") - setFlow("mock ready; waiting for input") - logEvent("Connected") - logEvent("Warmup complete") - ensureMockChunks() -} - -let panelZIndex = 10 - -function bringPanelToFront(panel) { - panelZIndex += 1 - panel.style.zIndex = String(panelZIndex) -} - -function makePanelMovable(panel, handle) { - handle.classList.add("panelDragHandle") - - const collapseButton = document.createElement("button") - collapseButton.type = "button" - collapseButton.className = "panelCollapseButton" - collapseButton.textContent = "\u2013" - collapseButton.setAttribute("aria-expanded", "true") - collapseButton.setAttribute("aria-label", "Collapse panel") - collapseButton.addEventListener("pointerdown", (event) => { - event.stopPropagation() - }) - collapseButton.addEventListener("click", (event) => { - event.stopPropagation() - const collapsed = panel.classList.toggle("is-collapsed") - collapseButton.textContent = collapsed ? "+" : "\u2013" - collapseButton.setAttribute("aria-expanded", collapsed ? "false" : "true") - collapseButton.setAttribute("aria-label", collapsed ? "Expand panel" : "Collapse panel") - }) - handle.appendChild(collapseButton) - - let dragging = false - let pointerId = null - let startX = 0 - let startY = 0 - let startLeft = 0 - let startTop = 0 - - const stageOf = () => panel.offsetParent || document.body - - handle.addEventListener("pointerdown", (event) => { - if (event.button !== 0) { - return - } - bringPanelToFront(panel) - const stageRect = stageOf().getBoundingClientRect() - const panelRect = panel.getBoundingClientRect() - startLeft = panelRect.left - stageRect.left - startTop = panelRect.top - stageRect.top - panel.classList.add("is-floating") - panel.style.left = `${startLeft}px` - panel.style.top = `${startTop}px` - startX = event.clientX - startY = event.clientY - dragging = true - pointerId = event.pointerId - handle.setPointerCapture(pointerId) - event.preventDefault() - }) - - handle.addEventListener("pointermove", (event) => { - if (!dragging || event.pointerId !== pointerId) { - return - } - const stageRect = stageOf().getBoundingClientRect() - const maxLeft = Math.max(0, stageRect.width - panel.offsetWidth) - const maxTop = Math.max(0, stageRect.height - panel.offsetHeight) - const nextLeft = Math.min(Math.max(0, startLeft + (event.clientX - startX)), maxLeft) - const nextTop = Math.min(Math.max(0, startTop + (event.clientY - startY)), maxTop) - panel.style.left = `${nextLeft}px` - panel.style.top = `${nextTop}px` - event.preventDefault() - }) - - const endDrag = () => { - if (!dragging) { - return - } - if (pointerId !== null && handle.hasPointerCapture(pointerId)) { - handle.releasePointerCapture(pointerId) - } - dragging = false - pointerId = null - } - handle.addEventListener("pointerup", endDrag) - handle.addEventListener("pointercancel", endDrag) - handle.addEventListener("lostpointercapture", endDrag) - - panel.addEventListener("pointerdown", () => { - bringPanelToFront(panel) - }) -} - -function setupPanelChrome() { - const panels = [ - { selector: ".statusCard", handle: ".panelLabel" }, - { selector: "#sceneCard", handle: ".panelLabel" }, - { selector: ".controlCard", handle: "h2" }, - { selector: ".logCard", handle: "h2" }, - ] - for (const entry of panels) { - const panel = document.querySelector(entry.selector) - if (!panel) { - continue - } - const handle = panel.querySelector(entry.handle) - if (!handle) { - continue - } - makePanelMovable(panel, handle) - } -} - -function initialize() { - document.body.dataset.status = "idle" - setFirstFrameInputMode("url") - if (mockMode) { - document.body.classList.add("mock-mode") - connectButton.textContent = "Start Mock Session" - logEvent("mock mode ready", { source: "client" }) - } else { - logEvent("viewer ready", { source: "client" }) - } - setFlow("waiting") - renderMetrics() - attachPointerControls() - setupPanelChrome() - void loadInitialScene() - window.requestAnimationFrame(drawMockScene) - startVideoFrameMonitor() -} - -connectButton.addEventListener("click", () => { - void connectSession() -}) -clearEventButton.addEventListener("click", () => { - sendTextEvent(activeEventId || "clear", "clear") -}) -uploadModeButton.addEventListener("click", () => { - if (initialSceneLocked) { - return - } - setFirstFrameInputMode("upload") - if (!selectedFirstFrameFile) { - firstFrameName.textContent = defaultFirstFrameName() - } - releaseAllKeys() -}) -urlModeButton.addEventListener("click", () => { - if (initialSceneLocked) { - return - } - setFirstFrameInputMode("url") - releaseAllKeys() -}) -firstFrameInput.addEventListener("change", () => { - if (initialSceneLocked) { - return - } - setFirstFrameInputMode("upload") - const [file] = firstFrameInput.files - selectedFirstFrameFile = file || null - firstFrameSelectionCommitted = false - if (selectedFirstFrameUrl) { - URL.revokeObjectURL(selectedFirstFrameUrl) - selectedFirstFrameUrl = null - } - if (selectedFirstFrameFile) { - selectedFirstFrameUrl = URL.createObjectURL(selectedFirstFrameFile) - firstFrameName.textContent = selectedFirstFrameFile.name - clearFirstFrameUrlInput() - setFirstFrameUrlStatus("Image not updated", "pending") - } else { - firstFrameName.textContent = defaultFirstFrameName() - setFirstFrameUrlStatus() - } - updateReadyPreview() -}) -firstFrameUrlInput.addEventListener("input", () => { - if (initialSceneLocked) { - return - } - setFirstFrameInputMode("url") - if (selectedFirstFrameFile) { - clearSelectedFirstFrameFile() - } - firstFrameUrlEdited = true - if (!selectedFirstFrameFile) { - firstFrameName.textContent = firstFrameUrlInput.value.trim() - ? "Upload Image" - : defaultFirstFrameName() - } - setFirstFrameUrlStatus( - firstFrameUrlInput.value.trim() ? "URL not updated" : "", - "pending" - ) -}) -firstFrameUrlUpdateButton.addEventListener("click", () => { - void updateFirstFrameInput() -}) -promptInput.addEventListener("input", () => { - if (initialSceneLocked) { - return - } - promptEdited = true -}) -addTextEventButton.addEventListener("click", () => { - if (initialSceneLocked) { - return - } - textEventDrafts.push(makeTextEventDraft({ label: "", prompt: "" })) - markTextEventsEdited() - renderTextEventEditor() - releaseAllKeys() -}) -firstFrameUrlInput.addEventListener("focus", releaseAllKeys) -promptInput.addEventListener("focus", releaseAllKeys) -addTextEventButton.addEventListener("focus", releaseAllKeys) -remoteVideo.addEventListener("loadedmetadata", updateMetricsFromVideo) -remoteVideo.addEventListener("playing", () => { - setVideoVisible(true) - updateMetricsFromVideo() -}) -remoteVideo.addEventListener("emptied", () => { - setVideoVisible(false) -}) -window.addEventListener("keydown", handleKeyDown) -window.addEventListener("keyup", handleKeyUp) -window.addEventListener("blur", releaseAllKeys) -window.addEventListener("pagehide", () => { - disconnectSession() -}) -window.addEventListener("beforeunload", () => { - disconnectSession() -}) - -initialize() diff --git a/integrations/lingbot/pyproject.toml b/integrations/lingbot/pyproject.toml index b14cd2c91..211069042 100644 --- a/integrations/lingbot/pyproject.toml +++ b/integrations/lingbot/pyproject.toml @@ -60,10 +60,8 @@ exclude = ["tests"] [tool.setuptools.package-data] "lingbot.webrtc.web" = [ - "*.html", - "*.css", - "*.js", - "assets/*.svg", + "adapter.css", + "adapter.js", ] [tool.uv] diff --git a/integrations/lingbot/tests/test_demo_api.py b/integrations/lingbot/tests/test_demo_api.py index c9446b414..b2c560545 100644 --- a/integrations/lingbot/tests/test_demo_api.py +++ b/integrations/lingbot/tests/test_demo_api.py @@ -443,6 +443,7 @@ def fake_create_packaged_app(**kwargs: Any) -> web.Application: app_calls.append(kwargs) app = web.Application() app[SESSION_MANAGER_KEY] = kwargs["session_manager"] + kwargs["configure_app"](app) return app monkeypatch.setattr( @@ -476,6 +477,12 @@ def fake_create_packaged_app(**kwargs: Any) -> web.Application: ) assert app_calls[0]["preload_name"] == "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()} + assert "/api/session/initial_scene" in route_paths + assert "/api/session/first_frame" in route_paths + assert "/api/session/input" in route_paths def test_lingbot_webrtc_demo_serves_through_shared_runner( @@ -488,6 +495,7 @@ def test_lingbot_webrtc_demo_serves_through_shared_runner( def fake_create_packaged_app(**kwargs: Any) -> web.Application: app = web.Application() app[SESSION_MANAGER_KEY] = kwargs["session_manager"] + kwargs["configure_app"](app) return app def fake_server_runner(**kwargs: Any) -> None: diff --git a/integrations/lingbot/tests/test_server_routes.py b/integrations/lingbot/tests/test_server_routes.py index fbb024bb9..6c4acd49f 100644 --- a/integrations/lingbot/tests/test_server_routes.py +++ b/integrations/lingbot/tests/test_server_routes.py @@ -133,7 +133,9 @@ def test_create_app_keeps_package_web_resource_materialized() -> None: assert len(static_resources) == 1 web_dir = static_resources[0].get_info()["directory"] assert web_dir.is_dir() - assert "Lingbot WebRTC Viewer" in (web_dir / "request_session.html").read_text() + assert ( + "FlashDreams WebRTC Drive" in (web_dir / "request_session.html").read_text() + ) finally: app[PACKAGE_RESOURCE_STACK_KEY].close() @@ -181,7 +183,26 @@ async def test_request_session_serves_html() -> None: response = await client.get("/request_session") body = await response.text() assert response.status == 200 - assert "Lingbot WebRTC Viewer" in body + assert "FlashDreams WebRTC Drive" in body + finally: + await client.close() + + +@pytest.mark.asyncio +async def test_lingbot_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: "Lingbot"' in body + assert "/api/session/initial_scene" in body + assert '{ key: "w"' not in body + assert '{ key: "q"' in body + assert "enablePostprocess" not in body + assert "RTCPeerConnection" not in body finally: await client.close() diff --git a/integrations/omnidreams/omnidreams/demo/webrtc.py b/integrations/omnidreams/omnidreams/demo/webrtc.py index 911c38c72..32c9e9a25 100644 --- a/integrations/omnidreams/omnidreams/demo/webrtc.py +++ b/integrations/omnidreams/omnidreams/demo/webrtc.py @@ -149,6 +149,7 @@ def create_omnidreams_webrtc_app( 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, diff --git a/integrations/omnidreams/omnidreams/webrtc/server.py b/integrations/omnidreams/omnidreams/webrtc/server.py index bc50bf49c..f577dc3df 100644 --- a/integrations/omnidreams/omnidreams/webrtc/server.py +++ b/integrations/omnidreams/omnidreams/webrtc/server.py @@ -56,6 +56,7 @@ ) WEB_DIR_RESOURCE = files("flashdreams.serving.webrtc").joinpath("web") +MODEL_WEB_DIR_RESOURCE = files("omnidreams.webrtc").joinpath("web") class _OmnidreamsSessionManager(WebRTCSessionManager, Protocol): @@ -231,6 +232,7 @@ def create_app( 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, diff --git a/integrations/omnidreams/omnidreams/webrtc/web/adapter.js b/integrations/omnidreams/omnidreams/webrtc/web/adapter.js new file mode 100644 index 000000000..7ae561fe4 --- /dev/null +++ b/integrations/omnidreams/omnidreams/webrtc/web/adapter.js @@ -0,0 +1,7 @@ +// 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/omnidreams/webrtc/web/assets/horizontal-dark.svg b/integrations/omnidreams/omnidreams/webrtc/web/assets/horizontal-dark.svg deleted file mode 100644 index 89b68f9d1..000000000 --- a/integrations/omnidreams/omnidreams/webrtc/web/assets/horizontal-dark.svg +++ /dev/null @@ -1,192 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/integrations/omnidreams/omnidreams/webrtc/web/assets/horizontal-light.svg b/integrations/omnidreams/omnidreams/webrtc/web/assets/horizontal-light.svg deleted file mode 100644 index a491910db..000000000 --- a/integrations/omnidreams/omnidreams/webrtc/web/assets/horizontal-light.svg +++ /dev/null @@ -1,182 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/integrations/omnidreams/pyproject.toml b/integrations/omnidreams/pyproject.toml index eafd80a1c..4e05b745c 100644 --- a/integrations/omnidreams/pyproject.toml +++ b/integrations/omnidreams/pyproject.toml @@ -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" = ["*.html", "*.css", "*.js", "assets/*.svg"] +"omnidreams.webrtc.web" = ["adapter.js"] "omnidreams.interactive_drive" = [ "configs/*.yaml", "configs/wheels/*.yaml", diff --git a/integrations/omnidreams/tests/test_demo_api.py b/integrations/omnidreams/tests/test_demo_api.py index d9411475a..80928c5cf 100644 --- a/integrations/omnidreams/tests/test_demo_api.py +++ b/integrations/omnidreams/tests/test_demo_api.py @@ -427,6 +427,7 @@ def fake_create_packaged_webrtc_app(**kwargs: Any) -> web.Application: "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 diff --git a/integrations/omnidreams/tests/test_webrtc_runtime.py b/integrations/omnidreams/tests/test_webrtc_runtime.py index 1a11f18f8..0776203f1 100644 --- a/integrations/omnidreams/tests/test_webrtc_runtime.py +++ b/integrations/omnidreams/tests/test_webrtc_runtime.py @@ -953,18 +953,23 @@ async def test_postprocess_options_exposes_only_launch_preset() -> None: def test_webrtc_ui_posts_selected_postprocess_preset() -> None: - web_dir = files("flashdreams.serving.webrtc").joinpath("web") - html = web_dir.joinpath("request_session.html").read_text(encoding="utf-8") - javascript = web_dir.joinpath("request_session.js").read_text(encoding="utf-8") + 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 'id="postprocessField"' in html - assert "hidden" in html - assert 'id="postprocessSelect"' in html assert 'fetch("/api/postprocess/options")' in javascript assert 'fetch("/api/session/input"' in javascript - assert "postprocessControlAvailable" in javascript - assert "postprocessField.hidden = !postprocessControlAvailable" 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 diff --git a/integrations/omnidreams/tests/test_webrtc_server_routes.py b/integrations/omnidreams/tests/test_webrtc_server_routes.py index bddb824bd..d1104057e 100644 --- a/integrations/omnidreams/tests/test_webrtc_server_routes.py +++ b/integrations/omnidreams/tests/test_webrtc_server_routes.py @@ -154,10 +154,8 @@ async def test_request_session_uses_lingbot_aligned_viewer_shell() -> None: assert "Connect Session" in body assert 'id="logState"' in body assert "World Model" in body - for key in ("w", "a", "s", "d"): - assert f'data-control-key="{key}"' in body - for key in ("q", "e", "i", "j", "k", "l"): - assert f'data-control-key="{key}"' not in body + assert 'id="controlRows"' in body + assert 'id="modelStatusSlot"' in body finally: await client.close() @@ -204,14 +202,15 @@ async def test_static_js_requests_recvonly_video_transceiver() -> None: @pytest.mark.asyncio -async def test_static_js_keeps_omnidreams_controls_and_lingbot_status_helpers() -> None: +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 allowedKeys = new Set(["w", "a", "s", "d"])' in body + 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 @@ -219,6 +218,23 @@ async def test_static_js_keeps_omnidreams_controls_and_lingbot_status_helpers() 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() From 6d14cf9ab223137feb59617adffbb5990574fcef Mon Sep 17 00:00:00 2001 From: aidanfnv Date: Fri, 7 Aug 2026 17:50:54 -0700 Subject: [PATCH 13/19] Continue Lingbot porting after webRTC refactor (#428) * Use inferencesession * Remove legacy lingbot, finish port * Address greptile comments * Address greptile comments * Address greptile comments * Address greptile comments --- docs/inference_runtime_api_design.md | 38 +- .../flashdreams/serving/webrtc/manager.py | 396 ++++++++++++++++- .../flashdreams/serving/webrtc/runtime.py | 31 ++ flashdreams/tests/test_webrtc_manager.py | 289 ++++++++++++- flashdreams/tests/test_webrtc_serving.py | 147 +++++++ .../lingbot/lingbot/webrtc/session.py | 279 +++++++++++- .../lingbot/tests/test_webrtc_runtime.py | 282 ++++++------ .../tests/test_webrtc_session_branch.py | 407 ++++++++++++++++++ 8 files changed, 1673 insertions(+), 196 deletions(-) create mode 100644 integrations/lingbot/tests/test_webrtc_session_branch.py diff --git a/docs/inference_runtime_api_design.md b/docs/inference_runtime_api_design.md index 6bd1018d9..260b24b24 100644 --- a/docs/inference_runtime_api_design.md +++ b/docs/inference_runtime_api_design.md @@ -44,10 +44,8 @@ model. ## Current Implementation Plan Implementation should happen on an experimental integration branch. PRs for this -work should target that branch until the API shape and OmniDreams migration are -working well enough to merge to `main` together. LingBot migration is deferred -to a separate follow-up after the OmniDreams path has clarified the shared demo -API shape. +work should target that branch until the API shape and migrated demo paths are +working well enough to merge to `main` together. The experimental branch can temporarily break or simplify command-line options while the demos are being moved to the new API. The required outcome for this @@ -62,7 +60,7 @@ Initial scope: - support selectable output modes such as MP4, JPEG/MJPEG stream, WebRTC, and headless/null where appropriate; - use or update benchmark tooling to verify the migrated OmniDreams demo; -- defer broader model migrations, hosted execution, full autotune, and polished +- defer additional model migrations, hosted execution, full autotune, and polished metrics until the first branch proves the API shape. ## Task Tracker @@ -75,13 +73,37 @@ Initial scope: | T3 | Complete | `CanonicalInputs`, `InferenceInput`, schemas, and mapping boundary. | Yes, after T1 direction is agreed. | T1. | Models can declare required global/per-step inputs, and mappings can convert canonical inputs into inference inputs. | | T4 | Complete | `ModelRunner`, `InferenceRuntime`, and `InferenceSession` skeleton. | Partly. | T1. | A minimal standard loop can initialize a runtime, run at least one sequential session, and close cleanly. | | T5 | Planned | Output mode selection. | Yes, after the result/output shape is agreed. | T1, T4. | A run can choose output behavior such as MP4, JPEG/MJPEG stream, WebRTC, benchmark artifact, or headless/null without changing model code. | -| T6 | Deferred | LingBot migration. | Yes, but out of scope for this branch. | T2, T3, T4. | LingBot runs through the new API path with its event inputs mapped into model inputs. | +| T6 | Partially complete | LingBot migration and live model-input cleanup. | Yes. | T2, T3, T4. | Scene/prompt/first-frame model inputs flow through `InferenceInput.global_conditioning` or a typed model-input object, static runtime settings remain in runtime config, LingBot uses the shared WebRTC hook shape, and any retained per-model WebRTC/demo wrappers are deliberate compatibility shims. | | T7 | Partially complete | OmniDreams migration. | Yes, once T2-T4 have a usable skeleton. | T2, T3, T4. | OmniDreams replay and WebRTC run through the shared demo API path; remaining work is output/stat integration, legacy demo retirement, and cleanup. | | T8 | Partially complete | Benchmark/smoke verification for OmniDreams. | Preparation can run early; final gate is late. | T5, T7. | Existing or updated benchmark tooling can run the migrated OmniDreams demo and produce enough evidence that it still works. | | T9 | Planned | Metrics and profiling normalization for the branch. | Yes, but final integration is late. | T4, T5, T8. | Basic canonical metrics are emitted for migrated demos; deeper metrics can remain follow-up work. | | T10 | Planned | CLI compatibility, legacy retirement, and migration cleanup. | Yes, after demo migrations start. | T5, T7, T8. | Required demo commands are restored or replaced, old interactive-drive and old OmniDreams demo/server paths are removed or reduced to compatibility shims, code used only by retired demos is removed, and user-facing docs/notes match the branch behavior. | | T11 | Planned | Stabilize and merge experimental branch to `main`. | No, final integration step. | T5, T7-T10. | OmniDreams passes agreed smoke/benchmark checks, review feedback is addressed, and the branch can merge as one API transition. | +Current LingBot migration status: + +- LingBot has partial runtime/session plumbing, but live WebRTC still retains + the segment-based `generate_chunk(segments, frame_times)` entry point while + the model-input boundary is cleaned up. +- Scene, prompt, and first-frame style model inputs should move through + `InferenceInput.global_conditioning` or a typed object carried there instead + of being hidden inside runtime config. +- Runtime config should keep static execution settings: pipeline config, + device, resolution/FPS, warmup, encoder options, movement speeds, and + cache/layout options. +- Browser-only session options, such as OmniDreams postprocess preset + selection, should remain pending session input unless they become true model + conditioning. +- LingBot still delegates parts of the WebRTC path through + `lingbot.webrtc.server.create_app()` and `LingbotWebRTCSessionManager`. + Follow-up work should move it to the same hook shape as OmniDreams: + `WebRTCManagerOptions`, `WebRTCAppExtension`, and shared route/resource + helpers. +- Full realtime `UserInputs` / `InputMapping` integration can be deferred + until after the live model-input handling is explicit. +- Later cleanup should remove or reduce old per-model WebRTC server wrappers + and obsolete demo code once both demos are fully on the shared path. + Current OmniDreams migration status: - The shared `flashdreams.runtime.demo` API and OmniDreams demo adapter exist. @@ -106,8 +128,8 @@ Suggested parallel split: stay coherent; - one person owns T5/T8/T9, because outputs, benchmarks, and metrics are tightly related; -- LingBot should be tracked as a separate follow-up once OmniDreams has settled - the shared demo API shape; +- one person owns T6's LingBot follow-up: live model inputs, shared WebRTC hook + migration, and cleanup of retained per-model wrappers; - one person should track branch health, CLI compatibility, and merge readiness. ## Architecture diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index 97c0c578e..e35e70263 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -11,7 +11,7 @@ import json from collections import deque from collections.abc import Set as AbstractSet -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from enum import IntEnum from typing import Any, Generic, TypeVar @@ -24,7 +24,17 @@ from loguru import logger from flashdreams.infra.video_output import VideoStepResult -from flashdreams.serving.realtime.input import KeyboardResampler +from flashdreams.runtime.inputs import ( + InferenceInput, + TimeWindow, + UserInputEvent, + UserInputs, +) +from flashdreams.serving.realtime.input import ( + DEFAULT_SUPPORTED_KEYS, + KeyboardResampler, + normalize_key, +) from flashdreams.serving.webrtc.encoders import ( DefaultRTCEncoder, VideoEncoder, @@ -63,6 +73,10 @@ # How often the liveness watchdog wakes to re-check the elapsed-since-last-message. _CLIENT_LIVENESS_CHECK_INTERVAL_S = 1.0 _DEFAULT_PERF_LOG_INTERVAL_CHUNKS = 5 +_MAX_SESSION_USER_EVENTS = 1024 +"""Maximum unconsumed raw events kept for an ``InferenceSession`` step.""" +_RELEASE_USER_EVENT_TYPES = frozenset({"key_up"}) +_KEY_USER_EVENT_TYPES = frozenset({"key_down", "key_up"}) _RuntimeT = TypeVar("_RuntimeT", bound=WebRTCSessionRuntime) _RuntimeConfigT = TypeVar("_RuntimeConfigT", bound=WebRTCRuntimeConfig) @@ -91,6 +105,8 @@ class WebRTCControlSignal(IntEnum): ACTION_STEP = 2 CLOSE = 3 EVENT = 4 + SESSION_STEP = 5 + """One step driven by mapped ``InferenceInput`` rather than pose segments.""" EXIT = 99 @@ -107,6 +123,14 @@ class ManagedWebRTCSession: generation_task: asyncio.Task[Any] | None = None first_action_received: asyncio.Event = field(default_factory=asyncio.Event) pending_action_arrivals: deque[float] = field(default_factory=deque) + inference_session: Any | None = None + """Active ``InferenceSession``; ``None`` means call ``runtime.generate_chunk``.""" + session_steps_completed: int = 0 + session_input_state_advanced: bool = False + user_events: deque[UserInputEvent] = field(default_factory=deque) + """Raw user events awaiting canonicalization, oldest first.""" + coalesced_release_events: dict[str, UserInputEvent] = field(default_factory=dict) + """Overflow key releases, coalesced by normalized key.""" last_client_message_at: float = 0.0 liveness_task: asyncio.Task[Any] | None = None closed: bool = False @@ -349,6 +373,287 @@ 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, *, @@ -373,6 +678,45 @@ async def _handle_event_message( ) return False + if managed_session.inference_session is not None: + # On the session branch a text event is just another user event: + # the mapping turns it into a session-global conditioning update + # applied by the next step, so there is no separate runtime call. + clears = state in clear_states + try: + event_payload = self._validate_user_event_payload( + managed_session=managed_session, + event_type="text_event", + payload={ + "event_id": None if clears else event_id, + "state": state, + }, + ) + self._record_user_event( + managed_session=managed_session, + timestamp_s=asyncio.get_running_loop().time(), + event_type="text_event", + payload=event_payload, + ) + except Exception as exc: + if channel is not None: + self._send_json(channel, make_error_payload(str(exc))) + return False + if channel is not None: + active_event_id = event_payload.get("event_id") + ack_event_id = ( + None if active_event_id is None else str(active_event_id) + ) + self._send_json( + channel, + make_event_ack_payload( + event_id=ack_event_id, + state=str(event_payload.get("state", state)), + result={"active_event_id": ack_event_id}, + ), + ) + return True + trigger_event = getattr(managed_session.runtime, "trigger_event", None) if not callable(trigger_event): if channel is not None: @@ -485,6 +829,11 @@ async def _create_answer_with_runtime_ready_locked( resampler=resampler, last_client_message_at=loop.time(), ) + session_runtime: Any = self._runtime + if self._drives_inference_session(session_runtime): + managed_session.inference_session = ( + await session_runtime.start_inference_session() + ) self._active_session = managed_session if enable_liveness_watchdog: managed_session.liveness_task = asyncio.create_task( @@ -707,6 +1056,18 @@ async def _handle_datachannel_message( # resampler's ``next_chunk_start_v`` so virtual-time comparisons in # ``KeyboardResampler.sample_chunk`` are well-defined. arrival_t = asyncio.get_running_loop().time() + if managed_session.inference_session is not None: + try: + self._record_user_event( + managed_session=managed_session, + timestamp_s=arrival_t, + event_type="key_down" if event == "keydown" else "key_up", + payload={"key": key}, + ) + except Exception as exc: + self._send_json(channel, make_error_payload(str(exc))) + if event != "keyup": + return managed_session.resampler.on_edge(arrival_t=arrival_t, event=event, key=key) managed_session.pending_action_arrivals.append(arrival_t) # Releases the generation worker, which blocks on this until the @@ -771,15 +1132,20 @@ async def _generation_worker( # Catch the virtual clock up to wall if it has fallen more # than one chunk behind so end-to-end latency stays bounded. - # Held-key continuity is preserved because ``sample_chunk`` - # folds every event below the new window start into the - # carried state. + # The segment branch folds skipped edges through the resampler; + # the session branch first advances its input canonicalizer + # across the skipped raw-event window. now = loop.time() - lag = now - (resampler.next_chunk_start_v + chunk_duration) - if lag > chunk_duration: - resampler.next_chunk_start_v = now - chunk_duration + self._catch_up_input_clock( + managed_session=managed_session, + now=now, + chunk_duration=chunk_duration, + ) 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 consumed_action_arrivals: list[float] = [] @@ -791,9 +1157,17 @@ async def _generation_worker( managed_session.pending_action_arrivals.popleft() ) try: - result = await runtime.generate_chunk( - segments=segments, frame_times=frame_times - ) + 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 + ) except Exception as exc: logger.exception("Chunk generation failed.") channel = managed_session.control_channel diff --git a/flashdreams/flashdreams/serving/webrtc/runtime.py b/flashdreams/flashdreams/serving/webrtc/runtime.py index 049574a13..d402657e1 100644 --- a/flashdreams/flashdreams/serving/webrtc/runtime.py +++ b/flashdreams/flashdreams/serving/webrtc/runtime.py @@ -9,6 +9,10 @@ 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 flashdreams.serving.realtime.input import PoseSegment @@ -64,6 +68,33 @@ def trigger_event( ) -> dict[str, Any] | Awaitable[dict[str, Any]]: ... +class WebRTCInferenceSessionRuntime(Protocol): + """Optional runtime capability for driving an ``InferenceSession``. + + 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. + """ + + async def start_inference_session(self) -> InferenceSession: ... + + @property + def input_mapping(self) -> InputMapping: ... + + @property + def input_canonicalizer(self) -> InputCanonicalizer: ... + + @property + def input_source_schema(self) -> UserInputSchema: ... + + class WebRTCServerLifecycle(Protocol): """Distributed worker lifecycle used by the shared WebRTC serve loop.""" diff --git a/flashdreams/tests/test_webrtc_manager.py b/flashdreams/tests/test_webrtc_manager.py index b7dca0582..e4537edb0 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -107,6 +107,14 @@ def sample_chunk( return [(0.0, 0.0, frozenset({"w"}))], [0.0] +class _RecordingResampler(_FakeResampler): + def __init__(self) -> None: + self.edges: list[tuple[float, str, str]] = [] + + def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: + self.edges.append((arrival_t, event, key)) + + class _CountingVideoTrack(_FakeVideoTrack): async def enqueue_chunk(self, chunk: Any) -> int: return int(chunk.shape[0]) @@ -117,6 +125,10 @@ def _model_name(self) -> str: return "fake-model" +class _WOnlyTestManager(_BaseTestManager): + _resampler_supported_keys = frozenset({"w"}) + + def _make_manager( manager_cls: type[BaseWebRTCSessionManager], runtime: Any ) -> BaseWebRTCSessionManager: @@ -186,6 +198,276 @@ def _managed_session( return managed, video_track, peer, channel +def test_record_user_event_rejects_full_queue( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(manager_module, "_MAX_SESSION_USER_EVENTS", 2) + runtime = object() + manager = _make_manager(_BaseTestManager, runtime) + managed, _video_track, _peer, _channel = _managed_session(runtime) + + for index in range(2): + manager._record_user_event( + managed_session=managed, + timestamp_s=float(index), + event_type="key_down", + payload={"key": "w"}, + ) + + with pytest.raises(RuntimeError, match="Too many queued WebRTC user events"): + manager._record_user_event( + managed_session=managed, + timestamp_s=2.0, + event_type="key_down", + payload={"key": "w"}, + ) + + assert len(managed.user_events) == 2 + + +def test_record_user_event_keeps_release_when_queue_full( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(manager_module, "_MAX_SESSION_USER_EVENTS", 1) + runtime = object() + manager = _make_manager(_BaseTestManager, runtime) + managed, _video_track, _peer, _channel = _managed_session(runtime) + manager._record_user_event( + managed_session=managed, + timestamp_s=0.0, + event_type="key_down", + payload={"key": "ArrowUp"}, + ) + + manager._record_user_event( + managed_session=managed, + timestamp_s=0.1, + event_type="key_up", + payload={"key": "w"}, + ) + + assert len(managed.user_events) == 1 + assert managed.user_events[0].event_type == "key_up" + assert managed.user_events[0].payload["key"] == "w" + + +def test_record_user_event_does_not_evict_unrelated_event_for_release( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(manager_module, "_MAX_SESSION_USER_EVENTS", 1) + runtime = object() + manager = _make_manager(_BaseTestManager, runtime) + managed, _video_track, _peer, _channel = _managed_session(runtime) + manager._record_user_event( + managed_session=managed, + timestamp_s=0.0, + event_type="text_event", + payload={"event_id": "storm"}, + ) + + manager._record_user_event( + managed_session=managed, + timestamp_s=0.1, + event_type="key_up", + payload={"key": "w"}, + ) + manager._record_user_event( + managed_session=managed, + timestamp_s=0.2, + event_type="key_up", + payload={"key": "w"}, + ) + + assert [(event.event_type, dict(event.payload)) for event in managed.user_events] == [ + ("text_event", {"event_id": "storm"}) + ] + assert len(managed.user_events) == 1 + assert set(managed.coalesced_release_events) == {"w"} + assert managed.coalesced_release_events["w"].timestamp_s == pytest.approx(0.2) + assert [ + (event.event_type, dict(event.payload)) + for event in manager._pending_user_events(managed) + ] == [ + ("text_event", {"event_id": "storm"}), + ("key_up", {"key": "w"}), + ] + + +def test_record_user_event_ignores_unsupported_key_events( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(manager_module, "_MAX_SESSION_USER_EVENTS", 1) + runtime = object() + manager = _make_manager(_WOnlyTestManager, runtime) + managed, _video_track, _peer, _channel = _managed_session(runtime) + manager._record_user_event( + managed_session=managed, + timestamp_s=0.0, + event_type="text_event", + payload={"event_id": "storm"}, + ) + + manager._record_user_event( + managed_session=managed, + timestamp_s=0.1, + event_type="key_up", + payload={"key": "z"}, + ) + manager._record_user_event( + managed_session=managed, + timestamp_s=0.2, + event_type="key_down", + payload={"key": "z"}, + ) + + assert [(event.event_type, dict(event.payload)) for event in managed.user_events] == [ + ("text_event", {"event_id": "storm"}) + ] + + +def test_catch_up_input_clock_advances_session_input_state() -> None: + class _RecordingCanonicalizer: + def __init__(self) -> None: + self.windows: list[tuple[float, float]] = [] + self.event_batches: list[list[tuple[float, str]]] = [] + + def canonicalize( + self, + user_inputs: Any, + *, + window: Any, + source_schema: Any, + ) -> object: + del source_schema + self.windows.append((window.start_s, window.end_s)) + self.event_batches.append( + [ + (event.timestamp_s, event.event_type) + for event in user_inputs.events + ] + ) + return object() + + canonicalizer = _RecordingCanonicalizer() + runtime = SimpleNamespace( + input_canonicalizer=canonicalizer, + input_source_schema=object(), + ) + manager = _make_manager(_BaseTestManager, runtime) + managed, _video_track, _peer, _channel = _managed_session(runtime) + managed.inference_session = object() + managed.resampler.next_chunk_start_v = 0.0 + manager._record_user_event( + managed_session=managed, + timestamp_s=0.5, + event_type="key_up", + payload={"key": "w"}, + ) + manager._record_user_event( + managed_session=managed, + timestamp_s=2.0, + event_type="key_down", + payload={"key": "w"}, + ) + + manager._catch_up_input_clock( + managed_session=managed, + now=3.0, + chunk_duration=1.0, + ) + + assert managed.resampler.next_chunk_start_v == pytest.approx(2.0) + assert managed.session_input_state_advanced + assert canonicalizer.windows == [(0.0, 2.0)] + assert canonicalizer.event_batches == [[(0.5, "key_up"), (2.0, "key_down")]] + assert [(event.timestamp_s, event.event_type) for event in managed.user_events] == [ + (pytest.approx(2.0), "key_down") + ] + + +def test_catch_up_input_clock_snaps_legacy_path_without_canonicalizer() -> None: + runtime = object() + manager = _make_manager(_BaseTestManager, runtime) + managed, _video_track, _peer, _channel = _managed_session(runtime) + managed.resampler.next_chunk_start_v = 0.0 + + manager._catch_up_input_clock( + managed_session=managed, + now=3.0, + chunk_duration=1.0, + ) + + assert managed.resampler.next_chunk_start_v == pytest.approx(2.0) + + +@pytest.mark.asyncio +async def test_action_keydown_reports_error_when_user_event_queue_full( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(manager_module, "_MAX_SESSION_USER_EVENTS", 1) + runtime = object() + manager = _make_manager(_BaseTestManager, runtime) + managed, _video_track, _peer, channel = _managed_session(runtime) + managed.inference_session = object() + managed.first_action_received.clear() + manager._record_user_event( + managed_session=managed, + timestamp_s=0.0, + event_type="key_down", + payload={"key": "w"}, + ) + + await manager._handle_datachannel_message( + managed_session=managed, + raw_message='{"type":"action","action":{"event":"keydown","key":"w"}}', + ) + + assert len(managed.user_events) == 1 + assert not managed.first_action_received.is_set() + assert [json.loads(message) for message in channel.messages] == [ + { + "type": "error", + "message": ( + "Too many queued WebRTC user events; wait for inference to catch up." + ), + } + ] + + +@pytest.mark.asyncio +async def test_action_keyup_updates_state_when_user_event_queue_full( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(manager_module, "_MAX_SESSION_USER_EVENTS", 1) + runtime = object() + manager = _make_manager(_BaseTestManager, runtime) + managed, _video_track, _peer, channel = _managed_session(runtime) + managed.inference_session = object() + managed.first_action_received.clear() + resampler = _RecordingResampler() + managed.resampler = resampler # ty:ignore[invalid-assignment] + manager._record_user_event( + managed_session=managed, + timestamp_s=0.0, + event_type="key_down", + payload={"key": "w"}, + ) + + await manager._handle_datachannel_message( + managed_session=managed, + raw_message='{"type":"action","action":{"event":"keyup","key":"w"}}', + ) + + assert [(event.event_type, dict(event.payload)) for event in managed.user_events] == [ + ("key_up", {"key": "w"}) + ] + assert managed.first_action_received.is_set() + assert len(managed.pending_action_arrivals) == 1 + assert len(resampler.edges) == 1 + assert resampler.edges[0][1:] == ("keyup", "w") + assert channel.messages == [] + + @pytest.mark.asyncio async def test_generation_worker_closes_session_when_flag_set() -> None: class _ClosingRuntime: @@ -479,13 +761,6 @@ class _WsadManager(_BaseTestManager): @pytest.mark.asyncio async def test_step_action_starts_generation_without_key_edge() -> None: - class _RecordingResampler(_FakeResampler): - def __init__(self) -> None: - self.edges: list[tuple[float, str, str]] = [] - - def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: - self.edges.append((arrival_t, event, key)) - runtime = SimpleNamespace() manager = _make_manager(_BaseTestManager, runtime) managed, _video_track, _peer, _channel = _managed_session(runtime) diff --git a/flashdreams/tests/test_webrtc_serving.py b/flashdreams/tests/test_webrtc_serving.py index c7b8a61b7..c06d434aa 100644 --- a/flashdreams/tests/test_webrtc_serving.py +++ b/flashdreams/tests/test_webrtc_serving.py @@ -3,8 +3,10 @@ from __future__ import annotations +import json from contextlib import nullcontext from importlib.resources import files +from typing import Any import numpy as np import pytest @@ -19,6 +21,10 @@ KeyboardState, ) from flashdreams.serving.webrtc.media import tensor_chunk_to_rgb_frames +from flashdreams.serving.webrtc.manager import ( + BaseWebRTCSessionManager, + ManagedWebRTCSession, +) from flashdreams.serving.webrtc.messages import ( make_chunk_done_payload, make_error_payload, @@ -119,6 +125,147 @@ async def shutdown(self) -> None: self.shutdown_calls += 1 +class _FakeCloseable: + async def close(self) -> None: + return + + +class _FakeControlChannel: + def __init__(self) -> None: + self.messages: list[dict[str, object]] = [] + + def send(self, payload: str) -> None: + decoded = json.loads(payload) + assert isinstance(decoded, dict) + self.messages.append(decoded) + + +class _Manager(BaseWebRTCSessionManager[Any, object]): + def _model_name(self) -> str: + return "fake" + + +def _managed_session_with_channel( + runtime: object, +) -> tuple[ManagedWebRTCSession, _FakeControlChannel]: + channel = _FakeControlChannel() + managed_session = ManagedWebRTCSession( + runtime=runtime, + video_track=_FakeCloseable(), # ty:ignore[invalid-argument-type] + video_encoder=_FakeCloseable(), # ty:ignore[invalid-argument-type] + peer_connection=_FakeCloseable(), + resampler=KeyboardResampler(fps=30, start_v=0.0), + control_channel=channel, + ) + return managed_session, channel + + +@pytest.mark.asyncio +async def test_event_message_dispatches_to_runtime_and_acknowledges() -> None: + class _FakeRuntime: + def __init__(self) -> None: + self.calls: list[tuple[str, str]] = [] + + async def trigger_event( + self, *, event_id: str, state: str + ) -> dict[str, object]: + self.calls.append((event_id, state)) + return {"active_event_id": event_id} + + runtime = _FakeRuntime() + manager = _Manager(runtime=runtime, runtime_config=object(), fps=30) + managed_session, channel = _managed_session_with_channel(runtime) + + 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_preserves_ack_fields() -> None: + class _FakeRuntime: + def __init__(self) -> None: + self.calls: list[tuple[str, str]] = [] + + 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, + } + + runtime = _FakeRuntime() + manager = _Manager(runtime=runtime, runtime_config=object(), fps=30) + managed_session, channel = _managed_session_with_channel(runtime) + + 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() -> None: + class _FakeRuntime: + def __init__(self) -> None: + self.calls = 0 + + async def trigger_event( + self, *, event_id: str, state: str + ) -> dict[str, object]: + del event_id, state + self.calls += 1 + return {} + + runtime = _FakeRuntime() + manager = _Manager(runtime=runtime, runtime_config=object(), fps=30) + managed_session, channel = _managed_session_with_channel(runtime) + + await manager._handle_datachannel_message( + managed_session=managed_session, + raw_message='{"type":"event","state":"trigger"}', + ) + + assert runtime.calls == 0 + assert channel.messages == [ + { + "type": "error", + "message": ( + "Event payload must include non-empty 'event_id' " + "unless state clears the active event." + ), + } + ] + assert not managed_session.first_action_received.is_set() + + def test_packaged_webrtc_app_keeps_resource_materialized(tmp_path) -> None: (tmp_path / "request_session.html").write_text( "session", encoding="utf-8" diff --git a/integrations/lingbot/lingbot/webrtc/session.py b/integrations/lingbot/lingbot/webrtc/session.py index 62dbb8442..e9b79c283 100644 --- a/integrations/lingbot/lingbot/webrtc/session.py +++ b/integrations/lingbot/lingbot/webrtc/session.py @@ -24,6 +24,7 @@ import re import socket import ssl +import threading import urllib.parse from dataclasses import dataclass, field from pathlib import Path @@ -58,6 +59,20 @@ WebRTCControlSignal, ) from flashdreams.serving.webrtc.server import SessionBusyError +from flashdreams.runtime.canonical import InputCanonicalizer +from flashdreams.runtime.inputs import ( + InferenceInput, + UserInputCapability, + UserInputSchema, +) +from flashdreams.runtime.types import StepRequest, StepResult +from lingbot.input_mapping import ( + FIELD_CAMERA_INTRINSICS, + FIELD_CAMERA_TRAJECTORY, + KeyboardToCameraCommand, + LingbotInputMapping, + TextEventSelection, +) from lingbot.encoder.utils import preprocess_example_poses _INTRINSICS_REFERENCE_HEIGHT = 480 @@ -603,7 +618,11 @@ def __init__(self, config: LingbotRuntimeConfig | None = None) -> None: self._prompt: str | None = None self._base_text_embeddings: torch.Tensor | None = None self._event_embeddings: dict[str, torch.Tensor] = {} + self._prompt_embeddings: dict[str, torch.Tensor] = {} self._active_event_id: str | None = None + self._input_mapping: LingbotInputMapping | 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 @@ -675,6 +694,100 @@ async def trigger_event( 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, *, @@ -684,16 +797,14 @@ async def generate_chunk( """Generate one autoregressive chunk from a piecewise-constant timeline. Args: - segments: Piecewise-constant keyboard-state segments - covering the chunk's virtual-time window; produced by - :meth:`KeyboardResampler.sample_chunk`. - 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. + 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: - :class:`VideoStepResult` carrying the produced video chunk - and the post-generation pipeline stats. + Video chunk and post-generation pipeline stats. Raises: LingbotRuntimeError: Runtime is closed or not initialized. @@ -766,6 +877,12 @@ def _generate_chunk_sync_all_ranks( ) -> 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, @@ -859,6 +976,7 @@ def _precompute_event_embeddings_sync( ) -> None: if not text_events: self._event_embeddings = {} + self._prompt_embeddings = {} return event_ids = [event.event_id for event in text_events] if len(event_ids) != len(set(event_ids)): @@ -869,6 +987,13 @@ def _precompute_event_embeddings_sync( event_id: embeddings[index : index + 1].contiguous() for index, event_id in enumerate(event_ids) } + # The session branch receives a prompt rather than an event id, so keep + # a prompt-keyed view of the same tensors. Without it a live text event + # would pay a text-encoder pass mid-rollout. + self._prompt_embeddings = { + prompt: self._event_embeddings[event_id] + for prompt, event_id in zip(prompts, event_ids, strict=True) + } def _build_base_intrinsics(self) -> torch.Tensor: if self._device is None: @@ -1068,6 +1193,9 @@ def _reset_rollout_sync( text=[self._prompt], image=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: @@ -1168,12 +1296,34 @@ def _generate_one_chunk_sync( poses_t = torch.from_numpy(poses).to(device=self._device, dtype=torch.float32) poses_t = poses_t.view(num_frames, 4, 4) intrinsics_t = self._base_intrinsics.view(1, 4).repeat(num_frames, 1) + return self._generate_from_camera_inputs( + poses=poses_t, + intrinsics=intrinsics_t, + num_frames=num_frames, + ) + + def _generate_from_camera_inputs( + self, + *, + poses: torch.Tensor, + intrinsics: torch.Tensor, + num_frames: int, + ) -> VideoStepResult: + """Generate one chunk from an already-resolved camera trajectory. + + Shared by the segment path and the mapped-input session path so both + reach the model through identical conditioning. + """ + if self._pipeline is None or self._cache is None: + raise LingbotRuntimeError("Runtime is not initialized.") + if self._device is None: + raise LingbotRuntimeError("Runtime device is not initialized.") from lingbot.encoder.camctrl import CamCtrlInput # noqa: PLC0415 camctrl_input = CamCtrlInput( - intrinsics=intrinsics_t, - poses=poses_t, + intrinsics=intrinsics.to(device=self._device, dtype=torch.float32), + poses=poses.to(device=self._device, dtype=torch.float32), world_scale=self._world_scale, ) video_chunk = self._pipeline.generate( @@ -1196,6 +1346,115 @@ def _generate_one_chunk_sync( self.autoregressive_index += 1 return result + def _step_sync(self, inputs: InferenceInput) -> StepResult: + """Generate one chunk from mapped model inputs.""" + if self._pipeline is None or self._cache is None: + raise LingbotRuntimeError("Runtime is not initialized.") + num_frames = int( + self._pipeline.get_num_output_frames(self.autoregressive_index) + ) + self._apply_conditioning_update_sync(inputs) + poses = _require_camera_tensor( + inputs, FIELD_CAMERA_TRAJECTORY, expected_shape=(num_frames, 4, 4) + ) + intrinsics = _require_camera_tensor( + inputs, FIELD_CAMERA_INTRINSICS, expected_shape=(num_frames, 4) + ) + step_index = self.autoregressive_index + result = self._generate_from_camera_inputs( + poses=poses, + intrinsics=intrinsics, + num_frames=num_frames, + ) + return StepResult( + step_index=step_index, + output=result, + frame_count=result.num_frames, + metrics=result.stats or {}, + ) + + def _apply_conditioning_update_sync(self, inputs: InferenceInput) -> None: + """Apply a text-event prompt swap requested by the mapping.""" + prompt = inputs.global_conditioning.get("prompt") + if prompt is None or prompt == self._prompt: + return + embeddings = self._prompt_embeddings.get(prompt) + if embeddings is None: + embeddings = self._encode_text_embeddings_sync([prompt]) + self._replace_rollout_text_embeddings(embeddings) + self._prompt = prompt + self._active_event_id = next( + ( + event_id + for event_id, tensor in self._event_embeddings.items() + if tensor is embeddings + ), + None, + ) + + +def _require_camera_tensor( + inputs: InferenceInput, + name: str, + *, + expected_shape: tuple[int, ...], +) -> torch.Tensor: + """Return one required per-step camera tensor, shape-checked.""" + if name not in inputs.step: + raise LingbotRuntimeError( + f"Lingbot step inputs are missing {name!r}; the selected input " + f"mapping must produce it for every step." + ) + value = inputs.step[name] + if not isinstance(value, torch.Tensor): + value = torch.as_tensor(np.asarray(value), dtype=torch.float32) + if tuple(value.shape) != expected_shape: + raise LingbotRuntimeError( + f"Lingbot step input {name!r} must have shape {expected_shape}, got " + f"{tuple(value.shape)}." + ) + return value + + +LINGBOT_WEBRTC_SOURCE_SCHEMA = UserInputSchema( + capabilities=( + 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"}) + ), + ), + description="Lingbot WebRTC data-channel input.", +) + + +class LingbotWebRTCInferenceSession: + """``InferenceSession`` view of a live Lingbot WebRTC rollout. + + The rollout itself is owned by :class:`LingbotInferenceRuntime`; this only + adapts it to the runtime-API stepping surface so the shared manager can + drive it with mapped inputs. + """ + + def __init__(self, *, runtime: LingbotInferenceRuntime) -> None: + self._runtime = runtime + + def next_step_request(self) -> StepRequest | None: + return self._runtime._next_step_request_sync() + + def step(self, inputs: InferenceInput) -> StepResult: + return self._runtime._step_blocking(inputs) + + def reset(self, inputs: InferenceInput | None = None) -> None: + raise LingbotRuntimeError( + "Reset a Lingbot WebRTC rollout through the runtime's session " + "lifecycle, not through the inference session." + ) + + def close(self) -> None: + # The runtime outlives the session and is closed by the serve loop. + return None + _ManagedLingbotSession = ManagedWebRTCSession diff --git a/integrations/lingbot/tests/test_webrtc_runtime.py b/integrations/lingbot/tests/test_webrtc_runtime.py index 4153efb1b..e38061a38 100644 --- a/integrations/lingbot/tests/test_webrtc_runtime.py +++ b/integrations/lingbot/tests/test_webrtc_runtime.py @@ -17,12 +17,16 @@ import asyncio import ipaddress -import json from pathlib import Path from typing import cast import pytest import torch +from lingbot.input_mapping import ( + KeyboardToCameraCommand, + LingbotInputMapping, + TextEventSelection, +) from lingbot.webrtc import session from lingbot.webrtc.session import ( LingbotRuntimeConfig, @@ -30,6 +34,9 @@ ) 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 pytestmark = pytest.mark.ci_cpu @@ -42,16 +49,6 @@ async def close(self) -> None: self.closed = True -class _FakeControlChannel: - def __init__(self) -> None: - self.messages: list[dict[str, object]] = [] - - def send(self, payload: str) -> None: - decoded = json.loads(payload) - assert isinstance(decoded, dict) - self.messages.append(decoded) - - class _FakeVideoEncoder: """Minimal ``VideoEncoder``-shaped stub for ``_ManagedLingbotSession`` construction. Enough to satisfy the dataclass field; the tests here do @@ -570,149 +567,60 @@ def _fail_remote_fetch(image_url: str) -> object: assert runtime._prompt == "follow a coastal highway" -@pytest.mark.asyncio -async def test_event_message_dispatches_to_runtime_and_acknowledges( - monkeypatch: pytest.MonkeyPatch, -) -> None: - class _FakeRuntime: +def test_apply_conditioning_update_swaps_precomputed_text_embeddings() -> None: + class _FakeTransformer: def __init__(self) -> None: - self.calls: list[tuple[str, str]] = [] - - async def trigger_event( - self, *, event_id: str, state: str - ) -> dict[str, object]: - self.calls.append((event_id, state)) - return {"active_event_id": event_id} - - monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) - manager = LingbotWebRTCSessionManager( - runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) - ) - runtime = _FakeRuntime() - channel = _FakeControlChannel() - managed_session = session._ManagedLingbotSession( - 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"}', - ) + self.calls: list[tuple[object, torch.Tensor]] = [] - 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() + def replace_text_embeddings( + self, cache: object, text_embeddings: torch.Tensor + ) -> None: + self.calls.append((cache, text_embeddings)) + class _FakeDiffusionModel: + def __init__(self) -> None: + self.transformer = _FakeTransformer() -@pytest.mark.asyncio -async def test_clear_event_message_does_not_require_event_id_and_preserves_ack_fields( - monkeypatch: pytest.MonkeyPatch, -) -> None: - class _FakeRuntime: + class _FakePipeline: def __init__(self) -> None: - self.calls: list[tuple[str, str]] = [] - - 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, - } + self.diffusion_model = _FakeDiffusionModel() - monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) - manager = LingbotWebRTCSessionManager( - runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) - ) - runtime = _FakeRuntime() - channel = _FakeControlChannel() - managed_session = session._ManagedLingbotSession( - 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, + 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, + } - await manager._handle_datachannel_message( - managed_session=managed_session, - raw_message='{"type":"event","state":"clear"}', + runtime._apply_conditioning_update_sync( + InferenceInput(global_conditioning={"prompt": "a glowing portal opens"}) ) - 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.calls = 0 - - async def trigger_event( - self, *, event_id: str, state: str - ) -> dict[str, object]: - del event_id, state - self.calls += 1 - return {} + transformer = runtime._pipeline.diffusion_model.transformer + assert runtime._active_event_id == "portal" + assert runtime._prompt == "a glowing portal opens" + assert transformer.calls == [(transformer_cache, event_text)] - monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) - manager = LingbotWebRTCSessionManager( - runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) - ) - runtime = _FakeRuntime() - channel = _FakeControlChannel() - managed_session = session._ManagedLingbotSession( - 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, + runtime._apply_conditioning_update_sync( + InferenceInput(global_conditioning={"prompt": "base prompt"}) ) - await manager._handle_datachannel_message( - managed_session=managed_session, - raw_message='{"type":"event","state":"trigger"}', - ) - - assert runtime.calls == 0 - assert channel.messages == [ - { - "type": "error", - "message": ( - "Event payload must include non-empty 'event_id' " - "unless state clears the active event." - ), - } - ] - assert not managed_session.first_action_received.is_set() + assert runtime._active_event_id is None + assert runtime._prompt == "base prompt" + assert transformer.calls[-1] == (transformer_cache, base_text) def test_trigger_event_sync_swaps_precomputed_text_embeddings() -> None: @@ -763,6 +671,36 @@ def __init__(self) -> None: assert transformer.calls[-1] == (transformer_cache, base_text) +def test_validate_user_event_rejects_invalid_text_events() -> None: + runtime = session.LingbotInferenceRuntime( + config=LingbotRuntimeConfig( + device="cpu", + warmup_chunks=0, + text_events=(), + ) + ) + runtime._event_embeddings = {"portal": torch.ones((1, 2, 3))} + + assert runtime.validate_user_event( + event_type="text_event", + payload={"event_id": "portal", "state": "trigger"}, + ) == {"event_id": "portal", "state": "trigger"} + assert runtime.validate_user_event( + event_type="text_event", + payload={"event_id": None, "state": "clear"}, + ) == {"event_id": None, "state": "clear"} + with pytest.raises(ValueError, match="Unknown event_id='unknown'"): + runtime.validate_user_event( + event_type="text_event", + payload={"event_id": "unknown", "state": "trigger"}, + ) + with pytest.raises(ValueError, match="Event state must be one of"): + runtime.validate_user_event( + event_type="text_event", + payload={"event_id": "portal", "state": "explode"}, + ) + + def test_reset_rollout_precomputes_session_text_events() -> None: class _FakePipeline: def __init__(self) -> None: @@ -935,15 +873,52 @@ async def _fake_loopback_warmup( async def test_loopback_warmup_drives_session_generation( monkeypatch: pytest.MonkeyPatch, ) -> None: + class _FakeInferenceSession: + def __init__(self, runtime: "_FakeRuntime") -> None: + self._runtime = runtime + + def next_step_request(self) -> StepRequest: + step_index = self._runtime.step_index + return StepRequest( + step_index=step_index, + metadata={"num_frames": 1, "frame_start": step_index}, + ) + + def step(self, inputs: InferenceInput) -> StepResult: + chunk_index = self._runtime.step_index + self._runtime.step_index += 1 + self._runtime.generated_inputs.append(inputs) + return StepResult( + step_index=chunk_index, + output=VideoStepResult( + chunk_index=chunk_index, + num_frames=1, + video_chunk=torch.zeros((1, 1, 1, 3, 2, 2), dtype=torch.uint8), + stats=None, + ), + frame_count=1, + ) + class _FakeRuntime: def __init__(self, config: LingbotRuntimeConfig) -> None: self.config = config self.initialize_calls = 0 self.reset_calls = 0 self.close_calls = 0 - self.generated_segments: list[ - list[tuple[float, float, frozenset[str]]] - ] = [] + self.step_index = 0 + self.generated_inputs: list[InferenceInput] = [] + self.input_canonicalizer = InputCanonicalizer( + [KeyboardToCameraCommand(), TextEventSelection()] + ) + self.input_mapping = LingbotInputMapping( + fps=30, + base_intrinsics=torch.tensor([416.0, 416.0, 416.0, 240.0]), + world_scale=1.0, + text_event_prompts={}, + ) + self.input_mapping.set_base_prompt("warmup prompt") + self.input_source_schema = session.LINGBOT_WEBRTC_SOURCE_SCHEMA + self._active_event_id = None async def initialize(self) -> None: self.initialize_calls += 1 @@ -960,21 +935,8 @@ def peek_steady_chunk_num_frames(self) -> int: 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 start_inference_session(self) -> _FakeInferenceSession: + return _FakeInferenceSession(self) async def close(self) -> None: self.close_calls += 1 @@ -1000,7 +962,7 @@ def _fake_runtime_factory(config: LingbotRuntimeConfig) -> _FakeRuntime: assert fake_runtime is not None assert fake_runtime.initialize_calls == 1 assert fake_runtime.reset_calls == 1 - assert len(fake_runtime.generated_segments) == 2 + assert len(fake_runtime.generated_inputs) == 2 assert not manager.has_active_session() diff --git a/integrations/lingbot/tests/test_webrtc_session_branch.py b/integrations/lingbot/tests/test_webrtc_session_branch.py new file mode 100644 index 000000000..64e1728e5 --- /dev/null +++ b/integrations/lingbot/tests/test_webrtc_session_branch.py @@ -0,0 +1,407 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The manager's ``InferenceSession`` branch must preserve camera controls. + +The session branch buffers raw events, canonicalizes them over the chunk +window, and maps them into per-step ``InferenceInput``. The resulting camera +trajectory must match the direct resampler/integrator reference, or moving +LingBot's live path onto the runtime API would silently change how it drives. +""" + +from __future__ import annotations + +import json +from typing import Any + +import numpy as np +import pytest +import torch +from lingbot.input_mapping import ( + FIELD_CAMERA_INTRINSICS, + FIELD_CAMERA_TRAJECTORY, + KeyboardToCameraCommand, + LingbotInputMapping, + TextEventSelection, +) +from lingbot.webrtc.session import LINGBOT_WEBRTC_SOURCE_SCHEMA + +from flashdreams.infra.video_output import VideoStepResult +from flashdreams.runtime.inputs import InferenceInput, TimeWindow +from flashdreams.runtime.canonical import InputCanonicalizer +from flashdreams.runtime.types import StepRequest, StepResult +from flashdreams.serving.realtime.input import KeyboardResampler +from flashdreams.serving.webrtc.controls import CameraPoseIntegrator +from flashdreams.serving.webrtc.manager import ( + BaseWebRTCSessionManager, + ManagedWebRTCSession, +) + +pytestmark = pytest.mark.ci_cpu + +_FPS = 16 +_NUM_FRAMES = 4 +_BASE_INTRINSICS = torch.tensor([416.0, 416.0, 416.0, 240.0]) + + +class _FakeRuntimeConfig: + video_width = 64 + video_height = 64 + warmup_chunks = 0 + warmup_timeout_s = 1.0 + + +class _FakeSession: + """Records what the manager hands the model.""" + + def __init__(self) -> None: + self.steps: list[InferenceInput] = [] + self._index = 0 + + def next_step_request(self) -> StepRequest: + return StepRequest( + step_index=self._index, + metadata={ + "num_frames": _NUM_FRAMES, + "frame_start": self._index * _NUM_FRAMES, + }, + ) + + def step(self, inputs: InferenceInput) -> StepResult: + self.steps.append(inputs) + index = self._index + self._index += 1 + return StepResult( + step_index=index, + output=VideoStepResult( + chunk_index=index, + video_chunk=torch.zeros(_NUM_FRAMES, 3, 4, 4), + layout="tchw", + num_frames=_NUM_FRAMES, + stats={}, + ), + frame_count=_NUM_FRAMES, + ) + + +class _FakeRuntime: + def __init__(self, *, text_event_prompts: dict[str, str] | None = None) -> None: + self.text_event_prompts = dict(text_event_prompts or {}) + self.input_canonicalizer = InputCanonicalizer( + [KeyboardToCameraCommand(), TextEventSelection()] + ) + self.input_mapping = LingbotInputMapping( + fps=_FPS, + base_intrinsics=_BASE_INTRINSICS, + world_scale=1.0, + text_event_prompts=text_event_prompts, + ) + self.input_mapping.set_base_prompt("a calm street") + self.input_source_schema = LINGBOT_WEBRTC_SOURCE_SCHEMA + self.session = _FakeSession() + + async def start_inference_session(self) -> _FakeSession: + return self.session + + def validate_user_event( + self, *, event_type: str, payload: dict[str, Any] + ) -> dict[str, Any] | None: + if event_type != "text_event": + return payload + state = str(payload.get("state", "trigger")).strip().lower() or "trigger" + clear_states = {"clear", "release", "off", "none"} + trigger_states = {"trigger", "hold", "on"} + if state not in clear_states and state not in trigger_states: + raise ValueError(f"Unsupported text event state {state!r}.") + event_id = payload.get("event_id") + if event_id is None or state in clear_states: + return {"event_id": None, "state": state} + event_id = str(event_id) + if event_id not in self.text_event_prompts: + raise ValueError(f"Unknown Lingbot text event_id={event_id!r}.") + return {"event_id": event_id, "state": state} + + +class _Manager(BaseWebRTCSessionManager[Any, Any]): + def _model_name(self) -> str: + return "fake" + + +class _FakeControlChannel: + def __init__(self) -> None: + self.messages: list[dict[str, object]] = [] + + def send(self, payload: str) -> None: + decoded = json.loads(payload) + assert isinstance(decoded, dict) + self.messages.append(decoded) + + +def _managed_session(runtime: _FakeRuntime) -> ManagedWebRTCSession: + return ManagedWebRTCSession( + runtime=runtime, + video_track=None, + video_encoder=None, + peer_connection=None, + resampler=KeyboardResampler(fps=_FPS, start_v=0.0), + inference_session=runtime.session, + ) + + +def _manager(runtime: _FakeRuntime) -> _Manager: + return _Manager(runtime=runtime, runtime_config=_FakeRuntimeConfig(), fps=_FPS) + + +def _reference_poses( + edges: list[tuple[float, str, str]], *, chunks: int +) -> np.ndarray: + resampler = KeyboardResampler(fps=_FPS, start_v=0.0) + integrator = CameraPoseIntegrator() + for timestamp_s, event, key in edges: + resampler.on_edge(arrival_t=timestamp_s, event=event, key=key) + poses = [] + for _ in range(chunks): + segments, frame_times = resampler.sample_chunk(_NUM_FRAMES) + poses.append( + integrator.integrate_chunk(segments=segments, frame_times=frame_times) + ) + return np.concatenate(poses) + + +def _session_branch_poses( + edges: list[tuple[float, str, str]], *, chunks: int +) -> np.ndarray: + runtime = _FakeRuntime() + manager = _manager(runtime) + managed = _managed_session(runtime) + for timestamp_s, event, key in edges: + manager._record_user_event( + managed_session=managed, + timestamp_s=timestamp_s, + event_type="key_down" if event == "keydown" else "key_up", + payload={"key": key}, + ) + + poses = [] + for chunk_index in range(chunks): + start_s = chunk_index * _NUM_FRAMES / _FPS + end_s = (chunk_index + 1) * _NUM_FRAMES / _FPS + window = TimeWindow(start_s=start_s, end_s=end_s) + request = runtime.session.next_step_request() + from dataclasses import replace + + step_inputs = manager._build_step_inputs( + managed_session=managed, + request=replace(request, user_input_window=window), + window=window, + ) + runtime.session.step(step_inputs) + manager._prune_consumed_user_events(managed, before_s=start_s) + poses.append(step_inputs.step[FIELD_CAMERA_TRAJECTORY].numpy()) + return np.concatenate(poses) + + +@pytest.mark.parametrize( + "edges", + [ + pytest.param([(0.0, "keydown", "w")], id="hold_forward"), + pytest.param( + [(0.0, "keydown", "w"), (0.13, "keydown", "a")], id="mid_chunk_turn" + ), + pytest.param( + [(0.0, "keydown", "w"), (0.25, "keydown", "d")], id="chunk_boundary" + ), + pytest.param( + [(0.01, "keydown", "w"), (0.04, "keyup", "w"), (0.08, "keydown", "w")], + id="rapid_toggle", + ), + pytest.param([], id="idle"), + ], +) +def test_session_branch_matches_reference_camera_integration( + edges: list[tuple[float, str, str]], +) -> None: + reference = _reference_poses(edges, chunks=3) + session_branch = _session_branch_poses(edges, chunks=3) + + assert reference.shape == session_branch.shape + np.testing.assert_allclose(session_branch, reference, atol=1e-5) + + +def test_session_branch_supplies_intrinsics_for_every_step() -> None: + runtime = _FakeRuntime() + manager = _manager(runtime) + managed = _managed_session(runtime) + manager._record_user_event( + managed_session=managed, + timestamp_s=0.0, + event_type="key_down", + payload={"key": "w"}, + ) + window = TimeWindow(start_s=0.0, end_s=_NUM_FRAMES / _FPS) + + step_inputs = manager._build_step_inputs( + managed_session=managed, + request=runtime.session.next_step_request(), + window=window, + ) + + assert step_inputs.step[FIELD_CAMERA_INTRINSICS].shape == (_NUM_FRAMES, 4) + assert torch.allclose( + step_inputs.step[FIELD_CAMERA_INTRINSICS][0], _BASE_INTRINSICS + ) + + +def test_consumed_events_are_pruned() -> None: + runtime = _FakeRuntime() + manager = _manager(runtime) + managed = _managed_session(runtime) + for index in range(5): + manager._record_user_event( + managed_session=managed, + timestamp_s=index * 0.1, + event_type="key_down", + payload={"key": "w"}, + ) + + manager._prune_consumed_user_events(managed, before_s=0.25) + + # Held-key state lives in the converter, so consumed events are safe to drop + # and must be, or a long session's buffer grows without bound. + assert [event.timestamp_s for event in managed.user_events] == [ + pytest.approx(0.3), + pytest.approx(0.4), + ] + + +@pytest.mark.asyncio +async def test_catch_up_window_clears_release_and_renders_latest_step() -> None: + runtime = _FakeRuntime() + manager = _manager(runtime) + managed = _managed_session(runtime) + manager._record_user_event( + managed_session=managed, + timestamp_s=0.1, + event_type="key_down", + payload={"key": "w"}, + ) + manager._record_user_event( + managed_session=managed, + timestamp_s=0.2, + event_type="key_up", + payload={"key": "w"}, + ) + manager._record_user_event( + managed_session=managed, + timestamp_s=0.8, + event_type="key_down", + payload={"key": "w"}, + ) + + manager._catch_up_input_clock( + managed_session=managed, + now=1.0, + chunk_duration=_NUM_FRAMES / _FPS, + ) + await manager._step_inference_session( + managed_session=managed, + window=TimeWindow(start_s=0.75, end_s=1.0), + ) + + assert [(event.timestamp_s, event.event_type) for event in managed.user_events] == [ + (pytest.approx(0.8), "key_down") + ] + poses = runtime.session.steps[0].step[FIELD_CAMERA_TRAJECTORY] + assert not torch.allclose(poses[:, :3, 3], torch.zeros_like(poses[:, :3, 3])) + + +@pytest.mark.asyncio +async def test_text_event_becomes_a_buffered_user_event() -> None: + runtime = _FakeRuntime(text_event_prompts={"storm": "a violent storm"}) + manager = _manager(runtime) + managed = _managed_session(runtime) + assert not hasattr(runtime, "trigger_event") + + handled = await manager._handle_event_message( + managed_session=managed, + payload={"event_id": "storm", "state": "trigger"}, + ) + + assert handled is True + assert [event.event_type for event in managed.user_events] == ["text_event"] + assert managed.user_events[0].payload["event_id"] == "storm" + + # A text event can itself be the first interaction, so it is stamped just + # before the resampler re-anchors its clock. Chunk 0 must still see it. + anchor = managed.user_events[0].timestamp_s + 0.05 + await manager._step_inference_session( + managed_session=managed, + window=TimeWindow(start_s=anchor, end_s=anchor + _NUM_FRAMES / _FPS), + ) + + assert len(runtime.session.steps) == 1 + assert runtime.session.steps[0].global_conditioning["prompt"] == "a violent storm" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("payload", "message"), + [ + pytest.param( + {"event_id": "unknown", "state": "trigger"}, + "Unknown Lingbot text event_id='unknown'", + id="unknown_event", + ), + pytest.param( + {"event_id": "storm", "state": "explode"}, + "Unsupported text event state 'explode'", + id="bad_state", + ), + ], +) +async def test_text_event_rejects_invalid_payload_before_ack( + payload: dict[str, str], + message: str, +) -> None: + runtime = _FakeRuntime(text_event_prompts={"storm": "a violent storm"}) + manager = _manager(runtime) + managed = _managed_session(runtime) + channel = _FakeControlChannel() + managed.control_channel = channel + + handled = await manager._handle_event_message( + managed_session=managed, + payload=payload, + ) + + assert handled is False + assert list(managed.user_events) == [] + assert channel.messages[0]["type"] == "error" + assert message in str(channel.messages[0]["message"]) + assert len(channel.messages) == 1 + + +def test_real_lingbot_runtime_selects_the_session_branch() -> None: + """The shipped runtime must be session-capable while retaining segment stepping.""" + from lingbot.webrtc.session import LingbotInferenceRuntime, LingbotRuntimeConfig + + runtime = LingbotInferenceRuntime(config=LingbotRuntimeConfig(device="cpu")) + + assert BaseWebRTCSessionManager._drives_inference_session(runtime) is True + assert callable(runtime.generate_chunk) + + +def test_session_start_requires_an_initialized_rollout() -> None: + """Starting a session before reset must fail loudly, not silently no-op.""" + import asyncio + + from lingbot.webrtc.session import ( + LingbotInferenceRuntime, + LingbotRuntimeConfig, + LingbotRuntimeError, + ) + + runtime = LingbotInferenceRuntime(config=LingbotRuntimeConfig(device="cpu")) + + with pytest.raises(LingbotRuntimeError, match="input mapping is not initialized"): + asyncio.run(runtime.start_inference_session()) From 209f936c9935da68a1695440a42adbb1061ff657 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Fri, 7 Aug 2026 21:19:10 -0700 Subject: [PATCH 14/19] Unify serving paths (#429) * Canonicalize generated video step results * Simplify video output stream consumption * Pass step results through WebRTC delivery * Run serving runtimes on thread-affine workers * Unify Lingbot model session execution * Unify OmniDreams model session execution * Make WebRTC manager capabilities explicit * Define explicit WebRTC app adapter contracts * Route outputs through integration capabilities * Drive WebRTC generation with step requests * Fix serving type-check regressions * Apply repository-wide lint fixes * Skip unavailable Transformer Engine in CPU tests * Record serving architecture validation * Unify WebRTC session manager implementations * Consolidate WebRTC runtime lifecycle * refactor(omnidreams): use shared WebRTC demo APIs * refactor(omnidreams): remove legacy WebRTC implementation * refactor(omnidreams): remove WebRTC postprocessing * Consolidate WebRTC demo integrations * Simplify WebRTC demo launch path * Share demo application lifecycle --- ...ntime_serving_architecture_improvements.md | 417 +++++ .../flashdreams/infra/postprocess/__init__.py | 2 + .../flashdreams/infra/postprocess/stream.py | 35 +- flashdreams/flashdreams/infra/results.py | 136 ++ flashdreams/flashdreams/infra/runner.py | 8 +- flashdreams/flashdreams/infra/time.py | 32 + flashdreams/flashdreams/infra/video_output.py | 310 ++-- flashdreams/flashdreams/runtime/__init__.py | 2 + .../flashdreams/runtime/demo/__init__.py | 5 +- flashdreams/flashdreams/runtime/demo/app.py | 77 +- flashdreams/flashdreams/runtime/demo/spec.py | 18 +- .../flashdreams/runtime/demo/webrtc.py | 240 +-- flashdreams/flashdreams/runtime/inputs.py | 21 +- flashdreams/flashdreams/runtime/runner.py | 5 +- flashdreams/flashdreams/runtime/types.py | 22 +- .../flashdreams/runtime/video_output.py | 59 +- flashdreams/flashdreams/runtime/worker.py | 108 ++ .../flashdreams/serving/output_targets.py | 248 +-- .../flashdreams/serving/realtime/media.py | 22 +- .../flashdreams/serving/webrtc/encoders.py | 16 +- .../flashdreams/serving/webrtc/manager.py | 559 ++----- .../flashdreams/serving/webrtc/media.py | 34 +- .../flashdreams/serving/webrtc/nvenc.py | 44 +- .../flashdreams/serving/webrtc/runtime.py | 316 +++- .../serving/webrtc/web/mock_ui_server.py | 17 +- .../serving/webrtc/web/request_session.js | 30 +- flashdreams/tests/test_encoders.py | 151 +- flashdreams/tests/test_output_targets.py | 76 +- flashdreams/tests/test_rope_kernel.py | 33 +- flashdreams/tests/test_runtime_demo_api.py | 128 +- .../tests/test_runtime_video_output.py | 18 +- flashdreams/tests/test_runtime_worker.py | 97 ++ flashdreams/tests/test_runtime_worker_gpu.py | 61 + flashdreams/tests/test_video_output.py | 217 ++- flashdreams/tests/test_webrtc_manager.py | 187 ++- flashdreams/tests/test_webrtc_serving.py | 4 +- .../causal_forcing/causal_forcing/runner.py | 36 +- .../cosmos_predict2/cosmos_predict2/runner.py | 28 +- .../fastvideo_causal_wan22/runner.py | 36 +- integrations/flashvsr/flashvsr/runner.py | 48 +- .../hy_worldplay/hy_worldplay/runner.py | 40 +- integrations/lingbot/lingbot/demo/adapter.py | 112 +- .../lingbot/lingbot/demo/{cli.py => app.py} | 57 +- integrations/lingbot/lingbot/demo/spec.py | 6 +- integrations/lingbot/lingbot/demo/webrtc.py | 129 +- integrations/lingbot/lingbot/input_mapping.py | 14 +- integrations/lingbot/lingbot/model_session.py | 130 ++ .../lingbot/lingbot/output_targets.py | 82 + integrations/lingbot/lingbot/runner.py | 7 +- integrations/lingbot/lingbot/runtime.py | 187 +-- integrations/lingbot/lingbot/webrtc/server.py | 72 +- .../lingbot/lingbot/webrtc/session.py | 548 ++---- .../lingbot/lingbot/webrtc/web/adapter.js | 9 + integrations/lingbot/pyproject.toml | 2 +- integrations/lingbot/tests/test_demo_api.py | 137 +- .../tests/test_distributed_server_main.py | 12 +- .../lingbot/tests/test_input_mapping.py | 18 +- .../lingbot/tests/test_runtime_gpu.py | 21 +- .../tests/test_runtime_session_inputs.py | 6 +- .../lingbot/tests/test_server_routes.py | 2 +- integrations/lingbot/tests/test_smoke.py | 15 +- .../lingbot/tests/test_webrtc_runtime.py | 300 +++- .../tests/test_webrtc_runtime_distributed.py | 2 +- .../omnidreams/omnidreams/demo/README.md | 6 +- .../omnidreams/omnidreams/demo/adapter.py | 124 +- .../omnidreams/demo/{cli.py => app.py} | 56 +- .../omnidreams/omnidreams/demo/replay.py | 75 +- .../omnidreams/omnidreams/demo/spec.py | 9 +- .../omnidreams/omnidreams/demo/web/adapter.js | 17 + .../omnidreams/omnidreams/demo/webrtc.py | 697 ++++++-- .../world_model/flashdreams_adapter.py | 116 +- .../omnidreams/omnidreams/model_session.py | 167 ++ .../omnidreams/omnidreams/output_targets.py | 143 ++ integrations/omnidreams/omnidreams/runner.py | 40 +- integrations/omnidreams/omnidreams/scenes.py | 346 +++- .../omnidreams/omnidreams/webrtc/__init__.py | 4 - .../omnidreams/omnidreams/webrtc/server.py | 377 ----- .../omnidreams/omnidreams/webrtc/session.py | 1189 ------------- .../omnidreams/webrtc/web/adapter.js | 7 - integrations/omnidreams/pyproject.toml | 8 +- .../test_world_model_adapter.py | 1 + .../omnidreams/tests/test_demo_api.py | 253 ++- .../omnidreams/tests/test_nvenc_smoke.py | 8 +- .../omnidreams/tests/test_webrtc_runtime.py | 1482 ----------------- .../tests/test_webrtc_server_routes.py | 324 ---- .../self_forcing/self_forcing/runner.py | 36 +- integrations/wan21/wan21/runner.py | 28 +- 87 files changed, 4953 insertions(+), 6371 deletions(-) create mode 100644 docs/inference_runtime_serving_architecture_improvements.md create mode 100644 flashdreams/flashdreams/infra/results.py create mode 100644 flashdreams/flashdreams/infra/time.py create mode 100644 flashdreams/flashdreams/runtime/worker.py create mode 100644 flashdreams/tests/test_runtime_worker.py create mode 100644 flashdreams/tests/test_runtime_worker_gpu.py rename integrations/lingbot/lingbot/demo/{cli.py => app.py} (88%) create mode 100644 integrations/lingbot/lingbot/model_session.py create mode 100644 integrations/lingbot/lingbot/output_targets.py rename integrations/omnidreams/omnidreams/demo/{cli.py => app.py} (85%) create mode 100644 integrations/omnidreams/omnidreams/demo/web/adapter.js create mode 100644 integrations/omnidreams/omnidreams/model_session.py create mode 100644 integrations/omnidreams/omnidreams/output_targets.py delete mode 100644 integrations/omnidreams/omnidreams/webrtc/__init__.py delete mode 100644 integrations/omnidreams/omnidreams/webrtc/server.py delete mode 100644 integrations/omnidreams/omnidreams/webrtc/session.py delete mode 100644 integrations/omnidreams/omnidreams/webrtc/web/adapter.js delete mode 100644 integrations/omnidreams/tests/test_webrtc_runtime.py delete mode 100644 integrations/omnidreams/tests/test_webrtc_server_routes.py 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()}" From a7ea381092f95fd2434b22e3b142397bd3b3a6eb Mon Sep 17 00:00:00 2001 From: jarcherNV Date: Sat, 8 Aug 2026 16:33:44 -0700 Subject: [PATCH 15/19] Restore shared WebRTC manager input helpers (#432) * Restore shared WebRTC manager input helpers * fix(webrtc): restore session completion and lingbot adapter Treat completed inference sessions as terminal in the shared WebRTC manager, restore LingBot's session-branch adapter against the current model-session core, and update affected WebRTC tests for the current result and manager APIs. --- flashdreams/flashdreams/runtime/worker.py | 15 +- .../flashdreams/serving/webrtc/manager.py | 348 +++++++++++++++++- .../flashdreams/serving/webrtc/runtime.py | 1 + flashdreams/tests/test_webrtc_manager.py | 129 ++++++- flashdreams/tests/test_webrtc_serving.py | 16 +- .../lingbot/lingbot/webrtc/session.py | 126 +++++-- .../lingbot/tests/test_webrtc_runtime.py | 169 +++------ .../tests/test_webrtc_session_branch.py | 137 ++++++- 8 files changed, 730 insertions(+), 211 deletions(-) diff --git a/flashdreams/flashdreams/runtime/worker.py b/flashdreams/flashdreams/runtime/worker.py index 6e0c17eb2..af4576c09 100644 --- a/flashdreams/flashdreams/runtime/worker.py +++ b/flashdreams/flashdreams/runtime/worker.py @@ -7,7 +7,7 @@ import asyncio from concurrent.futures import ThreadPoolExecutor -from typing import Any, Callable, TypeVar +from typing import Any, Callable, TypeVar, cast import torch @@ -63,6 +63,19 @@ async def call( future.add_done_callback(_consume_exception) raise + def call_blocking( + self, + func: Callable[..., _T], + /, + *args: Any, + **kwargs: Any, + ) -> _T: + """Run one callable from synchronous code on the owned worker thread.""" + if not self._accepting: + raise RuntimeError("runtime worker is closed") + future = self._executor.submit(_invoke, func, args, kwargs) + return cast(_T, future.result()) + async def close(self) -> None: """Drain submitted work and stop accepting lifecycle calls.""" async with self._close_lock: diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index 30fcbc56c..a7c6c87a4 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -10,6 +10,7 @@ import inspect import json from collections import deque +from collections.abc import Mapping from collections.abc import Set as AbstractSet from dataclasses import dataclass, field, replace from typing import Any, Generic, TypeVar @@ -22,9 +23,18 @@ ) from loguru import logger -from flashdreams.runtime.inputs import TimeWindow +from flashdreams.runtime.inputs import ( + InferenceInput, + TimeWindow, + UserInputEvent, + UserInputs, +) from flashdreams.runtime.types import StepRequest, StepResult -from flashdreams.serving.realtime.input import KeyboardResampler +from flashdreams.serving.realtime.input import ( + DEFAULT_SUPPORTED_KEYS, + KeyboardResampler, + normalize_key, +) from flashdreams.serving.webrtc.encoders import ( DefaultRTCEncoder, VideoEncoder, @@ -73,6 +83,10 @@ _RuntimeConfigT = TypeVar("_RuntimeConfigT", bound=WebRTCRuntimeConfig) +class _InferenceSessionExhausted(RuntimeError): + """Raised when an ``InferenceSession`` reports normal completion.""" + + def _summarize_sdp_candidates(sdp: str) -> str: candidates = [ line.removeprefix("a=candidate:") @@ -107,18 +121,22 @@ def _summarize_sdp_candidates(sdp: str) -> str: ) -def _stat_float(stats: dict[str, float], name: str, default: float = 0.0) -> float: +def _stat_float( + stats: Mapping[str, float | int], name: str, default: float = 0.0 +) -> float: value = stats.get(name) if value is None: return default return float(value) -def _stat_ms(stats: dict[str, float], name: str, default_ms: float = 0.0) -> float: +def _stat_ms( + stats: Mapping[str, float | int], name: str, default_ms: float = 0.0 +) -> float: return _stat_float(stats, name, default_ms / 1e3) * 1e3 -def _stat_int(stats: dict[str, float], name: str) -> int: +def _stat_int(stats: Mapping[str, float | int], name: str) -> int: return int(round(_stat_float(stats, name))) @@ -239,14 +257,24 @@ def _make_resampler(self, *, start_v: float) -> KeyboardResampler: def _make_resampler_at_fps( self, *, start_v: float, fps: float ) -> KeyboardResampler: - if self.supported_control_keys is None: + supported_control_keys = self._effective_supported_control_keys() + if supported_control_keys is None: return KeyboardResampler(fps=fps, start_v=start_v) return KeyboardResampler( fps=fps, start_v=start_v, - supported_keys=self.supported_control_keys, + supported_keys=supported_control_keys, ) + def _effective_supported_control_keys(self) -> frozenset[str] | None: + supported_control_keys = self.supported_control_keys + if supported_control_keys is not None: + return frozenset(supported_control_keys) + legacy_supported_keys = getattr(self, "_resampler_supported_keys", None) + if legacy_supported_keys is None: + return None + return frozenset(legacy_supported_keys) + @staticmethod def _positive_int_runtime_value(value: Any, *, label: str) -> int: try: @@ -425,9 +453,7 @@ async def _handle_event_message( return False if channel is not None: active_event_id = event_payload.get("event_id") - ack_event_id = ( - None if active_event_id is None else str(active_event_id) - ) + ack_event_id = None if active_event_id is None else str(active_event_id) self._send_json( channel, make_event_ack_payload( @@ -467,6 +493,273 @@ async def _handle_event_message( ) return True + @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 use the same monotonic clock as the realtime resampler so + chunk ``TimeWindow`` filtering and raw data-channel events agree. + """ + 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._effective_supported_control_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.""" + 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, + ) -> StepResult: + """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 _InferenceSessionExhausted() + if request.step_index == 0 and not managed_session.session_input_state_advanced: + 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) + if not isinstance(result, StepResult): + raise TypeError( + "Inference session steps must produce StepResult, got " + f"{type(result).__name__}." + ) + self._prune_consumed_user_events(managed_session, before_s=window.start_s) + managed_session.session_steps_completed += 1 + return result + + 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, + ) + def has_active_session(self) -> bool: return self._active_session is not None and not self._active_session.closed @@ -885,7 +1178,7 @@ async def _generation_worker( chunk_start_v = resampler.next_chunk_start_v segments, frame_times = resampler.sample_chunk(input_num_frames) chunk_end_v = resampler.next_chunk_start_v - request = replace( + segment_request = replace( request, user_input_window=TimeWindow( start_s=chunk_start_v, @@ -901,15 +1194,32 @@ async def _generation_worker( managed_session.pending_action_arrivals.popleft() ) try: - 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}." + 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.step( + request=segment_request, + segments=segments, + frame_times=frame_times, ) + if result.step_index != segment_request.step_index: + raise RuntimeError( + "Runtime result step does not match its request: " + f"requested {segment_request.step_index}, " + f"got {result.step_index}." + ) + except _InferenceSessionExhausted: + logger.info( + "Inference session reported completion; closing WebRTC session." + ) + await self.close_active_session() + return except Exception as exc: logger.exception("Chunk generation failed.") channel = managed_session.control_channel diff --git a/flashdreams/flashdreams/serving/webrtc/runtime.py b/flashdreams/flashdreams/serving/webrtc/runtime.py index d80141f6b..70c16fb38 100644 --- a/flashdreams/flashdreams/serving/webrtc/runtime.py +++ b/flashdreams/flashdreams/serving/webrtc/runtime.py @@ -36,6 +36,7 @@ class WebRTCControlSignal(IntEnum): ACTION_STEP = 2 CLOSE = 3 EVENT = 4 + SESSION_STEP = 5 EXIT = 99 diff --git a/flashdreams/tests/test_webrtc_manager.py b/flashdreams/tests/test_webrtc_manager.py index e6089e869..8e4839e5e 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -6,7 +6,7 @@ import asyncio import json from types import SimpleNamespace -from typing import Any +from typing import Any, cast import pytest import torch @@ -147,6 +147,15 @@ def _make_manager( ) +def test_drives_inference_session_detects_session_runtime() -> None: + class _SessionRuntime: + async def start_inference_session(self) -> object: + return object() + + assert BaseWebRTCSessionManager._drives_inference_session(_SessionRuntime()) + assert not BaseWebRTCSessionManager._drives_inference_session(object()) + + def test_runtime_frame_timing_contract() -> None: class _Runtime: def peek_input_fps(self) -> float: @@ -295,9 +304,9 @@ def test_record_user_event_does_not_evict_unrelated_event_for_release( payload={"key": "w"}, ) - assert [(event.event_type, dict(event.payload)) for event in managed.user_events] == [ - ("text_event", {"event_id": "storm"}) - ] + assert [ + (event.event_type, dict(event.payload)) for event in managed.user_events + ] == [("text_event", {"event_id": "storm"})] assert len(managed.user_events) == 1 assert set(managed.coalesced_release_events) == {"w"} assert managed.coalesced_release_events["w"].timestamp_s == pytest.approx(0.2) @@ -337,9 +346,9 @@ def test_record_user_event_ignores_unsupported_key_events( payload={"key": "z"}, ) - assert [(event.event_type, dict(event.payload)) for event in managed.user_events] == [ - ("text_event", {"event_id": "storm"}) - ] + assert [ + (event.event_type, dict(event.payload)) for event in managed.user_events + ] == [("text_event", {"event_id": "storm"})] def test_catch_up_input_clock_advances_session_input_state() -> None: @@ -358,10 +367,7 @@ def canonicalize( del source_schema self.windows.append((window.start_s, window.end_s)) self.event_batches.append( - [ - (event.timestamp_s, event.event_type) - for event in user_inputs.events - ] + [(event.timestamp_s, event.event_type) for event in user_inputs.events] ) return object() @@ -475,9 +481,9 @@ async def test_action_keyup_updates_state_when_user_event_queue_full( raw_message='{"type":"action","action":{"event":"keyup","key":"w"}}', ) - assert [(event.event_type, dict(event.payload)) for event in managed.user_events] == [ - ("key_up", {"key": "w"}) - ] + assert [ + (event.event_type, dict(event.payload)) for event in managed.user_events + ] == [("key_up", {"key": "w"})] assert managed.first_action_received.is_set() assert len(managed.pending_action_arrivals) == 1 assert len(resampler.edges) == 1 @@ -485,6 +491,57 @@ async def test_action_keyup_updates_state_when_user_event_queue_full( assert channel.messages == [] +@pytest.mark.asyncio +async def test_session_event_message_validates_records_and_activates() -> None: + class _Runtime: + def __init__(self) -> None: + self.validate_calls: list[tuple[str, dict[str, Any]]] = [] + + def validate_user_event( + self, *, event_type: str, payload: dict[str, Any] + ) -> dict[str, Any]: + self.validate_calls.append((event_type, dict(payload))) + return { + "event_id": payload["event_id"], + "state": payload["state"], + "validated": True, + } + + runtime = _Runtime() + manager = _make_manager(_BaseTestManager, runtime) + managed, _video_track, _peer, channel = _managed_session(runtime) + managed.inference_session = object() + managed.first_action_received.clear() + + await manager._handle_datachannel_message( + managed_session=managed, + raw_message=json.dumps( + {"type": "event", "event_id": "storm", "state": "trigger"} + ), + ) + + assert runtime.validate_calls == [ + ("text_event", {"event_id": "storm", "state": "trigger"}) + ] + assert [ + (event.event_type, dict(event.payload)) for event in managed.user_events + ] == [ + ( + "text_event", + {"event_id": "storm", "state": "trigger", "validated": True}, + ) + ] + assert managed.first_action_received.is_set() + assert [json.loads(message) for message in channel.messages] == [ + { + "type": "event_ack", + "event_id": "storm", + "state": "trigger", + "active_event_id": "storm", + } + ] + + @pytest.mark.asyncio async def test_generation_worker_closes_session_when_flag_set() -> None: class _ClosingRuntime: @@ -560,6 +617,50 @@ async def step( assert not peer.closed +@pytest.mark.asyncio +async def test_generation_worker_closes_completed_inference_session_without_retry() -> ( + None +): + class _CompletedSession: + def __init__(self) -> None: + self.calls = 0 + + def next_step_request(self) -> StepRequest | None: + self.calls += 1 + return None + + def step(self, inputs: Any) -> StepResult: + del inputs + raise AssertionError("completed sessions must not be stepped") + + class _SessionRuntime: + def __init__(self) -> None: + self.session = _CompletedSession() + + def next_step_request(self) -> StepRequest: + return _step_request() + + runtime = _SessionRuntime() + manager = _make_manager(_BaseTestManager, runtime) + managed, video_track, peer, channel = _managed_session(runtime) + managed.inference_session = runtime.session + resampler = cast(_FakeResampler, managed.resampler) + resampler.dt = 1.0 / 30.0 + resampler.next_chunk_start_v = asyncio.get_running_loop().time() + manager._active_session = managed + + task = asyncio.create_task(manager._generation_worker(managed_session=managed)) + managed.generation_task = task + await asyncio.wait_for(task, timeout=5.0) + + assert runtime.session.calls == 1 + assert not manager.has_active_session() + assert managed.closed + assert video_track.closed + assert peer.closed + assert channel.messages == [] + + @pytest.mark.asyncio async def test_chunk_done_payload_includes_model_and_extra() -> None: class _OneChunkRuntime: diff --git a/flashdreams/tests/test_webrtc_serving.py b/flashdreams/tests/test_webrtc_serving.py index e376f7eb2..21d789fe0 100644 --- a/flashdreams/tests/test_webrtc_serving.py +++ b/flashdreams/tests/test_webrtc_serving.py @@ -20,11 +20,11 @@ KeyboardResampler, KeyboardState, ) -from flashdreams.serving.webrtc.media import tensor_chunk_to_rgb_frames from flashdreams.serving.webrtc.manager import ( BaseWebRTCSessionManager, ManagedWebRTCSession, ) +from flashdreams.serving.webrtc.media import tensor_chunk_to_rgb_frames from flashdreams.serving.webrtc.messages import ( make_chunk_done_payload, make_error_payload, @@ -140,7 +140,7 @@ def send(self, payload: str) -> None: self.messages.append(decoded) -class _Manager(BaseWebRTCSessionManager[Any, object]): +class _Manager(BaseWebRTCSessionManager[Any, Any]): def _model_name(self) -> str: return "fake" @@ -173,7 +173,9 @@ async def trigger_event( return {"active_event_id": event_id} runtime = _FakeRuntime() - manager = _Manager(runtime=runtime, runtime_config=object(), fps=30) + manager = _Manager( + runtime=runtime, runtime_config=object(), fps=30, identity="fake" + ) managed_session, channel = _managed_session_with_channel(runtime) await manager._handle_datachannel_message( @@ -211,7 +213,9 @@ async def trigger_event( } runtime = _FakeRuntime() - manager = _Manager(runtime=runtime, runtime_config=object(), fps=30) + manager = _Manager( + runtime=runtime, runtime_config=object(), fps=30, identity="fake" + ) managed_session, channel = _managed_session_with_channel(runtime) await manager._handle_datachannel_message( @@ -245,7 +249,9 @@ async def trigger_event( return {} runtime = _FakeRuntime() - manager = _Manager(runtime=runtime, runtime_config=object(), fps=30) + manager = _Manager( + runtime=runtime, runtime_config=object(), fps=30, identity="fake" + ) managed_session, channel = _managed_session_with_channel(runtime) await manager._handle_datachannel_message( diff --git a/integrations/lingbot/lingbot/webrtc/session.py b/integrations/lingbot/lingbot/webrtc/session.py index a45c3b2e1..0f95b6618 100644 --- a/integrations/lingbot/lingbot/webrtc/session.py +++ b/integrations/lingbot/lingbot/webrtc/session.py @@ -39,7 +39,13 @@ from flashdreams.core.io.disk import default_flashdreams_cache_dir from flashdreams.infra.config import derive_config from flashdreams.infra.video_output import VideoOutputStream -from flashdreams.runtime import StepResult +from flashdreams.runtime.canonical import InputCanonicalizer +from flashdreams.runtime.inputs import ( + InferenceInput, + UserInputCapability, + UserInputSchema, +) +from flashdreams.runtime.types import StepRequest, StepResult from flashdreams.serving.webrtc.controls import ( CameraPoseIntegrator, PoseSegment, @@ -54,13 +60,7 @@ ThreadAffineDistributedWebRTCRuntime, ) from flashdreams.serving.webrtc.server import SessionBusyError -from flashdreams.runtime.canonical import InputCanonicalizer -from flashdreams.runtime.inputs import ( - InferenceInput, - UserInputCapability, - UserInputSchema, -) -from flashdreams.runtime.types import StepRequest, StepResult +from lingbot.encoder.utils import preprocess_example_poses from lingbot.input_mapping import ( FIELD_CAMERA_INTRINSICS, FIELD_CAMERA_TRAJECTORY, @@ -68,7 +68,6 @@ LingbotInputMapping, TextEventSelection, ) -from lingbot.encoder.utils import preprocess_example_poses from lingbot.model_session import LingbotModelSessionCore _INTRINSICS_REFERENCE_HEIGHT = 480 @@ -632,6 +631,90 @@ async def trigger_event( 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.""" + if self._base_intrinsics is None: + self._input_mapping = None + self._input_canonicalizer = None + return + self._input_canonicalizer = InputCanonicalizer( + [KeyboardToCameraCommand(), TextEventSelection()] + ) + 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 mapped-input chunk for the session branch.""" + if self._model_session is None: + raise LingbotRuntimeError("Runtime is not initialized.") + step_index = self._model_session.step_index + num_frames = self._model_session.next_num_frames() + return StepRequest( + step_index=step_index, + metadata={ + "input_frame_count": num_frames, + "num_frames": num_frames, + "frame_start": step_index * num_frames, + }, + ) + + def _step_blocking(self, inputs: InferenceInput) -> StepResult: + """Run one mapped step from synchronous ``InferenceSession`` code.""" + if self._closed: + raise LingbotRuntimeError("Session is closed.") + with self._sync_step_lock: + if self._closed: + raise LingbotRuntimeError("Session is closed.") + return self._worker.call_blocking(self._step_sync_all_ranks, inputs) + # 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 @@ -672,6 +755,10 @@ def _steady_output_frame_count(self) -> int: self._pipeline.get_num_output_frames(self._STEADY_STATE_AR_PROBE_INDEX) ) + @distributed_op(WebRTCControlSignal.SESSION_STEP) + def _step_sync_all_ranks(self, inputs: InferenceInput) -> StepResult: + return self._step_sync(inputs) + @distributed_op(WebRTCControlSignal.EVENT) def _trigger_event_sync_all_ranks( self, @@ -1044,16 +1131,14 @@ def _generate_from_camera_inputs( poses: torch.Tensor, intrinsics: torch.Tensor, num_frames: int, - ) -> VideoStepResult: + ) -> StepResult: """Generate one chunk from an already-resolved camera trajectory. Shared by the segment path and the mapped-input session path so both reach the model through identical conditioning. """ - 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.") - if self._device is None: - raise LingbotRuntimeError("Runtime device is not initialized.") from lingbot.encoder.camctrl import CamCtrlInput # noqa: PLC0415 @@ -1073,11 +1158,9 @@ def _generate_from_camera_inputs( def _step_sync(self, inputs: InferenceInput) -> StepResult: """Generate one chunk from mapped model inputs.""" - 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.") - num_frames = int( - self._pipeline.get_num_output_frames(self.autoregressive_index) - ) + num_frames = self._model_session.next_num_frames() self._apply_conditioning_update_sync(inputs) poses = _require_camera_tensor( inputs, FIELD_CAMERA_TRAJECTORY, expected_shape=(num_frames, 4, 4) @@ -1085,18 +1168,11 @@ def _step_sync(self, inputs: InferenceInput) -> StepResult: intrinsics = _require_camera_tensor( inputs, FIELD_CAMERA_INTRINSICS, expected_shape=(num_frames, 4) ) - step_index = self.autoregressive_index - result = self._generate_from_camera_inputs( + return self._generate_from_camera_inputs( poses=poses, intrinsics=intrinsics, num_frames=num_frames, ) - return StepResult( - step_index=step_index, - output=result, - frame_count=result.num_frames, - metrics=result.stats or {}, - ) def _apply_conditioning_update_sync(self, inputs: InferenceInput) -> None: """Apply a text-event prompt swap requested by the mapping.""" diff --git a/integrations/lingbot/tests/test_webrtc_runtime.py b/integrations/lingbot/tests/test_webrtc_runtime.py index 2c2f2ac7f..2b70dc366 100644 --- a/integrations/lingbot/tests/test_webrtc_runtime.py +++ b/integrations/lingbot/tests/test_webrtc_runtime.py @@ -22,6 +22,11 @@ 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 ( @@ -31,6 +36,8 @@ from flashdreams.infra.video_output import VideoOutputStream from flashdreams.runtime import StepRequest, StepResult +from flashdreams.runtime.canonical import InputCanonicalizer +from flashdreams.runtime.inputs import InferenceInput from flashdreams.serving.webrtc import runtime as webrtc_runtime from flashdreams.serving.webrtc.manager import ( BaseWebRTCSessionManager, @@ -595,117 +602,36 @@ def replace_text_embeddings( ) -> None: self.calls.append((cache, text_embeddings)) - 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: + class _FakeDiffusionModel: def __init__(self) -> None: self.transformer = _FakeTransformer() - 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: + class _FakePipeline: def __init__(self) -> None: self.diffusion_model = _FakeDiffusionModel() - async def trigger_event( - self, *, event_id: str, state: str - ) -> dict[str, object]: - del event_id, state - self.calls += 1 - return {} - - monkeypatch.setattr(session, "LingbotInferenceRuntime", _fake_runtime_factory) - manager = create_lingbot_webrtc_session_manager( - runtime_config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0) + runtime = session.LingbotInferenceRuntime( + config=LingbotRuntimeConfig( + device="cpu", + warmup_chunks=0, + text_events=(), + ) ) - 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_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() + _attach_model_session(runtime, runtime._pipeline, 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, + } + + runtime._apply_conditioning_update_sync( + InferenceInput(global_conditioning={"prompt": "a glowing portal opens"}) ) transformer = runtime._pipeline.diffusion_model.transformer @@ -981,22 +907,21 @@ def next_step_request(self) -> StepRequest: step_index = self._runtime.step_index return StepRequest( step_index=step_index, - metadata={"num_frames": 1, "frame_start": step_index}, + metadata={ + "input_frame_count": 1, + "num_frames": 1, + "frame_start": step_index, + }, ) def step(self, inputs: InferenceInput) -> StepResult: chunk_index = self._runtime.step_index self._runtime.step_index += 1 self._runtime.generated_inputs.append(inputs) - return StepResult( + return StepResult.from_video_chunk( step_index=chunk_index, - output=VideoStepResult( - chunk_index=chunk_index, - num_frames=1, - video_chunk=torch.zeros((1, 1, 1, 3, 2, 2), dtype=torch.uint8), - stats=None, - ), - frame_count=1, + video_chunk=torch.zeros((1, 1, 1, 3, 2, 2), dtype=torch.uint8), + layout="bvtchw", ) class _FakeRuntime: @@ -1037,24 +962,12 @@ def peek_steady_output_num_frames(self) -> int: def next_step_request(self) -> StepRequest: return StepRequest( - step_index=len(self.generated_segments), + step_index=self.step_index, 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 start_inference_session(self) -> _FakeInferenceSession: + return _FakeInferenceSession(self) async def close(self) -> None: self.close_calls += 1 diff --git a/integrations/lingbot/tests/test_webrtc_session_branch.py b/integrations/lingbot/tests/test_webrtc_session_branch.py index 64e1728e5..8cf8e6cda 100644 --- a/integrations/lingbot/tests/test_webrtc_session_branch.py +++ b/integrations/lingbot/tests/test_webrtc_session_branch.py @@ -26,9 +26,8 @@ ) from lingbot.webrtc.session import LINGBOT_WEBRTC_SOURCE_SCHEMA -from flashdreams.infra.video_output import VideoStepResult -from flashdreams.runtime.inputs import InferenceInput, TimeWindow from flashdreams.runtime.canonical import InputCanonicalizer +from flashdreams.runtime.inputs import InferenceInput, TimeWindow from flashdreams.runtime.types import StepRequest, StepResult from flashdreams.serving.realtime.input import KeyboardResampler from flashdreams.serving.webrtc.controls import CameraPoseIntegrator @@ -62,6 +61,7 @@ def next_step_request(self) -> StepRequest: return StepRequest( step_index=self._index, metadata={ + "input_frame_count": _NUM_FRAMES, "num_frames": _NUM_FRAMES, "frame_start": self._index * _NUM_FRAMES, }, @@ -71,16 +71,10 @@ def step(self, inputs: InferenceInput) -> StepResult: self.steps.append(inputs) index = self._index self._index += 1 - return StepResult( + return StepResult.from_video_chunk( step_index=index, - output=VideoStepResult( - chunk_index=index, - video_chunk=torch.zeros(_NUM_FRAMES, 3, 4, 4), - layout="tchw", - num_frames=_NUM_FRAMES, - stats={}, - ), - frame_count=_NUM_FRAMES, + video_chunk=torch.zeros(_NUM_FRAMES, 3, 4, 4), + layout="tchw", ) @@ -137,24 +131,41 @@ def send(self, payload: str) -> None: self.messages.append(decoded) +class _FakeCloseable: + async def close(self) -> None: + return + + +class _FakeVideoEncoder: + fps = _FPS + backend = "fake" + prefers_codec: str | None = None + + def close(self) -> None: + return + + def _managed_session(runtime: _FakeRuntime) -> ManagedWebRTCSession: return ManagedWebRTCSession( runtime=runtime, - video_track=None, - video_encoder=None, - peer_connection=None, + video_track=_FakeCloseable(), # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), # ty:ignore[invalid-argument-type] + peer_connection=_FakeCloseable(), resampler=KeyboardResampler(fps=_FPS, start_v=0.0), inference_session=runtime.session, ) def _manager(runtime: _FakeRuntime) -> _Manager: - return _Manager(runtime=runtime, runtime_config=_FakeRuntimeConfig(), fps=_FPS) + return _Manager( + runtime=runtime, + runtime_config=_FakeRuntimeConfig(), + fps=_FPS, + identity="fake", + ) -def _reference_poses( - edges: list[tuple[float, str, str]], *, chunks: int -) -> np.ndarray: +def _reference_poses(edges: list[tuple[float, str, str]], *, chunks: int) -> np.ndarray: resampler = KeyboardResampler(fps=_FPS, start_v=0.0) integrator = CameraPoseIntegrator() for timestamp_s, event, key in edges: @@ -388,7 +399,95 @@ def test_real_lingbot_runtime_selects_the_session_branch() -> None: runtime = LingbotInferenceRuntime(config=LingbotRuntimeConfig(device="cpu")) assert BaseWebRTCSessionManager._drives_inference_session(runtime) is True - assert callable(runtime.generate_chunk) + assert callable(runtime.step) + assert callable(runtime.start_inference_session) + + +def test_real_lingbot_inference_session_steps_on_runtime_worker() -> None: + import asyncio + + from lingbot.model_session import LingbotModelSessionCore + from lingbot.webrtc.session import LingbotInferenceRuntime, LingbotRuntimeConfig + + from flashdreams.infra.video_output import VideoOutputStream + + class _FakeTransformer: + pass + + class _FakeDiffusionModel: + def __init__(self) -> None: + self.transformer = _FakeTransformer() + + class _FakePipeline: + def __init__(self) -> None: + self.diffusion_model = _FakeDiffusionModel() + self.generated_inputs: list[object] = [] + + def get_num_output_frames(self, autoregressive_index: int) -> int: + del autoregressive_index + return _NUM_FRAMES + + def generate( + self, + *, + autoregressive_index: int, + cache: object, + input: object, + ) -> torch.Tensor: + del autoregressive_index, cache + self.generated_inputs.append(input) + return torch.zeros((_NUM_FRAMES, 3, 4, 4), dtype=torch.uint8) + + def finalize( + self, *, autoregressive_index: int, cache: object + ) -> dict[str, float]: + del autoregressive_index, cache + return {"decode_s": 0.1} + + runtime = LingbotInferenceRuntime( + config=LingbotRuntimeConfig(device="cpu", warmup_chunks=0, text_events=()) + ) + pipeline = _FakePipeline() + runtime._pipeline = pipeline + runtime._model_session = LingbotModelSessionCore( + pipeline=pipeline, + output_stream_factory=lambda: VideoOutputStream( + postprocess_stream=None, + output_layout="tchw", + ), + ) + runtime._model_session._cache = object() + runtime._input_canonicalizer = InputCanonicalizer( + [KeyboardToCameraCommand(), TextEventSelection()] + ) + runtime._input_mapping = LingbotInputMapping( + fps=_FPS, + base_intrinsics=_BASE_INTRINSICS, + world_scale=1.0, + text_event_prompts={}, + ) + + try: + inference_session = asyncio.run(runtime.start_inference_session()) + request = inference_session.next_step_request() + result = inference_session.step( + InferenceInput( + step={ + FIELD_CAMERA_TRAJECTORY: torch.eye(4).repeat(_NUM_FRAMES, 1, 1), + FIELD_CAMERA_INTRINSICS: _BASE_INTRINSICS.repeat(_NUM_FRAMES, 1), + }, + ) + ) + finally: + asyncio.run(runtime.close()) + + assert request is not None + assert request.step_index == 0 + assert request.metadata["input_frame_count"] == _NUM_FRAMES + assert result.step_index == 0 + assert result.frame_count == _NUM_FRAMES + assert result.metrics["decode_s"] == pytest.approx(0.1) + assert len(pipeline.generated_inputs) == 1 def test_session_start_requires_an_initialized_rollout() -> None: From 6239fdcf0c96c20114b40682bcc3729bde1922ea Mon Sep 17 00:00:00 2001 From: jarcherNV Date: Sun, 9 Aug 2026 23:55:55 -0700 Subject: [PATCH 16/19] unify demo runtime through OmniDreams migration (#433) unify demo runtime through OmniDreams migration Add the shared demo runtime stack for replay, WebRTC, input providers, output sinks, timing, warmup, metrics, and error handling. Migrate OmniDreams null, precomputed MP4, Ludus MP4, and default WebRTC paths onto the shared runtime, add GPU CI coverage for those paths, and clean up the OmniDreams demo layout while keeping legacy fallback compatibility. --- .github/workflows/omnidreams-demo-runtime.yml | 328 ++++ flashdreams/flashdreams/runtime/__init__.py | 14 +- flashdreams/flashdreams/runtime/config.py | 8 + .../flashdreams/runtime/demo/__init__.py | 152 +- flashdreams/flashdreams/runtime/demo/app.py | 11 +- .../flashdreams/runtime/demo/drivers.py | 1011 +++++++++++++ flashdreams/flashdreams/runtime/demo/host.py | 174 +++ .../flashdreams/runtime/demo/outputs.py | 266 +++- .../flashdreams/runtime/demo/pipeline.py | 75 + .../flashdreams/runtime/demo/replay.py | 554 ++++++- .../flashdreams/runtime/demo/run_modes.py | 535 +++++++ .../runtime/demo/session_inputs.py | 184 +++ flashdreams/flashdreams/runtime/demo/spec.py | 17 +- .../flashdreams/runtime/demo/timing.py | 383 +++++ .../flashdreams/runtime/demo/validation.py | 200 +++ flashdreams/flashdreams/runtime/metrics.py | 224 ++- flashdreams/flashdreams/runtime/runner.py | 645 ++++++-- flashdreams/flashdreams/runtime/types.py | 93 +- flashdreams/flashdreams/runtime/worker.py | 91 +- .../flashdreams/serving/realtime/timing.py | 50 + .../flashdreams/serving/webrtc/encoders.py | 49 +- .../flashdreams/serving/webrtc/manager.py | 938 +++++++++++- .../flashdreams/serving/webrtc/media.py | 36 +- .../flashdreams/serving/webrtc/nvenc.py | 109 ++ .../flashdreams/serving/webrtc/runtime.py | 1 + .../flashdreams/serving/webrtc/server.py | 8 +- .../flashdreams/serving/webrtc/services.py | 1161 ++++++++++++++ .../flashdreams/serving/webrtc/warmup.py | 29 +- .../serving/webrtc/web/mock_ui_server.py | 16 +- .../serving/webrtc/web/request_session.html | 2 +- .../serving/webrtc/web/request_session.js | 9 +- flashdreams/tests/test_demo_runtime_host.py | 218 +++ .../tests/test_demo_runtime_output_sinks.py | 221 +++ .../test_demo_runtime_realtime_driver.py | 916 ++++++++++++ .../tests/test_demo_runtime_run_modes.py | 774 ++++++++++ flashdreams/tests/test_demo_runtime_timing.py | 237 +++ .../tests/test_demo_runtime_validation.py | 680 +++++++++ .../tests/test_demo_runtime_vertical_slice.py | 1328 +++++++++++++++++ flashdreams/tests/test_demo_runtime_warmup.py | 416 ++++++ flashdreams/tests/test_encoders.py | 21 +- .../tests/test_inference_runtime_api.py | 154 ++ .../tests/test_realtime_timing_metrics.py | 71 + flashdreams/tests/test_runtime_demo_api.py | 184 ++- flashdreams/tests/test_runtime_runner.py | 94 +- flashdreams/tests/test_runtime_worker.py | 35 +- flashdreams/tests/test_webrtc_manager.py | 291 +++- flashdreams/tests/test_webrtc_services.py | 751 ++++++++++ flashdreams/tests/test_webrtc_serving.py | 39 +- flashdreams/tests/test_webrtc_warmup.py | 154 ++ integrations/lingbot/tests/test_demo_api.py | 7 +- .../omnidreams/omnidreams/demo/README.md | 65 +- .../omnidreams/omnidreams/demo/__init__.py | 16 + .../omnidreams/omnidreams/demo/adapter.py | 225 ++- .../omnidreams/omnidreams/demo/app.py | 90 +- .../omnidreams/omnidreams/demo/providers.py | 692 +++++++++ .../omnidreams/omnidreams/demo/replay.py | 245 +-- .../omnidreams/omnidreams/demo/runtime.py | 353 +++++ .../omnidreams/omnidreams/demo/spec.py | 282 +++- .../demo/traces/ludus_forward_sweep_60s.json | 86 ++ .../omnidreams/demo/web/adapter.css | 16 + .../omnidreams/omnidreams/demo/web/adapter.js | 1 + .../omnidreams/omnidreams/demo/webrtc.py | 583 ++------ .../omnidreams/demo/webrtc_config.py | 86 ++ .../omnidreams/demo/webrtc_legacy.py | 718 +++++++++ integrations/omnidreams/pyproject.toml | 2 +- .../omnidreams/tests/test_demo_api.py | 1259 ++++++++++++++-- 66 files changed, 17642 insertions(+), 1041 deletions(-) create mode 100644 .github/workflows/omnidreams-demo-runtime.yml create mode 100644 flashdreams/flashdreams/runtime/demo/drivers.py create mode 100644 flashdreams/flashdreams/runtime/demo/host.py create mode 100644 flashdreams/flashdreams/runtime/demo/pipeline.py create mode 100644 flashdreams/flashdreams/runtime/demo/run_modes.py create mode 100644 flashdreams/flashdreams/runtime/demo/session_inputs.py create mode 100644 flashdreams/flashdreams/runtime/demo/timing.py create mode 100644 flashdreams/flashdreams/runtime/demo/validation.py create mode 100644 flashdreams/flashdreams/serving/webrtc/services.py create mode 100644 flashdreams/tests/test_demo_runtime_host.py create mode 100644 flashdreams/tests/test_demo_runtime_output_sinks.py create mode 100644 flashdreams/tests/test_demo_runtime_realtime_driver.py create mode 100644 flashdreams/tests/test_demo_runtime_run_modes.py create mode 100644 flashdreams/tests/test_demo_runtime_timing.py create mode 100644 flashdreams/tests/test_demo_runtime_validation.py create mode 100644 flashdreams/tests/test_demo_runtime_vertical_slice.py create mode 100644 flashdreams/tests/test_demo_runtime_warmup.py create mode 100644 flashdreams/tests/test_realtime_timing_metrics.py create mode 100644 flashdreams/tests/test_webrtc_services.py create mode 100644 flashdreams/tests/test_webrtc_warmup.py create mode 100644 integrations/omnidreams/omnidreams/demo/providers.py create mode 100644 integrations/omnidreams/omnidreams/demo/runtime.py create mode 100644 integrations/omnidreams/omnidreams/demo/traces/ludus_forward_sweep_60s.json create mode 100644 integrations/omnidreams/omnidreams/demo/web/adapter.css create mode 100644 integrations/omnidreams/omnidreams/demo/webrtc_config.py create mode 100644 integrations/omnidreams/omnidreams/demo/webrtc_legacy.py diff --git a/.github/workflows/omnidreams-demo-runtime.yml b/.github/workflows/omnidreams-demo-runtime.yml new file mode 100644 index 000000000..1f494fa83 --- /dev/null +++ b/.github/workflows/omnidreams-demo-runtime.yml @@ -0,0 +1,328 @@ +name: OmniDreams Demo Runtime + +on: + push: + branches: + - main + - "pull-request/[0-9]+" + paths: + - ".github/workflows/omnidreams-demo-runtime.yml" + - "pyproject.toml" + - "uv.lock" + - "flashdreams/pyproject.toml" + - "flashdreams/flashdreams/core/**" + - "flashdreams/flashdreams/infra/**" + - "flashdreams/flashdreams/runtime/**" + - "flashdreams/flashdreams/serving/**" + - "flashdreams/flashdreams/recipes/taehv/**" + - "flashdreams/flashdreams/recipes/wan/**" + - "integrations/omnidreams/**" + - "integrations/lingbot/**" + workflow_dispatch: + +permissions: + contents: read + +jobs: + demo-runtime: + name: null, precomputed MP4, and Ludus MP4 + runs-on: linux-amd64-gpu-rtxpro6000-latest-2 + timeout-minutes: 180 + defaults: + run: + shell: bash + container: + image: nvidia/cuda:13.2.1-cudnn-devel-ubuntu24.04 + options: --gpus all + env: + UV_PROJECT_ENVIRONMENT: /tmp/flashdreams-venv + UV_LINK_MODE: copy + UV_PYTHON: "3.10" + MAX_JOBS: 8 + ARTIFACT_DIR: artifacts/omnidreams_demo_runtime + NULL_BLOCKS: "10" + PRECOMPUTED_BLOCKS: "75" + LUDUS_BLOCKS: "76" + FPS: "30" + EXAMPLE_DATA_UUID: 239560dc-33d1-11ef-9720-00044bcbccac + LUDUS_SCENE_UUID: 0d404ff7-2b66-498c-b047-1ed8cded60d4 + LUDUS_TRACE: integrations/omnidreams/omnidreams/demo/traces/ludus_forward_sweep_60s.json + EXPECTED_WIDTH: "1280" + EXPECTED_HEIGHT: "704" + MIN_DURATION_SECONDS: "18" + MAX_DURATION_SECONDS: "22" + steps: + - name: Detect GPU architecture + id: gpu-arch + run: | + nvidia-smi + compute_cap=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | head -1 | tr -d '[:space:]') + arch=$(echo "${compute_cap}" | tr -d '.') + echo "arch=${arch}" >> "$GITHUB_OUTPUT" + echo "Detected GPU compute capability: ${compute_cap} -> sm_${arch}" + + - name: Checkout + uses: actions/checkout@v4 + + - name: Install system dependencies + run: | + apt-get update -qq + DEBIAN_FRONTEND=noninteractive apt-get install -y -qq --no-install-recommends \ + python3 python3-dev python3-venv \ + ffmpeg \ + gcc g++ ninja-build \ + libnccl-dev \ + curl git ca-certificates jq unzip + rm -rf /var/lib/apt/lists/* + + - name: Setup proxy cache + uses: nv-gha-runners/setup-proxy-cache@main + + - name: Setup uv + uses: astral-sh/setup-uv@v6 + with: + enable-cache: true + cache-suffix: "omnidreams-demo-runtime-sm${{ steps.gpu-arch.outputs.arch }}" + prune-cache: false + + - name: Install dependencies + env: + NVTE_CUDA_ARCHS: ${{ steps.gpu-arch.outputs.arch }} + run: | + uv venv --clear + uv sync --locked --extra dev + + - name: Verify GPU availability + run: nvidia-smi + + - name: Run OmniDreams demo modes + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + set -uo pipefail + + log_dir="${ARTIFACT_DIR}/logs" + output_dir="${ARTIFACT_DIR}/outputs" + summary="${ARTIFACT_DIR}/summary.md" + status_file="${ARTIFACT_DIR}/command-status.env" + mkdir -p "${log_dir}" "${output_dir}" + : > "${status_file}" + + odemo() { + uv run --no-sync --package flashdreams-omnidreams omnidreams-demo "$@" + } + + run_demo() { + local name="$1" + shift + local log="${log_dir}/${name}.log" + + { + printf '$' + printf ' %q' "$@" + printf '\n\n' + "$@" + } 2>&1 | tee "${log}" + + local rc="${PIPESTATUS[0]}" + echo "${name}=${rc}" >> "${status_file}" + echo "${name} exit code: ${rc}" | tee -a "${summary}" + return 0 + } + + { + echo "# OmniDreams Demo Runtime CI" + echo + echo "| Mode | Expected blocks | Output |" + echo "| --- | ---: | --- |" + echo "| null | ${NULL_BLOCKS} | none |" + echo "| precomputed MP4 | ${PRECOMPUTED_BLOCKS} | omnidreams-demo-precomputed-20s.mp4 |" + echo "| Ludus MP4 | ${LUDUS_BLOCKS} | omnidreams-demo-ludus-20s.mp4 |" + echo + echo "## Command Status" + } > "${summary}" + + run_demo null \ + odemo replay \ + --output-mode null \ + --device cuda:0 \ + --total-blocks "${NULL_BLOCKS}" + + run_demo precomputed-mp4 \ + odemo replay \ + --device cuda:0 \ + --example-data \ + --example-data-uuid "${EXAMPLE_DATA_UUID}" \ + --total-blocks "${PRECOMPUTED_BLOCKS}" \ + --fps "${FPS}" \ + --output "${output_dir}/omnidreams-demo-precomputed-20s.mp4" + + run_demo ludus-mp4 \ + odemo replay \ + --conditioning-mode ludus-scene-driving \ + --keyboard-trace "${LUDUS_TRACE}" \ + --device cuda:0 \ + --scene-uuid "${LUDUS_SCENE_UUID}" \ + --seed 42 \ + --total-blocks "${LUDUS_BLOCKS}" \ + --output "${output_dir}/omnidreams-demo-ludus-20s.mp4" + + - name: Validate OmniDreams demo artifacts + run: | + set -euo pipefail + + log_dir="${ARTIFACT_DIR}/logs" + output_dir="${ARTIFACT_DIR}/outputs" + probe_dir="${ARTIFACT_DIR}/ffprobe" + summary="${ARTIFACT_DIR}/summary.md" + status_file="${ARTIFACT_DIR}/command-status.env" + mkdir -p "${probe_dir}" + + status_of() { + awk -F= -v name="$1" '$1 == name { print $2 }' "${status_file}" + } + + assert_exit_zero() { + local name="$1" + local rc + rc="$(status_of "${name}")" + if [ "${rc}" != "0" ]; then + echo "${name} command failed with exit code ${rc}" >&2 + exit 1 + fi + } + + assert_clean_log() { + local log="$1" + if grep -En "ERROR|Traceback|Exception|status=failed|Run failed|failed run" "${log}"; then + echo "failure marker found in ${log}" >&2 + exit 1 + fi + } + + assert_log_contains() { + local log="$1" + local pattern="$2" + local label="$3" + if ! grep -Eq "${pattern}" "${log}"; then + echo "expected ${label} in ${log}" >&2 + exit 1 + fi + } + + validate_mp4() { + local mode="$1" + local mp4="$2" + local metadata="${probe_dir}/${mode}.json" + + if [ ! -s "${mp4}" ]; then + echo "expected non-empty MP4 at ${mp4}" >&2 + exit 1 + fi + + ffprobe \ + -v error \ + -select_streams v:0 \ + -show_entries stream=width,height,r_frame_rate,avg_frame_rate,nb_frames,duration:format=duration \ + -of json \ + "${mp4}" > "${metadata}" + + local stream_count width height duration + stream_count="$(jq '.streams | length' "${metadata}")" + width="$(jq -r '.streams[0].width // ""' "${metadata}")" + height="$(jq -r '.streams[0].height // ""' "${metadata}")" + duration="$(jq -r '.streams[0].duration // .format.duration // "0"' "${metadata}")" + + if [ "${stream_count}" -lt 1 ]; then + echo "ffprobe found no video stream in ${mp4}" >&2 + exit 1 + fi + + if [ "${width}" != "${EXPECTED_WIDTH}" ] || [ "${height}" != "${EXPECTED_HEIGHT}" ]; then + echo "unexpected ${mode} resolution ${width}x${height}; expected ${EXPECTED_WIDTH}x${EXPECTED_HEIGHT}" >&2 + exit 1 + fi + + awk \ + -v duration="${duration}" \ + -v min_duration="${MIN_DURATION_SECONDS}" \ + -v max_duration="${MAX_DURATION_SECONDS}" \ + 'BEGIN { + if ((duration + 0) < min_duration || (duration + 0) > max_duration) { + exit 1 + } + }' || { + echo "unexpected ${mode} duration ${duration}s; expected ${MIN_DURATION_SECONDS}-${MAX_DURATION_SECONDS}s" >&2 + exit 1 + } + } + + null_log="${log_dir}/null.log" + precomputed_log="${log_dir}/precomputed-mp4.log" + ludus_log="${log_dir}/ludus-mp4.log" + + assert_exit_zero null + assert_exit_zero precomputed-mp4 + assert_exit_zero ludus-mp4 + + assert_clean_log "${null_log}" + assert_clean_log "${precomputed_log}" + assert_clean_log "${ludus_log}" + + assert_log_contains "${null_log}" "AR 9 encode" "null final AR block" + assert_log_contains "${null_log}" "OmniDreams demo replay step 9 frames=" "null final replay step" + assert_log_contains "${null_log}" "Loaded OmniDreams demo HDMaps shape=.*views=1" "null precomputed HDMaps" + + assert_log_contains "${precomputed_log}" "AR 74 encode" "precomputed final AR block" + assert_log_contains "${precomputed_log}" "OmniDreams demo replay step 74 frames=" "precomputed final replay step" + assert_log_contains "${precomputed_log}" "Loaded OmniDreams demo HDMaps shape=.*views=1" "precomputed HDMaps" + + assert_log_contains "${ludus_log}" "AR 75 encode" "Ludus final AR block" + assert_log_contains "${ludus_log}" "OmniDreams demo replay step 75 frames=" "Ludus final replay step" + assert_log_contains "${ludus_log}" "ludus_backend=cuda" "Ludus CUDA backend" + assert_log_contains "${ludus_log}" "trace_events=[1-9][0-9]*" "nonzero Ludus trace events" + + validate_mp4 precomputed-mp4 "${output_dir}/omnidreams-demo-precomputed-20s.mp4" + validate_mp4 ludus-mp4 "${output_dir}/omnidreams-demo-ludus-20s.mp4" + + { + echo + echo "## Validation" + echo + echo "- All commands exited zero." + echo "- Logs contained expected final AR blocks and provider markers." + echo "- MP4 outputs were non-empty and passed ffprobe stream checks." + } >> "${summary}" + + - name: Trim uv cache for upload + if: always() + run: | + cache_dir="${UV_CACHE_DIR:-/github/home/.cache/uv}" + echo "=== Cache size before trim ===" + du -sh "${cache_dir}" || true + du -sh "${cache_dir}"/*/ 2>/dev/null || true + + rm -rf "${cache_dir}/wheels-v6" + rm -rf "${cache_dir}/archive-v0" + + find "${cache_dir}/git-v0/checkouts" \ + \( -name "build" -o -name "*.egg-info" -o -name "__pycache__" \) \ + -type d -exec rm -rf {} + 2>/dev/null || true + + rm -rf "${cache_dir}/sdists-v9/editable" + + echo "" + echo "=== Cache size after trim ===" + du -sh "${cache_dir}" || true + du -sh "${cache_dir}"/*/ 2>/dev/null || true + echo "" + echo "=== Cached built wheels (sdists-v9) ===" + find "${cache_dir}/sdists-v9" -name "*.whl" -exec ls -lh {} \; 2>/dev/null || true + + - name: Upload demo runtime artifacts + if: always() + uses: actions/upload-artifact@v4 + with: + name: omnidreams-demo-runtime + path: ${{ env.ARTIFACT_DIR }} + if-no-files-found: ignore diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py index fb6eb4b05..6e205823c 100644 --- a/flashdreams/flashdreams/runtime/__init__.py +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -52,14 +52,20 @@ from flashdreams.runtime.metrics import ( InMemoryMetricsRecorder, MetricsRecorder, + MetricsSnapshot, NullMetricsRecorder, RuntimeMetricSample, ) from flashdreams.runtime.output import NullOutputTarget, OutputArtifact, OutputTarget from flashdreams.runtime.runner import run_inference_session -from flashdreams.runtime.types import StepRequest, StepResult +from flashdreams.runtime.types import ( + StepRequest, + StepRequirements, + StepResult, + step_requirements_from_request, +) from flashdreams.runtime.video_output import Mp4VideoOutputTarget -from flashdreams.runtime.worker import ThreadAffineRuntimeWorker +from flashdreams.runtime.worker import ModelExecutionWorker, ThreadAffineRuntimeWorker __all__ = [ "CanonicalInputs", @@ -90,7 +96,9 @@ "KeyboardToDriverCommand", "MappingCompatibility", "MetricsRecorder", + "MetricsSnapshot", "ModelAdapter", + "ModelExecutionWorker", "Mp4VideoOutputTarget", "NullMetricsRecorder", "NullOutputTarget", @@ -100,10 +108,12 @@ "RuntimeMetricSample", "ScriptedModality", "StepRequest", + "StepRequirements", "StepResult", "TimeWindow", "ThreadAffineRuntimeWorker", "run_inference_session", + "step_requirements_from_request", "undeclared_inference_inputs", "UserInputCapability", "UserInputEvent", diff --git a/flashdreams/flashdreams/runtime/config.py b/flashdreams/flashdreams/runtime/config.py index 4b8752f13..f1c0c2a0c 100644 --- a/flashdreams/flashdreams/runtime/config.py +++ b/flashdreams/flashdreams/runtime/config.py @@ -61,6 +61,9 @@ class InferenceConfig: cache_policy: str | None = None """Optional cache policy selector; ``None`` leaves the choice to the adapter.""" + seed: int | None = None + """Optional seed used when resolving deterministic demo/runtime behavior.""" + runtime_options: Mapping[str, Any] = field(default_factory=dict) """Adapter/backend-specific runtime options.""" @@ -70,6 +73,11 @@ class InferenceConfig: def __post_init__(self) -> None: if not self.model_id.strip(): raise ValueError("InferenceConfig.model_id must be non-empty.") + if self.seed is not None: + if isinstance(self.seed, bool) or not isinstance(self.seed, int): + raise TypeError("InferenceConfig.seed must be an integer.") + if self.seed < 0: + raise ValueError("InferenceConfig.seed must be >= 0.") object.__setattr__( self, "runtime_options", freeze_mapping(self.runtime_options) ) diff --git a/flashdreams/flashdreams/runtime/demo/__init__.py b/flashdreams/flashdreams/runtime/demo/__init__.py index 7a0535556..60187237c 100644 --- a/flashdreams/flashdreams/runtime/demo/__init__.py +++ b/flashdreams/flashdreams/runtime/demo/__init__.py @@ -3,11 +3,70 @@ """Experimental shared demo API above the inference runtime API.""" -from flashdreams.runtime.demo.outputs import build_output_target -from flashdreams.runtime.demo.replay import run_replay_demo +from flashdreams.runtime.demo.drivers import ( + CLEANUP_TIMEOUT_S, + BatchSessionDriver, + DriverInvariantError, + RealtimeSessionDriver, + run_demo_session, + run_demo_session_async, + shielded_session_cleanup, + uncancel_current_task, +) +from flashdreams.runtime.demo.host import ( + ModelWarmupPlan, + RuntimeHost, + WarmupSessionInputs, +) +from flashdreams.runtime.demo.outputs import ( + Mp4OutputSink, + NullOutputSink, + OutputDecision, + OutputSink, + SessionInfo, + build_output_sink, + build_output_target, +) +from flashdreams.runtime.demo.pipeline import StepOutcome, StepPipeline +from flashdreams.runtime.demo.replay import OutputSinkFactory, run_replay_demo +from flashdreams.runtime.demo.run_modes import ( + AsyncSessionDriver, + BenchmarkErrorPolicy, + DefaultErrorPolicy, + ErrorAction, + InMemorySessionMetricsRecorder, + MetricsSnapshot, + Mp4ErrorPolicy, + NativeWindowErrorPolicy, + NoopTransportService, + NullErrorPolicy, + RunContext, + RunMode, + RunModeCapabilities, + RunModeWarmup, + RunResult, + RunSummary, + SessionDriver, + SessionEdges, + SingleSessionAdmissionPolicy, + WebRTCErrorPolicy, + build_model_warmup_plan, + warmup_run_context, +) +from flashdreams.runtime.demo.session_inputs import ( + BatchInputSource, + ControlDecision, + InputSource, + ModelInputProvider, + PreparedStep, + ProviderCapabilities, + RealtimeInputSource, + UserInputWindow, +) from flashdreams.runtime.demo.spec import ( DemoAdapter, DemoSpec, + ModelWarmupAdapter, Mp4OutputSpec, NullOutputSpec, OutputSpec, @@ -15,16 +74,105 @@ WebRTCAppResources, WebRTCOutputSpec, ) +from flashdreams.runtime.demo.timing import ( + SPARSE_KEY_SEGMENTS_METADATA_KEY, + ActivationPolicy, + ActivationResult, + ActivationSignal, + AlwaysActiveActivationPolicy, + CatchUpDecision, + CatchUpPolicy, + DeterministicClock, + KeyboardRealtimeInputSource, + RealtimeClock, + RealtimeWindowResult, + ResamplerRealtimeClock, + SignalActivationPolicy, + input_frame_count_from_request, +) +from flashdreams.runtime.demo.validation import ( + ResolvedRunCapabilities, + resolve_run_capabilities, + validate_resolved_run, +) __all__ = [ + "BatchInputSource", + "BatchSessionDriver", + "CLEANUP_TIMEOUT_S", + "ControlDecision", + "DefaultErrorPolicy", "DemoAdapter", "DemoSpec", + "DriverInvariantError", + "ErrorAction", + "AsyncSessionDriver", + "ActivationPolicy", + "ActivationResult", + "ActivationSignal", + "AlwaysActiveActivationPolicy", + "BenchmarkErrorPolicy", + "InMemorySessionMetricsRecorder", + "InputSource", + "CatchUpDecision", + "CatchUpPolicy", + "DeterministicClock", + "MetricsSnapshot", + "ModelWarmupAdapter", + "ModelWarmupPlan", + "ModelInputProvider", + "KeyboardRealtimeInputSource", + "Mp4ErrorPolicy", + "Mp4OutputSink", "Mp4OutputSpec", + "NativeWindowErrorPolicy", + "NoopTransportService", "NullOutputSpec", + "NullOutputSink", + "NullErrorPolicy", + "OutputDecision", + "OutputSinkFactory", "OutputSpec", + "OutputSink", "PreparedScenario", + "PreparedStep", + "ProviderCapabilities", + "RealtimeInputSource", + "RealtimeClock", + "RealtimeSessionDriver", + "RealtimeWindowResult", + "ResolvedRunCapabilities", + "ResamplerRealtimeClock", + "RunContext", + "RunMode", + "RunModeCapabilities", + "RunModeWarmup", + "RunResult", + "RunSummary", + "RuntimeHost", + "SessionEdges", + "SessionDriver", + "SessionInfo", + "SignalActivationPolicy", + "SingleSessionAdmissionPolicy", + "SPARSE_KEY_SEGMENTS_METADATA_KEY", + "StepOutcome", + "StepPipeline", + "UserInputWindow", + "WarmupSessionInputs", "WebRTCAppResources", + "WebRTCErrorPolicy", "WebRTCOutputSpec", + "build_output_sink", "build_output_target", + "build_model_warmup_plan", + "input_frame_count_from_request", + "resolve_run_capabilities", + "run_demo_session", + "run_demo_session_async", "run_replay_demo", + "shielded_session_cleanup", + "uncancel_current_task", + "validate_resolved_run", + "warmup_run_context", ] diff --git a/flashdreams/flashdreams/runtime/demo/app.py b/flashdreams/flashdreams/runtime/demo/app.py index b84e59163..d44dd81bb 100644 --- a/flashdreams/flashdreams/runtime/demo/app.py +++ b/flashdreams/flashdreams/runtime/demo/app.py @@ -6,6 +6,7 @@ from __future__ import annotations import argparse +import sys from abc import ABC, abstractmethod from typing import Any @@ -29,10 +30,18 @@ def main(self, argv: list[str] | None = None) -> None: configure_logging() args = self.parse_args(argv) if args.command == "replay": - run_replay_demo( + result = run_replay_demo( spec=self.replay_spec(args), adapter=self.replay_adapter(), ) + if result.status != "completed": + reason = result.reason or ( + str(result.error) if result.error is not None else None + ) + if reason is None: + reason = f"Replay demo ended with status {result.status!r}." + print(reason, file=sys.stderr) + raise SystemExit(1) return if args.command == "webrtc": context = initialize_cuda_distributed( diff --git a/flashdreams/flashdreams/runtime/demo/drivers.py b/flashdreams/flashdreams/runtime/demo/drivers.py new file mode 100644 index 000000000..9bdd40db8 --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/drivers.py @@ -0,0 +1,1011 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Session drivers and helpers for demo runtime vertical slices.""" + +from __future__ import annotations + +import asyncio +import inspect +from typing import Any, cast + +from flashdreams.runtime.interfaces import InferenceSession +from flashdreams.runtime.types import ( + StepRequest, + StepRequirements, + step_requirements_from_request, +) + +from .host import RuntimeHost +from .outputs import SessionInfo +from .pipeline import StepPipeline +from .run_modes import ( + DriverStatus, + RunContext, + RunMode, + RunResult, + SessionEdges, + SessionReservation, +) +from .session_inputs import BatchInputSource, ModelInputProvider +from .spec import DemoAdapter, DemoSpec, PreparedScenario +from .timing import ActivationPolicy, RealtimeClock +from .validation import resolve_run_capabilities, validate_resolved_run + +CLEANUP_TIMEOUT_S = 30.0 +_MODEL_CLEANUP_FAILED_REASON = "model-affine cleanup failed" +_MODEL_CLEANUP_TIMED_OUT_REASON = "model-affine cleanup timed out" + + +class DriverInvariantError(RuntimeError): + """A driver invariant was violated; this is a driver bug, not a run result.""" + + +class BatchSessionDriver: + """Minimal finite-session driver for Phase 2 fake-model coverage.""" + + def run_one_session( + self, + *, + host: RuntimeHost, + provider: ModelInputProvider, + session_edges: SessionEdges, + pipeline: StepPipeline, + ) -> RunResult: + session: InferenceSession | None = None + final_status: DriverStatus = "completed" + final_reason: str | None = None + final_error: Exception | None = None + invariant_closed = False + setup_ok = False + try: + try: + initial_input = host.call(provider.prepare_initial_input) + session = host.call(host.start_session, initial_input) + session_info = host.call(_session_info, session) + session_edges.output_sink.open(session_info) + setup_ok = True + except Exception as exc: + action = session_edges.error_policy.handle_setup_error(exc) + if action.drop_chunk or action.result_status == "completed": + raise DriverInvariantError( + "Setup failures must resolve to failed or skipped." + ) from exc + session_edges.metrics.record_error(exc, action) + final_status = action.result_status + final_reason = str(exc) + final_error = exc if action.result_status == "failed" else None + + input_source = cast(BatchInputSource, session_edges.input_source) + while setup_ok: + if session is None: + raise DriverInvariantError("setup_ok was set without a session.") + try: + if session_edges.input_source.is_finished(): + break + request = _next_step_requirements(host=host, session=session) + if request is None: + break + user_window = input_source.next_window(request) + outcome = host.call( + pipeline.execute_step, + request=request, + user_window=user_window, + provider=provider, + session=session, + output=session_edges.output_sink, + metrics=session_edges.metrics, + ) + if outcome.control.reset: + host.call(session.reset, outcome.control.reset_input) + if not outcome.control.provider_already_reset: + host.call(provider.reset, outcome.control.reset_input) + continue + if outcome.control.close_session: + break + if outcome.output.should_stop: + break + except DriverInvariantError: + raise + except Exception as exc: + action = session_edges.error_policy.handle(exc) + session_edges.metrics.record_error(exc, action) + if action.drop_chunk: + continue + final_status = action.result_status + final_reason = str(exc) + final_error = exc if action.result_status == "failed" else None + break + except DriverInvariantError as exc: + if session is not None: + _close_on_host_best_effort( + host=host, + close=session.close, + session_edges=session_edges, + ) + _close_on_host_best_effort( + host=host, + close=provider.close, + session_edges=session_edges, + ) + session_edges.close_result( + status="failed", + reason=str(exc), + error=exc, + ) + invariant_closed = True + raise + except Exception as exc: + final_status = "failed" + final_reason = str(exc) + final_error = exc + finally: + if not invariant_closed: + if session is not None: + _close_on_host_best_effort( + host=host, + close=session.close, + session_edges=session_edges, + ) + _close_on_host_best_effort( + host=host, + close=provider.close, + session_edges=session_edges, + ) + + return session_edges.close_result( + status=final_status, + reason=final_reason, + error=final_error, + ) + + +class RealtimeSessionDriver: + """Async realtime session driver built on shared Phase 5 primitives.""" + + cleanup_timeout_s: float + + def __init__(self, *, cleanup_timeout_s: float = CLEANUP_TIMEOUT_S) -> None: + if cleanup_timeout_s <= 0: + raise ValueError("cleanup_timeout_s must be > 0.") + self.cleanup_timeout_s = float(cleanup_timeout_s) + + async def run_one_session( + self, + *, + host: RuntimeHost, + provider: ModelInputProvider, + session_edges: SessionEdges, + pipeline: StepPipeline, + ) -> RunResult: + session: InferenceSession | None = None + final_status: DriverStatus = "completed" + final_reason: str | None = None + final_error: Exception | None = None + setup_ok = False + generation = 0 + first_step_started = False + invariant_error: DriverInvariantError | None = None + try: + activation, clock = _realtime_activation_and_clock(session_edges) + input_source = _realtime_input_source(session_edges) + activation_result = await activation.wait_until_active(clock) + if not activation_result.activated: + final_status = "not_activated" + final_reason = activation_result.reason + elif not session_edges.transport.is_active(): + final_status = "not_activated" + final_reason = "transport closed before first step" + else: + try: + initial_input = await host.call_async( + provider.prepare_initial_input + ) + session = await host.call_async(host.start_session, initial_input) + session_info = await host.call_async(_session_info, session) + session_edges.output_sink.open(session_info) + session_edges.output_sink.begin_generation(generation) + setup_ok = True + except Exception as exc: + action = session_edges.error_policy.handle_setup_error(exc) + if action.drop_chunk or action.result_status == "completed": + raise DriverInvariantError( + "Setup failures must resolve to failed or skipped." + ) from exc + session_edges.metrics.record_error(exc, action) + final_status = action.result_status + final_reason = str(exc) + final_error = exc if action.result_status == "failed" else None + + while setup_ok: + if session is None: + raise DriverInvariantError("setup_ok was set without a session.") + if not session_edges.transport.is_active(): + if not first_step_started: + final_status = "not_activated" + final_reason = "transport closed before first step" + break + try: + request = await _next_step_requirements_async( + host=host, + session=session, + ) + if request is None: + break + window_result = await input_source.next_realtime_window( + request=request, + clock=clock, + ) + session_edges.metrics.record_catch_up(window_result.catch_up) + if ( + not session_edges.transport.is_active() + and not first_step_started + ): + final_status = "not_activated" + final_reason = "transport closed before first step" + break + outcome = await host.call_async( + pipeline.execute_step, + request=request, + user_window=window_result.window, + provider=provider, + session=session, + output=session_edges.output_sink, + metrics=session_edges.metrics, + ) + first_step_started = True + if outcome.control.reset: + await host.call_async( + session.reset, + outcome.control.reset_input, + ) + if not outcome.control.provider_already_reset: + await host.call_async( + provider.reset, + outcome.control.reset_input, + ) + generation += 1 + session_edges.output_sink.begin_generation(generation) + continue + if outcome.control.close_session: + break + if outcome.output.should_stop: + break + if outcome.output.backpressure_s > 0: + await clock.apply_backpressure(outcome.output.backpressure_s) + except DriverInvariantError: + raise + except Exception as exc: + action = session_edges.error_policy.handle(exc) + session_edges.metrics.record_error(exc, action) + if action.close_session: + final_status = action.result_status + final_reason = str(exc) + final_error = exc if action.result_status == "failed" else None + break + if action.drop_chunk: + continue + final_status = "failed" + final_reason = str(exc) + final_error = exc + break + except asyncio.CancelledError: + uncancel_current_task() + final_status = "cancelled" + final_reason = ( + "cancelled before first step" if session is None else "cancelled" + ) + final_error = None + except DriverInvariantError as exc: + invariant_error = exc + final_status = "failed" + final_reason = str(exc) + final_error = exc + except Exception as exc: + final_status = "failed" + final_reason = str(exc) + final_error = exc + + result = await shielded_session_cleanup( + host=host, + session=session, + provider=provider, + session_edges=session_edges, + status=final_status, + reason=final_reason, + error=final_error, + timeout_s=self.cleanup_timeout_s, + ) + if invariant_error is not None: + raise invariant_error + return result + + +def _mark_host_cleanup_failed(host: RuntimeHost, exc: Exception | None = None) -> None: + host.mark_unhealthy(_MODEL_CLEANUP_FAILED_REASON, exc) + + +def run_demo_session( + *, + context: RunContext, + spec: DemoSpec, + scenario: PreparedScenario, + adapter: DemoAdapter, + run_mode: RunMode, + pipeline: StepPipeline, + reservation: SessionReservation | None = None, +) -> RunResult: + """Run one prepared demo session through a selected run mode.""" + if reservation is None: + reservation = context.admission.try_reserve() + if reservation is None: + result = RunResult.rejected(reason="busy") + context.run_metrics.record_session(result) + return result + + provider: Any | None = None + session_edges: SessionEdges | None = None + driver_started = False + try: + create_provider = getattr(adapter, "create_model_input_provider") + provider = context.host.call(create_provider, spec, scenario) + run_mode.validate_session( + spec=spec, + scenario=scenario, + adapter=adapter, + provider=provider, + ) + session_edges = run_mode.create_session_edges( + context=context, + spec=spec, + scenario=scenario, + provider=provider, + adapter=adapter, + ) + resolved_capabilities = resolve_run_capabilities( + spec=spec, + provider=provider, + session_edges=session_edges, + ) + validate_resolved_run( + spec=spec, + adapter=adapter, + provider=provider, + run_mode=run_mode, + session_edges=session_edges, + resolved=resolved_capabilities, + ) + if session_edges.is_closed: + raise DriverInvariantError( + "RunMode returned already closed SessionEdges; session edges " + "must not be reused." + ) + driver = run_mode.select_driver() + driver_started = True + result = _run_sync_driver( + driver=driver, + host=context.host, + provider=provider, + session_edges=session_edges, + pipeline=pipeline, + ) + context.run_metrics.record_session(result) + return result + except DriverInvariantError as exc: + _record_run_session_error(context, exc) + if provider is not None and not driver_started: + _close_partial_provider_sync( + context=context, + provider=provider, + session_edges=session_edges, + ) + if session_edges is not None and ( + driver_started or not session_edges.is_closed + ): + result = session_edges.close_result( + status="failed", + reason=str(exc), + error=exc, + ) + context.run_metrics.record_session(result) + raise + except Exception as exc: + _record_run_session_error(context, exc) + if provider is not None and not driver_started: + _close_partial_provider_sync( + context=context, + provider=provider, + session_edges=session_edges, + ) + if session_edges is not None and ( + driver_started or not session_edges.is_closed + ): + result = session_edges.close_result( + status="failed", + reason=str(exc), + error=exc, + ) + else: + result = RunResult(status="failed", reason=str(exc), error=exc) + context.run_metrics.record_session(result) + return result + finally: + reservation.release() + + +async def run_demo_session_async( + *, + context: RunContext, + spec: DemoSpec, + scenario: PreparedScenario, + adapter: DemoAdapter, + run_mode: RunMode, + pipeline: StepPipeline, + reservation: SessionReservation | None = None, +) -> RunResult: + """Run one prepared async/realtime demo session through a selected run mode.""" + if reservation is None: + reservation = context.admission.try_reserve() + if reservation is None: + result = RunResult.rejected(reason="busy") + context.run_metrics.record_session(result) + return result + + provider: Any | None = None + session_edges: SessionEdges | None = None + try: + try: + create_provider = getattr(adapter, "create_model_input_provider") + provider = await context.host.call_async(create_provider, spec, scenario) + run_mode.validate_session( + spec=spec, + scenario=scenario, + adapter=adapter, + provider=provider, + ) + session_edges = run_mode.create_session_edges( + context=context, + spec=spec, + scenario=scenario, + provider=provider, + adapter=adapter, + ) + resolved_capabilities = resolve_run_capabilities( + spec=spec, + provider=provider, + session_edges=session_edges, + ) + validate_resolved_run( + spec=spec, + adapter=adapter, + provider=provider, + run_mode=run_mode, + session_edges=session_edges, + resolved=resolved_capabilities, + ) + if session_edges.is_closed: + raise DriverInvariantError( + "RunMode returned already closed SessionEdges; session edges " + "must not be reused." + ) + driver = run_mode.select_driver() + result = await _run_async_driver( + driver=driver, + host=context.host, + provider=provider, + session_edges=session_edges, + pipeline=pipeline, + ) + context.run_metrics.record_session(result) + return result + except asyncio.CancelledError: + uncancel_current_task() + result = await _close_partial_session_async( + context=context, + provider=provider, + session_edges=session_edges, + status="cancelled", + reason="cancelled during session assembly", + error=None, + close_provider=_needs_partial_provider_cleanup(session_edges), + ) + context.run_metrics.record_session(result) + return result + except DriverInvariantError as exc: + _record_run_session_error(context, exc) + should_record_session = session_edges is not None + result = await _close_partial_session_async( + context=context, + provider=provider, + session_edges=session_edges, + status="failed", + reason=str(exc), + error=exc, + close_provider=_needs_partial_provider_cleanup(session_edges), + ) + if should_record_session: + context.run_metrics.record_session(result) + raise + except Exception as exc: + _record_run_session_error(context, exc) + result = await _close_partial_session_async( + context=context, + provider=provider, + session_edges=session_edges, + status="failed", + reason=str(exc), + error=exc, + close_provider=_needs_partial_provider_cleanup(session_edges), + ) + context.run_metrics.record_session(result) + return result + finally: + reservation.release() + + +def _session_info(session: InferenceSession) -> SessionInfo: + session_info = getattr(session, "session_info", None) + if not callable(session_info): + return SessionInfo() + value = session_info() + if not isinstance(value, SessionInfo): + raise TypeError( + "session.session_info() must return SessionInfo, " + f"got {type(value).__name__}." + ) + return value + + +def _next_step_requirements( + *, + host: RuntimeHost, + session: InferenceSession, +) -> StepRequirements | None: + next_requirements = getattr(session, "next_step_requirements", None) + if callable(next_requirements): + return _coerce_step_requirements(host.call(next_requirements)) + + next_request = getattr(session, "next_step_request", None) + if not callable(next_request): + raise TypeError( + "InferenceSession must provide next_step_requirements() or " + "legacy next_step_request()." + ) + return _coerce_step_requirements(host.call(next_request)) + + +async def _next_step_requirements_async( + *, + host: RuntimeHost, + session: InferenceSession, +) -> StepRequirements | None: + next_requirements = getattr(session, "next_step_requirements", None) + if callable(next_requirements): + return _coerce_step_requirements(await host.call_async(next_requirements)) + + next_request = getattr(session, "next_step_request", None) + if not callable(next_request): + raise TypeError( + "InferenceSession must provide next_step_requirements() or " + "legacy next_step_request()." + ) + return _coerce_step_requirements(await host.call_async(next_request)) + + +def _coerce_step_requirements(value: object) -> StepRequirements | None: + if value is None: + return None + if isinstance(value, StepRequirements): + return value + if isinstance(value, StepRequest): + return step_requirements_from_request(value) + raise TypeError( + "Session next-step method must return StepRequirements, legacy " + f"StepRequest, or None; got {type(value).__name__}." + ) + + +def _close_safely(close: Any, session_edges: SessionEdges) -> bool: + try: + close() + except Exception as exc: + session_edges.record_cleanup_error(exc) + return False + return True + + +def _close_on_host_best_effort( + *, + host: RuntimeHost, + close: Any, + session_edges: SessionEdges, +) -> bool: + try: + cleanup_succeeded = host.call(_close_safely, close, session_edges) + except Exception as exc: + # If the host/worker is already unavailable, do not fall back to calling + # model-affine cleanup directly on the caller thread. Record the loss and + # let close_result finalize output, transport, and metrics. + session_edges.record_cleanup_error(exc) + _mark_host_cleanup_failed(host, exc) + return False + if not cleanup_succeeded: + _mark_host_cleanup_failed(host) + return False + return True + + +def _close_partial_provider_sync( + *, + context: RunContext, + provider: Any, + session_edges: SessionEdges | None, +) -> None: + if session_edges is not None: + _close_on_host_best_effort( + host=context.host, + close=provider.close, + session_edges=session_edges, + ) + return + try: + cleanup_succeeded = context.host.call( + _close_run_provider_safely, + provider.close, + context, + ) + except Exception as exc: + _record_run_cleanup_error(context, exc) + _mark_host_cleanup_failed(context.host, exc) + return + if not cleanup_succeeded: + _mark_host_cleanup_failed(context.host) + + +def _close_run_provider_safely(close: Any, context: RunContext) -> bool: + try: + close() + except Exception as exc: + _record_run_cleanup_error(context, exc) + return False + return True + + +async def shielded_session_cleanup( + *, + host: RuntimeHost, + session: InferenceSession | None, + provider: ModelInputProvider, + session_edges: SessionEdges, + status: DriverStatus, + reason: str | None, + error: Exception | None, + timeout_s: float = CLEANUP_TIMEOUT_S, +) -> RunResult: + """Close realtime session resources exactly once without leaking cancellation.""" + + if timeout_s <= 0: + session_edges.record_cleanup_error(ValueError("timeout_s must be > 0.")) + return session_edges.close_result(status=status, reason=reason, error=error) + + async def cleanup() -> RunResult: + unhealthy_reason = await _close_model_resources_async( + host=host, + session=session, + provider=provider, + session_edges=session_edges, + timeout_s=timeout_s, + ) + if unhealthy_reason is not None: + host.mark_unhealthy(unhealthy_reason) + return session_edges.close_result( + status=status, + reason=reason, + error=error, + ) + + try: + cleanup_task = asyncio.create_task(cleanup()) + except RuntimeError as exc: + session_edges.record_cleanup_error(exc) + return session_edges.close_result(status=status, reason=reason, error=error) + + session_edges.cleanup_tasks.add(cleanup_task) + try: + while not cleanup_task.done(): + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + uncancel_current_task() + continue + except Exception: + break + return _cleanup_result(cleanup_task, session_edges, status, reason, error) + finally: + session_edges.cleanup_tasks.discard(cleanup_task) + + +def _run_sync_driver( + *, + driver: object, + host: RuntimeHost, + provider: ModelInputProvider, + session_edges: SessionEdges, + pipeline: StepPipeline, +) -> RunResult: + run_one_session = getattr(driver, "run_one_session", None) + if not callable(run_one_session): + raise TypeError( + "RunMode.select_driver() must return an object with run_one_session(...)." + ) + result = run_one_session( + host=host, + provider=provider, + session_edges=session_edges, + pipeline=pipeline, + ) + if inspect.isawaitable(result): + raise TypeError( + "run_demo_session(...) requires a synchronous session driver; " + "use run_demo_session_async(...) for async drivers." + ) + if not isinstance(result, RunResult): + raise TypeError( + "Session driver run_one_session(...) must return RunResult, " + f"got {type(result).__name__}." + ) + return result + + +async def _run_async_driver( + *, + driver: object, + host: RuntimeHost, + provider: ModelInputProvider, + session_edges: SessionEdges, + pipeline: StepPipeline, +) -> RunResult: + run_one_session = getattr(driver, "run_one_session", None) + if not callable(run_one_session): + raise TypeError( + "RunMode.select_driver() must return an object with run_one_session(...)." + ) + result = run_one_session( + host=host, + provider=provider, + session_edges=session_edges, + pipeline=pipeline, + ) + if not inspect.isawaitable(result): + raise TypeError("run_demo_session_async(...) requires an async session driver.") + resolved = await result + if not isinstance(resolved, RunResult): + raise TypeError( + "Async session driver run_one_session(...) must return RunResult, " + f"got {type(resolved).__name__}." + ) + return resolved + + +async def _close_partial_session_async( + *, + context: RunContext, + provider: Any | None, + session_edges: SessionEdges | None, + status: DriverStatus, + reason: str | None, + error: Exception | None, + close_provider: bool, +) -> RunResult: + if provider is not None and close_provider and session_edges is not None: + return await shielded_session_cleanup( + host=context.host, + session=None, + provider=provider, + session_edges=session_edges, + status=status, + reason=reason, + error=error, + ) + if provider is not None and close_provider: + await _close_provider_async( + context=context, + provider=provider, + session_edges=session_edges, + ) + if session_edges is not None: + return session_edges.close_result(status=status, reason=reason, error=error) + return RunResult(status=status, reason=reason, error=error) + + +def _needs_partial_provider_cleanup(session_edges: SessionEdges | None) -> bool: + return session_edges is None or not session_edges.is_closed + + +async def _close_provider_async( + *, + context: RunContext, + provider: Any, + session_edges: SessionEdges | None, +) -> None: + try: + close_task = asyncio.create_task(context.host.call_async(provider.close)) + except RuntimeError as close_exc: + _record_provider_cleanup_error( + context=context, + session_edges=session_edges, + exc=close_exc, + ) + return + + while not close_task.done(): + try: + await asyncio.shield(close_task) + except asyncio.CancelledError: + uncancel_current_task() + continue + except Exception: + break + + try: + await close_task + except asyncio.CancelledError: + uncancel_current_task() + _record_provider_cleanup_error( + context=context, + session_edges=session_edges, + exc=RuntimeError("provider cleanup was cancelled"), + ) + except Exception as close_exc: + _record_provider_cleanup_error( + context=context, + session_edges=session_edges, + exc=close_exc, + ) + + +def _record_provider_cleanup_error( + *, + context: RunContext, + session_edges: SessionEdges | None, + exc: Exception, +) -> None: + if session_edges is not None: + session_edges.record_cleanup_error(exc) + else: + _record_run_cleanup_error(context, exc) + # Partial async assembly may only have a provider to close. If that + # model-affine cleanup fails, quarantine the host instead of admitting a new + # session onto a worker that may still own model resources. + _mark_host_cleanup_failed(context.host, exc) + + +def _record_run_cleanup_error(context: RunContext, exc: Exception) -> None: + try: + context.run_metrics.record_cleanup_error(exc) + except Exception: + return + + +def _record_run_session_error(context: RunContext, exc: Exception) -> None: + try: + context.run_metrics.record_session_error(exc) + except Exception: + return + + +async def _close_model_resources_async( + *, + host: RuntimeHost, + session: InferenceSession | None, + provider: ModelInputProvider, + session_edges: SessionEdges, + timeout_s: float, +) -> str | None: + try: + resources_closed = await asyncio.wait_for( + host.call_async( + _close_model_resources_safely, + session.close if session is not None else None, + provider.close, + session_edges, + ), + timeout=timeout_s, + ) + except asyncio.TimeoutError as exc: + # Keep provider cleanup ordered behind session cleanup on the model worker. + # A timed-out session close may still hold CUDA/Triton state, so running + # provider cleanup on another thread or replacing the worker is unsafe. + # The caller marks the host unhealthy so future sessions reject instead. + session_edges.record_orphaned_cleanup(exc) + return _MODEL_CLEANUP_TIMED_OUT_REASON + except Exception as exc: + session_edges.record_cleanup_error(exc) + return _MODEL_CLEANUP_FAILED_REASON + if not resources_closed: + return _MODEL_CLEANUP_FAILED_REASON + return None + + +def _close_model_resources_safely( + session_close: Any | None, + provider_close: Any, + session_edges: SessionEdges, +) -> bool: + resources_closed = True + # Session and provider close are intentionally ordered on the model worker. + # If session close hangs, timeout handling records orphaned cleanup and + # quarantines the host rather than moving provider close to another thread. + if session_close is not None: + resources_closed = _close_safely(session_close, session_edges) + return _close_safely(provider_close, session_edges) and resources_closed + + +def _cleanup_result( + cleanup_task: asyncio.Task[RunResult], + session_edges: SessionEdges, + status: DriverStatus, + reason: str | None, + error: Exception | None, +) -> RunResult: + if cleanup_task.done() and not cleanup_task.cancelled(): + exc = cleanup_task.exception() + if exc is None: + return cleanup_task.result() + if isinstance(exc, Exception): + session_edges.record_cleanup_error(exc) + else: + session_edges.record_cleanup_error( + RuntimeError(f"cleanup failed with {type(exc).__name__}") + ) + return session_edges.close_result(status=status, reason=reason, error=error) + + +def _realtime_activation_and_clock( + session_edges: SessionEdges, +) -> tuple[ActivationPolicy, RealtimeClock]: + activation = session_edges.activation + if activation is None: + raise DriverInvariantError( + "RealtimeSessionDriver requires SessionEdges.activation." + ) + clock = session_edges.clock + if not isinstance(clock, RealtimeClock): + raise DriverInvariantError("RealtimeSessionDriver requires a RealtimeClock.") + return activation, clock + + +def _realtime_input_source(session_edges: SessionEdges) -> Any: + input_source = session_edges.input_source + next_realtime_window = getattr(input_source, "next_realtime_window", None) + if not callable(next_realtime_window): + raise DriverInvariantError( + "RealtimeSessionDriver requires a RealtimeInputSource." + ) + return input_source + + +def uncancel_current_task() -> None: + task = asyncio.current_task() + if task is None: + return + uncancel = getattr(task, "uncancel", None) + if not callable(uncancel): + return + cancelling = getattr(task, "cancelling", None) + if not callable(cancelling): + return + while cancelling(): + uncancel() + + +__all__ = [ + "BatchSessionDriver", + "CLEANUP_TIMEOUT_S", + "DriverInvariantError", + "RealtimeSessionDriver", + "run_demo_session", + "run_demo_session_async", + "shielded_session_cleanup", + "uncancel_current_task", +] diff --git a/flashdreams/flashdreams/runtime/demo/host.py b/flashdreams/flashdreams/runtime/demo/host.py new file mode 100644 index 000000000..6a6e93ce5 --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/host.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Runtime host and model-execution boundary for shared demos.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import TypeVar + +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.inputs import InferenceInput +from flashdreams.runtime.interfaces import InferenceRuntime, InferenceSession +from flashdreams.runtime.worker import ModelExecutionWorker + +_T = TypeVar("_T") + + +@dataclass(frozen=True, kw_only=True, slots=True) +class WarmupSessionInputs: + """Inputs used to warm one temporary runtime session.""" + + initial_input: InferenceInput + step_inputs: Sequence[InferenceInput] = () + + def __post_init__(self) -> None: + object.__setattr__(self, "step_inputs", tuple(self.step_inputs)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class ModelWarmupPlan: + """Host-owned model warmup plan built by a demo adapter or run mode.""" + + sessions: Sequence[WarmupSessionInputs] = () + measured: bool = False + metadata: Mapping[str, object] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__(self, "sessions", tuple(self.sessions)) + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +class RuntimeHost: + """Own one runtime and the worker used for model-affine calls.""" + + def __init__( + self, + runtime: InferenceRuntime, + *, + worker: ModelExecutionWorker | None = None, + is_control_rank: bool = True, + worker_loop: Callable[[], None] | None = None, + ) -> None: + self._runtime = runtime + self._worker = worker or ModelExecutionWorker() + self._is_control_rank = is_control_rank + self._worker_loop = worker_loop + self._healthy = True + self._closed = False + self._unhealthy_reason: str | None = None + self._unhealthy_error: Exception | None = None + + @property + def runtime(self) -> InferenceRuntime: + """Return the hosted runtime.""" + return self._runtime + + @property + def worker(self) -> ModelExecutionWorker: + """Return the host's model-execution worker.""" + return self._worker + + @property + def is_control_rank(self) -> bool: + """Whether this process owns run modes, providers, sinks, and metrics.""" + return self._is_control_rank + + @property + def is_healthy(self) -> bool: + """Return whether admission should continue accepting sessions.""" + return self._healthy and not self._closed + + @property + def unhealthy_reason(self) -> str | None: + """Return the first latched unhealthy reason, if any.""" + return self._unhealthy_reason + + @property + def unhealthy_error(self) -> Exception | None: + """Return the first latched unhealthy error, if any.""" + return self._unhealthy_error + + def mark_unhealthy( + self, + reason: str = "marked unhealthy", + error: Exception | None = None, + ) -> None: + """Latch the host as unhealthy without overwriting the first reason.""" + if not self._healthy: + return + self._healthy = False + self._unhealthy_reason = reason + self._unhealthy_error = error + + def preload(self) -> None: + """Initialize optional distributed state and preload runtime resources.""" + self._call_optional_runtime_hook("initialize_distributed") + self._call_optional_runtime_hook("preload") + + def warmup(self, plan: ModelWarmupPlan | None = None) -> None: + """Run warmup sessions through the same worker boundary as real sessions.""" + plan = plan or ModelWarmupPlan() + for warmup_session in plan.sessions: + session = self.call(self.start_session, warmup_session.initial_input) + try: + for step_input in warmup_session.step_inputs: + self.call(session.step, step_input) + finally: + self.call(session.close) + + def call(self, func: Callable[..., _T], /, *args: object, **kwargs: object) -> _T: + """Run one model-affine callable synchronously on the worker.""" + self._require_open() + return self._worker.call_blocking(func, *args, **kwargs) + + async def call_async( + self, + func: Callable[..., _T], + /, + *args: object, + **kwargs: object, + ) -> _T: + """Run model-affine work without blocking realtime event loops.""" + self._require_open() + return await self._worker.call(func, *args, **kwargs) + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + """Start one inference session through the hosted runtime.""" + self._require_open() + return self._runtime.start_session(inputs) + + def run_worker_loop(self) -> None: + """Serve control-rank work on non-control ranks until runtime shutdown.""" + worker_loop = self._worker_loop + if worker_loop is None: + worker_loop = getattr(self._runtime, "run_worker_loop", None) + if worker_loop is None: + worker_loop = getattr(self._runtime, "wait_for_termination", None) + if callable(worker_loop): + worker_loop() + + def close(self) -> None: + """Close runtime-owned state and stop the model-execution worker.""" + if self._closed: + return + try: + self._worker.call_blocking(self._runtime.close) + self._call_optional_runtime_hook("close_distributed") + finally: + self._closed = True + self._worker.close_blocking() + + def _call_optional_runtime_hook(self, name: str) -> None: + hook = getattr(self._runtime, name, None) + if callable(hook): + self.call(hook) + + def _require_open(self) -> None: + if self._closed: + raise RuntimeError("runtime host is closed") + + +__all__ = ["ModelWarmupPlan", "RuntimeHost", "WarmupSessionInputs"] diff --git a/flashdreams/flashdreams/runtime/demo/outputs.py b/flashdreams/flashdreams/runtime/demo/outputs.py index 421ec3bb4..866acbb00 100644 --- a/flashdreams/flashdreams/runtime/demo/outputs.py +++ b/flashdreams/flashdreams/runtime/demo/outputs.py @@ -1,18 +1,270 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Shared demo output-target construction.""" +"""Shared demo output contracts and output construction.""" from __future__ import annotations +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field from pathlib import Path +from typing import Literal, Protocol, runtime_checkable -from flashdreams.runtime.output import NullOutputTarget, OutputTarget +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.infra.runner_io import ( + DEFAULT_RUNNER_INSTALL_HINT, + write_video_tensor, +) +from flashdreams.infra.video_output import VideoResultCollector, prepare_video_for_mp4 +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.output import NullOutputTarget, OutputArtifact, OutputTarget +from flashdreams.runtime.types import StepResult from flashdreams.runtime.video_output import Mp4VideoOutputTarget, VideoWriter from .spec import Mp4OutputSpec, NullOutputSpec, OutputSpec, WebRTCOutputSpec +@dataclass(frozen=True, kw_only=True, slots=True) +class SessionInfo: + """Output-facing metadata known after session setup.""" + + output_layout: str | None = None + steady_output_frame_count: int | None = None + metadata: Mapping[str, object] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.output_layout is not None and not self.output_layout.strip(): + raise ValueError("SessionInfo.output_layout must be non-empty when set.") + if ( + self.steady_output_frame_count is not None + and self.steady_output_frame_count < 0 + ): + raise ValueError( + "SessionInfo.steady_output_frame_count must be >= 0 when set." + ) + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class OutputDecision: + """Flow-control decision returned by an output sink after one step.""" + + should_stop: bool = False + dropped: bool = False + drop_policy: Literal["none", "drop_newest", "drop_oldest"] = "none" + backpressure_s: float = 0.0 + metadata: Mapping[str, object] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.drop_policy not in {"none", "drop_newest", "drop_oldest"}: + raise ValueError(f"Unsupported drop_policy={self.drop_policy!r}.") + if self.backpressure_s < 0: + raise ValueError("OutputDecision.backpressure_s must be >= 0.") + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@runtime_checkable +class OutputSink(Protocol): + """Consumes generated session outputs for a demo run mode.""" + + produces_artifacts: bool + + def open(self, session_info: SessionInfo) -> None: + """Prepare output resources for a session.""" + ... + + def begin_generation(self, generation: int) -> None: + """Start an output generation, discarding stale live output if needed.""" + ... + + def write(self, result: StepResult) -> OutputDecision: + """Consume one generated result and return output flow-control state.""" + ... + + def close(self) -> Sequence[OutputArtifact]: + """Finalize output resources and return produced artifacts.""" + ... + + +@dataclass(slots=True) +class NullOutputSink: + """Output sink for headless runs and fake-model vertical-slice tests.""" + + store_results: bool = False + produces_artifacts: bool = False + output_count: int = field(default=0, init=False) + results: list[Mapping[str, object]] = field(default_factory=list, init=False) + opened: bool = field(default=False, init=False) + closed: bool = field(default=False, init=False) + session_info: SessionInfo | None = field(default=None, init=False) + generation: int | None = field(default=None, init=False) + + def open(self, session_info: SessionInfo) -> None: + self.session_info = session_info + self.output_count = 0 + self.results.clear() + self.opened = True + self.closed = False + + def begin_generation(self, generation: int) -> None: + if generation < 0: + raise ValueError("generation must be >= 0.") + self.generation = generation + + def write(self, result: StepResult) -> OutputDecision: + if not self.opened or self.closed: + raise RuntimeError("Cannot write to a closed output sink.") + self.output_count += 1 + if self.store_results: + self.results.append(_result_record(result)) + return OutputDecision() + + def close(self) -> Sequence[OutputArtifact]: + self.closed = True + return () + + +@dataclass(slots=True) +class Mp4OutputSink: + """MP4 artifact sink for shared demo drivers.""" + + output_path: Path + fps: int | float + output_layout: VideoTensorLayout = "bvtchw" + writer: VideoWriter = field(default=write_video_tensor, repr=False) + install_hint: str = DEFAULT_RUNNER_INSTALL_HINT + move_to_cpu: bool = True + enabled: bool = True + produces_artifacts: bool = True + _opened: bool = field(default=False, init=False, repr=False) + _closed: bool = field(default=True, init=False, repr=False) + _collector: VideoResultCollector | None = field( + default=None, + init=False, + repr=False, + ) + _artifacts: tuple[OutputArtifact, ...] | None = field( + default=None, + init=False, + repr=False, + ) + session_info: SessionInfo | None = field(default=None, init=False) + + def __post_init__(self) -> None: + if float(self.fps) <= 0: + raise ValueError("Mp4OutputSink.fps must be > 0.") + self.output_path = Path(self.output_path) + + def open(self, session_info: SessionInfo) -> None: + self.session_info = session_info + self._collector = VideoResultCollector( + output_layout=self.output_layout, + enabled=self.enabled, + move_to_cpu=self.move_to_cpu, + ) + self._artifacts = None + self._opened = True + self._closed = False + + def begin_generation(self, generation: int) -> None: + if generation < 0: + raise ValueError("generation must be >= 0.") + + def write(self, result: StepResult) -> OutputDecision: + if not self._opened or self._closed or self._collector is None: + raise RuntimeError("Cannot write to a closed output sink.") + if result.layout is None: + raise TypeError("Mp4OutputSink requires a video StepResult with layout.") + if result.layout != self.output_layout: + raise ValueError( + "Mp4OutputSink received layout " + f"{result.layout!r}; expected {self.output_layout!r}." + ) + self._collector.add(result) + return OutputDecision() + + def close(self) -> Sequence[OutputArtifact]: + if self._artifacts is not None: + return self._artifacts + if self._collector is None: + self._opened = False + self._closed = True + self._artifacts = () + return self._artifacts + + collector = self._collector + self._collector = None + self._opened = False + self._closed = True + video = collector.finish() + if video is None: + self._artifacts = () + return self._artifacts + 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=writable_layout, + install_hint=self.install_hint, + ) + self._artifacts = ( + OutputArtifact( + kind="video/mp4", + uri=str(path), + metadata={ + "fps": self.fps, + "source_layout": self.output_layout, + "shape": tuple(int(dim) for dim in video.shape), + "stats_history": tuple(collector.stats_history), + }, + ), + ) + return self._artifacts + + +def build_output_sink( + output: OutputSpec, + *, + mp4_writer: VideoWriter | None = None, +) -> OutputSink: + """Build a shared demo output sink from a demo output spec.""" + if isinstance(output, NullOutputSpec): + return NullOutputSink(store_results=output.store_results) + if isinstance(output, Mp4OutputSpec): + writer = mp4_writer or write_video_tensor + return Mp4OutputSink( + output_path=Path(output.path), + fps=output.fps, + output_layout=output.output_layout, + writer=writer, + move_to_cpu=output.move_to_cpu, + ) + if isinstance(output, WebRTCOutputSpec): + raise ValueError("WebRTC output requires a realtime transport sink.") + raise TypeError(f"Unsupported demo output spec: {type(output).__name__}.") + + +def _result_record(result: StepResult) -> Mapping[str, object]: + record: dict[str, object] = { + "step_index": result.step_index, + "frame_count": result.frame_count, + "metrics": dict(result.metrics), + "metadata": dict(result.metadata), + } + if result.layout is not None: + record["layout"] = result.layout + if result.output_window is not None: + record["output_window"] = ( + result.output_window.start_s, + result.output_window.end_s, + ) + return freeze_mapping(record) + + def build_output_target( output: OutputSpec, *, @@ -42,4 +294,12 @@ def build_output_target( raise TypeError(f"Unsupported demo output spec: {type(output).__name__}.") -__all__ = ["build_output_target"] +__all__ = [ + "Mp4OutputSink", + "NullOutputSink", + "OutputDecision", + "OutputSink", + "SessionInfo", + "build_output_sink", + "build_output_target", +] diff --git a/flashdreams/flashdreams/runtime/demo/pipeline.py b/flashdreams/flashdreams/runtime/demo/pipeline.py new file mode 100644 index 000000000..5a9da08b0 --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/pipeline.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared per-step pipeline for demo session drivers.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from flashdreams.runtime.interfaces import InferenceSession +from flashdreams.runtime.types import StepRequirements, StepResult + +from .outputs import OutputDecision, OutputSink +from .run_modes import SessionMetricsRecorder +from .session_inputs import ControlDecision, ModelInputProvider, UserInputWindow + + +@dataclass(frozen=True, kw_only=True, slots=True) +class StepOutcome: + """Combined output and control result from one shared model step.""" + + output: OutputDecision = field(default_factory=OutputDecision) + control: ControlDecision = field(default_factory=ControlDecision) + + +class StepPipeline: + """Shared invariant for provider conversion, model step, output, and metrics.""" + + def execute_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + provider: ModelInputProvider, + session: InferenceSession, + output: OutputSink, + metrics: SessionMetricsRecorder, + ) -> StepOutcome: + prepared = provider.prepare_step( + request=request, + user_window=user_window, + ) + if prepared.control.reset or prepared.control.close_session: + metrics.record_control( + request=request, + user_window=user_window, + control=prepared.control, + ) + return StepOutcome(control=prepared.control) + if prepared.inference_input is None: + raise RuntimeError("ModelInputProvider returned no inference input.") + + result = session.step(prepared.inference_input) + if not isinstance(result, StepResult): + raise TypeError( + "InferenceSession.step must return StepResult, " + f"got {type(result).__name__}." + ) + decision = output.write(result) + if not isinstance(decision, OutputDecision): + raise TypeError( + "OutputSink.write must return OutputDecision, " + f"got {type(decision).__name__}." + ) + metrics.record_step( + request=request, + user_window=user_window, + inference_input=prepared.inference_input, + result=result, + decision=decision, + ) + return StepOutcome(output=decision) + + +__all__ = ["StepOutcome", "StepPipeline"] diff --git a/flashdreams/flashdreams/runtime/demo/replay.py b/flashdreams/flashdreams/runtime/demo/replay.py index 18b873254..fa2f507a9 100644 --- a/flashdreams/flashdreams/runtime/demo/replay.py +++ b/flashdreams/flashdreams/runtime/demo/replay.py @@ -5,28 +5,72 @@ from __future__ import annotations +import math from collections.abc import Callable, Sequence +from flashdreams.runtime.canonical import InputCanonicalizer +from flashdreams.runtime.config import InferenceConfig +from flashdreams.runtime.inputs import ( + CanonicalInputs, + CanonicalInputSchema, + InferenceInput, + InferenceInputSchema, + TimeWindow, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.interfaces import InferenceRuntime, InferenceSession +from flashdreams.runtime.mapping import ( + DeclaresMappingSchema, + InputMapping, + check_mapping_compatibility, +) from flashdreams.runtime.metrics import MetricsRecorder, NullMetricsRecorder from flashdreams.runtime.output import OutputArtifact, OutputTarget -from flashdreams.runtime.runner import run_inference_session +from flashdreams.runtime.types import ( + StepRequest, + StepRequirements, + StepResult, + step_requirements_from_request, +) -from .outputs import build_output_target -from .spec import DemoAdapter, DemoSpec, OutputSpec, WebRTCOutputSpec +from .drivers import BatchSessionDriver, run_demo_session +from .host import ModelWarmupPlan, RuntimeHost +from .outputs import OutputDecision, OutputSink, build_output_sink, build_output_target +from .pipeline import StepPipeline +from .run_modes import ( + Mp4ErrorPolicy, + NullErrorPolicy, + RunContext, + RunModeCapabilities, + RunResult, + SessionEdges, + SingleSessionAdmissionPolicy, +) +from .session_inputs import PreparedStep, ProviderCapabilities, UserInputWindow +from .spec import ( + DemoAdapter, + DemoSpec, + OutputSpec, + PreparedScenario, + WebRTCOutputSpec, +) OutputTargetFactory = Callable[[OutputSpec], OutputTarget] InferenceSessionRunner = Callable[..., Sequence[OutputArtifact]] +OutputSinkFactory = Callable[[OutputSpec], OutputSink] def run_replay_demo( *, spec: DemoSpec, adapter: DemoAdapter, - output_target_factory: OutputTargetFactory = build_output_target, + output_target_factory: OutputTargetFactory | None = None, + output_sink_factory: OutputSinkFactory = build_output_sink, metrics: MetricsRecorder | None = None, - runner: InferenceSessionRunner = run_inference_session, -) -> tuple[OutputArtifact, ...]: - """Run one prepared demo scenario through the shared runtime runner.""" + runner: InferenceSessionRunner | None = None, +) -> RunResult: + """Run one prepared replay scenario through the shared batch demo path.""" _require_supported_mode( mode=spec.input_mode, supported=adapter.supported_input_modes(), @@ -55,12 +99,87 @@ def run_replay_demo( if spec.config is None: raise RuntimeError("DemoSpec.config was not initialized.") + if runner is not None: + return _run_replay_demo_with_compat_runner( + spec=spec, + adapter=adapter, + prepared=prepared, + mapping=mapping, + output_target_factory=output_target_factory or build_output_target, + metrics=metrics, + runner=runner or _default_inference_session_runner(), + ) + + if output_target_factory is not None: + output_sink_factory = _output_target_sink_factory(output_target_factory) + + return _run_replay_demo_with_run_mode( + spec=spec, + adapter=adapter, + prepared=prepared, + mapping=mapping, + output_sink_factory=output_sink_factory, + metrics=metrics, + ) + + +def _output_target_sink_factory( + output_target_factory: OutputTargetFactory, +) -> OutputSinkFactory: + def create_output_sink(output_spec: OutputSpec) -> "_OutputTargetSink": + return _OutputTargetSink(output_target_factory(output_spec)) + + return create_output_sink + + +class _OutputTargetSink: + produces_artifacts = True + + def __init__(self, output: OutputTarget) -> None: + self._output = output + self._closed = True + self._artifacts: tuple[OutputArtifact, ...] | None = None + + def open(self, session_info: object) -> None: + del session_info + self._output.open() + self._closed = False + self._artifacts = None + + def begin_generation(self, generation: int) -> None: + del generation + + def write(self, result: StepResult) -> OutputDecision: + self._output.write(result) + return OutputDecision() + + def close(self) -> Sequence[OutputArtifact]: + if self._artifacts is not None: + return self._artifacts + if self._closed: + self._artifacts = () + return self._artifacts + self._closed = True + self._artifacts = tuple(self._output.close()) + return self._artifacts + + +def _run_replay_demo_with_compat_runner( + *, + spec: DemoSpec, + adapter: DemoAdapter, + prepared: "PreparedScenario", + mapping: InputMapping, + output_target_factory: OutputTargetFactory, + metrics: MetricsRecorder | None, + runner: InferenceSessionRunner, +) -> RunResult: output = output_target_factory(spec.output) metrics_recorder = metrics or NullMetricsRecorder() - return tuple( + artifacts = tuple( runner( adapter=adapter, - config=spec.config, + config=_require_config(spec), mapping=mapping, canonicalizer=prepared.canonicalizer, source_schema=prepared.source_schema, @@ -70,6 +189,422 @@ def run_replay_demo( metrics=metrics_recorder, ) ) + return RunResult(status="completed", artifacts=artifacts) + + +def _run_replay_demo_with_run_mode( + *, + spec: DemoSpec, + adapter: DemoAdapter, + prepared: "PreparedScenario", + mapping: InputMapping, + output_sink_factory: OutputSinkFactory, + metrics: MetricsRecorder | None, +) -> RunResult: + config = _require_config(spec) + _validate_replay_mapping( + adapter=adapter, + config=config, + mapping=mapping, + source_schema=prepared.source_schema, + canonicalizer=prepared.canonicalizer, + ) + request_state = _ReplayStepRequestState() + runtime = _ReplayRuntimeAdapter( + runtime=adapter.create_runtime(config), + request_state=request_state, + ) + host = RuntimeHost(runtime) + mode = _ReplayRunMode( + request_state=request_state, + output_sink_factory=output_sink_factory, + run_metrics=metrics or NullMetricsRecorder(), + ) + replay_adapter = _ReplayProviderAdapter( + adapter=adapter, + mapping=mapping, + request_state=request_state, + ) + context = mode.create_run_context( + spec=spec, + adapter=replay_adapter, + host=host, + model_warmup_plan=ModelWarmupPlan(), + ) + try: + return run_demo_session( + context=context, + spec=spec, + scenario=prepared, + adapter=replay_adapter, + run_mode=mode, + pipeline=StepPipeline(), + ) + finally: + context.close() + host.close() + + +class _ReplayRunMode: + name = "replay" + capabilities = RunModeCapabilities( + requires_finite_input=True, + supports_artifacts=True, + ) + + def __init__( + self, + *, + request_state: "_ReplayStepRequestState", + output_sink_factory: OutputSinkFactory, + run_metrics: MetricsRecorder, + ) -> None: + self._request_state = request_state + self._output_sink_factory = output_sink_factory + self._run_metrics = run_metrics + + def validate_run(self, *, spec: DemoSpec, adapter: DemoAdapter) -> None: + del spec, adapter + + def validate_session( + self, + *, + spec: DemoSpec, + scenario: "PreparedScenario", + adapter: DemoAdapter, + provider: object, + ) -> None: + del spec, scenario, adapter, provider + + def create_run_context( + self, + *, + spec: DemoSpec, + adapter: DemoAdapter, + host: RuntimeHost, + model_warmup_plan: ModelWarmupPlan, + ) -> RunContext: + del spec, adapter + return RunContext( + host=host, + run_metrics=self._run_metrics, + admission=SingleSessionAdmissionPolicy( + health_check=lambda: host.is_healthy + ), + model_warmup_plan=model_warmup_plan, + ) + + def create_session_edges( + self, + *, + context: RunContext, + spec: DemoSpec, + scenario: "PreparedScenario", + provider: object, + adapter: DemoAdapter, + ) -> SessionEdges: + del provider, adapter + return SessionEdges( + input_source=_ReplayBatchInputSource( + scenario=scenario, + request_state=self._request_state, + ), + output_sink=self._output_sink_factory(spec.output), + cleanup_tasks=context.cleanup_tasks, + error_policy=( + Mp4ErrorPolicy() if spec.output.mode == "mp4" else NullErrorPolicy() + ), + ) + + def select_driver(self) -> BatchSessionDriver: + return BatchSessionDriver() + + +class _ReplayProviderAdapter: + def __init__( + self, + *, + adapter: DemoAdapter, + mapping: InputMapping, + request_state: "_ReplayStepRequestState", + ) -> None: + self._adapter = adapter + self._mapping = mapping + self._request_state = request_state + + @property + def model_id(self) -> str: + return self._adapter.model_id + + @property + def inference_input_schema(self) -> InferenceInputSchema: + return self._adapter.inference_input_schema + + @property + def canonical_input_schema(self) -> CanonicalInputSchema | None: + return self._adapter.canonical_input_schema + + def default_input_mapping(self) -> InputMapping | None: + return self._adapter.default_input_mapping() + + def supported_input_modes(self) -> tuple[str, ...]: + return self._adapter.supported_input_modes() + + def supported_output_modes(self) -> tuple[str, ...]: + return self._adapter.supported_output_modes() + + def validate_config(self, config: InferenceConfig) -> None: + self._adapter.validate_config(config) + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + return self._adapter.create_runtime(config) + + def prepare_scenario(self, spec: DemoSpec) -> "PreparedScenario": + return self._adapter.prepare_scenario(spec) + + def create_model_input_provider( + self, + spec: DemoSpec, + scenario: "PreparedScenario", + ) -> object: + create_provider = getattr(self._adapter, "create_model_input_provider", None) + if callable(create_provider): + return create_provider(spec, scenario) + return _ReplayMappingModelInputProvider( + adapter=self._adapter, + scenario=scenario, + mapping=self._mapping, + request_state=self._request_state, + ) + + +class _ReplayRuntimeAdapter: + def __init__( + self, + *, + runtime: InferenceRuntime, + request_state: "_ReplayStepRequestState", + ) -> None: + self._runtime = runtime + self._request_state = request_state + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + return _ReplaySessionAdapter( + session=self._runtime.start_session(inputs), + request_state=self._request_state, + ) + + def close(self) -> None: + self._runtime.close() + + +class _ReplaySessionAdapter: + def __init__( + self, + *, + session: InferenceSession, + request_state: "_ReplayStepRequestState", + ) -> None: + self._session = session + self._request_state = request_state + + def next_step_requirements(self) -> StepRequirements | None: + next_requirements = getattr(self._session, "next_step_requirements", None) + if callable(next_requirements): + value = next_requirements() + self._request_state.clear() + return value + + request = self._session.next_step_request() + if request is None: + self._request_state.clear() + return None + self._request_state.store(request) + return step_requirements_from_request( + request, + allow_user_input_window=True, + ) + + def next_step_request(self) -> StepRequest | None: + return self._session.next_step_request() + + def step(self, inputs: InferenceInput) -> StepResult: + return self._session.step(inputs) + + def reset(self, inputs: InferenceInput | None = None) -> None: + self._session.reset(inputs) + + def close(self) -> None: + self._session.close() + + +class _ReplayStepRequestState: + def __init__(self) -> None: + self._request: StepRequest | None = None + + def store(self, request: StepRequest) -> None: + self._request = request + + def request_for_window(self, step_index: int) -> StepRequest | None: + request = self._request + if request is None: + return None + if request.step_index != step_index: + raise RuntimeError( + "Replay input source request mismatch: " + f"expected step {request.step_index}, got {step_index}." + ) + return request + + def consume_for_step(self, request: StepRequirements) -> StepRequest: + legacy_request = self.request_for_window(request.step_index) + if legacy_request is not None: + self._request = None + return legacy_request + return StepRequest( + step_index=request.step_index, + inference_input_schema=request.inference_input_schema, + metadata=request.metadata, + ) + + def clear(self) -> None: + self._request = None + + +class _ReplayBatchInputSource: + is_finite = True + is_deterministic = True + + def __init__( + self, + *, + scenario: "PreparedScenario", + request_state: _ReplayStepRequestState, + ) -> None: + self.user_input_schema = scenario.source_schema + self._user_inputs = scenario.user_inputs + self._request_state = request_state + + def is_finished(self) -> bool: + return False + + def next_window(self, request: StepRequirements) -> UserInputWindow: + legacy_request = self._request_state.request_for_window(request.step_index) + window = ( + legacy_request.user_input_window if legacy_request is not None else None + ) + if window is None: + window = _all_user_inputs_window(self._user_inputs) + return UserInputWindow( + start_s=window.start_s, + end_s=window.end_s, + inputs=self._user_inputs, + ) + + +class _ReplayMappingModelInputProvider: + def __init__( + self, + *, + adapter: DemoAdapter, + scenario: "PreparedScenario", + mapping: InputMapping, + request_state: _ReplayStepRequestState, + ) -> None: + self.capabilities = ProviderCapabilities( + supports_recorded_input=True, + deterministic_given_inputs=True, + user_input_schema=scenario.source_schema, + inference_input_schema=adapter.inference_input_schema, + ) + self._scenario = scenario + self._mapping = mapping + self._request_state = request_state + self._step_base_inputs = InferenceInput( + step=scenario.initial_inputs.step, + metadata=scenario.initial_inputs.metadata, + ) + + def prepare_initial_input(self) -> InferenceInput: + self._scenario.canonicalizer.reset() + return self._mapping.map_global_conditioning_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=self._scenario.initial_inputs, + ) + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + legacy_request = self._request_state.consume_for_step(request) + canonical_inputs = self._scenario.canonicalizer.canonicalize( + self._scenario.user_inputs, + window=TimeWindow(start_s=user_window.start_s, end_s=user_window.end_s), + source_schema=self._scenario.source_schema, + ) + return PreparedStep( + inference_input=self._mapping.map_step_inputs( + canonical_inputs=canonical_inputs, + inference_input=self._step_base_inputs, + request=legacy_request, + ) + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self._scenario.canonicalizer.reset() + + def close(self) -> None: + return None + + +def _validate_replay_mapping( + *, + adapter: DemoAdapter, + config: InferenceConfig, + mapping: InputMapping, + source_schema: UserInputSchema, + canonicalizer: InputCanonicalizer, +) -> None: + adapter.validate_config(config) + canonical_schema = canonicalizer.canonical_schema(source_schema) + if isinstance(mapping, DeclaresMappingSchema): + compatibility = check_mapping_compatibility( + canonical_schema=canonical_schema, + inference_input_schema=adapter.inference_input_schema, + mapping_schema=mapping.mapping_schema, + ) + compatibility.raise_if_incompatible() + mapping.validate( + canonical_schema=canonical_schema, + inference_input_schema=adapter.inference_input_schema, + ) + + +def _require_config(spec: DemoSpec) -> InferenceConfig: + if spec.config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + return spec.config + + +def _all_user_inputs_window(user_inputs: UserInputs) -> TimeWindow: + if not user_inputs.events: + return TimeWindow(start_s=0.0, end_s=3600.0) + return TimeWindow( + start_s=0.0, + end_s=max( + 3600.0, + math.nextafter(user_inputs.events[-1].timestamp_s, math.inf), + ), + ) + + +def _default_inference_session_runner() -> InferenceSessionRunner: + from flashdreams.runtime.runner import run_inference_session + + return run_inference_session def _require_supported_mode( @@ -88,6 +623,7 @@ def _require_supported_mode( __all__ = [ "InferenceSessionRunner", + "OutputSinkFactory", "OutputTargetFactory", "run_replay_demo", ] diff --git a/flashdreams/flashdreams/runtime/demo/run_modes.py b/flashdreams/flashdreams/runtime/demo/run_modes.py new file mode 100644 index 000000000..cb93e4e7f --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/run_modes.py @@ -0,0 +1,535 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run/session result and policy helpers for demo session drivers.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from threading import Lock +from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable + +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.metrics import ( + InMemoryMetricsRecorder, + MetricsRecorder, + MetricsSnapshot, +) +from flashdreams.runtime.output import OutputArtifact + +from .host import ModelWarmupPlan, WarmupSessionInputs +from .outputs import OutputSink + +if TYPE_CHECKING: + from .host import RuntimeHost + from .pipeline import StepPipeline + from .session_inputs import InputSource, ModelInputProvider + from .spec import DemoAdapter, DemoSpec, PreparedScenario + from .timing import ActivationPolicy, DeterministicClock, RealtimeClock + +SessionStatus = Literal[ + "completed", + "failed", + "skipped", + "cancelled", + "rejected", + "not_activated", +] + +DriverStatus = Literal[ + "completed", + "failed", + "skipped", + "cancelled", + "not_activated", +] + + +@dataclass(frozen=True, kw_only=True, slots=True) +class RunResult: + """Outcome of one demo session.""" + + __hash__ = None + + status: SessionStatus + artifacts: Sequence[OutputArtifact] = () + metrics: MetricsSnapshot | None = None + reason: str | None = None + error: Exception | None = None + + @classmethod + def rejected(cls, reason: str) -> "RunResult": + """Admission refused the session. The only no-session result helper.""" + return cls(status="rejected", reason=reason) + + def __post_init__(self) -> None: + object.__setattr__(self, "artifacts", tuple(self.artifacts)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class RunSummary: + """Summary for a run context after one or more sessions.""" + + metrics: MetricsSnapshot + sessions: Sequence[RunResult] = () + + def __post_init__(self) -> None: + object.__setattr__(self, "sessions", tuple(self.sessions)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class ErrorAction: + """Driver policy decision for an operational error.""" + + close_session: bool = True + drop_chunk: bool = False + continue_next_scenario: bool = False + result_status: Literal["completed", "failed", "skipped"] = "failed" + + +class DefaultErrorPolicy: + """Default policy: operational errors fail the current session.""" + + def handle_setup_error(self, exc: Exception) -> ErrorAction: + del exc + return ErrorAction(result_status="failed") + + def handle(self, exc: Exception) -> ErrorAction: + del exc + return ErrorAction(result_status="failed") + + +@runtime_checkable +class ErrorPolicy(Protocol): + """Maps driver-observed exceptions to session outcomes.""" + + def handle_setup_error(self, exc: Exception) -> ErrorAction: ... + + def handle(self, exc: Exception) -> ErrorAction: ... + + +class Mp4ErrorPolicy(DefaultErrorPolicy): + """Abort MP4 sessions on setup or step errors.""" + + +class NullErrorPolicy(DefaultErrorPolicy): + """Abort headless/null sessions on setup or step errors.""" + + +class NativeWindowErrorPolicy(DefaultErrorPolicy): + """Abort native-window sessions unless a future UI policy overrides it.""" + + +class BenchmarkErrorPolicy(DefaultErrorPolicy): + """Close failed scenarios while letting benchmark loops continue.""" + + def handle_setup_error(self, exc: Exception) -> ErrorAction: + del exc + return ErrorAction(result_status="failed", continue_next_scenario=True) + + def handle(self, exc: Exception) -> ErrorAction: + del exc + return ErrorAction(result_status="failed", continue_next_scenario=True) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class WebRTCErrorPolicy: + """Drop configured recoverable realtime errors, otherwise close the session.""" + + recoverable_exception_types: tuple[type[Exception], ...] = () + + def handle_setup_error(self, exc: Exception) -> ErrorAction: + del exc + return ErrorAction(result_status="failed") + + def handle(self, exc: Exception) -> ErrorAction: + if self.recoverable_exception_types and isinstance( + exc, self.recoverable_exception_types + ): + return ErrorAction( + close_session=False, + drop_chunk=True, + result_status="failed", + ) + return ErrorAction(result_status="failed") + + +@dataclass(frozen=True, kw_only=True, slots=True) +class RunModeCapabilities: + """Run-mode requirements and output/transport capabilities.""" + + realtime: bool = False + requires_finite_input: bool = False + supports_backpressure: bool = False + supports_interactive_events: bool = False + supports_artifacts: bool = False + + +SessionMetricsRecorder = MetricsRecorder +InMemorySessionMetricsRecorder = InMemoryMetricsRecorder + + +class NoopTransportService: + """Idempotent placeholder transport for batch sessions.""" + + def __init__(self) -> None: + self.closed = False + + def is_active(self) -> bool: + return not self.closed + + def close(self) -> None: + self.closed = True + + +@runtime_checkable +class TransportService(Protocol): + """Per-session transport lifecycle hook.""" + + def is_active(self) -> bool: ... + + def close(self) -> None: ... + + +@runtime_checkable +class SessionReservation(Protocol): + """Admission reservation for one session.""" + + def release(self) -> None: ... + + +class SingleSessionAdmissionPolicy: + """Atomic single-session admission policy.""" + + def __init__(self, *, health_check: Any | None = None) -> None: + self._lock = Lock() + self._reserved = False + self._health_check = health_check + + def try_reserve(self) -> SessionReservation | None: + with self._lock: + if self._reserved or not self._is_healthy(): + return None + self._reserved = True + return _SingleSessionReservation(self) + + def _release(self) -> None: + with self._lock: + self._reserved = False + + def _is_healthy(self) -> bool: + if self._health_check is None: + return True + return bool(self._health_check()) + + +class _SingleSessionReservation: + def __init__(self, policy: SingleSessionAdmissionPolicy) -> None: + self._policy = policy + self._released = False + self.release_count = 0 + + def release(self) -> None: + if self._released: + return + self._released = True + self.release_count += 1 + self._policy._release() + + +@runtime_checkable +class AdmissionPolicy(Protocol): + """Atomically reserves session capacity or rejects.""" + + def try_reserve(self) -> SessionReservation | None: ... + + +@runtime_checkable +class SessionDriver(Protocol): + """Synchronous one-session driver selected by a run mode.""" + + def run_one_session( + self, + *, + host: "RuntimeHost", + provider: "ModelInputProvider", + session_edges: "SessionEdges", + pipeline: "StepPipeline", + ) -> RunResult: ... + + +@runtime_checkable +class AsyncSessionDriver(Protocol): + """Async one-session driver selected by realtime run modes.""" + + async def run_one_session( + self, + *, + host: "RuntimeHost", + provider: "ModelInputProvider", + session_edges: "SessionEdges", + pipeline: "StepPipeline", + ) -> RunResult: ... + + +@dataclass(slots=True) +class RunContext: + """Run-scoped services shared by one or more demo sessions.""" + + host: "RuntimeHost" + run_metrics: SessionMetricsRecorder + admission: AdmissionPolicy + model_warmup_plan: ModelWarmupPlan = field(default_factory=ModelWarmupPlan) + services: Mapping[str, object] = field(default_factory=dict) + cleanup_tasks: set[asyncio.Task[RunResult]] = field(default_factory=set) + + def __post_init__(self) -> None: + self.services = freeze_mapping(self.services) + + def close(self) -> RunSummary: + if self.cleanup_tasks: + raise RuntimeError( + "Pending session cleanup tasks; async runs must await close_async()." + ) + for service in self.services.values(): + close = getattr(service, "close", None) + if callable(close): + try: + close() + except Exception as exc: + self.run_metrics.record_cleanup_error(exc) + return RunSummary( + metrics=self.run_metrics.close(), + sessions=tuple(getattr(self.run_metrics, "sessions", ())), + ) + + async def close_async(self) -> RunSummary: + while self.cleanup_tasks: + pending = tuple(self.cleanup_tasks) + await asyncio.gather(*pending, return_exceptions=True) + self.cleanup_tasks.difference_update(pending) + return self.close() + + +@dataclass(slots=True) +class SessionEdges: + """Per-session input/output/policy bundle consumed by drivers.""" + + input_source: "InputSource" + output_sink: OutputSink + cleanup_tasks: set[asyncio.Task[RunResult]] + metrics: SessionMetricsRecorder = field( + default_factory=InMemorySessionMetricsRecorder + ) + error_policy: ErrorPolicy = field(default_factory=DefaultErrorPolicy) + transport: TransportService = field(default_factory=NoopTransportService) + clock: "RealtimeClock | DeterministicClock | None" = None + activation: "ActivationPolicy | None" = None + _closed_result: RunResult | None = field(default=None, init=False, repr=False) + + @property + def is_closed(self) -> bool: + """Return whether ``close_result(...)`` has already finalized this session.""" + return self._closed_result is not None + + def record_cleanup_error(self, exc: Exception) -> None: + """Record a cleanup error without letting metrics failures block teardown.""" + try: + self.metrics.record_cleanup_error(exc) + except Exception: + return + + def record_orphaned_cleanup(self, exc: Exception) -> None: + """Record timed-out worker cleanup without blocking teardown.""" + try: + self.metrics.record_orphaned_cleanup(exc) + except Exception: + return + + def close_result( + self, + *, + status: DriverStatus = "completed", + reason: str | None = None, + error: Exception | None = None, + ) -> RunResult: + """Idempotently close output, transport, and metrics once.""" + if self._closed_result is not None: + return self._closed_result + + artifacts: Sequence[OutputArtifact] = () + try: + artifacts = tuple(self.output_sink.close()) + except Exception as exc: + self.record_cleanup_error(exc) + try: + self.transport.close() + except Exception as exc: + self.record_cleanup_error(exc) + try: + metrics = self.metrics.close() + except Exception as exc: + metrics = MetricsSnapshot(errors=(f"metrics.close failed: {exc}",)) + self._closed_result = RunResult( + status=status, + artifacts=artifacts, + metrics=metrics, + reason=reason, + error=error, + ) + return self._closed_result + + +@runtime_checkable +class RunMode(Protocol): + """Run/session construction strategy consumed by shared helpers.""" + + name: str + capabilities: RunModeCapabilities + + def validate_run( + self, + *, + spec: "DemoSpec", + adapter: "DemoAdapter", + ) -> None: ... + + def validate_session( + self, + *, + spec: "DemoSpec", + scenario: "PreparedScenario", + adapter: "DemoAdapter", + provider: "ModelInputProvider", + ) -> None: ... + + def create_run_context( + self, + *, + spec: "DemoSpec", + adapter: "DemoAdapter", + host: "RuntimeHost", + model_warmup_plan: ModelWarmupPlan, + ) -> RunContext: ... + + def create_session_edges( + self, + *, + context: RunContext, + spec: "DemoSpec", + scenario: "PreparedScenario", + provider: "ModelInputProvider", + adapter: "DemoAdapter", + ) -> SessionEdges: ... + + def select_driver(self) -> SessionDriver | AsyncSessionDriver: ... + + +@runtime_checkable +class RunModeWarmup(Protocol): + """Optional run-mode warmup for output or transport services.""" + + def warmup_context( + self, + *, + context: RunContext, + spec: "DemoSpec", + scenario: "PreparedScenario", + adapter: "DemoAdapter", + ) -> None: ... + + +def build_model_warmup_plan( + *, + host: "RuntimeHost", + adapter: "DemoAdapter", + spec: "DemoSpec", + scenario: "PreparedScenario", +) -> ModelWarmupPlan: + """Build a host-owned warmup plan through the model-affine worker.""" + + create_sessions = getattr(adapter, "create_model_warmup_sessions", None) + if create_sessions is None: + return ModelWarmupPlan() + if not callable(create_sessions): + raise TypeError( + "Demo adapter create_model_warmup_sessions attribute must be callable." + ) + sessions = host.call(create_sessions, spec, scenario) + return ModelWarmupPlan(sessions=_coerce_warmup_sessions(sessions)) + + +def warmup_run_context( + *, + context: RunContext, + spec: "DemoSpec", + scenario: "PreparedScenario", + adapter: "DemoAdapter", + run_mode: object, +) -> None: + """Run model warmup, then optional output/transport warmup for a context.""" + + context.host.warmup(context.model_warmup_plan) + warmup_context = getattr(run_mode, "warmup_context", None) + if warmup_context is None: + return + if not callable(warmup_context): + raise TypeError("RunMode.warmup_context attribute must be callable.") + warmup_context( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + ) + + +def _coerce_warmup_sessions(value: object) -> tuple[WarmupSessionInputs, ...]: + if not isinstance(value, Sequence): + raise TypeError( + "Demo adapter create_model_warmup_sessions(...) must return a sequence " + f"of WarmupSessionInputs, got {type(value).__name__}." + ) + sessions: list[WarmupSessionInputs] = [] + for session in value: + if not isinstance(session, WarmupSessionInputs): + raise TypeError( + "Demo adapter create_model_warmup_sessions(...) must return only " + f"WarmupSessionInputs, got {type(session).__name__}." + ) + sessions.append(session) + return tuple(sessions) + + +__all__ = [ + "AdmissionPolicy", + "AsyncSessionDriver", + "BenchmarkErrorPolicy", + "DefaultErrorPolicy", + "DriverStatus", + "ErrorAction", + "ErrorPolicy", + "InMemorySessionMetricsRecorder", + "MetricsSnapshot", + "Mp4ErrorPolicy", + "NativeWindowErrorPolicy", + "NoopTransportService", + "NullErrorPolicy", + "RunContext", + "RunMode", + "RunModeCapabilities", + "RunModeWarmup", + "RunResult", + "RunSummary", + "SessionEdges", + "SessionDriver", + "SessionMetricsRecorder", + "SessionReservation", + "SessionStatus", + "SingleSessionAdmissionPolicy", + "TransportService", + "WebRTCErrorPolicy", + "build_model_warmup_plan", + "warmup_run_context", +] diff --git a/flashdreams/flashdreams/runtime/demo/session_inputs.py b/flashdreams/flashdreams/runtime/demo/session_inputs.py new file mode 100644 index 000000000..dda42995c --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/session_inputs.py @@ -0,0 +1,184 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Input-source and model-input-provider contracts for demo sessions.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.inputs import ( + InferenceInput, + InferenceInputSchema, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.types import StepRequirements + +if TYPE_CHECKING: + from .timing import RealtimeClock, RealtimeWindowResult + + +@dataclass(frozen=True, kw_only=True, slots=True) +class ProviderCapabilities: + """Model-provider capabilities used to validate run-mode compatibility.""" + + supports_realtime_clock: bool = False + supports_recorded_input: bool = False + supports_reset: bool = False + deterministic_given_inputs: bool = False + user_input_schema: UserInputSchema = field(default_factory=UserInputSchema) + inference_input_schema: InferenceInputSchema = field( + default_factory=InferenceInputSchema + ) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class ControlDecision: + """Provider-authored control request for the current session.""" + + reset: bool = False + close_session: bool = False + reset_input: InferenceInput | None = None + provider_already_reset: bool = False + reason: str | None = None + + def __post_init__(self) -> None: + if self.reason is not None and not self.reason.strip(): + raise ValueError("ControlDecision.reason must be non-empty when set.") + + +@dataclass(frozen=True, kw_only=True, slots=True) +class UserInputWindow: + """User/app inputs selected by a driver for one model step.""" + + __hash__ = None + + start_s: float + end_s: float + frame_times: Sequence[float] = () + inputs: UserInputs = field(default_factory=UserInputs) + control: ControlDecision | None = None + metadata: Mapping[str, object] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not math.isfinite(self.start_s) or self.start_s < 0: + raise ValueError("UserInputWindow.start_s must be finite and >= 0.") + if not math.isfinite(self.end_s) or self.end_s < self.start_s: + raise ValueError("UserInputWindow.end_s must be finite and >= start_s.") + previous = -math.inf + for frame_time in self.frame_times: + if not math.isfinite(float(frame_time)): + raise ValueError("UserInputWindow.frame_times must be finite.") + if float(frame_time) < previous: + raise ValueError( + "UserInputWindow.frame_times must be sorted in ascending order." + ) + previous = float(frame_time) + object.__setattr__( + self, "frame_times", tuple(float(t) for t in self.frame_times) + ) + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class PreparedStep: + """Model-facing input plus optional provider-authored control decision.""" + + __hash__ = None + + inference_input: InferenceInput | None = None + control: ControlDecision = field(default_factory=ControlDecision) + + +@runtime_checkable +class InputSource(Protocol): + """Facts common to every demo session input source.""" + + is_finite: bool + is_deterministic: bool + user_input_schema: UserInputSchema + + def is_finished(self) -> bool: + """Return whether the driver should stop requesting windows.""" + ... + + +@runtime_checkable +class BatchInputSource(InputSource, Protocol): + """Finite input source consumed by the batch driver.""" + + def next_window(self, request: StepRequirements) -> UserInputWindow: + """Return the next batch input window for ``request``.""" + ... + + +@runtime_checkable +class RealtimeInputSource(InputSource, Protocol): + """Realtime input source consumed by a future realtime driver.""" + + async def next_realtime_window( + self, + *, + request: StepRequirements, + clock: "RealtimeClock", + ) -> "RealtimeWindowResult": + """Return the next realtime window result. + + The concrete realtime result shape lands with the realtime clock phase. + Keeping this protocol separate now prevents batch sources from stubbing + async behavior they never serve. + """ + ... + + +@runtime_checkable +class ModelInputProvider(Protocol): + """Model-owned conversion from user windows into model-facing inputs.""" + + capabilities: ProviderCapabilities + + def prepare_initial_input(self) -> InferenceInput: + """Prepare session-global model inputs.""" + ... + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + """Prepare one model step from a driver-owned user input window.""" + ... + + def reset(self, inputs: InferenceInput | None = None) -> None: + """Reset provider-owned session state. + + Implementations must be idempotent so driver cleanup and reset control + paths can safely converge after failures. + """ + ... + + def close(self) -> None: + """Release provider-owned resources. + + Implementations must be idempotent and tolerate cleanup after partial + setup or earlier reset failures. + """ + ... + + +__all__ = [ + "BatchInputSource", + "ControlDecision", + "InputSource", + "ModelInputProvider", + "PreparedStep", + "ProviderCapabilities", + "RealtimeInputSource", + "UserInputWindow", +] diff --git a/flashdreams/flashdreams/runtime/demo/spec.py b/flashdreams/flashdreams/runtime/demo/spec.py index 6ba652f38..d7a6a309b 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 Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass, field, replace from pathlib import Path from typing import Any, Literal, Protocol, TypeAlias @@ -18,6 +18,8 @@ from flashdreams.runtime.interfaces import ModelAdapter from flashdreams.runtime.mapping import InputMapping +from .host import WarmupSessionInputs + @dataclass(frozen=True, kw_only=True, slots=True) class NullOutputSpec: @@ -170,9 +172,22 @@ def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: ... +class ModelWarmupAdapter(Protocol): + """Optional adapter hook for model-affine runtime warmup inputs.""" + + def create_model_warmup_sessions( + self, + spec: DemoSpec, + scenario: PreparedScenario, + ) -> Sequence[WarmupSessionInputs]: + """Return temporary synthetic or loopback sessions for model warmup.""" + ... + + __all__ = [ "DemoAdapter", "DemoSpec", + "ModelWarmupAdapter", "Mp4OutputSpec", "NullOutputSpec", "OutputSpec", diff --git a/flashdreams/flashdreams/runtime/demo/timing.py b/flashdreams/flashdreams/runtime/demo/timing.py new file mode 100644 index 000000000..d48e1e213 --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/timing.py @@ -0,0 +1,383 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Realtime activation, clock, and input-window primitives for demo run modes.""" + +from __future__ import annotations + +import asyncio +import math +import time +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass, field +from typing import Literal, Protocol, runtime_checkable + +from flashdreams.runtime.inputs import UserInputs, UserInputSchema +from flashdreams.runtime.types import StepRequirements + +from .session_inputs import UserInputWindow + +CatchUpPolicy = Literal["drop", "fold", "compress"] + +SPARSE_KEY_SEGMENTS_METADATA_KEY = "sparse_key_segments" + + +@dataclass(frozen=True, kw_only=True, slots=True) +class CatchUpDecision: + """How a realtime clock bounded stale virtual input time.""" + + skipped_s: float = 0.0 + skipped_windows: int = 0 + input_policy: CatchUpPolicy | None = None + reason: str | None = None + + def __post_init__(self) -> None: + if not math.isfinite(self.skipped_s) or self.skipped_s < 0.0: + raise ValueError("CatchUpDecision.skipped_s must be finite and >= 0.") + if self.skipped_windows < 0: + raise ValueError("CatchUpDecision.skipped_windows must be >= 0.") + if self.input_policy not in {None, "drop", "fold", "compress"}: + raise ValueError( + f"Unsupported catch-up input_policy={self.input_policy!r}." + ) + if self.reason is not None and not self.reason.strip(): + raise ValueError("CatchUpDecision.reason must be non-empty when set.") + + +@dataclass(frozen=True, kw_only=True, slots=True) +class RealtimeWindowResult: + """Realtime input window plus any catch-up decision that preceded it.""" + + window: UserInputWindow + catch_up: CatchUpDecision = field(default_factory=CatchUpDecision) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class ActivationResult: + """Result of waiting for a realtime activation gate.""" + + activated: bool + reason: str | None = None + + def __post_init__(self) -> None: + if self.reason is not None and not self.reason.strip(): + raise ValueError("ActivationResult.reason must be non-empty when set.") + + +@runtime_checkable +class DeterministicClock(Protocol): + """Clock facts for finite deterministic run modes.""" + + is_realtime: bool + is_deterministic: bool + + +@runtime_checkable +class RealtimeClock(Protocol): + """Realtime virtual clock used by realtime drivers and input sources.""" + + is_realtime: bool + is_deterministic: bool + + def now(self) -> float: ... + + def anchor(self, wall_time_s: float) -> None: ... + + async def wait_until_window_end(self, end_s: float) -> None: ... + + async def apply_backpressure(self, requested_s: float) -> None: ... + + def catch_up( + self, + *, + request: StepRequirements, + max_lag_s: float, + policy: CatchUpPolicy, + ) -> CatchUpDecision: ... + + +@runtime_checkable +class ActivationPolicy(Protocol): + """Wait until a realtime session should start generating.""" + + timeout_s: float | None + + async def wait_until_active( + self, + clock: RealtimeClock | DeterministicClock, + ) -> ActivationResult: ... + + +@runtime_checkable +class ActivationSignal(Protocol): + """Event-like object accepted by ``SignalActivationPolicy``.""" + + def is_set(self) -> bool: ... + + async def wait(self) -> object: ... + + +@dataclass(slots=True) +class AlwaysActiveActivationPolicy: + """Activation policy for batch/null modes or already-ready realtime modes.""" + + timeout_s: float | None = None + anchor_clock: bool = False + + async def wait_until_active( + self, + clock: RealtimeClock | DeterministicClock, + ) -> ActivationResult: + _anchor_if_realtime(clock, anchor=self.anchor_clock) + return ActivationResult(activated=True) + + +@dataclass(slots=True) +class SignalActivationPolicy: + """Activate when any supplied signal fires, with optional timeout.""" + + signals: Sequence[ActivationSignal] + timeout_s: float | None = None + timeout_reason: str = "activation timed out" + anchor_clock: bool = True + + def __post_init__(self) -> None: + if not self.signals: + raise ValueError("SignalActivationPolicy.signals must be non-empty.") + self.signals = tuple(self.signals) + if self.timeout_s is not None and self.timeout_s <= 0.0: + raise ValueError("SignalActivationPolicy.timeout_s must be > 0 when set.") + if not self.timeout_reason.strip(): + raise ValueError("SignalActivationPolicy.timeout_reason must be non-empty.") + + async def wait_until_active( + self, + clock: RealtimeClock | DeterministicClock, + ) -> ActivationResult: + if any(signal.is_set() for signal in self.signals): + _anchor_if_realtime(clock, anchor=self.anchor_clock) + return ActivationResult(activated=True) + + tasks = [asyncio.create_task(signal.wait()) for signal in self.signals] + try: + done, pending = await asyncio.wait( + tasks, + timeout=self.timeout_s, + return_when=asyncio.FIRST_COMPLETED, + ) + if not done: + return ActivationResult( + activated=False, + reason=self.timeout_reason, + ) + for task in done: + task.result() + _anchor_if_realtime(clock, anchor=self.anchor_clock) + return ActivationResult(activated=True) + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + +SparseKeySegment = tuple[float, float, frozenset[str]] + + +class _RealtimeTimeline(Protocol): + dt: float + next_chunk_start_v: float + + +class _SparseInputResampler(_RealtimeTimeline, Protocol): + def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: ... + + def reset(self, *, start_v: float) -> None: ... + + def sample_chunk( + self, + num_frames: int, + ) -> tuple[Sequence[SparseKeySegment], Sequence[float]]: ... + + +@dataclass(slots=True) +class ResamplerRealtimeClock: + """Realtime clock that reuses ``KeyboardResampler``'s virtual timeline.""" + + resampler: _RealtimeTimeline + now_fn: Callable[[], float] = time.monotonic + sleep_fn: Callable[[float], Awaitable[None]] = asyncio.sleep + is_realtime: bool = True + is_deterministic: bool = False + _pending_backpressure_s: float = field(default=0.0, init=False, repr=False) + + @property + def pending_backpressure_s(self) -> float: + return self._pending_backpressure_s + + def now(self) -> float: + return float(self.now_fn()) + + def anchor(self, wall_time_s: float) -> None: + if not math.isfinite(wall_time_s): + raise ValueError("wall_time_s must be finite.") + self.resampler.next_chunk_start_v = float(wall_time_s) + self._pending_backpressure_s = 0.0 + + async def wait_until_window_end(self, end_s: float) -> None: + if not math.isfinite(end_s): + raise ValueError("end_s must be finite.") + delay_s = float(end_s) - self.now() + if delay_s > 0.0: + await self.sleep_fn(delay_s) + + async def apply_backpressure(self, requested_s: float) -> None: + if not math.isfinite(requested_s) or requested_s < 0.0: + raise ValueError("requested_s must be finite and >= 0.") + self._pending_backpressure_s += float(requested_s) + + def catch_up( + self, + *, + request: StepRequirements, + max_lag_s: float, + policy: CatchUpPolicy, + ) -> CatchUpDecision: + if policy != "fold": + raise NotImplementedError( + f"Catch-up policy {policy!r} has no existing resampler analog yet." + ) + if not math.isfinite(max_lag_s) or max_lag_s < 0.0: + raise ValueError("max_lag_s must be finite and >= 0.") + + input_frame_count = input_frame_count_from_request(request) + chunk_duration_s = input_frame_count * float(self.resampler.dt) + if chunk_duration_s <= 0.0: + raise ValueError("Realtime resampler dt must produce a positive window.") + + effective_now_s = self.now() + self._pending_backpressure_s + self._pending_backpressure_s = 0.0 + current_start_s = float(self.resampler.next_chunk_start_v) + lag_s = effective_now_s - (current_start_s + chunk_duration_s) + if lag_s <= max_lag_s: + return CatchUpDecision() + + latest_start_s = effective_now_s - chunk_duration_s + if latest_start_s <= current_start_s: + return CatchUpDecision() + + skipped_s = latest_start_s - current_start_s + skipped_windows = max(1, math.ceil(skipped_s / chunk_duration_s)) + self.resampler.next_chunk_start_v = latest_start_s + return CatchUpDecision( + skipped_s=skipped_s, + skipped_windows=skipped_windows, + input_policy=policy, + reason="lag exceeded max_lag_s", + ) + + +@dataclass(slots=True) +class KeyboardRealtimeInputSource: + """Realtime input source backed by the existing keyboard resampler.""" + + resampler: _SparseInputResampler + max_lag_s: float | None = None + catch_up_policy: CatchUpPolicy = "fold" + is_finite: bool = False + is_deterministic: bool = False + user_input_schema: UserInputSchema = field(default_factory=UserInputSchema) + + def __post_init__(self) -> None: + if self.max_lag_s is not None and ( + not math.isfinite(self.max_lag_s) or self.max_lag_s < 0.0 + ): + raise ValueError( + "KeyboardRealtimeInputSource.max_lag_s must be finite and >= 0." + ) + if self.catch_up_policy != "fold": + raise NotImplementedError( + f"Catch-up policy {self.catch_up_policy!r} has no existing " + "KeyboardResampler analog yet." + ) + + def is_finished(self) -> bool: + return False + + def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: + self.resampler.on_edge(arrival_t=arrival_t, event=event, key=key) + + def reset(self, *, start_v: float) -> None: + self.resampler.reset(start_v=start_v) + + async def next_realtime_window( + self, + *, + request: StepRequirements, + clock: RealtimeClock, + ) -> RealtimeWindowResult: + input_frame_count = input_frame_count_from_request(request) + chunk_duration_s = input_frame_count * self.resampler.dt + window_end_s = self.resampler.next_chunk_start_v + chunk_duration_s + await clock.wait_until_window_end(window_end_s) + catch_up = clock.catch_up( + request=request, + max_lag_s=self.max_lag_s + if self.max_lag_s is not None + else chunk_duration_s, + policy=self.catch_up_policy, + ) + start_s = self.resampler.next_chunk_start_v + segments, frame_times = self.resampler.sample_chunk(input_frame_count) + end_s = self.resampler.next_chunk_start_v + window = UserInputWindow( + start_s=start_s, + end_s=end_s, + frame_times=tuple(frame_times), + inputs=UserInputs(), + metadata={SPARSE_KEY_SEGMENTS_METADATA_KEY: tuple(segments)}, + ) + return RealtimeWindowResult(window=window, catch_up=catch_up) + + +def input_frame_count_from_request(request: StepRequirements) -> int: + """Return the positive input frame count declared by a step requirement.""" + + value = request.input_frame_count + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError("StepRequirements.input_frame_count must be an integer.") + parsed = value + if parsed <= 0: + raise ValueError("StepRequirements.input_frame_count must be > 0.") + return parsed + + +def _anchor_if_realtime( + clock: RealtimeClock | DeterministicClock, + *, + anchor: bool, +) -> None: + if not anchor or not getattr(clock, "is_realtime", False): + return + now = getattr(clock, "now", None) + clock_anchor = getattr(clock, "anchor", None) + if callable(now) and callable(clock_anchor): + clock_anchor(float(now())) + + +__all__ = [ + "ActivationPolicy", + "ActivationResult", + "ActivationSignal", + "AlwaysActiveActivationPolicy", + "CatchUpDecision", + "CatchUpPolicy", + "DeterministicClock", + "KeyboardRealtimeInputSource", + "RealtimeClock", + "RealtimeWindowResult", + "ResamplerRealtimeClock", + "SPARSE_KEY_SEGMENTS_METADATA_KEY", + "SignalActivationPolicy", + "input_frame_count_from_request", +] diff --git a/flashdreams/flashdreams/runtime/demo/validation.py b/flashdreams/flashdreams/runtime/demo/validation.py new file mode 100644 index 000000000..c5058b148 --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/validation.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Capability resolution and validation for shared demo runs.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass + +from flashdreams.runtime.inputs import UserInputCapability, UserInputSchema + +from .run_modes import RunMode, RunModeCapabilities, SessionEdges +from .session_inputs import ( + BatchInputSource, + ModelInputProvider, + ProviderCapabilities, + RealtimeInputSource, +) +from .spec import DemoAdapter, DemoSpec + + +@dataclass(frozen=True, kw_only=True, slots=True) +class ResolvedRunCapabilities: + """Capabilities of one concrete provider/run-mode/session-edges pairing.""" + + finite: bool + deterministic: bool + realtime: bool + resettable: bool + produces_artifacts: bool + + +def resolve_run_capabilities( + *, + spec: DemoSpec, + provider: ModelInputProvider, + session_edges: SessionEdges, +) -> ResolvedRunCapabilities: + """Resolve concrete run capabilities from provider, edges, and config.""" + + provider_capabilities = _provider_capabilities(provider) + clock = session_edges.clock + realtime = bool(getattr(clock, "is_realtime", False)) + deterministic_clock = ( + bool(getattr(clock, "is_deterministic", False)) if clock is not None else True + ) + config = spec.config + seeded = config is not None and config.seed is not None + return ResolvedRunCapabilities( + finite=bool(session_edges.input_source.is_finite), + deterministic=( + provider_capabilities.deterministic_given_inputs + and bool(session_edges.input_source.is_deterministic) + and deterministic_clock + and seeded + ), + realtime=realtime, + resettable=provider_capabilities.supports_reset, + produces_artifacts=bool(session_edges.output_sink.produces_artifacts), + ) + + +def validate_resolved_run( + *, + spec: DemoSpec, + adapter: DemoAdapter, + provider: ModelInputProvider, + run_mode: RunMode, + session_edges: SessionEdges, + resolved: ResolvedRunCapabilities, +) -> None: + """Reject structurally incompatible provider/input/run-mode combinations.""" + + del spec, adapter + provider_capabilities = _provider_capabilities(provider) + run_mode_capabilities = _run_mode_capabilities(run_mode) + _validate_input_source_shape( + run_mode_capabilities=run_mode_capabilities, + session_edges=session_edges, + resolved=resolved, + ) + _validate_provider_modes( + provider_capabilities=provider_capabilities, + run_mode_capabilities=run_mode_capabilities, + resolved=resolved, + ) + _validate_user_input_schema( + provider_schema=provider_capabilities.user_input_schema, + source_schema=_input_source_user_input_schema(session_edges.input_source), + ) + + +def _validate_input_source_shape( + *, + run_mode_capabilities: RunModeCapabilities, + session_edges: SessionEdges, + resolved: ResolvedRunCapabilities, +) -> None: + input_source = session_edges.input_source + if run_mode_capabilities.realtime: + if not isinstance(input_source, RealtimeInputSource): + raise ValueError("Realtime run modes require a RealtimeInputSource.") + if session_edges.clock is None or not resolved.realtime: + raise ValueError("Realtime run modes require a realtime clock.") + return + if not isinstance(input_source, BatchInputSource): + raise ValueError("Batch run modes require a BatchInputSource.") + if resolved.realtime: + raise ValueError("Batch run modes cannot use a realtime clock.") + + +def _validate_provider_modes( + *, + provider_capabilities: ProviderCapabilities, + run_mode_capabilities: RunModeCapabilities, + resolved: ResolvedRunCapabilities, +) -> None: + if run_mode_capabilities.realtime and not ( + provider_capabilities.supports_realtime_clock + ): + raise ValueError("Provider does not support realtime input.") + if run_mode_capabilities.requires_finite_input: + if not resolved.finite: + raise ValueError("Run mode requires finite input.") + if not provider_capabilities.supports_recorded_input: + raise ValueError("Provider does not support recorded input.") + if resolved.produces_artifacts and not run_mode_capabilities.supports_artifacts: + raise ValueError("Run mode does not support artifact output.") + if ( + run_mode_capabilities.supports_interactive_events + and not provider_capabilities.supports_realtime_clock + ): + raise ValueError("Interactive run mode requires realtime provider support.") + + +def _validate_user_input_schema( + *, + provider_schema: UserInputSchema, + source_schema: UserInputSchema, +) -> None: + missing = _missing_capabilities( + required=provider_schema.declared_capabilities(), + provided=source_schema, + ) + if missing: + names = ", ".join( + f"{capability.event_type}[{','.join(sorted(capability.payload_fields))}]" + for capability in missing + ) + raise ValueError( + f"Input source does not satisfy provider raw user input schema: {names}." + ) + + +def _missing_capabilities( + *, + required: Sequence[UserInputCapability], + provided: UserInputSchema, +) -> tuple[UserInputCapability, ...]: + return tuple( + capability for capability in required if not provided.supports(capability) + ) + + +def _provider_capabilities(provider: ModelInputProvider) -> ProviderCapabilities: + capabilities = getattr(provider, "capabilities", None) + if not isinstance(capabilities, ProviderCapabilities): + raise TypeError( + "ModelInputProvider.capabilities must be a ProviderCapabilities " + f"instance, got {type(capabilities).__name__}." + ) + return capabilities + + +def _run_mode_capabilities(run_mode: RunMode) -> RunModeCapabilities: + capabilities = getattr(run_mode, "capabilities", None) + if not isinstance(capabilities, RunModeCapabilities): + raise TypeError( + "RunMode.capabilities must be a RunModeCapabilities instance, " + f"got {type(capabilities).__name__}." + ) + return capabilities + + +def _input_source_user_input_schema(input_source: object) -> UserInputSchema: + schema = getattr(input_source, "user_input_schema", None) + if not isinstance(schema, UserInputSchema): + raise TypeError( + "InputSource.user_input_schema must be a UserInputSchema instance, " + f"got {type(schema).__name__}." + ) + return schema + + +__all__ = [ + "ResolvedRunCapabilities", + "resolve_run_capabilities", + "validate_resolved_run", +] diff --git a/flashdreams/flashdreams/runtime/metrics.py b/flashdreams/flashdreams/runtime/metrics.py index 4286204f6..8a8fd6b1b 100644 --- a/flashdreams/flashdreams/runtime/metrics.py +++ b/flashdreams/flashdreams/runtime/metrics.py @@ -6,7 +6,8 @@ from __future__ import annotations import math -from collections.abc import Mapping +from collections import Counter, defaultdict +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from typing import Any, Protocol, runtime_checkable @@ -45,9 +46,31 @@ def __post_init__(self) -> None: object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) +@dataclass(frozen=True, kw_only=True, slots=True) +class MetricsSnapshot: + """Closed session or run metrics summary.""" + + counters: Mapping[str, int | float] = field(default_factory=dict) + timings: Mapping[str, Sequence[float]] = field(default_factory=dict) + session_statuses: Sequence[str] = () + errors: Sequence[str] = () + + def __post_init__(self) -> None: + object.__setattr__(self, "counters", freeze_mapping(self.counters)) + object.__setattr__( + self, + "timings", + freeze_mapping( + {key: tuple(values) for key, values in self.timings.items()} + ), + ) + object.__setattr__(self, "session_statuses", tuple(self.session_statuses)) + object.__setattr__(self, "errors", tuple(self.errors)) + + @runtime_checkable class MetricsRecorder(Protocol): - """Collector for runtime metrics.""" + """Collector for runtime, session, and run metrics.""" def record(self, sample: RuntimeMetricSample) -> None: """Record one metric sample.""" @@ -64,7 +87,53 @@ def record_timing( """Record one timing sample in seconds.""" ... - def close(self) -> None: + def record_step( + self, + *, + request: object, + user_window: object, + inference_input: object, + result: object, + decision: object, + ) -> None: + """Record one successful model step.""" + ... + + def record_control( + self, + *, + request: object, + user_window: object, + control: object, + ) -> None: + """Record one provider-authored control decision.""" + ... + + def record_error(self, exc: Exception, action: object) -> None: + """Record a driver-observed operational error.""" + ... + + def record_catch_up(self, decision: object) -> None: + """Record a realtime catch-up decision.""" + ... + + def record_cleanup_error(self, exc: Exception) -> None: + """Record a cleanup failure without interrupting teardown.""" + ... + + def record_orphaned_cleanup(self, exc: Exception) -> None: + """Record cleanup that timed out and is still queued.""" + ... + + def record_session(self, result: object) -> None: + """Record one closed session result.""" + ... + + def record_session_error(self, exc: Exception) -> None: + """Record diagnostic session assembly failure detail.""" + ... + + def close(self) -> MetricsSnapshot: """Finalize metric collection.""" ... @@ -74,6 +143,14 @@ class InMemoryMetricsRecorder: """Simple metrics recorder useful for tests, smoke runs, and adapters.""" samples: list[RuntimeMetricSample] = field(default_factory=list) + step_count: int = 0 + control_count: int = 0 + catch_up_count: int = 0 + errors: list[str] = field(default_factory=list) + cleanup_errors: list[str] = field(default_factory=list) + orphaned_cleanup_errors: list[str] = field(default_factory=list) + session_errors: list[str] = field(default_factory=list) + sessions: list[object] = field(default_factory=list) closed: bool = False def record(self, sample: RuntimeMetricSample) -> None: @@ -100,8 +177,96 @@ def record_timing( ) ) - def close(self) -> None: + def record_step( + self, + *, + request: object, + user_window: object, + inference_input: object, + result: object, + decision: object, + ) -> None: + del request, user_window, inference_input, result, decision + if not self.closed: + self.step_count += 1 + + def record_control( + self, + *, + request: object, + user_window: object, + control: object, + ) -> None: + del request, user_window, control + if not self.closed: + self.control_count += 1 + + def record_error(self, exc: Exception, action: object) -> None: + del action + if not self.closed: + self.errors.append(str(exc)) + + def record_catch_up(self, decision: object) -> None: + del decision + if not self.closed: + self.catch_up_count += 1 + + def record_cleanup_error(self, exc: Exception) -> None: + if not self.closed: + self.cleanup_errors.append(str(exc)) + + def record_orphaned_cleanup(self, exc: Exception) -> None: + if not self.closed: + self.orphaned_cleanup_errors.append(str(exc)) + + def record_session(self, result: object) -> None: + if not self.closed: + self.sessions.append(result) + + def record_session_error(self, exc: Exception) -> None: + if not self.closed: + self.session_errors.append(str(exc)) + + def close(self) -> MetricsSnapshot: self.closed = True + return self.snapshot() + + def snapshot(self) -> MetricsSnapshot: + timings: defaultdict[str, list[float]] = defaultdict(list) + for sample in self.samples: + if sample.category == "timing": + timings[sample.name].append(float(sample.value)) + session_statuses = tuple( + str(getattr(result, "status", "unknown")) for result in self.sessions + ) + session_status_counts = Counter(session_statuses) + return MetricsSnapshot( + counters={ + "samples": len(self.samples), + "steps": self.step_count, + "controls": self.control_count, + "catch_ups": self.catch_up_count, + "sessions": len(self.sessions), + "errors": len(self.errors), + "cleanup_errors": len(self.cleanup_errors), + "orphaned_cleanup_errors": len(self.orphaned_cleanup_errors), + "session_errors": len(self.session_errors), + **{ + f"sessions.{status}": count + for status, count in sorted(session_status_counts.items()) + }, + }, + timings=timings, + session_statuses=session_statuses, + errors=tuple( + ( + *self.errors, + *self.cleanup_errors, + *self.orphaned_cleanup_errors, + *self.session_errors, + ) + ), + ) class NullMetricsRecorder: @@ -120,5 +285,52 @@ def record_timing( ) -> None: del name, duration_s, step_index, metadata - def close(self) -> None: - return None + def record_step( + self, + *, + request: object, + user_window: object, + inference_input: object, + result: object, + decision: object, + ) -> None: + del request, user_window, inference_input, result, decision + + def record_control( + self, + *, + request: object, + user_window: object, + control: object, + ) -> None: + del request, user_window, control + + def record_error(self, exc: Exception, action: object) -> None: + del exc, action + + def record_catch_up(self, decision: object) -> None: + del decision + + def record_cleanup_error(self, exc: Exception) -> None: + del exc + + def record_orphaned_cleanup(self, exc: Exception) -> None: + del exc + + def record_session(self, result: object) -> None: + del result + + def record_session_error(self, exc: Exception) -> None: + del exc + + def close(self) -> MetricsSnapshot: + return MetricsSnapshot() + + +__all__ = [ + "InMemoryMetricsRecorder", + "MetricsRecorder", + "MetricsSnapshot", + "NullMetricsRecorder", + "RuntimeMetricSample", +] diff --git a/flashdreams/flashdreams/runtime/runner.py b/flashdreams/flashdreams/runtime/runner.py index c73a4f67c..f9a83411f 100644 --- a/flashdreams/flashdreams/runtime/runner.py +++ b/flashdreams/flashdreams/runtime/runner.py @@ -6,6 +6,8 @@ from __future__ import annotations import math +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any from flashdreams.runtime.canonical import InputCanonicalizer from flashdreams.runtime.config import InferenceConfig @@ -13,6 +15,7 @@ CanonicalInputs, CanonicalInputSchema, InferenceInput, + InferenceInputSchema, TimeWindow, UserInputs, UserInputSchema, @@ -29,7 +32,18 @@ ) from flashdreams.runtime.metrics import MetricsRecorder from flashdreams.runtime.output import OutputArtifact, OutputTarget -from flashdreams.runtime.types import StepResult +from flashdreams.runtime.types import ( + StepRequest, + StepRequirements, + StepResult, + step_requirements_from_request, +) + +if TYPE_CHECKING: + from flashdreams.runtime.demo.host import RuntimeHost + from flashdreams.runtime.demo.outputs import OutputDecision, SessionInfo + from flashdreams.runtime.demo.run_modes import RunResult + from flashdreams.runtime.demo.session_inputs import PreparedStep, UserInputWindow _DEFAULT_SESSION_HORIZON_S = 3600.0 @@ -46,73 +60,567 @@ def run_inference_session( output: OutputTarget, metrics: MetricsRecorder, ) -> tuple[OutputArtifact, ...]: - """Run one sequential inference session through the standard loop. + """Run one sequential inference session through the shared batch driver. - This v0 loop intentionally handles one adapter/runtime/session, one selected - input mapping, one replay/live input batch, one output target, and one - metrics recorder. It is synchronous and owns only orchestration. + The signature and failure semantics remain compatible with the original + replay runner while the implementation delegates the step loop to the shared + demo runtime pipeline. """ - runtime: InferenceRuntime | None = None - session: InferenceSession | None = None - output_opened = False - output_artifacts: tuple[OutputArtifact, ...] = () + return _run_inference_session_with_shared_batch( + adapter=adapter, + config=config, + mapping=mapping, + canonicalizer=canonicalizer, + source_schema=source_schema, + user_inputs=user_inputs, + initial_inputs=initial_inputs, + output=output, + metrics=metrics, + ) + + +def _run_inference_session_with_shared_batch( + *, + adapter: ModelAdapter, + config: InferenceConfig, + mapping: InputMapping, + canonicalizer: InputCanonicalizer, + source_schema: UserInputSchema, + user_inputs: UserInputs, + initial_inputs: InferenceInput, + output: OutputTarget, + metrics: MetricsRecorder, +) -> tuple[OutputArtifact, ...]: + lifecycle = _LegacyBatchLifecycle() + request_state = _LegacyStepRequestState() + runtime = _LegacyLazyRuntime( + adapter=adapter, + config=config, + lifecycle=lifecycle, + request_state=request_state, + ) + from flashdreams.runtime.demo.host import RuntimeHost + + host = RuntimeHost(runtime) + metrics_recorder = _LegacyRunnerMetricsRecorder(metrics) + output_sink = _LegacyOutputTargetSink( + output=output, + lifecycle=lifecycle, + host=host, + ) primary_error: BaseException | None = None try: - adapter.validate_config(config) - canonical_schema = canonicalizer.canonical_schema(source_schema) - _check_declared_mapping_compatibility( - mapping=mapping, - canonical_schema=canonical_schema, + _validate_legacy_runner_inputs( adapter=adapter, + config=config, + mapping=mapping, + canonicalizer=canonicalizer, + source_schema=source_schema, ) - mapping.validate( - canonical_schema=canonical_schema, + provider = _LegacyMappedModelInputProvider( + mapping=mapping, + canonicalizer=canonicalizer, + source_schema=source_schema, + user_inputs=user_inputs, + initial_inputs=initial_inputs, inference_input_schema=adapter.inference_input_schema, + request_state=request_state, ) - canonicalizer.reset() - mapped_initial_inputs = mapping.map_global_conditioning_inputs( - canonical_inputs=CanonicalInputs(), - inference_input=initial_inputs, + input_source = _LegacyBatchInputSource( + source_schema=source_schema, + user_inputs=user_inputs, + request_state=request_state, ) - runtime = adapter.create_runtime(config) - session = runtime.start_session(mapped_initial_inputs) - output.open() - output_opened = True - step_base_inputs = InferenceInput( - step=initial_inputs.step, - metadata=initial_inputs.metadata, + result = _run_shared_batch_session( + host=host, + provider=provider, + input_source=input_source, + output_sink=output_sink, + metrics=metrics_recorder, ) - - while (request := session.next_step_request()) is not None: - step_inputs = mapping.map_step_inputs( - canonical_inputs=canonicalizer.canonicalize( - user_inputs, - window=request.user_input_window - or _all_user_inputs_window(user_inputs), - source_schema=source_schema, - ), - inference_input=step_base_inputs, - request=request, - ) - result = session.step(step_inputs) - output.write(result) - _record_timing_metrics(metrics, result) + _raise_legacy_runner_error( + result=result, + output_sink=output_sink, + metrics=metrics_recorder, + ) + return tuple(result.artifacts) except BaseException as exc: primary_error = exc + if not metrics_recorder.closed: + _close_metrics_suppressing_secondary(metrics_recorder) raise finally: - cleanup_error, output_artifacts = _close_run_resources( - output=output if output_opened else None, - session=session, - runtime=runtime, + try: + host.close() + except BaseException: + if primary_error is None: + raise + + +def _validate_legacy_runner_inputs( + *, + adapter: ModelAdapter, + config: InferenceConfig, + mapping: InputMapping, + canonicalizer: InputCanonicalizer, + source_schema: UserInputSchema, +) -> CanonicalInputSchema: + adapter.validate_config(config) + canonical_schema = canonicalizer.canonical_schema(source_schema) + _check_declared_mapping_compatibility( + mapping=mapping, + canonical_schema=canonical_schema, + adapter=adapter, + ) + mapping.validate( + canonical_schema=canonical_schema, + inference_input_schema=adapter.inference_input_schema, + ) + return canonical_schema + + +def _run_shared_batch_session( + *, + host: RuntimeHost, + provider: "_LegacyMappedModelInputProvider", + input_source: "_LegacyBatchInputSource", + output_sink: "_LegacyOutputTargetSink", + metrics: "_LegacyRunnerMetricsRecorder", +) -> RunResult: + from flashdreams.runtime.demo.drivers import BatchSessionDriver + from flashdreams.runtime.demo.pipeline import StepPipeline + from flashdreams.runtime.demo.run_modes import SessionEdges + + return BatchSessionDriver().run_one_session( + host=host, + provider=provider, + session_edges=SessionEdges( + input_source=input_source, + output_sink=output_sink, + cleanup_tasks=set(), metrics=metrics, + ), + pipeline=StepPipeline(), + ) + + +class _LegacyLazyRuntime: + """Create the legacy runtime only after global inputs are mapped.""" + + def __init__( + self, + *, + adapter: ModelAdapter, + config: InferenceConfig, + lifecycle: "_LegacyBatchLifecycle", + request_state: "_LegacyStepRequestState", + ) -> None: + self._adapter = adapter + self._config = config + self._lifecycle = lifecycle + self._request_state = request_state + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + runtime = self._lifecycle.runtime + if runtime is None: + runtime = self._adapter.create_runtime(self._config) + self._lifecycle.set_runtime(runtime) + session = runtime.start_session(inputs) + self._lifecycle.set_session(session) + return _LegacySessionAdapter( + session=session, + request_state=self._request_state, + ) + + def close(self) -> None: + self._lifecycle.close_runtime_direct() + + +class _LegacyBatchLifecycle: + """Own legacy model resources whose close order is caller-visible.""" + + def __init__(self) -> None: + self.runtime: InferenceRuntime | None = None + self.session: InferenceSession | None = None + self._session_close_attempted = False + self._runtime_close_attempted = False + + def set_runtime(self, runtime: InferenceRuntime) -> None: + self.runtime = runtime + + def set_session(self, session: InferenceSession) -> None: + self.session = session + self._session_close_attempted = False + + def close_session_via_host(self, host: RuntimeHost) -> None: + host.call(self.close_session_direct) + + def close_runtime_via_host(self, host: RuntimeHost) -> None: + host.call(self.close_runtime_direct) + + def close_session_direct(self) -> None: + if self.session is None or self._session_close_attempted: + return + self._session_close_attempted = True + self.session.close() + + def close_runtime_direct(self) -> None: + if self.runtime is None or self._runtime_close_attempted: + return + self._runtime_close_attempted = True + self.runtime.close() + + +class _LegacySessionAdapter: + """Expose old sessions through the new StepRequirements boundary.""" + + def __init__( + self, + *, + session: InferenceSession, + request_state: "_LegacyStepRequestState", + ) -> None: + self._session = session + self._request_state = request_state + + def next_step_requirements(self) -> StepRequirements | None: + request = self._session.next_step_request() + if request is None: + self._request_state.clear() + return None + self._request_state.store(request) + return step_requirements_from_request( + request, + allow_user_input_window=True, + ) + + def next_step_request(self) -> StepRequest | None: + return self._session.next_step_request() + + def session_info(self) -> SessionInfo: + from flashdreams.runtime.demo.outputs import SessionInfo + + session_info = getattr(self._session, "session_info", None) + if not callable(session_info): + return SessionInfo() + value = session_info() + if not isinstance(value, SessionInfo): + raise TypeError( + "session.session_info() must return SessionInfo, " + f"got {type(value).__name__}." + ) + return value + + def step(self, inputs: InferenceInput) -> StepResult: + return self._session.step(inputs) + + def reset(self, inputs: InferenceInput | None = None) -> None: + self._session.reset(inputs) + + def close(self) -> None: + return None + + +class _LegacyStepRequestState: + """Share the current legacy request between the session, source, and provider.""" + + def __init__(self) -> None: + self._request: StepRequest | None = None + + def store(self, request: StepRequest) -> None: + self._request = request + + def require_for_window(self, step_index: int) -> StepRequest: + request = self._request + if request is None: + raise RuntimeError("Legacy input source has no active step request.") + if request.step_index != step_index: + raise RuntimeError( + "Legacy input source request mismatch: " + f"expected step {request.step_index}, got {step_index}." + ) + return request + + def consume_for_step(self, step_index: int) -> StepRequest: + request = self.require_for_window(step_index) + self._request = None + return request + + def clear(self) -> None: + self._request = None + + +class _LegacyBatchInputSource: + is_finite = True + is_deterministic = True + + def __init__( + self, + *, + source_schema: UserInputSchema, + user_inputs: UserInputs, + request_state: _LegacyStepRequestState, + ) -> None: + self.user_input_schema = source_schema + self._user_inputs = user_inputs + self._request_state = request_state + + def is_finished(self) -> bool: + return False + + def next_window(self, request: StepRequirements) -> UserInputWindow: + from flashdreams.runtime.demo.session_inputs import UserInputWindow + + legacy_request = self._request_state.require_for_window(request.step_index) + window = legacy_request.user_input_window or _all_user_inputs_window( + self._user_inputs + ) + return UserInputWindow( + start_s=window.start_s, + end_s=window.end_s, + inputs=self._user_inputs, + ) + + +class _LegacyMappedModelInputProvider: + def __init__( + self, + *, + mapping: InputMapping, + canonicalizer: InputCanonicalizer, + source_schema: UserInputSchema, + user_inputs: UserInputs, + initial_inputs: InferenceInput, + inference_input_schema: InferenceInputSchema, + request_state: _LegacyStepRequestState, + ) -> None: + from flashdreams.runtime.demo.session_inputs import ProviderCapabilities + + self.capabilities = ProviderCapabilities( + supports_recorded_input=True, + deterministic_given_inputs=True, + user_input_schema=source_schema, + inference_input_schema=inference_input_schema, + ) + self._mapping = mapping + self._canonicalizer = canonicalizer + self._source_schema = source_schema + self._user_inputs = user_inputs + self._initial_inputs = initial_inputs + self._request_state = request_state + self._step_base_inputs = InferenceInput( + step=initial_inputs.step, + metadata=initial_inputs.metadata, + ) + + def prepare_initial_input(self) -> InferenceInput: + self._canonicalizer.reset() + return self._mapping.map_global_conditioning_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=self._initial_inputs, ) - if cleanup_error is not None and primary_error is None: + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + from flashdreams.runtime.demo.session_inputs import PreparedStep + + legacy_request = self._request_state.consume_for_step(request.step_index) + canonical_inputs = self._canonicalizer.canonicalize( + self._user_inputs, + window=TimeWindow(start_s=user_window.start_s, end_s=user_window.end_s), + source_schema=self._source_schema, + ) + return PreparedStep( + inference_input=self._mapping.map_step_inputs( + canonical_inputs=canonical_inputs, + inference_input=self._step_base_inputs, + request=legacy_request, + ) + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self._canonicalizer.reset() + + def close(self) -> None: + return None + + +class _LegacyOutputTargetSink: + produces_artifacts = True + + def __init__( + self, + *, + output: OutputTarget, + lifecycle: _LegacyBatchLifecycle, + host: RuntimeHost, + ) -> None: + self._output = output + self._lifecycle = lifecycle + self._host = host + self._opened = False + self._closed = False + self._artifacts: tuple[OutputArtifact, ...] = () + self.cleanup_error: BaseException | None = None + + def open(self, session_info: SessionInfo) -> None: + del session_info + self._output.open() + self._opened = True + self._closed = False + self._artifacts = () + self.cleanup_error = None + + def begin_generation(self, generation: int) -> None: + del generation + + def write(self, result: StepResult) -> OutputDecision: + from flashdreams.runtime.demo.outputs import OutputDecision + + self._output.write(result) + return OutputDecision() + + def close(self) -> Sequence[OutputArtifact]: + if self._closed: + if self.cleanup_error is not None: + raise self.cleanup_error + return self._artifacts + + self._closed = True + cleanup_error: BaseException | None = None + artifacts: tuple[OutputArtifact, ...] = () + + def remember_error(exc: BaseException) -> None: + nonlocal cleanup_error + if cleanup_error is None: + cleanup_error = exc + + if self._opened: + try: + artifacts = tuple(self._output.close()) + except BaseException as exc: + remember_error(exc) + try: + self._lifecycle.close_session_via_host(self._host) + except BaseException as exc: + remember_error(exc) + try: + self._lifecycle.close_runtime_via_host(self._host) + except BaseException as exc: + remember_error(exc) + + self._artifacts = artifacts + self.cleanup_error = cleanup_error + if cleanup_error is not None: raise cleanup_error + return self._artifacts + + +class _LegacyRunnerMetricsRecorder: + def __init__(self, metrics: MetricsRecorder) -> None: + self._metrics = metrics + self.closed = False + self.close_error: BaseException | None = None - return output_artifacts + def record(self, sample: Any) -> None: + self._metrics.record(sample) + + def record_timing( + self, + name: str, + duration_s: float, + *, + step_index: int | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> None: + self._metrics.record_timing( + name, + duration_s, + step_index=step_index, + metadata=metadata, + ) + + def record_step( + self, + *, + request: object, + user_window: object, + inference_input: object, + result: object, + decision: object, + ) -> None: + del request, user_window, inference_input, decision + if isinstance(result, StepResult): + _record_timing_metrics(self._metrics, result) + + def record_control( + self, + *, + request: object, + user_window: object, + control: object, + ) -> None: + del request, user_window, control + + def record_error(self, exc: Exception, action: object) -> None: + del exc, action + + def record_catch_up(self, decision: object) -> None: + del decision + + def record_cleanup_error(self, exc: Exception) -> None: + del exc + + def record_orphaned_cleanup(self, exc: Exception) -> None: + del exc + + def record_session(self, result: object) -> None: + del result + + def record_session_error(self, exc: Exception) -> None: + del exc + + def close(self) -> Any: + if self.closed: + if self.close_error is not None: + raise self.close_error + return None + self.closed = True + try: + return self._metrics.close() + except BaseException as exc: + self.close_error = exc + raise + + +def _raise_legacy_runner_error( + *, + result: RunResult, + output_sink: _LegacyOutputTargetSink, + metrics: _LegacyRunnerMetricsRecorder, +) -> None: + if result.error is not None: + raise result.error + if output_sink.cleanup_error is not None: + raise output_sink.cleanup_error + if metrics.close_error is not None: + raise metrics.close_error + + +def _close_metrics_suppressing_secondary( + metrics: _LegacyRunnerMetricsRecorder, +) -> None: + try: + metrics.close() + except BaseException: + return def _check_declared_mapping_compatibility( @@ -158,45 +666,4 @@ def _record_timing_metrics( ) -def _close_run_resources( - *, - output: OutputTarget | None, - session: InferenceSession | None, - runtime: InferenceRuntime | None, - metrics: MetricsRecorder, -) -> tuple[BaseException | None, tuple[OutputArtifact, ...]]: - cleanup_error: BaseException | None = None - artifacts: tuple[OutputArtifact, ...] = () - - def remember_error(exc: BaseException) -> None: - nonlocal cleanup_error - if cleanup_error is None: - cleanup_error = exc - - if output is not None: - try: - artifacts = tuple(output.close()) - except BaseException as exc: - remember_error(exc) - - if session is not None: - try: - session.close() - except BaseException as exc: - remember_error(exc) - - if runtime is not None: - try: - runtime.close() - except BaseException as exc: - remember_error(exc) - - try: - metrics.close() - except BaseException as exc: - remember_error(exc) - - return cleanup_error, artifacts - - __all__ = ["run_inference_session"] diff --git a/flashdreams/flashdreams/runtime/types.py b/flashdreams/flashdreams/runtime/types.py index 4130a925f..49ee958a0 100644 --- a/flashdreams/flashdreams/runtime/types.py +++ b/flashdreams/flashdreams/runtime/types.py @@ -13,6 +13,65 @@ from flashdreams.runtime._utils import freeze_mapping from flashdreams.runtime.inputs import InferenceInputSchema, TimeWindow +_STEP_REQUIREMENTS_INPUT_COUNT_METADATA_KEY = "input_frame_count" +_STEP_REQUIREMENTS_STEADY_OUTPUT_COUNT_METADATA_KEY = "steady_output_frame_count" +_STEP_REQUIREMENTS_USER_INPUT_METADATA_KEYS = frozenset( + { + "input_window", + "user_input", + "user_input_window", + "user_inputs", + } +) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class StepRequirements: + """Model-authored per-step requirements consumed by shared demo drivers.""" + + __hash__ = None + + step_index: int + input_frame_count: int = 1 + steady_output_frame_count: int | None = None + inference_input_schema: InferenceInputSchema | None = None + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if isinstance(self.step_index, bool) or not isinstance(self.step_index, int): + raise TypeError("StepRequirements.step_index must be an integer.") + if self.step_index < 0: + raise ValueError("StepRequirements.step_index must be >= 0.") + if isinstance(self.input_frame_count, bool) or not isinstance( + self.input_frame_count, int + ): + raise TypeError("StepRequirements.input_frame_count must be an integer.") + if self.input_frame_count <= 0: + raise ValueError("StepRequirements.input_frame_count must be > 0.") + if self.steady_output_frame_count is not None: + if isinstance(self.steady_output_frame_count, bool) or not isinstance( + self.steady_output_frame_count, int + ): + raise TypeError( + "StepRequirements.steady_output_frame_count must be an integer." + ) + if self.steady_output_frame_count < 0: + raise ValueError( + "StepRequirements.steady_output_frame_count must be >= 0." + ) + user_input_keys = sorted( + key + for key in self.metadata + if key in _STEP_REQUIREMENTS_USER_INPUT_METADATA_KEYS + ) + if user_input_keys: + joined = ", ".join(user_input_keys) + raise ValueError( + "StepRequirements.metadata must not include driver-owned user " + f"input keys: {joined}." + ) + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + @dataclass(frozen=True, kw_only=True, slots=True) class StepRequest: @@ -36,4 +95,36 @@ def __post_init__(self) -> None: object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) -__all__ = ["StepRequest", "StepResult"] +def step_requirements_from_request( + request: StepRequest, + *, + allow_user_input_window: bool = False, +) -> StepRequirements: + """Adapt a legacy ``StepRequest`` that did not carry driver-owned inputs.""" + + if request.user_input_window is not None and not allow_user_input_window: + raise ValueError( + "StepRequest.user_input_window cannot be adapted to StepRequirements; " + "user input windows are driver-owned." + ) + metadata = dict(request.metadata) + input_frame_count = metadata.pop(_STEP_REQUIREMENTS_INPUT_COUNT_METADATA_KEY, 1) + steady_output_frame_count = metadata.pop( + _STEP_REQUIREMENTS_STEADY_OUTPUT_COUNT_METADATA_KEY, + None, + ) + return StepRequirements( + step_index=request.step_index, + input_frame_count=input_frame_count, + steady_output_frame_count=steady_output_frame_count, + inference_input_schema=request.inference_input_schema, + metadata=metadata, + ) + + +__all__ = [ + "StepRequest", + "StepRequirements", + "StepResult", + "step_requirements_from_request", +] diff --git a/flashdreams/flashdreams/runtime/worker.py b/flashdreams/flashdreams/runtime/worker.py index af4576c09..e523b6e6a 100644 --- a/flashdreams/flashdreams/runtime/worker.py +++ b/flashdreams/flashdreams/runtime/worker.py @@ -6,15 +6,17 @@ from __future__ import annotations import asyncio +import threading from concurrent.futures import ThreadPoolExecutor from typing import Any, Callable, TypeVar, cast import torch _T = TypeVar("_T") +_EXECUTOR_FUTURE_POLL_INTERVAL_S = 0.01 -class ThreadAffineRuntimeWorker: +class ModelExecutionWorker: """Run ordered runtime lifecycle calls on one owned OS thread. CUDA graphs, Triton launchers, and some backend contexts are thread-local. @@ -38,14 +40,23 @@ def __init__( thread_name_prefix=thread_name, initializer=self._initialize_thread, ) + self._state_lock = threading.Lock() self._accepting = True self._closed = False - self._close_lock = asyncio.Lock() + self._thread_id: int | None = None @property def closed(self) -> bool: return self._closed + @property + def worker_thread_id(self) -> int | None: + return self._thread_id + + @property + def is_worker_thread(self) -> bool: + return self._thread_id == threading.get_ident() + async def call( self, func: Callable[..., _T], @@ -54,11 +65,11 @@ async def call( **kwargs: Any, ) -> _T: """Run one callable after all previously submitted worker calls.""" - if not self._accepting: - raise RuntimeError("runtime worker is closed") + self._require_not_worker_thread() + self._require_accepting() future = self._submit(func, args, kwargs) try: - return await asyncio.shield(future) + return await _await_executor_future(future) except asyncio.CancelledError: future.add_done_callback(_consume_exception) raise @@ -71,21 +82,32 @@ def call_blocking( **kwargs: Any, ) -> _T: """Run one callable from synchronous code on the owned worker thread.""" - if not self._accepting: - raise RuntimeError("runtime worker is closed") + self._require_not_worker_thread() + self._require_accepting() future = self._executor.submit(_invoke, func, args, kwargs) return cast(_T, future.result()) 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 + self._require_not_worker_thread() + if not self._begin_close(): + return + try: + barrier = asyncio.wrap_future(self._executor.submit(_noop)) + await _await_executor_future(barrier) + finally: + self._finish_close() + + def close_blocking(self) -> None: + """Synchronous close for non-async setup and teardown paths.""" + self._require_not_worker_thread() + if not self._begin_close(): + return + try: + barrier = self._executor.submit(_noop) + barrier.result() + finally: + self._finish_close() def _submit( self, @@ -97,9 +119,36 @@ def _submit( return loop.run_in_executor(self._executor, _invoke, func, args, kwargs) def _initialize_thread(self) -> None: + self._thread_id = threading.get_ident() if self._device is not None and self._device.type == "cuda": torch.cuda.set_device(self._device) + def _require_accepting(self) -> None: + if not self._accepting: + raise RuntimeError("runtime worker is closed") + + def _require_not_worker_thread(self) -> None: + if self.is_worker_thread: + raise RuntimeError( + "Cannot dispatch to the model execution worker from its own " + "thread; call the function directly." + ) + + def _begin_close(self) -> bool: + with self._state_lock: + if self._closed: + return False + self._accepting = False + return True + + def _finish_close(self) -> None: + self._executor.shutdown(wait=True, cancel_futures=False) + with self._state_lock: + self._closed = True + + +ThreadAffineRuntimeWorker = ModelExecutionWorker + def _invoke( func: Callable[..., _T], @@ -113,9 +162,19 @@ def _noop() -> None: return +async def _await_executor_future(future: asyncio.Future[_T]) -> _T: + """Await an executor future without relying on a single cross-thread wakeup.""" + while not future.done(): + await asyncio.wait( + {future}, + timeout=_EXECUTOR_FUTURE_POLL_INTERVAL_S, + ) + return future.result() + + def _consume_exception(future: asyncio.Future[Any]) -> None: if not future.cancelled(): future.exception() -__all__ = ["ThreadAffineRuntimeWorker"] +__all__ = ["ModelExecutionWorker", "ThreadAffineRuntimeWorker"] diff --git a/flashdreams/flashdreams/serving/realtime/timing.py b/flashdreams/flashdreams/serving/realtime/timing.py index 9891cfcdc..30147f97d 100644 --- a/flashdreams/flashdreams/serving/realtime/timing.py +++ b/flashdreams/flashdreams/serving/realtime/timing.py @@ -11,6 +11,8 @@ from threading import Lock from typing import Protocol +from flashdreams.runtime.metrics import MetricsRecorder + TraceComponentValue = str | int | float | bool | None @@ -402,6 +404,36 @@ def summarize_chunk_history(chunks: Iterable[ChunkTimes]) -> RecentTimingSummary ) +def record_chunk_timing_metrics( + metrics: MetricsRecorder, + chunk: ChunkTimes, +) -> None: + """Record available chunk timing durations to a session metrics recorder.""" + + _record_stage_timing_metrics( + metrics, + prefix="realtime.chunk", + durations_ms=chunk.stage_durations_ms(), + step_index=chunk.chunk_index, + ) + + +def record_video_model_timing_metrics( + metrics: MetricsRecorder, + timings: VideoModelTimings, + *, + chunk_index: int | None = None, +) -> None: + """Record backend-visible video model stage durations to session metrics.""" + + _record_stage_timing_metrics( + metrics, + prefix="realtime.model", + durations_ms=timings.stage_durations_ms(), + step_index=chunk_index, + ) + + class RollingChunkTimingSummary: def __init__(self, capacity: int) -> None: self._chunks: deque[ChunkTimes] = deque(maxlen=capacity) @@ -518,6 +550,24 @@ def _add_optional_trace_range( ) +def _record_stage_timing_metrics( + metrics: MetricsRecorder, + *, + prefix: str, + durations_ms: Mapping[str, float], + step_index: int | None, +) -> None: + for stage_name, duration_ms in durations_ms.items(): + try: + metrics.record_timing( + f"{prefix}.{stage_name}", + float(duration_ms) / 1000.0, + step_index=step_index, + ) + except Exception: + return + + def _summarize_values(values: list[float]) -> StageDurationSummary: ordered = sorted(values) count = len(ordered) diff --git a/flashdreams/flashdreams/serving/webrtc/encoders.py b/flashdreams/flashdreams/serving/webrtc/encoders.py index 5eea69991..b07d9324c 100644 --- a/flashdreams/flashdreams/serving/webrtc/encoders.py +++ b/flashdreams/flashdreams/serving/webrtc/encoders.py @@ -20,7 +20,7 @@ import importlib.util from dataclasses import dataclass -from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, Literal, Protocol, cast, runtime_checkable import torch from aiortc import MediaStreamTrack @@ -69,6 +69,20 @@ class VideoEncoder(Protocol): def create_track(self, *, maxsize: int) -> BufferedVideoTrack | NVENCVideoTrack: ... + def prepare_chunk_payload( + self, + result: StepResult, + track: MediaStreamTrack, + ) -> object: ... + + async def deliver_prepared_chunk( + self, + payload: object, + track: MediaStreamTrack, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: ... + async def deliver_chunk( self, result: StepResult, @@ -110,10 +124,24 @@ def create_track(self, *, maxsize: int) -> BufferedVideoTrack: return BufferedVideoTrack(fps=self.fps, maxsize=maxsize) - async def deliver_chunk( + def prepare_chunk_payload( self, result: StepResult, track: MediaStreamTrack, + ) -> tuple[object, ...]: + from flashdreams.serving.webrtc.media import BufferedVideoTrack + + if not isinstance(track, BufferedVideoTrack): + raise TypeError( + "DefaultRTCEncoder requires a BufferedVideoTrack; got " + f"{type(track).__name__}. Create it via encoder.create_track()." + ) + return track.prepare_result_frames(result) + + async def deliver_prepared_chunk( + self, + payload: object, + track: MediaStreamTrack, *, force_keyframe: bool = False, ) -> ChunkDeliveryResult: @@ -128,7 +156,9 @@ async def deliver_chunk( "DefaultRTCEncoder requires a BufferedVideoTrack; got " f"{type(track).__name__}. Create it via encoder.create_track()." ) - enqueued = await track.enqueue_result(result) + if not isinstance(payload, tuple): + raise TypeError("DefaultRTCEncoder payload must be a tuple of RGB frames.") + enqueued = await track.enqueue_frames(cast(Any, payload)) return ChunkDeliveryResult( backend=self.backend, num_frames=enqueued, @@ -136,6 +166,19 @@ async def deliver_chunk( encode_ms=0.0, ) + async def deliver_chunk( + self, + result: StepResult, + track: MediaStreamTrack, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + return await self.deliver_prepared_chunk( + self.prepare_chunk_payload(result, track), + track, + force_keyframe=force_keyframe, + ) + def close(self) -> None: return diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index a7c6c87a4..4291cd637 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -10,10 +10,10 @@ import inspect import json from collections import deque -from collections.abc import Mapping +from collections.abc import Callable, Mapping from collections.abc import Set as AbstractSet from dataclasses import dataclass, field, replace -from typing import Any, Generic, TypeVar +from typing import Any, Generic, TypeVar, cast from aiortc import ( RTCConfiguration, @@ -23,12 +23,37 @@ ) from loguru import logger +from flashdreams.runtime.demo import ( + SPARSE_KEY_SEGMENTS_METADATA_KEY, + DemoSpec, + InMemorySessionMetricsRecorder, + ModelInputProvider, + PreparedScenario, + PreparedStep, + ProviderCapabilities, + ResamplerRealtimeClock, + RunContext, + RunResult, + RuntimeHost, + SessionEdges, + SessionInfo, + SingleSessionAdmissionPolicy, + StepPipeline, + UserInputWindow, + WebRTCErrorPolicy, + WebRTCOutputSpec, + run_demo_session_async, +) from flashdreams.runtime.inputs import ( + CanonicalInputSchema, InferenceInput, + InferenceInputSchema, TimeWindow, UserInputEvent, UserInputs, + UserInputSchema, ) +from flashdreams.runtime.mapping import InputMapping from flashdreams.runtime.types import StepRequest, StepResult from flashdreams.serving.realtime.input import ( DEFAULT_SUPPORTED_KEYS, @@ -37,7 +62,9 @@ ) from flashdreams.serving.webrtc.encoders import ( DefaultRTCEncoder, + EncoderBackend, VideoEncoder, + select_encoder, ) from flashdreams.serving.webrtc.media import BufferedVideoTrack, NVENCVideoTrack from flashdreams.serving.webrtc.messages import ( @@ -55,6 +82,18 @@ WebRTCSessionRuntime, ) from flashdreams.serving.webrtc.server import SessionBusyError +from flashdreams.serving.webrtc.services import ( + WEBRTC_SKIPPED_INPUTS_METADATA_KEY, + WEBRTC_SKIPPED_WINDOW_METADATA_KEY, + WEBRTC_USER_INPUT_SCHEMA, + ThreadSafeWebRTCOutputBridge, + WebRTCActivationPolicy, + WebRTCChunkDelivery, + WebRTCInputSource, + WebRTCOutputSink, + WebRTCRunMode, + WebRTCTransportService, +) from flashdreams.serving.webrtc.warmup import ( run_loopback_warmup_session, wait_for_ice_gathering_complete, @@ -78,8 +117,12 @@ """Maximum unconsumed raw events kept for an ``InferenceSession`` step.""" _RELEASE_USER_EVENT_TYPES = frozenset({"key_up"}) _KEY_USER_EVENT_TYPES = frozenset({"key_down", "key_up"}) +_SESSION_INPUT_KEY = "webrtc_session_input" +_STEP_REQUEST_KEY = "webrtc_step_request" +_SEGMENTS_KEY = "webrtc_segments" +_FRAME_TIMES_KEY = "webrtc_frame_times" -_RuntimeT = TypeVar("_RuntimeT", bound=WebRTCSessionRuntime) +_RuntimeT = TypeVar("_RuntimeT") _RuntimeConfigT = TypeVar("_RuntimeConfigT", bound=WebRTCRuntimeConfig) @@ -140,6 +183,421 @@ def _stat_int(stats: Mapping[str, float | int], name: str) -> int: return int(round(_stat_float(stats, name))) +def _runtime_drives_inference_session(runtime: Any) -> bool: + return callable(getattr(runtime, "start_inference_session", None)) + + +def _run_on_event_loop(loop: asyncio.AbstractEventLoop, awaitable: Any) -> Any: + """Run one legacy async WebRTC runtime call from a RuntimeHost worker.""" + return asyncio.run_coroutine_threadsafe(awaitable, loop).result() + + +def _step_request_from_requirements( + request: Any, + *, + window: TimeWindow, +) -> StepRequest: + metadata = dict(getattr(request, "metadata", {})) + metadata["input_frame_count"] = request.input_frame_count + steady_output_frame_count = getattr(request, "steady_output_frame_count", None) + if steady_output_frame_count is not None: + metadata["steady_output_frame_count"] = steady_output_frame_count + return StepRequest( + step_index=request.step_index, + inference_input_schema=getattr(request, "inference_input_schema", None), + user_input_window=window, + metadata=metadata, + ) + + +def _encoder_backend_from_config(value: object) -> EncoderBackend: + backend = str(value) + if backend not in {"auto", "default", "nvenc"}: + raise ValueError( + f"encoder_backend must be 'auto', 'default', or 'nvenc', got {backend!r}." + ) + return cast(EncoderBackend, backend) + + +def _gpu_id_from_device_spec(device_spec: str) -> int: + if not device_spec.startswith("cuda"): + return 0 + _prefix, separator, index = device_spec.partition(":") + if not separator or not index: + return 0 + try: + return int(index) + except ValueError: + return 0 + + +class _LegacyWebRTCRuntimeAdapter: + """Shared compatibility adapter from old async WebRTC runtimes to RuntimeHost.""" + + def __init__(self, *, runtime: Any, loop: asyncio.AbstractEventLoop) -> None: + self._runtime = runtime + self._loop = loop + + def reset_for_new_session(self, session_input: Any = None) -> None: + _run_on_event_loop( + self._loop, + self._runtime.reset_for_new_session(session_input=session_input), + ) + + def start_session(self, inputs: InferenceInput) -> "_LegacyWebRTCSessionAdapter": + del inputs + inference_session = None + if _runtime_drives_inference_session(self._runtime): + inference_session = _run_on_event_loop( + self._loop, + self._runtime.start_inference_session(), + ) + return _LegacyWebRTCSessionAdapter( + runtime=self._runtime, + inference_session=inference_session, + loop=self._loop, + ) + + def close(self) -> None: + # The underlying async runtime is owned by BaseWebRTCSessionManager and + # closed from shutdown(); RuntimeHost only owns this adapter's worker. + return + + +class _LegacyWebRTCSessionAdapter: + """RuntimeHost-facing session view over a legacy WebRTC runtime/session.""" + + def __init__( + self, + *, + runtime: Any, + inference_session: Any | None, + loop: asyncio.AbstractEventLoop, + ) -> None: + self._runtime = runtime + self._inference_session = inference_session + self._loop = loop + + def session_info(self) -> SessionInfo: + steady_frames: int | None = None + try: + steady_frames = int(self._runtime.peek_steady_output_num_frames()) + except Exception: + steady_frames = None + return SessionInfo(steady_output_frame_count=steady_frames) + + def next_step_request(self) -> StepRequest | None: + if self._inference_session is not None: + return self._inference_session.next_step_request() + return self._runtime.next_step_request() + + def step(self, inputs: InferenceInput) -> StepResult: + if self._inference_session is not None: + result = self._inference_session.step(inputs) + else: + result = _run_on_event_loop( + self._loop, + self._runtime.step( + request=inputs.step[_STEP_REQUEST_KEY], + segments=list(inputs.step[_SEGMENTS_KEY]), + frame_times=list(inputs.step[_FRAME_TIMES_KEY]), + ), + ) + request = inputs.step[_STEP_REQUEST_KEY] + if result.step_index != request.step_index: + raise RuntimeError( + "Runtime result step does not match its request: " + f"requested {request.step_index}, got {result.step_index}." + ) + if not isinstance(result, StepResult): + raise TypeError( + "WebRTC session steps must produce StepResult, got " + f"{type(result).__name__}." + ) + return result + + def reset(self, inputs: InferenceInput | None = None) -> None: + session_input = None + if inputs is not None: + session_input = inputs.global_conditioning.get(_SESSION_INPUT_KEY) + _run_on_event_loop( + self._loop, + self._runtime.reset_for_new_session(session_input=session_input), + ) + if self._inference_session is not None: + self._inference_session = _run_on_event_loop( + self._loop, + self._runtime.start_inference_session(), + ) + + def close(self) -> None: + close = getattr(self._inference_session, "close", None) + if callable(close): + close() + + +class _LegacyWebRTCModelInputProvider: + """Shared provider used until model-specific WebRTC providers land.""" + + def __init__(self, *, runtime: Any, session_input: Any = None) -> None: + self._runtime = runtime + self._session_input = session_input + self._uses_inference_session = _runtime_drives_inference_session(runtime) + self._session_input_state_advanced = False + self.capabilities = ProviderCapabilities( + supports_realtime_clock=True, + supports_reset=True, + deterministic_given_inputs=False, + user_input_schema=self._user_input_schema(), + ) + + def prepare_initial_input(self) -> InferenceInput: + if self._session_input is None: + return InferenceInput() + return InferenceInput( + global_conditioning={_SESSION_INPUT_KEY: self._session_input} + ) + + def prepare_step( + self, + *, + request: Any, + user_window: UserInputWindow, + ) -> PreparedStep: + if self._uses_inference_session: + return PreparedStep( + inference_input=self._prepare_inference_session_step( + request=request, + user_window=user_window, + ) + ) + return PreparedStep( + inference_input=self._prepare_segment_step( + request=request, + user_window=user_window, + ) + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self._session_input_state_advanced = False + + def close(self) -> None: + return + + def _user_input_schema(self) -> UserInputSchema: + schema = getattr(self._runtime, "input_source_schema", None) + if isinstance(schema, UserInputSchema): + return schema + return WEBRTC_USER_INPUT_SCHEMA + + def _prepare_inference_session_step( + self, + *, + request: Any, + user_window: UserInputWindow, + ) -> InferenceInput: + self._advance_skipped_input_state(user_window) + window_start = user_window.start_s + if request.step_index == 0 and not self._session_input_state_advanced: + window_start = 0.0 + window = TimeWindow(start_s=window_start, end_s=user_window.end_s) + canonical_inputs = self._runtime.input_canonicalizer.canonicalize( + user_window.inputs, + window=window, + source_schema=self._runtime.input_source_schema, + ) + mapping = self._runtime.input_mapping + inference_input = InferenceInput( + metadata={ + **dict(user_window.metadata), + "frame_times": tuple(user_window.frame_times), + "window_start_s": window.start_s, + "window_end_s": window.end_s, + } + ) + return mapping.map_step_inputs( + canonical_inputs=canonical_inputs, + inference_input=inference_input, + request=_step_request_from_requirements(request, window=window), + ) + + def _advance_skipped_input_state(self, user_window: UserInputWindow) -> None: + skipped_inputs = user_window.metadata.get(WEBRTC_SKIPPED_INPUTS_METADATA_KEY) + skipped_window = user_window.metadata.get(WEBRTC_SKIPPED_WINDOW_METADATA_KEY) + if not isinstance(skipped_inputs, UserInputs): + return + if not isinstance(skipped_window, tuple) or len(skipped_window) != 2: + return + start_value, end_value = skipped_window + if not isinstance(start_value, int | float) or not isinstance( + end_value, + int | float, + ): + return + start_s = float(start_value) + end_s = float(end_value) + if end_s <= start_s: + return + self._runtime.input_canonicalizer.canonicalize( + skipped_inputs, + window=TimeWindow(start_s=start_s, end_s=end_s), + source_schema=self._runtime.input_source_schema, + ) + self._session_input_state_advanced = True + + @staticmethod + def _prepare_segment_step( + *, + request: Any, + user_window: UserInputWindow, + ) -> InferenceInput: + segments = user_window.metadata.get(SPARSE_KEY_SEGMENTS_METADATA_KEY) + if not isinstance(segments, tuple): + raise RuntimeError("WebRTC user window is missing resampled key segments.") + window = TimeWindow(start_s=user_window.start_s, end_s=user_window.end_s) + return InferenceInput( + step={ + _STEP_REQUEST_KEY: _step_request_from_requirements( + request, + window=window, + ), + _SEGMENTS_KEY: tuple(segments), + _FRAME_TIMES_KEY: tuple(user_window.frame_times), + } + ) + + +class _LegacyWebRTCDemoAdapter: + """Minimal adapter for the shared helper while WebRTC providers migrate.""" + + model_id: str + inference_input_schema = InferenceInputSchema() + canonical_input_schema = CanonicalInputSchema() + + def __init__( + self, + *, + runtime: Any, + identity: str, + session_input: Any = None, + ) -> None: + self._runtime = runtime + self.model_id = identity + self._session_input = session_input + + def supported_input_modes(self) -> tuple[str, ...]: + return ("webrtc",) + + def supported_output_modes(self) -> tuple[str, ...]: + return ("webrtc",) + + def default_input_mapping(self) -> InputMapping | None: + return None + + def validate_config(self, config: Any) -> None: + if config.model_id != self.model_id: + raise ValueError( + f"Expected WebRTC model_id={self.model_id!r}, got {config.model_id!r}." + ) + + def create_runtime(self, config: Any) -> Any: + self.validate_config(config) + return self._runtime + + def prepare_scenario(self, spec: Any) -> PreparedScenario: + del spec + return PreparedScenario(initial_inputs=self._initial_inputs()) + + def create_model_input_provider( + self, + spec: Any, + scenario: PreparedScenario, + ) -> _LegacyWebRTCModelInputProvider: + del spec, scenario + return _LegacyWebRTCModelInputProvider( + runtime=self._runtime, + session_input=self._session_input, + ) + + def _initial_inputs(self) -> InferenceInput: + if self._session_input is None: + return InferenceInput() + return InferenceInput( + global_conditioning={_SESSION_INPUT_KEY: self._session_input} + ) + + +class _ManagedWebRTCSessionEdgeFactory: + """Build shared realtime edges for one negotiated peer connection.""" + + def __init__( + self, + *, + manager: "BaseWebRTCSessionManager[Any, Any]", + managed_session: "ManagedWebRTCSession", + loop: asyncio.AbstractEventLoop, + ) -> None: + self._manager = manager + self._managed_session = managed_session + self._loop = loop + + def create_session_edges( + self, + *, + context: RunContext, + spec: Any, + scenario: PreparedScenario, + provider: ModelInputProvider, + adapter: Any, + ) -> SessionEdges: + del spec, scenario, provider, adapter + input_source = self._managed_session.input_source + transport = self._managed_session.transport + if input_source is None or transport is None: + raise RuntimeError("Managed WebRTC session is missing shared edges.") + bridge = ThreadSafeWebRTCOutputBridge( + loop=self._loop, + video_encoder=self._managed_session.video_encoder, + video_track=self._managed_session.video_track, + on_chunk_delivery=self._on_chunk_delivery, + on_error=self._on_delivery_error, + ) + return SessionEdges( + input_source=input_source, + output_sink=WebRTCOutputSink(bridge=bridge), + cleanup_tasks=context.cleanup_tasks, + metrics=InMemorySessionMetricsRecorder(), + error_policy=WebRTCErrorPolicy(), + transport=transport, + clock=ResamplerRealtimeClock( + resampler=self._managed_session.resampler, + now_fn=self._loop.time, + sleep_fn=asyncio.sleep, + ), + activation=WebRTCActivationPolicy( + input_source=input_source, + transport=transport, + ), + ) + + def _on_chunk_delivery(self, chunk: WebRTCChunkDelivery) -> None: + self._manager._handle_shared_chunk_delivery( + managed_session=self._managed_session, + chunk=chunk, + ) + + def _on_delivery_error(self, exc: BaseException) -> None: + self._manager._handle_shared_delivery_error( + managed_session=self._managed_session, + exc=exc, + ) + if self._manager.fatal_generation_errors: + self._loop.call_soon_threadsafe( + lambda: asyncio.create_task(self._manager.close_active_session()) + ) + + @dataclass(slots=True) class ManagedWebRTCSession: """Per-session state for the single active WebRTC peer connection.""" @@ -152,6 +610,9 @@ class ManagedWebRTCSession: control_channel: Any | None = None generation_task: asyncio.Task[Any] | None = None first_action_received: asyncio.Event = field(default_factory=asyncio.Event) + input_source: WebRTCInputSource | None = None + transport: WebRTCTransportService | None = None + reservation: Any | None = None pending_action_arrivals: deque[float] = field(default_factory=deque) inference_session: Any | None = None """Active ``InferenceSession``; ``None`` means call ``runtime.generate_chunk``.""" @@ -189,6 +650,11 @@ async def close(self) -> None: self.generation_task.cancel() with contextlib.suppress(asyncio.CancelledError): await self.generation_task + if self.generation_task is None or self.generation_task.done(): + reservation = self.reservation + self.reservation = None + if reservation is not None: + reservation.release() self.generation_task = None await self.video_track.close() @@ -212,6 +678,11 @@ def __init__( supported_control_keys: AbstractSet[str] | None = None, fatal_generation_errors: bool = False, client_liveness_timeout_s: float = DEFAULT_CLIENT_LIVENESS_TIMEOUT_S, + shared_host: RuntimeHost | None = None, + shared_adapter: Any | None = None, + shared_spec: DemoSpec | None = None, + shared_scenario: PreparedScenario | None = None, + shared_pipeline_factory: Callable[[], StepPipeline] | None = None, ) -> None: if client_liveness_timeout_s <= 0: raise ValueError("client_liveness_timeout_s must be > 0") @@ -234,6 +705,15 @@ def __init__( self._preload_lock = asyncio.Lock() self._session_lock = asyncio.Lock() self._pending_session_input: Any = None + self._shared_runtime_adapter: _LegacyWebRTCRuntimeAdapter | None = None + self._shared_host: RuntimeHost | None = shared_host + self._owns_shared_host = shared_host is not None + self._shared_context: RunContext | None = None + self._shared_adapter = shared_adapter + self._shared_spec = shared_spec + self._shared_scenario = shared_scenario + self._shared_pipeline_factory = shared_pipeline_factory + self._shared_video_encoder: VideoEncoder | None = None @property def pending_session_input(self) -> Any: @@ -296,8 +776,11 @@ def _positive_float_runtime_value(value: Any, *, label: str) -> float: return parsed def _runtime_input_fps(self, runtime: Any) -> float: + peek_input_fps = getattr(runtime, "peek_input_fps", None) + if not callable(peek_input_fps): + return float(self.fps) return self._positive_float_runtime_value( - runtime.peek_input_fps(), + peek_input_fps(), label="peek_input_fps", ) @@ -315,9 +798,22 @@ def _runtime_next_step_request(self, runtime: Any) -> tuple[StepRequest, int]: return request, input_num_frames def _runtime_steady_output_num_frames(self, runtime: Any) -> int: + peek_output_frames = getattr(runtime, "peek_steady_output_num_frames", None) + if callable(peek_output_frames): + return self._positive_int_runtime_value( + peek_output_frames(), + label="peek_steady_output_num_frames", + ) + pipeline = getattr(runtime, "pipeline", None) + get_num_frames = getattr(pipeline, "get_num_frames", None) + if callable(get_num_frames): + return self._positive_int_runtime_value( + get_num_frames(1), + label="pipeline.get_num_frames(1)", + ) return self._positive_int_runtime_value( - runtime.peek_steady_output_num_frames(), - label="peek_steady_output_num_frames", + 1, + label="fallback steady output frame count", ) def _resolve_video_encoder(self) -> VideoEncoder: @@ -330,10 +826,60 @@ def _resolve_video_encoder(self) -> VideoEncoder: transparently get the software path without having to opt in. """ encoder = getattr(self._runtime, "video_encoder", None) + if encoder is None: + encoder = self._shared_video_encoder if encoder is None: encoder = DefaultRTCEncoder(fps=self.fps) return encoder + def _shared_run_context(self, loop: asyncio.AbstractEventLoop) -> RunContext: + if self._shared_context is not None: + return self._shared_context + host = self._shared_host + if host is None: + runtime_adapter = _LegacyWebRTCRuntimeAdapter( + runtime=self._runtime, + loop=loop, + ) + host = RuntimeHost(runtime_adapter) + self._shared_runtime_adapter = runtime_adapter + self._shared_host = host + self._shared_context = RunContext( + host=host, + run_metrics=InMemorySessionMetricsRecorder(), + admission=SingleSessionAdmissionPolicy( + health_check=lambda: host.is_healthy + ), + ) + return self._shared_context + + def _shared_demo_spec(self) -> DemoSpec: + return DemoSpec( + model_id=self.identity, + input_mode="webrtc", + output=WebRTCOutputSpec( + fps=self.fps, + video_width=self.runtime_config.video_width, + video_height=self.runtime_config.video_height, + warmup_chunks=self.runtime_config.warmup_chunks, + warmup_timeout_s=self.runtime_config.warmup_timeout_s, + client_liveness_timeout_s=self.client_liveness_timeout_s, + ), + ) + + async def _reset_runtime_for_session( + self, + *, + context: RunContext, + session_input: Any, + ) -> None: + reset = getattr(context.host.runtime, "reset_for_new_session", None) + if not callable(reset): + if self._shared_adapter is not None: + return + raise RuntimeError("WebRTC runtime adapter cannot reset sessions.") + await context.host.call_async(reset, session_input) + def _prefer_h264_video_codec(self, *, transceiver: Any) -> None: """Constrain the transceiver's codec preferences to H.264 variants. @@ -496,7 +1042,7 @@ async def _handle_event_message( @staticmethod def _drives_inference_session(runtime: Any) -> bool: """Return whether ``runtime`` should be driven through ``InferenceSession``.""" - return callable(getattr(runtime, "start_inference_session", None)) + return _runtime_drives_inference_session(runtime) def _record_user_event( self, @@ -769,14 +1315,46 @@ def is_runtime_ready(self) -> bool: async def preload_runtime(self) -> None: async with self._preload_lock: if not self._runtime_ready: - await self._runtime.initialize() + initialize = getattr(self._runtime, "initialize", None) + if callable(initialize): + result = initialize() + if inspect.isawaitable(result): + await result + elif self._shared_host is not None: + await asyncio.to_thread(self._shared_host.preload) self._runtime_ready = True + self._initialize_shared_video_encoder() if not self._warmup_complete: await self._run_loopback_warmup_session( num_chunks=self.runtime_config.warmup_chunks ) self._warmup_complete = True + def _initialize_shared_video_encoder(self) -> None: + if self._shared_video_encoder is not None: + return + if getattr(self._runtime, "video_encoder", None) is not None: + return + encoder_backend = getattr(self.runtime_config, "encoder_backend", None) + if encoder_backend is None: + return + backend = _encoder_backend_from_config(encoder_backend) + device_spec = str(getattr(self.runtime_config, "device", "")) + device_type = device_spec.split(":", maxsplit=1)[0] + if device_type != "cuda" and backend == "auto": + backend = "default" + if device_type != "cuda" and backend == "nvenc": + raise RuntimeError("encoder_backend='nvenc' requires a CUDA device.") + self._shared_video_encoder = select_encoder( + backend=backend, + width=self.runtime_config.video_width, + height=self.runtime_config.video_height, + fps=self.fps, + bitrate=int(getattr(self.runtime_config, "encoder_bitrate_bps", 6_000_000)), + gpu_id=_gpu_id_from_device_spec(device_spec), + gop=int(getattr(self.runtime_config, "encoder_gop", self.fps)), + ) + async def create_answer(self, *, offer_sdp: str, offer_type: str) -> dict[str, str]: if not self._runtime_ready or not self._warmup_complete: await self.preload_runtime() @@ -808,46 +1386,61 @@ async def _create_answer_with_runtime_ready_locked( if not self._runtime_ready: raise RuntimeError("Runtime is not initialized.") - 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 - # is throttled to the consumer's drain rate. AR step 0 emits fewer - # frames than steady state; sizing to it would force a per-chunk - # stall, so we size to the steady-state count. - num_frames = self._runtime_steady_output_num_frames(self._runtime) - video_encoder = self._resolve_video_encoder() - video_track = video_encoder.create_track(maxsize=num_frames) - # Use ``addTransceiver`` (not ``addTrack``) so we can constrain the - # SDP m-line's codec list via ``setCodecPreferences`` when the - # encoder emits pre-encoded H.264 packets. - video_transceiver = peer_connection.addTransceiver( - video_track, - direction="sendonly", - ) - if video_encoder.prefers_codec == "h264": - self._prefer_h264_video_codec(transceiver=video_transceiver) - # Start the resampler's virtual clock at 0; the real anchor is set - # in the ``on_datachannel`` handler so chunk 0's window starts when - # input can actually arrive. - resampler = self._make_resampler_at_fps( - start_v=0.0, - fps=self._runtime_input_fps(self._runtime), - ) loop = asyncio.get_running_loop() - managed_session = ManagedWebRTCSession( - runtime=self._runtime, - video_track=video_track, - video_encoder=video_encoder, - peer_connection=peer_connection, - resampler=resampler, - last_client_message_at=loop.time(), - ) - session_runtime: Any = self._runtime - if self._drives_inference_session(session_runtime): - managed_session.inference_session = ( - await session_runtime.start_inference_session() + context = self._shared_run_context(loop) + reservation = context.admission.try_reserve() + if reservation is None: + raise SessionBusyError(self.busy_message) + try: + await self._reset_runtime_for_session( + context=context, + session_input=session_input, + ) + except Exception: + reservation.release() + raise + + try: + peer_connection = RTCPeerConnection(rtc_configuration) + # Bounded queue sized to one *steady-state* chunk so the producer + # is throttled to the consumer's drain rate. AR step 0 emits fewer + # frames than steady state; sizing to it would force a per-chunk + # stall, so we size to the steady-state count. + num_frames = self._runtime_steady_output_num_frames(self._runtime) + video_encoder = self._resolve_video_encoder() + video_track = video_encoder.create_track(maxsize=num_frames) + # Use ``addTransceiver`` (not ``addTrack``) so we can constrain the + # SDP m-line's codec list via ``setCodecPreferences`` when the + # encoder emits pre-encoded H.264 packets. + video_transceiver = peer_connection.addTransceiver( + video_track, + direction="sendonly", + ) + if video_encoder.prefers_codec == "h264": + self._prefer_h264_video_codec(transceiver=video_transceiver) + # Start the resampler's virtual clock at 0; the real anchor is set + # in the ``on_datachannel`` handler so chunk 0's window starts when + # input can actually arrive. + resampler = self._make_resampler_at_fps( + start_v=0.0, + fps=self._runtime_input_fps(self._runtime), ) + input_source = WebRTCInputSource(resampler=resampler) + transport = WebRTCTransportService(loop=loop) + managed_session = ManagedWebRTCSession( + runtime=self._runtime, + video_track=video_track, + video_encoder=video_encoder, + peer_connection=peer_connection, + resampler=resampler, + input_source=input_source, + transport=transport, + reservation=reservation, + last_client_message_at=loop.time(), + ) + except Exception: + reservation.release() + raise self._active_session = managed_session if enable_liveness_watchdog: managed_session.liveness_task = asyncio.create_task( @@ -858,10 +1451,13 @@ async def _create_answer_with_runtime_ready_locked( def on_datachannel(channel: Any) -> None: managed_session.control_channel = channel # Re-anchor the resampler at channel open. The real - # virtual-clock anchor happens in ``_generation_worker`` once - # the first keyboard event arrives. + # virtual-clock anchor happens in ``WebRTCActivationPolicy`` once + # the first browser event activates the shared realtime driver. channel_open_v = asyncio.get_running_loop().time() - managed_session.resampler.reset(start_v=channel_open_v) + if managed_session.input_source is not None: + managed_session.input_source.reset(start_v=channel_open_v) + else: + managed_session.resampler.reset(start_v=channel_open_v) @channel.on("message") def on_message(message: Any) -> None: @@ -872,15 +1468,21 @@ def on_message(message: Any) -> None: ) ) - # Spawn the generation worker once the channel is wired up so + # Spawn the shared realtime session once the channel is wired up so # ``chunk_done`` notifications have a channel to land on. managed_session.generation_task = asyncio.create_task( - self._generation_worker(managed_session=managed_session) + self._run_realtime_driver_session( + managed_session=managed_session, + context=context, + session_input=session_input, + ) ) @channel.on("close") def on_close() -> None: logger.info("Control data channel closed; closing active session.") + if managed_session.transport is not None: + managed_session.transport.disconnect("data channel closed") asyncio.create_task(self.close_active_session()) @peer_connection.on("connectionstatechange") @@ -993,15 +1595,37 @@ async def _client_liveness_watchdog( async def shutdown(self) -> None: await self.close_active_session() - await self._runtime.close() + if self._shared_context is not None: + await self._shared_context.close_async() + if self._shared_host is not None: + await asyncio.to_thread(self._shared_host.close) + self._shared_context = None + self._shared_host = None + self._shared_runtime_adapter = None + if self._shared_video_encoder is not None: + self._shared_video_encoder.close() + self._shared_video_encoder = None + if not self._owns_shared_host: + close = getattr(self._runtime, "close", None) + if callable(close): + result = close() + if inspect.isawaitable(result): + await result self._runtime_ready = False self._warmup_complete = False def wait_for_termination(self) -> None: - self._runtime.wait_for_termination() + wait = getattr(self._runtime, "wait_for_termination", None) + if callable(wait): + wait() + return + if self._shared_host is not None: + self._shared_host.run_worker_loop() def send_exit_signal(self) -> None: - self._runtime.send_exit_signal() + send = getattr(self._runtime, "send_exit_signal", None) + if callable(send): + send() async def _handle_datachannel_message( self, @@ -1013,6 +1637,10 @@ async def _handle_datachannel_message( if channel is None or managed_session.closed: return managed_session.last_client_message_at = asyncio.get_running_loop().time() + if managed_session.transport is not None: + managed_session.transport.mark_client_message( + managed_session.last_client_message_at + ) if not isinstance(raw_message, str): self._send_json(channel, make_error_payload("Expected text payload.")) @@ -1029,6 +1657,12 @@ async def _handle_datachannel_message( channel, make_error_payload("Payload must be a JSON object.") ) return + if managed_session.input_source is not None: + await self._handle_shared_datachannel_payload( + managed_session=managed_session, + payload=payload, + ) + return message_type = str(payload.get("type", "")).strip().lower() if message_type == MESSAGE_TYPE_HEARTBEAT: return @@ -1106,6 +1740,202 @@ async def _handle_datachannel_message( # user actually interacts. Idempotent once already set. managed_session.first_action_received.set() + async def _handle_shared_datachannel_payload( + self, + *, + managed_session: ManagedWebRTCSession, + payload: dict[str, Any], + ) -> None: + channel = managed_session.control_channel + input_source = managed_session.input_source + if channel is None or input_source is None: + return + message_type = str(payload.get("type", "")).strip().lower() + if message_type == MESSAGE_TYPE_HEARTBEAT: + return + if message_type == MESSAGE_TYPE_DISCONNECT: + logger.info("Client requested disconnect; closing active session.") + if managed_session.transport is not None: + managed_session.transport.disconnect("client disconnected") + await self.close_active_session() + return + if message_type == MESSAGE_TYPE_EVENT: + handled = self._record_shared_event_payload( + managed_session=managed_session, + payload=payload, + ) + if handled: + managed_session.first_action_received.set() + return + result = input_source.handle_browser_payload( + payload, + timestamp_s=asyncio.get_running_loop().time(), + ) + if result.kind == "error": + self._send_json(channel, make_error_payload(result.error or "Bad input.")) + return + if result.activated: + managed_session.first_action_received.set() + + def _record_shared_event_payload( + self, + *, + managed_session: ManagedWebRTCSession, + payload: dict[str, Any], + ) -> bool: + channel = managed_session.control_channel + input_source = managed_session.input_source + if channel is None or input_source is None: + return False + event_id = str(payload.get("event_id", payload.get("id", ""))).strip() + state = str(payload.get("state", "trigger")).strip().lower() or "trigger" + clear_states = {"clear", "release", "off", "none"} + if not event_id and state not in clear_states: + self._send_json( + channel, + make_error_payload( + ( + "Event payload must include non-empty 'event_id' " + "unless state clears the active event." + ), + ), + ) + return False + clears = state in clear_states + try: + event_payload = self._validate_user_event_payload( + managed_session=managed_session, + event_type="text_event", + payload={ + "event_id": None if clears else event_id, + "state": state, + }, + ) + active_event_id = event_payload.get("event_id") + source_event_id = None if active_event_id is None else str(active_event_id) + input_source.record_user_event( + timestamp_s=asyncio.get_running_loop().time(), + event_type="text_event", + payload=event_payload, + source_event_id=source_event_id, + ) + except Exception as exc: + self._send_json(channel, make_error_payload(str(exc))) + return False + active_event_id = event_payload.get("event_id") + ack_event_id = None if active_event_id is None else str(active_event_id) + self._send_json( + channel, + make_event_ack_payload( + event_id=ack_event_id, + state=str(event_payload.get("state", state)), + result={"active_event_id": ack_event_id}, + ), + ) + return True + + async def _run_realtime_driver_session( + self, + *, + managed_session: ManagedWebRTCSession, + context: RunContext, + session_input: Any, + ) -> None: + adapter = self._shared_adapter + spec = self._shared_spec + scenario = self._shared_scenario + if adapter is None or spec is None: + adapter = _LegacyWebRTCDemoAdapter( + runtime=self._runtime, + identity=self.identity, + session_input=session_input, + ) + spec = self._shared_demo_spec() + if scenario is None: + scenario = adapter.prepare_scenario(spec) + run_mode = WebRTCRunMode( + edge_factory=_ManagedWebRTCSessionEdgeFactory( + manager=self, + managed_session=managed_session, + loop=asyncio.get_running_loop(), + ) + ) + try: + result = await run_demo_session_async( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=run_mode, + pipeline=( + self._shared_pipeline_factory() + if self._shared_pipeline_factory is not None + else StepPipeline() + ), + reservation=managed_session.reservation, + ) + if result.status == "completed": + logger.info("Shared WebRTC session completed.") + else: + logger.warning( + "Shared WebRTC session ended with status={} reason={}", + result.status, + result.reason, + ) + if result.status != "completed" and result.reason: + channel = managed_session.control_channel + if channel is not None: + self._send_json(channel, make_error_payload(result.reason)) + finally: + managed_session.reservation = None + if self._active_session is managed_session: + await self.close_active_session() + + def _handle_shared_chunk_delivery( + self, + *, + managed_session: ManagedWebRTCSession, + chunk: WebRTCChunkDelivery, + ) -> None: + channel = managed_session.control_channel + if channel is None or managed_session.closed: + return + delivery = chunk.delivery + enqueued_frames = int(getattr(delivery, "num_frames", chunk.frame_count)) + encode_ms = float(getattr(delivery, "encode_ms", 0.0)) + play_ms = chunk.frame_count * 1000.0 / managed_session.video_track.fps + queue_depth = managed_session.video_track.qsize() + self._send_json( + channel, + make_chunk_done_payload( + chunk_index=chunk.step_index, + num_frames=chunk.frame_count, + enqueued_frames=enqueued_frames, + fps=managed_session.video_track.fps, + width=self.runtime_config.video_width, + height=self.runtime_config.video_height, + model=self.identity, + gen_ms=_stat_ms(chunk.metrics, "model_step_s"), + enqueue_ms=encode_ms, + play_ms=play_ms, + queue_depth=queue_depth, + lag_ms=0.0, + control_latency_ms=None, + consumed_actions=0, + extra=chunk.metadata, + ), + ) + + def _handle_shared_delivery_error( + self, + *, + managed_session: ManagedWebRTCSession, + exc: BaseException, + ) -> None: + channel = managed_session.control_channel + if channel is not None: + self._send_json(channel, make_error_payload(str(exc))) + async def _generation_worker( self, *, managed_session: ManagedWebRTCSession ) -> None: diff --git a/flashdreams/flashdreams/serving/webrtc/media.py b/flashdreams/flashdreams/serving/webrtc/media.py index 470569132..25fa09437 100644 --- a/flashdreams/flashdreams/serving/webrtc/media.py +++ b/flashdreams/flashdreams/serving/webrtc/media.py @@ -81,16 +81,37 @@ def maxsize(self) -> int: def qsize(self) -> int: return self._frames.qsize() - async def enqueue_result(self, result: StepResult) -> int: + def prepare_result_frames(self, result: StepResult) -> tuple[np.ndarray, ...]: + if self._closed: + return () + return tuple(self._frame_converter(result)) + + async def enqueue_frames(self, frames: Sequence[np.ndarray]) -> int: if self._closed: return 0 - frames = await asyncio.to_thread(self._frame_converter, result) for i, frame in enumerate(frames): if self._closed: return i await self._frames.put(frame) return len(frames) + async def enqueue_result(self, result: StepResult) -> int: + if self._closed: + return 0 + frames = await asyncio.to_thread(self.prepare_result_frames, result) + return await self.enqueue_frames(frames) + + async def flush(self) -> None: + """Drop queued frames while keeping the RTP timestamp sequence alive.""" + if self._closed: + return + while True: + try: + self._frames.get_nowait() + except asyncio.QueueEmpty: + break + self._next_deadline_s = None + async def recv(self) -> VideoFrame: if self._closed: raise MediaStreamError @@ -234,6 +255,17 @@ def enqueue_encoded_packet_nowait(self, packet: Packet) -> bool: self._packets.put_nowait(packet) return True + async def flush(self) -> None: + """Drop queued encoded packets while preserving the open media track.""" + if self._closed: + return + while True: + try: + self._packets.get_nowait() + except asyncio.QueueEmpty: + break + self._next_deadline_s = None + async def recv(self) -> Packet: if self._closed: raise MediaStreamError diff --git a/flashdreams/flashdreams/serving/webrtc/nvenc.py b/flashdreams/flashdreams/serving/webrtc/nvenc.py index 3a4cc93dc..94a52200f 100644 --- a/flashdreams/flashdreams/serving/webrtc/nvenc.py +++ b/flashdreams/flashdreams/serving/webrtc/nvenc.py @@ -23,6 +23,7 @@ import contextlib import time from collections.abc import Callable +from dataclasses import dataclass from fractions import Fraction from typing import TYPE_CHECKING, Any @@ -30,6 +31,7 @@ from aiortc import MediaStreamTrack from av.packet import Packet from loguru import logger +from torch import Tensor from flashdreams.runtime import StepResult from flashdreams.serving.webrtc.encoders import ChunkDeliveryResult @@ -57,6 +59,13 @@ _RTP_VIDEO_CLOCK = 90_000 +@dataclass(frozen=True, slots=True) +class NVENCChunkPayload: + """Encoder-owned CUDA frames prepared before async delivery is scheduled.""" + + frames: Tensor + + def _payload_contains_nal_type(payload: bytes, nal_type: int) -> bool: """Scan an Annex-B H.264 payload for the presence of a specific NAL type.""" i = 0 @@ -240,6 +249,35 @@ def create_track(self, *, maxsize: int) -> NVENCVideoTrack: return NVENCVideoTrack(fps=self.fps, maxsize=maxsize) + def prepare_chunk_payload( + self, + result: StepResult, + track: MediaStreamTrack, + ) -> NVENCChunkPayload: + from flashdreams.serving.webrtc.media import NVENCVideoTrack + + if not isinstance(track, NVENCVideoTrack): + raise TypeError( + "PyNvHardwareEncoder requires an NVENCVideoTrack; got " + f"{type(track).__name__}. Create it via encoder.create_track()." + ) + return NVENCChunkPayload(frames=_result_to_abgr_frames(result)) + + async def deliver_prepared_chunk( + self, + payload: object, + track: MediaStreamTrack, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + if not isinstance(payload, NVENCChunkPayload): + raise TypeError("PyNvHardwareEncoder payload must be an NVENCChunkPayload.") + return await self._deliver_prepared_frames( + payload.frames, + track, + force_keyframe=force_keyframe, + ) + async def deliver_chunk( self, result: StepResult, @@ -297,6 +335,63 @@ def _stream(packet: Packet) -> None: encode_ms=encode_ms, ) + async def _deliver_prepared_frames( + self, + frames: Tensor, + track: MediaStreamTrack, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + from flashdreams.serving.webrtc.media import NVENCVideoTrack + + if not isinstance(track, NVENCVideoTrack): + raise TypeError( + "PyNvHardwareEncoder requires an NVENCVideoTrack; got " + f"{type(track).__name__}. Create it via encoder.create_track()." + ) + loop = asyncio.get_running_loop() + emitted = 0 + enqueued = 0 + + def _stream(packet: Packet) -> None: + nonlocal emitted, enqueued + emitted += 1 + enqueue = track.enqueue_encoded_packet(packet) + try: + future = asyncio.run_coroutine_threadsafe( + enqueue, + loop, + ) + except RuntimeError: + enqueue.close() + return + try: + accepted = future.result() + except Exception: + return + if accepted: + enqueued += 1 + + _num_frames, num_keyframes, encode_ms = await asyncio.to_thread( + self.encode_frames_sync, + frames, + force_keyframe=force_keyframe, + on_packet=_stream, + ) + if enqueued < emitted: + logger.debug( + "NVENC track closed while enqueueing encoded chunk; " + "enqueued {} of {} packet(s).", + enqueued, + emitted, + ) + return ChunkDeliveryResult( + backend=self.backend, + num_frames=enqueued, + num_keyframes=num_keyframes, + encode_ms=encode_ms, + ) + def encode_chunk_sync( self, result: StepResult, @@ -311,6 +406,20 @@ def encode_chunk_sync( just to get access to the emitted packets. """ frames = _result_to_abgr_frames(result) + return self.encode_frames_sync( + frames, + force_keyframe=force_keyframe, + on_packet=on_packet, + ) + + def encode_frames_sync( + self, + frames: Tensor, + *, + force_keyframe: bool = False, + on_packet: Callable[[Packet], None] | None = None, + ) -> tuple[int, int, float]: + """Encode preconverted ``ABGR`` frames for prepared async delivery.""" if not frames.is_cuda: raise ValueError("expected CUDA tensor for hardware encode path") num_frames = frames.shape[0] diff --git a/flashdreams/flashdreams/serving/webrtc/runtime.py b/flashdreams/flashdreams/serving/webrtc/runtime.py index 70c16fb38..32c7d4365 100644 --- a/flashdreams/flashdreams/serving/webrtc/runtime.py +++ b/flashdreams/flashdreams/serving/webrtc/runtime.py @@ -37,6 +37,7 @@ class WebRTCControlSignal(IntEnum): CLOSE = 3 EVENT = 4 SESSION_STEP = 5 + SESSION_CLOSE = 6 EXIT = 99 diff --git a/flashdreams/flashdreams/serving/webrtc/server.py b/flashdreams/flashdreams/serving/webrtc/server.py index bf257701d..7a2c6694d 100644 --- a/flashdreams/flashdreams/serving/webrtc/server.py +++ b/flashdreams/flashdreams/serving/webrtc/server.py @@ -91,10 +91,12 @@ async def healthz(request: web.Request) -> web.StreamResponse: ) async def ui_config(_: web.Request) -> web.StreamResponse: - adapter_module = None + payload: dict[str, str | None] = {"adapter_module": None} if model_web_dir is not None and (model_web_dir / "adapter.js").is_file(): - adapter_module = "/model-static/adapter.js?v=model-ui-v1" - return web.json_response({"adapter_module": adapter_module}) + payload["adapter_module"] = "/model-static/adapter.js?v=model-ui-v2" + if model_web_dir is not None and (model_web_dir / "adapter.css").is_file(): + payload["model_stylesheet"] = "/model-static/adapter.css?v=model-ui-v2" + return web.json_response(payload) async def on_startup(app: web.Application) -> None: manager = app[SESSION_MANAGER_KEY] diff --git a/flashdreams/flashdreams/serving/webrtc/services.py b/flashdreams/flashdreams/serving/webrtc/services.py new file mode 100644 index 000000000..3ce38e4be --- /dev/null +++ b/flashdreams/flashdreams/serving/webrtc/services.py @@ -0,0 +1,1161 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Session-edge services for the shared WebRTC demo run mode. + +These classes are the Phase 12 decomposition layer: they translate WebRTC +transport facts into the shared demo runtime contracts without owning model +execution. The production manager still uses its legacy execution hook until +the realtime-driver adoption phase. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import inspect +import json +import math +import threading +from collections import deque +from collections.abc import Callable, Coroutine, Mapping, MutableSet, Sequence +from concurrent.futures import Future +from dataclasses import dataclass, field +from typing import Any, Literal, Protocol, runtime_checkable + +from flashdreams.runtime import ( + StepRequirements, + StepResult, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.demo import ( + AsyncSessionDriver, + DemoAdapter, + DemoSpec, + InMemorySessionMetricsRecorder, + ModelInputProvider, + ModelWarmupPlan, + OutputDecision, + PreparedScenario, + RealtimeSessionDriver, + RealtimeWindowResult, + RunContext, + RunModeCapabilities, + RunResult, + RuntimeHost, + SessionEdges, + SessionInfo, + SingleSessionAdmissionPolicy, + StepPipeline, + UserInputWindow, + WebRTCErrorPolicy, + input_frame_count_from_request, + run_demo_session_async, +) +from flashdreams.runtime.demo.timing import ( + SPARSE_KEY_SEGMENTS_METADATA_KEY, + ActivationResult, + CatchUpPolicy, + DeterministicClock, + RealtimeClock, +) + +from .messages import ( + MESSAGE_TYPE_ACTION, + MESSAGE_TYPE_DISCONNECT, + MESSAGE_TYPE_EVENT, + MESSAGE_TYPE_HEARTBEAT, +) +from .server import SessionBusyError + +WebRTCMessageKind = Literal[ + "action", + "disconnect", + "event", + "heartbeat", + "error", +] + +WebRTCDropPolicy = Literal["none", "drop_newest", "drop_oldest"] + +_CLEAR_EVENT_STATES = frozenset({"clear", "release", "off", "none"}) +WEBRTC_SKIPPED_INPUTS_METADATA_KEY = "webrtc_skipped_inputs" +WEBRTC_SKIPPED_WINDOW_METADATA_KEY = "webrtc_skipped_window" + + +@dataclass(frozen=True, kw_only=True, slots=True) +class WebRTCMessageResult: + """Result of translating one browser data-channel message.""" + + kind: WebRTCMessageKind + activated: bool = False + error: str | None = None + + +@dataclass(frozen=True, kw_only=True, slots=True) +class WebRTCOfferRequest: + """Browser SDP offer passed to the shared offer/session handler.""" + + sdp: str + type: str + + def __post_init__(self) -> None: + if not self.sdp.strip(): + raise ValueError("WebRTCOfferRequest.sdp must be non-empty.") + if not self.type.strip(): + raise ValueError("WebRTCOfferRequest.type must be non-empty.") + + +@dataclass(frozen=True, kw_only=True, slots=True) +class WebRTCOutputBridgeDecision: + """Immediate delivery decision from a nonblocking WebRTC output bridge.""" + + accepted: bool = True + should_stop: bool = False + dropped: bool = False + drop_policy: WebRTCDropPolicy = "none" + backpressure_s: float = 0.0 + metadata: Mapping[str, object] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.drop_policy not in {"none", "drop_newest", "drop_oldest"}: + raise ValueError(f"Unsupported drop_policy={self.drop_policy!r}.") + if not math.isfinite(self.backpressure_s) or self.backpressure_s < 0.0: + raise ValueError( + "WebRTCOutputBridgeDecision.backpressure_s must be finite and >= 0." + ) + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class WebRTCChunkDelivery: + """Completed WebRTC chunk delivery plus the model chunk summary.""" + + delivery: object + step_index: int + frame_count: int + generation: int + force_keyframe: bool + metadata: Mapping[str, object] = field(default_factory=dict) + metrics: Mapping[str, float | int] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__(self, "metadata", freeze_mapping(self.metadata)) + object.__setattr__(self, "metrics", freeze_mapping(self.metrics)) + + +@runtime_checkable +class WebRTCOfferAnswerer(Protocol): + """Creates an SDP answer after the shared session task has been scheduled.""" + + async def create_answer( + self, + *, + offer: WebRTCOfferRequest, + session_task: asyncio.Task[RunResult], + ) -> Mapping[str, str]: ... + + +@runtime_checkable +class BlockingPreparationService(Protocol): + """Runs blocking scenario preparation outside the aiohttp event loop.""" + + async def run( + self, + func: Callable[..., object], + *args: object, + **kwargs: object, + ) -> object: ... + + +@runtime_checkable +class WebRTCOutputBridge(Protocol): + """Thread-safe bridge from model-worker output writes to WebRTC delivery.""" + + def begin_generation(self, generation: int) -> None: ... + + def submit_chunk( + self, + result: StepResult, + *, + generation: int, + force_keyframe: bool = False, + ) -> WebRTCOutputBridgeDecision: ... + + def close(self) -> None: ... + + +@runtime_checkable +class WebRTCSessionEdgeFactory(Protocol): + """Builds per-peer shared session edges on the WebRTC control rank.""" + + def create_session_edges( + self, + *, + context: RunContext, + spec: DemoSpec, + scenario: PreparedScenario, + provider: ModelInputProvider, + adapter: DemoAdapter, + ) -> SessionEdges: ... + + +class AsyncioBlockingPreparationService: + """Default blocking-prep service backed by ``asyncio.to_thread``.""" + + async def run( + self, + func: Callable[..., object], + *args: object, + **kwargs: object, + ) -> object: + return await asyncio.to_thread(func, *args, **kwargs) + + +class WebRTCTransportService: + """Idempotent per-peer transport lifecycle for realtime session edges.""" + + def __init__( + self, + *, + loop: asyncio.AbstractEventLoop | None = None, + on_close: Callable[[str | None], None] | None = None, + ) -> None: + self._loop = loop + self._on_close = on_close + self._closed_signal = _ThreadSafeActivationSignal(loop=loop) + self._lock = threading.Lock() + self._closed = False + self._close_reason: str | None = None + self._close_count = 0 + self.last_client_message_at: float | None = None + + @property + def close_count(self) -> int: + """Number of effective transport closes, after idempotency.""" + + return self._close_count + + @property + def close_reason(self) -> str | None: + return self._close_reason + + @property + def closed_signal(self) -> "_ThreadSafeActivationSignal": + return self._closed_signal + + def mark_client_message(self, timestamp_s: float) -> None: + if not math.isfinite(timestamp_s) or timestamp_s < 0.0: + raise ValueError("timestamp_s must be finite and >= 0.") + self.last_client_message_at = float(timestamp_s) + + def is_active(self) -> bool: + return not self._closed + + def disconnect(self, reason: str = "client disconnected") -> None: + self.close(reason=reason) + + def close(self, reason: str | None = None) -> None: + callback: Callable[[str | None], None] | None = None + with self._lock: + if self._closed: + return + self._closed = True + self._close_reason = reason + self._close_count += 1 + callback = self._on_close + self._closed_signal.set() + if callback is not None: + callback(reason) + + +@dataclass(slots=True) +class WebRTCActivationPolicy: + """Activate on the first browser action/event or stop on disconnect.""" + + input_source: "WebRTCInputSource" + transport: WebRTCTransportService + timeout_s: float | None = None + timeout_reason: str = "activation timed out" + anchor_clock: bool = True + + def __post_init__(self) -> None: + if self.timeout_s is not None and self.timeout_s <= 0.0: + raise ValueError("timeout_s must be > 0 when set.") + if not self.timeout_reason.strip(): + raise ValueError("timeout_reason must be non-empty.") + + async def wait_until_active( + self, + clock: RealtimeClock | DeterministicClock, + ) -> ActivationResult: + if self.input_source.activation_signal.is_set(): + self._anchor(clock) + return ActivationResult(activated=True) + if not self.transport.is_active(): + return ActivationResult( + activated=False, + reason=self.transport.close_reason or "transport closed", + ) + + activation_task = asyncio.create_task( + self.input_source.activation_signal.wait() + ) + closed_task = asyncio.create_task(self.transport.closed_signal.wait()) + try: + done, pending = await asyncio.wait( + {activation_task, closed_task}, + timeout=self.timeout_s, + return_when=asyncio.FIRST_COMPLETED, + ) + if not done: + return ActivationResult( + activated=False, + reason=self.timeout_reason, + ) + for task in done: + task.result() + if not self.transport.is_active(): + return ActivationResult( + activated=False, + reason=self.transport.close_reason or "transport closed", + ) + self._anchor(clock) + return ActivationResult(activated=True) + finally: + for task in (activation_task, closed_task): + if not task.done(): + task.cancel() + await asyncio.gather( + activation_task, + closed_task, + return_exceptions=True, + ) + + def _anchor(self, clock: RealtimeClock | DeterministicClock) -> None: + if not self.anchor_clock or not clock.is_realtime: + return + now = getattr(clock, "now", None) + anchor = getattr(clock, "anchor", None) + if callable(now) and callable(anchor): + anchor(now()) + + +@dataclass(slots=True) +class WebRTCInputSource: + """Realtime source fed by browser data-channel events.""" + + resampler: Any + max_lag_s: float | None = None + catch_up_policy: CatchUpPolicy = "fold" + user_input_schema: UserInputSchema = field( + default_factory=lambda: WEBRTC_USER_INPUT_SCHEMA + ) + is_finite: bool = False + is_deterministic: bool = False + _activation_signal: "_ThreadSafeActivationSignal" = field( + default_factory=lambda: _ThreadSafeActivationSignal(), + init=False, + repr=False, + ) + _events: deque[UserInputEvent] = field( + default_factory=deque, + init=False, + repr=False, + ) + + def __post_init__(self) -> None: + if self.max_lag_s is not None and ( + not math.isfinite(self.max_lag_s) or self.max_lag_s < 0.0 + ): + raise ValueError("max_lag_s must be finite and >= 0.") + if self.catch_up_policy != "fold": + raise NotImplementedError( + f"Catch-up policy {self.catch_up_policy!r} has no WebRTC analog yet." + ) + + @property + def activation_signal(self) -> "_ThreadSafeActivationSignal": + return self._activation_signal + + def is_finished(self) -> bool: + return False + + def reset(self, *, start_v: float) -> None: + self.resampler.reset(start_v=start_v) + self._events.clear() + self._activation_signal.clear() + + def handle_browser_message( + self, + raw_message: object, + *, + timestamp_s: float, + ) -> WebRTCMessageResult: + """Translate one browser data-channel message into typed user inputs.""" + + if not isinstance(raw_message, str): + return WebRTCMessageResult(kind="error", error="Expected text payload.") + try: + payload = json.loads(raw_message) + except json.JSONDecodeError: + return WebRTCMessageResult(kind="error", error="Invalid JSON payload.") + if not isinstance(payload, dict): + return WebRTCMessageResult( + kind="error", + error="Payload must be a JSON object.", + ) + return self.handle_browser_payload(payload, timestamp_s=timestamp_s) + + def handle_browser_payload( + self, + payload: Mapping[str, object], + *, + timestamp_s: float, + ) -> WebRTCMessageResult: + message_type = str(payload.get("type", "")).strip().lower() + if message_type == MESSAGE_TYPE_HEARTBEAT: + return WebRTCMessageResult(kind="heartbeat") + if message_type == MESSAGE_TYPE_DISCONNECT: + return WebRTCMessageResult(kind="disconnect") + if message_type == MESSAGE_TYPE_EVENT: + return self._record_text_event(payload, timestamp_s=timestamp_s) + if message_type == MESSAGE_TYPE_ACTION: + action_payload = payload.get("action", payload) + if not isinstance(action_payload, Mapping): + return WebRTCMessageResult( + kind="error", + error="'action' must be an object.", + ) + return self._record_action( + {str(key): value for key, value in action_payload.items()}, + timestamp_s=timestamp_s, + ) + return WebRTCMessageResult( + kind="error", + error=( + "Unsupported message type, expected " + "'action', 'event', 'heartbeat', or 'disconnect'." + ), + ) + + def record_user_event( + self, + *, + timestamp_s: float, + event_type: str, + payload: Mapping[str, object], + source_event_id: str | None = None, + activate: bool = True, + ) -> None: + event = UserInputEvent( + timestamp_s=timestamp_s, + event_type=event_type, + payload=dict(payload), + source="webrtc", + source_event_id=source_event_id, + ) + self.user_input_schema.validate_event(event) + self._events.append(event) + if activate: + self._activation_signal.set() + + async def next_realtime_window( + self, + *, + request: StepRequirements, + clock: RealtimeClock, + ) -> RealtimeWindowResult: + input_frame_count = input_frame_count_from_request(request) + chunk_duration_s = input_frame_count * float(self.resampler.dt) + if chunk_duration_s <= 0.0: + raise ValueError("Realtime resampler dt must produce a positive window.") + + window_end_s = float(self.resampler.next_chunk_start_v) + chunk_duration_s + await clock.wait_until_window_end(window_end_s) + pre_catch_up_start_s = float(self.resampler.next_chunk_start_v) + catch_up = clock.catch_up( + request=request, + max_lag_s=self.max_lag_s + if self.max_lag_s is not None + else chunk_duration_s, + policy=self.catch_up_policy, + ) + start_s = float(self.resampler.next_chunk_start_v) + segments, frame_times = self.resampler.sample_chunk(input_frame_count) + end_s = float(self.resampler.next_chunk_start_v) + metadata: dict[str, object] = { + SPARSE_KEY_SEGMENTS_METADATA_KEY: tuple(segments), + } + if start_s > pre_catch_up_start_s: + metadata[WEBRTC_SKIPPED_INPUTS_METADATA_KEY] = UserInputs( + events=self._events_for_window(pre_catch_up_start_s, start_s) + ) + metadata[WEBRTC_SKIPPED_WINDOW_METADATA_KEY] = ( + pre_catch_up_start_s, + start_s, + ) + window = RealtimeWindowResult( + window=_user_input_window( + start_s=start_s, + end_s=end_s, + frame_times=tuple(frame_times), + inputs=UserInputs(events=self._events_for_window(start_s, end_s)), + metadata=metadata, + ), + catch_up=catch_up, + ) + self._prune_events(before_s=start_s) + return window + + def _record_action( + self, + payload: Mapping[str, object], + *, + timestamp_s: float, + ) -> WebRTCMessageResult: + event = str(payload.get("event", "")).strip().lower() + if event == "step": + self._activation_signal.set() + return WebRTCMessageResult(kind="action", activated=True) + if event not in {"keydown", "keyup"}: + return WebRTCMessageResult( + kind="error", + error=f"Unsupported event={event!r}; expected 'keydown' or 'keyup'.", + ) + key = str(payload.get("key", "")).strip() + if not key: + return WebRTCMessageResult( + kind="error", + error="Action payload must include non-empty 'key'.", + ) + self.resampler.on_edge(arrival_t=timestamp_s, event=event, key=key) + self.record_user_event( + timestamp_s=timestamp_s, + event_type="key_down" if event == "keydown" else "key_up", + payload={"key": key}, + ) + return WebRTCMessageResult(kind="action", activated=True) + + def _record_text_event( + self, + payload: Mapping[str, object], + *, + timestamp_s: float, + ) -> WebRTCMessageResult: + state = str(payload.get("state", "trigger")).strip().lower() or "trigger" + event_id = str(payload.get("event_id", payload.get("id", ""))).strip() + clears = state in _CLEAR_EVENT_STATES + if not event_id and not clears: + return WebRTCMessageResult( + kind="error", + error=( + "Event payload must include non-empty 'event_id' unless state " + "clears the active event." + ), + ) + active_event_id = None if clears else event_id + self.record_user_event( + timestamp_s=timestamp_s, + event_type="text_event", + payload={"event_id": active_event_id, "state": state}, + source_event_id=active_event_id, + ) + return WebRTCMessageResult(kind="event", activated=True) + + def _events_for_window( + self, + start_s: float, + end_s: float, + ) -> tuple[UserInputEvent, ...]: + return tuple( + sorted( + ( + event + for event in self._events + if start_s <= event.timestamp_s < end_s + ), + key=lambda event: event.timestamp_s, + ) + ) + + def _prune_events(self, *, before_s: float) -> None: + self._events = deque( + event for event in self._events if event.timestamp_s >= before_s + ) + + +class WebRTCOutputSink: + """Output sink that schedules WebRTC media delivery without blocking.""" + + produces_artifacts = False + + def __init__(self, *, bridge: WebRTCOutputBridge) -> None: + self._bridge = bridge + self._opened = False + self._closed = True + self._bridge_closed = False + self._generation = 0 + self._force_keyframe = False + self.session_info: SessionInfo | None = None + + def open(self, session_info: SessionInfo) -> None: + self.session_info = session_info + self._opened = True + self._closed = False + self._generation = 0 + self._force_keyframe = True + self._bridge.begin_generation(0) + + def begin_generation(self, generation: int) -> None: + if generation < 0: + raise ValueError("generation must be >= 0.") + self._generation = generation + self._force_keyframe = True + self._bridge.begin_generation(generation) + + def write(self, result: StepResult) -> OutputDecision: + if not self._opened or self._closed: + raise RuntimeError("Cannot write to a closed output sink.") + decision = self._bridge.submit_chunk( + result, + generation=self._generation, + force_keyframe=self._force_keyframe, + ) + self._force_keyframe = False + return OutputDecision( + should_stop=decision.should_stop, + dropped=decision.dropped, + drop_policy=decision.drop_policy, + backpressure_s=decision.backpressure_s, + metadata=decision.metadata, + ) + + def close(self) -> Sequence[Any]: + if self._bridge_closed: + return () + self._closed = True + self._opened = False + self._bridge.close() + self._bridge_closed = True + return () + + +class ThreadSafeWebRTCOutputBridge: + """Schedule async encoder delivery from any thread without blocking writes.""" + + def __init__( + self, + *, + loop: asyncio.AbstractEventLoop, + video_encoder: Any, + video_track: Any, + max_pending_chunks: int = 2, + close_track: bool = True, + on_delivery: Callable[[object], None] | None = None, + on_chunk_delivery: Callable[[WebRTCChunkDelivery], None] | None = None, + on_error: Callable[[BaseException], None] | None = None, + ) -> None: + if max_pending_chunks <= 0: + raise ValueError("max_pending_chunks must be > 0.") + self._loop = loop + self._video_encoder = video_encoder + self._video_track = video_track + self._max_pending_chunks = max_pending_chunks + self._close_track = close_track + self._on_delivery = on_delivery + self._on_chunk_delivery = on_chunk_delivery + self._on_error = on_error + self._pending: dict[Future[WebRTCChunkDelivery], int] = {} + self._lock = threading.Lock() + self._closed = False + self._generation = 0 + + @property + def pending_count(self) -> int: + with self._lock: + return len(self._pending) + + def begin_generation(self, generation: int) -> None: + if generation < 0: + raise ValueError("generation must be >= 0.") + with self._lock: + if self._closed or generation <= self._generation: + return + self._generation = generation + stale = tuple( + future + for future, future_generation in self._pending.items() + if future_generation < generation + ) + for future in stale: + future.cancel() + self._schedule_track_flush() + + def submit_chunk( + self, + result: StepResult, + *, + generation: int, + force_keyframe: bool = False, + ) -> WebRTCOutputBridgeDecision: + prepare = getattr(self._video_encoder, "prepare_chunk_payload", None) + deliver = getattr(self._video_encoder, "deliver_prepared_chunk", None) + if not callable(prepare) or not callable(deliver): + raise TypeError( + "ThreadSafeWebRTCOutputBridge requires a video encoder with " + "prepare_chunk_payload(...) and deliver_prepared_chunk(...)." + ) + with self._lock: + if self._closed: + return WebRTCOutputBridgeDecision( + accepted=False, + should_stop=True, + dropped=True, + drop_policy="drop_newest", + metadata={"reason": "closed"}, + ) + if generation < self._generation: + return WebRTCOutputBridgeDecision( + accepted=False, + dropped=True, + drop_policy="drop_newest", + metadata={"reason": "stale generation"}, + ) + if len(self._pending) >= self._max_pending_chunks: + return WebRTCOutputBridgeDecision( + accepted=False, + dropped=True, + drop_policy="drop_newest", + metadata={"reason": "pending queue full"}, + ) + payload = prepare(result, self._video_track) + chunk = WebRTCChunkDelivery( + delivery=None, + step_index=result.step_index, + frame_count=result.frame_count, + generation=generation, + force_keyframe=force_keyframe, + metadata=result.metadata, + metrics=result.metrics, + ) + with self._lock: + if self._closed: + return WebRTCOutputBridgeDecision( + accepted=False, + should_stop=True, + dropped=True, + drop_policy="drop_newest", + metadata={"reason": "closed"}, + ) + if generation < self._generation: + return WebRTCOutputBridgeDecision( + accepted=False, + dropped=True, + drop_policy="drop_newest", + metadata={"reason": "stale generation"}, + ) + if len(self._pending) >= self._max_pending_chunks: + return WebRTCOutputBridgeDecision( + accepted=False, + dropped=True, + drop_policy="drop_newest", + metadata={"reason": "pending queue full"}, + ) + future = asyncio.run_coroutine_threadsafe( + self._deliver( + payload, + chunk=chunk, + generation=generation, + force_keyframe=force_keyframe, + ), + self._loop, + ) + self._pending[future] = generation + future.add_done_callback(self._on_done) + + return WebRTCOutputBridgeDecision( + accepted=True, + backpressure_s=self._track_backpressure_s(), + ) + + def close(self) -> None: + with self._lock: + if self._closed: + return + self._closed = True + pending = tuple(self._pending) + for future in pending: + future.cancel() + if self._close_track: + self._schedule_track_close() + + async def _deliver( + self, + payload: object, + *, + chunk: WebRTCChunkDelivery, + generation: int, + force_keyframe: bool, + ) -> WebRTCChunkDelivery: + with self._lock: + if self._closed or generation < self._generation: + raise asyncio.CancelledError + delivery = await self._video_encoder.deliver_prepared_chunk( + payload, + self._video_track, + force_keyframe=force_keyframe, + ) + with self._lock: + if self._closed or generation < self._generation: + stale_after_delivery = True + else: + stale_after_delivery = False + if stale_after_delivery: + self._schedule_track_flush() + raise asyncio.CancelledError + return WebRTCChunkDelivery( + delivery=delivery, + step_index=chunk.step_index, + frame_count=chunk.frame_count, + generation=chunk.generation, + force_keyframe=chunk.force_keyframe, + metadata=chunk.metadata, + metrics=chunk.metrics, + ) + + def _on_done(self, future: Future[WebRTCChunkDelivery]) -> None: + with self._lock: + self._pending.pop(future, None) + if future.cancelled(): + return + try: + result = future.result() + except BaseException as exc: + if self._on_error is not None: + self._on_error(exc) + return + if self._on_delivery is not None: + self._on_delivery(result.delivery) + if self._on_chunk_delivery is not None: + self._on_chunk_delivery(result) + + def _track_backpressure_s(self) -> float: + qsize = getattr(self._video_track, "qsize", None) + fps = getattr(self._video_track, "fps", None) or getattr( + self._video_encoder, + "fps", + None, + ) + if not callable(qsize) or fps is None: + return 0.0 + try: + queue_depth = int(qsize()) + frames_per_second = float(fps) + except (TypeError, ValueError): + return 0.0 + if frames_per_second <= 0.0: + return 0.0 + return max(0.0, queue_depth / frames_per_second) + + def _schedule_track_close(self) -> None: + close = getattr(self._video_track, "close", None) + if not callable(close): + return + try: + result = close() + if inspect.isawaitable(result): + asyncio.run_coroutine_threadsafe(result, self._loop) + except BaseException as exc: + if self._on_error is not None: + self._on_error(exc) + + def _schedule_track_flush(self) -> None: + flush = getattr(self._video_track, "flush", None) + if not callable(flush): + return + try: + result = flush() + if inspect.isawaitable(result): + asyncio.run_coroutine_threadsafe(result, self._loop) + except BaseException as exc: + if self._on_error is not None: + self._on_error(exc) + + +class WebRTCRunMode: + """Shared realtime run mode that delegates peer-specific edges to WebRTC.""" + + name = "webrtc" + capabilities = RunModeCapabilities( + realtime=True, + supports_backpressure=True, + supports_interactive_events=True, + ) + + def __init__( + self, + *, + edge_factory: WebRTCSessionEdgeFactory, + blocking_preparation: BlockingPreparationService | None = None, + driver: AsyncSessionDriver | None = None, + error_policy: WebRTCErrorPolicy | None = None, + ) -> None: + self._edge_factory = edge_factory + self._blocking_preparation = ( + blocking_preparation or AsyncioBlockingPreparationService() + ) + self._driver = driver or RealtimeSessionDriver() + self._error_policy = error_policy or WebRTCErrorPolicy() + + @property + def blocking_preparation(self) -> BlockingPreparationService: + return self._blocking_preparation + + @property + def error_policy(self) -> WebRTCErrorPolicy: + return self._error_policy + + def validate_run(self, *, spec: DemoSpec, adapter: DemoAdapter) -> None: + del adapter + if spec.output.mode != "webrtc": + raise ValueError("WebRTCRunMode requires WebRTC output.") + + def validate_session( + self, + *, + spec: DemoSpec, + scenario: PreparedScenario, + adapter: DemoAdapter, + provider: ModelInputProvider, + ) -> None: + del spec, scenario, adapter + if not provider.capabilities.supports_realtime_clock: + raise ValueError("WebRTC providers must support realtime clocks.") + + def create_run_context( + self, + *, + spec: DemoSpec, + adapter: DemoAdapter, + host: RuntimeHost, + model_warmup_plan: ModelWarmupPlan, + ) -> RunContext: + del spec, adapter + services: dict[str, object] = {} + if host.is_control_rank: + services["blocking_preparation"] = self._blocking_preparation + return RunContext( + host=host, + run_metrics=InMemorySessionMetricsRecorder(), + admission=SingleSessionAdmissionPolicy( + health_check=lambda: host.is_control_rank and host.is_healthy + ), + model_warmup_plan=model_warmup_plan, + services=services, + ) + + def create_session_edges( + self, + *, + context: RunContext, + spec: DemoSpec, + scenario: PreparedScenario, + provider: ModelInputProvider, + adapter: DemoAdapter, + ) -> SessionEdges: + if not context.host.is_control_rank: + raise RuntimeError("WebRTC session edges are control-rank only.") + edges = self._edge_factory.create_session_edges( + context=context, + spec=spec, + scenario=scenario, + provider=provider, + adapter=adapter, + ) + if not isinstance(edges, SessionEdges): + raise TypeError( + "WebRTC edge factory must return SessionEdges, " + f"got {type(edges).__name__}." + ) + return edges + + def select_driver(self) -> AsyncSessionDriver: + return self._driver + + +class WebRTCSessionOfferHandler: + """Reserve, prepare, and launch one WebRTC session before SDP negotiation.""" + + def __init__( + self, + *, + context: RunContext, + spec: DemoSpec, + adapter: DemoAdapter, + run_mode: WebRTCRunMode, + answerer: WebRTCOfferAnswerer, + pipeline: StepPipeline | None = None, + session_helper: Callable[..., Coroutine[Any, Any, RunResult]] | None = None, + busy_message: str = "Another WebRTC session is already active.", + session_tasks: MutableSet[asyncio.Task[RunResult]] | None = None, + ) -> None: + self._context = context + self._spec = spec + self._adapter = adapter + self._run_mode = run_mode + self._answerer = answerer + self._pipeline = pipeline or StepPipeline() + self._session_helper = session_helper or run_demo_session_async + self._busy_message = busy_message + self._session_tasks = session_tasks if session_tasks is not None else set() + + async def handle_offer( + self, + *, + offer_sdp: str, + offer_type: str, + ) -> Mapping[str, str]: + if not self._context.host.is_control_rank: + raise RuntimeError("WebRTC offers are handled only on the control rank.") + reservation = self._context.admission.try_reserve() + if reservation is None: + raise SessionBusyError(self._busy_message) + + task: asyncio.Task[RunResult] | None = None + try: + scenario = await self._run_blocking_prepare(self._spec) + task = asyncio.create_task( + self._session_helper( + context=self._context, + spec=self._spec, + scenario=scenario, + adapter=self._adapter, + run_mode=self._run_mode, + pipeline=self._pipeline, + reservation=reservation, + ) + ) + self._track_task(task) + answer = await self._answerer.create_answer( + offer=WebRTCOfferRequest(sdp=offer_sdp, type=offer_type), + session_task=task, + ) + return dict(answer) + except Exception: + if task is not None: + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + reservation.release() + raise + + async def _run_blocking_prepare(self, spec: DemoSpec) -> PreparedScenario: + service = self._run_mode.blocking_preparation + result = await service.run(self._adapter.prepare_scenario, spec) + if not isinstance(result, PreparedScenario): + raise TypeError( + "DemoAdapter.prepare_scenario must return PreparedScenario, " + f"got {type(result).__name__}." + ) + return result + + def _track_task(self, task: asyncio.Task[RunResult]) -> None: + self._session_tasks.add(task) + task.add_done_callback(self._discard_task) + + def _discard_task(self, task: asyncio.Task[RunResult]) -> None: + self._session_tasks.discard(task) + if task.cancelled(): + return + with contextlib.suppress(Exception): + task.exception() + + +WEBRTC_USER_INPUT_SCHEMA = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="key_down", + input_modality="keyboard", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_type="key_up", + input_modality="keyboard", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_type="text_event", + input_modality="text", + payload_fields=frozenset({"event_id", "state"}), + ), + ), + description="browser WebRTC data-channel events", +) + + +def _user_input_window( + *, + start_s: float, + end_s: float, + frame_times: Sequence[float], + inputs: UserInputs, + metadata: Mapping[str, object], +) -> UserInputWindow: + return UserInputWindow( + start_s=start_s, + end_s=end_s, + frame_times=frame_times, + inputs=inputs, + metadata=metadata, + ) + + +class _ThreadSafeActivationSignal: + def __init__(self, *, loop: asyncio.AbstractEventLoop | None = None) -> None: + self._loop = loop + self._event = asyncio.Event() + + def is_set(self) -> bool: + return self._event.is_set() + + async def wait(self) -> object: + return await self._event.wait() + + def set(self) -> None: + if self._event.is_set(): + return + if self._loop is not None and self._loop.is_running(): + self._loop.call_soon_threadsafe(self._event.set) + return + self._event.set() + + def clear(self) -> None: + self._event.clear() + + +__all__ = [ + "AsyncioBlockingPreparationService", + "BlockingPreparationService", + "ThreadSafeWebRTCOutputBridge", + "WEBRTC_USER_INPUT_SCHEMA", + "WEBRTC_SKIPPED_INPUTS_METADATA_KEY", + "WEBRTC_SKIPPED_WINDOW_METADATA_KEY", + "WebRTCActivationPolicy", + "WebRTCInputSource", + "WebRTCMessageResult", + "WebRTCOfferAnswerer", + "WebRTCOfferRequest", + "WebRTCOutputBridge", + "WebRTCOutputBridgeDecision", + "WebRTCChunkDelivery", + "WebRTCOutputSink", + "WebRTCRunMode", + "WebRTCSessionEdgeFactory", + "WebRTCSessionOfferHandler", + "WebRTCTransportService", +] diff --git a/flashdreams/flashdreams/serving/webrtc/warmup.py b/flashdreams/flashdreams/serving/webrtc/warmup.py index 05f3d1c10..5f59c4a86 100644 --- a/flashdreams/flashdreams/serving/webrtc/warmup.py +++ b/flashdreams/flashdreams/serving/webrtc/warmup.py @@ -13,6 +13,8 @@ from aiortc.mediastreams import MediaStreamError from loguru import logger as loguru_logger +from .messages import MESSAGE_TYPE_CHUNK_DONE, MESSAGE_TYPE_ERROR + class CreateAnswerCallback(Protocol): async def __call__(self, *, offer_sdp: str, offer_type: str) -> dict[str, str]: ... @@ -47,14 +49,30 @@ async def run_loopback_warmup_session( client_peer.addTransceiver("video", direction="recvonly") channel_open = asyncio.Event() warmup_done = asyncio.Event() + warmup_failure: str | None = None received_chunks = 0 drain_tasks: set[asyncio.Task[Any]] = set() heartbeat_task: asyncio.Task[Any] | None = None + def fail_warmup(reason: str) -> None: + nonlocal warmup_failure + if warmup_done.is_set(): + return + warmup_failure = reason + warmup_done.set() + @control_channel.on("open") def on_open() -> None: channel_open.set() + @control_channel.on("close") + def on_close() -> None: + if received_chunks < num_chunks: + fail_warmup( + f"{label} loopback warmup data channel closed before warmup " + f"completed ({received_chunks}/{num_chunks} chunk(s))." + ) + @control_channel.on("message") def on_message(message: Any) -> None: nonlocal received_chunks @@ -64,7 +82,14 @@ def on_message(message: Any) -> None: payload = json.loads(message) except json.JSONDecodeError: return - if not isinstance(payload, dict) or payload.get("type") != "chunk_done": + if not isinstance(payload, dict): + return + message_type = payload.get("type") + if message_type == MESSAGE_TYPE_ERROR: + message_text = str(payload.get("message", "unknown error")) + fail_warmup(f"{label} loopback warmup failed: {message_text}") + return + if message_type != MESSAGE_TYPE_CHUNK_DONE: return received_chunks += 1 logger.info( @@ -109,6 +134,8 @@ def on_track(track: Any) -> None: for action_payload in action_payloads: control_channel.send(json.dumps(action_payload)) await asyncio.wait_for(warmup_done.wait(), timeout=warmup_timeout_s) + if warmup_failure is not None: + raise RuntimeError(warmup_failure) finally: if heartbeat_task is not None: heartbeat_task.cancel() diff --git a/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py b/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py index 735179a4a..0c62855b5 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py +++ b/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py @@ -83,13 +83,15 @@ def do_HEAD(self) -> None: def _serve_ui_config(self) -> bool: if urlsplit(self.path).path != "/api/ui/config": return False - adapter_module = ( - "/model-static/adapter.js?v=model-ui-v1" - if self.model_web_dir is not None - and (self.model_web_dir / "adapter.js").is_file() - else None - ) - payload = json.dumps({"adapter_module": adapter_module}).encode("utf-8") + ui_config: dict[str, str | None] = {"adapter_module": None} + if self.model_web_dir is not None: + if (self.model_web_dir / "adapter.js").is_file(): + ui_config["adapter_module"] = "/model-static/adapter.js?v=model-ui-v2" + if (self.model_web_dir / "adapter.css").is_file(): + ui_config["model_stylesheet"] = ( + "/model-static/adapter.css?v=model-ui-v2" + ) + payload = json.dumps(ui_config).encode("utf-8") self.send_response(200) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(payload))) diff --git a/flashdreams/flashdreams/serving/webrtc/web/request_session.html b/flashdreams/flashdreams/serving/webrtc/web/request_session.html index ee42c82a0..ad158656e 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/request_session.html +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.html @@ -87,6 +87,6 @@

Client Logs

- + diff --git a/flashdreams/flashdreams/serving/webrtc/web/request_session.js b/flashdreams/flashdreams/serving/webrtc/web/request_session.js index 95c340b2d..a13384bc5 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/request_session.js +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.js @@ -291,10 +291,14 @@ const modelContext = { async function loadModelAdapter() { let adapter = {} + const stylesheetHrefs = new Set() try { const response = await fetch("/api/ui/config") if (response.ok) { const config = await response.json() + if (typeof config.model_stylesheet === "string" && config.model_stylesheet) { + stylesheetHrefs.add(config.model_stylesheet) + } if (typeof config.adapter_module === "string" && config.adapter_module) { const module = await import(config.adapter_module) if (module.default && typeof module.default === "object") { @@ -308,9 +312,12 @@ async function loadModelAdapter() { modelAdapter = adapter if (typeof adapter.stylesheet === "string" && adapter.stylesheet) { + stylesheetHrefs.add(adapter.stylesheet) + } + for (const href of stylesheetHrefs) { const stylesheet = document.createElement("link") stylesheet.rel = "stylesheet" - stylesheet.href = adapter.stylesheet + stylesheet.href = href document.head.append(stylesheet) } const modelControls = Array.isArray(adapter.controls) ? adapter.controls : [] diff --git a/flashdreams/tests/test_demo_runtime_host.py b/flashdreams/tests/test_demo_runtime_host.py new file mode 100644 index 000000000..b14dafcb6 --- /dev/null +++ b/flashdreams/tests/test_demo_runtime_host.py @@ -0,0 +1,218 @@ +# 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 ( + InferenceInput, + InferenceSession, + StepRequest, + StepResult, +) +from flashdreams.runtime.demo import ModelWarmupPlan, RuntimeHost, WarmupSessionInputs + +pytestmark = pytest.mark.ci_cpu + + +def test_runtime_host_latches_health_and_runs_lifecycle_in_order() -> None: + runtime = _LifecycleRuntime() + host = RuntimeHost(runtime) + error = RuntimeError("runtime wedged") + + host.mark_unhealthy("first failure", error) + host.mark_unhealthy("second failure") + + assert not host.is_healthy + assert host.unhealthy_reason == "first failure" + assert host.unhealthy_error is error + + initial_input = InferenceInput(global_conditioning={"session": "warmup"}) + step_inputs = ( + InferenceInput(step={"step": 0}), + InferenceInput(step={"step": 1}), + ) + host.preload() + host.warmup( + ModelWarmupPlan( + sessions=( + WarmupSessionInputs( + initial_input=initial_input, + step_inputs=step_inputs, + ), + ), + ) + ) + host.close() + host.close() + + assert runtime.events == [ + "initialize_distributed", + "preload", + ("start_session", initial_input), + ("step", step_inputs[0]), + ("step", step_inputs[1]), + "session.close", + "runtime.close", + "close_distributed", + ] + assert not host.is_healthy + with pytest.raises(RuntimeError, match="closed"): + host.call(lambda: None) + + +@pytest.mark.asyncio +async def test_runtime_host_call_async_does_not_block_event_loop() -> None: + runtime = _LifecycleRuntime() + host = RuntimeHost(runtime) + loop_thread_id = threading.get_ident() + started = threading.Event() + release = threading.Event() + heartbeat_ticks = 0 + + def _slow_model_call() -> int: + started.set() + assert release.wait(timeout=2.0) + return threading.get_ident() + + async def _heartbeat_until_done(task: asyncio.Task[int]) -> None: + nonlocal heartbeat_ticks + while not task.done(): + heartbeat_ticks += 1 + await asyncio.sleep(0) + + try: + model_task = asyncio.create_task(host.call_async(_slow_model_call)) + assert await asyncio.to_thread(started.wait, 1.0) + heartbeat_task = asyncio.create_task(_heartbeat_until_done(model_task)) + for _ in range(5): + await asyncio.sleep(0) + assert heartbeat_ticks > 0 + + release.set() + worker_thread_id = await model_task + await heartbeat_task + finally: + host.close() + + assert worker_thread_id != loop_thread_id + + +def test_runtime_host_reentrant_sync_and_async_dispatch_raise() -> None: + host = RuntimeHost(_LifecycleRuntime()) + + def _nested_sync_dispatch() -> None: + host.call(lambda: None) + + def _nested_async_dispatch() -> None: + async def _dispatch() -> None: + await host.call_async(lambda: None) + + asyncio.run(_dispatch()) + + try: + with pytest.raises(RuntimeError, match="own thread"): + host.call(_nested_sync_dispatch) + with pytest.raises(RuntimeError, match="own thread"): + host.call(_nested_async_dispatch) + finally: + host.close() + + +def test_non_control_rank_setup_returns_after_worker_loop_without_demo_edges() -> None: + runtime = _LifecycleRuntime() + host = RuntimeHost( + runtime, + is_control_rank=False, + worker_loop=runtime.run_worker_loop, + ) + constructed: list[str] = [] + + result = _fake_run_setup(host, constructed) + + assert result == "worker-rank" + assert constructed == [] + assert runtime.events == [ + "initialize_distributed", + "preload", + "run_worker_loop", + "runtime.close", + "close_distributed", + ] + + +def test_control_rank_setup_reaches_demo_assembly() -> None: + runtime = _LifecycleRuntime() + host = RuntimeHost(runtime) + constructed: list[str] = [] + + try: + result = _fake_run_setup(host, constructed) + finally: + host.close() + + assert result == "control-rank" + assert constructed == ["run_mode", "provider", "input_source", "output_sink"] + assert runtime.events[:2] == ["initialize_distributed", "preload"] + assert "run_worker_loop" not in runtime.events + + +def _fake_run_setup(host: RuntimeHost, constructed: list[str]) -> str: + host.preload() + if not host.is_control_rank: + host.run_worker_loop() + host.close() + return "worker-rank" + + constructed.extend(["run_mode", "provider", "input_source", "output_sink"]) + return "control-rank" + + +class _LifecycleRuntime: + def __init__(self) -> None: + self.events: list[object] = [] + + def initialize_distributed(self) -> None: + self.events.append("initialize_distributed") + + def preload(self) -> None: + self.events.append("preload") + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + self.events.append(("start_session", inputs)) + return _LifecycleSession(self.events) + + def run_worker_loop(self) -> None: + self.events.append("run_worker_loop") + + def close(self) -> None: + self.events.append("runtime.close") + + def close_distributed(self) -> None: + self.events.append("close_distributed") + + +class _LifecycleSession: + def __init__(self, events: list[object]) -> None: + self._events = events + self._next_step = 0 + + def next_step_request(self) -> StepRequest | None: + request = StepRequest(step_index=self._next_step) + self._next_step += 1 + return request + + def step(self, inputs: InferenceInput) -> StepResult: + self._events.append(("step", inputs)) + return StepResult(step_index=self._next_step, output=None) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self._next_step = 0 + + def close(self) -> None: + self._events.append("session.close") diff --git a/flashdreams/tests/test_demo_runtime_output_sinks.py b/flashdreams/tests/test_demo_runtime_output_sinks.py new file mode 100644 index 000000000..147e86407 --- /dev/null +++ b/flashdreams/tests/test_demo_runtime_output_sinks.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import fields, is_dataclass +from pathlib import Path +from typing import Any + +import pytest +import torch + +from flashdreams.runtime import OutputArtifact, StepResult, TimeWindow +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSink, + Mp4OutputSpec, + NullOutputSink, + NullOutputSpec, + OutputDecision, + SessionInfo, + WebRTCOutputSpec, + build_output_sink, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_mp4_output_sink_writes_artifact_and_close_is_idempotent( + tmp_path: Path, +) -> None: + writer_calls: list[dict[str, Any]] = [] + + def fake_writer( + video: torch.Tensor, + path: Path, + *, + fps: int | float, + layout: str, + install_hint: str, + ) -> Path: + del install_hint + writer_calls.append( + { + "shape": tuple(video.shape), + "path": path, + "fps": fps, + "layout": layout, + } + ) + return path + + sink = Mp4OutputSink( + output_path=tmp_path / "out.mp4", + fps=24, + writer=fake_writer, + move_to_cpu=False, + ) + sink.open(SessionInfo(output_layout="bvtchw", steady_output_frame_count=1)) + sink.begin_generation(0) + + decision = sink.write( + StepResult.from_video_chunk( + step_index=2, + video_chunk=torch.zeros((1, 2, 3, 3, 4, 5)), + layout="bvtchw", + metrics={"model_step_s": 0.25}, + output_window=TimeWindow(start_s=1.0, end_s=2.0), + ) + ) + artifacts = tuple(sink.close()) + second_close = tuple(sink.close()) + + assert decision == OutputDecision() + assert artifacts == second_close + assert artifacts == ( + OutputArtifact( + kind="video/mp4", + uri=str(tmp_path / "out.mp4"), + metadata={ + "fps": 24, + "source_layout": "bvtchw", + "shape": (1, 2, 3, 3, 4, 5), + "stats_history": ( + { + "step_index": 2, + "frames": 3, + "model_step_s": 0.25, + "output_start_s": 1.0, + "output_end_s": 2.0, + }, + ), + }, + ), + ) + assert writer_calls == [ + { + "shape": (3, 4, 10, 3), + "path": tmp_path / "out.mp4", + "fps": 24, + "layout": "thwc", + } + ] + + +def test_output_sink_is_built_from_demo_spec(tmp_path: Path) -> None: + def fake_writer(*args: Any, **kwargs: Any) -> Path: + del args, kwargs + return tmp_path / "demo.mp4" + + spec = DemoSpec( + model_id="fake-demo", + input_mode="replay", + output=Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=12), + ) + + mp4_sink = build_output_sink(spec.output, mp4_writer=fake_writer) + null_sink = build_output_sink(NullOutputSpec(store_results=True)) + + assert isinstance(mp4_sink, Mp4OutputSink) + assert mp4_sink.output_path == tmp_path / "demo.mp4" + assert mp4_sink.fps == 12 + assert mp4_sink.writer is fake_writer + assert isinstance(null_sink, NullOutputSink) + assert null_sink.store_results + with pytest.raises(ValueError, match="realtime transport sink"): + build_output_sink(WebRTCOutputSpec()) + + +def test_sinks_do_not_retain_step_result_references(tmp_path: Path) -> None: + mp4_sink = Mp4OutputSink( + output_path=tmp_path / "out.mp4", + fps=24, + writer=lambda *args: tmp_path / "out.mp4", + move_to_cpu=False, + ) + mp4_sink.open(SessionInfo(output_layout="bvtchw", steady_output_frame_count=1)) + mp4_result = StepResult.from_video_chunk( + step_index=0, + video_chunk=torch.zeros((1, 1, 1, 3, 2, 2)), + layout="bvtchw", + ) + + mp4_sink.write(mp4_result) + + null_sink = NullOutputSink(store_results=True) + null_sink.open(SessionInfo()) + null_result = StepResult( + step_index=1, + output=object(), + frame_count=2, + metrics={"model_step_s": 0.1}, + metadata={"source": "fake"}, + ) + + null_sink.write(null_result) + + assert not _object_graph_contains(mp4_sink, mp4_result) + assert not _object_graph_contains(null_sink, null_result) + assert null_sink.results == [ + { + "step_index": 1, + "frame_count": 2, + "metrics": {"model_step_s": 0.1}, + "metadata": {"source": "fake"}, + } + ] + + +def test_null_output_sink_records_steps_without_artifacts() -> None: + sink = NullOutputSink(store_results=True) + sink.open(SessionInfo(output_layout="fake-video", steady_output_frame_count=1)) + + sink.write(StepResult(step_index=0, output="first", frame_count=1)) + sink.write(StepResult(step_index=1, output="second", frame_count=2)) + artifacts = tuple(sink.close()) + + assert artifacts == () + assert tuple(sink.close()) == () + assert sink.output_count == 2 + assert sink.results == [ + {"step_index": 0, "frame_count": 1, "metrics": {}, "metadata": {}}, + {"step_index": 1, "frame_count": 2, "metrics": {}, "metadata": {}}, + ] + + +def _object_graph_contains( + root: object, + needle: object, + *, + seen: set[int] | None = None, +) -> bool: + if root is needle: + return True + if seen is None: + seen = set() + root_id = id(root) + if root_id in seen: + return False + seen.add(root_id) + if root is None or isinstance(root, str | bytes | int | float | bool | Path): + return False + if isinstance(root, torch.Tensor): + return False + if callable(root): + return False + if isinstance(root, Mapping): + return any( + _object_graph_contains(key, needle, seen=seen) + or _object_graph_contains(value, needle, seen=seen) + for key, value in root.items() + ) + if isinstance(root, list | tuple | set | frozenset): + return any(_object_graph_contains(value, needle, seen=seen) for value in root) + if is_dataclass(root): + return any( + _object_graph_contains(getattr(root, field.name), needle, seen=seen) + for field in fields(root) + ) + return False diff --git a/flashdreams/tests/test_demo_runtime_realtime_driver.py b/flashdreams/tests/test_demo_runtime_realtime_driver.py new file mode 100644 index 000000000..5c79b83f7 --- /dev/null +++ b/flashdreams/tests/test_demo_runtime_realtime_driver.py @@ -0,0 +1,916 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +import time +from collections.abc import Callable, Sequence +from typing import Any, Literal, cast + +import pytest + +from flashdreams.runtime import ( + InferenceInput, + InferenceRuntime, + StepRequest, + StepRequirements, + StepResult, + UserInputSchema, +) +from flashdreams.runtime.demo import ( + ActivationResult, + DriverInvariantError, + ErrorAction, + InMemorySessionMetricsRecorder, + OutputDecision, + PreparedStep, + ProviderCapabilities, + RealtimeSessionDriver, + RealtimeWindowResult, + RunContext, + RunResult, + RuntimeHost, + SessionEdges, + SessionInfo, + SingleSessionAdmissionPolicy, + StepPipeline, + UserInputWindow, + shielded_session_cleanup, +) + +pytestmark = pytest.mark.ci_cpu + + +@pytest.mark.asyncio +async def test_realtime_driver_non_activation_returns_not_activated() -> None: + runtime = _FakeRealtimeRuntime(session=_FakeRealtimeSession(num_steps=1)) + host = RuntimeHost(runtime) + provider = _FakeRealtimeProvider() + output = _RecordingOutputSink() + transport = _RecordingTransport() + metrics = InMemorySessionMetricsRecorder() + edges = _edges( + input_source=_RealtimeInputSource(), + output=output, + transport=transport, + metrics=metrics, + activation=_ActivationPolicy(ActivationResult(activated=False, reason="idle")), + ) + + try: + result = await RealtimeSessionDriver().run_one_session( + host=host, + provider=provider, + session_edges=edges, + pipeline=StepPipeline(), + ) + finally: + host.close() + + assert result.status == "not_activated" + assert result.reason == "idle" + assert runtime.start_session_inputs == [] + assert provider.prepare_initial_count == 0 + assert provider.close_count == 1 + assert output.close_count == 1 + assert transport.close_count == 1 + assert metrics.closed + + +@pytest.mark.asyncio +async def test_realtime_driver_transport_close_before_first_step_is_not_activated() -> ( + None +): + session = _FakeRealtimeSession(num_steps=1) + runtime = _FakeRealtimeRuntime(session=session) + host = RuntimeHost(runtime) + transport = _RecordingTransport() + edges = _edges( + input_source=_RealtimeInputSource(transport_to_close=transport), + transport=transport, + ) + + try: + result = await RealtimeSessionDriver().run_one_session( + host=host, + provider=_FakeRealtimeProvider(), + session_edges=edges, + pipeline=StepPipeline(), + ) + finally: + host.close() + + assert result.status == "not_activated" + assert result.reason == "transport closed before first step" + assert session.step_inputs == [] + + +@pytest.mark.asyncio +async def test_repeated_cancellation_during_cleanup_still_closes_edges() -> None: + entered_window = asyncio.Event() + session = _FakeRealtimeSession(num_steps=1) + runtime = _FakeRealtimeRuntime(session=session) + host = RuntimeHost(runtime) + provider = _FakeRealtimeProvider(close_delay_s=0.05) + output = _RecordingOutputSink() + transport = _RecordingTransport() + metrics = InMemorySessionMetricsRecorder() + edges = _edges( + input_source=_RealtimeInputSource( + entered=entered_window, + wait_forever=True, + ), + output=output, + transport=transport, + metrics=metrics, + ) + task = asyncio.create_task( + RealtimeSessionDriver().run_one_session( + host=host, + provider=provider, + session_edges=edges, + pipeline=StepPipeline(), + ) + ) + await entered_window.wait() + + task.cancel() + await asyncio.sleep(0) + task.cancel() + result = await task + host.close() + + assert result.status == "cancelled" + assert result.reason == "cancelled" + assert session.close_count == 1 + assert provider.close_count == 1 + assert output.close_count == 1 + assert transport.close_count == 1 + assert metrics.closed + assert not edges.cleanup_tasks + + +@pytest.mark.asyncio +async def test_cancelled_realtime_driver_inside_timeout_returns_result() -> None: + timeout_context = getattr(asyncio, "timeout", None) + if timeout_context is None: + pytest.skip("asyncio.timeout is unavailable on this Python version.") + entered_window = asyncio.Event() + runtime = _FakeRealtimeRuntime(session=_FakeRealtimeSession(num_steps=1)) + host = RuntimeHost(runtime) + edges = _edges( + input_source=_RealtimeInputSource( + entered=entered_window, + wait_forever=True, + ) + ) + + async def timeout_after_window_entry(timeout: Any) -> None: + await entered_window.wait() + timeout.reschedule(asyncio.get_running_loop().time()) + + try: + async with timeout_context(None) as timeout: + timeout_task = asyncio.create_task(timeout_after_window_entry(timeout)) + result = await RealtimeSessionDriver().run_one_session( + host=host, + provider=_FakeRealtimeProvider(), + session_edges=edges, + pipeline=StepPipeline(), + ) + timeout_task.cancel() + await asyncio.gather(timeout_task, return_exceptions=True) + finally: + host.close() + + assert entered_window.is_set() + assert result.status == "cancelled" + + +@pytest.mark.asyncio +async def test_realtime_driver_invariant_finalizes_edges_before_reraising() -> None: + runtime = _FakeRealtimeRuntime(session=_FakeRealtimeSession(num_steps=1)) + host = RuntimeHost(runtime) + provider = _FakeRealtimeProvider(fail_initial=RuntimeError("bad setup policy")) + output = _RecordingOutputSink() + transport = _RecordingTransport() + metrics = InMemorySessionMetricsRecorder() + edges = _edges( + output=output, + transport=transport, + metrics=metrics, + error_policy=_SetupPolicy(result_status="completed"), + ) + + try: + with pytest.raises(DriverInvariantError, match="Setup failures") as raised: + await RealtimeSessionDriver().run_one_session( + host=host, + provider=provider, + session_edges=edges, + pipeline=StepPipeline(), + ) + finally: + host.close() + + result = edges.close_result() + assert result.status == "failed" + assert result.error is raised.value + assert provider.close_count == 1 + assert output.close_count == 1 + assert transport.close_count == 1 + assert metrics.closed + + +@pytest.mark.asyncio +async def test_realtime_step_invariant_reraises_without_error_policy() -> None: + runtime = _FakeRealtimeRuntime(session=_FakeRealtimeSession(num_steps=1)) + host = RuntimeHost(runtime) + output = _RecordingOutputSink() + transport = _RecordingTransport() + metrics = InMemorySessionMetricsRecorder() + edges = _edges( + output=output, + transport=transport, + metrics=metrics, + error_policy=_DropOutputErrorPolicy(), + ) + + try: + with pytest.raises(DriverInvariantError, match="step invariant"): + await RealtimeSessionDriver().run_one_session( + host=host, + provider=_FakeRealtimeProvider(), + session_edges=edges, + pipeline=_InvariantPipeline(), + ) + finally: + host.close() + + assert metrics.errors == [] + assert output.close_count == 1 + assert transport.close_count == 1 + assert metrics.closed + + +@pytest.mark.asyncio +async def test_run_context_close_async_drains_registered_cleanup_task() -> None: + runtime = _FakeRealtimeRuntime(session=_FakeRealtimeSession(num_steps=1)) + host = RuntimeHost(runtime) + context = RunContext( + host=host, + run_metrics=InMemorySessionMetricsRecorder(), + admission=SingleSessionAdmissionPolicy(), + ) + edges = _edges(cleanup_tasks=context.cleanup_tasks) + cleanup_task = asyncio.create_task( + shielded_session_cleanup( + host=host, + session=runtime.session, + provider=_FakeRealtimeProvider(close_delay_s=0.02), + session_edges=edges, + status="cancelled", + reason="test", + error=None, + ) + ) + await asyncio.sleep(0) + + summary = await context.close_async() + cleanup_result = await cleanup_task + + assert cleanup_task.done() + assert cleanup_result.status == "cancelled" + assert not context.cleanup_tasks + assert edges.is_closed + assert summary.metrics.counters["sessions"] == 0 + + +@pytest.mark.asyncio +async def test_shielded_cleanup_never_raises_and_returns_result_on_close_errors() -> ( + None +): + runtime = _FakeRealtimeRuntime( + session=_FakeRealtimeSession(num_steps=1, fail_close=RuntimeError("session")) + ) + host = RuntimeHost(runtime) + provider = _FakeRealtimeProvider(fail_close=RuntimeError("provider")) + metrics = InMemorySessionMetricsRecorder() + edges = _edges(metrics=metrics) + + try: + result = await shielded_session_cleanup( + host=host, + session=runtime.session, + provider=provider, + session_edges=edges, + status="failed", + reason="test failure", + error=RuntimeError("original"), + ) + assert not host.is_healthy + assert host.unhealthy_reason == "model-affine cleanup failed" + finally: + host.close() + + assert result.status == "failed" + assert result.reason == "test failure" + assert metrics.closed + assert metrics.cleanup_errors == ["session", "provider"] + + +@pytest.mark.asyncio +async def test_shielded_cleanup_timeout_bounds_shutdown() -> None: + host = _NeverReturningHost() + session = _FakeRealtimeSession(num_steps=1) + provider = _FakeRealtimeProvider() + metrics = InMemorySessionMetricsRecorder() + edges = _edges(metrics=metrics) + + result = await shielded_session_cleanup( + host=cast(RuntimeHost, host), + session=session, + provider=provider, + session_edges=edges, + status="cancelled", + reason="timeout test", + error=None, + timeout_s=0.001, + ) + + assert result.status == "cancelled" + assert host.unhealthy_reason == "model-affine cleanup timed out" + assert host.close_targets == [session, provider] + assert provider.close_count == 0 + assert metrics.cleanup_errors == [] + assert len(metrics.orphaned_cleanup_errors) == 1 + assert metrics.closed + + +@pytest.mark.asyncio +async def test_shielded_cleanup_dispatch_failure_marks_host_unhealthy() -> None: + host = _RejectingHost(RuntimeError("worker rejected cleanup")) + session = _FakeRealtimeSession(num_steps=1) + provider = _FakeRealtimeProvider() + metrics = InMemorySessionMetricsRecorder() + edges = _edges(metrics=metrics) + + result = await shielded_session_cleanup( + host=cast(RuntimeHost, host), + session=session, + provider=provider, + session_edges=edges, + status="cancelled", + reason="dispatch failure", + error=None, + timeout_s=0.001, + ) + + assert result.status == "cancelled" + assert host.unhealthy_reason == "model-affine cleanup failed" + assert host.cleanup_dispatch_count == 1 + assert session.close_count == 0 + assert provider.close_count == 0 + assert metrics.cleanup_errors == ["worker rejected cleanup"] + assert metrics.orphaned_cleanup_errors == [] + assert metrics.closed + + +@pytest.mark.asyncio +async def test_realtime_driver_applies_backpressure_through_clock() -> None: + session = _FakeRealtimeSession(num_steps=2) + runtime = _FakeRealtimeRuntime(session=session) + host = RuntimeHost(runtime) + clock = _RecordingRealtimeClock() + metrics = InMemorySessionMetricsRecorder() + output = _RecordingOutputSink( + decisions=( + OutputDecision(backpressure_s=0.25), + OutputDecision(should_stop=True), + ) + ) + edges = _edges(clock=clock, output=output, metrics=metrics) + + try: + result = await RealtimeSessionDriver().run_one_session( + host=host, + provider=_FakeRealtimeProvider(), + session_edges=edges, + pipeline=StepPipeline(), + ) + finally: + host.close() + + assert result.status == "completed" + assert clock.backpressure == [0.25] + assert metrics.catch_up_count == 2 + assert len(output.results) == 2 + + +@pytest.mark.asyncio +async def test_realtime_driver_calls_step_pipeline_on_runtime_host() -> None: + session = _FakeRealtimeSession(num_steps=1) + runtime = _FakeRealtimeRuntime(session=session) + host = _RecordingRuntimeHost(runtime) + edges = _edges( + output=_RecordingOutputSink(decisions=(OutputDecision(should_stop=True),)) + ) + + try: + result = await RealtimeSessionDriver().run_one_session( + host=host, + provider=_FakeRealtimeProvider(), + session_edges=edges, + pipeline=StepPipeline(), + ) + finally: + host.close() + + assert result.status == "completed" + assert "execute_step" in host.async_calls + assert "prepare_step" not in host.async_calls + assert "step" not in host.async_calls + + +@pytest.mark.asyncio +async def test_slow_fake_model_step_does_not_block_event_loop() -> None: + session = _FakeRealtimeSession(num_steps=1, step_delay_s=0.05) + runtime = _FakeRealtimeRuntime(session=session) + host = RuntimeHost(runtime) + edges = _edges( + output=_RecordingOutputSink(decisions=(OutputDecision(should_stop=True),)) + ) + ticks = 0 + finished = False + + async def heartbeat() -> None: + nonlocal ticks + while not finished: + ticks += 1 + await asyncio.sleep(0.005) + + heartbeat_task = asyncio.create_task(heartbeat()) + try: + result = await RealtimeSessionDriver().run_one_session( + host=host, + provider=_FakeRealtimeProvider(), + session_edges=edges, + pipeline=StepPipeline(), + ) + finally: + finished = True + await heartbeat_task + host.close() + + assert result.status == "completed" + assert ticks >= 2 + + +@pytest.mark.asyncio +async def test_realtime_driver_fatal_model_error_returns_failed() -> None: + session = _FakeRealtimeSession(num_steps=1, fail_step=0) + runtime = _FakeRealtimeRuntime(session=session) + host = RuntimeHost(runtime) + metrics = InMemorySessionMetricsRecorder() + edges = _edges(metrics=metrics) + + try: + result = await RealtimeSessionDriver().run_one_session( + host=host, + provider=_FakeRealtimeProvider(), + session_edges=edges, + pipeline=StepPipeline(), + ) + finally: + host.close() + + assert result.status == "failed" + assert result.reason == "step failed" + assert isinstance(result.error, RuntimeError) + assert session.close_count == 1 + assert metrics.errors == ["step failed"] + + +@pytest.mark.asyncio +async def test_realtime_driver_can_drop_recoverable_output_error() -> None: + session = _FakeRealtimeSession(num_steps=2) + runtime = _FakeRealtimeRuntime(session=session) + host = RuntimeHost(runtime) + output = _RecordingOutputSink( + fail_first_write=RuntimeError("output queue full"), + decisions=(OutputDecision(should_stop=True),), + ) + metrics = InMemorySessionMetricsRecorder() + edges = _edges( + output=output, + metrics=metrics, + error_policy=_DropOutputErrorPolicy(), + ) + + try: + result = await RealtimeSessionDriver().run_one_session( + host=host, + provider=_FakeRealtimeProvider(), + session_edges=edges, + pipeline=StepPipeline(), + ) + finally: + host.close() + + assert result.status == "completed" + assert metrics.errors == ["output queue full"] + assert [step.step_index for step in output.results] == [1] + assert len(session.step_inputs) == 2 + + +def _edges( + *, + input_source: "_RealtimeInputSource | None" = None, + output: "_RecordingOutputSink | None" = None, + transport: "_RecordingTransport | None" = None, + metrics: InMemorySessionMetricsRecorder | None = None, + activation: "_ActivationPolicy | None" = None, + clock: "_RecordingRealtimeClock | None" = None, + cleanup_tasks: set[asyncio.Task[RunResult]] | None = None, + error_policy: Any | None = None, +) -> SessionEdges: + return SessionEdges( + input_source=input_source or _RealtimeInputSource(), + output_sink=output + or _RecordingOutputSink(decisions=(OutputDecision(should_stop=True),)), + cleanup_tasks=cleanup_tasks or set(), + metrics=metrics or InMemorySessionMetricsRecorder(), + error_policy=error_policy or _DefaultTestErrorPolicy(), + transport=transport or _RecordingTransport(), + clock=clock or _RecordingRealtimeClock(), + activation=activation or _ActivationPolicy(ActivationResult(activated=True)), + ) + + +def _window(index: int) -> UserInputWindow: + start_s = float(index) + return UserInputWindow( + start_s=start_s, + end_s=start_s + 1.0, + frame_times=(start_s + 1.0,), + ) + + +class _ActivationPolicy: + timeout_s: float | None = None + + def __init__(self, result: ActivationResult) -> None: + self.result = result + self.calls = 0 + + async def wait_until_active(self, clock: Any) -> ActivationResult: + del clock + self.calls += 1 + await asyncio.sleep(0) + return self.result + + +class _RecordingRealtimeClock: + is_realtime = True + is_deterministic = False + + def __init__(self) -> None: + self.backpressure: list[float] = [] + self.anchors: list[float] = [] + + def now(self) -> float: + return 0.0 + + def anchor(self, wall_time_s: float) -> None: + self.anchors.append(wall_time_s) + + async def wait_until_window_end(self, end_s: float) -> None: + del end_s + + async def apply_backpressure(self, requested_s: float) -> None: + self.backpressure.append(requested_s) + await asyncio.sleep(0) + + def catch_up(self, **kwargs: Any) -> object: + del kwargs + return object() + + +class _RealtimeInputSource: + is_finite = False + is_deterministic = False + user_input_schema = UserInputSchema() + + def __init__( + self, + *, + entered: asyncio.Event | None = None, + wait_forever: bool = False, + transport_to_close: "_RecordingTransport | None" = None, + ) -> None: + self.entered = entered + self.wait_forever = wait_forever + self.transport_to_close = transport_to_close + self.requests: list[StepRequirements] = [] + + def is_finished(self) -> bool: + return False + + async def next_realtime_window( + self, + *, + request: StepRequirements, + clock: Any, + ) -> RealtimeWindowResult: + del clock + self.requests.append(request) + if self.entered is not None: + self.entered.set() + if self.wait_forever: + await asyncio.Event().wait() + if self.transport_to_close is not None: + self.transport_to_close.close() + return RealtimeWindowResult(window=_window(request.step_index)) + + +class _FakeRealtimeProvider: + capabilities = ProviderCapabilities( + supports_realtime_clock=True, + supports_reset=True, + deterministic_given_inputs=False, + ) + + def __init__( + self, + *, + fail_initial: Exception | None = None, + fail_close: Exception | None = None, + close_delay_s: float = 0.0, + ) -> None: + self.fail_initial = fail_initial + self.fail_close = fail_close + self.close_delay_s = close_delay_s + self.prepare_initial_count = 0 + self.close_count = 0 + self.reset_inputs: list[InferenceInput | None] = [] + + def prepare_initial_input(self) -> InferenceInput: + if self.fail_initial is not None: + raise self.fail_initial + self.prepare_initial_count += 1 + return InferenceInput(global_conditioning={"prompt": "realtime"}) + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + return PreparedStep( + inference_input=InferenceInput( + step={ + "request_step": request.step_index, + "window": (user_window.start_s, user_window.end_s), + } + ) + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + self.reset_inputs.append(inputs) + + def close(self) -> None: + if self.close_delay_s: + time.sleep(self.close_delay_s) + self.close_count += 1 + if self.fail_close is not None: + raise self.fail_close + + +class _FakeRealtimeRuntime: + def __init__(self, *, session: "_FakeRealtimeSession") -> None: + self.session = session + self.start_session_inputs: list[InferenceInput] = [] + self.close_count = 0 + + def start_session(self, inputs: InferenceInput) -> "_FakeRealtimeSession": + self.start_session_inputs.append(inputs) + return self.session + + def close(self) -> None: + self.close_count += 1 + + +class _FakeRealtimeSession: + def __init__( + self, + *, + num_steps: int, + fail_step: int | None = None, + fail_close: Exception | None = None, + step_delay_s: float = 0.0, + ) -> None: + self.num_steps = num_steps + self.fail_step = fail_step + self.fail_close = fail_close + self.step_delay_s = step_delay_s + self.next_request_index = 0 + self.step_inputs: list[InferenceInput] = [] + self.close_count = 0 + + def session_info(self) -> SessionInfo: + return SessionInfo(output_layout="fake-realtime", steady_output_frame_count=1) + + def next_step_requirements(self) -> StepRequirements | None: + if self.next_request_index >= self.num_steps: + return None + request = StepRequirements(step_index=self.next_request_index) + self.next_request_index += 1 + return request + + def next_step_request(self) -> StepRequest | None: + raise AssertionError("demo driver should request StepRequirements") + + def step(self, inputs: InferenceInput) -> StepResult: + step_index = len(self.step_inputs) + if self.step_delay_s: + time.sleep(self.step_delay_s) + if self.fail_step == step_index: + raise RuntimeError("step failed") + self.step_inputs.append(inputs) + return StepResult( + step_index=step_index, + output=f"frame-{step_index}", + frame_count=1, + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self.next_request_index = 0 + self.step_inputs.clear() + + def close(self) -> None: + self.close_count += 1 + if self.fail_close is not None: + raise self.fail_close + + +class _RecordingRuntimeHost(RuntimeHost): + def __init__(self, runtime: InferenceRuntime) -> None: + super().__init__(runtime) + self.async_calls: list[str] = [] + + async def call_async( + self, + func: Callable[..., Any], + /, + *args: object, + **kwargs: object, + ) -> Any: + self.async_calls.append(getattr(func, "__name__", type(func).__name__)) + return await super().call_async(func, *args, **kwargs) + + +class _InvariantPipeline(StepPipeline): + def execute_step(self, **kwargs: object) -> Any: + del kwargs + raise DriverInvariantError("step invariant") + + +class _RecordingOutputSink: + produces_artifacts = False + + def __init__( + self, + *, + decisions: Sequence[OutputDecision] = (), + fail_first_write: Exception | None = None, + ) -> None: + self.decisions = list(decisions) + self.fail_first_write = fail_first_write + self.opened_with: SessionInfo | None = None + self.generations: list[int] = [] + self.results: list[StepResult] = [] + self.close_count = 0 + self.write_attempts = 0 + + def open(self, session_info: SessionInfo) -> None: + self.opened_with = session_info + + def begin_generation(self, generation: int) -> None: + self.generations.append(generation) + + def write(self, result: StepResult) -> OutputDecision: + self.write_attempts += 1 + if self.fail_first_write is not None: + exc = self.fail_first_write + self.fail_first_write = None + raise exc + self.results.append(result) + if self.decisions: + return self.decisions.pop(0) + return OutputDecision() + + def close(self) -> Sequence[Any]: + self.close_count += 1 + return () + + +class _RecordingTransport: + def __init__(self) -> None: + self.active = True + self.close_count = 0 + + def is_active(self) -> bool: + return self.active + + def close(self) -> None: + self.active = False + self.close_count += 1 + + +class _DefaultTestErrorPolicy: + def handle_setup_error(self, exc: Exception) -> ErrorAction: + del exc + return ErrorAction(result_status="failed") + + def handle(self, exc: Exception) -> ErrorAction: + del exc + return ErrorAction(result_status="failed") + + +class _SetupPolicy(_DefaultTestErrorPolicy): + def __init__( + self, + *, + result_status: Literal["completed", "failed", "skipped"], + ) -> None: + self.result_status = result_status + + def handle_setup_error(self, exc: Exception) -> ErrorAction: + del exc + return ErrorAction(result_status=self.result_status) + + +class _DropOutputErrorPolicy(_DefaultTestErrorPolicy): + def handle(self, exc: Exception) -> ErrorAction: + del exc + return ErrorAction( + close_session=False, + drop_chunk=True, + result_status="failed", + ) + + +class _NeverReturningHost: + def __init__(self) -> None: + self.unhealthy_reason: str | None = None + self.close_targets: list[Any] = [] + + async def call_async( + self, + func: Callable[..., Any], + /, + *args: object, + **kwargs: object, + ) -> Any: + del func, kwargs + for arg in args: + if callable(arg): + close = cast(Callable[[], None], arg) + self.close_targets.append(getattr(close, "__self__", close)) + await asyncio.Event().wait() + + def mark_unhealthy( + self, + reason: str = "marked unhealthy", + error: Exception | None = None, + ) -> None: + del error + self.unhealthy_reason = reason + + +class _RejectingHost: + def __init__(self, exc: Exception) -> None: + self.exc = exc + self.cleanup_dispatch_count = 0 + self.unhealthy_reason: str | None = None + + async def call_async( + self, + func: Callable[..., Any], + /, + *args: object, + **kwargs: object, + ) -> Any: + del func, args, kwargs + self.cleanup_dispatch_count += 1 + raise self.exc + + def mark_unhealthy( + self, + reason: str = "marked unhealthy", + error: Exception | None = None, + ) -> None: + del error + self.unhealthy_reason = reason diff --git a/flashdreams/tests/test_demo_runtime_run_modes.py b/flashdreams/tests/test_demo_runtime_run_modes.py new file mode 100644 index 000000000..8742eb0e9 --- /dev/null +++ b/flashdreams/tests/test_demo_runtime_run_modes.py @@ -0,0 +1,774 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +import contextlib +import threading +from collections.abc import Callable, Coroutine, Mapping, Sequence +from pathlib import Path +from typing import Any, cast + +import pytest + +from flashdreams.runtime import ( + CanonicalInputSchema, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InferenceRuntime, + InferenceSession, + InputMapping, + StepRequirements, + UserInputSchema, +) +from flashdreams.runtime.demo import ( + AsyncSessionDriver, + BenchmarkErrorPolicy, + DemoSpec, + DriverInvariantError, + InMemorySessionMetricsRecorder, + ModelWarmupPlan, + Mp4ErrorPolicy, + Mp4OutputSpec, + NativeWindowErrorPolicy, + NullErrorPolicy, + NullOutputSpec, + OutputDecision, + PreparedScenario, + ProviderCapabilities, + RunContext, + RunModeCapabilities, + RunResult, + RuntimeHost, + SessionDriver, + SessionEdges, + SessionInfo, + StepPipeline, + UserInputWindow, + WebRTCErrorPolicy, + WebRTCOutputSpec, + run_demo_session, + run_demo_session_async, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_fake_mp4_run_mode_calls_session_helper_once(tmp_path: Path) -> None: + spec = DemoSpec( + model_id="fake-demo", + input_mode="replay", + output=Mp4OutputSpec(path=tmp_path / "fake.mp4", fps=12), + ) + adapter = _FakeAdapter() + mode = _FakeRunMode(name="mp4", driver=_ClosingSyncDriver()) + runtime = _UnusedRuntime() + helper_calls: list[DemoSpec] = [] + + result = _run_fake_single_session_mode( + spec=spec, + adapter=adapter, + mode=mode, + runtime=runtime, + helper=lambda **kwargs: _record_sync_helper(helper_calls, **kwargs), + ) + + assert result.status == "completed" + assert helper_calls == [spec] + assert len(mode.created_edges) == 1 + assert mode.created_edges[0].is_closed + assert mode.created_edges[0].cleanup_tasks is mode.require_context().cleanup_tasks + assert adapter.providers[0].close_count == 1 + assert mode.admission.reservations[0].release_count == 1 + + +def test_fake_benchmark_run_mode_calls_helper_once_per_scenario() -> None: + specs = [ + DemoSpec( + model_id="fake-demo", + input_mode="replay", + output=NullOutputSpec(), + scenario=f"scenario-{index}", + ) + for index in range(2) + ] + adapter = _FakeAdapter() + mode = _FakeRunMode(name="benchmark", driver=_ClosingSyncDriver()) + helper_calls: list[DemoSpec] = [] + + results = _run_fake_benchmark_mode( + specs=specs, + adapter=adapter, + mode=mode, + runtime=_UnusedRuntime(), + helper=lambda **kwargs: _record_sync_helper(helper_calls, **kwargs), + ) + + assert [result.status for result in results] == ["completed", "completed"] + assert helper_calls == specs + assert adapter.prepare_scenario_calls == specs + assert len(adapter.providers) == 2 + assert len({id(provider) for provider in adapter.providers}) == 2 + assert all(provider.close_count == 1 for provider in adapter.providers) + assert len(mode.created_edges) == 2 + assert len({id(edges) for edges in mode.created_edges}) == 2 + assert all(edges.is_closed for edges in mode.created_edges) + assert all( + reservation.release_count == 1 for reservation in mode.admission.reservations + ) + + +@pytest.mark.asyncio +async def test_fake_webrtc_offer_reserves_before_prepare_or_negotiation() -> None: + spec = DemoSpec( + model_id="fake-demo", + input_mode="keyboard-driving", + output=WebRTCOutputSpec(port=8081), + ) + events: list[str] = [] + blocking_io = _BlockingIOService(events) + webrtc = _FakeWebRTCService(events) + adapter = _FakeAdapter(events=events) + mode = _FakeRunMode( + name="webrtc", + driver=_ClosingAsyncDriver(), + admission=_RecordingAdmission(events=events), + services={"blocking_io": blocking_io, "webrtc": webrtc}, + ) + context = mode.create_run_context( + spec=spec, + adapter=adapter, + host=RuntimeHost(_UnusedRuntime()), + model_warmup_plan=ModelWarmupPlan(), + ) + helper_calls: list[DemoSpec] = [] + + answer = await _handle_fake_webrtc_offer( + context=context, + spec=spec, + adapter=adapter, + mode=mode, + helper=lambda **kwargs: _record_async_helper(helper_calls, **kwargs), + events=events, + ) + + assert answer == "answer" + assert events.index("admission.reserve") < events.index("blocking_io.run") + assert events.index("blocking_io.run") < events.index("webrtc.answer") + assert blocking_io.run_count == 1 + assert helper_calls == [spec] + assert mode.admission.reservations[0].release_count == 1 + assert adapter.providers[0].close_count == 1 + assert mode.created_edges[0].is_closed + + +@pytest.mark.asyncio +async def test_async_session_cancellation_shields_pre_edge_provider_cleanup() -> None: + spec = DemoSpec( + model_id="fake-demo", + input_mode="keyboard-driving", + output=WebRTCOutputSpec(port=8081), + ) + provider = _BlockingCloseProvider() + adapter = _BlockingCloseAdapter(provider=provider) + mode = _CancelBeforeEdgesRunMode(name="webrtc", driver=_ClosingAsyncDriver()) + context = mode.create_run_context( + spec=spec, + adapter=adapter, + host=RuntimeHost(_UnusedRuntime()), + model_warmup_plan=ModelWarmupPlan(), + ) + scenario = adapter.prepare_scenario(spec) + task = asyncio.create_task( + run_demo_session_async( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=mode, + pipeline=StepPipeline(), + ) + ) + + close_started = await asyncio.to_thread(provider.close_started.wait, 1.0) + assert close_started + task.cancel() + provider.release_close.set() + try: + result = await asyncio.wait_for(task, timeout=1.0) + finally: + provider.release_close.set() + context.host.close() + + assert result.status == "cancelled" + assert result.reason == "cancelled during session assembly" + assert provider.close_count == 1 + assert mode.created_edges == [] + assert mode.admission.reservations[0].release_count == 1 + run_metrics = cast(InMemorySessionMetricsRecorder, context.run_metrics) + assert run_metrics.sessions == [result] + + +def test_run_demo_session_rejects_reused_closed_session_edges() -> None: + spec = DemoSpec( + model_id="fake-demo", + input_mode="replay", + output=NullOutputSpec(), + ) + adapter = _FakeAdapter() + mode = _ReusingRunMode(name="mp4", driver=_ClosingSyncDriver()) + context = mode.create_run_context( + spec=spec, + adapter=adapter, + host=RuntimeHost(_UnusedRuntime()), + model_warmup_plan=ModelWarmupPlan(), + ) + scenario = adapter.prepare_scenario(spec) + + first = run_demo_session( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=mode, + pipeline=StepPipeline(), + ) + with pytest.raises(DriverInvariantError, match="must not be reused"): + run_demo_session( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=mode, + pipeline=StepPipeline(), + ) + + assert first.status == "completed" + run_metrics = cast(InMemorySessionMetricsRecorder, context.run_metrics) + assert len(run_metrics.sessions) == 1 + assert run_metrics.sessions[0] is first + assert adapter.providers[1].close_count == 1 + + +@pytest.mark.asyncio +async def test_run_context_close_async_drains_cleanup_tasks() -> None: + metrics = InMemorySessionMetricsRecorder() + context = RunContext( + host=RuntimeHost(_UnusedRuntime()), + run_metrics=metrics, + admission=_RecordingAdmission(events=[]), + ) + task = asyncio.create_task(_finished_cleanup_result()) + context.cleanup_tasks.add(task) + + with pytest.raises(RuntimeError, match="Pending session cleanup tasks"): + context.close() + + summary = await context.close_async() + + assert not context.cleanup_tasks + assert task.done() + assert summary.metrics.counters["sessions"] == 0 + assert metrics.closed + + +def test_error_policy_implementations_keep_setup_failures_terminal() -> None: + exc = RuntimeError("setup failed") + policies = ( + Mp4ErrorPolicy(), + BenchmarkErrorPolicy(), + WebRTCErrorPolicy(recoverable_exception_types=(RuntimeError,)), + NativeWindowErrorPolicy(), + NullErrorPolicy(), + ) + + for policy in policies: + action = policy.handle_setup_error(exc) + assert action.result_status == "failed" + assert action.close_session + assert not action.drop_chunk + + +def test_benchmark_error_policy_marks_failed_scenario_continuable() -> None: + action = BenchmarkErrorPolicy().handle(RuntimeError("scenario failed")) + + assert action.result_status == "failed" + assert action.close_session + assert action.continue_next_scenario + assert not action.drop_chunk + + +def test_webrtc_error_policy_can_drop_recoverable_step_errors() -> None: + action = WebRTCErrorPolicy( + recoverable_exception_types=(RuntimeError,), + ).handle(RuntimeError("output queue full")) + + assert action.result_status == "failed" + assert not action.close_session + assert action.drop_chunk + assert not action.continue_next_scenario + + +def _run_fake_single_session_mode( + *, + spec: DemoSpec, + adapter: "_FakeAdapter", + mode: "_FakeRunMode", + runtime: "_UnusedRuntime", + helper: Callable[..., RunResult], +) -> RunResult: + mode.validate_run(spec=spec, adapter=adapter) + scenario = adapter.prepare_scenario(spec) + context = mode.create_run_context( + spec=spec, + adapter=adapter, + host=RuntimeHost(runtime), + model_warmup_plan=ModelWarmupPlan(), + ) + mode.warmup_context( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + ) + return helper( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=mode, + pipeline=StepPipeline(), + ) + + +def _run_fake_benchmark_mode( + *, + specs: Sequence[DemoSpec], + adapter: "_FakeAdapter", + mode: "_FakeRunMode", + runtime: "_UnusedRuntime", + helper: Callable[..., RunResult], +) -> list[RunResult]: + mode.validate_run(spec=specs[0], adapter=adapter) + context = mode.create_run_context( + spec=specs[0], + adapter=adapter, + host=RuntimeHost(runtime), + model_warmup_plan=ModelWarmupPlan(), + ) + results: list[RunResult] = [] + for spec in specs: + scenario = adapter.prepare_scenario(spec) + results.append( + helper( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=mode, + pipeline=StepPipeline(), + ) + ) + return results + + +async def _handle_fake_webrtc_offer( + *, + context: RunContext, + spec: DemoSpec, + adapter: "_FakeAdapter", + mode: "_FakeRunMode", + helper: Callable[..., Coroutine[Any, Any, RunResult]], + events: list[str], +) -> str: + events.append("handler.start") + reservation = context.admission.try_reserve() + if reservation is None: + return "busy" + + task: asyncio.Task[RunResult] | None = None + try: + blocking_io = cast(_BlockingIOService, context.services["blocking_io"]) + scenario = await blocking_io.run(adapter.prepare_scenario, spec) + task = asyncio.create_task( + helper( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=mode, + pipeline=StepPipeline(), + reservation=reservation, + ) + ) + webrtc = cast(_FakeWebRTCService, context.services["webrtc"]) + return await webrtc.answer(task) + except Exception: + if task is not None: + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + reservation.release() + raise + + +def _record_sync_helper(calls: list[DemoSpec], **kwargs: Any) -> RunResult: + calls.append(kwargs["spec"]) + return run_demo_session(**kwargs) + + +async def _record_async_helper(calls: list[DemoSpec], **kwargs: Any) -> RunResult: + calls.append(kwargs["spec"]) + return await run_demo_session_async(**kwargs) + + +async def _finished_cleanup_result() -> RunResult: + await asyncio.sleep(0) + return RunResult.rejected(reason="test cleanup") + + +class _FakeAdapter: + model_id = "fake-demo" + inference_input_schema = InferenceInputSchema() + canonical_input_schema = CanonicalInputSchema() + + def __init__(self, *, events: list[str] | None = None) -> None: + self.events = events + self.prepare_scenario_calls: list[DemoSpec] = [] + self.providers: list[_FakeProvider] = [] + + def supported_input_modes(self) -> tuple[str, ...]: + return ("replay", "keyboard-driving") + + def supported_output_modes(self) -> tuple[str, ...]: + return ("null", "mp4", "webrtc") + + def default_input_mapping(self) -> InputMapping: + return IdentityInputMapping() + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.model_id: + raise ValueError(f"Unsupported model_id={config.model_id!r}.") + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + return _UnusedRuntime() + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + if self.events is not None: + self.events.append("adapter.prepare_scenario") + self.prepare_scenario_calls.append(spec) + return PreparedScenario(initial_inputs=InferenceInput()) + + def create_model_input_provider( + self, + spec: DemoSpec, + scenario: PreparedScenario, + ) -> "_FakeProvider": + del spec, scenario + provider = _FakeProvider() + self.providers.append(provider) + return provider + + +class _BlockingCloseAdapter(_FakeAdapter): + def __init__(self, *, provider: "_BlockingCloseProvider") -> None: + super().__init__() + self.provider = provider + + def create_model_input_provider( + self, + spec: DemoSpec, + scenario: PreparedScenario, + ) -> "_BlockingCloseProvider": + del spec, scenario + self.providers.append(self.provider) + return self.provider + + +class _FakeProvider: + capabilities = ProviderCapabilities( + supports_recorded_input=True, + supports_reset=True, + deterministic_given_inputs=True, + ) + + def __init__(self) -> None: + self.close_count = 0 + + def close(self) -> None: + self.close_count += 1 + + +class _BlockingCloseProvider(_FakeProvider): + def __init__(self) -> None: + super().__init__() + self.close_started = threading.Event() + self.release_close = threading.Event() + + def close(self) -> None: + self.close_started.set() + if not self.release_close.wait(timeout=1.0): + raise RuntimeError("timed out waiting to release provider close") + super().close() + + +class _FakeRunMode: + def __init__( + self, + *, + name: str, + driver: SessionDriver | AsyncSessionDriver, + admission: "_RecordingAdmission | None" = None, + services: Mapping[str, object] | None = None, + ) -> None: + self.name = name + self.driver = driver + self.created_edges: list[SessionEdges] = [] + self.validate_run_count = 0 + self.warmup_count = 0 + self.capabilities = RunModeCapabilities(requires_finite_input=True) + self.admission = admission or _RecordingAdmission(events=[]) + self.services = services or {} + self.context: RunContext | None = None + + def require_context(self) -> RunContext: + if self.context is None: + raise AssertionError("Run context was not created.") + return self.context + + def validate_run(self, *, spec: DemoSpec, adapter: Any) -> None: + del spec, adapter + self.validate_run_count += 1 + + def validate_session( + self, + *, + spec: DemoSpec, + scenario: Any, + adapter: Any, + provider: Any, + ) -> None: + del spec, scenario, adapter, provider + + def create_run_context( + self, + *, + spec: DemoSpec, + adapter: Any, + host: RuntimeHost, + model_warmup_plan: ModelWarmupPlan, + ) -> RunContext: + del spec, adapter + self.context = RunContext( + host=host, + run_metrics=InMemorySessionMetricsRecorder(), + admission=self.admission, + model_warmup_plan=model_warmup_plan, + services=self.services, + ) + return self.context + + def create_session_edges( + self, + *, + context: RunContext, + spec: DemoSpec, + scenario: Any, + provider: Any, + adapter: Any, + ) -> SessionEdges: + del spec, scenario, provider, adapter + edges = SessionEdges( + input_source=_FinishedInputSource(), + output_sink=_RecordingOutputSink(), + cleanup_tasks=context.cleanup_tasks, + metrics=InMemorySessionMetricsRecorder(), + transport=_RecordingTransport(), + ) + self.created_edges.append(edges) + return edges + + def select_driver(self) -> SessionDriver | AsyncSessionDriver: + return self.driver + + def warmup_context( + self, + *, + context: RunContext, + spec: DemoSpec, + scenario: Any, + adapter: Any, + ) -> None: + del context, spec, scenario, adapter + self.warmup_count += 1 + + +class _ReusingRunMode(_FakeRunMode): + def __init__( + self, + *, + name: str, + driver: SessionDriver | AsyncSessionDriver, + ) -> None: + super().__init__(name=name, driver=driver) + self._edges: SessionEdges | None = None + + def create_session_edges( + self, + *, + context: RunContext, + spec: DemoSpec, + scenario: Any, + provider: Any, + adapter: Any, + ) -> SessionEdges: + if self._edges is None: + self._edges = super().create_session_edges( + context=context, + spec=spec, + scenario=scenario, + provider=provider, + adapter=adapter, + ) + return self._edges + + +class _CancelBeforeEdgesRunMode(_FakeRunMode): + def validate_session( + self, + *, + spec: DemoSpec, + scenario: Any, + adapter: Any, + provider: Any, + ) -> None: + del spec, scenario, adapter, provider + raise asyncio.CancelledError + + +class _ClosingSyncDriver: + def run_one_session( + self, + *, + host: RuntimeHost, + provider: Any, + session_edges: SessionEdges, + pipeline: StepPipeline, + ) -> RunResult: + del host, pipeline + provider.close() + return session_edges.close_result(status="completed") + + +class _ClosingAsyncDriver: + async def run_one_session( + self, + *, + host: RuntimeHost, + provider: Any, + session_edges: SessionEdges, + pipeline: StepPipeline, + ) -> RunResult: + del pipeline + await host.call_async(provider.close) + return session_edges.close_result(status="completed") + + +class _RecordingAdmission: + def __init__(self, *, events: list[str]) -> None: + self.events = events + self.reservations: list[_RecordingReservation] = [] + + def try_reserve(self) -> "_RecordingReservation": + self.events.append("admission.reserve") + reservation = _RecordingReservation() + self.reservations.append(reservation) + return reservation + + +class _RecordingReservation: + def __init__(self) -> None: + self.release_count = 0 + + def release(self) -> None: + if self.release_count: + return + self.release_count += 1 + + +class _FinishedInputSource: + is_finite = True + is_deterministic = True + user_input_schema = UserInputSchema() + + def is_finished(self) -> bool: + return True + + def next_window(self, request: StepRequirements) -> UserInputWindow: + del request + return UserInputWindow(start_s=0.0, end_s=0.0) + + +class _RecordingOutputSink: + produces_artifacts = False + + def __init__(self) -> None: + self.close_count = 0 + + def open(self, session_info: SessionInfo) -> None: + del session_info + + def begin_generation(self, generation: int) -> None: + del generation + + def write(self, result: Any) -> OutputDecision: + del result + return OutputDecision() + + def close(self) -> Sequence[Any]: + self.close_count += 1 + return () + + +class _RecordingTransport: + def close(self) -> None: + return + + def is_active(self) -> bool: + return True + + +class _BlockingIOService: + def __init__(self, events: list[str]) -> None: + self.events = events + self.run_count = 0 + + async def run(self, func: Callable[..., Any], *args: object) -> Any: + self.run_count += 1 + self.events.append("blocking_io.run") + await asyncio.sleep(0) + return func(*args) + + +class _FakeWebRTCService: + def __init__(self, events: list[str]) -> None: + self.events = events + + async def answer(self, task: asyncio.Task[RunResult]) -> str: + self.events.append("webrtc.answer") + result = await task + assert result.status == "completed" + return "answer" + + +class _UnusedRuntime: + def start_session(self, inputs: InferenceInput) -> InferenceSession: + del inputs + raise AssertionError("The fake Phase 4 drivers do not start sessions.") + + def close(self) -> None: + return diff --git a/flashdreams/tests/test_demo_runtime_timing.py b/flashdreams/tests/test_demo_runtime_timing.py new file mode 100644 index 000000000..98e3cf56d --- /dev/null +++ b/flashdreams/tests/test_demo_runtime_timing.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +from typing import Any, cast + +import pytest + +from flashdreams.runtime import StepRequirements, UserInputSchema +from flashdreams.runtime.demo import NullOutputSink, RunResult, SessionEdges +from flashdreams.runtime.demo.timing import ( + SPARSE_KEY_SEGMENTS_METADATA_KEY, + CatchUpDecision, + CatchUpPolicy, + KeyboardRealtimeInputSource, + ResamplerRealtimeClock, + SignalActivationPolicy, +) +from flashdreams.serving.realtime.input import KeyboardResampler + +pytestmark = pytest.mark.ci_cpu + + +@pytest.mark.asyncio +async def test_signal_activation_waits_for_first_input_and_anchors_clock() -> None: + event = asyncio.Event() + resampler = KeyboardResampler(fps=30.0, start_v=0.0) + clock = ResamplerRealtimeClock(resampler=resampler, now_fn=lambda: 12.0) + policy = SignalActivationPolicy(signals=(event,), timeout_s=1.0) + + wait_task = asyncio.create_task(policy.wait_until_active(clock)) + await asyncio.sleep(0) + + assert not wait_task.done() + + event.set() + result = await wait_task + + assert result.activated + assert result.reason is None + assert resampler.next_chunk_start_v == pytest.approx(12.0) + + +@pytest.mark.asyncio +async def test_activation_timeout_can_close_edges_as_not_activated() -> None: + event = asyncio.Event() + resampler = KeyboardResampler(fps=30.0, start_v=0.0) + clock = ResamplerRealtimeClock(resampler=resampler, now_fn=lambda: 12.0) + policy = SignalActivationPolicy( + signals=(event,), + timeout_s=0.001, + timeout_reason="no first input", + ) + cleanup_tasks: set[asyncio.Task[RunResult]] = set() + edges = SessionEdges( + input_source=_OpenRealtimeInputSource(), + output_sink=NullOutputSink(), + cleanup_tasks=cleanup_tasks, + activation=policy, + clock=clock, + ) + + activation = await policy.wait_until_active(clock) + result = edges.close_result( + status="not_activated", + reason=activation.reason, + ) + + assert not activation.activated + assert activation.reason == "no first input" + assert result.status == "not_activated" + assert result.reason == "no first input" + assert edges.is_closed + assert resampler.next_chunk_start_v == pytest.approx(0.0) + + +def test_resampler_clock_catch_up_bounds_latency() -> None: + resampler = KeyboardResampler(fps=1.0, start_v=0.0) + clock = ResamplerRealtimeClock(resampler=resampler, now_fn=lambda: 5.0) + + decision = clock.catch_up( + request=_request(input_frame_count=1), + max_lag_s=1.0, + policy="fold", + ) + + assert decision == CatchUpDecision( + skipped_s=4.0, + skipped_windows=4, + input_policy="fold", + reason="lag exceeded max_lag_s", + ) + assert resampler.next_chunk_start_v == pytest.approx(4.0) + + +@pytest.mark.asyncio +async def test_realtime_input_source_matches_resampler_for_recorded_trace() -> None: + expected_resampler = _resampler_with_recorded_trace() + expected_resampler.next_chunk_start_v = 2.0 + expected_segments, expected_frame_times = expected_resampler.sample_chunk(2) + source_resampler = _resampler_with_recorded_trace() + sleep = _RecordingSleep() + clock = ResamplerRealtimeClock( + resampler=source_resampler, + now_fn=lambda: 3.0, + sleep_fn=sleep, + ) + source = KeyboardRealtimeInputSource(resampler=source_resampler) + + result = await source.next_realtime_window( + request=_request(input_frame_count=2), + clock=clock, + ) + + assert sleep.delays == [] + assert result.catch_up == CatchUpDecision( + skipped_s=2.0, + skipped_windows=2, + input_policy="fold", + reason="lag exceeded max_lag_s", + ) + assert result.window.start_s == pytest.approx(2.0) + assert result.window.end_s == pytest.approx(3.0) + assert result.window.frame_times == tuple(expected_frame_times) + assert result.window.metadata[SPARSE_KEY_SEGMENTS_METADATA_KEY] == tuple( + expected_segments + ) + + +@pytest.mark.asyncio +async def test_backpressure_is_clock_adjustment_not_blocking_sleep() -> None: + resampler = KeyboardResampler(fps=1.0, start_v=0.0) + sleep = _RecordingSleep() + clock = ResamplerRealtimeClock( + resampler=resampler, + now_fn=lambda: 2.2, + sleep_fn=sleep, + ) + + await clock.apply_backpressure(0.3) + decision = clock.catch_up( + request=_request(input_frame_count=1), + max_lag_s=1.0, + policy="fold", + ) + + assert sleep.delays == [] + assert clock.pending_backpressure_s == pytest.approx(0.0) + assert decision.skipped_s == pytest.approx(1.5) + assert decision.skipped_windows == 2 + assert decision.input_policy == "fold" + assert resampler.next_chunk_start_v == pytest.approx(1.5) + + +@pytest.mark.asyncio +async def test_window_floor_sleeps_only_when_virtual_time_is_ahead() -> None: + resampler = KeyboardResampler(fps=1.0, start_v=0.0) + sleep = _RecordingSleep() + clock = ResamplerRealtimeClock( + resampler=resampler, + now_fn=lambda: 1.0, + sleep_fn=sleep, + ) + + await clock.wait_until_window_end(1.25) + await clock.wait_until_window_end(0.75) + + assert sleep.delays == [0.25] + + +@pytest.mark.parametrize("policy", ["drop", "compress"]) +def test_keyboard_resampler_defers_unsupported_catch_up_policies( + policy: str, +) -> None: + resampler = KeyboardResampler(fps=1.0, start_v=0.0) + clock = ResamplerRealtimeClock(resampler=resampler, now_fn=lambda: 5.0) + unsupported_policy = cast(CatchUpPolicy, policy) + + with pytest.raises(NotImplementedError, match="no existing resampler analog"): + clock.catch_up( + request=_request(input_frame_count=1), + max_lag_s=1.0, + policy=unsupported_policy, + ) + + with pytest.raises(NotImplementedError, match="KeyboardResampler analog"): + KeyboardRealtimeInputSource( + resampler=resampler, + catch_up_policy=unsupported_policy, + ) + + +def _request(*, input_frame_count: int) -> StepRequirements: + return StepRequirements( + step_index=0, + input_frame_count=input_frame_count, + ) + + +def _resampler_with_recorded_trace() -> KeyboardResampler: + resampler = KeyboardResampler(fps=2.0, start_v=0.0) + for arrival_t, event, key in ( + (0.25, "keydown", "w"), + (1.25, "keydown", "a"), + (2.25, "keyup", "w"), + (2.75, "keydown", "d"), + ): + resampler.on_edge(arrival_t=arrival_t, event=event, key=key) + return resampler + + +class _RecordingSleep: + def __init__(self) -> None: + self.delays: list[float] = [] + + async def __call__(self, delay_s: float) -> None: + self.delays.append(delay_s) + + +class _OpenRealtimeInputSource: + is_finite = False + is_deterministic = False + user_input_schema = UserInputSchema() + + def is_finished(self) -> bool: + return False + + async def next_realtime_window( + self, + *, + request: StepRequirements, + clock: Any, + ) -> object: + del request, clock + raise AssertionError("Activation timeout must close before requesting input.") diff --git a/flashdreams/tests/test_demo_runtime_validation.py b/flashdreams/tests/test_demo_runtime_validation.py new file mode 100644 index 000000000..c5780bd72 --- /dev/null +++ b/flashdreams/tests/test_demo_runtime_validation.py @@ -0,0 +1,680 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +import pytest + +import flashdreams.runtime.demo as demo_api +from flashdreams.runtime import ( + DRIVER_COMMAND, + CanonicalInputs, + CanonicalInputSchema, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InferenceRuntime, + InputCanonicalizer, + InputField, + InputMapping, + InputMappingSchema, + KeyboardToDriverCommand, + OutputArtifact, + StepRequest, + StepRequirements, + StepResult, + TimeWindow, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + NullOutputSpec, + OutputDecision, + PreparedScenario, + PreparedStep, + ProviderCapabilities, + ResolvedRunCapabilities, + RunContext, + RunModeCapabilities, + RunResult, + RuntimeHost, + SessionEdges, + SessionInfo, + UserInputWindow, + resolve_run_capabilities, + validate_resolved_run, +) +from flashdreams.runtime.demo.timing import RealtimeWindowResult + +pytestmark = pytest.mark.ci_cpu + +KEY_SCHEMA = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="key_down", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_type="key_up", + payload_fields=frozenset({"key"}), + ), + ) +) + + +def test_provider_capabilities_declare_raw_and_inference_schemas() -> None: + capabilities = ProviderCapabilities( + supports_recorded_input=True, + deterministic_given_inputs=True, + user_input_schema=KEY_SCHEMA, + inference_input_schema=InferenceInputSchema( + step_fields=(InputField(name="driver_command"),), + ), + ) + + assert capabilities.user_input_schema.supports( + UserInputCapability(event_type="key_down", payload_fields=frozenset({"key"})) + ) + assert capabilities.inference_input_schema.missing_step(InferenceInput()) == ( + "driver_command", + ) + + +def test_provider_that_wraps_mapping_canonicalizes_raw_inputs_first() -> None: + mapping = _DriverCommandMapping() + provider = _MappingBackedProvider(mapping=mapping, source_schema=KEY_SCHEMA) + source = _BatchInputSource(user_input_schema=KEY_SCHEMA) + edges = _edges(input_source=source) + run_mode = _RunMode( + RunModeCapabilities(requires_finite_input=True, supports_artifacts=True) + ) + resolved = resolve_run_capabilities( + spec=_spec(seed=7, output=Mp4OutputSpec(path="out.mp4", fps=12)), + provider=provider, + session_edges=edges, + ) + + validate_resolved_run( + spec=_spec(seed=7), + adapter=_Adapter(), + provider=provider, + run_mode=run_mode, + session_edges=edges, + resolved=resolved, + ) + prepared = provider.prepare_step( + request=StepRequirements(step_index=4), + user_window=UserInputWindow( + start_s=0.0, + end_s=1.0, + inputs=UserInputs( + events=( + UserInputEvent( + timestamp_s=0.1, + event_type="key_down", + payload={"key": "w"}, + ), + ) + ), + ), + ) + + assert mapping.validated_with is not None + assert mapping.validated_with[0] == CanonicalInputSchema( + modalities=(DRIVER_COMMAND,), + description=KEY_SCHEMA.description, + ) + assert prepared.inference_input is not None + assert prepared.inference_input.step["request_step"] == 4 + assert prepared.inference_input.step["driver_command"]["throttle"] == 1.0 + + +def test_raw_user_schema_validation_is_not_canonical_schema_validation() -> None: + provider = _Provider( + ProviderCapabilities( + supports_recorded_input=True, + user_input_schema=KEY_SCHEMA, + inference_input_schema=InferenceInputSchema( + step_fields=(InputField(name="driver_command"),), + ), + ) + ) + source = _BatchInputSource( + user_input_schema=UserInputSchema(event_types=frozenset({"key_down"})) + ) + run_mode = _RunMode( + RunModeCapabilities(requires_finite_input=True, supports_artifacts=True) + ) + edges = _edges(input_source=source) + resolved = resolve_run_capabilities( + spec=_spec(seed=1), + provider=provider, + session_edges=edges, + ) + + with pytest.raises(ValueError, match="raw user input schema"): + validate_resolved_run( + spec=_spec(seed=1), + adapter=_Adapter(), + provider=provider, + run_mode=run_mode, + session_edges=edges, + resolved=resolved, + ) + + +def test_mp4_rejects_provider_without_recorded_input_support() -> None: + provider = _Provider( + ProviderCapabilities( + supports_recorded_input=False, + supports_realtime_clock=True, + ) + ) + run_mode = _RunMode( + RunModeCapabilities(requires_finite_input=True, supports_artifacts=True) + ) + edges = _edges(output_sink=_OutputSink(produces_artifacts=True)) + resolved = resolve_run_capabilities( + spec=_spec(seed=1, output=Mp4OutputSpec(path="out.mp4", fps=12)), + provider=provider, + session_edges=edges, + ) + + with pytest.raises(ValueError, match="recorded input"): + validate_resolved_run( + spec=_spec(seed=1), + adapter=_Adapter(), + provider=provider, + run_mode=run_mode, + session_edges=edges, + resolved=resolved, + ) + + +def test_webrtc_rejects_provider_without_realtime_input_support() -> None: + provider = _Provider(ProviderCapabilities(supports_recorded_input=True)) + run_mode = _RunMode( + RunModeCapabilities( + realtime=True, + supports_backpressure=True, + supports_interactive_events=True, + ) + ) + edges = _edges(input_source=_RealtimeInputSource(), clock=_Clock(realtime=True)) + resolved = resolve_run_capabilities( + spec=_spec(seed=1), + provider=provider, + session_edges=edges, + ) + + with pytest.raises(ValueError, match="realtime input"): + validate_resolved_run( + spec=_spec(seed=1), + adapter=_Adapter(), + provider=provider, + run_mode=run_mode, + session_edges=edges, + resolved=resolved, + ) + + +def test_realtime_run_mode_rejects_batch_input_source() -> None: + provider = _Provider(ProviderCapabilities(supports_realtime_clock=True)) + run_mode = _RunMode(RunModeCapabilities(realtime=True)) + edges = _edges(input_source=_BatchInputSource(), clock=_Clock(realtime=True)) + resolved = resolve_run_capabilities( + spec=_spec(seed=1), + provider=provider, + session_edges=edges, + ) + + with pytest.raises(ValueError, match="RealtimeInputSource"): + validate_resolved_run( + spec=_spec(seed=1), + adapter=_Adapter(), + provider=provider, + run_mode=run_mode, + session_edges=edges, + resolved=resolved, + ) + + +def test_determinism_resolves_from_provider_source_clock_and_seed() -> None: + provider = _Provider( + ProviderCapabilities( + supports_recorded_input=True, + supports_reset=True, + deterministic_given_inputs=True, + ) + ) + + deterministic = resolve_run_capabilities( + spec=_spec(seed=123), + provider=provider, + session_edges=_edges(clock=_Clock(deterministic=True)), + ) + unseeded = resolve_run_capabilities( + spec=_spec(seed=None), + provider=provider, + session_edges=_edges(clock=_Clock(deterministic=True)), + ) + nondeterministic_source = resolve_run_capabilities( + spec=_spec(seed=123), + provider=provider, + session_edges=_edges( + input_source=_BatchInputSource(deterministic=False), + clock=_Clock(deterministic=True), + ), + ) + + assert deterministic == ResolvedRunCapabilities( + finite=True, + deterministic=True, + realtime=False, + resettable=True, + produces_artifacts=True, + ) + assert not unseeded.deterministic + assert not nondeterministic_source.deterministic + + +def test_no_general_purpose_input_mapping_provider_is_exported() -> None: + assert not hasattr(demo_api, "InputMappingProvider") + + +def test_reset_control_updates_provider_and_session_together() -> None: + reset_input = InferenceInput(global_conditioning={"prompt": "reset"}) + provider = _ResettingProvider(reset_input=reset_input) + session = _ResettableSession(num_steps=2) + edges = _edges(input_source=_BatchInputSource(num_windows=2)) + + result = demo_api.BatchSessionDriver().run_one_session( + host=RuntimeHost(_Runtime(session=session)), + provider=provider, + session_edges=edges, + pipeline=demo_api.StepPipeline(), + ) + + assert result.status == "completed" + assert session.reset_inputs == [reset_input] + assert provider.reset_inputs == [reset_input] + assert len(session.step_inputs) == 1 + + +def _spec( + *, + seed: int | None, + output: Any | None = None, +) -> DemoSpec: + return DemoSpec( + model_id="fake-demo", + input_mode="replay", + output=output or NullOutputSpec(), + config=InferenceConfig(model_id="fake-demo", seed=seed), + ) + + +def _edges( + *, + input_source: Any | None = None, + output_sink: Any | None = None, + clock: Any | None = None, +) -> SessionEdges: + return SessionEdges( + input_source=input_source or _BatchInputSource(), + output_sink=output_sink or _OutputSink(produces_artifacts=True), + cleanup_tasks=set(), + clock=clock, + ) + + +class _Adapter: + model_id = "fake-demo" + inference_input_schema = InferenceInputSchema() + canonical_input_schema = None + + def supported_input_modes(self) -> tuple[str, ...]: + return ("replay",) + + def supported_output_modes(self) -> tuple[str, ...]: + return ("null", "mp4") + + def default_input_mapping(self) -> InputMapping: + return IdentityInputMapping() + + def validate_config(self, config: InferenceConfig) -> None: + del config + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + del config + raise NotImplementedError + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + del spec + return PreparedScenario(initial_inputs=InferenceInput()) + + +class _Provider: + def __init__(self, capabilities: ProviderCapabilities) -> None: + self.capabilities = capabilities + + def prepare_initial_input(self) -> InferenceInput: + return InferenceInput() + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + del request, user_window + return PreparedStep(inference_input=InferenceInput()) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + + def close(self) -> None: + return + + +class _MappingBackedProvider(_Provider): + def __init__( + self, *, mapping: "_DriverCommandMapping", source_schema: UserInputSchema + ) -> None: + self.mapping = mapping + self.canonicalizer = InputCanonicalizer((KeyboardToDriverCommand(),)) + self.source_schema = source_schema + capabilities = ProviderCapabilities( + supports_recorded_input=True, + supports_reset=True, + deterministic_given_inputs=True, + user_input_schema=source_schema, + inference_input_schema=InferenceInputSchema( + step_fields=(InputField(name="driver_command"),), + ), + ) + super().__init__(capabilities) + self.mapping.validate( + canonical_schema=self.canonicalizer.canonical_schema(source_schema), + inference_input_schema=capabilities.inference_input_schema, + ) + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + canonical_inputs = self.canonicalizer.canonicalize( + user_window.inputs, + window=TimeWindow(start_s=user_window.start_s, end_s=user_window.end_s), + source_schema=self.source_schema, + ) + inference_input = self.mapping.map_step_inputs( + canonical_inputs=canonical_inputs, + inference_input=InferenceInput(), + request=StepRequest(step_index=request.step_index), + ) + return PreparedStep(inference_input=inference_input) + + +class _DriverCommandMapping: + mapping_schema = InputMappingSchema( + name="driver-command", + consumes=(DRIVER_COMMAND,), + produces_step=(InputField(name="driver_command"),), + ) + + def __init__(self) -> None: + self.validated_with: ( + tuple[ + CanonicalInputSchema | None, + InferenceInputSchema | None, + ] + | None + ) = None + + def validate( + self, + *, + canonical_schema: CanonicalInputSchema | None = None, + inference_input_schema: InferenceInputSchema | None = None, + ) -> None: + if canonical_schema is not None and not canonical_schema.supports( + DRIVER_COMMAND + ): + raise ValueError("mapping cannot be fed") + if inference_input_schema is not None: + inference_input_schema.require_step( + InferenceInput(step={"driver_command": object()}) + ) + self.validated_with = (canonical_schema, inference_input_schema) + + def map_global_conditioning_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + ) -> InferenceInput: + del canonical_inputs + return inference_input + + def map_step_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + request: StepRequest, + ) -> InferenceInput: + del inference_input + return InferenceInput( + step={ + "driver_command": canonical_inputs.values["driver_command"], + "request_step": request.step_index, + } + ) + + +class _RunMode: + name = "fake" + + def __init__(self, capabilities: RunModeCapabilities) -> None: + self.capabilities = capabilities + + def validate_run(self, *, spec: DemoSpec, adapter: Any) -> None: + del spec, adapter + + def validate_session( + self, + *, + spec: DemoSpec, + scenario: PreparedScenario, + adapter: Any, + provider: Any, + ) -> None: + del spec, scenario, adapter, provider + + def create_run_context( + self, + *, + spec: DemoSpec, + adapter: Any, + host: RuntimeHost, + model_warmup_plan: Any, + ) -> RunContext: + del spec, adapter, model_warmup_plan + return RunContext( + host=host, + run_metrics=demo_api.InMemorySessionMetricsRecorder(), + admission=demo_api.SingleSessionAdmissionPolicy(), + ) + + def create_session_edges( + self, + *, + context: RunContext, + spec: DemoSpec, + scenario: PreparedScenario, + provider: Any, + adapter: Any, + ) -> SessionEdges: + del context, spec, scenario, provider, adapter + return _edges() + + def select_driver(self) -> Any: + raise NotImplementedError + + +class _BatchInputSource: + is_finite = True + + def __init__( + self, + *, + user_input_schema: UserInputSchema | None = None, + deterministic: bool = True, + num_windows: int = 1, + ) -> None: + self.user_input_schema = user_input_schema or UserInputSchema() + self.is_deterministic = deterministic + self.num_windows = num_windows + self.index = 0 + + def is_finished(self) -> bool: + return self.index >= self.num_windows + + def next_window(self, request: StepRequirements) -> UserInputWindow: + del request + self.index += 1 + return UserInputWindow(start_s=0.0, end_s=1.0) + + +class _RealtimeInputSource: + is_finite = False + is_deterministic = False + user_input_schema = UserInputSchema() + + def is_finished(self) -> bool: + return False + + async def next_realtime_window( + self, + *, + request: StepRequirements, + clock: Any, + ) -> RealtimeWindowResult: + del request, clock + return RealtimeWindowResult(window=UserInputWindow(start_s=0.0, end_s=1.0)) + + +class _Clock: + def __init__(self, *, realtime: bool = False, deterministic: bool = True) -> None: + self.is_realtime = realtime + self.is_deterministic = deterministic + + +class _OutputSink: + def __init__(self, *, produces_artifacts: bool) -> None: + self.produces_artifacts = produces_artifacts + + def open(self, session_info: SessionInfo) -> None: + del session_info + + def begin_generation(self, generation: int) -> None: + del generation + + def write(self, result: StepResult) -> OutputDecision: + del result + return OutputDecision() + + def close(self) -> Sequence[OutputArtifact]: + return () + + +class _ResettingProvider(_Provider): + def __init__(self, *, reset_input: InferenceInput) -> None: + self.reset_input = reset_input + self.prepare_count = 0 + self.reset_inputs: list[InferenceInput | None] = [] + super().__init__( + ProviderCapabilities( + supports_recorded_input=True, + supports_reset=True, + deterministic_given_inputs=True, + ) + ) + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + del request, user_window + self.prepare_count += 1 + if self.prepare_count == 1: + return PreparedStep( + control=demo_api.ControlDecision( + reset=True, + reset_input=self.reset_input, + ) + ) + return PreparedStep(inference_input=InferenceInput(step={"after_reset": True})) + + def reset(self, inputs: InferenceInput | None = None) -> None: + self.reset_inputs.append(inputs) + + +class _Runtime: + def __init__(self, *, session: "_ResettableSession") -> None: + self.session = session + + def start_session(self, inputs: InferenceInput) -> "_ResettableSession": + del inputs + return self.session + + def close(self) -> None: + return + + +class _ResettableSession: + def __init__(self, *, num_steps: int) -> None: + self.num_steps = num_steps + self.next_request_index = 0 + self.reset_inputs: list[InferenceInput | None] = [] + self.step_inputs: list[InferenceInput] = [] + + def session_info(self) -> SessionInfo: + return SessionInfo() + + def next_step_requirements(self) -> StepRequirements | None: + if self.next_request_index >= self.num_steps: + return None + request = StepRequirements(step_index=self.next_request_index) + self.next_request_index += 1 + return request + + def next_step_request(self) -> StepRequest | None: + requirements = self.next_step_requirements() + if requirements is None: + return None + return StepRequest(step_index=requirements.step_index) + + def step(self, inputs: InferenceInput) -> StepResult: + self.step_inputs.append(inputs) + return StepResult(step_index=len(self.step_inputs) - 1, output=None) + + def reset(self, inputs: InferenceInput | None = None) -> None: + self.reset_inputs.append(inputs) + self.next_request_index = 0 + + def close(self) -> None: + return diff --git a/flashdreams/tests/test_demo_runtime_vertical_slice.py b/flashdreams/tests/test_demo_runtime_vertical_slice.py new file mode 100644 index 000000000..1f430a459 --- /dev/null +++ b/flashdreams/tests/test_demo_runtime_vertical_slice.py @@ -0,0 +1,1328 @@ +# 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 +from collections.abc import Callable, Sequence +from typing import Any, Literal, cast + +import pytest + +import flashdreams.runtime.demo.drivers as drivers_module +from flashdreams.runtime import ( + CanonicalInputSchema, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InferenceRuntime, + InferenceSession, + InputMapping, + OutputArtifact, + StepRequest, + StepRequirements, + StepResult, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.demo import ( + BatchSessionDriver, + ControlDecision, + DemoSpec, + DriverInvariantError, + ErrorAction, + InMemorySessionMetricsRecorder, + ModelWarmupPlan, + NullOutputSpec, + OutputDecision, + PreparedScenario, + PreparedStep, + ProviderCapabilities, + RunContext, + RunModeCapabilities, + RunResult, + RuntimeHost, + SessionEdges, + SessionInfo, + SingleSessionAdmissionPolicy, + StepPipeline, + UserInputWindow, + run_demo_session, + run_demo_session_async, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_step_pipeline_passes_provider_input_to_session_and_sink() -> None: + provider = _FakeVideoModelInputProvider() + session = _FakeVideoSession(num_steps=1) + output = _RecordingOutputSink() + output.open(SessionInfo(output_layout="fake-video", steady_output_frame_count=1)) + metrics = InMemorySessionMetricsRecorder() + request = StepRequirements(step_index=0) + user_window = _window(0) + + outcome = StepPipeline().execute_step( + request=request, + user_window=user_window, + provider=provider, + session=session, + output=output, + metrics=metrics, + ) + + assert outcome == _empty_step_outcome() + assert session.step_inputs == provider.prepared_step_inputs + assert [result.output for result in output.results] == ["frame-0"] + assert metrics.step_count == 1 + + +def test_batch_driver_runs_fake_video_demo_through_runtime_host() -> None: + session = _FakeVideoSession(num_steps=2) + runtime = _FakeVideoRuntime(session=session) + host = _RecordingRuntimeHost(runtime) + provider = _FakeVideoModelInputProvider() + output = _RecordingOutputSink() + metrics = InMemorySessionMetricsRecorder() + edges = SessionEdges( + input_source=_FakeBatchInputSource(num_windows=2), + output_sink=output, + cleanup_tasks=set(), + metrics=metrics, + ) + + result = BatchSessionDriver().run_one_session( + host=host, + provider=provider, + session_edges=edges, + pipeline=StepPipeline(), + ) + + assert result.status == "completed" + assert result.metrics is not None + assert result.metrics.counters["steps"] == 2 + assert runtime.start_session_inputs == [provider.initial_input] + assert [dict(inputs.step) for inputs in session.step_inputs] == [ + {"request_step": 0, "window": (0.0, 1.0)}, + {"request_step": 1, "window": (1.0, 2.0)}, + ] + assert [result.output for result in output.results] == ["frame-0", "frame-1"] + assert output.opened_with == SessionInfo( + output_layout="fake-video", + steady_output_frame_count=1, + ) + assert session.close_count == 1 + assert provider.close_count == 1 + assert host.calls.count("execute_step") == 2 + assert "prepare_initial_input" in host.calls + assert "start_session" in host.calls + assert "prepare_step" not in host.calls + assert "step" not in host.calls + + +def test_batch_driver_cleanup_failure_marks_host_unhealthy() -> None: + session = _FakeVideoSession(num_steps=1) + runtime = _FakeVideoRuntime(session=session) + host = RuntimeHost(runtime) + provider = _FakeVideoModelInputProvider( + fail_close=RuntimeError("provider close failed") + ) + metrics = InMemorySessionMetricsRecorder() + + try: + result = BatchSessionDriver().run_one_session( + host=host, + provider=provider, + session_edges=SessionEdges( + input_source=_FakeBatchInputSource(num_windows=1), + output_sink=_RecordingOutputSink(), + cleanup_tasks=set(), + metrics=metrics, + ), + pipeline=StepPipeline(), + ) + + assert result.status == "completed" + assert not host.is_healthy + assert host.unhealthy_reason == "model-affine cleanup failed" + assert provider.close_count == 1 + assert metrics.cleanup_errors == ["provider close failed"] + finally: + host.close() + + +def test_batch_driver_slices_windows_from_step_requirements() -> None: + session = _FakeVideoSession(num_steps=2, input_frame_counts=(3, 2)) + runtime = _FakeVideoRuntime(session=session) + host = _RecordingRuntimeHost(runtime) + provider = _FakeVideoModelInputProvider() + input_source = _SlicingBatchInputSource(fps=2.0, num_windows=2) + + result = BatchSessionDriver().run_one_session( + host=host, + provider=provider, + session_edges=SessionEdges( + input_source=input_source, + output_sink=_RecordingOutputSink(), + cleanup_tasks=set(), + metrics=InMemorySessionMetricsRecorder(), + ), + pipeline=StepPipeline(), + ) + + assert result.status == "completed" + assert [request.step_index for request in input_source.next_window_requests] == [ + 0, + 1, + ] + assert [ + request.input_frame_count for request in input_source.next_window_requests + ] == [3, 2] + assert input_source.windows == [ + _window_with_frame_times(start_s=0.0, frame_times=(0.0, 0.5, 1.0)), + _window_with_frame_times(start_s=1.5, frame_times=(1.5, 2.0)), + ] + assert [dict(inputs.step) for inputs in session.step_inputs] == [ + {"request_step": 0, "window": (0.0, 1.5)}, + {"request_step": 1, "window": (1.5, 2.5)}, + ] + + +def test_run_demo_session_builds_edges_and_records_session_once() -> None: + session = _FakeVideoSession(num_steps=1) + runtime = _FakeVideoRuntime(session=session) + run_metrics = InMemorySessionMetricsRecorder() + context = _run_context(runtime, run_metrics=run_metrics) + provider = _FakeVideoModelInputProvider() + adapter = _FakeDemoAdapter(provider=provider) + output = _RecordingOutputSink() + factory_calls: list[tuple[DemoSpec, PreparedScenario]] = [] + run_mode = _FakeRunMode( + input_source=_FakeBatchInputSource(num_windows=1), + output_sink_factory=lambda spec, scenario: _record_output_factory_call( + factory_calls, + spec, + scenario, + output, + ), + ) + spec = _spec() + scenario = _scenario() + + result = run_demo_session( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=run_mode, + pipeline=StepPipeline(), + ) + + assert result.status == "completed" + assert adapter.provider_calls == [(spec, scenario)] + assert factory_calls == [(spec, scenario)] + assert run_metrics.sessions == [result] + assert len(run_metrics.sessions) == 1 + new_reservation = context.admission.try_reserve() + assert new_reservation is not None + new_reservation.release() + + +def test_busy_admission_returns_rejected_and_records_once() -> None: + runtime = _FakeVideoRuntime(session=_FakeVideoSession(num_steps=1)) + admission = SingleSessionAdmissionPolicy() + held = admission.try_reserve() + assert held is not None + run_metrics = InMemorySessionMetricsRecorder() + context = _run_context(runtime, admission=admission, run_metrics=run_metrics) + adapter = _FakeDemoAdapter(provider=_FakeVideoModelInputProvider()) + + result = run_demo_session( + context=context, + spec=_spec(), + scenario=_scenario(), + adapter=adapter, + run_mode=_FakeRunMode(input_source=_FakeBatchInputSource(num_windows=1)), + pipeline=StepPipeline(), + ) + + held.release() + assert result == RunResult.rejected(reason="busy") + assert run_metrics.sessions == [result] + assert adapter.provider_calls == [] + assert runtime.start_session_inputs == [] + + +def test_setup_failure_returns_failed_before_runtime_session_creation() -> None: + runtime = _FakeVideoRuntime(session=_FakeVideoSession(num_steps=1)) + provider = _FakeVideoModelInputProvider( + fail_initial=ValueError("invalid provider compatibility") + ) + metrics = InMemorySessionMetricsRecorder() + + result = BatchSessionDriver().run_one_session( + host=RuntimeHost(runtime), + provider=provider, + session_edges=SessionEdges( + input_source=_FakeBatchInputSource(num_windows=1), + output_sink=_RecordingOutputSink(), + cleanup_tasks=set(), + metrics=metrics, + ), + pipeline=StepPipeline(), + ) + + assert result.status == "failed" + assert isinstance(result.error, ValueError) + assert result.reason == "invalid provider compatibility" + assert runtime.start_session_inputs == [] + assert provider.close_count == 1 + assert metrics.errors == ["invalid provider compatibility"] + + +def test_output_sink_open_failure_returns_failed_before_step_loop() -> None: + session = _FakeVideoSession(num_steps=1) + metrics = InMemorySessionMetricsRecorder() + output = _RecordingOutputSink(fail_open=RuntimeError("open failed")) + + result = BatchSessionDriver().run_one_session( + host=RuntimeHost(_FakeVideoRuntime(session=session)), + provider=_FakeVideoModelInputProvider(), + session_edges=SessionEdges( + input_source=_FakeBatchInputSource(num_windows=1), + output_sink=output, + cleanup_tasks=set(), + metrics=metrics, + ), + pipeline=StepPipeline(), + ) + + assert result.status == "failed" + assert result.reason == "open failed" + assert metrics.errors == ["open failed"] + assert session.step_inputs == [] + assert output.results == [] + assert output.close_count == 1 + + +def test_run_demo_session_closes_provider_when_validation_fails() -> None: + runtime = _FakeVideoRuntime(session=_FakeVideoSession(num_steps=1)) + run_metrics = InMemorySessionMetricsRecorder() + context = _run_context(runtime, run_metrics=run_metrics) + provider = _FakeVideoModelInputProvider() + + result = run_demo_session( + context=context, + spec=_spec(), + scenario=_scenario(), + adapter=_FakeDemoAdapter(provider=provider), + run_mode=_FakeRunMode( + input_source=_FakeBatchInputSource(num_windows=1), + validate_error=ValueError("provider incompatible"), + ), + pipeline=StepPipeline(), + ) + + assert result.status == "failed" + assert result.reason == "provider incompatible" + assert provider.close_count == 1 + assert runtime.start_session_inputs == [] + assert run_metrics.sessions == [result] + assert run_metrics.session_errors == ["provider incompatible"] + snapshot = run_metrics.close() + assert snapshot.counters["sessions"] == 1 + assert snapshot.counters["session_errors"] == 1 + assert snapshot.session_statuses == ("failed",) + + +def test_run_demo_session_keeps_failure_when_run_cleanup_metrics_fail() -> None: + runtime = _FakeVideoRuntime(session=_FakeVideoSession(num_steps=1)) + run_metrics = _FailingCleanupMetrics() + context = _run_context(runtime, run_metrics=run_metrics) + provider = _FakeVideoModelInputProvider( + fail_close=RuntimeError("provider close failed") + ) + + result = run_demo_session( + context=context, + spec=_spec(), + scenario=_scenario(), + adapter=_FakeDemoAdapter(provider=provider), + run_mode=_FakeRunMode( + input_source=_FakeBatchInputSource(num_windows=1), + validate_error=ValueError("provider incompatible"), + ), + pipeline=StepPipeline(), + ) + + assert result.status == "failed" + assert result.reason == "provider incompatible" + assert provider.close_count == 1 + assert not context.host.is_healthy + assert context.host.unhealthy_reason == "model-affine cleanup failed" + assert run_metrics.cleanup_error_attempts == 1 + assert run_metrics.sessions == [result] + assert runtime.start_session_inputs == [] + + +@pytest.mark.asyncio +async def test_run_demo_session_async_keeps_failure_when_run_cleanup_metrics_fail() -> ( + None +): + runtime = _FakeVideoRuntime(session=_FakeVideoSession(num_steps=1)) + run_metrics = _FailingCleanupMetrics() + context = _run_context(runtime, run_metrics=run_metrics) + provider = _FakeVideoModelInputProvider( + fail_close=RuntimeError("provider close failed") + ) + + result = await run_demo_session_async( + context=context, + spec=_spec(), + scenario=_scenario(), + adapter=_FakeDemoAdapter(provider=provider), + run_mode=_FakeRunMode( + input_source=_FakeBatchInputSource(num_windows=1), + validate_error=ValueError("provider incompatible"), + ), + pipeline=StepPipeline(), + ) + + assert result.status == "failed" + assert result.reason == "provider incompatible" + assert provider.close_count == 1 + assert not context.host.is_healthy + assert context.host.unhealthy_reason == "model-affine cleanup failed" + assert run_metrics.cleanup_error_attempts == 1 + assert run_metrics.sessions == [result] + assert runtime.start_session_inputs == [] + + +@pytest.mark.asyncio +async def test_run_demo_session_async_invariant_cancellation_finalizes_edges() -> None: + close_entered = threading.Event() + release_close = threading.Event() + runtime = _FakeVideoRuntime(session=_FakeVideoSession(num_steps=1)) + run_metrics = InMemorySessionMetricsRecorder() + context = _run_context(runtime, run_metrics=run_metrics) + provider = _BlockingCloseVideoModelInputProvider( + close_entered=close_entered, + release_close=release_close, + ) + output = _RecordingOutputSink() + transport = _RecordingTransport() + session_metrics = InMemorySessionMetricsRecorder() + select_error = DriverInvariantError("select driver invariant") + task = asyncio.create_task( + run_demo_session_async( + context=context, + spec=_spec(), + scenario=_scenario(), + adapter=_FakeDemoAdapter(provider=provider), + run_mode=_FakeRunMode( + input_source=_FakeBatchInputSource(num_windows=1), + output_sink=output, + metrics=session_metrics, + transport=transport, + select_error=select_error, + ), + pipeline=StepPipeline(), + ) + ) + + try: + assert await asyncio.to_thread(close_entered.wait, 2.0) + task.cancel() + await asyncio.sleep(0) + release_close.set() + with pytest.raises( + DriverInvariantError, match="select driver invariant" + ) as raised: + await task + finally: + release_close.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + context.host.close() + + assert raised.value is select_error + assert provider.close_count == 1 + assert output.close_count == 1 + assert transport.close_count == 1 + assert session_metrics.closed + assert len(run_metrics.sessions) == 1 + recorded = cast(RunResult, run_metrics.sessions[0]) + assert recorded.status == "failed" + assert recorded.error is select_error + assert runtime.start_session_inputs == [] + + +@pytest.mark.asyncio +async def test_run_demo_session_async_cancels_before_driver_owns_cleanup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runtime = _FakeVideoRuntime(session=_FakeVideoSession(num_steps=1)) + run_metrics = InMemorySessionMetricsRecorder() + context = _run_context(runtime, run_metrics=run_metrics) + provider = _FakeVideoModelInputProvider() + output = _RecordingOutputSink() + transport = _RecordingTransport() + session_metrics = InMemorySessionMetricsRecorder() + driver_boundary_reached = asyncio.Event() + release_driver = asyncio.Event() + + async def fake_run_async_driver( + *, + driver: object, + host: RuntimeHost, + provider: Any, + session_edges: SessionEdges, + pipeline: StepPipeline, + ) -> RunResult: + del driver, host, provider, session_edges, pipeline + driver_boundary_reached.set() + await release_driver.wait() + return RunResult(status="completed") + + monkeypatch.setattr( + drivers_module, + "_run_async_driver", + fake_run_async_driver, + ) + task = asyncio.create_task( + run_demo_session_async( + context=context, + spec=_spec(), + scenario=_scenario(), + adapter=_FakeDemoAdapter(provider=provider), + run_mode=_FakeRunMode( + input_source=_FakeBatchInputSource(num_windows=1), + output_sink=output, + metrics=session_metrics, + transport=transport, + ), + pipeline=StepPipeline(), + ) + ) + + try: + await asyncio.wait_for(driver_boundary_reached.wait(), timeout=2.0) + task.cancel() + result = await task + finally: + release_driver.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + context.host.close() + + assert result.status == "cancelled" + assert result.reason == "cancelled during session assembly" + assert provider.close_count == 1 + assert output.close_count == 1 + assert transport.close_count == 1 + assert session_metrics.closed + assert run_metrics.sessions == [result] + assert runtime.start_session_inputs == [] + + +def test_setup_failure_can_return_skipped_but_not_completed() -> None: + skipped = BatchSessionDriver().run_one_session( + host=RuntimeHost(_FakeVideoRuntime(session=_FakeVideoSession(num_steps=1))), + provider=_FakeVideoModelInputProvider(fail_initial=RuntimeError("skip me")), + session_edges=SessionEdges( + input_source=_FakeBatchInputSource(num_windows=1), + output_sink=_RecordingOutputSink(), + cleanup_tasks=set(), + error_policy=_SetupPolicy(result_status="skipped"), + ), + pipeline=StepPipeline(), + ) + assert skipped.status == "skipped" + assert skipped.error is None + + provider = _FakeVideoModelInputProvider(fail_initial=RuntimeError("bad policy")) + output = _RecordingOutputSink() + transport = _RecordingTransport() + metrics = InMemorySessionMetricsRecorder() + edges = SessionEdges( + input_source=_FakeBatchInputSource(num_windows=1), + output_sink=output, + cleanup_tasks=set(), + metrics=metrics, + error_policy=_SetupPolicy(result_status="completed"), + transport=transport, + ) + + with pytest.raises(DriverInvariantError, match="Setup failures"): + BatchSessionDriver().run_one_session( + host=RuntimeHost(_FakeVideoRuntime(session=_FakeVideoSession(num_steps=1))), + provider=provider, + session_edges=edges, + pipeline=StepPipeline(), + ) + + assert output.close_count == 1 + assert transport.close_count == 1 + assert metrics.closed + assert provider.close_count == 1 + + +def test_batch_driver_invariant_finalizes_edges_when_host_closed() -> None: + runtime = _FakeVideoRuntime(session=_FakeVideoSession(num_steps=1)) + host = RuntimeHost(runtime) + host.close() + provider = _FakeVideoModelInputProvider() + output = _RecordingOutputSink() + transport = _RecordingTransport() + metrics = InMemorySessionMetricsRecorder() + edges = SessionEdges( + input_source=_FakeBatchInputSource(num_windows=1), + output_sink=output, + cleanup_tasks=set(), + metrics=metrics, + error_policy=_SetupPolicy(result_status="completed"), + transport=transport, + ) + + with pytest.raises(DriverInvariantError, match="Setup failures"): + BatchSessionDriver().run_one_session( + host=host, + provider=provider, + session_edges=edges, + pipeline=StepPipeline(), + ) + + assert edges.is_closed + assert output.close_count == 1 + assert transport.close_count == 1 + assert metrics.closed + assert metrics.cleanup_errors == ["runtime host is closed"] + assert provider.close_count == 0 + + +def test_batch_driver_ordinary_cleanup_finalizes_edges_when_host_closed() -> None: + session = _FakeVideoSession(num_steps=1) + runtime = _FakeVideoRuntime(session=session) + host = _ClosingAfterStepRuntimeHost(runtime) + provider = _FakeVideoModelInputProvider() + output = _RecordingOutputSink() + transport = _RecordingTransport() + metrics = InMemorySessionMetricsRecorder() + + result = BatchSessionDriver().run_one_session( + host=host, + provider=provider, + session_edges=SessionEdges( + input_source=_FakeBatchInputSource(num_windows=1), + output_sink=output, + cleanup_tasks=set(), + metrics=metrics, + transport=transport, + ), + pipeline=StepPipeline(), + ) + + assert result.status == "completed" + assert result.metrics is not None + assert result.metrics.counters["steps"] == 1 + assert result.metrics.counters["cleanup_errors"] == 2 + assert output.close_count == 1 + assert transport.close_count == 1 + assert metrics.closed + assert metrics.cleanup_errors == [ + "runtime host is closed", + "runtime host is closed", + ] + assert session.close_count == 0 + assert provider.close_count == 0 + + +def test_batch_driver_invariant_finalizes_edges_when_cleanup_metrics_fail() -> None: + runtime = _FakeVideoRuntime(session=_FakeVideoSession(num_steps=1)) + host = RuntimeHost(runtime) + host.close() + provider = _FakeVideoModelInputProvider() + output = _RecordingOutputSink() + transport = _RecordingTransport() + metrics = _FailingCleanupMetrics() + edges = SessionEdges( + input_source=_FakeBatchInputSource(num_windows=1), + output_sink=output, + cleanup_tasks=set(), + metrics=metrics, + error_policy=_SetupPolicy(result_status="completed"), + transport=transport, + ) + + with pytest.raises(DriverInvariantError, match="Setup failures") as raised: + BatchSessionDriver().run_one_session( + host=host, + provider=provider, + session_edges=edges, + pipeline=StepPipeline(), + ) + + result = edges.close_result() + assert result.status == "failed" + assert result.error is raised.value + assert output.close_count == 1 + assert transport.close_count == 1 + assert metrics.closed + assert metrics.cleanup_error_attempts == 1 + assert provider.close_count == 0 + + +def test_run_demo_session_closes_edges_when_driver_invariant_escapes() -> None: + runtime = _FakeVideoRuntime(session=_FakeVideoSession(num_steps=1)) + run_metrics = InMemorySessionMetricsRecorder() + context = _run_context(runtime, run_metrics=run_metrics) + provider = _FakeVideoModelInputProvider( + fail_initial=RuntimeError("bad setup policy") + ) + output = _RecordingOutputSink() + transport = _RecordingTransport() + session_metrics = InMemorySessionMetricsRecorder() + + with pytest.raises(DriverInvariantError, match="Setup failures"): + run_demo_session( + context=context, + spec=_spec(), + scenario=_scenario(), + adapter=_FakeDemoAdapter(provider=provider), + run_mode=_FakeRunMode( + input_source=_FakeBatchInputSource(num_windows=1), + output_sink=output, + metrics=session_metrics, + transport=transport, + error_policy=_SetupPolicy(result_status="completed"), + ), + pipeline=StepPipeline(), + ) + + assert output.close_count == 1 + assert transport.close_count == 1 + assert session_metrics.closed + assert provider.close_count == 1 + assert len(run_metrics.sessions) == 1 + recorded = cast(RunResult, run_metrics.sessions[0]) + assert recorded.status == "failed" + assert isinstance(recorded.error, DriverInvariantError) + + +def test_input_source_finished_error_returns_failed_not_completed() -> None: + metrics = InMemorySessionMetricsRecorder() + + result = BatchSessionDriver().run_one_session( + host=RuntimeHost(_FakeVideoRuntime(session=_FakeVideoSession(num_steps=1))), + provider=_FakeVideoModelInputProvider(), + session_edges=SessionEdges( + input_source=_FakeBatchInputSource( + num_windows=1, + fail_is_finished=RuntimeError("input source failed"), + ), + output_sink=_RecordingOutputSink(), + cleanup_tasks=set(), + metrics=metrics, + ), + pipeline=StepPipeline(), + ) + + assert result.status == "failed" + assert result.reason == "input source failed" + assert metrics.errors == ["input source failed"] + + +def test_step_failure_returns_failed_from_driver() -> None: + session = _FakeVideoSession(num_steps=1, fail_step=0) + output = _RecordingOutputSink() + + result = BatchSessionDriver().run_one_session( + host=RuntimeHost(_FakeVideoRuntime(session=session)), + provider=_FakeVideoModelInputProvider(), + session_edges=SessionEdges( + input_source=_FakeBatchInputSource(num_windows=1), + output_sink=output, + cleanup_tasks=set(), + ), + pipeline=StepPipeline(), + ) + + assert result.status == "failed" + assert isinstance(result.error, RuntimeError) + assert result.reason == "step failed" + assert output.results == [] + assert session.close_count == 1 + + +def test_session_edges_close_result_is_idempotent_and_first_result_wins() -> None: + output = _RecordingOutputSink( + artifacts=(OutputArtifact(kind="test/artifact", uri="memory://artifact"),) + ) + transport = _RecordingTransport() + metrics = InMemorySessionMetricsRecorder() + edges = SessionEdges( + input_source=_FakeBatchInputSource(num_windows=0), + output_sink=output, + cleanup_tasks=set(), + metrics=metrics, + transport=transport, + ) + first_error = RuntimeError("first") + + first = edges.close_result( + status="failed", + reason="first", + error=first_error, + ) + second = edges.close_result(status="completed") + + assert second is first + assert first.status == "failed" + assert first.reason == "first" + assert first.error is first_error + assert tuple(first.artifacts) == ( + OutputArtifact(kind="test/artifact", uri="memory://artifact"), + ) + assert output.close_count == 1 + assert transport.close_count == 1 + assert metrics.closed + + +def test_output_sink_close_failure_records_cleanup_error_without_losing_result() -> ( + None +): + session = _FakeVideoSession(num_steps=1) + metrics = InMemorySessionMetricsRecorder() + output = _RecordingOutputSink(fail_close=RuntimeError("close failed")) + + result = BatchSessionDriver().run_one_session( + host=RuntimeHost(_FakeVideoRuntime(session=session)), + provider=_FakeVideoModelInputProvider(), + session_edges=SessionEdges( + input_source=_FakeBatchInputSource(num_windows=1), + output_sink=output, + cleanup_tasks=set(), + metrics=metrics, + ), + pipeline=StepPipeline(), + ) + + assert result.status == "completed" + assert result.reason is None + assert result.metrics is not None + assert result.metrics.counters["steps"] == 1 + assert result.metrics.counters["cleanup_errors"] == 1 + assert result.metrics.errors == ("close failed",) + assert output.close_count == 1 + + +def test_run_result_rejected_is_the_only_convenience_constructor() -> None: + constructors = { + name + for name, value in RunResult.__dict__.items() + if isinstance(value, classmethod) + } + + assert constructors == {"rejected"} + assert RunResult.rejected(reason="busy").status == "rejected" + + +def _empty_step_outcome() -> Any: + from flashdreams.runtime.demo import StepOutcome + + return StepOutcome(output=OutputDecision(), control=ControlDecision()) + + +def _window(index: int) -> UserInputWindow: + start_s = float(index) + return UserInputWindow( + start_s=start_s, + end_s=start_s + 1.0, + frame_times=(start_s + 1.0,), + inputs=UserInputs(), + ) + + +def _window_with_frame_times( + *, + start_s: float, + frame_times: Sequence[float], +) -> UserInputWindow: + return UserInputWindow( + start_s=start_s, + end_s=start_s + len(frame_times) * 0.5, + frame_times=frame_times, + inputs=UserInputs(), + ) + + +def _spec() -> DemoSpec: + return DemoSpec( + model_id="fake-video-demo", + input_mode="replay", + output=NullOutputSpec(), + config=InferenceConfig(model_id="fake-video-demo"), + ) + + +def _scenario() -> PreparedScenario: + return PreparedScenario(initial_inputs=InferenceInput()) + + +def _run_context( + runtime: _FakeVideoRuntime, + *, + admission: SingleSessionAdmissionPolicy | None = None, + run_metrics: InMemorySessionMetricsRecorder | None = None, +) -> RunContext: + host = RuntimeHost(runtime) + return RunContext( + host=host, + run_metrics=run_metrics or InMemorySessionMetricsRecorder(), + admission=admission + or SingleSessionAdmissionPolicy(health_check=lambda: host.is_healthy), + ) + + +def _record_output_factory_call( + calls: list[tuple[DemoSpec, PreparedScenario]], + spec: DemoSpec, + scenario: PreparedScenario, + output: "_RecordingOutputSink", +) -> "_RecordingOutputSink": + calls.append((spec, scenario)) + return output + + +class _FakeVideoModelInputProvider: + capabilities = ProviderCapabilities( + supports_recorded_input=True, + supports_reset=True, + deterministic_given_inputs=True, + ) + + def __init__( + self, + *, + fail_initial: Exception | None = None, + fail_close: Exception | None = None, + ) -> None: + self.fail_initial = fail_initial + self.fail_close = fail_close + self.initial_input = InferenceInput( + global_conditioning={"prompt": "fake video prompt"} + ) + self.prepared_step_inputs: list[InferenceInput] = [] + self.reset_inputs: list[InferenceInput | None] = [] + self.close_count = 0 + + def prepare_initial_input(self) -> InferenceInput: + if self.fail_initial is not None: + raise self.fail_initial + return self.initial_input + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + inference_input = InferenceInput( + step={ + "request_step": request.step_index, + "window": (user_window.start_s, user_window.end_s), + } + ) + self.prepared_step_inputs.append(inference_input) + return PreparedStep(inference_input=inference_input) + + def reset(self, inputs: InferenceInput | None = None) -> None: + self.reset_inputs.append(inputs) + + def close(self) -> None: + self.close_count += 1 + if self.fail_close is not None: + raise self.fail_close + + +class _BlockingCloseVideoModelInputProvider(_FakeVideoModelInputProvider): + def __init__( + self, + *, + close_entered: threading.Event, + release_close: threading.Event, + ) -> None: + super().__init__() + self.close_entered = close_entered + self.release_close = release_close + + def close(self) -> None: + self.close_count += 1 + self.close_entered.set() + assert self.release_close.wait(timeout=2.0) + + +class _FakeBatchInputSource: + is_finite = True + is_deterministic = True + user_input_schema = UserInputSchema() + + def __init__( + self, + *, + num_windows: int, + fail_is_finished: Exception | None = None, + ) -> None: + self.windows = [_window(index) for index in range(num_windows)] + self.fail_is_finished = fail_is_finished + self.next_window_requests: list[StepRequirements] = [] + self.index = 0 + + def is_finished(self) -> bool: + if self.fail_is_finished is not None: + raise self.fail_is_finished + return self.index >= len(self.windows) + + def next_window(self, request: StepRequirements) -> UserInputWindow: + self.next_window_requests.append(request) + window = self.windows[self.index] + self.index += 1 + return window + + +class _SlicingBatchInputSource: + is_finite = True + is_deterministic = True + user_input_schema = UserInputSchema() + + def __init__(self, *, fps: float, num_windows: int) -> None: + self.fps = fps + self.num_windows = num_windows + self.next_window_requests: list[StepRequirements] = [] + self.windows: list[UserInputWindow] = [] + self.window_index = 0 + self.next_frame_index = 0 + + def is_finished(self) -> bool: + return self.window_index >= self.num_windows + + def next_window(self, request: StepRequirements) -> UserInputWindow: + self.next_window_requests.append(request) + start_frame = self.next_frame_index + self.next_frame_index += request.input_frame_count + self.window_index += 1 + frame_times = tuple( + frame_index / self.fps + for frame_index in range(start_frame, self.next_frame_index) + ) + window = _window_with_frame_times( + start_s=start_frame / self.fps, + frame_times=frame_times, + ) + self.windows.append(window) + return window + + +class _FakeVideoRuntime: + def __init__(self, *, session: "_FakeVideoSession") -> None: + self.session = session + self.start_session_inputs: list[InferenceInput] = [] + self.close_count = 0 + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + self.start_session_inputs.append(inputs) + return self.session + + def close(self) -> None: + self.close_count += 1 + + +class _FakeVideoSession: + def __init__( + self, + *, + num_steps: int, + input_frame_counts: Sequence[int] | None = None, + fail_step: int | None = None, + ) -> None: + self.num_steps = num_steps + self.input_frame_counts = tuple(input_frame_counts or (1,) * num_steps) + self.fail_step = fail_step + self.next_request_index = 0 + self.step_inputs: list[InferenceInput] = [] + self.close_count = 0 + + def session_info(self) -> SessionInfo: + return SessionInfo(output_layout="fake-video", steady_output_frame_count=1) + + def next_step_requirements(self) -> StepRequirements | None: + if self.next_request_index >= self.num_steps: + return None + request = StepRequirements( + step_index=self.next_request_index, + input_frame_count=self.input_frame_counts[self.next_request_index], + ) + self.next_request_index += 1 + return request + + def next_step_request(self) -> StepRequest | None: + raise AssertionError("demo driver should request StepRequirements") + + def step(self, inputs: InferenceInput) -> StepResult: + step_index = len(self.step_inputs) + if self.fail_step == step_index: + raise RuntimeError("step failed") + self.step_inputs.append(inputs) + return StepResult( + step_index=step_index, + output=f"frame-{step_index}", + frame_count=1, + metrics={"model_step_s": 0.01}, + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self.next_request_index = 0 + self.step_inputs.clear() + + def close(self) -> None: + self.close_count += 1 + + +class _RecordingRuntimeHost(RuntimeHost): + def __init__(self, runtime: _FakeVideoRuntime) -> None: + super().__init__(runtime) + self.calls: list[str] = [] + + def call(self, func: Callable[..., Any], /, *args: object, **kwargs: object) -> Any: + self.calls.append(getattr(func, "__name__", type(func).__name__)) + return super().call(func, *args, **kwargs) + + +class _ClosingAfterStepRuntimeHost(_RecordingRuntimeHost): + def call(self, func: Callable[..., Any], /, *args: object, **kwargs: object) -> Any: + result = super().call(func, *args, **kwargs) + if getattr(func, "__name__", type(func).__name__) == "execute_step": + self.close() + return result + + +class _RecordingOutputSink: + produces_artifacts = True + + def __init__( + self, + *, + artifacts: Sequence[OutputArtifact] = (), + decision: OutputDecision | None = None, + fail_open: Exception | None = None, + fail_close: Exception | None = None, + ) -> None: + self.artifacts = tuple(artifacts) + self.decision = decision or OutputDecision() + self.fail_open = fail_open + self.fail_close = fail_close + self.opened_with: SessionInfo | None = None + self.results: list[StepResult] = [] + self.close_count = 0 + + def open(self, session_info: SessionInfo) -> None: + if self.fail_open is not None: + raise self.fail_open + self.opened_with = session_info + + def begin_generation(self, generation: int) -> None: + del generation + + def write(self, result: StepResult) -> OutputDecision: + self.results.append(result) + return self.decision + + def close(self) -> Sequence[OutputArtifact]: + self.close_count += 1 + if self.fail_close is not None: + raise self.fail_close + return self.artifacts + + +class _RecordingTransport: + def __init__(self) -> None: + self.close_count = 0 + + def is_active(self) -> bool: + return self.close_count == 0 + + def close(self) -> None: + self.close_count += 1 + + +class _FailingCleanupMetrics(InMemorySessionMetricsRecorder): + cleanup_error_attempts: int + + def __init__(self) -> None: + super().__init__() + self.cleanup_error_attempts = 0 + + def record_cleanup_error(self, exc: Exception) -> None: + del exc + self.cleanup_error_attempts += 1 + raise RuntimeError("cleanup metrics failed") + + +class _SetupPolicy: + def __init__( + self, + *, + result_status: Literal["completed", "failed", "skipped"], + ) -> None: + self.result_status = result_status + + def handle_setup_error(self, exc: Exception) -> ErrorAction: + del exc + return ErrorAction(result_status=self.result_status) + + def handle(self, exc: Exception) -> ErrorAction: + del exc + return ErrorAction(result_status="failed") + + +class _FakeDemoAdapter: + model_id = "fake-video-demo" + inference_input_schema = InferenceInputSchema() + canonical_input_schema = CanonicalInputSchema() + + def __init__(self, *, provider: _FakeVideoModelInputProvider) -> None: + self.provider = provider + self.provider_calls: list[tuple[DemoSpec, PreparedScenario]] = [] + + def supported_input_modes(self) -> tuple[str, ...]: + return ("replay",) + + def supported_output_modes(self) -> tuple[str, ...]: + return ("null",) + + def default_input_mapping(self) -> InputMapping: + return IdentityInputMapping() + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.model_id: + raise ValueError(f"Unsupported model_id={config.model_id!r}.") + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + del config + raise NotImplementedError("FakeVideoDemo uses an explicit RuntimeHost.") + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + del spec + return _scenario() + + def create_model_input_provider( + self, + spec: DemoSpec, + scenario: PreparedScenario, + ) -> _FakeVideoModelInputProvider: + self.provider_calls.append((spec, scenario)) + return self.provider + + +class _FakeRunMode: + name = "fake" + + def __init__( + self, + *, + input_source: _FakeBatchInputSource, + output_sink: _RecordingOutputSink | None = None, + output_sink_factory: ( + Callable[[DemoSpec, PreparedScenario], _RecordingOutputSink] | None + ) = None, + metrics: InMemorySessionMetricsRecorder | None = None, + transport: _RecordingTransport | None = None, + error_policy: _SetupPolicy | None = None, + validate_error: Exception | None = None, + select_error: Exception | None = None, + ) -> None: + self.input_source = input_source + self.output_sink = output_sink or _RecordingOutputSink() + self.output_sink_factory = output_sink_factory + self.metrics = metrics or InMemorySessionMetricsRecorder() + self.transport = transport + self.error_policy = error_policy + self.validate_error = validate_error + self.select_error = select_error + self.capabilities = RunModeCapabilities( + requires_finite_input=True, + supports_artifacts=True, + ) + + def validate_run( + self, + *, + spec: DemoSpec, + adapter: Any, + ) -> None: + del spec, adapter + + def create_run_context( + self, + *, + spec: DemoSpec, + adapter: Any, + host: RuntimeHost, + model_warmup_plan: ModelWarmupPlan, + ) -> RunContext: + del spec, adapter + return RunContext( + host=host, + run_metrics=InMemorySessionMetricsRecorder(), + admission=SingleSessionAdmissionPolicy( + health_check=lambda: host.is_healthy + ), + model_warmup_plan=model_warmup_plan, + ) + + def validate_session( + self, + *, + spec: DemoSpec, + scenario: PreparedScenario, + adapter: Any, + provider: Any, + ) -> None: + del spec, scenario, adapter, provider + if self.validate_error is not None: + raise self.validate_error + + def create_session_edges( + self, + *, + context: RunContext, + spec: DemoSpec, + scenario: PreparedScenario, + provider: Any, + adapter: Any, + ) -> SessionEdges: + del provider, adapter + output_sink = ( + self.output_sink_factory(spec, scenario) + if self.output_sink_factory is not None + else self.output_sink + ) + return SessionEdges( + input_source=self.input_source, + output_sink=output_sink, + cleanup_tasks=context.cleanup_tasks, + metrics=self.metrics, + error_policy=self.error_policy or _SetupPolicy(result_status="failed"), + transport=self.transport or _RecordingTransport(), + ) + + def select_driver(self) -> BatchSessionDriver: + if self.select_error is not None: + raise self.select_error + return BatchSessionDriver() diff --git a/flashdreams/tests/test_demo_runtime_warmup.py b/flashdreams/tests/test_demo_runtime_warmup.py new file mode 100644 index 000000000..4c3951ed0 --- /dev/null +++ b/flashdreams/tests/test_demo_runtime_warmup.py @@ -0,0 +1,416 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import threading +from typing import Any + +import pytest + +from flashdreams.runtime import ( + CanonicalInputSchema, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InferenceRuntime, + InferenceSession, + InputMapping, + StepRequest, + StepRequirements, + StepResult, +) +from flashdreams.runtime.demo import ( + DemoSpec, + InMemorySessionMetricsRecorder, + ModelWarmupPlan, + NullOutputSpec, + PreparedScenario, + PreparedStep, + ProviderCapabilities, + RunContext, + RuntimeHost, + UserInputWindow, + WarmupSessionInputs, + build_model_warmup_plan, + warmup_run_context, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_model_warmup_plan_uses_temporary_provider_on_worker_thread() -> None: + setup_thread_id = threading.get_ident() + runtime = _WarmupRuntime() + host = RuntimeHost(runtime) + adapter = _WarmupAdapter(warmup_steps=2) + spec = _spec() + scenario = adapter.prepare_scenario(spec) + + try: + plan = build_model_warmup_plan( + host=host, + adapter=adapter, + spec=spec, + scenario=scenario, + ) + real_provider = host.call(adapter.create_model_input_provider, spec, scenario) + finally: + host.close() + + warmup_provider = adapter.warmup_providers[0] + assert adapter.warmup_thread_id == host.worker.worker_thread_id + assert adapter.warmup_thread_id != setup_thread_id + assert warmup_provider is not real_provider + assert warmup_provider.close_count == 1 + assert real_provider.close_count == 0 + assert plan == ModelWarmupPlan( + sessions=( + WarmupSessionInputs( + initial_input=InferenceInput( + global_conditioning={"provider": "warmup"} + ), + step_inputs=( + InferenceInput(step={"provider": "warmup", "step": 0}), + InferenceInput(step={"provider": "warmup", "step": 1}), + ), + ), + ), + ) + + +def test_runtime_host_warmup_uses_runtime_session_api() -> None: + runtime = _WarmupRuntime() + host = RuntimeHost(runtime) + adapter = _WarmupAdapter(warmup_steps=2) + spec = _spec() + scenario = adapter.prepare_scenario(spec) + + try: + plan = build_model_warmup_plan( + host=host, + adapter=adapter, + spec=spec, + scenario=scenario, + ) + host.warmup(plan) + finally: + host.close() + + assert runtime.events[:4] == [ + ( + "start_session", + InferenceInput(global_conditioning={"provider": "warmup"}), + ), + ("step", InferenceInput(step={"provider": "warmup", "step": 0})), + ("step", InferenceInput(step={"provider": "warmup", "step": 1})), + "session.close", + ] + + +def test_run_mode_warmup_context_warms_transport_without_model_session() -> None: + runtime = _WarmupRuntime() + host = RuntimeHost(runtime) + adapter = _WarmupAdapter(warmup_steps=0) + spec = _spec() + scenario = adapter.prepare_scenario(spec) + transport = _TransportWarmupService() + mode = _TransportWarmupRunMode() + context = RunContext( + host=host, + run_metrics=InMemorySessionMetricsRecorder(), + admission=_Admission(), + model_warmup_plan=ModelWarmupPlan(), + services={"transport": transport}, + ) + + try: + warmup_run_context( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=mode, + ) + + assert runtime.events == [] + assert mode.warmup_calls == 1 + assert transport.warmup_calls == 1 + finally: + host.close() + + +def test_model_warmup_is_excluded_from_run_metrics() -> None: + runtime = _WarmupRuntime() + host = RuntimeHost(runtime) + adapter = _WarmupAdapter(warmup_steps=1) + spec = _spec() + scenario = adapter.prepare_scenario(spec) + metrics = InMemorySessionMetricsRecorder() + + try: + plan = build_model_warmup_plan( + host=host, + adapter=adapter, + spec=spec, + scenario=scenario, + ) + context = RunContext( + host=host, + run_metrics=metrics, + admission=_Admission(), + model_warmup_plan=plan, + ) + + warmup_run_context( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=object(), + ) + + assert runtime.events[:3] == [ + ( + "start_session", + InferenceInput(global_conditioning={"provider": "warmup"}), + ), + ("step", InferenceInput(step={"provider": "warmup", "step": 0})), + "session.close", + ] + assert metrics.sessions == [] + assert metrics.step_count == 0 + assert metrics.control_count == 0 + finally: + host.close() + + +def test_adapter_without_model_warmup_hook_gets_empty_plan() -> None: + host = RuntimeHost(_WarmupRuntime()) + adapter = _NoWarmupAdapter() + spec = _spec() + scenario = adapter.prepare_scenario(spec) + + try: + plan = build_model_warmup_plan( + host=host, + adapter=adapter, + spec=spec, + scenario=scenario, + ) + finally: + host.close() + + assert plan == ModelWarmupPlan() + + +def _spec() -> DemoSpec: + return DemoSpec( + model_id="fake-demo", + input_mode="replay", + output=NullOutputSpec(), + ) + + +def _scenario() -> PreparedScenario: + return PreparedScenario(initial_inputs=InferenceInput()) + + +class _WarmupAdapter: + model_id = "fake-demo" + inference_input_schema = InferenceInputSchema() + canonical_input_schema = CanonicalInputSchema() + + def __init__(self, *, warmup_steps: int) -> None: + self.warmup_steps = warmup_steps + self.warmup_thread_id: int | None = None + self.warmup_providers: list[_WarmupProvider] = [] + self.real_providers: list[_WarmupProvider] = [] + + def supported_input_modes(self) -> tuple[str, ...]: + return ("replay",) + + def supported_output_modes(self) -> tuple[str, ...]: + return ("null",) + + def default_input_mapping(self) -> InputMapping: + return IdentityInputMapping() + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.model_id: + raise ValueError(f"Unsupported model_id={config.model_id!r}.") + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + return _WarmupRuntime() + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + del spec + return _scenario() + + def create_model_warmup_sessions( + self, + spec: DemoSpec, + scenario: PreparedScenario, + ) -> tuple[WarmupSessionInputs, ...]: + del spec, scenario + self.warmup_thread_id = threading.get_ident() + provider = _WarmupProvider(name="warmup") + self.warmup_providers.append(provider) + try: + initial_input = provider.prepare_initial_input() + step_inputs = [] + for step_index in range(self.warmup_steps): + prepared = provider.prepare_step( + request=StepRequirements(step_index=step_index), + user_window=UserInputWindow( + start_s=float(step_index), + end_s=float(step_index + 1), + ), + ) + if prepared.inference_input is None: + raise RuntimeError("Warmup provider returned no step input.") + step_inputs.append(prepared.inference_input) + return ( + WarmupSessionInputs( + initial_input=initial_input, + step_inputs=tuple(step_inputs), + ), + ) + finally: + provider.close() + + def create_model_input_provider( + self, + spec: DemoSpec, + scenario: PreparedScenario, + ) -> "_WarmupProvider": + del spec, scenario + provider = _WarmupProvider(name="real") + self.real_providers.append(provider) + return provider + + +class _NoWarmupAdapter: + model_id = "fake-demo" + inference_input_schema = InferenceInputSchema() + canonical_input_schema = CanonicalInputSchema() + + def supported_input_modes(self) -> tuple[str, ...]: + return ("replay",) + + def supported_output_modes(self) -> tuple[str, ...]: + return ("null",) + + def default_input_mapping(self) -> InputMapping: + return IdentityInputMapping() + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.model_id: + raise ValueError(f"Unsupported model_id={config.model_id!r}.") + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + return _WarmupRuntime() + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + del spec + return _scenario() + + +class _WarmupProvider: + capabilities = ProviderCapabilities(supports_recorded_input=True) + + def __init__(self, *, name: str) -> None: + self.name = name + self.close_count = 0 + + def prepare_initial_input(self) -> InferenceInput: + return InferenceInput(global_conditioning={"provider": self.name}) + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + del user_window + return PreparedStep( + inference_input=InferenceInput( + step={"provider": self.name, "step": request.step_index} + ) + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + + def close(self) -> None: + self.close_count += 1 + + +class _WarmupRuntime: + def __init__(self) -> None: + self.events: list[object] = [] + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + self.events.append(("start_session", inputs)) + return _WarmupSession(events=self.events) + + def close(self) -> None: + self.events.append("runtime.close") + + +class _WarmupSession: + def __init__(self, *, events: list[object]) -> None: + self.events = events + self.next_step = 0 + + def next_step_request(self) -> StepRequest | None: + request = StepRequest(step_index=self.next_step) + self.next_step += 1 + return request + + def step(self, inputs: InferenceInput) -> StepResult: + self.events.append(("step", inputs)) + return StepResult(step_index=self.next_step, output=None) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self.next_step = 0 + + def close(self) -> None: + self.events.append("session.close") + + +class _TransportWarmupService: + def __init__(self) -> None: + self.warmup_calls = 0 + + def warmup(self) -> None: + self.warmup_calls += 1 + + +class _TransportWarmupRunMode: + def __init__(self) -> None: + self.warmup_calls = 0 + + def warmup_context( + self, + *, + context: RunContext, + spec: DemoSpec, + scenario: PreparedScenario, + adapter: Any, + ) -> None: + del spec, scenario, adapter + transport = context.services["transport"] + if not isinstance(transport, _TransportWarmupService): + raise TypeError("Expected fake transport warmup service.") + transport.warmup() + self.warmup_calls += 1 + + +class _Admission: + def try_reserve(self) -> None: + return None diff --git a/flashdreams/tests/test_encoders.py b/flashdreams/tests/test_encoders.py index 610217807..06fa38e85 100644 --- a/flashdreams/tests/test_encoders.py +++ b/flashdreams/tests/test_encoders.py @@ -30,7 +30,7 @@ import asyncio import sys import threading -from collections.abc import Callable +from collections.abc import Callable, Sequence from fractions import Fraction from types import ModuleType, SimpleNamespace from unittest.mock import MagicMock, patch @@ -371,10 +371,18 @@ class _FakeBufferedVideoTrack: def __init__(self) -> None: self.enqueued_results: list[StepResult] = [] + self.enqueued_frames: list[object] = [] - async def enqueue_result(self, result: StepResult) -> int: + def prepare_result_frames(self, result: StepResult) -> tuple[object, ...]: self.enqueued_results.append(result) - return result.frame_count + return tuple(object() for _ in range(result.frame_count)) + + async def enqueue_frames(self, frames: Sequence[object]) -> int: + self.enqueued_frames.extend(frames) + return len(frames) + + async def enqueue_result(self, result: StepResult) -> int: + return await self.enqueue_frames(self.prepare_result_frames(result)) class TestDefaultRTCEncoderDeliver: @@ -405,6 +413,7 @@ async def test_deliver_chunk_returns_frames_from_track( assert result.num_frames == 4 assert result.num_keyframes == 0 assert fake_track.enqueued_results == [step_result] + assert len(fake_track.enqueued_frames) == 4 @pytest.mark.parametrize( ("layout", "shape"), @@ -429,7 +438,7 @@ async def test_software_conversion_uses_declared_layout( await track.close() @pytest.mark.asyncio - async def test_software_path_defers_host_conversion_to_track(self) -> None: + async def test_software_path_prepares_host_frames_with_track(self) -> None: from flashdreams.serving.webrtc.media import BufferedVideoTrack source = torch.zeros((2, 3, 2, 2), dtype=torch.uint8) @@ -447,7 +456,9 @@ def _converter(delivered: StepResult) -> list[np.ndarray]: 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) + encoder = DefaultRTCEncoder(fps=30) + payload = encoder.prepare_chunk_payload(step_result, track) + delivery = await encoder.deliver_prepared_chunk(payload, track) assert delivery.num_frames == 2 assert seen == [step_result] diff --git a/flashdreams/tests/test_inference_runtime_api.py b/flashdreams/tests/test_inference_runtime_api.py index 42f75d688..531f2f218 100644 --- a/flashdreams/tests/test_inference_runtime_api.py +++ b/flashdreams/tests/test_inference_runtime_api.py @@ -4,6 +4,7 @@ from __future__ import annotations from dataclasses import fields +from types import SimpleNamespace from typing import Any, cast import pytest @@ -16,15 +17,19 @@ InferenceInputSchema, InMemoryMetricsRecorder, InputField, + MetricsSnapshot, + NullMetricsRecorder, NullOutputTarget, OutputArtifact, RuntimeMetricSample, StepRequest, + StepRequirements, StepResult, TimeWindow, UserInputEvent, UserInputs, UserInputSchema, + step_requirements_from_request, ) pytestmark = pytest.mark.ci_cpu @@ -38,11 +43,13 @@ def test_inference_config_keeps_runtime_settings_separate() -> None: backend="local", precision="bf16", compile=False, + seed=123, runtime_options={"chunk_size": 3}, ) assert config.model_id == "lingbot-world" assert config.preset_id == "fast-taehv" + assert config.seed == 123 assert config.runtime_options["chunk_size"] == 3 assert denied_app_fields.isdisjoint(field.name for field in fields(InferenceConfig)) with pytest.raises(TypeError): @@ -54,6 +61,13 @@ def test_inference_config_rejects_empty_model_id() -> None: InferenceConfig(model_id=" ") +def test_inference_config_rejects_invalid_seed() -> None: + with pytest.raises(TypeError, match="seed"): + InferenceConfig(model_id="fake", seed=True) + with pytest.raises(ValueError, match="seed"): + InferenceConfig(model_id="fake", seed=-1) + + @pytest.mark.parametrize( ("factory", "match"), [ @@ -67,6 +81,12 @@ def test_inference_config_rejects_empty_model_id() -> None: ), (lambda: UserInputEvent(timestamp_s=0.0, event_type=" "), "event_type"), (lambda: StepRequest(step_index=-1), "step_index"), + (lambda: StepRequirements(step_index=-1), "step_index"), + (lambda: StepRequirements(step_index=0, input_frame_count=0), "input_frame"), + ( + lambda: StepRequirements(step_index=0, steady_output_frame_count=-1), + "steady_output", + ), (lambda: StepResult(step_index=-1), "step_index"), (lambda: StepResult(step_index=0, frame_count=-1), "frame_count"), (lambda: RuntimeMetricSample(name=" ", value=1.0), "name"), @@ -185,6 +205,66 @@ def test_identity_input_mapping_leaves_inference_input_unchanged() -> None: ) +def test_step_requirements_adapt_legacy_request_metadata() -> None: + schema = InferenceInputSchema(step_fields=(InputField(name="camera_poses"),)) + request = StepRequest( + step_index=3, + inference_input_schema=schema, + metadata={ + "input_frame_count": 4, + "steady_output_frame_count": 2, + "model": "fake-video-demo", + }, + ) + + requirements = step_requirements_from_request(request) + + assert requirements == StepRequirements( + step_index=3, + input_frame_count=4, + steady_output_frame_count=2, + inference_input_schema=schema, + metadata={"model": "fake-video-demo"}, + ) + with pytest.raises(TypeError): + cast(Any, requirements.metadata)["model"] = "changed" + + +def test_step_requirements_keep_user_inputs_driver_owned() -> None: + requirements = StepRequirements(step_index=0, metadata={"model": "fake"}) + + assert not hasattr(requirements, "user_input_window") + with pytest.raises(ValueError, match="driver-owned user input"): + StepRequirements(step_index=0, metadata={"user_inputs": UserInputs()}) + with pytest.raises(ValueError, match="driver-owned"): + step_requirements_from_request( + StepRequest( + step_index=0, + user_input_window=TimeWindow(start_s=0.0, end_s=1.0), + ) + ) + + +def test_step_requirements_can_drop_legacy_user_window_when_source_owns_it() -> None: + request = StepRequest( + step_index=2, + user_input_window=TimeWindow(start_s=1.0, end_s=2.0), + metadata={"input_frame_count": 3, "model": "fake"}, + ) + + requirements = step_requirements_from_request( + request, + allow_user_input_window=True, + ) + + assert requirements == StepRequirements( + step_index=2, + input_frame_count=3, + metadata={"model": "fake"}, + ) + assert not hasattr(requirements, "user_input_window") + + def test_null_output_target_counts_and_optionally_stores_results() -> None: target = NullOutputTarget(store_results=True) result = StepResult(step_index=0, output=b"frame") @@ -233,6 +313,80 @@ def test_in_memory_metrics_recorder_uses_seconds_for_timing() -> None: assert sample.unit == "s" assert sample.category == "timing" assert sample.step_index == 2 + snapshot = recorder.close() + assert isinstance(snapshot, MetricsSnapshot) + assert recorder.closed + assert snapshot.counters["samples"] == 1 + assert snapshot.timings["model_step"] == (pytest.approx(0.125),) + + +def test_in_memory_metrics_recorder_rolls_up_sessions_and_diagnostics() -> None: + recorder = InMemoryMetricsRecorder() + + recorder.record_session(SimpleNamespace(status="completed")) + recorder.record_session_error(RuntimeError("assembly failed")) + recorder.record_error(RuntimeError("step failed"), object()) + recorder.record_catch_up(object()) + recorder.record_cleanup_error(RuntimeError("cleanup failed")) + recorder.record_orphaned_cleanup(RuntimeError("orphaned cleanup")) + snapshot = recorder.close() + + assert snapshot.counters["sessions"] == 1 + assert snapshot.counters["sessions.completed"] == 1 + assert snapshot.counters.get("sessions.failed", 0) == 0 + assert snapshot.counters["session_errors"] == 1 + assert snapshot.counters["catch_ups"] == 1 + assert snapshot.session_statuses == ("completed",) + assert snapshot.errors == ( + "step failed", + "cleanup failed", + "orphaned cleanup", + "assembly failed", + ) + + +def test_cancelled_session_rollup_does_not_count_as_failed() -> None: + recorder = InMemoryMetricsRecorder() + + recorder.record_session(SimpleNamespace(status="cancelled")) + snapshot = recorder.close() + + assert snapshot.counters["sessions"] == 1 + assert snapshot.counters["sessions.cancelled"] == 1 + assert snapshot.counters.get("sessions.failed", 0) == 0 + assert snapshot.session_statuses == ("cancelled",) + + +def test_null_metrics_recorder_keeps_old_and_new_calls_noop() -> None: + recorder = NullMetricsRecorder() + + recorder.record(RuntimeMetricSample(name="runtime", value=1.0)) + recorder.record_timing("model_step", 0.125, step_index=2) + recorder.record_step( + request=object(), + user_window=object(), + inference_input=object(), + result=object(), + decision=object(), + ) + recorder.record_control( + request=object(), + user_window=object(), + control=object(), + ) + recorder.record_error(RuntimeError("step failed"), object()) + recorder.record_catch_up(object()) + recorder.record_cleanup_error(RuntimeError("cleanup failed")) + recorder.record_orphaned_cleanup(RuntimeError("orphaned cleanup")) + recorder.record_session(SimpleNamespace(status="failed")) + recorder.record_session_error(RuntimeError("assembly failed")) + snapshot = recorder.close() + + assert isinstance(snapshot, MetricsSnapshot) + assert snapshot.counters == {} + assert snapshot.timings == {} + assert snapshot.session_statuses == () + assert snapshot.errors == () def test_timing_metric_samples_must_use_seconds() -> None: diff --git a/flashdreams/tests/test_realtime_timing_metrics.py b/flashdreams/tests/test_realtime_timing_metrics.py new file mode 100644 index 000000000..db157f195 --- /dev/null +++ b/flashdreams/tests/test_realtime_timing_metrics.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from flashdreams.runtime import InMemoryMetricsRecorder +from flashdreams.serving.realtime.timing import ( + ChunkTimes, + VideoModelTimings, + record_chunk_timing_metrics, + record_video_model_timing_metrics, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_chunk_timing_records_feed_session_metrics() -> None: + chunk = ChunkTimes.create( + chunk_index=2, + input_sample_time=0.0, + request_time=0.010, + request_poses_ready_time=0.030, + intended_present_times=[0.100], + ) + chunk.chunk_render_start_time = 0.050 + chunk.chunk_ready_time = 0.090 + chunk.frames[0].image_ready_time = 0.110 + chunk.frames[0].present_time = 0.140 + metrics = InMemoryMetricsRecorder() + + record_chunk_timing_metrics(metrics, chunk) + snapshot = metrics.close() + + assert snapshot.timings["realtime.chunk.input_to_request"] == ( + pytest.approx(0.010), + ) + assert snapshot.timings["realtime.chunk.request_to_poses_ready"] == ( + pytest.approx(0.020), + ) + assert snapshot.timings["realtime.chunk.queue_wait"] == (pytest.approx(0.020),) + assert snapshot.timings["realtime.chunk.chunk_render"] == (pytest.approx(0.040),) + assert metrics.samples[0].step_index == 2 + + +def test_video_model_timing_records_feed_session_metrics() -> None: + timings = VideoModelTimings( + condition_start_time=1.0, + condition_ready_time=1.010, + model_start_time=1.020, + model_ready_time=1.070, + cache_update_start_time=1.075, + cache_update_ready_time=1.080, + decode_start_time=1.085, + decode_ready_time=1.095, + merge_start_time=1.100, + merge_ready_time=1.115, + ) + metrics = InMemoryMetricsRecorder() + + record_video_model_timing_metrics(metrics, timings, chunk_index=3) + snapshot = metrics.close() + + assert snapshot.timings["realtime.model.condition"] == (pytest.approx(0.010),) + assert snapshot.timings["realtime.model.model"] == (pytest.approx(0.050),) + assert snapshot.timings["realtime.model.cache_update"] == (pytest.approx(0.005),) + assert snapshot.timings["realtime.model.decode"] == (pytest.approx(0.010),) + assert snapshot.timings["realtime.model.merge"] == (pytest.approx(0.015),) + assert snapshot.timings["realtime.model.total"] == (pytest.approx(0.115),) + assert metrics.samples[0].step_index == 3 diff --git a/flashdreams/tests/test_runtime_demo_api.py b/flashdreams/tests/test_runtime_demo_api.py index 6cdf772e5..e40b62b51 100644 --- a/flashdreams/tests/test_runtime_demo_api.py +++ b/flashdreams/tests/test_runtime_demo_api.py @@ -3,6 +3,7 @@ from __future__ import annotations +import argparse from collections.abc import Sequence from pathlib import Path from types import SimpleNamespace @@ -36,14 +37,21 @@ ) from flashdreams.runtime.demo import ( DemoSpec, + Mp4OutputSink, Mp4OutputSpec, + NullOutputSink, NullOutputSpec, + OutputSink, + OutputSpec, PreparedScenario, + RunResult, WebRTCAppResources, WebRTCOutputSpec, + build_output_sink, build_output_target, run_replay_demo, ) +from flashdreams.runtime.demo.app import DemoApplication from flashdreams.runtime.demo.webrtc import ( serve_webrtc_demo, ) @@ -52,7 +60,42 @@ pytestmark = pytest.mark.ci_cpu -def test_replay_demo_uses_shared_runner() -> None: +def test_replay_demo_uses_shared_batch_path_by_default() -> None: + adapter = _FakeDemoAdapter() + sinks: list[OutputSink] = [] + + def output_sink_factory(output_spec: OutputSpec) -> OutputSink: + sink = build_output_sink(output_spec) + sinks.append(sink) + return sink + + spec = DemoSpec( + model_id="fake-demo", + scenario="valid-scenario", + input_mode="replay", + output=NullOutputSpec(), + ) + + result = run_replay_demo( + spec=spec, + adapter=adapter, + metrics=NullMetricsRecorder(), + output_sink_factory=output_sink_factory, + ) + + assert result.status == "completed" + assert result.artifacts == () + assert len(sinks) == 1 + assert isinstance(sinks[0], NullOutputSink) + assert adapter.create_runtime_called + assert adapter.runtime is not None + assert adapter.runtime.closed + assert adapter.runtime.session is not None + assert adapter.runtime.session.closed + assert adapter.prepare_scenario_calls == [spec] + + +def test_replay_demo_keeps_compat_runner_injection() -> None: adapter = _FakeDemoAdapter() output = _RecordingOutputTarget() calls: list[dict[str, Any]] = [] @@ -68,7 +111,7 @@ def fake_runner(**kwargs: Any) -> Sequence[OutputArtifact]: output=NullOutputSpec(), ) - artifacts = run_replay_demo( + result = run_replay_demo( spec=spec, adapter=adapter, output_target_factory=lambda output_spec: output, @@ -76,7 +119,10 @@ def fake_runner(**kwargs: Any) -> Sequence[OutputArtifact]: runner=fake_runner, ) - assert artifacts == (OutputArtifact(kind="test/artifact", uri="memory://artifact"),) + assert result == RunResult( + status="completed", + artifacts=(OutputArtifact(kind="test/artifact", uri="memory://artifact"),), + ) assert len(calls) == 1 assert calls[0]["adapter"] is adapter assert calls[0]["config"] == spec.config @@ -90,8 +136,9 @@ def fake_runner(**kwargs: Any) -> Sequence[OutputArtifact]: assert not adapter.create_runtime_called -def test_replay_demo_builds_output_target_from_spec(tmp_path: Path) -> None: +def test_replay_demo_builds_output_sink_from_spec(tmp_path: Path) -> None: writer_calls: list[dict[str, Any]] = [] + sinks: list[OutputSink] = [] def fake_writer( video: torch.Tensor, @@ -119,18 +166,24 @@ def fake_writer( output=Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=12), ) - artifacts = run_replay_demo( + result = run_replay_demo( spec=spec, adapter=_FakeDemoAdapter(video_output=True), - output_target_factory=lambda output_spec: build_output_target( - output_spec, - mp4_writer=fake_writer, + output_sink_factory=lambda output_spec: _record_output_sink( + sinks, + build_output_sink( + output_spec, + mp4_writer=fake_writer, + ), ), ) - assert len(artifacts) == 1 - assert artifacts[0].kind == "video/mp4" - assert artifacts[0].uri == str(tmp_path / "demo.mp4") + assert result.status == "completed" + assert len(sinks) == 1 + assert isinstance(sinks[0], Mp4OutputSink) + assert len(result.artifacts) == 1 + assert result.artifacts[0].kind == "video/mp4" + assert result.artifacts[0].uri == str(tmp_path / "demo.mp4") assert writer_calls == [ { "shape": (2, 2, 2, 3), @@ -141,6 +194,63 @@ def fake_writer( ] +def test_replay_demo_mp4_sink_matches_legacy_output_target_payload( + tmp_path: Path, +) -> None: + def writer(records: list[dict[str, Any]]): + def fake_writer( + video: torch.Tensor, + path: Path, + *, + fps: int | float, + layout: str, + install_hint: str, + ) -> Path: + del install_hint + records.append( + { + "bytes": video.detach().cpu().numpy().tobytes(), + "shape": tuple(video.shape), + "path": path, + "fps": fps, + "layout": layout, + } + ) + return path + + return fake_writer + + spec = DemoSpec( + model_id="fake-demo", + scenario="valid-scenario", + input_mode="replay", + output=Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=12), + ) + sink_records: list[dict[str, Any]] = [] + target_records: list[dict[str, Any]] = [] + + sink_result = run_replay_demo( + spec=spec, + adapter=_FakeDemoAdapter(video_output=True), + output_sink_factory=lambda output_spec: build_output_sink( + output_spec, + mp4_writer=writer(sink_records), + ), + ) + target_result = run_replay_demo( + spec=spec, + adapter=_FakeDemoAdapter(video_output=True), + output_target_factory=lambda output_spec: build_output_target( + output_spec, + mp4_writer=writer(target_records), + ), + ) + + assert sink_result.status == "completed" + assert target_result.status == "completed" + assert sink_records == target_records + + def test_replay_demo_fails_before_runtime_creation_when_scenario_invalid() -> None: adapter = _FakeDemoAdapter(scenario_valid=False) output_factory_calls = 0 @@ -170,6 +280,18 @@ def output_factory(output_spec: object) -> OutputTarget: assert output_factory_calls == 0 +def test_replay_demo_step_failure_exits_nonzero_and_prints_reason( + capsys: pytest.CaptureFixture[str], +) -> None: + app = _ReplayOnlyDemoApplication(adapter=_FakeDemoAdapter(fail_step=0)) + + with pytest.raises(SystemExit) as raised: + app.main(["replay"]) + + assert raised.value.code == 1 + assert "step failed" in capsys.readouterr().err + + def test_demo_adapter_declares_supported_modes() -> None: adapter = _FakeDemoAdapter( input_modes=("replay",), @@ -291,6 +413,36 @@ def map_step_inputs( ) +def _record_output_sink(sinks: list[OutputSink], sink: OutputSink) -> OutputSink: + sinks.append(sink) + return sink + + +class _ReplayOnlyDemoApplication(DemoApplication): + def __init__(self, *, adapter: "_FakeDemoAdapter") -> None: + self._adapter = adapter + + def parse_args(self, argv: list[str] | None = None) -> argparse.Namespace: + del argv + return argparse.Namespace(command="replay") + + def replay_spec(self, args: argparse.Namespace) -> DemoSpec: + del args + return DemoSpec( + model_id="fake-demo", + scenario="valid-scenario", + input_mode="replay", + output=NullOutputSpec(), + ) + + def replay_adapter(self) -> "_FakeDemoAdapter": + return self._adapter + + def serve_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: + del args, context + raise AssertionError("webrtc should not run") + + class _FakeDemoAdapter: model_id = "fake-demo" inference_input_schema = InferenceInputSchema( @@ -304,11 +456,13 @@ def __init__( *, scenario_valid: bool = True, video_output: bool = False, + fail_step: int | None = None, input_modes: tuple[str, ...] = ("replay",), output_modes: tuple[str, ...] = ("null", "mp4"), ) -> None: self._scenario_valid = scenario_valid self._video_output = video_output + self._fail_step = fail_step self._input_modes = input_modes self._output_modes = output_modes self.mapping = _ChunkIndexMapping() @@ -344,6 +498,7 @@ def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: self.runtime = _FakeRuntime( inference_input_schema=self.inference_input_schema, video_output=self._video_output, + fail_step=self._fail_step, ) return self.runtime @@ -360,9 +515,11 @@ def __init__( *, inference_input_schema: InferenceInputSchema, video_output: bool, + fail_step: int | None, ) -> None: self._inference_input_schema = inference_input_schema self._video_output = video_output + self._fail_step = fail_step self.session: _FakeSession | None = None self.closed = False @@ -371,6 +528,7 @@ def start_session(self, inputs: InferenceInput) -> InferenceSession: self.session = _FakeSession( inference_input_schema=self._inference_input_schema, video_output=self._video_output, + fail_step=self._fail_step, ) return self.session @@ -384,9 +542,11 @@ def __init__( *, inference_input_schema: InferenceInputSchema, video_output: bool, + fail_step: int | None, ) -> None: self._inference_input_schema = inference_input_schema self._video_output = video_output + self._fail_step = fail_step self.step_index = 0 self.closed = False @@ -403,6 +563,8 @@ def next_step_request(self) -> StepRequest | None: def step(self, inputs: InferenceInput) -> StepResult: self._inference_input_schema.require_step(inputs) + if self._fail_step == self.step_index: + raise RuntimeError("step failed") if self._video_output: result = StepResult.from_video_chunk( step_index=self.step_index, diff --git a/flashdreams/tests/test_runtime_runner.py b/flashdreams/tests/test_runtime_runner.py index b755ff48a..7ba92e836 100644 --- a/flashdreams/tests/test_runtime_runner.py +++ b/flashdreams/tests/test_runtime_runner.py @@ -8,6 +8,7 @@ import pytest +import flashdreams.runtime.runner as runner_module from flashdreams.runtime import ( DRIVER_COMMAND, CanonicalInputs, @@ -25,6 +26,7 @@ InputField, InputMapping, InputMappingSchema, + MetricsSnapshot, NullOutputTarget, OutputArtifact, RuntimeMetricSample, @@ -73,6 +75,57 @@ def test_run_inference_session_completes_two_step_run() -> None: assert metrics.closed +def test_run_inference_session_delegates_to_shared_batch_helper( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[Mapping[str, object]] = [] + artifact = OutputArtifact(kind="test/artifact", uri="memory://artifact") + + def _fake_helper(**kwargs: object) -> tuple[OutputArtifact, ...]: + calls.append(kwargs) + return (artifact,) + + monkeypatch.setattr( + runner_module, + "_run_inference_session_with_shared_batch", + _fake_helper, + ) + adapter = _FakeAdapter() + config = InferenceConfig(model_id="fake-model") + mapping = _ChunkIndexMapping() + canonicalizer = InputCanonicalizer() + source_schema = UserInputSchema() + user_inputs = UserInputs() + initial_inputs = InferenceInput(global_conditioning={"prompt": "drive forward"}) + output = NullOutputTarget() + metrics = InMemoryMetricsRecorder() + + artifacts = runner_module.run_inference_session( + adapter=adapter, + config=config, + mapping=mapping, + canonicalizer=canonicalizer, + source_schema=source_schema, + user_inputs=user_inputs, + initial_inputs=initial_inputs, + output=output, + metrics=metrics, + ) + + assert artifacts == (artifact,) + assert len(calls) == 1 + call = calls[0] + assert call["adapter"] is adapter + assert call["config"] is config + assert call["mapping"] is mapping + assert call["canonicalizer"] is canonicalizer + assert call["source_schema"] is source_schema + assert call["user_inputs"] is user_inputs + assert call["initial_inputs"] is initial_inputs + assert call["output"] is output + assert call["metrics"] is metrics + + def test_runner_preserves_initial_step_inputs_for_identity_mapping() -> None: adapter = _FakeAdapter() @@ -656,5 +709,44 @@ def record_timing( ) ) - def close(self) -> None: + def record_step( + self, + *, + request: object, + user_window: object, + inference_input: object, + result: object, + decision: object, + ) -> None: + del request, user_window, inference_input, result, decision + + def record_control( + self, + *, + request: object, + user_window: object, + control: object, + ) -> None: + del request, user_window, control + + def record_error(self, exc: Exception, action: object) -> None: + del exc, action + + def record_catch_up(self, decision: object) -> None: + del decision + + def record_cleanup_error(self, exc: Exception) -> None: + del exc + + def record_orphaned_cleanup(self, exc: Exception) -> None: + del exc + + def record_session(self, result: object) -> None: + del result + + def record_session_error(self, exc: Exception) -> None: + del exc + + def close(self) -> MetricsSnapshot: self._events.append("metrics.close") + return MetricsSnapshot() diff --git a/flashdreams/tests/test_runtime_worker.py b/flashdreams/tests/test_runtime_worker.py index f6fbf84aa..d76558005 100644 --- a/flashdreams/tests/test_runtime_worker.py +++ b/flashdreams/tests/test_runtime_worker.py @@ -8,11 +8,15 @@ import pytest -from flashdreams.runtime import ThreadAffineRuntimeWorker +from flashdreams.runtime import ModelExecutionWorker, ThreadAffineRuntimeWorker pytestmark = pytest.mark.ci_cpu +def test_model_execution_worker_keeps_legacy_worker_alias() -> None: + assert ThreadAffineRuntimeWorker is ModelExecutionWorker + + @pytest.mark.asyncio async def test_worker_preserves_order_and_thread_affinity() -> None: worker = ThreadAffineRuntimeWorker(thread_name="test-runtime") @@ -95,3 +99,32 @@ async def test_worker_sets_cuda_device_when_thread_starts( await worker.close() assert [str(device) for device in seen] == ["cuda:3"] + + +def test_blocking_worker_call_is_not_reentrant() -> None: + worker = ModelExecutionWorker() + + def _nested_dispatch() -> None: + worker.call_blocking(lambda: None) + + try: + with pytest.raises(RuntimeError, match="own thread"): + worker.call_blocking(_nested_dispatch) + finally: + worker.close_blocking() + + +def test_async_worker_call_is_not_reentrant_from_worker_thread() -> None: + worker = ModelExecutionWorker() + + def _nested_async_dispatch() -> None: + async def _dispatch() -> None: + await worker.call(lambda: None) + + asyncio.run(_dispatch()) + + try: + with pytest.raises(RuntimeError, match="own thread"): + worker.call_blocking(_nested_async_dispatch) + finally: + worker.close_blocking() diff --git a/flashdreams/tests/test_webrtc_manager.py b/flashdreams/tests/test_webrtc_manager.py index 8e4839e5e..787a8ee7b 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -11,7 +11,16 @@ import pytest import torch -from flashdreams.runtime import StepRequest, StepResult +from flashdreams.runtime import ( + InferenceInput, + StepRequest, + StepRequirements, + StepResult, + UserInputEvent, + UserInputs, +) +from flashdreams.runtime.demo import RunResult +from flashdreams.runtime.demo.timing import SPARSE_KEY_SEGMENTS_METADATA_KEY 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 @@ -20,6 +29,12 @@ ManagedWebRTCSession, ) from flashdreams.serving.webrtc.server import SessionBusyError +from flashdreams.serving.webrtc.services import ( + WEBRTC_SKIPPED_INPUTS_METADATA_KEY, + WEBRTC_SKIPPED_WINDOW_METADATA_KEY, + WebRTCInputSource, + WebRTCTransportService, +) pytestmark = pytest.mark.ci_cpu @@ -83,6 +98,29 @@ async def deliver_chunk( encode_ms=0.1, ) + def prepare_chunk_payload( + self, + result: StepResult, + track: Any, + ) -> StepResult: + del track + return result + + async def deliver_prepared_chunk( + self, + payload: object, + track: Any, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + if not isinstance(payload, StepResult): + raise TypeError("fake payload must be StepResult") + return await self.deliver_chunk( + payload, + track, + force_keyframe=force_keyframe, + ) + def close(self) -> None: return @@ -122,6 +160,28 @@ def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: self.edges.append((arrival_t, event, key)) +class _SharedResampler: + def __init__(self, *, start_v: float = 0.0, dt: float = 0.001) -> None: + self.next_chunk_start_v = start_v + self.dt = dt + self.edges: list[tuple[float, str, str]] = [] + + def reset(self, *, start_v: float) -> None: + self.next_chunk_start_v = start_v + self.edges.clear() + + def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: + self.edges.append((arrival_t, event, key)) + + def sample_chunk( + self, num_frames: int + ) -> tuple[list[tuple[float, float, frozenset[str]]], list[float]]: + start = self.next_chunk_start_v + end = start + num_frames * self.dt + self.next_chunk_start_v = end + return [(start, end, frozenset({"w"}))], [end] + + class _CountingVideoTrack(_FakeVideoTrack): async def enqueue_result(self, result: StepResult) -> int: return result.frame_count @@ -423,6 +483,94 @@ def test_catch_up_input_clock_snaps_legacy_path_without_canonicalizer() -> None: assert managed.resampler.next_chunk_start_v == pytest.approx(2.0) +def test_legacy_provider_advances_skipped_webrtc_input_state() -> None: + class _RecordingCanonicalizer: + def __init__(self) -> None: + self.windows: list[tuple[float, float]] = [] + self.event_batches: list[list[str]] = [] + + def canonicalize( + self, + user_inputs: UserInputs, + *, + window: Any, + source_schema: Any, + ) -> object: + del source_schema + self.windows.append((window.start_s, window.end_s)) + self.event_batches.append( + [event.event_type for event in user_inputs.events] + ) + return object() + + class _RecordingMapping: + def __init__(self) -> None: + self.inference_inputs: list[InferenceInput] = [] + + def map_step_inputs( + self, + *, + canonical_inputs: object, + inference_input: InferenceInput, + request: StepRequest, + ) -> InferenceInput: + del canonical_inputs + self.inference_inputs.append(inference_input) + return InferenceInput(step={"mapped_step": request.step_index}) + + mapping = _RecordingMapping() + runtime = SimpleNamespace( + start_inference_session=lambda: object(), + input_canonicalizer=_RecordingCanonicalizer(), + input_source_schema=object(), + input_mapping=mapping, + ) + provider = manager_module._LegacyWebRTCModelInputProvider(runtime=runtime) + skipped_inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.5, + event_type="key_down", + payload={"key": "w"}, + ), + ) + ) + current_inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=2.5, + event_type="key_up", + payload={"key": "w"}, + ), + ) + ) + + prepared = provider.prepare_step( + request=StepRequirements(step_index=0, input_frame_count=1), + user_window=manager_module.UserInputWindow( + start_s=2.0, + end_s=3.0, + frame_times=(2.25, 2.75), + inputs=current_inputs, + metadata={ + SPARSE_KEY_SEGMENTS_METADATA_KEY: ((2.0, 3.0, frozenset({"w"})),), + WEBRTC_SKIPPED_INPUTS_METADATA_KEY: skipped_inputs, + WEBRTC_SKIPPED_WINDOW_METADATA_KEY: (0.0, 2.0), + }, + ), + ) + + assert prepared.inference_input == InferenceInput(step={"mapped_step": 0}) + assert runtime.input_canonicalizer.windows == [(0.0, 2.0), (2.0, 3.0)] + assert runtime.input_canonicalizer.event_batches == [["key_down"], ["key_up"]] + assert mapping.inference_inputs[0].metadata["frame_times"] == (2.25, 2.75) + assert mapping.inference_inputs[0].metadata["window_start_s"] == 2.0 + assert mapping.inference_inputs[0].metadata["window_end_s"] == 3.0 + assert mapping.inference_inputs[0].metadata[SPARSE_KEY_SEGMENTS_METADATA_KEY] == ( + (2.0, 3.0, frozenset({"w"})), + ) + + @pytest.mark.asyncio async def test_action_keydown_reports_error_when_user_event_queue_full( monkeypatch: pytest.MonkeyPatch, @@ -841,6 +989,147 @@ class _FrequentLogManager(_BaseTestManager): assert perf_logs[0][1][-2:] == (13, 512) +@pytest.mark.asyncio +async def test_realtime_driver_session_uses_shared_step_pipeline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pipeline_calls = 0 + original_pipeline = manager_module.StepPipeline + + class _RecordingPipeline(original_pipeline): + def execute_step( + self, + *, + request: StepRequirements, + user_window: Any, + provider: Any, + session: Any, + output: Any, + metrics: Any, + ) -> Any: + nonlocal pipeline_calls + pipeline_calls += 1 + return original_pipeline.execute_step( + self, + request=request, + user_window=user_window, + provider=provider, + session=session, + output=output, + metrics=metrics, + ) + + class _SharedRuntime: + def __init__(self) -> None: + self.step_requests = 0 + self.step_calls: list[tuple[int, list[Any], list[float]]] = [] + + async def reset_for_new_session(self, session_input: Any = None) -> None: + del session_input + + def next_step_request(self) -> StepRequest | None: + if self.step_requests > 0: + return None + self.step_requests += 1 + return _step_request(step_index=0, input_frame_count=1) + + async def step( + self, + *, + request: StepRequest, + segments: list[Any], + frame_times: list[float], + ) -> StepResult: + self.step_calls.append((request.step_index, segments, frame_times)) + return StepResult(step_index=request.step_index, output="ok", frame_count=1) + + def peek_input_fps(self) -> float: + return 30.0 + + def peek_steady_output_num_frames(self) -> int: + return 1 + + monkeypatch.setattr(manager_module, "StepPipeline", _RecordingPipeline) + runtime = _SharedRuntime() + manager = _make_manager(_BaseTestManager, runtime) + context = manager._shared_run_context(asyncio.get_running_loop()) + reservation = context.admission.try_reserve() + assert reservation is not None + resampler = _SharedResampler(start_v=asyncio.get_running_loop().time()) + input_source = WebRTCInputSource(resampler=resampler) + input_source.handle_browser_payload( + {"type": "action", "action": {"event": "step"}}, + timestamp_s=asyncio.get_running_loop().time(), + ) + managed, video_track, peer, channel = _managed_session(runtime) + managed.resampler = resampler # ty:ignore[invalid-assignment] + managed.input_source = input_source + managed.transport = WebRTCTransportService(loop=asyncio.get_running_loop()) + managed.reservation = reservation + manager._active_session = managed + + managed.generation_task = asyncio.create_task( + manager._run_realtime_driver_session( + managed_session=managed, + context=context, + session_input=None, + ) + ) + await asyncio.wait_for(managed.generation_task, timeout=5.0) + + assert pipeline_calls == 1 + assert runtime.step_calls + assert runtime.step_calls[0][0] == 0 + assert not manager.has_active_session() + assert video_track.closed + assert peer.closed + chunk_done = [ + json.loads(message) + for message in channel.messages + if json.loads(message).get("type") == "chunk_done" + ] + assert len(chunk_done) == 1 + assert chunk_done[0]["model"] == "fake-model" + + +@pytest.mark.asyncio +async def test_realtime_driver_session_reports_non_completed_result( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def fake_run_demo_session_async(**kwargs: Any) -> RunResult: + del kwargs + return RunResult( + status="not_activated", + reason="transport closed before first step", + ) + + monkeypatch.setattr( + manager_module, + "run_demo_session_async", + fake_run_demo_session_async, + ) + runtime = SimpleNamespace() + manager = _make_manager(_BaseTestManager, runtime) + context = manager._shared_run_context(asyncio.get_running_loop()) + reservation = context.admission.try_reserve() + assert reservation is not None + managed, _video_track, _peer, channel = _managed_session(runtime) + managed.reservation = reservation + manager._active_session = managed + + await manager._run_realtime_driver_session( + managed_session=managed, + context=context, + session_input=None, + ) + + assert json.loads(channel.messages[0]) == { + "type": "error", + "message": "transport closed before first step", + } + assert not manager.has_active_session() + + @pytest.mark.asyncio async def test_create_answer_raises_busy_with_subclass_message() -> None: manager = _make_manager( diff --git a/flashdreams/tests/test_webrtc_services.py b/flashdreams/tests/test_webrtc_services.py new file mode 100644 index 000000000..bda6eee4c --- /dev/null +++ b/flashdreams/tests/test_webrtc_services.py @@ -0,0 +1,751 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +import json +import threading +from collections.abc import Mapping, Sequence +from typing import Any + +import pytest + +from flashdreams.runtime import ( + CanonicalInputSchema, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InferenceRuntime, + InferenceSession, + InputMapping, + StepRequirements, + StepResult, +) +from flashdreams.runtime.demo import ( + DemoAdapter, + DemoSpec, + InMemorySessionMetricsRecorder, + ModelInputProvider, + ModelWarmupPlan, + OutputDecision, + PreparedScenario, + ProviderCapabilities, + RealtimeSessionDriver, + RunContext, + RunResult, + RuntimeHost, + SessionEdges, + SessionInfo, + StepPipeline, + UserInputWindow, + WebRTCOutputSpec, + run_demo_session_async, +) +from flashdreams.runtime.demo.timing import ResamplerRealtimeClock +from flashdreams.serving.webrtc.server import SessionBusyError +from flashdreams.serving.webrtc.services import ( + AsyncioBlockingPreparationService, + ThreadSafeWebRTCOutputBridge, + WebRTCActivationPolicy, + WebRTCInputSource, + WebRTCOfferRequest, + WebRTCOutputSink, + WebRTCRunMode, + WebRTCSessionOfferHandler, + WebRTCTransportService, +) + +pytestmark = pytest.mark.ci_cpu + + +@pytest.mark.asyncio +async def test_webrtc_offer_handler_calls_shared_session_helper() -> None: + spec = _webrtc_spec() + adapter = _FakeAdapter() + mode = WebRTCRunMode(edge_factory=_FinishedEdgeFactory()) + context = mode.create_run_context( + spec=spec, + adapter=adapter, + host=RuntimeHost(_UnusedRuntime()), + model_warmup_plan=ModelWarmupPlan(), + ) + answerer = _RecordingAnswerer() + helper_calls: list[DemoSpec] = [] + handler = WebRTCSessionOfferHandler( + context=context, + spec=spec, + adapter=adapter, + run_mode=mode, + answerer=answerer, + session_helper=lambda **kwargs: _record_completed_helper( + helper_calls, + **kwargs, + ), + ) + + answer = await handler.handle_offer(offer_sdp="v=0\r\n", offer_type="offer") + + assert answer == {"sdp": "answer-sdp", "type": "answer"} + assert helper_calls == [spec] + assert answerer.offers == [WebRTCOfferRequest(sdp="v=0\r\n", type="offer")] + + +@pytest.mark.asyncio +async def test_webrtc_busy_rejects_before_prepare_provider_or_answer() -> None: + spec = _webrtc_spec() + adapter = _FakeAdapter() + mode = WebRTCRunMode(edge_factory=_FinishedEdgeFactory()) + context = RunContext( + host=RuntimeHost(_UnusedRuntime()), + run_metrics=InMemorySessionMetricsRecorder(), + admission=_BusyAdmission(), + ) + answerer = _RecordingAnswerer() + handler = WebRTCSessionOfferHandler( + context=context, + spec=spec, + adapter=adapter, + run_mode=mode, + answerer=answerer, + ) + + with pytest.raises(SessionBusyError): + await handler.handle_offer(offer_sdp="v=0\r\n", offer_type="offer") + + assert adapter.prepare_thread_id is None + assert adapter.providers == [] + assert answerer.offers == [] + + +@pytest.mark.asyncio +async def test_webrtc_scenario_prepare_runs_off_event_loop_thread() -> None: + loop_thread_id = threading.get_ident() + spec = _webrtc_spec() + adapter = _FakeAdapter() + + result = await AsyncioBlockingPreparationService().run( + adapter.prepare_scenario, + spec, + ) + + assert isinstance(result, PreparedScenario) + assert adapter.prepare_thread_id is not None + assert adapter.prepare_thread_id != loop_thread_id + + +@pytest.mark.asyncio +async def test_webrtc_input_source_emits_typed_user_inputs() -> None: + resampler = _FakeResampler(dt=0.1, start_v=0.0) + source = WebRTCInputSource(resampler=resampler) + source.handle_browser_message( + json.dumps({"type": "action", "action": {"event": "keydown", "key": "w"}}), + timestamp_s=0.05, + ) + source.handle_browser_message( + json.dumps({"type": "event", "event_id": "prompt-1", "state": "trigger"}), + timestamp_s=0.06, + ) + clock = ResamplerRealtimeClock( + resampler=resampler, + now_fn=lambda: 0.2, + sleep_fn=_record_sleep, + ) + + result = await source.next_realtime_window( + request=StepRequirements(step_index=0, input_frame_count=2), + clock=clock, + ) + + assert source.activation_signal.is_set() + assert resampler.edges == [(0.05, "keydown", "w")] + assert result.window.start_s == pytest.approx(0.0) + assert result.window.end_s == pytest.approx(0.2) + assert [event.event_type for event in result.window.inputs.events] == [ + "key_down", + "text_event", + ] + assert result.window.inputs.events[0].payload == {"key": "w"} + assert result.window.inputs.events[1].payload == { + "event_id": "prompt-1", + "state": "trigger", + } + + +@pytest.mark.asyncio +async def test_webrtc_output_sink_uses_nonblocking_threadsafe_bridge() -> None: + loop = asyncio.get_running_loop() + encoder = _BlockingEncoder() + track = _FakeVideoTrack() + deliveries: list[object] = [] + bridge = ThreadSafeWebRTCOutputBridge( + loop=loop, + video_encoder=encoder, + video_track=track, + on_delivery=deliveries.append, + ) + sink = WebRTCOutputSink(bridge=bridge) + sink.open(SessionInfo()) + step_result = StepResult(step_index=0, frame_count=1) + + decision = sink.write(step_result) + + assert isinstance(decision, OutputDecision) + assert not decision.dropped + assert encoder.prepared_payloads == [step_result.step_index] + await asyncio.wait_for(encoder.started.wait(), timeout=1.0) + assert not encoder.release.is_set() + assert bridge.pending_count == 1 + + encoder.release.set() + await asyncio.wait_for(encoder.done.wait(), timeout=1.0) + await asyncio.sleep(0) + + assert deliveries == ["delivered"] + assert encoder.delivered_payloads == [{"step_index": step_result.step_index}] + assert bridge.pending_count == 0 + sink.close() + + +@pytest.mark.asyncio +async def test_webrtc_output_bridge_prepares_payload_before_async_delivery() -> None: + loop = asyncio.get_running_loop() + encoder = _BlockingEncoder() + track = _FakeVideoTrack() + bridge = ThreadSafeWebRTCOutputBridge( + loop=loop, + video_encoder=encoder, + video_track=track, + ) + sink = WebRTCOutputSink(bridge=bridge) + sink.open(SessionInfo()) + + decision = sink.write(StepResult(step_index=7, frame_count=1)) + + assert not decision.dropped + assert encoder.prepared_payloads == [7] + assert encoder.delivered_payloads == [] + + encoder.release.set() + await asyncio.wait_for(encoder.done.wait(), timeout=1.0) + await asyncio.sleep(0) + + assert encoder.delivered_payloads == [{"step_index": 7}] + sink.close() + + +@pytest.mark.asyncio +async def test_webrtc_output_bridge_drops_full_queue_before_payload_prepare() -> None: + loop = asyncio.get_running_loop() + encoder = _BlockingEncoder() + bridge = ThreadSafeWebRTCOutputBridge( + loop=loop, + video_encoder=encoder, + video_track=_FakeVideoTrack(), + max_pending_chunks=1, + ) + sink = WebRTCOutputSink(bridge=bridge) + sink.open(SessionInfo()) + + first = sink.write(StepResult(step_index=0, frame_count=1)) + second = sink.write(StepResult(step_index=1, frame_count=1)) + + assert not first.dropped + assert second.dropped + assert second.drop_policy == "drop_newest" + assert encoder.prepared_payloads == [0] + + encoder.release.set() + await asyncio.wait_for(encoder.done.wait(), timeout=1.0) + sink.close() + + +@pytest.mark.asyncio +async def test_webrtc_output_bridge_generation_reset_cancels_stale_delivery() -> None: + loop = asyncio.get_running_loop() + encoder = _BlockingEncoder() + track = _FakeVideoTrack() + deliveries: list[object] = [] + chunk_deliveries: list[int] = [] + bridge = ThreadSafeWebRTCOutputBridge( + loop=loop, + video_encoder=encoder, + video_track=track, + on_delivery=deliveries.append, + on_chunk_delivery=lambda chunk: chunk_deliveries.append(chunk.step_index), + ) + sink = WebRTCOutputSink(bridge=bridge) + sink.open(SessionInfo()) + + first = sink.write(StepResult(step_index=0, frame_count=1)) + await asyncio.wait_for(encoder.started.wait(), timeout=1.0) + sink.begin_generation(1) + for _ in range(10): + if track.flush_count: + break + await asyncio.sleep(0) + second = sink.write(StepResult(step_index=1, frame_count=1)) + + assert not first.dropped + assert not second.dropped + assert track.flush_count == 1 + + encoder.release.set() + await asyncio.wait_for(encoder.done.wait(), timeout=1.0) + await asyncio.sleep(0) + + assert deliveries == ["delivered"] + assert chunk_deliveries == [1] + sink.close() + + +@pytest.mark.asyncio +async def test_disconnect_closes_transport_and_releases_reservation_once() -> None: + spec = _webrtc_spec() + adapter = _FakeAdapter() + transport_closed: list[str | None] = [] + transport = WebRTCTransportService(on_close=transport_closed.append) + mode = WebRTCRunMode( + edge_factory=_DisconnectedEdgeFactory(transport=transport), + driver=RealtimeSessionDriver(cleanup_timeout_s=1.0), + ) + context = mode.create_run_context( + spec=spec, + adapter=adapter, + host=RuntimeHost(_SessionRuntime()), + model_warmup_plan=ModelWarmupPlan(), + ) + reservation = context.admission.try_reserve() + assert reservation is not None + transport.disconnect("browser disconnect") + + result = await _record_async_helper( + [], + context=context, + spec=spec, + scenario=adapter.prepare_scenario(spec), + adapter=adapter, + run_mode=mode, + pipeline=StepPipeline(), + reservation=reservation, + ) + transport.close("cleanup close") + + assert result.status == "not_activated" + assert result.reason == "browser disconnect" + assert transport.close_count == 1 + assert transport_closed == ["browser disconnect"] + assert reservation.release_count == 1 # ty:ignore[unresolved-attribute] + assert adapter.providers[0].close_count == 1 + + +def test_webrtc_run_mode_objects_are_control_rank_only() -> None: + spec = _webrtc_spec() + adapter = _FakeAdapter() + edge_factory = _FinishedEdgeFactory() + mode = WebRTCRunMode(edge_factory=edge_factory) + worker_context = mode.create_run_context( + spec=spec, + adapter=adapter, + host=RuntimeHost(_UnusedRuntime(), is_control_rank=False), + model_warmup_plan=ModelWarmupPlan(), + ) + + assert worker_context.services == {} + assert worker_context.admission.try_reserve() is None + with pytest.raises(RuntimeError, match="control-rank only"): + mode.create_session_edges( + context=worker_context, + spec=spec, + scenario=adapter.prepare_scenario(spec), + provider=_FakeProvider(), + adapter=adapter, + ) + + control_context = mode.create_run_context( + spec=spec, + adapter=adapter, + host=RuntimeHost(_UnusedRuntime(), is_control_rank=True), + model_warmup_plan=ModelWarmupPlan(), + ) + assert set(control_context.services) == {"blocking_preparation"} + assert control_context.admission.try_reserve() is not None + + +def _webrtc_spec() -> DemoSpec: + return DemoSpec( + model_id="fake-demo", + input_mode="keyboard-driving", + output=WebRTCOutputSpec(port=8081), + ) + + +async def _record_sleep(delay_s: float) -> None: + del delay_s + + +async def _record_async_helper( + calls: list[DemoSpec], + **kwargs: Any, +) -> RunResult: + calls.append(kwargs["spec"]) + return await run_demo_session_async(**kwargs) + + +async def _record_completed_helper( + calls: list[DemoSpec], + **kwargs: Any, +) -> RunResult: + calls.append(kwargs["spec"]) + return RunResult(status="completed") + + +class _FakeAdapter: + model_id = "fake-demo" + inference_input_schema = InferenceInputSchema() + canonical_input_schema = CanonicalInputSchema() + + def __init__(self) -> None: + self.prepare_thread_id: int | None = None + self.providers: list[_FakeProvider] = [] + + def supported_input_modes(self) -> tuple[str, ...]: + return ("keyboard-driving",) + + def supported_output_modes(self) -> tuple[str, ...]: + return ("webrtc",) + + def default_input_mapping(self) -> InputMapping: + return IdentityInputMapping() + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.model_id: + raise ValueError(f"Unsupported model_id={config.model_id!r}.") + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + return _SessionRuntime() + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + self.prepare_thread_id = threading.get_ident() + assert spec.model_id == self.model_id + return PreparedScenario(initial_inputs=InferenceInput()) + + def create_model_input_provider( + self, + spec: DemoSpec, + scenario: PreparedScenario, + ) -> "_FakeProvider": + del spec, scenario + provider = _FakeProvider() + self.providers.append(provider) + return provider + + +class _FakeProvider: + capabilities = ProviderCapabilities( + supports_realtime_clock=True, + supports_reset=True, + deterministic_given_inputs=False, + ) + + def __init__(self) -> None: + self.close_count = 0 + + def prepare_initial_input(self) -> InferenceInput: + return InferenceInput() + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> Any: + del request, user_window + raise AssertionError("disconnected tests must stop before step prep") + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + + def close(self) -> None: + self.close_count += 1 + + +class _SessionRuntime: + def start_session(self, inputs: InferenceInput) -> InferenceSession: + del inputs + return _NeverSteppedSession() + + def close(self) -> None: + return + + +class _UnusedRuntime: + def start_session(self, inputs: InferenceInput) -> InferenceSession: + del inputs + raise AssertionError("runtime should not be used") + + def close(self) -> None: + return + + +class _NeverSteppedSession: + def next_step_requirements(self) -> StepRequirements | None: + return StepRequirements(step_index=0, input_frame_count=1) + + def next_step_request(self) -> None: + return None + + def step(self, inputs: InferenceInput) -> StepResult: + del inputs + raise AssertionError("disconnected tests must stop before stepping") + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + + def close(self) -> None: + return + + def session_info(self) -> SessionInfo: + return SessionInfo() + + +class _FinishedEdgeFactory: + def __init__(self) -> None: + self.edges: list[SessionEdges] = [] + + def create_session_edges( + self, + *, + context: RunContext, + spec: DemoSpec, + scenario: PreparedScenario, + provider: ModelInputProvider, + adapter: DemoAdapter, + ) -> SessionEdges: + del spec, scenario, provider, adapter + edges = SessionEdges( + input_source=_FinishedRealtimeInputSource(), + output_sink=_RecordingOutputSink(), + cleanup_tasks=context.cleanup_tasks, + metrics=InMemorySessionMetricsRecorder(), + transport=WebRTCTransportService(), + clock=_InstantClock(), + activation=_AlreadyActive(), + ) + self.edges.append(edges) + return edges + + +class _DisconnectedEdgeFactory: + def __init__(self, *, transport: WebRTCTransportService) -> None: + self.transport = transport + + def create_session_edges( + self, + *, + context: RunContext, + spec: DemoSpec, + scenario: PreparedScenario, + provider: ModelInputProvider, + adapter: DemoAdapter, + ) -> SessionEdges: + del spec, scenario, provider, adapter + resampler = _FakeResampler(dt=0.1, start_v=0.0) + source = WebRTCInputSource(resampler=resampler) + return SessionEdges( + input_source=source, + output_sink=_RecordingOutputSink(), + cleanup_tasks=context.cleanup_tasks, + metrics=InMemorySessionMetricsRecorder(), + transport=self.transport, + clock=ResamplerRealtimeClock( + resampler=resampler, + now_fn=lambda: 0.0, + sleep_fn=_record_sleep, + ), + activation=WebRTCActivationPolicy( + input_source=source, + transport=self.transport, + ), + ) + + +class _FinishedRealtimeInputSource: + is_finite = False + is_deterministic = False + user_input_schema = _FakeProvider.capabilities.user_input_schema + + def is_finished(self) -> bool: + return True + + async def next_realtime_window( + self, + *, + request: StepRequirements, + clock: Any, + ) -> Any: + del request, clock + raise AssertionError("finished async driver should not request windows") + + +class _RecordingOutputSink: + produces_artifacts = False + + def __init__(self) -> None: + self.close_count = 0 + + def open(self, session_info: SessionInfo) -> None: + del session_info + + def begin_generation(self, generation: int) -> None: + del generation + + def write(self, result: StepResult) -> OutputDecision: + del result + return OutputDecision() + + def close(self) -> Sequence[Any]: + self.close_count += 1 + return () + + +class _AlreadyActive: + timeout_s = None + + async def wait_until_active(self, clock: Any) -> Any: + del clock + return type("Activation", (), {"activated": True, "reason": None})() + + +class _InstantClock: + is_realtime = True + is_deterministic = False + + def now(self) -> float: + return 0.0 + + def anchor(self, wall_time_s: float) -> None: + del wall_time_s + + async def wait_until_window_end(self, end_s: float) -> None: + del end_s + + async def apply_backpressure(self, requested_s: float) -> None: + del requested_s + + def catch_up( + self, + *, + request: StepRequirements, + max_lag_s: float, + policy: str, + ) -> Any: + del request, max_lag_s, policy + return type("CatchUp", (), {"skipped_s": 0.0})() + + +class _BusyAdmission: + def try_reserve(self) -> None: + return None + + +class _RecordingAnswerer: + def __init__(self) -> None: + self.offers: list[WebRTCOfferRequest] = [] + + async def create_answer( + self, + *, + offer: WebRTCOfferRequest, + session_task: asyncio.Task[RunResult], + ) -> Mapping[str, str]: + self.offers.append(offer) + result = await session_task + assert result.status == "completed" + return {"sdp": "answer-sdp", "type": "answer"} + + +class _FakeResampler: + def __init__(self, *, dt: float, start_v: float) -> None: + self.dt = dt + self.next_chunk_start_v = start_v + self.edges: list[tuple[float, str, str]] = [] + + def reset(self, *, start_v: float) -> None: + self.next_chunk_start_v = start_v + self.edges.clear() + + def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: + self.edges.append((arrival_t, event, key)) + + def sample_chunk( + self, + num_frames: int, + ) -> tuple[tuple[tuple[float, float, frozenset[str]], ...], tuple[float, ...]]: + start = self.next_chunk_start_v + frame_times = tuple(start + index * self.dt for index in range(num_frames)) + end = start + num_frames * self.dt + self.next_chunk_start_v = end + return (((start, end, frozenset({"w"})),), frame_times) + + +class _BlockingEncoder: + fps = 30 + + def __init__(self) -> None: + self.started = asyncio.Event() + self.release = asyncio.Event() + self.done = asyncio.Event() + self.prepared_payloads: list[int] = [] + self.delivered_payloads: list[object] = [] + + def prepare_chunk_payload( + self, + result: StepResult, + track: Any, + ) -> object: + del track + self.prepared_payloads.append(result.step_index) + return {"step_index": result.step_index} + + async def deliver_prepared_chunk( + self, + payload: object, + track: Any, + *, + force_keyframe: bool = False, + ) -> str: + del track, force_keyframe + self.delivered_payloads.append(payload) + self.started.set() + await self.release.wait() + self.done.set() + return "delivered" + + async def deliver_chunk( + self, + result: StepResult, + track: Any, + *, + force_keyframe: bool = False, + ) -> str: + return await self.deliver_prepared_chunk( + self.prepare_chunk_payload(result, track), + track, + force_keyframe=force_keyframe, + ) + + +class _FakeVideoTrack: + fps = 30 + + def __init__(self) -> None: + self.flush_count = 0 + + def qsize(self) -> int: + return 0 + + async def flush(self) -> None: + self.flush_count += 1 diff --git a/flashdreams/tests/test_webrtc_serving.py b/flashdreams/tests/test_webrtc_serving.py index 21d789fe0..06e003496 100644 --- a/flashdreams/tests/test_webrtc_serving.py +++ b/flashdreams/tests/test_webrtc_serving.py @@ -332,7 +332,7 @@ def test_shared_viewer_exposes_model_extension_slots() -> None: html = web_dir.joinpath("request_session.html").read_text(encoding="utf-8") javascript = web_dir.joinpath("request_session.js").read_text(encoding="utf-8") - assert "/static/request_session.js?v=shared-webrtc-v3" in html + assert "/static/request_session.js?v=shared-webrtc-v4" in html for slot in ( "modelStageSlot", "modelStatusSlot", @@ -341,6 +341,8 @@ def test_shared_viewer_exposes_model_extension_slots() -> None: ): assert f'id="{slot}"' in html assert 'fetch("/api/ui/config")' in javascript + assert "config.model_stylesheet" in javascript + assert "stylesheetHrefs" in javascript assert "await modelAdapter?.beforeConnect?.(modelContext)" in javascript assert "sendCommand: sendModelCommand" in javascript assert 'id="postprocessField"' in html @@ -399,7 +401,7 @@ async def test_packaged_webrtc_app_serves_model_adapter(tmp_path) -> None: try: config_response = await client.get("/api/ui/config") assert await config_response.json() == { - "adapter_module": "/model-static/adapter.js?v=model-ui-v1" + "adapter_module": "/model-static/adapter.js?v=model-ui-v2" } adapter_response = await client.get("/model-static/adapter.js") assert adapter_response.status == 200 @@ -408,6 +410,39 @@ async def test_packaged_webrtc_app_serves_model_adapter(tmp_path) -> None: await client.close() +@pytest.mark.asyncio +async def test_packaged_webrtc_app_serves_model_stylesheet(tmp_path) -> None: + shared_dir = tmp_path / "shared" + model_dir = tmp_path / "model" + shared_dir.mkdir() + model_dir.mkdir() + (shared_dir / "request_session.html").write_text("session") + (model_dir / "adapter.css").write_text(".stageVideo { object-fit: contain; }") + app = create_packaged_webrtc_app( + web_resource=shared_dir, + model_web_resource=model_dir, + session_manager=_FakeSessionManager(), + request_session_url="http://127.0.0.1:8080/request_session", + preload_name="Test", + as_file_fn=lambda resource: nullcontext(resource), + ) + client = TestClient(TestServer(app)) + await client.start_server() + try: + config_response = await client.get("/api/ui/config") + assert await config_response.json() == { + "adapter_module": None, + "model_stylesheet": "/model-static/adapter.css?v=model-ui-v2", + } + stylesheet_response = await client.get("/model-static/adapter.css") + assert stylesheet_response.status == 200 + assert ( + await stylesheet_response.text() == ".stageVideo { object-fit: contain; }" + ) + finally: + await client.close() + + def test_webrtc_message_helpers_preserve_public_payload_shape() -> None: assert make_error_payload("boom") == {"type": "error", "message": "boom"} assert make_event_ack_payload( diff --git a/flashdreams/tests/test_webrtc_warmup.py b/flashdreams/tests/test_webrtc_warmup.py new file mode 100644 index 000000000..53ed09f1b --- /dev/null +++ b/flashdreams/tests/test_webrtc_warmup.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from flashdreams.serving.webrtc import warmup as warmup_module +from flashdreams.serving.webrtc.messages import make_error_payload +from flashdreams.serving.webrtc.warmup import run_loopback_warmup_session + +pytestmark = pytest.mark.ci_cpu + + +@pytest.mark.asyncio +async def test_loopback_warmup_fails_on_server_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + channel = _FakeLoopbackChannel( + incoming_on_send=(make_error_payload("shared driver failed"),) + ) + _install_fake_loopback_peer(monkeypatch, channel=channel) + + with pytest.raises(RuntimeError, match="shared driver failed"): + await run_loopback_warmup_session( + num_chunks=1, + warmup_timeout_s=1.0, + create_answer=_fake_create_answer, + action_payloads=(_step_action(),), + ) + + +@pytest.mark.asyncio +async def test_loopback_warmup_fails_on_early_channel_close( + monkeypatch: pytest.MonkeyPatch, +) -> None: + channel = _FakeLoopbackChannel(close_on_send=True) + _install_fake_loopback_peer(monkeypatch, channel=channel) + + with pytest.raises(RuntimeError, match=r"0/1 chunk"): + await run_loopback_warmup_session( + num_chunks=1, + warmup_timeout_s=1.0, + create_answer=_fake_create_answer, + action_payloads=(_step_action(),), + ) + + +async def _fake_create_answer(*, offer_sdp: str, offer_type: str) -> dict[str, str]: + del offer_sdp, offer_type + return {"sdp": "answer-sdp", "type": "answer"} + + +def _step_action() -> dict[str, object]: + return {"type": "action", "action": {"event": "step"}} + + +def _install_fake_loopback_peer( + monkeypatch: pytest.MonkeyPatch, + *, + channel: "_FakeLoopbackChannel", +) -> None: + monkeypatch.setattr( + warmup_module, + "RTCPeerConnection", + lambda _configuration: _FakeLoopbackPeer(channel=channel), + ) + + async def wait_for_ice_gathering_complete(*args: Any, **kwargs: Any) -> None: + del args, kwargs + + monkeypatch.setattr( + warmup_module, + "wait_for_ice_gathering_complete", + wait_for_ice_gathering_complete, + ) + + +class _FakeLoopbackPeer: + iceGatheringState = "complete" + + def __init__(self, *, channel: "_FakeLoopbackChannel") -> None: + self.localDescription: Any | None = None + self._channel = channel + + def createDataChannel(self, *args: Any, **kwargs: Any) -> "_FakeLoopbackChannel": + del args, kwargs + return self._channel + + def addTransceiver(self, *args: Any, **kwargs: Any) -> None: + del args, kwargs + + def on(self, event_name: str) -> Any: + del event_name + + def decorator(callback: Any) -> Any: + return callback + + return decorator + + async def createOffer(self) -> Any: + return SimpleNamespace(sdp="offer-sdp", type="offer") + + async def setLocalDescription(self, description: Any) -> None: + self.localDescription = description + + async def setRemoteDescription(self, description: Any) -> None: + del description + self._channel.open() + + async def close(self) -> None: + self._channel.close() + + +class _FakeLoopbackChannel: + readyState = "open" + + def __init__( + self, + *, + incoming_on_send: tuple[dict[str, str], ...] = (), + close_on_send: bool = False, + ) -> None: + self._incoming_on_send = list(incoming_on_send) + self._close_on_send = close_on_send + self._handlers: dict[str, Any] = {} + + def on(self, event_name: str) -> Any: + def decorator(callback: Any) -> Any: + self._handlers[event_name] = callback + return callback + + return decorator + + def send(self, message: str) -> None: + del message + if self._incoming_on_send: + self._handlers["message"]( + warmup_module.json.dumps(self._incoming_on_send.pop(0)) + ) + if self._close_on_send: + self.close() + + def open(self) -> None: + self._handlers["open"]() + + def close(self) -> None: + self.readyState = "closed" + close_handler = self._handlers.get("close") + if close_handler is not None: + close_handler() diff --git a/integrations/lingbot/tests/test_demo_api.py b/integrations/lingbot/tests/test_demo_api.py index 2cd5edfac..6682dfa5e 100644 --- a/integrations/lingbot/tests/test_demo_api.py +++ b/integrations/lingbot/tests/test_demo_api.py @@ -135,14 +135,17 @@ def fake_runner(**kwargs: Any) -> Sequence[OutputArtifact]: ), ) - artifacts = run_replay_demo( + result = run_replay_demo( spec=spec, adapter=adapter, output_target_factory=lambda output_spec: output, runner=fake_runner, ) - assert artifacts == (OutputArtifact(kind="video/mp4", uri="memory://lingbot"),) + assert result.status == "completed" + assert result.artifacts == ( + OutputArtifact(kind="video/mp4", uri="memory://lingbot"), + ) assert len(calls) == 1 assert calls[0]["adapter"] is adapter assert calls[0]["config"] == spec.config diff --git a/integrations/omnidreams/omnidreams/demo/README.md b/integrations/omnidreams/omnidreams/demo/README.md index ce23b48cd..9dab46ed8 100644 --- a/integrations/omnidreams/omnidreams/demo/README.md +++ b/integrations/omnidreams/omnidreams/demo/README.md @@ -8,21 +8,46 @@ SPDX-License-Identifier: Apache-2.0 This folder contains the experimental OmniDreams demo built on `flashdreams.runtime.demo`. -Run commands from the FlashDreams workspace root: +Run commands from the FlashDreams workspace root. The following setup was used +for remote GPU validation on GB300: ```bash cd /path/to/flashdreams export HF_TOKEN= +export CUDA_HOME=/usr/local/cuda-13.1 +export CUDA_PATH="$CUDA_HOME" +export PATH="$CUDA_HOME/bin:$PATH" +export LD_LIBRARY_PATH="$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}" +hash -r +"$CUDA_HOME/bin/nvcc" --version + +uv sync --python 3.12 --package flashdreams-omnidreams --no-dev ``` -## MP4 Replay +## Null Replay -Generate an MP4 from the bundled single-view sample data: +Run a short replay without writing video output: + +```bash +uv run --python 3.12 --package flashdreams-omnidreams omnidreams-demo replay \ + --output-mode null \ + --device cuda:0 \ + --total-blocks 10 +``` + +## Precomputed MP4 Replay + +Generate an MP4 from bundled single-view sample data and pre-rendered HDMaps: ```bash mkdir -p outputs -uv run --package flashdreams-omnidreams omnidreams-demo replay \ - --output outputs/omnidreams-demo.mp4 +uv run --python 3.12 --package flashdreams-omnidreams omnidreams-demo replay \ + --device cuda:0 \ + --example-data \ + --example-data-uuid 239560dc-33d1-11ef-9720-00044bcbccac \ + --total-blocks 225 \ + --fps 30 \ + --output outputs/omnidreams-demo-precomputed-1min.mp4 ``` This replay path mirrors the benchmark runner path: it uses a prompt, first @@ -30,20 +55,24 @@ frame, and pre-rendered HDMap video. It does not load a Ludus scene or render HDMaps at runtime. The demo defaults to the stable non-perf OmniDreams preset used by the benchmark path. -To provide benchmark-style assets explicitly: +Pass `--example-data-uuid ` to select another bundled single-view sample, +or `--no-example-data` to require explicit asset paths. + +## Ludus MP4 Replay + +Generate an MP4 by rendering HDMap conditioning from a recorded keyboard trace: ```bash -uv run --package flashdreams-omnidreams omnidreams-demo replay \ - --prompt "Driving scene from a front-facing car camera." \ - --hdmap-video-paths /path/to/camera_front_wide_120fov_hdmap.mp4 \ - --first-frame-paths /path/to/first_frame.png \ - --camera-names camera_front_wide_120fov \ - --output outputs/omnidreams-demo.mp4 +uv run --python 3.12 --package flashdreams-omnidreams omnidreams-demo replay \ + --conditioning-mode ludus-scene-driving \ + --keyboard-trace integrations/omnidreams/omnidreams/demo/traces/ludus_forward_sweep_60s.json \ + --device cuda:0 \ + --scene-uuid 0d404ff7-2b66-498c-b047-1ed8cded60d4 \ + --seed 42 \ + --total-blocks 226 \ + --output outputs/omnidreams-demo--ludus-1min.mp4 ``` -Pass `--example-data-uuid ` to select another bundled single-view sample, -or `--no-example-data` to require explicit asset paths. - The `omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf` preset remains an explicit `--preset-id` opt-in. It should become the default only after the compile/cache behavior is reliable enough for the demo path. @@ -55,9 +84,11 @@ 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 \ +uv run --python 3.12 --package flashdreams-omnidreams omnidreams-demo webrtc \ --host 0.0.0.0 \ - --port 8082 + --port 8089 \ + --device cuda:0 \ + --scene-uuid 0d404ff7-2b66-498c-b047-1ed8cded60d4 ``` The scene UUID is optional; when omitted, the runtime uses the default diff --git a/integrations/omnidreams/omnidreams/demo/__init__.py b/integrations/omnidreams/omnidreams/demo/__init__.py index 6fa3a9b21..a3d646989 100644 --- a/integrations/omnidreams/omnidreams/demo/__init__.py +++ b/integrations/omnidreams/omnidreams/demo/__init__.py @@ -4,17 +4,33 @@ """Experimental OmniDreams demo adapter built on ``flashdreams.runtime.demo``.""" from omnidreams.demo.adapter import OmnidreamsDemoAdapter +from omnidreams.demo.providers import ( + LudusSceneConditioningProvider, + PrecomputedHDMapProvider, +) from omnidreams.demo.spec import ( DEFAULT_OMNIDREAMS_PRESET, + OMNIDREAMS_CONDITIONING_LUDUS, + OMNIDREAMS_CONDITIONING_MODES, + OMNIDREAMS_CONDITIONING_PRECOMPUTED, OMNIDREAMS_MODEL_ID, + OmnidreamsKeyboardTraceEvent, + OmnidreamsLudusReplayScenario, OmnidreamsReplayScenario, OmnidreamsWebRTCScenario, ) __all__ = [ "DEFAULT_OMNIDREAMS_PRESET", + "OMNIDREAMS_CONDITIONING_LUDUS", + "OMNIDREAMS_CONDITIONING_MODES", + "OMNIDREAMS_CONDITIONING_PRECOMPUTED", "OMNIDREAMS_MODEL_ID", + "LudusSceneConditioningProvider", "OmnidreamsDemoAdapter", + "OmnidreamsKeyboardTraceEvent", + "OmnidreamsLudusReplayScenario", "OmnidreamsReplayScenario", "OmnidreamsWebRTCScenario", + "PrecomputedHDMapProvider", ] diff --git a/integrations/omnidreams/omnidreams/demo/adapter.py b/integrations/omnidreams/omnidreams/demo/adapter.py index a1ca8e8f7..4dc7d1c87 100644 --- a/integrations/omnidreams/omnidreams/demo/adapter.py +++ b/integrations/omnidreams/omnidreams/demo/adapter.py @@ -5,8 +5,9 @@ from __future__ import annotations +import math from collections.abc import Callable -from typing import Any +from typing import Any, cast from omnidreams.config import OMNIDREAMS_CONFIGS, OMNIDREAMS_RUNNERS @@ -17,28 +18,45 @@ InferenceInput, InferenceInputSchema, InputCanonicalizer, - InputField, UserInputSchema, ) from flashdreams.runtime.demo import ( DemoSpec, - Mp4OutputSpec, PreparedScenario, ) +from flashdreams.runtime.demo.session_inputs import ModelInputProvider from flashdreams.runtime.interfaces import InferenceRuntime +from flashdreams.serving.webrtc.services import WEBRTC_USER_INPUT_SCHEMA -from .replay import ( - OmnidreamsReplayRuntime, - OmnidreamsReplayRuntimeOptions, +from .providers import ( + LudusSceneConditioningProvider, + PrecomputedHDMapProvider, + keyboard_driving_user_input_schema, + precomputed_hdmap_inference_input_schema, +) +from .runtime import ( + OmnidreamsRuntime, + OmnidreamsRuntimeOptions, PipelineFactory, ) from .spec import ( DEFAULT_OMNIDREAMS_PRESET, + DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID, + OMNIDREAMS_CONDITIONING_LUDUS, + OMNIDREAMS_CONDITIONING_MODES, + OMNIDREAMS_CONDITIONING_PRECOMPUTED, OMNIDREAMS_MODEL_ID, + LudusBackendName, + OmnidreamsLudusReplayScenario, + OmnidreamsWebRTCScenario, + conditioning_mode_from_scenario, + resolve_ludus_replay_scenario, resolve_replay_scenario, + resolve_webrtc_scenario, ) -ReplayRuntimeFactory = Callable[..., InferenceRuntime] +RuntimeFactory = Callable[..., InferenceRuntime] +ReplayRuntimeFactory = RuntimeFactory class OmnidreamsDemoAdapter: @@ -47,10 +65,19 @@ class OmnidreamsDemoAdapter: def __init__( self, *, - replay_runtime_factory: ReplayRuntimeFactory = OmnidreamsReplayRuntime, + runtime_factory: RuntimeFactory | None = None, + replay_runtime_factory: ReplayRuntimeFactory | None = None, pipeline_factory: PipelineFactory | None = None, ) -> None: - self._replay_runtime_factory = replay_runtime_factory + if runtime_factory is not None and replay_runtime_factory is not None: + raise ValueError( + "Specify either runtime_factory or replay_runtime_factory, not both." + ) + self._runtime_factory = ( + runtime_factory + if runtime_factory is not None + else replay_runtime_factory or OmnidreamsRuntime + ) self._pipeline_factory = pipeline_factory self._mapping = IdentityInputMapping() @@ -60,15 +87,7 @@ def model_id(self) -> str: @property def inference_input_schema(self) -> InferenceInputSchema: - return InferenceInputSchema( - global_conditioning_fields=( - InputField( - name="scenario", - input_modality="omnidreams/replay-scenario", - description="Resolved OmniDreams replay scenario.", - ), - ) - ) + return precomputed_hdmap_inference_input_schema() @property def canonical_input_schema(self) -> CanonicalInputSchema | None: @@ -78,31 +97,57 @@ def default_input_mapping(self) -> IdentityInputMapping: return self._mapping def supported_input_modes(self) -> tuple[str, ...]: - return ("replay",) + return ("replay", "keyboard-driving") def supported_output_modes(self) -> tuple[str, ...]: - return ("mp4",) + return ("mp4", "null", "webrtc") + + def supported_conditioning_modes(self) -> tuple[str, ...]: + return OMNIDREAMS_CONDITIONING_MODES def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + if spec.output.mode not in self.supported_output_modes(): + raise ValueError( + "OmniDreams demo supports output modes " + f"{self.supported_output_modes()}, got {spec.output.mode!r}." + ) + if spec.input_mode == "keyboard-driving": + if spec.output.mode != "webrtc": + raise ValueError( + "OmniDreams keyboard-driving input currently requires " + f"output.mode='webrtc', got {spec.output.mode!r}." + ) + return self._prepare_webrtc_scenario(spec) if spec.input_mode != "replay": raise ValueError( - "OmniDreams prepare_scenario currently supports only " - f"input_mode='replay', got {spec.input_mode!r}." + "OmniDreams prepare_scenario supports input modes " + f"{self.supported_input_modes()}, got {spec.input_mode!r}." + ) + if spec.output.mode == "webrtc": + raise ValueError("OmniDreams replay input does not support WebRTC output.") + conditioning_mode = conditioning_mode_from_scenario(spec.scenario) + if conditioning_mode == OMNIDREAMS_CONDITIONING_PRECOMPUTED: + scenario = resolve_replay_scenario( + spec.scenario, + default_prompt=self._default_replay_prompt(spec.config), + ) + source_schema = UserInputSchema(description="fixed OmniDreams replay input") + elif conditioning_mode == OMNIDREAMS_CONDITIONING_LUDUS: + scenario = resolve_ludus_replay_scenario(spec.scenario) + source_schema = keyboard_driving_user_input_schema() + else: + raise ValueError( + f"Unsupported OmniDreams conditioning mode: {conditioning_mode!r}." ) - if not isinstance(spec.output, Mp4OutputSpec): - raise ValueError("OmniDreams replay demo currently requires MP4 output.") - scenario = resolve_replay_scenario( - spec.scenario, - default_prompt=self._default_replay_prompt(spec.config), - ) return PreparedScenario( initial_inputs=InferenceInput( global_conditioning={"scenario": scenario}, ), - source_schema=UserInputSchema(description="fixed OmniDreams replay input"), + source_schema=source_schema, canonicalizer=InputCanonicalizer(), mapping=self._mapping, metadata={ + "conditioning_mode": conditioning_mode, "model_id": self.model_id, "preset_id": self._preset_id(spec.config), "num_views": len(scenario.camera_names), @@ -119,11 +164,111 @@ def validate_config(self, config: InferenceConfig) -> None: def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: self.validate_config(config) - return self._replay_runtime_factory( + return self._runtime_factory( config=config, - options=OmnidreamsReplayRuntimeOptions( + options=OmnidreamsRuntimeOptions( pipeline_config=self._pipeline_config(config), pipeline_factory=self._pipeline_factory, + release_oneshot_encoders_after_cache_init=( + _bool_runtime_option( + config.runtime_options, + "release_oneshot_encoders_after_cache_init", + True, + ) + ), + ), + ) + + def create_model_input_provider( + self, + spec: DemoSpec, + scenario: PreparedScenario, + ) -> ModelInputProvider: + if spec.input_mode not in {"replay", "keyboard-driving"}: + raise ValueError( + "OmniDreams providers support input modes " + f"{self.supported_input_modes()}, got {spec.input_mode!r}." + ) + if spec.config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + conditioning_mode = str( + scenario.metadata.get( + "conditioning_mode", + conditioning_mode_from_scenario(spec.scenario), + ) + ) + if conditioning_mode == OMNIDREAMS_CONDITIONING_LUDUS: + return LudusSceneConditioningProvider( + scenario=scenario, + config=spec.config, + ) + if conditioning_mode != OMNIDREAMS_CONDITIONING_PRECOMPUTED: + raise ValueError( + f"Unsupported OmniDreams conditioning mode: {conditioning_mode!r}." + ) + return PrecomputedHDMapProvider( + scenario=scenario, + config=spec.config, + ) + + def _prepare_webrtc_scenario(self, spec: DemoSpec) -> PreparedScenario: + if spec.config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + scenario = self._webrtc_ludus_scenario( + resolve_webrtc_scenario(spec.scenario), + spec=spec, + ) + return PreparedScenario( + initial_inputs=InferenceInput( + global_conditioning={"scenario": scenario}, + ), + source_schema=WEBRTC_USER_INPUT_SCHEMA, + canonicalizer=InputCanonicalizer(), + mapping=self._mapping, + metadata={ + "conditioning_mode": OMNIDREAMS_CONDITIONING_LUDUS, + "model_id": self.model_id, + "preset_id": self._preset_id(spec.config), + "num_views": 1, + }, + ) + + def _webrtc_ludus_scenario( + self, + scenario: OmnidreamsWebRTCScenario, + *, + spec: DemoSpec, + ) -> OmnidreamsLudusReplayScenario: + config = spec.config + if config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + output = spec.output + fps = int(getattr(output, "fps", 30)) + video_height = int(getattr(output, "video_height", 704)) + video_width = int(getattr(output, "video_width", 1280)) + options = config.runtime_options + return OmnidreamsLudusReplayScenario( + keyboard_events=(), + scene_dir=scenario.scene_dir, + scene_uuid=scenario.scene_uuid or DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID, + scene_variant=scenario.scene_variant, + camera_name=scenario.camera_name, + total_blocks=int( + options.get( + "total_blocks", + options.get("webrtc_total_blocks", 2_147_483_647), + ) + ), + pixel_height=video_height, + pixel_width=video_width, + fps=fps, + move_speed_per_s=float(options.get("move_speed_per_s", 6.0)), + rotate_speed_rad_per_s=float( + options.get("rotate_speed_rad_per_s", math.radians(35.0)) + ), + ludus_backend=cast( + LudusBackendName, + str(options.get("ludus_backend", "cuda")), ), ) @@ -153,7 +298,25 @@ def _default_replay_prompt(self, config: InferenceConfig | None) -> str: return "" if runner is None else str(getattr(runner, "prompt", "")) +def _bool_runtime_option( + options: Any, + name: str, + default: bool, +) -> bool: + value = options.get(name, default) + if isinstance(value, bool): + return value + if isinstance(value, str): + lowered = value.strip().lower() + if lowered in {"1", "true", "yes", "on"}: + return True + if lowered in {"0", "false", "no", "off"}: + return False + return bool(value) + + __all__ = [ "OmnidreamsDemoAdapter", "ReplayRuntimeFactory", + "RuntimeFactory", ] diff --git a/integrations/omnidreams/omnidreams/demo/app.py b/integrations/omnidreams/omnidreams/demo/app.py index 1366643f9..ac1641917 100644 --- a/integrations/omnidreams/omnidreams/demo/app.py +++ b/integrations/omnidreams/omnidreams/demo/app.py @@ -6,6 +6,7 @@ from __future__ import annotations import argparse +import math from pathlib import Path from typing import Any @@ -15,6 +16,7 @@ from flashdreams.runtime.demo import ( DemoSpec, Mp4OutputSpec, + NullOutputSpec, WebRTCOutputSpec, ) from flashdreams.runtime.demo.app import DemoApplication @@ -22,6 +24,10 @@ from .adapter import OmnidreamsDemoAdapter from .spec import ( DEFAULT_OMNIDREAMS_PRESET, + DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID, + OMNIDREAMS_CONDITIONING_LUDUS, + OMNIDREAMS_CONDITIONING_MODES, + OMNIDREAMS_CONDITIONING_PRECOMPUTED, OMNIDREAMS_MODEL_ID, OmnidreamsWebRTCScenario, ) @@ -33,13 +39,32 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: ) subparsers = parser.add_subparsers(dest="command", required=True) - replay = subparsers.add_parser("replay", help="Run an MP4 replay demo.") + replay = subparsers.add_parser("replay", help="Run a finite replay demo.") replay.add_argument("--preset-id", default=DEFAULT_OMNIDREAMS_PRESET) replay.add_argument("--device", default="cuda") + replay.add_argument("--seed", type=int, default=42) + replay.add_argument( + "--conditioning-mode", + choices=OMNIDREAMS_CONDITIONING_MODES, + default=OMNIDREAMS_CONDITIONING_PRECOMPUTED, + ) replay.add_argument("--prompt", default=None) replay.add_argument("--hdmap-video-paths", type=_split_paths, default=()) replay.add_argument("--first-frame-paths", type=_split_paths, default=()) replay.add_argument("--camera-names", type=_split_strings, default=()) + replay.add_argument("--keyboard-trace", type=Path, default=None) + replay.add_argument("--scene-path", type=Path, default=None) + replay.add_argument("--scene-dir", type=Path, default=None) + replay.add_argument("--scene-uuid", default=DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID) + replay.add_argument("--scene-variant", default="default") + replay.add_argument("--camera-name", default="camera_front_wide_120fov") + replay.add_argument("--move-speed-per-s", type=float, default=6.0) + replay.add_argument( + "--rotate-speed-rad-per-s", + type=float, + default=math.radians(35.0), + ) + replay.add_argument("--ludus-backend", choices=("cuda", "vulkan"), default="cuda") replay.add_argument( "--example-data", action=argparse.BooleanOptionalAction, @@ -54,7 +79,8 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: replay.add_argument("--pixel-height", type=int, default=704) replay.add_argument("--pixel-width", type=int, default=1280) replay.add_argument("--fps", type=int, default=30) - replay.add_argument("--output", type=Path, required=True) + replay.add_argument("--output-mode", choices=("mp4", "null"), default="mp4") + replay.add_argument("--output", type=Path, default=None) webrtc = subparsers.add_parser("webrtc", help="Serve a WebRTC driving demo.") webrtc.add_argument("--preset-id", default=DEFAULT_OMNIDREAMS_PRESET) @@ -74,7 +100,21 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: webrtc.add_argument("--client-liveness-timeout-s", type=float, default=10.0) webrtc.add_argument("--debug-serve-hdmaps", action="store_true") webrtc.add_argument("--prefer-sw-encoder", action="store_true") - return parser.parse_args(argv) + args = parser.parse_args(argv) + if args.command == "replay": + if args.output_mode == "mp4" and args.output is None: + parser.error("replay --output is required when --output-mode=mp4.") + if args.output_mode == "null" and args.output is not None: + parser.error("replay --output is only valid when --output-mode=mp4.") + if ( + args.conditioning_mode == OMNIDREAMS_CONDITIONING_LUDUS + and args.keyboard_trace is None + ): + parser.error( + "replay --keyboard-trace is required when " + "--conditioning-mode=ludus-scene-driving." + ) + return args class OmnidreamsDemoApplication(DemoApplication): @@ -108,6 +148,7 @@ def main(argv: list[str] | None = None) -> None: def _replay_spec(args: argparse.Namespace) -> DemoSpec: scenario: dict[str, object] = { + "conditioning_mode": args.conditioning_mode, "example_data": args.example_data, "example_data_uuid": args.example_data_uuid, "total_blocks": args.total_blocks, @@ -117,27 +158,56 @@ def _replay_spec(args: argparse.Namespace) -> DemoSpec: } if args.prompt: scenario["prompt"] = args.prompt - if args.hdmap_video_paths: - scenario["hdmap_video_paths"] = args.hdmap_video_paths - if args.first_frame_paths: - scenario["first_frame_paths"] = args.first_frame_paths - if args.camera_names: - scenario["camera_names"] = args.camera_names + if args.conditioning_mode == OMNIDREAMS_CONDITIONING_LUDUS: + scenario.update( + { + "keyboard_trace_path": args.keyboard_trace, + "scene_path": args.scene_path, + "scene_dir": args.scene_dir, + "scene_uuid": args.scene_uuid, + "scene_variant": args.scene_variant, + "camera_name": args.camera_name, + "move_speed_per_s": args.move_speed_per_s, + "rotate_speed_rad_per_s": args.rotate_speed_rad_per_s, + "ludus_backend": args.ludus_backend, + } + ) + else: + if args.hdmap_video_paths: + scenario["hdmap_video_paths"] = args.hdmap_video_paths + if args.first_frame_paths: + scenario["first_frame_paths"] = args.first_frame_paths + if args.camera_names: + scenario["camera_names"] = args.camera_names return DemoSpec( model_id=OMNIDREAMS_MODEL_ID, preset_id=args.preset_id, input_mode="replay", scenario=scenario, - output=Mp4OutputSpec(path=args.output, fps=args.fps), + output=_replay_output_spec(args), config=InferenceConfig( model_id=OMNIDREAMS_MODEL_ID, preset_id=args.preset_id, device=args.device, + seed=args.seed, + runtime_options={"seed": args.seed}, ), ) +def _replay_output_spec(args: argparse.Namespace) -> Mp4OutputSpec | NullOutputSpec: + if args.output_mode == "mp4": + if args.output is None: + raise ValueError("OmniDreams MP4 replay requires --output.") + return Mp4OutputSpec(path=args.output, fps=args.fps) + if args.output_mode == "null": + return NullOutputSpec() + raise ValueError( + f"Unsupported OmniDreams replay output mode: {args.output_mode!r}." + ) + + def _webrtc_spec(args: argparse.Namespace, *, device: str) -> DemoSpec: return DemoSpec( model_id=OMNIDREAMS_MODEL_ID, diff --git a/integrations/omnidreams/omnidreams/demo/providers.py b/integrations/omnidreams/omnidreams/demo/providers.py new file mode 100644 index 000000000..eff6de760 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/providers.py @@ -0,0 +1,692 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OmniDreams model-input providers for shared demo run modes.""" + +from __future__ import annotations + +import contextlib +import os +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import torch.distributed as dist +from loguru import logger +from omnidreams.runner import _load_video + +from flashdreams.infra.runner_io import ( + DEFAULT_RUNNER_INSTALL_HINT, + load_first_frame_tensor, +) +from flashdreams.runtime.config import InferenceConfig +from flashdreams.runtime.demo import ( + PreparedScenario, + PreparedStep, + ProviderCapabilities, + UserInputWindow, +) +from flashdreams.runtime.demo.session_inputs import ControlDecision +from flashdreams.runtime.demo.timing import SPARSE_KEY_SEGMENTS_METADATA_KEY +from flashdreams.runtime.inputs import ( + InferenceInput, + InferenceInputSchema, + InputField, + UserInputCapability, + UserInputSchema, +) +from flashdreams.runtime.types import StepRequirements +from flashdreams.serving.realtime.input import ( + WSAD_SUPPORTED_KEYS, + CameraPoseIntegrator, + KeyboardResampler, + PoseSegment, +) + +from .spec import ( + DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID, + OmnidreamsLudusReplayScenario, + OmnidreamsReplayScenario, +) + + +class PrecomputedHDMapProvider: + """Prepare fixed OmniDreams HDMap conditioning for replay-style runs.""" + + def __init__( + self, + *, + scenario: PreparedScenario, + config: InferenceConfig, + ) -> None: + self._scenario = _precomputed_scenario_from_prepared(scenario) + self._device = _device_from_config(config) + self._dtype = torch.bfloat16 + self._frame_start = 0 + self._closed = False + self.capabilities = ProviderCapabilities( + supports_recorded_input=True, + supports_reset=True, + deterministic_given_inputs=True, + user_input_schema=scenario.source_schema, + inference_input_schema=precomputed_hdmap_inference_input_schema(), + ) + self._hdmap_videos: torch.Tensor | None = self._load_hdmaps() + + def prepare_initial_input(self) -> InferenceInput: + self._require_open() + scenario = self._scenario + first_frames = [ + load_first_frame_tensor( + path, + pixel_height=scenario.pixel_height, + pixel_width=scenario.pixel_width, + device=self._device, + dtype=self._dtype, + allow_video=True, + install_hint=DEFAULT_RUNNER_INSTALL_HINT, + ) + for path in scenario.first_frame_paths + ] + return InferenceInput( + global_conditioning={ + "scenario": scenario, + "prompt": [list(scenario.prompts)], + "first_frame": torch.stack(first_frames, dim=0).unsqueeze(0), + }, + metadata={"view_names": tuple(scenario.camera_names)}, + ) + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + del user_window + self._require_open() + hdmap_videos = self._require_hdmaps() + frame_end = self._frame_start + request.input_frame_count + if frame_end > hdmap_videos.shape[2]: + return PreparedStep( + control=ControlDecision( + close_session=True, + reason="OmniDreams precomputed HDMap input exhausted.", + ) + ) + + frame_start = self._frame_start + self._frame_start = frame_end + return PreparedStep( + inference_input=InferenceInput( + step={"hdmap": hdmap_videos[:, :, frame_start:frame_end]}, + metadata={ + "hdmap_frame_start": frame_start, + "hdmap_frame_end": frame_end, + }, + ) + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self._require_open() + self._frame_start = 0 + + def close(self) -> None: + if self._closed: + return + self._closed = True + self._hdmap_videos = None + + def _load_hdmaps(self) -> torch.Tensor: + scenario = self._scenario + videos = [ + _load_video( + path, + pixel_height=scenario.pixel_height, + pixel_width=scenario.pixel_width, + device=self._device, + dtype=self._dtype, + ) + for path in scenario.hdmap_video_paths + ] + hdmap_videos = torch.stack(videos, dim=0).unsqueeze(0) + if _is_rank_zero(): + logger.info( + "Loaded OmniDreams demo HDMaps shape={} views={}", + tuple(hdmap_videos.shape), + len(scenario.camera_names), + ) + return hdmap_videos + + def _require_hdmaps(self) -> torch.Tensor: + hdmap_videos = self._hdmap_videos + if hdmap_videos is None: + raise RuntimeError("OmniDreams precomputed HDMap provider is closed.") + return hdmap_videos + + def _require_open(self) -> None: + if self._closed: + raise RuntimeError("OmniDreams precomputed HDMap provider is closed.") + + +class LudusSceneConditioningProvider: + """Render finite Ludus keyboard-driving traces into OmniDreams HDMaps.""" + + def __init__( + self, + *, + scenario: PreparedScenario, + config: InferenceConfig, + ) -> None: + self._scenario = _ludus_scenario_from_prepared(scenario) + self._device = _device_from_config(config) + self._dtype = torch.bfloat16 + self._closed = False + self._scene: Any | None = None + self._rasterizer: Any | None = None + self._pose_integrator: CameraPoseIntegrator | None = None + self._keyboard_resampler: KeyboardResampler | None = None + self._next_timestamp_us = 0 + self._step_index = 0 + self.capabilities = ProviderCapabilities( + supports_realtime_clock=True, + supports_recorded_input=True, + supports_reset=True, + deterministic_given_inputs=True, + user_input_schema=scenario.source_schema, + inference_input_schema=precomputed_hdmap_inference_input_schema(), + ) + + def prepare_initial_input(self) -> InferenceInput: + self._require_open() + scene = self._ensure_scene_loaded() + return InferenceInput( + global_conditioning={ + "scenario": self._scenario, + "prompt": [[str(scene.prompt)]], + "first_frame": _initial_rgb_tensor( + scene.initial_rgb, + device=self._device, + dtype=self._dtype, + ), + }, + metadata={ + "view_names": self._scenario.camera_names, + "scene_id": str(getattr(scene, "scene_id", "")), + }, + ) + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + self._require_open() + scenario = self._scenario + if request.step_index >= scenario.total_blocks: + return PreparedStep( + control=ControlDecision( + close_session=True, + reason="OmniDreams Ludus replay input exhausted.", + ) + ) + + self._ensure_scene_loaded() + pose_integrator = self._require_pose_integrator() + rasterizer = self._require_rasterizer() + segments, frame_times = self._sample_controls( + request=request, + user_window=user_window, + ) + rig_poses_world = pose_integrator.integrate_chunk( + segments=segments, + frame_times=frame_times, + ) + timestamps_us = self._consume_timestamps(request.input_frame_count) + raster_chunk = rasterizer.render_chunk( + rig_poses_world=rig_poses_world, + timestamps_us=timestamps_us, + ) + hdmap = _condition_frames_tensor( + raster_chunk.frames, + device=self._device, + dtype=self._dtype, + ) + self._step_index += 1 + return PreparedStep( + inference_input=InferenceInput( + step={"hdmap": hdmap}, + metadata={ + "frame_timestamps_us": tuple(int(t) for t in timestamps_us), + "keyboard_segments": _segments_metadata(segments), + "camera_name": scenario.camera_name, + "scene_uuid": scenario.scene_uuid, + }, + ) + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self._require_open() + if self._scene is not None: + self._reset_driving_state(self._scene) + else: + self._step_index = 0 + self._next_timestamp_us = 0 + + def close(self) -> None: + if self._closed: + return + self._closed = True + rasterizer = self._rasterizer + self._rasterizer = None + self._scene = None + self._pose_integrator = None + self._keyboard_resampler = None + _close_rasterizer(rasterizer) + + def _ensure_scene_loaded(self) -> Any: + if self._scene is not None: + return self._scene + scenario = self._scenario + scene_path = _resolve_ludus_scene_path(scenario) + scene = _load_ludus_scene_bundle(scenario, scene_path) + rasterizer = _new_ludus_rasterizer(scenario) + try: + rasterizer.load_scene(scene) + except Exception: + with contextlib.suppress(Exception): + _close_rasterizer(rasterizer) + raise + self._scene = scene + self._rasterizer = rasterizer + self._reset_driving_state(scene) + if _is_rank_zero(): + logger.info( + "Loaded OmniDreams Ludus replay scene={} camera={} trace_events={}", + scene_path, + scenario.camera_name, + len(scenario.keyboard_events), + ) + return scene + + def _reset_driving_state(self, scene: Any) -> None: + scenario = self._scenario + pose_integrator = CameraPoseIntegrator( + move_speed_per_s=scenario.move_speed_per_s, + rotate_speed_rad_per_s=scenario.rotate_speed_rad_per_s, + coordinate_system="FLU", + ) + pose_integrator.reset(np.asarray(scene.initial_rig_to_world, dtype=np.float32)) + keyboard_resampler = KeyboardResampler( + fps=float(scenario.fps), + supported_keys=WSAD_SUPPORTED_KEYS, + ) + for event in scenario.keyboard_events: + keyboard_resampler.on_edge( + arrival_t=event.timestamp_s, + event=event.event, + key=event.key, + ) + self._pose_integrator = pose_integrator + self._keyboard_resampler = keyboard_resampler + self._next_timestamp_us = int(scene.initial_timestamp_us) + self._step_index = 0 + + def _sample_controls( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> tuple[list[PoseSegment], list[float]]: + raw_segments = user_window.metadata.get(SPARSE_KEY_SEGMENTS_METADATA_KEY) + if isinstance(raw_segments, tuple): + frame_times = list(user_window.frame_times) + if len(frame_times) != request.input_frame_count: + raise RuntimeError( + "OmniDreams Ludus realtime window frame_times length does " + "not match the requested input frame count." + ) + return [_pose_segment(segment) for segment in raw_segments], frame_times + if raw_segments is not None: + raise RuntimeError( + "OmniDreams Ludus realtime key segments metadata must be a tuple." + ) + return self._require_keyboard_resampler().sample_chunk( + request.input_frame_count + ) + + def _consume_timestamps(self, num_frames: int) -> np.ndarray: + step_us = int(round(1_000_000 / float(self._scenario.fps))) + timestamps = np.array( + [ + self._next_timestamp_us + frame_index * step_us + for frame_index in range(num_frames) + ], + dtype=np.int64, + ) + self._next_timestamp_us += num_frames * step_us + return timestamps + + def _require_rasterizer(self) -> Any: + if self._rasterizer is None: + raise RuntimeError("OmniDreams Ludus rasterizer is not initialized.") + return self._rasterizer + + def _require_pose_integrator(self) -> CameraPoseIntegrator: + if self._pose_integrator is None: + raise RuntimeError("OmniDreams Ludus pose integrator is not initialized.") + return self._pose_integrator + + def _require_keyboard_resampler(self) -> KeyboardResampler: + if self._keyboard_resampler is None: + raise RuntimeError( + "OmniDreams Ludus keyboard resampler is not initialized." + ) + return self._keyboard_resampler + + def _require_open(self) -> None: + if self._closed: + raise RuntimeError("OmniDreams Ludus conditioning provider is closed.") + + +def keyboard_driving_user_input_schema() -> UserInputSchema: + return UserInputSchema( + capabilities=( + UserInputCapability( + event_type="keydown", + input_modality="keyboard", + payload_fields=frozenset({"key"}), + description="Keyboard key press edge.", + ), + UserInputCapability( + event_type="keyup", + input_modality="keyboard", + payload_fields=frozenset({"key"}), + description="Keyboard key release edge.", + ), + ), + description="Recorded or realtime WSAD keyboard driving controls.", + ) + + +def precomputed_hdmap_inference_input_schema() -> InferenceInputSchema: + return InferenceInputSchema( + global_conditioning_fields=( + InputField( + name="prompt", + input_modality="omnidreams/prompt", + description="OmniDreams prompt batch.", + ), + InputField( + name="first_frame", + input_modality="video/frame", + description="Initial OmniDreams conditioning frame tensor.", + ), + InputField( + name="scenario", + required=False, + input_modality="omnidreams/replay-scenario", + description="Resolved OmniDreams replay scenario metadata.", + ), + ), + step_fields=( + InputField( + name="hdmap", + input_modality="omnidreams/hdmap-video", + frequency_consumed="per_step", + description="Per-step HDMap conditioning chunk.", + ), + ), + ) + + +def _precomputed_scenario_from_prepared( + scenario: PreparedScenario, +) -> OmnidreamsReplayScenario: + value = scenario.initial_inputs.global_conditioning.get("scenario") + if not isinstance(value, OmnidreamsReplayScenario): + raise TypeError( + "OmniDreams precomputed HDMap provider requires " + "initial_inputs.global_conditioning['scenario'] to be an " + "OmnidreamsReplayScenario." + ) + return value + + +def _ludus_scenario_from_prepared( + scenario: PreparedScenario, +) -> OmnidreamsLudusReplayScenario: + value = scenario.initial_inputs.global_conditioning.get("scenario") + if not isinstance(value, OmnidreamsLudusReplayScenario): + raise TypeError( + "OmniDreams Ludus conditioning provider requires " + "initial_inputs.global_conditioning['scenario'] to be an " + "OmnidreamsLudusReplayScenario." + ) + return value + + +def _resolve_ludus_scene_path(scenario: OmnidreamsLudusReplayScenario) -> Path: + if scenario.scene_path is not None: + if not scenario.scene_path.exists(): + raise FileNotFoundError( + f"OmniDreams Ludus scene_path missing: {scenario.scene_path}" + ) + return scenario.scene_path + if scenario.scene_dir is not None: + return _resolve_local_ludus_scene_path(scenario) + + from omnidreams.scenes import hf_hub_download_scene # noqa: PLC0415 + + return hf_hub_download_scene( + scenario.scene_uuid or DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID, + scenario.scene_variant, + ) + + +def _resolve_local_ludus_scene_path(scenario: OmnidreamsLudusReplayScenario) -> Path: + scene_dir = scenario.scene_dir + if scene_dir is None: + raise RuntimeError("OmniDreams Ludus scene_dir is unexpectedly unset.") + if scene_dir.is_file(): + return scene_dir + if not scene_dir.is_dir(): + raise FileNotFoundError(f"OmniDreams Ludus scene_dir missing: {scene_dir}") + + candidates = _local_ludus_scene_candidates(scenario) + for candidate in candidates: + if candidate.is_file(): + return candidate + archives = sorted(scene_dir.glob("*.usdz")) + if scenario.scene_uuid is None and len(archives) == 1: + return archives[0] + expected = ", ".join(path.name for path in candidates) + raise FileNotFoundError( + f"No OmniDreams Ludus USDZ scene archive found in {scene_dir}. " + f"Expected one of: {expected}." + ) + + +def _local_ludus_scene_candidates( + scenario: OmnidreamsLudusReplayScenario, +) -> tuple[Path, ...]: + scene_dir = scenario.scene_dir + if scene_dir is None or scenario.scene_uuid is None: + return () + + from omnidreams.scenes import ( # noqa: PLC0415 + normalise_scene_uuid, + scene_variant_suffix, + ) + + bare_uuid = normalise_scene_uuid(scenario.scene_uuid) + suffix = scene_variant_suffix(scenario.scene_variant) + stems = [f"clipgt-{bare_uuid}{suffix}", f"{bare_uuid}{suffix}"] + if suffix: + stems.extend((f"clipgt-{bare_uuid}", bare_uuid)) + return tuple(scene_dir / f"{stem}.usdz" for stem in dict.fromkeys(stems)) + + +def _load_ludus_scene_bundle( + scenario: OmnidreamsLudusReplayScenario, + scene_path: Path, +) -> Any: + from omnidreams.interactive_drive.scene_loader import ( # noqa: PLC0415 + load_scene_bundle, + ) + + return load_scene_bundle( + scene_path=scene_path, + camera_name=scenario.camera_name, + variant=scenario.scene_variant, + prompt_override=scenario.prompt, + raster=_ludus_raster_config(scenario), + ) + + +def _new_ludus_rasterizer(scenario: OmnidreamsLudusReplayScenario) -> Any: + from omnidreams.interactive_drive.rasterizer import ( # noqa: PLC0415 + LudusConditionRasterizer, + ) + + return LudusConditionRasterizer(_ludus_raster_config(scenario), bev=None) + + +def _ludus_raster_config(scenario: OmnidreamsLudusReplayScenario) -> Any: + from omnidreams.interactive_drive.config import RasterConfig # noqa: PLC0415 + + return RasterConfig( + width=scenario.pixel_width, + height=scenario.pixel_height, + ludus_backend=scenario.ludus_backend, + ) + + +def _initial_rgb_tensor( + frame: object, + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + tensor = torch.from_numpy(_rgb_hwc_uint8(frame)) + tensor = tensor.permute(2, 0, 1).unsqueeze(0).unsqueeze(0).unsqueeze(2) + return _to_model_range(tensor, device=device, dtype=dtype) + + +def _condition_frames_tensor( + frames: tuple[object, ...], + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + cuda_video = _condition_cuda_video(frames) + if cuda_video is not None: + tensor = cuda_video.permute(0, 3, 1, 2).unsqueeze(0).unsqueeze(0) + return _to_model_range(tensor, device=device, dtype=dtype) + video = np.stack( + [_rgb_hwc_uint8(_frame_rgb(frame)) for frame in frames], + axis=0, + ) + tensor = torch.from_numpy(np.ascontiguousarray(video)) + tensor = tensor.permute(0, 3, 1, 2).unsqueeze(0).unsqueeze(0) + return _to_model_range(tensor, device=device, dtype=dtype) + + +def _condition_cuda_video(frames: tuple[object, ...]) -> torch.Tensor | None: + tensors: list[torch.Tensor] = [] + for frame in frames: + to_cuda_tensor = getattr(_frame_rgb(frame), "to_cuda_tensor", None) + if not callable(to_cuda_tensor): + return None + try: + tensor = to_cuda_tensor() + except RuntimeError: + return None + if ( + not torch.is_tensor(tensor) + or not tensor.is_cuda + or tensor.dtype != torch.uint8 + or tensor.ndim != 3 + or tensor.shape[-1] < 3 + ): + return None + tensors.append(tensor[..., :3]) + return torch.stack(tensors, dim=0) + + +def _frame_rgb(frame: object) -> object: + return getattr(frame, "rgb_host_uint8", frame) + + +def _rgb_hwc_uint8(frame: object) -> np.ndarray: + if torch.is_tensor(frame): + array = frame.detach().cpu().numpy() + else: + array = np.asarray(frame, dtype=np.uint8) + if array.ndim != 3 or array.shape[-1] < 3: + raise ValueError( + "OmniDreams Ludus rendered frames must be HWC RGB/RGBA uint8 arrays." + ) + return np.ascontiguousarray(np.array(array[..., :3], dtype=np.uint8, copy=True)) + + +def _to_model_range( + tensor: torch.Tensor, + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + return tensor.to(device=device, dtype=dtype) / 127.5 - 1.0 + + +def _segments_metadata( + segments: list[PoseSegment], +) -> tuple[tuple[float, float, tuple[str, ...]], ...]: + return tuple( + (float(start), float(end), tuple(sorted(keys))) for start, end, keys in segments + ) + + +def _pose_segment(value: object) -> PoseSegment: + if not isinstance(value, tuple) or len(value) != 3: + raise RuntimeError("OmniDreams Ludus key segment must be a 3-tuple.") + start, end, keys = value + if not isinstance(start, int | float) or not isinstance(end, int | float): + raise RuntimeError("OmniDreams Ludus key segment bounds must be numeric.") + if not isinstance(keys, frozenset | set | tuple | list): + raise RuntimeError("OmniDreams Ludus key segment keys must be a sequence.") + return (float(start), float(end), frozenset(str(key) for key in keys)) + + +def _close_rasterizer(rasterizer: Any | None) -> None: + if rasterizer is None: + return + close = getattr(rasterizer, "cleanup", None) or getattr( + rasterizer, + "close", + None, + ) + if callable(close): + close() + + +def _device_from_config(config: InferenceConfig) -> torch.device: + if dist.is_initialized(): + return torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}") + return torch.device(config.device or "cuda") + + +def _is_rank_zero() -> bool: + return not dist.is_initialized() or dist.get_rank() == 0 + + +__all__ = [ + "LudusSceneConditioningProvider", + "PrecomputedHDMapProvider", + "keyboard_driving_user_input_schema", + "precomputed_hdmap_inference_input_schema", +] diff --git a/integrations/omnidreams/omnidreams/demo/replay.py b/integrations/omnidreams/omnidreams/demo/replay.py index 8ccb58650..b6c6d91e2 100644 --- a/integrations/omnidreams/omnidreams/demo/replay.py +++ b/integrations/omnidreams/omnidreams/demo/replay.py @@ -1,246 +1,29 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""OmniDreams replay runtime for the shared demo runner.""" +"""Compatibility aliases for the OmniDreams runtime module.""" from __future__ import annotations -import os -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 -from flashdreams.infra.postprocess import VideoTensorLayout -from flashdreams.infra.runner_io import ( - DEFAULT_RUNNER_INSTALL_HINT, - load_first_frame_tensor, +from .runtime import ( + OmnidreamsRuntime, + OmnidreamsRuntimeOptions, + OmnidreamsSession, + OmnidreamsSessionScenario, + PipelineFactory, ) -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 -from flashdreams.runtime.types import StepRequest, StepResult - -from .spec import OmnidreamsReplayScenario - -PipelineFactory = Callable[[Any, str], Any] - - -@dataclass(frozen=True, kw_only=True, slots=True) -class OmnidreamsReplayRuntimeOptions: - """Construction knobs for the replay runtime.""" - - pipeline_config: Any - pipeline_factory: PipelineFactory | None = None - output_layout: VideoTensorLayout = "bvtchw" - - -class OmnidreamsReplayRuntime: - """Heavyweight OmniDreams runtime consumed by ``run_inference_session``.""" - - def __init__( - self, - *, - config: InferenceConfig, - options: OmnidreamsReplayRuntimeOptions, - ) -> None: - self.config = config - self.options = options - if _is_torchrun_env() and not dist.is_initialized(): - init_distributed() - - if dist.is_initialized(): - self.local_rank = int(os.environ.get("LOCAL_RANK", "0")) - self.world_size = dist.get_world_size() - self.global_rank = dist.get_rank() - device = f"cuda:{self.local_rank}" - else: - self.local_rank = 0 - self.world_size = 1 - self.global_rank = 0 - device = config.device or "cuda" - - self.is_rank_zero = self.global_rank == 0 - factory = options.pipeline_factory or _default_pipeline_factory - self.pipeline = factory(options.pipeline_config, device) - - def start_session(self, inputs: InferenceInput) -> InferenceSession: - scenario = _scenario_from_inputs(inputs) - return OmnidreamsReplaySession( - pipeline=self.pipeline, - scenario=scenario, - device=torch.device(f"cuda:{self.local_rank}") - if dist.is_initialized() - else torch.device(self.config.device or "cuda"), - is_rank_zero=self.is_rank_zero, - output_layout=self.options.output_layout, - ) - - def close(self) -> None: - pipeline = getattr(self, "pipeline", None) - if pipeline is not None: - close = getattr(pipeline, "close", None) - if callable(close): - close() - del self.pipeline - device = torch.device(self.config.device or "cuda") - if device.type == "cuda" and torch.cuda.is_available(): - torch.cuda.empty_cache() - - -class OmnidreamsReplaySession: - """One MP4 replay rollout over a prepared scenario.""" - - def __init__( - self, - *, - pipeline: Any, - scenario: OmnidreamsReplayScenario, - device: torch.device, - is_rank_zero: bool, - output_layout: VideoTensorLayout, - ) -> None: - self.pipeline = pipeline - self.scenario = scenario - self.device = device - self.is_rank_zero = is_rank_zero - self.output_layout = output_layout - self.dtype = torch.bfloat16 - self._closed = False - self._frame_start = 0 - 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) - if dist.is_initialized(): - dist.barrier() - - def next_step_request(self) -> StepRequest | None: - if self._closed: - return None - step_index = self._model_session.step_index - if step_index >= self.scenario.total_blocks: - return None - 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=step_index) - - def step(self, inputs: InferenceInput) -> StepResult: - del inputs - if self._closed: - raise RuntimeError("OmniDreams replay session is closed.") - - 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=[{}, {})", - step_index, - self._frame_start, - frame_end, - ) - result = self._model_session.step( - self._hdmap_videos[:, :, self._frame_start : frame_end] - ) - self._frame_start = frame_end - 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.") - self._model_session.reset(self._initialize_cache) - self._frame_start = 0 - - def close(self) -> None: - self._closed = True - self._model_session.close() - - def _initialize_cache(self) -> Any: - scenario = self.scenario - first_frames = [ - load_first_frame_tensor( - path, - pixel_height=scenario.pixel_height, - pixel_width=scenario.pixel_width, - device=self.device, - dtype=self.dtype, - allow_video=True, - install_hint=DEFAULT_RUNNER_INSTALL_HINT, - ) - for path in scenario.first_frame_paths - ] - first_frames_t = torch.stack(first_frames, dim=0).unsqueeze(0) - cache = self.pipeline.initialize_cache( - text=[list(scenario.prompts)], - image=first_frames_t, - view_names=list(scenario.camera_names), - ) - release = getattr(self.pipeline, "release_oneshot_encoders", None) - if callable(release): - release() - return cache - - def _load_hdmaps(self) -> torch.Tensor: - scenario = self.scenario - videos = [ - _load_video( - path, - pixel_height=scenario.pixel_height, - pixel_width=scenario.pixel_width, - device=self.device, - dtype=self.dtype, - ) - for path in scenario.hdmap_video_paths - ] - # [B=1, V, T, C, H, W] - hdmap_videos = torch.stack(videos, dim=0).unsqueeze(0) - if self.is_rank_zero: - logger.info( - "Loaded OmniDreams demo HDMaps shape={} views={}", - tuple(hdmap_videos.shape), - len(scenario.camera_names), - ) - return hdmap_videos - - -def _default_pipeline_factory(pipeline_config: Any, device: str) -> Any: - return pipeline_config.setup().to(device=device).eval() - - -def _scenario_from_inputs(inputs: InferenceInput) -> OmnidreamsReplayScenario: - scenario = inputs.global_conditioning.get("scenario") - if not isinstance(scenario, OmnidreamsReplayScenario): - raise TypeError( - "OmniDreams replay runtime requires global_conditioning['scenario'] " - "to be an OmnidreamsReplayScenario." - ) - return scenario - - -def _is_torchrun_env() -> bool: - return "RANK" in os.environ and "WORLD_SIZE" in os.environ +OmnidreamsReplayRuntimeOptions = OmnidreamsRuntimeOptions +OmnidreamsReplayRuntime = OmnidreamsRuntime +OmnidreamsReplaySession = OmnidreamsSession __all__ = [ "OmnidreamsReplayRuntime", "OmnidreamsReplayRuntimeOptions", "OmnidreamsReplaySession", + "OmnidreamsRuntime", + "OmnidreamsRuntimeOptions", + "OmnidreamsSession", + "OmnidreamsSessionScenario", "PipelineFactory", ] diff --git a/integrations/omnidreams/omnidreams/demo/runtime.py b/integrations/omnidreams/omnidreams/demo/runtime.py new file mode 100644 index 000000000..1ee99ad31 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/runtime.py @@ -0,0 +1,353 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OmniDreams runtime/session contracts for shared demo run modes.""" + +from __future__ import annotations + +import os +from collections.abc import Callable, Sequence +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 flashdreams.core.distributed import init as init_distributed +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.infra.runner_io import ( + DEFAULT_RUNNER_INSTALL_HINT, + load_first_frame_tensor, +) +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 +from flashdreams.runtime.types import StepRequest, StepRequirements, StepResult + +from .spec import OmnidreamsLudusReplayScenario, OmnidreamsReplayScenario + +OmnidreamsSessionScenario = OmnidreamsReplayScenario | OmnidreamsLudusReplayScenario + +PipelineFactory = Callable[[Any, str], Any] + + +@dataclass(frozen=True, kw_only=True, slots=True) +class OmnidreamsRuntimeOptions: + """Construction knobs for the OmniDreams runtime.""" + + pipeline_config: Any + pipeline_factory: PipelineFactory | None = None + output_layout: VideoTensorLayout = "bvtchw" + release_oneshot_encoders_after_cache_init: bool = True + + +class OmnidreamsRuntime: + """Heavyweight OmniDreams runtime consumed by shared demo run modes.""" + + def __init__( + self, + *, + config: InferenceConfig, + options: OmnidreamsRuntimeOptions, + ) -> None: + self.config = config + self.options = options + if _is_torchrun_env() and not dist.is_initialized(): + init_distributed() + + if dist.is_initialized(): + self.local_rank = int(os.environ.get("LOCAL_RANK", "0")) + self.world_size = dist.get_world_size() + self.global_rank = dist.get_rank() + device = f"cuda:{self.local_rank}" + else: + self.local_rank = 0 + self.world_size = 1 + self.global_rank = 0 + device = config.device or "cuda" + + self.is_rank_zero = self.global_rank == 0 + factory = options.pipeline_factory or _default_pipeline_factory + self.pipeline = factory(options.pipeline_config, device) + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + scenario = _scenario_from_inputs(inputs) + return OmnidreamsSession( + pipeline=self.pipeline, + scenario=scenario, + initial_inputs=inputs, + device=torch.device(f"cuda:{self.local_rank}") + if dist.is_initialized() + else torch.device(self.config.device or "cuda"), + is_rank_zero=self.is_rank_zero, + output_layout=self.options.output_layout, + rollout_seed=self.config.seed, + release_oneshot_encoders_after_cache_init=( + self.options.release_oneshot_encoders_after_cache_init + ), + ) + + def close(self) -> None: + pipeline = getattr(self, "pipeline", None) + if pipeline is not None: + close = getattr(pipeline, "close", None) + if callable(close): + close() + del self.pipeline + device = torch.device(self.config.device or "cuda") + if device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.empty_cache() + + +class OmnidreamsSession: + """One OmniDreams rollout over a prepared scenario.""" + + def __init__( + self, + *, + pipeline: Any, + scenario: OmnidreamsSessionScenario, + initial_inputs: InferenceInput, + device: torch.device, + is_rank_zero: bool, + output_layout: VideoTensorLayout, + rollout_seed: int | None, + release_oneshot_encoders_after_cache_init: bool, + ) -> None: + self.pipeline = pipeline + self.scenario = scenario + self._initial_inputs = initial_inputs + self.device = device + self.is_rank_zero = is_rank_zero + self.output_layout = output_layout + self.rollout_seed = rollout_seed + self.release_oneshot_encoders_after_cache_init = ( + release_oneshot_encoders_after_cache_init + ) + self.dtype = torch.bfloat16 + self._closed = False + 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) + if self.device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.synchronize(device=self.device) + if dist.is_initialized(): + dist.barrier() + + def next_step_requirements(self) -> StepRequirements | None: + if self._closed: + return None + step_index = self._model_session.step_index + if step_index >= self.scenario.total_blocks: + return None + num_frames = self._model_session.next_num_frames() + return StepRequirements( + step_index=step_index, + input_frame_count=num_frames, + ) + + def next_step_request(self) -> StepRequest | None: + requirements = self.next_step_requirements() + if requirements is None: + return None + metadata = dict(requirements.metadata) + metadata["input_frame_count"] = requirements.input_frame_count + if requirements.steady_output_frame_count is not None: + metadata["steady_output_frame_count"] = ( + requirements.steady_output_frame_count + ) + return StepRequest( + step_index=requirements.step_index, + inference_input_schema=requirements.inference_input_schema, + metadata=metadata, + ) + + def step(self, inputs: InferenceInput) -> StepResult: + if self._closed: + raise RuntimeError("OmniDreams replay session is closed.") + + step_index = self._model_session.step_index + num_frames = self._model_session.next_num_frames() + hdmap = _hdmap_from_inputs(inputs) + if hdmap.shape[2] != num_frames: + raise ValueError( + "OmniDreams step HDMap frame count mismatch: " + f"expected {num_frames}, got {hdmap.shape[2]}." + ) + logger.info( + "OmniDreams demo replay step {} frames={}", + step_index, + num_frames, + ) + return self._model_session.step(hdmap) + + 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.") + self._initial_inputs = inputs + self._model_session.reset(self._initialize_cache) + + def close(self) -> None: + self._closed = True + self._model_session.close() + + def _initialize_cache(self) -> Any: + scenario = self.scenario + _seed_pipeline_for_rollout(self.pipeline, self.rollout_seed) + cache = self.pipeline.initialize_cache( + text=_prompt_from_inputs(self._initial_inputs, scenario), + image=_first_frame_from_inputs( + self._initial_inputs, + scenario=scenario, + device=self.device, + dtype=self.dtype, + ), + view_names=_view_names_from_inputs(self._initial_inputs, scenario), + ) + if self.release_oneshot_encoders_after_cache_init: + release = getattr(self.pipeline, "release_oneshot_encoders", None) + if callable(release): + release() + return cache + + +def _default_pipeline_factory(pipeline_config: Any, device: str) -> Any: + return pipeline_config.setup().to(device=device).eval() + + +def _scenario_from_inputs(inputs: InferenceInput) -> OmnidreamsSessionScenario: + scenario = inputs.global_conditioning.get("scenario") + if not isinstance( + scenario, + (OmnidreamsReplayScenario, OmnidreamsLudusReplayScenario), + ): + raise TypeError( + "OmniDreams replay runtime requires global_conditioning['scenario'] " + "to be an OmnidreamsReplayScenario or OmnidreamsLudusReplayScenario." + ) + return scenario + + +def _prompt_from_inputs( + inputs: InferenceInput, + scenario: OmnidreamsSessionScenario, +) -> list[list[str]]: + prompt = inputs.global_conditioning.get("prompt") + if prompt is None: + if not scenario.prompts: + raise ValueError( + "OmniDreams initial prompt is required when the scenario does " + "not carry fallback prompts." + ) + return [list(scenario.prompts)] + if isinstance(prompt, str): + return [[prompt]] + if isinstance(prompt, Sequence): + values = list(prompt) + if all(isinstance(value, str) for value in values): + return [[str(value) for value in values]] + batches: list[list[str]] = [] + for batch in values: + if not isinstance(batch, Sequence) or isinstance(batch, str): + raise TypeError( + "OmniDreams initial prompt batches must be string sequences." + ) + batches.append([str(item) for item in batch]) + return batches + raise TypeError( + "OmniDreams initial prompt must be a string or sequence of strings." + ) + + +def _first_frame_from_inputs( + inputs: InferenceInput, + *, + scenario: OmnidreamsSessionScenario, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + first_frame = inputs.global_conditioning.get("first_frame") + if isinstance(first_frame, torch.Tensor): + return first_frame + if first_frame is not None: + raise TypeError("OmniDreams initial first_frame must be a torch.Tensor.") + first_frame_paths = getattr(scenario, "first_frame_paths", ()) + if not first_frame_paths: + raise ValueError( + "OmniDreams initial first_frame tensor is required when the " + "scenario does not carry fallback first_frame_paths." + ) + first_frames = [ + load_first_frame_tensor( + path, + pixel_height=scenario.pixel_height, + pixel_width=scenario.pixel_width, + device=device, + dtype=dtype, + allow_video=True, + install_hint=DEFAULT_RUNNER_INSTALL_HINT, + ) + for path in first_frame_paths + ] + return torch.stack(first_frames, dim=0).unsqueeze(0) + + +def _seed_pipeline_for_rollout(pipeline: Any, seed: int | None) -> None: + if seed is None: + return + diffusion_model = getattr(pipeline, "diffusion_model", None) + rng = getattr(diffusion_model, "rng", None) + if rng is None: + return + rng.manual_seed(int(seed)) + + +def _view_names_from_inputs( + inputs: InferenceInput, + scenario: OmnidreamsSessionScenario, +) -> list[str]: + value = inputs.metadata.get("view_names") or inputs.global_conditioning.get( + "view_names" + ) + if value is None: + return list(scenario.camera_names) + if isinstance(value, str): + return [value] + if isinstance(value, Sequence): + return [str(item) for item in value] + raise TypeError("OmniDreams view_names metadata must be a string sequence.") + + +def _hdmap_from_inputs(inputs: InferenceInput) -> torch.Tensor: + hdmap = inputs.step.get("hdmap") + if not isinstance(hdmap, torch.Tensor): + raise TypeError("OmniDreams session step requires step['hdmap'] tensor.") + if hdmap.ndim != 6: + raise ValueError( + "OmniDreams step['hdmap'] must have shape [B, V, T, C, H, W], " + f"got {tuple(hdmap.shape)}." + ) + return hdmap + + +def _is_torchrun_env() -> bool: + return "RANK" in os.environ and "WORLD_SIZE" in os.environ + + +__all__ = [ + "OmnidreamsRuntime", + "OmnidreamsRuntimeOptions", + "OmnidreamsSession", + "OmnidreamsSessionScenario", + "PipelineFactory", +] diff --git a/integrations/omnidreams/omnidreams/demo/spec.py b/integrations/omnidreams/omnidreams/demo/spec.py index 6a4f147cf..8df653ecd 100644 --- a/integrations/omnidreams/omnidreams/demo/spec.py +++ b/integrations/omnidreams/omnidreams/demo/spec.py @@ -5,10 +5,12 @@ from __future__ import annotations +import json +import math from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path -from typing import Any +from typing import Any, Literal, TypeAlias, cast from omnidreams.runner import ( DEFAULT_EXAMPLE_DATA_UUID_1V, @@ -22,6 +24,43 @@ DEFAULT_OMNIDREAMS_PRESET = "omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae" OMNIDREAMS_MODEL_ID = "omnidreams" DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID = "0d404ff7-2b66-498c-b047-1ed8cded60d4" +OMNIDREAMS_CONDITIONING_PRECOMPUTED = "precomputed-hdmap" +OMNIDREAMS_CONDITIONING_LUDUS = "ludus-scene-driving" +OMNIDREAMS_CONDITIONING_MODES = ( + OMNIDREAMS_CONDITIONING_PRECOMPUTED, + OMNIDREAMS_CONDITIONING_LUDUS, +) +LudusBackendName: TypeAlias = Literal["cuda", "vulkan"] + +_KEY_EVENT_ALIASES = { + "down": "keydown", + "key_down": "keydown", + "keyup": "keyup", + "up": "keyup", + "key_up": "keyup", +} + + +@dataclass(frozen=True, kw_only=True, slots=True) +class OmnidreamsKeyboardTraceEvent: + """One recorded keyboard edge in a finite Ludus replay trace.""" + + timestamp_s: float + event: str + key: str + + def __post_init__(self) -> None: + timestamp_s = float(self.timestamp_s) + if not math.isfinite(timestamp_s) or timestamp_s < 0: + raise ValueError( + "OmnidreamsKeyboardTraceEvent.timestamp_s must be finite and >= 0." + ) + key = str(self.key).strip().lower() + if not key: + raise ValueError("OmnidreamsKeyboardTraceEvent.key must be non-empty.") + object.__setattr__(self, "timestamp_s", timestamp_s) + object.__setattr__(self, "event", _normalize_key_event_name(self.event)) + object.__setattr__(self, "key", key) @dataclass(frozen=True, kw_only=True, slots=True) @@ -69,6 +108,86 @@ def __post_init__(self) -> None: ) +@dataclass(frozen=True, kw_only=True, slots=True) +class OmnidreamsLudusReplayScenario: + """Resolved Ludus scene plus a finite recorded keyboard trace.""" + + keyboard_events: tuple[OmnidreamsKeyboardTraceEvent, ...] + scene_path: Path | None = None + scene_dir: Path | None = None + scene_uuid: str | None = DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID + scene_variant: str = SCENE_VARIANT_DEFAULT + camera_name: str = "camera_front_wide_120fov" + prompt: str | None = None + total_blocks: int = 60 + pixel_height: int = DEFAULT_VIDEO_HEIGHT + pixel_width: int = DEFAULT_VIDEO_WIDTH + fps: int = 30 + move_speed_per_s: float = 6.0 + rotate_speed_rad_per_s: float = math.radians(35.0) + ludus_backend: LudusBackendName = "cuda" + + @property + def camera_names(self) -> tuple[str, ...]: + return (self.camera_name,) + + @property + def prompts(self) -> tuple[str, ...]: + return () if self.prompt is None else (self.prompt,) + + def __post_init__(self) -> None: + if self.scene_path is not None: + object.__setattr__(self, "scene_path", Path(self.scene_path)) + if self.scene_dir is not None: + object.__setattr__(self, "scene_dir", Path(self.scene_dir)) + if not (self.scene_path or self.scene_dir or self.scene_uuid): + raise ValueError( + "OmnidreamsLudusReplayScenario requires scene_path, " + "scene_dir, or scene_uuid." + ) + if not self.scene_variant.strip(): + raise ValueError("OmnidreamsLudusReplayScenario.scene_variant is required.") + if not self.camera_name.strip(): + raise ValueError("OmnidreamsLudusReplayScenario.camera_name is required.") + if self.total_blocks <= 0: + raise ValueError("OmnidreamsLudusReplayScenario.total_blocks must be > 0.") + if self.pixel_height <= 0 or self.pixel_width <= 0: + raise ValueError( + "OmnidreamsLudusReplayScenario pixel dimensions must be > 0." + ) + if self.fps <= 0: + raise ValueError("OmnidreamsLudusReplayScenario.fps must be > 0.") + if self.move_speed_per_s <= 0: + raise ValueError( + "OmnidreamsLudusReplayScenario.move_speed_per_s must be > 0." + ) + if self.rotate_speed_rad_per_s <= 0: + raise ValueError( + "OmnidreamsLudusReplayScenario.rotate_speed_rad_per_s must be > 0." + ) + object.__setattr__( + self, + "ludus_backend", + _ludus_backend_name(self.ludus_backend), + ) + previous_timestamp_s = -math.inf + normalized_events: list[OmnidreamsKeyboardTraceEvent] = [] + for event in self.keyboard_events: + normalized = ( + event + if isinstance(event, OmnidreamsKeyboardTraceEvent) + else _keyboard_trace_event(event) + ) + if normalized.timestamp_s < previous_timestamp_s: + raise ValueError( + "OmnidreamsLudusReplayScenario.keyboard_events must be sorted " + "by non-decreasing timestamp_s." + ) + previous_timestamp_s = normalized.timestamp_s + normalized_events.append(normalized) + object.__setattr__(self, "keyboard_events", tuple(normalized_events)) + + @dataclass(frozen=True, kw_only=True, slots=True) class OmnidreamsWebRTCScenario: """Scene/options for the shared WebRTC demo path.""" @@ -89,6 +208,33 @@ def __post_init__(self) -> None: raise ValueError("OmnidreamsWebRTCScenario.camera_name is required.") +def conditioning_mode_from_scenario(value: Any) -> str: + """Return the resolved OmniDreams replay conditioning mode.""" + if isinstance(value, OmnidreamsLudusReplayScenario): + return OMNIDREAMS_CONDITIONING_LUDUS + if isinstance(value, OmnidreamsReplayScenario): + return OMNIDREAMS_CONDITIONING_PRECOMPUTED + if value is None or not isinstance(value, Mapping): + return OMNIDREAMS_CONDITIONING_PRECOMPUTED + + mode = ( + str(value.get("conditioning_mode", OMNIDREAMS_CONDITIONING_PRECOMPUTED)) + .strip() + .lower() + ) + if mode in {"precomputed", "hdmap", "precomputed-hdmaps"}: + mode = OMNIDREAMS_CONDITIONING_PRECOMPUTED + if mode in {"ludus", "keyboard-driving", "ludus-keyboard"}: + mode = OMNIDREAMS_CONDITIONING_LUDUS + if mode not in OMNIDREAMS_CONDITIONING_MODES: + supported = ", ".join(OMNIDREAMS_CONDITIONING_MODES) + raise ValueError( + f"Unsupported OmniDreams conditioning_mode={mode!r}. " + f"Supported modes: {supported}." + ) + return mode + + def resolve_replay_scenario( value: Any, *, @@ -151,6 +297,40 @@ def resolve_replay_scenario( ) +def resolve_ludus_replay_scenario(value: Any) -> OmnidreamsLudusReplayScenario: + """Normalize a user/demo scenario into a Ludus recorded-trace scenario.""" + if isinstance(value, OmnidreamsLudusReplayScenario): + _require_optional_existing_path(value.scene_path, label="scene_path") + return value + if value is None: + value = {} + if not isinstance(value, Mapping): + raise TypeError( + "OmniDreams Ludus replay scenario must be an " + "OmnidreamsLudusReplayScenario, a mapping, or None." + ) + return OmnidreamsLudusReplayScenario( + keyboard_events=_keyboard_trace_events(value), + scene_path=_optional_path(value.get("scene_path")), + scene_dir=_optional_path(value.get("scene_dir")), + scene_uuid=_optional_string( + 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")), + prompt=_optional_string(value.get("prompt")), + total_blocks=int(value.get("total_blocks", 60)), + pixel_height=int(value.get("pixel_height", DEFAULT_VIDEO_HEIGHT)), + pixel_width=int(value.get("pixel_width", DEFAULT_VIDEO_WIDTH)), + fps=int(value.get("fps", 30)), + move_speed_per_s=float(value.get("move_speed_per_s", 6.0)), + rotate_speed_rad_per_s=float( + value.get("rotate_speed_rad_per_s", math.radians(35.0)) + ), + ludus_backend=_ludus_backend_name(value.get("ludus_backend", "cuda")), + ) + + def resolve_webrtc_scenario(value: Any) -> OmnidreamsWebRTCScenario: """Normalize a user/demo scenario into a WebRTC scenario.""" if value is None: @@ -205,6 +385,68 @@ def _resolve_example_data_default(value: Mapping[str, Any]) -> bool: ) +def _keyboard_trace_events( + value: Mapping[str, Any], +) -> tuple[OmnidreamsKeyboardTraceEvent, ...]: + events_value = value.get("keyboard_events") + if events_value is None: + trace_path = _optional_path(value.get("keyboard_trace_path")) + if trace_path is None: + return () + if not trace_path.exists(): + raise FileNotFoundError( + f"OmniDreams keyboard_trace_path missing: {trace_path}" + ) + loaded = json.loads(trace_path.read_text(encoding="utf-8")) + events_value = ( + loaded.get("events", ()) if isinstance(loaded, Mapping) else loaded + ) + if isinstance(events_value, (str, bytes)) or not isinstance(events_value, Sequence): + raise TypeError("OmniDreams keyboard trace must be a sequence of events.") + return tuple(_keyboard_trace_event(event) for event in events_value) + + +def _keyboard_trace_event(value: Any) -> OmnidreamsKeyboardTraceEvent: + if isinstance(value, OmnidreamsKeyboardTraceEvent): + return value + if not isinstance(value, Mapping): + raise TypeError( + "OmniDreams keyboard trace events must be mappings or " + "OmnidreamsKeyboardTraceEvent instances." + ) + timestamp = _first_present(value, ("timestamp_s", "time_s", "timestamp", "t")) + if timestamp is None: + raise ValueError("OmniDreams keyboard trace event missing timestamp_s.") + event = _first_present(value, ("event", "event_type", "type")) + if event is None: + raise ValueError("OmniDreams keyboard trace event missing event.") + key = value.get("key") + if key is None: + raise ValueError("OmniDreams keyboard trace event missing key.") + return OmnidreamsKeyboardTraceEvent( + timestamp_s=float(timestamp), + event=str(event), + key=str(key), + ) + + +def _first_present(value: Mapping[str, Any], keys: tuple[str, ...]) -> Any: + for key in keys: + if key in value: + return value[key] + return None + + +def _normalize_key_event_name(value: str) -> str: + event = str(value).strip().lower() + event = _KEY_EVENT_ALIASES.get(event, event) + if event not in {"keydown", "keyup"}: + raise ValueError( + "OmniDreams keyboard trace event must be 'keydown' or 'keyup'." + ) + return event + + def _bool_value(value: Any) -> bool: if isinstance(value, bool): return value @@ -248,6 +490,37 @@ def _string_tuple(value: Any) -> tuple[str, ...]: raise TypeError(f"Expected string or string sequence, got {type(value).__name__}.") +def _optional_path(value: Any) -> Path | None: + if value is None or value == "": + return None + return Path(value) + + +def _optional_string(value: Any) -> str | None: + if value is None: + return None + text = str(value).strip() + return text or None + + +def _ludus_backend_name(value: Any) -> LudusBackendName: + backend = str(value).strip().lower() + if backend not in {"cuda", "vulkan"}: + raise ValueError( + "OmnidreamsLudusReplayScenario.ludus_backend must be 'cuda' or 'vulkan'." + ) + return cast(LudusBackendName, backend) + + +def _require_optional_existing_path(path: Path | None, *, label: str) -> None: + if path is None: + return + if not path.exists(): + raise FileNotFoundError( + f"OmniDreams Ludus replay scenario missing {label}: {path}" + ) + + def _require_existing_paths(paths: tuple[Path, ...], *, label: str) -> None: if not paths: raise ValueError(f"OmniDreams replay scenario requires {label}.") @@ -262,9 +535,16 @@ def _require_existing_paths(paths: tuple[Path, ...], *, label: str) -> None: __all__ = [ "DEFAULT_OMNIDREAMS_PRESET", "DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID", + "OMNIDREAMS_CONDITIONING_LUDUS", + "OMNIDREAMS_CONDITIONING_MODES", + "OMNIDREAMS_CONDITIONING_PRECOMPUTED", "OMNIDREAMS_MODEL_ID", + "OmnidreamsKeyboardTraceEvent", + "OmnidreamsLudusReplayScenario", "OmnidreamsReplayScenario", "OmnidreamsWebRTCScenario", + "conditioning_mode_from_scenario", + "resolve_ludus_replay_scenario", "resolve_replay_scenario", "resolve_webrtc_scenario", ] diff --git a/integrations/omnidreams/omnidreams/demo/traces/ludus_forward_sweep_60s.json b/integrations/omnidreams/omnidreams/demo/traces/ludus_forward_sweep_60s.json new file mode 100644 index 000000000..0f66222b6 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/traces/ludus_forward_sweep_60s.json @@ -0,0 +1,86 @@ +{ + "name": "ludus_forward_sweep_60s", + "description": "Deterministic WSAD trace for OmniDreams Ludus replay MP4 validation.", + "events": [ + { + "timestamp_s": 0.0, + "event": "keydown", + "key": "w" + }, + { + "timestamp_s": 6.0, + "event": "keydown", + "key": "d" + }, + { + "timestamp_s": 10.0, + "event": "keyup", + "key": "d" + }, + { + "timestamp_s": 14.0, + "event": "keydown", + "key": "a" + }, + { + "timestamp_s": 18.0, + "event": "keyup", + "key": "a" + }, + { + "timestamp_s": 24.0, + "event": "keyup", + "key": "w" + }, + { + "timestamp_s": 25.0, + "event": "keydown", + "key": "w" + }, + { + "timestamp_s": 30.0, + "event": "keydown", + "key": "d" + }, + { + "timestamp_s": 34.0, + "event": "keyup", + "key": "d" + }, + { + "timestamp_s": 38.0, + "event": "keydown", + "key": "a" + }, + { + "timestamp_s": 42.0, + "event": "keyup", + "key": "a" + }, + { + "timestamp_s": 48.0, + "event": "keyup", + "key": "w" + }, + { + "timestamp_s": 50.0, + "event": "keydown", + "key": "w" + }, + { + "timestamp_s": 55.0, + "event": "keydown", + "key": "d" + }, + { + "timestamp_s": 58.0, + "event": "keyup", + "key": "d" + }, + { + "timestamp_s": 60.0, + "event": "keyup", + "key": "w" + } + ] +} diff --git a/integrations/omnidreams/omnidreams/demo/web/adapter.css b/integrations/omnidreams/omnidreams/demo/web/adapter.css new file mode 100644 index 000000000..f4cb8868b --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/web/adapter.css @@ -0,0 +1,16 @@ +/* +SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 +*/ + +/* Keep the browser from enlarging OmniDreams' native 1280x704 stream. */ +.stageVideo { + inset: 50% auto auto 50%; + width: min(100vw, 1280px, 181.82vh); + height: auto; + max-height: min(100vh, 704px); + aspect-ratio: 1280 / 704; + transform: translate(-50%, -50%); + object-fit: contain; + object-position: center; +} diff --git a/integrations/omnidreams/omnidreams/demo/web/adapter.js b/integrations/omnidreams/omnidreams/demo/web/adapter.js index d07fb8cc1..37d19a299 100644 --- a/integrations/omnidreams/omnidreams/demo/web/adapter.js +++ b/integrations/omnidreams/omnidreams/demo/web/adapter.js @@ -3,6 +3,7 @@ export default { modelName: "OmniDreams", + stylesheet: "/model-static/adapter.css?v=model-ui-v2", controls: [ { label: "Drive / Turn", diff --git a/integrations/omnidreams/omnidreams/demo/webrtc.py b/integrations/omnidreams/omnidreams/demo/webrtc.py index 00857ea85..e5d8e2a84 100644 --- a/integrations/omnidreams/omnidreams/demo/webrtc.py +++ b/integrations/omnidreams/omnidreams/demo/webrtc.py @@ -13,458 +13,61 @@ # 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.""" +"""OmniDreams browser hooks for the shared WebRTC demo runtime.""" from __future__ import annotations -import tempfile -import time +import os from collections.abc import Callable -from dataclasses import dataclass, replace -from pathlib import Path +from dataclasses import replace 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 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 import InferenceConfig +from flashdreams.runtime.demo import ( + DemoSpec, + RuntimeHost, + 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.encoders import EncoderBackend +from flashdreams.serving.webrtc.controls import WSAD_SUPPORTED_KEYS from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager -from flashdreams.serving.webrtc.runtime import ThreadAffineDistributedWebRTCRuntime from flashdreams.serving.webrtc.server import create_webrtc_app +from .adapter import OmnidreamsDemoAdapter, RuntimeFactory from .spec import ( DEFAULT_OMNIDREAMS_PRESET, - DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID, OMNIDREAMS_MODEL_ID, resolve_webrtc_scenario, ) +from .webrtc_config import OmnidreamsWebRTCModelRuntimeConfig WebRTCRuntimeFactory = Callable[..., Any] - - -class OmnidreamsWebRTCModelRuntimeError(RuntimeError): - """Raised when the OmniDreams demo runtime is used incorrectly.""" - - -@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__( - 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." - ) - 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, - ) - - 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, - ) - - 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 _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 _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 +SharedRuntimeFactory = RuntimeFactory def serve_omnidreams_webrtc_demo( *, spec: DemoSpec, world_rank: int = 0, - runtime_factory: WebRTCRuntimeFactory = OmnidreamsWebRTCModelRuntime, + runtime_factory: WebRTCRuntimeFactory | None = None, + shared_runtime_factory: SharedRuntimeFactory | None = None, 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.""" + """Create OmniDreams' runtime and serve it through shared WebRTC transport.""" + if runtime_factory is not None and shared_runtime_factory is not None: + raise ValueError( + "Specify either legacy runtime_factory or shared_runtime_factory, not both." + ) if spec.input_mode != "keyboard-driving": raise ValueError( "OmniDreams WebRTC requires input_mode='keyboard-driving', " @@ -480,7 +83,49 @@ def serve_omnidreams_webrtc_demo( f"OmniDreams WebRTC requires model_id={OMNIDREAMS_MODEL_ID!r}, " f"got {config.model_id!r}." ) + scenario = resolve_webrtc_scenario(spec.scenario) + runtime_config = _webrtc_runtime_config( + output=spec.output, + config=config, + scenario=scenario, + ) + if _should_use_legacy_webrtc_path( + scenario=scenario, + runtime_factory=runtime_factory, + ): + from .webrtc_legacy import ( # noqa: PLC0415 + OmnidreamsWebRTCModelRuntime, + _serve_legacy_omnidreams_webrtc_demo, + ) + + return _serve_legacy_omnidreams_webrtc_demo( + spec=spec, + output=spec.output, + runtime_config=runtime_config, + runtime_factory=runtime_factory or OmnidreamsWebRTCModelRuntime, + world_rank=world_rank, + create_app_fn=create_app_fn, + server_runner=server_runner, + ) + + return _serve_shared_omnidreams_webrtc_demo( + spec=_shared_webrtc_spec(spec, runtime_config=runtime_config), + output=spec.output, + runtime_config=runtime_config, + shared_runtime_factory=shared_runtime_factory, + world_rank=world_rank, + create_app_fn=create_app_fn, + server_runner=server_runner, + ) + + +def _webrtc_runtime_config( + *, + output: WebRTCOutputSpec, + config: InferenceConfig, + scenario: Any, +) -> OmnidreamsWebRTCModelRuntimeConfig: preset_id = _preset_id(config) seed = _option(config, "seed", 42) runtime_config = OmnidreamsWebRTCModelRuntimeConfig( @@ -491,17 +136,57 @@ def serve_omnidreams_webrtc_demo( 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, + video_height=output.video_height, + video_width=output.video_width, + fps=output.fps, camera_name=scenario.camera_name, - warmup_chunks=spec.output.warmup_chunks, - warmup_timeout_s=spec.output.warmup_timeout_s, + warmup_chunks=output.warmup_chunks, + warmup_timeout_s=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) + return _apply_runtime_options(runtime_config, config.runtime_options) + + +def _should_use_legacy_webrtc_path( + *, + scenario: Any, + runtime_factory: WebRTCRuntimeFactory | None, +) -> bool: + if runtime_factory is not None: + return True + if bool(getattr(scenario, "debug_serve_hdmaps", False)): + logger.info( + "Using the legacy OmniDreams WebRTC path because debug HDMap " + "streaming is still implemented by the compatibility facade." + ) + return True + if _distributed_world_size() > 1: + logger.info( + "Using the legacy OmniDreams WebRTC path for multi-rank serving; " + "shared RuntimeHost distributed fan-out is not yet complete." + ) + return True + return False + + +def _serve_shared_omnidreams_webrtc_demo( + *, + spec: DemoSpec, + output: WebRTCOutputSpec, + runtime_config: OmnidreamsWebRTCModelRuntimeConfig, + shared_runtime_factory: SharedRuntimeFactory | None, + world_rank: int, + create_app_fn: CreateWebRTCApp, + server_runner: RunWebRTCServer, +) -> object: + adapter = OmnidreamsDemoAdapter(runtime_factory=shared_runtime_factory) + prepared = adapter.prepare_scenario(spec) + config = spec.config + if config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + runtime = adapter.create_runtime(config) + host = RuntimeHost(runtime) manager = BaseWebRTCSessionManager( runtime=runtime, runtime_config=runtime_config, @@ -511,12 +196,16 @@ def serve_omnidreams_webrtc_demo( warmup_label="OmniDreams WebRTC", supported_control_keys=WSAD_SUPPORTED_KEYS, fatal_generation_errors=True, - client_liveness_timeout_s=spec.output.client_liveness_timeout_s, + client_liveness_timeout_s=output.client_liveness_timeout_s, + shared_host=host, + shared_adapter=adapter, + shared_spec=spec, + shared_scenario=prepared, ) from importlib.resources import files return serve_webrtc_demo( - output=spec.output, + output=output, model_id=spec.model_id, session_manager=manager, app_resources=WebRTCAppResources( @@ -529,6 +218,36 @@ def serve_omnidreams_webrtc_demo( ) +def _shared_webrtc_spec( + spec: DemoSpec, + *, + runtime_config: OmnidreamsWebRTCModelRuntimeConfig, +) -> DemoSpec: + config = spec.config + if config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + runtime_options = dict(config.runtime_options) + runtime_options.update( + { + "pipeline_config": runtime_config.pipeline_config, + "seed": runtime_config.seed, + "move_speed_per_s": runtime_config.move_speed_per_s, + "rotate_speed_rad_per_s": runtime_config.rotate_speed_rad_per_s, + "release_oneshot_encoders_after_cache_init": False, + } + ) + return replace( + spec, + config=replace( + config, + preset_id=runtime_config.pipeline_config_name, + device=runtime_config.device, + seed=runtime_config.seed, + runtime_options=runtime_options, + ), + ) + + def _preset_id(config: InferenceConfig | None) -> str: return ( DEFAULT_OMNIDREAMS_PRESET @@ -556,6 +275,13 @@ def _option(config: InferenceConfig, name: str, default: Any) -> Any: return config.runtime_options.get(name, default) +def _distributed_world_size() -> int: + try: + return int(os.environ.get("WORLD_SIZE", "1")) + except ValueError: + return 1 + + def _apply_runtime_options( runtime_config: OmnidreamsWebRTCModelRuntimeConfig, options: Any, @@ -576,9 +302,8 @@ def _apply_runtime_options( __all__ = [ - "OmnidreamsWebRTCModelRuntime", "OmnidreamsWebRTCModelRuntimeConfig", - "OmnidreamsWebRTCModelRuntimeError", + "SharedRuntimeFactory", "WebRTCRuntimeFactory", "serve_omnidreams_webrtc_demo", ] diff --git a/integrations/omnidreams/omnidreams/demo/webrtc_config.py b/integrations/omnidreams/omnidreams/demo/webrtc_config.py new file mode 100644 index 000000000..272123cbd --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/webrtc_config.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared OmniDreams WebRTC runtime configuration.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from omnidreams.scenes import SCENE_VARIANT_DEFAULT + +from flashdreams.serving.webrtc.encoders import EncoderBackend + +from .runtime import PipelineFactory +from .spec import DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID + + +@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 = math.radians(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.""" + + pipeline_factory: PipelineFactory | None = None + """Optional test/runtime override for constructing the shared pipeline.""" + + +__all__ = ["OmnidreamsWebRTCModelRuntimeConfig"] diff --git a/integrations/omnidreams/omnidreams/demo/webrtc_legacy.py b/integrations/omnidreams/omnidreams/demo/webrtc_legacy.py new file mode 100644 index 000000000..0d7768636 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/webrtc_legacy.py @@ -0,0 +1,718 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Legacy OmniDreams WebRTC compatibility facade.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Any + +import torch +from loguru import logger +from omnidreams.transformer import CosmosTransformerConfig + +from flashdreams.core.distributed.rank_orchestration import distributed_op +from flashdreams.runtime import ( + CanonicalInputs, + CanonicalInputSchema, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InputCanonicalizer, + StepRequest, + StepRequirements, + StepResult, + TimeWindow, + step_requirements_from_request, +) +from flashdreams.runtime.demo import ( + DemoSpec, + PreparedScenario, + SessionInfo, + UserInputWindow, + WebRTCAppResources, + WebRTCOutputSpec, +) +from flashdreams.runtime.demo.timing import SPARSE_KEY_SEGMENTS_METADATA_KEY +from flashdreams.runtime.demo.webrtc import ( + CreateWebRTCApp, + RunWebRTCServer, + serve_webrtc_demo, +) +from flashdreams.serving.webrtc.controls import WSAD_SUPPORTED_KEYS, PoseSegment +from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager +from flashdreams.serving.webrtc.runtime import ( + ThreadAffineDistributedWebRTCRuntime, + WebRTCControlSignal, +) +from flashdreams.serving.webrtc.services import WEBRTC_USER_INPUT_SCHEMA + +from .providers import LudusSceneConditioningProvider +from .runtime import OmnidreamsRuntime, OmnidreamsRuntimeOptions +from .spec import ( + DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID, + OMNIDREAMS_MODEL_ID, + OmnidreamsLudusReplayScenario, +) +from .webrtc_config import OmnidreamsWebRTCModelRuntimeConfig + +WebRTCRuntimeFactory = Callable[..., Any] +_WEBRTC_SESSION_TOTAL_BLOCKS = 2_147_483_647 +_WEBRTC_STEP_REQUEST_KEY = "omnidreams_webrtc_step_request" + + +class OmnidreamsWebRTCModelRuntimeError(RuntimeError): + """Raised when the OmniDreams demo runtime is used incorrectly.""" + + +class OmnidreamsWebRTCModelRuntime( + ThreadAffineDistributedWebRTCRuntime[ + OmnidreamsWebRTCModelRuntimeConfig, + None, + ] +): + """Compatibility WebRTC facade over the shared OmniDreams runtime/session.""" + + def __init__(self, *, config: OmnidreamsWebRTCModelRuntimeConfig) -> None: + super().__init__( + config=config, + runtime_error_type=OmnidreamsWebRTCModelRuntimeError, + thread_name="omnidreams-demo-runtime", + ) + # The shared WebRTC input source emits normalized runtime events + # (``key_down``/``key_up``). The Ludus provider consumes sparse + # resampler metadata on this transitional path, so keep validation + # aligned with the WebRTC source rather than the replay trace schema. + self.input_source_schema = WEBRTC_USER_INPUT_SCHEMA + self.input_canonicalizer = InputCanonicalizer() + self.input_mapping = _OmnidreamsWebRTCInputMapping() + self._runtime: OmnidreamsRuntime | None = None + self._active_provider: LudusSceneConditioningProvider | None = None + self._active_session: Any | None = None + self._debug_session: _OmnidreamsHDMapDebugSession | None = None + self._steady_output_frame_count_value = 1 + + def _is_runtime_initialized(self) -> bool: + return self._runtime is not None + + def _runtime_step_index(self) -> int: + requirements = self._next_step_requirements_sync() + if requirements is None: + return 0 + return requirements.step_index + + def _next_input_frame_count(self) -> int: + requirements = self._next_step_requirements_sync() + if requirements is None: + raise OmnidreamsWebRTCModelRuntimeError( + "OmniDreams WebRTC session is complete." + ) + return requirements.input_frame_count + + def _steady_output_frame_count(self) -> int: + return self._steady_output_frame_count_value + + def _initialize_sync(self) -> None: + if self._runtime is not None: + return + if self._device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for OmniDreams WebRTC inference.") + _validate_single_view_pipeline_config( + pipeline_config_name=self.config.pipeline_config_name, + pipeline_config=self.config.pipeline_config, + ) + logger.info( + "Setting up shared OmniDreams runtime {} on {} for WebRTC.", + self.config.pipeline_config_name, + self._device, + ) + self._runtime = OmnidreamsRuntime( + config=self._inference_config(), + options=OmnidreamsRuntimeOptions( + pipeline_config=self.config.pipeline_config, + pipeline_factory=self.config.pipeline_factory, + # WebRTC warms the same long-lived runtime before real browser + # sessions. Keep prompt/image encoders available for later + # peer connections until Phase 14 replaces loopback warmup with + # first-class model/runtime warmup. + release_oneshot_encoders_after_cache_init=False, + ), + ) + self._initialize_video_encoder_sync() + + def _reset_rollout_sync(self, session_input: None = None) -> None: + del session_input + self._close_active_session_sync() + runtime = self._require_runtime() + scenario = self._session_scenario() + prepared = PreparedScenario( + initial_inputs=InferenceInput(global_conditioning={"scenario": scenario}), + source_schema=self.input_source_schema, + metadata={ + "conditioning_mode": "ludus-scene-driving", + "model_id": OMNIDREAMS_MODEL_ID, + "preset_id": self.config.pipeline_config_name, + }, + ) + provider = LudusSceneConditioningProvider( + scenario=prepared, + config=self._inference_config(), + ) + try: + initial_input = provider.prepare_initial_input() + session = runtime.start_session(initial_input) + except Exception: + provider.close() + raise + self._active_provider = provider + if self.config.debug_serve_hdmaps: + self._debug_session = _OmnidreamsHDMapDebugSession( + pipeline=runtime.pipeline, + scenario=scenario, + ) + self._active_session = self._debug_session + else: + self._debug_session = None + self._active_session = session + self._steady_output_frame_count_value = _steady_output_frame_count( + self._active_session, + fallback_pipeline=runtime.pipeline, + ) + + def _generate_one_chunk_sync( + self, + *, + segments: list[PoseSegment], + frame_times: list[float], + ) -> StepResult: + request = self._next_step_request_sync() + if request is None: + raise OmnidreamsWebRTCModelRuntimeError( + "OmniDreams WebRTC session is complete." + ) + inputs = self.input_mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput( + metadata={ + SPARSE_KEY_SEGMENTS_METADATA_KEY: tuple(segments), + "frame_times": tuple(frame_times), + "window_start_s": request.step_index / float(self.config.fps), + "window_end_s": (request.step_index + len(frame_times)) + / float(self.config.fps), + } + ), + request=request, + ) + return self._step_active_session_sync(inputs) + + def _close_sync(self) -> None: + self._close_active_session_sync() + runtime = self._runtime + self._runtime = None + if runtime is not None: + runtime.close() + if self._device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.synchronize(device=self._device) + torch.cuda.empty_cache() + + async def start_inference_session(self) -> "_OmnidreamsWebRTCInferenceSession": + self._require_open_and_initialized() + if not await self._worker.call(self._has_active_session_sync): + await self.reset_for_new_session() + return _OmnidreamsWebRTCInferenceSession(self) + + def _next_step_request_sync(self) -> StepRequest | None: + requirements = self._next_step_requirements_sync() + if requirements is None: + return None + metadata = dict(requirements.metadata) + metadata["input_frame_count"] = requirements.input_frame_count + if requirements.steady_output_frame_count is not None: + metadata["steady_output_frame_count"] = ( + requirements.steady_output_frame_count + ) + return StepRequest( + step_index=requirements.step_index, + inference_input_schema=requirements.inference_input_schema, + metadata=metadata, + ) + + def _next_step_requirements_sync(self) -> StepRequirements | None: + session = self._require_active_session() + next_requirements = getattr(session, "next_step_requirements", None) + if callable(next_requirements): + result = next_requirements() + else: + next_request = session.next_step_request() + if next_request is None: + return None + result = step_requirements_from_request(next_request) + if result is None: + return None + if not isinstance(result, StepRequirements): + raise TypeError( + "OmniDreams WebRTC session requirements must be StepRequirements, " + f"got {type(result).__name__}." + ) + return result + + def _session_info_sync(self) -> SessionInfo: + return SessionInfo( + output_layout="bvtchw", + steady_output_frame_count=self._steady_output_frame_count(), + metadata={"model_id": OMNIDREAMS_MODEL_ID}, + ) + + def _step_active_session_sync(self, inputs: InferenceInput) -> StepResult: + provider = self._require_active_provider() + session = self._require_active_session() + request = _request_from_step_inputs(inputs) + requirements = step_requirements_from_request( + request, + allow_user_input_window=True, + ) + window = _user_window_from_step_inputs( + inputs, + request=request, + input_frame_count=requirements.input_frame_count, + ) + prepared = provider.prepare_step(request=requirements, user_window=window) + if prepared.control.close_session: + raise OmnidreamsWebRTCModelRuntimeError( + prepared.control.reason or "OmniDreams WebRTC input is exhausted." + ) + if prepared.control.reset: + reset_input = prepared.control.reset_input + session.reset(reset_input) + if not prepared.control.provider_already_reset: + provider.reset(reset_input) + raise OmnidreamsWebRTCModelRuntimeError( + prepared.control.reason or "OmniDreams WebRTC session reset requested." + ) + if prepared.inference_input is None: + raise OmnidreamsWebRTCModelRuntimeError( + "OmniDreams WebRTC provider returned no inference input." + ) + result = session.step(prepared.inference_input) + if not isinstance(result, StepResult): + raise TypeError( + "OmniDreams WebRTC session steps must produce StepResult, got " + f"{type(result).__name__}." + ) + return result + + @distributed_op(WebRTCControlSignal.SESSION_STEP) + def _step_active_session_sync_all_ranks( + self, + inputs: InferenceInput, + ) -> StepResult: + return self._step_active_session_sync(inputs) + + @distributed_op(WebRTCControlSignal.SESSION_CLOSE) + def _close_active_session_sync_all_ranks(self) -> None: + self._close_active_session_sync() + + def _close_active_session_sync(self) -> None: + session = self._active_session + provider = self._active_provider + self._active_session = None + self._debug_session = None + self._active_provider = None + first_error: Exception | None = None + close_session = getattr(session, "close", None) + if callable(close_session): + try: + close_session() + except Exception as exc: + first_error = exc + if provider is not None: + try: + provider.close() + except Exception as exc: + if first_error is None: + first_error = exc + if first_error is not None: + raise first_error + + def _has_active_session_sync(self) -> bool: + return self._active_session is not None and self._active_provider is not None + + def _require_runtime(self) -> OmnidreamsRuntime: + if self._runtime is None: + raise OmnidreamsWebRTCModelRuntimeError("Runtime is not initialized.") + return self._runtime + + def _require_active_session(self) -> Any: + if self._active_session is None: + raise OmnidreamsWebRTCModelRuntimeError( + "OmniDreams WebRTC session is not initialized." + ) + return self._active_session + + def _require_active_provider(self) -> LudusSceneConditioningProvider: + if self._active_provider is None: + raise OmnidreamsWebRTCModelRuntimeError( + "OmniDreams WebRTC provider is not initialized." + ) + return self._active_provider + + def _inference_config(self) -> InferenceConfig: + return InferenceConfig( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=self.config.pipeline_config_name, + device=str(self.config.device), + seed=self.config.seed, + runtime_options={"seed": self.config.seed}, + ) + + def _session_scenario(self) -> OmnidreamsLudusReplayScenario: + return OmnidreamsLudusReplayScenario( + keyboard_events=(), + scene_dir=self.config.scene_dir, + scene_uuid=self.config.scene_uuid or DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID, + scene_variant=self.config.scene_variant, + camera_name=self.config.camera_name, + total_blocks=_WEBRTC_SESSION_TOTAL_BLOCKS, + pixel_height=self.config.video_height, + pixel_width=self.config.video_width, + fps=self.config.fps, + move_speed_per_s=self.config.move_speed_per_s, + rotate_speed_rad_per_s=self.config.rotate_speed_rad_per_s, + ) + + +class _OmnidreamsWebRTCInputMapping: + """Carry shared WebRTC window facts into the OmniDreams session facade.""" + + def validate( + self, + *, + canonical_schema: CanonicalInputSchema | None = None, + inference_input_schema: InferenceInputSchema | None = None, + ) -> None: + del canonical_schema, inference_input_schema + + def map_global_conditioning_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + ) -> InferenceInput: + del canonical_inputs + return inference_input + + def map_step_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + request: StepRequest, + ) -> InferenceInput: + del canonical_inputs + step = dict(inference_input.step) + step[_WEBRTC_STEP_REQUEST_KEY] = request + return InferenceInput( + global_conditioning=inference_input.global_conditioning, + step=step, + metadata=inference_input.metadata, + ) + + +class _OmnidreamsWebRTCInferenceSession: + """Synchronous session proxy consumed by the shared WebRTC compatibility path.""" + + def __init__(self, runtime: OmnidreamsWebRTCModelRuntime) -> None: + self._runtime = runtime + self._closed = False + + def session_info(self) -> SessionInfo: + self._require_open() + return self._runtime._worker.call_blocking(self._runtime._session_info_sync) + + def next_step_requirements(self) -> StepRequirements | None: + self._require_open() + return self._runtime._worker.call_blocking( + self._runtime._next_step_requirements_sync + ) + + def next_step_request(self) -> StepRequest | None: + self._require_open() + return self._runtime._worker.call_blocking( + self._runtime._next_step_request_sync + ) + + def step(self, inputs: InferenceInput) -> StepResult: + self._require_open() + return self._runtime._worker.call_blocking( + self._runtime._step_active_session_sync_all_ranks, + inputs, + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self._require_open() + self._runtime._worker.call_blocking(self._runtime._reset_rollout_sync_all_ranks) + + def close(self) -> None: + if self._closed: + return + self._closed = True + self._runtime._worker.call_blocking( + self._runtime._close_active_session_sync_all_ranks + ) + + def _require_open(self) -> None: + if self._closed: + raise OmnidreamsWebRTCModelRuntimeError( + "OmniDreams WebRTC inference session is closed." + ) + + +class _OmnidreamsHDMapDebugSession: + """Session-shaped debug path that streams rendered Ludus HDMaps.""" + + def __init__( + self, *, pipeline: Any, scenario: OmnidreamsLudusReplayScenario + ) -> None: + self._pipeline = pipeline + self._scenario = scenario + self._step_index = 0 + self._closed = False + + def session_info(self) -> SessionInfo: + return SessionInfo( + output_layout="bvtchw", + steady_output_frame_count=self._steady_output_frame_count(), + metadata={"stream": "hdmap"}, + ) + + def next_step_requirements(self) -> StepRequirements | None: + if self._closed or self._step_index >= self._scenario.total_blocks: + return None + return StepRequirements( + step_index=self._step_index, + input_frame_count=self._num_frames(self._step_index), + steady_output_frame_count=self._steady_output_frame_count(), + ) + + def next_step_request(self) -> StepRequest | None: + requirements = self.next_step_requirements() + if requirements is None: + return None + return StepRequest( + step_index=requirements.step_index, + metadata={ + "input_frame_count": requirements.input_frame_count, + "steady_output_frame_count": requirements.steady_output_frame_count, + }, + ) + + def step(self, inputs: InferenceInput) -> StepResult: + requirements = self.next_step_requirements() + if requirements is None: + raise OmnidreamsWebRTCModelRuntimeError( + "OmniDreams WebRTC debug session is complete." + ) + hdmap = inputs.step.get("hdmap") + if not isinstance(hdmap, torch.Tensor): + raise TypeError("OmniDreams WebRTC debug session requires step['hdmap'].") + result = StepResult.from_video_chunk( + step_index=requirements.step_index, + video_chunk=hdmap.detach(), + layout="bvtchw", + metadata={"stream": "hdmap"}, + ) + self._step_index += 1 + return result + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self._step_index = 0 + self._closed = False + + def close(self) -> None: + self._closed = True + + def _steady_output_frame_count(self) -> int: + return self._num_frames(1) + + def _num_frames(self, step_index: int) -> int: + get_num_frames = getattr(self._pipeline, "get_num_frames", None) + if not callable(get_num_frames): + return 1 + return int(get_num_frames(step_index)) + + +def _request_from_step_inputs(inputs: InferenceInput) -> StepRequest: + request = inputs.step.get(_WEBRTC_STEP_REQUEST_KEY) + if not isinstance(request, StepRequest): + raise TypeError( + "OmniDreams WebRTC step input is missing the shared StepRequest." + ) + return request + + +def _user_window_from_step_inputs( + inputs: InferenceInput, + *, + request: StepRequest, + input_frame_count: int, +) -> UserInputWindow: + frame_times = _frame_times_from_metadata(inputs.metadata, input_frame_count) + segments = _segments_from_metadata(inputs.metadata) + window = request.user_input_window or TimeWindow( + start_s=float(inputs.metadata.get("window_start_s", 0.0)), + end_s=float(inputs.metadata.get("window_end_s", frame_times[-1])), + ) + return UserInputWindow( + start_s=window.start_s, + end_s=window.end_s, + frame_times=frame_times, + metadata={SPARSE_KEY_SEGMENTS_METADATA_KEY: segments}, + ) + + +def _frame_times_from_metadata( + metadata: Mapping[str, object], + input_frame_count: int, +) -> tuple[float, ...]: + value = metadata.get("frame_times") + if not isinstance(value, tuple): + raise OmnidreamsWebRTCModelRuntimeError( + "OmniDreams WebRTC step input is missing frame_times metadata." + ) + frame_times = tuple( + _float_metadata_value(frame_time, label="frame_times") for frame_time in value + ) + if len(frame_times) != input_frame_count: + raise OmnidreamsWebRTCModelRuntimeError( + "OmniDreams WebRTC frame_times length does not match " + f"input_frame_count={input_frame_count}." + ) + return frame_times + + +def _segments_from_metadata(metadata: Mapping[str, object]) -> tuple[PoseSegment, ...]: + value = metadata.get(SPARSE_KEY_SEGMENTS_METADATA_KEY) + if not isinstance(value, tuple): + raise OmnidreamsWebRTCModelRuntimeError( + "OmniDreams WebRTC step input is missing resampled key segments." + ) + segments: list[PoseSegment] = [] + for segment in value: + if not isinstance(segment, tuple) or len(segment) != 3: + raise OmnidreamsWebRTCModelRuntimeError( + "OmniDreams WebRTC key segments must be 3-tuples." + ) + start_s, end_s, keys = segment + if not isinstance(keys, frozenset | set | tuple | list): + raise OmnidreamsWebRTCModelRuntimeError( + "OmniDreams WebRTC key segment keys must be a sequence." + ) + segments.append( + ( + _float_metadata_value(start_s, label="segment start"), + _float_metadata_value(end_s, label="segment end"), + frozenset(str(key) for key in keys), + ) + ) + return tuple(segments) + + +def _float_metadata_value(value: object, *, label: str) -> float: + if isinstance(value, bool) or not isinstance(value, int | float): + raise OmnidreamsWebRTCModelRuntimeError( + f"OmniDreams WebRTC {label} metadata must be numeric." + ) + return float(value) + + +def _steady_output_frame_count(session: Any, *, fallback_pipeline: Any) -> int: + session_info = getattr(session, "session_info", None) + if callable(session_info): + value = session_info() + if isinstance(value, SessionInfo) and value.steady_output_frame_count: + return int(value.steady_output_frame_count) + get_num_frames = getattr(fallback_pipeline, "get_num_frames", None) + if callable(get_num_frames): + return int(get_num_frames(1)) + return 1 + + +def _validate_single_view_pipeline_config( + *, + pipeline_config_name: str, + pipeline_config: Any, +) -> None: + diffusion_model = getattr(pipeline_config, "diffusion_model", None) + transformer_cfg = getattr(diffusion_model, "transformer", None) + if transformer_cfg is None: + return + if not isinstance(transformer_cfg, CosmosTransformerConfig): + raise TypeError( + "OmniDreams WebRTC requires a CosmosTransformerConfig pipeline." + ) + if transformer_cfg.num_views != 1: + raise ValueError( + "OmniDreams WebRTC supports only single-view configs; " + f"{pipeline_config_name!r} has num_views={transformer_cfg.num_views}." + ) + + +def _serve_legacy_omnidreams_webrtc_demo( + *, + spec: DemoSpec, + output: WebRTCOutputSpec, + runtime_config: OmnidreamsWebRTCModelRuntimeConfig, + runtime_factory: WebRTCRuntimeFactory, + world_rank: int, + create_app_fn: CreateWebRTCApp, + server_runner: RunWebRTCServer, +) -> object: + 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=output.client_liveness_timeout_s, + ) + from importlib.resources import files + + return serve_webrtc_demo( + output=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, + ) + + +__all__ = [ + "OmnidreamsWebRTCModelRuntime", + "OmnidreamsWebRTCModelRuntimeError", + "WebRTCRuntimeFactory", + "_serve_legacy_omnidreams_webrtc_demo", +] diff --git a/integrations/omnidreams/pyproject.toml b/integrations/omnidreams/pyproject.toml index 221cb042f..fe48be188 100644 --- a/integrations/omnidreams/pyproject.toml +++ b/integrations/omnidreams/pyproject.toml @@ -138,7 +138,7 @@ exclude = ["tests"] # workspace editable. Editable installs pick these up from the source # tree automatically. [tool.setuptools.package-data] -"omnidreams.demo" = ["web/adapter.js"] +"omnidreams.demo" = ["web/adapter.js", "web/adapter.css"] "omnidreams.interactive_drive" = [ "configs/*.yaml", "configs/wheels/*.yaml", diff --git a/integrations/omnidreams/tests/test_demo_api.py b/integrations/omnidreams/tests/test_demo_api.py index 1d2f81ef3..3d083535a 100644 --- a/integrations/omnidreams/tests/test_demo_api.py +++ b/integrations/omnidreams/tests/test_demo_api.py @@ -3,51 +3,83 @@ from __future__ import annotations +import asyncio +import json +import sys from collections.abc import Sequence from pathlib import Path from types import SimpleNamespace from typing import Any +import numpy as np import omnidreams.demo as demo_package import omnidreams.demo.spec as spec_module import pytest +import tomli as tomllib import torch from aiohttp import web from omnidreams.config import OMNIDREAMS_RUNNERS from omnidreams.demo import ( DEFAULT_OMNIDREAMS_PRESET, + OMNIDREAMS_CONDITIONING_LUDUS, + OMNIDREAMS_CONDITIONING_PRECOMPUTED, OMNIDREAMS_MODEL_ID, + LudusSceneConditioningProvider, OmnidreamsDemoAdapter, + OmnidreamsLudusReplayScenario, OmnidreamsReplayScenario, OmnidreamsWebRTCScenario, + PrecomputedHDMapProvider, ) from omnidreams.demo.app import _replay_spec, _webrtc_spec, parse_args from omnidreams.demo.replay import ( OmnidreamsReplayRuntime, OmnidreamsReplayRuntimeOptions, + OmnidreamsReplaySession, +) +from omnidreams.demo.runtime import ( + OmnidreamsRuntime, + OmnidreamsRuntimeOptions, + OmnidreamsSession, ) from omnidreams.demo.webrtc import ( - OmnidreamsWebRTCModelRuntime, OmnidreamsWebRTCModelRuntimeConfig, + _should_use_legacy_webrtc_path, serve_omnidreams_webrtc_demo, ) from flashdreams.runtime import ( + CanonicalInputs, InferenceConfig, InferenceInput, OutputArtifact, OutputTarget, StepRequest, + StepRequirements, StepResult, ) from flashdreams.runtime.demo import ( DemoSpec, Mp4OutputSpec, + NullOutputSpec, + OutputDecision, + PreparedScenario, + RuntimeHost, + SessionInfo, + UserInputWindow, WebRTCOutputSpec, ) from flashdreams.runtime.demo.replay import run_replay_demo -from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager +from flashdreams.runtime.demo.timing import SPARSE_KEY_SEGMENTS_METADATA_KEY +from flashdreams.serving.webrtc.manager import ( + BaseWebRTCSessionManager, + ManagedWebRTCSession, +) from flashdreams.serving.webrtc.server import SESSION_MANAGER_KEY +from flashdreams.serving.webrtc.services import ( + WebRTCInputSource, + WebRTCTransportService, +) pytestmark = pytest.mark.ci_cpu @@ -59,12 +91,139 @@ def test_omnidreams_demo_defaults_to_stable_non_perf_preset() -> None: assert not args.preset_id.endswith("-perf") -def test_omnidreams_demo_adapter_declares_replay_modes_only() -> None: +def test_omnidreams_replay_cli_builds_null_output_spec() -> None: + args = parse_args(["replay", "--output-mode", "null"]) + + spec = _replay_spec(args) + + assert spec.input_mode == "replay" + assert isinstance(spec.output, NullOutputSpec) + assert spec.config is not None + assert spec.config.model_id == OMNIDREAMS_MODEL_ID + + +def test_omnidreams_replay_cli_builds_ludus_conditioning_spec( + tmp_path: Path, +) -> None: + trace_path = tmp_path / "trace.json" + trace_path.write_text( + json.dumps( + { + "events": [ + {"timestamp_s": 0.0, "event": "keydown", "key": "w"}, + {"timestamp_s": 0.5, "event": "keyup", "key": "w"}, + ] + } + ), + encoding="utf-8", + ) + output_path = tmp_path / "demo.mp4" + + args = parse_args( + [ + "replay", + "--conditioning-mode", + OMNIDREAMS_CONDITIONING_LUDUS, + "--keyboard-trace", + str(trace_path), + "--scene-uuid", + "scene-1", + "--scene-variant", + "rain", + "--camera-name", + "camera_front_wide_120fov", + "--seed", + "123", + "--total-blocks", + "3", + "--output", + str(output_path), + ] + ) + + spec = _replay_spec(args) + + assert spec.input_mode == "replay" + assert isinstance(spec.scenario, dict) + scenario = spec.scenario + assert scenario["conditioning_mode"] == OMNIDREAMS_CONDITIONING_LUDUS + assert scenario["keyboard_trace_path"] == trace_path + assert scenario["scene_uuid"] == "scene-1" + assert scenario["scene_variant"] == "rain" + assert scenario["total_blocks"] == 3 + assert isinstance(spec.output, Mp4OutputSpec) + assert spec.output.path == output_path + assert spec.config is not None + assert spec.config.seed == 123 + assert spec.config.runtime_options["seed"] == 123 + + +def test_omnidreams_demo_adapter_declares_shared_modes() -> None: adapter = OmnidreamsDemoAdapter() assert adapter.model_id == OMNIDREAMS_MODEL_ID - assert adapter.supported_input_modes() == ("replay",) - assert adapter.supported_output_modes() == ("mp4",) + assert adapter.supported_input_modes() == ("replay", "keyboard-driving") + assert adapter.supported_output_modes() == ("mp4", "null", "webrtc") + assert adapter.supported_conditioning_modes() == ( + OMNIDREAMS_CONDITIONING_PRECOMPUTED, + OMNIDREAMS_CONDITIONING_LUDUS, + ) + assert [ + field.name + for field in adapter.inference_input_schema.global_conditioning_fields + ] == ["prompt", "first_frame", "scenario"] + assert [field.name for field in adapter.inference_input_schema.step_fields] == [ + "hdmap" + ] + + +def test_omnidreams_runtime_keeps_replay_aliases() -> None: + assert OmnidreamsReplayRuntime is OmnidreamsRuntime + assert OmnidreamsReplayRuntimeOptions is OmnidreamsRuntimeOptions + assert OmnidreamsReplaySession is OmnidreamsSession + + +def test_omnidreams_demo_adapter_accepts_shared_runtime_factory() -> None: + runtime = _FactoryRuntime() + pipeline_config = object() + calls: list[dict[str, Any]] = [] + + def pipeline_factory(config_value: Any, device: str) -> Any: + del config_value, device + return object() + + def runtime_factory(**kwargs: Any) -> Any: + calls.append(kwargs) + return runtime + + adapter = OmnidreamsDemoAdapter( + runtime_factory=runtime_factory, + pipeline_factory=pipeline_factory, + ) + config = InferenceConfig( + model_id=OMNIDREAMS_MODEL_ID, + runtime_options={"pipeline_config": pipeline_config}, + ) + + assert adapter.create_runtime(config) is runtime + assert len(calls) == 1 + assert calls[0]["config"] == config + options = calls[0]["options"] + assert isinstance(options, OmnidreamsRuntimeOptions) + assert options.pipeline_config is pipeline_config + assert options.pipeline_factory is pipeline_factory + + +def test_omnidreams_demo_adapter_rejects_ambiguous_runtime_factories() -> None: + def runtime_factory(**kwargs: Any) -> _FactoryRuntime: + del kwargs + return _FactoryRuntime() + + with pytest.raises(ValueError, match="runtime_factory"): + OmnidreamsDemoAdapter( + runtime_factory=runtime_factory, + replay_runtime_factory=runtime_factory, + ) def test_omnidreams_demo_does_not_import_legacy_webrtc_package() -> None: @@ -107,14 +266,17 @@ def fake_runner(**kwargs: Any) -> Sequence[OutputArtifact]: ), ) - artifacts = run_replay_demo( + result = run_replay_demo( spec=spec, adapter=adapter, output_target_factory=lambda output_spec: output, runner=fake_runner, ) - assert artifacts == (OutputArtifact(kind="video/mp4", uri="memory://omnidreams"),) + assert result.status == "completed" + assert result.artifacts == ( + OutputArtifact(kind="video/mp4", uri="memory://omnidreams"), + ) assert len(calls) == 1 assert calls[0]["adapter"] is adapter assert calls[0]["config"] == spec.config @@ -212,31 +374,364 @@ def test_omnidreams_replay_cli_can_disable_example_data(tmp_path: Path) -> None: OmnidreamsDemoAdapter().prepare_scenario(spec) -def test_omnidreams_replay_runtime_generates_video_step_result( +def test_omnidreams_precomputed_hdmap_provider_prepares_inputs( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import omnidreams.demo.providers as providers_module + + hdmap = tmp_path / "hdmap.mp4" + first_frame = tmp_path / "first.png" + hdmap.write_bytes(b"fake") + first_frame.write_bytes(b"fake") + loaded_hdmap = torch.arange(3 * 3 * 2 * 2).reshape(3, 3, 2, 2) + monkeypatch.setattr( + providers_module, + "load_first_frame_tensor", + lambda *args, **kwargs: torch.ones(1, 3, 2, 2), + ) + monkeypatch.setattr( + providers_module, + "_load_video", + lambda *args, **kwargs: loaded_hdmap, + ) + adapter = OmnidreamsDemoAdapter() + spec = _replay_demo_spec( + tmp_path=tmp_path, + hdmap=hdmap, + first_frame=first_frame, + total_blocks=2, + ) + prepared = adapter.prepare_scenario(spec) + + provider = adapter.create_model_input_provider(spec, prepared) + + assert isinstance(provider, PrecomputedHDMapProvider) + initial = provider.prepare_initial_input() + scenario = initial.global_conditioning["scenario"] + assert isinstance(scenario, OmnidreamsReplayScenario) + assert initial.global_conditioning["prompt"] == [["drive"]] + assert initial.global_conditioning["first_frame"].shape == (1, 1, 1, 3, 2, 2) + assert initial.metadata["view_names"] == ("camera_front_wide_120fov",) + + step = provider.prepare_step( + request=StepRequirements(step_index=0, input_frame_count=2), + user_window=UserInputWindow(start_s=0.0, end_s=1.0), + ) + + assert step.inference_input is not None + hdmap_chunk = step.inference_input.step["hdmap"] + assert isinstance(hdmap_chunk, torch.Tensor) + assert hdmap_chunk.shape == (1, 1, 2, 3, 2, 2) + torch.testing.assert_close(hdmap_chunk[0, 0], loaded_hdmap[:2]) + + exhausted = provider.prepare_step( + request=StepRequirements(step_index=1, input_frame_count=2), + user_window=UserInputWindow(start_s=1.0, end_s=2.0), + ) + + assert exhausted.inference_input is None + assert exhausted.control.close_session is True + provider.reset() + reset_step = provider.prepare_step( + request=StepRequirements(step_index=0, input_frame_count=1), + user_window=UserInputWindow(start_s=0.0, end_s=1.0), + ) + assert reset_step.inference_input is not None + torch.testing.assert_close( + reset_step.inference_input.step["hdmap"][0, 0], + loaded_hdmap[:1], + ) + provider.close() + + +def test_omnidreams_ludus_provider_prepares_deterministic_hdmaps( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: - import omnidreams.demo.replay as replay_module + _scene, rasterizers = _install_fake_ludus_provider_dependencies(monkeypatch) + scene_path = tmp_path / "scene.usdz" + scene_path.write_bytes(b"fake") + adapter = OmnidreamsDemoAdapter() + spec = _ludus_replay_demo_spec( + tmp_path=tmp_path, + scene_path=scene_path, + total_blocks=2, + ) + prepared = adapter.prepare_scenario(spec) + + provider = adapter.create_model_input_provider(spec, prepared) + + assert isinstance(provider, LudusSceneConditioningProvider) + initial = provider.prepare_initial_input() + scenario = initial.global_conditioning["scenario"] + assert isinstance(scenario, OmnidreamsLudusReplayScenario) + assert scenario.camera_names == ("camera_front_wide_120fov",) + assert initial.global_conditioning["prompt"] == [["city scene"]] + assert initial.global_conditioning["first_frame"].shape == (1, 1, 1, 3, 2, 2) + assert initial.metadata["view_names"] == ("camera_front_wide_120fov",) + + first = provider.prepare_step( + request=StepRequirements(step_index=0, input_frame_count=2), + user_window=UserInputWindow(start_s=0.0, end_s=2 / 30), + ) + + assert first.inference_input is not None + first_hdmap = first.inference_input.step["hdmap"] + assert isinstance(first_hdmap, torch.Tensor) + assert first_hdmap.shape == (1, 1, 2, 3, 2, 2) + assert first.inference_input.metadata["frame_timestamps_us"] == (1_000, 34_333) + assert first.inference_input.metadata["keyboard_segments"] == ( + (0.0, 2 / 30, ("w",)), + ) + assert len(rasterizers) == 1 + assert rasterizers[0].calls[0]["timestamps_us"] == (1_000, 34_333) + assert rasterizers[0].calls[0]["rig_poses_world"].shape == (2, 4, 4) + assert rasterizers[0].calls[0]["rig_poses_world"][0, 0, 3] > 0 + + provider.reset() + reset_first = provider.prepare_step( + request=StepRequirements(step_index=0, input_frame_count=2), + user_window=UserInputWindow(start_s=0.0, end_s=2 / 30), + ) + + assert reset_first.inference_input is not None + torch.testing.assert_close(reset_first.inference_input.step["hdmap"], first_hdmap) + + provider.reset() + realtime_first = provider.prepare_step( + request=StepRequirements(step_index=0, input_frame_count=2), + user_window=UserInputWindow( + start_s=0.0, + end_s=2 / 30, + frame_times=(1 / 30, 2 / 30), + metadata={ + SPARSE_KEY_SEGMENTS_METADATA_KEY: ((0.0, 2 / 30, frozenset({"w"})),) + }, + ), + ) + + assert realtime_first.inference_input is not None + torch.testing.assert_close( + realtime_first.inference_input.step["hdmap"], + first_hdmap, + ) + provider.close() + assert rasterizers[0].closed is True + + +def test_omnidreams_replay_run_mode_uses_precomputed_provider( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import omnidreams.demo.providers as providers_module hdmap = tmp_path / "hdmap.mp4" first_frame = tmp_path / "first.png" hdmap.write_bytes(b"fake") first_frame.write_bytes(b"fake") + loaded_hdmap = torch.arange(2 * 3 * 2 * 2).reshape(2, 3, 2, 2) pipeline = _FakeOmnidreamsPipeline() + sink = _RecordingOutputSink() monkeypatch.setattr( - replay_module, + providers_module, "load_first_frame_tensor", lambda *args, **kwargs: torch.zeros(1, 3, 2, 2), ) monkeypatch.setattr( - replay_module, + providers_module, "_load_video", - lambda *args, **kwargs: torch.zeros(2, 3, 2, 2), + lambda *args, **kwargs: loaded_hdmap, + ) + adapter = OmnidreamsDemoAdapter( + pipeline_factory=lambda pipeline_config, device: pipeline, + ) + spec = _replay_demo_spec( + tmp_path=tmp_path, + hdmap=hdmap, + first_frame=first_frame, + total_blocks=2, + ) + + result = run_replay_demo( + spec=spec, + adapter=adapter, + output_sink_factory=lambda output_spec: sink, ) - runtime = OmnidreamsReplayRuntime( + assert result.status == "completed" + assert result.artifacts == ( + OutputArtifact(kind="video/mp4", uri="memory://omnidreams"), + ) + assert [result.step_index for result in sink.results] == [0, 1] + assert pipeline.initialize_cache_calls == [ + { + "text": [["drive"]], + "image_shape": (1, 1, 1, 3, 2, 2), + "view_names": ["camera_front_wide_120fov"], + } + ] + assert len(pipeline.generated_hdmaps) == 2 + torch.testing.assert_close(pipeline.generated_hdmaps[0][0, 0], loaded_hdmap[:1]) + torch.testing.assert_close(pipeline.generated_hdmaps[1][0, 0], loaded_hdmap[1:2]) + + +def test_omnidreams_replay_run_mode_uses_ludus_provider( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_fake_ludus_provider_dependencies(monkeypatch) + scene_path = tmp_path / "scene.usdz" + scene_path.write_bytes(b"fake") + pipeline = _FakeOmnidreamsPipeline() + sink = _RecordingOutputSink() + adapter = OmnidreamsDemoAdapter( + pipeline_factory=lambda pipeline_config, device: pipeline, + ) + spec = _ludus_replay_demo_spec( + tmp_path=tmp_path, + scene_path=scene_path, + total_blocks=2, + ) + + result = run_replay_demo( + spec=spec, + adapter=adapter, + output_sink_factory=lambda output_spec: sink, + ) + + assert result.status == "completed" + assert result.artifacts == ( + OutputArtifact(kind="video/mp4", uri="memory://omnidreams"), + ) + assert [result.step_index for result in sink.results] == [0, 1] + assert pipeline.initialize_cache_calls == [ + { + "text": [["city scene"]], + "image_shape": (1, 1, 1, 3, 2, 2), + "view_names": ["camera_front_wide_120fov"], + } + ] + assert len(pipeline.generated_hdmaps) == 2 + assert pipeline.generated_hdmaps[0].shape == (1, 1, 1, 3, 2, 2) + assert pipeline.generated_hdmaps[1].shape == (1, 1, 1, 3, 2, 2) + assert not torch.equal(pipeline.generated_hdmaps[0], pipeline.generated_hdmaps[1]) + + +def test_omnidreams_replay_null_output_uses_precomputed_provider( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import omnidreams.demo.providers as providers_module + + hdmap = tmp_path / "hdmap.mp4" + first_frame = tmp_path / "first.png" + hdmap.write_bytes(b"fake") + first_frame.write_bytes(b"fake") + loaded_hdmap = torch.arange(2 * 3 * 2 * 2).reshape(2, 3, 2, 2) + pipeline = _FakeOmnidreamsPipeline() + monkeypatch.setattr( + providers_module, + "load_first_frame_tensor", + lambda *args, **kwargs: torch.zeros(1, 3, 2, 2), + ) + monkeypatch.setattr( + providers_module, + "_load_video", + lambda *args, **kwargs: loaded_hdmap, + ) + adapter = OmnidreamsDemoAdapter( + pipeline_factory=lambda pipeline_config, device: pipeline, + ) + spec = _replay_demo_spec( + tmp_path=tmp_path, + hdmap=hdmap, + first_frame=first_frame, + total_blocks=2, + output=NullOutputSpec(), + ) + + result = run_replay_demo(spec=spec, adapter=adapter) + + assert result.status == "completed" + assert result.artifacts == () + assert pipeline.initialize_cache_calls == [ + { + "text": [["drive"]], + "image_shape": (1, 1, 1, 3, 2, 2), + "view_names": ["camera_front_wide_120fov"], + } + ] + assert len(pipeline.generated_hdmaps) == 2 + torch.testing.assert_close(pipeline.generated_hdmaps[0][0, 0], loaded_hdmap[:1]) + torch.testing.assert_close(pipeline.generated_hdmaps[1][0, 0], loaded_hdmap[1:2]) + + +def test_omnidreams_replay_output_target_path_uses_precomputed_provider( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import omnidreams.demo.providers as providers_module + + hdmap = tmp_path / "hdmap.mp4" + first_frame = tmp_path / "first.png" + hdmap.write_bytes(b"fake") + first_frame.write_bytes(b"fake") + loaded_hdmap = torch.arange(1 * 3 * 2 * 2).reshape(1, 3, 2, 2) + pipeline = _FakeOmnidreamsPipeline() + output = _RecordingOutputTarget() + monkeypatch.setattr( + providers_module, + "load_first_frame_tensor", + lambda *args, **kwargs: torch.zeros(1, 3, 2, 2), + ) + monkeypatch.setattr( + providers_module, + "_load_video", + lambda *args, **kwargs: loaded_hdmap, + ) + adapter = OmnidreamsDemoAdapter( + pipeline_factory=lambda pipeline_config, device: pipeline, + ) + spec = _replay_demo_spec( + tmp_path=tmp_path, + hdmap=hdmap, + first_frame=first_frame, + total_blocks=1, + ) + + result = run_replay_demo( + spec=spec, + adapter=adapter, + output_target_factory=lambda output_spec: output, + ) + + assert result.status == "completed" + assert [result.step_index for result in output.results] == [0] + assert len(pipeline.generated_hdmaps) == 1 + torch.testing.assert_close(pipeline.generated_hdmaps[0][0, 0], loaded_hdmap) + + +def test_omnidreams_replay_runtime_generates_video_step_result( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import omnidreams.demo.runtime as runtime_module + + hdmap = tmp_path / "hdmap.mp4" + first_frame = tmp_path / "first.png" + hdmap.write_bytes(b"fake") + first_frame.write_bytes(b"fake") + pipeline = _FakeOmnidreamsPipeline() + monkeypatch.setattr( + runtime_module, + "load_first_frame_tensor", + lambda *args, **kwargs: torch.zeros(1, 3, 2, 2), + ) + + runtime = OmnidreamsRuntime( config=InferenceConfig(model_id=OMNIDREAMS_MODEL_ID, device="cpu"), - options=OmnidreamsReplayRuntimeOptions( + options=OmnidreamsRuntimeOptions( pipeline_config=object(), pipeline_factory=lambda pipeline_config, device: pipeline, ), @@ -254,11 +749,17 @@ def test_omnidreams_replay_runtime_generates_video_step_result( session = runtime.start_session( InferenceInput(global_conditioning={"scenario": scenario}) ) + assert isinstance(session, OmnidreamsSession) + requirements = session.next_step_requirements() + assert isinstance(requirements, StepRequirements) + assert requirements.step_index == 0 + assert requirements.input_frame_count == 1 request = session.next_step_request() assert request is not None assert request.step_index == 0 - result = session.step(InferenceInput()) + assert request.metadata["input_frame_count"] == 1 + result = session.step(InferenceInput(step={"hdmap": torch.zeros(1, 1, 1, 3, 2, 2)})) assert result.step_index == 0 assert result.frame_count == 1 @@ -267,6 +768,7 @@ def test_omnidreams_replay_runtime_generates_video_step_result( 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.released_encoders is True assert pipeline.initialize_cache_calls == [ { "text": [["drive"]], @@ -341,7 +843,10 @@ 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: + legacy_module_name = "omnidreams.demo.webrtc_legacy" + sys.modules.pop(legacy_module_name, None) pipeline_config = object() + runtime = _FactoryRuntime() spec = DemoSpec( model_id=OMNIDREAMS_MODEL_ID, preset_id=DEFAULT_OMNIDREAMS_PRESET, @@ -350,7 +855,6 @@ def test_omnidreams_webrtc_demo_uses_shared_manager_with_model_config() -> None: scene_uuid="scene-1", scene_variant="rain", camera_name="camera_front_wide_120fov", - debug_serve_hdmaps=True, prefer_sw_encoder=True, ), output=WebRTCOutputSpec( @@ -371,32 +875,121 @@ def test_omnidreams_webrtc_demo_uses_shared_manager_with_model_config() -> None: ) calls: list[dict[str, Any]] = [] + runtime_calls: list[dict[str, Any]] = [] + + def shared_runtime_factory(**kwargs: Any) -> Any: + runtime_calls.append(kwargs) + return runtime + serve_omnidreams_webrtc_demo( spec=spec, world_rank=1, - runtime_factory=_FakeWebRTCRuntime, + shared_runtime_factory=shared_runtime_factory, 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._runtime is runtime + assert isinstance(manager._shared_host, RuntimeHost) + assert isinstance(manager._shared_adapter, OmnidreamsDemoAdapter) + assert isinstance(manager._shared_scenario, PreparedScenario) + assert manager.runtime_config.pipeline_config is pipeline_config + assert manager.runtime_config.pipeline_config_name == DEFAULT_OMNIDREAMS_PRESET + assert manager.runtime_config.scene_uuid == "scene-1" + assert manager.runtime_config.scene_variant == "rain" + assert manager.runtime_config.seed == 123 + assert manager.runtime_config.device == "cuda:7" + assert manager.runtime_config.video_width == 64 + assert manager.runtime_config.video_height == 32 + assert manager.runtime_config.fps == 24 + assert manager.runtime_config.debug_serve_hdmaps is False + assert manager.runtime_config.encoder_backend == "default" assert manager.identity == DEFAULT_OMNIDREAMS_PRESET + assert len(runtime_calls) == 1 + runtime_config = runtime_calls[0]["config"] + assert runtime_config.seed == 123 + assert runtime_config.runtime_options["pipeline_config"] is pipeline_config + assert ( + runtime_config.runtime_options["release_oneshot_encoders_after_cache_init"] + is False + ) + options = runtime_calls[0]["options"] + assert isinstance(options, OmnidreamsRuntimeOptions) + assert options.release_oneshot_encoders_after_cache_init is False + scenario = manager._shared_scenario.initial_inputs.global_conditioning["scenario"] + assert isinstance(scenario, OmnidreamsLudusReplayScenario) + assert scenario.scene_uuid == "scene-1" + assert scenario.scene_variant == "rain" + assert scenario.pixel_width == 64 + assert scenario.pixel_height == 32 + assert scenario.fps == 24 assert calls[0]["host"] == "0.0.0.0" assert calls[0]["port"] == 8082 + assert legacy_module_name not in sys.modules + + +def test_omnidreams_webrtc_demo_keeps_legacy_runtime_factory_path() -> None: + spec = DemoSpec( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + input_mode="keyboard-driving", + scenario=OmnidreamsWebRTCScenario(), + output=WebRTCOutputSpec( + host="0.0.0.0", + port=8082, + fps=24, + video_width=64, + video_height=32, + warmup_chunks=0, + warmup_timeout_s=1.0, + ), + config=InferenceConfig( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + device="cuda:7", + runtime_options={"pipeline_config": object(), "seed": 123}, + ), + ) + + 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 manager.runtime_config is runtime.config + assert runtime.config.debug_serve_hdmaps is False + + +def test_omnidreams_webrtc_demo_keeps_legacy_fallback_gates( + monkeypatch: pytest.MonkeyPatch, +) -> None: + assert _should_use_legacy_webrtc_path( + scenario=OmnidreamsWebRTCScenario(), + runtime_factory=_FakeWebRTCRuntime, + ) + assert _should_use_legacy_webrtc_path( + scenario=OmnidreamsWebRTCScenario(debug_serve_hdmaps=True), + runtime_factory=None, + ) + + monkeypatch.setenv("WORLD_SIZE", "2") + assert _should_use_legacy_webrtc_path( + scenario=OmnidreamsWebRTCScenario(), + runtime_factory=None, + ) + + monkeypatch.setenv("WORLD_SIZE", "not-an-int") + assert not _should_use_legacy_webrtc_path( + scenario=OmnidreamsWebRTCScenario(), + runtime_factory=None, + ) def test_omnidreams_webrtc_demo_installs_model_assets_without_routes( @@ -439,7 +1032,7 @@ def fake_create_packaged_webrtc_app(**kwargs: Any) -> web.Application: app = serve_omnidreams_webrtc_demo( spec=spec, - runtime_factory=_FakeWebRTCRuntime, + shared_runtime_factory=lambda **kwargs: _FactoryRuntime(), server_runner=lambda **kwargs: None, ) @@ -453,6 +1046,26 @@ def fake_create_packaged_webrtc_app(**kwargs: Any) -> web.Application: assert app_calls[0]["configure_app"] is None +def test_omnidreams_webrtc_adapter_caps_video_display_size() -> None: + web_dir = Path(demo_package.__file__).resolve().parent / "web" + adapter_js = (web_dir / "adapter.js").read_text(encoding="utf-8") + adapter_css = (web_dir / "adapter.css").read_text(encoding="utf-8") + + assert 'stylesheet: "/model-static/adapter.css?v=model-ui-v2"' in adapter_js + assert ".stageVideo" in adapter_css + assert "1280px" in adapter_css + assert "704px" in adapter_css + assert "calc(" not in adapter_css + assert "object-fit: contain" in adapter_css + + pyproject = Path(__file__).resolve().parents[1] / "pyproject.toml" + with pyproject.open("rb") as fh: + meta = tomllib.load(fh) + package_data = meta["tool"]["setuptools"]["package-data"]["omnidreams.demo"] + assert "web/adapter.js" in package_data + assert "web/adapter.css" in package_data + + def test_omnidreams_webrtc_demo_serves_through_shared_runner( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -495,7 +1108,7 @@ def fake_server_runner(**kwargs: Any) -> None: app = serve_omnidreams_webrtc_demo( spec=spec, world_rank=0, - runtime_factory=_FakeWebRTCRuntime, + shared_runtime_factory=lambda **kwargs: _FactoryRuntime(), server_runner=fake_server_runner, ) @@ -508,59 +1121,402 @@ def fake_server_runner(**kwargs: Any) -> None: @pytest.mark.asyncio -async def test_omnidreams_demo_runtime_generates_directly_from_controls() -> None: +async def test_omnidreams_webrtc_runtime_uses_shared_session( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from omnidreams.demo.webrtc_legacy import OmnidreamsWebRTCModelRuntime + + _scene, rasterizers = _install_fake_ludus_provider_dependencies(monkeypatch) + scene_path = tmp_path / "scene.usdz" + scene_path.write_bytes(b"fake") + pipeline = _VariableFrameOmnidreamsPipeline((2, 3)) config = OmnidreamsWebRTCModelRuntimeConfig( pipeline_config_name="fake", pipeline_config=object(), + pipeline_factory=lambda pipeline_config, device: pipeline, + scene_dir=scene_path, device="cpu", fps=30, + video_height=2, + video_width=2, 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], + await runtime.initialize() + await runtime.reset_for_new_session() + session = await runtime.start_inference_session() + + first_request = session.next_step_request() + assert first_request is not None + assert first_request.metadata["input_frame_count"] == 2 + first = session.step( + runtime.input_mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput( + metadata={ + SPARSE_KEY_SEGMENTS_METADATA_KEY: ( + (0.0, 2 / 30, frozenset({"w"})), + ), + "frame_times": (1 / 30, 2 / 30), + "window_start_s": 0.0, + "window_end_s": 2 / 30, + } + ), + request=first_request, + ) + ) + second_request = session.next_step_request() + assert second_request is not None + assert second_request.metadata["input_frame_count"] == 3 + second = session.step( + runtime.input_mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=InferenceInput( + metadata={ + SPARSE_KEY_SEGMENTS_METADATA_KEY: ( + (2 / 30, 5 / 30, frozenset({"d"})), + ), + "frame_times": (3 / 30, 4 / 30, 5 / 30), + "window_start_s": 2 / 30, + "window_end_s": 5 / 30, + } + ), + request=second_request, + ) ) 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 isinstance(session, OmnidreamsSession) is False + assert pipeline.initialize_cache_calls == [ + { + "text": [["city scene"]], + "image_shape": (1, 1, 1, 3, 2, 2), + "view_names": ["camera_front_wide_120fov"], + } ] - assert wrapper.finalized == [0, 1] + assert [tuple(hdmap.shape) for hdmap in pipeline.generated_hdmaps] == [ + (1, 1, 2, 3, 2, 2), + (1, 1, 3, 3, 2, 2), + ] + assert rasterizers[0].calls[0]["timestamps_us"] == (1_000, 34_333) + assert rasterizers[0].calls[1]["timestamps_us"] == (67_666, 100_999, 134_332) + session.close() await runtime.close() + assert rasterizers[0].closed is True + + +@pytest.mark.asyncio +async def test_omnidreams_webrtc_runtime_keeps_encoders_after_warmup_session( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from omnidreams.demo.webrtc_legacy import OmnidreamsWebRTCModelRuntime + + _install_fake_ludus_provider_dependencies(monkeypatch) + scene_path = tmp_path / "scene.usdz" + scene_path.write_bytes(b"fake") + pipeline = _FailsIfEncodersReleasedOmnidreamsPipeline() + config = OmnidreamsWebRTCModelRuntimeConfig( + pipeline_config_name="fake", + pipeline_config=object(), + pipeline_factory=lambda pipeline_config, device: pipeline, + scene_dir=scene_path, + device="cpu", + fps=30, + video_height=2, + video_width=2, + warmup_chunks=0, + ) + runtime = OmnidreamsWebRTCModelRuntime(config=config) + await runtime.initialize() + + await runtime.reset_for_new_session() + warmup_session = await runtime.start_inference_session() + warmup_session.close() + await runtime.reset_for_new_session() + browser_session = await runtime.start_inference_session() + + assert pipeline.released_encoders is False + assert len(pipeline.initialize_cache_calls) == 2 + browser_session.close() + await runtime.close() + + +@pytest.mark.asyncio +async def test_omnidreams_webrtc_manager_drives_shared_session( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _scene, rasterizers = _install_fake_ludus_provider_dependencies(monkeypatch) + scene_path = tmp_path / "scene.usdz" + scene_path.write_bytes(b"fake") + pipeline = _VariableFrameOmnidreamsPipeline((2, 3)) + adapter = OmnidreamsDemoAdapter( + pipeline_factory=lambda pipeline_config, device: pipeline, + ) + spec = DemoSpec( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + input_mode="keyboard-driving", + scenario=OmnidreamsWebRTCScenario( + scene_dir=scene_path, + scene_uuid="scene-1", + camera_name="camera_front_wide_120fov", + ), + output=WebRTCOutputSpec( + fps=30, + video_width=2, + video_height=2, + warmup_chunks=0, + warmup_timeout_s=1.0, + ), + config=InferenceConfig( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + device="cpu", + seed=123, + runtime_options={ + "pipeline_config": object(), + "seed": 123, + "release_oneshot_encoders_after_cache_init": False, + }, + ), + ) + prepared = adapter.prepare_scenario(spec) + assert spec.config is not None + runtime = adapter.create_runtime(spec.config) + runtime_config = OmnidreamsWebRTCModelRuntimeConfig( + pipeline_config_name=DEFAULT_OMNIDREAMS_PRESET, + pipeline_config=object(), + scene_dir=scene_path, + scene_uuid="scene-1", + device="cpu", + fps=30, + video_height=2, + video_width=2, + warmup_chunks=0, + ) + host = RuntimeHost(runtime) + manager = BaseWebRTCSessionManager( + runtime=runtime, + runtime_config=runtime_config, + fps=runtime_config.fps, + identity=runtime_config.pipeline_config_name, + supported_control_keys=frozenset({"w", "a", "s", "d"}), + shared_host=host, + shared_adapter=adapter, + shared_spec=spec, + shared_scenario=prepared, + ) + manager._runtime_ready = True + loop = asyncio.get_running_loop() + context = manager._shared_run_context(loop) + reservation = context.admission.try_reserve() + assert reservation is not None + resampler = _FakeWebRTCResampler(start_v=loop.time(), fps=runtime_config.fps) + input_source = WebRTCInputSource(resampler=resampler) + input_source.handle_browser_payload( + {"type": "action", "action": {"event": "step"}}, + timestamp_s=loop.time(), + ) + transport = WebRTCTransportService(loop=loop) + channel = _FakeWebRTCChannel() + managed_session = ManagedWebRTCSession( + runtime=runtime, + video_track=_FakeWebRTCVideoTrack(fps=runtime_config.fps), # ty:ignore[invalid-argument-type] + video_encoder=_FakeWebRTCVideoEncoder(), # ty:ignore[invalid-argument-type] + peer_connection=_FakeWebRTCPeerConnection(), + resampler=resampler, # ty:ignore[invalid-argument-type] + control_channel=channel, + input_source=input_source, + transport=transport, + reservation=reservation, + last_client_message_at=loop.time(), + ) + manager._active_session = managed_session + try: + managed_session.generation_task = asyncio.create_task( + manager._run_realtime_driver_session( + managed_session=managed_session, + context=context, + session_input=None, + ) + ) + chunk = await _wait_for_chunk_done(channel) + + assert chunk["type"] == "chunk_done" + assert chunk["model"] == DEFAULT_OMNIDREAMS_PRESET + assert chunk["num_frames"] == 2 + assert [tuple(hdmap.shape) for hdmap in pipeline.generated_hdmaps][:1] == [ + (1, 1, 2, 3, 2, 2) + ] + assert rasterizers[0].calls[0]["timestamps_us"] == (1_000, 34_333) + assert isinstance( + managed_session.input_source, + WebRTCInputSource, + ) + finally: + transport.close("test complete") + await manager.shutdown() class _RecordingOutputTarget: + def __init__(self) -> None: + self.results: list[StepResult] = [] + def open(self) -> None: return None def write(self, result: StepResult) -> None: - del result + self.results.append(result) def close(self) -> Sequence[OutputArtifact]: return () +def _replay_demo_spec( + *, + tmp_path: Path, + hdmap: Path, + first_frame: Path, + total_blocks: int, + output: Mp4OutputSpec | NullOutputSpec | None = None, +) -> DemoSpec: + return DemoSpec( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + input_mode="replay", + scenario={ + "prompt": "drive", + "hdmap_video_paths": (hdmap,), + "first_frame_paths": (first_frame,), + "camera_names": ("camera_front_wide_120fov",), + "total_blocks": total_blocks, + "pixel_height": 2, + "pixel_width": 2, + "fps": 30, + }, + output=output or Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=30), + config=InferenceConfig( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + device="cpu", + runtime_options={"pipeline_config": object()}, + ), + ) + + +def _ludus_replay_demo_spec( + *, + tmp_path: Path, + scene_path: Path, + total_blocks: int, + output: Mp4OutputSpec | NullOutputSpec | None = None, +) -> DemoSpec: + return DemoSpec( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + input_mode="replay", + scenario={ + "conditioning_mode": OMNIDREAMS_CONDITIONING_LUDUS, + "keyboard_events": ( + {"timestamp_s": 0.0, "event": "keydown", "key": "w"}, + {"timestamp_s": 0.5, "event": "keyup", "key": "w"}, + ), + "scene_path": scene_path, + "scene_variant": "default", + "camera_name": "camera_front_wide_120fov", + "total_blocks": total_blocks, + "pixel_height": 2, + "pixel_width": 2, + "fps": 30, + }, + output=output or Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=30), + config=InferenceConfig( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + device="cpu", + seed=123, + runtime_options={"pipeline_config": object(), "seed": 123}, + ), + ) + + +def _install_fake_ludus_provider_dependencies( + monkeypatch: pytest.MonkeyPatch, +) -> tuple[SimpleNamespace, list["_FakeLudusRasterizer"]]: + import omnidreams.demo.providers as providers_module + + scene = SimpleNamespace( + scene_id="fake-scene", + prompt="city scene", + initial_rgb=np.zeros((2, 2, 3), dtype=np.uint8), + initial_rig_to_world=np.eye(4, dtype=np.float32), + initial_timestamp_us=1_000, + ) + rasterizers: list[_FakeLudusRasterizer] = [] + + def fake_load_scene_bundle(*args: Any, **kwargs: Any) -> SimpleNamespace: + del args, kwargs + return scene + + def fake_new_rasterizer(*args: Any, **kwargs: Any) -> "_FakeLudusRasterizer": + del args, kwargs + rasterizer = _FakeLudusRasterizer() + rasterizers.append(rasterizer) + return rasterizer + + monkeypatch.setattr( + providers_module, + "_load_ludus_scene_bundle", + fake_load_scene_bundle, + ) + monkeypatch.setattr( + providers_module, + "_new_ludus_rasterizer", + fake_new_rasterizer, + ) + return scene, rasterizers + + +class _RecordingOutputSink: + produces_artifacts = True + + def __init__(self) -> None: + self.session_info: SessionInfo | None = None + self.results: list[StepResult] = [] + self.closed = False + + def open(self, session_info: SessionInfo) -> None: + self.session_info = session_info + + def begin_generation(self, generation: int) -> None: + del generation + + def write(self, result: StepResult) -> OutputDecision: + self.results.append(result) + return OutputDecision() + + def close(self) -> Sequence[OutputArtifact]: + self.closed = True + return (OutputArtifact(kind="video/mp4", uri="memory://omnidreams"),) + + +class _FactoryRuntime: + def start_session(self, inputs: InferenceInput) -> Any: + del inputs + raise NotImplementedError + + def close(self) -> None: + return None + + class _FakeOmnidreamsPipeline: def __init__(self) -> None: self.initialize_cache_calls: list[dict[str, Any]] = [] + self.generated_hdmaps: list[torch.Tensor] = [] self.released_encoders = False def initialize_cache( @@ -593,7 +1549,8 @@ def generate( cache: object, hdmap: torch.Tensor, ) -> torch.Tensor: - del cache, hdmap + del cache + self.generated_hdmaps.append(hdmap.detach().clone()) return torch.full((1, 1, 1, 3, 2, 2), float(autoregressive_index)) def finalize(self, *, autoregressive_index: int, cache: object) -> dict[str, float]: @@ -601,61 +1558,171 @@ def finalize(self, *, autoregressive_index: int, cache: object) -> dict[str, flo return {"denoise_s": 0.25} -class _FakeRenderer: +class _VariableFrameOmnidreamsPipeline(_FakeOmnidreamsPipeline): + def __init__(self, frame_counts: tuple[int, ...]) -> None: + super().__init__() + self._frame_counts = frame_counts + + def get_num_frames(self, autoregressive_index: int) -> int: + return self._frame_counts[ + min(autoregressive_index, len(self._frame_counts) - 1) + ] + + def generate( + self, + *, + autoregressive_index: int, + cache: object, + hdmap: torch.Tensor, + ) -> torch.Tensor: + del cache + self.generated_hdmaps.append(hdmap.detach().clone()) + frame_count = self.get_num_frames(autoregressive_index) + return torch.full((1, 1, frame_count, 3, 2, 2), float(autoregressive_index)) + + +class _FailsIfEncodersReleasedOmnidreamsPipeline(_FakeOmnidreamsPipeline): + def initialize_cache( + self, + *, + text: list[list[str]], + image: torch.Tensor, + view_names: list[str], + ) -> object: + if self.released_encoders: + raise AssertionError("encoders were released before the next session") + return super().initialize_cache( + text=text, + image=image, + view_names=view_names, + ) + + +class _FakeLudusRasterizer: def __init__(self) -> None: + self.loaded_scene: object | None = None + self.calls: list[dict[str, Any]] = [] self.closed = False + def load_scene(self, scene: object) -> None: + self.loaded_scene = scene + + def render_chunk( + self, + *, + rig_poses_world: np.ndarray, + timestamps_us: np.ndarray, + ) -> SimpleNamespace: + self.calls.append( + { + "rig_poses_world": np.array(rig_poses_world, copy=True), + "timestamps_us": tuple(int(t) for t in timestamps_us), + } + ) + frames = [] + for timestamp_us in timestamps_us: + value = int(timestamp_us % 251) + frames.append( + SimpleNamespace( + rgb_host_uint8=np.full((2, 2, 3), value, dtype=np.uint8) + ) + ) + return SimpleNamespace(frames=tuple(frames)) + def cleanup(self) -> None: self.closed = True -class _FakeConditioningWrapper: - initial_frame_chunk_size = 2 - frame_chunk_size = 3 +class _FakeWebRTCResampler: + def __init__(self, *, start_v: float, fps: int) -> None: + self.dt = 1.0 / fps + self.next_chunk_start_v = start_v - def __init__(self) -> None: - self.calls: list[tuple[str, tuple[int, ...], list[int]]] = [] - self.finalized: list[int] = [] - self.cleaned = False + def reset(self, *, start_v: float) -> None: + self.next_chunk_start_v = start_v + + def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: + del arrival_t, event, key + + def sample_chunk( + self, + num_frames: int, + ) -> tuple[list[tuple[float, float, frozenset[str]]], list[float]]: + start = self.next_chunk_start_v + frame_times = [start + (index + 1) * self.dt for index in range(num_frames)] + end = frame_times[-1] + self.next_chunk_start_v = end + return [(start, end, frozenset({"w"}))], frame_times + + +class _FakeWebRTCVideoTrack: + def __init__(self, *, fps: int) -> None: + self.fps = fps + self.closed = False + self.enqueued: list[StepResult] = [] - def start_generation(self, **kwargs: Any) -> SimpleNamespace: - return self._output("start", kwargs=kwargs, frame_count=2, step_index=0) + async def enqueue_result(self, result: StepResult) -> int: + self.enqueued.append(result) + return result.frame_count - def continue_generation(self, **kwargs: Any) -> SimpleNamespace: - return self._output("continue", kwargs=kwargs, frame_count=3, step_index=1) + def qsize(self) -> int: + return 0 - def _output( + async def close(self) -> None: + self.closed = True + + +class _FakeWebRTCVideoEncoder: + backend = "fake" + prefers_codec: str | None = None + + def prepare_chunk_payload(self, result: StepResult, track: Any) -> StepResult: + del track + return result + + async def deliver_prepared_chunk( self, - operation: str, + payload: object, + track: Any, *, - kwargs: dict[str, Any], - frame_count: int, - step_index: int, + force_keyframe: bool = False, ) -> 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()) + del force_keyframe + if not isinstance(payload, StepResult): + raise TypeError("Fake WebRTC encoder expected a StepResult payload.") 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}, + num_frames=await track.enqueue_result(payload), + encode_ms=0.0, ) - 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 _FakeWebRTCPeerConnection: + def __init__(self) -> None: + self.closed = False + + async def close(self) -> None: + self.closed = True + + +class _FakeWebRTCChannel: + def __init__(self) -> None: + self.messages: list[str] = [] + + def send(self, message: str) -> None: + self.messages.append(message) + + +async def _wait_for_chunk_done(channel: _FakeWebRTCChannel) -> dict[str, Any]: + for _ in range(100): + chunk_done = [ + json.loads(message) + for message in channel.messages + if json.loads(message).get("type") == "chunk_done" + ] + if chunk_done: + return chunk_done[0] + await asyncio.sleep(0.01) + pytest.fail("Timed out waiting for WebRTC chunk_done.") class _FakeWebRTCRuntime: From 23f182cacd34484278a348e4659f1b2fc8e845af Mon Sep 17 00:00:00 2001 From: jarcherNV Date: Mon, 10 Aug 2026 18:43:10 -0700 Subject: [PATCH 17/19] Migrate LingBot demos onto unified runtime (#435) Migrate LingBot demos onto unified runtime Routes LingBot replay and WebRTC paths through the shared demo runtime stack, adds provider/shared WebRTC parity coverage, and documents the new demo commands. --- .../flashdreams/serving/webrtc/manager.py | 8 +- integrations/lingbot/README.md | 47 +- integrations/lingbot/lingbot/demo/__init__.py | 2 + integrations/lingbot/lingbot/demo/adapter.py | 111 ++++- .../lingbot/lingbot/demo/providers.py | 229 +++++++++ integrations/lingbot/lingbot/demo/webrtc.py | 79 +++- integrations/lingbot/lingbot/input_mapping.py | 6 + integrations/lingbot/lingbot/runner.py | 59 ++- integrations/lingbot/tests/test_demo_api.py | 290 +++++++++++- .../lingbot/tests/test_demo_providers.py | 433 ++++++++++++++++++ integrations/lingbot/tests/test_smoke.py | 51 ++- 11 files changed, 1248 insertions(+), 67 deletions(-) create mode 100644 integrations/lingbot/lingbot/demo/providers.py create mode 100644 integrations/lingbot/tests/test_demo_providers.py diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index 4291cd637..9aa22653d 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -33,7 +33,6 @@ ProviderCapabilities, ResamplerRealtimeClock, RunContext, - RunResult, RuntimeHost, SessionEdges, SessionInfo, @@ -79,7 +78,6 @@ from flashdreams.serving.webrtc.runtime import ( WebRTCControlSignal, WebRTCRuntimeConfig, - WebRTCSessionRuntime, ) from flashdreams.serving.webrtc.server import SessionBusyError from flashdreams.serving.webrtc.services import ( @@ -681,6 +679,7 @@ def __init__( shared_host: RuntimeHost | None = None, shared_adapter: Any | None = None, shared_spec: DemoSpec | None = None, + shared_spec_factory: Callable[[Any], DemoSpec] | None = None, shared_scenario: PreparedScenario | None = None, shared_pipeline_factory: Callable[[], StepPipeline] | None = None, ) -> None: @@ -711,6 +710,7 @@ def __init__( self._shared_context: RunContext | None = None self._shared_adapter = shared_adapter self._shared_spec = shared_spec + self._shared_spec_factory = shared_spec_factory self._shared_scenario = shared_scenario self._shared_pipeline_factory = shared_pipeline_factory self._shared_video_encoder: VideoEncoder | None = None @@ -1844,6 +1844,7 @@ async def _run_realtime_driver_session( adapter = self._shared_adapter spec = self._shared_spec scenario = self._shared_scenario + spec_factory = self._shared_spec_factory if adapter is None or spec is None: adapter = _LegacyWebRTCDemoAdapter( runtime=self._runtime, @@ -1851,6 +1852,9 @@ async def _run_realtime_driver_session( session_input=session_input, ) spec = self._shared_demo_spec() + elif spec_factory is not None and session_input is not None: + spec = spec_factory(session_input) + scenario = None if scenario is None: scenario = adapter.prepare_scenario(spec) run_mode = WebRTCRunMode( diff --git a/integrations/lingbot/README.md b/integrations/lingbot/README.md index b781e4a21..e0029e48a 100644 --- a/integrations/lingbot/README.md +++ b/integrations/lingbot/README.md @@ -98,6 +98,51 @@ uv run torchrun --nproc_per_node=4 --no-python flashdreams-run \ lingbot-world-fast --example-data True --total-blocks 21 ``` +## Run (shared demo API) + +The current unified-demo-runtime path is exposed through `lingbot-demo`. +From the repository root on a CUDA machine: + +```bash +export HF_TOKEN= + +uv sync --python 3.12 --package flashdreams-lingbot --no-dev +``` + +Generate a short MP4 from the bundled example assets: + +```bash +uv run --python 3.12 --package flashdreams-lingbot lingbot-demo replay \ + --device cuda:0 \ + --preset-id lingbot-world-v2-14b-causal-fast-taehv-window15-sink3 \ + --example-idx 0 \ + --total-blocks 10 \ + --fps 16 \ + --pixel-height 352 \ + --pixel-width 640 \ + --output outputs/lingbot-demo-replay.mp4 +``` + +Serve the shared WebRTC demo: + +```bash +uv run --python 3.12 --package flashdreams-lingbot lingbot-demo webrtc \ + --host 0.0.0.0 \ + --port 8089 \ + --device cuda:0 \ + --preset-id lingbot-world-v2-14b-causal-fast-taehv-window15-sink3 \ + --warmup-chunks 8 \ + --fps 16 \ + --video-height 352 \ + --video-width 640 \ + --example-idx 0 +``` + +Then open: + +- [http://localhost:8089/request_session](http://localhost:8089/request_session) +- [http://localhost:8089/healthz](http://localhost:8089/healthz) + ## Programmatic access Access via runner. @@ -143,7 +188,7 @@ for i in range(total_blocks): generated_chunks.append(video_chunk.cpu()) # each chunk is [T, C, H, W] ``` -## Run (WebRTC interactive demo) +## Run (compatibility WebRTC server) The `lingbot.webrtc` subpackage exposes a minimal WebRTC server that binds the integration pipeline to keyboard input over a DataChannel and streams the diff --git a/integrations/lingbot/lingbot/demo/__init__.py b/integrations/lingbot/lingbot/demo/__init__.py index a3da57837..08e4f9d1c 100644 --- a/integrations/lingbot/lingbot/demo/__init__.py +++ b/integrations/lingbot/lingbot/demo/__init__.py @@ -4,6 +4,7 @@ """Experimental Lingbot demo adapter built on ``flashdreams.runtime.demo``.""" from lingbot.demo.adapter import LingbotDemoAdapter +from lingbot.demo.providers import LingbotInputProvider from lingbot.demo.spec import ( DEFAULT_LINGBOT_PRESET, LINGBOT_MODEL_ID, @@ -15,6 +16,7 @@ "DEFAULT_LINGBOT_PRESET", "LINGBOT_MODEL_ID", "LingbotDemoAdapter", + "LingbotInputProvider", "LingbotReplayInputs", "LingbotWebRTCScenario", ] diff --git a/integrations/lingbot/lingbot/demo/adapter.py b/integrations/lingbot/lingbot/demo/adapter.py index f117cbfa7..1424f3eb0 100644 --- a/integrations/lingbot/lingbot/demo/adapter.py +++ b/integrations/lingbot/lingbot/demo/adapter.py @@ -18,6 +18,7 @@ DemoSpec, Mp4OutputSpec, PreparedScenario, + WebRTCOutputSpec, ) from flashdreams.runtime.interfaces import InferenceRuntime from lingbot.input_mapping import ( @@ -25,16 +26,22 @@ TextEventSelection, ) from lingbot.runtime import ( + FIELD_FPS, + FIELD_PIXEL_HEIGHT, + FIELD_PIXEL_WIDTH, + FIELD_TOTAL_BLOCKS, LingbotModelAdapter, LingbotReplayRuntime, PipelineFactory, inference_input_from_replay_inputs, ) +from .providers import LingbotInputProvider from .spec import ( resolve_replay_inputs, resolve_text_event_prompts, resolve_user_input_events, + resolve_webrtc_scenario, ) ReplayRuntimeFactory = Callable[..., InferenceRuntime] @@ -55,27 +62,37 @@ def __init__( ) def supported_input_modes(self) -> tuple[str, ...]: - return ("replay",) + return ("replay", "keyboard-driving") def supported_output_modes(self) -> tuple[str, ...]: - return ("mp4",) + return ("mp4", "webrtc") def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: - if spec.input_mode != "replay": + if spec.input_mode == "replay": + if not isinstance(spec.output, Mp4OutputSpec): + raise ValueError("Lingbot replay demo currently requires MP4 output.") + scenario = spec.scenario + live_camera = False + elif spec.input_mode == "keyboard-driving": + if not isinstance(spec.output, WebRTCOutputSpec): + raise ValueError( + "Lingbot keyboard-driving demo requires WebRTC output." + ) + scenario = _keyboard_driving_scenario(spec, output=spec.output) + live_camera = True + else: raise ValueError( - "Lingbot prepare_scenario currently supports only " - f"input_mode='replay', got {spec.input_mode!r}." + "Lingbot prepare_scenario supports input_mode='replay' or " + f"'keyboard-driving', got {spec.input_mode!r}." ) - if not isinstance(spec.output, Mp4OutputSpec): - raise ValueError("Lingbot replay demo currently requires MP4 output.") replay_inputs = resolve_replay_inputs( - spec.scenario, - default_prompt=self.default_replay_prompt(spec.config), + scenario, + default_prompt=_default_prompt(self, spec), ) - text_event_prompts = resolve_text_event_prompts(spec.scenario) - user_inputs = resolve_user_input_events(spec.scenario) - if _camera_source(spec.scenario) == "events": + text_event_prompts = resolve_text_event_prompts(scenario) + user_inputs = resolve_user_input_events(scenario) + if live_camera or _camera_source(scenario) == "events": # Live control still needs the scenario's calibration, so the trace # is loaded for its intrinsics and world scale and then discarded # as a trajectory source. @@ -99,7 +116,11 @@ def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: return PreparedScenario( initial_inputs=inference_input_from_replay_inputs(replay_inputs), user_inputs=user_inputs, - source_schema=_source_schema(user_inputs), + source_schema=_source_schema( + user_inputs, + include_keyboard=live_camera, + include_text_events=live_camera and bool(text_event_prompts), + ), canonicalizer=_canonicalizer(text_event_prompts), mapping=mapping, metadata={ @@ -108,6 +129,17 @@ def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: }, ) + def create_model_input_provider( + self, + spec: DemoSpec, + scenario: PreparedScenario, + ) -> LingbotInputProvider: + del spec + return LingbotInputProvider( + scenario=scenario, + inference_input_schema=self.inference_input_schema, + ) + def _camera_source(scenario: Any) -> str: if isinstance(scenario, Mapping): @@ -115,6 +147,48 @@ def _camera_source(scenario: Any) -> str: return "trace" +def _keyboard_driving_scenario( + spec: DemoSpec, + *, + output: WebRTCOutputSpec, +) -> Mapping[str, Any]: + webrtc_scenario = resolve_webrtc_scenario(spec.scenario) + scenario: dict[str, Any] = ( + dict(spec.scenario) if isinstance(spec.scenario, Mapping) else {} + ) + scenario.setdefault("camera_source", "events") + scenario.setdefault("example_data", True) + scenario.setdefault("example_idx", webrtc_scenario.example_idx) + scenario.setdefault(FIELD_TOTAL_BLOCKS, _total_blocks_default(spec)) + scenario.setdefault(FIELD_PIXEL_HEIGHT, output.video_height) + scenario.setdefault(FIELD_PIXEL_WIDTH, output.video_width) + scenario.setdefault(FIELD_FPS, output.fps) + config = spec.config + if config is not None and "text_events" not in scenario: + text_events = config.runtime_options.get("text_events") + if text_events is not None: + scenario["text_events"] = text_events + return scenario + + +def _total_blocks_default(spec: DemoSpec) -> int: + config = spec.config + if config is not None: + total_blocks = config.runtime_options.get("total_blocks") + if total_blocks is not None: + return int(total_blocks) + return 1_000_000 + + +def _default_prompt(adapter: LingbotDemoAdapter, spec: DemoSpec) -> str: + config = spec.config + if config is not None: + default_prompt = config.runtime_options.get("default_prompt") + if default_prompt is not None: + return str(default_prompt) + return adapter.default_replay_prompt(config) + + _KEY_EVENT_TYPES = frozenset({"key_down", "key_up"}) _KEYBOARD_CAPABILITIES = ( @@ -128,7 +202,12 @@ def _camera_source(scenario: Any) -> str: ) -def _source_schema(user_inputs: UserInputs) -> UserInputSchema: +def _source_schema( + user_inputs: UserInputs, + *, + include_keyboard: bool = False, + include_text_events: bool = False, +) -> UserInputSchema: """Declare what this scenario's event source can provide. Capabilities describe the source, not the particular trace. A keyboard @@ -140,9 +219,9 @@ def _source_schema(user_inputs: UserInputs) -> UserInputSchema: """ observed = {event.event_type for event in user_inputs.events} capabilities: list[UserInputCapability] = [] - if observed & _KEY_EVENT_TYPES: + if include_keyboard or observed & _KEY_EVENT_TYPES: capabilities.extend(_KEYBOARD_CAPABILITIES) - if "text_event" in observed: + if include_text_events or "text_event" in observed: capabilities.append(_TEXT_EVENT_CAPABILITY) for event_type in sorted(observed - _KEY_EVENT_TYPES - {"text_event"}): payload_fields: frozenset[str] = frozenset() diff --git a/integrations/lingbot/lingbot/demo/providers.py b/integrations/lingbot/lingbot/demo/providers.py new file mode 100644 index 000000000..b11db6834 --- /dev/null +++ b/integrations/lingbot/lingbot/demo/providers.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot model-input providers for shared demo run modes.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from flashdreams.runtime import ( + CanonicalInputs, + InferenceInput, + InferenceInputSchema, + StepRequest, + StepRequirements, + TimeWindow, + UserInputs, +) +from flashdreams.runtime.demo import ( + PreparedScenario, + PreparedStep, + ProviderCapabilities, + UserInputWindow, +) +from flashdreams.serving.webrtc.services import ( + WEBRTC_SKIPPED_INPUTS_METADATA_KEY, + WEBRTC_SKIPPED_WINDOW_METADATA_KEY, +) +from lingbot.input_mapping import CAMERA_COMMAND, LingbotInputMapping +from lingbot.runtime import LingbotModelAdapter + + +class LingbotInputProvider: + """Convert shared user-input windows into Lingbot model inputs. + + Lingbot's existing mapping API still accepts the legacy ``StepRequest`` + shape, because the runtime owns frame-start metadata today. This provider is + the model-owned bridge from shared demo drivers to that mapping boundary; + the bridge stays here so WebRTC transport code never has to know Lingbot's + camera or prompt semantics. + """ + + def __init__( + self, + *, + scenario: PreparedScenario, + inference_input_schema: InferenceInputSchema | None = None, + ) -> None: + mapping = scenario.mapping + if not isinstance(mapping, LingbotInputMapping): + raise TypeError( + "LingbotInputProvider requires PreparedScenario.mapping to be " + f"LingbotInputMapping, got {type(mapping).__name__}." + ) + if inference_input_schema is None: + inference_input_schema = LingbotModelAdapter().inference_input_schema + + self.capabilities = ProviderCapabilities( + supports_realtime_clock=_supports_realtime_clock(mapping), + supports_recorded_input=True, + supports_reset=True, + deterministic_given_inputs=True, + user_input_schema=scenario.source_schema, + inference_input_schema=inference_input_schema, + ) + self._scenario = scenario + self._mapping = mapping + self._step_base_inputs = InferenceInput( + step=scenario.initial_inputs.step, + metadata=scenario.initial_inputs.metadata, + ) + self._next_frame_start = 0 + self._closed = False + + def prepare_initial_input(self) -> InferenceInput: + self._require_open() + self._reset_state() + return self._mapping.map_global_conditioning_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=self._scenario.initial_inputs, + ) + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + self._require_open() + if user_window.control is not None: + return PreparedStep(control=user_window.control) + + self._advance_skipped_input_state(user_window) + legacy_request = self._legacy_step_request( + request=request, + user_window=user_window, + ) + assert legacy_request.user_input_window is not None + canonical_inputs = self._scenario.canonicalizer.canonicalize( + user_window.inputs, + window=legacy_request.user_input_window, + source_schema=self._scenario.source_schema, + ) + inference_input = self._mapping.map_step_inputs( + canonical_inputs=canonical_inputs, + inference_input=self._step_base_inputs, + request=legacy_request, + ) + self._next_frame_start = _required_int( + legacy_request.metadata, + "frame_start", + ) + _required_positive_int(legacy_request.metadata, "num_frames") + return PreparedStep(inference_input=inference_input) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self._require_open() + self._reset_state() + + def close(self) -> None: + self._closed = True + + def _legacy_step_request( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> StepRequest: + metadata: dict[str, Any] = dict(request.metadata) + metadata["num_frames"] = _metadata_positive_int( + metadata, + "num_frames", + default=request.input_frame_count, + ) + metadata["frame_start"] = _metadata_int( + metadata, + "frame_start", + default=self._next_frame_start, + ) + return StepRequest( + step_index=request.step_index, + inference_input_schema=request.inference_input_schema, + user_input_window=TimeWindow( + start_s=user_window.start_s, + end_s=user_window.end_s, + ), + metadata=metadata, + ) + + def _reset_state(self) -> None: + self._scenario.canonicalizer.reset() + self._mapping.reset() + self._next_frame_start = 0 + + def _advance_skipped_input_state(self, user_window: UserInputWindow) -> None: + skipped_inputs = user_window.metadata.get(WEBRTC_SKIPPED_INPUTS_METADATA_KEY) + skipped_window = user_window.metadata.get(WEBRTC_SKIPPED_WINDOW_METADATA_KEY) + if not isinstance(skipped_inputs, UserInputs): + return + if not isinstance(skipped_window, tuple) or len(skipped_window) != 2: + return + start_value, end_value = skipped_window + if not isinstance(start_value, int | float) or not isinstance( + end_value, + int | float, + ): + return + start_s = float(start_value) + end_s = float(end_value) + if end_s <= start_s: + return + self._scenario.canonicalizer.canonicalize( + skipped_inputs, + window=TimeWindow(start_s=start_s, end_s=end_s), + source_schema=self._scenario.source_schema, + ) + + def _require_open(self) -> None: + if self._closed: + raise RuntimeError("LingbotInputProvider is closed.") + + +def _supports_realtime_clock(mapping: LingbotInputMapping) -> bool: + return any( + modality.name == CAMERA_COMMAND.name + for modality in mapping.mapping_schema.consumes + ) + + +def _metadata_int( + metadata: Mapping[str, Any], + name: str, + *, + default: int, +) -> int: + if name not in metadata: + return default + return _required_int(metadata, name) + + +def _metadata_positive_int( + metadata: Mapping[str, Any], + name: str, + *, + default: int, +) -> int: + if name not in metadata: + if default <= 0: + raise ValueError(f"StepRequirements.{name} fallback must be > 0.") + return default + return _required_positive_int(metadata, name) + + +def _required_int(metadata: Mapping[str, Any], name: str) -> int: + value = metadata[name] + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"Step metadata {name!r} must be an integer.") + return value + + +def _required_positive_int(metadata: Mapping[str, Any], name: str) -> int: + value = _required_int(metadata, name) + if value <= 0: + raise ValueError(f"Step metadata {name!r} must be > 0.") + return value + + +__all__ = ["LingbotInputProvider"] diff --git a/integrations/lingbot/lingbot/demo/webrtc.py b/integrations/lingbot/lingbot/demo/webrtc.py index 4efa4152e..854326905 100644 --- a/integrations/lingbot/lingbot/demo/webrtc.py +++ b/integrations/lingbot/lingbot/demo/webrtc.py @@ -6,6 +6,7 @@ from __future__ import annotations from collections.abc import Callable +from dataclasses import replace from importlib.resources import files from typing import Any @@ -19,7 +20,13 @@ 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.demo import LingbotDemoAdapter from lingbot.runtime import ( + FIELD_FPS, + FIELD_PIXEL_HEIGHT, + FIELD_PIXEL_WIDTH, + FIELD_PROMPT, + FIELD_TOTAL_BLOCKS, LingbotModelAdapter, build_lingbot_webrtc_runtime_config, ) @@ -27,7 +34,6 @@ from lingbot.webrtc.session import ( LingbotInferenceRuntime, LingbotRuntimeConfig, - create_lingbot_webrtc_session_manager, ) from .spec import resolve_webrtc_scenario @@ -80,11 +86,30 @@ def serve_lingbot_webrtc_demo( runtime_options=config.runtime_options, ) runtime = runtime_factory(config=runtime_config) - manager = create_lingbot_webrtc_session_manager( + demo_adapter = LingbotDemoAdapter() + shared_spec = _shared_webrtc_spec( + spec, + runtime_config=runtime_config, + example_idx=scenario.example_idx, + ) + prepared = demo_adapter.prepare_scenario(shared_spec) + manager = BaseWebRTCSessionManager( runtime=runtime, runtime_config=runtime_config, fps=spec.output.fps, + identity=runtime_config.config_name, + busy_message="A Lingbot session is already active.", + warmup_label="Lingbot WebRTC", client_liveness_timeout_s=spec.output.client_liveness_timeout_s, + shared_adapter=demo_adapter, + shared_spec=shared_spec, + shared_spec_factory=lambda session_input: _shared_webrtc_spec( + spec, + runtime_config=runtime_config, + example_idx=scenario.example_idx, + session_input=session_input, + ), + shared_scenario=prepared, ) return serve_webrtc_demo( output=spec.output, @@ -105,6 +130,56 @@ def _option(config: InferenceConfig, name: str, default: Any) -> Any: return config.runtime_options.get(name, default) +def _shared_webrtc_spec( + spec: DemoSpec, + *, + runtime_config: LingbotRuntimeConfig, + example_idx: int, + session_input: Any = None, +) -> DemoSpec: + config = spec.config + if config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + runtime_options = dict(config.runtime_options) + runtime_options.update( + { + "default_prompt": runtime_config.default_prompt, + "pipeline_config": runtime_config.pipeline_config, + "seed": runtime_config.seed, + } + ) + scenario: dict[str, Any] = { + "camera_source": "events", + "example_data": True, + "example_idx": example_idx, + "text_events": runtime_config.text_events, + FIELD_TOTAL_BLOCKS: int(_option(config, "total_blocks", 1_000_000)), + FIELD_PIXEL_HEIGHT: runtime_config.video_height, + FIELD_PIXEL_WIDTH: runtime_config.video_width, + FIELD_FPS: runtime_config.fps, + } + # Browser-provided first-frame payloads stay runtime/session-owned because + # they can be bytes or remote payloads. The provider only needs the active + # prompt/catalog plus example-data calibration for live camera mapping. + prompt = getattr(session_input, "prompt", None) + if prompt: + scenario[FIELD_PROMPT] = str(prompt) + text_events = getattr(session_input, "text_events", None) + if text_events is not None: + scenario["text_events"] = text_events + return replace( + spec, + scenario=scenario, + config=replace( + config, + preset_id=runtime_config.config_name, + device=runtime_config.device, + seed=runtime_config.seed, + runtime_options=runtime_options, + ), + ) + + __all__ = [ "WebRTCRuntimeFactory", "serve_lingbot_webrtc_demo", diff --git a/integrations/lingbot/lingbot/input_mapping.py b/integrations/lingbot/lingbot/input_mapping.py index 0814d9819..245b0a112 100644 --- a/integrations/lingbot/lingbot/input_mapping.py +++ b/integrations/lingbot/lingbot/input_mapping.py @@ -622,6 +622,12 @@ def set_base_prompt(self, prompt: str) -> None: """Record the rollout prompt restored when a text event is cleared.""" self._base_prompt = prompt + def reset(self) -> None: + """Reset state accumulated while mapping a rollout.""" + self._applied_event_id = None + if self._integrator is not None: + self._integrator.reset() + def _pose_segments( command: Mapping[str, Any], diff --git a/integrations/lingbot/lingbot/runner.py b/integrations/lingbot/lingbot/runner.py index 7e4a139c4..9a4ffe0f0 100644 --- a/integrations/lingbot/lingbot/runner.py +++ b/integrations/lingbot/lingbot/runner.py @@ -24,9 +24,10 @@ from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig -from flashdreams.runtime import InputCanonicalizer, UserInputs, UserInputSchema -from flashdreams.runtime.metrics import NullMetricsRecorder -from flashdreams.runtime.runner import run_inference_session +from flashdreams.infra.runner_io import runner_artifact_path +from flashdreams.runtime.demo import DemoSpec, Mp4OutputSpec, OutputSpec +from flashdreams.runtime.demo.replay import run_replay_demo +from lingbot.demo import LingbotDemoAdapter from lingbot.example_data import ( EXAMPLE_DATA_AVAILABLE_IDXS, EXAMPLE_DATA_PROMPT_AVAILABLE_IDXS, @@ -37,10 +38,9 @@ LingbotWorldInferencePipeline, ) from lingbot.runtime import ( - LingbotModelAdapter, + LINGBOT_MODEL_ID, LingbotRunnerOutputTarget, inference_config_from_runner_config, - inference_input_from_replay_inputs, replay_inputs_from_runner_config, ) @@ -166,7 +166,7 @@ def _fill_example_data_defaults(self) -> None: def run(self) -> None: """Drive an AR rollout through the Lingbot runtime API path.""" cfg = self.config - adapter = LingbotModelAdapter() + adapter = LingbotDemoAdapter() inference_config = inference_config_from_runner_config( cfg, device=f"cuda:{self.local_rank}" if self.world_size > 1 else cfg.device, @@ -176,22 +176,37 @@ def run(self) -> None: cfg, is_rank_zero=self.is_rank_zero, ) - initial_inputs = inference_input_from_replay_inputs(replay_inputs) - output_target = LingbotRunnerOutputTarget( - output_stream=self.create_video_output_stream(fps=cfg.fps), - output_dir=cfg.output_dir, - runner_name=cfg.runner_name, - fps=cfg.fps, + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=str(cfg.pipeline.name), + input_mode="replay", + output=Mp4OutputSpec( + path=runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4"), + fps=cfg.fps, + output_layout=cfg.postprocess_output_layout or "tchw", + ), + scenario=replay_inputs, + config=inference_config, ) - mapping = adapter.create_input_mapping(replay_inputs) - run_inference_session( + + def _output_target_factory( + output_spec: OutputSpec, + ) -> LingbotRunnerOutputTarget: + del output_spec + return LingbotRunnerOutputTarget( + output_stream=self.create_video_output_stream(fps=cfg.fps), + output_dir=cfg.output_dir, + runner_name=cfg.runner_name, + fps=cfg.fps, + ) + + result = run_replay_demo( + spec=spec, adapter=adapter, - config=inference_config, - mapping=mapping, - canonicalizer=InputCanonicalizer(), - source_schema=UserInputSchema(description="Lingbot runner fixed inputs"), - user_inputs=UserInputs(), - initial_inputs=initial_inputs, - output=output_target, - metrics=NullMetricsRecorder(), + output_target_factory=_output_target_factory, ) + if result.status != "completed": + raise RuntimeError( + f"Lingbot runner failed with status {result.status!r}: " + f"{result.reason or result.error or 'unknown error'}" + ) diff --git a/integrations/lingbot/tests/test_demo_api.py b/integrations/lingbot/tests/test_demo_api.py index 6682dfa5e..8b33b865b 100644 --- a/integrations/lingbot/tests/test_demo_api.py +++ b/integrations/lingbot/tests/test_demo_api.py @@ -15,6 +15,7 @@ DEFAULT_LINGBOT_PRESET, LINGBOT_MODEL_ID, LingbotDemoAdapter, + LingbotInputProvider, LingbotReplayInputs, LingbotWebRTCScenario, ) @@ -37,7 +38,11 @@ FIELD_TOTAL_BLOCKS, inference_input_from_replay_inputs, ) -from lingbot.webrtc.session import LingbotRuntimeConfig +from lingbot.webrtc.session import ( + LingbotRuntimeConfig, + LingbotSessionInput, + TextEventSpec, +) from flashdreams.runtime import ( CanonicalInputs, @@ -46,11 +51,15 @@ OutputArtifact, OutputTarget, StepRequest, + StepRequirements, StepResult, + UserInputEvent, + UserInputs, ) from flashdreams.runtime.demo import ( DemoSpec, Mp4OutputSpec, + UserInputWindow, WebRTCOutputSpec, ) from flashdreams.runtime.demo.replay import run_replay_demo @@ -71,18 +80,45 @@ def _write_camera_assets(poses: Path, intrinsics: Path, *, frames: int = 64) -> ) +def _patch_lingbot_webrtc_example( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + *, + example_idx: int = 0, +) -> Path: + """Provide a local example-data directory for shared WebRTC preparation.""" + import lingbot.runtime as runtime_module + + example_dir = tmp_path / f"example-{example_idx:02d}" + example_dir.mkdir() + (example_dir / "image.jpg").write_bytes(b"fake") + _write_camera_assets(example_dir / "poses.npy", example_dir / "intrinsics.npy") + (example_dir / "prompt.txt").write_text("drive through a forest\n") + + def fake_download(*, is_rank_zero: bool, example_idx: int) -> Path: + del is_rank_zero, example_idx + return example_dir + + monkeypatch.setattr( + runtime_module, + "ensure_example_data_downloaded", + fake_download, + ) + return example_dir + + def test_lingbot_demo_defaults_to_interactive_preset() -> None: args = parse_args(["replay", "--output", "demo.mp4"]) assert args.preset_id == "lingbot-world-fast-taehv-window15-sink3" -def test_lingbot_demo_adapter_declares_replay_modes_only() -> None: +def test_lingbot_demo_adapter_declares_shared_demo_modes() -> None: adapter = LingbotDemoAdapter() assert adapter.model_id == LINGBOT_MODEL_ID - assert adapter.supported_input_modes() == ("replay",) - assert adapter.supported_output_modes() == ("mp4",) + assert adapter.supported_input_modes() == ("replay", "keyboard-driving") + assert adapter.supported_output_modes() == ("mp4", "webrtc") fields = { field.name for field in adapter.inference_input_schema.global_conditioning_fields @@ -155,6 +191,101 @@ def fake_runner(**kwargs: Any) -> Sequence[OutputArtifact]: assert inputs[FIELD_TOTAL_BLOCKS] == 1 +def test_lingbot_replay_demo_run_mode_uses_model_provider( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lingbot.runtime as runtime_module + + import flashdreams.runtime.demo.replay as replay_module + + image = tmp_path / "image.jpg" + poses = tmp_path / "poses.npy" + intrinsics = tmp_path / "intrinsics.npy" + image.write_bytes(b"fake") + _write_camera_assets(poses, intrinsics, frames=16) + pipeline = _FakeLingbotPipeline() + driver_calls: list[dict[str, Any]] = [] + pipeline_calls: list[dict[str, Any]] = [] + + monkeypatch.setattr( + runtime_module, + "load_first_frame_tensor", + lambda *args, **kwargs: torch.zeros(1, 3, 2, 2), + ) + + class RecordingBatchSessionDriver(replay_module.BatchSessionDriver): + def run_one_session(self, **kwargs: Any) -> Any: + driver_calls.append(kwargs) + return super().run_one_session(**kwargs) + + class RecordingStepPipeline(replay_module.StepPipeline): + def execute_step(self, **kwargs: Any) -> Any: + pipeline_calls.append(kwargs) + return super().execute_step(**kwargs) + + monkeypatch.setattr( + replay_module, + "BatchSessionDriver", + RecordingBatchSessionDriver, + ) + monkeypatch.setattr(replay_module, "StepPipeline", RecordingStepPipeline) + + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + input_mode="replay", + scenario={ + "prompt": "drive through a city", + "image_path": image, + "pose_path": poses, + "intrinsic_path": intrinsics, + "total_blocks": 1, + }, + output=Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=16, output_layout="tchw"), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + device="cpu", + runtime_options={"pipeline_config": object()}, + ), + ) + expected_mapping = LingbotDemoAdapter().create_input_mapping( + LingbotReplayInputs( + prompt="drive through a city", + first_frame_path=image, + camera_poses_path=poses, + camera_intrinsics_path=intrinsics, + total_blocks=1, + ) + ) + + def pipeline_factory(pipeline_config: object, device: str) -> _FakeLingbotPipeline: + del pipeline_config, device + return pipeline + + adapter = LingbotDemoAdapter(pipeline_factory=pipeline_factory) + + result = run_replay_demo( + spec=spec, + adapter=adapter, + output_target_factory=lambda output_spec: _RecordingOutputTarget(), + ) + + assert result.status == "completed" + assert len(driver_calls) == 1 + assert len(pipeline_calls) == 1 + assert isinstance(pipeline_calls[0]["provider"], LingbotInputProvider) + assert pipeline.generate_calls == [ + { + "autoregressive_index": 0, + "intrinsics_shape": (1, 4), + "poses_shape": (1, 4, 4), + "world_scale": pytest.approx(expected_mapping.camera_trace.world_scale), + } + ] + + def test_lingbot_replay_invalid_scenario_fails_before_runtime_creation( tmp_path: Path, ) -> None: @@ -382,7 +513,48 @@ def test_lingbot_webrtc_cli_builds_keyboard_driving_spec() -> None: assert spec.config.runtime_options["example_idx"] == 2 -def test_lingbot_webrtc_demo_uses_existing_manager_with_model_config() -> None: +def test_lingbot_adapter_prepares_public_webrtc_scenario_as_live_provider( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + _patch_lingbot_webrtc_example(monkeypatch, tmp_path) + adapter = LingbotDemoAdapter() + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + input_mode="keyboard-driving", + scenario=LingbotWebRTCScenario(example_idx=0), + output=WebRTCOutputSpec(fps=16, video_width=64, video_height=32), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + runtime_options={ + "text_events": [ + { + "event_id": "storm", + "label": "Storm", + "prompt": "A storm moves through the scene.", + } + ] + }, + ), + ) + + prepared = adapter.prepare_scenario(spec) + provider = adapter.create_model_input_provider(spec, prepared) + + assert isinstance(provider, LingbotInputProvider) + assert provider.capabilities.supports_realtime_clock is True + assert {"key_down", "key_up", "text_event"}.issubset( + prepared.source_schema.declared_event_types() + ) + + +def test_lingbot_webrtc_demo_uses_shared_manager_with_model_config( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + _patch_lingbot_webrtc_example(monkeypatch, tmp_path, example_idx=2) pipeline_config = object() spec = DemoSpec( model_id=LINGBOT_MODEL_ID, @@ -418,6 +590,17 @@ def test_lingbot_webrtc_demo_uses_existing_manager_with_model_config() -> None: runtime = manager._runtime assert isinstance(runtime, _FakeWebRTCRuntime) assert type(manager) is BaseWebRTCSessionManager + assert isinstance(manager._shared_adapter, LingbotDemoAdapter) + assert manager._shared_spec is not None + assert manager._shared_spec.input_mode == "keyboard-driving" + assert isinstance(manager._shared_spec.output, WebRTCOutputSpec) + assert manager._shared_scenario is not None + provider = manager._shared_adapter.create_model_input_provider( + manager._shared_spec, + manager._shared_scenario, + ) + assert isinstance(provider, LingbotInputProvider) + assert provider.capabilities.supports_realtime_clock is True assert manager.runtime_config is runtime.config assert runtime.config.pipeline_config is pipeline_config assert runtime.config.config_name == DEFAULT_LINGBOT_PRESET @@ -426,18 +609,113 @@ def test_lingbot_webrtc_demo_uses_existing_manager_with_model_config() -> None: assert runtime.config.video_width == 64 assert runtime.config.video_height == 32 assert runtime.config.fps == 24 + assert runtime.config.warmup_chunks == 0 + assert runtime.config.warmup_timeout_s == 1.0 assert runtime.config.encoder_backend == "default" assert runtime.config.example_data_dir.name == "02" + assert isinstance(spec.output, WebRTCOutputSpec) + assert manager.client_liveness_timeout_s == spec.output.client_liveness_timeout_s assert manager.identity == DEFAULT_LINGBOT_PRESET assert calls[0]["host"] == "0.0.0.0" assert calls[0]["port"] == 8080 +def test_lingbot_webrtc_shared_provider_reflects_pending_session_input( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + _patch_lingbot_webrtc_example(monkeypatch, tmp_path) + pipeline_config = object() + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + input_mode="keyboard-driving", + scenario=LingbotWebRTCScenario(example_idx=0), + output=WebRTCOutputSpec( + fps=16, + video_width=64, + video_height=32, + warmup_chunks=8, + warmup_timeout_s=2.5, + ), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + runtime_options={"pipeline_config": pipeline_config}, + ), + ) + calls: list[dict[str, Any]] = [] + pending = LingbotSessionInput( + prompt="drive through a custom city", + text_events=( + TextEventSpec( + event_id="rain", + label="Rain", + prompt="heavy rain falls on the road", + ), + ), + ) + + serve_lingbot_webrtc_demo( + spec=spec, + runtime_factory=_FakeWebRTCRuntime, + server_runner=lambda **kwargs: calls.append(kwargs), + ) + + manager = calls[0]["session_manager"] + assert callable(manager._shared_spec_factory) + session_spec = manager._shared_spec_factory(pending) + assert isinstance(session_spec.output, WebRTCOutputSpec) + assert session_spec.output.video_width == 64 + assert session_spec.output.video_height == 32 + assert session_spec.output.warmup_chunks == 8 + prepared = manager._shared_adapter.prepare_scenario(session_spec) + provider = manager._shared_adapter.create_model_input_provider( + session_spec, + prepared, + ) + + initial = provider.prepare_initial_input() + assert initial.global_conditioning[FIELD_PROMPT] == "drive through a custom city" + assert initial.global_conditioning[FIELD_PIXEL_HEIGHT] == 32 + assert initial.global_conditioning[FIELD_PIXEL_WIDTH] == 64 + assert {"key_down", "key_up", "text_event"}.issubset( + prepared.source_schema.declared_event_types() + ) + prepared_step = provider.prepare_step( + request=StepRequirements( + step_index=0, + input_frame_count=4, + metadata={"frame_start": 0, "num_frames": 4}, + ), + user_window=UserInputWindow( + start_s=0.0, + end_s=0.25, + inputs=UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, + event_type="text_event", + payload={"event_id": "rain"}, + ), + ) + ), + ), + ) + assert prepared_step.inference_input is not None + assert ( + prepared_step.inference_input.global_conditioning[FIELD_PROMPT] + == "heavy rain falls on the road" + ) + + def test_lingbot_webrtc_demo_uses_shared_viewer_shell( monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, ) -> None: import flashdreams.runtime.demo.webrtc as shared_webrtc_module + _patch_lingbot_webrtc_example(monkeypatch, tmp_path) app_calls: list[dict[str, Any]] = [] def fake_create_packaged_app(**kwargs: Any) -> web.Application: @@ -492,9 +770,11 @@ def fake_create_packaged_app(**kwargs: Any) -> web.Application: def test_lingbot_webrtc_demo_serves_through_shared_runner( monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, ) -> None: import flashdreams.runtime.demo.webrtc as shared_webrtc_module + _patch_lingbot_webrtc_example(monkeypatch, tmp_path) server_calls: list[dict[str, Any]] = [] def fake_create_packaged_app(**kwargs: Any) -> web.Application: diff --git a/integrations/lingbot/tests/test_demo_providers.py b/integrations/lingbot/tests/test_demo_providers.py new file mode 100644 index 000000000..db7425da9 --- /dev/null +++ b/integrations/lingbot/tests/test_demo_providers.py @@ -0,0 +1,433 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +import torch +from lingbot.demo import DEFAULT_LINGBOT_PRESET, LINGBOT_MODEL_ID, LingbotDemoAdapter +from lingbot.demo.providers import LingbotInputProvider +from lingbot.input_mapping import ( + FIELD_CAMERA_TRAJECTORY, + FIELD_TOTAL_CAMERA_FRAMES, + LingbotInputMapping, +) +from lingbot.runtime import ( + FIELD_FIRST_FRAME_PATH, + FIELD_PROMPT, + FIELD_TOTAL_BLOCKS, +) + +from flashdreams.runtime import ( + InferenceConfig, + InferenceInput, + StepRequest, + StepRequirements, + TimeWindow, + UserInputEvent, + UserInputs, +) +from flashdreams.runtime.demo import DemoSpec, Mp4OutputSpec, PreparedScenario +from flashdreams.runtime.demo.session_inputs import UserInputWindow +from flashdreams.serving.webrtc.services import ( + WEBRTC_SKIPPED_INPUTS_METADATA_KEY, + WEBRTC_SKIPPED_WINDOW_METADATA_KEY, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_lingbot_provider_initial_input_matches_mapping_path(tmp_path: Path) -> None: + adapter = LingbotDemoAdapter() + expected = _prepared_scenario(tmp_path, adapter=adapter) + actual = _prepared_scenario(tmp_path, adapter=adapter) + assert isinstance(expected.mapping, LingbotInputMapping) + + provider = LingbotInputProvider( + scenario=actual, + inference_input_schema=adapter.inference_input_schema, + ) + + expected_initial = expected.mapping.map_global_conditioning_inputs( + canonical_inputs=expected.canonicalizer.canonicalize( + UserInputs(), + window=TimeWindow(start_s=0.0, end_s=0.0), + source_schema=expected.source_schema, + ), + inference_input=expected.initial_inputs, + ) + actual_initial = provider.prepare_initial_input() + + assert actual_initial.global_conditioning == expected_initial.global_conditioning + assert actual_initial.step == expected_initial.step + assert actual_initial.metadata == expected_initial.metadata + assert actual_initial.global_conditioning[FIELD_PROMPT] == "drive through a city" + assert actual_initial.global_conditioning[FIELD_FIRST_FRAME_PATH] == ( + tmp_path / "image.jpg" + ) + assert actual_initial.global_conditioning[FIELD_TOTAL_BLOCKS] == 2 + assert FIELD_TOTAL_CAMERA_FRAMES in actual_initial.global_conditioning + + +def test_lingbot_provider_trace_steps_match_mapping_path(tmp_path: Path) -> None: + adapter = LingbotDemoAdapter() + expected = _prepared_scenario(tmp_path, adapter=adapter) + actual = _prepared_scenario(tmp_path, adapter=adapter) + provider = LingbotInputProvider( + scenario=actual, + inference_input_schema=adapter.inference_input_schema, + ) + + provider.prepare_initial_input() + expected_first = _legacy_step(expected, step_index=0, frame_start=0, num_frames=4) + actual_first = _provider_step( + provider, + step_index=0, + frame_start=0, + num_frames=4, + inputs=actual.user_inputs, + ) + expected_second = _legacy_step(expected, step_index=1, frame_start=4, num_frames=4) + actual_second = _provider_step( + provider, + step_index=1, + frame_start=4, + num_frames=4, + inputs=actual.user_inputs, + ) + + assert actual_first.global_conditioning == expected_first.global_conditioning + assert torch.allclose( + actual_first.step[FIELD_CAMERA_TRAJECTORY], + expected_first.step[FIELD_CAMERA_TRAJECTORY], + ) + assert torch.allclose( + actual_second.step[FIELD_CAMERA_TRAJECTORY], + expected_second.step[FIELD_CAMERA_TRAJECTORY], + ) + assert not torch.allclose( + actual_first.step[FIELD_CAMERA_TRAJECTORY], + actual_second.step[FIELD_CAMERA_TRAJECTORY], + ) + + +def test_lingbot_provider_uses_driver_user_window_inputs(tmp_path: Path) -> None: + adapter = LingbotDemoAdapter() + scenario_events = ({"t": 10.0, "type": "key_down", "key": "a"},) + expected = _prepared_scenario( + tmp_path, + adapter=adapter, + camera_source="events", + events=scenario_events, + ) + actual = _prepared_scenario( + tmp_path, + adapter=adapter, + camera_source="events", + events=scenario_events, + ) + provider = LingbotInputProvider( + scenario=actual, + inference_input_schema=adapter.inference_input_schema, + ) + window_inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, + event_type="key_down", + payload={"key": "w"}, + ), + ) + ) + + provider.prepare_initial_input() + expected_step = _legacy_step( + expected, + step_index=0, + frame_start=0, + num_frames=4, + inputs=window_inputs, + ) + actual_step = _provider_step( + provider, + step_index=0, + frame_start=0, + num_frames=4, + inputs=window_inputs, + ) + + poses = actual_step.step[FIELD_CAMERA_TRAJECTORY] + assert torch.allclose(poses, expected_step.step[FIELD_CAMERA_TRAJECTORY]) + assert not torch.allclose(poses[0], poses[-1]) + + +def test_lingbot_provider_folds_webrtc_skipped_inputs_into_state( + tmp_path: Path, +) -> None: + adapter = LingbotDemoAdapter() + scenario_events = ({"t": 10.0, "type": "key_down", "key": "a"},) + with_skip = _prepared_scenario( + tmp_path, + adapter=adapter, + camera_source="events", + events=scenario_events, + ) + idle = _prepared_scenario( + tmp_path, + adapter=adapter, + camera_source="events", + events=scenario_events, + ) + provider = LingbotInputProvider( + scenario=with_skip, + inference_input_schema=adapter.inference_input_schema, + ) + idle_provider = LingbotInputProvider( + scenario=idle, + inference_input_schema=adapter.inference_input_schema, + ) + skipped_inputs = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, + event_type="key_down", + payload={"key": "w"}, + ), + ) + ) + + provider.prepare_initial_input() + idle_provider.prepare_initial_input() + actual = _provider_step( + provider, + step_index=0, + frame_start=4, + num_frames=4, + inputs=UserInputs(), + metadata={ + WEBRTC_SKIPPED_INPUTS_METADATA_KEY: skipped_inputs, + WEBRTC_SKIPPED_WINDOW_METADATA_KEY: (0.0, 0.25), + }, + ) + expected_idle = _provider_step( + idle_provider, + step_index=0, + frame_start=4, + num_frames=4, + inputs=UserInputs(), + ) + + poses = actual.step[FIELD_CAMERA_TRAJECTORY] + assert not torch.allclose(poses, expected_idle.step[FIELD_CAMERA_TRAJECTORY]) + assert not torch.allclose(poses[0], poses[-1]) + + +def test_lingbot_provider_reset_clears_text_event_state(tmp_path: Path) -> None: + adapter = LingbotDemoAdapter() + actual = _prepared_scenario( + tmp_path, + adapter=adapter, + camera_source="events", + text_events={"storm": "a violent storm"}, + events=( + {"t": 10.0, "type": "key_down", "key": "w"}, + {"t": 10.1, "type": "text_event", "event_id": "storm"}, + ), + ) + provider = LingbotInputProvider( + scenario=actual, + inference_input_schema=adapter.inference_input_schema, + ) + first_inputs = _window_inputs_with_text_event(timestamp_s=0.0) + + provider.prepare_initial_input() + first = _provider_step( + provider, + step_index=0, + frame_start=0, + num_frames=4, + inputs=first_inputs, + ) + repeated = _provider_step( + provider, + step_index=1, + frame_start=4, + num_frames=4, + inputs=UserInputs(), + ) + provider.reset() + after_reset = _provider_step( + provider, + step_index=0, + frame_start=0, + num_frames=4, + inputs=first_inputs, + ) + + assert first.global_conditioning[FIELD_PROMPT] == "a violent storm" + assert repeated.global_conditioning == {} + assert after_reset.global_conditioning[FIELD_PROMPT] == "a violent storm" + + +@pytest.mark.parametrize( + ("metadata", "match"), + [ + ({"frame_start": "0", "num_frames": 4}, "frame_start"), + ({"frame_start": 0.5, "num_frames": 4}, "frame_start"), + ({"frame_start": 0, "num_frames": "4"}, "num_frames"), + ({"frame_start": 0, "num_frames": 4.5}, "num_frames"), + ], +) +def test_lingbot_provider_rejects_non_integer_frame_metadata( + tmp_path: Path, + metadata: dict[str, object], + match: str, +) -> None: + provider = LingbotInputProvider( + scenario=_prepared_scenario(tmp_path, adapter=LingbotDemoAdapter()) + ) + provider.prepare_initial_input() + + with pytest.raises(TypeError, match=match): + provider.prepare_step( + request=StepRequirements( + step_index=0, + input_frame_count=4, + metadata=metadata, + ), + user_window=UserInputWindow( + start_s=0.0, + end_s=0.25, + inputs=UserInputs(), + ), + ) + + +def _prepared_scenario( + tmp_path: Path, + *, + adapter: LingbotDemoAdapter, + camera_source: str = "trace", + text_events: dict[str, str] | None = None, + events: tuple[dict[str, Any], ...] = (), +) -> PreparedScenario: + image = tmp_path / "image.jpg" + poses = tmp_path / "poses.npy" + intrinsics = tmp_path / "intrinsics.npy" + image.write_bytes(b"fake") + _write_camera_assets(poses, intrinsics) + scenario: dict[str, Any] = { + "prompt": "drive through a city", + "image_path": image, + "pose_path": poses, + "intrinsic_path": intrinsics, + "camera_source": camera_source, + "total_blocks": 2, + } + if text_events is not None: + scenario["text_events"] = text_events + if events: + scenario["events"] = list(events) + return adapter.prepare_scenario( + DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + input_mode="replay", + scenario=scenario, + output=Mp4OutputSpec(path=tmp_path / "demo.mp4", fps=16), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + runtime_options={"pipeline_config": object()}, + ), + ) + ) + + +def _write_camera_assets(poses: Path, intrinsics: Path, *, frames: int = 32) -> None: + trajectory = np.tile(np.eye(4, dtype=np.float32), (frames, 1, 1)) + trajectory[:, 2, 3] = np.arange(frames, dtype=np.float32) + np.save(poses, trajectory) + np.save( + intrinsics, + np.tile(np.array([416.0, 416.0, 416.0, 240.0], dtype=np.float32), (frames, 1)), + ) + + +def _legacy_step( + scenario: PreparedScenario, + *, + step_index: int, + frame_start: int, + num_frames: int, + inputs: UserInputs | None = None, +) -> InferenceInput: + mapping = scenario.mapping + assert isinstance(mapping, LingbotInputMapping) + window = TimeWindow( + start_s=frame_start / 16, + end_s=(frame_start + num_frames) / 16, + ) + return mapping.map_step_inputs( + canonical_inputs=scenario.canonicalizer.canonicalize( + scenario.user_inputs if inputs is None else inputs, + window=window, + source_schema=scenario.source_schema, + ), + inference_input=InferenceInput( + step=scenario.initial_inputs.step, + metadata=scenario.initial_inputs.metadata, + ), + request=StepRequest( + step_index=step_index, + user_input_window=window, + metadata={"frame_start": frame_start, "num_frames": num_frames}, + ), + ) + + +def _provider_step( + provider: LingbotInputProvider, + *, + step_index: int, + frame_start: int, + num_frames: int, + inputs: UserInputs, + metadata: dict[str, object] | None = None, +) -> InferenceInput: + prepared = provider.prepare_step( + request=StepRequirements( + step_index=step_index, + input_frame_count=num_frames, + metadata={"frame_start": frame_start, "num_frames": num_frames}, + ), + user_window=UserInputWindow( + start_s=frame_start / 16, + end_s=(frame_start + num_frames) / 16, + inputs=inputs, + metadata={} if metadata is None else metadata, + ), + ) + assert prepared.inference_input is not None + return prepared.inference_input + + +def _window_inputs_with_text_event(*, timestamp_s: float) -> UserInputs: + return UserInputs( + events=( + UserInputEvent( + timestamp_s=timestamp_s, + event_type="key_down", + payload={"key": "w"}, + ), + UserInputEvent( + timestamp_s=timestamp_s, + event_type="text_event", + payload={"event_id": "storm"}, + ), + ) + ) diff --git a/integrations/lingbot/tests/test_smoke.py b/integrations/lingbot/tests/test_smoke.py index a787431a3..171bacb89 100644 --- a/integrations/lingbot/tests/test_smoke.py +++ b/integrations/lingbot/tests/test_smoke.py @@ -18,6 +18,7 @@ from __future__ import annotations import sys +from collections.abc import Callable from pathlib import Path from typing import cast @@ -42,11 +43,8 @@ example_data_dirname, ) from lingbot.runtime import ( - FIELD_FIRST_FRAME_PATH, - FIELD_PROMPT, - FIELD_TOTAL_BLOCKS, LINGBOT_MODEL_ID, - LingbotModelAdapter, + LingbotReplayInputs, LingbotRunnerOutputTarget, ) from lingbot.transformer import ( @@ -57,7 +55,8 @@ from flashdreams.infra.config import derive_config from flashdreams.infra.runner import RunnerConfig -from flashdreams.runtime import InferenceConfig, InferenceInput +from flashdreams.runtime import InferenceConfig +from flashdreams.runtime.demo import DemoSpec, Mp4OutputSpec, RunResult pytestmark = pytest.mark.ci_cpu @@ -189,7 +188,7 @@ def test_runner_delegates_to_runtime_api_with_direct_inputs( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - """Keep the CLI runner on the new runtime path, not the old rollout loop.""" + """Keep the CLI runner on the shared replay runtime path.""" image = tmp_path / "image.jpg" poses = tmp_path / "poses.npy" intrinsics = tmp_path / "intrinsics.npy" @@ -206,15 +205,16 @@ def test_runner_delegates_to_runtime_api_with_direct_inputs( intrinsic_path=intrinsics, total_blocks=1, device="cpu", + output_dir=tmp_path, ), ) pipeline = object() output_stream = object() captured: dict[str, object] = {} - def _fake_run_inference_session(**kwargs: object) -> tuple[object, ...]: + def _fake_run_replay_demo(**kwargs: object) -> RunResult: captured.update(kwargs) - return () + return RunResult(status="completed") monkeypatch.setattr( runner, @@ -223,8 +223,8 @@ def _fake_run_inference_session(**kwargs: object) -> tuple[object, ...]: ) monkeypatch.setattr( runner_mod, - "run_inference_session", - _fake_run_inference_session, + "run_replay_demo", + _fake_run_replay_demo, ) runner.config = runner_config runner.pipeline = pipeline @@ -234,19 +234,32 @@ def _fake_run_inference_session(**kwargs: object) -> tuple[object, ...]: runner.run() - assert isinstance(captured["adapter"], LingbotModelAdapter) - config = captured["config"] + spec = captured["spec"] + assert isinstance(spec, DemoSpec) + assert spec.model_id == LINGBOT_MODEL_ID + assert spec.input_mode == "replay" + assert spec.preset_id == str(runner_config.pipeline.name) + config = spec.config assert isinstance(config, InferenceConfig) assert config.model_id == LINGBOT_MODEL_ID assert config.device == "cpu" assert config.runtime_options["pipeline"] is pipeline - 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 - output = captured["output"] + scenario = spec.scenario + assert isinstance(scenario, LingbotReplayInputs) + assert scenario.prompt == "drive through a city" + assert scenario.first_frame_path == image + assert scenario.total_blocks == 1 + output_spec = spec.output + assert isinstance(output_spec, Mp4OutputSpec) + assert output_spec.path == tmp_path / f"{runner_config.runner_name}.mp4" + assert output_spec.fps == runner_config.fps + assert output_spec.output_layout == "tchw" + output_target_factory = cast( + Callable[[Mp4OutputSpec], LingbotRunnerOutputTarget], + captured["output_target_factory"], + ) + assert callable(output_target_factory) + output = output_target_factory(output_spec) assert isinstance(output, LingbotRunnerOutputTarget) assert output.output_stream is output_stream From b982790eb03a8633296af4070fd3f2d708b1ac3c Mon Sep 17 00:00:00 2001 From: jarcherNV Date: Mon, 10 Aug 2026 22:25:32 -0700 Subject: [PATCH 18/19] Migrate LingBot demo onto unified runtime (#437) Migrate LingBot demo onto unified runtime Port LingBot replay, null, MP4, and WebRTC demo paths to the shared demo runtime, add focused CI coverage for null and MP4 outputs, and update demo docs with the validated commands. Clean up shared keyboard/input ownership, preserve legacy compatibility where still needed, and improve shared WebRTC responsiveness for interactive demos. --- .github/workflows/lingbot-demo-runtime.yml | 334 +++++++++++++++ .../flashdreams/core/checkpoint/load.py | 362 +++++++++++++++- .../recipes/wan/transformer/wan21.py | 55 ++- flashdreams/flashdreams/runtime/__init__.py | 22 + flashdreams/flashdreams/runtime/canonical.py | 2 +- .../flashdreams/runtime/demo/__init__.py | 8 +- flashdreams/flashdreams/runtime/demo/app.py | 6 +- .../flashdreams/runtime/demo/bootstrap.py | 98 +++++ .../flashdreams/runtime/demo/timing.py | 74 ++-- .../flashdreams/runtime/demo/webrtc.py | 114 +---- flashdreams/flashdreams/runtime/keyboard.py | 136 ++++++ .../flashdreams/serving/realtime/input.py | 401 ++---------------- .../flashdreams/serving/webrtc/bootstrap.py | 88 +--- .../flashdreams/serving/webrtc/controls.py | 8 +- .../flashdreams/serving/webrtc/demo.py | 113 +++++ .../flashdreams/serving/webrtc/manager.py | 92 ++-- .../flashdreams/serving/webrtc/runtime.py | 9 +- .../flashdreams/serving/webrtc/services.py | 107 ++++- .../serving/webrtc/web/request_session.js | 12 - flashdreams/tests/test_checkpoint_loading.py | 111 +++++ flashdreams/tests/test_demo_runtime_timing.py | 53 +-- flashdreams/tests/test_realtime_serving.py | 4 +- flashdreams/tests/test_runtime_demo_api.py | 2 +- flashdreams/tests/test_webrtc_manager.py | 91 ++-- flashdreams/tests/test_webrtc_services.py | 62 ++- flashdreams/tests/test_webrtc_serving.py | 53 +-- integrations/lingbot/README.md | 14 + integrations/lingbot/lingbot/config.py | 4 +- integrations/lingbot/lingbot/controls.py | 283 ++++++++++++ integrations/lingbot/lingbot/demo/adapter.py | 7 +- integrations/lingbot/lingbot/demo/app.py | 34 +- integrations/lingbot/lingbot/demo/webrtc.py | 4 +- integrations/lingbot/lingbot/input_mapping.py | 5 +- integrations/lingbot/lingbot/runtime.py | 5 + .../lingbot/lingbot/webrtc/session.py | 17 +- .../lingbot/lingbot/webrtc/web/adapter.css | 19 +- .../lingbot/lingbot/webrtc/web/adapter.js | 21 +- integrations/lingbot/tests/test_controls.py | 6 +- integrations/lingbot/tests/test_demo_api.py | 121 +++++- .../lingbot/tests/test_keyboard_parity.py | 2 +- integrations/lingbot/tests/test_smoke.py | 8 + .../tests/test_webrtc_session_branch.py | 6 +- .../omnidreams/omnidreams/demo/controls.py | 305 +++++++++++++ .../omnidreams/omnidreams/demo/providers.py | 94 ++-- .../omnidreams/omnidreams/demo/webrtc.py | 6 +- .../omnidreams/demo/webrtc_legacy.py | 18 +- .../interactive_drive/input/keyboard.py | 4 +- .../omnidreams/tests/test_demo_api.py | 144 ++++++- 48 files changed, 2617 insertions(+), 927 deletions(-) create mode 100644 .github/workflows/lingbot-demo-runtime.yml create mode 100644 flashdreams/flashdreams/runtime/demo/bootstrap.py create mode 100644 flashdreams/flashdreams/runtime/keyboard.py create mode 100644 flashdreams/flashdreams/serving/webrtc/demo.py create mode 100644 flashdreams/tests/test_checkpoint_loading.py create mode 100644 integrations/lingbot/lingbot/controls.py create mode 100644 integrations/omnidreams/omnidreams/demo/controls.py diff --git a/.github/workflows/lingbot-demo-runtime.yml b/.github/workflows/lingbot-demo-runtime.yml new file mode 100644 index 000000000..98b575f6e --- /dev/null +++ b/.github/workflows/lingbot-demo-runtime.yml @@ -0,0 +1,334 @@ +name: LingBot Demo Runtime + +on: + push: + branches: + - main + - "pull-request/[0-9]+" + paths: + - ".github/workflows/lingbot-demo-runtime.yml" + - "pyproject.toml" + - "uv.lock" + - "flashdreams/pyproject.toml" + - "flashdreams/flashdreams/core/**" + - "flashdreams/flashdreams/infra/**" + - "flashdreams/flashdreams/runtime/**" + - "flashdreams/flashdreams/serving/**" + - "flashdreams/flashdreams/recipes/taehv/**" + - "flashdreams/flashdreams/recipes/wan/**" + - "flashdreams/tests/test_webrtc_*.py" + - "integrations/lingbot/**" + workflow_dispatch: + +permissions: + contents: read + +jobs: + demo-runtime: + name: null and MP4 + runs-on: linux-amd64-gpu-rtxpro6000-latest-2 + timeout-minutes: 180 + defaults: + run: + shell: bash + container: + image: nvidia/cuda:13.2.1-cudnn-devel-ubuntu24.04 + options: --gpus all + env: + UV_PROJECT_ENVIRONMENT: /tmp/flashdreams-venv + UV_LINK_MODE: copy + UV_PYTHON: "3.12" + MAX_JOBS: 8 + HF_HOME: /tmp/huggingface + FLASHDREAMS_CACHE_DIR: /tmp/flashdreams-cache + # Streaming avoids the old duplicate merged safetensors cache. CI uses + # the generic reserve so the model-specific 200 GiB first-run budget does + # not reject runners that have enough room for the streamed shards. + FLASHDREAMS_MIN_CACHE_FREE_GB: "20" + ARTIFACT_DIR: artifacts/lingbot_demo_runtime + PRESET_ID: lingbot-world-v2-14b-causal-fast-taehv-window15-sink3 + BLOCKS: "5" + FPS: "16" + WIDTH: "640" + HEIGHT: "352" + EXPECTED_WIDTH: "640" + EXPECTED_HEIGHT: "352" + MIN_DURATION_SECONDS: "3" + MAX_DURATION_SECONDS: "5" + steps: + - name: Detect GPU architecture + id: gpu-arch + run: | + nvidia-smi + compute_cap=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | head -1 | tr -d '[:space:]') + arch=$(echo "${compute_cap}" | tr -d '.') + echo "arch=${arch}" >> "$GITHUB_OUTPUT" + echo "Detected GPU compute capability: ${compute_cap} -> sm_${arch}" + + - name: Checkout + uses: actions/checkout@v4 + + - name: Install system dependencies + run: | + apt-get update -qq + DEBIAN_FRONTEND=noninteractive apt-get install -y -qq --no-install-recommends \ + python3 python3-dev python3-venv \ + ffmpeg \ + gcc g++ ninja-build \ + libnccl-dev \ + curl git ca-certificates jq unzip + rm -rf /var/lib/apt/lists/* + + - name: Setup proxy cache + uses: nv-gha-runners/setup-proxy-cache@main + + - name: Setup uv + uses: astral-sh/setup-uv@v6 + with: + enable-cache: true + cache-suffix: "lingbot-demo-runtime-sm${{ steps.gpu-arch.outputs.arch }}" + prune-cache: false + + - name: Install dependencies + env: + NVTE_CUDA_ARCHS: ${{ steps.gpu-arch.outputs.arch }} + run: | + uv venv --clear + uv sync --locked --package flashdreams-lingbot --no-dev + + - name: Verify GPU availability + run: nvidia-smi + + - name: Run LingBot demo modes + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + set -uo pipefail + + log_dir="${ARTIFACT_DIR}/logs" + output_dir="${ARTIFACT_DIR}/outputs" + summary="${ARTIFACT_DIR}/summary.md" + status_file="${ARTIFACT_DIR}/command-status.env" + mkdir -p "${log_dir}" "${output_dir}" + : > "${status_file}" + + ldemo() { + uv run --no-sync --package flashdreams-lingbot lingbot-demo "$@" + } + + run_demo() { + local name="$1" + shift + local log="${log_dir}/${name}.log" + + { + printf '$' + printf ' %q' "$@" + printf '\n\n' + "$@" + } 2>&1 | tee "${log}" + + local rc="${PIPESTATUS[0]}" + echo "${name}=${rc}" >> "${status_file}" + echo "${name} exit code: ${rc}" | tee -a "${summary}" + return 0 + } + + { + echo "# LingBot Demo Runtime CI" + echo + echo "| Mode | Expected blocks | Output |" + echo "| --- | ---: | --- |" + echo "| null | ${BLOCKS} | none |" + echo "| MP4 | ${BLOCKS} | lingbot-demo-replay.mp4 |" + echo + echo "## Command Status" + } > "${summary}" + + run_demo null \ + ldemo replay \ + --device cuda:0 \ + --preset-id "${PRESET_ID}" \ + --example-idx 0 \ + --total-blocks "${BLOCKS}" \ + --fps "${FPS}" \ + --pixel-height "${HEIGHT}" \ + --pixel-width "${WIDTH}" \ + --output-mode null + + run_demo mp4 \ + ldemo replay \ + --device cuda:0 \ + --preset-id "${PRESET_ID}" \ + --example-idx 0 \ + --total-blocks "${BLOCKS}" \ + --fps "${FPS}" \ + --pixel-height "${HEIGHT}" \ + --pixel-width "${WIDTH}" \ + --output "${output_dir}/lingbot-demo-replay.mp4" + + - name: Validate LingBot demo artifacts + run: | + set -euo pipefail + + log_dir="${ARTIFACT_DIR}/logs" + output_dir="${ARTIFACT_DIR}/outputs" + probe_dir="${ARTIFACT_DIR}/ffprobe" + summary="${ARTIFACT_DIR}/summary.md" + status_file="${ARTIFACT_DIR}/command-status.env" + mkdir -p "${probe_dir}" + + status_of() { + awk -F= -v name="$1" '$1 == name { print $2 }' "${status_file}" + } + + assert_exit_zero() { + local name="$1" + local rc + rc="$(status_of "${name}")" + if [ "${rc}" != "0" ]; then + echo "${name} command failed with exit code ${rc}" >&2 + exit 1 + fi + } + + assert_clean_log() { + local log="$1" + if grep -En "ERROR|Traceback|Exception|status=failed|Run failed|failed run" "${log}"; then + echo "failure marker found in ${log}" >&2 + exit 1 + fi + } + + assert_log_contains() { + local log="$1" + local pattern="$2" + local label="$3" + if ! grep -Eq "${pattern}" "${log}"; then + echo "expected ${label} in ${log}" >&2 + exit 1 + fi + } + + assert_log_not_contains() { + local log="$1" + local pattern="$2" + local label="$3" + if grep -Eq "${pattern}" "${log}"; then + echo "unexpected ${label} in ${log}" >&2 + exit 1 + fi + } + + validate_mp4() { + local mode="$1" + local mp4="$2" + local metadata="${probe_dir}/${mode}.json" + + if [ ! -s "${mp4}" ]; then + echo "expected non-empty MP4 at ${mp4}" >&2 + exit 1 + fi + + ffprobe \ + -v error \ + -select_streams v:0 \ + -show_entries stream=width,height,r_frame_rate,avg_frame_rate,nb_frames,duration:format=duration \ + -of json \ + "${mp4}" > "${metadata}" + + local stream_count width height duration + stream_count="$(jq '.streams | length' "${metadata}")" + width="$(jq -r '.streams[0].width // ""' "${metadata}")" + height="$(jq -r '.streams[0].height // ""' "${metadata}")" + duration="$(jq -r '.streams[0].duration // .format.duration // "0"' "${metadata}")" + + if [ "${stream_count}" -lt 1 ]; then + echo "ffprobe found no video stream in ${mp4}" >&2 + exit 1 + fi + + if [ "${width}" != "${EXPECTED_WIDTH}" ] || [ "${height}" != "${EXPECTED_HEIGHT}" ]; then + echo "unexpected ${mode} resolution ${width}x${height}; expected ${EXPECTED_WIDTH}x${EXPECTED_HEIGHT}" >&2 + exit 1 + fi + + awk \ + -v duration="${duration}" \ + -v min_duration="${MIN_DURATION_SECONDS}" \ + -v max_duration="${MAX_DURATION_SECONDS}" \ + 'BEGIN { + if ((duration + 0) < min_duration || (duration + 0) > max_duration) { + exit 1 + } + }' || { + echo "unexpected ${mode} duration ${duration}s; expected ${MIN_DURATION_SECONDS}-${MAX_DURATION_SECONDS}s" >&2 + exit 1 + } + } + + null_log="${log_dir}/null.log" + mp4_log="${log_dir}/mp4.log" + + assert_exit_zero null + assert_exit_zero mp4 + + assert_clean_log "${null_log}" + assert_clean_log "${mp4_log}" + + assert_log_contains "${null_log}" "Streaming sharded safetensors checkpoint" "null streaming checkpoint load" + assert_log_contains "${null_log}" "Finished streaming .* safetensors shard" "null streamed checkpoint completion" + assert_log_contains "${null_log}" "AR 4 encode" "null final AR block" + assert_log_contains "${null_log}" "Lingbot runtime step 4 frames=" "null final replay step" + assert_log_not_contains "${null_log}" "Loading merged sharded checkpoint from cache|Saved merged sharded checkpoint" "merged safetensors cache usage" + + assert_log_contains "${mp4_log}" "Streaming sharded safetensors checkpoint" "MP4 streaming checkpoint load" + assert_log_contains "${mp4_log}" "Finished streaming .* safetensors shard" "MP4 streamed checkpoint completion" + assert_log_contains "${mp4_log}" "AR 4 encode" "MP4 final AR block" + assert_log_contains "${mp4_log}" "Lingbot runtime step 4 frames=" "MP4 final replay step" + assert_log_not_contains "${mp4_log}" "Loading merged sharded checkpoint from cache|Saved merged sharded checkpoint" "merged safetensors cache usage" + + validate_mp4 mp4 "${output_dir}/lingbot-demo-replay.mp4" + + { + echo + echo "## Validation" + echo + echo "- Null and MP4 commands exited zero." + echo "- Logs contained expected streaming-checkpoint and final AR-step markers." + echo "- Logs did not contain merged-safetensors cache markers." + echo "- MP4 output was non-empty and passed ffprobe stream checks." + } >> "${summary}" + + - name: Trim uv cache for upload + if: always() + run: | + cache_dir="${UV_CACHE_DIR:-/github/home/.cache/uv}" + echo "=== Cache size before trim ===" + du -sh "${cache_dir}" || true + du -sh "${cache_dir}"/*/ 2>/dev/null || true + + rm -rf "${cache_dir}/wheels-v6" + rm -rf "${cache_dir}/archive-v0" + + find "${cache_dir}/git-v0/checkouts" \ + \( -name "build" -o -name "*.egg-info" -o -name "__pycache__" \) \ + -type d -exec rm -rf {} + 2>/dev/null || true + + rm -rf "${cache_dir}/sdists-v9/editable" + + echo "" + echo "=== Cache size after trim ===" + du -sh "${cache_dir}" || true + du -sh "${cache_dir}"/*/ 2>/dev/null || true + echo "" + echo "=== Cached built wheels (sdists-v9) ===" + find "${cache_dir}/sdists-v9" -name "*.whl" -exec ls -lh {} \; 2>/dev/null || true + + - name: Upload demo runtime artifacts + if: always() + uses: actions/upload-artifact@v4 + with: + name: lingbot-demo-runtime + path: ${{ env.ARTIFACT_DIR }} + if-no-files-found: ignore diff --git a/flashdreams/flashdreams/core/checkpoint/load.py b/flashdreams/flashdreams/core/checkpoint/load.py index 4d11f7642..4432e1be5 100644 --- a/flashdreams/flashdreams/core/checkpoint/load.py +++ b/flashdreams/flashdreams/core/checkpoint/load.py @@ -20,14 +20,16 @@ import io import json import os +import time from collections.abc import Callable, Mapping -from concurrent.futures import ProcessPoolExecutor +from concurrent.futures import ThreadPoolExecutor from typing import Literal, overload from urllib.parse import unquote, urlparse import torch from huggingface_hub import hf_hub_download, try_to_load_from_cache from loguru import logger +from safetensors import safe_open from safetensors.torch import load as load_safetensors from safetensors.torch import load_file as load_safetensors_file from safetensors.torch import save_file as save_safetensors @@ -239,7 +241,7 @@ def _safetensors_device(map_location: str | torch.device) -> str: def _hf_hub_download_shard_task( args: tuple[str, str, str | None, str], ) -> tuple[str, str]: - """Picklable worker: download one shard; used by ProcessPoolExecutor.""" + """Download or resolve one Hugging Face shard.""" repo_id, shard_file, subfolder, revision = args settings: dict[str, object] = { "repo": repo_id, @@ -275,7 +277,7 @@ def _parallel_hf_hub_download_shards( subfolder: str | None, revision: str, ) -> dict[str, str]: - """Download unique shard files in parallel processes; returns shard -> local path.""" + """Download unique shard files in parallel workers; returns shard -> local path.""" if not shard_files: return {} if len(shard_files) == 1: @@ -297,10 +299,10 @@ def _parallel_hf_hub_download_shards( work = [(repo_id, s, subfolder, revision) for s in shard_files] logger.info( f"Downloading {len(shard_files)} Hugging Face safetensors shards " - f"with up to {max_workers} parallel processes" + f"with up to {max_workers} parallel workers" ) shard_to_path: dict[str, str] = {} - with ProcessPoolExecutor(max_workers=max_workers) as pool: + with ThreadPoolExecutor(max_workers=max_workers) as pool: for shard_file, path in pool.map(_hf_hub_download_shard_task, work): shard_to_path[shard_file] = path return shard_to_path @@ -737,12 +739,334 @@ def _load_checkpoint_from_local( ) -> dict[str, torch.Tensor]: """Load checkpoint from local filesystem.""" if ext == ".safetensors": - with open(path, "rb") as f: - return load_safetensors(f.read()) + return load_safetensors_file(path, device=_safetensors_device(map_location)) else: return torch.load(path, map_location=map_location, weights_only=False) +def _copy_checkpoint_tensor(destination: torch.Tensor, source: torch.Tensor) -> int: + """Copy one checkpoint tensor into ``destination`` with bounded staging.""" + checkpoint_bytes = source.numel() * source.element_size() + if destination.device.type != "cpu": + staged = source.to(dtype=destination.dtype) + if staged.data_ptr() == source.data_ptr(): + staged = staged.clone() + destination.copy_(staged) + if destination.device.type == "cuda": + # Keep the CPU staging buffer alive until CUDA has consumed it. + torch.cuda.synchronize(destination.device) + del staged + return checkpoint_bytes + + destination.copy_(source.to(device=destination.device, dtype=destination.dtype)) + return checkpoint_bytes + + +def _stream_safetensors_into_model( + model: torch.nn.Module, + path: str, +) -> torch.nn.Module: + """Copy a safetensors checkpoint into a model with bounded host residency.""" + model_state = model.state_dict() + + with safe_open(path, framework="pt", device="cpu") as source: + checkpoint_keys = set(source.keys()) + model_keys = set(model_state) + missing = sorted(model_keys - checkpoint_keys) + unexpected = sorted(checkpoint_keys - model_keys) + if missing or unexpected: + details = [] + if missing: + details.append(f"Missing key(s): {', '.join(missing[:20])}") + if unexpected: + details.append(f"Unexpected key(s): {', '.join(unexpected[:20])}") + raise RuntimeError( + f"Checkpoint does not match {type(model).__name__}: " + + "; ".join(details) + ) + + for name, destination in model_state.items(): + source_shape = tuple(source.get_slice(name).get_shape()) + if source_shape != tuple(destination.shape): + raise RuntimeError( + f"Checkpoint tensor {name!r} has shape {source_shape}, " + f"expected {tuple(destination.shape)}" + ) + + with torch.no_grad(): + for name, destination in model_state.items(): + tensor = source.get_tensor(name) + try: + _copy_checkpoint_tensor(destination, tensor) + finally: + del tensor + + return model + + +def _stream_sharded_safetensors_into_model( + model: torch.nn.Module, + *, + weight_map: Mapping[str, str], + resolve_shard_path: Callable[[str], str], +) -> torch.nn.Module: + """Copy a sharded safetensors checkpoint into a model one shard at a time.""" + model_state = model.state_dict() + checkpoint_keys = set(weight_map) + model_keys = set(model_state) + missing = sorted(model_keys - checkpoint_keys) + unexpected = sorted(checkpoint_keys - model_keys) + if missing or unexpected: + details = [] + if missing: + details.append(f"Missing key(s): {', '.join(missing[:20])}") + if unexpected: + details.append(f"Unexpected key(s): {', '.join(unexpected[:20])}") + raise RuntimeError( + f"Checkpoint does not match {type(model).__name__}: " + "; ".join(details) + ) + + keys_by_shard: dict[str, list[str]] = {} + for tensor_name, shard_file in weight_map.items(): + keys_by_shard.setdefault(shard_file, []).append(tensor_name) + + shard_files = sorted(keys_by_shard) + destination_devices = sorted( + {str(tensor.device) for tensor in model_state.values()} + ) + destination_dtypes = sorted({str(tensor.dtype) for tensor in model_state.values()}) + logger.info( + "Streaming sharded safetensors into {}: {} shard(s), {} tensor(s), " + "destination devices={}, dtypes={}", + type(model).__name__, + len(shard_files), + len(weight_map), + destination_devices, + destination_dtypes, + ) + + for shard_index, shard_file in enumerate(shard_files, start=1): + shard_path = resolve_shard_path(shard_file) + tensor_names = keys_by_shard[shard_file] + shard_size_gib = os.path.getsize(shard_path) / 1024**3 + started = time.perf_counter() + logger.info( + "Validating safetensors shard {}/{}: {} tensors, {:.2f} GiB, {}", + shard_index, + len(shard_files), + len(tensor_names), + shard_size_gib, + shard_file, + ) + with safe_open(shard_path, framework="pt", device="cpu") as source: + shard_keys = set(source.keys()) + for name in tensor_names: + if name not in shard_keys: + raise KeyError( + f"Key {name!r} missing from shard {shard_file!r} " + f"(path {shard_path!r})" + ) + source_shape = tuple(source.get_slice(name).get_shape()) + destination = model_state[name] + if source_shape != tuple(destination.shape): + raise RuntimeError( + f"Checkpoint tensor {name!r} has shape {source_shape}, " + f"expected {tuple(destination.shape)}" + ) + logger.info( + "Validated safetensors shard {}/{} in {:.1f}s: {}", + shard_index, + len(shard_files), + time.perf_counter() - started, + shard_file, + ) + + total_copied_bytes = 0 + total_started = time.perf_counter() + with torch.no_grad(): + for shard_index, shard_file in enumerate(shard_files, start=1): + shard_path = resolve_shard_path(shard_file) + tensor_names = keys_by_shard[shard_file] + shard_copied_bytes = 0 + started = time.perf_counter() + logger.info( + "Streaming safetensors shard {}/{} into model: {} tensors, {}", + shard_index, + len(shard_files), + len(tensor_names), + shard_file, + ) + with safe_open(shard_path, framework="pt", device="cpu") as source: + for name in tensor_names: + tensor = source.get_tensor(name) + try: + tensor_bytes = _copy_checkpoint_tensor( + model_state[name], tensor + ) + shard_copied_bytes += tensor_bytes + total_copied_bytes += tensor_bytes + finally: + del tensor + elapsed = time.perf_counter() - started + throughput = ( + shard_copied_bytes / 1024**3 / elapsed if elapsed > 0 else float("inf") + ) + logger.info( + "Streamed safetensors shard {}/{} in {:.1f}s: {:.2f} GiB copied " + "({:.2f} GiB/s), {}", + shard_index, + len(shard_files), + elapsed, + shard_copied_bytes / 1024**3, + throughput, + shard_file, + ) + + elapsed = time.perf_counter() - total_started + throughput = total_copied_bytes / 1024**3 / elapsed if elapsed > 0 else float("inf") + logger.info( + "Finished streaming {} safetensors shard(s) in {:.1f}s: {:.2f} GiB copied " + "({:.2f} GiB/s)", + len(shard_files), + elapsed, + total_copied_bytes / 1024**3, + throughput, + ) + + return model + + +def _stream_sharded_safetensors_index_into_model( + checkpoint_path: str, + *, + model: torch.nn.Module, + checkpoint_min_free_gb: float | None, +) -> torch.nn.Module | None: + """Stream a safetensors index checkpoint into ``model`` without merging.""" + if checkpoint_path.startswith("s3://"): + return None + + if _is_huggingface_checkpoint_url(checkpoint_path): + repo_id, index_filename, subfolder, revision = ( + _parse_huggingface_checkpoint_url(checkpoint_path) + ) + logger.info( + f"Streaming sharded safetensors checkpoint from Hugging Face: " + f"{checkpoint_path}" + ) + settings: dict[str, object] = { + "repo": repo_id, + "filename": index_filename, + "revision": revision, + } + _preflight_checkpoint_cache_requirement( + label="Hugging Face sharded checkpoint cache", + min_free_gb=checkpoint_min_free_gb, + settings=settings, + ) + min_bytes = _preflight_hf_cache( + label="Hugging Face checkpoint index cache", + settings=settings, + ) + try: + index_local = hf_hub_download( + repo_id=repo_id, + filename=index_filename, + subfolder=subfolder, + revision=revision, + ) + except Exception as exc: + _raise_hf_cache_disk_error( + exc, + label="Hugging Face checkpoint index cache", + required_bytes=min_bytes, + settings=settings, + ) + raise + with open(index_local) as f: + index = json.load(f) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError( + f"Invalid or empty weight_map in safetensors index: {index_local}" + ) + + unique_shards = sorted(set(weight_map.values())) + shard_to_path = _parallel_hf_hub_download_shards( + repo_id=repo_id, + shard_files=unique_shards, + subfolder=subfolder, + revision=revision, + ) + + def resolve_shard_path(shard_file: str) -> str: + return shard_to_path[shard_file] + + return _stream_sharded_safetensors_into_model( + model, + weight_map=weight_map, + resolve_shard_path=resolve_shard_path, + ) + + if not os.path.isfile(checkpoint_path): + raise FileNotFoundError( + f"Sharded safetensors index not found: {checkpoint_path}" + ) + logger.info( + f"Streaming sharded safetensors checkpoint from local index: {checkpoint_path}" + ) + with open(checkpoint_path) as f: + index = json.load(f) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError( + f"Invalid or empty weight_map in safetensors index: {checkpoint_path}" + ) + base_dir = os.path.dirname(os.path.abspath(checkpoint_path)) + + def resolve_shard_path(shard_file: str) -> str: + return os.path.join(base_dir, shard_file) + + return _stream_sharded_safetensors_into_model( + model, + weight_map=weight_map, + resolve_shard_path=resolve_shard_path, + ) + + +def _resolve_streamable_safetensors_path( + checkpoint_path: str, + *, + local_cache_dir: str, + checkpoint_min_free_gb: float | None, +) -> str | None: + """Resolve a locally available safetensors file for streaming model loads.""" + if _is_sharded_safetensors_index_checkpoint(checkpoint_path): + if checkpoint_path.startswith("s3://"): + return None + cache_path = _sharded_safetensors_merge_cache_path( + checkpoint_path, local_cache_dir + ) + if os.path.exists(cache_path): + logger.info(f"Streaming merged sharded checkpoint from cache: {cache_path}") + return cache_path + return None + + if _get_checkpoint_extension(checkpoint_path) != ".safetensors": + return None + if _is_huggingface_checkpoint_url(checkpoint_path): + return _download_checkpoint_from_huggingface_url( + checkpoint_path, + checkpoint_min_free_gb=checkpoint_min_free_gb, + ) + if checkpoint_path.startswith("s3://"): + cache_path = os.path.join( + local_cache_dir, checkpoint_path.removeprefix("s3://") + ) + return cache_path if os.path.exists(cache_path) else None + return checkpoint_path + + def _load_checkpoint_from_s3( s3_path: str, ext: str, @@ -837,8 +1161,9 @@ def load_checkpoint( Args: checkpoint_path: ``s3://`` URI, local path, or HF URL. Single-file or DCP directory. - model: Model to load weights into. Required for DCP. Optional for - single-file: when provided, ``load_state_dict`` is called. + model: Model to load weights into. Required for DCP. Cached + safetensors are streamed into a provided model; other single-file + formats use ``load_state_dict``. checkpoint_type: ``"auto"``, ``"single"``, or ``"distributed"``. local_cache_dir: Directory for caches. credential_path: S3 credentials path. @@ -873,6 +1198,25 @@ def load_checkpoint( checkpoint_type = "distributed" if checkpoint_type == "single": + if model is not None: + if _is_sharded_safetensors_index_checkpoint(checkpoint_path): + streamed_model = _stream_sharded_safetensors_index_into_model( + checkpoint_path, + model=model, + checkpoint_min_free_gb=checkpoint_min_free_gb, + ) + if streamed_model is not None: + logger.info(f"Streamed checkpoint into model: {checkpoint_path}") + return streamed_model + stream_path = _resolve_streamable_safetensors_path( + checkpoint_path, + local_cache_dir=local_cache_dir, + checkpoint_min_free_gb=checkpoint_min_free_gb, + ) + if stream_path is not None: + _stream_safetensors_into_model(model, stream_path) + logger.info(f"Streamed checkpoint into model: {checkpoint_path}") + return model state_dict = load_single_checkpoint( checkpoint_path=checkpoint_path, local_cache_dir=local_cache_dir, diff --git a/flashdreams/flashdreams/recipes/wan/transformer/wan21.py b/flashdreams/flashdreams/recipes/wan/transformer/wan21.py index c1ae0ab04..0022f1fc2 100644 --- a/flashdreams/flashdreams/recipes/wan/transformer/wan21.py +++ b/flashdreams/flashdreams/recipes/wan/transformer/wan21.py @@ -157,6 +157,16 @@ class Wan21TransformerConfig(TransformerConfig): """Pre-load state-dict remap (e.g. Self-Forcing's ``generator_ema.model.…`` layout).""" + stream_checkpoint: bool = False + """Load cached safetensors directly into the model with bounded host residency.""" + + init_device: str | None = None + """Optional device used for initial network parameter allocation. + + Large streaming-checkpoint models can set this to the final runtime device + so the module is not first materialized as fp32 CPU tensors. + """ + batch_shape: tuple[int, ...] = (1,) """Batch dims of the latent (excluding the L, D dims).""" @@ -273,19 +283,29 @@ def __init__(self, config: Wan21TransformerConfig) -> None: self._output_height: int | None = None self._output_width: int | None = None - self.network = config.network.setup() - self.network = self.network.to(dtype=config.dtype) + self.network = self._setup_network(config) self.network.eval() self.network.set_context_parallel_group(cp_group=self._cp_group) if config.checkpoint_path is not None: - state_dict = load_checkpoint( - config.checkpoint_path, - checkpoint_min_free_gb=config.checkpoint_min_free_gb, - ) - if config.state_dict_transform is not None: - state_dict = config.state_dict_transform(state_dict) - self.network.load_state_dict(state_dict) + if config.stream_checkpoint: + if config.state_dict_transform is not None: + raise ValueError( + "stream_checkpoint does not support state_dict_transform" + ) + load_checkpoint( + config.checkpoint_path, + model=self.network, + checkpoint_min_free_gb=config.checkpoint_min_free_gb, + ) + else: + state_dict = load_checkpoint( + config.checkpoint_path, + checkpoint_min_free_gb=config.checkpoint_min_free_gb, + ) + if config.state_dict_transform is not None: + state_dict = config.state_dict_transform(state_dict) + self.network.load_state_dict(state_dict) self.network.update_parameters_after_loading_checkpoint() if config.compile_network: @@ -308,6 +328,23 @@ def __init__(self, config: Wan21TransformerConfig) -> None: self._cuda_graph_dispatch.uncond_call or self.network ) + @staticmethod + def _setup_network(config: Wan21TransformerConfig) -> WanDiTNetwork: + init_device = ( + None if config.init_device is None else torch.device(config.init_device) + ) + if init_device is None: + return config.network.setup().to(dtype=config.dtype) + + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(config.dtype) + with torch.device(init_device): + network = config.network.setup() + finally: + torch.set_default_dtype(previous_dtype) + return network.to(device=init_device, dtype=config.dtype) + @property def latent_shape(self) -> tuple[int, ...]: """Per-rank post-patchify latent shape ``[*batch_shape, L/cp, D]``. diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py index 6e205823c..3280e6862 100644 --- a/flashdreams/flashdreams/runtime/__init__.py +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -38,6 +38,18 @@ InferenceSession, ModelAdapter, ) +from flashdreams.runtime.keyboard import ( + DEFAULT_SUPPORTED_KEYS, + DRIVING_SUPPORTED_KEYS, + KEY_ALIASES, + WSAD_SUPPORTED_KEYS, + ImageRequest, + KeyboardState, + PromptRequest, + ResetRequest, + SparseInputSnapshot, + normalize_key, +) from flashdreams.runtime.mapping import ( DeclaresMappingSchema, IdentityInputMapping, @@ -76,8 +88,10 @@ "combine_mapping_schemas", "DeclaresMappingSchema", "DEFAULT_DRIVING_BINDINGS", + "DEFAULT_SUPPORTED_KEYS", "DeviceConverter", "DeviceConverterSchema", + "DRIVING_SUPPORTED_KEYS", "DRIVER_COMMAND", "ExecutionBackend", "IdentityInputMapping", @@ -93,6 +107,9 @@ "InputMapping", "InputMappingSchema", "InputPhase", + "ImageRequest", + "KEY_ALIASES", + "KeyboardState", "KeyboardToDriverCommand", "MappingCompatibility", "MetricsRecorder", @@ -105,8 +122,11 @@ "OutputArtifact", "OutputTarget", "Precision", + "PromptRequest", + "ResetRequest", "RuntimeMetricSample", "ScriptedModality", + "SparseInputSnapshot", "StepRequest", "StepRequirements", "StepResult", @@ -119,5 +139,7 @@ "UserInputEvent", "UserInputs", "UserInputSchema", + "WSAD_SUPPORTED_KEYS", + "normalize_key", "validate_phase", ] diff --git a/flashdreams/flashdreams/runtime/canonical.py b/flashdreams/flashdreams/runtime/canonical.py index 55f333ce7..bf3960380 100644 --- a/flashdreams/flashdreams/runtime/canonical.py +++ b/flashdreams/flashdreams/runtime/canonical.py @@ -37,7 +37,7 @@ UserInputs, UserInputSchema, ) -from flashdreams.serving.realtime.input import KeyboardState, normalize_key +from flashdreams.runtime.keyboard import KeyboardState, normalize_key DriverBindings = Mapping[str, frozenset[str]] diff --git a/flashdreams/flashdreams/runtime/demo/__init__.py b/flashdreams/flashdreams/runtime/demo/__init__.py index 60187237c..8d61e31bd 100644 --- a/flashdreams/flashdreams/runtime/demo/__init__.py +++ b/flashdreams/flashdreams/runtime/demo/__init__.py @@ -75,7 +75,6 @@ WebRTCOutputSpec, ) from flashdreams.runtime.demo.timing import ( - SPARSE_KEY_SEGMENTS_METADATA_KEY, ActivationPolicy, ActivationResult, ActivationSignal, @@ -83,8 +82,9 @@ CatchUpDecision, CatchUpPolicy, DeterministicClock, - KeyboardRealtimeInputSource, RealtimeClock, + RealtimeEventInputSource, + RealtimeEventResampler, RealtimeWindowResult, ResamplerRealtimeClock, SignalActivationPolicy, @@ -121,7 +121,6 @@ "ModelWarmupAdapter", "ModelWarmupPlan", "ModelInputProvider", - "KeyboardRealtimeInputSource", "Mp4ErrorPolicy", "Mp4OutputSink", "Mp4OutputSpec", @@ -139,6 +138,8 @@ "ProviderCapabilities", "RealtimeInputSource", "RealtimeClock", + "RealtimeEventInputSource", + "RealtimeEventResampler", "RealtimeSessionDriver", "RealtimeWindowResult", "ResolvedRunCapabilities", @@ -155,7 +156,6 @@ "SessionInfo", "SignalActivationPolicy", "SingleSessionAdmissionPolicy", - "SPARSE_KEY_SEGMENTS_METADATA_KEY", "StepOutcome", "StepPipeline", "UserInputWindow", diff --git a/flashdreams/flashdreams/runtime/demo/app.py b/flashdreams/flashdreams/runtime/demo/app.py index d44dd81bb..b763e4eaf 100644 --- a/flashdreams/flashdreams/runtime/demo/app.py +++ b/flashdreams/flashdreams/runtime/demo/app.py @@ -14,12 +14,12 @@ 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 ( +from flashdreams.runtime.demo.bootstrap import ( configure_logging, initialize_cuda_distributed, ) +from flashdreams.runtime.demo.replay import run_replay_demo +from flashdreams.runtime.demo.spec import DemoAdapter, DemoSpec class DemoApplication(ABC): diff --git a/flashdreams/flashdreams/runtime/demo/bootstrap.py b/flashdreams/flashdreams/runtime/demo/bootstrap.py new file mode 100644 index 000000000..6b1e68f5d --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/bootstrap.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared process bootstrap for demo applications.""" + +from __future__ import annotations + +import logging +import os +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist + + +@dataclass(frozen=True, slots=True) +class DistributedDemoContext: + """CUDA/distributed launch context for a demo process.""" + + device: torch.device + world_rank: int + world_size: int + + +def configure_logging(*, world_rank: int | None = None) -> None: + from flashdreams.core.distributed import configure_loguru_for_distributed + + configure_loguru_for_distributed(world_rank=world_rank) + for logger_name in ("aioice", "aioice.ice", "aiortc"): + logging.getLogger(logger_name).setLevel(logging.WARNING) + + +def _distributed_init() -> None: + from flashdreams.core.distributed import init as distributed_init + + distributed_init() + + +def initialize_cuda_distributed( + *, + default_device: str | torch.device = "cuda:0", + distributed_init_fn: Callable[[], object] | None = None, + configure_logging_fn: Callable[..., None] = configure_logging, + torch_module: Any = torch, + dist_module: Any = dist, +) -> DistributedDemoContext: + """Initialize CUDA and optional torch.distributed for demo serving.""" + if not torch_module.cuda.is_available(): + raise RuntimeError("CUDA is required for inference in the demo server.") + + has_rank = "RANK" in os.environ + has_world_size = "WORLD_SIZE" in os.environ + if has_rank != has_world_size: + raise RuntimeError( + "Distributed launch expects both RANK and WORLD_SIZE to be set." + ) + + distributed_launch = has_rank and has_world_size + if distributed_launch: + if distributed_init_fn is None: + distributed_init_fn = _distributed_init + distributed_init_fn() + world_rank = dist_module.get_rank() + world_size = dist_module.get_world_size() + else: + world_rank = 0 + world_size = 1 + + device_count = torch_module.cuda.device_count() + if device_count < 1: + raise RuntimeError("CUDA device count must be >= 1 for inference.") + if distributed_launch: + local_rank = world_rank % device_count + torch_device = torch_module.device(f"cuda:{local_rank}") + else: + torch_device = torch_module.device(default_device) + if torch_device.type != "cuda": + raise RuntimeError( + f"CUDA device is required for inference, got {torch_device}." + ) + if torch_device.index is None: + torch_device = torch_module.device("cuda:0") + torch_module.cuda.set_device(torch_device) + configure_logging_fn(world_rank=world_rank) + return DistributedDemoContext( + device=torch_device, + world_rank=world_rank, + world_size=world_size, + ) + + +__all__ = [ + "DistributedDemoContext", + "configure_logging", + "initialize_cuda_distributed", +] diff --git a/flashdreams/flashdreams/runtime/demo/timing.py b/flashdreams/flashdreams/runtime/demo/timing.py index d48e1e213..e6609cba9 100644 --- a/flashdreams/flashdreams/runtime/demo/timing.py +++ b/flashdreams/flashdreams/runtime/demo/timing.py @@ -19,8 +19,6 @@ CatchUpPolicy = Literal["drop", "fold", "compress"] -SPARSE_KEY_SEGMENTS_METADATA_KEY = "sparse_key_segments" - @dataclass(frozen=True, kw_only=True, slots=True) class CatchUpDecision: @@ -181,28 +179,60 @@ async def wait_until_active( await asyncio.gather(*tasks, return_exceptions=True) -SparseKeySegment = tuple[float, float, frozenset[str]] - - class _RealtimeTimeline(Protocol): dt: float next_chunk_start_v: float - -class _SparseInputResampler(_RealtimeTimeline, Protocol): - def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: ... - def reset(self, *, start_v: float) -> None: ... def sample_chunk( self, num_frames: int, - ) -> tuple[Sequence[SparseKeySegment], Sequence[float]]: ... + ) -> Sequence[float]: ... + + +@dataclass(slots=True) +class RealtimeEventResampler: + """Transport-neutral realtime window timeline. + + The shared driver only owns virtual time and frame sample locations. Raw + browser/native events are stored as :class:`UserInputs`; model providers + decide how to interpret them. + """ + + fps: float + start_v: float = 0.0 + next_chunk_start_v: float = field(init=False) + _dt: float = field(init=False, repr=False) + + def __post_init__(self) -> None: + if self.fps <= 0: + raise ValueError("fps must be > 0") + self._dt = 1.0 / float(self.fps) + self.next_chunk_start_v = float(self.start_v) + + @property + def dt(self) -> float: + return self._dt + + def reset(self, *, start_v: float) -> None: + self.next_chunk_start_v = float(start_v) + + def sample_chunk(self, num_frames: int) -> tuple[float, ...]: + if num_frames < 1: + raise ValueError("num_frames must be >= 1") + chunk_start_v = self.next_chunk_start_v + chunk_end_v = chunk_start_v + num_frames * self._dt + frame_times = tuple( + chunk_start_v + (index + 1) * self._dt for index in range(num_frames) + ) + self.next_chunk_start_v = chunk_end_v + return frame_times @dataclass(slots=True) class ResamplerRealtimeClock: - """Realtime clock that reuses ``KeyboardResampler``'s virtual timeline.""" + """Realtime clock that reuses a resampler's virtual timeline.""" resampler: _RealtimeTimeline now_fn: Callable[[], float] = time.monotonic @@ -245,7 +275,7 @@ def catch_up( ) -> CatchUpDecision: if policy != "fold": raise NotImplementedError( - f"Catch-up policy {policy!r} has no existing resampler analog yet." + f"Catch-up policy {policy!r} has no existing timeline analog yet." ) if not math.isfinite(max_lag_s) or max_lag_s < 0.0: raise ValueError("max_lag_s must be finite and >= 0.") @@ -278,10 +308,10 @@ def catch_up( @dataclass(slots=True) -class KeyboardRealtimeInputSource: - """Realtime input source backed by the existing keyboard resampler.""" +class RealtimeEventInputSource: + """Realtime input source backed by raw event windows.""" - resampler: _SparseInputResampler + resampler: _RealtimeTimeline max_lag_s: float | None = None catch_up_policy: CatchUpPolicy = "fold" is_finite: bool = False @@ -293,20 +323,17 @@ def __post_init__(self) -> None: not math.isfinite(self.max_lag_s) or self.max_lag_s < 0.0 ): raise ValueError( - "KeyboardRealtimeInputSource.max_lag_s must be finite and >= 0." + "RealtimeEventInputSource.max_lag_s must be finite and >= 0." ) if self.catch_up_policy != "fold": raise NotImplementedError( f"Catch-up policy {self.catch_up_policy!r} has no existing " - "KeyboardResampler analog yet." + "event-window analog yet." ) def is_finished(self) -> bool: return False - def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: - self.resampler.on_edge(arrival_t=arrival_t, event=event, key=key) - def reset(self, *, start_v: float) -> None: self.resampler.reset(start_v=start_v) @@ -328,14 +355,13 @@ async def next_realtime_window( policy=self.catch_up_policy, ) start_s = self.resampler.next_chunk_start_v - segments, frame_times = self.resampler.sample_chunk(input_frame_count) + frame_times = self.resampler.sample_chunk(input_frame_count) end_s = self.resampler.next_chunk_start_v window = UserInputWindow( start_s=start_s, end_s=end_s, frame_times=tuple(frame_times), inputs=UserInputs(), - metadata={SPARSE_KEY_SEGMENTS_METADATA_KEY: tuple(segments)}, ) return RealtimeWindowResult(window=window, catch_up=catch_up) @@ -373,11 +399,11 @@ def _anchor_if_realtime( "CatchUpDecision", "CatchUpPolicy", "DeterministicClock", - "KeyboardRealtimeInputSource", + "RealtimeEventInputSource", + "RealtimeEventResampler", "RealtimeClock", "RealtimeWindowResult", "ResamplerRealtimeClock", - "SPARSE_KEY_SEGMENTS_METADATA_KEY", "SignalActivationPolicy", "input_frame_count_from_request", ] diff --git a/flashdreams/flashdreams/runtime/demo/webrtc.py b/flashdreams/flashdreams/runtime/demo/webrtc.py index b9444a197..2e835a6e9 100644 --- a/flashdreams/flashdreams/runtime/demo/webrtc.py +++ b/flashdreams/flashdreams/runtime/demo/webrtc.py @@ -1,114 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Shared WebRTC demo construction.""" +"""Deprecated location for WebRTC demo construction helpers. -from __future__ import annotations - -from collections.abc import Callable -from importlib.resources import files -from pathlib import Path -from typing import Any - -from aiohttp import web - -from flashdreams.serving.webrtc.bootstrap import run_webrtc_server -from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager -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 serve_webrtc_demo( - *, - output: WebRTCOutputSpec, - model_id: str, - session_manager: BaseWebRTCSessionManager[Any, Any], - app_resources: WebRTCAppResources, - world_rank: int = 0, - create_app_fn: CreateWebRTCApp = create_webrtc_app, - server_runner: RunWebRTCServer = run_webrtc_server, -) -> web.Application | None: - """Serve a prepared model WebRTC runtime through the shared transport.""" - app = ( - _create_app( - output=output, - model_id=model_id, - app_resources=app_resources, - session_manager=session_manager, - create_app_fn=create_app_fn, - ) - if world_rank == 0 - else None - ) - server_runner( - world_rank=world_rank, - session_manager=session_manager, - app=app, - host=output.host, - port=output.port, - ) - return app - - -def _create_app( - *, - output: WebRTCOutputSpec, - model_id: str, - app_resources: WebRTCAppResources, - session_manager: BaseWebRTCSessionManager[Any, Any], - create_app_fn: CreateWebRTCApp, -) -> web.Application: - if output.web_dir is not None: - return _build_webrtc_app( - output=output, - session_manager=session_manager, - create_app_fn=create_app_fn, - preload_name=output.preload_name or app_resources.preload_name or model_id, - ) - 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, - cleanup_callback=close_package_resources, - ) - - -def _build_webrtc_app( - *, - output: WebRTCOutputSpec, - session_manager: BaseWebRTCSessionManager[Any, Any], - create_app_fn: CreateWebRTCApp, - preload_name: str, -) -> web.Application: - if output.web_dir is None: - raise ValueError("WebRTC app creation requires output.web_dir.") - return create_app_fn( - web_dir=Path(output.web_dir), - session_manager=session_manager, - request_session_url=_request_session_url(output), - preload_name=preload_name, - ) - - -def _request_session_url(output: WebRTCOutputSpec) -> str: - host = "127.0.0.1" if output.host in {"0.0.0.0", "::"} else output.host - return f"http://{host}:{output.port}{output.request_session_path}" +The concrete server helper lives in the WebRTC serving package so the runtime +demo API does not depend on transport infrastructure. +""" +from __future__ import annotations -__all__ = [ - "CreateWebRTCApp", - "RunWebRTCServer", - "serve_webrtc_demo", -] +__all__: list[str] = [] diff --git a/flashdreams/flashdreams/runtime/keyboard.py b/flashdreams/flashdreams/runtime/keyboard.py new file mode 100644 index 000000000..5ef83fc82 --- /dev/null +++ b/flashdreams/flashdreams/runtime/keyboard.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Keyboard state helpers shared by runtime input canonicalizers.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +DEFAULT_SUPPORTED_KEYS = frozenset({"w", "a", "s", "d", "q", "e", "i", "k", "j", "l"}) +DRIVING_SUPPORTED_KEYS = frozenset( + {"w", "a", "s", "d", "up", "down", "left", "right", "space"} +) +WSAD_SUPPORTED_KEYS = frozenset({"w", "a", "s", "d"}) +KEY_ALIASES = { + "arrowup": "w", + "arrowleft": "a", + "arrowdown": "s", + "arrowright": "d", +} + + +@dataclass(frozen=True, slots=True) +class ResetRequest: + """Transport-neutral request to reset the realtime rollout.""" + + reason: str | None = None + request_id: str | None = None + + +@dataclass(frozen=True, slots=True) +class PromptRequest: + """Transport-neutral prompt update request.""" + + prompt: str + negative_prompt: str | None = None + request_id: str | None = None + + +@dataclass(frozen=True, slots=True) +class ImageRequest: + """Transport-neutral image update request.""" + + data: bytes + content_type: str + request_id: str | None = None + + +@dataclass(frozen=True, slots=True) +class SparseInputSnapshot: + """Sparse input state sampled at a realtime loop boundary.""" + + timestamp_s: float + pressed_keys: frozenset[str] = field(default_factory=frozenset) + effective_keys: frozenset[str] = field(default_factory=frozenset) + reset: ResetRequest | None = None + prompt: PromptRequest | None = None + image: ImageRequest | None = None + + +def normalize_key(key: str) -> str: + normalized = key.strip().lower() + return KEY_ALIASES.get(normalized, normalized) + + +@dataclass(slots=True) +class KeyboardState: + pressed_keys: set[str] = field(default_factory=set) + supported_keys: frozenset[str] = DEFAULT_SUPPORTED_KEYS + _press_order: dict[str, int] = field(default_factory=dict) + _press_counter: int = 0 + + def apply_event(self, *, event: str, key: str) -> bool: + normalized_key = normalize_key(key) + if normalized_key not in self.supported_keys: + return False + + normalized_event = event.strip().lower() + if normalized_event == "keydown": + self.pressed_keys.add(normalized_key) + self._press_counter += 1 + self._press_order[normalized_key] = self._press_counter + return True + if normalized_event == "keyup": + self.pressed_keys.discard(normalized_key) + self._press_order.pop(normalized_key, None) + return True + return False + + def snapshot(self) -> frozenset[str]: + return frozenset(self.pressed_keys) + + def sparse_snapshot(self, *, timestamp_s: float) -> SparseInputSnapshot: + return SparseInputSnapshot( + timestamp_s=timestamp_s, + pressed_keys=self.snapshot(), + effective_keys=self.resolved_effective_keys(), + ) + + def _latest_pressed(self, keys: tuple[str, ...]) -> str | None: + latest_key: str | None = None + latest_idx = -1 + for key in keys: + if key not in self.pressed_keys: + continue + idx = self._press_order.get(key, -1) + if idx >= latest_idx: + latest_idx = idx + latest_key = key + return latest_key + + def resolved_effective_keys(self) -> frozenset[str]: + effective: set[str] = set() + for key in ( + self._latest_pressed(("w", "s")), + self._latest_pressed(("a", "d", "j", "l")), + self._latest_pressed(("q", "e")), + self._latest_pressed(("i", "k")), + ): + if key is not None: + effective.add(key) + return frozenset(key for key in effective if key in self.supported_keys) + + +__all__ = [ + "DEFAULT_SUPPORTED_KEYS", + "DRIVING_SUPPORTED_KEYS", + "ImageRequest", + "KEY_ALIASES", + "KeyboardState", + "PromptRequest", + "ResetRequest", + "SparseInputSnapshot", + "WSAD_SUPPORTED_KEYS", + "normalize_key", +] diff --git a/flashdreams/flashdreams/serving/realtime/input.py b/flashdreams/flashdreams/serving/realtime/input.py index 6baf92dde..800450e6e 100644 --- a/flashdreams/flashdreams/serving/realtime/input.py +++ b/flashdreams/flashdreams/serving/realtime/input.py @@ -1,385 +1,32 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Input state and sparse-control helpers for realtime serving.""" +"""Compatibility re-exports for transport-neutral realtime input containers.""" from __future__ import annotations -from collections import deque -from dataclasses import dataclass, field -from typing import Literal - -import numpy as np - -DEFAULT_SUPPORTED_KEYS = frozenset({"w", "a", "s", "d", "q", "e", "i", "k", "j", "l"}) -DRIVING_SUPPORTED_KEYS = frozenset( - {"w", "a", "s", "d", "up", "down", "left", "right", "space"} +from flashdreams.runtime.keyboard import ( + DEFAULT_SUPPORTED_KEYS, + DRIVING_SUPPORTED_KEYS, + KEY_ALIASES, + WSAD_SUPPORTED_KEYS, + ImageRequest, + KeyboardState, + PromptRequest, + ResetRequest, + SparseInputSnapshot, + normalize_key, ) -WSAD_SUPPORTED_KEYS = frozenset({"w", "a", "s", "d"}) -KEY_ALIASES = { - "arrowup": "w", - "arrowleft": "a", - "arrowdown": "s", - "arrowright": "d", -} - - -@dataclass(frozen=True, slots=True) -class ResetRequest: - """Transport-neutral request to reset the realtime rollout.""" - - reason: str | None = None - request_id: str | None = None - - -@dataclass(frozen=True, slots=True) -class PromptRequest: - """Transport-neutral prompt update request.""" - - prompt: str - negative_prompt: str | None = None - request_id: str | None = None - - -@dataclass(frozen=True, slots=True) -class ImageRequest: - """Transport-neutral image update request.""" - - data: bytes - content_type: str - request_id: str | None = None - - -@dataclass(frozen=True, slots=True) -class SparseInputSnapshot: - """Sparse input state sampled at a realtime loop boundary.""" - - timestamp_s: float - pressed_keys: frozenset[str] = field(default_factory=frozenset) - effective_keys: frozenset[str] = field(default_factory=frozenset) - reset: ResetRequest | None = None - prompt: PromptRequest | None = None - image: ImageRequest | None = None - - -def normalize_key(key: str) -> str: - normalized = key.strip().lower() - return KEY_ALIASES.get(normalized, normalized) - - -@dataclass(slots=True) -class KeyboardState: - pressed_keys: set[str] = field(default_factory=set) - supported_keys: frozenset[str] = DEFAULT_SUPPORTED_KEYS - _press_order: dict[str, int] = field(default_factory=dict) - _press_counter: int = 0 - - def apply_event(self, *, event: str, key: str) -> bool: - normalized_key = normalize_key(key) - if normalized_key not in self.supported_keys: - return False - - normalized_event = event.strip().lower() - if normalized_event == "keydown": - self.pressed_keys.add(normalized_key) - self._press_counter += 1 - self._press_order[normalized_key] = self._press_counter - return True - if normalized_event == "keyup": - self.pressed_keys.discard(normalized_key) - self._press_order.pop(normalized_key, None) - return True - return False - - def snapshot(self) -> frozenset[str]: - return frozenset(self.pressed_keys) - - def sparse_snapshot(self, *, timestamp_s: float) -> SparseInputSnapshot: - return SparseInputSnapshot( - timestamp_s=timestamp_s, - pressed_keys=self.snapshot(), - effective_keys=self.resolved_effective_keys(), - ) - - def _latest_pressed(self, keys: tuple[str, ...]) -> str | None: - latest_key: str | None = None - latest_idx = -1 - for key in keys: - if key not in self.pressed_keys: - continue - idx = self._press_order.get(key, -1) - if idx >= latest_idx: - latest_idx = idx - latest_key = key - return latest_key - - def resolved_effective_keys(self) -> frozenset[str]: - effective: set[str] = set() - for key in ( - self._latest_pressed(("w", "s")), - self._latest_pressed(("a", "d", "j", "l")), - self._latest_pressed(("q", "e")), - self._latest_pressed(("i", "k")), - ): - if key is not None: - effective.add(key) - return frozenset(key for key in effective if key in self.supported_keys) - - -PoseSegment = tuple[float, float, frozenset[str]] - - -class KeyboardResampler: - """Resample sparse keydown/keyup edges into a chunk timeline.""" - - def __init__( - self, - *, - fps: float, - start_v: float = 0.0, - supported_keys: frozenset[str] = DEFAULT_SUPPORTED_KEYS, - ) -> None: - if fps <= 0: - raise ValueError("fps must be > 0") - self._fps = float(fps) - self._dt = 1.0 / self._fps - self._supported_keys = supported_keys - self.next_chunk_start_v = start_v - self._event_log: deque[tuple[float, dict[str, str]]] = deque() - self._carried_state = KeyboardState(supported_keys=supported_keys) - - @property - def fps(self) -> float: - return self._fps - - @property - def dt(self) -> float: - return self._dt - - def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: - self._event_log.append((arrival_t, {"event": event, "key": key})) - - def sample_chunk(self, num_frames: int) -> tuple[list[PoseSegment], list[float]]: - if num_frames < 1: - raise ValueError("num_frames must be >= 1") - - chunk_start_v = self.next_chunk_start_v - chunk_end_v = chunk_start_v + num_frames * self._dt - - while self._event_log and self._event_log[0][0] < chunk_start_v: - _, payload = self._event_log.popleft() - self._carried_state.apply_event(**payload) - - segments: list[PoseSegment] = [] - prev_t = chunk_start_v - prev_state = self._carried_state.resolved_effective_keys() - while self._event_log and self._event_log[0][0] <= chunk_end_v: - event_t, payload = self._event_log.popleft() - if event_t > prev_t: - segments.append((prev_t, event_t, prev_state)) - self._carried_state.apply_event(**payload) - prev_state = self._carried_state.resolved_effective_keys() - prev_t = event_t - if prev_t < chunk_end_v: - segments.append((prev_t, chunk_end_v, prev_state)) - elif not segments: - segments.append((chunk_start_v, chunk_end_v, prev_state)) - - frame_times = [chunk_start_v + (i + 1) * self._dt for i in range(num_frames)] - self.next_chunk_start_v = chunk_end_v - return segments, frame_times - - def reset(self, *, start_v: float) -> None: - self._event_log.clear() - self._carried_state = KeyboardState(supported_keys=self._supported_keys) - self.next_chunk_start_v = start_v - - def event_log_size(self) -> int: - return len(self._event_log) - - -def _rotation_matrix(axis: str, angle_rad: float) -> np.ndarray: - cos_t = np.float32(np.cos(angle_rad)) - sin_t = np.float32(np.sin(angle_rad)) - if axis == "x": - return np.array( - [ - [1.0, 0.0, 0.0], - [0.0, cos_t, -sin_t], - [0.0, sin_t, cos_t], - ], - dtype=np.float32, - ) - if axis == "y": - return np.array( - [ - [cos_t, 0.0, sin_t], - [0.0, 1.0, 0.0], - [-sin_t, 0.0, cos_t], - ], - dtype=np.float32, - ) - if axis == "z": - return np.array( - [ - [cos_t, -sin_t, 0.0], - [sin_t, cos_t, 0.0], - [0.0, 0.0, 1.0], - ], - dtype=np.float32, - ) - return np.eye(3, dtype=np.float32) - - -@dataclass(slots=True) -class CameraPoseIntegrator: - """Integrate a piecewise-constant keyboard timeline into a camera trajectory.""" - - move_speed_per_s: float = 0.8 - rotate_speed_rad_per_s: float = float(np.deg2rad(32.0)) - pitch_limit_rad: float = float(np.deg2rad(85.0)) - coordinate_system: Literal["RDF", "FLU"] = "RDF" - _current_pose: np.ndarray = field( - default_factory=lambda: np.eye(4, dtype=np.float32), - ) - _current_pitch: float = 0.0 - - def __post_init__(self) -> None: - if self.coordinate_system not in {"RDF", "FLU"}: - raise ValueError( - "coordinate_system must be 'RDF' (right-down-forward) " - "or 'FLU' (forward-left-up)" - ) - - def reset(self, pose: np.ndarray | None = None) -> None: - if pose is None: - self._current_pose = np.eye(4, dtype=np.float32) - self._current_pitch = 0.0 - return - if pose.shape != (4, 4): - raise ValueError(f"Expected pose shape (4, 4), got {pose.shape}") - self._current_pose = pose.astype(np.float32, copy=True) - if self.coordinate_system == "FLU": - self._current_pitch = float(np.arcsin(np.clip(pose[2, 0], -1.0, 1.0))) - else: - self._current_pitch = float(np.arctan2(pose[2, 1], pose[1, 1])) - - def current_pose(self) -> np.ndarray: - return self._current_pose.copy() - - def _advance(self, *, state: frozenset[str], duration: float) -> None: - if duration <= 0: - return - - yaw_rate = 0.0 - if self.coordinate_system == "FLU": - if "a" in state or "j" in state: - yaw_rate += self.rotate_speed_rad_per_s - if "d" in state or "l" in state: - yaw_rate -= self.rotate_speed_rad_per_s - else: - if "a" in state or "j" in state: - yaw_rate -= self.rotate_speed_rad_per_s - if "d" in state or "l" in state: - yaw_rate += self.rotate_speed_rad_per_s - pitch_rate = 0.0 - if "i" in state: - pitch_rate += self.rotate_speed_rad_per_s - if "k" in state: - pitch_rate -= self.rotate_speed_rad_per_s - - yaw_delta = yaw_rate * duration - pitch_delta = pitch_rate * duration - - new_pitch = self._current_pitch + pitch_delta - if -self.pitch_limit_rad <= new_pitch <= self.pitch_limit_rad: - self._current_pitch = new_pitch - else: - pitch_delta = 0.0 - - rot = self._current_pose[:3, :3] - trans = self._current_pose[:3, 3] - if self.coordinate_system == "FLU": - rot_pitch = _rotation_matrix("y", -pitch_delta) - rot_yaw = _rotation_matrix("z", yaw_delta) - else: - rot_pitch = _rotation_matrix("x", pitch_delta) - rot_yaw = _rotation_matrix("y", yaw_delta) - rot_new = rot_yaw @ rot @ rot_pitch - - forward_rate = 0.0 - if "w" in state: - forward_rate += self.move_speed_per_s - if "s" in state: - forward_rate -= self.move_speed_per_s - right_rate = 0.0 - if "e" in state: - right_rate += self.move_speed_per_s - if "q" in state: - right_rate -= self.move_speed_per_s - - if self.coordinate_system == "FLU": - vec_forward = rot_new[:, 0] - vec_right = -rot_new[:, 1] - forward_flat = np.array( - [vec_forward[0], vec_forward[1], 0.0], dtype=np.float32 - ) - right_flat = np.array([vec_right[0], vec_right[1], 0.0], dtype=np.float32) - else: - vec_right = rot_new[:, 0] - vec_forward = rot_new[:, 2] - forward_flat = np.array( - [vec_forward[0], 0.0, vec_forward[2]], dtype=np.float32 - ) - right_flat = np.array([vec_right[0], 0.0, vec_right[2]], dtype=np.float32) - forward_norm = np.linalg.norm(forward_flat) - right_norm = np.linalg.norm(right_flat) - if forward_norm > 0: - forward_flat /= forward_norm - if right_norm > 0: - right_flat /= right_norm - - move_vec = forward_flat * (forward_rate * duration) + right_flat * ( - right_rate * duration - ) - self._current_pose = np.eye(4, dtype=np.float32) - self._current_pose[:3, :3] = rot_new - self._current_pose[:3, 3] = trans + move_vec - - def integrate_chunk( - self, - *, - segments: list[PoseSegment], - frame_times: list[float], - ) -> np.ndarray: - if not segments: - raise ValueError("segments must be non-empty") - if not frame_times: - raise ValueError("frame_times must be non-empty") - chunk_start = segments[0][0] - chunk_end = segments[-1][1] - if any( - frame_times[i] >= frame_times[i + 1] for i in range(len(frame_times) - 1) - ): - raise ValueError("frame_times must be strictly increasing") - if frame_times[0] < chunk_start - 1e-9 or frame_times[-1] > chunk_end + 1e-9: - raise ValueError( - "frame_times must lie within the chunk window " - f"[{chunk_start}, {chunk_end}]" - ) - - poses: list[np.ndarray] = [] - cur_t = chunk_start - ft_idx = 0 - for _, seg_end, seg_state in segments: - while ft_idx < len(frame_times) and frame_times[ft_idx] <= seg_end: - target_t = frame_times[ft_idx] - self._advance(state=seg_state, duration=target_t - cur_t) - cur_t = target_t - poses.append(self._current_pose.copy()) - ft_idx += 1 - if seg_end > cur_t: - self._advance(state=seg_state, duration=seg_end - cur_t) - cur_t = seg_end - return np.stack(poses, axis=0).astype(np.float32) +__all__ = [ + "DEFAULT_SUPPORTED_KEYS", + "DRIVING_SUPPORTED_KEYS", + "ImageRequest", + "KEY_ALIASES", + "KeyboardState", + "PromptRequest", + "ResetRequest", + "SparseInputSnapshot", + "WSAD_SUPPORTED_KEYS", + "normalize_key", +] diff --git a/flashdreams/flashdreams/serving/webrtc/bootstrap.py b/flashdreams/flashdreams/serving/webrtc/bootstrap.py index 9a5405a81..353e1a88c 100644 --- a/flashdreams/flashdreams/serving/webrtc/bootstrap.py +++ b/flashdreams/flashdreams/serving/webrtc/bootstrap.py @@ -6,96 +6,22 @@ from __future__ import annotations import gc -import logging -import os -from collections.abc import Callable -from dataclasses import dataclass -from typing import Any import torch import torch.distributed as dist from aiohttp import web from loguru import logger +from flashdreams.runtime.demo.bootstrap import ( + DistributedDemoContext as WebRTCDistributedContext, +) +from flashdreams.runtime.demo.bootstrap import ( + configure_logging, + initialize_cuda_distributed, +) from flashdreams.serving.webrtc.runtime import WebRTCServerLifecycle -@dataclass(frozen=True, slots=True) -class WebRTCDistributedContext: - """CUDA/distributed launch context for a WebRTC demo server.""" - - device: torch.device - world_rank: int - world_size: int - - -def configure_logging(*, world_rank: int | None = None) -> None: - from flashdreams.core.distributed import configure_loguru_for_distributed - - configure_loguru_for_distributed(world_rank=world_rank) - for logger_name in ("aioice", "aioice.ice", "aiortc"): - logging.getLogger(logger_name).setLevel(logging.WARNING) - - -def _distributed_init() -> None: - from flashdreams.core.distributed import init as distributed_init - - distributed_init() - - -def initialize_cuda_distributed( - *, - default_device: str | torch.device = "cuda:0", - distributed_init_fn: Callable[[], object] | None = None, - configure_logging_fn: Callable[..., None] = configure_logging, - torch_module: Any = torch, - dist_module: Any = dist, -) -> WebRTCDistributedContext: - """Initialize CUDA and optional torch.distributed for WebRTC serving.""" - if not torch_module.cuda.is_available(): - raise RuntimeError("CUDA is required for inference in the WebRTC server.") - - has_rank = "RANK" in os.environ - has_world_size = "WORLD_SIZE" in os.environ - if has_rank != has_world_size: - raise RuntimeError( - "Distributed launch expects both RANK and WORLD_SIZE to be set." - ) - - distributed_launch = has_rank and has_world_size - if distributed_launch: - if distributed_init_fn is None: - distributed_init_fn = _distributed_init - distributed_init_fn() - world_rank = dist_module.get_rank() - world_size = dist_module.get_world_size() - else: - world_rank = 0 - world_size = 1 - - device_count = torch_module.cuda.device_count() - if device_count < 1: - raise RuntimeError("CUDA device count must be >= 1 for inference.") - if distributed_launch: - local_rank = world_rank % device_count - torch_device = torch_module.device(f"cuda:{local_rank}") - else: - torch_device = torch_module.device(default_device) - if torch_device.type != "cuda": - raise RuntimeError( - f"CUDA device is required for inference, got {torch_device}." - ) - if torch_device.index is None: - torch_device = torch_module.device("cuda:0") - torch_module.cuda.set_device(torch_device) - configure_logging_fn(world_rank=world_rank) - return WebRTCDistributedContext( - device=torch_device, - world_rank=world_rank, - world_size=world_size, - ) - - def run_webrtc_server( *, world_rank: int, diff --git a/flashdreams/flashdreams/serving/webrtc/controls.py b/flashdreams/flashdreams/serving/webrtc/controls.py index 7f1d8461e..ecc2d3c34 100644 --- a/flashdreams/flashdreams/serving/webrtc/controls.py +++ b/flashdreams/flashdreams/serving/webrtc/controls.py @@ -5,15 +5,12 @@ from __future__ import annotations -from flashdreams.serving.realtime.input import ( +from flashdreams.runtime.keyboard import ( DEFAULT_SUPPORTED_KEYS, KEY_ALIASES, WSAD_SUPPORTED_KEYS, - CameraPoseIntegrator, ImageRequest, - KeyboardResampler, KeyboardState, - PoseSegment, PromptRequest, ResetRequest, SparseInputSnapshot, @@ -24,11 +21,8 @@ "DEFAULT_SUPPORTED_KEYS", "KEY_ALIASES", "WSAD_SUPPORTED_KEYS", - "CameraPoseIntegrator", "ImageRequest", - "KeyboardResampler", "KeyboardState", - "PoseSegment", "PromptRequest", "ResetRequest", "SparseInputSnapshot", diff --git a/flashdreams/flashdreams/serving/webrtc/demo.py b/flashdreams/flashdreams/serving/webrtc/demo.py new file mode 100644 index 000000000..0e4b281ab --- /dev/null +++ b/flashdreams/flashdreams/serving/webrtc/demo.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared WebRTC demo construction.""" + +from __future__ import annotations + +from collections.abc import Callable +from importlib.resources import files +from pathlib import Path +from typing import Any + +from aiohttp import web + +from flashdreams.runtime.demo.spec import WebRTCAppResources, WebRTCOutputSpec +from flashdreams.serving.webrtc.bootstrap import run_webrtc_server +from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager +from flashdreams.serving.webrtc.server import ( + close_package_resources, + create_packaged_webrtc_app, + create_webrtc_app, +) + +CreateWebRTCApp = Callable[..., web.Application] +RunWebRTCServer = Callable[..., None] + + +def serve_webrtc_demo( + *, + output: WebRTCOutputSpec, + model_id: str, + session_manager: BaseWebRTCSessionManager[Any, Any], + app_resources: WebRTCAppResources, + world_rank: int = 0, + create_app_fn: CreateWebRTCApp = create_webrtc_app, + server_runner: RunWebRTCServer = run_webrtc_server, +) -> web.Application | None: + """Serve a prepared model WebRTC runtime through the shared transport.""" + app = ( + _create_app( + output=output, + model_id=model_id, + app_resources=app_resources, + session_manager=session_manager, + create_app_fn=create_app_fn, + ) + if world_rank == 0 + else None + ) + server_runner( + world_rank=world_rank, + session_manager=session_manager, + app=app, + host=output.host, + port=output.port, + ) + return app + + +def _create_app( + *, + output: WebRTCOutputSpec, + model_id: str, + app_resources: WebRTCAppResources, + session_manager: BaseWebRTCSessionManager[Any, Any], + create_app_fn: CreateWebRTCApp, +) -> web.Application: + if output.web_dir is not None: + return _build_webrtc_app( + output=output, + session_manager=session_manager, + create_app_fn=create_app_fn, + preload_name=output.preload_name or app_resources.preload_name or model_id, + ) + 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, + cleanup_callback=close_package_resources, + ) + + +def _build_webrtc_app( + *, + output: WebRTCOutputSpec, + session_manager: BaseWebRTCSessionManager[Any, Any], + create_app_fn: CreateWebRTCApp, + preload_name: str, +) -> web.Application: + if output.web_dir is None: + raise ValueError("WebRTC app creation requires output.web_dir.") + return create_app_fn( + web_dir=Path(output.web_dir), + session_manager=session_manager, + request_session_url=_request_session_url(output), + preload_name=preload_name, + ) + + +def _request_session_url(output: WebRTCOutputSpec) -> str: + host = "127.0.0.1" if output.host in {"0.0.0.0", "::"} else output.host + return f"http://{host}:{output.port}{output.request_session_path}" + + +__all__ = [ + "CreateWebRTCApp", + "RunWebRTCServer", + "serve_webrtc_demo", +] diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index 9aa22653d..59177742e 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -24,13 +24,13 @@ from loguru import logger from flashdreams.runtime.demo import ( - SPARSE_KEY_SEGMENTS_METADATA_KEY, DemoSpec, InMemorySessionMetricsRecorder, ModelInputProvider, PreparedScenario, PreparedStep, ProviderCapabilities, + RealtimeEventResampler, ResamplerRealtimeClock, RunContext, RuntimeHost, @@ -52,13 +52,9 @@ UserInputs, UserInputSchema, ) +from flashdreams.runtime.keyboard import DEFAULT_SUPPORTED_KEYS, normalize_key from flashdreams.runtime.mapping import InputMapping from flashdreams.runtime.types import StepRequest, StepResult -from flashdreams.serving.realtime.input import ( - DEFAULT_SUPPORTED_KEYS, - KeyboardResampler, - normalize_key, -) from flashdreams.serving.webrtc.encoders import ( DefaultRTCEncoder, EncoderBackend, @@ -119,6 +115,7 @@ _STEP_REQUEST_KEY = "webrtc_step_request" _SEGMENTS_KEY = "webrtc_segments" _FRAME_TIMES_KEY = "webrtc_frame_times" +_LEGACY_SPARSE_KEY_SEGMENTS_METADATA_KEY = "sparse_key_segments" _RuntimeT = TypeVar("_RuntimeT") _RuntimeConfigT = TypeVar("_RuntimeConfigT", bound=WebRTCRuntimeConfig) @@ -450,7 +447,7 @@ def _prepare_segment_step( request: Any, user_window: UserInputWindow, ) -> InferenceInput: - segments = user_window.metadata.get(SPARSE_KEY_SEGMENTS_METADATA_KEY) + segments = user_window.metadata.get(_LEGACY_SPARSE_KEY_SEGMENTS_METADATA_KEY) if not isinstance(segments, tuple): raise RuntimeError("WebRTC user window is missing resampled key segments.") window = TimeWindow(start_s=user_window.start_s, end_s=user_window.end_s) @@ -604,7 +601,8 @@ class ManagedWebRTCSession: video_track: BufferedVideoTrack | NVENCVideoTrack video_encoder: VideoEncoder peer_connection: Any - resampler: KeyboardResampler + resampler: RealtimeEventResampler + legacy_segment_resampler: Any | None = None control_channel: Any | None = None generation_task: asyncio.Task[Any] | None = None first_action_received: asyncio.Event = field(default_factory=asyncio.Event) @@ -682,6 +680,7 @@ def __init__( shared_spec_factory: Callable[[Any], DemoSpec] | None = None, shared_scenario: PreparedScenario | None = None, shared_pipeline_factory: Callable[[], StepPipeline] | None = None, + legacy_segment_resampler_factory: Callable[..., Any] | None = None, ) -> None: if client_liveness_timeout_s <= 0: raise ValueError("client_liveness_timeout_s must be > 0") @@ -714,6 +713,7 @@ def __init__( self._shared_scenario = shared_scenario self._shared_pipeline_factory = shared_pipeline_factory self._shared_video_encoder: VideoEncoder | None = None + self._legacy_segment_resampler_factory = legacy_segment_resampler_factory @property def pending_session_input(self) -> Any: @@ -731,19 +731,35 @@ def set_pending_session_input(self, session_input: Any) -> None: raise SessionBusyError(self.busy_message) self._pending_session_input = session_input - def _make_resampler(self, *, start_v: float) -> KeyboardResampler: + def _make_resampler(self, *, start_v: float) -> RealtimeEventResampler: return self._make_resampler_at_fps(start_v=start_v, fps=self.fps) def _make_resampler_at_fps( self, *, start_v: float, fps: float - ) -> KeyboardResampler: + ) -> RealtimeEventResampler: + return RealtimeEventResampler(fps=fps, start_v=start_v) + + def _make_legacy_segment_resampler_at_fps( + self, *, start_v: float, fps: float + ) -> Any: + factory = self._legacy_segment_resampler_factory + if factory is None: + raise RuntimeError( + "Legacy WebRTC segment runtimes require " + "legacy_segment_resampler_factory." + ) supported_control_keys = self._effective_supported_control_keys() - if supported_control_keys is None: - return KeyboardResampler(fps=fps, start_v=start_v) - return KeyboardResampler( - fps=fps, - start_v=start_v, - supported_keys=supported_control_keys, + kwargs: dict[str, object] = { + "fps": fps, + "start_v": start_v, + } + if supported_control_keys is not None: + kwargs["supported_keys"] = supported_control_keys + return factory(**kwargs) + + def _needs_legacy_segment_metadata(self) -> bool: + return self._shared_adapter is None and not _runtime_drives_inference_session( + self._runtime ) def _effective_supported_control_keys(self) -> frozenset[str] | None: @@ -1425,7 +1441,21 @@ async def _create_answer_with_runtime_ready_locked( start_v=0.0, fps=self._runtime_input_fps(self._runtime), ) - input_source = WebRTCInputSource(resampler=resampler) + legacy_segment_resampler = None + if self._needs_legacy_segment_metadata(): + legacy_segment_resampler = self._make_legacy_segment_resampler_at_fps( + start_v=0.0, + fps=self._runtime_input_fps(self._runtime), + ) + input_source = WebRTCInputSource( + resampler=resampler, + legacy_segment_resampler=legacy_segment_resampler, + legacy_segments_metadata_key=( + _LEGACY_SPARSE_KEY_SEGMENTS_METADATA_KEY + if legacy_segment_resampler is not None + else None + ), + ) transport = WebRTCTransportService(loop=loop) managed_session = ManagedWebRTCSession( runtime=self._runtime, @@ -1433,6 +1463,7 @@ async def _create_answer_with_runtime_ready_locked( video_encoder=video_encoder, peer_connection=peer_connection, resampler=resampler, + legacy_segment_resampler=legacy_segment_resampler, input_source=input_source, transport=transport, reservation=reservation, @@ -1718,9 +1749,8 @@ async def _handle_datachannel_message( ) return - # Stamp arrival on the same monotonic clock that seeds the - # resampler's ``next_chunk_start_v`` so virtual-time comparisons in - # ``KeyboardResampler.sample_chunk`` are well-defined. + # Stamp arrival on the same monotonic clock that seeds the realtime + # window clock so user-input windows can be compared directly. arrival_t = asyncio.get_running_loop().time() if managed_session.inference_session is not None: try: @@ -1734,7 +1764,9 @@ async def _handle_datachannel_message( self._send_json(channel, make_error_payload(str(exc))) if event != "keyup": return - managed_session.resampler.on_edge(arrival_t=arrival_t, event=event, key=key) + legacy_resampler = managed_session.legacy_segment_resampler + if legacy_resampler is not None: + legacy_resampler.on_edge(arrival_t=arrival_t, event=event, key=key) managed_session.pending_action_arrivals.append(arrival_t) # Releases the generation worker, which blocks on this until the # user actually interacts. Idempotent once already set. @@ -1943,14 +1975,13 @@ def _handle_shared_delivery_error( async def _generation_worker( self, *, managed_session: ManagedWebRTCSession ) -> None: - """Drive back-to-back chunk generation aligned to the resampler clock. + """Drive back-to-back chunk generation aligned to the realtime clock. Sits idle until the first keyboard event arrives, then drives the chunk loop. Each iteration waits for wallclock to catch up to the - *end* of the next chunk's virtual window, samples the chunk's - 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 + *end* of the next chunk's virtual window, hands legacy segment data and + frame times to the runtime, and pushes generated frames into the video + track. The track's bounded queue then paces the loop to playback via backpressure on ``BufferedVideoTrack.enqueue_result``. """ loop = asyncio.get_running_loop() @@ -2010,8 +2041,15 @@ async def _generation_worker( t_before_gen = loop.time() chunk_start_v = resampler.next_chunk_start_v - segments, frame_times = resampler.sample_chunk(input_num_frames) + frame_times = list(resampler.sample_chunk(input_num_frames)) chunk_end_v = resampler.next_chunk_start_v + segments: list[Any] = [] + legacy_resampler = managed_session.legacy_segment_resampler + if legacy_resampler is not None: + legacy_resampler.next_chunk_start_v = chunk_start_v + segments, frame_times = legacy_resampler.sample_chunk( + input_num_frames + ) segment_request = replace( request, user_input_window=TimeWindow( diff --git a/flashdreams/flashdreams/serving/webrtc/runtime.py b/flashdreams/flashdreams/serving/webrtc/runtime.py index 32c7d4365..1472807ca 100644 --- a/flashdreams/flashdreams/serving/webrtc/runtime.py +++ b/flashdreams/flashdreams/serving/webrtc/runtime.py @@ -20,7 +20,6 @@ ) 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, @@ -90,7 +89,7 @@ async def step( self, *, request: StepRequest, - segments: list[PoseSegment], + segments: list[Any], frame_times: list[float], ) -> StepResult: ... @@ -203,7 +202,7 @@ async def step( self, *, request: StepRequest, - segments: list[PoseSegment], + segments: list[Any], frame_times: list[float], ) -> StepResult: self._require_open_and_initialized(session=True) @@ -289,7 +288,7 @@ def _reset_rollout_sync_all_ranks( @distributed_op(WebRTCControlSignal.ACTION_STEP) def _generate_chunk_sync_all_ranks( self, - segments: list[PoseSegment], + segments: list[Any], frame_times: list[float], ) -> StepResult: return self._generate_one_chunk_sync(segments=segments, frame_times=frame_times) @@ -325,7 +324,7 @@ def _reset_rollout_sync( def _generate_one_chunk_sync( self, *, - segments: list[PoseSegment], + segments: list[Any], frame_times: list[float], ) -> StepResult: ... diff --git a/flashdreams/flashdreams/serving/webrtc/services.py b/flashdreams/flashdreams/serving/webrtc/services.py index 3ce38e4be..cfcb11029 100644 --- a/flashdreams/flashdreams/serving/webrtc/services.py +++ b/flashdreams/flashdreams/serving/webrtc/services.py @@ -57,7 +57,6 @@ run_demo_session_async, ) from flashdreams.runtime.demo.timing import ( - SPARSE_KEY_SEGMENTS_METADATA_KEY, ActivationResult, CatchUpPolicy, DeterministicClock, @@ -350,6 +349,8 @@ class WebRTCInputSource: """Realtime source fed by browser data-channel events.""" resampler: Any + legacy_segment_resampler: Any | None = None + legacy_segments_metadata_key: str | None = None max_lag_s: float | None = None catch_up_policy: CatchUpPolicy = "fold" user_input_schema: UserInputSchema = field( @@ -387,6 +388,8 @@ def is_finished(self) -> bool: def reset(self, *, start_v: float) -> None: self.resampler.reset(start_v=start_v) + if self.legacy_segment_resampler is not None: + self.legacy_segment_resampler.reset(start_v=start_v) self._events.clear() self._activation_signal.clear() @@ -486,11 +489,18 @@ async def next_realtime_window( policy=self.catch_up_policy, ) start_s = float(self.resampler.next_chunk_start_v) - segments, frame_times = self.resampler.sample_chunk(input_frame_count) + frame_times = tuple(self.resampler.sample_chunk(input_frame_count)) end_s = float(self.resampler.next_chunk_start_v) - metadata: dict[str, object] = { - SPARSE_KEY_SEGMENTS_METADATA_KEY: tuple(segments), - } + metadata: dict[str, object] = {} + legacy_resampler = self.legacy_segment_resampler + legacy_metadata_key = self.legacy_segments_metadata_key + if legacy_resampler is not None and legacy_metadata_key is not None: + legacy_resampler.next_chunk_start_v = start_s + segments, legacy_frame_times = legacy_resampler.sample_chunk( + input_frame_count + ) + metadata[legacy_metadata_key] = tuple(segments) + frame_times = tuple(legacy_frame_times) if start_s > pre_catch_up_start_s: metadata[WEBRTC_SKIPPED_INPUTS_METADATA_KEY] = UserInputs( events=self._events_for_window(pre_catch_up_start_s, start_s) @@ -533,7 +543,12 @@ def _record_action( kind="error", error="Action payload must include non-empty 'key'.", ) - self.resampler.on_edge(arrival_t=timestamp_s, event=event, key=key) + if self.legacy_segment_resampler is not None: + self.legacy_segment_resampler.on_edge( + arrival_t=timestamp_s, + event=event, + key=key, + ) self.record_user_event( timestamp_s=timestamp_s, event_type="key_down" if event == "keydown" else "key_up", @@ -671,6 +686,7 @@ def __init__( self._on_chunk_delivery = on_chunk_delivery self._on_error = on_error self._pending: dict[Future[WebRTCChunkDelivery], int] = {} + self._delivery_lock = asyncio.Lock() self._lock = threading.Lock() self._closed = False self._generation = 0 @@ -727,12 +743,12 @@ def submit_chunk( metadata={"reason": "stale generation"}, ) if len(self._pending) >= self._max_pending_chunks: - return WebRTCOutputBridgeDecision( - accepted=False, - dropped=True, - drop_policy="drop_newest", - metadata={"reason": "pending queue full"}, - ) + stale = self._pop_pending_locked() + else: + stale = () + if stale: + self._cancel_stale_deliveries(stale) + self._schedule_track_flush() payload = prepare(result, self._video_track) chunk = WebRTCChunkDelivery( delivery=None, @@ -760,11 +776,27 @@ def submit_chunk( metadata={"reason": "stale generation"}, ) if len(self._pending) >= self._max_pending_chunks: + stale = self._pop_pending_locked() + else: + stale = () + if stale: + self._cancel_stale_deliveries(stale) + self._schedule_track_flush() + with self._lock: + if self._closed: + return WebRTCOutputBridgeDecision( + accepted=False, + should_stop=True, + dropped=True, + drop_policy="drop_newest", + metadata={"reason": "closed"}, + ) + if generation < self._generation: return WebRTCOutputBridgeDecision( accepted=False, dropped=True, drop_policy="drop_newest", - metadata={"reason": "pending queue full"}, + metadata={"reason": "stale generation"}, ) future = asyncio.run_coroutine_threadsafe( self._deliver( @@ -805,11 +837,16 @@ async def _deliver( with self._lock: if self._closed or generation < self._generation: raise asyncio.CancelledError - delivery = await self._video_encoder.deliver_prepared_chunk( - payload, - self._video_track, - force_keyframe=force_keyframe, - ) + async with self._delivery_lock: + with self._lock: + if self._closed or generation < self._generation: + raise asyncio.CancelledError + await self._flush_full_track_queue(frame_count=chunk.frame_count) + delivery = await self._video_encoder.deliver_prepared_chunk( + payload, + self._video_track, + force_keyframe=force_keyframe, + ) with self._lock: if self._closed or generation < self._generation: stale_after_delivery = True @@ -844,6 +881,18 @@ def _on_done(self, future: Future[WebRTCChunkDelivery]) -> None: if self._on_chunk_delivery is not None: self._on_chunk_delivery(result) + def _pop_pending_locked(self) -> tuple[Future[WebRTCChunkDelivery], ...]: + pending = tuple(self._pending) + self._pending.clear() + return pending + + @staticmethod + def _cancel_stale_deliveries( + futures: Sequence[Future[WebRTCChunkDelivery]], + ) -> None: + for future in futures: + future.cancel() + def _track_backpressure_s(self) -> float: qsize = getattr(self._video_track, "qsize", None) fps = getattr(self._video_track, "fps", None) or getattr( @@ -862,6 +911,28 @@ def _track_backpressure_s(self) -> float: return 0.0 return max(0.0, queue_depth / frames_per_second) + async def _flush_full_track_queue(self, *, frame_count: int) -> None: + if frame_count <= 0: + return + qsize = getattr(self._video_track, "qsize", None) + flush = getattr(self._video_track, "flush", None) + if not callable(qsize) or not callable(flush): + return + try: + queue_depth = int(qsize()) + except (TypeError, ValueError): + return + if queue_depth < frame_count: + return + + # WebRTC is interactive: a full track queue means a whole generated + # chunk is stale relative to the latest input window. Drop that queued + # media before enqueueing the current chunk so visual latency stays + # bounded instead of preserving every frame. + result = flush() + if inspect.isawaitable(result): + await result + def _schedule_track_close(self) -> None: close = getattr(self._video_track, "close", None) if not callable(close): diff --git a/flashdreams/flashdreams/serving/webrtc/web/request_session.js b/flashdreams/flashdreams/serving/webrtc/web/request_session.js index a13384bc5..76b7507f9 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/request_session.js +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.js @@ -605,15 +605,6 @@ function enqueueAction(action) { } } -function enqueueHeldKeyRepeats() { - const heldKeys = Array.from(activeKeys).sort((a, b) => { - return (heldKeyOrder.get(a) || 0) - (heldKeyOrder.get(b) || 0) - }) - for (const key of heldKeys) { - enqueueAction({ event: "keydown", key }) - } -} - function setKeyHeld(key, source, held) { const normalized = normalizeKey(key) if (!allowedKeys.has(normalized)) { @@ -695,9 +686,6 @@ function handleControlMessage(rawMessage) { logEvent(parts.join(", ")) setStatus(activeKeys.size > 0 ? "Generating" : "Waiting", activeKeys.size > 0 ? "generating" : "waiting") setFlow(`chunk ${payload.chunk_index} complete`) - if (activeKeys.size > 0) { - enqueueHeldKeyRepeats() - } modelAdapter?.onControlMessage?.(payload, modelContext) return } diff --git a/flashdreams/tests/test_checkpoint_loading.py b/flashdreams/tests/test_checkpoint_loading.py new file mode 100644 index 000000000..f3b129bd1 --- /dev/null +++ b/flashdreams/tests/test_checkpoint_loading.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Checkpoint loading behavior tests.""" + +from __future__ import annotations + +import importlib +import json +from pathlib import Path +from typing import Any + +import pytest +import torch +from safetensors.torch import save_file as save_safetensors_file + +pytestmark = pytest.mark.ci_cpu + + +def test_local_safetensors_uses_file_backed_loader( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Load local safetensors without materializing the file as bytes.""" + checkpoint_load = importlib.import_module("flashdreams.core.checkpoint.load") + checkpoint_path = tmp_path / "weights.safetensors" + expected = {"weight": torch.ones(2)} + calls: list[tuple[str, str]] = [] + + def fake_load_file(path: str, *, device: str) -> dict[str, torch.Tensor]: + calls.append((path, device)) + return expected + + def reject_bytes_load(_data: bytes) -> dict[str, torch.Tensor]: + pytest.fail("safetensors checkpoints must use the file-backed loader") + + monkeypatch.setattr(checkpoint_load, "load_safetensors_file", fake_load_file) + monkeypatch.setattr(checkpoint_load, "load_safetensors", reject_bytes_load) + + actual = checkpoint_load.load_single_checkpoint( + str(checkpoint_path), + map_location=torch.device("cpu"), + ) + + assert actual is expected + assert calls == [(str(checkpoint_path), "cpu")] + + +def test_safetensors_model_load_streams_without_full_state_dict( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Stream safetensors tensors directly into a materialized model.""" + checkpoint_load = importlib.import_module("flashdreams.core.checkpoint.load") + checkpoint_path = tmp_path / "weights.safetensors" + expected = torch.arange(6, dtype=torch.float32).view(2, 3) + save_safetensors_file({"weight": expected}, checkpoint_path) + model = torch.nn.Linear(3, 2, bias=False) + + def reject_full_load(*_args: Any, **_kwargs: Any) -> None: + pytest.fail("model loads must not materialize the complete state dict") + + monkeypatch.setattr(checkpoint_load, "load_safetensors_file", reject_full_load) + + actual = checkpoint_load.load_checkpoint(str(checkpoint_path), model=model) + + assert actual is model + torch.testing.assert_close(model.weight, expected) + + +def test_sharded_safetensors_model_load_streams_without_merged_state_dict( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Stream indexed safetensors shards into a model without merging first.""" + checkpoint_load = importlib.import_module("flashdreams.core.checkpoint.load") + shard_a = tmp_path / "model-00001-of-00002.safetensors" + shard_b = tmp_path / "model-00002-of-00002.safetensors" + index_path = tmp_path / "model.safetensors.index.json" + expected_weight = torch.arange(6, dtype=torch.float32).view(2, 3) + expected_bias = torch.tensor([3.0, 4.0], dtype=torch.float32) + save_safetensors_file({"weight": expected_weight}, shard_a) + save_safetensors_file({"bias": expected_bias}, shard_b) + index_path.write_text( + json.dumps( + { + "metadata": {"total_size": 0}, + "weight_map": { + "weight": shard_a.name, + "bias": shard_b.name, + }, + } + ), + encoding="utf-8", + ) + model = torch.nn.Linear(3, 2) + + def reject_merge(*_args: Any, **_kwargs: Any) -> None: + pytest.fail("sharded model loads must not materialize a merged state dict") + + monkeypatch.setattr( + checkpoint_load, + "_load_sharded_safetensors_index_checkpoint", + reject_merge, + ) + + actual = checkpoint_load.load_checkpoint(str(index_path), model=model) + + assert actual is model + torch.testing.assert_close(model.weight, expected_weight) + torch.testing.assert_close(model.bias, expected_bias) diff --git a/flashdreams/tests/test_demo_runtime_timing.py b/flashdreams/tests/test_demo_runtime_timing.py index 98e3cf56d..e0aa9cd0d 100644 --- a/flashdreams/tests/test_demo_runtime_timing.py +++ b/flashdreams/tests/test_demo_runtime_timing.py @@ -11,14 +11,13 @@ from flashdreams.runtime import StepRequirements, UserInputSchema from flashdreams.runtime.demo import NullOutputSink, RunResult, SessionEdges from flashdreams.runtime.demo.timing import ( - SPARSE_KEY_SEGMENTS_METADATA_KEY, CatchUpDecision, CatchUpPolicy, - KeyboardRealtimeInputSource, + RealtimeEventInputSource, + RealtimeEventResampler, ResamplerRealtimeClock, SignalActivationPolicy, ) -from flashdreams.serving.realtime.input import KeyboardResampler pytestmark = pytest.mark.ci_cpu @@ -26,7 +25,7 @@ @pytest.mark.asyncio async def test_signal_activation_waits_for_first_input_and_anchors_clock() -> None: event = asyncio.Event() - resampler = KeyboardResampler(fps=30.0, start_v=0.0) + resampler = RealtimeEventResampler(fps=30.0, start_v=0.0) clock = ResamplerRealtimeClock(resampler=resampler, now_fn=lambda: 12.0) policy = SignalActivationPolicy(signals=(event,), timeout_s=1.0) @@ -46,7 +45,7 @@ async def test_signal_activation_waits_for_first_input_and_anchors_clock() -> No @pytest.mark.asyncio async def test_activation_timeout_can_close_edges_as_not_activated() -> None: event = asyncio.Event() - resampler = KeyboardResampler(fps=30.0, start_v=0.0) + resampler = RealtimeEventResampler(fps=30.0, start_v=0.0) clock = ResamplerRealtimeClock(resampler=resampler, now_fn=lambda: 12.0) policy = SignalActivationPolicy( signals=(event,), @@ -77,7 +76,7 @@ async def test_activation_timeout_can_close_edges_as_not_activated() -> None: def test_resampler_clock_catch_up_bounds_latency() -> None: - resampler = KeyboardResampler(fps=1.0, start_v=0.0) + resampler = RealtimeEventResampler(fps=1.0, start_v=0.0) clock = ResamplerRealtimeClock(resampler=resampler, now_fn=lambda: 5.0) decision = clock.catch_up( @@ -96,18 +95,15 @@ def test_resampler_clock_catch_up_bounds_latency() -> None: @pytest.mark.asyncio -async def test_realtime_input_source_matches_resampler_for_recorded_trace() -> None: - expected_resampler = _resampler_with_recorded_trace() - expected_resampler.next_chunk_start_v = 2.0 - expected_segments, expected_frame_times = expected_resampler.sample_chunk(2) - source_resampler = _resampler_with_recorded_trace() +async def test_realtime_input_source_emits_transport_neutral_window() -> None: + source_resampler = RealtimeEventResampler(fps=2.0, start_v=0.0) sleep = _RecordingSleep() clock = ResamplerRealtimeClock( resampler=source_resampler, now_fn=lambda: 3.0, sleep_fn=sleep, ) - source = KeyboardRealtimeInputSource(resampler=source_resampler) + source = RealtimeEventInputSource(resampler=source_resampler) result = await source.next_realtime_window( request=_request(input_frame_count=2), @@ -123,15 +119,14 @@ async def test_realtime_input_source_matches_resampler_for_recorded_trace() -> N ) assert result.window.start_s == pytest.approx(2.0) assert result.window.end_s == pytest.approx(3.0) - assert result.window.frame_times == tuple(expected_frame_times) - assert result.window.metadata[SPARSE_KEY_SEGMENTS_METADATA_KEY] == tuple( - expected_segments - ) + assert result.window.frame_times == pytest.approx((2.5, 3.0)) + assert result.window.inputs.events == () + assert result.window.metadata == {} @pytest.mark.asyncio async def test_backpressure_is_clock_adjustment_not_blocking_sleep() -> None: - resampler = KeyboardResampler(fps=1.0, start_v=0.0) + resampler = RealtimeEventResampler(fps=1.0, start_v=0.0) sleep = _RecordingSleep() clock = ResamplerRealtimeClock( resampler=resampler, @@ -156,7 +151,7 @@ async def test_backpressure_is_clock_adjustment_not_blocking_sleep() -> None: @pytest.mark.asyncio async def test_window_floor_sleeps_only_when_virtual_time_is_ahead() -> None: - resampler = KeyboardResampler(fps=1.0, start_v=0.0) + resampler = RealtimeEventResampler(fps=1.0, start_v=0.0) sleep = _RecordingSleep() clock = ResamplerRealtimeClock( resampler=resampler, @@ -171,22 +166,22 @@ async def test_window_floor_sleeps_only_when_virtual_time_is_ahead() -> None: @pytest.mark.parametrize("policy", ["drop", "compress"]) -def test_keyboard_resampler_defers_unsupported_catch_up_policies( +def test_realtime_event_source_defers_unsupported_catch_up_policies( policy: str, ) -> None: - resampler = KeyboardResampler(fps=1.0, start_v=0.0) + resampler = RealtimeEventResampler(fps=1.0, start_v=0.0) clock = ResamplerRealtimeClock(resampler=resampler, now_fn=lambda: 5.0) unsupported_policy = cast(CatchUpPolicy, policy) - with pytest.raises(NotImplementedError, match="no existing resampler analog"): + with pytest.raises(NotImplementedError, match="no existing timeline analog"): clock.catch_up( request=_request(input_frame_count=1), max_lag_s=1.0, policy=unsupported_policy, ) - with pytest.raises(NotImplementedError, match="KeyboardResampler analog"): - KeyboardRealtimeInputSource( + with pytest.raises(NotImplementedError, match="event-window analog"): + RealtimeEventInputSource( resampler=resampler, catch_up_policy=unsupported_policy, ) @@ -199,18 +194,6 @@ def _request(*, input_frame_count: int) -> StepRequirements: ) -def _resampler_with_recorded_trace() -> KeyboardResampler: - resampler = KeyboardResampler(fps=2.0, start_v=0.0) - for arrival_t, event, key in ( - (0.25, "keydown", "w"), - (1.25, "keydown", "a"), - (2.25, "keyup", "w"), - (2.75, "keydown", "d"), - ): - resampler.on_edge(arrival_t=arrival_t, event=event, key=key) - return resampler - - class _RecordingSleep: def __init__(self) -> None: self.delays: list[float] = [] diff --git a/flashdreams/tests/test_realtime_serving.py b/flashdreams/tests/test_realtime_serving.py index 33b21476a..f52bf3d29 100644 --- a/flashdreams/tests/test_realtime_serving.py +++ b/flashdreams/tests/test_realtime_serving.py @@ -9,14 +9,14 @@ import pytest import torch -from flashdreams.serving.realtime.frame_bus import LatestFrameBus -from flashdreams.serving.realtime.input import ( +from flashdreams.runtime.keyboard import ( ImageRequest, KeyboardState, PromptRequest, ResetRequest, SparseInputSnapshot, ) +from flashdreams.serving.realtime.frame_bus import LatestFrameBus from flashdreams.serving.realtime.media import ( encode_rgb_frame_to_jpeg, rgb_array_to_uint8_frames, diff --git a/flashdreams/tests/test_runtime_demo_api.py b/flashdreams/tests/test_runtime_demo_api.py index e40b62b51..9600f19d0 100644 --- a/flashdreams/tests/test_runtime_demo_api.py +++ b/flashdreams/tests/test_runtime_demo_api.py @@ -52,7 +52,7 @@ run_replay_demo, ) from flashdreams.runtime.demo.app import DemoApplication -from flashdreams.runtime.demo.webrtc import ( +from flashdreams.serving.webrtc.demo import ( serve_webrtc_demo, ) from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager diff --git a/flashdreams/tests/test_webrtc_manager.py b/flashdreams/tests/test_webrtc_manager.py index 787a8ee7b..bafe2b1a8 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -20,9 +20,8 @@ UserInputs, ) from flashdreams.runtime.demo import RunResult -from flashdreams.runtime.demo.timing import SPARSE_KEY_SEGMENTS_METADATA_KEY +from flashdreams.runtime.keyboard import WSAD_SUPPORTED_KEYS 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 from flashdreams.serving.webrtc.manager import ( BaseWebRTCSessionManager, @@ -145,41 +144,50 @@ 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] + def reset(self, *, start_v: float) -> None: + self.next_chunk_start_v = start_v + def sample_chunk(self, num_frames: int) -> list[float]: + start = self.next_chunk_start_v + frame_times = [start + index * self.dt for index in range(num_frames)] + self.next_chunk_start_v = start + num_frames * self.dt + return frame_times + + +class _RecordingLegacySegmentResampler: + dt = 0.0 + next_chunk_start_v = 0.0 -class _RecordingResampler(_FakeResampler): def __init__(self) -> None: self.edges: list[tuple[float, str, str]] = [] + def reset(self, *, start_v: float) -> None: + self.next_chunk_start_v = start_v + self.edges.clear() + def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: self.edges.append((arrival_t, event, key)) + 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 _SharedResampler: def __init__(self, *, start_v: float = 0.0, dt: float = 0.001) -> None: self.next_chunk_start_v = start_v self.dt = dt - self.edges: list[tuple[float, str, str]] = [] def reset(self, *, start_v: float) -> None: self.next_chunk_start_v = start_v - self.edges.clear() - def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: - self.edges.append((arrival_t, event, key)) - - def sample_chunk( - self, num_frames: int - ) -> tuple[list[tuple[float, float, frozenset[str]]], list[float]]: + def sample_chunk(self, num_frames: int) -> list[float]: start = self.next_chunk_start_v end = start + num_frames * self.dt self.next_chunk_start_v = end - return [(start, end, frozenset({"w"}))], [end] + return [end] class _CountingVideoTrack(_FakeVideoTrack): @@ -553,7 +561,9 @@ def map_step_inputs( frame_times=(2.25, 2.75), inputs=current_inputs, metadata={ - SPARSE_KEY_SEGMENTS_METADATA_KEY: ((2.0, 3.0, frozenset({"w"})),), + manager_module._LEGACY_SPARSE_KEY_SEGMENTS_METADATA_KEY: ( + (2.0, 3.0, frozenset({"w"})), + ), WEBRTC_SKIPPED_INPUTS_METADATA_KEY: skipped_inputs, WEBRTC_SKIPPED_WINDOW_METADATA_KEY: (0.0, 2.0), }, @@ -566,9 +576,9 @@ def map_step_inputs( assert mapping.inference_inputs[0].metadata["frame_times"] == (2.25, 2.75) assert mapping.inference_inputs[0].metadata["window_start_s"] == 2.0 assert mapping.inference_inputs[0].metadata["window_end_s"] == 3.0 - assert mapping.inference_inputs[0].metadata[SPARSE_KEY_SEGMENTS_METADATA_KEY] == ( - (2.0, 3.0, frozenset({"w"})), - ) + assert mapping.inference_inputs[0].metadata[ + manager_module._LEGACY_SPARSE_KEY_SEGMENTS_METADATA_KEY + ] == ((2.0, 3.0, frozenset({"w"})),) @pytest.mark.asyncio @@ -615,8 +625,8 @@ async def test_action_keyup_updates_state_when_user_event_queue_full( managed, _video_track, _peer, channel = _managed_session(runtime) managed.inference_session = object() managed.first_action_received.clear() - resampler = _RecordingResampler() - managed.resampler = resampler # ty:ignore[invalid-assignment] + resampler = _RecordingLegacySegmentResampler() + managed.legacy_segment_resampler = resampler manager._record_user_event( managed_session=managed, timestamp_s=0.0, @@ -904,7 +914,7 @@ async def step( managed, _video_track, _peer, channel = _managed_session(runtime) resampler = _SplitResampler() managed.video_track = _CountingVideoTrack() # ty:ignore[invalid-assignment] - managed.resampler = resampler # ty:ignore[invalid-assignment] + managed.legacy_segment_resampler = resampler runtime.managed_session = managed manager._active_session = managed @@ -1056,7 +1066,13 @@ def peek_steady_output_num_frames(self) -> int: reservation = context.admission.try_reserve() assert reservation is not None resampler = _SharedResampler(start_v=asyncio.get_running_loop().time()) - input_source = WebRTCInputSource(resampler=resampler) + input_source = WebRTCInputSource( + resampler=resampler, + legacy_segment_resampler=_RecordingLegacySegmentResampler(), + legacy_segments_metadata_key=( + manager_module._LEGACY_SPARSE_KEY_SEGMENTS_METADATA_KEY + ), + ) input_source.handle_browser_payload( {"type": "action", "action": {"event": "step"}}, timestamp_s=asyncio.get_running_loop().time(), @@ -1146,24 +1162,17 @@ async def test_create_answer_raises_busy_with_subclass_message() -> None: await manager.create_answer(offer_sdp="x", offer_type="offer") -def test_make_resampler_honors_supported_keys() -> None: - wsad = _make_manager( +def test_supported_key_payload_honors_configured_keys() -> None: + wsad_manager = _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. - assert wsad_segments[0][2] == frozenset() + ) + assert not wsad_manager._supports_key_payload({"key": "q"}) + assert wsad_manager._supports_key_payload({"key": "ArrowUp"}) - default = _make_manager( - _BaseTestManager, runtime=SimpleNamespace() - )._make_resampler(start_v=1.0) - default.on_edge(arrival_t=0.5, event="keydown", key="q") - default_segments, _ = default.sample_chunk(num_frames=1) - # The default key set (used by Lingbot) accepts 'q'. - assert default_segments[0][2] == frozenset({"q"}) + default_manager = _make_manager(_BaseTestManager, runtime=SimpleNamespace()) + assert default_manager._supports_key_payload({"key": "q"}) @pytest.mark.asyncio @@ -1172,8 +1181,8 @@ async def test_step_action_starts_generation_without_key_edge() -> None: manager = _make_manager(_BaseTestManager, runtime) managed, _video_track, _peer, _channel = _managed_session(runtime) managed.first_action_received.clear() - resampler = _RecordingResampler() - managed.resampler = resampler # ty:ignore[invalid-assignment] + resampler = _RecordingLegacySegmentResampler() + managed.legacy_segment_resampler = resampler await manager._handle_datachannel_message( managed_session=managed, diff --git a/flashdreams/tests/test_webrtc_services.py b/flashdreams/tests/test_webrtc_services.py index bda6eee4c..4c1a6138e 100644 --- a/flashdreams/tests/test_webrtc_services.py +++ b/flashdreams/tests/test_webrtc_services.py @@ -159,9 +159,10 @@ async def test_webrtc_input_source_emits_typed_user_inputs() -> None: ) assert source.activation_signal.is_set() - assert resampler.edges == [(0.05, "keydown", "w")] assert result.window.start_s == pytest.approx(0.0) assert result.window.end_s == pytest.approx(0.2) + assert result.window.frame_times == pytest.approx((0.1, 0.2)) + assert result.window.metadata == {} assert [event.event_type for event in result.window.inputs.events] == [ "key_down", "text_event", @@ -236,25 +237,55 @@ async def test_webrtc_output_bridge_prepares_payload_before_async_delivery() -> @pytest.mark.asyncio -async def test_webrtc_output_bridge_drops_full_queue_before_payload_prepare() -> None: +async def test_webrtc_output_bridge_replaces_full_pending_queue() -> None: loop = asyncio.get_running_loop() encoder = _BlockingEncoder() + track = _FakeVideoTrack() bridge = ThreadSafeWebRTCOutputBridge( loop=loop, video_encoder=encoder, - video_track=_FakeVideoTrack(), + video_track=track, max_pending_chunks=1, ) sink = WebRTCOutputSink(bridge=bridge) sink.open(SessionInfo()) first = sink.write(StepResult(step_index=0, frame_count=1)) + await asyncio.wait_for(encoder.started.wait(), timeout=1.0) second = sink.write(StepResult(step_index=1, frame_count=1)) assert not first.dropped - assert second.dropped - assert second.drop_policy == "drop_newest" - assert encoder.prepared_payloads == [0] + assert not second.dropped + assert encoder.prepared_payloads == [0, 1] + for _ in range(10): + if track.flush_count: + break + await asyncio.sleep(0) + assert track.flush_count == 1 + + encoder.release.set() + await asyncio.wait_for(encoder.done.wait(), timeout=1.0) + sink.close() + + +@pytest.mark.asyncio +async def test_webrtc_output_bridge_flushes_full_track_queue_before_delivery() -> None: + loop = asyncio.get_running_loop() + encoder = _BlockingEncoder() + track = _FakeVideoTrack(queue_depth=2) + bridge = ThreadSafeWebRTCOutputBridge( + loop=loop, + video_encoder=encoder, + video_track=track, + ) + sink = WebRTCOutputSink(bridge=bridge) + sink.open(SessionInfo()) + + decision = sink.write(StepResult(step_index=0, frame_count=2)) + + assert not decision.dropped + await asyncio.wait_for(encoder.started.wait(), timeout=1.0) + assert track.flush_count == 1 encoder.release.set() await asyncio.wait_for(encoder.done.wait(), timeout=1.0) @@ -671,24 +702,21 @@ class _FakeResampler: def __init__(self, *, dt: float, start_v: float) -> None: self.dt = dt self.next_chunk_start_v = start_v - self.edges: list[tuple[float, str, str]] = [] def reset(self, *, start_v: float) -> None: self.next_chunk_start_v = start_v - self.edges.clear() - - def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: - self.edges.append((arrival_t, event, key)) def sample_chunk( self, num_frames: int, - ) -> tuple[tuple[tuple[float, float, frozenset[str]], ...], tuple[float, ...]]: + ) -> tuple[float, ...]: start = self.next_chunk_start_v - frame_times = tuple(start + index * self.dt for index in range(num_frames)) + frame_times = tuple( + start + (index + 1) * self.dt for index in range(num_frames) + ) end = start + num_frames * self.dt self.next_chunk_start_v = end - return (((start, end, frozenset({"w"})),), frame_times) + return frame_times class _BlockingEncoder: @@ -741,11 +769,13 @@ async def deliver_chunk( class _FakeVideoTrack: fps = 30 - def __init__(self) -> None: + def __init__(self, *, queue_depth: int = 0) -> None: + self.queue_depth = queue_depth self.flush_count = 0 def qsize(self) -> int: - return 0 + return self.queue_depth async def flush(self) -> None: self.flush_count += 1 + self.queue_depth = 0 diff --git a/flashdreams/tests/test_webrtc_serving.py b/flashdreams/tests/test_webrtc_serving.py index 06e003496..71b1f48e8 100644 --- a/flashdreams/tests/test_webrtc_serving.py +++ b/flashdreams/tests/test_webrtc_serving.py @@ -14,12 +14,8 @@ from aiohttp import web from aiohttp.test_utils import TestClient, TestServer -from flashdreams.serving.webrtc.controls import ( - WSAD_SUPPORTED_KEYS, - CameraPoseIntegrator, - KeyboardResampler, - KeyboardState, -) +from flashdreams.runtime.demo import RealtimeEventResampler +from flashdreams.runtime.keyboard import WSAD_SUPPORTED_KEYS, KeyboardState from flashdreams.serving.webrtc.manager import ( BaseWebRTCSessionManager, ManagedWebRTCSession, @@ -48,49 +44,6 @@ def test_wsad_keyboard_state_rejects_non_driving_keys() -> None: assert state.resolved_effective_keys() == frozenset({"w"}) -def test_wsad_resampler_preserves_held_key() -> None: - resampler = KeyboardResampler( - fps=30, - start_v=1.0, - supported_keys=WSAD_SUPPORTED_KEYS, - ) - resampler.on_edge(arrival_t=0.5, event="keydown", key="w") - - segments, frame_times = resampler.sample_chunk(num_frames=2) - - assert segments == [(1.0, 1.0 + 2 / 30, frozenset({"w"}))] - assert frame_times == pytest.approx([1.0 + 1 / 30, 1.0 + 2 / 30]) - - -def test_camera_pose_integrator_flu_uses_driving_axes() -> None: - integrator = CameraPoseIntegrator( - move_speed_per_s=2.0, - rotate_speed_rad_per_s=float(np.pi / 2), - coordinate_system="FLU", - ) - - integrator.reset() - poses = integrator.integrate_chunk( - segments=[(0.0, 1.0, frozenset({"w"}))], - frame_times=[1.0], - ) - assert poses[-1][:3, 3] == pytest.approx([2.0, 0.0, 0.0]) - - integrator.reset() - poses = integrator.integrate_chunk( - segments=[(0.0, 1.0, frozenset({"a"}))], - frame_times=[1.0], - ) - assert poses[-1][:3, 0] == pytest.approx([0.0, 1.0, 0.0], abs=1e-6) - - integrator.reset() - poses = integrator.integrate_chunk( - segments=[(0.0, 1.0, frozenset({"d"}))], - frame_times=[1.0], - ) - assert poses[-1][:3, 0] == pytest.approx([0.0, -1.0, 0.0], abs=1e-6) - - def test_tensor_chunk_to_rgb_frames_supports_omnidreams_layout() -> None: chunk = torch.zeros((1, 1, 2, 3, 4, 5), dtype=torch.uint8) chunk[0, 0, 1, 0] = 255 @@ -154,7 +107,7 @@ def _managed_session_with_channel( video_track=_FakeCloseable(), # ty:ignore[invalid-argument-type] video_encoder=_FakeCloseable(), # ty:ignore[invalid-argument-type] peer_connection=_FakeCloseable(), - resampler=KeyboardResampler(fps=30, start_v=0.0), + resampler=RealtimeEventResampler(fps=30, start_v=0.0), control_channel=channel, ) return managed_session, channel diff --git a/integrations/lingbot/README.md b/integrations/lingbot/README.md index e0029e48a..71c71f91c 100644 --- a/integrations/lingbot/README.md +++ b/integrations/lingbot/README.md @@ -123,6 +123,20 @@ uv run --python 3.12 --package flashdreams-lingbot lingbot-demo replay \ --output outputs/lingbot-demo-replay.mp4 ``` +Run the same replay path without writing video: + +```bash +uv run --python 3.12 --package flashdreams-lingbot lingbot-demo replay \ + --device cuda:0 \ + --preset-id lingbot-world-v2-14b-causal-fast-taehv-window15-sink3 \ + --example-idx 0 \ + --total-blocks 10 \ + --fps 16 \ + --pixel-height 352 \ + --pixel-width 640 \ + --output-mode null +``` + Serve the shared WebRTC demo: ```bash diff --git a/integrations/lingbot/lingbot/config.py b/integrations/lingbot/lingbot/config.py index 72c647f56..ef8530d40 100644 --- a/integrations/lingbot/lingbot/config.py +++ b/integrations/lingbot/lingbot/config.py @@ -76,6 +76,7 @@ in_dim=16 + 4 + 16, ), checkpoint_path=LINGBOT_WORLD_V1_CHECKPOINT_PATH, + stream_checkpoint=True, # Single-rollout layout: tensors flow through the stack as # ``[T, C, H, W]`` (or ``[T, ...]``) with no leading batch/view dim. batch_shape=(), @@ -136,7 +137,8 @@ ) # LingBot-World v2 uses the same architecture and runtime as v1. The -# transformer checkpoint is the only model-level substitution. +# transformer checkpoint is the only model-level substitution; it inherits +# the bounded checkpoint loader from the v1 base config. PIPELINE_LINGBOT_WORLD_V2_14B_CAUSAL_FAST = derive_config( PIPELINE_LINGBOT_WORLD_FAST, name="lingbot-world-v2-14b-causal-fast", diff --git a/integrations/lingbot/lingbot/controls.py b/integrations/lingbot/lingbot/controls.py new file mode 100644 index 000000000..da7cbf16c --- /dev/null +++ b/integrations/lingbot/lingbot/controls.py @@ -0,0 +1,283 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lingbot-owned keyboard segmenting and camera pose integration.""" + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass, field +from typing import Literal + +import numpy as np + +from flashdreams.runtime.keyboard import DEFAULT_SUPPORTED_KEYS, KeyboardState + +PoseSegment = tuple[float, float, frozenset[str]] + + +class KeyboardResampler: + """Resample sparse keydown/keyup edges into a Lingbot camera timeline.""" + + def __init__( + self, + *, + fps: float, + start_v: float = 0.0, + supported_keys: frozenset[str] = DEFAULT_SUPPORTED_KEYS, + ) -> None: + if fps <= 0: + raise ValueError("fps must be > 0") + self._fps = float(fps) + self._dt = 1.0 / self._fps + self._supported_keys = supported_keys + self.next_chunk_start_v = start_v + self._event_log: deque[tuple[float, dict[str, str]]] = deque() + self._carried_state = KeyboardState(supported_keys=supported_keys) + + @property + def fps(self) -> float: + return self._fps + + @property + def dt(self) -> float: + return self._dt + + def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: + entry = (arrival_t, {"event": event, "key": key}) + if not self._event_log or arrival_t >= self._event_log[-1][0]: + self._event_log.append(entry) + return + for index, (event_t, _) in enumerate(self._event_log): + if arrival_t < event_t: + self._event_log.insert(index, entry) + return + self._event_log.append(entry) + + def sample_chunk(self, num_frames: int) -> tuple[list[PoseSegment], list[float]]: + if num_frames < 1: + raise ValueError("num_frames must be >= 1") + + chunk_start_v = self.next_chunk_start_v + chunk_end_v = chunk_start_v + num_frames * self._dt + + while self._event_log and self._event_log[0][0] < chunk_start_v: + _, payload = self._event_log.popleft() + self._carried_state.apply_event(**payload) + + segments: list[PoseSegment] = [] + prev_t = chunk_start_v + prev_state = self._carried_state.resolved_effective_keys() + while self._event_log and self._event_log[0][0] <= chunk_end_v: + event_t, payload = self._event_log.popleft() + if event_t > prev_t: + segments.append((prev_t, event_t, prev_state)) + self._carried_state.apply_event(**payload) + prev_state = self._carried_state.resolved_effective_keys() + prev_t = event_t + if prev_t < chunk_end_v: + segments.append((prev_t, chunk_end_v, prev_state)) + elif not segments: + segments.append((chunk_start_v, chunk_end_v, prev_state)) + + frame_times = [chunk_start_v + (i + 1) * self._dt for i in range(num_frames)] + self.next_chunk_start_v = chunk_end_v + return segments, frame_times + + def reset(self, *, start_v: float) -> None: + self._event_log.clear() + self._carried_state = KeyboardState(supported_keys=self._supported_keys) + self.next_chunk_start_v = start_v + + def event_log_size(self) -> int: + return len(self._event_log) + + +def _rotation_matrix(axis: str, angle_rad: float) -> np.ndarray: + cos_t = np.float32(np.cos(angle_rad)) + sin_t = np.float32(np.sin(angle_rad)) + if axis == "x": + return np.array( + [ + [1.0, 0.0, 0.0], + [0.0, cos_t, -sin_t], + [0.0, sin_t, cos_t], + ], + dtype=np.float32, + ) + if axis == "y": + return np.array( + [ + [cos_t, 0.0, sin_t], + [0.0, 1.0, 0.0], + [-sin_t, 0.0, cos_t], + ], + dtype=np.float32, + ) + if axis == "z": + return np.array( + [ + [cos_t, -sin_t, 0.0], + [sin_t, cos_t, 0.0], + [0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + return np.eye(3, dtype=np.float32) + + +@dataclass(slots=True) +class CameraPoseIntegrator: + """Integrate a piecewise-constant keyboard timeline into a camera path.""" + + move_speed_per_s: float = 0.8 + rotate_speed_rad_per_s: float = float(np.deg2rad(32.0)) + pitch_limit_rad: float = float(np.deg2rad(85.0)) + coordinate_system: Literal["RDF", "FLU"] = "RDF" + _current_pose: np.ndarray = field( + default_factory=lambda: np.eye(4, dtype=np.float32), + ) + _current_pitch: float = 0.0 + + def __post_init__(self) -> None: + if self.coordinate_system not in {"RDF", "FLU"}: + raise ValueError( + "coordinate_system must be 'RDF' (right-down-forward) " + "or 'FLU' (forward-left-up)" + ) + + def reset(self, pose: np.ndarray | None = None) -> None: + if pose is None: + self._current_pose = np.eye(4, dtype=np.float32) + self._current_pitch = 0.0 + return + if pose.shape != (4, 4): + raise ValueError(f"Expected pose shape (4, 4), got {pose.shape}") + self._current_pose = pose.astype(np.float32, copy=True) + if self.coordinate_system == "FLU": + self._current_pitch = float(np.arcsin(np.clip(pose[2, 0], -1.0, 1.0))) + else: + self._current_pitch = float(np.arctan2(pose[2, 1], pose[1, 1])) + + def current_pose(self) -> np.ndarray: + return self._current_pose.copy() + + def _advance(self, *, state: frozenset[str], duration: float) -> None: + if duration <= 0: + return + + yaw_rate = 0.0 + if self.coordinate_system == "FLU": + if "a" in state or "j" in state: + yaw_rate += self.rotate_speed_rad_per_s + if "d" in state or "l" in state: + yaw_rate -= self.rotate_speed_rad_per_s + else: + if "a" in state or "j" in state: + yaw_rate -= self.rotate_speed_rad_per_s + if "d" in state or "l" in state: + yaw_rate += self.rotate_speed_rad_per_s + pitch_rate = 0.0 + if "i" in state: + pitch_rate += self.rotate_speed_rad_per_s + if "k" in state: + pitch_rate -= self.rotate_speed_rad_per_s + + yaw_delta = yaw_rate * duration + pitch_delta = pitch_rate * duration + + new_pitch = self._current_pitch + pitch_delta + if -self.pitch_limit_rad <= new_pitch <= self.pitch_limit_rad: + self._current_pitch = new_pitch + else: + pitch_delta = 0.0 + + rot = self._current_pose[:3, :3] + trans = self._current_pose[:3, 3] + if self.coordinate_system == "FLU": + rot_pitch = _rotation_matrix("y", -pitch_delta) + rot_yaw = _rotation_matrix("z", yaw_delta) + else: + rot_pitch = _rotation_matrix("x", pitch_delta) + rot_yaw = _rotation_matrix("y", yaw_delta) + rot_new = rot_yaw @ rot @ rot_pitch + + forward_rate = 0.0 + if "w" in state: + forward_rate += self.move_speed_per_s + if "s" in state: + forward_rate -= self.move_speed_per_s + right_rate = 0.0 + if "e" in state: + right_rate += self.move_speed_per_s + if "q" in state: + right_rate -= self.move_speed_per_s + + if self.coordinate_system == "FLU": + vec_forward = rot_new[:, 0] + vec_right = -rot_new[:, 1] + forward_flat = np.array( + [vec_forward[0], vec_forward[1], 0.0], dtype=np.float32 + ) + right_flat = np.array([vec_right[0], vec_right[1], 0.0], dtype=np.float32) + else: + vec_right = rot_new[:, 0] + vec_forward = rot_new[:, 2] + forward_flat = np.array( + [vec_forward[0], 0.0, vec_forward[2]], dtype=np.float32 + ) + right_flat = np.array([vec_right[0], 0.0, vec_right[2]], dtype=np.float32) + forward_norm = np.linalg.norm(forward_flat) + right_norm = np.linalg.norm(right_flat) + if forward_norm > 0: + forward_flat /= forward_norm + if right_norm > 0: + right_flat /= right_norm + + move_vec = forward_flat * (forward_rate * duration) + right_flat * ( + right_rate * duration + ) + self._current_pose = np.eye(4, dtype=np.float32) + self._current_pose[:3, :3] = rot_new + self._current_pose[:3, 3] = trans + move_vec + + def integrate_chunk( + self, + *, + segments: list[PoseSegment], + frame_times: list[float], + ) -> np.ndarray: + if not segments: + raise ValueError("segments must be non-empty") + if not frame_times: + raise ValueError("frame_times must be non-empty") + chunk_start = segments[0][0] + chunk_end = segments[-1][1] + if any( + frame_times[i] >= frame_times[i + 1] for i in range(len(frame_times) - 1) + ): + raise ValueError("frame_times must be strictly increasing") + if frame_times[0] < chunk_start - 1e-9 or frame_times[-1] > chunk_end + 1e-9: + raise ValueError( + "frame_times must lie within the chunk window " + f"[{chunk_start}, {chunk_end}]" + ) + + poses: list[np.ndarray] = [] + cur_t = chunk_start + ft_idx = 0 + for _, seg_end, seg_state in segments: + while ft_idx < len(frame_times) and frame_times[ft_idx] <= seg_end: + target_t = frame_times[ft_idx] + self._advance(state=seg_state, duration=target_t - cur_t) + cur_t = target_t + poses.append(self._current_pose.copy()) + ft_idx += 1 + if seg_end > cur_t: + self._advance(state=seg_state, duration=seg_end - cur_t) + cur_t = seg_end + + return np.stack(poses, axis=0).astype(np.float32) + + +__all__ = ["CameraPoseIntegrator", "KeyboardResampler", "PoseSegment"] diff --git a/integrations/lingbot/lingbot/demo/adapter.py b/integrations/lingbot/lingbot/demo/adapter.py index 1424f3eb0..fedd0a2c2 100644 --- a/integrations/lingbot/lingbot/demo/adapter.py +++ b/integrations/lingbot/lingbot/demo/adapter.py @@ -17,6 +17,7 @@ from flashdreams.runtime.demo import ( DemoSpec, Mp4OutputSpec, + NullOutputSpec, PreparedScenario, WebRTCOutputSpec, ) @@ -65,12 +66,12 @@ def supported_input_modes(self) -> tuple[str, ...]: return ("replay", "keyboard-driving") def supported_output_modes(self) -> tuple[str, ...]: - return ("mp4", "webrtc") + return ("mp4", "null", "webrtc") def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: if spec.input_mode == "replay": - if not isinstance(spec.output, Mp4OutputSpec): - raise ValueError("Lingbot replay demo currently requires MP4 output.") + if not isinstance(spec.output, (Mp4OutputSpec, NullOutputSpec)): + raise ValueError("Lingbot replay demo requires MP4 or null output.") scenario = spec.scenario live_camera = False elif spec.input_mode == "keyboard-driving": diff --git a/integrations/lingbot/lingbot/demo/app.py b/integrations/lingbot/lingbot/demo/app.py index df6b9c703..8aa495dff 100644 --- a/integrations/lingbot/lingbot/demo/app.py +++ b/integrations/lingbot/lingbot/demo/app.py @@ -13,6 +13,7 @@ from flashdreams.runtime.demo import ( DemoSpec, Mp4OutputSpec, + NullOutputSpec, WebRTCOutputSpec, ) from flashdreams.runtime.demo.app import DemoApplication @@ -48,7 +49,7 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: ) subparsers = parser.add_subparsers(dest="command", required=True) - replay = subparsers.add_parser("replay", help="Run an MP4 replay demo.") + replay = subparsers.add_parser("replay", help="Run a finite replay demo.") replay.add_argument("--preset-id", "--config-name", default=DEFAULT_LINGBOT_PRESET) replay.add_argument("--device", default="cuda") replay.add_argument("--prompt", default=None) @@ -81,7 +82,8 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: replay.add_argument("--pixel-height", type=int, default=DEFAULT_PIXEL_HEIGHT) replay.add_argument("--pixel-width", type=int, default=DEFAULT_PIXEL_WIDTH) replay.add_argument("--fps", type=int, default=DEFAULT_FPS) - replay.add_argument("--output", type=Path, required=True) + replay.add_argument("--output-mode", choices=("mp4", "null"), default="mp4") + replay.add_argument("--output", type=Path, default=None) webrtc = subparsers.add_parser("webrtc", help="Serve a WebRTC driving demo.") webrtc.add_argument("--preset-id", "--config-name", default=DEFAULT_LINGBOT_PRESET) @@ -109,7 +111,13 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: default=0, choices=EXAMPLE_DATA_AVAILABLE_IDXS, ) - return parser.parse_args(argv) + args = parser.parse_args(argv) + if args.command == "replay": + if args.output_mode == "mp4" and args.output is None: + parser.error("replay --output is required when --output-mode=mp4.") + if args.output_mode == "null" and args.output is not None: + parser.error("replay --output is only valid when --output-mode=mp4.") + return args class LingbotDemoApplication(DemoApplication): @@ -176,11 +184,7 @@ def _replay_spec(args: argparse.Namespace) -> DemoSpec: preset_id=args.preset_id, input_mode="replay", scenario=scenario, - output=Mp4OutputSpec( - path=args.output, - fps=args.fps, - output_layout="tchw", - ), + output=_replay_output_spec(args), config=InferenceConfig( model_id=LINGBOT_MODEL_ID, preset_id=args.preset_id, @@ -189,6 +193,20 @@ def _replay_spec(args: argparse.Namespace) -> DemoSpec: ) +def _replay_output_spec(args: argparse.Namespace) -> Mp4OutputSpec | NullOutputSpec: + if args.output_mode == "mp4": + if args.output is None: + raise ValueError("Lingbot MP4 replay requires --output.") + return Mp4OutputSpec( + path=args.output, + fps=args.fps, + output_layout="tchw", + ) + if args.output_mode == "null": + return NullOutputSpec() + raise ValueError(f"Unsupported Lingbot replay output mode: {args.output_mode!r}.") + + def _webrtc_spec( args: argparse.Namespace, *, diff --git a/integrations/lingbot/lingbot/demo/webrtc.py b/integrations/lingbot/lingbot/demo/webrtc.py index 854326905..99ba80874 100644 --- a/integrations/lingbot/lingbot/demo/webrtc.py +++ b/integrations/lingbot/lingbot/demo/webrtc.py @@ -12,12 +12,12 @@ from flashdreams.runtime import InferenceConfig from flashdreams.runtime.demo import DemoSpec, WebRTCAppResources, WebRTCOutputSpec -from flashdreams.runtime.demo.webrtc import ( +from flashdreams.serving.webrtc.bootstrap import run_webrtc_server +from flashdreams.serving.webrtc.demo 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.demo import LingbotDemoAdapter diff --git a/integrations/lingbot/lingbot/input_mapping.py b/integrations/lingbot/lingbot/input_mapping.py index 245b0a112..2963e2895 100644 --- a/integrations/lingbot/lingbot/input_mapping.py +++ b/integrations/lingbot/lingbot/input_mapping.py @@ -43,12 +43,11 @@ UserInputCapability, UserInputs, ) +from flashdreams.runtime.keyboard import DEFAULT_SUPPORTED_KEYS, KeyboardState 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 ( +from lingbot.controls import ( CameraPoseIntegrator, - KeyboardState, PoseSegment, ) diff --git a/integrations/lingbot/lingbot/runtime.py b/integrations/lingbot/lingbot/runtime.py index 501376b7f..ef7b549ef 100644 --- a/integrations/lingbot/lingbot/runtime.py +++ b/integrations/lingbot/lingbot/runtime.py @@ -17,6 +17,7 @@ from loguru import logger from flashdreams.core.distributed import init as init_distributed +from flashdreams.infra.config import derive_config from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner_io import ( load_first_frame_tensor, @@ -954,6 +955,10 @@ def _apply_webrtc_runtime_options( def _default_pipeline_factory(pipeline_config: Any, device: str) -> Any: + pipeline_config = derive_config( + base_config=pipeline_config, + diffusion_model=dict(transformer=dict(init_device=device)), + ) return pipeline_config.setup().to(device=device).eval() diff --git a/integrations/lingbot/lingbot/webrtc/session.py b/integrations/lingbot/lingbot/webrtc/session.py index 0f95b6618..03c1d2b90 100644 --- a/integrations/lingbot/lingbot/webrtc/session.py +++ b/integrations/lingbot/lingbot/webrtc/session.py @@ -27,7 +27,7 @@ import urllib.parse from dataclasses import dataclass, field from pathlib import Path -from typing import Any +from typing import Any, cast import cv2 import numpy as np @@ -46,10 +46,6 @@ UserInputSchema, ) from flashdreams.runtime.types import StepRequest, StepResult -from flashdreams.serving.webrtc.controls import ( - CameraPoseIntegrator, - PoseSegment, -) from flashdreams.serving.webrtc.encoders import EncoderBackend from flashdreams.serving.webrtc.manager import ( DEFAULT_CLIENT_LIVENESS_TIMEOUT_S, @@ -60,6 +56,7 @@ ThreadAffineDistributedWebRTCRuntime, ) from flashdreams.serving.webrtc.server import SessionBusyError +from lingbot.controls import CameraPoseIntegrator, PoseSegment from lingbot.encoder.utils import preprocess_example_poses from lingbot.input_mapping import ( FIELD_CAMERA_INTRINSICS, @@ -798,7 +795,10 @@ def _initialize_sync(self) -> None: enable_sync_and_profile=True, diffusion_model=dict( seed=rollout_seed, - transformer=dict(compile_network=self.config.compile_network), + transformer=dict( + compile_network=self.config.compile_network, + init_device=str(self._device), + ), ), ) self._pipeline = pipeline_config.setup().to(device=self._device) @@ -1095,7 +1095,7 @@ def _close_sync(self) -> None: def _generate_one_chunk_sync( self, *, - segments: list[PoseSegment], + segments: list[Any], frame_times: list[float], ) -> StepResult: if ( @@ -1113,8 +1113,9 @@ def _generate_one_chunk_sync( ) if not segments: raise LingbotRuntimeError(f"Chunk={step_index} received empty segments.") + pose_segments = cast(list[PoseSegment], segments) poses = self.pose_integrator.integrate_chunk( - segments=segments, frame_times=frame_times + segments=pose_segments, frame_times=frame_times ) poses_t = torch.from_numpy(poses).to(device=self._device, dtype=torch.float32) poses_t = poses_t.view(num_frames, 4, 4) diff --git a/integrations/lingbot/lingbot/webrtc/web/adapter.css b/integrations/lingbot/lingbot/webrtc/web/adapter.css index 61e8116ee..30ed19e51 100644 --- a/integrations/lingbot/lingbot/webrtc/web/adapter.css +++ b/integrations/lingbot/lingbot/webrtc/web/adapter.css @@ -4,13 +4,26 @@ .firstFramePreview { position: absolute; inset: 0; - width: 100%; - height: 100%; - object-fit: cover; opacity: 0; transition: opacity 220ms ease; } +.stageVideo, +.firstFramePreview { + inset: 50% auto auto 50%; + width: min( + 100vw, + var(--lingbot-video-width, 832px), + var(--lingbot-video-width-from-vh, 179.31vh) + ); + height: auto; + max-height: min(100vh, var(--lingbot-video-height, 464px)); + aspect-ratio: var(--lingbot-video-aspect, 832 / 464); + transform: translate(-50%, -50%); + object-fit: contain; + object-position: center; +} + body.is-ready-preview .firstFramePreview { opacity: 1; } diff --git a/integrations/lingbot/lingbot/webrtc/web/adapter.js b/integrations/lingbot/lingbot/webrtc/web/adapter.js index b323c32be..91389585e 100644 --- a/integrations/lingbot/lingbot/webrtc/web/adapter.js +++ b/integrations/lingbot/lingbot/webrtc/web/adapter.js @@ -319,10 +319,29 @@ function applyInitialScene(scene) { } renderEventControls() context.setModelName(scene.model || "Lingbot") + applyVideoSizing(scene.resolution) context.setResolution(scene.resolution?.width, scene.resolution?.height) updatePreview() } +function applyVideoSizing(resolution) { + const width = Number(resolution?.width) + const height = Number(resolution?.height) + if ( + !Number.isFinite(width) || + !Number.isFinite(height) || + width <= 0 || + height <= 0 + ) { + return + } + const style = document.documentElement.style + style.setProperty("--lingbot-video-width", `${width}px`) + style.setProperty("--lingbot-video-height", `${height}px`) + style.setProperty("--lingbot-video-width-from-vh", `${(width / height) * 100}vh`) + style.setProperty("--lingbot-video-aspect", `${width} / ${height}`) +} + function mockInitialScene() { return { prompt: "Drive through a cinematic city street at sunset.", @@ -476,7 +495,7 @@ function attachListeners() { export default { modelName: "Lingbot", - stylesheet: new URL("./adapter.css", import.meta.url).href, + stylesheet: new URL("./adapter.css?v=lingbot-video-size-v2", import.meta.url).href, controls, async mount(sharedContext) { diff --git a/integrations/lingbot/tests/test_controls.py b/integrations/lingbot/tests/test_controls.py index 108536e25..cf3aaa1ae 100644 --- a/integrations/lingbot/tests/test_controls.py +++ b/integrations/lingbot/tests/test_controls.py @@ -19,14 +19,14 @@ import numpy as np import pytest - -from flashdreams.serving.webrtc.controls import ( +from lingbot.controls import ( CameraPoseIntegrator, KeyboardResampler, - KeyboardState, PoseSegment, ) +from flashdreams.runtime.keyboard import KeyboardState + pytestmark = pytest.mark.ci_cpu ## KeyboardState basics (unchanged from the old design) diff --git a/integrations/lingbot/tests/test_demo_api.py b/integrations/lingbot/tests/test_demo_api.py index 8b33b865b..98444a7c5 100644 --- a/integrations/lingbot/tests/test_demo_api.py +++ b/integrations/lingbot/tests/test_demo_api.py @@ -59,6 +59,7 @@ from flashdreams.runtime.demo import ( DemoSpec, Mp4OutputSpec, + NullOutputSpec, UserInputWindow, WebRTCOutputSpec, ) @@ -118,7 +119,7 @@ def test_lingbot_demo_adapter_declares_shared_demo_modes() -> None: 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_output_modes() == ("mp4", "null", "webrtc") fields = { field.name for field in adapter.inference_input_schema.global_conditioning_fields @@ -137,6 +138,66 @@ def test_lingbot_demo_adapter_declares_shared_demo_modes() -> None: assert step_fields == {FIELD_CAMERA_TRAJECTORY, FIELD_CAMERA_INTRINSICS} +def test_lingbot_replay_cli_builds_null_output_spec() -> None: + args = parse_args(["replay", "--output-mode", "null", "--total-blocks", "1"]) + + spec = _replay_spec(args) + + assert spec.model_id == LINGBOT_MODEL_ID + assert spec.input_mode == "replay" + assert isinstance(spec.output, NullOutputSpec) + assert isinstance(spec.scenario, dict) + assert spec.scenario[FIELD_TOTAL_BLOCKS] == 1 + + +def test_lingbot_replay_cli_requires_output_only_for_mp4(tmp_path: Path) -> None: + parse_args(["replay", "--output", str(tmp_path / "demo.mp4")]) + + with pytest.raises(SystemExit): + parse_args(["replay"]) + with pytest.raises(SystemExit): + parse_args( + [ + "replay", + "--output-mode", + "null", + "--output", + str(tmp_path / "demo.mp4"), + ] + ) + + +def test_lingbot_replay_adapter_accepts_null_output(tmp_path: Path) -> None: + image = tmp_path / "image.jpg" + poses = tmp_path / "poses.npy" + intrinsics = tmp_path / "intrinsics.npy" + image.write_bytes(b"fake") + _write_camera_assets(poses, intrinsics) + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + input_mode="replay", + scenario={ + "prompt": "drive through a city", + "image_path": image, + "pose_path": poses, + "intrinsic_path": intrinsics, + "total_blocks": 1, + }, + output=NullOutputSpec(), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + runtime_options={"pipeline_config": object()}, + ), + ) + + prepared = LingbotDemoAdapter().prepare_scenario(spec) + + assert prepared.initial_inputs.global_conditioning[FIELD_FIRST_FRAME_PATH] == image + assert prepared.initial_inputs.global_conditioning[FIELD_TOTAL_BLOCKS] == 1 + + def test_lingbot_replay_demo_uses_shared_runner(tmp_path: Path) -> None: image = tmp_path / "image.jpg" poses = tmp_path / "poses.npy" @@ -286,6 +347,60 @@ def pipeline_factory(pipeline_config: object, device: str) -> _FakeLingbotPipeli ] +def test_lingbot_replay_demo_runs_with_null_output( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lingbot.runtime as runtime_module + + image = tmp_path / "image.jpg" + poses = tmp_path / "poses.npy" + intrinsics = tmp_path / "intrinsics.npy" + image.write_bytes(b"fake") + _write_camera_assets(poses, intrinsics, frames=16) + pipeline = _FakeLingbotPipeline() + monkeypatch.setattr( + runtime_module, + "load_first_frame_tensor", + lambda *args, **kwargs: torch.zeros(1, 3, 2, 2), + ) + + spec = DemoSpec( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + input_mode="replay", + scenario={ + "prompt": "drive through a city", + "image_path": image, + "pose_path": poses, + "intrinsic_path": intrinsics, + "total_blocks": 1, + "pixel_height": 2, + "pixel_width": 2, + }, + output=NullOutputSpec(), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=DEFAULT_LINGBOT_PRESET, + device="cpu", + runtime_options={"pipeline_config": object()}, + ), + ) + + def pipeline_factory(pipeline_config: object, device: str) -> _FakeLingbotPipeline: + del pipeline_config, device + return pipeline + + result = run_replay_demo( + spec=spec, + adapter=LingbotDemoAdapter(pipeline_factory=pipeline_factory), + ) + + assert result.status == "completed" + assert result.artifacts == () + assert len(pipeline.generate_calls) == 1 + + def test_lingbot_replay_invalid_scenario_fails_before_runtime_creation( tmp_path: Path, ) -> None: @@ -713,7 +828,7 @@ def test_lingbot_webrtc_demo_uses_shared_viewer_shell( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - import flashdreams.runtime.demo.webrtc as shared_webrtc_module + import flashdreams.serving.webrtc.demo as shared_webrtc_module _patch_lingbot_webrtc_example(monkeypatch, tmp_path) app_calls: list[dict[str, Any]] = [] @@ -772,7 +887,7 @@ def test_lingbot_webrtc_demo_serves_through_shared_runner( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - import flashdreams.runtime.demo.webrtc as shared_webrtc_module + import flashdreams.serving.webrtc.demo as shared_webrtc_module _patch_lingbot_webrtc_example(monkeypatch, tmp_path) server_calls: list[dict[str, Any]] = [] diff --git a/integrations/lingbot/tests/test_keyboard_parity.py b/integrations/lingbot/tests/test_keyboard_parity.py index b233d7266..fcbf924f4 100644 --- a/integrations/lingbot/tests/test_keyboard_parity.py +++ b/integrations/lingbot/tests/test_keyboard_parity.py @@ -16,6 +16,7 @@ import numpy as np import pytest import torch +from lingbot.controls import CameraPoseIntegrator, KeyboardResampler from lingbot.input_mapping import ( FIELD_CAMERA_TRAJECTORY, KeyboardToCameraCommand, @@ -32,7 +33,6 @@ UserInputs, UserInputSchema, ) -from flashdreams.serving.webrtc.controls import CameraPoseIntegrator, KeyboardResampler pytestmark = pytest.mark.ci_cpu diff --git a/integrations/lingbot/tests/test_smoke.py b/integrations/lingbot/tests/test_smoke.py index 171bacb89..f51d0c2bd 100644 --- a/integrations/lingbot/tests/test_smoke.py +++ b/integrations/lingbot/tests/test_smoke.py @@ -297,6 +297,14 @@ def test_lingbot_configs_carry_documented_checkpoint_disk_requirement() -> None: ) +def test_lingbot_configs_enable_streaming_checkpoint_load() -> None: + """Use bounded checkpoint loading for every LingBot model preset.""" + for cfg in RUNNER_CONFIGS.values(): + transformer = cfg.pipeline.diffusion_model.transformer + assert isinstance(transformer, LingbotWorldTransformerConfig) + assert transformer.stream_checkpoint + + def test_v2_only_replaces_the_v1_checkpoint() -> None: """Derive the v2 model by replacing only the v1 checkpoint and slug.""" expected = derive_config( diff --git a/integrations/lingbot/tests/test_webrtc_session_branch.py b/integrations/lingbot/tests/test_webrtc_session_branch.py index 8cf8e6cda..1a37d60f6 100644 --- a/integrations/lingbot/tests/test_webrtc_session_branch.py +++ b/integrations/lingbot/tests/test_webrtc_session_branch.py @@ -17,6 +17,7 @@ import numpy as np import pytest import torch +from lingbot.controls import CameraPoseIntegrator, KeyboardResampler from lingbot.input_mapping import ( FIELD_CAMERA_INTRINSICS, FIELD_CAMERA_TRAJECTORY, @@ -27,10 +28,9 @@ from lingbot.webrtc.session import LINGBOT_WEBRTC_SOURCE_SCHEMA from flashdreams.runtime.canonical import InputCanonicalizer +from flashdreams.runtime.demo import RealtimeEventResampler from flashdreams.runtime.inputs import InferenceInput, TimeWindow from flashdreams.runtime.types import StepRequest, StepResult -from flashdreams.serving.realtime.input import KeyboardResampler -from flashdreams.serving.webrtc.controls import CameraPoseIntegrator from flashdreams.serving.webrtc.manager import ( BaseWebRTCSessionManager, ManagedWebRTCSession, @@ -151,7 +151,7 @@ def _managed_session(runtime: _FakeRuntime) -> ManagedWebRTCSession: video_track=_FakeCloseable(), # ty:ignore[invalid-argument-type] video_encoder=_FakeVideoEncoder(), # ty:ignore[invalid-argument-type] peer_connection=_FakeCloseable(), - resampler=KeyboardResampler(fps=_FPS, start_v=0.0), + resampler=RealtimeEventResampler(fps=_FPS, start_v=0.0), inference_session=runtime.session, ) diff --git a/integrations/omnidreams/omnidreams/demo/controls.py b/integrations/omnidreams/omnidreams/demo/controls.py new file mode 100644 index 000000000..cedc384ae --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/controls.py @@ -0,0 +1,305 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OmniDreams-owned keyboard segmenting and camera pose integration.""" + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass, field +from typing import Literal + +import numpy as np + +from flashdreams.runtime.keyboard import WSAD_SUPPORTED_KEYS, KeyboardState + +PoseSegment = tuple[float, float, frozenset[str]] +SPARSE_KEY_SEGMENTS_METADATA_KEY = "sparse_key_segments" +"""Legacy OmniDreams WebRTC metadata key retained for debug/multi-rank paths.""" + + +class KeyboardResampler: + """Resample sparse keydown/keyup edges into an OmniDreams driving timeline.""" + + def __init__( + self, + *, + fps: float, + start_v: float = 0.0, + supported_keys: frozenset[str] = WSAD_SUPPORTED_KEYS, + ) -> None: + if fps <= 0: + raise ValueError("fps must be > 0") + self._fps = float(fps) + self._dt = 1.0 / self._fps + self._supported_keys = supported_keys + self.next_chunk_start_v = start_v + self._event_log: deque[tuple[float, dict[str, str]]] = deque() + self._carried_state = KeyboardState(supported_keys=supported_keys) + + @property + def fps(self) -> float: + return self._fps + + @property + def dt(self) -> float: + return self._dt + + def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: + entry = (arrival_t, {"event": event, "key": key}) + if not self._event_log or arrival_t >= self._event_log[-1][0]: + self._event_log.append(entry) + return + for index, (event_t, _) in enumerate(self._event_log): + if arrival_t < event_t: + self._event_log.insert(index, entry) + return + self._event_log.append(entry) + + def advance_to(self, end_v: float) -> None: + """Fold queued edges into carried state without producing frame samples.""" + if end_v < self.next_chunk_start_v: + raise ValueError("end_v must be >= next_chunk_start_v") + + chunk_start_v = self.next_chunk_start_v + while self._event_log and self._event_log[0][0] < chunk_start_v: + _, payload = self._event_log.popleft() + self._carried_state.apply_event(**payload) + while self._event_log and self._event_log[0][0] <= end_v: + _, payload = self._event_log.popleft() + self._carried_state.apply_event(**payload) + self.next_chunk_start_v = end_v + + def sample_chunk(self, num_frames: int) -> tuple[list[PoseSegment], list[float]]: + if num_frames < 1: + raise ValueError("num_frames must be >= 1") + + chunk_start_v = self.next_chunk_start_v + chunk_end_v = chunk_start_v + num_frames * self._dt + + while self._event_log and self._event_log[0][0] < chunk_start_v: + _, payload = self._event_log.popleft() + self._carried_state.apply_event(**payload) + + segments: list[PoseSegment] = [] + prev_t = chunk_start_v + prev_state = self._carried_state.resolved_effective_keys() + while self._event_log and self._event_log[0][0] <= chunk_end_v: + event_t, payload = self._event_log.popleft() + if event_t > prev_t: + segments.append((prev_t, event_t, prev_state)) + self._carried_state.apply_event(**payload) + prev_state = self._carried_state.resolved_effective_keys() + prev_t = event_t + if prev_t < chunk_end_v: + segments.append((prev_t, chunk_end_v, prev_state)) + elif not segments: + segments.append((chunk_start_v, chunk_end_v, prev_state)) + + frame_times = [chunk_start_v + (i + 1) * self._dt for i in range(num_frames)] + self.next_chunk_start_v = chunk_end_v + return segments, frame_times + + def reset(self, *, start_v: float) -> None: + self._event_log.clear() + self._carried_state = KeyboardState(supported_keys=self._supported_keys) + self.next_chunk_start_v = start_v + + def event_log_size(self) -> int: + return len(self._event_log) + + +def _rotation_matrix(axis: str, angle_rad: float) -> np.ndarray: + cos_t = np.float32(np.cos(angle_rad)) + sin_t = np.float32(np.sin(angle_rad)) + if axis == "x": + return np.array( + [ + [1.0, 0.0, 0.0], + [0.0, cos_t, -sin_t], + [0.0, sin_t, cos_t], + ], + dtype=np.float32, + ) + if axis == "y": + return np.array( + [ + [cos_t, 0.0, sin_t], + [0.0, 1.0, 0.0], + [-sin_t, 0.0, cos_t], + ], + dtype=np.float32, + ) + if axis == "z": + return np.array( + [ + [cos_t, -sin_t, 0.0], + [sin_t, cos_t, 0.0], + [0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + return np.eye(3, dtype=np.float32) + + +@dataclass(slots=True) +class CameraPoseIntegrator: + """Integrate a piecewise-constant keyboard timeline into a camera path.""" + + move_speed_per_s: float = 0.8 + rotate_speed_rad_per_s: float = float(np.deg2rad(32.0)) + pitch_limit_rad: float = float(np.deg2rad(85.0)) + coordinate_system: Literal["RDF", "FLU"] = "RDF" + _current_pose: np.ndarray = field( + default_factory=lambda: np.eye(4, dtype=np.float32), + ) + _current_pitch: float = 0.0 + + def __post_init__(self) -> None: + if self.coordinate_system not in {"RDF", "FLU"}: + raise ValueError( + "coordinate_system must be 'RDF' (right-down-forward) " + "or 'FLU' (forward-left-up)" + ) + + def reset(self, pose: np.ndarray | None = None) -> None: + if pose is None: + self._current_pose = np.eye(4, dtype=np.float32) + self._current_pitch = 0.0 + return + if pose.shape != (4, 4): + raise ValueError(f"Expected pose shape (4, 4), got {pose.shape}") + self._current_pose = pose.astype(np.float32, copy=True) + if self.coordinate_system == "FLU": + self._current_pitch = float(np.arcsin(np.clip(pose[2, 0], -1.0, 1.0))) + else: + self._current_pitch = float(np.arctan2(pose[2, 1], pose[1, 1])) + + def current_pose(self) -> np.ndarray: + return self._current_pose.copy() + + def _advance(self, *, state: frozenset[str], duration: float) -> None: + if duration <= 0: + return + + yaw_rate = 0.0 + if self.coordinate_system == "FLU": + if "a" in state or "j" in state: + yaw_rate += self.rotate_speed_rad_per_s + if "d" in state or "l" in state: + yaw_rate -= self.rotate_speed_rad_per_s + else: + if "a" in state or "j" in state: + yaw_rate -= self.rotate_speed_rad_per_s + if "d" in state or "l" in state: + yaw_rate += self.rotate_speed_rad_per_s + pitch_rate = 0.0 + if "i" in state: + pitch_rate += self.rotate_speed_rad_per_s + if "k" in state: + pitch_rate -= self.rotate_speed_rad_per_s + + yaw_delta = yaw_rate * duration + pitch_delta = pitch_rate * duration + + new_pitch = self._current_pitch + pitch_delta + if -self.pitch_limit_rad <= new_pitch <= self.pitch_limit_rad: + self._current_pitch = new_pitch + else: + pitch_delta = 0.0 + + rot = self._current_pose[:3, :3] + trans = self._current_pose[:3, 3] + if self.coordinate_system == "FLU": + rot_pitch = _rotation_matrix("y", -pitch_delta) + rot_yaw = _rotation_matrix("z", yaw_delta) + else: + rot_pitch = _rotation_matrix("x", pitch_delta) + rot_yaw = _rotation_matrix("y", yaw_delta) + rot_new = rot_yaw @ rot @ rot_pitch + + forward_rate = 0.0 + if "w" in state: + forward_rate += self.move_speed_per_s + if "s" in state: + forward_rate -= self.move_speed_per_s + right_rate = 0.0 + if "e" in state: + right_rate += self.move_speed_per_s + if "q" in state: + right_rate -= self.move_speed_per_s + + if self.coordinate_system == "FLU": + vec_forward = rot_new[:, 0] + vec_right = -rot_new[:, 1] + forward_flat = np.array( + [vec_forward[0], vec_forward[1], 0.0], dtype=np.float32 + ) + right_flat = np.array([vec_right[0], vec_right[1], 0.0], dtype=np.float32) + else: + vec_right = rot_new[:, 0] + vec_forward = rot_new[:, 2] + forward_flat = np.array( + [vec_forward[0], 0.0, vec_forward[2]], dtype=np.float32 + ) + right_flat = np.array([vec_right[0], 0.0, vec_right[2]], dtype=np.float32) + forward_norm = np.linalg.norm(forward_flat) + right_norm = np.linalg.norm(right_flat) + if forward_norm > 0: + forward_flat /= forward_norm + if right_norm > 0: + right_flat /= right_norm + + move_vec = forward_flat * (forward_rate * duration) + right_flat * ( + right_rate * duration + ) + self._current_pose = np.eye(4, dtype=np.float32) + self._current_pose[:3, :3] = rot_new + self._current_pose[:3, 3] = trans + move_vec + + def integrate_chunk( + self, + *, + segments: list[PoseSegment], + frame_times: list[float], + ) -> np.ndarray: + if not segments: + raise ValueError("segments must be non-empty") + if not frame_times: + raise ValueError("frame_times must be non-empty") + chunk_start = segments[0][0] + chunk_end = segments[-1][1] + if any( + frame_times[i] >= frame_times[i + 1] for i in range(len(frame_times) - 1) + ): + raise ValueError("frame_times must be strictly increasing") + if frame_times[0] < chunk_start - 1e-9 or frame_times[-1] > chunk_end + 1e-9: + raise ValueError( + "frame_times must lie within the chunk window " + f"[{chunk_start}, {chunk_end}]" + ) + + poses: list[np.ndarray] = [] + cur_t = chunk_start + ft_idx = 0 + for _, seg_end, seg_state in segments: + while ft_idx < len(frame_times) and frame_times[ft_idx] <= seg_end: + target_t = frame_times[ft_idx] + self._advance(state=seg_state, duration=target_t - cur_t) + cur_t = target_t + poses.append(self._current_pose.copy()) + ft_idx += 1 + if seg_end > cur_t: + self._advance(state=seg_state, duration=seg_end - cur_t) + cur_t = seg_end + + return np.stack(poses, axis=0).astype(np.float32) + + +__all__ = [ + "CameraPoseIntegrator", + "KeyboardResampler", + "PoseSegment", + "SPARSE_KEY_SEGMENTS_METADATA_KEY", + "WSAD_SUPPORTED_KEYS", +] diff --git a/integrations/omnidreams/omnidreams/demo/providers.py b/integrations/omnidreams/omnidreams/demo/providers.py index eff6de760..27c3784e5 100644 --- a/integrations/omnidreams/omnidreams/demo/providers.py +++ b/integrations/omnidreams/omnidreams/demo/providers.py @@ -28,22 +28,26 @@ UserInputWindow, ) from flashdreams.runtime.demo.session_inputs import ControlDecision -from flashdreams.runtime.demo.timing import SPARSE_KEY_SEGMENTS_METADATA_KEY from flashdreams.runtime.inputs import ( InferenceInput, InferenceInputSchema, InputField, UserInputCapability, + UserInputs, UserInputSchema, ) from flashdreams.runtime.types import StepRequirements -from flashdreams.serving.realtime.input import ( +from flashdreams.serving.webrtc.services import ( + WEBRTC_SKIPPED_INPUTS_METADATA_KEY, + WEBRTC_SKIPPED_WINDOW_METADATA_KEY, +) + +from .controls import ( WSAD_SUPPORTED_KEYS, CameraPoseIntegrator, KeyboardResampler, PoseSegment, ) - from .spec import ( DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID, OmnidreamsLudusReplayScenario, @@ -342,22 +346,71 @@ def _sample_controls( request: StepRequirements, user_window: UserInputWindow, ) -> tuple[list[PoseSegment], list[float]]: - raw_segments = user_window.metadata.get(SPARSE_KEY_SEGMENTS_METADATA_KEY) - if isinstance(raw_segments, tuple): - frame_times = list(user_window.frame_times) - if len(frame_times) != request.input_frame_count: - raise RuntimeError( - "OmniDreams Ludus realtime window frame_times length does " - "not match the requested input frame count." - ) - return [_pose_segment(segment) for segment in raw_segments], frame_times - if raw_segments is not None: + frame_times = list(user_window.frame_times) + if frame_times and len(frame_times) != request.input_frame_count: raise RuntimeError( - "OmniDreams Ludus realtime key segments metadata must be a tuple." + "OmniDreams Ludus realtime window frame_times length does not " + "match the requested input frame count." ) - return self._require_keyboard_resampler().sample_chunk( + resampler = self._require_keyboard_resampler() + self._advance_skipped_input_state(user_window=user_window, resampler=resampler) + # Realtime/WebRTC windows carry explicit frame times on the driver's + # clock. Batch replay windows do not, so the provider-owned trace + # resampler must keep advancing from its prior chunk instead. + if frame_times: + resampler.next_chunk_start_v = user_window.start_s + self._record_keyboard_events(resampler, user_window.inputs) + segments, sampled_frame_times = resampler.sample_chunk( request.input_frame_count ) + if not frame_times: + frame_times = sampled_frame_times + return segments, frame_times + + def _advance_skipped_input_state( + self, + *, + user_window: UserInputWindow, + resampler: KeyboardResampler, + ) -> None: + skipped_inputs = user_window.metadata.get(WEBRTC_SKIPPED_INPUTS_METADATA_KEY) + skipped_window = user_window.metadata.get(WEBRTC_SKIPPED_WINDOW_METADATA_KEY) + if not isinstance(skipped_inputs, UserInputs): + return + if not isinstance(skipped_window, tuple) or len(skipped_window) != 2: + return + start_value, end_value = skipped_window + if not isinstance(start_value, int | float) or not isinstance( + end_value, + int | float, + ): + return + start_s = float(start_value) + end_s = float(end_value) + if end_s <= start_s: + return + resampler.next_chunk_start_v = start_s + self._record_keyboard_events(resampler, skipped_inputs) + resampler.advance_to(end_s) + + @staticmethod + def _record_keyboard_events( + resampler: KeyboardResampler, + inputs: UserInputs, + ) -> None: + for event in inputs.events: + if event.event_type not in {"key_down", "key_up", "keydown", "keyup"}: + continue + key = event.payload.get("key") + if not isinstance(key, str): + continue + resampler.on_edge( + arrival_t=event.timestamp_s, + event="keydown" + if event.event_type in {"key_down", "keydown"} + else "keyup", + key=key, + ) def _consume_timestamps(self, num_frames: int) -> np.ndarray: step_us = int(round(1_000_000 / float(self._scenario.fps))) @@ -651,17 +704,6 @@ def _segments_metadata( ) -def _pose_segment(value: object) -> PoseSegment: - if not isinstance(value, tuple) or len(value) != 3: - raise RuntimeError("OmniDreams Ludus key segment must be a 3-tuple.") - start, end, keys = value - if not isinstance(start, int | float) or not isinstance(end, int | float): - raise RuntimeError("OmniDreams Ludus key segment bounds must be numeric.") - if not isinstance(keys, frozenset | set | tuple | list): - raise RuntimeError("OmniDreams Ludus key segment keys must be a sequence.") - return (float(start), float(end), frozenset(str(key) for key in keys)) - - def _close_rasterizer(rasterizer: Any | None) -> None: if rasterizer is None: return diff --git a/integrations/omnidreams/omnidreams/demo/webrtc.py b/integrations/omnidreams/omnidreams/demo/webrtc.py index e5d8e2a84..fb3d12628 100644 --- a/integrations/omnidreams/omnidreams/demo/webrtc.py +++ b/integrations/omnidreams/omnidreams/demo/webrtc.py @@ -32,17 +32,17 @@ WebRTCAppResources, WebRTCOutputSpec, ) -from flashdreams.runtime.demo.webrtc import ( +from flashdreams.serving.webrtc.bootstrap import run_webrtc_server +from flashdreams.serving.webrtc.demo import ( CreateWebRTCApp, RunWebRTCServer, serve_webrtc_demo, ) -from flashdreams.serving.webrtc.bootstrap import run_webrtc_server -from flashdreams.serving.webrtc.controls import WSAD_SUPPORTED_KEYS from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager from flashdreams.serving.webrtc.server import create_webrtc_app from .adapter import OmnidreamsDemoAdapter, RuntimeFactory +from .controls import WSAD_SUPPORTED_KEYS from .spec import ( DEFAULT_OMNIDREAMS_PRESET, OMNIDREAMS_MODEL_ID, diff --git a/integrations/omnidreams/omnidreams/demo/webrtc_legacy.py b/integrations/omnidreams/omnidreams/demo/webrtc_legacy.py index 0d7768636..401d9158e 100644 --- a/integrations/omnidreams/omnidreams/demo/webrtc_legacy.py +++ b/integrations/omnidreams/omnidreams/demo/webrtc_legacy.py @@ -18,7 +18,7 @@ from __future__ import annotations from collections.abc import Callable, Mapping -from typing import Any +from typing import Any, cast import torch from loguru import logger @@ -46,13 +46,11 @@ WebRTCAppResources, WebRTCOutputSpec, ) -from flashdreams.runtime.demo.timing import SPARSE_KEY_SEGMENTS_METADATA_KEY -from flashdreams.runtime.demo.webrtc import ( +from flashdreams.serving.webrtc.demo import ( CreateWebRTCApp, RunWebRTCServer, serve_webrtc_demo, ) -from flashdreams.serving.webrtc.controls import WSAD_SUPPORTED_KEYS, PoseSegment from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager from flashdreams.serving.webrtc.runtime import ( ThreadAffineDistributedWebRTCRuntime, @@ -60,6 +58,12 @@ ) from flashdreams.serving.webrtc.services import WEBRTC_USER_INPUT_SCHEMA +from .controls import ( + SPARSE_KEY_SEGMENTS_METADATA_KEY, + WSAD_SUPPORTED_KEYS, + KeyboardResampler, + PoseSegment, +) from .providers import LudusSceneConditioningProvider from .runtime import OmnidreamsRuntime, OmnidreamsRuntimeOptions from .spec import ( @@ -195,9 +199,10 @@ def _reset_rollout_sync(self, session_input: None = None) -> None: def _generate_one_chunk_sync( self, *, - segments: list[PoseSegment], + segments: list[Any], frame_times: list[float], ) -> StepResult: + pose_segments = cast(list[PoseSegment], segments) request = self._next_step_request_sync() if request is None: raise OmnidreamsWebRTCModelRuntimeError( @@ -207,7 +212,7 @@ def _generate_one_chunk_sync( canonical_inputs=CanonicalInputs(), inference_input=InferenceInput( metadata={ - SPARSE_KEY_SEGMENTS_METADATA_KEY: tuple(segments), + SPARSE_KEY_SEGMENTS_METADATA_KEY: tuple(pose_segments), "frame_times": tuple(frame_times), "window_start_s": request.step_index / float(self.config.fps), "window_end_s": (request.step_index + len(frame_times)) @@ -693,6 +698,7 @@ def _serve_legacy_omnidreams_webrtc_demo( supported_control_keys=WSAD_SUPPORTED_KEYS, fatal_generation_errors=True, client_liveness_timeout_s=output.client_liveness_timeout_s, + legacy_segment_resampler_factory=KeyboardResampler, ) from importlib.resources import files diff --git a/integrations/omnidreams/omnidreams/interactive_drive/input/keyboard.py b/integrations/omnidreams/omnidreams/interactive_drive/input/keyboard.py index a8083a659..ea8adecda 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/input/keyboard.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/input/keyboard.py @@ -11,11 +11,11 @@ VehicleState, ) -from flashdreams.serving.realtime.input import ( +from flashdreams.runtime.keyboard import ( DRIVING_SUPPORTED_KEYS, normalize_key, ) -from flashdreams.serving.realtime.input import ( +from flashdreams.runtime.keyboard import ( KeyboardState as RealtimeKeyboardState, ) diff --git a/integrations/omnidreams/tests/test_demo_api.py b/integrations/omnidreams/tests/test_demo_api.py index 3d083535a..846e78f96 100644 --- a/integrations/omnidreams/tests/test_demo_api.py +++ b/integrations/omnidreams/tests/test_demo_api.py @@ -6,7 +6,7 @@ import asyncio import json import sys -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from pathlib import Path from types import SimpleNamespace from typing import Any @@ -32,6 +32,7 @@ PrecomputedHDMapProvider, ) from omnidreams.demo.app import _replay_spec, _webrtc_spec, parse_args +from omnidreams.demo.controls import SPARSE_KEY_SEGMENTS_METADATA_KEY from omnidreams.demo.replay import ( OmnidreamsReplayRuntime, OmnidreamsReplayRuntimeOptions, @@ -57,6 +58,8 @@ StepRequest, StepRequirements, StepResult, + UserInputEvent, + UserInputs, ) from flashdreams.runtime.demo import ( DemoSpec, @@ -70,13 +73,14 @@ WebRTCOutputSpec, ) from flashdreams.runtime.demo.replay import run_replay_demo -from flashdreams.runtime.demo.timing import SPARSE_KEY_SEGMENTS_METADATA_KEY from flashdreams.serving.webrtc.manager import ( BaseWebRTCSessionManager, ManagedWebRTCSession, ) from flashdreams.serving.webrtc.server import SESSION_MANAGER_KEY from flashdreams.serving.webrtc.services import ( + WEBRTC_SKIPPED_INPUTS_METADATA_KEY, + WEBRTC_SKIPPED_WINDOW_METADATA_KEY, WebRTCInputSource, WebRTCTransportService, ) @@ -473,7 +477,11 @@ def test_omnidreams_ludus_provider_prepares_deterministic_hdmaps( first = provider.prepare_step( request=StepRequirements(step_index=0, input_frame_count=2), - user_window=UserInputWindow(start_s=0.0, end_s=2 / 30), + user_window=UserInputWindow( + start_s=0.0, + end_s=2 / 30, + frame_times=(1 / 30, 2 / 30), + ), ) assert first.inference_input is not None @@ -492,7 +500,11 @@ def test_omnidreams_ludus_provider_prepares_deterministic_hdmaps( provider.reset() reset_first = provider.prepare_step( request=StepRequirements(step_index=0, input_frame_count=2), - user_window=UserInputWindow(start_s=0.0, end_s=2 / 30), + user_window=UserInputWindow( + start_s=0.0, + end_s=2 / 30, + frame_times=(1 / 30, 2 / 30), + ), ) assert reset_first.inference_input is not None @@ -505,9 +517,15 @@ def test_omnidreams_ludus_provider_prepares_deterministic_hdmaps( start_s=0.0, end_s=2 / 30, frame_times=(1 / 30, 2 / 30), - metadata={ - SPARSE_KEY_SEGMENTS_METADATA_KEY: ((0.0, 2 / 30, frozenset({"w"})),) - }, + inputs=UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, + event_type="key_down", + payload={"key": "w"}, + ), + ) + ), ), ) @@ -520,6 +538,95 @@ def test_omnidreams_ludus_provider_prepares_deterministic_hdmaps( assert rasterizers[0].closed is True +def test_omnidreams_ludus_provider_advances_replay_trace_without_frame_times( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _scene, _rasterizers = _install_fake_ludus_provider_dependencies(monkeypatch) + scene_path = tmp_path / "scene.usdz" + scene_path.write_bytes(b"fake") + adapter = OmnidreamsDemoAdapter() + spec = _ludus_replay_demo_spec( + tmp_path=tmp_path, + scene_path=scene_path, + total_blocks=2, + keyboard_events=( + {"timestamp_s": 0.0, "event": "keydown", "key": "w"}, + {"timestamp_s": 0.08, "event": "keydown", "key": "d"}, + ), + ) + prepared = adapter.prepare_scenario(spec) + provider = adapter.create_model_input_provider(spec, prepared) + assert isinstance(provider, LudusSceneConditioningProvider) + provider.prepare_initial_input() + replay_window = UserInputWindow(start_s=0.0, end_s=3600.0) + + first = provider.prepare_step( + request=StepRequirements(step_index=0, input_frame_count=2), + user_window=replay_window, + ) + second = provider.prepare_step( + request=StepRequirements(step_index=1, input_frame_count=2), + user_window=replay_window, + ) + + assert first.inference_input is not None + assert second.inference_input is not None + assert first.inference_input.metadata["keyboard_segments"] == ( + (0.0, 2 / 30, ("w",)), + ) + assert second.inference_input.metadata["keyboard_segments"] == ( + (2 / 30, 0.08, ("w",)), + (0.08, 4 / 30, ("d", "w")), + ) + + +def test_omnidreams_ludus_provider_folds_webrtc_skipped_inputs( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _scene, _rasterizers = _install_fake_ludus_provider_dependencies(monkeypatch) + scene_path = tmp_path / "scene.usdz" + scene_path.write_bytes(b"fake") + adapter = OmnidreamsDemoAdapter() + spec = _ludus_replay_demo_spec( + tmp_path=tmp_path, + scene_path=scene_path, + total_blocks=1, + ) + prepared = adapter.prepare_scenario(spec) + provider = adapter.create_model_input_provider(spec, prepared) + assert isinstance(provider, LudusSceneConditioningProvider) + provider.prepare_initial_input() + + skipped_release = UserInputs( + events=( + UserInputEvent( + timestamp_s=0.1, + event_type="key_up", + payload={"key": "w"}, + ), + ) + ) + step = provider.prepare_step( + request=StepRequirements(step_index=0, input_frame_count=2), + user_window=UserInputWindow( + start_s=0.25, + end_s=0.25 + 2 / 30, + frame_times=(0.25 + 1 / 30, 0.25 + 2 / 30), + metadata={ + WEBRTC_SKIPPED_INPUTS_METADATA_KEY: skipped_release, + WEBRTC_SKIPPED_WINDOW_METADATA_KEY: (0.0, 0.25), + }, + ), + ) + + assert step.inference_input is not None + assert step.inference_input.metadata["keyboard_segments"] == ( + (0.25, 0.25 + 2 / 30, ()), + ) + + def test_omnidreams_replay_run_mode_uses_precomputed_provider( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -995,7 +1102,7 @@ def test_omnidreams_webrtc_demo_keeps_legacy_fallback_gates( def test_omnidreams_webrtc_demo_installs_model_assets_without_routes( monkeypatch: pytest.MonkeyPatch, ) -> None: - import flashdreams.runtime.demo.webrtc as shared_webrtc_module + import flashdreams.serving.webrtc.demo as shared_webrtc_module app_calls: list[dict[str, Any]] = [] @@ -1069,7 +1176,7 @@ def test_omnidreams_webrtc_adapter_caps_video_display_size() -> None: def test_omnidreams_webrtc_demo_serves_through_shared_runner( monkeypatch: pytest.MonkeyPatch, ) -> None: - import flashdreams.runtime.demo.webrtc as shared_webrtc_module + import flashdreams.serving.webrtc.demo as shared_webrtc_module server_calls: list[dict[str, Any]] = [] @@ -1413,6 +1520,10 @@ def _ludus_replay_demo_spec( tmp_path: Path, scene_path: Path, total_blocks: int, + keyboard_events: Sequence[Mapping[str, object]] = ( + {"timestamp_s": 0.0, "event": "keydown", "key": "w"}, + {"timestamp_s": 0.5, "event": "keyup", "key": "w"}, + ), output: Mp4OutputSpec | NullOutputSpec | None = None, ) -> DemoSpec: return DemoSpec( @@ -1421,10 +1532,7 @@ def _ludus_replay_demo_spec( input_mode="replay", scenario={ "conditioning_mode": OMNIDREAMS_CONDITIONING_LUDUS, - "keyboard_events": ( - {"timestamp_s": 0.0, "event": "keydown", "key": "w"}, - {"timestamp_s": 0.5, "event": "keyup", "key": "w"}, - ), + "keyboard_events": tuple(keyboard_events), "scene_path": scene_path, "scene_variant": "default", "camera_name": "camera_front_wide_120fov", @@ -1641,18 +1749,12 @@ def __init__(self, *, start_v: float, fps: int) -> None: def reset(self, *, start_v: float) -> None: self.next_chunk_start_v = start_v - def on_edge(self, *, arrival_t: float, event: str, key: str) -> None: - del arrival_t, event, key - - def sample_chunk( - self, - num_frames: int, - ) -> tuple[list[tuple[float, float, frozenset[str]]], list[float]]: + def sample_chunk(self, num_frames: int) -> list[float]: start = self.next_chunk_start_v frame_times = [start + (index + 1) * self.dt for index in range(num_frames)] end = frame_times[-1] self.next_chunk_start_v = end - return [(start, end, frozenset({"w"}))], frame_times + return frame_times class _FakeWebRTCVideoTrack: From c3d6e6633e9824b5c1957b4e3946e9dc294d597e Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Tue, 11 Aug 2026 18:42:01 +0000 Subject: [PATCH 19/19] Migrate FlashVSR demos to runtime API Signed-off-by: Gangzheng Tong --- .../flashdreams/serving/webrtc/server.py | 4 +- .../serving/webrtc/web/mock_ui_server.py | 4 +- .../serving/webrtc/web/request_session.css | 4 + .../serving/webrtc/web/request_session.html | 6 +- .../serving/webrtc/web/request_session.js | 12 + flashdreams/tests/test_webrtc_serving.py | 11 +- integrations/flashvsr/README.md | 45 ++ integrations/flashvsr/flashvsr/config.py | 4 +- integrations/flashvsr/flashvsr/corrector.py | 7 +- .../flashvsr/flashvsr/demo/__init__.py | 26 + .../flashvsr/flashvsr/demo/adapter.py | 158 ++++ integrations/flashvsr/flashvsr/demo/app.py | 317 ++++++++ .../flashvsr/flashvsr/demo/providers.py | 137 ++++ integrations/flashvsr/flashvsr/demo/server.py | 281 +++++++ integrations/flashvsr/flashvsr/demo/spec.py | 198 +++++ .../flashvsr/flashvsr/demo/web/adapter.css | 151 ++++ .../flashvsr/flashvsr/demo/web/adapter.js | 211 +++++ integrations/flashvsr/flashvsr/demo/webrtc.py | 207 +++++ integrations/flashvsr/flashvsr/runtime.py | 729 ++++++++++++++++++ integrations/flashvsr/pyproject.toml | 4 +- .../flashvsr/tests/test_runtime_api.py | 495 ++++++++++++ .../flashvsr/tests/test_webrtc_upload.py | 420 ++++++++++ uv.lock | 4 +- 23 files changed, 3417 insertions(+), 18 deletions(-) create mode 100644 integrations/flashvsr/flashvsr/demo/__init__.py create mode 100644 integrations/flashvsr/flashvsr/demo/adapter.py create mode 100644 integrations/flashvsr/flashvsr/demo/app.py create mode 100644 integrations/flashvsr/flashvsr/demo/providers.py create mode 100644 integrations/flashvsr/flashvsr/demo/server.py create mode 100644 integrations/flashvsr/flashvsr/demo/spec.py create mode 100644 integrations/flashvsr/flashvsr/demo/web/adapter.css create mode 100644 integrations/flashvsr/flashvsr/demo/web/adapter.js create mode 100644 integrations/flashvsr/flashvsr/demo/webrtc.py create mode 100644 integrations/flashvsr/flashvsr/runtime.py create mode 100644 integrations/flashvsr/tests/test_runtime_api.py create mode 100644 integrations/flashvsr/tests/test_webrtc_upload.py diff --git a/flashdreams/flashdreams/serving/webrtc/server.py b/flashdreams/flashdreams/serving/webrtc/server.py index 7a2c6694d..860640a6f 100644 --- a/flashdreams/flashdreams/serving/webrtc/server.py +++ b/flashdreams/flashdreams/serving/webrtc/server.py @@ -93,9 +93,9 @@ async def healthz(request: web.Request) -> web.StreamResponse: async def ui_config(_: web.Request) -> web.StreamResponse: payload: dict[str, str | None] = {"adapter_module": None} if model_web_dir is not None and (model_web_dir / "adapter.js").is_file(): - payload["adapter_module"] = "/model-static/adapter.js?v=model-ui-v2" + payload["adapter_module"] = "/model-static/adapter.js?v=model-ui-v4" if model_web_dir is not None and (model_web_dir / "adapter.css").is_file(): - payload["model_stylesheet"] = "/model-static/adapter.css?v=model-ui-v2" + payload["model_stylesheet"] = "/model-static/adapter.css?v=model-ui-v4" return web.json_response(payload) async def on_startup(app: web.Application) -> None: diff --git a/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py b/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py index 0c62855b5..867dd3e02 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py +++ b/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py @@ -86,10 +86,10 @@ def _serve_ui_config(self) -> bool: ui_config: dict[str, str | None] = {"adapter_module": None} if self.model_web_dir is not None: if (self.model_web_dir / "adapter.js").is_file(): - ui_config["adapter_module"] = "/model-static/adapter.js?v=model-ui-v2" + ui_config["adapter_module"] = "/model-static/adapter.js?v=model-ui-v4" if (self.model_web_dir / "adapter.css").is_file(): ui_config["model_stylesheet"] = ( - "/model-static/adapter.css?v=model-ui-v2" + "/model-static/adapter.css?v=model-ui-v4" ) payload = json.dumps(ui_config).encode("utf-8") self.send_response(200) diff --git a/flashdreams/flashdreams/serving/webrtc/web/request_session.css b/flashdreams/flashdreams/serving/webrtc/web/request_session.css index 237804419..53325e01b 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/request_session.css +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.css @@ -298,6 +298,10 @@ body[data-status="generating"] .connectButton { padding: 18px 20px 20px; } +.controlCard[hidden] { + display: none; +} + .controlCard h2, .logCard h2 { display: flex; diff --git a/flashdreams/flashdreams/serving/webrtc/web/request_session.html b/flashdreams/flashdreams/serving/webrtc/web/request_session.html index ad158656e..bd2bdcd86 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/request_session.html +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.html @@ -9,7 +9,7 @@ FlashDreams WebRTC Drive - +
@@ -47,7 +47,7 @@

FlashDreams WebRTC Drive

-
+
- + diff --git a/flashdreams/flashdreams/serving/webrtc/web/request_session.js b/flashdreams/flashdreams/serving/webrtc/web/request_session.js index 76b7507f9..e61f88407 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/request_session.js +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.js @@ -11,6 +11,7 @@ const mockMode = new URLSearchParams(window.location.search).has("mock") * @property {{postprocess?: boolean}=} capabilities * @property {(context: Object) => (void|Promise)=} mount * @property {(context: Object) => (void|Promise)=} beforeConnect + * @property {(context: Object) => void=} onConnect * @property {(action: Object, context: Object) => void=} onActionSent * @property {(payload: Object, context: Object) => boolean=} onControlMessage * @property {(visible: boolean, context: Object) => void=} onVideoVisibilityChanged @@ -34,6 +35,7 @@ const postprocessSelect = document.getElementById("postprocessSelect") const modelStageSlot = document.getElementById("modelStageSlot") const modelStatusSlot = document.getElementById("modelStatusSlot") const modelPanelSlot = document.getElementById("modelPanelSlot") +const controlCard = document.getElementById("controlCard") const modelControlSlot = document.getElementById("modelControlSlot") const controlRows = document.getElementById("controlRows") @@ -162,6 +164,11 @@ function setVideoVisible(visible) { modelAdapter?.onVideoVisibilityChanged?.(visible, modelContext) } +function syncControlCardVisibility() { + const isEmpty = controlRows.childElementCount === 0 && modelControlSlot.childElementCount === 0 + controlCard.hidden = isEmpty +} + function renderControls(groups) { controlRows.replaceChildren() allowedKeys = new Set() @@ -193,6 +200,7 @@ function renderControls(groups) { controlRows.append(row) } controlButtons = Array.from(controlRows.querySelectorAll("[data-control-key]")) + syncControlCardVisibility() } function setPostprocessDisabled(disabled) { @@ -275,6 +283,8 @@ const modelContext = { logEvent, releaseControls: releaseAllKeys, sendCommand: sendModelCommand, + setFlow, + setStatus, setModelName(name) { if (typeof name === "string" && name) { metrics.model = name @@ -339,6 +349,7 @@ async function loadModelAdapter() { } } await adapter.mount?.(modelContext) + syncControlCardVisibility() } function renderMetrics() { @@ -901,6 +912,7 @@ async function connectSession() { setFlow("connected; waiting for input") logEvent("control data channel open") startHeartbeat() + modelAdapter?.onConnect?.(modelContext) } channel.onclose = () => { connected = false diff --git a/flashdreams/tests/test_webrtc_serving.py b/flashdreams/tests/test_webrtc_serving.py index 71b1f48e8..1fb07cbef 100644 --- a/flashdreams/tests/test_webrtc_serving.py +++ b/flashdreams/tests/test_webrtc_serving.py @@ -285,7 +285,7 @@ def test_shared_viewer_exposes_model_extension_slots() -> None: html = web_dir.joinpath("request_session.html").read_text(encoding="utf-8") javascript = web_dir.joinpath("request_session.js").read_text(encoding="utf-8") - assert "/static/request_session.js?v=shared-webrtc-v4" in html + assert "/static/request_session.js?v=shared-webrtc-v5" in html for slot in ( "modelStageSlot", "modelStatusSlot", @@ -296,8 +296,13 @@ def test_shared_viewer_exposes_model_extension_slots() -> None: assert 'fetch("/api/ui/config")' in javascript assert "config.model_stylesheet" in javascript assert "stylesheetHrefs" in javascript + assert 'id="controlCard"' in html + assert "syncControlCardVisibility" in javascript assert "await modelAdapter?.beforeConnect?.(modelContext)" in javascript + assert "modelAdapter?.onConnect?.(modelContext)" in javascript assert "sendCommand: sendModelCommand" in javascript + assert "setFlow," in javascript + assert "setStatus," in javascript assert 'id="postprocessField"' in html assert 'fetch("/api/postprocess/options")' in javascript assert "@typedef {Object} WebRTCModelAdapter" in javascript @@ -354,7 +359,7 @@ async def test_packaged_webrtc_app_serves_model_adapter(tmp_path) -> None: try: config_response = await client.get("/api/ui/config") assert await config_response.json() == { - "adapter_module": "/model-static/adapter.js?v=model-ui-v2" + "adapter_module": "/model-static/adapter.js?v=model-ui-v4" } adapter_response = await client.get("/model-static/adapter.js") assert adapter_response.status == 200 @@ -385,7 +390,7 @@ async def test_packaged_webrtc_app_serves_model_stylesheet(tmp_path) -> None: config_response = await client.get("/api/ui/config") assert await config_response.json() == { "adapter_module": None, - "model_stylesheet": "/model-static/adapter.css?v=model-ui-v2", + "model_stylesheet": "/model-static/adapter.css?v=model-ui-v4", } stylesheet_response = await client.get("/model-static/adapter.css") assert stylesheet_response.status == 200 diff --git a/integrations/flashvsr/README.md b/integrations/flashvsr/README.md index 9a0ea644b..548476463 100644 --- a/integrations/flashvsr/README.md +++ b/integrations/flashvsr/README.md @@ -52,6 +52,51 @@ export HF_TOKEN= export HF_HOME=~/.cache/huggingface # default ``` +## Runtime API demos + +The `flashvsr-demo` command uses a native `InferenceRuntime` and +`InferenceSession`. Its model input provider decodes the source video and +supplies each cold/steady frame chunk through `InferenceInput.step`; the +legacy runner, postprocessor, and gRPC session code are not involved. + +Run the full model with no output artifact: + +```bash +uv run flashvsr-demo replay \ + --input /path/to/input.mp4 \ + --output-mode null +``` + +Write the upscaled result to MP4: + +```bash +uv run flashvsr-demo replay \ + --input /path/to/input.mp4 \ + --output-mode mp4 \ + --output /tmp/flashvsr-output.mp4 +``` + +For a quick eager smoke test, add `--chunk-size 8 --no-compile +--no-cuda-graph --color-corrector torch`. + +Start the shared WebRTC viewer without a server-side input: + +```bash +uv run flashvsr-demo webrtc \ + --host 0.0.0.0 \ + --port 8082 +``` + +Open `http://:8082/request_session`, choose an MP4 in the **Input +Video** field, then connect. The browser uploads and decodes the video before +negotiating the WebRTC session. When the control channel is ready, click **Start +FlashVSR** to begin playback. Model construction is deferred until the uploaded +video supplies its dimensions. + +`--input /path/to/input.mp4` remains available as an optional server-side +fallback; choosing a browser file overrides it for the next session. Use +`--no-loop-input` for one finite pass. + ## Single-GPU Run Once installed, the slug is discovered automatically by `flashdreams-run`: diff --git a/integrations/flashvsr/flashvsr/config.py b/integrations/flashvsr/flashvsr/config.py index 2b158511e..22f28b038 100644 --- a/integrations/flashvsr/flashvsr/config.py +++ b/integrations/flashvsr/flashvsr/config.py @@ -231,13 +231,13 @@ def build_flashvsr_v1_1( scale=scale, projector_checkpoint_path=checkpoint_path["encoder"], use_compile=compile_network, - use_cuda_graph=True, + use_cuda_graph=use_cuda_graph, dtype=dtype, ), decoder=FlashVSRDecoderConfig( tcdecoder_checkpoint_path=checkpoint_path["decoder"], use_compile=compile_network, - use_cuda_graph=True, + use_cuda_graph=use_cuda_graph, color_corrector_implementation=color_corrector_implementation, dtype=dtype, ), diff --git a/integrations/flashvsr/flashvsr/corrector.py b/integrations/flashvsr/flashvsr/corrector.py index 3eac071a7..0c5384ade 100644 --- a/integrations/flashvsr/flashvsr/corrector.py +++ b/integrations/flashvsr/flashvsr/corrector.py @@ -87,9 +87,10 @@ def _calc_mean_std( ) -> Tuple[torch.Tensor, torch.Tensor]: assert feat.dim() == 4, "feat must be (N, C, H, W)" N, C = feat.shape[:2] - var = feat.view(N, C, -1).var(dim=2, unbiased=False) + eps - std = var.sqrt().view(N, C, 1, 1) - mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1) + flat = feat.reshape(N, C, -1) + var = flat.var(dim=2, unbiased=False) + eps + std = var.sqrt().reshape(N, C, 1, 1) + mean = flat.mean(dim=2).reshape(N, C, 1, 1) return mean, std diff --git a/integrations/flashvsr/flashvsr/demo/__init__.py b/integrations/flashvsr/flashvsr/demo/__init__.py new file mode 100644 index 000000000..d64b059e4 --- /dev/null +++ b/integrations/flashvsr/flashvsr/demo/__init__.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FlashVSR demos built on the shared inference runtime API.""" + +from flashvsr.demo.adapter import FlashVSRDemoAdapter +from flashvsr.demo.providers import FlashVSRVideoInputProvider +from flashvsr.demo.spec import ( + DEFAULT_FLASHVSR_INPUT_URL, + FlashVSRVideoScenario, + PreparedFlashVSRVideo, +) +from flashvsr.runtime import ( + DEFAULT_FLASHVSR_PRESET, + FLASHVSR_MODEL_ID, +) + +__all__ = [ + "DEFAULT_FLASHVSR_INPUT_URL", + "DEFAULT_FLASHVSR_PRESET", + "FLASHVSR_MODEL_ID", + "FlashVSRDemoAdapter", + "FlashVSRVideoInputProvider", + "FlashVSRVideoScenario", + "PreparedFlashVSRVideo", +] diff --git a/integrations/flashvsr/flashvsr/demo/adapter.py b/integrations/flashvsr/flashvsr/demo/adapter.py new file mode 100644 index 000000000..1637b6c39 --- /dev/null +++ b/integrations/flashvsr/flashvsr/demo/adapter.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FlashVSR model/demo adapter for shared replay and WebRTC run modes.""" + +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path +from typing import Any + +from flashdreams.runtime import InferenceInput, InputCanonicalizer, UserInputSchema +from flashdreams.runtime.demo import DemoSpec, PreparedScenario, WebRTCOutputSpec +from flashvsr.runtime import ( + FIELD_CHUNK_SIZE, + FIELD_FPS, + FIELD_INPUT_HEIGHT, + FIELD_INPUT_WIDTH, + FIELD_TAIL_POLICY, + FIELD_TOTAL_FRAMES, + FlashVSRModelAdapter, +) + +from .providers import PREPARED_VIDEO_METADATA_KEY, FlashVSRVideoInputProvider +from .spec import ( + FlashVSRVideoScenario, + PreparedFlashVSRVideo, + prepare_video_source, + resolve_video_scenario, +) + + +class FlashVSRDemoAdapter(FlashVSRModelAdapter): + """Prepare decoded videos for native runtime API demo execution.""" + + def supported_input_modes(self) -> tuple[str, ...]: + return ("replay",) + + def supported_output_modes(self) -> tuple[str, ...]: + return ("mp4", "null", "webrtc") + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + if spec.input_mode != "replay": + raise ValueError( + "FlashVSR demos require input_mode='replay'; WebRTC replays or " + "loops that decoded source through the realtime transport." + ) + if spec.output.mode not in self.supported_output_modes(): + raise ValueError(f"Unsupported FlashVSR output mode: {spec.output.mode!r}.") + if spec.config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + self.validate_config(spec.config) + + if isinstance(spec.scenario, PreparedFlashVSRVideo): + prepared = spec.scenario + video_scenario = prepared.scenario + else: + video_scenario = resolve_video_scenario(spec.scenario) + prepared = prepare_video_source( + video_scenario, + scale=self._scale_for_spec(spec), + ) + if video_scenario.loop_input and not isinstance(spec.output, WebRTCOutputSpec): + raise ValueError( + "FlashVSR loop_input is supported only with WebRTC output." + ) + + global_conditioning: dict[str, Any] = { + FIELD_INPUT_HEIGHT: prepared.input_height, + FIELD_INPUT_WIDTH: prepared.input_width, + FIELD_FPS: prepared.fps, + FIELD_CHUNK_SIZE: video_scenario.chunk_size, + FIELD_TAIL_POLICY: video_scenario.tail_policy, + } + if not video_scenario.loop_input: + global_conditioning[FIELD_TOTAL_FRAMES] = prepared.total_frames + return PreparedScenario( + initial_inputs=InferenceInput(global_conditioning=global_conditioning), + source_schema=UserInputSchema(description="decoded FlashVSR video source"), + canonicalizer=InputCanonicalizer(), + mapping=self.default_input_mapping(), + metadata={ + PREPARED_VIDEO_METADATA_KEY: prepared, + "model_id": self.model_id, + "preset_id": self.preset_id(spec.config), + "input_path": str(prepared.resolved_path), + "target_height": prepared.target_height, + "target_width": prepared.target_width, + }, + ) + + def prepare_uploaded_video( + self, + spec: DemoSpec, + *, + upload_path: Path, + original_name: str, + ) -> PreparedFlashVSRVideo: + """Decode one uploaded WebRTC video at the server playback rate. + + Args: + spec: Base WebRTC demo specification. + upload_path: Server-generated temporary upload path. + original_name: Sanitized browser filename used only for metadata. + + Returns: + Decoded CPU video ready for a model input provider. + + Raises: + ValueError: The specification is not WebRTC output. + """ + if not isinstance(spec.output, WebRTCOutputSpec): + raise ValueError("FlashVSR uploads require WebRTC output.") + scenario = resolve_video_scenario(spec.scenario) + uploaded_scenario = replace( + scenario, + input_path=upload_path, + fps=float(spec.output.fps), + ) + prepared = prepare_video_source( + uploaded_scenario, + scale=self._scale_for_spec(spec), + ) + display_scenario = replace( + uploaded_scenario, + input_path=original_name, + ) + return replace( + prepared, + scenario=display_scenario, + resolved_path=Path(original_name), + ) + + def create_model_input_provider( + self, + spec: DemoSpec, + scenario: PreparedScenario, + ) -> FlashVSRVideoInputProvider: + del spec + return FlashVSRVideoInputProvider( + scenario=scenario, + inference_input_schema=self.inference_input_schema, + ) + + def _scale_for_spec(self, spec: DemoSpec) -> int: + config = spec.config + if config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + pipeline_config = self.pipeline_config(config) + configured_scale = getattr( + getattr(pipeline_config, "encoder", None), + "scale", + 2, + ) + return int(config.runtime_options.get("scale", configured_scale)) + + +__all__ = ["FlashVSRDemoAdapter", "FlashVSRVideoScenario"] diff --git a/integrations/flashvsr/flashvsr/demo/app.py b/integrations/flashvsr/flashvsr/demo/app.py new file mode 100644 index 000000000..c5b4a0001 --- /dev/null +++ b/integrations/flashvsr/flashvsr/demo/app.py @@ -0,0 +1,317 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Command-line entry point for native FlashVSR runtime API demos.""" + +from __future__ import annotations + +import argparse +import math +from pathlib import Path +from typing import Any + +from loguru import logger + +from flashdreams.infra.runner_io import read_video_fps, resolve_input_path +from flashdreams.runtime import InferenceConfig +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + NullOutputSpec, + WebRTCOutputSpec, +) +from flashdreams.runtime.demo.app import DemoApplication +from flashvsr.runtime import DEFAULT_FLASHVSR_PRESET, FLASHVSR_MODEL_ID + +from .adapter import FlashVSRDemoAdapter +from .spec import ( + DEFAULT_FLASHVSR_INPUT_URL, + FLASHVSR_INPUT_CACHE_DIR, + FlashVSRVideoScenario, +) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + """Parse replay or WebRTC demo arguments.""" + parser = argparse.ArgumentParser( + description="FlashVSR demos using the native flashdreams.runtime API." + ) + subparsers = parser.add_subparsers(dest="command", required=True) + + replay = subparsers.add_parser("replay", help="Upscale a finite MP4 input.") + _add_model_arguments( + replay, + default_device="cuda", + default_input=DEFAULT_FLASHVSR_INPUT_URL, + ) + replay.add_argument("--output-mode", choices=("mp4", "null"), default="mp4") + replay.add_argument("--output", type=Path, default=None) + + webrtc = subparsers.add_parser( + "webrtc", + help="Upload an MP4 in the browser and stream the upscaled result.", + ) + _add_model_arguments( + webrtc, + default_device="cuda:0", + default_input=None, + ) + webrtc.add_argument("--host", default="0.0.0.0") + webrtc.add_argument("--port", type=int, default=8082) + webrtc.add_argument( + "--loop-input", + action=argparse.BooleanOptionalAction, + default=True, + help="Loop the selected source until the browser disconnects (default: true).", + ) + webrtc.add_argument("--warmup-chunks", type=int, default=1) + webrtc.add_argument("--warmup-timeout-s", type=float, default=600.0) + webrtc.add_argument("--client-liveness-timeout-s", type=float, default=30.0) + webrtc.add_argument( + "--prefer-sw-encoder", + action="store_true", + help=( + "Use aiortc's software encoder. Browser-upload mode already selects " + "this resolution-agnostic backend." + ), + ) + + args = parser.parse_args(argv) + if args.command == "replay": + if args.output_mode == "mp4" and args.output is None: + parser.error("replay --output is required when --output-mode=mp4.") + if args.output_mode == "null" and args.output is not None: + parser.error("replay --output is valid only when --output-mode=mp4.") + return args + + +def _add_model_arguments( + parser: argparse.ArgumentParser, + *, + default_device: str, + default_input: str | None, +) -> None: + parser.add_argument("--preset-id", default=DEFAULT_FLASHVSR_PRESET) + parser.add_argument( + "--input", + "--input-path", + dest="input_path", + default=default_input, + help=( + "Optional server-side input video. WebRTC can instead upload an MP4 " + "from the browser." + ), + ) + parser.add_argument("--device", default=default_device) + parser.add_argument("--chunk-size", type=int, choices=(8, 16), default=16) + parser.add_argument( + "--fps", + type=float, + default=None, + help=( + "Override playback FPS. Defaults to source metadata, or 30 for " + "upload-only WebRTC startup." + ), + ) + parser.add_argument("--scale", type=int, choices=(2, 4), default=2) + parser.add_argument( + "--crop-region", + choices=("none", "bottom_half", "top_half"), + default="none", + ) + parser.add_argument("--tail-policy", choices=("drop", "pad"), default="drop") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--compile", + action=argparse.BooleanOptionalAction, + default=None, + help="Override the preset's torch.compile setting.", + ) + parser.add_argument( + "--cuda-graph", + action=argparse.BooleanOptionalAction, + default=None, + help="Override CUDA graphs for the encoder, DiT, and decoder.", + ) + parser.add_argument( + "--color-corrector", + choices=("cuda", "torch"), + default=None, + help="Override the preset's color-correction implementation.", + ) + + +class FlashVSRDemoApplication(DemoApplication): + """Dispatch FlashVSR finite replay and realtime WebRTC demos.""" + + 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) -> FlashVSRDemoAdapter: + return FlashVSRDemoAdapter() + + def serve_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: + from .webrtc import serve_flashvsr_webrtc_demo + + serve_flashvsr_webrtc_demo( + spec=_webrtc_spec(args, device=str(context.device)), + world_rank=context.world_rank, + ) + + +_APPLICATION = FlashVSRDemoApplication() + + +def main(argv: list[str] | None = None) -> None: + """Run the FlashVSR demo application.""" + _APPLICATION.main(argv) + + +def _replay_spec(args: argparse.Namespace) -> DemoSpec: + fps = _resolve_demo_fps(args, webrtc=False) + return DemoSpec( + model_id=FLASHVSR_MODEL_ID, + preset_id=args.preset_id, + input_mode="replay", + scenario=_video_scenario(args, loop_input=False, fps=fps), + output=_replay_output_spec(args, fps=fps), + config=_inference_config(args, device=args.device, fps=fps), + ) + + +def _replay_output_spec( + args: argparse.Namespace, + *, + fps: float, +) -> Mp4OutputSpec | NullOutputSpec: + if args.output_mode == "null": + return NullOutputSpec() + if args.output is None: + raise ValueError("FlashVSR MP4 replay requires --output.") + return Mp4OutputSpec( + path=args.output, + fps=fps, + output_layout="bcthw", + ) + + +def _webrtc_spec(args: argparse.Namespace, *, device: str) -> DemoSpec: + fps = _resolve_demo_fps(args, webrtc=True) + return DemoSpec( + model_id=FLASHVSR_MODEL_ID, + preset_id=args.preset_id, + input_mode="replay", + scenario=_video_scenario(args, loop_input=args.loop_input, fps=fps), + # Upload-only startup has no resolution yet. The adapter replaces + # these placeholders with the decoded target dimensions per session. + output=WebRTCOutputSpec( + host=args.host, + port=args.port, + fps=int(fps), + video_width=128, + video_height=128, + warmup_chunks=args.warmup_chunks, + warmup_timeout_s=args.warmup_timeout_s, + client_liveness_timeout_s=args.client_liveness_timeout_s, + preload_name="FlashVSR", + ), + config=_inference_config( + args, + device=device, + fps=fps, + extra_options={"prefer_sw_encoder": args.prefer_sw_encoder}, + ), + ) + + +def _video_scenario( + args: argparse.Namespace, + *, + loop_input: bool, + fps: float, +) -> FlashVSRVideoScenario: + return FlashVSRVideoScenario( + input_path=args.input_path, + chunk_size=args.chunk_size, + fps=fps, + crop_region=args.crop_region, + tail_policy=args.tail_policy, + loop_input=loop_input, + ) + + +def _inference_config( + args: argparse.Namespace, + *, + device: str, + fps: float, + extra_options: dict[str, Any] | None = None, +) -> InferenceConfig: + runtime_options: dict[str, Any] = { + "fps": fps, + "chunk_size": args.chunk_size, + "scale": args.scale, + } + if args.cuda_graph is not None: + runtime_options["use_cuda_graph"] = args.cuda_graph + if args.color_corrector is not None: + runtime_options["color_corrector_implementation"] = args.color_corrector + if extra_options: + runtime_options.update(extra_options) + return InferenceConfig( + model_id=FLASHVSR_MODEL_ID, + preset_id=args.preset_id, + device=device, + seed=args.seed, + compile=args.compile, + runtime_options=runtime_options, + ) + + +def _resolve_demo_fps( + args: argparse.Namespace, + *, + webrtc: bool, +) -> float: + fps = args.fps + if fps is None and args.input_path is None: + fps = 30.0 + elif fps is None: + resolved_path = resolve_input_path( + args.input_path, + cache_dir=FLASHVSR_INPUT_CACHE_DIR, + ) + try: + fps = float(read_video_fps(resolved_path)) + except Exception: + logger.warning("Could not read input fps; using 30 fps.") + fps = 30.0 + fps = float(fps) + if fps <= 0: + raise ValueError("FlashVSR fps must be > 0.") + if not webrtc: + return fps + transport_fps = int(round(fps)) + if transport_fps <= 0: + raise ValueError("FlashVSR WebRTC fps must round to at least 1.") + if not math.isclose(fps, transport_fps, rel_tol=0.0, abs_tol=1e-6): + logger.warning( + "WebRTC requires integer fps; using {} instead of {}.", + transport_fps, + fps, + ) + return float(transport_fps) + + +if __name__ == "__main__": + main() + + +__all__ = [ + "FlashVSRDemoApplication", + "main", + "parse_args", +] diff --git a/integrations/flashvsr/flashvsr/demo/providers.py b/integrations/flashvsr/flashvsr/demo/providers.py new file mode 100644 index 000000000..d0aa01b35 --- /dev/null +++ b/integrations/flashvsr/flashvsr/demo/providers.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""New-API model input provider for decoded FlashVSR videos.""" + +from __future__ import annotations + +import torch + +from flashdreams.runtime import InferenceInput, InferenceInputSchema +from flashdreams.runtime.demo import ( + ControlDecision, + PreparedScenario, + PreparedStep, + ProviderCapabilities, + UserInputWindow, +) +from flashdreams.runtime.types import StepRequirements +from flashvsr.runtime import ( + FIELD_VALID_FRAME_COUNT, + FIELD_VIDEO_CHUNK, +) + +from .spec import PreparedFlashVSRVideo + +PREPARED_VIDEO_METADATA_KEY = "prepared_video" + + +class FlashVSRVideoInputProvider: + """Slice one decoded video into exact model-requested frame chunks.""" + + def __init__( + self, + *, + scenario: PreparedScenario, + inference_input_schema: InferenceInputSchema, + ) -> None: + prepared = scenario.metadata.get(PREPARED_VIDEO_METADATA_KEY) + if not isinstance(prepared, PreparedFlashVSRVideo): + raise TypeError( + "FlashVSR prepared scenario is missing its decoded video source." + ) + self.prepared = prepared + self._initial_input = scenario.initial_inputs + self._cursor = 0 + self._closed = False + self.capabilities = ProviderCapabilities( + supports_realtime_clock=True, + supports_recorded_input=True, + supports_reset=True, + deterministic_given_inputs=True, + user_input_schema=scenario.source_schema, + inference_input_schema=inference_input_schema, + ) + + def prepare_initial_input(self) -> InferenceInput: + self._require_open() + return self._initial_input + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + del user_window + self._require_open() + requested = request.input_frame_count + expected_valid = int(request.metadata.get(FIELD_VALID_FRAME_COUNT, requested)) + source_start = self._cursor + if self.prepared.scenario.loop_input: + chunk = self._looping_chunk(requested) + valid = requested + else: + chunk, valid = self._finite_chunk(requested) + if valid <= 0: + return PreparedStep( + control=ControlDecision( + close_session=True, + reason="FlashVSR source video is exhausted.", + ) + ) + if valid != expected_valid: + raise RuntimeError( + "FlashVSR provider/session frame-count disagreement: " + f"session requested {expected_valid} valid frames, provider has {valid}." + ) + return PreparedStep( + inference_input=InferenceInput( + step={FIELD_VIDEO_CHUNK: chunk}, + metadata={ + FIELD_VALID_FRAME_COUNT: valid, + "source_frame_start": source_start, + }, + ) + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self._require_open() + self._cursor = 0 + + def close(self) -> None: + self._closed = True + + def _finite_chunk(self, requested: int) -> tuple[torch.Tensor, int]: + video = self.prepared.video + total = self.prepared.total_frames + valid = min(requested, total - self._cursor) + if valid <= 0: + return video[:, :, :0], 0 + chunk = video[:, :, self._cursor : self._cursor + valid] + self._cursor += valid + if valid == requested: + return chunk, valid + if self.prepared.scenario.tail_policy != "pad": + return chunk, valid + padding = chunk[:, :, -1:].expand(-1, -1, requested - valid, -1, -1) + return torch.cat((chunk, padding), dim=2), valid + + def _looping_chunk(self, requested: int) -> torch.Tensor: + video = self.prepared.video + total = self.prepared.total_frames + indices = (torch.arange(requested) + self._cursor) % total + chunk = video.index_select(2, indices) + self._cursor = (self._cursor + requested) % total + return chunk + + def _require_open(self) -> None: + if self._closed: + raise RuntimeError("FlashVSR video input provider is closed.") + + +__all__ = [ + "PREPARED_VIDEO_METADATA_KEY", + "FlashVSRVideoInputProvider", +] diff --git a/integrations/flashvsr/flashvsr/demo/server.py b/integrations/flashvsr/flashvsr/demo/server.py new file mode 100644 index 000000000..00ddb1016 --- /dev/null +++ b/integrations/flashvsr/flashvsr/demo/server.py @@ -0,0 +1,281 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""HTTP upload routes for FlashVSR WebRTC sessions.""" + +from __future__ import annotations + +import asyncio +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Protocol +from urllib.parse import unquote + +from aiohttp import web +from aiohttp.multipart import BodyPartReader + +from flashdreams.runtime.demo import DemoSpec +from flashdreams.serving.webrtc.server import SessionBusyError + +from .adapter import FlashVSRDemoAdapter +from .spec import PreparedFlashVSRVideo + +MAX_UPLOAD_VIDEO_BYTES = 512 * 1024 * 1024 +"""Maximum accepted MP4 upload size.""" + +_UPLOAD_ROUTE = "/api/session/input" +_OFFER_ROUTE = "/api/webrtc/offer" +_ACCEPTED_VIDEO_CONTENT_TYPES = frozenset( + { + "application/mp4", + "application/octet-stream", + "video/mp4", + } +) + + +class FlashVSRSessionManager(Protocol): + """Manager surface needed by the upload controller.""" + + @property + def pending_session_input(self) -> Any: ... + + def has_active_session(self) -> bool: ... + + def set_pending_session_input(self, session_input: Any) -> None: ... + + +@dataclass(frozen=True, slots=True) +class FlashVSRWebRTCSessionInput: + """Decoded input staged for one WebRTC offer.""" + + prepared_video: PreparedFlashVSRVideo + """CPU video consumed by the native model input provider.""" + + original_name: str + """Sanitized browser filename used for display metadata.""" + + +class FlashVSRUploadController: + """Validate, decode, and stage uploaded videos for the next session.""" + + def __init__( + self, + *, + manager: FlashVSRSessionManager, + adapter: FlashVSRDemoAdapter, + spec: DemoSpec, + default_video: PreparedFlashVSRVideo | None, + ) -> None: + self.manager = manager + self.adapter = adapter + self.spec = spec + self.default_video = default_video + + def has_session_input(self) -> bool: + """Return whether an upload or server-side fallback is available.""" + return ( + isinstance( + self.manager.pending_session_input, + FlashVSRWebRTCSessionInput, + ) + or self.default_video is not None + ) + + def status_payload(self) -> dict[str, Any]: + """Return browser-facing input availability and video metadata.""" + pending = self.manager.pending_session_input + if isinstance(pending, FlashVSRWebRTCSessionInput): + return { + "upload_required": False, + "has_default_input": self.default_video is not None, + "input_source": "uploaded", + "filename": pending.original_name, + **_video_metadata(pending.prepared_video), + } + payload: dict[str, Any] = { + "upload_required": self.default_video is None, + "has_default_input": self.default_video is not None, + "input_source": "server" if self.default_video is not None else None, + } + if self.default_video is not None: + payload.update(_video_metadata(self.default_video)) + return payload + + def stage_uploaded_video( + self, + *, + upload_path: Path, + original_name: str, + ) -> dict[str, Any]: + """Decode an uploaded MP4 and stage it for the next WebRTC offer.""" + if self.manager.has_active_session(): + raise SessionBusyError("A FlashVSR session is already active.") + prepared = self.adapter.prepare_uploaded_video( + self.spec, + upload_path=upload_path, + original_name=original_name, + ) + session_input = FlashVSRWebRTCSessionInput( + prepared_video=prepared, + original_name=original_name, + ) + self.manager.set_pending_session_input(session_input) + return { + "upload_required": False, + "has_default_input": self.default_video is not None, + "input_source": "uploaded", + "filename": original_name, + **_video_metadata(prepared), + } + + +FLASHVSR_UPLOAD_CONTROLLER_KEY = web.AppKey( + "flashvsr_upload_controller", + FlashVSRUploadController, +) + + +def configure_flashvsr_webrtc_app( + app: web.Application, + *, + controller: FlashVSRUploadController, +) -> None: + """Register FlashVSR upload routes and offer validation.""" + app[FLASHVSR_UPLOAD_CONTROLLER_KEY] = controller + app.router.add_get(_UPLOAD_ROUTE, _session_input_status) + app.router.add_post(_UPLOAD_ROUTE, _session_input_upload) + app.middlewares.append(_require_session_input) + + +@web.middleware +async def _require_session_input( + request: web.Request, + handler: Any, +) -> web.StreamResponse: + controller = request.app[FLASHVSR_UPLOAD_CONTROLLER_KEY] + if ( + request.path == _OFFER_ROUTE + and not controller.manager.has_active_session() + and not controller.has_session_input() + ): + raise web.HTTPBadRequest( + reason="Upload an MP4 before connecting the FlashVSR session." + ) + return await handler(request) + + +async def _session_input_status(request: web.Request) -> web.StreamResponse: + controller = request.app[FLASHVSR_UPLOAD_CONTROLLER_KEY] + return web.json_response(controller.status_payload()) + + +async def _session_input_upload(request: web.Request) -> web.StreamResponse: + if not request.content_type.startswith("multipart/"): + raise web.HTTPBadRequest(reason="Expected a multipart MP4 upload.") + try: + reader = await request.multipart() + except Exception as exc: + raise web.HTTPBadRequest(reason="Expected a multipart MP4 upload.") from exc + + video_field: BodyPartReader | None = None + while True: + field = await reader.next() + if field is None: + break + if not isinstance(field, BodyPartReader): + continue + if field.name != "video": + await field.release() + continue + if video_field is not None: + raise web.HTTPBadRequest(reason="Upload exactly one MP4 video.") + video_field = field + break + + if video_field is None or not video_field.filename: + raise web.HTTPBadRequest(reason="Upload an MP4 in the 'video' field.") + original_name = _sanitize_filename(video_field.filename) + if Path(original_name).suffix.lower() != ".mp4": + raise web.HTTPBadRequest(reason="Uploaded video must use the .mp4 extension.") + content_type = ( + video_field.headers.get( + "Content-Type", + "application/octet-stream", + ) + .partition(";")[0] + .strip() + .lower() + ) + if content_type not in _ACCEPTED_VIDEO_CONTENT_TYPES: + raise web.HTTPBadRequest(reason="Uploaded video must be an MP4.") + + controller = request.app[FLASHVSR_UPLOAD_CONTROLLER_KEY] + if controller.manager.has_active_session(): + raise web.HTTPConflict(reason="A FlashVSR session is already active.") + + with tempfile.TemporaryDirectory(prefix="flashvsr-upload-") as temp_dir: + upload_path = Path(temp_dir) / "input.mp4" + await _stream_uploaded_video(video_field, upload_path) + try: + payload = await asyncio.to_thread( + controller.stage_uploaded_video, + upload_path=upload_path, + original_name=original_name, + ) + except SessionBusyError as exc: + raise web.HTTPConflict(reason=str(exc)) from exc + except Exception as exc: + raise web.HTTPBadRequest(reason=f"Invalid uploaded MP4: {exc}") from exc + return web.json_response(payload) + + +async def _stream_uploaded_video( + field: BodyPartReader, + upload_path: Path, +) -> None: + total_bytes = 0 + with upload_path.open("wb") as stream: + while True: + chunk = await field.read_chunk(size=1024 * 1024) + if not chunk: + break + total_bytes += len(chunk) + if total_bytes > MAX_UPLOAD_VIDEO_BYTES: + raise web.HTTPRequestEntityTooLarge( + max_size=MAX_UPLOAD_VIDEO_BYTES, + actual_size=total_bytes, + ) + stream.write(chunk) + if total_bytes == 0: + raise web.HTTPBadRequest(reason="Uploaded MP4 is empty.") + + +def _sanitize_filename(value: str) -> str: + filename = Path(unquote(value).replace("\\", "/")).name.strip() + return filename or "upload.mp4" + + +def _video_metadata(prepared: PreparedFlashVSRVideo) -> dict[str, Any]: + return { + "fps": prepared.fps, + "num_frames": prepared.total_frames, + "input_resolution": { + "width": prepared.input_width, + "height": prepared.input_height, + }, + "resolution": { + "width": prepared.target_width, + "height": prepared.target_height, + }, + } + + +__all__ = [ + "FLASHVSR_UPLOAD_CONTROLLER_KEY", + "MAX_UPLOAD_VIDEO_BYTES", + "FlashVSRUploadController", + "FlashVSRWebRTCSessionInput", + "configure_flashvsr_webrtc_app", +] diff --git a/integrations/flashvsr/flashvsr/demo/spec.py b/integrations/flashvsr/flashvsr/demo/spec.py new file mode 100644 index 000000000..93ec13d64 --- /dev/null +++ b/integrations/flashvsr/flashvsr/demo/spec.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed video-source specifications for FlashVSR demos.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal, cast + +import torch +from loguru import logger + +from flashdreams.infra.runner_io import ( + read_video_fps, + read_video_rgb, + resolve_input_path, + rgb_video_to_normalized_tensor, +) +from flashvsr.runtime import TailPolicy + +DEFAULT_FLASHVSR_INPUT_URL = ( + "https://raw.githubusercontent.com/OpenImagingLab/FlashVSR/main/" + "examples/WanVSR/inputs/example1.mp4" +) +FLASHVSR_INPUT_CACHE_DIR = ( + Path(os.path.expanduser(os.getenv("FLASHDREAMS_CACHE_DIR", "~/.cache/flashdreams"))) + / "flashvsr" +) + +CropRegion = Literal["none", "bottom_half", "top_half"] + + +@dataclass(frozen=True, kw_only=True, slots=True) +class FlashVSRVideoScenario: + """User-facing source-video and chunking options.""" + + input_path: str | Path | None = DEFAULT_FLASHVSR_INPUT_URL + """Optional server-side source; a missing value requires a WebRTC upload.""" + + chunk_size: Literal[8, 16] = 16 + fps: float | None = None + crop_region: CropRegion = "none" + tail_policy: TailPolicy = "drop" + loop_input: bool = False + + def __post_init__(self) -> None: + if self.chunk_size not in {8, 16}: + raise ValueError("FlashVSR scenario chunk_size must be 8 or 16.") + if self.fps is not None and self.fps <= 0: + raise ValueError("FlashVSR scenario fps must be > 0 when provided.") + if self.crop_region not in {"none", "bottom_half", "top_half"}: + raise ValueError( + "FlashVSR crop_region must be 'none', 'bottom_half', or 'top_half'." + ) + if self.tail_policy not in {"drop", "pad"}: + raise ValueError("FlashVSR tail_policy must be 'drop' or 'pad'.") + + +@dataclass(frozen=True, slots=True, eq=False) +class PreparedFlashVSRVideo: + """Decoded CPU source video plus model/output shape facts.""" + + scenario: FlashVSRVideoScenario + resolved_path: Path + video: torch.Tensor + input_height: int + input_width: int + target_height: int + target_width: int + fps: float + + @property + def total_frames(self) -> int: + """Return the number of decoded source frames.""" + return int(self.video.shape[2]) + + +def resolve_video_scenario(value: Any) -> FlashVSRVideoScenario: + """Normalize a public demo scenario into a FlashVSR video scenario.""" + if value is None: + return FlashVSRVideoScenario() + if isinstance(value, FlashVSRVideoScenario): + return value + if isinstance(value, str | Path): + return FlashVSRVideoScenario(input_path=value) + if not isinstance(value, Mapping): + raise TypeError( + "FlashVSR scenario must be a path, mapping, FlashVSRVideoScenario, or None." + ) + return FlashVSRVideoScenario( + input_path=value.get( + "input_path", value.get("input", DEFAULT_FLASHVSR_INPUT_URL) + ), + chunk_size=cast(Literal[8, 16], int(value.get("chunk_size", 16))), + fps=(None if value.get("fps") is None else float(value["fps"])), + crop_region=cast(CropRegion, str(value.get("crop_region", "none"))), + tail_policy=cast(TailPolicy, str(value.get("tail_policy", "drop"))), + loop_input=bool(value.get("loop_input", False)), + ) + + +def prepare_video_source( + scenario: FlashVSRVideoScenario, + *, + scale: int, +) -> PreparedFlashVSRVideo: + """Resolve, decode, crop, and normalize one low-resolution input video.""" + if scenario.input_path is None: + raise ValueError( + "No FlashVSR input video is configured. Upload an MP4 in the " + "WebRTC UI or launch with --input." + ) + resolved_path = resolve_input_path( + scenario.input_path, + cache_dir=FLASHVSR_INPUT_CACHE_DIR, + ) + if not resolved_path.is_file(): + raise FileNotFoundError( + f"FlashVSR input video does not exist: {scenario.input_path!r}." + ) + logger.info("Reading FlashVSR demo input {}.", resolved_path) + video_rgb = read_video_rgb(resolved_path) + if video_rgb.ndim != 4 or video_rgb.shape[-1] != 3: + raise ValueError( + "FlashVSR input decoder must produce [T,H,W,3] RGB frames, " + f"got {tuple(video_rgb.shape)}." + ) + if video_rgb.shape[0] <= 0: + raise ValueError("FlashVSR input video contains no frames.") + if scenario.crop_region != "none": + height = int(video_rgb.shape[1]) + half = height // 2 + if half <= 0: + raise ValueError("FlashVSR input is too short to crop vertically.") + if scenario.crop_region == "bottom_half": + video_rgb = video_rgb[:, height - half :, :, :] + else: + video_rgb = video_rgb[:, :half, :, :] + + _, height, width, _ = video_rgb.shape + target_height = (height * scale // 128) * 128 + target_width = (width * scale // 128) * 128 + if target_height <= 0 or target_width <= 0: + raise ValueError( + "FlashVSR input is too small after cropping: " + f"input={height}x{width}, scale={scale}; each scaled axis must be >= 128." + ) + fps = scenario.fps + if fps is None: + try: + fps = float(read_video_fps(resolved_path)) + except Exception: + logger.warning("Could not read input fps; using 30 fps.") + fps = 30.0 + video = ( + rgb_video_to_normalized_tensor( + video_rgb, + device=torch.device("cpu"), + dtype=torch.float32, + ) + .permute(1, 0, 2, 3) + .unsqueeze(0) + ) + cold_frames = 5 if scenario.chunk_size == 8 else 13 + if ( + not scenario.loop_input + and scenario.tail_policy == "drop" + and video.shape[2] < cold_frames + ): + raise ValueError( + f"FlashVSR input has {video.shape[2]} frames; chunk_size=" + f"{scenario.chunk_size} needs at least {cold_frames}." + ) + return PreparedFlashVSRVideo( + scenario=scenario, + resolved_path=resolved_path, + video=video.contiguous(), + input_height=height, + input_width=width, + target_height=target_height, + target_width=target_width, + fps=float(fps), + ) + + +__all__ = [ + "CropRegion", + "DEFAULT_FLASHVSR_INPUT_URL", + "FLASHVSR_INPUT_CACHE_DIR", + "FlashVSRVideoScenario", + "PreparedFlashVSRVideo", + "prepare_video_source", + "resolve_video_scenario", +] diff --git a/integrations/flashvsr/flashvsr/demo/web/adapter.css b/integrations/flashvsr/flashvsr/demo/web/adapter.css new file mode 100644 index 000000000..4c90b7c9a --- /dev/null +++ b/integrations/flashvsr/flashvsr/demo/web/adapter.css @@ -0,0 +1,151 @@ +/* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. */ +/* SPDX-License-Identifier: Apache-2.0 */ + +.stageVideo { + inset: 50% auto auto 50%; + width: 100vw; + height: 100vh; + transform: translate(-50%, -50%); + object-fit: contain; + object-position: center; +} + +.flashvsrUploadCard { + position: absolute; + top: clamp(120px, 15vh, 168px); + left: clamp(18px, 3vw, 52px); + display: grid; + gap: 1rem; + width: min(31rem, calc(100vw - 36px)); + max-height: min(440px, calc(100vh - 224px)); + padding: 1.25rem; + overflow: auto; +} + +.flashvsrUploadHeader { + display: flex; + align-items: start; + justify-content: space-between; + gap: 1rem; +} + +.flashvsrUploadHeader h2 { + margin: 0.3rem 0 0; + color: var(--text); + font-size: 1.25rem; + line-height: 1.2; +} + +.flashvsrUploadBadge { + padding: 0.32rem 0.55rem; + border: 1px solid rgba(99, 216, 255, 0.36); + border-radius: 999px; + background: rgba(99, 216, 255, 0.10); + color: var(--cyan); + font-size: 0.7rem; + font-weight: 800; + letter-spacing: 0.08em; +} + +.flashvsrUploadHint { + margin: 0; + color: var(--muted); + font-size: 0.85rem; + line-height: 1.45; +} + +.flashvsrUploadControl { + display: grid; + gap: 0.55rem; + padding: 0.85rem; + border: 1px solid rgba(255, 255, 255, 0.13); + border-radius: 8px; + background: rgba(255, 255, 255, 0.045); +} + +.flashvsrFieldLabel { + color: var(--muted); + font-size: 0.72rem; + font-weight: 750; + letter-spacing: 0.08em; + text-transform: uppercase; +} + +.flashvsrVideoInput { + width: 100%; + min-width: 0; + color: var(--text); + font: 0.82rem/1.3 inherit; +} + +.flashvsrVideoInput::file-selector-button { + min-height: 34px; + margin-right: 0.75rem; + padding: 0 0.8rem; + border: 1px solid rgba(255, 255, 255, 0.20); + border-radius: 6px; + background: rgba(255, 255, 255, 0.08); + color: var(--text); + cursor: pointer; + font: inherit; + font-weight: 750; +} + +.flashvsrVideoInput::file-selector-button:hover { + background: rgba(255, 255, 255, 0.14); +} + +.flashvsrVideoInput:disabled { + cursor: not-allowed; + opacity: 0.58; +} + +.flashvsrUploadStatus { + min-height: 1.25rem; + color: var(--muted); + font-size: 0.82rem; + line-height: 1.4; +} + +.flashvsrUploadStatus[data-state="error"] { + color: var(--danger); +} + +.flashvsrUploadStatus[data-state="pending"] { + color: var(--warning); +} + +.flashvsrUploadStatus[data-state="ready"], +.flashvsrUploadStatus[data-state="running"] { + color: var(--accent-strong); +} + +.flashvsrStartButton { + width: 100%; + min-height: 42px; + border: 1px solid rgba(142, 240, 28, 0.55); + border-radius: 7px; + background: rgba(142, 240, 28, 0.16); + color: var(--text); + cursor: pointer; + font-weight: 800; +} + +.flashvsrStartButton:hover:not(:disabled) { + background: rgba(142, 240, 28, 0.25); +} + +.flashvsrStartButton:disabled { + cursor: not-allowed; + opacity: 0.45; +} + +@media (max-width: 900px) { + .flashvsrUploadCard { + top: 270px; + right: 18px; + left: 18px; + width: auto; + max-height: 440px; + } +} diff --git a/integrations/flashvsr/flashvsr/demo/web/adapter.js b/integrations/flashvsr/flashvsr/demo/web/adapter.js new file mode 100644 index 000000000..6870f270c --- /dev/null +++ b/integrations/flashvsr/flashvsr/demo/web/adapter.js @@ -0,0 +1,211 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +const START_ACTION = { + type: "action", + action: { event: "step" }, +} + +let context = null +let uploadCard = null +let uploadInput = null +let uploadStatus = null +let startButton = null +let uploadRequired = true +let hasDefaultInput = false +let inputReady = false +let connected = false +let started = false + +function makeUploadCard() { + const panel = document.createElement("section") + panel.className = "flashvsrUploadCard overlayPanel" + panel.setAttribute("aria-label", "FlashVSR video input") + panel.setAttribute("aria-busy", "false") + panel.innerHTML = + '
' + + '
FlashVSR' + + '

Upscale a video

' + + 'MP4' + + "
" + + '

Choose a source video, connect the session, then start FlashVSR.

' + + '" + + '' + + '' + return panel +} + +function setStatus(message, state = "idle") { + uploadStatus.textContent = message + uploadStatus.dataset.state = state +} + +function updateStartButton() { + startButton.disabled = !connected || !inputReady || started + startButton.textContent = started ? "FlashVSR Running" : "Start FlashVSR" + startButton.setAttribute("aria-busy", started ? "true" : "false") +} + +function showInputPrompt() { + const [file] = uploadInput.files + if (file) { + setStatus("Selected " + file.name + ". Connect the session to upload it.", "pending") + } else if (hasDefaultInput) { + setStatus("Server input ready. Connect the session or choose an MP4 to override it.", "ready") + } else { + setStatus("Choose an MP4 before connecting.", "pending") + } +} + +function applyVideoMetadata(payload) { + const width = Number(payload?.resolution?.width) + const height = Number(payload?.resolution?.height) + if (Number.isFinite(width) && Number.isFinite(height) && width > 0 && height > 0) { + context.setResolution(width, height) + } +} + +async function loadInputStatus() { + const response = await fetch("/api/session/input") + if (!response.ok) { + throw new Error("input status failed (" + response.status + ")") + } + const payload = await response.json() + uploadRequired = Boolean(payload.upload_required) + hasDefaultInput = Boolean(payload.has_default_input) + inputReady = !uploadRequired + applyVideoMetadata(payload) + showInputPrompt() +} + +async function uploadSelectedVideo() { + const [file] = uploadInput.files + if (!file) { + if (uploadRequired) { + throw new Error("Choose an MP4 before connecting.") + } + inputReady = true + return + } + if (!file.name.toLowerCase().endsWith(".mp4")) { + throw new Error("Choose an .mp4 video.") + } + + setStatus("Uploading and decoding " + file.name + "...", "pending") + const form = new FormData() + form.append("video", file, file.name) + const response = await fetch("/api/session/input", { + method: "POST", + body: form, + }) + if (!response.ok) { + const text = (await response.text()).trim().replace(/^\d+:\s*/, "") + throw new Error(text || "video upload failed (" + response.status + ")") + } + const payload = await response.json() + uploadRequired = false + hasDefaultInput = Boolean(payload.has_default_input) + inputReady = true + applyVideoMetadata(payload) + setStatus( + "Uploaded " + + payload.num_frames + + " frames at " + + payload.resolution.width + + "x" + + payload.resolution.height + + ". Waiting for connection.", + "ready", + ) + context.logEvent("uploaded " + file.name, { source: "client" }) +} + +function startFlashVSR() { + if (!connected || !inputReady || started) { + return + } + started = true + updateStartButton() + if (!context.sendCommand(START_ACTION, "start FlashVSR")) { + started = false + updateStartButton() + setStatus("The WebRTC control channel is not ready yet.", "error") + return + } + context.setFlow("FlashVSR running") + setStatus("FlashVSR is running.", "running") +} + +export default { + modelName: "FlashVSR", + stylesheet: new URL("./adapter.css?v=flashvsr-upload-v3", import.meta.url).href, + + async mount(sharedContext) { + context = sharedContext + uploadCard = makeUploadCard() + context.slots.panel.append(uploadCard) + uploadInput = uploadCard.querySelector(".flashvsrVideoInput") + uploadStatus = uploadCard.querySelector(".flashvsrUploadStatus") + startButton = uploadCard.querySelector(".flashvsrStartButton") + updateStartButton() + uploadInput.addEventListener("focus", context.releaseControls) + uploadInput.addEventListener("change", () => { + const [file] = uploadInput.files + uploadRequired = !file && !hasDefaultInput + inputReady = !file && hasDefaultInput + showInputPrompt() + updateStartButton() + context.releaseControls() + }) + startButton.addEventListener("click", startFlashVSR) + try { + await loadInputStatus() + } catch (error) { + setStatus(error.message, "error") + context.logEvent("input status unavailable: " + error.message, { + source: "client", + level: "error", + }) + } + }, + + async beforeConnect() { + connected = false + started = false + uploadCard.setAttribute("aria-busy", "true") + uploadInput.disabled = true + updateStartButton() + await uploadSelectedVideo() + }, + + onConnect() { + connected = true + uploadCard.setAttribute("aria-busy", "false") + uploadInput.disabled = true + updateStartButton() + context.setFlow("ready; click Start FlashVSR") + setStatus("Connected. Select Start FlashVSR to begin.", "ready") + }, + + onControlMessage(payload) { + if (payload.type === "chunk_done" && started) { + context.setStatus("Generating", "generating") + context.setFlow("FlashVSR running; chunk " + payload.chunk_index + " complete") + } + return false + }, + + onDisconnect() { + connected = false + started = false + inputReady = hasDefaultInput + uploadRequired = !hasDefaultInput + uploadCard.setAttribute("aria-busy", "false") + uploadInput.disabled = false + updateStartButton() + showInputPrompt() + }, +} diff --git a/integrations/flashvsr/flashvsr/demo/webrtc.py b/integrations/flashvsr/flashvsr/demo/webrtc.py new file mode 100644 index 000000000..c241dea36 --- /dev/null +++ b/integrations/flashvsr/flashvsr/demo/webrtc.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native FlashVSR runtime hooks for the shared WebRTC server.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, replace +from importlib.resources import files +from typing import Any, Literal + +from loguru import logger + +from flashdreams.runtime.demo import ( + DemoSpec, + PreparedScenario, + RuntimeHost, + WebRTCAppResources, + WebRTCOutputSpec, +) +from flashdreams.serving.webrtc.bootstrap import run_webrtc_server +from flashdreams.serving.webrtc.demo import ( + CreateWebRTCApp, + RunWebRTCServer, + serve_webrtc_demo, +) +from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager +from flashdreams.serving.webrtc.server import create_webrtc_app +from flashvsr.runtime import FlashVSRInferenceRuntime + +from .adapter import FlashVSRDemoAdapter +from .providers import PREPARED_VIDEO_METADATA_KEY +from .server import ( + FlashVSRUploadController, + FlashVSRWebRTCSessionInput, + configure_flashvsr_webrtc_app, +) +from .spec import PreparedFlashVSRVideo, resolve_video_scenario + +RuntimeFactory = Callable[..., Any] + + +@dataclass(frozen=True, kw_only=True, slots=True) +class FlashVSRWebRTCRuntimeConfig: + """Transport facts consumed by the shared WebRTC session manager.""" + + pipeline_config_name: str + device: str + video_height: int + video_width: int + fps: int + warmup_chunks: int + warmup_timeout_s: float + encoder_backend: Literal["auto", "default", "nvenc"] = "auto" + encoder_bitrate_bps: int = 6_000_000 + encoder_gop: int = 30 + + +def serve_flashvsr_webrtc_demo( + *, + spec: DemoSpec, + world_rank: int = 0, + runtime_factory: RuntimeFactory = FlashVSRInferenceRuntime, + create_app_fn: CreateWebRTCApp = create_webrtc_app, + server_runner: RunWebRTCServer = run_webrtc_server, +) -> object: + """Serve uploaded or server-side videos through the native runtime.""" + if spec.input_mode != "replay": + raise ValueError("FlashVSR WebRTC requires input_mode='replay'.") + if not isinstance(spec.output, WebRTCOutputSpec): + raise ValueError("FlashVSR WebRTC requires WebRTC output.") + if spec.config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + + adapter = FlashVSRDemoAdapter(runtime_factory=runtime_factory) + scenario = resolve_video_scenario(spec.scenario) + default_scenario: PreparedScenario | None = None + default_video: PreparedFlashVSRVideo | None = None + if scenario.input_path is not None: + default_scenario = adapter.prepare_scenario(spec) + candidate = default_scenario.metadata.get(PREPARED_VIDEO_METADATA_KEY) + if not isinstance(candidate, PreparedFlashVSRVideo): + raise TypeError("Prepared FlashVSR scenario is missing its decoded video.") + default_video = candidate + + output = spec.output + if default_scenario is not None: + output = replace( + output, + video_height=int(default_scenario.metadata["target_height"]), + video_width=int(default_scenario.metadata["target_width"]), + ) + runtime_options = dict(spec.config.runtime_options) + runtime_options.update( + { + "fps": output.fps, + "chunk_size": scenario.chunk_size, + } + ) + config = replace(spec.config, runtime_options=runtime_options) + shared_spec = replace( + spec, + scenario=scenario, + output=output, + config=config, + ) + runtime = adapter.create_runtime(config) + host = RuntimeHost(runtime) + preset_id = adapter.preset_id(config) + startup_warmup_chunks = output.warmup_chunks if default_scenario is not None else 0 + if default_scenario is None and output.warmup_chunks > 0: + logger.info( + "No startup FlashVSR input; deferring model construction until upload." + ) + runtime_config = FlashVSRWebRTCRuntimeConfig( + pipeline_config_name=preset_id, + device=config.device or "cuda:0", + video_height=output.video_height, + video_width=output.video_width, + fps=output.fps, + warmup_chunks=startup_warmup_chunks, + warmup_timeout_s=output.warmup_timeout_s, + # Browser uploads can change resolution between sessions. aiortc's + # software path accepts that; NVENC is bound to its startup dimensions. + encoder_backend="default", + encoder_gop=output.fps, + ) + manager = BaseWebRTCSessionManager( + runtime=runtime, + runtime_config=runtime_config, + fps=output.fps, + identity=preset_id, + busy_message="A FlashVSR session is already active.", + warmup_label="FlashVSR WebRTC", + fatal_generation_errors=True, + client_liveness_timeout_s=output.client_liveness_timeout_s, + shared_host=host, + shared_adapter=adapter, + shared_spec=shared_spec, + shared_spec_factory=lambda session_input: _uploaded_session_spec( + shared_spec, + session_input=session_input, + ), + shared_scenario=default_scenario, + ) + upload_controller = FlashVSRUploadController( + manager=manager, + adapter=adapter, + spec=shared_spec, + default_video=default_video, + ) + return serve_webrtc_demo( + output=output, + model_id=spec.model_id, + session_manager=manager, + app_resources=WebRTCAppResources( + model_web_resource=files("flashvsr.demo").joinpath("web"), + preload_name="FlashVSR", + configure_app=lambda app: configure_flashvsr_webrtc_app( + app, + controller=upload_controller, + ), + ), + world_rank=world_rank, + create_app_fn=create_app_fn, + server_runner=server_runner, + ) + + +def _uploaded_session_spec( + spec: DemoSpec, + *, + session_input: Any, +) -> DemoSpec: + if not isinstance(session_input, FlashVSRWebRTCSessionInput): + raise TypeError("FlashVSR WebRTC requires a decoded video upload.") + if not isinstance(spec.output, WebRTCOutputSpec): + raise ValueError("FlashVSR WebRTC requires WebRTC output.") + if spec.config is None: + raise RuntimeError("DemoSpec.config was not initialized.") + prepared = session_input.prepared_video + output = replace( + spec.output, + video_height=prepared.target_height, + video_width=prepared.target_width, + ) + runtime_options = dict(spec.config.runtime_options) + runtime_options.update( + { + "fps": output.fps, + "chunk_size": prepared.scenario.chunk_size, + } + ) + return replace( + spec, + scenario=prepared, + output=output, + config=replace(spec.config, runtime_options=runtime_options), + ) + + +__all__ = [ + "FlashVSRWebRTCRuntimeConfig", + "RuntimeFactory", + "serve_flashvsr_webrtc_demo", +] diff --git a/integrations/flashvsr/flashvsr/runtime.py b/integrations/flashvsr/flashvsr/runtime.py new file mode 100644 index 000000000..997a1fd19 --- /dev/null +++ b/integrations/flashvsr/flashvsr/runtime.py @@ -0,0 +1,729 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native runtime API implementation for FlashVSR video upscaling.""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, Literal, cast + +import torch +import torch.distributed as dist +from loguru import logger + +from flashdreams.core.distributed import init as init_distributed +from flashdreams.infra.config import derive_config +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.infra.time import TimeWindow +from flashdreams.runtime import ( + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InputField, +) +from flashdreams.runtime.demo.outputs import SessionInfo +from flashdreams.runtime.interfaces import InferenceRuntime, InferenceSession +from flashdreams.runtime.types import StepRequest, StepRequirements, StepResult + +FLASHVSR_MODEL_ID = "flashvsr" +DEFAULT_FLASHVSR_PRESET = "flashvsr-v1.1-sparse-ratio-2.0" + +FIELD_INPUT_HEIGHT = "input_height" +FIELD_INPUT_WIDTH = "input_width" +FIELD_FPS = "fps" +FIELD_TOTAL_FRAMES = "total_frames" +FIELD_CHUNK_SIZE = "chunk_size" +FIELD_TAIL_POLICY = "tail_policy" +FIELD_VIDEO_CHUNK = "video_chunk" +FIELD_VALID_FRAME_COUNT = "valid_frame_count" + +TailPolicy = Literal["drop", "pad"] +PipelineFactory = Callable[[Any, str], Any] + +_CHUNK_MODES: dict[int, tuple[int, int]] = { + 8: (5, 8), + 16: (13, 16), +} + + +@dataclass(frozen=True, kw_only=True, slots=True) +class FlashVSRSessionInputs: + """Session-wide shape and timing information for one input video.""" + + input_height: int + input_width: int + fps: float + chunk_size: Literal[8, 16] + total_frames: int | None = None + tail_policy: TailPolicy = "drop" + + def __post_init__(self) -> None: + if self.input_height <= 0 or self.input_width <= 0: + raise ValueError("FlashVSR input dimensions must be > 0.") + if self.fps <= 0: + raise ValueError("FlashVSR fps must be > 0.") + if self.chunk_size not in _CHUNK_MODES: + raise ValueError("FlashVSR chunk_size must be 8 or 16.") + if self.total_frames is not None and self.total_frames <= 0: + raise ValueError("FlashVSR total_frames must be > 0 when provided.") + if self.tail_policy not in {"drop", "pad"}: + raise ValueError("FlashVSR tail_policy must be 'drop' or 'pad'.") + + @property + def cold_frame_count(self) -> int: + """Return the raw frame count required by the first model step.""" + return _CHUNK_MODES[self.chunk_size][0] + + @property + def steady_frame_count(self) -> int: + """Return the raw frame count required after the first model step.""" + return _CHUNK_MODES[self.chunk_size][1] + + +@dataclass(frozen=True, kw_only=True, slots=True) +class FlashVSRRuntimeOptions: + """Construction options for the reusable FlashVSR runtime.""" + + pipeline_config: Any + sparse_ratio: float = 2.0 + scale: Literal[2, 4] = 2 + pipeline: Any | None = None + pipeline_factory: PipelineFactory | None = None + output_layout: VideoTensorLayout = "bcthw" + compile_network: bool | None = None + use_cuda_graph: bool | None = None + color_corrector_implementation: Literal["cuda", "torch"] | None = None + + def __post_init__(self) -> None: + if self.sparse_ratio <= 0: + raise ValueError("FlashVSR sparse_ratio must be > 0.") + if self.scale not in {2, 4}: + raise ValueError("FlashVSR scale must be 2 or 4.") + if self.output_layout != "bcthw": + raise ValueError("FlashVSR native runtime output_layout must be 'bcthw'.") + + +class FlashVSRModelAdapter: + """Expose FlashVSR through the transport-neutral inference runtime API.""" + + def __init__( + self, + *, + runtime_factory: Callable[..., InferenceRuntime] | None = None, + pipeline_factory: PipelineFactory | None = None, + ) -> None: + self._runtime_factory = runtime_factory or FlashVSRInferenceRuntime + self._pipeline_factory = pipeline_factory + self._mapping = IdentityInputMapping() + + @property + def model_id(self) -> str: + return FLASHVSR_MODEL_ID + + @property + def inference_input_schema(self) -> InferenceInputSchema: + return InferenceInputSchema( + description="FlashVSR streaming RGB video inputs.", + global_conditioning_fields=( + InputField( + name=FIELD_INPUT_HEIGHT, + input_modality="pixel-height", + frequency_consumed="once", + ), + InputField( + name=FIELD_INPUT_WIDTH, + input_modality="pixel-width", + frequency_consumed="once", + ), + InputField( + name=FIELD_FPS, + input_modality="fps", + frequency_consumed="once", + ), + InputField( + name=FIELD_CHUNK_SIZE, + input_modality="frame-count", + frequency_consumed="once", + ), + InputField( + name=FIELD_TOTAL_FRAMES, + required=False, + input_modality="frame-count", + frequency_consumed="once", + description="Absent for an unbounded looping source.", + ), + InputField( + name=FIELD_TAIL_POLICY, + input_modality="policy", + frequency_consumed="once", + ), + ), + step_fields=( + InputField( + name=FIELD_VIDEO_CHUNK, + input_modality="video/rgb", + frequency_consumed="per_step", + metadata={ + "shape": "[B,3,T,H,W]", + "layout": "bcthw", + "range": "[-1,1]", + }, + description="One normalized low-resolution RGB chunk.", + ), + ), + ) + + @property + def canonical_input_schema(self) -> None: + return None + + def default_input_mapping(self) -> IdentityInputMapping: + return self._mapping + + def preset_id(self, config: InferenceConfig | None) -> str: + """Return the requested preset or the stable FlashVSR default.""" + if config is None or config.preset_id is None: + return DEFAULT_FLASHVSR_PRESET + return config.preset_id + + def pipeline_config(self, config: InferenceConfig) -> Any: + """Resolve the dimension-independent pipeline scaffold.""" + custom = config.runtime_options.get("pipeline_config") + if custom is not None: + return custom + from flashvsr.config import RUNNER_CONFIGS # noqa: PLC0415 + + preset_id = self.preset_id(config) + try: + return RUNNER_CONFIGS[preset_id].pipeline + except KeyError as exc: + supported = ", ".join(sorted(RUNNER_CONFIGS)) + raise ValueError( + f"Unsupported FlashVSR preset_id={preset_id!r}. " + f"Supported presets: {supported}." + ) from exc + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.model_id: + raise ValueError( + f"FlashVSR adapter requires model_id={self.model_id!r}, " + f"got {config.model_id!r}." + ) + self.pipeline_config(config) + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + pipeline_config = self.pipeline_config(config) + options = config.runtime_options + configured_scale = getattr( + getattr(pipeline_config, "encoder", None), + "scale", + 2, + ) + scale = int(options.get("scale", configured_scale)) + sparse_ratio = options.get("sparse_ratio") + if sparse_ratio is None: + from flashvsr.config import RUNNER_CONFIGS # noqa: PLC0415 + + runner = RUNNER_CONFIGS.get(self.preset_id(config)) + sparse_ratio = getattr(runner, "sparse_ratio", 2.0) + return self._runtime_factory( + config=config, + options=FlashVSRRuntimeOptions( + pipeline_config=pipeline_config, + sparse_ratio=float(sparse_ratio), + scale=cast(Literal[2, 4], scale), + pipeline=options.get("pipeline"), + pipeline_factory=self._pipeline_factory, + compile_network=config.compile, + use_cuda_graph=_optional_bool(options.get("use_cuda_graph")), + color_corrector_implementation=_optional_color_corrector( + options.get("color_corrector_implementation") + ), + ), + ) + + +class FlashVSRInferenceRuntime: + """Own one reusable, resolution-specific FlashVSR model pipeline.""" + + def __init__( + self, + *, + config: InferenceConfig, + options: FlashVSRRuntimeOptions, + ) -> None: + self.config = config + self.options = options + if _is_torchrun_env() and not dist.is_initialized(): + init_distributed() + if dist.is_initialized(): + self.local_rank = int(os.environ.get("LOCAL_RANK", "0")) + self.world_size = dist.get_world_size() + self.global_rank = dist.get_rank() + self.device = torch.device(f"cuda:{self.local_rank}") + else: + self.local_rank = 0 + self.world_size = 1 + self.global_rank = 0 + self.device = torch.device(config.device or "cuda") + if self.world_size != 1: + raise NotImplementedError( + "The native FlashVSR demo runtime currently supports one GPU. " + "Use the existing full-attention runner for context parallelism." + ) + + self.is_rank_zero = self.global_rank == 0 + self.pipeline: Any | None = options.pipeline + self._owns_pipeline = options.pipeline is None + self._pipeline_shape: tuple[int, int] | None = None + self._reusable_cache: Any | None = None + self._active_session: FlashVSRInferenceSession | None = None + self._closed = False + + def preload(self) -> None: + """Keep the lazy runtime host hook explicit; dimensions arrive per session.""" + if self._closed: + raise RuntimeError("FlashVSR runtime is closed.") + + def peek_input_fps(self) -> float: + """Return transport timing before a concrete session has started.""" + return float(self.config.runtime_options.get("fps", 30.0)) + + def peek_steady_output_num_frames(self) -> int: + """Return the steady chunk size used to bound WebRTC output queues.""" + return int(self.config.runtime_options.get("chunk_size", 16)) + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + if self._closed: + raise RuntimeError("FlashVSR runtime is closed.") + if self._active_session is not None: + raise RuntimeError("FlashVSR runtime already has an active session.") + session_inputs = session_inputs_from_inference_input(inputs) + pipeline = self._pipeline_for(session_inputs) + cache = self._acquire_cache(pipeline) + _reset_pipeline_rng(pipeline, self.config.seed) + session = FlashVSRInferenceSession( + pipeline=pipeline, + cache=cache, + inputs=session_inputs, + device=self.device, + output_layout=self.options.output_layout, + rollout_seed=self.config.seed, + on_close=self._release_session, + ) + self._active_session = session + return session + + def close(self) -> None: + if self._closed: + return + self._closed = True + if self._active_session is not None: + self._active_session.close() + pipeline = self.pipeline + self.pipeline = None + self._reusable_cache = None + if self._owns_pipeline and pipeline is not None: + close = getattr(pipeline, "close", None) + if callable(close): + close() + if self.device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.empty_cache() + + def _pipeline_for(self, inputs: FlashVSRSessionInputs) -> Any: + shape = (inputs.input_height, inputs.input_width) + if self.pipeline is not None: + if self._pipeline_shape is None: + self._pipeline_shape = shape + elif self._pipeline_shape != shape: + if not self._owns_pipeline: + raise ValueError( + "An injected FlashVSR pipeline cannot change resolution: " + f"loaded {self._pipeline_shape}, requested {shape}." + ) + self._dispose_pipeline() + if self.pipeline is not None: + return self.pipeline + + pipeline_config = self._derive_pipeline_config(inputs) + factory = self.options.pipeline_factory or _default_pipeline_factory + logger.info( + "Building native FlashVSR runtime for input={}x{} scale={}.", + inputs.input_height, + inputs.input_width, + self.options.scale, + ) + self.pipeline = factory(pipeline_config, str(self.device)) + self._pipeline_shape = shape + return self.pipeline + + def _derive_pipeline_config(self, inputs: FlashVSRSessionInputs) -> Any: + scale = self.options.scale + target_height = (inputs.input_height * scale // 128) * 128 + target_width = (inputs.input_width * scale // 128) * 128 + if target_height <= 0 or target_width <= 0: + raise ValueError( + "FlashVSR scaled dimensions must each contain a 128-pixel block; " + f"got input={inputs.input_height}x{inputs.input_width}, scale={scale}." + ) + topk_ratio = ( + self.options.sparse_ratio * 768 * 1280 / (target_height * target_width) + ) + encoder_updates: dict[str, Any] = { + "input_H": inputs.input_height, + "input_W": inputs.input_width, + "scale": scale, + } + decoder_updates: dict[str, Any] = {} + transformer_updates: dict[str, Any] = {"topk_ratio": topk_ratio} + if self.options.compile_network is not None: + enabled = self.options.compile_network + encoder_updates["use_compile"] = enabled + decoder_updates["use_compile"] = enabled + transformer_updates["compile_network"] = enabled + if self.options.use_cuda_graph is not None: + enabled = self.options.use_cuda_graph + encoder_updates["use_cuda_graph"] = enabled + decoder_updates["use_cuda_graph"] = enabled + transformer_updates["use_cuda_graph"] = enabled + if self.options.color_corrector_implementation is not None: + decoder_updates["color_corrector_implementation"] = ( + self.options.color_corrector_implementation + ) + diffusion_updates: dict[str, Any] = {"transformer": transformer_updates} + if self.config.seed is not None: + diffusion_updates["seed"] = int(self.config.seed) + return derive_config( + self.options.pipeline_config, + encoder=encoder_updates, + decoder=decoder_updates, + diffusion_model=diffusion_updates, + ) + + def _acquire_cache(self, pipeline: Any) -> Any: + cache = self._reusable_cache + self._reusable_cache = None + if cache is None: + return pipeline.initialize_cache() + reset = getattr(pipeline, "reset_cache_in_place", None) + if not callable(reset): + return pipeline.initialize_cache() + reset(cache) + return cache + + def _release_session( + self, + session: "FlashVSRInferenceSession", + cache: Any, + ) -> None: + if self._active_session is session: + self._active_session = None + if not self._closed: + self._reusable_cache = cache + + def _dispose_pipeline(self) -> None: + pipeline = self.pipeline + self.pipeline = None + self._pipeline_shape = None + self._reusable_cache = None + if pipeline is not None: + close = getattr(pipeline, "close", None) + if callable(close): + close() + + +class FlashVSRInferenceSession: + """Run one FlashVSR video stream from new-API frame chunks.""" + + def __init__( + self, + *, + pipeline: Any, + cache: Any, + inputs: FlashVSRSessionInputs, + device: torch.device, + output_layout: VideoTensorLayout, + rollout_seed: int | None, + on_close: Callable[["FlashVSRInferenceSession", Any], None], + ) -> None: + self.pipeline = pipeline + self.cache = cache + self.inputs = inputs + self.device = device + self.output_layout = output_layout + self.rollout_seed = rollout_seed + self._on_close = on_close + self._step_index = 0 + self._frame_start = 0 + self._closed = False + + def session_info(self) -> SessionInfo: + """Return video-output facts known after cache initialization.""" + return SessionInfo( + output_layout=self.output_layout, + steady_output_frame_count=self.inputs.steady_frame_count, + metadata={ + "input_height": self.inputs.input_height, + "input_width": self.inputs.input_width, + }, + ) + + def next_step_requirements(self) -> StepRequirements | None: + if self._closed: + return None + requested = ( + self.inputs.cold_frame_count + if self._step_index == 0 + else self.inputs.steady_frame_count + ) + valid = requested + if self.inputs.total_frames is not None: + remaining = self.inputs.total_frames - self._frame_start + if remaining <= 0: + return None + if remaining < requested: + if self.inputs.tail_policy == "drop": + return None + valid = remaining + return StepRequirements( + step_index=self._step_index, + input_frame_count=requested, + steady_output_frame_count=self.inputs.steady_frame_count, + inference_input_schema=FlashVSRModelAdapter().inference_input_schema, + metadata={ + FIELD_VALID_FRAME_COUNT: valid, + "frame_start": self._frame_start, + }, + ) + + def next_step_request(self) -> StepRequest | None: + """Expose the protocol-compatible request alongside the native requirements.""" + requirements = self.next_step_requirements() + if requirements is None: + return None + metadata = dict(requirements.metadata) + metadata["input_frame_count"] = requirements.input_frame_count + metadata["steady_output_frame_count"] = requirements.steady_output_frame_count + return StepRequest( + step_index=requirements.step_index, + inference_input_schema=requirements.inference_input_schema, + metadata=metadata, + ) + + def step(self, inputs: InferenceInput) -> StepResult: + if self._closed: + raise RuntimeError("FlashVSR inference session is closed.") + request = self.next_step_requirements() + if request is None: + raise RuntimeError("FlashVSR inference session has no remaining step.") + valid = int( + inputs.metadata.get( + FIELD_VALID_FRAME_COUNT, + request.metadata[FIELD_VALID_FRAME_COUNT], + ) + ) + expected_valid = int(request.metadata[FIELD_VALID_FRAME_COUNT]) + if valid != expected_valid: + raise ValueError( + "FlashVSR valid frame count mismatch: " + f"expected {expected_valid}, got {valid}." + ) + video = _require_video_chunk( + inputs, + expected_frames=request.input_frame_count, + expected_height=self.inputs.input_height, + expected_width=self.inputs.input_width, + ) + dtype = self.pipeline.diffusion_model.dtype + video = video.to(device=self.device, dtype=dtype) + step_index = self._step_index + output = self.pipeline.generate( + autoregressive_index=step_index, + cache=self.cache, + input=video, + ) + stats = self.pipeline.finalize( + autoregressive_index=step_index, + cache=self.cache, + ) + if output.ndim != 5 or output.shape[2] < valid: + raise ValueError( + "FlashVSR pipeline output must be [B,C,T,H,W] with at least " + f"{valid} frames; got {tuple(output.shape)}." + ) + if output.shape[2] != valid: + output = output[:, :, :valid] + frame_start = self._frame_start + frame_end = frame_start + valid + self._frame_start = frame_end + self._step_index += 1 + return StepResult.from_video_chunk( + step_index=step_index, + video_chunk=output, + layout=self.output_layout, + output_window=TimeWindow( + start_s=frame_start / self.inputs.fps, + end_s=frame_end / self.inputs.fps, + ), + metrics=_numeric_metrics(stats), + metadata={ + "input_frame_count": request.input_frame_count, + FIELD_VALID_FRAME_COUNT: valid, + "resolution": { + "width": int(output.shape[-1]), + "height": int(output.shape[-2]), + }, + }, + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + if self._closed: + raise RuntimeError("FlashVSR inference session is closed.") + if inputs is not None: + replacement = session_inputs_from_inference_input(inputs) + if replacement != self.inputs: + raise ValueError( + "FlashVSR session reset cannot change video shape or timing." + ) + reset = getattr(self.pipeline, "reset_cache_in_place", None) + if not callable(reset): + raise RuntimeError( + "FlashVSR pipeline does not support in-place cache reset." + ) + reset(self.cache) + _reset_pipeline_rng(self.pipeline, self.rollout_seed) + self._step_index = 0 + self._frame_start = 0 + + def close(self) -> None: + if self._closed: + return + self._closed = True + self._on_close(self, self.cache) + + +def session_inputs_from_inference_input( + inputs: InferenceInput, +) -> FlashVSRSessionInputs: + """Validate and decode session-global FlashVSR inputs.""" + values = inputs.global_conditioning + required = ( + FIELD_INPUT_HEIGHT, + FIELD_INPUT_WIDTH, + FIELD_FPS, + FIELD_CHUNK_SIZE, + FIELD_TAIL_POLICY, + ) + missing = tuple(name for name in required if name not in values) + if missing: + raise ValueError(f"Missing FlashVSR global conditioning field(s): {missing}.") + chunk_size = int(values[FIELD_CHUNK_SIZE]) + return FlashVSRSessionInputs( + input_height=int(values[FIELD_INPUT_HEIGHT]), + input_width=int(values[FIELD_INPUT_WIDTH]), + fps=float(values[FIELD_FPS]), + chunk_size=cast(Literal[8, 16], chunk_size), + total_frames=( + None + if values.get(FIELD_TOTAL_FRAMES) is None + else int(values[FIELD_TOTAL_FRAMES]) + ), + tail_policy=cast(TailPolicy, str(values[FIELD_TAIL_POLICY])), + ) + + +def _require_video_chunk( + inputs: InferenceInput, + *, + expected_frames: int, + expected_height: int, + expected_width: int, +) -> torch.Tensor: + value = inputs.step.get(FIELD_VIDEO_CHUNK) + if not isinstance(value, torch.Tensor): + raise TypeError( + f"FlashVSR step input {FIELD_VIDEO_CHUNK!r} must be a torch.Tensor." + ) + expected = (1, 3, expected_frames, expected_height, expected_width) + if tuple(value.shape) != expected: + raise ValueError( + f"FlashVSR video chunk must have shape {expected}, got {tuple(value.shape)}." + ) + if not value.is_floating_point(): + raise TypeError("FlashVSR video chunk must use a floating-point dtype.") + return value + + +def _default_pipeline_factory(pipeline_config: Any, device: str) -> Any: + return pipeline_config.setup().to(device=device).eval() + + +def _reset_pipeline_rng(pipeline: Any, seed: int | None) -> None: + if seed is None: + return + rng = getattr(getattr(pipeline, "diffusion_model", None), "rng", None) + if rng is not None: + rng.manual_seed(int(seed)) + + +def _numeric_metrics(value: Any) -> dict[str, float | int]: + if not isinstance(value, dict): + return {} + return { + str(key): metric + for key, metric in value.items() + if isinstance(metric, int | float) and not isinstance(metric, bool) + } + + +def _optional_bool(value: Any) -> bool | None: + if value is None or isinstance(value, bool): + return value + if isinstance(value, str): + lowered = value.strip().lower() + if lowered in {"1", "true", "yes", "on"}: + return True + if lowered in {"0", "false", "no", "off"}: + return False + return bool(value) + + +def _optional_color_corrector( + value: Any, +) -> Literal["cuda", "torch"] | None: + if value is None: + return None + if value not in {"cuda", "torch"}: + raise ValueError("color_corrector_implementation must be 'cuda' or 'torch'.") + return cast(Literal["cuda", "torch"], value) + + +def _is_torchrun_env() -> bool: + return "RANK" in os.environ and "WORLD_SIZE" in os.environ + + +__all__ = [ + "DEFAULT_FLASHVSR_PRESET", + "FIELD_CHUNK_SIZE", + "FIELD_FPS", + "FIELD_INPUT_HEIGHT", + "FIELD_INPUT_WIDTH", + "FIELD_TAIL_POLICY", + "FIELD_TOTAL_FRAMES", + "FIELD_VALID_FRAME_COUNT", + "FIELD_VIDEO_CHUNK", + "FLASHVSR_MODEL_ID", + "FlashVSRInferenceRuntime", + "FlashVSRInferenceSession", + "FlashVSRModelAdapter", + "FlashVSRRuntimeOptions", + "FlashVSRSessionInputs", + "PipelineFactory", + "TailPolicy", + "session_inputs_from_inference_input", +] diff --git a/integrations/flashvsr/pyproject.toml b/integrations/flashvsr/pyproject.toml index 68f75d610..c5fe45b0a 100644 --- a/integrations/flashvsr/pyproject.toml +++ b/integrations/flashvsr/pyproject.toml @@ -24,7 +24,7 @@ description = "FlashVSR streaming video super-resolution (LR projector + Wan 2.1 readme = "README.md" requires-python = ">=3.10" dependencies = [ - "flashdreams", + "flashdreams[serving]", "grpcio>=1.80", "mediapy>=1.1", "Pillow>=10", @@ -44,6 +44,7 @@ dev = [ ] [project.scripts] +flashvsr-demo = "flashvsr.demo.app:main" flashvsr-grpc-server = "flashvsr.grpc.uplift_server:main" flashvsr-grpc-client = "flashvsr.grpc.uplift_client:main" flashvsr-feed-frames = "flashvsr.grpc.uplift_client:run_continuous_client" @@ -71,6 +72,7 @@ exclude = ["tests"] [tool.setuptools.package-data] "flashvsr" = ["csrc/*.cu"] +"flashvsr.demo" = ["web/adapter.js", "web/adapter.css"] # Pytest's ``configfile`` resolution picks the nearest ``pyproject.toml`` to # the test files, so when running ``pytest integrations/flashvsr/tests`` diff --git a/integrations/flashvsr/tests/test_runtime_api.py b/integrations/flashvsr/tests/test_runtime_api.py new file mode 100644 index 000000000..955b50c16 --- /dev/null +++ b/integrations/flashvsr/tests/test_runtime_api.py @@ -0,0 +1,495 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +import torch + +from flashdreams.runtime import InferenceConfig, InferenceInput, StepRequirements +from flashdreams.runtime.demo import ( + DemoSpec, + NullOutputSpec, + PreparedScenario, + UserInputWindow, + WebRTCOutputSpec, +) +from flashdreams.runtime.demo.replay import run_replay_demo +from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager +from flashvsr.demo import ( + FLASHVSR_MODEL_ID, + FlashVSRDemoAdapter, + FlashVSRVideoInputProvider, + FlashVSRVideoScenario, + PreparedFlashVSRVideo, +) +from flashvsr.demo.app import _replay_spec, _webrtc_spec, parse_args +from flashvsr.demo.providers import PREPARED_VIDEO_METADATA_KEY +from flashvsr.demo.spec import prepare_video_source +from flashvsr.demo.webrtc import serve_flashvsr_webrtc_demo +from flashvsr.runtime import ( + FIELD_CHUNK_SIZE, + FIELD_FPS, + FIELD_INPUT_HEIGHT, + FIELD_INPUT_WIDTH, + FIELD_TAIL_POLICY, + FIELD_TOTAL_FRAMES, + FIELD_VALID_FRAME_COUNT, + FIELD_VIDEO_CHUNK, + FlashVSRModelAdapter, +) + +pytestmark = pytest.mark.ci_cpu + + +class _FakePipeline: + def __init__(self) -> None: + self.diffusion_model = SimpleNamespace( + dtype=torch.float32, + rng=torch.Generator().manual_seed(0), + ) + self.generated: list[tuple[int, tuple[int, ...]]] = [] + self.finalized: list[int] = [] + self.cache = SimpleNamespace(reset_count=0) + + def initialize_cache(self) -> Any: + return self.cache + + def reset_cache_in_place(self, cache: Any) -> None: + assert cache is self.cache + cache.reset_count += 1 + + def generate( + self, + *, + autoregressive_index: int, + cache: Any, + input: torch.Tensor, + ) -> torch.Tensor: + assert cache is self.cache + self.generated.append((autoregressive_index, tuple(input.shape))) + return input + 0.25 + + def finalize(self, *, autoregressive_index: int, cache: Any) -> dict[str, float]: + assert cache is self.cache + self.finalized.append(autoregressive_index) + return {"total_ms": 2.5} + + +class _FakeRuntime: + def __init__(self, *, config: InferenceConfig, options: Any) -> None: + self.config = config + self.options = options + + def preload(self) -> None: + return + + def peek_input_fps(self) -> float: + return float(self.config.runtime_options["fps"]) + + def peek_steady_output_num_frames(self) -> int: + return int(self.config.runtime_options["chunk_size"]) + + def start_session(self, inputs: InferenceInput) -> Any: + raise AssertionError(f"test server should not start a session: {inputs}") + + def close(self) -> None: + return + + +def _pipeline_config() -> Any: + return SimpleNamespace(encoder=SimpleNamespace(scale=2)) + + +def _initial_inputs( + *, + total_frames: int | None = 13, + tail_policy: str = "drop", +) -> InferenceInput: + values: dict[str, Any] = { + FIELD_INPUT_HEIGHT: 64, + FIELD_INPUT_WIDTH: 64, + FIELD_FPS: 20.0, + FIELD_CHUNK_SIZE: 8, + FIELD_TAIL_POLICY: tail_policy, + } + if total_frames is not None: + values[FIELD_TOTAL_FRAMES] = total_frames + return InferenceInput(global_conditioning=values) + + +def _prepared_video( + *, + frames: int = 6, + loop_input: bool = False, + tail_policy: str = "pad", +) -> PreparedFlashVSRVideo: + video = torch.arange(frames, dtype=torch.float32).view(1, 1, frames, 1, 1) + video = video.expand(1, 3, frames, 64, 64).contiguous() + scenario = FlashVSRVideoScenario( + input_path="memory.mp4", + chunk_size=8, + fps=20.0, + tail_policy=tail_policy, + loop_input=loop_input, + ) + return PreparedFlashVSRVideo( + scenario=scenario, + resolved_path=Path("memory.mp4"), + video=video, + input_height=64, + input_width=64, + target_height=128, + target_width=128, + fps=20.0, + ) + + +def test_adapter_declares_native_video_inputs_and_demo_modes() -> None: + adapter = FlashVSRDemoAdapter() + + assert adapter.model_id == FLASHVSR_MODEL_ID + assert adapter.supported_input_modes() == ("replay",) + assert adapter.supported_output_modes() == ("mp4", "null", "webrtc") + global_fields = { + field.name + for field in adapter.inference_input_schema.global_conditioning_fields + } + step_fields = {field.name for field in adapter.inference_input_schema.step_fields} + assert { + FIELD_INPUT_HEIGHT, + FIELD_INPUT_WIDTH, + FIELD_FPS, + FIELD_CHUNK_SIZE, + FIELD_TAIL_POLICY, + }.issubset(global_fields) + assert step_fields == {FIELD_VIDEO_CHUNK} + + +def test_native_runtime_session_requests_and_processes_cold_then_steady() -> None: + pipeline = _FakePipeline() + config = InferenceConfig( + model_id=FLASHVSR_MODEL_ID, + device="cpu", + seed=7, + runtime_options={ + "pipeline_config": _pipeline_config(), + "pipeline": pipeline, + "chunk_size": 8, + "fps": 20.0, + }, + ) + runtime = FlashVSRModelAdapter().create_runtime(config) + session = runtime.start_session(_initial_inputs()) + + first = session.next_step_requirements() + assert first is not None + assert first.input_frame_count == 5 + first_result = session.step( + InferenceInput( + step={FIELD_VIDEO_CHUNK: torch.zeros(1, 3, 5, 64, 64)}, + metadata={FIELD_VALID_FRAME_COUNT: 5}, + ) + ) + second = session.next_step_requirements() + assert second is not None + assert second.input_frame_count == 8 + second_result = session.step( + InferenceInput( + step={FIELD_VIDEO_CHUNK: torch.zeros(1, 3, 8, 64, 64)}, + metadata={FIELD_VALID_FRAME_COUNT: 8}, + ) + ) + + assert first_result.layout == "bcthw" + assert first_result.frame_count == 5 + assert first_result.output_window is not None + assert first_result.output_window.end_s == pytest.approx(0.25) + assert first_result.metadata["resolution"] == {"width": 64, "height": 64} + assert second_result.frame_count == 8 + assert second_result.output_window is not None + assert second_result.output_window.start_s == pytest.approx(0.25) + assert second_result.metrics["total_ms"] == 2.5 + assert session.next_step_requirements() is None + assert pipeline.generated == [ + (0, (1, 3, 5, 64, 64)), + (1, (1, 3, 8, 64, 64)), + ] + assert pipeline.finalized == [0, 1] + + session.close() + second_session = runtime.start_session(_initial_inputs(total_frames=5)) + assert pipeline.cache.reset_count == 1 + second_session.close() + runtime.close() + + +def test_session_rejects_invalid_tail_metadata_before_model_execution() -> None: + pipeline = _FakePipeline() + runtime = FlashVSRModelAdapter().create_runtime( + InferenceConfig( + model_id=FLASHVSR_MODEL_ID, + device="cpu", + runtime_options={ + "pipeline_config": _pipeline_config(), + "pipeline": pipeline, + "chunk_size": 8, + "fps": 20.0, + }, + ) + ) + session = runtime.start_session(_initial_inputs(total_frames=3, tail_policy="pad")) + + with pytest.raises(ValueError, match="valid frame count mismatch"): + session.step( + InferenceInput( + step={FIELD_VIDEO_CHUNK: torch.zeros(1, 3, 5, 64, 64)}, + metadata={FIELD_VALID_FRAME_COUNT: 5}, + ) + ) + + request = session.next_step_requirements() + assert request is not None + assert request.step_index == 0 + assert pipeline.generated == [] + assert pipeline.finalized == [] + session.close() + runtime.close() + + +def test_provider_loops_short_source_to_exact_requested_shape() -> None: + prepared = _prepared_video(frames=3, loop_input=True) + scenario = PreparedScenario( + initial_inputs=_initial_inputs(total_frames=None), + metadata={PREPARED_VIDEO_METADATA_KEY: prepared}, + ) + provider = FlashVSRVideoInputProvider( + scenario=scenario, + inference_input_schema=FlashVSRModelAdapter().inference_input_schema, + ) + + step = provider.prepare_step( + request=StepRequirements( + step_index=0, + input_frame_count=5, + metadata={FIELD_VALID_FRAME_COUNT: 5}, + ), + user_window=UserInputWindow(start_s=0.0, end_s=0.25), + ) + + assert step.inference_input is not None + chunk = step.inference_input.step[FIELD_VIDEO_CHUNK] + assert tuple(chunk.shape) == (1, 3, 5, 64, 64) + assert chunk[0, 0, :, 0, 0].tolist() == [0.0, 1.0, 2.0, 0.0, 1.0] + + +def test_shared_replay_uses_provider_for_padded_tail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import flashvsr.demo.adapter as adapter_module + + prepared = _prepared_video(frames=6, tail_policy="pad") + monkeypatch.setattr( + adapter_module, "prepare_video_source", lambda *a, **k: prepared + ) + pipeline = _FakePipeline() + spec = DemoSpec( + model_id=FLASHVSR_MODEL_ID, + input_mode="replay", + scenario=prepared.scenario, + output=NullOutputSpec(), + config=InferenceConfig( + model_id=FLASHVSR_MODEL_ID, + device="cpu", + seed=0, + runtime_options={ + "pipeline_config": _pipeline_config(), + "pipeline": pipeline, + "chunk_size": 8, + "fps": 20.0, + }, + ), + ) + + result = run_replay_demo(spec=spec, adapter=FlashVSRDemoAdapter()) + + assert result.status == "completed" + assert pipeline.generated == [ + (0, (1, 3, 5, 64, 64)), + (1, (1, 3, 8, 64, 64)), + ] + + +def test_prepare_video_source_normalizes_and_derives_target_shape( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + import flashvsr.demo.spec as spec_module + + path = tmp_path / "input.mp4" + path.write_bytes(b"fixture") + pixels = np.full((5, 64, 96, 3), 255, dtype=np.uint8) + monkeypatch.setattr(spec_module, "resolve_input_path", lambda *a, **k: path) + monkeypatch.setattr(spec_module, "read_video_rgb", lambda _: pixels) + + prepared = prepare_video_source( + FlashVSRVideoScenario(input_path=path, chunk_size=8, fps=24.0), + scale=2, + ) + + assert tuple(prepared.video.shape) == (1, 3, 5, 64, 96) + assert torch.all(prepared.video == 1) + assert (prepared.target_height, prepared.target_width) == (128, 128) + + +def test_cli_builds_null_and_webrtc_specs(tmp_path: Path) -> None: + replay_args = parse_args( + [ + "replay", + "--input", + str(tmp_path / "input.mp4"), + "--output-mode", + "null", + "--chunk-size", + "8", + "--fps", + "24", + "--no-compile", + "--no-cuda-graph", + "--color-corrector", + "torch", + ] + ) + replay_spec = _replay_spec(replay_args) + assert isinstance(replay_spec.output, NullOutputSpec) + assert replay_spec.config is not None + assert replay_spec.config.compile is False + assert replay_spec.config.runtime_options["use_cuda_graph"] is False + assert ( + replay_spec.config.runtime_options["color_corrector_implementation"] == "torch" + ) + + webrtc_args = parse_args( + [ + "webrtc", + "--input", + str(tmp_path / "input.mp4"), + "--port", + "9090", + "--fps", + "29.97", + "--warmup-chunks", + "1", + "--prefer-sw-encoder", + ] + ) + webrtc_spec = _webrtc_spec(webrtc_args, device="cuda:3") + assert isinstance(webrtc_spec.output, WebRTCOutputSpec) + assert webrtc_spec.output.port == 9090 + assert isinstance(webrtc_spec.scenario, FlashVSRVideoScenario) + assert webrtc_spec.scenario.loop_input is True + assert webrtc_spec.scenario.fps == 30.0 + assert webrtc_spec.output.fps == 30 + assert webrtc_spec.config is not None + assert webrtc_spec.config.device == "cuda:3" + assert webrtc_spec.config.runtime_options["fps"] == 30.0 + + +def test_cli_uses_source_fps_when_not_overridden( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + import flashvsr.demo.app as app_module + + input_path = tmp_path / "input.mp4" + monkeypatch.setattr( + app_module, + "resolve_input_path", + lambda *args, **kwargs: input_path, + ) + monkeypatch.setattr(app_module, "read_video_fps", lambda path: 24.0) + + args = parse_args(["replay", "--input", str(input_path), "--output-mode", "null"]) + spec = _replay_spec(args) + + assert isinstance(spec.scenario, FlashVSRVideoScenario) + assert spec.scenario.fps == 24.0 + assert spec.config is not None + assert spec.config.runtime_options["fps"] == 24.0 + + +def test_webrtc_uses_native_shared_host_and_resolved_output_shape( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import flashvsr.demo.adapter as adapter_module + + prepared = _prepared_video(frames=5, loop_input=True) + monkeypatch.setattr( + adapter_module, "prepare_video_source", lambda *a, **k: prepared + ) + calls: list[dict[str, Any]] = [] + spec = DemoSpec( + model_id=FLASHVSR_MODEL_ID, + input_mode="replay", + scenario=prepared.scenario, + output=WebRTCOutputSpec( + host="127.0.0.1", + port=8088, + fps=20, + video_height=1, + video_width=1, + warmup_chunks=0, + ), + config=InferenceConfig( + model_id=FLASHVSR_MODEL_ID, + device="cpu", + runtime_options={ + "pipeline_config": _pipeline_config(), + "chunk_size": 8, + "fps": 20.0, + "prefer_sw_encoder": True, + }, + ), + ) + + result = serve_flashvsr_webrtc_demo( + spec=spec, + world_rank=1, + runtime_factory=_FakeRuntime, + server_runner=lambda **kwargs: calls.append(kwargs), + ) + + assert result is None + assert len(calls) == 1 + manager = calls[0]["session_manager"] + assert type(manager) is BaseWebRTCSessionManager + assert manager._shared_host is not None + assert isinstance(manager._shared_adapter, FlashVSRDemoAdapter) + assert manager._shared_scenario is not None + assert manager.runtime_config.video_height == 128 + assert manager.runtime_config.video_width == 128 + assert manager.runtime_config.encoder_backend == "default" + assert calls[0]["host"] == "127.0.0.1" + assert calls[0]["port"] == 8088 + manager._shared_host.close() + + +def test_replay_cli_requires_output_only_for_mp4(tmp_path: Path) -> None: + parse_args(["replay", "--output", str(tmp_path / "output.mp4")]) + with pytest.raises(SystemExit): + parse_args(["replay"]) + with pytest.raises(SystemExit): + parse_args( + [ + "replay", + "--output-mode", + "null", + "--output", + str(tmp_path / "output.mp4"), + ] + ) diff --git a/integrations/flashvsr/tests/test_webrtc_upload.py b/integrations/flashvsr/tests/test_webrtc_upload.py new file mode 100644 index 000000000..6670261fd --- /dev/null +++ b/integrations/flashvsr/tests/test_webrtc_upload.py @@ -0,0 +1,420 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +from aiohttp import FormData, web +from aiohttp.test_utils import TestClient, TestServer + +from flashdreams.runtime import InferenceConfig +from flashdreams.runtime.demo import DemoSpec, WebRTCOutputSpec +from flashdreams.serving.webrtc.server import PACKAGE_RESOURCE_STACK_KEY +from flashvsr.corrector import FlashVSRColorCorrector +from flashvsr.demo.app import _webrtc_spec, parse_args +from flashvsr.demo.providers import PREPARED_VIDEO_METADATA_KEY +from flashvsr.demo.server import ( + FlashVSRUploadController, + FlashVSRWebRTCSessionInput, + configure_flashvsr_webrtc_app, +) +from flashvsr.demo.spec import FlashVSRVideoScenario, PreparedFlashVSRVideo +from flashvsr.demo.webrtc import serve_flashvsr_webrtc_demo +from flashvsr.runtime import FLASHVSR_MODEL_ID + +pytestmark = pytest.mark.ci_cpu + + +class _FakeSessionManager: + def __init__(self) -> None: + self.pending_session_input: Any = None + self.active = False + + def has_active_session(self) -> bool: + return self.active + + def set_pending_session_input(self, session_input: Any) -> None: + self.pending_session_input = session_input + + +class _FakeUploadAdapter: + def __init__(self, prepared: PreparedFlashVSRVideo) -> None: + self.prepared = prepared + self.upload_bytes: bytes | None = None + self.upload_path: Path | None = None + self.error: Exception | None = None + + def prepare_uploaded_video( + self, + spec: DemoSpec, + *, + upload_path: Path, + original_name: str, + ) -> PreparedFlashVSRVideo: + del spec, original_name + self.upload_path = upload_path + self.upload_bytes = upload_path.read_bytes() + if self.error is not None: + raise self.error + return self.prepared + + +class _FakeRuntime: + def __init__(self, *, config: InferenceConfig, options: Any) -> None: + self.config = config + self.options = options + self.closed = False + + def preload(self) -> None: + return + + def peek_input_fps(self) -> float: + return float(self.config.runtime_options["fps"]) + + def peek_steady_output_num_frames(self) -> int: + return int(self.config.runtime_options["chunk_size"]) + + def start_session(self, inputs: Any) -> Any: + raise AssertionError(f"test should not start a session: {inputs}") + + def close(self) -> None: + self.closed = True + + +def _pipeline_config() -> Any: + return SimpleNamespace(encoder=SimpleNamespace(scale=2)) + + +def _prepared_video( + *, + filename: str = "upload.mp4", + frames: int = 5, + input_height: int = 64, + input_width: int = 96, + target_height: int = 128, + target_width: int = 128, + fps: float = 20.0, +) -> PreparedFlashVSRVideo: + scenario = FlashVSRVideoScenario( + input_path=filename, + chunk_size=8, + fps=fps, + loop_input=True, + ) + return PreparedFlashVSRVideo( + scenario=scenario, + resolved_path=Path(filename), + video=torch.zeros(1, 3, frames, input_height, input_width), + input_height=input_height, + input_width=input_width, + target_height=target_height, + target_width=target_width, + fps=fps, + ) + + +def _upload_spec(*, input_path: str | Path | None = None) -> DemoSpec: + return DemoSpec( + model_id=FLASHVSR_MODEL_ID, + input_mode="replay", + scenario=FlashVSRVideoScenario( + input_path=input_path, + chunk_size=8, + fps=20.0, + loop_input=True, + ), + output=WebRTCOutputSpec( + host="127.0.0.1", + port=8088, + fps=20, + video_height=128, + video_width=128, + warmup_chunks=1, + ), + config=InferenceConfig( + model_id=FLASHVSR_MODEL_ID, + device="cpu", + runtime_options={ + "pipeline_config": _pipeline_config(), + "chunk_size": 8, + "fps": 20.0, + }, + ), + ) + + +async def _build_upload_client( + controller: FlashVSRUploadController, +) -> TestClient: + app = web.Application() + + async def offer(_: web.Request) -> web.StreamResponse: + return web.json_response({"sdp": "answer", "type": "answer"}) + + app.router.add_post("/api/webrtc/offer", offer) + configure_flashvsr_webrtc_app(app, controller=controller) + client = TestClient(TestServer(app)) + await client.start_server() + return client + + +def test_webrtc_cli_allows_browser_upload_without_input() -> None: + args = parse_args(["webrtc"]) + spec = _webrtc_spec(args, device="cuda:0") + + assert args.input_path is None + assert isinstance(spec.scenario, FlashVSRVideoScenario) + assert spec.scenario.input_path is None + assert spec.scenario.fps == 30.0 + assert isinstance(spec.output, WebRTCOutputSpec) + assert spec.output.fps == 30 + + +def test_webrtc_server_defers_warmup_without_default_input() -> None: + calls: list[dict[str, Any]] = [] + app = serve_flashvsr_webrtc_demo( + spec=_upload_spec(), + runtime_factory=_FakeRuntime, + server_runner=lambda **kwargs: calls.append(kwargs), + ) + + assert app is not None + manager = calls[0]["session_manager"] + assert manager._shared_scenario is None + assert callable(manager._shared_spec_factory) + assert manager.runtime_config.warmup_chunks == 0 + assert manager.runtime_config.encoder_backend == "default" + routes = {resource.canonical for resource in app.router.resources()} + assert "/api/session/input" in routes + + manager._shared_host.close() + app[PACKAGE_RESOURCE_STACK_KEY].close() + + +def test_uploaded_spec_uses_prepared_video_and_output_dimensions() -> None: + calls: list[dict[str, Any]] = [] + app = serve_flashvsr_webrtc_demo( + spec=_upload_spec(), + runtime_factory=_FakeRuntime, + server_runner=lambda **kwargs: calls.append(kwargs), + ) + assert app is not None + manager = calls[0]["session_manager"] + prepared = _prepared_video(target_height=256, target_width=384) + session_spec = manager._shared_spec_factory( + FlashVSRWebRTCSessionInput( + prepared_video=prepared, + original_name="upload.mp4", + ) + ) + + assert session_spec.scenario is prepared + assert isinstance(session_spec.output, WebRTCOutputSpec) + assert session_spec.output.video_height == 256 + assert session_spec.output.video_width == 384 + scenario = manager._shared_adapter.prepare_scenario(session_spec) + assert scenario.metadata[PREPARED_VIDEO_METADATA_KEY] is prepared + + manager._shared_host.close() + app[PACKAGE_RESOURCE_STACK_KEY].close() + + +@pytest.mark.asyncio +async def test_offer_requires_upload_when_no_default_input() -> None: + manager = _FakeSessionManager() + controller = FlashVSRUploadController( + manager=manager, + adapter=_FakeUploadAdapter(_prepared_video()), + spec=_upload_spec(), + default_video=None, + ) + client = await _build_upload_client(controller) + try: + status_response = await client.get("/api/session/input") + status = await status_response.json() + assert status_response.status == 200 + assert status["upload_required"] is True + + offer = await client.post( + "/api/webrtc/offer", + json={"sdp": "offer", "type": "offer"}, + ) + assert offer.status == 400 + assert "Upload an MP4" in await offer.text() + finally: + await client.close() + + +@pytest.mark.asyncio +async def test_mp4_upload_is_decoded_staged_and_temp_file_removed() -> None: + prepared = _prepared_video(target_height=256, target_width=384) + adapter = _FakeUploadAdapter(prepared) + manager = _FakeSessionManager() + controller = FlashVSRUploadController( + manager=manager, + adapter=adapter, + spec=_upload_spec(), + default_video=None, + ) + client = await _build_upload_client(controller) + try: + form = FormData() + form.add_field( + "video", + b"fake-mp4", + filename="../uploaded.mp4", + content_type="video/mp4", + ) + response = await client.post("/api/session/input", data=form) + payload = await response.json() + + assert response.status == 200 + assert payload["input_source"] == "uploaded" + assert payload["filename"] == "uploaded.mp4" + assert payload["resolution"] == {"width": 384, "height": 256} + assert adapter.upload_bytes == b"fake-mp4" + assert adapter.upload_path is not None + assert not adapter.upload_path.exists() + assert isinstance(manager.pending_session_input, FlashVSRWebRTCSessionInput) + assert manager.pending_session_input.prepared_video is prepared + finally: + await client.close() + + +@pytest.mark.asyncio +async def test_upload_rejects_empty_wrong_type_and_decode_failure() -> None: + adapter = _FakeUploadAdapter(_prepared_video()) + manager = _FakeSessionManager() + controller = FlashVSRUploadController( + manager=manager, + adapter=adapter, + spec=_upload_spec(), + default_video=None, + ) + client = await _build_upload_client(controller) + try: + wrong_type = FormData() + wrong_type.add_field( + "video", + b"not-video", + filename="input.txt", + content_type="text/plain", + ) + response = await client.post("/api/session/input", data=wrong_type) + assert response.status == 400 + + empty = FormData() + empty.add_field( + "video", + b"", + filename="empty.mp4", + content_type="video/mp4", + ) + response = await client.post("/api/session/input", data=empty) + assert response.status == 400 + + adapter.error = ValueError("decoder rejected container") + corrupt = FormData() + corrupt.add_field( + "video", + b"corrupt", + filename="corrupt.mp4", + content_type="video/mp4", + ) + response = await client.post("/api/session/input", data=corrupt) + assert response.status == 400 + assert "decoder rejected container" in await response.text() + finally: + await client.close() + + +@pytest.mark.asyncio +async def test_upload_returns_conflict_while_session_is_active() -> None: + manager = _FakeSessionManager() + manager.active = True + controller = FlashVSRUploadController( + manager=manager, + adapter=_FakeUploadAdapter(_prepared_video()), + spec=_upload_spec(), + default_video=None, + ) + client = await _build_upload_client(controller) + try: + form = FormData() + form.add_field( + "video", + b"fake-mp4", + filename="input.mp4", + content_type="video/mp4", + ) + response = await client.post("/api/session/input", data=form) + assert response.status == 409 + finally: + await client.close() + + +def test_torch_adain_accepts_center_cropped_noncontiguous_video() -> None: + torch.manual_seed(0) + content = torch.randn(1, 3, 8, 8, 16)[:, :, :5, :, 1:15] + style = torch.randn(1, 3, 8, 8, 16)[:, :, :5, :, 1:15] + assert not content.is_contiguous() + assert not style.is_contiguous() + corrector = FlashVSRColorCorrector(implementation="torch") + + expected = corrector( + content.contiguous(), + style.contiguous(), + method="adain", + ) + actual = corrector(content, style, method="adain") + + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + + +def test_flashvsr_browser_adapter_uploads_before_connect() -> None: + adapter_path = ( + Path(__file__).parents[1] / "flashvsr" / "demo" / "web" / "adapter.js" + ) + stylesheet_path = adapter_path.with_name("adapter.css") + source = adapter_path.read_text() + stylesheet = stylesheet_path.read_text() + + assert 'type="file"' in source + assert "video/mp4" in source + assert "/api/session/input" in source + assert 'class="flashvsrStartButton"' in source + assert 'action: { event: "step" }' in source + assert "beforeConnect" in source + assert "onConnect" in source + assert "sendCommand(START_ACTION" in source + assert 'startButton.addEventListener("click"' in source + assert " controls," not in source + assert "RTCPeerConnection" not in source + on_connect = source.split("onConnect() {", maxsplit=1)[1].split("},", maxsplit=1)[0] + assert "sendCommand" not in on_connect + assert ".flashvsrUploadCard" in stylesheet + assert ".flashvsrStartButton" in stylesheet + assert "position: absolute" in stylesheet + + +def test_shared_webrtc_shell_hides_empty_controls_and_notifies_connect() -> None: + web_dir = ( + Path(__file__).parents[3] + / "flashdreams" + / "flashdreams" + / "serving" + / "webrtc" + / "web" + ) + + assert 'id="controlCard"' in (web_dir / "request_session.html").read_text() + source = (web_dir / "request_session.js").read_text() + assert "syncControlCardVisibility" in source + assert "modelAdapter?.onConnect?.(modelContext)" in source + assert "setFlow," in source + assert "setStatus," in source diff --git a/uv.lock b/uv.lock index 4d0a32ba5..d5295f82f 100644 --- a/uv.lock +++ b/uv.lock @@ -1149,7 +1149,7 @@ name = "flashdreams-flashvsr" version = "0.1.0" source = { editable = "integrations/flashvsr" } dependencies = [ - { name = "flashdreams" }, + { name = "flashdreams", extra = ["serving"] }, { name = "grpcio" }, { name = "mediapy" }, { name = "pillow" }, @@ -1164,7 +1164,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "flashdreams", editable = "flashdreams" }, + { name = "flashdreams", extras = ["serving"], editable = "flashdreams" }, { name = "grpcio", specifier = ">=1.80" }, { name = "grpcio-tools", marker = "extra == 'dev'", specifier = ">=1.80" }, { name = "mediapy", specifier = ">=1.1" },