vision: images ride the DSpark speculative loop

recipe/server.py

@@ -17,6 +17,9 @@ import time
 import threading
 import queue
 import uuid
+import io
+import base64
+import urllib.request
 
 os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
 os.environ.setdefault("MASTER_PORT", "29500")
@@ -38,12 +41,13 @@ from torch.nn.attention import SDPBackend, sdpa_kernel
 # implement (NotImplementedError). It crashes the cache save. So we can't cache the dynamo frontend;
 # the ~200s warmup tracing stays. (AOTInductor would avoid dynamo entirely — separate effort.)
 from types import SimpleNamespace
-from typing import List, Optional
+from typing import Any, List, Optional
 import json as _json
+from PIL import Image
 from fastapi import FastAPI
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
+from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer, TorchAoConfig
 from torchao.quantization import (
     Float8DynamicActivationFloat8WeightConfig, Float8WeightOnlyConfig, PerRow)
 from torchao.quantization.quantize_.common.kernel_preference import KernelPreference
@@ -112,6 +116,10 @@ EV = FP8Gemma4DSparkEvaluator(0, _args)
 print(f"[dspark] TIMING model+quant load: {time.time()-_T0:.1f}s", flush=True)
 EV.confidence_head_recorder = None  # metrics recorder is only started inside evaluate(); we bypass it
 TOK = EV.tokenizer
+# The multimodal processor (image + text). gemma-4-12B is a VLM; DSpark's target processes the image in
+# the prefill (pixel_values), the KV cache carries it, and the draft speculates over image-aware hidden
+# states — so vision rides the same speculative loop as text. Only loaded once here.
+PROC = AutoProcessor.from_pretrained(TARGET)
 STOP = resolve_stop_token_ids(EV.target_model, TOK)  # gemma real eos set (e.g. [1,106,50]), not a guess
 _LOCK = threading.Lock()
 
@@ -156,7 +164,38 @@ app = FastAPI()
 
 class Msg(BaseModel):
     role: str
-    content: str
+    content: Any  # str, or OpenAI multimodal list: [{"type":"text",...}, {"type":"image_url",...}]
+
+
+def _load_image(url: str) -> Image.Image:
+    """Load a PIL image from a data: URI or an http(s) URL."""
+    if url.startswith("data:"):
+        return Image.open(io.BytesIO(base64.b64decode(url.split(",", 1)[1]))).convert("RGB")
+    return Image.open(io.BytesIO(urllib.request.urlopen(url, timeout=20).read())).convert("RGB")
+
+
+def _build_chat(messages):
+    """Turn request messages into (chat-for-template, has_images). Accepts plain string content or an
+    OpenAI multimodal content list (text parts + image_url/image parts), which become gemma image parts."""
+    chat, has_images = [], False
+    for m in messages:
+        c = m.content
+        if isinstance(c, str):
+            chat.append({"role": m.role, "content": c})
+            continue
+        parts = []
+        for part in (c or []):
+            t = part.get("type") if isinstance(part, dict) else None
+            if t == "text":
+                parts.append({"type": "text", "text": part.get("text", "")})
+            elif t in ("image_url", "image"):
+                url = part.get("image_url", {}).get("url") if t == "image_url" else (part.get("image") or part.get("url"))
+                parts.append({"type": "image", "image": _load_image(url)})
+                has_images = True
+            elif isinstance(part, str):
+                parts.append({"type": "text", "text": part})
+        chat.append({"role": m.role, "content": parts})
+    return chat, has_images
 
 
 class ChatReq(BaseModel):
@@ -196,13 +235,25 @@ def _split_think(ids, thinking):
 
 @app.post("/v1/chat/completions")
 def chat(req: ChatReq):
-    msgs = [{"role": m.role, "content": m.content} for m in req.messages]
     # thinking on by default; a `*-nothink` model alias (routed to this same backend) turns it off
     thinking = THINKING and not (req.model or "").endswith("-nothink")
-    ids = TOK.apply_chat_template(
-        msgs, add_generation_prompt=True, enable_thinking=thinking,
-        return_tensors="pt", return_dict=True,
-    )["input_ids"].to(EV.device)
+    chat_msgs, has_images = _build_chat(req.messages)
+    if has_images:
+        # multimodal: the processor emits input_ids + pixel_values/mm_token_type_ids/image_position_ids;
+        # the extra keys ride to the DSpark prefill (prefill_mm) so the target embeds the image and the
+        # draft speculates over image-aware hidden states — vision on the same speculative loop as text.
+        enc = PROC.apply_chat_template(
+            chat_msgs, add_generation_prompt=True, tokenize=True, return_dict=True,
+            return_tensors="pt", enable_thinking=thinking,
+        ).to(EV.device)
+        ids = enc["input_ids"]
+        prefill_mm = {k: enc[k] for k in enc if k not in ("input_ids", "attention_mask")}
+    else:
+        ids = TOK.apply_chat_template(
+            chat_msgs, add_generation_prompt=True, enable_thinking=thinking,
+            return_tensors="pt", return_dict=True,
+        )["input_ids"].to(EV.device)
+        prefill_mm = None
     cid = "chatcmpl-" + uuid.uuid4().hex[:12]
     created = int(time.time())
     model = req.model or MODEL_NAME
@@ -228,7 +279,7 @@ def chat(req: ChatReq):
                     EV.args.temperature = temp
                     with sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]), torch.no_grad():
                         EV.generate_one_sample(
-                            input_ids=ids, stop_token_ids=STOP,
+                            input_ids=ids, stop_token_ids=STOP, prefill_mm=prefill_mm,
                             stream_callback=lambda t: q.put(t[0].tolist()),
                         )
             except Exception as e:  # surface generation errors to the stream instead of hanging
@@ -271,7 +322,7 @@ def chat(req: ChatReq):
         EV.args.temperature = temp
         t = time.time()
         with sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]), torch.no_grad():
-            res = EV.generate_one_sample(input_ids=ids, stop_token_ids=STOP)
+            res = EV.generate_one_sample(input_ids=ids, stop_token_ids=STOP, prefill_mm=prefill_mm)
         dt = time.time() - t
     gen = res.output_ids[0, res.num_input_tokens:].tolist()
     reasoning, content = _split_think(gen, thinking)