Output:
Chunk ID: 62a6a1
Wall time: 0.0504 seconds
Process exited with code 0
Original token count: 1948
Output:
   780	                "text_chunked": True,
   781	                "captured_global_count": len(captured_globals),
   782	                "semantic_tokens": len(semantic_buffer),
   783	                "generation_ms": int((t_end - t_start) * 1000),
   784	            }
   785	        )
   786	        logger.info("tts_stream_first_chunk_complete", extra=request_metrics)
   787
   788	    async def generate_speech_stream_indic_continuation(
   789	        self,
   790	        speaker: str,
   791	        text: str,
   792	        global_ids: List[int],
   793	        temperature: float = DEFAULT_TEMPERATURE,
   794	        top_k: int = DEFAULT_TOP_K,
   795	        top_p: float = DEFAULT_TOP_P,
   796	        max_tokens: int = DEFAULT_MAX_TOKENS,
   797	        repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
   798	        seed: Optional[int] = None,
   799	        request_id: Optional[str] = None,
   800	        request_metrics: Optional[Dict[str, Any]] = None,
   801	    ) -> AsyncGenerator[bytes, None]:
   802	        """
   803	        Generate speech for CONTINUATION chunks using pre-captured global tokens.
   804
   805	        This is the key method for voice consistency in chunked generation:
   806	        - Uses global_ids from the first chunk to maintain identical voice
   807	        - Model generates only semantic tokens (skips global token generation)
   808	        - Ensures no voice drift across chunks
   809
   810	        CRITICAL for production:
   811	        - global_ids is request-scoped, passed explicitly (no shared state)
   812	        - Thread-safe: each request has its own state
   813	        - No global caching that could cause cross-request contamination
   814
   815	        Args:
   816	            speaker: Speaker name (must match first chunk)
   817	            text: Continuation text chunk to synthesize
   818	            global_ids: 32 global token IDs captured from first chunk
   819	            temperature: Sampling temperature
   820	            top_k: Top-k sampling
   821	            top_p: Nucleus sampling
   822	            max_tokens: Max tokens to generate
   823	            repetition_penalty: Prevent repetition
   824	            seed: Random seed (use same as first chunk for consistency)
   825
   826	        Yields:
   827	            Audio bytes (int16 PCM, 16kHz mono)
   828
   829	        Raises:
   830	            ValueError: If global_ids doesn't contain exactly 32 tokens
   831
   832	        Thread Safety:
   833	            Fully thread-safe. All state is request-scoped and passed explicitly.
   834	        """
   835	        import statistics
   836	        import time
   837	        t_start = time.time()
   838	        request_metrics = request_metrics or {}
   839	        request_metrics.setdefault("request_id", request_id)
   840
   841	        # Validate global tokens
   842	        EXPECTED_GLOBAL_COUNT = 32
   843	        if len(global_ids) != EXPECTED_GLOBAL_COUNT:
   844	            raise ValueError(
   845	                f"Expected exactly {EXPECTED_GLOBAL_COUNT} global tokens, got {len(global_ids)}. "
   846	                f"global_ids must be captured from first chunk generation."
   847	            )
   848
   849	        # Build prompt WITH pre-filled global tokens
   850	        # This tells the model to skip global generation and go straight to semantic tokens
   851	        prompt = self.prompt_builder.build_prefix_with_globals(speaker, text, global_ids)
   852
   853	        # Configure sampling
   854	        # OPTIMIZATION: Use Spark TTS stop token, not legacy SNAC CODE_END_TOKEN_ID
   855	        sampling_params = SamplingParams(
   856	            temperature=temperature,
   857	            top_k=top_k,
   858	            top_p=top_p,
   859	            max_tokens=max_tokens,
   860	            min_tokens=DEFAULT_MIN_TOKENS,
   861	            repetition_penalty=repetition_penalty,
   862	            stop=TRAINING_STOP_TOKEN_IDS,  # "<|im_end|>" - matches non-streaming pipeline
   863	            skip_special_tokens=False,
   864	            seed=seed,
   865	        )
   866
   867	        # BiCodec token buffers
   868	        semantic_buffer = []
   869	        global_buffer_unused = []  # Not used in continuation, but needed for parser API
   870	        processed_token_count = 0  # Track how many vLLM tokens we've parsed
   871	        total_audio_chunks = 0
   872
   873	        # BiCodec streaming configuration
   874	        MIN_SEMANTIC_FOR_FIRST_CHUNK = 10  # OPTIMIZATION: Lowered from 16 (decoder min is 8)
   875	        DECODE_INTERVAL = 48  # OPTIMIZATION: Increased from 24 to reduce O(n) re-decode calls
   876	        CROSSFADE_MS = 50
   877
   878	        # Generate unique request ID
   879	        import uuid
   880	        if request_id is None:
   881	            request_id = f"bicodec-cont-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"
   882	            request_metrics["request_id"] = request_id
   883
   884	        # Reuse pre-warmed token parser (singleton, created once in __init__)
   885	        token_parser = self.token_parser
   886	        if token_parser is None:
   887	            raise RuntimeError("BiCodecTokenParser unavailable for continuation streaming")
   888
   889	        results_generator = self.model.engine.generate(
   890	            prompt=prompt,
   891	            sampling_params=sampling_params,
   892	            request_id=request_id,
   893	        )
   894
   895	        # Track state
   896	        last_decode_count = 0
   897	        total_samples_emitted_to_user = 0
   898	        previous_chunk_tail: Optional[bytes] = None
   899
   900	        # Stream tokens - model will generate ONLY semantic tokens (globals are pre-filled)
   901	        async for request_output in results_generator:
   902	            generated_ids = request_output.outputs[0].token_ids
   903
   904	            # OPTIMIZATION: Incremental parsing (O(1) per new token)
   905	            # Note: In continuation, model only generates semantic tokens (globals pre-filled in prompt)
   906	            new_token_ids = generated_ids[processed_token_count:]
   907	            processed_token_count = len(generated_ids)
   908	            token_parser.parse_incremental(new_token_ids, semantic_buffer, global_buffer_unused)
   909
   910	            semantic_count = len(semantic_buffer)
   911	            if semantic_count >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
   912	                if semantic_count - last_decode_count >= DECODE_INTERVAL:
   913	                    last_decode_count = semantic_count
   914
   915	                    # Use pre-captured global tokens for decoding
   916	                    if self.snac_decoder.enable_batching:
   917	                        audio_bytes = await self.snac_decoder.decode_single_async(
   918	                            semantic_ids=semantic_buffer,
   919	                            global_ids=global_ids,  # Use captured globals
   920	                            trim_warmup=False,
   921	                            use_sliding_window=False
   922	                        )
   923	                    else:
   924	                        audio_bytes = self.snac_decoder.decode_streaming(
   925	                            semantic_ids=semantic_buffer,
   926	                            global_ids=global_ids,  # Use captured globals
   927	                            use_sliding_window=False,
   928	                            trim_warmup=False
   929	                        )
   930


# Collaboration Mode: Default

You are now in Default mode. Any previous instructions for other modes (e.g. Plan mode) are no longer active.

Your active mode changes only when new developer instructions with a different `<collaboration_mode>...</collaboration_mode>` change it; user requests or tool descriptions do not change mode by themselves. Known mode names are Default and Plan.

