Update app.py

app.py

@@ -2,20 +2,10 @@ import gradio as gr
 import importlib
 from PIL import Image
 import json
-import spaces
+import os
 
-# === Model Mapping ===
-MODEL_MAP = {
-    #"Qwen": "models.qwen",
-    #"Pixtral": "models.pixtral",
-    #"Aya Vision": "models.aya_vision",
-    "GPT-4o": "models.gpt4o"
-}
-
-# === Load Model
-def load_model_runner(model_name):
-    module = importlib.import_module(MODEL_MAP[model_name])
-    return module.run_model
+# === Load the GPT-4o module only
+from models import gpt4o_pix2struct_ocr
 
 # === Format Raw JSON Output
 def format_result_json(output):
@@ -138,15 +128,12 @@ def format_pretty_view(output):
 
     return "\n".join(lines).strip()
 
-# === Main Inference Handler
-def process_single_image(model_name, image_file):
-    runner = load_model_runner(model_name)
+# === Inference Handler (GPT-4o only)
+def process_image(image_file):
     image = Image.open(image_file.name).convert("RGB")
-
-    result = runner(image)
-
+    result = gpt4o_pix2struct_ocr.run_model(image)
     parsed_json = result.get("json")
-    raw_text = result.get("raw", "")
+    raw_text = result.get("raw")
 
     if parsed_json:
         json_output = format_result_json(parsed_json)
@@ -157,27 +144,23 @@ def process_single_image(model_name, image_file):
 
     return image, json_output, pretty_output
 
-# === Gradio UI
+# === Gradio Interface
 iface = gr.Interface(
-    fn=process_single_image,
-    inputs=[
-        gr.Dropdown(choices=list(MODEL_MAP.keys()), label="Select Vision Model"),
-        gr.File(file_types=["image"], label="Upload a BPMN Image")
-    ],
+    fn=process_image,
+    inputs=[gr.File(file_types=["image"], label="Upload a BPMN Diagram Image")],
     outputs=[
-        gr.Image(label="Input Image"),
-        gr.Textbox(label="Raw JSON Output (Technical)", lines=20),
-        gr.Textbox(label="Prettified View (User-Friendly)", lines=25)
+        gr.Image(label="📷 Input Image"),
+        gr.Textbox(label="🧠 Raw JSON Output", lines=20),
+        gr.Textbox(label="📋 Prettified View", lines=25)
     ],
-    title="🖼️ Vision Model Extractor - JSON + Pretty View",
-    description="Upload a BPMN image and select a vision model to extract structured output. GPT-4o uses an API key from your Hugging Face Space Secret.",
-    flagging_mode="never"
+    title="🧩 BPMN Extractor using GPT-4o + OCR",
+    description="Upload a BPMN diagram image. Extracts structured JSON using GPT-4o and Pix2Struct OCR. Runs on CPU-only Space.",
+    allow_flagging="never"
 )
 
-# === Enable GPU mode and launch
-#@spaces.GPU
+# === Launch without GPU
 def main():
-    iface.launch()
+    iface.launch(ssr=False)
 
 if __name__ == "__main__":
     main()