"""
DeepSeek-OCR Gradio Space (CPU)
-------------------------------
Upload an image and run it through deepseek-ai/DeepSeek-OCR using the
built-in "Free OCR" prompt mode (plain text extraction, no markdown/
grounding boxes). A "Markdown" mode is also included since it's a
one-line prompt swap.

This version runs on CPU only (no CUDA, no flash-attention). The model
stays in bfloat16, matching hardcoded dtype casts inside its own custom
code — PyTorch supports bfloat16 tensors on CPU, just without hardware
acceleration, so it works but is slow. DeepSeek-OCR is a vision-language
model, so expect anywhere from ~30s to a few minutes per image depending
on resolution preset and CPU. Use the "Tiny" or "Small" resolution
presets for faster (if lower-fidelity) results.
"""

import os
import re
import shutil
import tempfile

import gradio as gr
import torch
from transformers import AutoModel, AutoTokenizer

MODEL_NAME = "deepseek-ai/DeepSeek-OCR"

# Force CPU regardless of what's visible to torch.
os.environ["CUDA_VISIBLE_DEVICES"] = ""
DEVICE = "cpu"

# ---------------------------------------------------------------------------
# Model loading (once, at startup)
# ---------------------------------------------------------------------------
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)

print("Loading model on CPU (this can take a while, and first run also "
      "downloads several GB of weights)...")
model = AutoModel.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    use_safetensors=True,
)
# NOTE: don't cast to float32 here. DeepSeek-OCR's own forward pass hardcodes
# casting image inputs to bfloat16 internally (independent of the model's
# actual dtype), so the model's weights must also be bfloat16 or you get a
# "Input type (BFloat16) and bias type (float) should be the same" crash.
# PyTorch CPU does support bfloat16 tensors/ops (just without hardware
# acceleration), so this works, it's just slow.
model = model.eval().to(DEVICE).to(torch.bfloat16)

# ---------------------------------------------------------------------------
# DeepSeek-OCR's own `model.infer()` method has `.cuda()` calls hardcoded
# directly in its custom modeling code (independent of what device the model
# itself was loaded on), so it crashes on CPU-only hardware with
# "Found no NVIDIA driver". Since the model and all its tensors already live
# on CPU here, we monkey-patch `.cuda()` into a no-op that just returns the
# tensor/module unchanged — this makes those hardcoded calls harmless instead
# of failing.
if not torch.cuda.is_available():
    torch.Tensor.cuda = lambda self, *args, **kwargs: self
    torch.nn.Module.cuda = lambda self, *args, **kwargs: self

# Preset resolution modes exposed by the model's own infer() signature
SIZE_PRESETS = {
    "Tiny (512, fastest)": dict(base_size=512, image_size=512, crop_mode=False),
    "Small (640)": dict(base_size=640, image_size=640, crop_mode=False),
    "Base (1024, high quality)": dict(base_size=1024, image_size=1024, crop_mode=False),
    "Large (1280, best quality)": dict(base_size=1280, image_size=1280, crop_mode=False),
    "Gundam (1024/640 crop, best for documents)": dict(base_size=1024, image_size=640, crop_mode=True),
}

TASK_PROMPTS = {
    "Free OCR (plain text)": "<image>\nFree OCR. ",
    "Markdown (structure + tables)": "<image>\n<|grounding|>Convert the document to markdown. ",
}


def run_ocr(image, task_choice, size_choice, progress=gr.Progress(track_tqdm=True)):
    if image is None:
        return "Please upload an image first."

    prompt = TASK_PROMPTS[task_choice]
    preset = SIZE_PRESETS[size_choice]

    # model.infer() reads the image from a file path and writes its result to
    # output_path. Point both at a temp dir.
    tmp_dir = tempfile.mkdtemp(prefix="deepseek_ocr_")
    image_path = os.path.join(tmp_dir, "input.png")
    image.convert("RGB").save(image_path)

    try:
        with torch.no_grad():
            # IMPORTANT: model.infer()'s own source code only `return`s the
            # extracted text when save_results=True (or test_compress/
            # eval_mode=True); with all three False/default it computes
            # everything but returns None. So we must pass save_results=True
            # and then read the text back from the file it writes
            # (output_path/result.mmd) rather than trusting the return value.
            model.infer(
                tokenizer,
                prompt=prompt,
                image_file=image_path,
                output_path=tmp_dir,
                base_size=preset["base_size"],
                image_size=preset["image_size"],
                crop_mode=preset["crop_mode"],
                save_results=True,
            )

        result_file = os.path.join(tmp_dir, "result.mmd")
        if not os.path.exists(result_file):
            return ("No output file was produced. Try a different resolution "
                    "preset, or check the Space logs for errors.")

        with open(result_file, "r", encoding="utf-8") as f:
            text = f.read()

        # Free OCR output is plain text already; Markdown/grounding output
        # can contain <|ref|>...<|/ref|> and <|det|>[[...]]<|/det|> tags —
        # strip those so the box shows clean readable text either way.
        text = re.sub(r"<\|ref\|>.*?<\|/ref\|>", "", text, flags=re.DOTALL)
        text = re.sub(r"<\|det\|>\[\[.*?\]\]<\|/det\|>", "", text, flags=re.DOTALL)
        text = text.strip()

        return text if text else "No text was detected in this image."
    finally:
        shutil.rmtree(tmp_dir, ignore_errors=True)


with gr.Blocks(title="DeepSeek-OCR (CPU)") as demo:
    gr.Markdown(
        """
        # DeepSeek-OCR — Free OCR (CPU)
        Upload an image (a photo, screenshot, scanned page, etc.) and extract
        its text using [deepseek-ai/DeepSeek-OCR](https://huggingface.co/deepseek-ai/DeepSeek-OCR).

        ⚠️ Running on **CPU** — this is noticeably slower than GPU. Start with
        the "Tiny" or "Small" resolution preset if you want faster results.
        """
    )

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload image")
            task_dropdown = gr.Dropdown(
                choices=list(TASK_PROMPTS.keys()),
                value="Free OCR (plain text)",
                label="Task",
            )
            size_dropdown = gr.Dropdown(
                choices=list(SIZE_PRESETS.keys()),
                value="Gundam (1024/640 crop, best for documents)",
                label="Resolution mode (lower = faster on CPU)",
            )
            run_button = gr.Button("Run OCR", variant="primary")

        with gr.Column():
            output_text = gr.Textbox(
                label="Extracted text", lines=25
            )

    run_button.click(
        fn=run_ocr,
        inputs=[image_input, task_dropdown, size_dropdown],
        outputs=output_text,
    )

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