import gradio as gr
import torch
import os
from PIL import Image

MODEL_ID = "OdiaGenAIOCR/odia-ocr-qwen-finetuned_v3"
CHECKPOINT = "checkpoint-1300"

_model = None
_processor = None

def load_model():
    global _model, _processor
    if _model is not None:
        return _model, _processor

    from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
    from peft import PeftModel

    BASE_MODEL = "Qwen/Qwen2.5-VL-7B-Instruct"
    device = "cuda" if torch.cuda.is_available() else "cpu"

    _processor = AutoProcessor.from_pretrained(
        MODEL_ID,
        revision=CHECKPOINT,
        trust_remote_code=True,
    )

    _model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        MODEL_ID,
        revision=CHECKPOINT,
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        device_map="auto" if device == "cuda" else None,
        trust_remote_code=True,
    )
    if device == "cpu":
        _model = _model.to(device)
    _model.eval()
    return _model, _processor


def ocr_predict(image: Image.Image):
    if image is None:
        return "Please upload an image."

    device = "cuda" if torch.cuda.is_available() else "cpu"
    try:
        model, processor = load_model()
    except Exception as e:
        return f"Model load error: {e}"

    from qwen_vl_utils import process_vision_info

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": "Please perform OCR on this image and output only the recognized text."},
            ],
        }
    ]

    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    ).to(device)

    with torch.no_grad():
        generated_ids = model.generate(**inputs, max_new_tokens=512)

    generated_ids_trimmed = [
        out_ids[len(in_ids):]
        for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    return output_text[0]


iface = gr.Interface(
    fn=ocr_predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Textbox(label="Recognized Odia Text"),
    title="OdiaGenAI OCR",
    description="Upload Odia images or scanned documents for OCR using the fine-tuned Qwen2.5-VL model.",
    examples=[],
)

iface.launch()