Fix HunYuan-VL image assignment during batched generation

src/transformers/models/hunyuan_vl/modeling_hunyuan_vl.py

@@ -979,7 +979,12 @@ def get_image_features(
     ) -> tuple | BaseModelOutputWithPooling:
         vision_dtype = next(self.vision_tower.parameters()).dtype
         pixel_values = pixel_values.to(vision_dtype)
-        return self.vision_tower(pixel_values, grid_thw=image_grid_thw, **kwargs)
+        image_outputs = self.vision_tower(pixel_values, grid_thw=image_grid_thw, **kwargs)
+        # Include row separators and image boundary tokens in each split.
+        merged_h, merged_w = (image_grid_thw[:, 1:] // self.config.vision_config.spatial_merge_size).unbind(-1)
+        split_sizes = (merged_h * (merged_w + 1) + 2).tolist()
+        image_outputs.pooler_output = torch.split(image_outputs.pooler_output[0], split_sizes)
+        return image_outputs
 
     def get_placeholder_mask(
         self,
@@ -1064,7 +1069,9 @@ def forward(
             mm_encoder_outputs["image"] = self.get_image_features(pixel_values, image_grid_thw, return_dict=True)
 
         if mm_encoder_outputs.get("image") is not None:
-            image_embeds = mm_encoder_outputs["image"].pooler_output.to(inputs_embeds.device, inputs_embeds.dtype)
+            image_embeds = torch.cat(mm_encoder_outputs["image"].pooler_output, dim=0).to(
+                inputs_embeds.device, inputs_embeds.dtype
+            )
             image_mask = self.get_placeholder_mask(
                 input_ids,
                 inputs_embeds=inputs_embeds,

src/transformers/models/hunyuan_vl/modular_hunyuan_vl.py

@@ -1224,7 +1224,12 @@ def get_image_features(
     ) -> tuple | BaseModelOutputWithPooling:
         vision_dtype = next(self.vision_tower.parameters()).dtype
         pixel_values = pixel_values.to(vision_dtype)
-        return self.vision_tower(pixel_values, grid_thw=image_grid_thw, **kwargs)
+        image_outputs = self.vision_tower(pixel_values, grid_thw=image_grid_thw, **kwargs)
+        # Include row separators and image boundary tokens in each split.
+        merged_h, merged_w = (image_grid_thw[:, 1:] // self.config.vision_config.spatial_merge_size).unbind(-1)
+        split_sizes = (merged_h * (merged_w + 1) + 2).tolist()
+        image_outputs.pooler_output = torch.split(image_outputs.pooler_output[0], split_sizes)
+        return image_outputs
 
     def get_placeholder_mask(
         self,
@@ -1281,7 +1286,9 @@ def forward(
             mm_encoder_outputs["image"] = self.get_image_features(pixel_values, image_grid_thw, return_dict=True)
 
         if mm_encoder_outputs.get("image") is not None:
-            image_embeds = mm_encoder_outputs["image"].pooler_output.to(inputs_embeds.device, inputs_embeds.dtype)
+            image_embeds = torch.cat(mm_encoder_outputs["image"].pooler_output, dim=0).to(
+                inputs_embeds.device, inputs_embeds.dtype
+            )
             image_mask = self.get_placeholder_mask(
                 input_ids,
                 inputs_embeds=inputs_embeds,

src/transformers/models/idefics2/modeling_idefics2.py

@@ -836,6 +836,7 @@ def get_image_features(
         # Remove padding images - padding images are full 0.
         nb_values_per_image = pixel_values.shape[1:].numel()
         real_images_inds = (pixel_values == 0.0).sum(dim=(-1, -2, -3)) != nb_values_per_image
+        num_real_images_per_batch = real_images_inds.view(batch_size, num_images).sum(dim=1)
         pixel_values = pixel_values[real_images_inds].contiguous()
 
         # Handle the vision attention mask
@@ -861,9 +862,12 @@ def get_image_features(
         image_hidden_states = image_outputs.last_hidden_state
 
         # Modality projection & resampling
-        image_outputs.pooler_output = self.connector(
+        image_features = self.connector(
             image_hidden_states, attention_mask=patch_attention_mask.view(pixel_values.size(0), -1)
         )
+        image_outputs.pooler_output = [
+            f.reshape(-1, f.shape[-1]) for f in image_features.split(num_real_images_per_batch.tolist())
+        ]
 
         return image_outputs
 
@@ -938,6 +942,8 @@ def forward(
             image_hidden_states = None
 
         if image_hidden_states is not None:
+            if isinstance(image_hidden_states, (list, tuple)):
+                image_hidden_states = torch.cat(image_hidden_states, dim=0)
             image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=inputs_embeds.device)
             # When we generate, we don't want to replace the potential image_token_id that we generated by images
             # that simply don't exist

src/transformers/models/idefics3/modeling_idefics3.py

@@ -580,6 +580,7 @@ def get_image_features(
         # Remove padding images - padding images are full 0.
         nb_values_per_image = pixel_values.shape[1:].numel()
         real_images_inds = (pixel_values == 0.0).sum(dim=(-1, -2, -3)) != nb_values_per_image
+        num_real_images_per_batch = real_images_inds.view(batch_size, num_images).sum(dim=1)
         pixel_values = pixel_values[real_images_inds].contiguous()
 
         # Handle the vision attention mask
@@ -607,7 +608,9 @@ def get_image_features(
 
         # Modality projection & resampling
         image_features = self.connector(image_hidden_states)
-        image_outputs.pooler_output = image_features
+        image_outputs.pooler_output = [
+            f.reshape(-1, f.shape[-1]) for f in image_features.split(num_real_images_per_batch.tolist())
+        ]
 
         return image_outputs
 
@@ -679,6 +682,8 @@ def forward(
             image_hidden_states = None
 
         if image_hidden_states is not None:
+            if isinstance(image_hidden_states, (list, tuple)):
+                image_hidden_states = torch.cat(image_hidden_states, dim=0)
             image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=inputs_embeds.device)
             # When we generate, we don't want to replace the potential image_token_id that we generated by images
             # that simply don't exist

src/transformers/models/llava_next_video/modeling_llava_next_video.py

@@ -640,7 +640,7 @@ def get_video_features(
         # Same as image features except that video has pooling layer
         video_features = self.vision_resampler(selected_video_features)
         video_features = self.multi_modal_projector(video_features)
-        video_outputs.pooler_output = torch.split(video_features, frames, dim=0)
+        video_outputs.pooler_output = video_features.unflatten(0, (batch_size, frames)).flatten(1, 2)
         return video_outputs
 
 

src/transformers/models/llava_next_video/modular_llava_next_video.py

@@ -339,7 +339,7 @@ def get_video_features(
         # Same as image features except that video has pooling layer
         video_features = self.vision_resampler(selected_video_features)
         video_features = self.multi_modal_projector(video_features)
-        video_outputs.pooler_output = torch.split(video_features, frames, dim=0)
+        video_outputs.pooler_output = video_features.unflatten(0, (batch_size, frames)).flatten(1, 2)
         return video_outputs
 
     def get_placeholder_mask(

src/transformers/models/modernvbert/modeling_modernvbert.py

@@ -241,7 +241,8 @@ def inputs_merger(
         - The merging happens so that we obtain the following sequence: `vector_tok_1 vector_tok_2 vector_tok_3 vector_fake_tok_around_image {sequence of image_seq_len image hidden states} vector_fake_toke_around_image vector_tok_4`. That sequence is fed to the LM.
         - To fit the format of that sequence, `input_ids`, `inputs_embeds`, `attention_mask` are all 3 adapted to insert the image hidden states.
         """
-        _, patch_size, _ = image_hidden_states.shape
+        patch_size = self.config.vision_config.patch_size
+        image_hidden_states = image_hidden_states.view(-1, patch_size, image_hidden_states.shape[-1])
 
         if input_ids is None:
             image_mask = inputs_embeds == self.get_input_embeddings()(
@@ -297,6 +298,7 @@ def get_image_features(
 
         # If no images, leave one empty image.
         real_images_inds[0] |= ~torch.any(real_images_inds)
+        num_real_images_per_batch = real_images_inds.view(batch_size, num_images).sum(dim=1)
 
         pixel_values = pixel_values[real_images_inds].contiguous()
         # Handle the vision attention mask
@@ -323,7 +325,9 @@ def get_image_features(
 
         # Modality projection & resampling
         image_features = self.connector(image_hidden_states)
-        image_outputs.pooler_output = image_features
+        image_outputs.pooler_output = [
+            f.reshape(-1, f.shape[-1]) for f in image_features.split(num_real_images_per_batch.tolist())
+        ]
 
         return image_outputs
 
@@ -369,6 +373,7 @@ def forward(
 
         # Merge image and text embeddings
         if image_hidden_states is not None:
+            image_hidden_states = torch.cat(image_hidden_states, dim=0)
             image_hidden_states = image_hidden_states.to(dtype=inputs_embeds.dtype, device=inputs_embeds.device)
             inputs_embeds = self.inputs_merger(
                 input_ids=input_ids, inputs_embeds=inputs_embeds, image_hidden_states=image_hidden_states

src/transformers/models/modernvbert/modular_modernvbert.py

@@ -296,6 +296,7 @@ def forward(
 
         # Merge image and text embeddings
         if image_hidden_states is not None:
+            image_hidden_states = torch.cat(image_hidden_states, dim=0)
             image_hidden_states = image_hidden_states.to(dtype=inputs_embeds.dtype, device=inputs_embeds.device)
             inputs_embeds = self.inputs_merger(
                 input_ids=input_ids, inputs_embeds=inputs_embeds, image_hidden_states=image_hidden_states

src/transformers/models/smolvlm/modeling_smolvlm.py

@@ -478,7 +478,8 @@ def inputs_merger(
         - The merging happens so that we obtain the following sequence: `vector_tok_1 vector_tok_2 vector_tok_3 vector_fake_tok_around_image {sequence of image_seq_len image hidden states} vector_fake_toke_around_image vector_tok_4`. That sequence is fed to the LM.
         - To fit the format of that sequence, `input_ids`, `inputs_embeds`, `attention_mask` are all 3 adapted to insert the image hidden states.
         """
-        _, patch_size, _ = image_hidden_states.shape
+        patch_size = self.config.vision_config.patch_size
+        image_hidden_states = image_hidden_states.view(-1, patch_size, image_hidden_states.shape[-1])
 
         if input_ids is None:
             image_mask = inputs_embeds == self.get_input_embeddings()(
@@ -534,6 +535,7 @@ def get_image_features(
 
         # If no images, leave one empty image.
         real_images_inds[0] |= ~torch.any(real_images_inds)
+        num_real_images_per_batch = real_images_inds.view(batch_size, num_images).sum(dim=1)
 
         pixel_values = pixel_values[real_images_inds].contiguous()
         # Handle the vision attention mask
@@ -560,7 +562,9 @@ def get_image_features(
 
         # Modality projection & resampling
         image_features = self.connector(image_hidden_states)
-        image_outputs.pooler_output = image_features
+        image_outputs.pooler_output = [
+            f.reshape(-1, f.shape[-1]) for f in image_features.split(num_real_images_per_batch.tolist())
+        ]
 
         return image_outputs
 
@@ -624,6 +628,8 @@ def forward(
             image_hidden_states = None
 
         if image_hidden_states is not None:
+            if isinstance(image_hidden_states, (list, tuple)):
+                image_hidden_states = torch.cat(image_hidden_states, dim=0)
             image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=inputs_embeds.device)
             # When we generate, we don't want to replace the potential image_token_id that we generated by images
             # that simply don't exist

src/transformers/models/smolvlm/modular_smolvlm.py

@@ -110,7 +110,8 @@ class SmolVLMModel(Idefics3Model):
     def inputs_merger(
         self, input_ids: torch.LongTensor, inputs_embeds: torch.Tensor, image_hidden_states: torch.Tensor
     ):
-        _, patch_size, _ = image_hidden_states.shape
+        patch_size = self.config.vision_config.patch_size
+        image_hidden_states = image_hidden_states.view(-1, patch_size, image_hidden_states.shape[-1])
 
         if input_ids is None:
             image_mask = inputs_embeds == self.get_input_embeddings()(
@@ -166,6 +167,7 @@ def get_image_features(
 
         # If no images, leave one empty image.
         real_images_inds[0] |= ~torch.any(real_images_inds)
+        num_real_images_per_batch = real_images_inds.view(batch_size, num_images).sum(dim=1)
 
         pixel_values = pixel_values[real_images_inds].contiguous()
         # Handle the vision attention mask
@@ -192,7 +194,9 @@ def get_image_features(
 
         # Modality projection & resampling
         image_features = self.connector(image_hidden_states)
-        image_outputs.pooler_output = image_features
+        image_outputs.pooler_output = [
+            f.reshape(-1, f.shape[-1]) for f in image_features.split(num_real_images_per_batch.tolist())
+        ]
 
         return image_outputs
 
@@ -252,6 +256,8 @@ def forward(
             image_hidden_states = None
 
         if image_hidden_states is not None:
+            if isinstance(image_hidden_states, (list, tuple)):
+                image_hidden_states = torch.cat(image_hidden_states, dim=0)
             image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=inputs_embeds.device)
             # When we generate, we don't want to replace the potential image_token_id that we generated by images
             # that simply don't exist

tests/models/chameleon/test_modeling_chameleon.py

@@ -245,6 +245,10 @@ def test_batching_equivalence(self):
     def test_get_image_features_output(self, return_dict: bool | None):
         pass
 
+    @unittest.skip("Skip get_image_features tests as those are tested via ChameleonVision2SeqModelTest instead")
+    def test_get_image_features_pooler_output(self):
+        pass
+
     @unittest.skip("Skip get_image_features tests as those are tested via ChameleonVision2SeqModelTest instead")
     def test_get_image_features_hidden_states(self):
         pass

tests/models/gemma4/test_modeling_gemma4.py

@@ -288,6 +288,10 @@ def test_get_image_features_attentions(self):
     def test_get_image_features_output(self, return_dict: bool | None):
         pass
 
+    @unittest.skip("The tester has no image in input dict")
+    def test_get_image_features_pooler_output(self):
+        pass
+
     @unittest.skip("The tester has no videos in input dict")
     def test_get_video_features_hidden_states(self):
         pass

tests/models/gemma4_unified/test_modeling_gemma4_unified.py

@@ -244,6 +244,7 @@ def skip_mm_output_format(self):
             "test_get_image_features_hidden_states",
             "test_get_image_features_attentions",
             "test_get_image_features_output",
+            "test_get_image_features_pooler_output",
             "test_get_video_features_hidden_states",
             "test_get_video_features_attentions",
             "test_get_video_features_output",

tests/models/inkling/test_modeling_inkling.py

@@ -176,6 +176,10 @@ def test_get_image_features_attentions(self):
     def test_get_image_features_output(self, return_dict: bool | None):
         pass
 
+    @unittest.skip("The tester has no image in input dict")
+    def test_get_image_features_pooler_output(self):
+        pass
+
     @unittest.skip("The tester has no videos in input dict")
     def test_get_video_features_hidden_states(self):
         pass

tests/models/nemotron_h_omni/test_modeling_nemotron_h_omni.py

@@ -347,6 +347,11 @@ def test_get_image_features_output(self, return_dict: bool | None):
     def test_get_video_features_output(self, return_dict: bool | None):
         pass
 
+    @unittest.skip("FIXME raushan - should apply temporal factor while processing!")
+    # prob this model doesn't even work with beam/sampling atm, low prio
+    def test_get_video_features_pooler_output(self):
+        pass
+
     def test_mismatching_num_image_tokens(self):
         # packed `pixel_values` have no batch dimension, so images are dropped or added by their patches
         config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
@@ -435,6 +440,11 @@ def test_get_image_features_output(self, return_dict: bool | None):
     def test_get_video_features_output(self, return_dict: bool | None):
         pass
 
+    @unittest.skip("FIXME raushan - should apply temporal factor while processing!")
+    # prob this model doesn't even work with beam/sampling atm, low prio
+    def test_get_video_features_pooler_output(self):
+        pass
+
     @unittest.skip(reason="Audio tester isnt supported yet")
     def test_generate_from_multimodal_encoder_outputs_and_raw_data(self):
         pass