From 5e462c68d06c5cae358b2f822d3360b3bf578105 Mon Sep 17 00:00:00 2001 From: claude <183307934+dxqb@users.noreply.github.com> Date: Sat, 20 Jun 2026 12:20:46 +0200 Subject: [PATCH 1/3] Fix Kohya UNet LoRA key conversion for conv_in/conv_out/time_embedding _convert_unet_lora_key() had no mapping for these three top-level UNet submodules, so Kohya-format keys touching them (e.g. lora_unet_conv_in, lora_unet_time_embed_0/2) came out as conv.in/conv.out/time.embed.0/2 instead of conv_in/conv_out/time_embedding.linear_1/2, and were reported as unexpected keys instead of being applied. --- src/diffusers/loaders/lora_conversion_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index 7c522f46a255..22505662a10a 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -278,6 +278,10 @@ def _convert_unet_lora_key(key): diffusers_name = diffusers_name.replace("proj.in", "proj_in") diffusers_name = diffusers_name.replace("proj.out", "proj_out") diffusers_name = diffusers_name.replace("emb.layers", "time_emb_proj") + diffusers_name = diffusers_name.replace("conv.in", "conv_in") + diffusers_name = diffusers_name.replace("conv.out", "conv_out") + diffusers_name = diffusers_name.replace("time.embed.0", "time_embedding.linear_1") + diffusers_name = diffusers_name.replace("time.embed.2", "time_embedding.linear_2") # SDXL specific conversions. if "emb" in diffusers_name and "time.emb.proj" not in diffusers_name: From f4656cc6e785c69a5c3bee70048af0ef7af1ac23 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 26 Jun 2026 15:10:27 +0200 Subject: [PATCH 2/3] Handle both sgm and diffusers spellings for conv/time_embedding keys The initial fix mapped conv_in/conv_out in the diffusers spelling (conv.in/ conv.out) and time_embedding in the sgm spelling (time_embed.0/.2), so neither SD1.x nor SDXL was fully covered. Add the missing spellings: - sgm conv_in/conv_out: input_blocks.0.0 / out.2 (kohya-ss SDXL sgm UNet), mapped before the block renames so input_blocks.0.0 does not become down_blocks.0.0. - diffusers time_embedding: time_embedding.linear_1/2 (kohya-ss trains SD1.x on the diffusers UNet). Verified against kohya-ss source (sdxl_original_unet.py, networks/lora.py) and the diffusers UNet module names; regression set unchanged. Co-Authored-By: Claude Opus 4.8 --- src/diffusers/loaders/lora_conversion_utils.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index 22505662a10a..b29bcd6d8741 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -263,6 +263,12 @@ def _convert_unet_lora_key(key): """ diffusers_name = key.replace("lora_unet_", "").replace("_", ".") + # kohya-ss trains SDXL on its own sgm/LDM UNet, so conv_in / conv_out arrive as + # input_blocks.0.0 / out.2. Map these before the block renames below, otherwise + # input_blocks.0.0 would become down_blocks.0.0 instead of conv_in. + diffusers_name = diffusers_name.replace("input.blocks.0.0", "conv_in") + diffusers_name = diffusers_name.replace("out.2", "conv_out") + # Replace common U-Net naming patterns. diffusers_name = diffusers_name.replace("input.blocks", "down_blocks") diffusers_name = diffusers_name.replace("down.blocks", "down_blocks") @@ -282,6 +288,11 @@ def _convert_unet_lora_key(key): diffusers_name = diffusers_name.replace("conv.out", "conv_out") diffusers_name = diffusers_name.replace("time.embed.0", "time_embedding.linear_1") diffusers_name = diffusers_name.replace("time.embed.2", "time_embedding.linear_2") + # kohya-ss trains SD 1.x on the diffusers UNet (not the sgm UNet it uses for SDXL), + # so the time-embedding MLP keeps the diffusers spelling time_embedding.linear_N + # rather than the sgm time_embed.N handled above. + diffusers_name = diffusers_name.replace("time.embedding.linear.1", "time_embedding.linear_1") + diffusers_name = diffusers_name.replace("time.embedding.linear.2", "time_embedding.linear_2") # SDXL specific conversions. if "emb" in diffusers_name and "time.emb.proj" not in diffusers_name: From c4ede93f997f4716f7f1da133e32f3f98bda53cf Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 26 Jun 2026 17:31:34 +0200 Subject: [PATCH 3/3] Map SDXL sgm label_emb LoRA keys and pass UNet top-level modules through The conv_in/conv_out/time_embedding fix only reached _convert_unet_lora_key; for the SDXL sgm UNet those keys never got there, because _maybe_map_sgm_blocks_to_diffusers treats every non-text key as a down/mid/up block. The top-level modules that live outside that block structure (time_embed, label_emb, out = conv_out, and input_blocks.0.0 = conv_in) hit the "layer not supported" raise, or crashed the inner block-index int() parse. - Pass those top-level modules through unchanged so _convert_unet_lora_key maps them, instead of block-remapping or raising. - Map the sgm label_emb (SDXL added-conditioning MLP) to diffusers add_embedding: label_emb.0.0/0.2 -> add_embedding.linear_1/2, before the SDXL index-strip heuristic that would otherwise collapse the layer index. All additions follow the kohya/sgm naming pattern and are no-ops on real kohya-ss files (which contain none of these top-level UNet LoRA keys); verified end-to-end loading a full SDXL sgm UNet LoRA into the diffusers pipeline with no unexpected/missing adapter keys. Co-Authored-By: Claude Opus 4.8 --- src/diffusers/loaders/lora_conversion_utils.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index b29bcd6d8741..ac907065ff3a 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -71,6 +71,12 @@ def _maybe_map_sgm_blocks_to_diffusers(state_dict, unet_config, delimiter="_", b for layer in all_keys: if "text" in layer: new_state_dict[layer] = state_dict.pop(layer) + elif not any(p in layer for p in sgm_patterns) or f"input_blocks{delimiter}0{delimiter}0" in layer: + # SDXL's sgm UNet has modules outside the input/middle/output block structure that + # _convert_unet_lora_key maps directly: time_embed, label_emb, out (out.2 = conv_out) + # and input_blocks.0.0 (= conv_in). Pass these through instead of block-remapping + # (conv_in's input_blocks.0 would otherwise be parsed as a down-block) or raising. + new_state_dict[layer] = state_dict.pop(layer) else: layer_id = int(layer.split(delimiter)[:block_slice_pos][-1]) if sgm_patterns[0] in layer: @@ -288,6 +294,10 @@ def _convert_unet_lora_key(key): diffusers_name = diffusers_name.replace("conv.out", "conv_out") diffusers_name = diffusers_name.replace("time.embed.0", "time_embedding.linear_1") diffusers_name = diffusers_name.replace("time.embed.2", "time_embedding.linear_2") + # sgm label_emb (SDXL added-conditioning MLP) -> diffusers add_embedding. Map before the + # SDXL index-strip heuristic below, which would otherwise collapse the layer index. + diffusers_name = diffusers_name.replace("label.emb.0.0", "add_embedding.linear_1") + diffusers_name = diffusers_name.replace("label.emb.0.2", "add_embedding.linear_2") # kohya-ss trains SD 1.x on the diffusers UNet (not the sgm UNet it uses for SDXL), # so the time-embedding MLP keeps the diffusers spelling time_embedding.linear_N # rather than the sgm time_embed.N handled above.