"""Load a training checkpoint and save a lean inference-only model file.

Keeps only the fp16 model weights (plus the config needed to rebuild the
architecture). Drops optimizer state, scaler state, and everything else.

Usage:
    python quantize.py --input ckpt.pt --output evotalk_fp16.pt
"""

import os
import argparse

import torch

from model import EvoTalk


def quantize(input_path, output_path):
    ckpt = torch.load(input_path, map_location="cpu", weights_only=False)

    config = ckpt.get("config", None)
    if config is None:
        raise ValueError("checkpoint has no config; cannot rebuild the model")

    model = EvoTalk(config)
    model.load_state_dict(ckpt["model"])
    model.eval()

    state = {
        "model": {k: v.half() for k, v in model.state_dict().items()},
        "config": config,
        "step": ckpt.get("step", 0),
    }

    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    torch.save(state, output_path)
    print(f"saved fp16 weights -> {output_path}")
    print(f"weights: {sum(v.numel() for v in state['model'].values()) / 1e6:.2f}M "
          f"({os.path.getsize(output_path) / 1e6:.1f} MB)")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=str, default="ckpt.pt")
    parser.add_argument("--output", type=str, default="evotalk_fp16.pt")
    args = parser.parse_args()

    quantize(args.input, args.output)