draw-things-chat

import fs from "fs";
import path from "path";
import { createRequire } from "module";
import { decompress as fpzipDecompress } from "../fpzip/decompress.js";

// createRequire needs a valid file URL or path.
// On Hub's production.js (CJS), __filename is always defined.
let _require: NodeRequire;
try {
  const base = typeof __filename === "string" ? __filename
    : process.argv[1] ?? process.cwd();
  _require = createRequire(base);
} catch {
  _require = (typeof globalThis.require === "function" ? globalThis.require : require) as NodeRequire;
}

// ─── @ffmpeg/core module interface ────────────────────────────────────────────
type FFmpegFS = {
  writeFile(path: string, data: Uint8Array): void;
  readFile(path: string, opts: { encoding: "binary" }): Uint8Array;
  unlink(path: string): void;
};
type FFmpegCoreModule = {
  exec(...args: string[]): number;
  setLogger(fn: (msg: { type: string; message: string }) => void): void;
  reset(): void;
  FS: FFmpegFS;
};

// ─── Singleton loader ─────────────────────────────────────────────────────────
//
// @ffmpeg/core is compiled for browser Worker environments (ENVIRONMENT_IS_WORKER=true).
// It references `self.location.href` during init. We mock `self` as a new Function
// parameter so the Worker-targeted setup doesn't crash in Node.js / LM Studio.
//
let _corePromise: Promise<FFmpegCoreModule> | null = null;

export function loadCore(): Promise<FFmpegCoreModule> {
  if (_corePromise) return _corePromise;
  _corePromise = (async () => {
    const coreJsPath = _require.resolve("@ffmpeg/core");
    const wasmPath = coreJsPath.replace(/\.js$/, ".wasm");

    const code = fs.readFileSync(coreJsPath, "utf8");
    const wasmBinary = new Uint8Array(fs.readFileSync(wasmPath));

    const mockSelf = {
      location: { href: "file://" + path.dirname(coreJsPath).replace(/\\/g, "/") + "/" },
    };

    type CoreFactory = (opts: Record<string, unknown>) => Promise<FFmpegCoreModule>;
    const factory = new Function(
      "require",
      "__dirname",
      "__filename",
      "self",
      code + "\nreturn createFFmpegCore;\n"
    ) as unknown as (
      req: NodeRequire,
      dir: string,
      file: string,
      self: typeof mockSelf
    ) => CoreFactory;

    const createFFmpegCore = factory(
      _require,
      path.dirname(coreJsPath),
      coreJsPath,
      mockSelf
    );

    return createFFmpegCore({
      wasmBinary,
      print: () => {},
      printErr: () => {},
    });
  })();
  return _corePromise;
}

// ─── NNC tensor → ffmpeg f32le interleaved conversion ────────────────────────
//
// Draw Things serialises audio via s4nnc Tensor<Float>.data(using: []).
// That call prepends a binary header before the raw float samples:
//
//   [4 bytes]  UInt32 identifier  (= 0 for uncompressed)
//   [64 bytes] ccv_nnc_tensor_param_t  (type / format / datatype / reserved / dim[12])
//   [rest]     raw fpzip stream (no zlib wrapper — fpzipAndZipEncode delegates to
//              fpzipEncode for float data, which skips the zip layer)
//
// ffmpeg -f f32le -ac 2 expects INTERLEAVED layout: [L₀ R₀  L₁ R₁ … L_{T-1} R_{T-1}].
// The fpzip-decompressed tensor is NC(2, T) planar: [ch0[0..T-1], ch1[0..T-1]].

const NNC_HEADER_SIZE = 4 /* UInt32 identifier */ + 64 /* ccv_nnc_tensor_param_t */;
const FPZIP_IDENTIFIER = 0xf7217;

