Fix Kohya UNet LoRA key conversion for conv_in/conv_out/time_embedding (#14006)

* 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.

* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

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:
@@ -263,6 +269,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")
@@ -278,6 +290,19 @@ 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")
+    # 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.
+    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: