backend/magpie_adapter.py

import os
from pathlib import Path
from typing import Dict, Optional

import soundfile as sf

from backend.synthesis_catalog import MAGPIE_LANGUAGES, MAGPIE_MODEL, MAGPIE_SPEAKERS
from backend.types import VoiceConfig

try:
    import spaces
except ImportError:
    class _SpacesShim:
        @staticmethod
        def GPU(fn=None, **_kwargs):
            def decorate(inner):
                return inner

            if fn is not None:
                return decorate(fn)
            return decorate

    spaces = _SpacesShim()

MAGPIE_SPEAKER_IDS = {
    "John": 0,
    "Sofia": 1,
    "Aria": 2,
    "Jason": 3,
    "Leo": 4,
}
MAGPIE_SUPPORTED_LANGUAGES = {language["value"] for language in MAGPIE_LANGUAGES}


class MagpieAdapter:
    def __init__(
        self,
        repo_id: str = "nvidia/magpie_tts_multilingual_357m",
        checkpoint_filename: str = "magpie_tts_multilingual_357m.nemo",
        codec_model_path: str = "nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps",
    ) -> None:
        self.repo_id = repo_id
        self.checkpoint_filename = checkpoint_filename
        self.codec_model_path = codec_model_path
        self._model = None
        self._engine = "unloaded"
        self._load_error: Optional[Exception] = None

    def _checkpoint_path(self) -> str:
        from huggingface_hub import hf_hub_download

        return hf_hub_download(
            repo_id=self.repo_id,
            filename=self.checkpoint_filename,
            token=os.environ.get("HF_TOKEN"),
        )

    def _load_model(self):
        if self._model is not None:
            return self._model
        try:
            import torch
            from nemo.collections.tts.modules.magpietts_inference.utils import (
                ModelLoadConfig,
                load_magpie_model,
            )
        except Exception as exc:
            self._engine = "load_failed"
            self._load_error = exc
            return None

        config = ModelLoadConfig(
            nemo_file=self._checkpoint_path(),
            codecmodel_path=self.codec_model_path,
            legacy_codebooks=False,
            legacy_text_conditioning=False,
            hparams_from_wandb=None,
        )
        model, _ = load_magpie_model(config)
        model.eval()
        if torch.cuda.is_available():
            model.cuda()
        self._model = model
        self._engine = "magpie"
        self._load_error = None
        return self._model

    def speaker_index_for(self, speaker: str) -> int:
        if speaker not in MAGPIE_SPEAKER_IDS:
            raise ValueError(f"Unsupported Magpie speaker: {speaker}")
        return MAGPIE_SPEAKER_IDS[speaker]

    @spaces.GPU(duration=300)
    def synthesize(
        self,
        *,
        text: str,
        output_path: Path,
        voice_config: VoiceConfig,
        diffusion_steps: int,
        speed: float,
        language: Optional[str] = None,
    ) -> Dict[str, object]:
        del diffusion_steps, speed
        output_path.parent.mkdir(parents=True, exist_ok=True)
        speaker = voice_config.speaker or "Sofia"
        speaker_index = self.speaker_index_for(speaker)
        target_language = voice_config.language or language or "en"
        if target_language not in MAGPIE_SUPPORTED_LANGUAGES:
            raise ValueError(f"Unsupported Magpie language: {target_language}")

        model = self._load_model()
        if model is None:
            detail = f"{type(self._load_error).__name__}: {self._load_error}" if self._load_error else "unknown error"
            raise ValueError(
                "Magpie TTS runtime is unavailable. "
                "This app must install NVIDIA NeMo Magpie dependencies to synthesize real speech. "
                f"Model load failed with {detail}."
            )

        cleaned_text = text.strip()
        if cleaned_text and cleaned_text[-1] not in ".?!":
            cleaned_text = f"{cleaned_text}."
        audio, audio_len = model.do_tts(
            cleaned_text,
            language=target_language,
            apply_TN=voice_config.apply_text_normalization,
            speaker_index=speaker_index,
        )
        waveform = audio[0, : audio_len[0]].detach().cpu().numpy()
        sample_rate = int(getattr(model, "sample_rate", 22050))
        sf.write(str(output_path), waveform, sample_rate)
        duration_seconds = int(round(len(waveform) / sample_rate))
        return {
            "duration_seconds": max(1, duration_seconds),
            "sample_rate": sample_rate,
            "backend": "local",
            "model": MAGPIE_MODEL,
            "engine": self._engine,
        }