Universal-DeepSeek-OCR-2

Universal DeepSeek-OCR 2 – CPU, MPS, CUDA Support

This repository uses the weights from the original DeepSeek-OCR 2 and modifies model to support inference on different devices such as CPU and MPS (Apple Metal GPU). By default runs on CPU.

• Link to the original DeepSeek-OCR-2 model

• Link to the original DeepSeek-OCR-2 repository

<div align="center"> <img src="https://github.com/deepseek-ai/DeepSeek-V2/blob/main/figures/logo.svg?raw=true" width="60%" alt="DeepSeek AI" /> </div> <p align="center"> <img src="assets/fig1.png" style="width: 900px" align=center> </p> <p align="center"> <a href="">Explore more human-like visual encoding.</a> </p>

Usage

Sample code available at: https://github.com/Dogacel/Universal-DeepSeek-OCR-2

mamba create -n deepseek-ocr-2 python=3.12.9
mamba activate deepseek-ocr-2

pip install torch==2.6.0  torchvision Pillow transformers==4.46.3 tokenizers==0.20.3 einops addict easydict

from transformers import AutoModel, AutoTokenizer
import torch

model_name = '.'

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True, use_safetensors=True)
model = model.eval().to("cpu").to(torch.float16)

# prompt = "<image>\nFree OCR. "
prompt = "<image>\n<|grounding|>Convert the document to markdown. "
image_file = 'samples/paper.png'
output_path = 'tmp'

res = model.infer(
    tokenizer, 
    prompt=prompt, 
    image_file=image_file, 
    output_path = output_path, 
    base_size = 1024, 
    image_size = 768, 
    crop_mode = True, 
    save_results = True,
    test_compress = True,
)

To change device type, you should update two things,

- model.eval().to("cpu").to(torch.float16)
+ model = model.eval().to("mps").to(torch.float16)

res = model.infer(
    tokenizer, 
    prompt=prompt, 
    image_file=image_file, 
    output_path = output_path, 
    base_size = 1024, 
    image_size = 768, 
    crop_mode = True, 
    save_results = True,
    test_compress = True,
+   device = "mps",
+   dtype = torch.float16,
)

For CUDA, you should also use bfloat16 to get as close as possible to the original implementation.

- model.eval().to("cpu").to(torch.float16)
+ model = model.eval().to("cuda").to(torch.bfloat16)

res = model.infer(
    tokenizer, 
    prompt=prompt, 
    image_file=image_file, 
    output_path = output_path, 
    base_size = 1024, 
    image_size = 768, 
    crop_mode = True, 
    save_results = True,
    test_compress = True,
+   device = "cuda",
+   dtype = torch.bfloat16,
)