Expanding inputs for image tokens in BLIP-2 breaks batch inference.
Author: pspdada
Pull request: no
Labels: name: bug

name: Vision

### System Info

- `transformers` version: 4.45.2
- Platform: Linux-5.15.0-120-generic-x86_64-with-glibc2.35
- Python version: 3.10.15
- Huggingface_hub version: 0.25.2
- Safetensors version: 0.4.5
- Accelerate version: 1.0.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.0+cu124 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?:  None
- Using GPU in script?: None
- GPU type: NVIDIA A100-PCIE-40GB

### Who can help?

@zucchini-nlp 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I tried to use the model `Salesforce/blip2-flan-t5-xl`, but encountered a warning.

> Expanding inputs for image tokens in BLIP-2 should be done in processing. Please follow instruction here (https://gist.github.com/zucchini-nlp/e9f20b054fa322f84ac9311d9ab67042) to update your BLIP-2 model. Using processors without these attributes in the config is deprecated and will throw an error in v4.47.

The code mentioned there that need to be added is:
```
# Load your model and processor and run the following to update BLIP-2 model
# It will update file in your repo by adding new args in configs and resizing embedding layer
# Then you'll be able to run BLIP-2 without warnings/errors
from transformers import AddedToken

processor.num_query_tokens = model.config.num_query_tokens
image_token = AddedToken("<image>", normalized=False, special=True)
processor.tokenizer.add_tokens([image_token], special_tokens=True)

model.resize_token_embeddings(len(processor.tokenizer), pad_to_multiple_of=64) # pad for efficient computation
model.config.image_token_index = len(processor.tokenizer) - 1

model.push_to_hub("YOUR-REPO")
processor.push_to_hub("YOUR-REPO")
```

After adding the code mentioned there, I received error when perform batch inference. What should I do?

```
import torch
from torch import device
from transformers import AutoProcessor, Blip2Processor, Blip2ForConditionalGeneration, AddedToken
from PIL import Image
from PIL.ImageFile import ImageFile

from model.utils.init_model import *


class BLIP2Generator:

    def __init__(self, model_dir: str = '', device: str = 'cpu'):
        self.model_dir = model_dir
        self.device = device
        self.model, self.processor = self._create()

    def _create(self) -> tuple[Blip2ForConditionalGeneration, Blip2Processor]:
        """
        Create a BLIP2 model and processor for text generation based on image input.

        Returns:
            A tuple containing the model and processor.
        """
        model_name = f"Salesforce/blip2-flan-t5-xl"

        processor = AutoProcessor.from_pretrained(
            model_name,
            cache_dir=self.model_dir,
        )
        model= Blip2ForConditionalGeneration.from_pretrained(
            model_name,
            cache_dir=self.model_dir,
        ).to(self.device)

        # Following https://gist.github.com/zucchini-nlp/e9f20b054fa322f84ac9311d9ab67042#file-update_blip-py
        processor.num_query_tokens = model.config.num_query_tokens
        image_token = AddedToken("<image>", normalized=False, special=True)
        processor.tokenizer.add_tokens([image_token], special_tokens=True)

        # Resize token embeddings for efficient computation
        model.resize_token_embeddings(len(processor.tokenizer), pad_to_multiple_of=64)
        model.config.image_token_index = len(processor.tokenizer) - 1

        return model, processor

    def gen(
        self,
        images: ImageFile | list[ImageFile],
        questions: str | list[str] = "Please describe this image in detail.",
        max_new_tokens: int = 200,
    ) -> str:
        if not isinstance(images, list):
            images = [images]
        if not isinstance(questions, list):
            questions = [questions] * len(images)

        assert len(images) == len(questions), "Number of images and questions must match."
        print(f"Processing {len(images)} images and {len(questions)} questions.")
        with torch.inference_mode():
            inputs = self.processor(
                images=images,
                text=questions,
                return_tensors="pt",
                return_token_type_ids=False,
                padding=True,
            ).to(self.device, torch.bfloat16)

            # Debug: Print input shapes for checking
            print(f"Pixel values shape: {inputs['pixel_values'].shape}")
            print(f"Input IDs shape: {inputs['input_ids'].shape}")

            output = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
            )

            decoded_output = self.processor.batch_decode(
                output,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
            )
        return decoded_output if len(decoded_output) > 1 else decoded_output[0]


def demo_blip2():
    blip2 = BLIP2Generator(model_dir=model_dir, device=device)
    image_1 = Image.open('./demo/material/1.jpg')
    image_2 = Image.open('./demo/material/2.jpg')

    question_1 = "How many birds in the picture?"
    question_2 = "What is on the hand of the girl?"
    result = blip2.gen([image_1, image_2], questions=[question_1, question_2])
    print(result)

```


The output is:
```
Processing 2 images and 2 questions.
Pixel values shape: torch.Size([2, 3, 224, 224])
Input IDs shape: torch.Size([1, 42])
Traceback (most recent call last):
  File "/root/llm-project/LVLM/demo.py", line 4, in <module>
    demo_blip2()
  File "/root/llm-project/LVLM/demo/blip2.py", line 16, in demo_blip2
    result = blip2.gen([image_1, image_2], questions=[question_1, question_2])
  File "/root/llm-project/LVLM/model/blip2_generator.py", line 102, in gen
    output = self.model.generate(
  File "/root/anaconda3/envs/LVLM/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/root/anaconda3/envs/LVLM/lib/python3.10/site-packages/transformers/models/blip_2/modeling_blip_2.py", line 2324, in generate
    inputs_embeds[special_image_mask] = language_model_inputs.flatten()
RuntimeError: shape mismatch: value tensor of shape [131072] cannot be broadcast to indexing result of shape [65536]
```

I provide the same number of images and texts (equally 2), but the processor don't work as espected, since 
`Pixel values shape: torch.Size([2, 3, 224, 224]) Input IDs shape: torch.Size([1, 42])”`
After remove the code mentioned ahead, the batch inference get correctly down, and get
`Pixel values shape: torch.Size([2, 3, 224, 224]) Input IDs shape: torch.Size([2, 10])`



### Expected behavior

What should I do?

Comments:
The input_ids tensor may lack a value equal to self.config.image_token_index, resulting in an empty special_image_mask.
 inputs_embeds[special_image_mask] = language_model_inputs.flatten()
    ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^
RuntimeError: shape mismatch: value tensor of shape [65536] cannot be broadcast to indexing result of shape [0]

How to solve the problem.

> The input_ids tensor may lack a value equal to self.config.image_token_index, resulting in an empty special_image_mask. inputs_embeds[special_image_mask] = language_model_inputs.flatten() ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^ RuntimeError: shape mismatch: value tensor of shape [65536] cannot be broadcast to indexing result of shape [0]
> 
> How to solve the problem.

Hi, I haven't encountered the problem you mentioned. You can ask zucchini-nlp by @ ing or open an issue and provide more detailed information.