Update app.py

app.py

@@ -1,86 +1,102 @@
 import gradio as gr
-import torch
+import requests
+import base64
 import json
-import os
 from PIL import Image
-from transformers import AutoProcessor, AutoModelForVision2Seq
-from qwen_vl_utils import process_vision_info
-
-# Use 3B to stay within 16GB RAM limit
-MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct" 
-
-print("Loading processor...")
-processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
-
-print(f"Loading model {MODEL_ID} on CPU...")
-# AutoModelForVision2Seq is the correct class for Vision-Language models
-model = AutoModelForVision2Seq.from_pretrained(
-    MODEL_ID,
-    trust_remote_code=True,
-    torch_dtype=torch.float32, 
-    low_cpu_mem_usage=True,
-    device_map="cpu"
-)
-model.eval()
-print("Model loaded successfully!")
+import io
+import os
+
+OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
+
+MODEL = "qwen/qwen-2.5-vl-72b-instruct"
+
+def image_to_base64(image: Image.Image):
+    buf = io.BytesIO()
+    image.save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode()
 
 def extract_document(image: Image.Image):
     if image is None:
-        return {"error": "Please upload an image."}
-
-    # Format prompt for Qwen2.5-VL
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image", "image": image},
-                {"type": "text", "text": "Extract all data from this document as a clean JSON object. Include document_type, fields, and tables."},
-            ],
-        }
-    ]
-
-    # Prepare inputs
-    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("cpu")
-
-    # Generate output (Takes ~1 min on CPU)
-    with torch.no_grad():
-        generated_ids = model.generate(**inputs, max_new_tokens=1024)
-        
-    generated_ids_trimmed = [
-        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-    ]
-    output_text = processor.batch_decode(
-        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-    )[0]
+        return {"error": "No image uploaded"}
+
+    img_b64 = image_to_base64(image)
+
+    prompt = """
+You are a universal document understanding system.
+
+Rules:
+- Be document-agnostic
+- Detect document type if possible
+- Extract ALL visible structured data
+- Extract tables completely (columns + rows)
+- Preserve numbers exactly
+- Use null for missing values
+- Do NOT hallucinate
+- Return ONLY valid JSON
+
+Schema:
+{
+  "document_type": string | null,
+  "confidence": number (0-1),
+  "summary": string,
+  "fields": { "<key>": "<value | null>" },
+  "tables": [
+    {
+      "table_name": string,
+      "columns": [string],
+      "rows": [[string | number | null]]
+    }
+  ]
+}
+"""
+
+    payload = {
+        "model": MODEL,
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "input_text", "text": prompt},
+                    {
+                        "type": "input_image",
+                        "image_base64": img_b64
+                    }
+                ]
+            }
+        ],
+        "temperature": 0
+    }
+
+    headers = {
+        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
+        "Content-Type": "application/json"
+    }
+
+    r = requests.post(
+        "https://openrouter.ai/api/v1/chat/completions",
+        headers=headers,
+        json=payload,
+        timeout=120
+    )
+
+    response = r.json()
+    text = response["choices"][0]["message"]["content"]
 
     try:
-        # Simple extraction of JSON block
-        start = output_text.find("{")
-        end = output_text.rfind("}") + 1
-        return json.loads(output_text[start:end])
-    except:
-        return {"raw_output": output_text}
-
-with gr.Blocks() as demo:
-    gr.Markdown("# 📄 DocAI — Universal Document Intelligence")
-    gr.Markdown("Using Qwen2.5-VL-3B on CPU.")
-    
+        start = text.find("{")
+        end = text.rfind("}") + 1
+        return json.loads(text[start:end])
+    except Exception:
+        return {"raw_output": text}
+
+with gr.Blocks(title="DocAI – Universal Document Extractor") as demo:
+    gr.Markdown("# 📄 DocAI – Universal Document Intelligence")
+    gr.Markdown("Vision-LLM powered. No templates. Any document.")
+
     with gr.Row():
-        with gr.Column():
-            input_img = gr.Image(type="pil", label="Upload Document")
-            submit_btn = gr.Button("Extract Data", variant="primary")
-        with gr.Column():
-            output_json = gr.JSON(label="Extracted JSON")
+        img = gr.Image(type="pil", label="Upload document")
+        out = gr.JSON(label="Extracted JSON")
 
-    submit_btn.click(fn=extract_document, inputs=input_img, outputs=output_json)
+    gr.Button("Extract").click(extract_document, img, out)
 
-demo.launch()
+demo.launch()