audio-engine-hub/app/engines/xtts.py

"""
NovaAi – TTS-Engine-Hub
engines/xtts.py
Version: v0.1.0

Description:
  Coqui XTTS v2 engine adapter.
  Supports multilingual synthesis and voice cloning via reference audio.
"""

import os
import asyncio
import logging
import torch
from .engine_base import TTSEngineBase
from app.config import settings

logger = logging.getLogger(__name__)

class XTTSEngine(TTSEngineBase):
    def __init__(self):
        logger.debug("XTTSEngine __init__ started.")
        self.device = "cpu"
        if torch.cuda.is_available():
            logger.debug("CUDA is available.")
            if settings.XTTS_DEVICE == "cuda":
                self.device = "cuda"
                logger.debug(f"XTTS_DEVICE setting is 'cuda'. Using CUDA.")
            else:
                logger.debug(f"XTTS_DEVICE setting is '{settings.XTTS_DEVICE}'. Falling back to CPU despite CUDA availability.")
        else:
            logger.debug("CUDA is not available. Using CPU.")

        self.model = None
        self.tts = None
        
        # Verify license acceptance
        if not settings.XTTS_ACCEPT_LICENSE:
            logger.warning("XTTS license not accepted. Engine will not load. Set XTTS_ACCEPT_LICENSE=true in .env")
            return

        try:
            from TTS.api import TTS
            logger.debug("Coqui TTS library imported successfully.")
        except ImportError:
            logger.error("Coqui TTS library not found. Install 'TTS' via pip.")
            return

        logger.info(f"Initializing XTTS v2 on {self.device}...")
        try:
            # Set environment variable to bypass TTS library's interactive license prompt
            # This tells the TTS library that we agree to the terms
            os.environ['COQUI_TOS_AGREED'] = '1'

            # Initialize TTS with the model name.
            # This will download the model if not present.
            # We use the official model name.
            logger.debug(f"Calling TTS('tts_models/multilingual/multi-dataset/xtts_v2').to({self.device})...")
            self.tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(self.device)
            logger.info("XTTS v2 model loaded successfully.")
        except Exception as e:
            logger.error(f"Failed to load XTTS model: {e}", exc_info=True) # exc_info=True to log traceback
            self.tts = None
        logger.debug("XTTSEngine __init__ finished.")

    def list_models(self):
        return ["xtts_v2"]

    def list_voices(self, model: str = None):
        """
        Returns a list of available reference audio files (speakers) 
        found in the VOICES_DIR.
        """
        voices_dir = settings.VOICES_DIR
        if not os.path.exists(voices_dir):
            return ["default"]
        
        # List .wav files in the voices directory
        voices = [f for f in os.listdir(voices_dir) if f.lower().endswith(".wav")]
        return sorted(voices) if voices else ["default"]

    def healthcheck(self):
        if not settings.XTTS_ACCEPT_LICENSE:
            return {"status": "license_not_accepted", "detail": "Set XTTS_ACCEPT_LICENSE=true"}
        if self.tts is None:
            return {"status": "error", "detail": "Model not loaded"}
        return {"status": "ok", "device": self.device}

    async def synthesize(self, text: str, speaker: str = None, model: str = None, fmt: str = "wav"):
        """
        Synthesize speech using XTTS v2.
        
        Args:
            text: Text to synthesize.
            speaker: Filename of the reference audio in VOICES_DIR (e.g., "my_voice.wav").
            model: Ignored (only xtts_v2 supported).
            fmt: Output format (wav by default).
        """
        if not self.tts:
            raise RuntimeError("XTTS engine is not initialized or license not accepted.")

        # Resolve speaker/reference audio
        voices_dir = settings.VOICES_DIR
        if not os.path.exists(voices_dir):
             os.makedirs(voices_dir, exist_ok=True)

        # precise path handling
        speaker_wav = None
        if speaker and speaker != "default":
             potential_path = os.path.join(voices_dir, speaker)
             if os.path.exists(potential_path):
                 speaker_wav = potential_path
             else:
                 # Check if speaker has extension, if not try adding .wav
                 if not speaker.lower().endswith(".wav"):
                      potential_path_ext = os.path.join(voices_dir, f"{speaker}.wav")
                      if os.path.exists(potential_path_ext):
                          speaker_wav = potential_path_ext
        
        # Fallback if no valid speaker provided - XTTS NEEDS a speaker reference.
        # We'll use a default sample if provided, or fail.
        # Ideally, we should ship a default reference.
        if not speaker_wav:
             # Try to find *any* wav file in the dir to use as default
             available = self.list_voices()
             if available and available[0] != "default":
                 speaker_wav = os.path.join(voices_dir, available[0])
                 logger.warning(f"No valid speaker '{speaker}' found. Using first available: {available[0]}")
             else:
                 raise ValueError("XTTS requires a reference audio file (speaker). Please upload a .wav file to app/asset/voices/")

        # Output file
        import tempfile
        fd, output_path = tempfile.mkstemp(suffix=".wav", prefix="xtts_")
        os.close(fd)

        # Run synthesis in thread pool to avoid blocking event loop
        # XTTS API: tts.tts_to_file(text=..., speaker_wav=..., language=..., file_path=...)
        # We need to detect language or default to English ("en")
        # For now, we hardcode "en" or try to auto-detect if the library supports it, 
        # but tts_to_file usually requires language for multilingual models.
        language = "en" # TODO: Add language parameter to API or auto-detect
        
        logger.info(f"Synthesizing with XTTS. Speaker: {os.path.basename(speaker_wav)}, Lang: {language}")

        try:
            await asyncio.to_thread(
                self.tts.tts_to_file,
                text=text,
                speaker_wav=speaker_wav,
                language=language,
                file_path=output_path
            )
        except Exception as e:
            logger.error(f"XTTS synthesis failed: {e}")
            if os.path.exists(output_path):
                os.remove(output_path)
            raise RuntimeError(f"XTTS synthesis failed: {str(e)}")

        return output_path

    async def selftest(self):
        try:
            # Check if we have at least one reference voice
            voices = self.list_voices()
            if not voices or voices == ["default"]:
                return {"selftest": False, "error": "No reference voices found in asset/voices", "engine": "xtts"}
            
            test_voice = voices[0]
            output = await self.synthesize("XTTS selftest.", speaker=test_voice)
            
            if os.path.exists(output) and os.path.getsize(output) > 0:
                os.remove(output)
                return {"selftest": True, "engine": "xtts"}
            return {"selftest": False, "error": "Output file empty or missing", "engine": "xtts"}
        except Exception as e:
             return {"selftest": False, "error": str(e), "engine": "xtts"}