diff --git a/.DS_Store b/.DS_Store
new file mode 100644
index 0000000..24e4bf2
Binary files /dev/null and b/.DS_Store differ
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">
+
+
+
Max Batch Characters
+
+
Maximum characters per batch for sentence splitting (default: 600)
+
+
Crossfade Duration (ms)
+
+
Crossfade duration between audio batches in milliseconds (default: 30)
+
-
Play
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: