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

Commit
38 days ago
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>
Author
Parents
Loading