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/.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/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 new file mode 100644 index 000000000..260b24b24 --- /dev/null +++ b/docs/inference_runtime_api_design.md @@ -0,0 +1,786 @@ + + +# 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; +- `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; +- 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. + +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 +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 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 +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 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 OmniDreams demo; +- defer additional model migrations, hosted execution, full autotune, and polished + metrics until the first branch proves the API shape. + +## Task Tracker + +| 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`, `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 | 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 | 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. +- 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: + +- 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 + related; +- 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 + +```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 + 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 global and per-step InferenceInput 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: global conditioning, 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 + global conditioning: 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 global conditioning 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 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: 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. | + +## 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 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: + +- keyboard keydown/keyup events; +- reset requests; +- prompt or image selection/update events; +- 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. + +## CanonicalInputs And InferenceInput + +Inputs move through three layers: + +```text +UserInputs -> CanonicalInputs -> InferenceInput + raw canonicalized encoded +``` + +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. + +`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. + +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 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, +control tensors, event markers, or model-specific fields. + +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, 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 +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 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 + +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 +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 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 +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, 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 +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, 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 +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 +`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 or mapper + set and check obvious compatibility between the app event source and the + model; +- 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 +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. +- 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. + +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. +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 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. +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. + +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 +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 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. + +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 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 + +- `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 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`. 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. +- `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 `CanonicalInputs` to `InferenceInput` 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 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? +- 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 + 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; +- `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; +- 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. diff --git a/docs/inference_runtime_inputs_implementation.md b/docs/inference_runtime_inputs_implementation.md new file mode 100644 index 000000000..75460cced --- /dev/null +++ b/docs/inference_runtime_inputs_implementation.md @@ -0,0 +1,289 @@ + + +# 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 +- `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`. + +## 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. 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 + +The encoded layer splits model-facing inputs into two slots: + +- **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_conditioning", "step"]`. The phase names the +`InferenceInput` slot the caller provides. + +`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`. + +## Global Conditioning Is Session-Global State + +`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 InferenceInput, InferenceInputSchema, InputField + +schema = InferenceInputSchema( + global_conditioning_fields=( + InputField(name="prompt"), + InputField(name="scene_id"), + ) +) +schema.require_global_conditioning( + InferenceInput(global_conditioning={"prompt": "drive", "scene_id": "town_02"}) +) + +step_with_prompt_update = InferenceInput( + global_conditioning={"prompt": "heavy rain"}, + step={"steering": 0.0}, +) +``` + +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 + +`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_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 +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. + +`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 +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 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 +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. +- **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. + +## 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 \ + flashdreams/tests/test_runtime_runner.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_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/docs/inference_runtime_supported_inputs_inventory.md b/docs/inference_runtime_supported_inputs_inventory.md new file mode 100644 index 000000000..f82f3192b --- /dev/null +++ b/docs/inference_runtime_supported_inputs_inventory.md @@ -0,0 +1,354 @@ + + +# 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 global conditioning: prompt text plus latent/output height and width + derived from run config. +- 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 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 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 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 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 global conditioning: prompt text, first-frame tensor, base text + embeddings, precomputed text-event embeddings, base intrinsics, and world + scale. +- 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. + +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 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: + +- 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 global conditioning: prompt list, first-frame tensor, view names; or + precomputed text/image/negative-text embeddings. +- 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 global conditioning: scene data, renderer, first-frame tensor, prompt, + camera calibration/extrinsics, initial ego pose, and initial timestamp. +- 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. + +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 global conditioning: scene bundle, selected camera, prompt, initial + RGB frame, initial rig pose, and initial timestamp. +- 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 global conditioning: synthetic transformer context, optional negative + context, height, and width. +- 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 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`: + +- 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 global conditioning: decoder context such as prompt, fps, + `save_stage1`, refiner seed, sink size, and streaming refiner window/block + parameters. +- 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 + 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 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 +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 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, `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 +`[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 adapter 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 embeddings 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 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 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`. Global conditioning is session-global state; `step` is the + payload for one generated chunk or frame window. +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, 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 + 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 `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`. +- 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_conditioning_fields=( + 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"), + ), +) +``` + +```python +omnidreams_model = InferenceInputSchema( + description="omnidreams", + global_conditioning_fields=( + 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, + 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"), + ), +) +``` + +```python +hy_worldplay_model = InferenceInputSchema( + description="hy-worldplay", + global_conditioning_fields=( + 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"), + InputField(name="memory_config", frequency_consumed="per_step"), + ), +) +``` + +```python +sana_wm_model = InferenceInputSchema( + description="sana-wm", + global_conditioning_fields=( + 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", + input_modality="c2w_sequence", + frequency_consumed="per_step", + metadata={"shape": "[F,4,4]", "coordinates": "opencv_c2w"}, + ), + InputField( + name="camera_intrinsics_vec4", + required=False, + input_modality="intrinsics_vec4_sequence", + frequency_consumed="per_step", + 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/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/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/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 2ef2fc296..060adbd31 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: @@ -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,19 +186,16 @@ def create_video_output_stream( self, *, fps: float | None = None, - move_to_cpu: bool = True, - ) -> RunnerVideoOutputStream: - """Create the standard runner video output stream for one rollout.""" + ) -> VideoOutputStream: + """Create the standard post-processing 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, - 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 930bcf572..199c36b5c 100644 --- a/flashdreams/flashdreams/infra/video_output.py +++ b/flashdreams/flashdreams/infra/video_output.py @@ -18,14 +18,17 @@ from __future__ import annotations from collections.abc import Mapping -from dataclasses import dataclass, field -from typing import Any +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: @@ -41,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)]) @@ -132,152 +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 RunnerVideoOutputStream: - """Post-process, collect, and summarize runner video chunks.""" +class VideoOutputStream: + """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, - ) -> None: - """Process one generated chunk and collect it when this rank writes output.""" + 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) + 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, + output_window=output_window, + metrics=metrics, + metadata=result_metadata, + ) - def finish(self) -> Tensor | None: - """Flush post-processing and return the collected rank-zero video.""" + def finish(self) -> StepResult | None: + """Close the stream and return a post-processing tail, when present.""" 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() + 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, + metadata={"postprocess_tail": True}, + ) + + +class VideoResultCollector: + """Collect video results for persistence or composed presentation.""" - def _append_if_nonempty(self, output: Tensor) -> None: - if not self.collect_output or output.shape[self._time_dim] == 0: + def __init__( + self, + *, + 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 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) @@ -288,12 +272,43 @@ def _collected_output(self) -> Tensor | None: return output +def prepare_video_for_mp4( + video: Tensor, + *, + layout: VideoTensorLayout | 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, 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 video layout for MP4: {layout!r}") + + __all__ = [ "LazyRGBFrame", - "RunnerVideoOutputStream", - "VideoStepResult", + "VideoOutputStream", + "VideoResultCollector", "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/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 new file mode 100644 index 000000000..3280e6862 --- /dev/null +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -0,0 +1,145 @@ +# 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.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, + CanonicalInputs, + CanonicalInputSchema, + CanonicalModality, + InferenceInput, + InferenceInputSchema, + InputField, + InputPhase, + TimeWindow, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, + validate_phase, +) +from flashdreams.runtime.interfaces import ( + InferenceRuntime, + 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, + InputMapping, + InputMappingSchema, + MappingCompatibility, + check_mapping_compatibility, + check_mapping_set_compatibility, + combine_mapping_schemas, + undeclared_inference_inputs, +) +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, + StepRequirements, + StepResult, + step_requirements_from_request, +) +from flashdreams.runtime.video_output import Mp4VideoOutputTarget +from flashdreams.runtime.worker import ModelExecutionWorker, ThreadAffineRuntimeWorker + +__all__ = [ + "CanonicalInputs", + "CanonicalInputSchema", + "CanonicalModality", + "check_mapping_compatibility", + "check_mapping_set_compatibility", + "combine_mapping_schemas", + "DeclaresMappingSchema", + "DEFAULT_DRIVING_BINDINGS", + "DEFAULT_SUPPORTED_KEYS", + "DeviceConverter", + "DeviceConverterSchema", + "DRIVING_SUPPORTED_KEYS", + "DRIVER_COMMAND", + "ExecutionBackend", + "IdentityInputMapping", + "InferenceConfig", + "InferenceInput", + "InferenceInputSchema", + "InferenceRuntime", + "InferenceSession", + "InMemoryMetricsRecorder", + "INPUT_PHASES", + "InputCanonicalizer", + "InputField", + "InputMapping", + "InputMappingSchema", + "InputPhase", + "ImageRequest", + "KEY_ALIASES", + "KeyboardState", + "KeyboardToDriverCommand", + "MappingCompatibility", + "MetricsRecorder", + "MetricsSnapshot", + "ModelAdapter", + "ModelExecutionWorker", + "Mp4VideoOutputTarget", + "NullMetricsRecorder", + "NullOutputTarget", + "OutputArtifact", + "OutputTarget", + "Precision", + "PromptRequest", + "ResetRequest", + "RuntimeMetricSample", + "ScriptedModality", + "SparseInputSnapshot", + "StepRequest", + "StepRequirements", + "StepResult", + "TimeWindow", + "ThreadAffineRuntimeWorker", + "run_inference_session", + "step_requirements_from_request", + "undeclared_inference_inputs", + "UserInputCapability", + "UserInputEvent", + "UserInputs", + "UserInputSchema", + "WSAD_SUPPORTED_KEYS", + "normalize_key", + "validate_phase", +] 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/canonical.py b/flashdreams/flashdreams/runtime/canonical.py new file mode 100644 index 000000000..bf3960380 --- /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.runtime.keyboard 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/config.py b/flashdreams/flashdreams/runtime/config.py new file mode 100644 index 000000000..f1c0c2a0c --- /dev/null +++ b/flashdreams/flashdreams/runtime/config.py @@ -0,0 +1,84 @@ +# 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.""" + + 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.""" + + 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.") + 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) + ) + object.__setattr__(self, "resource_hints", freeze_mapping(self.resource_hints)) diff --git a/flashdreams/flashdreams/runtime/demo/__init__.py b/flashdreams/flashdreams/runtime/demo/__init__.py new file mode 100644 index 000000000..8d61e31bd --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/__init__.py @@ -0,0 +1,178 @@ +# 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.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, + PreparedScenario, + WebRTCAppResources, + WebRTCOutputSpec, +) +from flashdreams.runtime.demo.timing import ( + ActivationPolicy, + ActivationResult, + ActivationSignal, + AlwaysActiveActivationPolicy, + CatchUpDecision, + CatchUpPolicy, + DeterministicClock, + RealtimeClock, + RealtimeEventInputSource, + RealtimeEventResampler, + 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", + "Mp4ErrorPolicy", + "Mp4OutputSink", + "Mp4OutputSpec", + "NativeWindowErrorPolicy", + "NoopTransportService", + "NullOutputSpec", + "NullOutputSink", + "NullErrorPolicy", + "OutputDecision", + "OutputSinkFactory", + "OutputSpec", + "OutputSink", + "PreparedScenario", + "PreparedStep", + "ProviderCapabilities", + "RealtimeInputSource", + "RealtimeClock", + "RealtimeEventInputSource", + "RealtimeEventResampler", + "RealtimeSessionDriver", + "RealtimeWindowResult", + "ResolvedRunCapabilities", + "ResamplerRealtimeClock", + "RunContext", + "RunMode", + "RunModeCapabilities", + "RunModeWarmup", + "RunResult", + "RunSummary", + "RuntimeHost", + "SessionEdges", + "SessionDriver", + "SessionInfo", + "SignalActivationPolicy", + "SingleSessionAdmissionPolicy", + "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 new file mode 100644 index 000000000..b763e4eaf --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/app.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared command lifecycle for model demo applications.""" + +from __future__ import annotations + +import argparse +import sys +from abc import ABC, abstractmethod +from typing import Any + +import torch +import torch.distributed as dist + +from flashdreams.core.distributed import init as distributed_init +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): + """Base command application shared by model replay and WebRTC demos.""" + + 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": + 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( + 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}") + + @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.""" + + @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/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/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 new file mode 100644 index 000000000..866acbb00 --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/outputs.py @@ -0,0 +1,305 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""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.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, + *, + 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__ = [ + "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 new file mode 100644 index 000000000..fa2f507a9 --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/replay.py @@ -0,0 +1,629 @@ +# 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 + +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.types import ( + StepRequest, + StepRequirements, + StepResult, + step_requirements_from_request, +) + +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 | None = None, + output_sink_factory: OutputSinkFactory = build_output_sink, + metrics: MetricsRecorder | None = None, + 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(), + 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.") + + 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() + artifacts = tuple( + runner( + adapter=adapter, + config=_require_config(spec), + 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, + ) + ) + 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( + *, + 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", + "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 new file mode 100644 index 000000000..d7a6a309b --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/spec.py @@ -0,0 +1,197 @@ +# 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 Callable, Mapping, Sequence +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 + +from .host import WarmupSessionInputs + + +@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 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.""" + + __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): + """Transport-neutral model adapter consumed by demo runners.""" + + 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.""" + ... + + +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", + "PreparedScenario", + "WebRTCOutputSpec", + "WebRTCAppResources", +] diff --git a/flashdreams/flashdreams/runtime/demo/timing.py b/flashdreams/flashdreams/runtime/demo/timing.py new file mode 100644 index 000000000..e6609cba9 --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/timing.py @@ -0,0 +1,409 @@ +# 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"] + + +@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) + + +class _RealtimeTimeline(Protocol): + dt: float + next_chunk_start_v: float + + def reset(self, *, start_v: float) -> None: ... + + def sample_chunk( + self, + num_frames: int, + ) -> 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 a resampler'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 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.") + + 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 RealtimeEventInputSource: + """Realtime input source backed by raw event windows.""" + + resampler: _RealtimeTimeline + 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( + "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 " + "event-window analog yet." + ) + + def is_finished(self) -> bool: + return False + + 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 + 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(), + ) + 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", + "RealtimeEventInputSource", + "RealtimeEventResampler", + "RealtimeClock", + "RealtimeWindowResult", + "ResamplerRealtimeClock", + "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/demo/webrtc.py b/flashdreams/flashdreams/runtime/demo/webrtc.py new file mode 100644 index 000000000..2e835a6e9 --- /dev/null +++ b/flashdreams/flashdreams/runtime/demo/webrtc.py @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deprecated location for WebRTC demo construction helpers. + +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__: list[str] = [] diff --git a/flashdreams/flashdreams/runtime/inputs.py b/flashdreams/flashdreams/runtime/inputs.py new file mode 100644 index 000000000..70260de2e --- /dev/null +++ b/flashdreams/flashdreams/runtime/inputs.py @@ -0,0 +1,449 @@ +# 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, Literal, cast + +from flashdreams.infra.time import TimeWindow +from flashdreams.runtime._utils import freeze_mapping + +InputPhase = Literal["global_conditioning", "step"] + +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_conditioning' or 'step', got {value!r}." + ) + return cast(InputPhase, value) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class InputField: + """Lightweight schema field for user snapshots or model inputs. + + ``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 + input_modality: str | None = None + frequency_consumed: 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 + input_modality: 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 + 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 + ) + + +@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, ...] = () + capabilities: tuple[UserInputCapability, ...] = () + 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.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``.""" + 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 InferenceInputSchema: + """Minimal metadata for global conditioning and per-step inputs.""" + + global_conditioning_fields: tuple[InputField, ...] = () + """Model inputs carried in the global conditioning slot.""" + + step_fields: tuple[InputField, ...] = () + """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_conditioning_fields + if validate_phase(phase) == "global_conditioning" + 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 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_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_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}" + ) + + def require_step(self, inputs: "InferenceInput") -> 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 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 + + 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. 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. + """ + + __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, "global_conditioning", freeze_mapping(self.global_conditioning) + ) + object.__setattr__(self, "step", freeze_mapping(self.step)) + object.__setattr__(self, "metadata", freeze_mapping(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_conditioning" + else self.step + ) + + +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..5c5054dc4 --- /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 ( + CanonicalInputSchema, + InferenceInput, + InferenceInputSchema, +) +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: + """Return the next step's runtime request, or ``None`` when complete.""" + ... + + def step(self, inputs: InferenceInput) -> StepResult: + """Run one sequential inference step.""" + ... + + def reset(self, inputs: InferenceInput | 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: InferenceInput) -> InferenceSession: + """Create an isolated session from global conditioning 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, 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 + def model_id(self) -> str: + """Stable identity for the model adapter or runtime integration.""" + ... + + @property + def inference_input_schema(self) -> InferenceInputSchema: + """Model-facing global conditioning and per-step input requirements.""" + ... + + @property + 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 canonical-to-model mapping.""" + ... + + 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/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/runtime/mapping.py b/flashdreams/flashdreams/runtime/mapping.py new file mode 100644 index 000000000..6dfb5cc45 --- /dev/null +++ b/flashdreams/flashdreams/runtime/mapping.py @@ -0,0 +1,384 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Input mapping boundary from canonical inputs to encoded inference inputs.""" + +from __future__ import annotations + +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 ( + INPUT_PHASES, + CanonicalInputs, + CanonicalInputSchema, + CanonicalModality, + InferenceInput, + InferenceInputSchema, + InputField, + InputPhase, +) +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, + *, + canonical_schema: CanonicalInputSchema | None = None, + inference_input_schema: InferenceInputSchema | None = None, + ) -> None: + """Fail early for obvious app, event-source, and model mismatches.""" + ... + + def map_global_conditioning_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + ) -> InferenceInput: + """Build global conditioning inputs for session start or reset.""" + ... + + def map_step_inputs( + self, + *, + canonical_inputs: CanonicalInputs, + inference_input: InferenceInput, + request: StepRequest, + ) -> InferenceInput: + """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, + *, + 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, 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_conditioning: 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_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``.""" + 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 + input_modality_ok = ( + produced.input_modality is None + or required.input_modality is None + or produced.input_modality == required.input_modality + ) + return input_modality_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_conditioning": [], + "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_conditioning=tuple(produces["global_conditioning"]), + 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_conditioning_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/metrics.py b/flashdreams/flashdreams/runtime/metrics.py new file mode 100644 index 000000000..8a8fd6b1b --- /dev/null +++ b/flashdreams/flashdreams/runtime/metrics.py @@ -0,0 +1,336 @@ +# 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 import Counter, defaultdict +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 + + +@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)) + + +@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, session, and run 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 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.""" + ... + + +@dataclass(slots=True) +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: + 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 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: + """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 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/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/runner.py b/flashdreams/flashdreams/runtime/runner.py new file mode 100644 index 000000000..f9a83411f --- /dev/null +++ b/flashdreams/flashdreams/runtime/runner.py @@ -0,0 +1,669 @@ +# 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 collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any + +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, + 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 ( + 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 + + +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 shared batch driver. + + 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. + """ + + 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: + _validate_legacy_runner_inputs( + adapter=adapter, + config=config, + mapping=mapping, + canonicalizer=canonicalizer, + source_schema=source_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, + ) + input_source = _LegacyBatchInputSource( + source_schema=source_schema, + user_inputs=user_inputs, + request_state=request_state, + ) + result = _run_shared_batch_session( + host=host, + provider=provider, + input_source=input_source, + output_sink=output_sink, + metrics=metrics_recorder, + ) + _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: + 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, + ) + + 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 + + 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( + *, + 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, + ) + + +__all__ = ["run_inference_session"] diff --git a/flashdreams/flashdreams/runtime/types.py b/flashdreams/flashdreams/runtime/types.py new file mode 100644 index 000000000..49ee958a0 --- /dev/null +++ b/flashdreams/flashdreams/runtime/types.py @@ -0,0 +1,130 @@ +# 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.infra.results import StepResult +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: + """Per-step runtime request emitted by an inference session. + + 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 + + step_index: int + inference_input_schema: InferenceInputSchema | 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)) + + +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/video_output.py b/flashdreams/flashdreams/runtime/video_output.py new file mode 100644 index 000000000..c92cf2032 --- /dev/null +++ b/flashdreams/flashdreams/runtime/video_output.py @@ -0,0 +1,103 @@ +# 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 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.output import OutputArtifact +from flashdreams.runtime.types import StepResult + +VideoWriter = Callable[..., Path] + + +@dataclass(slots=True) +class Mp4VideoOutputTarget: + """Write layout-aware runtime step results 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 + enabled: bool = True + _opened: bool = field(default=False, init=False, repr=False) + _collector: VideoResultCollector | None = field( + default=None, + init=False, + repr=False, + ) + + @property + def closed(self) -> bool: + return not self._opened + + def open(self) -> None: + self._collector = VideoResultCollector( + output_layout=self.output_layout, + 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._collector is None: + raise RuntimeError("Cannot write to a closed output target.") + if result.layout is None: + raise TypeError( + "Mp4VideoOutputTarget requires a video StepResult with layout." + ) + if result.layout != self.output_layout: + raise ValueError( + "Mp4VideoOutputTarget received layout " + f"{result.layout!r}; expected {self.output_layout!r}." + ) + self._collector.add(result) + + def close(self) -> Sequence[OutputArtifact]: + if self._collector is None: + self._opened = False + return () + + collector = self._collector + self._collector = None + self._opened = False + video = collector.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, + "shape": tuple(int(dim) for dim in video.shape), + "stats_history": tuple(collector.stats_history), + }, + ), + ) + + +__all__ = ["Mp4VideoOutputTarget"] diff --git a/flashdreams/flashdreams/runtime/worker.py b/flashdreams/flashdreams/runtime/worker.py new file mode 100644 index 000000000..e523b6e6a --- /dev/null +++ b/flashdreams/flashdreams/runtime/worker.py @@ -0,0 +1,180 @@ +# 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 +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 ModelExecutionWorker: + """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._state_lock = threading.Lock() + self._accepting = True + self._closed = False + 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], + /, + *args: Any, + **kwargs: Any, + ) -> _T: + """Run one callable after all previously submitted worker calls.""" + self._require_not_worker_thread() + self._require_accepting() + future = self._submit(func, args, kwargs) + try: + return await _await_executor_future(future) + except asyncio.CancelledError: + 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.""" + 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.""" + 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, + 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: + 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], + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> _T: + return func(*args, **kwargs) + + +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__ = ["ModelExecutionWorker", "ThreadAffineRuntimeWorker"] diff --git a/flashdreams/flashdreams/serving/output_targets.py b/flashdreams/flashdreams/serving/output_targets.py index b7e7be981..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,179 +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 = [ - "--config_name", - _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))) - _append_webrtc_bind_args(argv, options) - return OutputTargetSpec( - mode="webrtc", - label="LingBot WebRTC server", - module="lingbot.webrtc.server", - 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/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/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/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/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/encoders.py b/flashdreams/flashdreams/serving/webrtc/encoders.py index bb27c7f6c..b07d9324c 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 @@ -22,12 +20,14 @@ 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 from loguru import logger +from flashdreams.runtime import StepResult + if TYPE_CHECKING: from flashdreams.serving.webrtc.media import BufferedVideoTrack, NVENCVideoTrack @@ -69,9 +69,23 @@ 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, - chunk: torch.Tensor, + result: StepResult, track: MediaStreamTrack, *, force_keyframe: bool = False, @@ -110,9 +124,23 @@ def create_track(self, *, maxsize: int) -> BufferedVideoTrack: return BufferedVideoTrack(fps=self.fps, maxsize=maxsize) - async def deliver_chunk( + def prepare_chunk_payload( self, - chunk: torch.Tensor, + 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, @@ -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_chunk(chunk) + 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 e752be2a4..59177742e 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 Callable, Mapping from collections.abc import Set as AbstractSet -from dataclasses import dataclass, field -from enum import IntEnum -from typing import Any, Generic, TypeVar +from dataclasses import dataclass, field, replace +from typing import Any, Generic, TypeVar, cast from aiortc import ( RTCConfiguration, @@ -23,10 +23,43 @@ ) from loguru import logger -from flashdreams.serving.realtime.input import KeyboardResampler +from flashdreams.runtime.demo import ( + DemoSpec, + InMemorySessionMetricsRecorder, + ModelInputProvider, + PreparedScenario, + PreparedStep, + ProviderCapabilities, + RealtimeEventResampler, + ResamplerRealtimeClock, + RunContext, + 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.keyboard import DEFAULT_SUPPORTED_KEYS, normalize_key +from flashdreams.runtime.mapping import InputMapping +from flashdreams.runtime.types import StepRequest, StepResult 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 ( @@ -39,12 +72,22 @@ make_event_ack_payload, ) from flashdreams.serving.webrtc.runtime import ( + WebRTCControlSignal, WebRTCRuntimeConfig, - WebRTCSessionRuntime, - WebRTCStepResult, - make_webrtc_step_result, ) 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, @@ -54,8 +97,7 @@ "BaseWebRTCSessionManager", "ManagedWebRTCSession", "WebRTCControlSignal", - "WebRTCStepResult", - "make_webrtc_step_result", + "StepResult", ] # Close the active session if no client heartbeat/control message arrives @@ -65,35 +107,490 @@ # 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 - -_RuntimeT = TypeVar("_RuntimeT", bound=WebRTCSessionRuntime) +_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"}) +_SESSION_INPUT_KEY = "webrtc_session_input" +_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) -def _stat_float(stats: dict[str, float], name: str, default: float = 0.0) -> float: +class _InferenceSessionExhausted(RuntimeError): + """Raised when an ``InferenceSession`` reports normal completion.""" + + +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: 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))) -class WebRTCControlSignal(IntEnum): - """Rank-orchestration signals shared by the single-session runtimes.""" +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 + - INITIALIZE = 0 - RESET_SESSION = 1 - ACTION_STEP = 2 - CLOSE = 3 - EVENT = 4 - EXIT = 99 +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(_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) + 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) @@ -104,11 +601,23 @@ 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) + 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``.""" + 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 @@ -137,6 +646,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() @@ -146,11 +660,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__( @@ -159,12 +668,33 @@ 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, + 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, + legacy_segment_resampler_factory: Callable[..., Any] | None = None, ) -> 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 @@ -172,36 +702,75 @@ def __init__( self._active_session: ManagedWebRTCSession | None = None self._preload_lock = asyncio.Lock() self._session_lock = asyncio.Lock() - - def _model_name(self) -> str: - """Human-readable model identifier reported in ``chunk_done``.""" - raise NotImplementedError - - def _peek_pending_session_input(self) -> Any: - """Session input applied to the next ``create_answer`` (or ``None``).""" - return None - - 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 _make_resampler(self, *, start_v: float) -> KeyboardResampler: + 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_spec_factory = shared_spec_factory + 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: + """Input that will be applied to the next successfully negotiated session.""" + return self._pending_session_input + + @property + def runtime(self) -> _RuntimeT: + """Model runtime driven by this transport manager.""" + return self._runtime + + 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) -> 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: - if self._resampler_supported_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), + ) -> 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() + 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: + 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: @@ -223,55 +792,110 @@ 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) + 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( + peek_input_fps(), + label="peek_input_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", + 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): + peek_output_frames = getattr(runtime, "peek_steady_output_num_frames", None) + if callable(peek_output_frames): return self._positive_int_runtime_value( - method(), + 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_chunk_num_frames(), - label="peek_steady_chunk_num_frames", + 1, + label="fallback steady output frame count", ) - 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. """ 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. @@ -329,7 +953,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. @@ -341,16 +965,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 {} - async def _handle_event_message( self, *, @@ -375,6 +989,43 @@ 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: @@ -404,6 +1055,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 _runtime_drives_inference_session(runtime) + + 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 @@ -413,29 +1331,61 @@ 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() 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( @@ -448,45 +1398,80 @@ 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.") - - await self._reset_runtime_for_session(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), - ) + raise RuntimeError("Runtime is not initialized.") + 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(), - ) + 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), + ) + 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, + video_track=video_track, + video_encoder=video_encoder, + peer_connection=peer_connection, + resampler=resampler, + legacy_segment_resampler=legacy_segment_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( @@ -497,10 +1482,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: @@ -511,15 +1499,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") @@ -531,11 +1525,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) @@ -549,7 +1558,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.") @@ -559,13 +1571,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, ) @@ -614,15 +1626,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, @@ -634,6 +1668,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.")) @@ -650,6 +1688,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 @@ -705,28 +1749,240 @@ 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() - managed_session.resampler.on_edge(arrival_t=arrival_t, event=event, key=key) + 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 + 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. 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 + spec_factory = self._shared_spec_factory + 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() + 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( + 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: - """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 - backpressure on ``BufferedVideoTrack.enqueue_chunk``. + *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() runtime = managed_session.runtime @@ -758,8 +2014,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. @@ -773,17 +2029,34 @@ 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() - segments, frame_times = resampler.sample_chunk(input_num_frames) + chunk_start_v = resampler.next_chunk_start_v + 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( + start_s=chunk_start_v, + end_s=chunk_end_v, + ), + ) consumed_action_arrivals: list[float] = [] while ( managed_session.pending_action_arrivals @@ -793,21 +2066,44 @@ 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.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 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, ) @@ -816,7 +2112,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 @@ -824,17 +2120,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={} " @@ -845,7 +2140,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, @@ -880,9 +2175,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, @@ -897,13 +2192,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, @@ -911,7 +2206,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..25fa09437 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,16 +81,37 @@ def maxsize(self) -> int: def qsize(self) -> int: return self._frames.qsize() - async def enqueue_chunk(self, video_chunk: torch.Tensor) -> 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, video_chunk) 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 @@ -220,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 ade33cb49..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,7 +31,9 @@ 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 # Runtime imports ``PyNvVideoCodec`` unconditionally (the isolation @@ -56,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 @@ -71,13 +81,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 +102,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] → @@ -258,9 +249,38 @@ 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, - chunk: torch.Tensor, + result: StepResult, track: MediaStreamTrack, *, force_keyframe: bool = False, @@ -297,7 +317,64 @@ 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, + ) + 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, + ) + + 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, ) @@ -317,18 +394,34 @@ 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) + 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] 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 79fa9e0e2..1472807ca 100644 --- a/flashdreams/flashdreams/serving/webrtc/runtime.py +++ b/flashdreams/flashdreams/serving/webrtc/runtime.py @@ -1,47 +1,43 @@ # 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 -from collections.abc import Awaitable, Mapping -from typing import Any, Protocol +import asyncio +from abc import ABC, abstractmethod +from collections.abc import Awaitable +from enum import IntEnum +from typing import Any, Generic, Protocol, TypeVar import torch - -from flashdreams.infra.postprocess import VideoTensorLayout -from flashdreams.infra.video_output import VideoStepResult, infer_video_num_frames -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 {}), - ) +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.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 + SESSION_STEP = 5 + SESSION_CLOSE = 6 + EXIT = 99 class WebRTCRuntimeConfig(Protocol): @@ -53,37 +49,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, *, - segments: list[PoseSegment], + request: StepRequest, + segments: list[Any], frame_times: list[float], - ) -> WebRTCStepResult: ... + ) -> StepResult: ... async def close(self) -> None: ... @@ -96,13 +104,229 @@ def trigger_event( ) -> dict[str, Any] | Awaitable[dict[str, Any]]: ... -class WebRTCServerLifecycle(Protocol): - """Distributed worker lifecycle used by the shared WebRTC serve loop.""" +_ConfigT = TypeVar("_ConfigT", bound=ThreadAffineWebRTCRuntimeConfig) +_SessionInputT = TypeVar("_SessionInputT") - def send_exit_signal(self) -> None: ... - def wait_for_termination(self) -> None: ... +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. + """ + + MASTER_RANK = 0 + + 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 is_master(self) -> bool: + return self.rank == self.MASTER_RANK + + @property + 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[Any], + 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[Any], + 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() + + @abstractmethod + def _is_runtime_initialized(self) -> bool: ... + + @abstractmethod + def _runtime_step_index(self) -> int: ... + + @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[Any], + 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/server.py b/flashdreams/flashdreams/serving/webrtc/server.py index a9e386533..860640a6f 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,14 @@ 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-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-v4" + return web.json_response(payload) + async def on_startup(app: web.Application) -> None: manager = app[SESSION_MANAGER_KEY] logger.info("Preloading {} runtime on startup.", preload_name) @@ -104,7 +113,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 +129,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 +149,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/services.py b/flashdreams/flashdreams/serving/webrtc/services.py new file mode 100644 index 000000000..cfcb11029 --- /dev/null +++ b/flashdreams/flashdreams/serving/webrtc/services.py @@ -0,0 +1,1232 @@ +# 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 ( + 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 + 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( + 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) + if self.legacy_segment_resampler is not None: + self.legacy_segment_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) + frame_times = tuple(self.resampler.sample_chunk(input_frame_count)) + end_s = float(self.resampler.next_chunk_start_v) + 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) + ) + 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'.", + ) + 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", + 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._delivery_lock = asyncio.Lock() + 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: + 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, + 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: + 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": "stale generation"}, + ) + 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 + 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 + 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 _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( + 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) + + 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): + 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/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/integrations/lingbot/lingbot/webrtc/web/assets/horizontal-dark.svg b/flashdreams/flashdreams/serving/webrtc/web/assets/horizontal-dark.svg similarity index 100% rename from integrations/lingbot/lingbot/webrtc/web/assets/horizontal-dark.svg rename to flashdreams/flashdreams/serving/webrtc/web/assets/horizontal-dark.svg diff --git a/integrations/lingbot/lingbot/webrtc/web/assets/horizontal-light.svg b/flashdreams/flashdreams/serving/webrtc/web/assets/horizontal-light.svg similarity index 100% rename from integrations/lingbot/lingbot/webrtc/web/assets/horizontal-light.svg rename to flashdreams/flashdreams/serving/webrtc/web/assets/horizontal-light.svg diff --git a/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py b/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py new file mode 100644 index 000000000..867dd3e02 --- /dev/null +++ b/flashdreams/flashdreams/serving/webrtc/web/mock_ui_server.py @@ -0,0 +1,160 @@ +# 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. + +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 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") + + +class MockUIRequestHandler(SimpleHTTPRequestHandler): + """Serve the static viewer without preloading a model runtime.""" + + def __init__( + self, + request: socket | tuple[bytes, socket], + client_address: tuple[str, int], + server: BaseServer, + *, + directory: str | PathLike[str] | None = None, + model_web_dir: Path | None = None, + ) -> None: + self.model_web_dir = model_web_dir + super().__init__( + request, + client_address, + server, + directory=directory, + ) + + def _rewrite_path(self) -> bool: + path = urlsplit(self.path).path + if path == "/": + self.send_response(302) + self.send_header("Location", "/request_session?mock=1") + self.end_headers() + return True + if path == "/request_session": + self.path = "/request_session.html" + elif path.startswith("/static/"): + self.path = "/" + path.removeprefix("/static/") + 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 + 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-v4" + if (self.model_web_dir / "adapter.css").is_file(): + ui_config["model_stylesheet"] = ( + "/model-static/adapter.css?v=model-ui-v4" + ) + 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))) + 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), + 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" + ) + try: + server.serve_forever() + except KeyboardInterrupt: + print("\nStopping mock UI server.") + finally: + server.server_close() + + +if __name__ == "__main__": + main() diff --git a/integrations/omnidreams/omnidreams/webrtc/web/request_session.css b/flashdreams/flashdreams/serving/webrtc/web/request_session.css similarity index 97% rename from integrations/omnidreams/omnidreams/webrtc/web/request_session.css rename to flashdreams/flashdreams/serving/webrtc/web/request_session.css index 890c36147..53325e01b 100644 --- a/integrations/omnidreams/omnidreams/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; } @@ -282,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/integrations/omnidreams/omnidreams/webrtc/web/request_session.html b/flashdreams/flashdreams/serving/webrtc/web/request_session.html similarity index 70% rename from integrations/omnidreams/omnidreams/webrtc/web/request_session.html rename to flashdreams/flashdreams/serving/webrtc/web/request_session.html index 263a1b299..bd2bdcd86 100644 --- a/integrations/omnidreams/omnidreams/webrtc/web/request_session.html +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.html @@ -8,16 +8,17 @@ - Omnidreams WebRTC Drive - + FlashDreams WebRTC Drive +
-
-

