Repository: huggingface/transformers
Pull request: [`fix`] Use `last_hidden_state` key from `get_image_features` for llama4

# What does this PR do?

Resolves https://github.com/huggingface/transformers/pull/42564#issuecomment-3874606093
#42564 updated `get_image_features` for Llama4, but it erroneously started using `pooler_output` instead of the previous `last_hidden_state`. Additionally, the Llama4VisionModel output has been updated to `BaseModelOutputWithPooling` to match many vision models.

## Reproducer
```python
import torch

from transformers import AutoProcessor, Llama4ForConditionalGeneration

model_id = "hf-internal-testing/tiny-random-llama4"
processor = AutoProcessor.from_pretrained(model_id)
model = Llama4ForConditionalGeneration.from_pretrained(
    model_id,
    attn_implementation="sdpa",  # flex attention / flash_attention_2 do not work, debugging...
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

# Fix meta tensor: convert to float32 with random weights
if model.vision_model.rotary_embedding.freqs_ci.is_meta:
    shape = model.vision_model.rotary_embedding.freqs_ci.shape
    # Note: Ideally should compute proper RoPE frequencies, but using random as requested
    model.vision_model.rotary_embedding.freqs_ci = torch.randn(*shape, dtype=torch.float32, device=model.device)

url1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg"
url2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/datasets/cat_style_layout.png"
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "url": url1},
            {"type": "image", "url": url2},
            {"type": "text", "text": "Can you describe how these two images are similar, and how they differ?"},
        ]
    },
]

inputs = processor.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device)

outputs = model.generate(
    **inputs,
    max_new_tokens=32,
)

response = processor.batch_decode(outputs[:, inputs["input_ids"].shape[-1]:])[0]
print(response)
print(outputs[0])
```
### On `main`:
```python
Traceback (most recent call last):
  File "c:\code\transformers\demo_llama4.py", line 41, in <module>
    outputs = model.generate(
              ^^^^^^^^^^^^^^^
  File "C:\Users\tom\.conda\envs\transformers\Lib\site-packages\torch\utils\_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "C:\code\transformers\src\transformers\generation\utils.py", line 2638, in generate
    result = decoding_method(
             ^^^^^^^^^^^^^^^^
  File "C:\code\transformers\src\transformers\generation\utils.py", line 2833, in _sample
    outputs = self._prefill(input_ids, generation_config, model_kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\code\transformers\src\transformers\generation\utils.py", line 3822, in _prefill
    return self(**model_inputs, return_dict=True)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\tom\.conda\envs\transformers\Lib\site-packages\torch\nn\modules\module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\tom\.conda\envs\transformers\Lib\site-packages\torch\nn\modules\module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\code\transformers\src\transformers\utils\generic.py", line 1016, in wrapper
    outputs = func(self, *args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\code\transformers\src\transformers\models\llama4\modeling_llama4.py", line 1327, in forward
    ).pooler_output
      ^^^^^^^^^^^^^
AttributeError: 'BaseModelOutput' object has no attribute 'pooler_output'
```
#### On this PR:
```
广场 meno बतздрав 공기грамمعграмمعграмمعграмمعграмمعграмمعграмمعграмمعграмمعграмمعграмمعграмمعграмمعграм
tensor([200000, 200005,   1556,  ...,  96359,  20938,  96359], device='cuda:0')
```
(Gibberish obviously as this is a tiny-random model, but it works now!)

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] 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?
@zucchini-nlp @vasqu 

Thank you to @Mecoli1219 for reporting this.

- Tom Aarsen

Changed files:
filename: src/transformers/models/llama4/modeling_llama4.py
patch: @@ -1080,7 +1080,7 @@ def forward(
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
         **kwargs,
-    ) -> BaseModelOutput | tuple[torch.Tensor, ...]:
+    ) -> BaseModelOutputWithPooling | tuple[torch.Tensor, ...]:
         r"""
 
         Example:
@@ -1166,7 +1166,7 @@ def forward(
         if not return_dict:
             return tuple(v for v in [hidden_state, hidden_states, attentions] if v is not None)
 
-        return BaseModelOutput(
+        return BaseModelOutputWithPooling(
             last_hidden_state=hidden_state,
             hidden_states=hidden_states,
             attentions=attentions,
@@ -1324,7 +1324,7 @@ def forward(
                 pixel_values=pixel_values,
                 vision_feature_select_strategy=vision_feature_select_strategy,
                 return_dict=True,
-            ).pooler_output
+            ).last_hidden_state
 
             vision_flat = image_features.view(-1, image_features.size(-1))
             projected_vision_flat = self.multi_modal_projector(vision_flat).to(