From 64e8e7b86035f1f08aca0d8d95dfc139ab692384 Mon Sep 17 00:00:00 2001 From: Kedi Wu Date: Wed, 21 Oct 2026 12:19:04 -0700 Subject: [PATCH 1/5] Re-add the SGLang notebooks for the Simulation cookbooks Four notebooks serve the Simulation models through the published lmsysorg/sglang:dev image (sgl-project/sglang#39475): Sim-Bimanual, Sim-Transfer, Sim-Auto, and Transfer-Auto (WSM-only and joint camera plus LiDAR). SGLang opens inputs by server-visible path instead of uploads, so the notebooks stage their inputs under the output folder inside the mounted cookbook root and the server command mounts that root. - helpers/sim_auto_sglang_helpers.py builds the Sim-Auto multipart form; helpers/auto_helpers.py gains the SGLang request, asynchronous job, control staging, and LiDAR validation paths; bimanual/bimanual_helpers.py encodes the action payload within the 1 MiB form-field limit. - CPU tests cover the request contracts and the notebook flows against a mock server; the cookbook workflow lints the new helper. - The Simulation README, the Auto overview, and the Transfer-Auto guide link the notebooks. Migrated from nvidia-cosmos/cosmos-framework-private branch readd-sim-sglang (ede4d0e9) as one change relative to its base 84a653ac, since the public history was synced by squash in #292. Ported to this repository: Hugging Face revisions follow main as in #294, because the previously pinned commits no longer exist, and the README keeps main's Framework links beside the new SGLang ones. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Kedi Wu --- .../cookbooks/cosmos3/simulation/README.md | 11 +- .../cosmos3/simulation/auto/README.md | 2 + .../cookbooks/cosmos3/simulation/auto/auto.md | 76 +- .../simulation/auto/helpers/auto_helpers.py | 243 +++++- .../auto/helpers/sim_auto_sglang_helpers.py | 58 ++ .../auto/run_auto_with_sglang.ipynb | 534 ++++++++++++ .../auto/run_auto_with_vllm_omni.ipynb | 2 +- .../auto/run_sim-auto_with_sglang.ipynb | 763 ++++++++++++++++++ .../simulation/auto/tests/test_auto_media.py | 13 + .../auto/tests/test_daytime_sample.py | 181 ++++- .../auto/tests/test_sim_auto_helpers.py | 183 +++++ .../simulation/bimanual/bimanual_helpers.py | 354 ++++++++ .../bimanual/run_bimanual_with_sglang.ipynb | 483 +++++++++++ .../run_bimanual_with_vllm_omni.ipynb | 2 +- .../transfer/run_transfer_with_sglang.ipynb | 381 +++++++++ .../run_transfer_with_vllm_omni.ipynb | 3 +- .github/workflows/validate-cookbooks.yml | 2 +- 17 files changed, 3263 insertions(+), 28 deletions(-) create mode 100644 .agents/skills/cookbooks/cosmos3/simulation/auto/helpers/sim_auto_sglang_helpers.py create mode 100644 .agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_sglang.ipynb create mode 100644 .agents/skills/cookbooks/cosmos3/simulation/auto/run_sim-auto_with_sglang.ipynb create mode 100644 .agents/skills/cookbooks/cosmos3/simulation/bimanual/bimanual_helpers.py create mode 100644 .agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_sglang.ipynb create mode 100644 .agents/skills/cookbooks/cosmos3/simulation/transfer/run_transfer_with_sglang.ipynb diff --git a/.agents/skills/cookbooks/cosmos3/simulation/README.md b/.agents/skills/cookbooks/cosmos3/simulation/README.md index 910bd3e58..7e7dc709d 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/README.md +++ b/.agents/skills/cookbooks/cosmos3/simulation/README.md @@ -8,10 +8,10 @@ spatial control video. | Model | Generates | From | Start here | Notebooks | | --- | --- | --- | --- | --- | -| [Cosmos3-Nano-Sim-Bimanual](https://huggingface.co/nvidia/Cosmos3-Nano-Sim-Bimanual) | Video rollout of bimanual robot manipulation | Initial RGB frame, task description, time-aligned action trajectory | [Bimanual guide](bimanual/bimanual.md) | [Framework](bimanual/bimanual.md), [vLLM-Omni](bimanual/run_bimanual_with_vllm_omni.ipynb) | -| [Cosmos3-Nano-Transfer-Auto](https://huggingface.co/nvidia/Cosmos3-Nano-Transfer-Auto) | Synchronized multi-camera driving video, joint mode adds LiDAR | Per-camera WSM controls and captions, HD-map LiDAR controls for joint mode | [Auto overview](auto/README.md), [Transfer-Auto guide](auto/auto.md) | [Framework](auto/run_auto_with_cosmos_framework.ipynb), [vLLM-Omni](auto/run_auto_with_vllm_omni.ipynb) | -| [Cosmos3-Nano-Sim-Auto](https://huggingface.co/nvidia/Cosmos3-Nano-Sim-Auto) | Same outputs as Transfer-Auto, generated chunk by chunk | Same inputs as Transfer-Auto | [Auto overview](auto/README.md) | [Framework](auto/run_sim-auto_with_cosmos_framework.ipynb), [vLLM-Omni](auto/run_sim_auto_with_vllm_omni.ipynb) | -| [Cosmos3-Nano-Sim-Transfer](https://huggingface.co/nvidia/Cosmos3-Nano-Sim-Transfer) | 480p video | Caption and one edge, blur, depth, or segmentation control video | [Sim-Transfer guide](transfer/transfer.md) | [Framework](transfer/transfer.md), [vLLM-Omni](transfer/run_transfer_with_vllm_omni.ipynb) | +| [Cosmos3-Nano-Sim-Bimanual](https://huggingface.co/nvidia/Cosmos3-Nano-Sim-Bimanual) | Video rollout of bimanual robot manipulation | Initial RGB frame, task description, time-aligned action trajectory | [Bimanual guide](bimanual/bimanual.md) | [Framework](bimanual/bimanual.md), [vLLM-Omni](bimanual/run_bimanual_with_vllm_omni.ipynb), [SGLang](bimanual/run_bimanual_with_sglang.ipynb) | +| [Cosmos3-Nano-Transfer-Auto](https://huggingface.co/nvidia/Cosmos3-Nano-Transfer-Auto) | Synchronized multi-camera driving video, joint mode adds LiDAR | Per-camera WSM controls and captions, HD-map LiDAR controls for joint mode | [Auto overview](auto/README.md), [Transfer-Auto guide](auto/auto.md) | [Framework](auto/run_auto_with_cosmos_framework.ipynb), [vLLM-Omni](auto/run_auto_with_vllm_omni.ipynb), [SGLang](auto/run_auto_with_sglang.ipynb) | +| [Cosmos3-Nano-Sim-Auto](https://huggingface.co/nvidia/Cosmos3-Nano-Sim-Auto) | Same outputs as Transfer-Auto, generated chunk by chunk | Same inputs as Transfer-Auto | [Auto overview](auto/README.md) | [Framework](auto/run_sim-auto_with_cosmos_framework.ipynb), [vLLM-Omni](auto/run_sim_auto_with_vllm_omni.ipynb), [SGLang](auto/run_sim-auto_with_sglang.ipynb) | +| [Cosmos3-Nano-Sim-Transfer](https://huggingface.co/nvidia/Cosmos3-Nano-Sim-Transfer) | 480p video | Caption and one edge, blur, depth, or segmentation control video | [Sim-Transfer guide](transfer/transfer.md) | [Framework](transfer/transfer.md), [vLLM-Omni](transfer/run_transfer_with_vllm_omni.ipynb), [SGLang](transfer/run_transfer_with_sglang.ipynb) | To drive the Auto models with your own scene, render WSM controls with [generate_wsm.ipynb](auto/generate_wsm.ipynb). @@ -37,6 +37,9 @@ distilled student: a fixed four-step sampler with a rolling key/value cache. Bimanual and Sim-Transfer add model-specific setup steps in their guides. - **vLLM-Omni:** Docker with the NVIDIA Container Toolkit. All four notebooks use the `vllm/vllm-omni:cosmos3` image. +- **SGLang:** Docker with the NVIDIA Container Toolkit. The four SGLang notebooks use + the published `lmsysorg/sglang:dev` image; the Auto and Sim-Transfer notebooks read + their inputs from the mounted cookbook root. If you pass `--defaults-file` to Framework inference, it **replaces** the built-in sampling defaults rather than merging with them. Start from the defaults file the diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/README.md b/.agents/skills/cookbooks/cosmos3/simulation/auto/README.md index 2ac34984f..7b50f60a2 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/auto/README.md +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/README.md @@ -11,7 +11,9 @@ controls. Joint generation also needs numerical HD-map LiDAR controls. | Native Transfer-Auto | [Framework](run_auto_with_cosmos_framework.ipynb) | Seven cameras, 297 frames, 720p; optional LiDAR | | Native Sim-Auto | [Framework](run_sim-auto_with_cosmos_framework.ipynb) | Seven cameras, 297 frames, 720p; optional LiDAR | | Served Transfer-Auto | [vLLM-Omni](run_auto_with_vllm_omni.ipynb) | Seven cameras, 201 frames, 480p; WSM-only and joint examples | +| Served Transfer-Auto | [SGLang](run_auto_with_sglang.ipynb) | Seven cameras, 201 frames, 480p; WSM-only and joint examples from the published `lmsysorg/sglang:dev` image | | Served Sim-Auto | [vLLM-Omni](run_sim_auto_with_vllm_omni.ipynb) | Seven cameras, 297 frames, 480p; RGB and joint examples | +| Served Sim-Auto | [SGLang](run_sim-auto_with_sglang.ipynb) | Seven cameras, 297 frames, 480p; RGB example from the published `lmsysorg/sglang:dev` image | | Render your controls | [WSM renderer](generate_wsm.ipynb) | Controls from your own RDS-HQ scene and calibration | | Sim-Auto with FlashDreams | [FlashDreams](https://github.com/NVIDIA/flashdreams/integrations_v2/cosmos3_nano_sim_auto) | Integrate with [FlashDreams](https://flashdreams.org/main/models/cosmos3_nano_sim_auto/) for accelerated inference | diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/auto.md b/.agents/skills/cookbooks/cosmos3/simulation/auto/auto.md index f9ab75b1a..a3d569c54 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/auto/auto.md +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/auto.md @@ -16,12 +16,17 @@ model), and [Cosmos3-Nano-Sim-Transfer](../transfer/transfer.md). ## Supported backends -This cookbook supports two runtimes: +This cookbook supports three runtimes: - **Cosmos Framework** for export, native PyTorch inference. - **vLLM-Omni** for an OpenAI-compatible video server. The runnable [Multiview Simulation notebook](run_auto_with_vllm_omni.ipynb) uses the published WSM controls and the synchronous video endpoint, followed by a joint WSM/RGB/LiDAR example that retrieves both generated RGB and numeric LiDAR through asynchronous jobs. +- **SGLang** for asynchronous WSM and joint camera/LiDAR generation. Its + [Multiview Simulation notebook](run_auto_with_sglang.ipynb) uses server-visible input paths + and downloads generated camera videos and numeric LiDAR through the Videos API. The + [Sim-Auto SGLang notebook](run_sim-auto_with_sglang.ipynb) runs the causal student's + seven-camera sample through the same API. These examples consume the same fixed camera order and WSM-per-view contract. The commands, request fields, and environment are runtime-specific. @@ -295,6 +300,19 @@ you submit generation requests. The complete client is in ### Send the multipart request +For SGLang generation, use [`run_auto_with_sglang.ipynb`](run_auto_with_sglang.ipynb). +Call `prepare_request(..., backend="sglang")` to preserve absolute `control_path` and `vision_path` values +and return multipart form fields with an empty upload list. Both backends encode +`extra_params` as a JSON string. The control files must be readable +at the same paths on the server; mount the checkout and prepared outputs into the server +container at those paths. SGLang does not accept per-camera `control_reference_index` fields. +Then call `submit_wsm(..., async_mode=True, timeout=3600)`. The helper submits multipart form data to `POST /v1/videos`, +polls `GET /v1/videos/{id}`, and downloads `/content` when the job completes. It saves the +latest status in `job.json` and raises on failure, cancellation, or polling timeout. +For joint LiDAR, use `submit_joint(..., backend="sglang")` and the +[SGLang LiDAR contract](#joint-rgb-and-lidar-with-sglang) below. The helpers default to the +vLLM-Omni multipart contract and synchronous WSM endpoint shown below. + `POST /v1/videos/sync` accepts one control per camera as repeated multipart `input_references` parts. In the JSON-encoded `extra_params`, each ordered view replaces its local `control_path` with a zero-based `control_reference_index` pointing at the corresponding upload. Each view @@ -394,6 +412,62 @@ metadata against the job descriptor. It displays the camera grid and saves `sample_outputs.json`. Generated LiDAR remains in the downloaded safetensors file with its metadata sidecar for further analysis. +### Joint RGB and LiDAR with SGLang + +The LiDAR section of [`run_auto_with_sglang.ipynb`](run_auto_with_sglang.ipynb) uses the same +daytime sample and [joint request](assets/daytime_scene/transfer_auto_joint.json). `prepare_request(..., backend="sglang")` retains each camera's absolute +`control_path` and `vision_path`, retains the numeric `lidar.control_path`, and translates the +manifest's `lidar.return_output` to SGLang's `lidar.decode`. It returns an empty upload list; +the server reads the files directly. Mount the checkout and prepared outputs at the same +absolute paths inside the server container. The joint checkpoint must include `lidar_vae/`; +the published `lmsysorg/sglang:dev` image runs the LiDAR VAE on its built-in FlexAttention kernel. + +With generated LiDAR enabled, the JSON-encoded `extra_params` contains: + +```json +{ + "lidar": { + "control_path": "/absolute/path/to/cosmos-framework/.agents/skills/cookbooks/cosmos3/simulation/outputs/sglang/auto/lidar/prepared/lidar_control.safetensors", + "decode": true + } +} +``` + +The joint preparation cell converts the sample's `control.tar` into the FP32 tensor +`frames[3,100,128,1800]` under the output folder. SGLang selects its first +67 sweeps for 201 camera frames at 30 FPS and a 10 Hz LiDAR checkpoint, and handles padding +and normalization. `submit_joint(..., backend="sglang")` submits `POST /v1/videos`, polls the +job, and downloads the camera video from `/content` and numeric LiDAR from +`/content?variant=rangemap`. Setting the manifest's `return_output=false` produces +`decode=false`: LiDAR still conditions camera generation, but no numeric artifact is returned. +The notebook keeps output enabled for its validation cells. + +SGLang's generated safetensors contains three tensors, each shaped `[67,128,1800]`: + +| Tensor | Dtype | Meaning | +| --- | --- | --- | +| `range_m` | FP32 | Range in metres; invalid rays are zero | +| `intensity` | FP16 | Intensity in [0,1]; invalid rays are zero | +| `validity` | uint8 | Binary validity mask | + +`validate_sglang_lidar()` checks these tensors and embedded metadata. The notebook compares +them with the job descriptor's `sweeps`, `fps`, `height`, `width`, `min_range_m`, and +`max_range_m`. SGLang's generated file uses a different layout from the input's `frames` tensor. + +Default paths below are relative to `.agents/skills/cookbooks/cosmos3/simulation/`: + +| Artifact | Path | +| --- | --- | +| Input LiDAR | `outputs/sglang/auto/lidar/prepared/lidar_control.safetensors` (converted locally) | +| Generated LiDAR | `outputs/sglang/auto/lidar/multiview_camera_major.lidar.safetensors` | +| Generated descriptor | `outputs/sglang/auto/lidar/multiview_camera_major.lidar.json` | + +The notebook prints the resolved input and generated paths. Its server URL defaults to +`SGLANG_BASE_URL` (`http://127.0.0.1:30000`); optional authentication uses `SGLANG_API_KEY`. +`COSMOS3_OUTPUT_DIR` and `AUTO_SAMPLE_ROOT` select the output folder and a local sample copy, as in +the vLLM-Omni notebook. Because the server opens inputs by path, the notebook stages each RGB +condition under the output folder; keep that folder inside the mounted cookbook root. + ### vLLM-Omni troubleshooting - **`missing views[0].prompt`:** supply a nonempty `prompt` in every camera view, even for diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/helpers/auto_helpers.py b/.agents/skills/cookbooks/cosmos3/simulation/auto/helpers/auto_helpers.py index 268217e46..0298440ba 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/auto/helpers/auto_helpers.py +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/helpers/auto_helpers.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -"""CPU-only helpers for the Cosmos Framework and vLLM-Omni Auto notebooks.""" +"""CPU-only helpers for the Cosmos Framework, vLLM-Omni, and SGLang Auto notebooks.""" from __future__ import annotations @@ -87,8 +87,16 @@ def probe_video(path: Path) -> dict: } -def prepare_controls(views: list[dict], base_dir: Path, output_dir: Path, frames: int, fps: float) -> list[Path]: - """Decode synchronized controls and trim their first window without resizing.""" +def prepare_controls( + views: list[dict], base_dir: Path, output_dir: Path, frames: int, fps: float, *, materialize: bool = False +) -> list[Path]: + """Decode synchronized controls and trim their first window without resizing. + + A control that already has ``frames`` frames is returned in place. With + ``materialize`` it is copied into ``output_dir`` instead, so every returned path + lies under ``output_dir``; a server that opens inputs by path sees only what is + mounted there. + """ if frames < 1 or (frames - 1) % 4 or not math.isfinite(fps) or fps <= 0: raise ValueError("Use a positive 4k+1 camera frame count and positive finite FPS.") sources = [require_asset(base_dir / view["control_path"]) for view in views] @@ -130,6 +138,10 @@ def prepare_controls(views: list[dict], base_dir: Path, output_dir: Path, frames "+faststart", str(target), ) + elif materialize: + target = output_dir / f"{view['camera_key']}.mp4" + if not (target.exists() and target.samefile(source)): + shutil.copyfile(source, target) actual = probe_video(target) if any( actual[key] != expected @@ -158,6 +170,30 @@ def validate_rgb_conditions(views: list[dict], base_dir: Path) -> None: raise ValueError(f"Cannot decode RGB first frame: {path}: {result.stderr.decode(errors='replace')}") +def stage_conditions(views: list[dict], output_dir: Path, *, field: str = "vision_path") -> list[Path]: + """Copy each camera's condition file into ``output_dir`` as ````. + + SGLang opens inputs by path on the server, and a downloaded sample lives in the + client's Hugging Face cache. Staging the files under the mounted output folder + gives the server one absolute path that exists on both sides. The copy keeps the + suffix of the path named in the request: a Hugging Face snapshot entry is a + symlink into an extension-less blob store, and the server's video reader picks + its decoder by extension. + Every call rewrites the copies, so a changed condition is picked up even when + its size did not change. + """ + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + staged = [] + for view in views: + named = Path(view[field]) + source = require_asset(named) + target = output_dir / f"{view['camera_key']}{named.suffix}" + shutil.copyfile(source, target) + staged.append(target) + return staged + + def validate_lidar(path: Path, sweeps: int, *, exact: bool = False) -> dict: """Read numeric sweeps incrementally; never load a GPU library or normalize inputs.""" import numpy as np @@ -184,6 +220,59 @@ def validate_lidar(path: Path, sweeps: int, *, exact: bool = False) -> dict: return {"shape": shape, "dtype": "float32", "metadata": metadata} +def validate_sglang_lidar(path: Path, sweeps: int, *, fps: float = 10.0) -> dict: + """Validate SGLang's generated range map incrementally on the CPU.""" + import numpy as np + from safetensors import safe_open + + if sweeps < 1 or not math.isfinite(fps) or fps <= 0: + raise ValueError("Expected positive LiDAR sweep count and finite positive FPS.") + require_asset(path) + dtypes = {"range_m": "F32", "intensity": "F16", "validity": "U8"} + shape = [sweeps, 128, 1800] + with safe_open(path, framework="numpy") as handle: + if set(handle.keys()) != set(dtypes): + raise ValueError("SGLang LiDAR must contain range_m, intensity, and validity tensors.") + tensors = {name: handle.get_slice(name) for name in dtypes} + for name, tensor in tensors.items(): + if tensor.get_shape() != shape or tensor.get_dtype() != dtypes[name]: + raise ValueError(f"Expected {name} with shape {shape} and dtype {dtypes[name]}.") + metadata = handle.metadata() or {} + try: + actual_fps = float(metadata["fps"]) + min_range = float(metadata["min_range_m"]) + max_range = float(metadata["max_range_m"]) + invalid_range = float(metadata["invalid_range_m"]) + except (KeyError, ValueError) as exc: + raise ValueError("Missing or invalid SGLang LiDAR metadata.") from exc + if ( + metadata.get("layout") != "[sweeps, beams, azimuth]" + or actual_fps != fps + or not math.isfinite(min_range) + or not math.isfinite(max_range) + or not 0 <= min_range < max_range + or invalid_range != 0 + ): + raise ValueError(f"Unexpected SGLang LiDAR metadata: {metadata}.") + for index in range(sweeps): + ranges = tensors["range_m"][index] + intensity = tensors["intensity"][index] + validity = tensors["validity"][index] + if ( + not np.isfinite(ranges).all() + or not np.isfinite(intensity).all() + or (ranges < 0).any() + or (intensity < 0).any() + or (intensity > 1).any() + or not np.isin(validity, [0, 1]).all() + ): + raise ValueError(f"Invalid SGLang LiDAR values in sweep {index}.") + valid = validity == 1 + if (ranges[valid] <= 0).any() or (ranges[~valid] != 0).any() or (intensity[~valid] != 0).any(): + raise ValueError(f"Invalid masked SGLang LiDAR values in sweep {index}.") + return {"shape": shape, "dtypes": dtypes, "metadata": metadata} + + def split_cameras( path: Path, output_dir: Path, camera_order: tuple, frames: int, fps: float, *, width: int = 832, height: int = 480 ) -> dict: @@ -300,8 +389,32 @@ def validate_framework_outputs(output_dir: Path, manifest: dict) -> dict: # vLLM-Omni request preparation and HTTP client. -def prepare_request(manifest: dict, base_dir: Path, model: str) -> tuple[dict, list[Path]]: - """Translate local manifest paths to the existing multipart Videos API contract.""" +def check_model(client, expected: str | None = None) -> str: + """Return the served model ID from ``GET /v1/models``, requiring ``expected`` when given.""" + response = client.get("/v1/models") + response.raise_for_status() + ids = [item["id"] for item in response.json().get("data", []) if item.get("id")] + if expected: + if expected not in ids: + raise ValueError(f"Expected served model {expected!r}; found {ids}.") + return expected + if len(ids) != 1: + raise ValueError(f"Set the expected model explicitly; server returned {ids}.") + return ids[0] + + +def prepare_request( + manifest: dict, base_dir: Path, model: str, *, backend: str = "vllm_omni" +) -> tuple[dict, list[Path]]: + """Build a Videos API request and its upload list. + + Both backends use multipart form fields. vLLM-Omni uploads every file and + refers to it by a reference index; SGLang reads each file from an absolute + path visible to the server, so the upload list stays empty. For SGLang, + ``lidar.return_output`` is translated to ``lidar.decode``. + """ + if backend not in ("vllm_omni", "sglang"): + raise ValueError(f"Unknown backend {backend!r}; use 'vllm_omni' or 'sglang'.") for field in ("negative_prompt", "negative_metadata_mode"): if manifest.get(field) is not None: raise ValueError( @@ -334,19 +447,38 @@ def prepare_request(manifest: dict, base_dir: Path, model: str) -> tuple[dict, l raise ValueError("Supply local paths, not upload indexes, in the input manifest.") if f"{role}_path" in view: path = require_asset(base_dir / view.pop(f"{role}_path")) - view[f"{role}_reference_index"] = len(paths) - paths.append(path) + if backend == "sglang": + view[f"{role}_path"] = str(path) + else: + view[f"{role}_reference_index"] = len(paths) + paths.append(path) if "lidar" in extra: - if extra["lidar"].get("return_output") is not True: + lidar = extra["lidar"] + if backend == "sglang": + if set(lidar) - {"control_path", "return_output", "decode"}: + raise ValueError("SGLang LiDAR accepts control_path and decode (or the return_output alias).") + decode = lidar.get("decode", lidar.get("return_output", False)) + if not isinstance(decode, bool) or ( + "return_output" in lidar + and (not isinstance(lidar["return_output"], bool) or lidar["return_output"] != decode) + ): + raise ValueError("LiDAR decode/return_output must be booleans and agree when both are supplied.") + elif lidar.get("return_output") is not True: raise ValueError("The joint walkthrough requires lidar.return_output=true.") + reference_suffix = "path" if backend == "sglang" else "reference_index" if any( - "vision_reference_index" not in view or "control_reference_index" not in view + f"vision_{reference_suffix}" not in view or f"control_{reference_suffix}" not in view for view in extra["multiview"]["views"] ): raise ValueError("The joint I2V example requires WSM and RGB for every camera.") - path = require_asset(base_dir / extra["lidar"].pop("control_path")) - extra["lidar"]["control_reference_index"] = len(paths) - paths.append(path) + if backend == "sglang": + lidar.pop("return_output", None) + lidar["decode"] = decode + lidar["control_path"] = str(require_asset(base_dir / lidar["control_path"])) + else: + path = require_asset(base_dir / lidar.pop("control_path")) + lidar["control_reference_index"] = len(paths) + paths.append(path) if len(paths) > 23 or any(path.stat().st_size > 512 * 1024**2 for path in paths): raise ValueError("Videos API accepts at most 23 files, each no larger than 512 MiB.") form = { @@ -367,6 +499,11 @@ def prepare_request(manifest: dict, base_dir: Path, model: str) -> tuple[dict, l def _post(client, endpoint: str, form: dict, paths: list[Path]): + """Send multipart form fields with optional file uploads.""" + if not paths: + # A None filename makes each part a regular field. Passing only data= + # would use URL encoding, which SGLang's Videos API does not parse. + return client.post(endpoint, files=[(key, (None, value)) for key, value in form.items()]) with ExitStack() as stack: files = [ ( @@ -405,7 +542,37 @@ def _download(client, url: str, path: Path) -> None: temporary.unlink(missing_ok=True) -def submit_joint(client, form: dict, paths: list[Path], output_dir: Path, *, timeout: float = 3600) -> dict: +def submit_wsm( + client, form: dict, paths: list[Path], output_dir: Path, *, async_mode: bool = False, timeout: float = 3600 +) -> Path: + """Save the WSM video, optionally submitting and polling an async job. + + For SGLang, prepare the request with backend="sglang" and set async_mode=True. + The call blocks until the video is saved; timeout limits async job polling. + """ + path = output_dir / "multiview_camera_major.mp4" + if async_mode: + job = _submit_video_job(client, form, paths, output_dir, timeout=timeout) + _download(client, f"/v1/videos/{job['id']}/content", path) + return path + output_dir.mkdir(parents=True, exist_ok=True) + response = _post(client, "/v1/videos/sync", form, paths) + _check_response(response, output_dir) + if not response.content: + raise ValueError("The synchronous video response was empty.") + path.write_bytes(response.content) + return path + + +def _submit_video_job( + client, form: dict, paths: list[Path], output_dir: Path, *, timeout: float, interval: float = 10 +) -> dict: + """Submit an async Videos API job and poll until it completes. + + Every job response is saved as ``job.json``. A failed or cancelled job saves + the response as ``response_error.txt`` and raises; a polling timeout raises + but leaves the job running on the server. + """ if not math.isfinite(timeout) or timeout <= 0: raise ValueError("Polling timeout must be positive and finite.") output_dir.mkdir(parents=True, exist_ok=True) @@ -417,9 +584,9 @@ def submit_joint(client, form: dict, paths: list[Path], output_dir: Path, *, tim print(f"Submitted {job['id']} at {client.base_url.join(job_url)}") while True: (output_dir / "job.json").write_text(json.dumps(job, indent=2) + "\n") - if job.get("status") == "failed": + if job.get("status") in {"failed", "cancelled"}: (output_dir / "response_error.txt").write_text(json.dumps(job, indent=2)) - raise RuntimeError(f"Video generation failed: {job.get('error')}") + raise RuntimeError(f"Video generation {job['status']}: {job.get('error')}") _check_response(response, output_dir) if job.get("status") == "completed": break @@ -428,18 +595,42 @@ def submit_joint(client, form: dict, paths: list[Path], output_dir: Path, *, tim remaining = deadline - time.monotonic() if remaining <= 0: raise TimeoutError(f"Polling timed out; job remains at {client.base_url.join(job_url)}. See job.json.") - time.sleep(min(1, remaining)) + time.sleep(min(interval, remaining)) response = client.get(job_url, timeout=max(0.001, min(30, deadline - time.monotonic()))) try: job = response.json() except ValueError: _check_response(response, output_dir) raise - if not job.get("lidar"): + return job + + +def submit_joint( + client, form: dict, paths: list[Path], output_dir: Path, *, timeout: float = 3600, backend: str = "vllm_omni" +) -> dict: + """Save a joint job's video and, for SGLang decode requests, its native LiDAR artifact. + + vLLM-Omni returns only the LiDAR descriptor; download its numeric output from /lidar separately. + """ + return_lidar = True + if backend == "sglang": + return_lidar = json.loads(form["extra_params"])["lidar"].get("decode", False) + job = _submit_video_job(client, form, paths, output_dir, timeout=timeout) + job_url = f"/v1/videos/{job['id']}" + if return_lidar and not job.get("lidar"): raise ValueError("Completed job is missing its requested LiDAR artifact; see job.json.") + if return_lidar and backend == "sglang" and not job["lidar"].get("files", {}).get("rangemap"): + raise ValueError("Completed job is missing its LiDAR rangemap file; see job.json.") video_path = output_dir / "multiview_camera_major.mp4" _download(client, job_url + "/content", video_path) - return {"video": video_path, "descriptor": job["lidar"], "job_id": job["id"]} + if backend != "sglang": + return {"video": video_path, "descriptor": job["lidar"], "job_id": job["id"]} + if not return_lidar: + return {"video": video_path, "lidar": None, "descriptor": None, "job_id": job["id"]} + lidar_path = output_dir / "multiview_camera_major.lidar.safetensors" + _download(client, job_url + "/content?variant=rangemap", lidar_path) + lidar_path.with_suffix(".json").write_text(json.dumps(job["lidar"], indent=2) + "\n") + return {"video": video_path, "lidar": lidar_path, "descriptor": job["lidar"], "job_id": job["id"]} def download_lidar(client, job: dict, output_path: Path) -> tuple[Path, dict]: @@ -467,3 +658,19 @@ def validate_lidar_output(path: Path, descriptor: dict, sweeps: int, *, lidar_fp if info["metadata"].get(field) != descriptor.get(field): raise ValueError(f"Embedded LiDAR metadata {field!r} differs from the job descriptor.") return info + + +def save_outputs(output_dir: Path, video: Path, cameras: dict, manifest: dict, *, lidar: dict | None = None) -> None: + """Write ``sample_outputs.json`` for a camera-major response split into per-camera files.""" + summary = { + "camera_major_response": str(video), + "camera_order": list(cameras), + "frames_per_view": manifest["num_frames"], + "fps": manifest["fps"], + "files_by_camera": {key: str(path) for key, path in cameras.items()}, + } + if lidar is not None: + summary["job_id"] = lidar["job_id"] + if lidar["lidar"] is not None: + summary.update(lidar_file=str(lidar["lidar"]), lidar=lidar["descriptor"]) + (output_dir / "sample_outputs.json").write_text(json.dumps(summary, indent=2) + "\n") diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/helpers/sim_auto_sglang_helpers.py b/.agents/skills/cookbooks/cosmos3/simulation/auto/helpers/sim_auto_sglang_helpers.py new file mode 100644 index 000000000..a9ee1e201 --- /dev/null +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/helpers/sim_auto_sglang_helpers.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 +"""CPU-only request helpers for the seven-camera Sim-Auto SGLang notebook. + +SGLang serves the same ``multiview.views[]`` request as vLLM-Omni on the shared daytime +sample, but opens every WSM control, RGB condition, and LiDAR control from a path the +server can read instead of accepting uploads, and runs each rollout as one asynchronous +Videos API job without session fields. Manifest loading and the input checks are shared +with the vLLM-Omni helpers. +""" + +from __future__ import annotations + +import copy +import json +from pathlib import Path + +from helpers.auto_helpers import require_asset +from helpers.sim_auto_vllm_omni_helpers import CAMERA_ORDER, MODEL_ID, load_request, require_inputs + +__all__ = ["CAMERA_ORDER", "MODEL_ID", "load_request", "prepare_request", "require_inputs"] + + +def _number(value: float) -> str: + """Multipart form text for a number: ``30.0`` travels as ``30``.""" + return str(int(value)) if float(value).is_integer() else str(value) + + +def prepare_request(manifest: dict, model: str = MODEL_ID) -> tuple[dict, list[Path]]: + """Translate a loaded Sim-Auto request to SGLang's multipart Videos API fields. + + Returns the form fields and an empty upload list: every ``control_path`` and + ``vision_path``, and a joint request's ``lidar.control_path``, travel absolute inside + the JSON-encoded ``extra_params`` and must be readable at that path on the server. + The form carries no session fields, so two calls for the same manifest produce the + same request. + """ + require_inputs(manifest) + multiview = copy.deepcopy(manifest["multiview"]) + multiview["num_frames"] = manifest["num_frames"] + for view in multiview["views"]: + for role in ("control", "vision"): + view[f"{role}_path"] = str(require_asset(Path(view[f"{role}_path"]))) + extra = {"wsm": {}, "multiview": multiview, "resolution": "480", "aspect_ratio": "16,9"} + if "lidar" in manifest: + extra["lidar"] = {"control_path": str(require_asset(Path(manifest["lidar"]["control_path"]))), "decode": True} + form = { + "model": model, + "prompt": manifest["prompt"], + "size": "832x480", + "num_frames": str(manifest["num_frames"]), + "fps": _number(manifest["fps"]), + "seed": str(manifest["seed"]), + "guidance_scale": "1.0", + "num_inference_steps": "4", + "extra_params": json.dumps(extra), + } + return form, [] diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_sglang.ipynb b/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_sglang.ipynb new file mode 100644 index 000000000..5754caa99 --- /dev/null +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_sglang.ipynb @@ -0,0 +1,534 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "auto-sglang-license", + "metadata": {}, + "source": [ + "" + ] + }, + { + "cell_type": "markdown", + "id": "auto-sglang-intro", + "metadata": {}, + "source": [ + "# Cosmos3-Nano-Transfer-Auto with SGLang\n", + "\n", + "Cosmos3-Nano-Transfer-Auto generates a **synchronized multi-camera driving clip** from one World Scenario Model (WSM) control video per camera. The WSM controls render the scene's lane lines, road surface, and vehicle boxes from each camera's viewpoint, and a caption per camera describes its appearance. Joint requests also condition every camera on a first RGB frame and on a numeric LiDAR range map, and **generate numeric LiDAR** alongside the cameras.\n", + "\n", + "This notebook runs the two examples of the [Transfer-Auto guide](auto.md) on the shared **daytime sample**, one scene recorded by seven cameras, through one SGLang server. Each runs at **201 frames per camera, 30 FPS, and 832 × 480**:\n", + "\n", + "1. **WSM-only:** seven WSM controls produce seven camera videos.\n", + "2. **Joint RGB and LiDAR:** seven WSM controls, seven first-frame RGB conditions, and a numeric HD-map LiDAR control produce seven camera videos and 67 numeric LiDAR sweeps.\n", + "\n", + "Both run as asynchronous Videos API jobs. SGLang opens the inputs from paths the server can read, so the notebook writes every prepared input under its output folder inside the mounted cookbook root and uploads nothing. The examples are independent: run **Shared setup**, then either one.\n", + "\n", + "| Component | Version |\n", + "| --- | --- |\n", + "| SGLang | The published [`lmsysorg/sglang:dev`](https://hub.docker.com/r/lmsysorg/sglang/tags) image, which carries the Cosmos3 Multiview-AV pipeline from [sgl-project/sglang#39475](https://github.com/sgl-project/sglang/pull/39475); the tag moves, so pull it before starting the server |\n", + "| `nvidia/Cosmos3-Nano-Transfer-Auto` | [`main`](https://huggingface.co/nvidia/Cosmos3-Nano-Transfer-Auto/tree/main) |\n", + "| Daytime sample | [`main`](https://huggingface.co/nvidia/Cosmos3-Nano-Transfer-Auto/tree/main/assets/sample_data) of the same repository, under `assets/sample_data` |" + ] + }, + { + "cell_type": "markdown", + "id": "auto-sglang-prerequisites", + "metadata": {}, + "source": [ + "## 1. Prerequisites\n", + "\n", + "**Server host:** Linux with Docker, the NVIDIA Container Toolkit, and at least one NVIDIA GPU. `HF_TOKEN` needs read access to `nvidia/Cosmos3-Nano-Transfer-Auto`. The server opens the inputs by path, so the cookbook checkout, including the output folder this notebook writes, must be on the server host or on a filesystem it shares with this notebook, at one absolute path.\n", + "\n", + "**Notebook client:** Python 3.10 or newer, and a Hugging Face token with read access to `nvidia/Cosmos3-Nano-Transfer-Auto`, which also publishes the sample; run `hf auth login` or set `HF_TOKEN`. In the notebook kernel environment, install:\n", + "\n", + "```bash\n", + "pip install httpx pillow numpy safetensors matplotlib ipython huggingface_hub\n", + "```\n", + "\n", + "Install FFmpeg with the `libx264` encoder on the client host; the package also includes FFprobe. On Ubuntu/Debian:\n", + "\n", + "```bash\n", + "sudo apt-get update\n", + "sudo apt-get install -y ffmpeg\n", + "```\n", + "\n", + "On macOS with Homebrew, use `brew install ffmpeg`. Verify that both `ffmpeg -version` and `ffprobe -version` work.\n", + "\n", + "**Inputs:** The examples use the shared daytime sample published with the model under `assets/sample_data`: one scene with seven synchronized cameras, each with an RGB recording, a WSM control, and a caption, plus a sparse HD-map LiDAR control archive. Step 3 downloads the files these examples need, about 100 MiB, into the Hugging Face cache and checks them against the sample's `SHA256SUMS`. To use a copy you already have, set `AUTO_SAMPLE_ROOT` to its `sample_data` folder. The joint preparation cell converts `daytime_scene/lidar/control.tar` into a local `prepared/lidar_control.safetensors` under the output folder; see the [format notes](assets/daytime_scene/README.md)." + ] + }, + { + "cell_type": "markdown", + "id": "auto-sglang-server", + "metadata": {}, + "source": [ + "## 2. Start the server\n", + "\n", + "The server runs in the published SGLang image, `lmsysorg/sglang:dev`, which includes the Cosmos3 Multiview-AV pipeline, the joint LiDAR encoder and decoder, and the asynchronous numeric-output API. On the server host, set `COSMOS_ROOT` to the absolute cookbook root (`.agents/skills`), which step 3 prints as `Cookbook root`. `SGLANG_HF_CACHE` is a host directory for the Hugging Face cache. The container runs as root by default, so the setup lines create that directory world-writable; alternatively, start the container with your host user ID. Pull the image, then run the server in a terminal and leave it running while using the notebook; both examples use it:\n", + "\n", + "```bash\n", + "docker pull lmsysorg/sglang:dev\n", + "export COSMOS_ROOT=/absolute/path/to/cosmos-framework/.agents/skills\n", + "export SGLANG_HF_CACHE=\"${SGLANG_HF_CACHE:-$HOME/.cache/sglang-huggingface}\"\n", + "mkdir -p \"$SGLANG_HF_CACHE\"\n", + "chmod 777 \"$SGLANG_HF_CACHE\"\n", + "docker run --init --runtime nvidia --gpus all \\\n", + " -v $SGLANG_HF_CACHE:/root/.cache/huggingface \\\n", + " -v \"${COSMOS_ROOT}:${COSMOS_ROOT}\" \\\n", + " -e HF_TOKEN=$HF_TOKEN \\\n", + " -p 30000:30000 \\\n", + " --ipc=host \\\n", + " lmsysorg/sglang:dev \\\n", + " sglang serve \\\n", + " --model-path nvidia/Cosmos3-Nano-Transfer-Auto \\\n", + " --revision main \\\n", + " --host 0.0.0.0 --port 30000\n", + "```\n", + "\n", + "| Setting | Effect |\n", + "| --- | --- |\n", + "| `-v \"${COSMOS_ROOT}:${COSMOS_ROOT}\"` | Mounts the cookbook root at its host path. The requests name every prepared WSM control, staged RGB condition, and converted LiDAR control by absolute path under the output folder, which must be visible at the same path. |\n", + "| `-v $SGLANG_HF_CACHE:/root/.cache/huggingface` | Keeps the downloaded checkpoint across container restarts. The setup lines default the directory to `~/.cache/sglang-huggingface`, create it, and make it world-writable. |\n", + "| `-e HF_TOKEN=$HF_TOKEN` | Read access to the gated checkpoint. |\n", + "| `--revision` | Loads the checkpoint from `main`. The first start downloads it into the mounted cache. |\n", + "| `docker pull` | The `dev` tag moves with SGLang's main branch; pulling before each start picks up the current image. |\n", + "| `--gpus all` | Select one GPU with `--gpus '\"device=0\"'`. |\n", + "| `--host 0.0.0.0` | Accepts connections from outside the container. For a remote notebook, set `SGLANG_BASE_URL`, or forward the port with `ssh -L 30000:127.0.0.1:30000 `. |\n", + "\n", + "The camera attention backend is selected automatically: FlashAttention-4 block-sparse kernels on Hopper and Blackwell, FlexAttention's Triton kernels elsewhere; `SGLANG_DIFFUSION_COSMOS3_MULTIVIEW_ATTENTION_BACKEND` (`triton` or `fa4`) overrides the choice. Joint LiDAR generation requires the checkpoint's `lidar_vae/` component; the LiDAR VAE runs on SGLang's built-in FlexAttention kernel, and `SGLANG_DIFFUSION_COSMOS3_LIDAR_USE_NATTEN=1` switches to `natten` when that package is installed in the image. The HTTP port opens only after the weights are loaded and the warmup request finishes, so wait for the `Uvicorn running` line." + ] + }, + { + "cell_type": "markdown", + "id": "auto-sglang-client", + "metadata": {}, + "source": [ + "## 3. Shared setup\n", + "\n", + "Run the notebook from anywhere inside `.agents/skills`, or set `COSMOS_COOKBOOK_ROOT` to it. These optional environment variables override the defaults:\n", + "\n", + "| Variable | Default | Purpose |\n", + "| --- | --- | --- |\n", + "| `SGLANG_BASE_URL` | `http://127.0.0.1:30000` | Server address |\n", + "| `SGLANG_API_KEY` | unset | Bearer token, if the server was started with `--api-key` |\n", + "| `COSMOS3_OUTPUT_DIR` | `.agents/skills/cookbooks/cosmos3/simulation/outputs/sglang/auto` | Output folder; each example writes to its own subfolder. Keep it inside the mounted cookbook root: the server reads every prepared input from it. |\n", + "| `AUTO_SAMPLE_ROOT` | unset | A local `sample_data` folder to use instead of downloading the sample |\n", + "\n", + "The next cell sets the configuration and imports the CPU-only request, media, and HTTP helpers from [`helpers/auto_helpers.py`](helpers/auto_helpers.py) and the sample helpers from [`helpers/daytime_sample.py`](helpers/daytime_sample.py); none needs GPU libraries. Then fetch the sample and check the server." + ] + }, + { + "cell_type": "code", + "id": "auto-sglang-config", + "metadata": {}, + "source": [ + "from pathlib import Path\n", + "import copy\n", + "import json\n", + "import os\n", + "import sys\n", + "import time\n", + "\n", + "import httpx\n", + "from IPython.display import HTML, display\n", + "\n", + "\n", + "def find_auto_root() -> Path:\n", + " configured = os.environ.get(\"COSMOS_COOKBOOK_ROOT\")\n", + " start = Path(configured).expanduser().resolve() if configured else Path.cwd().resolve()\n", + " for root in [start] if configured else [start, *start.parents]:\n", + " for candidate in (root, root / \"auto\", root / \"cookbooks\" / \"cosmos3\" / \"simulation\" / \"auto\", root / \".agents\" / \"skills\" / \"cookbooks\" / \"cosmos3\" / \"simulation\" / \"auto\"):\n", + " if (candidate / \"helpers\" / \"auto_helpers.py\").is_file():\n", + " return candidate\n", + " raise FileNotFoundError(\"Run inside .agents/skills or set COSMOS_COOKBOOK_ROOT to it.\")\n", + "\n", + "\n", + "AUTO_ROOT = find_auto_root()\n", + "SIMULATION_ROOT = AUTO_ROOT.parent\n", + "COOKBOOK_ROOT = SIMULATION_ROOT.parents[2] # The .agents/skills folder the server mounts.\n", + "if str(AUTO_ROOT) not in sys.path:\n", + " sys.path.insert(0, str(AUTO_ROOT))\n", + "\n", + "from helpers.auto_helpers import (\n", + " camera_grid,\n", + " check_model,\n", + " prepare_controls,\n", + " prepare_request,\n", + " probe_video,\n", + " save_outputs,\n", + " split_cameras,\n", + " stage_conditions,\n", + " submit_joint,\n", + " submit_wsm,\n", + " validate_lidar,\n", + " validate_rgb_conditions,\n", + " validate_sglang_lidar,\n", + ")\n", + "from helpers.daytime_sample import SAMPLE_CAMERAS, SAMPLE_REPO, SAMPLE_REVISION, fetch_sample, resolve_request\n", + "\n", + "MODEL_ID = \"nvidia/Cosmos3-Nano-Transfer-Auto\"\n", + "MODEL_REVISION = \"main\"\n", + "SGLANG_BASE_URL = os.environ.get(\"SGLANG_BASE_URL\", \"http://127.0.0.1:30000\").rstrip(\"/\")\n", + "SGLANG_API_KEY = os.environ.get(\"SGLANG_API_KEY\")\n", + "HEADERS = {\"Authorization\": f\"Bearer {SGLANG_API_KEY}\"} if SGLANG_API_KEY else {}\n", + "HTTP_TIMEOUT = 600 # Per HTTP request: submission, each poll, and each download.\n", + "GENERATION_TIMEOUT = 3600 # Job wait, from submission to completion.\n", + "OUTPUT_DIR = Path(\n", + " os.environ.get(\"COSMOS3_OUTPUT_DIR\", SIMULATION_ROOT / \"outputs\" / \"sglang\" / \"auto\")\n", + ").expanduser().resolve()\n", + "AUTO_SAMPLE_ROOT = os.environ.get(\"AUTO_SAMPLE_ROOT\") # Unset: download the sample from Hugging Face\n", + "CAPTION_FOLDER = \"captions_with_timestamp\" # Alternative: \"captions\"\n", + "NEGATIVE_PROMPT_PATH = COOKBOOK_ROOT / \"cosmos3-nano-transfer-auto\" / \"inference\" / \"assets\" / \"negative-prompt.txt\"\n", + "\n", + "OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n", + "print(\"Cookbook root:\", COOKBOOK_ROOT)\n", + "print(\"Server:\", SGLANG_BASE_URL, \"(with API key)\" if SGLANG_API_KEY else \"\")\n", + "print(\"Model:\", MODEL_ID, \"@\", MODEL_REVISION)\n", + "print(\"Sample:\", AUTO_SAMPLE_ROOT or f\"{SAMPLE_REPO} @ {SAMPLE_REVISION}\")\n", + "print(\"Outputs:\", OUTPUT_DIR)\n", + "if not OUTPUT_DIR.is_relative_to(COOKBOOK_ROOT):\n", + " print(f\"Note: {OUTPUT_DIR} is outside the cookbook root; the server must see it at the same path.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "auto-sglang-sample-heading", + "metadata": {}, + "source": [ + "### Download and verify the sample\n", + "\n", + "Download the seven cameras' RGB recordings, WSM controls, both caption styles, and the HD-map LiDAR control archive from `assets/sample_data` on Hugging Face, and verify them against `SHA256SUMS`; later runs reuse the Hugging Face cache. No model weights are downloaded here.\n", + "\n", + "`CAPTION_FOLDER=\"captions_with_timestamp\"` selects timestamped captions, which better match training and empirically improve inference. Set it to `\"captions\"` for the original dense paragraphs." + ] + }, + { + "cell_type": "code", + "id": "auto-sglang-sample", + "metadata": {}, + "source": [ + "SAMPLE_ROOT = fetch_sample(AUTO_SAMPLE_ROOT)\n", + "print(\"Verified sample:\", SAMPLE_ROOT)\n", + "print(\"Cameras:\", \", \".join(SAMPLE_CAMERAS))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "auto-sglang-verify-heading", + "metadata": {}, + "source": [ + "### Verify the server\n", + "\n", + "Wait for the server to finish starting, then confirm that it serves the Transfer-Auto model. The cell retries connection errors and server errors for up to 30 minutes, so you can run it while the server is still loading. `/v1/models` reports the model name, not its revision; the `--revision` in the server command pins that." + ] + }, + { + "cell_type": "code", + "id": "auto-sglang-verify", + "metadata": {}, + "source": [ + "deadline = time.monotonic() + 1800\n", + "while True:\n", + " try:\n", + " with httpx.Client(base_url=SGLANG_BASE_URL, headers=HEADERS, timeout=30) as client:\n", + " served = check_model(client, MODEL_ID)\n", + " break\n", + " except httpx.TransportError as error:\n", + " reason = type(error).__name__\n", + " except httpx.HTTPStatusError as error:\n", + " if error.response.status_code < 500:\n", + " raise\n", + " reason = f\"HTTP {error.response.status_code}\"\n", + " if time.monotonic() > deadline:\n", + " raise TimeoutError(f\"No ready server at {SGLANG_BASE_URL} after 1800s ({reason}).\")\n", + " time.sleep(5)\n", + "print(\"Server ready; serving\", served)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "auto-sglang-wsm-start", + "metadata": {}, + "source": [ + "## 4. WSM-only example\n", + "\n", + "### Prepare the inputs\n", + "\n", + "The [WSM request](assets/daytime_scene/transfer_auto_wsm.json) asks for **201 frames per camera at 30 FPS** with `resolution: \"480\"` and `aspect_ratio: \"16,9\"`, seed 0, 35 steps, guidance 3, and flow shift 10, matching the recommended Transfer-Auto baseline. It lists the sample's seven cameras in the checkpoint's camera order, each with its WSM control and caption file. The server identifies each camera by its `camera_key`, so the seven keep their trained camera IDs 0 to 6; the generated video follows the request order.\n", + "\n", + "`resolve_request` resolves the paths under the sample and replaces each caption file with its text, which covers all 300 source frames. The Videos API also takes a top-level caption; the request copies the front wide camera's. The notebook adds the inference skill's recommended `per_view_negative_prompt`. The runtime adds camera labels and control emphasis.\n", + "\n", + "Every control is 1280 × 720 with 300 frames at 30 FPS. The next cell decodes every control, checks that their timelines are synchronized, and trims their first 201 frames without resizing into the output folder; controls that already have 201 frames are copied there unchanged. `prepare_request(..., backend=\"sglang\")` keeps each prepared control's absolute path in the JSON-encoded `extra_params` and returns no uploads; SGLang does not take `control_reference_index` fields." + ] + }, + { + "cell_type": "code", + "id": "auto-sglang-wsm-prepare", + "metadata": {}, + "source": [ + "WSM_REQUEST = AUTO_ROOT / \"assets\" / \"daytime_scene\" / \"transfer_auto_wsm.json\"\n", + "WSM_OUTPUT_DIR = OUTPUT_DIR / \"wsm\"\n", + "\n", + "wsm_manifest = resolve_request(WSM_REQUEST, SAMPLE_ROOT, caption_folder=CAPTION_FOLDER)\n", + "wsm_manifest[\"per_view_negative_prompt\"] = NEGATIVE_PROMPT_PATH.read_text().strip()\n", + "wsm_views = wsm_manifest[\"multiview\"][\"views\"]\n", + "assert tuple(view[\"camera_key\"] for view in wsm_views) == SAMPLE_CAMERAS\n", + "assert wsm_manifest[\"model_mode\"] == \"text2video\" and not any(\"vision_path\" in view for view in wsm_views)\n", + "wsm_controls = prepare_controls(\n", + " wsm_views, SAMPLE_ROOT, WSM_OUTPUT_DIR / \"prepared_controls\",\n", + " wsm_manifest[\"num_frames\"], wsm_manifest[\"fps\"], materialize=True,\n", + ")\n", + "wsm_request = copy.deepcopy(wsm_manifest)\n", + "for view, control in zip(wsm_request[\"multiview\"][\"views\"], wsm_controls, strict=True):\n", + " view[\"control_path\"] = str(control)\n", + "wsm_form, wsm_uploads = prepare_request(wsm_request, SAMPLE_ROOT, MODEL_ID, backend=\"sglang\")\n", + "assert not wsm_uploads # SGLang opens each control_path itself; nothing is uploaded.\n", + "assert all(Path(view[\"control_path\"]).is_absolute() for view in json.loads(wsm_form[\"extra_params\"])[\"multiview\"][\"views\"])\n", + "WSM_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n", + "(WSM_OUTPUT_DIR / \"request.json\").write_text(json.dumps(wsm_form, indent=2, ensure_ascii=False) + \"\\n\")\n", + "for camera_key, control in zip(SAMPLE_CAMERAS, wsm_controls):\n", + " info = probe_video(control)\n", + " print(f\"{camera_key:28} {info['frames']} frames, {info['width']}×{info['height']}, {float(info['fps']):g} FPS\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "auto-sglang-wsm-generate", + "metadata": {}, + "source": [ + "### Generate\n", + "\n", + "`submit_wsm(..., async_mode=True)` sends the form fields to `POST /v1/videos`, polls `GET /v1/videos/{id}` every ten seconds, and downloads the finished MP4 from `/content`. It saves the latest job status as `job.json`; a rejected submission or a failed job leaves `response_error.txt`, and a polling timeout leaves the job running on the server at the printed URL. The request carries the model, top-level caption, canvas `size`, `num_frames`, `fps`, `seed`, `num_inference_steps`, `guidance_scale`, and `flow_shift` from the request file, and in `extra_params` the `wsm: {}` hint, the `multiview.views` list with each camera's key, caption, and absolute `control_path`, the `per_view_negative_prompt`, and the `resolution` and `aspect_ratio`. The first job also compiles kernels, so it takes longer than later ones." + ] + }, + { + "cell_type": "code", + "id": "auto-sglang-wsm-submit", + "metadata": {}, + "source": [ + "with httpx.Client(base_url=SGLANG_BASE_URL, headers=HEADERS, timeout=HTTP_TIMEOUT) as client:\n", + " started = time.monotonic()\n", + " wsm_video = submit_wsm(client, wsm_form, wsm_uploads, WSM_OUTPUT_DIR, async_mode=True, timeout=GENERATION_TIMEOUT)\n", + " wsm_elapsed = round(time.monotonic() - started, 1)\n", + "print(f\"{wsm_elapsed:.1f}s {wsm_video}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "auto-sglang-wsm-validate-heading", + "metadata": {}, + "source": [ + "### Split and view the camera videos\n", + "\n", + "The result is one camera-major MP4: 201 frames of the first camera, then 201 of the second, through the seventh. `split_cameras` checks its `7 × 201 = 1407` frames, FPS, and canvas, then writes one labeled video per camera; `save_outputs` records them in `sample_outputs.json`. The folder then holds:\n", + "\n", + "```text\n", + "outputs/sglang/auto/wsm/\n", + " request.json # The exact form fields that were sent\n", + " job.json # The last job response\n", + " prepared_controls/ # The seven trimmed controls the server read\n", + " multiview_camera_major.mp4 # The downloaded result\n", + " vision_viewNN_.mp4 # One video per camera\n", + " sample_outputs.json # Camera order, frame count, FPS, and the per-camera files\n", + "```\n", + "\n", + "Inspect the grid for control adherence, cross-camera consistency, and visual quality, and see [Validate the rollout](auto.md#validate-the-rollout)." + ] + }, + { + "cell_type": "code", + "id": "auto-sglang-wsm-validate", + "metadata": {}, + "source": [ + "wsm_cameras = split_cameras(wsm_video, WSM_OUTPUT_DIR, SAMPLE_CAMERAS, wsm_manifest[\"num_frames\"], wsm_manifest[\"fps\"])\n", + "save_outputs(WSM_OUTPUT_DIR, wsm_video, wsm_cameras, wsm_manifest)\n", + "display(HTML(camera_grid(wsm_cameras)))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "auto-sglang-lidar-start", + "metadata": {}, + "source": [ + "## 5. Joint RGB and LiDAR example\n", + "\n", + "### Prepare the inputs\n", + "\n", + "The [joint request](assets/daytime_scene/transfer_auto_joint.json) uses the same server, cameras, and captions and asks for the same **201 frames per camera at 30 FPS**, 832 × 480, seed 0, 35 steps, guidance 3, flow shift 10, and control guidance 1. Each camera adds the sample's RGB recording under `vision_path`; with `condition_video_as_image: true`, only its first decoded frame conditions generation. The request also adds a numeric LiDAR control with `return_output: true`, which the SGLang request carries as `lidar.decode: true`.\n", + "\n", + "`resolve_request` converts the published `daytime_scene/lidar/control.tar` to a local `prepared/lidar_control.safetensors` under the output folder: FP32 `frames[3,100,128,1800]` with range in metres, unit intensity at valid returns, and validity. The 201-frame request consumes its first 67 sweeps at 10 Hz; the next cell checks them one at a time, and the server truncates the trailing sweeps and applies its own padding and normalization.\n", + "\n", + "The cell then decodes each RGB condition's first frame, trims the WSM controls to 201 frames, or copies them unchanged when they already have 201 frames, and copies each RGB condition into the output folder, because SGLang opens inputs by path and the sample lives in this notebook's Hugging Face cache. `prepare_request(..., backend=\"sglang\")` keeps every path absolute and returns no uploads." + ] + }, + { + "cell_type": "code", + "id": "auto-sglang-lidar-prepare", + "metadata": {}, + "source": [ + "LIDAR_REQUEST = AUTO_ROOT / \"assets\" / \"daytime_scene\" / \"transfer_auto_joint.json\"\n", + "LIDAR_OUTPUT_DIR = OUTPUT_DIR / \"lidar\"\n", + "LIDAR_FPS = 10.0\n", + "\n", + "lidar_manifest = resolve_request(\n", + " LIDAR_REQUEST, SAMPLE_ROOT, caption_folder=CAPTION_FOLDER,\n", + " lidar_output=LIDAR_OUTPUT_DIR / \"prepared\" / \"lidar_control.safetensors\",\n", + ")\n", + "lidar_manifest[\"per_view_negative_prompt\"] = NEGATIVE_PROMPT_PATH.read_text().strip()\n", + "lidar_views = lidar_manifest[\"multiview\"][\"views\"]\n", + "assert tuple(view[\"camera_key\"] for view in lidar_views) == SAMPLE_CAMERAS\n", + "assert lidar_manifest[\"multiview\"][\"condition_video_as_image\"] is True\n", + "assert lidar_manifest[\"lidar\"][\"return_output\"] is True\n", + "assert lidar_manifest[\"resolution\"] == \"480\" and lidar_manifest[\"aspect_ratio\"] == \"16,9\"\n", + "lidar_sweeps = round(lidar_manifest[\"num_frames\"] * LIDAR_FPS / lidar_manifest[\"fps\"])\n", + "\n", + "lidar_input = Path(lidar_manifest[\"lidar\"][\"control_path\"])\n", + "lidar_input_info = validate_lidar(lidar_input, lidar_sweeps)\n", + "validate_rgb_conditions(lidar_views, SAMPLE_ROOT)\n", + "lidar_controls = prepare_controls(\n", + " lidar_views, SAMPLE_ROOT, LIDAR_OUTPUT_DIR / \"prepared_controls\",\n", + " lidar_manifest[\"num_frames\"], lidar_manifest[\"fps\"], materialize=True,\n", + ")\n", + "lidar_conditions = stage_conditions(lidar_views, LIDAR_OUTPUT_DIR / \"staged_conditions\")\n", + "lidar_request = copy.deepcopy(lidar_manifest)\n", + "for view, control, condition in zip(lidar_request[\"multiview\"][\"views\"], lidar_controls, lidar_conditions, strict=True):\n", + " view[\"control_path\"] = str(control)\n", + " view[\"vision_path\"] = str(condition)\n", + "lidar_form, lidar_uploads = prepare_request(lidar_request, SAMPLE_ROOT, MODEL_ID, backend=\"sglang\")\n", + "lidar_extra = json.loads(lidar_form[\"extra_params\"])\n", + "assert not lidar_uploads\n", + "assert lidar_extra[\"lidar\"] == {\"control_path\": str(lidar_input.resolve()), \"decode\": True}\n", + "assert all(\n", + " \"control_path\" in view and \"vision_path\" in view and not any(key.endswith(\"_reference_index\") for key in view)\n", + " for view in lidar_extra[\"multiview\"][\"views\"]\n", + ")\n", + "LIDAR_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n", + "(LIDAR_OUTPUT_DIR / \"request.json\").write_text(json.dumps(lidar_form, indent=2, ensure_ascii=False) + \"\\n\")\n", + "rgb_conditions = dict(zip(SAMPLE_CAMERAS, lidar_conditions))\n", + "print(f\"LiDAR input: {lidar_input_info['shape']} {lidar_input_info['dtype']}; using {lidar_sweeps} sweeps at {LIDAR_FPS:g} Hz\")\n", + "print(f\"Prepared {len(lidar_controls)} WSM controls, {len(lidar_conditions)} RGB conditions, and the LiDAR control under {LIDAR_OUTPUT_DIR}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "auto-sglang-lidar-submit-heading", + "metadata": {}, + "source": [ + "### Submit the job and download both artifacts\n", + "\n", + "`submit_joint(..., backend=\"sglang\")` sends the form fields to `POST /v1/videos`, polls `GET /v1/videos/{id}` until the job completes, then downloads the camera-major MP4 from `/content` and the numeric LiDAR from `/content?variant=rangemap`, saving the job's LiDAR descriptor beside it as `multiview_camera_major.lidar.json`. The latest job status is saved as `job.json`; a rejected submission or a failed job leaves `response_error.txt`, and a polling timeout leaves the job on the server at the printed URL. Rerunning the cell submits a new job.\n", + "\n", + "SGLang also accepts `lidar.decode: false` for LiDAR-conditioned camera generation without numeric output; this walkthrough keeps decoding enabled so the next cell can inspect the generated tensor." + ] + }, + { + "cell_type": "code", + "id": "auto-sglang-lidar-submit", + "metadata": {}, + "source": [ + "with httpx.Client(base_url=SGLANG_BASE_URL, headers=HEADERS, timeout=HTTP_TIMEOUT) as client:\n", + " started = time.monotonic()\n", + " lidar_result = submit_joint(client, lidar_form, lidar_uploads, LIDAR_OUTPUT_DIR, timeout=GENERATION_TIMEOUT, backend=\"sglang\")\n", + " lidar_elapsed = round(time.monotonic() - started, 1)\n", + "print(f\"Job {lidar_result['job_id']} completed in {lidar_elapsed:.1f}s\")\n", + "print(\"Camera video:\", lidar_result[\"video\"])\n", + "print(\"Generated LiDAR:\", lidar_result[\"lidar\"])" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "auto-sglang-lidar-validate-heading", + "metadata": {}, + "source": [ + "### Validate generated outputs and display cameras\n", + "\n", + "The camera-major MP4 must contain `7 × 201 = 1407` frames. LiDAR uses its own 10 Hz timeline: 67 sweeps. SGLang writes three tensors of shape `[67,128,1800]`: FP32 `range_m`, FP16 `intensity`, and uint8 `validity`, with the layout, FPS, and range limits in the file's metadata. This differs from the input's single `frames` tensor.\n", + "\n", + "`validate_sglang_lidar` checks the tensor shapes, dtypes, physical values, and metadata; the cell then compares them with the job descriptor, splits the camera videos, writes `sample_outputs.json`, and displays the camera grid. Valid shapes and values do not establish that the LiDAR agrees with the HD map or the cameras; inspect the generated sweeps before relying on them. The folder then holds:\n", + "\n", + "```text\n", + "outputs/sglang/auto/lidar/\n", + " request.json # The exact form fields that were sent\n", + " job.json # The last job response, including the LiDAR descriptor\n", + " prepared/lidar_control.safetensors # The converted LiDAR control the server read\n", + " prepared_controls/ # The seven trimmed WSM controls\n", + " staged_conditions/ # The seven RGB conditions, copied from the sample\n", + " multiview_camera_major.mp4 # The camera-major video from /content\n", + " vision_viewNN_.mp4 # One video per camera\n", + " multiview_camera_major.lidar.safetensors # The numeric LiDAR from /content?variant=rangemap\n", + " multiview_camera_major.lidar.json # Its descriptor\n", + " sample_outputs.json # Camera order, files, job ID, and the LiDAR descriptor\n", + "```" + ] + }, + { + "cell_type": "code", + "id": "auto-sglang-lidar-validate", + "metadata": {}, + "source": [ + "lidar_output_info = validate_sglang_lidar(lidar_result[\"lidar\"], lidar_sweeps, fps=LIDAR_FPS)\n", + "lidar_descriptor = lidar_result[\"descriptor\"]\n", + "assert lidar_descriptor[\"sweeps\"] == lidar_sweeps\n", + "assert lidar_descriptor[\"fps\"] == LIDAR_FPS\n", + "assert [lidar_descriptor[\"sweeps\"], lidar_descriptor[\"height\"], lidar_descriptor[\"width\"]] == lidar_output_info[\"shape\"]\n", + "for field in (\"fps\", \"min_range_m\", \"max_range_m\"):\n", + " assert float(lidar_output_info[\"metadata\"][field]) == float(lidar_descriptor[field]), field\n", + "\n", + "lidar_cameras = split_cameras(lidar_result[\"video\"], LIDAR_OUTPUT_DIR, SAMPLE_CAMERAS, lidar_manifest[\"num_frames\"], lidar_manifest[\"fps\"])\n", + "save_outputs(LIDAR_OUTPUT_DIR, lidar_result[\"video\"], lidar_cameras, lidar_manifest, lidar=lidar_result)\n", + "print(f\"LiDAR: {lidar_output_info['shape']} at {lidar_descriptor['fps']:g} Hz; split into {len(lidar_cameras)} cameras\")\n", + "display(HTML(camera_grid(lidar_cameras)))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "auto-sglang-next-steps", + "metadata": {}, + "source": [ + "## 6. Next steps\n", + "\n", + "- **Use your own scene:** copy a request file from [`assets/daytime_scene/`](assets/daytime_scene), point its paths at your data, then set `WSM_REQUEST` or `LIDAR_REQUEST` to the copy and `SAMPLE_ROOT` to the folder its paths are relative to. List any of the checkpoint's eleven cameras in its camera order, `CAMERA_ORDER` in [`helpers/auto_helpers.py`](helpers/auto_helpers.py), each with a WSM control and a caption file or inline `prompt`. For the joint example, add a `vision_path` per camera and a numeric LiDAR control; [`helpers/lidar_rangemap.py`](helpers/lidar_rangemap.py) converts a sparse HD-map range-map archive. Every prepared input lands under the output folder, so keep it inside the mounted cookbook root.\n", + "- **Validate further:** apply the [shared Simulation checklist](../README.md#validate-a-rollout) and the [Transfer-Auto checks](auto.md#validate-the-rollout).\n", + "- **Troubleshooting:**\n", + " - *Missing inputs:* rerun the sample download in step 3 and the example's **Prepare the inputs** cell; `AUTO_SAMPLE_ROOT` must point at a complete `sample_data` folder.\n", + " - *HTTP error on submission:* inspect `response_error.txt` and the server log. A path the server cannot open means the checkout or the output folder is not mounted at the same absolute path inside the container.\n", + " - *Job `failed`:* `job.json` and `response_error.txt` carry the server's error message; the server log has the traceback.\n", + " - *Timeout:* the job keeps running on the server at the printed URL; `GET /v1/videos/{id}` reports its status. Raise `GENERATION_TIMEOUT` before resubmitting.\n", + "- **Other runtimes:** the [vLLM-Omni notebook](run_auto_with_vllm_omni.ipynb) runs the same two examples with uploads, and the [Cosmos Framework notebook](run_auto_with_cosmos_framework.ipynb) runs them natively from the published JSONL. For the autoregressive student of this model, see the [Sim-Auto SGLang notebook](run_sim-auto_with_sglang.ipynb)." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_vllm_omni.ipynb b/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_vllm_omni.ipynb index ad948b2fa..9fb552b82 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_vllm_omni.ipynb +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_vllm_omni.ipynb @@ -1062,7 +1062,7 @@ "\n", "- **Use your own scene:** copy a request file from [`assets/daytime_scene/`](assets/daytime_scene), point its paths at your data, then set `WSM_REQUEST` or `LIDAR_REQUEST` to the copy and `SAMPLE_ROOT` to the folder its paths are relative to. List any of the checkpoint's eleven cameras in its camera order, `CAMERA_ORDER` in [`helpers/auto_helpers.py`](helpers/auto_helpers.py), each with a WSM control and a caption file or inline `prompt`. The [WSM generation notebook](generate_wsm.ipynb) writes a native request for the Framework notebooks; serving uses the templates above. For the joint example, add a `vision_path` per camera and a numeric LiDAR control; [`helpers/lidar_rangemap.py`](helpers/lidar_rangemap.py) converts a sparse HD-map range-map archive.\n", "- **Validate further:** apply the [shared Simulation checklist](../README.md#validate-a-rollout) and the [Transfer-Auto checks](auto.md#validate-the-rollout). If a request fails, see [vLLM-Omni troubleshooting](auto.md#vllm-omni-troubleshooting).\n", - "- **Other runtimes:** run the same media with the published native JSONL in the [Cosmos Framework notebook](run_auto_with_cosmos_framework.ipynb). For the autoregressive student of this model, see the [Sim-Auto notebook](run_sim-auto_with_cosmos_framework.ipynb).\n", + "- **Other runtimes:** run the same media with the published native JSONL in the [Cosmos Framework notebook](run_auto_with_cosmos_framework.ipynb), or serve them through SGLang with the [SGLang notebook](run_auto_with_sglang.ipynb). For the autoregressive student of this model, see the [Sim-Auto notebook](run_sim-auto_with_cosmos_framework.ipynb).\n", "\n", "For custom caption filenames, pass `caption_folder=None` to the request helper to preserve your paths.\n" ] diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/run_sim-auto_with_sglang.ipynb b/.agents/skills/cookbooks/cosmos3/simulation/auto/run_sim-auto_with_sglang.ipynb new file mode 100644 index 000000000..9c4941020 --- /dev/null +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/run_sim-auto_with_sglang.ipynb @@ -0,0 +1,763 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "sim-auto-sglang-00", + "metadata": {}, + "source": [ + "" + ] + }, + { + "cell_type": "markdown", + "id": "sim-auto-sglang-01", + "metadata": {}, + "source": [ + "# Cosmos3-Nano-Sim-Auto with SGLang\n", + "\n", + "Cosmos3-Nano-Sim-Auto generates **synchronized multi-camera driving video** from one World Scenario Model (WSM) control video and one RGB first-frame condition per camera. The WSM controls render the scene's lane lines, road surface, and vehicle boxes from each camera's viewpoint; the RGB frame sets each camera's starting appearance; a caption per camera describes the scene. It is the causal, distilled student of Cosmos3-Nano-Transfer-Auto: it generates every camera in 0.4-second chunks with a rolling key/value cache and a fixed **four-step** sampler at guidance 1.\n", + "\n", + "This notebook runs the RGB example of the shared **daytime sample**, one scene recorded by seven cameras, through one SGLang server at **297 frames per camera, 30 FPS, 832 × 480, and seed 0** on one GPU. SGLang reads the inputs from paths the server can open and runs the rollout as one asynchronous Videos API job.\n", + "\n", + "| Component | Version |\n", + "| --- | --- |\n", + "| SGLang | The published [`lmsysorg/sglang:dev`](https://hub.docker.com/r/lmsysorg/sglang/tags) image, which carries the Cosmos3 Sim-Auto pipeline from [sgl-project/sglang#39475](https://github.com/sgl-project/sglang/pull/39475); the tag moves, so pull it before starting the server |\n", + "| `nvidia/Cosmos3-Nano-Sim-Auto` | [`main`](https://huggingface.co/nvidia/Cosmos3-Nano-Sim-Auto/tree/main) |\n", + "| Daytime sample | [`main`](https://huggingface.co/nvidia/Cosmos3-Nano-Transfer-Auto/tree/main/assets/sample_data) of `nvidia/Cosmos3-Nano-Transfer-Auto`, under `assets/sample_data` |\n", + "\n", + "Steps 1 and 2 run on the server host. Steps 3 to 6 run in this notebook." + ] + }, + { + "cell_type": "markdown", + "id": "sim-auto-sglang-02", + "metadata": {}, + "source": [ + "## 1. Prerequisites\n", + "\n", + "**Server host:** Linux with Docker, the NVIDIA Container Toolkit, and one NVIDIA GPU with enough memory for the model. The SGLang rollout runs on FlexAttention's Triton kernels, so it does not need the FlashAttention-4 kernels that the [Framework notebook](run_sim-auto_with_cosmos_framework.ipynb) requires. The host driver must support the CUDA runtime of the `lmsysorg/sglang:dev` image. `HF_TOKEN` needs read access to `nvidia/Cosmos3-Nano-Sim-Auto`. The server opens the inputs by path, so the cookbook checkout, including the output folder this notebook writes, must be on the server host or on a filesystem it shares with this notebook, at one absolute path.\n", + "\n", + "**Notebook client:** Python 3.10 or newer, and a Hugging Face token with read access to `nvidia/Cosmos3-Nano-Transfer-Auto`, which publishes the sample; run `hf auth login` or set `HF_TOKEN`. In the notebook kernel environment, install:\n", + "\n", + "```bash\n", + "pip install httpx pillow numpy matplotlib ipython huggingface_hub\n", + "```\n", + "\n", + "Install FFmpeg on the client host; the package also includes FFprobe. On Ubuntu/Debian:\n", + "\n", + "```bash\n", + "sudo apt-get update\n", + "sudo apt-get install -y ffmpeg\n", + "```\n", + "\n", + "On macOS with Homebrew, use `brew install ffmpeg`. Verify that both `ffmpeg -version` and `ffprobe -version` work.\n", + "\n", + "**Inputs:** The example uses the shared daytime sample published with Transfer-Auto under `assets/sample_data`: one scene with seven synchronized cameras, each with an RGB recording, a WSM control, and a caption. Step 3 downloads the files this example needs, about 100 MiB, into the Hugging Face cache and checks them against the sample's `SHA256SUMS`. To use a copy you already have, set `AUTO_SAMPLE_ROOT` to its `sample_data` folder.\n", + "\n", + "Nothing is uploaded. Step 4 writes the trimmed WSM controls and a copy of each RGB condition under the output folder, inside the mounted cookbook root, and the request names those files by absolute path." + ] + }, + { + "cell_type": "markdown", + "id": "sim-auto-sglang-03", + "metadata": {}, + "source": [ + "## 2. Start the server\n", + "\n", + "The server runs in the published SGLang image, `lmsysorg/sglang:dev`, which includes the Cosmos3 Sim-Auto pipeline. On the server host, set `COSMOS_ROOT` to the absolute cookbook root (`.agents/skills`), which step 3 prints as `Cookbook root`. `SGLANG_HF_CACHE` is a host directory for the Hugging Face cache. The container runs as root by default, so the setup lines create that directory world-writable; alternatively, start the container with your host user ID. Pull the image, then run the server in a terminal and leave it running while using the notebook:\n", + "\n", + "```bash\n", + "docker pull lmsysorg/sglang:dev\n", + "export COSMOS_ROOT=/absolute/path/to/cosmos-framework/.agents/skills\n", + "export SGLANG_HF_CACHE=\"${SGLANG_HF_CACHE:-$HOME/.cache/sglang-huggingface}\"\n", + "mkdir -p \"$SGLANG_HF_CACHE\"\n", + "chmod 777 \"$SGLANG_HF_CACHE\"\n", + "docker run --init --runtime nvidia --gpus all \\\n", + " -v $SGLANG_HF_CACHE:/root/.cache/huggingface \\\n", + " -v \"${COSMOS_ROOT}:${COSMOS_ROOT}\" \\\n", + " -e HF_TOKEN=$HF_TOKEN \\\n", + " -p 30000:30000 \\\n", + " --ipc=host \\\n", + " lmsysorg/sglang:dev \\\n", + " sglang serve \\\n", + " --model-path nvidia/Cosmos3-Nano-Sim-Auto \\\n", + " --revision main \\\n", + " --host 0.0.0.0 --port 30000\n", + "```\n", + "\n", + "| Setting | Effect |\n", + "| --- | --- |\n", + "| `-v \"${COSMOS_ROOT}:${COSMOS_ROOT}\"` | Mounts the cookbook root at its host path. The request names every WSM control and RGB condition by absolute path, and the trimmed controls and staged RGB conditions that step 4 writes under the output folder must be visible at the same path. |\n", + "| `-v $SGLANG_HF_CACHE:/root/.cache/huggingface` | Keeps the downloaded checkpoint across container restarts. The setup lines default the directory to `~/.cache/sglang-huggingface`, create it, and make it world-writable. |\n", + "| `-e HF_TOKEN=$HF_TOKEN` | Read access to the gated checkpoint. |\n", + "| `--revision` | Loads the checkpoint from `main`. The first start downloads it into the mounted cache. |\n", + "| `docker pull` | The `dev` tag moves with SGLang's main branch; pulling before each start picks up the current image. |\n", + "| `--gpus all` | Runs on one GPU; select one with `--gpus '\"device=0\"'`. The Sim-Auto pipeline rejects sequence parallelism; `--tp-size` shards the transformer across GPUs if the model does not fit on one. |\n", + "| `--host 0.0.0.0` | Accepts connections from outside the container. For a remote notebook, set `SGLANG_BASE_URL`, or forward the port with `ssh -L 30000:127.0.0.1:30000 `. |\n", + "\n", + "The HTTP port opens only after the weights are loaded and the warmup request finishes, so wait for the `Uvicorn running` line. The start-up log also prints `Cosmos3-Nano-Sim-Auto pipeline stages created`, followed by the checkpoint's chunk, ring, and sampler settings. This pipeline runs no text or video guardrails. Generation runs as an asynchronous job, so no server-side request timeout needs raising." + ] + }, + { + "cell_type": "markdown", + "id": "sim-auto-sglang-04", + "metadata": {}, + "source": [ + "## 3. Set up the client\n", + "\n", + "Run the notebook from anywhere inside `.agents/skills`, or set `COSMOS_COOKBOOK_ROOT` to it. These optional environment variables override the defaults:\n", + "\n", + "| Variable | Default | Purpose |\n", + "| --- | --- | --- |\n", + "| `SGLANG_BASE_URL` | `http://127.0.0.1:30000` | Server address |\n", + "| `SGLANG_API_KEY` | unset | Bearer token, if the server was started with `--api-key` |\n", + "| `COSMOS3_OUTPUT_DIR` | `.agents/skills/cookbooks/cosmos3/simulation/outputs/sglang/sim_auto` | Output folder. Keep it inside the mounted cookbook root: the server reads the trimmed controls and staged RGB conditions from it. |\n", + "| `AUTO_SAMPLE_ROOT` | unset | A local `sample_data` folder to use instead of downloading the sample |\n", + "\n", + "The next cell sets the configuration and imports the Sim-Auto input helpers from [`helpers/sim_auto_sglang_helpers.py`](helpers/sim_auto_sglang_helpers.py) and the sample helpers from [`helpers/daytime_sample.py`](helpers/daytime_sample.py). The cell after it defines the client helpers for server readiness, Videos API jobs, and output validation; neither needs GPU libraries. Then fetch the sample." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sim-auto-sglang-05", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import copy\n", + "import json\n", + "import os\n", + "import sys\n", + "import time\n", + "\n", + "import httpx\n", + "import matplotlib.pyplot as plt\n", + "from IPython.display import HTML, display\n", + "\n", + "\n", + "def find_auto_root() -> Path:\n", + " configured = os.environ.get(\"COSMOS_COOKBOOK_ROOT\")\n", + " start = Path(configured).expanduser().resolve() if configured else Path.cwd().resolve()\n", + " for root in [start] if configured else [start, *start.parents]:\n", + " for candidate in (root, root / \"auto\", root / \"cookbooks\" / \"cosmos3\" / \"simulation\" / \"auto\", root / \".agents\" / \"skills\" / \"cookbooks\" / \"cosmos3\" / \"simulation\" / \"auto\"):\n", + " if (candidate / \"helpers\" / \"sim_auto_sglang_helpers.py\").is_file():\n", + " return candidate\n", + " raise FileNotFoundError(\"Run inside .agents/skills or set COSMOS_COOKBOOK_ROOT to it.\")\n", + "\n", + "\n", + "AUTO_ROOT = find_auto_root()\n", + "SIMULATION_ROOT = AUTO_ROOT.parent\n", + "COOKBOOK_ROOT = SIMULATION_ROOT.parents[2] # The .agents/skills folder the server mounts.\n", + "if str(AUTO_ROOT) not in sys.path:\n", + " sys.path.insert(0, str(AUTO_ROOT))\n", + "\n", + "from helpers.auto_helpers import prepare_controls, split_cameras, stage_conditions, validate_rgb_conditions\n", + "from helpers.daytime_sample import SAMPLE_REPO, SAMPLE_REVISION, fetch_sample\n", + "from helpers.sim_auto_sglang_helpers import CAMERA_ORDER, MODEL_ID, load_request, prepare_request, require_inputs\n", + "\n", + "MODEL_REVISION = \"main\"\n", + "SGLANG_BASE_URL = os.environ.get(\"SGLANG_BASE_URL\", \"http://127.0.0.1:30000\").rstrip(\"/\")\n", + "SGLANG_API_KEY = os.environ.get(\"SGLANG_API_KEY\")\n", + "HEADERS = {\"Authorization\": f\"Bearer {SGLANG_API_KEY}\"} if SGLANG_API_KEY else {}\n", + "HTTP_TIMEOUT = 600 # Per HTTP request: submission, each poll, and the MP4 download.\n", + "GENERATION_TIMEOUT = 3600 # Job wait, from submission to completion.\n", + "OUTPUT_DIR = Path(\n", + " os.environ.get(\"COSMOS3_OUTPUT_DIR\", SIMULATION_ROOT / \"outputs\" / \"sglang\" / \"sim_auto\")\n", + ").expanduser().resolve()\n", + "AUTO_SAMPLE_ROOT = os.environ.get(\"AUTO_SAMPLE_ROOT\") # Unset: download the sample from Hugging Face\n", + "CAPTION_FOLDER = \"captions_with_timestamp\" # Alternative: \"captions\"\n", + "\n", + "OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n", + "print(\"Cookbook root:\", COOKBOOK_ROOT)\n", + "print(\"Server:\", SGLANG_BASE_URL, \"(with API key)\" if SGLANG_API_KEY else \"\")\n", + "print(\"Model:\", MODEL_ID, \"@\", MODEL_REVISION)\n", + "print(\"Sample:\", AUTO_SAMPLE_ROOT or f\"{SAMPLE_REPO} @ {SAMPLE_REVISION}\")\n", + "print(\"Outputs:\", OUTPUT_DIR)\n", + "if not OUTPUT_DIR.is_relative_to(COOKBOOK_ROOT):\n", + " print(f\"Note: {OUTPUT_DIR} is outside the cookbook root; the server must see it at the same path.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sim-auto-sglang-06", + "metadata": {}, + "outputs": [], + "source": [ + "# Client helpers. The same cell appears in each Simulation serving notebook; keep the copies in sync.\n", + "import base64\n", + "import html\n", + "import json\n", + "import shutil\n", + "import subprocess\n", + "import tempfile\n", + "import time\n", + "from fractions import Fraction\n", + "from io import BytesIO\n", + "from pathlib import Path\n", + "\n", + "import httpx\n", + "from PIL import Image\n", + "\n", + "_LFS_POINTER = b\"version https://git-lfs.github.com/spec/v1\"\n", + "\n", + "\n", + "def require_media_tools() -> None:\n", + " missing = [name for name in (\"ffmpeg\", \"ffprobe\") if not shutil.which(name)]\n", + " if missing:\n", + " raise RuntimeError(f\"Install FFmpeg and add {' and '.join(missing)} to PATH.\")\n", + "\n", + "\n", + "def require_file(path: Path) -> Path:\n", + " \"\"\"Resolve an input file and reject missing, empty, or unhydrated Git LFS files.\"\"\"\n", + " path = Path(path).expanduser().resolve()\n", + " if not path.is_file() or path.stat().st_size == 0:\n", + " raise FileNotFoundError(f\"Missing or empty input: {path}\")\n", + " with path.open(\"rb\") as stream:\n", + " if stream.read(len(_LFS_POINTER)) == _LFS_POINTER:\n", + " raise ValueError(f\"Unhydrated Git LFS pointer: {path}. Run git lfs pull in the cookbook checkout.\")\n", + " return path\n", + "\n", + "\n", + "def wait_for_model(client: httpx.Client, model: str, *, timeout: float = 1800, interval: float = 5) -> list[str]:\n", + " \"\"\"Poll ``GET /v1/models`` until the server answers, then require ``model``.\n", + "\n", + " Connection errors and 5xx responses mean the server is still starting and are\n", + " retried until ``timeout``. Other HTTP errors, such as a missing API key, and a\n", + " server that answers without ``model`` fail immediately.\n", + " \"\"\"\n", + " deadline = time.monotonic() + timeout\n", + " waiting = False\n", + " while True:\n", + " try:\n", + " response = client.get(\"/v1/models\", timeout=min(30, timeout))\n", + " if response.status_code < 500:\n", + " break\n", + " reason = f\"HTTP {response.status_code}\"\n", + " except httpx.TransportError as error:\n", + " reason = type(error).__name__\n", + " if time.monotonic() + interval > deadline:\n", + " raise TimeoutError(f\"No ready server at {client.base_url} after {timeout:.0f}s ({reason}).\")\n", + " if not waiting:\n", + " print(f\"Waiting for the server at {client.base_url} ({reason}) ...\")\n", + " waiting = True\n", + " time.sleep(interval)\n", + " response.raise_for_status()\n", + " served = [item[\"id\"] for item in response.json().get(\"data\", []) if item.get(\"id\")]\n", + " if model not in served:\n", + " raise ValueError(f\"Expected served model {model!r}; the server at {client.base_url} serves {served}.\")\n", + " return served\n", + "\n", + "\n", + "def submit_video_job(client: httpx.Client, form: dict, output_dir: Path, *, timeout: float, interval: float = 10) -> dict:\n", + " \"\"\"Send ``POST /v1/videos`` as multipart form fields and poll the job until it completes.\n", + "\n", + " The form is saved as ``request.json`` beside the outputs so the request can be\n", + " inspected or replayed, and every job response as ``job.json``. A rejected\n", + " submission or a failed job saves the server's response as ``response_error.txt``\n", + " and raises. A polling timeout raises too, but leaves the job running on the\n", + " server at the printed URL.\n", + " \"\"\"\n", + " output_dir = Path(output_dir)\n", + " output_dir.mkdir(parents=True, exist_ok=True)\n", + " (output_dir / \"request.json\").write_text(json.dumps(form, indent=2, ensure_ascii=False) + \"\\n\")\n", + " error_path = output_dir / \"response_error.txt\"\n", + " error_path.unlink(missing_ok=True)\n", + " job_path = output_dir / \"job.json\"\n", + " job_path.unlink(missing_ok=True)\n", + " # A None filename makes each part a regular field; the Videos API parses multipart, not URL-encoded, forms.\n", + " response = client.post(\"/v1/videos\", files=[(key, (None, value)) for key, value in form.items()])\n", + " if response.is_error:\n", + " error_path.write_text(response.text)\n", + " raise RuntimeError(f\"HTTP {response.status_code} from /v1/videos; see {error_path}\\n{response.text[:2000]}\")\n", + " job = response.json()\n", + " job_url = f\"/v1/videos/{job['id']}\"\n", + " deadline = time.monotonic() + timeout\n", + " print(f\"Submitted job {job['id']} at {client.base_url.join(job_url)}\")\n", + " while True:\n", + " job_path.write_text(json.dumps(job, indent=2) + \"\\n\")\n", + " status = job.get(\"status\")\n", + " if status == \"completed\":\n", + " return job\n", + " if status in {\"failed\", \"cancelled\"}:\n", + " error_path.write_text(json.dumps(job, indent=2) + \"\\n\")\n", + " raise RuntimeError(f\"Video generation {status}: {job.get('error')}; see {error_path}\")\n", + " if status not in {\"queued\", \"in_progress\"}:\n", + " raise RuntimeError(f\"Unexpected job status {status!r}; see {job_path}\")\n", + " remaining = deadline - time.monotonic()\n", + " if remaining <= 0:\n", + " raise TimeoutError(\n", + " f\"No result after {timeout:.0f}s; the job remains at {client.base_url.join(job_url)}. See {job_path}.\"\n", + " )\n", + " time.sleep(min(interval, remaining))\n", + " response = client.get(job_url)\n", + " if response.is_error:\n", + " error_path.write_text(response.text)\n", + " raise RuntimeError(f\"HTTP {response.status_code} from {job_url}; see {error_path}\\n{response.text[:2000]}\")\n", + " job = response.json()\n", + "\n", + "\n", + "def download_video(client: httpx.Client, job: dict, output_path: Path) -> Path:\n", + " \"\"\"Save a completed job's ``GET /v1/videos/{id}/content`` MP4 atomically.\"\"\"\n", + " output_path = Path(output_path)\n", + " temporary = output_path.with_name(output_path.name + \".part\")\n", + " try:\n", + " with client.stream(\"GET\", f\"/v1/videos/{job['id']}/content\", headers={\"Accept\": \"video/mp4\"}) as response:\n", + " if response.is_error:\n", + " response.read()\n", + " raise RuntimeError(f\"HTTP {response.status_code} from the content download\\n{response.text[:2000]}\")\n", + " content_type = response.headers.get(\"content-type\", \"\")\n", + " if not content_type.startswith(\"video/mp4\"):\n", + " raise ValueError(f\"Expected a video/mp4 response; got {content_type!r}.\")\n", + " with temporary.open(\"wb\") as stream:\n", + " for chunk in response.iter_bytes():\n", + " stream.write(chunk)\n", + " if temporary.stat().st_size == 0:\n", + " raise ValueError(\"The content download was empty.\")\n", + " temporary.replace(output_path)\n", + " finally:\n", + " temporary.unlink(missing_ok=True)\n", + " return output_path\n", + "\n", + "\n", + "def probe_video(path: Path) -> dict:\n", + " \"\"\"Return the first video stream's size, average FPS, and decoded frame count.\"\"\"\n", + " result = subprocess.run(\n", + " [\n", + " \"ffprobe\", \"-v\", \"error\", \"-count_frames\", \"-select_streams\", \"v:0\",\n", + " \"-show_entries\", \"stream=width,height,avg_frame_rate,nb_read_frames\", \"-of\", \"json\", str(path),\n", + " ],\n", + " check=True, capture_output=True, text=True,\n", + " )\n", + " streams = json.loads(result.stdout).get(\"streams\", [])\n", + " if not streams:\n", + " raise ValueError(f\"No video stream in {path}\")\n", + " stream = streams[0]\n", + " return {\n", + " \"width\": int(stream[\"width\"]),\n", + " \"height\": int(stream[\"height\"]),\n", + " \"fps\": Fraction(stream[\"avg_frame_rate\"]),\n", + " \"frames\": int(stream[\"nb_read_frames\"]),\n", + " }\n", + "\n", + "\n", + "def validate_video(path: Path, *, num_frames: int, fps: float, width: int, height: int) -> dict:\n", + " \"\"\"Require the requested frame count, FPS, and size, and an error-free full decode.\"\"\"\n", + " info = probe_video(path)\n", + " problems = []\n", + " if info[\"frames\"] != num_frames:\n", + " problems.append(f\"{info['frames']} frames, expected {num_frames}\")\n", + " if info[\"fps\"] != Fraction(fps):\n", + " problems.append(f\"{float(info['fps']):g} FPS, expected {fps:g}\")\n", + " if (info[\"width\"], info[\"height\"]) != (width, height):\n", + " problems.append(f\"{info['width']}x{info['height']}, expected {width}x{height}\")\n", + " decode = subprocess.run(\n", + " [\"ffmpeg\", \"-nostdin\", \"-v\", \"error\", \"-xerror\", \"-i\", str(path), \"-f\", \"null\", \"-\"],\n", + " capture_output=True, text=True,\n", + " )\n", + " if decode.returncode or decode.stderr.strip():\n", + " problems.append(f\"decode failed: {decode.stderr.strip()[:500]}\")\n", + " if problems:\n", + " raise ValueError(f\"{path}: \" + \"; \".join(problems))\n", + " return {**info, \"fps\": float(info[\"fps\"])}\n", + "\n", + "\n", + "def read_frame(path: Path, index: int) -> Image.Image:\n", + " \"\"\"Decode one zero-based frame as RGB.\"\"\"\n", + " result = subprocess.run(\n", + " [\n", + " \"ffmpeg\", \"-nostdin\", \"-v\", \"error\", \"-i\", str(path),\n", + " \"-vf\", rf\"select=eq(n\\,{index})\", \"-frames:v\", \"1\", \"-f\", \"image2pipe\", \"-vcodec\", \"png\", \"-\",\n", + " ],\n", + " check=True, capture_output=True,\n", + " )\n", + " if not result.stdout:\n", + " raise ValueError(f\"{path} has no frame {index}\")\n", + " return Image.open(BytesIO(result.stdout)).convert(\"RGB\")\n", + "\n", + "\n", + "def preview_video(path: Path, *, width: int = 416, frames: int | None = None) -> bytes:\n", + " \"\"\"Re-encode a small H.264 preview, optionally trimmed to ``frames``, for embedding in a notebook.\"\"\"\n", + " with tempfile.TemporaryDirectory() as directory:\n", + " preview = Path(directory) / \"preview.mp4\"\n", + " limit = [\"-frames:v\", str(frames)] if frames else []\n", + " subprocess.run(\n", + " [\n", + " \"ffmpeg\", \"-nostdin\", \"-v\", \"error\", \"-i\", str(path), \"-an\", *limit,\n", + " \"-vf\", f\"scale={width}:-2\", \"-c:v\", \"libx264\", \"-crf\", \"28\", \"-preset\", \"fast\",\n", + " \"-pix_fmt\", \"yuv420p\", \"-movflags\", \"+faststart\", str(preview),\n", + " ],\n", + " check=True, capture_output=True,\n", + " )\n", + " return preview.read_bytes()\n", + "\n", + "\n", + "def video_row(videos: dict[str, Path], *, width: int = 416, frames: int | None = None) -> str:\n", + " \"\"\"HTML that embeds labeled previews side by side; they autoplay muted and loop together.\n", + "\n", + " Each video is re-encoded at ``width`` and trimmed to ``frames``, so a long or\n", + " high-resolution input stays small and loops in step with the others.\n", + " \"\"\"\n", + " cells = []\n", + " for label, path in videos.items():\n", + " data = base64.b64encode(preview_video(path, width=width, frames=frames)).decode(\"ascii\")\n", + " cells.append(\n", + " f'
{html.escape(label)}
'\n", + " f'
'\n", + " )\n", + " return '
' + \"\".join(cells) + \"
\"\n", + "\n", + "\n", + "def write_sample_outputs(output_dir: Path, samples: list[dict]) -> Path:\n", + " \"\"\"Write ``sample_outputs.json``: one record per generated sample, each with a ``status``.\"\"\"\n", + " for sample in samples:\n", + " if \"status\" not in sample:\n", + " raise ValueError(\"Every sample record needs a status.\")\n", + " path = Path(output_dir) / \"sample_outputs.json\"\n", + " path.write_text(json.dumps(samples, indent=2, default=str) + \"\\n\")\n", + " return path\n", + "\n", + "\n", + "require_media_tools()" + ] + }, + { + "cell_type": "markdown", + "id": "sim-auto-sglang-sample-heading", + "metadata": {}, + "source": [ + "### Download and verify the sample\n", + "\n", + "Download the seven cameras' RGB recordings, WSM controls, and both caption styles from `assets/sample_data` on Hugging Face, and verify them against `SHA256SUMS`; later runs reuse the Hugging Face cache. No model weights are downloaded here.\n", + "\n", + "`CAPTION_FOLDER=\"captions_with_timestamp\"` selects timestamped captions, which better match training and empirically improve inference. Set it to `\"captions\"` for the original dense paragraphs." + ] + }, + { + "cell_type": "code", + "id": "sim-auto-sglang-sample", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "SAMPLE_ROOT = fetch_sample(AUTO_SAMPLE_ROOT)\n", + "print(\"Verified sample:\", SAMPLE_ROOT)\n", + "print(\"Cameras:\", \", \".join(CAMERA_ORDER))" + ] + }, + { + "cell_type": "markdown", + "id": "sim-auto-sglang-07", + "metadata": {}, + "source": [ + "### Verify the server\n", + "\n", + "Wait for the server to finish starting, then confirm that it serves the Sim-Auto model. The cell polls for up to 30 minutes, so you can run it while the server is still loading. `/v1/models` reports the model name, not its revision; the `--revision` in the server command pins that." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sim-auto-sglang-08", + "metadata": {}, + "outputs": [], + "source": [ + "with httpx.Client(base_url=SGLANG_BASE_URL, headers=HEADERS) as client:\n", + " served = wait_for_model(client, MODEL_ID, timeout=1800)\n", + "print(\"Server ready; serving\", served)" + ] + }, + { + "cell_type": "markdown", + "id": "sim-auto-sglang-09", + "metadata": {}, + "source": [ + "## 4. Prepare the inputs\n", + "\n", + "The [RGB request](assets/daytime_scene/sim_auto.json) asks for **297 frames per camera at 30 FPS** with `resolution: \"480\"`, seed **0**, and `condition_video_as_image: true`. It lists the sample's seven cameras in the checkpoint's camera order, each with its WSM control, RGB recording, and caption file; every file is named after its camera, such as `daytime_scene/wsm/camera_front_wide_120fov.mp4`. The server identifies each camera by its `camera_key`, so the seven keep their trained camera IDs 0 to 6; the generated video follows the request order. The checkpoint fixes guidance 1 and four steps.\n", + "\n", + "`load_request` resolves the paths under the sample and replaces each caption file with its text, which covers all 300 source frames. The Videos API also takes a top-level caption; the request copies the front wide camera's.\n", + "\n", + "Every source video has 300 frames at 30 FPS and 1280 × 720. The server resizes and center-crops inputs to its 832 × 480 output.\n", + "\n", + "The next cell checks the request and reports every missing file at once. It decodes each RGB condition's first frame, which is all that generation uses. It fully decodes the WSM controls, checks that their timelines are synchronized, and trims their first 297 frames without resizing into the output folder; controls that already have 297 frames are copied there unchanged. Because SGLang opens inputs by path and the sample lives in this notebook's Hugging Face cache, the cell also copies each RGB condition into the output folder, so every path in the request lies inside the mounted cookbook root. To use WSM controls from another directory, such as a different render of this scene, set `CONTROL_DIR` to a folder with one `.mp4` per camera." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sim-auto-sglang-10", + "metadata": {}, + "outputs": [], + "source": [ + "REQUEST_PATH = AUTO_ROOT / \"assets\" / \"daytime_scene\" / \"sim_auto.json\"\n", + "CONTROL_DIR = None # Or a directory with one .mp4 WSM control per camera\n", + "\n", + "manifest = load_request(REQUEST_PATH, SAMPLE_ROOT, caption_folder=CAPTION_FOLDER, control_dir=Path(CONTROL_DIR) if CONTROL_DIR else None)\n", + "require_inputs(manifest)\n", + "views = manifest[\"multiview\"][\"views\"]\n", + "validate_rgb_conditions(views, SAMPLE_ROOT)\n", + "\n", + "sample_dir = OUTPUT_DIR / \"rgb\"\n", + "controls = prepare_controls(views, SAMPLE_ROOT, sample_dir / \"prepared_controls\", manifest[\"num_frames\"], manifest[\"fps\"], materialize=True)\n", + "conditions = stage_conditions(views, sample_dir / \"staged_conditions\")\n", + "prepared = copy.deepcopy(manifest)\n", + "for view, control, condition in zip(prepared[\"multiview\"][\"views\"], controls, conditions, strict=True):\n", + " view[\"control_path\"] = str(control)\n", + " view[\"vision_path\"] = str(condition)\n", + "rgb_conditions = dict(zip(CAMERA_ORDER, conditions))\n", + "wsm_controls = dict(zip(CAMERA_ORDER, controls))\n", + "print(f\"{manifest['name']}: {manifest['num_frames']} frames per camera, {manifest['fps']:g} FPS, seed {manifest['seed']}\")\n", + "for camera_key, control, condition in zip(CAMERA_ORDER, controls, conditions):\n", + " print(f\" {camera_key:26} WSM {control.name}, RGB {condition.name}\")" + ] + }, + { + "cell_type": "markdown", + "id": "sim-auto-sglang-11", + "metadata": {}, + "source": [ + "### Preview the inputs\n", + "\n", + "The first frame of each camera's WSM control and of its RGB condition." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sim-auto-sglang-12", + "metadata": {}, + "outputs": [], + "source": [ + "def camera_frames(paths: dict[str, Path], index: int, title: str) -> None:\n", + " \"\"\"Show one frame of each camera in a 2 × 4 grid.\"\"\"\n", + " fig, axes = plt.subplots(2, 4, figsize=(16, 5), layout=\"constrained\")\n", + " for ax in axes.flat:\n", + " ax.axis(\"off\")\n", + " for ax, (camera_key, path) in zip(axes.flat, paths.items()):\n", + " ax.imshow(read_frame(path, index))\n", + " ax.set_title(camera_key, fontsize=9)\n", + " fig.suptitle(title)\n", + " display(fig)\n", + " plt.close(fig)\n", + "\n", + "\n", + "camera_frames(wsm_controls, 0, \"WSM controls — frame 0\")\n", + "camera_frames(rgb_conditions, 0, \"RGB conditions — frame 0\")" + ] + }, + { + "cell_type": "markdown", + "id": "sim-auto-sglang-13", + "metadata": {}, + "source": [ + "## 5. Generate\n", + "\n", + "The example submits one `POST /v1/videos` job for the complete rollout, polls `GET /v1/videos/{id}` until it completes, and downloads the finished MP4 from `GET /v1/videos/{id}/content`. The request is multipart form fields with no file uploads:\n", + "\n", + "- **Form fields:** the model, top-level caption, canvas `size`, `num_frames`, `fps`, and `seed`, taken from the request file, plus the checkpoint's fixed `guidance_scale` of 1 and four `num_inference_steps`. The server rejects other values for the last two instead of ignoring them.\n", + "- **`extra_params` (JSON):** `wsm: {}`; `multiview.views` with each camera's key, caption, and the absolute `control_path` and `vision_path` the server opens; `multiview.condition_video_as_image: true`; `resolution` and `aspect_ratio`. SGLang takes no `*_reference_index` fields and no session fields: each job is one complete rollout, and the model stays loaded between jobs.\n", + "\n", + "The next cell builds the request and prints it, with the captions shortened and the paths relative to the cookbook root." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sim-auto-sglang-14", + "metadata": {}, + "outputs": [], + "source": [ + "form, uploads = prepare_request(prepared, MODEL_ID)\n", + "extra_params = json.loads(form[\"extra_params\"])\n", + "assert not uploads # SGLang opens control_path and vision_path itself; nothing is uploaded.\n", + "assert [view[\"camera_key\"] for view in extra_params[\"multiview\"][\"views\"]] == list(CAMERA_ORDER)\n", + "assert all(\n", + " Path(view[key]).is_absolute() and not any(field.endswith(\"_reference_index\") for field in view)\n", + " for view in extra_params[\"multiview\"][\"views\"]\n", + " for key in (\"control_path\", \"vision_path\")\n", + ")\n", + "\n", + "\n", + "def short(text: str) -> str:\n", + " return text if len(text) <= 80 else text[:80] + \" ...\"\n", + "\n", + "\n", + "def portable(path: str) -> str:\n", + " return str(Path(path).relative_to(COOKBOOK_ROOT)) if Path(path).is_relative_to(COOKBOOK_ROOT) else path\n", + "\n", + "\n", + "for view in extra_params[\"multiview\"][\"views\"]:\n", + " view[\"prompt\"] = short(view[\"prompt\"])\n", + " for key in (\"control_path\", \"vision_path\"):\n", + " view[key] = portable(view[key])\n", + "printable = {**form, \"prompt\": short(form[\"prompt\"]), \"extra_params\": extra_params}\n", + "print(json.dumps(printable, indent=2, ensure_ascii=False))" + ] + }, + { + "cell_type": "markdown", + "id": "sim-auto-sglang-15", + "metadata": {}, + "source": [ + "Generate the rollout. The first job takes longer than later ones, and the cell prints the job URL as soon as the server accepts the request. Each run of this cell submits a new job; a failed run leaves no video for the next step to validate." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sim-auto-sglang-16", + "metadata": {}, + "outputs": [], + "source": [ + "video = None\n", + "form, _ = prepare_request(prepared, MODEL_ID)\n", + "with httpx.Client(base_url=SGLANG_BASE_URL, headers=HEADERS, timeout=HTTP_TIMEOUT) as client:\n", + " started = time.monotonic()\n", + " job = submit_video_job(client, form, sample_dir, timeout=GENERATION_TIMEOUT)\n", + " video = download_video(client, job, sample_dir / \"multiview_camera_major.mp4\")\n", + " elapsed_seconds = round(time.monotonic() - started, 1)\n", + "print(f\"{elapsed_seconds:.1f}s, including polling and download {video}\")" + ] + }, + { + "cell_type": "markdown", + "id": "sim-auto-sglang-17", + "metadata": {}, + "source": [ + "## 6. Validate and view\n", + "\n", + "The result is one camera-major MP4: 297 frames of the first camera, then 297 of the second, through the seventh. Decode it completely, check its `7 × 297 = 2079` frames, FPS, and canvas, then split it into one labeled video per camera. The folder then holds:\n", + "\n", + "```text\n", + "outputs/sglang/sim_auto/rgb/\n", + " request.json # The exact form fields that were sent\n", + " job.json # The last job response\n", + " prepared_controls/ # The seven trimmed WSM controls the server read\n", + " staged_conditions/ # The seven RGB conditions, copied from the sample for the server\n", + " multiview_camera_major.mp4 # The downloaded result\n", + " vision_viewNN_.mp4 # One video per camera\n", + " sample_outputs.json # Status, timing, job, inputs, cameras, and the decoded video properties\n", + "```\n", + "\n", + "A rejected submission or a failed job leaves `response_error.txt` with the server's response instead of the MP4. A video that decodes at the right size does not establish control adherence, cross-camera consistency, or visual quality; inspect the frames and playback below, and see the [shared Simulation checklist](../README.md#validate-a-rollout)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sim-auto-sglang-18", + "metadata": {}, + "outputs": [], + "source": [ + "if video is None:\n", + " raise RuntimeError(\"Complete generation before validating outputs.\")\n", + "num_frames, fps = manifest[\"num_frames\"], manifest[\"fps\"]\n", + "info = validate_video(video, num_frames=len(CAMERA_ORDER) * num_frames, fps=fps, width=832, height=480)\n", + "cameras = split_cameras(video, sample_dir, CAMERA_ORDER, num_frames, fps)\n", + "write_sample_outputs(sample_dir, [{\n", + " \"name\": manifest[\"name\"],\n", + " \"status\": \"success\",\n", + " \"job_id\": job[\"id\"],\n", + " \"model_revision\": MODEL_REVISION,\n", + " \"sample_revision\": SAMPLE_REVISION,\n", + " \"video\": video,\n", + " \"request\": sample_dir / \"request.json\",\n", + " \"job\": sample_dir / \"job.json\",\n", + " \"inputs\": [\n", + " view[key] for view in json.loads(form[\"extra_params\"])[\"multiview\"][\"views\"] for key in (\"control_path\", \"vision_path\")\n", + " ],\n", + " \"elapsed_seconds\": elapsed_seconds,\n", + " \"frames_per_camera\": num_frames,\n", + " \"cameras\": cameras,\n", + " **info,\n", + "}])\n", + "print(f\"{info['frames']} camera-major frames, {info['width']}×{info['height']}, {info['fps']:g} FPS; \"\n", + " f\"full decode passed; split into {len(cameras)} cameras of {num_frames} frames\")" + ] + }, + { + "cell_type": "markdown", + "id": "sim-auto-sglang-19", + "metadata": {}, + "source": [ + "### Compare with the conditions\n", + "\n", + "The first generated frame of each camera, to compare with its RGB condition above, and the final frame. The first frame should closely follow the condition; later frames should follow the WSM controls' motion and layout and keep objects consistent where neighboring cameras overlap." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sim-auto-sglang-20", + "metadata": {}, + "outputs": [], + "source": [ + "camera_frames(cameras, 0, \"SGLang output — frame 0\")\n", + "camera_frames(cameras, num_frames - 1, f\"SGLang output — frame {num_frames - 1}\")" + ] + }, + { + "cell_type": "markdown", + "id": "sim-auto-sglang-21", + "metadata": {}, + "source": [ + "### Play the videos\n", + "\n", + "Each camera plays muted and looped. The players embed small re-encoded previews, so saved notebook outputs stay compact; the full-resolution MP4s remain in the output folder." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sim-auto-sglang-22", + "metadata": {}, + "outputs": [], + "source": [ + "display(HTML(video_row(cameras, width=320)))" + ] + }, + { + "cell_type": "markdown", + "id": "sim-auto-sglang-23", + "metadata": {}, + "source": [ + "## 7. Next steps\n", + "\n", + "- **Use your own scene:** copy [`sim_auto.json`](assets/daytime_scene/sim_auto.json), point its paths at your data, then set `REQUEST_PATH` to the copy and `SAMPLE_ROOT` to the folder its paths are relative to. Keep the seven cameras in the order above, each with a `control_path`, a `vision_path`, and a caption file or inline `prompt`. Keep `num_frames` at 4k + 1, guidance 1, and four steps. The prepared controls and staged conditions always land under the output folder, so keep it inside the mounted cookbook root. For a different WSM render of this scene, set `CONTROL_DIR` instead.\n", + "- **Validate further:** apply the [shared Simulation checklist](../README.md#validate-a-rollout). Check camera alignment and temporal consistency before using results downstream, and keep run-specific notes, including the server's GPU and the checkpoint revision it resolved from `main`, alongside your generated outputs.\n", + "- **Troubleshooting:**\n", + " - *Missing inputs:* rerun the sample download in step 3 and then **Prepare the inputs**; `AUTO_SAMPLE_ROOT` must point at a complete `sample_data` folder.\n", + " - *HTTP error on submission:* inspect `response_error.txt` and the server log. The pipeline needs a WSM control and an RGB condition for every camera, cameras from the checkpoint's eleven, guidance 1, four steps, and `condition_video_as_image: true`. A path the server cannot open means the checkout or the output folder is not mounted at the same absolute path inside the container.\n", + " - *Job `failed`:* `job.json` and `response_error.txt` carry the server's error message; the server log has the traceback.\n", + " - *`Cosmos3-Nano-Sim-Auto requires transformer/config.json['cosmos3_nano_sim_auto']` at start-up:* the checkpoint is not a Sim-Auto export with the SGLang artifact. Use the published checkpoint on `main`, or an export produced with `--cosmos3-nano-sim-auto`.\n", + " - *Timeout:* the job keeps running on the server at the printed URL; `GET /v1/videos/{id}` reports its status. Raise `GENERATION_TIMEOUT` before resubmitting.\n", + "- **Other models and runtimes:** [Cosmos3-Nano-Transfer-Auto with SGLang](run_auto_with_sglang.ipynb) generates all eleven cameras with full diffusion sampling, with an optional joint LiDAR example. The [vLLM-Omni notebook](run_sim_auto_with_vllm_omni.ipynb) serves this sample through vLLM-Omni's synchronous endpoint, and the [Framework notebook](run_sim-auto_with_cosmos_framework.ipynb) runs Sim-Auto natively, including joint RGB and LiDAR generation." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_auto_media.py b/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_auto_media.py index b70aefdc9..5ba913628 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_auto_media.py +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_auto_media.py @@ -78,6 +78,19 @@ def test_trim_preserves_frame_rate_count_geometry_and_sequence(self): for output, reference, width in zip(prepared, references, (64, 80), strict=True): self.assert_video(output, reference, fps, width) + def test_exact_length_controls_are_copied_only_when_materialized(self): + fps, frames = 30, 201 + source = self.make_video("exact.mp4", [20 + i % 211 for i in range(frames)], fps) + reference = self.frame_levels(source) + views = [{"camera_key": "camera_0", "control_path": source.name}] + in_place = prepare_controls(views, self.root, self.root / "prepared", frames, fps) + self.assertEqual(in_place, [source.resolve()]) + self.assertEqual(list((self.root / "prepared").iterdir()), []) + copied = prepare_controls(views, self.root, self.root / "materialized", frames, fps, materialize=True) + self.assertEqual(copied, [self.root / "materialized" / "camera_0.mp4"]) + self.assertEqual(copied[0].read_bytes(), source.read_bytes()) + self.assert_video(copied[0], reference, fps) + def test_split_preserves_camera_boundaries_and_each_frames_timeline(self): for fps in (30, 24): with self.subTest(fps=fps): diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_daytime_sample.py b/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_daytime_sample.py index 9d121a4f6..f0e15e192 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_daytime_sample.py +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_daytime_sample.py @@ -1,5 +1,6 @@ """CPU checks for the shared daytime sample requests and the LiDAR range-map conversion.""" +import contextlib import copy import hashlib import io @@ -16,7 +17,9 @@ AUTO_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(AUTO_ROOT)) -from helpers.auto_helpers import CAMERA_ORDER, prepare_request +from helpers.auto_helpers import ( + CAMERA_ORDER, prepare_request, stage_conditions, submit_joint, validate_sglang_lidar, +) from helpers.daytime_sample import SAMPLE_CAMERAS, read_caption, resolve_request, sample_files, verify_sample from helpers.lidar_rangemap import convert_control, pool_nearest_return @@ -194,6 +197,182 @@ def test_inline_requests_pass_through_with_resolved_paths(self): self.assertEqual(manifest["multiview"]["views"][0]["control_path"], str(self.root / "wsm.mp4")) +class SglangRequestTests(unittest.TestCase): + """The SGLang branch of the shared request helpers: server-visible paths, async jobs, LiDAR artifacts.""" + + def setUp(self): + temporary = tempfile.TemporaryDirectory() + self.addCleanup(temporary.cleanup) + self.root = make_sample(Path(temporary.name).resolve()) + converter = patch("helpers.daytime_sample.convert_control", side_effect=mock_convert) + converter.start() + self.addCleanup(converter.stop) + self.lidar = self.root / "prepared" / "lidar_control.safetensors" + + def resolve(self, name): + return resolve_request(ASSET_DIR / name, self.root, caption_folder="captions", lidar_output=self.lidar) + + def test_sglang_requests_carry_server_visible_paths_instead_of_uploads(self): + for name in ("transfer_auto_wsm.json", "transfer_auto_joint.json"): + with self.subTest(request=name): + template = json.loads((ASSET_DIR / name).read_text()) + manifest = self.resolve(name) + original = copy.deepcopy(manifest) + form, uploads = prepare_request(manifest, self.root, "model", backend="sglang") + vllm_form, vllm_uploads = prepare_request(manifest, self.root, "model") + self.assertEqual(manifest, original) + self.assertEqual(uploads, []) + self.assertEqual(len(vllm_uploads), 15 if "lidar" in template else 7) + self.assertEqual( + {key: value for key, value in form.items() if key != "extra_params"}, + {key: value for key, value in vllm_form.items() if key != "extra_params"}, + ) + extra = json.loads(form["extra_params"]) + self.assertNotIn("model_mode", extra) + for source, view in zip(template["multiview"]["views"], extra["multiview"]["views"], strict=True): + self.assertEqual([key for key in view if key.endswith("_reference_index")], []) + self.assertEqual(view["prompt"], f"Daytime caption for {source['camera_key']}.") + for role in ("control", "vision"): + if f"{role}_path" in source: + self.assertEqual(Path(view[f"{role}_path"]), self.root / source[f"{role}_path"]) + else: + self.assertNotIn(f"{role}_path", view) + if "lidar" in template: + self.assertEqual(extra["lidar"], {"control_path": str(self.lidar), "decode": True}) + self.assertEqual(extra["control_guidance"], 1) + else: + self.assertNotIn("lidar", extra) + + def test_sglang_lidar_decode_follows_return_output_and_unknown_backends_reject(self): + manifest = self.resolve("transfer_auto_joint.json") + manifest["lidar"]["return_output"] = False + form, _ = prepare_request(manifest, self.root, "model", backend="sglang") + self.assertFalse(json.loads(form["extra_params"])["lidar"]["decode"]) + with self.assertRaisesRegex(ValueError, "return_output=true"): + prepare_request(manifest, self.root, "model") + manifest["lidar"]["decode"] = True + with self.assertRaisesRegex(ValueError, "agree"): + prepare_request(manifest, self.root, "model", backend="sglang") + with self.assertRaisesRegex(ValueError, "backend"): + prepare_request(self.resolve("transfer_auto_wsm.json"), self.root, "model", backend="trt") + + def test_stage_conditions_copies_each_camera_under_its_key(self): + views = self.resolve("transfer_auto_joint.json")["multiview"]["views"] + staged = stage_conditions(views, self.root / "staged") + self.assertEqual([path.name for path in staged], [f"{view['camera_key']}.mp4" for view in views]) + for view, path in zip(views, staged, strict=True): + self.assertEqual(path.read_bytes(), Path(view["vision_path"]).read_bytes()) + self.assertEqual(stage_conditions(views, self.root / "staged"), staged) # Idempotent. + changed = Path(views[0]["vision_path"]) + changed.write_bytes(b"X" * changed.stat().st_size) # Same length, different content. + self.assertEqual(stage_conditions(views, self.root / "staged")[0].read_bytes(), changed.read_bytes()) + Path(views[0]["vision_path"]).unlink() + with self.assertRaises(FileNotFoundError): + stage_conditions(views, self.root / "staged-again") + + def test_stage_conditions_keeps_the_request_suffix_for_snapshot_symlinks(self): + # A Hugging Face snapshot names files normally but symlinks them into an extension-less blob store. + views = self.resolve("transfer_auto_joint.json")["multiview"]["views"] + blobs = self.root / "blobs" + blobs.mkdir() + for index, view in enumerate(views): + named = Path(view["vision_path"]) + blob = blobs / f"{index:064x}" + named.replace(blob) + named.symlink_to(blob) + staged = stage_conditions(views, self.root / "staged") + self.assertEqual([path.name for path in staged], [f"{view['camera_key']}.mp4" for view in views]) + for view, path in zip(views, staged, strict=True): + self.assertFalse(path.is_symlink()) + self.assertEqual(path.read_bytes(), Path(view["vision_path"]).read_bytes()) + + def test_sglang_joint_job_downloads_the_video_and_rangemap(self): + import httpx + + manifest = self.resolve("transfer_auto_joint.json") + form, uploads = prepare_request(manifest, self.root, "model", backend="sglang") + descriptor = {"sweeps": 67, "fps": 10.0, "files": {"rangemap": "/srv/job-1.lidar.safetensors"}} + polls = [] + + def handler(request): + path = request.url.path + if path == "/v1/videos" and request.method == "POST": + body = request.read() + self.assertIn(b'name="extra_params"', body) + self.assertNotIn(b'name="input_references"', body) + return httpx.Response(200, json={"id": "job-1", "status": "queued"}) + if path == "/v1/videos/job-1": + polls.append(request.method) + return httpx.Response(200, json={"id": "job-1", "status": "completed", "lidar": descriptor}) + if path == "/v1/videos/job-1/content": + if request.url.params.get("variant") == "rangemap": + return httpx.Response(200, content=b"mock rangemap") + return httpx.Response(200, content=b"mock video", headers={"content-type": "video/mp4"}) + self.fail(f"Unexpected request {request.method} {request.url}") + + output_dir = self.root / "out" + with ( + patch("time.sleep"), + contextlib.redirect_stdout(io.StringIO()), + httpx.Client(base_url="http://test", transport=httpx.MockTransport(handler)) as client, + ): + result = submit_joint(client, form, uploads, output_dir, backend="sglang") + self.assertEqual((result["job_id"], result["descriptor"]), ("job-1", descriptor)) + self.assertEqual(result["video"].read_bytes(), b"mock video") + self.assertEqual(result["lidar"].read_bytes(), b"mock rangemap") + self.assertEqual(json.loads(result["lidar"].with_suffix(".json").read_text()), descriptor) + self.assertEqual(json.loads((output_dir / "job.json").read_text())["status"], "completed") + self.assertEqual(polls, ["GET"]) + + # The vLLM-Omni contract is unchanged: uploads travel as input_references and no rangemap is fetched. + vllm_form, vllm_uploads = prepare_request(manifest, self.root, "model") + seen = [] + + def vllm_handler(request): + seen.append(request.url.path) + if request.method == "POST": + self.assertIn(b'name="input_references"', request.read()) + return httpx.Response(200, json={"id": "job-2", "status": "completed", "lidar": descriptor}) + return httpx.Response(200, content=b"mock video", headers={"content-type": "video/mp4"}) + + with ( + patch("time.sleep"), + contextlib.redirect_stdout(io.StringIO()), + httpx.Client(base_url="http://test", transport=httpx.MockTransport(vllm_handler)) as client, + ): + result = submit_joint(client, vllm_form, vllm_uploads, self.root / "vllm") + self.assertEqual(result, { + "video": self.root / "vllm" / "multiview_camera_major.mp4", "descriptor": descriptor, "job_id": "job-2", + }) + self.assertEqual(seen, ["/v1/videos", "/v1/videos/job-2/content"]) + + def test_validate_sglang_lidar_checks_layout_metadata_and_masked_values(self): + from safetensors.numpy import save_file + + sweeps = 2 + validity = np.zeros((sweeps, 128, 1800), dtype=np.uint8) + validity[:, :64] = 1 + ranges = np.where(validity == 1, 12.5, 0.0).astype(np.float32) + intensity = np.where(validity == 1, 0.25, 0.0).astype(np.float16) + metadata = { + "layout": "[sweeps, beams, azimuth]", "fps": "10.0", + "min_range_m": "0.5", "max_range_m": "120.0", "invalid_range_m": "0", + } + path = self.root / "generated.safetensors" + save_file({"range_m": ranges, "intensity": intensity, "validity": validity}, path, metadata=metadata) + info = validate_sglang_lidar(path, sweeps) + self.assertEqual((info["shape"], info["metadata"]["fps"]), ([sweeps, 128, 1800], "10.0")) + with self.assertRaisesRegex(ValueError, "metadata"): + validate_sglang_lidar(path, sweeps, fps=30.0) + with self.assertRaisesRegex(ValueError, "shape"): + validate_sglang_lidar(path, sweeps + 1) + broken = ranges.copy() + broken[0, 0, 0] = 0.0 # A valid ray must have a positive range. + save_file({"range_m": broken, "intensity": intensity, "validity": validity}, path, metadata=metadata) + with self.assertRaisesRegex(ValueError, "masked"): + validate_sglang_lidar(path, sweeps) + + class LidarRangemapTests(unittest.TestCase): def setUp(self): temporary = tempfile.TemporaryDirectory() diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_sim_auto_helpers.py b/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_sim_auto_helpers.py index aaba503f9..9229875f9 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_sim_auto_helpers.py +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_sim_auto_helpers.py @@ -24,6 +24,7 @@ from helpers.sim_auto_vllm_omni_helpers import ( CAMERA_ORDER, MODEL_ID, load_request, prepare_request, require_inputs, ) +from helpers.sim_auto_sglang_helpers import prepare_request as prepare_sglang_request from test_daytime_sample import make_sample, mock_convert ASSET_DIR = AUTO_ROOT / "assets" / "daytime_scene" @@ -297,5 +298,187 @@ def stub_media(scope): exec(cells["sim-auto-18"], scope) +class SimAutoSglangRequestTests(SimAutoFixture, unittest.TestCase): + def test_server_visible_paths_without_uploads_or_sessions(self): + manifest = load_request(self.request_path, self.sample_root) + original = copy.deepcopy(manifest) + form, paths = prepare_sglang_request(manifest) + extra = json.loads(form["extra_params"]) + self.assertEqual(manifest, original) + self.assertEqual(paths, []) + self.assertTrue(all(isinstance(value, str) for value in form.values())) + self.assertEqual(form["model"], MODEL_ID) + self.assertEqual(form["prompt"], f"Daytime caption for {CAMERA_ORDER[0]}.") + self.assertEqual((form["guidance_scale"], form["num_inference_steps"], form["seed"]), ("1.0", "4", "0")) + self.assertEqual((form["num_frames"], form["fps"], form["size"]), ("297", "30", "832x480")) + for field in ("negative_prompt", "flow_shift", "enable_frame_interpolation"): + self.assertNotIn(field, form) + self.assertEqual(set(extra), {"wsm", "multiview", "resolution", "aspect_ratio"}) + self.assertEqual((extra["wsm"], extra["resolution"], extra["aspect_ratio"]), ({}, "480", "16,9")) + self.assertEqual(set(extra["multiview"]), {"views", "condition_video_as_image", "num_frames"}) + self.assertTrue(extra["multiview"]["condition_video_as_image"]) + self.assertEqual(extra["multiview"]["num_frames"], 297) + for source, mapped in zip(self.manifest["multiview"]["views"], extra["multiview"]["views"], strict=True): + self.assertEqual(set(mapped), {"camera_key", "prompt", "control_path", "vision_path"}) + self.assertEqual(mapped["camera_key"], source["camera_key"]) + self.assertEqual(mapped["prompt"], f"Daytime caption for {source['camera_key']}.") + for role in ("control", "vision"): + self.assertEqual(Path(mapped[f"{role}_path"]), self.sample_root / source[f"{role}_path"]) + self.assertTrue(Path(mapped[f"{role}_path"]).is_absolute()) + self.assertEqual(prepare_sglang_request(manifest)[0], form) # No per-request session identifiers. + + def test_joint_request_names_the_converted_lidar_control(self): + lidar = self.root / "lidar_control.safetensors" + manifest = load_request(self.joint_path, self.sample_root, lidar_output=lidar) + form, paths = prepare_sglang_request(manifest) + extra = json.loads(form["extra_params"]) + self.assertEqual(paths, []) + self.assertEqual(extra["lidar"], {"control_path": str(lidar), "decode": True}) + self.assertEqual(form["fps"], "30") + lidar.unlink() + with self.assertRaisesRegex(ValueError, "LiDAR control"): + prepare_sglang_request(manifest) + + def test_contract_and_missing_inputs_reject_before_submission(self): + manifest = load_request(self.request_path, self.sample_root) + for field, value in (("guidance_scale", 3.0), ("num_inference_steps", 8), ("num_frames", 296)): + with self.assertRaises(ValueError): + prepare_sglang_request({**manifest, field: value}) + reordered = load_request(self.request_path, self.sample_root) + reordered["multiview"]["views"].reverse() + with self.assertRaisesRegex(ValueError, "canonical"): + prepare_sglang_request(reordered) + Path(manifest["multiview"]["views"][0]["control_path"]).write_bytes(b"version https://git-lfs.github.com/spec/v1\n") + with self.assertRaisesRegex(ValueError, "Git LFS pointer"): + prepare_sglang_request(manifest) + + +class SimAutoSglangNotebookTests(SimAutoFixture, unittest.TestCase): + def test_execution_order_jobs_and_failed_run_cannot_reuse_prior_video(self): + notebook = json.loads((AUTO_ROOT / "run_sim-auto_with_sglang.ipynb").read_text()) + cells = ["".join(cell["source"]) for cell in notebook["cells"] if cell["cell_type"] == "code"] + bundled = 'REQUEST_PATH = AUTO_ROOT / "assets" / "daytime_scene" / "sim_auto.json"' + self.assertEqual(sum(bundled in cell for cell in cells), 1) + cells = [cell.replace(bundled, f"REQUEST_PATH = Path({str(self.request_path)!r})") for cell in cells] + generate = next(cell for cell in cells if "job = submit_video_job(" in cell) + validate = next(cell for cell in cells if "Complete generation" in cell) + import IPython.display + + original_client = httpx.Client + state = {"fail": None, "jobs": 0, "polls": 0} + + def handler(request): + path = request.url.path + if path == "/v1/models": + return httpx.Response(200, json={"data": [{"id": MODEL_ID}]}) + if path == "/v1/videos" and request.method == "POST": + body = request.read() + self.assertIn(b'name="extra_params"', body) + self.assertIn(b'name="num_inference_steps"', body) + self.assertNotIn(b'name="input_references"', body) + self.assertNotIn(b"_reference_index", body) + self.assertNotIn(b"session_id", body) + if state["fail"] == "submit": + return httpx.Response(400, text="mock submission rejected") + state["jobs"] += 1 + return httpx.Response(200, json={"id": f"job-{state['jobs']}", "status": "queued", "progress": 0}) + job_id = f"job-{state['jobs']}" + if path == f"/v1/videos/{job_id}" and request.method == "GET": + state["polls"] += 1 + if state["fail"] == "job": + return httpx.Response(200, json={"id": job_id, "status": "failed", "error": {"message": "mock rollout failed"}}) + return httpx.Response(200, json={"id": job_id, "status": "completed", "progress": 100}) + if path == f"/v1/videos/{job_id}/content" and request.method == "GET": + return httpx.Response(200, content=b"mock video", headers={"content-type": "video/mp4"}) + self.fail(f"Unexpected request {request.method} {path}") + + def client_factory(*args, **kwargs): + return original_client(*args, transport=httpx.MockTransport(handler), **kwargs) + + def split(video, output, cameras, frames, fps): + self.assertEqual((frames, fps), (297, 30)) + self.assertEqual(tuple(cameras), CAMERA_ORDER) + paths = {camera: output / f"{camera}.mp4" for camera in cameras} + for path in paths.values(): + path.write_bytes(b"mock camera") + return paths + + def stub_media(scope): + # Replace the FFmpeg-backed notebook helpers; the fixture videos are placeholders. + scope.update( + validate_video=lambda path, **expected: { + "frames": expected["num_frames"], "fps": float(expected["fps"]), + "width": expected["width"], "height": expected["height"], + }, + read_frame=lambda path, index: PILImage.new("RGB", (8, 8)), + video_row=lambda videos, **kwargs: "", + ) + + from helpers import auto_helpers + environment = { + "COSMOS_COOKBOOK_ROOT": str(AUTO_ROOT.parents[3]), + "SGLANG_BASE_URL": "http://test", "SGLANG_API_KEY": "", + "COSMOS3_OUTPUT_DIR": str(self.root / "runs"), + "AUTO_SAMPLE_ROOT": str(self.sample_root), + } + controls = [self.sample_root / view["control_path"] for view in self.manifest["multiview"]["views"]] + with ( + contextlib.redirect_stdout(io.StringIO()), + patch.dict("os.environ", environment), + patch("time.sleep"), + patch.object(IPython.display, "display"), + patch.object(IPython.display, "HTML"), + patch.object(httpx, "Client", side_effect=client_factory), + patch.object(auto_helpers, "validate_rgb_conditions"), + patch.object(auto_helpers, "prepare_controls", return_value=controls) as prepare_controls, + patch.object(auto_helpers, "split_cameras", side_effect=split), + ): + scope = {} + for cell in cells: + exec(cell, scope) + if "def validate_video" in cell: + stub_media(scope) + self.assertEqual(scope["SAMPLE_ROOT"], self.sample_root) + # SGLang opens the controls by path, so exact-length controls must be copied under the mount. + self.assertTrue(prepare_controls.call_args.kwargs.get("materialize")) + sample_dir = scope["sample_dir"] + staged = sample_dir / "staged_conditions" + self.assertEqual({path.name for path in staged.iterdir()}, {f"{camera}.mp4" for camera in CAMERA_ORDER}) + record, = json.loads((sample_dir / "sample_outputs.json").read_text()) + self.assertEqual((record["status"], record["job_id"], len(record["cameras"])), ("success", "job-1", 7)) + self.assertEqual(len(record["inputs"]), 14) + self.assertEqual(state["polls"], 1) # Queued at submission, completed at the first poll. + self.assertEqual(json.loads((sample_dir / "job.json").read_text())["status"], "completed") + request = json.loads((sample_dir / "request.json").read_text()) + extra = json.loads(request["extra_params"]) + self.assertNotIn("session_id", extra) + for view, control in zip(extra["multiview"]["views"], controls, strict=True): + self.assertEqual(Path(view["control_path"]), control) + self.assertEqual(Path(view["vision_path"]), staged / f"{view['camera_key']}.mp4") + self.assertEqual(view["prompt"], f"Daytime caption for {view['camera_key']}.") + + exec(generate, scope) + exec(validate, scope) + record, = json.loads((sample_dir / "sample_outputs.json").read_text()) + self.assertEqual(record["job_id"], "job-2") + + state["fail"] = "submit" + with self.assertRaisesRegex(RuntimeError, "HTTP 400"): + exec(generate, scope) + self.assertIsNone(scope["video"]) + self.assertEqual((sample_dir / "response_error.txt").read_text(), "mock submission rejected") + with self.assertRaisesRegex(RuntimeError, "Complete generation"): + exec(validate, scope) + + state["fail"] = "job" + with self.assertRaisesRegex(RuntimeError, "Video generation failed"): + exec(generate, scope) + self.assertIsNone(scope["video"]) + self.assertEqual(json.loads((sample_dir / "job.json").read_text())["status"], "failed") + self.assertIn("mock rollout failed", (sample_dir / "response_error.txt").read_text()) + with self.assertRaisesRegex(RuntimeError, "Complete generation"): + exec(validate, scope) + + if __name__ == "__main__": unittest.main() diff --git a/.agents/skills/cookbooks/cosmos3/simulation/bimanual/bimanual_helpers.py b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/bimanual_helpers.py new file mode 100644 index 000000000..b17e5b13a --- /dev/null +++ b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/bimanual_helpers.py @@ -0,0 +1,354 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 +"""CPU-only input preparation and HTTP client for the SGLang Bimanual notebook. + +This client targets the unified (schema-5) Cosmos3-Nano-Sim-Bimanual export: 59-value +rows that are already normalized, with one embodiment name per row. The server +validates the checkpoint contract, pads rows to the model width and routes every row +to its embodiment domain. +""" + +from __future__ import annotations + +import json +import math +import subprocess +import time +from dataclasses import dataclass +from io import BytesIO +from pathlib import Path + +import numpy as np +from PIL import Image, ImageOps, UnidentifiedImageError + +# Keep the server's order for ties when selecting the closest image aspect. +CANVASES = ((640, 640), (736, 544), (544, 736), (832, 480), (480, 832)) +VIEWPOINTS = {"ego_view", "third_person_view", "wrist_view", "concat_view"} +UNIFIED_ACTION_SPACE = "normalized_unified_v1" +UNIFIED_ACTION_DIM = 59 +# Keys of the earlier native (29-value) recipe that the unified export does not take. +_RETIRED_KEYS = ("camera_trajectory", "camera_pose_convention", "camera_translation_scale", "camera_action_normalization") + + +@dataclass(frozen=True) +class BimanualInput: + prompt: str + image: Image.Image + action: np.ndarray + domain_names: list[str] + extra_args: dict + num_frames: int + fps: int + seed: int + height: int + width: int + output_height: int + output_width: int + metadata: dict + + +def _integer(value, name: str, minimum: int = 1) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + raise ValueError(f"{name} must be an integer >= {minimum}.") + return value + + +def _asset(value, base_dir: Path) -> Path: + if not isinstance(value, (str, Path)) or not str(value).strip(): + raise ValueError("Supply a local image/video or action file path.") + path = (base_dir / Path(value).expanduser()).resolve() + if not path.is_file() or path.stat().st_size == 0: + raise FileNotFoundError(f"Missing or empty input: {path}") + with path.open("rb") as stream: + if stream.read(100).startswith(b"version https://git-lfs.github.com/spec/v1"): + raise ValueError(f"Unhydrated Git LFS pointer: {path}. Run git lfs pull.") + return path + + +def first_media_image(path: Path) -> Image.Image: + """Read a local image or extract the first video frame with FFmpeg.""" + try: + with Image.open(path) as image: + return ImageOps.exif_transpose(image).convert("RGB") + except UnidentifiedImageError: + result = subprocess.run( + [ + "ffmpeg", "-nostdin", "-v", "error", "-i", str(path), + "-frames:v", "1", "-f", "image2pipe", "-vcodec", "png", "-", + ], + capture_output=True, + check=True, + ) + if not result.stdout: + raise ValueError(f"No conditioning frame could be decoded from {path}.") + with Image.open(BytesIO(result.stdout)) as image: + return image.convert("RGB") + + +def _numeric_array(value, name: str) -> np.ndarray: + try: + values = np.asarray(value, dtype=object) + if any( + isinstance(item, (bool, np.bool_)) or not isinstance(item, (int, float, np.integer, np.floating)) + for item in values.flat + ): + raise ValueError(f"{name} must contain real numbers, not booleans or strings.") + array = values.astype(np.float64) + except (TypeError, OverflowError) as exc: + raise ValueError(f"Invalid numeric {name}.") from exc + if not np.isfinite(array).all(): + raise ValueError(f"{name} must contain only finite values.") + return array + + +def load_action_sidecar(path: Path, *, default_domain: str | None = None) -> tuple[np.ndarray, list[str]]: + """Rows and one embodiment name per row from an action file. + + The file is either the imaginaire4 unified sidecar object + ``{"action_space": "normalized_unified_v1", "action": rows, "domain_names": names}`` + or a bare array of rows, in which case every row uses ``default_domain``. + ``domain_names`` may hold one name for the whole clip or one name per row. + """ + payload = json.loads(path.read_text(encoding="utf-8")) + if isinstance(payload, dict): + space = payload.get("action_space") + if space != UNIFIED_ACTION_SPACE: + raise ValueError(f"{path}: action_space must be {UNIFIED_ACTION_SPACE!r}, got {space!r}.") + if "action" not in payload: + raise ValueError(f"{path}: the sidecar object needs an `action` list.") + rows = _numeric_array(payload["action"], "actions") + names = payload.get("domain_names") + if names is None: + if default_domain is None: + raise ValueError(f"{path}: the sidecar has no domain_names and the record sets no domain_name.") + names = [default_domain] + if isinstance(names, str): + names = [names] + else: + rows = _numeric_array(payload, "actions") + if default_domain is None: + raise ValueError(f"{path}: a bare action array needs domain_name in the record.") + names = [default_domain] + if rows.ndim != 2: + raise ValueError(f"{path}: actions must be a list of rows.") + if not isinstance(names, list) or not names or any(not isinstance(name, str) or not name.strip() for name in names): + raise ValueError(f"{path}: domain_names must be a non-empty list of embodiment names.") + if len(names) == 1: + names = names * len(rows) + if len(names) != len(rows): + raise ValueError(f"{path}: {len(names)} domain names for {len(rows)} action rows; supply one per row.") + return rows, [name.strip() for name in names] + + +def domain_runs(names: list[str]) -> list[tuple[str, int, int]]: + """Consecutive (name, start_row, end_row) segments; end is exclusive.""" + runs: list[tuple[str, int, int]] = [] + start = 0 + for index in range(1, len(names) + 1): + if index == len(names) or names[index] != names[start]: + runs.append((names[start], start, index)) + start = index + return runs + + +def prepare_bimanual_input(record: dict, *, base_dir: Path) -> BimanualInput: + """Prepare the notebook's unified-schema example without a runtime package.""" + num_frames = _integer(record.get("num_frames", 901), "num_frames", 2) + if (num_frames - 1) % 4: + raise ValueError("The Bimanual recipe requires num_frames = 4k+1.") + fps = _integer(record.get("fps", 30), "fps") + seed = _integer(record.get("seed", 42), "seed", 0) + if str(record.get("resolution", record.get("image_size", 480))) != "480": + raise ValueError("This Bimanual helper supports the checkpoint's 480p canvases.") + if record.get("model_mode", "forward_dynamics") != "forward_dynamics": + raise ValueError("This cookbook uses an initial image and actions for forward dynamics.") + if record.get("autoregressive") is False: + raise ValueError("Bimanual generation is autoregressive.") + retired = [key for key in _RETIRED_KEYS if key in record] + if retired or "action_space" in record: + raise ValueError( + f"{retired or ['action_space']} belong to the earlier 29-value recipe. The unified export takes " + f"pre-normalized {UNIFIED_ACTION_DIM}-value rows in a {UNIFIED_ACTION_SPACE!r} sidecar; convert camera " + "trajectories to unified rows before submitting them." + ) + normalization = record.get("action_normalization", "none") + if normalization != "none": + raise ValueError("Unified sidecars are already normalized; action_normalization must be 'none'.") + if any(record.get(key) is not None for key in ("action", "actions")): + raise ValueError("Point action_path at the sidecar file; inline rows are not supported here.") + + action_path = _asset(record.get("action_path"), base_dir) + rows, domain_names = load_action_sidecar(action_path, default_domain=record.get("domain_name")) + if rows.shape != (num_frames - 1, UNIFIED_ACTION_DIM): + raise ValueError(f"Expected {num_frames - 1}x{UNIFIED_ACTION_DIM} unified actions, got {rows.shape}.") + with np.errstate(over="ignore"): + action = rows.astype(np.float32) + if not np.isfinite(action).all(): + raise ValueError("Actions must be finite in float32.") + if record.get("raw_action_dim", UNIFIED_ACTION_DIM) != UNIFIED_ACTION_DIM: + raise ValueError(f"raw_action_dim must be {UNIFIED_ACTION_DIM} for the unified export.") + if _integer(record.get("action_chunk_size", num_frames - 1), "action_chunk_size") != num_frames - 1: + raise ValueError("action_chunk_size must equal num_frames - 1.") + domain = record.get("domain_name", domain_names[0]) + if not isinstance(domain, str) or not domain.strip(): + raise ValueError("domain_name must be an embodiment name.") + + vision_path = _asset(record.get("vision_path"), base_dir) + image = first_media_image(vision_path) + height, width = record.get("height"), record.get("width") + if height is None and width is None: + width, height = min(CANVASES, key=lambda size: abs(image.height / image.width - size[1] / size[0])) + elif height is None or width is None: + raise ValueError("height and width must be supplied together.") + else: + height, width = _integer(height, "height"), _integer(width, "width") + if (width, height) not in CANVASES: + raise ValueError(f"Choose a trained 480p canvas (width, height) from {CANVASES}.") + scale = min(width / image.width, height / image.height, 1.0) + output_width = max(int(image.width * scale + 0.5) // 16, 1) * 16 + output_height = max(int(image.height * scale + 0.5) // 16, 1) * 16 + viewpoint = record.get("view_point", "ego_view") + if viewpoint not in VIEWPOINTS: + raise ValueError(f"Unsupported view_point: {viewpoint!r}.") + prompt = record.get("prompt", "") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string.") + extra_args = { + "action_mode": "forward_dynamics", + "domain_name": domain, + "domain_names": domain_names, + "action_normalization": "none", + "action_view_point": viewpoint, + } + # The server defaults to the export's caption format (plain text for the + # unified checkpoint); a record may still pin it explicitly. + if "format_prompt_as_json" in record: + if not isinstance(record["format_prompt_as_json"], bool): + raise ValueError("format_prompt_as_json must be a boolean.") + extra_args["format_prompt_as_json"] = record["format_prompt_as_json"] + metadata = { + "vision_path": str(vision_path), "action_source": str(action_path), + "action_space": UNIFIED_ACTION_SPACE, + "domain_runs": [{"domain_name": name, "rows": [start, end]} for name, start, end in domain_runs(domain_names)], + "num_frames": num_frames, "fps": fps, "seed": seed, + "canvas": {"height": height, "width": width}, + "output": {"height": output_height, "width": output_width}, + **{key: value for key, value in extra_args.items() if key != "domain_names"}, + } + return BimanualInput( + prompt=prompt.strip(), image=image, action=action, domain_names=domain_names, extra_args=extra_args, + num_frames=num_frames, fps=fps, seed=seed, height=height, width=width, + output_height=output_height, output_width=output_width, metadata=metadata, + ) + + +MAX_FORM_FIELD_BYTES = 1024 * 1024 # The server parses each multipart form field within this limit. + + +def encode_extra_params(prepared: BimanualInput) -> str: + """JSON for the ``extra_params`` form field, with the action rows in their shortest float32 form. + + Rows written with Python's double ``repr`` can exceed the server's 1 MiB form-field + limit for long rollouts. The shortest decimal that round-trips each float32 value is + about half the size and reproduces the same tensor on the server; the result is checked + before it is sent. + """ + payload = dict(prepared.extra_args) + try: + import orjson + except ImportError: + payload["action"] = [ + [float(np.format_float_positional(value, unique=True, trim="-")) for value in row] + for row in prepared.action + ] + text = json.dumps(payload, separators=(",", ":")) + else: + payload["action"] = prepared.action + text = orjson.dumps(payload, option=orjson.OPT_SERIALIZE_NUMPY).decode("utf-8") + decoded = np.asarray(json.loads(text)["action"], dtype=np.float32) + if decoded.shape != prepared.action.shape or not np.array_equal(decoded, prepared.action): + raise ValueError("Action serialization must preserve every float32 value.") + size = len(text.encode("utf-8")) + if size > MAX_FORM_FIELD_BYTES: + raise ValueError( + f"The action payload is {size:,} bytes, above the server's {MAX_FORM_FIELD_BYTES:,}-byte multipart-field " + "limit; shorten the rollout or raise the server's form-field limit." + ) + return text + + +def submit_bimanual( + client, prepared: BimanualInput, model: str, output_path: Path, *, timeout: float = 3600, poll_interval: float = 5 +) -> Path: + """Submit multipart image/actions, poll an async job, and stream its video. + + The client is an httpx.Client with the server base_url and optional auth. + timeout bounds job polling; the client's timeout applies to upload/download. + The step count and guidance are fixed by the distilled checkpoint, so the + request does not send them. + """ + if any(not math.isfinite(value) or value <= 0 for value in (timeout, poll_interval)): + raise ValueError("Polling timeout and interval must be positive and finite.") + output_path.parent.mkdir(parents=True, exist_ok=True) + error_path = output_path.with_suffix(".error.txt") + + def check(response): + if response.is_error: + error_path.write_text(response.text) + response.raise_for_status() + + form = { + "model": model, "prompt": prepared.prompt, + "size": f"{prepared.width}x{prepared.height}", + "num_frames": str(prepared.num_frames), "fps": str(prepared.fps), + "seed": str(prepared.seed), + "extra_params": encode_extra_params(prepared), + } + output_path.with_suffix(".request.json").write_text(json.dumps(form, indent=2) + "\n") + output_path.with_suffix(".metadata.json").write_text(json.dumps(prepared.metadata, indent=2) + "\n") + image_buffer = BytesIO() + prepared.image.save(image_buffer, format="PNG") + response = client.post( + "/v1/videos", data=form, + files={"input_reference": ("conditioning.png", image_buffer.getvalue(), "image/png")}, + ) + check(response) + job = response.json() + job_url = f"/v1/videos/{job['id']}" + deadline = time.monotonic() + timeout + print(f"Submitted {job['id']} at {client.base_url.join(job_url)}") + while True: + output_path.with_suffix(".job.json").write_text(json.dumps(job, indent=2) + "\n") + status = job.get("status") + if status in {"failed", "cancelled"}: + error_path.write_text(json.dumps(job, indent=2) + "\n") + raise RuntimeError(f"Video generation {status}: {job.get('error')}") + if status == "completed": + break + if status not in {"queued", "in_progress"}: + raise RuntimeError(f"Unexpected video job status: {job}") + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"Polling timed out; job remains at {client.base_url.join(job_url)}.") + time.sleep(min(poll_interval, remaining)) + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"Polling timed out; job remains at {client.base_url.join(job_url)}.") + response = client.get(job_url, timeout=min(30, remaining)) + check(response) + job = response.json() + temporary = output_path.with_suffix(output_path.suffix + ".part") + try: + with client.stream("GET", job_url + "/content") as response: + if response.is_error: + response.read() + check(response) + with temporary.open("wb") as stream: + for chunk in response.iter_bytes(): + stream.write(chunk) + if temporary.stat().st_size == 0: + raise ValueError("The generated video response was empty.") + temporary.replace(output_path) + finally: + temporary.unlink(missing_ok=True) + return output_path diff --git a/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_sglang.ipynb b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_sglang.ipynb new file mode 100644 index 000000000..e8e37bd0d --- /dev/null +++ b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_sglang.ipynb @@ -0,0 +1,483 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "bimanual-sglang-license", + "metadata": {}, + "source": [ + "" + ] + }, + { + "cell_type": "markdown", + "id": "bimanual-sglang-intro", + "metadata": {}, + "source": [ + "# Cosmos3-Nano-Sim-Bimanual with SGLang\n", + "\n", + "Replay the published `agibot_mixed_028` example of `nvidia/Cosmos3-Nano-Sim-Bimanual` through a persistent SGLang server: AgiBot manipulation, camera movement, then AgiBot manipulation again. Its 900 normalized unified 59D actions produce **901 frames at 30 FPS** with seed **2026** from the published initial image, a 640×360 frame. The server pads the frame into the trained 832×480 canvas for VAE encoding and drops the padded latents, so the output keeps the image's size on the 16-pixel latent grid, 640×352. The notebook uses an HTTP client and local CPU helpers; the SGLang server owns the model and checkpoint configuration.\n", + "\n", + "The model predicts the visual result of supplied actions. It does not choose robot actions or feed generated observations back into a policy.\n", + "\n", + "| Component | Version |\n", + "| --- | --- |\n", + "| SGLang | The published [`lmsysorg/sglang:dev`](https://hub.docker.com/r/lmsysorg/sglang/tags) image, which carries the Cosmos Dreams pipeline from [sgl-project/sglang#39475](https://github.com/sgl-project/sglang/pull/39475); the tag moves, so pull it before starting the server |\n", + "| Model and examples | The Hugging Face revisions selected by `COSMOS3_MODEL_REVISION` and `COSMOS3_EXAMPLE_REVISION`, by default the repository's current default revision |" + ] + }, + { + "cell_type": "markdown", + "id": "bimanual-sglang-contract", + "metadata": {}, + "source": [ + "## 1. Model and runtime contract\n", + "\n", + "The server loads the exported `nvidia/Cosmos3-Nano-Sim-Bimanual` artifact. Its `model_index.json` must select `Cosmos3NanoSimBimanualPipeline`.\n", + "\n", + "This is an **autoregressive, self-forcing DMD-distilled, 480p** checkpoint. SGLang reads the causal rollout and sampling configuration from the export: the chunk size in latent frames, the fixed-step denoising schedule, and the history mode. The start-up log prints the resolved values. Keep the exported `transformer/config.json` and `checkpoint.json` intact.\n", + "\n", + "The `conditioning` contract (schema 5) must be preserved. It defines the 59-value `unified_v1` row that every embodiment shares, the embodiment table (`agibotworld` is domain 15, `camera_pose` is domain 2), the `backward_chunk_anchored_16f` pose convention, and the `global_asinh_unified_v1` normalizer that the published sidecars already apply. The server pads each row to 64 values and routes it to the domain named for that row. Chunk and window sizes count latent frames; one latent frame corresponds to four video frames. The training loader sent prompts as plain text, so the server tokenizes the prompt verbatim unless a request asks for the JSON caption.\n", + "\n", + "The server runs in the published SGLang image, `lmsysorg/sglang:dev`. Pull it once on the server host:\n", + "\n", + "```bash\n", + "docker pull lmsysorg/sglang:dev\n", + "```\n", + "\n", + "The notebook client needs FFmpeg with `ffprobe` for output verification." + ] + }, + { + "cell_type": "markdown", + "id": "bimanual-sglang-server", + "metadata": {}, + "source": [ + "## 2. Start the server\n", + "\n", + "`SGLANG_HF_CACHE` is a host directory for the Hugging Face cache, and `HF_TOKEN` needs read access to the checkpoint. The container runs as root by default, so the setup lines create that directory world-writable; alternatively, start the container with your host user ID. `COSMOS3_MODEL_REVISION` selects the checkpoint revision; leave it unset for the repository's default. Run this in a terminal and leave it running while using the notebook:\n", + "\n", + "```bash\n", + "export SGLANG_HF_CACHE=\"${SGLANG_HF_CACHE:-$HOME/.cache/sglang-huggingface}\"\n", + "mkdir -p \"$SGLANG_HF_CACHE\"\n", + "chmod 777 \"$SGLANG_HF_CACHE\"\n", + "docker run --init --runtime nvidia --gpus all \\\n", + " -v $SGLANG_HF_CACHE:/root/.cache/huggingface \\\n", + " -e HF_TOKEN=$HF_TOKEN \\\n", + " -p 30000:30000 \\\n", + " --ipc=host \\\n", + " lmsysorg/sglang:dev \\\n", + " sglang serve \\\n", + " --model-path nvidia/Cosmos3-Nano-Sim-Bimanual \\\n", + " --revision \"${COSMOS3_MODEL_REVISION:-main}\" \\\n", + " --host 0.0.0.0 --port 30000\n", + "```\n", + "\n", + "The first start downloads the checkpoint into the mounted cache. The HTTP port opens only after the weights are loaded and the warmup requests finish, so wait for the `Uvicorn running` line. The start-up log also prints the resolved checkpoint, the denoising steps per chunk, and the history mode. The notebook uploads the initial image and sends the actions inline, so the server needs no access to the client's files. For a remote notebook, set `SGLANG_BASE_URL`, or forward the port with `ssh -L 30000:127.0.0.1:30000 `." + ] + }, + { + "cell_type": "markdown", + "id": "bimanual-sglang-client", + "metadata": {}, + "source": [ + "## 3. Set up the client\n", + "\n", + "Install these packages in the notebook kernel environment, and authenticate to Hugging Face with read access to `nvidia/Cosmos3-Nano-Sim-Bimanual`; run `hf auth login` or set `HF_TOKEN`:\n", + "\n", + "```bash\n", + "pip install numpy Pillow httpx matplotlib ipython huggingface_hub orjson\n", + "```\n", + "\n", + "`orjson` is optional: it writes the action payload compactly, and without it the helper falls back to the standard library with the same result.\n", + "\n", + "Install FFmpeg with `ffprobe` on the client as well. The local [`bimanual_helpers.py`](bimanual_helpers.py) prepares the image and actions with NumPy and Pillow, then submits HTTP requests. The client does not import an inference runtime or download checkpoint weights; SGLang validates the exported model and conditioning contract on the server.\n", + "\n", + "| Variable | Default | Purpose |\n", + "| --- | --- | --- |\n", + "| `SGLANG_BASE_URL` | `http://127.0.0.1:30000` | Server address |\n", + "| `SGLANG_API_KEY` | unset | Bearer token, if the server was started with `--api-key` |\n", + "| `BIMANUAL_OUTPUT_DIR` | `outputs/sglang/bimanual` under this notebook's folder | Output folder |\n", + "| `COSMOS3_MODEL_REVISION` | `main` | Checkpoint revision; set it to the same value the server was started with |\n", + "| `COSMOS3_EXAMPLE_REVISION` | `main` | Revision of the published manifests, images, and action sidecars |\n", + "| `COSMOS3_BIMANUAL_MANIFEST` | `examples/inference/sim_bimanual_mixed_action_hf.jsonl` | Manifest to read the example row from; `examples/inference/sim_bimanual_unifiedv1_hf.jsonl` holds the single-domain examples |\n", + "| `COSMOS3_BIMANUAL_ACTION` | `agibot_mixed_028_normalized_mixed_action.json` | Basename of the row's action sidecar, which selects the example |\n", + "\n", + "Run the notebook from the cookbook root (`.agents/skills`), the simulation folder, or this `bimanual` folder." + ] + }, + { + "cell_type": "code", + "id": "bimanual-sglang-config", + "metadata": {}, + "source": [ + "from pathlib import Path\n", + "from fractions import Fraction\n", + "from urllib.parse import unquote\n", + "import json\n", + "import math\n", + "import os\n", + "import shutil\n", + "import subprocess\n", + "import sys\n", + "import time\n", + "\n", + "import httpx\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from IPython.display import Video, display\n", + "\n", + "start = Path.cwd().resolve()\n", + "SIMULATION_ROOT = next(\n", + " (\n", + " candidate\n", + " for root in (start, *start.parents)\n", + " for candidate in (root, root / \"cookbooks\" / \"cosmos3\" / \"simulation\", root / \".agents\" / \"skills\" / \"cookbooks\" / \"cosmos3\" / \"simulation\")\n", + " if (candidate / \"bimanual\" / \"bimanual_helpers.py\").is_file()\n", + " ),\n", + " None,\n", + ")\n", + "if SIMULATION_ROOT is None:\n", + " raise FileNotFoundError(\"Run inside .agents/skills, which contains cookbooks/cosmos3/simulation/bimanual/bimanual_helpers.py.\")\n", + "BIMANUAL_ROOT = SIMULATION_ROOT / \"bimanual\"\n", + "if str(BIMANUAL_ROOT) not in sys.path:\n", + " sys.path.insert(0, str(BIMANUAL_ROOT))\n", + "\n", + "MODEL_ID = \"nvidia/Cosmos3-Nano-Sim-Bimanual\"\n", + "MODEL_REVISION = os.environ.get(\"COSMOS3_MODEL_REVISION\", \"main\")\n", + "EXAMPLE_REVISION = os.environ.get(\"COSMOS3_EXAMPLE_REVISION\", \"main\")\n", + "MANIFEST_NAME = os.environ.get(\"COSMOS3_BIMANUAL_MANIFEST\", \"examples/inference/sim_bimanual_mixed_action_hf.jsonl\")\n", + "ACTION_BASENAME = os.environ.get(\"COSMOS3_BIMANUAL_ACTION\", \"agibot_mixed_028_normalized_mixed_action.json\")\n", + "SGLANG_BASE_URL = os.environ.get(\"SGLANG_BASE_URL\", \"http://127.0.0.1:30000\").rstrip(\"/\")\n", + "SGLANG_API_KEY = os.environ.get(\"SGLANG_API_KEY\")\n", + "HEADERS = {\"Authorization\": f\"Bearer {SGLANG_API_KEY}\"} if SGLANG_API_KEY else {}\n", + "HTTP_TIMEOUT = 3600.0\n", + "OUTPUT_DIR = Path(\n", + " os.environ.get(\"BIMANUAL_OUTPUT_DIR\", BIMANUAL_ROOT / \"outputs\" / \"sglang\" / \"bimanual\")\n", + ").expanduser().resolve()\n", + "\n", + "assert shutil.which(\"ffmpeg\") and shutil.which(\"ffprobe\"), \"Install FFmpeg and ffprobe.\"\n", + "OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n", + "print(\"Server:\", SGLANG_BASE_URL, \"(with API key)\" if SGLANG_API_KEY else \"\")\n", + "print(\"Model:\", MODEL_ID, \"@\", MODEL_REVISION)\n", + "print(\"Example:\", MANIFEST_NAME, \"/\", ACTION_BASENAME, \"@\", EXAMPLE_REVISION)\n", + "print(\"Outputs:\", OUTPUT_DIR)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "bimanual-sglang-verify-heading", + "metadata": {}, + "source": [ + "### Verify `/v1/models`\n", + "\n", + "Check that the server is ready and advertises the Bimanual model, then display its model list." + ] + }, + { + "cell_type": "code", + "id": "bimanual-sglang-verify", + "metadata": {}, + "source": [ + "with httpx.Client(base_url=SGLANG_BASE_URL, headers=HEADERS, timeout=30) as client:\n", + " models_response = client.get(\"/v1/models\")\n", + " models_response.raise_for_status()\n", + " models_payload = models_response.json()\n", + "model_ids = [item[\"id\"] for item in models_payload.get(\"data\", []) if item.get(\"id\")]\n", + "if not model_ids:\n", + " raise RuntimeError(f\"The server returned no model IDs: {models_payload}\")\n", + "if MODEL_ID not in model_ids:\n", + " raise RuntimeError(f\"Expected {MODEL_ID}; the server advertises {model_ids}\")\n", + "ACTIVE_MODEL = MODEL_ID\n", + "print(\"active model:\", ACTIVE_MODEL)\n", + "print(json.dumps(models_payload, indent=2))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "bimanual-sglang-imports", + "metadata": {}, + "source": [ + "from huggingface_hub import hf_hub_download\n", + "\n", + "from bimanual_helpers import domain_runs, prepare_bimanual_input, submit_bimanual\n", + "\n", + "# Preparation is local CPU work. The server applies its checkpoint contract." + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "bimanual-sglang-download-heading", + "metadata": {}, + "source": [ + "## 4. Download the published example\n", + "\n", + "The manifest, by default `examples/inference/sim_bimanual_mixed_action_hf.jsonl`, comes from the model repository at `COSMOS3_EXAMPLE_REVISION`. Each row names its initial image and action sidecar by Hugging Face URL; the next cell selects the row whose sidecar is `COSMOS3_BIMANUAL_ACTION`, by default `agibot_mixed_028`, downloads both files into the Hugging Face cache, and keeps the row's frame count, FPS, seed, resolution, action width, and prompt format for the request. The action values are already normalized and validity-masked; the server pads the 59 input channels to the model's 64.\n", + "\n", + "To use your own recording, set these before starting the notebook; the published row still supplies the frame count, FPS, seed, and prompt format:\n", + "\n", + "| Variable | Replaces |\n", + "| --- | --- |\n", + "| `ACTION_PATH` | The sidecar: a JSON object with `action_space: \"normalized_unified_v1\"`, a `900 × 59` `action` array, and `domain_names`, one per row or one for the whole clip. |\n", + "| `VISION_PATH` | The initial observation: an image, or a video whose first frame is used. |\n", + "| `HUMANOID_PROMPT` | The task description, sent as plain text. |" + ] + }, + { + "cell_type": "code", + "id": "bimanual-sglang-download", + "metadata": {}, + "source": [ + "manifest_path = Path(hf_hub_download(repo_id=MODEL_ID, filename=MANIFEST_NAME, revision=EXAMPLE_REVISION))\n", + "rows = [json.loads(line) for line in manifest_path.read_text().splitlines() if line.strip()]\n", + "row = next(item for item in rows if item[\"action_path\"].endswith(\"/\" + ACTION_BASENAME))\n", + "\n", + "\n", + "def download_hosted_asset(url: str) -> Path:\n", + " prefix = f\"https://huggingface.co/{MODEL_ID}/resolve/main/\"\n", + " if not url.startswith(prefix):\n", + " raise ValueError(f\"Expected a published {prefix} URL, got {url!r}\")\n", + " return Path(hf_hub_download(repo_id=MODEL_ID, filename=unquote(url.removeprefix(prefix)), revision=EXAMPLE_REVISION))\n", + "\n", + "\n", + "VISION_PATH = Path(os.environ.get(\"VISION_PATH\") or download_hosted_asset(row[\"vision_path\"])).expanduser().resolve()\n", + "ACTION_PATH = Path(os.environ.get(\"ACTION_PATH\") or download_hosted_asset(row[\"action_path\"])).expanduser().resolve()\n", + "PROMPT = os.environ.get(\"HUMANOID_PROMPT\", row[\"prompt\"])\n", + "record = {key: row[key] for key in (\"num_frames\", \"fps\", \"seed\", \"resolution\", \"raw_action_dim\", \"format_prompt_as_json\") if key in row}\n", + "record.update(prompt=PROMPT, vision_path=str(VISION_PATH), action_path=str(ACTION_PATH))\n", + "\n", + "for path in (ACTION_PATH, VISION_PATH):\n", + " if not path.is_file():\n", + " raise FileNotFoundError(f\"Missing input; check the download or your override: {path}\")\n", + "print(f\"{row['name']}: {record['num_frames']} frames at {record['fps']} FPS, seed {record['seed']}\")\n", + "print(\"Actions:\", ACTION_PATH)\n", + "print(\"Initial image:\", VISION_PATH)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "bimanual-sglang-actions-heading", + "metadata": {}, + "source": [ + "### Action structure\n", + "\n", + "The sidecar has `action_space: \"normalized_unified_v1\"`, a `900×59` action array, and one domain name per row. The first 300 rows use `agibotworld`, the next 300 use `camera_pose`, and the last 300 return to `agibotworld`. The checkpoint maps these to domain IDs 15, 2, and 15. The initial image plus 900 action rows produce 901 frames. The 59 columns are grouped as follows (Python slices exclude the end index):\n", + "\n", + "| Columns | Contents |\n", + "| --- | --- |\n", + "| `0:9` | Ego/camera translation and rotation |\n", + "| `9:18` | Right-wrist translation and rotation |\n", + "| `18:33` | Right fingertips |\n", + "| `33` | Right-gripper width |\n", + "| `34:43` | Left-wrist translation and rotation |\n", + "| `43:58` | Left fingertips |\n", + "| `58` | Left-gripper width |\n", + "\n", + "Poses use the checkpoint's `backward_chunk_anchored_16f` convention and six rotation values (the first two matrix columns). Values are already globally asinh-normalized; unavailable channels are masked to zero after normalization. Do not normalize again: the request sends `action_normalization=\"none\"`, the server pads 59 channels to 64, and camera rows may only carry values in columns `0:9`." + ] + }, + { + "cell_type": "code", + "id": "bimanual-sglang-actions", + "metadata": {}, + "source": [ + "sidecar = json.loads(ACTION_PATH.read_text(encoding=\"utf-8\"))\n", + "assert sidecar[\"action_space\"] == \"normalized_unified_v1\", sidecar.get(\"action_space\")\n", + "actions = sidecar[\"action\"]\n", + "domain_names = sidecar[\"domain_names\"]\n", + "if isinstance(domain_names, str):\n", + " domain_names = [domain_names] * len(actions)\n", + "assert isinstance(actions, list) and len(actions) == record[\"num_frames\"] - 1, \"expected num_frames - 1 action rows\"\n", + "assert len(domain_names) == len(actions), \"one domain name is required per action row\"\n", + "for row_index, values in enumerate(actions):\n", + " assert isinstance(values, list) and len(values) == 59, f\"row {row_index}: expected 59 values\"\n", + " for value in values:\n", + " assert isinstance(value, (int, float)) and not isinstance(value, bool), f\"row {row_index}: non-numeric value\"\n", + " assert math.isfinite(value), f\"row {row_index}: non-finite value\"\n", + "print(f\"Validated {len(actions)}×{len(actions[0])} finite actions from {ACTION_PATH}\")\n", + "print(\"Domain sequence:\", [(name, end - start) for name, start, end in domain_runs(domain_names)])" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "bimanual-sglang-prepare-heading", + "metadata": {}, + "source": [ + "### Prepare the input\n", + "\n", + "The initial image provides the starting observation. The 600 robot actions and 300 camera actions drive the following frames, producing a 901-frame video. The helper keeps the rows as they are and sends one embodiment name per row; the server applies no further normalization. It also reports the canvas the server encodes and the expected output size: for the published 640×360 image, 832×480 and 640×352." + ] + }, + { + "cell_type": "code", + "id": "bimanual-sglang-prepare", + "metadata": {}, + "source": [ + "humanoid_input = prepare_bimanual_input(record, base_dir=ACTION_PATH.parent)\n", + "assert humanoid_input.action.shape == (record[\"num_frames\"] - 1, 59)\n", + "assert humanoid_input.extra_args[\"action_normalization\"] == \"none\"\n", + "assert humanoid_input.num_frames == len(humanoid_input.domain_names) + 1\n", + "print(f\"Prepared {row['name']}: {humanoid_input.image.width}×{humanoid_input.image.height} input, \"\n", + " f\"{humanoid_input.num_frames} frames at {humanoid_input.fps} FPS, seed {humanoid_input.seed}\")\n", + "print(\"Canvas the server encodes:\", (humanoid_input.width, humanoid_input.height))\n", + "print(\"Expected output size:\", (humanoid_input.output_width, humanoid_input.output_height))\n", + "print(\"Domain runs:\", domain_runs(humanoid_input.domain_names))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "bimanual-sglang-preview-heading", + "metadata": {}, + "source": [ + "### Preview the inputs\n", + "\n", + "The image is the starting observation. The heatmap shows all 59 normalized action channels over the 30-second sequence, grouped as in the table above. Dashed lines mark the domain switches at 10 and 20 seconds." + ] + }, + { + "cell_type": "code", + "id": "bimanual-sglang-preview", + "metadata": {}, + "source": [ + "action_values = humanoid_input.action\n", + "groups = [\n", + " (0, 9, \"Ego/camera pose\"), (9, 18, \"Right-wrist pose\"),\n", + " (18, 33, \"Right fingertips\"), (33, 34, \"Right gripper\"),\n", + " (34, 43, \"Left-wrist pose\"), (43, 58, \"Left fingertips\"),\n", + " (58, 59, \"Left gripper\"),\n", + "]\n", + "fig, (image_ax, action_ax) = plt.subplots(1, 2, figsize=(14, 5), layout=\"constrained\")\n", + "image_ax.imshow(humanoid_input.image)\n", + "image_ax.set(title=\"Humanoid: initial observation\")\n", + "image_ax.axis(\"off\")\n", + "limit = max(float(np.abs(action_values).max()), 1e-6)\n", + "heatmap = action_ax.imshow(\n", + " action_values.T, aspect=\"auto\", cmap=\"RdBu_r\", vmin=-limit, vmax=limit,\n", + " extent=(0, len(action_values) / humanoid_input.fps, 58.5, -0.5),\n", + ")\n", + "action_ax.set_yticks([(start + end - 1) / 2 for start, end, _ in groups],\n", + " [label for _, _, label in groups])\n", + "for _, end, _ in groups[:-1]:\n", + " action_ax.axhline(end - 0.5, color=\"black\", linewidth=0.5)\n", + "for _, start, _ in domain_runs(humanoid_input.domain_names)[1:]:\n", + " action_ax.axvline(start / humanoid_input.fps, color=\"black\", linestyle=\"--\", linewidth=1)\n", + "action_ax.set(title=f\"{row['name']}: {action_values.shape[0]} × {action_values.shape[1]}\", xlabel=\"Time (seconds)\")\n", + "fig.colorbar(heatmap, ax=action_ax, label=\"Checkpoint-normalized value\", shrink=0.8)\n", + "display(fig)\n", + "plt.close(fig)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "bimanual-sglang-generate-heading", + "metadata": {}, + "source": [ + "## 5. Generate with `POST /v1/videos`\n", + "\n", + "The request uploads the initial image as `input_reference` and sends the complete action sequence with its per-row domain names in JSON-encoded `extra_params`. The `size` field names the 832×480 canvas: the server pads the 640×360 image at the bottom and right to encode it, then drops the padded latent rows and columns, so the transformer and the output use only the 640×352 content region. Per-row domain names switch from AgiBot to camera and back within this single request. The step count and guidance are fixed by the distilled checkpoint, so the request does not send them. SGLang returns a job ID; the helper polls `GET /v1/videos/{id}` and downloads `/content` when generation completes.\n", + "\n", + "The helper writes the action rows in their shortest float32 form and checks that the payload stays within the server's 1 MiB multipart-field limit; a longer rollout than that limit allows fails before anything is sent. It waits up to `HTTP_TIMEOUT` seconds and saves request and job JSON sidecars beside the MP4. HTTP and generation errors are saved to an `.error.txt` sidecar. Rerunning this cell creates a new job." + ] + }, + { + "cell_type": "code", + "id": "bimanual-sglang-generate", + "metadata": {}, + "source": [ + "prepared = humanoid_input\n", + "output = OUTPUT_DIR / row[\"name\"] / f\"{row['name']}.mp4\"\n", + "started = time.monotonic()\n", + "with httpx.Client(base_url=SGLANG_BASE_URL, headers=HEADERS, timeout=HTTP_TIMEOUT) as client:\n", + " output = submit_bimanual(\n", + " client, prepared, ACTIVE_MODEL, output,\n", + " timeout=HTTP_TIMEOUT, poll_interval=5,\n", + " )\n", + "elapsed_seconds = time.monotonic() - started\n", + "print(f\"{row['name']}: {elapsed_seconds:.1f}s — {output}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "bimanual-sglang-verify-output-heading", + "metadata": {}, + "source": [ + "## 6. Verify and view the output\n", + "\n", + "Decode the generated video and check that it contains the requested frame count at the requested FPS and the expected output size. For the published image, the output keeps the 640-pixel width; its 360-pixel height becomes 352, the largest multiple of the 16-pixel latent stride, because the padded latent rows are dropped rather than decoded.\n", + "\n", + "Review the three phases: robot manipulation, a camera position/orientation loop, then resumed robot manipulation. The domain switches are at approximately 10 and 20 seconds." + ] + }, + { + "cell_type": "code", + "id": "bimanual-sglang-verify-output", + "metadata": {}, + "source": [ + "probe = subprocess.run([\n", + " \"ffprobe\", \"-v\", \"error\", \"-count_frames\", \"-select_streams\", \"v:0\",\n", + " \"-show_entries\", \"stream=width,height,avg_frame_rate,nb_read_frames\", \"-of\", \"json\", str(output),\n", + "], check=True, capture_output=True, text=True)\n", + "stream = json.loads(probe.stdout)[\"streams\"][0]\n", + "assert int(stream[\"nb_read_frames\"]) == prepared.num_frames, stream\n", + "assert Fraction(stream[\"avg_frame_rate\"]) == Fraction(str(prepared.fps)), stream\n", + "assert (int(stream[\"width\"]), int(stream[\"height\"])) == (prepared.output_width, prepared.output_height), stream\n", + "subprocess.run([\"ffmpeg\", \"-nostdin\", \"-v\", \"error\", \"-xerror\", \"-i\", str(output), \"-f\", \"null\", \"-\"], check=True)\n", + "print(stream)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "id": "bimanual-sglang-play", + "metadata": {}, + "source": [ + "display(Video(str(output), embed=True, html_attributes=\"controls loop\"))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "bimanual-sglang-next-steps", + "metadata": {}, + "source": [ + "## 7. Next steps\n", + "\n", + "- **Run the other published examples:** set `COSMOS3_BIMANUAL_MANIFEST` to `examples/inference/sim_bimanual_unifiedv1_hf.jsonl` and `COSMOS3_BIMANUAL_ACTION` to a sidecar basename from it, such as `agibot_001_normalized_unifiedv1.json`, then restart the notebook.\n", + "- **Use your own recording:** set `ACTION_PATH` to a `normalized_unified_v1` sidecar, `VISION_PATH` to the initial image or a video whose first frame is the initial observation, and `HUMANOID_PROMPT` to the task description, then restart the notebook. Keep `num_frames - 1` equal to the number of action rows and divisible by four, and do not normalize the values again. The [Bimanual guide](bimanual.md#download-the-model-and-examples) describes the sidecar contract.\n", + "- **Validate further:** apply the [shared Simulation checklist](../README.md#validate-a-rollout). Inspect the complete rollout for frozen frames, discontinuities, and disagreement with the supplied actions, including both domain transitions around 10 and 20 seconds. Use a new output directory when changing the checkpoint, inputs, or server settings.\n", + "- **Other runtimes:** the [vLLM-Omni notebook](run_bimanual_with_vllm_omni.ipynb) serves the same published example through the synchronous Videos API, and the [Bimanual guide](bimanual.md#run-inference) runs it natively with Cosmos Framework." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_vllm_omni.ipynb b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_vllm_omni.ipynb index be2611559..c43918487 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_vllm_omni.ipynb +++ b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_vllm_omni.ipynb @@ -364,7 +364,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Inspect the complete rollout for frozen frames, discontinuities, and disagreement with the supplied actions. In the default mixed-domain example, inspect the transitions around 10 and 20 seconds. For other inputs, derive transition times from their action rows and FPS. Use a new output directory when changing the checkpoint, inputs, or server settings." + "Inspect the complete rollout for frozen frames, discontinuities, and disagreement with the supplied actions. In the default mixed-domain example, inspect the transitions around 10 and 20 seconds. For other inputs, derive transition times from their action rows and FPS. Use a new output directory when changing the checkpoint, inputs, or server settings. The [SGLang notebook](run_bimanual_with_sglang.ipynb) runs the published mixed-action example through SGLang's asynchronous Videos API." ] } ], diff --git a/.agents/skills/cookbooks/cosmos3/simulation/transfer/run_transfer_with_sglang.ipynb b/.agents/skills/cookbooks/cosmos3/simulation/transfer/run_transfer_with_sglang.ipynb new file mode 100644 index 000000000..6d6170970 --- /dev/null +++ b/.agents/skills/cookbooks/cosmos3/simulation/transfer/run_transfer_with_sglang.ipynb @@ -0,0 +1,381 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "transfer-sglang-00", + "metadata": {}, + "source": [ + "" + ] + }, + { + "cell_type": "markdown", + "id": "transfer-sglang-01", + "metadata": {}, + "source": [ + "# Sim-Transfer with SGLang\n", + "\n", + "Cosmos3-Nano-Sim-Transfer generates a **480p video from a caption and a control video**. Each sample uses one of four control modalities:\n", + "\n", + "- **Depth** guides scene geometry and depth structure.\n", + "- **Edge** guides object boundaries and outlines.\n", + "- **Blur** guides coarse appearance, layout, and motion.\n", + "- **Segmentation** (`seg`) guides the layout of regions in the scene.\n", + "\n", + "The caption describes the desired scene and appearance, while the control video guides its spatial structure and motion. The distilled causal model generates latents autoregressively with **four denoising steps**, retaining control and generated-latent K/V history across the rollout.\n", + "\n", + "This notebook runs the four published HF `example2` controls of the [Transfer cookbook](transfer.md) through one SGLang server at **121 frames, 832 × 480, and 30 FPS**, with the saved `example2` prompt. It checks each output's frame count, dimensions, and FPS, then plays the generated videos.\n", + "\n", + "| Component | Version |\n", + "| --- | --- |\n", + "| SGLang | The published [`lmsysorg/sglang:dev`](https://hub.docker.com/r/lmsysorg/sglang/tags) image, which carries the Cosmos Dreams Transfer pipeline from [sgl-project/sglang#39475](https://github.com/sgl-project/sglang/pull/39475); the tag moves, so pull it before starting the server |\n", + "| `nvidia/Cosmos3-Nano-Sim-Transfer` | [`main`](https://huggingface.co/nvidia/Cosmos3-Nano-Sim-Transfer/tree/main) |\n", + "| `example2` sample | [`main`](https://huggingface.co/nvidia/Cosmos3-Nano-Sim-Transfer/tree/main/assets/sample_data/example2) of the same repository, under `assets/sample_data/example2` |" + ] + }, + { + "cell_type": "markdown", + "id": "transfer-sglang-02", + "metadata": {}, + "source": [ + "## Shared setup\n", + "\n", + "Install the notebook dependencies, and authenticate to Hugging Face with read access to `nvidia/Cosmos3-Nano-Sim-Transfer`, which publishes both the checkpoint and the sample; run `hf auth login` or set `HF_TOKEN`:\n", + "\n", + "```bash\n", + "pip install requests huggingface_hub matplotlib pillow ipython\n", + "```\n", + "\n", + "Install FFmpeg on the client host; the package also includes FFprobe.\n", + "\n", + "**Inputs:** the notebook downloads the saved `example2` prompt and its four precomputed controls from the same Hugging Face repository as the model. It writes them under `simulation/outputs/sglang/transfer/inputs`, inside the cookbook root, because the server does not receive uploads: a Sim-Transfer request names each control by absolute `control_path`, and the Docker command below mounts the cookbook root at its host path. Run the notebook from that checkout on the server host, or use the same absolute checkout path on both hosts. Later runs reuse the downloaded files." + ] + }, + { + "cell_type": "code", + "id": "transfer-sglang-03", + "metadata": {}, + "source": [ + "from fractions import Fraction\n", + "from pathlib import Path\n", + "import json\n", + "import os\n", + "import shutil\n", + "import subprocess\n", + "import time\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import requests\n", + "from IPython.display import Video, display\n", + "\n", + "MODEL_ID = \"nvidia/Cosmos3-Nano-Sim-Transfer\"\n", + "MODEL_REVISION = \"main\"\n", + "ASSET_REVISION = \"main\"\n", + "BASE_URL = os.environ.get(\"SGLANG_BASE_URL\", \"http://127.0.0.1:30000\").rstrip(\"/\")\n", + "assert shutil.which(\"ffmpeg\") and shutil.which(\"ffprobe\"), \"Install FFmpeg and ffprobe.\"\n", + "root_override = os.environ.get(\"COSMOS_COOKBOOK_ROOT\")\n", + "candidates = [Path(root_override)] if root_override else [Path.cwd(), *Path.cwd().parents]\n", + "COOKBOOK_ROOT = next((p for p in candidates if (p / \"cookbooks/cosmos3/simulation/transfer/transfer.md\").is_file()), None)\n", + "assert COOKBOOK_ROOT, \"Run inside .agents/skills or set COSMOS_COOKBOOK_ROOT to it.\"\n", + "COOKBOOK_ROOT = COOKBOOK_ROOT.resolve()\n", + "SIM_ROOT = COOKBOOK_ROOT / \"cookbooks/cosmos3/simulation\"\n", + "OUTPUT_DIR = Path(os.environ.get(\"TRANSFER_OUTPUT_DIR\", SIM_ROOT / \"outputs\" / \"sglang\" / \"transfer\")).expanduser().resolve()\n", + "OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n", + "print(\"Cookbook root:\", COOKBOOK_ROOT)\n", + "print(\"Checkpoint:\", MODEL_ID, \"@\", MODEL_REVISION)\n", + "print(\"Inputs:\", MODEL_ID, \"@\", ASSET_REVISION)\n", + "print(\"Outputs:\", OUTPUT_DIR)\n", + "if not OUTPUT_DIR.is_relative_to(COOKBOOK_ROOT):\n", + " print(f\"Note: {OUTPUT_DIR} is outside the cookbook root; the server must see it at the same path.\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "transfer-sglang-04", + "metadata": {}, + "source": [ + "### Start the server\n", + "\n", + "The server runs in the published SGLang image, `lmsysorg/sglang:dev`, and loads the exported `nvidia/Cosmos3-Nano-Sim-Transfer` artifact from `main`, matching `MODEL_REVISION` above. One server runs all four examples. `SGLANG_HF_CACHE` is a host directory for the Hugging Face cache, and `HF_TOKEN` needs read access to the checkpoint. The container runs as root by default, so the setup lines create that directory world-writable; alternatively, start the container with your host user ID. Run this in a terminal and leave it running while using the notebook:\n", + "\n", + "```bash\n", + "docker pull lmsysorg/sglang:dev\n", + "export COSMOS_ROOT=/absolute/path/to/cosmos-framework/.agents/skills\n", + "export SGLANG_HF_CACHE=\"${SGLANG_HF_CACHE:-$HOME/.cache/sglang-huggingface}\"\n", + "mkdir -p \"$SGLANG_HF_CACHE\"\n", + "chmod 777 \"$SGLANG_HF_CACHE\"\n", + "docker run --init --runtime nvidia --gpus all \\\n", + " -v $SGLANG_HF_CACHE:/root/.cache/huggingface \\\n", + " -v \"${COSMOS_ROOT}:${COSMOS_ROOT}\" \\\n", + " -e HF_TOKEN=$HF_TOKEN \\\n", + " -p 30000:30000 \\\n", + " --ipc=host \\\n", + " lmsysorg/sglang:dev \\\n", + " sglang serve \\\n", + " --model-path nvidia/Cosmos3-Nano-Sim-Transfer \\\n", + " --revision main \\\n", + " --host 0.0.0.0\n", + "```\n", + "\n", + "The first start downloads the checkpoint into the mounted cache. The HTTP port opens only after the weights are loaded and the warmup requests finish, so wait for the `Uvicorn running` line. The container mounts the cookbook root (`.agents/skills`) at its host path because the notebook sends each control video as a server-visible `control_path` rather than uploading it; the downloaded `example2` controls live under that root in the output folder. The next cell waits until `/v1/models` lists the model." + ] + }, + { + "cell_type": "code", + "id": "transfer-sglang-05", + "metadata": {}, + "source": [ + "deadline = time.monotonic() + 1800\n", + "while True:\n", + " try:\n", + " response = requests.get(f\"{BASE_URL}/v1/models\", timeout=10)\n", + " response.raise_for_status()\n", + " models = response.json()\n", + " break\n", + " except requests.RequestException:\n", + " if time.monotonic() >= deadline:\n", + " raise TimeoutError(\"Model server did not become ready; inspect the SGLang container log.\")\n", + " time.sleep(3)\n", + "assert MODEL_ID in [m[\"id\"] for m in models[\"data\"]], models\n", + "print(\"Server ready:\", BASE_URL)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "transfer-sglang-06", + "metadata": {}, + "source": [ + "## Prepare the input\n", + "\n", + "The four precomputed `example2` controls and the saved appearance prompt come from `nvidia/Cosmos3-Nano-Sim-Transfer`, under `assets/sample_data/example2`. Each request uses **121 frames, 832 × 480, 30 FPS, and seed 2025**.\n", + "\n", + "| Control | Precomputed HF asset | What it preserves |\n", + "| --- | --- | --- |\n", + "| Depth | `example2_depth.mp4` | Scene geometry and depth structure |\n", + "| Edge | `example2_edge.mp4` | Object boundaries and outlines |\n", + "| Blur | `example2_blur.mp4` | Coarse appearance, layout, and motion |\n", + "| Segmentation | `example2_segmentation.mp4` | Semantic region layout |\n", + "\n", + "The next cell downloads the files under the output folder, compacts the saved prompt into one JSON string, and probes every control for its frame count and rate. To derive edge or blur controls from your own RGB video instead, follow the [Cosmos Framework workflow](transfer.md#prepare-the-input)." + ] + }, + { + "cell_type": "code", + "id": "transfer-sglang-07", + "metadata": {}, + "source": [ + "from huggingface_hub import snapshot_download\n", + "\n", + "HINTS = (\"depth\", \"edge\", \"blur\", \"seg\")\n", + "NUM_FRAMES, FPS, WIDTH, HEIGHT, SEED = 121, 30, 832, 480, 2025\n", + "FILES = {\"depth\": \"example2_depth.mp4\", \"edge\": \"example2_edge.mp4\", \"blur\": \"example2_blur.mp4\", \"seg\": \"example2_segmentation.mp4\"}\n", + "\n", + "\n", + "def require_file(path: Path) -> Path:\n", + " path = Path(path).resolve()\n", + " if not path.is_file() or path.stat().st_size == 0:\n", + " raise FileNotFoundError(f\"Missing or empty input: {path}\")\n", + " return path\n", + "\n", + "\n", + "def probe_video(path):\n", + " result = subprocess.run([\"ffprobe\", \"-v\", \"error\", \"-count_frames\", \"-select_streams\", \"v:0\", \"-show_entries\", \"stream=width,height,avg_frame_rate,nb_read_frames\", \"-of\", \"json\", str(path)], check=True, capture_output=True, text=True)\n", + " return json.loads(result.stdout)[\"streams\"][0]\n", + "\n", + "\n", + "sample_root = Path(snapshot_download(\n", + " MODEL_ID, revision=ASSET_REVISION, allow_patterns=[\"assets/sample_data/example2/*\"], local_dir=OUTPUT_DIR / \"inputs\",\n", + ")) / \"assets/sample_data/example2\"\n", + "prompt = json.dumps(json.loads(require_file(sample_root / \"example2_upsampled_prompt.json\").read_text()))\n", + "cases = {}\n", + "for hint in HINTS:\n", + " control = require_file(sample_root / FILES[hint])\n", + " metadata = probe_video(control)\n", + " assert int(metadata[\"nb_read_frames\"]) >= NUM_FRAMES, metadata\n", + " assert abs(float(Fraction(metadata[\"avg_frame_rate\"])) - FPS) < 0.1, \"Use controls at the recipe FPS (29.97/30 accepted).\"\n", + " cases[hint] = dict(name=f\"example2_{hint}\", hint=hint, control=control, prompt=prompt, num_frames=NUM_FRAMES,\n", + " width=WIDTH, height=HEIGHT, fps=FPS, seed=SEED, input_metadata=metadata)\n", + " print(f\"{hint}: {NUM_FRAMES} frames, {WIDTH}×{HEIGHT}, {FPS} FPS, seed {SEED}; control {metadata['width']}×{metadata['height']}, \"\n", + " f\"{metadata['nb_read_frames']} source frames; prompt {len(prompt):,} characters\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "transfer-sglang-08", + "metadata": {}, + "source": [ + "### Preview controls\n", + "\n", + "First, middle, and final requested frames of each precomputed control." + ] + }, + { + "cell_type": "code", + "id": "transfer-sglang-09", + "metadata": {}, + "source": [ + "def read_frame(path, frame_index):\n", + " from io import BytesIO\n", + " from PIL import Image\n", + " result = subprocess.run([\n", + " \"ffmpeg\", \"-nostdin\", \"-v\", \"error\", \"-i\", str(path),\n", + " \"-vf\", rf\"select=eq(n\\,{frame_index})\", \"-frames:v\", \"1\", \"-f\", \"image2pipe\", \"-vcodec\", \"png\", \"-\",\n", + " ], check=True, capture_output=True)\n", + " return Image.open(BytesIO(result.stdout)).convert(\"RGB\")\n", + "\n", + "fig, axes = plt.subplots(len(cases), 3, figsize=(12, 11), layout=\"constrained\")\n", + "for row, case in zip(axes, cases.values()):\n", + " for ax, frame_index in zip(row, (0, case[\"num_frames\"] // 2, case[\"num_frames\"] - 1)):\n", + " ax.imshow(read_frame(case[\"control\"], frame_index))\n", + " ax.set_title(f\"{case['hint']} — frame {frame_index}\")\n", + " ax.axis(\"off\")\n", + "display(fig)\n", + "plt.close(fig)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "transfer-sglang-10", + "metadata": {}, + "source": [ + "## Generate the four examples\n", + "\n", + "Generate a video from each caption and control pair using four denoising steps, with prompt emphasis on. Each example is one asynchronous job that saves an MP4 in `OUTPUT_DIR`.\n", + "\n", + "The notebook submits each job to `POST /v1/videos`, polls `GET /v1/videos/{id}` every five seconds, and downloads `GET /v1/videos/{id}/content` once the job completes. `GENERATION_TIMEOUT` limits the wait to 30 minutes per example. The latest job response is saved beside each video as `.job.json`; failed or cancelled jobs and HTTP errors show their error details." + ] + }, + { + "cell_type": "code", + "id": "transfer-sglang-11", + "metadata": {}, + "source": [ + "GENERATION_TIMEOUT = 1800 # Maximum job wait per example, in seconds.\n", + "POLL_INTERVAL = 10\n", + "\n", + "outputs = {}\n", + "elapsed_seconds = {}\n", + "for name, case in cases.items():\n", + " extra_params = dict(emphasize_control_in_prompt=True)\n", + " form = {\n", + " \"model\": MODEL_ID, \"prompt\": case[\"prompt\"],\n", + " \"size\": f\"{case['width']}x{case['height']}\",\n", + " \"num_frames\": str(case[\"num_frames\"]), \"fps\": str(case[\"fps\"]),\n", + " \"seed\": str(case[\"seed\"]), \"num_inference_steps\": \"4\",\n", + " \"guidance_scale\": \"1.0\",\n", + " \"control_hint\": case[\"hint\"],\n", + " \"control_path\": str(case[\"control\"]),\n", + " \"extra_params\": json.dumps(extra_params),\n", + " }\n", + " started = time.monotonic()\n", + " response = requests.post(\n", + " f\"{BASE_URL}/v1/videos\",\n", + " files=[(key, (None, value)) for key, value in form.items()],\n", + " headers={\"Accept\": \"application/json\"}, timeout=60,\n", + " )\n", + " if not response.ok:\n", + " raise RuntimeError(f\"{name}: submission HTTP {response.status_code}\\n{response.text[:4000]}\")\n", + " job = response.json()\n", + " job_id = job.get(\"id\")\n", + " if not job_id:\n", + " raise RuntimeError(f\"{name}: no video job ID in response: {job}\")\n", + " job_path = OUTPUT_DIR / f\"{name}.job.json\"\n", + " deadline = time.monotonic() + GENERATION_TIMEOUT\n", + " while True:\n", + " job_path.write_text(json.dumps(job, indent=2))\n", + " status = job.get(\"status\")\n", + " print(f\"{name}: {job_id} — {status} ({job.get('progress', 0)}%)\")\n", + " if status == \"completed\":\n", + " break\n", + " if status in {\"failed\", \"cancelled\", \"canceled\"}:\n", + " raise RuntimeError(f\"{name}: generation {status}\\n{json.dumps(job, indent=2)}\")\n", + " remaining = deadline - time.monotonic()\n", + " if remaining <= 0:\n", + " raise TimeoutError(f\"{name}: timed out waiting for job {job_id}; see {job_path}\")\n", + " time.sleep(min(POLL_INTERVAL, remaining))\n", + " remaining = deadline - time.monotonic()\n", + " if remaining <= 0:\n", + " raise TimeoutError(f\"{name}: timed out waiting for job {job_id}; see {job_path}\")\n", + " response = requests.get(\n", + " f\"{BASE_URL}/v1/videos/{job_id}\", timeout=min(30, remaining),\n", + " )\n", + " if not response.ok:\n", + " raise RuntimeError(f\"{name}: polling HTTP {response.status_code}\\n{response.text[:4000]}\")\n", + " job = response.json()\n", + "\n", + " response = requests.get(\n", + " f\"{BASE_URL}/v1/videos/{job_id}/content\",\n", + " headers={\"Accept\": \"video/mp4\"}, timeout=300,\n", + " )\n", + " if not response.ok:\n", + " raise RuntimeError(f\"{name}: download HTTP {response.status_code}\\n{response.text[:4000]}\")\n", + " if not response.headers.get(\"content-type\", \"\").startswith(\"video/mp4\") or not response.content:\n", + " raise RuntimeError(f\"{name}: expected a nonempty MP4 response, got {response.headers}\")\n", + " output = OUTPUT_DIR / f\"{name}.mp4\"\n", + " temporary_output = output.with_suffix(\".mp4.part\")\n", + " try:\n", + " temporary_output.write_bytes(response.content)\n", + " temporary_output.replace(output)\n", + " finally:\n", + " temporary_output.unlink(missing_ok=True)\n", + " outputs[name] = output\n", + " elapsed_seconds[name] = time.monotonic() - started\n", + " print(f\"{name}: {elapsed_seconds[name]:.1f}s — {output}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "transfer-sglang-12", + "metadata": {}, + "source": [ + "## View generated videos\n", + "\n", + "Play the depth, edge, blur, and segmentation results below." + ] + }, + { + "cell_type": "code", + "id": "transfer-sglang-13", + "metadata": {}, + "source": [ + "for name, output in outputs.items():\n", + " print(name)\n", + " display(Video(str(output), embed=True, width=640, html_attributes=\"controls loop\"))" + ], + "execution_count": null, + "outputs": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.10" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/.agents/skills/cookbooks/cosmos3/simulation/transfer/run_transfer_with_vllm_omni.ipynb b/.agents/skills/cookbooks/cosmos3/simulation/transfer/run_transfer_with_vllm_omni.ipynb index c083cf5b2..14bb07964 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/transfer/run_transfer_with_vllm_omni.ipynb +++ b/.agents/skills/cookbooks/cosmos3/simulation/transfer/run_transfer_with_vllm_omni.ipynb @@ -685,7 +685,8 @@ "\n", "- **Use your own control:** replace one path in `files` or build a case with your caption and absolute control path. To derive edge or blur from RGB, follow the [Cosmos Framework preprocessing](transfer.md#prepare-the-input).\n", "- **Validate further:** apply the [shared Simulation checklist](../README.md#validate-a-rollout) and the [Sim-Transfer checks](transfer.md#validate-the-rollout) for control adherence and temporal consistency.\n", - "- **Cosmos Framework:** use the native [Hugging Face example2 recipe](transfer.md#cosmos-framework) for direct-resize inference and locally computed edge/blur controls." + "- **Cosmos Framework:** use the native [Hugging Face example2 recipe](transfer.md#cosmos-framework) for direct-resize inference and locally computed edge/blur controls.\n", + "- **SGLang:** the [SGLang notebook](run_transfer_with_sglang.ipynb) runs the same four `example2` controls through SGLang's asynchronous Videos API." ] } ], diff --git a/.github/workflows/validate-cookbooks.yml b/.github/workflows/validate-cookbooks.yml index 817291cca..06cfd4159 100644 --- a/.github/workflows/validate-cookbooks.yml +++ b/.github/workflows/validate-cookbooks.yml @@ -116,7 +116,7 @@ jobs: - name: Check Sim-Auto request and HTTP helpers run: | - ruff check --output-format=github --select=E9,F63,F7,F82 "$AUTO/helpers/sim_auto_vllm_omni_helpers.py" "$AUTO/tests/test_sim_auto_helpers.py" + ruff check --output-format=github --select=E9,F63,F7,F82 "$AUTO/helpers/sim_auto_vllm_omni_helpers.py" "$AUTO/helpers/sim_auto_sglang_helpers.py" "$AUTO/tests/test_sim_auto_helpers.py" python -m unittest discover -s "$AUTO/tests" -p 'test_sim_auto_helpers.py' -v - name: Check media timelines with system FFmpeg From f772755f1cc70414e8023556039b10cb10b3c5ef Mon Sep 17 00:00:00 2001 From: Kedi Wu Date: Wed, 21 Oct 2026 13:52:05 -0700 Subject: [PATCH 2/5] Document multi-GPU Ulysses sequence parallelism for the Auto SGLang servers The Sim-Auto and Transfer-Auto SGLang notebooks now say how to use several GPUs: add --num-gpus N --ulysses-degree N --sp-degree N to sglang serve. Both pipelines accept Ulysses sequence parallelism only, the published checkpoints have 8 key/value heads, so N is 2, 4, or 8, and the server never detects the GPU count on its own. Sim-Auto may combine it with --tp-size under the same divisibility rule; Transfer-Auto may add --enable-cfg-parallel for the two guidance branches. This also drops the outdated claim that Sim-Auto rejects sequence parallelism. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Kedi Wu --- .../cosmos3/simulation/auto/run_auto_with_sglang.ipynb | 2 +- .../cosmos3/simulation/auto/run_sim-auto_with_sglang.ipynb | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_sglang.ipynb b/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_sglang.ipynb index 5754caa99..34ce1213c 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_sglang.ipynb +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_sglang.ipynb @@ -94,7 +94,7 @@ "| `-e HF_TOKEN=$HF_TOKEN` | Read access to the gated checkpoint. |\n", "| `--revision` | Loads the checkpoint from `main`. The first start downloads it into the mounted cache. |\n", "| `docker pull` | The `dev` tag moves with SGLang's main branch; pulling before each start picks up the current image. |\n", - "| `--gpus all` | Select one GPU with `--gpus '\"device=0\"'`. |\n", + "| `--gpus all` | Makes every GPU visible; select one with `--gpus '\"device=0\"'`. With N GPUs, add `--num-gpus N --ulysses-degree N --sp-degree N` to `sglang serve`: Ulysses sequence parallelism splits the camera-major sequence across the GPUs and speeds up generation. The Multiview pipeline supports no tensor or ring parallelism, so pass both degrees explicitly, and N must divide its 8 key/value heads (2, 4, or 8). `--enable-cfg-parallel` additionally runs the two guidance branches on separate GPU groups and doubles the GPUs required. |\n", "| `--host 0.0.0.0` | Accepts connections from outside the container. For a remote notebook, set `SGLANG_BASE_URL`, or forward the port with `ssh -L 30000:127.0.0.1:30000 `. |\n", "\n", "The camera attention backend is selected automatically: FlashAttention-4 block-sparse kernels on Hopper and Blackwell, FlexAttention's Triton kernels elsewhere; `SGLANG_DIFFUSION_COSMOS3_MULTIVIEW_ATTENTION_BACKEND` (`triton` or `fa4`) overrides the choice. Joint LiDAR generation requires the checkpoint's `lidar_vae/` component; the LiDAR VAE runs on SGLang's built-in FlexAttention kernel, and `SGLANG_DIFFUSION_COSMOS3_LIDAR_USE_NATTEN=1` switches to `natten` when that package is installed in the image. The HTTP port opens only after the weights are loaded and the warmup request finishes, so wait for the `Uvicorn running` line." diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/run_sim-auto_with_sglang.ipynb b/.agents/skills/cookbooks/cosmos3/simulation/auto/run_sim-auto_with_sglang.ipynb index 9c4941020..0cf9a0958 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/auto/run_sim-auto_with_sglang.ipynb +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/run_sim-auto_with_sglang.ipynb @@ -93,7 +93,7 @@ "| `-e HF_TOKEN=$HF_TOKEN` | Read access to the gated checkpoint. |\n", "| `--revision` | Loads the checkpoint from `main`. The first start downloads it into the mounted cache. |\n", "| `docker pull` | The `dev` tag moves with SGLang's main branch; pulling before each start picks up the current image. |\n", - "| `--gpus all` | Runs on one GPU; select one with `--gpus '\"device=0\"'`. The Sim-Auto pipeline rejects sequence parallelism; `--tp-size` shards the transformer across GPUs if the model does not fit on one. |\n", + "| `--gpus all` | Makes every GPU visible; select one with `--gpus '\"device=0\"'`. With N GPUs, add `--num-gpus N --ulysses-degree N --sp-degree N` to `sglang serve`: Ulysses sequence parallelism splits every chunk across the GPUs and speeds up generation. Sim-Auto supports no other sequence parallelism, so pass both degrees explicitly, and N must divide the checkpoint's 8 key/value heads (2, 4, or 8). `--tp-size` can be combined when the model does not fit on one GPU, as long as `tp_size × ulysses_degree` still divides 8. |\n", "| `--host 0.0.0.0` | Accepts connections from outside the container. For a remote notebook, set `SGLANG_BASE_URL`, or forward the port with `ssh -L 30000:127.0.0.1:30000 `. |\n", "\n", "The HTTP port opens only after the weights are loaded and the warmup request finishes, so wait for the `Uvicorn running` line. The start-up log also prints `Cosmos3-Nano-Sim-Auto pipeline stages created`, followed by the checkpoint's chunk, ring, and sampler settings. This pipeline runs no text or video guardrails. Generation runs as an asynchronous job, so no server-side request timeout needs raising." From 10a763e8eba435934fc410da9691a47afc572bde Mon Sep 17 00:00:00 2001 From: Kedi Wu Date: Wed, 21 Oct 2026 22:47:58 -0700 Subject: [PATCH 3/5] add test cases for bimanual helpers --- .../bimanual/tests/test_bimanual_helpers.py | 569 ++++++++++++++++++ .github/workflows/validate-cookbooks.yml | 11 +- 2 files changed, 578 insertions(+), 2 deletions(-) create mode 100644 .agents/skills/cookbooks/cosmos3/simulation/bimanual/tests/test_bimanual_helpers.py diff --git a/.agents/skills/cookbooks/cosmos3/simulation/bimanual/tests/test_bimanual_helpers.py b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/tests/test_bimanual_helpers.py new file mode 100644 index 000000000..b5a1f8842 --- /dev/null +++ b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/tests/test_bimanual_helpers.py @@ -0,0 +1,569 @@ +"""CPU checks for the SGLang Bimanual cookbook: input preparation, payload encoding, and the HTTP client.""" + +import ast +import contextlib +import io +import itertools +import json +import re +import shutil +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import httpx +import numpy as np +from PIL import Image + +BIMANUAL_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(BIMANUAL_ROOT)) + +import bimanual_helpers +from bimanual_helpers import ( + CANVASES, + UNIFIED_ACTION_DIM, + UNIFIED_ACTION_SPACE, + BimanualInput, + domain_runs, + encode_extra_params, + first_media_image, + load_action_sidecar, + prepare_bimanual_input, + submit_bimanual, +) + +try: + import orjson +except ImportError: # The helper falls back to the standard library; both paths are covered when orjson is present. + orjson = None + +NUM_FRAMES = 9 # A short 4k+1 rollout with eight action rows. +DOMAINS = ["agibotworld"] * 4 + ["camera_pose"] * 2 + ["agibotworld"] * 2 +MODEL = "nvidia/Cosmos3-Nano-Sim-Bimanual" + + +def make_rows(count: int = NUM_FRAMES - 1, seed: int = 0) -> list[list[float]]: + """Pre-normalized double-precision rows, as the published sidecars store them.""" + rows = np.random.default_rng(seed).uniform(-1, 1, (count, UNIFIED_ACTION_DIM)) + rows[0, :4] = (0.0, 0.0, 1.1982105263541598e-07, 1.0) # The first values of the published sample. + return rows.tolist() + + +def parse_multipart(request: httpx.Request) -> dict[str, tuple[str | None, str | None, bytes]]: + """Return each multipart part as name -> (filename, content type, body).""" + boundary = request.headers["content-type"].split("boundary=", 1)[1].encode() + parts = {} + for chunk in request.read().split(b"--" + boundary)[1:-1]: + header, body = chunk.removeprefix(b"\r\n").split(b"\r\n\r\n", 1) + fields = {key.decode(): value.decode() for key, value in re.findall(rb'(\w+)="([^"]*)"', header)} + content_type = re.search(rb"Content-Type: ([^\r\n]+)", header) + parts[fields["name"]] = ( + fields.get("filename"), content_type.group(1).decode() if content_type else None, body.removesuffix(b"\r\n"), + ) + return parts + + +class BimanualFixture(unittest.TestCase): + def setUp(self): + temporary = tempfile.TemporaryDirectory() + self.addCleanup(temporary.cleanup) + self.root = Path(temporary.name).resolve() + self.image_path = self.root / "initial.png" + Image.new("RGB", (640, 360), (40, 80, 120)).save(self.image_path) # The published example's image size. + self.rows = make_rows() + self.sidecar_path = self.write_sidecar("actions.json", self.rows) + # The notebook builds its record from the published manifest row plus the downloaded paths. + self.record = { + "num_frames": NUM_FRAMES, "fps": 30, "seed": 2026, "resolution": "480", "raw_action_dim": 59, + "format_prompt_as_json": False, "prompt": " Place the cup in the sink. ", + "vision_path": "initial.png", "action_path": "actions.json", + } + + def write_sidecar(self, name: str, rows, domain_names=DOMAINS) -> Path: + """Write a unified sidecar; ``domain_names=None`` omits the field.""" + payload = {"action_space": UNIFIED_ACTION_SPACE, "action": rows, "contract": "slow_loop_middle_third_v1"} + if domain_names is not None: + payload["domain_names"] = domain_names + path = self.root / name + path.write_text(json.dumps(payload)) + return path + + def prepare(self, **overrides) -> BimanualInput: + return prepare_bimanual_input({**self.record, **overrides}, base_dir=self.root) + + +class PreparationTests(BimanualFixture): + def test_published_shape_record_prepares_canvas_output_size_and_metadata(self): + record = dict(self.record) + prepared = prepare_bimanual_input(record, base_dir=self.root) + self.assertEqual(record, self.record) # The record is read, not rewritten. + self.assertEqual(prepared.prompt, "Place the cup in the sink.") + self.assertEqual((prepared.image.size, prepared.image.mode), ((640, 360), "RGB")) + self.assertEqual((prepared.action.dtype, prepared.action.shape), (np.float32, (NUM_FRAMES - 1, UNIFIED_ACTION_DIM))) + np.testing.assert_array_equal(prepared.action, np.asarray(self.rows, dtype=np.float32)) + self.assertEqual(prepared.domain_names, DOMAINS) + self.assertEqual((prepared.num_frames, prepared.fps, prepared.seed), (NUM_FRAMES, 30, 2026)) + self.assertEqual((prepared.width, prepared.height), (832, 480)) # The 480p canvas closest to 16:9. + self.assertEqual((prepared.output_width, prepared.output_height), (640, 352)) # Content region in 16-pixel steps. + self.assertEqual(prepared.extra_args, { + "action_mode": "forward_dynamics", "domain_name": "agibotworld", "domain_names": DOMAINS, + "action_normalization": "none", "action_view_point": "ego_view", "format_prompt_as_json": False, + }) + self.assertEqual(prepared.metadata, { + "vision_path": str(self.image_path), "action_source": str(self.sidecar_path), + "action_space": UNIFIED_ACTION_SPACE, + "domain_runs": [ + {"domain_name": "agibotworld", "rows": [0, 4]}, + {"domain_name": "camera_pose", "rows": [4, 6]}, + {"domain_name": "agibotworld", "rows": [6, 8]}, + ], + "num_frames": NUM_FRAMES, "fps": 30, "seed": 2026, + "canvas": {"height": 480, "width": 832}, "output": {"height": 352, "width": 640}, + "action_mode": "forward_dynamics", "domain_name": "agibotworld", "action_normalization": "none", + "action_view_point": "ego_view", "format_prompt_as_json": False, + }) + # Without format_prompt_as_json the server's export default applies, so the field is not sent. + record.pop("format_prompt_as_json") + self.assertNotIn("format_prompt_as_json", prepare_bimanual_input(record, base_dir=self.root).extra_args) + + def test_integral_floats_are_accepted_for_integer_fields(self): + # The published manifest rows carry fps: 30; a record written with 30.0 names the same clock. + prepared = self.prepare(num_frames=float(NUM_FRAMES), fps=30.0, seed=2026.0, height=480.0, width=832.0) + for value, expected in ( + (prepared.num_frames, NUM_FRAMES), (prepared.fps, 30), (prepared.seed, 2026), + (prepared.height, 480), (prepared.width, 832), + ): + self.assertIs(type(value), int) + self.assertEqual(value, expected) + self.assertEqual(json.dumps(prepared.metadata["fps"]), "30") + self.assertEqual(self.prepare(seed=0).seed, 0) + for field, value, message in ( + ("fps", 29.97, "fps must be an integer >= 1"), + ("fps", True, "fps must be an integer >= 1"), + ("fps", "30", "fps must be an integer >= 1"), + ("fps", float("nan"), "fps must be an integer >= 1"), + ("fps", float("inf"), "fps must be an integer >= 1"), + ("fps", 0, "fps must be an integer >= 1"), + ("seed", -1, "seed must be an integer >= 0"), + ("num_frames", 1, "num_frames must be an integer >= 2"), + ("num_frames", 8, r"num_frames = 4k\+1"), + ("action_chunk_size", 7.5, "action_chunk_size must be an integer >= 1"), + ("height", 480.5, "height must be an integer >= 1"), + ): + overrides = {field: value} if field != "height" else {"height": value, "width": 832} + with self.subTest(field=field, value=value), self.assertRaisesRegex(ValueError, message): + self.prepare(**overrides) + + def test_action_rows_must_match_the_frame_count_and_unified_width(self): + with self.assertRaisesRegex(ValueError, rf"Expected 12x{UNIFIED_ACTION_DIM} unified actions, got \(8, {UNIFIED_ACTION_DIM}\)"): + self.prepare(num_frames=13) + self.write_sidecar("narrow.json", [row[:29] for row in self.rows]) + with self.assertRaisesRegex(ValueError, r"got \(8, 29\)"): + self.prepare(action_path="narrow.json") + with self.assertRaisesRegex(ValueError, "action_chunk_size must equal num_frames - 1"): + self.prepare(action_chunk_size=NUM_FRAMES) + self.assertEqual(self.prepare(action_chunk_size=NUM_FRAMES - 1).num_frames, NUM_FRAMES) + with self.assertRaisesRegex(ValueError, f"raw_action_dim must be {UNIFIED_ACTION_DIM}"): + self.prepare(raw_action_dim=29) + self.write_sidecar("overflow.json", [[1e39, *row[1:]] for row in self.rows]) # Finite in float64, inf in float32. + with self.assertRaisesRegex(ValueError, "finite in float32"): + self.prepare(action_path="overflow.json") + + def test_native_recipe_keys_inline_rows_and_contract_mismatches_reject(self): + for overrides, message in ( + ({"camera_trajectory": []}, "earlier 29-value recipe"), + ({"camera_pose_convention": "backward_chunk_anchored_16f"}, "earlier 29-value recipe"), + ({"camera_translation_scale": 1.0}, "earlier 29-value recipe"), + ({"camera_action_normalization": "none"}, "earlier 29-value recipe"), + ({"action_space": UNIFIED_ACTION_SPACE}, "earlier 29-value recipe"), + ({"action_normalization": "global_asinh_unified_v1"}, "action_normalization must be 'none'"), + ({"action": self.rows}, "inline rows are not supported"), + ({"actions": self.rows}, "inline rows are not supported"), + ({"model_mode": "inverse_dynamics"}, "forward dynamics"), + ({"autoregressive": False}, "autoregressive"), + ({"resolution": "720"}, "480p canvases"), + ({"resolution": 720}, "480p canvases"), + ({"view_point": "top_view"}, "Unsupported view_point: 'top_view'"), + ({"prompt": ["Place the cup."]}, "prompt must be a string"), + ({"format_prompt_as_json": "false"}, "format_prompt_as_json must be a boolean"), + ({"domain_name": 15}, "domain_name must be an embodiment name"), + ({"height": 480}, "height and width must be supplied together"), + ({"width": 832}, "height and width must be supplied together"), + ({"height": 720, "width": 1280}, "trained 480p canvas"), + ({"height": 544, "width": 544}, "trained 480p canvas"), # Mixes the sides of two canvases. + ): + with self.subTest(**overrides), self.assertRaisesRegex(ValueError, message): + self.prepare(**overrides) + # Accepted spellings of the same contract. + for overrides in ( + {"resolution": 480}, {"model_mode": "forward_dynamics"}, {"autoregressive": True}, + {"action_normalization": "none"}, {"action_condition": True, "aspect_ratio": "16,9", "num_steps": 4}, + {"view_point": "third_person_view"}, {"format_prompt_as_json": True}, + ): + with self.subTest(**overrides): + self.assertEqual(self.prepare(**overrides).num_frames, NUM_FRAMES) + without_resolution = {key: value for key, value in self.record.items() if key != "resolution"} + self.assertEqual(prepare_bimanual_input(without_resolution, base_dir=self.root).height, 480) # Defaults to 480. + self.assertEqual(prepare_bimanual_input({**without_resolution, "image_size": "480"}, base_dir=self.root).height, 480) + with self.assertRaisesRegex(ValueError, "480p canvases"): + prepare_bimanual_input({**without_resolution, "image_size": 720}, base_dir=self.root) + + def test_canvas_follows_the_image_aspect_unless_a_trained_canvas_is_named(self): + for size, canvas, output in ( + ((640, 360), (832, 480), (640, 352)), # Published example: 16:9 content, cropped to 16-pixel multiples. + ((1920, 1080), (832, 480), (832, 464)), # Larger than the canvas: scaled down, height rounded to 16. + ((360, 640), (480, 832), (352, 640)), # Portrait. + ((500, 500), (640, 640), (496, 496)), # Square. + ((1472, 1088), (736, 544), (736, 544)), # Twice the 4:3-ish canvas fills it exactly. + ): + with self.subTest(size=size): + name = f"image_{size[0]}x{size[1]}.png" + Image.new("RGB", size).save(self.root / name) + prepared = self.prepare(vision_path=name) + self.assertEqual((prepared.width, prepared.height), canvas) + self.assertEqual((prepared.output_width, prepared.output_height), output) + self.assertEqual(prepared.metadata["canvas"], {"height": canvas[1], "width": canvas[0]}) + for width, height in CANVASES: + with self.subTest(canvas=(width, height)): + prepared = self.prepare(height=height, width=width) + self.assertEqual((prepared.width, prepared.height), (width, height)) + prepared = self.prepare(height=544, width=736) # A named canvas never upscales the image. + self.assertEqual(self.prepare(height=832, width=480).metadata["canvas"], {"height": 832, "width": 480}) # Portrait. + self.assertEqual((prepared.output_width, prepared.output_height), (640, 352)) + + def test_sidecar_domain_names_broadcast_or_fall_back_to_the_record(self): + expected_rows = np.asarray(self.rows) + single = self.write_sidecar("single.json", self.rows, "camera_pose") + rows, names = load_action_sidecar(single) + np.testing.assert_array_equal(rows, expected_rows) + self.assertEqual(names, ["camera_pose"] * (NUM_FRAMES - 1)) + one = self.write_sidecar("one.json", self.rows, [" agibotworld "]) + self.assertEqual(load_action_sidecar(one)[1], ["agibotworld"] * (NUM_FRAMES - 1)) # Stripped, then broadcast. + unnamed = self.write_sidecar("unnamed.json", self.rows, None) + self.assertEqual(load_action_sidecar(unnamed, default_domain="agibotworld")[1], ["agibotworld"] * (NUM_FRAMES - 1)) + with self.assertRaisesRegex(ValueError, "no domain_names"): + load_action_sidecar(unnamed) + bare = self.root / "bare.json" + bare.write_text(json.dumps(self.rows)) + rows, names = load_action_sidecar(bare, default_domain="agibotworld") + np.testing.assert_array_equal(rows, expected_rows) + self.assertEqual(names, ["agibotworld"] * (NUM_FRAMES - 1)) + with self.assertRaisesRegex(ValueError, "bare action array needs domain_name"): + load_action_sidecar(bare) + # Through a record: domain_name fills in missing names and labels the request. + prepared = self.prepare(action_path="unnamed.json", domain_name="agibotworld") + self.assertEqual(prepared.extra_args["domain_name"], "agibotworld") + self.assertEqual(prepared.metadata["domain_runs"], [{"domain_name": "agibotworld", "rows": [0, NUM_FRAMES - 1]}]) + with self.assertRaisesRegex(ValueError, "no domain_names"): + self.prepare(action_path="unnamed.json") + prepared = self.prepare(domain_name="camera_pose") # The sidecar's per-row names still route every row. + self.assertEqual((prepared.extra_args["domain_name"], prepared.domain_names), ("camera_pose", DOMAINS)) + + def test_sidecar_contract_errors_name_the_file_and_numeric_errors_name_the_actions(self): + unified = {"action_space": UNIFIED_ACTION_SPACE, "action": self.rows, "domain_names": DOMAINS} + for name, payload, message in ( + ("space.json", {**unified, "action_space": "native_v0"}, "action_space must be 'normalized_unified_v1', got 'native_v0'"), + ("none.json", {**unified, "action_space": None}, "got None"), + ("missing.json", {"action_space": UNIFIED_ACTION_SPACE, "domain_names": DOMAINS}, "needs an `action` list"), + ("flat.json", {**unified, "action": self.rows[0]}, "actions must be a list of rows"), + ("count.json", {**unified, "domain_names": DOMAINS[:3]}, "3 domain names for 8 action rows"), + ("blank.json", {**unified, "domain_names": ["agibotworld", " "] * 4}, "non-empty list of embodiment names"), + ("typed.json", {**unified, "domain_names": [15] * 8}, "non-empty list of embodiment names"), + ("empty.json", {**unified, "domain_names": []}, "non-empty list of embodiment names"), + ): + with self.subTest(file=name): + path = self.root / name + path.write_text(json.dumps(payload)) + with self.assertRaisesRegex(ValueError, message) as raised: + load_action_sidecar(path) + self.assertIn(str(path), str(raised.exception)) + for name, first, message in ( + ("bool.json", True, "actions must contain real numbers"), + ("text.json", "0.5", "actions must contain real numbers"), + ("nan.json", float("nan"), "actions must contain only finite values"), + ("inf.json", float("-inf"), "actions must contain only finite values"), + ): + with self.subTest(file=name): + path = self.write_sidecar(name, [[first, *self.rows[0][1:]], *self.rows[1:]]) + with self.assertRaisesRegex(ValueError, message): + load_action_sidecar(path) + scalar = self.write_sidecar("scalar.json", "0.5", "agibotworld") # A string where the rows should be. + with self.assertRaisesRegex(ValueError, "actions must contain real numbers"): + load_action_sidecar(scalar) + + def test_domain_runs_segments_consecutive_embodiments(self): + self.assertEqual(domain_runs([]), []) + self.assertEqual(domain_runs(["agibotworld"]), [("agibotworld", 0, 1)]) + self.assertEqual(domain_runs(DOMAINS), [("agibotworld", 0, 4), ("camera_pose", 4, 6), ("agibotworld", 6, 8)]) + + def test_inputs_must_exist_be_non_empty_and_hydrated(self): + with self.assertRaisesRegex(FileNotFoundError, "Missing or empty input"): + self.prepare(action_path="absent.json") + (self.root / "empty.json").write_bytes(b"") + with self.assertRaisesRegex(FileNotFoundError, "Missing or empty input"): + self.prepare(action_path="empty.json") + (self.root / "pointer.png").write_bytes(b"version https://git-lfs.github.com/spec/v1\noid sha256:0\nsize 1\n") + with self.assertRaisesRegex(ValueError, "Unhydrated Git LFS pointer"): + self.prepare(vision_path="pointer.png") + for value in (None, "", " ", 7): + with self.subTest(value=value), self.assertRaisesRegex(ValueError, "Supply a local"): + self.prepare(vision_path=value) + # Relative paths resolve against base_dir; absolute paths are used as given. + self.assertEqual(self.prepare(vision_path=str(self.image_path)).metadata["vision_path"], str(self.image_path)) + nested = self.root / "nested" + nested.mkdir() + self.assertEqual( + prepare_bimanual_input({**self.record, "vision_path": "../initial.png", "action_path": "../actions.json"}, base_dir=nested).image.size, + (640, 360), + ) + + def test_first_media_image_honors_exif_orientation_and_decodes_the_first_video_frame(self): + rotated = self.root / "rotated.jpg" + exif = Image.Exif() + exif[0x0112] = 6 # Orientation: rotate 90 degrees clockwise for display. + Image.new("RGB", (40, 20)).save(rotated, exif=exif) + self.assertEqual(first_media_image(rotated).size, (20, 40)) + self.assertEqual((first_media_image(self.image_path).size, first_media_image(self.image_path).mode), ((640, 360), "RGB")) + if not all(shutil.which(name) for name in ("ffmpeg", "ffprobe")): + self.skipTest("Install ffmpeg to decode video conditioning frames.") + with self.assertRaises((subprocess.CalledProcessError, ValueError)): + first_media_image(self.sidecar_path) # Neither an image nor a video. + video = self.root / "clip.mp4" + width, height = 64, 48 + frames = b"".join(bytes([level]) * (width * height) for level in (200, 20, 20)) + subprocess.run( + [ + "ffmpeg", "-v", "error", "-y", "-f", "rawvideo", "-pixel_format", "gray", + "-video_size", f"{width}x{height}", "-framerate", "30", "-i", "pipe:0", + "-an", "-c:v", "libx264", "-crf", "0", "-pix_fmt", "yuv420p", str(video), + ], + input=frames, capture_output=True, check=True, + ) + frame = first_media_image(video) + self.assertEqual((frame.size, frame.mode), ((width, height), "RGB")) + self.assertAlmostEqual(float(np.asarray(frame).mean()), 200, delta=4) # The first frame, not a later dark one. + prepared = self.prepare(vision_path=video.name) + self.assertEqual((prepared.width, prepared.height), (736, 544)) # 4:3 lands on the 736x544 canvas. + self.assertEqual((prepared.output_width, prepared.output_height), (width, height)) + + +class EncodingTests(BimanualFixture): + def encodings(self, prepared: BimanualInput) -> dict[str, str]: + """The extra_params text from orjson, when installed, and from the standard-library fallback.""" + texts = {} + if orjson is not None: + texts["orjson"] = encode_extra_params(prepared) + with patch.dict(sys.modules, {"orjson": None}): # Makes ``import orjson`` raise ImportError. + texts["json"] = encode_extra_params(prepared) + return texts + + def test_extra_params_round_trip_float32_rows_in_compact_form(self): + prepared = self.prepare() + naive = json.dumps({**prepared.extra_args, "action": prepared.action.astype(float).tolist()}, separators=(",", ":")) + for encoder, text in self.encodings(prepared).items(): + with self.subTest(encoder=encoder): + payload = json.loads(text) + self.assertEqual({key: value for key, value in payload.items() if key != "action"}, prepared.extra_args) + np.testing.assert_array_equal(np.asarray(payload["action"], dtype=np.float32), prepared.action) + self.assertNotIn(" ", text) + self.assertLess(len(text), 0.75 * len(naive)) # Shortest float32 decimals rather than double repr. + + def test_published_rollout_length_stays_within_the_form_field_limit(self): + # The published mixed-action example has 900 rows: 300 robot, 300 camera, 300 robot. + frames = 901 + names = ["agibotworld"] * 300 + ["camera_pose"] * 300 + ["agibotworld"] * 300 + self.write_sidecar("long.json", make_rows(frames - 1, seed=1), names) + prepared = self.prepare(num_frames=frames, action_path="long.json") + self.assertEqual(prepared.metadata["domain_runs"], [ + {"domain_name": "agibotworld", "rows": [0, 300]}, {"domain_name": "camera_pose", "rows": [300, 600]}, + {"domain_name": "agibotworld", "rows": [600, 900]}, + ]) + naive = json.dumps({**prepared.extra_args, "action": prepared.action.astype(float).tolist()}, separators=(",", ":")) + self.assertGreater(len(naive.encode()), bimanual_helpers.MAX_FORM_FIELD_BYTES) # Double repr would be rejected. + for encoder, text in self.encodings(prepared).items(): + with self.subTest(encoder=encoder): + self.assertLessEqual(len(text.encode()), bimanual_helpers.MAX_FORM_FIELD_BYTES) + np.testing.assert_array_equal(np.asarray(json.loads(text)["action"], dtype=np.float32), prepared.action) + with patch.object(bimanual_helpers, "MAX_FORM_FIELD_BYTES", 4096), self.assertRaisesRegex(ValueError, "multipart-field"): + encode_extra_params(prepared) + + def test_lossy_serialization_is_detected_before_sending(self): + def truncate(value, **kwargs): + return f"{float(value):.3f}" + + prepared = self.prepare() + with ( + patch.dict(sys.modules, {"orjson": None}), + patch.object(np, "format_float_positional", truncate), + self.assertRaisesRegex(ValueError, "preserve every float32 value"), + ): + encode_extra_params(prepared) + + +class SubmitTests(BimanualFixture): + def setUp(self): + super().setUp() + self.prepared = self.prepare(fps=30.0) # An integral float travels as the form field "30". + self.output = self.root / "out" / "sample" / "sample.mp4" + self.requests = [] + self.on_submit = lambda: httpx.Response(200, json={"id": "job-7", "status": "queued"}) + self.polls = iter([{"id": "job-7", "status": "completed"}]) + self.on_content = lambda: httpx.Response(200, content=b"mock video", headers={"content-type": "video/mp4"}) + + def handler(self, request: httpx.Request) -> httpx.Response: + self.requests.append((request.method, request.url.path)) + if (request.method, request.url.path) == ("POST", "/v1/videos"): + self.submission = request + return self.on_submit() + if (request.method, request.url.path) == ("GET", "/v1/videos/job-7"): + item = next(self.polls) + return item if isinstance(item, httpx.Response) else httpx.Response(200, json=item) + if (request.method, request.url.path) == ("GET", "/v1/videos/job-7/content"): + return self.on_content() + raise AssertionError(f"Unexpected request {request.method} {request.url}") + + def submit(self, **kwargs): + with ( + patch("time.sleep") as sleep, + contextlib.redirect_stdout(io.StringIO()) as stdout, + httpx.Client(base_url="http://test", transport=httpx.MockTransport(self.handler)) as client, + ): + result = submit_bimanual(client, self.prepared, MODEL, self.output, **kwargs) + return result, sleep, stdout.getvalue() + + def test_multipart_submission_polls_the_job_and_streams_the_video_with_sidecars(self): + self.polls = iter([{"id": "job-7", "status": "in_progress", "progress": 50}, {"id": "job-7", "status": "completed", "progress": 100}]) + result, sleep, printed = self.submit(timeout=120, poll_interval=5) + self.assertEqual(result, self.output) + self.assertEqual(self.output.read_bytes(), b"mock video") + self.assertIn("Submitted job-7 at http://test/v1/videos/job-7", printed) + self.assertEqual(self.requests, [ + ("POST", "/v1/videos"), ("GET", "/v1/videos/job-7"), ("GET", "/v1/videos/job-7"), ("GET", "/v1/videos/job-7/content"), + ]) + self.assertEqual([call.args[0] for call in sleep.call_args_list], [5, 5]) + + parts = parse_multipart(self.submission) + form = {name: body.decode() for name, (_, _, body) in parts.items() if name != "input_reference"} + expected_form = { + "model": MODEL, "prompt": "Place the cup in the sink.", "size": "832x480", + "num_frames": "9", "fps": "30", "seed": "2026", "extra_params": encode_extra_params(self.prepared), + } + self.assertEqual(form, expected_form) + for name in ("num_frames", "fps", "seed"): # Integer text: SGLang parses these fields as integers. + self.assertTrue(form[name].isdigit(), form[name]) + filename, content_type, png = parts["input_reference"] + self.assertEqual((filename, content_type), ("conditioning.png", "image/png")) + with Image.open(io.BytesIO(png)) as image: + self.assertEqual((image.format, image.size, image.getpixel((0, 0))), ("PNG", (640, 360), (40, 80, 120))) + extra = json.loads(form["extra_params"]) + self.assertEqual({key: value for key, value in extra.items() if key != "action"}, self.prepared.extra_args) + np.testing.assert_array_equal(np.asarray(extra["action"], dtype=np.float32), self.prepared.action) + + self.assertEqual(json.loads(self.output.with_suffix(".request.json").read_text()), expected_form) + self.assertEqual(json.loads(self.output.with_suffix(".metadata.json").read_text()), self.prepared.metadata) + self.assertEqual(json.loads(self.output.with_suffix(".job.json").read_text()), {"id": "job-7", "status": "completed", "progress": 100}) + self.assertEqual( + sorted(path.name for path in self.output.parent.iterdir()), + ["sample.job.json", "sample.metadata.json", "sample.mp4", "sample.request.json"], + ) + + def test_a_job_completed_at_submission_needs_no_polling(self): + self.on_submit = lambda: httpx.Response(200, json={"id": "job-7", "status": "completed"}) + self.polls = iter([]) + result, sleep, _ = self.submit() + sleep.assert_not_called() + self.assertEqual(result.read_bytes(), b"mock video") + self.assertEqual([path for _, path in self.requests], ["/v1/videos", "/v1/videos/job-7/content"]) + + def test_rejected_submissions_failed_jobs_and_poll_errors_save_the_error_and_raise(self): + error_path = self.output.with_suffix(".error.txt") + self.on_submit = lambda: httpx.Response(400, text="mock submission rejected") + with self.assertRaises(httpx.HTTPStatusError): + self.submit() + self.assertEqual(error_path.read_text(), "mock submission rejected") + self.assertFalse(self.output.exists()) + self.assertTrue(self.output.with_suffix(".request.json").exists()) # Saved before sending, for replay. + + self.on_submit = lambda: httpx.Response(200, json={"id": "job-7", "status": "queued"}) + self.polls = iter([{"id": "job-7", "status": "failed", "error": {"message": "mock rollout failed"}}]) + with self.assertRaisesRegex(RuntimeError, "Video generation failed: {'message': 'mock rollout failed'}"): + self.submit() + self.assertEqual(json.loads(error_path.read_text())["error"], {"message": "mock rollout failed"}) + self.assertEqual(json.loads(self.output.with_suffix(".job.json").read_text())["status"], "failed") + self.assertNotIn(("GET", "/v1/videos/job-7/content"), self.requests) + + for status, message in (("cancelled", "Video generation cancelled"), ("exploded", "Unexpected video job status")): + with self.subTest(status=status): + self.polls = iter([{"id": "job-7", "status": status}]) + with self.assertRaisesRegex(RuntimeError, message): + self.submit() + self.assertFalse(self.output.exists()) + + self.polls = iter([httpx.Response(503, text="mock poll outage")]) + with self.assertRaises(httpx.HTTPStatusError): + self.submit() + self.assertEqual(error_path.read_text(), "mock poll outage") + + def test_download_failures_leave_no_partial_video(self): + self.on_submit = lambda: httpx.Response(200, json={"id": "job-7", "status": "completed"}) + self.on_content = lambda: httpx.Response(500, text="mock download failed") + with self.assertRaises(httpx.HTTPStatusError): + self.submit() + self.assertEqual(self.output.with_suffix(".error.txt").read_text(), "mock download failed") + self.on_content = lambda: httpx.Response(200, content=b"") + with self.assertRaisesRegex(ValueError, "empty"): + self.submit() + self.assertEqual([path.name for path in self.output.parent.iterdir() if ".mp4" in path.name], []) + + def test_polling_timeout_raises_and_leaves_the_job_running(self): + self.polls = itertools.repeat({"id": "job-7", "status": "in_progress"}) + clock = SimpleNamespace(monotonic=Mock(side_effect=itertools.count(100.0, 0.3)), sleep=Mock()) + with ( + patch.object(bimanual_helpers, "time", clock), + contextlib.redirect_stdout(io.StringIO()), + httpx.Client(base_url="http://test", transport=httpx.MockTransport(self.handler)) as client, + self.assertRaisesRegex(TimeoutError, "job remains at http://test/v1/videos/job-7"), + ): + submit_bimanual(client, self.prepared, MODEL, self.output, timeout=1, poll_interval=5) + # Deadline 101.0: sleeps shrink to the remaining time, one poll fits, then the deadline passes. + self.assertEqual([round(call.args[0], 6) for call in clock.sleep.call_args_list], [0.7, 0.1]) + self.assertEqual(self.requests, [("POST", "/v1/videos"), ("GET", "/v1/videos/job-7")]) # No cancellation. + self.assertEqual(json.loads(self.output.with_suffix(".job.json").read_text())["status"], "in_progress") + self.assertFalse(self.output.with_suffix(".error.txt").exists()) + + def test_polling_parameters_must_be_positive_and_finite(self): + for kwargs in ( + {"timeout": 0}, {"timeout": -5}, {"timeout": float("inf")}, {"timeout": float("nan")}, + {"poll_interval": 0}, {"poll_interval": float("inf")}, + ): + with ( + self.subTest(**kwargs), + httpx.Client(base_url="http://test", transport=httpx.MockTransport(self.handler)) as client, + self.assertRaisesRegex(ValueError, "positive and finite"), + ): + submit_bimanual(client, self.prepared, MODEL, self.output, **kwargs) + self.assertEqual(self.requests, []) + self.assertFalse(self.output.parent.exists()) + + +class NotebookContractTests(unittest.TestCase): + def test_notebook_imports_names_the_helper_module_defines(self): + notebook = json.loads((BIMANUAL_ROOT / "run_bimanual_with_sglang.ipynb").read_text()) + imported = set() + for cell in notebook["cells"]: + if cell["cell_type"] != "code": + continue + for node in ast.walk(ast.parse("".join(cell["source"]))): + if isinstance(node, ast.ImportFrom) and node.module == "bimanual_helpers": + imported.update(alias.name for alias in node.names) + self.assertLessEqual({"domain_runs", "prepare_bimanual_input", "submit_bimanual"}, imported) + for name in sorted(imported): + self.assertTrue(callable(getattr(bimanual_helpers, name, None)), name) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/workflows/validate-cookbooks.yml b/.github/workflows/validate-cookbooks.yml index 06cfd4159..ec96b8456 100644 --- a/.github/workflows/validate-cookbooks.yml +++ b/.github/workflows/validate-cookbooks.yml @@ -2,7 +2,8 @@ # SPDX-License-Identifier: OpenMDW-1.1 # Checks that every cookbook and evaluation notebook parses and lints, then runs the -# Auto cookbook's CPU tests (request helpers, media timelines, synthetic RDS-HQ scene). +# Auto cookbook's CPU tests (request helpers, media timelines, synthetic RDS-HQ scene) +# and the Bimanual cookbook's CPU tests (input preparation, payload encoding, HTTP client). name: Validate cookbooks on: push: @@ -15,6 +16,7 @@ permissions: contents: read env: AUTO: .agents/skills/cookbooks/cosmos3/simulation/auto + BIMANUAL: .agents/skills/cookbooks/cosmos3/simulation/bimanual jobs: validate-cookbooks: name: Validate cookbooks @@ -29,7 +31,7 @@ jobs: - name: Install notebook checks run: >- python -m pip install --disable-pip-version-check - nbformat==5.11.1 ruff==0.16.4 httpx numpy attrs safetensors ipython matplotlib huggingface_hub imageio-ffmpeg==0.6.0 + nbformat==5.11.1 ruff==0.16.4 httpx numpy pillow orjson attrs safetensors ipython matplotlib huggingface_hub imageio-ffmpeg==0.6.0 - name: Install media tools run: | @@ -119,6 +121,11 @@ jobs: ruff check --output-format=github --select=E9,F63,F7,F82 "$AUTO/helpers/sim_auto_vllm_omni_helpers.py" "$AUTO/helpers/sim_auto_sglang_helpers.py" "$AUTO/tests/test_sim_auto_helpers.py" python -m unittest discover -s "$AUTO/tests" -p 'test_sim_auto_helpers.py' -v + - name: Check Bimanual input preparation and HTTP helpers + run: | + ruff check --output-format=github --select=E9,F63,F7,F82 "$BIMANUAL/bimanual_helpers.py" "$BIMANUAL/tests/test_bimanual_helpers.py" + python -m unittest discover -s "$BIMANUAL/tests" -p 'test_bimanual_helpers.py' -v + - name: Check media timelines with system FFmpeg run: | ffmpeg -version From def2c2e740a9f2f03d213391c969252793c8332e Mon Sep 17 00:00:00 2001 From: Kedi Wu Date: Wed, 21 Oct 2026 22:51:26 -0700 Subject: [PATCH 4/5] cast fps to integer for bimanual request, backend param shared check, add note about which commit the notebook is validated against --- .../cookbooks/cosmos3/simulation/README.md | 4 +- .../simulation/auto/helpers/auto_helpers.py | 18 ++++- .../auto/run_auto_with_sglang.ipynb | 56 +++++++------- .../auto/run_sim-auto_with_sglang.ipynb | 4 +- .../auto/tests/test_daytime_sample.py | 3 + .../simulation/bimanual/bimanual_helpers.py | 10 ++- .../bimanual/run_bimanual_with_sglang.ipynb | 76 +++++++++++-------- .../transfer/run_transfer_with_sglang.ipynb | 38 +++++----- 8 files changed, 121 insertions(+), 88 deletions(-) diff --git a/.agents/skills/cookbooks/cosmos3/simulation/README.md b/.agents/skills/cookbooks/cosmos3/simulation/README.md index 7e7dc709d..b4b622feb 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/README.md +++ b/.agents/skills/cookbooks/cosmos3/simulation/README.md @@ -39,7 +39,9 @@ distilled student: a fixed four-step sampler with a rolling key/value cache. the `vllm/vllm-omni:cosmos3` image. - **SGLang:** Docker with the NVIDIA Container Toolkit. The four SGLang notebooks use the published `lmsysorg/sglang:dev` image; the Auto and Sim-Transfer notebooks read - their inputs from the mounted cookbook root. + their inputs from the mounted cookbook root. The `dev` tag moves with SGLang's main + branch; the notebooks were validated against SGLang commit [`ab580b2e`](https://github.com/sgl-project/sglang/commit/ab580b2e060059c0f8387b598a9f648dcec13f93), + published as the fixed tag `lmsysorg/sglang:nightly-dev-20261020-ab580b2e`. If you pass `--defaults-file` to Framework inference, it **replaces** the built-in sampling defaults rather than merging with them. Start from the defaults file the diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/helpers/auto_helpers.py b/.agents/skills/cookbooks/cosmos3/simulation/auto/helpers/auto_helpers.py index 0298440ba..7fcccb8f8 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/auto/helpers/auto_helpers.py +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/helpers/auto_helpers.py @@ -387,7 +387,17 @@ def validate_framework_outputs(output_dir: Path, manifest: dict) -> dict: return result -# vLLM-Omni request preparation and HTTP client. +# vLLM-Omni and SGLang request preparation and HTTP client. + +BACKENDS = ("vllm_omni", "sglang") + + +def _check_backend(backend: str) -> str: + """Return ``backend`` when it names a supported Videos API server.""" + if backend not in BACKENDS: + raise ValueError(f"Unknown backend {backend!r}; use {' or '.join(map(repr, BACKENDS))}.") + return backend + def check_model(client, expected: str | None = None) -> str: """Return the served model ID from ``GET /v1/models``, requiring ``expected`` when given.""" @@ -406,15 +416,14 @@ def check_model(client, expected: str | None = None) -> str: def prepare_request( manifest: dict, base_dir: Path, model: str, *, backend: str = "vllm_omni" ) -> tuple[dict, list[Path]]: - """Build a Videos API request and its upload list. + """Build a Videos API request and its upload list for one of ``BACKENDS``. Both backends use multipart form fields. vLLM-Omni uploads every file and refers to it by a reference index; SGLang reads each file from an absolute path visible to the server, so the upload list stays empty. For SGLang, ``lidar.return_output`` is translated to ``lidar.decode``. """ - if backend not in ("vllm_omni", "sglang"): - raise ValueError(f"Unknown backend {backend!r}; use 'vllm_omni' or 'sglang'.") + backend = _check_backend(backend) for field in ("negative_prompt", "negative_metadata_mode"): if manifest.get(field) is not None: raise ValueError( @@ -612,6 +621,7 @@ def submit_joint( vLLM-Omni returns only the LiDAR descriptor; download its numeric output from /lidar separately. """ + backend = _check_backend(backend) return_lidar = True if backend == "sglang": return_lidar = json.loads(form["extra_params"])["lidar"].get("decode", False) diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_sglang.ipynb b/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_sglang.ipynb index 34ce1213c..a2034d5fe 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_sglang.ipynb +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/run_auto_with_sglang.ipynb @@ -93,7 +93,7 @@ "| `-v $SGLANG_HF_CACHE:/root/.cache/huggingface` | Keeps the downloaded checkpoint across container restarts. The setup lines default the directory to `~/.cache/sglang-huggingface`, create it, and make it world-writable. |\n", "| `-e HF_TOKEN=$HF_TOKEN` | Read access to the gated checkpoint. |\n", "| `--revision` | Loads the checkpoint from `main`. The first start downloads it into the mounted cache. |\n", - "| `docker pull` | The `dev` tag moves with SGLang's main branch; pulling before each start picks up the current image. |\n", + "| `docker pull` | The `dev` tag moves with SGLang's main branch; pulling before each start picks up the current image. This notebook was validated against the image built from SGLang commit [`ab580b2e`](https://github.com/sgl-project/sglang/commit/ab580b2e060059c0f8387b598a9f648dcec13f93), published as the fixed tag `lmsysorg/sglang:nightly-dev-20261020-ab580b2e`; use that tag in place of `dev` to reproduce the validated run. |\n", "| `--gpus all` | Makes every GPU visible; select one with `--gpus '\"device=0\"'`. With N GPUs, add `--num-gpus N --ulysses-degree N --sp-degree N` to `sglang serve`: Ulysses sequence parallelism splits the camera-major sequence across the GPUs and speeds up generation. The Multiview pipeline supports no tensor or ring parallelism, so pass both degrees explicitly, and N must divide its 8 key/value heads (2, 4, or 8). `--enable-cfg-parallel` additionally runs the two guidance branches on separate GPU groups and doubles the GPUs required. |\n", "| `--host 0.0.0.0` | Accepts connections from outside the container. For a remote notebook, set `SGLANG_BASE_URL`, or forward the port with `ssh -L 30000:127.0.0.1:30000 `. |\n", "\n", @@ -121,8 +121,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "auto-sglang-config", "metadata": {}, + "outputs": [], "source": [ "from pathlib import Path\n", "import copy\n", @@ -190,9 +192,7 @@ "print(\"Outputs:\", OUTPUT_DIR)\n", "if not OUTPUT_DIR.is_relative_to(COOKBOOK_ROOT):\n", " print(f\"Note: {OUTPUT_DIR} is outside the cookbook root; the server must see it at the same path.\")" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -208,15 +208,15 @@ }, { "cell_type": "code", + "execution_count": null, "id": "auto-sglang-sample", "metadata": {}, + "outputs": [], "source": [ "SAMPLE_ROOT = fetch_sample(AUTO_SAMPLE_ROOT)\n", "print(\"Verified sample:\", SAMPLE_ROOT)\n", "print(\"Cameras:\", \", \".join(SAMPLE_CAMERAS))" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -230,8 +230,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "auto-sglang-verify", "metadata": {}, + "outputs": [], "source": [ "deadline = time.monotonic() + 1800\n", "while True:\n", @@ -249,9 +251,7 @@ " raise TimeoutError(f\"No ready server at {SGLANG_BASE_URL} after 1800s ({reason}).\")\n", " time.sleep(5)\n", "print(\"Server ready; serving\", served)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -271,8 +271,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "auto-sglang-wsm-prepare", "metadata": {}, + "outputs": [], "source": [ "WSM_REQUEST = AUTO_ROOT / \"assets\" / \"daytime_scene\" / \"transfer_auto_wsm.json\"\n", "WSM_OUTPUT_DIR = OUTPUT_DIR / \"wsm\"\n", @@ -297,9 +299,7 @@ "for camera_key, control in zip(SAMPLE_CAMERAS, wsm_controls):\n", " info = probe_video(control)\n", " print(f\"{camera_key:28} {info['frames']} frames, {info['width']}×{info['height']}, {float(info['fps']):g} FPS\")" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -313,17 +313,17 @@ }, { "cell_type": "code", + "execution_count": null, "id": "auto-sglang-wsm-submit", "metadata": {}, + "outputs": [], "source": [ "with httpx.Client(base_url=SGLANG_BASE_URL, headers=HEADERS, timeout=HTTP_TIMEOUT) as client:\n", " started = time.monotonic()\n", " wsm_video = submit_wsm(client, wsm_form, wsm_uploads, WSM_OUTPUT_DIR, async_mode=True, timeout=GENERATION_TIMEOUT)\n", " wsm_elapsed = round(time.monotonic() - started, 1)\n", "print(f\"{wsm_elapsed:.1f}s {wsm_video}\")" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -349,15 +349,15 @@ }, { "cell_type": "code", + "execution_count": null, "id": "auto-sglang-wsm-validate", "metadata": {}, + "outputs": [], "source": [ "wsm_cameras = split_cameras(wsm_video, WSM_OUTPUT_DIR, SAMPLE_CAMERAS, wsm_manifest[\"num_frames\"], wsm_manifest[\"fps\"])\n", "save_outputs(WSM_OUTPUT_DIR, wsm_video, wsm_cameras, wsm_manifest)\n", "display(HTML(camera_grid(wsm_cameras)))" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -377,8 +377,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "auto-sglang-lidar-prepare", "metadata": {}, + "outputs": [], "source": [ "LIDAR_REQUEST = AUTO_ROOT / \"assets\" / \"daytime_scene\" / \"transfer_auto_joint.json\"\n", "LIDAR_OUTPUT_DIR = OUTPUT_DIR / \"lidar\"\n", @@ -421,9 +423,7 @@ "rgb_conditions = dict(zip(SAMPLE_CAMERAS, lidar_conditions))\n", "print(f\"LiDAR input: {lidar_input_info['shape']} {lidar_input_info['dtype']}; using {lidar_sweeps} sweeps at {LIDAR_FPS:g} Hz\")\n", "print(f\"Prepared {len(lidar_controls)} WSM controls, {len(lidar_conditions)} RGB conditions, and the LiDAR control under {LIDAR_OUTPUT_DIR}\")" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -439,8 +439,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "auto-sglang-lidar-submit", "metadata": {}, + "outputs": [], "source": [ "with httpx.Client(base_url=SGLANG_BASE_URL, headers=HEADERS, timeout=HTTP_TIMEOUT) as client:\n", " started = time.monotonic()\n", @@ -449,9 +451,7 @@ "print(f\"Job {lidar_result['job_id']} completed in {lidar_elapsed:.1f}s\")\n", "print(\"Camera video:\", lidar_result[\"video\"])\n", "print(\"Generated LiDAR:\", lidar_result[\"lidar\"])" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -481,8 +481,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "auto-sglang-lidar-validate", "metadata": {}, + "outputs": [], "source": [ "lidar_output_info = validate_sglang_lidar(lidar_result[\"lidar\"], lidar_sweeps, fps=LIDAR_FPS)\n", "lidar_descriptor = lidar_result[\"descriptor\"]\n", @@ -496,9 +498,7 @@ "save_outputs(LIDAR_OUTPUT_DIR, lidar_result[\"video\"], lidar_cameras, lidar_manifest, lidar=lidar_result)\n", "print(f\"LiDAR: {lidar_output_info['shape']} at {lidar_descriptor['fps']:g} Hz; split into {len(lidar_cameras)} cameras\")\n", "display(HTML(camera_grid(lidar_cameras)))" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/run_sim-auto_with_sglang.ipynb b/.agents/skills/cookbooks/cosmos3/simulation/auto/run_sim-auto_with_sglang.ipynb index 0cf9a0958..1fd7bb2e2 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/auto/run_sim-auto_with_sglang.ipynb +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/run_sim-auto_with_sglang.ipynb @@ -92,7 +92,7 @@ "| `-v $SGLANG_HF_CACHE:/root/.cache/huggingface` | Keeps the downloaded checkpoint across container restarts. The setup lines default the directory to `~/.cache/sglang-huggingface`, create it, and make it world-writable. |\n", "| `-e HF_TOKEN=$HF_TOKEN` | Read access to the gated checkpoint. |\n", "| `--revision` | Loads the checkpoint from `main`. The first start downloads it into the mounted cache. |\n", - "| `docker pull` | The `dev` tag moves with SGLang's main branch; pulling before each start picks up the current image. |\n", + "| `docker pull` | The `dev` tag moves with SGLang's main branch; pulling before each start picks up the current image. This notebook was validated against the image built from SGLang commit [`ab580b2e`](https://github.com/sgl-project/sglang/commit/ab580b2e060059c0f8387b598a9f648dcec13f93), published as the fixed tag `lmsysorg/sglang:nightly-dev-20261020-ab580b2e`; use that tag in place of `dev` to reproduce the validated run. |\n", "| `--gpus all` | Makes every GPU visible; select one with `--gpus '\"device=0\"'`. With N GPUs, add `--num-gpus N --ulysses-degree N --sp-degree N` to `sglang serve`: Ulysses sequence parallelism splits every chunk across the GPUs and speeds up generation. Sim-Auto supports no other sequence parallelism, so pass both degrees explicitly, and N must divide the checkpoint's 8 key/value heads (2, 4, or 8). `--tp-size` can be combined when the model does not fit on one GPU, as long as `tp_size × ulysses_degree` still divides 8. |\n", "| `--host 0.0.0.0` | Accepts connections from outside the container. For a remote notebook, set `SGLANG_BASE_URL`, or forward the port with `ssh -L 30000:127.0.0.1:30000 `. |\n", "\n", @@ -438,9 +438,9 @@ }, { "cell_type": "code", + "execution_count": null, "id": "sim-auto-sglang-sample", "metadata": {}, - "execution_count": null, "outputs": [], "source": [ "SAMPLE_ROOT = fetch_sample(AUTO_SAMPLE_ROOT)\n", diff --git a/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_daytime_sample.py b/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_daytime_sample.py index f0e15e192..ac9692884 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_daytime_sample.py +++ b/.agents/skills/cookbooks/cosmos3/simulation/auto/tests/test_daytime_sample.py @@ -255,6 +255,9 @@ def test_sglang_lidar_decode_follows_return_output_and_unknown_backends_reject(s prepare_request(manifest, self.root, "model", backend="sglang") with self.assertRaisesRegex(ValueError, "backend"): prepare_request(self.resolve("transfer_auto_wsm.json"), self.root, "model", backend="trt") + with self.assertRaisesRegex(ValueError, "backend"): + submit_joint(None, form, [], self.root / "unsubmitted", backend="trt") + self.assertFalse((self.root / "unsubmitted").exists()) # Rejected before any request or output. def test_stage_conditions_copies_each_camera_under_its_key(self): views = self.resolve("transfer_auto_joint.json")["multiview"]["views"] diff --git a/.agents/skills/cookbooks/cosmos3/simulation/bimanual/bimanual_helpers.py b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/bimanual_helpers.py index b17e5b13a..868112dcb 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/bimanual/bimanual_helpers.py +++ b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/bimanual_helpers.py @@ -48,9 +48,15 @@ class BimanualInput: def _integer(value, name: str, minimum: int = 1) -> int: - if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + """Return ``value`` as an ``int``; integral floats such as ``30.0`` are accepted, booleans are not.""" + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not float(value).is_integer() + or value < minimum + ): raise ValueError(f"{name} must be an integer >= {minimum}.") - return value + return int(value) def _asset(value, base_dir: Path) -> Path: diff --git a/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_sglang.ipynb b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_sglang.ipynb index e8e37bd0d..79d74593f 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_sglang.ipynb +++ b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_sglang.ipynb @@ -39,12 +39,14 @@ "\n", "The `conditioning` contract (schema 5) must be preserved. It defines the 59-value `unified_v1` row that every embodiment shares, the embodiment table (`agibotworld` is domain 15, `camera_pose` is domain 2), the `backward_chunk_anchored_16f` pose convention, and the `global_asinh_unified_v1` normalizer that the published sidecars already apply. The server pads each row to 64 values and routes it to the domain named for that row. Chunk and window sizes count latent frames; one latent frame corresponds to four video frames. The training loader sent prompts as plain text, so the server tokenizes the prompt verbatim unless a request asks for the JSON caption.\n", "\n", - "The server runs in the published SGLang image, `lmsysorg/sglang:dev`. Pull it once on the server host:\n", + "The server runs in the published SGLang image, `lmsysorg/sglang:dev`. Pull it on the server host before each start:\n", "\n", "```bash\n", "docker pull lmsysorg/sglang:dev\n", "```\n", "\n", + "The `dev` tag moves with SGLang's main branch. This notebook was validated against the image built from SGLang commit [`ab580b2e`](https://github.com/sgl-project/sglang/commit/ab580b2e060059c0f8387b598a9f648dcec13f93), published as the fixed tag `lmsysorg/sglang:nightly-dev-20261020-ab580b2e`; use that tag in place of `dev` to reproduce the validated run.\n", + "\n", "The notebook client needs FFmpeg with `ffprobe` for output verification." ] }, @@ -108,8 +110,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "bimanual-sglang-config", "metadata": {}, + "outputs": [], "source": [ "from pathlib import Path\n", "from fractions import Fraction\n", @@ -162,9 +166,7 @@ "print(\"Model:\", MODEL_ID, \"@\", MODEL_REVISION)\n", "print(\"Example:\", MANIFEST_NAME, \"/\", ACTION_BASENAME, \"@\", EXAMPLE_REVISION)\n", "print(\"Outputs:\", OUTPUT_DIR)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -178,8 +180,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "bimanual-sglang-verify", "metadata": {}, + "outputs": [], "source": [ "with httpx.Client(base_url=SGLANG_BASE_URL, headers=HEADERS, timeout=30) as client:\n", " models_response = client.get(\"/v1/models\")\n", @@ -193,23 +197,21 @@ "ACTIVE_MODEL = MODEL_ID\n", "print(\"active model:\", ACTIVE_MODEL)\n", "print(json.dumps(models_payload, indent=2))" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "code", + "execution_count": null, "id": "bimanual-sglang-imports", "metadata": {}, + "outputs": [], "source": [ "from huggingface_hub import hf_hub_download\n", "\n", "from bimanual_helpers import domain_runs, prepare_bimanual_input, submit_bimanual\n", "\n", "# Preparation is local CPU work. The server applies its checkpoint contract." - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -231,8 +233,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "bimanual-sglang-download", "metadata": {}, + "outputs": [], "source": [ "manifest_path = Path(hf_hub_download(repo_id=MODEL_ID, filename=MANIFEST_NAME, revision=EXAMPLE_REVISION))\n", "rows = [json.loads(line) for line in manifest_path.read_text().splitlines() if line.strip()]\n", @@ -258,9 +262,7 @@ "print(f\"{row['name']}: {record['num_frames']} frames at {record['fps']} FPS, seed {record['seed']}\")\n", "print(\"Actions:\", ACTION_PATH)\n", "print(\"Initial image:\", VISION_PATH)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -286,8 +288,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "bimanual-sglang-actions", "metadata": {}, + "outputs": [], "source": [ "sidecar = json.loads(ACTION_PATH.read_text(encoding=\"utf-8\"))\n", "assert sidecar[\"action_space\"] == \"normalized_unified_v1\", sidecar.get(\"action_space\")\n", @@ -304,9 +308,7 @@ " assert math.isfinite(value), f\"row {row_index}: non-finite value\"\n", "print(f\"Validated {len(actions)}×{len(actions[0])} finite actions from {ACTION_PATH}\")\n", "print(\"Domain sequence:\", [(name, end - start) for name, start, end in domain_runs(domain_names)])" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -320,8 +322,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "bimanual-sglang-prepare", "metadata": {}, + "outputs": [], "source": [ "humanoid_input = prepare_bimanual_input(record, base_dir=ACTION_PATH.parent)\n", "assert humanoid_input.action.shape == (record[\"num_frames\"] - 1, 59)\n", @@ -332,9 +336,7 @@ "print(\"Canvas the server encodes:\", (humanoid_input.width, humanoid_input.height))\n", "print(\"Expected output size:\", (humanoid_input.output_width, humanoid_input.output_height))\n", "print(\"Domain runs:\", domain_runs(humanoid_input.domain_names))" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -348,8 +350,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "bimanual-sglang-preview", "metadata": {}, + "outputs": [], "source": [ "action_values = humanoid_input.action\n", "groups = [\n", @@ -377,9 +381,7 @@ "fig.colorbar(heatmap, ax=action_ax, label=\"Checkpoint-normalized value\", shrink=0.8)\n", "display(fig)\n", "plt.close(fig)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -395,8 +397,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "bimanual-sglang-generate", "metadata": {}, + "outputs": [], "source": [ "prepared = humanoid_input\n", "output = OUTPUT_DIR / row[\"name\"] / f\"{row['name']}.mp4\"\n", @@ -408,9 +412,7 @@ " )\n", "elapsed_seconds = time.monotonic() - started\n", "print(f\"{row['name']}: {elapsed_seconds:.1f}s — {output}\")" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -426,8 +428,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "bimanual-sglang-verify-output", "metadata": {}, + "outputs": [], "source": [ "probe = subprocess.run([\n", " \"ffprobe\", \"-v\", \"error\", \"-count_frames\", \"-select_streams\", \"v:0\",\n", @@ -439,19 +443,17 @@ "assert (int(stream[\"width\"]), int(stream[\"height\"])) == (prepared.output_width, prepared.output_height), stream\n", "subprocess.run([\"ffmpeg\", \"-nostdin\", \"-v\", \"error\", \"-xerror\", \"-i\", str(output), \"-f\", \"null\", \"-\"], check=True)\n", "print(stream)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "code", + "execution_count": null, "id": "bimanual-sglang-play", "metadata": {}, + "outputs": [], "source": [ "display(Video(str(output), embed=True, html_attributes=\"controls loop\"))" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -469,13 +471,21 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3", + "display_name": ".venv", "language": "python", "name": "python3" }, "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", "name": "python", - "pygments_lexer": "ipython3" + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.3" } }, "nbformat": 4, diff --git a/.agents/skills/cookbooks/cosmos3/simulation/transfer/run_transfer_with_sglang.ipynb b/.agents/skills/cookbooks/cosmos3/simulation/transfer/run_transfer_with_sglang.ipynb index 6d6170970..f1dc5e176 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/transfer/run_transfer_with_sglang.ipynb +++ b/.agents/skills/cookbooks/cosmos3/simulation/transfer/run_transfer_with_sglang.ipynb @@ -54,8 +54,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "transfer-sglang-03", "metadata": {}, + "outputs": [], "source": [ "from fractions import Fraction\n", "from pathlib import Path\n", @@ -88,9 +90,7 @@ "print(\"Outputs:\", OUTPUT_DIR)\n", "if not OUTPUT_DIR.is_relative_to(COOKBOOK_ROOT):\n", " print(f\"Note: {OUTPUT_DIR} is outside the cookbook root; the server must see it at the same path.\")" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -120,13 +120,17 @@ " --host 0.0.0.0\n", "```\n", "\n", + "The `dev` tag moves with SGLang's main branch, so pull it before each start. This notebook was validated against the image built from SGLang commit [`ab580b2e`](https://github.com/sgl-project/sglang/commit/ab580b2e060059c0f8387b598a9f648dcec13f93), published as the fixed tag `lmsysorg/sglang:nightly-dev-20261020-ab580b2e`; use that tag in place of `dev` to reproduce the validated run.\n", + "\n", "The first start downloads the checkpoint into the mounted cache. The HTTP port opens only after the weights are loaded and the warmup requests finish, so wait for the `Uvicorn running` line. The container mounts the cookbook root (`.agents/skills`) at its host path because the notebook sends each control video as a server-visible `control_path` rather than uploading it; the downloaded `example2` controls live under that root in the output folder. The next cell waits until `/v1/models` lists the model." ] }, { "cell_type": "code", + "execution_count": null, "id": "transfer-sglang-05", "metadata": {}, + "outputs": [], "source": [ "deadline = time.monotonic() + 1800\n", "while True:\n", @@ -141,9 +145,7 @@ " time.sleep(3)\n", "assert MODEL_ID in [m[\"id\"] for m in models[\"data\"]], models\n", "print(\"Server ready:\", BASE_URL)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -166,8 +168,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "transfer-sglang-07", "metadata": {}, + "outputs": [], "source": [ "from huggingface_hub import snapshot_download\n", "\n", @@ -202,9 +206,7 @@ " width=WIDTH, height=HEIGHT, fps=FPS, seed=SEED, input_metadata=metadata)\n", " print(f\"{hint}: {NUM_FRAMES} frames, {WIDTH}×{HEIGHT}, {FPS} FPS, seed {SEED}; control {metadata['width']}×{metadata['height']}, \"\n", " f\"{metadata['nb_read_frames']} source frames; prompt {len(prompt):,} characters\")" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -218,8 +220,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "transfer-sglang-09", "metadata": {}, + "outputs": [], "source": [ "def read_frame(path, frame_index):\n", " from io import BytesIO\n", @@ -238,9 +242,7 @@ " ax.axis(\"off\")\n", "display(fig)\n", "plt.close(fig)" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -256,8 +258,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "transfer-sglang-11", "metadata": {}, + "outputs": [], "source": [ "GENERATION_TIMEOUT = 1800 # Maximum job wait per example, in seconds.\n", "POLL_INTERVAL = 10\n", @@ -330,9 +334,7 @@ " outputs[name] = output\n", " elapsed_seconds[name] = time.monotonic() - started\n", " print(f\"{name}: {elapsed_seconds[name]:.1f}s — {output}\")" - ], - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", @@ -346,15 +348,15 @@ }, { "cell_type": "code", + "execution_count": null, "id": "transfer-sglang-13", "metadata": {}, + "outputs": [], "source": [ "for name, output in outputs.items():\n", " print(name)\n", " display(Video(str(output), embed=True, width=640, html_attributes=\"controls loop\"))" - ], - "execution_count": null, - "outputs": [] + ] } ], "metadata": { From aae6241b2f8b43ff8f304c5c33c953fa31de2b87 Mon Sep 17 00:00:00 2001 From: Kedi Wu Date: Thu, 22 Oct 2026 17:19:57 -0700 Subject: [PATCH 5/5] fix description for bimanual notebook --- .../bimanual/run_bimanual_with_sglang.ipynb | 12 ++++++------ .../bimanual/tests/test_bimanual_helpers.py | 2 +- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_sglang.ipynb b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_sglang.ipynb index 79d74593f..beecb8e21 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_sglang.ipynb +++ b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/run_bimanual_with_sglang.ipynb @@ -16,7 +16,7 @@ "source": [ "# Cosmos3-Nano-Sim-Bimanual with SGLang\n", "\n", - "Replay the published `agibot_mixed_028` example of `nvidia/Cosmos3-Nano-Sim-Bimanual` through a persistent SGLang server: AgiBot manipulation, camera movement, then AgiBot manipulation again. Its 900 normalized unified 59D actions produce **901 frames at 30 FPS** with seed **2026** from the published initial image, a 640×360 frame. The server pads the frame into the trained 832×480 canvas for VAE encoding and drops the padded latents, so the output keeps the image's size on the 16-pixel latent grid, 640×352. The notebook uses an HTTP client and local CPU helpers; the SGLang server owns the model and checkpoint configuration.\n", + "Replay the published `agibot_mixed_028` example of `nvidia/Cosmos3-Nano-Sim-Bimanual` through a persistent SGLang server: a mixed-domain sequence in which AgiBot manipulation rows alternate with camera-pose rows, each row tagged with its embodiment. Its 900 normalized unified 59D actions produce **901 frames at 30 FPS** with seed **2026** from the published initial image, a 640×360 frame. The server pads the frame into the trained 832×480 canvas for VAE encoding and drops the padded latents, so the output keeps the image's size on the 16-pixel latent grid, 640×352. The notebook uses an HTTP client and local CPU helpers; the SGLang server owns the model and checkpoint configuration.\n", "\n", "The model predicts the visual result of supplied actions. It does not choose robot actions or feed generated observations back into a policy.\n", "\n", @@ -271,7 +271,7 @@ "source": [ "### Action structure\n", "\n", - "The sidecar has `action_space: \"normalized_unified_v1\"`, a `900×59` action array, and one domain name per row. The first 300 rows use `agibotworld`, the next 300 use `camera_pose`, and the last 300 return to `agibotworld`. The checkpoint maps these to domain IDs 15, 2, and 15. The initial image plus 900 action rows produce 901 frames. The 59 columns are grouped as follows (Python slices exclude the end index):\n", + "The sidecar has `action_space: \"normalized_unified_v1\"`, a `900×59` action array, and one domain name per row: `agibotworld` for manipulation rows and `camera_pose` for camera-motion rows, which the checkpoint maps to domain IDs 15 and 2. The rows alternate between the two domains in runs of varying length, and the next cell prints that run sequence as `(domain, row count)` pairs. At the release revision of the model repository the layout is a short camera lead-in, two long blocks that switch at rows 300 and 600, and several brief alternations over the final six seconds; a sidecar with a different layout needs no notebook changes, because the domain names travel with the request. The initial image plus 900 action rows produce 901 frames. The 59 columns are grouped as follows (Python slices exclude the end index):\n", "\n", "| Columns | Contents |\n", "| --- | --- |\n", @@ -317,7 +317,7 @@ "source": [ "### Prepare the input\n", "\n", - "The initial image provides the starting observation. The 600 robot actions and 300 camera actions drive the following frames, producing a 901-frame video. The helper keeps the rows as they are and sends one embodiment name per row; the server applies no further normalization. It also reports the canvas the server encodes and the expected output size: for the published 640×360 image, 832×480 and 640×352." + "The initial image provides the starting observation. The 900 action rows, each routed to the domain named for it, drive the following frames, producing a 901-frame video. The helper keeps the rows as they are and sends one embodiment name per row; the server applies no further normalization. The cell also prints the domain runs as `(domain, first_row, end_row)` with the end excluded, which locate every switch in the output. It also reports the canvas the server encodes and the expected output size: for the published 640×360 image, 832×480 and 640×352." ] }, { @@ -345,7 +345,7 @@ "source": [ "### Preview the inputs\n", "\n", - "The image is the starting observation. The heatmap shows all 59 normalized action channels over the 30-second sequence, grouped as in the table above. Dashed lines mark the domain switches at 10 and 20 seconds." + "The image is the starting observation. The heatmap shows all 59 normalized action channels over the 30-second sequence, grouped as in the table above. Dashed lines mark every domain switch; divide a run's first row by the FPS to get its time in seconds." ] }, { @@ -423,7 +423,7 @@ "\n", "Decode the generated video and check that it contains the requested frame count at the requested FPS and the expected output size. For the published image, the output keeps the 640-pixel width; its 360-pixel height becomes 352, the largest multiple of the 16-pixel latent stride, because the padded latent rows are dropped rather than decoded.\n", "\n", - "Review the three phases: robot manipulation, a camera position/orientation loop, then resumed robot manipulation. The domain switches are at approximately 10 and 20 seconds." + "Review the rollout against the domain runs printed in step 4: manipulation should follow the `agibotworld` rows and camera motion the `camera_pose` rows, with each switch at its run's first row divided by the FPS. At the release revision the two longest blocks switch at 10 and 20 seconds, after a camera lead-in of under a second, and the final six seconds alternate several times." ] }, { @@ -464,7 +464,7 @@ "\n", "- **Run the other published examples:** set `COSMOS3_BIMANUAL_MANIFEST` to `examples/inference/sim_bimanual_unifiedv1_hf.jsonl` and `COSMOS3_BIMANUAL_ACTION` to a sidecar basename from it, such as `agibot_001_normalized_unifiedv1.json`, then restart the notebook.\n", "- **Use your own recording:** set `ACTION_PATH` to a `normalized_unified_v1` sidecar, `VISION_PATH` to the initial image or a video whose first frame is the initial observation, and `HUMANOID_PROMPT` to the task description, then restart the notebook. Keep `num_frames - 1` equal to the number of action rows and divisible by four, and do not normalize the values again. The [Bimanual guide](bimanual.md#download-the-model-and-examples) describes the sidecar contract.\n", - "- **Validate further:** apply the [shared Simulation checklist](../README.md#validate-a-rollout). Inspect the complete rollout for frozen frames, discontinuities, and disagreement with the supplied actions, including both domain transitions around 10 and 20 seconds. Use a new output directory when changing the checkpoint, inputs, or server settings.\n", + "- **Validate further:** apply the [shared Simulation checklist](../README.md#validate-a-rollout). Inspect the complete rollout for frozen frames, discontinuities, and disagreement with the supplied actions, including every domain transition in the printed domain runs. Use a new output directory when changing the checkpoint, inputs, or server settings.\n", "- **Other runtimes:** the [vLLM-Omni notebook](run_bimanual_with_vllm_omni.ipynb) serves the same published example through the synchronous Videos API, and the [Bimanual guide](bimanual.md#run-inference) runs it natively with Cosmos Framework." ] } diff --git a/.agents/skills/cookbooks/cosmos3/simulation/bimanual/tests/test_bimanual_helpers.py b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/tests/test_bimanual_helpers.py index b5a1f8842..2e7b4535b 100644 --- a/.agents/skills/cookbooks/cosmos3/simulation/bimanual/tests/test_bimanual_helpers.py +++ b/.agents/skills/cookbooks/cosmos3/simulation/bimanual/tests/test_bimanual_helpers.py @@ -373,7 +373,7 @@ def test_extra_params_round_trip_float32_rows_in_compact_form(self): self.assertLess(len(text), 0.75 * len(naive)) # Shortest float32 decimals rather than double repr. def test_published_rollout_length_stays_within_the_form_field_limit(self): - # The published mixed-action example has 900 rows: 300 robot, 300 camera, 300 robot. + # 900 rows alternating between the robot and camera domains, the length of the published mixed-action example. frames = 901 names = ["agibotworld"] * 300 + ["camera_pose"] * 300 + ["agibotworld"] * 300 self.write_sidecar("long.json", make_rows(frames - 1, seed=1), names)