streaming_server.py

#!/usr/bin/env python3
"""
Orpheus TTS Streaming Server

High-performance TTS server with separate LLM and SNAC decoder processes.
Optimized for RTX 4090 with FP8 TensorRT-LLM.

Architecture:
    Client Request → LLM Process (FP8) → Token Stream → SNAC Decoder (FP16) → Audio Stream

This design achieves:
- Minimal TTFB (~150ms target)
- 8-12 concurrent real-time streams on RTX 4090
- No blocking between LLM and audio decoding

Usage:
    python streaming_server.py --config config/model_config.yaml
"""

import os
import sys
import time
import uuid
import yaml
import asyncio
import logging
from pathlib import Path
from typing import Optional, List, Dict, Any, AsyncGenerator
from contextlib import asynccontextmanager
from dataclasses import dataclass
import io
import wave

import numpy as np
import uvicorn
from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import StreamingResponse, Response
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

# Add src to path
sys.path.insert(0, str(Path(__file__).parent / "src"))

from llm_engine import OrpheusLLMEngine, LLMProcessManager, GenerationConfig
from snac_decoder import SNACDecoder, StreamingAudioDecoder, SAMPLE_RATE

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


# Global components
llm_manager: Optional[LLMProcessManager] = None
audio_decoder: Optional[StreamingAudioDecoder] = None
config: Dict[str, Any] = {}


def load_config(config_path: str) -> Dict[str, Any]:
    """Load configuration."""
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan manager."""
    global llm_manager, audio_decoder, config
    
    config_path = os.environ.get("CONFIG_PATH", "config/model_config.yaml")
    config = load_config(config_path)
    
    logger.info("=" * 60)
    logger.info("Orpheus TTS Streaming Server")
    logger.info("RTX 4090 Optimized (FP8 LLM + FP16 SNAC)")
    logger.info("=" * 60)
    
    model_config = config['model']
    server_config = config['server']
    
    # Initialize LLM process manager
    logger.info("Initializing LLM engine...")
    llm_manager = LLMProcessManager(
        engine_dir=model_config['trt_engine_path'],
        tokenizer_path=model_config['hf_model_path'],
        max_concurrent=server_config['max_concurrent_requests'],
    )
    
    try:
        await llm_manager.initialize()
        logger.info("✓ LLM engine ready")
    except Exception as e:
        logger.error(f"✗ LLM engine failed: {e}")
        llm_manager = None
    
    # Initialize SNAC decoder
    logger.info("Initializing SNAC decoder...")
    snac_path = str(Path(model_config['hf_model_path']).parent / "snac_24khz")
    audio_decoder = StreamingAudioDecoder(
        snac_path=snac_path,
        buffer_frames=2,
    )
    
    try:
        audio_decoder.decoder.load()
        logger.info("✓ SNAC decoder ready")
    except Exception as e:
        logger.error(f"✗ SNAC decoder failed: {e}")
        audio_decoder = None
    
    logger.info("=" * 60)
    logger.info(f"Server ready on {server_config['host']}:{server_config['port']}")
    logger.info(f"Max concurrent: {server_config['max_concurrent_requests']}")
    logger.info("=" * 60)
    
    yield
    
    # Cleanup
    if llm_manager:
        await llm_manager.shutdown()
    logger.info("Server shutdown complete")


# Create FastAPI app
app = FastAPI(
    title="Orpheus TTS Streaming Server",
    description="High-performance TTS with FP8 TensorRT-LLM",
    version="2.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# Request models
class TTSRequest(BaseModel):
    """TTS synthesis request."""
    prompt: str = Field(..., description="Text to synthesize")
    voice: str = Field(default="tara")
    max_tokens: int = Field(default=4096)
    temperature: float = Field(default=0.6, ge=0.0, le=2.0)
    top_p: float = Field(default=0.95, ge=0.0, le=1.0)
    top_k: int = Field(default=50, ge=1)
    repetition_penalty: float = Field(default=1.1, ge=1.0, le=2.0)
    stop_token_ids: List[int] = Field(default=[128258, 128009])
    request_id: Optional[str] = None
    streaming: bool = Field(default=True)


class HealthResponse(BaseModel):
    status: str
    llm_ready: bool
    snac_ready: bool
    active_requests: int
    gpu_memory_used_gb: float


# Helper function to create WAV header for streaming
def create_wav_header(sample_rate: int = 24000, channels: int = 1, bits: int = 16) -> bytes:
    """Create WAV header for streaming (unknown length)."""
    # Use maximum possible length placeholder
    data_size = 0x7FFFFFFF
    file_size = data_size + 36
    
    header = b'RIFF'
    header += file_size.to_bytes(4, 'little')
    header += b'WAVE'
    header += b'fmt '
    header += (16).to_bytes(4, 'little')  # Subchunk1Size
    header += (1).to_bytes(2, 'little')   # AudioFormat (PCM)
    header += channels.to_bytes(2, 'little')
    header += sample_rate.to_bytes(4, 'little')
    header += (sample_rate * channels * bits // 8).to_bytes(4, 'little')  # ByteRate
    header += (channels * bits // 8).to_bytes(2, 'little')  # BlockAlign
    header += bits.to_bytes(2, 'little')
    header += b'data'
    header += data_size.to_bytes(4, 'little')
    
    return header


@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint."""
    import torch
    
    gpu_memory = 0
    if torch.cuda.is_available():
        gpu_memory = torch.cuda.memory_allocated() / (1024**3)
    
    return HealthResponse(
        status="healthy" if llm_manager and audio_decoder else "degraded",
        llm_ready=llm_manager is not None,
        snac_ready=audio_decoder is not None,
        active_requests=llm_manager.active_requests if llm_manager else 0,
        gpu_memory_used_gb=round(gpu_memory, 2),
    )


