chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Add Mistral Small 4 (Pixtral) support (#20708)

Browse Source 
Signed-off-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: Alex Nails <alexnails@radixark.ai>
Co-authored-by: Dimitrios Bariamis <12195802+dbari@users.noreply.github.com>
Co-authored-by: dbari <dbari@users.noreply.github.com>
This commit is contained in:
Xinyuan Tong
2026-03-18 21:15:32 +00:00
committed by GitHub
parent df1d046de2
commit 6b8a6545b2
18 changed files with 360 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

113
python/sglang/srt/multimodal/processors/pixtral.py
View File

@@ -1,11 +1,12 @@
import asyncio
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
from sglang.srt.models.pixtral import (
PixtralForConditionalGeneration,
PixtralVisionModel,
@@ -20,63 +21,47 @@ class PixtralProcessor(BaseMultimodalProcessor):
models = [PixtralVisionModel, PixtralForConditionalGeneration]
PAD_TOKEN = "<pad>"
IMG_BREAK_TOKEN_ID = 12
IMG_END_TOKEN_ID = 13
def get_patch_grid_size(
self,
*,
image_width: int,
image_height: int,
) -> tuple[int, int]:
max_width = max_height = self.image_size
patch_width = patch_height = self.patch_size
ratio = max(image_width / max_width, image_height / max_height)
if ratio > 1:
image_width = int(math.floor(image_width / ratio))
image_height = int(math.floor(image_height / ratio))
nrows, ncols = _get_pixtral_hf_num_image_tokens(
(image_height, image_width),
(patch_height, patch_width),
)
return ncols, nrows
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
)
# Instantiate the patcher logic helper using the class defined above
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 hasattr(self.vision_config, "spatial_merge_size"):
self._processor.spatial_merge_size = self.vision_config.spatial_merge_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=_processor.image_token,
image_token=self.image_token,
image_token_id=self.IM_TOKEN_ID,
).build(_processor)
_processor.tokenizer.add_special_tokens(
tokenizer.add_special_tokens(
{
"pad_token": getattr(hf_config, "pad_token", self.PAD_TOKEN),
}
)
async def _resize(self, image):
num_w_tokens, num_h_tokens = self.get_patch_grid_size(
image_width=image.size[0],
image_height=image.size[1],
)
new_size = (num_w_tokens * self.patch_size, num_h_tokens * self.patch_size)
return image.resize(new_size)
async def process_mm_data_async(
self,
image_data: List[Union[str, bytes]],
@@ -92,16 +77,58 @@ class PixtralProcessor(BaseMultimodalProcessor):
return_text=True,
)
if mm_data.images:
resize_tasks = [self._resize(image) for image in mm_data.images]
mm_data.images = await asyncio.gather(*resize_tasks)
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
)
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 {
"mm_items": mm_items,
"input_ids": input_ids.tolist(),
"im_token_id": self.IM_TOKEN_ID,
"im_token": self._processor.image_token,
"im_token": self.image_token,
}