Output:
Chunk ID: 62c500
Wall time: 0.0515 seconds
Process exited with code 0
Original token count: 1813
Output:
"""
Structured logging for Veena3 TTS service.

Features:
- JSON-formatted logs for structured queries
- Request ID tracking via context
- Lifecycle event helpers (request_received, first_audio_emitted, etc.)
- NO PII in logs (text content stored in Supabase, not logs)

Usage:
    from veena3modal.shared.logging import get_logger, log_event, set_request_context

    logger = get_logger(__name__)
    set_request_context(request_id="req-123")
    log_event(logger, "request_received", text_length=100, speaker="lipakshi")
"""

import logging
import json
import sys
import traceback
from datetime import datetime, timezone
from typing import Any, Dict, Optional
from contextvars import ContextVar

# Context variable for request-scoped data
_request_context: ContextVar[Dict[str, Any]] = ContextVar('request_context', default={})

def set_request_context(request_id: Optional[str] = None, **kwargs) -> None:
    """
    Set request context for current async task.

    Args:
        request_id: Unique request identifier
        **kwargs: Additional context fields (user_id, etc.)
    """
    ctx = {"request_id": request_id, **kwargs}
    _request_context.set(ctx)

def get_request_context() -> Dict[str, Any]:
    """Get current request context."""
    return _request_context.get()

def clear_request_context() -> None:
    """Clear request context after request completes."""
    _request_context.set({})

class JSONFormatter(logging.Formatter):
    """
    JSON log formatter for structured logging.

    Output format:
    {
        "timestamp": "2024-01-01T12:00:00.000Z",
        "level": "INFO",
        "logger": "veena3modal.api",
        "message": "Request received",
        "request_id": "req-123",
        ...extra fields...
    }
    """

    # Fields to exclude from extra (already handled separately)
    RESERVED_ATTRS = {
        'args', 'asctime', 'created', 'exc_info', 'exc_text', 'filename',
        'funcName', 'levelname', 'levelno', 'lineno', 'module', 'msecs',
        'message', 'msg', 'name', 'pathname', 'process', 'processName',
        'relativeCreated', 'stack_info', 'thread', 'threadName',
        'taskName',  # Python 3.12+
    }

    def format(self, record: logging.LogRecord) -> str:
        """Format log record as JSON."""
        # Build base log entry
        log_entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage(),
        }

        # Add request context if available
        ctx = get_request_context()
        if ctx.get("request_id"):
            log_entry["request_id"] = ctx["request_id"]

        # Add extra fields from record
        for key, value in record.__dict__.items():
            if key not in self.RESERVED_ATTRS and not key.startswith('_'):
                try:
                    # Ensure value is JSON serializable
                    json.dumps(value)
                    log_entry[key] = value
                except (TypeError, ValueError):
                    log_entry[key] = str(value)

        # Add exception info if present
        if record.exc_info:
            log_entry["exception"] = ''.join(traceback.format_exception(*record.exc_info))

        return json.dumps(log_entry, default=str)

def get_logger(name: str) -> logging.Logger:
    """
    Get a logger configured for structured JSON output.

    Args:
        name: Logger name (typically __name__)

    Returns:
        Configured logger instance
    """
    logger = logging.getLogger(name)

    # Only configure if no handlers exist (avoid duplicate setup)
    if not logger.handlers:
        handler = logging.StreamHandler(sys.stdout)
        handler.setFormatter(JSONFormatter())
        logger.addHandler(handler)
        logger.setLevel(logging.INFO)

    return logger

def log_event(
    logger: logging.Logger,
    event_type: str,
    level: int = logging.INFO,
    request_id: Optional[str] = None,
    **kwargs
) -> None:
    """
    Log a structured event.

    Args:
        logger: Logger instance
        event_type: Event name (e.g., "request_received", "first_audio_emitted")
        level: Log level (default: INFO)
        request_id: Override request_id from context
        **kwargs: Event-specific fields
    """
    extra = {"event": event_type, **kwargs}

    if request_id:
        extra["request_id"] = request_id

    logger.log(level, event_type, extra=extra)

def create_lifecycle_event(
    event_type: str,
    request_id: str,
    **kwargs
) -> Dict[str, Any]:
    """
    Create a lifecycle event dictionary.

    Lifecycle events:
    - request_received: TTS request started
    - auth_validated: API key validated
    - generation_started: Model inference began
    - first_audio_emitted: First audio chunk sent (TTFB)
    - request_completed: Request finished successfully
    - request_failed: Request failed with error

    Args:
        event_type: Event name
        request_id: Request identifier
        **kwargs: Event-specific fields

    Returns:
        Event dictionary ready for logging
    """
    return {
        "event": event_type,
        "request_id": request_id,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        **kwargs
    }

# Pre-configured logger for the TTS service
tts_logger = get_logger("veena3modal.tts")

def log_request_received(
    request_id: str,
    text_length: int,
    speaker: str,
    stream: bool,
    format: str,
    **kwargs
) -> None:
    """Log request_received lifecycle event."""
    log_event(
        tts_logger,
        "request_received",
        request_id=request_id,
        text_length=text_length,
        speaker=speaker,
        stream=stream,
        format=format,
        **kwargs
    )

def log_first_audio_emitted(
    request_id: str,
    ttfb_ms: int,
    chunk_size_bytes: int,
    **kwargs
) -> None:
    """Log first_audio_emitted lifecycle event (TTFB marker)."""
    log_event(
        tts_logger,
        "first_audio_emitted",
        request_id=request_id,
        ttfb_ms=ttfb_ms,
        chunk_size_bytes=chunk_size_bytes,
        **kwargs
    )

def log_request_completed(
    request_id: str,
    status_code: int,
    total_duration_ms: int,
    audio_duration_seconds: float,
    rtf: float,
    chunks_sent: int = 1,
    **kwargs
) -> None:
    """Log request_completed lifecycle event."""
    log_event(
        tts_logger,
        "request_completed",
        request_id=request_id,
        status_code=status_code,
        total_duration_ms=total_duration_ms,
        audio_duration_seconds=audio_duration_seconds,
        rtf=rtf,
        chunks_sent=chunks_sent,
        **kwargs
    )

def log_request_failed(
    request_id: str,
    status_code: int,
    error_code: str,
    error_message: str,
    **kwargs
) -> None:
    """Log request_failed lifecycle event."""
    log_event(
        tts_logger,
        "request_failed",
        level=logging.ERROR,
        request_id=request_id,
        status_code=status_code,
        error_code=error_code,
        error_message=error_message,
        **kwargs
    )


Output:
Chunk ID: f2c72a
Wall time: 0.0507 seconds
Process exited with code 0
Original token count: 2084
Output:
"""
Prometheus metrics for Veena3 TTS service.

Features:
- Request counters (by speaker, format, status)
- TTFB histogram
- RTF histogram
- Audio duration histogram
- Model load time gauge

Usage:
    from veena3modal.shared.metrics import record_request_received, record_ttfb

    record_request_received(speaker="lipakshi", stream=True, format="wav")
    record_ttfb(ttfb_seconds=0.25, speaker="lipakshi", stream=True)
"""

