| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Optimized CPU image processor for Kimi-K2.5/K2.6 vision chunks.""" |
| 4 | |
| 5 | import io |
| 6 | import json |
| 7 | import math |
| 8 | from typing import Any |
| 9 | |
| 10 | import numpy as np |
| 11 | import pybase64 as base64 |
| 12 | import torch |
| 13 | from PIL import Image |
| 14 | from transformers.image_processing_utils import BaseImageProcessor, BatchFeature |
| 15 | from transformers.utils import TensorType |
| 16 | |
| 17 | from vllm.utils.import_utils import is_numba_available |
| 18 | from vllm.utils.jit_monitor import numba_workqueue_threading_layer |
| 19 | |
| 20 | if is_numba_available(): |
| 21 | from numba import njit, prange |
| 22 | |
| 23 | @njit(parallel=True, cache=True) |
| 24 | def _write_fused_patches( |
| 25 | frames: np.ndarray, |
| 26 | out: np.ndarray, |
| 27 | out_offset: int, |
| 28 | new_h: int, |
| 29 | new_w: int, |
| 30 | padded_h: int, |
| 31 | padded_w: int, |
| 32 | patch_size: int, |
| 33 | normalize_lut: np.ndarray, |
| 34 | ) -> None: |
| 35 | # frames: [T, new_h, new_w, 3] uint8, without padding. |
| 36 | # out: [total_patches, 3, patch_size, patch_size] float32. |
| 37 | t_size = frames.shape[0] |
| 38 | patch_h = padded_h // patch_size |
| 39 | patch_w = padded_w // patch_size |
| 40 | total = t_size * padded_h * padded_w * 3 |
| 41 | hwc = padded_h * padded_w * 3 |
| 42 | wc = padded_w * 3 |
| 43 | |
| 44 | for linear in prange(total): |
| 45 | t = linear // hwc |
| 46 | rem = linear - t * hwc |
| 47 | y = rem // wc |
| 48 | rem = rem - y * wc |
| 49 | x = rem // 3 |
| 50 | c = rem - x * 3 |
| 51 | |
| 52 | value = frames[t, y, x, c] if y < new_h and x < new_w else 0 |
| 53 | |
| 54 | patch_idx = ( |
| 55 | out_offset |
| 56 | + t * patch_h * patch_w |
| 57 | + (y // patch_size) * patch_w |
| 58 | + (x // patch_size) |
| 59 | ) |
| 60 | out[patch_idx, c, y % patch_size, x % patch_size] = normalize_lut[value, c] |
| 61 | |
| 62 | else: |
| 63 | |
| 64 | def _write_fused_patches(*args: Any, **kwargs: Any) -> None: |
| 65 | raise RuntimeError("numba is required for fused Kimi image preprocessing") |
| 66 | |
| 67 | |
| 68 | def navit_resize_image( |
| 69 | width: int, |
| 70 | height: int, |
| 71 | patch_size: int, |
| 72 | merge_kernel_size: int, |
| 73 | in_patch_limit: int, |
| 74 | patch_limit_on_one_side: int, |
| 75 | fixed_output_tokens: int | None, |
| 76 | ) -> dict[str, int]: |
| 77 | s1 = math.sqrt( |
| 78 | in_patch_limit |
| 79 | / (max(1.0, width // patch_size) * max(1.0, height // patch_size)) |
| 80 | ) |
| 81 | s2 = patch_limit_on_one_side * patch_size / width |
| 82 | s3 = patch_limit_on_one_side * patch_size / height |
| 83 | scale = min(1.0, s1, s2, s3) |
| 84 | new_w = min(max(1, int(width * scale)), patch_limit_on_one_side * patch_size) |
| 85 | new_h = min(max(1, int(height * scale)), patch_limit_on_one_side * patch_size) |
| 86 | |
| 87 | factor = merge_kernel_size * patch_size |
| 88 | pad_height = (factor - new_h % factor) % factor |
| 89 | pad_width = (factor - new_w % factor) % factor |
| 90 | |
| 91 | if fixed_output_tokens is not None: |
| 92 | num_tokens = fixed_output_tokens |
| 93 | else: |
| 94 | token_height = (new_h + pad_height) // factor |
| 95 | token_width = (new_w + pad_width) // factor |
| 96 | num_tokens = token_height * token_width |
| 97 | |
| 98 | return { |
| 99 | "num_tokens": num_tokens, |
| 100 | "new_width": new_w, |
| 101 | "new_height": new_h, |
| 102 | "pad_width": pad_width, |
| 103 | "pad_height": pad_height, |
| 104 | "sampled_nframes": 1, |
| 105 | } |
| 106 | |
| 107 | |
| 108 | def navit_resize_video( |
| 109 | width: int, |
| 110 | height: int, |
| 111 | nframes: int, |
| 112 | avg_fps: float, |
| 113 | sample_fps: float, |
| 114 | patch_size: int, |
| 115 | merge_kernel_size: int, |
| 116 | in_patch_limit_each_frame: int, |
| 117 | patch_limit_on_one_side: int, |
| 118 | in_patch_limit_total: int | None, |
| 119 | max_num_frames_each_video: int | None, |
| 120 | fixed_output_tokens_each_frame: int | None, |
| 121 | ) -> dict[str, int]: |
| 122 | sample_fps = min(sample_fps, avg_fps) |
| 123 | sampled_nframes = max(round(nframes * sample_fps / avg_fps), 1) |
| 124 | if max_num_frames_each_video is not None: |
| 125 | sampled_nframes = min(sampled_nframes, max_num_frames_each_video) |
| 126 | |
| 127 | if in_patch_limit_total is not None: |
| 128 | in_patch_limit_each_frame = min( |
| 129 | round(in_patch_limit_total / sampled_nframes), |
| 130 | in_patch_limit_each_frame, |
| 131 | ) |
| 132 | |
| 133 | ret = navit_resize_image( |
| 134 | width, |
| 135 | height, |
| 136 | patch_size, |
| 137 | merge_kernel_size, |
| 138 | in_patch_limit_each_frame, |
| 139 | patch_limit_on_one_side, |
| 140 | fixed_output_tokens_each_frame, |
| 141 | ) |
| 142 | ret["sampled_nframes"] = sampled_nframes |
| 143 | return ret |
| 144 | |
| 145 | |
| 146 | def _to_pil(data: Any) -> Image.Image: |
| 147 | if hasattr(data, "media") and hasattr(data, "original_bytes"): |
| 148 | data = data.media |
| 149 | if isinstance(data, Image.Image): |
| 150 | return data if data.mode == "RGB" else data.convert("RGB") |
| 151 | if isinstance(data, str): |
| 152 | if data.startswith("data:"): |
| 153 | raw_base64 = data.split(",", 1)[1] |
| 154 | return Image.open(io.BytesIO(base64.b64decode(raw_base64))).convert("RGB") |
| 155 | return Image.open(data).convert("RGB") |
| 156 | if isinstance(data, bytes): |
| 157 | return Image.open(io.BytesIO(data)).convert("RGB") |
| 158 | raise ValueError(f"Unsupported data type: {type(data)}") |
| 159 | |
| 160 | |
| 161 | def _ensure_media_type(media: dict[str, Any]) -> dict[str, Any]: |
| 162 | if media["type"] == "image": |
| 163 | media["image"] = _to_pil(media["image"]) |
| 164 | return media |
| 165 | if media["type"] == "video_chunk": |
| 166 | media["video_chunk"] = [_to_pil(frame) for frame in media["video_chunk"]] |
| 167 | return media |
| 168 | raise ValueError(f"Unsupported media type: {media['type']}") |
| 169 | |
| 170 | |
| 171 | class KimiK25FusedVisionProcessor(BaseImageProcessor): |
| 172 | model_type = "kimi_k25" |
| 173 | |
| 174 | def __init__(self, media_proc_cfg: dict[str, Any], **kwargs: Any) -> None: |
| 175 | super().__init__(**kwargs) |
| 176 | media_proc_cfg = dict(media_proc_cfg) |
| 177 | merge_kernel_size = media_proc_cfg["merge_kernel_size"] |
| 178 | if isinstance(merge_kernel_size, (list, tuple)): |
| 179 | media_proc_cfg["merge_kernel_size"] = int(merge_kernel_size[0]) |
| 180 | self.media_proc_cfg = media_proc_cfg |
| 181 | self.num_frames_per_chunk = media_proc_cfg["temporal_merge_kernel_size"] |
| 182 | values = np.arange(256, dtype=np.float32)[:, None] |
| 183 | image_mean = np.asarray(media_proc_cfg["image_mean"], dtype=np.float32) |
| 184 | image_std_inv = 1.0 / np.asarray(media_proc_cfg["image_std"], dtype=np.float32) |
| 185 | self.normalize_lut = (values / 255.0 - image_mean[None, :]) * image_std_inv[ |
| 186 | None, : |
| 187 | ] |
| 188 | |
| 189 | def media_tokens_calculator(self, media: dict[str, Any]) -> int: |
| 190 | media = _ensure_media_type(media) |
| 191 | ret = self.get_resize_config(media) |
| 192 | return ret["num_tokens"] |
| 193 | |
| 194 | def get_resize_config(self, media_input: dict[str, Any]) -> dict[str, int]: |
| 195 | if media_input["type"] == "image": |
| 196 | width, height = media_input["image"].size |
| 197 | return navit_resize_image( |
| 198 | width, |
| 199 | height, |
| 200 | self.media_proc_cfg["patch_size"], |
| 201 | self.media_proc_cfg["merge_kernel_size"], |
| 202 | self.media_proc_cfg["in_patch_limit"], |
| 203 | self.media_proc_cfg["patch_limit_on_one_side"], |
| 204 | self.media_proc_cfg["fixed_output_tokens"], |
| 205 | ) |
| 206 | |
| 207 | if media_input["type"] == "video_chunk": |
| 208 | frame = media_input["video_chunk"][0] |
| 209 | width, height = frame.size |
| 210 | num_frames = len(media_input["video_chunk"]) |
| 211 | in_patch_limit_each_frame = self.media_proc_cfg["in_patch_limit_each_frame"] |
| 212 | if in_patch_limit_each_frame is None: |
| 213 | in_patch_limit_each_frame = self.media_proc_cfg["in_patch_limit"] |
| 214 | |
| 215 | return navit_resize_video( |
| 216 | width, |
| 217 | height, |
| 218 | num_frames, |
| 219 | 1.0, |
| 220 | math.inf, |
| 221 | self.media_proc_cfg["patch_size"], |
| 222 | self.media_proc_cfg["merge_kernel_size"], |
| 223 | in_patch_limit_each_frame, |
| 224 | self.media_proc_cfg["patch_limit_on_one_side"], |
| 225 | self.media_proc_cfg["in_patch_limit_video"], |
| 226 | None, |
| 227 | self.media_proc_cfg["fixed_output_tokens"], |
| 228 | ) |
| 229 | |
| 230 | raise ValueError(f"Unsupported type: {media_input['type']}") |
| 231 | |
| 232 | @staticmethod |
| 233 | def resize_image(image: Image.Image, new_width: int, new_height: int) -> np.ndarray: |
| 234 | image = image.resize((new_width, new_height), resample=Image.Resampling.BICUBIC) |
| 235 | return np.asarray(image) |
| 236 | |
| 237 | def preprocess( |
| 238 | self, |
| 239 | medias: list[dict[str, Any]], |
| 240 | return_tensors: str | TensorType | None = None, |
| 241 | ) -> BatchFeature: |
| 242 | if not isinstance(medias, list): |
| 243 | medias = [medias] |
| 244 | if not medias: |
| 245 | return BatchFeature(data={}, tensor_type=return_tensors) |
| 246 | |
| 247 | if njit is None: |
| 248 | raise RuntimeError("numba is required for fused Kimi image preprocessing") |
| 249 | |
| 250 | patch_size = int(self.media_proc_cfg["patch_size"]) |
| 251 | prepared = [] |
| 252 | grid_thws_np = np.empty((len(medias), 3), dtype=np.int64) |
| 253 | total_patches = 0 |
| 254 | |
| 255 | for idx, item in enumerate(medias): |
| 256 | item = _ensure_media_type(item) |
| 257 | resize_config = self.get_resize_config(item) |
| 258 | new_width = resize_config["new_width"] |
| 259 | new_height = resize_config["new_height"] |
| 260 | pad_width = resize_config["pad_width"] |
| 261 | pad_height = resize_config["pad_height"] |
| 262 | padded_width = new_width + pad_width |
| 263 | padded_height = new_height + pad_height |
| 264 | |
| 265 | if item["type"] == "image": |
| 266 | image_np = self.resize_image(item["image"], new_width, new_height) |
| 267 | frames = image_np[np.newaxis, ...] |
| 268 | elif item["type"] == "video_chunk": |
| 269 | frames = np.stack( |
| 270 | [ |
| 271 | self.resize_image(frame, new_width, new_height) |
| 272 | for frame in item["video_chunk"] |
| 273 | ], |
| 274 | axis=0, |
| 275 | ) |
| 276 | else: |
| 277 | raise ValueError(f"Unsupported type: {item['type']}") |
| 278 | |
| 279 | t_size = frames.shape[0] |
| 280 | grid_h = padded_height // patch_size |
| 281 | grid_w = padded_width // patch_size |
| 282 | grid_thws_np[idx, 0] = t_size |
| 283 | grid_thws_np[idx, 1] = grid_h |
| 284 | grid_thws_np[idx, 2] = grid_w |
| 285 | |
| 286 | num_patches = t_size * grid_h * grid_w |
| 287 | prepared.append( |
| 288 | ( |
| 289 | frames, |
| 290 | new_height, |
| 291 | new_width, |
| 292 | padded_height, |
| 293 | padded_width, |
| 294 | num_patches, |
| 295 | ) |
| 296 | ) |
| 297 | total_patches += num_patches |
| 298 | |
| 299 | pixel_values_np = np.empty( |
| 300 | (total_patches, 3, patch_size, patch_size), dtype=np.float32 |
| 301 | ) |
| 302 | out_offset = 0 |
| 303 | with numba_workqueue_threading_layer(): |
| 304 | for ( |
| 305 | frames, |
| 306 | new_height, |
| 307 | new_width, |
| 308 | padded_height, |
| 309 | padded_width, |
| 310 | num_patches, |
| 311 | ) in prepared: |
| 312 | _write_fused_patches( |
| 313 | frames, |
| 314 | pixel_values_np, |
| 315 | out_offset, |
| 316 | new_height, |
| 317 | new_width, |
| 318 | padded_height, |
| 319 | padded_width, |
| 320 | patch_size, |
| 321 | self.normalize_lut, |
| 322 | ) |
| 323 | out_offset += num_patches |
| 324 | |
| 325 | data = { |
| 326 | "pixel_values": torch.from_numpy(pixel_values_np), |
| 327 | "grid_thws": torch.from_numpy(grid_thws_np), |
| 328 | } |
| 329 | return BatchFeature(data=data, tensor_type=return_tensors) |
| 330 | |
| 331 | def __repr__(self): |
| 332 | return f"KimiK25FusedVisionProcessor(media_proc_cfg={self.media_proc_cfg})" |
| 333 | |
| 334 | def to_dict(self) -> dict[str, Any]: |
| 335 | output = super().to_dict() |
| 336 | output["media_proc_cfg"] = self.media_proc_cfg |
| 337 | if "media_processor" in output: |
| 338 | del output["media_processor"] |
| 339 | return output |
| 340 | |
| 341 | @classmethod |
| 342 | def from_dict(cls, config_dict: dict[str, Any], **kwargs): |
| 343 | config = config_dict.copy() |
| 344 | media_proc_cfg = config.pop("media_proc_cfg", {}) |
| 345 | return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs) |
| 346 | |
| 347 | def to_json_string(self): |
| 348 | dictionary = self.to_dict() |
| 349 | for key, value in dictionary.items(): |
| 350 | if hasattr(value, "tolist"): |
| 351 | dictionary[key] = value.tolist() |
| 352 | return json.dumps(dictionary, indent=2, sort_keys=True) + "\n" |