Skip to content

Fix LoKr key mapping for AI-Toolkit Chroma/Flux models #5

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

Open
Valheol wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix LoKr key mapping for AI-Toolkit Chroma/Flux models #5

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
69 changes: 59 additions & 10 deletions module/load_lora.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,14 +26,63 @@
def standardize_lora_key_format(lora_sd):
new_sd = {}
for k, v in lora_sd.items():
# aitoolkit/lycoris format
if k.startswith("lycoris_blocks_"):
# AI-Toolkit Chroma/Flux format for double_blocks and single_blocks
# Convert: lycoris_double_blocks_0_img_attn_proj.lokr_w1
# To: diffusion_model.double_blocks.0.img_attn.proj.lokr_w1
if k.startswith("lycoris_double_blocks_") or k.startswith("lycoris_single_blocks_"):
# Split on the first dot to separate base key from LoKr component
parts = k.split('.', 1)
base_key = parts[0] # e.g., lycoris_double_blocks_0_img_attn_proj
suffix = '.' + parts[1] if len(parts) > 1 else '' # e.g., .lokr_w1 or .alpha

# Remove lycoris_ prefix
base_key = base_key.replace('lycoris_', '')

# Split into components: ['double', 'blocks', '0', 'img', 'attn', 'proj']
components = base_key.split('_')

# Reconstruct the key
if components[0] in ['double', 'single'] and len(components) > 1 and components[1] == 'blocks':
block_type = components[0] + '_blocks' # double_blocks or single_blocks
block_idx = components[2] if len(components) > 2 else '0'

# Build the remaining path by processing components individually
if len(components) > 3:
# Build path intelligently based on known patterns
path_parts = components[3:]

# Check for attention patterns: img/txt + attn + proj/qkv
if len(path_parts) >= 3 and path_parts[1] == 'attn' and path_parts[0] in ['img', 'txt']:
# e.g., ['img', 'attn', 'proj'] -> 'img_attn.proj'
remaining = f"{path_parts[0]}_attn.{'.'.join(path_parts[2:])}"
# Check for MLP patterns: img/txt + mlp + number
elif len(path_parts) >= 3 and path_parts[1] == 'mlp' and path_parts[0] in ['img', 'txt']:
# e.g., ['img', 'mlp', '0'] -> 'img_mlp.0'
remaining = f"{path_parts[0]}_mlp.{'.'.join(path_parts[2:])}"
# Check for norm patterns: img/txt + norm1/norm2
elif len(path_parts) >= 2 and path_parts[0] in ['img', 'txt'] and 'norm' in path_parts[1]:
# e.g., ['img', 'norm1'] -> 'img_norm1'
remaining = '_'.join(path_parts)
# Single block linear layers: linear1, linear2
elif len(path_parts) == 1:
remaining = path_parts[0]
else:
# Fallback: join with dots
remaining = '.'.join(path_parts)
else:
remaining = ''

k = f"diffusion_model.{block_type}.{block_idx}.{remaining}{suffix}"

# aitoolkit/lycoris format (for other models)
elif k.startswith("lycoris_blocks_"):
k = k.replace("lycoris_blocks_", "blocks.")
k = k.replace("_cross_attn_", ".cross_attn.")
k = k.replace("_self_attn_", ".self_attn.")
k = k.replace("_ffn_net_0_proj", ".ffn.0")
k = k.replace("_ffn_net_2", ".ffn.2")
k = k.replace("to_out_0", "o")

# Diffusers format
if k.startswith('transformer.'):
k = k.replace('transformer.', 'diffusion_model.')
Expand Down Expand Up @@ -172,14 +221,14 @@ def standardize_lora_key_format(lora_sd):
k = k.replace('.to_v.', '.v.')
k = k.replace('.to_out.0.', '.o.')

if "img_attn.proj" in k:
k = k.replace("img_attn.proj", "img_attn_proj")
if "img_attn.qkv" in k:
k = k.replace("img_attn.qkv", "img_attn_qkv")
if "txt_attn.proj" in k:
k = k.replace("txt_attn.proj", "txt_attn_proj")
if "txt_attn.qkv" in k:
k = k.replace("txt_attn.qkv", "txt_attn_qkv")
# NOTE: Removed img_attn/txt_attn dot-to-underscore conversions (lines 224-231)
# These were converting the CORRECT format (img_attn.proj) to INCORRECT format (img_attn_proj)
# for Chroma/Flux models. The key mapping in model_lora_keys_unet() line 371:
# key_map[k[:-len(".weight")]] = k
# already handles dot notation correctly, mapping:
# "diffusion_model.double_blocks.0.img_attn.proj" -> "diffusion_model.double_blocks.0.img_attn.proj.weight"
# So we keep the dot notation that our Chroma conversion creates.

new_sd[k] = v
return new_sd

Expand Down