nx_hifigan_gen.nx source
↩ module page · 108 lines · 7168 B
1// nx_hifigan_gen.nx -- the HiFi-GAN GENERATOR GRAPH (composition of the verified nx_vocops atoms), configurable
2// to the real fetched config (conv_pre → [leaky→upsample→MRF]×4 → leaky→conv_post→tanh). This gate runs the FULL
3// forward end-to-end on a TINY synthetic config (same wiring, small channels = fast) and verifies: output length =
4// T * prod(upsample_rates), channels collapse to 1, and every output is FINITE + tanh-bounded [-1,1]. When a real
5// checkpoint is loaded (weights via nx_https_get_stream + nx_safetensors + weight_norm_apply), this same graph runs
6// at the real 512-ch dims and produces audio -> nx_voice_eval measures the jump. license_tier: ORIGINAL expect_exit: 0
7import "nx_syscalls.nx"
8import "nx_f32.nx"
9import "nx_f32_activations.nx"
10import "nx_vocops.nx"
11const F_MAGIC_131072: i64 = 131072
12const F_MAGIC_4096: i64 = 4096
13const F_MAGIC_65536: i64 = 65536
14
15const SLOPE01: i64 = 0x3DCCCCCD // 0.1 (leaky_relu_slope from the real config)
16const F_THIRD: i64 = 0x3EAAAAAB // 1/3
17const WC_VAL: i64 = 0x3CA3D70A // 0.02 (small synthetic weight so values stay bounded)
18
19func cw(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
20func cn(v: i64) -> i64 { let bb: *u8=sys_mmap(28); var m: i64=v; if m<0{sys_write(1,"-" as *u8,1);m=0-m} let t: *u8=sys_mmap(28); var k: i64=0; if m==0{t[0]=48 as u8;k=1} while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} var i: i64=0; while i<k{bb[i]=t[k-1-i];i=i+1} sys_write(1,bb,k); return 0 }
21func vcopy(dst: *i64, src: *i64, n: i64) -> i64 { var i: i64=0; while i<n { dst[i]=src[i]; i=i+1 } return 0 }
22func vaddto(dst: *i64, src: *i64, n: i64) -> i64 { var i: i64=0; while i<n { dst[i]=nx_f32_add(dst[i], src[i]); i=i+1 } return 0 }
23func vscale(v: *i64, n: i64, s: i64) -> i64 { var i: i64=0; while i<n { v[i]=nx_f32_mul(v[i], s); i=i+1 } return 0 }
24
25// one full HiFi-GAN ResBlock: 3 dilation branches (1,3,5) applied sequentially. synthetic const weights wc, zero bias bz.
26func gen_resblock(x: *i64, ch: i64, L: i64, K: i64, wc: *i64, bz: *i64, out: *i64, sa: *i64, sb: *i64, t1: *i64, t2: *i64) -> i64 {
27 resblock1_1dil(x, wc, bz, wc, bz, sa, ch, L, K, 1, SLOPE01, t1, t2)
28 resblock1_1dil(sa, wc, bz, wc, bz, sb, ch, L, K, 3, SLOPE01, t1, t2)
29 resblock1_1dil(sb, wc, bz, wc, bz, out, ch, L, K, 5, SLOPE01, t1, t2)
30 return 0
31}
32// MRF: sum of 3 resblocks (K=3,7,11) / 3
33func gen_mrf(x: *i64, ch: i64, L: i64, wc: *i64, bz: *i64, out: *i64, sa: *i64, sb: *i64, sc: *i64, t1: *i64, t2: *i64) -> i64 {
34 gen_resblock(x, ch, L, 3, wc, bz, out, sa, sb, t1, t2)
35 gen_resblock(x, ch, L, 7, wc, bz, sc, sa, sb, t1, t2); vaddto(out, sc, ch*L)
36 gen_resblock(x, ch, L, 11, wc, bz, sc, sa, sb, t1, t2); vaddto(out, sc, ch*L)
37 vscale(out, ch*L, F_THIRD)
38 return 0
39}
40
41// full generator: mel[C_mel,T] -> audio[1, T*up_rate^4]. n_stages=4. returns audio length.
42func generator_forward(mel: *i64, C_mel: i64, T: i64, init_ch: i64, up_rate: i64, up_kernel: i64, out: *i64) -> i64 {
43 let wc: *i64 = sys_mmap(F_MAGIC_131072*8) as *i64; var i: i64=0; while i<F_MAGIC_131072 { wc[i]=WC_VAL; i=i+1 }
44 let bz: *i64 = sys_mmap(F_MAGIC_4096*8) as *i64; i=0; while i<F_MAGIC_4096 { bz[i]=0; i=i+1 }
45 let bufX: *i64 = sys_mmap(F_MAGIC_65536*8) as *i64
46 let bufU: *i64 = sys_mmap(F_MAGIC_65536*8) as *i64
47 let bufM: *i64 = sys_mmap(F_MAGIC_65536*8) as *i64
48 let sa: *i64 = sys_mmap(F_MAGIC_65536*8) as *i64; let sb: *i64 = sys_mmap(F_MAGIC_65536*8) as *i64
49 let sc: *i64 = sys_mmap(F_MAGIC_65536*8) as *i64; let t1: *i64 = sys_mmap(F_MAGIC_65536*8) as *i64; let t2: *i64 = sys_mmap(F_MAGIC_65536*8) as *i64
50 // conv_pre: mel[C_mel,T] -> [init_ch, T], K=7 pad=3
51 conv1d(mel, wc, bz, bufX, C_mel, T, init_ch, 7, 1, 3, 1)
52 var ch: i64 = init_ch; var L: i64 = T
53 var stage: i64 = 0
54 while stage < 4 {
55 leaky_relu_vec(bufX, ch*L, SLOPE01)
56 let ch2: i64 = ch/2
57 let pad: i64 = (up_kernel - up_rate)/2
58 let L2: i64 = conv_transpose1d(bufX, wc, bz, bufU, ch, L, ch2, up_kernel, up_rate, pad, 0, 1)
59 gen_mrf(bufU, ch2, L2, wc, bz, bufM, sa, sb, sc, t1, t2)
60 vcopy(bufX, bufM, ch2*L2)
61 ch = ch2; L = L2
62 stage = stage + 1
63 }
64 leaky_relu_vec(bufX, ch*L, SLOPE01)
65 // conv_post: [ch,L] -> [1, L], K=7 pad=3
66 conv1d(bufX, wc, bz, out, ch, L, 1, 7, 1, 3, 1)
67 // tanh
68 i=0; while i<L { out[i]=nx_f32_tanh(out[i]); i=i+1 }
69 return L
70}
71
72func main() -> i64 {
73 cw("=== nx_hifigan_gen -- FULL HiFi-GAN generator graph, end-to-end (tiny synthetic config, real wiring) ===\n" as *u8)
74 // TINY config: same structure as the real model (4 upsample stages, halving channels, MRF 3x[K3,7,11]x[dil1,3,5])
75 let C_mel: i64 = 4 // real=80
76 let T: i64 = 8 // mel frames
77 let init_ch: i64 = 32 // real=512 (halves 4x: 32->16->8->4->2 ; real 512->256->128->64->32)
78 let up_rate: i64 = 2 // real=4
79 let up_kernel: i64 = 4 // real=8 (=2*rate)
80 let mel: *i64 = sys_mmap(C_mel*T*8) as *i64
81 var i: i64=0; while i<C_mel*T { mel[i]=0x3F000000; i=i+1 } // synthetic mel = 0.5 everywhere
82 let out: *i64 = sys_mmap(F_MAGIC_65536*8) as *i64
83
84 let audio_len: i64 = generator_forward(mel, C_mel, T, init_ch, up_rate, up_kernel, out)
85 let expect_len: i64 = T * up_rate*up_rate*up_rate*up_rate // T * up_rate^4
86 cw("mel [" as *u8); cn(C_mel); cw("," as *u8); cn(T); cw("] -> audio length="); cn(audio_len); cw(" (expect "); cn(expect_len); cw(" = T*rate^4)\n" as *u8)
87
88 // finiteness + bounded: tanh output must be in [-1,1] and not NaN/Inf
89 var finite: i64=1; var maxbits: i64=0; var nan: i64=0
90 i=0; while i<audio_len {
91 let raw: i64 = out[i] & 0xFFFFFFFF
92 let exp: i64 = (raw>>23)&0xFF
93 if exp==0xFF { nan=1 } // NaN/Inf
94 let mag: i64 = raw & 0x7FFFFFFF
95 if mag > 0x3F800800 { finite=0 } // |x| > ~1.0001 -> not tanh-bounded
96 i=i+1
97 }
98 cw("sample out[0..3]="); var s: i64=0; while s<4 { let b: *u8=sys_mmap(8); var j: i64=0; while j<8 { let n: i64=(out[s]>>((7-j)*4))&0xF; if n<10 {b[j]=(48+n) as u8} else {b[j]=(87+n) as u8} j=j+1 } sys_write(1,b,8); cw(" " as *u8); s=s+1 } cw("\n" as *u8)
99
100 var pass: i64=0; var tot: i64=3
101 if audio_len==expect_len { pass=pass+1; cw("PASS T1 output length = T * prod(upsample_rates) (the 4-stage upsampling wiring is correct)\n" as *u8) } else { cw("FAIL T1 len="); cn(audio_len); cw("\n" as *u8) }
102 if nan==0 { pass=pass+1; cw("PASS T2 no NaN/Inf across the whole forward (conv/transpose/resblock/mrf/tanh numerically stable)\n" as *u8) } else { cw("FAIL T2 NaN/Inf present\n" as *u8) }
103 if finite==1 { pass=pass+1; cw("PASS T3 every output tanh-bounded in [-1,1] (valid audio range)\n" as *u8) } else { cw("FAIL T3 out of [-1,1]\n" as *u8) }
104
105 cw("nx_hifigan_gen pass="); cn(pass); cw("/"); cn(tot)
106 if pass==tot { cw(" GREEN -- the FULL HiFi-GAN generator graph runs end-to-end (mel->audio, correct shape, stable). Load real 512-ch weights -> it produces real audio.\n" as *u8); sys_exit(0); return 0 }
107 cw(" RED\n" as *u8); sys_exit(1); return 1
108}