""" TTS service for generating multi-voice audiobooks using Kokoro TTS. """ import logging import re import io from typing import Dict, List, Tuple, Optional import numpy as np import soundfile as sf from kokoro import KPipeline logger = logging.getLogger(__name__) class TTSService: """ Service for text-to-speech synthesis using Kokoro TTS. Handles pipeline initialization, text segmentation, voice assignment, and audio concatenation for multi-voice audiobook generation. """ # Default voice for unmapped roles DEFAULT_VOICE = "af_bella" # Sample rate for Kokoro TTS SAMPLE_RATE = 24000 # Role marker patterns COLON_PATTERN = r'^([A-Z][a-zA-Z\s]+?):\s*(.+)$' BRACKET_PATTERN = r'^\[([A-Z][a-zA-Z\s]+?)\]\s*(.+)$' def __init__(self): """Initializes TTS service with empty pipeline cache.""" self.pipelines: Dict[str, KPipeline] = {} logger.info("TTSService initialized") def get_pipeline(self, lang_code: str) -> KPipeline: """ Returns or creates KPipeline for given language. Caches pipelines to avoid reloading. Args: lang_code: Language code (a, b, f, j) Returns: KPipeline instance for the language Raises: ValueError: If language code is not supported """ supported_langs = ['a', 'b', 'f', 'j'] if lang_code not in supported_langs: raise ValueError( f"Unsupported language code: {lang_code}. " f"Supported codes: {supported_langs}" ) # Return cached pipeline if available if lang_code in self.pipelines: logger.debug(f"Using cached pipeline for language: {lang_code}") return self.pipelines[lang_code] # Create new pipeline logger.info(f"Initializing KPipeline for language: {lang_code}") try: pipeline = KPipeline(lang_code=lang_code) self.pipelines[lang_code] = pipeline logger.info(f"Successfully initialized pipeline for language: {lang_code}") return pipeline except Exception as e: logger.error(f"Failed to initialize pipeline for language {lang_code}: {e}") raise def segment_text( self, text: str, role_voice_map: Dict[str, str] ) -> List[Tuple[str, str, str]]: """ Segments text into (role_id, voice, text_content) tuples. Identifies role markers and assigns segments accordingly. Args: text: Input text to segment role_voice_map: Mapping of role IDs to voice IDs Returns: List of tuples: (role_id, voice_id, text_content) """ segments = [] lines = text.split('\n') for line in lines: line = line.strip() if not line: continue # Try colon pattern first colon_match = re.match(self.COLON_PATTERN, line, re.MULTILINE) if colon_match: role_name = colon_match.group(1).strip() content = colon_match.group(2).strip() voice = role_voice_map.get(role_name, self.DEFAULT_VOICE) segments.append((role_name, voice, content)) continue # Try bracket pattern bracket_match = re.match(self.BRACKET_PATTERN, line) if bracket_match: role_name = bracket_match.group(1).strip() content = bracket_match.group(2).strip() voice = role_voice_map.get(role_name, self.DEFAULT_VOICE) segments.append((role_name, voice, content)) continue # No role marker found - use default narrator voice = role_voice_map.get("Narrator", self.DEFAULT_VOICE) segments.append(("Narrator", voice, line)) logger.info(f"Segmented text into {len(segments)} segments") return segments def _chunk_text(self, text: str, max_length: int = 500) -> List[str]: """ Splits long text into manageable chunks. Splits on sentence boundaries when possible. Args: text: Text to chunk max_length: Maximum characters per chunk Returns: List of text chunks """ if len(text) <= max_length: return [text] chunks = [] # Split on sentence boundaries sentences = re.split(r'([.!?]+\s+)', text) current_chunk = "" for i in range(0, len(sentences), 2): sentence = sentences[i] punctuation = sentences[i + 1] if i + 1 < len(sentences) else "" full_sentence = sentence + punctuation if len(current_chunk) + len(full_sentence) <= max_length: current_chunk += full_sentence else: if current_chunk: chunks.append(current_chunk.strip()) current_chunk = full_sentence if current_chunk: chunks.append(current_chunk.strip()) return chunks if chunks else [text] def synthesize_segment( self, pipeline: KPipeline, text: str, voice: str, speed: float = 1.0 ) -> np.ndarray: """ Synthesizes a single text segment with specified voice. Handles long text by chunking. Args: pipeline: KPipeline instance text: Text to synthesize voice: Kokoro voice ID speed: Playback speed (0.5-1.5) Returns: Audio samples as numpy array """ # Chunk long text chunks = self._chunk_text(text) audio_chunks = [] for chunk in chunks: try: # Generate audio for chunk audio, _ = pipeline(chunk, voice=voice, speed=speed) audio_chunks.append(audio) except Exception as e: logger.error(f"Error synthesizing chunk with voice {voice}: {e}") raise # Concatenate chunks if multiple if len(audio_chunks) > 1: return np.concatenate(audio_chunks) elif audio_chunks: return audio_chunks[0] else: return np.array([], dtype=np.float32) def concatenate_audio(self, segments: List[np.ndarray]) -> np.ndarray: """ Concatenates multiple audio segments into single array. Ensures consistent sample rate (24kHz). Args: segments: List of audio arrays Returns: Concatenated audio array """ if not segments: return np.array([], dtype=np.float32) # Filter out empty segments valid_segments = [s for s in segments if len(s) > 0] if not valid_segments: return np.array([], dtype=np.float32) # Concatenate all segments concatenated = np.concatenate(valid_segments) logger.info(f"Concatenated {len(valid_segments)} audio segments") return concatenated async def generate_audiobook( self, text: str, lang: str, role_voice_map: Dict[str, str], speed: float = 1.0 ) -> bytes: """ Main orchestration method for audiobook generation. Segments text, synthesizes each segment, concatenates, and returns audio bytes. Args: text: Input text lang: Language code role_voice_map: Role to voice ID mapping speed: Playback speed Returns: Audio data as bytes (WAV format) """ logger.info( f"Generating audiobook: lang={lang}, speed={speed}, " f"roles={len(role_voice_map)}" ) # Get or create pipeline for language pipeline = self.get_pipeline(lang) # Segment text by roles segments = self.segment_text(text, role_voice_map) if not segments: raise ValueError("No segments to synthesize") # Synthesize each segment audio_segments = [] for role_id, voice, content in segments: logger.debug(f"Synthesizing segment for role '{role_id}' with voice '{voice}'") audio = self.synthesize_segment(pipeline, content, voice, speed) audio_segments.append(audio) # Concatenate all audio final_audio = self.concatenate_audio(audio_segments) # Convert to bytes (WAV format) buffer = io.BytesIO() sf.write(buffer, final_audio, self.SAMPLE_RATE, format='WAV') buffer.seek(0) audio_bytes = buffer.read() logger.info(f"Generated audiobook: {len(audio_bytes)} bytes") return audio_bytes