Refactor `output_hidden_states` to allow index selection
Author: hlky
Pull request: no
Labels: name: Feature request

name: Vision

### Feature request


Models such as Llava use CLIPVision with `output_hidden_states=True` then select the index of `hidden_states` that it needs e.g. `image_outputs.hidden_states[vision_feature_layer]`.

`output_hidden_states` could be refactored to `output_hidden_states: Optional[Union[bool, int]] = None` and `hidden_states` output changed to `hidden_states: Optional[Union[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]] = None`.

The layer index would need to be normalized to account for negative index:
```python
if type(output_hidden_states) is int and output_hidden_states < 0:
    output_hidden_states = min(len(self.layers), len(self.layers) + output_hidden_states + 1) # or config.num_hidden_layers
```

`CLIPEncoder` could be changed from
```python
for idx, encoder_layer in enumerate(self.layers):
    if output_hidden_states:
        encoder_states = encoder_states + (hidden_states,)
```
to
```python
for idx, encoder_layer in enumerate(self.layers):
    if type(output_hidden_states) is int and output_hidden_states == idx:
        encoder_states = hidden_states
    elif output_hidden_states:
        encoder_states = encoder_states + (hidden_states,)
```
and after the loop changed from
```python
if output_hidden_states:
    encoder_states = encoder_states + (hidden_states,)
```
to
```python
if type(output_hidden_states) is int and output_hidden_states == len(self.layers):
    encoder_states = hidden_states
elif output_hidden_states:
    encoder_states = encoder_states + (hidden_states,)
```

Models such as Llava would then be able to do:
```python
image_outputs = self.vision_tower(pixel_values, output_hidden_states=vision_feature_layer)
selected_image_feature = image_outputs.hidden_states
```

### Motivation

Memory efficiency, as per the [comment](https://github.com/huggingface/transformers/blob/68049b17a6bb4c9b0d499e9e77121effa2f5a6c0/src/transformers/models/llava/modeling_llava.py#L454) in Llava `"this is not memory efficient at all (output_hidden_states=True) will save all the hidden stated."`

This would also improve memory efficiency in CLIP Text, specifically in Diffusers pipelines such as Stable Diffusion XL where the penultimate layer is used or when `clip skip` is used.


### Your contribution

I can submit a PR.

Comments:
cc maybe @amyeroberts or @qubvel 

Hi @hlky, thanks for this proposal! 

We can definitely think about ways we could make things more efficient here and this suggestion is a good place to start. 

In terms of the proposal, a few comments: 
* I'd prefer it if `output_hidden_states` could be a list, rather than an int. This way, a subset of encoder_hidden_states can be selected and the type of encoder_hidden_states remains consistent (a tuple). This will be more flexible for different use cases. 
* We'll have to be careful with how this interacts with other parameters which select for features, namely with `out_indices` in backbones
* nit - the preferred way to check types in python is using `isinstance` over `type(x) == y` as it's faster and takes account of things like inheritance. 

If you'd like to open a PR with this, we'd be happy to review! 🤗  

Thanks for the comments!

- A list would indeed be better, VIPLlava currently selects a list of features so it will be useful there.
- I'll keep that in mind, we can certainly extend this to other models and `out_indices` in backbones.
- That was to account for Python's `bool` being a subclass of `int`, `isinstance(True, int)` returns `True`

I'll submit a PR soon.