Output:
Chunk ID: 0097a8
Wall time: 0.0515 seconds
Process exited with code 0
Original token count: 9202
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()

        if t_first_output is None:
            t_first_output = now
            prefill_ms = (t_first_output - t_gen_start) * 1000
            profile.stage("3a_vllm_prefill").record(prefill_ms)

        # Track per-iteration token rate
        current_count = len(request_output.outputs[0].token_ids)
        new_tokens = current_count - prev_token_count
        if new_tokens > 0 and iteration_count > 1:
            # Estimate instantaneous tokens/sec
            # (note: vLLM may batch multiple tokens per yield)
            dt = now - t_gen_start
            if dt > 0:
                profile.token_generation_rates.append(current_count / dt)
        prev_token_count = current_count
        final_output = request_output

    t_gen_end = time.perf_counter()
    total_gen_ms = (t_gen_end - t_gen_start) * 1000
    decode_only_ms = total_gen_ms - (profile.stages.get("3a_vllm_prefill", StageTimer("")).total_ms)
    profile.stage("3b_vllm_decode").record(decode_only_ms)
    profile.stage("3_vllm_total").record(total_gen_ms)
    profile.generated_tokens = len(final_output.outputs[0].token_ids) if final_output else 0

    if not final_output:
        profile.total_ms = (time.perf_counter() - t_total_start) * 1000
        return profile

    generated_text = final_output.outputs[0].text

    # === Stage 4: Token Extraction (regex) ===
    t0 = time.perf_counter()
    semantic_matches = re.findall(r"<\|bicodec_semantic_(\d+)\|>", generated_text)
    semantic_ids = [int(t) for t in semantic_matches]
    global_matches = re.findall(r"<\|bicodec_global_(\d+)\|>", generated_text)
    global_ids = [int(t) for t in global_matches]
    profile.stage("4_token_extraction").record((time.perf_counter() - t0) * 1000)
    profile.semantic_tokens = len(semantic_ids)
    profile.global_tokens = len(global_ids)

    if not semantic_ids or not global_ids:
        profile.total_ms = (time.perf_counter() - t_total_start) * 1000
        return profile

    # === Stage 5: Token Validation ===
    t0 = time.perf_counter()
    runtime.bicodec_decoder.validate_tokens(semantic_ids, global_ids)
    profile.stage("5_token_validation").record((time.perf_counter() - t0) * 1000)

    # === Stage 6: BiCodec Decode (GPU) ===
    t0 = time.perf_counter()
    audio_bytes = runtime.bicodec_decoder.decode_to_bytes(semantic_ids, global_ids)
    decode_ms = (time.perf_counter() - t0) * 1000
    profile.stage("6_bicodec_decode").record(decode_ms)
    profile.decode_latencies.append(decode_ms)

    if audio_bytes is None:
        profile.total_ms = (time.perf_counter() - t_total_start) * 1000
        return profile
#!/usr/bin/env python3
"""
Local Stress Test for Veena3 TTS Server.

Progressive load ramp with GPU monitoring, latency percentiles,
error categorization, and throughput measurement.

Unlike the Modal load tests, this runs locally with:
- No network latency (localhost)
- Direct GPU monitoring (nvidia-smi + torch.cuda)
- Higher concurrency ceiling (single-machine, no rate limits)
- Real streaming TTFB measurement (byte-level timing)

Usage:
    # Full suite (ramps 1 → 5 → 10 → 20 → 50 → 100 concurrent)
    python scripts/stress_test_local.py

    # Quick smoke test (1 → 5 → 10 only)
    python scripts/stress_test_local.py --quick

    # Custom ramp
    python scripts/stress_test_local.py --levels 1,5,10,20

    # Streaming only
    python scripts/stress_test_local.py --stream-only

    # Non-streaming only
    python scripts/stress_test_local.py --no-stream

    # Custom endpoint
    python scripts/stress_test_local.py --url http://localhost:8080
"""

import argparse
import asyncio
import json
import os
import statistics
import subprocess
import sys
import time
import threading
from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any, Tuple

import httpx

# === Configuration ===

DEFAULT_URL = "http://localhost:8000"
GENERATE_PATH = "/v1/tts/generate"

# Test texts at different lengths to stress different parts of the pipeline
TEST_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 quality "
        "with moderate length input that should complete in reasonable time."
    ),
    "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. She picked up her pen "
        "and began to write. The words flowed like water from a broken dam."
    ),
}

SPEAKERS = ["lipakshi", "vardan", "reet", "Nandini", "krishna", "anika"]

# === GPU Monitoring ===

@dataclass
class GPUSnapshot:
    timestamp: float
    memory_used_mb: float
    memory_total_mb: float
    gpu_utilization_pct: float
    temperature_c: float

class GPUMonitor:
    """Background thread that polls nvidia-smi for GPU stats."""

    def __init__(self, interval_seconds: float = 1.0):
        self.interval = interval_seconds
        self.snapshots: List[GPUSnapshot] = []
        self._running = False
        self._thread: Optional[threading.Thread] = None

    def start(self):
        self._running = True
        self._thread = threading.Thread(target=self._poll_loop, daemon=True)
        self._thread.start()

    def stop(self):
        self._running = False
        if self._thread:
            self._thread.join(timeout=5)

    def _poll_loop(self):
        while self._running:
            try:
                result = subprocess.run(
                    [
                        "nvidia-smi",
                        "--query-gpu=memory.used,memory.total,utilization.gpu,temperature.gpu",
                        "--format=csv,noheader,nounits",
                    ],
                    capture_output=True, text=True, timeout=5,
                )
                if result.returncode == 0:
                    parts = result.stdout.strip().split(", ")
                    if len(parts) >= 4:
                        self.snapshots.append(GPUSnapshot(
                            timestamp=time.time(),
                            memory_used_mb=float(parts[0]),
                            memory_total_mb=float(parts[1]),
                            gpu_utilization_pct=float(parts[2]),
                            temperature_c=float(parts[3]),
                        ))
            except Exception:
                pass
            time.sleep(self.interval)

    def summary(self) -> Dict[str, Any]:
        if not self.snapshots:
            return {"error": "no GPU data collected"}

        mem_used = [s.memory_used_mb for s in self.snapshots]
        gpu_util = [s.gpu_utilization_pct for s in self.snapshots]
        temps = [s.temperature_c for s in self.snapshots]

        return {
            "samples": len(self.snapshots),
            "memory_mb": {
                "min": min(mem_used),
                "max": max(mem_used),
                "avg": statistics.mean(mem_used),
                "total": self.snapshots[0].memory_total_mb,
            },
            "gpu_util_pct": {
                "min": min(gpu_util),
                "max": max(gpu_util),
                "avg": statistics.mean(gpu_util),
            },
            "temperature_c": {
                "min": min(temps),
                "max": max(temps),
                "avg": statistics.mean(temps),
            },
        }

