Paul Bailey
e6f4e2ea32
- Created TTSService class with pipeline caching per language - Implemented text segmentation to identify role markers (colon and bracket patterns) - Implemented segment synthesis with voice assignment - Implemented audio concatenation at 24 kHz sample rate - Added fallback to default voice (af_bella) for unmapped roles - Supports language codes: a, b, f, j - Handles long text by chunking on sentence boundaries - Added comprehensive unit tests (13 tests, all passing) Requirements: 10.1, 10.3, 10.4, 6.2, 6.3, 6.4, 6.5, 6.6
280 lines
9.0 KiB
Python
280 lines
9.0 KiB
Python
"""
|
|
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
|