使用 PAI-Blade 优化 Stable Diffusion 推理流程 (二)
背景
上一篇中,我们使用了 PAI-Blade 优化了 diffusers 中 Stable Diffusion 模型。本篇,我们继续介绍使用 PAI-Blade 优化 LoRA 和 Controlnet 的推理流程。相关优化已经同样在 registry.cn-beijing.aliyuncs.com/blade_demo/blade_diffusion
镜像中可以直接使用。同时,我们将介绍 Stable-Diffusion-webui 中集成 PAI-Blade 优化的方法。
LoRA 优化
PAI-Blade 优化 LoRA 的方式,与前文方法基本相同。包括:加载模型、优化模型、替换原始模型。以下仅介绍与前文不同的部分。
首先,加载 Stable DIffusion 模型后,需要加载 LoRA 权重。
pipe.unet.load_attn_procs("lora/")
使用 LoRA 时,用户可能需要切换不同的 LoRA 权重,尝试不同的风格。因此,PAI-Blade 需要在优化配置中,传入freeze_module=False
,使得优化过程中,不对权重进行编译优化,从而不影响模型加载权重的功能。通过这种方式,PAI-Blade 优化后的模型,依然可以使用pipe.unet.load_attn_procs()
方式加载 LoRA 的权重,而不需要重新编译优化。
由于模型权重未进行优化流程,一些对常量的优化无法进行,因此会损失部分优化空间。为了解决性能受损的问题,PAI-Blade 中,使用了部分 patch,对原始模型进行 python 层级的替换,使得模型更适合 PAI-Blade 优化。通过在优化前,使用 torch_blade.monkey_patch
优化 Stable Diffusion 模型中的 unet 和 vae 部分,能更好的发挥 PAI-Blade 能力。
如果没有 LoRA 权重切换的需求,可以忽略上述步骤,获得更快的推理速度。
Benchmark
我们在 A100/A10 上测试了上述对 LoRA 优化的结果,测试模型为 runwayml/stable-diffusion-v1-5
,测试采样步数为 50。
ControlNet 适配
根据 ControlNet 的模型结构图以及diffusers中ControlNet实现,可以将 ControlNet 的推理分为两部分。
ControlNet 部分,其 input blocks 和 mid block 结构与 Stable DiffusionUnet 的前半部分相同,剩余部分为卷积。ControlNet 所有输出传入到 Stable DIffusion 的 Unet 中,作为输入;
Stable Diffusion 的 Unet 除了原始输入外,额外增加了 ControlNet 的输出作为输入。
根据上述特点,我们可以做出以下的优化:
首先,优化 ControlNet,
在优化 unet 模型时,由于 torch2.0 之前的版本,torch.jit.trace 不支持使用 dict 作为输入,所以我们使用 Wrapper 包装 Unet 后便于 trace 和优化。同时,使用优化后的 ControlNet 执行一次推理,将其输出添加到 Unet 输入中。
结合上述功能,可以同时实现:
LoRA 权重替换;
ControlNet 权重替换,来使用不同 ControlNet model。
benchmark
我们在 A100/A10 上测试了上述对 ControlNet 优化的结果,测试模型为 runwayml/stable-diffusion-v1-5
,测试采样步数为 50。
小结
在上述部分,我们使用了 PAI-Blade 优化了 Stable DIffusion 模型的 encoder、unet、decoder 部分,大幅降低推理延时的同时,减少了显存占用,从而降低 Stable DIffusion 模型推理成本。同时,PAI-Blade 支持了 LoRA、ControlNet 等常用功能,扩展了 PAI-Blade 的实用性。
webui 适配
stable-diffusion-webui 是 Stable DIffusion 非常热门的应用,PAI-Blade 同样提供了对其优化支持。目前,PAI-Blade 已经支持了模型权重切换、LoRA、ControlNet 等 webui 中常用的功能,同时通过 extension 的形式集成,可以便于用户使用。目前,相关优化已经集成到 PAI-EAS 的 eas-registry.cn-hangzhou.cr.aliyuncs.com/pai-eas/sdwebui-inference:0.0.2-py310-gpu-cu117-ubuntu2204-blade
镜像,可以通过 PAI_EAS 直接体验 PAI-Blade 的优化能力。
下面介绍该插件中,PAI-Blade 在 webui 中优化方式和性能。webui 优化原理与 diffusers 大致相同,以下是几个主要不同点:
分模块优化 Unet 和 ControlNet
由于 webui 中,ControlNet 需要逐个调用 Unet 的子模块,为了兼顾 ControlNet,PAI-Blade 并没有像 diffusers 中一样,优化整个 Unet 和 ControlNet。而是采取逐个子模块优化的方法,将 Unet、ControlNet 中所有的 down blocks、 mid block、up blocks 分别进行优化和替换。经过测试,此种优化方式几乎不影响模型推理速度。
不冻结权重
webui 的网页上,可以快捷的切换模型权重。因此,PAI-Blade 采取和 diffusers 中 LoRA 优化同样的方法,不对权重进行优化。
LoRA 优化
webui 中,多个 LoRA 会逐个调用 LoRA 计算,计算时间随 LoRA 数量增多而变长。PAI-Blade 在加载 LoRA 权重时,将多个 LoRA 的权重与 scale 预先 fuse,减少了运行时的开销。加载和 fuse 的开销,经测试可忽略不计。
Benchmark
我们在 A10 上测试了 webui 中,Stable DIffusion V1 模型在 batch size 为 1,分辨率为 512*512 条件下的推理速度。由于 webui 中涉及到网络传输等模型无关部分的延迟,因此本部分只测试了模型部分耗时。结果如下:
由该表可知,webui 在 eager 和 xformers 模式下,推理时间随 LoRA 数量增加而延长,而 PAI-Blade 将所有 LoRA 的权重融合到基础模型,所以推理时间与 LoRA 数量无关。
总结
这两篇文章中,我们介绍了 PAI-Blade 在 Stable DIffusion 模型上的优化经验,目前已经支持了 Diffusers 和 Stable-DIffusion-webui 两种主流推理方式。
我们调研了相关公开的竞品对 Stable Diffusion 的支持情况,结果如下:
根据公开性能数字和业务实测,PAI-Blade 对 Stable DIffusion 模型,不仅支持最为全面,同时性能和显存使用也是最佳水平。
目前 PAI-Blade 已经陆续在相关业务中上线使用。接下来,我们将继续优化性能,完善相关功能支持。欢迎大家交流、联系和合作~
版权声明: 本文为 InfoQ 作者【阿里云大数据AI技术】的原创文章。
原文链接:【http://xie.infoq.cn/article/5ce6645686a749c3bb11c4866】。文章转载请联系作者。
评论