nx_f32_vae_decode_tiny.nx source
↩ module page · 118 lines · 5333 B
1// nx_f32_vae_decode_tiny.nx -- first end-to-end sovereign f32 VAE-DECODE (latent -> RGB), the R4/R7
2// integration milestone of the sd-server -> Nishi migration.
3//
4// Composes the now-gated sovereign f32 bricks into the canonical VAE-decoder shape:
5// latent[1,Cz,Hl,Wl]
6// -> conv_in (Cz->Cm, 3x3 pad1) nx_f32_conv2d
7// -> ResBlock (Cm) nx_f32_resblock (GN->SiLU->Conv x2 + residual)
8// -> upsample (2x) nx_f32_upsample
9// -> conv_out (Cm->3, 3x3 pad1) nx_f32_conv2d
10// -> RGB[1,3,2*Hl,2*Wl]
11//
12// This is the "decode half" of txt2img, proven to COMPOSE + run end-to-end + be deterministic on the
13// software-f32 tier with no third-party. The full Flux/Z-Image decoder is this shape scaled up (16-ch
14// latent, more ResBlocks, 3 upsample stages) with REAL dequantized GGUF weights -- the next rungs wire
15// nx_gguf_load_model + nx_q4k_to_f32 in. Here the weights are fixed small values to prove the pipeline.
16// license_tier: ORIGINAL
17import "nx_syscalls.nx"
18import "nx_f32.nx"
19import "nx_f32_div.nx"
20import "nx_f32_cvt.nx"
21import "nx_f32_conv2d.nx"
22import "nx_f32_resblock.nx"
23import "nx_f32_upsample.nx"
24
25func nx_f32vd_fill(buf: *i64, n: i64, val: i64) -> i64 {
26 var i: i64 = 0
27 while i < n { buf[i] = val; i = i + 1 }
28 return 0
29}
30
31// Run the tiny decode once: latent -> rgb. Returns 0 ok, else a stage-coded error (100s conv_in,
32// 200s resblock, 400s conv_out). Cz/Cm/Hl/Wl fixed by the caller's buffers; G divides Cm.
33func nx_f32_vae_decode_tiny_run(latent: *i64, Cz: i64, Cm: i64, Hl: i64, Wl: i64, G: i64,
34 conv_in_w: *i64, c1w: *i64, c2w: *i64,
35 g1: *i64, b1: *i64, g2: *i64, b2: *i64,
36 conv_out_w: *i64,
37 mid1: *i64, mid2: *i64, up: *i64, scr: *i64, rgb: *i64) -> i64 {
38 // conv_in: [1,Cz,Hl,Wl] -> [1,Cm,Hl,Wl]
39 let e1: i64 = nx_f32_conv2d_forward(latent, 1, Cz, Hl, Wl, conv_in_w, Cm, 3, 3, 1, 1, 0 as *i64, mid1)
40 if e1 != NX_F32CV_OK { return 100 + e1 }
41 // ResBlock(Cm) -> mid2
42 let e2: i64 = nx_f32_resblock_forward(mid1, 1, Cm, Hl, Wl, G, g1, b1, c1w, 0 as *i64, g2, b2, c2w, 0 as *i64, mid2, scr)
43 if e2 != NX_F32RB_OK { return 200 + e2 }
44 // upsample 2x: [1,Cm,Hl,Wl] -> [1,Cm,2Hl,2Wl]
45 let e3: i64 = nx_f32_upsample_nn(mid2, 1, Cm, Hl, Wl, 2, up)
46 if e3 != NX_F32US_OK { return 300 + e3 }
47 // conv_out: [1,Cm,2Hl,2Wl] -> [1,3,2Hl,2Wl]
48 let e4: i64 = nx_f32_conv2d_forward(up, 1, Cm, Hl * 2, Wl * 2, conv_out_w, 3, 3, 3, 1, 1, 0 as *i64, rgb)
49 if e4 != NX_F32CV_OK { return 400 + e4 }
50 return 0
51}
52
53// ===== Self-test (inline integration gate) ========================
54// (1) the full decode runs end-to-end with no stage error
55// (2) DETERMINISTIC: two runs on the same input/weights -> bit-identical RGB
56// (3) SIGNAL: the output is not all-zero (data actually flowed through every stage)
57func main() -> i64 {
58 let Cz: i64 = 4
59 let Cm: i64 = 8
60 let Hl: i64 = 2
61 let Wl: i64 = 2
62 let G: i64 = 4 // divides Cm=8 (C_per_group=2)
63 let H2: i64 = Hl * 2
64 let W2: i64 = Wl * 2
65 let rgb_n: i64 = 3 * H2 * W2 // 48
66
67 let latent: *i64 = sys_mmap(Cz * Hl * Wl * 8) as *i64
68 let conv_in_w: *i64 = sys_mmap(Cm * Cz * 9 * 8) as *i64
69 let c1w: *i64 = sys_mmap(Cm * Cm * 9 * 8) as *i64
70 let c2w: *i64 = sys_mmap(Cm * Cm * 9 * 8) as *i64
71 let g1: *i64 = sys_mmap(Cm * 8) as *i64
72 let b1: *i64 = sys_mmap(Cm * 8) as *i64
73 let g2: *i64 = sys_mmap(Cm * 8) as *i64
74 let b2: *i64 = sys_mmap(Cm * 8) as *i64
75 let conv_out_w: *i64 = sys_mmap(3 * Cm * 9 * 8) as *i64
76 let mid1: *i64 = sys_mmap(Cm * Hl * Wl * 8) as *i64
77 let mid2: *i64 = sys_mmap(Cm * Hl * Wl * 8) as *i64
78 let up: *i64 = sys_mmap(Cm * H2 * W2 * 8) as *i64
79 let scr: *i64 = sys_mmap(Cm * H2 * W2 * 8) as *i64
80 let rgb1: *i64 = sys_mmap(rgb_n * 8) as *i64
81 let rgb2: *i64 = sys_mmap(rgb_n * 8) as *i64
82
83 let one: i64 = nx_i32_to_f32(1)
84 let p1: i64 = nx_f32_div(one, nx_i32_to_f32(10)) // 0.1
85
86 // fixed weights: small convs, unit GroupNorm scale, zero shift
87 nx_f32vd_fill(conv_in_w, Cm * Cz * 9, p1)
88 nx_f32vd_fill(c1w, Cm * Cm * 9, p1)
89 nx_f32vd_fill(c2w, Cm * Cm * 9, p1)
90 nx_f32vd_fill(conv_out_w, 3 * Cm * 9, p1)
91 nx_f32vd_fill(g1, Cm, one)
92 nx_f32vd_fill(b1, Cm, 0)
93 nx_f32vd_fill(g2, Cm, one)
94 nx_f32vd_fill(b2, Cm, 0)
95 // latent values 1..(Cz*Hl*Wl)
96 var i: i64 = 0
97 while i < Cz * Hl * Wl { latent[i] = nx_i32_to_f32(i + 1); i = i + 1 }
98
99 // (1) run end-to-end -> rgb1
100 let r1: i64 = nx_f32_vae_decode_tiny_run(latent, Cz, Cm, Hl, Wl, G,
101 conv_in_w, c1w, c2w, g1, b1, g2, b2, conv_out_w, mid1, mid2, up, scr, rgb1)
102 if r1 != 0 { return r1 }
103
104 // (2) run again -> rgb2 ; deterministic => bit-identical
105 let r2: i64 = nx_f32_vae_decode_tiny_run(latent, Cz, Cm, Hl, Wl, G,
106 conv_in_w, c1w, c2w, g1, b1, g2, b2, conv_out_w, mid1, mid2, up, scr, rgb2)
107 if r2 != 0 { return 50 }
108 i = 0
109 while i < rgb_n { if rgb1[i] != rgb2[i] { return 60 } i = i + 1 }
110
111 // (3) signal: output not all-zero
112 var nz: i64 = 0
113 i = 0
114 while i < rgb_n { if rgb1[i] != 0 { nz = 1 } i = i + 1 }
115 if nz == 0 { return 70 }
116
117 return 0
118}