Update app.py

app.py

@@ -1,7 +1,86 @@
-from fastapi import FastAPI
+import io
+import torch
+from PIL import Image
+from threading import Thread
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.responses import JSONResponse
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+from transformers.generation.streamers import TextIteratorStreamer
 
 app = FastAPI()
 
-@app.get("/")
-def greet_json():
-    return {"Hello": "World!"}
+# Setup device
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+# Load model and processor
+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,
+    )
+    .to(device)
+    .eval()
+)
+
+
+def generate_response(image, prompt_text, **kwargs):
+    images = [image]
+
+    messages = [
+        {
+            "role": "user",
+            "content": [{"type": "image"} for _ in images]
+            + [{"type": "text", "text": prompt_text}],
+        }
+    ]
+    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = processor(text=prompt, images=images, 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),
+    }
+
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    output = ""
+    for chunk in streamer:
+        output += chunk.replace("<|im_end|>", "")
+    return output.strip()
+
+
+@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")
+
+    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})