Stable Diffusion Webui源码剖析

1、关键python依赖

（1）xformers：优化加速方案。它可以对模型进行适当的优化来加速图片生成并降低显存占用。缺点是输出图像不稳定，有可能比不开Xformers略差。

（2）GFPGAN：它是腾讯开源的人脸修复算法，利用预先训练号的面部GAN（如styleGAN2）中封装的丰富多样的先验因素进行盲脸（blind face）修复，旨在开发用于现实世界人脸修复的实用算法。

（3）CLIP：Contrastive Language-Image Pre-Training，多模态方向的算法。可以训练出一个可以处理图像和文本的模型，从而使得模型可以同时理解图像和对图像的描述。

（4）OPEN-CLIP：一个开源的clip实现。

（5）Pyngrok：Ngrok工具的python实现，可以实现内网穿透

2、核心目录文件

（1）sd根目录下的repositories

存放算法源码

1）stable-diffusion-stability-ai：sd算法

2）taming-transformers：高分辨率图像合成算法

3）k-diffusion：扩散算法

4）CodeFormer：图片高清修复算法

5）BLIP：多模态算法

（2）sd根目录/models

存放模型文件

3、Gradio使用说明

【stable diffusion webui源码解析】-界面篇ui.py - 知乎

sd是基于gradio构建的，它是python库，仅需几行代码就可以构造一个html界面。

测试例子：

gr.Interface是只有左右分列的布局，它有3个输入参数：

参数1：处理函数，根据inputs中传入的组件按照顺序对应到函数的入参

参数2：组件信息

参数3：输出的数据类型

4、webui之模型处理流程

（1）cleanup_models函数move模型文件

将models目录下的文件移到相关子目录下，比如ckpt文件和safetensors文件放到Stable-diffusion子目录下。

（2）启动SD模型setup_model流程

该模型位于：/data/work/xiehao/stable-diffusion-webui/models/Stable-diffusion

主要是通过list_models函数遍历所有的模型的信息并存到checkpoint_alisases中。

第1步，查看sd/models/Stable-diffusion下是否有cpkt和safetensors结尾的文件，有则放入model_list列表中，没有则从hugginface下载模型。

第2步，通过CheckpointInfo函数检查model_list中每个模型的checkpoint信息。如果是safetensors文件，通过read_metadata_from_safetensors读取文件信息。Safetensors模型的参数都存放在json中，把键值对读出来存放到metadata字段中。

第3步，最后把每个模型根据{id : 模型对象}的键值对存放到checkpoint_alisases全局变量中。

（3）启动codeformer模型的setup_model流程

该模型位于：/data/work/xiehao/stable-diffusion-webui/models/Codeformer

主要将Codeformer初始化之后的实例放到shared.face_restorers列表中。在此过程中并没有将模型参数装载到Codeformer网络中。

（4）启动GFPGAN模型的setup_model流程

（5）遍历并加载内置的upscaler算法

这些算法位于：/data/work/xiehao/stable-diffusion-webui/modules

遍历该目录下_model.py结尾的文件，通过importlib.import_module()进行加载，这一步未看到实际作用。

初始化以下放大算法[<class 'modules.upscaler.UpscalerNone'>, <class 'modules.upscaler.UpscalerLanczos'>, <class 'modules.upscaler.UpscalerNearest'>, <class 'modules.esrgan_model.UpscalerESRGAN'>, <class 'modules.realesrgan_model.UpscalerRealESRGAN'>]，其中第1个没任何算法，第2-4是img.resize()方法实现的，第5、6个需要单独加载模型，数据都以UpscalerData格式存放，其中该对象的local_data_path存放了模型的本地地址信息。

比如：shared.sd_upscalers[5].local_data_path为：

'/data/work/xiehao/stable-diffusion-webui/models/RealESRGAN/RealESRGAN_x4plus_anime_6B.pth'

（6）加载py执行脚本load_scripts

遍历sd根目录/scripts下的py脚本以及 extensions下各扩展组件的py脚本，放到scripts_list变量中，格式如下：ScriptFile(basedir='/data/work/xiehao/stable-diffusion-webui/extensions/sd-webui-controlnet', filename='processor.py', path='/data/work/xiehao/stable-diffusion-webui/extensions/sd-webui-controlnet/scripts/processor.py')

遍历并导入scripts_list中的类型为Script或ScriptPostprocessing的py文件：

Load_module(path)加载第三方组件时可能会输出日志信息：

（7）遍历VAE模型

目前没有装任何vae模型

（8）加载模型load_model

Select_checkpoint()函数，获取sd模型信息，majicmixRealistic_v4.safetensors/majicmixRealistic_v4.safetensors [d819c8be6b]

do_inpainting_hijack函数。设置PLMSSampler的p_sample_plms。关于该方法，重建图片的反向去噪过程的每一步的图片都应用了该方法。

get_checkpoint_state_dict函数。如果是safetensors则使用safetensors.torch.load_file加载模型参数，否则使用torch.load加载模型参数。加载到pl_sd的dict类型变量中。

pl_sd字典做进一步处理：如果最外层是state_dict的key，则取该key下的value。此时pl_sd下就是模型各个节点名及对应的weights值。然后替换下面的key值：

find_checkpoint_config函数。先从模型目录下找下yaml配置文件，如果没有则执行guess_model_config_from_state_dict函数，即从模型参数中获取模型配置，最后返回/data/work/xiehao/stable-diffusion-webui/configs/v1-inference.yaml作为配置文件，信息如下：

接着用OmegaConf.load加载yaml文件，然后通过/data/work/xiehao/stable-diffusion-webui/repositories/stable-diffusion-stability-ai/ldm/util.py(82)instantiate_from_config()加载yaml信息获得model。具体步骤为：

步骤1，通过yaml的target信息，可以知道model为ldm.models.diffusion.ddpm的LatentDiffusion类。模型的源码位于：sd根目录/modules/models/diffusion/ddpm_edit.py。

步骤2，通过getattr(module的obj，class_name)获取model的类。

load_model_weights函数，将模型参数加载到模型中。通过model.load_state_dict(state_dict, strict=False)加载。因为程序参数no_half为false，所以模型量化需要从float32变为半精度tensor，half()的时候不对vae模块做处理。Vae模块为model.first_stage_model部分，所以先存到一个临时变量，half()量化完成后再赋值回去。Vae最后再单独变为float16。然后把模型放到cuda上。

Hijack函数，处理用户输入的embedding信息。假如给一个初始值，通过SD会生成未知的东西，我们通过添加额外的信息（比如prompts）让sd朝着我们想要的方向生成东西，这个就是劫持的功能，劫持是在embeddings层的。模型的embedding类为：transformers.models.clip.modeling_clip.CLIPTextEmbeddings，它的token_embeddings类为：torch.nn.modules.sparse.Embedding。

针对prompts的embedding处理类为：FrozenCLIPEmbedderWithCustomWords。约有4.9W个token。然后针对token的权重进行处理，普通单词为1.0, 中括号则除以1.1，小括号则乘以1.1.

指定优化方法apply_optimizations，通过xformers工具优化sd模型中的CrossAttention。（跨注意力机制是一种扩展自注意力机制的技术。自注意力机制是一种通过计算查询query、键key和值value之间的关联度来为输入序列中的每个元素分配权重的方法，而跨注意力机制则通过引入额外的输入序列来融合两个不同来源的信息以实现更准确的建模）。

load_textual_inversion_embeddings函数，加载根目录/embeddings下的embedding文件。加载[('/data/work/xiehao/stable-diffusion-webui/embeddings', <modules.textual_inversion.textual_inversion.DirWithTextualInversionEmbeddings object at 0x7ff2900b39d0>)]两个下的embeddings信息。比如：badhandv4、easynegative、EasyNegativeV2、ng_deepnegative_v1_75t等。

model_loaded_callback函数，遍历callback_map['callbacks_model_loaded']所有的回调函数，然后把sd_model模型传进去依次执行这些回调函数。比如/data/work/xiehao/stable-diffusion-webui/extensions/a1111-sd-webui-tagcomplete/scripts/tag_autocomplete_helper.py的get_embeddings方法，/data/work/xiehao/stable-diffusion-webui/extensions-builtin/Lora/scripts/lora_script.py的assign_lora_names_to_compvis_modules方法。