Add train_narrator_tts.py

train_narrator_tts.py

@@ -20,9 +20,10 @@ 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}/
-  - HFDataset accesses row["audio"]["array"] (numpy) + row["audio"]["sampling_rate"]
-  - Audio decode via soundfile (not torchcodec) -- so store paths, not arrays
+  - 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)
@@ -55,7 +56,6 @@ MAX_SECS = 12.0
 
 DATASET_NAME = "narrator"
 EXP_ARCH = "F5TTS_v1_Base"
-WAVS_DIR = Path("/tmp/f5_work/wavs")
 
 
 def _lib_root() -> Path:
@@ -75,7 +75,7 @@ def _pkg_ckpt_root() -> Path:
 
 
 def prepare_data() -> None:
-    """Download LJSpeech, write WAVs, build HF dataset + vocab at paths F5-TTS expects."""
+    """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")
 
@@ -87,9 +87,12 @@ def prepare_data() -> None:
             tar.extractall("/tmp/", filter="data")
         print(f"Extracted -> {lj_root}")
 
-    WAVS_DIR.mkdir(parents=True, exist_ok=True)
-
-    wav_paths: list[str] = []
+    # Store audio as pre-decoded float arrays — NO hf_datasets.Audio feature.
+    # Using Audio feature triggers torchcodec at read time, which requires FFmpeg
+    # system libs that aren't available in HF Jobs. Storing plain dicts sidesteps this:
+    # row["audio"]["array"] is a Python list[float], row["audio"]["sampling_rate"] is int.
+    # F5-TTS CustomDataset calls torch.FloatTensor(row["audio"]["array"]) which accepts lists.
+    audio_data: list[dict] = []
     texts: list[str] = []
     durations: list[float] = []
     count = 0
@@ -119,9 +122,7 @@ def prepare_data() -> None:
                 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))
+            audio_data.append({"array": arr.tolist(), "sampling_rate": TARGET_SR})
             texts.append(text)
             durations.append(len(arr) / TARGET_SR)
             count += 1
@@ -129,17 +130,12 @@ def prepare_data() -> None:
     print(f"Prepared {count} samples")
 
     # F5-TTS dataset.py builds the path as data/{dataset_name}_{tokenizer}/raw/
-    # The base model (F5TTS_v1_Base) was trained with --tokenizer pinyin (2546-token vocab).
-    # We MUST use pinyin to keep text_embed.weight shape [2546, 512] matching the checkpoint.
-    # Pinyin tokenizer handles English ASCII fine (maps each char through its vocab).
+    # 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)
 
-    # Pass file PATHS (strings) to Audio feature — encode_example for a string stores
-    # {"path": str, "bytes": None} without torchcodec. Decode at training time uses soundfile.
-    ds = hf_datasets.Dataset.from_dict({"audio": wav_paths, "text": texts})
-    ds = ds.cast_column("audio", hf_datasets.Audio(sampling_rate=TARGET_SR))
+    ds = hf_datasets.Dataset.from_dict({"audio": audio_data, "text": texts})
     ds.save_to_disk(str(raw_dir))
     print(f"Saved dataset -> {raw_dir}")