AudioLDM 2 推理提速十倍：Diffusers 优化教程

AudioLDM 2 在 AudioLDM 2: Learning Holistic Audio Generation with Self-supervised Pretraining 中被提出，作者为 Haohe Liu 等人。AudioLDM 2 以文本提示作为输入，并预测相应的音频。它可以生成逼真的音效、人类语音和音乐。

虽然生成的音频质量很高，但使用原始实现进行推理非常慢：一个 10 秒的音频样本需要超过 30 秒才能生成。这是由于多种因素造成的，包括深度多阶段建模方法、庞大的检查点大小以及未优化的代码。

在这篇博客文章中，我们展示了如何在 Hugging Face 🧨 Diffusers 库中使用 AudioLDM 2，探索一系列代码优化，如半精度、flash attention 和编译，以及模型优化，如调度器选择和负面提示，将推理时间减少超过 10 倍，而输出音频的质量下降最小。这篇博客文章还附带了一个更精简的 Colab notebook，其中包含所有代码但解释较少。

读到最后，了解如何在短短 1 秒内生成一个 10 秒的音频样本！

模型概述

受 Stable Diffusion 启发，AudioLDM 2 是一个文本到音频的潜在扩散模型（LDM），它从文本嵌入中学习连续的音频表示。

整体生成过程总结如下：

1. 给定文本输入 x\boldsymbol{x}，使用两个文本编码器模型来计算文本嵌入：CLAP 的文本分支，以及 Flan-T5 的文本编码器

E1=CLAP(x);E2=T5(x) \boldsymbol{E}_{1} = \text{CLAP}\left(\boldsymbol{x} \right); \quad \boldsymbol{E}_{2} = \text{T5}\left(\boldsymbol{x}\right)

CLAP 文本嵌入被训练为与相应音频样本的嵌入对齐，而 Flan-T5 嵌入则更好地表示文本的语义。

1. 这些文本嵌入通过各自的线性投影被投影到一个共享的嵌入空间：

P1=WCLAPE1;P2=WT5E2 \boldsymbol{P}_{1} = \boldsymbol{W}_{\text{CLAP}} \boldsymbol{E}_{1}; \quad \boldsymbol{P}_{2} = \boldsymbol{W}_{\text{T5}}\boldsymbol{E}_{2}

在 diffusers 实现中，这些投影由 AudioLDM2ProjectionModel 定义。

1. 使用 GPT2 语言模型（LM）以自回归方式生成 NN 个新嵌入向量序列，条件是投影后的 CLAP 和 Flan-T5 嵌入：

E~i=GPT2(P1,P2,E~1:i−1)for i=1,…,N \tilde{\boldsymbol{E}}_{i} = \text{GPT2}\left(\boldsymbol{P}_{1}, \boldsymbol{P}_{2}, \tilde{\boldsymbol{E}}_{1:i-1}\right) \qquad \text{for } i=1,\dots,N

1. 生成的嵌入向量 E~1:N\tilde{\boldsymbol{E}}_{1:N} 和 Flan-T5 文本嵌入 E2\boldsymbol{E}_{2} 被用作 LDM 中的交叉注意力条件，LDM 通过反向扩散过程对随机潜在变量进行去噪。LDM 在反向扩散过程中总共运行 TT 个推理步骤：

zt=LDM(zt−1∣E~1:N,E2)for t=1,…,T \boldsymbol{z}_{t} = \text{LDM}\left(\boldsymbol{z}_{t-1} | \tilde{\boldsymbol{E}}_{1:N}, \boldsymbol{E}_{2}\right) \qquad \text{for } t = 1, \dots, T

其中初始潜变量 z0\boldsymbol{z}_{0} 从正态分布 N(0,I)\mathcal{N} \left(\boldsymbol{0}, \boldsymbol{I} \right) 中采样得到。 LDM 的 UNet 的独特之处在于它接收两组交叉注意力嵌入，一组是来自 GPT2 语言模型的 E~1:N\tilde{\boldsymbol{E}}_{1:N}，另一组是来自 Flan-T5 的 E2\boldsymbol{E}_{2}，而大多数其他 LDM 只使用一组交叉注意力条件。

1. 最终去噪后的潜变量 zT\boldsymbol{z}_{T} 被传递给 VAE 解码器以恢复 Mel 频谱图 s\boldsymbol{s}：

s=VAEdec(zT) \boldsymbol{s} = \text{VAE}_{\text{dec}} \left(\boldsymbol{z}_{T}\right)

1. Mel 频谱图被传递给声码器以获得输出音频波形 y\mathbf{y}：

y=Vocoder(s) \boldsymbol{y} = \text{Vocoder}\left(\boldsymbol{s}\right)

下图展示了文本输入如何通过文本条件模型传递，其中两个提示嵌入被用作 LDM 中的交叉条件：

关于 AudioLDM 2 模型训练方式的完整细节，读者可参阅 AudioLDM 2 论文。

