chore: import upstream snapshot with attribution

python/sglang/srt/multimodal/processors/pixtral.py

@@ -0,0 +1,134 @@
+import math
+from typing import List, Union
+
+from transformers import PreTrainedTokenizerBase
+from transformers.models.pixtral.image_processing_pixtral import (
+    _num_image_tokens as _get_pixtral_hf_num_image_tokens,
+)
+
+from sglang.srt.managers.schedule_batch import Modality, MultimodalProcessorOutput
+from sglang.srt.models.pixtral import (
+    PixtralForConditionalGeneration,
+    PixtralVisionModel,
+)
+from sglang.srt.multimodal.processors.base_processor import (
+    BaseMultimodalProcessor,
+    MultimodalSpecialTokens,
+)
+
+
+class PixtralProcessor(BaseMultimodalProcessor):
+    models = [PixtralVisionModel, PixtralForConditionalGeneration]
+    gpu_image_decode = False  # Pixtral processes loaded image as PIL image explicitly
+
+    PAD_TOKEN = "<pad>"
+    DEFAULT_IMAGE_TOKEN = "[IMG]"
+
+    def __init__(self, hf_config, server_args, _processor, *args, **kwargs):
+        super().__init__(hf_config, server_args, _processor, *args, **kwargs)
+        self.IM_TOKEN_ID = getattr(
+            hf_config, "image_token_index", PixtralVisionModel.DEFAULT_IMAGE_TOKEN_ID
+        )
+
+        self.vision_config = hf_config.vision_config
+        self.image_size = self.vision_config.image_size
+        self.patch_size = self.vision_config.patch_size
+
+        # spatial_merge_size may live on vision_config (Mistral native) or
+        # on the top-level config (HF native Mistral3Config).
+        self._spatial_merge_size = getattr(
+            self.vision_config,
+            "spatial_merge_size",
+            getattr(hf_config, "spatial_merge_size", 1),
+        )
+
+        self._processor.patch_size = self.patch_size
+        if self._spatial_merge_size > 1:
+            self._processor.spatial_merge_size = self._spatial_merge_size
+
+        tokenizer = (
+            _processor
+            if isinstance(_processor, PreTrainedTokenizerBase)
+            else _processor.tokenizer
+        )
+        self.image_token = getattr(_processor, "image_token", self.DEFAULT_IMAGE_TOKEN)
+
+        self.mm_tokens = MultimodalSpecialTokens(
+            image_token=self.image_token,
+            image_token_id=self.IM_TOKEN_ID,
+        ).build(_processor)
+        tokenizer.add_special_tokens(
+            {
+                "pad_token": getattr(hf_config, "pad_token", self.PAD_TOKEN),
+            }
+        )
+
+    async def process_mm_data_async(
+        self,
+        image_data: List[Union[str, bytes]],
+        input_text,
+        request_obj,
+        *args,
+        **kwargs,
+    ):
+        mm_data = await self.load_mm_data(
+            prompt=input_text,
+            multimodal_tokens=self.mm_tokens,
+            image_data=image_data,
+            return_text=True,
+        )
+        if mm_data.images:
+            effective_patch = self.patch_size * self._spatial_merge_size
+            image_nrows = []
+            for img in mm_data.images:
+                w, h = img.size
+                ratio = max(w / self.image_size, h / self.image_size)
+                if ratio > 1:
+                    w = int(math.floor(w / ratio))
+                    h = int(math.floor(h / ratio))
+                nrows, _ = _get_pixtral_hf_num_image_tokens(
+                    (h, w), (effective_patch, effective_patch)
+                )
+                image_nrows.append(nrows)
+
+            mm_items, input_ids, _ = self.process_and_combine_mm_data(
+                mm_data, self.mm_tokens
+            )
+
+            # For multi-image: split single IMAGE mm_item into per-image items
+            if len(mm_data.images) > 1:
+                from sglang.srt.managers.schedule_batch import MultimodalDataItem
+
+                old_item = next(
+                    item for item in mm_items if item.modality == Modality.IMAGE
+                )
+                all_offsets = old_item.offsets
+                old_feature = old_item.feature
+                old_image_sizes = getattr(old_item, "image_sizes", None)
+
+                mm_items = [
+                    item for item in mm_items if item.modality != Modality.IMAGE
+                ]
+                offset_idx = 0
+                for i, img in enumerate(mm_data.images):
+                    nr = image_nrows[i]
+                    item_offsets = all_offsets[offset_idx : offset_idx + nr]
+                    offset_idx += nr
+                    new_item = MultimodalDataItem(modality=Modality.IMAGE)
+                    new_item.feature = old_feature[i : i + 1]
+                    new_item.offsets = item_offsets
+                    if old_image_sizes is not None:
+                        new_item.model_specific_data["image_sizes"] = old_image_sizes[
+                            i : i + 1
+                        ]
+                    mm_items.append(new_item)
+        else:
+            mm_items, input_ids, _ = self.process_and_combine_mm_data(
+                mm_data, self.mm_tokens
+            )
+
+        return MultimodalProcessorOutput(
+            mm_items=mm_items,
+            input_ids=input_ids.tolist(),
+            im_token_id=self.IM_TOKEN_ID,
+        )