inference_server.py

#!/usr/bin/env python3
"""High-concurrency async TTS inference server with pre-computed speaker embeddings.

Uses vLLM-Omni's AsyncOmni engine for concurrent request handling with the
two-stage Qwen3-TTS pipeline (AR Talker -> Code2Wav). Pre-computed speaker
embeddings are loaded into GPU memory at startup for zero-overhead voice cloning.

Usage:
    # Start the server
    python inference_server.py

    # With custom settings
    TTS_MODEL_PATH=Qwen/Qwen3-TTS-12Hz-1.7B-Base \\
    TTS_PROFILES_DIR=voice_profiles \\
    TTS_PORT=8090 \\
    python inference_server.py

API:
    POST /synthesize       - Generate audio from text + speaker_id
    POST /synthesize/stream - Streaming audio generation
    GET  /speakers         - List available speaker profiles
    GET  /health           - Health check
"""

from __future__ import annotations

import asyncio
import base64
import io
import logging
import os
import struct
import sys
import time
import uuid
from contextlib import asynccontextmanager
from typing import Any, Optional

import numpy as np
import soundfile as sf
import torch
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, Response, StreamingResponse
from pydantic import BaseModel, Field

# vLLM-Omni requires spawn for multi-process workers.
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")

import inference_config as cfg
from embedding_cache import EmbeddingCacheManager

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Request / Response models
# ---------------------------------------------------------------------------


class SynthesizeRequest(BaseModel):
    speaker_id: str = Field(..., description="Speaker profile ID (filename stem)")
    text: str = Field(..., min_length=1, description="Text to synthesize")
    language: str = Field(default="Auto", description="Language: Auto, Chinese, English, etc.")
    response_format: str = Field(default="wav", description="Output format: wav, pcm, mp3, flac")
    max_new_tokens: int = Field(default=2048, ge=1, le=4096, description="Max codec tokens")
    stream: bool = Field(default=False, description="Enable streaming audio output")


class SpeakerInfo(BaseModel):
    speaker_id: str
    embedding_dim: int
    metadata: dict[str, str]


class HealthResponse(BaseModel):
    status: str
    model: str
    speakers_loaded: int
    cache_stats: dict


# ---------------------------------------------------------------------------
# Prompt length estimation
# ---------------------------------------------------------------------------

_ESTIMATOR_CACHE: dict[str, Any] = {}


def _estimate_prompt_len(additional_information: dict[str, Any], model_name: str) -> int:
    """Estimate placeholder prompt_token_ids length for the AR Talker."""
    try:
        from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
            Qwen3TTSConfig,
        )
        from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
            Qwen3TTSTalkerForConditionalGeneration,
        )

        if model_name not in _ESTIMATOR_CACHE:
            from transformers import AutoTokenizer

            tok = AutoTokenizer.from_pretrained(
                model_name, trust_remote_code=True, padding_side="left"
            )
            config = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
            _ESTIMATOR_CACHE[model_name] = (tok, getattr(config, "talker_config", None))

        tok, tcfg = _ESTIMATOR_CACHE[model_name]
        task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
        return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
            additional_information=additional_information,
            task_type=task_type,
            tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
            codec_language_id=getattr(tcfg, "codec_language_id", None),
            spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
        )
    except Exception as exc:
        logger.warning("Prompt length estimation failed, using fallback 2048: %s", exc)
        return 2048


# ---------------------------------------------------------------------------
# Audio encoding helpers
# ---------------------------------------------------------------------------


def encode_wav(audio_np: np.ndarray, sr: int) -> bytes:
    """Encode float32 audio to WAV bytes."""
    buf = io.BytesIO()
    sf.write(buf, audio_np, samplerate=sr, format="WAV")
    return buf.getvalue()


def encode_pcm(audio_np: np.ndarray) -> bytes:
    """Encode float32 audio to raw 16-bit PCM bytes."""
    pcm_int16 = np.clip(audio_np * 32767.0, -32768, 32767).astype(np.int16)
    return pcm_int16.tobytes()


