OneFlow 將 Stable Diffusion的推理性能推向了一個(gè)全新的SOTA

OneFlow 將 Stable Diffusion 的推理性能推向了一個(gè)全新的 SOTA。

第一輛汽車(chē)誕生之初，時(shí)速只有 16 公里，甚至不如馬車(chē)跑得快，很長(zhǎng)一段時(shí)間，汽車(chē)尷尬地像一種“很酷的玩具”。人工智能作圖的出現(xiàn)也是如此。

AI 作圖一開(kāi)始的 “風(fēng)格化” 本身就為 “玩” 而生，大家普遍興致勃勃地嘗試頭像生成、磨皮，但很快就失去興趣。直到擴(kuò)散模型的降臨，才給 AI 作圖帶來(lái)質(zhì)變，讓人們看到了 “AI 轉(zhuǎn)成生產(chǎn)力” 的曙光：畫(huà)家、設(shè)計(jì)師不用絞盡腦汁思考色彩、構(gòu)圖，只要告訴 Diffusion 模型想要什么，就能言出法隨般地生成高質(zhì)量圖片。

然而，與汽車(chē)一樣，如果擴(kuò)散模型生成圖片時(shí)“馬力不足”，那就沒(méi)法擺脫玩具的標(biāo)簽，成為人類手中真正的生產(chǎn)工具。

起初，AI 作圖需要幾天，再縮減到幾十分鐘，再到幾分鐘，出圖時(shí)間在不斷加速，問(wèn)題是，究竟快到什么程度，才會(huì)在專業(yè)的美術(shù)從業(yè)者甚至普通大眾之間普及開(kāi)來(lái)？

顯然，現(xiàn)在還無(wú)法給出具體答案。即便如此，可以確定的是 AI 作圖在技術(shù)和速度上的突破，很可能已經(jīng)接近甚至超過(guò)閾值，因?yàn)檫@一次，OneFlow 帶來(lái)了字面意義上 “一秒出圖” 的 Stable Diffusion 模型。

OneFlow Stable Diffusion 使用地址：https://github.com/Oneflow-Inc/diffusers/wiki/How-to-Run-OneFlow-Stable-Diffusion

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

比快更快，OneFlow 一馬當(dāng)先

下面的圖表分別展示了在 A100 (PCIe 40GB / SXM 80GB)、RTX 2080 和 T4 不同類型的 GPU 硬件上，分別使用 PyTorch, TensorRT, AITemplate 和 OneFlow 四種深度學(xué)習(xí)框架或者編譯器，對(duì) Stable Diffusion 進(jìn)行推理時(shí)的性能表現(xiàn)。

對(duì)于 A100 顯卡，無(wú)論是 PCIe 40GB 的配置還是 SXM 80GB 的配置，OneFlow 的性能可以在目前的最優(yōu)性能之上繼續(xù)提升 15% 以上。

特別是在 SXM 80GB A100 上，OneFlow 首次讓 Stable Diffusion 的推理速度達(dá)到了 50it/s 以上，首次把生成一張圖片需要采樣 50 輪的時(shí)間降到 1 秒以內(nèi)，是當(dāng)之無(wú)愧的性能之王。

在 T4 推理卡上，由于 AITemplate 暫不支持 Stable Diffsuion，相比于目前 SOTA 性能的 TensorRT，OneFlow 的性能是它的 1.5 倍。

而在 RTX2080 上，TensorRT 在編譯 Stable Diffsuion 時(shí)會(huì) OOM ，相比于目前 SOTA 性能的 PyTorch，OneFlow 的性能是它的 2.25 倍。

綜上，在各種硬件以及更多框架的對(duì)比中，OneFlow 都將 Stable Diffusion 的推理性能推向了一個(gè)全新的 SOTA。

生成圖片展示

利用 OneFlow 版的 Stable Diffusion，你可以把天馬行空的想法很快轉(zhuǎn)化成藝術(shù)圖片，譬如：

以假亂真的陽(yáng)光、沙灘和椰樹(shù)：

倉(cāng)鼠救火員、長(zhǎng)兔耳朵的狗子：

在火星上吃火鍋：

未來(lái)異世界 AI：

集齊 OneFlow 七龍珠：

圖片均基于 OneFlow 版 Stable Diffusion 生成。如果你一時(shí)沒(méi)有好的 idea，可以在 lexica 上參考一下廣大網(wǎng)友的創(chuàng)意，不僅有生成圖片還提供了對(duì)應(yīng)的描述文字。

無(wú)縫兼容 PyTorch 生態(tài)，實(shí)現(xiàn)一鍵模型遷移

想體驗(yàn) OneFlow Stable Diffusion？只需要修改三行代碼，你就可以將 HuggingFace 中的 PyTorch Stable Diffusion 模型改為 OneFlow 模型，分別是將 import torch 改為 import oneflow as torch 和將 StableDiffusionPipeline 改為 OneFlowStableDiffusionPipeline：

之所以能這么輕松遷移模型，是因?yàn)?OneFlow Stable Diffusion 有兩個(gè)出色的特性：

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

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

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

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

使用 OneFlow 運(yùn)行 Stable Diffusion

在 docker 中使用 OneFlow 運(yùn)行 StableDiffusion 模型生成圖片：

docker run --rm -it --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -v ${HF_HOME}:${HF_HOME} -v ${PWD}:${PWD} -w ${PWD} -e HF_HOME=${HF_HOME} -e HUGGING_FACE_HUB_TOKEN=${HUGGING_FACE_HUB_TOKEN} oneflowinc/oneflow-sd:cu112 python3 /demos/oneflow-t2i.py # --prompt "a photo of an astronaut riding a horse on mars"

更詳盡的使用方法請(qǐng)參考：https://github.com/Oneflow-Inc/diffusers/wiki/How-to-Run-OneFlow-Stable-Diffusion

后續(xù)工作

后續(xù) OneFlow 團(tuán)隊(duì)將積極推動(dòng) OneFlow 的 diffusers（https://github.com/Oneflow-Inc/diffusers.git） 和 transformers（https://github.com/Oneflow-Inc/transformers.git） 的 fork 倉(cāng)庫(kù)內(nèi)容合并到 huggingface 上游的的對(duì)應(yīng)倉(cāng)庫(kù)。這也是 OneFlow 首次以 transformers/diffusers 的后端的形式開(kāi)發(fā)模型，歡迎各位開(kāi)發(fā)者朋友在 GitHub 上反饋意見(jiàn)。

值得一提的是，在優(yōu)化和加速 Stable Diffusion 模型的過(guò)程中使用了 OneFlow 自研編譯器，不僅讓 PyTorch 前端搭建的 Stable Diffusion 在 NVIDIA GPU 上跑得更快，而且也可以讓這樣的模型在國(guó)產(chǎn) AI 芯片和 GPU 上跑得更快，這些將在之后的文章中揭秘技術(shù)細(xì)節(jié)。

審核編輯 ：李倩