# === Request Result ===

@dataclass
class RequestResult:
    success: bool
    status_code: int
    latency_ms: float  # Total wall-clock time
    ttfb_ms: float = 0.0  # Time to first byte (from HTTP response start)
    server_ttfb_ms: float = 0.0  # Server-reported TTFB (X-TTFB-ms header)
    audio_bytes: int = 0
    audio_seconds: float = 0.0
    error: Optional[str] = None
    speaker: str = ""
    text_length: int = 0
    stream: bool = False

# === Load Level Result ===

@dataclass
class LevelResult:
    level_name: str
    concurrency: int
    total_requests: int
    mode: str  # "non-streaming" or "streaming"
    text_category: str  # "short", "medium", "long", "mixed"
    results: List[RequestResult] = field(default_factory=list)
    wall_time_seconds: float = 0.0
    gpu_summary: Dict[str, Any] = field(default_factory=dict)

    @property
    def successes(self) -> List[RequestResult]:
        return [r for r in self.results if r.success]

    @property
    def failures(self) -> List[RequestResult]:
        return [r for r in self.results if not r.success]

    @property
    def success_rate(self) -> float:
        return len(self.successes) / len(self.results) if self.results else 0

    @property
    def throughput_rps(self) -> float:
        return self.total_requests / self.wall_time_seconds if self.wall_time_seconds > 0 else 0

    def _percentile(self, values: List[float], pct: float) -> float:
        if not values:
            return 0.0
        s = sorted(values)
        idx = int(len(s) * pct / 100)
        return s[min(idx, len(s) - 1)]

    @property
    def latency_stats(self) -> Dict[str, float]:
        lats = [r.latency_ms for r in self.successes]
        if not lats:
            return {}
        return {
            "min": min(lats),
            "avg": statistics.mean(lats),
            "p50": self._percentile(lats, 50),
            "p95": self._percentile(lats, 95),
            "p99": self._percentile(lats, 99),
            "max": max(lats),
        }

    @property
    def ttfb_stats(self) -> Dict[str, float]:
        ttfbs = [r.ttfb_ms for r in self.successes if r.ttfb_ms > 0]
        if not ttfbs:
            return {}
        return {
            "min": min(ttfbs),
            "avg": statistics.mean(ttfbs),
            "p50": self._percentile(ttfbs, 50),
            "p95": self._percentile(ttfbs, 95),
            "max": max(ttfbs),
        }

    @property
    def audio_duration_total(self) -> float:
        return sum(r.audio_seconds for r in self.successes)

    @property
    def effective_rtf(self) -> float:
        """Wall time / total audio produced. Lower = better."""
        total_audio = self.audio_duration_total
        return self.wall_time_seconds / total_audio if total_audio > 0 else 0

    def error_summary(self) -> Dict[str, int]:
        errors: Dict[str, int] = {}
        for r in self.failures:
            key = r.error or f"HTTP {r.status_code}"
            errors[key[:80]] = errors.get(key[:80], 0) + 1
        return errors

# === HTTP Client ===

