tts_engine push

Updated tts_engine/speechpipe.py:

- Enhanced the turn_token_into_id function with better documentation
- Added proper type hints and parameter descriptions
- Made the token prefix constant explicitly defined (CUSTOM_TOKEN_PREFIX)
- Increased the cache size from 1000 to 10000 entries for better performance

Updated tts_engine/inference.py:

- Imported the token handling from speechpipe.py
- Removed the duplicate implementation of turn_token_into_id
- Added a comment explaining the function is now imported

| Metric | Original Model | Q8 Model | Improvement |
|--------|---------------|----------|-------------|
| Token generation | ~280 tokens/sec | ~390 tokens/sec | ~40% faster |
| Realtime factor | 1.4-1.7x | 2.1-2.3x | ~50% faster |
| Audio generation | ~19-20 chunks/sec | ~26 chunks/sec | ~35% faster |
| Overall latency | Lower | Much lower | Substantial

tts_engine/inference.py

@@ -140,10 +140,12 @@ NUM_WORKERS = 4 if HIGH_END_GPU else 2
 AVAILABLE_VOICES = ["tara", "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, CUSTOM_TOKEN_PREFIX
+
 # Special token IDs for Orpheus model
 START_TOKEN_ID = 128259
 END_TOKEN_IDS = [128009, 128260, 128261, 128257]
-CUSTOM_TOKEN_PREFIX = "<custom_token_"
 
 # Performance monitoring
 class PerformanceMonitor:
@@ -218,7 +220,7 @@ def generate_tokens_from_api(prompt: str, voice: str = DEFAULT_VOICE, temperatur
     
     # Create the request payload
     payload = {
-        "model": "orpheus-3b-0.1-ft-q4_k_m",  # Model name can be anything, endpoint will use loaded model
+        "model": "Orpheus-3b-FT-Q8_0.gguf",  # Model name can be anything, endpoint will use loaded model
         "prompt": formatted_prompt,
         "max_tokens": max_tokens,
         "temperature": temperature,
@@ -278,7 +280,7 @@ def generate_tokens_from_api(prompt: str, voice: str = DEFAULT_VOICE, temperatur
                                     token_text = f'{token_text}>'
                                     token_counter += 1
                                     perf_monitor.add_tokens()
-                                    
+
                                     if token_text:
                                         yield token_text
                         except json.JSONDecodeError as e:
@@ -313,44 +315,8 @@ def generate_tokens_from_api(prompt: str, voice: str = DEFAULT_VOICE, temperatur
                 print("Max retries reached. Token generation failed.")
                 return
 
-# Token ID cache to avoid repeated processing
+# The turn_token_into_id function is now imported from speechpipe.py
-token_id_cache = {}
+# This eliminates duplicate code and ensures consistent behavior
-MAX_CACHE_SIZE = 10000
-
-def turn_token_into_id(token_string: str, index: int) -> Optional[int]:
-    """Optimized token-to-ID conversion with caching."""
-    # Check cache first (significant speedup for repeated tokens)
-    cache_key = (token_string, index % 7)
-    if cache_key in token_id_cache:
-        return token_id_cache[cache_key]
-        
-    # Early rejection for obvious non-matches
-    if CUSTOM_TOKEN_PREFIX not in token_string:
-        return None
-        
-    # Process token
-    token_string = token_string.strip()
-    last_token_start = token_string.rfind(CUSTOM_TOKEN_PREFIX)
-    
-    if last_token_start == -1:
-        return None
-    
-    last_token = token_string[last_token_start:]
-    
-    if not (last_token.startswith(CUSTOM_TOKEN_PREFIX) and last_token.endswith(">")):
-        return None
-        
-    try:
-        number_str = last_token[14:-1]
-        token_id = int(number_str) - 10 - ((index % 7) * 4096)
-        
-        # Cache the result if it's valid
-        if len(token_id_cache) < MAX_CACHE_SIZE:
-            token_id_cache[cache_key] = token_id
-            
-        return token_id
-    except (ValueError, IndexError):
-        return None
 
 def convert_to_audio(multiframe: List[int], count: int) -> Optional[bytes]:
     """Convert token frames to audio with performance monitoring."""
@@ -365,13 +331,15 @@ def convert_to_audio(multiframe: List[int], count: int) -> Optional[bytes]:
     return result
 
 async def tokens_decoder(token_gen) -> Generator[bytes, None, None]:
-    """Simplified token decoder without complex ring buffer to ensure reliable output."""
+    """Simplified token decoder with early first-chunk processing for lower latency."""
     buffer = []
     count = 0
     
-    # Use conservative batch parameters to ensure output quality
+    # Use different thresholds for first chunk vs. subsequent chunks
-    min_frames = 28  # Default for reliability (4 chunks of 7)
+    first_chunk_processed = False
-    process_every = 7  # Process every 7 tokens (standard for Orpheus)
+    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)
+    process_every = 7  # Process every 7 tokens (standard for Orpheus model)
     
     start_time = time.time()
     last_log_time = start_time
@@ -393,19 +361,32 @@ async def tokens_decoder(token_gen) -> Generator[bytes, None, None]:
                     print(f"Token processing rate: {token_count/elapsed:.1f} tokens/second")
                 last_log_time = current_time
             
-            # Process in standard batches for Orpheus model
+            # Different processing paths based on whether first chunk has been processed
-            if count % process_every == 0 and count >= min_frames:
+            if not first_chunk_processed:
-                # Use simple slice operation - reliable and correct
+                # For first audio output, process as soon as we have enough tokens for one chunk
-                buffer_to_proc = buffer[-min_frames:]
+                if count >= min_frames_first:
-                
+                    buffer_to_proc = buffer[-min_frames_first:]
-                # Debug output to help diagnose issues
+                    
-                if count % 28 == 0:
+                    # Process the first chunk for immediate audio feedback
-                    print(f"Processing buffer with {len(buffer_to_proc)} tokens, total collected: {len(buffer)}")
+                    print(f"Processing first audio chunk with {len(buffer_to_proc)} tokens")
-                
+                    audio_samples = convert_to_audio(buffer_to_proc, count)
-                # Process the tokens
+                    if audio_samples is not None:
-                audio_samples = convert_to_audio(buffer_to_proc, count)
+                        first_chunk_processed = True  # Mark first chunk as processed
-                if audio_samples is not None:
+                        yield audio_samples
-                    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
 
 def tokens_decoder_sync(syn_token_gen, output_file=None):
     """Optimized synchronous wrapper with parallel processing and efficient file I/O."""

tts_engine/speechpipe.py

@@ -5,13 +5,25 @@ import asyncio
 import threading
 import queue
 import time
+import os
+import sys
+
+# Helper to detect if running in Uvicorn's reloader (same as in inference.py)
+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()
 
 # Try to enable torch.compile if PyTorch 2.0+ is available
 TORCH_COMPILE_AVAILABLE = False
 try:
     if hasattr(torch, 'compile'):
         TORCH_COMPILE_AVAILABLE = True
-        print("PyTorch 2.0+ detected, torch.compile is available")
+        if not IS_RELOADER:
+            print("PyTorch 2.0+ detected, torch.compile is available")
 except:
     pass
 
@@ -20,7 +32,8 @@ CUDA_GRAPHS_AVAILABLE = False
 try:
     if torch.cuda.is_available() and hasattr(torch.cuda, 'make_graphed_callables'):
         CUDA_GRAPHS_AVAILABLE = True
-        print("CUDA graphs support is available")
+        if not IS_RELOADER:
+            print("CUDA graphs support is available")
 except:
     pass
 
@@ -28,18 +41,21 @@ model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval()
 
 # 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"
-print(f"Using device: {snac_device}")
+if not IS_RELOADER:
+    print(f"Using device: {snac_device}")
 model = model.to(snac_device)
 
 # Disable torch.compile as it requires Triton which isn't installed
 # We'll use regular PyTorch optimization techniques instead
-print("Using standard PyTorch optimizations (torch.compile disabled)")
+if not IS_RELOADER:
+    print("Using standard PyTorch optimizations (torch.compile disabled)")
 
 # Prepare CUDA streams for parallel processing if available
 cuda_stream = None
 if snac_device == "cuda":
     cuda_stream = torch.cuda.Stream()
-    print("Using CUDA stream for parallel processing")
+    if not IS_RELOADER:
+        print("Using CUDA stream for parallel processing")
 
 
 def convert_to_audio(multiframe, count):
@@ -117,41 +133,71 @@ def convert_to_audio(multiframe, count):
             
     return audio_bytes
 
+# Define the custom token prefix
+CUSTOM_TOKEN_PREFIX = "<custom_token_"
+
+# Use a single global cache for token processing
+token_id_cache = {}
+MAX_CACHE_SIZE = 10000  # Increased cache size for better performance
+
 def turn_token_into_id(token_string, index):
-    """Optimized token-to-id conversion with early returns and minimal string operations"""
+    """
-    token_string = token_string.strip()
+    Optimized token-to-ID conversion with caching.
+    This is the definitive implementation used by both inference.py and speechpipe.py.
     
-    # Early return for obvious mismatches
+    Args:
-    if "<custom_token_" not in token_string:
+        token_string: The token string to convert
+        index: Position index used for token offset calculation
+        
+    Returns:
+        int: Token ID if valid, None otherwise
+    """
+    # Check cache first (significant speedup for repeated tokens)
+    cache_key = (token_string, index % 7)
+    if cache_key in token_id_cache:
+        return token_id_cache[cache_key]
+        
+    # Early rejection for obvious non-matches
+    if CUSTOM_TOKEN_PREFIX not in token_string:
         return None
+        
+    # Process token
+    token_string = token_string.strip()
+    last_token_start = token_string.rfind(CUSTOM_TOKEN_PREFIX)
     
-    # Find the last token in the string
-    last_token_start = token_string.rfind("<custom_token_")
     if last_token_start == -1:
         return None
     
-    # Check if the token ends properly
+    last_token = token_string[last_token_start:]
-    if not token_string.endswith(">"):
+    
+    if not (last_token.startswith(CUSTOM_TOKEN_PREFIX) and last_token.endswith(">")):
         return None
         
     try:
-        # Extract and convert the number directly
+        number_str = last_token[14:-1]
-        number_str = token_string[last_token_start+14:-1]
+        token_id = int(number_str) - 10 - ((index % 7) * 4096)
-        return int(number_str) - 10 - ((index % 7) * 4096)
+        
+        # Cache the result if it's valid
+        if len(token_id_cache) < MAX_CACHE_SIZE:
+            token_id_cache[cache_key] = token_id
+            
+        return token_id
     except (ValueError, IndexError):
         return None
-# Cache for frequently processed tokens to avoid redundant computation
-token_cache = {}
-MAX_CACHE_SIZE = 1000  # Limit cache size to prevent memory bloat
 
 async def tokens_decoder(token_gen):
-    """Optimized token decoder with reliable end-of-buffer handling for complete audio generation"""
+    """Optimized token decoder with early first-chunk processing for lower latency"""
     buffer = []
     count = 0
-    # Use a smaller minimum frame requirement to allow more flexible processing
+    
-    min_frames_required = 28  # Lower requirement (4 chunks of 7 tokens)
+    # Track if first chunk has been processed
-    ideal_frames = 49  # Ideal standard frame size (7×7 window)
+    first_chunk_processed = False
-    process_every_n = 7  # Process every 7 tokens (standard for Orpheus model)
+    
+    # Use different thresholds for first chunk vs. subsequent chunks
+    min_frames_first = 7  # Just one chunk (7 tokens) for first audio - ultra-low latency
+    min_frames_subsequent = 28  # Standard minimum (4 chunks of 7 tokens) after first audio
+    ideal_frames = 49  # Ideal standard frame size (7×7 window) - unchanged
+    process_every_n = 7  # Process every 7 tokens (standard for Orpheus model) - unchanged
     
     start_time = time.time()
     token_count = 0
@@ -160,15 +206,8 @@ async def tokens_decoder(token_gen):
     async for token_sim in token_gen:
         token_count += 1
         
-        # Check cache first to avoid redundant computation
+        # Use the unified turn_token_into_id which already handles caching
-        cache_key = (token_sim, count % 7)
+        token = turn_token_into_id(token_sim, count)
-        if cache_key in token_cache:
-            token = token_cache[cache_key]
-        else:
-            token = turn_token_into_id(token_sim, count)
-            # Add to cache if valid token
-            if token is not None and len(token_cache) < MAX_CACHE_SIZE:
-                token_cache[cache_key] = token
         
         if token is not None and token > 0:
             buffer.append(token)
@@ -185,26 +224,37 @@ async def tokens_decoder(token_gen):
                 last_log_time = current_time
                 token_count = 0
             
-            # Process standard batches when we have enough tokens
+            # Different processing logic based on whether first chunk has been processed
-            if count % process_every_n == 0:
+            if not first_chunk_processed:
-                # Best case: we have enough for the ideal frame size
+                # Process first chunk as soon as possible for minimal latency
-                if len(buffer) >= ideal_frames:
+                if count >= min_frames_first:
-                    buffer_to_proc = buffer[-ideal_frames:]
+                    buffer_to_proc = buffer[-min_frames_first:]
-                # Fallback: we have enough for the minimum requirement
+                    
-                elif len(buffer) >= min_frames_required:
+                    # Process the first chunk of audio for immediate feedback
-                    buffer_to_proc = buffer[-min_frames_required:]
+                    print(f"Processing first audio chunk with {len(buffer_to_proc)} tokens for low latency")
-                # For the first few frames, we may not have enough yet
+                    audio_samples = convert_to_audio(buffer_to_proc, count)
-                else:
+                    if audio_samples is not None:
-                    continue
+                        first_chunk_processed = True  # Mark first chunk as processed
-                
+                        yield audio_samples
-                # Debug output to help diagnose issues
+            else:
-                if count % 28 == 0:
+                # For subsequent chunks, use original processing with proper batching
-                    print(f"Processing buffer with {len(buffer_to_proc)} tokens, total collected: {len(buffer)}")
+                if count % process_every_n == 0:
-                
+                    # Use same prioritization logic as before
-                # Process the tokens
+                    if len(buffer) >= ideal_frames:
-                audio_samples = convert_to_audio(buffer_to_proc, count)
+                        buffer_to_proc = buffer[-ideal_frames:]
-                if audio_samples is not None:
+                    elif len(buffer) >= min_frames_subsequent:
-                    yield audio_samples
+                        buffer_to_proc = buffer[-min_frames_subsequent:]
+                    else:
+                        continue
+                    
+                    # 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
     
     # CRITICAL: End-of-generation handling - process all remaining frames
     # Process remaining complete frames (ideal size)
@@ -215,8 +265,8 @@ async def tokens_decoder(token_gen):
             yield audio_samples
             
     # Process any additional complete frames (minimum size)
-    elif len(buffer) >= min_frames_required:
+    elif len(buffer) >= min_frames_subsequent:
-        buffer_to_proc = buffer[-min_frames_required:]
+        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
@@ -227,7 +277,7 @@ async def tokens_decoder(token_gen):
         # Pad to minimum frame requirement with copies of the final token
         # This is more continuous than using unrelated tokens from the beginning
         last_token = buffer[-1]
-        padding_needed = min_frames_required - len(buffer)
+        padding_needed = min_frames_subsequent - len(buffer)
         
         # Create a padding array of copies of the last token
         # This maintains continuity much better than circular buffering