Omnidreams WebRTC Drive

+
+

FlashDreams WebRTC Drive

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

Omnidreams WebRTC Drive

Idle - + +
Flow waiting
-
+
+ +
@@ -86,12 +81,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 75% rename from integrations/omnidreams/omnidreams/webrtc/web/request_session.js rename to flashdreams/flashdreams/serving/webrtc/web/request_session.js index 11b7302e8..e61f88407 100644 --- a/integrations/omnidreams/omnidreams/webrtc/web/request_session.js +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.js @@ -1,6 +1,23 @@ // 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") + +/** + * @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 {(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 + * @property {(context: Object) => void=} onDisconnect + */ + const connectButton = document.getElementById("connectButton") const statusText = document.getElementById("statusText") const flowText = document.getElementById("flowText") @@ -15,9 +32,13 @@ 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 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") -const allowedKeys = new Set(["w", "a", "s", "d"]) const keyAliases = new Map([ ["arrowup", "w"], ["arrowleft", "a"], @@ -32,6 +53,11 @@ const pendingActions = [] const maxPendingActions = 32 const heartbeatIntervalMs = 2000 +let allowedKeys = new Set() +let controlButtons = [] +/** @type {WebRTCModelAdapter|null} */ +let modelAdapter = null + let peerConnection = null let controlChannel = null let statsTimer = null @@ -41,7 +67,7 @@ let inferenceInFlight = false let connected = false let disconnecting = false let heldKeySequence = 0 -let postprocessControlAvailable = false +let postprocessAvailable = false const metrics = { fps: null, @@ -50,7 +76,7 @@ const metrics = { rttMs: null, resolution: null, step: null, - model: "Omnidreams", + model: "World Model", } function normalizeKey(rawKey) { @@ -58,6 +84,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 }) } @@ -83,7 +124,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 { @@ -120,44 +161,84 @@ function setFlow(message) { function setVideoVisible(visible) { document.body.classList.toggle("has-video", 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() + 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) + } + controlButtons = Array.from(controlRows.querySelectorAll("[data-control-key]")) + syncControlCardVisibility() } function setPostprocessDisabled(disabled) { - postprocessSelect.disabled = disabled || !postprocessControlAvailable + postprocessSelect.disabled = disabled || !postprocessAvailable } async function loadPostprocessOptions() { - const response = await fetch("/api/postprocess/options") - 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() { +async function configurePostprocessSession() { + if (!postprocessAvailable || mockMode) { + return + } const postprocessPreset = postprocessSelect.value const response = await fetch("/api/session/input", { method: "POST", @@ -168,10 +249,107 @@ 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, + setFlow, + setStatus, + 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 = {} + 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") { + adapter = module.default + } + } + } + } catch (error) { + logEvent(`model UI unavailable: ${error.message}`, { source: "client", level: "error" }) + } + + 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 = href + document.head.append(stylesheet) + } + const modelControls = Array.isArray(adapter.controls) ? adapter.controls : [] + renderControls(modelControls) + if (typeof adapter.modelName === "string") { + modelContext.setModelName(adapter.modelName) + } + if (adapter.capabilities?.postprocess === 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) + syncControlCardVisibility() } function renderMetrics() { @@ -181,7 +359,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) { @@ -423,6 +601,7 @@ function sendControlAction(action) { action, }) ) + modelAdapter?.onActionSent?.(action, modelContext) recordActionSent(action) setStatus("Generating", "generating") setFlow(`sent ${actionLabel(action)}, waiting=${inferenceInFlight}`) @@ -437,15 +616,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)) { @@ -527,9 +697,11 @@ 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 + } + + if (modelAdapter?.onControlMessage?.(payload, modelContext)) { return } @@ -633,7 +805,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 +827,7 @@ async function dumpPeerStats(reason) { } console.groupEnd() } catch (error) { - console.warn("[Omnidreams WebRTC] getStats failed", error) + console.warn("[FlashDreams WebRTC] getStats failed", error) } } @@ -696,6 +868,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" })) @@ -725,7 +898,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 @@ -738,9 +912,11 @@ async function connectSession() { setFlow("connected; waiting for input") logEvent("control data channel open") startHeartbeat() + modelAdapter?.onConnect?.(modelContext) } channel.onclose = () => { connected = false + setPostprocessDisabled(false) if (document.body.dataset.status !== "error") { setStatus("Closed", "idle") } @@ -748,6 +924,7 @@ async function connectSession() { logEvent("control data channel closed", { source: "client" }) stopHeartbeat() stopStatsPolling() + modelAdapter?.onDisconnect?.(modelContext) if (!disconnecting && pc.connectionState !== "closed") { pc.close() } @@ -785,6 +962,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") @@ -840,10 +1018,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 @@ -857,6 +1039,9 @@ function handleKeyDown(event) { } function handleKeyUp(event) { + if (isEditableControlTarget(event.target)) { + return + } const key = normalizeKey(event.key) if (!allowedKeys.has(key)) { return @@ -906,20 +1091,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", () => { @@ -943,4 +1124,4 @@ window.addEventListener("beforeunload", () => { disconnectSession() }) -initialize() +void initialize() 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_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_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_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..e0aa9cd0d --- /dev/null +++ b/flashdreams/tests/test_demo_runtime_timing.py @@ -0,0 +1,220 @@ +# 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 ( + CatchUpDecision, + CatchUpPolicy, + RealtimeEventInputSource, + RealtimeEventResampler, + ResamplerRealtimeClock, + SignalActivationPolicy, +) + +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 = 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) + + 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 = RealtimeEventResampler(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 = RealtimeEventResampler(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_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 = RealtimeEventInputSource(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 == 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 = RealtimeEventResampler(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 = RealtimeEventResampler(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_realtime_event_source_defers_unsupported_catch_up_policies( + policy: str, +) -> None: + 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 timeline analog"): + clock.catch_up( + request=_request(input_frame_count=1), + max_lag_s=1.0, + policy=unsupported_policy, + ) + + with pytest.raises(NotImplementedError, match="event-window analog"): + RealtimeEventInputSource( + 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, + ) + + +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 20d2183c6..06fa38e85 100644 --- a/flashdreams/tests/test_encoders.py +++ b/flashdreams/tests/test_encoders.py @@ -30,16 +30,19 @@ 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 +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,111 @@ class _FakeBufferedVideoTrack: aiortc runtime dependencies without breaking the isinstance check).""" def __init__(self) -> None: - self.enqueued_chunks: list = [] + self.enqueued_results: list[StepResult] = [] + self.enqueued_frames: list[object] = [] - async def enqueue_chunk(self, chunk) -> int: - self.enqueued_chunks.append(chunk) - return 4 + def prepare_result_frames(self, result: StepResult) -> tuple[object, ...]: + self.enqueued_results.append(result) + 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: + @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] + assert len(fake_track.enqueued_frames) == 4 + + @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_prepares_host_frames_with_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) + 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] + 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 +506,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 +586,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 +679,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 +697,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 +732,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 +781,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_inference_runtime_api.py b/flashdreams/tests/test_inference_runtime_api.py new file mode 100644 index 000000000..531f2f218 --- /dev/null +++ b/flashdreams/tests/test_inference_runtime_api.py @@ -0,0 +1,399 @@ +# 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 types import SimpleNamespace +from typing import Any, cast + +import pytest + +from flashdreams.runtime import ( + CanonicalInputs, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InMemoryMetricsRecorder, + InputField, + MetricsSnapshot, + NullMetricsRecorder, + NullOutputTarget, + OutputArtifact, + RuntimeMetricSample, + StepRequest, + StepRequirements, + StepResult, + TimeWindow, + UserInputEvent, + UserInputs, + UserInputSchema, + step_requirements_from_request, +) + +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, + 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): + 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=" ") + + +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"), + [ + (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: 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"), + (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_schema_validates_global_conditioning_and_step_payloads() -> None: + schema = InferenceInputSchema( + global_conditioning_fields=( + InputField(name="prompt"), + InputField(name="global_conditioning_frame"), + ), + step_fields=(InputField(name="camera_poses"),), + ) + inputs = InferenceInput( + global_conditioning={"prompt": "drive", "global_conditioning_frame": object()} + ) + + schema.require_global_conditioning(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_inference_input_unchanged() -> None: + mapping = IdentityInputMapping() + inference_input = InferenceInput( + global_conditioning={"prompt": "fixed"}, step={"hdmap": object()} + ) + request = StepRequest(step_index=0) + + assert ( + mapping.map_global_conditioning_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=inference_input, + ) + is inference_input + ) + assert ( + mapping.map_step_inputs( + canonical_inputs=CanonicalInputs(), + inference_input=inference_input, + request=request, + ) + is inference_input + ) + + +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") + + 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 + 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: + with pytest.raises(ValueError, match="unit='s'"): + RuntimeMetricSample( + name="model_step", + value=12.5, + unit="ms", + category="timing", + ) diff --git a/flashdreams/tests/test_output_targets.py b/flashdreams/tests/test_output_targets.py index 171e04786..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,9 +77,10 @@ def test_lingbot_webrtc_target_translates_runner_config() -> None: ), ) - assert spec.module == "lingbot.webrtc.server" + assert spec.module == "lingbot.demo.app" assert spec.argv == ( - "--config_name", + "webrtc", + "--preset-id", "lingbot-world-fast", "--device", "cuda:1", @@ -83,14 +90,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", ) @@ -108,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_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_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_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_canonical.py b/flashdreams/tests/test_runtime_canonical.py new file mode 100644 index 000000000..cfe2e3646 --- /dev/null +++ b/flashdreams/tests/test_runtime_canonical.py @@ -0,0 +1,591 @@ +# 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 reaches session start without touching canonicalization.""" + inputs = InferenceInput( + global_conditioning={"prompt": "heavy rain"}, + step={"steering": 0.0}, + ) + + assert inputs.global_conditioning["prompt"] == "heavy rain" + assert inputs.step["steering"] == 0.0 + + +# --- 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_demo_api.py b/flashdreams/tests/test_runtime_demo_api.py new file mode 100644 index 000000000..9600f19d0 --- /dev/null +++ b/flashdreams/tests/test_runtime_demo_api.py @@ -0,0 +1,650 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import argparse +from collections.abc import Sequence +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +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, + 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.serving.webrtc.demo import ( + serve_webrtc_demo, +) +from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager + +pytestmark = pytest.mark.ci_cpu + + +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]] = [] + + 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(), + ) + + result = run_replay_demo( + spec=spec, + adapter=adapter, + output_target_factory=lambda output_spec: output, + metrics=NullMetricsRecorder(), + runner=fake_runner, + ) + + 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 + 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_sink_from_spec(tmp_path: Path) -> None: + writer_calls: list[dict[str, Any]] = [] + sinks: list[OutputSink] = [] + + 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), + ) + + result = run_replay_demo( + spec=spec, + adapter=_FakeDemoAdapter(video_output=True), + output_sink_factory=lambda output_spec: _record_output_sink( + sinks, + build_output_sink( + output_spec, + mp4_writer=fake_writer, + ), + ), + ) + + 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), + "path": tmp_path / "demo.mp4", + "fps": 12, + "layout": "thwc", + } + ] + + +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 + + 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_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",), + output_modes=("null", "mp4"), + ) + + assert adapter.supported_input_modes() == ("replay",) + assert adapter.supported_output_modes() == ("null", "mp4") + + 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_serves_a_prepared_session_manager() -> None: + 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, + ), + ) + + 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: + 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, + ) + + +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( + 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, + 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() + 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 + + 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, + fail_step=self._fail_step, + ) + 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 + + +class _FakeRuntime: + def __init__( + self, + *, + 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 + + 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, + fail_step=self._fail_step, + ) + return self.session + + def close(self) -> None: + self.closed = True + + +class _FakeSession: + def __init__( + self, + *, + 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 + + 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) + 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, + 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: + 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 + + 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: + def __init__(self, config: Any) -> None: + self.config = config + + async def initialize(self) -> None: + return None + + async def reset_for_new_session(self, *, session_input: Any = None) -> None: + del session_input + + def peek_input_fps(self) -> float: + return 24.0 + + def peek_steady_output_num_frames(self) -> int: + return 1 + + 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 request, 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 new file mode 100644 index 000000000..b774371f4 --- /dev/null +++ b/flashdreams/tests/test_runtime_input_mapping.py @@ -0,0 +1,509 @@ +# 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 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 + +from typing import Any + +import pytest + +from flashdreams.runtime import ( + DRIVER_COMMAND, + 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", + input_modality="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_conditioning=(InputField(name="prompt", input_modality="text"),), +) +FRAME_MAPPING = InputMappingSchema( + name="conditioning-frame", + produces_global_conditioning=( + 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_conditioning_fields=(InputField(name="prompt", input_modality="text"),), + step_fields=( + InputField(name="steering"), + InputField(name="camera_delta", required=False), + ), +) + + +# --- user input events and windowing ------------------------------------ + + +def test_session_start_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_input_modality_mismatch_blocks_capability_match() -> None: + source = UserInputSchema( + capabilities=( + UserInputCapability(event_type="prompt_set", input_modality="embedding"), + ) + ) + + assert not source.supports( + UserInputCapability(event_type="prompt_set", input_modality="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_conditioning", "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_conditioning") 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_conditioning")["prompt"] == "drive" + assert inputs.for_phase("step")["steering"] == 0.25 + + +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", + frequency_consumed="once", + metadata={"coordinates": "opencv_c2w"}, + ) + + assert field.frequency_consumed == "once" + 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_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: + 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_global_conditioning_mapping_matches_global_conditioning_field() -> None: + model = InferenceInputSchema( + global_conditioning_fields=( + InputField(name="camera_trajectory", frequency_consumed="per_step"), + ) + ) + mapping = InputMappingSchema( + name="trajectory", + produces_global_conditioning=( + InputField(name="camera_trajectory", frequency_consumed="once"), + ), + ) + + compatibility = check_mapping_compatibility( + canonical_schema=CANONICAL_ALL, + inference_input_schema=model, + mapping_schema=mapping, + ) + + assert compatibility.can_drive + + +def test_unspecified_input_modality_stays_permissive() -> None: + model = InferenceInputSchema( + global_conditioning_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_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_conditioning=( + InputField(name="prompt", metadata={"source": "a"}), + ), + ) + second = InputMappingSchema( + name="b", + produces_global_conditioning=( + InputField(name="prompt", metadata={"source": "b", "extra": "kept"}), + ), + ) + + combined = combine_mapping_schemas((first, second)) + + assert len(combined.produces_global_conditioning) == 1 + metadata = combined.produces_global_conditioning[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 diff --git a/flashdreams/tests/test_runtime_runner.py b/flashdreams/tests/test_runtime_runner.py new file mode 100644 index 000000000..7ba92e836 --- /dev/null +++ b/flashdreams/tests/test_runtime_runner.py @@ -0,0 +1,752 @@ +# 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.runner as runner_module +from flashdreams.runtime import ( + DRIVER_COMMAND, + CanonicalInputs, + CanonicalInputSchema, + CanonicalModality, + DeviceConverterSchema, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InferenceRuntime, + InferenceSession, + InMemoryMetricsRecorder, + InputCanonicalizer, + InputField, + InputMapping, + InputMappingSchema, + MetricsSnapshot, + 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_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() + + 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 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_video_output.py b/flashdreams/tests/test_runtime_video_output.py new file mode 100644 index 000000000..71a4271dd --- /dev/null +++ b/flashdreams/tests/test_runtime_video_output.py @@ -0,0 +1,85 @@ +# 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.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="video StepResult"): + 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.from_video_chunk( + step_index=3, + 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), + ) + ) + + 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"] == ( + { + "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..d76558005 --- /dev/null +++ b/flashdreams/tests/test_runtime_worker.py @@ -0,0 +1,130 @@ +# 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 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") + 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"] + + +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_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 4b598cb63..4cf70b72b 100644 --- a/flashdreams/tests/test_video_output.py +++ b/flashdreams/tests/test_video_output.py @@ -5,72 +5,102 @@ from __future__ import annotations +from typing import Any, cast + import pytest import torch from flashdreams.infra.video_output import ( LazyRGBFrame, - RunnerVideoOutputStream, - VideoStepResult, + VideoOutputStream, + VideoResultCollector, infer_video_num_frames, lazy_rgb_frames_from_video_tensor, + prepare_video_for_mp4, video_tensor_to_hwc_uint8, ) -from flashdreams.serving.webrtc.manager import WebRTCStepResult, make_webrtc_step_result +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_webrtc_step_result_uses_shared_video_result_contract() -> None: - result = WebRTCStepResult( - chunk_index=1, - num_frames=2, +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)), - stats=None, + layout="tchw", + metadata=metadata, + metrics=metrics, ) - assert isinstance(result, VideoStepResult) - assert result.num_frames == 2 + 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_make_webrtc_step_result_preserves_tensor_and_infers_frames() -> None: +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", + ) - result = make_webrtc_step_result( - chunk_index=4, - video_chunk=video, - layout="tchw", - stats={"decode_ms": 1.5}, + result = output_stream.process( + video, + autoregressive_index=4, + metrics={"decode_ms": 1.5}, ) - assert isinstance(result, WebRTCStepResult) - 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: @@ -98,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", ) @@ -111,51 +141,134 @@ 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_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) - 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 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_runner_video_output_stream_collects_noop_chunks_without_postprocess() -> None: - output_stream = RunnerVideoOutputStream( +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) assert torch.equal(output[:, :, :2], first) assert torch.equal(output[:, :, 2:], second) + + +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=cast(Any, _TailPostprocess()), + output_layout="bcthw", + ) + result = output_stream.process( + torch.zeros((1, 3, 2, 4, 5)), + autoregressive_index=6, + ) + + 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 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 e213d0760..bafe2b1a8 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -6,20 +6,34 @@ import asyncio import json from types import SimpleNamespace -from typing import Any +from typing import Any, cast import pytest import torch +from flashdreams.runtime import ( + InferenceInput, + StepRequest, + StepRequirements, + StepResult, + UserInputEvent, + UserInputs, +) +from flashdreams.runtime.demo import RunResult +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, ManagedWebRTCSession, - WebRTCStepResult, ) 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 @@ -33,14 +47,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 +74,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 +83,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, @@ -76,6 +97,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 @@ -100,6 +144,30 @@ class _FakeResampler: dt = 0.0 next_chunk_start_v = 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 + + 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]]: @@ -107,39 +175,74 @@ def sample_chunk( 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 + + 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 + end = start + num_frames * self.dt + self.next_chunk_start_v = end + return [end] + + 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): + _resampler_supported_keys = frozenset({"w"}) 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_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: + 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 @@ -148,8 +251,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 @@ -162,7 +265,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 @@ -186,27 +292,436 @@ 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) + + +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={ + 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), + }, + ), + ) + + 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[ + manager_module._LEGACY_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, +) -> 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 = _RecordingLegacySegmentResampler() + managed.legacy_segment_resampler = resampler + 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_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: 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 - ) -> WebRTCStepResult: - 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 @@ -229,13 +744,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 - ) -> WebRTCStepResult: - 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: @@ -260,34 +775,74 @@ async def generate_chunk( 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: 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 - ) -> WebRTCStepResult: - 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 WebRTCStepResult( - 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 @@ -334,21 +889,24 @@ def __init__(self) -> None: def peek_input_fps(self) -> float: return 6.0 - def peek_next_input_num_frames(self) -> int: - return 2 - - async def generate_chunk( - self, *, segments: Any, frame_times: list[float] - ) -> WebRTCStepResult: - del segments + def next_step_request(self) -> StepRequest: + return _step_request(input_frame_count=2) + + 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 WebRTCStepResult( - 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() @@ -356,7 +914,7 @@ async def generate_chunk( 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 @@ -392,22 +950,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 - ) -> WebRTCStepResult: - 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 WebRTCStepResult( - 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, @@ -442,11 +1000,159 @@ 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" +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 - manager = _make_manager(_BusyManager, runtime=SimpleNamespace()) + 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, + 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(), + ) + 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( + _BaseTestManager, + runtime=SimpleNamespace(), + busy_message="custom busy message", + ) manager._runtime_ready = True manager._warmup_complete = True existing, *_ = _managed_session(runtime=SimpleNamespace()) @@ -456,42 +1162,27 @@ class _BusyManager(_BaseTestManager): await manager.create_answer(offer_sdp="x", offer_type="offer") -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 +def test_supported_key_payload_honors_configured_keys() -> None: + wsad_manager = _make_manager( + _BaseTestManager, + runtime=SimpleNamespace(), + supported_control_keys=WSAD_SUPPORTED_KEYS, ) - 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 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) 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 new file mode 100644 index 000000000..4c1a6138e --- /dev/null +++ b/flashdreams/tests/test_webrtc_services.py @@ -0,0 +1,781 @@ +# 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 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", + ] + 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_replaces_full_pending_queue() -> None: + loop = asyncio.get_running_loop() + encoder = _BlockingEncoder() + track = _FakeVideoTrack() + bridge = ThreadSafeWebRTCOutputBridge( + loop=loop, + video_encoder=encoder, + 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 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) + 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 + + def reset(self, *, start_v: float) -> None: + self.next_chunk_start_v = start_v + + def sample_chunk( + self, + num_frames: int, + ) -> tuple[float, ...]: + start = self.next_chunk_start_v + 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 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, *, queue_depth: int = 0) -> None: + self.queue_depth = queue_depth + self.flush_count = 0 + + def qsize(self) -> int: + 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 e1d1c9fc1..1fb07cbef 100644 --- a/flashdreams/tests/test_webrtc_serving.py +++ b/flashdreams/tests/test_webrtc_serving.py @@ -3,7 +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 @@ -11,11 +14,11 @@ 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, ) from flashdreams.serving.webrtc.media import tensor_chunk_to_rgb_frames from flashdreams.serving.webrtc.messages import ( @@ -41,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 @@ -118,6 +78,153 @@ 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, Any]): + 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=RealtimeEventResampler(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, identity="fake" + ) + 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, identity="fake" + ) + 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, identity="fake" + ) + 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" @@ -173,6 +280,37 @@ def _raise_creation_failure(**_kwargs) -> web.Application: assert closed +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-v5" in html + for slot in ( + "modelStageSlot", + "modelStatusSlot", + "modelPanelSlot", + "modelControlSlot", + ): + assert f'id="{slot}"' in html + 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 + assert "adapter.capabilities?.postprocess === true" in javascript + assert "renderControls(modelControls)" in javascript + assert "/api/session/initial_scene" not in javascript + + @pytest.mark.asyncio async def test_packaged_webrtc_app_serves_common_routes(tmp_path) -> None: (tmp_path / "request_session.html").write_text( @@ -200,6 +338,69 @@ 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-v4" + } + 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() + + +@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-v4", + } + 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/causal_forcing/causal_forcing/runner.py b/integrations/causal_forcing/causal_forcing/runner.py index 1a362bec9..574ec0561 100644 --- a/integrations/causal_forcing/causal_forcing/runner.py +++ b/integrations/causal_forcing/causal_forcing/runner.py @@ -28,19 +28,18 @@ 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, WanInferencePipelineCache, ) +from flashdreams.runtime.video_output import Mp4VideoOutputTarget __all__ = [ "CausalForcingI2VRunnerConfig", @@ -168,29 +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, + ) + ) - generated = output_stream.finish() - if generated is None: + tail = output_stream.finish() + if tail is not None: + output_target.write(tail) + artifacts = output_target.close() + if not artifacts: 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_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 8d4733601..d37bbad69 100644 --- a/integrations/cosmos_predict2/cosmos_predict2/runner.py +++ b/integrations/cosmos_predict2/cosmos_predict2/runner.py @@ -28,19 +28,18 @@ 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, CosmosInferencePipelineCache, ) +from flashdreams.runtime.video_output import Mp4VideoOutputTarget __all__ = [ "Cosmos2I2VRunner", @@ -138,27 +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) - generated = output_stream.finish() - if generated 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 - - 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_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 a1aa9f560..781ca58e5 100644 --- a/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py +++ b/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py @@ -26,16 +26,15 @@ 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, WanInferencePipelineCache, ) +from flashdreams.runtime.video_output import Mp4VideoOutputTarget __all__ = [ "FastvideoCausalWan22T2VRunnerConfig", @@ -128,30 +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, + ) + ) - generated = output_stream.finish() - if generated is None: + tail = output_stream.finish() + if tail is not None: + output_target.write(tail) + artifacts = output_target.close() + if not artifacts: 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_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/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/runner.py b/integrations/flashvsr/flashvsr/runner.py index e1707dd14..9f1f40cb0 100644 --- a/integrations/flashvsr/flashvsr/runner.py +++ b/integrations/flashvsr/flashvsr/runner.py @@ -30,15 +30,14 @@ 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 flashdreams.runtime.video_output import Mp4VideoOutputTarget from flashvsr.encoder import FlashVSREncoder from flashvsr.pipeline import ( FlashVSRPipeline, @@ -425,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( @@ -434,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. @@ -444,30 +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, + ) ) - generated = output_stream.finish() - if generated is None: + tail = output_stream.finish() + if tail is not None: + output_target.write(tail) + artifacts = output_target.close() + if not artifacts: 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_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/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/integrations/hy_worldplay/hy_worldplay/runner.py b/integrations/hy_worldplay/hy_worldplay/runner.py index a97faee4e..d29a9030d 100644 --- a/integrations/hy_worldplay/hy_worldplay/runner.py +++ b/integrations/hy_worldplay/hy_worldplay/runner.py @@ -31,13 +31,12 @@ 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 +from flashdreams.runtime.video_output import Mp4VideoOutputTarget __all__ = [ "DEFAULT_PROMPT", @@ -341,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() @@ -352,24 +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) - video = output_stream.finish() + output_target.write( + output_stream.process( + chunk, + autoregressive_index=ar_idx, + metrics=stats, + ) + ) elapsed = time.time() - start_time - if video is None: + tail = output_stream.finish() + if tail is not None: + output_target.write(tail) + artifacts = output_target.close() + if not artifacts: 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") + 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/README.md b/integrations/lingbot/README.md index b781e4a21..71c71f91c 100644 --- a/integrations/lingbot/README.md +++ b/integrations/lingbot/README.md @@ -98,6 +98,65 @@ 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 +``` + +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 +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 +202,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/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/__init__.py b/integrations/lingbot/lingbot/demo/__init__.py new file mode 100644 index 000000000..08e4f9d1c --- /dev/null +++ b/integrations/lingbot/lingbot/demo/__init__.py @@ -0,0 +1,22 @@ +# 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.providers import LingbotInputProvider +from lingbot.demo.spec import ( + DEFAULT_LINGBOT_PRESET, + LINGBOT_MODEL_ID, + LingbotReplayInputs, + LingbotWebRTCScenario, +) + +__all__ = [ + "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 new file mode 100644 index 000000000..fedd0a2c2 --- /dev/null +++ b/integrations/lingbot/lingbot/demo/adapter.py @@ -0,0 +1,259 @@ +# 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 ( + InputCanonicalizer, + UserInputCapability, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + NullOutputSpec, + PreparedScenario, + WebRTCOutputSpec, +) +from flashdreams.runtime.interfaces import InferenceRuntime +from lingbot.input_mapping import ( + KeyboardToCameraCommand, + 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] + + +class LingbotDemoAdapter(LingbotModelAdapter): + """Model-owned Lingbot adapter consumed by shared demo launchers.""" + + def __init__( + self, + *, + replay_runtime_factory: ReplayRuntimeFactory = LingbotReplayRuntime, + pipeline_factory: PipelineFactory | None = None, + ) -> None: + super().__init__( + runtime_factory=replay_runtime_factory, + pipeline_factory=pipeline_factory, + ) + + def supported_input_modes(self) -> tuple[str, ...]: + return ("replay", "keyboard-driving") + + def supported_output_modes(self) -> tuple[str, ...]: + return ("mp4", "null", "webrtc") + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + if spec.input_mode == "replay": + 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": + 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 supports input_mode='replay' or " + f"'keyboard-driving', got {spec.input_mode!r}." + ) + + replay_inputs = resolve_replay_inputs( + scenario, + default_prompt=_default_prompt(self, spec), + ) + 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. + 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 live 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, + include_keyboard=live_camera, + include_text_events=live_camera and bool(text_event_prompts), + ), + canonicalizer=_canonicalizer(text_event_prompts), + mapping=mapping, + metadata={ + "model_id": self.model_id, + "preset_id": self.preset_id(spec.config), + }, + ) + + 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): + return str(scenario.get("camera_source", "trace")) + 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 = ( + 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, + *, + 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 + 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 include_keyboard or observed & _KEY_EVENT_TYPES: + capabilities.extend(_KEYBOARD_CAPABILITIES) + 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() + 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", +] diff --git a/integrations/lingbot/lingbot/demo/app.py b/integrations/lingbot/lingbot/demo/app.py new file mode 100644 index 000000000..8aa495dff --- /dev/null +++ b/integrations/lingbot/lingbot/demo/app.py @@ -0,0 +1,250 @@ +# 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 +from typing import Any + +from flashdreams.runtime import InferenceConfig +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + NullOutputSpec, + WebRTCOutputSpec, +) +from flashdreams.runtime.demo.app import DemoApplication +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 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) + 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-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) + 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, + ) + 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): + """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, + ) + + 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, + ), + world_rank=context.world_rank, + ) + + +_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: + 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=_replay_output_spec(args), + config=InferenceConfig( + model_id=LINGBOT_MODEL_ID, + preset_id=args.preset_id, + device=args.device, + ), + ) + + +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, + *, + 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/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/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..8fd15494e --- /dev/null +++ b/integrations/lingbot/lingbot/demo/spec.py @@ -0,0 +1,165 @@ +# 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..99ba80874 --- /dev/null +++ b/integrations/lingbot/lingbot/demo/webrtc.py @@ -0,0 +1,186 @@ +# 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 collections.abc import Callable +from dataclasses import replace +from importlib.resources import files +from typing import Any + +from flashdreams.runtime import InferenceConfig +from flashdreams.runtime.demo import DemoSpec, 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 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, +) +from lingbot.webrtc.server import configure_lingbot_webrtc_app +from lingbot.webrtc.session import ( + LingbotInferenceRuntime, + LingbotRuntimeConfig, +) + +from .spec import resolve_webrtc_scenario + +WebRTCRuntimeFactory = Callable[..., Any] + + +def serve_lingbot_webrtc_demo( + *, + spec: DemoSpec, + 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) + 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, + 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) + + +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/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..2963e2895 --- /dev/null +++ b/integrations/lingbot/lingbot/input_mapping.py @@ -0,0 +1,675 @@ +# 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.keyboard import DEFAULT_SUPPORTED_KEYS, KeyboardState +from flashdreams.runtime.mapping import InputMappingSchema +from flashdreams.runtime.types import StepRequest +from lingbot.controls import ( + CameraPoseIntegrator, + PoseSegment, +) + +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 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], + *, + 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/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 de269ae78..9a4ffe0f0 100644 --- a/integrations/lingbot/lingbot/runner.py +++ b/integrations/lingbot/lingbot/runner.py @@ -20,33 +20,35 @@ 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.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, + ensure_example_data_downloaded, + example_data_dirname, ) from lingbot.pipeline import ( LingbotWorldInferencePipeline, ) +from lingbot.runtime import ( + LINGBOT_MODEL_ID, + LingbotRunnerOutputTarget, + inference_config_from_runner_config, + replay_inputs_from_runner_config, +) __all__ = [ + "EXAMPLE_DATA_AVAILABLE_IDXS", "LingbotWorldRunnerConfig", "LingbotWorldRunner", + "example_data_dirname", ] @@ -58,67 +60,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): @@ -127,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 @@ -222,118 +164,49 @@ 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 = LingbotDemoAdapter() + 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.", + replay_inputs = replay_inputs_from_runner_config( + cfg, + is_rank_zero=self.is_rank_zero, ) - - 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, + 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, ) - 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)}" + 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, ) - 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, - fps=cfg.fps, - layout="tchw", - install_hint="Install the lingbot plugin: pip install flashdreams-lingbot.", + result = run_replay_demo( + spec=spec, + adapter=adapter, + output_target_factory=_output_target_factory, ) - logger.info( - f"[{cfg.runner_name}] wrote video {tuple(video.shape)} " - f"-> {video_path.resolve()}" - ) - - 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()}" + 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/lingbot/runtime.py b/integrations/lingbot/lingbot/runtime.py new file mode 100644 index 000000000..ef7b549ef --- /dev/null +++ b/integrations/lingbot/lingbot/runtime.py @@ -0,0 +1,1082 @@ +# 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 +from collections.abc import Callable, Mapping +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any, cast + +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.config import derive_config +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.infra.runner_io import ( + load_first_frame_tensor, + runner_artifact_path, + write_runner_stats, +) +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 +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, +) +from lingbot.model_session import LingbotModelSessionCore + +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) + 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=cast(VideoTensorLayout, output_layout), + ), + ) + + 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._frame_start = 0 + self._active_prompt = session_inputs.prompt + 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(): + 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.inputs.total_blocks: + return None + 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=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._model_session.step_index + num_frames = self._model_session.next_num_frames() + 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, + ) + result = self._model_session.step( + camctrl_input, + output_window=TimeWindow( + start_s=frame_start / self.inputs.fps, + end_s=frame_end / self.inputs.fps, + ), + ) + 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.") + self._active_prompt = self.inputs.prompt + self._reset_model_session() + 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) + 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._model_session.step_index, + ) + + def close(self) -> None: + self._closed = True + self._model_session.close() + + def _reset_model_session(self) -> None: + 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, + ) + self._model_session.reset( + prompt=self.inputs.prompt, + first_frames=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: VideoOutputStream + output_dir: Path + runner_name: str + 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.") + if result.layout is None: + raise TypeError( + "LingbotRunnerOutputTarget requires a video StepResult with layout." + ) + 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 + 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(), + ) + stats_history = artifacts[0].metadata.get("stats_history", ()) + if stats_history: + stats_path = write_runner_stats( + self.output_dir, + self.runner_name, + list(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: + pipeline_config = derive_config( + base_config=pipeline_config, + diffusion_model=dict(transformer=dict(init_device=device)), + ) + return pipeline_config.setup().to(device=device).eval() + + +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..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,35 +50,46 @@ from flashdreams.serving.webrtc.server import ( close_package_resources as _close_package_resources, ) -from lingbot.runner import ( +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, + LingbotInferenceRuntime, LingbotRuntimeConfig, LingbotSessionInput, - LingbotWebRTCSessionManager, + LingbotWebRTCSessionController, + create_lingbot_webrtc_session_manager, normalize_prompt_text, 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 -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: @@ -105,6 +118,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, @@ -162,34 +181,68 @@ def create_app( *, request_session_url: str, session_manager: WebRTCSessionManager | None = None, + session_controller: LingbotSessionController | None = None, ) -> 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) + 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, + 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_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, + *, + 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) @@ -308,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, @@ -317,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( @@ -335,34 +388,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, ) @@ -411,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 c0b59680b..03c1d2b90 100644 --- a/integrations/lingbot/lingbot/webrtc/session.py +++ b/integrations/lingbot/lingbot/webrtc/session.py @@ -17,17 +17,17 @@ from __future__ import annotations -import asyncio import http.client import io import ipaddress import re import socket import ssl +import threading 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 @@ -35,31 +35,37 @@ 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.serving.webrtc.controls import ( - CameraPoseIntegrator, - PoseSegment, -) -from flashdreams.serving.webrtc.encoders import ( - EncoderBackend, - VideoEncoder, - select_encoder, +from flashdreams.infra.video_output import VideoOutputStream +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.encoders import EncoderBackend from flashdreams.serving.webrtc.manager import ( DEFAULT_CLIENT_LIVENESS_TIMEOUT_S, BaseWebRTCSessionManager, - ManagedWebRTCSession, WebRTCControlSignal, - WebRTCStepResult, - make_webrtc_step_result, +) +from flashdreams.serving.webrtc.runtime import ( + 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, + FIELD_CAMERA_TRAJECTORY, + KeyboardToCameraCommand, + LingbotInputMapping, + TextEventSelection, +) +from lingbot.model_session import LingbotModelSessionCore _INTRINSICS_REFERENCE_HEIGHT = 480 _INTRINSICS_REFERENCE_WIDTH = 832 @@ -476,6 +482,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) @@ -569,84 +577,36 @@ 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._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 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 - - 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" @@ -654,65 +614,103 @@ 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 generate_chunk( - self, - *, - segments: list[PoseSegment], - frame_times: list[float], - ) -> WebRTCStepResult: - """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. - - Returns: - :class:`WebRTCStepResult` carrying the produced video chunk - and the post-generation pipeline stats. - - Raises: - LingbotRuntimeError: Runtime is closed or not initialized. + 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("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 + 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) - def peek_next_chunk_num_frames(self) -> int: - """Return the number of frames the next chunk's pipeline call will emit. + @property + def input_mapping(self) -> LingbotInputMapping: + if self._input_mapping is None: + raise LingbotRuntimeError("Runtime input mapping is not initialized.") + return self._input_mapping - 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: + @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.") - return int(self._pipeline.get_num_output_frames(self.autoregressive_index)) + 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 @@ -722,7 +720,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 @@ -741,23 +752,9 @@ 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], - ) -> WebRTCStepResult: - 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: + return self._step_sync(inputs) @distributed_op(WebRTCControlSignal.EVENT) def _trigger_event_sync_all_ranks( @@ -767,23 +764,21 @@ 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 - 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(): raise RuntimeError("CUDA is required for Lingbot runtime.") @@ -796,47 +791,27 @@ 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, - 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) + 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.") @@ -849,6 +824,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)): @@ -859,10 +835,15 @@ 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: - 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) @@ -977,8 +958,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. @@ -1034,13 +1013,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 @@ -1052,23 +1027,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" @@ -1102,9 +1073,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 @@ -1112,141 +1083,232 @@ 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() def _generate_one_chunk_sync( self, *, - segments: list[PoseSegment], + segments: list[Any], frame_times: list[float], - ) -> WebRTCStepResult: + ) -> 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.") + 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) 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, + ) -> 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._model_session is None: + raise LingbotRuntimeError("Runtime 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( - autoregressive_index=self.autoregressive_index, - cache=self._cache, - 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", - 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: + """Generate one chunk from mapped model inputs.""" + if self._pipeline is None or self._model_session is None: + raise LingbotRuntimeError("Runtime is not initialized.") + 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) + ) + intrinsics = _require_camera_tensor( + inputs, FIELD_CAMERA_INTRINSICS, expected_shape=(num_frames, 4) + ) + return self._generate_from_camera_inputs( + poses=poses, + intrinsics=intrinsics, + num_frames=num_frames, + ) -_ManagedLingbotSession = ManagedWebRTCSession + 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, + ) -class LingbotWebRTCSessionManager( - BaseWebRTCSessionManager[LingbotInferenceRuntime, LingbotRuntimeConfig] -): - """Owns one active WebRTC session and forwards actions into Lingbot runtime.""" +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 - _busy_message = "A Lingbot session is already active." - _warmup_label = "Lingbot WebRTC" - _runtime_error_types = (LingbotRuntimeError,) - def __init__( - self, - *, - 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() - fps = runtime_config.fps if fps is None else fps - if fps <= 0: - raise ValueError("fps must be > 0") - super().__init__( - runtime=LingbotInferenceRuntime(config=runtime_config), - runtime_config=runtime_config, - fps=fps, - client_liveness_timeout_s=client_liveness_timeout_s, +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." ) - self._pending_session_input: LingbotSessionInput | None = None - def _model_name(self) -> str: - return self.runtime_config.config_name + def close(self) -> None: + # The runtime outlives the session and is closed by the serve loop. + return None - 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 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, + ) - def _clear_pending_session_input(self) -> None: - self._pending_session_input = None - async def _reset_runtime_for_session( - self, session_input: LingbotSessionInput | None +class LingbotWebRTCSessionController: + """Own Lingbot browser inputs and preview data outside the transport manager.""" + + def __init__( + self, + manager: BaseWebRTCSessionManager[ + LingbotInferenceRuntime, + LingbotRuntimeConfig, + ], ) -> 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) @@ -1256,11 +1318,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( @@ -1271,27 +1333,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, @@ -1311,8 +1373,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( @@ -1327,11 +1389,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 @@ -1379,15 +1441,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.css b/integrations/lingbot/lingbot/webrtc/web/adapter.css new file mode 100644 index 000000000..30ed19e51 --- /dev/null +++ b/integrations/lingbot/lingbot/webrtc/web/adapter.css @@ -0,0 +1,325 @@ +/* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. */ +/* SPDX-License-Identifier: Apache-2.0 */ + +.firstFramePreview { + position: absolute; + inset: 0; + 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; +} + +.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..91389585e --- /dev/null +++ b/integrations/lingbot/lingbot/webrtc/web/adapter.js @@ -0,0 +1,554 @@ +// 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: "Drive / Turn", + keys: [ + { key: "w", label: "Forward" }, + { key: "a", label: "Turn left" }, + { key: "s", label: "Backward" }, + { key: "d", label: "Turn right" }, + ], + }, + { + 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") + 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.", + 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?v=lingbot-video-size-v2", 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/mock_ui_server.py b/integrations/lingbot/lingbot/webrtc/web/mock_ui_server.py deleted file mode 100644 index fecdadd6c..000000000 --- a/integrations/lingbot/lingbot/webrtc/web/mock_ui_server.py +++ /dev/null @@ -1,78 +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. - -from __future__ import annotations - -import argparse -from functools import partial -from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer -from importlib.resources import as_file, files -from urllib.parse import urlsplit - -WEB_DIR_RESOURCE = files("lingbot.webrtc").joinpath("web") - - -class MockUIRequestHandler(SimpleHTTPRequestHandler): - """Serve the static viewer without preloading the Lingbot runtime.""" - - def _rewrite_path(self) -> bool: - path = urlsplit(self.path).path - if path == "/": - self.send_response(302) - self.send_header("Location", "/request_session?mock=1") - self.end_headers() - return True - if path == "/request_session": - self.path = "/request_session.html" - elif path.startswith("/static/"): - self.path = "/" + path.removeprefix("/static/") - return False - - def do_GET(self) -> None: - if self._rewrite_path(): - return - super().do_GET() - - def do_HEAD(self) -> None: - if self._rewrite_path(): - return - super().do_HEAD() - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Serve the Lingbot 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() - - -def main() -> None: - args = parse_args() - with as_file(WEB_DIR_RESOURCE) as web_dir: - 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" - ) - try: - server.serve_forever() - except KeyboardInterrupt: - print("\nStopping mock UI server.") - finally: - server.server_close() - - -if __name__ == "__main__": - main() 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 f9e39cd92..f0175fb35 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.app: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 @@ -57,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_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 new file mode 100644 index 000000000..98444a7c5 --- /dev/null +++ b/integrations/lingbot/tests/test_demo_api.py @@ -0,0 +1,1027 @@ +# 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 numpy as np +import pytest +import torch +from aiohttp import web +from lingbot.demo import ( + DEFAULT_LINGBOT_PRESET, + LINGBOT_MODEL_ID, + LingbotDemoAdapter, + LingbotInputProvider, + LingbotReplayInputs, + LingbotWebRTCScenario, +) +from lingbot.demo.app import _replay_spec, _webrtc_spec, parse_args +from lingbot.demo.replay import ( + LingbotReplayRuntime, + LingbotReplayRuntimeOptions, +) +from lingbot.demo.webrtc import serve_lingbot_webrtc_demo +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, + LingbotSessionInput, + TextEventSpec, +) + +from flashdreams.runtime import ( + CanonicalInputs, + InferenceConfig, + InferenceInput, + OutputArtifact, + OutputTarget, + StepRequest, + StepRequirements, + StepResult, + UserInputEvent, + UserInputs, +) +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + NullOutputSpec, + UserInputWindow, + WebRTCOutputSpec, +) +from flashdreams.runtime.demo.replay import run_replay_demo +from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager +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 _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_shared_demo_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", "null", "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_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" + 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}, + ), + ) + + result = run_replay_demo( + spec=spec, + adapter=adapter, + output_target_factory=lambda output_spec: output, + runner=fake_runner, + ) + + 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 + 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_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_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: + 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, 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 + 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_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, + 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}, + ), + ) + + 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 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 + 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.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.serving.webrtc.demo 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: + app_calls.append(kwargs) + app = web.Application() + app[SESSION_MANAGER_KEY] = kwargs["session_manager"] + kwargs["configure_app"](app) + return app + + monkeypatch.setattr( + shared_webrtc_module, "create_packaged_webrtc_app", fake_create_packaged_app + ) + 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()}, + ), + ) + + 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 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"] == "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 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, + tmp_path: Path, +) -> None: + import flashdreams.serving.webrtc.demo 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: + 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( + shared_webrtc_module, "create_packaged_webrtc_app", fake_create_packaged_app + ) + 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()}, + ), + ) + + 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]["app"] is app + assert server_calls[0]["host"] == "0.0.0.0" + assert server_calls[0]["port"] == 8080 + assert type(server_calls[0]["session_manager"]) is BaseWebRTCSessionManager + + +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_input_fps(self) -> float: + return 16.0 + + def peek_steady_output_num_frames(self) -> int: + return 1 + + 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 request, 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_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_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 new file mode 100644 index 000000000..87ffb0b63 --- /dev/null +++ b/integrations/lingbot/tests/test_input_mapping.py @@ -0,0 +1,452 @@ +# 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 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 + # 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..fcbf924f4 --- /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.controls import CameraPoseIntegrator, KeyboardResampler +from lingbot.input_mapping import ( + FIELD_CAMERA_TRAJECTORY, + KeyboardToCameraCommand, + LingbotInputMapping, +) + +from flashdreams.runtime import ( + InferenceInput, + InputCanonicalizer, + StepRequest, + TimeWindow, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, +) + +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..b5124c091 --- /dev/null +++ b/integrations/lingbot/tests/test_runtime_gpu.py @@ -0,0 +1,262 @@ +# 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.runtime import ( + CanonicalInputs, + InferenceConfig, + InferenceInput, + StepResult, +) + +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, 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] + + +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 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 new file mode 100644 index 000000000..27a9c4b93 --- /dev/null +++ b/integrations/lingbot/tests/test_runtime_session_inputs.py @@ -0,0 +1,397 @@ +# 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 lingbot.runtime as runtime_module +import pytest +import torch +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"]] + 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( + 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_server_routes.py b/integrations/lingbot/tests/test_server_routes.py index fbb024bb9..167f147a1 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"' 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/lingbot/tests/test_smoke.py b/integrations/lingbot/tests/test_smoke.py index bf9c1ab1a..f51d0c2bd 100644 --- a/integrations/lingbot/tests/test_smoke.py +++ b/integrations/lingbot/tests/test_smoke.py @@ -18,12 +18,15 @@ from __future__ import annotations import sys +from collections.abc import Callable 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 +42,11 @@ LingbotWorldRunnerConfig, example_data_dirname, ) +from lingbot.runtime import ( + LINGBOT_MODEL_ID, + LingbotReplayInputs, + LingbotRunnerOutputTarget, +) from lingbot.transformer import ( LINGBOT_WORLD_MIN_CHECKPOINT_FREE_GB, LingbotWorldTransformer, @@ -47,9 +55,23 @@ from flashdreams.infra.config import derive_config from flashdreams.infra.runner import RunnerConfig +from flashdreams.runtime import InferenceConfig +from flashdreams.runtime.demo import DemoSpec, Mp4OutputSpec, RunResult 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 +99,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 +127,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 +184,86 @@ 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 shared replay runtime path.""" + 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", + output_dir=tmp_path, + ), + ) + pipeline = object() + output_stream = object() + captured: dict[str, object] = {} + + def _fake_run_replay_demo(**kwargs: object) -> RunResult: + captured.update(kwargs) + return RunResult(status="completed") + + monkeypatch.setattr( + runner, + "create_video_output_stream", + lambda **_kwargs: output_stream, + ) + monkeypatch.setattr( + runner_mod, + "run_replay_demo", + _fake_run_replay_demo, + ) + runner.config = runner_config + runner.pipeline = pipeline + runner.local_rank = 0 + runner.world_size = 1 + runner.is_rank_zero = True + + runner.run() + + 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 + 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 + + def test_runners_dict_is_non_empty() -> None: """Plugin must expose at least one runner.""" assert RUNNER_CONFIGS, "RUNNER_CONFIGS is empty" @@ -195,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_runtime.py b/integrations/lingbot/tests/test_webrtc_runtime.py index 97a6212a9..2b70dc366 100644 --- a/integrations/lingbot/tests/test_webrtc_runtime.py +++ b/integrations/lingbot/tests/test_webrtc_runtime.py @@ -17,24 +17,54 @@ 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.model_session import LingbotModelSessionCore from lingbot.webrtc import session from lingbot.webrtc.session import ( LingbotRuntimeConfig, - LingbotWebRTCSessionManager, + create_lingbot_webrtc_session_manager, ) -from flashdreams.infra.postprocess import VideoTensorLayout -from flashdreams.serving.webrtc.manager import WebRTCStepResult +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, + 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 @@ -43,18 +73,8 @@ 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`` + """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.""" @@ -72,19 +92,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." - ) - assert LingbotWebRTCSessionManager._warmup_label == "Lingbot WebRTC" - assert LingbotWebRTCSessionManager._runtime_error_types == ( - session.LingbotRuntimeError, + manager = create_lingbot_webrtc_session_manager( + runtime_config=LingbotRuntimeConfig(device="cpu") ) - # 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: @@ -106,7 +120,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) ) @@ -119,7 +133,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) ) @@ -141,7 +157,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", @@ -173,7 +189,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 +221,16 @@ 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_process( + _stream: VideoOutputStream, video_chunk: object, **kwargs: object + ) -> StepResult: + captured["video_chunk"] = video_chunk captured.update(kwargs) - stats = cast(dict[str, float] | None, kwargs["stats"]) - layout = cast(VideoTensorLayout | None, kwargs["layout"]) - return WebRTCStepResult( - chunk_index=0, - num_frames=2, + return StepResult.from_video_chunk( + step_index=0, video_chunk=torch.zeros((2, 3, 4, 5)), - stats=stats, - layout=layout, + layout="tchw", + metrics={"total_ms": 3.0}, ) runtime = session.LingbotInferenceRuntime( @@ -223,12 +239,12 @@ def _fake_make_webrtc_step_result(**kwargs: object) -> WebRTCStepResult: 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( - session, - "make_webrtc_step_result", - _fake_make_webrtc_step_result, + VideoOutputStream, + "process", + _fake_process, ) result = runtime._generate_one_chunk_sync( @@ -237,11 +253,13 @@ 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") + 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( @@ -394,11 +412,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 @@ -441,9 +460,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", @@ -453,10 +473,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()] @@ -501,16 +521,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"] @@ -519,8 +540,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( @@ -571,149 +592,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() + _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, + } - 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 {} - - 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, - ) + 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)] - await manager._handle_datachannel_message( - managed_session=managed_session, - raw_message='{"type":"event","state":"trigger"}', + runtime._apply_conditioning_update_sync( + InferenceInput(global_conditioning={"prompt": "base prompt"}) ) - 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: @@ -746,7 +678,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} @@ -764,6 +696,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: @@ -792,6 +754,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, @@ -830,7 +793,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 @@ -858,7 +821,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]] = [] @@ -908,18 +871,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) ) @@ -936,15 +899,51 @@ 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={ + "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.from_video_chunk( + step_index=chunk_index, + video_chunk=torch.zeros((1, 1, 1, 3, 2, 2), dtype=torch.uint8), + layout="bvtchw", + ) + 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 @@ -955,28 +954,21 @@ 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 generate_chunk( - self, - *, - segments: list[tuple[float, float, frozenset[str]]], - frame_times: list[float], - ) -> WebRTCStepResult: - del frame_times - chunk_index = len(self.generated_segments) - self.generated_segments.append(segments) - return WebRTCStepResult( - chunk_index=chunk_index, - num_frames=1, - video_chunk=torch.zeros((1, 1, 1, 3, 2, 2), dtype=torch.uint8), - stats=None, + def next_step_request(self) -> StepRequest: + return StepRequest( + step_index=self.step_index, + metadata={"input_frame_count": 1}, ) + async def start_inference_session(self) -> _FakeInferenceSession: + return _FakeInferenceSession(self) + async def close(self) -> None: self.close_calls += 1 @@ -988,7 +980,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, @@ -1001,7 +993,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() @@ -1036,7 +1028,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) ) @@ -1053,7 +1045,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 @@ -1088,10 +1080,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] @@ -1116,13 +1108,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] @@ -1149,12 +1141,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/lingbot/tests/test_webrtc_session_branch.py b/integrations/lingbot/tests/test_webrtc_session_branch.py new file mode 100644 index 000000000..1a37d60f6 --- /dev/null +++ b/integrations/lingbot/tests/test_webrtc_session_branch.py @@ -0,0 +1,506 @@ +# 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.controls import CameraPoseIntegrator, KeyboardResampler +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.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.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={ + "input_frame_count": _NUM_FRAMES, + "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.from_video_chunk( + step_index=index, + video_chunk=torch.zeros(_NUM_FRAMES, 3, 4, 4), + layout="tchw", + ) + + +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) + + +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=_FakeCloseable(), # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), # ty:ignore[invalid-argument-type] + peer_connection=_FakeCloseable(), + resampler=RealtimeEventResampler(fps=_FPS, start_v=0.0), + inference_session=runtime.session, + ) + + +def _manager(runtime: _FakeRuntime) -> _Manager: + return _Manager( + runtime=runtime, + runtime_config=_FakeRuntimeConfig(), + fps=_FPS, + identity="fake", + ) + + +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.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: + """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()) diff --git a/integrations/omnidreams/omnidreams/demo/README.md b/integrations/omnidreams/omnidreams/demo/README.md new file mode 100644 index 000000000..9dab46ed8 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/README.md @@ -0,0 +1,97 @@ + + +# OmniDreams Shared Demo API + +This folder contains the experimental OmniDreams demo built on +`flashdreams.runtime.demo`. + +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 +``` + +## Null Replay + +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 --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 +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. + +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 --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 +``` + +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 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 --python 3.12 --package flashdreams-omnidreams omnidreams-demo webrtc \ + --host 0.0.0.0 \ + --port 8089 \ + --device cuda:0 \ + --scene-uuid 0d404ff7-2b66-498c-b047-1ed8cded60d4 +``` + +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..a3d646989 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/__init__.py @@ -0,0 +1,36 @@ +# 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.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 new file mode 100644 index 000000000..4dc7d1c87 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/adapter.py @@ -0,0 +1,322 @@ +# 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 + +import math +from collections.abc import Callable +from typing import Any, cast + +from omnidreams.config import OMNIDREAMS_CONFIGS, OMNIDREAMS_RUNNERS + +from flashdreams.runtime import ( + CanonicalInputSchema, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InputCanonicalizer, + UserInputSchema, +) +from flashdreams.runtime.demo import ( + DemoSpec, + 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 .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, +) + +RuntimeFactory = Callable[..., InferenceRuntime] +ReplayRuntimeFactory = RuntimeFactory + + +class OmnidreamsDemoAdapter: + """Model-owned OmniDreams adapter consumed by shared demo launchers.""" + + def __init__( + self, + *, + runtime_factory: RuntimeFactory | None = None, + replay_runtime_factory: ReplayRuntimeFactory | None = None, + pipeline_factory: PipelineFactory | None = None, + ) -> None: + 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() + + @property + def model_id(self) -> str: + return OMNIDREAMS_MODEL_ID + + @property + def inference_input_schema(self) -> InferenceInputSchema: + return precomputed_hdmap_inference_input_schema() + + @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", "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 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}." + ) + return PreparedScenario( + initial_inputs=InferenceInput( + global_conditioning={"scenario": scenario}, + ), + 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), + }, + ) + + 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._runtime_factory( + config=config, + 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")), + ), + ) + + 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 _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 new file mode 100644 index 000000000..ac1641917 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/app.py @@ -0,0 +1,253 @@ +# 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 +import math +from pathlib import Path +from typing import Any + +from omnidreams.runner import DEFAULT_EXAMPLE_DATA_UUID_1V + +from flashdreams.runtime import InferenceConfig +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + NullOutputSpec, + WebRTCOutputSpec, +) +from flashdreams.runtime.demo.app import DemoApplication + +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, +) + + +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 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, + 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-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) + 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("--prefer-sw-encoder", action="store_true") + 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): + """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)), + world_rank=context.world_rank, + ) + + +_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: + 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, + "pixel_height": args.pixel_height, + "pixel_width": args.pixel_width, + "fps": args.fps, + } + if args.prompt: + scenario["prompt"] = args.prompt + 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=_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, + 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, + 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/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 new file mode 100644 index 000000000..27c3784e5 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/providers.py @@ -0,0 +1,734 @@ +# 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.inputs import ( + InferenceInput, + InferenceInputSchema, + InputField, + UserInputCapability, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.types import StepRequirements +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, + 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]]: + frame_times = list(user_window.frame_times) + if frame_times and len(frame_times) != request.input_frame_count: + raise RuntimeError( + "OmniDreams Ludus realtime window frame_times length does not " + "match the requested input frame count." + ) + 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))) + 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 _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 new file mode 100644 index 000000000..b6c6d91e2 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/replay.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compatibility aliases for the OmniDreams runtime module.""" + +from __future__ import annotations + +from .runtime import ( + OmnidreamsRuntime, + OmnidreamsRuntimeOptions, + OmnidreamsSession, + OmnidreamsSessionScenario, + PipelineFactory, +) + +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 new file mode 100644 index 000000000..8df653ecd --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/spec.py @@ -0,0 +1,550 @@ +# 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 + +import json +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal, TypeAlias, cast + +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 + +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) +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 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.""" + + 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" + debug_serve_hdmaps: bool = False + 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 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, + *, + 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_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: + 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_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)), + 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 _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 + 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 _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}.") + 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", + "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 new file mode 100644 index 000000000..37d19a299 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/web/adapter.js @@ -0,0 +1,18 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +export default { + modelName: "OmniDreams", + stylesheet: "/model-static/adapter.css?v=model-ui-v2", + 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 new file mode 100644 index 000000000..fb3d12628 --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/webrtc.py @@ -0,0 +1,309 @@ +# 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. + +"""OmniDreams browser hooks for the shared WebRTC demo runtime.""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from dataclasses import replace +from typing import Any + +from loguru import logger +from omnidreams.config import OMNIDREAMS_CONFIGS + +from flashdreams.runtime import InferenceConfig +from flashdreams.runtime.demo import ( + DemoSpec, + 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 .adapter import OmnidreamsDemoAdapter, RuntimeFactory +from .controls import WSAD_SUPPORTED_KEYS +from .spec import ( + DEFAULT_OMNIDREAMS_PRESET, + OMNIDREAMS_MODEL_ID, + resolve_webrtc_scenario, +) +from .webrtc_config import OmnidreamsWebRTCModelRuntimeConfig + +WebRTCRuntimeFactory = Callable[..., Any] +SharedRuntimeFactory = RuntimeFactory + + +def serve_omnidreams_webrtc_demo( + *, + spec: DemoSpec, + world_rank: int = 0, + 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 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', " + 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) + 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( + 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=output.video_height, + video_width=output.video_width, + fps=output.fps, + camera_name=scenario.camera_name, + 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", + ) + 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, + 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, + shared_host=host, + shared_adapter=adapter, + shared_spec=spec, + shared_scenario=prepared, + ) + 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, + ) + + +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 + if config is None or config.preset_id is None + else config.preset_id + ) + + +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 _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, +) -> 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__ = [ + "OmnidreamsWebRTCModelRuntimeConfig", + "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..401d9158e --- /dev/null +++ b/integrations/omnidreams/omnidreams/demo/webrtc_legacy.py @@ -0,0 +1,724 @@ +# 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, cast + +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.serving.webrtc.demo import ( + CreateWebRTCApp, + RunWebRTCServer, + serve_webrtc_demo, +) +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 .controls import ( + SPARSE_KEY_SEGMENTS_METADATA_KEY, + WSAD_SUPPORTED_KEYS, + KeyboardResampler, + PoseSegment, +) +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[Any], + frame_times: list[float], + ) -> StepResult: + pose_segments = cast(list[PoseSegment], segments) + 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(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)) + / 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, + legacy_segment_resampler_factory=KeyboardResampler, + ) + 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/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/omnidreams/interactive_drive/world_model/flashdreams_adapter.py b/integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py index e61b4904c..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,7 +32,7 @@ VideoPostprocessChainConfig, VideoPostprocessStream, ) -from flashdreams.infra.video_output import lazy_rgb_frames_from_video_tensor +from flashdreams.infra.video_output import VideoOutputStream PipelineFactory = Callable[[WorldModelManifest, WorldModelProfileConfig], Any] _VIEW_NAMES = ["camera_front_wide_120fov"] @@ -493,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._postprocess_stream: VideoPostprocessStream | None = None @property def pipeline(self) -> Any: @@ -607,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( @@ -621,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: " @@ -630,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._postprocess_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: " @@ -657,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), + result = model_session.step( + self._condition_tensor(condition_frames), + delay_finalization=True, ) - video = self._postprocess_video( - video, autoregressive_index=self._next_block_index - ) - 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( @@ -681,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( @@ -693,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: @@ -709,43 +695,37 @@ 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", 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", ) - 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." - ) - return processed - def _close_postprocess_stream(self) -> None: - if self._postprocess_stream is None: - return - self._postprocess_stream.finish() - self._postprocess_stream = None + 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 self._model_session def _initialize_cache(self, initial_rgb: object, prompt: str) -> Any: if self.manifest.synthetic_model: @@ -843,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 2332d1c65..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, @@ -50,6 +51,7 @@ 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).""" @@ -154,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 @@ -391,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 @@ -403,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,9 +438,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 = write_video_tensor( canvas, video_path, fps=cfg.output_fps, @@ -442,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 c62724422..000000000 --- a/integrations/omnidreams/omnidreams/webrtc/server.py +++ /dev/null @@ -1,375 +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("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, - 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 86e4d020b..000000000 --- a/integrations/omnidreams/omnidreams/webrtc/session.py +++ /dev/null @@ -1,1179 +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.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, - WebRTCStepResult, - make_webrtc_step_result, -) -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._postprocess_stream: VideoPostprocessStream | None = 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], - ) -> WebRTCStepResult: - 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], - ) -> WebRTCStepResult: - 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 - self._postprocess_stream = VideoPostprocessStream( - postprocess=postprocess, - output_layout="bvtchw", - fps=self.config.fps, - per_view=False, - world_size=world_size, - ) - 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 - - def _generate_one_chunk_sync( - self, - *, - segments: list[PoseSegment], - frame_times: list[float], - ) -> WebRTCStepResult: - 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 not serve_hdmaps and self._postprocess_stream is not None: - video_chunk = self._postprocess_stream.process( - video_chunk, - autoregressive_index=self.autoregressive_index, - ) - - 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 - - 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/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 99d35d73e..fe48be188 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 @@ -77,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]", ] @@ -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.app: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. @@ -133,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.demo" = ["web/adapter.js", "web/adapter.css"] "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 new file mode 100644 index 000000000..846e78f96 --- /dev/null +++ b/integrations/omnidreams/tests/test_demo_api.py @@ -0,0 +1,1866 @@ +# 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 sys +from collections.abc import Mapping, 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.controls import SPARSE_KEY_SEGMENTS_METADATA_KEY +from omnidreams.demo.replay import ( + OmnidreamsReplayRuntime, + OmnidreamsReplayRuntimeOptions, + OmnidreamsReplaySession, +) +from omnidreams.demo.runtime import ( + OmnidreamsRuntime, + OmnidreamsRuntimeOptions, + OmnidreamsSession, +) +from omnidreams.demo.webrtc import ( + OmnidreamsWebRTCModelRuntimeConfig, + _should_use_legacy_webrtc_path, + serve_omnidreams_webrtc_demo, +) + +from flashdreams.runtime import ( + CanonicalInputs, + InferenceConfig, + InferenceInput, + OutputArtifact, + OutputTarget, + StepRequest, + StepRequirements, + StepResult, + UserInputEvent, + UserInputs, +) +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, + 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, +) + +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_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", "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: + 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: + 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}, + ), + ) + + result = run_replay_demo( + spec=spec, + adapter=adapter, + output_target_factory=lambda output_spec: output, + runner=fake_runner, + ) + + 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 + 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_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: + _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, + frame_times=(1 / 30, 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, + frame_times=(1 / 30, 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), + inputs=UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, + event_type="key_down", + payload={"key": "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_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, +) -> 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( + 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, + ) + + 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": [["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=OmnidreamsRuntimeOptions( + 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}) + ) + 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 + 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 + 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.released_encoders is True + 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: + 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, + input_mode="keyboard-driving", + scenario=OmnidreamsWebRTCScenario( + scene_uuid="scene-1", + scene_variant="rain", + camera_name="camera_front_wide_120fov", + 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}, + ), + ) + + 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, + shared_runtime_factory=shared_runtime_factory, + server_runner=lambda **kwargs: calls.append(kwargs), + ) + + manager = calls[0]["session_manager"] + assert type(manager) is BaseWebRTCSessionManager + 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( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import flashdreams.serving.webrtc.demo 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"] + if configure_app := kwargs["configure_app"]: + configure_app(app) + return app + + monkeypatch.setattr( + shared_webrtc_module, + "create_packaged_webrtc_app", + fake_create_packaged_webrtc_app, + ) + 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()}, + ), + ) + + app = serve_omnidreams_webrtc_demo( + spec=spec, + shared_runtime_factory=lambda **kwargs: _FactoryRuntime(), + server_runner=lambda **kwargs: None, + ) + + 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/demo/web") + 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: + import flashdreams.serving.webrtc.demo 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"] + if configure_app := kwargs["configure_app"]: + configure_app(app) + return app + + def fake_server_runner(**kwargs: Any) -> None: + server_calls.append(kwargs) + + monkeypatch.setattr( + shared_webrtc_module, + "create_packaged_webrtc_app", + fake_create_packaged_webrtc_app, + ) + 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()}, + ), + ) + + app = serve_omnidreams_webrtc_demo( + spec=spec, + world_rank=0, + shared_runtime_factory=lambda **kwargs: _FactoryRuntime(), + server_runner=fake_server_runner, + ) + + assert len(server_calls) == 1 + assert server_calls[0]["world_rank"] == 0 + assert server_calls[0]["app"] is app + assert server_calls[0]["host"] == "0.0.0.0" + assert server_calls[0]["port"] == 8082 + assert type(server_calls[0]["session_manager"]) is BaseWebRTCSessionManager + + +@pytest.mark.asyncio +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) + 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 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 [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: + 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, + 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( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + input_mode="replay", + scenario={ + "conditioning_mode": OMNIDREAMS_CONDITIONING_LUDUS, + "keyboard_events": tuple(keyboard_events), + "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( + 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 + 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]: + del autoregressive_index, cache + return {"denoise_s": 0.25} + + +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 _FakeWebRTCResampler: + def __init__(self, *, start_v: float, fps: int) -> None: + self.dt = 1.0 / fps + self.next_chunk_start_v = start_v + + 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 + 1) * self.dt for index in range(num_frames)] + end = frame_times[-1] + self.next_chunk_start_v = end + return frame_times + + +class _FakeWebRTCVideoTrack: + def __init__(self, *, fps: int) -> None: + self.fps = fps + self.closed = False + self.enqueued: list[StepResult] = [] + + async def enqueue_result(self, result: StepResult) -> int: + self.enqueued.append(result) + return result.frame_count + + def qsize(self) -> int: + return 0 + + 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, + payload: object, + track: Any, + *, + force_keyframe: bool = False, + ) -> SimpleNamespace: + del force_keyframe + if not isinstance(payload, StepResult): + raise TypeError("Fake WebRTC encoder expected a StepResult payload.") + return SimpleNamespace( + num_frames=await track.enqueue_result(payload), + encode_ms=0.0, + ) + + +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: + 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_input_fps(self) -> float: + return 30.0 + + def peek_steady_output_num_frames(self) -> int: + return 1 + + 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 request, 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/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 9f46e6be7..000000000 --- a/integrations/omnidreams/tests/test_webrtc_runtime.py +++ /dev/null @@ -1,1475 +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.serving.webrtc.controls import ( - WSAD_SUPPORTED_KEYS, - CameraPoseIntegrator, -) -from flashdreams.serving.webrtc.encoders import ( - 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 - -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._postprocess_stream = postprocess_stream # ty:ignore[invalid-assignment] - - 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._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._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: - web_dir = files("omnidreams.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 '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 "postprocess_preset: postprocessPreset" in javascript - - -@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], - ) -> WebRTCStepResult: - del frame_times - chunk_index = len(self.generated_segments) - self.generated_segments.append(segments) - return WebRTCStepResult( - 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], - ) -> WebRTCStepResult: - 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 f052aac10..000000000 --- a/integrations/omnidreams/tests/test_webrtc_server_routes.py +++ /dev/null @@ -1,308 +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 ( - "Omnidreams 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 "Omnidreams 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 - 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 - 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_omnidreams_controls_and_lingbot_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 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_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 da01d9787..3e4857c7f 100644 --- a/integrations/self_forcing/self_forcing/runner.py +++ b/integrations/self_forcing/self_forcing/runner.py @@ -26,16 +26,15 @@ 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, WanInferencePipelineCache, ) +from flashdreams.runtime.video_output import Mp4VideoOutputTarget __all__ = [ "SelfForcingT2VRunnerConfig", @@ -127,29 +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, + ) + ) - generated = output_stream.finish() - if generated is None: + tail = output_stream.finish() + if tail is not None: + output_target.write(tail) + artifacts = output_target.close() + if not artifacts: 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_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 7e0c1fe5d..2098c746b 100644 --- a/integrations/wan21/wan21/runner.py +++ b/integrations/wan21/wan21/runner.py @@ -28,19 +28,18 @@ 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, WanInferencePipelineCache, ) +from flashdreams.runtime.video_output import Mp4VideoOutputTarget __all__ = [ "Wan21I2VRunnerConfig", @@ -173,17 +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) - generated = output_stream.finish() - if generated 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 - - # 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_artifact = artifacts[0] + video_path = Path(video_artifact.uri) logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " @@ -191,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()}" diff --git a/uv.lock b/uv.lock index a077bb380..d5295f82f 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]] @@ -1265,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" }, @@ -1280,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" }, @@ -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" }, ]