| |
| """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) |
| 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() |
|
|
