AudioLDM 2 优化指南：使用 Hugging Face Diffusers 减少推理时间

Hugging Face 发布了 AudioLDM 2 的优化指南，AudioLDM 2 是一种能够生成逼真的音效、人声和音乐的文本到音频潜在扩散模型。通过在 diffusers 库中应用代码和模型优化的组合，10 秒音频样本的推理时间可以从原始实现的超过 30 秒降低到不到 1 秒。

AudioLDM 2 模型架构

AudioLDM 2 是一种文本到音频潜在扩散模型（LDM），通过从文本嵌入学习连续表示来生成音频。生成过程遵循多阶段流水线：

文本编码：模型使用两个编码器——CLAP 的文本分支（用于音频对齐嵌入）和 Flan-T5 的文本编码器（用于语义表示）——来计算文本嵌入。
投射：这些嵌入通过 AudioLDM2ProjectionModel 进行线性投射进入共享嵌入空间。
自回归生成：GPT2 语言模型在投射后的 CLAP 和 Flan-T5 嵌入条件下生成 N 个嵌入向量的序列。
潜在扩散：基于 UNet 的 LDM 在 $T$ 步推理过程中对随机潜在变量进行去噪。与大多数 LDM 不同，AudioLDM 2 的 UNet 使用两组交叉注意力嵌入：来自 GPT2 的和来自 Flan-T5 的。
解码：最终去噪的潜在变量通过 VAE 解码器恢复为梅尔频谱图，然后由声码器转换为音频波形。

可用的模型检查点

检查点	任务	模型大小	训练数据（小时）
`cvssp/audioldm2`	文本到音频	1.1B	1150k
`cvssp/audioldm2-music`	文本到音乐	1.1B	665k
`cvssp/audioldm2-large`	文本到音频	1.5B	1150k

推理速度优化

为了解决原始实现的推理速度慢的问题，Hugging Face 在 diffusers 库中确定了四种主要的优化技术：

1. Flash Attention (SDPA)

使用 PyTorch 2.0 或更高版本时，diffusers 库会自动启用 torch.nn.functional.scaled_dot_product_attention（SDPA）。这提供了一种类似 Flash Attention 的内存高效注意力操作，在不改变输出质量的情况下减少计算时间。

2. 半精度 (float16)

将模型权重和计算从 float32 转换为 float16（半精度）可显著降低 GPU 内存使用并提升推理速度，对音频质量的影响不可察觉。这是通过在 .from_pretrained 调用中传递 torch_dtype=torch.float16 实现的。

3. Torch 编译

将管道中计算最昂贵的 UNet 部分用 torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) 包装起来，可以获得显著的加速。虽然第一次推理运行由于编译开销会很慢（可能长达 2 分钟），但所有后续生成都会显著加快。

4. 高效调度器

将默认的 DDIMScheduler（通常需要 200 步）替换为性能更高的调度器（如 DPMSolverMultistepScheduler），使模型仅在 20-25 步推理内即可达到相似质量。此优化结合其他优化可将 10 秒样本的生成时间降低到不到 1 秒。

内存管理和长音频

生成长音频样本（例如 150 秒）或使用更大的检查点（如 audioldm2-large）会增加潜在变量的宽度，这可能导致 CUDA 内存不足（OOM）错误，因为交叉注意力内存随序列长度的平方增长。

为了缓解此问题，Hugging Face 建议通过 pipe.enable_model_cpu_offload() 进行 CPU 卸载。此技术仅将当前活动的模型组件保留在 GPU 上，其余部分卸载到 CPU，从而在 GPU 内存有限的情况下实现长音频生成和使用更大模型，对推理时间的影响极小。