File size: 10,727 Bytes
267d7ad | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 | """Unlimited-OCR MLX Inference Pipeline.
Complete inference pipeline for document OCR using MLX acceleration on Apple Silicon.
Usage:
python inference.py --model_dir ./unlimited-ocr-mlx-weights --image document.jpg --output ./output
"""
import os
import sys
import json
import time
import argparse
from typing import Optional, List
import numpy as np
import mlx.core as mx
from .config import UnlimitedOCRConfig
from .model import UnlimitedOCRModel
from .image_processing import load_image, preprocess_image, build_input
def load_tokenizer(model_dir: str):
"""Load tokenizer files from the model directory."""
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(
model_dir,
trust_remote_code=True,
use_fast=False,
)
return tokenizer
def load_model(model_dir: str) -> UnlimitedOCRModel:
"""Load the MLX model with converted weights."""
# Load config
config_path = os.path.join(model_dir, "config.json")
with open(config_path) as f:
config_dict = json.load(f)
config = UnlimitedOCRConfig.from_original_config(config_dict)
# Load weights
weights_path = os.path.join(model_dir, "model.safetensors")
if not os.path.exists(weights_path):
raise FileNotFoundError(
f"MLX weights not found at {weights_path}. "
"Run convert.py first to convert from PyTorch."
)
import safetensors.torch
st_weights = safetensors.torch.load_file(weights_path, device="cpu")
weights = {}
for k, v in st_weights.items():
weights[k] = mx.array(v.float().numpy())
# Create model and load weights
model = UnlimitedOCRModel(config)
model.load_weights(list(weights.items()))
mx.eval(model.parameters())
print(f"Model loaded with {sum(v.size for v in weights.values()):,} parameters")
return model
def create_attention_mask(seq_len: int) -> mx.array:
"""Create causal attention mask."""
mask = mx.tril(mx.ones((seq_len, seq_len), dtype=mx.bool_))
mask = mx.where(mask, 0.0, float('-inf'))
return mask[None, None, :, :]
def format_conversation(prompt: str, image_path: str) -> List[dict]:
"""Format conversation for the model."""
return [
{
"role": "User",
"content": f"<image_placeholder>\n{prompt}",
"images": [image_path],
},
{"role": "Assistant", "content": ""},
]
class UnlimitedOCRInference:
"""High-level inference interface for Unlimited-OCR MLX."""
def __init__(self, model_dir: str):
self.model_dir = model_dir
self.model = None
self.tokenizer = None
def load(self):
"""Load model and tokenizer."""
print("Loading model...")
self.model = load_model(self.model_dir)
print("Loading tokenizer...")
self.tokenizer = load_tokenizer(self.model_dir)
print("Ready!")
return self
def encode_text(self, text: str, bos: bool = True) -> List[int]:
"""Encode text to token IDs."""
tokens = self.tokenizer.encode(text, add_special_tokens=False)
if bos:
tokens = [self.tokenizer.bos_token_id] + tokens
return tokens
def decode_text(self, token_ids: List[int]) -> str:
"""Decode token IDs to text."""
return self.tokenizer.decode(token_ids, skip_special_tokens=True)
def process_image(self, image_path: str):
"""Load and preprocess an image."""
image = load_image(image_path)
if image is None:
raise ValueError(f"Cannot load image: {image_path}")
return preprocess_image(
image,
base_size=1024,
image_size=640,
crop_mode=True,
)
def infer_single(
self,
image_path: str,
prompt: str = "document parsing.",
output_dir: Optional[str] = None,
max_length: int = 32768,
temperature: float = 0.0,
base_size: int = 1024,
image_size: int = 640,
crop_mode: bool = True,
) -> str:
"""Run OCR inference on a single image.
Args:
image_path: Path to the input image
prompt: OCR prompt
output_dir: Output directory for results
max_length: Maximum generation length
temperature: Sampling temperature (0 = greedy)
base_size: Base image size for global view
image_size: Tile size for patches
crop_mode: Whether to use dynamic tiling
Returns:
Generated OCR text
"""
if self.model is None:
self.load()
# Create output directory
if output_dir:
os.makedirs(output_dir, exist_ok=True)
os.makedirs(os.path.join(output_dir, "images"), exist_ok=True)
# Format conversation
conversation = [
{"role": "User", "content": f"<image_placeholder>\n{prompt}"},
{"role": "Assistant", "content": ""},
]
# Build text prompt
from .image_processing import load_image as _load
text_parts = []
for msg in conversation:
role = msg["role"]
content = msg["content"]
if role == "User":
text_parts.append(f"User: {content}")
elif role == "Assistant":
text_parts.append(f"Assistant: {content}")
full_prompt = "\n".join(text_parts)
prompt_tokens = self.encode_text(full_prompt)
# Process image
image = _load(image_path)
if image is None:
raise ValueError(f"Cannot load image: {image_path}")
patches_arr, orig_arr, crop_shape = preprocess_image(
image, base_size=base_size, image_size=image_size, crop_mode=crop_mode
)
# Convert to MLX arrays
patches_mx = mx.array(patches_arr) if patches_arr.shape[0] > 0 else None
orig_mx = mx.array(orig_arr)
# Compute number of image tokens from vision encoder output shape
# SAM: 1024 → 64x64 → 16x16 spatial after net_3
# CLIP+concat: → 256 spatial tokens, 2048 dim
# Projector: → 256 spatial tokens, 1280 dim
# With crop_mode: local (grid of 256 each) + global (256) + separators
if crop_mode and patches_arr.shape[0] > 0:
w_crop, h_crop = crop_shape
n_local_tokens = w_crop * h_crop * 272 # 256 + newline=16 tokens, roughly
n_global_tokens = 272 # 256 + 16 newlines + separator
n_image_tokens = n_local_tokens + n_global_tokens
else:
n_image_tokens = 272 # 256 + 16 newlines + separator
# Build input with image token masks
# The model replaces token 0 with image features
# We need to compute images_seq_mask properly
# For simplicity: insert image tokens at the start after BOS
input_ids = prompt_tokens.copy()
total_image_feats = n_image_tokens
# Mask: True where image features should be placed
# After first BOS token, insert image features
seq_mask = np.zeros(len(input_ids) + total_image_feats, dtype=bool)
# Mark image positions (right after the first <image_placeholder> token)
image_start = 1 # After BOS
seq_mask[image_start:image_start + total_image_feats] = True
# Extend input_ids with placeholder positions
extended_ids = input_ids[:1] + [0] * total_image_feats + input_ids[1:]
print(f"Input: {len(extended_ids)} tokens, {total_image_feats} image tokens")
print("Running OCR inference...")
start_time = time.time()
# Prepare model inputs
input_ids_mx = mx.array([extended_ids], dtype=mx.int32)
images_seq_mask_mx = mx.array([seq_mask], dtype=bool)
# Prepare image tensor in the format the model expects
# [patches, original]
image_tensor = [patches_mx, orig_mx]
images = [image_tensor]
images_spatial_crop = [crop_shape] if crop_mode else [(1, 1)]
# Generate
output_ids = self.model.generate(
input_ids=input_ids_mx,
images=images,
images_seq_mask=images_seq_mask_mx,
images_spatial_crop=images_spatial_crop,
max_length=max_length,
temperature=temperature,
eos_token_id=self.tokenizer.eos_token_id,
)
elapsed = time.time() - start_time
tokens_generated = output_ids.shape[1] - len(extended_ids)
tps = tokens_generated / elapsed if elapsed > 0 else 0
# Decode
output_tokens = output_ids[0].tolist()
result = self.decode_text(output_tokens)
print(f"\n=== OCR Result ({tokens_generated} tokens, {elapsed:.1f}s, {tps:.1f} t/s) ===")
print(result)
if output_dir:
result_path = os.path.join(output_dir, "result.txt")
with open(result_path, "w", encoding="utf-8") as f:
f.write(result)
print(f"Saved result to {result_path}")
return result
def main():
parser = argparse.ArgumentParser(description="Unlimited-OCR MLX Inference")
parser.add_argument("--model_dir", type=str, required=True,
help="Directory containing MLX weights and tokenizer")
parser.add_argument("--image", type=str, required=True,
help="Path to input image")
parser.add_argument("--prompt", type=str, default="document parsing.",
help="OCR prompt")
parser.add_argument("--output", type=str, default="./output",
help="Output directory")
parser.add_argument("--max_length", type=int, default=32768,
help="Maximum generation length")
parser.add_argument("--temperature", type=float, default=0.0,
help="Sampling temperature")
parser.add_argument("--base_size", type=int, default=1024,
help="Base image size")
parser.add_argument("--image_size", type=int, default=640,
help="Tile image size")
parser.add_argument("--no_crop", action="store_true",
help="Disable dynamic tiling (use base mode)")
args = parser.parse_args()
engine = UnlimitedOCRInference(args.model_dir)
result = engine.infer_single(
image_path=args.image,
prompt=args.prompt,
output_dir=args.output,
max_length=args.max_length,
temperature=args.temperature,
base_size=args.base_size,
image_size=args.image_size,
crop_mode=not args.no_crop,
)
if __name__ == "__main__":
main()
|