kimi_k3_vision_processing.py

"""Image processor class for Kimi-K3.
"""

import json
from typing import Any, Dict, Optional, Union

import numpy as np
import torch
from PIL import Image
from transformers.image_processing_utils import (BaseImageProcessor,
                                                 BatchFeature)
from transformers.utils import TensorType

from .media_utils import (MediaInput, TransparentBgConfig, _to_tensor,
                          ensure_media_type, image_to_np, navit_patchify,
                          navit_resize_image, normalize)


class KimiK3VisionProcessor(BaseImageProcessor):
    model_type = "kimi_k3"

    def __init__(
        self,
        media_proc_cfg: dict,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.media_proc_cfg = media_proc_cfg

    @property
    def _transparent_bg_config(self) -> Optional[TransparentBgConfig]:
        cfg = self.media_proc_cfg.get("transparent_bg_config")
        if cfg is None:
            return None
        if isinstance(cfg, TransparentBgConfig):
            return cfg
        return TransparentBgConfig(**cfg)

    @property
    def _transparent_bg_fill_stage(self) -> str:
        return self.media_proc_cfg.get("transparent_bg_fill_stage",
                                       "before_resize")

    def media_tokens_calculator(self, media: MediaInput):
        media = ensure_media_type(
            media,
            transparent_bg_config=self._transparent_bg_config,
            transparent_bg_fill_stage=self._transparent_bg_fill_stage,
        )
        ret = self.get_resize_config(media)
        return ret['num_tokens']

    @classmethod
    def make_image_prompt(cls, width: int, height: int) -> str:
        """Build the K3 image placeholder with resolution info."""
        return (f"<|media_begin|>image {width}x{height}"
                f"<|media_content|><|media_pad|><|media_end|>")

    def get_resize_config(self, media_input: MediaInput) -> dict:
        if media_input['type'] == 'image':
            w, h = media_input['image'].size
            ret = navit_resize_image(
                w, h, self.media_proc_cfg['patch_size'],
                self.media_proc_cfg['merge_kernel_size'],
                self.media_proc_cfg['in_patch_limit'],
                self.media_proc_cfg['patch_limit_on_one_side'],
                self.media_proc_cfg['fixed_output_tokens'])
            return ret
        else:
            raise ValueError("Unsupported type: {}".format(
                media_input['type']))

    def resize_image(self, image: Image.Image, new_width: int, new_height: int,
                     pad_width: int, pad_height: int) -> np.ndarray:
        image_np = image_to_np(
            image,
            (new_width, new_height),
            "resize",
            transparent_bg_config=self._transparent_bg_config,
            transparent_bg_fill_stage=self._transparent_bg_fill_stage,
        )
        image_np = np.pad(
            image_np,
            ((0, pad_height), (0, pad_width), (0, 0)),
            mode="constant",
            constant_values=0,
        )
        return image_np

    def preprocess(
        self,
        medias: list[MediaInput],
        return_tensors: Optional[Union[str, TensorType]] = None,
    ) -> BatchFeature:
        """
        Preprocess a atom vision input (images) into model-ready tensors.

        Args:
            medias: List of MediaInput.
            return_tensors: Desired output format ('pt', 'np', 'tf', or None).

        Returns:
            BatchFeature containing 'pixel_values' and 'grid_thws' tensors.
        """
        if not isinstance(medias, list):
            medias = [medias]
        if medias:
            pixel_values = []
            for item in medias:
                item = ensure_media_type(
                    item,
                    transparent_bg_config=self._transparent_bg_config,
                    transparent_bg_fill_stage=self._transparent_bg_fill_stage,
                )
                resize_config = self.get_resize_config(item)
                new_width, new_height, pad_width, pad_height = resize_config[
                    'new_width'], resize_config['new_height'], resize_config[
                        'pad_width'], resize_config['pad_height']
                if item['type'] == 'image':
                    image = item['image']
                    image_np = self.resize_image(image, new_width, new_height,
                                                 pad_width, pad_height)
                    pixel_values.append(np.expand_dims(image_np, axis=0))
                else:
                    raise ValueError("Unsupported type: {}".format(
                        item['type']))
            normalized_pixel_values = []
            image_std_inv = 1.0 / np.array(self.media_proc_cfg['image_std'])
            image_mean = np.array(self.media_proc_cfg['image_mean'])
            for pixels in pixel_values:
                pixels = normalize(pixels, image_mean, image_std_inv)
                pixels_and_thw = navit_patchify(
                    pixels,
                    self.media_proc_cfg['patch_size'],
                )
                normalized_pixel_values.append(pixels_and_thw)

            pixel_values = torch.cat([
                _to_tensor(pixel_value['pixel_values'])
                for pixel_value in normalized_pixel_values
            ])
            grid_thws = torch.cat([
                _to_tensor(pixel_value['grid_thw'],
                           dtype=torch.int64).unsqueeze(0)
                for pixel_value in normalized_pixel_values
            ])

            data = {
                'pixel_values': pixel_values,
                'grid_thws': grid_thws,
            }

        else:
            data = {}

        return BatchFeature(data=data, tensor_type=return_tensors)

    def __repr__(self):
        return f"KimiK3VisionProcessor(media_proc_cfg={self.media_proc_cfg})"

    def to_dict(self) -> Dict[str, Any]:
        output = super().to_dict()
        output["media_proc_cfg"] = self.media_proc_cfg
        if "media_processor" in output:
            del output["media_processor"]
        return output

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any], **kwargs):
        config = config_dict.copy()
        media_proc_cfg = config.pop("media_proc_cfg", {})
        return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs)

    def to_json_string(self):
        dictionary = self.to_dict()
        for key, value in dictionary.items():
            if hasattr(value, 'tolist'):
                dictionary[key] = value.tolist()
        return json.dumps(dictionary, indent=2, sort_keys=True) + "\n"