【Bug已解决】FLUX kohya LoRA conversion crashes onfinal_layer(KeyError on missing adaLN_modulation_1; Incompatible keys on final_layer alphas) 解决方案

一、现象长什么样

把 Kohya 格式的 FLUX LoRA 转成 diffusers 格式时，只要 LoRA 覆盖了 FLUX 的final_layer（输出层，含adaLN_modulation自适应层归一化），转换就会崩：

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

报错：

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

或者走了另一条路径但 alpha 对不上：

RuntimeError： Error(s) in loading state_dict for FluxTransformer2DModel: Missing key(s): "proj_out.lora_up.weight" Unexpected key(s): "final_layer.adaLN_modulation_1.lora_up.weight"

现象总结：FLUX 的final_layer包含adaLN_modulation_1/adaLN_modulation_2这类自适应层归一化权重，而 Kohya→diffusers 转换器的 key 正则只覆盖 transformer blocks 的注意力/前馈，没覆盖final_layer及其adaLN_modulation，于是要么 KeyError，要么 key 映射到了错误的proj_out导致 Incompatible keys。

二、背景

FLUX 的 transformer 在双块（double_blocks）/单块（single_blocks）之后有一个final_layer，负责把最后隐状态投影回像素空间。它的结构里有一个adaLN_modulation小 MLP，由时间步/文本条件调制，包含两个线性：adaLN_modulation_1、adaLN_modulation_2。

Kohya 对 FLUX 产 LoRA 时，会把final_layer.adaLN_modulation_1这类也纳入（尤其当训练脚本把 final_layer 标为可训练）。但 diffusers 的 FLUX 转换器在写 key 映射时，注意力只列举了double_blocks_*/single_blocks_*，漏掉了final_layer_*。于是：

• 转换器正则匹配不到final_layer_adaLN_modulation_1，走「未知 key」分支 →KeyError；
• 或者有人强行把final_layer映射到proj_out（FLUX 确实有个proj_out），但adaLN_modulation和proj_out是不同模块，映射后权重形状/语义全错 → Incompatible keys。

三、根因

根因两点：

1. 转换器 key 正则漏掉final_layer家族：adaLN_modulation_1/adaLN_modulation_2没有在UNET_TO_DIFFUSERS（FLUX 版）的映射表里，匹配不到就 KeyError。
2. final_layer与proj_out被错误混为一谈：有人把整个final_layer前缀映射到proj_out，但 FLUX 的final_layer含adaLN_modulation（调制）和linear/proj多个子模块，不能整体平移到proj_out。

本质：转换器的 FLUX 模块家族白名单漏了final_layer及其自适应调制子层，且把它的语义和proj_out混淆。

四、最小可运行复现

用标准库复现「转换器正则漏掉 final_layer」：

import re # FLUX 转换器（不完整）：只管 double/single blocks FLUX_MAP = { r"double_blocks_\d+_img_attn_proj$": "attn.to_q", r"single_blocks_\d+_linear$": "attn.to_q", } def convert_flux(kohya_key: str): base = kohya_key.replace("lora_unet_", "").replace(".lora_up.weight", "") for pat, repl in FLUX_MAP.items(): if re.search(pat, base): return f"transformer.{base}.lora.up.weight".replace(base, repl) raise KeyError(f"Cannot find corresponding diffusers key for '{kohya_key}'") try: convert_flux("lora_unet_final_layer_adaLN_modulation_1.lora_up.weight") except KeyError as e: print("KeyError：", e)

要复现「错误映射到 proj_out」：把final_layer整体替换成proj_out，转换后的 key 是transformer.proj_out.adaLN_modulation_1.lora.up.weight——而真实 FLUX 里proj_out没有adaLN_modulation子模块，加载即 Incompatible keys。

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

最小修复：在 FLUX 转换器里补上final_layer家族的正确映射（区分adaLN_modulation与linear/proj），且绝不整体平移到proj_out：

import re FLUX_FINAL_LAYER_MAP = { r"final_layer_adaLN_modulation_1$": "final_layer.adaLN_modulation.1", r"final_layer_adaLN_modulation_2$": "final_layer.adaLN_modulation.2", r"final_layer_linear$": "final_layer.linear", } def convert_flux_key(kohya_key: str): base = kohya_key.replace("lora_unet_", "") suffix = ".lora.up.weight" if kohya_key.endswith(".lora_down.weight"): suffix = ".lora.down.weight" elif kohya_key.endswith(".alpha"): suffix = ".alpha" stem = base.replace(".lora_up.weight", "").replace(".lora_down.weight", "").replace(".alpha", "") # 先试 final_layer 家族 for pat, repl in FLUX_FINAL_LAYER_MAP.items(): if re.search(pat, stem): return f"transformer.{repl}{suffix}" # 再试 double/single blocks（原有逻辑） if re.search(r"double_blocks_\d+", stem) or re.search(r"single_blocks_\d+", stem): return f"transformer.{stem}{suffix}" raise KeyError(f"未识别的 FLUX key: {kohya_key}")

这样final_layer.adaLN_modulation.1这类 key 能被正确翻译，且不会误进proj_out。

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

