探索 SDXL 的簡單最佳化方法

Stable Diffusion XL (SDXL) 是 Stability AI 推出的最新潛在擴散模型，用於生成高質量、超逼真的影像。它克服了以往 Stable Diffusion 模型在處理手部、文字以及空間構圖正確性等方面的挑戰。此外，SDXL 的上下文感知能力更強，只需較少的提示詞就能生成更美觀的影像。

然而，所有這些改進都是以模型體積顯著增大為代價的。到底大了多少？SDXL 基礎模型有 35 億個引數（尤其是 UNet 部分），大約是之前 Stable Diffusion 模型的三倍大。

為了探索如何最佳化 SDXL 的推理速度和記憶體使用，我們在 A100 GPU (40 GB) 上進行了一些測試。每次推理執行，我們生成 4 張影像並重復 3 次。在計算推理延遲時，我們只考慮 3 次迭代中的最後一次。

所以，如果你直接開箱即用地執行全精度的 SDXL，並使用預設的注意力機制，它將消耗 28GB 記憶體，耗時 72.2 秒！

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0").to("cuda")
pipe.unet.set_default_attn_processor()

這很不實用，並且可能會拖慢你的速度，因為你通常需要生成不止 4 張影像。而且，如果你沒有更強大的 GPU，就會遇到那個令人沮喪的記憶體不足錯誤資訊。那麼，我們該如何最佳化 SDXL 以提高推理速度並減少其記憶體使用呢？

在 🤗 Diffusers 中，我們有許多最佳化技巧和技術來幫助你執行像 SDXL 這樣記憶體密集型的模型，接下來我們將向你展示如何做！我們將重點關注兩個方面：推理速度和記憶體。

推理速度

擴散是一個隨機過程，所以不能保證你一次就能得到滿意的影像。通常情況下，你需要多次執行推理並進行迭代，這就是為什麼最佳化速度至關重要。本節將重點介紹如何使用更低精度的權重，並結合記憶體高效的注意力和 PyTorch 2.0 中的 torch.compile 來提升速度並減少推理時間。

更低精度

模型權重以特定的精度儲存，該精度表示為一種浮點資料型別。標準的浮點資料型別是 float32 (fp32)，它可以精確表示大範圍的浮點數。對於推理而言，通常不需要那麼高的精度，所以你可以使用 float16 (fp16)，它能表示的浮點數範圍較窄。這意味著 fp16 佔用的儲存空間只有 fp32 的一半，而且由於計算更簡單，速度是 fp32 的兩倍。此外，現代 GPU 卡有專門最佳化的硬體來執行 fp16 計算，使其速度更快。

在 🤗 Diffusers 中，你可以透過指定 torch.dtype 引數在載入模型時轉換權重，從而在推理時使用 fp16。

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")
pipe.unet.set_default_attn_processor()

與完全未經最佳化的 SDXL 流水線相比，使用 fp16 僅需 21.7GB 記憶體，耗時 14.8 秒。你幾乎將推理速度提升了整整一分鐘！

記憶體高效注意力

Transformer 模組中使用的注意力塊可能是一個巨大的瓶頸，因為隨著輸入序列變長，記憶體會呈二次方增長。這會迅速佔用大量記憶體，並導致你收到記憶體不足的錯誤資訊。😬

記憶體高效的注意力演算法旨在減輕計算注意力的記憶體負擔，無論是透過利用稀疏性還是分塊技術。這些最佳化演算法過去主要以需要單獨安裝的第三方庫的形式提供。但從 PyTorch 2.0 開始，情況不再如此。PyTorch 2 引入了縮放點積注意力 (SDPA)，它提供了Flash Attention、記憶體高效注意力 (xFormers) 以及一個 C++ 實現的 PyTorch 融合實現。SDPA 可能是加速推理最簡單的方法：如果你正在使用 PyTorch ≥ 2.0 和 🤗 Diffusers，它預設是自動啟用的！

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")

與完全未經最佳化的 SDXL 流水線相比，使用 fp16 和 SDPA 所需的記憶體量相同，但推理時間縮短至 11.4 秒。讓我們以此為新的基準，來比較其他最佳化方法。

torch.compile

PyTorch 2.0 還引入了 torch.compile API，用於將你的 PyTorch 程式碼即時 (JIT) 編譯成更最佳化的推理核心。與其他編譯器解決方案不同，torch.compile 對現有程式碼的改動極小，只需用該函式包裝你的模型即可。

透過 mode 引數，你可以在編譯時針對記憶體開銷或推理速度進行最佳化，這為你提供了更大的靈活性。

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

與之前的基準 (fp16 + SDPA) 相比，用 torch.compile 包裝 UNet 後，推理時間縮短至 10.2 秒。

