nx_genrole.nx source
↩ module page · 183 lines · 7948 B
1// nx_genrole.nx -- THE single answer to "what is this file". One classifier, three consumers.
2//
3// This exists because the rule had already been written twice: nx_gen_modelid's main() and
4// nx_gen_zoo's zo_classify() each carried their own copy of the tensor-name probes. The zoo's own
5// header even claimed it "reuses nx_gen_modelid's probes rather than re-deriving them, because two
6// copies of a classification rule drift and the drift is silent" -- and then re-derived them.
7// ★★★★★ A COMMENT ASSERTING DRY IS NOT DRY. The second copy was already missing the adapter-family
8// resolution (kohya vs peft prefixes) that modelid had learned, so a LoRA the front door named
9// "z-image" the zoo named only "adapter-lora(kohya)". They had ALREADY drifted.
10//
11// ★ FILENAME AND EXTENSION ARE CLAIMS; THE TENSOR NAMES ARE EVIDENCE. Probing names is the only
12// honest signature: a .safetensors DiT can be Z-Image or FLUX, a .gguf can be a DiT or an LLM, and
13// a file sitting in the diffusion/ folder can be an adapter that is not a model at all.
14//
15// It REFUSES to guess. UNKNOWN is a verdict, never defaulted to the family we happen to support --
16// a wrong family runs to completion and produces a wrong picture with no error anywhere.
17// license_tier: ORIGINAL
18
19import "nx_syscalls.nx"
20import "nx_le.nx"
21import "nx_f32.nx"
22import "nx_f32_div.nx"
23import "nx_f32_cvt.nx"
24import "nx_f16.nx"
25import "nx_strconv.nx"
26import "nx_genweights.nx"
27import "nx_genarch.nx"
28
29// role
30const ROLE_UNKNOWN: i64 = 0
31const ROLE_DIT: i64 = 1
32const ROLE_VAE_DEC: i64 = 2
33const ROLE_TEXTENC: i64 = 3
34const ROLE_ADAPTER: i64 = 4
35// family
36const FAM_UNKNOWN: i64 = 0
37const FAM_ZIMAGE: i64 = 1
38const FAM_FLUX: i64 = 2
39const FAM_TAESD: i64 = 3
40const FAM_SDVAE: i64 = 4
41const FAM_LLMGGUF: i64 = 5
42// adapter kind
43const ADK_NONE: i64 = 0
44const ADK_KOHYA: i64 = 1
45const ADK_PEFT: i64 = 2
46const ADK_LOKR: i64 = 3
47
48// out[] slots
49const RL_ROLE: i64 = 0
50const RL_FAMILY: i64 = 1
51const RL_RUNNABLE: i64 = 2
52const RL_ADKIND: i64 = 3
53const RL_SLOTS: i64 = 8
54
55func _rl_len(s: *u8) -> i64 {
56 var n: i64 = 0
57 while s[n] != (0 as u8) { n = n + 1 }
58 return n
59}
60func nx_role_has(gw: *i64, name: *u8) -> i64 {
61 if nx_gw_find(gw, name, _rl_len(name)) < 0 { return 0 }
62 return 1
63}
64
65func nx_role_name(r: i64) -> *u8 {
66 if r == ROLE_DIT { return "dit" as *u8 }
67 if r == ROLE_VAE_DEC { return "vae-decoder" as *u8 }
68 if r == ROLE_TEXTENC { return "text-encoder" as *u8 }
69 if r == ROLE_ADAPTER { return "adapter" as *u8 }
70 return "unknown" as *u8
71}
72func nx_family_name(f: i64) -> *u8 {
73 if f == FAM_ZIMAGE { return "z-image" as *u8 }
74 if f == FAM_FLUX { return "flux" as *u8 }
75 if f == FAM_TAESD { return "taesd" as *u8 }
76 if f == FAM_SDVAE { return "sd-vae" as *u8 }
77 if f == FAM_LLMGGUF { return "llm-gguf" as *u8 }
78 return "unknown" as *u8
79}
80func nx_adkind_name(k: i64) -> *u8 {
81 if k == ADK_KOHYA { return "lora(kohya)" as *u8 }
82 if k == ADK_PEFT { return "lora(peft)" as *u8 }
83 if k == ADK_LOKR { return "lokr" as *u8 }
84 return "none" as *u8
85}
86
87// Which family does an ADAPTER target? An adapter carries the BASE model's layer names, so the
88// same probes that identify a checkpoint identify what the adapter was trained against.
89// ⚠ Probe the STACK, not one module: an adapter only carries the layers it was trained on, so
90// "layers.0.attention.qkv" misses one that starts at adaLN_modulation.
91// ★ ONE ROLE CAN HAVE SEVERAL NAMING CONVENTIONS; PROBE FOR EACH, DO NOT ASSUME ONE WON.
92func _rl_adapter_family(gw: *i64) -> i64 {
93 if nx_role_has(gw, "diffusion_model.double_blocks.0." as *u8) == 1 { return FAM_FLUX }
94 if nx_role_has(gw, "diffusion_model.layers.0." as *u8) == 1 { return FAM_ZIMAGE }
95 if nx_role_has(gw, "lora_unet_double_blocks_0_" as *u8) == 1 { return FAM_FLUX }
96 if nx_role_has(gw, "lora_unet_layers_0_" as *u8) == 1 { return FAM_ZIMAGE }
97 if nx_role_has(gw, "lora_unet_context_refiner_0_" as *u8) == 1 { return FAM_ZIMAGE }
98 return FAM_UNKNOWN
99}
100
101// Classify an already-opened container. Fills out[RL_*]. Returns the role.
102func nx_role_probe(gw: *i64, out: *i64) -> i64 {
103 var i: i64 = 0
104 while i < RL_SLOTS { out[i] = 0; i = i + 1 }
105
106 // ---- ADAPTERS FIRST ----
107 // A LoRA carries the same layer names as its base, so probing for a base architecture first
108 // classifies an adapter as a checkpoint. Realism_Engine_Klein_V2 sits in the diffusion/ folder
109 // and is not a model at all: it is a LoKr adapter for a Flux base.
110 // ★ A FILE IN THE MODELS FOLDER IS NOT NECESSARILY A MODEL.
111 var adk: i64 = ADK_NONE
112 if nx_role_has(gw, ".lokr_w1" as *u8) == 1 { adk = ADK_LOKR }
113 if nx_role_has(gw, ".lora_A.weight" as *u8) == 1 { adk = ADK_PEFT }
114 if nx_role_has(gw, ".lora_down.weight" as *u8) == 1 { adk = ADK_KOHYA }
115 if adk != ADK_NONE {
116 out[RL_ROLE] = ROLE_ADAPTER
117 out[RL_FAMILY] = _rl_adapter_family(gw)
118 out[RL_ADKIND] = adk
119 // LoRA folds sovereignly (nx_genlora); LoKr does not decompose the same way and is not
120 // implemented -- reporting it runnable would fold nothing and silently change no pixels.
121 if adk == ADK_LOKR { out[RL_RUNNABLE] = 0 } else { out[RL_RUNNABLE] = 1 }
122 return ROLE_ADAPTER
123 }
124
125 // ---- checkpoints, most specific first ----
126 if nx_role_has(gw, "model.diffusion_model.layers.0.attention.qkv.weight" as *u8) == 1 {
127 out[RL_ROLE] = ROLE_DIT
128 out[RL_FAMILY] = FAM_ZIMAGE
129 // cap_embedder is what distinguishes a real Z-Image DiT from a look-alike stack, and the
130 // architecture must be DERIVABLE -- a DiT whose shapes we cannot read is not runnable.
131 if nx_role_has(gw, "model.diffusion_model.cap_embedder.1.weight" as *u8) == 1 {
132 let arch: *i64 = sys_mmap(NX_ARCH_SLOTS * 8 + 64) as *i64
133 if nx_arch_probe(gw, arch) == 0 { out[RL_RUNNABLE] = 1 }
134 }
135 return ROLE_DIT
136 }
137 if nx_role_has(gw, "diffusion_model.double_blocks.0.img_attn.qkv.weight" as *u8) == 1 {
138 out[RL_ROLE] = ROLE_DIT
139 out[RL_FAMILY] = FAM_FLUX
140 return ROLE_DIT
141 }
142 if nx_role_has(gw, "model.diffusion_model.double_blocks.0.img_attn.qkv.weight" as *u8) == 1 {
143 out[RL_ROLE] = ROLE_DIT
144 out[RL_FAMILY] = FAM_FLUX
145 return ROLE_DIT
146 }
147 // TAESD: a tiny all-conv decoder; layer 2 is a residual TAEBlock
148 if nx_role_has(gw, "decoder.layers.2.conv.0.weight" as *u8) == 1 {
149 out[RL_ROLE] = ROLE_VAE_DEC
150 out[RL_FAMILY] = FAM_TAESD
151 out[RL_RUNNABLE] = 1
152 return ROLE_VAE_DEC
153 }
154 if nx_role_has(gw, "decoder.conv_in.weight" as *u8) == 1 {
155 out[RL_ROLE] = ROLE_VAE_DEC
156 out[RL_FAMILY] = FAM_SDVAE
157 return ROLE_VAE_DEC
158 }
159 if nx_role_has(gw, "token_embd.weight" as *u8) == 1 {
160 out[RL_ROLE] = ROLE_TEXTENC
161 out[RL_FAMILY] = FAM_LLMGGUF
162 // third-party by design today, but PRESENT and usable -- runnable is about the pipeline
163 // being able to execute, not about who wrote the weights.
164 out[RL_RUNNABLE] = 1
165 return ROLE_TEXTENC
166 }
167 return ROLE_UNKNOWN
168}
169
170// Why is this part not runnable? Named reasons, because a bare NO is not actionable.
171func nx_role_blocker(out: *i64) -> *u8 {
172 if out[RL_RUNNABLE] == 1 { return "" as *u8 }
173 if out[RL_ROLE] == ROLE_ADAPTER {
174 if out[RL_ADKIND] == ADK_LOKR { return "lokr fold not implemented" as *u8 }
175 return "adapter family unresolved" as *u8
176 }
177 if out[RL_ROLE] == ROLE_DIT {
178 if out[RL_FAMILY] == FAM_FLUX { return "flux double/single-block arch not implemented" as *u8 }
179 return "z-image arch not derivable (no cap_embedder or unreadable shapes)" as *u8
180 }
181 if out[RL_ROLE] == ROLE_VAE_DEC { return "sd-vae conv+attn decoder not implemented" as *u8 }
182 return "file did not match any known signature" as *u8
183}