draw-things-chat

import{readFile,writeFile}from"fs/promises";import{inflateSync}from"zlib";import{PNG}from"pngjs";import{Float16Array}from"@petamoriken/float16";import{decompress as fpzipDecompress}from"../fpzip/decompress.js";function parseArgs(argv){const args={};for(let i=0;i<argv.length;i++){const a=argv[i];if(a==="--help"||a==="-h"){args.help=true;continue}if(a==="--raw"){args.raw=true;continue}if(a==="--in"){args.in=argv[++i];continue}if(a==="--out"){args.out=argv[++i];continue}if(a==="--concat"){args.concat=[];while(i+1<argv.length&&!argv[i+1].startsWith("--")){args.concat.push(argv[++i])}continue}throw new Error(`Unknown arg: ${a}`)}return args}function usage(){console.log(["Usage:","  npx tsx helpers/GRPCBin2PNG.ts --in <input.bin> --out <output.png>","  npx tsx helpers/GRPCBin2PNG.ts --concat <part1.bin> <part2.bin> ... --out <output.png>","","Options:","  --raw    Treat input as uncompressed Float16 tensor (no zip/fpzip)","  -h, --help"].join("\n"))}function clampByte(v){if(v<0)return 0;if(v>255)return 255;return v|0}function mapFloatToU8(v){return clampByte(Math.round((v+1)*127.5))}function readHeaderU32LE(data,word){const o=word*4;return data[o]|data[o+1]<<8|data[o+2]<<16|data[o+3]<<24}async function main(){const args=parseArgs(process.argv.slice(2));if(args.help){usage();return}if(!args.out)throw new Error("Missing --out");if(!args.in&&(!args.concat||args.concat.length===0)){throw new Error("Missing --in or --concat")}const inputBuffers=args.concat?.length?await Promise.all(args.concat.map(p=>readFile(p))):[await readFile(args.in)];const input=Uint8Array.from(Buffer.concat(inputBuffers));if(input.byteLength<68)throw new Error("Input too small to be a tensor");const magic=readHeaderU32LE(input,0);const height=readHeaderU32LE(input,6)>>>0;const width=readHeaderU32LE(input,7)>>>0;const channels=readHeaderU32LE(input,8)>>>0;const isCompressed=magic===1012247&&!args.raw;const payload=input.subarray(68);let rgba;if(isCompressed){let fpzipBytes=payload;try{fpzipBytes=new Uint8Array(inflateSync(Buffer.from(payload)))}catch{}let floats;try{floats=await fpzipDecompress(fpzipBytes)}catch{floats=await fpzipDecompress(payload)}const pixels=width*height;rgba=new Uint8Array(pixels*4);for(let i=0;i<pixels;i++){const out=i*4;if(channels===1){const v=mapFloatToU8(floats[i]??0);rgba[out+0]=v;rgba[out+1]=v;rgba[out+2]=v;rgba[out+3]=255;continue}const base=i*channels;rgba[out+0]=mapFloatToU8(floats[base+0]??0);rgba[out+1]=mapFloatToU8(floats[base+1]??0);rgba[out+2]=mapFloatToU8(floats[base+2]??0);rgba[out+3]=channels===4?mapFloatToU8(floats[base+3]??1):255}}else{const values=(input.byteLength-68)/2;if(!Number.isInteger(values))throw new Error("Invalid raw Float16 payload length");const f16=new Float16Array(input.buffer,input.byteOffset+68,values);const pixels=width*height;rgba=new Uint8Array(pixels*4);for(let i=0;i<pixels;i++){const out=i*4;if(channels===1){const v=mapFloatToU8(f16[i]??0);rgba[out+0]=v;rgba[out+1]=v;rgba[out+2]=v;rgba[out+3]=255;continue}const base=i*channels;rgba[out+0]=mapFloatToU8(f16[base+0]??0);rgba[out+1]=mapFloatToU8(f16[base+1]??0);rgba[out+2]=mapFloatToU8(f16[base+2]??0);rgba[out+3]=channels===4?mapFloatToU8(f16[base+3]??1):255}}const png=new PNG({width,height});png.data=Buffer.from(rgba);const outPng=PNG.sync.write(png);await writeFile(args.out,outPng)}main().catch(err=>{console.error(err?.stack??String(err));process.exit(1)});