## request_user_input availability

The `request_user_input` tool is unavailable in Default mode. If you call it while in Default mode, it will return an error.

If a decision is necessary and cannot be discovered from local context, ask the user directly. However, in Default mode you should strongly prefer executing the user's request rather than stopping to ask questions.

Summary: auto

## Skills
A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list of skills that can be used. Each entry includes a name, description, and file path so you can open the source for full instructions when using a specific skill.
### Available skills
- skill-creator: Guide for creating effective skills. This skill should be used when users want to create a new skill (or update an existing skill) that extends Codex's capabilities with specialized knowledge, workflows, or tool integrations. (file: /home/ubuntu/.codex/skills/.system/skill-creator/SKILL.md)
- skill-installer: Install Codex skills into $CODEX_HOME/skills from a curated list or a GitHub repo path. Use when a user asks to list installable skills, install a curated skill, or install a skill from another repo (including private repos). (file: /home/ubuntu/.codex/skills/.system/skill-installer/SKILL.md)
### How to use skills
- Discovery: The list above is the skills available in this session (name + description + file path). Skill bodies live on disk at the listed paths.
- Trigger rules: If the user names a skill (with `$SkillName` or plain text) OR the task clearly matches a skill's description shown above, you must use that skill for that turn. Multiple mentions mean use them all. Do not carry skills across turns unless re-mentioned.
- Missing/blocked: If a named skill isn't in the list or the path can't be read, say so briefly and continue with the best fallback.
- How to use a skill (progressive disclosure):
  1) After deciding to use a skill, open its `SKILL.md`. Read only enough to follow the workflow.
  2) When `SKILL.md` references relative paths (e.g., `scripts/foo.py`), resolve them relative to the skill directory listed above first, and only consider other paths if needed.
  3) If `SKILL.md` points to extra folders such as `references/`, load only the specific files needed for the request; don't bulk-load everything.
  4) If `scripts/` exist, prefer running or patching them instead of retyping large code blocks.
  5) If `assets/` or templates exist, reuse them instead of recreating from scratch.
- Coordination and sequencing:
  - If multiple skills apply, choose the minimal set that covers the request and state the order you'll use them.
  - Announce which skill(s) you're using and why (one short line). If you skip an obvious skill, say why.
- Context hygiene:
  - Keep context small: summarize long sections instead of pasting them; only load extra files when needed.
  - Avoid deep reference-chasing: prefer opening only files directly linked from `SKILL.md` unless you're blocked.
  - When variants exist (frameworks, providers, domains), pick only the relevant reference file(s) and note that choice.
- Safety and fallback: If a skill can't be applied cleanly (missing files, unclear instructions), state the issue, pick the next-best approach, and continue.


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


Output:
Chunk ID: bb4b74
Wall time: 0.0508 seconds
Process exited with code 0
Original token count: 995
Output:
   930
   931	                    if audio_bytes:
   932	                        total_samples_decoded = len(audio_bytes) // 2
   933
   934	                        if total_samples_decoded > total_samples_emitted_to_user:
   935	                            new_bytes_start = total_samples_emitted_to_user * 2
   936	                            new_audio_bytes = audio_bytes[new_bytes_start:]
   937
   938	                            to_emit, previous_chunk_tail = crossfade_bytes_int16(
   939	                                previous_chunk_tail,
   940	                                new_audio_bytes,
   941	                                sample_rate_hz=16000,
   942	                                crossfade_ms=CROSSFADE_MS,
   943	                            )
   944
   945	                            samples_emitted_in_this_chunk = len(to_emit) // 2
   946	                            total_samples_emitted_to_user += samples_emitted_in_this_chunk
   947
   948	                            if to_emit:
   949	                                total_audio_chunks += 1
   950	                                yield to_emit
   951
   952	        # Final chunk processing
   953	        if len(semantic_buffer) > last_decode_count:
   954	            audio_bytes = self.snac_decoder.decode_streaming(
   955	                semantic_ids=semantic_buffer,
   956	                global_ids=global_ids,  # Use captured globals
   957	                use_sliding_window=False,
   958	                trim_warmup=False
   959	            )
   960	            if audio_bytes:
   961	                total_samples_decoded = len(audio_bytes) // 2
   962
   963	                if total_samples_decoded > total_samples_emitted_to_user:
   964	                    new_bytes_start = total_samples_emitted_to_user * 2
   965	                    new_audio_bytes = audio_bytes[new_bytes_start:]
   966
   967	                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
   968	                        previous_chunk_tail,
   969	                        new_audio_bytes,
   970	                        sample_rate_hz=16000,
   971	                        crossfade_ms=CROSSFADE_MS,
   972	                    )
   973
   974	                    samples_emitted_in_final = len(to_emit) // 2
   975	                    total_samples_emitted_to_user += samples_emitted_in_final
   976
   977	                    if to_emit:
   978	                        total_audio_chunks += 1
   979	                        yield to_emit
   980
   981	        t_end = time.time()
   982	        request_metrics.update(
   983	            {
   984	                "request_id": request_id,
   985	                "speaker": speaker,
   986	                "audio_duration_seconds": total_samples_emitted_to_user / 16000,
   987	                "audio_chunks": total_audio_chunks,
   988	                "semantic_tokens": len(semantic_buffer),
   989	                "global_tokens": len(global_ids),
   990	                "generation_ms": int((t_end - t_start) * 1000),
   991	            }
   992	        )
   993	        logger.info("tts_stream_continuation_complete", extra=request_metrics)
   994
   995	        # Flush remaining tail
   996	        if previous_chunk_tail:
   997	            total_audio_chunks += 1
   998	            tail_samples = len(previous_chunk_tail) // 2
   999	            total_samples_emitted_to_user += tail_samples
  1000	            yield previous_chunk_tail
  1001
  1002	        t_end = time.time()
  1003	        audio_duration_s = total_samples_emitted_to_user / 16000
  1004	        print(f"🎵 Continuation chunk complete: {audio_duration_s:.2f}s audio, {total_audio_chunks} chunks (using cached globals)")


Summary: auto


Input:
*** Begin Patch
*** Update File: /home/ubuntu/veenaModal/veena3modal/core/streaming_pipeline.py
@@
         if len(semantic_buffer) > last_decode_count:
             audio_bytes = self.snac_decoder.decode_streaming(
                 semantic_ids=semantic_buffer,
                 global_ids=global_ids,  # Use captured globals
                 use_sliding_window=False,
                 trim_warmup=False
             )
@@
                         if to_emit:
                             total_audio_chunks += 1
                             yield to_emit
-
-        t_end = time.time()
-        request_metrics.update(
-            {
-                "request_id": request_id,
-                "speaker": speaker,
-                "audio_duration_seconds": total_samples_emitted_to_user / 16000,
-                "audio_chunks": total_audio_chunks,
-                "semantic_tokens": len(semantic_buffer),
-                "global_tokens": len(global_ids),
-                "generation_ms": int((t_end - t_start) * 1000),
-            }
-        )
-        logger.info("tts_stream_continuation_complete", extra=request_metrics)
-
+
         # Flush remaining tail
         if previous_chunk_tail:
             total_audio_chunks += 1
             tail_samples = len(previous_chunk_tail) // 2
             total_samples_emitted_to_user += tail_samples
             yield previous_chunk_tail
-
-        t_end = time.time()
-        audio_duration_s = total_samples_emitted_to_user / 16000
-        print(f"🎵 Continuation chunk complete: {audio_duration_s:.2f}s audio, {total_audio_chunks} chunks (using cached globals)")
+
+        t_end = time.time()
+        request_metrics.update(
+            {
+                "request_id": request_id,
+                "speaker": speaker,
+                "audio_duration_seconds": total_samples_emitted_to_user / 16000,
+                "audio_chunks": total_audio_chunks,
+                "semantic_tokens": len(semantic_buffer),
+                "global_tokens": len(global_ids),
+                "generation_ms": int((t_end - t_start) * 1000),
+            }
+        )
+        logger.info("tts_stream_continuation_complete", extra=request_metrics)
*** End Patch


Output: {"output":"Success. Updated the following files:\nM /home/ubuntu/veenaModal/veena3modal/core/streaming_pipeline.py\n","metadata":{"exit_code":0,"duration_seconds":0.0}}


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


Output:
Chunk ID: a83aa0
Wall time: 0.0509 seconds
Process exited with code 0
Original token count: 1285
Output:
   640
   641	        # BiCodec token buffers
   642	        semantic_buffer = []
   643	        global_buffer = []
   644	        processed_token_count = 0  # Track how many vLLM tokens we've parsed
   645	        total_audio_chunks = 0
   646
   647	        # BiCodec streaming configuration
   648	        EXPECTED_GLOBAL_COUNT = 32
   649	        MIN_SEMANTIC_FOR_FIRST_CHUNK = 10  # OPTIMIZATION: Lowered from 16 (decoder min is 8)
   650	        DECODE_INTERVAL = 48  # OPTIMIZATION: Increased from 24 to reduce O(n) re-decode calls
   651	        CROSSFADE_MS = 50
   652
   653	        # Generate unique request ID
   654	        import uuid
   655	        if request_id is None:
   656	            request_id = f"bicodec-first-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"
   657	            request_metrics["request_id"] = request_id
   658
   659	        # Reuse pre-warmed token parser (singleton, created once in __init__)
   660	        token_parser = self.token_parser
   661	        if token_parser is None:
   662	            raise RuntimeError("BiCodecTokenParser unavailable for first-chunk streaming")
   663
   664	        results_generator = self.model.engine.generate(
   665	            prompt=prompt,
   666	            sampling_params=sampling_params,
   667	            request_id=request_id,
   668	        )
   669
   670	        # Track state
   671	        last_decode_count = 0
   672	        total_samples_emitted_to_user = 0
   673	        previous_chunk_tail: Optional[bytes] = None
   674	        captured_globals: List[int] = []  # Will be populated once we have 32 globals
   675
   676	        # Stream tokens with BiCodec extraction
   677	        async for request_output in results_generator:
   678	            generated_ids = request_output.outputs[0].token_ids
   679
   680	            # OPTIMIZATION: Incremental parsing (O(1) per new token)
   681	            new_token_ids = generated_ids[processed_token_count:]
   682	            processed_token_count = len(generated_ids)
   683	            token_parser.parse_incremental(new_token_ids, semantic_buffer, global_buffer)
   684
   685	            # Capture global tokens once we have all 32
   686	            if not captured_globals and len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
   687	                captured_globals = global_buffer[:EXPECTED_GLOBAL_COUNT]
   688
   689	            # Process if we have all global tokens
   690	            if len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
   691	                semantic_count = len(semantic_buffer)
   692	                if semantic_count >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
   693	                    if semantic_count - last_decode_count >= DECODE_INTERVAL:
   694	                        last_decode_count = semantic_count
   695
   696	                        all_semantic = semantic_buffer
   697	                        window_global = global_buffer[:EXPECTED_GLOBAL_COUNT]
   698
   699	                        if self.snac_decoder.enable_batching:
   700	                            audio_bytes = await self.snac_decoder.decode_single_async(
   701	                                semantic_ids=all_semantic,
   702	                                global_ids=window_global,
   703	                                trim_warmup=False,
   704	                                use_sliding_window=False
   705	                            )
   706	                        else:
   707	                            audio_bytes = self.snac_decoder.decode_streaming(
   708	                                semantic_ids=all_semantic,
   709	                                global_ids=window_global,
   710	                                use_sliding_window=False,
   711	                                trim_warmup=False
   712	                            )
   713
   714	                        if audio_bytes:
   715	                            total_samples_decoded = len(audio_bytes) // 2
   716
   717	                            if total_samples_decoded > total_samples_emitted_to_user:
   718	                                new_bytes_start = total_samples_emitted_to_user * 2
   719	                                new_audio_bytes = audio_bytes[new_bytes_start:]
   720
   721	                                to_emit, previous_chunk_tail = crossfade_bytes_int16(
   722	                                    previous_chunk_tail,
   723	                                    new_audio_bytes,
   724	                                    sample_rate_hz=16000,
   725	                                    crossfade_ms=CROSSFADE_MS,
   726	                                )
   727
   728	                                samples_emitted_in_this_chunk = len(to_emit) // 2
   729	                                total_samples_emitted_to_user += samples_emitted_in_this_chunk
   730


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


Output:
Chunk ID: 094d2b
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 1004
Output:
   560	                "generation_ms": int(total_wall_ms),
   561	                "audio_duration_seconds": audio_duration_s,
   562	                "audio_chunks": total_audio_chunks,
   563	                "samples_emitted": total_samples_emitted_to_user,
   564	                "text_chunked": False,
   565	                "concurrency_mode": "single",
   566	            }
   567	        )
   568	        if emit_timing:
   569	            logger.info("tts_stream_indic_complete", extra=request_metrics)
   570
   571	    async def generate_speech_stream_indic_first_chunk(
   572	        self,
   573	        speaker: str,
   574	        text: str,
   575	        temperature: float = DEFAULT_TEMPERATURE,
   576	        top_k: int = DEFAULT_TOP_K,
   577	        top_p: float = DEFAULT_TOP_P,
   578	        max_tokens: int = DEFAULT_MAX_TOKENS,
   579	        repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
   580	        seed: Optional[int] = None,
   581	        request_id: Optional[str] = None,
   582	        request_metrics: Optional[Dict[str, Any]] = None,
   583	    ) -> AsyncGenerator[tuple, None]:
   584	        """
   585	        Generate speech for the FIRST chunk of multi-chunk text, capturing global tokens.
   586
   587	        This method is specifically for chunked generation: it yields (audio_bytes, global_ids)
   588	        tuples instead of just audio_bytes. The caller captures global_ids from the first
   589	        yield and passes them to generate_speech_stream_indic_continuation() for subsequent chunks.
   590
   591	        Use Case (chunked streaming with voice consistency):
   592	            globals_captured = None
   593	            async for audio_bytes, global_ids in pipeline.generate_speech_stream_indic_first_chunk(...):
   594	                if globals_captured is None and global_ids:
   595	                    globals_captured = global_ids  # Capture once
   596	                yield audio_bytes
   597	            # Now use globals_captured for subsequent chunks
   598
   599	        Args:
   600	            speaker: Speaker name
   601	            text: First text chunk to synthesize
   602	            temperature: Sampling temperature
   603	            top_k: Top-k sampling
   604	            top_p: Nucleus sampling
   605	            max_tokens: Max tokens to generate
   606	            repetition_penalty: Prevent repetition
   607	            seed: Random seed for reproducibility
   608
   609	        Yields:
   610	            Tuple of (audio_bytes, global_ids)
   611	            - audio_bytes: Raw PCM audio (int16, 16kHz)
   612	            - global_ids: List of 32 global token IDs (populated after first decode, else empty)
   613
   614	        Thread Safety:
   615	            This method is stateless and thread-safe. Each call creates its own
   616	            request-scoped state. No global state is modified or shared.
   617	        """
   618	        import time
   619	        import statistics
   620	        t_start = time.time()
   621	        request_metrics = request_metrics or {}
   622	        request_metrics.setdefault("request_id", request_id)
   623
   624	        # Build prompt using Indic prompt builder
   625	        prompt = self.prompt_builder.build_prefix(speaker, text)
   626
   627	        # Configure sampling
   628	        # OPTIMIZATION: Use Spark TTS stop token, not legacy SNAC CODE_END_TOKEN_ID
   629	        sampling_params = SamplingParams(
   630	            temperature=temperature,
   631	            top_k=top_k,
   632	            top_p=top_p,
   633	            max_tokens=max_tokens,
   634	            min_tokens=DEFAULT_MIN_TOKENS,
   635	            repetition_penalty=repetition_penalty,
   636	            stop=TRAINING_STOP_TOKEN_IDS,  # "<|im_end|>" - matches non-streaming pipeline
   637	            skip_special_tokens=False,
   638	            seed=seed,
   639	        )
   640


