nvidia/MiniMax-M3-NVFP4

1# Copyright 2023-2024 SGLang Team2# Licensed under the Apache License, Version 2.0 (the "License");3"""4MiniMax VL family HuggingFace-compatible VideoProcessor.5"""6 7import math8from typing import List, Optional, Tuple, Union9 10import torch11import torchvision12from torchvision.transforms import InterpolationMode13from transformers import BatchFeature14from transformers.image_utils import PILImageResampling, SizeDict15from transformers.processing_utils import (16    Unpack,17    VideosKwargs,18)19from transformers.utils import TensorType20from transformers.video_processing_utils import BaseVideoProcessor21from transformers.video_utils import group_videos_by_shape, reorder_videos22 23MAX_RATIO = 20024 25 26def round_by_factor(number: int, factor: int) -> int:27    return round(number / factor) * factor28 29 30def ceil_by_factor(number: int, factor: int) -> int:31    return math.ceil(number / factor) * factor32 33 34def floor_by_factor(number: int, factor: int) -> int:35    return math.floor(number / factor) * factor36 37 38def smart_resize(39    height: int,40    width: int,41    factor: int = 28,42    min_pixels: int = 4 * 28 * 28,43    max_pixels: int = 451584,44) -> tuple[int, int]:45    if max(height, width) / min(height, width) > MAX_RATIO:46        raise ValueError(47            f"absolute aspect ratio must be smaller than {MAX_RATIO}, "48            f"got {max(height, width) / min(height, width)}"49        )50    h_bar = max(factor, round_by_factor(height, factor))51    w_bar = max(factor, round_by_factor(width, factor))52    if h_bar * w_bar > max_pixels:53        beta = math.sqrt((height * width) / max_pixels)54        h_bar = floor_by_factor(height / beta, factor)55        w_bar = floor_by_factor(width / beta, factor)56    elif h_bar * w_bar < min_pixels:57        beta = math.sqrt(min_pixels / (height * width))58        h_bar = ceil_by_factor(height * beta, factor)59        w_bar = ceil_by_factor(width * beta, factor)60    return h_bar, w_bar61 62 63class MiniMaxM3VLVideoProcessorKwargs(VideosKwargs, total=False):64    patch_size: int65    temporal_patch_size: int66    merge_size: int67    min_pixels: int68    max_pixels: int69    total_pixels: int70    min_frames: int71    max_frames: int72    fps: float | int73 74 75class MiniMaxM3VLVideoProcessor(BaseVideoProcessor):76    do_resize = True77    resample = PILImageResampling.BICUBIC78    size = {"height": 672, "width": 672}79    default_to_square = False80    do_rescale = True81    rescale_factor = 1 / 25582    do_normalize = True83    image_mean = [0.48145466, 0.4578275, 0.40821073]84    image_std = [0.26862954, 0.26130258, 0.27577711]85    do_convert_rgb = True86    do_sample_frames = False87    patch_size = 1488    temporal_patch_size = 289    merge_size = 290    min_pixels = 4 * 28 * 2891    max_pixels = 768 * 28 * 28                  # 602,11292    total_pixels = int(64000 * 28 * 28 * 0.9)   # ~45M, ~64k tokens budget93    fps = 1.094    min_frames = 495    max_frames = 76896    valid_kwargs = MiniMaxM3VLVideoProcessorKwargs97    model_input_names = ["pixel_values_videos", "video_grid_thw"]98 99    def __init__(self, **kwargs: Unpack[MiniMaxM3VLVideoProcessorKwargs]):100        super().__init__(**kwargs)101 102    def _preprocess(103        self,104        videos: List[torch.Tensor],105        do_convert_rgb: bool,106        do_resize: bool,107        size: SizeDict,108        resample: PILImageResampling | InterpolationMode | int | None,109        do_rescale: bool,110        rescale_factor: float,111        do_normalize: bool,112        image_mean: float | List[float] | None,113        image_std: float | List[float] | None,114        patch_size: int,115        temporal_patch_size: int,116        merge_size: int,117        min_pixels: int,118        max_pixels: int,119        return_tensors: str | TensorType | None = None,120        **kwargs,121    ) -> BatchFeature:122        grouped_videos, grouped_videos_index = group_videos_by_shape(videos)123        resized_videos_grouped = {}124        factor = patch_size * merge_size125        for shape, stacked_videos in grouped_videos.items():126            batch_size, num_frames, channels, height, width = stacked_videos.shape127            resized_height, resized_width = height, width128            if do_resize:129                resized_height, resized_width = smart_resize(130                    height, width, factor=factor,131                    min_pixels=min_pixels, max_pixels=max_pixels,132                )133                stacked_videos = stacked_videos.view(134                    batch_size * num_frames, channels, height, width135                )136                stacked_videos = self.resize(137                    stacked_videos,138                    size=SizeDict(height=resized_height, width=resized_width),139                    resample=resample,140                )141                stacked_videos = stacked_videos.view(142                    batch_size,143                    num_frames,144                    channels,145                    resized_height,146                    resized_width,147                )148            resized_videos_grouped[shape] = stacked_videos149        resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index)150 151        grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos)152        processed_videos_grouped = {}153        processed_grids = {}154        for shape, stacked_videos in grouped_videos.items():155            resized_height, resized_width = stacked_videos.shape[-2:]156            patches = self.rescale_and_normalize(157                stacked_videos,158                do_rescale,159                rescale_factor,160                do_normalize,161                image_mean,162                image_std,163            )164 165            if pad := -patches.shape[1] % temporal_patch_size:166                repeats = patches[:, -1:].expand(-1, pad, -1, -1, -1)167                patches = torch.cat([patches, repeats], dim=1)168 169            batch_size, grid_t, channels = patches.shape[:3]170            grid_t = grid_t // temporal_patch_size171            grid_h, grid_w = resized_height // patch_size, resized_width // patch_size172 173            patches = patches.view(174                batch_size,175                grid_t,176                temporal_patch_size,177                channels,178                grid_h // merge_size,179                merge_size,180                patch_size,181                grid_w // merge_size,182                merge_size,183                patch_size,184            )185            patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)186            flatten_patches = patches.reshape(187                batch_size,188                grid_t * grid_h * grid_w,189                channels * temporal_patch_size * patch_size * patch_size,190            )191 192            processed_videos_grouped[shape] = flatten_patches193            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size194 195        processed_videos = reorder_videos(196            processed_videos_grouped, grouped_videos_index197        )198        processed_grids = reorder_videos(processed_grids, grouped_videos_index)199        pixel_values_videos = torch.cat(processed_videos, dim=0)200        video_grid_thw = torch.tensor(processed_grids, dtype=torch.long)201 202        return BatchFeature(203            data={204                "pixel_values_videos": pixel_values_videos,205                "video_grid_thw": video_grid_thw,206            },207            tensor_type=return_tensors,208        )209