#!/usr/bin/env python3 """ JSON-manifest driven MFA alignment using the "Libri" alignment pipeline. MFA CLI resolution is delegated to ``mfa_align.resolve_mfa_bin`` (active conda env first, then legacy ``mfa_env`` sidecar). Prefer installing MFA in the same env as Whisper/PhonePaint when possible. Why this exists: - `mfa_align_libri.py` contains the correct model-pair + enhanced-dictionary (G2P) fallback logic. - `mfa_align_cli.py` is a wrapper around `mfa_align.py` + `merge_mfa_json.py`, which does *not* use that newer logic. This script reuses the core pipeline from `mfa_align_libri.py`, but swaps the input stage: - instead of LibriSpeech tree + *.trans.txt, we accept a JSON manifest specifying audio + text. Input JSON formats supported (auto-detected): 1) Dict mapping: { "audio.wav": {"text": "hello world"} , ... } 2) Dict mapping: { "audio.wav": "hello world", ... } 3) List of objects: [ {"audio": "audio.wav", "text": "hello world"}, ... ] 4) List of strings: [ "audio.wav", ... ] (requires --sidecar_text_ext, default: .txt) In this mode, transcript is read from a sidecar file next to the audio, e.g. audio.wav + audio.txt Typical usage: python3 mfa_align_json.py --audio_dir /path/to/audio --input_json manifest.json --output_json aligned.json \\ --mfa_model_pairs english_mfa:english_mfa --use_g2p_fallback --dump_pre_g2p """ from __future__ import annotations import argparse import json import os import shutil import tempfile from pathlib import Path from typing import Dict, List, Tuple, Iterable, Any from textgrid import IntervalTier, TextGrid # pip install textgrid from mfa_align import ( _ffmpeg_convert, parse_model_pairs, run_incremental_alignment_with_model_pairs, detect_fail_words_from_textgrid, patch_fail_words_from_fallback, run_mfa, extract_words_from_texts, detect_oov_words, generate_g2p_pronunciations, create_enhanced_dictionary, ) def _pick_tiers(tg: TextGrid) -> Tuple[IntervalTier | None, IntervalTier | None]: phones_tier = None words_tier = None for t in tg.tiers: if not isinstance(t, IntervalTier): continue lname = (t.name or "").lower() if ("phone" in lname) or ("segment" in lname): phones_tier = t if "word" in lname: words_tier = t return phones_tier, words_tier def _build_flat_aligned( out_dir: Path, included_wavs: List[str], wav_to_key: Dict[str, str] ) -> Dict[str, dict]: """Create JSON with top-level words and phones (no nesting).""" aligned: Dict[str, dict] = {} for wav_name in included_wavs: tg_path = out_dir / Path(wav_name).with_suffix(".TextGrid") if not tg_path.exists(): continue tg = TextGrid() tg.read(tg_path) phones_tier, words_tier = _pick_tiers(tg) if words_tier is None: continue words_dict: Dict[str, dict] = {} phones_dict: Dict[str, dict] = {} for w in words_tier: w_text = (w.mark or "").strip() if not w_text: continue words_dict[str(len(words_dict))] = { "xmin": round(w.minTime, 3), "xmax": round(w.maxTime, 3), "text": w_text, } if phones_tier is not None: for p in phones_tier: p_text = (p.mark or "").strip() if not p_text: continue phones_dict[str(len(phones_dict))] = { "xmin": round(p.minTime, 3), "xmax": round(p.maxTime, 3), "text": p_text, } top_key = wav_to_key.get(wav_name, wav_name) aligned[top_key] = {"words": words_dict, "phones": phones_dict} return aligned def _normalize_manifest( manifest: Any, *, audio_dir: Path, sidecar_text_ext: str, ) -> Dict[str, Dict[str, str]]: """ Normalize supported manifest formats into: { key: {"audio": str, "text": str} } Where key is what will appear in the output JSON (typically the audio filename). """ def read_sidecar_text(audio_path: Path) -> str: sidecar = audio_path.with_suffix(sidecar_text_ext) if not sidecar.exists(): raise FileNotFoundError(f"Missing sidecar transcript: {sidecar}") return sidecar.read_text(encoding="utf-8").strip() out: Dict[str, Dict[str, str]] = {} if isinstance(manifest, dict): for k, v in manifest.items(): if isinstance(v, dict): text = v.get("text", "") audio = v.get("audio", k) else: text = str(v) audio = k if not str(text).strip(): # allow empty text entries but skip them continue out[str(k)] = {"audio": str(audio), "text": str(text)} return out if isinstance(manifest, list): for item in manifest: if isinstance(item, str): audio = item audio_path = Path(audio) if not audio_path.is_absolute(): audio_path = audio_dir / audio_path text = read_sidecar_text(audio_path) out[item] = {"audio": item, "text": text} elif isinstance(item, dict): audio = item.get("audio") or item.get("path") or item.get("file") text = item.get("text") if audio is None: raise ValueError(f"Manifest object missing 'audio': {item}") if text is None: audio_path = Path(str(audio)) if not audio_path.is_absolute(): audio_path = audio_dir / audio_path text = read_sidecar_text(audio_path) key = str(item.get("key") or audio) if not str(text).strip(): continue out[key] = {"audio": str(audio), "text": str(text)} else: raise ValueError(f"Unsupported manifest list item type: {type(item)}") return out raise ValueError(f"Unsupported manifest type: {type(manifest)}") def _iter_unique_ids(keys: Iterable[str]) -> Iterable[Tuple[str, str]]: """ Yield (key, utt_id) pairs where utt_id is safe/unique for MFA corpus filenames. Uses audio stem as base, disambiguates with _{n} if needed. """ seen: Dict[str, int] = {} for key in keys: base = Path(key).stem or "utt" n = seen.get(base, 0) seen[base] = n + 1 utt_id = base if n == 0 else f"{base}_{n}" yield key, utt_id def main() -> int: ap = argparse.ArgumentParser( description="Align audio specified by a JSON manifest using the mfa_align_libri pipeline." ) ap.add_argument("--audio_dir", type=Path, required=True, help="Base directory for relative audio paths.") ap.add_argument( "--input_json", type=Path, required=True, help="Manifest JSON (supported formats described in module docstring).", ) ap.add_argument("--output_json", type=Path, required=True, help="Output alignment JSON path.") ap.add_argument( "--wav_output", type=Path, help="Optional directory to save converted 16k mono WAVs (mirrors relative path under audio_dir when possible).", ) ap.add_argument( "--sidecar_text_ext", default=".txt", help="Used only when manifest entries do not contain text (default: .txt).", ) ap.add_argument( "--gain", type=float, default=None, help="Optional volume gain (dB) applied when converting audio to 16k mono WAV.", ) ap.add_argument("-j", "--threads", type=int, default=os.cpu_count() or 4) ap.add_argument( "--omp_threads", type=int, default=1, help="OpenMP/BLAS threads to use inside MFA/ffmpeg subprocesses (default: 1).", ) # Must match `mfa_align.py`'s argument shape: nargs='*', action='append' # so `parse_model_pairs()` receives a list-of-lists. ap.add_argument("--mfa_model_pairs", nargs="*", action="append") ap.add_argument("--mfa_dict_fallbacks", nargs="+", default=["english_mfa"]) ap.add_argument("--mfa_acoustic_fallbacks", nargs="+", default=["english_mfa"]) ap.add_argument("--use_g2p_fallback", action="store_true", help="After english_mfa, G2P only for utterances with fail/OOV words.") ap.add_argument("--g2p_model", default="english_us_mfa") ap.add_argument("--pre_g2p_output_json", type=Path, help="Where to write a pre-G2P snapshot JSON.") ap.add_argument("--dump_pre_g2p", action="store_true", help="Enable pre-G2P snapshot emission.") ap.add_argument("--fallback_tolerance_ms", type=float, default=50.0) ap.add_argument("--fail_word_eps_ms", type=float, default=10.0) ap.add_argument("--debug_keep_tmp", action="store_true") ap.add_argument("--update", action="store_true", help="Skip entries already present in output_json.") args = ap.parse_args() args.audio_dir = args.audio_dir.resolve() if not args.audio_dir.exists(): raise FileNotFoundError(f"--audio_dir not found: {args.audio_dir}") existing_keys: set[str] = set() if args.update and args.output_json.exists(): try: existing = json.loads(args.output_json.read_text(encoding="utf-8")) if isinstance(existing, dict): existing_keys = set(existing.keys()) except Exception: existing_keys = set() manifest_raw = json.loads(args.input_json.read_text(encoding="utf-8")) norm = _normalize_manifest( manifest_raw, audio_dir=args.audio_dir, sidecar_text_ext=args.sidecar_text_ext ) if args.update and existing_keys: norm = {k: v for k, v in norm.items() if k not in existing_keys} if not norm: print("All entries already present in output_json. Nothing to do.") return 0 model_pairs = parse_model_pairs(args) print(f"Using model pairs: {' → '.join([f'{a}+{d}' for a, d in model_pairs])}") # Prepare temp corpus if args.debug_keep_tmp: tmp_dir = Path(tempfile.mkdtemp(prefix="mfa_json_")) cleanup = False else: tmp_dir = Path(tempfile.mkdtemp(prefix="mfa_json_")) cleanup = True text_data: Dict[str, dict] = {} included_wavs: List[str] = [] wav_to_key: Dict[str, str] = {} try: # Isolate MFA working dirs per run, but reuse cached pretrained models via symlink os.environ["MFA_FALLBACK_TOL"] = f"{max(0.0, (args.fallback_tolerance_ms / 1000.0)):.3f}" mfa_tmp = tmp_dir / "mfa_tmp" mfa_root = tmp_dir / "mfa_root" mfa_cache = mfa_root / "cache" mfa_tmp.mkdir(parents=True, exist_ok=True) mfa_root.mkdir(parents=True, exist_ok=True) mfa_cache.mkdir(parents=True, exist_ok=True) os.environ["MFA_TEMP_DIR"] = str(mfa_tmp) os.environ["MFA_ROOT"] = str(mfa_root) os.environ["MFA_ROOT_DIR"] = str(mfa_root) os.environ["MFA_CACHE_DIR"] = str(mfa_cache) pretrained_candidates = [ Path.home() / "Documents" / "MFA" / "pretrained_models", Path.home() / ".local" / "share" / "Montreal-Forced-Alignment" / "pretrained_models", ] for cand in pretrained_candidates: if cand.exists(): target = mfa_root / "pretrained_models" if not target.exists(): try: target.symlink_to(cand, target_is_directory=True) except Exception: pass break omp_threads = max(1, int(args.omp_threads)) os.environ["OMP_NUM_THREADS"] = str(omp_threads) os.environ.setdefault("OPENBLAS_NUM_THREADS", str(omp_threads)) os.environ.setdefault("MKL_NUM_THREADS", str(omp_threads)) os.environ.setdefault("NUMEXPR_NUM_THREADS", str(omp_threads)) # Build the MFA corpus (wav + lab files) keys = list(norm.keys()) for key, utt_id in _iter_unique_ids(keys): entry = norm[key] audio_rel = Path(entry["audio"]) audio_path = audio_rel if audio_rel.is_absolute() else (args.audio_dir / audio_rel) if not audio_path.exists(): print(f"[warn] missing audio: {audio_path} (key={key}); skipping") continue text = (entry.get("text") or "").strip() if not text: print(f"[warn] empty transcript for key={key}; skipping") continue wav_name = f"{utt_id}.wav" dst_wav = tmp_dir / wav_name _ffmpeg_convert(audio_path, dst_wav, gain_db=args.gain) (tmp_dir / f"{utt_id}.lab").write_text(text, encoding="utf-8") # Optional WAV export: mirror relative path under audio_dir when possible if args.wav_output: args.wav_output.mkdir(parents=True, exist_ok=True) try: rel = audio_path.resolve().relative_to(args.audio_dir) export = args.wav_output / rel export = export.with_suffix(".wav") except Exception: export = args.wav_output / f"{utt_id}.wav" export.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(dst_wav, export) text_data[wav_name] = {"text": text} included_wavs.append(wav_name) wav_to_key[wav_name] = key if not included_wavs: raise RuntimeError("No valid (audio,text) entries found after filtering.") pre_g2p_path = None if args.use_g2p_fallback and args.dump_pre_g2p: if args.pre_g2p_output_json is None: pre_g2p_path = args.output_json.with_name(args.output_json.stem + "_pre_g2p.json") else: pre_g2p_path = args.pre_g2p_output_json out_dir, _stats = run_incremental_alignment_with_model_pairs( tmp_dir, model_pairs, args.threads, text_data=text_data, use_g2p=False, g2p_model=args.g2p_model, pre_g2p_json_path=pre_g2p_path, included_fns=included_wavs, ) # Always dump a fresh pre-G2P snapshot derived from TextGrids (same behavior as mfa_align_libri) if pre_g2p_path is not None: pre_aligned = _build_flat_aligned(out_dir, included_wavs, wav_to_key) pre_g2p_path.parent.mkdir(parents=True, exist_ok=True) pre_g2p_path.write_text( json.dumps(pre_aligned, indent=2, ensure_ascii=False), encoding="utf-8" ) print(f"[pre-G2P] JSON snapshot written to: {pre_g2p_path}") # Detect fail_words per file, optionally run G2P-only fallback on a subset, and patch only fail segments eps_s = max(0.0, float(args.fail_word_eps_ms) / 1000.0) fail_words_by_wav: Dict[str, List[int]] = {} fail_wavs: List[str] = [] for wav_name in included_wavs: tg_path = out_dir / Path(wav_name).with_suffix(".TextGrid") if not tg_path.exists(): continue idxs = detect_fail_words_from_textgrid(tg_path, eps_s=eps_s) if idxs: fail_words_by_wav[wav_name] = idxs fail_wavs.append(wav_name) if args.use_g2p_fallback and fail_wavs: g2p_corpus = tmp_dir / "corpus_g2p_fail_only" g2p_corpus.mkdir(parents=True, exist_ok=True) text_data_fail: Dict[str, dict] = {} for wav_name in fail_wavs: src_wav = tmp_dir / wav_name dst_wav = g2p_corpus / wav_name src_lab = tmp_dir / Path(wav_name).with_suffix(".lab").name dst_lab = g2p_corpus / Path(wav_name).with_suffix(".lab").name if not dst_wav.exists(): try: os.symlink(src_wav, dst_wav) except Exception: shutil.copy2(src_wav, dst_wav) if src_lab.exists() and not dst_lab.exists(): try: os.symlink(src_lab, dst_lab) except Exception: shutil.copy2(src_lab, dst_lab) text_data_fail[wav_name] = text_data[wav_name] primary_acoustic, primary_dict = model_pairs[0] tmp_g2p_dir = g2p_corpus / "_g2p" tmp_g2p_dir.mkdir(exist_ok=True) all_words = extract_words_from_texts(text_data_fail) oov_words = detect_oov_words(all_words, primary_dict) g2p_out_dir = None if oov_words and generate_g2p_pronunciations(oov_words, args.g2p_model, tmp_g2p_dir): enhanced = create_enhanced_dictionary( primary_dict, tmp_g2p_dir / "g2p_pronunciations.dict", tmp_g2p_dir, ) if enhanced is not None and enhanced.exists(): g2p_out_dir = run_mfa( g2p_corpus, primary_acoustic, str(enhanced), args.threads, out_dir=(g2p_corpus / "out_g2p"), tag="g2p", ) if g2p_out_dir is not None: patch_fail_words_from_fallback( out_dir, g2p_out_dir, fail_words_by_wav, eps_s=eps_s, convert_arpabet=False, ) aligned = _build_flat_aligned(out_dir, included_wavs, wav_to_key) # Merge with existing output if --update was provided if args.update and args.output_json.exists(): try: existing = json.loads(args.output_json.read_text(encoding="utf-8")) if isinstance(existing, dict): existing.update(aligned) aligned = existing except Exception: pass args.output_json.parent.mkdir(parents=True, exist_ok=True) args.output_json.write_text(json.dumps(aligned, indent=2, ensure_ascii=False), encoding="utf-8") print(f"[done] wrote {len(aligned)} entries to {args.output_json}") return 0 finally: if cleanup: shutil.rmtree(tmp_dir, ignore_errors=True) else: print(f"[keep] temp at {tmp_dir}") if __name__ == "__main__": raise SystemExit(main())