import re
from typing import Optional

# Try to import prometheus_client, fall back to no-op if not installed
try:
    from prometheus_client import (
        Counter, Histogram, Gauge,
        CollectorRegistry, generate_latest, CONTENT_TYPE_LATEST
    )
    PROMETHEUS_AVAILABLE = True
except ImportError:
    PROMETHEUS_AVAILABLE = False

# Module-level registry (singleton)
_registry: Optional['CollectorRegistry'] = None

# Metrics instances (lazy initialized)
_requests_total: Optional['Counter'] = None
_requests_completed: Optional['Counter'] = None
_requests_failed: Optional['Counter'] = None
_ttfb_seconds: Optional['Histogram'] = None
_rtf: Optional['Histogram'] = None
_audio_duration_seconds: Optional['Histogram'] = None
_chunks_sent: Optional['Histogram'] = None
_model_load_seconds: Optional['Gauge'] = None
_model_loaded: Optional['Gauge'] = None

def get_metrics_registry() -> Optional['CollectorRegistry']:
    """Get or create the metrics registry."""
    global _registry

    if not PROMETHEUS_AVAILABLE:
        return None

    if _registry is None:
        _registry = CollectorRegistry()
        _initialize_metrics()

    return _registry

def _initialize_metrics() -> None:
    """Initialize all metrics with the registry."""
    global _requests_total, _requests_completed, _requests_failed
    global _ttfb_seconds, _rtf, _audio_duration_seconds, _chunks_sent
    global _model_load_seconds, _model_loaded

    if not PROMETHEUS_AVAILABLE or _registry is None:
        return

    # Request counters
    _requests_total = Counter(
        'veena3_tts_requests_total',
        'Total TTS requests received',
        ['speaker', 'stream', 'format'],
        registry=_registry
    )

    _requests_completed = Counter(
        'veena3_tts_requests_completed_total',
        'Total TTS requests completed successfully',
        ['speaker', 'stream', 'status_code'],
        registry=_registry
    )

    _requests_failed = Counter(
        'veena3_tts_requests_failed_total',
        'Total TTS requests failed',
        ['speaker', 'error_code', 'status_code'],
        registry=_registry
    )

    # Performance histograms
    _ttfb_seconds = Histogram(
        'veena3_tts_ttfb_seconds',
        'Time to first byte in seconds',
        ['speaker', 'stream'],
        buckets=(0.1, 0.25, 0.5, 0.75, 1.0, 1.5, 2.0, 3.0, 5.0, 10.0),
        registry=_registry
    )

    _rtf = Histogram(
        'veena3_tts_rtf',
        'Real-time factor (generation_time / audio_duration)',
        ['speaker'],
        buckets=(0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.5, 2.0),
        registry=_registry
    )

    _audio_duration_seconds = Histogram(
        'veena3_tts_audio_duration_seconds',
        'Generated audio duration in seconds',
        ['speaker'],
        buckets=(1, 2, 5, 10, 20, 30, 60, 120, 300),
        registry=_registry
    )

    _chunks_sent = Histogram(
        'veena3_tts_chunks_sent',
        'Number of audio chunks sent in streaming mode',
        ['speaker'],
        buckets=(1, 2, 5, 10, 20, 50, 100),
        registry=_registry
    )

    # Model metrics
    _model_load_seconds = Gauge(
        'veena3_tts_model_load_seconds',
        'Time taken to load the model',
        ['model_version'],
        registry=_registry
    )

    _model_loaded = Gauge(
        'veena3_tts_model_loaded',
        'Whether the model is loaded (1) or not (0)',
        ['model_version'],
        registry=_registry
    )

def sanitize_label(value: str) -> str:
    """
    Sanitize a string for use as a Prometheus label.

    Prometheus labels must match [a-zA-Z_][a-zA-Z0-9_]*
    """
    # Replace non-alphanumeric chars (except underscore) with underscore
    sanitized = re.sub(r'[^a-zA-Z0-9_]', '_', str(value))
    # Ensure doesn't start with a digit
    if sanitized and sanitized[0].isdigit():
        sanitized = '_' + sanitized
    return sanitized

def record_request_received(
    speaker: str,
    stream: bool,
    format: str,
) -> None:
    """Record a request received."""
    if _requests_total is None:
        get_metrics_registry()  # Initialize if needed

    if _requests_total is not None:
        _requests_total.labels(
            speaker=sanitize_label(speaker),
            stream=str(stream).lower(),
            format=sanitize_label(format),
        ).inc()

def record_request_completed(
    status_code: int,
    duration_seconds: float,
    speaker: str,
    stream: bool,
) -> None:
    """Record a successful request completion."""
    if _requests_completed is None:
        get_metrics_registry()

    if _requests_completed is not None:
        _requests_completed.labels(
            speaker=sanitize_label(speaker),
            stream=str(stream).lower(),
            status_code=str(status_code),
        ).inc()

def record_request_failed(
    status_code: int,
    error_code: str,
    speaker: str,
) -> None:
    """Record a failed request."""
    if _requests_failed is None:
        get_metrics_registry()

    if _requests_failed is not None:
        _requests_failed.labels(
            speaker=sanitize_label(speaker),
            error_code=sanitize_label(error_code),
            status_code=str(status_code),
        ).inc()

def record_ttfb(
    ttfb_seconds: float,
    speaker: str,
    stream: bool,
) -> None:
    """Record time to first byte."""
    if _ttfb_seconds is None:
        get_metrics_registry()

    if _ttfb_seconds is not None:
        _ttfb_seconds.labels(
            speaker=sanitize_label(speaker),
            stream=str(stream).lower(),
        ).observe(ttfb_seconds)

def record_rtf(
    rtf: float,
    speaker: str,
) -> None:
    """Record real-time factor."""
    if _rtf is None:
        get_metrics_registry()

    if _rtf is not None:
        _rtf.labels(
            speaker=sanitize_label(speaker),
        ).observe(rtf)

def record_audio_duration(
    duration_seconds: float,
    speaker: str,
) -> None:
    """Record generated audio duration."""
    if _audio_duration_seconds is None:
        get_metrics_registry()

    if _audio_duration_seconds is not None:
        _audio_duration_seconds.labels(
            speaker=sanitize_label(speaker),
        ).observe(duration_seconds)

def record_chunks_sent(
    chunks: int,
    speaker: str,
) -> None:
    """Record number of chunks sent in streaming."""
    if _chunks_sent is None:
        get_metrics_registry()

    if _chunks_sent is not None:
        _chunks_sent.labels(
            speaker=sanitize_label(speaker),
        ).observe(chunks)

