Output:
Chunk ID: f0b4df
Wall time: 0.0531 seconds
Process exited with code 0
Original token count: 2226
Output:
"""
Modal entrypoint for Veena3 TTS service.

Deploys a GPU-backed ASGI app with:
- True streaming TTS
- Autoscaling (0 to N containers)
- Per-container concurrency (vLLM batching)
- Memory snapshots for fast cold starts

Usage:
    # Deploy to Modal
    modal deploy veena3modal/app.py

    # Serve locally (for testing)
    modal serve veena3modal/app.py

    # Run a single function
    modal run veena3modal/app.py::tts_api
"""

from __future__ import annotations

import os
import modal

# Get the repo root for adding local sources
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# === Image Definition ===
# Full dependencies for TTS inference

image = (
    modal.Image.debian_slim(python_version="3.10")
    .apt_install(
        "ffmpeg",  # Audio encoding (opus, mp3, flac, mulaw)
        "git",     # For pip git installs
    )
    .pip_install(
        # Core ML/AI
        "torch>=2.0.0",
        "torchaudio>=2.0.0",
        "transformers>=4.35.0",
        "vllm>=0.4.0",

        # FastAPI
        "fastapi[standard]>=0.100.0",
        "uvicorn[standard]>=0.23.0",

        # Audio processing
        "numpy>=1.24.0",
        "scipy>=1.10.0",
        "librosa>=0.10.0",
        "soundfile>=0.12.0",

        # Utils
        "pydantic>=2.0.0",
        "prometheus_client>=0.17.0",
        "supabase>=2.0.0",

        # vLLM dependencies
        "einops>=0.6.0",
        "einx>=0.2.0",

        # SparkTTS dependencies (from external/sparktts/requirements.txt)
        "omegaconf>=2.3.0",
        "safetensors>=0.5.0",
        "soxr>=0.5.0",

        # AP-BWE (super-resolution) dependencies
        "matplotlib>=3.8.0",
        "natsort>=8.0.0",
        "joblib>=1.0.0",
    )
    .run_commands(
        # Ensure CUDA is properly configured
        "pip install --no-cache-dir nvidia-ml-py3 || true"
    )
    # Environment settings
    .env({
        "AUTH_BYPASS_MODE": "true",  # Disable auth for testing
        # sparktts is at /root/external/sparktts, AP-BWE at /root/external/AP-BWE
        "PYTHONPATH": "/root/external/sparktts:/root/external/AP-BWE:/root",
        "_IMAGE_BUILD_VERSION": "2025-12-25-sr-chunking-support",  # Add SR support to generate_speech_chunked
    })
    # Add sparktts package - create proper structure by copying files
    # The original has symlinks that don't work with Modal mounts
    .run_commands(
        # Create sparktts package directory
        "mkdir -p /root/external/sparktts/sparktts/models",
        "mkdir -p /root/external/sparktts/sparktts/utils",
    )
    .add_local_dir(
        local_path=os.path.join(REPO_ROOT, "external", "sparktts", "sparktts"),
        remote_path="/root/external/sparktts/sparktts",
        copy=True,
    )
    .add_local_dir(
        local_path=os.path.join(REPO_ROOT, "external", "sparktts", "models"),
        remote_path="/root/external/sparktts/sparktts/models",  # Put models inside sparktts package
        copy=True,
    )
    .add_local_dir(
        local_path=os.path.join(REPO_ROOT, "external", "sparktts", "utils"),
        remote_path="/root/external/sparktts/sparktts/utils",  # Put utils inside sparktts package
        copy=True,
    )
    # Add AP-BWE for super-resolution (16kHz -> 48kHz)
    .add_local_dir(
        local_path=os.path.join(REPO_ROOT, "external", "AP-BWE"),
        remote_path="/root/external/AP-BWE",
        copy=True,
    )
    # Non-copy local mounts (MUST BE LAST - no build steps after this!)
    .add_local_python_source("veena3modal")
)

# === Volume for Model Weights ===
# Models should be uploaded to this volume before deployment

try:
    model_volume = modal.Volume.from_name("veena3-models", create_if_missing=True)
except Exception:
    model_volume = None  # Will fail at runtime if volume not available

# === Secrets (Optional) ===
# Secrets are optional - the app will run in bypass mode without them
# Create secrets in Modal dashboard: https://modal.com/secrets
# Or use modal.Secret.from_dotenv() for local testing

# === App Definition ===

app = modal.App(
    name="veena3-tts",
    image=image,
)

@app.cls(
    # === GPU Selection ===
    # L40S: Best cost/performance ratio for TTS inference
    # - 48GB VRAM, $1.50/hr (vs A100-80GB at $4.50/hr)
    # - Handles ~12 req/s with <2s p95 latency per container
    # - Memory-bound workload, so faster GPUs have diminishing returns
    gpu="L40S",

    # === CPU Allocation (OPTIMIZATION Dec 2025) ===
    # Explicitly reserve 4 vCPUs for the Python driver process
    # Prevents CPU throttling during vLLM scheduler + streaming loop
    # Modal default is only 0.125 cores which can bottleneck streaming
    cpu=4.0,

    volumes={"/models": model_volume} if model_volume else {},
    secrets=[modal.Secret.from_name("veena3-secrets")],  # Supabase credentials

    # === Autoscaling Configuration ===
    # See .cursor/scaling.md for detailed analysis
    #
    # LOAD TEST RESULTS (Dec 2025):
    # - Cold boot: ~49s (memory snapshot captures imports, not GPU model)
    # - MAX SUSTAINABLE: 280 req/s @ 250 concurrent users (P95=1417ms, 100% success)
    # - Breaking point: 300+ users (429 errors start)
    # - Per container: ~14 req/s (280/20 containers)
    #
    min_containers=1,        # Keep 1 warm 24/7 for instant response
    max_containers=50,       # ~280 req/s capacity at P95 < 3s
    buffer_containers=1,     # REQUIRED: 49s cold boot too slow for burst
    scaledown_window=300,    # 5 min idle before scaledown (user changed from 120)

    # === Timeouts ===
    timeout=120,             # 2 min max per request (reduced from 10 min)
    startup_timeout=600,     # 10 min for model loading (reduced from 20 min)

    # === Memory Optimization ===
    # Memory snapshot for faster cold boots
    # Note: GPU snapshot is experimental (alpha) - disabled for stability
    enable_memory_snapshot=True,
)
# === Per-Container Concurrency ===
# vLLM continuous batching handles concurrent requests efficiently
# Increased from 12 to 20 to maximize GPU utilization while keeping P95 < 3s
@modal.concurrent(max_inputs=20, target_inputs=15)
class TTSService:
    """
    TTS service class with model lifecycle management.

    Uses @modal.enter for model loading (once per container).
    Memory snapshot captures imports and non-GPU state for faster restarts.
    """

    @modal.enter()
    def load_model(self):
        """
        Load TTS model when container starts.

        Called once per container lifecycle.
        Memory snapshot captures import state but not GPU memory.
        """
        import os
        import logging

        logging.basicConfig(level=logging.INFO)
        logger = logging.getLogger(__name__)

        logger.info("🚀 Starting TTS model loading...")

        # Set model paths from Volume mount or env vars
        model_path = os.environ.get(
            "MODEL_PATH",
            "/models/spark_tts_4speaker"
        )
        sr_path = os.environ.get(
            "AP_BWE_CHECKPOINT_DIR",
            "/models/ap_bwe/16kto48k"
        )

        # Check if model exists
        if not os.path.exists(model_path):
            logger.warning(f"Model path not found: {model_path}")
            logger.warning("TTS will return 503 until model is available")
            return

        try:
            from veena3modal.services.tts_runtime import initialize_runtime

            # Initialize with GPU memory optimization
            self.runtime = initialize_runtime(
                model_path=model_path,
                sr_checkpoint_dir=sr_path if os.path.exists(sr_path) else None,
                device="cuda",
                gpu_memory_utilization=0.85,
                enable_sr=os.path.exists(sr_path),
            )

            logger.info(f"✅ TTS model loaded: {self.runtime.model_version}")

        except Exception as e:
            logger.error(f"❌ Model loading failed: {e}")
            import traceback
            traceback.print_exc()

    @modal.asgi_app()
    def serve(self):
        """
        Return FastAPI ASGI app for serving requests.
        """
        from veena3modal.api.fastapi_app import create_app
        return create_app()

# === Standalone Function (Alternative Deployment) ===
# Use this if you don't need @modal.enter lifecycle
# Note: Prefer TTSService class for production (better model lifecycle management)