@app.get("/v1/voices")
async def list_voices():
    """List available voices."""
    return {
        "voices": config.get('voices', {}).get('available', []),
        "default": config.get('voices', {}).get('default', 'tara'),
    }


@app.get("/v1/stats")
async def get_stats():
    """Get server statistics."""
    if not llm_manager:
        raise HTTPException(503, "Server not ready")
    
    stats = llm_manager.engine.get_stats()
    stats["active_requests"] = llm_manager.active_requests
    stats["max_concurrent"] = llm_manager.max_concurrent
    
    return stats


@app.post("/v1/audio/speech")
async def synthesize_speech(request: TTSRequest):
    """
    Synthesize speech from text with streaming.
    
    The audio is generated in real-time as tokens are produced,
    providing minimal time-to-first-byte.
    """
    if not llm_manager or not audio_decoder:
        raise HTTPException(503, "Server not ready")
    
    request_id = request.request_id or str(uuid.uuid4())
    start_time = time.perf_counter()
    
    logger.info(f"[{request_id[:8]}] Request: {len(request.prompt)} chars, voice={request.voice}")
    
    gen_config = GenerationConfig(
        max_tokens=request.max_tokens,
        temperature=request.temperature,
        top_p=request.top_p,
        top_k=request.top_k,
        repetition_penalty=request.repetition_penalty,
        stop_token_ids=request.stop_token_ids,
    )
    
    if request.streaming:
        async def audio_stream():
            """Generate audio stream from tokens."""
            first_audio = True
            total_samples = 0
            
            # Send WAV header first
            yield create_wav_header(SAMPLE_RATE)
            
            # Start streaming decoder
            decoder = StreamingAudioDecoder(
                snac_path=str(Path(config['model']['hf_model_path']).parent / "snac_24khz"),
            )
            decoder.start()
            
            try:
                async for token_chunk in llm_manager.generate(
                    request.prompt,
                    request.voice,
                    gen_config,
                    streaming=True,
                ):
                    # Decode tokens to audio
                    audio = decoder.add_tokens(token_chunk)
                    
                    if audio is not None and len(audio) > 0:
                        if first_audio:
                            ttfb = (time.perf_counter() - start_time) * 1000
                            logger.info(f"[{request_id[:8]}] TTFB: {ttfb:.1f}ms")
                            first_audio = False
                        
                        # Convert to 16-bit PCM
                        audio_int16 = (np.clip(audio, -1, 1) * 32767).astype(np.int16)
                        yield audio_int16.tobytes()
                        total_samples += len(audio)
                
                # Flush remaining audio
                decoder.stop()
                
            except Exception as e:
                logger.error(f"[{request_id[:8]}] Stream error: {e}")
                raise
            
            total_time = time.perf_counter() - start_time
            audio_duration = total_samples / SAMPLE_RATE
            rtf = audio_duration / total_time if total_time > 0 else 0
            
            logger.info(
                f"[{request_id[:8]}] Complete: {audio_duration:.1f}s audio "
                f"in {total_time:.2f}s (RTF: {rtf:.1f}x)"
            )
        
        return StreamingResponse(
            audio_stream(),
            media_type="audio/wav",
            headers={
                "X-Request-ID": request_id,
                "Transfer-Encoding": "chunked",
            }
        )
    
    else:
        # Non-streaming: generate all at once
        try:
            all_tokens = []
            async for chunk in llm_manager.generate(
                request.prompt,
                request.voice,
                gen_config,
                streaming=False,
            ):
                all_tokens.extend(chunk)
            
            # Decode all tokens
            audio = audio_decoder.decoder.decode(all_tokens)
            
            # Convert to WAV
            audio_int16 = (np.clip(audio, -1, 1) * 32767).astype(np.int16)
            
            buffer = io.BytesIO()
            with wave.open(buffer, 'wb') as wav:
                wav.setnchannels(1)
                wav.setsampwidth(2)
                wav.setframerate(SAMPLE_RATE)
                wav.writeframes(audio_int16.tobytes())
            
            buffer.seek(0)
            audio_bytes = buffer.read()
            
            total_time = time.perf_counter() - start_time
            logger.info(f"[{request_id[:8]}] Complete in {total_time:.2f}s")
            
            return Response(
                content=audio_bytes,
                media_type="audio/wav",
                headers={
                    "X-Request-ID": request_id,
                    "X-Generation-Time": str(total_time),
                }
            )
            
        except Exception as e:
            logger.error(f"[{request_id[:8]}] Error: {e}")
            raise HTTPException(500, str(e))