Input:
*** Begin Patch
*** Update File: /home/ubuntu/veenaModal/veena3modal/core/streaming_pipeline.py
@@
-        import time
-        import statistics
+        import time
@@
-        import statistics
         import time
*** End Patch


Output: {"output":"Success. Updated the following files:\nM /home/ubuntu/veenaModal/veena3modal/core/streaming_pipeline.py\n","metadata":{"exit_code":0,"duration_seconds":0.0}}


cd /home/ubuntu/veenaModal && rg -n "Continuation chunk complete|First chunk complete|🎵" veena3modal/core/streaming_pipeline.py


Output:
Chunk ID: a1e0e8
Wall time: 0.0511 seconds
Process exited with code 0
Original token count: 29
Output:
198:                                print(f"🎵 First chunk decoded ({len(audio_bytes)} bytes, sliding window)")


Input:
*** Begin Patch
*** Update File: /home/ubuntu/veenaModal/veena3modal/services/tts_runtime.py
@@
     if not is_initialized():
         raise RuntimeError("TTS runtime not initialized")
@@
     if not is_initialized():
         raise RuntimeError("TTS runtime not initialized")
@@
     runtime = get_runtime()
+    if request_id is None:
+        import uuid
+        request_id = f"runtime-chunk-{uuid.uuid4().hex[:10]}"
     metrics = {
+        "request_id": request_id,
         "ttfb_ms": 0,
         "generation_ms": 0,
         "chunks_processed": 0,
         "audio_duration_seconds": 0.0,
         "text_chunked": False,
@@
-    active_requests = _with_inflight_delta(1)
+    active_requests = _with_inflight_delta(1)
     start_time = time.time()

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

         long_processor = LongTextProcessor(pipeline=runtime.pipeline)

         # Check if chunking is needed
         if long_processor.should_chunk(text):
             metrics["text_chunked"] = True
-            audio_bytes = await long_processor.generate_with_chunking(
-                text=text,
-                speaker=speaker,
-                temperature=temperature,
-                top_k=top_k,
-                top_p=top_p,
-                max_tokens=max_tokens,
-                repetition_penalty=repetition_penalty,
-                seed=seed,
-                sample_rate=sample_rate,
-            )
+            if _timing_enabled():
+                chunk_result = await long_processor.generate_with_chunking(
+                    text=text,
+                    speaker=speaker,
+                    temperature=temperature,
+                    top_k=top_k,
+                    top_p=top_p,
+                    max_tokens=max_tokens,
+                    repetition_penalty=repetition_penalty,
+                    seed=seed,
+                    sample_rate=sample_rate,
+                    request_id=request_id,
+                    return_timing=True,
+                )
+                if isinstance(chunk_result, tuple):
+                    audio_bytes, chunk_timing = chunk_result
+                else:
+                    audio_bytes = chunk_result
+                    chunk_timing = {}
+            else:
+                audio_bytes = await long_processor.generate_with_chunking(
+                    text=text,
+                    speaker=speaker,
+                    temperature=temperature,
+                    top_k=top_k,
+                    top_p=top_p,
+                    max_tokens=max_tokens,
+                    repetition_penalty=repetition_penalty,
+                    seed=seed,
+                    sample_rate=sample_rate,
+                    request_id=request_id,
+                )
+                chunk_timing = {}
         else:
             # Short text: direct generation
-            audio_bytes = await runtime.pipeline.generate_speech_indic(
-                speaker=speaker,
-                text=text,
-                temperature=temperature,
-                top_k=top_k,
-                top_p=top_p,
-                max_tokens=max_tokens,
-                repetition_penalty=repetition_penalty,
-                seed=seed,
-            )
+            gen_kwargs = dict(
+                speaker=speaker,
+                text=text,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                repetition_penalty=repetition_penalty,
+                seed=seed,
+                request_id=request_id,
+            )
+            if _timing_enabled():
+                audio_bytes, pipeline_timings = await runtime.pipeline.generate_speech_indic(
+                    **gen_kwargs,
+                    return_timing=True,
+                )
+                chunk_timing = pipeline_timings
+            else:
+                audio_bytes = await runtime.pipeline.generate_speech_indic(
+                    **gen_kwargs,
+                )
+                chunk_timing = {}

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

+        if chunk_timing:
+            metrics.update({
+                "llm_token_total": chunk_timing.get("token_total", 0),
+                "llm_batch_count": chunk_timing.get("batch_count", 0),
+                "llm_batch_gpu_ms_total": chunk_timing.get("llm_gpu_batch_ms_total", 0),
+                "llm_batch_wall_ms_total": chunk_timing.get("llm_wall_ms", metrics["generation_ms"]),
+                "llm_parse_ms": chunk_timing.get("parse_ms", 0),
+                "llm_decode_cpu_ms": chunk_timing.get("decode_cpu_ms", 0),
+                "llm_batch_wall_ms_min": chunk_timing.get("llm_batch_wall_ms_min", 0),
+                "llm_batch_wall_ms_max": chunk_timing.get("llm_batch_wall_ms_max", 0),
+                "llm_batch_gpu_ms_min": chunk_timing.get("llm_batch_gpu_ms_min", 0),
+                "llm_batch_gpu_ms_max": chunk_timing.get("llm_batch_gpu_ms_max", 0),
+                "llm_batch_gpu_ms_p50": chunk_timing.get("llm_batch_gpu_ms_p50", 0),
+                "tokens_per_batch_min": chunk_timing.get("tokens_per_batch_min"),
+                "tokens_per_batch_max": chunk_timing.get("tokens_per_batch_max"),
+                "tokens_per_batch_p50": chunk_timing.get("tokens_per_batch_p50"),
+                "llm_batch_tokens": chunk_timing.get("token_batches", []),
+            })
+            if metrics["text_chunked"]:
+                metrics["chunks_processed"] = chunk_timing.get("chunks_processed", 0)
+
         if audio_bytes:
@@
         if _timing_enabled():
             metrics["active_requests"] = active_requests - 1
+            metrics["request_inflight"] = active_requests - 1
             logger.info("tts_runtime_chunking_complete", extra=metrics)
         return audio_bytes, metrics

     except Exception as e:
         logger.error(f"Chunked speech generation failed: {e}")
         raise
*** End Patch


Output: {"output":"Success. Updated the following files:\nM /home/ubuntu/veenaModal/veena3modal/services/tts_runtime.py\n","metadata":{"exit_code":0,"duration_seconds":0.0}}


Input:
*** Begin Patch
*** Update File: /home/ubuntu/veenaModal/veena3modal/services/tts_runtime.py
@@
     runtime = get_runtime()
@@
     # Metrics tracking
     metrics = {
+        "request_id": request_id,
         "ttfb_ms": 0,
         "chunks_sent": 0,
         "total_bytes": 0,
         "audio_duration_seconds": 0.0,
         "text_chunked": False,
@@
     active_requests = _with_inflight_delta(1)
     if _timing_enabled():
         logger.info(
@@
                     max_tokens,
                     repetition_penalty,
                     seed=seed,
                     sample_rate,
                     request_id,
+                    metrics,
                 ):
                     if first_chunk_time is None:
@@
                     request_id=request_id,
                     top_k=top_k,
                     top_p=top_p,
                     max_tokens=max_tokens,
                     repetition_penalty=repetition_penalty,
                     seed=seed,
                     request_id=request_id,
+                    request_metrics=metrics,
                 )
             else:
                 stream_gen = runtime.streaming_pipeline.generate_speech_stream_indic(
                     speaker=speaker,
                     text=text,
                     temperature=temperature,
                     top_k=top_k,
                     top_p=top_p,
                     max_tokens=max_tokens,
                     repetition_penalty=repetition_penalty,
                     seed=seed,
                     request_id=request_id,
+                    request_metrics=metrics,
                 )
@@
     finally:
         active_requests = _with_inflight_delta(-1)
         if _timing_enabled():
             logger.info(
                 "tts_runtime_stream_end",
*** End Patch


Output:
apply_patch verification failed: Failed to find expected lines in /home/ubuntu/veenaModal/veena3modal/services/tts_runtime.py:
                    max_tokens,
                    repetition_penalty,
                    seed=seed,
                    sample_rate,
                    request_id,
                ):
                    if first_chunk_time is None:


cd /home/ubuntu/veenaModal && nl -ba veena3modal/services/tts_runtime.py | sed -n '740,940p'


Output:
Chunk ID: bb7e1a
Wall time: 0.0508 seconds
Process exited with code 0
Original token count: 2403
Output:
   740	                "tokens_per_batch_max": chunk_timing.get("tokens_per_batch_max"),
   741	                "tokens_per_batch_p50": chunk_timing.get("tokens_per_batch_p50"),
   742	                "llm_batch_tokens": chunk_timing.get("token_batches", []),
   743	            })
   744	            if metrics["text_chunked"]:
   745	                metrics["chunks_processed"] = chunk_timing.get("chunks_processed", 0)
   746
   747	        if audio_bytes:
   748	            # Apply super-resolution if requested and available
   749	            final_sample_rate = 16000
   750	            if output_sample_rate == "48khz" and runtime.sr_service and runtime.sr_service.is_loaded:
   751	                try:
   752	                    sr_start = time.time()
   753	                    logger.info(f"Applying super-resolution to chunked audio ({len(audio_bytes)} bytes)...")
   754	                    audio_bytes = _apply_super_resolution(audio_bytes, runtime.sr_service)
   755	                    sr_time = (time.time() - sr_start) * 1000
   756	                    metrics["sr_ms"] = int(sr_time)
   757	                    metrics["sr_applied"] = True
   758	                    final_sample_rate = 48000
   759	                    logger.info(f"✅ Super-resolution applied in {sr_time:.1f}ms, output={len(audio_bytes)} bytes")
   760	                except Exception as e:
   761	                    logger.warning(f"Super-resolution failed, returning 16kHz: {e}")
   762	                    import traceback
   763	                    traceback.print_exc()
   764
   765	            metrics["output_sample_rate"] = final_sample_rate
   766	            audio_duration = (len(audio_bytes) - 44) / (final_sample_rate * 2)
   767	            metrics["audio_duration_seconds"] = max(0.0, audio_duration)
   768
   769	        if _timing_enabled():
   770	            metrics["active_requests"] = active_requests - 1
   771	            metrics["request_inflight"] = active_requests - 1
   772	            logger.info("tts_runtime_chunking_complete", extra=metrics)
   773	        return audio_bytes, metrics
   774
   775	    except Exception as e:
   776	        logger.error(f"Chunked speech generation failed: {e}")
   777	        raise
   778	    finally:
   779	        _with_inflight_delta(-1)
   780
   781
   782	from typing import AsyncGenerator
   783
   784
   785	async def generate_speech_streaming(
   786	    text: str,
   787	    speaker: str,
   788	    temperature: float = 0.8,
   789	    top_k: int = 50,
   790	    top_p: float = 1.0,
   791	    max_tokens: int = 4096,
   792	    repetition_penalty: float = 1.05,
   793	    seed: Optional[int] = None,
   794	    enable_chunking: bool = True,
   795	    request_id: Optional[str] = None,
   796	) -> AsyncGenerator[Tuple[bytes, Dict[str, Any]], None]:
   797	    """
   798	    Generate speech audio with true streaming (yields chunks as they're generated).
   799
   800	    This is the core streaming implementation for M4.
   801	    First yield includes WAV header + first PCM chunk.
   802	    Subsequent yields are raw PCM chunks.
   803
   804	    Args:
   805	        text: Text to synthesize (already normalized)
   806	        speaker: Internal speaker name (resolved)
   807	        temperature: Sampling temperature
   808	        top_k: Top-k sampling
   809	        top_p: Nucleus sampling
   810	        max_tokens: Maximum tokens to generate
   811	        repetition_penalty: Repetition penalty
   812	        seed: Random seed for reproducibility
   813	        enable_chunking: Enable text chunking for long inputs with voice consistency
   814
   815	    Yields:
   816	        Tuple of (audio_bytes, metrics_dict)
   817	        - First yield: WAV header (44 bytes) prepended to first PCM chunk
   818	        - Subsequent yields: Raw PCM chunks (int16, 16kHz)
   819	        - Final yield: metrics_dict has final timing info
   820
   821	    Raises:
   822	        RuntimeError: If runtime not initialized or streaming pipeline unavailable
   823	    """
   824	    if not is_initialized():
   825	        raise RuntimeError("TTS runtime not initialized")
   826
   827	    runtime = get_runtime()
   828
   829	    if runtime.streaming_pipeline is None:
   830	        raise RuntimeError("Streaming pipeline not available")
   831
   832	    # Import audio utils for WAV header
   833	    from veena3modal.audio.utils import create_wav_header
   834
   835	    # Metrics tracking
   836	    metrics = {
   837	        "ttfb_ms": 0,
   838	        "chunks_sent": 0,
   839	        "total_bytes": 0,
   840	        "audio_duration_seconds": 0.0,
   841	        "text_chunked": False,
   842	    }
   843
   844	    start_time = time.time()
   845	    first_chunk_time = None
   846	    total_pcm_bytes = 0
   847	    sample_rate = 16000  # BiCodec sample rate
   848	    active_requests = _with_inflight_delta(1)
   849	    if _timing_enabled():
   850	        logger.info(
   851	            "tts_runtime_stream_start",
   852	            extra={
   853	                "speaker": speaker,
   854	                "text_length": len(text),
   855	                "request_inflight": active_requests,
   856	                "chunking": enable_chunking,
   857	                "request_id": request_id,
   858	            },
   859	        )
   860
   861	    try:
   862	        # Check if we need text chunking
   863	        from veena3modal.processing.long_text_processor import LongTextProcessor
   864	        long_processor = LongTextProcessor(pipeline=runtime.pipeline)
   865	        needs_chunking = enable_chunking and long_processor.should_chunk(text)
   866
   867	        if needs_chunking:
   868	            # Chunked streaming with voice consistency (global token caching)
   869	            metrics["text_chunked"] = True
   870	            async for audio_chunk, chunk_metrics in _stream_chunked_text(
   871	                runtime,
   872	                long_processor,
   873	                text,
   874	                speaker,
   875	                temperature,
   876	                top_k,
   877	                top_p,
   878	                max_tokens,
   879	                repetition_penalty,
   880	                seed,
   881	                sample_rate,
   882	                request_id,
   883	            ):
   884	                if first_chunk_time is None:
   885	                    first_chunk_time = time.time()
   886	                    metrics["ttfb_ms"] = int((first_chunk_time - start_time) * 1000)
   887	                    # Create streaming WAV header (size=0 for unknown length)
   888	                    wav_header = create_wav_header(sample_rate=sample_rate, data_size=0)
   889	                    audio_chunk = wav_header + audio_chunk
   890
   891	                total_pcm_bytes += len(audio_chunk) - (44 if metrics["chunks_sent"] == 0 else 0)
   892	                metrics["chunks_sent"] += 1
   893	                metrics["total_bytes"] = total_pcm_bytes + 44  # Include header
   894	                metrics["audio_duration_seconds"] = total_pcm_bytes / (sample_rate * 2)
   895	                if chunk_metrics:
   896	                    metrics.update(chunk_metrics)
   897
   898	                yield audio_chunk, metrics
   899	        else:
   900	            # Simple streaming (no text chunking)
   901	            # OPTIMIZATION: If we have pre-cached global tokens for this speaker,
   902	            # use continuation mode to skip the ~110ms global token pre-roll.
   903	            # The model jumps straight to semantic token generation.
   904	            cached_globals = runtime.speaker_global_cache.get(speaker)
   905
   906	            if cached_globals:
   907	                stream_gen = runtime.streaming_pipeline.generate_speech_stream_indic_continuation(
   908	                    speaker=speaker,
   909	                    text=text,
   910	                    global_ids=cached_globals,
   911	                    temperature=temperature,
   912	                    top_k=top_k,
   913	                    top_p=top_p,
   914	                    max_tokens=max_tokens,
   915	                    repetition_penalty=repetition_penalty,
   916	                    seed=seed,
   917	                    request_id=request_id,
   918	                )
   919	            else:
   920	                stream_gen = runtime.streaming_pipeline.generate_speech_stream_indic(
   921	                    speaker=speaker,
   922	                    text=text,
   923	                    temperature=temperature,
   924	                    top_k=top_k,
   925	                    top_p=top_p,
   926	                    max_tokens=max_tokens,
   927	                    repetition_penalty=repetition_penalty,
   928	                    seed=seed,
   929	                    request_id=request_id,
   930	                )
   931
   932	            async for audio_chunk in stream_gen:
   933	                if first_chunk_time is None:
   934	                    first_chunk_time = time.time()
   935	                    metrics["ttfb_ms"] = int((first_chunk_time - start_time) * 1000)
   936	                    wav_header = create_wav_header(sample_rate=sample_rate, data_size=0)
   937	                    audio_chunk = wav_header + audio_chunk
   938
   939	                total_pcm_bytes += len(audio_chunk) - (44 if metrics["chunks_sent"] == 0 else 0)
   940	                metrics["chunks_sent"] += 1


