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";

type Args = {
  in?: string;
  out?: string;
  raw?: boolean;
  concat?: string[];
  help?: boolean;
};

function parseArgs(argv: string[]): Args {
  const args: 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() {
  // Intentionally minimal and compatible with the existing Bazel tool flags.
  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: number) {
  if (v < 0) return 0;
  if (v > 255) return 255;
  return v | 0;
}

function mapFloatToU8(v: number) {
  // NNC/Draw Things uses: FloatType(v) * 2 / 255 - 1 for encoding
  // Inverse: (v + 1) * 255 / 2 = (v + 1) * 127.5
  return clampByte(Math.round((v + 1) * 127.5));
}

function readHeaderU32LE(data: Uint8Array, word: number) {
  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) as 1 | 3 | 4;

  const isCompressed = magic === 1012247 && !args.raw;
  const payload = input.subarray(68);

  let rgba: Uint8Array;

  if (isCompressed) {
    // Try zlib-inflate first (".zip"), then fpzip (".fpzip").
    let fpzipBytes: Uint8Array = payload;

    try {
      fpzipBytes = new Uint8Array(inflateSync(Buffer.from(payload)));
    } catch {
      // Not zlib-compressed (or failed inflate) - fall back to raw fpzip bytes.
    }

    let floats: Float32Array;
    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 {
    // Raw Float16 payload.
    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 as any).sync.write(png as any);
  await writeFile(args.out, outPng);
}

main().catch((err) => {
  console.error(err?.stack ?? String(err));
  process.exit(1);
});