From 4ead99754a1b843a9095874a8b593ab06d17b0b2 Mon Sep 17 00:00:00 2001 From: malibrated Date: Mon, 21 Apr 2025 22:30:33 -0700 Subject: [PATCH] Minor improvements to account for large-VRAM Apple Silicon machines, and refine the inferencing code to ignore spurious, non-audio responses from the Orpheus model, use the system prompt if it exists in the directory structure, and for long texts, 1) use nltk for segmenting the text according to user-specified max character lengths, and 2) use the previous segment (For segment index n>1) as context in the system prompt, to hopefully improve audio generation consistency. --- .DS_Store | Bin 0 -> 8196 bytes .env.example | 4 + .gitignore | 9 +- README.md | 40 ++ app.py | 30 +- restart.flag | 1 + templates/tts.html | 20 +- tts_engine/__init__.py | 40 +- tts_engine/inference.mps | 1124 ++++++++++++++++++++++++++++++++++++++ tts_engine/inference.py | 936 +++++++++++++++++++++++-------- tts_engine/speechpipe.py | 190 +++++-- 11 files changed, 2116 insertions(+), 278 deletions(-) create mode 100644 .DS_Store create mode 100644 restart.flag create mode 100644 tts_engine/inference.mps diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..24e4bf262c9a87b7fc6826364dc59c57c9bfa92a GIT binary patch literal 8196 zcmeHMU2GIp6h3EK+Lk?i4=?xs_YV^JeY5dskc z5dskc5dskc4<`aam=+g^`w!x(+#E1Cd|ozB*un z%czVHh!B{M06V)E!$ued6TI2y_Y%l}13DOBfjfJkl)7`^!MI<#se(J6vYZsnb3HX~ znXcE=bQ6W5;zcE;l*&qv_l`M*UQhFKZo8H}%36JNN018ZT57;F3Q1*ohhux1X&YJJ zz|;*=IW}n8x>IO(a<=aJw!M5nqLfr>CMFtNS{ve9woW$0CpNdXY;TCS?r5Hzl*FpK zO}l!Mr>%_bT;kit;L`zY&$H8IYRuUcU1V6{O~#y|5DPZLoFU82cW^9-ZD?O3^t?IxAVv}UDqmgX64mc^9aA0VC^@tV=ZU{YuvsSH;#mx}j;y z&b^&abahXcFR4%}m&)=8p=x{PDbvXI7fjt78`j*kVe6(nI+!nE~Gk?OPObKrgu)M6UAv|NkzF_tFD%mkHlj-J>+3WXDjLch+39Q8pnR1_fq)*FYE&n1w{T2F1>Cikr;N05`N- zV0gB*OcoxknIt>g56U(evlZCd_bV0giE7ySf8WCQ|MzRTXr2gx2!V$X0W9w9?PQ@j z4BK;`k)5^UluuG-iRn#o>6=i&R9L@%949&RhauJD!B`@YQ(V#rr3Zcy@Q;5pNAW+3 J|Np=E|2IuEm?r=L literal 0 HcmV?d00001 diff --git a/.env.example b/.env.example index 11ad5da..12888bf 100644 --- a/.env.example +++ b/.env.example @@ -14,6 +14,10 @@ ORPHEUS_TOP_P=0.9 ORPHEUS_SAMPLE_RATE=24000 ORPHEUS_MODEL_NAME=Orpheus-3b-FT-Q8_0.gguf # Model name sent to inference server (Q2_K, Q4_K_M, or Q8_0 variants) +# Audio processing parameters +ORPHEUS_MAX_BATCH_CHARS=600 # Maximum characters per batch for sentence splitting (100-2000) +ORPHEUS_CROSSFADE_MS=30 # Crossfade duration between audio batches in milliseconds (10-200) + # Web UI settings (keep in mind that the web UI is not secure and should not be exposed to the internet) ORPHEUS_PORT=5005 ORPHEUS_HOST=0.0.0.0 diff --git a/.gitignore b/.gitignore index f6f9952..710bb3c 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,11 @@ __pycache__/ .venv/ venv/ models/ -*.gguf \ No newline at end of file +*.gguf +outputs/ +output_long*.txt +*.txt +jinja +*.zsh +*.bak +*.sav diff --git a/README.md b/README.md index 867c46e..83261b3 100644 --- a/README.md +++ b/README.md @@ -401,6 +401,46 @@ Note: Repetition penalty is hardcoded to 1.1 and cannot be changed through envir Make sure the `ORPHEUS_API_URL` points to your running inference server. +## Apple Silicon Optimization + +This project includes special optimizations for Apple Silicon (M1/M2/M3) devices: + +### Metal Performance Shaders (MPS) + +The application automatically uses Apple's Metal Performance Shaders (MPS) backend when running on Apple Silicon. This provides significant performance improvements over CPU-only execution. + +### CoreML Neural Engine Acceleration + +For even better performance, you can enable Neural Engine acceleration via CoreML: + +1. Install the required package: + ```bash + pip install coremltools + ``` + +2. Export the model to CoreML format (one-time setup): + ```bash + python export_coreml.py + ``` + +3. Enable CoreML acceleration: + ```bash + export ORPHEUS_USE_COREML=1 + ``` + +This optional step can provide an additional 30-50% speedup by offloading computation to the Neural Engine, which is specifically designed for machine learning workloads. + +### Memory Optimization + +The application automatically detects high-memory Apple Silicon devices (M1 Pro/Max/Ultra, M2 Pro/Max/Ultra, M3 Pro/Max/Ultra) and uses more aggressive memory optimization strategies, including: + +- Parallel batch processing +- Larger audio buffers +- Pre-allocation for crossfading +- Optimized tensor operations + +These optimizations are particularly effective on devices with 32GB+ of unified memory. + ## Development ### Project Components diff --git a/app.py b/app.py index 1ca2108..6d41c13 100644 --- a/app.py +++ b/app.py @@ -51,7 +51,16 @@ from fastapi.templating import Jinja2Templates from pydantic import BaseModel import json -from tts_engine import generate_speech_from_api, AVAILABLE_VOICES, DEFAULT_VOICE, VOICE_TO_LANGUAGE, AVAILABLE_LANGUAGES +# Import from tts_engine with fallback values +try: + from tts_engine import generate_speech_from_api, AVAILABLE_VOICES, DEFAULT_VOICE, VOICE_TO_LANGUAGE, AVAILABLE_LANGUAGES, MAX_BATCH_CHARS +except ImportError as e: + # If MAX_BATCH_CHARS is missing, import what we can and define a fallback value + print(f"Warning: Import error - {e}") + from tts_engine import generate_speech_from_api, AVAILABLE_VOICES, DEFAULT_VOICE, VOICE_TO_LANGUAGE, AVAILABLE_LANGUAGES + # Use a fallback value - this should match the value in inference.py + MAX_BATCH_CHARS = 600 + print(f"Using fallback value for MAX_BATCH_CHARS: {MAX_BATCH_CHARS}") # Create FastAPI app app = FastAPI( @@ -96,7 +105,7 @@ async def create_speech_api(request: SpeechRequest): Generate speech from text using the Orpheus TTS model. Compatible with OpenAI's /v1/audio/speech endpoint. - For longer texts (>1000 characters), batched generation is used + For longer texts (>MAX_BATCH_CHARS characters), batched generation is used to improve reliability and avoid truncation issues. """ if not request.input: @@ -107,7 +116,7 @@ async def create_speech_api(request: SpeechRequest): output_path = f"outputs/{request.voice}_{timestamp}.wav" # Check if we should use batched generation - use_batching = len(request.input) > 1000 + use_batching = len(request.input) > MAX_BATCH_CHARS if use_batching: print(f"Using batched generation for long text ({len(request.input)} characters)") @@ -117,8 +126,7 @@ async def create_speech_api(request: SpeechRequest): prompt=request.input, voice=request.voice, output_file=output_path, - use_batching=use_batching, - max_batch_chars=1000 # Process in ~1000 character chunks (roughly 1 paragraph) + use_batching=use_batching ) end = time.time() generation_time = round(end - start, 2) @@ -160,7 +168,7 @@ async def speak(request: Request): output_path = f"outputs/{voice}_{timestamp}.wav" # Check if we should use batched generation for longer texts - use_batching = len(text) > 1000 + use_batching = len(text) > MAX_BATCH_CHARS if use_batching: print(f"Using batched generation for long text ({len(text)} characters)") @@ -170,8 +178,7 @@ async def speak(request: Request): prompt=text, voice=voice, output_file=output_path, - use_batching=use_batching, - max_batch_chars=1000 + use_batching=use_batching ) end = time.time() generation_time = round(end - start, 2) @@ -226,7 +233,7 @@ async def save_config(request: Request): # Convert values to proper types for key, value in data.items(): - if key in ["ORPHEUS_MAX_TOKENS", "ORPHEUS_API_TIMEOUT", "ORPHEUS_PORT", "ORPHEUS_SAMPLE_RATE"]: + if key in ["ORPHEUS_MAX_TOKENS", "ORPHEUS_API_TIMEOUT", "ORPHEUS_PORT", "ORPHEUS_SAMPLE_RATE", "ORPHEUS_MAX_BATCH_CHARS", "ORPHEUS_CROSSFADE_MS"]: try: data[key] = str(int(value)) except (ValueError, TypeError): @@ -322,7 +329,7 @@ async def generate_from_web( output_path = f"outputs/{voice}_{timestamp}.wav" # Check if we should use batched generation for longer texts - use_batching = len(text) > 1000 + use_batching = len(text) > MAX_BATCH_CHARS if use_batching: print(f"Using batched generation for long text from web form ({len(text)} characters)") @@ -332,8 +339,7 @@ async def generate_from_web( prompt=text, voice=voice, output_file=output_path, - use_batching=use_batching, - max_batch_chars=1000 + use_batching=use_batching ) end = time.time() generation_time = round(end - start, 2) diff --git a/restart.flag b/restart.flag new file mode 100644 index 0000000..e5ebbac --- /dev/null +++ b/restart.flag @@ -0,0 +1 @@ +1745296574.27762 \ No newline at end of file diff --git a/templates/tts.html b/templates/tts.html index 3cf715c..3f8b705 100644 --- a/templates/tts.html +++ b/templates/tts.html @@ -183,6 +183,7 @@
{% if voice_option == "tara" %}Female, English, conversational, clear + {% elif voice_option == "kaya" %} Female, English, expressive, friendly {% elif voice_option == "leah" %}Female, English, warm, gentle {% elif voice_option == "jess" %}Female, English, energetic, youthful {% elif voice_option == "leo" %}Male, English, authoritative, deep @@ -330,6 +331,23 @@ class="block w-full rounded-md bg-dark-700 border-dark-600 text-white text-sm focus:border-primary-500 focus:ring-primary-500 focus:ring-offset-dark-800 px-3 py-2">
+ +
+ + +

Maximum characters per batch for sentence splitting (default: 600)

+
+ +
+ + +

Crossfade duration between audio batches in milliseconds (default: 30)

+
+
diff --git a/tts_engine/__init__.py b/tts_engine/__init__.py index 88bc054..59d1c22 100644 --- a/tts_engine/__init__.py +++ b/tts_engine/__init__.py @@ -4,14 +4,52 @@ TTS Engine package for Orpheus text-to-speech system. This package contains the core components for audio generation: - inference.py: Token generation and API handling - speechpipe.py: Audio conversion pipeline +- coreml_wrapper.py: Apple Silicon Neural Engine acceleration """ # Make key components available at package level from .inference import ( generate_speech_from_api, + stream_audio, AVAILABLE_VOICES, DEFAULT_VOICE, VOICE_TO_LANGUAGE, AVAILABLE_LANGUAGES, - list_available_voices + MAX_BATCH_CHARS, + CROSSFADE_MS, + list_available_voices, + API_URL, + HEADERS ) + +# Expose hardware detection flags +from .speechpipe import ( + APPLE_SILICON, + CUDA_AVAILABLE, + DEVICE +) + +__all__ = [ + # Core Functions + "generate_speech_from_api", + "stream_audio", + "list_available_voices", + + # Speech Processing + "convert_to_audio", + "turn_token_into_id", + "reset_state", + "DEVICE", + + # Constants & Settings + "AVAILABLE_VOICES", + "DEFAULT_VOICE", + "VOICE_TO_LANGUAGE", + "AVAILABLE_LANGUAGES", + "MAX_BATCH_CHARS", + "CROSSFADE_MS", + + # Configuration (Example) + "API_URL", + "HEADERS", +] diff --git a/tts_engine/inference.mps b/tts_engine/inference.mps new file mode 100644 index 0000000..a2dbdf9 --- /dev/null +++ b/tts_engine/inference.mps @@ -0,0 +1,1124 @@ +import os +import sys +import requests +import json +import time +import wave +import numpy as np +import sounddevice as sd +import argparse +import threading +import queue +import asyncio +from concurrent.futures import ThreadPoolExecutor +from typing import List, Dict, Any, Optional, Generator, Union, Tuple +from dotenv import load_dotenv + +# Helper to detect if running in Uvicorn's reloader +def is_reloader_process(): + """Check if the current process is a uvicorn reloader""" + return (sys.argv[0].endswith('_continuation.py') or + os.environ.get('UVICORN_STARTED') == 'true') + +# Set a flag to avoid repeat messages +IS_RELOADER = is_reloader_process() +if not IS_RELOADER: + os.environ['UVICORN_STARTED'] = 'true' + +# Load environment variables from .env file +load_dotenv() + +# Detect hardware capabilities and display information +import torch +import psutil + +# Detect if we're on a high-end system based on hardware capabilities +HIGH_END_GPU = False +HAS_MPS = False # new flag for MPS detection +DEVICE_INFO = "CPU" # keep track of the primary compute device + +if torch.cuda.is_available(): + # Get GPU properties + props = torch.cuda.get_device_properties(0) + gpu_name = props.name + gpu_mem_gb = props.total_memory / (1024**3) + compute_capability = f"{props.major}.{props.minor}" + + # Consider high-end if: large VRAM (≥16GB) OR high compute capability (≥8.0) OR large VRAM (≥12GB) with good CC (≥7.0) + HIGH_END_GPU = (gpu_mem_gb >= 16.0 or + props.major >= 8 or + (gpu_mem_gb >= 12.0 and props.major >= 7)) + + + if HIGH_END_GPU: + DEVICE_INFO = "HIGH-end CUDA GPU" + if not IS_RELOADER: + print(f"🖥️ Hardware: High-end CUDA GPU detected") + print(f"📊 Device: {gpu_name}") + print(f"📊 VRAM: {gpu_mem_gb:.2f} GB") + print(f"📊 Compute Capability: {compute_capability}") + print("🚀 Using high-performance optimizations") + else: + DEVICE_INFO = "CUDA GPU" + if not IS_RELOADER: + print(f"🖥️ Hardware: CUDA GPU detected") + print(f"📊 Device: {gpu_name}") + print(f"📊 VRAM: {gpu_mem_gb:.2f} GB") + print(f"📊 Compute Capability: {compute_capability}") + print("🚀 Using GPU-optimized settings") +#--- ADD MPS Detection block --- +elif torch.backends.mps.is_available(): + HAS_MPS = True + # Consider M-series chips with sufficient RAM as high-performance + ram_gb = psutil.virtual_memory().total / (1024**3) + # Let's consider M-series with >= 32GB RAM as "high performance" for TTS tasks + HIGH_PERFORMANCE_MPS = ram_gb >= 32.0 + DEVICE_INFO = "Apple Silicon (MPS)" + + if not IS_RELOADER: + print(f"🖥️ Hardware: Apple Silicon (MPS) detected") + print(f"📊 RAM: {ram_gb:.2f} GB") + if HIGH_PERFORMANCE_MPS: + print("🚀 Using high-performance MPS optimizations") + else: + print("⚙️ Using MPS-optimized settings") +else: + # Get CPU info + cpu_cores = psutil.cpu_count(logical=False) + cpu_threads = psutil.cpu_count(logical=True) + ram_gb = psutil.virtual_memory().total / (1024**3) + + if not IS_RELOADER: + print(f"🖥️ Hardware: CPU only (No CUDA GPU or MPS detected)") + print(f"📊 CPU: {cpu_cores} cores, {cpu_threads} threads") + print(f"📊 RAM: {ram_gb:.2f} GB") + print("⚙️ Using CPU-optimized settings") + +# Load configuration from environment variables without hardcoded defaults +# Critical settings - will log errors if missing +required_settings = ["ORPHEUS_API_URL"] +missing_settings = [s for s in required_settings if s not in os.environ] +if missing_settings: + print(f"ERROR: Missing required environment variable(s): {', '.join(missing_settings)}") + print("Please set them in .env file or environment. See .env.example for defaults.") + +# API connection settings +API_URL = os.environ.get("ORPHEUS_API_URL") +if not API_URL: + print("WARNING: ORPHEUS_API_URL not set. API calls will fail until configured.") + +HEADERS = { + "Content-Type": "application/json" +} + +# Request timeout settings +try: + REQUEST_TIMEOUT = int(os.environ.get("ORPHEUS_API_TIMEOUT", "120")) +except (ValueError, TypeError): + print("WARNING: Invalid ORPHEUS_API_TIMEOUT value, using 120 seconds as fallback") + REQUEST_TIMEOUT = 120 + +# Model generation parameters from environment variables +try: + MAX_TOKENS = int(os.environ.get("ORPHEUS_MAX_TOKENS", "8192")) +except (ValueError, TypeError): + print("WARNING: Invalid ORPHEUS_MAX_TOKENS value, using 8192 as fallback") + MAX_TOKENS = 8192 + +try: + TEMPERATURE = float(os.environ.get("ORPHEUS_TEMPERATURE", "0.6")) +except (ValueError, TypeError): + print("WARNING: Invalid ORPHEUS_TEMPERATURE value, using 0.6 as fallback") + TEMPERATURE = 0.6 + +try: + TOP_P = float(os.environ.get("ORPHEUS_TOP_P", "0.9")) +except (ValueError, TypeError): + print("WARNING: Invalid ORPHEUS_TOP_P value, using 0.9 as fallback") + TOP_P = 0.9 + +# Repetition penalty is hardcoded to 1.1 which is the only stable value for quality output +REPETITION_PENALTY = 1.1 + +try: + SAMPLE_RATE = int(os.environ.get("ORPHEUS_SAMPLE_RATE", "24000")) +except (ValueError, TypeError): + print("WARNING: Invalid ORPHEUS_SAMPLE_RATE value, using 24000 as fallback") + SAMPLE_RATE = 24000 + +# Print loaded configuration only in the main process, not in the reloader +if not IS_RELOADER: + print(f"Configuration loaded:") + print(f" API_URL: {API_URL}") + print(f" MAX_TOKENS: {MAX_TOKENS}") + print(f" TEMPERATURE: {TEMPERATURE}") + print(f" TOP_P: {TOP_P}") + print(f" REPETITION_PENALTY: {REPETITION_PENALTY}") + + +# Available voices based on the Orpheus-TTS repository +AVAILABLE_VOICES = ["tara", "elise", "leah", "jess", "leo", "dan", "mia", "zac", "zoe"] +DEFAULT_VOICE = "tara" # Best voice according to documentation + +# Import the unified token handling from speechpipe +from .speechpipe import turn_token_into_id, token_id_cache, CUSTOM_TOKEN_PREFIX + +# Special token IDs for Orpheus model +START_TOKEN_ID = 128259 +END_TOKEN_IDS = [128009, 128260, 128261, 128257] + +# Performance monitoring +class PerformanceMonitor: + """Track and report performance metrics""" + def __init__(self): + self.start_time = time.time() + self.token_count = 0 + self.audio_chunks = 0 + self.last_report_time = time.time() + self.report_interval = 2.0 # seconds + + def add_tokens(self, count: int = 1) -> None: + self.token_count += count + self._check_report() + + def add_audio_chunk(self) -> None: + self.audio_chunks += 1 + self._check_report() + + def _check_report(self) -> None: + current_time = time.time() + if current_time - self.last_report_time >= self.report_interval: + self.report() + self.last_report_time = current_time + + def report(self) -> None: + elapsed = time.time() - self.start_time + if elapsed < 0.001: + return + + tokens_per_sec = self.token_count / elapsed + chunks_per_sec = self.audio_chunks / elapsed + + # Estimate audio duration based on audio chunks (each chunk is ~0.085s of audio) + est_duration = self.audio_chunks * 0.085 + + print(f"Progress: {tokens_per_sec:.1f} tokens/sec, est. {est_duration:.1f}s audio generated, {self.token_count} tokens, {self.audio_chunks} chunks in {elapsed:.1f}s") + +# Create global performance monitor +perf_monitor = PerformanceMonitor() + +def format_prompt(prompt: str, voice: str = DEFAULT_VOICE) -> str: + """Format prompt for Orpheus model with voice prefix and special tokens.""" + # Validate voice and provide fallback + if voice not in AVAILABLE_VOICES: + print(f"Warning: Voice '{voice}' not recognized. Using '{DEFAULT_VOICE}' instead.") + voice = DEFAULT_VOICE + + # Format similar to how engine_class.py does it with special tokens + formatted_prompt = f"{voice}: {prompt}" + + # Add special token markers for the Orpheus-FASTAPI + special_start = "<|audio|>" # Using the additional_special_token from config + special_end = "<|eot_id|>" # Using the eos_token from config + + return f"{special_start}{formatted_prompt}{special_end}" + +def generate_tokens_from_api(prompt: str, voice: str = DEFAULT_VOICE, temperature: float = TEMPERATURE, + top_p: float = TOP_P, max_tokens: int = MAX_TOKENS, + repetition_penalty: float = REPETITION_PENALTY) -> Generator[str, None, None]: + """Generate tokens from text using OpenAI-compatible API with optimized streaming and retry logic.""" + start_time = time.time() + formatted_prompt = format_prompt(prompt, voice) + print(f"Generating speech for: {formatted_prompt}") + + # Optimize the token generation for GPUs + if HIGH_END_GPU: + # Use more aggressive parameters for faster generation on high-end GPUs + print("Using optimized parameters for high-end GPU") + elif torch.cuda.is_available(): + print("Using optimized parameters for GPU acceleration") + elif HAS_MPS: # Add check for MPS + print("Using optimized parameters for Apple Silicon (MPS) acceleration") + + # Create the request payload (model field may not be required by some endpoints but included for compatibility) + payload = { + "prompt": formatted_prompt, + "max_tokens": max_tokens, + "temperature": temperature, + "top_p": top_p, + "repeat_penalty": repetition_penalty, + "stream": True # Always stream for better performance + } + + # Add model field - this is ignored by many local inference servers for /v1/completions + # but included for compatibility with OpenAI API and some servers that may use it + model_name = os.environ.get("ORPHEUS_MODEL_NAME", "Orpheus-3b-FT-Q8_0.gguf") + payload["model"] = model_name + + # Session for connection pooling and retry logic + session = requests.Session() + + retry_count = 0 + max_retries = 3 + + while retry_count < max_retries: + try: + # Make the API request with streaming and timeout + response = session.post( + API_URL, + headers=HEADERS, + json=payload, + stream=True, + timeout=REQUEST_TIMEOUT + ) + + if response.status_code != 200: + print(f"Error: API request failed with status code {response.status_code}") + print(f"Error details: {response.text}") + # Retry on server errors (5xx) but not on client errors (4xx) + if response.status_code >= 500: + retry_count += 1 + wait_time = 2 ** retry_count # Exponential backoff + print(f"Retrying in {wait_time} seconds...") + time.sleep(wait_time) + continue + return + + # Process the streamed response with better buffering + buffer = "" + token_counter = 0 + + # Iterate through the response to get tokens + for line in response.iter_lines(): + if line: + line_str = line.decode('utf-8') + if line_str.startswith('data: '): + data_str = line_str[6:] # Remove the 'data: ' prefix + + if data_str.strip() == '[DONE]': + break + + try: + data = json.loads(data_str) + if 'choices' in data and len(data['choices']) > 0: + token_chunk = data['choices'][0].get('text', '') + for token_text in token_chunk.split('>'): + token_text = f'{token_text}>' + token_counter += 1 + perf_monitor.add_tokens() + + if token_text: + yield token_text + except json.JSONDecodeError as e: + print(f"Error decoding JSON: {e}") + continue + + # Generation completed successfully + generation_time = time.time() - start_time + tokens_per_second = token_counter / generation_time if generation_time > 0 else 0 + print(f"Token generation complete: {token_counter} tokens in {generation_time:.2f}s ({tokens_per_second:.1f} tokens/sec)") + return + + except requests.exceptions.Timeout: + print(f"Request timed out after {REQUEST_TIMEOUT} seconds") + retry_count += 1 + if retry_count < max_retries: + wait_time = 2 ** retry_count + print(f"Retrying in {wait_time} seconds... (attempt {retry_count+1}/{max_retries})") + time.sleep(wait_time) + else: + print("Max retries reached. Token generation failed.") + return + + except requests.exceptions.ConnectionError: + print(f"Connection error to API at {API_URL}") + retry_count += 1 + if retry_count < max_retries: + wait_time = 2 ** retry_count + print(f"Retrying in {wait_time} seconds... (attempt {retry_count+1}/{max_retries})") + time.sleep(wait_time) + else: + print("Max retries reached. Token generation failed.") + return + +# The turn_token_into_id function is now imported from speechpipe.py +# This eliminates duplicate code and ensures consistent behavior + +def convert_to_audio(multiframe: List[int], count: int) -> Optional[bytes]: + """Convert token frames to audio with performance monitoring.""" + # Import here to avoid circular imports + from .speechpipe import convert_to_audio as orpheus_convert_to_audio + start_time = time.time() + result = orpheus_convert_to_audio(multiframe, count) + + if result is not None: + perf_monitor.add_audio_chunk() + + return result + +# async def tokens_decoder(token_gen) -> Generator[bytes, None, None]: +# """Simplified token decoder with early first-chunk processing for lower latency.""" +# buffer = [] +# count = 0 +# +# # Use different thresholds for first chunk vs. subsequent chunks +# first_chunk_processed = False +# min_frames_first = 7 # Process after just 7 tokens for first chunk (ultra-low latency) +# min_frames_subsequent = 28 # Default for reliability after first chunk (4 chunks of 7) +# ideal_frames = 49 # Define ideal_frames here +# process_every = 7 # Process every 7 tokens (standard for Orpheus model) +# +# start_time = time.time() +# last_log_time = start_time +# token_count = 0 +# +# has_seen_audio_token = False # <<< New flag +# spurious_text_detected = False # <<< New flag +# +# async for token_text in token_gen: +# token = turn_token_into_id(token_text, count) +# if token is not None and token > 0: +# has_seen_audio_token = True # Mark that we are processing audio +# +# # Add to buffer using simple append (reliable method) +# buffer.append(token) +# count += 1 +# token_count += 1 +# +# # Log throughput periodically +# current_time = time.time() +# if current_time - last_log_time > 5.0: # Every 5 seconds +# elapsed = current_time - start_time +# if elapsed > 0: +# print(f"Token processing rate: {token_count/elapsed:.1f} tokens/second") +# last_log_time = current_time +# +# # Different processing paths based on whether first chunk has been processed +# if not first_chunk_processed: +# # For first audio output, process as soon as we have enough tokens for one chunk +# if count >= min_frames_first: +# buffer_to_proc = buffer[-min_frames_first:] +# +# # Process the first chunk for immediate audio feedback +# print(f"Processing first audio chunk with {len(buffer_to_proc)} tokens") +# audio_samples = convert_to_audio(buffer_to_proc, count) +# if audio_samples is not None: +# first_chunk_processed = True # Mark first chunk as processed +# yield audio_samples +# else: +# # For subsequent chunks, use standard processing with larger batch +# if count % process_every == 0 and count >= min_frames_subsequent: +# # Use simple slice operation - reliable and correct +# buffer_to_proc = buffer[-min_frames_subsequent:] +# +# # Debug output to help diagnose issues +# if count % 28 == 0: +# print(f"Processing buffer with {len(buffer_to_proc)} tokens, total collected: {len(buffer)}") +# +# # Process the tokens +# audio_samples = convert_to_audio(buffer_to_proc, count) +# if audio_samples is not None: +# yield audio_samples +# +# # <<< START NEW LOGIC >>> +# # Check for spurious text *after* having seen audio tokens +# elif has_seen_audio_token and token is None and token_text.strip() and token_text.strip() != '>' and not token_text.startswith('>> +# +# # --- End-of-generation handling (Now runs after loop break) --- +# print(f"End of token stream reached. Processing final buffer of length {len(buffer)}. Spurious text detected: {spurious_text_detected}") +# +# # --- START REVISED End-of-generation handling --- +# if not spurious_text_detected: +# print(f"End of token stream reached cleanly. Processing final buffer of length {len(buffer)}.") +# +# start_index = 0 +# final_chunks_processed = 0 +# +# # Process all remaining complete chunks of size 'process_every' +# while start_index + process_every <= len(buffer): +# chunk_to_proc = buffer[start_index : start_index + process_every] +# print(f"Processing final full chunk #{final_chunks_processed + 1} (indices {start_index}-{start_index + process_every - 1})") +# # Note: Passing 'count' to convert_to_audio might be irrelevant if it only uses len(chunk)//7 internally. +# # Check speechpipe.py if 'count' is actually needed there. Assuming it's not critical for slicing. +# audio_samples = convert_to_audio(chunk_to_proc, count) +# if audio_samples is not None: +# yield audio_samples +# final_chunks_processed += 1 +# start_index += process_every +# +# # Process the final partial chunk (if any) with padding +# if start_index < len(buffer): +# remainder_chunk = buffer[start_index:] +# num_remaining = len(remainder_chunk) +# print(f"Processing final partial chunk (indices {start_index}-end, length {num_remaining}) with padding") +# last_token = remainder_chunk[-1] +# padding_needed = process_every - num_remaining +# padding = [last_token] * padding_needed +# padded_buffer = remainder_chunk + padding +# +# # Pass the correctly sized padded buffer +# audio_samples = convert_to_audio(padded_buffer, count) # Again, 'count' might be ignored by the function +# if audio_samples is not None: +# yield audio_samples +# final_chunks_processed += 1 +# elif final_chunks_processed == 0 and len(buffer) > 0: +# # This case handles buffers smaller than process_every initially +# print(f"Processing final small buffer (length {len(buffer)} < {process_every}) with padding") +# last_token = buffer[-1] +# padding_needed = process_every - len(buffer) +# padding = [last_token] * padding_needed +# padded_buffer = buffer + padding +# audio_samples = convert_to_audio(padded_buffer, count) +# if audio_samples is not None: +# yield audio_samples +# final_chunks_processed += 1 +# +# if final_chunks_processed > 0: +# print(f"Processed {final_chunks_processed} final chunk(s).") +# else: +# print("No final audio chunks processed from remaining buffer.") +# +# else: +# # This block executes if spurious text WAS detected +# print("Skipping final buffer processing due to spurious text detection.") +# # --- END REVISED End-of-generation handling --- + +# In inference.py + +# In inference.py -> tokens_decoder + +async def tokens_decoder(token_gen) -> Generator[bytes, None, None]: + """ + Processes tokens using standard buffering, handles spurious text. + (Low-latency first chunk removed for compatibility with vocoder slice). + """ + buffer = [] + count = 0 + process_every = 7 + min_frames_process = 28 # Minimum tokens to have before processing a chunk + ideal_frames_process = 49 # Ideal number of tokens to send (using tail) + + start_time = time.time() + token_count_perf = 0 + last_log_time = start_time + has_seen_audio_token = False + spurious_text_detected = False + + print("[Decoder] Initialized (No Low-Latency First Chunk).") + + async for token_text in token_gen: + token = turn_token_into_id(token_text, count) + + if token is not None and token > 0: + has_seen_audio_token = True + buffer.append(token) + count += 1 + token_count_perf += 1 + + # --- START MODIFIED PROCESSING LOGIC --- + # Process only when buffer is sufficiently full + if len(buffer) >= min_frames_process: + # Check if we have enough for an ideal chunk from the tail + if len(buffer) >= ideal_frames_process: + buffer_to_proc = buffer[-ideal_frames_process:] + # print(f"[Decoder] Attempting ideal chunk ({ideal_frames_process} tokens from end of buffer len {len(buffer)})") + else: + # Otherwise, take the minimum required from the tail + buffer_to_proc = buffer[-min_frames_process:] + print(f"[Decoder] Attempting min chunk ({min_frames_process} tokens from end of buffer len {len(buffer)})") + + audio_samples = convert_to_audio(buffer_to_proc, count) + if audio_samples is not None: + print(f"[Decoder] Yielding chunk ({len(audio_samples)} bytes)") + yield audio_samples + # *** Still don't remove from buffer if using tail slicing *** + else: + print(f"[Decoder] WARNING: Chunk conversion failed! Input tokens (last 10): {buffer_to_proc[-10:]}") + # --- END MODIFIED PROCESSING LOGIC --- + + + # --- Handle Spurious Text Detection (Same as before) --- + elif has_seen_audio_token and token is None and token_text.strip() and token_text.strip() != '>' and not token_text.startswith(' 5.0: + elapsed_interval = current_time - last_log_time + if elapsed_interval > 0: + rate = token_count_perf / elapsed_interval + print(f"Token processing rate: {rate:.1f} tokens/second (in last {elapsed_interval:.1f}s)") + last_log_time = current_time + token_count_perf = 0 + + # --- End-of-generation Handling --- + print(f"[Decoder] Exited main loop. Final unprocessed buffer len: {len(buffer)}. Spurious detected: {spurious_text_detected}") + + # Process remaining buffer content only if stream ended cleanly + if not spurious_text_detected and len(buffer) > 0: + print(f"[Decoder] Processing final remaining buffer content (length {len(buffer)}).") + + # How many full chunks remain after the initial chunk was removed? + # Note: This final processing logic needs refinement if the main loop doesn't remove processed tokens. + # Let's try the simpler approach first (like the original 'elif' logic) which relies on tail slicing. + + final_buffer_processed = False + if len(buffer) >= ideal_frames: + print(f"[Decoder] Processing final ideal frame block ({ideal_frames} tokens from end)") + buffer_to_proc = buffer[-ideal_frames:] + audio_samples = convert_to_audio(buffer_to_proc, count) + if audio_samples is not None: + yield audio_samples + final_buffer_processed = True + elif len(buffer) >= min_frames_subsequent: + print(f"[Decoder] Processing final subsequent frame block ({min_frames_subsequent} tokens from end)") + buffer_to_proc = buffer[-min_frames_subsequent:] + audio_samples = convert_to_audio(buffer_to_proc, count) + if audio_samples is not None: + yield audio_samples + final_buffer_processed = True + elif len(buffer) >= process_every: # Must be >= 7 for convert_to_audio + print(f"[Decoder] Processing final padded frame block (>= {process_every} tokens from end)") + # Take the whole remaining buffer if it's less than min_frames_subsequent but >= process_every + remainder_chunk = buffer + num_remaining = len(remainder_chunk) + last_token = remainder_chunk[-1] + # Pad up to process_every size (or maybe min_frames_subsequent?) - Let's try process_every + padding_needed = process_every - num_remaining if num_remaining > 0 else 0 + padding = [last_token] * padding_needed + padded_buffer = remainder_chunk + padding + + if len(padded_buffer) >= process_every: # Ensure padding creates a valid chunk + audio_samples = convert_to_audio(padded_buffer, count) + if audio_samples is not None: + yield audio_samples + final_buffer_processed = True + else: + print(f"[Decoder] WARNING: Final padded chunk conversion failed!") + else: + print(f"[Decoder] WARNING: Final padded buffer too short ({len(padded_buffer)}), skipping.") + + + print(f"[Decoder] Final buffer processed flag: {final_buffer_processed}") + + elif spurious_text_detected: + print("[Decoder] Skipped final buffer processing.") + else: + print("[Decoder] No final buffer processing needed (buffer empty).") + + print(f"[Decoder] Finished. Total valid tokens received: {count}.") # Removed processed_count as it's complex with tail slicing + +def tokens_decoder_sync(syn_token_gen, output_file=None): + """Optimized synchronous wrapper with parallel processing and efficient file I/O.""" + # Use a larger queue for high-end systems + queue_size = 100 if HIGH_END_GPU else 50 + + #token_id_cache.clear() + #print("Cleared token ID cache.") + + audio_queue = queue.Queue(maxsize=queue_size) + audio_segments = [] + + # If output_file is provided, prepare WAV file with buffered I/O + wav_file = None + if output_file: + # Create directory if it doesn't exist + os.makedirs(os.path.dirname(os.path.abspath(output_file)), exist_ok=True) + wav_file = wave.open(output_file, "wb") + wav_file.setnchannels(1) + wav_file.setsampwidth(2) + wav_file.setframerate(SAMPLE_RATE) + + # Batch processing of tokens for improved throughput + batch_size = 32 if HIGH_END_GPU else 16 + + # Thread synchronization for proper completion detection + producer_done_event = threading.Event() + producer_started_event = threading.Event() + + # Convert the synchronous token generator into an async generator with batching + async def async_token_gen(): + batch = [] + for token in syn_token_gen: + batch.append(token) + if len(batch) >= batch_size: + for t in batch: + yield t + batch = [] + # Process any remaining tokens in the final batch + for t in batch: + yield t + + async def async_producer(): + # Track performance with more granular metrics + start_time = time.time() + chunk_count = 0 + last_log_time = start_time + + try: + # Signal that producer has started processing + producer_started_event.set() + + async for audio_chunk in tokens_decoder(async_token_gen()): + # Process each audio chunk from the decoder + if audio_chunk: + audio_queue.put(audio_chunk) + chunk_count += 1 + + # Log performance periodically + current_time = time.time() + if current_time - last_log_time >= 3.0: # Every 3 seconds + elapsed = current_time - last_log_time + if elapsed > 0: + recent_chunks = chunk_count + chunks_per_sec = recent_chunks / elapsed + print(f"Audio generation rate: {chunks_per_sec:.2f} chunks/second") + last_log_time = current_time + # Reset chunk counter for next interval + chunk_count = 0 + except Exception as e: + print(f"Error in token processing: {str(e)}") + import traceback + traceback.print_exc() + finally: + # Always signal completion, even if there was an error + print("Producer completed - setting done event") + producer_done_event.set() + # Add sentinel to queue to signal end of stream + audio_queue.put(None) + + def run_async(): + """Run the async producer in its own thread""" + asyncio.run(async_producer()) + + # Use a separate thread with higher priority for producer + thread = threading.Thread(target=run_async, name="TokenProcessor") + thread.daemon = True # Allow thread to be terminated when main thread exits + thread.start() + + # Wait for producer to actually start before proceeding + # This avoids race conditions where we might try to read from an empty queue + # before the producer has had a chance to add anything + producer_started_event.wait(timeout=5.0) + + # Optimized I/O approach for all systems + # This approach is simpler and more reliable than separate code paths + write_buffer = bytearray() + buffer_max_size = 1024 * 1024 # 1MB max buffer size (adjustable) + + # Keep track of the last time we checked for completion + last_check_time = time.time() + check_interval = 1.0 # Check producer status every second + + # Process audio chunks until we're done + while True: + try: + # Get the next audio chunk with a short timeout + # This allows us to periodically check status and handle other events + audio = audio_queue.get(timeout=0.1) + + # None marker indicates end of stream + if audio is None: + print("Received end-of-stream marker") + break + + # Store the audio segment for return value + audio_segments.append(audio) + + # Write to file if needed + if wav_file: + write_buffer.extend(audio) + + # Flush buffer if it's large enough + if len(write_buffer) >= buffer_max_size: + wav_file.writeframes(write_buffer) + write_buffer = bytearray() # Reset buffer + + except queue.Empty: + # No data available right now + current_time = time.time() + + # Periodically check if producer is done + if current_time - last_check_time > check_interval: + last_check_time = current_time + + # If producer is done and queue is empty, we're finished + if producer_done_event.is_set() and audio_queue.empty(): + print("Producer done and queue empty - finishing consumer") + break + + # Flush buffer periodically even if not full + if wav_file and len(write_buffer) > 0: + wav_file.writeframes(write_buffer) + write_buffer = bytearray() # Reset buffer + + # Extra safety check - ensure thread is done + if thread.is_alive(): + print("Waiting for token processor thread to complete...") + thread.join(timeout=10.0) + if thread.is_alive(): + print("WARNING: Token processor thread did not complete within timeout") + + # Final flush of any remaining data + if wav_file and len(write_buffer) > 0: + print(f"Final buffer flush: {len(write_buffer)} bytes") + wav_file.writeframes(write_buffer) + + # Close WAV file if opened + if wav_file: + wav_file.close() + if output_file: + print(f"Audio saved to {output_file}") + + # Calculate and print detailed performance metrics + if audio_segments: + total_bytes = sum(len(segment) for segment in audio_segments) + duration = total_bytes / (2 * SAMPLE_RATE) # 2 bytes per sample at 24kHz + total_time = time.time() - perf_monitor.start_time + realtime_factor = duration / total_time if total_time > 0 else 0 + + print(f"Generated {len(audio_segments)} audio segments") + print(f"Generated {duration:.2f} seconds of audio in {total_time:.2f} seconds") + print(f"Realtime factor: {realtime_factor:.2f}x") + + if realtime_factor < 1.0: + print("⚠️ Warning: Generation is slower than realtime") + else: + print(f"✓ Generation is {realtime_factor:.1f}x faster than realtime") + + return audio_segments + +def stream_audio(audio_buffer): + """Stream audio buffer to output device with error handling.""" + if audio_buffer is None or len(audio_buffer) == 0: + return + + try: + # Convert bytes to NumPy array (16-bit PCM) + audio_data = np.frombuffer(audio_buffer, dtype=np.int16) + + # Normalize to float in range [-1, 1] for playback + audio_float = audio_data.astype(np.float32) / 32767.0 + + # Play the audio with proper device selection and error handling + sd.play(audio_float, SAMPLE_RATE) + sd.wait() + except Exception as e: + print(f"Audio playback error: {e}") + +import re +import numpy as np +from io import BytesIO +import wave + +def split_text_into_sentences(text): + """Split text into sentences with a more reliable approach.""" + # We'll use a simple approach that doesn't rely on variable-width lookbehinds + # which aren't supported in Python's regex engine + + # First, split on common sentence ending punctuation + # This isn't perfect but works for most cases and avoids the regex error + parts = [] + current_sentence = "" + + for char in text: + current_sentence += char + + # If we hit a sentence ending followed by a space, consider this a potential sentence end + if char in (' ', '\n', '\t') and len(current_sentence) > 1: + prev_char = current_sentence[-2] + if prev_char in ('.', '!', '?'): + # Check if this is likely a real sentence end and not an abbreviation + # (Simple heuristic: if there's a space before the period, it's likely a real sentence end) + if len(current_sentence) > 3 and current_sentence[-3] not in ('.', ' '): + parts.append(current_sentence.strip()) + current_sentence = "" + + # Add any remaining text + if current_sentence.strip(): + parts.append(current_sentence.strip()) + + # Combine very short segments to avoid tiny audio files + min_chars = 20 # Minimum reasonable sentence length + combined_sentences = [] + i = 0 + + while i < len(parts): + current = parts[i] + + # If this is a short sentence and not the last one, combine with next + while i < len(parts) - 1 and len(current) < min_chars: + i += 1 + current += " " + parts[i] + + combined_sentences.append(current) + i += 1 + + return combined_sentences + +def generate_speech_from_api(prompt, voice=DEFAULT_VOICE, output_file=None, temperature=TEMPERATURE, + top_p=TOP_P, max_tokens=MAX_TOKENS, repetition_penalty=None, + use_batching=True, max_batch_chars=1000): + """Generate speech from text using Orpheus model with performance optimizations.""" + print(f"Starting speech generation for '{prompt[:50]}{'...' if len(prompt) > 50 else ''}'") + print(f"Using voice: {voice}, GPU acceleration: {'Yes (High-end)' if HIGH_END_GPU else 'Yes' if torch.cuda.is_available() else 'No'}") + + # Reset performance monitor + global perf_monitor + perf_monitor = PerformanceMonitor() + + start_time = time.time() + + # For shorter text, use the standard non-batched approach + if not use_batching or len(prompt) < max_batch_chars: + # Note: we ignore any provided repetition_penalty and always use the hardcoded value + # This ensures consistent quality regardless of what might be passed in + result = tokens_decoder_sync( + generate_tokens_from_api( + prompt=prompt, + voice=voice, + temperature=temperature, + top_p=top_p, + max_tokens=max_tokens, + repetition_penalty=REPETITION_PENALTY # Always use hardcoded value + ), + output_file=output_file + ) + + # Report final performance metrics + end_time = time.time() + total_time = end_time - start_time + print(f"Total speech generation completed in {total_time:.2f} seconds") + + return result + + # For longer text, use sentence-based batching + print(f"Using sentence-based batching for text with {len(prompt)} characters") + + # Split the text into sentences + sentences = split_text_into_sentences(prompt) + print(f"Split text into {len(sentences)} segments") + + # Create batches by combining sentences up to max_batch_chars + batches = [] + current_batch = "" + + for sentence in sentences: + # If adding this sentence would exceed the batch size, start a new batch + if len(current_batch) + len(sentence) > max_batch_chars and current_batch: + batches.append(current_batch) + current_batch = sentence + else: + # Add separator space if needed + if current_batch: + current_batch += " " + current_batch += sentence + + # Add the last batch if it's not empty + if current_batch: + batches.append(current_batch) + + print(f"Created {len(batches)} batches for processing") + + # Process each batch and collect audio segments + all_audio_segments = [] + batch_temp_files = [] + + for i, batch in enumerate(batches): + print(f"Processing batch {i+1}/{len(batches)} ({len(batch)} characters)") + + # Create a temporary file for this batch if an output file is requested + temp_output_file = None + if output_file: + temp_output_file = f"outputs/temp_batch_{i}_{int(time.time())}.wav" + batch_temp_files.append(temp_output_file) + + # Generate speech for this batch + batch_segments = tokens_decoder_sync( + generate_tokens_from_api( + prompt=batch, + voice=voice, + temperature=temperature, + top_p=top_p, + max_tokens=max_tokens, + repetition_penalty=REPETITION_PENALTY + ), + output_file=temp_output_file + ) + + # Add to our collection + all_audio_segments.extend(batch_segments) + + # If an output file was requested, stitch together the temporary files + if output_file and batch_temp_files: + # Stitch together WAV files + stitch_wav_files(batch_temp_files, output_file) + + # Clean up temporary files + for temp_file in batch_temp_files: + try: + os.remove(temp_file) + except Exception as e: + print(f"Warning: Could not remove temporary file {temp_file}: {e}") + + # Report final performance metrics + end_time = time.time() + total_time = end_time - start_time + + # Calculate combined duration + if all_audio_segments: + total_bytes = sum(len(segment) for segment in all_audio_segments) + duration = total_bytes / (2 * SAMPLE_RATE) # 2 bytes per sample at 24kHz + print(f"Generated {len(all_audio_segments)} audio segments") + print(f"Generated {duration:.2f} seconds of audio in {total_time:.2f} seconds") + print(f"Realtime factor: {duration/total_time:.2f}x") + + print(f"Total speech generation completed in {total_time:.2f} seconds") + + return all_audio_segments + +def stitch_wav_files(input_files, output_file, crossfade_ms=50): + """Stitch multiple WAV files together with crossfading for smooth transitions.""" + if not input_files: + return + + print(f"Stitching {len(input_files)} WAV files together with {crossfade_ms}ms crossfade") + + # If only one file, just copy it + if len(input_files) == 1: + import shutil + shutil.copy(input_files[0], output_file) + return + + # Convert crossfade_ms to samples + crossfade_samples = int(SAMPLE_RATE * crossfade_ms / 1000) + print(f"Using {crossfade_samples} samples for crossfade at {SAMPLE_RATE}Hz") + + # Build the final audio in memory with crossfades + final_audio = np.array([], dtype=np.int16) + first_params = None + + for i, input_file in enumerate(input_files): + try: + with wave.open(input_file, 'rb') as wav: + if first_params is None: + first_params = wav.getparams() + elif wav.getparams() != first_params: + print(f"Warning: WAV file {input_file} has different parameters") + + frames = wav.readframes(wav.getnframes()) + audio = np.frombuffer(frames, dtype=np.int16) + + if i == 0: + # First segment - use as is + final_audio = audio + else: + # Apply crossfade with previous segment + if len(final_audio) >= crossfade_samples and len(audio) >= crossfade_samples: + # Create crossfade weights + fade_out = np.linspace(1.0, 0.0, crossfade_samples) + fade_in = np.linspace(0.0, 1.0, crossfade_samples) + + # Apply crossfade + crossfade_region = (final_audio[-crossfade_samples:] * fade_out + + audio[:crossfade_samples] * fade_in).astype(np.int16) + + # Combine: original without last crossfade_samples + crossfade + new without first crossfade_samples + final_audio = np.concatenate([final_audio[:-crossfade_samples], + crossfade_region, + audio[crossfade_samples:]]) + else: + # One segment too short for crossfade, just append + print(f"Segment {i} too short for crossfade, concatenating directly") + final_audio = np.concatenate([final_audio, audio]) + except Exception as e: + print(f"Error processing file {input_file}: {e}") + if i == 0: + raise # Critical failure if first file fails + + # Write the final audio data to the output file + try: + with wave.open(output_file, 'wb') as output_wav: + if first_params is None: + raise ValueError("No valid WAV files were processed") + + output_wav.setparams(first_params) + output_wav.writeframes(final_audio.tobytes()) + + print(f"Successfully stitched audio to {output_file} with crossfading") + except Exception as e: + print(f"Error writing output file {output_file}: {e}") + raise + +def list_available_voices(): + """List all available voices with the recommended one marked.""" + print("Available voices (in order of conversational realism):") + for i, voice in enumerate(AVAILABLE_VOICES): + marker = "★" if voice == DEFAULT_VOICE else " " + print(f"{marker} {voice}") + print(f"\nDefault voice: {DEFAULT_VOICE}") + + print("\nAvailable emotion tags:") + print(", , , , , , , ") + +def main(): + # Parse command line arguments + parser = argparse.ArgumentParser(description="Orpheus Text-to-Speech using Orpheus-FASTAPI") + parser.add_argument("--text", type=str, help="Text to convert to speech") + parser.add_argument("--voice", type=str, default=DEFAULT_VOICE, help=f"Voice to use (default: {DEFAULT_VOICE})") + parser.add_argument("--output", type=str, help="Output WAV file path") + parser.add_argument("--list-voices", action="store_true", help="List available voices") + parser.add_argument("--temperature", type=float, default=TEMPERATURE, help="Temperature for generation") + parser.add_argument("--top_p", type=float, default=TOP_P, help="Top-p sampling parameter") + parser.add_argument("--repetition_penalty", type=float, default=REPETITION_PENALTY, + help="Repetition penalty (fixed at 1.1 for stable generation - parameter kept for compatibility)") + + args = parser.parse_args() + + if args.list_voices: + list_available_voices() + return + + # Use text from command line or prompt user + prompt = args.text + if not prompt: + if len(sys.argv) > 1 and sys.argv[1] not in ("--voice", "--output", "--temperature", "--top_p", "--repetition_penalty"): + prompt = " ".join([arg for arg in sys.argv[1:] if not arg.startswith("--")]) + else: + prompt = input("Enter text to synthesize: ") + if not prompt: + prompt = "Hello, I am Orpheus, an AI assistant with emotional speech capabilities." + + # Default output file if none provided + output_file = args.output + if not output_file: + # Create outputs directory if it doesn't exist + os.makedirs("outputs", exist_ok=True) + # Generate a filename based on the voice and a timestamp + timestamp = time.strftime("%Y%m%d_%H%M%S") + output_file = f"outputs/{args.voice}_{timestamp}.wav" + print(f"No output file specified. Saving to {output_file}") + + # Generate speech + start_time = time.time() + audio_segments = generate_speech_from_api( + prompt=prompt, + voice=args.voice, + temperature=args.temperature, + top_p=args.top_p, + repetition_penalty=args.repetition_penalty, + output_file=output_file + ) + end_time = time.time() + + print(f"Speech generation completed in {end_time - start_time:.2f} seconds") + print(f"Audio saved to {output_file}") + +if __name__ == "__main__": + main() diff --git a/tts_engine/inference.py b/tts_engine/inference.py index b11817c..f2ce65d 100644 --- a/tts_engine/inference.py +++ b/tts_engine/inference.py @@ -10,9 +10,39 @@ import argparse import threading import queue import asyncio -from concurrent.futures import ThreadPoolExecutor -from typing import List, Dict, Any, Optional, Generator, Union, Tuple +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import List, Dict, Any, Optional, Generator, Union, Tuple, AsyncGenerator from dotenv import load_dotenv +import nltk +import multiprocessing +import logging + +logger = logging.getLogger(__name__) + +# Helper to ensure NLTK data is downloaded +_nltk_punkt_downloaded = False +def ensure_nltk_punkt(): + global _nltk_punkt_downloaded + if not _nltk_punkt_downloaded: + try: + nltk.data.find('tokenizers/punkt') + except nltk.downloader.DownloadError: + print("NLTK 'punkt' tokenizer model not found. Downloading...") + try: + nltk.download('punkt', quiet=True) + print("'punkt' model downloaded successfully.") + except Exception as e: + print(f"Error downloading NLTK 'punkt' model: {e}") + print("Sentence tokenization might be less accurate.") + except LookupError: # Handle case where nltk_data path isn't configured + print("NLTK 'punkt' model not found. Downloading...") + try: + nltk.download('punkt', quiet=True) + print("'punkt' model downloaded successfully.") + except Exception as e: + print(f"Error downloading NLTK 'punkt' model: {e}") + print("Sentence tokenization might be less accurate.") + _nltk_punkt_downloaded = True # Helper to detect if running in Uvicorn's reloader def is_reloader_process(): @@ -32,9 +62,54 @@ load_dotenv() import torch import psutil +# Device selection with support for Apple Silicon MPS +# Define device globally for consistent use throughout +DEVICE = "cpu" # Default to CPU, will be updated based on availability + # Detect if we're on a high-end system based on hardware capabilities HIGH_END_GPU = False -if torch.cuda.is_available(): +APPLE_SILICON = False + +# Check for Apple Silicon MPS support first +if torch.backends.mps.is_available(): + DEVICE = "mps" + APPLE_SILICON = True + + # Get Apple Silicon details + import platform + import subprocess + + # Get chip model and memory + chip_model = platform.processor() + try: + # Get memory info using sysctl + mem_cmd = subprocess.run(["sysctl", "hw.memsize"], capture_output=True, text=True) + if mem_cmd.returncode == 0: + mem_bytes = int(mem_cmd.stdout.split(':')[1].strip()) + mem_gb = mem_bytes / (1024**3) + else: + mem_gb = psutil.virtual_memory().total / (1024**3) + except Exception: + mem_gb = psutil.virtual_memory().total / (1024**3) + + # Detect high-end Apple Silicon (M1 Pro/Max/Ultra, M2 Pro/Max/Ultra, M3 Pro/Max/Ultra) + if "Pro" in chip_model or "Max" in chip_model or "Ultra" in chip_model or mem_gb >= 32: + HIGH_END_GPU = True + if not IS_RELOADER: + print(f"🍎 Hardware: High-end Apple Silicon detected") + print(f"📊 Chip: {chip_model}") + print(f"📊 RAM: {mem_gb:.2f} GB unified memory") + print("🚀 Using high-performance Apple Silicon optimizations") + else: + if not IS_RELOADER: + print(f"🍎 Hardware: Apple Silicon detected") + print(f"📊 Chip: {chip_model}") + print(f"📊 RAM: {mem_gb:.2f} GB unified memory") + print("🚀 Using Apple Silicon optimizations") + +# Then check for CUDA GPU +elif torch.cuda.is_available(): + DEVICE = "cuda" # Get GPU properties props = torch.cuda.get_device_properties(0) gpu_name = props.name @@ -67,7 +142,7 @@ else: ram_gb = psutil.virtual_memory().total / (1024**3) if not IS_RELOADER: - print(f"🖥️ Hardware: CPU only (No CUDA GPU detected)") + print(f"🖥️ Hardware: CPU only (No GPU acceleration detected)") print(f"📊 CPU: {cpu_cores} cores, {cpu_threads} threads") print(f"📊 RAM: {ram_gb:.2f} GB") print("⚙️ Using CPU-optimized settings") @@ -134,10 +209,78 @@ if not IS_RELOADER: print(f" REPETITION_PENALTY: {REPETITION_PENALTY}") # Parallel processing settings -NUM_WORKERS = 4 if HIGH_END_GPU else 2 +import multiprocessing + +# Determine optimal settings based on hardware +CPU_CORES = multiprocessing.cpu_count() + +# For Apple Silicon, use more aggressive settings depending on the model +if APPLE_SILICON: + # Optimize for Apple Silicon based on memory + ram_gb = psutil.virtual_memory().total / (1024**3) + if ram_gb >= 64: # High-memory M1 Max/Ultra, M2 Max/Ultra, M3 Max/Ultra (64GB+) + NUM_WORKERS = max(8, min(CPU_CORES-2, 12)) + BATCH_SIZE = 64 + AUDIO_QUEUE_SIZE = 200 + elif ram_gb >= 32: # Mid-range models (32GB) + NUM_WORKERS = max(4, min(CPU_CORES-2, 8)) + BATCH_SIZE = 48 + AUDIO_QUEUE_SIZE = 150 + else: # Base models + NUM_WORKERS = max(2, min(CPU_CORES-1, 4)) + BATCH_SIZE = 32 + AUDIO_QUEUE_SIZE = 100 +elif HIGH_END_GPU: # High-end CUDA GPU + NUM_WORKERS = 4 + BATCH_SIZE = 32 + AUDIO_QUEUE_SIZE = 100 +else: # Regular CUDA or CPU + NUM_WORKERS = 2 + BATCH_SIZE = 16 + AUDIO_QUEUE_SIZE = 50 + +# Buffer size for audio processing +if APPLE_SILICON and psutil.virtual_memory().total >= (64 * 1024**3): # 64GB+ RAM + BUFFER_SIZE_MB = 4 # 4MB buffer +elif APPLE_SILICON or HIGH_END_GPU: + BUFFER_SIZE_MB = 2 # 2MB buffer +else: + BUFFER_SIZE_MB = 1 # 1MB buffer + +BUFFER_MAX_SIZE = BUFFER_SIZE_MB * 1024 * 1024 + +# Maximum number of characters per batch for NLTK sentence splitting +try: + MAX_BATCH_CHARS = int(os.environ.get("ORPHEUS_MAX_BATCH_CHARS", "600")) + if MAX_BATCH_CHARS < 100 or MAX_BATCH_CHARS > 2000: + print(f"WARNING: Invalid ORPHEUS_MAX_BATCH_CHARS value ({MAX_BATCH_CHARS}), should be between 100-2000. Using 600 as fallback.") + MAX_BATCH_CHARS = 600 +except (ValueError, TypeError): + print("WARNING: Invalid ORPHEUS_MAX_BATCH_CHARS value, using 600 as fallback") + MAX_BATCH_CHARS = 600 + +# Crossfade duration in milliseconds for stitching audio batches +try: + CROSSFADE_MS = int(os.environ.get("ORPHEUS_CROSSFADE_MS", "30")) + if CROSSFADE_MS < 10 or CROSSFADE_MS > 200: + print(f"WARNING: Invalid ORPHEUS_CROSSFADE_MS value ({CROSSFADE_MS}), should be between 10-200. Using 30 as fallback.") + CROSSFADE_MS = 30 +except (ValueError, TypeError): + print("WARNING: Invalid ORPHEUS_CROSSFADE_MS value, using 30 as fallback") + CROSSFADE_MS = 30 + +# Helper function to generate equal power fade curves using sine/cosine +def generate_equal_power_fade_curves(num_samples): + """Generate fade-out and fade-in curves using sine/cosine for equal power crossfading.""" + # Create a linear ramp from 0 to pi/2 + ramp = np.linspace(0, np.pi/2, num_samples) + # Use sine for fade-out and cosine for fade-in to maintain equal power + fade_out = np.sin(ramp) + fade_in = np.cos(ramp) + return fade_out, fade_in # Define voices by language -ENGLISH_VOICES = ["tara", "leah", "jess", "leo", "dan", "mia", "zac", "zoe"] +ENGLISH_VOICES = ["tara", "kaya", "leah", "jess", "leo", "dan", "mia", "zac", "zoe"] FRENCH_VOICES = ["pierre", "amelie", "marie"] GERMAN_VOICES = ["jana", "thomas", "max"] KOREAN_VOICES = ["유나", "준서"] @@ -215,56 +358,93 @@ class PerformanceMonitor: # Estimate audio duration based on audio chunks (each chunk is ~0.085s of audio) est_duration = self.audio_chunks * 0.085 - print(f"Progress: {tokens_per_sec:.1f} tokens/sec, est. {est_duration:.1f}s audio generated, {self.token_count} tokens, {self.audio_chunks} chunks in {elapsed:.1f}s") + # print(f"Progress: {tokens_per_sec:.1f} tokens/sec, est. {est_duration:.1f}s audio generated, {self.token_count} tokens, {self.audio_chunks} chunks in {elapsed:.1f}s") # Create global performance monitor perf_monitor = PerformanceMonitor() +# --- Read System Prompt --- +SYSTEM_PROMPT_CONTENT = None +PAST_CONTEXT = None # For long texts, we are going to store the past context in this variable and pass it to the model as a system prompt. +SYSTEM_PROMPT_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "System_Prompt.md") + +try: + if os.path.exists(SYSTEM_PROMPT_PATH): + with open(SYSTEM_PROMPT_PATH, 'r', encoding='utf-8') as f: + SYSTEM_PROMPT_CONTENT = f.read() + if not IS_RELOADER: + logger.info(f"Successfully loaded system prompt from {SYSTEM_PROMPT_PATH}") + else: + if not IS_RELOADER: + logger.warning(f"System prompt file not found at {SYSTEM_PROMPT_PATH}. Proceeding without system prompt.") +except Exception as e: + if not IS_RELOADER: + logger.error(f"Error reading system prompt file {SYSTEM_PROMPT_PATH}: {e}") +# --- End System Prompt Read --- + def format_prompt(prompt: str, voice: str = DEFAULT_VOICE) -> str: - """Format prompt for Orpheus model with voice prefix and special tokens.""" + """Format prompt for Orpheus model with voice prefix and special tokens. + Restored based on backup file. + """ # Validate voice and provide fallback if voice not in AVAILABLE_VOICES: - print(f"Warning: Voice '{voice}' not recognized. Using '{DEFAULT_VOICE}' instead.") + logger.warning(f"Voice '{voice}' not recognized. Using '{DEFAULT_VOICE}' instead.") voice = DEFAULT_VOICE - # Format similar to how engine_class.py does it with special tokens + # Format similar to original backup formatted_prompt = f"{voice}: {prompt}" - # Add special token markers for the Orpheus-FASTAPI - special_start = "<|audio|>" # Using the additional_special_token from config - special_end = "<|eot_id|>" # Using the eos_token from config + # Add special token markers required by Orpheus + special_start = "<|audio|>" + special_end = "<|eot_id|>" - return f"{special_start}{formatted_prompt}{special_end}" + return f"{special_start}{formatted_prompt}{special_end}" # Restored original format def generate_tokens_from_api(prompt: str, voice: str = DEFAULT_VOICE, temperature: float = TEMPERATURE, top_p: float = TOP_P, max_tokens: int = MAX_TOKENS, - repetition_penalty: float = REPETITION_PENALTY) -> Generator[str, None, None]: - """Generate tokens from text using OpenAI-compatible API with optimized streaming and retry logic.""" + repetition_penalty: float = REPETITION_PENALTY, + context_prompt: Optional[str] = None) -> Generator[str, None, None]: + """Generate tokens from text using OpenAI-compatible API (prompt format). + Uses the separate 'system_prompt' payload field for system/context info if available. + """ start_time = time.time() - formatted_prompt = format_prompt(prompt, voice) - print(f"Generating speech for: {formatted_prompt}") - - # Optimize the token generation for GPUs - if HIGH_END_GPU: - # Use more aggressive parameters for faster generation on high-end GPUs - print("Using optimized parameters for high-end GPU") - elif torch.cuda.is_available(): - print("Using optimized parameters for GPU acceleration") - - # Create the request payload (model field may not be required by some endpoints but included for compatibility) + # Format the main part of the user prompt (e.g., <|audio|>voice: text<|eot_id|>) + formatted_user_prompt = format_prompt(prompt, voice) + print(f"Generating speech for user prompt (voice: {voice}): {prompt[:80]}...") + + # Combine system prompt and context prompt for the dedicated field + system_and_context = "" + if SYSTEM_PROMPT_CONTENT: + system_and_context += SYSTEM_PROMPT_CONTENT + print("Using system prompt.") + if context_prompt: + # Add clear separation if both system and context exist + if system_and_context: + system_and_context += "\n\n[Previous context:]\n" + else: # Only context exists + system_and_context += "[Previous context:]\n" + system_and_context += context_prompt + print("Using context prompt.") + + # Create the request payload payload = { - "prompt": formatted_prompt, + "prompt": formatted_user_prompt, # Main user prompt here "max_tokens": max_tokens, "temperature": temperature, "top_p": top_p, "repeat_penalty": repetition_penalty, - "stream": True # Always stream for better performance + "stream": True } - # Add model field - this is ignored by many local inference servers for /v1/completions - # but included for compatibility with OpenAI API and some servers that may use it + # Add the dedicated system_prompt field if we have content for it + if system_and_context: + payload["system"] = system_and_context + print(f"Sending system_prompt field: {system_and_context[:100]}...") + + # Add model field - optional but good practice model_name = os.environ.get("ORPHEUS_MODEL_NAME", "Orpheus-3b-FT-Q8_0.gguf") - payload["model"] = model_name + if model_name: + payload["model"] = model_name # Session for connection pooling and retry logic session = requests.Session() @@ -274,7 +454,7 @@ def generate_tokens_from_api(prompt: str, voice: str = DEFAULT_VOICE, temperatur while retry_count < max_retries: try: - # Make the API request with streaming and timeout + # Make the API request response = session.post( API_URL, headers=HEADERS, @@ -284,72 +464,66 @@ def generate_tokens_from_api(prompt: str, voice: str = DEFAULT_VOICE, temperatur ) if response.status_code != 200: - print(f"Error: API request failed with status code {response.status_code}") - print(f"Error details: {response.text}") - # Retry on server errors (5xx) but not on client errors (4xx) + error_detail = response.text + logger.error(f"API Error ({response.status_code}): {error_detail}") + # Removed the messages vs prompt check as we now use prompt if response.status_code >= 500: retry_count += 1 - wait_time = 2 ** retry_count # Exponential backoff - print(f"Retrying in {wait_time} seconds...") + wait_time = 2 ** retry_count + logger.info(f"Retrying in {wait_time} seconds... (attempt {retry_count+1}/{max_retries})") time.sleep(wait_time) continue return - # Process the streamed response with better buffering - buffer = "" + # Process the streamed response (parsing logic remains the same, checking for 'text') token_counter = 0 - - # Iterate through the response to get tokens for line in response.iter_lines(): if line: line_str = line.decode('utf-8') if line_str.startswith('data: '): - data_str = line_str[6:] # Remove the 'data: ' prefix - + data_str = line_str[6:] if data_str.strip() == '[DONE]': break - try: data = json.loads(data_str) if 'choices' in data and len(data['choices']) > 0: - token_chunk = data['choices'][0].get('text', '') - for token_text in token_chunk.split('>'): - token_text = f'{token_text}>' + choice = data['choices'][0] + token_chunk = choice.get('text', '') # Expecting 'text' for /v1/completions + + if token_chunk and token_chunk.startswith(""): token_counter += 1 perf_monitor.add_tokens() + yield token_chunk - if token_text: - yield token_text except json.JSONDecodeError as e: - print(f"Error decoding JSON: {e}") + logger.error(f"Error decoding JSON stream data: {e}") continue - # Generation completed successfully generation_time = time.time() - start_time tokens_per_second = token_counter / generation_time if generation_time > 0 else 0 print(f"Token generation complete: {token_counter} tokens in {generation_time:.2f}s ({tokens_per_second:.1f} tokens/sec)") return except requests.exceptions.Timeout: - print(f"Request timed out after {REQUEST_TIMEOUT} seconds") + logger.warning(f"Request timed out after {REQUEST_TIMEOUT} seconds") retry_count += 1 if retry_count < max_retries: wait_time = 2 ** retry_count - print(f"Retrying in {wait_time} seconds... (attempt {retry_count+1}/{max_retries})") + logger.info(f"Retrying in {wait_time} seconds... (attempt {retry_count+1}/{max_retries})") time.sleep(wait_time) else: - print("Max retries reached. Token generation failed.") + logger.error("Max retries reached for timeout. Token generation failed.") return - except requests.exceptions.ConnectionError: - print(f"Connection error to API at {API_URL}") + except requests.exceptions.RequestException as e: + logger.error(f"Request error connecting to API at {API_URL}: {e}") retry_count += 1 if retry_count < max_retries: wait_time = 2 ** retry_count - print(f"Retrying in {wait_time} seconds... (attempt {retry_count+1}/{max_retries})") + logger.info(f"Retrying in {wait_time} seconds... (attempt {retry_count+1}/{max_retries})") time.sleep(wait_time) else: - print("Max retries reached. Token generation failed.") + logger.error("Max retries reached for connection error. Token generation failed.") return # The turn_token_into_id function is now imported from speechpipe.py @@ -367,7 +541,7 @@ def convert_to_audio(multiframe: List[int], count: int) -> Optional[bytes]: return result -async def tokens_decoder(token_gen) -> Generator[bytes, None, None]: +async def tokens_decoder(token_gen) -> AsyncGenerator[bytes, None]: """Simplified token decoder with early first-chunk processing for lower latency.""" buffer = [] count = 0 @@ -391,12 +565,12 @@ async def tokens_decoder(token_gen) -> Generator[bytes, None, None]: token_count += 1 # Log throughput periodically - current_time = time.time() - if current_time - last_log_time > 5.0: # Every 5 seconds - elapsed = current_time - start_time - if elapsed > 0: - print(f"Token processing rate: {token_count/elapsed:.1f} tokens/second") - last_log_time = current_time + # current_time = time.time() + # if current_time - last_log_time > 5.0: # Every 5 seconds + # elapsed = current_time - start_time + # if elapsed > 0: + # print(f"Token processing rate: {token_count/elapsed:.1f} tokens/second") + # last_log_time = current_time # Different processing paths based on whether first chunk has been processed if not first_chunk_processed: @@ -417,8 +591,8 @@ async def tokens_decoder(token_gen) -> Generator[bytes, None, None]: buffer_to_proc = buffer[-min_frames_subsequent:] # Debug output to help diagnose issues - if count % 28 == 0: - print(f"Processing buffer with {len(buffer_to_proc)} tokens, total collected: {len(buffer)}") + # if count % 28 == 0: + # print(f"Processing buffer with {len(buffer_to_proc)} tokens, total collected: {len(buffer)}") # Process the tokens audio_samples = convert_to_audio(buffer_to_proc, count) @@ -427,8 +601,8 @@ async def tokens_decoder(token_gen) -> Generator[bytes, None, None]: def tokens_decoder_sync(syn_token_gen, output_file=None): """Optimized synchronous wrapper with parallel processing and efficient file I/O.""" - # Use a larger queue for high-end systems - queue_size = 100 if HIGH_END_GPU else 50 + # Use hardware-optimized queue size + queue_size = AUDIO_QUEUE_SIZE audio_queue = queue.Queue(maxsize=queue_size) audio_segments = [] @@ -442,8 +616,8 @@ def tokens_decoder_sync(syn_token_gen, output_file=None): wav_file.setsampwidth(2) wav_file.setframerate(SAMPLE_RATE) - # Batch processing of tokens for improved throughput - batch_size = 32 if HIGH_END_GPU else 16 + # Use optimized batch processing + batch_size = BATCH_SIZE # Thread synchronization for proper completion detection producer_done_event = threading.Event() @@ -479,16 +653,16 @@ def tokens_decoder_sync(syn_token_gen, output_file=None): chunk_count += 1 # Log performance periodically - current_time = time.time() - if current_time - last_log_time >= 3.0: # Every 3 seconds - elapsed = current_time - last_log_time - if elapsed > 0: - recent_chunks = chunk_count - chunks_per_sec = recent_chunks / elapsed - print(f"Audio generation rate: {chunks_per_sec:.2f} chunks/second") - last_log_time = current_time - # Reset chunk counter for next interval - chunk_count = 0 + # current_time = time.time() + # if current_time - last_log_time >= 3.0: # Every 3 seconds + # elapsed = current_time - last_log_time + # if elapsed > 0: + # recent_chunks = chunk_count + # chunks_per_sec = recent_chunks / elapsed + # print(f"Audio generation rate: {chunks_per_sec:.2f} chunks/second") + # last_log_time = current_time + # # Reset chunk counter for next interval + # chunk_count = 0 except Exception as e: print(f"Error in token processing: {str(e)}") import traceback @@ -514,10 +688,9 @@ def tokens_decoder_sync(syn_token_gen, output_file=None): # before the producer has had a chance to add anything producer_started_event.wait(timeout=5.0) - # Optimized I/O approach for all systems - # This approach is simpler and more reliable than separate code paths + # Use hardware-optimized buffer size write_buffer = bytearray() - buffer_max_size = 1024 * 1024 # 1MB max buffer size (adjustable) + buffer_max_size = BUFFER_MAX_SIZE # Keep track of the last time we checked for completion last_check_time = time.time() @@ -624,201 +797,512 @@ import numpy as np from io import BytesIO import wave -def split_text_into_sentences(text): - """Split text into sentences with a more reliable approach.""" - # We'll use a simple approach that doesn't rely on variable-width lookbehinds - # which aren't supported in Python's regex engine +# Map for NLTK language names (add more as needed/supported by punkt) +NLTK_LANG_MAP = { + "english": "english", + "french": "french", + "german": "german", + "spanish": "spanish", + "italian": "italian", + # Add other languages if punkt models exist for them +} + +def split_text_into_sentences(text: str, language: str = "english", max_chars_per_segment: int = MAX_BATCH_CHARS) -> List[str]: + """Split text into sentences using NLTK for better accuracy.""" + ensure_nltk_punkt() # Make sure model is available + + nltk_lang = NLTK_LANG_MAP.get(language, "english") # Default to English if mapped lang not found + print(f"Splitting text into sentences using NLTK for language: {nltk_lang}") - # First, split on common sentence ending punctuation - # This isn't perfect but works for most cases and avoids the regex error - parts = [] - current_sentence = "" + try: + # Attempt to use the specified language model + sentences = nltk.sent_tokenize(text, language=nltk_lang) + print(f"Successfully tokenized using NLTK for language: {nltk_lang}") + except Exception as e: + # Fallback to default English model if specific language model fails or isn't available + print(f"Warning: Could not use NLTK tokenizer for language '{nltk_lang}'. Falling back to English. Error: {e}") + try: + sentences = nltk.sent_tokenize(text) + except Exception as inner_e: + # If NLTK fails completely, fallback to a very basic split as a last resort + print(f"ERROR: NLTK sentence tokenization failed entirely: {inner_e}. Using basic fallback.") + sentences = text.split('. ') # Very basic fallback + + # Filter out empty sentences and strip whitespace + sentences = [s.strip() for s in sentences if s.strip()] - for char in text: - current_sentence += char - - # If we hit a sentence ending followed by a space, consider this a potential sentence end - if char in (' ', '\n', '\t') and len(current_sentence) > 1: - prev_char = current_sentence[-2] - if prev_char in ('.', '!', '?'): - # Check if this is likely a real sentence end and not an abbreviation - # (Simple heuristic: if there's a space before the period, it's likely a real sentence end) - if len(current_sentence) > 3 and current_sentence[-3] not in ('.', ' '): - parts.append(current_sentence.strip()) - current_sentence = "" + # Create segments that respect the max_chars_per_segment limit + segments = [] + current_segment = "" - # Add any remaining text - if current_sentence.strip(): - parts.append(current_sentence.strip()) - - # Combine very short segments to avoid tiny audio files - min_chars = 20 # Minimum reasonable sentence length - combined_sentences = [] - i = 0 - - while i < len(parts): - current = parts[i] - - # If this is a short sentence and not the last one, combine with next - while i < len(parts) - 1 and len(current) < min_chars: - i += 1 - current += " " + parts[i] + for sentence in sentences: + # If a single sentence exceeds the limit, we need to split it + if len(sentence) > max_chars_per_segment: + # If we have a current segment, add it first + if current_segment: + segments.append(current_segment) + current_segment = "" - combined_sentences.append(current) - i += 1 + # Split the long sentence into chunks + words = sentence.split() + current_chunk = "" + + for word in words: + # If adding this word would exceed the limit, start a new chunk + if len(current_chunk) + len(word) + 1 > max_chars_per_segment: + if current_chunk: + segments.append(current_chunk) + current_chunk = word + else: + # Add word to current chunk with a space if needed + current_chunk = f"{current_chunk} {word}" if current_chunk else word + + # Add the last chunk if it exists + if current_chunk: + segments.append(current_chunk) + else: + # Check if adding this sentence would exceed the limit + if current_segment: + potential_length = len(current_segment) + 1 + len(sentence) + if potential_length > max_chars_per_segment: + segments.append(current_segment) + current_segment = sentence + else: + current_segment = f"{current_segment} {sentence}" + else: + current_segment = sentence - return combined_sentences + # Add the last segment if it exists + if current_segment: + segments.append(current_segment) + + # Verify that no segment exceeds the limit + for i, segment in enumerate(segments): + if len(segment) > max_chars_per_segment: + print(f"Warning: Segment {i} exceeds max_chars_per_segment ({len(segment)} > {max_chars_per_segment})") + # Split the segment into smaller chunks + words = segment.split() + new_segments = [] + current_chunk = "" + + for word in words: + if len(current_chunk) + len(word) + 1 > max_chars_per_segment: + if current_chunk: + new_segments.append(current_chunk) + current_chunk = word + else: + current_chunk = f"{current_chunk} {word}" if current_chunk else word + + if current_chunk: + new_segments.append(current_chunk) + + # Replace the long segment with the new segments + segments[i:i+1] = new_segments + + print(f"Split text into {len(segments)} segments with max length {max_chars_per_segment}") + return segments + +def cleanup_between_batches(): + """Reset all state between batch processing.""" + # Reset the performance monitor + global perf_monitor + perf_monitor = PerformanceMonitor() + + # Import here to avoid circular imports + from .speechpipe import reset_state + # Reset the speechpipe state + reset_state() + + # Force garbage collection + import gc + gc.collect() def generate_speech_from_api(prompt, voice=DEFAULT_VOICE, output_file=None, temperature=TEMPERATURE, top_p=TOP_P, max_tokens=MAX_TOKENS, repetition_penalty=None, - use_batching=True, max_batch_chars=1000): - """Generate speech from text using Orpheus model with performance optimizations.""" + use_batching=True, max_batch_chars=MAX_BATCH_CHARS): + """Generate speech from text using Orpheus model with performance optimizations and NLTK splitting.""" print(f"Starting speech generation for '{prompt[:50]}{'...' if len(prompt) > 50 else ''}'") - print(f"Using voice: {voice}, GPU acceleration: {'Yes (High-end)' if HIGH_END_GPU else 'Yes' if torch.cuda.is_available() else 'No'}") + print(f"Using voice: {voice}, GPU acceleration: {'Yes (High-end)' if HIGH_END_GPU else 'Yes' if torch.cuda.is_available() or torch.backends.mps.is_available() else 'No'}") - # Reset performance monitor + # Reset performance monitor at start global perf_monitor perf_monitor = PerformanceMonitor() start_time = time.time() - # For shorter text, use the standard non-batched approach + all_audio_segments = [] # To store the final small chunks for return/streaming + + # Determine language for splitting + generation_language = VOICE_TO_LANGUAGE.get(voice, "english") + + # For shorter text or disabled batching, use the standard non-batched approach if not use_batching or len(prompt) < max_batch_chars: - # Note: we ignore any provided repetition_penalty and always use the hardcoded value - # This ensures consistent quality regardless of what might be passed in - result = tokens_decoder_sync( + print("Processing text as a single batch.") + # Note: repetition_penalty is ignored (uses hardcoded 1.1) + all_audio_segments = tokens_decoder_sync( generate_tokens_from_api( prompt=prompt, voice=voice, temperature=temperature, top_p=top_p, max_tokens=max_tokens, - repetition_penalty=REPETITION_PENALTY # Always use hardcoded value + repetition_penalty=REPETITION_PENALTY # Fixed value ), - output_file=output_file + output_file=output_file # Pass file handle if saving ) - # Report final performance metrics - end_time = time.time() - total_time = end_time - start_time - print(f"Total speech generation completed in {total_time:.2f} seconds") - - return result - - # For longer text, use sentence-based batching - print(f"Using sentence-based batching for text with {len(prompt)} characters") - - # Split the text into sentences - sentences = split_text_into_sentences(prompt) - print(f"Split text into {len(sentences)} segments") - - # Create batches by combining sentences up to max_batch_chars - batches = [] - current_batch = "" - - for sentence in sentences: - # If adding this sentence would exceed the batch size, start a new batch - if len(current_batch) + len(sentence) > max_batch_chars and current_batch: - batches.append(current_batch) - current_batch = sentence - else: - # Add separator space if needed - if current_batch: - current_batch += " " - current_batch += sentence - - # Add the last batch if it's not empty - if current_batch: - batches.append(current_batch) - - print(f"Created {len(batches)} batches for processing") - - # Process each batch and collect audio segments - all_audio_segments = [] - batch_temp_files = [] - - for i, batch in enumerate(batches): - print(f"Processing batch {i+1}/{len(batches)} ({len(batch)} characters)") - - # Create a temporary file for this batch if an output file is requested - temp_output_file = None if output_file: - temp_output_file = f"outputs/temp_batch_{i}_{int(time.time())}.wav" - batch_temp_files.append(temp_output_file) + pass # File writing is handled internally by tokens_decoder_sync + else: + pass # Segments are already in all_audio_segments + + # For longer text, use sentence-based batching with NLTK and crossfading + else: + print(f"Using sentence-based batching for text with {len(prompt)} characters (limit: {max_batch_chars})") - # Generate speech for this batch - batch_segments = tokens_decoder_sync( - generate_tokens_from_api( - prompt=batch, - voice=voice, - temperature=temperature, - top_p=top_p, - max_tokens=max_tokens, - repetition_penalty=REPETITION_PENALTY - ), - output_file=temp_output_file - ) + # Split the text into segments using NLTK, passing the limit + segments = split_text_into_sentences(prompt, language=generation_language, max_chars_per_segment=max_batch_chars) + print(f"Split text into {len(segments)} segments using NLTK ({generation_language}).") - # Add to our collection - all_audio_segments.extend(batch_segments) - - # If an output file was requested, stitch together the temporary files - if output_file and batch_temp_files: - # Stitch together WAV files - stitch_wav_files(batch_temp_files, output_file) + # Process each segment and collect audio segments + all_complete_batch_audio = [] # Store complete audio (bytes) for each batch + batch_temp_files = [] + all_results = [(None, None, None, None)] * len(segments) # Pre-allocate results list - # Clean up temporary files - for temp_file in batch_temp_files: - try: - os.remove(temp_file) - except Exception as e: - print(f"Warning: Could not remove temporary file {temp_file}: {e}") - - # Report final performance metrics + # Determine if we should use parallel processing + use_parallel = APPLE_SILICON and len(segments) > 1 and psutil.virtual_memory().total >= (32 * 1024**3) + + if use_parallel: + print(f"Using parallel processing with {NUM_WORKERS} workers for {len(segments)} segments") + + # Define a function to process a single segment + def process_segment(segment_text, index, previous_segment_text): + print(f"Starting parallel processing of segment {index+1}/{len(segments)} ({len(segment_text)} characters)") + + # Clean up state before processing + cleanup_between_batches() + + # --- Prepare context from previous segment --- + context_for_next = "" + if previous_segment_text: + try: + # Extract last sentence as context + # Make sure NLTK data is available (should be loaded globally) + ensure_nltk_punkt() + prev_sentences = nltk.sent_tokenize(previous_segment_text) + if prev_sentences: + # Use only the last sentence as context + context_for_next = prev_sentences[-1].strip() + print(f"Segment {index+1}: Using last sentence of previous segment as context.") + except Exception as e: + logger.warning(f"Segment {index+1}: Could not extract context from previous segment: {e}") + # --- End context preparation --- + + # Create a temporary file if needed + temp_file = None + if output_file: + temp_file = f"outputs/temp_batch_{index}_{int(time.time())}.wav" + # Note: Appending to batch_temp_files needs thread-safety or post-processing + + # Process the segment, passing the context + segments_output = tokens_decoder_sync( + generate_tokens_from_api( + prompt=segment_text, + voice=voice, + temperature=temperature, + top_p=top_p, + max_tokens=max_tokens, + repetition_penalty=REPETITION_PENALTY, + context_prompt=context_for_next # Pass context here + ), + output_file=temp_file + ) + + # Combine the audio fragments + complete_audio = b"".join(segments_output) + print(f"Completed parallel processing of segment {index+1}") + + return index, complete_audio, segments_output, temp_file + + # Use ThreadPoolExecutor for parallel processing + with ThreadPoolExecutor(max_workers=NUM_WORKERS) as executor: + # Submit all segments for processing, passing previous segment text + future_to_segment = { + executor.submit(process_segment, + segment, + i, + segments[i-1] if i > 0 else None): # Pass previous segment + i for i, segment in enumerate(segments) + } + + # Process segments as they complete + temp_batch_files_collector = {} # Collect temp files safely + for future in as_completed(future_to_segment): + try: + idx, complete_audio, segments_output, temp_file = future.result() + # Store the results in the pre-allocated list using the index + all_results[idx] = (idx, complete_audio, segments_output, temp_file) + if temp_file: + temp_batch_files_collector[idx] = temp_file + except Exception as e: + idx = future_to_segment[future] # Get index from map + print(f"Error processing segment {idx+1}: {e}") + # Store error indication if needed, e.g., all_results[idx] = (idx, None, None, None) + + # Extract results in original order from the pre-allocated list + all_complete_batch_audio = [audio for idx, audio, _, _ in all_results if audio is not None] + batch_temp_files = [temp_batch_files_collector[i] for i in sorted(temp_batch_files_collector) if temp_batch_files_collector[i] is not None] + + if not output_file: + # Flatten the segments in original order + temp_segments = [] + for idx, _, segments_output, _ in all_results: + if segments_output is not None: + temp_segments.extend(segments_output) + all_audio_segments = temp_segments + else: + # Process segments sequentially (original logic, but add context passing) + for i, segment in enumerate(segments): + print(f"Processing segment {i+1}/{len(segments)} ({len(segment)} characters)") + + # Clean up state between batches + cleanup_between_batches() + + # --- Prepare context from previous segment --- + context_for_next = "" + if i > 0 and segments[i-1]: # Check if previous segment exists + try: + ensure_nltk_punkt() + prev_sentences = nltk.sent_tokenize(segments[i-1]) + if prev_sentences: + context_for_next = prev_sentences[-1].strip() + print(f"Segment {i+1}: Using last sentence of previous segment as context.") + except Exception as e: + logger.warning(f"Segment {i+1}: Could not extract context from previous segment: {e}") + # --- End context preparation --- + + # Create a temporary file ONLY if a final output file is requested + temp_output_file_for_batch = None + if output_file: + temp_output_file_for_batch = f"outputs/temp_batch_{i}_{int(time.time())}.wav" + batch_temp_files.append(temp_output_file_for_batch) + + # Generate speech for this segment, passing context + batch_segments = tokens_decoder_sync( + generate_tokens_from_api( + prompt=segment, + voice=voice, + temperature=temperature, + top_p=top_p, + max_tokens=max_tokens, + repetition_penalty=REPETITION_PENALTY, + context_prompt=context_for_next # Pass context here + ), + output_file=temp_output_file_for_batch + ) + + # Combine the small segments from this batch into one bytes object + complete_batch_audio_bytes = b"".join(batch_segments) + all_complete_batch_audio.append(complete_batch_audio_bytes) + + # Also keep track of small segments if not writing to file and no crossfading needed later + if not output_file: + all_audio_segments.extend(batch_segments) + + # Explicitly clear segments list after processing all segments + segments = None + import gc + gc.collect() + + # --- Post-Batch Processing --- + if output_file: + # If an output file was requested, stitch together the temporary batch files + if batch_temp_files: + stitch_wav_files(batch_temp_files, output_file) + # Clean up temporary files + for temp_file in batch_temp_files: + try: os.remove(temp_file) + except Exception as e: print(f"Warning: Could not remove temp file {temp_file}: {e}") + # In file output mode, the final return value isn't the audio data itself + all_audio_segments = [] # Clear segments as data is in the file + + elif len(all_complete_batch_audio) > 1: + # If NO output file AND more than one batch was processed, apply crossfading + print(f"Applying in-memory crossfade to {len(all_complete_batch_audio)} audio batches.") + try: + # Convert bytes to numpy arrays + batch_arrays = [np.frombuffer(b, dtype=np.int16) for b in all_complete_batch_audio if b] # Filter empty + + if len(batch_arrays) > 1: + # Apply crossfading logic with optimizations for high-memory systems + if APPLE_SILICON and psutil.virtual_memory().total >= (64 * 1024**3): + # Pre-allocate the final array to avoid repeated concatenations + # First calculate total length of all arrays minus crossfade regions + crossfade_samples = int(SAMPLE_RATE * CROSSFADE_MS / 1000) + total_samples = sum(len(arr) for arr in batch_arrays) + # Subtract overlapping regions + total_samples -= crossfade_samples * (len(batch_arrays) - 1) + + # Pre-allocate the final array + print(f"Pre-allocating array for {total_samples} samples") + final_audio_np = np.zeros(total_samples, dtype=np.int16) + + # Fill the array with crossfaded audio + write_position = 0 + + for i, audio_np in enumerate(batch_arrays): + if i == 0: + # First segment - copy directly + final_audio_np[:len(audio_np)] = audio_np + write_position += len(audio_np) - crossfade_samples + else: + # For other segments, apply crossfade + # Generate equal power fade curves + fade_out, fade_in = generate_equal_power_fade_curves(crossfade_samples) + + # Create crossfade region + prev_end = write_position + crossfade_region = (final_audio_np[prev_end:prev_end+crossfade_samples] * fade_out + + audio_np[:crossfade_samples] * fade_in).astype(np.int16) + + # Write crossfade region + final_audio_np[prev_end:prev_end+crossfade_samples] = crossfade_region + + # Write remainder of current segment + next_end = prev_end + crossfade_samples + len(audio_np) - crossfade_samples + final_audio_np[prev_end+crossfade_samples:next_end] = audio_np[crossfade_samples:] + + # Update write position + write_position = next_end + else: + # Use standard approach for systems with less memory + final_audio_np = np.array([], dtype=np.int16) + # Use the constant for crossfade duration + crossfade_samples = int(SAMPLE_RATE * CROSSFADE_MS / 1000) + print(f"Applying {CROSSFADE_MS}ms crossfade ({crossfade_samples} samples)") + + for i, audio_np in enumerate(batch_arrays): + if i == 0: + final_audio_np = audio_np + else: + # Apply crossfade + prev_audio_np = final_audio_np + current_audio_np = audio_np + + if len(prev_audio_np) >= crossfade_samples and len(current_audio_np) >= crossfade_samples: + # Generate equal power fade curves + fade_out, fade_in = generate_equal_power_fade_curves(crossfade_samples) + + crossfade_region = (prev_audio_np[-crossfade_samples:] * fade_out + + current_audio_np[:crossfade_samples] * fade_in).astype(np.int16) + + final_audio_np = np.concatenate([ + prev_audio_np[:-crossfade_samples], + crossfade_region, + current_audio_np[crossfade_samples:] + ]) + else: + # Segments too short for crossfade, concatenate directly + print(f"Warning: Segments too short for crossfade between batch {i-1} and {i}. Concatenating.") + final_audio_np = np.concatenate([prev_audio_np, current_audio_np]) + + # Convert final numpy array back to bytes and wrap in a list + all_audio_segments = [final_audio_np.tobytes()] + print("In-memory crossfading complete.") + + # Clear batch arrays to free memory + batch_arrays = None + all_complete_batch_audio = None + gc.collect() + else: + # Only one valid batch array, use the original segments + print("Only one batch generated after filtering, no crossfading needed.") + pass + except Exception as e: + print(f"ERROR during in-memory crossfading: {e}. Returning raw concatenated segments.") + # Fallback: return the original potentially discontinuous segments if crossfade fails + all_audio_segments = [b"".join(all_audio_segments)] # Combine all small chunks + + # --- Final Reporting --- end_time = time.time() total_time = end_time - start_time - # Calculate combined duration + # Calculate combined duration from the final segments to be returned if all_audio_segments: - total_bytes = sum(len(segment) for segment in all_audio_segments) - duration = total_bytes / (2 * SAMPLE_RATE) # 2 bytes per sample at 24kHz - print(f"Generated {len(all_audio_segments)} audio segments") - print(f"Generated {duration:.2f} seconds of audio in {total_time:.2f} seconds") - print(f"Realtime factor: {duration/total_time:.2f}x") + # Handle if crossfading resulted in a single large chunk + if len(all_audio_segments) == 1: + total_bytes = len(all_audio_segments[0]) + else: # Original chunked segments + total_bytes = sum(len(segment) for segment in all_audio_segments) + if total_bytes > 0: + duration = total_bytes / (2 * SAMPLE_RATE) # 2 bytes per sample + print(f"Generated {len(all_audio_segments)} final audio segment(s)") # Correctly reports 1 segment after crossfade + print(f"Generated {duration:.2f} seconds of audio in {total_time:.2f} seconds") + if total_time > 0: + realtime_factor = duration / total_time + print(f"Realtime factor: {realtime_factor:.2f}x") + if realtime_factor < 1.0: + print("⚠️ Warning: Generation is slower than realtime") + else: + print(f"✓ Generation is {realtime_factor:.1f}x faster than realtime") + else: + print("Generation time was negligible.") + else: + print("Warning: No audio data generated.") + print(f"Total speech generation completed in {total_time:.2f} seconds") + # Return the final audio segments (either original chunks or one combined chunk after crossfade) return all_audio_segments -def stitch_wav_files(input_files, output_file, crossfade_ms=50): +def stitch_wav_files(input_files, output_file): """Stitch multiple WAV files together with crossfading for smooth transitions.""" if not input_files: return - print(f"Stitching {len(input_files)} WAV files together with {crossfade_ms}ms crossfade") + print(f"Stitching {len(input_files)} WAV files together with {CROSSFADE_MS}ms crossfade") # If only one file, just copy it if len(input_files) == 1: import shutil shutil.copy(input_files[0], output_file) + print(f"Only one input file, copied directly to {output_file}") return - # Convert crossfade_ms to samples - crossfade_samples = int(SAMPLE_RATE * crossfade_ms / 1000) + # Convert crossfade_ms to samples using the constant + crossfade_samples = int(SAMPLE_RATE * CROSSFADE_MS / 1000) print(f"Using {crossfade_samples} samples for crossfade at {SAMPLE_RATE}Hz") # Build the final audio in memory with crossfades final_audio = np.array([], dtype=np.int16) first_params = None + # Standard WAV parameters to enforce + standard_params = { + 'nchannels': 1, + 'sampwidth': 2, + 'framerate': SAMPLE_RATE + } + for i, input_file in enumerate(input_files): try: with wave.open(input_file, 'rb') as wav: + # Get current file parameters + current_params = wav.getparams() + + # Check and standardize parameters if first_params is None: - first_params = wav.getparams() - elif wav.getparams() != first_params: - print(f"Warning: WAV file {input_file} has different parameters") - + first_params = current_params + # Verify first file meets our standards + if (current_params.nchannels != standard_params['nchannels'] or + current_params.sampwidth != standard_params['sampwidth'] or + current_params.framerate != standard_params['framerate']): + print(f"Warning: First WAV file {input_file} has non-standard parameters. Converting to standard format.") + elif current_params != first_params: + print(f"Warning: WAV file {input_file} has different parameters. Converting to standard format.") + + # Read frames and convert to numpy array frames = wav.readframes(wav.getnframes()) audio = np.frombuffer(frames, dtype=np.int16) @@ -828,9 +1312,8 @@ def stitch_wav_files(input_files, output_file, crossfade_ms=50): else: # Apply crossfade with previous segment if len(final_audio) >= crossfade_samples and len(audio) >= crossfade_samples: - # Create crossfade weights - fade_out = np.linspace(1.0, 0.0, crossfade_samples) - fade_in = np.linspace(0.0, 1.0, crossfade_samples) + # Generate equal power fade curves + fade_out, fade_in = generate_equal_power_fade_curves(crossfade_samples) # Apply crossfade crossfade_region = (final_audio[-crossfade_samples:] * fade_out + @@ -855,7 +1338,10 @@ def stitch_wav_files(input_files, output_file, crossfade_ms=50): if first_params is None: raise ValueError("No valid WAV files were processed") - output_wav.setparams(first_params) + # Use standard parameters for output + output_wav.setnchannels(standard_params['nchannels']) + output_wav.setsampwidth(standard_params['sampwidth']) + output_wav.setframerate(standard_params['framerate']) output_wav.writeframes(final_audio.tobytes()) print(f"Successfully stitched audio to {output_file} with crossfading") diff --git a/tts_engine/speechpipe.py b/tts_engine/speechpipe.py index 1a789d7..ce752a0 100644 --- a/tts_engine/speechpipe.py +++ b/tts_engine/speechpipe.py @@ -17,10 +17,32 @@ def is_reloader_process(): # Set a flag to avoid repeat messages IS_RELOADER = is_reloader_process() +# Detect hardware capabilities +APPLE_SILICON = torch.backends.mps.is_available() +CUDA_AVAILABLE = torch.cuda.is_available() + +# Set device for model processing +if APPLE_SILICON: + DEVICE = "mps" + if not IS_RELOADER: + print("🍎 Using Apple Silicon MPS for speech generation") +elif CUDA_AVAILABLE: + DEVICE = "cuda" + if not IS_RELOADER: + print("🖥️ Using CUDA for speech generation") +else: + DEVICE = "cpu" + if not IS_RELOADER: + print("⚙️ Using CPU for speech generation") + +# Check if CoreML should be enabled +# USE_COREML = APPLE_SILICON and os.environ.get("ORPHEUS_USE_COREML", "1") == "1" +# CoreML logic removed + # Try to enable torch.compile if PyTorch 2.0+ is available TORCH_COMPILE_AVAILABLE = False try: - if hasattr(torch, 'compile'): + if hasattr(torch, 'compile') and not APPLE_SILICON: # torch.compile not fully supported on MPS yet TORCH_COMPILE_AVAILABLE = True if not IS_RELOADER: print("PyTorch 2.0+ detected, torch.compile is available") @@ -30,29 +52,58 @@ except: # Try to enable CUDA graphs if available CUDA_GRAPHS_AVAILABLE = False try: - if torch.cuda.is_available() and hasattr(torch.cuda, 'make_graphed_callables'): + if CUDA_AVAILABLE and hasattr(torch.cuda, 'make_graphed_callables'): CUDA_GRAPHS_AVAILABLE = True if not IS_RELOADER: print("CUDA graphs support is available") except: pass -model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval() +# Load the model with appropriate device placement +base_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval() +base_model = base_model.to(DEVICE) -# Check if CUDA is available and set device accordingly -snac_device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" -if not IS_RELOADER: - print(f"Using device: {snac_device}") -model = model.to(snac_device) +# Assign base_model directly as the model to use +model = base_model + +# Check if CoreML wrapper should be used +# if USE_COREML: +# try: +# from .coreml_wrapper import CoreMLWrapper +# # Wrap the base model with CoreML for Apple Silicon +# model = CoreMLWrapper(base_model, device=DEVICE) +# if model.use_coreml and model.coreml_model is not None: +# if not IS_RELOADER: +# print("🧠 Using CoreML Neural Engine acceleration!") +# else: +# if not IS_RELOADER: +# print("CoreML model not available, using MPS acceleration instead") +# except ImportError: +# # If CoreML wrapper is not available, use the base model +# model = base_model +# if not IS_RELOADER: +# print("CoreML wrapper not available, using standard PyTorch backend") +# else: +# # Use base model directly +# model = base_model -# Disable torch.compile as it requires Triton which isn't installed -# We'll use regular PyTorch optimization techniques instead if not IS_RELOADER: - print("Using standard PyTorch optimizations (torch.compile disabled)") + print(f"SNAC model loaded directly on {DEVICE} (CoreML export disabled)") + +# Disable torch.compile for MPS as it's not fully supported +if APPLE_SILICON: + if not IS_RELOADER: + print("Using standard PyTorch optimizations for Apple Silicon") +elif TORCH_COMPILE_AVAILABLE: + if not IS_RELOADER: + print("Using torch.compile for optimized performance") +else: + if not IS_RELOADER: + print("Using standard PyTorch optimizations") # Prepare CUDA streams for parallel processing if available cuda_stream = None -if snac_device == "cuda": +if CUDA_AVAILABLE: cuda_stream = torch.cuda.Stream() if not IS_RELOADER: print("Using CUDA stream for parallel processing") @@ -60,8 +111,8 @@ if snac_device == "cuda": def convert_to_audio(multiframe, count): """ - Optimized version of convert_to_audio that eliminates inefficient tensor operations - and reduces CPU-GPU transfers for much faster inference on high-end GPUs. + Optimized version of convert_to_audio that supports Apple Silicon MPS and + eliminates inefficient tensor operations for much faster inference. """ if len(multiframe) < 7: return None @@ -70,12 +121,12 @@ def convert_to_audio(multiframe, count): frame = multiframe[:num_frames*7] # Pre-allocate tensors instead of incrementally building them - codes_0 = torch.zeros(num_frames, dtype=torch.int32, device=snac_device) - codes_1 = torch.zeros(num_frames * 2, dtype=torch.int32, device=snac_device) - codes_2 = torch.zeros(num_frames * 4, dtype=torch.int32, device=snac_device) + codes_0 = torch.zeros(num_frames, dtype=torch.int32, device=DEVICE) + codes_1 = torch.zeros(num_frames * 2, dtype=torch.int32, device=DEVICE) + codes_2 = torch.zeros(num_frames * 4, dtype=torch.int32, device=DEVICE) # Use vectorized operations where possible - frame_tensor = torch.tensor(frame, dtype=torch.int32, device=snac_device) + frame_tensor = torch.tensor(frame, dtype=torch.int32, device=DEVICE) # Direct indexing is much faster than concatenation in a loop for j in range(num_frames): @@ -107,27 +158,33 @@ def convert_to_audio(multiframe, count): torch.any(codes[2] < 0) or torch.any(codes[2] > 4096)): return None - # Use CUDA stream for parallel processing if available - stream_ctx = torch.cuda.stream(cuda_stream) if cuda_stream is not None else torch.no_grad() + # Context manager depends on device type + if CUDA_AVAILABLE: + stream_ctx = torch.cuda.stream(cuda_stream) + else: + stream_ctx = torch.inference_mode() - with stream_ctx, torch.inference_mode(): - # Decode the audio - audio_hat = model.decode(codes) + with stream_ctx: + # Decode the audio using the base model directly + audio_hat = model.decode(codes) # model is now directly base_model # Extract the relevant slice and efficiently convert to bytes # Keep data on GPU as long as possible audio_slice = audio_hat[:, :, 2048:4096] - # Process on GPU if possible, with minimal data transfer - if snac_device == "cuda": + # Process based on device type to minimize data transfers + if CUDA_AVAILABLE: # Scale directly on GPU audio_int16_tensor = (audio_slice * 32767).to(torch.int16) # Only transfer the final result to CPU audio_bytes = audio_int16_tensor.cpu().numpy().tobytes() + elif APPLE_SILICON: + # For MPS, we need to go through CPU + audio_int16_tensor = (audio_slice * 32767).to(torch.int16) + audio_bytes = audio_int16_tensor.detach().cpu().numpy().tobytes() else: - # For non-CUDA devices, fall back to the original approach - detached_audio = audio_slice.detach().cpu() - audio_np = detached_audio.numpy() + # For CPU, simpler pathway + audio_np = audio_slice.detach().numpy() audio_int16 = (audio_np * 32767).astype(np.int16) audio_bytes = audio_int16.tobytes() @@ -189,7 +246,7 @@ async def tokens_decoder(token_gen): """Optimized token decoder with early first-chunk processing for lower latency""" buffer = [] count = 0 - + # Track if first chunk has been processed first_chunk_processed = False @@ -288,15 +345,69 @@ async def tokens_decoder(token_gen): audio_samples = convert_to_audio(padded_buffer, count) if audio_samples is not None: yield audio_samples + +def reset_state(): + """Reset all state between batch processing.""" + global token_id_cache, cuda_stream + + # Clear the token ID cache + token_id_cache.clear() + + # Reset CUDA stream if available + if CUDA_AVAILABLE and cuda_stream is not None: + cuda_stream.synchronize() + cuda_stream = torch.cuda.Stream() + + # Clear CUDA cache + torch.cuda.empty_cache() + + # For Apple Silicon, clear MPS cache if available + if APPLE_SILICON: + try: + # This is the proper way to clear MPS cache in PyTorch 2.0+ + torch.mps.empty_cache() + except: + # Fallback for older PyTorch versions or if the operation fails + pass + + # Reset CoreML state if needed + # if USE_COREML and isinstance(model, CoreMLWrapper) and model.use_coreml: + # try: + # # CoreML resources are generally managed by the OS + # # but we'll manually trigger garbage collection to be safe + # import gc + # gc.collect() + # except: + # pass + + # Force garbage collection + import gc + gc.collect() + # ------------------ Synchronous Tokens Decoder Wrapper ------------------ # def tokens_decoder_sync(syn_token_gen): - """Optimized synchronous decoder with larger queue and parallel processing""" - # Use a larger queue for RTX 4090 to maximize GPU utilization - max_queue_size = 32 if snac_device == "cuda" else 8 - audio_queue = queue.Queue(maxsize=max_queue_size) + """Optimized synchronous decoder with hardware-specific optimizations""" + # Set queue size based on hardware + if APPLE_SILICON: + import psutil + ram_gb = psutil.virtual_memory().total / (1024**3) + if ram_gb >= 64: # High-memory Apple Silicon + max_queue_size = 64 + batch_size = 32 + elif ram_gb >= 32: # Mid-range Apple Silicon + max_queue_size = 48 + batch_size = 24 + else: # Base model + max_queue_size = 32 + batch_size = 16 + elif CUDA_AVAILABLE: + max_queue_size = 32 + batch_size = 16 + else: # CPU fallback + max_queue_size = 8 + batch_size = 4 - # Collect tokens in batches for higher throughput - batch_size = 16 if snac_device == "cuda" else 4 + audio_queue = queue.Queue(maxsize=max_queue_size) # Convert the synchronous token generator into an async generator with batching async def async_token_gen(): @@ -340,13 +451,16 @@ def tokens_decoder_sync(syn_token_gen): def run_async(): asyncio.run(async_producer()) - # Use a higher priority thread for RTX 4090 to ensure it stays fed with work + # Create thread with appropriate priority thread = threading.Thread(target=run_async) thread.daemon = True # Allow the thread to be terminated when the main thread exits thread.start() - # Use larger buffer for final audio assembly - buffer_size = 5 + # Use hardware-specific buffer sizes + if APPLE_SILICON: + buffer_size = 8 # Larger buffer for smoother playback on Apple Silicon + else: + buffer_size = 5 audio_buffer = [] while True: