Implement TTS service with Kokoro integration
- 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
This commit is contained in:
279
backend/services/tts_service.py
Normal file
279
backend/services/tts_service.py
Normal file
@@ -0,0 +1,279 @@
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user