Update app.py

app.py

@@ -3,16 +3,22 @@ import argparse
 import torch
 from donut import DonutModel
 from pdf2image import convert_from_path
+from poppler import load_from_file, PageRenderer
 
 def demo_process(pdf_file):
+    pdf_document = load_from_file(pdf_file.name)
+    page_1 = pdf_document.create_page(0)
+
+    renderer = PageRenderer()
+    image = renderer.render_page(page_1)
+    image_data = image.data
+    output = model.inference(image=image_data, prompt=task_prompt)["predictions"][0]
+    return output
     
-    images = convert_from_path(pdf_file.name) 
-    # output = model.inference(image=input_img, prompt=task_prompt)["predictions"][0]
-    return str(len([1,2,3,4,5]))
-# model = DonutModel.from_pretrained("uartimcs/donut-invoice-extract")
-# model.eval()
+model = DonutModel.from_pretrained("uartimcs/donut-invoice-extract")
+model.eval()
 task_name = "SGSInvoice"
 task_prompt = f"<s_{task_name}>"
 
-demo = gr.Interface(fn=demo_process,inputs=gr.File(label="Upload PDF"),outputs="text", title=f"Donut 🍩 demonstration for `{task_name}` task",)
+demo = gr.Interface(fn=demo_process,inputs=gr.File(label="Upload PDF"),outputs="json", title=f"Donut 🍩 demonstration for `{task_name}` task",)
 demo.launch()