def encode_audio(audio_np: np.ndarray, sr: int, fmt: str) -> tuple[bytes, str]:
    """Encode audio to the requested format. Returns (bytes, media_type)."""
    if fmt == "pcm":
        return encode_pcm(audio_np), "audio/pcm"
    elif fmt == "mp3":
        buf = io.BytesIO()
        sf.write(buf, audio_np, samplerate=sr, format="MP3", subtype=None)
        return buf.getvalue(), "audio/mpeg"
    elif fmt == "flac":
        buf = io.BytesIO()
        sf.write(buf, audio_np, samplerate=sr, format="FLAC")
        return buf.getvalue(), "audio/flac"
    else:  # wav (default)
        return encode_wav(audio_np, sr), "audio/wav"


# ---------------------------------------------------------------------------
# Core generation logic
# ---------------------------------------------------------------------------


def build_prompt(
    text: str,
    embedding: torch.Tensor,
    language: str,
    max_new_tokens: int,
) -> dict[str, Any]:
    """Build an Omni prompt dict with a pre-computed speaker embedding."""
    additional_information: dict[str, Any] = {
        "task_type": ["Base"],
        "text": [text],
        "language": [language],
        "x_vector_only_mode": [True],
        "non_streaming_mode": [True],  # All text in prefill → fewer decode ops
        # Top-level tensor key: bypasses voice_clone_prompt serialization issue.
        "ref_spk_embedding": embedding,
        "max_new_tokens": [max_new_tokens],
    }
    ph_len = _estimate_prompt_len(additional_information, cfg.MODEL_PATH)
    return {
        "prompt_token_ids": [0] * ph_len,
        "additional_information": additional_information,
    }


def extract_audio_from_output(output: Any) -> tuple[np.ndarray, int]:
    """Extract (audio_numpy, sample_rate) from an OmniRequestOutput."""
    # Navigate to multimodal output depending on output structure.
    mm = None
    if hasattr(output, "request_output") and output.request_output is not None:
        ro = output.request_output
        if hasattr(ro, "outputs") and ro.outputs:
            mm = ro.outputs[0].multimodal_output
        elif hasattr(ro, "multimodal_output"):
            mm = ro.multimodal_output
    if mm is None and hasattr(output, "multimodal_output"):
        mm = output.multimodal_output
    if not mm:
        raise ValueError("No multimodal output in response")

    audio_key = "audio" if "audio" in mm else "model_outputs"
    audio_tensor = mm[audio_key]
    sr_val = mm.get("sr", 24000)
    if hasattr(sr_val, "item"):
        sr = sr_val.item()
    elif isinstance(sr_val, list):
        sr = int(sr_val[-1])
    else:
        sr = int(sr_val)

    if isinstance(audio_tensor, list):
        audio_tensor = torch.cat(audio_tensor, dim=-1)

    audio_np = audio_tensor.float().detach().cpu().numpy()
    if audio_np.ndim > 1:
        audio_np = audio_np.flatten()

    return audio_np, sr


