【Bug已解决】Kohya->diffusers LoRA key converter fails for Chroma and Flux1, if LoRA is not attention-only 解决方案

一、现象长什么样

diffusers 提供convert_kohya_lora_to_diffusers把 Kohya 格式的 LoRA（常见来自 sd-scripts / Kohya_ss GUI）转成 diffusers 可直接load_lora_weights的格式。当 LoRA 是「纯注意力」时一切正常，但只要 LoRA 里包含非注意力模块（如 MLP、norm/modulation 层），转 Chroma 或 FLUX.1 权重就会失败：

from diffusers.loaders import LoraLoaderMixin LoraLoaderMixin.lora_state_dict( None, "kohya_chroma_mlp_lora.safetensors" )

报错：

KeyError： "Cannot find corresponding diffusers key for 'lora_unet_double_blocks_0_img_mlp_fc1.lora_up.weight'"

或者转换脚本静默跳过这些 key，加载后效果不对：

WARNING: skipping unmapped key lora_unet_double_blocks_0_img_modulation_lin.weight

现象总结：Kohya→diffusers 转换器只为 FLUX/Chroma 的注意力模块（to_q/k/v/out、img_attn、txt_attn）写了 key 映射正则，而 Chroma/FLUX.1 的 LoRA 可能包含mlp、modulation、norm等非注意力模块的 key，这些 key 没有映射规则，于是KeyError或被静默丢弃。

二、背景

Kohya 的 LoRA key 命名约定形如：

• lora_unet_<block>_<module>.lora_up.weight/lora_down.weight/alpha
• 对 FLUX.1，<module>可能是double_blocks_0_img_attn_proj（注意力）、也可能是double_blocks_0_img_mlp_fc1（MLP）、double_blocks_0_img_modulation_lin（调制）、single_blocks_3_modulation_lin等。

diffusers 侧的对应命名是transformer.transformer_blocks.X.attn.to_q.lora.up.weight这种嵌套路径。转换器的工作就是把 Kohya 的扁平名「翻译」成 diffusers 的嵌套名。

问题在于：原转换器只覆盖注意力相关的<module>片段，遇到mlp_fc1/modulation_lin/norm这类非注意力片段时，正则匹配不到，于是 KeyError 或 skip。Chroma 因为结构里 MLP/modulation 占比更大，纯注意力 LoRA 很少见，几乎必踩。

三、根因

根因两点：

1. key 映射正则只覆盖注意力模块：转换器内部用一个UNET_TO_DIFFUSERS映射表 + 正则，只列举了to_q/to_k/to_v/to_out/proj/linear等注意力关键词，没列mlp、modulation、norm_linear、final_layer等。
2. 没有「未知模块即报错」的兜底策略：转换器遇到匹配不上的 key，有的版本直接raise KeyError，有的版本warn + skip。两种都不好——前者阻断转换，后者悄悄丢权重导致效果错。正确做法是先把所有可能的模块家族补全进映射，再对真正未知的 key 才报错。

本质：转换器的「模块家族白名单」不全，且对非注意力模块没有统一翻译规则。

四、最小可运行复现

用标准库复现「正则只匹配注意力，mlp key 漏掉」：

import re # 原转换器的（不完整）映射：只管注意力 ATTENTION_ONLY = { r"to_q$": "to_q", r"to_k$": "to_k", r"to_v$": "to_v", r"to_out$": "to_out.0", r"img_attn_proj$": "attn.to_q", # 简化示意 } def convert_key(kohya_key: str): base = kohya_key.replace("lora_unet_", "").replace(".lora_up.weight", "") for pat, repl in ATTENTION_ONLY.items(): if re.search(pat, base): return f"transformer.transformer_blocks.{base}.lora.up.weight".replace(base, repl) raise KeyError(f"Cannot find corresponding diffusers key for '{kohya_key}'") # Chroma 的 MLP LoRA key try: convert_key("lora_unet_double_blocks_0_img_mlp_fc1.lora_up.weight") except KeyError as e: print("KeyError：", e) # 因为 mlp_fc1 不在注意力白名单

要复现「静默 skip」，把raise换成warnings.warn + return None，转换后该权重丢失。

五、解决方案（第一层：最小直接修复）

最小修复：把 MLP / modulation / norm 等非注意力模块的映射片段补进转换正则，并对未知 key 显式报错（而不是静默丢）：

import re # 补全后的模块家族映射（含非注意力） MODULE_MAP = { # 注意力 r"img_attn_proj$": "attn.to_q", r"img_attn_qkv$": "attn.to_qkv", r"txt_attn_proj$": "attn.to_q", r"to_q$": "to_q", r"to_k$": "to_k", r"to_v$": "to_v", r"to_out$": "to_out.0", # 非注意力：MLP r"img_mlp_fc1$": "ff.net.0.proj", r"img_mlp_fc2$": "ff.net.2", r"txt_mlp_fc1$": "ff.net.0.proj", r"txt_mlp_fc2$": "ff.net.2", # 非注意力：modulation / norm r"img_modulation_lin$": "norm_linear", r"txt_modulation_lin$": "norm_linear", r"modulation_lin$": "norm_linear", r"final_layer_norm_linear$": "norm_out.linear", } def convert_key(kohya_key: str): base = kohya_key.replace("lora_unet_", "").replace(".lora_up.weight", "") base = base.replace(".lora_down.weight", "").replace(".alpha", "") for pat, repl in MODULE_MAP.items(): if re.search(pat, base): diff = base[: base.rfind(re.findall(pat, base)[0])] + repl suffix = ".alpha" if kohya_key.endswith(".alpha") else ".lora.up.weight" return f"transformer.transformer_blocks.{diff}{suffix}" # 真正未知才报错，绝不静默丢 raise KeyError(f"未识别的 Kohya key（请补充映射）: '{kohya_key}'")

