openvoice_transfer.py

"""
OpenVoice v2 — tone-colour voice transfer.

Wraps MyShell's ToneColorConverter to stamp a parent's voice timbre
onto MMS-synthesised audio.

Usage pattern:
  1. At profile upload time:
       se = extract_se("parent_voice.wav")
       torch.save(se, "profile.pt")

  2. At audio generation time:
       output_wav = transfer_voice(mms_wav_bytes, target_se_tensor)
"""

from __future__ import annotations

import os
import sys
import tempfile
import types
from typing import Optional

import torch

# Lazy-loaded singletons
_converter = None


def load_converter(ckpt_dir: str, device: str = "cpu") -> None:
    """
    Load the OpenVoice v2 ToneColorConverter from *ckpt_dir*.
    Expected files:
      - <ckpt_dir>/config.json
      - <ckpt_dir>/checkpoint.pth
    """
    global _converter

    from openvoice.api import ToneColorConverter

    config_path = os.path.join(ckpt_dir, "config.json")
    ckpt_path = os.path.join(ckpt_dir, "checkpoint.pth")

    print(f"[OpenVoice] Loading converter from {ckpt_dir} …")
    # OpenVoice imports wavmark inside the converter constructor even though
    # watermarking is not needed for voice transfer. A no-op loader avoids an
    # extra model and works around its broken enable_watermark keyword path.
    wavmark_stub = types.ModuleType("wavmark")

    class DisabledWatermark:
        def to(self, _device):
            return None

    wavmark_stub.load_model = DisabledWatermark
    sys.modules["wavmark"] = wavmark_stub

    converter = ToneColorConverter(config_path, device=device)
    converter.load_ckpt(ckpt_path)

    _converter = converter
    print("[OpenVoice] Converter ready.")


def _ensure_loaded() -> None:
    if _converter is None:
        raise RuntimeError(
            "OpenVoice converter is not loaded. Call load_converter() at startup."
        )


def extract_se(audio_path: str) -> torch.Tensor:
    """
    Extract the speaker embedding (tone colour) from *audio_path*.
    Returns a torch.Tensor that can be saved with torch.save().
    """
    _ensure_loaded()
    return _converter.extract_se(audio_path)


def transfer_voice(
    source_wav_bytes: bytes,
    target_se: torch.Tensor,
    tau: float = 0.3,
) -> bytes:
    """
    Stamp *target_se* (parent's voice) onto *source_wav_bytes* (MMS output).

    *tau* controls transfer strength (0.0 = source unchanged, 1.0 = maximum
    speaker similarity). Default 0.3 balances naturalness vs. speaker match.

    Returns WAV bytes with the parent's voice characteristics applied.
    """
    _ensure_loaded()

    # Write MMS output to a temp file
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as src_f:
        src_f.write(source_wav_bytes)
        src_path = src_f.name

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

    try:
        # Extract speaker embedding from MMS source audio
        src_se = _converter.extract_se(src_path)

        _converter.convert(
            audio_src_path=src_path,
            src_se=src_se,
            tgt_se=target_se,
            output_path=out_path,
            tau=tau,
        )

        with open(out_path, "rb") as f:
            return f.read()
    finally:
        for p in (src_path, out_path):
            try:
                os.unlink(p)
            except OSError:
                pass