图像生成：SD lora加载代码详解与实现

文章目录

前言
一、SD模型介绍
二、模型加载
- 1. 模型架构加载
- 2. safetensors权重加载
- 3. lora权重加载
三、Name匹配
四、权重融合
- 1、构建net类
- 2、匹配lora weight和model weight
- 3、基于lora权重创建lora模块
- 4、权重融合
五、整体pipeline
总结

前言

SD中lora的加载相信都不陌生，但是大家大多数都是利用SD webUI加载lora，本文主要梳理一下SD webUI中lora加载的代码逻辑。关于lora的原理，可以参考我之前的博客——图像生成：SD LoRA模型详解

一、SD模型介绍

SD model结构一般分为几个部分，如下：
在这里插入图片描述

SD webui使用pytorch lightning搭建，了解pl的同学可能知道，模型的相关配置一般都写在yaml文件中，因此其实可以根据yaml文件来判断模型的基本结构，类似如下：

      unet_config:target: ldm.modules.diffusionmodules.openaimodel.UNetModelparams:image_size: 32 # unusedin_channels: 4out_channels: 4model_channels: 320attention_resolutions: [ 4, 2, 1 ]num_res_blocks: 2channel_mult: [ 1, 2, 4, 4 ]num_heads: 8use_spatial_transformer: Truetransformer_depth: 1context_dim: 768use_checkpoint: Truelegacy: Falsefirst_stage_config:target: ldm.models.autoencoder.AutoencoderKLparams:embed_dim: 4monitor: val/rec_lossddconfig:double_z: truez_channels: 4resolution: 256in_channels: 3out_ch: 3ch: 128ch_mult:- 1- 2- 4- 4num_res_blocks: 2attn_resolutions: []dropout: 0.0lossconfig:target: torch.nn.Identitycond_stage_config:target: ldm.modules.encoders.modules.FrozenCLIPEmbedder

二、模型加载

因为整个模型分为VAE，CLIP，Unet等多个部分，lora模型一般是对这几个部分进行权重修改，有的可能只修改了Unet，有的可能多个部分都修改了，这个是由当时lora的训练师冻结的模块决定的。本文接下来主要以Unet部分的加载和lora修改为例来介绍，其它模块类似。

1. 模型架构加载

首先根据yaml文件加载模型架构（这里我只保留了Unet的配置文件，然后进行加载）：

unet_config_path = '/data/wangyx/工程处理/sd/unet.yaml'
diff_model_config = OmegaConf.load(unet_config_path)  
unet_config = diff_model_config.model.unet_config
diffusion_model = instantiate_from_config(unet_config)

2. safetensors权重加载

这里以比较火的chilloutmix作为example 。( ⁼̴̀ .̫ ⁼̴ )✧

#可视化权重结构
import torch
from safetensors.torch import load_file, save_file
from safetensors import safe_open
model_path = '/stable-diffusion-webui/stable-diffusion-webui/models/Stable-diffusion/chilloutmix_NiPrunedFp32Fix.safetensors'
tensors = {}
with safe_open(model_path, framework="pt", device='cpu') as f:for k in f.keys():tensors[k] = f.get_tensor(k)

获取权重后即可把对应权重文件加载到模型结构中。

3. lora权重加载

lora的权重加载直接用safetensor加载即可

lora_path = '**/sdxl_lcm_lora.safetensors'
pl_sd = safetensors.torch.load_file(lora_path)

三、Name匹配

打印出lora权重名字和模型每一层的名字后其实可以发现，其实都是不对应的，因此需要手动将他们匹配起来，在SD webUI中的Lora中部分使用下面这样一个函数完成权重匹配

