| """Image and video processor class for Kimi-K3. |
| """ |
|
|
| import base64 |
| import io |
| import json |
| import os |
| 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 transformers.video_utils import load_video, read_video_pyav |
|
|
| from .media_utils import (MediaInput, TransparentBgConfig, VideoChunkInput, |
| _to_tensor, ensure_media_type, image_to_np, |
| navit_patchify, navit_resize_image, normalize, |
| timestamp_as_str) |
|
|
|
|
| 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 |
| self.num_frames_per_chunk = media_proc_cfg[ |
| 'temporal_merge_kernel_size'] |
|
|
| @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|>") |
|
|
| @classmethod |
| def make_chunk_prompt(cls, timestamp_text: str) -> str: |
| """Build one timestamped K3 video-chunk placeholder.""" |
| return (f"{timestamp_text}<|media_begin|>video<|media_content|>" |
| f"<|media_pad|><|media_end|>") |
|
|
| @staticmethod |
| def _normalize_video_source(video_source): |
| if isinstance(video_source, bytes): |
| return io.BytesIO(video_source) |
| if isinstance(video_source, os.PathLike): |
| return os.fspath(video_source) |
| if isinstance(video_source, str) and video_source.startswith('data:'): |
| header, encoded = video_source.split(',', 1) |
| if not header.startswith('data:video/') or ';base64' not in header: |
| raise ValueError( |
| "Video data URLs must contain base64-encoded video data.") |
| return io.BytesIO(base64.b64decode(encoded, validate=True)) |
| return video_source |
|
|
| def split_video_chunks(self, video_source) -> list[VideoChunkInput]: |
| """Decode, uniformly sample, and split a video into temporal chunks.""" |
| config = self.media_proc_cfg |
|
|
| def sample_indices(metadata, **_): |
| total_frames = int(metadata.total_num_frames) |
| video_fps = float(metadata.fps or 0) |
| if total_frames <= 0 or video_fps <= 0: |
| raise ValueError( |
| "Video must contain frames and report a positive FPS.") |
|
|
| sample_fps = min(float(config['sample_fps']), video_fps) |
| sampled_nframes = max( |
| round(total_frames * sample_fps / video_fps), 1) |
| max_frames = config.get('max_num_frames_each_video') |
| if max_frames is not None: |
| sampled_nframes = min(sampled_nframes, int(max_frames)) |
| sampled_nframes = min(sampled_nframes, total_frames) |
| return np.linspace(0, total_frames - 1, |
| sampled_nframes).round().astype(int) |
|
|
| source = self._normalize_video_source(video_source) |
| if isinstance(source, io.BytesIO): |
| sampled_frames, metadata = read_video_pyav( |
| source, sample_indices_fn=sample_indices) |
| else: |
| sampled_frames, metadata = load_video( |
| source, backend="pyav", sample_indices_fn=sample_indices) |
|
|
| if isinstance(sampled_frames, torch.Tensor): |
| sampled_frames = sampled_frames.cpu().numpy() |
| frames = [Image.fromarray(np.asarray(frame).astype(np.uint8)) |
| for frame in sampled_frames] |
| if not frames: |
| raise ValueError("Video sampling did not return any frames.") |
|
|
| per_frame_limit = int(config.get('in_patch_limit_each_frame') |
| or config['in_patch_limit']) |
| total_patch_limit = config.get('in_patch_limit_video') |
| if total_patch_limit is not None: |
| per_frame_limit = min( |
| per_frame_limit, |
| max(1, round(int(total_patch_limit) / len(frames))), |
| ) |
|
|
| frame_indices = metadata.frames_indices |
| video_fps = float(metadata.fps) |
| chunks = [] |
| for start in range(0, len(frames), self.num_frames_per_chunk): |
| timestamp = timestamp_as_str( |
| frame_indices[start] / video_fps, |
| config['timestamp_mode'], |
| ) |
| chunks.append( |
| VideoChunkInput( |
| type='video_chunk', |
| video_chunk=frames[start:start + self.num_frames_per_chunk], |
| prompt=self.make_chunk_prompt(timestamp), |
| in_patch_limit=per_frame_limit, |
| )) |
| return chunks |
|
|
| 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 |
| elif media_input['type'] == 'video_chunk': |
| frame = media_input['video_chunk'][0] |
| return navit_resize_image( |
| frame.width, frame.height, |
| self.media_proc_cfg['patch_size'], |
| self.media_proc_cfg['merge_kernel_size'], |
| media_input['in_patch_limit'], |
| self.media_proc_cfg['patch_limit_on_one_side'], |
| self.media_proc_cfg['fixed_output_tokens']) |
| 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 images and video chunks into model-ready tensors. |
| |
| Video audio tracks are not processed. |
| |
| 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)) |
| elif item['type'] == 'video_chunk': |
| pixels = [ |
| self.resize_image(frame, new_width, new_height, |
| pad_width, pad_height) |
| for frame in item['video_chunk'] |
| ] |
| pixel_values.append(np.stack(pixels, 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" |
|
|
