draw-things-chat

// PNG iTXt chunk injection for generation metadata.
// Writes an XML:com.adobe.xmp iTXt chunk matching the format that the Draw Things app uses,
// so generated images are compatible with draw-things-index and external tools (CivitAI, etc.).
import os from "node:os";

export type PngXmpParams = {
  prompt?: string;
  negativePrompt?: string;
  model?: string;
  width?: number;
  height?: number;
  steps?: number;
  seed?: number;
  seedMode?: string;
  sampler?: string;
  guidanceScale?: number;
  strength?: number;
  shift?: number;
  loras?: Array<{ file: string; weight?: number }>;
  /** Absolute paths to source images (img2img / edit). Normalized to ~/... in output. */
  sources?: string[];
  mode?: string;
  isVideoFrame?: boolean;
};

// CRC32 table (IEEE 802.3 polynomial, same as used by zlib / PNG spec)
const CRC_TABLE = (() => {
  const table = new Int32Array(256);
  for (let n = 0; n < 256; n++) {
    let c = n;
    for (let k = 0; k < 8; k++) {
      c = c & 1 ? (0xedb88320 ^ (c >>> 1)) : (c >>> 1);
    }
    table[n] = c;
  }
  return table;
})();

function crc32(data: Buffer): number {
  let crc = 0xffffffff;
  for (let i = 0; i < data.length; i++) {
    crc = CRC_TABLE[(crc ^ data[i]!) & 0xff]! ^ (crc >>> 8);
  }
  return (crc ^ 0xffffffff) >>> 0;
}

function normPath(absPath: string): string {
  const home = os.homedir();
  return absPath.startsWith(home) ? "~" + absPath.slice(home.length) : absPath;
}

