Add offload support to Bark
Author: ylacombe
Pull request: yes

# What does this PR do?

Bark is a TTS model recently added by #24086. 

It is made of is of 4 main sub-models, which are sequentially called during generation (`BarkModel.generate`). When a sub-model is used, the other models remain unused, taking up a lot of space on the GPU.

This PR thus propose a simple, yet effective, handmade offloading of sub-models.

```python
from transformers import BarkModel, BarkProcessor

# no need to load the model onto GPU, it will be done inside the generate function
model = BarkModel.from_pretrained("suno/bark")
processor = BarkProcessor.from_pretrained("suno/bark")

# works if a GPU is available. Throws a warning if not, and uses default behavior.
device = "cuda"

# input must be put onto the right device, i.e onto the GPU.
input = processor("Hey, it's a test").to(device)

# one simple additional argument
output = model.generate(**input, offload_to_cpu = True)

# output is loaded onto GPU as well
```

With this PR, GPU footprint is around 60% lower, while being less than 10% slower, based on a benchmark I've done and that will be shared soon (`batch_size = 1` on a single TITAN RTX).


TODO:
- [x] write tests


## Who can review?

Hey @sanchit-gandhi and @amyeroberts, I'm tagging you because you were the ones reviewing  #24086 and because the changes are really small! 
I'm wondering whether the code I've written conforms to the transformer paradigm and whether I need to raise additional warnings or errors in extreme cases!
Many thanks!


Comments:
_The documentation is not available anymore as the PR was closed or merged._

Why not leverage Accelerate here as is done in Diffusers pipeline? (cc @patrickvonplaten )

Hi @sgugger, thank you for taking part! 
Diffusers pipeline is using [`accelerate's cpu_offload`](https://github.com/huggingface/accelerate/blob/69e4c3c54da3201eda288b500d138761e7a5221c/src/accelerate/big_modeling.py#L146) which itself is using [`accelerate's dispatch model`](https://github.com/huggingface/accelerate/blob/69e4c3c54da3201eda288b500d138761e7a5221c/src/accelerate/big_modeling.py#L146).  

The only issue here is that accelerate's cpu offload is done whenever calling the forward pass of the model. Here, some sub-model are auto-regressive so it would (if I'm not wrong) offload/onload over and over again while forwarding during the `generate` call. This would be sub-optimal and time-consuming, and would remove the benefits of my version of `offload_to_cpu`.

BTW, I'm open to suggestions on how to best make it work with accelerate, if it's possible! @muellerzr (hi !) or @sgugger , do you have some ideas?


No Diffusers uses [`cpu_offload_with_hook`](https://github.com/huggingface/accelerate/blob/69e4c3c54da3201eda288b500d138761e7a5221c/src/accelerate/big_modeling.py#L194) which gives you a hook to pass along to the next model in the pipeline. This way you can have the auto-regressive model in the middle be called several times and only when we go to the next model is it offloaded back to the CPU, which looks like what you are doing here in this PR.

Nice, it's indeed what I'm doing @sgugger , many thanks for your help! I'll adapt the PR

Here an example: https://github.com/huggingface/diffusers/blob/fa356bd4da2593c2b91f76c1f63b6238249ec001/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L225

@sgugger , thanks for the quick review! I've applied your last nit comments.



Hi @sgugger and @sanchit-gandhi , again thanks for the review! 
I've applied your last nit comments.
Don't hesitate to reach out to me if I need to refactor or improve something!