Fix Audio object handling in HF datasets (has .array attr, not a plain dict)

colab_run.py

@@ -0,0 +1,170 @@
+"""
+Colab runner for Inflect-Nano smoothness fine-tuning.
+
+Single script that handles the full pipeline with graceful checkpointing —
+safe to re-run after interruptions.  Each step is skipped if its output
+already exists.
+
+Usage in Colab:
+    !wget https://your-server/colab_run.py
+    !python colab_run.py
+"""
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+# ---------------------------------------------------------------------------
+# Config — change these
+# ---------------------------------------------------------------------------
+REPO_URL = "https://git.quietidiot.com/michael-treadgold/speech-nano.git"
+REPO_DIR = Path("/content/speech-nano")
+DATASET = "MikhailT/cmu-arctic"
+SPEAKER = "rms"                # "rms", "bdl", "jmk", "awb", "ksp" (CMU ARCTIC male speakers)
+TRAIN_STEPS = 5000
+BATCH_SIZE = 8
+
+# ---------------------------------------------------------------------------
+# Step helpers
+# ---------------------------------------------------------------------------
+def step(name: str) -> None:
+    print(f"\n{'='*60}\n  {name}\n{'='*60}")
+
+def run(cmd: str, **kwargs) -> bool:
+    print(f"  $ {cmd}")
+    result = subprocess.run(cmd, shell=True, **kwargs)
+    return result.returncode == 0
+
+def already(path: Path | str) -> bool:
+    p = Path(path)
+    ok = p.exists()
+    if ok:
+        print(f"  ⏭  Already exists: {p}")
+    return ok
+
+def pip_installed(pkg: str) -> bool:
+    result = subprocess.run(
+        [sys.executable, "-c", f"import {pkg}"],
+        capture_output=True, text=True, timeout=30,
+    )
+    return result.returncode == 0
+
+
+def main() -> None:
+    start = time.time()
+
+    # ---- Step 1: Clone or pull repo ----
+    step("1/6  Clone / pull repo")
+    if REPO_DIR.exists():
+        print(f"  Repo exists, pulling latest...")
+        subprocess.run(["git", "-C", str(REPO_DIR), "pull"], check=False)
+        subprocess.run(["git", "-C", str(REPO_DIR), "lfs", "pull"], check=False)
+    else:
+        run(f"git clone {REPO_URL} {REPO_DIR}")
+    sys.path.insert(0, str(REPO_DIR))
+    sys.path.insert(0, str(REPO_DIR / "third_party" / "tiny_tts_frontend"))
+    os.chdir(REPO_DIR)
+    print(f"  Working dir: {REPO_DIR}")
+
+    # ---- Step 2: Install dependencies ----
+    step("2/6  Install dependencies")
+    pkgs = "torch torchaudio soundfile numpy g2p_en transformers gradio numba scipy datasets"
+    if not pip_installed("g2p_en"):
+        run(f"{sys.executable} -m pip install -q {pkgs}")
+    else:
+        print("  ⏭  Already installed (g2p_en found)")
+
+    import nltk
+    nltk.download("averaged_perceptron_tagger_eng", quiet=True)
+    nltk.download("cmudict", quiet=True)
+    print("  NLTK data OK")
+
+    # ---- Step 3: Download model weights ----
+    step("3/6  Download model weights")
+    WEIGHTS_DIR = REPO_DIR / "weights"
+    AC_WEIGHTS = WEIGHTS_DIR / "inflect_nano_v1_acoustic.pt"
+    VO_WEIGHTS = WEIGHTS_DIR / "inflect_nano_v1_vocoder.pt"
+
+    if AC_WEIGHTS.exists() and AC_WEIGHTS.stat().st_size > 1000:
+        print(f"  ⏭  Acoustic weights OK ({AC_WEIGHTS.stat().st_size / 1e6:.1f} MB)")
+    else:
+        run(f"{sys.executable} download_weights.py")
+
+    # ---- Step 4: Preprocess dataset ----
+    step("4/6  Preprocess dataset")
+    JSONL = Path(f"/content/durations_{SPEAKER}.jsonl")
+
+    if JSONL.exists() and JSONL.stat().st_size > 100:
+        print(f"  ⏭  JSONL exists ({JSONL.stat().st_size / 1e6:.2f} MB)")
+    else:
+        ok = run(
+            f"{sys.executable} preprocess_dataset.py hf "
+            f"--dataset {DATASET} --split {SPEAKER} "
+            f"--out {JSONL} --voice-id mark"
+        )
+        if not ok:
+            print("  ❌  Preprocessing failed. Check error above.")
+            sys.exit(1)
+
+    # ---- Step 5: Train ----
+    step("5/6  Train smoothness model")
+    OUT_DIR = Path("/content/checkpoints/smooth-v1")
+    LATEST_CKPT = OUT_DIR / "inflect-smooth-latest.pt"
+
+    import torch
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"  Device: {device}  ({torch.cuda.get_device_name(0) if device == 'cuda' else 'CPU'})")
+
+    if LATEST_CKPT.exists():
+        print(f"  ⏭  Checkpoint exists — adding {TRAIN_STEPS} more steps")
+        resume_flag = "--resume"
+    else:
+        print(f"  Starting fresh — {TRAIN_STEPS} steps")
+        resume_flag = ""
+
+    ok = run(
+        f"{sys.executable} -m inflect_nano.train_smooth "
+        f"--durations-jsonl {JSONL} "
+        f"--out-dir {OUT_DIR} "
+        f"--init-checkpoint {AC_WEIGHTS} "
+        f"--vocoder-checkpoint {VO_WEIGHTS} "
+        f"--steps {TRAIN_STEPS} "
+        f"--batch-size {BATCH_SIZE} "
+        f"--device {device} "
+        f"{resume_flag}"
+    )
+    if not ok:
+        print("  ❌  Training failed. Check error above.")
+        sys.exit(1)
+
+    # ---- Step 6: Generate sample ----
+    step("6/6  Generate sample")
+    from inference import load_acoustic, load_vocoder, synthesize
+    import soundfile as sf
+    import numpy as np
+
+    ac_path = LATEST_CKPT if LATEST_CKPT.exists() else AC_WEIGHTS
+    acoustic, speakers, ap = load_acoustic(ac_path, torch.device(device))
+    vocoder, vp = load_vocoder(VO_WEIGHTS, torch.device(device))
+    print(f"  Acoustic: {ap:,} params  Vocoder: {vp:,} params  Total: {ap+vp:,}")
+
+    text = "Every man is destined to die, but his work echoes through the ages."
+    audio = synthesize(
+        text, acoustic, vocoder, speakers, torch.device(device),
+        smooth_prosody=True, mel_smooth_sigma=0.8,
+    )
+    out = Path("/content/sample_smooth.wav")
+    sf.write(str(out), audio, 24000, subtype="PCM_16")
+
+    elapsed = (time.time() - start) / 60
+    print(f"\n{'='*60}")
+    print(f"  Done in {elapsed:.0f}m.  Sample: {out}  ({audio.size/24000:.1f}s)")
+    print(f"  Checkpoints: {OUT_DIR}")
+    print(f"{'='*60}")
+
+
+if __name__ == "__main__":
+    main()

