Add UnivNet Vocoder Model for Tortoise TTS Diffusers Integration
Author: dg845
Pull request: yes

# What does this PR do?

This PR adds the UnivNet GAN vocoder model ([paper](https://arxiv.org/pdf/2106.07889.pdf), [code](https://github.com/mindslab-ai/univnet)) to `transformers`, which is the vocoder used in the Tortoise TTS text-to-speech model ([paper](https://arxiv.org/pdf/2305.07243.pdf), [code](https://github.com/neonbjb/tortoise-tts)) which is currently being integrated into `diffusers`. See [this issue](https://github.com/huggingface/diffusers/issues/3891) in `diffusers`.

![univnet_model_architecture](https://github.com/huggingface/transformers/assets/58458699/8a33190a-cd52-4e81-a6ed-2fc921b9f86f)

## 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#start-contributing-pull-requests),
      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.
- [x] 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).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@sanchit-gandhi 
@susnato 


Comments:
Hi @dg845 if you are planning to add it to the `models` folder, then I think it should have a doc file(`univnet.md`) in the docs.

For now I've added the UnivNet code to the `/src/transformers/models/univnet/` directory. @sanchit-gandhi, since the UnivNet model isn't technically a transformer model (in that it doesn't use any attention mechanisms), is this the best place to put it? For example, the [`SpeechT5HifiGan`](https://huggingface.co/docs/transformers/main/model_doc/speecht5#transformers.SpeechT5HifiGan) vocoder is in `/src/transformers/models/speecht5/` along with the other SpeechT5 models, but I assume most of the other Tortoise TTS code will go into `diffusers` rather than `transformers`.

Nice start @dg845! Yep fine to have it as a standalone model - we have ResNet in transformers as well which is not strictly attention-based.

Hi @sanchit-gandhi, I think the PR is ready for review.

The following are the differences between the [`SpeechT5HifiGan`](https://huggingface.co/docs/transformers/main/model_doc/speecht5#transformers.SpeechT5HifiGan) and the `UnivNetGan` model:

- The `SpeechT5HifiGan` outer residual blocks* (that is, [`HifiGanResidualBlock`](https://github.com/huggingface/transformers/blob/main/src/transformers/models/speecht5/modeling_speecht5.py#L3074)) upsamples the number of hidden channels between each outer residual block, but the `UnivNetGan` outer residual blocks* (`UnivNetLVCBlock`) keep the number of hidden channels constant.
- Although the structures of the inner residual blocks (for UnivNet, the `UnivNetLVCResidualBlock` module) are similar: `LReLU` => dilated `Conv1d` => `LReLU` => `Conv1d` => skip connection, the UnivNet model uses a [location variable convolutional layer](https://arxiv.org/pdf/2102.10815.pdf) followed by a [gated activation unit](https://proceedings.neurips.cc/paper_files/paper/2016/file/b1301141feffabac455e1f90a7de2054-Paper.pdf) in place of the second `Conv1d` layer.
- Accordingly, each outer residual block (`UnivNetLVCBlock`) in UnivNet has a kernel predictor residual network (`UnivNetKernelPredictor`) to predict the kernels and biases for the location variable convolutional layer in each inner residual block in the main resnet.
- In addition to a conditioning log-mel `spectrogram`, UnivNet takes in a noise sequence as input. The `noise_waveform` is the input to the "main" resnet (e.g. the stack of `UnivNetLVCResidualBlock`s), while the conditioning `spectrogram` is the input to the kernel predictor in each `UnivNetLVCBlock`.

(*) "Outer residual block" is a bit of a misnomer, since for both blocks in question (`HifiGanResidualBlock`, `UnivNetLVCBlock`) there's no skip connection between the input to the block and the main computation in the block.

Also, I'm not sure why `utils/check_table.py` is failing. I ran `make fix-copies` to create a table entry for UnivNet in `docs/source/en/index.md`, and then added a checkmark for PyTorch support, but for some reason `check_table.py` doesn't seem to like that.

> Also, I'm not sure why `utils/check_table.py` is failing.

`utils/check_table.py` is no longer failing after I merged `main` into the PR branch. Running `make fix-copies` adds an entry for UnivNet, but I'm not sure why it doesn't add a checkmark in the "PyTorch support" column, perhaps the model information is mis-configured.

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_24799). All of your documentation changes will be reflected on that endpoint.

All of the current tests (including the integration tests) should be passing now. The code to calculate reference slices from the [reference implementation](https://github.com/maum-ai/univnet) is in [this repository](https://github.com/dg845/univnet/tree/inference-test).

Hi @ArthurZucker, thanks for reviewing! I think I have addressed all of the review comments. In particular, the model class names should be consistently in the form `UnivNetGan*` and the model implementation should be simplified using the configs (although I am currently keeping the `apply_weight_norm`/`remove_weight_norm` methods, see https://github.com/huggingface/transformers/pull/24799#discussion_r1326830801).

From my side, I have the following questions:

- https://github.com/huggingface/transformers/pull/24799#discussion_r1327008418 (regarding using `global_rng = random.Random()` with `floats_list`)
- https://github.com/huggingface/transformers/pull/24799#discussion_r1327894425 (regarding isolating copied feature extractor tests)
- https://github.com/huggingface/transformers/pull/24799#discussion_r1327892862 (regarding the proper org to move the checkpoints to)
- https://github.com/huggingface/transformers/pull/24799#discussion_r1326841344 (regarding the usage of `_no_split_modules` for `accelerate` support)

Gentle ping @ArthurZucker for a follow up review.

Hey! Sorry I'll get this done by the end of the week 😉 

Thanks for iterating @dg845! Requesting a final review from @ArthurZucker when he gets the chance 🤗

Hi @sanchit-gandhi, I have a question about the design of a potential `batch_decode` method for the feature extractor (as a follow up to https://github.com/huggingface/transformers/pull/24799#discussion_r1344041876). The simplest implementation of a `batch_decode` method would be something like

```python
class UnivNetGanFeatureExtractor(SequenceFeatureExtractor):
    ...
    def batch_decode(self, audio):
        return audio[..., :-(self.pad_end_length * self.hop_length)]
    ...
```

which directly follows the current usage example. The downside here is that this `batch_decode` method should only be called if we padded the end of the spectrogram when calling the feature extractor, but it doesn't have a way to know whether this is the case. Having something like a `padding_mask` argument would give it the necessary information it needs, but UnivNet doesn't use any attention mechanisms so the feature extractor currently doesn't produce an attention mask and `UnivNetGan.forward` currently doesn't take a mask argument. I'm currently considering several designs and am not sure which one is best:

1. The feature extractor does not output a `padding_mask` and `UnivNetGan.forward` does not take a `padding_mask` argument. `batch_decode` is implemented as above, where it only strips the end padding and should only be called if `pad_end=True` when the feature extractor was called.
2. The feature extractor outputs a `padding_mask` and `UnivNetGan.forward` takes a `padding_mask` argument, but does nothing with it. The `batch_decode` method takes in both `audio` and `padding_mask` arguments and uses the mask to remove padding from the `audio` output.
3. The feature extractor outputs a `padding_mask` and `UnivNetGan.forward` takes a `padding_mask` argument and uses it to remove padding from the output waveform right before returning it. There is no need for a `batch_decode` method in this case.

I'm leaning toward (3) because it's the most user-friendly, since all padding is handled under the hood and there's no need to additionally call `feature_extractor.batch_decode`.


Thanks for the comprehensive summary @dg845! In the proposed implementation for 3, would the model still return a PyTorch tensor of outputs? If we have different padding for each element in the batch, we'll need to strip each audio by different amounts, so it won't be possible to output a PyTorch tensor from the model? The sequence length dimensions will be different for each element in the batch, so we can't output a tensor with the same seq len dim across items in the batch. We try and keep a `tensor in -> tensor out` approach to the modelling code to make it easy to compile and handle putting/removing inputs/outputs from torch device respectively.

So I think we'll need to go with 2, and have the feature extractor to strip this padding a return a list of numpy arrays (i.e. a ragged list of audio arrays). WDYT? 

Hi @dg845 and @sanchit-gandhi, IMO, we should lean towards a solution that looks like 3, but still allows to return Pytorch tensors. It would go like this:
- if `return_dict=False`, it returns a tuple `(waveform, waveform_length)` (or similar naming) with `waveform` of shape `(batch_size, max_waveform_length)` and `waveform_length` of shape `(batch_size,)` which stores the waveform lengths that you would compute from the `padding_mask`
- if `return_dict=True`, returns a `UnivNetOutput` which inherits from `ModelOuput`, with `waveform` and `waveform_length` as keys. It would look a bit like [VitsModelOuput](https://github.com/huggingface/transformers/blob/e8fdd7875def7be59e2c9b823705fbf003163ea0/src/transformers/models/vits/modeling_vits.py#L52) but with less keys!

WDYT of going for something like that ?

In that way, it would look a lot like the regular outputs of NLP models in transformers, while keeping our requirements. 

@sanchit-gandhi @ylacombe thanks for the input! I have added an implementation based on https://github.com/huggingface/transformers/pull/24799#issuecomment-1755905748 with the following details:

- `UnivNetFeatureExtractor.__call__` can return a `padding_mask` argument (the `attention_mask` output from `SequenceFeatureExtractor.pad`). When `pad_end = True` we will pad with audio silence instead of spectrogram silence following https://github.com/huggingface/transformers/pull/24799#discussion_r1325158242.
- `UnivNetModel.forward` accepts a `padding_mask` argument and will use it to calculate the length of each original unpadded waveform and return it as `waveform_lengths`. (The audio output `waveforms` will always be a batched, padded tensor of waveforms).
- `UnivNetFeatureExtractor.batch_decode` takes in the outputs of `UnivNetModel.forward` and returns a ragged list of 1D waveform arrays with padding removed if `waveform_lengths` is available.

Nice - I think this is a good design! Gently pinging @ArthurZucker for a final review here when you get the chance 🙌

@ArthurZucker @sanchit-gandhi @ylacombe I think one thing left to resolve is where to put the UnivNet model checkpoint (currently at [`dg845/univnet-dev`](https://huggingface.co/dg845/univnet-dev)). I'm not sure which org to put the model checkpoint under since the [original paper](https://arxiv.org/pdf/2106.07889.pdf) is from Kakao, but the checkpoint is from an [unofficial implementation](https://github.com/maum-ai/univnet) by maum.ai (see https://github.com/huggingface/transformers/pull/24799#discussion_r1325159440, https://github.com/huggingface/transformers/pull/24799#discussion_r1327892862).

We could reach out to maum.ai to ask them if they can create and org (if not already) and host the weights

@ArthurZucker Sounds good! I believe that they have an org at https://huggingface.co/maum-ai.

Hi @ArthurZucker @sanchit-gandhi @ylacombe, is there anything I can do to help out with transferring the checkpoint weights? (As a note, the checkpoint weights are currently stored at [`dg845/univnet-dev`](https://huggingface.co/dg845/univnet-dev) [with the model card written] and this is the checkpoint identifier used e.g. in the integration tests.)

Hi @dg845, I've contacted some people from maum-ai to move the weights to their organization (without any response yet)!

Hi @ArthurZucker @sanchit-gandhi @ylacombe would it be possible to merge this PR? @susnato and I have made a lot of progress on the tortoise-tts PR over at `diffusers`: https://github.com/huggingface/diffusers/pull/4106 and it would be helpful to have this PR merged to test the pipeline with the UnivNet vocoder.

Hey both! Yeah no problem, let's use the current path for the checkpoints and merge for now as they are slow to respond! 

Last nit is, would you mind rebasing on main to make sure you have the correct styling? 🙏🏻 

Hi @ArthurZucker, I have rebased on `main` and the CI is green :).

Thanks a lot!