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"{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)