模型記憶體佔用

如今的模型越來越大，要將它們裝入記憶體成為一項挑戰。本節重點介紹如何減少這些龐大模型的記憶體佔用，以便你可以在消費級 GPU 上執行它們。這些技術包括 CPU 解除安裝、分步解碼潛在表示為影像而非一次性完成，以及使用蒸餾版的自動編碼器。

模型 CPU 解除安裝

模型解除安裝透過將 UNet 載入到 GPU 記憶體中，而將擴散模型的其他元件（文字編碼器、VAE）載入到 CPU 上來節省記憶體。這樣，UNet 可以在 GPU 上進行多次迭代，直到不再需要它為止。

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
)
pipe.enable_model_cpu_offload()

與基準相比，現在需要 20.2GB 記憶體，為你節省了 1.5GB 記憶體。

順序 CPU 解除安裝

另一種可以為你節省更多記憶體但會降低推理速度的解除安裝方式是順序 CPU 解除安裝。它不是解除安裝整個模型（如 UNet），而是將儲存在不同 UNet 子模組中的模型權重解除安裝到 CPU，並且僅在前向傳播之前才載入到 GPU。本質上，你每次只加載模型的一部分，從而可以節省更多記憶體。唯一的缺點是它明顯更慢，因為你需要多次載入和解除安裝子模組。

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
)
pipe.enable_sequential_cpu_offload()

與基準相比，這需要 19.9GB 記憶體，但推理時間增加到 67 秒。

切片

在 SDXL 中，變分自編碼器 (VAE) 將精煉後的潛在表示（由 UNet 預測）解碼為逼真的影像。這一步的記憶體需求與預測的影像數量（批次大小）成正比。根據影像解析度和可用的 GPU視訊記憶體，它可能會非常佔用記憶體。

這就是“切片”技術發揮作用的地方。待解碼的輸入張量被分割成多個切片，解碼計算分幾步完成。這節省了記憶體並允許使用更大的批次大小。

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")
pipe.enable_vae_slicing()

透過切片計算，我們將記憶體減少到 15.4GB。如果再加入順序 CPU 解除安裝，記憶體可以進一步減少到 11.45GB，這讓你可以在每個提示詞下生成 4 張 (1024x1024) 影像。然而，使用順序解除安裝時，推理延遲也會增加。

快取計算

任何文字條件的影像生成模型通常都使用文字編碼器來從輸入提示中計算嵌入。SDXL 使用了兩個文字編碼器！這對推理延遲有相當大的影響。然而，由於這些嵌入在整個反向擴散過程中保持不變，我們可以預先計算它們並在後續過程中重複使用。這樣，在計算完文字嵌入後，我們就可以將文字編碼器從記憶體中移除。

首先，載入文字編碼器及其對應的分詞器，並從輸入提示中計算嵌入

tokenizers = [tokenizer, tokenizer_2]
text_encoders = [text_encoder, text_encoder_2]

(
    prompt_embeds,
    negative_prompt_embeds,
    pooled_prompt_embeds,
    negative_pooled_prompt_embeds
) = encode_prompt(tokenizers, text_encoders, prompt)

接下來，清空 GPU 記憶體以移除文字編碼器

del text_encoder, text_encoder_2, tokenizer, tokenizer_2
flush()

現在，這些嵌入可以直接用於 SDXL 流水線了

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    text_encoder=None,
    text_encoder_2=None,
    tokenizer=None,
    tokenizer_2=None,
    torch_dtype=torch.float16,
).to("cuda")

call_args = dict(
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds,
        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
        num_images_per_prompt=num_images_per_prompt,
        num_inference_steps=num_inference_steps,
)
image = pipe(**call_args).images[0]

結合 SDPA 和 fp16，我們可以將記憶體減少到 21.9GB。上面討論的其他用於最佳化記憶體的技術也可以與快取計算一起使用。

微型自動編碼器

如前所述，VAE 將潛在表示解碼為影像。很自然地，這一步直接受限於 VAE 的大小。所以，讓我們用一個更小的自動編碼器吧！由 madebyollin 製作的微型自動編碼器，在 Hub 上可用，大小僅為 10MB，它是從 SDXL 使用的原始 VAE 蒸餾而來的。

from diffusers import AutoencoderTiny

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
)
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.float16)
pipe.to("cuda")

透過這種設定，我們將記憶體需求減少到 15.6GB，同時還降低了推理延遲。

結論

總結一下我們最佳化所帶來的節省：

我們希望這些最佳化能讓你輕鬆執行自己喜歡的流水線。快來試試這些技術，並與我們分享你的影像吧！🤗

---

致謝：感謝 Pedro Cuenca 對草稿提出的有益審閱。