pipeline_scripts/inference.py

"""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)