File size: 2,346 Bytes
ef2542c e9c970f ef2542c e9c970f ef2542c e9c970f ef2542c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | #!/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()
|