使用 ONNX Runtime 和 Olive 加速 SD Turbo 和 SDXL Turbo 推理
介紹
SD Turbo 和 SDXL Turbo 是兩個快速的生成式文字到影像模型，能夠在短短一步內生成可用的影像，這比以前的 Stable Diffusion 模型通常需要的 30 多步是一個顯著的改進。SD Turbo 是 Stable Diffusion 2.1 的蒸餾版本，SDXL Turbo 是 SDXL 1.0 的蒸餾版本。我們之前展示過如何使用 ONNX Runtime 加速 Stable Diffusion 推理。ONNX Runtime 不僅在使用 SD Turbo 和 SDXL Turbo 時提供效能優勢，而且還使模型可以在 Python 之外的語言中使用，例如 C# 和 Java。
效能提升
在這篇文章中，我們將介紹 ONNX Runtime CUDA 和 TensorRT 執行提供程式中的最佳化，這些最佳化顯著加速了 SD Turbo 和 SDXL Turbo 在 NVIDIA GPU 上的推理。
在所有測試的（批次大小，步數）組合中，ONNX Runtime 都優於 PyTorch，SDXL Turbo 模型的吞吐量提升高達 229%，SD Turbo 模型吞吐量提升高達 120%。ONNX Runtime CUDA 在動態形狀方面表現出色，但在靜態形狀方面也比 PyTorch 有顯著改進。
如何執行 SD Turbo 和 SDXL Turbo
為了使用 ONNX Runtime CUDA 執行提供程式加速推理，請訪問我們在 Hugging Face 上最佳化的 SD Turbo 和 SDXL Turbo 版本。
這些模型由 Olive 生成，Olive 是一款易於使用的、硬體感知型的模型最佳化工具。請注意，為了獲得最佳效能，必須透過命令列啟用 fp16 VAE，如共享的最佳化版本所示。有關如何使用 Hugging Face 上託管的 ONNX 檔案執行 SD 和 SDXL 管道的說明，請參閱 SD Turbo 使用示例 和 SDXL Turbo 使用示例。
要使用 ONNX Runtime TensorRT 執行提供程式加速推理，請按照此處的說明進行操作。
以下是使用 SDXL Turbo 模型並由文字提示引導生成影像的示例
python3 demo_txt2img_xl.py \
--version xl-turbo \
"little cute gremlin wearing a jacket, cinematic, vivid colors, intricate masterpiece, golden ratio, highly detailed"
圖 1. 使用 SDXL Turbo 透過文字提示生成的穿著夾克的小可愛小妖精影像。
請注意，示例影像是在 4 個步驟中生成的，這表明 SD Turbo 和 SDXL Turbo 能夠以比以前的 Stable Diffusion 模型更少的步驟生成可用的影像。
要以使用者友好的方式試用 Stable Diffusion 模型，請參閱我們的 Automatic1111 SD WebUI 的 ONNX Runtime 擴充套件。此擴充套件可以在 NVIDIA GPU 上最佳化執行 Stable Diffusion UNet 模型，並使用 ONNX Runtime CUDA 執行提供程式對經過 Olive 最佳化的模型進行推理。目前，此擴充套件僅針對 Stable Diffusion 1.5 進行了最佳化。SD Turbo 和 SDXL Turbo 模型也可以使用，但效能最佳化仍在進行中。
C# 和 Java 中 Stable Diffusion 的應用
利用 ONNX Runtime 的跨平臺、效能和可用性優勢，社群成員也貢獻了自己的 Stable Diffusion 示例和 UI 工具，並使用 ONNX Runtime 進行推理。
這些社群貢獻包括 OnnxStack，這是一個 .NET 庫，它基於我們之前的 C# 教程構建，為使用者在使用 C# 和 ONNX Runtime 進行推理時，提供了各種 Stable Diffusion 模型的多種功能。
此外，Oracle 還發布了一個使用 Java 的 Stable Diffusion 示例，它在 ONNX Runtime 之上執行推理。這個專案也基於我們的 C# 教程。
基準測試結果
我們使用 A100-SXM4-80GB 的 Standard_ND96amsr_A100_v4 VM 和配備 RTX-4090 GPU 的 聯想桌上型電腦（WSL Ubuntu 20.04）對 SD Turbo 和 SDXL Turbo 模型進行了基準測試，以使用 LCM 排程器和 fp16 模型生成 512x512 解析度的影像。結果根據以下規格測量：
- onnxruntime-gpu==1.17.0（從原始碼構建）
- torch==2.1.0a0+32f93b1
- tensorrt==8.6.1
- transformers==4.36.0
- diffusers==0.24.0
- onnx==1.14.1
- onnx-graphsurgeon==0.3.27
- polygraphy==0.49.0
要復現這些結果，我們建議使用“使用示例”部分中連結的說明。
由於 SDXL Turbo 的原始 VAE 無法以 fp16 精度執行，我們在測試 SDXL Turbo 時使用了 sdxl-vae-fp16-fix。它的輸出與原始 VAE 略有差異，但解碼後的影像對於大多數用途來說足夠接近。
靜態形狀的 PyTorch 管道已應用通道優先記憶體格式和帶 reduce-overhead 模式的 torch.compile。
以下圖表顯示了每秒影像吞吐量與不同（批次大小，步數）組合下各種框架的對比。值得注意的是，每個條形上方的標籤表示相對於 Torch Compile 的加速百分比——例如，在第一個圖表中，對於（批次，步數）組合（4，1），ORT_TRT（靜態）比 Torch（編譯）快 31%。
我們選擇使用 1 步和 4 步，因為 SD Turbo 和 SDXL Turbo 都可以在最短 1 步內生成可用的影像，但通常在 3-5 步內生成質量最佳的影像。
SDXL Turbo
下面的圖表展示了 SDXL Turbo 模型在靜態和動態形狀下的每秒影像吞吐量。結果是在 A100-SXM4-80GB GPU 上針對不同的（批次大小，步數）組合收集的。對於動態形狀，TensorRT 引擎支援批次大小 1 到 8 和影像大小 512x512 到 768x768，但它針對批次大小 1 和影像大小 512x512 進行了最佳化。
SD Turbo
接下來的兩張圖表展示了 SD Turbo 模型在 A100-SXM4-80GB GPU 上靜態和動態形狀下的每秒影像吞吐量。
最後一組圖表展示了 SD Turbo 模型在 RTX-4090 GPU 上靜態和動態形狀下的每秒影像吞吐量。在此動態形狀測試中，TensorRT 引擎是為批次大小 1 到 8（針對批次大小 1 最佳化）和固定影像大小 512x512 構建的，因為記憶體限制。
ONNX Runtime 下 SD Turbo 和 SDXL Turbo 的速度如何？
這些結果表明，在靜態和動態形狀下，ONNX Runtime 在所有所示的（批次，步數）組合中，其 CUDA 和 TensorRT 執行提供程式均顯著優於 PyTorch。這個結論適用於兩種模型尺寸（SD Turbo 和 SDXL Turbo）以及兩種測試的 GPU。值得注意的是，ONNX Runtime 與 CUDA（動態形狀）在（批次，步數）組合（1，4）下比 Torch Eager 快 229%。
此外，由於在大多數（批次，步數）組合下 ORT_TRT 吞吐量高於相應的 ORT_CUDA 吞吐量，使用 TensorRT 執行提供程式的 ONNX Runtime 在靜態形狀方面表現稍好。當用戶在圖定義時知道批次和影像大小（例如，使用者只計劃生成批次大小為 1 且影像大小為 512x512 的影像）時，通常首選靜態形狀。在這些情況下，靜態形狀具有更快的效能。但是，如果使用者決定切換到不同的批次和/或影像大小，TensorRT 必須建立一個新的引擎（這意味著磁碟上的引擎檔案數量會增加一倍）並切換引擎（這意味著載入新引擎會花費額外時間）。
另一方面，在使用 A100-SXM4-80GB GPU 時，ONNX Runtime 與 CUDA 執行提供程式通常是 SD Turbo 和 SDXL Turbo 模型動態形狀的更好選擇，但在使用 RTX-4090 GPU 時，ONNX Runtime 與 TensorRT 執行提供程式在大多數（批次，步數）組合下動態形狀的表現稍好。使用動態形狀的好處是，當批次和影像大小直到圖執行時才可知時（例如，為一個影像執行批次大小為 1 且影像大小為 512x512，為另一個影像執行批次大小為 4 且影像大小為 512x768），使用者可以更快地執行推理。在這些情況下使用動態形狀時，使用者只需構建並儲存一個引擎，而無需在推理期間切換引擎。
GPU 最佳化
除了我們之前的 Stable Diffusion 部落格中介紹的技術外，ONNX Runtime 還應用了以下最佳化措施，以獲得本文所述的 SD Turbo 和 SDXL Turbo 結果：
- 啟用靜態形狀輸入的 CUDA 圖。
- 新增 Flash Attention V2。
- 刪除文字編碼器中的額外輸出（保留由 clip_skip 引數指定的隱藏狀態輸出）。
- 新增 SkipGroupNorm 融合，將組歸一化與其前面的 Add 節點融合。
此外，我們還增加了對新功能的支援，包括用於潛在一致性模型（LCM）的 LoRA 權重。
後續步驟
未來，我們計劃繼續改進 Stable Diffusion 工作，更新演示以支援新功能，如 IP Adapter 和 Stable Video Diffusion。ControlNet 支援也將很快推出。
我們還在努力最佳化 SD Turbo 和 SDXL Turbo 在我們現有 Stable Diffusion web UI 擴充套件上的效能，並計劃幫助將這兩種模型支援新增到 ONNX Runtime 社群成員開發的 Windows UI 中。
此外，關於如何使用 C# 和 ONNX Runtime 執行 SD Turbo 和 SDXL Turbo 的教程即將釋出。在此期間，請檢視我們之前關於 Stable Diffusion 的教程。
資源
檢視本文討論的一些資源
- SD Turbo：Olive 最佳化的 ONNX Runtime CUDA 模型，託管在 Hugging Face 上。
- SDXL Turbo：Olive 最佳化的 ONNX Runtime CUDA 模型，託管在 Hugging Face 上。
- Stable Diffusion GPU 最佳化：ONNX Runtime GitHub 倉庫中關於使用 NVIDIA GPU 最佳化 Stable Diffusion 的說明。
- Automatic1111 SD WebUI 的 ONNX Runtime 擴充套件：啟用在 NVIDIA GPU 上最佳化執行 Stable Diffusion UNet 模型的擴充套件。
- OnnxStack：社群貢獻的 .NET 庫，支援使用 C# 和 ONNX Runtime 進行 Stable Diffusion 推理。
- SD4J (Java 中的 Stable Diffusion)：Oracle 提供的 Java 和 ONNX Runtime Stable Diffusion 示例。
- 使用 C# 和 ONNX Runtime 進行 Stable Diffusion 推理：之前釋出的 C# 教程。
