pipeline/packages/tts/src/whisper-align.ts

import { spawn } from "node:child_process";
import { fileURLToPath } from "node:url";
import { dirname, join } from "node:path";
import type { WordTimestamp } from "@pipeline/shared";

const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);

const ALIGN_SCRIPT = join(__dirname, "..", "scripts", "align.py");

export interface WhisperAlignOptions {
  audioFilePath: string;
  text?: string;
  language?: string;
  model?: string;
  timeout?: number;
}

interface WhisperWord {
  word: string;
  start: number;
  end: number;
}

interface AlignResult {
  ok: boolean;
  words?: WhisperWord[];
  error?: string;
}

/**
 * Run Whisper forced alignment on a TTS-generated audio file.
 * Returns word-level timestamps mapped to the original narration text.
 * Falls back to null if whisper is unavailable.
 */
export async function alignWithWhisper(
  options: WhisperAlignOptions
): Promise<WordTimestamp[] | null> {
  const args = [
    ALIGN_SCRIPT,
    options.audioFilePath,
    ...(options.language ? ["--language", options.language] : []),
    ...(options.model ? ["--model", options.model] : []),
  ];

  const result = await runPython(args, options.timeout ?? 120_000);
  if (!result.ok || !result.words || result.words.length === 0) {
    return null;
  }

  return mapWhisperToOriginal(result.words, options.text);
}

function runPython(args: string[], timeout: number): Promise<AlignResult> {
  return new Promise((resolve) => {
    const proc = spawn("python3", args, {
      timeout,
      stdio: ["ignore", "pipe", "pipe"],
    });

    let stdout = "";
    let stderr = "";
    proc.stdout.on("data", (d: Buffer) => (stdout += d));
    proc.stderr.on("data", (d: Buffer) => (stderr += d));

    proc.on("close", () => {
      try {
        resolve(JSON.parse(stdout) as AlignResult);
      } catch {
        resolve({ ok: false, error: stderr || stdout || "parse error" });
      }
    });

    proc.on("error", (err) => {
      resolve({ ok: false, error: err.message });
    });
  });
}

/**
 * Map whisper's transcription timestamps back to the original text's words.
 *
 * Strategy: build a character-to-time mapping from whisper output, then
 * interpolate each original word/character into that timeline. This works
 * well for TTS audio where the spoken content closely matches the script.
 */
function mapWhisperToOriginal(
  whisperWords: WhisperWord[],
  originalText?: string
): WordTimestamp[] {
  if (!originalText || whisperWords.length === 0) {
    return whisperWords.map((w) => ({
      word: w.word,
      startSeconds: w.start,
      endSeconds: w.end,
    }));
  }

  // Build character-level time array from whisper output
  const charTimes: number[] = [];
  for (const w of whisperWords) {
    const perChar = (w.end - w.start) / Math.max(w.word.length, 1);
    for (let i = 0; i < w.word.length; i++) {
      charTimes.push(w.start + i * perChar);
    }
  }

  if (charTimes.length === 0) {
    return whisperWords.map((w) => ({
      word: w.word,
      startSeconds: w.start,
      endSeconds: w.end,
    }));
  }

  // Split original text into display units
  const isChinese = /[一-鿿]/.test(originalText);
  const units = isChinese
    ? originalText.split("").filter((c) => c.trim())
    : originalText.split(/\s+/).filter((w) => w);

  const totalOrigChars = units.reduce((sum, u) => sum + u.length, 0);

  let charOffset = 0;
  return units.map((unit) => {
    const startRatio = charOffset / totalOrigChars;
    const endRatio = (charOffset + unit.length) / totalOrigChars;

    const startIdx = Math.min(
      Math.floor(startRatio * charTimes.length),
      charTimes.length - 1
    );
    const endIdx = Math.min(
      Math.floor(endRatio * charTimes.length),
      charTimes.length - 1
    );

    const startSeconds = charTimes[startIdx];
    const endSeconds = charTimes[Math.max(endIdx, startIdx)];

    charOffset += unit.length;

    return { word: unit, startSeconds, endSeconds };
  });
}