@app.function(
    gpu="L40S",  # Best cost/performance for TTS
    cpu=4.0,     # Explicit CPU allocation (OPTIMIZATION Dec 2025)
    volumes={"/models": model_volume} if model_volume else {},
    secrets=[modal.Secret.from_name("veena3-secrets")],
    min_containers=0,        # Scale to zero (no min for standalone)
    max_containers=5,        # Lower cap for standalone


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


rg -n "csv|json\(|json.dump|concurrency|stress|percentile|argparse" /home/ubuntu/veenaModal/scripts -g'*.py'


Output:
Chunk ID: dceb38
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 2451
Output:
/home/ubuntu/veenaModal/scripts/profile_pipeline.py:15:For concurrency: measures contention points under parallel load.
/home/ubuntu/veenaModal/scripts/profile_pipeline.py:22:    python scripts/profile_pipeline.py --concurrency 1,5,10,20
/home/ubuntu/veenaModal/scripts/profile_pipeline.py:33:import argparse
/home/ubuntu/veenaModal/scripts/profile_pipeline.py:567:def print_concurrency_summary(profiles: List[RequestProfile], concurrency: int, stream: bool):
/home/ubuntu/veenaModal/scripts/profile_pipeline.py:576:    print(f"  CONCURRENCY={concurrency} ({mode}), {len(profiles)} requests")
/home/ubuntu/veenaModal/scripts/profile_pipeline.py:673:    concurrency_levels = [int(x) for x in args.concurrency.split(",")]
/home/ubuntu/veenaModal/scripts/profile_pipeline.py:676:    if not args.concurrency_only:
/home/ubuntu/veenaModal/scripts/profile_pipeline.py:694:    for conc in concurrency_levels:
/home/ubuntu/veenaModal/scripts/profile_pipeline.py:704:            print_concurrency_summary(profiles, conc, stream_mode)
/home/ubuntu/veenaModal/scripts/profile_pipeline.py:711:    parser = argparse.ArgumentParser(description="Veena3 TTS Pipeline Profiler")
/home/ubuntu/veenaModal/scripts/profile_pipeline.py:714:    parser.add_argument("--concurrency", default="1,5,10,20", help="Concurrency levels (comma-separated)")
/home/ubuntu/veenaModal/scripts/profile_pipeline.py:717:    parser.add_argument("--concurrency-only", action="store_true", help="Skip single-request profiles")
/home/ubuntu/veenaModal/scripts/stress_test_production.py:12:Runs the same concurrency ramp as the baseline test for direct comparison.
/home/ubuntu/veenaModal/scripts/stress_test_production.py:15:import argparse
/home/ubuntu/veenaModal/scripts/stress_test_production.py:88:# === GPU Monitor (reused from stress_test_local.py) ===
/home/ubuntu/veenaModal/scripts/stress_test_production.py:121:                     "--format=csv,noheader,nounits"],
/home/ubuntu/veenaModal/scripts/stress_test_production.py:220:def percentile(values, pct):
/home/ubuntu/veenaModal/scripts/stress_test_production.py:227:async def run_level(base_url, num_requests, concurrency, stream, gpu_monitor=None):
/home/ubuntu/veenaModal/scripts/stress_test_production.py:228:    sem = asyncio.Semaphore(concurrency)
/home/ubuntu/veenaModal/scripts/stress_test_production.py:243:    limits = httpx.Limits(max_connections=concurrency + 10, max_keepalive_connections=concurrency)
/home/ubuntu/veenaModal/scripts/stress_test_production.py:265:    label = f"{concurrency}c-production-{mode}"
/home/ubuntu/veenaModal/scripts/stress_test_production.py:269:    print(f"  {num_requests} requests, {concurrency} concurrent, {mode}")
/home/ubuntu/veenaModal/scripts/stress_test_production.py:282:              f"p50={percentile(lats,50):.0f}  p95={percentile(lats,95):.0f}  "
/home/ubuntu/veenaModal/scripts/stress_test_production.py:283:              f"p99={percentile(lats,99):.0f}  max={max(lats):.0f}")
/home/ubuntu/veenaModal/scripts/stress_test_production.py:288:              f"p50={percentile(ttfbs,50):.0f}  p95={percentile(ttfbs,95):.0f}  max={max(ttfbs):.0f}")
/home/ubuntu/veenaModal/scripts/stress_test_production.py:306:        "level": label, "concurrency": concurrency, "requests": num_requests,
/home/ubuntu/veenaModal/scripts/stress_test_production.py:310:        "latency_ms": {"avg": statistics.mean(lats), "p50": percentile(lats, 50),
/home/ubuntu/veenaModal/scripts/stress_test_production.py:311:                        "p95": percentile(lats, 95), "p99": percentile(lats, 99),
/home/ubuntu/veenaModal/scripts/stress_test_production.py:313:        "ttfb_ms": {"avg": statistics.mean(ttfbs), "p50": percentile(ttfbs, 50),
/home/ubuntu/veenaModal/scripts/stress_test_production.py:314:                     "p95": percentile(ttfbs, 95), "max": max(ttfbs)} if ttfbs else {},
/home/ubuntu/veenaModal/scripts/stress_test_production.py:325:            h = r.json()
/home/ubuntu/veenaModal/scripts/stress_test_production.py:398:        json.dump(all_results, f, indent=2, default=str)
/home/ubuntu/veenaModal/scripts/stress_test_production.py:403:    parser = argparse.ArgumentParser(description="Veena3 Production Stress Test")
/home/ubuntu/veenaModal/scripts/stress_test_production.py:406:    parser.add_argument("--output", default="stress_test_production.json")
/home/ubuntu/veenaModal/scripts/final_asr_verification.py:93:    result = response.json()
/home/ubuntu/veenaModal/scripts/comprehensive_validation.py:165:            data = response.json()
/home/ubuntu/veenaModal/scripts/comprehensive_validation.py:393:def test_concurrency():
/home/ubuntu/veenaModal/scripts/comprehensive_validation.py:398:    concurrency_levels = [1, 3, 5]
/home/ubuntu/veenaModal/scripts/comprehensive_validation.py:402:    for n_concurrent in concurrency_levels:
/home/ubuntu/veenaModal/scripts/comprehensive_validation.py:501:        ("Concurrency", test_concurrency),
/home/ubuntu/veenaModal/scripts/elevenlabs_asr_analysis.py:58:        result = response.json()
/home/ubuntu/veenaModal/scripts/validation/final_validation.py:22:    data = response.json()
/home/ubuntu/veenaModal/scripts/validate_modal_endpoint.py:50:    print(f"Response: {json.dumps(response.json(), indent=2)}")
/home/ubuntu/veenaModal/scripts/validate_modal_endpoint.py:53:        data = response.json()
/home/ubuntu/veenaModal/scripts/validate_modal_endpoint.py:236:            ws.send(json.dumps({
/home/ubuntu/veenaModal/scripts/elevenlabs_verify_fix.py:105:    return response.json()
/home/ubuntu/veenaModal/scripts/stress_test_local.py:5:Progressive load ramp with GPU monitoring, latency percentiles,
/home/ubuntu/veenaModal/scripts/stress_test_local.py:11:- Higher concurrency ceiling (single-machine, no rate limits)
/home/ubuntu/veenaModal/scripts/stress_test_local.py:16:    python scripts/stress_test_local.py
/home/ubuntu/veenaModal/scripts/stress_test_local.py:19:    python scripts/stress_test_local.py --quick
/home/ubuntu/veenaModal/scripts/stress_test_local.py:22:    python scripts/stress_test_local.py --levels 1,5,10,20
/home/ubuntu/veenaModal/scripts/stress_test_local.py:25:    python scripts/stress_test_local.py --stream-only
/home/ubuntu/veenaModal/scripts/stress_test_local.py:28:    python scripts/stress_test_local.py --no-stream
/home/ubuntu/veenaModal/scripts/stress_test_local.py:31:    python scripts/stress_test_local.py --url http://localhost:8080
/home/ubuntu/veenaModal/scripts/stress_test_local.py:34:import argparse
/home/ubuntu/veenaModal/scripts/stress_test_local.py:53:# Test texts at different lengths to stress different parts of the pipeline
/home/ubuntu/veenaModal/scripts/stress_test_local.py:111:                        "--format=csv,noheader,nounits",
/home/ubuntu/veenaModal/scripts/stress_test_local.py:180:    concurrency: int
/home/ubuntu/veenaModal/scripts/stress_test_local.py:204:    def _percentile(self, values: List[float], pct: float) -> float:
/home/ubuntu/veenaModal/scripts/stress_test_local.py:219:            "p50": self._percentile(lats, 50),
/home/ubuntu/veenaModal/scripts/stress_test_local.py:220:            "p95": self._percentile(lats, 95),
/home/ubuntu/veenaModal/scripts/stress_test_local.py:221:            "p99": self._percentile(lats, 99),
/home/ubuntu/veenaModal/scripts/stress_test_local.py:233:            "p50": self._percentile(ttfbs, 50),
/home/ubuntu/veenaModal/scripts/stress_test_local.py:234:            "p95": self._percentile(ttfbs, 95),
/home/ubuntu/veenaModal/scripts/stress_test_local.py:401:    concurrency: int,
/home/ubuntu/veenaModal/scripts/stress_test_local.py:408:    level_name = f"{concurrency}c-{text_category}-{mode}"
/home/ubuntu/veenaModal/scripts/stress_test_local.py:412:        concurrency=concurrency,
/home/ubuntu/veenaModal/scripts/stress_test_local.py:418:    semaphore = asyncio.Semaphore(concurrency)
/home/ubuntu/veenaModal/scripts/stress_test_local.py:440:        max_connections=concurrency + 10,
/home/ubuntu/veenaModal/scripts/stress_test_local.py:441:        max_keepalive_connections=concurrency,
/home/ubuntu/veenaModal/scripts/stress_test_local.py:473:    print(f"  {level.total_requests} requests, {level.concurrency} concurrent, {level.mode}")
/home/ubuntu/veenaModal/scripts/stress_test_local.py:545:async def run_stress_test(args):
/home/ubuntu/veenaModal/scripts/stress_test_local.py:546:    """Run the full stress test suite."""
/home/ubuntu/veenaModal/scripts/stress_test_local.py:554:            health = resp.json()
/home/ubuntu/veenaModal/scripts/stress_test_local.py:564:    # Parse concurrency levels
/home/ubuntu/veenaModal/scripts/stress_test_local.py:582:    warmup = await run_level(base_url, num_requests=2, concurrency=1, text_category="short", stream=False)
/home/ubuntu/veenaModal/scripts/stress_test_local.py:598:            num_requests = max(conc * 3, 10)  # At least 3x concurrency or 10
/home/ubuntu/veenaModal/scripts/stress_test_local.py:608:                concurrency=conc,
/home/ubuntu/veenaModal/scripts/stress_test_local.py:640:            "concurrency": level.concurrency,
/home/ubuntu/veenaModal/scripts/stress_test_local.py:656:        json.dump(data, f, indent=2, default=str)
/home/ubuntu/veenaModal/scripts/stress_test_local.py:660:    parser = argparse.ArgumentParser(
/home/ubuntu/veenaModal/scripts/stress_test_local.py:662:        formatter_class=argparse.RawDescriptionHelpFormatter,
/home/ubuntu/veenaModal/scripts/stress_test_local.py:665:    parser.add_argument("--levels", default="1,5,10,20,50", help="Comma-separated concurrency levels (default: 1,5,10,20,50)")
/home/ubuntu/veenaModal/scripts/stress_test_local.py:670:    parser.add_argument("--output", default="stress_test_results.json", help="Output JSON path")
/home/ubuntu/veenaModal/scripts/stress_test_local.py:677:    asyncio.run(run_stress_test(args))


Summary: auto


sed -n '1,260p' /home/ubuntu/veenaModal/scripts/profile_pipeline.py


Output:
Chunk ID: 1c6f88
Wall time: 0.0518 seconds
Process exited with code 0
Original token count: 2188
Output:
#!/usr/bin/env python3
"""
Detailed Pipeline Profiler for Veena3 TTS.

Instruments every stage of the TTS pipeline with microsecond precision:
- Prompt building
- vLLM prefill (prompt encoding)
- vLLM decode (per-token generation)
- Token parsing (BiCodec extraction)
- BiCodec audio decode (GPU)
- Crossfade / audio post-processing
- WAV header creation

For streaming: tracks token-by-token timing to find bottleneck.
For concurrency: measures contention points under parallel load.

Usage:
    # Single request breakdown
    python scripts/profile_pipeline.py

    # Concurrency profiling (1, 5, 10, 20 parallel requests)
    python scripts/profile_pipeline.py --concurrency 1,5,10,20

    # Long text profiling
    python scripts/profile_pipeline.py --text long

    # Streaming profiling
    python scripts/profile_pipeline.py --stream
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import os
import re
import sys
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

# === Path setup (same as local_server.py) ===
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
EXTERNAL_DIR = os.path.join(REPO_ROOT, "external")
for path in [os.path.join(EXTERNAL_DIR, "sparktts"), os.path.join(EXTERNAL_DIR, "AP-BWE"), REPO_ROOT]:
    if path not in sys.path:
        sys.path.insert(0, path)

os.environ["AUTH_BYPASS_MODE"] = "true"

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("profiler")

# Test texts
TEXTS = {
    "short": "Hello, this is a quick test.",
    "medium": (
        "The quick brown fox jumps over the lazy dog. "
        "This sentence tests basic text normalization and generation."
    ),
    "long": (
        "In the heart of a bustling city, where towering skyscrapers cast long shadows "
        "over crowded streets, there lived an old bookkeeper named Margaret. She had spent "
        "forty years cataloging the stories of others, yet never found time to write her own. "
        "One rainy Tuesday afternoon, she discovered a leather-bound journal tucked behind "
        "a shelf of dusty encyclopedias. Its pages were blank, waiting."
    ),
}

# === Timing Utilities ===

@dataclass
class StageTimer:
    """Accumulated timing for a named pipeline stage."""
    name: str
    calls: int = 0
    total_ms: float = 0.0
    min_ms: float = float("inf")
    max_ms: float = 0.0
    samples: List[float] = field(default_factory=list)

    def record(self, ms: float):
        self.calls += 1
        self.total_ms += ms
        self.min_ms = min(self.min_ms, ms)
        self.max_ms = max(self.max_ms, ms)
        self.samples.append(ms)

    @property
    def avg_ms(self) -> float:
        return self.total_ms / self.calls if self.calls > 0 else 0

    def report_line(self) -> str:
        if self.calls == 0:
            return f"  {self.name:<35} (not called)"
        if self.calls == 1:
            return f"  {self.name:<35} {self.total_ms:>8.2f}ms"
        return (
            f"  {self.name:<35} {self.total_ms:>8.2f}ms total  "
            f"avg={self.avg_ms:.2f}ms  min={self.min_ms:.2f}ms  "
            f"max={self.max_ms:.2f}ms  calls={self.calls}"
        )

@dataclass
class RequestProfile:
    """Complete timing profile for one TTS request."""
    request_id: str
    text: str
    speaker: str
    stream: bool
    stages: Dict[str, StageTimer] = field(default_factory=dict)

    # Metadata filled during profiling
    total_ms: float = 0.0
    prompt_tokens: int = 0
    generated_tokens: int = 0
    semantic_tokens: int = 0
    global_tokens: int = 0
    audio_bytes: int = 0
    audio_seconds: float = 0.0

    # Streaming-specific
    ttfb_ms: float = 0.0
    chunks_emitted: int = 0
    token_generation_rates: List[float] = field(default_factory=list)  # tokens/sec samples
    decode_latencies: List[float] = field(default_factory=list)  # ms per BiCodec decode

    def stage(self, name: str) -> StageTimer:
        if name not in self.stages:
            self.stages[name] = StageTimer(name=name)
        return self.stages[name]

    def report(self) -> str:
        lines = [
            f"\n{'=' * 80}",
            f"  REQUEST PROFILE: {self.request_id}",
            f"  text={len(self.text)} chars, speaker={self.speaker}, stream={self.stream}",
            f"{'=' * 80}",
            f"  TOTAL: {self.total_ms:.1f}ms",
            f"  Generated: {self.generated_tokens} tokens ({self.semantic_tokens} semantic, {self.global_tokens} global)",
            f"  Audio: {self.audio_seconds:.2f}s ({self.audio_bytes} bytes)",
        ]
        if self.stream:
            lines.append(f"  TTFB: {self.ttfb_ms:.1f}ms, Chunks: {self.chunks_emitted}")
        if self.generated_tokens > 0 and self.total_ms > 0:
            tps = self.generated_tokens / (self.total_ms / 1000)
            lines.append(f"  Token gen rate: {tps:.1f} tok/s")
        if self.audio_seconds > 0:
            rtf = (self.total_ms / 1000) / self.audio_seconds
            lines.append(f"  RTF: {rtf:.3f}")

        lines.append(f"\n  {'STAGE BREAKDOWN':}")
        lines.append(f"  {'-' * 76}")

        # Order stages by total time (descending)
        sorted_stages = sorted(self.stages.values(), key=lambda s: s.total_ms, reverse=True)
        for s in sorted_stages:
            lines.append(s.report_line())

        # Token generation rate over time (streaming)
        if self.token_generation_rates:
            lines.append(f"\n  TOKEN GEN RATE (tok/s over time):")
            # Show first 5, middle, last 5
            rates = self.token_generation_rates
            if len(rates) <= 12:
                lines.append(f"    {' → '.join(f'{r:.0f}' for r in rates)}")
            else:
                first = ' → '.join(f'{r:.0f}' for r in rates[:5])
                last = ' → '.join(f'{r:.0f}' for r in rates[-5:])
                lines.append(f"    {first} → ... → {last}")
            import statistics
            lines.append(f"    avg={statistics.mean(rates):.0f}  min={min(rates):.0f}  max={max(rates):.0f}")

        # BiCodec decode latencies
        if self.decode_latencies:
            lines.append(f"\n  BICODEC DECODE LATENCIES (ms per call):")
            import statistics
            lines.append(
                f"    avg={statistics.mean(self.decode_latencies):.2f}  "
                f"min={min(self.decode_latencies):.2f}  "
                f"max={max(self.decode_latencies):.2f}  "
                f"calls={len(self.decode_latencies)}"
            )

        lines.append(f"{'=' * 80}")
        return "\n".join(lines)

# === Non-Streaming Profiler ===

async def profile_non_streaming(
    runtime,
    text: str,
    speaker: str = "lipakshi",
    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,
) -> RequestProfile:
    """Profile a single non-streaming request with stage-level timing."""
    from vllm import SamplingParams
    from veena3modal.core.constants import TRAINING_STOP_TOKEN_IDS, AUDIO_SAMPLE_RATE
    from veena3modal.audio.utils import add_wav_header

    profile = RequestProfile(
        request_id=f"prof-{uuid.uuid4().hex[:8]}",
        text=text,
        speaker=speaker,
        stream=False,
    )

    t_total_start = time.perf_counter()

    # === Stage 1: Prompt Building ===
    t0 = time.perf_counter()
    prompt = runtime.prompt_builder.build_prefix(speaker, text)
    profile.stage("1_prompt_build").record((time.perf_counter() - t0) * 1000)
    profile.prompt_tokens = len(runtime.model.tokenizer.encode(prompt, add_special_tokens=False))

    # === Stage 2: Sampling Config ===
    t0 = time.perf_counter()
    sampling_params = SamplingParams(
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        max_tokens=max_tokens,
        repetition_penalty=repetition_penalty,
        stop=TRAINING_STOP_TOKEN_IDS,
        skip_special_tokens=False,
        seed=seed,
    )
    profile.stage("2_sampling_config").record((time.perf_counter() - t0) * 1000)

    # === Stage 3: vLLM Token Generation ===
    request_id = f"prof-{uuid.uuid4().hex[:8]}"
    t_gen_start = time.perf_counter()
    t_first_output = None
    iteration_count = 0
    prev_token_count = 0

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

    final_output = None
    async for request_output in results_generator:
        iteration_count += 1
        now = time.perf_counter()


I’m set to patch three runtime layers in one pass: streaming pipeline instrumentation, request-id propagation to streaming endpoints, and a new local stress benchmark script that logs per-request, per-batch, and per-decode timing plus GPU metrics.


nl -ba /home/ubuntu/veenaModal/veena3modal/core/streaming_pipeline.py | sed -n '1,260p'


Output:
Chunk ID: a5f6de
Wall time: 0.0528 seconds
Process exited with code 0
Original token count: 2936
Output:
     1	"""
     2	Veena3 Streaming Pipeline - Sliding Window Approach
     3
     4	Implements streaming for both SNAC and BiCodec tokens.
     5
     6	For BiCodec (Spark TTS):
     7	1. Generate tokens from vLLM
     8	2. Parse NEW token IDs incrementally (O(1) per token via cache)
     9	3. Buffer semantic + global tokens
    10	4. Apply sliding window (every N token pairs)
    11	5. Decode and stream audio chunks
    12
    13	OPTIMIZATION (Dec 2025):
    14	- Replaced O(n²) pattern (decode-all + regex-all per iteration)
    15	- Now uses incremental token parsing with BiCodecTokenParser
    16	- ~10x CPU reduction in streaming hot loop
    17
    18	For SNAC (legacy):
    19	1. Filter SNAC token IDs directly
    20	2. Apply sliding window
    21	3. Decode and stream
    22	"""
    23
    24	import asyncio
    25	import os
    26	import re
    27	from statistics import median, mean
    28	import torch
    29	from typing import AsyncGenerator, Optional, List
    30	from vllm import SamplingParams
    31	from veena3modal.audio.crossfade import crossfade_bytes_int16
    32	from veena3modal.core.token_utils import BiCodecTokenParser
    33
    34	from veena3modal.core.constants import (
    35	    CODE_END_TOKEN_ID,
    36	    CODE_START_TOKEN_ID,
    37	    SNAC_MIN_ID,
    38	    SNAC_MAX_ID,
    39	    TRAINING_STOP_TOKEN_IDS,
    40	    DEFAULT_TEMPERATURE,
    41	    DEFAULT_TOP_K,
    42	    DEFAULT_TOP_P,
    43	    DEFAULT_MAX_TOKENS,
    44	    DEFAULT_MIN_TOKENS,
    45	    DEFAULT_REPETITION_PENALTY,
    46	    DEFAULT_SEED,
    47	)
    48
    49
    50	def _detailed_timing_enabled() -> bool:
    51	    """Enable per-batch token/decode timing logs."""
    52	    return os.getenv("VEENA3_TIMING_LOGS", "").lower() in {"1", "true", "yes", "on"}
    53
    54
    55	class Veena3SlidingWindowPipeline:
    56	    """
    57	    Streaming TTS pipeline using sliding window approach.
    58
    59	    This eliminates choppy audio and popping artifacts by:
    60	    - Decoding overlapping 28-token windows (4 frames)
    61	    - Keeping only the middle 2048 samples from each decode
    62	    - Creating natural continuity between chunks
    63
    64	    Based on the official Canopy Labs implementation.
    65	    """
    66
    67	    def __init__(
    68	        self,
    69	        model,
    70	        prompt_builder,
    71	        snac_decoder,
    72	    ):
    73	        """
    74	        Initialize sliding window streaming pipeline.
    75
    76	        Args:
    77	            model: Veena3Model instance
    78	            prompt_builder: Veena3PromptBuilder instance
    79	            snac_decoder: SNACDecoder instance (with batching enabled)
    80	        """
    81	        self.model = model
    82	        self.prompt_builder = prompt_builder
    83	        self.snac_decoder = snac_decoder
    84
    85	        # OPTIMIZATION: Pre-warm BiCodecTokenParser once at init, not per-request
    86	        # Saves ~123ms per streaming TTFB by avoiding repeated 166K-entry vocab iteration
    87	        tokenizer = getattr(model, "tokenizer", None)
    88	        if tokenizer is None:
    89	            tokenizer = getattr(model.engine, "tokenizer", None)
    90	        self.token_parser = BiCodecTokenParser(tokenizer) if tokenizer else None
    91
    92	        print(f"🌊 Veena3SlidingWindowPipeline initialized (sliding window: 28 tokens)")
    93
    94	    async def generate_speech_stream(
    95	        self,
    96	        description: str,
    97	        text: str,
    98	        temperature: float = DEFAULT_TEMPERATURE,
    99	        top_p: float = DEFAULT_TOP_P,
   100	        max_tokens: int = DEFAULT_MAX_TOKENS,
   101	        repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
   102	        seed: Optional[int] = None,
   103	    ) -> AsyncGenerator[bytes, None]:
   104	        """
   105	        Generate speech audio with sliding window streaming.
   106
   107	        Yields audio chunks using overlapping windows for smooth playback.
   108
   109	        Args:
   110	            description: Character/voice description
   111	            text: Text to synthesize (with optional <emotion> tags)
   112	            temperature: Sampling temperature
   113	            top_p: Nucleus sampling
   114	            max_tokens: Max SNAC tokens to generate
   115	            repetition_penalty: Prevent loops
   116
   117	        Yields:
   118	            Audio bytes (int16 PCM, 24kHz mono)
   119	        """
   120	        print(f"\n🌊 Sliding window streaming generation")
   121	        print(f"📝 Description: {description[:80]}...")
   122	        print(f"💬 Text: {text}")
   123
   124	        # Build prompt
   125	        prompt = self.prompt_builder.build_prefix(description, text)
   126	        print(f"✅ Prompt built ({len(prompt)} chars)")
   127
   128	        # Configure sampling
   129	        sampling_params = SamplingParams(
   130	            temperature=temperature,
   131	            top_p=top_p,
   132	            max_tokens=max_tokens,
   133	            min_tokens=DEFAULT_MIN_TOKENS,
   134	            repetition_penalty=repetition_penalty,
   135	            stop_token_ids=[CODE_END_TOKEN_ID],  # Only audio EOS
   136	            seed=seed,  # None = random, int = reproducible
   137	        )
   138
   139	        print(f"🎲 Sampling: temp={temperature}, top_p={top_p}, sliding_window=28 tokens")
   140
   141	        # Token buffer - keeps ALL tokens (not chunked)
   142	        token_buffer = []
   143	        total_tokens_generated = 0
   144	        total_audio_chunks = 0
   145
   146	        # Generate tokens with vLLM (streaming)
   147	        print(f"🔮 Starting token generation...")
   148
   149	        # Generate unique request ID for concurrent streaming support
   150	        import uuid
   151	        import time
   152	        request_id = f"slide-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"
   153	        results_generator = self.model.engine.generate(
   154	            prompt=prompt,
   155	            sampling_params=sampling_params,
   156	            request_id=request_id,
   157	        )
   158
   159	        # Stream tokens with sliding window
   160	        async for request_output in results_generator:
   161	            # Extract generated token IDs
   162	            generated_ids = request_output.outputs[0].token_ids
   163
   164	            # Process only new tokens
   165	            new_tokens = generated_ids[total_tokens_generated:]
   166	            total_tokens_generated = len(generated_ids)
   167
   168	            # Filter and buffer SNAC tokens
   169	            for token_id in new_tokens:
   170	                if SNAC_MIN_ID <= token_id <= SNAC_MAX_ID:
   171	                    token_buffer.append(token_id)
   172
   173	                    # Process every 7 tokens (1 frame) when we have enough for sliding window
   174	                    # Official approach: once count > 27, take last 28 tokens
   175	                    if len(token_buffer) % 7 == 0 and len(token_buffer) > 27:
   176	                        # Sliding window: take last 28 tokens (4 frames)
   177	                        window_tokens = token_buffer[-28:]
   178
   179	                        # Decode with sliding window mode (returns middle 2048 samples only)
   180	                        if self.snac_decoder.enable_batching:
   181	                            audio_bytes = await self.snac_decoder.decode_single_async(
   182	                                window_tokens,
   183	                                trim_warmup=False,  # Sliding window handles trimming
   184	                                use_sliding_window=True  # CRITICAL: Use sliding window mode
   185	                            )
   186	                        else:
   187	                            audio_bytes = self.snac_decoder.decode_to_bytes(
   188	                                window_tokens,
   189	                                trim_warmup=False,
   190	                                use_sliding_window=True
   191	                            )
   192
   193	                        if audio_bytes:
   194	                            total_audio_chunks += 1
   195	                            if total_audio_chunks == 1:
   196	                                print(f"🎵 First chunk decoded ({len(audio_bytes)} bytes, sliding window)")
   197	                            yield audio_bytes
   198
   199	        # Note: No final chunk processing needed - sliding window handles all tokens
   200	        # as they come in (every 7 tokens after the first 28)
   201
   202	        print(f"✅ Sliding window streaming complete: {total_tokens_generated} tokens → {total_audio_chunks} audio chunks")
   203
   204	    def _extract_bicodec_tokens_from_text(self, text: str) -> tuple[List[int], List[int]]:
   205	        """
   206	        Extract BiCodec semantic and global tokens from generated text.
   207
   208	        Args:
   209	            text: Generated text containing BiCodec token markers
   210
   211	        Returns:
   212	            Tuple of (semantic_ids, global_ids)
   213	        """
   214	        semantic_matches = re.findall(r"<\|bicodec_semantic_(\d+)\|>", text)
   215	        semantic_ids = [int(t) for t in semantic_matches]
   216
   217	        global_matches = re.findall(r"<\|bicodec_global_(\d+)\|>", text)
   218	        global_ids = [int(t) for t in global_matches]
   219
   220	        return semantic_ids, global_ids
   221
   222	    async def generate_speech_stream_indic(
   223	        self,
   224	        speaker: str,
   225	        text: str,
   226	        temperature: float = DEFAULT_TEMPERATURE,
   227	        top_k: int = DEFAULT_TOP_K,  # Added for Spark TTS parity with non-streaming
   228	        top_p: float = DEFAULT_TOP_P,
   229	        max_tokens: int = DEFAULT_MAX_TOKENS,
   230	        repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
   231	        seed: Optional[int] = None,
   232	        request_id: Optional[str] = None,
   233	    ) -> AsyncGenerator[bytes, None]:
   234	        """
   235	        Generate speech audio with sliding window streaming for Indic model.
   236
   237	        Yields audio chunks using overlapping windows for smooth playback.
   238
   239	        Args:
   240	            speaker: Speaker name (one of 12 predefined speakers)
   241	            text: Text to synthesize with inline emotion tags
   242	                Examples:
   243	                - "Hello! Welcome."
   244	                - "<laugh> Hello there!"
   245	                - "नमस्ते! <excited> आज का दिन बहुत अच्छा है।"
   246	            temperature: Sampling temperature
   247	            top_p: Nucleus sampling
   248	            max_tokens: Max SNAC tokens to generate
   249	            repetition_penalty: Prevent loops
   250	            seed: Random seed for reproducibility
   251
   252	        Yields:
   253	            Audio bytes (int16 PCM, 16kHz mono - BiCodec)
   254	        """
   255	        import time
   256	        t_start = time.time()
   257	        emit_timing = _detailed_timing_enabled()
   258
   259	        # Build prompt using Indic prompt builder
   260	        prompt = self.prompt_builder.build_prefix(speaker, text)


nl -ba /home/ubuntu/veenaModal/veena3modal/core/streaming_pipeline.py | sed -n '260,560p'


Output:
Chunk ID: 0b4cbb
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 4412
Output:
   260	        prompt = self.prompt_builder.build_prefix(speaker, text)
   261
   262	        # Configure sampling (matching non-streaming pipeline parameters)
   263	        # OPTIMIZATION: Use Spark TTS stop token, not legacy SNAC CODE_END_TOKEN_ID (128258)
   264	        # which doesn't exist in the Spark vocab and would cause generation to run to max_tokens
   265	        sampling_params = SamplingParams(
   266	            temperature=temperature,
   267	            top_k=top_k,
   268	            top_p=top_p,
   269	            max_tokens=max_tokens,
   270	            min_tokens=DEFAULT_MIN_TOKENS,
   271	            repetition_penalty=repetition_penalty,
   272	            stop=TRAINING_STOP_TOKEN_IDS,  # "<|im_end|>" - matches non-streaming pipeline
   273	            skip_special_tokens=False,  # Keep BiCodec tokens in output
   274	            seed=seed,
   275	        )
   276
   277	        # BiCodec token buffers (separate for semantic and global)
   278	        semantic_buffer = []
   279	        global_buffer = []
   280	        processed_token_count = 0  # Track how many vLLM tokens we've parsed
   281	        total_audio_chunks = 0
   282
   283	        # BiCodec streaming configuration
   284	        EXPECTED_GLOBAL_COUNT = 32  # BiCodec uses 32 fixed global tokens (per paper)
   285	        MIN_SEMANTIC_FOR_FIRST_CHUNK = 10  # OPTIMIZATION: Lowered from 16 (decoder min is 8, -120ms TTFB)
   286	        # OPTIMIZATION: Increased from 24 to 48 to reduce O(n) re-decode calls by 2x.
   287	        # At 50 TPS, 48 tokens = ~960ms audio per chunk. Cuts total decode calls roughly in half.
   288	        DECODE_INTERVAL = 48
   289	        CROSSFADE_MS = 50
   290	        # OPTIMIZATION: Windowed decode - only decode last WINDOW_SIZE tokens instead of ALL.
   291	        # BiCodec WaveGenerator has finite receptive field (~64 tokens based on kernel sizes).
   292	        # Window must be >= receptive_field + DECODE_INTERVAL to avoid boundary artifacts.
   293	        WINDOW_SIZE = 128  # 128 tokens covers receptive field + margin for crossfade
   294
   295	        # Generate tokens with vLLM (streaming)
   296	        if request_id is None:
   297	            import uuid
   298	            request_id = f"bicodec-stream-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"
   299
   300	        # Reuse pre-warmed token parser (singleton, created once in __init__)
   301	        token_parser = self.token_parser
   302
   303	        t_gen_start = time.time()
   304	        batch_wall_ms: List[float] = []
   305	        batch_token_counts: List[int] = []
   306	        decode_calls = 0
   307	        decode_wall_ms = 0.0
   308	        decode_gpu_ms = 0.0
   309	        parse_wall_ms = 0.0
   310
   311	        results_generator = self.model.engine.generate(
   312	            prompt=prompt,
   313	            sampling_params=sampling_params,
   314	            request_id=request_id,
   315	        )
   316
   317	        # Track timing
   318	        t_first_token = None
   319	        t_first_chunk_ready = None
   320	        t_first_chunk_decoded = None
   321	        iteration_count = 0
   322	        last_decode_count = 0
   323
   324	        # Track audio samples EMITTED to the user (not just decoded)
   325	        # BiCodec produces ~320 samples per semantic token at 16kHz
   326	        # CRITICAL: This tracks samples that have been YIELDED, not held in tail
   327	        total_samples_emitted_to_user = 0
   328	        # Keep a tail to enable crossfading between consecutive emitted chunks
   329	        previous_chunk_tail: Optional[bytes] = None
   330
   331	        # Stream tokens with BiCodec extraction
   332	        async for request_output in results_generator:
   333	            iteration_count += 1
   334	            t_iter = time.time()
   335
   336	            # Track first token received
   337	            if t_first_token is None:
   338	                t_first_token = t_iter
   339
   340	            # Extract generated token IDs
   341	            generated_ids = request_output.outputs[0].token_ids
   342
   343	            parse_start = time.time()
   344	            # OPTIMIZATION: Only process NEW tokens (incremental parsing)
   345	            # Old code: tokenizer.decode(ALL) + re.findall(ALL) = O(n²) per stream
   346	            # New code: parse only new tokens = O(1) per token
   347	            new_token_ids = generated_ids[processed_token_count:]
   348	            processed_token_count = len(generated_ids)
   349
   350	            # Parse new tokens and append to buffers (in-place modification)
   351	            token_parser.parse_incremental(new_token_ids, semantic_buffer, global_buffer)
   352	            parse_wall_ms += (time.time() - parse_start) * 1000
   353	            batch_token_counts.append(len(new_token_ids))
   354	            batch_wall_ms.append(0 if t_iter is None else (time.time() - t_iter) * 1000)
   355
   356	            # TRUE STREAMING with BiCodec two-phase generation:
   357	            # Phase 1: Generate 32 global tokens (pre-roll)
   358	            # Phase 2: Stream semantic tokens incrementally
   359
   360	            # Check if we have the 32 global tokens (pre-roll complete)
   361	            if len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
   362	                # Pre-roll complete! Now we can stream semantic tokens
   363	                if t_first_chunk_ready is None and len(semantic_buffer) >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
   364	                    t_first_chunk_ready = time.time()
   365
   366	                # Stream semantic tokens incrementally
   367	                # Only decode if we have enough NEW semantic tokens since last decode
   368	                semantic_count = len(semantic_buffer)
   369	                if semantic_count >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
   370	                    # Check if we should decode (every DECODE_INTERVAL new semantic tokens)
   371	                    if semantic_count - last_decode_count >= DECODE_INTERVAL:
   372	                        last_decode_count = semantic_count
   373
   374	                        # OPTIMIZATION: Windowed decode instead of full re-decode
   375	                        # Only decode the last WINDOW_SIZE tokens (O(1) per chunk instead of O(n))
   376	                        # For first few chunks when total tokens < WINDOW_SIZE, decode all
   377	                        window_global = global_buffer[:EXPECTED_GLOBAL_COUNT]
   378
   379	                        if len(semantic_buffer) <= WINDOW_SIZE:
   380	                            # Early phase: decode all tokens (window not yet relevant)
   381	                            decode_semantic = semantic_buffer
   382	                        else:
   383	                            # Windowed: only decode last WINDOW_SIZE tokens
   384	                            decode_semantic = semantic_buffer[-WINDOW_SIZE:]
   385
   386	                        t_decode_start = time.time()
   387	                        if emit_timing and torch.cuda.is_available():
   388	                            torch.cuda.synchronize()
   389	                            decode_start_event = torch.cuda.Event(enable_timing=True)
   390	                            decode_end_event = torch.cuda.Event(enable_timing=True)
   391	                            decode_start_event.record()
   392	                        audio_bytes = await self.snac_decoder.decode_single_async(
   393	                            semantic_ids=decode_semantic,
   394	                            global_ids=window_global,
   395	                            trim_warmup=False,
   396	                            use_sliding_window=False,
   397	                        )
   398	                        t_decode_end = time.time()
   399	                        decode_calls += 1
   400	                        decode_wall_ms += (t_decode_end - t_decode_start) * 1000
   401	                        if emit_timing and torch.cuda.is_available():
   402	                            decode_end_event.record()
   403	                            torch.cuda.synchronize()
   404	                            decode_gpu_ms += decode_start_event.elapsed_time(decode_end_event)
   405
   406	                        if audio_bytes:
   407	                            total_samples_decoded = len(audio_bytes) // 2
   408
   409	                            if len(semantic_buffer) <= WINDOW_SIZE:
   410	                                # Early phase: extract new samples from full decode
   411	                                if total_samples_decoded > total_samples_emitted_to_user:
   412	                                    new_bytes_start = total_samples_emitted_to_user * 2
   413	                                    new_audio_bytes = audio_bytes[new_bytes_start:]
   414	                                else:
   415	                                    new_audio_bytes = b""
   416	                            else:
   417	                                # Windowed phase: take only the tail portion of the window decode
   418	                                # The last DECODE_INTERVAL tokens' worth of audio is the new content
   419	                                # Each token produces ~320 samples at 16kHz/50TPS
   420	                                new_samples_approx = DECODE_INTERVAL * 320
   421	                                if total_samples_decoded > new_samples_approx:
   422	                                    new_bytes_start = (total_samples_decoded - new_samples_approx) * 2
   423	                                else:
   424	                                    new_bytes_start = 0
   425	                                new_audio_bytes = audio_bytes[new_bytes_start:]
   426
   427	                            if new_audio_bytes:
   428	                                to_emit, previous_chunk_tail = crossfade_bytes_int16(
   429	                                    previous_chunk_tail,
   430	                                    new_audio_bytes,
   431	                                    sample_rate_hz=16000,
   432	                                    crossfade_ms=CROSSFADE_MS,
   433	                                )
   434
   435	                                samples_emitted_in_this_chunk = len(to_emit) // 2
   436	                                total_samples_emitted_to_user += samples_emitted_in_this_chunk
   437
   438	                                if to_emit:
   439	                                    total_audio_chunks += 1
   440	                                    if t_first_chunk_decoded is None:
   441	                                        t_first_chunk_decoded = t_decode_end
   442
   443	                                    if emit_timing:
   444	                                        logger.info(
   445	                                            "tts_stream_decode_batch",
   446	                                            extra={
   447	                                                "speaker": speaker,
   448	                                                "request_id": request_id,
   449	                                                "batch_index": decode_calls,
   450	                                                "decode_wall_ms": int((t_decode_end - t_decode_start) * 1000),
   451	                                                "decode_gpu_ms": int(decode_start_event.elapsed_time(decode_end_event))
   452	                                                if torch.cuda.is_available() else None,
   453	                                                "semantic_count": len(semantic_buffer),
   454	                                            },
   455	                                        )
   456
   457	                                    yield to_emit
   458
   459	        # Final chunk: decode any remaining semantic tokens
   460	        # Use the same 32 global tokens (they get pooled)
   461	        if len(semantic_buffer) > last_decode_count and len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
   462	            audio_bytes = self.snac_decoder.decode_streaming(
   463	                semantic_ids=semantic_buffer,  # All semantic tokens
   464	                global_ids=global_buffer[:EXPECTED_GLOBAL_COUNT],  # Use first 32 global tokens
   465	                use_sliding_window=False,  # Get all samples
   466	                trim_warmup=False
   467	            )
   468	            if audio_bytes:
   469	                # Only yield NEW samples
   470	                total_samples_decoded = len(audio_bytes) // 2
   471
   472	                if total_samples_decoded > total_samples_emitted_to_user:
   473	                    new_bytes_start = total_samples_emitted_to_user * 2
   474	                    new_audio_bytes = audio_bytes[new_bytes_start:]
   475
   476	                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
   477	                        previous_chunk_tail,
   478	                        new_audio_bytes,
   479	                        sample_rate_hz=16000,
   480	                        crossfade_ms=CROSSFADE_MS,
   481	                    )
   482
   483	                    samples_emitted_in_final = len(to_emit) // 2
   484	                    total_samples_emitted_to_user += samples_emitted_in_final
   485
   486	                    if to_emit:
   487	                        total_audio_chunks += 1
   488	                        yield to_emit
   489
   490	        # Flush any remaining tail that was held back for crossfade
   491	        if previous_chunk_tail:
   492	            total_audio_chunks += 1
   493	            tail_samples = len(previous_chunk_tail) // 2
   494	            total_samples_emitted_to_user += tail_samples
   495	            yield previous_chunk_tail
   496
   497	        # Log completion summary
   498	        t_end = time.time()
   499	        audio_duration_s = total_samples_emitted_to_user / 16000  # 16kHz sample rate
   500	        ttfb_ms = (t_first_chunk_decoded - t_start) * 1000 if t_first_chunk_decoded else 0
   501	        print(f"🎵 Streaming complete: {audio_duration_s:.2f}s audio, TTFB: {ttfb_ms:.0f}ms, {total_audio_chunks} chunks")
   502
   503	    async def generate_speech_stream_indic_first_chunk(
   504	        self,
   505	        speaker: str,
   506	        text: str,
   507	        temperature: float = DEFAULT_TEMPERATURE,
   508	        top_k: int = DEFAULT_TOP_K,
   509	        top_p: float = DEFAULT_TOP_P,
   510	        max_tokens: int = DEFAULT_MAX_TOKENS,
   511	        repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
   512	        seed: Optional[int] = None,
   513	    ) -> AsyncGenerator[tuple, None]:
   514	        """
   515	        Generate speech for the FIRST chunk of multi-chunk text, capturing global tokens.
   516
   517	        This method is specifically for chunked generation: it yields (audio_bytes, global_ids)
   518	        tuples instead of just audio_bytes. The caller captures global_ids from the first
   519	        yield and passes them to generate_speech_stream_indic_continuation() for subsequent chunks.
   520
   521	        Use Case (chunked streaming with voice consistency):
   522	            globals_captured = None
   523	            async for audio_bytes, global_ids in pipeline.generate_speech_stream_indic_first_chunk(...):
   524	                if globals_captured is None and global_ids:
   525	                    globals_captured = global_ids  # Capture once
   526	                yield audio_bytes
   527	            # Now use globals_captured for subsequent chunks
   528
   529	        Args:
   530	            speaker: Speaker name
   531	            text: First text chunk to synthesize
   532	            temperature: Sampling temperature
   533	            top_k: Top-k sampling
   534	            top_p: Nucleus sampling
   535	            max_tokens: Max tokens to generate
   536	            repetition_penalty: Prevent repetition
   537	            seed: Random seed for reproducibility
   538
   539	        Yields:
   540	            Tuple of (audio_bytes, global_ids)
   541	            - audio_bytes: Raw PCM audio (int16, 16kHz)
   542	            - global_ids: List of 32 global token IDs (populated after first decode, else empty)
   543
   544	        Thread Safety:
   545	            This method is stateless and thread-safe. Each call creates its own
   546	            request-scoped state. No global state is modified or shared.
   547	        """
   548	        import time
   549	        t_start = time.time()
   550
   551	        # Build prompt using Indic prompt builder
   552	        prompt = self.prompt_builder.build_prefix(speaker, text)
   553
   554	        # Configure sampling
   555	        # OPTIMIZATION: Use Spark TTS stop token, not legacy SNAC CODE_END_TOKEN_ID
   556	        sampling_params = SamplingParams(
   557	            temperature=temperature,
   558	            top_k=top_k,
   559	            top_p=top_p,
   560	            max_tokens=max_tokens,


nl -ba /home/ubuntu/veenaModal/veena3modal/core/streaming_pipeline.py | sed -n '560,940p'


Output:
Chunk ID: afb142
Wall time: 0.0511 seconds
Process exited with code 0
Original token count: 4467
Output:
   560	            max_tokens=max_tokens,
   561	            min_tokens=DEFAULT_MIN_TOKENS,
   562	            repetition_penalty=repetition_penalty,
   563	            stop=TRAINING_STOP_TOKEN_IDS,  # "<|im_end|>" - matches non-streaming pipeline
   564	            skip_special_tokens=False,
   565	            seed=seed,
   566	        )
   567
   568	        # BiCodec token buffers
   569	        semantic_buffer = []
   570	        global_buffer = []
   571	        processed_token_count = 0  # Track how many vLLM tokens we've parsed
   572	        total_audio_chunks = 0
   573
   574	        # BiCodec streaming configuration
   575	        EXPECTED_GLOBAL_COUNT = 32
   576	        MIN_SEMANTIC_FOR_FIRST_CHUNK = 10  # OPTIMIZATION: Lowered from 16 (decoder min is 8)
   577	        DECODE_INTERVAL = 48  # OPTIMIZATION: Increased from 24 to reduce O(n) re-decode calls
   578	        CROSSFADE_MS = 50
   579
   580	        # Generate unique request ID
   581	        import uuid
   582	        request_id = f"bicodec-first-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"
   583
   584	        # Reuse pre-warmed token parser (singleton, created once in __init__)
   585	        token_parser = self.token_parser
   586
   587	        results_generator = self.model.engine.generate(
   588	            prompt=prompt,
   589	            sampling_params=sampling_params,
   590	            request_id=request_id,
   591	        )
   592
   593	        # Track state
   594	        last_decode_count = 0
   595	        total_samples_emitted_to_user = 0
   596	        previous_chunk_tail: Optional[bytes] = None
   597	        captured_globals: List[int] = []  # Will be populated once we have 32 globals
   598
   599	        # Stream tokens with BiCodec extraction
   600	        async for request_output in results_generator:
   601	            generated_ids = request_output.outputs[0].token_ids
   602
   603	            # OPTIMIZATION: Incremental parsing (O(1) per new token)
   604	            new_token_ids = generated_ids[processed_token_count:]
   605	            processed_token_count = len(generated_ids)
   606	            token_parser.parse_incremental(new_token_ids, semantic_buffer, global_buffer)
   607
   608	            # Capture global tokens once we have all 32
   609	            if not captured_globals and len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
   610	                captured_globals = global_buffer[:EXPECTED_GLOBAL_COUNT]
   611
   612	            # Process if we have all global tokens
   613	            if len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
   614	                semantic_count = len(semantic_buffer)
   615	                if semantic_count >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
   616	                    if semantic_count - last_decode_count >= DECODE_INTERVAL:
   617	                        last_decode_count = semantic_count
   618
   619	                        all_semantic = semantic_buffer
   620	                        window_global = global_buffer[:EXPECTED_GLOBAL_COUNT]
   621
   622	                        if self.snac_decoder.enable_batching:
   623	                            audio_bytes = await self.snac_decoder.decode_single_async(
   624	                                semantic_ids=all_semantic,
   625	                                global_ids=window_global,
   626	                                trim_warmup=False,
   627	                                use_sliding_window=False
   628	                            )
   629	                        else:
   630	                            audio_bytes = self.snac_decoder.decode_streaming(
   631	                                semantic_ids=all_semantic,
   632	                                global_ids=window_global,
   633	                                use_sliding_window=False,
   634	                                trim_warmup=False
   635	                            )
   636
   637	                        if audio_bytes:
   638	                            total_samples_decoded = len(audio_bytes) // 2
   639
   640	                            if total_samples_decoded > total_samples_emitted_to_user:
   641	                                new_bytes_start = total_samples_emitted_to_user * 2
   642	                                new_audio_bytes = audio_bytes[new_bytes_start:]
   643
   644	                                to_emit, previous_chunk_tail = crossfade_bytes_int16(
   645	                                    previous_chunk_tail,
   646	                                    new_audio_bytes,
   647	                                    sample_rate_hz=16000,
   648	                                    crossfade_ms=CROSSFADE_MS,
   649	                                )
   650
   651	                                samples_emitted_in_this_chunk = len(to_emit) // 2
   652	                                total_samples_emitted_to_user += samples_emitted_in_this_chunk
   653
   654	                                if to_emit:
   655	                                    total_audio_chunks += 1
   656	                                    # Yield tuple: (audio_bytes, global_ids)
   657	                                    yield (to_emit, captured_globals)
   658
   659	        # Final chunk processing
   660	        if len(semantic_buffer) > last_decode_count and len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
   661	            audio_bytes = self.snac_decoder.decode_streaming(
   662	                semantic_ids=semantic_buffer,
   663	                global_ids=global_buffer[:EXPECTED_GLOBAL_COUNT],
   664	                use_sliding_window=False,
   665	                trim_warmup=False
   666	            )
   667	            if audio_bytes:
   668	                total_samples_decoded = len(audio_bytes) // 2
   669
   670	                if total_samples_decoded > total_samples_emitted_to_user:
   671	                    new_bytes_start = total_samples_emitted_to_user * 2
   672	                    new_audio_bytes = audio_bytes[new_bytes_start:]
   673
   674	                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
   675	                        previous_chunk_tail,
   676	                        new_audio_bytes,
   677	                        sample_rate_hz=16000,
   678	                        crossfade_ms=CROSSFADE_MS,
   679	                    )
   680
   681	                    samples_emitted_in_final = len(to_emit) // 2
   682	                    total_samples_emitted_to_user += samples_emitted_in_final
   683
   684	                    if to_emit:
   685	                        total_audio_chunks += 1
   686	                        yield (to_emit, captured_globals)
   687
   688	        # Flush remaining tail
   689	        if previous_chunk_tail:
   690	            total_audio_chunks += 1
   691	            tail_samples = len(previous_chunk_tail) // 2
   692	            total_samples_emitted_to_user += tail_samples
   693	            yield (previous_chunk_tail, captured_globals)
   694
   695	        t_end = time.time()
   696	        audio_duration_s = total_samples_emitted_to_user / 16000
   697	        print(f"🎵 First chunk complete: {audio_duration_s:.2f}s audio, {total_audio_chunks} chunks, captured {len(captured_globals)} globals")
   698
   699	    async def generate_speech_stream_indic_continuation(
   700	        self,
   701	        speaker: str,
   702	        text: str,
   703	        global_ids: List[int],
   704	        temperature: float = DEFAULT_TEMPERATURE,
   705	        top_k: int = DEFAULT_TOP_K,
   706	        top_p: float = DEFAULT_TOP_P,
   707	        max_tokens: int = DEFAULT_MAX_TOKENS,
   708	        repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
   709	        seed: Optional[int] = None,
   710	    ) -> AsyncGenerator[bytes, None]:
   711	        """
   712	        Generate speech for CONTINUATION chunks using pre-captured global tokens.
   713
   714	        This is the key method for voice consistency in chunked generation:
   715	        - Uses global_ids from the first chunk to maintain identical voice
   716	        - Model generates only semantic tokens (skips global token generation)
   717	        - Ensures no voice drift across chunks
   718
   719	        CRITICAL for production:
   720	        - global_ids is request-scoped, passed explicitly (no shared state)
   721	        - Thread-safe: each request has its own state
   722	        - No global caching that could cause cross-request contamination
   723
   724	        Args:
   725	            speaker: Speaker name (must match first chunk)
   726	            text: Continuation text chunk to synthesize
   727	            global_ids: 32 global token IDs captured from first chunk
   728	            temperature: Sampling temperature
   729	            top_k: Top-k sampling
   730	            top_p: Nucleus sampling
   731	            max_tokens: Max tokens to generate
   732	            repetition_penalty: Prevent repetition
   733	            seed: Random seed (use same as first chunk for consistency)
   734
   735	        Yields:
   736	            Audio bytes (int16 PCM, 16kHz mono)
   737
   738	        Raises:
   739	            ValueError: If global_ids doesn't contain exactly 32 tokens
   740
   741	        Thread Safety:
   742	            Fully thread-safe. All state is request-scoped and passed explicitly.
   743	        """
   744	        import time
   745	        t_start = time.time()
   746
   747	        # Validate global tokens
   748	        EXPECTED_GLOBAL_COUNT = 32
   749	        if len(global_ids) != EXPECTED_GLOBAL_COUNT:
   750	            raise ValueError(
   751	                f"Expected exactly {EXPECTED_GLOBAL_COUNT} global tokens, got {len(global_ids)}. "
   752	                f"global_ids must be captured from first chunk generation."
   753	            )
   754
   755	        # Build prompt WITH pre-filled global tokens
   756	        # This tells the model to skip global generation and go straight to semantic tokens
   757	        prompt = self.prompt_builder.build_prefix_with_globals(speaker, text, global_ids)
   758
   759	        # Configure sampling
   760	        # OPTIMIZATION: Use Spark TTS stop token, not legacy SNAC CODE_END_TOKEN_ID
   761	        sampling_params = SamplingParams(
   762	            temperature=temperature,
   763	            top_k=top_k,
   764	            top_p=top_p,
   765	            max_tokens=max_tokens,
   766	            min_tokens=DEFAULT_MIN_TOKENS,
   767	            repetition_penalty=repetition_penalty,
   768	            stop=TRAINING_STOP_TOKEN_IDS,  # "<|im_end|>" - matches non-streaming pipeline
   769	            skip_special_tokens=False,
   770	            seed=seed,
   771	        )
   772
   773	        # BiCodec token buffers
   774	        semantic_buffer = []
   775	        global_buffer_unused = []  # Not used in continuation, but needed for parser API
   776	        processed_token_count = 0  # Track how many vLLM tokens we've parsed
   777	        total_audio_chunks = 0
   778
   779	        # BiCodec streaming configuration
   780	        MIN_SEMANTIC_FOR_FIRST_CHUNK = 10  # OPTIMIZATION: Lowered from 16 (decoder min is 8)
   781	        DECODE_INTERVAL = 48  # OPTIMIZATION: Increased from 24 to reduce O(n) re-decode calls
   782	        CROSSFADE_MS = 50
   783
   784	        # Generate unique request ID
   785	        import uuid
   786	        request_id = f"bicodec-cont-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"
   787
   788	        # Reuse pre-warmed token parser (singleton, created once in __init__)
   789	        token_parser = self.token_parser
   790
   791	        results_generator = self.model.engine.generate(
   792	            prompt=prompt,
   793	            sampling_params=sampling_params,
   794	            request_id=request_id,
   795	        )
   796
   797	        # Track state
   798	        last_decode_count = 0
   799	        total_samples_emitted_to_user = 0
   800	        previous_chunk_tail: Optional[bytes] = None
   801
   802	        # Stream tokens - model will generate ONLY semantic tokens (globals are pre-filled)
   803	        async for request_output in results_generator:
   804	            generated_ids = request_output.outputs[0].token_ids
   805
   806	            # OPTIMIZATION: Incremental parsing (O(1) per new token)
   807	            # Note: In continuation, model only generates semantic tokens (globals pre-filled in prompt)
   808	            new_token_ids = generated_ids[processed_token_count:]
   809	            processed_token_count = len(generated_ids)
   810	            token_parser.parse_incremental(new_token_ids, semantic_buffer, global_buffer_unused)
   811
   812	            semantic_count = len(semantic_buffer)
   813	            if semantic_count >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
   814	                if semantic_count - last_decode_count >= DECODE_INTERVAL:
   815	                    last_decode_count = semantic_count
   816
   817	                    # Use pre-captured global tokens for decoding
   818	                    if self.snac_decoder.enable_batching:
   819	                        audio_bytes = await self.snac_decoder.decode_single_async(
   820	                            semantic_ids=semantic_buffer,
   821	                            global_ids=global_ids,  # Use captured globals
   822	                            trim_warmup=False,
   823	                            use_sliding_window=False
   824	                        )
   825	                    else:
   826	                        audio_bytes = self.snac_decoder.decode_streaming(
   827	                            semantic_ids=semantic_buffer,
   828	                            global_ids=global_ids,  # Use captured globals
   829	                            use_sliding_window=False,
   830	                            trim_warmup=False
   831	                        )
   832
   833	                    if audio_bytes:
   834	                        total_samples_decoded = len(audio_bytes) // 2
   835
   836	                        if total_samples_decoded > total_samples_emitted_to_user:
   837	                            new_bytes_start = total_samples_emitted_to_user * 2
   838	                            new_audio_bytes = audio_bytes[new_bytes_start:]
   839
   840	                            to_emit, previous_chunk_tail = crossfade_bytes_int16(
   841	                                previous_chunk_tail,
   842	                                new_audio_bytes,
   843	                                sample_rate_hz=16000,
   844	                                crossfade_ms=CROSSFADE_MS,
   845	                            )
   846
   847	                            samples_emitted_in_this_chunk = len(to_emit) // 2
   848	                            total_samples_emitted_to_user += samples_emitted_in_this_chunk
   849
   850	                            if to_emit:
   851	                                total_audio_chunks += 1
   852	                                yield to_emit
   853
   854	        # Final chunk processing
   855	        if len(semantic_buffer) > last_decode_count:
   856	            audio_bytes = self.snac_decoder.decode_streaming(
   857	                semantic_ids=semantic_buffer,
   858	                global_ids=global_ids,  # Use captured globals
   859	                use_sliding_window=False,
   860	                trim_warmup=False
   861	            )
   862	            if audio_bytes:
   863	                total_samples_decoded = len(audio_bytes) // 2
   864
   865	                if total_samples_decoded > total_samples_emitted_to_user:
   866	                    new_bytes_start = total_samples_emitted_to_user * 2
   867	                    new_audio_bytes = audio_bytes[new_bytes_start:]
   868
   869	                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
   870	                        previous_chunk_tail,
   871	                        new_audio_bytes,
   872	                        sample_rate_hz=16000,
   873	                        crossfade_ms=CROSSFADE_MS,
   874	                    )
   875
   876	                    samples_emitted_in_final = len(to_emit) // 2
   877	                    total_samples_emitted_to_user += samples_emitted_in_final
   878
   879	                    if to_emit:
   880	                        total_audio_chunks += 1
   881	                        yield to_emit
   882
   883	        # Flush remaining tail
   884	        if previous_chunk_tail:
   885	            total_audio_chunks += 1
   886	            tail_samples = len(previous_chunk_tail) // 2
   887	            total_samples_emitted_to_user += tail_samples
   888	            yield previous_chunk_tail
   889
   890	        t_end = time.time()
   891	        audio_duration_s = total_samples_emitted_to_user / 16000
   892	        print(f"🎵 Continuation chunk complete: {audio_duration_s:.2f}s audio, {total_audio_chunks} chunks (using cached globals)")


nl -ba /home/ubuntu/veenaModal/veena3modal/services/tts_runtime.py | sed -n '1,260p'


Output:
Chunk ID: e7e295
Wall time: 0.0515 seconds
Process exited with code 0
Original token count: 2863
Output:
     1	"""
     2	TTS runtime wrapper (container-scoped singleton).
     3
     4	In Modal, each container should load the model once and then handle many requests.
     5	This wrapper owns:
     6	- SparkTTSModel (vLLM engine)
     7	- IndicPromptBuilder
     8	- BiCodecDecoder
     9	- SparkTTSPipeline + Veena3SlidingWindowPipeline
    10	- optional SuperResolutionService
    11
    12	Imports framework-agnostic inference + processing code from `veena3modal/core`,
    13	`veena3modal/processing`, and `veena3modal/audio`.
    14	"""
    15
    16	from __future__ import annotations
    17
    18	import os
    19	import sys
    20	import time
    21	import logging
    22	from threading import Lock
    23	from dataclasses import dataclass, field
    24	from typing import Any, Dict, List, Optional, Tuple
    25
    26	logger = logging.getLogger(__name__)
    27
    28	_PROFILE_REQUEST_TIMING = "VEENA3_TIMING_LOGS"
    29	_inflight_counter = 0
    30	_inflight_lock = Lock()
    31
    32
    33	def _timing_enabled() -> bool:
    34	    """Return True when detailed request timing should be emitted."""
    35	    return os.getenv(_PROFILE_REQUEST_TIMING, "").lower() in {"1", "true", "yes", "on"}
    36
    37
    38	def _with_inflight_delta(delta: int) -> int:
    39	    """Update in-flight request counter and return the new value."""
    40	    global _inflight_counter
    41	    with _inflight_lock:
    42	        _inflight_counter = max(0, _inflight_counter + delta)
    43	        return _inflight_counter
    44
    45
    46	@dataclass
    47	class TTSRuntime:
    48	    """
    49	    Holds long-lived, per-container inference objects.
    50
    51	    Thread Safety:
    52	        This class is designed for async concurrency within a single container.
    53	        The vLLM engine handles internal batching and scheduling.
    54	        Do NOT share across processes without proper synchronization.
    55	    """
    56	    model: Any = None
    57	    pipeline: Any = None
    58	    streaming_pipeline: Any = None
    59	    prompt_builder: Any = None
    60	    bicodec_decoder: Any = None
    61	    sr_service: Optional[Any] = None
    62
    63	    model_version: str = "not_loaded"
    64	    is_loaded: bool = False
    65	    load_time_ms: float = 0.0
    66
    67	    # Configuration
    68	    model_path: str = ""
    69	    bicodec_path: str = ""
    70	    sr_checkpoint_dir: Optional[str] = None
    71	    device: str = "cuda"
    72
    73	    # OPTIMIZATION: Pre-computed global tokens per speaker (eliminates ~110ms pre-roll in streaming)
    74	    # Maps speaker_name -> list of 32 global token IDs captured at startup
    75	    speaker_global_cache: Dict[str, List[int]] = field(default_factory=dict)
    76
    77
    78	# Module-level singleton (per-container)
    79	_runtime: Optional[TTSRuntime] = None
    80
    81
    82	def get_runtime() -> Optional[TTSRuntime]:
    83	    """Get the current TTS runtime singleton."""
    84	    return _runtime
    85
    86
    87	def is_initialized() -> bool:
    88	    """Check if the runtime is initialized and ready."""
    89	    return _runtime is not None and _runtime.is_loaded
    90
    91
    92	def initialize_runtime(
    93	    model_path: Optional[str] = None,
    94	    bicodec_path: Optional[str] = None,
    95	    sr_checkpoint_dir: Optional[str] = None,
    96	    device: str = "cuda",
    97	    hf_token: Optional[str] = None,
    98	    gpu_memory_utilization: float = 0.25,
    99	    enable_sr: bool = False,
   100	    num_engines: int = 1,
   101	) -> TTSRuntime:
   102	    """
   103	    Initialize the TTS runtime with all components.
   104
   105	    This should be called once per container (e.g., in Modal's @modal.enter).
   106
   107	    Args:
   108	        model_path: Path to Spark TTS model (env: SPARK_TTS_MODEL_PATH)
   109	        bicodec_path: Path to BiCodec model (env: BICODEC_MODEL_PATH, defaults to model_path)
   110	        sr_checkpoint_dir: Path to super-resolution checkpoints (env: AP_BWE_CHECKPOINT_DIR)
   111	        device: Device for inference (cuda/cpu)
   112	        hf_token: HuggingFace token for private models (env: HF_TOKEN)
   113	        gpu_memory_utilization: vLLM GPU memory fraction (default: 0.25)
   114	        enable_sr: Enable super-resolution service
   115	        num_engines: Number of vLLM engine instances (default: 1, set to 2-3 for Tier 3 optimization)
   116
   117	    Returns:
   118	        Initialized TTSRuntime instance
   119
   120	    Raises:
   121	        RuntimeError: If initialization fails
   122	    """
   123	    global _runtime
   124
   125	    start_time = time.time()
   126
   127	    # Resolve paths from env vars if not provided
   128	    model_path = model_path or os.environ.get(
   129	        'SPARK_TTS_MODEL_PATH',
   130	        os.environ.get('MODEL_PATH', '/models/spark_tts_4speaker')
   131	    )
   132	    bicodec_path = bicodec_path or os.environ.get(
   133	        'BICODEC_MODEL_PATH',
   134	        model_path  # BiCodec is usually in the same directory
   135	    )
   136	    sr_checkpoint_dir = sr_checkpoint_dir or os.environ.get('AP_BWE_CHECKPOINT_DIR')
   137	    hf_token = hf_token or os.environ.get('HF_TOKEN') or os.environ.get('HUGGING_FACE_HUB_TOKEN')
   138
   139	    logger.info(f"Initializing TTS runtime...")
   140	    logger.info(f"  Model path: {model_path}")
   141	    logger.info(f"  BiCodec path: {bicodec_path}")
   142	    logger.info(f"  SR enabled: {enable_sr}, path: {sr_checkpoint_dir}")
   143
   144	    try:
   145	        # Import inference components (vendored into veena3modal for Modal deployment)
   146	        from veena3modal.core.model_loader import SparkTTSModel
   147	        from veena3modal.core.pipeline import SparkTTSPipeline
   148	        from veena3modal.core.bicodec_decoder import BiCodecDecoder
   149	        from veena3modal.processing.prompt_builder import IndicPromptBuilder
   150
   151	        # Load SparkTTS model with vLLM engine
   152	        if num_engines > 1:
   153	            # TIER 3 OPTIMIZATION: Multiple vLLM engines on same GPU
   154	            # Each engine gets gpu_memory_utilization / num_engines to share GPU fairly.
   155	            # E.g., 3 engines at 0.08 each = 0.24 total, ~8GB per engine, ~24GB total.
   156	            from veena3modal.core.multi_engine import create_multi_engine_model
   157	            per_engine_mem = gpu_memory_utilization / num_engines
   158	            logger.info(f"Loading {num_engines} vLLM engines ({per_engine_mem:.2f} GPU mem each)...")
   159	            model = create_multi_engine_model(
   160	                model_path=model_path,
   161	                num_engines=num_engines,
   162	                hf_token=hf_token,
   163	                gpu_memory_per_engine=per_engine_mem,
   164	            )
   165	        else:
   166	            logger.info("Loading SparkTTS model with vLLM...")
   167	            model = SparkTTSModel(
   168	                model_path=model_path,
   169	                hf_token=hf_token,
   170	                gpu_memory_utilization=gpu_memory_utilization,
   171	            )
   172
   173	        # Initialize prompt builder
   174	        logger.info("Initializing prompt builder...")
   175	        prompt_builder = IndicPromptBuilder(
   176	            tokenizer=model.tokenizer,
   177	            model=model,
   178	        )
   179
   180	        # Initialize BiCodec decoder
   181	        logger.info("Initializing BiCodec decoder...")
   182	        bicodec_decoder = BiCodecDecoder(
   183	            device=device,
   184	            model_path=bicodec_path,
   185	        )
   186
   187	        # Initialize pipeline
   188	        logger.info("Initializing TTS pipeline...")
   189	        pipeline = SparkTTSPipeline(
   190	            model=model,
   191	            prompt_builder=prompt_builder,
   192	            bicodec_decoder=bicodec_decoder,
   193	        )
   194
   195	        # Initialize streaming pipeline (for M4)
   196	        streaming_pipeline = None
   197	        try:
   198	            from veena3modal.core.streaming_pipeline import Veena3SlidingWindowPipeline
   199	            logger.info("Initializing streaming pipeline...")
   200	            # NOTE: Parameter is named 'snac_decoder' for legacy reasons, but works with BiCodecDecoder
   201	            streaming_pipeline = Veena3SlidingWindowPipeline(
   202	                model=model,
   203	                prompt_builder=prompt_builder,
   204	                snac_decoder=bicodec_decoder,  # BiCodecDecoder is interface-compatible
   205	            )
   206	        except ImportError as e:
   207	            logger.warning(f"Streaming pipeline not available: {e}")
   208
   209	        # Initialize super-resolution (optional)
   210	        sr_service = None
   211	        if enable_sr and sr_checkpoint_dir:
   212	            try:
   213	                from veena3modal.core.super_resolution import SuperResolutionService
   214	                logger.info(f"Initializing super-resolution from {sr_checkpoint_dir}...")
   215	                sr_service = SuperResolutionService(checkpoint_dir=sr_checkpoint_dir)
   216	                # Load the model explicitly
   217	                if sr_service.load_model(device=device):
   218	                    logger.info("✅ Super-resolution model loaded successfully")
   219	                else:
   220	                    logger.warning("Super-resolution model failed to load")
   221	                    sr_service = None
   222	            except Exception as e:
   223	                logger.warning(f"Super-resolution not available: {e}")
   224
   225	        # Determine model version
   226	        model_version = os.path.basename(model_path.rstrip('/'))
   227	        if not model_version:
   228	            model_version = "spark-tts"
   229
   230	        load_time = (time.time() - start_time) * 1000
   231
   232	        # Create runtime
   233	        _runtime = TTSRuntime(
   234	            model=model,
   235	            pipeline=pipeline,
   236	            streaming_pipeline=streaming_pipeline,
   237	            prompt_builder=prompt_builder,
   238	            bicodec_decoder=bicodec_decoder,
   239	            sr_service=sr_service,
   240	            model_version=model_version,
   241	            is_loaded=True,
   242	            load_time_ms=load_time,
   243	            model_path=model_path,
   244	            bicodec_path=bicodec_path,
   245	            sr_checkpoint_dir=sr_checkpoint_dir,
   246	            device=device,
   247	        )
   248
   249	        logger.info(f"✅ TTS runtime initialized in {load_time:.0f}ms")
   250	        logger.info(f"   Model version: {model_version}")
   251
   252	        # OPTIMIZATION: Pre-compute global tokens for all speakers
   253	        # Disabled with multi-engine to avoid async event loop conflicts
   254	        if num_engines <= 1:
   255	            _precompute_speaker_globals(_runtime)
   256	        else:
   257	            logger.info("Speaker globals cache skipped (multi-engine mode)")
   258
   259	        # Update FastAPI app with model version
   260	        try:


nl -ba /home/ubuntu/veenaModal/veena3modal/services/tts_runtime.py | sed -n '260,520p'