function escapeXml(s: string): string {
  return s
    .replace(/&/g, "&amp;")
    .replace(/</g, "&lt;")
    .replace(/>/g, "&gt;")
    .replace(/"/g, "&quot;");
}

function basename(filePath: string): string {
  return filePath.replace(/^.*[\\/]/, "");
}

function buildA1111Description(p: PngXmpParams): string {
  const lines: string[] = [];

  if (p.prompt) lines.push(p.prompt);
  if (p.negativePrompt) lines.push(`Negative prompt: ${p.negativePrompt}`);

  const info: string[] = [];
  if (typeof p.steps === "number") info.push(`Steps: ${p.steps}`);
  if (p.sampler) info.push(`Sampler: ${p.sampler}`);
  if (typeof p.guidanceScale === "number") info.push(`Guidance Scale: ${p.guidanceScale}`);
  if (typeof p.seed === "number") info.push(`Seed: ${p.seed}`);
  if (typeof p.width === "number" && typeof p.height === "number") {
    info.push(`Size: ${p.width}x${p.height}`);
  }
  if (p.model) info.push(`Model: ${p.model}`);
  if (typeof p.strength === "number") info.push(`Strength: ${p.strength}`);
  if (p.seedMode) info.push(`Seed Mode: ${p.seedMode}`);
  if (typeof p.shift === "number") info.push(`Shift: ${p.shift}`);
  if (info.length > 0) lines.push(info.join(", "));

  if (p.loras) {
    p.loras.forEach((l, i) => {
      const idx = i + 1;
      const name = basename(l.file);
      let entry = `LoRA ${idx} Model: ${name}`;
      if (typeof l.weight === "number") entry += `, LoRA ${idx} Weight: ${l.weight}`;
      lines.push(entry);
    });
  }

  if (p.sources && p.sources.length > 0) {
    lines.push(`Source: ${p.sources.map(normPath).join(", ")}`);
  }

  if (p.isVideoFrame) lines.push("Note: Last frame of generated video.");

  return lines.join("\n");
}

function buildUserCommentJson(p: PngXmpParams): string {
  const obj: Record<string, unknown> = {};

  if (p.prompt != null) obj["c"] = p.prompt;
  if (p.negativePrompt != null) obj["uc"] = p.negativePrompt;
  if (p.model) obj["model"] = p.model;
  if (p.sampler) obj["sampler"] = p.sampler;
  if (typeof p.guidanceScale === "number") obj["scale"] = p.guidanceScale;
  if (typeof p.seed === "number") obj["seed"] = p.seed;
  if (p.seedMode) obj["seed_mode"] = p.seedMode;
  if (typeof p.shift === "number") obj["shift"] = p.shift;
  if (typeof p.width === "number" && typeof p.height === "number") {
    obj["size"] = `${p.width}x${p.height}`;
  }
  if (typeof p.steps === "number") obj["steps"] = p.steps;
  if (typeof p.strength === "number") obj["strength"] = p.strength;
  if (p.loras && p.loras.length > 0) {
    obj["lora"] = p.loras.map((l) => {
      const entry: Record<string, unknown> = { model: basename(l.file) };
      if (typeof l.weight === "number") entry["weight"] = l.weight;
      return entry;
    });
  }
  if (p.sources && p.sources.length > 0) {
    obj["sources"] = p.sources.map(normPath);
  }
  if (p.mode) obj["mode"] = p.mode;
  if (p.isVideoFrame) obj["is_video_frame"] = true;
  obj["generated_by"] = "draw-things-chat";

  return JSON.stringify(obj);
}

function buildXmpString(p: PngXmpParams): string {
  const desc = escapeXml(buildA1111Description(p));
  const uc = buildUserCommentJson(p); // JSON in XML text content — no escaping needed/used by Draw Things
  return [
    `<x:xmpmeta xmlns:x="adobe:ns:meta/" x:xmptk="XMP Core 6.0.0">`,
    `   <rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">`,
    `      <rdf:Description rdf:about=""`,
    `            xmlns:dc="http://purl.org/dc/elements/1.1/"`,
    `            xmlns:xmp="http://ns.adobe.com/xap/1.0/"`,
    `            xmlns:exif="http://ns.adobe.com/exif/1.0/">`,
    `         <dc:description>`,
    `            <rdf:Alt>`,
    `               <rdf:li xml:lang="x-default">${desc}</rdf:li>`,
    `            </rdf:Alt>`,
    `         </dc:description>`,
    `         <xmp:CreatorTool>draw-things-chat</xmp:CreatorTool>`,
    `         <exif:UserComment>`,
    `            <rdf:Alt>`,
    `               <rdf:li xml:lang="x-default">${uc}</rdf:li>`,
    `            </rdf:Alt>`,
    `         </exif:UserComment>`,
    `      </rdf:Description>`,
    `   </rdf:RDF>`,
    `</x:xmpmeta>`,
  ].join("\n");
}

function buildITxtChunk(keyword: string, text: string): Buffer {
  // iTXt data: keyword\0 | comprFlag(0) | comprMethod(0) | langTag\0 | transKeyword\0 | text
  const keyBuf = Buffer.from(keyword, "utf8");
  const textBuf = Buffer.from(text, "utf8");
  const data = Buffer.concat([
    keyBuf,
    Buffer.from([0x00, 0x00, 0x00, 0x00, 0x00]),
    textBuf,
  ]);

  const typeAndData = Buffer.concat([Buffer.from("iTXt", "ascii"), data]);
  const crc = crc32(typeAndData);

  const chunk = Buffer.allocUnsafe(4 + 4 + data.length + 4);
  chunk.writeUInt32BE(data.length, 0);
  chunk.write("iTXt", 4, "ascii");
  data.copy(chunk, 8);
  chunk.writeUInt32BE(crc, 8 + data.length);

  return chunk;
}

const PNG_SIG = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]);
// IHDR is always 25 bytes: 4 (length field) + 4 ("IHDR") + 13 (data) + 4 (CRC)
const IHDR_BYTES = 25;

/**
 * Injects an XMP iTXt chunk into a PNG buffer, positioned right after the IHDR chunk.
 * Returns the original buffer unchanged if it is not a valid PNG.
 */
export function injectXmpIntoBuffer(pngBuf: Buffer, params: PngXmpParams): Buffer {
  if (pngBuf.length < PNG_SIG.length + IHDR_BYTES) return pngBuf;
  if (!pngBuf.subarray(0, PNG_SIG.length).equals(PNG_SIG)) return pngBuf;

  const xmp = buildXmpString(params);
  const chunk = buildITxtChunk("XML:com.adobe.xmp", xmp);
  const insertAt = PNG_SIG.length + IHDR_BYTES;

  return Buffer.concat([
    pngBuf.subarray(0, insertAt),
    chunk,
    pngBuf.subarray(insertAt),
  ]);
}