from __future__ import annotations

import gradio as gr
import torch
from PIL import Image
from transformers import AutoModel, AutoProcessor


MODEL_ID = "YOUR_USER/YOUR_MODEL"
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True)
model.eval()


def recognize(image: Image.Image) -> str:
    if image is None:
        return ""
    with torch.no_grad():
        inputs = processor(images=image, return_tensors="pt")
        logits = model(**inputs).logits
    return processor.batch_decode(logits)[0]


demo = gr.Interface(
    fn=recognize,
    inputs=gr.Image(type="pil", label="Input image"),
    outputs=gr.Textbox(label="Predicted text"),
    title="Ukrainian OCR / ICR",
)

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