train_narrator_tts.py

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "f5-tts>=1.0.0",
#   "datasets>=2.20.0",
#   "huggingface_hub>=0.24.0",
#   "soundfile>=0.12.1",
#   "librosa>=0.10.1",
#   "numpy<2.0",
#   "torch>=2.1.0",
#   "accelerate>=0.30.0",
# ]
# ///
"""Fine-tune F5-TTS on LJSpeech -> female narrator voice for lili.

Base:    SWivid/F5-TTS  (F5TTS_v1_Base, auto-downloaded when --finetune and no --pretrain)
Dataset: LJSpeech 1.1 downloaded directly from keithito.com  (22kHz, resampled to 24kHz)
Output:  pushed to HF_TTS_MODEL_REPO

Key F5-TTS internals learned from source inspection:
  - Dataset path:   {f5_tts_pkg}/../../data/{dataset_name}/  (load_from_disk)
  - Checkpoint out: {f5_tts_pkg}/../../ckpts/{dataset_name}/
  - CustomDataset accesses row["audio"]["array"] + row["audio"]["sampling_rate"]
  - Store audio as plain dict (no hf_datasets.Audio feature) to avoid torchcodec/FFmpeg dep
  - --tokenizer_path only respected when --tokenizer custom (not char/pinyin)
  - Must use --tokenizer pinyin to match base model text_embed shape [2546, 512]

Env vars:
  HF_TOKEN           — HF write token (injected as --secrets HF_TOKEN)
  HF_TTS_MODEL_REPO  — destination repo, e.g. 4moha/lili-narrator-v1
  NUM_SAMPLES        — optional, default 300
"""

import csv
import json
import os
import subprocess
import sys
import tarfile
import urllib.request
from pathlib import Path

import datasets as hf_datasets
import numpy as np
import soundfile as sf
from huggingface_hub import HfApi

HF_TOKEN = os.environ["HF_TOKEN"]
MODEL_REPO = os.environ.get("HF_TTS_MODEL_REPO", "4moha/lili-narrator-v1")
NUM_SAMPLES = int(os.environ.get("NUM_SAMPLES", "300"))

LJSPEECH_URL = "https://data.keithito.com/data/speech/LJSpeech-1.1.tar.bz2"
TARGET_SR = 24000
MIN_SECS = 2.0
MAX_SECS = 12.0

DATASET_NAME = "narrator"
EXP_ARCH = "F5TTS_v1_Base"
WAVS_DIR = Path("/tmp/f5_work/wavs")


def _lib_root() -> Path:
    """Returns .../lib/python3.x/ — where F5-TTS stores data/ and ckpts/ next to site-packages."""
    # sys.executable = .../bin/python  ->  parent = .../bin  ->  parent = env root
    # then down to lib/python3.x/
    ver = f"python{sys.version_info.major}.{sys.version_info.minor}"
    return Path(sys.executable).parent.parent / "lib" / ver


def _pkg_data_root() -> Path:
    return _lib_root() / "data"


def _pkg_ckpt_root() -> Path:
    return _lib_root() / "ckpts"


def prepare_data() -> None:
    """Download LJSpeech, build HF Arrow dataset + vocab at paths F5-TTS expects."""
    tar_path = Path("/tmp/ljspeech.tar.bz2")
    lj_root = Path("/tmp/LJSpeech-1.1")

    if not lj_root.exists():
        print(f"Downloading LJSpeech from {LJSPEECH_URL}...")
        urllib.request.urlretrieve(LJSPEECH_URL, str(tar_path))
        print("Extracting...")
        with tarfile.open(str(tar_path), "r:bz2") as tar:
            tar.extractall("/tmp/", filter="data")
        print(f"Extracted -> {lj_root}")

    WAVS_DIR.mkdir(parents=True, exist_ok=True)

    # F5-TTS CustomDataset.__getitem__ accesses row["audio_path"] (string path) and
    # row["text"]. It loads audio itself (soundfile/torchaudio) — no HF Audio feature,
    # no torchcodec, no FFmpeg dependency. Just store plain string paths + text.
    wav_paths: list[str] = []
    texts: list[str] = []
    durations: list[float] = []
    count = 0

    with open(str(lj_root / "metadata.csv"), encoding="utf-8") as f:
        for row in csv.reader(f, delimiter="|"):
            if count >= NUM_SAMPLES:
                break
            if len(row) < 3:
                continue
            clip_id, normalized_text = row[0], row[2]
            text = normalized_text.strip()
            if not text:
                continue
            wav_src = lj_root / "wavs" / f"{clip_id}.wav"
            if not wav_src.exists():
                continue

            arr, orig_sr = sf.read(str(wav_src))
            arr = arr.astype(np.float32)
            if arr.ndim > 1:
                arr = arr.mean(axis=1)
            duration = len(arr) / orig_sr
            if not (MIN_SECS <= duration <= MAX_SECS):
                continue
            if orig_sr != TARGET_SR:
                import librosa
                arr = librosa.resample(arr, orig_sr=orig_sr, target_sr=TARGET_SR)

            dest = WAVS_DIR / f"{clip_id}.wav"
            sf.write(str(dest), arr, TARGET_SR)
            wav_paths.append(str(dest))
            texts.append(text)
            durations.append(len(arr) / TARGET_SR)
            count += 1

    print(f"Prepared {count} samples")

    # F5-TTS dataset.py builds the path as data/{dataset_name}_{tokenizer}/raw/
    # Must use pinyin tokenizer to match base model's 2546-token text_embed shape.
    token_dir = _pkg_data_root() / f"{DATASET_NAME}_pinyin"
    raw_dir = token_dir / "raw"
    raw_dir.mkdir(parents=True, exist_ok=True)

    # CustomDataset.__getitem__ also reads row["duration"] per sample
    ds = hf_datasets.Dataset.from_dict({"audio_path": wav_paths, "text": texts, "duration": durations})
    ds.save_to_disk(str(raw_dir))
    print(f"Saved dataset -> {raw_dir}")

    # duration.json format confirmed from source: {"duration": [float, ...]}
    with (token_dir / "duration.json").open("w") as f:
        json.dump({"duration": durations}, f)
    print(f"Wrote duration.json ({len(durations)} entries) -> {token_dir / 'duration.json'}")

    # Copy the built-in pinyin vocab (2546 tokens) so get_tokenizer() can find it.
    # F5-TTS ships vocab.txt at site-packages/f5_tts/infer/examples/basic/vocab.txt
    site_packages = Path(sys.executable).parent.parent / "lib" / f"python{sys.version_info.major}.{sys.version_info.minor}" / "site-packages"
    builtin_vocab = site_packages / "f5_tts" / "infer" / "examples" / "vocab.txt"
    if builtin_vocab.exists():
        import shutil
        shutil.copy(builtin_vocab, token_dir / "vocab.txt")
        print(f"Copied built-in vocab -> {token_dir / 'vocab.txt'}")
    else:
        # Fall back: list what's in f5_tts to find vocab.txt
        print("Built-in vocab not found at expected path, searching...")
        for p in sorted((site_packages / "f5_tts").rglob("vocab.txt"))[:5]:
            print(f"  Found: {p}")
        raise FileNotFoundError(f"Cannot find F5-TTS built-in vocab.txt — searched {builtin_vocab}")


