diff --git a/cookbooks/cosmos3/generator/action/run_fd_with_sglang.ipynb b/cookbooks/cosmos3/generator/action/run_fd_with_sglang.ipynb index dc93da25..ce50bc1d 100644 --- a/cookbooks/cosmos3/generator/action/run_fd_with_sglang.ipynb +++ b/cookbooks/cosmos3/generator/action/run_fd_with_sglang.ipynb @@ -32,7 +32,6 @@ "docker run -d --init --name cosmos3-sglang-notebook \\\n", " --runtime nvidia --gpus '\"device=0\"' \\\n", " -e CUDA_DEVICE_ORDER=PCI_BUS_ID \\\n", - " -v ~/.cache/huggingface:/root/.cache/huggingface \\\n", " -v \"$PWD:/workspace\" \\\n", " -p 30000:30000 --ipc=host \\\n", " lmsysorg/sglang:dev \\\n", @@ -45,6 +44,27 @@ "```\n" ] }, + { + "cell_type": "markdown", + "id": "fdvl-install-md", + "metadata": {}, + "source": [ + "## Install Notebook Dependencies\n", + "\n", + "Run this cell in the notebook kernel before the data-loading and visualization cells.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fdvl-install-code", + "metadata": {}, + "outputs": [], + "source": [ + "%pip install matplotlib pandas pyarrow \"lerobot[dataset]\"\n", + "%pip install av==18.0.0\n" + ] + }, { "cell_type": "markdown", "id": "fdvl-vars-md", @@ -52,7 +72,15 @@ "source": [ "## Configure Notebook Variables\n", "\n", - "Run this cell after the SGLang server is available. It resolves local input/output paths and stores generated outputs under `outputs/cosmos3_action_sglang/` by default.\n" + "Run this cell after the SGLang server is available. It resolves local input/output paths and stores generated outputs under `outputs/cosmos3_action_sglang/` by default.\n", + "\n", + "Before running it, set `COSMOS3_REPO` to the root of a local checkout of the **cosmos-framework** repository:\n", + "\n", + "```bash\n", + "export COSMOS3_REPO=/path/to/cosmos-framework\n", + "```\n", + "\n", + "This must be the framework checkout that contains the `cosmos_framework/` Python package—not this `cosmos` cookbook repository, a model directory, or the `packages/cosmos3` model sources. The notebook adds this path to Python's import path to load `cosmos_framework`. Set it in the shell that launches Jupyter before running this cell.\n" ] }, { @@ -462,7 +490,7 @@ "def make_preview(src: Path, crf: int = 28) -> Path:\n", " \"\"\"Re-encode `src` to a compact, browser-friendly mp4 (cached).\"\"\"\n", " preview = src.with_name(f\"{src.stem}_preview.mp4\")\n", - " if not preview.exists():\n", + " if not preview.exists() or preview.stat().st_mtime < src.stat().st_mtime:\n", " subprocess.run(\n", " [FFMPEG, \"-y\", \"-loglevel\", \"error\", \"-i\", str(src),\n", " \"-c:v\", \"libx264\", \"-crf\", str(crf),\n", @@ -488,7 +516,9 @@ "source": [ "## Robotics\n", "\n", - "In this example, we show how to start from a LeRobot dataset of DROID and run **multiview** generation for robotics manipulation **autoregressively**.\n" + "In this example, we show how to start from a LeRobot dataset of DROID and run **multiview** generation for robotics manipulation **autoregressively**.\n", + "\n", + "> Note: Before running this LeRobot section, PyAV may need to be reinstalled as `av==15.1.0`.\n" ] }, { diff --git a/cookbooks/cosmos3/generator/action/run_fd_with_vllm_omni.ipynb b/cookbooks/cosmos3/generator/action/run_fd_with_vllm_omni.ipynb index acfea711..8bb25c69 100644 --- a/cookbooks/cosmos3/generator/action/run_fd_with_vllm_omni.ipynb +++ b/cookbooks/cosmos3/generator/action/run_fd_with_vllm_omni.ipynb @@ -31,6 +31,27 @@ "Forward dynamics predicts future visual observations from an initial image and an action trajectory. This notebook contains separate AV, robotics, UMI, and human hand-pose sections that each build their own input spec, run inference, and visualize generated videos.\n" ] }, + { + "cell_type": "markdown", + "id": "fdvl-install-md", + "metadata": {}, + "source": [ + "## Install Notebook Dependencies\n", + "\n", + "Run this cell in the notebook kernel before the data-loading and visualization cells.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fdvl-install-code", + "metadata": {}, + "outputs": [], + "source": [ + "%pip install matplotlib pandas pyarrow \"lerobot[dataset]\"\n", + "%pip install av==18.0.0\n" + ] + }, { "cell_type": "markdown", "id": "4100dead", @@ -96,7 +117,20 @@ "source": [ "## Configure Notebook Variables\n", "\n", - "Run this cell after the vLLM-Omni server is available. It resolves local input/output paths and stores generated outputs under `outputs/cosmos3_action_vllm/` by default.\n" + "Run this cell after the vLLM-Omni server is available. It resolves local input/output paths and stores generated outputs under `outputs/cosmos3_action_vllm/` by default.\n", + "\n", + "The configuration supports these environment variables:\n", + "\n", + "- `COSMOS3_VLLM_OUTPUT_ROOT`: optional override for the output directory. The notebook resolves this into the `COSMOS3_OUTPUT_ROOT` Python variable, which contains generated inputs, chunk outputs, and stitched videos.\n", + "- `COSMOS3_DROID_ROOT`: optional override for the DROID LeRobot dataset. Set it to a version-named dataset root such as `/path/to/droid_plus_lerobot_640x360_20260412`, containing the `success/` split. If unset, the notebook creates a versioned alias for the bundled `assets/droid_lerobot_example` sample.\n", + "\n", + "Before running it, set `COSMOS3_REPO` to the root of a local checkout of the **cosmos-framework** repository:\n", + "\n", + "```bash\n", + "export COSMOS3_REPO=/path/to/cosmos-framework\n", + "```\n", + "\n", + "This must be the framework checkout that contains the `cosmos_framework/` Python package—not this `cosmos` cookbook repository, a model directory, or the `packages/cosmos3` model sources. The notebook adds this path to Python's import path to load `cosmos_framework`. Set it in the shell that launches Jupyter before running this cell.\n" ] }, { @@ -108,6 +142,7 @@ "source": [ "from pathlib import Path\n", "import os\n", + "import json\n", "\n", "\n", "def find_repo_root(start: Path) -> Path:\n", @@ -140,7 +175,42 @@ "print(\"COSMOS3_REPO:\", COSMOS3_REPO)\n", "print(\"COSMOS3_INPUT_DIR:\", COSMOS3_INPUT_DIR)\n", "print(\"COSMOS3_OUTPUT_ROOT:\", COSMOS3_OUTPUT_ROOT)\n", - "print(\"COSMOS3_VLLM_BASE_URL:\", VLLM_BASE_URL)\n" + "print(\"COSMOS3_VLLM_BASE_URL:\", VLLM_BASE_URL)\n", + "\n", + "\n", + "_ACTION_VIEWPOINT_FRAMING = {\n", + " \"ego_view\": \"This video is captured from a first-person perspective looking at the scene.\",\n", + " \"third_person_view\": \"This video is captured from a third-person perspective looking towards the agent from the front.\",\n", + " \"wrist_view\": \"This video is captured from a wrist-mounted camera.\",\n", + " \"concat_view\": \"This video contains concatenated views from multiple camera perspectives.\",\n", + "}\n", + "_ACTION_ASPECT_RATIOS = {\"1,1\": 1.0, \"4,3\": 4 / 3, \"3,4\": 3 / 4, \"16,9\": 16 / 9, \"9,16\": 9 / 16}\n", + "\n", + "\n", + "def build_action_prompt(record: dict, *, width: int, height: int, default_prompt: str) -> str:\n", + " description = str(record.get(\"prompt\") or \"\").strip() or default_prompt\n", + " if not description.endswith((\".\", \"!\", \"?\")):\n", + " description += \".\"\n", + "\n", + " num_frames = int(record[\"action_chunk_size\"]) + 1\n", + " fps = int(record[\"fps\"])\n", + " duration_seconds = num_frames / fps\n", + " aspect_ratio = min(_ACTION_ASPECT_RATIOS, key=lambda key: abs(width / height - _ACTION_ASPECT_RATIOS[key]))\n", + " minutes, seconds = divmod(round(duration_seconds), 60)\n", + " action_prompt = {}\n", + " framing = _ACTION_VIEWPOINT_FRAMING.get(record.get(\"view_point\"))\n", + " if framing is not None:\n", + " action_prompt[\"cinematography\"] = {\"framing\": framing}\n", + " action_prompt.update(\n", + " {\n", + " \"actions\": [{\"time\": f\"0:00-{minutes}:{seconds:02d}\", \"description\": description}],\n", + " \"duration\": f\"{int(duration_seconds)}s\",\n", + " \"fps\": float(fps),\n", + " \"resolution\": {\"H\": height, \"W\": width},\n", + " \"aspect_ratio\": aspect_ratio,\n", + " }\n", + " )\n", + " return json.dumps(action_prompt)\n" ] }, { @@ -412,11 +482,12 @@ " \"domain_name\": record[\"domain_name\"],\n", " \"action_chunk_size\": record[\"action_chunk_size\"],\n", " \"image_size\": record[\"image_size\"],\n", - " \"view_point\": record[\"view_point\"],\n", " \"action\": json.loads(Path(record[\"action_path\"]).read_text()),\n", " \"guardrails\": False,\n", " }\n", - " prompt = str(record.get(\"prompt\") or \"\").strip() or \"A robot manipulates an object.\"\n", + " prompt = build_action_prompt(\n", + " record, width=input_width, height=input_height, default_prompt=\"A robot manipulates an object.\"\n", + " )\n", " form = {\n", " \"prompt\": prompt,\n", " \"num_frames\": record[\"action_chunk_size\"] + 1,\n", @@ -509,7 +580,7 @@ "def make_preview(src: Path, crf: int = 28) -> Path:\n", " \"\"\"Re-encode `src` to a compact, browser-friendly mp4 (cached).\"\"\"\n", " preview = src.with_name(f\"{src.stem}_preview.mp4\")\n", - " if not preview.exists():\n", + " if not preview.exists() or preview.stat().st_mtime < src.stat().st_mtime:\n", " subprocess.run(\n", " [FFMPEG, \"-y\", \"-loglevel\", \"error\", \"-i\", str(src),\n", " \"-c:v\", \"libx264\", \"-crf\", str(crf),\n", @@ -535,7 +606,11 @@ "source": [ "## Robotics\n", "\n", - "In this example, we show how to start from a LeRobot dataset of DROID and run **multiview** generation for robotics manipulation **autoregressively**.\n" + "In this example, we show how to start from a LeRobot dataset of DROID and run **multiview** generation for robotics manipulation **autoregressively**.\n", + "\n", + "> Note: Before running this LeRobot section, PyAV may need to be reinstalled as `av==15.1.0`.\n", + "\n", + "`DROIDLeRobotDataset` resolves its feature configuration from the dataset root name. The bundled sample uses a friendly asset name, so the setup cell exposes it through the matching `droid_plus_lerobot_640x360_20260412/success` layout. To use another dataset, set `COSMOS3_DROID_ROOT` to its version-named root directory.\n" ] }, { @@ -569,10 +644,39 @@ "\n", "from cosmos_framework.data.generator.action.datasets import DROIDLeRobotDataset\n", "\n", - "robotics_dataset_root = resolve_input(\"cookbooks/cosmos3/generator/action/assets/droid_lerobot_example\")\n", - "robotics_dataset = DROIDLeRobotDataset(root=robotics_dataset_root)\n", "robotics_num_chunks = 5\n", "robotics_chunk_length = 16\n", + "\n", + "# `DROIDLeRobotDataset` selects its feature configuration from the root directory\n", + "# name. The bundled sample comes from the 640x360 release but is stored under a\n", + "# friendly asset name, so expose it through the release layout expected by the loader.\n", + "droid_root_override = os.environ.get(\"COSMOS3_DROID_ROOT\")\n", + "if droid_root_override:\n", + " robotics_dataset_root = str(Path(droid_root_override).expanduser().resolve())\n", + "else:\n", + " bundled_droid_root = Path(resolve_input(\"cookbooks/cosmos3/generator/action/assets/droid_lerobot_example\"))\n", + " versioned_droid_root = (\n", + " COSMOS3_OUTPUT_ROOT\n", + " / \"datasets\"\n", + " / \"droid_plus_lerobot_640x360_20260412\"\n", + " )\n", + " success_root = versioned_droid_root / \"success\"\n", + " versioned_droid_root.mkdir(parents=True, exist_ok=True)\n", + " if success_root.is_symlink():\n", + " if success_root.resolve() != bundled_droid_root:\n", + " raise RuntimeError(f\"{success_root} points to {success_root.resolve()}, expected {bundled_droid_root}\")\n", + " elif success_root.exists():\n", + " if success_root.resolve() != bundled_droid_root:\n", + " raise RuntimeError(f\"{success_root} already exists and is not the bundled DROID sample\")\n", + " else:\n", + " success_root.symlink_to(bundled_droid_root, target_is_directory=True)\n", + " robotics_dataset_root = str(versioned_droid_root)\n", + "\n", + "robotics_dataset = DROIDLeRobotDataset(\n", + " root=robotics_dataset_root,\n", + " chunk_length=robotics_chunk_length,\n", + " use_success_only=True,\n", + ")\n", "robotics_chunk_starts = [chunk_idx * robotics_chunk_length for chunk_idx in range(robotics_num_chunks)]\n", "assert robotics_chunk_starts[-1] < len(robotics_dataset)\n", "\n", @@ -693,14 +797,15 @@ " \"domain_name\": record[\"domain_name\"],\n", " \"action_chunk_size\": record[\"action_chunk_size\"],\n", " \"image_size\": record[\"image_size\"],\n", - " \"view_point\": record[\"view_point\"],\n", " \"action\": json.loads(Path(record[\"action_path\"]).read_text()),\n", " \"guardrails\": False,\n", " }\n", " if disable_guardrails:\n", " extra_params[\"guardrails\"] = False\n", "\n", - " prompt = str(record.get(\"prompt\") or \"\").strip() or \" \"\n", + " prompt = build_action_prompt(\n", + " record, width=input_width, height=input_height, default_prompt=\"A robot manipulates an object.\"\n", + " )\n", " form = {\n", " \"prompt\": prompt,\n", " \"num_frames\": record[\"action_chunk_size\"] + 1,\n", @@ -1101,11 +1206,12 @@ " \"domain_name\": record[\"domain_name\"],\n", " \"action_chunk_size\": record[\"action_chunk_size\"],\n", " \"image_size\": record[\"image_size\"],\n", - " \"view_point\": record[\"view_point\"],\n", " \"action\": json.loads(Path(record[\"action_path\"]).read_text()),\n", " \"guardrails\": False,\n", " }\n", - " prompt = str(record.get(\"prompt\") or \"\").strip() or \"A robot manipulates an object.\"\n", + " prompt = build_action_prompt(\n", + " record, width=input_width, height=input_height, default_prompt=\"A robot manipulates an object.\"\n", + " )\n", " form = {\n", " \"prompt\": prompt,\n", " \"num_frames\": record[\"action_chunk_size\"] + 1,\n", diff --git a/cookbooks/cosmos3/generator/action/run_id_with_sglang.ipynb b/cookbooks/cosmos3/generator/action/run_id_with_sglang.ipynb index c10b4dff..967f30c0 100644 --- a/cookbooks/cosmos3/generator/action/run_id_with_sglang.ipynb +++ b/cookbooks/cosmos3/generator/action/run_id_with_sglang.ipynb @@ -38,7 +38,6 @@ "docker run -d --init --name cosmos3-sglang-notebook \\\n", " --runtime nvidia --gpus '\"device=0\"' \\\n", " -e CUDA_DEVICE_ORDER=PCI_BUS_ID \\\n", - " -v ~/.cache/huggingface:/root/.cache/huggingface \\\n", " -v \"$PWD:/workspace\" \\\n", " -p 30000:30000 --ipc=host \\\n", " lmsysorg/sglang:dev \\\n", @@ -51,6 +50,22 @@ "```\n" ] }, + { + "cell_type": "markdown", + "id": "idvl-vars-md", + "metadata": {}, + "source": [ + "## Configure Notebook Variables\n", + "\n", + "Before running the variables cell, set `COSMOS3_REPO` to the root of a local checkout of the **cosmos-framework** repository:\n", + "\n", + "```bash\n", + "export COSMOS3_REPO=/path/to/cosmos-framework\n", + "```\n", + "\n", + "This must be the framework checkout that contains the `cosmos_framework/` Python package—not this `cosmos` cookbook repository, a model directory, or the `packages/cosmos3` model sources. The notebook adds this path to Python's import path to load `cosmos_framework`. Set it in the shell that launches Jupyter before running the variables cell.\n" + ] + }, { "cell_type": "code", "execution_count": null, diff --git a/cookbooks/cosmos3/generator/action/run_id_with_vllm_omni.ipynb b/cookbooks/cosmos3/generator/action/run_id_with_vllm_omni.ipynb index fd3a6972..68be088d 100644 --- a/cookbooks/cosmos3/generator/action/run_id_with_vllm_omni.ipynb +++ b/cookbooks/cosmos3/generator/action/run_id_with_vllm_omni.ipynb @@ -86,6 +86,22 @@ "The request cells below use whichever Nano or Edge server is active at the configured endpoint. To keep outputs from separate model runs, set `COSMOS3_VLLM_OUTPUT_ROOT` to a different directory before running the configuration cell.\n" ] }, + { + "cell_type": "markdown", + "id": "idvl-vars-md", + "metadata": {}, + "source": [ + "## Configure Notebook Variables\n", + "\n", + "Before running the variables cell, set `COSMOS3_REPO` to the root of a local checkout of the **cosmos-framework** repository:\n", + "\n", + "```bash\n", + "export COSMOS3_REPO=/path/to/cosmos-framework\n", + "```\n", + "\n", + "This must be the framework checkout that contains the `cosmos_framework/` Python package—not this `cosmos` cookbook repository, a model directory, or the `packages/cosmos3` model sources. The notebook adds this path to Python's import path to load `cosmos_framework`. Set it in the shell that launches Jupyter before running the variables cell.\n" + ] + }, { "cell_type": "code", "execution_count": null, diff --git a/cookbooks/cosmos3/generator/action/run_policy_with_sglang.ipynb b/cookbooks/cosmos3/generator/action/run_policy_with_sglang.ipynb index c2744f98..8831a2d2 100644 --- a/cookbooks/cosmos3/generator/action/run_policy_with_sglang.ipynb +++ b/cookbooks/cosmos3/generator/action/run_policy_with_sglang.ipynb @@ -34,7 +34,13 @@ "source": [ "## Start SGLang Policy Server\n", "\n", - "Start the server in a terminal from the `cosmos` repo root.\n", + "Start the server in a terminal from the `cosmos` repo root. Before running it, set `COSMOS3_REPO` to the root of a local checkout of the **cosmos-framework** repository:\n", + "\n", + "```bash\n", + "export COSMOS3_REPO=/path/to/cosmos-framework\n", + "```\n", + "\n", + "This must be the framework checkout that contains the `cosmos_framework/` Python package—not this `cosmos` cookbook repository, a model directory, or the `packages/cosmos3` model sources. The server container mounts this checkout at `/workspace/cosmos-framework`.\n", "\n", "```bash\n", "docker rm -f cosmos3-sglang-policy-notebook 2>/dev/null || true\n", @@ -43,8 +49,7 @@ " --runtime nvidia --gpus '\"device=0\"' \\\n", " -e CUDA_DEVICE_ORDER=PCI_BUS_ID \\\n", " -e PYTHONPATH=/workspace/cosmos-framework \\\n", - " -v ~/.cache/huggingface:/root/.cache/huggingface \\\n", - " -v \"$PWD:/workspace\" \\\n", + " -v \"$COSMOS3_REPO:/workspace/cosmos-framework\" \\\n", " -p 30000:30000 --ipc=host \\\n", " lmsysorg/sglang:dev \\\n", " sglang serve \\\n", diff --git a/cookbooks/cosmos3/generator/action/run_policy_with_vllm_omni.ipynb b/cookbooks/cosmos3/generator/action/run_policy_with_vllm_omni.ipynb index 7f7affed..0b46cf8c 100644 --- a/cookbooks/cosmos3/generator/action/run_policy_with_vllm_omni.ipynb +++ b/cookbooks/cosmos3/generator/action/run_policy_with_vllm_omni.ipynb @@ -20,11 +20,11 @@ "\n", "Generator requires the Guardrail. Request access to the gated [nvidia/Cosmos-1.0-Guardrail](https://huggingface.co/nvidia/Cosmos-1.0-Guardrail) HF repository before running guarded examples. When running with guardrails enabled, also set and pass `HF_TOKEN` to the container; it must be a Hugging Face token authorized for that repository. To disable guardrails server-wide, pass `--no-guardrails` to `vllm serve`; alternatively, set `guardrails: false` in vLLM-Omni `extra_params` or `extra_args` to disable them per request.\n", "\n", - "## Overview\n", + "Running `nvidia/Cosmos3-Nano-Policy-DROID` also requires an up-to-date NumPy installation. Upgrade it in the environment that starts the model with `pip install --upgrade numpy`.\n", "\n", - "This notebook runs Cosmos3 Nano or Cosmos3-Edge **action policy** inference through vLLM-Omni using the checked-in DROID LeRobot sample under `assets/droid_lerobot_example`.\n", + "## Overview\n", "\n", - "It sends `POST /v1/videos` requests with a first frame and instruction, then retrieves a rollout video plus top-level `action` metadata." + "This notebook demonstrates one-shot rollout generation through the asynchronous `/v1/videos` API for Cosmos3 Nano or Cosmos3-Edge using the checked-in LeRobot sample under `assets/droid_lerobot_example`. It returns a 17-frame rollout and a `[16, 8]` DROID joint-position action chunk." ] }, { @@ -32,9 +32,17 @@ "id": "policy-server-md", "metadata": {}, "source": [ - "## Start vLLM-Omni Policy Server\n", + "## Start vLLM-Omni for One-Shot Policy Rollout\n", + "\n", + "Choose either the Cosmos3-Nano or Cosmos3-Edge DROID checkpoint and start it in a terminal from the `cosmos` repo root. Before running it, set `COSMOS3_REPO` to the root of a local checkout of the **cosmos-framework** repository:\n", + "\n", + "```bash\n", + "export COSMOS3_REPO=/path/to/cosmos-framework\n", + "```\n", + "\n", + "This must be the framework checkout that contains the `cosmos_framework/` Python package—not this `cosmos` cookbook repository, a model directory, or the `packages/cosmos3` model sources. The server container mounts this checkout at `/workspace/cosmos-framework`.\n", "\n", - "Choose either the Cosmos3-Nano or Cosmos3-Edge DROID policy server and start it in a terminal from the `cosmos` repo root. The commands are mutually exclusive: both use the same container name and publish the container's port `8000` to host port `8001`, which matches the notebook's default `http://localhost:8001` endpoint.\n", + "Both commands expose the generic image/video generation APIs, including the asynchronous `/v1/videos` policy-rollout path used first in this notebook. They publish the container's port `8000` to host port `8001`, matching the default `http://localhost:8001` endpoint.\n", "\n", "### Cosmos3-Nano-Policy-DROID\n", "\n", @@ -45,7 +53,7 @@ " --runtime nvidia --gpus '\"device=0\"' \\\n", " -e CUDA_DEVICE_ORDER=PCI_BUS_ID \\\n", " -e PYTHONPATH=/workspace/cosmos-framework \\\n", - " -v \"$PWD:/workspace\" \\\n", + " -v \"$COSMOS3_REPO:/workspace/cosmos-framework\" \\\n", " -p 8001:8000 --ipc=host \\\n", " vllm/vllm-omni:cosmos3 \\\n", " vllm serve nvidia/Cosmos3-Nano-Policy-DROID \\\n", @@ -69,7 +77,7 @@ " --runtime nvidia --gpus '\"device=0\"' \\\n", " -e CUDA_DEVICE_ORDER=PCI_BUS_ID \\\n", " -e PYTHONPATH=/workspace/cosmos-framework \\\n", - " -v \"$PWD:/workspace\" \\\n", + " -v \"$COSMOS3_REPO:/workspace/cosmos-framework\" \\\n", " -p 8001:8000 --ipc=host \\\n", " vllm/vllm-omni:cosmos3 \\\n", " vllm serve nvidia/Cosmos3-Edge-Policy-DROID \\\n", @@ -149,7 +157,7 @@ "import subprocess\n", "\n", "import numpy as np\n", - "from PIL import Image, ImageOps\n", + "from PIL import Image\n", "from IPython.display import display\n", "\n", "try:\n", @@ -182,18 +190,23 @@ "}\n", "frames = {key: Image.open(path).convert(\"RGB\") for key, path in frame_paths.items()}\n", "\n", - "# DROID policy uses a concatenated multi-view frame for the /v1/videos path.\n", - "target_w, target_h = 640, 540\n", - "top_h = target_h // 2\n", - "bottom_h = target_h - top_h\n", - "half_w = target_w // 2\n", - "wrist = ImageOps.fit(frames[\"observation/wrist_image_left\"], (target_w, top_h), method=Image.Resampling.BICUBIC)\n", - "left = ImageOps.fit(frames[\"observation/exterior_image_1_left\"], (half_w, bottom_h), method=Image.Resampling.BICUBIC)\n", - "right = ImageOps.fit(frames[\"observation/exterior_image_2_left\"], (half_w, bottom_h), method=Image.Resampling.BICUBIC)\n", + "# Match vLLM-Omni's compose_robolab_views geometry: preserve the 640x360\n", + "# wrist view on top and resize each exterior view to 320x180 below it.\n", + "wrist = frames[\"observation/wrist_image_left\"]\n", + "source_size = wrist.size\n", + "assert all(frame.size == source_size for frame in frames.values()), (\n", + " \"DROID camera frames must have matching dimensions\",\n", + " {key: frame.size for key, frame in frames.items()},\n", + ")\n", + "target_w = wrist.width\n", + "bottom_h, half_w = wrist.height // 2, wrist.width // 2\n", + "target_h = wrist.height + bottom_h\n", + "left = frames[\"observation/exterior_image_1_left\"].resize((half_w, bottom_h), Image.Resampling.BILINEAR)\n", + "right = frames[\"observation/exterior_image_2_left\"].resize((half_w, bottom_h), Image.Resampling.BILINEAR)\n", "policy_image = Image.new(\"RGB\", (target_w, target_h))\n", "policy_image.paste(wrist, (0, 0))\n", - "policy_image.paste(left, (0, top_h))\n", - "policy_image.paste(right, (half_w, top_h))\n", + "policy_image.paste(left, (0, wrist.height))\n", + "policy_image.paste(right, (half_w, wrist.height))\n", "policy_image_path = COSMOS3_INPUT_DIR / \"droid_policy_first_frame.png\"\n", "policy_image.save(policy_image_path)\n", "\n", @@ -211,9 +224,9 @@ "id": "policy-video-md", "metadata": {}, "source": [ - "## Run Policy Inference Through `/v1/videos`\n", + "## Run One-Shot Policy Inference Through `/v1/videos`\n", "\n", - "This path behaves like the forward/inverse vLLM action notebooks: it sends a multipart request to the OpenAI-compatible video API, polls the async job, writes the generated rollout video, and saves the predicted action from the response metadata." + "This path sends the correctly composed multiview image and instruction to the asynchronous video API, polls the job, writes the generated rollout video, and saves the raw `[16, 8]` DROID joint-position action chunk from the response metadata." ] }, { @@ -304,16 +317,19 @@ "def submit_policy_video(active_model: str) -> dict:\n", " run_dir = COSMOS3_POLICY_OUTPUT_DIR / \"video_api\"\n", " run_dir.mkdir(parents=True, exist_ok=True)\n", + " request_image_path = run_dir / \"policy_input.png\"\n", + " policy_image.save(request_image_path)\n", + " print(\"saved\", request_image_path)\n", "\n", - " input_width, input_height = Image.open(policy_image_path).size\n", + " input_width, input_height = Image.open(request_image_path).size\n", " target_width, target_height = closest_action_size(input_height, input_width)\n", " extra_params = {\n", " \"action_mode\": \"policy\",\n", " \"domain_name\": \"droid_lerobot\",\n", - " \"raw_action_dim\": 10,\n", + " # Released DROID policy checkpoints predict 7 joints + 1 gripper.\n", + " \"raw_action_dim\": 8,\n", " \"action_chunk_size\": 16,\n", " \"image_size\": 480,\n", - " \"view_point\": \"concat_view\",\n", " \"guardrails\": False,\n", " }\n", " is_edge_policy = \"Cosmos3-Edge-Policy-DROID\" in active_model\n", @@ -334,11 +350,11 @@ " \"extra_params\": json.dumps(extra_params),\n", " }\n", "\n", - " with policy_image_path.open(\"rb\") as image_file:\n", + " with request_image_path.open(\"rb\") as image_file:\n", " response = requests.post(\n", " f\"{VLLM_BASE_URL}/v1/videos\",\n", " data={key: str(value) for key, value in form.items()},\n", - " files={\"input_reference\": (policy_image_path.name, image_file, \"image/png\")},\n", + " files={\"input_reference\": (request_image_path.name, image_file, \"image/png\")},\n", " timeout=120,\n", " )\n", " if not response.ok:\n", @@ -366,6 +382,11 @@ " action = final.get(\"action\")\n", " if not action or \"data\" not in action:\n", " raise RuntimeError(f\"vLLM response did not include action data: {json.dumps(final, indent=2)}\")\n", + " action_array = np.asarray(action[\"data\"], dtype=np.float32)\n", + " if action_array.shape != (16, 8):\n", + " raise RuntimeError(f\"Expected a [16, 8] DROID action chunk, got {action_array.shape}\")\n", + " if not np.isfinite(action_array).all():\n", + " raise RuntimeError(\"DROID action response contains non-finite values\")\n", " (run_dir / \"action.json\").write_text(json.dumps(action, indent=2))\n", " sample_outputs = {\"outputs\": [{\"content\": {\"action\": action[\"data\"]}}]}\n", " (run_dir / \"sample_outputs.json\").write_text(json.dumps(sample_outputs, indent=2))\n", @@ -382,7 +403,7 @@ "\n", " print(\"saved\", run_dir / \"action.json\")\n", " print(\"action shape:\", action.get(\"shape\"), \"dtype:\", action.get(\"dtype\"), \"domain_id:\", action.get(\"domain_id\"))\n", - " return {\"initial\": initial, \"final\": final, \"run_dir\": run_dir, \"video_path\": video_path, \"action\": action}\n", + " return {\"initial\": initial, \"final\": final, \"run_dir\": run_dir, \"input_image_path\": request_image_path, \"video_path\": video_path, \"action\": action}\n", "\n", "\n", "active_vllm_model = check_vllm_server()\n", @@ -390,9 +411,9 @@ "policy_video_result = submit_policy_video(active_vllm_model)" ] }, - { - "cell_type": "markdown", - "id": "policy-preview-md", + { + "cell_type": "markdown", + "id": "policy-preview-md", "metadata": {}, "source": [ "## Inspect Video API Outputs\n",