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

from upsampler_theme import UPSAMPLER_THEME, UPSAMPLER_CSS, footer_html, header_html

MODEL_ID = "PaddlePaddle/PaddleOCR-VL"
# Pinned to the last transformers-4.x-compatible revision: the repo's current
# code straddles the v5 API break (v5 create_causal_mask naming with v4 rope
# tables) and imports on no released transformers version.
MODEL_REV = "7faba727a1c11ec07247fd41a52868ed966c7cb3"

PROMPTS = {
    "Text (OCR)": "OCR:",
    "Table": "Table Recognition:",
    "Formula": "Formula Recognition:",
    "Chart": "Chart Recognition:",
}

model = (
    AutoModelForCausalLM.from_pretrained(
        MODEL_ID, revision=MODEL_REV, trust_remote_code=True, torch_dtype=torch.bfloat16
    )
    .to("cuda")
    .eval()
)
processor = AutoProcessor.from_pretrained(MODEL_ID, revision=MODEL_REV, trust_remote_code=True)


def get_duration(image, task, progress=None):
    # 0.9B model: a page is well under 30s; scale a little with input size and
    # never over-request (anonymous quota pre-checks the requested duration).
    if image is None:
        return 10
    mp = (image.width * image.height) / 1_000_000
    return min(60, int(20 + 10 * mp))


@spaces.GPU(duration=get_duration)
def extract_text(image, task, progress=gr.Progress(track_tqdm=True)):
    if image is None:
        raise gr.Error("Please upload an image first.")
    image = image.convert("RGB")
    # Bound the longest side; the NaViT encoder handles native ratios, and
    # very large scans only slow decoding without helping accuracy.
    if max(image.size) > 2048:
        scale = 2048 / max(image.size)
        image = image.resize(
            (int(image.width * scale), int(image.height * scale)), Image.LANCZOS
        )
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": PROMPTS.get(task, "OCR:")},
            ],
        }
    ]
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to("cuda")
    with torch.inference_mode():
        outputs = model.generate(**inputs, max_new_tokens=2048, do_sample=False)
    generated = outputs[:, inputs["input_ids"].shape[1] :]
    text = processor.batch_decode(generated, skip_special_tokens=True)[0]
    return text.strip()


with gr.Blocks(theme=UPSAMPLER_THEME, css=UPSAMPLER_CSS) as demo:
    gr.HTML(
        header_html(
            "PaddleOCR-VL: Image to Text",
            "Extract text, tables, formulas, and chart data from any image",
        )
    )
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="pil", label="Image")
            task = gr.Dropdown(
                choices=list(PROMPTS.keys()),
                value="Text (OCR)",
                label="What to extract",
            )
            run = gr.Button("Extract Text", variant="primary")
        with gr.Column(scale=1):
            output = gr.Textbox(
                label="Extracted text",
                lines=18,
                show_copy_button=True,
            )
    run.click(extract_text, inputs=[input_image, task], outputs=output, api_name="extract_text")
    gr.HTML(
        footer_html(
            "PaddleOCR-VL is a 0.9B vision-language OCR model by PaddlePaddle that "
            "converts images to text in 109 languages, including handwriting, and "
            "reads tables, mathematical formulas, and charts from photos, scans, and "
            "screenshots.",
            "https://upsampler.com/free-image-to-text-no-signup",
            "free image to text converter",
        )
    )

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