Pixtral: vectorize patch embeddings and enable tests
Author: zucchini-nlp
Pull request: yes

# What does this PR do?

Continuation on the discussion from https://github.com/huggingface/transformers/pull/35110.

This PR gets rid of the loop over each image in the input for Pixtra and makes it more aligned with other VLMs. Now the model will pad the image on h/w dimensions and unpad it back after the vision patch embedding layer. That also helps us to get rid of extra dimension errors on processing code we've been having lately and remove the `BacthFeatureMix`

Test with the demo script from https://huggingface.co/mistral-community/pixtral-12b, as we don't have any test for pixtral as VLM at the moment. The generations match on text level

cc @Rocketknight1 wdyt about this? The design might need some changes as I had to make llava accept extra kwargs (`image_sizes`) to make the model work

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_35122). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@zucchini-nlp One concern I have about `padding the images before the vision backbone and then unpadding them back after` is that then you'll be always padding every image in the batch to the max size unnecessarily (especially when models like Pixtral support 10+ multi-image in one inference call), which IMO could be a bit problematic. At that point it maybe just better to flatten and concat pixel values to the batched tensor of shape (total_num_patches, C, H, W) since you will need to pass in additional information (image_sizes in this case) anyways.

@ywang96 hmm, indeed could induce some extra overhead by patching everything to "max-size", lemme see how much overhead comes from it

> better to flatten and concat pixel values to the batched tensor of shape (total_num_patches, C, H, W) s

Not sure I got this. Currently we return pixels of shape (batch-size, C, H, W). What is the difference between making it `total_num_patches`? IIUC you mean smth like flattening on `H*W` and concating on that dim similar to what Qwen2-VL does?

Made some evals for memory usage with inputs where all images are low resolution (100 x 100) and there is one with the highest resolution possible (1024 x 1024). The memory usage between old and the new processing logic is almost same, excluding memory used by model weights it is 

Old-processor -> 41470 MiB
New processor -> 41875 MiB


Higher batch size results in OOM in both cases in A100 80GB

```python
from PIL import Image
import torch
from transformers import AutoProcessor, LlavaForConditionalGeneration

model_id = "mistral-community/pixtral-12b"
processor = AutoProcessor.from_pretrained(model_id, use_fast=True)
processor.tokenizer.pad_token_id = processor.tokenizer.eos_token_id
processor.tokenizer.padding_side = "left"


url_dog = "https://picsum.photos/id/237/100/100"
url_mountain = "https://picsum.photos/seed/picsum/100/100"
url_stop = "https://picsum.photos/id/237/2000/2000"

chat = [
    {
      "role": "user", "content": [
        {"type": "text", "content": "Can this animal"}, 
        {"type": "image"},
        {"type": "text", "content": " live here?"}, 
        {"type": "image"},
      ]
    }
]
prompt = processor.apply_chat_template(chat)

chat_2 = [
    {
      "role": "user", "content": [
        {"type": "text", "content": "What do you see here"}, 
        {"type": "image"},
      ]
    }
]
prompt_2 = processor.apply_chat_template(chat_2)


begin = torch.cuda.memory_allocated()

prompts = [prompt] * 40 + [prompt_2]
images = [[url_dog, url_mountain]] * 40 + [[url_stop]]


inputs = processor(text=prompts, images=images, padding=True, return_tensors="pt")

model = LlavaForConditionalGeneration.from_pretrained(model_id, device_map="cuda:0", torch_dtype="float16")

inputs = inputs.to(model.device, torch.float16)
generate_ids = model.generate(**inputs, max_new_tokens=50)
output = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]

end = torch.cuda.max_memory_allocated()
print((end - begin) // 1024 ** 2, "MiB")
```



Ready for review! Quick update:

This PR modifies Pixtral model code to act similarly to other VLMs that divide each image into patches. So, now we pad the inputs to `max_patch_size` and will have a more standard 4D `pixel_values`. After running the patch embeddings, the `hidden_states` can be unpadded back. Similar method is used in llava-next based models and in Emu3 (maybe also Molmo 🤔 )

Plus, I enabled the tests that were skipped and fixed what needed to be fixed. All slow tests for Pixtral are passing + ran a sanity check by generating with "pixtral-12b" 

run slow pixtral

This comment contains run-slow, running the specified jobs: ['models/pixtral'] ...