| 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() |
| |
| 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(): |
| |
| |
| 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)} |
|
|