async def make_request(
    client: httpx.AsyncClient,
    base_url: str,
    text: str,
    speaker: str,
    stream: bool = False,
    timeout: float = 120.0,
) -> RequestResult:
    """Make a single TTS request with timing."""
    url = f"{base_url}{GENERATE_PATH}"
    payload = {
        "text": text,
        "speaker": speaker,
        "stream": stream,
        "format": "wav",
    }

    start = time.time()
    ttfb_time = 0.0

    try:
        if stream:
            # Streaming: measure true TTFB (time to first byte of audio)
            async with client.stream(
                "POST", url,
                json=payload,
                headers={"Content-Type": "application/json"},
                timeout=timeout,
            ) as response:
                first_byte = True
                total_bytes = 0
                async for chunk in response.aiter_bytes():
                    if first_byte:
                        ttfb_time = (time.time() - start) * 1000
                        first_byte = False
                    total_bytes += len(chunk)

                latency = (time.time() - start) * 1000

                if response.status_code == 200:
                    # Estimate audio duration from bytes (16kHz, 16-bit, mono, minus WAV header)
                    pcm_bytes = max(0, total_bytes - 44)
                    audio_seconds = pcm_bytes / (16000 * 2)

                    return RequestResult(
                        success=True,
                        status_code=200,
                        latency_ms=latency,
                        ttfb_ms=ttfb_time,
                        audio_bytes=total_bytes,
                        audio_seconds=audio_seconds,
                        speaker=speaker,
                        text_length=len(text),
                        stream=True,
                    )
                else:
                    return RequestResult(
                        success=False,
                        status_code=response.status_code,
                        latency_ms=latency,
                        ttfb_ms=ttfb_time,
                        error=f"HTTP {response.status_code}",
                        speaker=speaker,
#!/usr/bin/env python3
"""
Production Stress Test - Real-world sentences for Veena3 TTS.

Uses actual production-quality sentences across:
- Multiple languages (English, Hindi, Telugu)
- Various lengths (short greetings to long paragraphs)
- Emotion tags
- Mixed scripts
- Numbers, dates, special characters

Runs the same concurrency ramp as the baseline test for direct comparison.
"""

import argparse
import asyncio
import json
import os
import statistics
import subprocess
import sys
import time
import threading
from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any

import httpx

DEFAULT_URL = "http://localhost:8000"
GENERATE_PATH = "/v1/tts/generate"

# === Real production sentences ===
# These mimic actual API traffic: mixed languages, emotions, varied lengths

REAL_SENTENCES = {
    "en_short": [
        "Welcome to Veena, your personal voice assistant.",
        "The weather today is sunny with a high of thirty two degrees.",
        "Your order has been confirmed and will arrive by tomorrow.",
        "Please hold on, I am transferring you to a specialist.",
        "Thank you for calling. Have a wonderful day!",
    ],
    "en_medium": [
        "Good morning! I hope you are having a great start to your day. Let me walk you through the key highlights from yesterday's meeting.",
        "The quarterly revenue report shows a fifteen percent increase compared to last year, driven primarily by growth in the South Asian market.",
        "I would like to remind you that your subscription is expiring on March fifteenth. Would you like me to help you with the renewal process?",
        "According to the latest research, artificial intelligence is expected to contribute over fifteen trillion dollars to the global economy by twenty thirty.",
        "[excited] Congratulations! You have been selected for our premium membership program. Let me tell you about the amazing benefits.",
    ],
    "en_long": [
        "In the heart of Hyderabad, where the old city meets the new, there is a small bookshop that has been running for over fifty years. The owner, Ramesh uncle, knows every book by heart. He can tell you exactly which shelf to find your favorite author on, and he always has a recommendation ready. Last Tuesday, a young girl walked in looking for a book on astronomy. Ramesh uncle smiled and led her to the science section.",
        "The Indian space program has achieved remarkable milestones in the past decade. From the Mars Orbiter Mission in twenty fourteen to the Chandrayaan three lunar landing in twenty twenty three, ISRO has consistently demonstrated that world class space exploration does not require a massive budget. The upcoming Gaganyaan mission will mark India's first crewed spaceflight.",
    ],
    "hi_short": [
        "नमस्ते! आज का दिन बहुत अच्छा है।",
        "कृपया अपना नाम और फोन नंबर बताइए।",
        "आपका ऑर्डर सफलतापूर्वक प्लेस हो गया है।",
        "धन्यवाद! आपकी कॉल हमारे लिए महत्वपूर्ण है।",
    ],
    "hi_medium": [
        "[curious] क्या आप जानते हैं कि भारत दुनिया का सबसे बड़ा लोकतंत्र है? यहाँ एक सौ चालीस करोड़ से अधिक लोग रहते हैं।",
        "आज के मौसम का हाल बताते हैं। दिल्ली में तापमान बत्तीस डिग्री रहेगा और हल्की बारिश की संभावना है।",
    ],
    "te_short": [
        "నమస్కారం! మీరు ఎలా ఉన్నారు?",
        "మీ ఆర్డర్ విజయవంతంగా ప్లేస్ చేయబడింది।",
        "దయచేసి కొంచెం ఆగండి, మేము మీకు సహాయం చేస్తాము.",
    ],
    "emotion_mixed": [
        "[laughs] Oh that is so funny! I cannot believe you said that.",
        "[whispers] Let me tell you a secret. The new product launch is happening next week.",
        "[excited] Amazing news everyone! We just crossed one million users!",
        "[sighs] Unfortunately, the flight has been delayed by two hours. I apologize for the inconvenience.",
        "[angry] This is unacceptable! The service was supposed to be restored by noon today.",
    ],
}

# Flatten and tag for round-robin selection
ALL_SENTENCES = []
for category, sentences in REAL_SENTENCES.items():
    for s in sentences:
        ALL_SENTENCES.append({"text": s, "category": category, "length": len(s)})

SPEAKERS = ["lipakshi", "vardan", "reet", "Nandini", "krishna", "anika",
            "adarsh", "Nilay", "Aarvi", "Asha", "Bittu", "Mira"]

# === GPU Monitor (reused from stress_test_local.py) ===

@dataclass
class GPUSnapshot:
    timestamp: float
    memory_used_mb: float
    memory_total_mb: float
    gpu_utilization_pct: float
    temperature_c: float

class GPUMonitor:
    def __init__(self, interval_seconds: float = 1.0):
        self.interval = interval_seconds
        self.snapshots: List[GPUSnapshot] = []
        self._running = False
        self._thread = None

    def start(self):
        self._running = True
        self._thread = threading.Thread(target=self._poll, daemon=True)
        self._thread.start()

    def stop(self):
        self._running = False
        if self._thread:
            self._thread.join(timeout=5)

    def _poll(self):
        while self._running:
            try:
                r = subprocess.run(
                    ["nvidia-smi", "--query-gpu=memory.used,memory.total,utilization.gpu,temperature.gpu",
                     "--format=csv,noheader,nounits"],
                    capture_output=True, text=True, timeout=5)
                if r.returncode == 0:
                    parts = r.stdout.strip().split(", ")
                    if len(parts) >= 4:
                        self.snapshots.append(GPUSnapshot(
                            timestamp=time.time(),
                            memory_used_mb=float(parts[0]),
                            memory_total_mb=float(parts[1]),
                            gpu_utilization_pct=float(parts[2]),
                            temperature_c=float(parts[3]),
                        ))
            except Exception:
                pass
            time.sleep(self.interval)

    def summary(self):
        if not self.snapshots:
            return {}
        mem = [s.memory_used_mb for s in self.snapshots]
        util = [s.gpu_utilization_pct for s in self.snapshots]
        return {
            "memory_mb_avg": statistics.mean(mem),
            "memory_mb_peak": max(mem),
            "memory_mb_total": self.snapshots[0].memory_total_mb,
            "gpu_util_avg": statistics.mean(util),
            "gpu_util_peak": max(util),
            "temp_avg": statistics.mean(s.temperature_c for s in self.snapshots),
            "temp_peak": max(s.temperature_c for s in self.snapshots),
        }

# === Request Result ===

@dataclass
class RequestResult:
    success: bool
    status_code: int
    latency_ms: float
    ttfb_ms: float = 0.0
    audio_bytes: int = 0
    audio_seconds: float = 0.0
    error: Optional[str] = None
    text_length: int = 0
    category: str = ""
    stream: bool = False

# === HTTP Client ===

async def make_request(client, base_url, text, speaker, stream=False, timeout=120.0):
    url = f"{base_url}{GENERATE_PATH}"
    payload = {"text": text, "speaker": speaker, "stream": stream, "format": "wav"}
    start = time.time()

    try:
        if stream:
            async with client.stream("POST", url, json=payload,
                                     headers={"Content-Type": "application/json"},
                                     timeout=timeout) as resp:
                first_byte = True
                ttfb = 0.0
                total = 0
                async for chunk in resp.aiter_bytes():
                    if first_byte:
                        ttfb = (time.time() - start) * 1000
                        first_byte = False
                    total += len(chunk)
                latency = (time.time() - start) * 1000
                if resp.status_code == 200:
                    pcm = max(0, total - 44)
                    return RequestResult(True, 200, latency, ttfb, total,
                                        pcm / (16000 * 2), text_length=len(text), stream=True)
                return RequestResult(False, resp.status_code, latency, ttfb,
                                     error=f"HTTP {resp.status_code}", text_length=len(text), stream=True)
        else:
            resp = await client.post(url, json=payload,
                                     headers={"Content-Type": "application/json"},
                                     timeout=timeout)
            latency = (time.time() - start) * 1000
            audio_sec = 0.0
            h = resp.headers.get("X-Audio-Seconds")
            if h:
                audio_sec = float(h)
            elif resp.status_code == 200:
                audio_sec = max(0, len(resp.content) - 44) / (16000 * 2)
            if resp.status_code == 200:
                return RequestResult(True, 200, latency, latency, len(resp.content),
                                     audio_sec, text_length=len(text), stream=False)
            return RequestResult(False, resp.status_code, latency,
                                 error=resp.text[:200], text_length=len(text), stream=False)
    except httpx.TimeoutException:
        return RequestResult(False, 0, (time.time() - start) * 1000, error="TIMEOUT",
                             text_length=len(text), stream=stream)
    except Exception as e:
        return RequestResult(False, 0, (time.time() - start) * 1000, error=str(e)[:200],
                             text_length=len(text), stream=stream)

def percentile(values, pct):
    if not values:
        return 0
    s = sorted(values)
    return s[min(int(len(s) * pct / 100), len(s) - 1)]

async def run_level(base_url, num_requests, concurrency, stream, gpu_monitor=None):
    sem = asyncio.Semaphore(concurrency)
    results = []

    async def req(i, client):
        async with sem:
            item = ALL_SENTENCES[i % len(ALL_SENTENCES)]
            speaker = SPEAKERS[i % len(SPEAKERS)]
            r = await make_request(client, base_url, item["text"], speaker, stream)
            r.category = item["category"]
            return r

    if gpu_monitor:
        gpu_monitor.snapshots.clear()

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

    for r in raw:
        if isinstance(r, Exception):
            results.append(RequestResult(False, 0, 0, error=str(r)[:200], stream=stream))
        else:
            results.append(r)

    successes = [r for r in results if r.success]
    failures = [r for r in results if not r.success]
    lats = [r.latency_ms for r in successes]
    ttfbs = [r.ttfb_ms for r in successes if r.ttfb_ms > 0]
    audio_total = sum(r.audio_seconds for r in successes)
    rps = num_requests / wall if wall > 0 else 0
    rtf = wall / audio_total if audio_total > 0 else 0

    gpu = gpu_monitor.summary() if gpu_monitor else {}
    mode = "streaming" if stream else "non-streaming"
    label = f"{concurrency}c-production-{mode}"

    print(f"\n{'=' * 72}")
    print(f"  {label.upper()}")
    print(f"  {num_requests} requests, {concurrency} concurrent, {mode}")
    print(f"  Sentences: real production ({len(ALL_SENTENCES)} unique, multilingual)")
    print(f"{'=' * 72}")
    print(f"  Success:     {len(successes)}/{num_requests} ({len(successes)/num_requests:.0%})")
    print(f"  Failures:    {len(failures)}")
    print(f"  Wall time:   {wall:.2f}s")
    print(f"  Throughput:  {rps:.2f} req/s")
    print(f"  Audio total: {audio_total:.1f}s")
    print(f"  Eff. RTF:    {rtf:.3f}")

    if lats:
        print(f"\n  Latency (ms):")
        print(f"    min={min(lats):.0f}  avg={statistics.mean(lats):.0f}  "
              f"p50={percentile(lats,50):.0f}  p95={percentile(lats,95):.0f}  "
              f"p99={percentile(lats,99):.0f}  max={max(lats):.0f}")

    if ttfbs and stream:
        print(f"  TTFB (ms):")
        print(f"    min={min(ttfbs):.0f}  avg={statistics.mean(ttfbs):.0f}  "
              f"p50={percentile(ttfbs,50):.0f}  p95={percentile(ttfbs,95):.0f}  max={max(ttfbs):.0f}")

    if gpu:
        print(f"\n  GPU:")
        print(f"    Memory: {gpu['memory_mb_avg']:.0f}MB avg, {gpu['memory_mb_peak']:.0f}MB peak / {gpu['memory_mb_total']:.0f}MB")
        print(f"    Util:   {gpu['gpu_util_avg']:.0f}% avg, {gpu['gpu_util_peak']:.0f}% peak")
        print(f"    Temp:   {gpu['temp_avg']:.0f}C avg, {gpu['temp_peak']:.0f}C peak")

    errors = {}
    for r in failures:
        k = (r.error or f"HTTP {r.status_code}")[:60]
        errors[k] = errors.get(k, 0) + 1
    if errors:
        print(f"\n  Errors:")
        for e, c in sorted(errors.items(), key=lambda x: -x[1]):
            print(f"    [{c}x] {e}")

    return {
        "level": label, "concurrency": concurrency, "requests": num_requests,
        "mode": mode, "success_rate": len(successes) / num_requests if num_requests else 0,
        "throughput_rps": rps, "wall_time_s": wall, "eff_rtf": rtf,
        "audio_total_s": audio_total,
        "latency_ms": {"avg": statistics.mean(lats), "p50": percentile(lats, 50),
                        "p95": percentile(lats, 95), "p99": percentile(lats, 99),
                        "max": max(lats)} if lats else {},
        "ttfb_ms": {"avg": statistics.mean(ttfbs), "p50": percentile(ttfbs, 50),
                     "p95": percentile(ttfbs, 95), "max": max(ttfbs)} if ttfbs else {},
        "gpu": gpu, "errors": errors,
    }

async def main_async(args):
    base_url = args.url
    print(f"Checking server at {base_url}...")
    try:
        async with httpx.AsyncClient() as c:
            r = await c.get(f"{base_url}/v1/tts/health", timeout=10)
            h = r.json()
            print(f"  Status: {h.get('status')}, Model: {h.get('model_version')}, GPU: {h.get('gpu_available')}")
    except Exception as e:
        print(f"ERROR: {e}")
        sys.exit(1)

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

    all_results = []

    # Warmup with 3 real sentences
    print(f"\nWarming up (3 requests with real sentences)...")
    await run_level(base_url, 3, 1, stream=False, gpu_monitor=None)


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


cd /home/ubuntu/veenaModal && sed -n '1,360p' scripts/setup_local.sh && sed -n '1,340p' veena3modal/local_server.py && sed -n '1,260p' veena3modal/__main__.py


Output:
Chunk ID: fc33ee
Wall time: 0.0503 seconds
Process exited with code 0
Original token count: 5063
Output:
#!/usr/bin/env bash
# ============================================================================
# Veena3 TTS Local Setup
#
# Sets up everything needed to run the TTS server locally:
# 1. Creates/activates Python venv
# 2. Installs dependencies
# 3. Downloads model weights from HuggingFace
# 4. Validates GPU + CUDA availability
# 5. Runs a quick smoke test
#
# Usage:
#   bash scripts/setup_local.sh              # Full setup
#   bash scripts/setup_local.sh --skip-model # Skip model download
#   bash scripts/setup_local.sh --gpu-check  # Only check GPU
# ============================================================================

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
VENV_DIR="${REPO_ROOT}/venv"
MODEL_DIR="${REPO_ROOT}/models/spark_tts_4speaker"
HF_REPO="BayAreaBoys/spark_tts_4speaker"
PYTHON_VERSION_REQUIRED="3.10"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

log_info()  { echo -e "${BLUE}[INFO]${NC}  $*"; }
log_ok()    { echo -e "${GREEN}[OK]${NC}    $*"; }
log_warn()  { echo -e "${YELLOW}[WARN]${NC}  $*"; }
log_err()   { echo -e "${RED}[ERROR]${NC} $*"; }

# Parse flags
SKIP_MODEL=false
GPU_CHECK_ONLY=false
for arg in "$@"; do
    case "$arg" in
        --skip-model)  SKIP_MODEL=true ;;
        --gpu-check)   GPU_CHECK_ONLY=true ;;
    esac
done

# ============================================================================
# 1. GPU Check
# ============================================================================
echo ""
echo "============================================================"
echo "  Veena3 TTS Local Setup"
echo "============================================================"

log_info "Checking GPU..."
if command -v nvidia-smi &>/dev/null; then
    GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1)
    GPU_MEM=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader 2>/dev/null | head -1)
    GPU_DRIVER=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -1)
    CUDA_VER=$(nvidia-smi | grep "CUDA Version" | awk '{print $9}')
    log_ok "GPU: ${GPU_NAME} (${GPU_MEM})"
    log_ok "Driver: ${GPU_DRIVER}, CUDA: ${CUDA_VER}"
else
    log_warn "nvidia-smi not found. GPU inference will not work."
    log_warn "CPU-only mode available with --device cpu (very slow)"
fi

if [ "$GPU_CHECK_ONLY" = true ]; then
    exit 0
fi

# ============================================================================
# 2. Python Virtual Environment
# ============================================================================
log_info "Setting up Python virtual environment..."

if [ ! -d "$VENV_DIR" ]; then
    log_info "Creating venv at ${VENV_DIR}..."
    python3 -m venv "$VENV_DIR"
    log_ok "Virtual environment created"
else
    log_ok "Virtual environment already exists at ${VENV_DIR}"
