src/tts/voice_cloner.py

"""
VoiceCloner — OpenVoice V2 tone-color converter.

Sits downstream of WaxalTTSEngine: takes the base VITS audio and reshapes it
to match a target speaker's tone color.

Usage:
    cloner = VoiceCloner()
    cloner.preload()                          # background thread

    # After WaxalTTS produces (audio_np, sr) …
    se = cloner.extract_se(audio_np, sr)      # extract SE from user's mic audio
    result = cloner.convert(audio_np, sr, se) # returns (cloned_audio, sr) or None

The OpenVoice V2 checkpoint is downloaded from myshell-ai/openvoice-v2 on
HuggingFace Hub at first use (cached in data/openvoice_v2/).

Falls back gracefully (returns None) if openvoice is not installed or the
checkpoint download fails — in that case the caller uses the raw VITS output.
"""
from __future__ import annotations

import logging
import os
import tempfile
import threading
from pathlib import Path
from typing import Optional

import numpy as np

logger = logging.getLogger(__name__)

OV_HF_REPO  = "myshell-ai/openvoice-v2"
OV_CKPT_DIR = Path("data/openvoice_v2")
HF_TOKEN    = os.environ.get("HF_TOKEN")


class VoiceCloner:
    """
    Thin wrapper around OpenVoice V2 ToneColorConverter.

    Thread-safety: convert() holds _lock so parallel calls are serialised;
    the lock is released while waiting for subprocess/IO.
    """

    def __init__(self) -> None:
        self._lock      = threading.Lock()
        self._converter = None
        self._src_se    = None   # cached base-TTS source SE (computed on first convert)
        self._ready     = False
        self._error: Optional[str] = None

    def preload(self) -> None:
        threading.Thread(target=self._load, daemon=True).start()

    def get_status(self) -> str:
        if self._ready:  return "ready"
        if self._error:  return f"error: {self._error}"
        return "loading…"

    # ── Loading ───────────────────────────────────────────────────────────────

    def _load(self) -> None:
        try:
            from openvoice.api import ToneColorConverter   # noqa: F401 — validate import
            OV_CKPT_DIR.mkdir(parents=True, exist_ok=True)

            # Download checkpoint from HF Hub once, then use local cache
            converter_cfg = OV_CKPT_DIR / "converter" / "config.json"
            if not converter_cfg.exists():
                logger.info("VoiceCloner: downloading OpenVoice V2 from HF Hub …")
                from huggingface_hub import snapshot_download
                snapshot_download(
                    repo_id=OV_HF_REPO,
                    local_dir=str(OV_CKPT_DIR),
                    token=HF_TOKEN,
                )

            # Find config — repo layout may vary
            cfg_path = self._find_converter_config()
            if cfg_path is None:
                raise FileNotFoundError(
                    f"converter/config.json not found under {OV_CKPT_DIR}"
                )

            from openvoice.api import ToneColorConverter
            logger.info("VoiceCloner: loading ToneColorConverter from %s …", cfg_path)
            converter = ToneColorConverter(str(cfg_path), device="cpu")
            ckpt = cfg_path.parent / "checkpoint.pth"
            converter.load_ckpt(str(ckpt))

            with self._lock:
                self._converter = converter
                self._ready     = True
            logger.info("VoiceCloner: OpenVoice V2 ready")

        except Exception as exc:
            self._error = str(exc)
            logger.warning(
                "VoiceCloner: load failed — voice cloning disabled: %s", exc
            )

    def _find_converter_config(self) -> Optional[Path]:
        """Probe known checkpoint layouts to locate converter/config.json."""
        candidates = [
            OV_CKPT_DIR / "converter" / "config.json",
            OV_CKPT_DIR / "checkpoints_v2" / "converter" / "config.json",
        ]
        for p in candidates:
            if p.exists():
                return p
        # Walk one level deep as fallback
        for p in OV_CKPT_DIR.rglob("config.json"):
            if p.parent.name == "converter":
                return p
        return None

    # ── SE extraction ─────────────────────────────────────────────────────────

    def extract_se(self, audio_np: np.ndarray, sr: int) -> Optional[np.ndarray]:
        """
        Extract OpenVoice V2 tone-color SE from raw float32 audio.
        Returns a numpy array (shape depends on OV model, typically (1, 256)),
        or None if not ready.
        """
        if not self._ready:
            return None
        try:
            import soundfile as sf
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
                tmp = f.name
            sf.write(tmp, audio_np, sr)
            se = self._extract_se_from_file(tmp)
            Path(tmp).unlink(missing_ok=True)
            return se
        except Exception as exc:
            logger.debug("VoiceCloner.extract_se: %s", exc)
            return None

    def _extract_se_from_file(self, audio_path: str) -> Optional[np.ndarray]:
        try:
            from openvoice import se_extractor
            se, _ = se_extractor.get_se(
                audio_path,
                self._converter,
                target_dir=str(OV_CKPT_DIR / "tmp"),
                vad=False,
            )
            arr = se.cpu().numpy() if hasattr(se, "cpu") else np.array(se)
            return arr
        except Exception as exc:
            logger.debug("VoiceCloner._extract_se_from_file: %s", exc)
            return None

    # ── Voice conversion ──────────────────────────────────────────────────────

    def convert(
        self,
        audio_np: np.ndarray,
        sr: int,
        target_se: np.ndarray,
    ) -> Optional[tuple[np.ndarray, int]]:
        """
        Reshape audio to match the target speaker's tone color.

        Args:
            audio_np:  float32 audio from WaxalTTS (base voice).
            sr:        sample rate of audio_np.
            target_se: OpenVoice SE from SpeakerProfileManager (Individual or
                       Collective).

        Returns (cloned_audio_float32, sample_rate) or None if not ready.
        """
        if not self._ready:
            return None
        try:
            import soundfile as sf
            import torch

            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
                src_path = f.name
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
                out_path = f.name

            sf.write(src_path, audio_np, sr)

            with self._lock:
                # Extract source SE on first call, then cache it for the session
                if self._src_se is None:
                    se = self._extract_se_from_file(src_path)
                    if se is not None:
                        self._src_se = se

                if self._src_se is None:
                    logger.warning("VoiceCloner: could not extract source SE")
                    return None

                src_se_t = torch.tensor(self._src_se)
                tgt_se_t = torch.tensor(target_se)

                # Ensure batch dim matches what the converter expects
                if src_se_t.dim() == 1:
                    src_se_t = src_se_t.unsqueeze(0)
                if tgt_se_t.dim() == 1:
                    tgt_se_t = tgt_se_t.unsqueeze(0)

                self._converter.convert(
                    audio_src_path=src_path,
                    src_se=src_se_t,
                    tgt_se=tgt_se_t,
                    output_path=out_path,
                    message="@MyShell",
                )

            audio_out, out_sr = sf.read(out_path, dtype="float32")
            Path(src_path).unlink(missing_ok=True)
            Path(out_path).unlink(missing_ok=True)
            return audio_out.astype(np.float32), out_sr

        except Exception as exc:
            logger.error("VoiceCloner.convert: %s", exc)
            return None