Repository: huggingface/transformers
Pull request: Llava onevision: output align for tests and add `image_sizes` input param

In this PR, we do several things for llava_onevision model:
1. skip torch_exportable tests as it does not support it
2. unify expected output for cuda and xpu
3. add `image_sizes` param in `flash_attn_inference_equivalence` func to support VLM models like lighton_ocr and llava_onevision

Changed files:
filename: tests/models/lighton_ocr/test_modeling_lighton_ocr.py
patch: @@ -244,15 +244,11 @@ def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
         """
         inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels)
 
-        # Ensure image_sizes is properly set when pixel_values is present
-        if "pixel_values" in inputs_dict:
+        # Ensure image_sizes matches the batch size of pixel_values or input_ids
+        if "pixel_values" in inputs_dict and "image_sizes" in inputs_dict:
             batch_size = inputs_dict["pixel_values"].shape[0]
-            # If image_sizes is missing or None, infer from pixel_values
-            if "image_sizes" not in inputs_dict or inputs_dict["image_sizes"] is None:
-                _, _, height, width = inputs_dict["pixel_values"].shape
-                inputs_dict["image_sizes"] = [(height, width)] * batch_size
-            elif len(inputs_dict["image_sizes"]) != batch_size:
-                # Take only the first batch_size entries
+            # If image_sizes doesn't match batch size, adjust it
+            if len(inputs_dict["image_sizes"]) != batch_size:
                 inputs_dict["image_sizes"] = inputs_dict["image_sizes"][:batch_size]
 
         return inputs_dict

filename: tests/models/llava_onevision/test_modeling_llava_onevision.py
patch: @@ -355,7 +355,7 @@ def test_small_model_integration_test(self):
 
         EXPECTED_DECODED_TEXTS = Expectations(
             {
-                ("xpu", 3): 'user\n\nWhat do you see in this image?\nassistant\nThe image is a radar chart that compares the performance of different models in a specific task, likely related to natural language processing or machine learning. The chart is divided into several axes, each representing a different model or method. The models are color-coded and labeled with their respective names. The axes are labeled with terms such as "VQA," "GQA," "MQA," "VQAv2," "MM-Vet," "LLaVA-Bench," "LLaVA-1',
+                ("xpu", 3): 'user\n\nWhat do you see in this image?\nassistant\nThe image is a radar chart that compares the performance of different models in a specific task, likely related to natural language processing or machine learning. The chart is divided into several axes, each representing a different model or method. The models are color-coded and labeled with their respective names. The axes are labeled with terms such as "VQA," "GQA," "MQA," "VIZ," "TextVQA," "SQA-IMG," and "MQE." The radar chart shows',
                 ("cuda", 7): 'user\n\nWhat do you see in this image?\nassistant\nThe image is a radar chart that compares the performance of different models in a specific task, likely related to natural language processing or machine learning. The chart is divided into several axes, each representing a different model or method. The models are color-coded and labeled with their respective names. The axes are labeled with terms such as "VQA," "GQA," "MQA," "VQAv2," "MM-Vet," "LLaVA-Bench," "LLaVA-1',
                 ("cuda", 8): 'user\n\nWhat do you see in this image?\nassistant\nThe image is a radar chart that compares the performance of different models in a specific task, likely related to natural language processing or machine learning. The chart is divided into several axes, each representing a different model or method. The models are color-coded and labeled with their respective names. The axes are labeled with terms such as "VQA," "GQA," "MQA," "VIZ," "TextVQA," "SQA-IMG," and "MQE." The radar chart shows',
             }
@@ -483,8 +483,8 @@ def test_small_model_integration_test_multi_image_nested(self):
                     ],
                 ("xpu", 3): [
                     "user\nTell me about the french revolution.\nassistant\nThe French Revolution! A pivotal event in modern history that had a profound impact on the course of Western civilization. Here's a brief overview:\n\n**Background**\n\nIn the late 18th century,",
-                    'user\n\nWhat is the difference between these images?\nassistant\nThe image shows a traffic light with a stop sign in the foreground, while the other images show a car driving through a street intersection.',
-                    'user\n\nWhat do you see in this image?\nassistant\nThe image is a radar chart that represents the performance of different machine learning models in terms of their ability to predict the number of users who have been infected with COVID-19. The radar chart is'
+                    "user\n\nWhat is the difference between these images?\nassistant\nThe first image shows a stop sign with a traditional Chinese architectural background, while the second image displays a radar chart with various algorithms and models, including BLIP-2, InstructBLIP, Q",
+                    "user\n\nWhat do you see in this image?\nassistant\nThe image is a radar chart that compares the performance of different models in a specific task, likely related to natural language processing or machine learning. The chart is divided into several axes, each representing a different"
                     ],
             }
         )

filename: tests/test_modeling_common.py
patch: @@ -3263,6 +3263,9 @@ def flash_attn_inference_equivalence(
                     if "image_grid_thw" in inputs_dict:
                         continue
                     first_inputs["pixel_values"] = inputs_dict["pixel_values"][:1].to(torch.bfloat16)
+                # Some VLMs require image_sizes alongside pixel_values, e.g. lighton_ocr, llava_onevision
+                if "image_sizes" in inputs_dict:
+                    first_inputs["image_sizes"] = inputs_dict["image_sizes"][:1]
                 if model.config.is_encoder_decoder:
                     decoder_input_ids = inputs_dict.get("decoder_input_ids", first_inputs.get("input_ids"))
                     if decoder_input_ids is not None: