uncloseai-speech/scripts/whisper_refs.py

#!/usr/bin/env python3
"""Transcribe cloned-voices/*.wav with whisper-large-v3 and write ref_texts.

Outputs:
  - cloned-voices/whisper_refs.json         (raw transcripts keyed by voice name)
  - voice_to_speaker.default.yaml           (ref_text rewritten in place)
  - cloned-voices/voices_metadata.json      (ref_text rewritten in place)

F5-TTS cloning quality depends on ref_text matching the prosody of ref_audio
(commas, periods, casing). LibriSpeech ground-truth labels are ALL CAPS with
no punctuation, which is the wrong signal for a flow-matching TTS conditioned
on text. Whisper hears what F5 will hear, so its transcript is the better ref.

Run on a GPU host (4090 / 3090). Idempotent: rerun any time cloned-voices/
changes. Use `make whisper-refs` rather than calling this directly.
"""
from __future__ import annotations

import json
import os
import re
import sys
from pathlib import Path

import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline


REPO_ROOT = Path(__file__).resolve().parent.parent
VOICES_DIR = REPO_ROOT / "cloned-voices"
YAML_PATH = REPO_ROOT / "voice_to_speaker.default.yaml"
METADATA_PATH = VOICES_DIR / "voices_metadata.json"
WHISPER_JSON = VOICES_DIR / "whisper_refs.json"
MODEL_ID = os.environ.get("WHISPER_MODEL", "openai/whisper-large-v3")


def load_pipeline():
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device.startswith("cuda") else torch.float32
    print(f"[whisper_refs] loading {MODEL_ID} on {device} ({dtype})", flush=True)
    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        MODEL_ID, torch_dtype=dtype, low_cpu_mem_usage=True, use_safetensors=True
    ).to(device)
    processor = AutoProcessor.from_pretrained(MODEL_ID)
    return pipeline(
        "automatic-speech-recognition",
        model=model,
        tokenizer=processor.tokenizer,
        feature_extractor=processor.feature_extractor,
        torch_dtype=dtype,
        device=device,
        return_timestamps=False,
    )


def transcribe_all(asr) -> dict[str, str]:
    wavs = sorted(VOICES_DIR.glob("*.wav"))
    print(f"[whisper_refs] {len(wavs)} wavs to transcribe", flush=True)
    out: dict[str, str] = {}
    for i, wav in enumerate(wavs, 1):
        result = asr(
            str(wav),
            generate_kwargs={"language": "en", "task": "transcribe"},
        )
        text = result["text"].strip()
        text = re.sub(r"\s+", " ", text)
        voice = wav.stem
        out[voice] = text
        print(f"[whisper_refs] [{i:2d}/{len(wavs)}] {voice:10s} -> {text}", flush=True)
    return out


def rewrite_yaml(refs: dict[str, str]) -> int:
    """Surgical replace of ref_text values keyed by ref_audio filename.

    Avoids a full YAML round-trip so comments/order/whitespace stay untouched.
    Each voice block contains a `ref_audio: cloned-voices/<voice>.wav` line
    followed (next non-blank line) by `ref_text: "..."`. Match on the audio
    path and rewrite the next ref_text line.
    """
    lines = YAML_PATH.read_text().splitlines(keepends=True)
    audio_re = re.compile(r"^(\s*)ref_audio:\s*cloned-voices/([^\s.]+)\.wav\s*$")
    text_re = re.compile(r"^(\s*)ref_text:\s*.*$")
    changed = 0
    i = 0
    while i < len(lines):
        m = audio_re.match(lines[i].rstrip("\n"))
        if not m:
            i += 1
            continue
        voice = m.group(2)
        if voice not in refs:
            i += 1
            continue
        # find the next ref_text line within the same block (no blank line break)
        j = i + 1
        while j < len(lines) and lines[j].strip() != "":
            tm = text_re.match(lines[j].rstrip("\n"))
            if tm:
                indent = tm.group(1)
                escaped = refs[voice].replace("\\", "\\\\").replace('"', '\\"')
                lines[j] = f'{indent}ref_text: "{escaped}"\n'
                changed += 1
                break
            j += 1
        i = j + 1
    YAML_PATH.write_text("".join(lines))
    return changed


def rewrite_metadata(refs: dict[str, str]) -> int:
    data = json.loads(METADATA_PATH.read_text())
    changed = 0
    for voice, entry in data.items():
        if voice in refs and entry.get("ref_text") != refs[voice]:
            entry["ref_text"] = refs[voice]
            changed += 1
    METADATA_PATH.write_text(json.dumps(data, indent=2) + "\n")
    return changed


def main() -> int:
    if not VOICES_DIR.is_dir():
        print(f"[whisper_refs] no cloned-voices dir at {VOICES_DIR}", file=sys.stderr)
        return 2
    asr = load_pipeline()
    refs = transcribe_all(asr)
    WHISPER_JSON.write_text(json.dumps(refs, indent=2, ensure_ascii=False) + "\n")
    print(f"[whisper_refs] wrote {WHISPER_JSON} ({len(refs)} entries)", flush=True)
    yaml_changed = rewrite_yaml(refs)
    meta_changed = rewrite_metadata(refs)
    print(
        f"[whisper_refs] yaml: {yaml_changed} ref_text replaced | "
        f"metadata: {meta_changed} ref_text replaced",
        flush=True,
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())