#!/usr/bin/env python3
"""Create SDXS (random init) with a given config and save to ./transformer.

Defaults: 2560 features, 20 heads (full attention, kv=heads), 2+20+2 classic,
patch 2, 32 latent channels, multi-layer text (txtlayers=3).
The checkpoint layout is 6+8+6 (20 passes):
  python make_model.py --n-in 6 --n-mid 8 --n-out 6
"""
import argparse

from transformer_sdxs import SDXSTransformer, SDXSTransformerConfig


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--features", type=int, default=2560)
    ap.add_argument("--heads", type=int, default=20, help="attention heads (full attention, kv=heads)")
    ap.add_argument("--n-mid", type=int, default=20)
    ap.add_argument("--n-in", type=int, default=2)
    ap.add_argument("--n-out", type=int, default=2)
    ap.add_argument("--patch", type=int, default=2)
    ap.add_argument("--txtlayers", type=int, default=3, help="number of text layers to fuse (multi-layer, e.g. 3 = [2,12,22])")
    ap.add_argument("--txtdim", type=int, default=2048)
    ap.add_argument("--txtheads", type=int, default=8)
    ap.add_argument("--txtkvheads", type=int, default=8)
    ap.add_argument("--multiplier", type=int, default=4)
    ap.add_argument("--channels", type=int, default=32)
    ap.add_argument("--out", default="transformer")
    args = ap.parse_args()

    cfg = SDXSTransformerConfig(
        n_in=args.n_in, n_mid=args.n_mid, n_out=args.n_out,
        features=args.features, heads=args.heads, kvheads=args.heads,
        patch=args.patch, channels=args.channels, multiplier=args.multiplier,
        txtlayers=args.txtlayers, txtdim=args.txtdim, txtheads=args.txtheads,
        txtkvheads=args.txtkvheads,
    )
    m = SDXSTransformer(cfg)
    n = sum(p.numel() for p in m.parameters())
    import torch
    m.to(torch.bfloat16).save_pretrained(args.out)  # bf16, как чекпоинты трейна
    print(f"модель создана (случайный init) и сохранена в {args.out}:")
    print(f"  {args.features} features, {args.heads} heads (full attention, kv={args.heads}), "
          f"{cfg.n_in}+{cfg.n_mid}+{cfg.n_out} уникальных блоков = {m.num_passes} проходов, txtlayers={cfg.txtlayers}")
    print(f"  параметры: {n/1e9:.3f}B, bf16 {n*2/2**30:.2f} ГБ")


if __name__ == "__main__":
    main()