lora_diff.py

import argparse
from safetensors.torch import safe_open

def load_structure(path):
    """Return {key: (shape, dtype)} dict from safetensors file"""
    tensors = {}
    with safe_open(path, framework="pt") as f:
        for k in f.keys():
            t = f.get_tensor(k)
            tensors[k] = (tuple(t.shape), str(t.dtype))
    return tensors


def compare_loras(path1, path2, output_txt=None):
    print(f"Comparing:\n  1️⃣ {path1}\n  2️⃣ {path2}\n")

    t1 = load_structure(path1)
    t2 = load_structure(path2)

    keys1 = set(t1.keys())
    keys2 = set(t2.keys())

    only_in_1 = sorted(keys1 - keys2)
    only_in_2 = sorted(keys2 - keys1)
    common = sorted(keys1 & keys2)

    lines = []
    lines.append(f"Total tensors in 1: {len(t1)}")
    lines.append(f"Total tensors in 2: {len(t2)}\n")

    if only_in_1:
        lines.append("⚠️ Keys only in first file:")
        for k in only_in_1:
            lines.append(f"  {k}  [{t1[k][0]}]  {t1[k][1]}")
        lines.append("")

    if only_in_2:
        lines.append("⚠️ Keys only in second file:")
        for k in only_in_2:
            lines.append(f"  {k}  [{t2[k][0]}]  {t2[k][1]}")
        lines.append("")

    diffs = []
    for k in common:
        s1, d1 = t1[k]
        s2, d2 = t2[k]
        if s1 != s2 or d1 != d2:
            diffs.append((k, s1, d1, s2, d2))

    if diffs:
        lines.append("🔍 Mismatched tensors:")
        for k, s1, d1, s2, d2 in diffs:
            lines.append(f"  {k}")
            lines.append(f"    file1: {s1} {d1}")
            lines.append(f"    file2: {s2} {d2}")
        lines.append("")
    else:
        lines.append("✅ All shared tensors match in shape and dtype.\n")

    report = "\n".join(lines)
    print(report)

    if output_txt:
        with open(output_txt, "w", encoding="utf-8") as out:
            out.write(report)
        print(f"Result saved to: {output_txt}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Compare two LoRA .safetensors structures")
    parser.add_argument("--file1", required=True, help="Path to first .safetensors")
    parser.add_argument("--file2", required=True, help="Path to second .safetensors")
    parser.add_argument("--output", help="Optional path to save comparison result")
    args = parser.parse_args()

    compare_loras(args.file1, args.file2, args.output)