AudioLDM 2,但更快 ⚡️

釋出於2023年8月30日
在 GitHub 上更新
Open In Colab

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),它從文字嵌入中學習連續的音訊表示。

整體生成過程總結如下:

  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:i1)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 在反向擴散過程中執行總共 TT 推理步驟。

zt=LDM(zt1E~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 上找到。

模型權重 任務 模型大小 訓練資料 / 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]

使用負面提示時,推理時間沒有改變 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 函式提供了與 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 回溯(最近一次呼叫):在 <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 幀 /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 SrivastavSayak Paul 的建設性意見。頻譜圖圖片來源:瞭解 Mel 頻譜圖。波形圖片來源:阿爾託大學語音處理

社群

註冊登入 評論

© . This site is unofficial and not affiliated with Hugging Face, Inc.