Output:
Chunk ID: 1d6a65
Wall time: 0.0511 seconds
Process exited with code 0
Original token count: 1168
Output:
74:def _build_perf_headers(metrics: Dict[str, object]) -> Dict[str, str]:
105:    _add("X-LLM-Token-Total", llm_token_total)
106:    _add("X-LLM-Batch-Count", llm_batch_count)
107:    _add("X-LLM-Batch-Wall-Ms", llm_batch_wall_total)
108:    _add("X-LLM-Batch-GPU-Ms", llm_batch_gpu_total)
109:    _add("X-LLM-Batch-Wall-Ms-Min", _pick("llm_batch_wall_ms_min"))
110:    _add("X-LLM-Batch-Wall-Ms-Max", _pick("llm_batch_wall_ms_max"))
111:    _add("X-LLM-Batch-Wall-Ms-P50", _pick("llm_batch_wall_ms_p50"))
112:    _add("X-LLM-Batch-GPU-Ms-Min", _pick("llm_batch_gpu_ms_min"))
113:    _add("X-LLM-Batch-GPU-Ms-Max", _pick("llm_batch_gpu_ms_max"))
114:    _add("X-LLM-Batch-GPU-Ms-P50", _pick("llm_batch_gpu_ms_p50"))
115:    _add("X-LLM-Decode-Wall-Ms", llm_decode_wall_total)
116:    _add("X-LLM-Decode-GPU-Ms", llm_decode_gpu_total)
117:    _add("X-LLM-Decode-Wall-Ms-Min", _pick("llm_decode_wall_ms_min"))
118:    _add("X-LLM-Decode-Wall-Ms-Max", _pick("llm_decode_wall_ms_max"))
119:    _add("X-LLM-Decode-Wall-Ms-P50", _pick("llm_decode_wall_ms_p50"))
120:    _add("X-LLM-Decode-GPU-Ms-Min", _pick("llm_decode_gpu_ms_min"))
121:    _add("X-LLM-Decode-GPU-Ms-Max", _pick("llm_decode_gpu_ms_max"))
122:    _add("X-LLM-Decode-GPU-Ms-P50", _pick("llm_decode_gpu_ms_p50"))
123:    _add("X-LLM-Decode-Calls", _pick("decode_calls"))
124:    _add("X-LLM-Decode-CPU-Ms", decode_cpu_ms)
125:    _add("X-LLM-Parse-Ms", llm_parse_ms or parse_ms)
126:    _add("X-LLM-Parse-Avg-Ms", _pick("llm_parse_ms_avg"))
127:    _add("X-LLM-Tokens-Per-Batch-Min", _pick("tokens_per_batch_min"))
128:    _add("X-LLM-Tokens-Per-Batch-Max", _pick("tokens_per_batch_max"))
129:    _add("X-LLM-Tokens-Per-Batch-P50", _pick("tokens_per_batch_p50"))
130:    _add("X-Generation-Ms", generation_ms)
131:    _add("X-LLM-Generation-Ms", _pick("llm_wall_ms", generation_ms))
132:    _add("X-Request-Inflight", _pick("request_inflight", "active_requests"))
133:    _add("X-LLM-Text-Chunked", _pick("text_chunked"))
134:    _add("X-LLM-Chunks-Processed", _pick("chunks_processed"))
135:    _add("X-LLM-SR-Applied", _pick("sr_applied"))
139:            perf["X-LLM-Time-Per-Token-Ms"] = f"{float(generation_ms) / float(llm_token_total):.3f}"
145:            perf["X-LLM-Time-Per-Batch-Wall-Ms"] = f"{float(generation_ms or 0) / float(llm_batch_count):.3f}"
149:            perf["X-LLM-Time-Per-Batch-GPU-Ms"] = f"{float(llm_batch_gpu_total or 0) / float(llm_batch_count):.3f}"
194:    perf["X-Perf-Details"] = json.dumps(
285:                "X-Model-Version": model_version,
286:                "X-App-Version": _APP_VERSION,
302:    @app.post("/v1/tts/generate")
303:    async def tts_generate(request: Request):
327:                        message="API key required. Use 'Authorization: Bearer <key>' or 'X-API-Key: <key>' header.",
330:                    headers={"X-Request-ID": request_id},
347:                    headers={"X-Request-ID": request_id},
357:                    headers["X-Request-ID"] = request_id
388:                headers={"X-Request-ID": request_id},
402:                headers={"X-Request-ID": request_id},
446:                audio_bytes, metrics = await tts_runtime.generate_speech_chunked(
460:                audio_bytes, metrics = await tts_runtime.generate_speech(
483:                    headers={"X-Request-ID": request_id},
499:                "X-Request-ID": request_id,
500:                "X-Model-Version": model_version,
501:                "X-Format": "wav",
502:                "X-Sample-Rate": str(output_sample_rate),
503:                "X-Stream": "false",
504:                "X-Audio-Bytes": str(len(audio_bytes)),
505:                "X-Audio-Seconds": f"{audio_duration:.2f}",
506:                "X-TTFB-ms": str(metrics.get("ttfb_ms", 0)),
507:                "X-RTF": f"{rtf:.3f}",
508:                "X-Text-Chunked": str(metrics.get("text_chunked", False)).lower(),
509:                "X-SR-Applied": str(metrics.get("sr_applied", False)).lower(),
511:            headers.update(_build_perf_headers(metrics))
514:                headers["X-Seed"] = str(req.seed)
589:                    headers["X-Audio-Bytes"] = str(len(audio_bytes))
630:                headers={"X-Request-ID": request_id},
656:                headers={"X-Request-ID": request_id},
670:                headers={"X-Request-ID": request_id},
780:            "X-Request-ID": request_id,
781:            "X-Model-Version": model_version,
782:            "X-Format": "wav",
783:            "X-Sample-Rate": str(req.sample_rate),
784:            "X-Stream": "true",
785:            "X-Chunking-Enabled": str(req.chunking).lower(),
790:            headers["X-Seed"] = str(req.seed)


# 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.


sed -n '260,860p' veena3modal/api/fastapi_app.py


Output:
Chunk ID: ccd078
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 5634
Output:
            gpu_available = torch.cuda.is_available()
        except ImportError:
            pass  # torch not installed, skip GPU check

        # Compute uptime
        uptime = time.time() - _startup_time if _startup_time else 0

        # Determine overall status
        if model_loaded and gpu_available:
            status = "healthy"
        elif model_loaded or gpu_available:
            status = "degraded"
        else:
            status = "degraded"

        return JSONResponse(
            content={
                "status": status,
                "model_loaded": model_loaded,
                "model_version": model_version,
                "uptime_seconds": round(uptime, 2),
                "gpu_available": gpu_available,
                "app_version": _APP_VERSION,
            },
            headers={
                "X-Model-Version": model_version,
                "X-App-Version": _APP_VERSION,
            },
        )

    @app.get("/v1/tts/metrics")
    async def tts_metrics():
        """
        Prometheus metrics endpoint.

        Returns metrics in Prometheus text format for scraping.
        """
        return Response(
            content=get_metrics_text(),
            media_type=get_metrics_content_type(),
        )

    @app.post("/v1/tts/generate")
    async def tts_generate(request: Request):
        """
        Generate speech from text.

        Supports streaming (M4) and non-streaming (M3) modes.
        Currently implements non-streaming WAV generation only.
        """
        from veena3modal.api.schemas import TTSGenerateRequest, AudioFormat
        from veena3modal.services import tts_runtime

        request_id = str(uuid.uuid4())
        start_time = time.time()

        # === AUTH CHECK ===
        if FeatureFlags.is_auth_enabled():
            # Extract API key from headers
            headers_dict = dict(request.headers)
            api_key = extract_api_key(headers_dict)

            if not api_key:
                return JSONResponse(
                    status_code=401,
                    content=create_error_response(
                        code=ErrorCode.INVALID_API_KEY,
                        message="API key required. Use 'Authorization: Bearer <key>' or 'X-API-Key: <key>' header.",
                        request_id=request_id,
                    ),
                    headers={"X-Request-ID": request_id},
                )

            # Validate API key
            validator = get_api_validator()
            key_hash = hash_api_key(api_key)
            auth_result = validator.validate(key_hash)

            if not auth_result.is_valid:
                status_code = 401 if auth_result.error_code == "INVALID_API_KEY" else 403
                return JSONResponse(
                    status_code=status_code,
                    content=create_error_response(
                        code=ErrorCode(auth_result.error_code),
                        message=auth_result.error_message or "Authentication failed",
                        request_id=request_id,
                    ),
                    headers={"X-Request-ID": request_id},
                )

            # === RATE LIMIT CHECK ===
            if FeatureFlags.is_rate_limiting_enabled():
                rate_limiter = get_rate_limiter()
                allowed, remaining, reset_after = rate_limiter.check(key_hash)

                if not allowed:
                    headers = rate_limiter.get_headers(allowed, remaining, reset_after)
                    headers["X-Request-ID"] = request_id
                    return JSONResponse(
                        status_code=429,
                        content=create_error_response(
                            code=ErrorCode.RATE_LIMIT_EXCEEDED,
                            message=f"Rate limit exceeded. Retry after {int(reset_after) + 1} seconds.",
                            request_id=request_id,
                        ),
                        headers=headers,
                    )

        # Parse and validate request body
        try:
            body = await request.json()
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid JSON: {e}")

        try:
            req = TTSGenerateRequest(**body)
        except Exception as e:
            error_msg = str(e)
            # Extract validation error details
            return JSONResponse(
                status_code=400,
                content={
                    "error": {
                        "code": "VALIDATION_ERROR",
                        "message": error_msg,
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

        # Check if runtime is initialized
        if not tts_runtime.is_initialized():
            return JSONResponse(
                status_code=503,
                content={
                    "error": {
                        "code": "MODEL_NOT_LOADED",
                        "message": "TTS model not initialized. Please wait for warmup.",
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

        # Get resolved speaker
        speaker = req.get_resolved_speaker()

        # Get normalized text (apply text normalization if enabled)
        normalizer_func = None
        if req.normalize:
            try:
                from veena3modal.processing.text_normalizer import normalize_text
                normalizer_func = lambda t: normalize_text(t, verbose=req.normalize_verbose)
            except ImportError:
                logger.warning("TextNormalizer not available, skipping text normalization")

        text = req.get_normalized_text(normalizer_func=normalizer_func)

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

        # Log request received (lifecycle event + metrics)
        log_request_received(
            request_id=request_id,
            text_length=len(text),
            speaker=speaker,
            stream=req.stream,
            format=req.format,
        )
        metrics_request_received(speaker=speaker, stream=req.stream, format=req.format)

        # Streaming mode (M4)
        if req.stream:
            return await _handle_streaming_request(req, request_id, text, speaker, start_time)

        # Format encoding will be applied after generation (non-streaming only)
        # For streaming, only WAV is supported
        target_format = req.format

        try:
            # Get output sample rate - use_enum_values=True means req.output is already a string
            output_sr = req.output if hasattr(req, 'output') else "16khz"

            # Generate speech (non-streaming)
            if req.chunking:
                audio_bytes, metrics = await tts_runtime.generate_speech_chunked(
                    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,
                    sample_rate=req.sample_rate,
                    output_sample_rate=output_sr,
                    request_id=request_id,
                )
            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,
                    request_id=request_id,
                )

            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(),
            }
            headers.update(_build_perf_headers(metrics))

            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,
            )
            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,
                    request_id=request_id,
                ):
                    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,
                    **final_metrics,
                )
                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


Summary: auto


rg -n "def generate_speech\(|def generate_speech_chunked|def generate_speech_streaming|request_timeout|request_id|metrics|inflight|_timing_enabled|ttfb" veena3modal/services/tts_runtime.py -n


Output:
Chunk ID: 5031d2
Wall time: 0.0511 seconds
Process exited with code 0
Original token count: 1929
Output:
30:_inflight_counter = 0
31:_inflight_lock = Lock()
34:def _timing_enabled() -> bool:
51:def _with_inflight_delta(delta: int) -> int:
53:    global _inflight_counter
54:    with _inflight_lock:
55:        _inflight_counter = max(0, _inflight_counter + delta)
56:        return _inflight_counter
325:        request_id = f"warmup-{speaker_name}-{int(time.time() * 1000)}"
331:            request_id=request_id,
378:async def generate_speech(
388:    request_id: Optional[str] = None,
405:        Tuple of (audio_bytes, metrics_dict)
407:        metrics_dict: Dictionary with timing metrics
416:    metrics = {
417:        "request_id": request_id,
418:        "ttfb_ms": 0,
426:    if request_id is None:
428:        request_id = f"runtime-ns-{uuid.uuid4().hex[:10]}"
429:        metrics["request_id"] = request_id
431:    active_requests = _with_inflight_delta(1)
441:                "request_id": request_id,
443:                "request_inflight": active_requests,
464:        if _timing_enabled():
468:                request_id=request_id,
485:                request_id=request_id,
501:        metrics["generation_ms"] = int(generation_time * 1000)
502:        metrics["ttfb_ms"] = metrics["generation_ms"]  # Non-streaming: TTFB = total time
505:            metrics.update({
514:                "llm_wall_ms": pipeline_timings.get("llm_wall_ms", metrics["generation_ms"]),
515:                "ttfb_ms": int(pipeline_timings.get("ttfb_ms", metrics["ttfb_ms"])),
534:                    metrics["sr_ms"] = int(sr_time)
535:                    metrics["sr_applied"] = True
543:            metrics["output_sample_rate"] = sample_rate
547:            metrics["audio_duration_seconds"] = max(0.0, audio_duration)
549:        return audio_bytes, metrics
557:                "request_id": request_id,
563:        active_requests = _with_inflight_delta(-1)
564:        if _timing_enabled():
565:            metrics["request_inflight"] = active_requests
566:            metrics["active_requests"] = active_requests
569:                extra=metrics,
571:        if _timing_enabled():
573:                "tts_runtime_nonstream_inflight",
575:                    "request_id": request_id,
647:async def generate_speech_chunked(
658:    request_id: Optional[str] = None,
673:        Tuple of (audio_bytes, metrics_dict)
679:    if request_id is None:
681:        request_id = f"runtime-chunk-{uuid.uuid4().hex[:10]}"
682:    metrics = {
683:        "request_id": request_id,
684:        "ttfb_ms": 0,
693:    active_requests = _with_inflight_delta(1)
697:        if _timing_enabled():
702:                    "request_id": request_id,
704:                    "request_inflight": active_requests,
716:            metrics["text_chunked"] = True
717:            if _timing_enabled():
728:                    request_id=request_id,
747:                    request_id=request_id,
762:            if _timing_enabled():
766:                    request_id=request_id,
784:                    request_id=request_id,
801:        metrics["generation_ms"] = int(generation_time * 1000)
802:        metrics["ttfb_ms"] = metrics["generation_ms"]
805:            metrics.update({
809:                "llm_batch_wall_ms_total": chunk_timing.get("llm_wall_ms", metrics["generation_ms"]),
822:            if metrics["text_chunked"]:
823:                metrics["chunks_processed"] = chunk_timing.get("chunks_processed", 0)
834:                    metrics["sr_ms"] = int(sr_time)
835:                    metrics["sr_applied"] = True
843:            metrics["output_sample_rate"] = final_sample_rate
845:            metrics["audio_duration_seconds"] = max(0.0, audio_duration)
847:        if _timing_enabled():
848:            metrics["active_requests"] = active_requests - 1
849:            metrics["request_inflight"] = active_requests - 1
850:            logger.info("tts_runtime_chunking_complete", extra=metrics)
851:        return audio_bytes, metrics
857:        active_requests = _with_inflight_delta(-1)
858:        metrics["active_requests"] = active_requests
859:        metrics["request_inflight"] = active_requests
865:async def generate_speech_streaming(
875:    request_id: Optional[str] = None,
896:        Tuple of (audio_bytes, metrics_dict)
899:        - Final yield: metrics_dict has final timing info
909:    if request_id is None:
911:        request_id = f"runtime-stream-{uuid.uuid4().hex[:10]}"
920:    metrics = {
921:        "request_id": request_id,
923:        "ttfb_ms": 0,
930:    pipeline_metrics: Dict[str, Any] = {
931:        "request_id": request_id,
941:    active_requests = _with_inflight_delta(1)
942:    if _timing_enabled():
948:                "request_inflight": active_requests,
950:                "request_id": request_id,
962:            metrics["text_chunked"] = True
963:            async for audio_chunk, chunk_metrics in _stream_chunked_text(
975:                request_id,
976:                request_metrics=pipeline_metrics,
980:                    metrics["ttfb_ms"] = int((first_chunk_time - start_time) * 1000)
985:                total_pcm_bytes += len(audio_chunk) - (44 if metrics["chunks_sent"] == 0 else 0)
986:                metrics["chunks_sent"] += 1
987:                metrics["total_bytes"] = total_pcm_bytes + 44  # Include header
988:                metrics["audio_duration_seconds"] = total_pcm_bytes / (sample_rate * 2)
989:                if chunk_metrics:
990:                    metrics.update(chunk_metrics)
992:                yield audio_chunk, metrics
1011:                    request_id=request_id,
1012:                    request_metrics=pipeline_metrics,
1024:                    request_id=request_id,
1025:                    request_metrics=pipeline_metrics,
1031:                    metrics["ttfb_ms"] = int((first_chunk_time - start_time) * 1000)
1035:                total_pcm_bytes += len(audio_chunk) - (44 if metrics["chunks_sent"] == 0 else 0)
1036:                metrics["chunks_sent"] += 1
1037:                metrics["total_bytes"] = total_pcm_bytes + 44
1038:                metrics["audio_duration_seconds"] = total_pcm_bytes / (sample_rate * 2)
1039:                yield audio_chunk, metrics
1041:        active_requests = _with_inflight_delta(-1)
1042:        metrics.update(pipeline_metrics)
1043:        metrics["generation_ms"] = int((time.time() - start_time) * 1000)
1044:        metrics["text_chunked"] = bool(metrics.get("text_chunked", False))
1045:        if _timing_enabled():
1050:                    "request_id": request_id,
1051:                    "request_inflight": active_requests,
1052:                    "chunks_sent": metrics.get("chunks_sent", 0),
1053:                    "generation_ms": metrics["generation_ms"],
1070:    request_id: Optional[str] = None,
1071:    request_metrics: Optional[Dict[str, Any]] = None,
1088:    chunk_metrics = request_metrics if request_metrics is not None else {"request_id": request_id}
1089:    if "request_id" not in chunk_metrics:
1090:        chunk_metrics["request_id"] = request_id
1091:    if "chunks_processed" not in chunk_metrics:
1092:        chunk_metrics["chunks_processed"] = 0
1095:        chunk_metrics["chunks_processed"] = i + 1
1108:                request_id=request_id,
1109:                request_metrics=chunk_metrics,
1123:                    yield to_emit, chunk_metrics
1137:                    request_id=request_id,
1138:                    request_metrics=chunk_metrics,
1147:                        yield to_emit, chunk_metrics
1160:                    request_id=request_id,
1161:                    request_metrics=chunk_metrics,
1170:                        yield to_emit, chunk_metrics
1174:        yield previous_chunk_tail, chunk_metrics


sed -n '320,760p' veena3modal/services/tts_runtime.py


Output:
Chunk ID: 1b1f32
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 4018
Output:
    )

    async def _capture_globals_for_speaker(speaker_name: str) -> Optional[List[int]]:
        """Generate a short utterance and capture the 32 global tokens."""
        prompt = runtime.prompt_builder.build_prefix(speaker_name, test_text)
        request_id = f"warmup-{speaker_name}-{int(time.time() * 1000)}"

        final_output = None
        async for request_output in runtime.model.engine.generate(
            prompt=prompt,
            sampling_params=sampling_params,
            request_id=request_id,
        ):
            final_output = request_output

        if final_output is None:
            return None

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

        if len(global_ids) >= 32:
            return global_ids[:32]
        return None

    async def _run_all():
        for speaker_name in SPEAKER_MAP.keys():
            try:
                global_ids = await _capture_globals_for_speaker(speaker_name)
                if global_ids and len(global_ids) == 32:
                    runtime.speaker_global_cache[speaker_name] = global_ids
                    logger.info(f"  Cached globals for {speaker_name}")
                else:
                    logger.warning(f"  Failed to capture globals for {speaker_name}")
            except Exception as e:
                logger.warning(f"  Error caching globals for {speaker_name}: {e}")

    # Run async warmup in event loop
    try:
        loop = asyncio.get_event_loop()
        if loop.is_running():
            # Already in async context (unlikely at startup, but handle it)
            import concurrent.futures
            with concurrent.futures.ThreadPoolExecutor() as pool:
                pool.submit(lambda: asyncio.run(_run_all())).result()
        else:
            loop.run_until_complete(_run_all())
    except RuntimeError:
        # No event loop exists yet
        asyncio.run(_run_all())

    elapsed = (time.time() - t_start) * 1000
    cached = len(runtime.speaker_global_cache)
    total = len(SPEAKER_MAP)
    logger.info(f"Speaker globals cached: {cached}/{total} speakers in {elapsed:.0f}ms")

async def generate_speech(
    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,
    output_sample_rate: str = "16khz",
    request_id: Optional[str] = None,
) -> Tuple[Optional[bytes], Dict[str, Any]]:
    """
    Generate speech audio (non-streaming).

    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
        output_sample_rate: "16khz" or "48khz" (triggers super-resolution)

    Returns:
        Tuple of (audio_bytes, metrics_dict)
        audio_bytes: WAV audio bytes (16kHz or 48kHz, 16-bit PCM) or None if failed
        metrics_dict: Dictionary with timing metrics

    Raises:
        RuntimeError: If runtime not initialized
    """
    if not is_initialized():
        raise RuntimeError("TTS runtime not initialized")

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

    if request_id is None:
        import uuid
        request_id = f"runtime-ns-{uuid.uuid4().hex[:10]}"
        metrics["request_id"] = request_id

    active_requests = _with_inflight_delta(1)
    start_time = time.time()

    pipeline_timings = {}

    try:
        logger.info(
            "tts_runtime_nonstream_start",
            extra={
                "speaker": speaker,
                "request_id": request_id,
                "text_length": len(text),
                "request_inflight": active_requests,
                "temperature": temperature,
                "top_k": top_k,
                "top_p": top_p,
                "max_tokens": max_tokens,
            },
        )

        generation_timeout = _generation_timeout_seconds()

        gen_kwargs = dict(
            speaker=speaker,
            text=text,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            max_tokens=max_tokens,
            repetition_penalty=repetition_penalty,
            seed=seed,
        )

        if _timing_enabled():
            generation_coro = runtime.pipeline.generate_speech_indic(
                **gen_kwargs,
                return_timing=True,
                request_id=request_id,
            )
            if generation_timeout:
                try:
                    audio_bytes, pipeline_timings = await asyncio.wait_for(
                        generation_coro,
                        timeout=generation_timeout,
                    )
                except asyncio.TimeoutError as exc:
                    raise TimeoutError(
                        f"Non-stream generation timed out after {generation_timeout:.2f}s"
                    ) from exc
            else:
                audio_bytes, pipeline_timings = await generation_coro
        else:
            generation_coro = runtime.pipeline.generate_speech_indic(
                **gen_kwargs,
                request_id=request_id,
            )
            if generation_timeout:
                try:
                    audio_bytes = await asyncio.wait_for(
                        generation_coro,
                        timeout=generation_timeout,
                    )
                except asyncio.TimeoutError as exc:
                    raise TimeoutError(
                        f"Non-stream generation timed out after {generation_timeout:.2f}s"
                    ) from exc
            else:
                audio_bytes = await generation_coro

        generation_time = time.time() - start_time
        metrics["generation_ms"] = int(generation_time * 1000)
        metrics["ttfb_ms"] = metrics["generation_ms"]  # Non-streaming: TTFB = total time

        if pipeline_timings:
            metrics.update({
                "llm_token_total": pipeline_timings.get("token_total", 0),
                "llm_batch_count": pipeline_timings.get("batch_count", 0),
                "llm_batch_gpu_ms_total": pipeline_timings.get("llm_gpu_batch_ms_total", 0),
                "llm_batch_gpu_ms_min": pipeline_timings.get("llm_gpu_batch_ms_min", 0),
                "llm_batch_gpu_ms_max": pipeline_timings.get("llm_gpu_batch_ms_max", 0),
                "llm_batch_gpu_ms_p50": pipeline_timings.get("llm_gpu_batch_ms_p50", 0),
                "llm_decode_cpu_ms": pipeline_timings.get("decode_cpu_ms", 0),
                "llm_parse_ms": pipeline_timings.get("parse_ms", 0),
                "llm_wall_ms": pipeline_timings.get("llm_wall_ms", metrics["generation_ms"]),
                "ttfb_ms": int(pipeline_timings.get("ttfb_ms", metrics["ttfb_ms"])),
                "tokens_per_batch_min": pipeline_timings.get("tokens_per_batch_min"),
                "tokens_per_batch_max": pipeline_timings.get("tokens_per_batch_max"),
                "tokens_per_batch_p50": pipeline_timings.get("tokens_per_batch_p50"),
            })

        if audio_bytes:
            # Apply super-resolution if requested and available
            sample_rate = 16000
            logger.info(f"SR check: output_sample_rate={output_sample_rate!r}, "
                       f"sr_service={runtime.sr_service is not None}, "
                       f"is_loaded={runtime.sr_service.is_loaded if runtime.sr_service else False}")

            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 {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
                    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"] = sample_rate

            # Calculate audio duration
            audio_duration = (len(audio_bytes) - 44) / (sample_rate * 2)  # -44 for WAV header
            metrics["audio_duration_seconds"] = max(0.0, audio_duration)

        return audio_bytes, metrics

    except Exception as e:
        logger.error(f"Speech generation failed: {e}")
        logger.exception(
            "tts_runtime_nonstream_error",
            extra={
                "speaker": speaker,
                "request_id": request_id,
                "error": str(e),
            },
        )
        raise
    finally:
        active_requests = _with_inflight_delta(-1)
        if _timing_enabled():
            metrics["request_inflight"] = active_requests
            metrics["active_requests"] = active_requests
            logger.info(
                "tts_runtime_nonstream_complete",
                extra=metrics,
            )
        if _timing_enabled():
            logger.info(
                "tts_runtime_nonstream_inflight",
                extra={
                    "request_id": request_id,
                    "active_requests": active_requests,
                },
            )

def _apply_super_resolution(audio_bytes: bytes, sr_service) -> bytes:
    """
    Apply super-resolution to audio bytes.

    Args:
        audio_bytes: WAV audio at 16kHz
        sr_service: SuperResolutionService instance

    Returns:
        WAV audio at 48kHz
    """
    import numpy as np
    import struct
    import torch

    # Parse WAV header and extract PCM data
    if len(audio_bytes) < 44:
        raise ValueError("Invalid WAV data")

    # Extract PCM data (skip 44-byte WAV header)
    pcm_data = np.frombuffer(audio_bytes[44:], dtype=np.int16)

    # Convert to float32 for SR model
    audio_float = pcm_data.astype(np.float32) / 32768.0

    # Convert numpy to torch tensor for process_chunk
    audio_tensor = torch.from_numpy(audio_float)

    # Apply super-resolution using process_chunk (16kHz -> 48kHz)
    # process_chunk expects [batch, samples] or [samples] and returns [batch, samples]
    upsampled_tensor = sr_service.process_chunk(audio_tensor)

    # Convert back to numpy, squeeze batch dim if present
    upsampled = upsampled_tensor.squeeze().cpu().numpy()

    # Convert back to int16
    upsampled_int16 = np.clip(upsampled * 32768.0, -32768, 32767).astype(np.int16)

    # Create new WAV header for 48kHz
    sample_rate = 48000
    num_channels = 1
    bits_per_sample = 16
    byte_rate = sample_rate * num_channels * bits_per_sample // 8
    block_align = num_channels * bits_per_sample // 8
    data_size = len(upsampled_int16) * 2

    wav_header = struct.pack(
        '<4sI4s4sIHHIIHH4sI',
        b'RIFF',
        36 + data_size,
        b'WAVE',
        b'fmt ',
        16,  # fmt chunk size
        1,   # PCM format
        num_channels,
        sample_rate,
        byte_rate,
        block_align,
        bits_per_sample,
        b'data',
        data_size,
    )

    return wav_header + upsampled_int16.tobytes()

async def generate_speech_chunked(
    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,
    sample_rate: int = 16000,
    output_sample_rate: str = "16khz",
    request_id: Optional[str] = None,
) -> 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()
    if request_id is None:
        import uuid
        request_id = f"runtime-chunk-{uuid.uuid4().hex[:10]}"
    metrics = {
        "request_id": request_id,
        "ttfb_ms": 0,
        "generation_ms": 0,
        "chunks_processed": 0,
        "audio_duration_seconds": 0.0,
        "text_chunked": False,
        "sr_applied": False,
        "output_sample_rate": 16000,
    }

    active_requests = _with_inflight_delta(1)
    start_time = time.time()

    try:
        if _timing_enabled():
            logger.info(
                "tts_runtime_chunking_start",
                extra={
                    "speaker": speaker,
                    "request_id": request_id,
                    "text_length": len(text),
                    "request_inflight": active_requests,
                    "chunking": True,
                },
            )
        # Import long text processor
        from veena3modal.processing.long_text_processor import LongTextProcessor

        long_processor = LongTextProcessor(pipeline=runtime.pipeline)
        generation_timeout = _generation_timeout_seconds()

        # Check if chunking is needed
        if long_processor.should_chunk(text):
            metrics["text_chunked"] = True
            if _timing_enabled():
                chunk_result = 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,
                    request_id=request_id,
                    return_timing=True,
                )
                if isinstance(chunk_result, tuple):
                    audio_bytes, chunk_timing = chunk_result
                else:
                    audio_bytes = chunk_result
                    chunk_timing = {}
            else:
                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,
                    request_id=request_id,
                )
                chunk_timing = {}
        else:
            # Short text: direct generation
            gen_kwargs = dict(
                speaker=speaker,
                text=text,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                max_tokens=max_tokens,
                repetition_penalty=repetition_penalty,
                seed=seed,


sed -n '760,1160p' veena3modal/services/tts_runtime.py


Output:
Chunk ID: 6e89f5
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 4203
Output:
                seed=seed,
            )
            if _timing_enabled():
                generation_coro = runtime.pipeline.generate_speech_indic(
                    **gen_kwargs,
                    return_timing=True,
                    request_id=request_id,
                )
                if generation_timeout:
                    try:
                        audio_bytes, pipeline_timings = await asyncio.wait_for(
                            generation_coro,
                            timeout=generation_timeout,
                        )
                    except asyncio.TimeoutError as exc:
                        raise TimeoutError(
                            f"Chunked generation timed out after {generation_timeout:.2f}s"
                        ) from exc
                else:
                    audio_bytes, pipeline_timings = await generation_coro
                chunk_timing = pipeline_timings
            else:
                generation_coro = runtime.pipeline.generate_speech_indic(
                    **gen_kwargs,
                    request_id=request_id,
                )
                if generation_timeout:
                    try:
                        audio_bytes = await asyncio.wait_for(
                            generation_coro,
                            timeout=generation_timeout,
                        )
                    except asyncio.TimeoutError as exc:
                        raise TimeoutError(
                            f"Chunked generation timed out after {generation_timeout:.2f}s"
                        ) from exc
                else:
                    audio_bytes = await generation_coro
                chunk_timing = {}

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

        if chunk_timing:
            metrics.update({
                "llm_token_total": chunk_timing.get("token_total", 0),
                "llm_batch_count": chunk_timing.get("batch_count", 0),
                "llm_batch_gpu_ms_total": chunk_timing.get("llm_gpu_batch_ms_total", 0),
                "llm_batch_wall_ms_total": chunk_timing.get("llm_wall_ms", metrics["generation_ms"]),
                "llm_parse_ms": chunk_timing.get("parse_ms", 0),
                "llm_decode_cpu_ms": chunk_timing.get("decode_cpu_ms", 0),
                "llm_batch_wall_ms_min": chunk_timing.get("llm_batch_wall_ms_min", 0),
                "llm_batch_wall_ms_max": chunk_timing.get("llm_batch_wall_ms_max", 0),
                "llm_batch_gpu_ms_min": chunk_timing.get("llm_batch_gpu_ms_min", 0),
                "llm_batch_gpu_ms_max": chunk_timing.get("llm_batch_gpu_ms_max", 0),
                "llm_batch_gpu_ms_p50": chunk_timing.get("llm_batch_gpu_ms_p50", 0),
                "tokens_per_batch_min": chunk_timing.get("tokens_per_batch_min"),
                "tokens_per_batch_max": chunk_timing.get("tokens_per_batch_max"),
                "tokens_per_batch_p50": chunk_timing.get("tokens_per_batch_p50"),
                "llm_batch_tokens": chunk_timing.get("token_batches", []),
            })
            if metrics["text_chunked"]:
                metrics["chunks_processed"] = chunk_timing.get("chunks_processed", 0)

        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)

        if _timing_enabled():
            metrics["active_requests"] = active_requests - 1
            metrics["request_inflight"] = active_requests - 1
            logger.info("tts_runtime_chunking_complete", extra=metrics)
        return audio_bytes, metrics

    except Exception as e:
        logger.error(f"Chunked speech generation failed: {e}")
        raise
    finally:
        active_requests = _with_inflight_delta(-1)
        metrics["active_requests"] = active_requests
        metrics["request_inflight"] = active_requests

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,
    request_id: Optional[str] = None,
) -> 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 request_id is None:
        import uuid
        request_id = f"runtime-stream-{uuid.uuid4().hex[:10]}"

    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 = {
        "request_id": request_id,
        "speaker": speaker,
        "ttfb_ms": 0,
        "chunks_sent": 0,
        "total_bytes": 0,
        "audio_duration_seconds": 0.0,
        "text_chunked": False,
        "text_length": len(text),
    }
    pipeline_metrics: Dict[str, Any] = {
        "request_id": request_id,
        "speaker": speaker,
        "text_length": len(text),
        "text_chunked": False,
    }

    start_time = time.time()
    first_chunk_time = None
    total_pcm_bytes = 0
    sample_rate = 16000  # BiCodec sample rate
    active_requests = _with_inflight_delta(1)
    if _timing_enabled():
        logger.info(
            "tts_runtime_stream_start",
            extra={
                "speaker": speaker,
                "text_length": len(text),
                "request_inflight": active_requests,
                "chunking": enable_chunking,
                "request_id": request_id,
            },
        )

    try:
        # 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,
                request_id,
                request_metrics=pipeline_metrics,
            ):
                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)
                if chunk_metrics:
                    metrics.update(chunk_metrics)

                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:
                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,
                    request_id=request_id,
                    request_metrics=pipeline_metrics,
                )
            else:
                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,
                    request_id=request_id,
                    request_metrics=pipeline_metrics,
                )

            async for audio_chunk in stream_gen:
                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
    finally:
        active_requests = _with_inflight_delta(-1)
        metrics.update(pipeline_metrics)
        metrics["generation_ms"] = int((time.time() - start_time) * 1000)
        metrics["text_chunked"] = bool(metrics.get("text_chunked", False))
        if _timing_enabled():
            logger.info(
                "tts_runtime_stream_end",
                extra={
                    "speaker": speaker,
                    "request_id": request_id,
                    "request_inflight": active_requests,
                    "chunks_sent": metrics.get("chunks_sent", 0),
                    "generation_ms": metrics["generation_ms"],
                },
            )

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,
    request_id: Optional[str] = None,
    request_metrics: Optional[Dict[str, Any]] = None,
) -> 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 = request_metrics if request_metrics is not None else {"request_id": request_id}
    if "request_id" not in chunk_metrics:
        chunk_metrics["request_id"] = request_id
    if "chunks_processed" not in chunk_metrics:
        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,
                request_id=request_id,
                request_metrics=chunk_metrics,
            ):
                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,
                    request_id=request_id,
                    request_metrics=chunk_metrics,
                ):
                    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,
                    request_id=request_id,


sed -n '1160,1320p' veena3modal/services/tts_runtime.py


Output:
Chunk ID: cac72c
Wall time: 0.0514 seconds
Process exited with code 0
Original token count: 145
Output:
                    request_id=request_id,
                    request_metrics=chunk_metrics,
                ):
                    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


sed -n '1,260p' veena3modal/api/schemas.py


Output:
Chunk ID: 47b574
Wall time: 0.0511 seconds
Process exited with code 0
Original token count: 1906
Output:
"""
Pydantic schemas for TTS API requests/responses.

Ported from veena3srv/apps/api/serializers.py with full validation parity.
No Django dependencies - pure Pydantic for Modal deployment.
"""

from __future__ import annotations

import re
from enum import Enum
from typing import Any, Dict, List, Optional, Union
from uuid import UUID

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

# === Constants (mirrored from veena3srv/apps/inference/constants.py) ===
# These are duplicated here to avoid Django import chain; must stay in sync.

MAX_TEXT_LENGTH = 50000

# Spark TTS Speaker System - 12 Speakers
# Model: BayAreaBoys/spark_tts_4speaker (HuggingFace) has 12 speaker tokens
SPEAKER_MAP = {
    "lipakshi": 0, "vardan": 1, "reet": 2, "Nandini": 3,
    "krishna": 4, "anika": 5, "adarsh": 6, "Nilay": 7,
    "Aarvi": 8, "Asha": 9, "Bittu": 10, "Mira": 11,
}

# Friendly speaker name mappings (user-facing → internal)
FRIENDLY_SPEAKER_MAP = {
    "Mitra": "lipakshi", "Aaranya": "reet", "Taru": "Nandini",
    "Neer": "Nilay", "Dhruva": "vardan", "Ira": "anika",
    "Veda": "adarsh", "Aria": "krishna",
    "Aarvi": "Aarvi", "Asha": "Asha", "Bittu": "Bittu", "Mira": "Mira",
}

INDIC_SPEAKERS = list(SPEAKER_MAP.keys())
ALL_SPEAKER_NAMES = INDIC_SPEAKERS + list(FRIENDLY_SPEAKER_MAP.keys())

# Emotion tags (Spark TTS bracket format)
INDIC_EMOTION_TAGS = [
    "[angry]", "[curious]", "[excited]", "[giggle]", "[laughs harder]",
    "[laughs]", "[screams]", "[sighs]", "[sings]", "[whispers]"
]

# Legacy emotion tag mapping (<angle> → [bracket])
LEGACY_EMOTION_MAP = {
    "<angry>": "[angry]", "<curious>": "[curious]", "<excited>": "[excited]",
    "<giggle>": "[giggle]", "<laugh_harder>": "[laughs harder]",
    "<laugh>": "[laughs]", "<scream>": "[screams]", "<sigh>": "[sighs]",
    "<sing>": "[sings]", "<whisper>": "[whispers]"
}

# === Helper Functions ===

def resolve_speaker_name(name: str) -> str:
    """
    Resolve friendly speaker name to internal name.

    Examples:
        resolve_speaker_name("Mitra") -> "lipakshi"
        resolve_speaker_name("lipakshi") -> "lipakshi"

    Raises:
        ValueError: if name is not a valid speaker name
    """
    if name in FRIENDLY_SPEAKER_MAP:
        return FRIENDLY_SPEAKER_MAP[name]
    if name in INDIC_SPEAKERS:
        return name
    raise ValueError(
        f"Invalid speaker name: {name}. Valid names: {', '.join(ALL_SPEAKER_NAMES)}"
    )

# Emotion normalization mapping (various formats → [emotion])
_EMOTION_TO_TAG = {
    **{k.strip('<>'): v for k, v in LEGACY_EMOTION_MAP.items()},
    'laughing': '[laughs]', 'laugh': '[laughs]', 'laughs harder': '[laughs harder]',
    'laugh harder': '[laughs harder]', 'laugh_harder': '[laughs harder]',
    'sighs': '[sighs]', 'sigh': '[sighs]', 'giggles': '[giggle]', 'giggle': '[giggle]',
    'angry': '[angry]', 'excited': '[excited]', 'whispers': '[whispers]',
    'whisper': '[whispers]', 'screaming': '[screams]', 'scream': '[screams]',
    'screams': '[screams]', 'singing': '[sings]', 'sing': '[sings]',
    'sings': '[sings]', 'curious': '[curious]',
}

def normalize_emotion_tags(text: str) -> str:
    """
    Normalize emotion tags to Spark TTS bracket format.

    Conversions:
    - <laugh> → [laughs]
    - [laughing] → [laughs]
    - etc.
    """
    # Convert old <emotion> format to [emotion]
    angle_pattern = r'<([a-z_]+)>'

    def replace_angle(match):
        emotion = match.group(1).lower().strip()
        if emotion in _EMOTION_TO_TAG:
            return _EMOTION_TO_TAG[emotion]
        old_tag = f'<{emotion}>'
        if old_tag in LEGACY_EMOTION_MAP:
            return LEGACY_EMOTION_MAP[old_tag]
        return f'[{emotion}]'  # Unknown, convert to bracket anyway

    text = re.sub(angle_pattern, replace_angle, text)

    # Normalize [emotion] variants
    bracket_pattern = r'\[([^\]]+)\]'

    def replace_bracket(match):
        emotion = match.group(1).lower().strip()
        bracket_tag = f'[{emotion}]'
        if bracket_tag in INDIC_EMOTION_TAGS:
            return bracket_tag
        if emotion in _EMOTION_TO_TAG:
            return _EMOTION_TO_TAG[emotion]
        # Try singular/plural
        if emotion.endswith('s') and emotion[:-1] in _EMOTION_TO_TAG:
            return _EMOTION_TO_TAG[emotion[:-1]]
        if (emotion + 's') in _EMOTION_TO_TAG:
            return _EMOTION_TO_TAG[emotion + 's']
        return bracket_tag  # Keep as-is

    return re.sub(bracket_pattern, replace_bracket, text)

# === Enums ===

class AudioFormat(str, Enum):
    """Supported audio output formats."""
    WAV = "wav"
    OPUS = "opus"
    MP3 = "mp3"
    MULAW = "mulaw"
    FLAC = "flac"

class OutputSampleRate(str, Enum):
    """Output sample rate options."""
    SR_16KHZ = "16khz"
    SR_48KHZ = "48khz"

# === Request Schema ===

class TTSGenerateRequest(BaseModel):
    """
    Request schema for /v1/tts/generate.

    Mirrors validation rules from veena3srv/apps/api/serializers.py.
    """

    # Required
    text: str = Field(
        ...,
        min_length=1,
        max_length=MAX_TEXT_LENGTH,
        description="Text to synthesize (max 50,000 characters)"
    )

    # Speaker (required for indic_speakers model)
    speaker: Optional[str] = Field(
        None,
        max_length=50,
        description="Speaker name (internal or friendly)"
    )

    # Voice design (not currently supported)
    description: Optional[str] = Field(
        None,
        max_length=2000,
        description="Voice description (for voiceDesign model, NOT currently supported)"
    )
    voice_id: Optional[UUID] = Field(
        None,
        description="UUID of pre-created voice profile (for voiceDesign model)"
    )

    # Generation options
    model: Optional[str] = Field(
        None,
        max_length=50,
        description="Model to use (optional). Default: indic_speakers"
    )
    seed: Optional[int] = Field(
        None,
        ge=0,
        le=2**31 - 1,
        description="Random seed for reproducibility (0-2147483647)"
    )
    stream: bool = Field(
        False,
        description="Stream audio as it's generated"
    )
    format: AudioFormat = Field(
        AudioFormat.WAV,
        description="Audio output format"
    )
    sample_rate: int = Field(
        16000,
        description="Sample rate in Hz (default: 16000)"
    )

    # Advanced parameters
    temperature: float = Field(
        0.4,
        ge=0.0,
        le=2.0,
        description="Sampling temperature (0.0-2.0)"
    )
    top_k: int = Field(
        50,
        ge=1,
        le=100,
        description="Top-k sampling (1-100)"
    )
    top_p: float = Field(
        1.0,
        ge=0.0,
        le=1.0,
        description="Nucleus sampling (0.0-1.0)"
    )
    max_tokens: int = Field(
        4096,
        ge=128,
        le=4096,
        description="Maximum BiCodec tokens to generate"
    )
    repetition_penalty: float = Field(
        1.05,
        ge=1.0,
        le=2.0,
        description="Repetition penalty (1.0-2.0)"
    )

    # Preprocessing toggles
    normalize: bool = Field(
        True,
        description="Apply text normalization before TTS"
    )
    normalize_verbose: bool = Field(
        False,
        description="Return normalized text in X-Normalized-Text header"
    )
    chunking: bool = Field(
        True,
        description="Enable intelligent text chunking for long inputs"
    )

    # Output options
    output: OutputSampleRate = Field(
        OutputSampleRate.SR_16KHZ,


sed -n '1,260p' veena3modal/local_server.py


Output:
Chunk ID: bfca85
Wall time: 0.0514 seconds
Process exited with code 0
Original token count: 2380
Output:
"""
Local entry point for Veena3 TTS - runs the same FastAPI app without Modal.

Replicates what Modal's @modal.enter + @modal.asgi_app does:
1. Sets up PYTHONPATH for vendored deps (sparktts, AP-BWE)
2. Downloads model weights from HuggingFace if not present
3. Initializes TTS runtime (vLLM engine, BiCodec decoder, pipelines)
4. Serves the FastAPI app via uvicorn

Usage:
    # Basic (auto-downloads model, uses defaults)
    python -m veena3modal.local_server

    # Custom model path (skip download)
    python -m veena3modal.local_server --model-path /path/to/spark_tts_4speaker

    # Custom GPU memory + port
    python -m veena3modal.local_server --gpu-memory 0.5 --port 8080

    # With super-resolution (48kHz output)
    python -m veena3modal.local_server --enable-sr --sr-path /path/to/ap_bwe/16kto48k

    # CPU-only (no GPU, for testing only - very slow)
    python -m veena3modal.local_server --device cpu

Environment Variables (override defaults):
    MODEL_PATH              - Path to Spark TTS model directory
    HF_TOKEN                - HuggingFace token for private models
    AP_BWE_CHECKPOINT_DIR   - Path to super-resolution checkpoints
    AUTH_BYPASS_MODE        - "true" to disable auth (default for local)
    GPU_MEMORY_UTILIZATION  - vLLM GPU memory fraction (default: 0.25)
"""

from __future__ import annotations

import argparse
import logging
import os
import sys
import time

# === Path Setup (MUST happen before any veena3modal imports) ===
# Vendored deps: sparktts and AP-BWE live in external/ relative to repo root
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
EXTERNAL_DIR = os.path.join(REPO_ROOT, "external")

# sparktts package path: external/sparktts (contains sparktts/ subpackage)
SPARKTTS_PATH = os.path.join(EXTERNAL_DIR, "sparktts")
# AP-BWE path: external/AP-BWE
AP_BWE_PATH = os.path.join(EXTERNAL_DIR, "AP-BWE")

for path in [SPARKTTS_PATH, AP_BWE_PATH, REPO_ROOT]:
    if path not in sys.path:
        sys.path.insert(0, path)

# Also set PYTHONPATH for subprocesses (e.g., ffmpeg wrappers)
existing_pythonpath = os.environ.get("PYTHONPATH", "")
new_paths = f"{SPARKTTS_PATH}:{AP_BWE_PATH}:{REPO_ROOT}"
if existing_pythonpath:
    os.environ["PYTHONPATH"] = f"{new_paths}:{existing_pythonpath}"
else:
    os.environ["PYTHONPATH"] = new_paths

logger = logging.getLogger("veena3.local")

# Default local model directory (relative to repo)
DEFAULT_LOCAL_MODEL_DIR = os.path.join(REPO_ROOT, "models", "spark_tts_4speaker")

# Model sources (tried in order):
# 1. Local path (--model-path or MODEL_PATH env)
# 2. HuggingFace download (private repo: BayAreaBoys/spark_tts_4speaker)
# 3. Base Spark TTS model (public: SparkAudio/Spark-TTS-0.5B)
PRIVATE_HF_REPO = "BayAreaBoys/spark_tts_4speaker"
PUBLIC_FALLBACK_REPO = "SparkAudio/Spark-TTS-0.5B"

def download_model(
    hf_repo: str = PRIVATE_HF_REPO,
    local_dir: str = DEFAULT_LOCAL_MODEL_DIR,
    hf_token: str | None = None,
) -> str:
    """
    Download Spark TTS model from HuggingFace if not already present.

    Strategy:
    1. Try private repo (BayAreaBoys/spark_tts_4speaker) with HF token
    2. Fall back to public base model (SparkAudio/Spark-TTS-0.5B)
    3. Provide instructions for Modal volume download if both fail

    Returns the resolved local path.
    """
    # Check for any valid marker file (config.json at root, config.yaml, or LLM/config.json)
    marker_files = ["config.json", "config.yaml", "LLM/config.json"]
    if any(os.path.exists(os.path.join(local_dir, m)) for m in marker_files):
        logger.info(f"Model already present at {local_dir}")
        return local_dir

    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        logger.error(
            "huggingface_hub not installed. Install with: pip install huggingface-hub\n"
            "Or download model manually and pass --model-path"
        )
        sys.exit(1)

    os.makedirs(local_dir, exist_ok=True)

    # Try private repo first
    logger.info(f"Attempting download from {hf_repo}...")
    try:
        downloaded_path = snapshot_download(
            repo_id=hf_repo,
            local_dir=local_dir,
            token=hf_token,
            ignore_patterns=["*.bin", "training_args.bin", "optimizer.pt"],
        )
        logger.info(f"Model downloaded to {downloaded_path}")
        return downloaded_path
    except Exception as e:
        logger.warning(f"Private repo download failed: {e}")

    # Try public fallback
    if hf_repo != PUBLIC_FALLBACK_REPO:
        logger.info(f"Trying public base model: {PUBLIC_FALLBACK_REPO}...")
        try:
            downloaded_path = snapshot_download(
                repo_id=PUBLIC_FALLBACK_REPO,
                local_dir=local_dir,
                token=hf_token,
            )
            logger.info(f"Base model downloaded to {downloaded_path}")
            logger.warning(
                "Using public base SparkTTS model (no custom speakers). "
                "For full speaker support, download your fine-tuned model."
            )
            return downloaded_path
        except Exception as e2:
            logger.warning(f"Public model download also failed: {e2}")

    # All downloads failed - provide manual instructions
    logger.error(
        "\nModel download failed. Options:\n"
        "  1. Download from Modal volume:\n"
        "     modal volume get veena3-models spark_tts_4speaker models/spark_tts_4speaker\n\n"
        "  2. Download from HuggingFace (if you have access):\n"
        f"     huggingface-cli download {PRIVATE_HF_REPO} --local-dir {local_dir}\n\n"
        "  3. Use base SparkTTS model:\n"
        f"     huggingface-cli download {PUBLIC_FALLBACK_REPO} --local-dir {local_dir}\n\n"
        "  4. Specify a custom path:\n"
        "     python -m veena3modal.local_server --model-path /path/to/model"
    )
    sys.exit(1)

def resolve_model_paths(model_path: str) -> tuple[str, str]:
    """
    Resolve LLM and BiCodec paths from a Spark TTS model directory.

    Spark TTS model structure (from SparkAudio/Spark-TTS-0.5B):
      model_path/
        LLM/          <- Language model (vLLM loads this)
        BiCodec/      <- Audio tokenizer/decoder
        config.yaml   <- Top-level config

    Fine-tuned models may have LLM files at root (flat structure).
    This function auto-detects both layouts.

    Returns:
        (llm_path, bicodec_path) resolved paths
    """
    # Check for subdirectory layout (SparkAudio/Spark-TTS-0.5B style)
    llm_subdir = os.path.join(model_path, "LLM")
    bicodec_subdir = os.path.join(model_path, "BiCodec")

    if os.path.isdir(llm_subdir) and os.path.exists(os.path.join(llm_subdir, "config.json")):
        llm_path = llm_subdir
    else:
        # Flat layout: LLM files at root
        llm_path = model_path

    # BiCodec path is always the root (it contains BiCodec/ subdir)
    bicodec_path = model_path

    logger.info(f"Resolved model paths:")
    logger.info(f"  LLM (vLLM):    {llm_path}")
    logger.info(f"  BiCodec:       {bicodec_path}")

    return llm_path, bicodec_path

def initialize_local_runtime(
    model_path: str,
    device: str = "cuda",
    gpu_memory_utilization: float = 0.85,
    enable_sr: bool = False,
    sr_checkpoint_dir: str | None = None,
    hf_token: str | None = None,
    num_engines: int = 1,
):
    """
    Initialize TTS runtime for local serving.

    This replaces Modal's @modal.enter lifecycle hook.
    Loads: vLLM engine, BiCodec decoder, prompt builder, pipelines.
    """
    from veena3modal.services.tts_runtime import initialize_runtime

    # Resolve LLM vs BiCodec paths from model directory
    llm_path, bicodec_path = resolve_model_paths(model_path)

    start = time.time()
    logger.info("=" * 60)
    logger.info("Initializing Veena3 TTS Runtime (local)")
    logger.info(f"  Model:       {model_path}")
    logger.info(f"  LLM path:    {llm_path}")
    logger.info(f"  BiCodec:     {bicodec_path}")
    logger.info(f"  Device:      {device}")
    logger.info(f"  GPU Memory:  {gpu_memory_utilization:.0%} (total)")
    logger.info(f"  Engines:     {num_engines}")
    logger.info(f"  SR:          {'enabled' if enable_sr else 'disabled'}")
    logger.info("=" * 60)

    runtime = initialize_runtime(
        model_path=llm_path,
        bicodec_path=bicodec_path,
        sr_checkpoint_dir=sr_checkpoint_dir if enable_sr else None,
        device=device,
        hf_token=hf_token,
        gpu_memory_utilization=gpu_memory_utilization,
        enable_sr=enable_sr,
        num_engines=num_engines,
    )

    elapsed = time.time() - start
    logger.info(f"Runtime initialized in {elapsed:.1f}s")
    logger.info(f"  Model version: {runtime.model_version}")
    logger.info(f"  Streaming: {'available' if runtime.streaming_pipeline else 'unavailable'}")
    logger.info(f"  Super-resolution: {'loaded' if runtime.sr_service and runtime.sr_service.is_loaded else 'disabled'}")

    return runtime

def create_local_app():
    """
    Create the FastAPI app for local serving.

    Same app factory as Modal deployment, but called from local context.
    """
    from veena3modal.api.fastapi_app import create_app
    return create_app()

def parse_args():
    """Parse CLI arguments for local server configuration."""
    parser = argparse.ArgumentParser(
        description="Veena3 TTS Local Server - GPU-accelerated text-to-speech",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:


sed -n '260,520p' veena3modal/local_server.py


Output:
Chunk ID: 298366
Wall time: 0.0513 seconds
Process exited with code 0
Original token count: 1660
Output:
Examples:
  python -m veena3modal.local_server                          # Auto-download model, serve on :8000
  python -m veena3modal.local_server --port 8080              # Custom port
  python -m veena3modal.local_server --model-path ./my_model  # Use local model
  python -m veena3modal.local_server --gpu-memory 0.5         # Limit GPU memory
  python -m veena3modal.local_server --workers 2              # Multiple workers (careful: each loads model)
        """,
    )

    # Model configuration
    model_group = parser.add_argument_group("Model")
    model_group.add_argument(
        "--model-path",
        type=str,
        default=os.environ.get("MODEL_PATH", DEFAULT_LOCAL_MODEL_DIR),
        help=f"Path to Spark TTS model (default: {DEFAULT_LOCAL_MODEL_DIR})",
    )
    model_group.add_argument(
        "--hf-repo",
        type=str,
        default=PRIVATE_HF_REPO,
        help=f"HuggingFace repo to download from (default: {PRIVATE_HF_REPO})",
    )
    model_group.add_argument(
        "--hf-token",
        type=str,
        default=os.environ.get("HF_TOKEN"),
        help="HuggingFace token (default: from HF_TOKEN env var)",
    )
    model_group.add_argument(
        "--skip-download",
        action="store_true",
        help="Don't auto-download model (fail if not found)",
    )

    # GPU configuration
    gpu_group = parser.add_argument_group("GPU")
    gpu_group.add_argument(
        "--device",
        type=str,
        default="cuda",
        choices=["cuda", "cpu"],
        help="Device for inference (default: cuda)",
    )
    gpu_group.add_argument(
        "--gpu-memory",
        type=float,
        default=float(os.environ.get("GPU_MEMORY_UTILIZATION", "0.25")),
        help="vLLM GPU memory utilization fraction (default: 0.25)",
    )
    gpu_group.add_argument(
        "--num-engines",
        type=int,
        default=int(os.environ.get("NUM_ENGINES", "1")),
        help="Number of vLLM engine instances (default: 1, use 2-3 for multi-engine mode)",
    )

    # Super-resolution
    sr_group = parser.add_argument_group("Super-Resolution")
    sr_group.add_argument(
        "--enable-sr",
        action="store_true",
        default=False,
        help="Enable AP-BWE super-resolution (16kHz -> 48kHz)",
    )
    sr_group.add_argument(
        "--sr-path",
        type=str,
        default=os.environ.get("AP_BWE_CHECKPOINT_DIR"),
        help="Path to AP-BWE checkpoint directory",
    )

    # Server configuration
    server_group = parser.add_argument_group("Server")
    server_group.add_argument(
        "--host",
        type=str,
        default="0.0.0.0",
        help="Bind address (default: 0.0.0.0)",
    )
    server_group.add_argument(
        "--port",
        type=int,
        default=int(os.environ.get("PORT", "8000")),
        help="Port to serve on (default: 8000)",
    )
    server_group.add_argument(
        "--workers",
        type=int,
        default=1,
        help="Number of uvicorn workers (default: 1, each loads full model)",
    )
    server_group.add_argument(
        "--reload",
        action="store_true",
        default=False,
        help="Enable auto-reload for development (incompatible with --workers > 1)",
    )

    # Auth
    auth_group = parser.add_argument_group("Auth")
    auth_group.add_argument(
        "--auth",
        action="store_true",
        default=False,
        help="Enable API key authentication (default: disabled for local)",
    )

    # Logging
    log_group = parser.add_argument_group("Logging")
    log_group.add_argument(
        "--log-level",
        type=str,
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        help="Log level (default: INFO)",
    )

    return parser.parse_args()

def main():
    """
    Main entry point for local TTS server.

    Lifecycle:
    1. Parse args + configure logging
    2. Download model if needed
    3. Initialize TTS runtime (vLLM engine, decoders, pipelines)
    4. Create FastAPI app
    5. Run uvicorn
    """
    args = parse_args()

    # Configure logging
    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )

    # Auth bypass for local development (unless --auth flag is set)
    if not args.auth:
        os.environ["AUTH_BYPASS_MODE"] = "true"
        logger.info("Auth bypass enabled (local mode). Use --auth to require API keys.")
    else:
        os.environ.setdefault("AUTH_BYPASS_MODE", "false")

    # Resolve model path
    # Spark TTS root has config.yaml; LLM subdir has config.json; fine-tuned may have config.json at root
    model_path = args.model_path
    model_exists = any(
        os.path.exists(os.path.join(model_path, marker))
        for marker in ("config.json", "config.yaml", "LLM/config.json")
    )
    if not model_exists:
        if args.skip_download:
            logger.error(f"Model not found at {model_path} and --skip-download is set")
            sys.exit(1)
        model_path = download_model(
            hf_repo=args.hf_repo,
            local_dir=model_path,
            hf_token=args.hf_token,
        )

    # Initialize TTS runtime (this is the expensive step - loads model onto GPU)
    initialize_local_runtime(
        model_path=model_path,
        device=args.device,
        gpu_memory_utilization=args.gpu_memory,
        enable_sr=args.enable_sr,
        sr_checkpoint_dir=args.sr_path,
        hf_token=args.hf_token,
        num_engines=args.num_engines,
    )

    # Run uvicorn
    # NOTE: We use the factory string so uvicorn can find the app.
    # With --reload, uvicorn reimports, but the runtime singleton persists in memory.
    import uvicorn

    logger.info("")
    logger.info("=" * 60)
    logger.info(f"Veena3 TTS Local Server starting")
    logger.info(f"  URL: http://{args.host}:{args.port}")
    logger.info(f"  Health: http://{args.host}:{args.port}/v1/tts/health")
    logger.info(f"  Generate: POST http://{args.host}:{args.port}/v1/tts/generate")
    logger.info(f"  Metrics: http://{args.host}:{args.port}/v1/tts/metrics")
    logger.info(f"  Docs: http://{args.host}:{args.port}/docs")
    logger.info("=" * 60)

    uvicorn.run(
        "veena3modal.local_server:create_local_app",
        factory=True,
        host=args.host,
        port=args.port,
        workers=args.workers,
        reload=args.reload,
        log_level=args.log_level.lower(),
        # Performance: keep-alive for connection reuse under load
        timeout_keep_alive=30,
    )

if __name__ == "__main__":
    main()


sed -n '1,380p' scripts/stress_test_local.py