Implementation Pipeline of Stable Diffusion with ControlNet
@zilla0717
本文梳理了用ControlNet控制Stable Diffusion输出的实现思路。
分析对象
StableDiffusion WebUI
ControlNet Extension for StableDiffusion WebUI
ControlNet作为StableDiffusion WebUI的扩展,遵照其扩展开发规则。
参考资料
【StableDiffusion WebUI源码分析 — 知乎】
1. Gradio的基本用法
2. txt2img的实现
3. 模型加载的过程
4. 启动流程
5. 多语言的实现方式
6. 脚本的实现方式
7. 扩展的实现方式
8. Lora功能的实现方式
【StableDiffusion WebUI的Wiki】
【gradio UI component】
1. 实现扩展的一般流程
插件目录下,各文件、子目录作用如下:
-
install.py
:若有则自动执行,用于完成依赖库的安装。 - 子目录
scripts
:放py脚本,插件目录会被追加到sys.path
。建议脚本中用scripts.basedir()
来获取当前插件目录,因为用户可能重命名插件。 -
style.css
和子目录javascript
中的js文件会被加载到页面上。 -
preload.py
:若有,则在程序解析命令之前加载。在该文件里的preload
函数中追加与该扩展有关的命令行参数。如:
def preload(parser):
parser.add_argument("--wildcards-dir", type=str, default=None)
下面说明如何编写一个py脚本,以“旋转生成的图片”这一脚本为例(分析见注释)。
- import必要的包和函数(这部分不需要改动)
import modules.scripts as scripts
import gradio as gr
import os
from modules import images
from modules.processing import process_images, Processed
from modules.processing import Processed
from modules.shared import opts, cmd_opts, state
- 定义
Script
类,后续的title()
、show()
、ui()
、run()
都是该类的函数
class Script(scripts.Script)
-
title()
:定义脚本名称(显示在该插件的下拉菜单里)
def title(self):
return "Rotate Output"
-
show()
:其返回值控制该选项何时出现在下拉菜单
def show(self, is_img2img):
# 只有在img2img 界面才在下拉菜单显示该功能
return is_img2img
-
ui()
:定义这个脚本在UI上怎么展示,其返回值被用作参数
多数UI组件返回的是boolean。
def ui(self, is_img2img):
angle = gr.Slider(minimum=0.0, maximum=360.0, step=1, value=0,
label="Angle")
overwrite = gr.Checkbox(False, label="Overwrite existing files")
return [angle, overwrite]
-
run()
:获取UI传回的参数,做额外的计算过程
该函数在这个脚本在下拉菜单中被选中时被调用,它必须进行所有处理并返回带有结果的Processed
对象(与processing.process_images()
返回的结果相同)。
通常处理过程是调用process_images()
完成的。- 入参
-
p
(类型为StableDiffusionProcessing
的对象实例)
StableDiffusionProcessing
定义参见module/processing.py
,定义了它以及子类StableDiffusionProcessingTxt2Img
和StableDiffusionProcessingImg2Img
。 -
ui()
返回的参数
-
-
run()
内部可以自定义函数和引入额外的包。 - 对图片执行运算的函数以由
process_images()
返回的Processed
对象proc
和ui()
获取的参数 为入参,原始图片在proc.images
,返回处理后的proc
。
- 入参
def run(self, p, angle, overwrite):
def rotate(im, angle):
from PIL import Image
raf = im
if angle != 0:
raf = raf.rotate(angle, expand=True)
return raf
basename = ""
if(not overwrite):
if angle != 0:
basename += "rotated_" + str(angle)
else:
p.do_not_save_samples = True
proc = process_images(p)
for i in range(len(proc.images)):
proc.images[i] = rotate(proc.images[i], angle)
images.save_image(proc.images[i], p.outpath_samples, basename, proc.seed + i, proc.prompt, opts.samples_format, info= proc.info, p=p)
return proc
-
process()
:获取UI传回的参数,做额外的计算过程
该函数类似run()
,区别是它在开始执行总是可见的脚本前被调用,即在图像处理前被调用。
before_process_batch()
、process_batch()
、postprocess_batch()
等函数的作用见modules/scripts.py
。
2. ControlNet扩展的UI实现和回调方法
controlnet.py
的写法类似上面的例子,其ui()
实现如下:
def ui(self, is_img2img):
self.infotext_fields = []
self.paste_field_names = []
controls = ()
max_models = shared.opts.data.get("control_net_max_models_num", 1)
elem_id_tabname = ("img2img" if is_img2img else "txt2img") + "_controlnet"
with gr.Group(elem_id=elem_id_tabname):
with gr.Accordion(f"ControlNet {controlnet_version.version_flag}", open = False, elem_id="controlnet"):
if max_models > 1:
with gr.Tabs(elem_id=f"{elem_id_tabname}_tabs"):
for i in range(max_models):
with gr.Tab(f"ControlNet Unit {i}",
elem_classes=['cnet-unit-tab']):
controls += (self.uigroup(f"ControlNet-{i}", is_img2img, elem_id_tabname),)
else:
with gr.Column():
controls += (self.uigroup(f"ControlNet", is_img2img, elem_id_tabname),)
if shared.opts.data.get("control_net_sync_field_args", False):
for _, field_name in self.infotext_fields:
self.paste_field_names.append(field_name)
return controls
在api.py
中,可以看到 在web app启动(on_app_started
)时就会调用controlnet_api()
方法。
try:
import modules.script_callbacks as script_callbacks
script_callbacks.on_app_started(controlnet_api)
except:
pass
controlnet_api()
中定义了一些异步的方法(其中获取插件模型列表、版本、设置等信息的方法由GET请求调用,detect()
由POST请求调用),实现如下:
def controlnet_api(_: gr.Blocks, app: FastAPI):
@app.get("/controlnet/version")
async def version():
return {"version": external_code.get_api_version()}
@app.get("/controlnet/model_list")
async def model_list():
up_to_date_model_list = external_code.get_models(update=True)
logger.debug(up_to_date_model_list)
return {"model_list": up_to_date_model_list}
@app.get("/controlnet/module_list")
async def module_list(alias_names: bool = False):
_module_list = external_code.get_modules(alias_names)
logger.debug(_module_list)
return {
"module_list": _module_list,
"module_detail": external_code.get_modules_detail(alias_names)
}
@app.get("/controlnet/settings")
async def settings():
max_models_num = external_code.get_max_models_num()
return {"control_net_max_models_num":max_models_num}
cached_cn_preprocessors = global_state.cache_preprocessors(global_state.cn_preprocessor_modules)
@app.post("/controlnet/detect")
async def detect(
controlnet_module: str = Body("none", title='Controlnet Module'),
controlnet_input_images: List[str] = Body([], title='Controlnet Input Images'),
controlnet_processor_res: int = Body(512, title='Controlnet Processor Resolution'),
controlnet_threshold_a: float = Body(64, title='Controlnet Threshold a'),
controlnet_threshold_b: float = Body(64, title='Controlnet Threshold b')
):
controlnet_module = global_state.reverse_preprocessor_aliases.get(controlnet_module, controlnet_module)
if controlnet_module not in cached_cn_preprocessors:
raise HTTPException(
status_code=422, detail="Module not available")
if len(controlnet_input_images) == 0:
raise HTTPException(
status_code=422, detail="No image selected")
logger.info(f"Detecting {str(len(controlnet_input_images))} images with the {controlnet_module} module.")
results = []
processor_module = cached_cn_preprocessors[controlnet_module]
for input_image in controlnet_input_images:
img = external_code.to_base64_nparray(input_image)
results.append(processor_module(img, res=controlnet_processor_res, thr_a=controlnet_threshold_a, thr_b=controlnet_threshold_b)[0])
global_state.cn_preprocessor_unloadable.get(controlnet_module, lambda: None)()
results64 = list(map(encode_to_base64, results))
return {"images": results64, "info": "Success"}
3. ControlNet扩展的功能实现
原始的Stable Diffusion 由三个模型构成:text encoder模型(CLIPTextModel)、UNet模型和VAE 模型。ControlNet是在UNet网络上新增的旁路,用于增加额外的条件控制Stable Diffusion的输出。
在
controlnet.py
的Script
类的process()
中,实现了网络结构的注入。process()
在图像处理前被调用,此处unet
为原先网络的结构,UnetHook
为新定义的结构,通过UnetHook.hook()
改变原始的UNet。
sd_ldm = p.sd_model
unet = sd_ldm.model.diffusion_model
......
self.latest_network = UnetHook(lowvram=hook_lowvram)
self.latest_network.hook(model=unet, sd_ldm=sd_ldm, control_params=forward_params, process=p)
self.detected_map = detected_maps
self.post_processors = post_processors
UnetHook.hook方法,model即是原先的网络,hook方法先将原先的模型的forward方法保存起来(model._original_forward = model.forward),然后给它重新赋值,赋值为自行实现的forward2。
- 文本生成图片
text2img流程
text_embedding = text_encoder(prompt)
for i in steps:
predict_noise = unet(text_embedding, timestamp,latent)
latent_new = DDPM(latent, timestamp) # 求解器
img = vae_decoder(latent)
- img2img的流程
原始的img2img
如图片卡通风格转换
img_info = vae_encoder(img)
latent_init = handle(img_info)
其他类似text2img
unet 我们可以拆开为 uencoder和udecoder。
controlnet_information = contorlnet(controlnet_img, timestamp, latent,text_embedding )
encoder_info = uencoder(timestamp, latent,text_embedding)
信息融合:
decoder_input = controlnet_information * rate + encoder_info
predict_noise = decoder(decoder_input, timestamp, latent,text_embedding )
其他流程与text2img相同
img2paint(with mask)
要梳理什么:
- controlnet的pipeline具体实现,参考:onnxweb(一个repo)的diffusion 和 diffusers 的 controlnet
需要考虑的是? - controlnet的根据参数功能和实现(我有一版本,晚点发)
网友评论