diff --git a/cookbooks/cosmos3/reasoner/run_with_sglang.ipynb b/cookbooks/cosmos3/reasoner/run_with_sglang.ipynb new file mode 100644 index 00000000..3b526320 --- /dev/null +++ b/cookbooks/cosmos3/reasoner/run_with_sglang.ipynb @@ -0,0 +1,1374 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "license-header", + "metadata": {}, + "source": [ + "" + ] + }, + { + "cell_type": "markdown", + "id": "b4e72c4a", + "metadata": {}, + "source": [ + "# Cosmos3 Reasoner inference with SGLang\n", + "\n", + "This notebook serves the Cosmos3 Reasoner with an OpenAI-compatible **SGLang** server,\n", + "using the prebuilt `lmsysorg/sglang:dev` Docker image. It:\n", + "\n", + "1. Pulls the SGLang Docker image and installs the client dependencies.\n", + "2. Launches an OpenAI-compatible server for the model tier you choose.\n", + "3. Sends image and video reasoning requests with the `openai` client.\n", + "\n", + "The notebook can serve `nvidia/Cosmos3-Edge` on one GPU or `nvidia/Cosmos3-Super` across\n", + "4 GPUs with tensor parallelism. The query examples in the final section resolve the served\n", + "model dynamically." + ] + }, + { + "cell_type": "markdown", + "id": "3f8b961e", + "metadata": {}, + "source": [ + "## 1. Setup\n", + "\n", + "Serving runs entirely in Docker, so no local Python environment is required for the model.\n", + "You need:\n", + "\n", + "- Docker with the NVIDIA Container Toolkit (`--runtime nvidia` / `--gpus`).\n", + "\n", + "The Cosmos3-Super Reasoner checkpoint is gated on Hugging Face. If you choose Super,\n", + "authenticate with a token that has access after running the setup cells below. Both launch\n", + "cells mount `~/.cache/huggingface` into the container so downloaded weights are reused:\n", + "\n", + "```bash\n", + "hf auth login\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "abc67007", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import os\n", + "import subprocess\n", + "\n", + "\n", + "def find_repo_root() -> Path:\n", + " try:\n", + " return Path(\n", + " subprocess.check_output([\"git\", \"rev-parse\", \"--show-toplevel\"], text=True).strip()\n", + " ).resolve()\n", + " except Exception:\n", + " return Path.cwd().resolve()\n", + "\n", + "\n", + "COSMOS_ROOT = find_repo_root()\n", + "COSMOS_REASONER_ASSETS = COSMOS_ROOT / \"cookbooks\" / \"cosmos3\" / \"reasoner\" / \"assets\"\n", + "COSMOS3_MEDIA_ROOT = COSMOS_ROOT / \"cookbooks\" / \"cosmos3\"\n", + "\n", + "assert COSMOS_REASONER_ASSETS.exists(), COSMOS_REASONER_ASSETS\n", + "\n", + "os.environ[\"COSMOS_ROOT\"] = str(COSMOS_ROOT)\n", + "os.environ[\"COSMOS_REASONER_ASSETS\"] = str(COSMOS_REASONER_ASSETS)\n", + "os.environ[\"COSMOS3_MEDIA_ROOT\"] = str(COSMOS3_MEDIA_ROOT)\n", + "\n", + "\n", + "def asset_path(name: str) -> Path:\n", + " path = COSMOS_REASONER_ASSETS / name\n", + " if not path.exists():\n", + " raise FileNotFoundError(path)\n", + " return path\n", + "\n", + "\n", + "def asset_url(name: str) -> str:\n", + " return asset_path(name).resolve().as_uri()\n", + "\n", + "\n", + "print(\"cosmos root:\", COSMOS_ROOT)\n", + "print(\"Reasoner assets:\", COSMOS_REASONER_ASSETS)\n", + "print(\"Allowed media root:\", COSMOS3_MEDIA_ROOT)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "47bb9a58", + "metadata": {}, + "outputs": [], + "source": [ + "# Client-side dependencies for this notebook (not the model): the OpenAI client to\n", + "# send requests, and the Hugging Face CLI used by the `hf auth login` step above.\n", + "%pip install -q \"openai>=1.0\" pillow \"huggingface_hub[cli]\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "55d4497b", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "# Pull the prebuilt SGLang image with native Cosmos3 Reasoner support (~14 GB; first run only).\n", + "docker pull lmsysorg/sglang:dev" + ] + }, + { + "cell_type": "markdown", + "id": "7fc6b07a", + "metadata": {}, + "source": [ + "## 2. Launch a SGLang server\n", + "\n", + "Choose either the Edge or Super tier below. The commands are mutually exclusive: both use\n", + "the container name `cosmos3-reasoner-sglang` and publish the container's port 8000 to host\n", + "port **8001**, which matches the endpoint the query cells use.\n", + "\n", + "The first start downloads the weights and compiles CUDA graphs and can take several\n", + "minutes; the readiness cell below polls `/health` and waits automatically." + ] + }, + { + "cell_type": "markdown", + "id": "cosmos3-edge-vllm-heading", + "metadata": {}, + "source": [ + "### Cosmos3-Edge inference with SGLang\n", + "\n", + "`Cosmos3-Edge` is the smaller tier and runs on a single GPU. Run this cell to launch the server." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cosmos3-edge-vllm-launch", + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + ": \"${COSMOS_ROOT:?run the setup cell first}\"\n", + ": \"${COSMOS3_MEDIA_ROOT:?run the setup cell first}\"\n", + "CONTAINER=cosmos3-reasoner-sglang\n", + "\n", + "# setup hf cache dir\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", + "\n", + "# Remove any existing server first.\n", + "docker rm -f \"$CONTAINER\" 2>/dev/null || true\n", + "docker run -d --name \"$CONTAINER\" \\\n", + " --runtime nvidia --gpus '\"device=0\"' \\\n", + " -e HF_HOME=/root/.cache/huggingface \\\n", + " -v \"$SGLANG_HF_CACHE:/root/.cache/huggingface\" \\\n", + " -v \"$COSMOS_ROOT:$COSMOS_ROOT\" \\\n", + " -p 8001:8000 --ipc=host \\\n", + " lmsysorg/sglang:dev \\\n", + " sglang serve \\\n", + " --model-type llm \\\n", + " --model-path nvidia/Cosmos3-Edge \\\n", + " --tp-size 1 \\\n", + " --enable-multimodal \\\n", + " --host 0.0.0.0 \\\n", + " --port 8000\n", + "echo \"Cosmos3-Edge starting in container '$CONTAINER' — run the readiness cell next.\"" + ] + }, + { + "cell_type": "markdown", + "id": "cb382fed", + "metadata": {}, + "source": [ + "### Cosmos3-Super inference with SGLang\n", + "\n", + "`Cosmos3-Super` is the larger tier, served across **4 GPUs** with `--tensor-parallel-size 4`.\n", + "Run this cell to launch the server." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d53a0ead", + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + ": \"${COSMOS_ROOT:?run the setup cell first}\"\n", + ": \"${COSMOS3_MEDIA_ROOT:?run the setup cell first}\"\n", + "CONTAINER=cosmos3-reasoner-sglang\n", + "\n", + "# setup hf cache dir\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", + "\n", + "# Remove any existing server first.\n", + "docker rm -f \"$CONTAINER\" 2>/dev/null || true\n", + "docker run -d --name \"$CONTAINER\" \\\n", + " --runtime nvidia --gpus all \\\n", + " -e HF_HOME=/root/.cache/huggingface \\\n", + " -v \"$SGLANG_HF_CACHE:/root/.cache/huggingface\" \\\n", + " -v \"$COSMOS_ROOT:$COSMOS_ROOT\" \\\n", + " -p 8001:8000 --ipc=host \\\n", + " lmsysorg/sglang:dev \\\n", + " sglang serve \\\n", + " --model-type llm \\\n", + " --model-path nvidia/Cosmos3-Super \\\n", + " --tp-size 4 \\\n", + " --enable-multimodal \\\n", + " --host 0.0.0.0 \\\n", + " --port 8000\n", + "echo \"Cosmos3-Super starting in container '$CONTAINER' — run the readiness cell next.\"" + ] + }, + { + "cell_type": "markdown", + "id": "f2681425", + "metadata": {}, + "source": [ + "### Wait for the server to be ready\n", + "\n", + "Run this after launching either tier. It streams the container log and polls `/health`\n", + "until the server is ready, allowing up to 30 minutes on the first run while the weights\n", + "download and CUDA graphs compile." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4a6a4314", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "%%bash\n", + "set -uo pipefail\n", + "PORT=\"${SGLANG_PORT:-8001}\"\n", + "CONTAINER=cosmos3-reasoner-sglang\n", + "\n", + "echo \"Streaming server logs; waiting for http://127.0.0.1:${PORT}/health ...\"\n", + "docker logs -f \"$CONTAINER\" 2>&1 &\n", + "LOGPID=$!\n", + "\n", + "for i in $(seq 1 1800); do\n", + " if curl -fsS \"http://127.0.0.1:${PORT}/health\" >/dev/null 2>&1; then\n", + " echo; echo \"SGLang server is ready.\"\n", + " kill \"$LOGPID\" 2>/dev/null || true\n", + " exit 0\n", + " fi\n", + " if ! docker ps -q -f \"name=^${CONTAINER}$\" | grep -q .; then\n", + " echo; echo \"Container exited early (see logs above).\"\n", + " kill \"$LOGPID\" 2>/dev/null || true\n", + " exit 1\n", + " fi\n", + " sleep 2\n", + "done\n", + "\n", + "kill \"$LOGPID\" 2>/dev/null || true\n", + "echo \"Timed out waiting for SGLang server.\"\n", + "exit 1" + ] + }, + { + "cell_type": "markdown", + "id": "ce7c344c", + "metadata": {}, + "source": [ + "## 3. Query the model" + ] + }, + { + "cell_type": "markdown", + "id": "6b8dae47", + "metadata": {}, + "source": [ + "### Image Caption" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6c9851ca", + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "from IPython.display import Image, display\n", + "\n", + "client = openai.OpenAI(api_key=\"EMPTY\", base_url=\"http://localhost:8001/v1\")\n", + "MODEL = client.models.list().data[0].id\n", + "\n", + "image_path = asset_path(\"robot_153.jpg\")\n", + "image_url = image_path.resolve().as_uri()\n", + "\n", + "response = client.chat.completions.create(\n", + " model=MODEL,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"image_url\", \"image_url\": {\"url\": image_url}},\n", + " {\"type\": \"text\", \"text\": \"Caption the image in detail.\"},\n", + " ],\n", + " }\n", + " ],\n", + " max_tokens=4096,\n", + " seed=0,\n", + ")\n", + "display(Image(filename=str(image_path), width=512))\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "id": "1611a4c5", + "metadata": {}, + "source": [ + "### Video Caption" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ecc3a0d6", + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "from pathlib import Path\n", + "from IPython.display import Video, display\n", + "\n", + "prompt = \"Describe the video in detail.\"\n", + "\n", + "# Plain filesystem path (used for display)\n", + "video_path = str(asset_path(\"video_caption.mp4\"))\n", + "# file:// URL (used for the model request)\n", + "video_url = Path(video_path).resolve().as_uri()\n", + "\n", + "client = openai.OpenAI(api_key=\"EMPTY\", base_url=\"http://localhost:8001/v1\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=client.models.list().data[0].id,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"video_url\", \"video_url\": {\"url\": video_url}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " },\n", + " ],\n", + " max_tokens=4096,\n", + " extra_body={\"media_io_kwargs\": {\"video\": {\"num_frames\": -1, \"fps\": 4}}},\n", + ")\n", + "\n", + "# Display the input video (plain path, NOT the file:// URL)\n", + "display(Video(video_path, embed=True, width=640))\n", + "\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "id": "76977e81-5502-432a-93e7-6c036a8d3ea0", + "metadata": {}, + "source": [ + "### Temporal Localization" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0a23cde6-f552-4354-8509-8a914b0d0382", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "from IPython.display import Video, display\n", + "\n", + "# Plain filesystem path (used for display)\n", + "video_path = str(asset_path(\"temporal_localization_1.mp4\"))\n", + "\n", + "display(Video(video_path, embed=True, width=640))\n", + "\n", + "import openai\n", + "\n", + "prompt = (\n", + " \"\"\"List all action segments in the video. For each detected event, you must determine:\n", + "\n", + "Provide the result in json format with 'seconds' for time depiction for each event. Use keywords 'start', 'end' and 'caption' in the json output. Please list multiple events if applicable.\n", + "\n", + "```json\n", + "[\n", + "{\n", + " \"start\": t_start,\n", + " \"end\": t_end,\n", + " \"caption\": EVENT1\n", + "},\n", + "{\n", + " \"start\": t_start,\n", + " \"end\": t_end,\n", + " \"caption\": EVENT2\n", + "},\n", + "...\n", + "]\n", + "``` \"\"\"\n", + ")\n", + "video_url = asset_url(\"temporal_localization_1.mp4\")\n", + "\n", + "client = openai.OpenAI(\n", + " api_key=\"EMPTY\",\n", + " base_url=\"http://localhost:8001/v1\",\n", + ")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=client.models.list().data[0].id,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"video_url\", \"video_url\": {\"url\": video_url}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " },\n", + " ],\n", + " max_tokens=4096,\n", + " extra_body={\"media_io_kwargs\": {\"video\": {\"num_frames\": -1, \"fps\": 4}}},\n", + ")\n", + "\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6e394917-2872-440d-a514-8933a4704425", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "from IPython.display import Video, display\n", + "\n", + "# Plain filesystem path (used for display)\n", + "video_path = str(asset_path(\"temporal_localization_2.mp4\"))\n", + "\n", + "display(Video(video_path, embed=True, width=640))" + ] + }, + { + "cell_type": "markdown", + "id": "30321842-baa1-4b4e-bf75-19e085ffe927", + "metadata": {}, + "source": [ + "#### Event Timeline" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db0d2a8b-216d-41c7-a121-e4fc35b65bb0", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "prompt = (\n", + " \"Describe the notable events in the provided video. Provide the result in json format with 'mm:ss.ff' format for time depiction for each event.\"\n", + " \"Use keywords 'start', 'end' and 'caption' in the json output.\"\n", + ")\n", + "video_url = asset_url(\"temporal_localization_2.mp4\")\n", + "\n", + "client = openai.OpenAI(\n", + " api_key=\"EMPTY\",\n", + " base_url=\"http://localhost:8001/v1\",\n", + ")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=client.models.list().data[0].id,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"video_url\", \"video_url\": {\"url\": video_url}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " },\n", + " ],\n", + " max_tokens=4096,\n", + " extra_body={\"media_io_kwargs\": {\"video\": {\"num_frames\": -1, \"fps\": 4}}},\n", + ")\n", + "\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "id": "946e0382-b5ff-4cb7-a131-422681d7e63a", + "metadata": {}, + "source": [ + "#### Timestamp Query" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5d6f2813-d243-4071-8d85-bcd26901a9ca", + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "prompt = \"\"\"When is \"A man in a white sweater walks out of a room carrying a box, closes the door behind him, walks on the floor, and turns left at the end near the wall.\" depicted in the video? Please provide the result in json format with 'mm:ss.ff' format for time depiction for the event. Use keywords 'start', 'end' in the json output.\"\"\"\n", + "\n", + "response = client.chat.completions.create(\n", + " model=client.models.list().data[0].id,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"video_url\", \"video_url\": {\"url\": video_url}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " },\n", + " ],\n", + " max_tokens=4096,\n", + " extra_body={\"media_io_kwargs\": {\"video\": {\"num_frames\": -1, \"fps\": 4}}},\n", + ")\n", + "\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "id": "b8c4859e-4a4f-4ab2-846c-f116248295a5", + "metadata": {}, + "source": [ + "#### Interval Question" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6e10b4d2-6b4e-4185-b6d5-06a439f2e8fe", + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "prompt = \"What happened between 00:05.64 and 00:17.49?\"\n", + "response = client.chat.completions.create(\n", + " model=client.models.list().data[0].id,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"video_url\", \"video_url\": {\"url\": video_url}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " },\n", + " ],\n", + " max_tokens=4096,\n", + " extra_body={\"media_io_kwargs\": {\"video\": {\"num_frames\": -1, \"fps\": 4}}},\n", + ")\n", + "\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "id": "6b561503-fe30-4b15-bc59-7878fe1acc32", + "metadata": {}, + "source": [ + "### Embodied Reasoning\n", + "#### Robotics Next Action" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dade5f09-ee93-484f-9fbc-baa1a72aef0b", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "import openai\n", + "from pathlib import Path\n", + "from IPython.display import Video, display\n", + "\n", + "prompt = \"What can be the next immediate action? Answer the question using the following format: Your reasoning. Write your final answer immediately after the tag.\"\n", + "\n", + "# Plain filesystem path (used for display)\n", + "video_path = str(asset_path(\"robotics_next_action.mp4\"))\n", + "# file:// URL (used for the model request)\n", + "video_url = Path(video_path).resolve().as_uri()\n", + "\n", + "client = openai.OpenAI(api_key=\"EMPTY\", base_url=\"http://localhost:8001/v1\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=client.models.list().data[0].id,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"video_url\", \"video_url\": {\"url\": video_url}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " },\n", + " ],\n", + " max_tokens=4096,\n", + " extra_body={\"media_io_kwargs\": {\"video\": {\"num_frames\": -1, \"fps\": 4}}},\n", + ")\n", + "\n", + "# Display the input video (plain path, NOT the file:// URL)\n", + "display(Video(video_path, embed=True, width=640))\n", + "\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "id": "62acdd68-68e1-4dbe-a3d2-f31e0bceb023", + "metadata": {}, + "source": [ + "#### Drive Scene Next Action" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "080c7dfc-b392-49e6-8b94-6ff0868be0d7", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "import openai\n", + "from pathlib import Path\n", + "from IPython.display import Video, display\n", + "\n", + "prompt = \"You are an autonomous vehicle planning system. The video depicts the observation from the vehicle's camera. You need to observe the critical objects in the environment and reason your next action and the driving trajectory ahead.\"\n", + "\n", + "# Plain filesystem path (used for display)\n", + "video_path = str(asset_path(\"drive_scene_next_action.mp4\"))\n", + "# file:// URL (used for the model request)\n", + "video_url = Path(video_path).resolve().as_uri()\n", + "\n", + "client = openai.OpenAI(api_key=\"EMPTY\", base_url=\"http://localhost:8001/v1\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=client.models.list().data[0].id,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"video_url\", \"video_url\": {\"url\": video_url}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " },\n", + " ],\n", + " max_tokens=4096,\n", + " extra_body={\"media_io_kwargs\": {\"video\": {\"num_frames\": -1, \"fps\": 4}}},\n", + ")\n", + "\n", + "# Display the input video (plain path, NOT the file:// URL)\n", + "display(Video(video_path, embed=True, width=640))\n", + "\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "id": "47597159-f8b1-4d4f-a3fe-87074af41e22", + "metadata": {}, + "source": [ + "#### Robot Planning" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "61517853-1cf3-4f77-8dcb-e393f0851bbe", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import re\n", + "import openai\n", + "from pathlib import Path\n", + "from PIL import Image as PILImage, ImageDraw\n", + "from IPython.display import display\n", + "\n", + "client = openai.OpenAI(api_key=\"EMPTY\", base_url=\"http://localhost:8001/v1\")\n", + "MODEL = client.models.list().data[0].id\n", + "\n", + "image_path = str(asset_path(\"robot_planning.png\"))\n", + "image_url = Path(image_path).resolve().as_uri() # file:// URL for the model\n", + "\n", + "# Display the input image (scaled down to fit the cell)\n", + "preview = PILImage.open(image_path).convert(\"RGB\")\n", + "preview.thumbnail((768, 768))\n", + "display(preview)\n", + "\n", + "response = client.chat.completions.create(\n", + " model=MODEL,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"image_url\", \"image_url\": {\"url\": image_url}},\n", + " {\"type\": \"text\", \"text\": 'The task is to put flower into the red bottle. Generate a plan consisting of subtasks for accomplish the task.'},\n", + " ],\n", + " }\n", + " ],\n", + " max_tokens=4096,\n", + " seed=0,\n", + ")\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "id": "08d33a4e-d5fb-4bc1-9fef-56e549a41ff8", + "metadata": {}, + "source": [ + "#### Assisted Task Next Action" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a1ebcb33-0929-4e88-ac3a-85d29fb193a6", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "import openai\n", + "from pathlib import Path\n", + "from IPython.display import Video, display\n", + "\n", + "prompt = \"\"\"This is the overall task that the agent is trying to complete: \"The student exchanges the black ink cartridge of the printer.\"\n", + " In the video, the agent is trying to follow the instruction (a single step out of many to complete the overall task): \"place old ink_cartridge.\"\n", + " What should be the next action of the agent?\n", + " Answer the question using the following format:\n", + " \n", + " Your reasoning.\n", + " \n", + " Write your final answer immediately after the tag.\"\"\"\n", + "\n", + "# Plain filesystem path (used for display)\n", + "video_path = str(asset_path(\"assisted_task_next_action.mp4\"))\n", + "# file:// URL (used for the model request)\n", + "video_url = Path(video_path).resolve().as_uri()\n", + "\n", + "client = openai.OpenAI(api_key=\"EMPTY\", base_url=\"http://localhost:8001/v1\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=client.models.list().data[0].id,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"video_url\", \"video_url\": {\"url\": video_url}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " },\n", + " ],\n", + " max_tokens=4096,\n", + " extra_body={\"media_io_kwargs\": {\"video\": {\"num_frames\": -1, \"fps\": 4}}},\n", + ")\n", + "\n", + "# Display the input video (plain path, NOT the file:// URL)\n", + "display(Video(video_path, embed=True, width=640))\n", + "\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "id": "5cc324bb-28f5-43d3-82d8-7d6fb8223af1", + "metadata": {}, + "source": [ + "### Common Sense Reasoning" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "941b992b-e04b-4cc5-8477-d552ffe77f10", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "import openai\n", + "from pathlib import Path\n", + "from IPython.display import Video, display\n", + "\n", + "prompt = \"\"\"Can the countertop support the weight of the juicers?\n", + " Answer the question using the following format:\n", + "\n", + " \n", + " Your reasoning.\n", + " \n", + "\n", + " Write your final answer immediately after the tag.\"\"\"\n", + "\n", + "# Plain filesystem path (used for display)\n", + "video_path = str(asset_path(\"common_sense_reasoning.mp4\"))\n", + "# file:// URL (used for the model request)\n", + "video_url = Path(video_path).resolve().as_uri()\n", + "\n", + "client = openai.OpenAI(api_key=\"EMPTY\", base_url=\"http://localhost:8001/v1\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=client.models.list().data[0].id,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"video_url\", \"video_url\": {\"url\": video_url}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " },\n", + " ],\n", + " max_tokens=4096,\n", + " extra_body={\"media_io_kwargs\": {\"video\": {\"num_frames\": -1, \"fps\": 4}}},\n", + ")\n", + "\n", + "# Display the input video (plain path, NOT the file:// URL)\n", + "display(Video(video_path, embed=True, width=640))\n", + "\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "id": "d7b596e7-6bce-4ce9-8e59-eff036e12c94", + "metadata": {}, + "source": [ + "### 2D Grounding" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "70441035-22da-4d38-9ff8-5e5e76be32d4", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "import json\n", + "import re\n", + "import openai\n", + "from pathlib import Path\n", + "from PIL import Image as PILImage, ImageDraw\n", + "from IPython.display import display\n", + "\n", + "client = openai.OpenAI(api_key=\"EMPTY\", base_url=\"http://localhost:8001/v1\")\n", + "MODEL = client.models.list().data[0].id\n", + "\n", + "image_path = str(asset_path(\"grounding_2d.png\"))\n", + "image_url = Path(image_path).resolve().as_uri() # file:// URL for the model\n", + "\n", + "response = client.chat.completions.create(\n", + " model=MODEL,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"image_url\", \"image_url\": {\"url\": image_url}},\n", + " {\"type\": \"text\", \"text\": \"Locate the accurate bounding box of the load as a whole. Return a json.\"},\n", + " ],\n", + " }\n", + " ],\n", + " max_tokens=4096,\n", + " seed=0,\n", + ")\n", + "out = response.choices[0].message.content\n", + "print(out)\n", + "\n", + "\n", + "def parse_boxes(text):\n", + " \"\"\"Extract JSON box objects from the model text.\n", + "\n", + " The Reasoner replies in free-form text, so this tolerates a ```` block,\n", + " ``` fences, and any prose after the JSON (which varies between runs).\n", + " \"\"\"\n", + " if \"\" in text:\n", + " text = text.split(\"\")[-1]\n", + " text = re.sub(r\"```(?:json)?|```\", \"\", text).strip()\n", + " starts = [i for i in (text.find(\"[\"), text.find(\"{\")) if i != -1]\n", + " if not starts:\n", + " return []\n", + " data, _ = json.JSONDecoder().raw_decode(text[min(starts):])\n", + " return data if isinstance(data, list) else [data]\n", + "\n", + "\n", + "def box_coords(obj):\n", + " \"\"\"Return ``[x1, y1, x2, y2]`` from the varied shapes the model emits.\"\"\"\n", + " vals = obj if isinstance(obj, (list, tuple)) else (\n", + " obj.get(\"bbox_2d\") or obj.get(\"bbox\") or obj.get(\"box\") or obj.get(\"bounding_box\")\n", + " )\n", + " if vals is None and isinstance(obj, dict) and all(k in obj for k in (\"x1\", \"y1\", \"x2\")):\n", + " vals = [obj[\"x1\"], obj[\"y1\"], obj[\"x2\"], obj.get(\"y2\", obj.get(\"y3\"))]\n", + " if not vals or len(vals) < 4:\n", + " return None\n", + " try:\n", + " return [float(v) for v in vals[:4]]\n", + " except (TypeError, ValueError):\n", + " return None\n", + "\n", + "\n", + "# Draw boxes; coords are normalized to 0-1000\n", + "img = PILImage.open(image_path).convert(\"RGB\")\n", + "W, H = img.size\n", + "draw = ImageDraw.Draw(img)\n", + "\n", + "for obj in parse_boxes(out):\n", + " box = box_coords(obj)\n", + " if not box:\n", + " continue\n", + " x1, y1, x2, y2 = box\n", + " x1, x2 = x1 / 1000 * W, x2 / 1000 * W\n", + " y1, y2 = y1 / 1000 * H, y2 / 1000 * H\n", + " draw.rectangle([x1, y1, x2, y2], outline=\"red\", width=3)\n", + " label = (obj.get(\"label\") or obj.get(\"name\")) if isinstance(obj, dict) else None\n", + " if label:\n", + " draw.text((x1, max(0, y1 - 12)), str(label), fill=\"red\")\n", + "\n", + "# Display scaled down so a large image fits the cell\n", + "preview = img.copy()\n", + "preview.thumbnail((768, 768))\n", + "display(preview)" + ] + }, + { + "cell_type": "markdown", + "id": "02ec583a-b4d5-45fa-90b5-8651ffe1c543", + "metadata": {}, + "source": [ + "### Describe Anything" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1560bbcf-4067-4252-92bf-0b364e76e254", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "import json\n", + "import re\n", + "import openai\n", + "from pathlib import Path\n", + "from PIL import Image as PILImage, ImageDraw\n", + "from IPython.display import display\n", + "\n", + "client = openai.OpenAI(api_key=\"EMPTY\", base_url=\"http://localhost:8001/v1\")\n", + "MODEL = client.models.list().data[0].id\n", + "\n", + "image_path = str(asset_path(\"describe_anything.png\"))\n", + "image_url = Path(image_path).resolve().as_uri() # file:// URL for the model\n", + "\n", + "# Display the input image (scaled down to fit the cell)\n", + "preview = PILImage.open(image_path).convert(\"RGB\")\n", + "preview.thumbnail((768, 768))\n", + "display(preview)\n", + "\n", + "response = client.chat.completions.create(\n", + " model=MODEL,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"image_url\", \"image_url\": {\"url\": image_url}},\n", + " {\"type\": \"text\", \"text\": 'Please caption the notable attributes in the provided image. List and describe all marked subjects in the image with their categories and detailed captions using a json with keyword \"subject_id\", \"category\" and \"caption\".'},\n", + " ],\n", + " }\n", + " ],\n", + " max_tokens=4096,\n", + " seed=0,\n", + ")\n", + "print(response.choices[0].message.content)\n" + ] + }, + { + "cell_type": "markdown", + "id": "ae8c996b-ae6b-4ac9-acc4-5e9c09354c5f", + "metadata": {}, + "source": [ + "### Action CoT" + ] + }, + { + "cell_type": "markdown", + "id": "5e497778-d4f6-4db9-a1f9-ecb6363a7811", + "metadata": {}, + "source": [ + "#### Trajectory Coordinates" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b8ef9a53-8dc2-47d0-a370-1403c51eeae5", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "import json\n", + "import re\n", + "import openai\n", + "from pathlib import Path\n", + "from PIL import Image as PILImage, ImageDraw\n", + "from IPython.display import display\n", + "\n", + "client = openai.OpenAI(api_key=\"EMPTY\", base_url=\"http://localhost:8001/v1\")\n", + "MODEL = client.models.list().data[0].id\n", + "\n", + "image_path = str(asset_path(\"action_cot_trajectory.png\"))\n", + "image_url = Path(image_path).resolve().as_uri()\n", + "\n", + "prompt = \"\"\"You are given the task \"Move the pink bowl to the right\". Specify the 2D trajectory your end effector should follow in pixel space. Return the trajectory coordinates in JSON format like this: {\"point_2d\": [x, y], \"label\": \"gripper trajectory\"}.\n", + "Answer the question using the following format:\n", + "\n", + "\n", + "Your reasoning.\n", + "\n", + "\n", + "Write your final answer immediately after the tag.\n", + "\"\"\"\n", + "\n", + "response = client.chat.completions.create(\n", + " model=MODEL,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"image_url\", \"image_url\": {\"url\": image_url}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " }\n", + " ],\n", + " max_tokens=4096,\n", + " temperature=0.6,\n", + " top_p=0.95,\n", + " presence_penalty=0.0,\n", + " extra_body={\"top_k\": 20, \"repetition_penalty\": 1.0},\n", + ")\n", + "out = response.choices[0].message.content\n", + "print(out)\n", + "\n", + "\n", + "def parse_points(text):\n", + " \"\"\"Grab the JSON list of {point_2d, label} after the tag.\"\"\"\n", + " if \"\" in text:\n", + " text = text.split(\"\")[-1]\n", + " text = re.sub(r\"```(?:json)?\", \"\", text).strip().strip(\"`\").strip()\n", + " m = re.search(r\"\\[.*\\]\", text, re.DOTALL)\n", + " data = json.loads(m.group(0) if m else text)\n", + " return data if isinstance(data, list) else [data]\n", + "\n", + "\n", + "# Visualize the trajectory (points are in pixel space)\n", + "img = PILImage.open(image_path).convert(\"RGB\")\n", + "draw = ImageDraw.Draw(img)\n", + "W, H = img.size\n", + "\n", + "# coords are normalized to 0-1000 (per-axis) -> scale to pixels\n", + "pts = [(o[\"point_2d\"][0] / 1000 * W, o[\"point_2d\"][1] / 1000 * H)\n", + " for o in parse_points(out) if isinstance(o, dict) and \"point_2d\" in o]\n", + "if len(pts) > 1:\n", + " draw.line(pts, fill=\"lime\", width=5)\n", + "for i, (x, y) in enumerate(pts):\n", + " r = 12\n", + " draw.ellipse([x - r, y - r, x + r, y + r], fill=\"red\", outline=\"white\", width=3)\n", + " draw.text((x + 14, y - 14), str(i), fill=\"yellow\")\n", + "preview = img.copy()\n", + "preview.thumbnail((900, 900))\n", + "display(preview)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d3c3e0a0-288b-4080-b49a-b41e8a4ba14c", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "import json\n", + "import re\n", + "import openai\n", + "from pathlib import Path\n", + "from PIL import Image as PILImage, ImageDraw\n", + "from IPython.display import display\n", + "\n", + "client = openai.OpenAI(api_key=\"EMPTY\", base_url=\"http://localhost:8001/v1\")\n", + "MODEL = client.models.list().data[0].id\n", + "\n", + "image_path = str(asset_path(\"robot_planning.png\"))\n", + "image_url = Path(image_path).resolve().as_uri()\n", + "\n", + "prompt = \"\"\"You are given the task \"Put flower into the red bottle\". Specify the 2D trajectory your end effector should follow in pixel space. Return the trajectory coordinates in JSON format like this: {\"point_2d\": [x, y], \"label\": \"gripper trajectory\"}. \n", + "Answer the question using the following format:\n", + "\n", + " Your reasoning. \n", + "Write your final answer immediately after the tag.\n", + "\"\"\"\n", + "\n", + "response = client.chat.completions.create(\n", + " model=MODEL,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"image_url\", \"image_url\": {\"url\": image_url}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " }\n", + " ],\n", + " max_tokens=4096,\n", + " temperature=0.6,\n", + " top_p=0.95,\n", + " presence_penalty=0.0,\n", + " extra_body={\"top_k\": 20, \"repetition_penalty\": 1.0},\n", + ")\n", + "out = response.choices[0].message.content\n", + "print(out)\n", + "\n", + "\n", + "def parse_points(text):\n", + " \"\"\"Grab the JSON list of {point_2d, label} after the tag.\"\"\"\n", + " if \"\" in text:\n", + " text = text.split(\"\")[-1]\n", + " text = re.sub(r\"```(?:json)?\", \"\", text).strip().strip(\"`\").strip()\n", + " m = re.search(r\"\\[.*\\]\", text, re.DOTALL)\n", + " data = json.loads(m.group(0) if m else text)\n", + " return data if isinstance(data, list) else [data]\n", + "\n", + "\n", + "# Visualize the trajectory (points are in pixel space)\n", + "img = PILImage.open(image_path).convert(\"RGB\")\n", + "draw = ImageDraw.Draw(img)\n", + "W, H = img.size\n", + "\n", + "# coords are normalized to 0-1000 (per-axis) -> scale to pixels\n", + "pts = [(o[\"point_2d\"][0] / 1000 * W, o[\"point_2d\"][1] / 1000 * H)\n", + " for o in parse_points(out) if isinstance(o, dict) and \"point_2d\" in o]\n", + "if len(pts) > 1:\n", + " draw.line(pts, fill=\"lime\", width=5)\n", + "for i, (x, y) in enumerate(pts):\n", + " r = 12\n", + " draw.ellipse([x - r, y - r, x + r, y + r], fill=\"red\", outline=\"white\", width=3)\n", + " draw.text((x + 14, y - 14), str(i), fill=\"yellow\")\n", + "preview = img.copy()\n", + "preview.thumbnail((900, 900))\n", + "display(preview)" + ] + }, + { + "cell_type": "markdown", + "id": "3a1ef608-59e9-4e33-9d6b-25e95fb9c3b8", + "metadata": {}, + "source": [ + "#### Driving Scene" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5b5f1bd5-4dd1-443b-89fc-a92d46bdf82a", + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "from pathlib import Path\n", + "from IPython.display import Video, display\n", + "client = openai.OpenAI(api_key=\"EMPTY\", base_url=\"http://localhost:8001/v1\")\n", + "MODEL = client.models.list().data[0].id\n", + "video_path = str(asset_path(\"action_cot_driving_scene.mp4\"))\n", + "video_url = Path(video_path).resolve().as_uri()\n", + "prompt = \"\"\"The video depicts the observation from the vehicle's camera. You need to think step by step and identify the objects in the scene that are critical for safe navigation.\n", + "Answer the question using the following format:\n", + "\n", + "Your reasoning.\n", + "\n", + "Write your final answer immediately after the tag.\"\"\"\n", + "# Show the input video\n", + "display(Video(video_path, embed=True, width=640))\n", + "response = client.chat.completions.create(\n", + " model=MODEL,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"video_url\", \"video_url\": {\"url\": video_url}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " }\n", + " ],\n", + " max_tokens=4096,\n", + " temperature=0.6,\n", + " top_p=0.95,\n", + " presence_penalty=0.0,\n", + " extra_body={\"media_io_kwargs\": {\"video\": {\"num_frames\": -1, \"fps\": 4}}, \"top_k\": 20, \"repetition_penalty\": 1.0},\n", + ")\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "id": "3ba3fc77-3bcd-4509-a796-356ad20136ad", + "metadata": {}, + "source": [ + "### Physical Plausibility Analysis" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "192e522d-81d0-4ff0-a025-aba9636deeac", + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "from pathlib import Path\n", + "from IPython.display import Video, display\n", + "client = openai.OpenAI(api_key=\"EMPTY\", base_url=\"http://localhost:8001/v1\")\n", + "MODEL = client.models.list().data[0].id\n", + "video_path = str(asset_path(\"physical_plausibility.mp4\"))\n", + "video_url = Path(video_path).resolve().as_uri()\n", + "prompt = \"\"\"Is this video physically plausible/possible according to your understanding of e.g. object permanence, shape constancy (objects maintain shape over time), continuous trajectories of objects? Assume it is the normal laws of physics.\n", + "Your answer should be based on the events in the video and ignore the quality of the simulation engine. The rising wall is part of the experiment setup and should not be judged for plausibility.\n", + "(A) Possible\n", + "(B) Impossible\"\"\"\n", + "# Show the input video\n", + "display(Video(video_path, embed=True, width=640))\n", + "response = client.chat.completions.create(\n", + " model=MODEL,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"video_url\", \"video_url\": {\"url\": video_url}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " }\n", + " ],\n", + " max_tokens=4096,\n", + " extra_body={\"media_io_kwargs\": {\"video\": {\"num_frames\": -1, \"fps\": 4}}},\n", + ")\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "id": "80ff6302-8f3d-430a-b514-579aff17eb08", + "metadata": {}, + "source": [ + "### Situation Understanding" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "42d01738-02f5-45f6-aba2-f12eb53f722d", + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "from pathlib import Path\n", + "from IPython.display import Video, display\n", + "client = openai.OpenAI(api_key=\"EMPTY\", base_url=\"http://localhost:8001/v1\")\n", + "MODEL = client.models.list().data[0].id\n", + "video_path = str(asset_path(\"situation_understanding.mp4\"))\n", + "video_url = Path(video_path).resolve().as_uri()\n", + "prompt = \"What is the person doing with the skillet? What will the person likely do next in this situation?\"\n", + "# Show the input video\n", + "display(Video(video_path, embed=True, width=640))\n", + "response = client.chat.completions.create(\n", + " model=MODEL,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"video_url\", \"video_url\": {\"url\": video_url}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " }\n", + " ],\n", + " max_tokens=4096,\n", + " extra_body={\"media_io_kwargs\": {\"video\": {\"num_frames\": -1, \"fps\": 4}}},\n", + ")\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "id": "bbaff2a2", + "metadata": {}, + "source": [ + "## 4. Shut down the server\n", + "\n", + "When you are finished, stop and remove the container to free the GPUs and host port." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "89735fc4-5c42-410c-8697-878120d08f68", + "metadata": {}, + "outputs": [], + "source": [ + "!docker rm -f cosmos3-reasoner-sglang" + ] + } + ], + "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.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}