"""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"