Unlimited-OCR MLX

🚀 Unlimited-length document OCR model accelerated by Apple MLX framework, deeply optimized for Apple Silicon.

📖 Model Overview

Unlimited-OCR MLX is a high-precision OCR solution that fully migrates the Baidu PaddlePaddle team's Unlimited-OCR model to the Apple MLX framework.

Based on the DeepSeek-V2 architecture, combined with SAM-ViT-B + CLIP-L dual vision encoders, it can parse documents of any length in a single pass, implementing end-to-end text recognition and structured extraction.

✨ Core Features

Fixes to the MLX port of Unlimited-OCR-MLX

The original MLX port of LoJexLLM/Unlimited-OCR-MLX (as published on HuggingFace) does not produce correct OCR output out of the box. Several bugs in the vendored model code caused garbled or empty output. The fixes below were applied independently to both model.py (the library module used by loader.py / inference.py) and the model code embedded directly in pdf_to_markdown.py (the standalone script actually used for the OCR pipeline, which does not import model.py).

Correctness fixes

1. RMSNorm used the wrong normalization formula (root cause)

The ported code called:

mx.fast.rms_norm(x, 1.0 + self.weight, eps)

This is the Gemma-style RMSNorm convention, where the learned weight is centered at 0 and the model adds 1.0 before scaling. This model's checkpoint, however, stores weights centered around their final scale (mean roughly 0.24–0.67), matching the standard LLaMA/DeepSeek convention. Adding 1.0 + on top of already-scaled weights corrupted every hidden state passing through every norm layer, and the language model produced pure garbage output as a result.

Fixed to the standard form:

mx.fast.rms_norm(x, self.weight, self.eps)

with self.weight initialized to mx.ones((dims,)).

2. SAM vision encoder: relative position bias and position embedding interpolation were stubs

_get_rel_pos in the original port unconditionally returned a zero tensor — the windowed and global attention blocks in the SAM (Segment Anything) image encoder received no relative position signal at all. Likewise, the position-embedding resize routine was a fake reshape rather than an actual interpolation.

Rewrote both to match the standard SAM ImageEncoderViT:

• _get_rel_pos(q_size, k_size, rel_pos): resizes the relative position table to 2*max(q_size, k_size) - 1 via linear interpolation, then gathers by real relative coordinates (_interp_linear + coordinate math), instead of returning zeros.
• Position embedding interpolation (_interpolate_pos_embed / inline helper): does an actual row-then-column bilinear interpolation (align_corners=False convention) instead of a reshape.
• Relative position bias tables are sized correctly for each attention type: window attention uses window_size=14 → table size 2*14-1=27; global attention uses the full grid 64 (img_size=1024 / patch_size=16) → table size 2*64-1=127.

Without this fix the SAM encoder still ran without crashing, but visual features carried no positional information, degrading layout/table/image region recognition.

3. Image feature injection into the language model used a nonexistent API

The original code attempted:

inputs_embeds.at[idx].set(image_features)

mx.array has no .at[idx].set() method in MLX (this is a JAX idiom), so this either raised an error or silently did nothing depending on the MLX version, meaning image features never reached the language model.

Fixed by locating the first True position in the per-sequence image mask (the image-token block is contiguous) and rebuilding the embedding sequence via slicing and concatenation:

start = int(mx.argmax(mask.astype(mx.int32)).item())
n = img_feats.shape[0]
emb = mx.concatenate([emb[:start], img_feats, emb[start + n:]], axis=0)

n (the number of image feature rows) must equal the number of image placeholder positions reserved in the prompt — see the prompt-format note below.

Loading / integration fixes (outside the vendored model code)

These do not modify the ported architecture code itself, but were required to load the released checkpoint and drive the model correctly.

4. Weight name mismatch for the SAM neck

The checkpoint stores the SAM neck's nn.Sequential layers as sam_model.neck.<N>.<weight|bias>, while the MLX module tree expects sam_model.neck.layers.<N>.<weight|bias>. Added a regex remap step during loading:

m = re.match(r"(sam_model\.neck)\.(\d+)\.(weight|bias)$", key)
if m:
    key = f"{m.group(1)}.layers.{m.group(2)}.{m.group(3)}"

5. Conv2d weight layout mismatch (NCHW vs NHWC)

PyTorch stores 4D conv weights as [out, in, kh, kw] (NCHW-style); MLX convs expect [out, kh, kw, in] (NHWC-style). Added a transpose for any 4D tensor whose shape doesn't match the expected MLX shape:

if arr.ndim == 4 and arr.shape != exp:
    arr = np.transpose(arr, (0, 2, 3, 1))

After both remaps, all 2710 checkpoint tensors load into the model with zero dropped/unmatched keys.

6. Tokenizer loading crashed with AutoTokenizer

config.json's auto_map entry breaks AutoTokenizer.from_pretrained(...) on current transformers versions. Switched to loading the tokenizer file directly:

PreTrainedTokenizerFast(
    tokenizer_file="tokenizer.json",
    bos_token="<｜begin▁of▁sentence｜>",
    eos_token="<｜end▁of▁sentence｜>",
)

7. Prompt format — plain-text chat templates do not work

Using a conversational format like "User: ...\nAssistant:" causes the model to emit EOS immediately. The model instead expects a raw sequence of: [BOS] + [<image> placeholder] * N + encoded("\n" + prompt), with the per-position sequence mask (images_seq_mask) set to True exactly over the image placeholder block. N must equal the number of feature rows actually returned by encode_images(...) for that image — the image must be encoded before the input sequence is constructed, not the other way around.

