Unlimited-OCR 百度开源的一个OCR（光学字符识别）模型

Unlimited-OCR 是百度开源的高性能OCR模型，专注于一次性长时程文档解析。本教程将指导你在不同的环境中部署和使用该模型。

1. 部署准备

硬件要求

• GPU：推荐使用NVIDIA GPU（如V100、A100、H100等），至少需要 16GB显存。
• 内存：建议 32GB以上。
• 存储：模型文件约 7-10GB，需预留足够空间。

软件环境

• 操作系统：Linux（Ubuntu 20.04+）或 Windows WSL2。
• Python版本：3.12.3（项目测试版本）。
• CUDA版本：12.9（如使用vLLM的Hopper GPU专用镜像）或 CUDA 13.0（默认）。
• NVIDIA驱动：确保驱动版本与CUDA兼容。

主要依赖

torch==2.10.0
torchvision==0.25.0
transformers==4.57.1
Pillow==12.1.1
pymupdf==1.27.2.2  # 用于PDF处理
einops==0.8.2
addict==2.4.0
easydict==1.13
psutil==7.2.2
matplotlib==3.10.8

2. 部署方式一：使用 Transformers（适用于快速测试和研究）

这是最直接的部署方式，适合开发测试或小规模使用。

步骤1：创建虚拟环境并安装依赖

# 创建并激活虚拟环境
python3.12 -m venv ocr-env
source ocr-env/bin/activate

# 安装PyTorch（CUDA版本根据你的驱动选择）
pip install torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu129

# 安装其他依赖
pip install transformers==4.57.1 Pillow==12.1.1 pymupdf==1.27.2.2 \
    einops==0.8.2 addict==2.4.0 easydict==1.13 psutil==7.2.2 matplotlib==3.10.8

步骤2：下载模型并运行推理

创建一个Python脚本（如 ocr_demo.py）：

import os
import torch
from transformers import AutoModel, AutoTokenizer
import tempfile
import fitz  # PyMuPDF

# 1. 加载模型
model_name = 'baidu/Unlimited-OCR'
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(
    model_name,
    trust_remote_code=True,
    use_safetensors=True,
    torch_dtype=torch.bfloat16,
)
model = model.eval().cuda()

# 2. PDF转图像辅助函数
def pdf_to_images(pdf_path, dpi=300):
    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

# 3. 单张图片识别（使用gundam模式）
model.infer(
    tokenizer,
    prompt='<image>document parsing.',
    image_file='your_image.jpg',
    output_path='./output',
    base_size=1024,
    image_size=640,
    crop_mode=True,  # gundam模式
    max_length=32768,
    no_repeat_ngram_size=35,
    ngram_window=128,
    save_results=True,
)

# 4. 多页PDF识别（使用base模式）
pdf_images = pdf_to_images('your_document.pdf', dpi=300)
model.infer_multi(
    tokenizer,
    prompt='<image>Multi page parsing.',
    image_files=pdf_images,
    output_path='./output_pdf',
    image_size=1024,  # base模式
    max_length=32768,
    no_repeat_ngram_size=35,
    ngram_window=1024,
    save_results=True,
)

两种模式的选择：

• gundam模式：base_size=1024, image_size=640, crop_mode=True，适用于单张图像，可能通过裁剪策略优化处理。
• base模式：base_size=1024, image_size=1024, crop_mode=False，适用于多页文档或PDF，保证全图解析。

3. 部署方式二：使用 vLLM（适用于高性能生产环境）

vLLM 提供高效的推理服务，适合大规模部署。

使用预构建的Docker镜像（最便捷）

# 拉取镜像（根据GPU选择）
# 默认（CUDA 13.0）
docker pull vllm/vllm-openai:unlimited-ocr

# 或 Hopper GPU（CUDA 12.9）
docker pull vllm/vllm-openai:unlimited-ocr-cu129

# 启动服务
docker run --gpus all -p 8000:8000 \
    vllm/vllm-openai:unlimited-ocr \
    --model baidu/Unlimited-OCR

手动部署vLLM服务

# 安装vLLM（参考官方文档）
pip install vllm

