inference_processor.py

from __future__ import annotations

import json
import math
import os
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import numpy as np
from PIL import Image

import torch
import torch.nn.functional as F

from transformers import AutoTokenizer, BatchFeature, ProcessorMixin
from transformers.image_processing_utils import BaseImageProcessor
from transformers.utils import TensorType, cached_file


CONFIG_NAME = "config.json"
PREPROCESSOR_CONFIG_NAME = "preprocessor_config.json"
PROCESSOR_CONFIG_NAME = "processor_config.json"


ImageLike = Union[Image.Image, np.ndarray, torch.Tensor]


def _select_cached_file_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]:
    allowed = {
        "cache_dir",
        "force_download",
        "proxies",
        "token",
        "local_files_only",
        "revision",
        "subfolder",
    }
    out = {k: v for k, v in kwargs.items() if k in allowed}
    out.setdefault("_raise_exceptions_for_missing_entries", False)
    out.setdefault("_raise_exceptions_for_gated_repo", False)
    out.setdefault("_raise_exceptions_for_connection_errors", False)
    return out


def _resolve_repo_file(pretrained_model_name_or_path: Union[str, os.PathLike], filename: str, **kwargs) -> Optional[str]:
    path = str(pretrained_model_name_or_path)

    if os.path.isdir(path):
        candidate = os.path.join(path, filename)
        return candidate if os.path.exists(candidate) else None

    if os.path.isfile(path):
        return path if os.path.basename(path) == filename else None

    try:
        return cached_file(path, filename, **_select_cached_file_kwargs(kwargs))
    except Exception:
        return None


