"""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 # Kept in the public API for backward compatibility. 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"