scripts/voice_convert/cli.py

"""Voice-conversion stage entry point.

Reads English audio clips from --input, looks up each clip's target voice via
--voice-assignments, dispatches to the chosen --backend, writes converted WAVs
to --output, and emits a per-clip report JSON.

Usage:
    python scripts/voice_convert/cli.py \\
        --input data/voice_normalization/inputs \\
        --voice-bank data/voice_bank \\
        --voice-assignments data/voice_normalization/voice_assignments.json \\
        --backend seed_vc \\
        --output data/voice_normalization/outputs/seed_vc

To add a backend: subclass VoiceConverterBackend in backends/<name>.py and
register it in BACKEND_REGISTRY below.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from dataclasses import asdict
from pathlib import Path
from typing import Any

# Make the parent dir importable when run as a script.
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from voice_convert.base import ConversionResult, VoiceConverterBackend


# ── Backend registry ────────────────────────────────────────────────────────
# Lazy imports — each entry returns a backend INSTANCE on demand. Lazy so
# importing seed_vc (which pulls heavy deps) doesn't break the CLI when running
# with --backend rvc.

def _load_seed_vc() -> VoiceConverterBackend:
    from voice_convert.backends.seed_vc import SeedVcV2Backend
    return SeedVcV2Backend()


def _load_rvc() -> VoiceConverterBackend:
    from voice_convert.backends.rvc import RvcBackend
    return RvcBackend()


def _load_elevenlabs() -> VoiceConverterBackend:
    from voice_convert.backends.elevenlabs import ElevenLabsBackend
    return ElevenLabsBackend()


BACKEND_REGISTRY = {
    "seed_vc": _load_seed_vc,
    "rvc": _load_rvc,
    "elevenlabs": _load_elevenlabs,
}


# ── Helpers ────────────────────────────────────────────────────────────────

def load_voice_bank(voice_bank_dir: Path) -> dict[str, Path]:
    """Read voice_bank.json, return {voice_id: absolute path to wav}."""
    bank_json = voice_bank_dir / "voice_bank.json"
    if not bank_json.exists():
        raise FileNotFoundError(f"Voice bank metadata missing: {bank_json}")
    data = json.loads(bank_json.read_text(encoding="utf-8"))
    out = {}
    for v in data.get("voices", []):
        wav_path = voice_bank_dir / v["filename"]
        if not wav_path.exists():
            raise FileNotFoundError(f"Voice bank refers to missing wav: {wav_path}")
        out[v["id"]] = wav_path
    if not out:
        raise ValueError(f"Voice bank is empty: {bank_json}")
    return out


def load_voice_assignments(path: Path) -> dict[str, dict[str, Any]]:
    """Read voice_assignments.json. Returns {input_filename: {speaker_id, target_voice_id, ...}}."""
    if not path.exists():
        raise FileNotFoundError(f"Voice assignments missing: {path}")
    return json.loads(path.read_text(encoding="utf-8"))


def collect_input_wavs(input_dir: Path) -> list[Path]:
    return sorted(p for p in input_dir.glob("*.wav") if not p.name.startswith("_"))


# ── Main loop ──────────────────────────────────────────────────────────────

def run(args: argparse.Namespace) -> int:
    input_dir = Path(args.input)
    voice_bank_dir = Path(args.voice_bank)
    assignments_path = Path(args.voice_assignments)
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Load lookup tables
    voice_lookup = load_voice_bank(voice_bank_dir)
    assignments = load_voice_assignments(assignments_path)

    inputs = collect_input_wavs(input_dir)
    if not inputs:
        print(f"No WAV files found in {input_dir}", file=sys.stderr)
        return 1

    # Dispatch backend
    if args.backend not in BACKEND_REGISTRY:
        print(f"Unknown backend '{args.backend}'. Available: {list(BACKEND_REGISTRY)}",
              file=sys.stderr)
        return 1

    print(f"Loading backend: {args.backend}")
    backend = BACKEND_REGISTRY[args.backend]()
    backend.setup({
        "preserve_accent": args.preserve_accent,
        "preserve_emotion": args.preserve_emotion,
    })

    # Run conversions
    results: list[dict[str, Any]] = []
    t_total = time.time()

    for src in inputs:
        assignment = assignments.get(src.name)
        if not assignment:
            print(f"  SKIP {src.name} — no entry in voice_assignments.json")
            results.append({
                "input": src.name, "skipped": True,
                "reason": "no assignment",
            })
            continue

        target_voice_id = assignment["target_voice_id"]
        target_voice_path = voice_lookup.get(target_voice_id)
        if not target_voice_path:
            print(f"  SKIP {src.name} — voice '{target_voice_id}' not in bank")
            results.append({
                "input": src.name, "skipped": True,
                "reason": f"voice {target_voice_id} not in bank",
            })
            continue

        # Output filename: <input_stem>_voiced.wav  (e.g. spk_0_moss_tts → spk_0_moss_tts_voiced)
        out_path = output_dir / f"{src.stem}_voiced.wav"

        print(f"  {src.name} → {target_voice_id} → {out_path.name}")
        try:
            result = backend.convert(
                source_wav=src,
                target_voice_wav=target_voice_path,
                output_wav=out_path,
                preserve_accent=args.preserve_accent,
                preserve_emotion=args.preserve_emotion,
            )
            results.append({
                "input": src.name,
                "output": result.output_wav.name,
                "speaker_id": assignment.get("speaker_id"),
                "target_voice_id": target_voice_id,
                "method": result.method,
                "elapsed_seconds": result.elapsed_seconds,
                "metadata": result.metadata,
                "success": result.success,
                "error": result.error,
            })
        except Exception as e:
            print(f"    ERROR: {type(e).__name__}: {e}")
            results.append({
                "input": src.name,
                "speaker_id": assignment.get("speaker_id"),
                "target_voice_id": target_voice_id,
                "method": args.backend,
                "success": False,
                "error": f"{type(e).__name__}: {e}",
            })

    backend.cleanup()
    total_elapsed = round(time.time() - t_total, 2)

    # Write report
    report = {
        "backend": args.backend,
        "input_dir": str(input_dir),
        "output_dir": str(output_dir),
        "voice_bank": str(voice_bank_dir),
        "preserve_accent": args.preserve_accent,
        "preserve_emotion": args.preserve_emotion,
        "n_inputs": len(inputs),
        "n_succeeded": sum(1 for r in results if r.get("success")),
        "n_failed": sum(1 for r in results if r.get("success") is False),
        "n_skipped": sum(1 for r in results if r.get("skipped")),
        "total_elapsed_seconds": total_elapsed,
        "results": results,
    }
    report_path = output_dir / "voice_conversion_report.json"
    report_path.write_text(json.dumps(report, indent=2), encoding="utf-8")

    print(f"\nDone in {total_elapsed}s. "
          f"OK={report['n_succeeded']} Failed={report['n_failed']} Skipped={report['n_skipped']}")
    print(f"Report: {report_path}")
    return 0 if report["n_failed"] == 0 else 1


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Voice-conversion stage entry point.")
    p.add_argument("--input", required=True,
                   help="Folder containing input English WAVs to convert")
    p.add_argument("--voice-bank", required=True,
                   help="Folder containing voice_bank.json + reference WAVs")
    p.add_argument("--voice-assignments", required=True,
                   help="JSON mapping input filenames to target voice ids")
    p.add_argument("--backend", required=True, choices=list(BACKEND_REGISTRY),
                   help="Voice conversion backend to use")
    p.add_argument("--output", required=True,
                   help="Folder to write converted WAVs + report JSON")
    p.add_argument("--preserve-accent", type=lambda s: s.lower() in ("true", "1", "yes"),
                   default=True, help="Pass to backend if it supports accent transfer")
    p.add_argument("--preserve-emotion", type=lambda s: s.lower() in ("true", "1", "yes"),
                   default=True, help="Pass to backend if it supports emotion transfer")
    return p.parse_args()


if __name__ == "__main__":
    sys.exit(run(parse_args()))