import os
import json
import argparse

import torch
from diffusers import UNet2DConditionModel
from safetensors.torch import save_file


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--src", type=str, default="unet_old/config.json")
    parser.add_argument("--out", type=str, required=True)
    parser.add_argument("--layers-per-block", type=str, default=None)
    parser.add_argument("--transformer-layers", type=str, default=None)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    with open(args.src, "r", encoding="utf-8") as f:
        cfg = json.load(f)

    for k in ["_class_name", "_diffusers_version", "_name_or_path"]:
        cfg.pop(k, None)

    if args.layers_per_block is not None:
        val = args.layers_per_block
        try:
            cfg["layers_per_block"] = int(val)
        except ValueError:
            cfg["layers_per_block"] = [int(x) for x in val.split(",")]
    if args.transformer_layers is not None:
        cfg["transformer_layers_per_block"] = [int(x) for x in args.transformer_layers.split(",")]

    torch.manual_seed(args.seed)
    model = UNet2DConditionModel(**cfg)
    params = sum(p.numel() for p in model.parameters()) / 1e9
    print(f"Created {args.out}: {params:.3f}B params")

    os.makedirs(args.out, exist_ok=True)
    with open(os.path.join(args.out, "config.json"), "w", encoding="utf-8") as f:
        json.dump({
            "_class_name": "UNet2DConditionModel",
            "_diffusers_version": "0.37.1",
            "_name_or_path": args.out,
            **cfg,
        }, f, indent=2)

    save_file(model.state_dict(), os.path.join(args.out, "diffusion_pytorch_model.safetensors"))
    print("Saved")


if __name__ == "__main__":
    main()