Performance optimizations (pdf_to_markdown.py only)

Not correctness fixes, but changes required to make batch OCR of a 500+ page scanned book practical (~40–48 tokens/sec, up from an unusably slow baseline):

• Keep loaded weights in fp16. The original loader implicitly upcast weights to fp32, doubling model memory to ~12.4 GB and making decoding memory-bandwidth-bound. Weights are now kept in fp16 end-to-end; only transient math (softmax, RoPE, logit scaling) uses fp32. This was the single largest win — dense pages dropped from ~55s to ~4.4s.
• Vectorized the MoE decode path for short sequences. For T <= 32 (i.e. decoding one token at a time), expert routing is now done via mx.take/gather + einsum over stacked expert weights, avoiding the original per-token np.bincount (which forced a CPU sync every step) and a 64-iteration Python loop over experts. For large T (prefill), a separate grouped-by-expert loop path is retained to avoid holding all experts' intermediate activations in memory at once.
• Pipelined the generation loop. The decode loop issues mx.async_eval for the next token before checking the previous token for EOS, so the GPU is never left idle waiting on a CPU-side stop condition.

🏗️ Model Architecture

Input Image
    │
    ├──→ SAM-ViT-B (ViT-Base, 12 layers, 768 dims)
    │       │
    │       └──→ CLIP-L ViT (24 layers, 1024 dims)
    │                │
    │                └──→ Feature Concatenation [2048 dims]
    │                         │
    │                         └──→ Projection Layer Linear(2048→1280)
    │                                  │
    │                                  └──→ Image Feature Embedding
    │
    └──→ Text Tokens → Embedding
              │
              └──→ DeepSeek-V2 MoE Language Model (12 layers)
                        │
                        ├── Layer 0: Dense MLP (SwiGLU, 6848 dims)
                        ├── Layer 1-11: Mixture of Experts (64 Experts, Top-6 Routing)
                        └── Standard Multi-Head Attention + RoPE Positional Encoding
                              │
                              └──→ OCR Text Output

Core Specifications

🔧 Quick Start

Requirements

• macOS 14.0+ (Apple Silicon M1/M2/M3/M4)
• Python 3.10+
• MLX >= 0.20.0

Installation

pip install mlx mlx-lm safetensors transformers Pillow numpy

Model Download

# Download from Hugging Face
git lfs install
git clone https://huggingface.co/LoJexLLM/Unlimited-OCR-MLX

Python API

from unlimited_ocr_mlx import UnlimitedOCRInference

# Initialize engine
engine = UnlimitedOCRInference("./Unlimited-OCR-MLX")
engine.load()

# Single image OCR (high-precision dynamic tiling mode)
result = engine.infer_single(
    image_path="document.jpg",
    prompt="document parsing.",
    crop_mode=True,        # Enable dynamic tiling
    base_size=1024,        # Global view size
    image_size=640,        # Tile size
    max_length=32768,      # Max generation length
    temperature=0.0,       # Greedy decoding (high precision)
)

print(result)

Command Line

python -m unlimited_ocr_mlx.inference \
    --model_dir ./Unlimited-OCR-MLX \
    --image document.jpg \
    --prompt "document parsing." \
    --output ./ocr_results \
    --crop_mode \
    --base_size 1024 \
    --image_size 640

⚡ Performance Comparison

Measured performance on Apple M4 Pro (compared to original PyTorch MPS):

MLX fully leverages Apple Silicon's unified memory architecture and GPU/Neural Engine co-processing, delivering significant acceleration compared to the PyTorch MPS backend.

🎯 Inference Modes

1. Gundam Mode (High Precision)

• crop_mode=True, image_size=640
• Dynamic tiling + global view
• Suitable for high-precision document parsing

2. Base Mode (Fast)

• crop_mode=False, image_size=1024
• Single-scale global encoding
• Suitable for quick scanning of simple documents

📊 Precision Verification

The MLX version has undergone rigorous precision verification (256 random inputs, BF16→FP16 conversion):

• Cosine Similarity: > 0.999 (vs PyTorch original model)
• Token Match Rate: > 99.5% (same input, same output)
• Visual Feature Consistency: Structural Similarity (SSIM) > 0.998

📁 Model Files

Unlimited-OCR-MLX/
├── model.safetensors          # MLX weights file (FP16, ~6.2 GB)
├── config.json                # Model configuration
├── tokenizer.json             # Tokenizer
├── tokenizer_config.json      # Tokenizer config
├── special_tokens_map.json    # Special token mapping
├── unlimited_ocr_mlx/         # MLX implementation code
│   ├── model.py               #   Complete model definition
│   ├── config.py              #   Configuration management
│   ├── convert.py             #   Weight conversion tool
│   ├── inference.py           #   Inference pipeline
│   ├── image_processing.py    #   Image preprocessing
│   ├── loader.py              #   Weight loader
│   └── test_validation.py     #   Precision validation
├── README.md                  # This document
└── LICENSE                    # MIT License

🙏 Acknowledgements

• Original model: PaddlePaddle/Unlimited-OCR
• Baidu PaddlePaddle Team
• DeepSeek-OCR — Base architecture
• Apple MLX — Inference acceleration framework

📄 Citation

@misc{unlimited-ocr-mlx,
  title={Unlimited-OCR MLX: High-Precision OCR on Apple Silicon},
  author={PaddlePaddle MLX Community},
  year={2026},
  url={https://huggingface.co/LoJexLLM/Unlimited-OCR-MLX}
}

📜 License

This project is open source under the MIT License. Original model copyright belongs to the Baidu PaddlePaddle team.