fi

# Activate venv
source "${VENV_DIR}/bin/activate"
log_ok "Activated venv: $(which python)"
log_ok "Python version: $(python --version)"

# Upgrade pip
pip install --quiet --upgrade pip wheel setuptools

# ============================================================================
# 3. Install Dependencies
# ============================================================================
log_info "Installing dependencies..."

# Core requirements
pip install --quiet -r "${REPO_ROOT}/requirements.txt"

# FastAPI + uvicorn (not in requirements.txt since Modal bundles them)
pip install --quiet \
    "fastapi[standard]>=0.100.0" \
    "uvicorn[standard]>=0.23.0" \
    "prometheus_client>=0.17.0" \
    "supabase>=2.0.0"

# HuggingFace Hub for model download
pip install --quiet "huggingface-hub[cli]>=0.24.0"

# SparkTTS extra deps
pip install --quiet \
    "omegaconf>=2.3.0" \
    "safetensors>=0.5.0" \
    "soxr>=0.5.0"

# AP-BWE deps
pip install --quiet \
    "matplotlib>=3.8.0" \
    "natsort>=8.0.0" \
    "joblib>=1.0.0"

# Einops (vLLM + model deps)
pip install --quiet \
    "einops>=0.6.0" \
    "einx>=0.2.0"

log_ok "All dependencies installed"

# ============================================================================
# 4. Validate Key Imports
# ============================================================================
log_info "Validating imports..."

