Dogacel/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.
<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 easydictfrom 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,
)