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_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) 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) # 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_PORT=5005
ORPHEUS_HOST=0.0.0.0 ORPHEUS_HOST=0.0.0.0

9
.gitignore vendored
View file

@ -3,4 +3,11 @@ __pycache__/
.venv/ .venv/
venv/ venv/
models/ 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. 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 ## Development
### Project Components ### Project Components

30
app.py
View file

@ -51,7 +51,16 @@ from fastapi.templating import Jinja2Templates
from pydantic import BaseModel from pydantic import BaseModel
import json 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 # Create FastAPI app
app = FastAPI( app = FastAPI(
@ -96,7 +105,7 @@ async def create_speech_api(request: SpeechRequest):
Generate speech from text using the Orpheus TTS model. Generate speech from text using the Orpheus TTS model.
Compatible with OpenAI's /v1/audio/speech endpoint. 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. to improve reliability and avoid truncation issues.
""" """
if not request.input: if not request.input:
@ -107,7 +116,7 @@ async def create_speech_api(request: SpeechRequest):
output_path = f"outputs/{request.voice}_{timestamp}.wav" output_path = f"outputs/{request.voice}_{timestamp}.wav"
# Check if we should use batched generation # Check if we should use batched generation
use_batching = len(request.input) > 1000 use_batching = len(request.input) > MAX_BATCH_CHARS
if use_batching: if use_batching:
print(f"Using batched generation for long text ({len(request.input)} characters)") 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, prompt=request.input,
voice=request.voice, voice=request.voice,
output_file=output_path, output_file=output_path,
use_batching=use_batching, use_batching=use_batching
max_batch_chars=1000 # Process in ~1000 character chunks (roughly 1 paragraph)
) )
end = time.time() end = time.time()
generation_time = round(end - start, 2) generation_time = round(end - start, 2)
@ -160,7 +168,7 @@ async def speak(request: Request):
output_path = f"outputs/{voice}_{timestamp}.wav" output_path = f"outputs/{voice}_{timestamp}.wav"
# Check if we should use batched generation for longer texts # 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: if use_batching:
print(f"Using batched generation for long text ({len(text)} characters)") print(f"Using batched generation for long text ({len(text)} characters)")
@ -170,8 +178,7 @@ async def speak(request: Request):
prompt=text, prompt=text,
voice=voice, voice=voice,
output_file=output_path, output_file=output_path,
use_batching=use_batching, use_batching=use_batching
max_batch_chars=1000
) )
end = time.time() end = time.time()
generation_time = round(end - start, 2) generation_time = round(end - start, 2)
@ -226,7 +233,7 @@ async def save_config(request: Request):
# Convert values to proper types # Convert values to proper types
for key, value in data.items(): 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: try:
data[key] = str(int(value)) data[key] = str(int(value))
except (ValueError, TypeError): except (ValueError, TypeError):
@ -322,7 +329,7 @@ async def generate_from_web(
output_path = f"outputs/{voice}_{timestamp}.wav" output_path = f"outputs/{voice}_{timestamp}.wav"
# Check if we should use batched generation for longer texts # 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: if use_batching:
print(f"Using batched generation for long text from web form ({len(text)} characters)") 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, prompt=text,
voice=voice, voice=voice,
output_file=output_path, output_file=output_path,
use_batching=use_batching, use_batching=use_batching
max_batch_chars=1000
) )
end = time.time() end = time.time()
generation_time = round(end - start, 2) 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>
<div class="text-xs text-dark-300"> <div class="text-xs text-dark-300">
{% if voice_option == "tara" %}Female, English, conversational, clear {% 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 == "leah" %}Female, English, warm, gentle
{% elif voice_option == "jess" %}Female, English, energetic, youthful {% elif voice_option == "jess" %}Female, English, energetic, youthful
{% elif voice_option == "leo" %}Male, English, authoritative, deep {% 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"> 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> </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 --> <!-- Save button and restart button -->
<div class="col-span-1 md:col-span-2 mt-4 flex flex-col md:flex-row gap-4"> <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" <button id="save-config-btn" type="button"
@ -414,7 +432,7 @@
<div class="flex items-center space-x-4"> <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"> <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"> <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> </svg>
Play Play
</button> </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: This package contains the core components for audio generation:
- inference.py: Token generation and API handling - inference.py: Token generation and API handling
- speechpipe.py: Audio conversion pipeline - speechpipe.py: Audio conversion pipeline
- coreml_wrapper.py: Apple Silicon Neural Engine acceleration
""" """
# Make key components available at package level # Make key components available at package level
from .inference import ( from .inference import (
generate_speech_from_api, generate_speech_from_api,
stream_audio,
AVAILABLE_VOICES, AVAILABLE_VOICES,
DEFAULT_VOICE, DEFAULT_VOICE,
VOICE_TO_LANGUAGE, VOICE_TO_LANGUAGE,
AVAILABLE_LANGUAGES, 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 # Set a flag to avoid repeat messages
IS_RELOADER = is_reloader_process() 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 # Try to enable torch.compile if PyTorch 2.0+ is available
TORCH_COMPILE_AVAILABLE = False TORCH_COMPILE_AVAILABLE = False
try: 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 TORCH_COMPILE_AVAILABLE = True
if not IS_RELOADER: if not IS_RELOADER:
print("PyTorch 2.0+ detected, torch.compile is available") print("PyTorch 2.0+ detected, torch.compile is available")
@ -30,29 +52,58 @@ except:
# Try to enable CUDA graphs if available # Try to enable CUDA graphs if available
CUDA_GRAPHS_AVAILABLE = False CUDA_GRAPHS_AVAILABLE = False
try: 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 CUDA_GRAPHS_AVAILABLE = True
if not IS_RELOADER: if not IS_RELOADER:
print("CUDA graphs support is available") print("CUDA graphs support is available")
except: except:
pass 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 # Assign base_model directly as the model to use
snac_device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" model = base_model
if not IS_RELOADER:
print(f"Using device: {snac_device}") # Check if CoreML wrapper should be used
model = model.to(snac_device) # 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: 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 # Prepare CUDA streams for parallel processing if available
cuda_stream = None cuda_stream = None
if snac_device == "cuda": if CUDA_AVAILABLE:
cuda_stream = torch.cuda.Stream() cuda_stream = torch.cuda.Stream()
if not IS_RELOADER: if not IS_RELOADER:
print("Using CUDA stream for parallel processing") print("Using CUDA stream for parallel processing")
@ -60,8 +111,8 @@ if snac_device == "cuda":
def convert_to_audio(multiframe, count): def convert_to_audio(multiframe, count):
""" """
Optimized version of convert_to_audio that eliminates inefficient tensor operations Optimized version of convert_to_audio that supports Apple Silicon MPS and
and reduces CPU-GPU transfers for much faster inference on high-end GPUs. eliminates inefficient tensor operations for much faster inference.
""" """
if len(multiframe) < 7: if len(multiframe) < 7:
return None return None
@ -70,12 +121,12 @@ def convert_to_audio(multiframe, count):
frame = multiframe[:num_frames*7] frame = multiframe[:num_frames*7]
# Pre-allocate tensors instead of incrementally building them # Pre-allocate tensors instead of incrementally building them
codes_0 = torch.zeros(num_frames, 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=snac_device) codes_1 = torch.zeros(num_frames * 2, dtype=torch.int32, device=DEVICE)
codes_2 = torch.zeros(num_frames * 4, dtype=torch.int32, device=snac_device) codes_2 = torch.zeros(num_frames * 4, dtype=torch.int32, device=DEVICE)
# Use vectorized operations where possible # 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 # Direct indexing is much faster than concatenation in a loop
for j in range(num_frames): 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)): torch.any(codes[2] < 0) or torch.any(codes[2] > 4096)):
return None return None
# Use CUDA stream for parallel processing if available # Context manager depends on device type
stream_ctx = torch.cuda.stream(cuda_stream) if cuda_stream is not None else torch.no_grad() if CUDA_AVAILABLE:
stream_ctx = torch.cuda.stream(cuda_stream)
else:
stream_ctx = torch.inference_mode()
with stream_ctx, torch.inference_mode(): with stream_ctx:
# Decode the audio # Decode the audio using the base model directly
audio_hat = model.decode(codes) audio_hat = model.decode(codes) # model is now directly base_model
# Extract the relevant slice and efficiently convert to bytes # Extract the relevant slice and efficiently convert to bytes
# Keep data on GPU as long as possible # Keep data on GPU as long as possible
audio_slice = audio_hat[:, :, 2048:4096] audio_slice = audio_hat[:, :, 2048:4096]
# Process on GPU if possible, with minimal data transfer # Process based on device type to minimize data transfers
if snac_device == "cuda": if CUDA_AVAILABLE:
# Scale directly on GPU # Scale directly on GPU
audio_int16_tensor = (audio_slice * 32767).to(torch.int16) audio_int16_tensor = (audio_slice * 32767).to(torch.int16)
# Only transfer the final result to CPU # Only transfer the final result to CPU
audio_bytes = audio_int16_tensor.cpu().numpy().tobytes() 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: else:
# For non-CUDA devices, fall back to the original approach # For CPU, simpler pathway
detached_audio = audio_slice.detach().cpu() audio_np = audio_slice.detach().numpy()
audio_np = detached_audio.numpy()
audio_int16 = (audio_np * 32767).astype(np.int16) audio_int16 = (audio_np * 32767).astype(np.int16)
audio_bytes = audio_int16.tobytes() 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""" """Optimized token decoder with early first-chunk processing for lower latency"""
buffer = [] buffer = []
count = 0 count = 0
# Track if first chunk has been processed # Track if first chunk has been processed
first_chunk_processed = False first_chunk_processed = False
@ -288,15 +345,69 @@ async def tokens_decoder(token_gen):
audio_samples = convert_to_audio(padded_buffer, count) audio_samples = convert_to_audio(padded_buffer, count)
if audio_samples is not None: if audio_samples is not None:
yield audio_samples 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 ------------------ # # ------------------ Synchronous Tokens Decoder Wrapper ------------------ #
def tokens_decoder_sync(syn_token_gen): def tokens_decoder_sync(syn_token_gen):
"""Optimized synchronous decoder with larger queue and parallel processing""" """Optimized synchronous decoder with hardware-specific optimizations"""
# Use a larger queue for RTX 4090 to maximize GPU utilization # Set queue size based on hardware
max_queue_size = 32 if snac_device == "cuda" else 8 if APPLE_SILICON:
audio_queue = queue.Queue(maxsize=max_queue_size) 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 audio_queue = queue.Queue(maxsize=max_queue_size)
batch_size = 16 if snac_device == "cuda" else 4
# Convert the synchronous token generator into an async generator with batching # Convert the synchronous token generator into an async generator with batching
async def async_token_gen(): async def async_token_gen():
@ -340,13 +451,16 @@ def tokens_decoder_sync(syn_token_gen):
def run_async(): def run_async():
asyncio.run(async_producer()) 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 = threading.Thread(target=run_async)
thread.daemon = True # Allow the thread to be terminated when the main thread exits thread.daemon = True # Allow the thread to be terminated when the main thread exits
thread.start() thread.start()
# Use larger buffer for final audio assembly # Use hardware-specific buffer sizes
buffer_size = 5 if APPLE_SILICON:
buffer_size = 8 # Larger buffer for smoother playback on Apple Silicon
else:
buffer_size = 5
audio_buffer = [] audio_buffer = []
while True: while True: