Paul Bailey
cea97d3cd4
- Modified TTSRequest to accept optional segments array - Updated generate_audiobook endpoint to use segments if provided - Added generate_audiobook_from_segments method to TTSService - Ensured backward compatibility with text-based generation - Added comprehensive tests for segment-based generation - All tests passing (7/7) Requirements: 15.7
470 lines
16 KiB
Python
470 lines
16 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
|
|
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
|
|
|
|
async def generate_audiobook_from_segments(
|
|
self,
|
|
segments: List[Tuple[str, str]],
|
|
lang: str,
|
|
role_voice_map: Dict[str, str],
|
|
speed: float = 1.0
|
|
) -> bytes:
|
|
"""
|
|
Generates audiobook from pre-segmented script with speaker assignments.
|
|
|
|
This method is used when the user has reviewed and edited the script,
|
|
bypassing automatic segmentation.
|
|
|
|
Args:
|
|
segments: List of (speaker_name, text_content) tuples
|
|
lang: Language code
|
|
role_voice_map: Role to voice ID mapping
|
|
speed: Playback speed
|
|
|
|
Returns:
|
|
Audio data as bytes (WAV format)
|
|
|
|
Raises:
|
|
ValueError: If segments is empty or invalid
|
|
"""
|
|
logger.info(
|
|
f"Generating audiobook from {len(segments)} edited segments: "
|
|
f"lang={lang}, speed={speed}, roles={len(role_voice_map)}"
|
|
)
|
|
|
|
if not segments:
|
|
raise ValueError("No segments provided for synthesis")
|
|
|
|
# Get or create pipeline for language
|
|
pipeline = self.get_pipeline(lang)
|
|
|
|
# Convert segments to format with voice assignments
|
|
segments_with_voice = []
|
|
for speaker, text in segments:
|
|
voice = role_voice_map.get(speaker, self.DEFAULT_VOICE)
|
|
segments_with_voice.append((speaker, voice, text))
|
|
|
|
logger.info(f"Processing {len(segments_with_voice)} segments with voice assignments")
|
|
|
|
# Synthesize each segment
|
|
audio_segments = []
|
|
for role_id, voice, content in segments_with_voice:
|
|
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 from edited segments: {len(audio_bytes)} bytes")
|
|
return audio_bytes
|