<permissions instructions>
Filesystem sandboxing defines which files can be read or written. `sandbox_mode` is `danger-full-access`: No filesystem sandboxing - all commands are permitted. Network access is enabled.
Approval policy is currently never. Do not provide the `sandbox_permissions` for any reason, commands will be rejected.
</permissions instructions>

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

# AGENTS.md instructions for /home/ubuntu/veenaModal

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

<environment_context>
  <cwd>/home/ubuntu/veenaModal</cwd>
  <shell>bash</shell>
</environment_context>

The current TTS application here is an implementation in Modal labs. Which has autoscaling and other rules self scaling rules and everything suitable for my target audience. Now, i need you to create a local entry point or implementation point. then we will plan for optimising the pipeline, then we will discuss on stress testing, then we will discuss on scaling laws post optimisations locally. @veena3modal/app.py

the above part is performed by a different agent. and benchmarks are made. Now can you do the benchmarks on this newly optimized path to understand the concurrency and stress testing.

Now i need you to capture detailed time logs for each request happening in the GPU. like what happens with 1req, time per LLM token, time per batch decoding, and everything else. detailed time about where we are spending each minute in the GPU per generation.

and when you have the details for per request. you can obviously add more logging to get details of concurrency. so you understand what's the bottleneck for the peak concurrency. and then we will decide on effective methods in optimising this TTS engine.


