from transformers import AutoModel, AutoTokenizer
from peft import PeftModel
import torch

base = "deepseek-ai/DeepSeek-OCR"
lora = "giayphuyen/lora_model"

tokenizer = AutoTokenizer.from_pretrained(base, trust_remote_code=True)
model = AutoModel.from_pretrained(base, trust_remote_code=True).cuda()

model = PeftModel.from_pretrained(model, lora)
model = model.eval().to(torch.bfloat16)

prompt = "<image>\nExtract the text."
image = "your_image.jpg"

res = model.infer(
    tokenizer,
    prompt=prompt,
    image_file=image,
    output_path=".",
)
print(res)