0
  • 聊天消息
  • 系统消息
  • 评论与回复
登录后你可以
  • 下载海量资料
  • 学习在线课程
  • 观看技术视频
  • 写文章/发帖/加入社区
会员中心
创作中心

完善资料让更多小伙伴认识你,还能领取20积分哦,立即完善>

3天内不再提示

iPhone两秒出图,目前已知的最快移动端Stable Diffusion模型来了

3D视觉工坊 来源:机器之心 2023-06-12 15:25 次阅读

Stable Diffusion (SD)是当前最热门的文本到图像(text to image)生成扩散模型。尽管其强大的图像生成能力令人震撼,一个明显的不足是需要的计算资源巨大,推理速度很慢:以 SD-v1.5 为例,即使用半精度存储,其模型大小也有 1.7GB,近 10 亿参数,端上推理时间往往要接近 2min。

为了解决推理速度问题,学术界与业界已经开始对 SD 加速的研究,主要集中于两条路线:(1)减少推理步数,这条路线又可以分为两条子路线,一是通过提出更好的 noise scheduler 来减少步数,代表作是 DDIM [1],PNDM [2],DPM [3] 等;二是通过渐进式蒸馏(Progressive Distillation)来减少步数,代表作是 Progressive Distillation [4] 和 w-conditioning [5] 等。(2)工程技巧优化,代表作是 Qualcomm 通过 int8 量化 + 全栈式优化实现 SD-v1.5 在安卓手机上 15s 出图 [6],Google 通过端上 GPU 优化将 SD-v1.4 在三星手机上加速到 12s [7]。

尽管这些工作取得了长足的进步,但仍然不够快。

近日,Snap 研究院推出最新高性能 Stable Diffusion 模型,通过对网络结构、训练流程、损失函数全方位进行优化,在 iPhone 14 Pro 上实现 2 秒出图(512x512),且比 SD-v1.5 取得更好的 CLIP score。这是目前已知最快的端上 Stable Diffusion 模型!

5d25e23c-08ac-11ee-962d-dac502259ad0.png

论文地址:https://arxiv.org/pdf/2306.00980.pdf

Webpage: https://snap-research.github.io/SnapFusion

核心方法

Stable Diffusion 模型分为三部分:VAE encoder/decoder, text encoder, UNet,其中 UNet 无论是参数量还是计算量,都占绝对的大头,因此 SnapFusion 主要是对 UNet 进行优化。具体分为两部分:(1)UNet 结构上的优化:通过分析原有 UNet 的速度瓶颈,本文提出一套 UNet 结构自动评估、进化流程,得到了更为高效的 UNet 结构(称为 Efficient UNet)。(2)推理步数上的优化:众所周知,扩散模型在推理时是一个迭代的去噪过程,迭代的步数越多,生成图片的质量越高,但时间代价也随着迭代步数线性增加。为了减少步数并维持图片质量,我们提出一种 CFG-aware 蒸馏损失函数,在训练过程中显式考虑 CFG (Classifier-Free Guidance)的作用,这一损失函数被证明是提升 CLIP score 的关键!

下表是 SD-v1.5 与 SnapFusion 模型的概况对比,可见速度提升来源于 UNet 和 VAE decoder 两个部分,UNet 部分是大头。UNet 部分的改进有两方面,一是单次 latency 下降(1700ms -> 230ms,7.4x 加速),这是通过提出的 Efficient UNet 结构得到的;二是 Inference steps 降低(50 -> 8,6.25x 加速),这是通过提出的 CFG-aware Distillation 得到的。VAE decoder 的加速是通过结构化剪枝实现。

5d32e9a0-08ac-11ee-962d-dac502259ad0.png

下面着重介绍 Efficient UNet 的设计和 CFG-aware Distillation 损失函数的设计。

(1)Efficient UNet

我们通过分析 UNet 中的 Cross-Attention 和 ResNet 模块,定位速度瓶颈在于 Cross-Attention 模块(尤其是第一个 Downsample 阶段的 Cross-Attention),如下图所示。这个问题的根源是因为 attention 模块的复杂度跟特征图的 spatial size 成平方关系,在第一个 Downsample 阶段,特征图的 spatial size 仍然较大,导致计算复杂度高。

5d3f9132-08ac-11ee-962d-dac502259ad0.png

