take in iamges now

requirements.txt

@@ -4,5 +4,6 @@ torch
 transformers>=4.43.0
 accelerate
 sentencepiece
+Pillow
 pillow
 torchvision

server.py

@@ -1,7 +1,9 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
-import torch, os
+import torch, os, base64, io
+from typing import Any, Dict, List, Tuple
+from PIL import Image
 
 MODEL_ID = "osunlp/UGround-V1-72B"
 CACHE_DIR = (
@@ -27,21 +29,78 @@ app = FastAPI()
 async def root():
     return {"status": "ok"}
 
-class Message(BaseModel):
-    role: str
-    content: str
-
 class ChatCompletionRequest(BaseModel):
     model: str
-    messages: list[Message]
+    messages: List[Dict[str, Any]]
     max_tokens: int = 128
 
+def _decode_base64_image(data_url: str) -> Image.Image:
+    if data_url.startswith("data:"):
+        header, b64data = data_url.split(",", 1)
+    else:
+        b64data = data_url
+    img_bytes = base64.b64decode(b64data)
+    return Image.open(io.BytesIO(img_bytes)).convert("RGB")
+
+def _to_qwen_messages_and_images(messages: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], List[Any]]:
+    qwen_msgs: List[Dict[str, Any]] = []
+    images: List[Any] = []
+    for msg in messages:
+        role = msg.get("role", "user")
+        content = msg.get("content")
+        q_content: List[Dict[str, Any]] = []
+
+        if isinstance(content, str):
+            q_content.append({"type": "text", "text": content})
+        elif isinstance(content, list):
+            for part in content:
+                ptype = part.get("type")
+                if ptype == "text":
+                    text_val = part.get("text") or part.get("content") or ""
+                    q_content.append({"type": "text", "text": text_val})
+                elif ptype in ("image", "image_url"):
+                    # OpenAI style: {type:"image_url", image_url:{url:"..."}}
+                    url = part.get("image")
+                    if url is None and isinstance(part.get("image_url"), dict):
+                        url = part["image_url"].get("url")
+                    if isinstance(url, str) and url.startswith("data:image"):
+                        img = _decode_base64_image(url)
+                        images.append(img)
+                        q_content.append({"type": "image", "image": img})
+                    else:
+                        # URL or non-base64 string
+                        images.append(url)
+                        q_content.append({"type": "image", "image": url})
+        else:
+            # Unknown content; coerce to text
+            q_content.append({"type": "text", "text": str(content)})
+
+        qwen_msgs.append({"role": role, "content": q_content})
+
+    return qwen_msgs, images
+
 @app.post("/v1/chat/completions")
 async def chat_completions(req: ChatCompletionRequest):
-    prompt = "\n".join([m.content for m in req.messages])
-    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)
+    qwen_messages, image_inputs = _to_qwen_messages_and_images(req.messages)
+    prompt_text = processor.apply_chat_template(
+        qwen_messages, tokenize=False, add_generation_prompt=True
+    )
+    inputs = processor(
+        text=[prompt_text],
+        images=image_inputs if image_inputs else None,
+        padding=True,
+        return_tensors="pt",
+    )
+    inputs = inputs.to(model.device)
+
+    generated_ids = model.generate(**inputs, max_new_tokens=req.max_tokens)
+    trimmed = [
+        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    ]
+    output_texts = processor.batch_decode(
+        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )
+    text = output_texts[0] if output_texts else ""
 
     return {
         "id": "chatcmpl-uground72b",