# Python script. Merge two LoRA adapters.

import torch
from safetensors.torch import load_file, save_file
import math

# =====================================================
# Configuration
# =====================================================

# The two LORA adapters must come/trained from same base model. Must have same layers.
# Input
LORA1_PATH = "male_vocals_adapter_model.safetensors"
LORA2_PATH = "instrumental_adapter_model.safetensors"
# Output
OUTPUT_PATH = "adapter_model.safetensors"

# LORA strength for each input LORA
LORA1_STRENGTH = 0.8    #Try different value pairs e.g. 0.8  0.6  0.6  1.4
LORA2_STRENGTH = 0.8    #Try different value pairs e.g. 0.6  0.8  1.4  0.6

MAX_OUTPUT_RANK = None   # None = automatic, min of the two input LORA
AUTO_STABILIZE = False   #True or False , if the values in layers are to strong you can get odd sounds. Use True and the values will be adjusted. False, no adjusting. 

# OBS! When use LORA in ACE-Step1.5, you also need the adapter_config.json file from one of the input LORA. Not when running this script, just when loading and use.


# =====================================================


def rms(t):
    return torch.sqrt(torch.mean(t.float() ** 2))


def get_lora_pairs(state_dict):
    pairs = {}
    for key in state_dict:
        if "lora_A" in key:
            base = key.replace("lora_A.weight", "")
            b_key = base + "lora_B.weight"
            if b_key in state_dict:
                pairs[base] = (key, b_key)
    return pairs


def check_compatibility(sd1, sd2):
    pairs1 = get_lora_pairs(sd1)
    pairs2 = get_lora_pairs(sd2)

    if pairs1.keys() != pairs2.keys():
        raise ValueError("LoRA layers mismatch between files.")

    print(f"✓ Found {len(pairs1)} matching LoRA layers")
    return pairs1


def compute_delta(sd, a_key, b_key):
    A = sd[a_key]
    B = sd[b_key]

    # LoRA scaling
    rank = A.shape[0]
    alpha_key = a_key.replace("lora_A.weight", "alpha")

    alpha = sd.get(alpha_key, torch.tensor(rank))
    scale = alpha / rank

    delta = (B @ A) * scale
    return delta, rank


def factorize_to_lora(delta, rank):
    # SVD
    U, S, Vh = torch.linalg.svd(delta, full_matrices=False)

    U = U[:, :rank]
    S = S[:rank]
    Vh = Vh[:rank, :]

    B = U @ torch.diag(torch.sqrt(S))
    A = torch.diag(torch.sqrt(S)) @ Vh

    return A, B


def merge_loras(sd1, sd2):
    merged_sd = {}
    pairs = check_compatibility(sd1, sd2)

    for base, (a_key, b_key) in pairs.items():
        print(f"\n--- Merging layer: {base} ---")

        delta1, r1 = compute_delta(sd1, a_key, b_key)
        delta2, r2 = compute_delta(sd2, a_key, b_key)

        merged_delta = (LORA1_STRENGTH * delta1) + (LORA2_STRENGTH * delta2)

        rms1 = rms(delta1)
        rms2 = rms(delta2)
        rmsm = rms(merged_delta)

        print(f"RMS LORA1: {rms1:.6f}")
        print(f"RMS LORA2: {rms2:.6f}")
        print(f"RMS MERGED: {rmsm:.6f}")

        if AUTO_STABILIZE:
            max_allowed = max(rms1, rms2) * 1.2
            if rmsm > max_allowed:
                scale = max_allowed / rmsm
                merged_delta *= scale
                print(f"Stabilized by factor {scale:.4f}")

        output_rank = MAX_OUTPUT_RANK or min(r1, r2)

        A_new, B_new = factorize_to_lora(merged_delta, output_rank)

        dtype = sd1[a_key].dtype
        merged_sd[a_key] = A_new.to(dtype)
        merged_sd[b_key] = B_new.to(dtype)

        # alpha set same rank for stability
        alpha_key = a_key.replace("lora_A.weight", "alpha")
        merged_sd[alpha_key] = torch.tensor(output_rank, dtype=dtype)

        print(f"Output rank: {output_rank}")
        print(f"Max abs value: {torch.max(torch.abs(merged_delta)):.6f}")

    return merged_sd


def main():
    print("Loading LoRA 1...")
    sd1 = load_file(LORA1_PATH)

    print("Loading LoRA 2...")
    sd2 = load_file(LORA2_PATH)

    print("Merging...")
    merged = merge_loras(sd1, sd2)

    print("\nSaving merged LoRA...")
    save_file(merged, OUTPUT_PATH)

    print("\n✓ Merge complete.")
    print(f"Saved to: {OUTPUT_PATH}")


if __name__ == "__main__":
    main()