AudioLDM 2，但更快 ⚡️
'AudioLDM 2, but faster ⚡️'.
AudioLDM 2 是由刘浩鹤等人在《AudioLDM 2:学习具有自监督预训练的整体音频生成》中提出的。AudioLDM 2 以文本提示作为输入，并预测相应的音频。它可以生成逼真的声音效果、人声和音乐。
虽然生成的音频质量很高，但使用原始实现进行推理非常慢：生成10秒音频样本需要超过30秒的时间。这是由于多阶段深度建模方法、大型检查点大小和未优化的代码等因素的综合结果。
在本博客文章中，我们展示了如何在 Hugging Face 🧨 Diffusers 库中使用 AudioLDM 2，探索了一系列代码优化，例如半精度、闪存注意力和编译，以及模型优化，例如调度器选择和负提示，以在推理时间上减少超过10倍，并且输出音频质量的降低最小。本博客文章还附带有一个更简化的 Colab 笔记本，其中包含了所有的代码，但解释较少。
阅读到最后，了解如何在仅1秒钟内生成10秒音频样本！
模型概述
受稳定扩散的启发，AudioLDM 2 是一个文本到音频的潜在扩散模型（LDM），它从文本嵌入中学习连续的音频表示。
总体生成过程概括如下：
- 给定文本输入 x\boldsymbol{x}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) E1=CLAP(x);E2=T5(x)
CLAP 文本嵌入被训练成与相应音频样本的嵌入对齐，而 Flan-T5 嵌入则更好地表示文本的语义。
- 通过各自的线性投影将这些文本嵌入投影到共享的嵌入空间：
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} P1=WCLAPE1;P2=WT5E2
在 diffusers 实现中，这些投影由 AudioLDM2ProjectionModel 定义。
- 使用 GPT2 语言模型（LM）自回归地生成一系列 NNN 个新的嵌入向量，条件是投影的 CLAP 和 Flan-T5 嵌入：
Ei=GPT2(P1,P2,E1:i−1)for i=1,…,N \boldsymbol{E}_{i} = \text{GPT2}\left(\boldsymbol{P}_{1}, \boldsymbol{P}_{2}, \boldsymbol{E}_{1:i-1}\right) \qquad \text{for } i=1,\dots,N Ei=GPT2(P1,P2,E1:i−1)for i=1,…,N
- 生成的嵌入向量 E1:N\boldsymbol{E}_{1:N}E1:N 和 Flan-T5 文本嵌入 E2\boldsymbol{E}_{2}E2 用作 LDM 中的交叉注意力条件，通过逆扩散过程去噪一个随机潜变量。LDM 在逆扩散过程中运行总共 TTT 次推理步骤：
zt=LDM(zt−1∣E1:N,E2)for t=1,…,T \boldsymbol{z}_{t} = \text{LDM}\left(\boldsymbol{z}_{t-1} | \boldsymbol{E}_{1:N}, \boldsymbol{E}_{2}\right) \qquad \text{for } t = 1, \dots, T zt=LDM(zt−1∣E1:N,E2)for t=1,…,T
初始潜变量z0\boldsymbol{z}_{0}z0从正态分布N(0,I)\mathcal{N} \left(\boldsymbol{0}, \boldsymbol{I} \right)N(0,I)中抽取。LDM的UNet在于它使用了两个交叉注意力嵌入E1:N\boldsymbol{E}_{1:N}E1:N（来自GPT2语言模型）和E2\boldsymbol{E}_{2}E2（来自Flan-T5），而不是像其他大多数LDM一样只使用一个交叉注意力条件。
- 最终去噪的潜变量zT\boldsymbol{z}_{T}zT传递给VAE解码器以恢复Mel频谱图s\boldsymbol{s}s：
s=VAEdec(zT) \boldsymbol{s} = \text{VAE}_{\text{dec}} \left(\boldsymbol{z}_{T}\right) s=VAEdec(zT)
- Mel频谱图传递给声码器以获得输出音频波形y\mathbf{y}y：
y=Vocoder(s) \boldsymbol{y} = \text{Vocoder}\left(\boldsymbol{s}\right) y=Vocoder(s)
下图演示了文本输入如何通过文本调节模型，其中两个提示嵌入用作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)输出：
正在加载流水线组件...：100%|███████████████████████████████████████████| 11/11 [00:01<00:00, 7.62it/s]流水线可以像标准的PyTorch nn模块一样移动到GPU上：
pipe.to("cuda");太好了！我们将定义一个生成器并设置一个种子以确保可重现性。这样可以通过固定LDM模型中的起始潜变量来调整我们的提示并观察它们对生成结果的影响：
import torch
generator = torch.Generator("cuda").manual_seed(0)现在我们准备进行第一次生成！我们将在整个笔记本中使用相同的示例，其中我们将在固定的文本提示上进行音频生成，并在整个过程中使用相同的种子。参数audio_length_in_s控制生成音频的长度。默认为LDM训练时的音频长度（10.24秒）：
prompt = "巴西桑巴鼓声中，海浪轻轻拍打的声音"
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 = "低质量，平均质量。"
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\)；我们只是用负面输入替换了LDM的无条件输入。这意味着我们在音频质量上获得的任何改进都是免费的。
让我们来听一下生成的音频：
Audio(audio, rate=16000)您的浏览器不支持音频元素。
整体音频质量明显有所改善 – 噪音痕迹更少，音频听起来更加清晰。1{}^11请注意，实际上，从第一代到第二代，我们通常会看到推理时间的减少。这是由于第一次运行计算时发生的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将启用这些优化。要使用它，只需按照官方说明安装torch 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编译
为了获得额外的加速，我们可以使用新的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编译文档。
优化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 Traceback (most recent call last)
<ipython-input-33-c4cae6410ff5> 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内存不足。尝试分配1.95 GiB。GPU 0的总容量为14.75 GiB，其中1.66 GiB可用。进程414660使用了13.09 GiB的内存。分配的内存中，10.09 GiB由PyTorch分配，1.92 GiB由PyTorch保留但未分配。如果保留但未分配的内存较大，请尝试设置max_split_size_mb以避免碎片化。请参阅内存管理和PYTORCH_CUDA_ALLOC_CONF的文档。除非您有带有高内存的GPU，否则上面的代码可能会返回OOM错误。虽然AudioLDM 2流水线涉及多个组件，但只有在任何时候使用的模型必须在GPU上。其余的模块可以卸载到CPU上。这种技术称为CPU卸载，可以减少内存使用量，对推理时间的惩罚非常低。
我们可以使用函数enable_model_cpu_offload()在我们的流水线上启用CPU卸载：
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的大小超过了两倍（750M参数与350M相比），因此这种内存节省技巧在这里特别有益。
结论
在这篇博文中，我们展示了四种优化方法，这些方法在🧨 Diffusers中可以直接使用，将AudioLDM 2的生成时间从14秒减少到不到1秒。我们还强调了如何使用内存节省技巧，例如半精度和CPU卸载，以减少长音频样本或大检查点大小的峰值内存使用量。
由Sanchit Gandhi撰写的博文。非常感谢Vaibhav Srivastav和Sayak Paul提供的建设性评论。声谱图图像来源：了解Mel声谱图。波形图像来源：阿尔托语音处理。
