import spaces
import torch
import gradio as gr
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor

MODEL_ID = "NAMAA-Space/Qari-OCR-v0.3-VL-2B-Instruct"

# High max_pixels keeps small Arabic glyphs + tashkeel legible (OCR needs detail).
# 28 is Qwen2-VL's patch factor; 4096*28*28 ≈ 3.2M px ≈ a 300-DPI textbook page.
processor = AutoProcessor.from_pretrained(
    MODEL_ID, min_pixels=256 * 28 * 28, max_pixels=4096 * 28 * 28
)
# Loaded on CPU at startup — ZeroGPU only attaches a GPU inside @spaces.GPU functions.
model = Qwen2VLForConditionalGeneration.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)

PROMPT = (
    "Below is an image of one page of an Arabic school textbook. "
    "Transcribe ALL the Arabic text exactly as printed, preserving line breaks, "
    "headings, and right-to-left reading order. Keep diacritics (tashkeel) if present. "
    "Ignore any faint diagonal draft watermark. Output only the transcribed text."
)


@spaces.GPU(duration=120)
def ocr(image):
    if image is None:
        return ""
    model.to("cuda")
    messages = [{"role": "user", "content": [
        {"type": "image", "image": image},
        {"type": "text", "text": PROMPT},
    ]}]
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    # Feed the PIL image straight to the processor — avoids qwen-vl-utils/torchvision.
    inputs = processor(
        text=[text], images=[image], padding=True, return_tensors="pt",
    ).to("cuda")
    with torch.no_grad():
        generated = model.generate(**inputs, max_new_tokens=4096, do_sample=False)
    trimmed = [g[len(i):] for i, g in zip(inputs.input_ids, generated)]
    return processor.batch_decode(
        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]


with gr.Blocks(title="Qari-OCR — Arabic page OCR") as demo:
    gr.Markdown("## Qari-OCR — Arabic textbook page OCR\nUpload a page image, or call the `/ocr` API.")
    with gr.Row():
        inp = gr.Image(type="pil", label="Page image")
        out = gr.Textbox(label="Transcribed Arabic", lines=25, rtl=True)
    gr.Button("Run OCR", variant="primary").click(ocr, inp, out, api_name="ocr")

if __name__ == "__main__":
    demo.queue(max_size=16).launch()