Create vllm_streaming_inference.py

vllm_streaming_inference.py

@@ -0,0 +1,565 @@
+"""
+Maya-1-Voice VLLM Streaming Inference - Standalone Reference Implementation
+
+This is a complete, self-contained example for using Maya-1-Voice TTS model with VLLM and SNAC.
+Demonstrates streaming audio generation with sliding window approach for smooth playback.
+
+Requirements:
+    pip install vllm transformers torch snac numpy
+
+Usage:
+    python vllm_streaming_inference.py
+
+Author: Maya-1-Voice Team
+License: MIT
+"""
+
+import torch
+import numpy as np
+import asyncio
+from typing import List, Optional, AsyncGenerator
+from transformers import AutoTokenizer
+from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams
+from snac import SNAC
+
+
+# ============================================================================
+# CONSTANTS
+# ============================================================================
+
+# Special control tokens
+CODE_START_TOKEN_ID = 128257  # Start of Speech (SOS)
+CODE_END_TOKEN_ID = 128258    # End of Speech (EOS) - stop token for audio
+CODE_TOKEN_OFFSET = 128266    # Start of SNAC codes
+
+# SNAC token range (7 tokens per frame, 4096 codes per level)
+SNAC_MIN_ID = 128266
+SNAC_MAX_ID = 156937  # 128266 + (7 * 4096) - 1
+
+# SNAC configuration
+SNAC_MODEL_NAME = "hubertsiuzdak/snac_24khz"
+SNAC_SAMPLE_RATE = 24000
+SNAC_TOKENS_PER_FRAME = 7
+
+# Generation parameters
+DEFAULT_TEMPERATURE = 0.4
+DEFAULT_TOP_P = 0.9
+DEFAULT_MAX_TOKENS = 2000
+DEFAULT_MIN_TOKENS = 28  # At least 4 SNAC frames
+DEFAULT_REPETITION_PENALTY = 1.1
+
+
+# ============================================================================
+# SNAC DECODER
+# ============================================================================
+
+class SNACDecoder:
+    """
+    Decodes SNAC tokens (7-token frames) to audio waveforms.
+    
+    The unpacking logic converts flat 7-token frames back to hierarchical
+    3-level SNAC codes (matching the training preprocessing exactly).
+    """
+    
+    def __init__(self, device: str = "cuda"):
+        """Initialize SNAC decoder with 24kHz model."""
+        self.device = device
+        print(f"🎵 Loading SNAC 24kHz model to {device}...")
+        self.snac_model = SNAC.from_pretrained(SNAC_MODEL_NAME).eval().to(device)
+        print(f"✅ SNAC decoder initialized")
+    
+    def unpack_snac_from_7(self, vocab_ids: List[int]) -> List[List[int]]:
+        """
+        Unpack 7-token SNAC frames to 3 hierarchical levels.
+        
+        This is the EXACT INVERSE of training preprocessing.
+        
+        Frame structure (7 tokens per frame):
+        [slot0, slot1, slot2, slot3, slot4, slot5, slot6]
+        
+        Unpacking to [L1, L2, L3]:
+        - slot0 → L1[i]       (coarse: 1x rate)
+        - slot1 → L2[2*i]     (medium: 2x rate, even)
+        - slot2 → L3[4*i+0]   (fine: 4x rate)
+        - slot3 → L3[4*i+1]
+        - slot4 → L2[2*i+1]   (medium: odd)
+        - slot5 → L3[4*i+2]
+        - slot6 → L3[4*i+3]
+        
+        Args:
+            vocab_ids: List of SNAC token IDs (128266-156937), length divisible by 7
+        
+        Returns:
+            [L1, L2, L3] where L1=n, L2=2n, L3=4n elements
+        """
+        # Remove EOS token if present
+        if vocab_ids and vocab_ids[-1] == CODE_END_TOKEN_ID:
+            vocab_ids = vocab_ids[:-1]
+        
+        # Ensure complete frames
+        frames = len(vocab_ids) // SNAC_TOKENS_PER_FRAME
+        vocab_ids = vocab_ids[:frames * SNAC_TOKENS_PER_FRAME]
+        
+        if frames == 0:
+            return [[], [], []]
+        
+        l1, l2, l3 = [], [], []
+        
+        for i in range(frames):
+            slots = vocab_ids[i*7:(i+1)*7]
+            
+            # Subtract offset and mod 4096 to get original SNAC codes
+            l1.append((slots[0] - CODE_TOKEN_OFFSET) % 4096)
+            l2.extend([
+                (slots[1] - CODE_TOKEN_OFFSET) % 4096,  # Even
+                (slots[4] - CODE_TOKEN_OFFSET) % 4096,  # Odd
+            ])
+            l3.extend([
+                (slots[2] - CODE_TOKEN_OFFSET) % 4096,
+                (slots[3] - CODE_TOKEN_OFFSET) % 4096,
+                (slots[5] - CODE_TOKEN_OFFSET) % 4096,
+                (slots[6] - CODE_TOKEN_OFFSET) % 4096,
+            ])
+        
+        return [l1, l2, l3]
+    
+    @torch.inference_mode()
+    def decode(
+        self, 
+        snac_tokens: List[int], 
+        use_sliding_window: bool = False
+    ) -> Optional[np.ndarray]:
+        """
+        Decode SNAC tokens to audio waveform.
+        
+        Args:
+            snac_tokens: List of SNAC token IDs (7*n tokens)
+            use_sliding_window: If True, return only middle 2048 samples 
+                               (for smooth streaming without pops/clicks)
+        
+        Returns:
+            Audio waveform as float32 numpy array, 24kHz mono
+        """
+        if len(snac_tokens) < SNAC_TOKENS_PER_FRAME:
+            return None
+        
+        # Unpack to 3 hierarchical levels
+        levels = self.unpack_snac_from_7(snac_tokens)
+        
+        if not levels[0]:
+            return None
+        
+        # Convert to tensors
+        codes = [
+            torch.tensor(level, dtype=torch.long, device=self.device).unsqueeze(0)
+            for level in levels
+        ]
+        
+        # Decode through SNAC quantizer + decoder
+        z_q = self.snac_model.quantizer.from_codes(codes)
+        audio = self.snac_model.decoder(z_q)
+        
+        # Extract audio: [batch, 1, samples] → [samples]
+        audio = audio[0, 0].cpu().numpy()
+        
+        # Sliding window mode: keep middle 2048 samples only
+        # This eliminates popping/cracking in streaming by overlapping windows
+        if use_sliding_window and len(audio) >= 4096:
+            audio = audio[2048:4096]
+        
+        return audio
+    
+    def decode_to_bytes(
+        self, 
+        snac_tokens: List[int], 
+        use_sliding_window: bool = False
+    ) -> Optional[bytes]:
+        """
+        Decode SNAC tokens to audio bytes (int16 PCM).
+        
+        Args:
+            snac_tokens: List of SNAC token IDs
+            use_sliding_window: Use sliding window for smooth streaming
+        
+        Returns:
+            Audio as bytes (int16 PCM, 24kHz mono)
+        """
+        audio = self.decode(snac_tokens, use_sliding_window=use_sliding_window)
+        
+        if audio is None:
+            return None
+        
+        # Convert float32 to int16 PCM
+        audio_int16 = (audio * 32767).astype(np.int16)
+        return audio_int16.tobytes()
+
+
+# ============================================================================
+# CUSTOM LOGITS PROCESSOR
+# ============================================================================
+
+class OnlyAudioAfterSOS:
+    """
+    Restricts vocabulary to SNAC codes + EOS after SOS token.
+    
+    This prevents the model from generating text tokens during audio phase,
+    which would cause "hallucination" where the model repeats description text
+    instead of generating proper audio codes.
+    """
+    
+    def __init__(self):
+        self._seen_sos = False
+    
+    def __call__(
+        self,
+        prompt_token_ids: List[int],
+        generated_token_ids: List[int],
+        logits: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Apply constraint: after SOS, only allow SNAC codes + EOS.
+        
+        Args:
+            prompt_token_ids: Original prompt token IDs
+            generated_token_ids: Tokens generated so far
+            logits: Logits for next token [vocab_size]
+        
+        Returns:
+            Modified logits with masked tokens
+        """
+        # Check if SOS has been generated
+        if not self._seen_sos:
+            all_token_ids = prompt_token_ids + generated_token_ids
+            if CODE_START_TOKEN_ID in all_token_ids:
+                self._seen_sos = True
+            else:
+                return logits  # No constraint yet
+        
+        # Apply constraint: mask all tokens except SNAC codes + EOS
+        mask = torch.full_like(logits, float('-inf'))
+        mask[SNAC_MIN_ID:SNAC_MAX_ID + 1] = 0  # Allow SNAC codes
+        mask[CODE_END_TOKEN_ID] = 0            # Allow EOS
+        
+        return logits + mask
+    
+    def reset(self):
+        """Reset state for reuse across generations."""
+        self._seen_sos = False
+
+
+# ============================================================================
+# MAYA-1-VOICE MODEL
+# ============================================================================
+
+class Maya1VoiceModel:
+    """
+    Maya-1-Voice TTS Model with VLLM inference engine.
+    
+    Handles model loading, tokenizer initialization, and VLLM engine setup.
+    """
+    
+    def __init__(
+        self,
+        model_path: str,
+        dtype: str = "bfloat16",
+        max_model_len: int = 8192,
+        gpu_memory_utilization: float = 0.85,
+    ):
+        """
+        Initialize Maya-1-Voice model with VLLM.
+        
+        Args:
+            model_path: Path to model checkpoint (local or HuggingFace)
+            dtype: Model precision (bfloat16 recommended)
+            max_model_len: Maximum sequence length
+            gpu_memory_utilization: GPU memory fraction to use (0.0-1.0)
+        """
+        self.model_path = model_path
+        
+        print(f"🚀 Initializing Maya-1-Voice Model")
+        print(f"📁 Model: {model_path}")
+        print(f"🔢 Dtype: {dtype}")
+        
+        # Load tokenizer (must be from checkpoint with emotion tags)
+        print(f"📝 Loading tokenizer...")
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            trust_remote_code=True,
+        )
+        print(f"✅ Tokenizer loaded: {len(self.tokenizer)} tokens")
+        
+        # Initialize VLLM async engine
+        print(f"🔧 Initializing VLLM engine...")
+        engine_args = AsyncEngineArgs(
+            model=model_path,
+            tokenizer=model_path,
+            dtype=dtype,
+            max_model_len=max_model_len,
+            gpu_memory_utilization=gpu_memory_utilization,
+            trust_remote_code=True,
+        )
+        
+        self.engine = AsyncLLMEngine.from_engine_args(engine_args)
+        print(f"✅ VLLM engine ready")
+    
+    def build_prompt(self, description: str, text: str) -> str:
+        """
+        Build prompt in Maya-1-Voice format.
+        
+        Format: <description="..."> text
+        
+        The model expects:
+        1. Description of voice/character
+        2. Text to synthesize (optionally with <emotion> tags)
+        
+        Args:
+            description: Voice description 
+                Example: "Realistic male voice in the 30s age with american accent. 
+                         Normal pitch, warm timbre, conversational pacing."
+            text: Text to synthesize
+                Example: "Hello world! <excited> This is amazing!"
+        
+        Returns:
+            Formatted prompt string
+        """
+        return f'<description="{description}"> {text}'
+
+
+# ============================================================================
+# STREAMING PIPELINE
+# ============================================================================
+
+class Maya1VoiceStreamingPipeline:
+    """
+    Streaming TTS pipeline using sliding window approach.
+    
+    This generates smooth audio by:
+    1. Streaming tokens from VLLM as they're generated
+    2. Every 7 tokens, decoding the last 28 tokens (4 frames) - sliding window
+    3. Keeping only middle 2048 samples from each decode
+    4. Creating natural overlap between chunks for artifact-free playback
+    """
+    
+    def __init__(self, model: Maya1VoiceModel, snac_decoder: SNACDecoder):
+        """Initialize streaming pipeline."""
+        self.model = model
+        self.snac_decoder = snac_decoder
+        print(f"🌊 Maya-1-Voice Streaming Pipeline initialized")
+    
+    async def generate_speech_stream(
+        self,
+        description: str,
+        text: str,
+        temperature: float = DEFAULT_TEMPERATURE,
+        top_p: float = DEFAULT_TOP_P,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+        repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
+    ) -> AsyncGenerator[bytes, None]:
+        """
+        Generate speech audio with streaming.
+        
+        Args:
+            description: Voice/character description
+            text: Text to synthesize (with optional <emotion> tags)
+            temperature: Sampling temperature (lower = more stable)
+            top_p: Nucleus sampling
+            max_tokens: Max SNAC tokens to generate
+            repetition_penalty: Prevent repetition loops
+        
+        Yields:
+            Audio chunks as bytes (int16 PCM, 24kHz mono)
+        """
+        print(f"\n🌊 Starting streaming generation")
+        print(f"📝 Description: {description[:80]}...")
+        print(f"💬 Text: {text}")
+        
+        # Build prompt
+        prompt = self.model.build_prompt(description, text)
+        
+        # Configure sampling with custom logits processor
+        logits_processor = OnlyAudioAfterSOS()
+        
+        sampling_params = SamplingParams(
+            temperature=temperature,
+            top_p=top_p,
+            max_tokens=max_tokens,
+            min_tokens=DEFAULT_MIN_TOKENS,
+            repetition_penalty=repetition_penalty,
+            stop_token_ids=[CODE_END_TOKEN_ID],  # Stop on audio EOS
+            logits_processors=[logits_processor],  # Constrain to audio tokens
+        )
+        
+        print(f"🎲 Sampling: temp={temperature}, top_p={top_p}, max_tokens={max_tokens}")
+        
+        # Token buffer for sliding window
+        token_buffer = []
+        total_tokens = 0
+        total_chunks = 0
+        
+        # Generate with VLLM
+        import uuid
+        import time
+        request_id = f"maya1voice-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"
+        
+        results_generator = self.model.engine.generate(
+            prompt=prompt,
+            sampling_params=sampling_params,
+            request_id=request_id,
+        )
+        
+        # Stream tokens with sliding window decoding
+        async for request_output in results_generator:
+            generated_ids = request_output.outputs[0].token_ids
+            
+            # Process only new tokens
+            new_tokens = generated_ids[total_tokens:]
+            total_tokens = len(generated_ids)
+            
+            # Filter and buffer SNAC tokens only
+            for token_id in new_tokens:
+                if SNAC_MIN_ID <= token_id <= SNAC_MAX_ID:
+                    token_buffer.append(token_id)
+                    
+                    # Sliding window: process every 7 tokens when buffer > 27
+                    # Take last 28 tokens (4 frames) for smooth overlap
+                    if len(token_buffer) % 7 == 0 and len(token_buffer) > 27:
+                        window_tokens = token_buffer[-28:]
+                        
+                        # Decode with sliding window (returns middle 2048 samples)
+                        audio_bytes = self.snac_decoder.decode_to_bytes(
+                            window_tokens, 
+                            use_sliding_window=True
+                        )
+                        
+                        if audio_bytes:
+                            total_chunks += 1
+                            if total_chunks == 1:
+                                print(f"🎵 First chunk decoded ({len(audio_bytes)} bytes)")
+                            yield audio_bytes
+        
+        print(f"✅ Streaming complete: {total_tokens} tokens → {total_chunks} chunks")
+        
+        # Reset logits processor for next generation
+        logits_processor.reset()
+
+
+# ============================================================================
+# MAIN EXAMPLE
+# ============================================================================
+
+async def main():
+    """
+    Example usage of Maya-1-Voice streaming inference.
+    
+    This demonstrates:
+    1. Model initialization
+    2. SNAC decoder setup
+    3. Streaming generation
+    4. Audio chunk handling
+    """
+    
+    # Configuration
+    MODEL_PATH = "maya-research/maya-1-voice"  # Update with your model
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    print("=" * 80)
+    print("Maya-1-Voice VLLM Streaming Inference Example")
+    print("=" * 80)
+    
+    # Initialize model
+    model = Maya1VoiceModel(
+        model_path=MODEL_PATH,
+        dtype="bfloat16",
+        max_model_len=8192,
+        gpu_memory_utilization=0.85,
+    )
+    
+    # Initialize SNAC decoder
+    snac_decoder = SNACDecoder(device=DEVICE)
+    
+    # Create pipeline
+    pipeline = Maya1VoiceStreamingPipeline(model, snac_decoder)
+    
+    # Example 1: Professional voice
+    description = (
+        "Realistic male voice in the 30s age with american accent. "
+        "Normal pitch, warm timbre, conversational pacing, neutral tone delivery at med intensity."
+    )
+    text = "Hello! This is a test of the Maya-1-Voice text-to-speech system."
+    
+    print(f"\n{'='*80}")
+    print("Example 1: Professional Voice")
+    print(f"{'='*80}")
+    
+    audio_chunks = []
+    async for chunk in pipeline.generate_speech_stream(
+        description=description,
+        text=text,
+        temperature=0.4,
+        max_tokens=500,
+    ):
+        audio_chunks.append(chunk)
+        print(f"📦 Received chunk {len(audio_chunks)}: {len(chunk)} bytes")
+    
+    # Combine chunks
+    full_audio = b''.join(audio_chunks)
+    print(f"\n✅ Total audio: {len(full_audio)} bytes ({len(full_audio)//2} samples, {len(full_audio)/2/24000:.2f}s)")
+    
+    # Save audio (optional)
+    try:
+        import wave
+        output_file = "output_example1.wav"
+        with wave.open(output_file, 'wb') as wav:
+            wav.setnchannels(1)  # Mono
+            wav.setsampwidth(2)  # 16-bit
+            wav.setframerate(24000)  # 24kHz
+            wav.writeframes(full_audio)
+        print(f"💾 Saved to {output_file}")
+    except ImportError:
+        print(f"⚠️  Install 'wave' module to save audio files")
+    
+    # Example 2: Character voice with emotions
+    print(f"\n{'='*80}")
+    print("Example 2: Character Voice with Emotions")
+    print(f"{'='*80}")
+    
+    description = (
+        "Creative, dark_villain character. Male voice in their 40s with british accent. "
+        "Low pitch, gravelly timbre, slow pacing, angry tone at high intensity."
+    )
+    text = "The darkness isn't coming... <angry> it's already here!"
+    
+    audio_chunks = []
+    async for chunk in pipeline.generate_speech_stream(
+        description=description,
+        text=text,
+        temperature=0.5,
+        max_tokens=800,
+    ):
+        audio_chunks.append(chunk)
+        print(f"📦 Received chunk {len(audio_chunks)}: {len(chunk)} bytes")
+    
+    full_audio = b''.join(audio_chunks)
+    print(f"\n✅ Total audio: {len(full_audio)} bytes ({len(full_audio)//2} samples, {len(full_audio)/2/24000:.2f}s)")
+    
+    # Save audio
+    try:
+        import wave
+        output_file = "output_example2.wav"
+        with wave.open(output_file, 'wb') as wav:
+            wav.setnchannels(1)
+            wav.setsampwidth(2)
+            wav.setframerate(24000)
+            wav.writeframes(full_audio)
+        print(f"💾 Saved to {output_file}")
+    except ImportError:
+        pass
+    
+    print(f"\n{'='*80}")
+    print("🎉 Examples complete!")
+    print(f"{'='*80}")
+
+
+if __name__ == "__main__":
+    # Run async main
+    asyncio.run(main())