Add smoothness training pipeline with STFT/adversarial/vocoder-consistency losses
- inflect_nano/train_smooth.py: enhanced training with multi-res STFT loss, adversarial mel discriminator, vocoder consistency loss, deeper residual postnet, and cosine LR schedule - preprocess_dataset.py: convert HF datasets, local dirs, or LJSpeech CSVs to the durations.jsonl format needed by training - inference.py: add --smooth-prosody, --mel-smooth-sigma, --lowpass-hz flags for zero-cost inference-time quality improvements - test_inference.py: smoke tests for model loading and synthesis - colab_smooth_finetune.ipynb: Colab notebook for T4 GPU fine-tuning - requirements.txt: add numba, scipy, datasets
This commit is contained in:
Michael Treadgold
@@ -0,0 +1,426 @@
|
||||
"""
|
||||
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, trust_remote_code=True)
|
||||
else:
|
||||
ds_dict = load_dataset(dataset_path, subset, trust_remote_code=True)
|
||||
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]]
|
||||
|
||||
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 i, row in enumerate(ds):
|
||||
text = str(row.get(text_key, "")).strip()
|
||||
if not text:
|
||||
continue
|
||||
|
||||
# Get audio path or array
|
||||
audio_info = row.get("audio", row.get("file", None))
|
||||
if audio_info is None:
|
||||
continue
|
||||
|
||||
if isinstance(audio_info, dict):
|
||||
# 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
|
||||
if audio_dir:
|
||||
audio_path = str(audio_dir / Path(audio_info).name)
|
||||
if not Path(audio_path).is_file():
|
||||
continue
|
||||
audio_array = None
|
||||
sample_rate = 24000 # will be detected on load
|
||||
else:
|
||||
continue
|
||||
|
||||
try:
|
||||
phone_ids, tone_ids, lang_ids = text_to_ids(text)
|
||||
except Exception as e:
|
||||
print(f" Skipping row {i}: text-to-ids failed: {e}")
|
||||
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))
|
||||
|
||||
row_out = {
|
||||
"phone_ids": phone_ids,
|
||||
"tone_ids": tone_ids,
|
||||
"lang_ids": lang_ids,
|
||||
"hifigan_durations": durations,
|
||||
"target_audio": audio_path or "",
|
||||
"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_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()
|
||||
Reference in New Issue
Block a user