Output:
Chunk ID: 09af62
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 3756
Output:
                    temperature=req.temperature,
                    top_k=req.top_k,
                    top_p=req.top_p,
                    max_tokens=req.max_tokens,
                    repetition_penalty=req.repetition_penalty,
                    seed=req.seed,
                    sample_rate=req.sample_rate,
                    output_sample_rate=output_sr,
                )
            else:
                audio_bytes, metrics = await tts_runtime.generate_speech(
                    text=text,
                    speaker=speaker,
                    temperature=req.temperature,
                    top_k=req.top_k,
                    top_p=req.top_p,
                    max_tokens=req.max_tokens,
                    repetition_penalty=req.repetition_penalty,
                    seed=req.seed,
                    output_sample_rate=output_sr,
                )

            if audio_bytes is None:
                return JSONResponse(
                    status_code=500,
                    content={
                        "error": {
                            "code": "GENERATION_FAILED",
                            "message": "Audio generation failed. Check logs for details.",
                            "request_id": request_id,
                        }
                    },
                    headers={"X-Request-ID": request_id},
                )

            # Calculate RTF
            total_time = time.time() - start_time
            audio_duration = metrics.get("audio_duration_seconds", 0)
            rtf = total_time / audio_duration if audio_duration > 0 else 0

            # Build response headers
            runtime = tts_runtime.get_runtime()
            model_version = runtime.model_version if runtime else get_model_version()

            # Use actual output sample rate from metrics
            output_sample_rate = metrics.get("output_sample_rate", req.sample_rate)

            headers = {
                "X-Request-ID": request_id,
                "X-Model-Version": model_version,
                "X-Format": "wav",
                "X-Sample-Rate": str(output_sample_rate),
                "X-Stream": "false",
                "X-Audio-Bytes": str(len(audio_bytes)),
                "X-Audio-Seconds": f"{audio_duration:.2f}",
                "X-TTFB-ms": str(metrics.get("ttfb_ms", 0)),
                "X-RTF": f"{rtf:.3f}",
                "X-Text-Chunked": str(metrics.get("text_chunked", False)).lower(),
                "X-SR-Applied": str(metrics.get("sr_applied", False)).lower(),
            }

            if req.seed is not None:
                headers["X-Seed"] = str(req.seed)

            # Fire-and-forget sentence storage (non-blocking, after generation)
            sentence_store = _get_sentence_store()
            sentence_store.store_fire_and_forget(
                request_id=request_id,
                text=text,
                speaker=speaker,
                stream=False,
                format=req.format,
                temperature=req.temperature,
                top_k=req.top_k,
                top_p=req.top_p,
                max_tokens=req.max_tokens,
                repetition_penalty=req.repetition_penalty,
                seed=req.seed,
                text_chunked=metrics.get("text_chunked", False),
                ttfb_ms=metrics.get("ttfb_ms"),
                audio_duration_seconds=audio_duration,
            )

            # Log successful completion (lifecycle event + metrics)
            total_duration_ms = int(total_time * 1000)
            log_request_completed(
                request_id=request_id,
                status_code=200,
                total_duration_ms=total_duration_ms,
                audio_duration_seconds=audio_duration,
                rtf=rtf,
                chunks_sent=1,
            )
            metrics_request_completed(
                status_code=200,
                duration_seconds=total_time,
                speaker=speaker,
                stream=False,
            )
            record_ttfb(ttfb_seconds=metrics.get("ttfb_ms", 0) / 1000, speaker=speaker, stream=False)
            record_rtf(rtf=rtf, speaker=speaker)
            record_audio_duration(duration_seconds=audio_duration, speaker=speaker)

            clear_request_context()

            # Apply format encoding if not WAV
            media_type = "audio/wav"
            if target_format != AudioFormat.WAV.value:
                try:
                    from veena3modal.audio.encoder import (
                        OpusEncoder, MP3Encoder, MuLawEncoder, FLACEncoder
                    )

                    # Get sample rate from output (after potential SR)
                    sample_rate = metrics.get("output_sample_rate", 16000)

                    # Strip WAV header to get PCM data
                    pcm_data = audio_bytes[44:] if len(audio_bytes) > 44 else audio_bytes

                    if target_format == AudioFormat.OPUS.value:
                        encoder = OpusEncoder(sample_rate=sample_rate)
                        audio_bytes = encoder.encode(pcm_data, sample_rate)
                        media_type = "audio/opus"
                    elif target_format == AudioFormat.MP3.value:
                        encoder = MP3Encoder(sample_rate=sample_rate)
                        audio_bytes = encoder.encode(pcm_data, sample_rate)
                        media_type = "audio/mpeg"
                    elif target_format == AudioFormat.MULAW.value:
                        encoder = MuLawEncoder()
                        audio_bytes = encoder.encode(pcm_data, sample_rate)
                        media_type = "audio/x-wav"
                    elif target_format == AudioFormat.FLAC.value:
                        encoder = FLACEncoder(sample_rate=sample_rate)
                        audio_bytes = encoder.encode(pcm_data, sample_rate)
                        media_type = "audio/flac"

                    headers["X-Audio-Bytes"] = str(len(audio_bytes))
                    logger.info(f"Encoded audio to {target_format}: {len(audio_bytes)} bytes")
                except ImportError as e:
                    logger.warning(f"Audio encoder not available: {e}")
                    # Fall back to WAV
                except Exception as e:
                    logger.error(f"Format encoding failed: {e}")
                    # Fall back to WAV

            return Response(
                content=audio_bytes,
                media_type=media_type,
                headers=headers,
            )

        except Exception as e:
            logger.exception(f"TTS generation error: {e}")

            # Log failure (lifecycle event + metrics)
            log_request_failed(
                request_id=request_id,
                status_code=500,
                error_code="INTERNAL_ERROR",
                error_message=str(e),
            )
            metrics_request_failed(
                status_code=500,
                error_code="INTERNAL_ERROR",
                speaker=speaker,
            )
            clear_request_context()

            return JSONResponse(
                status_code=500,
                content={
                    "error": {
                        "code": "INTERNAL_ERROR",
                        "message": str(e),
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

    async def _handle_streaming_request(req, request_id: str, text: str, speaker: str, start_time: float):
        """
        Handle streaming TTS request.

        Returns a StreamingResponse that yields WAV header + PCM chunks.
        True streaming: first bytes sent ASAP before full audio is generated.
        """
        from fastapi.responses import StreamingResponse
        from veena3modal.api.schemas import AudioFormat
        from veena3modal.services import tts_runtime

        # Check if streaming pipeline is available
        runtime = tts_runtime.get_runtime()
        if runtime is None or runtime.streaming_pipeline is None:
            return JSONResponse(
                status_code=503,
                content={
                    "error": {
                        "code": "STREAMING_UNAVAILABLE",
                        "message": "Streaming pipeline not initialized.",
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

        # Non-WAV formats not implemented yet (M5)
        if req.format != AudioFormat.WAV.value:
            return JSONResponse(
                status_code=501,
                content={
                    "error": {
                        "code": "FORMAT_NOT_IMPLEMENTED",
                        "message": f"Streaming format '{req.format}' not yet implemented. Use format=wav.",
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

        # Set request context for structured logging
        set_request_context(request_id=request_id)

        # Create async generator for streaming response
        async def audio_stream_generator():
            """Yields audio chunks as they're generated."""
            final_metrics = {}
            first_chunk_sent = False
            chunks_count = 0
            try:
                async for audio_chunk, metrics in tts_runtime.generate_speech_streaming(
                    text=text,
                    speaker=speaker,
                    temperature=req.temperature,
                    top_k=req.top_k,
                    top_p=req.top_p,
                    max_tokens=req.max_tokens,
                    repetition_penalty=req.repetition_penalty,
                    seed=req.seed,
                    enable_chunking=req.chunking,
                ):
                    final_metrics = metrics
                    chunks_count += 1

                    # Fire-and-forget sentence storage after first chunk (don't block TTFB)
                    if not first_chunk_sent:
                        first_chunk_sent = True
                        ttfb_ms = metrics.get("ttfb_ms", 0)

                        # Log first audio emitted (TTFB marker)
                        log_first_audio_emitted(
                            request_id=request_id,
                            ttfb_ms=ttfb_ms,
                            chunk_size_bytes=len(audio_chunk),
                        )
                        record_ttfb(ttfb_seconds=ttfb_ms / 1000, speaker=speaker, stream=True)

                        sentence_store = _get_sentence_store()
                        sentence_store.store_fire_and_forget(
                            request_id=request_id,
                            text=text,
                            speaker=speaker,
                            stream=True,
                            format=req.format,
                            temperature=req.temperature,
                            top_k=req.top_k,
                            top_p=req.top_p,
                            max_tokens=req.max_tokens,
                            repetition_penalty=req.repetition_penalty,
                            seed=req.seed,
                            text_chunked=metrics.get("text_chunked", False),
                            ttfb_ms=ttfb_ms,
                            # audio_duration_seconds not known yet for streaming
                        )

                    yield audio_chunk

                # Stream completed - log metrics
                total_time = time.time() - start_time
                audio_duration = final_metrics.get("audio_duration_seconds", 0)
                rtf = total_time / audio_duration if audio_duration > 0 else 0

                log_request_completed(
                    request_id=request_id,
                    status_code=200,
                    total_duration_ms=int(total_time * 1000),
                    audio_duration_seconds=audio_duration,
                    rtf=rtf,
                    chunks_sent=chunks_count,
                )
                metrics_request_completed(
                    status_code=200,
                    duration_seconds=total_time,
                    speaker=speaker,
                    stream=True,
                )
                record_rtf(rtf=rtf, speaker=speaker)
                record_audio_duration(duration_seconds=audio_duration, speaker=speaker)
                record_chunks_sent(chunks=chunks_count, speaker=speaker)
                clear_request_context()
            except Exception as e:
                logger.exception(f"Streaming error: {e}")
                # Log streaming failure
                log_request_failed(
                    request_id=request_id,
                    status_code=500,
                    error_code="STREAMING_ERROR",
                    error_message=str(e),
                )
                metrics_request_failed(
                    status_code=500,
                    error_code="STREAMING_ERROR",
                    speaker=speaker,
                )
                clear_request_context()
                # Can't return error once streaming started
                # Best we can do is stop yielding
                return

        # Build response headers
        # Note: Some headers can't be set dynamically during streaming
        # TTFB and RTF will be logged but not in headers for streaming
        model_version = runtime.model_version if runtime else get_model_version()

        headers = {
            "X-Request-ID": request_id,
            "X-Model-Version": model_version,
            "X-Format": "wav",
            "X-Sample-Rate": str(req.sample_rate),
            "X-Stream": "true",
            "X-Chunking-Enabled": str(req.chunking).lower(),
            # Transfer-Encoding: chunked is set automatically by StreamingResponse
        }

        if req.seed is not None:
            headers["X-Seed"] = str(req.seed)

        return StreamingResponse(
            audio_stream_generator(),
            media_type="audio/wav",
            headers=headers,
        )

    # === WebSocket Support (M5b) ===
    try:
        from veena3modal.api.websocket_handler import add_websocket_routes
        add_websocket_routes(app)
        logger.info("WebSocket support enabled at /v1/tts/ws")
    except ImportError as e:
        logger.warning(f"WebSocket support not available: {e}")

    return app


Output:
Chunk ID: d0b779
Wall time: 0.0510 seconds
Process exited with code 0
Original token count: 3376
Output:
    repetition_penalty: float = 1.05,
    seed: Optional[int] = None,
    sample_rate: int = 16000,
    output_sample_rate: str = "16khz",
) -> Tuple[Optional[bytes], Dict[str, Any]]:
    """
    Generate speech with automatic text chunking for long inputs.

    Uses LongTextProcessor to split text and stitch audio.

    Args:
        text: Text to synthesize
        speaker: Internal speaker name
        ... (same as generate_speech)
        sample_rate: Internal generation sample rate (always 16000)
        output_sample_rate: "16khz" or "48khz" (triggers super-resolution)

    Returns:
        Tuple of (audio_bytes, metrics_dict)
    """
    if not is_initialized():
        raise RuntimeError("TTS runtime not initialized")

    runtime = get_runtime()
    metrics = {
        "ttfb_ms": 0,
        "generation_ms": 0,
        "chunks_processed": 0,
        "audio_duration_seconds": 0.0,
        "text_chunked": False,
        "sr_applied": False,
        "output_sample_rate": 16000,
    }

    start_time = time.time()

    try:
        # Import long text processor
        from veena3modal.processing.long_text_processor import LongTextProcessor

        long_processor = LongTextProcessor(pipeline=runtime.pipeline)

        # Check if chunking is needed
        if long_processor.should_chunk(text):
            metrics["text_chunked"] = True
            audio_bytes = await long_processor.generate_with_chunking(
                text=text,
                speaker=speaker,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                max_tokens=max_tokens,
                repetition_penalty=repetition_penalty,
                seed=seed,
                sample_rate=sample_rate,
            )
        else:
            # Short text: direct generation
            audio_bytes = await runtime.pipeline.generate_speech_indic(
                speaker=speaker,
                text=text,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                max_tokens=max_tokens,
                repetition_penalty=repetition_penalty,
                seed=seed,
            )

        generation_time = time.time() - start_time
        metrics["generation_ms"] = int(generation_time * 1000)
        metrics["ttfb_ms"] = metrics["generation_ms"]

        if audio_bytes:
            # Apply super-resolution if requested and available
            final_sample_rate = 16000
            if output_sample_rate == "48khz" and runtime.sr_service and runtime.sr_service.is_loaded:
                try:
                    sr_start = time.time()
                    logger.info(f"Applying super-resolution to chunked audio ({len(audio_bytes)} bytes)...")
                    audio_bytes = _apply_super_resolution(audio_bytes, runtime.sr_service)
                    sr_time = (time.time() - sr_start) * 1000
                    metrics["sr_ms"] = int(sr_time)
                    metrics["sr_applied"] = True
                    final_sample_rate = 48000
                    logger.info(f"✅ Super-resolution applied in {sr_time:.1f}ms, output={len(audio_bytes)} bytes")
                except Exception as e:
                    logger.warning(f"Super-resolution failed, returning 16kHz: {e}")
                    import traceback
                    traceback.print_exc()

            metrics["output_sample_rate"] = final_sample_rate
            audio_duration = (len(audio_bytes) - 44) / (final_sample_rate * 2)
            metrics["audio_duration_seconds"] = max(0.0, audio_duration)

        return audio_bytes, metrics

    except Exception as e:
        logger.error(f"Chunked speech generation failed: {e}")
        raise

from typing import AsyncGenerator

async def generate_speech_streaming(
    text: str,
    speaker: str,
    temperature: float = 0.8,
    top_k: int = 50,
    top_p: float = 1.0,
    max_tokens: int = 4096,
    repetition_penalty: float = 1.05,
    seed: Optional[int] = None,
    enable_chunking: bool = True,
) -> AsyncGenerator[Tuple[bytes, Dict[str, Any]], None]:
    """
    Generate speech audio with true streaming (yields chunks as they're generated).

    This is the core streaming implementation for M4.
    First yield includes WAV header + first PCM chunk.
    Subsequent yields are raw PCM chunks.

    Args:
        text: Text to synthesize (already normalized)
        speaker: Internal speaker name (resolved)
        temperature: Sampling temperature
        top_k: Top-k sampling
        top_p: Nucleus sampling
        max_tokens: Maximum tokens to generate
        repetition_penalty: Repetition penalty
        seed: Random seed for reproducibility
        enable_chunking: Enable text chunking for long inputs with voice consistency

    Yields:
        Tuple of (audio_bytes, metrics_dict)
        - First yield: WAV header (44 bytes) prepended to first PCM chunk
        - Subsequent yields: Raw PCM chunks (int16, 16kHz)
        - Final yield: metrics_dict has final timing info

    Raises:
        RuntimeError: If runtime not initialized or streaming pipeline unavailable
    """
    if not is_initialized():
        raise RuntimeError("TTS runtime not initialized")

    runtime = get_runtime()

    if runtime.streaming_pipeline is None:
        raise RuntimeError("Streaming pipeline not available")

    # Import audio utils for WAV header
    from veena3modal.audio.utils import create_wav_header

    # Metrics tracking
    metrics = {
        "ttfb_ms": 0,
        "chunks_sent": 0,
        "total_bytes": 0,
        "audio_duration_seconds": 0.0,
        "text_chunked": False,
    }

    start_time = time.time()
    first_chunk_time = None
    total_pcm_bytes = 0
    sample_rate = 16000  # BiCodec sample rate

    # Check if we need text chunking
    from veena3modal.processing.long_text_processor import LongTextProcessor
    long_processor = LongTextProcessor(pipeline=runtime.pipeline)
    needs_chunking = enable_chunking and long_processor.should_chunk(text)

    if needs_chunking:
        # Chunked streaming with voice consistency (global token caching)
        metrics["text_chunked"] = True
        async for audio_chunk, chunk_metrics in _stream_chunked_text(
            runtime, long_processor, text, speaker, temperature, top_k, top_p,
            max_tokens, repetition_penalty, seed, sample_rate
        ):
            # First chunk: prepend WAV header
            if first_chunk_time is None:
                first_chunk_time = time.time()
                metrics["ttfb_ms"] = int((first_chunk_time - start_time) * 1000)
                # Create streaming WAV header (size=0 for unknown length)
                wav_header = create_wav_header(sample_rate=sample_rate, data_size=0)
                audio_chunk = wav_header + audio_chunk

            total_pcm_bytes += len(audio_chunk) - (44 if metrics["chunks_sent"] == 0 else 0)
            metrics["chunks_sent"] += 1
            metrics["total_bytes"] = total_pcm_bytes + 44  # Include header
            metrics["audio_duration_seconds"] = total_pcm_bytes / (sample_rate * 2)

            yield audio_chunk, metrics
    else:
        # Simple streaming (no text chunking)
        # OPTIMIZATION: If we have pre-cached global tokens for this speaker,
        # use continuation mode to skip the ~110ms global token pre-roll.
        # The model jumps straight to semantic token generation.
        cached_globals = runtime.speaker_global_cache.get(speaker)

        if cached_globals:
            # Fast path: use cached globals, skip global token generation entirely
            stream_gen = runtime.streaming_pipeline.generate_speech_stream_indic_continuation(
                speaker=speaker,
                text=text,
                global_ids=cached_globals,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                max_tokens=max_tokens,
                repetition_penalty=repetition_penalty,
                seed=seed,
            )
        else:
            # Fallback: no cached globals, generate them inline (original path)
            stream_gen = runtime.streaming_pipeline.generate_speech_stream_indic(
                speaker=speaker,
                text=text,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                max_tokens=max_tokens,
                repetition_penalty=repetition_penalty,
                seed=seed,
            )

        async for audio_chunk in stream_gen:
            # First chunk: prepend WAV header
            if first_chunk_time is None:
                first_chunk_time = time.time()
                metrics["ttfb_ms"] = int((first_chunk_time - start_time) * 1000)
                wav_header = create_wav_header(sample_rate=sample_rate, data_size=0)
                audio_chunk = wav_header + audio_chunk

            total_pcm_bytes += len(audio_chunk) - (44 if metrics["chunks_sent"] == 0 else 0)
            metrics["chunks_sent"] += 1
            metrics["total_bytes"] = total_pcm_bytes + 44
            metrics["audio_duration_seconds"] = total_pcm_bytes / (sample_rate * 2)

            yield audio_chunk, metrics

async def _stream_chunked_text(
    runtime: TTSRuntime,
    long_processor,
    text: str,
    speaker: str,
    temperature: float,
    top_k: int,
    top_p: float,
    max_tokens: int,
    repetition_penalty: float,
    seed: Optional[int],
    sample_rate: int,
) -> AsyncGenerator[Tuple[bytes, Dict[str, Any]], None]:
    """
    Internal helper: stream chunked text with voice consistency.

    Uses global token caching from first chunk to maintain voice across chunks.
    """
    from veena3modal.audio.crossfade import crossfade_bytes_int16

    # Chunk text using the long text processor's chunker
    chunks = long_processor.chunk_text(text)

    if not chunks:
        return

    captured_globals: Optional[List[int]] = None
    previous_chunk_tail: Optional[bytes] = None
    chunk_metrics = {"chunks_processed": 0}

    for i, chunk_text in enumerate(chunks):
        chunk_metrics["chunks_processed"] = i + 1

        if i == 0:
            # First chunk: capture global tokens for voice consistency
            async for audio_bytes, global_ids in runtime.streaming_pipeline.generate_speech_stream_indic_first_chunk(
                speaker=speaker,
                text=chunk_text,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                max_tokens=max_tokens,
                repetition_penalty=repetition_penalty,
                seed=seed,
            ):
                if captured_globals is None and global_ids:
                    captured_globals = global_ids

                # Crossfade with previous tail (inter-chunk stitching)
                to_emit, previous_chunk_tail = crossfade_bytes_int16(
                    previous_chunk_tail,
                    audio_bytes,
                    sample_rate_hz=sample_rate,
                    crossfade_ms=50,
                )

                if to_emit:
                    yield to_emit, chunk_metrics
        else:
            # Continuation chunks: use captured globals for voice consistency
            if captured_globals is None:
                logger.warning(f"No captured globals for chunk {i+1}, using regular streaming")
                async for audio_bytes in runtime.streaming_pipeline.generate_speech_stream_indic(
                    speaker=speaker,
                    text=chunk_text,
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                    max_tokens=max_tokens,
                    repetition_penalty=repetition_penalty,
                    seed=seed,
                ):
                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
                        previous_chunk_tail,
                        audio_bytes,
                        sample_rate_hz=sample_rate,
                        crossfade_ms=50,
                    )
                    if to_emit:
                        yield to_emit, chunk_metrics
            else:
                # Use continuation method with cached globals
                async for audio_bytes in runtime.streaming_pipeline.generate_speech_stream_indic_continuation(
                    speaker=speaker,
                    text=chunk_text,
                    global_ids=captured_globals,
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                    max_tokens=max_tokens,
                    repetition_penalty=repetition_penalty,
                    seed=seed,
                ):
                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
                        previous_chunk_tail,
                        audio_bytes,
                        sample_rate_hz=sample_rate,
                        crossfade_ms=50,
                    )
                    if to_emit:
                        yield to_emit, chunk_metrics

    # Flush remaining tail
    if previous_chunk_tail:
        yield previous_chunk_tail, chunk_metrics


Output:
Chunk ID: d448de
Wall time: 0.0506 seconds
Process exited with code 0
Original token count: 4700
Output:
"""
Veena3 Streaming Pipeline - Sliding Window Approach

Implements streaming for both SNAC and BiCodec tokens.

For BiCodec (Spark TTS):
1. Generate tokens from vLLM
2. Parse NEW token IDs incrementally (O(1) per token via cache)
3. Buffer semantic + global tokens
4. Apply sliding window (every N token pairs)
5. Decode and stream audio chunks

OPTIMIZATION (Dec 2025):
- Replaced O(n²) pattern (decode-all + regex-all per iteration)
- Now uses incremental token parsing with BiCodecTokenParser
- ~10x CPU reduction in streaming hot loop

For SNAC (legacy):
1. Filter SNAC token IDs directly
2. Apply sliding window
3. Decode and stream
"""

import asyncio
import re
from typing import AsyncGenerator, Optional, List
from vllm import SamplingParams
from veena3modal.audio.crossfade import crossfade_bytes_int16
from veena3modal.core.token_utils import BiCodecTokenParser

from veena3modal.core.constants import (
    CODE_END_TOKEN_ID,
    CODE_START_TOKEN_ID,
    SNAC_MIN_ID,
    SNAC_MAX_ID,
    TRAINING_STOP_TOKEN_IDS,
    DEFAULT_TEMPERATURE,
    DEFAULT_TOP_K,
    DEFAULT_TOP_P,
    DEFAULT_MAX_TOKENS,
    DEFAULT_MIN_TOKENS,
    DEFAULT_REPETITION_PENALTY,
    DEFAULT_SEED,
)

class Veena3SlidingWindowPipeline:
    """
    Streaming TTS pipeline using sliding window approach.

    This eliminates choppy audio and popping artifacts by:
    - Decoding overlapping 28-token windows (4 frames)
    - Keeping only the middle 2048 samples from each decode
    - Creating natural continuity between chunks

    Based on the official Canopy Labs implementation.
    """

    def __init__(
        self,
        model,
        prompt_builder,
        snac_decoder,
    ):
        """
        Initialize sliding window streaming pipeline.

        Args:
            model: Veena3Model instance
            prompt_builder: Veena3PromptBuilder instance
            snac_decoder: SNACDecoder instance (with batching enabled)
        """
        self.model = model
        self.prompt_builder = prompt_builder
        self.snac_decoder = snac_decoder

        # OPTIMIZATION: Pre-warm BiCodecTokenParser once at init, not per-request
        # Saves ~123ms per streaming TTFB by avoiding repeated 166K-entry vocab iteration
        tokenizer = getattr(model, "tokenizer", None)
        if tokenizer is None:
            tokenizer = getattr(model.engine, "tokenizer", None)
        self.token_parser = BiCodecTokenParser(tokenizer) if tokenizer else None

        print(f"🌊 Veena3SlidingWindowPipeline initialized (sliding window: 28 tokens)")

    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,
        seed: Optional[int] = None,
    ) -> AsyncGenerator[bytes, None]:
        """
        Generate speech audio with sliding window streaming.

        Yields audio chunks using overlapping windows for smooth playback.

        Args:
            description: Character/voice description
            text: Text to synthesize (with optional <emotion> tags)
            temperature: Sampling temperature
            top_p: Nucleus sampling
            max_tokens: Max SNAC tokens to generate
            repetition_penalty: Prevent loops

        Yields:
            Audio bytes (int16 PCM, 24kHz mono)
        """
        print(f"\n🌊 Sliding window streaming generation")
        print(f"📝 Description: {description[:80]}...")
        print(f"💬 Text: {text}")

        # Build prompt
        prompt = self.prompt_builder.build_prefix(description, text)
        print(f"✅ Prompt built ({len(prompt)} chars)")

        # Configure sampling
        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],  # Only audio EOS
            seed=seed,  # None = random, int = reproducible
        )

        print(f"🎲 Sampling: temp={temperature}, top_p={top_p}, sliding_window=28 tokens")

        # Token buffer - keeps ALL tokens (not chunked)
        token_buffer = []
        total_tokens_generated = 0
        total_audio_chunks = 0

        # Generate tokens with vLLM (streaming)
        print(f"🔮 Starting token generation...")

        # Generate unique request ID for concurrent streaming support
        import uuid
        import time
        request_id = f"slide-{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
        async for request_output in results_generator:
            # Extract generated token IDs
            generated_ids = request_output.outputs[0].token_ids

            # Process only new tokens
            new_tokens = generated_ids[total_tokens_generated:]
            total_tokens_generated = len(generated_ids)

            # Filter and buffer SNAC tokens
            for token_id in new_tokens:
                if SNAC_MIN_ID <= token_id <= SNAC_MAX_ID:
                    token_buffer.append(token_id)

                    # Process every 7 tokens (1 frame) when we have enough for sliding window
                    # Official approach: once count > 27, take last 28 tokens
                    if len(token_buffer) % 7 == 0 and len(token_buffer) > 27:
                        # Sliding window: take last 28 tokens (4 frames)
                        window_tokens = token_buffer[-28:]

                        # Decode with sliding window mode (returns middle 2048 samples only)
                        if self.snac_decoder.enable_batching:
                            audio_bytes = await self.snac_decoder.decode_single_async(
                                window_tokens,
                                trim_warmup=False,  # Sliding window handles trimming
                                use_sliding_window=True  # CRITICAL: Use sliding window mode
                            )
                        else:
                            audio_bytes = self.snac_decoder.decode_to_bytes(
                                window_tokens,
                                trim_warmup=False,
                                use_sliding_window=True
                            )

                        if audio_bytes:
                            total_audio_chunks += 1
                            if total_audio_chunks == 1:
                                print(f"🎵 First chunk decoded ({len(audio_bytes)} bytes, sliding window)")
                            yield audio_bytes

        # Note: No final chunk processing needed - sliding window handles all tokens
        # as they come in (every 7 tokens after the first 28)

        print(f"✅ Sliding window streaming complete: {total_tokens_generated} tokens → {total_audio_chunks} audio chunks")

    def _extract_bicodec_tokens_from_text(self, text: str) -> tuple[List[int], List[int]]:
        """
        Extract BiCodec semantic and global tokens from generated text.

        Args:
            text: Generated text containing BiCodec token markers

        Returns:
            Tuple of (semantic_ids, global_ids)
        """
        semantic_matches = re.findall(r"<\|bicodec_semantic_(\d+)\|>", text)
        semantic_ids = [int(t) for t in semantic_matches]

        global_matches = re.findall(r"<\|bicodec_global_(\d+)\|>", text)
        global_ids = [int(t) for t in global_matches]

        return semantic_ids, global_ids

    async def generate_speech_stream_indic(
        self,
        speaker: str,
        text: str,
        temperature: float = DEFAULT_TEMPERATURE,
        top_k: int = DEFAULT_TOP_K,  # Added for Spark TTS parity with non-streaming
        top_p: float = DEFAULT_TOP_P,
        max_tokens: int = DEFAULT_MAX_TOKENS,
        repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
        seed: Optional[int] = None,
    ) -> AsyncGenerator[bytes, None]:
        """
        Generate speech audio with sliding window streaming for Indic model.

        Yields audio chunks using overlapping windows for smooth playback.

        Args:
            speaker: Speaker name (one of 12 predefined speakers)
            text: Text to synthesize with inline emotion tags
                Examples:
                - "Hello! Welcome."
                - "<laugh> Hello there!"
                - "नमस्ते! <excited> आज का दिन बहुत अच्छा है।"
            temperature: Sampling temperature
            top_p: Nucleus sampling
            max_tokens: Max SNAC tokens to generate
            repetition_penalty: Prevent loops
            seed: Random seed for reproducibility

        Yields:
            Audio bytes (int16 PCM, 16kHz mono - BiCodec)
        """
        import time
        t_start = time.time()

        # Build prompt using Indic prompt builder
        prompt = self.prompt_builder.build_prefix(speaker, text)

        # Configure sampling (matching non-streaming pipeline parameters)
        # OPTIMIZATION: Use Spark TTS stop token, not legacy SNAC CODE_END_TOKEN_ID (128258)
        # which doesn't exist in the Spark vocab and would cause generation to run to max_tokens
        sampling_params = SamplingParams(
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            max_tokens=max_tokens,
            min_tokens=DEFAULT_MIN_TOKENS,
            repetition_penalty=repetition_penalty,
            stop=TRAINING_STOP_TOKEN_IDS,  # "<|im_end|>" - matches non-streaming pipeline
            skip_special_tokens=False,  # Keep BiCodec tokens in output
            seed=seed,
        )

        # BiCodec token buffers (separate for semantic and global)
        semantic_buffer = []
        global_buffer = []
        processed_token_count = 0  # Track how many vLLM tokens we've parsed
        total_audio_chunks = 0

        # BiCodec streaming configuration
        EXPECTED_GLOBAL_COUNT = 32  # BiCodec uses 32 fixed global tokens (per paper)
        MIN_SEMANTIC_FOR_FIRST_CHUNK = 10  # OPTIMIZATION: Lowered from 16 (decoder min is 8, -120ms TTFB)
        # OPTIMIZATION: Increased from 24 to 48 to reduce O(n) re-decode calls by 2x.
        # At 50 TPS, 48 tokens = ~960ms audio per chunk. Cuts total decode calls roughly in half.
        DECODE_INTERVAL = 48
        CROSSFADE_MS = 50
        # OPTIMIZATION: Windowed decode - only decode last WINDOW_SIZE tokens instead of ALL.
        # BiCodec WaveGenerator has finite receptive field (~64 tokens based on kernel sizes).
        # Window must be >= receptive_field + DECODE_INTERVAL to avoid boundary artifacts.
        WINDOW_SIZE = 128  # 128 tokens covers receptive field + margin for crossfade

        # Generate tokens with vLLM (streaming)
        import uuid
        request_id = f"bicodec-stream-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"

        # Reuse pre-warmed token parser (singleton, created once in __init__)
        token_parser = self.token_parser

        t_gen_start = time.time()
        results_generator = self.model.engine.generate(
            prompt=prompt,
            sampling_params=sampling_params,
            request_id=request_id,
        )

        # Track timing
        t_first_token = None
        t_first_chunk_ready = None
        t_first_chunk_decoded = None
        iteration_count = 0
        last_decode_count = 0

        # Track audio samples EMITTED to the user (not just decoded)
        # BiCodec produces ~320 samples per semantic token at 16kHz
        # CRITICAL: This tracks samples that have been YIELDED, not held in tail
        total_samples_emitted_to_user = 0
        # Keep a tail to enable crossfading between consecutive emitted chunks
        previous_chunk_tail: Optional[bytes] = None

        # Stream tokens with BiCodec extraction
        async for request_output in results_generator:
            iteration_count += 1
            t_iter = time.time()

            # Track first token received
            if t_first_token is None:
                t_first_token = t_iter

            # Extract generated token IDs
            generated_ids = request_output.outputs[0].token_ids

            # OPTIMIZATION: Only process NEW tokens (incremental parsing)
            # Old code: tokenizer.decode(ALL) + re.findall(ALL) = O(n²) per stream
            # New code: parse only new tokens = O(1) per token
            new_token_ids = generated_ids[processed_token_count:]
            processed_token_count = len(generated_ids)

            # Parse new tokens and append to buffers (in-place modification)
            token_parser.parse_incremental(new_token_ids, semantic_buffer, global_buffer)

            # TRUE STREAMING with BiCodec two-phase generation:
            # Phase 1: Generate 32 global tokens (pre-roll)
            # Phase 2: Stream semantic tokens incrementally

            # Check if we have the 32 global tokens (pre-roll complete)
            if len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
                # Pre-roll complete! Now we can stream semantic tokens
                if t_first_chunk_ready is None and len(semantic_buffer) >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
                    t_first_chunk_ready = time.time()

                # Stream semantic tokens incrementally
                # Only decode if we have enough NEW semantic tokens since last decode
                semantic_count = len(semantic_buffer)
                if semantic_count >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
                    # Check if we should decode (every DECODE_INTERVAL new semantic tokens)
                    if semantic_count - last_decode_count >= DECODE_INTERVAL:
                        last_decode_count = semantic_count

                        # OPTIMIZATION: Windowed decode instead of full re-decode
                        # Only decode the last WINDOW_SIZE tokens (O(1) per chunk instead of O(n))
                        # For first few chunks when total tokens < WINDOW_SIZE, decode all
                        window_global = global_buffer[:EXPECTED_GLOBAL_COUNT]

                        if len(semantic_buffer) <= WINDOW_SIZE:
                            # Early phase: decode all tokens (window not yet relevant)
                            decode_semantic = semantic_buffer
                        else:
                            # Windowed: only decode last WINDOW_SIZE tokens
                            decode_semantic = semantic_buffer[-WINDOW_SIZE:]

                        t_decode_start = time.time()
                        audio_bytes = await self.snac_decoder.decode_single_async(
                            semantic_ids=decode_semantic,
                            global_ids=window_global,
                            trim_warmup=False,
                            use_sliding_window=False,
                        )
                        t_decode_end = time.time()

                        if audio_bytes:
                            total_samples_decoded = len(audio_bytes) // 2

                            if len(semantic_buffer) <= WINDOW_SIZE:
                                # Early phase: extract new samples from full decode
                                if total_samples_decoded > total_samples_emitted_to_user:
                                    new_bytes_start = total_samples_emitted_to_user * 2
                                    new_audio_bytes = audio_bytes[new_bytes_start:]
                                else:
                                    new_audio_bytes = b""
                            else:
                                # Windowed phase: take only the tail portion of the window decode
                                # The last DECODE_INTERVAL tokens' worth of audio is the new content
                                # Each token produces ~320 samples at 16kHz/50TPS
                                new_samples_approx = DECODE_INTERVAL * 320
                                if total_samples_decoded > new_samples_approx:
                                    new_bytes_start = (total_samples_decoded - new_samples_approx) * 2
                                else:
                                    new_bytes_start = 0
                                new_audio_bytes = audio_bytes[new_bytes_start:]

                            if new_audio_bytes:
                                to_emit, previous_chunk_tail = crossfade_bytes_int16(
                                    previous_chunk_tail,
                                    new_audio_bytes,
                                    sample_rate_hz=16000,
                                    crossfade_ms=CROSSFADE_MS,
                                )

                                samples_emitted_in_this_chunk = len(to_emit) // 2
                                total_samples_emitted_to_user += samples_emitted_in_this_chunk

                                if to_emit:
                                    total_audio_chunks += 1
                                    if t_first_chunk_decoded is None:
                                        t_first_chunk_decoded = t_decode_end

                                    yield to_emit

        # Final chunk: decode any remaining semantic tokens
        # Use the same 32 global tokens (they get pooled)
        if len(semantic_buffer) > last_decode_count and len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
            audio_bytes = self.snac_decoder.decode_streaming(
                semantic_ids=semantic_buffer,  # All semantic tokens
                global_ids=global_buffer[:EXPECTED_GLOBAL_COUNT],  # Use first 32 global tokens
                use_sliding_window=False,  # Get all samples
                trim_warmup=False
            )


Output:
Chunk ID: 8f4cb2
Wall time: 0.0516 seconds
Process exited with code 0
Original token count: 4905
Output:
            )
            if audio_bytes:
                # Only yield NEW samples
                total_samples_decoded = len(audio_bytes) // 2

                if total_samples_decoded > total_samples_emitted_to_user:
                    new_bytes_start = total_samples_emitted_to_user * 2
                    new_audio_bytes = audio_bytes[new_bytes_start:]

                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
                        previous_chunk_tail,
                        new_audio_bytes,
                        sample_rate_hz=16000,
                        crossfade_ms=CROSSFADE_MS,
                    )

                    samples_emitted_in_final = len(to_emit) // 2
                    total_samples_emitted_to_user += samples_emitted_in_final

                    if to_emit:
                        total_audio_chunks += 1
                        yield to_emit

        # Flush any remaining tail that was held back for crossfade
        if previous_chunk_tail:
            total_audio_chunks += 1
            tail_samples = len(previous_chunk_tail) // 2
            total_samples_emitted_to_user += tail_samples
            yield previous_chunk_tail

        # Log completion summary
        t_end = time.time()
        audio_duration_s = total_samples_emitted_to_user / 16000  # 16kHz sample rate
        ttfb_ms = (t_first_chunk_decoded - t_start) * 1000 if t_first_chunk_decoded else 0
        print(f"🎵 Streaming complete: {audio_duration_s:.2f}s audio, TTFB: {ttfb_ms:.0f}ms, {total_audio_chunks} chunks")

    async def generate_speech_stream_indic_first_chunk(
        self,
        speaker: str,
        text: str,
        temperature: float = DEFAULT_TEMPERATURE,
        top_k: int = DEFAULT_TOP_K,
        top_p: float = DEFAULT_TOP_P,
        max_tokens: int = DEFAULT_MAX_TOKENS,
        repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
        seed: Optional[int] = None,
    ) -> AsyncGenerator[tuple, None]:
        """
        Generate speech for the FIRST chunk of multi-chunk text, capturing global tokens.

        This method is specifically for chunked generation: it yields (audio_bytes, global_ids)
        tuples instead of just audio_bytes. The caller captures global_ids from the first
        yield and passes them to generate_speech_stream_indic_continuation() for subsequent chunks.

        Use Case (chunked streaming with voice consistency):
            globals_captured = None
            async for audio_bytes, global_ids in pipeline.generate_speech_stream_indic_first_chunk(...):
                if globals_captured is None and global_ids:
                    globals_captured = global_ids  # Capture once
                yield audio_bytes
            # Now use globals_captured for subsequent chunks

        Args:
            speaker: Speaker name
            text: First text chunk to synthesize
            temperature: Sampling temperature
            top_k: Top-k sampling
            top_p: Nucleus sampling
            max_tokens: Max tokens to generate
            repetition_penalty: Prevent repetition
            seed: Random seed for reproducibility

        Yields:
            Tuple of (audio_bytes, global_ids)
            - audio_bytes: Raw PCM audio (int16, 16kHz)
            - global_ids: List of 32 global token IDs (populated after first decode, else empty)

        Thread Safety:
            This method is stateless and thread-safe. Each call creates its own
            request-scoped state. No global state is modified or shared.
        """
        import time
        t_start = time.time()

        # Build prompt using Indic prompt builder
        prompt = self.prompt_builder.build_prefix(speaker, text)

        # Configure sampling
        # OPTIMIZATION: Use Spark TTS stop token, not legacy SNAC CODE_END_TOKEN_ID
        sampling_params = SamplingParams(
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            max_tokens=max_tokens,
            min_tokens=DEFAULT_MIN_TOKENS,
            repetition_penalty=repetition_penalty,
            stop=TRAINING_STOP_TOKEN_IDS,  # "<|im_end|>" - matches non-streaming pipeline
            skip_special_tokens=False,
            seed=seed,
        )

        # BiCodec token buffers
        semantic_buffer = []
        global_buffer = []
        processed_token_count = 0  # Track how many vLLM tokens we've parsed
        total_audio_chunks = 0

        # BiCodec streaming configuration
        EXPECTED_GLOBAL_COUNT = 32
        MIN_SEMANTIC_FOR_FIRST_CHUNK = 10  # OPTIMIZATION: Lowered from 16 (decoder min is 8)
        DECODE_INTERVAL = 48  # OPTIMIZATION: Increased from 24 to reduce O(n) re-decode calls
        CROSSFADE_MS = 50

        # Generate unique request ID
        import uuid
        request_id = f"bicodec-first-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"

        # Reuse pre-warmed token parser (singleton, created once in __init__)
        token_parser = self.token_parser

        results_generator = self.model.engine.generate(
            prompt=prompt,
            sampling_params=sampling_params,
            request_id=request_id,
        )

        # Track state
        last_decode_count = 0
        total_samples_emitted_to_user = 0
        previous_chunk_tail: Optional[bytes] = None
        captured_globals: List[int] = []  # Will be populated once we have 32 globals

        # Stream tokens with BiCodec extraction
        async for request_output in results_generator:
            generated_ids = request_output.outputs[0].token_ids

            # OPTIMIZATION: Incremental parsing (O(1) per new token)
            new_token_ids = generated_ids[processed_token_count:]
            processed_token_count = len(generated_ids)
            token_parser.parse_incremental(new_token_ids, semantic_buffer, global_buffer)

            # Capture global tokens once we have all 32
            if not captured_globals and len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
                captured_globals = global_buffer[:EXPECTED_GLOBAL_COUNT]

            # Process if we have all global tokens
            if len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
                semantic_count = len(semantic_buffer)
                if semantic_count >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
                    if semantic_count - last_decode_count >= DECODE_INTERVAL:
                        last_decode_count = semantic_count

                        all_semantic = semantic_buffer
                        window_global = global_buffer[:EXPECTED_GLOBAL_COUNT]

                        if self.snac_decoder.enable_batching:
                            audio_bytes = await self.snac_decoder.decode_single_async(
                                semantic_ids=all_semantic,
                                global_ids=window_global,
                                trim_warmup=False,
                                use_sliding_window=False
                            )
                        else:
                            audio_bytes = self.snac_decoder.decode_streaming(
                                semantic_ids=all_semantic,
                                global_ids=window_global,
                                use_sliding_window=False,
                                trim_warmup=False
                            )

                        if audio_bytes:
                            total_samples_decoded = len(audio_bytes) // 2

                            if total_samples_decoded > total_samples_emitted_to_user:
                                new_bytes_start = total_samples_emitted_to_user * 2
                                new_audio_bytes = audio_bytes[new_bytes_start:]

                                to_emit, previous_chunk_tail = crossfade_bytes_int16(
                                    previous_chunk_tail,
                                    new_audio_bytes,
                                    sample_rate_hz=16000,
                                    crossfade_ms=CROSSFADE_MS,
                                )

                                samples_emitted_in_this_chunk = len(to_emit) // 2
                                total_samples_emitted_to_user += samples_emitted_in_this_chunk

                                if to_emit:
                                    total_audio_chunks += 1
                                    # Yield tuple: (audio_bytes, global_ids)
                                    yield (to_emit, captured_globals)

        # Final chunk processing
        if len(semantic_buffer) > last_decode_count and len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
            audio_bytes = self.snac_decoder.decode_streaming(
                semantic_ids=semantic_buffer,
                global_ids=global_buffer[:EXPECTED_GLOBAL_COUNT],
                use_sliding_window=False,
                trim_warmup=False
            )
            if audio_bytes:
                total_samples_decoded = len(audio_bytes) // 2

                if total_samples_decoded > total_samples_emitted_to_user:
                    new_bytes_start = total_samples_emitted_to_user * 2
                    new_audio_bytes = audio_bytes[new_bytes_start:]

                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
                        previous_chunk_tail,
                        new_audio_bytes,
                        sample_rate_hz=16000,
                        crossfade_ms=CROSSFADE_MS,
                    )

                    samples_emitted_in_final = len(to_emit) // 2
                    total_samples_emitted_to_user += samples_emitted_in_final

                    if to_emit:
                        total_audio_chunks += 1
                        yield (to_emit, captured_globals)

        # Flush remaining tail
        if previous_chunk_tail:
            total_audio_chunks += 1
            tail_samples = len(previous_chunk_tail) // 2
            total_samples_emitted_to_user += tail_samples
            yield (previous_chunk_tail, captured_globals)

        t_end = time.time()
        audio_duration_s = total_samples_emitted_to_user / 16000
        print(f"🎵 First chunk complete: {audio_duration_s:.2f}s audio, {total_audio_chunks} chunks, captured {len(captured_globals)} globals")

    async def generate_speech_stream_indic_continuation(
        self,
        speaker: str,
        text: str,
        global_ids: List[int],
        temperature: float = DEFAULT_TEMPERATURE,
        top_k: int = DEFAULT_TOP_K,
        top_p: float = DEFAULT_TOP_P,
        max_tokens: int = DEFAULT_MAX_TOKENS,
        repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
        seed: Optional[int] = None,
    ) -> AsyncGenerator[bytes, None]:
        """
        Generate speech for CONTINUATION chunks using pre-captured global tokens.

        This is the key method for voice consistency in chunked generation:
        - Uses global_ids from the first chunk to maintain identical voice
        - Model generates only semantic tokens (skips global token generation)
        - Ensures no voice drift across chunks

        CRITICAL for production:
        - global_ids is request-scoped, passed explicitly (no shared state)
        - Thread-safe: each request has its own state
        - No global caching that could cause cross-request contamination

        Args:
            speaker: Speaker name (must match first chunk)
            text: Continuation text chunk to synthesize
            global_ids: 32 global token IDs captured from first chunk
            temperature: Sampling temperature
            top_k: Top-k sampling
            top_p: Nucleus sampling
            max_tokens: Max tokens to generate
            repetition_penalty: Prevent repetition
            seed: Random seed (use same as first chunk for consistency)

        Yields:
            Audio bytes (int16 PCM, 16kHz mono)

        Raises:
            ValueError: If global_ids doesn't contain exactly 32 tokens

        Thread Safety:
            Fully thread-safe. All state is request-scoped and passed explicitly.
        """
        import time
        t_start = time.time()

        # Validate global tokens
        EXPECTED_GLOBAL_COUNT = 32
        if len(global_ids) != EXPECTED_GLOBAL_COUNT:
            raise ValueError(
                f"Expected exactly {EXPECTED_GLOBAL_COUNT} global tokens, got {len(global_ids)}. "
                f"global_ids must be captured from first chunk generation."
            )

        # Build prompt WITH pre-filled global tokens
        # This tells the model to skip global generation and go straight to semantic tokens
        prompt = self.prompt_builder.build_prefix_with_globals(speaker, text, global_ids)

        # Configure sampling
        # OPTIMIZATION: Use Spark TTS stop token, not legacy SNAC CODE_END_TOKEN_ID
        sampling_params = SamplingParams(
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            max_tokens=max_tokens,
            min_tokens=DEFAULT_MIN_TOKENS,
            repetition_penalty=repetition_penalty,
            stop=TRAINING_STOP_TOKEN_IDS,  # "<|im_end|>" - matches non-streaming pipeline
            skip_special_tokens=False,
            seed=seed,
        )

        # BiCodec token buffers
        semantic_buffer = []
        global_buffer_unused = []  # Not used in continuation, but needed for parser API
        processed_token_count = 0  # Track how many vLLM tokens we've parsed
        total_audio_chunks = 0

        # BiCodec streaming configuration
        MIN_SEMANTIC_FOR_FIRST_CHUNK = 10  # OPTIMIZATION: Lowered from 16 (decoder min is 8)
        DECODE_INTERVAL = 48  # OPTIMIZATION: Increased from 24 to reduce O(n) re-decode calls
        CROSSFADE_MS = 50

        # Generate unique request ID
        import uuid
        request_id = f"bicodec-cont-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"

        # Reuse pre-warmed token parser (singleton, created once in __init__)
        token_parser = self.token_parser

        results_generator = self.model.engine.generate(
            prompt=prompt,
            sampling_params=sampling_params,
            request_id=request_id,
        )

        # Track state
        last_decode_count = 0
        total_samples_emitted_to_user = 0
        previous_chunk_tail: Optional[bytes] = None

        # Stream tokens - model will generate ONLY semantic tokens (globals are pre-filled)
        async for request_output in results_generator:
            generated_ids = request_output.outputs[0].token_ids

            # OPTIMIZATION: Incremental parsing (O(1) per new token)
            # Note: In continuation, model only generates semantic tokens (globals pre-filled in prompt)
            new_token_ids = generated_ids[processed_token_count:]
            processed_token_count = len(generated_ids)
            token_parser.parse_incremental(new_token_ids, semantic_buffer, global_buffer_unused)

            semantic_count = len(semantic_buffer)
            if semantic_count >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
                if semantic_count - last_decode_count >= DECODE_INTERVAL:
                    last_decode_count = semantic_count

                    # Use pre-captured global tokens for decoding
                    if self.snac_decoder.enable_batching:
                        audio_bytes = await self.snac_decoder.decode_single_async(
                            semantic_ids=semantic_buffer,
                            global_ids=global_ids,  # Use captured globals
                            trim_warmup=False,
                            use_sliding_window=False
                        )
                    else:
                        audio_bytes = self.snac_decoder.decode_streaming(
                            semantic_ids=semantic_buffer,
                            global_ids=global_ids,  # Use captured globals
                            use_sliding_window=False,
                            trim_warmup=False
                        )

                    if audio_bytes:
                        total_samples_decoded = len(audio_bytes) // 2

                        if total_samples_decoded > total_samples_emitted_to_user:
                            new_bytes_start = total_samples_emitted_to_user * 2
                            new_audio_bytes = audio_bytes[new_bytes_start:]

                            to_emit, previous_chunk_tail = crossfade_bytes_int16(
                                previous_chunk_tail,
                                new_audio_bytes,
                                sample_rate_hz=16000,
                                crossfade_ms=CROSSFADE_MS,
                            )

                            samples_emitted_in_this_chunk = len(to_emit) // 2
                            total_samples_emitted_to_user += samples_emitted_in_this_chunk

                            if to_emit:
                                total_audio_chunks += 1
                                yield to_emit

        # Final chunk processing
        if len(semantic_buffer) > last_decode_count:
            audio_bytes = self.snac_decoder.decode_streaming(
                semantic_ids=semantic_buffer,
                global_ids=global_ids,  # Use captured globals
                use_sliding_window=False,
                trim_warmup=False
            )
            if audio_bytes:
                total_samples_decoded = len(audio_bytes) // 2

                if total_samples_decoded > total_samples_emitted_to_user:
                    new_bytes_start = total_samples_emitted_to_user * 2
                    new_audio_bytes = audio_bytes[new_bytes_start:]

                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
                        previous_chunk_tail,
                        new_audio_bytes,
                        sample_rate_hz=16000,
                        crossfade_ms=CROSSFADE_MS,
                    )

                    samples_emitted_in_final = len(to_emit) // 2
                    total_samples_emitted_to_user += samples_emitted_in_final

                    if to_emit:
                        total_audio_chunks += 1
                        yield to_emit

        # Flush remaining tail
        if previous_chunk_tail:
            total_audio_chunks += 1
            tail_samples = len(previous_chunk_tail) // 2
            total_samples_emitted_to_user += tail_samples
            yield previous_chunk_tail

        t_end = time.time()
        audio_duration_s = total_samples_emitted_to_user / 16000
        print(f"🎵 Continuation chunk complete: {audio_duration_s:.2f}s audio, {total_audio_chunks} chunks (using cached globals)")


Output:
Chunk ID: 2d3a78
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 1800
Output:
"""
Spark TTS Model Loader

Loads Spark TTS model with vLLM engine and validates speakers/emotions.
Migrated from Veena3/Orpheus to Spark TTS architecture.
"""

import os
import torch
from typing import Optional, Dict, Any
from transformers import AutoTokenizer
from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams

from veena3modal.core.constants import (
    INDIC_EMOTION_TAGS,
    INDIC_SPEAKERS,
    SPEAKER_MAP,
    DEFAULT_MODEL_PATH,
    DEFAULT_MAX_MODEL_LEN,
    VLLM_CONFIG,
)

class SparkTTSModel:
    """
    Spark TTS Model with vLLM inference engine and BiCodec audio tokenizer.

    Model: BayAreaBoys/spark_tts_4speaker
    Architecture: Qwen2ForCausalLM with BiCodec audio tokenization
    Speakers: 12 predefined (mapped to speaker_0-11 as per training)
    Emotions: 10 tags in [bracket] format

    CRITICAL REQUIREMENTS:
    - Load tokenizer with special Spark TTS tokens
    - Configure vLLM with optimizations (prefix caching, CUDA graphs)
    - Support speaker mapping (lipakshi → speaker_0, etc.)
    """

    def __init__(
        self,
        model_path: str = None,
        model_type: str = "indic_speakers",  # Default to indic_speakers for backward compatibility
        hf_token: str = None,
        **engine_kwargs
    ):
        """
        Initialize Spark TTS model with vLLM.

        Args:
            model_path: Path to model directory (local or HF). If None, uses env var or default.
            model_type: Model type (kept for compatibility, always uses indic_speakers)
            hf_token: HuggingFace token for private models
            **engine_kwargs: Additional vLLM engine arguments (override defaults from VLLM_CONFIG)
        """
        self.model_type = "indic_speakers"  # Spark TTS only supports speaker-based

        # Use provided path, or check environment variable, or use default
        if model_path is None:
            model_path = os.environ.get(
                'SPARK_TTS_MODEL_PATH',
                DEFAULT_MODEL_PATH
            )

        self.model_path = model_path
        self.hf_token = hf_token

        print(f"🚀 Initializing Spark TTS Model")
        print(f"📁 Model path: {model_path}")
        print("🗣  Speakers: 12 predefined "
              "(lipakshi, vardan, reet, Nandini, krishna, anika, adarsh, "
              "Nilay, Aarvi, Asha, Bittu, Mira)")
        print(f"🎭 Emotions: 10 tags in [bracket] format")

        # Load tokenizer
        print(f"\n📝 Loading tokenizer from model...")
        tokenizer_kwargs = {
            'trust_remote_code': True,
        }

        # Add HF token for private repos
        if hf_token:
            tokenizer_kwargs['token'] = hf_token

        # Add emotion tags as special tokens
        special_tokens = {
            'additional_special_tokens': INDIC_EMOTION_TAGS
        }

        self.tokenizer = AutoTokenizer.from_pretrained(
            model_path,
            **tokenizer_kwargs
        )

        # Add special tokens
        self.tokenizer.add_special_tokens(special_tokens)

        print(f"✅ Tokenizer loaded: {len(self.tokenizer)} tokens")

        # Validate speakers and emotions
        self._validate_speakers()
        self._validate_emotion_tags()

        # Initialize vLLM engine with Spark TTS optimizations
        print(f"\n🔧 Initializing vLLM engine with optimizations...")

        # Start with default VLLM_CONFIG and override with kwargs
        engine_kwargs_dict = VLLM_CONFIG.copy()
        engine_kwargs_dict.update({
            'model': model_path,
            'tokenizer': model_path,
        })

        # Add HF token for private models
        if hf_token:
            os.environ['HUGGING_FACE_HUB_TOKEN'] = hf_token

        # Override with any user-provided kwargs
        engine_kwargs_dict.update(engine_kwargs)

        # Print configuration
        print(f"  - dtype: {engine_kwargs_dict['dtype']}")
        print(f"  - max_model_len: {engine_kwargs_dict['max_model_len']}")
        print(f"  - gpu_memory_utilization: {engine_kwargs_dict['gpu_memory_utilization']}")
        print(f"  - enable_prefix_caching: {engine_kwargs_dict.get('enable_prefix_caching', False)}")
        print(f"  - enforce_eager (disable CUDA graphs): {engine_kwargs_dict.get('enforce_eager', True)}")

        engine_args = AsyncEngineArgs(**engine_kwargs_dict)
        self.engine = AsyncLLMEngine.from_engine_args(engine_args)

        print(f"✅ vLLM engine initialized successfully!")
        print(f"\n🎉 Spark TTS Model ready for inference!\n")

    def _validate_emotion_tags(self):
        """
        Validate that emotion tags are properly configured for Spark TTS.

        For Spark TTS, emotion tags use [bracket] format and should be
        added as special tokens to the tokenizer.
        """
        print(f"\n🔍 Validating {len(INDIC_EMOTION_TAGS)} emotion tags for Spark TTS...")

        # Check if emotion tags can be tokenized
        for tag in INDIC_EMOTION_TAGS:
            token_ids = self.tokenizer.encode(tag, add_special_tokens=False)
            # Spark TTS emotion tags may be multiple tokens, that's okay
            if not token_ids:
                print(f"⚠  Warning: Emotion tag {tag} produced no tokens")

        print(f"✅ All {len(INDIC_EMOTION_TAGS)} emotion tags validated")

        # Print sample emotion tags
        print(f"\n📋 Sample emotion tags:")
        for i, tag in enumerate(INDIC_EMOTION_TAGS[:5]):
            token_ids = self.tokenizer.encode(tag, add_special_tokens=False)
            print(f"  {tag}: {len(token_ids)} tokens")
        if len(INDIC_EMOTION_TAGS) > 5:
            print(f"  ...")

    def _validate_speakers(self):
        """
        Validate that all speakers are properly mapped for Spark TTS.

        Spark TTS uses speaker tokens like <|speaker_0|> to <|speaker_7|>.
        User-facing names (lipakshi, vardan, etc.) are mapped to these IDs.
        """
        print(f"\n🔍 Validating {len(INDIC_SPEAKERS)} speakers for Spark TTS...")

        # Validate speaker mapping
        for speaker_name in INDIC_SPEAKERS:
            if speaker_name not in SPEAKER_MAP:
                print(f"⚠  Warning: Speaker {speaker_name} not in SPEAKER_MAP")
                continue

            speaker_id = SPEAKER_MAP[speaker_name]
            speaker_token = f'<|speaker_{speaker_id}|>'
            token_ids = self.tokenizer.encode(speaker_token, add_special_tokens=False)
            print(f"  {speaker_name} → {speaker_token}: {len(token_ids)} tokens")

        print(f"✅ All {len(INDIC_SPEAKERS)} speakers validated with mapping")

    def get_model_type(self) -> str:
        """Get the model type (always 'indic_speakers' for Spark TTS)."""
        return self.model_type

    def get_tokenizer(self):
        """Get the tokenizer."""
        return self.tokenizer

    def get_engine(self):
        """Get the vLLM engine."""
        return self.engine


