Update app.py

app.py

@@ -1,14 +1,20 @@
 import gradio as gr
-from pdf2image import convert_from_path
+import argparse
+import torch
+from PIL import Image
+from donut import DonutModel
+from io import BytesIO
 
-def pdf_to_jpg(pdf_path):
-    images = convert_from_path(pdf_path.name, first_page=1, last_page=1)
-    return images[0].save("first_page.jpg")
+def demo_process(pdf_file):
+    global model, task_name, task_prompt   
+    images = convert_from_path(pdf_file.name)    
+    # output = model.inference(image=input_img, prompt=task_prompt)["predictions"][0]
+    return str(len(images))
 
-interface = gr.Interface(
-    fn=pdf_to_jpg,
-    inputs=gr.File(label="Upload PDF"),
-    outputs="image"
-)
+model = DonutModel.from_pretrained("uartimcs/donut-invoice-extract")
+model.eval()
+task_name = "SGSInvoice"
+task_prompt = f"<s_{task_name}>"
 
-interface.launch()
+demo = gr.Interface(fn=demo_process,inputs=gr.File(label="Upload PDF"),outputs="json", title=f"Donut 🍩 demonstration for `{task_name}` task",)
+demo.launch()