from collections.abc import Generator
from threading import Thread

import gradio as gr
import spaces
import torch
from PIL import Image
from transformers import AutoModelForCausalLM, AutoProcessor, TextIteratorStreamer

MODEL_ID = "sbintuitions/sarashina2.2-ocr"

processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True, use_fast=False)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="cuda",
    dtype=torch.bfloat16,
    trust_remote_code=True,
)


@spaces.GPU(duration=90)
def run_ocr(image: Image.Image | None) -> Generator[tuple[str, str], None, None]:
    if image is None:
        yield "", ""
        return

    message = [{"role": "user", "content": [{"type": "image", "image": image}]}]
    inputs = processor.apply_chat_template(
        message,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)

    streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True, timeout=20.0)
    generate_kwargs = dict(
        **inputs,
        max_new_tokens=6000,
        temperature=0.0,
        top_p=0.95,
        repetition_penalty=1.2,
        use_cache=True,
        streamer=streamer,
    )

    exception_holder: list[Exception] = []

    def _generate() -> None:
        try:
            model.generate(**generate_kwargs)
        except Exception as e:  # noqa: BLE001
            exception_holder.append(e)

    thread = Thread(target=_generate)
    thread.start()

    result = ""
    for text in streamer:
        result += text
        yield result, result

    thread.join()
    if exception_holder:
        msg = f"Generation failed: {exception_holder[0]}"
        raise gr.Error(msg)


with gr.Blocks() as demo:
    gr.Markdown("# Sarashina2.2-OCR Demo")
    gr.Markdown(
        "Upload a document image to extract text using "
        "[sbintuitions/sarashina2.2-ocr](https://huggingface.co/sbintuitions/sarashina2.2-ocr)."
    )
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Document Image", type="pil")
            run_btn = gr.Button("Run OCR")
        with gr.Column():
            with gr.Tab("Rendered"):
                output_md = gr.Markdown(
                    label="Result",
                    latex_delimiters=[
                        {"left": "$$", "right": "$$", "display": True},
                        {"left": "$", "right": "$", "display": False},
                        {"left": "\\(", "right": "\\)", "display": False},
                        {"left": "\\[", "right": "\\]", "display": True},
                    ],
                )
            with gr.Tab("Raw"):
                output_text = gr.Textbox(label="Raw Markdown", lines=20)

    gr.on(
        triggers=[run_btn.click, image_input.upload],
        fn=run_ocr,
        inputs=image_input,
        outputs=[output_md, output_text],
    )

    gr.Examples(
        examples=[
            ["https://huggingface.co/sbintuitions/sarashina2.2-ocr/resolve/main/assets/sample1.jpeg"],
            ["https://huggingface.co/sbintuitions/sarashina2.2-ocr/resolve/main/assets/sample2.jpeg"],
            ["https://huggingface.co/sbintuitions/sarashina2.2-ocr/resolve/main/assets/sample3.jpeg"],
            ["https://huggingface.co/sbintuitions/sarashina2.2-ocr/resolve/main/assets/sample4.jpeg"],
        ],
        inputs=image_input,
        fn=run_ocr,
        outputs=[output_md, output_text],
    )

if __name__ == "__main__":
    demo.launch()