Create convert_simpletuner_lora.py

lora/convert_simpletuner_lora.py

@@ -0,0 +1,483 @@
+#!/usr/bin/env python3
+"""
+Convert SimpleTuner LoRA weights to diffusers-compatible format for AuraFlow.
+
+This script converts LoRA weights saved by SimpleTuner into a format that can be
+directly loaded by diffusers' load_lora_weights() method.
+
+Usage:
+    python convert_simpletuner_lora.py <input_lora.safetensors> <output_lora.safetensors>
+
+Example:
+    python convert_simpletuner_lora.py input_lora.safetensors diffusers_compatible_lora.safetensors
+"""
+
+import argparse
+import sys
+from pathlib import Path
+from typing import Dict
+
+import safetensors.torch
+import torch
+
+
+def detect_lora_format(state_dict: Dict[str, torch.Tensor]) -> str:
+    """
+    Detect the format of the LoRA state dict.
+    
+    Returns:
+        "peft" if already in PEFT/diffusers format
+        "mixed" if mixed format (some lora_A/B, some lora.down/up)
+        "simpletuner_transformer" if in SimpleTuner format with transformer prefix
+        "simpletuner_auraflow" if in SimpleTuner AuraFlow format
+        "kohya" if in Kohya format
+        "unknown" otherwise
+    """
+    keys = list(state_dict.keys())
+    
+    # Check the actual weight naming convention (lora_A/lora_B vs lora_down/lora_up)
+    has_lora_a_b = any((".lora_A." in k or ".lora_B." in k) for k in keys)
+    has_lora_down_up = any((".lora_down." in k or ".lora_up." in k) for k in keys)
+    has_lora_dot_down_up = any((".lora.down." in k or ".lora.up." in k) for k in keys)
+    
+    # Check prefixes
+    has_transformer_prefix = any(k.startswith("transformer.") for k in keys)
+    has_lora_transformer_prefix = any(k.startswith("lora_transformer_") for k in keys)
+    has_lora_unet_prefix = any(k.startswith("lora_unet_") for k in keys)
+    
+    # Mixed format: has both lora_A/B AND lora.down/up (SimpleTuner hybrid)
+    if has_transformer_prefix and has_lora_a_b and (has_lora_down_up or has_lora_dot_down_up):
+        return "mixed"
+    
+    # Pure PEFT format: transformer.* with ONLY lora_A/lora_B
+    if has_transformer_prefix and has_lora_a_b and not has_lora_down_up and not has_lora_dot_down_up:
+        return "peft"
+    
+    # SimpleTuner with transformer prefix but old naming: transformer.* with lora_down/lora_up
+    if has_transformer_prefix and (has_lora_down_up or has_lora_dot_down_up):
+        return "simpletuner_transformer"
+    
+    # SimpleTuner AuraFlow format: lora_transformer_* with lora_down/lora_up
+    if has_lora_transformer_prefix and has_lora_down_up:
+        return "simpletuner_auraflow"
+    
+    # Traditional Kohya format: lora_unet_* with lora_down/lora_up
+    if has_lora_unet_prefix and has_lora_down_up:
+        return "kohya"
+    
+    return "unknown"
+
+
+def convert_mixed_lora_to_diffusers(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """
+    Convert mixed LoRA format to pure PEFT format.
+    
+    SimpleTuner sometimes saves a hybrid format where some layers use lora_A/lora_B
+    and others use .lora.down./.lora.up. This converts all to lora_A/lora_B.
+    """
+    new_state_dict = {}
+    converted_count = 0
+    kept_count = 0
+    skipped_count = 0
+    renames = []
+    
+    # Get all keys
+    all_keys = sorted(state_dict.keys())
+    
+    print("\nProcessing keys:")
+    print("-" * 80)
+    
+    for key in all_keys:
+        # Already in correct format (lora_A or lora_B)
+        if ".lora_A." in key or ".lora_B." in key:
+            new_state_dict[key] = state_dict[key]
+            kept_count += 1
+        
+        # Needs conversion: .lora.down. -> .lora_A.
+        elif ".lora.down.weight" in key:
+            new_key = key.replace(".lora.down.weight", ".lora_A.weight")
+            new_state_dict[new_key] = state_dict[key]
+            renames.append((key, new_key))
+            converted_count += 1
+        
+        # Needs conversion: .lora.up. -> .lora_B.
+        elif ".lora.up.weight" in key:
+            new_key = key.replace(".lora.up.weight", ".lora_B.weight")
+            new_state_dict[new_key] = state_dict[key]
+            renames.append((key, new_key))
+            converted_count += 1
+        
+        # Skip alpha keys (not used in PEFT format)
+        elif ".alpha" in key:
+            skipped_count += 1
+            continue
+        
+        # Other keys (shouldn't happen, but keep them just in case)
+        else:
+            new_state_dict[key] = state_dict[key]
+            print(f"⚠ Warning: Unexpected key format: {key}")
+    
+    print(f"\nSummary:")
+    print(f"  ✓ Kept {kept_count} keys already in correct format (lora_A/lora_B)")
+    print(f"  ✓ Converted {converted_count} keys from .lora.down/.lora.up to lora_A/lora_B")
+    print(f"  ✓ Skipped {skipped_count} alpha keys")
+    
+    if renames:
+        print(f"\nRenames applied ({len(renames)} conversions):")
+        print("-" * 80)
+        for old_key, new_key in renames:
+            # Show the difference more clearly
+            if ".lora.down.weight" in old_key:
+                layer = old_key.replace(".lora.down.weight", "")
+                print(f"  {layer}")
+                print(f"    .lora.down.weight  →  .lora_A.weight")
+            elif ".lora.up.weight" in old_key:
+                layer = old_key.replace(".lora.up.weight", "")
+                print(f"  {layer}")
+                print(f"    .lora.up.weight    →  .lora_B.weight")
+    
+    return new_state_dict
+
+
+def convert_simpletuner_transformer_to_diffusers(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """
+    Convert SimpleTuner transformer format (already has transformer. prefix but uses lora_down/lora_up)
+    to diffusers PEFT format (transformer. prefix with lora_A/lora_B).
+    
+    This is a simpler conversion since the key structure is already correct.
+    """
+    new_state_dict = {}
+    renames = []
+    
+    # Get all unique LoRA layer base names (without .lora_down/.lora_up/.alpha suffix)
+    all_keys = list(state_dict.keys())
+    base_keys = set()
+    
+    for key in all_keys:
+        if ".lora_down.weight" in key:
+            base_key = key.replace(".lora_down.weight", "")
+            base_keys.add(base_key)
+    
+    print(f"\nFound {len(base_keys)} LoRA layers to convert")
+    print("-" * 80)
+    
+    # Convert each layer
+    for base_key in sorted(base_keys):
+        down_key = f"{base_key}.lora_down.weight"
+        up_key = f"{base_key}.lora_up.weight"
+        alpha_key = f"{base_key}.alpha"
+        
+        if down_key not in state_dict or up_key not in state_dict:
+            print(f"⚠ Warning: Missing weights for {base_key}")
+            continue
+        
+        down_weight = state_dict.pop(down_key)
+        up_weight = state_dict.pop(up_key)
+        
+        # Handle alpha scaling
+        has_alpha = False
+        if alpha_key in state_dict:
+            alpha = state_dict.pop(alpha_key)
+            lora_rank = down_weight.shape[0]
+            scale = alpha / lora_rank
+            
+            # Calculate scale_down and scale_up to preserve the scale value
+            scale_down = scale
+            scale_up = 1.0
+            while scale_down * 2 < scale_up:
+                scale_down *= 2
+                scale_up /= 2
+            
+            down_weight = down_weight * scale_down
+            up_weight = up_weight * scale_up
+            has_alpha = True
+        
+        # Store in PEFT format (lora_A = down, lora_B = up)
+        new_down_key = f"{base_key}.lora_A.weight"
+        new_up_key = f"{base_key}.lora_B.weight"
+        
+        new_state_dict[new_down_key] = down_weight
+        new_state_dict[new_up_key] = up_weight
+        
+        renames.append((down_key, new_down_key, has_alpha))
+        renames.append((up_key, new_up_key, has_alpha))
+    
+    # Check for any remaining keys
+    remaining = [k for k in state_dict.keys() if not k.startswith("text_encoder")]
+    if remaining:
+        print(f"⚠ Warning: {len(remaining)} keys were not converted: {remaining[:5]}")
+    
+    print(f"\nRenames applied ({len(renames)} conversions):")
+    print("-" * 80)
+    
+    # Group by layer
+    current_layer = None
+    for old_key, new_key, has_alpha in renames:
+        layer = old_key.replace(".lora_down.weight", "").replace(".lora_up.weight", "")
+        
+        if layer != current_layer:
+            alpha_str = " (alpha scaled)" if has_alpha else ""
+            print(f"\n  {layer}{alpha_str}")
+            current_layer = layer
+        
+        if ".lora_down.weight" in old_key:
+            print(f"    .lora_down.weight  →  .lora_A.weight")
+        elif ".lora_up.weight" in old_key:
+            print(f"    .lora_up.weight    →  .lora_B.weight")
+    
+    return new_state_dict
+
+
+def convert_simpletuner_auraflow_to_diffusers(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """
+    Convert SimpleTuner AuraFlow LoRA format to diffusers PEFT format.
+    
+    SimpleTuner typically saves LoRAs in a format similar to Kohya's sd-scripts,
+    but for transformer-based models like AuraFlow, the keys may differ.
+    """
+    new_state_dict = {}
+    
+    def _convert(original_key, diffusers_key, state_dict, new_state_dict):
+        """Helper to convert a single LoRA layer."""
+        down_key = f"{original_key}.lora_down.weight"
+        if down_key not in state_dict:
+            return False
+            
+        down_weight = state_dict.pop(down_key)
+        lora_rank = down_weight.shape[0]
+        
+        up_weight_key = f"{original_key}.lora_up.weight"
+        up_weight = state_dict.pop(up_weight_key)
+        
+        # Handle alpha scaling
+        alpha_key = f"{original_key}.alpha"
+        if alpha_key in state_dict:
+            alpha = state_dict.pop(alpha_key)
+            scale = alpha / lora_rank
+            
+            # Calculate scale_down and scale_up to preserve the scale value
+            scale_down = scale
+            scale_up = 1.0
+            while scale_down * 2 < scale_up:
+                scale_down *= 2
+                scale_up /= 2
+            
+            down_weight = down_weight * scale_down
+            up_weight = up_weight * scale_up
+        
+        # Store in PEFT format (lora_A = down, lora_B = up)
+        diffusers_down_key = f"{diffusers_key}.lora_A.weight"
+        new_state_dict[diffusers_down_key] = down_weight
+        new_state_dict[diffusers_down_key.replace(".lora_A.", ".lora_B.")] = up_weight
+        
+        return True
+    
+    # Get all unique LoRA layer names
+    all_unique_keys = {
+        k.replace(".lora_down.weight", "").replace(".lora_up.weight", "").replace(".alpha", "")
+        for k in state_dict
+        if ".lora_down.weight" in k or ".lora_up.weight" in k or ".alpha" in k
+    }
+    
+    # Process transformer blocks
+    for original_key in sorted(all_unique_keys):
+        if original_key.startswith("lora_transformer_single_transformer_blocks_"):
+            # Single transformer blocks
+            parts = original_key.split("lora_transformer_single_transformer_blocks_")[-1].split("_")
+            block_idx = int(parts[0])
+            diffusers_key = f"single_transformer_blocks.{block_idx}"
+            
+            # Map the rest of the key
+            remaining = "_".join(parts[1:])
+            if "attn_to_q" in remaining:
+                diffusers_key += ".attn.to_q"
+            elif "attn_to_k" in remaining:
+                diffusers_key += ".attn.to_k"
+            elif "attn_to_v" in remaining:
+                diffusers_key += ".attn.to_v"
+            elif "proj_out" in remaining:
+                diffusers_key += ".proj_out"
+            elif "proj_mlp" in remaining:
+                diffusers_key += ".proj_mlp"
+            elif "norm_linear" in remaining:
+                diffusers_key += ".norm.linear"
+            else:
+                print(f"Warning: Unhandled single block key pattern: {original_key}")
+                continue
+                
+        elif original_key.startswith("lora_transformer_transformer_blocks_"):
+            # Double transformer blocks
+            parts = original_key.split("lora_transformer_transformer_blocks_")[-1].split("_")
+            block_idx = int(parts[0])
+            diffusers_key = f"transformer_blocks.{block_idx}"
+            
+            # Map the rest of the key
+            remaining = "_".join(parts[1:])
+            if "attn_to_out_0" in remaining:
+                diffusers_key += ".attn.to_out.0"
+            elif "attn_to_add_out" in remaining:
+                diffusers_key += ".attn.to_add_out"
+            elif "attn_to_q" in remaining:
+                diffusers_key += ".attn.to_q"
+            elif "attn_to_k" in remaining:
+                diffusers_key += ".attn.to_k"
+            elif "attn_to_v" in remaining:
+                diffusers_key += ".attn.to_v"
+            elif "attn_add_q_proj" in remaining:
+                diffusers_key += ".attn.add_q_proj"
+            elif "attn_add_k_proj" in remaining:
+                diffusers_key += ".attn.add_k_proj"
+            elif "attn_add_v_proj" in remaining:
+                diffusers_key += ".attn.add_v_proj"
+            elif "ff_net_0_proj" in remaining:
+                diffusers_key += ".ff.net.0.proj"
+            elif "ff_net_2" in remaining:
+                diffusers_key += ".ff.net.2"
+            elif "ff_context_net_0_proj" in remaining:
+                diffusers_key += ".ff_context.net.0.proj"
+            elif "ff_context_net_2" in remaining:
+                diffusers_key += ".ff_context.net.2"
+            elif "norm1_linear" in remaining:
+                diffusers_key += ".norm1.linear"
+            elif "norm1_context_linear" in remaining:
+                diffusers_key += ".norm1_context.linear"
+            else:
+                print(f"Warning: Unhandled double block key pattern: {original_key}")
+                continue
+        
+        elif original_key.startswith("lora_te1_") or original_key.startswith("lora_te_"):
+            # Text encoder keys - handle separately
+            print(f"Found text encoder key: {original_key}")
+            continue
+        
+        else:
+            print(f"Warning: Unknown key pattern: {original_key}")
+            continue
+        
+        # Perform the conversion
+        _convert(original_key, diffusers_key, state_dict, new_state_dict)
+    
+    # Add "transformer." prefix to all keys
+    transformer_state_dict = {
+        f"transformer.{k}": v for k, v in new_state_dict.items() if not k.startswith("text_model.")
+    }
+    
+    # Check for remaining unconverted keys
+    if len(state_dict) > 0:
+        remaining_keys = [k for k in state_dict.keys() if not k.startswith("lora_te")]
+        if remaining_keys:
+            print(f"Warning: Some keys were not converted: {remaining_keys[:10]}")
+    
+    return transformer_state_dict
+
+
+def convert_lora(input_path: str, output_path: str) -> None:
+    """
+    Main conversion function.
+    
+    Args:
+        input_path: Path to input LoRA safetensors file
+        output_path: Path to output diffusers-compatible safetensors file
+    """
+    print(f"Loading LoRA from: {input_path}")
+    state_dict = safetensors.torch.load_file(input_path)
+    
+    print(f"Detecting LoRA format...")
+    format_type = detect_lora_format(state_dict)
+    print(f"Detected format: {format_type}")
+    
+    if format_type == "peft":
+        print("LoRA is already in diffusers-compatible PEFT format!")
+        print("No conversion needed. Copying file...")
+        import shutil
+        shutil.copy(input_path, output_path)
+        return
+    
+    elif format_type == "mixed":
+        print("Converting MIXED format LoRA to pure PEFT format...")
+        print("(Some layers use lora_A/B, others use .lora.down/.lora.up)")
+        converted_state_dict = convert_mixed_lora_to_diffusers(state_dict.copy())
+    
+    elif format_type == "simpletuner_transformer":
+        print("Converting SimpleTuner transformer format to diffusers...")
+        print("(has transformer. prefix but uses lora_down/lora_up naming)")
+        converted_state_dict = convert_simpletuner_transformer_to_diffusers(state_dict.copy())
+    
+    elif format_type == "simpletuner_auraflow":
+        print("Converting SimpleTuner AuraFlow format to diffusers...")
+        converted_state_dict = convert_simpletuner_auraflow_to_diffusers(state_dict.copy())
+    
+    elif format_type == "kohya":
+        print("Note: Detected Kohya format. This converter is optimized for AuraFlow.")
+        print("For other models, diffusers has built-in conversion.")
+        converted_state_dict = convert_simpletuner_auraflow_to_diffusers(state_dict.copy())
+    
+    else:
+        print("Error: Unknown LoRA format!")
+        print("Sample keys from the state dict:")
+        for i, key in enumerate(list(state_dict.keys())[:20]):
+            print(f"  {key}")
+        sys.exit(1)
+    
+    print(f"Saving converted LoRA to: {output_path}")
+    safetensors.torch.save_file(converted_state_dict, output_path)
+    
+    print("\nConversion complete!")
+    print(f"Original keys: {len(state_dict)}")
+    print(f"Converted keys: {len(converted_state_dict)}")
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert SimpleTuner LoRA to diffusers-compatible format",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Convert a SimpleTuner LoRA for AuraFlow
+  python convert_simpletuner_lora.py my_lora.safetensors diffusers_lora.safetensors
+  
+  # Check format without converting
+  python convert_simpletuner_lora.py my_lora.safetensors /tmp/test.safetensors
+        """
+    )
+    
+    parser.add_argument(
+        "input",
+        type=str,
+        help="Input LoRA file (SimpleTuner format)"
+    )
+    
+    parser.add_argument(
+        "output",
+        type=str,
+        help="Output LoRA file (diffusers-compatible format)"
+    )
+    
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Only detect format, don't convert"
+    )
+    
+    args = parser.parse_args()
+    
+    # Validate input file exists
+    if not Path(args.input).exists():
+        print(f"Error: Input file not found: {args.input}")
+        sys.exit(1)
+    
+    if args.dry_run:
+        print(f"Loading LoRA from: {args.input}")
+        state_dict = safetensors.torch.load_file(args.input)
+        format_type = detect_lora_format(state_dict)
+        print(f"Detected format: {format_type}")
+        print(f"\nSample keys ({min(10, len(state_dict))} of {len(state_dict)}):")
+        for key in list(state_dict.keys())[:10]:
+            print(f"  {key}")
+        return
+    
+    convert_lora(args.input, args.output)
+
+
+if __name__ == "__main__":
+    main()
+