from PIL import Image
import torch
from transformers import AutoModelForCausalLM, AutoProcessor
model_path = "Kimang18/paddleOCR_vl_Khmer_finetuned"
image_path = "khmer_text_region.png"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
image = Image.open(image_path).convert("RGB")
model = AutoModelForCausalLM.from_pretrained(
model_path, trust_remote_code=True, dtype=torch.bfloat16
).to(DEVICE).eval()
processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
messages = [
{"role": "user",
"content": [
{"type": "image", "image": image},
{"type": "text", "text": "OCR:"},
]
}
]
inputs = processor.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt"
).to(DEVICE)
outputs = model.generate(**inputs, max_new_tokens=1024)
outputs = processor.batch_decode(outputs, skip_special_tokens=True)[0]
print(outputs)