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 Processor, ImageProcessor, VideoProcessor.5"""6 7import math8import re9from typing import List, Optional, Tuple, Union10 11import torch12import torchvision13from torchvision.transforms import InterpolationMode14from transformers import BatchFeature15from transformers.image_processing_utils_fast import (16    BaseImageProcessorFast,17    group_images_by_shape,18    reorder_images,19)20from transformers.image_utils import PILImageResampling, SizeDict21from transformers.processing_utils import (22    ImagesKwargs,23    ProcessingKwargs,24    ProcessorMixin,25    Unpack,26    VideosKwargs,27)28from transformers.utils import TensorType29from transformers.video_processing_utils import BaseVideoProcessor30from transformers.video_utils import group_videos_by_shape, reorder_videos31 32 33class MiniMaxVLProcessorKwargs(ProcessingKwargs, total=False):34    _defaults = {35        "videos_kwargs": {36            "do_resize": False,37            "return_metadata": True,38        },39    }40 41 42class MiniMaxVLProcessor(ProcessorMixin):43    IMAGE_TOKEN = "]<]image[>["44    VIDEO_TOKEN = "]<]video[>["45    VISION_START_TOKEN = "]<]start of image[>["46    VISION_END_TOKEN = "]<]end of image[>["47 48    def __init__(49        self, image_processor=None, tokenizer=None, video_processor=None, **kwargs50    ):51        self.image_token_id = tokenizer.convert_tokens_to_ids(self.IMAGE_TOKEN)52        self.video_token_id = tokenizer.convert_tokens_to_ids(self.VIDEO_TOKEN)53        super().__init__(image_processor, tokenizer, video_processor)54        # Video expansion also uses image start/end tokens. Separate video55        # start/end tokens exist in the tokenizer, but the original MiniMax56        # serving path did not use them; keep that behavior for compatibility.57        self.vision_start_token_id = tokenizer.convert_tokens_to_ids(58            self.VISION_START_TOKEN59        )60        self.vision_end_token_id = tokenizer.convert_tokens_to_ids(61            self.VISION_END_TOKEN62        )63 64    def _prune_video_tokens(65        self,66        input_text: str,67        video_segments: List[int],68        video_token: str,69    ) -> str:70        """71        Prune video tokens by temporal_patch_size (e.g., 2:1).72 73        Expects the prompt to carry exactly sum(video_segments) video74        tokens — i.e. one token per *sampled* frame. Then drops token.75 76        Args:77            input_text: prompt with N video_tokens per segment78            video_segments: actual sampled frame count per video segment79            video_token: the video token string, e.g. ']<]video[>['80 81        Returns:82            Pruned input_text with ~N/temporal_patch_size tokens per segment.83        """84        # If no videos or temporal_patch_size <= 1, no pruning needed85        if not video_segments or self.video_processor.temporal_patch_size <= 1:86            return input_text87 88        # Split while keeping delimiters89        special_tokens = [video_token]  # , image_token]90        pattern = "|".join(map(re.escape, special_tokens))91        parts = re.split(f"({pattern})", input_text)92 93        def is_timestamp(text: str) -> bool:94            """Check if text ends with timestamp format like ']<]0.0 seconds[>['"""95            return (96                text.endswith("seconds[>[")97                or text.endswith("seconds[>[ ")98                or text.endswith("seconds [>[")99                or text.endswith("seconds [>[ ")100            )101 102        def extract_timestamp(text: str) -> str:103            """Extract timestamp text from the end, starting from ']<]'"""104            start_index = text.rfind("]<]")105            if start_index == -1:106                raise ValueError(f"Failed to extract timestamp: {text}")107            return text[start_index:]108 109        # Build new text with pruned video tokens110        final_parts = []111        current_seg_idx = 0  # Which video segment we're in112        frame_in_seg = 0  # Frame index within current segment113        last_timestamp_len = 0  # Length of timestamp to potentially remove114 115        for part in parts:116            if part == video_token:117                if current_seg_idx < len(video_segments):118                    if frame_in_seg % self.video_processor.temporal_patch_size == 0:119                        # Keep this video token120                        final_parts.append(part)121                        frame_in_seg += 1122                        if frame_in_seg >= video_segments[current_seg_idx]:123                            current_seg_idx += 1124                            frame_in_seg = 0125                        last_timestamp_len = 0126                    else:127                        # Skip this video token128                        frame_in_seg += 1129                        if frame_in_seg >= video_segments[current_seg_idx]:130                            current_seg_idx += 1131                            frame_in_seg = 0132                        # Remove the timestamp that was already appended133                        if last_timestamp_len > 0:134                            # Truncate the last part to remove timestamp135                            assert len(final_parts) > 0136                            final_parts[-1] = final_parts[-1][:-last_timestamp_len]137                            last_timestamp_len = 0138                else:139                    # No more video segments, keep as is140                    final_parts.append(part)141                    last_timestamp_len = 0142            else:143                # Text part144                final_parts.append(part)145                # Check if this text ends with a timestamp146                if is_timestamp(part):147                    last_timestamp_len = len(extract_timestamp(part))148                else:149                    last_timestamp_len = 0150 151        return "".join(final_parts)152 153    def __call__(154        self,155        images=None,156        text=None,157        videos=None,158        **kwargs: Unpack[MiniMaxVLProcessorKwargs],159    ) -> BatchFeature:160        output_kwargs = self._merge_kwargs(161            MiniMaxVLProcessorKwargs,162            tokenizer_init_kwargs=self.tokenizer.init_kwargs,163            **kwargs,164        )165 166        if images is not None:167            images_kwargs = output_kwargs["images_kwargs"]168            image_inputs = self.image_processor(images=images, **images_kwargs)169            image_grid_thw = image_inputs["image_grid_thw"]170 171        else:172            image_inputs = {}173            image_grid_thw = None174 175        if videos is not None:176            videos_kwargs = output_kwargs["videos_kwargs"]177            video_inputs = self.video_processor(videos=videos, **videos_kwargs)178            video_grid_thw = video_inputs["video_grid_thw"]179            if not kwargs.get("return_metadata"):180                video_metadata = video_inputs.pop("video_metadata")181            else:182                video_metadata = video_inputs["video_metadata"]183        else:184            video_inputs = {}185            video_grid_thw = None186 187        if not isinstance(text, list):188            text = [text]189        text = text.copy()190 191        # Expand image tokens192        if image_grid_thw is not None:193            merge_length = self.image_processor.merge_size**2194            placeholder = "]<]placeholder[>["195            index = 0196            for i in range(len(text)):197                while self.IMAGE_TOKEN in text[i]:198                    num_tokens = image_grid_thw[index].prod() // merge_length199                    text[i] = text[i].replace(200                        self.IMAGE_TOKEN,201                        self.VISION_START_TOKEN202                        + placeholder * num_tokens203                        + self.VISION_END_TOKEN,204                        1,205                    )206                    index += 1207                text[i] = text[i].replace(placeholder, self.IMAGE_TOKEN)208 209        # Expand video tokens210        if video_grid_thw is not None:211            merge_length = self.image_processor.merge_size**2212            placeholder = "]<]placeholder[>["213            index = 0214            for i in range(len(text)):215                while self.VIDEO_TOKEN in text[i]:216                    metadata = video_metadata[index]217                    grid_t = video_grid_thw[index][0]218                    frame_seqlen = video_grid_thw[index][1:].prod() // merge_length219 220                    video_placeholder = ""221                    for frame_idx in range(grid_t):222                        if (223                            metadata.fps is not None224                            and metadata.frames_indices is not None225                        ):226                            ts = (227                                metadata.frames_indices[228                                    min(229                                        frame_idx230                                        * self.video_processor.temporal_patch_size,231                                        len(metadata.frames_indices) - 1,232                                    )233                                ]234                                / metadata.fps235                            )236                            video_placeholder += f"]<]{ts:.1f} seconds[>["237                        video_placeholder += (238                            self.VISION_START_TOKEN239                            + placeholder * frame_seqlen240                            + self.VISION_END_TOKEN241                        )242 243                    text[i] = text[i].replace(self.VIDEO_TOKEN, video_placeholder, 1)244                    index += 1245                text[i] = text[i].replace(placeholder, self.VIDEO_TOKEN)246 247        # Tokenize248        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)249        text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])250 251        return BatchFeature(252            data={**text_inputs, **image_inputs, **video_inputs},253            tensor_type=return_tensors,254        )255