feat: Implement get_image_features method in Aria, Mistral3, and VipLlava models with updated parameters

Author: YushunXiang

src/transformers/models/aria/modeling_aria.py

@@ -1228,15 +1228,13 @@ def get_decoder(self):
     def get_image_features(
         self,
         pixel_values: torch.FloatTensor,
-        vision_feature_layer: Optional[Union[int, List[int]]] = None,
-        vision_feature_select_strategy: Optional[str] = None,
-        **kwargs,
+        pixel_mask: Optional[torch.FloatTensor] = None,
+        vision_feature_layer: int = -1,
     ):
         return self.model.get_image_features(
             pixel_values=pixel_values,
+            pixel_mask=pixel_mask,
             vision_feature_layer=vision_feature_layer,
-            vision_feature_select_strategy=vision_feature_select_strategy,
-            **kwargs,
         )
 
     # Make modules available throught conditional class for BC

src/transformers/models/aria/modular_aria.py

@@ -1497,6 +1497,18 @@ def forward(
     """
 )
 class AriaForConditionalGeneration(LlavaForConditionalGeneration):
+    def get_image_features(
+        self,
+        pixel_values: torch.FloatTensor,
+        pixel_mask: Optional[torch.FloatTensor] = None,
+        vision_feature_layer: int = -1,
+    ):
+        return self.model.get_image_features(
+            pixel_values=pixel_values,
+            pixel_mask=pixel_mask,
+            vision_feature_layer=vision_feature_layer,
+        )
+
     @can_return_tuple
     @auto_docstring
     def forward(

src/transformers/models/mistral3/modeling_mistral3.py

@@ -415,14 +415,14 @@ def get_decoder(self):
     def get_image_features(
         self,
         pixel_values: torch.FloatTensor,
+        image_sizes: torch.Tensor,
         vision_feature_layer: Optional[Union[int, List[int]]] = None,
-        vision_feature_select_strategy: Optional[str] = None,
         **kwargs,
     ):
         return self.model.get_image_features(
             pixel_values=pixel_values,
+            image_sizes=image_sizes,
             vision_feature_layer=vision_feature_layer,
-            vision_feature_select_strategy=vision_feature_select_strategy,
             **kwargs,
         )
 

src/transformers/models/mistral3/modular_mistral3.py

@@ -254,6 +254,20 @@ def forward(
 
 
 class Mistral3ForConditionalGeneration(LlavaForConditionalGeneration):
+    def get_image_features(
+        self,
+        pixel_values: torch.FloatTensor,
+        image_sizes: torch.Tensor,
+        vision_feature_layer: Optional[Union[int, List[int]]] = None,
+        **kwargs,
+    ):
+        return self.model.get_image_features(
+            pixel_values=pixel_values,
+            image_sizes=image_sizes,
+            vision_feature_layer=vision_feature_layer,
+            **kwargs,
+        )
+
     def forward(
         self,
         input_ids: torch.LongTensor = None,

src/transformers/models/vipllava/modeling_vipllava.py

@@ -333,18 +333,9 @@ def get_decoder(self):
         return self.model
 
     def get_image_features(
-        self,
-        pixel_values: torch.FloatTensor,
-        vision_feature_layer: Optional[Union[int, List[int]]] = None,
-        vision_feature_select_strategy: Optional[str] = None,
-        **kwargs,
+        self, pixel_values: torch.FloatTensor, vision_feature_layers: Optional[Union[int, List[int]]] = None
     ):
-        return self.model.get_image_features(
-            pixel_values=pixel_values,
-            vision_feature_layer=vision_feature_layer,
-            vision_feature_select_strategy=vision_feature_select_strategy,
-            **kwargs,
-        )
+        return self.model.get_image_features(pixel_values=pixel_values, vision_feature_layers=vision_feature_layers)
 
     # Make modules available throught conditional class for BC
     @property

src/transformers/models/vipllava/modular_vipllava.py

@@ -184,6 +184,13 @@ def forward(
 
 
 class VipLlavaForConditionalGeneration(LlavaForConditionalGeneration):
+    def get_image_features(
+        self, pixel_values: torch.FloatTensor, vision_feature_layers: Optional[Union[int, List[int]]] = None
+    ):
+        return self.model.get_image_features(
+            pixel_values=pixel_values, vision_feature_layers=vision_feature_layers
+        )
+
     def forward(
         self,
         input_ids: torch.LongTensor = None,