`nx_zip_read.nx` source

↩ module page · 101 lines · 6108 B
1// nx_zip_read.nx -- SOVEREIGN ZIP reader (the first half of the PyTorch .bin loader). A PyTorch .bin is a ZIP
2// archive of STORED (uncompressed) members: archive/data.pkl (the pickle describing the state_dict) + archive/
3// data/0, data/1, ... (raw tensor blobs) + archive/version. This finds the End-Of-Central-Directory record,
4// walks the central directory, and lists every member (name / method / size / local-header offset), plus locates
5// a member's raw data (STORED = data begins right after the local header). Tested on the REAL 50.7MB HiFi-GAN
6// pytorch_model.bin just downloaded. Next: a minimal pickle-VM over data.pkl. license_tier: ORIGINAL  expect_exit: 0
7import "nx_syscalls.nx"
8const K_MAGIC_67108864: i64 = 67108864
9
10func p(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
11func pn(v: i64) -> i64 { if v==0 { sys_write(1,"0" as *u8,1); return 0 } var m: i64=v; if m<0{sys_write(1,"-" as *u8,1);m=0-m} let d: *u8=sys_mmap(24); var k: i64=0; while m>0 {d[k]=(48+(m%10)) as u8; m=m/10; k=k+1} var i: i64=k-1; while i>=0 {let o: *u8=sys_mmap(1); o[0]=d[i]; sys_write(1,o,1); i=i-1} return 0 }
12func pc(c: i64) -> i64 { let o: *u8=sys_mmap(1); o[0]=c as u8; sys_write(1,o,1); return 0 }
13
14func u16le(b: *u8, o: i64) -> i64 { return (b[o]&0xff) | ((b[o+1]&0xff)<<8) }
15func u32le(b: *u8, o: i64) -> i64 { return (b[o]&0xff) | ((b[o+1]&0xff)<<8) | ((b[o+2]&0xff)<<16) | ((b[o+3]&0xff)<<24) }
16
17// find End-Of-Central-Directory (sig 0x06054b50 = PK\5\6), scanning backward from the end.
18func find_eocd(b: *u8, n: i64) -> i64 {
  var i: i64 = n - 22
  while i >= 0 {
      if b[i]==0x50 as u8 { if b[i+1]==0x4B as u8 { if b[i+2]==0x05 as u8 { if b[i+3]==0x06 as u8 { return i } } } }
      i = i - 1
  }
  return 0 - 1
25}
26// raw data offset of a STORED member given its local-header offset: skip local header (30) + fname + extra.
27func member_data_off(b: *u8, lho: i64) -> i64 {
  let fnl: i64 = u16le(b, lho + 26)
  let exl: i64 = u16le(b, lho + 28)
  return lho + 30 + fnl + exl
31}
32
33func main() -> i64 {
  p("=== nx_zip_read -- sovereign ZIP reader over the REAL HiFi-GAN pytorch_model.bin ===\n" as *u8)
  let path: *u8 = "/home/elderwesto/nx_stage/hifigan_pytorch_model.bin" as *u8
  let fd: i64 = sys_openat_rd(path)
  if fd < 0 { p("open failed (download the weights first)\n" as *u8); sys_exit(1); return 1 }
  let cap: i64 = K_MAGIC_67108864
  let buf: *u8 = sys_mmap(cap)
  var n: i64 = 0
  var go: i64 = 1
  while go == 1 {
      let r: i64 = sys_read(fd, (buf as i64 + n) as *u8, cap - n)
      if r <= 0 { go = 0 } else { n = n + r; if n >= cap { go = 0 } }
  }
  sys_close(fd)
  p("read "); pn(n); p(" bytes\n" as *u8)
48
  let eocd: i64 = find_eocd(buf, n)
  if eocd < 0 { p("verdict=FAIL (no EOCD -- not a ZIP)\n" as *u8); sys_exit(1); return 1 }
  let n_entries: i64 = u16le(buf, eocd + 10)
  let cd_size: i64 = u32le(buf, eocd + 12)
  let cd_off: i64 = u32le(buf, eocd + 16)
  p("EOCD @"); pn(eocd); p(": entries="); pn(n_entries); p(" cd_off="); pn(cd_off); p(" cd_size="); pn(cd_size); p("\n" as *u8)
55
  // walk the central directory
  var off: i64 = cd_off
  var e: i64 = 0
  var have_pkl: i64 = 0
  var have_data0: i64 = 0
  var stored_all: i64 = 1
  var pkl_lho: i64 = 0
  while e < n_entries {
      if buf[off]!=0x50 as u8 { p("  (bad CD sig at "); pn(off); p(")\n" as *u8); e = n_entries }
      else {
          let method: i64 = u16le(buf, off + 10)
          let usize:  i64 = u32le(buf, off + 24)
          let fnl:    i64 = u16le(buf, off + 28)
          let exl:    i64 = u16le(buf, off + 30)
          let cml:    i64 = u16le(buf, off + 32)
          let lho:    i64 = u32le(buf, off + 42)
          // print filename
          p("  ["); pn(e); p("] "); var j: i64=0; while j<fnl { pc(buf[off+46+j] as i64); j=j+1 }
          p("  method="); pn(method); p(" size="); pn(usize); p(" lho="); pn(lho); p("\n" as *u8)
          if method != 0 { stored_all = 0 }
          // detect data.pkl (ends with "data.pkl") and data/0
          if fnl >= 8 { var m: i64=1; var q: i64=0; let suf: *u8 = "data.pkl" as *u8; while q<8 { if buf[off+46+fnl-8+q]!=suf[q] { m=0 } q=q+1 } if m==1 { have_pkl=1; pkl_lho=lho } }
          if fnl >= 6 { var m2: i64=1; var q2: i64=0; let suf2: *u8 = "data/0" as *u8; while q2<6 { if buf[off+46+fnl-6+q2]!=suf2[q2] { m2=0 } q2=q2+1 } if m2==1 { have_data0=1 } }
          off = off + 46 + fnl + exl + cml
          e = e + 1
      }
  }
83
  // extract data.pkl start (should begin with pickle PROTO opcode 0x80)
  var pkl_ok: i64 = 0
  if have_pkl == 1 {
      let dof: i64 = member_data_off(buf, pkl_lho)
      p("data.pkl raw @"); pn(dof); p(" first bytes: "); var j: i64=0; while j<4 { let bb: *u8=sys_mmap(4); var nib: i64=(buf[dof+j]>>4)&0xF; if nib<10 {bb[0]=(48+nib) as u8} else {bb[0]=(87+nib) as u8} nib=buf[dof+j]&0xF; if nib<10 {bb[1]=(48+nib) as u8} else {bb[1]=(87+nib) as u8} bb[2]=32 as u8; sys_write(1,bb,3); j=j+1 } p("\n" as *u8)
      if buf[dof]==0x80 as u8 { pkl_ok = 1 }      // pickle PROTO opcode
  }
91
  var pass: i64=0; var tot: i64=4
  if n_entries > 0 { pass=pass+1; p("PASS ZIP central directory parsed ("); pn(n_entries); p(" members)\n" as *u8) } else { p("FAIL no entries\n" as *u8) }
  if have_pkl == 1 { pass=pass+1; p("PASS archive/data.pkl located (the state_dict pickle)\n" as *u8) } else { p("FAIL no data.pkl\n" as *u8) }
  if have_data0 == 1 { pass=pass+1; p("PASS archive/data/0 located (a raw tensor blob)\n" as *u8) } else { p("FAIL no data/0\n" as *u8) }
  if pkl_ok == 1 { pass=pass+1; p("PASS data.pkl begins with pickle PROTO (0x80) -> ready for the pickle-VM\n" as *u8) } else { p("FAIL data.pkl not a pickle\n" as *u8) }
97
  p("nx_zip_read pass="); pn(pass); p("/"); pn(tot)
  if pass==tot { p("  GREEN -- ZIP reader works on the real weights; members + tensor blobs + pickle located. Next: minimal pickle-VM over data.pkl -> state_dict.\n" as *u8); sys_exit(0); return 0 }
  p("  RED\n" as *u8); sys_exit(1); return 1
101}
nx_zip_read.nx source

`nx_zip_read.nx` source
