nx_genarch.nx source
↩ module page · 151 lines · 5914 B
1// nx_genarch.nx -- derive a DiT's architecture FROM THE MODEL FILE, so parts are hot-swappable.
2//
3// Operator, 2026-08-07: *"make sure that they are hot swappable with other models till we can pay
4// to make our own model"*. Z-Image Turbo is a MIX of parts (DiT + Qwen text encoder + VAE); the
5// right goal is not to rebuild each part but to be able to drop a different one in.
6//
7// A hardcoded architecture constant is the thing that prevents that. This lane has already been
8// burned four times by assumed constants -- layer count (30, not the 36 I assumed), norm_eps,
9// adaLN chunk order, SDPA scale -- and each was silently wrong rather than loudly wrong.
10// ★★★★★ AN ARCHITECTURE CONSTANT IN THE CODE IS A MODEL THAT CANNOT BE SWAPPED.
11//
12// Everything below is read off tensor SHAPES, because this GGUF carries zero metadata KV pairs
13// (measured: `kv 0`) -- so the shapes are the only self-description the file has.
14//
15// SLOTS
16// 0 hidden D from attention.qkv.weight dim0
17// 1 head_dim from attention.q_norm.weight length
18// 2 n_heads qkv_width / 3 / head_dim (q,k,v equal-width; verified, not assumed)
19// 3 qkv_width from attention.qkv.weight dim1
20// 4 ffn_dim from feed_forward.w1.weight dim1
21// 5 n_layers probed: layers.N until absent
22// 6 n_refiners probed: context_refiner.N until absent
23// 7 adaln_embed from adaLN_modulation.0.weight dim0
24// 8 n_chunks adaLN_modulation.0.weight dim1 / D
25// 9 out_dim from final_layer.linear.weight dim1 (= patch*patch*latent_channels)
26// 10 x_embed_in from x_embedder.weight dim0
27// license_tier: ORIGINAL
28
29import "nx_syscalls.nx"
30import "nx_strconv.nx"
31import "nx_genweights.nx"
32
33const NX_ARCH_SLOTS: i64 = 16
34const NX_ARCH_D: i64 = 0
35const NX_ARCH_HEAD_DIM: i64 = 1
36const NX_ARCH_N_HEADS: i64 = 2
37const NX_ARCH_QKV_W: i64 = 3
38const NX_ARCH_FFN_DIM: i64 = 4
39const NX_ARCH_N_LAYERS: i64 = 5
40const NX_ARCH_N_REFINE: i64 = 6
41const NX_ARCH_ADALN_EMB:i64 = 7
42const NX_ARCH_N_CHUNKS: i64 = 8
43const NX_ARCH_OUT_DIM: i64 = 9
44const NX_ARCH_XEMB_IN: i64 = 10
45
46func _ar_strlen(s: *u8) -> i64 {
47 var n: i64 = 0
48 while s[n] != (0 as u8) { n = n + 1 }
49 return n
50}
51
52// "model.diffusion_model.<stack>.<n>.<suffix>"
53func _ar_name(out: *u8, stack: *u8, n: i64, suffix: *u8) -> *u8 {
54 let pre: *u8 = "model.diffusion_model." as *u8
55 var o: i64 = 0
56 var i: i64 = 0
57 while pre[i] != (0 as u8) { out[o] = pre[i]; o = o + 1; i = i + 1 }
58 i = 0
59 while stack[i] != (0 as u8) { out[o] = stack[i]; o = o + 1; i = i + 1 }
60 out[o] = 0x2E; o = o + 1
61 let dec: *u8 = sys_mmap(32)
62 let nd: i64 = nx_strconv_format_i64(n, dec)
63 var k: i64 = 0
64 while k < nd { out[o] = dec[k]; o = o + 1; k = k + 1 }
65 out[o] = 0x2E; o = o + 1
66 i = 0
67 while suffix[i] != (0 as u8) { out[o] = suffix[i]; o = o + 1; i = i + 1 }
68 out[o] = 0
69 return out
70}
71
72// Count consecutive stack entries present in the file. Bounded so a malformed name scheme cannot
73// spin; probing beats dividing tensor_count by a guessed tensors-per-block, which would be an
74// assumption wearing a measurement's clothes.
75func _ar_count(gw: *i64, stack: *u8) -> i64 {
76 let nm: *u8 = sys_mmap(256)
77 var n: i64 = 0
78 var go: i64 = 1
79 while go == 1 {
80 _ar_name(nm, stack, n, "attention.qkv.weight" as *u8)
81 if nx_gw_find(gw, nm, _ar_strlen(nm)) < 0 { go = 0 }
82 else { n = n + 1; if n >= 512 { go = 0 } }
83 }
84 return n
85}
86
87func _ar_dim(gw: *i64, name: *u8, which: i64) -> i64 {
88 let idx: i64 = nx_gw_find(gw, name, _ar_strlen(name))
89 if idx < 0 { return 0 - 1 }
90 if which == 0 { return nx_gw_dim0(gw, idx) }
91 return nx_gw_dim1(gw, idx)
92}
93
94// Fill `arch`. Returns 0 on success, or a negative slot-specific verdict -- never a partial fill
95// with plausible defaults, because a default architecture is exactly the failure this file exists
96// to prevent.
97func nx_arch_probe(gw: *i64, arch: *i64) -> i64 {
98 var i: i64 = 0
99 while i < NX_ARCH_SLOTS { arch[i] = 0; i = i + 1 }
100 let nm: *u8 = sys_mmap(256)
101
102 _ar_name(nm, "layers" as *u8, 0, "attention.qkv.weight" as *u8)
103 let D: i64 = _ar_dim(gw, nm, 0)
104 let QW: i64 = _ar_dim(gw, nm, 1)
105 if D <= 0 { return 0 - 1 }
106 if QW <= 0 { return 0 - 2 }
107
108 _ar_name(nm, "layers" as *u8, 0, "attention.q_norm.weight" as *u8)
109 let HD: i64 = _ar_dim(gw, nm, 0)
110 if HD <= 0 { return 0 - 3 }
111
112 // q, k and v are equal width here. VERIFY it divides exactly rather than assume it.
113 if QW - (QW / (3 * HD)) * (3 * HD) != 0 { return 0 - 4 }
114 let NH: i64 = QW / (3 * HD)
115 if NH <= 0 { return 0 - 5 }
116
117 _ar_name(nm, "layers" as *u8, 0, "feed_forward.w1.weight" as *u8)
118 let FD: i64 = _ar_dim(gw, nm, 1)
119 if FD <= 0 { return 0 - 6 }
120
121 _ar_name(nm, "layers" as *u8, 0, "adaLN_modulation.0.weight" as *u8)
122 let AE: i64 = _ar_dim(gw, nm, 0)
123 let AW: i64 = _ar_dim(gw, nm, 1)
124 if AE <= 0 { return 0 - 7 }
125 if AW <= 0 { return 0 - 8 }
126 if AW - (AW / D) * D != 0 { return 0 - 9 }
127
128 let NL: i64 = _ar_count(gw, "layers" as *u8)
129 if NL <= 0 { return 0 - 10 }
130 let NR: i64 = _ar_count(gw, "context_refiner" as *u8)
131
132 let fl: *u8 = "model.diffusion_model.final_layer.linear.weight" as *u8
133 let OD: i64 = _ar_dim(gw, fl, 1)
134 if OD <= 0 { return 0 - 11 }
135 let xe: *u8 = "model.diffusion_model.x_embedder.weight" as *u8
136 let XI: i64 = _ar_dim(gw, xe, 0)
137 if XI <= 0 { return 0 - 12 }
138
139 arch[NX_ARCH_D] = D
140 arch[NX_ARCH_HEAD_DIM] = HD
141 arch[NX_ARCH_N_HEADS] = NH
142 arch[NX_ARCH_QKV_W] = QW
143 arch[NX_ARCH_FFN_DIM] = FD
144 arch[NX_ARCH_N_LAYERS] = NL
145 arch[NX_ARCH_N_REFINE] = NR
146 arch[NX_ARCH_ADALN_EMB] = AE
147 arch[NX_ARCH_N_CHUNKS] = AW / D
148 arch[NX_ARCH_OUT_DIM] = OD
149 arch[NX_ARCH_XEMB_IN] = XI
150 return 0
151}