[video ingestion]Enable cosmos3

video_ingestion_agent/Dockerfile.cosmos

@@ -0,0 +1,72 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Dockerfile for Video Ingestion Agent — Cosmos3 Reasoner variant.
+#
+# Serves only the Cosmos3-Nano Reasoner (VLM) tower via vLLM (`--hf-overrides`,
+# added automatically by scripts/serve.py); the diffusion generator tower is
+# never loaded (~8B-tier footprint).
+#
+# Single venv: vllm-cosmos3 pulls torch 2.11+cu130 and transformers 5.x; the
+# pipeline is installed on top of that (transformers 5.x works for SigLIP via
+# the get_*_features pooler_output handling in the code). vLLM server and the
+# pipeline therefore share one environment.
+#
+# Validated: python 3.13 / torch 2.11.0+cu130 / vllm 0.21.0 / vllm-cosmos3 0.1.0.
+# Weights are NOT bundled (downloaded at runtime via HF_TOKEN).
+#
+# Build: docker build -f Dockerfile.cosmos -t video_ingestion_agent:cosmos3 .
+# For a CUDA 12.8 node (unvalidated): --build-arg CUDA_VERSION=12.8.1
+#   --build-arg TORCH_BACKEND=cu128 --build-arg VLLM_VERSION=0.19.1
+
+ARG CUDA_VERSION=13.0.1
+ARG BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu24.04
+FROM ${BASE_IMAGE}
+
+# Same deps as the default Dockerfile, plus ninja-build: vllm-cosmos3 shells out
+# to `ninja` for kernel JIT at engine init, else EngineCore dies with
+# `FileNotFoundError: 'ninja'` during model load.
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+        python3 python3-venv python3-pip python3-dev python-is-python3 \
+        build-essential ninja-build \
+        ffmpeg \
+        ca-certificates curl git \
+        libgl1 libglib2.0-0 libxcb1 libsm6 libxext6 libxrender1 && \
+    rm -rf /var/lib/apt/lists/*
+
+# Triton needs ptxas to JIT vLLM's CUDA kernels.
+ENV TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
+
+# uv for fast installs.
+RUN pip install --no-cache-dir --break-system-packages uv
+
+WORKDIR /workspace/video_ingestion_agent
+
+ARG VLLM_VERSION=0.21.0
+ARG TORCH_BACKEND=cu130
+ARG VLLM_COSMOS3_GIT=git+https://github.com/NVIDIA/cosmos-framework.git#subdirectory=packages/vllm-cosmos3
+
+# 1. Cosmos vLLM stack FIRST — pins torch 2.11+cu130 + transformers 5.x.
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv venv .venv --python 3.13 --seed --managed-python && \
+    VIRTUAL_ENV=.venv uv pip install --torch-backend=${TORCH_BACKEND} \
+        "vllm==${VLLM_VERSION}" \
+        "vllm-cosmos3 @ ${VLLM_COSMOS3_GIT}" \
+        openai
+
+# 2. Pipeline package ON TOP — its deps (transformers>=4.40, torch>=2.0) are
+#    already satisfied by step 1, so the cu130 torch / transformers 5.x are kept.
+#    No `[server]` extra: vLLM is already provided by vllm-cosmos3.
+COPY pyproject.toml uv.lock README.md ./
+COPY src/ src/
+RUN --mount=type=cache,target=/root/.cache/uv \
+    VIRTUAL_ENV=.venv uv pip install --torch-backend=${TORCH_BACKEND} \
+        -e ".[local,benchmark,webapp]"
+
+ENV VIRTUAL_ENV="/workspace/video_ingestion_agent/.venv"
+ENV PATH="${VIRTUAL_ENV}/bin:${PATH}"
+
+# Remaining project files (scripts, configs, osmo_workflows, ...).
+# .venv/ is in .dockerignore so the COPY won't clobber the venv.
+COPY . /workspace/video_ingestion_agent/

video_ingestion_agent/configs/batch_ingestion.yaml

@@ -9,8 +9,7 @@
 # Model Configuration
 # ============================================================
 models:
-  # Visual language model for segmentation and verification
-  vlm_model: "Qwen/Qwen3-VL-8B-Instruct"
+  vlm_model: "Qwen/Qwen3-VL-8B-Instruct" # or "nvidia/Cosmos3-Nano"
   vlm_backend: "vllm"  # "local", "api", or "vllm"
   vlm_fps: 4             # Frame sampling rate for VLM
 
@@ -30,7 +29,10 @@ models:
   # vLLM backend settings (only used when vlm_backend: "vllm")
   vllm_url: "http://localhost:8000/v1"
   vllm_local_media: true  # Use file:// URLs (fastest, requires same machine)
-  vllm_tp_size: 8         # Tensor parallel across 8 GPUs (OSMO H100 node)
+  # Set to the node's GPU count: one shared vLLM server serves all workers, so
+  # it must span every GPU. (DP replicas would beat TP for an 8B model, but
+  # serve.py only wires --tensor-parallel-size.)
+  vllm_tp_size: 8
 
 # ============================================================
 # Paths

video_ingestion_agent/configs/ingestion.yaml

@@ -8,8 +8,7 @@
 # Model Configuration
 # ============================================================
 models:
-  # Visual language model for segmentation and verification
-  vlm_model: "Qwen/Qwen3-VL-8B-Instruct"
+  vlm_model: "Qwen/Qwen3-VL-8B-Instruct"  # or "nvidia/Cosmos3-Nano"
   vlm_backend: "vllm"  # "local", "api", or "vllm"
   vlm_fps: 4             # Frame sampling rate for VLM
 

video_ingestion_agent/configs/retrieval.yaml

@@ -4,8 +4,7 @@
 # Agent Configuration for Video Retrieval
 
 models:
-  # LLM for agent reasoning
-  llm_model: "Qwen/Qwen3-VL-8B-Instruct"
+  llm_model: "Qwen/Qwen3-VL-8B-Instruct"  # or "nvidia/Cosmos3-Nano"
   llm_backend: "vllm"  # "local" or "api"
 
   # Embedding model for semantic frame search

video_ingestion_agent/docs/pages/model_backends.rst

@@ -238,6 +238,65 @@ reachable, it raises a ``ConnectionError`` with instructions for starting the se
      - File path mode needs shared filesystem
 
 
+Cosmos3-Nano (Reasoner)
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+``nvidia/Cosmos3-Nano`` is an Omni (Mixture-of-Transformers) model. For this
+pipeline we serve **only its Reasoner (VLM) tower** — the diffusion *generator*
+tower is never loaded, so the footprint is ~8B-tier. It runs through the
+``vllm`` backend; switching to it is a one-line config change:
+
+.. code-block:: yaml
+
+   models:
+     vlm_model: "nvidia/Cosmos3-Nano"
+     vlm_backend: "vllm"
+
+``scripts/serve.py`` detects the ``cosmos3`` model name and automatically starts
+the server with ``--hf-overrides '{"architectures":
+["Cosmos3ReasonerForConditionalGeneration"]}'`` (plus ``--mm-encoder-tp-mode
+data --async-scheduling``), so only the Reasoner tower is loaded.
+
+.. important::
+
+   **Do not install Cosmos3 into your pipeline environment.** The server needs
+   the ``vllm-cosmos3`` plugin, which pulls ``vllm`` 0.21.x and a *transformers
+   5.x* fork that is incompatible with the pipeline's pinned ``transformers``
+   4.57. Use the prebuilt Cosmos3 image instead — it packages the vLLM server in
+   an isolated venv internally, so from your side it is still **one image, one
+   command**:
+
+   .. code-block:: bash
+
+      # Build once
+      docker build -f Dockerfile.cosmos -t video_ingestion_agent:cosmos3 .
+
+      # Run: server + pipeline, single container
+      docker run --rm --gpus all --network=host --ipc=host \
+        -v ~/.cache/huggingface:/root/.cache/huggingface \
+        -v /path/to/videos:/path/to/videos:ro \
+        -e HF_HUB_OFFLINE=1 \
+        video_ingestion_agent:cosmos3 \
+        bash -c 'python scripts/serve.py -c configs/ingestion.yaml && \
+                 python scripts/run_ingestion.py /path/to/videos/clip.mp4 \
+                   -c configs/ingestion.yaml -o runs/cosmos'
+
+   The image bridges its two internal venvs over the local HTTP API; OSMO's
+   ``start_vllm.sh`` wires this up automatically via ``VLLM_BIN``. ``--ipc=host``
+   gives vLLM enough shared memory, and mounting the HF cache avoids
+   re-downloading the ~33 GB checkpoint.
+
+**Validated combo:** CUDA 13 / torch 2.11+cu130 / vLLM 0.21.0 / vllm-cosmos3
+0.1.0 / Python 3.13. For a CUDA 12.8 node, rebuild with ``--build-arg
+TORCH_BACKEND=cu128 --build-arg VLLM_VERSION=0.19.1`` (per the model card, not
+validated here).
+
+**When to use:**
+
+- Physical-AI / robotics reasoning where Cosmos3 outperforms a generic VLM
+- A drop-in alternative to Qwen3-VL via the ``vllm`` backend (no code changes)
+
+
 api — NVIDIA Inference API
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

video_ingestion_agent/scripts/serve.py

@@ -98,6 +98,8 @@ def start_server(
     parsed = urlparse(api_url)
     port = parsed.port or 8000
 
+    is_cosmos3 = "cosmos3" in model_name.lower()
+
     # Build command
     cmd = [
         "vllm",
@@ -113,10 +115,29 @@ def start_server(
         str(gpu_memory_utilization),
         "--media-io-kwargs",
         '{"video": {"num_frames": -1}}',
-        "--mm-processor-kwargs",
-        '{"min_pixels": 262144, "max_pixels": 8388608}',
     ]
 
+    if is_cosmos3:
+        # Load only the Reasoner (VLM) tower of the Cosmos3 Omni model
+        # (requires the vllm-cosmos3 plugin).
+        cmd.extend(
+            [
+                "--hf-overrides",
+                '{"architectures": ["Cosmos3ReasonerForConditionalGeneration"]}',
+                "--mm-encoder-tp-mode",
+                "data",
+                "--async-scheduling",
+            ]
+        )
+    else:
+        # Qwen3-VL processor knobs (invalid for Cosmos3).
+        cmd.extend(
+            [
+                "--mm-processor-kwargs",
+                '{"min_pixels": 262144, "max_pixels": 8388608}',
+            ]
+        )
+
     if tensor_parallel_size > 1:
         cmd.extend(["--tensor-parallel-size", str(tensor_parallel_size)])
 

video_ingestion_agent/src/video_ingestion_agent/ingestion/entity_graph/extractors/entity_extractor.py

@@ -127,7 +127,8 @@ def extract_from_caption(
             model = self._get_model()
             response = model.generate_text(
                 conversation=conversation,
-                max_new_tokens=2048,
+                # Headroom for a <think> trace before the JSON (reasoning models).
+                max_new_tokens=4096,
                 temperature=0.0,
             )
 

video_ingestion_agent/src/video_ingestion_agent/ingestion/entity_graph/extractors/visual_extractor.py

@@ -210,7 +210,8 @@ def _generate_caption_for_chunk(self, frames: list[Frame]) -> str:
                 caption = vlm_model._model.generate_from_frames(
                     frames=images,
                     prompt=self.vlm_prompt,
-                    max_new_tokens=512,
+                    # Headroom for a <think> trace before the caption (reasoning models).
+                    max_new_tokens=2048,
                     temperature=0.0,
                 )
             else:
@@ -226,7 +227,8 @@ def _generate_caption_for_chunk(self, frames: list[Frame]) -> str:
                 ]
                 caption = vlm_model.generate_text(
                     conversation=conversation,
-                    max_new_tokens=512,
+                    # Headroom for a <think> trace before the caption (reasoning models).
+                    max_new_tokens=2048,
                     temperature=0.0,
                 )
 
@@ -288,6 +290,9 @@ def extract_embeddings(
                 # Extract embeddings
                 with torch.no_grad():
                     outputs = self.embedding_model.get_image_features(**inputs)
+                    # transformers >=5 returns a ModelOutput, not a tensor
+                    if hasattr(outputs, "pooler_output"):
+                        outputs = outputs.pooler_output
                     # Normalize embeddings
                     batch_embeddings = outputs / outputs.norm(dim=-1, keepdim=True)
                     batch_embeddings = batch_embeddings.cpu().numpy()
@@ -326,6 +331,9 @@ def extract_single_embedding(self, frame: Frame) -> np.ndarray:
 
         with torch.no_grad():
             outputs = self.embedding_model.get_image_features(**inputs)
+            # transformers >=5 returns a ModelOutput, not a tensor
+            if hasattr(outputs, "pooler_output"):
+                outputs = outputs.pooler_output
             embedding = outputs / outputs.norm(dim=-1, keepdim=True)
             embedding = embedding.cpu().numpy()[0]
 

video_ingestion_agent/src/video_ingestion_agent/ingestion/segmentation/dedup.py

@@ -160,7 +160,8 @@ def _ask_llm(
         fallback = self._merge_heuristic(prev, clip)
 
         try:
-            raw = self.model.generate_text(conversation, max_new_tokens=512, temperature=0.0)
+            # Headroom for a <think> trace before the merge JSON (reasoning models).
+            raw = self.model.generate_text(conversation, max_new_tokens=4096, temperature=0.0)
             data = _parse_llm_json(raw)
         except Exception:
             logger.warning(

video_ingestion_agent/src/video_ingestion_agent/retrieval/tools/search_frames.py

@@ -97,6 +97,9 @@ def encode_text(self, query: str) -> np.ndarray:
 
         with torch.no_grad():
             text_features = self._model.get_text_features(**inputs)
+            # transformers >=5 returns a ModelOutput, not a tensor
+            if hasattr(text_features, "pooler_output"):
+                text_features = text_features.pooler_output
             # Normalize
             text_features = text_features / text_features.norm(dim=-1, keepdim=True)
             embedding = text_features.cpu().numpy()[0]