Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sayakpaul merged 4 commits into from
Jun 28, 2026
Merged

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

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5e462c6
Fix Kohya UNet LoRA key conversion for conv_in/conv_out/time_embedding
dxqb Jun 20, 2026
f4656cc
Handle both sgm and diffusers spellings for conv/time_embedding keys
dxqb Jun 26, 2026
c4ede93
Map SDXL sgm label_emb LoRA keys and pass UNet top-level modules through
dxqb Jun 26, 2026
85ec21d
Merge branch 'main' into fix-unet-lora-conv-time-embedding
sayakpaul Jun 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions src/diffusers/loaders/lora_conversion_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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")
Expand All @@ -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:
Expand Down
Loading