draw-things-chat

import{readFile,writeFile}from"fs/promises";import{createHash}from"crypto";import{deflateSync}from"zlib";import{PNG}from"pngjs";import{setFloat16}from"@petamoriken/float16";import{compress as fpzipCompress}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==="--in"){args.in=argv[++i];continue}if(a==="--out"){args.out=argv[++i];continue}if(a==="--emit-mask"){args.emitMask=argv[++i];continue}if(a==="--no-compress"){args.noCompress=true;continue}if(a==="--fpzip-only"){continue}if(a==="--zlib-wrap"){args.zlibWrap=true;continue}throw new Error(`Unknown arg: ${a}`)}return args}function usage(){console.log(["Usage:","  npx tsx helpers/PNG2GRPCBin.ts --in <input.png> --out <image.bin> [--emit-mask <mask.bin>] [--no-compress] [--zlib-wrap]","","Output:","  Prints SHA-256 (hex) of the image blob (and mask blob, if emitted).","","Notes:","  Default compression is fpzip only (compatible with Draw Things).","  Use --zlib-wrap to add zlib compression on top of fpzip."].join("\n"))}function sha256Hex(data){return createHash("sha256").update(data).digest("hex")}function writeTensorHeader(dv,opts){dv.setUint32(0*4,opts.compressed?1012247:0,true);dv.setUint32(1*4,1,true);dv.setUint32(2*4,2,true);dv.setUint32(3*4,131072,true);dv.setUint32(5*4,1,true);dv.setUint32(6*4,opts.height,true);dv.setUint32(7*4,opts.width,true);dv.setUint32(8*4,opts.channels,true)}function writeMaskHeader(dv,opts){const header=[0,1,1,4096,0,opts.height,opts.width,0,0];header.forEach((v,i)=>dv.setUint32(i*4,v,true))}async function main(){const args=parseArgs(process.argv.slice(2));if(args.help){usage();return}if(!args.in)throw new Error("Missing --in");if(!args.out)throw new Error("Missing --out");if(args.noCompress&&args.zlibWrap){throw new Error("Cannot combine --no-compress and --zlib-wrap")}const pngIn=PNG.sync.read(await readFile(args.in));const{width,height,data}=pngIn;const channels=3;const nValues=width*height*channels;const floats=new Float32Array(nValues);for(let i=0;i<width*height;i++){const rgba=i*4;const base=i*3;floats[base+0]=data[rgba+0]*2/255-1;floats[base+1]=data[rgba+1]*2/255-1;floats[base+2]=data[rgba+2]*2/255-1}let tensor;if(args.noCompress){tensor=new Uint8Array(68+nValues*2);const dv=new DataView(tensor.buffer);writeTensorHeader(dv,{compressed:false,width,height,channels});for(let i=0;i<nValues;i++){setFloat16(dv,68+i*2,floats[i]??0,true)}}else{const fpzipBytes=await fpzipCompress(floats,16,{nx:channels,ny:width,nz:height,nf:1});const payload=args.zlibWrap?new Uint8Array(deflateSync(Buffer.from(fpzipBytes))):fpzipBytes;tensor=new Uint8Array(68+payload.byteLength);const dv=new DataView(tensor.buffer);writeTensorHeader(dv,{compressed:true,width,height,channels});tensor.set(payload,68)}await writeFile(args.out,tensor);console.log(sha256Hex(tensor));if(args.emitMask){const mask=new Uint8Array(68+width*height);const dv=new DataView(mask.buffer);writeMaskHeader(dv,{width,height});for(let i=0;i<width*height;i++){const alpha=data[i*4+3];mask[68+i]=alpha<128?2:0}await writeFile(args.emitMask,mask);console.log(`mask ${sha256Hex(mask)}`)}}main().catch(err=>{console.error(err?.stack??String(err));process.exit(1)});