Update app.py

app.py

@@ -1,47 +1,72 @@
-import gradio as gr
 import torch
-from transformers import Idefics3Processor, Idefics3ForConditionalGeneration
+from transformers import AutoProcessor, AutoModelForVision2Seq
 from PIL import Image
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import JSONResponse
+import io
 
-# Load model and processor
-model_id = "HuggingFaceM4/Idefics3-8B-Llama3"
+app = FastAPI()
 
-# FIX: Import the specific processor class directly to bypass the AutoProcessor bug
-processor = Idefics3Processor.from_pretrained(model_id, trust_remote_code=True)
+# SmolVLM - lightweight, works on CPU/low VRAM, same team as Idefics3
+model_id = "HuggingFaceTB/SmolVLM-Instruct"
 
-model = Idefics3ForConditionalGeneration.from_pretrained(
-    model_id, 
-    trust_remote_code=True, 
-    torch_dtype=torch.float32
+print("Loading processor...")
+processor = AutoProcessor.from_pretrained(model_id)
+
+print("Loading model...")
+model = AutoModelForVision2Seq.from_pretrained(
+    model_id,
+    torch_dtype=torch.float32,   # float32 for CPU
+    device_map="auto"
 )
+model.eval()
+print("Model ready!")
+
+
+@app.get("/")
+def root():
+    return {"status": "running", "model": model_id}
+
+
+@app.post("/extract")
+async def extract_text(file: UploadFile = File(...)):
+    # Read uploaded image
+    contents = await file.read()
+    image = Image.open(io.BytesIO(contents)).convert("RGB")
+
+    # Prompt for NEET question extraction
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image"},
+                {"type": "text", "text": (
+                    "Extract all text from this image exactly as it appears. "
+                    "Preserve question numbers, options (A, B, C, D), tables, "
+                    "and any mathematical or chemical expressions. "
+                    "Format clearly."
+                )}
+            ]
+        }
+    ]
 
-def process_image(image):
-    # Safety check for empty input
-    if image is None:
-        return "Please upload an image first."
-        
-    # Ensure image is PIL format
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image)
-            
-    image = image.convert("RGB")
-        
-    # Prepare inputs
-    messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Describe this image."}]}]
     prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
-    inputs = processor(text=prompt, images=image, return_tensors="pt")
-        
-    # Generate
-    generated_ids = model.generate(**inputs, max_new_tokens=500)
-    result = processor.batch_decode(generated_ids, skip_special_tokens=True)
-        
-    return result[0]
-
-# UI Setup
-demo = gr.Interface(
-    fn=process_image,
-    inputs=gr.Image(type="pil"),
-    outputs="text"
-)
 
-demo.launch()
+    inputs = processor(
+        text=prompt,
+        images=[image],
+        return_tensors="pt"
+    )
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=1024,
+            do_sample=False
+        )
+
+    # Decode only the generated part
+    generated = outputs[0][inputs["input_ids"].shape[1]:]
+    result = processor.decode(generated, skip_special_tokens=True)
+
+    return JSONResponse({"extracted_text": result})