作者｜梁德澎

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

(資料圖片)

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

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

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

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

歡迎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 進(jìn)行推理的性能表現(xiàn)。

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

其他硬件數(shù)據(jù)：

注：3090上的AIT數(shù)據(jù)由 IDEA 研究院提供

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

2

生成圖片展示

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

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

夢回江南，中國古代小鎮(zhèn)，唯美，插畫

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

古代建筑, 白雪紛飛

螺螄粉

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

3

無縫兼容 PyTorch 生態(tài)

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

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

OneFlowStableDiffusionPipeline.from_pretrained 能夠直接使用 PyTorch 權(quán)重。

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

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

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

4

動靜一體的編程體驗

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

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

不僅如此，OneFlow的nn.Graph模式基于MLIR實現(xiàn)了一系列計算圖的圖層優(yōu)化，譬如內(nèi)存布局、算子融合等。 ?

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

歡迎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范式轉(zhuǎn)變的一大步

2023年AI十大展望：GPT-4領(lǐng)銜大模型變革

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

OpenAI掌門Sam Altman：AI下一個發(fā)展階段

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

OneEmbedding:單卡訓(xùn)練TB級推薦模型不是夢

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

關(guān)鍵詞：