python -c "
import torch
print(f'  torch: {torch.__version__}, CUDA: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'  GPU: {torch.cuda.get_device_name(0)}')
    print(f'  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')
" || { log_err "PyTorch import failed"; exit 1; }

python -c "import vllm; print(f'  vllm: {vllm.__version__}')" \
    || { log_err "vLLM import failed"; exit 1; }

python -c "import fastapi; print(f'  fastapi: {fastapi.__version__}')" \
    || { log_err "FastAPI import failed"; exit 1; }

python -c "
import sys
sys.path.insert(0, '${REPO_ROOT}/external/sparktts')
import sparktts
print(f'  sparktts: available')
" || { log_warn "sparktts import failed (may need model files)"; }

log_ok "Core imports validated"

# ============================================================================
# 5. Download Model
# ============================================================================
if [ "$SKIP_MODEL" = false ]; then
    log_info "Checking model at ${MODEL_DIR}..."

    if [ -f "${MODEL_DIR}/config.json" ]; then
        log_ok "Model already present"
    else
        log_info "Downloading model from ${HF_REPO}..."
        log_info "This may take a few minutes on first run..."

        # Load HF token from .env if present
        if [ -f "${REPO_ROOT}/.env" ]; then
            export $(grep -v '^#' "${REPO_ROOT}/.env" | xargs)
        fi

        mkdir -p "${MODEL_DIR}"

        python -c "
from huggingface_hub import snapshot_download
snapshot_download(
    repo_id='${HF_REPO}',
    local_dir='${MODEL_DIR}',
    token='${HF_TOKEN:-}' or None,
)
print('Model download complete')
"
        if [ -f "${MODEL_DIR}/config.json" ]; then
            log_ok "Model downloaded successfully"
        else
            log_err "Model download may have failed - config.json not found"
            exit 1
        fi
    fi
else
    log_warn "Skipping model download (--skip-model)"
fi

# ============================================================================
# 6. Check ffmpeg (needed for audio encoding: opus, mp3, flac, mulaw)
# ============================================================================
log_info "Checking ffmpeg..."
if command -v ffmpeg &>/dev/null; then
    FFMPEG_VER=$(ffmpeg -version 2>&1 | head -1 | awk '{print $3}')
    log_ok "ffmpeg: ${FFMPEG_VER}"
else
    log_warn "ffmpeg not found. Install with: sudo apt install ffmpeg"
    log_warn "Audio format encoding (opus, mp3, flac, mulaw) will not work"
fi

# ============================================================================
# 7. Summary
# ============================================================================
echo ""
echo "============================================================"
echo "  Setup Complete!"
echo "============================================================"
echo ""
echo "  To activate the environment:"
echo "    source venv/bin/activate"
echo ""
echo "  To start the local TTS server:"
echo "    python -m veena3modal.local_server"
echo ""
echo "  To start with custom settings:"
echo "    python -m veena3modal.local_server --port 8080 --gpu-memory 0.5"
echo ""
echo "  Quick test after server starts:"
echo "    curl http://localhost:8000/v1/tts/health"
echo ""
echo "    curl -X POST http://localhost:8000/v1/tts/generate \\"
echo "      -H 'Content-Type: application/json' \\"
echo "      -d '{\"text\": \"Hello, this is a test.\", \"speaker\": \"Mitra\"}' \\"
echo "      --output test.wav"
echo ""
echo "============================================================"
"""
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,
        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(
"""Allow running as: python -m veena3modal (shortcut for local_server)."""
from veena3modal.local_server import main

if __name__ == "__main__":
    main()


Summary: auto


cd /home/ubuntu/veenaModal && rg -n "VEENA3_TIMING_LOGS|generate_speech|generate_speech_streaming|profile|stress|local" veena3modal/local_server.py && sed -n '340,760p' veena3modal/local_server.py


Output:
Chunk ID: d2b8bd
Wall time: 0.0513 seconds
Process exited with code 0
Original token count: 1596
Output:
12:    python -m veena3modal.local_server
15:    python -m veena3modal.local_server --model-path /path/to/spark_tts_4speaker
18:    python -m veena3modal.local_server --gpu-memory 0.5 --port 8080
21:    python -m veena3modal.local_server --enable-sr --sr-path /path/to/ap_bwe/16kto48k
24:    python -m veena3modal.local_server --device cpu
30:    AUTH_BYPASS_MODE        - "true" to disable auth (default for local)
64:logger = logging.getLogger("veena3.local")
66:# Default local model directory (relative to repo)
79:    local_dir: str = DEFAULT_LOCAL_MODEL_DIR,
90:    Returns the resolved local path.
94:    if any(os.path.exists(os.path.join(local_dir, m)) for m in marker_files):
95:        logger.info(f"Model already present at {local_dir}")
96:        return local_dir
107:    os.makedirs(local_dir, exist_ok=True)
114:            local_dir=local_dir,
129:                local_dir=local_dir,
147:        f"     huggingface-cli download {PRIVATE_HF_REPO} --local-dir {local_dir}\n\n"
149:        f"     huggingface-cli download {PUBLIC_FALLBACK_REPO} --local-dir {local_dir}\n\n"
151:        "     python -m veena3modal.local_server --model-path /path/to/model"
192:def initialize_local_runtime(
202:    Initialize TTS runtime for local serving.
214:    logger.info("Initializing Veena3 TTS Runtime (local)")
244:def create_local_app():
246:    Create the FastAPI app for local serving.
248:    Same app factory as Modal deployment, but called from local context.
255:    """Parse CLI arguments for local server configuration."""
261:  python -m veena3modal.local_server                          # Auto-download model, serve on :8000
262:  python -m veena3modal.local_server --port 8080              # Custom port
263:  python -m veena3modal.local_server --model-path ./my_model  # Use local model
264:  python -m veena3modal.local_server --gpu-memory 0.5         # Limit GPU memory
265:  python -m veena3modal.local_server --workers 2              # Multiple workers (careful: each loads model)
365:        help="Enable API key authentication (default: disabled for local)",
383:    Main entry point for local TTS server.
401:    # Auth bypass for local development (unless --auth flag is set)
404:        logger.info("Auth bypass enabled (local mode). Use --auth to require API keys.")
421:            local_dir=model_path,
426:    initialize_local_runtime(
452:        "veena3modal.local_server:create_local_app",
    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()


cd /home/ubuntu/veenaModal && sed -n '320,760p' scripts/profile_pipeline.py


Output:
Chunk ID: 346078
Wall time: 0.0507 seconds
Process exited with code 0
Original token count: 3870
Output:
        return profile

    # === Stage 7: WAV Header ===
    t0 = time.perf_counter()
    wav_bytes = add_wav_header(audio_bytes, sample_rate=AUDIO_SAMPLE_RATE)
    profile.stage("7_wav_header").record((time.perf_counter() - t0) * 1000)

    profile.audio_bytes = len(wav_bytes)
    profile.audio_seconds = (len(audio_bytes)) / (AUDIO_SAMPLE_RATE * 2)  # int16 = 2 bytes/sample
    profile.total_ms = (time.perf_counter() - t_total_start) * 1000

    return profile

# === Streaming Profiler ===

async def profile_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 streaming request with per-token and per-decode timing."""
    from vllm import SamplingParams
    from veena3modal.core.constants import CODE_END_TOKEN_ID, DEFAULT_MIN_TOKENS
    from veena3modal.core.token_utils import BiCodecTokenParser
    from veena3modal.audio.crossfade import crossfade_bytes_int16

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

    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)

    # === Stage 2: Parser Init ===
    t0 = time.perf_counter()
    tokenizer = runtime.model.tokenizer
    token_parser = BiCodecTokenParser(tokenizer)
    profile.stage("2_parser_init").record((time.perf_counter() - t0) * 1000)

    # === Stage 3: Sampling Config ===
    t0 = time.perf_counter()
    sampling_params = SamplingParams(
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        max_tokens=max_tokens,
        min_tokens=DEFAULT_MIN_TOKENS,
        repetition_penalty=repetition_penalty,
        stop_token_ids=[CODE_END_TOKEN_ID],
        seed=seed,
    )
    profile.stage("3_sampling_config").record((time.perf_counter() - t0) * 1000)

    # Buffers
    semantic_buffer: List[int] = []
    global_buffer: List[int] = []
    processed_token_count = 0

    EXPECTED_GLOBAL_COUNT = 32
    MIN_SEMANTIC_FIRST = 16
    DECODE_INTERVAL = 24

    request_id = f"prof-stream-{uuid.uuid4().hex[:8]}"

    # === Stage 4: vLLM Generation (streaming) ===
    t_gen_start = time.perf_counter()
    t_first_output = None
    t_globals_complete = None
    t_first_decode = None
    last_decode_count = 0
    total_samples_emitted = 0
    previous_chunk_tail: Optional[bytes] = None
    iteration_count = 0

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

    # Per-token timing: time between consecutive yields from vLLM
    t_last_yield = t_gen_start
    inter_token_times: List[float] = []

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

        # Track inter-yield timing (approximation of per-token time)
        inter_ms = (now - t_last_yield) * 1000
        inter_token_times.append(inter_ms)
        t_last_yield = now

        if t_first_output is None:
            t_first_output = now
            prefill_ms = (t_first_output - t_gen_start) * 1000
            profile.stage("4a_vllm_prefill").record(prefill_ms)

        generated_ids = request_output.outputs[0].token_ids

        # === Token Parsing ===
        t_parse = time.perf_counter()
        new_token_ids = generated_ids[processed_token_count:]
        processed_token_count = len(generated_ids)
        token_parser.parse_incremental(new_token_ids, semantic_buffer, global_buffer)
        profile.stage("4b_token_parsing").record((time.perf_counter() - t_parse) * 1000)

        # Track globals completion
        if t_globals_complete is None and len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
            t_globals_complete = now
            globals_ms = (t_globals_complete - t_gen_start) * 1000
            profile.stage("4c_global_token_gen").record(globals_ms)

        # === BiCodec Decode (when enough tokens) ===
        if len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
            sem_count = len(semantic_buffer)
            if sem_count >= MIN_SEMANTIC_FIRST and sem_count - last_decode_count >= DECODE_INTERVAL:
                last_decode_count = sem_count

                t_decode = time.perf_counter()
                audio_bytes = runtime.bicodec_decoder.decode_streaming(
                    semantic_ids=semantic_buffer,
                    global_ids=global_buffer[:EXPECTED_GLOBAL_COUNT],
                    use_sliding_window=False,
                    trim_warmup=False,
                )
                decode_ms = (time.perf_counter() - t_decode) * 1000
                profile.stage("5_bicodec_decode").record(decode_ms)
                profile.decode_latencies.append(decode_ms)

                if audio_bytes:
                    total_decoded_samples = len(audio_bytes) // 2
                    if total_decoded_samples > total_samples_emitted:
                        new_bytes = audio_bytes[total_samples_emitted * 2:]

                        t_cf = time.perf_counter()
                        to_emit, previous_chunk_tail = crossfade_bytes_int16(
                            previous_chunk_tail, new_bytes,
                            sample_rate_hz=16000, crossfade_ms=50,
                        )
                        profile.stage("6_crossfade").record((time.perf_counter() - t_cf) * 1000)

                        emitted_samples = len(to_emit) // 2
                        total_samples_emitted += emitted_samples

                        if to_emit:
                            profile.chunks_emitted += 1
                            if t_first_decode is None:
                                t_first_decode = time.perf_counter()
                                profile.ttfb_ms = (t_first_decode - t_total_start) * 1000

    t_gen_end = time.perf_counter()
    profile.stage("4_vllm_total").record((t_gen_end - t_gen_start) * 1000)

    # Final decode
    if len(semantic_buffer) > last_decode_count and len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
        t_decode = time.perf_counter()
        audio_bytes = runtime.bicodec_decoder.decode_streaming(
            semantic_ids=semantic_buffer,
            global_ids=global_buffer[:EXPECTED_GLOBAL_COUNT],
            use_sliding_window=False, trim_warmup=False,
        )
        decode_ms = (time.perf_counter() - t_decode) * 1000
        profile.stage("5_bicodec_decode").record(decode_ms)
        profile.decode_latencies.append(decode_ms)

        if audio_bytes:
            total_decoded_samples = len(audio_bytes) // 2
            if total_decoded_samples > total_samples_emitted:
                new_bytes = audio_bytes[total_samples_emitted * 2:]
                to_emit, previous_chunk_tail = crossfade_bytes_int16(
                    previous_chunk_tail, new_bytes, sample_rate_hz=16000, crossfade_ms=50,
                )
                total_samples_emitted += len(to_emit) // 2
                if to_emit:
                    profile.chunks_emitted += 1

    if previous_chunk_tail:
        total_samples_emitted += len(previous_chunk_tail) // 2
        profile.chunks_emitted += 1

    profile.generated_tokens = processed_token_count
    profile.semantic_tokens = len(semantic_buffer)
    profile.global_tokens = len(global_buffer)
    profile.audio_seconds = total_samples_emitted / 16000
    profile.audio_bytes = total_samples_emitted * 2
    profile.total_ms = (time.perf_counter() - t_total_start) * 1000

    # Compute token generation rates from inter-token times
    # Skip first (prefill) and compute tok/s for decode phase
    if len(inter_token_times) > 2:
        decode_times = inter_token_times[1:]  # Skip prefill
        for dt in decode_times:
            if dt > 0:
                profile.token_generation_rates.append(1000.0 / dt)  # ms → tok/s

    # Record inter-token timing stats
    if inter_token_times:
        import statistics
        decode_times = inter_token_times[1:] if len(inter_token_times) > 1 else inter_token_times
        if decode_times:
            profile.stage("4d_inter_token_avg").record(statistics.mean(decode_times))
            profile.stage("4e_inter_token_p99").record(
                sorted(decode_times)[int(len(decode_times) * 0.99)]
            )

    return profile

