将LoRA权重合并至Stable Diffusion原始模型的映射问题求助

问题梳理 #

作为Stable Diffusion新手，尝试手动合并LoRA权重到原始模型，文本编码器部分的代码如下：

for i in range(0, 12): 
 model_model[f"cond_stage_model.transformer.text_model.encoder.layers.{i}.mlp.fc1.weight"] +=
 model_lora[f"lora_te_text_model_encoder_layers_{i}_mlp_fc1.alpha"] *
torch.matmul( 
model_lora[f"lora_te_text_model_encoder_layers_{i}_mlp_fc1.lora_up.weight"].float(), 
model_lora[f"lora_te_text_model_encoder_layers_{i}_mlp_fc1.lora_down.weight"].float()
)

但遇到Unet部分的LoRA键无法匹配原始模型权重的问题，LoRA键示例：

lora_unet_down_blocks_0_attentions_0_proj_in.alpha
lora_unet_down_blocks_0_attentions_0_proj_in.lora_down.weight
lora_unet_down_blocks_0_attentions_0_proj_in.lora_up.weight

找到的原始模型中最接近的键为：

model.diffusion_model.input_blocks.0.0.weight
model.diffusion_model.input_blocks.1.1.proj_in.weight

对应关系与合并方案 #

1. Unet结构的命名差异 #

LoRA的Unet权重命名基于down_blocks/up_blocks/mid_block的拆分方式，而原始Stable Diffusion的Unet用的是input_blocks/middle_block/output_blocks，两者的对应关系需要明确：

• LoRA的down_blocks_0对应原始模型的input_blocks[1]（input_blocks[0]是无注意力的纯卷积输入层，没有对应的LoRA注意力模块）
• LoRA的down_blocks_0_attentions_0对应原始模型input_blocks[1][1]（input_blocks的每个带注意力的block结构为[卷积层, [归一化层, 注意力层]]，注意力层是第2个元素的第2个组件）

2. 正确的Unet LoRA合并代码示例 #

针对你给出的LoRA键，对应的合并代码如下：

# 处理lora_unet_down_blocks_0_attentions_0_proj_in
lora_key_prefix = "lora_unet_down_blocks_0_attentions_0_proj_in"
alpha = model_lora[f"{lora_key_prefix}.alpha"]
lora_up = model_lora[f"{lora_key_prefix}.lora_up.weight"].float()
lora_down = model_lora[f"{lora_key_prefix}.lora_down.weight"].float()

# 对应的原始模型权重键
original_key = "model.diffusion_model.input_blocks.1.1.proj_in.weight"
# 计算LoRA权重并合并
model_model[original_key] += alpha * torch.matmul(lora_up, lora_down)

3. 通用映射规则（扩展到其他Unet LoRA键） #

如果要批量处理所有Unet LoRA键，可以参考以下映射逻辑：

• LoRA的down_blocks_N_attentions_M_proj_X → 原始模型input_blocks[N+1][1].proj_X.weight（N从0开始，因为input_blocks[0]无注意力）
• LoRA的up_blocks_N_attentions_M_proj_X → 原始模型output_blocks[-(N+2)][1].proj_X.weight（up_blocks对应output_blocks的反向顺序）
• LoRA的mid_block_attentions_M_proj_X → 原始模型middle_block[1].proj_X.weight

4. 文本编码器合并的补充 #

你的文本编码器代码只处理了mlp.fc1层，完整的LoRA合并还需要覆盖mlp.fc2、self_attn.q_proj、self_attn.v_proj等层，否则会丢失部分LoRA效果，示例：

for i in range(0, 12):
    # 处理MLP fc1
    model_model[f"cond_stage_model.transformer.text_model.encoder.layers.{i}.mlp.fc1.weight"] += \
        model_lora[f"lora_te_text_model_encoder_layers_{i}_mlp_fc1.alpha"] * \
        torch.matmul(model_lora[f"lora_te_text_model_encoder_layers_{i}_mlp_fc1.lora_up.weight"].float(),
                    model_lora[f"lora_te_text_model_encoder_layers_{i}_mlp_fc1.lora_down.weight"].float())
    # 处理MLP fc2
    model_model[f"cond_stage_model.transformer.text_model.encoder.layers.{i}.mlp.fc2.weight"] += \
        model_lora[f"lora_te_text_model_encoder_layers_{i}_mlp_fc2.alpha"] * \
        torch.matmul(model_lora[f"lora_te_text_model_encoder_layers_{i}_mlp_fc2.lora_up.weight"].float(),
                    model_lora[f"lora_te_text_model_encoder_layers_{i}_mlp_fc2.lora_down.weight"].float())
    # 处理注意力Q层
    model_model[f"cond_stage_model.transformer.text_model.encoder.layers.{i}.self_attn.q_proj.weight"] += \
        model_lora[f"lora_te_text_model_encoder_layers_{i}_self_attn_q_proj.alpha"] * \
        torch.matmul(model_lora[f"lora_te_text_model_encoder_layers_{i}_self_attn_q_proj.lora_up.weight"].float(),
                    model_lora[f"lora_te_text_model_encoder_layers_{i}_self_attn_q_proj.lora_down.weight"].float())

内容的提问来源于stack exchange，提问作者Gaurish Trivedi