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
|
||||
163
backend/test_tts_service.py
Normal file
163
backend/test_tts_service.py
Normal file
@@ -0,0 +1,163 @@
|
||||
"""
|
||||
Basic tests for TTS service functionality.
|
||||
"""
|
||||
import pytest
|
||||
import numpy as np
|
||||
from services.tts_service import TTSService
|
||||
|
||||
|
||||
def test_tts_service_initialization():
|
||||
"""Test that TTSService initializes correctly."""
|
||||
service = TTSService()
|
||||
assert service is not None
|
||||
assert service.pipelines == {}
|
||||
assert service.DEFAULT_VOICE == "af_bella"
|
||||
assert service.SAMPLE_RATE == 24000
|
||||
|
||||
|
||||
def test_get_pipeline_caching():
|
||||
"""Test that pipelines are cached per language."""
|
||||
service = TTSService()
|
||||
|
||||
# Get pipeline for American English
|
||||
pipeline1 = service.get_pipeline('a')
|
||||
assert pipeline1 is not None
|
||||
assert 'a' in service.pipelines
|
||||
|
||||
# Get same pipeline again - should return cached version
|
||||
pipeline2 = service.get_pipeline('a')
|
||||
assert pipeline1 is pipeline2
|
||||
|
||||
|
||||
def test_get_pipeline_invalid_language():
|
||||
"""Test that invalid language codes raise ValueError."""
|
||||
service = TTSService()
|
||||
|
||||
with pytest.raises(ValueError, match="Unsupported language code"):
|
||||
service.get_pipeline('x')
|
||||
|
||||
|
||||
def test_segment_text_colon_pattern():
|
||||
"""Test text segmentation with colon pattern."""
|
||||
service = TTSService()
|
||||
text = "Alice: Hello there\nBob: Hi Alice"
|
||||
role_voice_map = {"Alice": "af_bella", "Bob": "am_adam"}
|
||||
|
||||
segments = service.segment_text(text, role_voice_map)
|
||||
|
||||
assert len(segments) == 2
|
||||
assert segments[0] == ("Alice", "af_bella", "Hello there")
|
||||
assert segments[1] == ("Bob", "am_adam", "Hi Alice")
|
||||
|
||||
|
||||
def test_segment_text_bracket_pattern():
|
||||
"""Test text segmentation with bracket pattern."""
|
||||
service = TTSService()
|
||||
text = "[Narrator] The story begins\n[Alice] Hello world"
|
||||
role_voice_map = {"Narrator": "af_sarah", "Alice": "af_bella"}
|
||||
|
||||
segments = service.segment_text(text, role_voice_map)
|
||||
|
||||
assert len(segments) == 2
|
||||
assert segments[0] == ("Narrator", "af_sarah", "The story begins")
|
||||
assert segments[1] == ("Alice", "af_bella", "Hello world")
|
||||
|
||||
|
||||
def test_segment_text_default_voice_fallback():
|
||||
"""Test that unmapped roles fall back to default voice."""
|
||||
service = TTSService()
|
||||
text = "Alice: Hello\nBob: Hi"
|
||||
role_voice_map = {"Alice": "af_bella"} # Bob not mapped
|
||||
|
||||
segments = service.segment_text(text, role_voice_map)
|
||||
|
||||
assert len(segments) == 2
|
||||
assert segments[0] == ("Alice", "af_bella", "Hello")
|
||||
assert segments[1] == ("Bob", "af_bella", "Hi") # Falls back to default
|
||||
|
||||
|
||||
def test_segment_text_no_markers():
|
||||
"""Test text without role markers uses Narrator."""
|
||||
service = TTSService()
|
||||
text = "This is plain text without markers"
|
||||
role_voice_map = {"Narrator": "af_sarah"}
|
||||
|
||||
segments = service.segment_text(text, role_voice_map)
|
||||
|
||||
assert len(segments) == 1
|
||||
assert segments[0] == ("Narrator", "af_sarah", "This is plain text without markers")
|
||||
|
||||
|
||||
def test_chunk_text_short():
|
||||
"""Test that short text is not chunked."""
|
||||
service = TTSService()
|
||||
text = "This is a short sentence."
|
||||
|
||||
chunks = service._chunk_text(text, max_length=500)
|
||||
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0] == text
|
||||
|
||||
|
||||
def test_chunk_text_long():
|
||||
"""Test that long text is chunked on sentence boundaries."""
|
||||
service = TTSService()
|
||||
text = "First sentence. Second sentence. Third sentence. Fourth sentence."
|
||||
|
||||
chunks = service._chunk_text(text, max_length=30)
|
||||
|
||||
assert len(chunks) > 1
|
||||
# Each chunk should contain complete sentences
|
||||
for chunk in chunks:
|
||||
assert chunk.strip()
|
||||
|
||||
|
||||
def test_concatenate_audio_empty():
|
||||
"""Test concatenating empty list returns empty array."""
|
||||
service = TTSService()
|
||||
|
||||
result = service.concatenate_audio([])
|
||||
|
||||
assert isinstance(result, np.ndarray)
|
||||
assert len(result) == 0
|
||||
|
||||
|
||||
def test_concatenate_audio_single():
|
||||
"""Test concatenating single segment."""
|
||||
service = TTSService()
|
||||
audio = np.array([0.1, 0.2, 0.3], dtype=np.float32)
|
||||
|
||||
result = service.concatenate_audio([audio])
|
||||
|
||||
assert isinstance(result, np.ndarray)
|
||||
assert len(result) == 3
|
||||
np.testing.assert_array_equal(result, audio)
|
||||
|
||||
|
||||
def test_concatenate_audio_multiple():
|
||||
"""Test concatenating multiple segments."""
|
||||
service = TTSService()
|
||||
audio1 = np.array([0.1, 0.2], dtype=np.float32)
|
||||
audio2 = np.array([0.3, 0.4], dtype=np.float32)
|
||||
audio3 = np.array([0.5, 0.6], dtype=np.float32)
|
||||
|
||||
result = service.concatenate_audio([audio1, audio2, audio3])
|
||||
|
||||
assert isinstance(result, np.ndarray)
|
||||
assert len(result) == 6
|
||||
expected = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6], dtype=np.float32)
|
||||
np.testing.assert_array_almost_equal(result, expected)
|
||||
|
||||
|
||||
def test_concatenate_audio_filters_empty():
|
||||
"""Test that empty segments are filtered out."""
|
||||
service = TTSService()
|
||||
audio1 = np.array([0.1, 0.2], dtype=np.float32)
|
||||
audio2 = np.array([], dtype=np.float32)
|
||||
audio3 = np.array([0.3, 0.4], dtype=np.float32)
|
||||
|
||||
result = service.concatenate_audio([audio1, audio2, audio3])
|
||||
|
||||
assert len(result) == 4
|
||||
expected = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32)
|
||||
np.testing.assert_array_almost_equal(result, expected)
|
||||
Reference in New Issue
Block a user