# === Concurrency Profiler ===

async def profile_concurrent(
    runtime,
    num_requests: int,
    text: str,
    speaker: str = "lipakshi",
    stream: bool = False,
) -> List[RequestProfile]:
    """Profile multiple concurrent requests."""

    speakers = ["lipakshi", "vardan", "reet", "Nandini", "krishna", "anika"]

    async def run_one(i: int) -> RequestProfile:
        spk = speakers[i % len(speakers)]
        if stream:
            return await profile_streaming(runtime, text, spk)
        else:
            return await profile_non_streaming(runtime, text, spk)

    profiles = await asyncio.gather(*[run_one(i) for i in range(num_requests)])
    return list(profiles)

def print_concurrency_summary(profiles: List[RequestProfile], concurrency: int, stream: bool):
    """Print aggregated summary for concurrent profiles."""
    import statistics

    mode = "streaming" if stream else "non-streaming"
    totals = [p.total_ms for p in profiles]
    wall_time = max(totals)

    print(f"\n{'#' * 80}")
    print(f"  CONCURRENCY={concurrency} ({mode}), {len(profiles)} requests")
    print(f"{'#' * 80}")
    print(f"  Wall time: {wall_time:.0f}ms")
    print(f"  Throughput: {len(profiles) / (wall_time / 1000):.2f} req/s")
    print(f"  Latency: avg={statistics.mean(totals):.0f}ms  "
          f"p50={sorted(totals)[len(totals)//2]:.0f}ms  "
          f"p99={sorted(totals)[int(len(totals)*0.99)]:.0f}ms")

    if stream:
        ttfbs = [p.ttfb_ms for p in profiles if p.ttfb_ms > 0]
        if ttfbs:
            print(f"  TTFB: avg={statistics.mean(ttfbs):.0f}ms  "
                  f"p50={sorted(ttfbs)[len(ttfbs)//2]:.0f}ms  "
                  f"p95={sorted(ttfbs)[int(len(ttfbs)*0.95)]:.0f}ms")

    # Aggregate stage timing
    stage_totals: Dict[str, List[float]] = {}
    for p in profiles:
        for name, s in p.stages.items():
            if name not in stage_totals:
                stage_totals[name] = []
            stage_totals[name].append(s.total_ms)

    print(f"\n  STAGE BREAKDOWN (avg across {len(profiles)} requests):")
    print(f"  {'-' * 76}")
    sorted_stages = sorted(stage_totals.items(), key=lambda x: -statistics.mean(x[1]))
    for name, times in sorted_stages:
        avg = statistics.mean(times)
        mx = max(times)
        mn = min(times)
        pct = (avg / statistics.mean(totals)) * 100 if totals else 0
        print(f"  {name:<35} avg={avg:>8.2f}ms  min={mn:>8.2f}ms  max={mx:>8.2f}ms  ({pct:.1f}%)")

    # BiCodec decode contention
    all_decode_lats = []
    for p in profiles:
        all_decode_lats.extend(p.decode_latencies)
    if all_decode_lats:
        print(f"\n  BICODEC DECODE (all calls across {len(profiles)} requests):")
        print(f"    total calls: {len(all_decode_lats)}")
        print(f"    avg={statistics.mean(all_decode_lats):.2f}ms  "
              f"p50={sorted(all_decode_lats)[len(all_decode_lats)//2]:.2f}ms  "
              f"p99={sorted(all_decode_lats)[int(len(all_decode_lats)*0.99)]:.2f}ms  "
              f"max={max(all_decode_lats):.2f}ms")

    # Token generation rates
    all_rates = []
    for p in profiles:
        all_rates.extend(p.token_generation_rates)
    if all_rates:
        print(f"\n  TOKEN GEN RATE (all samples):")
        print(f"    avg={statistics.mean(all_rates):.0f} tok/s  "
              f"min={min(all_rates):.0f}  max={max(all_rates):.0f}")

