Upload convert_lora.py

convert_lora.py

@@ -0,0 +1,78 @@
+import os
+import re
+from safetensors.torch import load_file, save_file
+
+def convert_key(key):
+    # 1. Remove the original prefix
+    new_key = key.replace("lora_unet__", "")
+    
+    # 2. Identify the parts.
+    parts = new_key.split(".")
+    block_name = parts[0]  # e.g., "layers_0_attention_to_k"
+    rest = parts[1:]       # e.g., ["lora_down", "weight"]
+
+    # --- FIX 1: Layers and Refiners (Force dot after number) ---
+    # The previous script missed the trailing underscore. 
+    # We replace "layers_0_" with "layers.0."
+    
+    # Fix "layers_X_" -> "layers.X."
+    block_name = re.sub(r'layers_(\d+)_', r'layers.\1.', block_name)
+    
+    # Fix "context_refiner_X_" -> "context_refiner.X."
+    block_name = re.sub(r'context_refiner_(\d+)_', r'context_refiner.\1.', block_name)
+    
+    # Fix "noise_refiner_X_" -> "noise_refiner.X."
+    block_name = re.sub(r'noise_refiner_(\d+)_', r'noise_refiner.\1.', block_name)
+
+    # --- FIX 2: Attention blocks ---
+    
+    # Fix "attention_to_k" -> "attention.to_k"
+    for t in ["to_k", "to_q", "to_v"]:
+        if f"_{t}" in block_name:
+            block_name = block_name.replace(f"_{t}", f".{t}")
+            
+    # Fix "attention_to_out_0" -> "attention.to_out.0"
+    # This is specific: ComfyUI expects "to_out.0"
+    if "_to_out" in block_name:
+        block_name = block_name.replace("_to_out", ".to_out")
+        # Handle the trailing number for output (e.g. to_out_0 -> to_out.0)
+        block_name = re.sub(r'\.to_out_(\d+)', r'.to_out.\1', block_name)
+
+    # --- FINALIZE ---
+    final_key = "diffusion_model." + ".".join([block_name] + rest)
+    return final_key
+
+# --- MAIN EXECUTION ---
+print("Looking for .safetensors files...")
+# Find original file (exclude previously converted ones to be safe)
+files = [f for f in os.listdir('.') if f.endswith('.safetensors') and 'converted' not in f and 'fixed' not in f]
+
+if not files:
+    print("Error: No original .safetensors file found.")
+    print("Make sure the original 'Z-Image-Fun-Lora-Distill-8-Steps.safetensors' is in this folder.")
+    exit()
+
+input_file = files[0]
+print(f"Processing: {input_file}")
+
+try:
+    tensors = load_file(input_file)
+    new_tensors = {}
+
+    print("Converting keys...")
+    # Print a preview of the first converted key to verify
+    first_key = list(tensors.keys())[0]
+    print(f"Preview: {first_key} \n      -> {convert_key(first_key)}")
+
+    for k, v in tensors.items():
+        new_k = convert_key(k)
+        new_tensors[new_k] = v
+        
+    output_file = input_file.replace(".safetensors", "_v3_fixed.safetensors")
+    save_file(new_tensors, output_file)
+
+    print(f"\nSUCCESS! Created: {output_file}")
+    print("Move this file to ComfyUI/models/loras/ and DELETE the old converted ones.")
+
+except Exception as e:
+    print(f"An error occurred: {e}")