alibaba-pai

/

MiniMax-H3-Acc-LoRAs

like 129

Follow

Alibaba-PAI 5.01k

NOTE: I Think it's not working. I just noticed it.

#!/usr/bin/env python3

import argparse
import re
import sys

try:
    from safetensors.torch import load_file, save_file
except ImportError:
    sys.exit("Missing dependency. Run: pip install safetensors")


def convert_key(key: str, add_prefix: bool) -> str | None:
    # Drop the PEFT adapter name segment, e.g. ".default" (or any adapter name)
    key = re.sub(r"\.lora_A\.[^.]+\.weight$", ".lora_A.weight", key)
    key = re.sub(r"\.lora_B\.[^.]+\.weight$", ".lora_B.weight", key)

    if key.endswith(".lora_A.weight"):
        key = key[: -len(".lora_A.weight")] + ".lora_down.weight"
    elif key.endswith(".lora_B.weight"):
        key = key[: -len(".lora_B.weight")] + ".lora_up.weight"
    else:
        # Not a recognized lora weight key (e.g. could be an alpha tensor) - pass through
        pass

    if add_prefix and not key.startswith("diffusion_model."):
        key = "diffusion_model." + key

    return key


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("input", help="path to input .safetensors LoRA")
    ap.add_argument("output", help="path to write converted .safetensors LoRA")
    ap.add_argument(
        "--no-prefix",
        action="store_true",
        help="do not prepend 'diffusion_model.' to converted keys",
    )
    args = ap.parse_args()

    add_prefix = not args.no_prefix

    tensors = load_file(args.input)
    converted = {}
    skipped = []

    for k, v in tensors.items():
        new_k = convert_key(k, add_prefix)
        if new_k in converted:
            print(f"[WARN] duplicate key after conversion, overwriting: {new_k}")
        converted[new_k] = v

    print(f"Converted {len(tensors)} keys.")
    print("Sample of converted keys:")
    for k in list(converted.keys())[:6]:
        print(" ", k)

    save_file(converted, args.output)
    print(f"\nSaved: {args.output}")


if __name__ == "__main__":
    main()