colab_smooth_finetune.ipynb

@@ -1,281 +1,47 @@
-# Inflect-Nano Smoothness Fine-Tuning — Google Colab Notebook
+{
-
+ "cells": [
-This notebook fine-tunes [Inflect-Nano-v1](https://huggingface.co/owensong/Inflect-Nano-v1) with enhanced smoothness losses.  
+  {
-Runs on a **free T4 GPU** in Colab.  Trains in ~3-6 hours.
+   "cell_type": "markdown",
-
+   "metadata": {},
----
+   "source": [
-
+    "# Inflect-Nano Smoothness Fine-Tuning\n",
-## Cell 1: Setup — clone repo, install deps
+    "Single-click pipeline. Each cell is safe to re-run \u2014 it skips steps already completed.\n",
-
+    "\n",
-```python
+    "**Runtime \u2192 Change runtime type \u2192 T4 GPU**"
-# @title Setup environment (run once)
+   ]
-import os, sys, subprocess
+  },
-from pathlib import Path
+  {
-
+   "cell_type": "code",
-REPO_URL = "https://huggingface.co/owensong/Inflect-Nano-v1"
+   "metadata": {
-REPO_DIR = "/content/Inflect-Nano-v1"
+    "id": "runner"
-
+   },
-# Clone
+   "source": [
-if not Path(REPO_DIR).exists():
+    "# Download the runner script and execute it\n",
-    !git clone {REPO_URL} {REPO_DIR}
+    "import os, sys, subprocess\n",
-else:
+    "from pathlib import Path\n",
-    %cd {REPO_DIR}
+    "\n",
-    !git pull
+    "REPO = Path(\"/content/speech-nano\")\n",
-
+    "if not REPO.exists():\n",
-%cd {REPO_DIR}
+    "    subprocess.run([\"git\", \"clone\", \"https://git.quietidiot.com/michael-treadgold/speech-nano.git\", str(REPO)], check=False)\n",
-
+    "else:\n",
-# Install deps (numba is needed by vendored frontend; scipy for lowpass)
+    "    subprocess.run([\"git\", \"-C\", str(REPO), \"pull\"], check=False)\n",
-!pip install -q torch torchaudio soundfile numpy g2p_en transformers gradio numba scipy datasets
+    "\n",
-
+    "sys.path.insert(0, str(REPO))\n",
-# Download NLTK data
+    "os.chdir(REPO)\n",
-import nltk
+    "\n",
-nltk.download('averaged_perceptron_tagger_eng', quiet=True)
+    "# Run the pipeline\n",
-nltk.download('cmudict', quiet=True)
+    "exec(open(\"colab_run.py\").read())"
-
+   ],
-print("✓ Setup complete")
+   "execution_count": null,
-print(f"  PyTorch {torch.__version__}  |  GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A'}")
+   "outputs": []
-```
+  }
-
+ ],
----
+ "metadata": {
-
+  "accelerator": "GPU",
-## Cell 2: Pick a dataset & preprocess
+  "colab": {
-
+   "provenance": []
-Choose one:
+  }
-
+ },
-| Dataset | Speakers | Size | Best for |
+ "nbformat": 4,
-|---------|----------|------|----------|
+ "nbformat_minor": 5
-| **CMU ARCTIC (rms)** | 1 US male | ~1.1K clips, ~1h | Fast fine-tune |
+}
-| **CMU ARCTIC (bdl)** | 1 US male | ~1.1K clips, ~1h | Fast fine-tune |
-| **LJSpeech** | 1 US female | 13.1K clips, ~24h | Best quality but female voice |
-
-```python
-# @title Select dataset and preprocess
-DATASET = "MikhailT/cmu-arctic"  # @param ["MikhailT/cmu-arctic", "keithito/lj_speech"]
-SPEAKER_SPLIT = "rms"           # @param ["rms", "bdl", "jmk", "awb", "ksp"] (for CMU ARCTIC)
-MAX_ROWS = 0                    # 0 = all rows, or set e.g. 500
-
-import sys
-sys.path.insert(0, str(Path.cwd()))
-
-from preprocess_dataset import process_hf_dataset
-
-OUT_JSONL = Path(f"/content/durations_{SPEAKER_SPLIT}.jsonl")
-
-print(f"Preprocessing {DATASET} / {SPEAKER_SPLIT} ...")
-n = process_hf_dataset(
-    dataset_path=DATASET,
-    output_jsonl=OUT_JSONL,
-    audio_dir=None,
-    split=SPEAKER_SPLIT,
-    max_rows=MAX_ROWS,
-    voice_id=SPEAKER_SPLIT,
-)
-print(f"✓ Wrote {n} rows to {OUT_JSONL}")
-```
-
----
-
-## Cell 3: Fine-tune with enhanced smoothness losses
-
-This uses the `train_smooth.py` module we added — it trains with:
-
-- Multi-resolution STFT loss (penalises buzz)
-- Adversarial mel discriminator (pushes toward realistic spectrograms)
-- Vocoder consistency loss (ensures mels work through the vocoder)
-- Deeper residual postnet
-- Cosine LR schedule with warmup
-
-```python
-# @title Run smoothness fine-tuning
-import torch
-import sys
-sys.path.insert(0, str(Path.cwd()))
-sys.path.insert(0, str(Path.cwd() / "third_party" / "tiny_tts_frontend"))
-
-from inflect_nano.train_smooth import train_smooth
-import argparse
-
-# Build args programmatically
-class Args:
-    durations_jsonl = OUT_JSONL
-    out_dir = Path("/content/checkpoints/smooth-v1")
-    max_rows = 0
-    steps = 5000              # 5K steps is ~2h on T4 for CMU ARCTIC
-    batch_size = 6
-    lr = 2e-4
-    weight_decay = 1e-4
-    warmup_steps = 500
-
-    # Architecture (keep same as original)
-    hidden = 168
-    encoder_layers = 5
-    decoder_layers = 6
-    decoder_ff_mult = 3
-    max_seconds = 12.0
-    max_frames = 1400
-    postnet_scale = 0.35      # Higher postnet influence
-    postnet_layers = 5        # Deeper residual postnet
-    postnet_kernel = 5
-    abs_frame_bins = 512
-
-    # Loss weights (tuned for smoothness)
-    mse_weight = 0.25
-    delta_weight = 0.25       # Higher spectral smoothness
-    accel_weight = 0.08       # Enable acceleration loss
-    duration_weight = 0.08
-    group_duration_weight = 0.02
-    energy_weight = 0.06
-    bright_weight = 0.06
-    pitch_weight = 0.06
-
-    # New smoothness losses
-    stft_weight = 0.15        # Multi-resolution STFT loss
-    predicted_prosody_mel_weight = 0.05
-    predicted_prosody_delta_weight = 0.03
-    robust_prosody_mix = 0.5
-    robust_prosody_mel_weight = 0.05
-    robust_prosody_delta_weight = 0.03
-    adv_mel_weight = 0.08     # Adversarial mel loss
-    fm_weight = 2.0
-
-    # Vocoder consistency
-    vocoder_checkpoint = Path(REPO_DIR) / "weights" / "inflect_nano_v1_vocoder.pt"
-    vocoder_wav_weight = 0.12
-    vocoder_mel_weight = 0.08
-
-    # Checkpointing
-    init_checkpoint = Path(REPO_DIR) / "weights" / "inflect_nano_v1_acoustic.pt"
-    save_interval = 1000
-    log_interval = 25
-    seed = 42
-    resume = False
-    preload_features = False  # T4 has GPU memory but not tons of RAM
-
-    # Misc
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    trainable = "all"
-    contextual_predictors = False
-    group_duration_planner = False
-    grad_clip = 5.0
-
-args = Args()
-
-print(f"Device: {args.device}")
-print(f"Checkpoint: {args.init_checkpoint}")
-print(f"Output dir: {args.out_dir}")
-print(f"Steps: {args.steps}  Batch: {args.batch_size}")
-
-train_smooth(args)
-```
-
----
-
-## Cell 4: Package the trained model for download
-
-```python
-# @title Export trained model
-import shutil, torch
-from pathlib import Path
-
-CHECKPOINT_DIR = Path("/content/checkpoints/smooth-v1")
-EXPORT_DIR = Path("/content/inflect-nano-smooth-export")
-
-# Find latest checkpoint
-checkpoints = sorted(CHECKPOINT_DIR.glob("inflect-smooth-*.pt"))
-if not checkpoints:
-    print("No checkpoints found!")
-else:
-    latest = checkpoints[-1]
-    print(f"Latest checkpoint: {latest.name}  ({latest.stat().st_size / 1e6:.1f} MB)")
-
-    ckpt = torch.load(latest, map_location="cpu")
-
-    EXPORT_DIR.mkdir(exist_ok=True)
-
-    # Save as inference-format checkpoint (same format as original)
-    acoustic_export = EXPORT_DIR / "inflect_nano_v1_acoustic_smooth.pt"
-    torch.save({
-        "model": ckpt["model"],
-        "config": ckpt.get("config", {}),  # will be loaded from original
-        "speakers": ckpt.get("speakers", {"mark": 0}),
-        "params": ckpt.get("params", 0),
-        "step": ckpt.get("step", 0),
-    }, acoustic_export)
-    print(f"Exported acoustic model: {acoustic_export}")
-
-    # Also copy the original vocoder (unchanged)
-    vocoder_src = Path(REPO_DIR) / "weights" / "inflect_nano_v1_vocoder.pt"
-    vocoder_dst = EXPORT_DIR / "inflect_nano_v1_vocoder.pt"
-    shutil.copy(vocoder_src, vocoder_dst)
-    print(f"Copied vocoder: {vocoder_dst}")
-
-    # Zip for download
-    !cd /content && zip -r inflect-nano-smooth.zip inflect-nano-smooth-export/
-    print(f"\n✓ Download: /content/inflect-nano-smooth.zip")
-```
-
----
-
-## Cell 5: Generate a sample
-
-```python
-# @title Test the fine-tuned model
-import sys, torch, numpy as np, soundfile as sf
-from pathlib import Path
-
-sys.path.insert(0, str(Path.cwd()))
-sys.path.insert(0, str(Path.cwd() / "third_party" / "tiny_tts_frontend"))
-
-from inference import load_acoustic, load_vocoder, synthesize
-
-TEXT = "Every man is destined to die, but his work echoes through the ages."  # @param {type:"string"}
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Load fine-tuned acoustic
-acoustic_path = EXPORT_DIR / "inflect_nano_v1_acoustic_smooth.pt"
-if not acoustic_path.exists():
-    acoustic_path = Path(REPO_DIR) / "weights" / "inflect_nano_v1_acoustic.pt"
-    print("Using original model (fine-tuned not found)")
-
-vocoder_path = Path(REPO_DIR) / "weights" / "inflect_nano_v1_vocoder.pt"
-
-acoustic, speakers, ap = load_acoustic(acoustic_path, device)
-vocoder, vp = load_vocoder(vocoder_path, device)
-
-print(f"Acoustic: {ap:,} params  Vocoder: {vp:,} params  Total: {ap+vp:,}")
-
-audio = synthesize(
-    TEXT, acoustic, vocoder, speakers, device,
-    smooth_prosody=True,
-    mel_smooth_sigma=0.8,
-)
-
-out_path = Path("/content/sample_smooth.wav")
-sf.write(str(out_path), audio, 24000, subtype="PCM_16")
-print(f"✓ Wrote {out_path}  ({audio.size / 24000:.1f}s)")
-```
-
----
-
-## Usage notes
-
-1. Upload this notebook to [Colab](https://colab.research.google.com/)
-2. Select **Runtime → Change runtime type → T4 GPU**
-3. Run cells 1-5 in order
-4. Download `inflect-nano-smooth.zip` from the Files panel
-
-### What the fine-tuning actually does
-
-| Loss | Weight | What it improves |
-|------|--------|-----------------|
-| Multi-resolution STFT | 0.15 | Spectral smoothness — directly penalises buzzy artifacts |
-| Adversarial mel | 0.08 | Realistic spectrogram texture |
-| Vocoder consistency (wav) | 0.12 | Ensures generated mels voice cleanly through vocoder |
-| Vocoder consistency (mel) | 0.08 | Mel-reconstruction fidelity |
-| Delta (spectral derivative) | 0.25 | Frame-to-frame smoothness |
-| Acceleration (2nd deriv) | 0.08 | Reduces jitter/stutter |
-| Prosody exposure bias | 0.05 | Trains with predicted (not reference) prosody |
-| Robust prosody mix | 0.05 | Mixed-mode prosody for stability |
-
-All new losses are **training-only** — the exported model has the same 4.6M params as the original (plus ~260K for the deeper postnet).

download_weights.py

@@ -0,0 +1,72 @@
+"""Download Inflect-Nano-v1 model weights from HuggingFace.
+
+Run this once after cloning the repo:
+    python download_weights.py
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+# Use requests or urllib — try both
+try:
+    import requests
+    _has_requests = True
+except ImportError:
+    import urllib.request
+    import urllib.error
+    _has_requests = False
+
+WEIGHTS_DIR = Path(__file__).resolve().parent / "weights"
+HF_BASE = "https://huggingface.co/owensong/Inflect-Nano-v1/resolve/main/weights"
+
+FILES = [
+    "inflect_nano_v1_acoustic.pt",
+    "inflect_nano_v1_vocoder.pt",
+]
+
+
+def download_file(url: str, dest: Path) -> None:
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    print(f"Downloading {dest.name}...", end=" ", flush=True)
+
+    if _has_requests:
+        resp = requests.get(url, stream=True)
+        resp.raise_for_status()
+        total = int(resp.headers.get("content-length", 0))
+        with dest.open("wb") as f:
+            downloaded = 0
+            for chunk in resp.iter_content(chunk_size=8192):
+                f.write(chunk)
+                downloaded += len(chunk)
+                if total:
+                    pct = downloaded * 100 // total
+                    print(f"\rDownloading {dest.name}... {pct}%", end="", flush=True)
+        print(" done.")
+    else:
+        urllib.request.urlretrieve(url, str(dest))
+        print("done.")
+
+
+def main() -> None:
+    missing = [f for f in FILES if not (WEIGHTS_DIR / f).is_file()]
+    if not missing:
+        print("All weights already present.")
+        return
+
+    print(f"Missing {len(missing)} weight file(s). Downloading from HuggingFace...\n")
+    for filename in missing:
+        url = f"{HF_BASE}/{filename}"
+        dest = WEIGHTS_DIR / filename
+        try:
+            download_file(url, dest)
+        except Exception as e:
+            print(f"\n  FAILED: {e}")
+            print(f"  Manual download: {url}")
+            sys.exit(1)
+
+    print(f"\nAll weights downloaded to {WEIGHTS_DIR}")
+
+
+if __name__ == "__main__":
+    main()

preprocess_dataset.py

@@ -133,85 +133,122 @@ def process_hf_dataset(
 
     print(f"Loading HF dataset: {dataset_path}  subset={subset}  split={split}")
     if split:
-        ds = load_dataset(dataset_path, subset, split=split, trust_remote_code=True)
+        ds = load_dataset(dataset_path, subset, split=split)
     else:
-        ds_dict = load_dataset(dataset_path, subset, trust_remote_code=True)
+        ds_dict = load_dataset(dataset_path, subset)
         splits = list(ds_dict.keys())
         print(f"Available splits: {splits}")
-        # Use first train split, or first available
         preferred = [s for s in splits if "train" in s.lower()]
         ds = ds_dict[preferred[0] if preferred else splits[0]]
 
+    # Debug: print first row to see available columns
+    first = next(iter(ds))
+    print(f"Dataset columns: {list(first.keys())}")
+    for k, v in first.items():
+        if isinstance(v, dict):
+            print(f"  {k}: dict with keys {list(v.keys())}")
+        elif isinstance(v, str):
+            print(f"  {k}: str ({len(v)} chars) = {v[:80]!r}")
+        else:
+            print(f"  {k}: {type(v).__name__}")
+
+    import tempfile
+
     mel_frontend = MelFrontend(HifiGanConfig(variant="v2plus"))
     output_jsonl.parent.mkdir(parents=True, exist_ok=True)
 
+    # Temp dir for audio extracted from HF arrays (cleaned up on exit)
+    audio_tmp = Path(tempfile.mkdtemp(prefix="inflect_audio_"))
+    print(f"Audio tmp dir: {audio_tmp}")
+
     count = 0
-    with output_jsonl.open("w", encoding="utf-8") as f:
+    skip_no_text = 0
-        for i, row in enumerate(ds):
+    skip_no_audio = 0
-            text = str(row.get(text_key, "")).strip()
+    skip_bad_ids = 0
-            if not text:
+    skip_bad_audio = 0
-                continue
+    try:
-
+        with output_jsonl.open("w", encoding="utf-8") as f:
-            # Get audio path or array
+            for i, row in enumerate(ds):
-            audio_info = row.get("audio", row.get("file", None))
+                text = str(row.get(text_key, "")).strip()
-            if audio_info is None:
+                if not text:
-                continue
+                    skip_no_text += 1
-
+                    if skip_no_text <= 3:
-            if isinstance(audio_info, dict):
+                        print(f"  Row {i}: no text (keys={list(row.keys())[:5]})")
-                # Audio is already loaded as array
-                audio_path = None
-                audio_array = audio_info.get("array")
-                sample_rate = audio_info.get("sampling_rate", 24000)
-                if audio_array is None:
                     continue
-            elif isinstance(audio_info, str):
+
-                audio_path = audio_info
+                # Get audio path or array
-                if audio_dir:
+                audio_info = row.get("audio", row.get("file", None))
-                    audio_path = str(audio_dir / Path(audio_info).name)
+                if audio_info is None:
-                if not Path(audio_path).is_file():
+                    skip_no_audio += 1
+                    if skip_no_audio <= 3:
+                        print(f"  Row {i}: no audio/file key (keys={list(row.keys())[:5]})")
                     continue
-                audio_array = None
-                sample_rate = 24000  # will be detected on load
-            else:
-                continue
 
-            try:
+                # HF Audio feature returns an object with .array / .sampling_rate / .path
-                phone_ids, tone_ids, lang_ids = text_to_ids(text)
+                if hasattr(audio_info, "array"):
-            except Exception as e:
+                    audio_array = audio_info["array"] if isinstance(audio_info, dict) else audio_info.array
-                print(f"  Skipping row {i}: text-to-ids failed: {e}")
+                    sr = audio_info["sampling_rate"] if isinstance(audio_info, dict) else audio_info.sampling_rate
-                continue
+                    audio_path = str(audio_tmp / f"utt_{i:06d}.wav")
+                    sf.write(audio_path, audio_array, sr, subtype="PCM_16")
+                elif isinstance(audio_info, dict):
+                    audio_array = audio_info.get("array")
+                    sr = audio_info.get("sampling_rate", 24000)
+                    if audio_array is None:
+                        skip_bad_audio += 1
+                        continue
+                    audio_path = str(audio_tmp / f"utt_{i:06d}.wav")
+                    sf.write(audio_path, audio_array, sr, subtype="PCM_16")
+                elif isinstance(audio_info, str):
+                    audio_path = audio_info
+                    if audio_dir:
+                        audio_path = str(audio_dir / Path(audio_info).name)
+                    if not Path(audio_path).is_file():
+                        skip_bad_audio += 1
+                        continue
+                else:
+                    skip_bad_audio += 1
+                    if skip_bad_audio <= 3:
+                        print(f"  Row {i}: unexpected audio type: {type(audio_info).__name__}")
+                    continue
 
-            if not phone_ids:
+                try:
-                continue
+                    phone_ids, tone_ids, lang_ids = text_to_ids(text)
+                except Exception as e:
+                    skip_bad_ids += 1
+                    if skip_bad_ids <= 3:
+                        print(f"  Row {i}: text-to-ids failed: {e}  text={text[:60]!r}")
+                    continue
+
+                if not phone_ids:
+                    continue
 
-            # Estimate durations
-            if audio_path:
                 durations = estimate_durations_uniform(phone_ids, audio_path, mel_frontend)
-            else:
-                # Uniform fallback: 8 frames per phone
-                durations = [8] * len(phone_ids)
 
-            speaker = str(row.get(speaker_key, voice_id))
+                speaker = str(row.get(speaker_key, voice_id))
 
-            row_out = {
+                row_out = {
-                "phone_ids": phone_ids,
+                    "phone_ids": phone_ids,
-                "tone_ids": tone_ids,
+                    "tone_ids": tone_ids,
-                "lang_ids": lang_ids,
+                    "lang_ids": lang_ids,
-                "hifigan_durations": durations,
+                    "hifigan_durations": durations,
-                "target_audio": audio_path or "",
+                    "target_audio": str(Path(audio_path).resolve()),
-                "speaker_id": hash(speaker) % 256,
+                    "speaker_id": hash(speaker) % 256,
-                "voice_id": speaker,
+                    "voice_id": speaker,
-            }
+                }
-            f.write(json.dumps(row_out, ensure_ascii=False) + "\n")
+                f.write(json.dumps(row_out, ensure_ascii=False) + "\n")
-            count += 1
+                count += 1
 
-            if max_rows > 0 and count >= max_rows:
+                if max_rows > 0 and count >= max_rows:
-                break
+                    break
 
-            if count % 100 == 0:
+                if count % 100 == 0:
-                print(f"  Processed {count} rows...")
+                    print(f"  Processed {count} rows...")
+    finally:
+        # Audio tmp files must survive until training is done, so we keep them
+        pass
 
     print(f"Wrote {count} rows to {output_jsonl}")
+    print(f"Skipped: no_text={skip_no_text}  no_audio={skip_no_audio}  bad_ids={skip_bad_ids}  bad_audio={skip_bad_audio}")
     return count