async function nncAudioToF32leInterleaved(nncData: Buffer): Promise<Buffer> {
  if (nncData.length <= NNC_HEADER_SIZE) {
    throw new Error(`NNC audio buffer too short: ${nncData.length} bytes`);
  }
  const identifier = nncData.readUInt32LE(0);
  if (identifier !== FPZIP_IDENTIFIER) {
    throw new Error(`Unexpected NNC audio identifier: 0x${identifier.toString(16)}`);
  }

  const fpzipStream = new Uint8Array(
    nncData.buffer,
    nncData.byteOffset + NNC_HEADER_SIZE,
    nncData.length - NNC_HEADER_SIZE,
  );

  const floats = await fpzipDecompress(fpzipStream);
  const samplesPerChannel = Math.floor(floats.length / 2);

  const interleaved = Buffer.allocUnsafe(samplesPerChannel * 8);
  const view = new DataView(interleaved.buffer);
  for (let i = 0; i < samplesPerChannel; i++) {
    view.setFloat32(i * 8,     floats[i],                    true);
    view.setFloat32(i * 8 + 4, floats[samplesPerChannel + i], true);
  }
  return interleaved;
}

/**
 * Assembles PNG frame buffers into an Apple ProRes 4444 .mov using
 * @ffmpeg/core WASM (fully in-process — no child_process spawn).
 *
 * @param frames          - Ordered array of PNG frame buffers (already trimmed)
 * @param fps             - Output frame rate
 * @param audioRaw        - Optional NNC-serialised LPCM audio (f32le, stereo).
 * @param audioSampleRate - Sample rate of the decoded audio in Hz (default: 24 000).
 *                         Use 48 000 for LTX-2.3 (BWE vocoder doubles the sample rate).
 * @returns Buffer containing the assembled .mov file
 */
export async function assembleVideo(
  frames: Buffer[],
  fps: number,
  audioRaw?: Buffer,
  audioSampleRate = 48_000
): Promise<Buffer> {
  if (frames.length === 0) {
    throw new Error("assembleVideo: frames array is empty");
  }

  const ffmpeg = await loadCore();
  ffmpeg.reset();

  const stderrLines: string[] = [];
  ffmpeg.setLogger((msg) => {
    if (msg.type === "stderr") stderrLines.push(msg.message);
  });

  const hasAudio = audioRaw && audioRaw.length > 0;

  try {
    for (let i = 0; i < frames.length; i++) {
      const name = `frame_${String(i + 1).padStart(4, "0")}.png`;
      ffmpeg.FS.writeFile(name, new Uint8Array(frames[i]));
    }

    if (hasAudio) {
      const audioInterleaved = await nncAudioToF32leInterleaved(audioRaw!);
      ffmpeg.FS.writeFile("audio.raw", new Uint8Array(audioInterleaved));
    }

    const args: string[] = [
      "-framerate", String(fps),
      "-i", "frame_%04d.png",
    ];

    if (hasAudio) {
      args.push(
        "-f", "f32le",
        "-ar", String(audioSampleRate),
        "-ac", "2",
        "-i", "audio.raw",
        "-c:v", "prores_ks",
        "-profile:v", "4444",
        "-pix_fmt", "yuva444p10le",
        "-c:a", "pcm_s32le",
      );
    } else {
      args.push(
        "-c:v", "prores_ks",
        "-profile:v", "4444",
        "-pix_fmt", "yuva444p10le",
      );
    }

    args.push("out.mov");

    const exitCode = ffmpeg.exec(...args);

    if (exitCode !== 0) {
      throw new Error(
        `ffmpeg.wasm exec failed (exit ${exitCode}): ${stderrLines.slice(-30).join("\n")}`
      );
    }

    const data = ffmpeg.FS.readFile("out.mov", { encoding: "binary" });
    return Buffer.from(data);
  } finally {
    for (let i = 0; i < frames.length; i++) {
      try {
        ffmpeg.FS.unlink(`frame_${String(i + 1).padStart(4, "0")}.png`);
      } catch {}
    }
    if (hasAudio) {
      try { ffmpeg.FS.unlink("audio.raw"); } catch {}
    }
    try {
      ffmpeg.FS.unlink("out.mov");
    } catch {}
    ffmpeg.setLogger(() => {});
  }
}