# 启动OpenAI兼容的API服务
python -m vllm.entrypoints.openai.api_server \
    --model baidu/Unlimited-OCR \
    --served-model-name Unlimited-OCR \
    --trust-remote-code \
    --max-model-len 32768 \
    --gpu-memory-utilization 0.9

服务启动后，可以通过 http://localhost:8000/v1 访问OpenAI兼容的API。

4. 部署方式三：使用 SGLang（适用于高级批处理）

SGLang 提供了灵活的批处理和并发控制。

步骤1：安装SGLang和相关依赖

# 创建uv管理的虚拟环境（推荐）
uv venv --python 3.12
source .venv/bin/activate

# 安装SGLang wheel（需要从项目仓库获取）
uv pip install wheel/sglang-0.0.0.dev11416+g92e8bb79e-py3-none-any.whl
uv pip install kernels==0.11.7 pymupdf==1.27.2.2

步骤2：启动SGLang服务器

python -m sglang.launch_server \
    --model baidu/Unlimited-OCR \
    --served-model-name Unlimited-OCR \
    --attention-backend fa3 \
    --page-size 1 \
    --mem-fraction-static 0.8 \
    --context-length 32768 \
    --enable-custom-logit-processor \
    --disable-overlap-schedule \
    --skip-server-warmup \
    --host 0.0.0.0 \
    --port 10000

步骤3：使用批量推理脚本

项目提供了 infer.py 脚本进行批处理：

# 处理图像目录
python infer.py \
    --image_dir ./examples/images \
    --output_dir ./outputs \
    --concurrency 8 \
    --image_mode gundam

# 处理PDF
python infer.py \
    --pdf ./examples/document.pdf \
    --output_dir ./outputs \
    --concurrency 8 \
    --image_mode gundam

# 可选参数
# --model_dir baidu/Unlimited-OCR  # 本地或Hugging Face模型ID
# --gpu 0                           # 指定GPU设备
# --server_log ./log/sglang_server.log

5. 结果后处理

模型输出可能包含 <|det|> 标记，需要清理以获得纯文本：

import re

DET_RE = re.compile(r'<\|det\|>([^<\s]+)(?:\s*\[[^\]]*\])?\s*<\|/det\|>(.*)', re.DOTALL)

def remove_det(raw: str) -> str:
    """移除<|det|>标记，按段落分组"""
    blocks = []
    cur = None
    for line in raw.splitlines():
        line = line.rstrip()
        if not line:
            continue
        m = DET_RE.match(line)
        if m:
            category, content = m.group(1).strip(), m.group(2).strip()
            if category == 'image':
                continue
            if cur is not None:
                blocks.append(cur)
            cur = [content] if content else []
            continue
        if cur is None:
            cur = []
        cur.append(line)
    if cur is not None:
        blocks.append(cur)
    return '\n\n'.join('\n'.join(b) for b in blocks).strip()

6. 常见问题与性能优化

显存不足

• 降低 --context-length 或 max_length 参数值。
• 使用 --gpu-memory-utilization 0.7（vLLM）或 --mem-fraction-static 0.7（SGLang）降低显存占用。
• 考虑使用 torch.float16 替代 bfloat16。

处理速度慢

• 使用 vLLM 或 SGLang 部署，它们针对推理进行了优化。
• 调整 batch_size 或 concurrency 参数。
• 确保使用支持 fa3 注意力后端的GPU（如Hopper架构）。

PDF解析问题

• 调整 dpi 参数（默认300），高DPI可能增加处理时间但提高识别精度。
• 对于扫描件PDF，确保图像质量足够。

模型下载慢

• 从 Hugging Face 或 ModelScope 镜像下载，设置 HF_ENDPOINT 环境变量：

  export HF_ENDPOINT=https://hf-mirror.com

总结

Unlimited-OCR 提供了三种主要的部署方式，满足不同场景需求：

• Transformers方式：适合快速测试和研究，易于集成到Python项目中。
• vLLM方式：适合生产环境，提供高性能API服务。
• SGLang方式：适合大规模批处理，支持灵活并发控制。

根据你的硬件资源、处理需求和熟悉程度，选择最适合的部署路径。对于初次尝试，建议从Transformers方式开始，验证模型在特定文档上的效果。对于需要稳定API服务的生产环境，vLLM是更好的选择。