Spaces:
Sleeping
Sleeping
| """Source-level span collection, inference invocation, and stitching.""" | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import logging | |
| import time | |
| from pathlib import Path | |
| from typing import Dict, List, Sequence, Tuple | |
| import numpy as np | |
| import soundfile as sf | |
| import soxr | |
| from pipeline_utils import per_segment_alignments | |
| from .constants import INFER_SCRIPT, OUTPUT_FS, REPO_ROOT, _SEGMENT_RE | |
| from .runtime import run_python | |
| from .spans import build_remapped_phone_intervals, build_spans_for_phone_map, collect_labeled_spans | |
| from .timing import _fmt_elapsed | |
| log = logging.getLogger(__name__) | |
| def _cmd_flag(name: str, value) -> List[str]: | |
| if value is None: | |
| return [] | |
| if isinstance(value, bool): | |
| return [name] if value else [] | |
| return [name, str(value)] | |
| def build_infer_command( | |
| args: argparse.Namespace, | |
| wav_path: Path, | |
| spans_path: Path, | |
| output_path: Path, | |
| ) -> List[str]: | |
| cmd: List[str] = [ | |
| "python", | |
| str(INFER_SCRIPT), | |
| "--wav", | |
| str(wav_path), | |
| "--spans", | |
| str(spans_path), | |
| "--checkpoint", | |
| str(args.checkpoint), | |
| "--output", | |
| str(output_path), | |
| ] | |
| cmd.extend(_cmd_flag("--save-resampled", args.save_resampled)) | |
| cmd.extend(_cmd_flag("--device", args.infer_device)) | |
| if args.max_utt_sec != float("inf"): | |
| cmd.extend(["--max_utt_sec", str(args.max_utt_sec)]) | |
| if args.ola: | |
| cmd.extend([ | |
| "--ola", | |
| "--ola_chunk_frames", str(args.ola_chunk_frames), | |
| "--ola_hop_factor", str(args.ola_hop_factor), | |
| "--ola_ctx_frames", str(args.ola_ctx_frames), | |
| ]) | |
| if args.windowed_attention: | |
| cmd.append("--windowed_attention") | |
| cmd.extend(["--context_frames", str(args.context_frames)]) | |
| cmd.extend(_cmd_flag("--normalize_vol", args.normalize_vol)) | |
| cmd.extend(_cmd_flag("--encodec-ckpt", args.encodec_ckpt)) | |
| if getattr(args, "full_phone_conditioning", False): | |
| cmd.append("--full_phone_conditioning") | |
| if getattr(args, "full_decode", False): | |
| cmd.append("--full-decode") | |
| return cmd | |
| def build_batch_infer_command( | |
| args: argparse.Namespace, | |
| manifest_path: Path, | |
| ) -> List[str]: | |
| """Like build_infer_command, but for one --batch_manifest covering every | |
| source file at once (checkpoint/EnCodec loaded a single time).""" | |
| cmd: List[str] = [ | |
| "python", | |
| str(INFER_SCRIPT), | |
| "--batch_manifest", | |
| str(manifest_path), | |
| "--checkpoint", | |
| str(args.checkpoint), | |
| ] | |
| cmd.extend(_cmd_flag("--device", args.infer_device)) | |
| if args.max_utt_sec != float("inf"): | |
| cmd.extend(["--max_utt_sec", str(args.max_utt_sec)]) | |
| if args.ola: | |
| cmd.extend([ | |
| "--ola", | |
| "--ola_chunk_frames", str(args.ola_chunk_frames), | |
| "--ola_hop_factor", str(args.ola_hop_factor), | |
| "--ola_ctx_frames", str(args.ola_ctx_frames), | |
| ]) | |
| if args.windowed_attention: | |
| cmd.append("--windowed_attention") | |
| cmd.extend(["--context_frames", str(args.context_frames)]) | |
| cmd.extend(_cmd_flag("--normalize_vol", args.normalize_vol)) | |
| cmd.extend(_cmd_flag("--encodec-ckpt", args.encodec_ckpt)) | |
| if getattr(args, "full_phone_conditioning", False): | |
| cmd.append("--full_phone_conditioning") | |
| if getattr(args, "full_decode", False): | |
| cmd.append("--full-decode") | |
| return cmd | |
| def segment_source_stem(name: str) -> str: | |
| """Map a segment filename back to its source wav stem (foo_segment_3.wav → foo).""" | |
| m = _SEGMENT_RE.search(name) | |
| if m: | |
| return name[:m.start()].rstrip("_") | |
| return Path(name).stem | |
| def segment_index(name: str) -> int: | |
| m = _SEGMENT_RE.search(name) | |
| return int(m.group(1)) if m else 0 | |
| def segments_by_source(segment_names: Sequence[str]) -> Dict[str, List[str]]: | |
| """Group segment filenames by source stem, ordered by segment index.""" | |
| grouped: Dict[str, List[tuple[int, str]]] = {} | |
| for name in segment_names: | |
| stem = segment_source_stem(name) | |
| grouped.setdefault(stem, []).append((segment_index(name), name)) | |
| return { | |
| stem: [name for _, name in sorted(items, key=lambda x: x[0])] | |
| for stem, items in grouped.items() | |
| } | |
| def source_order_from_segments(segment_names: Sequence[str]) -> List[str]: | |
| """Source stems in first-seen order (matches original file / segment order).""" | |
| order: List[str] = [] | |
| seen: set[str] = set() | |
| for name in segment_names: | |
| stem = segment_source_stem(name) | |
| if stem not in seen: | |
| seen.add(stem) | |
| order.append(stem) | |
| return order | |
| def segment_start_offsets(segment_names: Sequence[str], segments_dir: Path) -> Dict[str, float]: | |
| """Cumulative start time of each segment in the original source timeline.""" | |
| offsets: Dict[str, float] = {} | |
| t = 0.0 | |
| for name in segment_names: | |
| offsets[name] = t | |
| t += sf.info(str(segments_dir / name)).duration | |
| return offsets | |
| def resolve_source_wav(source_stem: str, input_dir: Path, ws: dict[str, Path]) -> Path: | |
| for root in (input_dir, ws["input"]): | |
| path = root / f"{source_stem}.wav" | |
| if path.is_file(): | |
| return path | |
| raise FileNotFoundError( | |
| f"Source wav not found for {source_stem!r} under {input_dir} or {ws['input']}" | |
| ) | |
| def _shift_spans(spans: list, offset_sec: float) -> list: | |
| """Add offset_sec to every span start/end field (seconds).""" | |
| shifted: list = [] | |
| for sp in spans: | |
| if isinstance(sp, dict): | |
| out = dict(sp) | |
| for key in ("start", "end", "t_start", "t_end"): | |
| if key in out and out[key] is not None: | |
| out[key] = float(out[key]) + offset_sec | |
| shifted.append(out) | |
| elif isinstance(sp, (list, tuple)) and len(sp) >= 2: | |
| shifted.append([float(sp[0]) + offset_sec, float(sp[1]) + offset_sec, *sp[2:]]) | |
| else: | |
| shifted.append(sp) | |
| return shifted | |
| def passthrough_wav(src: Path, dst: Path, *, out_fs: int = OUTPUT_FS) -> None: | |
| wav, sr = sf.read(str(src), dtype="float32", always_2d=False) | |
| if wav.ndim > 1: | |
| wav = wav.mean(axis=1) | |
| if sr != out_fs: | |
| wav = soxr.resample(wav, sr, out_fs, quality="VHQ").astype(np.float32) | |
| dst.parent.mkdir(parents=True, exist_ok=True) | |
| sf.write(str(dst), wav, out_fs) | |
| def collect_inpaint_jobs( | |
| ws: dict[str, Path], | |
| *, | |
| infer_args: argparse.Namespace, | |
| segment_names: Sequence[str], | |
| input_dir: Path, | |
| out_fs: int = OUTPUT_FS, | |
| ) -> Tuple[List[str], List[dict], List[int]]: | |
| """Build spans JSONs and return infer-script jobs without running inference. | |
| Returns ``(source_order, batch_jobs, span_counts)``. Missing segment | |
| alignments or zero matching spans abort with ``SystemExit`` (no silent | |
| passthrough). | |
| """ | |
| aligned = json.loads(ws["aligned"].read_text(encoding="utf-8")) | |
| per_segment = per_segment_alignments(aligned) | |
| # Prefer the dedicated per-segment file when present (survives writers that | |
| # only keep merged source keys in 03_aligned.json). | |
| per_seg_path = ws["aligned"].with_name("03_aligned_per_segment.json") | |
| if per_seg_path.is_file(): | |
| try: | |
| extra = json.loads(per_seg_path.read_text(encoding="utf-8")) | |
| except (json.JSONDecodeError, OSError) as exc: | |
| raise SystemExit( | |
| f"ERROR: could not read {per_seg_path}: {exc}" | |
| ) from exc | |
| if isinstance(extra, dict): | |
| for key, val in extra.items(): | |
| if _SEGMENT_RE.search(key) and isinstance(val, dict): | |
| per_segment[key] = val | |
| ws["spans"].mkdir(parents=True, exist_ok=True) | |
| ws["inpainted"].mkdir(parents=True, exist_ok=True) | |
| phone_map = ", ".join( | |
| f"{pi!r}->{po!r}" for pi, po in zip(infer_args.phones_in, infer_args.phones_out) | |
| ) | |
| grouped = segments_by_source(segment_names) | |
| source_order = source_order_from_segments(segment_names) | |
| batch_jobs: List[dict] = [] | |
| batch_spans_counts: List[int] = [] | |
| for source_stem in source_order: | |
| seg_names = grouped[source_stem] | |
| offsets = segment_start_offsets(seg_names, ws["segments"]) | |
| source_wav = resolve_source_wav(source_stem, input_dir, ws) | |
| out_path = ws["inpainted"] / f"{source_stem}.wav" | |
| source_dur = sf.info(str(source_wav)).duration | |
| all_spans: list = [] | |
| all_phone_intervals: list = [] | |
| for name in seg_names: | |
| chunk_path = ws["segments"] / name | |
| if not chunk_path.is_file(): | |
| raise FileNotFoundError(f"Missing segment wav: {chunk_path}") | |
| entry = per_segment.get(name) | |
| if entry is None: | |
| raise SystemExit( | |
| f"ERROR: No alignment for segment {name} (source {source_stem}) " | |
| f"in {ws['aligned']} (and {per_seg_path.name if per_seg_path.is_file() else 'no per-segment file'}). " | |
| f"Present keys: {sorted(aligned)}. Refusing silent passthrough — " | |
| f"re-run align or restore work/{source_stem}/03_aligned*.json." | |
| ) | |
| try: | |
| labeled = collect_labeled_spans( | |
| entry.get("phones", {}), wanted=infer_args.phones_in | |
| ) | |
| chunk_dur = sf.info(str(chunk_path)).duration | |
| spans = build_spans_for_phone_map( | |
| labeled, | |
| infer_args.phones_in, | |
| infer_args.phones_out, | |
| ext_sec=infer_args.ext, | |
| max_end=chunk_dur, | |
| force_ext=getattr(infer_args, "force_ext", True), | |
| ) | |
| except ValueError: | |
| continue | |
| all_spans.extend(_shift_spans(spans, offsets[name])) | |
| if getattr(infer_args, "full_phone_conditioning", False): | |
| intervals = build_remapped_phone_intervals( | |
| entry.get("phones", {}), | |
| infer_args.phones_in, | |
| infer_args.phones_out, | |
| ) | |
| all_phone_intervals.extend(_shift_spans(intervals, offsets[name])) | |
| if not all_spans: | |
| raise SystemExit( | |
| f"ERROR: No spans for source {source_stem} ({phone_map}); " | |
| f"refusing silent passthrough. Check --phones_in against the " | |
| f"alignment phones in {ws['aligned']} / {per_seg_path.name}." | |
| ) | |
| spans_payload: dict = {"spans": all_spans, "source": source_stem} | |
| if getattr(infer_args, "full_phone_conditioning", False): | |
| if not all_phone_intervals: | |
| raise SystemExit( | |
| f"--full_phone_conditioning set but no phone_intervals for {source_stem}" | |
| ) | |
| spans_payload["phone_intervals"] = all_phone_intervals | |
| spans_path = ws["spans"] / f"{source_stem}.json" | |
| spans_path.write_text( | |
| json.dumps(spans_payload, indent=2), | |
| encoding="utf-8", | |
| ) | |
| log.info( | |
| "Queued %d span(s) in %s (from %d segment(s), %.2f s, %s)", | |
| len(all_spans), source_stem, len(seg_names), source_dur, phone_map, | |
| ) | |
| job: dict = { | |
| "wav": str(source_wav), | |
| "spans": str(spans_path), | |
| "output": str(out_path), | |
| } | |
| if infer_args.save_resampled is not None: | |
| job["save_resampled"] = str(infer_args.save_resampled) | |
| batch_jobs.append(job) | |
| batch_spans_counts.append(len(all_spans)) | |
| return source_order, batch_jobs, batch_spans_counts | |
| def execute_infer_batch( | |
| *, | |
| infer_args: argparse.Namespace, | |
| batch_jobs: Sequence[dict], | |
| manifest_path: Path, | |
| ) -> None: | |
| """Write *batch_jobs* to *manifest_path* and run one infer subprocess.""" | |
| if not batch_jobs: | |
| return | |
| manifest_path.parent.mkdir(parents=True, exist_ok=True) | |
| manifest_path.write_text( | |
| json.dumps({"jobs": list(batch_jobs)}, indent=2), | |
| encoding="utf-8", | |
| ) | |
| log.info( | |
| "Inpainting %d source file(s) in one batch — checkpoint + EnCodec loaded once", | |
| len(batch_jobs), | |
| ) | |
| cmd = build_batch_infer_command(infer_args, manifest_path) | |
| t0 = time.perf_counter() | |
| run_python(cmd, cwd=REPO_ROOT) | |
| dt = time.perf_counter() - t0 | |
| log.info( | |
| "[timing] inpaint total — %d source file(s), %s (avg %s/file)", | |
| len(batch_jobs), | |
| _fmt_elapsed(dt), | |
| _fmt_elapsed(dt / len(batch_jobs)), | |
| ) | |
| def run_inpaint( | |
| ws: dict[str, Path], | |
| *, | |
| infer_args: argparse.Namespace, | |
| segment_names: Sequence[str], | |
| input_dir: Path, | |
| out_fs: int = OUTPUT_FS, | |
| ) -> List[str]: | |
| source_order, batch_jobs, batch_spans_counts = collect_inpaint_jobs( | |
| ws, | |
| infer_args=infer_args, | |
| segment_names=segment_names, | |
| input_dir=input_dir, | |
| out_fs=out_fs, | |
| ) | |
| if batch_jobs: | |
| phone_map = ", ".join( | |
| f"{pi!r}->{po!r}" | |
| for pi, po in zip(infer_args.phones_in, infer_args.phones_out) | |
| ) | |
| log.info( | |
| "Inpainting %d source file(s) (%d span(s) total, %s)", | |
| len(batch_jobs), sum(batch_spans_counts), phone_map, | |
| ) | |
| execute_infer_batch( | |
| infer_args=infer_args, | |
| batch_jobs=batch_jobs, | |
| manifest_path=ws["spans"] / "_batch_manifest.json", | |
| ) | |
| return source_order | |
| def stitch_wavs( | |
| paths: Sequence[Path], | |
| output_path: Path, | |
| *, | |
| crossfade_ms: float = 0.0, | |
| out_fs: int = OUTPUT_FS, | |
| ) -> None: | |
| if not paths: | |
| raise ValueError("No wav files to stitch") | |
| chunks: List[np.ndarray] = [] | |
| for p in paths: | |
| wav, sr = sf.read(str(p), dtype="float32", always_2d=False) | |
| if wav.ndim > 1: | |
| wav = wav.mean(axis=1) | |
| if sr != out_fs: | |
| wav = soxr.resample(wav, sr, out_fs, quality="VHQ").astype(np.float32) | |
| chunks.append(wav.astype(np.float32)) | |
| if crossfade_ms <= 0 or len(chunks) == 1: | |
| out = np.concatenate(chunks) | |
| else: | |
| fade = max(1, int(out_fs * crossfade_ms / 1000.0)) | |
| out = chunks[0] | |
| for nxt in chunks[1:]: | |
| if len(out) >= fade and len(nxt) >= fade: | |
| ramp = np.linspace(0.0, 1.0, fade, dtype=np.float32) | |
| blended = out[-fade:] * (1.0 - ramp) + nxt[:fade] * ramp | |
| out = np.concatenate([out[:-fade], blended, nxt[fade:]]) | |
| else: | |
| out = np.concatenate([out, nxt]) | |
| peak = float(np.max(np.abs(out))) if len(out) else 0.0 | |
| if peak > 1.0: | |
| out = out / peak | |
| output_path.parent.mkdir(parents=True, exist_ok=True) | |
| sf.write(str(output_path), out, out_fs) | |
| log.info("Wrote stitched output (%.2f s): %s", len(out) / out_fs, output_path) | |