import argparse
import gradio as gr
import torch
from PIL import Image
import re
from transformers import DonutProcessor, VisionEncoderDecoderModel

def demo_process(input_img, question=None):
    global processor, model

    input_img = Image.fromarray(input_img)
    pixel_values = processor(input_img, return_tensors="pt").pixel_values.to(device)

    if question:
        task_prompt = f"<s_{question}>"
        decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids.to(device)
    else:
        task_prompt = "<s_cord-v2>"
        decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids.to(device)

    with torch.no_grad():
        outputs = model.generate(
            pixel_values,
            decoder_input_ids=decoder_input_ids,
            max_length=1024,  
            early_stopping=True,
            pad_token_id=processor.tokenizer.pad_token_id,
            eos_token_id=processor.tokenizer.eos_token_id,
            use_cache=True,
            num_beams=1,
            bad_words_ids=[[processor.tokenizer.unk_token_id]],
            return_dict_in_generate=True,
        )

    seq = processor.batch_decode(outputs.sequences)[0]
    seq = seq.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
    seq = re.sub(r"<.*?>", "", seq, count=1).strip()
    seq = processor.token2json(seq)
    return seq

parser = argparse.ArgumentParser()
parser.add_argument("--task", type=str, default="cord-v2")
parser.add_argument("--pretrained_path", type=str, default="suthawadee/donut-demo_new")
args, left_argv = parser.parse_known_args()

processor = DonutProcessor.from_pretrained(args.pretrained_path)
model = VisionEncoderDecoderModel.from_pretrained(args.pretrained_path)
device = "cpu" if not torch.cuda.is_available() else "cuda"
model.to(device)
model.eval()

# เพิ่มตัวอย่างรูปภาพที่มีอยู่เพื่อทดสอบ
image1 = "8.jpg"
image2 = "15.jpg"

examples = [
    [Image.open(image1)],
    [Image.open(image2)]
]

def main(pretrained_path, examples):
    demo = gr.Interface(
        fn=demo_process,
        inputs=["image", "text"] if args.task == "docvqa" else "image",
        outputs="json",
        title="🇹🇭🧾ThaiReceipt",
        description="Upload image.",
        examples=examples
    )
    demo.launch(debug=True)

main(args.pretrained_path, examples)