This script provides a function to convert the key format of LoRA weights from Xlabs to the format used by Hugging Face's Diffusers library. The conversion is necessary for utilizing Xlabs LoRA weights in models compatible with the Diffusers ecosystem. The script reads a LoRA weight file in safetensors format, converts the keys using predefined …

convert_lora_keys.py

# This script provides a function to convert the key format of LoRA weights from Xlabs to the format used by 
# Hugging Face's Diffusers library. The conversion is necessary for utilizing Xlabs LoRA weights in models 
# compatible with the Diffusers ecosystem. The script reads a LoRA weight file in safetensors format, 
# converts the keys using predefined rules, and outputs the converted state dictionary.

import re
import torch
import safetensors.torch
from diffusers.loaders.single_file_utils import swap_scale_shift

def convert_lora_keys(old_state_dict):
    """
    Converts LoRA keys from Xlabs format to Hugging Face Diffusers format.
    
    Args:
        old_state_dict (dict): The original state dictionary with Xlabs keys.
    
    Returns:
        dict: A new state dictionary with converted keys suitable for Diffusers.
    """
    new_state_dict = {}
    
    for old_key, value in old_state_dict.items():
        # Handle double_blocks
        if old_key.startswith('double_blocks'):
            block_num = re.search(r'double_blocks\.(\d+)', old_key).group(1)
            new_key = f'transformer.transformer_blocks.{block_num}'
            
            if 'processor.proj_lora1' in old_key:
                new_key += '.attn.to_out.0'
            elif 'processor.proj_lora2' in old_key:
                new_key += '.attn.to_add_out'
            elif 'processor.qkv_lora1' in old_key:
                new_key += '.attn'
                if 'down' in old_key:
                    for proj in ['add_q_proj', 'add_k_proj', 'add_v_proj']:
                        proj_key = f'{new_key}.{proj}.lora_A.weight'
                        new_state_dict[proj_key] = value
                elif 'up' in old_key:
                    for proj in ['add_q_proj', 'add_k_proj', 'add_v_proj']:
                        proj_key = f'{new_key}.{proj}.lora_B.weight'
                        new_state_dict[proj_key] = value
                continue
            elif 'processor.qkv_lora2' in old_key:
                new_key += '.attn'
                if 'down' in old_key:
                    for proj in ['to_q', 'to_k', 'to_v']:
                        proj_key = f'{new_key}.{proj}.lora_A.weight'
                        new_state_dict[proj_key] = value
                elif 'up' in old_key:
                    for proj in ['to_q', 'to_k', 'to_v']:
                        proj_key = f'{new_key}.{proj}.lora_B.weight'
                        new_state_dict[proj_key] = value
                continue
            
            if 'down' in old_key:
                new_key += '.lora_A.weight'
            elif 'up' in old_key:
                new_key += '.lora_B.weight'
            
        # Handle single_blocks
        elif old_key.startswith('single_blocks'):
            block_num = re.search(r'single_blocks\.(\d+)', old_key).group(1)
            new_key = f'transformer.single_transformer_blocks.{block_num}'
            
            if 'proj_lora1' in old_key or 'proj_lora2' in old_key:
                new_key += '.proj_out'
            elif 'qkv_lora1' in old_key or 'qkv_lora2' in old_key:
                new_key += '.norm.linear'
            
            if 'down' in old_key:
                new_key += '.lora_A.weight'
            elif 'up' in old_key:
                new_key += '.lora_B.weight'
        
        else:
            # Handle other potential key patterns here
            new_key = old_key
        
        new_state_dict[new_key] = value
    
    return new_state_dict

# Load the original LoRA weights from a safetensors file
old_state_dict = safetensors.torch.load_file("realism_lora.safetensors")

# Print the original state dictionary keys and their shapes
for key, value in old_state_dict.items():
    print(f"{key} {list(value.shape)} {value.dtype}")

# Convert the keys to the Diffusers format
new_state_dict = convert_lora_keys(old_state_dict)

# Print the converted state dictionary keys and their shapes
for key, value in new_state_dict.items():
    print(f"{key} {list(value.shape)} {value.dtype}")

# Print the length of both dictionaries to ensure all keys were converted
print(len(old_state_dict), len(new_state_dict))

Sayak here from the diffusers team 👋

Thanks for working on this. I think you would also need to account for the fact that our implementation in diffusers doesn't do QKV projections all at once but rather individually. https://github.com/kohya-ss/sd-scripts/blob/sd3/networks/convert_flux_lora.py provides a good reference.

Oh, alright. Thank you for that insight.

Do you think I should still work on this since I see it has already been implemented?