使用 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能力。

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中涉及到網絡傳輸等模型無關部分的延遲,因此本部分只測試了模型部分耗時。結果如下:

steps

eager

xformers

PAI-Blade

no LoRAs

+ 2 LoRAs

ControlNet

no LoRAs

+ 2 LoRAs

ControlNet

any LoRAs

ControlNet

20

2.03

2.94

2.75

1.57

2.46

2.14

1.15

1.62

50

4.77

7.17

6.64

3.63

5.86

5.06

2.59

3.75

100

9.45

14.18

13.13

7.10

11.54

9.90

4.96

7.35

由該表可知,webui在eager和xformers模式下,推理時間隨LoRA數量增加而延長,而PAI-Blade將所有LoRA的權重融合到基礎模型,所以推理時間與LoRA數量無關。

總結

這兩篇文章中,我們介紹了PAI-Blade 在Stable DIffusion模型上的優化經驗,目前已經支持了Diffusers和Stable-DIffusion-webui 兩種主流推理方式。

我們調研了相關公開的競品對Stable Diffusion的支持情況,結果如下:

框架/模型 Base Model LoRA ControlNet webui
xformers
AITemplete
OneFlow
TensorRT
PAI-Blade

根據公開性能數字和業務實測,PAI-Blade對Stable DIffusion模型,不僅支持最爲全面,同時性能和顯存使用也是最佳水平。

原文鏈接

本文爲阿里雲原創內容,未經允許不得轉載。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章