“零”代码改动，静态编译让太乙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：

github.com/Oneflow-Inc…

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 进行推理的功用表现。

“零”代码改动，静态编译让太乙Stable Diffusion推理速度翻倍

“零”代码改动，静态编译让太乙Stable Diffusion推理速度翻倍

能够看到，关于 A100 显卡，无论是 PCIe 40GB 的装备仍是 SXM 80GB 的装备，OneFlow 的功用相比 PyTorch 能提高 1 倍以上，推理速度达到了 50it/s 以上，生成一张图片所需求的时间在 1 秒以内。

其他硬件数据：

“零”代码改动，静态编译让太乙Stable Diffusion推理速度翻倍

注：3090上的AIT数据由 IDEA 研究院提供

“零”代码改动，静态编译让太乙Stable Diffusion推理速度翻倍

“零”代码改动，静态编译让太乙Stable Diffusion推理速度翻倍

“零”代码改动，静态编译让太乙Stable Diffusion推理速度翻倍

“零”代码改动，静态编译让太乙Stable Diffusion推理速度翻倍

综上，在各种硬件的比照中，比照 PyTorch， OneFlow 能将太乙 Stable Diffusion 的推理功用提高 1 倍多。

2、生成图片展现

滔滔江水, 连绵不绝, 唯美, 插画

“零”代码改动，静态编译让太乙Stable Diffusion推理速度翻倍

长城, 清晨, 模糊, 唯美, 插画

“零”代码改动，静态编译让太乙Stable Diffusion推理速度翻倍

梦回江南，中国古代小镇，唯美，插画

“零”代码改动，静态编译让太乙Stable Diffusion推理速度翻倍

中国的未来城市, 科幻插画

“零”代码改动，静态编译让太乙Stable Diffusion推理速度翻倍

古代建筑, 白雪纷飞

“零”代码改动，静态编译让太乙Stable Diffusion推理速度翻倍

螺蛳粉

“零”代码改动，静态编译让太乙Stable Diffusion推理速度翻倍

注：上述图片均根据 OneFlow 版太乙 Stable Diffusion 生成

3、无缝兼容 PyTorch 生态

想体验 OneFlow 版的太乙 Stable Diffusion？只需求修正两行代码：

“零”代码改动，静态编译让太乙Stable Diffusion推理速度翻倍

之所以能这么轻松搬迁模型，是因为 OneFlow Stable Diffusion 有两个超卓的特性：

  1. OneFlowStableDiffusionPipeline.from_pretrained 能够直接使用 PyTorch 权重。
  2. 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：

github.com/Oneflow-Inc…

OneFlow 地址：
github.com/Oneflow-Inc…

欢迎 Star、试用 OneFlow 最新版别：
github.com/Oneflow-Inc…