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>