Minor improvements to account for large-VRAM Apple Silicon machines, and refine the inferencing code to ignore spurious, non-audio responses from the Orpheus model, use the system prompt if it exists in the directory structure, and for long texts, 1) use nltk for segmenting the text according to user-specified max character lengths, and 2) use the previous segment (For segment index n>1) as context in the system prompt, to hopefully improve audio generation consistency.

This commit is contained in:
malibrated 2025-04-21 22:30:33 -07:00
parent 70a5e06aef
commit 4ead99754a
11 changed files with 2116 additions and 278 deletions

BIN
.DS_Store vendored Normal file

Binary file not shown.

View file

@ -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

9
.gitignore vendored
View file

@ -3,4 +3,11 @@ __pycache__/
.venv/
venv/
models/
*.gguf
*.gguf
outputs/
output_long*.txt
*.txt
jinja
*.zsh
*.bak
*.sav

View file

@ -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

30
app.py
View file

@ -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)

1
restart.flag Normal file
View file

@ -0,0 +1 @@
1745296574.27762

View file

@ -183,6 +183,7 @@
</div>
<div class="text-xs text-dark-300">
{% 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">
</div>
<!-- Audio Generation Settings -->
<div>
<label for="max_batch_chars" class="block text-xs font-medium text-white mb-1">Max Batch Characters</label>
<input type="number" id="max_batch_chars" name="ORPHEUS_MAX_BATCH_CHARS" min="100" max="2000" step="50"
placeholder="600"
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">
<p class="text-xs text-purple-300 mt-1">Maximum characters per batch for sentence splitting (default: 600)</p>
</div>
<div>
<label for="crossfade_ms" class="block text-xs font-medium text-white mb-1">Crossfade Duration (ms)</label>
<input type="number" id="crossfade_ms" name="ORPHEUS_CROSSFADE_MS" min="10" max="200" step="1"
placeholder="30"
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">
<p class="text-xs text-purple-300 mt-1">Crossfade duration between audio batches in milliseconds (default: 30)</p>
</div>
<!-- Save button and restart button -->
<div class="col-span-1 md:col-span-2 mt-4 flex flex-col md:flex-row gap-4">
<button id="save-config-btn" type="button"
@ -414,7 +432,7 @@
<div class="flex items-center space-x-4">
<button id="play-btn" class="inline-flex items-center px-4 py-2 border border-primary-700 rounded-md shadow-sm text-sm font-medium text-white bg-primary-600 hover:bg-primary-700 focus:outline-none focus:ring-2 focus:ring-primary-500 focus:ring-offset-dark-800 focus:ring-offset-2">
<svg class="h-5 w-5 mr-2" fill="currentColor" viewBox="0 0 20 20">
<path fill-rule="evenodd" d="M10 18a8 8 0 100-16 8 8 0 000 16zM9.555 7.168A1 1 0 008 8v4a1 1 0 001.555.832l3-2a1 1 0 000-1.664l-3-2z" clip-rule="evenodd" />
<path fill-rule="evenodd" d="M10 18a8 8 0 11-16 0 8 8 0 0116 0zM9.555 7.168A1 1 0 008 8v4a1 1 0 001.555.832l3-2a1 1 0 000-1.664l-3-2z" clip-rule="evenodd" />
</svg>
Play
</button>

View file

@ -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",
]

1124
tts_engine/inference.mps Normal file

File diff suppressed because it is too large Load diff

File diff suppressed because it is too large Load diff

View file

@ -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: