Fal Wan Image Trainer lora to ComfyUI

fal_to_comfy_wan.py

from safetensors.torch import load_file, save_file
import argparse


def convert_lora_keys_for_wan_comfy(input_path, output_path):
    try:
        state_dict = load_file(input_path)
        new_state_dict = {}

        print(f"Successfully loaded LoRA from: {input_path}")
        print(f"Found {len(state_dict)} keys to convert.")

        for key, value in state_dict.items():
            new_key = key
            new_key = new_key.replace("transformer.", "diffusion_model.")

            # Attention block renaming
            new_key = new_key.replace(".attn1.", ".self_attn.")
            new_key = new_key.replace(".attn2.", ".cross_attn.")

            # Projection renaming (q, k, v, out)
            new_key = new_key.replace(".to_q.", ".q.")
            new_key = new_key.replace(".to_k.", ".k.")
            new_key = new_key.replace(".to_v.", ".v.")
            new_key = new_key.replace(".to_out.0.", ".o.")

            # Feed-forward network renaming
            new_key = new_key.replace(".ffn.net.0.proj.", ".ffn.0.")
            new_key = new_key.replace(".ffn.net.2.", ".ffn.2.")

            new_state_dict[new_key] = value

        print(f"\nConversion complete. {len(new_state_dict)} keys processed.")

        # Save the new state dictionary to the output file
        save_file(new_state_dict, output_path)
        print(f"Successfully converted LoRA and saved to: {output_path}")

        # Print a sample of the new keys for verification
        print("\n--- Sample of 10 New Keys ---")
        for i, key in enumerate(new_state_dict.keys()):
            if i >= 10:
                break
            print(f"Original: {list(state_dict.keys())[i]}")
            print(f"New -> : {key}\n")

    except Exception as e:
        print(f"An error occurred during conversion: {e}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Convert Wan LoRA keys for ComfyUI compatibility (non-fused)."
    )
    parser.add_argument(
        "input_path", type=str, help="Path to the input LoRA safetensor file."
    )
    parser.add_argument(
        "--output_path",
        type=str,
        help="Path to save the converted LoRA file. If not provided, it will be auto-generated.",
    )

    args = parser.parse_args()

    output_path = args.output_path
    if not output_path:
        input_filename = args.input_path.replace("\\", "/").split("/")[-1].split(".")[0]
        output_path = f"{input_filename}_converted.safetensors"

    convert_lora_keys_for_wan_comfy(args.input_path, output_path)