Add experimental server for self-hosted Hugging Face vision models

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

server/serve/run-docker.sh

@@ -0,0 +1,31 @@
+#!/usr/bin/env bash
+# Run transformers_server.py in Docker on an NVIDIA GPU machine.
+# EXPERIMENTAL: written for google/diffusiongemma-26B-A4B-it; a full agent run with it is not tested yet.
+#
+# Usage: IMAGE=<image> serve/run-docker.sh [start|stop|logs]
+#   IMAGE         a CUDA image with PyTorch and transformers >= 5.8 (required),
+#                 e.g. one you build FROM nvcr.io/nvidia/pytorch with `pip install -U transformers fastapi uvicorn`
+#   MODEL_REPO    Hugging Face repo, already downloaded to ~/.cache/huggingface (default: DiffusionGemma)
+#   MODEL_NAME    model id reported to clients;  PORT (8000);  IMAGE_TOKENS (280);  NAME (container name)
+set -euo pipefail
+cd "$(dirname "$0")"
+NAME=${NAME:-hidagent-model}
+PORT=${PORT:-8000}
+MODEL_REPO=${MODEL_REPO:-google/diffusiongemma-26B-A4B-it}
+HF=$HOME/.cache/huggingface
+
+case "${1:-start}" in
+  start)
+    : "${IMAGE:?set IMAGE to a CUDA image with PyTorch and transformers >= 5.8}"
+    # load from the snapshot directory: works offline, even when the cache has no refs/main
+    SNAPSHOT=$(ls -d "$HF/hub/models--${MODEL_REPO/\//--}"/snapshots/*/ | head -1)
+    docker rm -f "$NAME" >/dev/null 2>&1 || true
+    docker run -d --name "$NAME" --gpus all --ipc host --restart unless-stopped -p "$PORT:$PORT" \
+      -v "$HF:/root/.cache/huggingface:ro" -v "$PWD:/serve:ro" \
+      -e HF_HUB_OFFLINE=1 -e PORT="$PORT" -e IMAGE_TOKENS="${IMAGE_TOKENS:-280}" \
+      -e MODEL_ID="/root/.cache/huggingface${SNAPSHOT#$HF}" -e MODEL_NAME="${MODEL_NAME:-${MODEL_REPO##*/}}" \
+      --entrypoint python3 "$IMAGE" -u /serve/transformers_server.py
+    echo "started $NAME on :$PORT  (logs: $0 logs)" ;;
+  stop) docker rm -f "$NAME" ;;
+  logs) docker logs -f --tail 50 "$NAME" ;;
+esac

server/serve/transformers_server.py

@@ -0,0 +1,80 @@
+"""Minimal OpenAI-compatible server for a Hugging Face vision-language model (e.g. DiffusionGemma).
+
+Only what the agent needs: POST /v1/chat/completions with text and base64 image_url parts, GET /v1/models.
+Env: MODEL_ID (repo id or path), MODEL_NAME (id reported to clients), PORT (8000),
+     IMAGE_TOKENS (visual token budget; DiffusionGemma supports 70/140/280/560/1120).
+"""
+
+import asyncio
+import base64
+import io
+import os
+import re
+import time
+
+import torch
+import transformers
+import uvicorn
+from fastapi import FastAPI
+from PIL import Image
+
+MODEL_ID = os.environ.get("MODEL_ID", "google/diffusiongemma-26B-A4B-it")
+MODEL_NAME = os.environ.get("MODEL_NAME", MODEL_ID.rsplit("/", 1)[-1])
+IMAGE_TOKENS = int(os.environ.get("IMAGE_TOKENS", "280"))
+
+processor = transformers.AutoProcessor.from_pretrained(MODEL_ID)
+if IMAGE_TOKENS and hasattr(processor, "image_processor"):
+    processor.image_processor.max_soft_tokens = processor.image_processor.image_seq_length = IMAGE_TOKENS
+    processor.image_seq_length = IMAGE_TOKENS
+config = transformers.AutoConfig.from_pretrained(MODEL_ID)
+model_class = getattr(transformers, config.architectures[0], transformers.AutoModelForImageTextToText)
+model = model_class.from_pretrained(MODEL_ID, dtype="auto", device_map="auto").eval()
+
+app = FastAPI()
+lock = asyncio.Lock()  # one generation at a time
+
+
+def to_parts(content) -> list[dict]:
+    """OpenAI message content -> transformers chat parts (images first, as the model prefers)."""
+    if isinstance(content, str):
+        return [{"type": "text", "text": content}]
+    parts = []
+    for p in content:
+        if p["type"] == "text":
+            parts.append({"type": "text", "text": p["text"]})
+        elif p["type"] == "image_url":
+            data = base64.b64decode(p["image_url"]["url"].split(",", 1)[1])
+            parts.insert(0, {"type": "image", "image": Image.open(io.BytesIO(data)).convert("RGB")})
+    return parts
+
+
+def generate(messages: list[dict], max_tokens: int) -> str:
+    chat = [{"role": m["role"], "content": to_parts(m["content"])} for m in messages]
+    inputs = processor.apply_chat_template(chat, tokenize=True, add_generation_prompt=True,
+                                           return_dict=True, return_tensors="pt").to(model.device)
+    prompt = inputs["input_ids"][0]
+    with torch.inference_mode():
+        out = model.generate(**inputs, max_new_tokens=max_tokens)[0]
+    out = out.reshape(-1)
+    if len(out) > len(prompt) and torch.equal(out[:len(prompt)], prompt):  # the output echoes the prompt
+        out = out[len(prompt):]
+    text = processor.decode(out).rsplit("<channel|>", 1)[-1]  # drop an (empty) thinking channel
+    return re.sub(r"<\|?[a-z_]+\|?>", "", text).strip()  # and special tokens such as <turn|> <eos> <pad>
+
+
+@app.get("/v1/models")
+def models():
+    return {"object": "list", "data": [{"id": MODEL_NAME, "object": "model", "owned_by": "local"}]}
+
+
+@app.post("/v1/chat/completions")
+async def chat(body: dict):
+    async with lock:
+        text = await asyncio.to_thread(generate, body["messages"], int(body.get("max_tokens") or 512))
+    return {"id": f"chatcmpl-{time.time_ns()}", "object": "chat.completion", "created": int(time.time()),
+            "model": MODEL_NAME, "choices": [{"index": 0, "finish_reason": "stop",
+                                              "message": {"role": "assistant", "content": text}}]}
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "8000")))