image_processing_moonvit_v2.py

# Copyright 2025-2026 The Moonshot AI Team and HuggingFace Inc. team. All rights reserved.
#
# Extracted from moonshotai/Kimi-K3 for standalone MoonViT-V2 use.
# Licensed under the Kimi K3 License (see LICENSE in this repository).

"""Image / video processor for MoonViT-V2 (NaViT-style native resolution)."""

from __future__ import annotations

import math
from typing import Optional, Sequence, Union

import numpy as np
import torch
from PIL import Image
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
from transformers.image_utils import ImageInput, make_list_of_images, valid_images
from transformers.utils import TensorType


def _navit_resize(
    width: int,
    height: int,
    patch_size: int,
    merge_kernel_size: int,
    in_patch_limit: int,
    patch_limit_on_one_side: int,
):
    s1 = math.sqrt(
        in_patch_limit
        / (max(1.0, width // patch_size) * max(1.0, height // patch_size))
    )
    s2 = patch_limit_on_one_side * patch_size / width
    s3 = patch_limit_on_one_side * patch_size / height
    scale = min(1.0, s1, s2, s3)
    new_w = min(max(1, int(width * scale)), patch_limit_on_one_side * patch_size)
    new_h = min(max(1, int(height * scale)), patch_limit_on_one_side * patch_size)

    factor = merge_kernel_size * patch_size
    pad_height = (factor - new_h % factor) % factor
    pad_width = (factor - new_w % factor) % factor
    token_height = (new_h + pad_height) // factor
    token_width = (new_w + pad_width) // factor
    return {
        "new_width": new_w,
        "new_height": new_h,
        "pad_width": pad_width,
        "pad_height": pad_height,
        "num_tokens": token_height * token_width,
    }


def _patchify(pixel_values: np.ndarray, patch_size: int) -> dict:
    """pixel_values: (t, h, w, c) -> patches + grid_thw."""
    T, H, W, C = pixel_values.shape
    assert C == 3
    patches = pixel_values.reshape(
        T, H // patch_size, patch_size, W // patch_size, patch_size, C
    )
    patches = patches.transpose(0, 1, 3, 5, 2, 4)
    patches = patches.reshape(-1, C, patch_size, patch_size)
    grid_thw = np.array([T, H // patch_size, W // patch_size], dtype=np.int64)
    return {"pixel_values": patches, "grid_thw": grid_thw}


def _as_pil(image) -> Image.Image:
    if isinstance(image, Image.Image):
        return image.convert("RGB")
    return Image.fromarray(np.asarray(image)).convert("RGB")


class MoonViTV2ImageProcessor(BaseImageProcessor):
    model_type = "moonvit_v2"

    def __init__(
        self,
        patch_size: int = 14,
        merge_kernel_size: int = 2,
        in_patch_limit: int = 65536,
        patch_limit_on_one_side: int = 512,
        max_num_frames: int = 4,
        image_mean: tuple[float, float, float] = (0.5, 0.5, 0.5),
        image_std: tuple[float, float, float] = (0.5, 0.5, 0.5),
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.patch_size = patch_size
        self.merge_kernel_size = merge_kernel_size
        self.in_patch_limit = in_patch_limit
        self.patch_limit_on_one_side = patch_limit_on_one_side
        # Matches MoonViT-V2 init_pos_emb_time / temporal pos-emb capacity.
        self.max_num_frames = max_num_frames
        self.image_mean = list(image_mean)
        self.image_std = list(image_std)

    def _resize_config(self, image: Image.Image) -> dict:
        w, h = image.size
        return _navit_resize(
            w,
            h,
            self.patch_size,
            self.merge_kernel_size,
            self.in_patch_limit,
            self.patch_limit_on_one_side,
        )

    def _normalize_frame(self, image: Image.Image, cfg: dict) -> np.ndarray:
        image = image.resize(
            (cfg["new_width"], cfg["new_height"]), resample=Image.Resampling.BICUBIC
        )
        arr = np.asarray(image)
        if cfg["pad_height"] or cfg["pad_width"]:
            arr = np.pad(
                arr,
                ((0, cfg["pad_height"]), (0, cfg["pad_width"]), (0, 0)),
                mode="constant",
                constant_values=0,
            )
        mean = np.array(self.image_mean, dtype=np.float32)
        std_inv = 1.0 / np.array(self.image_std, dtype=np.float32)
        return (arr.astype(np.float32) / 255.0 - mean) * std_inv

    def _preprocess_one(self, image: Image.Image) -> tuple[np.ndarray, np.ndarray]:
        image = image.convert("RGB")
        cfg = self._resize_config(image)
        arr = self._normalize_frame(image, cfg)
        packed = _patchify(np.expand_dims(arr, axis=0), self.patch_size)
        return packed["pixel_values"], packed["grid_thw"]

    def preprocess(
        self,
        images: ImageInput,
        return_tensors: Optional[Union[str, TensorType]] = None,
    ) -> BatchFeature:
        """Preprocess one or more **images** (each becomes an independent sample with T=1)."""
        images = make_list_of_images(images)
        if not valid_images(images):
            raise ValueError(
                "Invalid image type. Must be PIL.Image.Image, numpy.ndarray, or torch.Tensor."
            )

        pixel_values, grid_thws = [], []
        for image in images:
            patches, grid_thw = self._preprocess_one(_as_pil(image))
            pixel_values.append(torch.from_numpy(patches))
            grid_thws.append(torch.from_numpy(grid_thw).unsqueeze(0))

        data = {
            "pixel_values": torch.cat(pixel_values, dim=0),
            "grid_thws": torch.cat(grid_thws, dim=0),
        }
        return BatchFeature(data=data, tensor_type=return_tensors)

    def preprocess_video(
        self,
        frames: Sequence[ImageInput],
        return_tensors: Optional[Union[str, TensorType]] = None,
    ) -> BatchFeature:
        """Preprocess a **video** as an ordered list of frames (one sample with T=len(frames)).

        All frames share the resize/pad config of the first frame so spatial grids align.
        ``T`` must be in ``[1, max_num_frames]`` (default 4, matching temporal pos-emb).
        """
        if not isinstance(frames, (list, tuple)) or len(frames) == 0:
            raise ValueError("`frames` must be a non-empty list/tuple of images.")
        if len(frames) > self.max_num_frames:
            raise ValueError(
                f"Got {len(frames)} frames, but max_num_frames={self.max_num_frames} "
                f"(MoonViT-V2 temporal pos-emb capacity)."
            )

        pil_frames = [_as_pil(f) for f in frames]
        cfg = self._resize_config(pil_frames[0])
        arrs = [self._normalize_frame(fr, cfg) for fr in pil_frames]
        pixels = np.stack(arrs, axis=0)  # (T, H, W, C)
        packed = _patchify(pixels, self.patch_size)

        data = {
            "pixel_values": torch.from_numpy(packed["pixel_values"]),
            "grid_thws": torch.from_numpy(packed["grid_thw"]).unsqueeze(0),
        }
        return BatchFeature(data=data, tensor_type=return_tensors)