FLUX kohya LoRA conversion crashes on `final_layer` (KeyError on missing adaLN_modulation_1; "Incompatible keys" on final_layer alphas)

[christopher5106]

Describe the bug

The kohya / sd-scripts FLUX LoRA converter (_convert_sd_scripts_to_ai_toolkit in lora_conversion_utils.py) crashes on LoRAs that include final_layer weights, in two distinct ways. Both reproduce on diffusers==0.38.0 and on current main.

Case A — KeyError when final_layer.linear has no adaLN_modulation_1 companion.
assign_remaining_weights pops the adaLN_modulation_1 source key unconditionally, so a LoRA that trained final_layer.linear but not final_layer.adaLN_modulation.1 (a real, if uncommon, kohya export) raises KeyError.

Case B — Incompatible keys detected when final_layer carries .alpha entries.
assign_remaining_weights consumes only lora_down/lora_up, never the .alpha keys. The leftover lora_unet_final_layer_*.alpha keys then reach the remaining_keys guard, which only tolerates lora_te* prefixes, so the (otherwise valid) LoRA is rejected.

Reproduction

import torch
from diffusers.loaders.lora_pipeline import FluxLoraLoaderMixin

RANK, HID = 4, 3072

def kohya(name, out, inp, alpha=True):
    d = {f"{name}.lora_down.weight": torch.zeros(RANK, inp),
         f"{name}.lora_up.weight":   torch.zeros(out, RANK)}
    if alpha:
        d[f"{name}.alpha"] = torch.tensor(float(RANK))
    return d

def base():  # one recognized block so detection routes to the kohya converter
    return kohya("lora_unet_double_blocks_0_img_attn_proj", HID, HID)

# Case A: final_layer.linear without adaLN_modulation_1  ->  KeyError
sd = base() | kohya("lora_unet_final_layer_linear", 64, HID)
FluxLoraLoaderMixin.lora_state_dict(sd)
# KeyError: 'lora_unet_final_layer_adaLN_modulation_1.lora_down.weight'

# Case B: final_layer linear + adaLN present, alphas included  ->  Incompatible keys
sd = base() | kohya("lora_unet_final_layer_linear", 64, HID) \
            | kohya("lora_unet_final_layer_adaLN_modulation_1", 2 * HID, HID)
FluxLoraLoaderMixin.lora_state_dict(sd)
# ValueError: Incompatible keys detected:
#   lora_unet_final_layer_linear.alpha, lora_unet_final_layer_adaLN_modulation_1.alpha

Root cause

In _convert_sd_scripts_to_ai_toolkit:

• assign_remaining_weights does value = source.pop(source_key) (L554) with no fallback — Case A. It also only handles lora_down/lora_up, so final_layer .alpha keys are never consumed.
• The unconsumed alphas then hit the remaining_keys guard, which raises unless every leftover key starts with lora_te/lora_te1 — Case B.

Suggested fix

• In assign_remaining_weights, skip an assignment whose source_key is absent (source.pop(source_key, None) + continue on None) so a lone final_layer.linear still converts.
• Consume final_layer .alpha keys alongside lora_down/lora_up (or strip them before the remaining_keys guard), as the per-block _convert_to_ai_toolkit path already does.

System Info

• diffusers 0.38.0 (and main @ 7bf0000)
• transformers 5.9.0, peft 0.19.1, torch 2.11+cu128, Python 3.11, Linux
• Model: black-forest-labs/FLUX.1 family (kohya/sd-scripts LoRAs with final_layer weights)

Who can help?

@sayakpaul @BenjaminBossan

[christopher5106]

This also blocks the ComfyUI / SwarmUI full-export FLUX LoRA path on main, not just hand-written kohya LoRAs — worth noting since that's a very common export format.

Those exports carry diffusion_model.final_layer.linear and diffusion_model.final_layer.adaLN_modulation.1 (so Case A doesn't fire), but they include .alpha for every module — so Case B does. On main (0.39.0.dev0) the ComfyUI branch correctly remaps diffusion_model.* → lora_unet_* and text_encoders.clip_l.transformer.* → lora_te_*, but conversion then fails on the unconsumed final_layer alphas, which also drags the otherwise-valid lora_te_* keys into the error message:

# diffusers 0.39.0.dev0 (main), real ComfyUI FLUX LoRA export
FluxLoraLoaderMixin.lora_state_dict(comfy_state_dict)
# ValueError: Incompatible keys detected:
#   lora_unet_final_layer_adaLN_modulation_1.alpha, lora_unet_final_layer_linear.alpha,
#   lora_te_text_model_encoder_layers_0_mlp_fc1.alpha, lora_te_text_model_encoder_layers_0_mlp_fc1.lora_down.weight, ...

Dropping just the final_layer keys, the exact same export converts cleanly on main — confirming the ComfyUI mapping itself is fine and final_layer (this issue) is the sole blocker:

comfy_no_final = {k: v for k, v in comfy_state_dict.items() if "final_layer" not in k}
sd = FluxLoraLoaderMixin.lora_state_dict(comfy_no_final)
# OK -> 988 transformer.* keys + 144 text_encoder.* keys

So fixing the final_layer alpha handling here would unblock the full ComfyUI/SwarmUI export as a side effect.

[christopher5106]

A third final_layer converter issue in _convert_sd_scripts_to_ai_toolkit, distinct from the two above — flagging it here since it's the same code path.

assign_remaining_weights applies its transform (here swap_scale_shift) to both lora_A (down) and lora_B (up):

for lora_key in ["lora_A", "lora_B"]:
    ...
    if transform:
        value = transform(value)   # applied to BOTH A and B

For the final_layer.adaLN_modulation_1 → norm_out.linear mapping that's wrong: the intended converted delta is swap(B @ A), and since swap_scale_shift permutes the output axis, swap(B @ A) = swap(B) @ A — only the up tensor should be swapped. lora_A has shape (rank, hidden), so swapping it permutes the rank axis and yields swap(B) @ swap(A) ≠ swap(B @ A), silently corrupting the norm_out LoRA for any FLUX LoRA that trains the final-layer adaLN modulation (ComfyUI/XLabs full exports do).

Minimal numeric repro:

import torch
from diffusers.loaders.lora_conversion_utils import swap_scale_shift
r, h = 4, 8
A = torch.randn(r, h)      # down (rank, hidden)
B = torch.randn(2*h, r)    # up   (2*hidden, rank)
assert torch.allclose(swap_scale_shift(B @ A), swap_scale_shift(B) @ A)        # intended, up-only  -> True
assert not torch.allclose(swap_scale_shift(B @ A), swap_scale_shift(B) @ swap_scale_shift(A))  # current -> mismatch

Fix: only transform the up tensor, e.g. gate it on lora_key == "lora_B" in assign_remaining_weights (if transform and lora_key == "lora_B":).

[lcheng321]

Fixed in #14001. All three cases addressed: KeyError, Incompatible keys, and swap_scale_shift on lora_A.