AudioLDM 2，加速?！
AudioLDM 2 由劉濠赫等人在 AudioLDM 2: Learning Holistic Audio Generation with Self-supervised Pretraining 一文中提出。AudioLDM 2 接受文本提示作為輸入并輸出對(duì)應(yīng)的音頻，其可用于生成逼真的聲效、人類語(yǔ)音以及音樂(lè)。
雖然生成的音頻質(zhì)量很高，但基于其原始實(shí)現(xiàn)進(jìn)行推理的速度非常慢: 生成一個(gè) 10 秒的音頻需要 30 秒以上的時(shí)間。慢的原因是多重的，包括其使用了多階段建模、checkpoint 較大以及代碼尚未優(yōu)化等。
本文將展示如何在 Hugging Face ?? Diffusers 庫(kù)中使用 AudioLDM 2，并在此基礎(chǔ)上探索一系列代碼優(yōu)化 (如半精度、Flash 注意力、圖編譯) 以及模型級(jí)優(yōu)化 (如選擇合適的調(diào)度器及反向提示)。最終我們將推理時(shí)間降低了?10 倍?多，且對(duì)輸出音頻質(zhì)量的影響最低。本文還附有一個(gè)更精簡(jiǎn)的 Colab notebook，這里面包含所有代碼但精簡(jiǎn)了很多文字部分。
最終，我們可以在短短 1 秒內(nèi)生成一個(gè) 10 秒的音頻！
模型概述
受 Stable Diffusion 的啟發(fā)，AudioLDM 2 是一種文生音頻的 _ 隱擴(kuò)散模型 (latent diffusion model，LDM)_，其可以將文本嵌入映射成連續(xù)的音頻表征。
大體的生成流程總結(jié)如下:
給定輸入文本?
，使用兩個(gè)文本編碼器模型來(lái)計(jì)算文本嵌入: CLAP 的文本分支，以及 Flan-T5 的文本編碼器。
CLAP 文本嵌入經(jīng)過(guò)訓(xùn)練，可以與對(duì)應(yīng)的音頻嵌入對(duì)齊，而 Flan-T5 嵌入可以更好地表征文本的語(yǔ)義。
這些文本嵌入通過(guò)各自的線性層投影到同一個(gè)嵌入空間:
在?
diffusers?實(shí)現(xiàn)中，這些投影由 AudioLDM2ProjectionModel 定義。使用 GPT2 語(yǔ)言模型 (LM) 基于 CLAP 和 Flan-T5 嵌入自回歸地生成一個(gè)含有?
?個(gè)嵌入向量的新序列:
以生成的嵌入向量?
?和 Flan-T5 文本嵌入?
?為條件，通過(guò) LDM 的反向擴(kuò)散過(guò)程對(duì)隨機(jī)隱變量進(jìn)行?去噪?。LDM 在反向擴(kuò)散過(guò)程中運(yùn)行??
?個(gè)步推理:
其中初始隱變量?
?是從正態(tài)分布?
?中采樣而得。LDM 的 UNet 的獨(dú)特之處在于它需要?兩組?交叉注意力嵌入，來(lái)自 GPT2 語(yǔ)言模型的?
?和來(lái)自 Flan-T5 的 ?
，而其他大多數(shù) LDM 只有一個(gè)交叉注意力條件。
把最終去噪后的隱變量?
?傳給 VAE 解碼器以恢復(fù)梅爾譜圖?
:
梅爾譜圖被傳給聲碼器 (vocoder) 以獲得輸出音頻波形?
:
下圖展示了文本輸入是如何作為條件傳遞給模型的，可以看到在 LDM 中兩個(gè)提示嵌入均被用作了交叉注意力的條件:
有關(guān)如何訓(xùn)練 AudioLDM 2 模型的完整的詳細(xì)信息，讀者可以參閱 AudioLDM 2 論文。
Hugging Face ?? Diffusers 提供了一個(gè)端到端的推理流水線類?AudioLDM2Pipeline?以將該模型的多階段生成過(guò)程包裝到單個(gè)可調(diào)用對(duì)象中，這樣用戶只需幾行代碼即可完成從文本生成音頻的過(guò)程。
AudioLDM 2 有三個(gè)變體。其中兩個(gè) checkpoint 適用于通用的文本到音頻生成任務(wù)，第三個(gè) checkpoint 專門(mén)針對(duì)文本到音樂(lè)生成。三個(gè)官方 checkpoint 的詳細(xì)信息請(qǐng)參見(jiàn)下表，這些 checkpoint 都可以在 Hugging Face Hub 上找到:
至此，我們已經(jīng)全面概述了 AudioLDM 2 生成的工作原理，接下來(lái)讓我們將這一理論付諸實(shí)踐！
加載流水線
我們以基礎(chǔ)版模型 cvssp/audioldm2 為例，首先使用?.from_pretrained?方法來(lái)加載整個(gè)管道，該方法會(huì)實(shí)例化管道并加載預(yù)訓(xùn)練權(quán)重:
輸出:
與 PyTorch 一樣，使用?to?方法將流水線移至 GPU:
現(xiàn)在，我們來(lái)定義一個(gè)隨機(jī)數(shù)生成器并固定一個(gè)種子，我們可以通過(guò)這種方式來(lái)固定 LDM 模型中的起始隱變量從而保證結(jié)果的可復(fù)現(xiàn)性，并可以觀察不同提示對(duì)生成過(guò)程和結(jié)果的影響:
現(xiàn)在，我們準(zhǔn)備好開(kāi)始第一次生成了！本文中的所有實(shí)驗(yàn)都會(huì)使用固定的文本提示以及相同的隨機(jī)種子來(lái)生成音頻，并比較不同方案的延時(shí)和效果。audio_length_in_s?參數(shù)主要控制所生成音頻的長(zhǎng)度，這里我們將其設(shè)置為默認(rèn)值，即 LDM 訓(xùn)練時(shí)的音頻長(zhǎng)度: 10.24 秒:
輸出:
酷！我們花了大約 13 秒最終生成出了音頻。我們來(lái)聽(tīng)一下:
請(qǐng)?閱讀原文?聽(tīng)取音頻附件
聽(tīng)起來(lái)跟我們的文字提示很吻合！質(zhì)量很好，但是有一些背景噪音。我們可以為流水線提供?反向提示 (negative prompt)，以防止其生成的音頻中含有某些不想要特征。這里，我們給模型一個(gè)反向提示，以防止模型生成低質(zhì)量的音頻。我們不設(shè)?audio_length_in_s?參數(shù)以使用其默認(rèn)值:
輸出:
使用反向提示??時(shí)，推理時(shí)間不變; 我們只需將 LDM 的無(wú)條件輸入替換為反向提示即可。這意味著我們?cè)谝纛l質(zhì)量方面獲得的任何收益都是免費(fèi)的。
我們聽(tīng)一下生成的音頻:
請(qǐng)?閱讀原文?聽(tīng)取音頻附件
顯然，整體音頻質(zhì)量有所改善 - 噪聲更少，并且音頻整體聽(tīng)起來(lái)更清晰。
?請(qǐng)注意，在實(shí)踐中，我們通常會(huì)看到第二次生成比第一次生成所需的推理時(shí)間有所減少。這是由于我們第一次運(yùn)行計(jì)算時(shí) CUDA 被“預(yù)熱”了。因此一般進(jìn)行基準(zhǔn)測(cè)試時(shí)我們會(huì)選擇第二次推理的時(shí)間作為結(jié)果。
優(yōu)化 1: Flash 注意力
PyTorch 2.0 及更高版本包含了一個(gè)優(yōu)化過(guò)的內(nèi)存高效的注意力機(jī)制的實(shí)現(xiàn)，用戶可通過(guò)?torch.nn.function.scaled_dot_product_attention?(SDPA) 函數(shù)來(lái)調(diào)用該優(yōu)化。該函數(shù)會(huì)根據(jù)輸入自動(dòng)使能多個(gè)內(nèi)置優(yōu)化，因此比普通的注意力實(shí)現(xiàn)運(yùn)行得更快、更節(jié)省內(nèi)存。總體而言，SDPA 函數(shù)的優(yōu)化與 Dao 等人在論文 Fast and Memory-Efficient Exact Attention with IO-Awareness 中所提出的?flash 注意力?類似。
如果安裝了 PyTorch 2.0 且?torch.nn.function.scaled_dot_product_attention?可用，Diffusers 將默認(rèn)啟用該函數(shù)。因此，僅需按照 官方說(shuō)明 安裝 torch 2.0 或更高版本，不需對(duì)流水線??作任何改動(dòng)，即能享受提速。
輸出:
有關(guān)在?diffusers?中使用 SDPA 的更多詳細(xì)信息，請(qǐng)參閱相應(yīng)的 文檔。
優(yōu)化 2: 半精度
默認(rèn)情況下，?AudioLDM2Pipeline?以 float32 (全) 精度方式加載模型權(quán)重。所有模型計(jì)算也以 float32 精度執(zhí)行。對(duì)推理而言，我們可以安全地將模型權(quán)重和計(jì)算轉(zhuǎn)換為 float16 (半) 精度，這能改善推理時(shí)間和 GPU 內(nèi)存，同時(shí)對(duì)生成質(zhì)量的影響微乎其微。
我們可以通過(guò)將?from_pretrained?的?torch_dtype?參數(shù)設(shè)為?torch.float16?來(lái)加載半精度權(quán)重:
我們運(yùn)行一下 float16 精度的生成，并聽(tīng)一下輸出:
輸出:
請(qǐng)?閱讀原文?聽(tīng)取音頻附件
音頻質(zhì)量與全精度生成基本沒(méi)有變化，推理加速了大約 2 秒。根據(jù)我們的經(jīng)驗(yàn)，使用具有 float16 精度的?diffusers?流水線，我們可以獲得顯著的推理加速而無(wú)明顯的音頻質(zhì)量下降。因此，我們建議默認(rèn)使用 float16 精度。
優(yōu)化 3: Torch Compile
為了獲得額外的加速，我們還可以使用新的?torch.compile?功能。由于在流水線中 UNet 通常計(jì)算成本最高，因此我們用?torch.compile?編譯一下 UNet，其余子模型 (文本編碼器和 VAE) 保持不變:
用?torch.compile?包裝 UNet 后，由于編譯 UNet 的開(kāi)銷，我們運(yùn)行第一步推理時(shí)通常會(huì)很慢。所以，我們先運(yùn)行一步流水線預(yù)熱，這樣后面真正運(yùn)行的時(shí)候就快了。請(qǐng)注意，第一次推理的編譯時(shí)間可能長(zhǎng)達(dá) 2 分鐘，請(qǐng)耐心等待！
輸出:
很棒！現(xiàn)在 UNet 已編譯完畢，現(xiàn)在可以以更快的速度運(yùn)行完整的擴(kuò)散過(guò)程了:
輸出:
只需 4 秒即可生成！在實(shí)踐中，你只需編譯 UNet 一次，然后就可以為后面的所有生成贏得一個(gè)更快的推理。這意味著編譯模型所花費(fèi)的時(shí)間可以由后續(xù)推理時(shí)間的收益所均攤。有關(guān)?torch.compile?的更多信息及選項(xiàng)，請(qǐng)參閱 torch compile 文檔。
優(yōu)化 4: 調(diào)度器
還有一個(gè)選項(xiàng)是減少推理步數(shù)。選擇更高效的調(diào)度器可以幫助減少步數(shù)，而不會(huì)犧牲輸出音頻質(zhì)量。你可以調(diào)用?schedulers.compatibles?屬性來(lái)查看哪些調(diào)度器與?AudioLDM2Pipeline?兼容:
輸出:
好！現(xiàn)在我們有一長(zhǎng)串的調(diào)度器備選??。默認(rèn)情況下，AudioLDM 2 使用?DDIMScheduler，其需要 200 個(gè)推理步才能生成高質(zhì)量的音頻。但是，性能更高的調(diào)度程序，例如?DPMSolverMultistepScheduler，只需?20-25 個(gè)推理步?即可獲得類似的結(jié)果。
讓我們看看如何將 AudioLDM 2 調(diào)度器從?DDIM?切換到?DPM Multistep?。我們需要使用?ConfigMixin.from_config()?方法以用原始?DDIMScheduler?的配置來(lái)加載?DPMSolverMultistepScheduler:
讓我們將推理步數(shù)設(shè)為 20，并使用新的調(diào)度器重新生成。由于 LDM 隱變量的形狀未更改，因此我們不必重編譯:
輸出:
這次只用了不到?1 秒?就生成了音頻！我們聽(tīng)下它的生成:
Audio(audio,?rate=16000)
請(qǐng)?閱讀原文?聽(tīng)取音頻附件
生成質(zhì)量與原來(lái)的基本相同，但只花了原來(lái)時(shí)間的一小部分！?? Diffusers 流水線是“可組合”的，這個(gè)設(shè)計(jì)允許你輕松地替換調(diào)度器或其他組件以獲得更高性能。
內(nèi)存消耗如何？
我們想要生成的音頻的長(zhǎng)度決定了 LDM 中待去噪的隱變量的?寬度?。由于 UNet 中交叉注意力層的內(nèi)存隨序列長(zhǎng)度 (寬度) 的平方而變化，因此生成非常長(zhǎng)的音頻可能會(huì)導(dǎo)致內(nèi)存不足錯(cuò)誤。我們還可以通過(guò) batch size 來(lái)控制生成的樣本數(shù)，進(jìn)而控制內(nèi)存使用。
如前所述，以 float16 半精度加載模型可以節(jié)省大量?jī)?nèi)存。使用 PyTorch 2.0 SDPA 也可以改善內(nèi)存占用，但這部分改善對(duì)超長(zhǎng)序列長(zhǎng)度來(lái)講可能不夠。
我們來(lái)試著生成一個(gè) 2.5 分鐘 (150 秒) 的音頻。我們通過(guò)設(shè)置?num_waveforms_per_prompt?=4?來(lái)生成 4 個(gè)候選音頻。一旦?num_waveforms_per_prompt?>1?，在生成的音頻和文本提示之間會(huì)有一個(gè)自動(dòng)評(píng)分機(jī)制: 將音頻和文本提示嵌入到 CLAP 音頻文本嵌入空間中，然后根據(jù)它們的余弦相似度得分進(jìn)行排名。生成的音頻中第?0?個(gè)音頻就是分?jǐn)?shù)“最高”的音頻。
由于我們更改了 UNet 中隱變量的寬度，因此我們必須使用新的隱變量形狀再執(zhí)行一次 torch 編譯。為了節(jié)省時(shí)間，我們就不編譯了，直接重新加載管道:
輸出:
除非你的 GPU 顯存很大，否則上面的代碼可能會(huì)返回 OOM 錯(cuò)誤。雖然 AudioLDM 2 流水線涉及多個(gè)組件，但任何時(shí)候只有當(dāng)前正在使用的模型必須在 GPU 上。其余模塊均可以卸載到 CPU。該技術(shù)稱為“CPU 卸載”，可大大減少顯存使用，且對(duì)推理時(shí)間的影響很小。
我們可以使用函數(shù) enable_model_cpu_offload() 在流水線上啟用 CPU 卸載:
調(diào)用 API 生成音頻的方式與以前相同:
輸出:
這樣，我們就可以生成 4 個(gè)各為 150 秒的樣本，所有這些都在一次流水線調(diào)用中完成！大版的 AudioLDM 2 checkpoint 比基礎(chǔ)版的 checkpoint 總內(nèi)存使用量更高，因?yàn)?UNet 的大小相差兩倍多 (750M 參數(shù)與 350M 參數(shù)相比)，因此這種節(jié)省內(nèi)存的技巧對(duì)大版的 checkpoint 特別有用。
總結(jié)
在本文中，我們展示了 ?? Diffusers 開(kāi)箱即用的四種優(yōu)化方法，并將 AudioLDM 2 的生成時(shí)間從 14 秒縮短到不到 1 秒。我們還重點(diǎn)介紹了如何使用內(nèi)存節(jié)省技巧 (例如半精度和 CPU 卸載) 來(lái)減少長(zhǎng)音頻樣本或大 checkpoint 場(chǎng)景下的峰值顯存使用量。
本文作者 Sanchit Gandhi 非常感謝 Vaibhav Srivastav 和 Sayak Paul 的建設(shè)性意見(jiàn)。頻譜圖圖像來(lái)自于 Getting to Know the Mel Spectrogram 一文，波形圖來(lái)自于 Aalto Speech Processing 一文。
?? 寶子們可以戳?閱讀原文?查看文中所有的外部鏈接喲！
英文原文:https://hf.co/blog/audioldm2
原文作者: Sanchit Gandhi
譯者: Matrix Yao (姚偉峰)，英特爾深度學(xué)習(xí)工程師，工作方向?yàn)?transformer-family 模型在各模態(tài)數(shù)據(jù)上的應(yīng)用及大規(guī)模模型的訓(xùn)練推理。
審校/排版: zhongdongy (阿東)
