#!/usr/bin/env python3
import argparse
import os
from pathlib import Path

import torch


def convert_state_dict_to_fp16(state_dict):
    converted = {}
    for key, value in state_dict.items():
        if isinstance(value, torch.Tensor) and value.is_floating_point():
            converted[key] = value.half()
        else:
            converted[key] = value
    return converted


def convert_checkpoint(src_path: Path, dst_path: Path):
    state_dict = torch.load(src_path, map_location="cpu", weights_only=True)
    converted = convert_state_dict_to_fp16(state_dict)
    torch.save(converted, dst_path)
    src_size = os.path.getsize(src_path)
    dst_size = os.path.getsize(dst_path)
    print(
        f"{src_path.name} -> {dst_path.name} | "
        f"{src_size / 1024**2:.2f} MiB -> {dst_size / 1024**2:.2f} MiB"
    )


def main():
    parser = argparse.ArgumentParser(description="Create fp16 checkpoint copies alongside the original fp32 files.")
    parser.add_argument(
        "--files",
        nargs="*",
        default=["llm.pt", "flow.pt", "hift.pt"],
        help="Checkpoint files to convert.",
    )
    args = parser.parse_args()

    for rel_path in args.files:
        src_path = Path(rel_path)
        if not src_path.exists():
            raise FileNotFoundError(f"Missing checkpoint: {src_path}")
        if src_path.suffix != ".pt":
            raise ValueError(f"Expected a .pt checkpoint, got: {src_path}")
        dst_path = src_path.with_suffix(f".fp16{src_path.suffix}")
        convert_checkpoint(src_path, dst_path)


if __name__ == "__main__":
    main()