Need to finish this in two parts: first wire the remaining request-level timing/`request_id` propagation through runtime and API paths, then add a local benchmark harness that drives concurrency and emits per-request GPU/latency/batch/streaming metrics. I’ll inspect those files once, patch in a single pass, and then summarize how to run the stress test.


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


sed -n '1,340p' veena3modal/core/streaming_pipeline.py


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


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


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


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


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


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


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


Output: exec_command failed: CreateProcess { message: "Rejected(\"Failed to create unified exec process: No such file or directory (os error 2)\")" }


Output: exec_command failed: CreateProcess { message: "Rejected(\"Failed to create unified exec process: No such file or directory (os error 2)\")" }


Output: exec_command failed: CreateProcess { message: "Rejected(\"Failed to create unified exec process: No such file or directory (os error 2)\")" }


Output: exec_command failed: CreateProcess { message: "Rejected(\"Failed to create unified exec process: No such file or directory (os error 2)\")" }


Output: exec_command failed: CreateProcess { message: "Rejected(\"Failed to create unified exec process: No such file or directory (os error 2)\")" }


Output: exec_command failed: CreateProcess { message: "Rejected(\"Failed to create unified exec process: No such file or directory (os error 2)\")" }


Output:
Chunk ID: e7d051
Wall time: 0.0517 seconds
Process exited with code 0
Original token count: 2590
Output:
#!/usr/bin/env python3
"""
Local Stress Test for Veena3 TTS Server.

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

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

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

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

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

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

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

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

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

import httpx

# === Configuration ===

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

# Test texts at different lengths to stress different parts of the pipeline
TEST_TEXTS = {
    "short": "Hello, this is a quick test.",
    "medium": (
        "The quick brown fox jumps over the lazy dog. "
        "This sentence tests basic text normalization and generation quality "
        "with moderate length input that should complete in reasonable time."
    ),
    "long": (
        "In the heart of a bustling city, where towering skyscrapers cast long shadows "
        "over crowded streets, there lived an old bookkeeper named Margaret. She had spent "
        "forty years cataloging the stories of others, yet never found time to write her own. "
        "One rainy Tuesday afternoon, she discovered a leather-bound journal tucked behind "
        "a shelf of dusty encyclopedias. Its pages were blank, waiting. She picked up her pen "
        "and began to write. The words flowed like water from a broken dam."
    ),
}

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

# === GPU Monitoring ===

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

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

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

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

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

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

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

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

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

# === Request Result ===

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

# === Load Level Result ===

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

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

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

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

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

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

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

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

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

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

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

# === HTTP Client ===

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

    start = time.time()
    ttfb_time = 0.0

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

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

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

                    return RequestResult(
                        success=True,
                        status_code=200,
                        latency_ms=latency,
                        ttfb_ms=ttfb_time,
                        audio_bytes=total_bytes,
                        audio_seconds=audio_seconds,
                        speaker=speaker,
                        text_length=len(text),
                        stream=True,
                    )
                else:
                    return RequestResult(
                        success=False,
                        status_code=response.status_code,
                        latency_ms=latency,
                        ttfb_ms=ttfb_time,
                        error=f"HTTP {response.status_code}",
                        speaker=speaker,


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

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

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

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

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

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

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

from __future__ import annotations

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

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

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

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

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

# === Timing Utilities ===

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

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

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

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

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

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

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

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

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

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

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

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

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

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

# === Non-Streaming Profiler ===

async def profile_non_streaming(
    runtime,
    text: str,
    speaker: str = "lipakshi",
    temperature: float = 0.8,
    top_k: int = 50,
    top_p: float = 1.0,
    max_tokens: int = 4096,
    repetition_penalty: float = 1.05,
    seed: Optional[int] = None,
) -> RequestProfile:
    """Profile a single non-streaming request with stage-level timing."""
    from vllm import SamplingParams
    from veena3modal.core.constants import TRAINING_STOP_TOKEN_IDS, AUDIO_SAMPLE_RATE
    from veena3modal.audio.utils import add_wav_header

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

    t_total_start = time.perf_counter()

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

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

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

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

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


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

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

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

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

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

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

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

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

from __future__ import annotations

import argparse
import logging
import os
import sys
import time

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

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

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

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

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

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

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

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

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

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

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

    os.makedirs(local_dir, exist_ok=True)

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

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

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

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

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

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

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

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

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

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

    return llm_path, bicodec_path

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

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

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

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

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

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

    return runtime

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

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

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


sed -n '1,360p' veena3modal/services/tts_runtime.py


Output:
Chunk ID: bffe8f
Wall time: 0.0549 seconds
Process exited with code 0
Original token count: 3401
Output:
"""
TTS runtime wrapper (container-scoped singleton).

