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:
parent
70a5e06aef
commit
4ead99754a
11 changed files with 2116 additions and 278 deletions
BIN
.DS_Store
vendored
Normal file
BIN
.DS_Store
vendored
Normal file
Binary file not shown.
|
|
@ -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
9
.gitignore
vendored
|
|
@ -3,4 +3,11 @@ __pycache__/
|
|||
.venv/
|
||||
venv/
|
||||
models/
|
||||
*.gguf
|
||||
*.gguf
|
||||
outputs/
|
||||
output_long*.txt
|
||||
*.txt
|
||||
jinja
|
||||
*.zsh
|
||||
*.bak
|
||||
*.sav
|
||||
|
|
|
|||
40
README.md
40
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
|
||||
|
|
|
|||
30
app.py
30
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)
|
||||
|
|
|
|||
1
restart.flag
Normal file
1
restart.flag
Normal file
|
|
@ -0,0 +1 @@
|
|||
1745296574.27762
|
||||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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
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
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Reference in a new issue