File size: 2,302 Bytes
df1fa7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | import io
import threading
import torch
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor
MODEL_ID = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"
MAX_IMAGES = 20
MAX_FILE_BYTES = 10 * 1024 * 1024
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageTextToText.from_pretrained(MODEL_ID, dtype=torch.float32).eval()
# ponytail: one generation at a time; free CPU Space has 2 vCPUs, parallel requests would just thrash.
lock = threading.Lock()
app = FastAPI(title="SmolVLM2 API")
@app.get("/")
def health():
return {"status": "ok", "model": MODEL_ID, "max_images": MAX_IMAGES}
@app.post("/generate")
def generate(
prompt: str = Form(...),
images: list[UploadFile] = File(default=[]),
max_new_tokens: int = Form(256, ge=1, le=1024),
):
if len(images) > MAX_IMAGES:
raise HTTPException(400, f"Max {MAX_IMAGES} images per request")
pil_images = []
for f in images:
data = f.file.read(MAX_FILE_BYTES + 1)
if len(data) > MAX_FILE_BYTES:
raise HTTPException(413, f"{f.filename} exceeds {MAX_FILE_BYTES // 2**20}MB")
try:
pil_images.append(Image.open(io.BytesIO(data)).convert("RGB"))
except Exception:
raise HTTPException(400, f"{f.filename} is not a valid image")
content = [{"type": "image"} for _ in pil_images] + [{"type": "text", "text": prompt}]
text = processor.apply_chat_template([{"role": "user", "content": content}], add_generation_prompt=True)
with lock, torch.inference_mode():
# Splitting turns each image into up to ~17 tiles: great for 1 image, but 20 images overflow the 8k context.
# Set on the shared processor (not a call kwarg, which transformers ignores), hence inside the lock.
processor.image_processor.do_image_splitting = len(pil_images) == 1
inputs = processor(text=[text], images=[pil_images] if pil_images else None, return_tensors="pt")
out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
answer = processor.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True)
return {"response": answer.strip(), "num_images": len(pil_images)}
|