"""权重瘦身: 只保留推理需要的部分。

推理的加载逻辑是「先 load model_state_dict, 再用 ema_state_dict.shadow 逐参数覆盖」,
所以把 EMA 影子权重合并进 model_state_dict 后只存一份, 推理结果逐位等价。

丢弃: optimizer_state_dict(最大头, AdamW 的 m/v 约 2x 参数量)、scheduler_state_dict、训练元信息。

用法: python tools_strip_ckpt.py <输入.pth> <输出.pth> [--no-ema]
"""
import os, sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import argparse, os, torch


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("src"); ap.add_argument("dst")
    ap.add_argument("--no-ema", action="store_true", help="不合并 EMA, 保留原始训练权重")
    a = ap.parse_args()

    ck = torch.load(a.src, map_location="cpu", weights_only=False)
    print(f"输入 {a.src}  ({os.path.getsize(a.src)/1e9:.2f} GB)")
    print("  顶层 key:", list(ck.keys()))
    msd = ck["model_state_dict"]
    shadow = (ck.get("ema_state_dict") or {}).get("shadow") if isinstance(ck.get("ema_state_dict"), dict) else None

    merged = 0
    if shadow and not a.no_ema:
        for k in msd:
            if k in shadow:
                msd[k] = shadow[k].to(msd[k].dtype); merged += 1
        print(f"  已合并 EMA 影子权重: {merged}/{len(msd)} 个张量")
    elif a.no_ema:
        print("  --no-ema: 保留原始训练权重")
    else:
        print("  ckpt 内无 EMA, 直接使用 model_state_dict")

    out = {"model_state_dict": msd}
    torch.save(out, a.dst)
    print(f"输出 {a.dst}  ({os.path.getsize(a.dst)/1e9:.2f} GB, "
          f"省 {100*(1-os.path.getsize(a.dst)/os.path.getsize(a.src)):.0f}%)")


if __name__ == "__main__":
    main()