import spaces
import torch
import gradio as gr
from PIL import Image
from transformers import AutoProcessor, AutoModelForImageTextToText
from peft import PeftModel

BASE_MODEL_ID = "google/gemma-3-12b-it"
ADAPTER_ID = "historyHulk/ModiTrans-12B-Gemma-Teacher"
PROMPT = "Translitrate the following Modi script to Devnagri script."
MAX_NEW_TOKENS = 350

# Model + processor are loaded once at startup and kept on CPU.
# They are moved to GPU inside the @spaces.GPU-decorated function,
# which is how ZeroGPU spaces work (GPU only attached per-call).
print("Loading processor...")
processor = AutoProcessor.from_pretrained(BASE_MODEL_ID)

print("Loading base model...")
base_model = AutoModelForImageTextToText.from_pretrained(
    BASE_MODEL_ID,
    torch_dtype=torch.bfloat16,
)

print("Loading LoRA adapter...")
model = PeftModel.from_pretrained(
    base_model,
    ADAPTER_ID,
    torch_dtype=torch.bfloat16,
)
model.eval()


@spaces.GPU
def transliterate(image: Image.Image) -> str:
    if image is None:
        return "Please upload an image of Modi script."

    device = "cuda"
    model.to(device)

    image = image.convert("RGB").resize((1024, 512))

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": PROMPT},
            ],
        },
    ]

    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to(device, dtype=torch.bfloat16)

    input_len = inputs["input_ids"].shape[-1]

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
        )

    generated = output_ids[0][input_len:]
    text = processor.decode(generated, skip_special_tokens=True)
    return text.strip()


with gr.Blocks(title="ModiTrans: Modi Script to Devanagari") as demo:
    gr.Markdown(
        """
        # ModiTrans — Modi Script to Devanagari Transliteration

        Upload a scanned image of historic **Modi script** text and this model
        will transliterate it into modern **Devanagari** script.

        Uses [`historyHulk/ModiTrans-12B-Gemma-Teacher`](https://huggingface.co/historyHulk/ModiTrans-12B-Gemma-Teacher),
        a LoRA adapter on `google/gemma-3-12b-it`, from the paper
        *"Historic Scripts to Modern Vision: A Novel Dataset and A VLM Framework
        for Transliteration of Modi Script to Devanagari"* (ICDAR 2025).

        > This is a gated model — the Space owner's `HF_TOKEN` must have accepted
        > access on the model page for inference to work.
        """
    )

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Modi Script Image")
            run_btn = gr.Button("Transliterate", variant="primary")
        with gr.Column():
            output_text = gr.Textbox(
                label="Devanagari Transliteration", lines=8
            )

    run_btn.click(fn=transliterate, inputs=image_input, outputs=output_text)
    image_input.change(fn=transliterate, inputs=image_input, outputs=output_text)

    gr.Markdown(
        """
        ---
        **Citation:** Kausadikar, H., Kale, T., Susladkar, O., Mittal, S.
        *Historic Scripts to Modern Vision.* ICDAR 2025 (Springer LNCS).
        """
    )

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