def convert_diffusers_name_to_compvis(key, is_sd2):def match(match_list, regex_text):regex = re_compiled.get(regex_text)if regex is None:regex = re.compile(regex_text)re_compiled[regex_text] = regexr = re.match(regex, key)if not r:return Falsematch_list.clear()match_list.extend([int(x) if re.match(re_digits, x) else x for x in r.groups()])return Truem = []if match(m, r"lora_unet_conv_in(.*)"):return f'diffusion_model_input_blocks_0_0{m[0]}'if match(m, r"lora_unet_conv_out(.*)"):return f'diffusion_model_out_2{m[0]}'if match(m, r"lora_unet_time_embedding_linear_(\d+)(.*)"):return f"diffusion_model_time_embed_{m[0] * 2 - 2}{m[1]}"if match(m, r"lora_unet_down_blocks_(\d+)_(attentions|resnets)_(\d+)_(.+)"):suffix = suffix_conversion.get(m[1], {}).get(m[3], m[3])return f"diffusion_model_input_blocks_{1 + m[0] * 3 + m[2]}_{1 if m[1] == 'attentions' else 0}_{suffix}"if match(m, r"lora_unet_mid_block_(attentions|resnets)_(\d+)_(.+)"):suffix = suffix_conversion.get(m[0], {}).get(m[2], m[2])return f"diffusion_model_middle_block_{1 if m[0] == 'attentions' else m[1] * 2}_{suffix}"if match(m, r"lora_unet_up_blocks_(\d+)_(attentions|resnets)_(\d+)_(.+)"):suffix = suffix_conversion.get(m[1], {}).get(m[3], m[3])return f"diffusion_model_output_blocks_{m[0] * 3 + m[2]}_{1 if m[1] == 'attentions' else 0}_{suffix}"if match(m, r"lora_unet_down_blocks_(\d+)_downsamplers_0_conv"):return f"diffusion_model_input_blocks_{3 + m[0] * 3}_0_op"if match(m, r"lora_unet_up_blocks_(\d+)_upsamplers_0_conv"):return f"diffusion_model_output_blocks_{2 + m[0] * 3}_{2 if m[0]>0 else 1}_conv"if match(m, r"lora_te_text_model_encoder_layers_(\d+)_(.+)"):if is_sd2:if 'mlp_fc1' in m[1]:return f"model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc1', 'mlp_c_fc')}"elif 'mlp_fc2' in m[1]:return f"model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc2', 'mlp_c_proj')}"else:return f"model_transformer_resblocks_{m[0]}_{m[1].replace('self_attn', 'attn')}"return f"transformer_text_model_encoder_layers_{m[0]}_{m[1]}"if match(m, r"lora_te2_text_model_encoder_layers_(\d+)_(.+)"):if 'mlp_fc1' in m[1]:return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc1', 'mlp_c_fc')}"elif 'mlp_fc2' in m[1]:return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc2', 'mlp_c_proj')}"else:return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('self_attn', 'attn')}"return keykey1 = 'lora_unet_down_blocks_0_downsamplers_0_conv'  #.alpha
new_key1 = convert_diffusers_name_to_compvis(key1, is_sd2=True)
print(new_key1)

通过这样一个函数即可把lora权重的名字替换成和模型名字一样的。

另外同时对SD原模型的层名字进行修改，并存在一个字典中：

def assign_network_names_to_compvis_modules(sd_model):network_layer_mapping = {}for name, module in sd_model.named_modules():network_name = name.replace(".", "_")network_layer_mapping[network_name] = modulemodule.network_layer_name = network_namesd_model.network_layer_mapping = network_layer_mapping

这样子lora和模型的名字就完成对应了。

四、权重融合

完成权重匹配后，就可以进行权重融合了，这里我将SD wenUI中的代码摘了一部分出来进行实现，从而更好理解原理。

1、构建net类

这里我调用了webUI中的net类，用于后续赋予相关属性

network_on_disk = NetworkOnDisk('name', '.pth')
net = Network('name', network_on_disk)
#这部分是建立一个空壳，用于后续操作

2、匹配lora weight和model weight

#创建nametuple保存权重
NetworkWeights = namedtuple('NetworkWeights', ['network_key', 'sd_key', 'w', 'sd_module'])
matched_networks = {}for key_network, weight in pl_sd.items():  #循环lora的每一项key_network_without_network_parts, network_part = key_network.split(".", 1)#如果是SDXL，那么isSD2需要选择为Truefkey = convert_diffusers_name_to_compvis(key_network_without_network_parts, True) key = fkey[16:]'''正常模型架构是model.diffusion_model, model.first_stage_model但是现在我们只加载了unet部分所以先去掉前半部分'''sd_module = diffusion_model.network_layer_mapping.get(key, None)#获取修改名字后对应的moduleif key not in matched_networks and sd_module is not None:  matched_networks[key] = NetworkWeights(network_key=key_network, sd_key=key, w={}, sd_module=sd_module)if sd_module is not None:matched_networks[key].w[network_part] = weight

通过上述代码就可以将匹配的Unet层权重和lora权重放在一个NetworkWeights组了，用于后续融合。代码中w存的是lora的权重，sd module存的是对应unet中的结构和权重。

3、基于lora权重创建lora模块

上面构建了matched_networks字典，每个key下对应了一组匹配好的lora权重和模型模块，接下来就是基于lora权重创建lora模块了，我们第一步创建了一个net空壳类，在这里赋予net.modules属性，并将创建好的lora模块赋予该属性。

for key, weights in matched_networks.items():print(key)print(weights)net_module = Nonefor nettype in module_types:net_module = nettype.create_module(net, weights)if net_module is not None:breaknet.modules[key] = net_module

这里的create_module来自于源码中的Lora/network_lora.py, 该函数主要是创建结构，并赋予权重，最后构建一个完整的lora模块。

4、权重融合

接下来就可以完成权重融合了，这里我以某一层为例

network_layer_name = 'output_blocks_11_1_transformer_blocks_0_attn1_to_k'
module = net.modules.get(network_layer_name, None)
print('get模块', module)
fb = matched_networks[network_layer_name].sd_module  
fb_weight = fb.weight
with torch.no_grad():updown, ex_bias = module.calc_updown(fb_weight) if len(fb_weight) == 4 and fb_weight.shape[1] == 9:# inpainting model. zero pad updown to make channel[1]  4 to 9updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5))fb_weight += updown