def record_model_load_time(
    duration_seconds: float,
    model_version: str,
) -> None:
    """Record model loading time."""
    if _model_load_seconds is None:
        get_metrics_registry()

    if _model_load_seconds is not None:
        _model_load_seconds.labels(
            model_version=sanitize_label(model_version),
        ).set(duration_seconds)

def set_model_loaded(
    loaded: bool,
    model_version: str,
) -> None:
    """Set model loaded status."""
    if _model_loaded is None:
        get_metrics_registry()

    if _model_loaded is not None:
        _model_loaded.labels(
            model_version=sanitize_label(model_version),
        ).set(1 if loaded else 0)

def get_metrics_text() -> str:
    """
    Get metrics in Prometheus text format.

    Returns:
        Prometheus text format metrics string
    """
    registry = get_metrics_registry()

    if not PROMETHEUS_AVAILABLE or registry is None:
        return "# Prometheus client not available\n"

    return generate_latest(registry).decode('utf-8')

def get_content_type() -> str:
    """Get the content type for Prometheus metrics."""
    if PROMETHEUS_AVAILABLE:
        return CONTENT_TYPE_LATEST
    return "text/plain; charset=utf-8"


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


**Planning detailed runtime instrumentation**


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


sed -n '280,620p' veena3modal/local_server.py


Output:
Chunk ID: 5791d2
Wall time: 0.0518 seconds
Process exited with code 0
Original token count: 2599
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:
  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,


Output:
Chunk ID: a6f6cb
Wall time: 0.0508 seconds
Process exited with code 0
Original token count: 1447
Output:
        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()


Summary: auto


**Planning detailed request metrics**


I’ve identified the core gap: the benchmark script expects detailed timing headers, but the API currently only returns basic totals. Next I’ll instrument the non-streaming inference path (LLM generation, incremental decode/parse, BiCodec GPU decode, inflight concurrency) and expose it via response headers/`x-perf-details`, then rerun stress tests.


sed -n '1,420p' veena3modal/core/bicodec_decoder.py


sed -n '1,380p' veena3modal/core/bicodec_batch_decoder.py


sed -n '1,320p' veena3modal/core/token_utils.py


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


sed -n '260,620p' veena3modal/processing/long_text_processor.py


Output:
Chunk ID: daabc1
Wall time: 0.0511 seconds
Process exited with code 0
Original token count: 3121
Output:
"""
Spark TTS BiCodec Decoder

Decodes BiCodec audio tokens (semantic + global) to audio waveforms.
Replaces SNAC decoder for Spark TTS model.
"""

import torch
import numpy as np
import sys
from pathlib import Path
from typing import List, Optional, Tuple

# Ensure `sparktts` (vendored under external/) is importable in local/dev runs.
# In Modal, PYTHONPATH already includes `/root/external/sparktts`.
try:
    from sparktts.models.audio_tokenizer import BiCodecTokenizer
except ImportError:  # pragma: no cover - environment-dependent
    _this_file = Path(__file__).resolve()
    _sparktts_path = None
    for _p in [_this_file.parent] + list(_this_file.parents):
        _candidate = _p / "external" / "sparktts"
        if _candidate.is_dir():
            _sparktts_path = str(_candidate)
            break
    if _sparktts_path and _sparktts_path not in sys.path:
        sys.path.insert(0, _sparktts_path)
    from sparktts.models.audio_tokenizer import BiCodecTokenizer

from veena3modal.core.constants import (
    BICODEC_TOKENIZER_PATH,
    AUDIO_SAMPLE_RATE,
)

