app/services/melo_engine.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import tempfile
import threading
import wave

import numpy as np

from app.config import Settings
from app.errors import OpenAICompatibleError
from app.services.supertonic_engine import SynthesisResult

try:
    from melo.api import TTS as MeloTTS
except ImportError:  # pragma: no cover
    MeloTTS = None


@dataclass(frozen=True)
class MeloVoiceSpec:
    alias: str
    canonical_name: str
    provider_voice_key: str
    source: str = "builtin"
    style_path: str = ""


class MeloEngine:
    public_model_ids = ["melo-en"]
    primary_model_id = "melo-en"
    default_sample_rate = 24000
    voice_aliases = {
        "default": MeloVoiceSpec(
            alias="default",
            canonical_name="EN Default",
            provider_voice_key="EN-Default",
        ),
        "us": MeloVoiceSpec(
            alias="us", canonical_name="EN US", provider_voice_key="EN-US"
        ),
        "br": MeloVoiceSpec(
            alias="br", canonical_name="EN BR", provider_voice_key="EN-BR"
        ),
        "india": MeloVoiceSpec(
            alias="india", canonical_name="EN India", provider_voice_key="EN_INDIA"
        ),
        "australia": MeloVoiceSpec(
            alias="australia", canonical_name="EN Australia", provider_voice_key="EN-AU"
        ),
    }

    def __init__(
        self, *, tts: object | None = None, tts_factory=None, device: str = "cpu"
    ) -> None:
        self._tts = tts
        self._tts_factory = tts_factory
        self.device = device
        self._init_lock = threading.Lock()

    @classmethod
    def from_settings(cls, settings: Settings) -> "MeloEngine":
        if MeloTTS is None:
            return cls(tts=None, tts_factory=None, device=settings.melo_device)
        return cls(
            tts_factory=lambda: MeloTTS(
                language="EN", device=settings.melo_device, use_hf=True
            ),
            device=settings.melo_device,
        )

    def is_available(self) -> bool:
        return self._tts is not None or self._tts_factory is not None

    def warmup(self) -> None:
        self._ensure_tts()

    def _ensure_tts(self):
        if self._tts is not None:
            return self._tts
        if self._tts_factory is None:
            raise OpenAICompatibleError(
                status_code=500,
                message="MeloTTS engine is not available.",
                error_type="server_error",
                code="engine_unavailable",
            )
        with self._init_lock:
            if self._tts is None:
                try:
                    self._tts = self._tts_factory()
                except Exception as exc:  # pragma: no cover
                    raise OpenAICompatibleError(
                        status_code=500,
                        message=f"Failed to initialize MeloTTS engine: {exc}",
                        error_type="server_error",
                        code="engine_init_failed",
                    ) from exc
        return self._tts

    def list_voice_bindings(self) -> list[dict[str, str]]:
        return [
            {
                "alias": spec.alias,
                "canonical_name": spec.canonical_name,
                "provider_voice_id": spec.provider_voice_key,
                "source": spec.source,
                "style_path": spec.style_path,
                "model": self.primary_model_id,
            }
            for spec in self.voice_aliases.values()
        ]

    def supports_voice(self, voice: str) -> bool:
        return voice.strip().lower() in self.voice_aliases

    def _resolve_speaker_id(self, voice: str) -> int:
        tts = self._ensure_tts()
        spk2id = getattr(
            getattr(getattr(tts, "hps", None), "data", None), "spk2id", None
        )
        if not isinstance(spk2id, dict) or not spk2id:
            raise OpenAICompatibleError(
                status_code=500,
                message="MeloTTS speaker map is unavailable.",
                error_type="server_error",
                code="speaker_map_unavailable",
            )

        spec = self.voice_aliases[voice.strip().lower()]
        if spec.provider_voice_key in spk2id:
            return int(spk2id[spec.provider_voice_key])

        if voice.strip().lower() == "default":
            for key in ("EN-Default", "EN-US"):
                if key in spk2id:
                    return int(spk2id[key])

            for key, speaker_id in spk2id.items():
                if str(key).upper().startswith("EN"):
                    return int(speaker_id)

        raise OpenAICompatibleError(
            status_code=500,
            message=f"MeloTTS could not resolve speaker '{voice}'.",
            param="voice",
            code="unsupported_voice",
        )

    def _read_waveform(self, output_path: Path) -> SynthesisResult:
        with wave.open(str(output_path), "rb") as wav_file:
            frames = wav_file.readframes(wav_file.getnframes())
            sample_rate = wav_file.getframerate()
            channels = wav_file.getnchannels()
            sample_width = wav_file.getsampwidth()

        dtype_map = {1: np.uint8, 2: np.int16, 4: np.int32}
        if sample_width not in dtype_map:
            raise OpenAICompatibleError(
                status_code=500,
                message=f"Unsupported MeloTTS sample width: {sample_width}",
                error_type="server_error",
                code="invalid_waveform",
            )

        waveform = np.frombuffer(frames, dtype=dtype_map[sample_width]).astype(
            np.float32
        )
        if sample_width == 1:
            waveform = (waveform - 128.0) / 128.0
        else:
            scale = float(2 ** (8 * sample_width - 1))
            waveform = waveform / scale
        if channels > 1:
            waveform = waveform.reshape(-1, channels)[:, 0]
        return SynthesisResult(waveform=waveform, sample_rate=sample_rate)

    def synthesize(
        self,
        *,
        text: str,
        voice: str,
        speed: float,
        quality: str,
        model_name: str,
        lang: str,
    ) -> SynthesisResult:
        del quality, model_name, lang
        normalized_voice = voice.strip().lower()
        if normalized_voice not in self.voice_aliases:
            raise OpenAICompatibleError(
                status_code=400,
                message=f"Unsupported MeloTTS voice '{voice}'.",
                param="voice",
                code="unsupported_voice",
            )
        if speed <= 0:
            raise OpenAICompatibleError(
                status_code=400,
                message="MeloTTS speed must be greater than 0.",
                param="speed",
                code="invalid_speed",
            )

        tts = self._ensure_tts()
        speaker_id = self._resolve_speaker_id(normalized_voice)
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as handle:
            output_path = Path(handle.name)
        try:
            tts.tts_to_file(text, speaker_id, str(output_path), speed=speed)
            return self._read_waveform(output_path)
        except Exception as exc:
            raise OpenAICompatibleError(
                status_code=500,
                message=f"Speech synthesis failed: {exc}",
                error_type="server_error",
                code="synthesis_failed",
            ) from exc
        finally:
            output_path.unlink(missing_ok=True)