import gradio as gr
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
    
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="json", title=f"Donut 🍩 demonstration for `{task_name}` task",)
demo.launch()