diff --git a/.agents/skills/cookbooks/cosmos3/README.md b/.agents/skills/cookbooks/cosmos3/README.md index bfa07955e..000fe8dff 100644 --- a/.agents/skills/cookbooks/cosmos3/README.md +++ b/.agents/skills/cookbooks/cosmos3/README.md @@ -11,6 +11,7 @@ backend you want to run and follow that one section. | [Diffusers](#diffusers) | Direct generation with `Cosmos3OmniPipeline` | Generator (Audiovisual) | | [TensorRT-LLM Generator](#tensorrt-llm-generator) | OpenAI-compatible VisualGen server (image/video/audio/action/transfer generation) | Generator (Audiovisual, Action, **Transfer**) | | [TensorRT-LLM Reasoner](#tensorrt-llm-reasoner) | OpenAI-compatible image/video reasoning server | Reasoner | +| [TensorRT-Edge-LLM](#tensorrt-edge-llm) | On-device TensorRT runtime for Cosmos3-Edge reasoning and DROID policy | Reasoner (Edge), Generator (Action / Policy) | | [Transformers](#transformers) | Hugging Face Transformers inference | Reasoner | | [vLLM](#vllm) | OpenAI-compatible reasoning server (image/video understanding) | Reasoner | | [vLLM-Omni](#vllm-omni) | OpenAI-compatible generation server (image/video/audio/action/transfer) | Generator (Audiovisual, Action, **Transfer**), Simulation ([Transfer-Auto](simulation/auto/auto.md), [Sim-Auto](simulation/auto/run_sim_auto_with_vllm_omni.ipynb), [Sim-Transfer](simulation/transfer/transfer.md#vllm-omni)) | @@ -434,6 +435,54 @@ The server exposes `/health` and the OpenAI-compatible API at [Reasoner TensorRT-LLM notebook](reasoner/run_with_tensorrt_llm.ipynb) for image and video requests. + +## TensorRT-Edge-LLM + +On-device **TensorRT** path for **Cosmos3-Edge** multimodal reasoning and +**Cosmos3-Edge-Policy-DROID** policy generation. This is a different stack from +[TensorRT-LLM Reasoner](#tensorrt-llm-reasoner) / [Generator](#tensorrt-llm-generator). +Reasoning exports ONNX, then builds engines on the target (Jetson Thor/Orin, +DRIVE, DGX Spark, or an x86 developer build). Policy on v0.11.0 builds engines +with the direct builder. + +Use [TensorRT-Edge-LLM](https://github.com/NVIDIA/TensorRT-Edge-LLM) **v0.11.0**. +Complete the upstream +[Installation](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/getting_started/installation.html) +guide first (Python package on the host; C++ runtime on the device). + +Policy examples need `-DBUILD_EXPERIMENTAL_MODELS=ON` so +`cosmos3_policy_inference` exists, and `-DENABLE_CUTE_DSL=ALL` for +[`tensorrt-edgellm-build`](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/getting_started/direct-engine-builder.html). +v0.11.0 `tensorrt-edgellm-export --task policy` writes `attention_pos_id` and +VAE `height` / `width` only. The v0.11.0 runtime binds `attention_position_ids` +and reads `builder_config.content_height` / `content_width`. The direct builder +writes that contract. `cosmos3_policy_inference` in this tag accepts +`--viewPoint concat_view`. + +Set a workspace and the Edge checkpoints. `EDGELLM_ROOT` matches the notebook +default: + +```bash +export EDGELLM_ROOT="${EDGELLM_ROOT:-$HOME/TensorRT-Edge-LLM}" +export POLICY_CHECKPOINT=nvidia/Cosmos3-Edge-Policy-DROID +export REASONING_CHECKPOINT=nvidia/Cosmos3-Edge +export ONNX_DIR="${ONNX_DIR:-$HOME/tensorrt-edgellm-workspace/Cosmos3-Edge/onnx}" +export ENGINE_DIR="${ENGINE_DIR:-$HOME/tensorrt-edgellm-workspace/Cosmos3-Edge/engines}" +``` + +Runnable cookbooks: + +- Reasoner: [`reasoner/run_with_trt_edge_llm.ipynb`](reasoner/run_with_trt_edge_llm.ipynb) + (engines under `$ENGINE_DIR/reasoning`) +- Policy: [`generator/action/run_policy_with_trt_edge_llm.ipynb`](generator/action/run_policy_with_trt_edge_llm.ipynb) + (engines under `$ENGINE_DIR/policy`) + +Upstream references: +[Cosmos3-Edge VLA guide](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/examples/vla/cosmos3.html) +for reasoning, and the +[direct engine builder](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/getting_started/direct-engine-builder.html) +for policy. + ## Transformers Local Python inference for the Cosmos3 Reasoner. This backend uses the diff --git a/.agents/skills/cookbooks/cosmos3/generator/action/README.md b/.agents/skills/cookbooks/cosmos3/generator/action/README.md index 1a4676959..8400c6d08 100644 --- a/.agents/skills/cookbooks/cosmos3/generator/action/README.md +++ b/.agents/skills/cookbooks/cosmos3/generator/action/README.md @@ -24,13 +24,15 @@ The rest of this doc shows how to run these modes on selected embodiments direct - [Run with vLLM-Omni](#run-with-vllm-omni) - [Quickstart](#quickstart-3) - [Notebook walkthrough](#vllm-omni-notebook-walkthrough) +- [Run with TensorRT-Edge-LLM](#run-with-tensorrt-edge-llm) - [Post-Train for Cosmos3-Nano-Policy-DROID](#post-train-for-cosmos3-nano-policy-droid) ## Overview -Examples are shown across native PyTorch (Cosmos Framework), Diffusers, -TensorRT-LLM, and vLLM-Omni. Every backend uses the sample -assets under [`assets/`](./assets) and covers three tasks: +Examples for native PyTorch (Cosmos Framework), Diffusers, TensorRT-LLM, +and vLLM-Omni use the sample assets under [`assets/`](./assets) and cover +the three tasks above. TensorRT-Edge-LLM covers the Edge policy task only, +with the same sample assets. Environment setup for all backends is centralized in the shared [Cosmos3 cookbooks environment setup](../../README.md) guide; each backend below @@ -292,6 +294,30 @@ write outputs under `outputs/cosmos3_action_vllm/`: - [`run_policy_with_vllm_omni.ipynb`](./run_policy_with_vllm_omni.ipynb) — policy inference for DROID through the async video API. +## Run with TensorRT-Edge-LLM + +### Quickstart + +On-device TensorRT path for **Cosmos3-Edge-Policy-DROID**. Set up the shared +[TensorRT-Edge-LLM](../../README.md#tensorrt-edge-llm) environment (upstream +install, C++ runtime with `-DBUILD_EXPERIMENTAL_MODELS=ON` and +`-DENABLE_CUTE_DSL=ALL`). Then follow +[`run_policy_with_trt_edge_llm.ipynb`](./run_policy_with_trt_edge_llm.ipynb): +`tensorrt-edgellm-build` → `cosmos3_policy_inference`. + +This is separate from workstation [vLLM-Omni](#run-with-vllm-omni) policy +serving. Multimodal reasoning on the same Edge-LLM runtime is documented in +[`../../reasoner/run_with_trt_edge_llm.ipynb`](../../reasoner/run_with_trt_edge_llm.ipynb). + +### Notebook walkthrough + +[`run_policy_with_trt_edge_llm.ipynb`](./run_policy_with_trt_edge_llm.ipynb) +builds the same 640×540 DROID `concat_view` as the vLLM-Omni / SGLang policy +notebooks (from `assets/droid_lerobot_example`), runs +`cosmos3_policy_inference` with `--viewPoint concat_view` and an 8-value +`--state`, and writes `action.json` (`[batch, chunk, action_dimension]`). +Optional `--steps` / `--seed` match the upstream VLA guide. + ## Post-Train for Cosmos3-Nano-Policy-DROID To reproduce our post-training recipe for [Cosmos3-Nano-Policy-DROID](https://huggingface.co/nvidia/Cosmos3-Nano-Policy-DROID), use the diff --git a/.agents/skills/cookbooks/cosmos3/generator/action/run_policy_with_trt_edge_llm.ipynb b/.agents/skills/cookbooks/cosmos3/generator/action/run_policy_with_trt_edge_llm.ipynb new file mode 100644 index 000000000..e92caa9b0 --- /dev/null +++ b/.agents/skills/cookbooks/cosmos3/generator/action/run_policy_with_trt_edge_llm.ipynb @@ -0,0 +1,417 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Cosmos3-Edge Policy with TensorRT-Edge-LLM\n", + "\n", + "This notebook runs **Cosmos3-Edge-Policy-DROID** with\n", + "[TensorRT-Edge-LLM](https://github.com/NVIDIA/TensorRT-Edge-LLM) (v0.11.0):\n", + "build the policy engines with the direct builder, then run\n", + "`cosmos3_policy_inference`.\n", + "\n", + "Policy sits under the **Generator / Action** cookbook surface (action chunks),\n", + "not the Reasoner text path. For Edge multimodal reasoning on the same runtime,\n", + "see [`../../reasoner/run_with_trt_edge_llm.ipynb`](../../reasoner/run_with_trt_edge_llm.ipynb).\n", + "\n", + "1. Points at a TensorRT-Edge-LLM v0.11.0 checkout built with\n", + " `-DBUILD_EXPERIMENTAL_MODELS=ON` and `-DENABLE_CUTE_DSL=ALL`, and a Python\n", + " environment with `pip install -e \".[builder]\"`.\n", + "2. Builds `und-prefill`, `gen`, and `vae-encoder` with\n", + " `tensorrt-edgellm-build` into `ENGINE_DIR/policy`. v0.11.0\n", + " `tensorrt-edgellm-export --task policy` writes the UND input\n", + " `attention_pos_id` and VAE `height` / `width` only. The v0.11.0 runtime\n", + " binds `attention_position_ids` and reads `builder_config.content_height` /\n", + " `content_width`. The direct builder writes that contract. DROID content is\n", + " 640×540 inside the 736×544 canvas.\n", + "3. Builds a 640×540 DROID `concat_view` from `assets/droid_lerobot_example`\n", + " (same layout as the vLLM-Omni / SGLang policy notebooks) and writes\n", + " `action.json`.\n", + "\n", + "Complete the upstream\n", + "[Installation](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/getting_started/installation.html)\n", + "and\n", + "[direct engine builder](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/getting_started/direct-engine-builder.html)\n", + "guides first.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Setup paths\n", + "\n", + "`EDGELLM_ROOT` defaults to `~/TensorRT-Edge-LLM`, the same default as the\n", + "shared cookbook README. It must contain\n", + "`build/experimental_models/cosmos3/examples/` after a CMake build with\n", + "`-DBUILD_EXPERIMENTAL_MODELS=ON` and `-DENABLE_CUTE_DSL=ALL`.\n", + "`POLICY_CHECKPOINT` is either a local checkpoint directory or the Hugging Face\n", + "repo id `nvidia/Cosmos3-Edge-Policy-DROID`. Policy engines go under\n", + "`ENGINE_DIR/policy`, separate from the reasoner engines in\n", + "`ENGINE_DIR/reasoning`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "66d95612", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import os\n", + "\n", + "\n", + "def find_repo_root(start: Path) -> Path:\n", + " for candidate in [start, *start.parents]:\n", + " if (candidate / \".agents\" / \"skills\" / \"cookbooks\" / \"cosmos3\").is_dir():\n", + " return candidate\n", + " if (candidate / \"cookbooks\" / \"cosmos3\").is_dir():\n", + " return candidate\n", + " raise FileNotFoundError(\"Could not locate the cosmos-framework repository root\")\n", + "\n", + "\n", + "def cookbook_root(repo: Path) -> Path:\n", + " for relative in (Path(\".agents/skills/cookbooks/cosmos3\"), Path(\"cookbooks/cosmos3\")):\n", + " candidate = repo / relative\n", + " if candidate.is_dir():\n", + " return candidate\n", + " raise FileNotFoundError(repo / \".agents/skills/cookbooks/cosmos3\")\n", + "\n", + "\n", + "COSMOS_ROOT = find_repo_root(Path.cwd().resolve())\n", + "DROID_ASSET_ROOT = (\n", + " cookbook_root(COSMOS_ROOT) / \"generator\" / \"action\" / \"assets\" / \"droid_lerobot_example\"\n", + ").resolve()\n", + "assert DROID_ASSET_ROOT.exists(), DROID_ASSET_ROOT\n", + "\n", + "EDGELLM_ROOT = Path(\n", + " os.environ.get(\"EDGELLM_ROOT\", Path.home() / \"TensorRT-Edge-LLM\")\n", + ").expanduser().resolve()\n", + "ENGINE_DIR = Path(\n", + " os.environ.get(\n", + " \"ENGINE_DIR\",\n", + " Path.home() / \"tensorrt-edgellm-workspace\" / \"Cosmos3-Edge\" / \"engines\",\n", + " )\n", + ").expanduser().resolve()\n", + "POLICY_CHECKPOINT = os.environ.get(\n", + " \"POLICY_CHECKPOINT\", \"nvidia/Cosmos3-Edge-Policy-DROID\"\n", + ")\n", + "POLICY_PROMPT = os.environ.get(\n", + " \"POLICY_PROMPT\",\n", + " \"Pick up the object and place it in the target container.\",\n", + ")\n", + "POLICY_STATE = os.environ.get(\n", + " \"POLICY_STATE\", \"0,0,0,0,0,0,0,0\"\n", + ")\n", + "POLICY_ENGINE_DIR = ENGINE_DIR / \"policy\"\n", + "WORK_DIR = POLICY_ENGINE_DIR / \"cookbook_work\"\n", + "WORK_DIR.mkdir(parents=True, exist_ok=True)\n", + "OUTPUT_JSON = WORK_DIR / \"action.json\"\n", + "\n", + "os.environ[\"COSMOS_ROOT\"] = str(COSMOS_ROOT)\n", + "os.environ[\"EDGELLM_ROOT\"] = str(EDGELLM_ROOT)\n", + "os.environ[\"ENGINE_DIR\"] = str(ENGINE_DIR)\n", + "os.environ[\"POLICY_ENGINE_DIR\"] = str(POLICY_ENGINE_DIR)\n", + "os.environ[\"POLICY_CHECKPOINT\"] = POLICY_CHECKPOINT\n", + "os.environ[\"DROID_ASSET_ROOT\"] = str(DROID_ASSET_ROOT)\n", + "os.environ[\"POLICY_PROMPT\"] = POLICY_PROMPT\n", + "os.environ[\"POLICY_STATE\"] = POLICY_STATE\n", + "os.environ[\"OUTPUT_JSON\"] = str(OUTPUT_JSON)\n", + "\n", + "print(\"cosmos root:\", COSMOS_ROOT)\n", + "print(\"EDGELLM_ROOT:\", EDGELLM_ROOT, \"(exists:\", EDGELLM_ROOT.exists(), \")\")\n", + "print(\"ENGINE_DIR:\", ENGINE_DIR)\n", + "print(\"POLICY_ENGINE_DIR:\", POLICY_ENGINE_DIR)\n", + "print(\"POLICY_CHECKPOINT:\", POLICY_CHECKPOINT)\n", + "print(\"DROID_ASSET_ROOT:\", DROID_ASSET_ROOT)\n", + "print(\"prompt:\", POLICY_PROMPT)\n", + "print(\"state:\", POLICY_STATE)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Build the policy engines\n", + "\n", + "Run this on the target GPU after `pip install -e \".[builder]\"` in the\n", + "TensorRT-Edge-LLM v0.11.0 checkout. The command reads the checkpoint directly\n", + "and writes `und_prefill/`, `gen/`, `vae_encoder/`, `text_tokenizer/`, and\n", + "`embed_tokens.safetensors` under `POLICY_ENGINE_DIR`.\n", + "\n", + "`--components` keeps the build on the policy towers. The checkpoint also\n", + "contains the reasoner configs, and `--components all` would build those too.\n", + "`--max-input-len 512` matches the v0.11.0 policy prompt budget. For DROID, the\n", + "builder sets `content_height=540` and `content_width=640` from\n", + "`checkpoint.json`. A repo id is downloaded with `huggingface-hub` (part of\n", + "`pip install -e \".[export]\"`); a local checkpoint directory skips that download.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + ": \"${EDGELLM_ROOT:?run the setup cell first}\"\n", + ": \"${POLICY_ENGINE_DIR:?run the setup cell first}\"\n", + ": \"${POLICY_CHECKPOINT:?run the setup cell first}\"\n", + "\n", + "if [[ -f \"$POLICY_CHECKPOINT/config.json\" ]]; then\n", + " MODEL_DIR=\"$POLICY_CHECKPOINT\"\n", + "else\n", + " MODEL_DIR=\"${CHECKPOINT_DIR:-$HOME/tensorrt-edgellm-workspace/Cosmos3-Edge/Cosmos3-Edge-Policy-DROID}\"\n", + " mkdir -p \"$MODEL_DIR\"\n", + " POLICY_CHECKPOINT=\"$POLICY_CHECKPOINT\" MODEL_DIR=\"$MODEL_DIR\" python3 - <<'PY'\n", + "import os\n", + "from huggingface_hub import snapshot_download\n", + "\n", + "snapshot_download(\n", + " repo_id=os.environ[\"POLICY_CHECKPOINT\"],\n", + " local_dir=os.environ[\"MODEL_DIR\"],\n", + ")\n", + "PY\n", + "fi\n", + "\n", + "cd \"$EDGELLM_ROOT\"\n", + "mkdir -p \"$POLICY_ENGINE_DIR\"\n", + "tensorrt-edgellm-build \\\n", + " --model-dir \"$MODEL_DIR\" \\\n", + " --engine-dir \"$POLICY_ENGINE_DIR\" \\\n", + " --components und-prefill,gen,vae-encoder \\\n", + " --max-input-len 512 \\\n", + " --max-batch-size 1\n", + "\n", + "python3 - <<'PY'\n", + "import json\n", + "import os\n", + "from pathlib import Path\n", + "\n", + "engine_dir = Path(os.environ[\"POLICY_ENGINE_DIR\"])\n", + "vae = json.loads((engine_dir / \"vae_encoder\" / \"config.json\").read_text())\n", + "und = json.loads((engine_dir / \"und_prefill\" / \"config.json\").read_text())\n", + "builder = vae[\"builder_config\"]\n", + "assert builder[\"content_height\"] == 540, builder\n", + "assert builder[\"content_width\"] == 640, builder\n", + "assert \"attention_position_ids\" in und[\"optimization_profile\"]\n", + "assert (engine_dir / \"embed_tokens.safetensors\").is_file()\n", + "assert (engine_dir / \"text_tokenizer\" / \"tokenizer.json\").is_file()\n", + "print(\n", + " \"policy contract: content\",\n", + " f\"{builder['content_width']}x{builder['content_height']}\",\n", + " \"canvas\",\n", + " f\"{builder['width']}x{builder['height']}\",\n", + ")\n", + "PY\n", + "\n", + "echo \"Built policy engines under $POLICY_ENGINE_DIR\"\n" + ] + }, + { + "cell_type": "markdown", + "id": "3b7a74e3", + "metadata": {}, + "source": [ + "## 3. Prepare a DROID observation\n", + "\n", + "This cell extracts the first frame from each checked-in DROID camera video and\n", + "builds the 640×540 `concat_view` used by the vLLM-Omni / SGLang policy\n", + "notebooks (wrist on top, two exterior views below). Pass that composite to\n", + "`--image`. The v0.11.0 runtime fits this 640×540 concat into the 736×544 canvas using\n", + "`content_width` / `content_height`. Pass the DROID multiview, not a reasoner\n", + "caption asset.\n", + "\n", + "Override `POLICY_PROMPT` in the setup cell if your instruction differs. The\n", + "upstream VLA guide's banana/bowl string is only a CLI placeholder.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e14dba81", + "metadata": {}, + "outputs": [], + "source": [ + "import subprocess\n", + "\n", + "from PIL import Image\n", + "\n", + "try:\n", + " import imageio_ffmpeg\n", + "except ImportError as exc:\n", + " raise RuntimeError(\n", + " \"Install imageio-ffmpeg in this notebook kernel: pip install imageio-ffmpeg\"\n", + " ) from exc\n", + "\n", + "FFMPEG = imageio_ffmpeg.get_ffmpeg_exe()\n", + "CAMERA_VIDEO_PATHS = {\n", + " \"observation/wrist_image_left\": DROID_ASSET_ROOT\n", + " / \"videos\"\n", + " / \"observation.image.wrist_image_left\"\n", + " / \"chunk-000\"\n", + " / \"file-000.mp4\",\n", + " \"observation/exterior_image_1_left\": DROID_ASSET_ROOT\n", + " / \"videos\"\n", + " / \"observation.image.exterior_image_1_left\"\n", + " / \"chunk-000\"\n", + " / \"file-000.mp4\",\n", + " \"observation/exterior_image_2_left\": DROID_ASSET_ROOT\n", + " / \"videos\"\n", + " / \"observation.image.exterior_image_2_left\"\n", + " / \"chunk-000\"\n", + " / \"file-000.mp4\",\n", + "}\n", + "for key, video_path in CAMERA_VIDEO_PATHS.items():\n", + " assert video_path.exists(), f\"missing {key}: {video_path}\"\n", + "\n", + "\n", + "def extract_first_frame(video_path: Path, out_path: Path) -> Path:\n", + " if not out_path.exists() or out_path.stat().st_mtime < video_path.stat().st_mtime:\n", + " subprocess.run(\n", + " [\n", + " FFMPEG,\n", + " \"-y\",\n", + " \"-loglevel\",\n", + " \"error\",\n", + " \"-i\",\n", + " str(video_path),\n", + " \"-frames:v\",\n", + " \"1\",\n", + " str(out_path),\n", + " ],\n", + " check=True,\n", + " )\n", + " return out_path\n", + "\n", + "\n", + "frame_paths = {\n", + " key: extract_first_frame(video_path, WORK_DIR / f\"policy_{key.split('/')[-1]}.png\")\n", + " for key, video_path in CAMERA_VIDEO_PATHS.items()\n", + "}\n", + "frames = {key: Image.open(path).convert(\"RGB\") for key, path in frame_paths.items()}\n", + "\n", + "# Same geometry as run_policy_with_vllm_omni.ipynb / run_policy_with_sglang.ipynb:\n", + "# native 640×360 wrist on top; exteriors resized to 320×180 on the bottom row.\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(\n", + " (half_w, bottom_h), Image.Resampling.BILINEAR\n", + ")\n", + "right = frames[\"observation/exterior_image_2_left\"].resize(\n", + " (half_w, bottom_h), Image.Resampling.BILINEAR\n", + ")\n", + "policy_image = Image.new(\"RGB\", (target_w, target_h))\n", + "policy_image.paste(wrist, (0, 0))\n", + "policy_image.paste(left, (0, wrist.height))\n", + "policy_image.paste(right, (half_w, wrist.height))\n", + "\n", + "OBSERVATION_IMAGE = WORK_DIR / \"droid_policy_first_frame.png\"\n", + "policy_image.save(OBSERVATION_IMAGE)\n", + "os.environ[\"OBSERVATION_IMAGE\"] = str(OBSERVATION_IMAGE)\n", + "\n", + "print(\"observation:\", OBSERVATION_IMAGE, policy_image.size)\n", + "print(\"prompt:\", POLICY_PROMPT)\n", + "try:\n", + " from IPython.display import display\n", + "\n", + " display(policy_image)\n", + "except Exception:\n", + " pass\n" + ] + }, + { + "cell_type": "markdown", + "id": "9ca81f19", + "metadata": {}, + "source": [ + "## 4. Run the policy\n", + "\n", + "`--image` is the 640×540 DROID concat from the previous cell. Pass\n", + "`--viewPoint concat_view` so the policy prompt matches that multiview layout\n", + "(the CLI default is `ego_view`). v0.11.0 `cosmos3_policy_inference` accepts\n", + "this flag. Policy-DROID engines are state-conditioned: `--state` is eight\n", + "comma-separated model-space values. The setup cell's `POLICY_STATE` is a\n", + "zero placeholder; replace it with the robot's current state. For an ordered\n", + "frame list, replace `--image`\n", + "with `--video frame_00.png,frame_01.png,frame_02.png`. The current policy\n", + "conditions on the most recent frame. Optional: `--steps` (denoise steps) and\n", + "`--seed`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e466b7c4", + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + ": \"${EDGELLM_ROOT:?run the setup cell first}\"\n", + ": \"${POLICY_ENGINE_DIR:?run the setup cell first}\"\n", + ": \"${OBSERVATION_IMAGE:?run the DROID observation cell first}\"\n", + ": \"${POLICY_PROMPT:?run the setup cell first}\"\n", + ": \"${POLICY_STATE:?run the setup cell first}\"\n", + ": \"${OUTPUT_JSON:?run the setup cell first}\"\n", + "\n", + "cd \"$EDGELLM_ROOT\"\n", + "./build/experimental_models/cosmos3/examples/cosmos3_policy_inference \\\n", + " --engineDir \"$POLICY_ENGINE_DIR\" \\\n", + " --image \"$OBSERVATION_IMAGE\" \\\n", + " --prompt \"$POLICY_PROMPT\" \\\n", + " --state \"$POLICY_STATE\" \\\n", + " --viewPoint concat_view \\\n", + " --output \"$OUTPUT_JSON\"\n", + "\n", + "echo \"Wrote $OUTPUT_JSON\"\n", + "python3 - <<'PY'\n", + "import json\n", + "import os\n", + "from pathlib import Path\n", + "data = json.loads(Path(os.environ[\"OUTPUT_JSON\"]).read_text())\n", + "print(json.dumps(data, indent=2)[:4000])\n", + "PY\n" + ] + }, + { + "cell_type": "markdown", + "id": "e55845c9", + "metadata": {}, + "source": [ + "## 5. Next steps\n", + "\n", + "- Override `POLICY_PROMPT` in the setup cell, or rebuild the concat with your\n", + " own DROID camera videos.\n", + "- Multimodal reasoning on TensorRT-Edge-LLM:\n", + " [`../../reasoner/run_with_trt_edge_llm.ipynb`](../../reasoner/run_with_trt_edge_llm.ipynb).\n", + "- OpenAI-compatible policy serving on workstation GPUs:\n", + " [`run_policy_with_vllm_omni.ipynb`](./run_policy_with_vllm_omni.ipynb).\n", + "- Cosmos Framework policy server:\n", + " [`run_policy_with_cosmos_framework.md`](./run_policy_with_cosmos_framework.md).\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/.agents/skills/cookbooks/cosmos3/reasoner/README.md b/.agents/skills/cookbooks/cosmos3/reasoner/README.md index e86219b58..d3b70f669 100644 --- a/.agents/skills/cookbooks/cosmos3/reasoner/README.md +++ b/.agents/skills/cookbooks/cosmos3/reasoner/README.md @@ -370,3 +370,24 @@ and walks through image and video examples covering detailed captioning, temporal localization, embodied reasoning, common-sense reasoning, 2D grounding, describe-anything, action CoT trajectories, driving scenes, physical plausibility, and situation understanding. + +## Run with TensorRT-Edge-LLM + +### Quickstart + +Set up the environment with the shared +[TensorRT-Edge-LLM setup](../README.md#tensorrt-edge-llm) (upstream install, +`EDGELLM_ROOT`, ONNX/engine workspace). This path targets **Cosmos3-Edge** on +Jetson Thor/Orin, DRIVE, DGX Spark, or an x86 developer build. It is separate +from the datacenter [TensorRT-LLM](#run-with-tensorrt-llm) server. + +[`run_with_trt_edge_llm.ipynb`](./run_with_trt_edge_llm.ipynb) exports +`nvidia/Cosmos3-Edge`, builds the LLM and visual engines, writes an image +reasoning `input.json` for `assets/robot_153.jpg`, and runs `llm_inference`. + +### Notebook walkthrough + +[`run_with_trt_edge_llm.ipynb`](./run_with_trt_edge_llm.ipynb) is a +command-driven walkthrough (export → engine build → image caption). For +**Cosmos3-Edge-Policy-DROID** on the same runtime, see +[`../generator/action/run_policy_with_trt_edge_llm.ipynb`](../generator/action/run_policy_with_trt_edge_llm.ipynb). diff --git a/.agents/skills/cookbooks/cosmos3/reasoner/run_with_trt_edge_llm.ipynb b/.agents/skills/cookbooks/cosmos3/reasoner/run_with_trt_edge_llm.ipynb new file mode 100644 index 000000000..94966eec4 --- /dev/null +++ b/.agents/skills/cookbooks/cosmos3/reasoner/run_with_trt_edge_llm.ipynb @@ -0,0 +1,295 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "c0762016", + "metadata": {}, + "source": [ + "# Cosmos3-Edge Reasoner with TensorRT-Edge-LLM\n", + "\n", + "This notebook runs **Cosmos3-Edge** multimodal reasoning with\n", + "[TensorRT-Edge-LLM](https://github.com/NVIDIA/TensorRT-Edge-LLM) (v0.11.0):\n", + "export the checkpoint to ONNX, build TensorRT engines, then run the C++\n", + "`llm_inference` binary.\n", + "\n", + "It is **not** the datacenter [TensorRT-LLM](./run_with_tensorrt_llm.ipynb)\n", + "OpenAI-compatible server path.\n", + "\n", + "1. Points at a TensorRT-Edge-LLM checkout and a local workspace for ONNX/engines.\n", + "2. Exports `nvidia/Cosmos3-Edge` with `--task reasoning`.\n", + "3. Builds the LLM and visual engines.\n", + "4. Runs an image caption request against the cookbook asset `robot_153.jpg`.\n", + "\n", + "Complete the upstream\n", + "[Installation](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/getting_started/installation.html)\n", + "first (Python export on an x86 host; C++ runtime on Jetson Thor/Orin, DRIVE,\n", + "DGX Spark, or an x86 developer build). The authoritative command sequence is\n", + "the upstream\n", + "[Cosmos3-Edge VLA guide](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/examples/vla/cosmos3.html).\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Setup paths\n", + "\n", + "Set `EDGELLM_ROOT` to your TensorRT-Edge-LLM checkout (the directory that\n", + "contains `build/` after you compile the runtime). Optional: override\n", + "`ONNX_DIR` / `ENGINE_DIR`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import os\n", + "import json\n", + "\n", + "\n", + "def find_repo_root(start: Path) -> Path:\n", + " for candidate in [start, *start.parents]:\n", + " if (candidate / \".agents\" / \"skills\" / \"cookbooks\" / \"cosmos3\").is_dir():\n", + " return candidate\n", + " if (candidate / \"cookbooks\" / \"cosmos3\").is_dir():\n", + " return candidate\n", + " raise FileNotFoundError(\"Could not locate the cosmos-framework repository root\")\n", + "\n", + "\n", + "def cookbook_root(repo: Path) -> Path:\n", + " for relative in (Path(\".agents/skills/cookbooks/cosmos3\"), Path(\"cookbooks/cosmos3\")):\n", + " candidate = repo / relative\n", + " if candidate.is_dir():\n", + " return candidate\n", + " raise FileNotFoundError(repo / \".agents/skills/cookbooks/cosmos3\")\n", + "\n", + "\n", + "COSMOS_ROOT = find_repo_root(Path.cwd().resolve())\n", + "REASONER_ASSETS = cookbook_root(COSMOS_ROOT) / \"reasoner\" / \"assets\"\n", + "assert REASONER_ASSETS.exists(), REASONER_ASSETS\n", + "\n", + "EDGELLM_ROOT = Path(\n", + " os.environ.get(\"EDGELLM_ROOT\", Path.home() / \"TensorRT-Edge-LLM\")\n", + ").expanduser().resolve()\n", + "ONNX_DIR = Path(\n", + " os.environ.get(\n", + " \"ONNX_DIR\",\n", + " Path.home() / \"tensorrt-edgellm-workspace\" / \"Cosmos3-Edge\" / \"onnx\",\n", + " )\n", + ").expanduser().resolve()\n", + "ENGINE_DIR = Path(\n", + " os.environ.get(\n", + " \"ENGINE_DIR\",\n", + " Path.home() / \"tensorrt-edgellm-workspace\" / \"Cosmos3-Edge\" / \"engines\",\n", + " )\n", + ").expanduser().resolve()\n", + "REASONING_CHECKPOINT = os.environ.get(\"REASONING_CHECKPOINT\", \"nvidia/Cosmos3-Edge\")\n", + "WORK_DIR = ENGINE_DIR / \"reasoning\" / \"cookbook_work\"\n", + "WORK_DIR.mkdir(parents=True, exist_ok=True)\n", + "\n", + "os.environ[\"COSMOS_ROOT\"] = str(COSMOS_ROOT)\n", + "os.environ[\"REASONER_ASSETS\"] = str(REASONER_ASSETS)\n", + "os.environ[\"EDGELLM_ROOT\"] = str(EDGELLM_ROOT)\n", + "os.environ[\"ONNX_DIR\"] = str(ONNX_DIR)\n", + "os.environ[\"ENGINE_DIR\"] = str(ENGINE_DIR)\n", + "os.environ[\"REASONING_CHECKPOINT\"] = REASONING_CHECKPOINT\n", + "os.environ[\"WORK_DIR\"] = str(WORK_DIR)\n", + "\n", + "print(\"cosmos root:\", COSMOS_ROOT)\n", + "print(\"EDGELLM_ROOT:\", EDGELLM_ROOT, \"(exists:\", EDGELLM_ROOT.exists(), \")\")\n", + "print(\"ONNX_DIR:\", ONNX_DIR)\n", + "print(\"ENGINE_DIR:\", ENGINE_DIR)\n", + "print(\"REASONING_CHECKPOINT:\", REASONING_CHECKPOINT)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Export the Reasoner on CPU\n", + "\n", + "Run this on the export host (after `pip install -e .` in the TensorRT-Edge-LLM\n", + "checkout). It downloads `nvidia/Cosmos3-Edge` on first use.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + ": \"${REASONING_CHECKPOINT:?run the setup cell first}\"\n", + ": \"${ONNX_DIR:?run the setup cell first}\"\n", + "\n", + "mkdir -p \"$ONNX_DIR/reasoning\"\n", + "tensorrt-edgellm-export \\\n", + " \"$REASONING_CHECKPOINT\" \\\n", + " \"$ONNX_DIR/reasoning\" \\\n", + " --task reasoning\n", + "\n", + "echo \"Exported reasoning ONNX under $ONNX_DIR/reasoning\"\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Build the LLM and visual engines\n", + "\n", + "Run these from the TensorRT-Edge-LLM tree after the C++ runtime is built\n", + "(`./build/...` binaries exist). Use the platform CMake flags from the upstream\n", + "installation guide; reasoning does **not** require\n", + "`-DBUILD_EXPERIMENTAL_MODELS=ON`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + ": \"${EDGELLM_ROOT:?run the setup cell first}\"\n", + ": \"${ONNX_DIR:?run the setup cell first}\"\n", + ": \"${ENGINE_DIR:?run the setup cell first}\"\n", + "\n", + "cd \"$EDGELLM_ROOT\"\n", + "mkdir -p \"$ENGINE_DIR/reasoning\"\n", + "\n", + "./build/examples/llm/llm_build \\\n", + " --onnxDir \"$ONNX_DIR/reasoning/llm\" \\\n", + " --engineDir \"$ENGINE_DIR/reasoning\" \\\n", + " --maxInputLen 2048 \\\n", + " --maxKVCacheCapacity 4096\n", + "\n", + "./build/examples/multimodal/visual_build \\\n", + " --onnxDir \"$ONNX_DIR/reasoning/visual\" \\\n", + " --engineDir \"$ENGINE_DIR/reasoning\"\n", + "\n", + "echo \"Built reasoning engines under $ENGINE_DIR/reasoning\"\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Write an image reasoning request\n", + "\n", + "The C++ `llm_inference` CLI accepts the standard\n", + "[image message format](https://nvidia.github.io/TensorRT-Edge-LLM/user_guide/format/input-format.html#message-fields).\n", + "This cell writes `input.json` pointing at the cookbook asset.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "image_path = (REASONER_ASSETS / \"robot_153.jpg\").resolve()\n", + "assert image_path.exists(), image_path\n", + "\n", + "input_path = WORK_DIR / \"input.json\"\n", + "output_path = WORK_DIR / \"output.json\"\n", + "\n", + "payload = {\n", + " \"batch_size\": 1,\n", + " \"max_generate_length\": 256,\n", + " \"requests\": [\n", + " {\n", + " \"messages\": [\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"image\", \"image\": str(image_path)},\n", + " {\"type\": \"text\", \"text\": \"Caption the image in detail.\"},\n", + " ],\n", + " }\n", + " ]\n", + " }\n", + " ],\n", + "}\n", + "input_path.write_text(json.dumps(payload, indent=2) + \"\\n\")\n", + "os.environ[\"INPUT_JSON\"] = str(input_path)\n", + "os.environ[\"OUTPUT_JSON\"] = str(output_path)\n", + "\n", + "print(\"wrote\", input_path)\n", + "print(\"image:\", image_path)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Run reasoning\n", + "\n", + "Inference reads the engines under `ENGINE_DIR/reasoning` and writes\n", + "`output.json`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + ": \"${EDGELLM_ROOT:?run the setup cell first}\"\n", + ": \"${ENGINE_DIR:?run the setup cell first}\"\n", + ": \"${INPUT_JSON:?run the input-json cell first}\"\n", + ": \"${OUTPUT_JSON:?run the input-json cell first}\"\n", + "\n", + "cd \"$EDGELLM_ROOT\"\n", + "./build/examples/llm/llm_inference \\\n", + " --engineDir \"$ENGINE_DIR/reasoning\" \\\n", + " --multimodalEngineDir \"$ENGINE_DIR/reasoning\" \\\n", + " --inputFile \"$INPUT_JSON\" \\\n", + " --outputFile \"$OUTPUT_JSON\"\n", + "\n", + "echo \"Wrote $OUTPUT_JSON\"\n", + "python3 - <<'PY'\n", + "import os\n", + "from pathlib import Path\n", + "out = Path(os.environ[\"OUTPUT_JSON\"])\n", + "print(out.read_text()[:4000])\n", + "PY\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Next steps\n", + "\n", + "- Change the prompt or image path in the input-json cell and re-run inference.\n", + "- For **Cosmos3-Edge-Policy-DROID** on TensorRT-Edge-LLM, see\n", + " [`../generator/action/run_policy_with_trt_edge_llm.ipynb`](../generator/action/run_policy_with_trt_edge_llm.ipynb)\n", + " (build the runtime with `-DBUILD_EXPERIMENTAL_MODELS=ON`).\n", + "- Prefer an OpenAI-compatible server on workstation GPUs? Use\n", + " [`run_with_vllm.ipynb`](./run_with_vllm.ipynb) or\n", + " [`run_with_tensorrt_llm.ipynb`](./run_with_tensorrt_llm.ipynb) (Nano/Super).\n", + "- Python-first Edge Reasoner without TensorRT engines:\n", + " [`run_with_transformers.ipynb`](./run_with_transformers.ipynb).\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}