class BiCodecDecoder:
    """
    BiCodec Decoder for Spark TTS.

    Decodes semantic and global tokens to audio waveforms using BiCodec.
    This replaces the SNAC decoder used in the old Indic Orpheus model.

    Supports streaming with sliding window approach (like SNAC).
    """

    def __init__(
        self,
        device: str = "cuda",
        model_path: str = BICODEC_TOKENIZER_PATH,
        enable_batching: bool = False,  # For compatibility with streaming pipeline
    ):
        """
        Initialize BiCodec decoder.

        Args:
            device: Device for BiCodec model (cuda/cpu)
            model_path: Path to BiCodec model checkpoint
            enable_batching: Enable async batching for streaming (not used yet)
        """
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        self.model_path = model_path
        self.enable_batching = enable_batching  # For streaming pipeline compatibility

        print(f"🎵 Loading BiCodec audio tokenizer from {model_path}...")
        print(f"   Device: {self.device}")

        # Initialize BiCodec tokenizer
        self.audio_tokenizer = BiCodecTokenizer(model_path, device=str(self.device))

        # OPTIMIZATION: Ensure model is on device once at init, not per-decode call
        self.audio_tokenizer.device = self.device
        self.audio_tokenizer.model.to(self.device)

        # NOTE: torch.compile disabled for BiCodec -- dynamic semantic sequence lengths
        # cause repeated recompilation (~120s each), worse than eager mode.
        # torch.compile works best with static shapes. If needed, pre-warm with
        # representative lengths or use mode="default" instead of "reduce-overhead".
        # The multi-engine approach (Tier 3) provides more impactful gains.

        print(f"✅ BiCodec decoder initialized (sample rate: {AUDIO_SAMPLE_RATE}Hz)")
        if enable_batching:
            print(f"   ℹ  Batching enabled (for streaming support)")

    def decode(
        self,
        semantic_ids: List[int],
        global_ids: List[int],
    ) -> Optional[np.ndarray]:
        """
        Decode BiCodec tokens to audio waveform.

        CRITICAL: BiCodec expects EXACTLY 32 global tokens!
        The decoder pools them via speaker_encoder → d_vector, then broadcasts
        via d_vector.unsqueeze(-1) across the time dimension.

        Args:
            semantic_ids: List of semantic token IDs (variable length)
            global_ids: List of global token IDs (MUST be exactly 32!)

        Returns:
            Audio waveform as numpy array (float32, 16kHz mono)
            Shape: (samples,)
            Returns None if decode fails
        """
        if not semantic_ids or not global_ids:
            print(f"⚠  Empty token lists: semantic={len(semantic_ids)}, global={len(global_ids)}")
            return None

        # Validate global token count
        EXPECTED_GLOBAL_TOKENS = 32
        if len(global_ids) != EXPECTED_GLOBAL_TOKENS:
            print(f"❌ BiCodec requires EXACTLY {EXPECTED_GLOBAL_TOKENS} global tokens, got {len(global_ids)}")
            # Try to fix it
            if len(global_ids) > EXPECTED_GLOBAL_TOKENS:
                global_ids = global_ids[:EXPECTED_GLOBAL_TOKENS]
                print(f"   └─ Truncated to {EXPECTED_GLOBAL_TOKENS}")
            else:
                global_ids = global_ids + [0] * (EXPECTED_GLOBAL_TOKENS - len(global_ids))
                print(f"   └─ Padded to {EXPECTED_GLOBAL_TOKENS}")

        try:
            # Convert to tensors following official Spark-TTS implementation
            # global_token_ids shape: (1, 32) - ALWAYS 32!
            # pred_semantic_ids shape: (1, N) - variable length
            pred_semantic = torch.tensor(semantic_ids).long().unsqueeze(0).to(self.device)
            pred_global = torch.tensor(global_ids).long().unsqueeze(0).to(self.device)

            # Device migration done once in __init__, not per-call
            # Decode through BiCodec
            # BiCodecTokenizer.detokenize() expects:
            #   - global_tokens: (batch, global_dim) -> will unsqueeze(1) internally
            #   - semantic_tokens: (batch, latent_dim)
            with torch.inference_mode():
                wav_np = self.audio_tokenizer.detokenize(
                    pred_global,  # (1, seq) - correct shape!
                    pred_semantic  # (1, seq)
                )

            return wav_np

        except Exception as e:
            print(f"❌ BiCodec decode error: {e}")
            import traceback
            traceback.print_exc()
            return None

    def decode_to_bytes(
        self,
        semantic_ids: List[int],
        global_ids: List[int],
    ) -> Optional[bytes]:
        """
        Decode BiCodec tokens to audio bytes (int16 PCM).

        Args:
            semantic_ids: List of semantic token IDs
            global_ids: List of global token IDs

        Returns:
            Audio as bytes (int16 PCM, 16kHz mono)
            Returns None if decode fails
        """
        audio = self.decode(semantic_ids, global_ids)

        if audio is None:
            return None

        # Clip and convert float32 to int16 PCM
        audio = np.clip(audio, -1.0, 1.0)
        audio_int16 = (audio * 32767).astype(np.int16)

        return audio_int16.tobytes()

    def validate_tokens(self, semantic_ids: List[int], global_ids: List[int]) -> bool:
        """
        Validate BiCodec tokens before decoding.

        Args:
            semantic_ids: List of semantic token IDs
            global_ids: List of global token IDs

        Returns:
            True if valid, False otherwise
        """
        # Check minimum length
        if not semantic_ids:
            print(f"❌ No semantic tokens")
            return False

        if not global_ids:
            print(f"❌ No global tokens")
            return False

        # Basic sanity checks
        if len(semantic_ids) < 1:
            print(f"❌ Too few semantic tokens: {len(semantic_ids)}")
            return False

        if len(global_ids) < 1:
            print(f"❌ Too few global tokens: {len(global_ids)}")
            return False

        return True

    def get_audio_duration(self, semantic_ids: List[int]) -> float:
        """
        Estimate audio duration from semantic tokens.

        Args:
            semantic_ids: List of semantic token IDs

        Returns:
            Estimated duration in seconds
        """
        # BiCodec generates audio based on semantic tokens
        # Approximate duration calculation (may need tuning)
        # This is an estimate and actual duration depends on model output
        estimated_samples = len(semantic_ids) * 320  # Rough estimate
        return estimated_samples / AUDIO_SAMPLE_RATE

    def decode_streaming(
        self,
        semantic_ids: List[int],
        global_ids: List[int],
        use_sliding_window: bool = False,
        trim_warmup: bool = False,
    ) -> Optional[bytes]:
        """
        Decode BiCodec tokens with streaming support.

        CRITICAL: BiCodec decoder expects EXACTLY 32 global tokens always!
        The decoder internally broadcasts d_vector via d_vector.unsqueeze(-1).

        From sparktts/models/bicodec.py line 184-186:
            d_vector = self.speaker_encoder.detokenize(global_tokens)  # Expects 32 tokens
            x = self.prenet(z_q, d_vector)
            x = x + d_vector.unsqueeze(-1)  # Broadcasts across time automatically!

        Args:
            semantic_ids: List of semantic token IDs (variable length, 50 TPS)
            global_ids: List of global token IDs (MUST be exactly 32 tokens!)
            use_sliding_window: If True, return only middle samples
            trim_warmup: Legacy parameter (not used for BiCodec)

        Returns:
            Audio bytes (int16 PCM) or None if decode fails
        """
        # Minimum semantic tokens needed for stable decode
        MIN_SEMANTIC_TOKENS = 8

        if len(semantic_ids) < MIN_SEMANTIC_TOKENS:
            print(f"⚠  Too few semantic tokens ({len(semantic_ids)}) for decode, need >= {MIN_SEMANTIC_TOKENS}")
            return None

        # CRITICAL: Ensure we have exactly 32 global tokens
        EXPECTED_GLOBAL_TOKENS = 32
        if len(global_ids) != EXPECTED_GLOBAL_TOKENS:
            print(f"⚠  WARNING: Got {len(global_ids)} global tokens, expected {EXPECTED_GLOBAL_TOKENS}")
            if len(global_ids) > EXPECTED_GLOBAL_TOKENS:
                # Truncate to first 32
                global_ids = global_ids[:EXPECTED_GLOBAL_TOKENS]
                print(f"   └─ Truncated to {EXPECTED_GLOBAL_TOKENS} tokens")
            else:
                # Pad with zeros if we somehow have fewer
                global_ids = global_ids + [0] * (EXPECTED_GLOBAL_TOKENS - len(global_ids))
                print(f"   └─ Padded to {EXPECTED_GLOBAL_TOKENS} tokens")

        # Decode with EXACTLY 32 global tokens + variable semantic tokens
        # The decoder handles broadcasting internally!
        audio = self.decode(semantic_ids, global_ids)

        if audio is None:
            return None

        # If using sliding window, return only the middle portion
        if use_sliding_window and len(audio) > 4096:
            # Keep middle samples for overlap-add streaming
            # BiCodec produces ~320 samples per semantic token (16kHz / 50 TPS)
            total_samples = len(audio)
            keep_samples = min(4096, total_samples // 2)
            start = (total_samples - keep_samples) // 2
            end = start + keep_samples
            audio = audio[start:end]

        # Convert to bytes
        audio = np.clip(audio, -1.0, 1.0)
        audio_int16 = (audio * 32767).astype(np.int16)

        return audio_int16.tobytes()

    async def decode_single_async(
        self,
        semantic_ids: List[int],
        global_ids: List[int],
        trim_warmup: bool = False,
        use_sliding_window: bool = False,
    ) -> Optional[bytes]:
        """
        Async wrapper for streaming decode - runs in executor to avoid blocking event loop.

        OPTIMIZATION: Previously ran synchronously, blocking all concurrent coroutines
        during GPU decode (~20-50ms). Now uses run_in_executor to unblock the event loop
        so other streams can progress while this decode runs on GPU.

        Args:
            semantic_ids: List of semantic token IDs
            global_ids: List of global token IDs
            trim_warmup: Legacy parameter (not used)
            use_sliding_window: Use sliding window mode

        Returns:
            Audio bytes (int16 PCM) or None
        """
        import asyncio
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            None,  # Default ThreadPoolExecutor
            self.decode_streaming,
            semantic_ids,
            global_ids,
            use_sliding_window,
            trim_warmup,
        )


