Paul Bailey
d18fddef57
- Added segment_text() method to OllamaService for AI-based text splitting - OllamaService now annotates each text segment with the appropriate speaker - Updated TTSService to use AI segmentation when available - Falls back to pattern-based segmentation when Ollama unavailable - AI intelligently assigns dialogue to speakers and narrative to Narrator - Added comprehensive test suite for AI segmentation - Verified AI correctly splits unmarked dialogue between multiple speakers This makes the tool much more powerful - users can now input plain narrative text without explicit role markers and the AI will intelligently assign voices to different characters and the narrator.
351 lines
12 KiB
Python
351 lines
12 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, 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
|
|
|
|
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_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
|