vl

Dockerfile

@@ -6,8 +6,8 @@ RUN apt-get update && apt-get install -y wget git && rm -rf /var/lib/apt/lists/*
 ENV HOME=/app
 # Use a writable directory for HF caches on Spaces
 ENV HF_HOME=/data/huggingface
-ENV TRANSFORMERS_CACHE=/data/huggingface
 ENV HF_HUB_DISABLE_PROGRESS_BARS=1
+ENV HF_HUB_CACHE=/data/huggingface
 
 WORKDIR /app
 

requirements.txt

@@ -1,6 +1,6 @@
 fastapi
 uvicorn[standard]
 torch
-transformers
+transformers>=4.43.0
 accelerate
 sentencepiece

server.py

@@ -1,22 +1,50 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM, AutoProcessor
 import os
 import torch
 
 MODEL_ID = "osunlp/UGround-V1-72B"
-CACHE_DIR = os.environ.get("TRANSFORMERS_CACHE", "/data/huggingface")
-
-# Load tokenizer & model (multi-GPU aware via accelerate)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, cache_dir=CACHE_DIR)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",  # automatically shards across GPUs
-    trust_remote_code=True,
-    cache_dir=CACHE_DIR
+CACHE_DIR = (
+    os.environ.get("HF_HUB_CACHE")
+    or os.environ.get("HF_HOME")
+    or "/data/huggingface"
 )
 
+# Inspect config and load appropriate stack
+config = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True, cache_dir=CACHE_DIR)
+
+is_qwen2_vl = getattr(config, "model_type", None) == "qwen2_vl" or (
+    config.__class__.__name__.lower().startswith("qwen2vl")
+)
+
+if is_qwen2_vl:
+    try:
+        from transformers import Qwen2VLForConditionalGeneration  # type: ignore
+    except Exception as e:
+        raise RuntimeError(
+            "Transformers version does not support Qwen2-VL. Please upgrade transformers to >=4.43."
+        ) from e
+    processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True, cache_dir=CACHE_DIR)
+    model = Qwen2VLForConditionalGeneration.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        trust_remote_code=True,
+        cache_dir=CACHE_DIR,
+    )
+    _use_processor = True
+else:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, cache_dir=CACHE_DIR)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",  # automatically shards across GPUs
+        trust_remote_code=True,
+        cache_dir=CACHE_DIR
+    )
+    _use_processor = False
+
 app = FastAPI()
 
 # OpenAI-style request schema
@@ -33,9 +61,14 @@ class ChatCompletionRequest(BaseModel):
 async def chat_completions(req: ChatCompletionRequest):
     # Concatenate messages into one prompt
     prompt = "\n".join([m.content for m in req.messages])
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=req.max_tokens)
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    if _use_processor:
+        inputs = processor(text=prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(**inputs, max_new_tokens=req.max_tokens)
+        text = processor.tokenizer.decode(outputs[0], skip_special_tokens=True)
+    else:
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(**inputs, max_new_tokens=req.max_tokens)
+        text = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
     return {
         "id": "chatcmpl-uground72b",