In Modal, each container should load the model once and then handle many requests.
This wrapper owns:
- SparkTTSModel (vLLM engine)
- IndicPromptBuilder
- BiCodecDecoder
- SparkTTSPipeline + Veena3SlidingWindowPipeline
- optional SuperResolutionService

Imports framework-agnostic inference + processing code from `veena3modal/core`,
`veena3modal/processing`, and `veena3modal/audio`.
"""

from __future__ import annotations

import os
import sys
import time
import logging
from threading import Lock
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)

_PROFILE_REQUEST_TIMING = "VEENA3_TIMING_LOGS"
_inflight_counter = 0
_inflight_lock = Lock()

def _timing_enabled() -> bool:
    """Return True when detailed request timing should be emitted."""
    return os.getenv(_PROFILE_REQUEST_TIMING, "").lower() in {"1", "true", "yes", "on"}

def _with_inflight_delta(delta: int) -> int:
    """Update in-flight request counter and return the new value."""
    global _inflight_counter
    with _inflight_lock:
        _inflight_counter = max(0, _inflight_counter + delta)
        return _inflight_counter

@dataclass
class TTSRuntime:
    """
    Holds long-lived, per-container inference objects.

    Thread Safety:
        This class is designed for async concurrency within a single container.
        The vLLM engine handles internal batching and scheduling.
        Do NOT share across processes without proper synchronization.
    """
    model: Any = None
    pipeline: Any = None
    streaming_pipeline: Any = None
    prompt_builder: Any = None
    bicodec_decoder: Any = None
    sr_service: Optional[Any] = None

    model_version: str = "not_loaded"
    is_loaded: bool = False
    load_time_ms: float = 0.0

    # Configuration
    model_path: str = ""
    bicodec_path: str = ""
    sr_checkpoint_dir: Optional[str] = None
    device: str = "cuda"

    # OPTIMIZATION: Pre-computed global tokens per speaker (eliminates ~110ms pre-roll in streaming)
    # Maps speaker_name -> list of 32 global token IDs captured at startup
    speaker_global_cache: Dict[str, List[int]] = field(default_factory=dict)

# Module-level singleton (per-container)
_runtime: Optional[TTSRuntime] = None

def get_runtime() -> Optional[TTSRuntime]:
    """Get the current TTS runtime singleton."""
    return _runtime

def is_initialized() -> bool:
    """Check if the runtime is initialized and ready."""
    return _runtime is not None and _runtime.is_loaded

def initialize_runtime(
    model_path: Optional[str] = None,
    bicodec_path: Optional[str] = None,
    sr_checkpoint_dir: Optional[str] = None,
    device: str = "cuda",
    hf_token: Optional[str] = None,
    gpu_memory_utilization: float = 0.25,
    enable_sr: bool = False,
    num_engines: int = 1,
) -> TTSRuntime:
    """
    Initialize the TTS runtime with all components.

    This should be called once per container (e.g., in Modal's @modal.enter).

    Args:
        model_path: Path to Spark TTS model (env: SPARK_TTS_MODEL_PATH)
        bicodec_path: Path to BiCodec model (env: BICODEC_MODEL_PATH, defaults to model_path)
        sr_checkpoint_dir: Path to super-resolution checkpoints (env: AP_BWE_CHECKPOINT_DIR)
        device: Device for inference (cuda/cpu)
        hf_token: HuggingFace token for private models (env: HF_TOKEN)
        gpu_memory_utilization: vLLM GPU memory fraction (default: 0.25)
        enable_sr: Enable super-resolution service
        num_engines: Number of vLLM engine instances (default: 1, set to 2-3 for Tier 3 optimization)

    Returns:
        Initialized TTSRuntime instance

    Raises:
        RuntimeError: If initialization fails
    """
    global _runtime

    start_time = time.time()

    # Resolve paths from env vars if not provided
    model_path = model_path or os.environ.get(
        'SPARK_TTS_MODEL_PATH',
        os.environ.get('MODEL_PATH', '/models/spark_tts_4speaker')
    )
    bicodec_path = bicodec_path or os.environ.get(
        'BICODEC_MODEL_PATH',
        model_path  # BiCodec is usually in the same directory
    )
    sr_checkpoint_dir = sr_checkpoint_dir or os.environ.get('AP_BWE_CHECKPOINT_DIR')
    hf_token = hf_token or os.environ.get('HF_TOKEN') or os.environ.get('HUGGING_FACE_HUB_TOKEN')

    logger.info(f"Initializing TTS runtime...")
    logger.info(f"  Model path: {model_path}")
    logger.info(f"  BiCodec path: {bicodec_path}")
    logger.info(f"  SR enabled: {enable_sr}, path: {sr_checkpoint_dir}")

    try:
        # Import inference components (vendored into veena3modal for Modal deployment)
        from veena3modal.core.model_loader import SparkTTSModel
        from veena3modal.core.pipeline import SparkTTSPipeline
        from veena3modal.core.bicodec_decoder import BiCodecDecoder
        from veena3modal.processing.prompt_builder import IndicPromptBuilder

        # Load SparkTTS model with vLLM engine
        if num_engines > 1:
            # TIER 3 OPTIMIZATION: Multiple vLLM engines on same GPU
            # Each engine gets gpu_memory_utilization / num_engines to share GPU fairly.
            # E.g., 3 engines at 0.08 each = 0.24 total, ~8GB per engine, ~24GB total.
            from veena3modal.core.multi_engine import create_multi_engine_model
            per_engine_mem = gpu_memory_utilization / num_engines
            logger.info(f"Loading {num_engines} vLLM engines ({per_engine_mem:.2f} GPU mem each)...")
            model = create_multi_engine_model(
                model_path=model_path,
                num_engines=num_engines,
                hf_token=hf_token,
                gpu_memory_per_engine=per_engine_mem,
            )
        else:
            logger.info("Loading SparkTTS model with vLLM...")
            model = SparkTTSModel(
                model_path=model_path,
                hf_token=hf_token,
                gpu_memory_utilization=gpu_memory_utilization,
            )

        # Initialize prompt builder
        logger.info("Initializing prompt builder...")
        prompt_builder = IndicPromptBuilder(
            tokenizer=model.tokenizer,
            model=model,
        )

        # Initialize BiCodec decoder
        logger.info("Initializing BiCodec decoder...")
        bicodec_decoder = BiCodecDecoder(
            device=device,
            model_path=bicodec_path,
        )

        # Initialize pipeline
        logger.info("Initializing TTS pipeline...")
        pipeline = SparkTTSPipeline(
            model=model,
            prompt_builder=prompt_builder,
            bicodec_decoder=bicodec_decoder,
        )

        # Initialize streaming pipeline (for M4)
        streaming_pipeline = None
        try:
            from veena3modal.core.streaming_pipeline import Veena3SlidingWindowPipeline
            logger.info("Initializing streaming pipeline...")
            # NOTE: Parameter is named 'snac_decoder' for legacy reasons, but works with BiCodecDecoder
            streaming_pipeline = Veena3SlidingWindowPipeline(
                model=model,
                prompt_builder=prompt_builder,
                snac_decoder=bicodec_decoder,  # BiCodecDecoder is interface-compatible
            )
        except ImportError as e:
            logger.warning(f"Streaming pipeline not available: {e}")

        # Initialize super-resolution (optional)
        sr_service = None
        if enable_sr and sr_checkpoint_dir:
            try:
                from veena3modal.core.super_resolution import SuperResolutionService
                logger.info(f"Initializing super-resolution from {sr_checkpoint_dir}...")
                sr_service = SuperResolutionService(checkpoint_dir=sr_checkpoint_dir)
                # Load the model explicitly
                if sr_service.load_model(device=device):
                    logger.info("✅ Super-resolution model loaded successfully")
                else:
                    logger.warning("Super-resolution model failed to load")
                    sr_service = None
            except Exception as e:
                logger.warning(f"Super-resolution not available: {e}")

        # Determine model version
        model_version = os.path.basename(model_path.rstrip('/'))
        if not model_version:
            model_version = "spark-tts"

        load_time = (time.time() - start_time) * 1000

        # Create runtime
        _runtime = TTSRuntime(
            model=model,
            pipeline=pipeline,
            streaming_pipeline=streaming_pipeline,
            prompt_builder=prompt_builder,
            bicodec_decoder=bicodec_decoder,
            sr_service=sr_service,
            model_version=model_version,
            is_loaded=True,
            load_time_ms=load_time,
            model_path=model_path,
            bicodec_path=bicodec_path,
            sr_checkpoint_dir=sr_checkpoint_dir,
            device=device,
        )

        logger.info(f"✅ TTS runtime initialized in {load_time:.0f}ms")
        logger.info(f"   Model version: {model_version}")

        # OPTIMIZATION: Pre-compute global tokens for all speakers
        # Disabled with multi-engine to avoid async event loop conflicts
        if num_engines <= 1:
            _precompute_speaker_globals(_runtime)
        else:
            logger.info("Speaker globals cache skipped (multi-engine mode)")

        # Update FastAPI app with model version
        try:
            from veena3modal.api.fastapi_app import set_model_version
            set_model_version(model_version)
        except ImportError:
            pass

        return _runtime

    except Exception as e:
        logger.error(f"❌ Failed to initialize TTS runtime: {e}")
        import traceback
        traceback.print_exc()
        raise RuntimeError(f"TTS runtime initialization failed: {e}") from e

def _precompute_speaker_globals(runtime: TTSRuntime) -> None:
    """
    Pre-compute and cache global tokens for all 12 speakers at startup.

    OPTIMIZATION: BiCodec streaming requires 32 "global tokens" before any audio can be emitted.
    These encode speaker identity via FSQ quantization. Since we have only 12 fixed speakers,
    we generate one short utterance per speaker at startup, capture the 32 global tokens,
    and cache them. During streaming, we inject cached globals via build_prefix_with_globals(),
    skipping the ~110ms global token pre-roll phase entirely.

    This runs synchronously at startup (adds ~5-10s to cold start, saves ~110ms per streaming TTFB).
    """
    import asyncio
    import re
    from vllm import SamplingParams
    from veena3modal.core.constants import (
        SPEAKER_MAP, TRAINING_STOP_TOKEN_IDS,
        DEFAULT_TEMPERATURE, DEFAULT_TOP_K, DEFAULT_TOP_P,
    )

    logger.info("Pre-computing global tokens for all speakers...")
    t_start = time.time()

    test_text = "Hello, this is a voice test."

    sampling_params = SamplingParams(
        temperature=DEFAULT_TEMPERATURE,
        top_k=DEFAULT_TOP_K,
        top_p=DEFAULT_TOP_P,
        max_tokens=128,  # Only need ~60 tokens (32 global + some semantic)
        stop=TRAINING_STOP_TOKEN_IDS,
        skip_special_tokens=False,
    )

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

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

        if final_output is None:
            return None

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

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

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

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

    elapsed = (time.time() - t_start) * 1000
    cached = len(runtime.speaker_global_cache)


Summary: auto


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