inference.py

"""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()