Update app.py

app.py

@@ -1,19 +1,24 @@
+import os
 import gradio as gr
-from transformers import AutoProcessor, AutoModelForSeq2SeqLM
+from transformers import Pix2StructProcessor, Pix2StructForConditionalGeneration
 from PIL import Image
 import json
 
-# Load DePlot (vision-to-table)
-deplot_model = AutoModelForSeq2SeqLM.from_pretrained("google/deplot")
-deplot_processor = AutoProcessor.from_pretrained("google/deplot")
+# Fix threading error
+os.environ["OMP_NUM_THREADS"] = "1"
+
+# Load DePlot
+model_id = "google/deplot"
+processor = Pix2StructProcessor.from_pretrained(model_id)
+model = Pix2StructForConditionalGeneration.from_pretrained(model_id)
 
 def extract_chart(image):
     # Step 1: Run DePlot
-    inputs = deplot_processor(images=image, text="Generate table from chart.", return_tensors="pt")
-    outputs = deplot_model.generate(**inputs, max_new_tokens=512)
-    table = deplot_processor.decode(outputs[0], skip_special_tokens=True)
+    inputs = processor(images=image, text="Generate table from chart.", return_tensors="pt")
+    predictions = model.generate(**inputs, max_new_tokens=512)
+    table = processor.decode(predictions[0], skip_special_tokens=True)
 
-    # Step 2: Create dummy structured JSON (later we plug in detection+OCR)
+    # Step 2: Dummy structured JSON
     structured_json = {
         "metadata": {"title": "Demo Chart", "chart_type": "bar", "confidence": 0.5},
         "axes": {"x_axis": {"label": "X", "ticks": []}, "y_axis": {"label": "Y", "ticks": []}},
@@ -21,7 +26,7 @@ def extract_chart(image):
         "legend": {"entries": []}
     }
 
-    # Step 3: Merge (for now, just attach DePlot’s output)
+    # Step 3: Merge outputs
     merged_output = {
         "structured_json": structured_json,
         "deplot_table": table,