Spaces:
Running on Zero
Running on Zero
| """Segmentation, transcription, and MFA/MAPS alignment backends.""" | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import logging | |
| import os | |
| import shutil | |
| import subprocess | |
| import sys | |
| import tempfile | |
| import time | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Dict, List, Optional, Sequence, Tuple | |
| from pipeline_utils import all_inputs_shorter_than, copy_as_single_segments, pick_asr_device_idx, write_aligned_json | |
| from unified_speech import segment_wav_directory | |
| from .constants import MAPS_ALIGN_SCRIPT, MFA_ALIGN_SCRIPT, MIN_SPLIT_SEC, REPO_ROOT, UNIFIED_SPEECH_SCRIPT | |
| from .runtime import resolve_mfa_bin, run_python | |
| log = logging.getLogger(__name__) | |
| def prepare_input(input_path: Path, ws: dict[str, Path]) -> Path: | |
| input_path = input_path.resolve() | |
| if input_path.is_dir(): | |
| return input_path | |
| ws["input"].mkdir(parents=True, exist_ok=True) | |
| dst = ws["input"] / input_path.name | |
| if not dst.exists() or dst.stat().st_mtime < input_path.stat().st_mtime: | |
| shutil.copy2(input_path, dst) | |
| return ws["input"] | |
| def run_segment( | |
| input_dir: Path, | |
| ws: dict[str, Path], | |
| *, | |
| silence_level: float, | |
| ) -> List[str]: | |
| """Applio-identical silence split in-process.""" | |
| ws["segments"].mkdir(parents=True, exist_ok=True) | |
| if all_inputs_shorter_than(input_dir, MIN_SPLIT_SEC): | |
| log.info( | |
| "All inputs shorter than %.1fs — skipping silence-based split", | |
| MIN_SPLIT_SEC, | |
| ) | |
| return copy_as_single_segments(input_dir, ws["segments"]) | |
| return segment_wav_directory( | |
| input_dir, | |
| ws["segments"], | |
| silence_level=silence_level, | |
| min_split_sec=MIN_SPLIT_SEC, | |
| ) | |
| def run_transcribe_whisper( | |
| ws: dict[str, Path], | |
| *, | |
| model_name: str, | |
| batch_size: int, | |
| language: str, | |
| device: str, | |
| device_idx: Optional[int] = None, | |
| ) -> None: | |
| cmd = [ | |
| "python", | |
| str(UNIFIED_SPEECH_SCRIPT), | |
| "transcribe", | |
| "--input-dir", | |
| str(ws["segments"]), | |
| "--output-json", | |
| str(ws["transcripts"]), | |
| "--model", | |
| model_name, | |
| "--language", | |
| language, | |
| "--device", | |
| device, | |
| "--batch-size", | |
| str(batch_size), | |
| ] | |
| if device.startswith("cuda"): | |
| idx = pick_asr_device_idx(device_idx) | |
| cmd.extend(["--device-idx", str(idx)]) | |
| run_python(cmd, cwd=REPO_ROOT) | |
| def run_transcribe(ws: dict[str, Path], args: argparse.Namespace) -> None: | |
| run_transcribe_whisper( | |
| ws, | |
| model_name=args.asr_model, | |
| batch_size=args.batch_size, | |
| language=args.language, | |
| device=args.asr_device, | |
| device_idx=args.asr_device_idx, | |
| ) | |
| class MfaAlignJob: | |
| """One MFA unit of work (usually one ``work/<stem>/`` tree).""" | |
| stem: str | |
| ws: dict[str, Path] | |
| segment_names: List[str] | |
| # Special ``--align-threads`` value: MFA -j = number of segment wavs for that input. | |
| ALIGN_THREADS_SEGMENTS = "segments" | |
| def parse_align_threads(value: str): | |
| """Argparse type: positive int, or ``segments`` / ``auto`` (= per-input segment count).""" | |
| key = str(value).strip().lower().replace("_", "-") | |
| if key in { | |
| "segments", "auto", "nseg", "n-segments", | |
| "per-segment", "per-segments", "per-file", | |
| }: | |
| return ALIGN_THREADS_SEGMENTS | |
| try: | |
| n = int(value) | |
| except ValueError as exc: | |
| raise argparse.ArgumentTypeError( | |
| f"--align-threads must be a positive int or 'segments', got {value!r}" | |
| ) from exc | |
| if n < 1: | |
| raise argparse.ArgumentTypeError("--align-threads must be >= 1") | |
| return n | |
| def resolve_mfa_align_threads( | |
| align_threads: object, | |
| segment_names: Sequence[str], | |
| ) -> int: | |
| """Map ``--align-threads`` to MFA ``--num_jobs`` for one input's segments.""" | |
| if align_threads == ALIGN_THREADS_SEGMENTS: | |
| return max(1, len(segment_names)) | |
| return max(1, int(align_threads)) # type: ignore[arg-type] | |
| def _transcript_text(entry: object) -> str: | |
| if isinstance(entry, dict): | |
| return str(entry.get("text") or "").strip() | |
| return str(entry or "").strip() | |
| def _mfa_align_env(env_mfa: str, *, threads_per_job: int) -> dict[str, str]: | |
| mfa_bin = resolve_mfa_bin(env_mfa) | |
| env = os.environ.copy() | |
| env["MFA_BIN"] = mfa_bin | |
| env["MFA_CONDA_ENV"] = env_mfa | |
| if threads_per_job <= 1: | |
| for key in ( | |
| "OMP_NUM_THREADS", | |
| "OPENBLAS_NUM_THREADS", | |
| "MKL_NUM_THREADS", | |
| "NUMEXPR_NUM_THREADS", | |
| ): | |
| env[key] = "1" | |
| return env | |
| def _mfa_align_cmd( | |
| *, | |
| audio_dir: Path, | |
| input_json: Path, | |
| output_json: Path, | |
| threads: int, | |
| ) -> List[str]: | |
| return [ | |
| sys.executable, | |
| str(MFA_ALIGN_SCRIPT), | |
| "--audio_dir", | |
| str(audio_dir), | |
| "--input_json", | |
| str(input_json), | |
| "--output_json", | |
| str(output_json), | |
| # english_mfa only; G2P runs only for files that still have fail/OOV words | |
| # after the primary dictionary pass (no english_us_arpa fallback pair). | |
| "--use_g2p_fallback", | |
| "--dump_pre_g2p", | |
| "--mfa_model_pairs", | |
| "english_mfa:english_mfa", | |
| "-j", | |
| str(max(1, int(threads))), | |
| ] | |
| def _finalize_mfa_workspace( | |
| ws: dict[str, Path], | |
| segment_names: Sequence[str], | |
| per_segment: dict, | |
| *, | |
| pre_g2p: Optional[dict] = None, | |
| ) -> None: | |
| """Write ``03_aligned*.json`` contracts from a per-segment MFA dict.""" | |
| per_segment_path = ws["aligned"].with_name("03_aligned_per_segment.json") | |
| per_segment_path.parent.mkdir(parents=True, exist_ok=True) | |
| per_segment_path.write_text( | |
| json.dumps(per_segment, indent=2, ensure_ascii=False) + "\n", | |
| encoding="utf-8", | |
| ) | |
| missing = [name for name in segment_names if name not in per_segment] | |
| if missing: | |
| raise RuntimeError(f"MFA alignment omitted segments: {missing}") | |
| write_aligned_json(per_segment, segment_names, ws["segments"], ws["aligned"]) | |
| if pre_g2p is not None: | |
| pre_path = per_segment_path.with_name("03_aligned_per_segment_pre_g2p.json") | |
| pre_path.write_text( | |
| json.dumps(pre_g2p, indent=2, ensure_ascii=False) + "\n", | |
| encoding="utf-8", | |
| ) | |
| write_aligned_json( | |
| pre_g2p, | |
| segment_names, | |
| ws["segments"], | |
| ws["aligned"].with_name("03_aligned_pre_g2p.json"), | |
| ) | |
| def run_align_mfa( | |
| ws: dict[str, Path], | |
| segment_names: Sequence[str], | |
| *, | |
| env_mfa: str, | |
| threads: int, | |
| ) -> None: | |
| """MFA JSON wrapper in-process; Kaldi work via MFA CLI (same env or fallback).""" | |
| if not MFA_ALIGN_SCRIPT.is_file(): | |
| raise SystemExit(f"MFA align script not found: {MFA_ALIGN_SCRIPT}") | |
| mfa_bin = resolve_mfa_bin(env_mfa) | |
| per_segment_path = ws["aligned"].with_name("03_aligned_per_segment.json") | |
| env = _mfa_align_env(env_mfa, threads_per_job=threads) | |
| cmd = _mfa_align_cmd( | |
| audio_dir=ws["segments"], | |
| input_json=ws["transcripts"], | |
| output_json=per_segment_path, | |
| threads=threads, | |
| ) | |
| log.info("Running MFA CLI (%s): %s", mfa_bin, " ".join(cmd)) | |
| subprocess.run( | |
| cmd, | |
| cwd=str(MFA_ALIGN_SCRIPT.parent), | |
| check=True, | |
| env=env, | |
| ) | |
| per_segment = json.loads(per_segment_path.read_text(encoding="utf-8")) | |
| pre_g2p_src = per_segment_path.with_name("03_aligned_per_segment_pre_g2p.json") | |
| pre = ( | |
| json.loads(pre_g2p_src.read_text(encoding="utf-8")) | |
| if pre_g2p_src.is_file() | |
| else None | |
| ) | |
| _finalize_mfa_workspace(ws, segment_names, per_segment, pre_g2p=pre) | |
| def run_align_mfa_batch_merged( | |
| jobs: Sequence[MfaAlignJob], | |
| *, | |
| env_mfa: str, | |
| threads: int, | |
| ) -> None: | |
| """Strategy 1: one ``mfa align`` over a merged corpus, then scatter JSON. | |
| Amortizes MFA cold-start across many short PhonePaint workspaces. Unique | |
| utterance keys encode ``stem`` + segment name so results can be remapped | |
| back to each ``03_aligned.json`` contract. | |
| """ | |
| if not jobs: | |
| return | |
| if len(jobs) == 1: | |
| job = jobs[0] | |
| run_align_mfa( | |
| job.ws, job.segment_names, env_mfa=env_mfa, threads=threads, | |
| ) | |
| return | |
| if not MFA_ALIGN_SCRIPT.is_file(): | |
| raise SystemExit(f"MFA align script not found: {MFA_ALIGN_SCRIPT}") | |
| # key → (stem, original segment wav name) | |
| key_map: Dict[str, Tuple[str, str]] = {} | |
| merged_manifest: Dict[str, dict] = {} | |
| for job in jobs: | |
| transcripts = json.loads(job.ws["transcripts"].read_text(encoding="utf-8")) | |
| if not isinstance(transcripts, dict): | |
| raise RuntimeError( | |
| f"Expected dict transcripts for stem={job.stem}; got {type(transcripts)}" | |
| ) | |
| for name in job.segment_names: | |
| audio_path = job.ws["segments"] / name | |
| if not audio_path.is_file(): | |
| raise FileNotFoundError(f"Missing segment wav: {audio_path}") | |
| text = _transcript_text(transcripts.get(name)) | |
| if not text: | |
| raise RuntimeError( | |
| f"Empty transcript for stem={job.stem} segment={name}" | |
| ) | |
| # Avoid path separators / collisions across stems. | |
| key = f"{job.stem}__{name}" | |
| if key in key_map: | |
| raise RuntimeError(f"Duplicate MFA batch key: {key}") | |
| key_map[key] = (job.stem, name) | |
| merged_manifest[key] = {"audio": str(audio_path.resolve()), "text": text} | |
| mfa_bin = resolve_mfa_bin(env_mfa) | |
| with tempfile.TemporaryDirectory(prefix="mfa_batch_merged_") as tmp: | |
| scratch = Path(tmp) | |
| audio_dir = scratch / "audio" | |
| audio_dir.mkdir() | |
| input_json = scratch / "merged_transcripts.json" | |
| output_json = scratch / "merged_aligned.json" | |
| input_json.write_text( | |
| json.dumps(merged_manifest, indent=2, ensure_ascii=False) + "\n", | |
| encoding="utf-8", | |
| ) | |
| env = _mfa_align_env(env_mfa, threads_per_job=threads) | |
| cmd = _mfa_align_cmd( | |
| audio_dir=audio_dir, | |
| input_json=input_json, | |
| output_json=output_json, | |
| threads=threads, | |
| ) | |
| log.info( | |
| "Running merged MFA for %d workspace(s) (%s): %s", | |
| len(jobs), mfa_bin, " ".join(cmd), | |
| ) | |
| subprocess.run( | |
| cmd, | |
| cwd=str(MFA_ALIGN_SCRIPT.parent), | |
| check=True, | |
| env=env, | |
| ) | |
| aligned_all = json.loads(output_json.read_text(encoding="utf-8")) | |
| pre_path = output_json.with_name(output_json.stem + "_pre_g2p.json") | |
| pre_all = ( | |
| json.loads(pre_path.read_text(encoding="utf-8")) | |
| if pre_path.is_file() | |
| else None | |
| ) | |
| by_stem: Dict[str, dict] = {job.stem: {} for job in jobs} | |
| pre_by_stem: Dict[str, dict] = {job.stem: {} for job in jobs} | |
| for key, payload in aligned_all.items(): | |
| mapped = key_map.get(key) | |
| if mapped is None: | |
| continue | |
| stem, name = mapped | |
| by_stem[stem][name] = payload | |
| if pre_all is not None: | |
| for key, payload in pre_all.items(): | |
| mapped = key_map.get(key) | |
| if mapped is None: | |
| continue | |
| stem, name = mapped | |
| pre_by_stem[stem][name] = payload | |
| for job in jobs: | |
| pre = pre_by_stem.get(job.stem) or None | |
| if pre is not None and not pre: | |
| pre = None | |
| _finalize_mfa_workspace( | |
| job.ws, job.segment_names, by_stem[job.stem], pre_g2p=pre, | |
| ) | |
| def run_align_mfa_jobs_parallel( | |
| jobs: Sequence[MfaAlignJob], | |
| *, | |
| env_mfa: str, | |
| max_parallel: int, | |
| threads_per_job: object = 1, | |
| poll_interval: float = 0.05, | |
| ) -> None: | |
| """Strategy 2: slot pool of MFA subprocesses (Applio libri launcher pattern). | |
| Uses ``Popen`` + poll (not ProcessPool) so this stays safe after the parent | |
| has already initialized CUDA (e.g. warm Whisper in ``run_demo_batched``). | |
| ``threads_per_job`` is MFA ``--num_jobs`` (-j) for each process. | |
| Use ``"segments"`` for per-stem ``-j = len(segments)``. Concurrent | |
| stems (``max_parallel > 1``) no longer override -j. | |
| """ | |
| if not jobs: | |
| return | |
| if not MFA_ALIGN_SCRIPT.is_file(): | |
| raise SystemExit(f"MFA align script not found: {MFA_ALIGN_SCRIPT}") | |
| workers = max(1, min(int(max_parallel), len(jobs))) | |
| per_segments = threads_per_job == ALIGN_THREADS_SEGMENTS | |
| if not per_segments: | |
| threads_fixed = max(1, int(threads_per_job)) # type: ignore[arg-type] | |
| if workers > 1 and threads_fixed > 1: | |
| log.warning( | |
| "max_parallel_align_jobs=%d with align_threads=%d " | |
| "(-j kept as requested; watch CPU oversubscription)", | |
| workers, threads_fixed, | |
| ) | |
| else: | |
| threads_fixed = None | |
| def _threads_for(job: MfaAlignJob) -> int: | |
| if per_segments: | |
| return resolve_mfa_align_threads(ALIGN_THREADS_SEGMENTS, job.segment_names) | |
| assert threads_fixed is not None | |
| return threads_fixed | |
| if workers <= 1: | |
| for job in jobs: | |
| j = _threads_for(job) | |
| log.info( | |
| "MFA serial stem=%s segments=%d -j %d", | |
| job.stem, len(job.segment_names), j, | |
| ) | |
| run_align_mfa( | |
| job.ws, job.segment_names, env_mfa=env_mfa, threads=j, | |
| ) | |
| return | |
| queue = list(jobs) | |
| processes: List[Tuple[subprocess.Popen, MfaAlignJob]] = [] | |
| log.info( | |
| "Running parallel MFA pool: %d job(s), max_parallel=%d, threads_per_job=%s", | |
| len(queue), workers, | |
| "segments" if per_segments else str(threads_fixed), | |
| ) | |
| while queue or processes: | |
| while queue and len(processes) < workers: | |
| job = queue.pop(0) | |
| j = _threads_for(job) | |
| per_segment_path = job.ws["aligned"].with_name( | |
| "03_aligned_per_segment.json" | |
| ) | |
| cmd = _mfa_align_cmd( | |
| audio_dir=job.ws["segments"], | |
| input_json=job.ws["transcripts"], | |
| output_json=per_segment_path, | |
| threads=j, | |
| ) | |
| env = _mfa_align_env(env_mfa, threads_per_job=j) | |
| log.info( | |
| "MFA pool start stem=%s segments=%d -j %d", | |
| job.stem, len(job.segment_names), j, | |
| ) | |
| proc = subprocess.Popen( | |
| cmd, env=env, cwd=str(MFA_ALIGN_SCRIPT.parent), | |
| ) | |
| processes.append((proc, job)) | |
| alive: List[Tuple[subprocess.Popen, MfaAlignJob]] = [] | |
| for proc, job in processes: | |
| ret = proc.poll() | |
| if ret is None: | |
| alive.append((proc, job)) | |
| elif ret != 0: | |
| raise RuntimeError( | |
| f"MFA align failed for stem={job.stem} exit={ret}" | |
| ) | |
| else: | |
| per_segment_path = job.ws["aligned"].with_name( | |
| "03_aligned_per_segment.json" | |
| ) | |
| per_segment = json.loads( | |
| per_segment_path.read_text(encoding="utf-8") | |
| ) | |
| pre_src = per_segment_path.with_name( | |
| "03_aligned_per_segment_pre_g2p.json" | |
| ) | |
| pre = ( | |
| json.loads(pre_src.read_text(encoding="utf-8")) | |
| if pre_src.is_file() | |
| else None | |
| ) | |
| _finalize_mfa_workspace( | |
| job.ws, job.segment_names, per_segment, pre_g2p=pre, | |
| ) | |
| processes = alive | |
| if processes or queue: | |
| time.sleep(poll_interval) | |
| def run_align_mfa_many( | |
| jobs: Sequence[MfaAlignJob], | |
| *, | |
| env_mfa: str, | |
| threads: object, | |
| max_parallel_align_jobs: int = 0, | |
| ) -> None: | |
| """Batch MFA for multiple workspaces. | |
| * ``max_parallel_align_jobs <= 0`` (default): Strategy 1 — merge corpora, | |
| one MFA invocation with ``--num_jobs=threads`` (best cold-start amortization). | |
| If ``threads == "segments"``, skip merge and run each stem serially with | |
| ``-j = len(segments)`` for that stem. | |
| * ``max_parallel_align_jobs >= 1``: Strategy 2 — up to N concurrent MFA | |
| processes. Each process keeps ``--align-threads`` as MFA ``-j`` | |
| (``segments`` → ``-j = len(that stem's segments)``). | |
| """ | |
| if not jobs: | |
| return | |
| per_segments = threads == ALIGN_THREADS_SEGMENTS | |
| if len(jobs) == 1: | |
| j = resolve_mfa_align_threads(threads, jobs[0].segment_names) | |
| run_align_mfa( | |
| jobs[0].ws, jobs[0].segment_names, env_mfa=env_mfa, threads=j, | |
| ) | |
| return | |
| if per_segments: | |
| # Never merge when -j tracks per-file segment count. | |
| max_p = int(max_parallel_align_jobs) | |
| if max_p <= 0: | |
| max_p = 1 | |
| run_align_mfa_jobs_parallel( | |
| jobs, | |
| env_mfa=env_mfa, | |
| max_parallel=max_p, | |
| threads_per_job=ALIGN_THREADS_SEGMENTS, | |
| ) | |
| return | |
| if int(max_parallel_align_jobs) >= 1: | |
| # N>=1 → process pool (N=1 is serial); always honor --align-threads. | |
| run_align_mfa_jobs_parallel( | |
| jobs, | |
| env_mfa=env_mfa, | |
| max_parallel=int(max_parallel_align_jobs), | |
| threads_per_job=max(1, int(threads)), # type: ignore[arg-type] | |
| ) | |
| else: | |
| run_align_mfa_batch_merged( | |
| jobs, env_mfa=env_mfa, threads=max(1, int(threads)), # type: ignore[arg-type] | |
| ) | |
| def run_align_maps( | |
| ws: dict[str, Path], | |
| segment_names: Sequence[str], | |
| *, | |
| model: Optional[Path], | |
| dictionary: Optional[Path], | |
| device: str, | |
| ensemble: bool, | |
| ) -> None: | |
| """MAPS align in the active interpreter (no conda hop).""" | |
| if not MAPS_ALIGN_SCRIPT.is_file(): | |
| raise SystemExit(f"MAPS align script not found: {MAPS_ALIGN_SCRIPT}") | |
| per_segment_path = ws["aligned"].with_name("03_aligned_per_segment.json") | |
| cmd = [ | |
| "python", | |
| str(MAPS_ALIGN_SCRIPT), | |
| "--audio_dir", | |
| str(ws["segments"]), | |
| "--input_json", | |
| str(ws["transcripts"]), | |
| "--output_json", | |
| str(per_segment_path), | |
| "--device", | |
| device, | |
| ] | |
| if model is not None: | |
| cmd.extend(["--model", str(model)]) | |
| if dictionary is not None: | |
| cmd.extend(["--dict", str(dictionary)]) | |
| if ensemble: | |
| cmd.append("--ensemble") | |
| run_python(cmd, cwd=MAPS_ALIGN_SCRIPT.parent) | |
| per_segment = json.loads(per_segment_path.read_text(encoding="utf-8")) | |
| missing = [name for name in segment_names if name not in per_segment] | |
| if missing: | |
| raise RuntimeError(f"MAPS alignment omitted segments: {missing}") | |
| write_aligned_json( | |
| per_segment, | |
| segment_names, | |
| ws["segments"], | |
| ws["aligned"], | |
| ) | |
| def run_align( | |
| ws: dict[str, Path], | |
| segment_names: Sequence[str], | |
| args: argparse.Namespace, | |
| ) -> None: | |
| if args.speech_backend == "maps": | |
| run_align_maps( | |
| ws, | |
| segment_names, | |
| model=args.maps_model, | |
| dictionary=args.maps_dict, | |
| device=args.maps_device, | |
| ensemble=args.maps_ensemble, | |
| ) | |
| else: | |
| run_align_mfa( | |
| ws, | |
| segment_names, | |
| env_mfa=args.env_mfa, | |
| threads=resolve_mfa_align_threads(args.align_threads, segment_names), | |
| ) | |