import gradio as gr
import torch
from PIL import Image

from donut import DonutModel

def demo_process(input_img):
    global pretrained_model, task_prompt, task_name
    # input_img = Image.fromarray(input_img)
    output = pretrained_model.inference(image=input_img, prompt=task_prompt)["predictions"][0]
    return output
task_name = "preparedFinetuneData_Bird"
# task_name = "cord-v2"
task_prompt = f"<s_{task_name}>"

image = Image.open("inv87.jpg")
image.save("inv87.jpg")
image = Image.open("inv17.jpg")
image.save("inv17.jpg")

PATH = 'epochs30_base_on_donut_base/'
pretrained_model = DonutModel.from_pretrained("doshan1250/p9OcrAiV2Bird")
pretrained_model.eval()

demo = gr.Interface(
    fn=demo_process,
    inputs= gr.Image(type="pil"),
    outputs="json",
    title=f"Goodarc p9 for `{task_name}` task, epochs30",
    description="""Goodarc p9 v2 訓練. 
    """,
    examples=[["inv87.jpg"], ["inv17.jpg"]],
    cache_examples=False,
)

demo.launch()