背景
上一篇中,我們使用了 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)
結合上述功能,可以同時實現:
- 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中涉及到網絡傳輸等模型無關部分的延遲,因此本部分只測試了模型部分耗時。結果如下:
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模型,不僅支持最爲全面,同時性能和顯存使用也是最佳水平。
本文爲阿里雲原創內容,未經允許不得轉載。