def _load_json_file(path: str) -> Dict[str, Any]:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def _load_image_processor_dict(pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> Dict[str, Any]:
    processor_path = _resolve_repo_file(pretrained_model_name_or_path, PROCESSOR_CONFIG_NAME, **kwargs)
    if processor_path is not None:
        processor_dict = _load_json_file(processor_path)
        nested = processor_dict.get("image_processor")
        if isinstance(nested, dict):
            return nested

    preprocessor_path = _resolve_repo_file(pretrained_model_name_or_path, PREPROCESSOR_CONFIG_NAME, **kwargs)
    if preprocessor_path is not None:
        return _load_json_file(preprocessor_path)

    config_path = _resolve_repo_file(pretrained_model_name_or_path, CONFIG_NAME, **kwargs)
    if config_path is not None:
        return _load_json_file(config_path)

    raise FileNotFoundError(
        f"Could not find {PREPROCESSOR_CONFIG_NAME}, {PROCESSOR_CONFIG_NAME}, or {CONFIG_NAME} in {pretrained_model_name_or_path!r}."
    )


class AnandaImageProcessor(BaseImageProcessor):
    """Image processor for Ananda OCR-style visual prefix inputs.

    Behavior mirrored from the development inference path:
      1. Convert to RGB / 3 channels.
      2. Convert to CHW float32 in [0, 1].
      3. Normalize with config mean/std.
      4. Pad H/W up to a multiple of patch_size.
      5. Pad again up to a multiple of patch_size * merge_factor.
      6. Emit `pixel_values` and `patch_attention_mask`.
    """

    model_input_names = ["pixel_values", "patch_attention_mask"]

    def __init__(
        self,
        patch_size: int = 16,
        merge_factor: int = 1,
        do_convert_rgb: bool = True,
        do_rescale: bool = True,
        rescale_factor: float = 1.0 / 255.0,
        do_normalize: bool = True,
        image_mean: Optional[Sequence[float]] = None,
        image_std: Optional[Sequence[float]] = None,
        pad_value: float = 0.0,
        processor_class: Optional[str] = "AnandaProcessor",
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.patch_size = int(patch_size)
        self.merge_factor = max(int(merge_factor), 1)
        self.do_convert_rgb = bool(do_convert_rgb)
        self.do_rescale = bool(do_rescale)
        self.rescale_factor = float(rescale_factor)
        self.do_normalize = bool(do_normalize)
        self.image_mean = list(image_mean) if image_mean is not None else [0.5, 0.5, 0.5]
        self.image_std = list(image_std) if image_std is not None else [0.5, 0.5, 0.5]
        self.pad_value = float(pad_value)
        self.processor_class = processor_class

    @classmethod
    def from_model_config(cls, model_config: Union[Dict[str, Any], Any]) -> "AnandaImageProcessor":
        if isinstance(model_config, dict):
            cfg = model_config
        else:
            cfg = vars(model_config)

        return cls(
            patch_size=int(cfg.get("patch_size", 16)),
            merge_factor=int(cfg.get("encoder_2d_merge_factor", 1)),
            image_mean=cfg.get("image_normalization_mean", [0.5, 0.5, 0.5]),
            image_std=cfg.get("image_normalization_std", [0.5, 0.5, 0.5]),
        )

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs: Any) -> "AnandaImageProcessor":
        config_dict = _load_image_processor_dict(pretrained_model_name_or_path, **kwargs)
        nested = config_dict.get("image_processor")
        if isinstance(nested, dict):
            config_dict = nested

        return cls(
            patch_size=int(config_dict.get("patch_size", 16)),
            merge_factor=int(config_dict.get("merge_factor", config_dict.get("encoder_2d_merge_factor", 1))),
            do_convert_rgb=bool(config_dict.get("do_convert_rgb", True)),
            do_rescale=bool(config_dict.get("do_rescale", True)),
            rescale_factor=float(config_dict.get("rescale_factor", 1.0 / 255.0)),
            do_normalize=bool(config_dict.get("do_normalize", True)),
            image_mean=config_dict.get("image_mean", config_dict.get("image_normalization_mean", [0.5, 0.5, 0.5])),
            image_std=config_dict.get("image_std", config_dict.get("image_normalization_std", [0.5, 0.5, 0.5])),
            pad_value=float(config_dict.get("pad_value", 0.0)),
            processor_class=config_dict.get("processor_class", "AnandaProcessor"),
        )

    def to_dict(self) -> Dict[str, Any]:
        return {
            "image_processor_type": self.__class__.__name__,
            "processor_class": self.processor_class,
            "auto_map": {
                "AutoImageProcessor": "inference_processor.AnandaImageProcessor",
                "AutoProcessor": "inference_processor.AnandaProcessor",
            },
            "patch_size": self.patch_size,
            "merge_factor": self.merge_factor,
            "do_convert_rgb": self.do_convert_rgb,
            "do_rescale": self.do_rescale,
            "rescale_factor": self.rescale_factor,
            "do_normalize": self.do_normalize,
            "image_mean": list(self.image_mean),
            "image_std": list(self.image_std),
            "pad_value": self.pad_value,
        }

    def save_pretrained(self, save_directory: Union[str, os.PathLike], **_: Any) -> List[str]:
        os.makedirs(save_directory, exist_ok=True)
        output_path = os.path.join(save_directory, PREPROCESSOR_CONFIG_NAME)
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(self.to_dict(), f, ensure_ascii=False, indent=2)
        return [output_path]

    @staticmethod
    def _ensure_list(images: Union[ImageLike, Sequence[ImageLike]]) -> List[ImageLike]:
        if isinstance(images, (list, tuple)):
            return list(images)
        return [images]

    def _to_chw_uint8(self, image: ImageLike) -> torch.Tensor:
        if isinstance(image, Image.Image):
            img = image.convert("RGB") if self.do_convert_rgb else image
            arr = np.array(img, dtype=np.uint8)
            tensor = torch.from_numpy(arr)
            if tensor.ndim == 2:
                tensor = tensor.unsqueeze(-1)
            tensor = tensor.permute(2, 0, 1).contiguous()
        elif isinstance(image, np.ndarray):
            arr = image
            if arr.ndim == 2:
                arr = arr[..., None]
            if arr.ndim != 3:
                raise ValueError(f"Expected 2D or 3D ndarray image, got shape={arr.shape}")
            tensor = torch.from_numpy(arr)
            if tensor.shape[0] in (1, 3, 4):
                pass
            elif tensor.shape[-1] in (1, 3, 4):
                tensor = tensor.permute(2, 0, 1)
            else:
                raise ValueError(f"Could not infer channel dimension from ndarray shape={arr.shape}")
            tensor = tensor.contiguous()
        elif torch.is_tensor(image):
            tensor = image.detach().cpu()
            if tensor.ndim == 2:
                tensor = tensor.unsqueeze(0)
            if tensor.ndim != 3:
                raise ValueError(f"Expected 2D or 3D tensor image, got shape={tuple(tensor.shape)}")
            if tensor.shape[0] in (1, 3, 4):
                pass
            elif tensor.shape[-1] in (1, 3, 4):
                tensor = tensor.permute(2, 0, 1)
            else:
                raise ValueError(f"Could not infer channel dimension from tensor shape={tuple(tensor.shape)}")
            tensor = tensor.contiguous()
        else:
            raise TypeError(f"Unsupported image type: {type(image)!r}")

        if tensor.shape[0] == 1:
            tensor = tensor.expand(3, -1, -1)
        elif tensor.shape[0] == 4:
            tensor = tensor[:3]
        elif tensor.shape[0] != 3:
            raise ValueError(f"Expected 1, 3, or 4 channels, got {tensor.shape[0]}")

        if tensor.dtype.is_floating_point:
            max_val = float(tensor.max().item()) if tensor.numel() else 0.0
            if max_val <= 1.0 + 1e-6:
                tensor = tensor * 255.0
            tensor = tensor.round().clamp_(0.0, 255.0).to(torch.uint8)
        else:
            tensor = tensor.clamp_(0, 255).to(torch.uint8)

        return tensor.contiguous()

    def _normalize(self, chw_u8: torch.Tensor) -> torch.Tensor:
        x = chw_u8.to(torch.float32)
        if self.do_rescale:
            x = x * self.rescale_factor
        mean = torch.tensor(self.image_mean, dtype=torch.float32).view(3, 1, 1)
        std = torch.tensor(self.image_std, dtype=torch.float32).view(3, 1, 1)
        if self.do_normalize:
            x = (x - mean) / std
        return x

    def _pad_to_patch_multiple(self, img: torch.Tensor) -> torch.Tensor:
        _, h, w = img.shape
        p = self.patch_size
        target_h = int(math.ceil(h / p) * p)
        target_w = int(math.ceil(w / p) * p)
        if target_h != h or target_w != w:
            img = F.pad(img, (0, target_w - w, 0, target_h - h), value=self.pad_value)
        return img

    def _pad_for_merge_factor(self, img_norm: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        if img_norm.ndim != 3:
            raise ValueError(f"Expected image tensor with shape (3,H,W), got {tuple(img_norm.shape)}")

        p = self.patch_size
        m = self.merge_factor
        base = p * m
        _, h, w = img_norm.shape

        if h % p != 0 or w % p != 0:
            raise ValueError(f"Image must be patch-multiple before merge padding, got H={h}, W={w}, patch_size={p}")

        target_h = int(math.ceil(h / base) * base)
        target_w = int(math.ceil(w / base) * base)

        ph, pw = h // p, w // p
        target_ph, target_pw = target_h // p, target_w // p

        mask_2d = torch.ones((ph, pw), dtype=torch.bool)
        if target_ph != ph or target_pw != pw:
            mask_2d = F.pad(mask_2d, (0, target_pw - pw, 0, target_ph - ph), value=False)

        if target_h != h or target_w != w:
            img_norm = F.pad(img_norm, (0, target_w - w, 0, target_h - h), value=self.pad_value)

        return img_norm, mask_2d.reshape(-1).to(torch.long)

    def _preprocess_single(self, image: ImageLike) -> Tuple[torch.Tensor, torch.Tensor]:
        chw_u8 = self._to_chw_uint8(image)
        img = self._normalize(chw_u8)
        img = self._pad_to_patch_multiple(img)
        return self._pad_for_merge_factor(img)

    def preprocess(
        self,
        images: Union[ImageLike, Sequence[ImageLike]],
        return_tensors: Optional[Union[str, TensorType]] = None,
        **_: Any,
    ) -> BatchFeature:
        image_list = self._ensure_list(images)
        if len(image_list) == 0:
            raise ValueError("`images` must contain at least one image")

        processed: List[torch.Tensor] = []
        patch_masks: List[torch.Tensor] = []
        for image in image_list:
            px, pm = self._preprocess_single(image)
            processed.append(px)
            patch_masks.append(pm)

        max_h = max(t.shape[1] for t in processed)
        max_w = max(t.shape[2] for t in processed)
        p = self.patch_size
        batch_patch_h = max_h // p
        batch_patch_w = max_w // p

        batch_pixels: List[torch.Tensor] = []
        batch_masks: List[torch.Tensor] = []
        for px, pm in zip(processed, patch_masks):
            _, h, w = px.shape
            ph, pw = h // p, w // p

            if h != max_h or w != max_w:
                px = F.pad(px, (0, max_w - w, 0, max_h - h), value=self.pad_value)

            pm_2d = pm.view(ph, pw).to(torch.bool)
            if ph != batch_patch_h or pw != batch_patch_w:
                pm_2d = F.pad(pm_2d, (0, batch_patch_w - pw, 0, batch_patch_h - ph), value=False)

            batch_pixels.append(px)
            batch_masks.append(pm_2d.reshape(-1).to(torch.long))

        data = {
            "pixel_values": torch.stack(batch_pixels, dim=0),
            "patch_attention_mask": torch.stack(batch_masks, dim=0),
        }
        return BatchFeature(data=data, tensor_type=return_tensors)

    __call__ = preprocess


class AnandaProcessor(ProcessorMixin):
    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "AutoImageProcessor"
    tokenizer_class = "AutoTokenizer"
    model_input_names = ["input_ids", "attention_mask", "pixel_values", "patch_attention_mask"]

    def __init__(self, image_processor: AnandaImageProcessor, tokenizer, **kwargs: Any) -> None:
        self.image_processor = image_processor
        self.tokenizer = tokenizer
        self.current_processor = self.image_processor
        self._in_target_context_manager = False
        super().__init__(image_processor, tokenizer, **kwargs)

    @classmethod
    def from_model_config(cls, tokenizer, model_config: Union[Dict[str, Any], Any]) -> "AnandaProcessor":
        image_processor = AnandaImageProcessor.from_model_config(model_config)
        return cls(image_processor=image_processor, tokenizer=tokenizer)

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: Union[str, os.PathLike],
        trust_remote_code: bool = True,
        **kwargs: Any,
    ) -> "AnandaProcessor":
        tokenizer = AutoTokenizer.from_pretrained(
            pretrained_model_name_or_path,
            trust_remote_code=trust_remote_code,
            **kwargs,
        )
        image_processor = AnandaImageProcessor.from_pretrained(pretrained_model_name_or_path, **kwargs)
        return cls(image_processor=image_processor, tokenizer=tokenizer)

    def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs: Any) -> List[str]:
        os.makedirs(save_directory, exist_ok=True)

        saved_files: List[str] = []
        saved_files.extend(self.image_processor.save_pretrained(save_directory))
        saved_files.extend(self.tokenizer.save_pretrained(save_directory))

        processor_dict = {
            "processor_class": self.__class__.__name__,
            "auto_map": {"AutoProcessor": "inference_processor.AnandaProcessor"},
            "image_processor": self.image_processor.to_dict(),
        }
        output_path = os.path.join(save_directory, PROCESSOR_CONFIG_NAME)
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(processor_dict, f, ensure_ascii=False, indent=2)
        saved_files.append(output_path)
        return saved_files

    def __call__(
        self,
        text: Optional[Union[str, Sequence[str]]] = None,
        images: Optional[Union[ImageLike, Sequence[ImageLike]]] = None,
        return_tensors: Optional[Union[str, TensorType]] = None,
        add_special_tokens: bool = True,
        **kwargs: Any,
    ) -> BatchFeature:
        if text is None and images is None:
            raise ValueError("At least one of `text` or `images` must be provided")

        encoding: Dict[str, Any] = {}

        if images is not None:
            image_features = self.image_processor(images=images, return_tensors=return_tensors)
            encoding.update(image_features)
            batch_size = int(image_features["pixel_values"].shape[0])
        else:
            batch_size = None

        if text is None:
            bos_id = self.tokenizer.bos_token_id
            eos_id = self.tokenizer.eos_token_id
            prompt_id = bos_id if bos_id is not None else eos_id
            if prompt_id is None:
                raise ValueError("Tokenizer must define bos_token_id or eos_token_id.")
            if batch_size is None:
                batch_size = 1

            input_ids = [[int(prompt_id)] for _ in range(batch_size)]
            attention_mask = [[1] for _ in range(batch_size)]
            if return_tensors == "pt" or return_tensors == TensorType.PYTORCH:
                encoding["input_ids"] = torch.tensor(input_ids, dtype=torch.long)
                encoding["attention_mask"] = torch.tensor(attention_mask, dtype=torch.long)
            else:
                encoding["input_ids"] = input_ids
                encoding["attention_mask"] = attention_mask
        else:
            text_encoding = self.tokenizer(
                text,
                add_special_tokens=add_special_tokens,
                return_tensors=return_tensors,
                **kwargs,
            )
            encoding.update(text_encoding)

        return BatchFeature(data=encoding, tensor_type=return_tensors)

    def batch_decode(self, *args: Any, **kwargs: Any):
        return self.tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args: Any, **kwargs: Any):
        return self.tokenizer.decode(*args, **kwargs)

    def apply_chat_template(self, *args: Any, **kwargs: Any):
        return self.tokenizer.apply_chat_template(*args, **kwargs)