@app.post("/predict")
async def predict(request: TTSRequest):
    """Baseten-compatible predict endpoint."""
    return await synthesize_speech(request)


@app.websocket("/v1/audio/stream")
async def websocket_stream(websocket: WebSocket):
    """
    WebSocket endpoint for real-time bidirectional streaming.
    
    Lower latency than HTTP for real-time voice applications.
    """
    await websocket.accept()
    
    if not llm_manager or not audio_decoder:
        await websocket.close(1011, "Server not ready")
        return
    
    request_id = str(uuid.uuid4())
    logger.info(f"[{request_id[:8]}] WebSocket connected")
    
    try:
        while True:
            data = await websocket.receive_json()
            
            text = data.get("text", "")
            voice = data.get("voice", "tara")
            
            if not text:
                await websocket.send_json({"error": "No text provided"})
                continue
            
            start_time = time.perf_counter()
            first_audio = True
            
            gen_config = GenerationConfig(
                max_tokens=data.get("max_tokens", 4096),
                temperature=data.get("temperature", 0.6),
                top_p=data.get("top_p", 0.95),
            )
            
            decoder = StreamingAudioDecoder(
                snac_path=str(Path(config['model']['hf_model_path']).parent / "snac_24khz"),
            )
            decoder.start()
            
            try:
                async for token_chunk in llm_manager.generate(text, voice, gen_config):
                    audio = decoder.add_tokens(token_chunk)
                    
                    if audio is not None and len(audio) > 0:
                        if first_audio:
                            ttfb = (time.perf_counter() - start_time) * 1000
                            await websocket.send_json({"ttfb_ms": ttfb})
                            first_audio = False
                        
                        audio_int16 = (np.clip(audio, -1, 1) * 32767).astype(np.int16)
                        await websocket.send_bytes(audio_int16.tobytes())
                
                decoder.stop()
                await websocket.send_json({"status": "done"})
                
            except Exception as e:
                await websocket.send_json({"error": str(e)})
            
            total_time = time.perf_counter() - start_time
            logger.info(f"[{request_id[:8]}] WebSocket request complete in {total_time:.2f}s")
            
    except WebSocketDisconnect:
        logger.info(f"[{request_id[:8]}] WebSocket disconnected")
    except Exception as e:
        logger.error(f"[{request_id[:8]}] WebSocket error: {e}")


def main():
    """Run the server."""
    import argparse
    
    parser = argparse.ArgumentParser(description="Orpheus TTS Streaming Server")
    parser.add_argument("--config", default="config/model_config.yaml")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8000)
    
    args = parser.parse_args()
    
    os.environ["CONFIG_PATH"] = args.config
    
    uvicorn.run(
        "streaming_server:app",
        host=args.host,
        port=args.port,
        log_level="info",
    )


if __name__ == "__main__":
    main()