import re
import json
import torch
import numpy as np
import gradio as gr
from PIL import Image
from transformers import DonutProcessor, VisionEncoderDecoderModel

auth_tok="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"


def demo_process(input_img):
    global processor, pretrained_model, task_prompt, task_name
    input_img = Image.fromarray(input_img)
    # prepare encoder inputs
    pixel_values = processor(input_img.convert("RGB"), return_tensors="pt").pixel_values
    pixel_values = pixel_values.to(device)
    # prepare decoder inputs
    task_prompt = "<s_lotoquine>"
    decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
    decoder_input_ids = decoder_input_ids.to(device)

    # autoregressively generate sequence
    outputs = pretrained_model.generate(
              pixel_values,
              decoder_input_ids=decoder_input_ids,
              max_length=pretrained_model.decoder.config.max_position_embeddings,
              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,
          )

    # turn into JSON
    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()  # remove first task start token
    seq = processor.token2json(seq)

    return seq



processor = DonutProcessor.from_pretrained("Aigle974/donut-lotoquine",use_auth_token=auth_tok)
pretrained_model = VisionEncoderDecoderModel.from_pretrained("Aigle974/donut-lotoquine",use_auth_token=auth_tok)
processor.feature_extractor.do_align_long_axis = True



device ="cuda" if torch.cuda.is_available() else "cpu"
pretrained_model.to(device)


pretrained_model.eval()

demo = gr.Interface(
    fn=demo_process,
    inputs="image",
    outputs="json",
    title=f"Lotoquine Automatic Extraction by Fab",
)
demo.launch()