Michael Treadgold
453 lines
16 KiB
Python
453 lines
16 KiB
Python
"""
|
|
Preprocess an audio+text dataset into the durations.jsonl format required
|
|
by Inflect-Nano training scripts.
|
|
|
|
Supports:
|
|
- HuggingFace datasets (CMU ARCTIC, LJSpeech, etc.)
|
|
- Local directory of .wav files with matching .txt transcriptions
|
|
- LJSpeech-format metadata.csv
|
|
|
|
Output: a .jsonl file with one JSON object per line containing:
|
|
phone_ids, tone_ids, lang_ids, hifigan_durations, target_audio, speaker_id, voice_id
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
import torchaudio
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parent
|
|
VENDORED_FRONTEND = REPO_ROOT / "third_party" / "tiny_tts_frontend"
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
sys.path.insert(0, str(VENDORED_FRONTEND))
|
|
|
|
from tiny_tts.nn import commons
|
|
from tiny_tts.text import phonemes_to_ids
|
|
from tiny_tts.text.english import grapheme_to_phoneme, normalize_text
|
|
from tiny_tts.utils import ADD_BLANK
|
|
|
|
from inflect_nano.text_cleaning import clean_tinytts_text
|
|
from inflect_nano.vocoder import HifiGanConfig, MelFrontend
|
|
|
|
|
|
def text_to_ids(text: str) -> tuple[list[int], list[int], list[int]]:
|
|
"""Convert English text to TinyTTS phone/tone/lang ID lists."""
|
|
cleaned = clean_tinytts_text(text)
|
|
normalized = normalize_text(cleaned)
|
|
phones, tones, _ = grapheme_to_phoneme(normalized)
|
|
phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
|
|
if ADD_BLANK:
|
|
phone_ids = commons.insert_blanks(phone_ids, 0)
|
|
tone_ids = commons.insert_blanks(tone_ids, 0)
|
|
lang_ids = commons.insert_blanks(lang_ids, 0)
|
|
return phone_ids, tone_ids, lang_ids
|
|
|
|
|
|
def estimate_durations_uniform(
|
|
phone_ids: list[int],
|
|
audio_path: str,
|
|
mel_frontend: MelFrontend,
|
|
sample_rate: int = 24000,
|
|
) -> list[int]:
|
|
"""Estimate mel-frame durations uniformly across all phones.
|
|
|
|
Loads the audio, extracts mel spectrogram, then distributes the total
|
|
frame count evenly across phones. This is a coarse estimate suitable
|
|
for decoder-only fine-tuning where the duration predictor is frozen.
|
|
"""
|
|
wav, sr = torchaudio.load(audio_path)
|
|
if wav.shape[0] > 1:
|
|
wav = wav.mean(dim=0, keepdim=True)
|
|
if sr != sample_rate:
|
|
wav = torchaudio.functional.resample(wav, sr, sample_rate)
|
|
wav = wav.clamp(-1, 1)
|
|
|
|
with torch.no_grad():
|
|
mel = mel_frontend(wav)
|
|
total_frames = mel.shape[-1]
|
|
|
|
# Count non-blank phones (id != 0)
|
|
visible_phones = [i for i, pid in enumerate(phone_ids) if pid != 0]
|
|
|
|
if not visible_phones:
|
|
return [0] * len(phone_ids)
|
|
|
|
# Distribute frames: each blank gets 1 frame, rest split evenly
|
|
blank_count = len(phone_ids) - len(visible_phones)
|
|
remaining = max(0, total_frames - blank_count)
|
|
base_dur = max(1, remaining // len(visible_phones))
|
|
remainder = remaining - base_dur * len(visible_phones)
|
|
|
|
durations = []
|
|
for pid in phone_ids:
|
|
if pid == 0:
|
|
durations.append(1)
|
|
else:
|
|
extra = 1 if remainder > 0 else 0
|
|
remainder = max(0, remainder - 1)
|
|
durations.append(base_dur + extra)
|
|
|
|
# Scale to match total_frames
|
|
current_sum = sum(durations)
|
|
if current_sum > 0 and current_sum != total_frames:
|
|
scale = total_frames / current_sum
|
|
scaled = [max(1, round(d * scale)) for d in durations]
|
|
# Fix rounding errors
|
|
diff = total_frames - sum(scaled)
|
|
for i in range(abs(diff)):
|
|
if diff > 0:
|
|
scaled[i % len(scaled)] += 1
|
|
else:
|
|
if scaled[i % len(scaled)] > 1:
|
|
scaled[i % len(scaled)] -= 1
|
|
durations = scaled
|
|
|
|
# Ensure no zero durations for non-blank phones
|
|
for i, pid in enumerate(phone_ids):
|
|
if pid != 0 and durations[i] < 1:
|
|
durations[i] = 1
|
|
|
|
return durations
|
|
|
|
|
|
def process_hf_dataset(
|
|
dataset_path: str,
|
|
output_jsonl: Path,
|
|
audio_dir: Path | None,
|
|
subset: str = "default",
|
|
split: str | None = None,
|
|
text_key: str = "text",
|
|
speaker_key: str = "speaker",
|
|
max_rows: int = 0,
|
|
voice_id: str = "speaker",
|
|
) -> int:
|
|
"""Process a HuggingFace dataset (loaded via `datasets` library)."""
|
|
from datasets import load_dataset
|
|
|
|
print(f"Loading HF dataset: {dataset_path} subset={subset} split={split}")
|
|
if split:
|
|
ds = load_dataset(dataset_path, subset, split=split)
|
|
else:
|
|
ds_dict = load_dataset(dataset_path, subset)
|
|
splits = list(ds_dict.keys())
|
|
print(f"Available splits: {splits}")
|
|
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
|
|
skip_no_text = 0
|
|
skip_no_audio = 0
|
|
skip_bad_ids = 0
|
|
try:
|
|
with output_jsonl.open("w", encoding="utf-8") as f:
|
|
for i, row in enumerate(ds):
|
|
text = str(row.get(text_key, "")).strip()
|
|
if not text:
|
|
skip_no_text += 1
|
|
if skip_no_text <= 3:
|
|
print(f" Row {i}: no text (keys={list(row.keys())[:5]})")
|
|
continue
|
|
|
|
# Get audio path or array
|
|
audio_info = row.get("audio", row.get("file", None))
|
|
if audio_info is None:
|
|
skip_no_audio += 1
|
|
if skip_no_audio <= 3:
|
|
print(f" Row {i}: no audio/file key (keys={list(row.keys())[:5]})")
|
|
continue
|
|
|
|
if isinstance(audio_info, dict):
|
|
# Audio loaded as array — save to temp file
|
|
audio_array = audio_info.get("array")
|
|
sr = audio_info.get("sampling_rate", 24000)
|
|
if audio_array is None:
|
|
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():
|
|
continue
|
|
else:
|
|
continue
|
|
|
|
try:
|
|
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
|
|
|
|
durations = estimate_durations_uniform(phone_ids, audio_path, mel_frontend)
|
|
|
|
speaker = str(row.get(speaker_key, voice_id))
|
|
|
|
row_out = {
|
|
"phone_ids": phone_ids,
|
|
"tone_ids": tone_ids,
|
|
"lang_ids": lang_ids,
|
|
"hifigan_durations": durations,
|
|
"target_audio": str(Path(audio_path).resolve()),
|
|
"speaker_id": hash(speaker) % 256,
|
|
"voice_id": speaker,
|
|
}
|
|
f.write(json.dumps(row_out, ensure_ascii=False) + "\n")
|
|
count += 1
|
|
|
|
if max_rows > 0 and count >= max_rows:
|
|
break
|
|
|
|
if count % 100 == 0:
|
|
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}")
|
|
return count
|
|
|
|
|
|
def process_local_dir(
|
|
audio_dir: Path,
|
|
output_jsonl: Path,
|
|
ext: str = ".wav",
|
|
text_ext: str = ".txt",
|
|
speaker: str = "speaker",
|
|
max_rows: int = 0,
|
|
) -> int:
|
|
"""Process a local directory of .wav files with matching .txt files."""
|
|
audio_files = sorted(audio_dir.glob(f"*{ext}"))
|
|
if not audio_files:
|
|
audio_files = sorted(audio_dir.rglob(f"*{ext}"))
|
|
|
|
mel_frontend = MelFrontend(HifiGanConfig(variant="v2plus"))
|
|
output_jsonl.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
count = 0
|
|
with output_jsonl.open("w", encoding="utf-8") as f:
|
|
for wav_path in audio_files:
|
|
txt_path = wav_path.with_suffix(text_ext)
|
|
if not txt_path.is_file():
|
|
continue
|
|
|
|
text = txt_path.read_text(encoding="utf-8").strip()
|
|
if not text:
|
|
continue
|
|
|
|
try:
|
|
phone_ids, tone_ids, lang_ids = text_to_ids(text)
|
|
except Exception as e:
|
|
print(f" Skipping {wav_path.name}: {e}")
|
|
continue
|
|
|
|
durations = estimate_durations_uniform(phone_ids, str(wav_path), mel_frontend)
|
|
|
|
row_out = {
|
|
"phone_ids": phone_ids,
|
|
"tone_ids": tone_ids,
|
|
"lang_ids": lang_ids,
|
|
"hifigan_durations": durations,
|
|
"target_audio": str(wav_path.resolve()),
|
|
"speaker_id": hash(speaker) % 256,
|
|
"voice_id": speaker,
|
|
}
|
|
f.write(json.dumps(row_out, ensure_ascii=False) + "\n")
|
|
count += 1
|
|
|
|
if max_rows > 0 and count >= max_rows:
|
|
break
|
|
|
|
if count % 100 == 0:
|
|
print(f" Processed {count} rows...")
|
|
|
|
print(f"Wrote {count} rows to {output_jsonl}")
|
|
return count
|
|
|
|
|
|
def process_metadata_csv(
|
|
csv_path: Path,
|
|
audio_dir: Path,
|
|
output_jsonl: Path,
|
|
sep: str = "|",
|
|
header: bool = False,
|
|
text_col: int = 1,
|
|
file_col: int = 0,
|
|
speaker: str = "speaker",
|
|
max_rows: int = 0,
|
|
) -> int:
|
|
"""Process an LJSpeech-style metadata.csv."""
|
|
mel_frontend = MelFrontend(HifiGanConfig(variant="v2plus"))
|
|
output_jsonl.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
count = 0
|
|
with csv_path.open("r", encoding="utf-8") as csv_file:
|
|
lines = csv_file.readlines()
|
|
if header:
|
|
lines = lines[1:]
|
|
|
|
with output_jsonl.open("w", encoding="utf-8") as out_f:
|
|
for line in lines:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
|
|
parts = line.split(sep)
|
|
if len(parts) < max(text_col, file_col) + 1:
|
|
continue
|
|
|
|
filename = parts[file_col].strip()
|
|
if not filename.endswith(".wav"):
|
|
filename += ".wav"
|
|
wav_path = audio_dir / filename
|
|
if not wav_path.is_file():
|
|
print(f" Missing: {wav_path}")
|
|
continue
|
|
|
|
text = parts[text_col].strip()
|
|
if not text:
|
|
continue
|
|
|
|
try:
|
|
phone_ids, tone_ids, lang_ids = text_to_ids(text)
|
|
except Exception as e:
|
|
print(f" Skipping {filename}: {e}")
|
|
continue
|
|
|
|
durations = estimate_durations_uniform(phone_ids, str(wav_path), mel_frontend)
|
|
|
|
row_out = {
|
|
"phone_ids": phone_ids,
|
|
"tone_ids": tone_ids,
|
|
"lang_ids": lang_ids,
|
|
"hifigan_durations": durations,
|
|
"target_audio": str(wav_path.resolve()),
|
|
"speaker_id": hash(speaker) % 256,
|
|
"voice_id": speaker,
|
|
}
|
|
out_f.write(json.dumps(row_out, ensure_ascii=False) + "\n")
|
|
count += 1
|
|
|
|
if max_rows > 0 and count >= max_rows:
|
|
break
|
|
|
|
if count % 100 == 0:
|
|
print(f" Processed {count} rows...")
|
|
|
|
print(f"Wrote {count} rows to {output_jsonl}")
|
|
return count
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser(
|
|
description="Preprocess a speech dataset for Inflect-Nano training.",
|
|
)
|
|
|
|
sub = ap.add_subparsers(dest="mode", required=True)
|
|
|
|
# HuggingFace mode
|
|
p_hf = sub.add_parser("hf", help="Process a HuggingFace dataset")
|
|
p_hf.add_argument("--dataset", required=True, help="HF dataset path, e.g. MikhailT/cmu-arctic")
|
|
p_hf.add_argument("--subset", default="default")
|
|
p_hf.add_argument("--split", default=None, help="e.g. 'rms' for CMU ARCTIC single speaker")
|
|
p_hf.add_argument("--text-key", default="text")
|
|
p_hf.add_argument("--speaker-key", default="speaker")
|
|
p_hf.add_argument("--audio-dir", type=Path, default=None, help="Local dir for audio if HF paths are remote")
|
|
p_hf.add_argument("--out", type=Path, required=True, help="Output .jsonl path")
|
|
p_hf.add_argument("--max-rows", type=int, default=0)
|
|
p_hf.add_argument("--voice-id", default="speaker")
|
|
|
|
# Local dir mode
|
|
p_dir = sub.add_parser("local", help="Process local directory of .wav+.txt pairs")
|
|
p_dir.add_argument("--audio-dir", type=Path, required=True)
|
|
p_dir.add_argument("--ext", default=".wav")
|
|
p_dir.add_argument("--text-ext", default=".txt")
|
|
p_dir.add_argument("--speaker", default="speaker")
|
|
p_dir.add_argument("--out", type=Path, required=True)
|
|
p_dir.add_argument("--max-rows", type=int, default=0)
|
|
|
|
# Metadata CSV mode (LJSpeech-style)
|
|
p_csv = sub.add_parser("csv", help="Process a metadata.csv file")
|
|
p_csv.add_argument("--csv", type=Path, required=True)
|
|
p_csv.add_argument("--audio-dir", type=Path, required=True)
|
|
p_csv.add_argument("--sep", default="|")
|
|
p_csv.add_argument("--header", action="store_true")
|
|
p_csv.add_argument("--file-col", type=int, default=0)
|
|
p_csv.add_argument("--text-col", type=int, default=1)
|
|
p_csv.add_argument("--speaker", default="speaker")
|
|
p_csv.add_argument("--out", type=Path, required=True)
|
|
p_csv.add_argument("--max-rows", type=int, default=0)
|
|
|
|
args = ap.parse_args()
|
|
|
|
if args.mode == "hf":
|
|
process_hf_dataset(
|
|
dataset_path=args.dataset,
|
|
output_jsonl=args.out,
|
|
audio_dir=args.audio_dir,
|
|
subset=args.subset,
|
|
split=args.split,
|
|
text_key=args.text_key,
|
|
speaker_key=args.speaker_key,
|
|
max_rows=args.max_rows,
|
|
voice_id=args.voice_id,
|
|
)
|
|
elif args.mode == "local":
|
|
process_local_dir(
|
|
audio_dir=args.audio_dir,
|
|
output_jsonl=args.out,
|
|
ext=args.ext,
|
|
text_ext=args.text_ext,
|
|
speaker=args.speaker,
|
|
max_rows=args.max_rows,
|
|
)
|
|
elif args.mode == "csv":
|
|
process_metadata_csv(
|
|
csv_path=args.csv,
|
|
audio_dir=args.audio_dir,
|
|
output_jsonl=args.out,
|
|
sep=args.sep,
|
|
header=args.header,
|
|
file_col=args.file_col,
|
|
text_col=args.text_col,
|
|
speaker=args.speaker,
|
|
max_rows=args.max_rows,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|