sglang 0.4.4.post2__py3-none-any.whl → 0.4.4.post4__py3-none-any.whl

sglang/srt/models/deepseek_vl2.py

@@ -11,7 +11,11 @@ from sglang.srt.configs.deepseekvl2 import (
 )
 from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
-from sglang.srt.managers.schedule_batch import MultimodalInputs
+from sglang.srt.managers.mm_utils import (
+    MultiModalityDataPaddingPatternImageTokens,
+    general_mm_embed_routine,
+)
+from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
@@ -150,7 +154,6 @@ class DeepseekVL2MlpProjector(nn.Module):
         return x
 
 
-# todo
 class DeepseekVL2ForCausalLM(nn.Module):
 
     def __init__(
@@ -215,32 +218,15 @@ class DeepseekVL2ForCausalLM(nn.Module):
         forward_batch: ForwardBatch,
         **kwargs: object,
     ):
-        input_embeds = self.language_model.model.embed_tokens(input_ids)
-        if (
-            forward_batch.forward_mode.is_extend()
-            and forward_batch.contains_image_inputs()
-        ):
-            extend_start_loc_cpu = forward_batch.extend_start_loc.cpu().numpy()
-            extend_seq_lens_cpu = forward_batch.extend_seq_lens.cpu().numpy()
-            for idx, image in enumerate(forward_batch.mm_inputs):
-                if image is None:
-                    continue
-                start_idx = extend_start_loc_cpu[idx]
-                end_idx = start_idx + extend_seq_lens_cpu[idx]
-                images_emb_mask = image.images_emb_mask.to(device="cuda")
-                image_features = self.get_image_feature(image)
-                input_embeds[start_idx:end_idx] = input_embeds[
-                    start_idx:end_idx
-                ].masked_scatter(images_emb_mask.unsqueeze(-1), image_features)
-
-        outputs = self.language_model.forward(
+        hs = general_mm_embed_routine(
             input_ids=input_ids,
             positions=positions,
             forward_batch=forward_batch,
-            input_embeds=input_embeds,
+            image_data_embedding_func=self.get_image_feature,
+            language_model=self.language_model,
         )
 
-        return outputs
+        return hs
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         stacked_params_mapping = [
@@ -263,94 +249,109 @@ class DeepseekVL2ForCausalLM(nn.Module):
                 weights_loader(param, loaded_weight)
 
     def pad_input_ids(self, input_ids: List[int], image_inputs: MultimodalInputs):
-        return input_ids
-
-    def get_image_feature(self, image_input: MultimodalInputs):
-        pixel_values = image_input.pixel_values.type(
-            next(self.vision.parameters()).dtype
-        ).to(device=next(self.vision.parameters()).device)
-        image_feature = self.vision.forward_features(pixel_values)
-        images_embeds = self.projector(image_feature)
-        _, hw, n_dim = images_embeds.shape
-        h = w = int(hw**0.5)
-        tile_index = 0
+        helper = MultiModalityDataPaddingPatternImageTokens(
+            image_token_id=image_inputs.im_token_id
+        )
+        return helper.pad_input_tokens(input_ids, image_inputs)
+
+    def get_image_feature(self, items: List[MultimodalDataItem]):
+
+        images_spatial_crop = torch.cat(
+            [item.image_spatial_crop for item in items], dim=0
+        )
+
+        assert images_spatial_crop.dim() == 3
+
+        # TODO: can it be batched ?
         images_in_this_batch = []
-        images_spatial_crop = image_input.image_spatial_crop
-        for jdx in range(images_spatial_crop.shape[1]):
-            num_width_tiles, num_height_tiles = images_spatial_crop[0, jdx]
-            if num_width_tiles == 0 or num_height_tiles == 0:
-                break
-            num_tiles_in_image = num_width_tiles * num_height_tiles
-
-            # [hw, D]
-            global_features = images_embeds[tile_index]
-
-            # [num_height_tiles * num_width_tiles, hw, D]
-            local_features = images_embeds[
-                tile_index + 1 : tile_index + 1 + num_tiles_in_image
-            ]
-            tile_index += num_tiles_in_image + 1
-
-            # format global and local features
-            # ----------------- global view add newline -----------------
-            # [hw, D] -> [h, w, D]
-            global_features = global_features.view(h, w, n_dim)
-
-            # [D]     -> [h, 1, D]
-            new_lines_in_global = repeat(self.image_newline, "d -> h 1 d", h=h)
-
-            # cat([h, w, D], [h, 1, D], dim=1) -> [h, w + 1, D]
-            global_features = torch.cat([global_features, new_lines_in_global], dim=1)
-
-            # [h, w + 1, D] -> [h * (w + 1), D]
-            global_features = global_features.view(-1, n_dim)
-
-            # ----------------- local view add newline -----------------
-            # [num_height_tiles * num_width_tiles, h * w, D] ->
-            # [num_height_tiles * h, num_width_tiles * w, D]
-            local_features = rearrange(
-                local_features,
-                "(th tw) (h w) d -> (th h) (tw w) d",
-                th=num_height_tiles,
-                tw=num_width_tiles,
-                h=h,
-                w=w,
+        for item in items:
+            assert item.pixel_values.dim() == 4
+            image_feature = self.vision.forward_features(
+                item.pixel_values.type(next(self.vision.parameters()).dtype).to(
+                    device=next(self.vision.parameters()).device
+                )
             )
+            images_embeds = self.projector(image_feature)
+            _, hw, n_dim = images_embeds.shape
+            h = w = int(hw**0.5)
+            tile_index = 0
+            for jdx in range(item.image_spatial_crop.shape[1]):
+                num_width_tiles, num_height_tiles = item.image_spatial_crop[0, jdx]
+                if num_width_tiles == 0 or num_height_tiles == 0:
+                    break
+                num_tiles_in_image = num_width_tiles * num_height_tiles
+
+                # [hw, D]
+                global_features = images_embeds[tile_index]
+
+                # [num_height_tiles * num_width_tiles, hw, D]
+                local_features = images_embeds[
+                    tile_index + 1 : tile_index + 1 + num_tiles_in_image
+                ]
+                tile_index += num_tiles_in_image + 1
 
-            # [D] -> [num_height_tiles * h, 1, D]
-            new_lines_in_local = repeat(
-                self.image_newline,
-                "d -> (th h) 1 d",
-                th=num_height_tiles,
-                h=h,
-            )
+                # format global and local features
+                # ----------------- global view add newline -----------------
+                # [hw, D] -> [h, w, D]
+                global_features = global_features.view(h, w, n_dim)
+
+                # [D]     -> [h, 1, D]
+                new_lines_in_global = repeat(self.image_newline, "d -> h 1 d", h=h)
 
-            # [num_height_tiles * h, num_width_tiles * w + 1, D]
-            local_features = torch.cat([local_features, new_lines_in_local], dim=1)
-
-            # [num_height_tiles * h, num_width_tiles * w + 1, D]
-            #   --> [(num_height_tiles * h) * (num_width_tiles * w + 1), D]
-            local_features = local_features.view(-1, n_dim)
-
-            # merge global and local tiles
-            if self.global_view_pos == "head":
-                global_local_features = torch.cat(
-                    [
-                        global_features,
-                        self.view_seperator[None, :],
-                        local_features,
-                    ]
+                # cat([h, w, D], [h, 1, D], dim=1) -> [h, w + 1, D]
+                global_features = torch.cat(
+                    [global_features, new_lines_in_global], dim=1
                 )
-            else:
-                global_local_features = torch.cat(
-                    [
-                        local_features,
-                        self.view_seperator[None, :],
-                        global_features,
-                    ]
+
+                # [h, w + 1, D] -> [h * (w + 1), D]
+                global_features = global_features.view(-1, n_dim)
+
+                # ----------------- local view add newline -----------------
+                # [num_height_tiles * num_width_tiles, h * w, D] ->
+                # [num_height_tiles * h, num_width_tiles * w, D]
+                local_features = rearrange(
+                    local_features,
+                    "(th tw) (h w) d -> (th h) (tw w) d",
+                    th=num_height_tiles,
+                    tw=num_width_tiles,
+                    h=h,
+                    w=w,
                 )
 
-            images_in_this_batch.append(global_local_features)
+                # [D] -> [num_height_tiles * h, 1, D]
+                new_lines_in_local = repeat(
+                    self.image_newline,
+                    "d -> (th h) 1 d",
+                    th=num_height_tiles,
+                    h=h,
+                )
+
+                # [num_height_tiles * h, num_width_tiles * w + 1, D]
+                local_features = torch.cat([local_features, new_lines_in_local], dim=1)
+
+                # [num_height_tiles * h, num_width_tiles * w + 1, D]
+                #   --> [(num_height_tiles * h) * (num_width_tiles * w + 1), D]
+                local_features = local_features.view(-1, n_dim)
+
+                # merge global and local tiles
+                if self.global_view_pos == "head":
+                    global_local_features = torch.cat(
+                        [
+                            global_features,
+                            self.view_seperator[None, :],
+                            local_features,
+                        ]
+                    )
+                else:
+                    global_local_features = torch.cat(
+                        [
+                            local_features,
+                            self.view_seperator[None, :],
+                            global_features,
+                        ]
+                    )
+
+                images_in_this_batch.append(global_local_features)
 
         return torch.cat(images_in_this_batch, dim=0)
 

sglang/srt/models/gemma3_causal.py

@@ -47,6 +47,12 @@ from sglang.srt.model_loader.weight_utils import (
 from sglang.srt.utils import add_prefix, make_layers
 
 
+# Aligned with HF's implementation, using sliding window inclusive with the last token
+# SGLang assumes exclusive
+def get_attention_sliding_window_size(config):
+    return config.sliding_window - 1
+
+
 # Adapted from:
 # https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/gemma3.py
 def extract_layer_index(prefix: str) -> int:
@@ -170,7 +176,7 @@ class Gemma3Attention(nn.Module):
             self.rope_scaling = {"rope_type": "default"}
             # FIXME(mick): idk why vllm does this
             # self.sliding_window = config.interleaved_sliding_window
-            self.sliding_window = config.sliding_window
+            self.sliding_window = get_attention_sliding_window_size(config)
         else:
             # Global attention. Use the values in config.json.
             self.rope_theta = config.rope_theta
@@ -184,6 +190,8 @@ class Gemma3Attention(nn.Module):
             num_kv_heads=self.num_kv_heads,
             layer_id=layer_id,
             logit_cap=getattr(self.config, "attn_logit_softcapping", None),
+            # Module must also define `get_attention_sliding_window_size` to correctly initialize
+            # attention backend in `ForwardBatch`.
             sliding_window_size=self.sliding_window,
             prefix=add_prefix("attn", prefix),
         )
@@ -609,6 +617,9 @@ class Gemma3ForCausalLM(PreTrainedModel):
     def get_input_embeddings(self) -> nn.Embedding:
         return self.model.embed_tokens
 
+    def get_attention_sliding_window_size(self):
+        return get_attention_sliding_window_size(self.config)
+
     def dtype(self) -> torch.dtype:
         return next(self.parameters()).dtype
 
@@ -621,7 +632,6 @@ class Gemma3ForCausalLM(PreTrainedModel):
         input_embeds: torch.Tensor = None,
         **kwargs,
     ) -> LogitsProcessor:
-
         hidden_states = self.model(
             input_ids, positions, forward_batch, input_embeds, **kwargs
         )

sglang/srt/models/gemma3_mm.py

@@ -21,14 +21,7 @@ from typing import Dict, Iterable, List, Optional, Set, Tuple, TypedDict
 
 import torch
 from torch import nn
-from transformers import (
-    AutoModel,
-    BatchFeature,
-    Gemma3Config,
-    Gemma3Processor,
-    PreTrainedModel,
-)
-from transformers.models.gemma3.processing_gemma3 import Gemma3ProcessorKwargs
+from transformers import AutoModel, Gemma3Config, PreTrainedModel
 
 from sglang.srt.hf_transformers_utils import get_processor
 from sglang.srt.layers.layernorm import Gemma3RMSNorm
@@ -38,7 +31,11 @@ from sglang.srt.managers.mm_utils import (
     MultiModalityDataPaddingPatternTokenPairs,
     general_mm_embed_routine,
 )
-from sglang.srt.managers.schedule_batch import MultimodalInputs
+from sglang.srt.managers.schedule_batch import (
+    MultimodalDataItem,
+    MultimodalInputs,
+    flatten_nested_list,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
@@ -268,17 +265,22 @@ class Gemma3ForConditionalGeneration(PreTrainedModel):
     def get_input_embeddings(self) -> nn.Embedding:
         return self.language_model.get_input_embeddings()
 
-    def get_image_feature(self, image_input: MultimodalInputs):
+    def get_attention_sliding_window_size(self):
+        """
+        This value is used to initialize attention backends in `ForwardBatch`.
+        """
+        return self.language_model.get_attention_sliding_window_size()
+
+    def get_image_feature(self, items: List[MultimodalDataItem]):
         """
         Projects the last hidden state from the vision model into language model space.
 
-        Args:
-            pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`)
-               The tensors corresponding to the input images.
         Returns:
             image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
         """
-        pixel_values = image_input.pixel_values
+        pixel_values = torch.stack(
+            flatten_nested_list([item.pixel_values for item in items]), dim=0
+        )
         pixel_values = pixel_values.to("cuda")
         pixel_values = pixel_values.to(dtype=self.language_model.dtype())
 
@@ -286,61 +288,6 @@ class Gemma3ForConditionalGeneration(PreTrainedModel):
         image_features = self.multi_modal_projector(vision_outputs)
         return image_features
 
-    def embed_mm_inputs(
-        self,
-        input_ids: torch.Tensor,
-        forward_batch: ForwardBatch,
-        image_input: MultimodalInputs,
-    ) -> torch.Tensor:
-        if input_ids is None:
-            raise ValueError("Unimplemented")
-        # boolean-masking image tokens
-        special_image_mask = torch.isin(
-            input_ids,
-            torch.tensor(image_input.pad_values, device=input_ids.device),
-        ).unsqueeze(-1)
-        num_image_tokens_in_input_ids = special_image_mask.sum()
-
-        inputs_embeds = None
-        if num_image_tokens_in_input_ids == 0:
-            inputs_embeds = self.get_input_embeddings()(input_ids)
-            return inputs_embeds
-        else:
-            # print(f"image tokens from input_ids: {inputs_embeds[special_image_mask].numel()}")
-            image_features = self.get_image_feature(image_input.pixel_values)
-
-            # print(f"image tokens from image embeddings: {image_features.numel()}")
-            num_image_tokens_in_embedding = (
-                image_features.shape[0] * image_features.shape[1]
-            )
-
-            if num_image_tokens_in_input_ids != num_image_tokens_in_embedding:
-                num_image = num_image_tokens_in_input_ids // image_features.shape[1]
-                image_features = image_features[:num_image, :]
-                logger.warning(
-                    f"Number of images does not match number of special image tokens in the input text. "
-                    f"Got {num_image_tokens_in_input_ids} image tokens in the text but {num_image_tokens_in_embedding} "
-                    "tokens from image embeddings."
-                )
-
-            # Important: clamp after extracting original image boundaries
-            input_ids.clamp_(min=0, max=self.vocab_size - 1)
-
-            inputs_embeds = self.get_input_embeddings()(input_ids)
-
-            special_image_mask = special_image_mask.expand_as(inputs_embeds).to(
-                inputs_embeds.device
-            )
-
-            image_features = image_features.to(
-                inputs_embeds.device, inputs_embeds.dtype
-            )
-            inputs_embeds = inputs_embeds.masked_scatter(
-                special_image_mask, image_features
-            )
-
-        return inputs_embeds
-
     @torch.no_grad()
     def forward(
         self,
@@ -399,22 +346,15 @@ class Gemma3ForConditionalGeneration(PreTrainedModel):
         else:
             llm_input_ids = input_ids
 
-        inputs_embeds = general_mm_embed_routine(
+        hs = general_mm_embed_routine(
             input_ids=llm_input_ids,
             forward_batch=forward_batch,
-            embed_tokens=self.get_input_embeddings(),
-            mm_data_embedding_func=self.get_image_feature,
-        )
-
-        outputs = self.language_model(
-            input_ids=None,
+            language_model=self.language_model,
+            image_data_embedding_func=self.get_image_feature,
             positions=positions,
-            forward_batch=forward_batch,
-            input_embeds=inputs_embeds,
-            **kwargs,
         )
 
-        return outputs
+        return hs
 
     def tie_weights(self):
         return self.language_model.tie_weights()

sglang/srt/models/llama.py

@@ -17,7 +17,7 @@
 """Inference-only LLaMA model compatible with HuggingFace weights."""
 
 import logging
-from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import torch
 from torch import nn
@@ -428,6 +428,9 @@ class LlamaForCausalLM(nn.Module):
         else:
             return self.pooler(hidden_states, forward_batch)
 
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.embed_tokens
+
     def get_hidden_dim(self, module_name):
         # return input_dim, output_dim
         if module_name in ["q_proj", "o_proj", "qkv_proj"]:

sglang/srt/models/llava.py

@@ -31,7 +31,7 @@ from transformers import (
 from transformers.models.llava.modeling_llava import LlavaMultiModalProjector
 
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
-from sglang.srt.managers.schedule_batch import MultimodalInputs
+from sglang.srt.managers.schedule_batch import Modality, MultimodalInputs
 from sglang.srt.mm_utils import (
     get_anyres_image_grid_shape,
     unpad_image,
@@ -42,17 +42,21 @@ from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.llama import LlamaForCausalLM
 from sglang.srt.models.mistral import MistralForCausalLM
 from sglang.srt.models.qwen2 import Qwen2ForCausalLM
-from sglang.srt.utils import add_prefix
+from sglang.srt.utils import add_prefix, flatten_nested_list
 
 
 class LlavaBaseForCausalLM(nn.Module):
     def pad_input_ids(self, input_ids: List[int], image_inputs: MultimodalInputs):
-        image_sizes, pad_values = image_inputs.image_sizes, image_inputs.pad_values
+        image_sizes = flatten_nested_list(
+            [item.image_sizes for item in image_inputs.mm_items]
+        )
+
+        pad_values = [item.pad_value for item in image_inputs.mm_items]
 
         # hardcode for spatial_unpad + anyres
-        if image_inputs.modalities is not None and (
-            "multi-images" in image_inputs.modalities
-            or "video" in image_inputs.modalities
+        if any(
+            item.modality == Modality.MULTI_IMAGES or item.modality == Modality.VIDEO
+            for item in image_inputs.mm_items
         ):
             image_aspect_ratio = "pad"
         else:
@@ -66,7 +70,7 @@ class LlavaBaseForCausalLM(nn.Module):
                     math.ceil(self.image_size / self.patch_size / 2) ** 2
                 )
             else:
-                new_image_feature_len = self.image_feature_len  # multiimage
+                new_image_feature_len = self.image_feature_len  # multi-image
 
             height = width = self.num_patches_per_side
             if "anyres" in image_aspect_ratio:
@@ -101,7 +105,7 @@ class LlavaBaseForCausalLM(nn.Module):
             # old_len + pad_len - 1, because we need to remove image_token_id
             input_ids = (
                 input_ids[:offset]
-                + [pad_values[image_idx]] * new_image_feature_len
+                + [pad_values[image_idx % len(pad_values)]] * new_image_feature_len
                 + input_ids[offset + 1 :]
             )
             offset_list.append(offset)
@@ -150,8 +154,8 @@ class LlavaBaseForCausalLM(nn.Module):
             modalities_list = []
             max_image_offset = []
             for im in image_inputs:
-                if im and im.modalities is not None:
-                    modalities_list.extend(im.modalities)
+                if im:
+                    modalities_list.extend([item.modality for item in im.mm_items])
                 if im and im.image_offsets:
                     max_image_offset.append(
                         np.max(np.array(im.image_offsets) + np.array(im.image_pad_len))
@@ -164,11 +168,19 @@ class LlavaBaseForCausalLM(nn.Module):
 
             if need_vision.any():
                 bs = forward_batch.batch_size
-                pixel_values = [
-                    image_inputs[i].pixel_values for i in range(bs) if need_vision[i]
-                ]
+                pixel_values = flatten_nested_list(
+                    [
+                        [item.pixel_values for item in image_inputs[i].mm_items]
+                        for i in range(bs)
+                        if need_vision[i]
+                    ]
+                )
                 image_sizes = [
-                    image_inputs[i].image_sizes for i in range(bs) if need_vision[i]
+                    flatten_nested_list(
+                        [item.image_sizes for item in image_inputs[i].mm_items]
+                    )
+                    for i in range(bs)
+                    if need_vision[i]
                 ]
 
                 ########## Encode Image ########
@@ -197,13 +209,13 @@ class LlavaBaseForCausalLM(nn.Module):
                     new_image_features = []
                     height = width = self.num_patches_per_side
                     for image_idx, image_feature in enumerate(image_features):
-                        if modalities_list[image_idx] == "image":
+                        if modalities_list[image_idx] == Modality.IMAGE:
                             image_aspect_ratio = (
                                 self.config.image_aspect_ratio
                             )  # single image
                         elif (
-                            modalities_list[image_idx] == "multi-images"
-                            or modalities_list[image_idx] == "video"
+                            modalities_list[image_idx] == Modality.MULTI_IMAGES
+                            or modalities_list[image_idx] == Modality.VIDEO
                         ):
                             image_aspect_ratio = "pad"  # multi image
                         # image_aspect_ratio = (
@@ -212,7 +224,7 @@ class LlavaBaseForCausalLM(nn.Module):
                         if (
                             image_feature.shape[0] > 1
                             and "anyres" in image_aspect_ratio
-                            and modalities_list[image_idx] == "image"
+                            and modalities_list[image_idx] == Modality.IMAGE
                         ):
                             base_image_feature = image_feature[0]
                             image_feature = image_feature[1:]
@@ -312,7 +324,7 @@ class LlavaBaseForCausalLM(nn.Module):
                             )
                             image_feature = image_feature.unsqueeze(0)
                         else:
-                            if modalities_list[image_idx] == "video":  # video
+                            if modalities_list[image_idx] == Modality.VIDEO:  # video
                                 # 2x2 pooling
                                 num_of_frames = image_feature.shape[0]
                                 image_feature = image_feature.view(

sglang/srt/models/llavavid.py

@@ -22,7 +22,7 @@ from transformers import CLIPVisionModel, LlavaConfig
 from transformers.models.llava.modeling_llava import LlavaMultiModalProjector
 
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
-from sglang.srt.managers.schedule_batch import MultimodalInputs
+from sglang.srt.managers.schedule_batch import MultimodalInputs, flatten_nested_list
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.llama import LlamaForCausalLM
@@ -58,7 +58,7 @@ class LlavaVidForCausalLM(nn.Module):
             )
 
     def pad_input_ids(self, input_ids: List[int], image_inputs: MultimodalInputs):
-        pad_values = image_inputs.pad_values
+        pad_values = [item.pad_value for item in image_inputs.mm_items]
         new_image_feature_len = self.image_feature_len
 
         pad_ids = pad_values * (
@@ -133,11 +133,19 @@ class LlavaVidForCausalLM(nn.Module):
             need_vision = start_positions <= np.array(max_image_offset)
 
             if need_vision.any():
-                pixel_values = [
-                    image_inputs[i].pixel_values for i in range(bs) if need_vision[i]
-                ]
+                pixel_values = flatten_nested_list(
+                    [
+                        [item.pixel_values for item in image_inputs[i].mm_items]
+                        for i in range(bs)
+                        if need_vision[i]
+                    ]
+                )
                 image_offsets = [
-                    image_inputs[i].image_offsets for i in range(bs) if need_vision[i]
+                    flatten_nested_list(
+                        [item.image_offsets for item in image_inputs[i].mm_items]
+                    )
+                    for i in range(bs)
+                    if need_vision[i]
                 ]
 
                 ########## Encode Image ########
@@ -246,7 +254,8 @@ class LlavaVidForCausalLM(nn.Module):
             "model.mm_projector.2": "multi_modal_projector.linear_2",
             "model.vision_resampler.mm_projector.0": "multi_modal_projector.linear_1",
             "model.vision_resampler.mm_projector.2": "multi_modal_projector.linear_2",
-            "model.vision_tower.vision_tower": "vision_tower",  # Update the vision tower weights if we find them in the checkpoint (it may be finetuned).
+            "model.vision_tower.vision_tower": "vision_tower",
+            # Update the vision tower weights if we find them in the checkpoint (it may be finetuned).
             "model.image_newline": "language_model.model.image_newline",
         }
         params_dict = dict(self.named_parameters())