Xiaomi-OCR-0

# Parse one document image to Markdown (adapted from the Xiaomi-OCR-0 model card).
# pip install "transformers==5.17.0" "torch==2.14.0" "torchvision==0.29.0" pillow
import torch
from transformers import AutoModelForImageTextToText, AutoProcessor

MODEL = "xiaomi-research/Xiaomi-OCR-0"
PROMPT = (
    "Extract all information from the main body of the document image and represent it "
    "in markdown format, ignoring headers and footers. Tables should be expressed in OTSL "
    "format, formulas in the document should be represented using LATEX format, and the "
    "parsing should be organized according to the reading order."
)

processor = AutoProcessor.from_pretrained(MODEL)
model = AutoModelForImageTextToText.from_pretrained(MODEL, dtype=torch.float16).to("cuda").eval()

messages = [{"role": "user", "content": [
    {"type": "image", "url": "page.png"},
    {"type": "text", "text": PROMPT},
]}]
inputs = processor.apply_chat_template(
    messages, tokenize=True, add_generation_prompt=True,
    return_dict=True, return_tensors="pt",
).to("cuda")

with torch.inference_mode():
    output = model.generate(**inputs, max_new_tokens=4096, do_sample=False)
print(processor.batch_decode(output[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0])