import json
import tempfile

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

MODEL_NAME = "Qwen/Qwen3.5-4B"

print("Loading model...")

processor = AutoProcessor.from_pretrained(MODEL_NAME)

model = AutoModelForMultimodalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float32,
    device_map="cpu"
)

print("Model Loaded!")


def extract_text(image):

    if image is None:
        return {}, None

    prompt = """
Extract all visible text from this image.

Return ONLY valid JSON.

{
    "ocr_text":"..."
}
"""

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

    text = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    inputs = processor(
        text=text,
        images=image,
        return_tensors="pt"
    )

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256
        )

    generated_ids = outputs[:, inputs["input_ids"].shape[-1]:]

    response = processor.batch_decode(
        generated_ids,
        skip_special_tokens=True
    )[0]

    try:
        data = json.loads(response)
    except Exception:
        data = {
            "ocr_text": response.strip()
        }

    temp = tempfile.NamedTemporaryFile(
        delete=False,
        suffix=".json",
        mode="w",
        encoding="utf-8"
    )

    json.dump(
        data,
        temp,
        indent=4,
        ensure_ascii=False
    )

    temp.close()

    return data, temp.name


demo = gr.Interface(
    fn=extract_text,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.JSON(label="OCR JSON"),
        gr.File(label="Download JSON")
    ],
    title="Qwen OCR Extractor",
    description="Upload an image and extract all visible text as JSON."
)

demo.launch()