Files
audiobook/backend/services/tts_service.py
Paul Bailey 271a127167 Implement audio generation endpoint
- Add POST /api/tts endpoint with full request validation
- Validate text, language code, speed range, and roles list
- Integrate TTSService to generate multi-voice audiobooks
- Fix Kokoro pipeline generator handling to properly collect audio
- Stream audio response with correct Content-Type (audio/wav)
- Add Content-Disposition header with filename (audiobook.wav)
- Handle errors gracefully with appropriate HTTP status codes
- Add comprehensive test suite for endpoint validation

Requirements: 6.1, 6.7, 6.8, 9.5
2025-12-09 01:21:39 +00:00

291 lines
9.5 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
result = pipeline(chunk, voice=voice, speed=speed)
# Kokoro returns a generator that yields audio arrays
# Collect all arrays and concatenate them
audio_parts = []
for part in result:
if isinstance(part, np.ndarray):
# Flatten to 1D if needed
audio_parts.append(part.flatten())
# 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_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