AudioLDM 2,但更快 ⚡️
AudioLDM 2 是由 Haohe Liu 等人在 AudioLDM 2: Learning Holistic Audio Generation with Self-supervised Pretraining 中提出的。AudioLDM 2 接收文本提示作为输入,并预测相应的音频。它可以生成逼真的音效、人类语音和音乐。
虽然生成的音频质量很高,但使用原始实现进行推理非常慢:生成一个10秒的音频样本需要30秒以上。这归因于多种因素,包括深度多阶段建模方法、大型检查点和未优化的代码。
在这篇博客文章中,我们将展示如何在 Hugging Face 🧨 Diffusers 库中使用 AudioLDM 2,探索一系列代码优化,如半精度、Flash Attention和编译,以及模型优化,如调度器选择和负面提示,以将推理时间缩短超过 10倍,同时对输出音频质量的影响极小。这篇博客文章还附带了一个更精简的 Colab 笔记本,其中包含所有代码但解释较少。
请阅读到最后,了解如何仅用1秒钟生成10秒的音频样本!
模型概述
受 Stable Diffusion 启发,AudioLDM 2 是一种文本到音频的 潜在扩散模型 (LDM),它从文本嵌入中学习连续的音频表示。
整体生成过程总结如下:
CLAP 文本嵌入旨在与相应音频样本的嵌入对齐,而 Flan-T5 嵌入则能更好地表示文本的语义。
- 这些文本嵌入通过独立的线性投影投射到共享嵌入空间:
在 diffusers
实现中,这些投影由 AudioLDM2ProjectionModel 定义。
- 一个 GPT2 语言模型 (LM) 用于自回归生成一个 个新嵌入向量序列,条件基于投影的 CLAP 和 Flan-T5 嵌入
- 生成的嵌入向量 和 Flan-T5 文本嵌入 在 LDM 中用作交叉注意力条件,通过反向扩散过程 去噪 随机潜变量。LDM 在反向扩散过程中运行总共 推理步骤。
其中初始潜在变量 从正态分布 中抽取。LDM 的 UNet 的独特之处在于它接收 两组 交叉注意力嵌入,一组来自 GPT2 语言模型 ,另一组来自 Flan-T5 ,而不是像大多数其他 LDM 那样只有一个交叉注意力条件。
- 最终去噪的潜在变量 传递给 VAE 解码器,以恢复 Mel 频谱图 。
- Mel 频谱图传递给声码器,以获得输出音频波形 。
下图展示了文本输入如何通过文本条件模型,其中两个提示嵌入用作 LDM 中的交叉条件。
有关 AudioLDM 2 模型如何训练的完整详细信息,读者可参考 AudioLDM 2 论文。
Hugging Face 🧨 Diffusers 提供了一个端到端推理管道类 AudioLDM2Pipeline
,它将这个多阶段生成过程封装成一个单一的可调用对象,让您只需几行代码即可从文本生成音频样本。
AudioLDM 2 有三个变体。其中两个检查点适用于一般的文本到音频生成任务。第三个检查点专门用于文本到音乐生成。下表提供了三个官方检查点的详细信息,所有这些检查点都可以在 Hugging Face Hub 上找到。
模型权重 | 任务 | 模型大小 | 训练数据 / h |
---|---|---|---|
cvssp/audioldm2 | 文本到音频 | 1.1B | 1150k |
cvssp/audioldm2-music | 文本到音乐 | 1.1B | 665k |
cvssp/audioldm2-large | 文本到音频 | 1.5B | 1150k |
现在我们已经对 AudioLDM 2 生成过程有了高层次的概述,让我们将这个理论付诸实践!
加载管道
在本教程中,我们将使用基础检查点 cvssp/audioldm2 中的预训练权重来初始化管道。我们可以使用 .from_pretrained
方法加载整个管道,它将实例化管道并加载预训练权重。
from diffusers import AudioLDM2Pipeline
model_id = "cvssp/audioldm2"
pipe = AudioLDM2Pipeline.from_pretrained(model_id)
输出
Loading pipeline components...: 100%|███████████████████████████████████████████| 11/11 [00:01<00:00, 7.62it/s]
可以将管道移到 GPU,就像标准 PyTorch nn 模块一样。
pipe.to("cuda");
太棒了!我们将定义一个生成器并设置一个种子以确保可复现性。这将允许我们调整提示并观察它们通过固定 LDM 模型中的起始潜在变量对生成的影响。
import torch
generator = torch.Generator("cuda").manual_seed(0)
现在我们准备进行第一次生成!我们将在整个笔记本中使用相同的运行示例,其中我们将音频生成条件限定为固定的文本提示,并始终使用相同的种子。audio_length_in_s
参数控制生成音频的长度。它默认为 LDM 训练时的音频长度(10.24秒)。
prompt = "The sound of Brazilian samba drums with waves gently crashing in the background"
audio = pipe(prompt, audio_length_in_s=10.24, generator=generator).audios[0]
输出
100%|███████████████████████████████████████████| 200/200 [00:13<00:00, 15.27it/s]
酷!这次运行大约花了13秒来生成。让我们听听输出的音频。
from IPython.display import Audio
Audio(audio, rate=16000)
听起来很像我们的文本提示!质量很好,但仍有背景噪音的痕迹。我们可以向管道提供一个负面提示,以阻止管道生成某些特征。在这种情况下,我们将传递一个负面提示,阻止模型在输出中生成低质量音频。我们将省略 audio_length_in_s
参数,并让它取其默认值。
negative_prompt = "Low quality, average quality."
audio = pipe(prompt, negative_prompt=negative_prompt, generator=generator.manual_seed(0)).audios[0]
输出
100%|███████████████████████████████████████████| 200/200 [00:12<00:00, 16.50it/s]
使用负面提示时,推理时间没有改变 ;我们只是用负面输入替换了 LDM 的无条件输入。这意味着我们获得的任何音频质量提升都是免费的。
我们来听听生成的音频。
Audio(audio, rate=16000)
音频的整体质量确实有所改善——噪音伪影减少了,音频通常听起来更清晰。 请注意,在实践中,我们通常会发现从第一次生成到第二次生成时,推理时间会有所减少。这是由于我们第一次运行计算时会发生 CUDA“预热”。第二次生成更能代表我们实际的推理时间。
优化 1:Flash Attention
PyTorch 2.0 及更高版本通过 torch.nn.functional.scaled_dot_product_attention
(SDPA) 函数包含了注意力操作的优化和内存高效实现。该函数根据输入自动应用多种内置优化,比普通的注意力实现运行更快且更节省内存。总的来说,SDPA 函数提供了与 Dao 等人在论文 Fast and Memory-Efficient Exact Attention with IO-Awareness 中提出的 Flash Attention 相似的行为。
如果安装了 PyTorch 2.0 并且 torch.nn.functional.scaled_dot_product_attention
可用,Diffusers 将默认启用这些优化。要使用它,只需按照 官方说明 安装 PyTorch 2.0 或更高版本,然后照常使用管道即可 🚀
audio = pipe(prompt, negative_prompt=negative_prompt, generator=generator.manual_seed(0)).audios[0]
输出
100%|███████████████████████████████████████████| 200/200 [00:12<00:00, 16.60it/s]
有关在 diffusers
中使用 SDPA 的更多详细信息,请参阅相应的 文档。
优化 2:半精度
默认情况下,AudioLDM2Pipeline
以 float32(全)精度加载模型权重。所有模型计算也以 float32 精度执行。对于推理,我们可以安全地将模型权重和计算转换为 float16(半)精度,这将改善推理时间和 GPU 内存,而对生成质量的影响微乎其微。
我们可以通过将 torch_dtype
参数传递给 .from_pretrained
来以 float16 精度加载权重。
pipe = AudioLDM2Pipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to("cuda");
我们以 float16 精度运行生成并听取音频输出。
audio = pipe(prompt, negative_prompt=negative_prompt, generator=generator.manual_seed(0)).audios[0]
Audio(audio, rate=16000)
输出
100%|███████████████████████████████████████████| 200/200 [00:09<00:00, 20.94it/s]
音频质量与全精度生成基本没有变化,推理速度提升了大约2秒。根据我们的经验,使用 float16 精度的 diffusers
管道,我们没有发现任何显著的音频质量下降,但却能持续获得显著的推理速度提升。因此,我们建议默认使用 float16 精度。
优化 3:Torch Compile
为了进一步加速,我们可以使用新的 torch.compile
特性。由于管道的 UNet 通常计算量最大,我们将 UNet 用 torch.compile
包装起来,而其余的子模型(文本编码器和 VAE)则保持不变。
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
在用 torch.compile
包装 UNet 后,我们运行的第一个推理步骤通常会很慢,这是由于编译 UNet 的前向传播的开销。让我们运行管道,执行编译步骤,以完成这个耗时较长的运行。请注意,第一个推理步骤可能需要长达2分钟才能编译完成,请耐心等待!
audio = pipe(prompt, negative_prompt=negative_prompt, generator=generator.manual_seed(0)).audios[0]
输出
100%|███████████████████████████████████████████| 200/200 [01:23<00:00, 2.39it/s]
太棒了!现在 UNet 已编译完成,我们可以运行完整的扩散过程并享受更快的推理优势。
audio = pipe(prompt, negative_prompt=negative_prompt, generator=generator.manual_seed(0)).audios[0]
输出
100%|███████████████████████████████████████████| 200/200 [00:04<00:00, 48.98it/s]
生成仅需4秒!实际上,您只需编译 UNet 一次,然后所有后续生成都将获得更快的推理速度。这意味着编译模型所需的时间可以通过后续推理时间的收益得到摊销。有关 torch.compile
的更多信息和选项,请参阅 torch compile 文档。
优化 4:调度器
另一个选项是减少推理步骤的数量。选择更高效的调度器可以帮助减少步骤数量,同时不牺牲输出音频质量。您可以通过调用 schedulers.compatibles
属性来查找哪些调度器与 AudioLDM2Pipeline
兼容。
pipe.scheduler.compatibles
输出
[diffusers.schedulers.scheduling_lms_discrete.LMSDiscreteScheduler,
diffusers.schedulers.scheduling_k_dpm_2_discrete.KDPM2DiscreteScheduler,
diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler,
diffusers.schedulers.scheduling_unipc_multistep.UniPCMultistepScheduler,
diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler,
diffusers.schedulers.scheduling_pndm.PNDMScheduler,
diffusers.schedulers.scheduling_dpmsolver_singlestep.DPMSolverSinglestepScheduler,
diffusers.schedulers.scheduling_heun_discrete.HeunDiscreteScheduler,
diffusers.schedulers.scheduling_ddpm.DDPMScheduler,
diffusers.schedulers.scheduling_deis_multistep.DEISMultistepScheduler,
diffusers.utils.dummy_torch_and_torchsde_objects.DPMSolverSDEScheduler,
diffusers.schedulers.scheduling_ddim.DDIMScheduler,
diffusers.schedulers.scheduling_k_dpm_2_ancestral_discrete.KDPM2AncestralDiscreteScheduler,
diffusers.schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteScheduler]
好的!我们有很长一串调度器可供选择 📝。默认情况下,AudioLDM 2 使用 DDIMScheduler
,需要 200 个推理步骤才能获得高质量的音频生成。然而,性能更好的调度器,如 DPMSolverMultistepScheduler
,仅需要 20-25 个推理步骤 即可达到相似的效果。
让我们看看如何将 AudioLDM 2 调度器从 DDIM 切换到 DPM Multistep。我们将使用 ConfigMixin.from_config()
方法从我们原始的 DDIMScheduler
配置中加载一个 DPMSolverMultistepScheduler
。
from diffusers import DPMSolverMultistepScheduler
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
我们将推理步数设置为 20,并使用新调度器重新运行生成。由于 LDM 潜在变量的形状没有改变,我们无需重复编译步骤。
audio = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=20, generator=generator.manual_seed(0)).audios[0]
输出
100%|███████████████████████████████████████████| 20/20 [00:00<00:00, 49.14it/s]
生成音频仅用了不到 1 秒!让我们听听生成的音频。
Audio(audio, rate=16000)
或多或少与我们的原始音频样本相同,但生成时间只有一小部分!🧨 Diffusers 管道设计为 可组合,让您可以轻松地将调度器和其他组件替换为性能更好的对应物。
内存如何?
我们想要生成的音频样本的长度决定了我们在 LDM 中去噪的潜在变量的 宽度。由于 UNet 中交叉注意力层的内存随着序列长度(宽度)的平方而变化,因此生成非常长的音频样本可能会导致内存不足错误。我们的批处理大小也决定了我们的内存使用量,控制了我们生成的样本数量。
我们已经提到,以 float16 半精度加载模型可以大幅节省内存。使用 PyTorch 2.0 SDPA 也能改善内存使用,但这对于极长的序列长度可能还不够。
让我们尝试生成一个时长 2.5 分钟(150 秒)的音频样本。我们还将通过设置 num_waveforms_per_prompt
=4
来生成 4 个候选音频。当 num_waveforms_per_prompt
>1
时,会在生成的音频和文本提示之间执行自动评分:音频和文本提示被嵌入到 CLAP 音频-文本嵌入空间中,然后根据它们的余弦相似度分数进行排名。我们可以将位置 0
的波形视为“最佳”波形。
由于我们改变了 UNet 中潜在变量的宽度,我们将不得不使用新的潜在变量形状执行另一个 torch 编译步骤。为了节省时间,我们将重新加载没有 torch 编译的管道,这样我们就不会在第一次遇到漫长的编译步骤。
pipe = AudioLDM2Pipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to("cuda")
audio = pipe(prompt, negative_prompt=negative_prompt, num_waveforms_per_prompt=4, audio_length_in_s=150, num_inference_steps=20, generator=generator.manual_seed(0)).audios[0]
输出: ```
OutOfMemoryError 回溯(最近一次调用):
23 帧 /usr/local/lib/python3.10/dist-packages/torch/nn/modules/linear.py 在 forward(self, input) 112 113 def forward(self, input: Tensor) -> Tensor: --> 114 return F.linear(input, self.weight, self.bias) 115 116 def extra_repr(self) -> str
内存不足错误:CUDA 内存不足。尝试分配 1.95 GiB。GPU 0 的总容量为 14.75 GiB,其中 1.66 GiB 可用。进程 414660 使用了 13.09 GiB 内存。在已分配的内存中,PyTorch 分配了 10.09 GiB,PyTorch 预留但未分配 1.92 GiB。如果预留但未分配的内存很大,请尝试设置 max_split_size_mb 以避免碎片化。请参阅内存管理和 PYTORCH_CUDA_ALLOC_CONF 的文档。
Unless you have a GPU with high RAM, the code above probably returned an OOM error. While the AudioLDM 2 pipeline involves
several components, only the model being used has to be on the GPU at any one time. The remainder of the modules can be
offloaded to the CPU. This technique, called *CPU offload*, can reduce memory usage, with a very low penalty to inference time.
We can enable CPU offload on our pipeline with the function [enable_model_cpu_offload()](https://huggingface.co/docs/diffusers/main/en/api/pipelines/audioldm2#diffusers.AudioLDM2Pipeline.enable_model_cpu_offload):
```python
pipe.enable_model_cpu_offload()
使用 CPU 卸载运行生成与之前相同。
audio = pipe(prompt, negative_prompt=negative_prompt, num_waveforms_per_prompt=4, audio_length_in_s=150, num_inference_steps=20, generator=generator.manual_seed(0)).audios[0]
输出
100%|███████████████████████████████████████████| 20/20 [00:36<00:00, 1.82s/it]
这样,我们就可以一次性生成四个样本,每个样本时长 150 秒!使用大型 AudioLDM 2 检查点将导致比基础检查点更高的整体内存使用量,因为 UNet 的大小是其两倍多(7.5 亿参数对 3.5 亿参数),因此这种内存节省技巧在这里特别有用。
结论
在这篇博客文章中,我们展示了 🧨 Diffusers 提供的四种开箱即用的优化方法,将 AudioLDM 2 的生成时间从 14 秒缩短到不到 1 秒。我们还强调了如何使用节省内存的技巧,例如半精度和 CPU 卸载,以减少长音频样本或大检查点大小的峰值内存使用量。
博客文章作者:Sanchit Gandhi。非常感谢 Vaibhav Srivastav 和 Sayak Paul 的建设性意见。频谱图图片来源:了解 Mel 频谱图。波形图片来源:阿尔托大学语音处理。