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

type Args = {
  in?: string;
  out?: string;
  emitMask?: string;
  noCompress?: boolean;
  zlibWrap?: boolean;
  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 === "--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;
    }

    // Legacy flag (now default behavior - fpzip without zlib wrapper)
    if (a === "--fpzip-only") {
      // This is now the default, flag kept for compatibility
      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: Uint8Array) {
  return createHash("sha256").update(data).digest("hex");
}

function writeTensorHeader(
  dv: DataView,
  opts: { compressed: boolean; width: number; height: number; channels: number }
) {
  // Matches src/imageHelpers.ts convertImageForRequest(), with the only difference
  // being header[0] which signals compression on server responses.
  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: DataView,
  opts: { width: number; height: number }
) {
  // Matches src/imageHelpers.ts convertImageToMask().
  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;

  // Draw Things expects image tensors in [-1, 1]. We emit RGB.
  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;
    // NNC/Draw Things uses: FloatType(v) * 2 / 255 - 1 (equivalent to v/127.5 - 1)
    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: Uint8Array;

  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 {
    // Compression pipeline: fpzip only (Draw Things expects raw fpzip, NOT zlib-wrapped)
    // precBits=16: Draw Things expects ~16-bit precision (matches Mac binary output size)
    // NNC fpzip dimension ordering: nx=channels, ny=width, nz=height (last dim first)
    const fpzipBytes = await fpzipCompress(floats, 16, {
      nx: channels, // Last dimension (channels)
      ny: width, // Second-to-last (width)
      nz: height, // Third-to-last (height)
      nf: 1, // Remaining dimensions
    });

    // Note: zlib wrapping was previously used but Draw Things expects raw fpzip.
    // Use --zlib-wrap flag if zlib wrapping is ever needed.
    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) {
    // If the PNG has alpha, derive a mask. Assumption: transparent => paint (2), opaque => retain (0).
    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);
});