Repository: huggingface/transformers
Pull request: Fix broken image inference for Fuyu model

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

When updating Transformers version in vLLM CI, we found fuyu's image inference is broken with gibberish:
https://buildkite.com/vllm/ci/builds/25473#01985c4d-d267-408b-87f5-5d77ae09d90c

- This PR fixed broken fuyu models to make sure image features are image_patches handled correctly.


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @eustlb
- graph models: @clefourrier

Library:

- flax: @gante and @Rocketknight1
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @zach-huggingface, @SunMarc and @qgallouedec
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @SunMarc @zach-huggingface
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @Rocketknight1
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Changed files:
filename: src/transformers/models/fuyu/modeling_fuyu.py
patch: @@ -225,7 +225,7 @@ def forward(
         if image_patches is not None:
             patch_embeddings = self.get_image_features(image_patches)
             patch_embeddings = torch.cat(patch_embeddings, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
-            special_image_mask = self.get_placeholder_tokens(
+            special_image_mask = self.get_placeholder_mask(
                 input_ids, inputs_embeds=inputs_embeds, image_features=patch_embeddings
             )
             inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, patch_embeddings)
@@ -379,6 +379,7 @@ def prepare_inputs_for_generation(
         inputs_embeds=None,
         image_patches=None,
         image_patches_indices=None,
+        cache_position=None,
         **kwargs,
     ):
         # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
@@ -390,10 +391,12 @@ def prepare_inputs_for_generation(
             inputs_embeds=inputs_embeds,
             image_patches=image_patches,
             image_patches_indices=image_patches_indices,
+            cache_position=cache_position,
             **kwargs,
         )
 
-        if past_key_values is not None:
+        if cache_position[0] != 0:
+            # set image_patches and image_patches_indices to `None` for decoding stage
             model_inputs["image_patches_indices"] = None
             model_inputs["image_patches"] = None
 

filename: tests/models/fuyu/test_modeling_fuyu.py
patch: @@ -13,19 +13,21 @@
 # limitations under the License.
 """Testing suite for the PyTorch Fuyu model."""
 
+import copy
 import io
 import unittest
 
 import pytest
 import requests
+import torch
 from parameterized import parameterized
 
 from transformers import FuyuConfig, is_torch_available, is_vision_available
-from transformers.testing_utils import require_torch, require_torch_accelerator, slow
+from transformers.testing_utils import require_torch, require_torch_accelerator, slow, torch_device
 from transformers.utils import cached_property
 
 from ...generation.test_utils import GenerationTesterMixin
-from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
+from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
 from ...test_pipeline_mixin import PipelineTesterMixin
 
 
@@ -47,6 +49,7 @@ def __init__(
         parent,
         batch_size=13,
         seq_length=7,
+        num_image_tokens=2,
         image_size=30,
         patch_size=15,
         num_channels=3,
@@ -67,12 +70,14 @@ def __init__(
         initializer_range=0.02,
         num_labels=3,
         num_choices=4,
-        pad_token_id=0,
+        pad_token_id=10,
+        image_token_id=1,
         scope=None,
     ):
         self.parent = parent
         self.batch_size = batch_size
-        self.seq_length = seq_length
+        self.num_image_tokens = num_image_tokens
+        self.seq_length = seq_length + num_image_tokens
         self.image_size = image_size
         self.patch_size = patch_size
         self.num_channels = num_channels
@@ -94,10 +99,15 @@ def __init__(
         self.num_labels = num_labels
         self.num_choices = num_choices
         self.pad_token_id = pad_token_id
+        self.image_token_id = image_token_id
         self.scope = scope
 
     def prepare_config_and_inputs(self):
+        config = self.get_config()
+
         input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
+        input_ids[input_ids == config.image_token_id] = self.pad_token_id
+        input_ids[:, : self.num_image_tokens] = config.image_token_id
 
         input_mask = None
         if self.use_input_mask:
@@ -109,8 +119,6 @@ def prepare_config_and_inputs(self):
             sequence_labels = ids_tensor([self.batch_size], self.type_sequence_label_size)
             token_labels = ids_tensor([self.batch_size, self.seq_length], self.num_labels)
 
-        config = self.get_config()
-
         return config, input_ids, input_mask, sequence_labels, token_labels
 
     def get_config(self):
@@ -128,6 +136,7 @@ def get_config(self):
             is_decoder=False,
             initializer_range=self.initializer_range,
             pad_token_id=self.pad_token_id,
+            image_token_id=self.image_token_id,
         )
 
     def prepare_config_and_inputs_for_common(self):
@@ -139,7 +148,10 @@ def prepare_config_and_inputs_for_common(self):
             sequence_labels,
             token_labels,
         ) = config_and_inputs
-        inputs_dict = {"input_ids": input_ids, "attention_mask": input_mask}
+        image_patches = floats_tensor(
+            [self.batch_size, self.num_image_tokens, config.num_channels * config.patch_size**2]
+        )
+        inputs_dict = {"input_ids": input_ids, "attention_mask": input_mask, "image_patches": image_patches}
         return config, inputs_dict
 
 
@@ -166,6 +178,27 @@ class FuyuModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin
     def setUp(self):
         self.model_tester = FuyuModelTester(self)
 
+    def test_mismatching_image_patches(self):
+        config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
+        for model_class in self.all_model_classes:
+            model = model_class(config).to(torch_device)
+            curr_input_dict = copy.deepcopy(input_dict)  # in=place modifications further
+
+            # two image token and two image
+            _ = model(**curr_input_dict)  # successful forward with no modifications
+
+            # remove one image but leave the image token in text
+            input_ids = curr_input_dict["input_ids"]
+            image_patches = curr_input_dict["image_patches"][1:, ...]
+            with self.assertRaises(ValueError):
+                _ = model(input_ids=input_ids, image_patches=image_patches)
+
+            # remove one image token from text
+            input_ids = curr_input_dict["input_ids"][2:]
+            image_patches = curr_input_dict["image_patches"]
+            with self.assertRaises(ValueError):
+                _ = model(input_ids=input_ids, image_patches=image_patches)
+
     @unittest.skip(
         reason="This architecture seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
     )
@@ -232,7 +265,7 @@ def default_processor(self):
 
     @cached_property
     def default_model(self):
-        return FuyuForCausalLM.from_pretrained("adept/fuyu-8b")
+        return FuyuForCausalLM.from_pretrained("adept/fuyu-8b", torch_dtype="float16", device_map=torch_device)
 
     def test_greedy_generation(self):
         processor = self.default_processor
@@ -243,7 +276,9 @@ def test_greedy_generation(self):
 
         text_prompt_coco_captioning = "Generate a coco-style caption.\n"
 
-        inputs = processor(images=image, text=text_prompt_coco_captioning, return_tensors="pt")
+        inputs = processor(images=image, text=text_prompt_coco_captioning, return_tensors="pt").to(
+            torch_device, torch.float16
+        )
         generated_ids = model.generate(**inputs, max_new_tokens=10)
 
         # take the last 8 tokens (in order to skip special \n\x04 characters) and decode them