“零”代碼改動，靜態編譯讓太乙Stable Diffusion推理速度翻倍

作者｜梁德澎

AI 作圖領域的工具一直不盡人意，直到去年 8 月 Stable Diffusion 開源，成為AI 圖像生成領域無可爭辯的劃時代模型。

(資料圖片)

為了提升其推理效率，OneFlow 首度將 Stable Diffusion 模型加速至“一秒出圖”時代，極大提升了文生圖的速度，在AIGC領域引發巨大反響，并得到了 Stability.ai 官方的支持。至今，OneFlow 還在不斷刷新 SOTA 紀錄。

不過，由于目前大部分團隊主要是基于翻譯 API + 英文 Stable Diffusion 模型進行開發，所以在使用中文獨特的敘事和表達時，英文版模型就很難給出正確匹配的圖片內容，這對部分國內用戶來說不太方便。

為了解決這一問題，國內的IDEA 研究院認知計算與自然語言研究中心（IDEA CCNL）也開源了第一個中文版本的“太乙 Stable Diffusion”，基于0.2億篩選過的中文圖文對訓練。上個月，太乙 Stable Diffusion 在 HuggingFace 上有近 15 萬下載量，是下載量最大的中文 Stable Diffusion。

近期，OneFlow 團隊為太乙 Stable Diffusion 適配了 OneFlow 后端，大大提升了推理性能，也可以做到一秒出圖。不少開發者好奇OneFlow使用了哪些優化“秘笈”，后文將進行簡要解讀。

歡迎Star、運行 OneFlow 版太乙 Stable Diffusion： ?

https://github.com/Oneflow-Inc/diffusers/wiki/How-to-Run-OneFlow-Stable-Diffusion#without-docker

1對比 PyTorch，OneFlow 將“太乙 Stable Diffusion”推理速度提升1倍以上

下面的圖表分別展示了在 A100 (PCIe 40GB / SXM 80GB)，V100 ( SXM2 32GB )， RTX 2080，RTX 3080 Ti，RTX 3090， 和 T4 不同類型的 GPU 硬件上分別使用 PyTorch, 和 OneFlow對 太乙 Stable Diffusion 進行推理的性能表現。

可以看到，對于 A100 顯卡，無論是 PCIe 40GB 的配置還是 SXM 80GB 的配置，OneFlow 的性能相比 PyTorch 能提升 1 倍以上，推理速度達到了 50it/s 以上，生成一張圖片所需要的時間在 1 秒以內。

其他硬件數據：

注：3090上的AIT數據由 IDEA 研究院提供

綜上，在各種硬件的對比中，對比 PyTorch， OneFlow 能將太乙 Stable Diffusion 的推理性能提升 1 倍多。

2

生成圖片展示

滔滔江水, 連綿不絕, 唯美, 插畫

長城, 清晨, 朦朧, 唯美, 插畫

夢回江南，中國古代小鎮，唯美，插畫

中國的未來城市, 科幻插畫

古代建筑, 白雪紛飛

螺螄粉

注：上述圖片均基于 OneFlow 版太乙 Stable Diffusion 生成

3

無縫兼容 PyTorch 生態

想體驗 OneFlow 版的太乙 Stable Diffusion？只需要修改兩行代碼：

之所以能這么輕松遷移模型，是因為 OneFlow Stable Diffusion 有兩個出色的特性：

OneFlowStableDiffusionPipeline.from_pretrained 能夠直接使用 PyTorch 權重。

OneFlow 本身的 API 和 PyTorch 對齊，因此 import oneflow as torch 之后，torch.autocast、torch.float16 等表達式完全不需要修改。

上述特性使得 OneFlow 兼容了 PyTorch 的生態，這不僅在 OneFlow 對 太乙 Stable Diffusion 的遷移中發揮了作用，也大大加速了 OneFlow 用戶遷移其它許多模型，比如在和 torchvision 對標的 flowvision 中，許多模型只需通過在 torchvision 模型文件中加入 import oneflow as torch 即可得到。

此外，OneFlow 還提供全局 “mock torch” 功能，在命令行運行 eval $(oneflow-mock-torch) 就可以讓接下來運行的所有 Python 腳本里的 import torch 都自動指向 oneflow。 ?

4

動靜一體的編程體驗

深度學習算法原型開發階段需要快速修改和調試，動態圖執行（Eager mode, define by run）最優。但在部署階段，模型已經固定下來，計算效率變得更重要，靜態圖執行（Lazy mode，define and run）可以借助編譯器做靜態優化來獲得更好的性能。因此，推理階段主要使用靜態圖模式。

最近，PyTorch 升級到2.0引入了compile()這個API，可以把一個模型或一個Module從動態圖執行變成靜態圖執行。OneFlow里也有一個類似的機制，不過接口名是nn.Graph()，它可以把傳入Module轉成靜態圖執行模式。

不僅如此，OneFlow的nn.Graph模式基于MLIR實現了一系列計算圖的圖層優化，譬如內存布局、算子融合等。 ?

這不僅使得計算圖表示的深度學習模型可以在各種硬件上達到最高性能，更重要的是，使得深度學習框架導入的計算圖更方便地在不同硬件之間實現遷移，有助于克服國產硬件軟件生態薄弱的問題。未來，我們將發布更多內容來揭示OneFlow深度學習編譯器的設計和實現。

歡迎Star、運行 OneFlow 版太乙 Stable Diffusion：

https://github.com/Oneflow-Inc/diffusers/wiki/How-to-Run-OneFlow-Stable-Diffusion#without-docker ?

OneFlow 地址：https://github.com/Oneflow-Inc/oneflow/ ?

其他人都在看

35張圖，直觀理解Stable Diffusion

ChatGPT的一小步，NLP范式轉變的一大步

2023年AI十大展望：GPT-4領銜大模型變革

李白：你的模型權重很不錯，可惜被我沒收了

OpenAI掌門Sam Altman：AI下一個發展階段

比快更快，開源Stable Diffusion刷新作圖速度

OneEmbedding:單卡訓練TB級推薦模型不是夢

歡迎Star、試用OneFlow最新版本：https://github.com/Oneflow-Inc/oneflow/https://github.com/Oneflow-Inc/oneflow/

關鍵詞：