关于人工智能:使用-PAIBlade-优化-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能力。

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的推理分为两局部。

1. ControlNet局部，其input blocks和  mid block 构造与Stable DiffusionUnet的前半部分雷同，残余局部为卷积。ControlNet所有输入传入到Stable DIffusion的Unet中，作为输出；
2. 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模型，不仅反对最为全面，同时性能和显存应用也是最佳程度。

目前PAI-Blade曾经陆续在相干业务中上线应用。接下来，咱们将持续优化性能，欠缺相干性能反对。欢送大家交换、分割和单干~