backend.py

"""Lazy OpenVoice V2 synthesis backend."""

from __future__ import annotations

import os
import tempfile
import threading
from pathlib import Path
from typing import Any, Protocol


class VoiceBackend(Protocol):
    name: str

    @property
    def loaded(self) -> bool: ...

    @property
    def supported_languages(self) -> tuple[str, ...]: ...

    def synthesize(
        self,
        text: str,
        language: str,
        reference_audio: Path,
        output_path: Path,
        speed: float,
    ) -> None: ...


class OpenVoiceBackend:
    name = "openvoice-v2"

    def __init__(self) -> None:
        root = Path(os.getenv("OPENVOICE_CHECKPOINT_ROOT", "checkpoints_v2"))
        self.converter_dir = root / "converter"
        self.source_paths = {
            "JP": Path(
                os.getenv(
                    "OPENVOICE_JP_SOURCE_SE",
                    str(root / "base_speakers" / "ses" / "jp.pth"),
                )
            ),
            "EN": Path(
                os.getenv(
                    "OPENVOICE_EN_SOURCE_SE",
                    str(root / "base_speakers" / "ses" / "en-us.pth"),
                )
            ),
            "ZH": Path(
                os.getenv(
                    "OPENVOICE_ZH_SOURCE_SE",
                    str(root / "base_speakers" / "ses" / "zh.pth"),
                )
            ),
        }
        self.device = os.getenv("OPENVOICE_DEVICE", "")
        self.watermark = os.getenv("OPENVOICE_WATERMARK", "@MyShell")
        self._converter: Any | None = None
        self._torch: Any | None = None
        self._tts_models: dict[str, Any] = {}
        self._source_embeddings: dict[str, Any] = {}
        self._load_lock = threading.Lock()
        self._inference_lock = threading.Lock()

    @property
    def loaded(self) -> bool:
        return self._converter is not None

    @property
    def supported_languages(self) -> tuple[str, ...]:
        return tuple(self.source_paths)

    def _load_converter(self) -> tuple[Any, Any]:
        if self._converter is None:
            with self._load_lock:
                if self._converter is None:
                    config = self.converter_dir / "config.json"
                    checkpoint = self.converter_dir / "checkpoint.pth"
                    if not config.is_file() or not checkpoint.is_file():
                        raise RuntimeError(
                            "OpenVoice converter checkpoints are missing under "
                            f"{self.converter_dir}"
                        )
                    import torch
                    from openvoice.api import ToneColorConverter

                    device = self.device or ("cuda:0" if torch.cuda.is_available() else "cpu")
                    converter = ToneColorConverter(str(config), device=device)
                    converter.load_ckpt(str(checkpoint))
                    self.device = device
                    self._torch = torch
                    self._converter = converter
        return self._converter, self._torch

    def _get_language_assets(self, language: str) -> tuple[Any, Any]:
        converter, torch = self._load_converter()
        del converter
        if language not in self.source_paths:
            raise ValueError(f"unsupported language: {language}")
        if language not in self._tts_models:
            with self._load_lock:
                if language not in self._tts_models:
                    source_path = self.source_paths[language]
                    if not source_path.is_file():
                        raise RuntimeError(
                            f"base-speaker embedding is missing for {language}: {source_path}"
                        )
                    from melo.api import TTS

                    self._tts_models[language] = TTS(language=language, device=self.device)
                    self._source_embeddings[language] = torch.load(
                        str(source_path), map_location=self.device, weights_only=True
                    )
        return self._tts_models[language], self._source_embeddings[language]

    def synthesize(
        self,
        text: str,
        language: str,
        reference_audio: Path,
        output_path: Path,
        speed: float,
    ) -> None:
        if not reference_audio.is_file():
            raise ValueError("reference audio does not exist")
        converter, _ = self._load_converter()
        tts_model, source_embedding = self._get_language_assets(language)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        from openvoice import se_extractor

        with self._inference_lock:
            target_embedding, _ = se_extractor.get_se(str(reference_audio), converter, vad=False)
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temporary:
                base_audio = Path(temporary.name)
            try:
                speaker_ids = getattr(tts_model.hps.data, "spk2id", {})
                speaker_id = next(iter(speaker_ids.values()), 0)
                tts_model.tts_to_file(text, speaker_id, str(base_audio), speed=speed)
                converter.convert(
                    audio_src_path=str(base_audio),
                    src_se=source_embedding,
                    tgt_se=target_embedding,
                    output_path=str(output_path),
                    message=self.watermark,
                )
            finally:
                base_audio.unlink(missing_ok=True)