| |
| """Create SDXSv3 (random init) with a given config and save to ./transformer. |
| |
| Defaults match the trained model (transformer/config.json): 2560 features, |
| 20 heads (full attention, kv=heads), 2 input + 12 middle (x2 loop) + 2 output, |
| patch 2, 32 latent channels, multi-layer text (txtlayers=3). |
| """ |
| import argparse |
|
|
| from transformer_sdxsv3 import SDXSv3Transformer, SDXSv3Config |
|
|
|
|
| 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=12) |
| 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("--loops", 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 = SDXSv3Config( |
| n_in=args.n_in, n_mid=args.n_mid, n_out=args.n_out, loops=args.loops, |
| 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 = SDXSv3Transformer(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} проходов, " |
| f"txtlayers={cfg.txtlayers}") |
| print(f" параметры: {n/1e9:.3f}B, bf16 {n*2/2**30:.2f} ГБ") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
