497 lines
22 KiB
Python
Executable File
497 lines
22 KiB
Python
Executable File
#!/usr/bin/env python
|
|
"""
|
|
Chatterbox Audiobook Generator
|
|
|
|
This script converts a text file into an audiobook using the Chatterbox TTS system.
|
|
It parses the text file into manageable chunks, generates audio for each chunk,
|
|
and assembles them into a complete audiobook.
|
|
"""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import gc
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import torch
|
|
from pathlib import Path
|
|
import uuid
|
|
|
|
# Import helper to fix Python path
|
|
import import_helper
|
|
|
|
# Import backend services
|
|
from backend.app.services.tts_service import TTSService
|
|
from backend.app.services.speaker_service import SpeakerManagementService
|
|
from backend.app.services.audio_manipulation_service import AudioManipulationService
|
|
from backend.app.config import DIALOG_GENERATED_DIR, TTS_TEMP_OUTPUT_DIR
|
|
|
|
class AudiobookGenerator:
|
|
def __init__(self, speaker_id, output_base_name, device="mps",
|
|
exaggeration=0.5, cfg_weight=0.5, temperature=0.8,
|
|
pause_between_sentences=0.5, pause_between_paragraphs=1.0,
|
|
keep_model_loaded=False, cleanup_interval=10, use_subprocess=False):
|
|
"""
|
|
Initialize the audiobook generator.
|
|
|
|
Args:
|
|
speaker_id: ID of the speaker to use
|
|
output_base_name: Base name for output files
|
|
device: Device to use for TTS (mps, cuda, cpu)
|
|
exaggeration: Controls expressiveness (0.0-1.0)
|
|
cfg_weight: Controls alignment with speaker characteristics (0.0-1.0)
|
|
temperature: Controls randomness in generation (0.0-1.0)
|
|
pause_between_sentences: Pause duration between sentences in seconds
|
|
pause_between_paragraphs: Pause duration between paragraphs in seconds
|
|
keep_model_loaded: If True, keeps model loaded across chunks (more efficient but uses more memory)
|
|
cleanup_interval: How often to perform deep cleanup when keep_model_loaded=True
|
|
use_subprocess: If True, uses separate processes for each chunk (slower but guarantees memory release)
|
|
"""
|
|
self.speaker_id = speaker_id
|
|
self.output_base_name = output_base_name
|
|
self.device = device
|
|
self.exaggeration = exaggeration
|
|
self.cfg_weight = cfg_weight
|
|
self.temperature = temperature
|
|
self.pause_between_sentences = pause_between_sentences
|
|
self.pause_between_paragraphs = pause_between_paragraphs
|
|
self.keep_model_loaded = keep_model_loaded
|
|
self.cleanup_interval = cleanup_interval
|
|
self.use_subprocess = use_subprocess
|
|
self.chunk_counter = 0
|
|
|
|
# Initialize services
|
|
self.tts_service = TTSService(device=device)
|
|
self.speaker_service = SpeakerManagementService()
|
|
self.audio_manipulator = AudioManipulationService()
|
|
|
|
# Create output directories
|
|
self.output_dir = DIALOG_GENERATED_DIR / output_base_name
|
|
self.output_dir.mkdir(parents=True, exist_ok=True)
|
|
self.temp_dir = TTS_TEMP_OUTPUT_DIR / output_base_name
|
|
self.temp_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Validate speaker
|
|
self._validate_speaker()
|
|
|
|
def _validate_speaker(self):
|
|
"""Validate that the specified speaker exists."""
|
|
speaker_info = self.speaker_service.get_speaker_by_id(self.speaker_id)
|
|
if not speaker_info:
|
|
raise ValueError(f"Speaker ID '{self.speaker_id}' not found.")
|
|
if not speaker_info.sample_path:
|
|
raise ValueError(f"Speaker ID '{self.speaker_id}' has no sample path defined.")
|
|
|
|
# Store speaker info for later use
|
|
self.speaker_info = speaker_info
|
|
|
|
def _cleanup_memory(self):
|
|
"""Force memory cleanup and garbage collection."""
|
|
print("Performing memory cleanup...")
|
|
|
|
# Force garbage collection multiple times for thorough cleanup
|
|
for _ in range(3):
|
|
gc.collect()
|
|
|
|
# Clear device-specific caches
|
|
if self.device == "cuda" and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.synchronize()
|
|
# Additional CUDA cleanup
|
|
try:
|
|
torch.cuda.reset_peak_memory_stats()
|
|
except:
|
|
pass
|
|
elif self.device == "mps" and torch.backends.mps.is_available():
|
|
if hasattr(torch.mps, "empty_cache"):
|
|
torch.mps.empty_cache()
|
|
if hasattr(torch.mps, "synchronize"):
|
|
torch.mps.synchronize()
|
|
# Try to free MPS memory more aggressively
|
|
try:
|
|
import os
|
|
# This forces MPS to release memory back to the system
|
|
if hasattr(torch.mps, "set_per_process_memory_fraction"):
|
|
current_allocated = torch.mps.current_allocated_memory() if hasattr(torch.mps, "current_allocated_memory") else 0
|
|
if current_allocated > 0:
|
|
torch.mps.empty_cache()
|
|
except:
|
|
pass
|
|
|
|
# Additional aggressive cleanup
|
|
if hasattr(torch, '_C') and hasattr(torch._C, '_cuda_clearCublasWorkspaces'):
|
|
try:
|
|
torch._C._cuda_clearCublasWorkspaces()
|
|
except:
|
|
pass
|
|
|
|
print("Memory cleanup completed.")
|
|
|
|
async def _generate_chunk_subprocess(self, chunk, segment_filename_base, speaker_sample_path):
|
|
"""
|
|
Generate a single chunk using cbx-generate.py in a subprocess.
|
|
This guarantees memory is released when the process exits.
|
|
"""
|
|
output_file = self.temp_dir / f"{segment_filename_base}.wav"
|
|
|
|
# Use cbx-generate.py for single chunk generation
|
|
cmd = [
|
|
sys.executable, "cbx-generate.py",
|
|
"--sample", str(speaker_sample_path),
|
|
"--output", str(output_file),
|
|
"--text", chunk,
|
|
"--device", self.device
|
|
]
|
|
|
|
print(f"Running subprocess: {' '.join(cmd[:4])} ... (text truncated)")
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
cmd,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=300, # 5 minute timeout per chunk
|
|
cwd=Path(__file__).parent # Run from project root
|
|
)
|
|
|
|
if result.returncode != 0:
|
|
raise RuntimeError(f"Subprocess failed: {result.stderr}")
|
|
|
|
if not output_file.exists():
|
|
raise RuntimeError(f"Output file not created: {output_file}")
|
|
|
|
print(f"Subprocess completed successfully: {output_file}")
|
|
return output_file
|
|
|
|
except subprocess.TimeoutExpired:
|
|
raise RuntimeError(f"Subprocess timed out after 5 minutes")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Subprocess error: {e}")
|
|
|
|
def split_text_into_chunks(self, text, max_length=300):
|
|
"""
|
|
Split text into chunks suitable for TTS processing.
|
|
|
|
This uses the same logic as the DialogProcessorService._split_text method
|
|
but adds additional paragraph handling.
|
|
"""
|
|
# Split text into paragraphs first
|
|
paragraphs = re.split(r'\n\s*\n', text)
|
|
paragraphs = [p.strip() for p in paragraphs if p.strip()]
|
|
|
|
all_chunks = []
|
|
|
|
for paragraph in paragraphs:
|
|
# Split paragraph into sentences
|
|
sentences = re.split(r'(?<=[.!?\u2026])\s+|(?<=[.!?\u2026])(?=[\"\')\]\}\u201d\u2019])|(?<=[.!?\u2026])$', paragraph.strip())
|
|
sentences = [s.strip() for s in sentences if s and s.strip()]
|
|
|
|
chunks = []
|
|
current_chunk = ""
|
|
|
|
for sentence in sentences:
|
|
if not sentence:
|
|
continue
|
|
if not current_chunk: # First sentence for this chunk
|
|
current_chunk = sentence
|
|
elif len(current_chunk) + len(sentence) + 1 <= max_length:
|
|
current_chunk += " " + sentence
|
|
else:
|
|
chunks.append(current_chunk)
|
|
current_chunk = sentence
|
|
|
|
if current_chunk: # Add the last chunk
|
|
chunks.append(current_chunk)
|
|
|
|
# Further split any chunks that are still too long
|
|
paragraph_chunks = []
|
|
for chunk in chunks:
|
|
if len(chunk) > max_length:
|
|
# Simple split by length if a sentence itself is too long
|
|
for i in range(0, len(chunk), max_length):
|
|
paragraph_chunks.append(chunk[i:i+max_length])
|
|
else:
|
|
paragraph_chunks.append(chunk)
|
|
|
|
# Add paragraph marker
|
|
if paragraph_chunks:
|
|
all_chunks.append({"type": "paragraph", "chunks": paragraph_chunks})
|
|
|
|
return all_chunks
|
|
|
|
async def generate_audiobook(self, text_file_path):
|
|
"""
|
|
Generate an audiobook from a text file.
|
|
|
|
Args:
|
|
text_file_path: Path to the text file to convert
|
|
|
|
Returns:
|
|
Path to the generated audiobook file
|
|
"""
|
|
# Read the text file
|
|
text_path = Path(text_file_path)
|
|
if not text_path.exists():
|
|
raise FileNotFoundError(f"Text file not found: {text_file_path}")
|
|
|
|
with open(text_path, 'r', encoding='utf-8') as f:
|
|
text = f.read()
|
|
|
|
print(f"Processing text file: {text_file_path}")
|
|
print(f"Text length: {len(text)} characters")
|
|
|
|
# Split text into chunks
|
|
paragraphs = self.split_text_into_chunks(text)
|
|
total_chunks = sum(len(p["chunks"]) for p in paragraphs)
|
|
print(f"Split into {len(paragraphs)} paragraphs with {total_chunks} total chunks")
|
|
|
|
# Generate audio for each chunk
|
|
segment_results = []
|
|
chunk_count = 0
|
|
|
|
# Pre-load model if keeping it loaded
|
|
if self.keep_model_loaded:
|
|
print("Pre-loading TTS model for batch processing...")
|
|
self.tts_service.load_model()
|
|
|
|
try:
|
|
for para_idx, paragraph in enumerate(paragraphs):
|
|
print(f"Processing paragraph {para_idx+1}/{len(paragraphs)}")
|
|
|
|
for chunk_idx, chunk in enumerate(paragraph["chunks"]):
|
|
chunk_count += 1
|
|
self.chunk_counter += 1
|
|
print(f" Generating audio for chunk {chunk_count}/{total_chunks}: {chunk[:50]}...")
|
|
|
|
# Generate unique filename for this chunk
|
|
segment_filename_base = f"{self.output_base_name}_p{para_idx}_c{chunk_idx}_{uuid.uuid4().hex[:8]}"
|
|
|
|
try:
|
|
# Get absolute speaker sample path
|
|
speaker_sample_path = Path(self.speaker_info.sample_path)
|
|
if not speaker_sample_path.is_absolute():
|
|
from backend.app.config import SPEAKER_DATA_BASE_DIR
|
|
speaker_sample_path = SPEAKER_DATA_BASE_DIR / speaker_sample_path
|
|
|
|
# Generate speech for this chunk
|
|
if self.use_subprocess:
|
|
# Use subprocess for guaranteed memory release
|
|
segment_output_path = await self._generate_chunk_subprocess(
|
|
chunk=chunk,
|
|
segment_filename_base=segment_filename_base,
|
|
speaker_sample_path=speaker_sample_path
|
|
)
|
|
else:
|
|
# Load model for this chunk (if not keeping loaded)
|
|
if not self.keep_model_loaded:
|
|
print("Loading TTS model...")
|
|
self.tts_service.load_model()
|
|
|
|
# Generate speech using the TTS service
|
|
segment_output_path = await self.tts_service.generate_speech(
|
|
text=chunk,
|
|
speaker_id=self.speaker_id,
|
|
speaker_sample_path=str(speaker_sample_path),
|
|
output_filename_base=segment_filename_base,
|
|
output_dir=self.temp_dir,
|
|
exaggeration=self.exaggeration,
|
|
cfg_weight=self.cfg_weight,
|
|
temperature=self.temperature
|
|
)
|
|
|
|
# Memory management strategy based on model lifecycle
|
|
if self.use_subprocess:
|
|
# No memory management needed - subprocess handles it
|
|
pass
|
|
elif self.keep_model_loaded:
|
|
# Light cleanup after each chunk
|
|
if self.chunk_counter % self.cleanup_interval == 0:
|
|
print(f"Performing periodic deep cleanup (chunk {self.chunk_counter})")
|
|
self._cleanup_memory()
|
|
else:
|
|
# Explicit memory cleanup after generation
|
|
self._cleanup_memory()
|
|
|
|
# Unload model after generation
|
|
print("Unloading TTS model...")
|
|
self.tts_service.unload_model()
|
|
|
|
# Additional memory cleanup after model unload
|
|
self._cleanup_memory()
|
|
|
|
# Add to segment results
|
|
segment_results.append({
|
|
"type": "speech",
|
|
"path": str(segment_output_path)
|
|
})
|
|
|
|
# Add pause between sentences
|
|
if chunk_idx < len(paragraph["chunks"]) - 1:
|
|
segment_results.append({
|
|
"type": "silence",
|
|
"duration": self.pause_between_sentences
|
|
})
|
|
|
|
except Exception as e:
|
|
print(f"Error generating speech for chunk: {e}")
|
|
# Ensure model is unloaded if there was an error and not using subprocess
|
|
if not self.use_subprocess:
|
|
if not self.keep_model_loaded and self.tts_service.model is not None:
|
|
print("Unloading TTS model after error...")
|
|
self.tts_service.unload_model()
|
|
# Force cleanup after error
|
|
self._cleanup_memory()
|
|
# Continue with next chunk
|
|
|
|
# Add longer pause between paragraphs
|
|
if para_idx < len(paragraphs) - 1:
|
|
segment_results.append({
|
|
"type": "silence",
|
|
"duration": self.pause_between_paragraphs
|
|
})
|
|
|
|
finally:
|
|
# Always unload model at the end if it was kept loaded
|
|
if self.keep_model_loaded and self.tts_service.model is not None:
|
|
print("Final cleanup: Unloading TTS model...")
|
|
self.tts_service.unload_model()
|
|
self._cleanup_memory()
|
|
|
|
# Concatenate all segments
|
|
print("Concatenating audio segments...")
|
|
concatenated_filename = f"{self.output_base_name}_audiobook.wav"
|
|
concatenated_path = self.output_dir / concatenated_filename
|
|
|
|
self.audio_manipulator.concatenate_audio_segments(
|
|
segment_results=segment_results,
|
|
output_concatenated_path=concatenated_path
|
|
)
|
|
|
|
# Create ZIP archive with all files
|
|
print("Creating ZIP archive...")
|
|
zip_filename = f"{self.output_base_name}_audiobook.zip"
|
|
zip_path = self.output_dir / zip_filename
|
|
|
|
# Collect all speech segment files
|
|
speech_segment_paths = [
|
|
Path(s["path"]) for s in segment_results
|
|
if s["type"] == "speech" and Path(s["path"]).exists()
|
|
]
|
|
|
|
self.audio_manipulator.create_zip_archive(
|
|
segment_file_paths=speech_segment_paths,
|
|
concatenated_audio_path=concatenated_path,
|
|
output_zip_path=zip_path
|
|
)
|
|
|
|
print(f"Audiobook generation complete!")
|
|
print(f"Audiobook file: {concatenated_path}")
|
|
print(f"ZIP archive: {zip_path}")
|
|
|
|
# Ensure model is unloaded at the end (just in case)
|
|
if self.tts_service.model is not None:
|
|
print("Final check: Unloading TTS model...")
|
|
self.tts_service.unload_model()
|
|
|
|
return concatenated_path
|
|
|
|
async def main():
|
|
parser = argparse.ArgumentParser(description="Generate an audiobook from a text file using Chatterbox TTS")
|
|
|
|
# Create a mutually exclusive group for the main operation vs listing speakers
|
|
group = parser.add_mutually_exclusive_group(required=True)
|
|
group.add_argument("--list-speakers", action="store_true", help="List available speakers and exit")
|
|
group.add_argument("text_file", nargs="?", help="Path to the text file to convert")
|
|
|
|
# Other arguments
|
|
parser.add_argument("--speaker", "-s", help="ID of the speaker to use")
|
|
parser.add_argument("--output", "-o", help="Base name for output files (default: derived from text filename)")
|
|
parser.add_argument("--device", default="mps", choices=["mps", "cuda", "cpu"], help="Device to use for TTS (default: mps)")
|
|
parser.add_argument("--exaggeration", type=float, default=0.5, help="Controls expressiveness (0.0-1.0, default: 0.5)")
|
|
parser.add_argument("--cfg-weight", type=float, default=0.5, help="Controls alignment with speaker (0.0-1.0, default: 0.5)")
|
|
parser.add_argument("--temperature", type=float, default=0.8, help="Controls randomness (0.0-1.0, default: 0.8)")
|
|
parser.add_argument("--sentence-pause", type=float, default=0.5, help="Pause between sentences in seconds (default: 0.5)")
|
|
parser.add_argument("--paragraph-pause", type=float, default=1.0, help="Pause between paragraphs in seconds (default: 1.0)")
|
|
parser.add_argument("--keep-model-loaded", action="store_true", help="Keep model loaded between chunks (faster but uses more memory)")
|
|
parser.add_argument("--cleanup-interval", type=int, default=10, help="How often to perform deep cleanup when keeping model loaded (default: 10)")
|
|
parser.add_argument("--force-cpu-on-oom", action="store_true", help="Automatically switch to CPU if MPS/CUDA runs out of memory")
|
|
parser.add_argument("--max-chunk-length", type=int, default=300, help="Maximum chunk length for text splitting (default: 300)")
|
|
parser.add_argument("--use-subprocess", action="store_true", help="Use separate processes for each chunk (guarantees memory release but slower)")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# List speakers if requested
|
|
if args.list_speakers:
|
|
speaker_service = SpeakerManagementService()
|
|
speakers = speaker_service.get_speakers()
|
|
print("Available speakers:")
|
|
for speaker in speakers:
|
|
print(f" {speaker.id}: {speaker.name}")
|
|
return
|
|
|
|
# Validate required arguments for audiobook generation
|
|
if not args.text_file:
|
|
parser.error("text_file is required when not using --list-speakers")
|
|
|
|
if not args.speaker:
|
|
parser.error("--speaker/-s is required when not using --list-speakers")
|
|
|
|
# Determine output base name if not provided
|
|
if not args.output:
|
|
text_path = Path(args.text_file)
|
|
args.output = text_path.stem
|
|
|
|
try:
|
|
# Create audiobook generator
|
|
generator = AudiobookGenerator(
|
|
speaker_id=args.speaker,
|
|
output_base_name=args.output,
|
|
device=args.device,
|
|
exaggeration=args.exaggeration,
|
|
cfg_weight=args.cfg_weight,
|
|
temperature=args.temperature,
|
|
pause_between_sentences=args.sentence_pause,
|
|
pause_between_paragraphs=args.paragraph_pause,
|
|
keep_model_loaded=args.keep_model_loaded,
|
|
cleanup_interval=args.cleanup_interval,
|
|
use_subprocess=args.use_subprocess
|
|
)
|
|
|
|
# Generate audiobook with automatic fallback
|
|
try:
|
|
await generator.generate_audiobook(args.text_file)
|
|
except (RuntimeError, torch.OutOfMemoryError) as e:
|
|
if args.force_cpu_on_oom and "out of memory" in str(e).lower() and args.device != "cpu":
|
|
print(f"\n⚠️ {args.device.upper()} out of memory: {e}")
|
|
print("🔄 Automatically switching to CPU and retrying...")
|
|
|
|
# Create new generator with CPU
|
|
generator = AudiobookGenerator(
|
|
speaker_id=args.speaker,
|
|
output_base_name=args.output,
|
|
device="cpu",
|
|
exaggeration=args.exaggeration,
|
|
cfg_weight=args.cfg_weight,
|
|
temperature=args.temperature,
|
|
pause_between_sentences=args.sentence_pause,
|
|
pause_between_paragraphs=args.paragraph_pause,
|
|
keep_model_loaded=args.keep_model_loaded,
|
|
cleanup_interval=args.cleanup_interval,
|
|
use_subprocess=args.use_subprocess
|
|
)
|
|
|
|
await generator.generate_audiobook(args.text_file)
|
|
print("✅ Successfully completed using CPU fallback!")
|
|
else:
|
|
raise
|
|
|
|
except Exception as e:
|
|
print(f"Error: {e}", file=sys.stderr)
|
|
return 1
|
|
|
|
return 0
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(asyncio.run(main()))
|