app.py

import gradio as gr
import os
import shutil
import argparse
import torch
from PIL import Image
from pdf2image import convert_from_path
from donut import DonutModel


def demo_process(doc):
    global model, task_prompt, task_name
    file_name = os.path.basename(doc)
    new_file_location = os.path.join(os.getcwd(),file_name)
    
    shutil.copyfile(doc, new_file_location)
    img = convert_from_path(new_file_location)[0]
    output = model.inference(image=img, prompt=task_prompt)["predictions"][0]
    return output
    
parser = argparse.ArgumentParser()
parser.add_argument("--task", type=str, default="SGSInvoice")
parser.add_argument("--pretrained_path", type=str, default="uartimcs/donut-invoice-extract")
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="json", title=f"Donut 🍩 demonstration for `{task_name}` task",)
demo.launch()