Output:
Chunk ID: 7a3195
Wall time: 0.0515 seconds
Process exited with code 0
Original token count: 2422
Output:
"""
BiCodec Batch Decoder - Batches concurrent decode requests into single GPU forward passes.

OPTIMIZATION: Instead of N sequential ~30ms decodes (one per concurrent stream),
batches them into 1 GPU call of ~35-40ms. For 10 concurrent streams, this saves
~260ms of serialized GPU time per decode round.

Architecture:
    Each streaming coroutine submits a decode request to the batch queue.
    A background worker collects requests within a timeout window,
    pads semantic sequences to equal length, runs a single batched decode,
    then dispatches results back to individual futures.

Usage:
    # At startup:
    batch_decoder = BiCodecBatchDecoder(bicodec_decoder, max_batch=16, timeout_ms=10)
    batch_decoder.start()

    # In streaming hot loop:
    audio_bytes = await batch_decoder.decode_async(semantic_ids, global_ids)

    # At shutdown:
    batch_decoder.stop()
"""

import asyncio
import logging
import time
import threading
from typing import List, Optional, Tuple
from dataclasses import dataclass

import numpy as np
import torch

logger = logging.getLogger(__name__)

@dataclass
class DecodeRequest:
    """A single decode request from a streaming coroutine."""
    semantic_ids: List[int]
    global_ids: List[int]
    future: asyncio.Future
    loop: asyncio.AbstractEventLoop
    submitted_at: float

class BiCodecBatchDecoder:
    """
    Batches concurrent BiCodec decode requests into single GPU forward passes.

    Thread Safety:
        - submit() is called from async coroutines (event loop thread)
        - _worker_loop() runs in a background thread
        - Communication via thread-safe asyncio.Queue
        - Results dispatched via loop.call_soon_threadsafe()
    """

    def __init__(
        self,
        bicodec_decoder,
        max_batch: int = 16,
        timeout_ms: float = 10.0,
    ):
        """
        Args:
            bicodec_decoder: BiCodecDecoder instance (used for single-item decode fallback)
            max_batch: Max requests to batch together
            timeout_ms: Max time to wait for batch to fill before executing
        """
        self.decoder = bicodec_decoder
        self.max_batch = max_batch
        self.timeout_ms = timeout_ms

        self._queue: asyncio.Queue = asyncio.Queue()
        self._running = False
        self._worker_thread: Optional[threading.Thread] = None

        # Stats
        self.total_batches = 0
        self.total_requests = 0
        self.total_batch_time_ms = 0.0

    def start(self):
        """Start the background batch worker."""
        self._running = True
        self._worker_thread = threading.Thread(target=self._worker_loop, daemon=True)
        self._worker_thread.start()
        logger.info(f"BiCodecBatchDecoder started (max_batch={self.max_batch}, timeout={self.timeout_ms}ms)")

    def stop(self):
        """Stop the background worker."""
        self._running = False
        if self._worker_thread:
            self._worker_thread.join(timeout=5)

    async def decode_async(
        self,
        semantic_ids: List[int],
        global_ids: List[int],
    ) -> Optional[bytes]:
        """
        Submit a decode request and wait for the batched result.

        This is the main API for streaming coroutines.
        Returns audio_bytes (int16 PCM) or None if decode failed.
        """
        loop = asyncio.get_event_loop()
        future = loop.create_future()

        request = DecodeRequest(
            semantic_ids=semantic_ids,
            global_ids=global_ids,
            future=future,
            loop=loop,
            submitted_at=time.time(),
        )

        await self._queue.put(request)
        return await future

    def _worker_loop(self):
        """Background thread that collects and executes batched decodes."""
        while self._running:
            batch: List[DecodeRequest] = []

            # Collect requests up to max_batch or timeout
            try:
                # Block for the first request
                first = asyncio.run_coroutine_threadsafe(
                    self._queue.get(),
                    asyncio.get_event_loop()
                ).result(timeout=1.0)
                batch.append(first)
            except Exception:
                continue

            # Try to fill the batch within timeout
            deadline = time.time() + self.timeout_ms / 1000
            while len(batch) < self.max_batch and time.time() < deadline:
                try:
                    req = asyncio.run_coroutine_threadsafe(
                        asyncio.wait_for(self._queue.get(), timeout=0.001),
                        asyncio.get_event_loop()
                    ).result(timeout=self.timeout_ms / 1000)
                    batch.append(req)
                except Exception:
                    break

            if not batch:
                continue

            # Execute the batch
            t_start = time.time()
            try:
                if len(batch) == 1:
                    # Single request: no padding needed
                    result = self.decoder.decode_streaming(
                        semantic_ids=batch[0].semantic_ids,
                        global_ids=batch[0].global_ids,
                        use_sliding_window=False,
                        trim_warmup=False,
                    )
                    batch[0].loop.call_soon_threadsafe(batch[0].future.set_result, result)
                else:
                    # Batched decode: pad sequences and run together
                    results = self._batched_decode(batch)
                    for req, result in zip(batch, results):
                        req.loop.call_soon_threadsafe(req.future.set_result, result)
            except Exception as e:
                for req in batch:
                    if not req.future.done():
                        req.loop.call_soon_threadsafe(req.future.set_exception, e)

            elapsed = (time.time() - t_start) * 1000
            self.total_batches += 1
            self.total_requests += len(batch)
            self.total_batch_time_ms += elapsed

    def _batched_decode(self, batch: List[DecodeRequest]) -> List[Optional[bytes]]:
        """
        Execute a batched BiCodec decode.

        Pads semantic sequences to equal length, runs a single forward pass,
        then splits results back to individual outputs.
        """
        # All global_ids should be exactly 32 tokens
        # Semantic sequences may vary in length - pad to max
        max_semantic_len = max(len(req.semantic_ids) for req in batch)
        batch_size = len(batch)

        # Build batched tensors
        semantic_batch = torch.zeros(batch_size, max_semantic_len, dtype=torch.long, device=self.decoder.device)
        global_batch = torch.zeros(batch_size, 32, dtype=torch.long, device=self.decoder.device)
        seq_lengths = []

        for i, req in enumerate(batch):
            sem_len = len(req.semantic_ids)
            semantic_batch[i, :sem_len] = torch.tensor(req.semantic_ids, dtype=torch.long)
            global_batch[i, :len(req.global_ids)] = torch.tensor(req.global_ids[:32], dtype=torch.long)
            seq_lengths.append(sem_len)

        # Run batched decode through BiCodec
        try:
            with torch.inference_mode():
                wav_batch = self.decoder.audio_tokenizer.detokenize(
                    global_batch,
                    semantic_batch,
                )
        except Exception as e:
            logger.error(f"Batched decode failed: {e}")
            # Fallback to sequential decode
            results = []
            for req in batch:
                result = self.decoder.decode_streaming(
                    semantic_ids=req.semantic_ids,
                    global_ids=req.global_ids,
                    use_sliding_window=False,
                    trim_warmup=False,
                )
                results.append(result)
            return results

        # Split results: trim each output to its actual length
        # BiCodec produces ~320 samples per semantic token
        results = []
        for i, (req, sem_len) in enumerate(zip(batch, seq_lengths)):
            try:
                audio = wav_batch[i] if wav_batch.ndim > 1 else wav_batch
                if isinstance(audio, torch.Tensor):
                    audio = audio.cpu().numpy()

                # Trim to actual length (remove padding artifacts)
                expected_samples = sem_len * 320
                audio = audio[:expected_samples] if len(audio) > expected_samples else audio

                # Convert to int16 PCM bytes
                audio = np.clip(audio, -1.0, 1.0)
                audio_int16 = (audio * 32767).astype(np.int16)
                results.append(audio_int16.tobytes())
            except Exception as e:
                logger.error(f"Batch result extraction failed for item {i}: {e}")
                results.append(None)

        return results

    def get_stats(self):
        """Return batch decoder statistics."""
        avg_batch_size = self.total_requests / self.total_batches if self.total_batches > 0 else 0
        avg_batch_time = self.total_batch_time_ms / self.total_batches if self.total_batches > 0 else 0
        return {
            "total_batches": self.total_batches,
            "total_requests": self.total_requests,
            "avg_batch_size": round(avg_batch_size, 1),
            "avg_batch_time_ms": round(avg_batch_time, 1),
        }