# === GPU Memory Snapshot ===

def gpu_snapshot() -> str:
    try:
        import torch
        if not torch.cuda.is_available():
            return "No GPU"
        alloc = torch.cuda.memory_allocated() / 1e9
        reserved = torch.cuda.memory_reserved() / 1e9
        total = torch.cuda.get_device_properties(0).total_memory / 1e9
        return f"GPU: {alloc:.1f}GB alloc, {reserved:.1f}GB reserved / {total:.1f}GB total"
    except Exception:
        return "GPU info unavailable"

# === Main ===

async def main_async(args):
    # Initialize runtime
    logger.info("Initializing TTS runtime for profiling...")
    logger.info(gpu_snapshot())

    from veena3modal.local_server import resolve_model_paths, DEFAULT_LOCAL_MODEL_DIR

    model_path = args.model_path
    llm_path, bicodec_path = resolve_model_paths(model_path)

    from veena3modal.services.tts_runtime import initialize_runtime, get_runtime, is_initialized

    if not is_initialized():
        initialize_runtime(
            model_path=llm_path,
            bicodec_path=bicodec_path,
            device="cuda",
            # gpu_memory_utilization uses VLLM_CONFIG default (0.25 after Tier 1 optimization)
        )

    runtime = get_runtime()
    logger.info(f"Runtime ready: {runtime.model_version}")
    logger.info(gpu_snapshot())

    text = TEXTS.get(args.text, TEXTS["short"])
    concurrency_levels = [int(x) for x in args.concurrency.split(",")]

    # === Single Request Profiles ===
    if not args.concurrency_only:
        logger.info(f"\n{'=' * 80}")
        logger.info(f"SINGLE REQUEST PROFILING (text={args.text}, {len(text)} chars)")
        logger.info(f"{'=' * 80}")

        # Non-streaming
        if not args.stream:
            logger.info("\nProfiling non-streaming...")
            p = await profile_non_streaming(runtime, text, "lipakshi")
            print(p.report())

        # Streaming
        if not args.no_stream:
            logger.info("\nProfiling streaming...")
            p = await profile_streaming(runtime, text, "lipakshi")
            print(p.report())

    # === Concurrency Profiles ===
    for conc in concurrency_levels:
        modes = []
        if not args.stream:
            modes.append(False)
        if not args.no_stream:
            modes.append(True)

        for stream_mode in modes:
            logger.info(f"\nProfiling {conc} concurrent {'streaming' if stream_mode else 'non-streaming'} requests...")
            profiles = await profile_concurrent(runtime, conc, text, stream=stream_mode)
            print_concurrency_summary(profiles, conc, stream_mode)

    logger.info(f"\n{gpu_snapshot()}")
    logger.info("Profiling complete.")

def main():
    parser = argparse.ArgumentParser(description="Veena3 TTS Pipeline Profiler")
    parser.add_argument("--model-path", default=os.path.join(REPO_ROOT, "models", "spark_tts_4speaker"))
    parser.add_argument("--text", default="short", choices=["short", "medium", "long"])
    parser.add_argument("--concurrency", default="1,5,10,20", help="Concurrency levels (comma-separated)")
    parser.add_argument("--stream", action="store_true", help="Stream mode only")
    parser.add_argument("--no-stream", action="store_true", help="Non-stream mode only")
    parser.add_argument("--concurrency-only", action="store_true", help="Skip single-request profiles")
    args = parser.parse_args()
    asyncio.run(main_async(args))

if __name__ == "__main__":
    main()


cd /home/ubuntu/veenaModal && rg -n "class LongTextProcessor|generate_with_chunking\(|pipeline\.generate_speech_indic\(|generate_speech_stream_indic_first_chunk|generate_speech_stream_indic_continuation" veena3modal/processing/long_text_processor.py veena3modal/services/tts_runtime.py


Output:
Chunk ID: 54f1b3
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 299
Output:
veena3modal/services/tts_runtime.py:450:            audio_bytes, pipeline_timings = await runtime.pipeline.generate_speech_indic(
veena3modal/services/tts_runtime.py:456:            audio_bytes = await runtime.pipeline.generate_speech_indic(
veena3modal/services/tts_runtime.py:660:            audio_bytes = await long_processor.generate_with_chunking(
veena3modal/services/tts_runtime.py:673:            audio_bytes = await runtime.pipeline.generate_speech_indic(
veena3modal/services/tts_runtime.py:847:                stream_gen = runtime.streaming_pipeline.generate_speech_stream_indic_continuation(
veena3modal/services/tts_runtime.py:934:            async for audio_bytes, global_ids in runtime.streaming_pipeline.generate_speech_stream_indic_first_chunk(
veena3modal/services/tts_runtime.py:983:                async for audio_bytes in runtime.streaming_pipeline.generate_speech_stream_indic_continuation(
veena3modal/processing/long_text_processor.py:54:class LongTextProcessor:
veena3modal/processing/long_text_processor.py:141:    async def generate_with_chunking(
veena3modal/processing/long_text_processor.py:222:            audio_bytes = await self.pipeline.generate_speech_indic(


cd /home/ubuntu/veenaModal && sed -n '1,260p' veena3modal/processing/long_text_processor.py


Output:
Chunk ID: 9c6acd
Wall time: 0.0514 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,


cd /home/ubuntu/veenaModal && sed -n '260,520p' veena3modal/processing/long_text_processor.py


Output:
Chunk ID: 1c18a1
Wall time: 0.0503 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


I’m now checking `veena3modal/app.py` and the local runtime entrypoint to anchor the benchmark tooling to the exact request path you want measured, so the logs map cleanly to the same code path you’ll optimize next.


cd /home/ubuntu/veenaModal && sed -n '1,360p' veena3modal/app.py