SmartManoj/Kimi-K3

1"""Image processor class for Kimi-K3.2"""3 4import json5from typing import Any, Dict, Optional, Union6 7import numpy as np8import torch9from PIL import Image10from transformers.image_processing_utils import (BaseImageProcessor,11                                                 BatchFeature)12from transformers.utils import TensorType13 14from .media_utils import (MediaInput, TransparentBgConfig, _to_tensor,15                          ensure_media_type, image_to_np, navit_patchify,16                          navit_resize_image, normalize)17 18 19class KimiK3VisionProcessor(BaseImageProcessor):20    model_type = "kimi_k3"21 22    def __init__(23        self,24        media_proc_cfg: dict,25        **kwargs,26    ):27        super().__init__(**kwargs)28        self.media_proc_cfg = media_proc_cfg29 30    @property31    def _transparent_bg_config(self) -> Optional[TransparentBgConfig]:32        cfg = self.media_proc_cfg.get("transparent_bg_config")33        if cfg is None:34            return None35        if isinstance(cfg, TransparentBgConfig):36            return cfg37        return TransparentBgConfig(**cfg)38 39    @property40    def _transparent_bg_fill_stage(self) -> str:41        return self.media_proc_cfg.get("transparent_bg_fill_stage",42                                       "before_resize")43 44    def media_tokens_calculator(self, media: MediaInput):45        media = ensure_media_type(46            media,47            transparent_bg_config=self._transparent_bg_config,48            transparent_bg_fill_stage=self._transparent_bg_fill_stage,49        )50        ret = self.get_resize_config(media)51        return ret['num_tokens']52 53    @classmethod54    def make_image_prompt(cls, width: int, height: int) -> str:55        """Build the K3 image placeholder with resolution info."""56        return (f"<|media_begin|>image {width}x{height}"57                f"<|media_content|><|media_pad|><|media_end|>")58 59    def get_resize_config(self, media_input: MediaInput) -> dict:60        if media_input['type'] == 'image':61            w, h = media_input['image'].size62            ret = navit_resize_image(63                w, h, self.media_proc_cfg['patch_size'],64                self.media_proc_cfg['merge_kernel_size'],65                self.media_proc_cfg['in_patch_limit'],66                self.media_proc_cfg['patch_limit_on_one_side'],67                self.media_proc_cfg['fixed_output_tokens'])68            return ret69        else:70            raise ValueError("Unsupported type: {}".format(71                media_input['type']))72 73    def resize_image(self, image: Image.Image, new_width: int, new_height: int,74                     pad_width: int, pad_height: int) -> np.ndarray:75        image_np = image_to_np(76            image,77            (new_width, new_height),78            "resize",79            transparent_bg_config=self._transparent_bg_config,80            transparent_bg_fill_stage=self._transparent_bg_fill_stage,81        )82        image_np = np.pad(83            image_np,84            ((0, pad_height), (0, pad_width), (0, 0)),85            mode="constant",86            constant_values=0,87        )88        return image_np89 90    def preprocess(91        self,92        medias: list[MediaInput],93        return_tensors: Optional[Union[str, TensorType]] = None,94    ) -> BatchFeature:95        """96        Preprocess a atom vision input (images) into model-ready tensors.97 98        Args:99            medias: List of MediaInput.100            return_tensors: Desired output format ('pt', 'np', 'tf', or None).101 102        Returns:103            BatchFeature containing 'pixel_values' and 'grid_thws' tensors.104        """105        if not isinstance(medias, list):106            medias = [medias]107        if medias:108            pixel_values = []109            for item in medias:110                item = ensure_media_type(111                    item,112                    transparent_bg_config=self._transparent_bg_config,113                    transparent_bg_fill_stage=self._transparent_bg_fill_stage,114                )115                resize_config = self.get_resize_config(item)116                new_width, new_height, pad_width, pad_height = resize_config[117                    'new_width'], resize_config['new_height'], resize_config[118                        'pad_width'], resize_config['pad_height']119                if item['type'] == 'image':120                    image = item['image']121                    image_np = self.resize_image(image, new_width, new_height,122                                                 pad_width, pad_height)123                    pixel_values.append(np.expand_dims(image_np, axis=0))124                else:125                    raise ValueError("Unsupported type: {}".format(126                        item['type']))127            normalized_pixel_values = []128            image_std_inv = 1.0 / np.array(self.media_proc_cfg['image_std'])129            image_mean = np.array(self.media_proc_cfg['image_mean'])130            for pixels in pixel_values:131                pixels = normalize(pixels, image_mean, image_std_inv)132                pixels_and_thw = navit_patchify(133                    pixels,134                    self.media_proc_cfg['patch_size'],135                )136                normalized_pixel_values.append(pixels_and_thw)137 138            pixel_values = torch.cat([139                _to_tensor(pixel_value['pixel_values'])140                for pixel_value in normalized_pixel_values141            ])142            grid_thws = torch.cat([143                _to_tensor(pixel_value['grid_thw'],144                           dtype=torch.int64).unsqueeze(0)145                for pixel_value in normalized_pixel_values146            ])147 148            data = {149                'pixel_values': pixel_values,150                'grid_thws': grid_thws,151            }152 153        else:154            data = {}155 156        return BatchFeature(data=data, tensor_type=return_tensors)157 158    def __repr__(self):159        return f"KimiK3VisionProcessor(media_proc_cfg={self.media_proc_cfg})"160 161    def to_dict(self) -> Dict[str, Any]:162        output = super().to_dict()163        output["media_proc_cfg"] = self.media_proc_cfg164        if "media_processor" in output:165            del output["media_processor"]166        return output167 168    @classmethod169    def from_dict(cls, config_dict: Dict[str, Any], **kwargs):170        config = config_dict.copy()171        media_proc_cfg = config.pop("media_proc_cfg", {})172        return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs)173 174    def to_json_string(self):175        dictionary = self.to_dict()176        for key, value in dictionary.items():177            if hasattr(value, 'tolist'):178                dictionary[key] = value.tolist()179        return json.dumps(dictionary, indent=2, sort_keys=True) + "\n"180