Output:
Chunk ID: 992101
Wall time: 0.0513 seconds
Process exited with code 0
Original token count: 1578
Output:
"""
BiCodec Token Utilities - Incremental Parsing

Provides O(1) per-token parsing instead of O(n) decode + regex scanning.

The key optimization: instead of:
    tokenizer.decode(all_tokens)  # O(n)
    re.findall(entire_text)       # O(n)

We do:
    for new_token_id in new_tokens:  # Only new tokens
        parsed = token_cache.get(token_id)  # O(1) lookup

This eliminates the O(n²) CPU burn in the streaming hot loop.
"""

from typing import Dict, Optional, Tuple, List
import re

class BiCodecTokenParser:
    """
    Caching parser for BiCodec tokens.

    Converts vLLM token IDs to (type, value) tuples efficiently.

    BiCodec tokens in the vocabulary look like:
        - <|bicodec_semantic_123|> -> ("semantic", 123)
        - <|bicodec_global_7|>     -> ("global", 7)

    Usage:
        parser = BiCodecTokenParser(tokenizer)

        # In hot loop (O(1) per token):
        for token_id in new_token_ids:
            result = parser.parse(token_id)
            if result:
                token_type, value = result
                if token_type == "semantic":
                    semantic_buffer.append(value)
                elif token_type == "global":
                    global_buffer.append(value)
    """

    # Regex patterns for parsing token strings
    SEMANTIC_PATTERN = re.compile(r"^<\|bicodec_semantic_(\d+)\|>$")
    GLOBAL_PATTERN = re.compile(r"^<\|bicodec_global_(\d+)\|>$")

    def __init__(self, tokenizer):
        """
        Initialize parser with tokenizer.

        Args:
            tokenizer: HuggingFace tokenizer with BiCodec special tokens
        """
        self.tokenizer = tokenizer
        # Cache: token_id -> Optional[Tuple[str, int]]
        # None means "not a BiCodec token" (cached negative)
        self._cache: Dict[int, Optional[Tuple[str, int]]] = {}

        # Pre-warm cache with known BiCodec token ranges if possible
        # This is optional but can speed up first-pass parsing
        self._prewarm_cache()

    def _prewarm_cache(self) -> None:
        """
        Pre-populate cache with BiCodec tokens from vocabulary.

        This avoids lazy cache misses during streaming.
        """
        try:
            vocab = self.tokenizer.get_vocab()
            for token_str, token_id in vocab.items():
                # Check if it's a BiCodec token
                if token_str.startswith("<|bicodec_"):
                    parsed = self._parse_token_string(token_str)
                    if parsed:
                        self._cache[token_id] = parsed
        except Exception:
            # If vocab access fails, we'll rely on lazy caching
            pass

    def _parse_token_string(self, token_str: str) -> Optional[Tuple[str, int]]:
        """
        Parse a token string into (type, value) tuple.

        Args:
            token_str: Token string like "<|bicodec_semantic_123|>"

        Returns:
            ("semantic", 123) or ("global", 7) or None if not BiCodec token
        """
        # Try semantic pattern
        match = self.SEMANTIC_PATTERN.match(token_str)
        if match:
            return ("semantic", int(match.group(1)))

        # Try global pattern
        match = self.GLOBAL_PATTERN.match(token_str)
        if match:
            return ("global", int(match.group(1)))

        return None

    def parse(self, token_id: int) -> Optional[Tuple[str, int]]:
        """
        Parse a single token ID into (type, value) tuple.

        O(1) for cached tokens, O(1) amortized for uncached.

        Args:
            token_id: vLLM token ID

        Returns:
            ("semantic", value) or ("global", value) or None
        """
        # Check cache first
        if token_id in self._cache:
            return self._cache[token_id]

        # Cache miss - decode and parse
        try:
            token_str = self.tokenizer.decode([token_id], skip_special_tokens=False)
            parsed = self._parse_token_string(token_str.strip())
            self._cache[token_id] = parsed
            return parsed
        except Exception:
            # Mark as non-BiCodec to avoid repeated failures
            self._cache[token_id] = None
            return None

    def parse_incremental(
        self,
        new_token_ids: List[int],
        semantic_buffer: List[int],
        global_buffer: List[int],
    ) -> Tuple[List[int], List[int]]:
        """
        Parse new tokens and append to existing buffers.

        This is the main method for incremental streaming:
        - Only processes NEW tokens (not entire history)
        - Modifies buffers in-place for efficiency
        - Returns the updated buffers (same objects)

        Args:
            new_token_ids: Only the new token IDs since last call
            semantic_buffer: Existing semantic token values (modified in-place)
            global_buffer: Existing global token values (modified in-place)

        Returns:
            Tuple of (semantic_buffer, global_buffer) - same objects, modified
        """
        for token_id in new_token_ids:
            parsed = self.parse(token_id)
            if parsed:
                token_type, value = parsed
                if token_type == "semantic":
                    semantic_buffer.append(value)
                elif token_type == "global":
                    global_buffer.append(value)

        return semantic_buffer, global_buffer

    def get_cache_stats(self) -> Dict[str, int]:
        """
        Return cache statistics for debugging.

        Returns:
            Dict with cache size and breakdown
        """
        total = len(self._cache)
        bicodec_count = sum(1 for v in self._cache.values() if v is not None)
        return {
            "total_cached": total,
            "bicodec_tokens": bicodec_count,
            "non_bicodec_tokens": total - bicodec_count,
        }

