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==="--profile"){args.profile=parseDecoderProfile(argv[++i]);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)","  --profile <direct-rgb|qwen-argb|flux2-latent32>","  -h, --help"].join("\n"))}function parseDecoderProfile(value){if(!value||value==="direct-rgb")return"direct-rgb";if(value==="qwen-argb")return"qwen-argb";if(value==="flux2-latent32")return"flux2-latent32";throw new Error("Unsupported decoder profile: "+value)}function clampByte(v){if(v<0)return 0;if(v>255)return 255;return v|0}function mapFloatToU8(v){return clampByte(Math.floor((v+1)*127.5))}function mapAlphaToU8(v){return clampByte(Math.round(v*255))}const flux2RgbProjection=[[.0058,.0495,-.0099,.2144,.0166,.0157,-.0398,-.0052,-.3527,-.0301,-.0107,.0746,.0156,-.0034,.0032,-.0939,.0018,.0284,-.0024,.1207,.0128,.0137,.0095,0,-.0465,.0095,.029,.022,-.0332,-.0085,-.0076,-.0111,-.0329],[.0113,.0443,.0096,.3009,-.0039,.0103,.0902,.0095,-.2712,-.0356,.0078,.009,.0169,-.004,.0181,-8e-4,.0043,.0056,-.0022,-.0026,.0101,-.0072,.0092,-.0077,-.0204,.0012,-.0034,.0169,-.0457,.0389,3e-4,-.046,-.0718],[.0073,.0836,.0644,.3652,-.0054,-.016,-.0235,.0109,-.1666,-.018,.0013,-.0941,.007,-.0114,.008,.0186,.0104,-.0127,-.003,.0065,.0142,-7e-4,-.0059,-.0049,-.0312,-.0066,.0025,-.0048,-.0468,.0609,-.0043,-.0614,-.0851]];function decodeFlux2Pixel(values,base,rgba,out){for(let channel=0;channel<3;channel++){const coefficients=flux2RgbProjection[channel];let value=coefficients[32];for(let index=0;index<32;index++){value+=coefficients[index]*(values[base+index]??0)}rgba[out+channel]=clampByte(Math.trunc(value*127.5+127.5))}rgba[out+3]=255}function decodePixel(profile,channels,values,base,rgba,out){if(profile==="flux2-latent32"&&channels===32){decodeFlux2Pixel(values,base,rgba,out);return}if(profile==="qwen-argb"&&channels===4){rgba[out+0]=mapFloatToU8(values[base+1]??0);rgba[out+1]=mapFloatToU8(values[base+2]??0);rgba[out+2]=mapFloatToU8(values[base+3]??0);rgba[out+3]=mapAlphaToU8(values[base+0]??1);return}rgba[out+0]=mapFloatToU8(values[base+0]??0);rgba[out+1]=mapFloatToU8(values[base+1]??0);rgba[out+2]=mapFloatToU8(values[base+2]??0);rgba[out+3]=255}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}const profile=args.profile??"direct-rgb";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;decodePixel(profile,channels,floats,base,rgba,out)}}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;decodePixel(profile,channels,f16,base,rgba,out)}}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)});

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==="--profile"){args.profile=parseDecoderProfile(argv[++i]);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)","  --profile <direct-rgb|qwen-argb|flux2-latent32>","  -h, --help"].join("\n"))}function parseDecoderProfile(value){if(!value||value==="direct-rgb")return"direct-rgb";if(value==="qwen-argb")return"qwen-argb";if(value==="flux2-latent32")return"flux2-latent32";throw new Error("Unsupported decoder profile: "+value)}function clampByte(v){if(v<0)return 0;if(v>255)return 255;return v|0}function mapFloatToU8(v){return clampByte(Math.floor((v+1)*127.5))}function mapAlphaToU8(v){return clampByte(Math.round(v*255))}const flux2RgbProjection=[[.0058,.0495,-.0099,.2144,.0166,.0157,-.0398,-.0052,-.3527,-.0301,-.0107,.0746,.0156,-.0034,.0032,-.0939,.0018,.0284,-.0024,.1207,.0128,.0137,.0095,0,-.0465,.0095,.029,.022,-.0332,-.0085,-.0076,-.0111,-.0329],[.0113,.0443,.0096,.3009,-.0039,.0103,.0902,.0095,-.2712,-.0356,.0078,.009,.0169,-.004,.0181,-8e-4,.0043,.0056,-.0022,-.0026,.0101,-.0072,.0092,-.0077,-.0204,.0012,-.0034,.0169,-.0457,.0389,3e-4,-.046,-.0718],[.0073,.0836,.0644,.3652,-.0054,-.016,-.0235,.0109,-.1666,-.018,.0013,-.0941,.007,-.0114,.008,.0186,.0104,-.0127,-.003,.0065,.0142,-7e-4,-.0059,-.0049,-.0312,-.0066,.0025,-.0048,-.0468,.0609,-.0043,-.0614,-.0851]];function decodeFlux2Pixel(values,base,rgba,out){for(let channel=0;channel<3;channel++){const coefficients=flux2RgbProjection[channel];let value=coefficients[32];for(let index=0;index<32;index++){value+=coefficients[index]*(values[base+index]??0)}rgba[out+channel]=clampByte(Math.trunc(value*127.5+127.5))}rgba[out+3]=255}function decodePixel(profile,channels,values,base,rgba,out){if(profile==="flux2-latent32"&&channels===32){decodeFlux2Pixel(values,base,rgba,out);return}if(profile==="qwen-argb"&&channels===4){rgba[out+0]=mapFloatToU8(values[base+1]??0);rgba[out+1]=mapFloatToU8(values[base+2]??0);rgba[out+2]=mapFloatToU8(values[base+3]??0);rgba[out+3]=mapAlphaToU8(values[base+0]??1);return}rgba[out+0]=mapFloatToU8(values[base+0]??0);rgba[out+1]=mapFloatToU8(values[base+1]??0);rgba[out+2]=mapFloatToU8(values[base+2]??0);rgba[out+3]=255}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}const profile=args.profile??"direct-rgb";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;decodePixel(profile,channels,floats,base,rgba,out)}}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;decodePixel(profile,channels,f16,base,rgba,out)}}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)});