# Collaboration Mode: Default

You are now in Default mode. Any previous instructions for other modes (e.g. Plan mode) are no longer active.

Your active mode changes only when new developer instructions with a different `<collaboration_mode>...</collaboration_mode>` change it; user requests or tool descriptions do not change mode by themselves. Known mode names are Default and Plan.

## request_user_input availability

The `request_user_input` tool is unavailable in Default mode. If you call it while in Default mode, it will return an error.

If a decision is necessary and cannot be discovered from local context, ask the user directly. However, in Default mode you should strongly prefer executing the user's request rather than stopping to ask questions.

Summary: auto

## Skills
A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list of skills that can be used. Each entry includes a name, description, and file path so you can open the source for full instructions when using a specific skill.
### Available skills
- skill-creator: Guide for creating effective skills. This skill should be used when users want to create a new skill (or update an existing skill) that extends Codex's capabilities with specialized knowledge, workflows, or tool integrations. (file: /home/ubuntu/.codex/skills/.system/skill-creator/SKILL.md)
- skill-installer: Install Codex skills into $CODEX_HOME/skills from a curated list or a GitHub repo path. Use when a user asks to list installable skills, install a curated skill, or install a skill from another repo (including private repos). (file: /home/ubuntu/.codex/skills/.system/skill-installer/SKILL.md)
### How to use skills
- Discovery: The list above is the skills available in this session (name + description + file path). Skill bodies live on disk at the listed paths.
- Trigger rules: If the user names a skill (with `$SkillName` or plain text) OR the task clearly matches a skill's description shown above, you must use that skill for that turn. Multiple mentions mean use them all. Do not carry skills across turns unless re-mentioned.
- Missing/blocked: If a named skill isn't in the list or the path can't be read, say so briefly and continue with the best fallback.
- How to use a skill (progressive disclosure):
  1) After deciding to use a skill, open its `SKILL.md`. Read only enough to follow the workflow.
  2) When `SKILL.md` references relative paths (e.g., `scripts/foo.py`), resolve them relative to the skill directory listed above first, and only consider other paths if needed.
  3) If `SKILL.md` points to extra folders such as `references/`, load only the specific files needed for the request; don't bulk-load everything.
  4) If `scripts/` exist, prefer running or patching them instead of retyping large code blocks.
  5) If `assets/` or templates exist, reuse them instead of recreating from scratch.
