nx_vitpose_inspect.nx source
↩ module page · 50 lines · 2943 B
1// nx_vitpose_inspect.nx -- dump the safetensors HEADER (tensor manifest) of the downloaded model to a file, WITHOUT
2// reading the multi-hundred-MB body: read 8-byte u64le header_len, then header_len bytes of JSON -> write to
3// models/vitpose_header.json. The manifest (every tensor name + dtype + shape + data_offsets) is the assembly
4// blueprint. Reads the model path from data/mp_fetch_dest.txt. expect_exit: 0
5import "nx_syscalls.nx"
6const K_MAGIC_4096: i64 = 4096
7const K_MAGIC_8388608: i64 = 8388608
8
9func w(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
10func wn(v: i64) -> i64 { var m: i64=v; if m<0{w("-" as *u8);m=0-m} let t:*u8=sys_mmap(24); 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; let o:*u8=sys_mmap(24); while i<k{o[i]=t[k-1-i];i=i+1} sys_write(1,o,k); return 0 }
11func u64le(buf: *u8, off: i64) -> i64 { var v: i64=0; var i: i64=0; while i<8 { v = v | ((buf[off+i]&0xff) << (i*8)); i=i+1 } return v }
12
13func readfile(path: *u8, buf: *u8, cap: i64) -> i64 {
14 let fd: i64 = sys_openat_rd(path)
15 if fd < 0 { return 0 - 1 }
16 var n: i64 = sys_read(fd, buf, cap - 1)
17 sys_close(fd)
18 if n < 0 { return 0 - 1 }
19 var go: i64 = 1
20 while go == 1 { go = 0; if n > 0 { let c: i64 = buf[n-1] as i64; if c == 10 { n=n-1; go=1 } if c == 13 { n=n-1; go=1 } if c == 32 { n=n-1; go=1 } } }
21 buf[n] = 0 as u8
22 return n
23}
24
25func main() -> i64 {
26 let path: *u8 = sys_mmap(K_MAGIC_4096)
27 if readfile("data/mp_fetch_dest.txt" as *u8, path, K_MAGIC_4096) < 0 { w("no dest file\n" as *u8); sys_exit(1); return 1 }
28 let fd: i64 = sys_openat_rd(path)
29 if fd < 0 { w("cannot open model: " as *u8); w(path); w("\n" as *u8); sys_exit(1); return 1 }
30 let hb: *u8 = sys_mmap(16)
31 var r8: i64 = sys_read(fd, hb, 8)
32 if r8 != 8 { w("short header-len read\n" as *u8); sys_close(fd); sys_exit(1); return 1 }
33 let hlen: i64 = u64le(hb, 0)
34 w("header_len=" as *u8); wn(hlen); w("\n" as *u8)
35 if hlen <= 0 { sys_close(fd); sys_exit(1); return 1 }
36 if hlen > K_MAGIC_8388608 { w("header suspiciously large\n" as *u8); sys_close(fd); sys_exit(1); return 1 }
37 let hdr: *u8 = sys_mmap(hlen + 16)
38 var got: i64 = 0
39 while got < hlen { let n: i64 = sys_read(fd, ((hdr as i64)+got) as *u8, hlen - got); if n <= 0 { got = hlen } else { got = got + n } }
40 sys_close(fd)
41 // write the header JSON to models/vitpose_header.json
42 let ofd: i64 = sys_openat_wr("/mnt/c/Users/elder/nishi-core/models/vitpose_header.json" as *u8, 0x1a4)
43 if ofd < 0 { w("cannot open output json\n" as *u8); sys_exit(1); return 1 }
44 var wr: i64 = 0
45 while wr < hlen { let n: i64 = sys_write(ofd, ((hdr as i64)+wr) as *u8, hlen - wr); if n <= 0 { wr = hlen } else { wr = wr + n } }
46 sys_close(ofd)
47 w("wrote models/vitpose_header.json (" as *u8); wn(hlen); w(" bytes)\n" as *u8)
48 sys_exit(0)
49 return 0
50}