import gradio as gr
import argparse
import torch
from donut import DonutModel
from pdf2image import convert_from_path

def demo_process(pdf_file):
    # images = convert_from_path(pdf_file)    
    # output = model.inference(image=input_img, prompt=task_prompt)["predictions"][0]
    return pdf_file.name

# 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.launch()