#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Unified TTS Processor - A Comprehensive Text-to-Speech System OVERVIEW: This module provides a unified interface for advanced text-to-speech synthesis using the BATONTTS model architecture. It supports two distinct inference modes that offer different levels of control over speech generation. ARCHITECTURE: - Main Model: BATONTTS-1.7B (based on Qwen3-1.7B-instruct) - Audio Generator: CosyVoice2-0.5B for speech token to waveform conversion - Inference Engine: vLLM for efficient large language model inference - Token Processing: Custom tokenization and offset handling for speech tokens TWO INFERENCE MODES: Mode 1 - Text-to-Speech (Automatic Feature Generation): Purpose: Generate speech from text input only Input: Text string Process: Text → Model generates word_features + speech_tokens → Audio Use Case: Simple TTS when you don't need fine-grained prosodic control Implementation: Matches mode1.py exactly for consistency Mode 2 - Text+Features-to-Speech (Manual Feature Control): Purpose: Generate speech with explicit prosodic control Input: Text string + word_features (JSON) Process: Text + Features → Model generates speech_tokens → Audio Use Case: Advanced TTS with precise control over pitch, duration, etc. Implementation: Matches mode2.py exactly for consistency KEY FEATURES: - Lazy model loading for memory efficiency - Voice cloning support via prompt audio - Configurable speech speed control - Comprehensive error handling and logging - GPU memory optimization - Tensor parallelism support for multi-GPU setups USAGE PATTERNS: # Basic text-to-speech tts = UnifiedTTS() tts.text_to_speech("Hello world", "output.wav") # Advanced prosodic control features = '[{"word": "Hello", "pitch_mean": 300, "duration": 0.5}]' tts.text_features_to_speech("Hello", features, "output.wav") # Cleanup when done tts.cleanup() """ import os import sys import json import torch import torchaudio import numpy as np import argparse import uuid import logging from typing import List, Optional, Tuple, Dict, Any from transformers import AutoTokenizer from vllm import LLM, SamplingParams from modelscope import snapshot_download cosyvoice_cache_dir = snapshot_download('iic/CosyVoice2-0.5B') # Add CosyVoice paths sys.path.append('third-party/CosyVoice') sys.path.append('third-party/Matcha-TTS') from cosyvoice.cli.cosyvoice import CosyVoice2 from cosyvoice.utils.file_utils import load_wav # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) class UnifiedTTS: """ Unified Text-to-Speech Processor This class provides a unified interface for two TTS modes: 1. Mode 1 (text_to_speech): Input text only, generates word_features and speech_token internally 2. Mode 2 (text_features_to_speech): Input text and word_features, generates speech_token only Architecture: - Uses vLLM for efficient model inference - Integrates CosyVoice2 for speech token to audio conversion - Supports both emotion-controlled and standard TTS generation - Provides easy-to-use API for external script integration Usage Example: tts = UnifiedTTS() # Mode 1: Text only success = tts.text_to_speech("Hello world", "output1.wav") # Mode 2: Text + features features = '[{"word": "Hello", "pitch_mean": 300}]' success = tts.text_features_to_speech("Hello", features, "output2.wav") """ def __init__(self, model_path: str = 'Yue-Wang/BatonTTS-1.7B', cosyvoice_model_dir: str = './pretrained_models/CosyVoice2-0.5B', prompt_audio_path: str = './prompt.wav', tensor_parallel_size: int = 1, gpu_memory_utilization: float = 0.7, fp16: bool = False): """ Initialize the Unified TTS Processor Args: model_path: Path to the main TTS model (Yue-Wang/BatonTTS-1.7B) cosyvoice_model_dir: Directory path to CosyVoice2 model prompt_audio_path: Path to prompt audio file for voice cloning tensor_parallel_size: Number of GPUs for tensor parallelism gpu_memory_utilization: GPU memory utilization ratio (0.0-1.0) fp16: Whether to use half precision for CosyVoice2 Implementation: 1. Store configuration parameters 2. Initialize model loading flags 3. Defer actual model loading until first inference call """ self.model_path = model_path self.cosyvoice_model_dir = cosyvoice_model_dir self.prompt_audio_path = prompt_audio_path self.tensor_parallel_size = tensor_parallel_size self.gpu_memory_utilization = gpu_memory_utilization self.fp16 = fp16 # Model components (loaded lazily) self.llm = None self.tokenizer = None self.cosyvoice = None self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.sample_rate = None # Cached prompt features self.prompt_token = None self.prompt_feat = None self.speaker_embedding = None # Special tokens and sampling parameters self.special_tokens = {} self.sampling_params = None # Model loading flag self.models_loaded = False logger.info(f"UnifiedTTS initialized with model: {model_path}") def _load_models(self): """ Load all required models for TTS processing Model Components: 1. vLLM: Large language model for text-to-speech token generation - Uses BATONTTS-1.7B model - Handles both Mode 1 (text→features+speech) and Mode 2 (text+features→speech) - Optimized with tensor parallelism for multi-GPU setups 2. Tokenizer: Text tokenization and encoding/decoding - Handles special tokens for TTS control (, , etc.) - Converts between text and token IDs for model input/output - Supports offset correction for speech token recovery 3. CosyVoice2: Speech token to waveform conversion - Converts discrete speech tokens to continuous audio waveforms - Supports voice cloning using prompt audio features - Configurable with half precision (fp16) for memory efficiency Performance Optimizations: - Tensor parallelism for multi-GPU inference - GPU memory utilization control (default 0.7) - Half precision support for CosyVoice2 - Lazy loading pattern to reduce initialization time Usage: This method is called automatically on first inference request. All components are initialized once and reused for subsequent calls. """ if self.models_loaded: return logger.info("Loading models...") # Load main TTS model with vLLM # vLLM provides optimized inference for large language models with features like: # - Continuous batching for improved throughput # - PagedAttention for memory efficiency # - Tensor parallelism for multi-GPU scaling logger.info(f"Loading main model from {self.model_path}") self.llm = LLM( model=self.model_path, tensor_parallel_size=self.tensor_parallel_size, gpu_memory_utilization=self.gpu_memory_utilization, trust_remote_code=True, dtype="bfloat16", max_model_len=2500, ) # Load tokenizer separately # The tokenizer handles conversion between text and token IDs # It's essential for preparing model inputs and processing outputs logger.info("Loading tokenizer...") self.tokenizer = AutoTokenizer.from_pretrained(self.model_path) # Load CosyVoice2 model # CosyVoice2 converts discrete speech tokens to continuous audio waveforms # It supports voice cloning using prompt audio features for consistent voice characteristics logger.info(f"Loading CosyVoice2 model from {self.cosyvoice_model_dir}") self.cosyvoice = CosyVoice2(cosyvoice_cache_dir, fp16=self.fp16, device="cuda" if torch.cuda.is_available() else "cpu") self.sample_rate = self.cosyvoice.sample_rate # Preload prompt audio features for voice cloning # This extracts and caches features from the prompt audio file # These features are used to maintain consistent voice characteristics self._preload_prompt_features() # Configure special tokens used for TTS control # Special tokens structure the input/output format for different TTS modes self._setup_special_tokens() # Configure sampling parameters for text generation # These parameters control the quality and diversity of generated speech tokens self._setup_sampling_params() self.models_loaded = True logger.info("All models loaded successfully!") def _preload_prompt_features(self): """ Preload prompt audio features for voice cloning Feature Extraction Process: 1. Load prompt audio file and resample to 16kHz 2. Extract speech tokens using CosyVoice2 frontend 3. Resample audio to model's native sample rate 4. Extract speech features 5. Extract speaker embedding for voice characteristics These features are cached and reused for all inference calls to maintain consistent voice characteristics. """ if os.path.exists(self.prompt_audio_path): try: # Load and process prompt audio prompt_speech = load_wav(self.prompt_audio_path, 16000) # Extract speech tokens self.prompt_token, _ = self.cosyvoice.frontend._extract_speech_token(prompt_speech) logger.info(f"Preloaded prompt token, shape: {self.prompt_token.shape}") # Extract speech features prompt_speech_resample = torchaudio.transforms.Resample( orig_freq=16000, new_freq=self.sample_rate )(prompt_speech) self.prompt_feat, _ = self.cosyvoice.frontend._extract_speech_feat(prompt_speech_resample) logger.info(f"Preloaded prompt feat, shape: {self.prompt_feat.shape}") # Extract speaker embedding self.speaker_embedding = self.cosyvoice.frontend._extract_spk_embedding(prompt_speech) logger.info(f"Preloaded speaker embedding, shape: {self.speaker_embedding.shape}") except Exception as e: logger.warning(f"Failed to load prompt audio: {e}") self._use_default_prompt_features() else: logger.warning(f"Prompt audio not found: {self.prompt_audio_path}") self._use_default_prompt_features() def _use_default_prompt_features(self): """ Use default (empty) prompt features when prompt audio is unavailable Default Features: - Empty speech token tensor - Zero-filled feature tensor with correct dimensions - Zero-filled speaker embedding with standard size These defaults allow the model to generate speech without voice cloning, using the model's default voice characteristics. """ self.prompt_token = torch.zeros(1, 0, dtype=torch.int32) self.prompt_feat = torch.zeros(1, 0, 80) self.speaker_embedding = torch.zeros(1, 192) logger.info("Using default prompt features") def _setup_special_tokens(self): """ Configure special tokens for TTS control and mode switching Special Tokens Purpose: - custom_token_0: Marks the start of speech token generation phase * In Mode 1: Signals transition from text processing to speech token output * In Mode 2: Indicates where speech tokens should be inserted - custom_token_1: Marks the end of speech token generation phase * Used as stop token to terminate speech token generation * Prevents model from generating beyond intended speech sequence - custom_token_2: Additional control token for future extensions * Reserved for advanced TTS control features * Currently used in Mode 2 for feature conditioning - eos_token_id: Standard end-of-sequence marker * Separates different parts of the input sequence * Used in prompt construction for both modes Token Format Compatibility: - Stores tokens as lists to match mode1.py format exactly - Maintains backward compatibility with dictionary format - Ensures consistent token handling across different inference modes Implementation Details: - Tokens are encoded using tokenizer's call method (matching mode1.py) - Each token is stored as a list of token IDs (usually single element) - Dictionary format provides easy access for legacy code """ # Store as lists (matching mode1.py format) # This ensures exact compatibility with the original mode1.py implementation self.custom_token_0_ids = self.tokenizer('').input_ids self.custom_token_1_ids = self.tokenizer('').input_ids self.custom_token_2_ids = self.tokenizer('').input_ids self.eos_token_id = self.tokenizer.eos_token_id # Maintain dictionary format for backward compatibility # This allows existing code to continue using self.special_tokens['token_name'] self.special_tokens = { 'custom_token_0': self.custom_token_0_ids[0], 'custom_token_1': self.custom_token_1_ids[0], 'custom_token_2': self.custom_token_2_ids[0], 'eos_token_id': self.eos_token_id } logger.info("Special tokens configured") def _setup_sampling_params(self): """ Configure sampling parameters for text generation (matches mode1.py exactly) Sampling Strategy: The parameters are carefully tuned for TTS generation to balance quality, diversity, and computational efficiency. These settings match mode1.py to ensure consistent behavior across different inference implementations. Parameter Details: - temperature (0.8): Controls randomness in token selection * Lower values (0.1-0.5): More deterministic, consistent output * Higher values (0.8-1.0): More creative, diverse output * 0.6 provides optimal balance for natural speech variation - top_p (1.0): Nucleus sampling threshold for quality control * Considers all tokens in the probability distribution * Allows full model expressiveness for speech token generation * Maintains coherent speech token sequences - max_tokens (2048): Maximum tokens to generate in one pass * Accommodates long speech sequences (up to ~20 seconds of audio) * Matches mode1.py configuration for consistency * Prevents runaway generation while allowing sufficient length - stop_token_ids: Tokens that terminate generation * Uses custom_token_1 to mark end of speech token sequence * Ensures clean separation between speech tokens and other content * Prevents model from generating beyond intended speech boundary - repetition_penalty (1.1): Reduces repetitive output patterns * Slightly penalizes recently generated tokens * Improves speech naturalness by reducing monotonous patterns * Maintains speech quality while encouraging variation Usage: These parameters are used by vLLM during text generation for both Mode 1 (text→speech tokens) and Mode 2 (text+features→speech tokens). """ self.sampling_params = SamplingParams( temperature=0.6, # Balanced randomness for natural speech variation top_p=1, # Full nucleus sampling for maximum expressiveness max_tokens=2048, # Maximum sequence length (matches mode1.py) stop_token_ids=[self.eos_token_id], # Stop at end of speech tokens repetition_penalty=1.1, # Slight penalty to reduce repetitive patterns ) logger.info("Sampling parameters configured") def _prepare_mode1_input(self, text: str) -> str: """ Prepare input prompt for Mode 1 inference (matches mode1.py exactly) Mode 1 Workflow: Mode 1 performs end-to-end text-to-speech conversion in a single pass. The model generates both intermediate features and final speech tokens from the input text, making it suitable for scenarios where no specific voice characteristics are required. Input Format Structure: + text + Component Breakdown: 1. : Speech generation trigger - Signals model to begin speech token generation - Marks the start of the input sequence - Model generates speech tokens after this marker 2. text: The input text to be converted to speech - Can be any natural language text - Processed by the model's text understanding capabilities - No length restrictions beyond model's context window 3. : End-of-text marker - Signals the end of text input to the model - Separates text content from control tokens - Essential for proper model interpretation Processing Steps: 1. Tokenize input text using the model's tokenizer 2. Construct sequence: custom_token_0 + text_tokens + custom_token_1 3. Decode token sequence back to text for vLLM processing Compatibility: This implementation exactly matches mode1.py's prepare_inference_prompt method to ensure identical behavior and output quality. Args: text: Input text to synthesize Returns: Formatted prompt string for vLLM inference """ # Encode text to tokens (matching mode1.py implementation) # Using return_tensors="pt" and converting to list for consistency text_tokens = self.tokenizer(text, return_tensors="pt").input_ids[0].tolist() # Construct input sequence: + text + # This exactly matches mode1.py's prepare_inference_prompt method # The sequence structure is critical for proper model interpretation input_sequence = ( self.custom_token_0_ids + # Speech generation trigger text_tokens + # Original text content self.custom_token_1_ids # End-of-text marker ) # Convert token sequence back to text for vLLM processing # vLLM expects text input, so we decode the token sequence prompt = self.tokenizer.decode(input_sequence, skip_special_tokens=False) return prompt def _prepare_mode2_input(self, text: str, word_features: str) -> str: """ Prepare input prompt for Mode 2 inference (matches mode2.py exactly) Mode 2 Input Format (from mode2.py): + text + + generated_features + Args: text: Input text to synthesize word_features: JSON string containing word-level prosodic features (called generated_features in mode2.py) Returns: Formatted prompt string for vLLM inference Implementation Logic: 1. Tokenize input text and generated_features using the model's tokenizer 2. Construct sequence: custom_token_0 + text_tokens + custom_token_1 + features_tokens + custom_token_2 3. Decode back to text format for vLLM processing 4. Return formatted prompt The model will generate speech_tokens ending with eos_token_id This format matches the exact implementation in mode2.py """ # Encode text and generated_features to tokens (matching mode2.py implementation) text_tokens = self.tokenizer(text, return_tensors="pt").input_ids[0].tolist() features_tokens = self.tokenizer(word_features, return_tensors="pt").input_ids[0].tolist() # Construct input sequence: custom_token_0 + text + custom_token_1 + generated_features + custom_token_2 # This exactly matches mode2.py's prepare_mode2_inference_prompt method input_sequence = ( [self.special_tokens['custom_token_0']] + text_tokens + [self.special_tokens['custom_token_1']] + features_tokens + [self.special_tokens['custom_token_2']] ) # Convert token sequence back to text for vLLM prompt = self.tokenizer.decode(input_sequence, skip_special_tokens=False) return prompt def _process_features(self, features: str) -> str: """ Process word features input (JSON parsing with fallback) Processing Logic: 1. Attempt to parse as JSON for validation 2. If successful, convert back to string representation 3. If parsing fails, use original string as-is 4. Log the processing result for debugging Args: features: Input features string (JSON or plain text) Returns: Processed features string This flexible approach handles both JSON-formatted features and plain text features, ensuring compatibility with various input formats while maintaining robustness. """ try: # Try JSON parsing for validation parsed_features = json.loads(features) processed_features = str(parsed_features) logger.debug("Successfully parsed features as JSON") return processed_features except (json.JSONDecodeError, TypeError, ValueError) as e: # Use original string if JSON parsing fails logger.debug(f"Using features as plain string: {e}") return features def _generate_response(self, prompt: str) -> Tuple[str, bool]: """ Generate model response using vLLM Generation Process: 1. Use vLLM's generate method with configured sampling parameters 2. Extract generated text from model output 3. Check if generation was truncated due to length limits 4. Return both generated text and truncation status Args: prompt: Formatted input prompt for the model Returns: Tuple of (generated_text, is_truncated) The truncation flag helps identify potential quality issues when the generation hits the maximum token limit. """ outputs = self.llm.generate([prompt], self.sampling_params) # Extract generated text and check truncation output = outputs[0] generated_text = output.outputs[0].text is_truncated = output.outputs[0].finish_reason == 'length' if is_truncated: logger.warning("Generation was truncated, may affect audio quality") return generated_text, is_truncated def _extract_mode1_features(self, generated_text: str, original_prompt: str) -> str: """ Extract generated features from Mode 1 output Mode 1 Output Format: According to the sequence format: custom_token_0 + text + custom_token_1 + generated_features + custom_token_2 + features + eos_token_id This method extracts the generated_features part from the model output. The generated_features are located between custom_token_1 and custom_token_2. Args: generated_text: Model generated text original_prompt: Original input prompt Returns: JSON string containing the generated features """ try: # Find the features section between custom_token_1 and custom_token_2 custom_token_1_str = self.tokenizer.decode(self.custom_token_1_ids, skip_special_tokens=False) custom_token_2_str = self.tokenizer.decode(self.custom_token_2_ids, skip_special_tokens=False) # Look for features in the generated text full_text = original_prompt + generated_text # Find the start of features (after custom_token_1) start_idx = full_text.find(custom_token_1_str) if start_idx == -1: logger.warning("Could not find custom_token_1 in generated text") return "" start_idx += len(custom_token_1_str) # Find the end of features (before custom_token_2) end_idx = full_text.find(custom_token_2_str, start_idx) if end_idx == -1: logger.warning("Could not find custom_token_2 in generated text") # If no custom_token_2 found, try to find eos_token eos_str = self.tokenizer.decode([self.eos_token_id], skip_special_tokens=False) end_idx = full_text.find(eos_str, start_idx) if end_idx == -1: # Take the rest of the text end_idx = len(full_text) # Extract features text features_text = full_text[start_idx:end_idx].strip() if features_text: logger.info(f"Extracted features: {features_text[:100]}...") return features_text else: logger.warning("No features found in generated text") return "" except Exception as e: logger.error(f"Error extracting mode1 features: {e}") return "" def _extract_mode1_outputs(self, generated_text: str, original_prompt: str) -> List[int]: """ Extract speech tokens from Mode 1 generation output (matches mode1.py exactly) Output Processing Workflow: Mode 1 generates a sequence containing both the original prompt and new speech tokens. This method extracts only the speech tokens that represent the audio content, filtering out text tokens and control tokens to prepare for audio synthesis. Token Offset Mechanism: The model uses a token offset system to distinguish between different token types: - Text tokens: IDs in lower range (0 to ~151669) - Speech tokens: IDs in higher range (151669+ with offset applied) - The offset (151669 + 100) is used to map speech tokens back to their original values Processing Steps: 1. Merge prompt and generated text to get complete sequence - Prompt contains: text + [EOS] + - Generated text contains: speech_tokens + - Full sequence represents the complete model output 2. Encode merged text to token IDs - Converts text representation back to numerical token IDs - Preserves all tokens including special tokens and speech tokens - Uses same tokenizer as model for consistency 3. Apply offset correction to recover speech tokens - Speech tokens are stored with offset (151669 + 100) added - Subtracting offset recovers original speech token values - Only tokens above offset threshold are considered speech tokens 4. Filter and return speech tokens - Removes text tokens, control tokens, and invalid tokens - Returns clean list of speech token IDs for audio synthesis Compatibility: This implementation exactly matches mode1.py's extract_output_audio_tokens method to ensure identical behavior, output format, and audio quality. Args: generated_text: Model generated text original_prompt: Original input prompt Returns: List of speech token integers (offset corrected) """ try: # Step 1: Merge original prompt and generated text (matching mode1.py exactly) # This gives us the complete sequence that the model processed full_text = original_prompt + generated_text if not full_text.strip(): return [] # Step 2: Encode text to token IDs (matching mode1.py tokenization) # Convert the text back to numerical token IDs for processing token_ids = self.tokenizer.encode(full_text, add_special_tokens=False) # Step 3: Apply offset correction to recover original speech tokens # The model stores speech tokens with an offset of (151669 + 100) # We subtract this offset to get the original speech token values speech_tokens = [] offset = 151669 + 100 # Exact same offset as mode1.py for token_id in token_ids: if token_id >= offset: # Recover original speech token by subtracting offset original_token = token_id - offset speech_tokens.append(original_token) logger.info(f"Extracted speech_tokens count: {len(speech_tokens)}") return speech_tokens except Exception as e: logger.error(f"Error extracting mode1 outputs: {e}") return [] def _extract_mode2_outputs(self, generated_text: str, original_prompt: str) -> List[int]: """ Extract speech_tokens from Mode 2 output (matches mode2.py exactly) Extraction Logic (from mode2.py extract_speech_tokens_from_mode2_output): 1. Encode generated text to token IDs using tokenizer 2. Apply offset correction to recover original speech token values 3. Filter tokens to ensure they're in valid speech token range Args: generated_text: Model generated text original_prompt: Original input prompt (unused in mode2.py, kept for consistency) Returns: List of speech token integers (offset corrected) This exactly matches the implementation in mode2.py's extract_speech_tokens_from_mode2_output method """ try: if not generated_text.strip(): return [] # Encode generated text to token IDs (matching mode2.py) token_ids = self.tokenizer.encode(generated_text, add_special_tokens=False) # Recover original speech_token values (subtract offset) speech_tokens = [] offset = 151669 + 100 # Consistent with mode2.py implementation for token_id in token_ids: if token_id >= offset: original_token = token_id - offset speech_tokens.append(original_token) logger.info(f"Extracted speech_tokens count: {len(speech_tokens)}") return speech_tokens except Exception as e: logger.error(f"Error extracting mode2 outputs: {e}") return [] def _text_to_speech_tokens(self, text: str) -> List[int]: """ Convert text representation to speech token integers Conversion Process: 1. Tokenize text using the model's tokenizer 2. Apply offset correction to recover original token values 3. Filter tokens to ensure they're in valid speech token range 4. Return list of corrected integer tokens Args: text: Text representation of speech tokens Returns: List of speech token integers The offset correction (151669 + 100) reverses the encoding applied during model training to map speech tokens to the model's vocabulary space. """ if not text.strip(): return [] try: # Encode text to token IDs token_ids = self.tokenizer.encode(text, add_special_tokens=False) # Apply offset correction to recover original speech tokens speech_tokens = [] offset = 151669 + 100 # Consistent with original implementation for token_id in token_ids: if token_id >= offset: original_token = token_id - offset speech_tokens.append(original_token) return speech_tokens except Exception as e: logger.error(f"Text to speech_token conversion failed: {e}") return [] def _convert_tokens_to_audio(self, speech_tokens: List[int], speed: float = 1.0) -> Optional[torch.Tensor]: """ Convert speech tokens to audio waveform Conversion Process: 1. Validate speech tokens and convert to tensor format 2. Generate unique inference ID for cache management 3. Call CosyVoice2's token2wav method with prompt features 4. Clean up inference cache to prevent memory leaks 5. Return generated audio tensor Args: speech_tokens: List of speech token integers speed: Speech speed multiplier (1.0 = normal speed) Returns: Audio tensor or None if conversion fails The cache management ensures proper cleanup of temporary data structures used during audio generation. """ if len(speech_tokens) == 0: logger.warning("Empty speech tokens provided") return None try: # Convert to tensor format speech_token_tensor = torch.tensor([speech_tokens], dtype=torch.int32) logger.debug(f"Converting speech tokens, shape: {speech_token_tensor.shape}") # Generate unique inference ID inference_uuid = str(uuid.uuid1()) # Initialize cache with self.cosyvoice.model.lock: self.cosyvoice.model.hift_cache_dict[inference_uuid] = None try: # Generate audio using CosyVoice2 tts_speech = self.cosyvoice.model.token2wav( token=speech_token_tensor, prompt_token=self.prompt_token, prompt_feat=self.prompt_feat, embedding=self.speaker_embedding, token_offset=0, uuid=inference_uuid, finalize=True, speed=speed ) return tts_speech.cpu() finally: # Clean up cache with self.cosyvoice.model.lock: if inference_uuid in self.cosyvoice.model.hift_cache_dict: self.cosyvoice.model.hift_cache_dict.pop(inference_uuid) except Exception as e: logger.error(f"Audio conversion failed: {e}") return None def _save_audio(self, audio_tensor: torch.Tensor, output_path: str) -> bool: """ Save audio tensor to file Saving Process: 1. Ensure output directory exists 2. Use torchaudio to save tensor as WAV file 3. Apply correct sample rate from CosyVoice2 model 4. Handle any file I/O errors gracefully Args: audio_tensor: Generated audio waveform tensor output_path: Target file path for saving Returns: True if saving successful, False otherwise The method ensures proper directory creation and error handling for robust file operations across different environments. """ # Ensure output directory exists try: os.makedirs(os.path.dirname(output_path), exist_ok=True) except Exception: pass try: # Save audio file torchaudio.save(output_path, audio_tensor, self.sample_rate) logger.info(f"Audio saved successfully to: {output_path}") return True except Exception as e: logger.error(f"Failed to save audio: {e}") return False def text_to_speech(self, text: str, output_path: str, speed: float = 1.0) -> bool: """ Mode 1: End-to-End Text-to-Speech Conversion (matches mode1.py exactly) Overview: Mode 1 provides a complete text-to-speech pipeline that generates speech directly from input text without requiring external features or prompt audio. The model internally generates both acoustic features and speech tokens in a single pass, making it ideal for general-purpose TTS applications. Key Advantages: - Single-pass inference for optimal performance - No dependency on external features or prompt audio - Consistent voice characteristics across different texts - Simplified API for straightforward TTS tasks - Direct compatibility with mode1.py implementation Technical Workflow: 1. Model Initialization: Ensures all components are loaded and ready - vLLM engine for language model inference - Tokenizer for text processing and token conversion - CosyVoice2 for speech synthesis from tokens 2. Input Validation and Preparation: Formats text according to Mode 1 protocol - Validates input text for proper content - Constructs prompt: text + [EOS] + - This format signals the model to generate speech tokens directly 3. Speech Token Generation: Uses vLLM for efficient model inference - Applies optimized sampling parameters for quality control - Generates speech token sequence ending with - Produces discrete tokens representing audio content 4. Token Extraction and Processing: Isolates speech tokens from output - Filters out input text and control tokens - Applies offset correction (151669 + 100) to recover original values - Validates token sequence for successful audio synthesis 5. Audio Synthesis: Converts tokens to waveform using CosyVoice2 - Transforms discrete speech tokens to continuous audio signal - Applies speed adjustment if specified (default 1.0 = normal) - Generates high-quality audio with natural voice characteristics 6. File Output: Saves audio to specified path with proper formatting - Supports common audio formats (WAV, MP3, etc.) - Maintains audio quality and sample rate consistency - Provides success/failure feedback for error handling Error Handling and Validation: - Comprehensive input validation for text content - Model loading failure detection and recovery - Token generation validation and error reporting - Audio synthesis error handling with detailed logging - File I/O error management with clear feedback Performance Optimizations: - Lazy model loading to reduce initialization overhead - Efficient token processing to minimize memory usage - Optimized sampling parameters for quality-speed balance - Streamlined audio conversion pipeline Compatibility: This implementation exactly matches mode1.py's behavior to ensure: - Identical input format and processing logic - Same token extraction and offset correction - Consistent audio quality and characteristics - Compatible output format and file structure Args: text: Input text to convert to speech - Supports natural language text in multiple languages - Handles punctuation, numbers, and special characters - No strict length limitations (within model context window) - Automatically processes formatting and normalization output_path: File path where generated audio will be saved - Supports various audio formats (WAV, MP3, FLAC, etc.) - Creates parent directories if they don't exist - Overwrites existing files at the specified path - Should include appropriate file extension speed: Speech speed multiplier for tempo control - 1.0 = normal speaking speed (default) - < 1.0 = slower speech (e.g., 0.8 for 20% slower) - > 1.0 = faster speech (e.g., 1.2 for 20% faster) - Applied during audio synthesis phase Returns: bool: Success status of the TTS operation - True: Audio successfully generated and saved - False: Operation failed (check logs for details) Raises: Exception: Caught and logged, returns False for graceful handling Usage Examples: # Basic text-to-speech success = tts.text_to_speech("Hello world", "output.wav") # With custom speed success = tts.text_to_speech( "Welcome to our service", "welcome.wav", speed=0.9 ) """ try: # Ensure models are loaded self._load_models() logger.info(f"Mode 1: Text to speech - {text}") # Validate input if not text.strip(): logger.error("Empty text provided") return False # Prepare input for Mode 1 prompt = self._prepare_mode1_input(text) # Generate response generated_text, is_truncated = self._generate_response(prompt) # Extract speech_tokens (Mode 1 now only returns speech_tokens, matching mode1.py) speech_tokens = self._extract_mode1_outputs(generated_text, prompt) if len(speech_tokens) == 0: logger.error("Failed to extract speech tokens") return False # Convert to audio audio_tensor = self._convert_tokens_to_audio(speech_tokens, speed) if audio_tensor is None: logger.error("Failed to convert tokens to audio") return False # Save audio file success = self._save_audio(audio_tensor, output_path) if success: logger.info(f"Mode 1 synthesis completed successfully!") return True else: logger.error("Failed to save audio file") return False except Exception as e: logger.error(f"Mode 1 synthesis failed: {e}") return False def text_to_speech_with_features(self, text: str, output_path: str, speed: float = 1.0) -> Tuple[bool, str]: """ Mode 1: End-to-End Text-to-Speech Conversion with Features Extraction This method extends the standard text_to_speech functionality to also return the generated features that the model produces internally. This is useful for understanding the prosodic characteristics the model assigns to the text. Args: text: Input text to convert to speech output_path: File path where generated audio will be saved speed: Speech speed multiplier for tempo control Returns: Tuple[bool, str]: (success_status, generated_features_json) - success_status: True if audio generation succeeded, False otherwise - generated_features_json: JSON string of generated features, empty if failed """ try: # Ensure models are loaded self._load_models() logger.info(f"Mode 1 with features: Text to speech - {text}") # Validate input if not text.strip(): logger.error("Empty text provided") return False, "" # Prepare input for Mode 1 prompt = self._prepare_mode1_input(text) # Generate response generated_text, is_truncated = self._generate_response(prompt) # Extract features from the generated text generated_features = self._extract_mode1_features(generated_text, prompt) # Extract speech_tokens (Mode 1 now only returns speech_tokens, matching mode1.py) speech_tokens = self._extract_mode1_outputs(generated_text, prompt) if len(speech_tokens) == 0: logger.error("Failed to extract speech tokens") return False, generated_features # Convert to audio audio_tensor = self._convert_tokens_to_audio(speech_tokens, speed) if audio_tensor is None: logger.error("Failed to convert tokens to audio") return False, generated_features # Save audio file success = self._save_audio(audio_tensor, output_path) if success: logger.info(f"Mode 1 synthesis with features completed successfully!") return True, generated_features else: logger.error("Failed to save audio file") return False, generated_features except Exception as e: logger.error(f"Mode 1 synthesis with features failed: {e}") return False, "" def text_features_to_speech(self, text: str, word_features: str, output_path: str, speed: float = 1.0) -> bool: """ Mode 2: Convert text + features to speech (uses provided features) Process Flow: 1. Load models if not already loaded 2. Prepare Mode 2 input format with text and features 3. Generate model response to get speech_tokens only 4. Extract speech_tokens from the generated output 5. Convert speech_tokens to audio waveform 6. Save audio to specified output path Args: text: Input text to synthesize word_features: Pre-generated word-level prosodic features output_path: Path to save the generated audio file speed: Speech speed multiplier (1.0 = normal) Returns: True if synthesis successful, False otherwise This mode is ideal when you have specific prosodic requirements and want precise control over speech characteristics. """ try: # Ensure models are loaded self._load_models() logger.info(f"Mode 2: Text + features to speech - {text}") logger.info(f"Features: {word_features[:100]}...") # Show first 100 chars # Validate inputs if not text.strip() or not word_features.strip(): logger.error("Empty text or features provided") return False # Prepare input for Mode 2 prompt = self._prepare_mode2_input(text, word_features) # Generate response generated_text, is_truncated = self._generate_response(prompt) # Extract speech_tokens (Mode 2 now matches mode2.py implementation) speech_tokens = self._extract_mode2_outputs(generated_text, prompt) if len(speech_tokens) == 0: logger.error("Failed to extract speech tokens") return False # Convert to audio audio_tensor = self._convert_tokens_to_audio(speech_tokens, speed) if audio_tensor is None: logger.error("Failed to convert tokens to audio") return False # Save audio file success = self._save_audio(audio_tensor, output_path) if success: logger.info(f"Mode 2 synthesis completed successfully!") return True else: logger.error("Failed to save audio file") return False except Exception as e: logger.error(f"Mode 2 synthesis failed: {e}") return False def cleanup(self): """ Clean up model resources and free memory Cleanup Process: 1. Delete vLLM model instance 2. Delete CosyVoice2 model instance 3. Clear CUDA cache if available 4. Reset loading flags and cached features This method should be called when the TTS instance is no longer needed to free up GPU memory for other processes. """ if self.llm: del self.llm self.llm = None if self.cosyvoice: del self.cosyvoice self.cosyvoice = None if torch.cuda.is_available(): torch.cuda.empty_cache() self.models_loaded = False logger.info("Model resources cleaned up") def main(): """ Command-line interface for testing both TTS modes CLI Features: 1. Supports both Mode 1 (text-only) and Mode 2 (text+features) 2. Uses default configurations from single_inference_tts.py 3. Provides comprehensive argument parsing for all parameters 4. Includes example test cases for both modes Usage Examples: # Mode 1: Text only python unified_tts.py --mode 1 --text "Hello world" --output "output1.wav" # Mode 2: Text + features python unified_tts.py --mode 2 --text "Hello world" --features '[{...}]' --output "output2.wav" The CLI provides a convenient way to test the TTS functionality and serves as an example for integration into other scripts. """ parser = argparse.ArgumentParser(description='Unified TTS Script - Two modes for text-to-speech synthesis') # Model configuration arguments parser.add_argument('--model_path', type=str, default='Yue-Wang/BatonTTS-1.7B', help='Path to the main TTS model') parser.add_argument('--cosyvoice_model_dir', type=str, default='./pretrained_models/CosyVoice2-0.5B', help='Directory path to CosyVoice2 model') parser.add_argument('--prompt_audio_path', type=str, default='./prompt.wav', help='Path to prompt audio file for voice cloning') # Mode selection and input arguments parser.add_argument('--mode', type=int, choices=[1, 2], default=1, help='TTS mode: 1=text only, 2=text+features') parser.add_argument('--text', type=str, default='Kids are talking by the door', help='Input text to synthesize') parser.add_argument('--features', type=str, default='[{"word": "Kids are talking","pitch_mean": 315,"pitch_slope": 90,"energy_rms": 0.005,"energy_slope": 25,"spectral_centroid": 2650},{"word": "by the door","pitch_mean": 360,"pitch_slope": -110,"energy_rms": 0.004,"energy_slope": -30,"spectral_centroid": 2900}]', help='Word-level features for Mode 2 (JSON format)') # Output and performance arguments parser.add_argument('--output', type=str, default='unified_output.wav', help='Output audio file path') parser.add_argument('--speed', type=float, default=1.0, help='Speech speed multiplier (1.0 = normal speed)') parser.add_argument('--tensor_parallel_size', type=int, default=1, help='Number of GPUs for tensor parallelism') parser.add_argument('--gpu_memory_utilization', type=float, default=0.7, help='GPU memory utilization ratio (0.0-1.0)') parser.add_argument('--fp16', action='store_true', help='Use half precision for CosyVoice2') args = parser.parse_args() # Initialize TTS processor tts = UnifiedTTS( model_path=args.model_path, cosyvoice_model_dir=args.cosyvoice_model_dir, prompt_audio_path=args.prompt_audio_path, tensor_parallel_size=args.tensor_parallel_size, gpu_memory_utilization=args.gpu_memory_utilization, fp16=args.fp16 ) try: # Execute synthesis based on selected mode if args.mode == 1: print(f"Running Mode 1: Text to Speech") print(f"Text: {args.text}") print(f"Output: {args.output}") success, features = tts.text_to_speech_with_features( text=args.text, output_path=args.output, speed=args.speed ) if success: print(f"Features: {features}") elif args.mode == 2: print(f"Running Mode 2: Text + Features to Speech") print(f"Text: {args.text}") print(f"Features: {args.features[:100]}...") # Show first 100 chars print(f"Output: {args.output}") success = tts.text_features_to_speech( text=args.text, word_features=args.features, output_path=args.output, speed=args.speed ) # Report results if success: print(f"\n✅ Synthesis completed successfully!") print(f"Audio file saved to: {args.output}") else: print(f"\n❌ Synthesis failed!") print(f"Please check the logs for error details.") sys.exit(1) except KeyboardInterrupt: print("\n⏹️ Synthesis interrupted by user") except Exception as e: print(f"\n💥 Unexpected error: {e}") sys.exit(1) finally: # Clean up resources tts.cleanup() print("🧹 Resources cleaned up") if __name__ == "__main__": main()