把「FLUX 各模块家族（含 final_layer 自适应调制）的 key 规则」收敛成一个 dataclass 单一真源：

from dataclasses import dataclass, field from typing import Dict @dataclass(frozen=True) class FluxFinalLayerLoraPolicy: """FLUX kohya LoRA 转换（含 final_layer）的单一真源。""" # final_layer 子模块：kohya 片段 -> diffusers 路径片段 final_layer_modules: Dict[str, str] = field(default_factory=lambda: { "final_layer_adaLN_modulation_1": "final_layer.adaLN_modulation.1", "final_layer_adaLN_modulation_2": "final_layer.adaLN_modulation.2", "final_layer_linear": "final_layer.linear", }) # 注意力块前缀 block_prefixes: tuple = ("double_blocks", "single_blocks") # 明确 NOT 映射到 proj_out 的模块（防止混淆） forbidden_proj_out_mapping: tuple = ("final_layer_adaLN_modulation",) def convert(self, kohya_key: str) -> str: stem = kohya_key.replace("lora_unet_", "") suffix = self._suffix(kohya_key) stem = stem.replace(".lora_up.weight", "").replace(".lora_down.weight", "").replace(".alpha", "") for frag, path in self.final_layer_modules.items(): if stem == frag or stem.startswith(frag): return f"transformer.{path}{suffix}" if any(stem.startswith(p + "_") for p in self.block_prefixes): return f"transformer.{stem}{suffix}" raise KeyError(f"未识别的 FLUX key: {kohya_key}") def _suffix(self, kohya_key: str) -> str: if kohya_key.endswith(".lora_down.weight"): return ".lora.down.weight" if kohya_key.endswith(".alpha"): return ".alpha" return ".lora.up.weight" def validate_no_proj_out_confusion(self, diffusers_key: str) -> bool: if "proj_out" in diffusers_key: return not any(f in diffusers_key for f in self.forbidden_proj_out_mapping) return True

转换主函数只调用policy.convert，且转换后跑validate_no_proj_out_confusion防回退到错误映射。

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

用 pytest 把「final_layer 可翻译 + 不混 proj_out + 未知 key 报错」固化成回归：

import pytest from mylib.flux_final_layer import FluxFinalLayerLoraPolicy POLICY = FluxFinalLayerLoraPolicy() def test_adaln_modulation_1_translates(): out = POLICY.convert("lora_unet_final_layer_adaLN_modulation_1.lora_up.weight") assert out == "transformer.final_layer.adaLN_modulation.1.lora.up.weight" def test_adaln_modulation_2_translates(): out = POLICY.convert("lora_unet_final_layer_adaLN_modulation_2.lora_down.weight") assert out.endswith("final_layer.adaLN_modulation.2.lora.down.weight") def test_linear_translates(): out = POLICY.convert("lora_unet_final_layer_linear.lora_up.weight") assert out.endswith("final_layer.linear.lora.up.weight") def test_not_mapped_to_proj_out(): key = POLICY.convert("lora_unet_final_layer_adaLN_modulation_1.lora_up.weight") assert POLICY.validate_no_proj_out_confusion(key) is True assert "proj_out" not in key def test_unknown_key_raises(): with pytest.raises(KeyError, match="未识别"): POLICY.convert("lora_unet_final_layer_unknown.lora_up.weight") def test_block_keys_still_work(): out = POLICY.convert("lora_unet_double_blocks_0_img_attn_proj.lora_up.weight") assert "double_blocks_0" in out

CI 把test_adaln_modulation_1_translates与test_not_mapped_to_proj_out作为 FLUX LoRA 转换的必过项，要求「任何 final_layer key 必须正确映射且不得进 proj_out」。

八、排查清单

FLUX kohya LoRA 转 final_layer 失败按顺序查：

1. 报错 key 是否含final_layer/adaLN_modulation？这些是转换器常漏的 final_layer 家族。
2. 是否被错误映射到了proj_out？adaLN_modulation和proj_out是不同模块，形状/语义都不该混。
3. 转换正则是否只列了double/single_blocks？是就补final_layer_modules映射。
4. adaLN_modulation_1/_2是否分别映射？Kohya 用数字后缀，diffusers 是adaLN_modulation.1/.2。
5. alpha标量是否一起转换？final_layer 的 alpha 也要映射到 diffusers 的 alpha 字段。
6. 转换后load_lora_weights是否真的注入了transformer.final_layer.*？确认权重落地而非空张量。

九、小结

「FLUX kohya LoRA conversion crashes on final_layer」本质是转换器的 FLUX 模块家族白名单漏了 final_layer 及其 adaLN_modulation 自适应调制子层，且把它和 proj_out 错误混淆，导致 KeyError 或 Incompatible keys。第一层补上 final_layer 家族的正确映射并明确禁止平移到 proj_out；第二层把 FLUX 各模块（含 final_layer）的 key 规则收敛到FluxFinalLayerLoraPolicy单一真源；第三层用 pytest 守住「final_layer 可翻译、不混 proj_out、未知 key 报错」。通用教训：**key 转换工具必须覆盖目标模型的所有叶子模块（包括输出层的自适应调制），且不能把结构不同的模块强行合并映射，否则要么 KeyError 要么静默错配。