| """Image and native clip-video processor for Kimi-K3 MoonViT3D.""" |
|
|
| import json |
| import math |
| 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 .media_utils import ( |
| MediaInput, |
| TransparentBgConfig, |
| VideoReader, |
| _to_tensor, |
| ensure_media_type, |
| estimate_motion_score, |
| image_to_np, |
| load_image_source, |
| navit_patchify, |
| navit_resize_image, |
| navit_resize_video_frame, |
| normalize, |
| sample_video_indices, |
| uniformly_limit_indices, |
| ) |
|
|
|
|
| def _env_int(names: tuple[str, ...], default: int) -> int: |
| for name in names: |
| value = os.getenv(name) |
| if value is not None: |
| return int(value) |
| return int(default) |
|
|
|
|
| class GroundingPIImageProcessor(BaseImageProcessor): |
| model_type = "groundingpi_backbone" |
|
|
| def __init__(self, media_proc_cfg: dict, **kwargs): |
| super().__init__(**kwargs) |
| self.media_proc_cfg = dict(media_proc_cfg) |
| self.patch_size = int(media_proc_cfg.get("patch_size", 14)) |
| self.merge_size = int(media_proc_cfg.get("merge_kernel_size", 2)) |
| image_max_tokens = os.getenv("IMAGE_MAX_TOKEN_NUM") |
| if image_max_tokens: |
| image_max_tokens = int(image_max_tokens) |
| if image_max_tokens < 1: |
| raise ValueError("IMAGE_MAX_TOKEN_NUM must be positive") |
| self.media_proc_cfg["in_patch_limit"] = ( |
| image_max_tokens * self.merge_size**2 |
| ) |
| self.max_output_tokens = ( |
| int(self.media_proc_cfg["in_patch_limit"]) // self.merge_size**2 |
| ) |
| self.video_clip_size = int( |
| os.getenv("GROUNDINGPI_VIDEO_CLIP_SIZE", media_proc_cfg.get("video_clip_size", 4)) |
| ) |
| if not 1 <= self.video_clip_size <= 4: |
| raise ValueError("Kimi-K3 MoonViT3D video_clip_size must be in [1, 4]") |
|
|
| self.video_default_fps = float(media_proc_cfg.get("video_default_fps", 2.0)) |
| self.video_train_fps = float( |
| os.getenv( |
| "GROUNDINGPI_VIDEO_TRAIN_FPS", |
| media_proc_cfg.get("video_train_fps", self.video_default_fps), |
| ) |
| ) |
| self.video_fps_candidates = tuple( |
| float(x) |
| for x in media_proc_cfg.get( |
| "video_fps_candidates", |
| media_proc_cfg.get("video_train_fps_choices", [1, 2, 4]), |
| ) |
| ) |
| if self.video_train_fps <= 0 or not self.video_fps_candidates: |
| raise ValueError("Video FPS configuration must contain positive values") |
| if any(fps <= 0 for fps in self.video_fps_candidates): |
| raise ValueError("video_fps_candidates must contain positive values") |
| self.video_high_motion_fps = float(media_proc_cfg.get("video_high_motion_fps", 4.0)) |
| self.video_auto_high_motion = bool(media_proc_cfg.get("video_auto_high_motion", True)) |
| self.video_high_motion_threshold = float( |
| media_proc_cfg.get("video_high_motion_threshold", 0.08) |
| ) |
| self.video_max_tokens = _env_int( |
| ("GROUNDINGPI_VIDEO_MAX_TOKENS", "VIDEO_MAX_TOKEN_NUM"), |
| media_proc_cfg.get("video_max_tokens", 4096), |
| ) |
| self.video_max_tokens_per_frame = _env_int( |
| ("GROUNDINGPI_VIDEO_MAX_TOKENS_PER_FRAME", "VIDEO_FRAME_MAX_TOKEN_NUM"), |
| media_proc_cfg.get("video_max_tokens_per_frame", 1024), |
| ) |
| self.video_max_frames = _env_int( |
| ("GROUNDINGPI_VIDEO_MAX_FRAMES", "VIDEO_MAX_FRAMES"), |
| media_proc_cfg.get( |
| "video_max_frames", |
| media_proc_cfg.get("max_num_frames_each_video", 64), |
| ), |
| ) |
| if self.video_max_frames < 1: |
| raise ValueError("video_max_frames must be positive") |
|
|
| def __call__(self, images=None, **kwargs): |
| """Accept image-only, video-only, or mixed processor calls.""" |
| return self.preprocess(images=images, **kwargs) |
|
|
| @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): |
| if media.get("type") == "video": |
| result = self.preprocess(videos=[media], return_tensors="pt") |
| grid = result["video_grid_thw"] |
| return int((grid[:, 1] * grid[:, 2] // self.merge_size**2).sum()) |
| media = ensure_media_type( |
| media, |
| transparent_bg_config=self._transparent_bg_config, |
| transparent_bg_fill_stage=self._transparent_bg_fill_stage, |
| ) |
| return self.get_resize_config(media)["num_tokens"] |
|
|
| @classmethod |
| def make_image_prompt(cls, width: int, height: int) -> str: |
| 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": |
| raise ValueError(f"Unsupported type: {media_input['type']}") |
| width, height = media_input["image"].size |
| input_patch_limit = int(self.media_proc_cfg["in_patch_limit"]) |
| while True: |
| result = navit_resize_image( |
| width, |
| height, |
| self.patch_size, |
| self.merge_size, |
| input_patch_limit, |
| self.media_proc_cfg["patch_limit_on_one_side"], |
| self.media_proc_cfg.get("fixed_output_tokens"), |
| ) |
| if result["num_tokens"] <= self.max_output_tokens: |
| return result |
| reduced_limit = int( |
| input_patch_limit * self.max_output_tokens / result["num_tokens"] |
| ) |
| input_patch_limit = min(input_patch_limit - 1, reduced_limit) |
|
|
| 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, |
| ) |
| return np.pad( |
| image_np, |
| ((0, pad_height), (0, pad_width), (0, 0)), |
| mode="constant", |
| constant_values=0, |
| ) |
|
|
| def _normalize_and_patchify(self, pixels: np.ndarray) -> dict[str, np.ndarray]: |
| image_std_inv = 1.0 / np.array(self.media_proc_cfg["image_std"]) |
| image_mean = np.array(self.media_proc_cfg["image_mean"]) |
| return navit_patchify( |
| normalize(pixels, image_mean, image_std_inv), |
| self.patch_size, |
| ) |
|
|
| def _preprocess_images(self, images) -> dict[str, torch.Tensor]: |
| if images is None: |
| return {} |
| if not isinstance(images, list): |
| images = [images] |
| patched = [] |
| for item in images: |
| if not isinstance(item, dict) and not isinstance(item, Image.Image): |
| item = load_image_source(item) |
| item = item if isinstance(item, dict) else {"type": "image", "image": item} |
| if not isinstance(item.get("image"), Image.Image): |
| item["image"] = load_image_source(item["image"]) |
| item = ensure_media_type( |
| item, |
| transparent_bg_config=self._transparent_bg_config, |
| transparent_bg_fill_stage=self._transparent_bg_fill_stage, |
| ) |
| resize = self.get_resize_config(item) |
| pixels = self.resize_image( |
| item["image"], |
| resize["new_width"], |
| resize["new_height"], |
| resize["pad_width"], |
| resize["pad_height"], |
| ) |
| patched.append(self._normalize_and_patchify(np.expand_dims(pixels, axis=0))) |
| if not patched: |
| return {} |
| return { |
| "pixel_values": torch.cat([_to_tensor(x["pixel_values"]) for x in patched]), |
| "image_grid_thw": torch.stack( |
| [_to_tensor(x["grid_thw"], dtype=torch.int64) for x in patched] |
| ), |
| } |
|
|
| def _choose_initial_fps( |
| self, |
| reader: VideoReader, |
| item: dict, |
| video_mode: str, |
| requested_fps: float | None, |
| high_motion: bool | None, |
| auto_high_motion: bool, |
| ) -> tuple[float, float | None, bool]: |
| item_fps = item.get("fps", requested_fps) |
| if item_fps is not None: |
| return float(item_fps), None, bool(item.get("high_motion", high_motion or False)) |
|
|
| if video_mode == "train": |
| return self.video_train_fps, None, False |
|
|
| item_high_motion = item.get("high_motion", high_motion) |
| motion_score = None |
| if item_high_motion is None and auto_high_motion: |
| probe_count = min(8, reader.frame_count) |
| probe_indices = np.linspace(0, reader.frame_count - 1, probe_count).round().astype(int).tolist() |
| motion_score = estimate_motion_score(reader.get_frames(probe_indices)) |
| item_high_motion = motion_score >= self.video_high_motion_threshold |
| item_high_motion = bool(item_high_motion) |
| fps = self.video_high_motion_fps if item_high_motion else self.video_default_fps |
| return fps, motion_score, item_high_motion |
|
|
| def _preprocess_one_video( |
| self, |
| raw_item, |
| *, |
| video_mode: str, |
| requested_fps: float | None, |
| high_motion: bool | None, |
| auto_high_motion: bool, |
| max_tokens: int, |
| max_tokens_per_frame: int, |
| max_frames: int, |
| ) -> tuple[list[dict[str, np.ndarray]], dict[str, Any]]: |
| item = dict(raw_item) if isinstance(raw_item, dict) else {"video": raw_item} |
| source = item.get("video", item.get("frames")) |
| if source is None: |
| raise ValueError("Video input dict requires a 'video' or 'frames' field") |
|
|
| item_max_tokens = int(item.get("max_tokens", max_tokens)) |
| item_frame_tokens = int(item.get("max_tokens_per_frame", max_tokens_per_frame)) |
| item_max_frames = int(item.get("max_frames", max_frames)) |
| if item_max_tokens < 1 or item_frame_tokens < 1: |
| raise ValueError("Video token limits must be positive") |
| if item_max_frames < 1: |
| raise ValueError("Video max_frames must be positive") |
| item_frame_tokens = min(item_frame_tokens, item_max_tokens) |
|
|
| source_fps = item.get("source_fps") |
| if source_fps is None and isinstance(source, (list, tuple)): |
| source_fps = item.get("fps", requested_fps or self.video_default_fps) |
| reader = VideoReader(source, source_fps=source_fps) |
| try: |
| initial_fps, motion_score, is_high_motion = self._choose_initial_fps( |
| reader, |
| item, |
| video_mode, |
| requested_fps, |
| high_motion, |
| auto_high_motion, |
| ) |
| first_frame = reader.get_frames([0])[0] |
| resize = navit_resize_video_frame( |
| first_frame.width, |
| first_frame.height, |
| self.patch_size, |
| self.merge_size, |
| int(self.media_proc_cfg.get("in_patch_limit_each_frame", 16384)), |
| int(self.media_proc_cfg["patch_limit_on_one_side"]), |
| item_frame_tokens, |
| ) |
| tokens_per_clip = int(resize["num_tokens"]) |
| max_clips = max(1, item_max_tokens // tokens_per_clip) |
|
|
| fps_candidates = sorted( |
| {float(initial_fps), *self.video_fps_candidates}, |
| reverse=True, |
| ) |
| fps_candidates = [fps for fps in fps_candidates if fps <= initial_fps] |
| if not fps_candidates: |
| fps_candidates = [initial_fps] |
|
|
| indices = None |
| selected_fps = fps_candidates[-1] |
| for candidate in fps_candidates: |
| candidate_indices = sample_video_indices( |
| reader.frame_count, reader.source_fps, candidate |
| ) |
| if math.ceil(len(candidate_indices) / self.video_clip_size) <= max_clips: |
| indices = candidate_indices |
| selected_fps = candidate |
| break |
| if indices is None: |
| indices = sample_video_indices( |
| reader.frame_count, reader.source_fps, fps_candidates[-1] |
| ) |
| indices = uniformly_limit_indices( |
| indices, min(item_max_frames, max_clips * self.video_clip_size) |
| ) |
| selected_fps = fps_candidates[-1] |
| else: |
| indices = uniformly_limit_indices(indices, item_max_frames) |
|
|
| frames = reader.get_frames(indices) |
| patched_clips = [] |
| clip_frame_counts = [] |
| for start in range(0, len(frames), self.video_clip_size): |
| clip = frames[start:start + self.video_clip_size] |
| clip_pixels = [ |
| self.resize_image( |
| frame, |
| resize["new_width"], |
| resize["new_height"], |
| resize["pad_width"], |
| resize["pad_height"], |
| ) |
| for frame in clip |
| ] |
| patched_clips.append(self._normalize_and_patchify(np.stack(clip_pixels))) |
| clip_frame_counts.append(len(clip)) |
|
|
| output_tokens = len(patched_clips) * tokens_per_clip |
| if output_tokens > item_max_tokens: |
| raise AssertionError( |
| f"Video token budget exceeded: {output_tokens} > {item_max_tokens}" |
| ) |
| metadata = { |
| "source_fps": float(reader.source_fps), |
| "duration_seconds": float(reader.duration), |
| "initial_fps": float(initial_fps), |
| "selected_fps": float(selected_fps), |
| "effective_fps": float(len(indices) / reader.duration), |
| "high_motion": is_high_motion, |
| "motion_score": motion_score, |
| "sampled_frames": len(indices), |
| "sampled_frame_indices": indices, |
| "clip_size": self.video_clip_size, |
| "clip_frame_counts": clip_frame_counts, |
| "num_clips": len(patched_clips), |
| "tokens_per_clip": tokens_per_clip, |
| "video_tokens": output_tokens, |
| "max_video_tokens": item_max_tokens, |
| "max_tokens_per_frame": item_frame_tokens, |
| "max_frames": item_max_frames, |
| "decode_backend": reader.backend, |
| "source_kind": reader.source_kind, |
| "source_bytes": reader.source_bytes, |
| "media_io_seconds": reader.io_seconds, |
| "resized_size": [ |
| resize["new_width"] + resize["pad_width"], |
| resize["new_height"] + resize["pad_height"], |
| ], |
| } |
| return patched_clips, metadata |
| finally: |
| reader.close() |
|
|
| def _preprocess_videos( |
| self, |
| videos, |
| *, |
| video_mode: str, |
| video_fps: float | None, |
| video_high_motion: bool | None, |
| video_auto_high_motion: bool, |
| video_max_tokens: int, |
| video_max_tokens_per_frame: int, |
| video_max_frames: int, |
| sampling_seed: int | None, |
| ) -> tuple[dict[str, torch.Tensor], list[dict[str, Any]]]: |
| if videos is None: |
| return {}, [] |
| if not isinstance(videos, list): |
| videos = [videos] |
| del sampling_seed |
| patched = [] |
| metadata = [] |
| clip_counts = [] |
| for item in videos: |
| item_patched, item_metadata = self._preprocess_one_video( |
| item, |
| video_mode=video_mode, |
| requested_fps=video_fps, |
| high_motion=video_high_motion, |
| auto_high_motion=video_auto_high_motion, |
| max_tokens=video_max_tokens, |
| max_tokens_per_frame=video_max_tokens_per_frame, |
| max_frames=video_max_frames, |
| ) |
| patched.extend(item_patched) |
| metadata.append(item_metadata) |
| clip_counts.append(len(item_patched)) |
| if not patched: |
| return {}, metadata |
| return { |
| "pixel_values_videos": torch.cat( |
| [_to_tensor(x["pixel_values"]) for x in patched] |
| ), |
| "video_grid_thw": torch.stack( |
| [_to_tensor(x["grid_thw"], dtype=torch.int64) for x in patched] |
| ), |
| "video_clip_counts": torch.tensor(clip_counts, dtype=torch.int64), |
| }, metadata |
|
|
| def preprocess( |
| self, |
| medias: Optional[list[MediaInput]] = None, |
| return_tensors: Optional[Union[str, TensorType]] = None, |
| images=None, |
| videos=None, |
| do_resize: Optional[bool] = None, |
| video_mode: str = "inference", |
| video_fps: float | None = None, |
| video_high_motion: bool | None = None, |
| video_auto_high_motion: bool | None = None, |
| video_max_tokens: int | None = None, |
| video_max_tokens_per_frame: int | None = None, |
| video_max_frames: int | None = None, |
| video_sampling_seed: int | None = None, |
| return_video_metadata: bool = False, |
| **kwargs, |
| ) -> BatchFeature: |
| del do_resize, kwargs |
| if video_mode not in {"train", "inference"}: |
| raise ValueError("video_mode must be 'train' or 'inference'") |
| if medias is not None and images is None and videos is None: |
| medias = medias if isinstance(medias, list) else [medias] |
| images = [m for m in medias if not isinstance(m, dict) or m.get("type", "image") == "image"] |
| videos = [m for m in medias if isinstance(m, dict) and m.get("type") == "video"] |
|
|
| data = self._preprocess_images(images) |
| video_data, metadata = self._preprocess_videos( |
| videos, |
| video_mode=video_mode, |
| video_fps=video_fps, |
| video_high_motion=video_high_motion, |
| video_auto_high_motion=( |
| self.video_auto_high_motion |
| if video_auto_high_motion is None |
| else bool(video_auto_high_motion) |
| ), |
| video_max_tokens=( |
| self.video_max_tokens if video_max_tokens is None else int(video_max_tokens) |
| ), |
| video_max_tokens_per_frame=( |
| self.video_max_tokens_per_frame |
| if video_max_tokens_per_frame is None |
| else int(video_max_tokens_per_frame) |
| ), |
| video_max_frames=( |
| self.video_max_frames |
| if video_max_frames is None |
| else int(video_max_frames) |
| ), |
| sampling_seed=video_sampling_seed, |
| ) |
| data.update(video_data) |
| batch = BatchFeature(data=data, tensor_type=return_tensors) |
| if return_video_metadata: |
| batch["video_metadata"] = metadata |
| return batch |
|
|
| def __repr__(self): |
| return f"GroundingPIImageProcessor(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 |
| output.pop("media_processor", None) |
| 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" |
|
|