- Coordination and sequencing:
  - If multiple skills apply, choose the minimal set that covers the request and state the order you'll use them.
  - Announce which skill(s) you're using and why (one short line). If you skip an obvious skill, say why.
- Context hygiene:
  - Keep context small: summarize long sections instead of pasting them; only load extra files when needed.
  - Avoid deep reference-chasing: prefer opening only files directly linked from `SKILL.md` unless you're blocked.
  - When variants exist (frameworks, providers, domains), pick only the relevant reference file(s) and note that choice.
- Safety and fallback: If a skill can't be applied cleanly (missing files, unclear instructions), state the issue, pick the next-best approach, and continue.


**Investigating missing detailed metrics**


rg -n "x-perf-details|x-llm-|llm_batch|decode_gpu|parse_ms|request_inflight|perf|timing" veena3modal -S


rg -n "timing|llm_|decode_|batch_|gpu_" scripts/stress_test_optimized_local.py -S


sed -n '520,980p' scripts/stress_test_optimized_local.py


sed -n '1,320p' veena3modal/shared/logging.py


sed -n '1,320p' veena3modal/shared/metrics.py


Output:
Chunk ID: 5d33d5
Wall time: 0.0514 seconds
Process exited with code 0
Original token count: 682
Output:
veena3modal/local_server.py:459:        # Performance: keep-alive for connection reuse under load
veena3modal/app.py:140:    # L40S: Best cost/performance ratio for TTS inference
veena3modal/app.py:255:    gpu="L40S",  # Best cost/performance for TTS
veena3modal/services/tts_runtime.py:375:        metrics_dict: Dictionary with timing metrics
veena3modal/services/tts_runtime.py:658:        - Final yield: metrics_dict has final timing info
veena3modal/audio/crossfade.py:32:    Perform equal-power crossfade between previous chunk tail and the start of a new chunk.
veena3modal/tests/integration/test_tts_generate.py:70:    # Cleanup (optional - runtime persists for performance)
veena3modal/tests/performance/__init__.py:1:"""Performance/benchmark tests for `veena3modal` (TTFB/RTF regressions)."""
veena3modal/tests/performance/test_benchmarks.py:2:Performance benchmarks for TTS service.
veena3modal/tests/performance/test_benchmarks.py:7:Run with: pytest veena3modal/tests/performance -v -m slow
veena3modal/tests/performance/test_benchmarks.py:140:        and technical terms like API and HTTP. Let's see how well it performs with
veena3modal/tests/performance/test_benchmarks.py:259:    Cache performance benchmarks.
veena3modal/tests/performance/test_benchmarks.py:262:    def test_rate_limiter_check_performance(self):
veena3modal/tests/performance/test_benchmarks.py:277:        print(f"\nRate limiter check performance:")
veena3modal/tests/performance/test_benchmarks.py:283:    def test_api_key_cache_performance(self):
veena3modal/tests/performance/test_benchmarks.py:302:        print(f"\nAPI key cache lookup performance:")
veena3modal/core/super_resolution.py:84:            # Performance tracking
veena3modal/core/super_resolution.py:296:        # Track performance
veena3modal/core/super_resolution.py:301:        # Log performance occasionally
veena3modal/core/super_resolution.py:305:                f"SR performance: {self.total_chunks_processed} chunks, "
veena3modal/core/super_resolution.py:312:        """Get performance statistics."""
veena3modal/shared/metrics.py:92:    # Performance histograms
veena3modal/core/streaming_pipeline.py:299:        # Track timing
veena3modal/tests/modal_live/test_deployed_endpoint.py:256:class TestPerformance:
veena3modal/tests/modal_live/test_deployed_endpoint.py:257:    """Performance tests against deployed Modal endpoint."""
veena3modal/tests/modal_live/test_load.py:145:    """Make a single TTS request and measure timing."""
veena3modal/tests/modal_live/test_load.py:393:            text="Load test sentence for performance measurement.",
veena3modal/tests/modal_live/debug_asr_issue.py:340:    # Test 1: Short text (should have near-perfect transcription)