这一改后，Chroma/FLUX.1 的 MLP、modulation LoRA key 都能正确翻译，且未知 key 会明确报错而非悄悄丢失。

六、解决方案（第二层：结构性改进）

把「每个模型支持哪些模块家族」收敛成一个 dataclass 单一真源，转换器按需加载对应模型的映射：

from dataclasses import dataclass, field from typing import Dict, List @dataclass(frozen=True) class KohyaLoraModulePolicy: """Kohya->diffusers LoRA 转换的模块家族单一真源。""" model_family: str # 该模型可能出现的模块家族（注意力 + 非注意力） module_families: Dict[str, str] = field(default_factory=dict) # 是否允许非注意力模块（False 时遇到 mlp/modulation 直接报错） allow_non_attention: bool = True def known_suffixes(self) -> List[str]: return list(self.module_families.keys()) def translate(self, module_fragment: str) -> str: for pat, repl in self.module_families.items(): if module_fragment.endswith(pat): return repl raise KeyError(f"未识别模块片段: {module_fragment}") CHROMA_POLICY = KohyaLoraModulePolicy( model_family="chroma", allow_non_attention=True, module_families={ "img_attn_proj": "attn.to_q", "img_mlp_fc1": "ff.net.0.proj", "img_mlp_fc2": "ff.net.2", "img_modulation_lin": "norm_linear", "txt_attn_proj": "attn.to_q", "txt_mlp_fc1": "ff.net.0.proj", "txt_modulation_lin": "norm_linear", }, ) FLUX1_POLICY = KohyaLoraModulePolicy( model_family="flux1", allow_non_attention=True, module_families={ "img_attn_proj": "attn.to_q", "img_mlp_fc1": "ff.net.0.proj", "img_mlp_fc2": "ff.net.2", "modulation_lin": "norm_linear", "final_layer_norm_linear": "norm_out.linear", }, )

转换主函数convert_kohya_lora_to_diffusers(state_dict, policy)只依赖传入的policy，新增模型家族只需补一个KohyaLoraModulePolicy实例，转换器主体不动。

七、解决方案（第三层：断言 / CI 守护）

用 pytest 把「非注意力 key 可翻译 + 未知 key 报错 + 不静默丢」固化成回归：

import pytest from mylib.kohya_convert import convert_key_v2, KohyaLoraModulePolicy, CHROMA_POLICY, FLUX1_POLICY def test_mlp_key_translates(): out = convert_key_v2("lora_unet_double_blocks_0_img_mlp_fc1.lora_up.weight", CHROMA_POLICY) assert "ff.net.0.proj" in out, "Chroma MLP LoRA key 应被翻译" assert out.endswith(".lora.up.weight") def test_modulation_key_translates(): out = convert_key_v2("lora_unet_double_blocks_0_img_modulation_lin.lora_up.weight", CHROMA_POLICY) assert "norm_linear" in out def test_unknown_key_raises_not_skipped(): with pytest.raises(KeyError, match="未识别"): convert_key_v2("lora_unet_double_blocks_0_unknown_xyz.lora_up.weight", CHROMA_POLICY) def test_flux1_final_layer_key(): out = convert_key_v2("lora_unet_final_layer_norm_linear.lora_up.weight", FLUX1_POLICY) assert "norm_out.linear" in out def test_roundtrip_preserved_count(): sd = { "lora_unet_double_blocks_0_img_attn_proj.lora_up.weight": None, "lora_unet_double_blocks_0_img_mlp_fc1.lora_up.weight": None, "lora_unet_double_blocks_0_img_modulation_lin.lora_up.weight": None, } converted = {convert_key_v2(k, CHROMA_POLICY): v for k, v in sd.items()} assert len(converted) == 3, "所有 key 都应被翻译，无一遗漏"

CI 把test_mlp_key_translates与test_unknown_key_raises_not_skipped作为 LoRA 转换模块的必过项，防止非注意力模块再次漏映射或被静默丢弃。

八、排查清单

Kohya→diffusers LoRA 转换失败按顺序查：

1. 报错 key 是否含mlp/modulation/norm/final_layer？这些是常见非注意力模块，原映射常漏。
2. 转换器是否只对注意力写正则？是就用MODULE_MAP把非注意力片段补全。
3. 未知 key 是被raise还是被warn + skip？静默 skip 会导致权重丢失、效果错，必须改成显式报错。
4. 模型家族（Chroma / FLUX.1）是否各自有不同的模块命名？用KohyaLoraModulePolicy分别描述。
5. alpha标量是否一起转换？LoRA 的alphakey 也要映射到 diffusers 的alpha字段，否则缩放比例错。
6. 转换后load_lora_weights是否真的注入了这些模块？用pipe.unet.double_blocks[0].ff.net[0].proj.lora_up.weight之类确认权重落地，而非空张量。

九、小结

「Kohya→diffusers LoRA key converter fails for Chroma and Flux1, if LoRA is not attention-only」本质是转换器的模块家族白名单不全，且对非注意力模块没有统一翻译规则，遇到未知 key 要么 KeyError 要么静默丢权重。第一层把 MLP/modulation/norm 等片段补进映射正则，并改成「未知即显式报错」；第二层把每个模型的模块家族收敛到KohyaLoraModulePolicy单一真源，转换器主体与模型无关；第三层用 pytest 守住「非注意力 key 可翻译、未知 key 报错、无遗漏」。通用教训：任何 key 转换/格式迁移工具，都必须把「源 schema 的全部字段家族」列为单一真源，否则只要源多出一个字段类别，转换就崩或静默失真。