为了优化 UNet 结构,我们提出一套 UNet 结构自动评估、进化流程:先对 UNet 进行鲁棒性训练(Robust Training),在训练中随机 drop 一些模块,以此来测试出每个模块对性能的真实影响,从而构建一个 “对 CLIP score 的影响 vs. latency” 的查找表;然后根据该查找表,优先去除对 CLIP score 影响不大同时又很耗时的模块。这一套流程是在线自动进行,完成之后,我们就得到了一个全新的 UNet 结构,称为 Efficient UNet。相比原版 UNet,实现 7.4x 加速且性能不降。

(2)CFG-aware Step Distillation

CFG(Classifier-Free Guidance)是 SD 推理阶段的必备技巧,可以大幅提升图片质量,非常关键!尽管已有工作对扩散模型进行步数蒸馏(Step Distillation)来加速 [4],但是它们没有在蒸馏训练中把 CFG 纳入优化目标,也就是说,蒸馏损失函数并不知道后面会用到 CFG。这一点根据我们的观察,在步数少的时候会严重影响 CLIP score。

为了解决这个问题,我们提出在计算蒸馏损失函数之前,先让 teacher 和 student 模型都进行 CFG,这样损失函数是在经过 CFG 之后的特征上计算,从而显式地考虑了不同 CFG scale 的影响。实验中我们发现,完全使用 CFG-aware Distillation 尽管可以提高 CLIP score, 但 FID 也明显变差。我们进而提出了一个随机采样方案来混合原来的 Step Distillation 损失函数和 CFG-aware Distillation 损失函数,实现了二者的优势共存,既显著提高了 CLIP score,同时 FID 也没有变差。这一步骤,实现进一步推理阶段加速 6.25 倍,实现总加速约 46 倍。

除了以上两个主要贡献,文中还有对 VAE decoder 的剪枝加速以及蒸馏流程上的精心设计,具体内容请参考论文。

实验结果

SnapFusion 对标 SD-v1.5 text to image 功能,目标是实现推理时间大幅缩减并维持图像质量不降,最能说明这一点的是下图:

5d50e6f8-08ac-11ee-962d-dac502259ad0.png

该图是在 MS COCO’14 验证集上随机选取 30K caption-image pairs 测算 CLIP score 和 FID。CLIP score 衡量图片与文本的语义吻合程度,越大越好;FID 衡量生成图片与真实图片之间的分布距离(一般被认为是生成图片多样性的度量),越小越好。图中不同的点是使用不同的 CFG scale 获得,每一个 CFG scale 对应一个数据点。从图中可见,我们的方法(红线)可以达到跟 SD-v1.5(蓝线)同样的最低 FID,同时,我们方法的 CLIP score 更好。值得注意的是,SD-v1.5 需要 1.4min 生成一张图片,而 SnapFusion 仅需要 1.84s,这也是目前我们已知最快的移动端 Stable Diffusion 模型!

下面是一些 SnapFusion 生成的样本:

5d65abba-08ac-11ee-962d-dac502259ad0.png

更多样本请参考文章附录。

除了这些主要结果,文中也展示了众多烧蚀分析(Ablation Study)实验,希望能为高效 SD 模型的研发提供参考经验:

(1)之前 Step Distillation 的工作通常采用渐进式方案 [4, 5],但我们发现,在 SD 模型上渐进式蒸馏并没有比直接蒸馏更有优势,且过程繁琐,因此我们在文中采用的是直接蒸馏方案。

5d7888e8-08ac-11ee-962d-dac502259ad0.png

(2)CFG 虽然可以大幅提升图像质量,但代价是推理成本翻倍。今年 CVPR’23 Award Candidate 的 On Distillation 一文 [5] 提出 w-conditioning,将 CFG 参数作为 UNet 的输入进行蒸馏(得到的模型叫做 w-conditioned UNet),从而在推理时省却 CFG 这一步,实现推理成本减半。但是我们发现,这样做其实会造成图片质量下降,CLIP score 降低(如下图中,四条 w-conditioned 线 CLIP score 均未超过 0.30, 劣于 SD-v1.5)。而我们的方法则可以减少步数,同时将 CLIP score 提高,得益于所提出的 CFG-aware 蒸馏损失函数!尤其值得主要的是,下图中绿线(w-conditioned, 16 steps)与橙线(Ours,8 steps)的推理代价是一样的,但明显橙线更优,说明我们的技术路线比 w-conditioning [5] 在蒸馏 CFG guided SD 模型上更为有效。

