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)