AudioLDM 2 優化指南：使用 Hugging Face Diffusers 縮短推論時間

Hugging Face 發布了 AudioLDM 2 的優化指南，這是一個文本轉音訊的潛在擴散模型 (latent diffusion model)，能夠生成逼真的音效、人類語音和音樂。透過在 diffusers 函式庫中結合程式碼與模型優化，將 10 秒音訊樣本的推論時間從原始實作的 30 秒以上縮短至不到 1 秒。

AudioLDM 2 模型架構

AudioLDM 2 是一個文本轉音訊的潛在擴散模型 (LDM)，透過從文本嵌入 (text embeddings) 中學習連續表示來生成音訊。生成過程遵循多階段流水線：

文本編碼 (Text Encoding)：模型使用兩個編碼器——CLAP 的文本分支（用於音訊對齊的嵌入）和 Flan-T5 的文本編碼器（用於語義表示）來計算文本嵌入。
投影 (Projection)：這些嵌入透過 AudioLDM2ProjectionModel 進行線性投影，進入共享嵌入空間。
自回歸生成 (Auto-regressive Generation)：一個 GPT2 語言模型根據投影後的 CLAP 和 Flan-T5 嵌入，生成一個包含 $N$ 個嵌入向量的序列。
潛在擴散 (Latent Diffusion)：一個基於 UNet 的 LDM 在 $T$ 個推論步驟中對隨機潛在變量進行去噪。與大多數 LDM 不同，AudioLDM 2 的 UNet 使用了兩組交叉注意力 (cross-attention) 嵌入：分別來自 GPT2 和 Flan-T5。
解碼 (Decoding)：最終去噪後的潛在變量透過 VAE 解碼器恢復為 Mel 頻譜圖，然後由聲碼器 (vocoder) 轉換為音訊波形。

可用的模型權重 (Checkpoints)

Checkpoint	任務	模型大小	訓練數據 (小時)
`cvssp/audioldm2`	Text-to-audio	1.1B	1150k
`cvssp/audioldm2-music`	Text-to-music	1.1B	665k
`cvssp/audioldm2-large`	Text-to-audio	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 Compile

使用 torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) 來封裝 UNet（流水線中計算量最大的部分）可以提供實質性的加速。雖然第一次推論運行會因為編譯開銷而較慢（可能長達 2 分鐘），但隨後的所有生成都會快得多。

4. 高效調度器 (Efficient Schedulers)

使用更高效的調度器（如 DPMSolverMultistepScheduler）取代預設的 DDIMScheduler（通常需要 200 個步驟），可以讓模型僅需 20-25 個推論步驟即可達到相似的品質。這項優化結合其他技術，可以將 10 秒樣本的生成時間縮短到 1 秒以內。

記憶體管理與長音訊

生成長音訊樣本（例如 150 秒）或使用較大的權重（如 audioldm2-large）會增加潛在變量的寬度，這可能導致 CUDA 顯存溢出 (OOM) 錯誤，因為交叉注意力的記憶體隨序列長度呈平方級增長。

為了緩解這一問題，Hugging Face 建議透過 pipe.enable_model_cpu_offload() 進行 CPU 卸載 (CPU offloading)。這項技術僅將當前活動的模型組件保留在 GPU 上，並將其餘部分卸載到 CPU，從而允許在 RAM 有限的 GPU 上生成長音訊並使用更大的模型，且對推論時間的影響極小。