5d86fdb0-08ac-11ee-962d-dac502259ad0.png

(3)既有 Step Distillation 的工作 [4, 5] 没有将原有的损失函数和蒸馏损失函数加在一起,熟悉图像分类知识蒸馏的朋友应该知道,这种设计直觉上来说是欠优的。于是我们提出把原有的损失函数加入到训练中,如下图所示,确实有效(小幅降低 FID)。

5d9456fe-08ac-11ee-962d-dac502259ad0.png

总结与未来工作

本文提出 SnapFusion,一种移动端高性能 Stable Diffusion 模型。SnapFusion 有两点核心贡献:(1)通过对现有 UNet 的逐层分析,定位速度瓶颈,提出一种新的高效 UNet 结构(Efficient UNet),可以等效替换原 Stable Diffusion 中的 UNet,实现 7.4x 加速;(2)对推理阶段的迭代步数进行优化,提出一种全新的步数蒸馏方案(CFG-aware Step Distillation),减少步数的同时可显著提升 CLIP score,实现 6.25x 加速。总体来说,SnapFusion 在 iPhone 14 Pro 上实现 2 秒内出图,这是目前已知最快的移动端 Stable Diffusion 模型。

未来工作:

1.SD 模型在多种图像生成场景中都可以使用,本文囿于时间,目前只关注了 text to image 这个核心任务,后期将跟进其他任务(如 inpainting,ControlNet 等等)。

2. 本文主要关注速度上的提升,并未对模型存储进行优化。我们相信所提出的 Efficient UNet 仍然具备压缩的空间,结合其他的高性能优化方法(如剪枝,量化),有望缩小存储,并将时间降低到 1 秒以内,离端上实时 SD 更进一步。

声明:本文内容及配图由入驻作者撰写或者入驻合作网站授权转载。文章观点仅代表作者本人,不代表电子发烧友网立场。文章及其配图仅供工程师学习之用,如有内容侵权或者其他违规问题,请联系本站处理。 举报投诉
  • iPhone
    +关注

    关注

    28

    文章

    13460

    浏览量

    201662
  • 模型
    +关注

    关注

    1

    文章

    3226

    浏览量

    48809
  • 网络结构
    +关注

    关注

    0

    文章

    48

    浏览量

    11077

原文标题:iPhone两秒出图,目前已知的最快移动端Stable Diffusion模型来了

文章出处:【微信号:3D视觉工坊,微信公众号:3D视觉工坊】欢迎添加关注!文章转载请注明出处。