Hugging Face 🧨 Diffusers 提供了一个端到端推理流水线类 AudioLDM2Pipeline，将这一多阶段生成过程封装为单个可调用对象，使你只需几行代码即可从文本生成音频样本。

AudioLDM 2 有三种变体。其中两个检查点适用于文本到音频生成的通用任务。第三个检查点专门针对文本到音乐生成进行训练。有关这三个官方检查点的详细信息，请参见下表，它们都可以在 Hugging Face Hub 上找到：

现在我们已经对 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]

流水线可以像标准 PyTorch nn 模块一样移动到 GPU 上：

pipe.to("cuda");

很好！我们将定义一个 Generator 并设置一个种子以确保可复现性。这样我们就可以通过固定 LDM 模型中的起始潜变量来调整提示词，并观察它们对生成结果的影响：

import torch

generator = torch.Generator("cuda").manual_seed(0)

现在我们准备进行第一次生成！在整个 notebook 中我们将沿用同一个运行示例，其中我们会以固定的文本提示为条件进行音频生成，并始终使用相同的随机种子。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)

听起来和我们的文本提示很吻合！音质不错，但仍带有背景噪声的瑕疵。我们可以为 pipeline 提供一个负面提示，以阻止 pipeline 生成某些特征。在本例中，我们将传入一个负面提示，以阻止模型在输出中生成低质量音频。我们将省略 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]

使用负面提示时，推理时间不变1{}^1；我们只是将 LDM 的无条件输入替换为负面输入。这意味着我们在音频质量上获得的任何提升都是免费的。

让我们听一下生成的音频：

Audio(audio, rate=16000)

整体音频质量确实有所改善——噪声瑕疵更少，音频听起来总体上更清晰。 1{}^1 请注意，在实践中，我们通常会看到从第一次生成到第二次生成时推理时间有所减少。这是由于我们第一次运行计算时发生的 CUDA“预热”。第二次生成能更好地衡量我们实际的推理时间。

优化 1：Flash Attention

PyTorch 2.0 及更高版本通过 torch.nn.functional.scaled_dot_product_attention（SDPA）函数包含了优化且内存高效的注意力运算实现。该函数会根据输入自动应用若干内置优化，并且比原版注意力实现运行得更快、更节省内存。总体而言，SDPA 函数的行为与 flash attention 类似，后者由 Dao 等人在论文 Fast and Memory-Efficient Exact Attention with IO-Awareness 中提出。

如果安装了 PyTorch 2.0 并且 torch.nn.functional.scaled_dot_product_attention 可用，这些优化将在 Diffusers 中默认启用。要使用它，只需按照官方说明安装 torch 2.0 或更高版本，然后照常使用 pipeline 即可 🚀

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 内存占用，而生成质量的变化几乎无法察觉。

我们可以通过向 .from_pretrained 传入 torch_dtype 参数，以 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 秒。 根据我们的经验，使用 diffusers 流水线搭配 float16 精度时，我们没有看到任何明显的音频质量下降， 但始终能获得显著的推理加速。因此，我们建议默认使用 float16 精度。

优化 3：Torch Compile

为了获得额外的加速，我们可以使用新的 torch.compile 功能。由于流水线中的 UNet 通常是 计算开销最大的部分，我们用 torch.compile 包装 unet，其余子模型（文本编码器 和 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 编译的 pipeline，这样我们就不会一开始就遭遇漫长的编译步骤：

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 Traceback (most recent call last) in <cell line: 5>() 3 pipe.to("cuda") 4 ----> 5 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]

23 frames /usr/local/lib/python3.10/dist-packages/torch/nn/modules/linear.py in 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:

OutOfMemoryError: CUDA out of memory. Tried to allocate 1.95 GiB. GPU 0 has a total capacty of 14.75 GiB of which 1.66 GiB is free. Process 414660 has 13.09 GiB memory in use. Of the allocated memory 10.09 GiB is allocated by PyTorch, and 1.92 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and 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 offload 运行生成与之前相同：

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]

这样一来，我们只需一次 pipeline 调用就能生成四个样本，每个时长 150 秒！使用大型 AudioLDM 2 checkpoint 会比基础 checkpoint 导致更高的总体内存占用，因为 UNet 的大小超过两倍（750M 参数对比 350M），所以这种内存节省技巧在这里尤其有益。

结论

在这篇博客文章中，我们展示了 🧨 Diffusers 开箱即用的四种优化方法，将 AudioLDM 2 的生成时间从 14 秒缩短到不到 1 秒。我们还重点介绍了如何采用内存节省技巧，例如半精度和 CPU offload，以降低长音频样本或大型 checkpoint 的峰值内存占用。

博客文章由 Sanchit Gandhi 撰写。非常感谢 Vaibhav Srivastav 和 Sayak Paul 提出的建设性意见。频谱图图片来源：Getting to Know the Mel Spectrogram。 波形图图片来源：Aalto Speech Processing。