# ---------------------------------------------------------------------------
# Server lifecycle and app
# ---------------------------------------------------------------------------


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Server startup / shutdown lifecycle."""
    # ── Startup ───────────────────────────────────────────────────────
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )
    logger.info("Starting TTS inference server...")

    # 1. CUDA performance flags
    if torch.cuda.is_available():
        torch.backends.cudnn.benchmark = True
        torch.backends.cuda.matmul.allow_tf32 = True

    # 2. Load embedding cache
    logger.info("Loading speaker embeddings from %s", cfg.PROFILES_DIR)
    cache = EmbeddingCacheManager(
        profiles_dir=cfg.PROFILES_DIR,
        device="cpu",  # Keep on CPU for serialization; vLLM handles GPU transfer.
    )
    n = cache.preload_all()
    if n == 0:
        logger.warning("No speaker embeddings found in %s", cfg.PROFILES_DIR)

    # 3. Initialize AsyncOmni engine
    from vllm_omni import AsyncOmni

    logger.info("Initializing AsyncOmni with model=%s", cfg.MODEL_PATH)
    engine = AsyncOmni(
        model=cfg.MODEL_PATH,
        stage_configs_path=cfg.STAGE_CONFIG_PATH,
        stage_init_timeout=cfg.STAGE_INIT_TIMEOUT,
    )

    # 4. Concurrency semaphore for backpressure
    semaphore = asyncio.Semaphore(cfg.MAX_CONCURRENT_REQUESTS)

    # 5. Request counter for unique IDs
    app.state.cache = cache
    app.state.engine = engine
    app.state.semaphore = semaphore
    app.state.request_count = 0
    app.state.active_requests = 0
    app.state.total_audio_seconds = 0.0
    app.state.total_latency_seconds = 0.0

    logger.info(
        "Server ready: %d speakers loaded, max_concurrent=%d",
        n,
        cfg.MAX_CONCURRENT_REQUESTS,
    )

    yield

    # ── Shutdown ──────────────────────────────────────────────────────
    logger.info("Shutting down server...")
    engine.shutdown()


app = FastAPI(
    title="Qwen3-TTS Inference Server",
    description="High-concurrency TTS with pre-computed speaker embeddings",
    version="1.0.0",
    lifespan=lifespan,
)


# ---------------------------------------------------------------------------
# API endpoints
# ---------------------------------------------------------------------------


@app.post("/synthesize")
async def synthesize(req: SynthesizeRequest):
    """Generate audio from text using a pre-computed speaker embedding.

    The speaker embedding is looked up from the cache by speaker_id,
    injected as a top-level tensor in additional_information, and passed
    through the vLLM-Omni two-stage pipeline (Talker -> Code2Wav).
    """
    start_time = time.perf_counter()

    # Validate speaker exists
    cache: EmbeddingCacheManager = app.state.cache
    try:
        embedding = cache.get(req.speaker_id)
    except KeyError as e:
        raise HTTPException(status_code=404, detail=str(e))

    if not req.text.strip():
        raise HTTPException(status_code=400, detail="Text cannot be empty")

    # Backpressure
    semaphore: asyncio.Semaphore = app.state.semaphore
    if semaphore.locked():
        logger.warning("Server at capacity (%d concurrent)", cfg.MAX_CONCURRENT_REQUESTS)

    async with semaphore:
        app.state.active_requests += 1
        app.state.request_count += 1
        request_id = f"tts-{app.state.request_count}-{uuid.uuid4().hex[:8]}"

        try:
            # Build prompt
            prompt = build_prompt(
                text=req.text,
                embedding=embedding,
                language=req.language,
                max_new_tokens=req.max_new_tokens,
            )

            # Generate via AsyncOmni
            engine = app.state.engine
            final_output = None

            async for output in engine.generate(
                prompt=prompt,
                request_id=request_id,
                sampling_params_list=None,
                output_modalities=["audio"],
            ):
                final_output = output

            if final_output is None:
                raise HTTPException(status_code=500, detail="No output from model")

            # Extract audio
            audio_np, sr = extract_audio_from_output(final_output)

            # Encode to requested format
            audio_bytes, media_type = encode_audio(audio_np, sr, req.response_format)

            # Update stats
            elapsed = time.perf_counter() - start_time
            audio_duration = len(audio_np) / sr
            app.state.total_audio_seconds += audio_duration
            app.state.total_latency_seconds += elapsed

            logger.info(
                "Request %s: %.1fs latency, %.1fs audio, speaker=%s, fmt=%s",
                request_id,
                elapsed,
                audio_duration,
                req.speaker_id,
                req.response_format,
            )

            return Response(
                content=audio_bytes,
                media_type=media_type,
                headers={
                    "X-Request-Id": request_id,
                    "X-Audio-Duration": f"{audio_duration:.3f}",
                    "X-Latency": f"{elapsed:.3f}",
                    "X-Sample-Rate": str(sr),
                },
            )

        except HTTPException:
            raise
        except asyncio.CancelledError:
            await engine.abort(request_id)
            raise HTTPException(status_code=499, detail="Client disconnected")
        except Exception as e:
            logger.exception("Generation failed for %s: %s", request_id, e)
            raise HTTPException(status_code=500, detail=f"Generation failed: {e}")
        finally:
            app.state.active_requests -= 1


@app.post("/synthesize/stream")
async def synthesize_stream(req: SynthesizeRequest):
    """Streaming audio generation via chunked transfer encoding.

    Returns raw PCM float32 audio chunks as they are produced by the
    Code2Wav stage. The sample rate is in the X-Sample-Rate header.
    """
    cache: EmbeddingCacheManager = app.state.cache
    try:
        embedding = cache.get(req.speaker_id)
    except KeyError as e:
        raise HTTPException(status_code=404, detail=str(e))

    if not req.text.strip():
        raise HTTPException(status_code=400, detail="Text cannot be empty")

    semaphore: asyncio.Semaphore = app.state.semaphore

    async def audio_chunk_generator():
        async with semaphore:
            app.state.active_requests += 1
            app.state.request_count += 1
            request_id = f"tts-stream-{app.state.request_count}-{uuid.uuid4().hex[:8]}"

            try:
                prompt = build_prompt(
                    text=req.text,
                    embedding=embedding,
                    language=req.language,
                    max_new_tokens=req.max_new_tokens,
                )

                engine = app.state.engine
                async for output in engine.generate(
                    prompt=prompt,
                    request_id=request_id,
                    sampling_params_list=None,
                    output_modalities=["audio"],
                ):
                    try:
                        audio_np, sr = extract_audio_from_output(output)
                        # Yield raw PCM float32 bytes for minimal latency.
                        yield audio_np.astype(np.float32).tobytes()
                    except Exception:
                        continue

            except asyncio.CancelledError:
                await engine.abort(request_id)
            except Exception as e:
                logger.exception("Streaming failed for %s: %s", request_id, e)
            finally:
                app.state.active_requests -= 1

    return StreamingResponse(
        audio_chunk_generator(),
        media_type="application/octet-stream",
        headers={
            "X-Audio-Format": "pcm_f32le",
            "X-Sample-Rate": "24000",
            "Transfer-Encoding": "chunked",
        },
    )


@app.get("/speakers", response_model=list[SpeakerInfo])
async def list_speakers():
    """List all available speaker profiles with metadata."""
    cache: EmbeddingCacheManager = app.state.cache
    speakers = []
    for sid in cache.list_speakers():
        speakers.append(
            SpeakerInfo(
                speaker_id=sid,
                embedding_dim=1024,
                metadata=cache.get_metadata(sid),
            )
        )
    return speakers


@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check with cache and server statistics."""
    cache: EmbeddingCacheManager = app.state.cache
    count = getattr(app.state, "request_count", 0)
    total_latency = getattr(app.state, "total_latency_seconds", 0.0)
    stats = cache.stats()
    stats["active_requests"] = getattr(app.state, "active_requests", 0)
    stats["total_requests"] = count
    stats["avg_latency_s"] = total_latency / count if count > 0 else 0.0
    stats["total_audio_seconds"] = getattr(app.state, "total_audio_seconds", 0.0)

    return HealthResponse(
        status="ok",
        model=cfg.MODEL_PATH,
        speakers_loaded=len(cache),
        cache_stats=stats,
    )


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------


def main():
    import uvicorn

    uvicorn.run(
        "inference_server:app",
        host=cfg.HOST,
        port=cfg.PORT,
        log_level="info",
        access_log=True,
        workers=1,  # Single worker: AsyncOmni manages its own concurrency.
    )


if __name__ == "__main__":
    main()