""" 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 from .ollama_service import OllamaService 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, ollama_service: Optional[OllamaService] = None): """Initializes TTS service with empty pipeline cache and optional Ollama service.""" self.pipelines: Dict[str, KPipeline] = {} self.ollama_service = ollama_service or OllamaService() 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, e, f, h, i, j, p, z) Returns: KPipeline instance for the language Raises: ValueError: If language code is not supported """ supported_langs = ['a', 'b', 'e', 'f', 'h', 'i', 'j', 'p', 'z'] 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 async def segment_text_with_ai( self, text: str, role_voice_map: Dict[str, str] ) -> Optional[List[Tuple[str, str, str]]]: """ Uses AI to segment text and assign speakers intelligently. 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) or None if AI unavailable """ try: # Check if Ollama is available is_available = await self.ollama_service.is_available() if not is_available: logger.info("Ollama not available for AI segmentation") return None # Get list of speakers from role_voice_map speakers = list(role_voice_map.keys()) # Request AI segmentation ai_segments = await self.ollama_service.segment_text(text, speakers) if not ai_segments: logger.warning("AI segmentation returned no segments") return None # Convert to format with voice assignments segments = [] for speaker, content in ai_segments: voice = role_voice_map.get(speaker, self.DEFAULT_VOICE) segments.append((speaker, voice, content)) logger.info(f"AI segmented text into {len(segments)} segments") return segments except Exception as e: logger.warning(f"AI segmentation failed: {e}") return None 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 using pattern matching. Identifies role markers and assigns segments accordingly. This is the fallback method when AI segmentation is not available. 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"Pattern-based segmentation produced {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 result = pipeline(chunk, voice=voice, speed=speed) # Kokoro returns a generator that yields Result objects # Each Result object has an .audio attribute (torch.Tensor) audio_parts = [] for part in result: # Extract audio tensor and convert to numpy if hasattr(part, 'audio'): audio_tensor = part.audio # Convert torch tensor to numpy and flatten audio_np = audio_tensor.cpu().numpy().flatten() audio_parts.append(audio_np) # Concatenate all parts if audio_parts: audio = np.concatenate(audio_parts) 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_script_segments( self, text: str, roles: List[str] ) -> List[Tuple[str, str]]: """ Generates script segments with speaker assignments using AI. This method is used for the script preview feature, allowing users to review and edit speaker assignments before generating audio. Args: text: Input text to segment roles: List of detected role IDs/names Returns: List of tuples: (speaker_name, text_content) Raises: Exception: If both AI and pattern-based segmentation fail """ logger.info(f"Generating script segments for {len(text)} chars with {len(roles)} roles") # Try AI segmentation first try: is_available = await self.ollama_service.is_available() if is_available: logger.info("Using AI for script segmentation") segments = await self.ollama_service.segment_text(text, roles) if segments: logger.info(f"AI segmentation produced {len(segments)} segments") return segments else: logger.warning("AI segmentation returned empty result") else: logger.info("Ollama not available for script segmentation") except Exception as e: logger.warning(f"AI segmentation failed: {e}") # Fall back to pattern-based segmentation logger.info("Using pattern-based segmentation as fallback") # Create a dummy role_voice_map for pattern-based segmentation # We don't need actual voices here, just role assignments role_voice_map = {role: self.DEFAULT_VOICE for role in roles} # Use existing segment_text method segments_with_voice = self.segment_text(text, role_voice_map) # Convert from (role, voice, text) to (role, text) segments = [(role, text) for role, voice, text in segments_with_voice] logger.info(f"Pattern-based segmentation produced {len(segments)} segments") return segments 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. Uses AI segmentation when available, falls back to pattern matching. 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) # Try AI segmentation first, fall back to pattern matching segments = await self.segment_text_with_ai(text, role_voice_map) if segments is None: logger.info("Using pattern-based segmentation as fallback") segments = self.segment_text(text, role_voice_map) else: logger.info("Using AI-powered segmentation") 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