def create_bicodec_parser(tokenizer) -> BiCodecTokenParser:
    """
    Factory function to create a BiCodec token parser.

    Args:
        tokenizer: HuggingFace tokenizer

    Returns:
        Configured BiCodecTokenParser instance
    """
    return BiCodecTokenParser(tokenizer)


Output:
Chunk ID: 49e05e
Wall time: 0.0513 seconds
Process exited with code 0
Original token count: 2491
Output:
"""
Long Text Processing Service

Handles text chunking and audio stitching for long text inputs.

Why This Approach:
- Text-based chunking: Preserves semantic meaning
- SNAC token sliding window: Separate layer for smooth audio
- Crossfade stitching: Seamless audio transitions between chunks
- Unlimited text: Can process any length text

Architecture:
1. Long text → IndicSentenceChunker → Text chunks
2. Each text chunk → TTS Model → Audio chunk
3. All audio chunks → Crossfade stitching → Final audio
"""

import asyncio
import numpy as np
import struct
from typing import List, Optional, AsyncGenerator
import logging

from veena3modal.processing.text_chunker import IndicSentenceChunker, crossfade_audio
from veena3modal.core.streaming_pipeline import Veena3SlidingWindowPipeline
from veena3modal.core.pipeline import SparkTTSPipeline

logger = logging.getLogger(__name__)

def create_wav_header(data_size: int, sample_rate: int = 24000, channels: int = 1, bits_per_sample: int = 16) -> bytes:
    """
    Create WAV file header for PCM data.

    Args:
        data_size: Size of PCM data in bytes
        sample_rate: Sample rate (Hz)
        channels: Number of channels
        bits_per_sample: Bits per sample

    Returns:
        WAV header bytes (44 bytes)
    """
    byte_rate = sample_rate * channels * bits_per_sample // 8
    block_align = channels * bits_per_sample // 8

    header = struct.pack('<4sI4s', b'RIFF', 36 + data_size, b'WAVE')
    header += struct.pack('<4sIHHIIHH', b'fmt ', 16, 1, channels, sample_rate, byte_rate, block_align, bits_per_sample)
    header += struct.pack('<4sI', b'data', data_size)

    return header

