app.py

import subprocess, sys, os, tempfile
from threading import Thread
from typing import Iterator, Tuple
import queue
import threading

_RUNTIME_PKGS = [
    "torch==2.10.0",
    "torchvision==0.25.0",
    "transformers==4.57.1",
]

print("Installing pinned runtime dependencies...")
subprocess.run(
    [sys.executable, "-m", "pip", "install", "--quiet", "--no-cache-dir"] + _RUNTIME_PKGS,
    check=True,
)
print("Runtime deps installed.")

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

MODEL_NAME = "baidu/Unlimited-OCR"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
print("Loading model (CPU)...")
model = AutoModel.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    use_safetensors=True,
    torch_dtype=torch.float32,
).eval()
print("Model ready (CPU).")


def pdf_to_images(pdf_path: str, dpi: int = 200) -> list[str]:
    import fitz
    doc = fitz.open(pdf_path)
    tmp_dir = tempfile.mkdtemp(prefix="pdf_ocr_")
    mat = fitz.Matrix(dpi / 72, dpi / 72)
    paths = []
    for i, page in enumerate(doc):
        out = os.path.join(tmp_dir, f"page_{i + 1:04d}.png")
        page.get_pixmap(matrix=mat).save(out)
        paths.append(out)
    doc.close()
    return paths


def _collect_output(out_dir: str) -> str:
    result = ""
    for fname in sorted(os.listdir(out_dir)):
        if fname.endswith((".txt", ".md")):
            with open(os.path.join(out_dir, fname), "r", encoding="utf-8") as f:
                result += f.read() + "\n"
    if not result:
        for fname in sorted(os.listdir(out_dir)):
            fpath = os.path.join(out_dir, fname)
            if os.path.isfile(fpath):
                try:
                    with open(fpath, "r", encoding="utf-8") as f:
                        result += f.read() + "\n"
                except Exception:
                    pass
    return result.strip()


class ThreadTargetedStdout:
    def __init__(self, target_thread, q, original_stdout):
        self.target_thread = target_thread
        self.q = q
        self.original_stdout = original_stdout

    def write(self, data):
        self.original_stdout.write(data)
        self.original_stdout.flush()
        if threading.current_thread() == self.target_thread:
            if data:
                lower_data = data.lower()
                if "tps:" in lower_data or "tokens/s" in lower_data:
                    return len(data)
                self.q.put(data)
        return len(data)

    def flush(self):
        self.original_stdout.flush()

    def __getattr__(self, name):
        return getattr(self.original_stdout, name)


def _run_ocr_internal(image_path: str, mode: str, prompt: str) -> str:
    out_dir = tempfile.mkdtemp(prefix="ocr_out_")

    if mode == "gundam":
        base_size, image_size, crop_mode, ngram_window = 1024, 640, True, 128
    else:
        base_size, image_size, crop_mode, ngram_window = 1024, 1024, False, 128

    _infer_kwargs = dict(
        prompt=f"<image>{prompt}",
        image_file=image_path,
        output_path=out_dir,
        base_size=base_size,
        image_size=image_size,
        crop_mode=crop_mode,
        max_length=8192,
        no_repeat_ngram_size=35,
        ngram_window=ngram_window,
        save_results=True,
    )

    q = queue.Queue()
    errors = []

    def _infer_thread():
        try:
            model.infer(tokenizer, **_infer_kwargs)
        except Exception as e:
            errors.append(str(e))

    thread = Thread(target=_infer_thread, daemon=True)
    original_stdout = sys.stdout
    targeted_stdout = ThreadTargetedStdout(thread, q, original_stdout)
    sys.stdout = targeted_stdout

    accumulated = ""
    try:
        thread.start()
        while thread.is_alive() or not q.empty():
            try:
                chunk = q.get(timeout=0.02)
                accumulated += chunk
            except queue.Empty:
                continue
    finally:
        sys.stdout = original_stdout
        thread.join()

    full_text = _collect_output(out_dir)

    if full_text:
        return full_text
    elif accumulated:
        return accumulated
    elif errors:
        raise RuntimeError(f"OCR failed: {', '.join(errors)}")
    return ""


def run_ocr(
    image: gr.Image,
    mode: str = "gundam",
    prompt: str = "document parsing.",
) -> str:
    """
    Run OCR on a single image. Returns extracted text.

    Args:
        image: The image to run OCR on.
        mode: 'gundam' for fast mode, 'base' for accurate mode.
        prompt: Instruction for the OCR model.
    """
    if image is None:
        return "Error: No image provided."

    tmp_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
    image.save(tmp_path)
    result = _run_ocr_internal(tmp_path, mode, prompt)
    os.unlink(tmp_path)
    return result


def run_ocr_pdf(
    pdf_file: gr.File,
    mode: str = "gundam",
    prompt: str = "document parsing.",
) -> str:
    """
    Run OCR on a PDF document. Converts each page to an image and runs OCR.
    Returns concatenated text from all pages.

    Args:
        pdf_file: The PDF file to run OCR on.
        mode: 'gundam' for fast mode, 'base' for accurate mode.
        prompt: Instruction for the OCR model.
    """
    if pdf_file is None:
        return "Error: No PDF file provided."

    pages = pdf_to_images(pdf_file, dpi=200)
    all_text = []
    for i, page_path in enumerate(pages):
        text = _run_ocr_internal(page_path, mode, prompt)
        all_text.append(f"--- Page {i + 1} ---\n{text}")
        os.unlink(page_path)
    return "\n\n".join(all_text)


with gr.Blocks(title="Unlimited OCR MCP") as demo:
    gr.Markdown("# Unlimited OCR MCP Server")
    gr.Markdown("Baidu's Unlimited-OCR model as an MCP tool. Send images or PDFs to extract text.")

    with gr.Tab("Image OCR"):
        image_input = gr.Image(label="Image", type="pil")
        mode_input = gr.Radio(["gundam", "base"], value="gundam", label="Mode")
        prompt_input = gr.Textbox(value="document parsing.", label="Prompt")
        image_btn = gr.Button("Run OCR")
        image_output = gr.Textbox(label="OCR Result", lines=20)
        image_btn.click(fn=run_ocr, inputs=[image_input, mode_input, prompt_input], outputs=image_output)

    with gr.Tab("PDF OCR"):
        pdf_input = gr.File(label="PDF File", file_types=[".pdf"])
        pdf_mode = gr.Radio(["gundam", "base"], value="gundam", label="Mode")
        pdf_prompt = gr.Textbox(value="document parsing.", label="Prompt")
        pdf_btn = gr.Button("Run OCR")
        pdf_output = gr.Textbox(label="OCR Result", lines=20)
        pdf_btn.click(fn=run_ocr_pdf, inputs=[pdf_input, pdf_mode, pdf_prompt], outputs=pdf_output)


demo.launch(mcp_server=True)