audiobook/backend/services/piper_tts_service.py

"""
Piper TTS service for generating multi-voice audiobooks.
"""
import logging
import io
import os
import subprocess
from typing import Dict, List, Tuple, Optional
import numpy as np
import soundfile as sf
from pydub import AudioSegment
from pydub.effects import speedup

logger = logging.getLogger(__name__)


class PiperTTSService:
    """
    Service for text-to-speech synthesis using Piper TTS.
    
    Piper is a fast, local neural TTS system that supports multiple voices
    and allows pitch modification through audio processing.
    """
    
    # Default voice for unmapped roles
    DEFAULT_VOICE = "en_US-lessac-medium"
    
    # Sample rate for Piper TTS (typically 22050)
    SAMPLE_RATE = 22050
    
    # Available Piper voices with their characteristics
    AVAILABLE_VOICES = {
        # English voices
        "en_US-lessac-medium": {"lang": "en", "gender": "male", "quality": "medium"},
        "en_US-libritts-high": {"lang": "en", "gender": "neutral", "quality": "high"},
        "en_US-amy-medium": {"lang": "en", "gender": "female", "quality": "medium"},
        "en_GB-alan-medium": {"lang": "en", "gender": "male", "quality": "medium"},
        "en_GB-alba-medium": {"lang": "en", "gender": "female", "quality": "medium"},
    }
    
    def __init__(self):
        """Initializes Piper TTS service."""
        self.models_dir = os.path.join(os.path.dirname(__file__), "..", "piper_models")
        os.makedirs(self.models_dir, exist_ok=True)
        logger.info("PiperTTSService initialized")
    
    def _check_piper_installed(self) -> bool:
        """Check if Piper is installed and accessible."""
        try:
            result = subprocess.run(
                ["piper", "--version"],
                capture_output=True,
                text=True,
                timeout=5
            )
            return result.returncode == 0
        except (subprocess.SubprocessError, FileNotFoundError):
            return False
    
    def _download_voice_model(self, voice: str) -> bool:
        """
        Download a Piper voice model if not already present.
        
        Args:
            voice: Voice ID to download
            
        Returns:
            True if model is available, False otherwise
        """
        model_path = os.path.join(self.models_dir, f"{voice}.onnx")
        config_path = os.path.join(self.models_dir, f"{voice}.onnx.json")
        
        # Check if model already exists
        if os.path.exists(model_path) and os.path.exists(config_path):
            return True
        
        logger.info(f"Downloading Piper voice model: {voice}")
        
        # Download using piper's built-in download
        try:
            subprocess.run(
                ["piper", "--model", voice, "--download-dir", self.models_dir],
                check=True,
                timeout=60
            )
            return os.path.exists(model_path)
        except subprocess.SubprocessError as e:
            logger.error(f"Failed to download voice model {voice}: {e}")
            return False
    
    def synthesize_with_pitch(
        self,
        text: str,
        voice: str,
        speed: float = 1.0,
        pitch_shift: int = 0
    ) -> np.ndarray:
        """
        Synthesizes text with Piper and applies pitch shifting.
        
        Args:
            text: Text to synthesize
            voice: Piper voice ID
            speed: Playback speed (0.5-2.0)
            pitch_shift: Semitones to shift pitch (-12 to +12)
            
        Returns:
            Audio samples as numpy array
        """
        # Ensure voice model is available
        if not self._download_voice_model(voice):
            logger.warning(f"Voice {voice} not available, using default")
            voice = self.DEFAULT_VOICE
            self._download_voice_model(voice)
        
        model_path = os.path.join(self.models_dir, f"{voice}.onnx")
        
        # Generate audio with Piper
        try:
            # Run piper command
            process = subprocess.Popen(
                ["piper", "--model", model_path, "--output-raw"],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
            
            # Send text and get audio
            stdout, stderr = process.communicate(input=text.encode('utf-8'), timeout=30)
            
            if process.returncode != 0:
                logger.error(f"Piper synthesis failed: {stderr.decode()}")
                return np.array([], dtype=np.float32)
            
            # Convert raw audio to numpy array
            audio_data = np.frombuffer(stdout, dtype=np.int16)
            audio_float = audio_data.astype(np.float32) / 32768.0
            
            # Apply pitch shift and speed change using pydub
            if pitch_shift != 0 or speed != 1.0:
                audio_float = self._apply_audio_effects(
                    audio_float,
                    self.SAMPLE_RATE,
                    speed,
                    pitch_shift
                )
            
            return audio_float
            
        except subprocess.TimeoutExpired:
            logger.error("Piper synthesis timed out")
            return np.array([], dtype=np.float32)
        except Exception as e:
            logger.error(f"Error in Piper synthesis: {e}")
            return np.array([], dtype=np.float32)
    
    def _apply_audio_effects(
        self,
        audio: np.ndarray,
        sample_rate: int,
        speed: float,
        pitch_shift: int
    ) -> np.ndarray:
        """
        Apply speed and pitch effects to audio using pydub.
        
        Args:
            audio: Audio samples
            sample_rate: Sample rate
            speed: Speed multiplier
            pitch_shift: Semitones to shift
            
        Returns:
            Processed audio samples
        """
        try:
            # Convert numpy to pydub AudioSegment
            audio_int16 = (audio * 32768.0).astype(np.int16)
            audio_bytes = audio_int16.tobytes()
            
            audio_segment = AudioSegment(
                audio_bytes,
                frame_rate=sample_rate,
                sample_width=2,
                channels=1
            )
            
            # Apply pitch shift
            if pitch_shift != 0:
                # Pitch shift by changing frame rate then resampling
                new_sample_rate = int(sample_rate * (2.0 ** (pitch_shift / 12.0)))
                audio_segment = audio_segment._spawn(
                    audio_segment.raw_data,
                    overrides={'frame_rate': new_sample_rate}
                )
                audio_segment = audio_segment.set_frame_rate(sample_rate)
            
            # Apply speed change
            if speed != 1.0:
                audio_segment = speedup(audio_segment, playback_speed=speed)
            
            # Convert back to numpy
            samples = np.array(audio_segment.get_array_of_samples())
            audio_float = samples.astype(np.float32) / 32768.0
            
            return audio_float
            
        except Exception as e:
            logger.error(f"Error applying audio effects: {e}")
            return audio
    
    def get_available_voices(self) -> Dict[str, Dict[str, str]]:
        """
        Returns dictionary of available Piper voices.
        
        Returns:
            Dict mapping voice IDs to their characteristics
        """
        return self.AVAILABLE_VOICES.copy()