{ "cells": [ { "cell_type": "markdown", "id": "458e701e", "metadata": {}, "source": [ "" ] }, { "cell_type": "markdown", "id": "28202840", "metadata": {}, "source": [ "# Cosmos3 Reasoner with Transformers\n", "\n", "This notebook runs Cosmos3 Reasoner inference directly with Hugging Face\n", "Transformers — a Python-first path with no server to launch.\n", "\n", "1. Sets up an isolated venv with the Cosmos3 Transformers integration.\n", "2. Registers a Jupyter kernel and loads the Reasoner in process.\n", "3. Runs image and video reasoning requests for **Nano**, **Super**, or **Edge**.\n", "\n", "**Nano / Super** use `Cosmos3OmniForConditionalGeneration` with the unified\n", "`nvidia/Cosmos3-Nano` or `nvidia/Cosmos3-Super` checkpoint.\n", "**Edge** uses `AutoModelForImageTextToText` (`Cosmos3EdgeForConditionalGeneration`)\n", "with `nvidia/Cosmos3-Edge` — do not load Edge with the Omni class.\n", "\n", "The integration loads **only the Reasoner tower** and returns text for text,\n", "image, and video understanding. It does not generate images, video, audio, or\n", "actions — use the Diffusers or vLLM-Omni cookbooks for those.\n", "\n", "This notebook installs Transformers from GitHub `main` so Edge works.\n", "Nano/Super also run on that build (PyPI `transformers>=5.11.0` is enough if you\n", "only need Nano/Super).\n", "\n", "Note: if you have already completed steps 1-4 and installed the\n", "`Cosmos3 Transformers (Python 3.13)` kernel, switch to that kernel, run the\n", "Restore Environment cell in step 4, then continue from step 5.\n" ] }, { "cell_type": "markdown", "id": "98a8252e", "metadata": {}, "source": [ "## 1. Prerequisites\n", "\n", "Use a Linux machine with NVIDIA GPU access, model access on Hugging Face, and\n", "either `uvx hf@latest auth login` or `HF_TOKEN` set.\n", "\n", "> **Headless servers:** if you see an error like `libxcb.so.1: cannot open shared\n", "> object file` when importing, install the required system libraries:\n", ">\n", "> ```bash\n", "> apt-get install -y libxcb1 libgl1 libglib2.0-0\n", "> ```\n", "\n", "> **uv version:** these notebooks need `uv >= 0.11.3`. Older versions do not\n", "> recognize newer `--torch-backend` values such as `cu130`. Upgrade with\n", "> `uv self update` if you hit version-related errors." ] }, { "cell_type": "markdown", "id": "f9f3020e", "metadata": {}, "source": [ "## 2. Configure Paths and Environment\n", "\n", "The defaults are relative to this `cosmos` checkout. Override any of these before\n", "running the next cell if needed:\n", "\n", "```bash\n", "export COSMOS3_TRANSFORMERS_VENV=/path/to/.venv-cosmos3-transformers\n", "export COSMOS3_TORCH_BACKEND=auto # or cu130 / cu128 to pin an explicit CUDA wheel\n", "export HF_HOME=/path/to/large/huggingface/cache\n", "export UV_LINK_MODE=copy\n", "export CUDA_VISIBLE_DEVICES=0\n", "```" ] }, { "cell_type": "code", "execution_count": null, "id": "38384104", "metadata": {}, "outputs": [], "source": [ "from pathlib import Path\n", "import os\n", "\n", "\n", "def find_repo_root(start: Path) -> Path:\n", " for path in [start, *start.parents]:\n", " if (path / \"README.md\").exists() and (path / \"cookbooks\").exists():\n", " return path\n", " return start\n", "\n", "\n", "def configure_transformers_environment() -> None:\n", " global COSMOS_ROOT\n", " global COSMOS_REASONER_ASSETS\n", " global COSMOS3_TRANSFORMERS_VENV\n", " global COSMOS3_TORCH_BACKEND\n", "\n", " COSMOS_ROOT = find_repo_root(Path.cwd().resolve())\n", " COSMOS_REASONER_ASSETS = COSMOS_ROOT / \"cookbooks\" / \"cosmos3\" / \"reasoner\" / \"assets\"\n", " COSMOS3_TRANSFORMERS_VENV = Path(\n", " os.environ.get(\"COSMOS3_TRANSFORMERS_VENV\", COSMOS_ROOT / \".venv-cosmos3-transformers\")\n", " ).resolve()\n", " COSMOS3_TORCH_BACKEND = os.environ.get(\"COSMOS3_TORCH_BACKEND\", \"auto\")\n", "\n", " os.environ[\"COSMOS3_TRANSFORMERS_VENV\"] = str(COSMOS3_TRANSFORMERS_VENV)\n", " os.environ[\"COSMOS3_TORCH_BACKEND\"] = COSMOS3_TORCH_BACKEND\n", " os.environ.setdefault(\"UV_LINK_MODE\", \"copy\")\n", "\n", " assert COSMOS_REASONER_ASSETS.exists(), COSMOS_REASONER_ASSETS\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", "configure_transformers_environment()\n", "print(\"cosmos root:\", COSMOS_ROOT)\n", "print(\"Reasoner assets:\", COSMOS_REASONER_ASSETS)\n", "print(\"Transformers venv:\", COSMOS3_TRANSFORMERS_VENV)\n", "print(\"Torch backend:\", COSMOS3_TORCH_BACKEND)" ] }, { "cell_type": "markdown", "id": "0ea0d464", "metadata": {}, "source": [ "## 3. Install Transformers Dependencies\n", "\n", "This cell creates the venv, installs dependencies (including Transformers from\n", "`main`, required for Cosmos3-Edge), and registers a Jupyter kernel so the model\n", "can run in process.\n", "\n", "Nano/Super Reasoner support first appears in Transformers `v5.11.0`. Edge support\n", "is on Transformers `main` until a PyPI release includes it.\n", "\n", "`--torch-backend` defaults to `auto`, which lets uv pick a CUDA build of\n", "`torch`/`torchvision` that matches your driver. Set `COSMOS3_TORCH_BACKEND=cu130`\n", "(or `cu128`) above to pin an explicit wheel.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "34acbe10", "metadata": {}, "outputs": [], "source": [ "%%bash\n", "set -euo pipefail\n", "\n", "if ! command -v uv >/dev/null 2>&1; then\n", " echo \"uv is not installed. Install it first: https://docs.astral.sh/uv/getting-started/installation/\"\n", " exit 1\n", "fi\n", "\n", "export UV_LINK_MODE=\"${UV_LINK_MODE:-copy}\"\n", "uv venv \"$COSMOS3_TRANSFORMERS_VENV\" --python 3.13 --seed --managed-python --allow-existing\n", "source \"$COSMOS3_TRANSFORMERS_VENV/bin/activate\"\n", "\n", "uv pip install --torch-backend=\"$COSMOS3_TORCH_BACKEND\" \\\n", " accelerate \\\n", " av \\\n", " ipykernel \\\n", " pillow \\\n", " \"safetensors>=0.8.0\" \\\n", " torch \\\n", " \"torchvision==0.25.0\" \\\n", " \"transformers @ git+https://github.com/huggingface/transformers.git\"\n", "\n", "\"$COSMOS3_TRANSFORMERS_VENV/bin/python\" -m ipykernel install --user \\\n", " --name cosmos3-transformers \\\n", " --display-name \"Cosmos3 Transformers (Python 3.13)\"\n", "\n", "echo\n", "echo \"Installed dependencies into: $COSMOS3_TRANSFORMERS_VENV\"\n", "echo \"Next: switch this notebook kernel to: Cosmos3 Transformers (Python 3.13)\"\n" ] }, { "cell_type": "markdown", "id": "72c1bd2d", "metadata": {}, "source": [ "### 3a. Optional: Install ModelOpt for FP8\n", "\n", "Skip this subsection for standard BF16 inference. To load the public Nano/Super `fp8`\n", "revision, run the cell below before switching kernels. `requests` is installed\n", "explicitly because ModelOpt `0.44.0` imports it but does not declare it as a dependency.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "96eb0e73", "metadata": {}, "outputs": [], "source": [ "%%bash\n", "set -euo pipefail\n", "\n", "source \"$COSMOS3_TRANSFORMERS_VENV/bin/activate\"\n", "uv pip install \\\n", " \"nvidia-modelopt==0.44.0\" \\\n", " requests\n" ] }, { "cell_type": "markdown", "id": "1149a05a", "metadata": {}, "source": [ "## 4. Select the Transformers Kernel\n", "\n", "The install cell registers the `Cosmos3 Transformers (Python 3.13)` Jupyter kernel.\n", "\n", "**Switch this notebook to that kernel before running the remaining Python cells**,\n", "then run the Restore Environment cell immediately below. It can take a moment for\n", "the new kernel to appear in the notebook interface." ] }, { "cell_type": "code", "execution_count": null, "id": "3d209e12", "metadata": {}, "outputs": [], "source": [ "# Run this cell immediately after switching to the Cosmos3 Transformers kernel.\n", "# It restores the same paths as the configure cell in step 2.\n", "from pathlib import Path\n", "import os\n", "\n", "\n", "def find_repo_root(start: Path) -> Path:\n", " for path in [start, *start.parents]:\n", " if (path / \"README.md\").exists() and (path / \"cookbooks\").exists():\n", " return path\n", " return start\n", "\n", "\n", "COSMOS_ROOT = find_repo_root(Path.cwd().resolve())\n", "COSMOS_REASONER_ASSETS = COSMOS_ROOT / \"cookbooks\" / \"cosmos3\" / \"reasoner\" / \"assets\"\n", "COSMOS3_TRANSFORMERS_VENV = Path(\n", " os.environ.get(\"COSMOS3_TRANSFORMERS_VENV\", COSMOS_ROOT / \".venv-cosmos3-transformers\")\n", ").resolve()\n", "os.environ[\"COSMOS3_TRANSFORMERS_VENV\"] = str(COSMOS3_TRANSFORMERS_VENV)\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", "print(\"cosmos root:\", COSMOS_ROOT)\n", "print(\"Reasoner assets:\", COSMOS_REASONER_ASSETS)" ] }, { "cell_type": "markdown", "id": "ab6a9c99", "metadata": {}, "source": [ "## 5. Verify GPU and Python Environment" ] }, { "cell_type": "code", "execution_count": null, "id": "948b294f", "metadata": {}, "outputs": [], "source": [ "import os\n", "import sys\n", "from pathlib import Path\n", "\n", "if \"COSMOS3_TRANSFORMERS_VENV\" not in os.environ:\n", " raise RuntimeError(\"Run the Restore Environment cell after switching to the Transformers kernel.\")\n", "\n", "expected_python = (Path(os.environ[\"COSMOS3_TRANSFORMERS_VENV\"]) / \"bin\" / \"python\").resolve()\n", "current_python = Path(sys.executable).resolve()\n", "print(\"kernel python:\", current_python)\n", "print(\"expected python:\", expected_python)\n", "if current_python != expected_python:\n", " raise RuntimeError(\n", " \"This notebook is not running inside the Transformers venv. \"\n", " \"Switch the kernel to 'Cosmos3 Transformers (Python 3.13)', then run the Restore Environment cell above.\"\n", " )\n", "\n", "import torch\n", "import transformers\n", "\n", "print(\"transformers:\", transformers.__version__)\n", "print(\"torch:\", torch.__version__)\n", "print(\"torch cuda:\", torch.version.cuda)\n", "print(\"cuda available:\", torch.cuda.is_available())\n", "print(\"device count:\", torch.cuda.device_count())\n", "if torch.cuda.is_available():\n", " print(\"device 0:\", torch.cuda.get_device_name(0))" ] }, { "cell_type": "markdown", "id": "0e00b6d3", "metadata": {}, "source": [ "## 6. Load the Reasoner\n", "\n", "Load the processor and model once, then reuse them for every request below.\n", "\n", "The default is the public BF16 checkpoint. To use the public Nano/Super FP8 weights,\n", "set `USE_MODELOPT_FP8 = True`. The model ID stays the same and Transformers\n", "downloads the `fp8` revision to the Hugging Face cache and loads it from the\n", "resolved snapshot path. ModelOpt restores that snapshot's calibrated E4M3 weights\n", "and static scales; this is not runtime quantization.\n", "\n", "Set `model_id` to `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super`, or\n", "`nvidia/Cosmos3-Edge`. `device_map=\"auto\"` places the model on the available\n", "GPU(s) and can shard Super across multiple GPUs when Accelerate is installed.\n", "\n", "> **First run downloads the checkpoint** (tens of GiB for Nano/Super; Edge is\n", "> smaller). Nano/Super load only the Reasoner tower from the unified checkpoint\n", "> into memory. Subsequent runs reuse the Hugging Face cache.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "937f9ede", "metadata": {}, "outputs": [], "source": [ "import torch\n", "\n", "USE_MODELOPT_FP8 = False\n", "model_id = \"nvidia/Cosmos3-Nano\" # or \"nvidia/Cosmos3-Super\" or \"nvidia/Cosmos3-Edge\"\n", "FP8_REVISION = \"fp8\"\n", "\n", "if USE_MODELOPT_FP8:\n", " # Enable ModelOpt restoration before importing the Transformers model class.\n", " # Importing this backend registers the real per-tensor FP8 GEMM implementation.\n", " from huggingface_hub import snapshot_download\n", " from modelopt.torch import opt as modelopt_opt\n", " from modelopt.torch.quantization.backends import fp8_per_tensor_gemm # noqa: F401\n", "\n", " modelopt_opt.enable_huggingface_checkpointing()\n", " checkpoint_path = snapshot_download(model_id, revision=FP8_REVISION)\n", "else:\n", " checkpoint_path = model_id\n", "\n", "from transformers import (\n", " AutoModelForImageTextToText,\n", " AutoProcessor,\n", " Cosmos3OmniForConditionalGeneration,\n", ")\n", "\n", "# Nano / Super: Omni class. Edge: Auto image-text API (Cosmos3EdgeForConditionalGeneration).\n", "processor = AutoProcessor.from_pretrained(checkpoint_path)\n", "if model_id == \"nvidia/Cosmos3-Edge\":\n", " model = AutoModelForImageTextToText.from_pretrained(\n", " checkpoint_path,\n", " dtype=torch.bfloat16,\n", " device_map=\"auto\",\n", " )\n", "else:\n", " model = Cosmos3OmniForConditionalGeneration.from_pretrained(\n", " checkpoint_path,\n", " dtype=torch.bfloat16,\n", " device_map=\"auto\",\n", " )\n", "\n", "\n", "def run_reasoner(content, fps=None, max_new_tokens=512):\n", " \"\"\"Run one Reasoner request. `content` is a chat content list (image/video/text blocks).\"\"\"\n", " messages = [{\"role\": \"user\", \"content\": content}]\n", " template_kwargs = dict(\n", " tokenize=True,\n", " add_generation_prompt=True,\n", " return_dict=True,\n", " return_tensors=\"pt\",\n", " )\n", " if fps is not None:\n", " template_kwargs[\"fps\"] = fps\n", "\n", " inputs = processor.apply_chat_template(messages, **template_kwargs).to(model.device, torch.bfloat16)\n", " generated_ids = model.generate(**inputs, do_sample=False, max_new_tokens=max_new_tokens)\n", " generated_ids_trimmed = [\n", " out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)\n", " ]\n", " return processor.batch_decode(\n", " generated_ids_trimmed,\n", " skip_special_tokens=True,\n", " clean_up_tokenization_spaces=False,\n", " )[0]\n", "\n", "\n", "print(\"Loaded\", model_id, \"as\", type(model).__name__)\n" ] }, { "cell_type": "markdown", "id": "0f80ec19", "metadata": {}, "source": [ "## 7. Image Reasoning" ] }, { "cell_type": "code", "execution_count": null, "id": "3cdd010c", "metadata": {}, "outputs": [], "source": [ "from IPython.display import Image, display\n", "\n", "image_path = asset_path(\"robot_153.jpg\")\n", "display(Image(filename=str(image_path), width=512))\n", "\n", "output = run_reasoner(\n", " [\n", " {\"type\": \"image\", \"path\": str(image_path.resolve())},\n", " {\"type\": \"text\", \"text\": \"Caption the image in detail.\"},\n", " ]\n", ")\n", "print(output)" ] }, { "cell_type": "markdown", "id": "88cadada", "metadata": {}, "source": [ "## 8. Video Reasoning\n", "\n", "Use a `video` content block and pass a frame sampling rate (`fps`) to the helper.\n", "\n", "> Video decoding uses the packages installed above. Transformers prints a\n", "> deprecation warning that it fell back to the `torchvision` decoder — this is\n", "> expected and harmless. To switch to the modern `torchcodec` decoder, install it\n", "> along with system FFmpeg libraries (`libavutil`/`libavcodec`)." ] }, { "cell_type": "code", "execution_count": null, "id": "e7cab369", "metadata": {}, "outputs": [], "source": [ "from IPython.display import Video, display\n", "\n", "video_path = asset_path(\"video_caption.mp4\")\n", "display(Video(str(video_path), embed=True, width=640))\n", "\n", "output = run_reasoner(\n", " [\n", " {\"type\": \"video\", \"path\": str(video_path.resolve())},\n", " {\"type\": \"text\", \"text\": \"Describe the notable events in this video.\"},\n", " ],\n", " fps=2,\n", ")\n", "print(output)" ] }, { "cell_type": "markdown", "id": "5d2a6b57", "metadata": {}, "source": [ "## 9. Next Steps\n", "\n", "- Run **Cosmos3-Super**: set `model_id = \"nvidia/Cosmos3-Super\"` in step 6 and\n", " re-run from there. `device_map=\"auto\"` shards it across multiple GPUs.\n", "- Run **Cosmos3-Edge**: set `model_id = \"nvidia/Cosmos3-Edge\"` in step 6 and\n", " re-run from there (uses `AutoModelForImageTextToText`, not the Omni class).\n", "- Run **ModelOpt FP8 Nano/Super**: set `USE_MODELOPT_FP8 = True` and re-run from step 6.\n", " The same model ID is loaded from its public `fp8` revision (Edge has none).\n", "- Try other Reasoner tasks (temporal localization, grounding, embodied reasoning)\n", " by changing the prompt and asset — see the\n", " [Reasoner Prompt Guide](./reasoner_prompt_guide.md).\n", "- Need an OpenAI-compatible server instead of in-process Python? See\n", " [`run_with_vllm.ipynb`](./run_with_vllm.ipynb) or [`run_with_nim.ipynb`](./run_with_nim.ipynb).\n" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "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.10.20" } }, "nbformat": 4, "nbformat_minor": 5 }