Gemma Video Agent

A multi-stage JAX/Keras video tracking pipeline built on TPU, anchored by Gemma 4 multimodal served via vLLM.
Submitted as part of the Google AI GDE TPU Sprint 2026 (vLLM topic #1 — Offline Batch Inference & Online API Server on TPU).

What this does

Given a video clip, the pipeline:

1. Samples frames with OpenCV.
2. Sends frames to Gemma 4 multimodal (served by vLLM on TPU) for scene analysis and per-frame trackable-object proposals.
3. Generates segmentation seeds via a configurable segmenter (Falcon-Perception locally; SAM 3 on KerasHub/TPU is the sprint target).
4. Refines per-frame detections with RF-DETR.
5. Stitches detections into stable tracks with a pure-Python IoU tracker (sample-position-aware, so it survives sparse 1 fps sampling of fast-motion video).
6. Renders an annotated MP4 plus structured analysis.json and tracks.json artifacts.

The same client code runs against a vLLM endpoint regardless of whether vLLM is on TPU, GPU, or CPU — only the --vllm-base-url changes. That's the whole point: the sprint deliverable is the TPU serving path, but local development stays cheap.

Architecture

Quick start — TPU (sprint target)

The full TPU runbook is in notebooks/RUN_ON_GCP_TPU.md. Skeleton:

# 1. Spin up TPU VM (use whatever zone has v5e-8 / v6e-8 capacity)
gcloud compute tpus tpu-vm create gemma-video-agent-tpu \
  --zone=us-central1-a --accelerator-type=v5litepod-8 \
  --version=v2-alpha-tpuv5-lite

# 2. SSH in
gcloud compute tpus tpu-vm ssh gemma-video-agent-tpu --zone=us-central1-a

# 3. Run vLLM serving Gemma 4 (recipe-validated on v5litepod-8 / v6e-8)
sudo docker run -itd --name gemma4 --privileged --network host --shm-size 16G \
  -v /dev/shm:/dev/shm \
  -e VLLM_ARGS="--model=google/gemma-4-E2B-it --tensor-parallel-size=8 \
                --max-model-len=16384 --max-num-batched-tokens=4096 \
                --disable_chunked_mm_input \
                --host=0.0.0.0 --port=8000" \
  vllm/vllm-tpu:nightly

# 4. From your laptop, point the agent at the TPU endpoint
mac-video-perception run \
  --video path/to/clip.mp4 \
  --output-dir runs/tpu \
  --scene-backend gemma-vllm \
  --vllm-base-url http://<tpu-vm-ip>:8000/v1 \
  --vllm-model google/gemma-4-E2B-it

Outputs: analysis.json, tracks.json, annotated.mp4.

Quick start — local CPU dev

For iteration without burning TPU credits. The integration code is identical; only the model & host change.

# Setup
python3 -m venv .venv && source .venv/bin/activate
pip install -e .[gemma,test,web]

# Option A: real Gemma 4 via transformers + Apple MPS (recommended on M-series Mac)
mac-video-perception run \
  --video path/to/clip.mp4 \
  --output-dir runs/local \
  --scene-backend gemma-transformers \
  --gemma-model google/gemma-4-E2B-it

# Option B: vLLM-CPU serving any OpenAI-compatible multimodal model
#   (Gemma 4 itself does NOT run on vLLM-CPU on Apple Silicon — head_dim=512 is
#   not supported by the ARM CPU attention kernel. Use a smaller model for CPU
#   dev or fall back to option A.)
vllm serve HuggingFaceTB/SmolVLM-Instruct \
  --host 127.0.0.1 --port 8000 --max-model-len 4096 --enforce-eager &

mac-video-perception run \
  --video path/to/clip.mp4 \
  --output-dir runs/local \
  --scene-backend gemma-vllm \
  --vllm-model HuggingFaceTB/SmolVLM-Instruct

Repo layout

src/video_perception/        # package source
  config.py                  # Pydantic models for all stage configs
  scene.py                   # Gemma4VLLMBackend, Gemma4TransformersBackend, HeuristicSceneBackend
  segmenters.py              # FalconPerceptionBackend, Sam31Backend (Mac); Sam3KerasBackend stub for TPU
  detectors.py               # RFDetrBackend
  tracking.py                # IoUTracker (sample-position-aware)
  pipeline.py                # VideoPerceptionPipeline orchestrator
  video.py                   # Frame extraction + annotated render
  cli.py                     # mac-video-perception CLI entrypoint
  web_app.py                 # Local FastAPI demo UI

scripts/
  e2e_vllm_local.py          # End-to-end runner against a vLLM HTTP endpoint (CPU/TPU)
  e2e_gemma4_local.py        # End-to-end runner using transformers Gemma 4 (Mac MPS)
  poke_vllm.py               # Single-image multimodal smoke test against vLLM

notebooks/
  gemma_video_agent_tpu.ipynb  # Standalone Kaggle notebook (KerasHub Gemma 4 path)
  RUN_ON_GCP_TPU.md            # Full GCP TPU + vLLM runbook
  kernel-metadata.json         # Kaggle kernel metadata

tests/                       # 22 passing pytest tests (incl. mock vLLM server)

SPRINT_ENTRY.md              # Sprint spreadsheet fields (drop-in)

TPU code in this repo

The "TPU code" the sprint entry refers to lives in:

• src/video_perception/scene.py::Gemma4VLLMBackend — the OpenAI-compatible client that talks to vLLM serving Gemma 4 on TPU. Validated end-to-end against vLLM-CPU + a mock OpenAI server (4 unit tests + 2 pipeline-selection tests, all passing).
• scripts/e2e_vllm_local.py — runs the full pipeline against any vLLM endpoint. Same script drives Mac-CPU vLLM and a GCP TPU VM vLLM endpoint; only --vllm-base-url changes.
• scripts/poke_vllm.py — single-image multimodal smoke test, useful right after vllm serve boots on the TPU VM.
• notebooks/gemma_video_agent_tpu.ipynb — alternate path: Gemma 4 via KerasHub directly on TPU (no vLLM), for the case where you want a single self-contained notebook.
• notebooks/RUN_ON_GCP_TPU.md — step-by-step GCP TPU VM setup + vLLM Docker run + how to drive the pipeline from the VM or from your laptop.

License

Apache 2.0