Update app.py

app.py

@@ -2,10 +2,20 @@ import gradio as gr
 import importlib
 from PIL import Image
 import json
-import os
+import spaces
 
-# === Load the GPT-4o module only
-from models import gpt4o_pix2struct_ocr
+# === 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
 
 # === Format Raw JSON Output
 def format_result_json(output):
@@ -128,12 +138,15 @@ def format_pretty_view(output):
 
     return "\n".join(lines).strip()
 
-# === Inference Handler (GPT-4o only)
-def process_image(image_file):
+# === Main Inference Handler
+def process_single_image(model_name, image_file):
+    runner = load_model_runner(model_name)
     image = Image.open(image_file.name).convert("RGB")
-    result = gpt4o_pix2struct_ocr.run_model(image)
+
+    result = runner(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)
@@ -144,23 +157,27 @@ def process_image(image_file):
 
     return image, json_output, pretty_output
 
-# === Gradio Interface
+# === Gradio UI
 iface = gr.Interface(
-    fn=process_image,
-    inputs=[gr.File(file_types=["image"], label="Upload a BPMN Diagram Image")],
+    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")
+    ],
     outputs=[
-        gr.Image(label="📷 Input Image"),
-        gr.Textbox(label="🧠 Raw JSON Output", lines=20),
-        gr.Textbox(label="📋 Prettified View", lines=25)
+        gr.Image(label="Input Image"),
+        gr.Textbox(label="Raw JSON Output (Technical)", lines=20),
+        gr.Textbox(label="Prettified View (User-Friendly)", lines=25)
     ],
-    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"
+    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"
 )
 
-# === Launch without GPU
+# === Enable GPU mode and launch
+#@spaces.GPU
 def main():
-    iface.launch(ssr=False)
+    iface.launch()
 
 if __name__ == "__main__":
     main()