Update app.py

app.py

@@ -6,12 +6,16 @@ from PIL import Image
 from pdf2image import convert_from_path
 from donut import DonutModel
 
-def demo_process(doc):
-    global model, task_prompt, task_name
+def process_pdf(doc):
     doc = os.path.join(os.getcwd(), doc)
     img = convert_from_path(doc)[0]
-    # output = model.inference(image=input_img, prompt=task_prompt)["predictions"][0]
-    return doc
+    return img
+
+def demo_process(doc):
+    global model, task_prompt, task_name
+    img = process_pdf(doc)
+    output = model.inference(image=img, prompt=task_prompt)["predictions"][0]
+    return output
     
 parser = argparse.ArgumentParser()
 parser.add_argument("--task", type=str, default="SGSInvoice")
@@ -20,7 +24,7 @@ args, left_argv = parser.parse_known_args()
 task_name = args.task
 task_prompt = f"<s_{task_name}>"
 
-# model = DonutModel.from_pretrained("uartimcs/donut-invoice-extract")
-# model.eval()
-demo = gr.Interface(fn=demo_process,inputs=gr.File(label="Upload PDF"),outputs="text", title=f"Donut 🍩 demonstration for `{task_name}` task",)
+model = DonutModel.from_pretrained("uartimcs/donut-invoice-extract")
+model.eval()
+demo = gr.Interface(fn=demo_process,inputs=gr.File(label="Upload PDF"),outputs="json", title=f"Donut 🍩 demonstration for `{task_name}` task",)
 demo.launch()