def _patch_torchaudio() -> None:
    """Write sitecustomize.py into the venv so torchaudio.load uses soundfile.

    torchaudio 2.6+ delegates load() to torchcodec which needs FFmpeg system libs.
    sitecustomize.py runs at Python startup before any module import, so our patch
    replaces torchaudio.load before F5-TTS's dataset.py ever calls it.
    """
    sc = _lib_root() / "site-packages" / "sitecustomize.py"
    sc.write_text(
        "try:\n"
        "    import torchaudio, soundfile as _sf, torch as _torch, numpy as _np\n"
        "    def _sf_load(uri, frame_offset=0, num_frames=-1, normalize=True,\n"
        "                 channels_first=True, format=None, backend=None):\n"
        "        arr, sr = _sf.read(str(uri), dtype='float32', always_2d=False)\n"
        "        if arr.ndim == 1: arr = arr[_np.newaxis, :]\n"
        "        elif channels_first: arr = arr.T\n"
        "        return _torch.from_numpy(arr.copy()), sr\n"
        "    torchaudio.load = _sf_load\n"
        "    print('sitecustomize: torchaudio.load patched to use soundfile')\n"
        "except Exception as _e: print(f'sitecustomize patch failed: {_e}')\n",
        encoding="utf-8",
    )
    print(f"Installed torchaudio soundfile patch -> {sc}")


def fine_tune() -> Path:
    _patch_torchaudio()
    env = {**os.environ, "HF_TOKEN": HF_TOKEN}

    cmd = [
        sys.executable, "-m", "f5_tts.train.finetune_cli",
        "--exp_name", EXP_ARCH,
        "--dataset_name", DATASET_NAME,
        "--finetune",                          # auto-downloads F5TTS_v1_Base base ckpt
        "--tokenizer", "pinyin",               # matches base model's 2546-token vocab (avoids embed size mismatch)
        "--epochs", "10",
        "--batch_size_per_gpu", "4",
        "--batch_size_type", "sample",         # use sample count, not frame count
        "--learning_rate", "1e-4",
        "--num_warmup_updates", "100",
        "--save_per_updates", "500",
        "--last_per_updates", "200",
    ]

    print("Running:", " ".join(cmd))
    result = subprocess.run(cmd, env=env, check=False)

    if result.returncode != 0:
        # Show what was actually written so we know if partial ckpts exist
        ckpt_root = _pkg_ckpt_root()
        print(f"Checkpoint root: {ckpt_root}")
        for p in sorted(ckpt_root.rglob("*"))[:30]:
            print(" ", p)
        raise RuntimeError(f"Fine-tune exited {result.returncode}")

    # Checkpoints land at {site-packages}/ckpts/{dataset_name}/
    ckpt_dir = _pkg_ckpt_root() / DATASET_NAME
    if not ckpt_dir.exists():
        # Fall back to any ckpts subdir
        ckpt_dir = _pkg_ckpt_root()
    return ckpt_dir


def push_model(ckpt_dir: Path) -> None:
    print(f"Pushing {ckpt_dir} -> {MODEL_REPO}...")
    api = HfApi(token=HF_TOKEN)
    api.create_repo(MODEL_REPO, exist_ok=True, private=True, repo_type="model")
    api.upload_folder(
        folder_path=str(ckpt_dir),
        repo_id=MODEL_REPO,
        commit_message=f"F5-TTS fine-tune: LJSpeech, {NUM_SAMPLES} samples, 10 epochs",
    )
    print(f"Done -> https://huggingface.co/{MODEL_REPO}")


if __name__ == "__main__":
    prepare_data()
    ckpt_dir = fine_tune()
    push_model(ckpt_dir)