Add AudioLDM 2

[sanchit-gandhi]

What does this PR do?

Adds AudioLDM 2 from the paper AudioLDM 2: A General Framework for Audio, Music, and Speech Generation

Architecture

1. CLAP: prompt embedding from text input
2. Flan-T5 Encoder: second prompt embedding from text input
3. Projection model: project the CLAP and T5 prompt embeddings to a shared space and insert special SOS/EOS tokens
4. Language model: generate 8 new hidden-states, conditional on the projected hidden-states. Use these as the final prompt embeddings for the diffusion model
5. UNet, Scheduler: custom UNet architecture (see below)
6. VAE: convert the final latents to a mel-spectrogram
7. Vocoder: convert the mel-spectrogram to an audio waveform

Steps 5-7 is the same as AudioLDM. The remainder are new.

Diagram:

UNet

The vanilla UNet 2D cross-attention layer looks as follows:

• Resnet Block (hidden_states)
• Cross-Attention Transformer Block (hidden_states, encoder_hidden_states)

=> this is the architecture that is used in diffusers when we run the UNet forward with arguments of (hidden_states, encoder_hidden_states)

AudioLDM 2 extends the vanilla UNet architecture to use an additional self-attention layer and two cross-attention layers:

• Resnet Block (hidden_states)
• Self-Attention Transformer Block with double self-attention (hidden_states)
• Cross-Attention Transformer Block 1 (hidden_states, encoder_hidden_states_1)
• Cross-Attention Transformer Block 2 (hidden_states, encoder_hidden_states_2)

=> here we use a different set of encoder hidden-states for cross-attention blocks 1 and 2. The first hidden-states are those obtained from the T5 model. The second hidden-states are those generated from the language model. Also, we don’t want to pass either of these encoder hidden-states to the self-attention layer, since it uses double self-attention.

Checklist

Weight conversion:

• UNet
• VAE
• CLAP
• Flan-T5 (HF Transformers)
• GPT2 (HF Transformers)
• HiFiGAN
• Audio MAE embeddings / projection layers

Forward pass:

• UNet
• VAE
• CLAP
• Flan-T5
• GPT2
• HiFiGAN
• Audio MAE embeddings / projection layers

Pipeline:

• Add pipeline
• Add tests
• Write docs

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[tin2tin]

If possible, then maybe consider doing a pruned version of the model, so it'll be able to run on 6 GB of VRAM?

[sanchit-gandhi]

It's just the slow integration tests and docs to go here! In the interest of time, would you be able to do a first pass of this @sayakpaul @williamberman to confirm that you're happy with the diffusers design? I can then go full send on finishing the last tests and tidying it up for merge tomorrow 🚀 Thanks both for your help so far!

[williamberman]

I think this is good if we need to make a separate class because this is a one off pipeline. I think it's possible @patrickvonplaten might have a different opinion but I don't think it's worth blocking because we can merge into the existing unet after the fact if we need to.

So let's just move forward here and I'll file an issue asking patrick to double check he's ok with it being a separate class when he's back from vacation :)

[williamberman]

#4658 follow up issue

[williamberman]

Few small requests but looks good!

[sanchit-gandhi]

Alright this is working nicely for the base, large and music variants of AudioLDM2! The TTS checkpoints require a VITS text encoder, which will be merged as part of huggingface/transformers#24085 Waiting on this before converting the TTS checkpoints. I've written the AudioLDM2 pipeline in such a way that the TTS models should be compatible directly! (possibly with some minor updates)

[patrickvonplaten]

Great PR and great reviews here - nice!