Output:
Chunk ID: dc7287
Wall time: 0.0516 seconds
Process exited with code 0
Original token count: 1503
Output:
4:This script is focused on collecting per-request GPU/LLM timing for local tuning:
6:- GPU timing breakdown (LLM batch wall / GPU wall, parse, decode)
58:    gpu_utilization_pct: float
102:                                gpu_utilization_pct=float(parts[2]),
118:        gpu_util = [s.gpu_utilization_pct for s in self.snapshots]
129:            "gpu_util_pct": {
130:                "min": min(gpu_util),
131:                "max": max(gpu_util),
132:                "avg": statistics.mean(gpu_util),
156:    timing: Dict[str, Any] = field(default_factory=dict)
170:    gpu_summary: Dict[str, Any] = field(default_factory=dict)
231:    def timing_stats(self, key: str) -> Dict[str, float]:
234:            v = r.timing.get(key)
248:    def timing_rate(self, numerator_key: str, denominator_key: str) -> Tuple[float, int]:
252:            n = r.timing.get(numerator_key)
253:            d = r.timing.get(denominator_key)
300:    """Parse server timing headers into numeric metric fields."""
314:        "x-llm-token-total": "llm_token_total",
315:        "x-llm-batch-count": "llm_batch_count",
316:        "x-llm-batch-wall-ms": "llm_batch_wall_ms_total",
317:        "x-llm-batch-gpu-ms": "llm_batch_gpu_ms_total",
318:        "x-llm-batch-wall-ms-min": "llm_batch_wall_ms_min",
319:        "x-llm-batch-wall-ms-max": "llm_batch_wall_ms_max",
320:        "x-llm-batch-wall-ms-p50": "llm_batch_wall_ms_p50",
321:        "x-llm-batch-gpu-ms-min": "llm_batch_gpu_ms_min",
322:        "x-llm-batch-gpu-ms-max": "llm_batch_gpu_ms_max",
323:        "x-llm-batch-gpu-ms-p50": "llm_batch_gpu_ms_p50",
324:        "x-llm-decode-wall-ms": "llm_decode_wall_ms_total",
325:        "x-llm-decode-gpu-ms": "llm_decode_gpu_ms_total",
326:        "x-llm-decode-wall-ms-min": "llm_decode_wall_ms_min",
327:        "x-llm-decode-wall-ms-max": "llm_decode_wall_ms_max",
328:        "x-llm-decode-wall-ms-p50": "llm_decode_wall_ms_p50",
329:        "x-llm-decode-gpu-ms-min": "llm_decode_gpu_ms_min",
330:        "x-llm-decode-gpu-ms-max": "llm_decode_gpu_ms_max",
331:        "x-llm-decode-gpu-ms-p50": "llm_decode_gpu_ms_p50",
332:        "x-llm-decode-calls": "decode_calls",
333:        "x-llm-decode-cpu-ms": "llm_decode_cpu_ms",
334:        "x-llm-parse-ms": "llm_parse_ms",
335:        "x-llm-parse-avg-ms": "llm_parse_ms_avg",
336:        "x-llm-tokens-per-batch-min": "tokens_per_batch_min",
337:        "x-llm-tokens-per-batch-max": "tokens_per_batch_max",
338:        "x-llm-tokens-per-batch-p50": "tokens_per_batch_p50",
343:        "x-llm-time-per-token-ms": "llm_time_per_token_ms",
344:        "x-llm-time-per-batch-wall-ms": "llm_time_per_batch_wall_ms",
345:        "x-llm-time-per-batch-gpu-ms": "llm_time_per_batch_gpu_ms",
346:        "x-llm-parse-avg-ms": "llm_parse_ms_avg",
364:    if "llm_token_total" in metrics and "generation_ms" in metrics and metrics["llm_token_total"]:
366:            metrics["llm_time_per_token_ms"] = float(metrics["generation_ms"]) / float(metrics["llm_token_total"])
370:    if "llm_batch_count" in metrics and "generation_ms" in metrics and metrics["llm_batch_count"]:
372:            metrics["llm_time_per_batch_wall_ms"] = float(metrics["generation_ms"]) / float(metrics["llm_batch_count"])
376:    if "llm_batch_count" in metrics and "llm_batch_gpu_ms_total" in metrics and metrics["llm_batch_count"]:
378:            metrics["llm_time_per_batch_gpu_ms"] = float(metrics["llm_batch_gpu_ms_total"]) / float(metrics["llm_batch_count"])
424:                timing = _parse_perf_headers(response.headers)
440:                        timing=timing,
441:                        request_inflight=timing.get("request_inflight"),
454:                    timing=timing,
456:                    request_inflight=timing.get("request_inflight"),
466:            timing = _parse_perf_headers(response.headers)
492:                    timing=timing,
493:                    request_inflight=timing.get("request_inflight"),
507:                timing=timing,
509:                request_inflight=timing.get("request_inflight"),
541:    gpu_monitor: Optional[GPUMonitor] = None,
573:    if gpu_monitor:
574:        gpu_monitor.snapshots.clear()
601:    if gpu_monitor:
602:        level.gpu_summary = gpu_monitor.summary()
639:        llm = level.timing_stats("generation_ms")
643:        wall = level.timing_stats("llm_batch_wall_ms_total")
647:        gpu = level.timing_stats("llm_batch_gpu_ms_total")
651:        parse = level.timing_stats("llm_parse_ms")
655:        decode = level.timing_stats("llm_decode_wall_ms_total")
660:        token_ms, token_count = level.timing_rate("generation_ms", "llm_token_total")
662:            print(f"  Token timing: avg {token_ms / token_count * 1000:.3f} us/token")
664:    if level.gpu_summary and "memory_mb" in level.gpu_summary:
665:        gpu_summary = level.gpu_summary
666:        mem = gpu_summary["memory_mb"]
667:        util = gpu_summary["gpu_util_pct"]
668:        temp = gpu_summary["temperature_c"]
693:        gpu_util = level.gpu_summary.get("gpu_util_pct", {}).get("avg", 0)
694:        gpu_mem = level.gpu_summary.get("memory_mb", {}).get("max", 0)
702:            f"{gpu_util:>4.0f}% "
703:            f"{gpu_mem:>7.0f}MB "
729:            "gpu": level.gpu_summary,
730:            "timing_averages": {
731:                "generation_ms": level.timing_stats("generation_ms"),
732:                "llm_batch_wall_ms_total": level.timing_stats("llm_batch_wall_ms_total"),
733:                "llm_batch_gpu_ms_total": level.timing_stats("llm_batch_gpu_ms_total"),
734:                "llm_parse_ms": level.timing_stats("llm_parse_ms"),
735:                "llm_decode_wall_ms_total": level.timing_stats("llm_decode_wall_ms_total"),
749:                    "timing": r.timing,
768:            print(f"  GPU: {health.get('gpu_available')}")
774:    gpu_monitor = GPUMonitor(interval_seconds=0.5)
775:    gpu_monitor.start()
788:        gpu_monitor=gpu_monitor,
802:            gpu_monitor=gpu_monitor,
810:    gpu_monitor.stop()


Output:
Chunk ID: bc357d
Wall time: 0.0507 seconds
Process exited with code 0
Original token count: 2933
Output:
            stream=stream,
        )
    except Exception as e:
        return RequestResult(
            success=False,
            status_code=0,
            latency_ms=(time.time() - start) * 1000,
            error=str(e)[:200],
            speaker=speaker,
            text_length=len(text),
            stream=stream,
        )

async def run_level(
    base_url: str,
    num_requests: int,
    concurrency: int,
    text_category: str,
    chunking: bool,
    stream: bool,
    gpu_monitor: Optional[GPUMonitor] = None,
) -> LevelResult:
    mode = "streaming" if stream else "non-streaming"
    level = LevelResult(
        level_name=f"{concurrency}c-{text_category}-{mode}",
        concurrency=concurrency,
        total_requests=num_requests,
        mode=mode,
        text_category=text_category,
        chunking=chunking,
    )

    semaphore = asyncio.Semaphore(concurrency)

    if text_category == "mixed":
        texts = list(TEST_TEXTS.values())
    else:
        texts = [TEST_TEXTS.get(text_category, TEST_TEXTS["short"])]

    async def limited_request(i: int, client: httpx.AsyncClient) -> RequestResult:
        async with semaphore:
            text = texts[i % len(texts)]
            speaker = SPEAKERS[i % len(SPEAKERS)]
            return await make_request(
                client=client,
                base_url=base_url,
                text=text,
                speaker=speaker,
                chunking=chunking,
                stream=stream,
            )

    if gpu_monitor:
        gpu_monitor.snapshots.clear()

    start_time = time.time()
    limits = httpx.Limits(
        max_connections=concurrency + 10,
        max_keepalive_connections=concurrency,
    )
    async with httpx.AsyncClient(limits=limits) as client:
        tasks = [limited_request(i, client) for i in range(num_requests)]
        responses = await asyncio.gather(*tasks, return_exceptions=True)

    level.wall_time_seconds = time.time() - start_time

    for item in responses:
        if isinstance(item, RequestResult):
            level.results.append(item)
        else:
            level.results.append(
                RequestResult(
                    success=False,
                    status_code=0,
                    latency_ms=0,
                    error=str(item)[:200],
                    stream=stream,
                )
            )

    if gpu_monitor:
        level.gpu_summary = gpu_monitor.summary()

    return level

def print_level_report(level: LevelResult):
    print(f"\n{'=' * 72}")
    print(f"  {level.level_name.upper()}")
    print(f"  {level.total_requests} requests, {level.concurrency} concurrent, {level.mode}, chunking={level.chunking}")
    print(f"{'=' * 72}")

    print(f"  Success:     {len(level.successes)}/{level.total_requests} ({level.success_rate:.0%})")
    print(f"  Wall time:   {level.wall_time_seconds:.2f}s")
    print(f"  Throughput:  {level.throughput_rps:.2f} req/s")
    print(f"  Audio total: {level.audio_duration_total:.2f}s")
    print(f"  Eff. RTF:    {level.effective_rtf:.3f}")

    if level.observed_inflight_max() is not None:
        print(f"  Peak inflight: {level.observed_inflight_max()}")

    latency = level.latency_stats
    if latency:
        print("\n  Latency (ms):")
        print(
            f"    min={latency['min']:.0f} avg={latency['avg']:.0f} "
            f"p50={latency['p50']:.0f} p95={latency['p95']:.0f} p99={latency['p99']:.0f} max={latency['max']:.0f}"
        )

    ttfb = level.ttfb_stats
    if ttfb and not level.mode.startswith("stream"):
        print("\n  Server TTFB (ms):")
        print(
            f"    min={ttfb['min']:.0f} avg={ttfb['avg']:.0f} "
            f"p50={ttfb['p50']:.0f} p95={ttfb['p95']:.0f} max={ttfb['max']:.0f}"
        )

    if level.mode == "non-streaming":
        llm = level.timing_stats("generation_ms")
        if llm:
            print("\n  Generation (ms):")
            print(f"    min={llm['min']:.0f} avg={llm['avg']:.0f} p50={llm['p50']:.0f} p95={llm['p95']:.0f} max={llm['max']:.0f}")
        wall = level.timing_stats("llm_batch_wall_ms_total")
        if wall:
            print("\n  LLM batch wall (ms):")
            print(f"    min={wall['min']:.0f} avg={wall['avg']:.0f} p50={wall['p50']:.0f} p95={wall['p95']:.0f} max={wall['max']:.0f}")
        gpu = level.timing_stats("llm_batch_gpu_ms_total")
        if gpu:
            print("\n  LLM batch GPU (ms):")
            print(f"    min={gpu['min']:.0f} avg={gpu['avg']:.0f} p50={gpu['p50']:.0f} p95={gpu['p95']:.0f} max={gpu['max']:.0f}")
        parse = level.timing_stats("llm_parse_ms")
        if parse:
            print("\n  Parse (ms):")
            print(f"    avg={parse['avg']:.0f} p95={parse['p95']:.0f}")
        decode = level.timing_stats("llm_decode_wall_ms_total")
        if decode:
            print("\n  Decode (ms):")
            print(f"    avg={decode['avg']:.0f} p95={decode['p95']:.0f}")

        token_ms, token_count = level.timing_rate("generation_ms", "llm_token_total")
        if token_count:
            print(f"  Token timing: avg {token_ms / token_count * 1000:.3f} us/token")

    if level.gpu_summary and "memory_mb" in level.gpu_summary:
        gpu_summary = level.gpu_summary
        mem = gpu_summary["memory_mb"]
        util = gpu_summary["gpu_util_pct"]
        temp = gpu_summary["temperature_c"]
        print("\n  GPU:")
        print(f"    Memory: {mem['avg']:.0f}MB avg, {mem['max']:.0f}MB peak / {mem['total']:.0f}MB")
        print(f"    Util:   {util['avg']:.0f}% avg, {util['max']:.0f}% peak")
        print(f"    Temp:   {temp['avg']:.0f}C avg, {temp['max']:.0f}C peak")

    errors = level.error_summary()
    if errors:
        print("\n  Errors:")
        for key, count in sorted(errors.items(), key=lambda kv: -kv[1]):
            print(f"    [{count}x] {key}")

def print_summary_table(levels: List[LevelResult]):
    print(f"\n{'=' * 110}")
    print("  STRESS BENCHMARK SUMMARY")
    print(f"{'=' * 110}")
    header = f"{'Level':<26} {'OK%':>5} {'RPS':>7} {'p50':>7} {'p95':>7} {'GPU%':>5} {'GPUmem':>8} {'RTF':>6} {'Inflight':>8}"
    print(header)
    print("-" * 110)

    for level in levels:
        lat = level.latency_stats
        p50 = f"{lat.get('p50', 0):.0f}" if lat else "---"
        p95 = f"{lat.get('p95', 0):.0f}" if lat else "---"
        gpu_util = level.gpu_summary.get("gpu_util_pct", {}).get("avg", 0)
        gpu_mem = level.gpu_summary.get("memory_mb", {}).get("max", 0)
        inflight = level.observed_inflight_max() or 0
        print(
            f"{level.level_name:<26} "
            f"{level.success_rate:>4.0%} "
            f"{level.throughput_rps:>7.2f} "
            f"{p50:>7} "
            f"{p95:>7} "
            f"{gpu_util:>4.0f}% "
            f"{gpu_mem:>7.0f}MB "
            f"{level.effective_rtf:>6.3f} "
            f"{inflight:>8d}"
        )

    print("=" * 110)

def save_results(levels: List[LevelResult], path: str):
    payload = []
    for level in levels:
        payload.append({
            "level": level.level_name,
            "concurrency": level.concurrency,
            "mode": level.mode,
            "chunking": level.chunking,
            "text_category": level.text_category,
            "success_rate": level.success_rate,
            "throughput_rps": level.throughput_rps,
            "wall_time_s": level.wall_time_seconds,
            "latency_ms": level.latency_stats,
            "ttfb_ms": level.ttfb_stats,
            "effective_rtf": level.effective_rtf,
            "audio_total_seconds": level.audio_duration_total,
            "observed_inflight_max": level.observed_inflight_max(),
            "errors": level.error_summary(),
            "gpu": level.gpu_summary,
            "timing_averages": {
                "generation_ms": level.timing_stats("generation_ms"),
                "llm_batch_wall_ms_total": level.timing_stats("llm_batch_wall_ms_total"),
                "llm_batch_gpu_ms_total": level.timing_stats("llm_batch_gpu_ms_total"),
                "llm_parse_ms": level.timing_stats("llm_parse_ms"),
                "llm_decode_wall_ms_total": level.timing_stats("llm_decode_wall_ms_total"),
            },
            "requests": [
                {
                    "request_id": r.request_id,
                    "success": r.success,
                    "status_code": r.status_code,
                    "latency_ms": r.latency_ms,
                    "ttfb_ms": r.ttfb_ms,
                    "server_ttfb_ms": r.server_ttfb_ms,
                    "audio_bytes": r.audio_bytes,
                    "audio_seconds": r.audio_seconds,
                    "request_inflight": r.request_inflight,
                    "error": r.error,
                    "timing": r.timing,
                }
                for r in level.results
            ],
        })

    with open(path, "w") as fd:
        json.dump(payload, fd, indent=2)

async def run_stress_test(args):
    base_url = args.url

    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(f"{base_url}/v1/tts/health", timeout=10)
            health = response.json()
            print(f"  Health: {health.get('status')}")
            print(f"  Version: {health.get('model_version')}")
            print(f"  GPU: {health.get('gpu_available')}")
    except Exception as e:
        print(f"ERROR: server not reachable at {base_url}: {e}")
        sys.exit(1)

    levels = [int(v) for v in args.levels.split(",") if v.strip()]
    gpu_monitor = GPUMonitor(interval_seconds=0.5)
    gpu_monitor.start()

    mode_label = "streaming" if args.stream else "non-streaming"
    results: List[LevelResult] = []

    print(f"\nWarming up (1 request, non-streaming)...")
    warmup = await run_level(
        base_url=base_url,
        num_requests=1,
        concurrency=1,
        text_category=args.text,
        chunking=args.chunking,
        stream=False,
        gpu_monitor=gpu_monitor,
    )
    print(f"  warmup ok: {warmup.success_rate:.0%}")

    for conc in levels:
        num_requests = max(conc * 3, 10)
        num_requests = min(num_requests, 200)
        level = await run_level(
            base_url=base_url,
            num_requests=num_requests,
            concurrency=conc,
            text_category=args.text,
            chunking=args.chunking,
            stream=args.stream,
            gpu_monitor=gpu_monitor,
        )
        print_level_report(level)
        results.append(level)
        if conc != levels[-1]:
            print("  cooldown 3s...")
            await asyncio.sleep(3)

    gpu_monitor.stop()
    print_summary_table(results)

    if args.output:
        save_results(results, args.output)
        print(f"\nResults written to {args.output}")

def main():
    parser = argparse.ArgumentParser(description="Detailed local optimized-path TTS benchmark")
    parser.add_argument("--url", default=DEFAULT_URL, help=f"Server URL (default: {DEFAULT_URL})")
    parser.add_argument("--levels", default="1,5,10,20,50", help="Comma-separated concurrency levels")
    parser.add_argument("--text", default="short", choices=["short", "medium", "long", "mixed"], help="Input text category")
    parser.add_argument("--chunking", action="store_true", help="Enable server-side chunking", default=True)
    parser.add_argument("--no-chunking", action="store_false", dest="chunking", help="Disable chunking")
    parser.add_argument("--stream", action="store_true", help="Use streaming endpoint")
    parser.add_argument("--output", default="stress_test_optimized_local.json", help="Output JSON path")

    args = parser.parse_args()
    asyncio.run(run_stress_test(args))

if __name__ == "__main__":
    main()