收藏 人收藏

    评论

    相关推荐

    Stable Diffusion的完整指南:核心基础知识、制作AI数字人视频和本地部署要求

    Stable Diffusion是一种扩散模型diffusion model)的变体,叫做“潜在扩散模型”(latent
    的头像 发表于 09-18 10:06 2444次阅读
    <b class='flag-5'>Stable</b> <b class='flag-5'>Diffusion</b>的完整指南:核心基础知识、制作AI数字人视频和本地部署要求

    uCOS-II中断延迟小于两秒

    */OSTimeDlyHMSM(0,0,2,0); /*延迟2s*/ EXTI_ClearITPendingBit(EXTI_Line9);} 中断里的延迟设为两秒了,为什么实际执行起来时间远远小于两秒啊??也没有其他的中断来抢占。求助!!
    发表于 04-28 02:59

    Stability AI开源图像生成模型Stable Diffusion

    Stable Diffusion 的很多用户已经公开发布了生成图像的样例,Stability AI 的首席开发者 Katherine Crowson 在推特上分享了许多图像。基于 AI 的图像合成可能会对艺术家和艺术领域带来一定的影响
    的头像 发表于 09-21 15:37 2952次阅读

    大脑视觉信号被Stable Diffusion复现图像!“人类的谋略和谎言不存在了”

    它不需要在复杂的深度学习模型上进行训练或做精细的微调,只需要做好fMRI(功能磁共振成像技术)成像到Stable Diffusion中潜在表征的简单线性映射关系就行。
    的头像 发表于 03-08 10:21 762次阅读

    大脑视觉信号被Stable Diffusion复现图像!

    它不需要在复杂的深度学习模型上进行训练或做精细的微调,只需要做好fMRI(功能磁共振成像技术)成像到Stable Diffusion中潜在表征的简单线性映射关系就行。
    的头像 发表于 03-16 10:26 813次阅读

    一文读懂Stable Diffusion教程,搭载高性能PC集群,实现生成式AI应用

    PC Farm、生成式AI和Stable Diffusion模型都是非常有用的工具和技术,可以帮助用户快速构建和管理计算机集群,生成高质量的数据和图像,提高模型的泛化能力和鲁棒性。
    的头像 发表于 05-01 07:47 2339次阅读
    一文读懂<b class='flag-5'>Stable</b> <b class='flag-5'>Diffusion</b>教程,搭载高性能PC集群,实现生成式AI应用

    使用OpenVINO™在算力魔方上加速stable diffusion模型

    Stable Diffusion 是 stability.ai 开源的 AI 图像生成模型,实现输入文字,生成图像的功能。Stable Diffus
    的头像 发表于 05-12 09:10 1462次阅读
    使用OpenVINO™在算力魔方上加速<b class='flag-5'>stable</b> <b class='flag-5'>diffusion</b><b class='flag-5'>模型</b>

    优化 Stable Diffusion 在 GKE 上的启动体验

    以下文章来源于谷歌云服务,作者 Google Cloud 背景 现如今随着 AIGC 这个话题越来越热,越来越多优秀的开源项目基于文生的 AI 模型如 MidJourney,Stable
    的头像 发表于 06-03 08:35 887次阅读

    基于一种移动高性能 Stable Diffusion 模型

      Stable Diffusion (SD)是当前最热门的文本到图像(text to image)生成扩散模型。尽管其强大的图像生成能力令人震撼,一个明显的不足是需要的计算资源巨大,推理速度很慢
    的头像 发表于 06-12 10:14 791次阅读
    基于一种<b class='flag-5'>移动</b><b class='flag-5'>端</b>高性能 <b class='flag-5'>Stable</b> <b class='flag-5'>Diffusion</b> <b class='flag-5'>模型</b>

    美格智能高算力AI模组成功运行Stable Diffusion模型

    研发团队成功在自研高算力模组上运行文生模型Stable Diffusion,而这仅仅是我们将思索变为实际行动的第一步,AIGC的大时代,才刚刚开始!
    的头像 发表于 07-21 11:17 836次阅读

    使用OpenVINO在Stable Diffusion V2.1上实现AI硬件加速的方法

    随着 AIGC 模型越来越强大,并取得了更惊人的结果,任意运行 AIGC 模型, 比如 Stable Diffusion,对于开发人员来说仍然面临一些挑战。首先,GPU 的安装设置需要
    的头像 发表于 07-21 11:47 2026次阅读
    使用OpenVINO在<b class='flag-5'>Stable</b> <b class='flag-5'>Diffusion</b> V2.1上实现AI硬件加速的方法

    树莓派能跑Stable Diffusion了?

    Stable Diffusion是一种文本到图像生成的大型深度学习模型,它可以根据文本的描述生成详细的图像,也可以用于其他任务,如图像修复、图像扩展、图像翻译等。
    的头像 发表于 07-26 11:46 1543次阅读

    Stable Diffusion的完整指南:核心基础知识、制作AI数字人视频和本地部署要求

    Stable Diffusion;stable;diffusion;SD;stablediffussion;sadtalker;PC集群;PC Farm;PC农场;GPU集群;GAN
    的头像 发表于 09-07 14:12 2199次阅读
    <b class='flag-5'>Stable</b> <b class='flag-5'>Diffusion</b>的完整指南:核心基础知识、制作AI数字人视频和本地部署要求

    Stability AI试图通过新的图像生成人工智能模型保持领先地位

    Stability AI的最新图像生成模型Stable Cascade承诺比其业界领先的前身Stable Diffusion更快、更强大,而Stab
    的头像 发表于 02-19 16:03 934次阅读
    Stability AI试图通过新的图像生成人工智能<b class='flag-5'>模型</b>保持领先地位

    如何开启Stable Diffusion WebUI模型推理部署

    如何开启Stable Diffusion WebUI模型推理部署
    的头像 发表于 12-11 20:13 97次阅读
    如何开启<b class='flag-5'>Stable</b> <b class='flag-5'>Diffusion</b> WebUI<b class='flag-5'>模型</b>推理部署