Output:
Chunk ID: d178ee
Wall time: 0.0516 seconds
Process exited with code 0
Original token count: 2906
Output:
#!/usr/bin/env python3
"""
Direct-runtime stress benchmark with detailed timing for optimized local TTS path.

This bypasses HTTP serving and calls `tts_runtime` directly so we can isolate
generation + decode bottlenecks under concurrency.
"""

from __future__ import annotations

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

# Path setup (same pattern 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)

# Keep local auth disabled to match local_server behavior
os.environ["AUTH_BYPASS_MODE"] = "true"

TEST_TEXTS = {
    "short": "Hello, this is a quick benchmark test.",
    "medium": (
        "The quick brown fox jumps over the lazy dog. "
        "This benchmark sentence checks token generation speed, decode latency, "
        "and concurrency behavior for Spark TTS."
    ),
    "long": (
        "In the center of a crowded city, a quiet storyteller documented voices from dawn to dusk. "
        "Every pause, every laugh, every whisper revealed a different cadence. "
        "She replayed each recording at night, aligning tone and timing until each phrase felt natural."
    ),
}

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

@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 = 0.5):
        self.interval = interval_seconds
        self.snapshots: List[GPUSnapshot] = []
        self._running = False
        self._thread: Optional[threading.Thread] = None

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

    def stop(self) -> None:
        self._running = False
        if self._thread:
            self._thread.join(timeout=3)

    def _poll_loop(self) -> None:
        while self._running:
            try:
                res = 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 res.returncode == 0 and res.stdout:
                    parts = [p.strip() for p in res.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 {"samples": 0, "error": "no GPU samples captured"}
        mem = [s.memory_used_mb for s in self.snapshots]
        util = [s.gpu_utilization_pct for s in self.snapshots]
        temp = [s.temperature_c for s in self.snapshots]
        return {
            "samples": len(self.snapshots),
            "memory_mb": {
                "min": min(mem),
                "avg": statistics.mean(mem),
                "max": max(mem),
                "total": self.snapshots[0].memory_total_mb,
            },
            "gpu_util_pct": {
                "min": min(util),
                "avg": statistics.mean(util),
                "max": max(util),
            },
            "temperature_c": {
                "min": min(temp),
                "avg": statistics.mean(temp),
                "max": max(temp),
            },
        }

@dataclass
class RequestResult:
    success: bool
    latency_ms: float
    audio_bytes: int
    audio_seconds: float
    speaker: str
    text_length: int
    error: Optional[str] = None
    timing: Dict[str, Any] = field(default_factory=dict)

@dataclass
class LevelResult:
    concurrency: int
    total_requests: int
    results: List[RequestResult] = field(default_factory=list)
    wall_time_s: 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.0

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

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

    def timing_stats(self, key: str) -> Dict[str, float]:
        vals: List[float] = []
        for r in self.successes:
            v = r.timing.get(key)
            if isinstance(v, (int, float)):
                vals.append(float(v))
        if not vals:
            return {}
        return {
            "avg": statistics.mean(vals),
            "p50": self.p(vals, 50),
            "p95": self.p(vals, 95),
            "max": max(vals),
        }

    def latency_stats(self) -> Dict[str, float]:
        vals = [r.latency_ms for r in self.successes]
        if not vals:
            return {}
        return {
            "avg": statistics.mean(vals),
            "p50": self.p(vals, 50),
            "p95": self.p(vals, 95),
            "max": max(vals),
        }

    def latency_over_metric_stats(self, key: str) -> Dict[str, float]:
        vals: List[float] = []
        for r in self.successes:
            base = r.timing.get(key)
            if isinstance(base, (int, float)):
                vals.append(max(0.0, float(r.latency_ms) - float(base)))
        if not vals:
            return {}
        return {
            "avg": statistics.mean(vals),
            "p50": self.p(vals, 50),
            "p95": self.p(vals, 95),
            "max": max(vals),
        }

def init_runtime(
    model_path: str,
    gpu_mem: float,
    num_engines: int,
    max_num_batched_tokens: Optional[int],
    max_num_seqs: Optional[int],
    enable_chunked_prefill: Optional[bool],
    enable_prefix_caching: Optional[bool],
    disable_log_stats: Optional[bool],
    enforce_eager: Optional[bool],
) -> None:
    from veena3modal.local_server import resolve_model_paths
    from veena3modal.services.tts_runtime import initialize_runtime, is_initialized

    llm_path, bicodec_path = resolve_model_paths(model_path)
    if not is_initialized():
        initialize_runtime(
            model_path=llm_path,
            bicodec_path=bicodec_path,
            device="cuda",
            gpu_memory_utilization=gpu_mem,
            num_engines=num_engines,
            max_num_batched_tokens=max_num_batched_tokens,
            max_num_seqs=max_num_seqs,
            enable_chunked_prefill=enable_chunked_prefill,
            enable_prefix_caching=enable_prefix_caching,
            disable_log_stats=disable_log_stats,
            enforce_eager=enforce_eager,
            precompute_speaker_globals=False,
        )

async def run_one_request(
    text: str,
    speaker: str,
    max_tokens: int,
    timeout_s: float,
    chunking: bool,
) -> RequestResult:
    from veena3modal.services import tts_runtime

    t0 = time.perf_counter()
    try:
        if chunking:
            coro = tts_runtime.generate_speech_chunked(
                text=text,
                speaker=speaker,
                max_tokens=max_tokens,
            )
        else:
            coro = tts_runtime.generate_speech(
                text=text,
                speaker=speaker,
                max_tokens=max_tokens,
            )
        audio_bytes, timing = await asyncio.wait_for(coro, timeout=timeout_s)
        elapsed_ms = (time.perf_counter() - t0) * 1000
        if not audio_bytes:
            return RequestResult(
                success=False,
                latency_ms=elapsed_ms,
                audio_bytes=0,
                audio_seconds=0.0,
                speaker=speaker,
                text_length=len(text),
                error="no_audio",
                timing=timing or {},
            )
        sample_rate = int((timing or {}).get("output_sample_rate", 16000))
        audio_seconds = max(0.0, (len(audio_bytes) - 44) / (sample_rate * 2))
        return RequestResult(
            success=True,
            latency_ms=elapsed_ms,
            audio_bytes=len(audio_bytes),
            audio_seconds=audio_seconds,
            speaker=speaker,
            text_length=len(text),
            timing=timing or {},
        )
    except asyncio.TimeoutError:
        return RequestResult(
            success=False,
            latency_ms=(time.perf_counter() - t0) * 1000,
            audio_bytes=0,
            audio_seconds=0.0,
            speaker=speaker,
            text_length=len(text),
            error="timeout",
        )
    except Exception as exc:
        return RequestResult(
            success=False,
            latency_ms=(time.perf_counter() - t0) * 1000,
            audio_bytes=0,
            audio_seconds=0.0,
            speaker=speaker,
            text_length=len(text),
            error=str(exc)[:180],
        )

async def run_level(
    concurrency: int,
    total_requests: int,
    text: str,
    max_tokens: int,
    timeout_s: float,
    chunking: bool,
    gpu_monitor: GPUMonitor,
) -> LevelResult:
    level = LevelResult(concurrency=concurrency, total_requests=total_requests)
    semaphore = asyncio.Semaphore(concurrency)

    async def wrapped(i: int) -> RequestResult:
        async with semaphore:
            speaker = SPEAKERS[i % len(SPEAKERS)]
            return await run_one_request(
                text=text,
                speaker=speaker,
                max_tokens=max_tokens,
                timeout_s=timeout_s,
                chunking=chunking,
            )

    gpu_monitor.snapshots.clear()
    t0 = time.perf_counter()
    level.results = await asyncio.gather(*[wrapped(i) for i in range(total_requests)])
    level.wall_time_s = time.perf_counter() - t0
    level.gpu_summary = gpu_monitor.summary()
    return level

def print_level(level: LevelResult) -> None:
    print(f"\n=== concurrency={level.concurrency} requests={level.total_requests} ===")
    print(f"success={len(level.successes)}/{level.total_requests} ({level.success_rate:.0%})")
    print(f"wall={level.wall_time_s:.2f}s throughput={level.throughput_rps:.2f} req/s")
    lat = level.latency_stats()
    if lat:
        print(f"latency_ms avg={lat['avg']:.0f} p50={lat['p50']:.0f} p95={lat['p95']:.0f} max={lat['max']:.0f}")

    keys = [
        "generation_ms",


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


**Identifying reusable scheduler knobs**


cd /home/ubuntu/veenaModal && sed -n '360,860p' scripts/stress_test_runtime_detailed.py


Output:
Chunk ID: 5d5ec7
Wall time: 0.0513 seconds
Process exited with code 0
Original token count: 3180
Output:
        "generation_ms",
        "llm_time_in_queue_ms",
        "llm_scheduler_ms",
        "llm_first_token_ms",
        "llm_request_lifecycle_ms",
        "llm_queued_to_scheduled_ms",
        "llm_scheduled_to_first_token_ms",
        "llm_first_to_last_token_ms",
        "llm_queued_to_last_token_ms",
        "timeline_llm_first_batch_ms",
        "timeline_llm_done_ms",
        "timeline_parse_done_ms",
        "timeline_bicodec_done_ms",
        "timeline_total_ms",
        "llm_batch_wall_ms_total",
        "llm_batch_gpu_ms_total",
        "llm_decode_wall_ms_total",
        "llm_decode_gpu_ms_total",
        "llm_parse_ms",
        "llm_time_per_token_ms",
        "bicodec_decode_wall_ms",
        "bicodec_decode_gpu_ms",
    ]
    for key in keys:
        st = level.timing_stats(key)
        if st:
            print(f"{key}: avg={st['avg']:.2f} p50={st['p50']:.2f} p95={st['p95']:.2f} max={st['max']:.2f}")

    tl_total = level.timing_stats("timeline_total_ms")
    llm_wall = level.timing_stats("llm_batch_wall_ms_total")
    bicodec = level.timing_stats("bicodec_decode_wall_ms")
    latency_over_timeline = level.latency_over_metric_stats("timeline_total_ms")
    latency_over_generation = level.latency_over_metric_stats("generation_ms")
    if latency_over_timeline:
        print(
            "latency_over_timeline_ms: avg={:.2f} p50={:.2f} p95={:.2f} max={:.2f}".format(
                latency_over_timeline["avg"],
                latency_over_timeline["p50"],
                latency_over_timeline["p95"],
                latency_over_timeline["max"],
            )
        )
    if latency_over_generation:
        print(
            "latency_over_generation_ms: avg={:.2f} p50={:.2f} p95={:.2f} max={:.2f}".format(
                latency_over_generation["avg"],
                latency_over_generation["p50"],
                latency_over_generation["p95"],
                latency_over_generation["max"],
            )
        )
    if tl_total and llm_wall:
        total_avg = tl_total.get("avg", 0.0)
        llm_avg = llm_wall.get("avg", 0.0)
        bicodec_avg = bicodec.get("avg", 0.0) if bicodec else 0.0
        other_avg = max(0.0, total_avg - llm_avg - bicodec_avg)
        if total_avg > 0:
            print(
                "timeline_share: llm={:.1f}% bicodec={:.1f}% other={:.1f}%".format(
                    (llm_avg / total_avg) * 100,
                    (bicodec_avg / total_avg) * 100,
                    (other_avg / total_avg) * 100,
                )
            )
        samples = [
            r for r in level.successes
            if isinstance(r.timing.get("timeline_total_ms"), (int, float))
        ]
        if samples:
            sample = sorted(samples, key=lambda r: r.latency_ms)[len(samples) // 2]
            t = sample.timing
            print(
                "timeline_sample_ms: first_batch={:.1f} llm_done={:.1f} parse_done={:.1f} "
                "bicodec_done={:.1f} total={:.1f}".format(
                    float(t.get("timeline_llm_first_batch_ms", 0.0)),
                    float(t.get("timeline_llm_done_ms", 0.0)),
                    float(t.get("timeline_parse_done_ms", 0.0)),
                    float(t.get("timeline_bicodec_done_ms", 0.0)),
                    float(t.get("timeline_total_ms", 0.0)),
                )
            )

    if "memory_mb" in level.gpu_summary:
        g = level.gpu_summary
        print(
            "gpu: util_avg={:.0f}% util_max={:.0f}% mem_avg={:.0f}MB mem_max={:.0f}MB".format(
                g["gpu_util_pct"]["avg"],
                g["gpu_util_pct"]["max"],
                g["memory_mb"]["avg"],
                g["memory_mb"]["max"],
            )
        )

    errors: Dict[str, int] = {}
    for r in level.failures:
        k = r.error or "failed"
        errors[k] = errors.get(k, 0) + 1
    for k, v in sorted(errors.items(), key=lambda kv: -kv[1]):
        print(f"error[{v}] {k}")

def save_results(path: str, levels: List[LevelResult], args: argparse.Namespace) -> None:
    payload = {
        "config": {
            "levels": args.levels,
            "text": args.text,
            "max_tokens": args.max_tokens,
            "timeout_s": args.timeout,
            "chunking": args.chunking,
            "requests_multiplier": args.requests_multiplier,
            "gpu_memory": args.gpu_memory,
            "num_engines": args.num_engines,
            "max_num_batched_tokens": args.max_num_batched_tokens,
            "max_num_seqs": args.max_num_seqs,
            "enable_chunked_prefill": args.enable_chunked_prefill,
            "disable_chunked_prefill": args.disable_chunked_prefill,
            "disable_prefix_caching": args.disable_prefix_caching,
            "disable_engine_stats_logs": args.disable_engine_stats_logs,
            "enforce_eager": args.enforce_eager,
            "enable_gpu_decode_timing": args.enable_gpu_decode_timing,
            "non_stream_final_only": args.non_stream_final_only,
            "timestamp": int(time.time()),
        },
        "levels": [],
    }
    for lvl in levels:
        payload["levels"].append(
            {
                "concurrency": lvl.concurrency,
                "total_requests": lvl.total_requests,
                "success_rate": lvl.success_rate,
                "throughput_rps": lvl.throughput_rps,
                "wall_time_s": lvl.wall_time_s,
                "latency_ms": lvl.latency_stats(),
                "gpu": lvl.gpu_summary,
                "timing": {
                    "generation_ms": lvl.timing_stats("generation_ms"),
                    "llm_time_in_queue_ms": lvl.timing_stats("llm_time_in_queue_ms"),
                    "llm_scheduler_ms": lvl.timing_stats("llm_scheduler_ms"),
                    "llm_first_token_ms": lvl.timing_stats("llm_first_token_ms"),
                    "llm_request_lifecycle_ms": lvl.timing_stats("llm_request_lifecycle_ms"),
                    "llm_queued_to_scheduled_ms": lvl.timing_stats("llm_queued_to_scheduled_ms"),
                    "llm_scheduled_to_first_token_ms": lvl.timing_stats("llm_scheduled_to_first_token_ms"),
                    "llm_first_to_last_token_ms": lvl.timing_stats("llm_first_to_last_token_ms"),
                    "llm_queued_to_last_token_ms": lvl.timing_stats("llm_queued_to_last_token_ms"),
                    "llm_batch_wall_ms_total": lvl.timing_stats("llm_batch_wall_ms_total"),
                    "llm_batch_gpu_ms_total": lvl.timing_stats("llm_batch_gpu_ms_total"),
                    "llm_decode_wall_ms_total": lvl.timing_stats("llm_decode_wall_ms_total"),
                    "llm_decode_gpu_ms_total": lvl.timing_stats("llm_decode_gpu_ms_total"),
                    "llm_parse_ms": lvl.timing_stats("llm_parse_ms"),
                    "llm_time_per_token_ms": lvl.timing_stats("llm_time_per_token_ms"),
                    "bicodec_decode_wall_ms": lvl.timing_stats("bicodec_decode_wall_ms"),
                    "bicodec_decode_gpu_ms": lvl.timing_stats("bicodec_decode_gpu_ms"),
                    "timeline_llm_first_batch_ms": lvl.timing_stats("timeline_llm_first_batch_ms"),
                    "timeline_llm_done_ms": lvl.timing_stats("timeline_llm_done_ms"),
                    "timeline_parse_done_ms": lvl.timing_stats("timeline_parse_done_ms"),
                    "timeline_bicodec_done_ms": lvl.timing_stats("timeline_bicodec_done_ms"),
                    "timeline_total_ms": lvl.timing_stats("timeline_total_ms"),
                    "latency_over_timeline_ms": lvl.latency_over_metric_stats("timeline_total_ms"),
                    "latency_over_generation_ms": lvl.latency_over_metric_stats("generation_ms"),
                },
                "requests": [
                    {
                        "success": r.success,
                        "latency_ms": r.latency_ms,
                        "audio_bytes": r.audio_bytes,
                        "audio_seconds": r.audio_seconds,
                        "speaker": r.speaker,
                        "text_length": r.text_length,
                        "error": r.error,
                        "timing": r.timing,
                    }
                    for r in lvl.results
                ],
            }
        )
    with open(path, "w", encoding="utf-8") as fd:
        json.dump(payload, fd, indent=2)

async def main_async(args: argparse.Namespace) -> None:
    text = TEST_TEXTS[args.text]
    levels = [int(x.strip()) for x in args.levels.split(",") if x.strip()]
    gpu_monitor = GPUMonitor(interval_seconds=0.5)
    gpu_monitor.start()

    # warmup
    await run_one_request(
        text=text,
        speaker=SPEAKERS[0],
        max_tokens=args.max_tokens,
        timeout_s=args.timeout,
        chunking=args.chunking,
    )

    results: List[LevelResult] = []
    for idx, c in enumerate(levels):
        n = max(args.min_requests, c * args.requests_multiplier)
        level = await run_level(
            concurrency=c,
            total_requests=n,
            text=text,
            max_tokens=args.max_tokens,
            timeout_s=args.timeout,
            chunking=args.chunking,
            gpu_monitor=gpu_monitor,
        )
        print_level(level)
        results.append(level)
        if idx < len(levels) - 1:
            await asyncio.sleep(2)

    gpu_monitor.stop()
    save_results(args.output, results, args)
    print(f"\nresults saved: {args.output}")

def main() -> None:
    parser = argparse.ArgumentParser(description="Direct runtime stress benchmark (detailed timings)")
    parser.add_argument("--levels", default="1,2,4,8", help="Comma-separated concurrency levels")
    parser.add_argument("--text", choices=["short", "medium", "long"], default="short")
    parser.add_argument("--max-tokens", type=int, default=128, help="max_tokens sent to generation")
    parser.add_argument("--timeout", type=float, default=120.0, help="Per-request timeout seconds")
    parser.add_argument("--requests-multiplier", type=int, default=4, help="Requests per level ~= concurrency * multiplier")
    parser.add_argument("--min-requests", type=int, default=8)
    parser.add_argument("--chunking", action="store_true", default=False, help="Use generate_speech_chunked")
    parser.add_argument("--gpu-memory", type=float, default=0.25)
    parser.add_argument("--num-engines", type=int, default=1, help="Number of vLLM engines")
    parser.add_argument("--max-num-batched-tokens", type=int, default=None)
    parser.add_argument("--max-num-seqs", type=int, default=None)
    parser.add_argument("--enable-chunked-prefill", action="store_true", default=None)
    parser.add_argument("--disable-chunked-prefill", action="store_true", default=False)
    parser.add_argument("--disable-prefix-caching", action="store_true", default=False)
    parser.add_argument("--disable-engine-stats-logs", action="store_true", default=False)
    parser.add_argument("--enforce-eager", action="store_true", default=False)
    parser.add_argument(
        "--non-stream-final-only",
        action="store_true",
        default=False,
        help="Use vLLM FINAL_ONLY output mode (less per-token driver overhead, less granular timeline)",
    )
    parser.add_argument(
        "--enable-gpu-decode-timing",
        action="store_true",
        default=False,
        help="Enable precise BiCodec GPU timing (adds synchronize overhead)",
    )
    parser.add_argument("--model-path", type=str, default="")
    parser.add_argument("--output", default="stress_runtime_detailed.json")
    args = parser.parse_args()

    from veena3modal.local_server import DEFAULT_LOCAL_MODEL_DIR
    model_path = args.model_path or DEFAULT_LOCAL_MODEL_DIR
    if args.enable_gpu_decode_timing:
        os.environ["VEENA3_PERF_GPU_TIMING"] = "1"
    else:
        os.environ.setdefault("VEENA3_PERF_GPU_TIMING", "0")
    if args.non_stream_final_only:
        os.environ["VEENA3_NON_STREAM_FINAL_ONLY"] = "1"
    else:
        os.environ.setdefault("VEENA3_NON_STREAM_FINAL_ONLY", "0")
    enable_chunked_prefill = args.enable_chunked_prefill
    if args.disable_chunked_prefill:
        enable_chunked_prefill = False
    init_runtime(
        model_path=model_path,
        gpu_mem=args.gpu_memory,
        num_engines=args.num_engines,
        max_num_batched_tokens=args.max_num_batched_tokens,
        max_num_seqs=args.max_num_seqs,
        enable_chunked_prefill=enable_chunked_prefill,
        enable_prefix_caching=(False if args.disable_prefix_caching else None),
        disable_log_stats=(True if args.disable_engine_stats_logs else None),
        enforce_eager=(True if args.enforce_eager else None),
    )

    asyncio.run(main_async(args))

if __name__ == "__main__":
    main()


Summary: auto


**Analyzing streaming decode concurrency**


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


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


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


Output:
Chunk ID: 8cd732
Wall time: 0.0509 seconds
Process exited with code 0
Original token count: 3256
Output:
"""
Spark TTS BiCodec Decoder

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            return wav_np

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

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

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

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

        if audio is None:
            return None

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

        return audio_int16.tobytes()

    async def decode_to_bytes_async(
        self,
        semantic_ids: List[int],
        global_ids: List[int],
    ) -> Optional[bytes]:
        """
        Async wrapper for decode_to_bytes.

        Runs decode in the default executor so the event loop stays responsive
        under high request concurrency.
        """
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(
            None,
            self.decode_to_bytes,
            semantic_ids,
            global_ids,
        )

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

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

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

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

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

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

        return True

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

        Args:
            semantic_ids: List of semantic token IDs

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

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

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

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

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

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

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

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

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

        if audio is None:
            return None

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

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

        return audio_int16.tobytes()

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

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

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

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


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

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

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

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

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

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

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

import numpy as np
import torch

logger = logging.getLogger(__name__)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            if not batch:
                continue

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

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

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

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

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

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

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

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

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

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

        return results

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


Output:
Chunk ID: 70cc09
Wall time: 0.0504 seconds
Process exited with code 0
Original token count: 1800
Output:
"""
Spark TTS Model Loader

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

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

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

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

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

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

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

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

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

        self.model_path = model_path
        self.hf_token = hf_token

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


**Inspecting batch decoder usage**


cd /home/ubuntu/veenaModal && rg -n "BiCodecBatchDecoder|decode_batch|batch_decoder|enable_batching" veena3modal -g '*.py'


Output:
Chunk ID: 7f3d43
Wall time: 0.0509 seconds
Process exited with code 0
Original token count: 342
Output:
veena3modal/core/streaming_pipeline.py:179:                        if self.snac_decoder.enable_batching:
veena3modal/core/streaming_pipeline.py:597:                        if self.snac_decoder.enable_batching:
veena3modal/core/streaming_pipeline.py:801:                    if self.snac_decoder.enable_batching:
veena3modal/core/bicodec_batch_decoder.py:16:    batch_decoder = BiCodecBatchDecoder(bicodec_decoder, max_batch=16, timeout_ms=10)
veena3modal/core/bicodec_batch_decoder.py:17:    batch_decoder.start()
veena3modal/core/bicodec_batch_decoder.py:20:    audio_bytes = await batch_decoder.decode_async(semantic_ids, global_ids)
veena3modal/core/bicodec_batch_decoder.py:23:    batch_decoder.stop()
veena3modal/core/bicodec_batch_decoder.py:49:class BiCodecBatchDecoder:
veena3modal/core/bicodec_batch_decoder.py:90:        logger.info(f"BiCodecBatchDecoder started (max_batch={self.max_batch}, timeout={self.timeout_ms}ms)")
veena3modal/core/bicodec_decoder.py:51:        enable_batching: bool = False,  # For compatibility with streaming pipeline
veena3modal/core/bicodec_decoder.py:59:            enable_batching: Enable async batching for streaming (not used yet)
veena3modal/core/bicodec_decoder.py:63:        self.enable_batching = enable_batching  # For streaming pipeline compatibility
veena3modal/core/bicodec_decoder.py:82:        if enable_batching:


**Planning vLLM scheduler tuning**


cd /home/ubuntu/veenaModal && sed -n '1,320p' veena3modal/core/constants.py


Output:
Chunk ID: a09d7c
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 1489
Output:
"""
Spark TTS Constants

Token IDs, special tokens, speakers, and emotions for Spark TTS model.
Migrated from Veena3/Orpheus to Spark TTS (BiCodec) architecture.
"""

import os

# Stop tokens for generation
# Spark TTS uses <|im_end|> as stop token
TRAINING_STOP_TOKEN_IDS = ["<|im_end|>"]  # String-based stop token for Spark TTS

# Spark TTS Speaker System - 12 Speakers from Training
# Maps user-facing speaker names to internal speaker IDs (speaker_0 to speaker_11)
# Model: BayAreaBoys/spark_tts_4speaker (HuggingFace) has 12 speaker tokens
SPEAKER_MAP = {
    "lipakshi": 0,    # speaker_0
    "vardan": 1,      # speaker_1
    "reet": 2,        # speaker_2
    "Nandini": 3,     # speaker_3
    "krishna": 4,     # speaker_4
    "anika": 5,       # speaker_5
    "adarsh": 6,      # speaker_6
    "Nilay": 7,       # speaker_7
    "Aarvi": 8,       # speaker_8
    "Asha": 9,        # speaker_9
    "Bittu": 10,      # speaker_10
    "Mira": 11,       # speaker_11
}

# Friendly speaker name mappings (user-facing → internal)
# Users can use these friendly names instead of internal names
FRIENDLY_SPEAKER_MAP = {
    "Mitra": "lipakshi",
    "Aaranya": "reet",
    "Taru": "Nandini",
    "Neer": "Nilay",
    "Dhruva": "vardan",
    "Ira": "anika",
    "Veda": "adarsh",
    "Aria": "krishna",
    # Direct mappings for new speakers
    "Aarvi": "Aarvi",
    "Asha": "Asha",
    "Bittu": "Bittu",
    "Mira": "Mira",
}

# Valid speaker names (for API validation)
# Include both internal names and friendly names
INDIC_SPEAKERS = list(SPEAKER_MAP.keys())
ALL_SPEAKER_NAMES = INDIC_SPEAKERS + list(FRIENDLY_SPEAKER_MAP.keys())

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

    Args:
        name: Speaker name (friendly or internal)

    Returns:
        Internal speaker name

    Examples:
        resolve_speaker_name("Mitra") -> "lipakshi"
        resolve_speaker_name("lipakshi") -> "lipakshi"
    """
    # Check if it's a friendly name
    if name in FRIENDLY_SPEAKER_MAP:
        return FRIENDLY_SPEAKER_MAP[name]

    # Check if it's already an internal name
    if name in INDIC_SPEAKERS:
        return name

    # Invalid speaker name
    raise ValueError(f"Invalid speaker name: {name}. Valid names: {', '.join(ALL_SPEAKER_NAMES)}")

# Spark TTS Emotion Tags (bracket format: [emotion])
# NOTE: API users must now use [emotion] instead of <emotion>
INDIC_EMOTION_TAGS = [
    "[angry]",
    "[curious]",
    "[excited]",
    "[giggle]",
    "[laughs harder]",
    "[laughs]",
    "[screams]",
    "[sighs]",
    "[sings]",
    "[whispers]"
]

# Legacy emotion tags for backward compatibility mapping
# Maps old <emotion> format to new [emotion] format
LEGACY_EMOTION_MAP = {
    "<angry>": "[angry]",
    "<curious>": "[curious]",
    "<excited>": "[excited]",
    "<giggle>": "[giggle]",
    "<laugh_harder>": "[laughs harder]",
    "<laugh>": "[laughs]",
    "<scream>": "[screams]",
    "<sigh>": "[sighs]",
    "<sing>": "[sings]",
    "<whisper>": "[whispers]"
}

# All emotion tags (for validation)
ALL_EMOTION_TAGS = INDIC_EMOTION_TAGS

# Model configuration
DEFAULT_MODEL_PATH = "BayAreaBoys/spark_tts_4speaker"
DEFAULT_MAX_MODEL_LEN = 4096  # Lower than old model (was 8192)

# BiCodec Audio Tokenizer
# Location: resolved from env var MODEL_PATH or default
BICODEC_TOKENIZER_PATH = os.environ.get('MODEL_PATH', os.environ.get('SPARK_TTS_MODEL_PATH', '/models/spark_tts_4speaker'))

# Audio configuration
AUDIO_SAMPLE_RATE = 16000  # Changed from 24kHz to 16kHz for BiCodec
AUDIO_CHANNELS = 1
AUDIO_BITS_PER_SAMPLE = 16

# Generation defaults for Spark TTS
DEFAULT_TEMPERATURE = 0.8  # Higher than old model (was 0.4)
DEFAULT_TOP_K = 50
DEFAULT_TOP_P = 1.0
DEFAULT_MAX_TOKENS = 4096  # Increased to handle complex multilingual generation
DEFAULT_MIN_TOKENS = 28  # Minimum tokens for generation (legacy compatibility)
DEFAULT_REPETITION_PENALTY = 1.0  # Repetition penalty (legacy compatibility)
DEFAULT_SEED = None  # None = random, set integer for reproducibility

# vLLM Configuration optimizations for Spark TTS
# NOTE: gpu_memory_utilization adjusted based on available memory
# Production: 0.85, Limited memory: 0.3-0.5
#
# OPTIMIZATION Dec 2025:
# - enable_chunked_prefill: Prevents long prompts from blocking concurrent streams
# - async_scheduling: DISABLED - caused EngineCore issues in production
# OPTIMIZATION Feb 2026:
# Reduced gpu_memory_utilization from 0.85 to 0.25
# Model is only 0.5B (~1.3GB weights). At 0.85, vLLM pre-allocated 65GB for KV cache
# (enough for 1,399 concurrent seqs) when peak actual need is ~100 concurrent.
# At 0.25: ~18GB KV cache, supports ~380 concurrent seqs (4x peak), frees ~48GB VRAM.
# This enables fitting on L4 (24GB) or even T4 (16GB) GPUs.
VLLM_CONFIG = {
    "dtype": "bfloat16",
    "gpu_memory_utilization": 0.25,  # 0.5B model needs <2GB; 0.25 gives ~380 concurrent seqs
    "max_model_len": 4096,
    "max_num_batched_tokens": 4096,
    "enable_prefix_caching": True,  # Cache common prompts for faster TTFB
    "enable_chunked_prefill": True,  # Chunk long prefills to avoid blocking concurrent streams
    "enforce_eager": False,  # Enable CUDA graphs for low latency
    "disable_log_stats": False,
    "trust_remote_code": True,
    "tensor_parallel_size": 1,
}

# Streaming configuration (may need adjustment for BiCodec)
STREAM_BUFFER_SIZE = 50  # Buffer size for streaming (BiCodec-specific, TBD)
BICODEC_BATCH_SIZE = 64
BICODEC_BATCH_TIMEOUT_MS = 15

# Backward compatibility: Old SNAC constants for files not yet migrated
# These are NOT used by Spark TTS but needed by legacy streaming code
# TODO: Remove once streaming_pipeline.py is fully migrated to BiCodec
CODE_START_TOKEN_ID = 128257  # Legacy SNAC constant
CODE_END_TOKEN_ID = 128258    # Legacy SNAC constant
SNAC_MIN_ID = 128266          # Legacy SNAC constant
SNAC_MAX_ID = 156937          # Legacy SNAC constant


**Designing detailed streaming metrics**


**Implementing env-configurable decode and detailed metrics**


**Preparing top imports patch**


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


Output:
Chunk ID: e60275
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 1959
Output:
     1	"""
     2	Veena3 Streaming Pipeline - Sliding Window Approach
     3
     4	Implements streaming for both SNAC and BiCodec tokens.
     5
     6	For BiCodec (Spark TTS):
     7	1. Generate tokens from vLLM
     8	2. Parse NEW token IDs incrementally (O(1) per token via cache)
     9	3. Buffer semantic + global tokens
    10	4. Apply sliding window (every N token pairs)
    11	5. Decode and stream audio chunks
    12
    13	OPTIMIZATION (Dec 2025):
    14	- Replaced O(n²) pattern (decode-all + regex-all per iteration)
    15	- Now uses incremental token parsing with BiCodecTokenParser
    16	- ~10x CPU reduction in streaming hot loop
    17
    18	For SNAC (legacy):
    19	1. Filter SNAC token IDs directly
    20	2. Apply sliding window
    21	3. Decode and stream
    22	"""
    23
    24	import asyncio
    25	import logging
    26	import os
    27	import re
    28	from typing import AsyncGenerator, Optional, List
    29	from vllm import SamplingParams
    30	from vllm.sampling_params import RequestOutputKind
    31	from veena3modal.audio.crossfade import crossfade_bytes_int16
    32	from veena3modal.core.token_utils import BiCodecTokenParser
    33
    34	from veena3modal.core.constants import (
    35	    CODE_END_TOKEN_ID,
    36	    CODE_START_TOKEN_ID,
    37	    SNAC_MIN_ID,
    38	    SNAC_MAX_ID,
    39	    TRAINING_STOP_TOKEN_IDS,
    40	    DEFAULT_TEMPERATURE,
    41	    DEFAULT_TOP_K,
    42	    DEFAULT_TOP_P,
    43	    DEFAULT_MAX_TOKENS,
    44	    DEFAULT_MIN_TOKENS,
    45	    DEFAULT_REPETITION_PENALTY,
    46	    DEFAULT_SEED,
    47	)
    48
    49	logger = logging.getLogger(__name__)
    50	_STREAM_OUTPUT_KIND_RAW = os.environ.get("VEENA3_STREAM_OUTPUT_KIND", "delta").strip().lower()
    51	_STREAM_USE_DELTA_OUTPUT = _STREAM_OUTPUT_KIND_RAW in {"delta", "d", "1", "true", "yes", "on"}
    52
    53
    54	class Veena3SlidingWindowPipeline:
    55	    """
    56	    Streaming TTS pipeline using sliding window approach.
    57
    58	    This eliminates choppy audio and popping artifacts by:
    59	    - Decoding overlapping 28-token windows (4 frames)
    60	    - Keeping only the middle 2048 samples from each decode
    61	    - Creating natural continuity between chunks
    62
    63	    Based on the official Canopy Labs implementation.
    64	    """
    65
    66	    def __init__(
    67	        self,
    68	        model,
    69	        prompt_builder,
    70	        snac_decoder,
    71	    ):
    72	        """
    73	        Initialize sliding window streaming pipeline.
    74
    75	        Args:
    76	            model: Veena3Model instance
    77	            prompt_builder: Veena3PromptBuilder instance
    78	            snac_decoder: SNACDecoder instance (with batching enabled)
    79	        """
    80	        self.model = model
    81	        self.prompt_builder = prompt_builder
    82	        self.snac_decoder = snac_decoder
    83
    84	        # OPTIMIZATION: Pre-warm BiCodecTokenParser once at init, not per-request
    85	        # Saves ~123ms per streaming TTFB by avoiding repeated 166K-entry vocab iteration
    86	        tokenizer = getattr(model, "tokenizer", None)
    87	        if tokenizer is None:
    88	            tokenizer = getattr(model.engine, "tokenizer", None)
    89	        self.token_parser = BiCodecTokenParser(tokenizer) if tokenizer else None
    90
    91	        print(f"🌊 Veena3SlidingWindowPipeline initialized (sliding window: 28 tokens)")
    92
    93	    async def generate_speech_stream(
    94	        self,
    95	        description: str,
    96	        text: str,
    97	        temperature: float = DEFAULT_TEMPERATURE,
    98	        top_p: float = DEFAULT_TOP_P,
    99	        max_tokens: int = DEFAULT_MAX_TOKENS,
   100	        repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
   101	        seed: Optional[int] = None,
   102	    ) -> AsyncGenerator[bytes, None]:
   103	        """
   104	        Generate speech audio with sliding window streaming.
   105
   106	        Yields audio chunks using overlapping windows for smooth playback.
   107
   108	        Args:
   109	            description: Character/voice description
   110	            text: Text to synthesize (with optional <emotion> tags)
   111	            temperature: Sampling temperature
   112	            top_p: Nucleus sampling
   113	            max_tokens: Max SNAC tokens to generate
   114	            repetition_penalty: Prevent loops
   115
   116	        Yields:
   117	            Audio bytes (int16 PCM, 24kHz mono)
   118	        """
   119	        logger.debug("sliding-window streaming generation started")
   120	        logger.debug("description_prefix=%s", description[:80])
   121	        logger.debug("text_len=%d", len(text))
   122
   123	        # Build prompt
   124	        prompt = self.prompt_builder.build_prefix(description, text)
   125	        logger.debug("prompt built chars=%d", len(prompt))
   126
   127	        # Configure sampling
   128	        sampling_params = SamplingParams(
   129	            temperature=temperature,
   130	            top_p=top_p,
   131	            max_tokens=max_tokens,
   132	            min_tokens=DEFAULT_MIN_TOKENS,
   133	            repetition_penalty=repetition_penalty,
   134	            stop_token_ids=[CODE_END_TOKEN_ID],  # Only audio EOS
   135	            seed=seed,  # None = random, int = reproducible
   136	        )
   137
   138	        logger.debug("sampling temp=%s top_p=%s sliding_window=28", temperature, top_p)
   139
   140	        # Token buffer - keeps ALL tokens (not chunked)
   141	        token_buffer = []
   142	        total_tokens_generated = 0
   143	        total_audio_chunks = 0
   144
   145	        # Generate tokens with vLLM (streaming)
   146	        print(f"🔮 Starting token generation...")
   147
   148	        # Generate unique request ID for concurrent streaming support
   149	        import uuid
   150	        import time
   151	        request_id = f"slide-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"
   152	        results_generator = self.model.engine.generate(
   153	            prompt=prompt,
   154	            sampling_params=sampling_params,
   155	            request_id=request_id,
   156	        )
   157
   158	        # Stream tokens with sliding window
   159	        async for request_output in results_generator:
   160	            # Extract generated token IDs
   161	            generated_ids = request_output.outputs[0].token_ids
   162
   163	            # Process only new tokens
   164	            new_tokens = generated_ids[total_tokens_generated:]
   165	            total_tokens_generated = len(generated_ids)
   166
   167	            # Filter and buffer SNAC tokens
   168	            for token_id in new_tokens:
   169	                if SNAC_MIN_ID <= token_id <= SNAC_MAX_ID:
   170	                    token_buffer.append(token_id)
   171
   172	                    # Process every 7 tokens (1 frame) when we have enough for sliding window
   173	                    # Official approach: once count > 27, take last 28 tokens
   174	                    if len(token_buffer) % 7 == 0 and len(token_buffer) > 27:
   175	                        # Sliding window: take last 28 tokens (4 frames)
   176	                        window_tokens = token_buffer[-28:]
   177
   178	                        # Decode with sliding window mode (returns middle 2048 samples only)
   179	                        if self.snac_decoder.enable_batching:
   180	                            audio_bytes = await self.snac_decoder.decode_single_async(


**Gathering context for patch**


cd /home/ubuntu/veenaModal && nl -ba veena3modal/core/streaming_pipeline.py | sed -n '180,620p'


Output:
Chunk ID: 6b8ad4
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 6079
Output:
   180	                            audio_bytes = await self.snac_decoder.decode_single_async(
   181	                                window_tokens,
   182	                                trim_warmup=False,  # Sliding window handles trimming
   183	                                use_sliding_window=True  # CRITICAL: Use sliding window mode
   184	                            )
   185	                        else:
   186	                            audio_bytes = self.snac_decoder.decode_to_bytes(
   187	                                window_tokens,
   188	                                trim_warmup=False,
   189	                                use_sliding_window=True
   190	                            )
   191
   192	                        if audio_bytes:
   193	                            total_audio_chunks += 1
   194	                            if total_audio_chunks == 1:
   195	                                logger.debug("first sliding-window chunk decoded bytes=%d", len(audio_bytes))
   196	                            yield audio_bytes
   197
   198	        # Note: No final chunk processing needed - sliding window handles all tokens
   199	        # as they come in (every 7 tokens after the first 28)
   200
   201	        logger.debug(
   202	            "sliding-window complete tokens=%d chunks=%d",
   203	            total_tokens_generated,
   204	            total_audio_chunks,
   205	        )
   206
   207	    def _extract_bicodec_tokens_from_text(self, text: str) -> tuple[List[int], List[int]]:
   208	        """
   209	        Extract BiCodec semantic and global tokens from generated text.
   210
   211	        Args:
   212	            text: Generated text containing BiCodec token markers
   213
   214	        Returns:
   215	            Tuple of (semantic_ids, global_ids)
   216	        """
   217	        semantic_matches = re.findall(r"<\|bicodec_semantic_(\d+)\|>", text)
   218	        semantic_ids = [int(t) for t in semantic_matches]
   219
   220	        global_matches = re.findall(r"<\|bicodec_global_(\d+)\|>", text)
   221	        global_ids = [int(t) for t in global_matches]
   222
   223	        return semantic_ids, global_ids
   224
   225	    async def generate_speech_stream_indic(
   226	        self,
   227	        speaker: str,
   228	        text: str,
   229	        temperature: float = DEFAULT_TEMPERATURE,
   230	        top_k: int = DEFAULT_TOP_K,  # Added for Spark TTS parity with non-streaming
   231	        top_p: float = DEFAULT_TOP_P,
   232	        max_tokens: int = DEFAULT_MAX_TOKENS,
   233	        repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
   234	        seed: Optional[int] = None,
   235	    ) -> AsyncGenerator[bytes, None]:
   236	        """
   237	        Generate speech audio with sliding window streaming for Indic model.
   238
   239	        Yields audio chunks using overlapping windows for smooth playback.
   240
   241	        Args:
   242	            speaker: Speaker name (one of 12 predefined speakers)
   243	            text: Text to synthesize with inline emotion tags
   244	                Examples:
   245	                - "Hello! Welcome."
   246	                - "<laugh> Hello there!"
   247	                - "नमस्ते! <excited> आज का दिन बहुत अच्छा है।"
   248	            temperature: Sampling temperature
   249	            top_p: Nucleus sampling
   250	            max_tokens: Max SNAC tokens to generate
   251	            repetition_penalty: Prevent loops
   252	            seed: Random seed for reproducibility
   253
   254	        Yields:
   255	            Audio bytes (int16 PCM, 16kHz mono - BiCodec)
   256	        """
   257	        import time
   258	        t_start = time.time()
   259
   260	        # Build prompt using Indic prompt builder
   261	        prompt = self.prompt_builder.build_prefix(speaker, text)
   262
   263	        # Configure sampling (matching non-streaming pipeline parameters)
   264	        # OPTIMIZATION: Use Spark TTS stop token, not legacy SNAC CODE_END_TOKEN_ID (128258)
   265	        # which doesn't exist in the Spark vocab and would cause generation to run to max_tokens
   266	        sampling_params = SamplingParams(
   267	            temperature=temperature,
   268	            top_k=top_k,
   269	            top_p=top_p,
   270	            max_tokens=max_tokens,
   271	            min_tokens=DEFAULT_MIN_TOKENS,
   272	            repetition_penalty=repetition_penalty,
   273	            stop=TRAINING_STOP_TOKEN_IDS,  # "<|im_end|>" - matches non-streaming pipeline
   274	            skip_special_tokens=False,  # Keep BiCodec tokens in output
   275	            seed=seed,
   276	            output_kind=(RequestOutputKind.DELTA if _STREAM_USE_DELTA_OUTPUT else RequestOutputKind.CUMULATIVE),
   277	        )
   278
   279	        # BiCodec token buffers (separate for semantic and global)
   280	        semantic_buffer = []
   281	        global_buffer = []
   282	        processed_token_count = 0  # Used for cumulative mode; ignored for delta mode.
   283	        total_audio_chunks = 0
   284
   285	        # BiCodec streaming configuration
   286	        EXPECTED_GLOBAL_COUNT = 32  # BiCodec uses 32 fixed global tokens (per paper)
   287	        MIN_SEMANTIC_FOR_FIRST_CHUNK = 10  # OPTIMIZATION: Lowered from 16 (decoder min is 8, -120ms TTFB)
   288	        # OPTIMIZATION: Increased from 24 to 48 to reduce O(n) re-decode calls by 2x.
   289	        # At 50 TPS, 48 tokens = ~960ms audio per chunk. Cuts total decode calls roughly in half.
   290	        DECODE_INTERVAL = 48
   291	        CROSSFADE_MS = 50
   292	        # OPTIMIZATION: Windowed decode - only decode last WINDOW_SIZE tokens instead of ALL.
   293	        # BiCodec WaveGenerator has finite receptive field (~64 tokens based on kernel sizes).
   294	        # Window must be >= receptive_field + DECODE_INTERVAL to avoid boundary artifacts.
   295	        WINDOW_SIZE = 128  # 128 tokens covers receptive field + margin for crossfade
   296
   297	        # Generate tokens with vLLM (streaming)
   298	        import uuid
   299	        request_id = f"bicodec-stream-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"
   300
   301	        # Reuse pre-warmed token parser (singleton, created once in __init__)
   302	        token_parser = self.token_parser
   303
   304	        t_gen_start = time.time()
   305	        results_generator = self.model.engine.generate(
   306	            prompt=prompt,
   307	            sampling_params=sampling_params,
   308	            request_id=request_id,
   309	        )
   310
   311	        # Track timing
   312	        t_first_token = None
   313	        t_first_chunk_ready = None
   314	        t_first_chunk_decoded = None
   315	        iteration_count = 0
   316	        last_decode_count = 0
   317
   318	        # Track audio samples EMITTED to the user (not just decoded)
   319	        # BiCodec produces ~320 samples per semantic token at 16kHz
   320	        # CRITICAL: This tracks samples that have been YIELDED, not held in tail
   321	        total_samples_emitted_to_user = 0
   322	        # Keep a tail to enable crossfading between consecutive emitted chunks
   323	        previous_chunk_tail: Optional[bytes] = None
   324
   325	        # Stream tokens with BiCodec extraction
   326	        async for request_output in results_generator:
   327	            iteration_count += 1
   328	            t_iter = time.time()
   329
   330	            # Track first token received
   331	            if t_first_token is None:
   332	                t_first_token = t_iter
   333
   334	            # Extract generated token IDs
   335	            generated_ids = request_output.outputs[0].token_ids
   336
   337	            # In DELTA mode token_ids are already incremental.
   338	            if _STREAM_USE_DELTA_OUTPUT:
   339	                new_token_ids = generated_ids
   340	            else:
   341	                new_token_ids = generated_ids[processed_token_count:]
   342	                processed_token_count = len(generated_ids)
   343
   344	            # Parse new tokens and append to buffers (in-place modification)
   345	            token_parser.parse_incremental(new_token_ids, semantic_buffer, global_buffer)
   346
   347	            # TRUE STREAMING with BiCodec two-phase generation:
   348	            # Phase 1: Generate 32 global tokens (pre-roll)
   349	            # Phase 2: Stream semantic tokens incrementally
   350
   351	            # Check if we have the 32 global tokens (pre-roll complete)
   352	            if len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
   353	                # Pre-roll complete! Now we can stream semantic tokens
   354	                if t_first_chunk_ready is None and len(semantic_buffer) >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
   355	                    t_first_chunk_ready = time.time()
   356
   357	                # Stream semantic tokens incrementally
   358	                # Only decode if we have enough NEW semantic tokens since last decode
   359	                semantic_count = len(semantic_buffer)
   360	                if semantic_count >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
   361	                    # Check if we should decode (every DECODE_INTERVAL new semantic tokens)
   362	                    if semantic_count - last_decode_count >= DECODE_INTERVAL:
   363	                        last_decode_count = semantic_count
   364
   365	                        # OPTIMIZATION: Windowed decode instead of full re-decode
   366	                        # Only decode the last WINDOW_SIZE tokens (O(1) per chunk instead of O(n))
   367	                        # For first few chunks when total tokens < WINDOW_SIZE, decode all
   368	                        window_global = global_buffer[:EXPECTED_GLOBAL_COUNT]
   369
   370	                        if len(semantic_buffer) <= WINDOW_SIZE:
   371	                            # Early phase: decode all tokens (window not yet relevant)
   372	                            decode_semantic = semantic_buffer
   373	                        else:
   374	                            # Windowed: only decode last WINDOW_SIZE tokens
   375	                            decode_semantic = semantic_buffer[-WINDOW_SIZE:]
   376
   377	                        t_decode_start = time.time()
   378	                        audio_bytes = await self.snac_decoder.decode_single_async(
   379	                            semantic_ids=decode_semantic,
   380	                            global_ids=window_global,
   381	                            trim_warmup=False,
   382	                            use_sliding_window=False,
   383	                        )
   384	                        t_decode_end = time.time()
   385
   386	                        if audio_bytes:
   387	                            total_samples_decoded = len(audio_bytes) // 2
   388
   389	                            if len(semantic_buffer) <= WINDOW_SIZE:
   390	                                # Early phase: extract new samples from full decode
   391	                                if total_samples_decoded > total_samples_emitted_to_user:
   392	                                    new_bytes_start = total_samples_emitted_to_user * 2
   393	                                    new_audio_bytes = audio_bytes[new_bytes_start:]
   394	                                else:
   395	                                    new_audio_bytes = b""
   396	                            else:
   397	                                # Windowed phase: take only the tail portion of the window decode
   398	                                # The last DECODE_INTERVAL tokens' worth of audio is the new content
   399	                                # Each token produces ~320 samples at 16kHz/50TPS
   400	                                new_samples_approx = DECODE_INTERVAL * 320
   401	                                if total_samples_decoded > new_samples_approx:
   402	                                    new_bytes_start = (total_samples_decoded - new_samples_approx) * 2
   403	                                else:
   404	                                    new_bytes_start = 0
   405	                                new_audio_bytes = audio_bytes[new_bytes_start:]
   406
   407	                            if new_audio_bytes:
   408	                                to_emit, previous_chunk_tail = crossfade_bytes_int16(
   409	                                    previous_chunk_tail,
   410	                                    new_audio_bytes,
   411	                                    sample_rate_hz=16000,
   412	                                    crossfade_ms=CROSSFADE_MS,
   413	                                )
   414
   415	                                samples_emitted_in_this_chunk = len(to_emit) // 2
   416	                                total_samples_emitted_to_user += samples_emitted_in_this_chunk
   417
   418	                                if to_emit:
   419	                                    total_audio_chunks += 1
   420	                                    if t_first_chunk_decoded is None:
   421	                                        t_first_chunk_decoded = t_decode_end
   422
   423	                                    yield to_emit
   424
   425	        # Final chunk: decode any remaining semantic tokens
   426	        # Use the same 32 global tokens (they get pooled)
   427	        if len(semantic_buffer) > last_decode_count and len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
   428	            audio_bytes = self.snac_decoder.decode_streaming(
   429	                semantic_ids=semantic_buffer,  # All semantic tokens
   430	                global_ids=global_buffer[:EXPECTED_GLOBAL_COUNT],  # Use first 32 global tokens
   431	                use_sliding_window=False,  # Get all samples
   432	                trim_warmup=False
   433	            )
   434	            if audio_bytes:
   435	                # Only yield NEW samples
   436	                total_samples_decoded = len(audio_bytes) // 2
   437
   438	                if total_samples_decoded > total_samples_emitted_to_user:
   439	                    new_bytes_start = total_samples_emitted_to_user * 2
   440	                    new_audio_bytes = audio_bytes[new_bytes_start:]
   441
   442	                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
   443	                        previous_chunk_tail,
   444	                        new_audio_bytes,
   445	                        sample_rate_hz=16000,
   446	                        crossfade_ms=CROSSFADE_MS,
   447	                    )
   448
   449	                    samples_emitted_in_final = len(to_emit) // 2
   450	                    total_samples_emitted_to_user += samples_emitted_in_final
   451
   452	                    if to_emit:
   453	                        total_audio_chunks += 1
   454	                        yield to_emit
   455
   456	        # Flush any remaining tail that was held back for crossfade
   457	        if previous_chunk_tail:
   458	            total_audio_chunks += 1
   459	            tail_samples = len(previous_chunk_tail) // 2
   460	            total_samples_emitted_to_user += tail_samples
   461	            yield previous_chunk_tail
   462
   463	        # Log completion summary
   464	        t_end = time.time()
   465	        audio_duration_s = total_samples_emitted_to_user / 16000  # 16kHz sample rate
   466	        ttfb_ms = (t_first_chunk_decoded - t_start) * 1000 if t_first_chunk_decoded else 0
   467	        logger.debug(
   468	            "streaming complete audio=%.2fs ttfb=%.0fms chunks=%d",
   469	            audio_duration_s,
   470	            ttfb_ms,
   471	            total_audio_chunks,
   472	        )
   473
   474	    async def generate_speech_stream_indic_first_chunk(
   475	        self,
   476	        speaker: str,
   477	        text: str,
   478	        temperature: float = DEFAULT_TEMPERATURE,
   479	        top_k: int = DEFAULT_TOP_K,
   480	        top_p: float = DEFAULT_TOP_P,
   481	        max_tokens: int = DEFAULT_MAX_TOKENS,
   482	        repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
   483	        seed: Optional[int] = None,
   484	    ) -> AsyncGenerator[tuple, None]:
   485	        """
   486	        Generate speech for the FIRST chunk of multi-chunk text, capturing global tokens.
   487
   488	        This method is specifically for chunked generation: it yields (audio_bytes, global_ids)
   489	        tuples instead of just audio_bytes. The caller captures global_ids from the first
   490	        yield and passes them to generate_speech_stream_indic_continuation() for subsequent chunks.
   491
   492	        Use Case (chunked streaming with voice consistency):
   493	            globals_captured = None
   494	            async for audio_bytes, global_ids in pipeline.generate_speech_stream_indic_first_chunk(...):
   495	                if globals_captured is None and global_ids:
   496	                    globals_captured = global_ids  # Capture once
   497	                yield audio_bytes
   498	            # Now use globals_captured for subsequent chunks
   499
   500	        Args:
   501	            speaker: Speaker name
   502	            text: First text chunk to synthesize
   503	            temperature: Sampling temperature
   504	            top_k: Top-k sampling
   505	            top_p: Nucleus sampling
   506	            max_tokens: Max tokens to generate
   507	            repetition_penalty: Prevent repetition
   508	            seed: Random seed for reproducibility
   509
   510	        Yields:
   511	            Tuple of (audio_bytes, global_ids)
   512	            - audio_bytes: Raw PCM audio (int16, 16kHz)
   513	            - global_ids: List of 32 global token IDs (populated after first decode, else empty)
   514
   515	        Thread Safety:
   516	            This method is stateless and thread-safe. Each call creates its own
   517	            request-scoped state. No global state is modified or shared.
   518	        """
   519	        import time
   520	        t_start = time.time()
   521
   522	        # Build prompt using Indic prompt builder
   523	        prompt = self.prompt_builder.build_prefix(speaker, text)
   524
   525	        # Configure sampling
   526	        # OPTIMIZATION: Use Spark TTS stop token, not legacy SNAC CODE_END_TOKEN_ID
   527	        sampling_params = SamplingParams(
   528	            temperature=temperature,
   529	            top_k=top_k,
   530	            top_p=top_p,
   531	            max_tokens=max_tokens,
   532	            min_tokens=DEFAULT_MIN_TOKENS,
   533	            repetition_penalty=repetition_penalty,
   534	            stop=TRAINING_STOP_TOKEN_IDS,  # "<|im_end|>" - matches non-streaming pipeline
   535	            skip_special_tokens=False,
   536	            seed=seed,
   537	            output_kind=(RequestOutputKind.DELTA if _STREAM_USE_DELTA_OUTPUT else RequestOutputKind.CUMULATIVE),
   538	        )
   539
   540	        # BiCodec token buffers
   541	        semantic_buffer = []
   542	        global_buffer = []
   543	        processed_token_count = 0  # Used for cumulative mode; ignored for delta mode.
   544	        total_audio_chunks = 0
   545
   546	        # BiCodec streaming configuration
   547	        EXPECTED_GLOBAL_COUNT = 32
   548	        MIN_SEMANTIC_FOR_FIRST_CHUNK = 10  # OPTIMIZATION: Lowered from 16 (decoder min is 8)
   549	        DECODE_INTERVAL = 48  # OPTIMIZATION: Increased from 24 to reduce O(n) re-decode calls
   550	        CROSSFADE_MS = 50
   551
   552	        # Generate unique request ID
   553	        import uuid
   554	        request_id = f"bicodec-first-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"
   555
   556	        # Reuse pre-warmed token parser (singleton, created once in __init__)
   557	        token_parser = self.token_parser
   558
   559	        results_generator = self.model.engine.generate(
   560	            prompt=prompt,
   561	            sampling_params=sampling_params,
   562	            request_id=request_id,
   563	        )
   564
   565	        # Track state
   566	        last_decode_count = 0
   567	        total_samples_emitted_to_user = 0
   568	        previous_chunk_tail: Optional[bytes] = None
   569	        captured_globals: List[int] = []  # Will be populated once we have 32 globals
   570
   571	        # Stream tokens with BiCodec extraction
   572	        async for request_output in results_generator:
   573	            generated_ids = request_output.outputs[0].token_ids
   574
   575	            # In DELTA mode token_ids are already incremental.
   576	            if _STREAM_USE_DELTA_OUTPUT:
   577	                new_token_ids = generated_ids
   578	            else:
   579	                new_token_ids = generated_ids[processed_token_count:]
   580	                processed_token_count = len(generated_ids)
   581	            token_parser.parse_incremental(new_token_ids, semantic_buffer, global_buffer)
   582
   583	            # Capture global tokens once we have all 32
   584	            if not captured_globals and len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
   585	                captured_globals = global_buffer[:EXPECTED_GLOBAL_COUNT]
   586
   587	            # Process if we have all global tokens
   588	            if len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
   589	                semantic_count = len(semantic_buffer)
   590	                if semantic_count >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
   591	                    if semantic_count - last_decode_count >= DECODE_INTERVAL:
   592	                        last_decode_count = semantic_count
   593
   594	                        all_semantic = semantic_buffer
   595	                        window_global = global_buffer[:EXPECTED_GLOBAL_COUNT]
   596
   597	                        if self.snac_decoder.enable_batching:
   598	                            audio_bytes = await self.snac_decoder.decode_single_async(
   599	                                semantic_ids=all_semantic,
   600	                                global_ids=window_global,
   601	                                trim_warmup=False,
   602	                                use_sliding_window=False
   603	                            )
   604	                        else:
   605	                            audio_bytes = self.snac_decoder.decode_streaming(
   606	                                semantic_ids=all_semantic,
   607	                                global_ids=window_global,
   608	                                use_sliding_window=False,
   609	                                trim_warmup=False
   610	                            )
   611
   612	                        if audio_bytes:
   613	                            total_samples_decoded = len(audio_bytes) // 2
   614
   615	                            if total_samples_decoded > total_samples_emitted_to_user:
   616	                                new_bytes_start = total_samples_emitted_to_user * 2
   617	                                new_audio_bytes = audio_bytes[new_bytes_start:]
   618
   619	                                to_emit, previous_chunk_tail = crossfade_bytes_int16(
   620	                                    previous_chunk_tail,


**Requesting continuation with numbering**


cd /home/ubuntu/veenaModal && nl -ba veena3modal/core/streaming_pipeline.py | sed -n '620,980p'