nx_safetensors_load.nx source

1// nx_safetensors_load.nx -- the STANDARD, reusable safetensors -> f32 LOADER library (importable, no main), the 2// model-port driver's stage-4/5 for safetensors models (ViTPose etc.). The existing nx_safetensors.nx is a 3// main()-based demo (bounded-header-read proof) and cannot be imported; this is the clean lib the port + any HF 4// safetensors model uses. Format: [0..8) u64le header_len N; [8..8+N) JSON header {"tname":{"dtype","shape", 5// "data_offsets":[a,b]},...}; [8+N..) data, tensor bytes at [8+N+a, 8+N+b). Reads F32 (bits direct), F16 6// (nx_f16_to_f32), BF16 (= high 16 bits of an f32 -> <<16) into an i64-held-f32 buffer. license_tier: ORIGINAL 7import "nx_syscalls.nx" 8import "nx_f32_cvt.nx" 9 10func stl_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 } 11func stl_u32le(buf: *u8, off: i64) -> i64 { return (buf[off]&0xff) | ((buf[off+1]&0xff)<<8) | ((buf[off+2]&0xff)<<16) | ((buf[off+3]&0xff)<<24) } 12func stl_u16le(buf: *u8, off: i64) -> i64 { return (buf[off]&0xff) | ((buf[off+1]&0xff)<<8) } 13 14func stl_streq(a: *u8, b: *u8) -> i64 { var i: i64=0; while a[i]!=(0 as u8){ if a[i]!=b[i]{return 0} i=i+1 } if b[i]!=(0 as u8){return 0} return 1 } 15 16// find needle (NUL-terminated) in buf[from..n); return start index or -1. 17func stl_find(buf: *u8, n: i64, from: i64, needle: *u8) -> i64 { 18 var m: i64=0; while needle[m]!=(0 as u8){m=m+1} 19 if m==0 { return 0-1 } 20 var i: i64=from 21 while i+m<=n { 22 var j: i64=0; var ok: i64=1 23 while j<m { if buf[i+j]!=needle[j] {ok=0; j=m} else {j=j+1} } 24 if ok==1 { return i } 25 i=i+1 26 } 27 return 0-1 28} 29 30// parse a base-10 integer at buf[pos]; endp[0] <- index past the last digit. 31func stl_parse_int(buf: *u8, pos: i64, endp: *i64) -> i64 { 32 var v: i64=0; var p: i64=pos; var go: i64=1 33 while go==1 { 34 let c: i64 = buf[p] as i64 35 if c>=48 { if c<=57 { v = v*10 + (c-48); p=p+1 } else { go=0 } } else { go=0 } 36 } 37 endp[0]=p 38 return v 39} 40 41func stl_header_len(buf: *u8) -> i64 { return stl_u64le(buf, 0) } 42func stl_data_start(buf: *u8) -> i64 { return 8 + stl_u64le(buf, 0) } 43 44// locate a tensor by name inside the JSON header. Fills dtype_out (NUL-term) + offs[0],offs[1] (data_offsets). 45// Returns 1 if found, 0 if not. hlen = stl_header_len(buf). 46func stl_tensor(buf: *u8, hlen: i64, name: *u8, dtype_out: *u8, offs: *i64) -> i64 { 47 let hend: i64 = 8 + hlen 48 let np: i64 = stl_find(buf, hend, 8, name) 49 if np < 0 { return 0 } 50 let dp: i64 = stl_find(buf, hend, np, "\"dtype\":\"" as *u8) 51 if dp < 0 { return 0 } 52 var i: i64 = dp + 9 // past `"dtype":"` 53 var k: i64 = 0 54 while buf[i] != (34 as u8) { dtype_out[k] = buf[i]; i = i + 1; k = k + 1 } 55 dtype_out[k] = 0 as u8 56 let op: i64 = stl_find(buf, hend, np, "\"data_offsets\":[" as *u8) 57 if op < 0 { return 0 } 58 let e0: *i64 = sys_mmap(8) as *i64 59 offs[0] = stl_parse_int(buf, op + 16, e0) // past `"data_offsets":[` 60 var p: i64 = e0[0] 61 while buf[p] < (48 as u8) { p = p + 1 } // skip comma/space to the second number 62 let e1: *i64 = sys_mmap(8) as *i64 63 offs[1] = stl_parse_int(buf, p, e1) 64 return 1 65} 66 67// read a tensor's values into out (i64-held f32 bits). Returns the element count, or -1 on unsupported dtype. 68func stl_read_f32(buf: *u8, data_start: i64, offs: *i64, dtype: *u8, out: *i64) -> i64 { 69 let bs: i64 = data_start + offs[0] 70 let nbytes: i64 = offs[1] - offs[0] 71 if stl_streq(dtype, "F32" as *u8) == 1 { 72 let n: i64 = nbytes / 4 73 var i: i64 = 0 74 while i < n { out[i] = stl_u32le(buf, bs + i*4); i = i + 1 } 75 return n 76 } 77 if stl_streq(dtype, "F16" as *u8) == 1 { 78 let n: i64 = nbytes / 2 79 var i: i64 = 0 80 while i < n { out[i] = nx_f16_to_f32(stl_u16le(buf, bs + i*2)); i = i + 1 } 81 return n 82 } 83 if stl_streq(dtype, "BF16" as *u8) == 1 { 84 let n: i64 = nbytes / 2 85 var i: i64 = 0 86 while i < n { out[i] = (stl_u16le(buf, bs + i*2)) << 16; i = i + 1 } // bf16 = high 16 bits of f32 87 return n 88 } 89 return 0 - 1 90}