关于stable-diffusion:使用-PAIBlade-优化-Stable-Diffusion-推理流程二

8次阅读

共计 3792 个字符,预计需要花费 10 分钟才能阅读完成。

背景

上一篇中,咱们应用了 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 能力。

from torch_blade.monkey_patch import patch_utils

patch_utils.patch_conv2d(pipe.vae.decoder)
patch_utils.patch_conv2d(pipe.unet)

opt_cfg = torch_blade.Config()
...
opt_cfg.freeze_module = False
with opt_cfg, torch.no_grad():
    ...

如果没有 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,

controlnet = torch_blade.optimize(pipe.controlnet, model_inputs=tuple(controlnet_inputs), allow_tracing=True)

在优化 unet 模型时,因为 torch2.0 之前的版本,torch.jit.trace 不反对应用 dict 作为输出,所以咱们应用 Wrapper 包装 Unet 后便于 trace 和优化。同时,应用优化后的 ControlNet 执行一次推理,将其输入增加到 Unet 输出中。

class UnetWrapper(torch.nn.Module):
    def __init__(self, unet):
        super().__init__()
        self.unet = unet

    def forward(
        self,
        sample,
        timestep,
        encoder_hidden_states,
        down_block_additional_residuals,
        mid_block_additional_residual,
    ):
        return self.unet(
            sample,
            timestep,
            encoder_hidden_states=encoder_hidden_states,
            down_block_additional_residuals=down_block_additional_residuals,
            mid_block_additional_residual=mid_block_additional_residual,
        )

...
down_block_res_samples, mid_block_res_sample = controlnet(*controlnet_inputs)
unet_inputs += [tuple(down_block_res_samples), mid_block_res_sample]
unet = torch_blade.optimize(UnetWrapper(pipe.unet).eval(), model_inputs=tuple(unet_inputs), allow_tracing=True)

联合上述性能,能够同时实现:

  1. LoRA 权重替换;
  2. 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 模型,不仅反对最为全面,同时性能和显存应用也是最佳程度。

原文链接

本文为阿里云原创内容,未经容许不得转载。

正文完
 0