使用 ONNX Runtime 與 Olive 加速 SD Turbo 與 SDXL Turbo 推論

TL;DR

Hugging Face 與 Microsoft 已使用 ONNX Runtime 與 Olive 優化工具對 SD Turbo 與 SDXL Turbo 推論進行優化，使 SDXL Turbo 的吞吐量提升最高可達 229%，SD Turbo 提升 120%，相較於 PyTorch。這些優化使得僅需一步即可進行高效能影像生成，並將模型的可用性擴展至 C#、Java 等非 Python 語言。

效能基準：ONNX Runtime 與 PyTorch 比較

ONNX Runtime 在所有測試的批次大小與步數上，於 NVIDIA GPU 上的 SD Turbo 與 SDXL Turbo 模型均顯著優於 PyTorch。

主要吞吐量提升

  • SDXL Turbo: 吞吐量提升最高達 229%，相較於 PyTorch。
  • SD Turbo: 吞吐量提升最高達 120%，相較於 PyTorch。
  • 執行提供者: CUDA 與 TensorRT 兩種執行提供者在靜態與動態形狀下皆相較於 PyTorch 有顯著提升。

硬體與設定結果

基準測試使用 A100-SXM4-80GB 與 RTX-4090 GPU，搭配 LCM Scheduler 與 fp16 模型進行。

  • 靜態 vs. 動態形狀: 當批次與影像尺寸在圖形定義時已知，靜態形狀通常提供更快的效能。動態形狀則更具彈性，允許使用者在執行期間變更批次與影像尺寸，而無需重新建構引擎。
  • GPU 特色: 在 A100 GPU 上，使用 CUDA 執行提供者的 ONNX Runtime 通常是動態形狀的較佳選擇。於 RTX-4090 GPU 上，TensorRT 執行提供者在動態形狀下通常表現稍佳。

技術優化與工具

效能提升來自 Olive 模型優化工具以及 ONNX Runtime 內部的特定 GPU 級別增強。

Olive 優化工具

模型使用 Olive 產生，Olive 是一款具備硬體感知的模型優化工具。為獲得最佳效能，必須透過指令列啟用 fp16 VAE。

GPU 專屬增強

除了標準的 Stable Diffusion 優化外，還實作了以下特定技術改進：

  • CUDA Graph: 為靜態形狀輸入啟用，以減少開銷。
  • Flash Attention V2: 整合以加速注意力機制。
  • 文字編碼器優化: 移除文字編碼器的額外輸出，只保留 clip_skip 參數指定的隱藏狀態輸出。
  • SkipGroupNorm 融合: 將群組正規化與前置的 Add 節點融合。
  • LoRA 支援: 為潛在一致性模型 (LCM) 添加 LoRA 權重支援。

跨平台可及性與整合

透過使用 ONNX Runtime，SD Turbo 與 SDXL Turbo 不再受限於 Python 環境，得以整合至更廣泛的軟體堆疊中。

語言支援

  • C#: 社群專案如 OnnxStack 提供 .NET 程式庫以進行 Stable Diffusion 推論。
  • Java: Oracle 已發布一個 Stable Diffusion 範例 (sd4j)，在 ONNX Runtime 上執行推論。

Web UI 整合

針對 Automatic1111 的 SD WebUI 的 ONNX Runtime 擴充套件，允許使用 CUDA 執行提供者與 Olive 優化模型，在 NVIDIA GPU 上優化執行 Stable Diffusion UNet 模型。

未來路線圖

Hugging Face 與 Microsoft 計畫將這些優化擴展至更多功能與模型，包括：

  • 功能支援: 整合 IP Adapter 與 ControlNet。
  • 模型支援: 擴展至 Stable Video Diffusion。
  • UI 改進: 進一步優化現有 Stable Diffusion Web UI 擴充套件中的 SD Turbo 與 SDXL Turbo 效能，並支援社群開發的 Windows UI。

Sources