class LongTextProcessor:
    """
    Process long texts by chunking, generating audio, and stitching.

    Features:
    - Intelligent text chunking (respects sentence boundaries, Indic languages)
    - Parallel audio generation (if desired)
    - Crossfade stitching for seamless audio
    - Memory efficient (processes chunks sequentially by default)

    Why text chunking, not token chunking:
    - Text maintains semantic meaning
    - Natural sentence boundaries = better prosody
    - Simpler implementation
    - Model handles tokenization internally
    - SNAC token sliding window is separate (already implemented)

    Max Input Length Reference (post-normalization):
    "इसके निर्माण में मुख्य वास्तुकार उस्ताद अहमद लाहौरी के नेतृत्व में लगभग
    20,000 कारीगरों और शिल्पकारों ने दिन-रात मेहनत की थी। इमारत के सामने बना
    'चारबाग' शैली का उद्यान और पानी की नहरें इसकी सुंदरता में चार चाँद लगा देते हैं।"
    = ~232 characters after normalization (numbers expanded to words)
    """

    # Maximum safe input length for model (characters)
    # Based on reference Hindi sentence (~232 chars after normalization)
    # Model struggles with longer inputs, causes token repetition
    MAX_MODEL_INPUT_LENGTH = 230  # ~232 chars is the reference max

    # Threshold for when to enable chunking (characters)
    # Trigger chunking when text approaches max safe length
    # Using 95% of max to leave some buffer
    CHUNKING_THRESHOLD = 220  # ~95% of MAX_MODEL_INPUT_LENGTH

    # Chunk size for long texts (characters)
    # Use 75% of max safe length to ensure each chunk processes cleanly
    # This gives model headroom for complex multilingual content
    CHUNK_SIZE = 170  # ~75% of MAX_MODEL_INPUT_LENGTH

    # Crossfade duration (milliseconds)
    CROSSFADE_MS = 50  # 50ms smooth transition

    def __init__(
        self,
        pipeline: SparkTTSPipeline,
        streaming_pipeline: Optional[Veena3SlidingWindowPipeline] = None
    ):
        """
        Initialize long text processor.

        Args:
            pipeline: Non-streaming TTS pipeline
            streaming_pipeline: Optional streaming pipeline (for streaming mode)
        """
        self.pipeline = pipeline
        self.streaming_pipeline = streaming_pipeline
        self.chunker = IndicSentenceChunker(max_chunk_length=self.CHUNK_SIZE)

        logger.info(
            f"LongTextProcessor initialized",
            extra={
                "max_model_input_length": self.MAX_MODEL_INPUT_LENGTH,
                "chunking_threshold": self.CHUNKING_THRESHOLD,
                "chunk_size": self.CHUNK_SIZE,
                "crossfade_ms": self.CROSSFADE_MS,
            }
        )

    def should_chunk(self, text: str) -> bool:
        """Determine if text should be chunked."""
        return len(text) > self.CHUNKING_THRESHOLD

    def chunk_text(self, text: str) -> list:
        """
        Chunk text using the internal chunker.

        Delegates to self.chunker.chunk_text() for actual chunking logic.
        This method is used by _stream_chunked_text in tts_runtime.

        Args:
            text: Text to chunk

        Returns:
            List of text chunks
        """
        return self.chunker.chunk_text(text)

    async def generate_with_chunking(
        self,
        text: str,
        speaker: str,
        temperature: float = 0.4,
        top_k: int = 50,  # Added for Spark TTS
        top_p: float = 0.9,
        max_tokens: int = 4096,
        repetition_penalty: float = 1.05,  # Prevent token repetition
        seed: Optional[int] = None,
        sample_rate: int = 16000,  # BiCodec uses 16kHz, not 24kHz
    ) -> Optional[bytes]:
        """
        Generate audio for long text with chunking and stitching.

        Process:
        1. Chunk text at natural boundaries
        2. Generate audio for each chunk
        3. Stitch audio with crossfade

        Args:
            text: Long text to synthesize
            speaker: Speaker name (lipakshi, reet, etc.)
            temperature: Sampling temperature
            top_p: Nucleus sampling
            max_tokens: Max tokens per chunk
            seed: Random seed (same seed = same voice across chunks)
            sample_rate: Audio sample rate

        Returns:
            Stitched audio bytes (int16 PCM) or None if generation fails
        """
        # Detect language mix for logging
        lang_info = self.chunker.detect_language_mix(text)

        logger.info(
            f"Processing long text with chunking",
            extra={
                "text_length": len(text),
                "primary_language": lang_info['primary'],
                "lang_mix": lang_info,
                "speaker": speaker,
            }
        )

        # 1. Chunk text
        text_chunks = self.chunker.chunk_text(text)

        logger.info(
            f"Text chunked into {len(text_chunks)} chunks",
            extra={
                "num_chunks": len(text_chunks),
                "chunk_sizes": [len(chunk) for chunk in text_chunks],
            }
        )

        # 2. Generate audio for each chunk
        audio_chunks = []

        # Log summary at start
        logger.info(
            f"🎵 Generating {len(text_chunks)} chunks for {len(text)} chars",
            extra={
                "total_chunks": len(text_chunks),
                "total_chars": len(text),
                "speaker": speaker,
            }
        )

        for i, chunk in enumerate(text_chunks):
            # Calculate appropriate max_tokens for this chunk
            # BiCodec semantic tokens: ~50 tokens/second, ~5 words/second = ~10 tokens/word
            # Rough estimate: 5 chars/word, so ~2 tokens/char for audio generation
            # Add buffer for safety: multiply by 3 to handle variations
            chunk_max_tokens = min(
                int(len(chunk) * 3),  # ~2 tokens/char * 1.5 safety buffer = 3
                4096  # Cap at model's max context length
            )
            chunk_max_tokens = max(chunk_max_tokens, 100)  # Ensure minimum

            # Generate audio for this chunk (returns WAV format with header)
            audio_bytes = await self.pipeline.generate_speech_indic(
                speaker=speaker,
                text=chunk,
                temperature=temperature,
                top_k=top_k,  # Added for Spark TTS
                top_p=top_p,
                max_tokens=chunk_max_tokens,  # Use calculated limit, not global max_tokens
                repetition_penalty=repetition_penalty,  # Prevent token repetition
                seed=seed,  # Use same seed for consistent voice across chunks
            )

            if audio_bytes is None:
                logger.error(
                    f"❌ Failed to generate audio for chunk {i+1}/{len(text_chunks)}",
                    extra={
                        "chunk_index": i,
                        "chunk_length": len(chunk),
                        "chunk_text_preview": chunk[:150],
                        "total_chunks": len(text_chunks),
                    }
                )
                return None

            # Strip WAV header (44 bytes) to get raw PCM data for stitching
            pcm_data = audio_bytes[44:] if len(audio_bytes) > 44 else audio_bytes
            audio_chunks.append(pcm_data)

        # 3. Stitch audio with crossfade
        stitched_pcm = self.stitch_audio_chunks(
            audio_chunks,
            sample_rate=sample_rate,
            crossfade_ms=self.CROSSFADE_MS,
        )

        # 4. Add WAV header to stitched PCM data
        wav_header = create_wav_header(
            data_size=len(stitched_pcm),
            sample_rate=sample_rate,
            channels=1,


Output:
Chunk ID: 785f4c
Wall time: 0.0513 seconds
Process exited with code 0
Original token count: 1086
Output:
            channels=1,
            bits_per_sample=16
        )
        final_audio = wav_header + stitched_pcm

        # Calculate total audio metrics
        total_audio_duration = len(stitched_pcm) / (sample_rate * 2)

        logger.info(
            f"✅ Generated {len(text_chunks)} chunks → {total_audio_duration:.1f}s audio",
            extra={
                "total_chunks": len(text_chunks),
                "audio_duration": total_audio_duration,
                "audio_size_mb": len(final_audio) / (1024 * 1024),
            }
        )

        return final_audio

    def stitch_audio_chunks(
        self,
        audio_chunks: List[bytes],
        sample_rate: int = 16000,  # BiCodec uses 16kHz, not 24kHz
        crossfade_ms: int = 50,
    ) -> bytes:
        """
        Stitch multiple audio chunks with crossfade transitions.

        Why crossfade:
        - Prevents pops/clicks at chunk boundaries
        - Smooth transitions between chunks
        - Professional audio quality

        Args:
            audio_chunks: List of audio bytes (int16 PCM)
            sample_rate: Sample rate (Hz)
            crossfade_ms: Crossfade duration (milliseconds)

        Returns:
            Stitched audio bytes
        """
        if not audio_chunks:
            return b''

        if len(audio_chunks) == 1:
            return audio_chunks[0]

        # Calculate crossfade samples
        crossfade_samples = int((crossfade_ms / 1000.0) * sample_rate)

        # Start with first chunk
        result = audio_chunks[0]

        # Stitch remaining chunks with crossfade
        for i in range(1, len(audio_chunks)):
            result = crossfade_audio(
                result,
                audio_chunks[i],
                crossfade_samples=crossfade_samples,
                sample_rate=sample_rate,
            )

            logger.debug(
                f"Crossfaded chunk {i+1}",
                extra={
                    "chunk_index": i,
                    "current_length_bytes": len(result),
                }
            )

        return result

    async def generate_with_chunking_streaming(
        self,
        text: str,
        speaker: str,
        temperature: float = 0.4,
        top_p: float = 0.9,
        max_tokens: int = 4096,
        seed: Optional[int] = None,
    ) -> AsyncGenerator[bytes, None]:
        """
        Generate audio for long text with streaming (experimental).

        Note: Streaming long texts is more complex:
        - Need to stream each chunk
        - Need to buffer for crossfade
        - Adds latency for stitching

        For now, recommend using non-streaming for long texts.

        Args:
            text: Long text to synthesize
            speaker: Speaker name
            temperature: Sampling temperature
            top_p: Nucleus sampling
            max_tokens: Max tokens per chunk
            repetition_penalty: Repetition penalty
            seed: Random seed

        Yields:
            Audio bytes (in chunks)
        """
        if self.streaming_pipeline is None:
            raise ValueError("Streaming pipeline not available")

        # Chunk text
        text_chunks = self.chunker.chunk_text(text)

        logger.info(
            f"Streaming long text with {len(text_chunks)} chunks",
            extra={"num_chunks": len(text_chunks)}
        )

        # For streaming, we process each chunk sequentially
        # Future optimization: overlap processing + streaming
        for i, chunk in enumerate(text_chunks):
            logger.info(f"Streaming chunk {i+1}/{len(text_chunks)}")

            # Stream this chunk
            async for audio_bytes in self.streaming_pipeline.generate_speech_stream_indic(
                speaker=speaker,
                text=chunk,
                temperature=temperature,
                top_p=top_p,
                max_tokens=max_tokens,
                seed=seed,
            ):
                yield audio_bytes

            # TODO: Add crossfade between chunks in streaming mode
            # This requires buffering the end of previous chunk
            # and beginning of next chunk


**Planning detailed GPU timing inspection**


rg -n "request_output\.metrics|\.metrics|llm_batch|gpu_ms|prompt_tokens|output_len|prefill|decode" scripts/profile_pipeline.py veena3modal -S


sed -n '260,760p' scripts/profile_pipeline.py


sed -n '760,1200p' scripts/profile_pipeline.py