Update main.py

main.py

@@ -1,94 +1,48 @@
-import io
-import torch
-from PIL import Image
-from fastapi import FastAPI, File, UploadFile, Form
-from fastapi.responses import JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
+from fastapi import FastAPI, Query
 from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
-from transformers.generation.streamers import TextIteratorStreamer
+from qwen_vl_utils import process_vision_info
+import torch
 
 app = FastAPI()
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
+checkpoint = "Qwen/Qwen2.5-VL-3B-Instruct"
+min_pixels = 256*28*28
+max_pixels = 1280*28*28
+processor = AutoProcessor.from_pretrained(
+    checkpoint,
+    min_pixels=min_pixels,
+    max_pixels=max_pixels
 )
-
-# Device config
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Load processor and model
-MODEL_ID = "nanonets/Nanonets-OCR-s"
-processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
 model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_ID,
-    trust_remote_code=True,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    checkpoint,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    # attn_implementation="flash_attention_2",
 )
-model = model.to(device).eval()
 
+@app.get("/")
+def read_root():
+    return {"message": "API is live. Use the /predict endpoint."}
 
-def generate_response(image: Image.Image, prompt_text: str, **kwargs) -> str:
+@app.get("/predict")
+def predict(image_url: str = Query(...), prompt: str = Query(...)):
     messages = [
-        {
-            "role": "user",
-            "content": [{"type": "image"}, {"type": "text", "text": prompt_text}],
-        }
+        {"role": "system", "content": "You are a helpful assistant with vision abilities."},
+        {"role": "user", "content": [{"type": "image", "image": image_url}, {"type": "text", "text": prompt}]},
     ]
-    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
-    inputs = processor(text=prompt, images=[image], return_tensors="pt").to(device)
-
-    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = {
-        **inputs,
-        "streamer": streamer,
-        "max_new_tokens": kwargs.get("max_new_tokens", 1024),
-        "temperature": kwargs.get("temperature", 0.6),
-        "top_p": kwargs.get("top_p", 0.9),
-        "top_k": kwargs.get("top_k", 50),
-        "repetition_penalty": kwargs.get("repetition_penalty", 1.2),
-    }
-
-    # 🧠 Direct call (no thread)
-    model.generate(**generation_kwargs)
-
-    output = ""
-    for chunk in streamer:
-        output += chunk.replace("<|im_end|>", "")
-    return output.strip()
-
-
-@app.get("/")
-def health():
-    return {"status": "running", "model": MODEL_ID}
-
-
-@app.post("/ocr/image")
-async def ocr_image(
-    prompt: str = Form(...),
-    image: UploadFile = File(...),
-    max_new_tokens: int = Form(1024),
-    temperature: float = Form(0.6),
-    top_p: float = Form(0.9),
-    top_k: int = Form(50),
-    repetition_penalty: float = Form(1.2),
-):
-    image_bytes = await image.read()
-    pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-
-    try:
-        result = generate_response(
-            image=pil_image,
-            prompt_text=prompt,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            repetition_penalty=repetition_penalty,
-        )
-        return JSONResponse(content={"result": result})
-    except Exception as e:
-        return JSONResponse(status_code=500, content={"error": str(e)})
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+    ).to(model.device)
+    with torch.no_grad():
+        generated_ids = model.generate(**inputs, max_new_tokens=128)
+    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+    output_texts = processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )
+    return {"response": output_texts[0]}