From 1e7792bd2e5c9cbae84cd2e2fefa2ffb79e9e590 Mon Sep 17 00:00:00 2001
From: sharath <sharath@hume.ai>
Date: Thu, 19 Mar 2026 21:26:57 +0000
Subject: [PATCH 1/4] streaming support

---
 README.md                | 136 ++++++++++++---
 tada/modules/__init__.py |   4 +-
 tada/modules/decoder.py  | 365 +++++++++++++++++++++++++++++++++++++++
 tada/modules/encoder.py  | 127 ++++++++++++++
 tada/modules/tada.py     | 109 ++++++++++--
 5 files changed, 705 insertions(+), 36 deletions(-)

diff --git a/README.md b/README.md
index 6125a7a..a4e8463 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,5 @@
 <!-- ---
-license: llama3.2
+license: mit
 language:
   - en
 tags:
@@ -30,6 +30,7 @@ TADA achieves high-fidelity synthesis and generation with a fraction of the comp
 ## Updates
 
 **March 2026**
+- **Streaming audio generation** — `generate(on_audio_chunk=callback)` for end-to-end streaming. Audio chunks are emitted during LLM generation. True TTFA ~350ms on GPU. No retraining needed.
 - Encoder no longer loaded inside `TadaForCausalLM` — saves ~2.5 GB VRAM. Load it separately only when encoding new prompts.
 - Added `EncoderOutput.save()` / `EncoderOutput.load()` for prompt caching — encode once, reuse without the encoder.
 - Default flow matching steps reduced from 20 to 10 (no perceptible quality loss, ~1.3x faster).
@@ -194,6 +195,120 @@ output = model.generate(
 )
 ```
 
+
+### Streaming Audio Generation
+
+By default, `generate()` produces the complete audio after all tokens are generated (non-streaming). To stream audio chunks as they are generated, pass an `on_audio_chunk` callback:
+
+```python
+import torch
+import torchaudio
+
+from tada.modules.encoder import Encoder, EncoderOutput
+from tada.modules.tada import TadaForCausalLM
+
+device = "cuda"
+encoder = Encoder.from_pretrained("HumeAI/tada-codec", subfolder="encoder").to(device)
+model = TadaForCausalLM.from_pretrained("HumeAI/tada-1b", torch_dtype=torch.bfloat16).to(device)
+
+audio, sample_rate = torchaudio.load("samples/ljspeech.wav")
+prompt = encoder(audio.to(device), sample_rate=sample_rate)
+
+# Stream audio chunks as they are generated
+def stream_audio(chunk: torch.Tensor, sample_rate: int):
+    # chunk shape: (num_samples,) at 24kHz
+    # Write to audio buffer, stream to speaker, send over network, etc.
+    print(f"Received {chunk.shape[-1] / sample_rate:.2f}s of audio")
+
+output = model.generate(
+    prompt=prompt,
+    text="Hello, this is a demonstration of streaming text to speech.",
+    on_audio_chunk=stream_audio,
+)
+# output.audio still contains the full concatenated audio
+```
+
+When `on_audio_chunk` is not provided (default), the existing non-streaming path runs unchanged.
+
+#### Application integration
+
+For production applications, connect the callback to your audio pipeline:
+
+```python
+import queue
+
+# Thread-safe audio queue for a web server
+audio_queue = queue.Queue()
+
+def stream_to_queue(chunk: torch.Tensor, sample_rate: int):
+    audio_bytes = (chunk.cpu().numpy() * 32767).astype("int16").tobytes()
+    audio_queue.put(audio_bytes)
+
+output = model.generate(
+    prompt=prompt,
+    text="Hello world",
+    on_audio_chunk=stream_to_queue,
+)
+audio_queue.put(None)  # Signal end of stream
+```
+
+```python
+# Write chunks directly to a WAV file as they arrive
+import soundfile as sf
+
+with sf.SoundFile("output.wav", mode="w", samplerate=24000, channels=1) as f:
+    def stream_to_file(chunk: torch.Tensor, sample_rate: int):
+        f.write(chunk.cpu().numpy())
+
+    output = model.generate(
+        prompt=prompt,
+        text="Hello world",
+        on_audio_chunk=stream_to_file,
+    )
+```
+
+For CPU or memory-constrained devices, reduce the CNN window size:
+
+```python
+output = model.generate(
+    prompt=prompt,
+    text="Hello world",
+    on_audio_chunk=stream_audio,
+    streaming_cnn_window_size=50,  # Default 100. Use 50 for <32GB RAM.
+)
+```
+
+#### How it works
+
+Streaming is end-to-end: audio chunks are emitted **during** LLM generation, not after. As the LLM generates each token, the acoustic features are immediately decoded into audio and delivered via the callback.
+
+The decoder has two stages:
+
+1. **Transformer**: Per-layer KV cache ensures each new token only computes Q and attends to cached K,V. This is bit-exact with the non-streaming path.
+
+2. **CNN (DACDecoder)**: A sliding window (default 100 frames) over accumulated transformer hidden states. The window provides 20 frames of left context and 15 frames of right lookahead for stable output. No model retraining required.
+
+#### Performance
+
+True time-to-first-audio (TTFA) measured from `generate()` call to first audio chunk:
+
+| Model | Device | Text | True TTFA | Total time | Non-streaming |
+|-------|--------|------|-----------|------------|---------------|
+| TADA-1B | GPU (H100) | Short | ~440ms | 0.47s | 0.77s |
+| TADA-1B | GPU (H100) | Medium | ~350ms | 1.06s | 0.86s |
+| TADA-1B | GPU (H100) | Long | ~330ms | 3.02s | 2.47s |
+| TADA-1B | CPU | Short | ~2.3s | 3.2s | 2.9s |
+| TADA-1B | CPU | Long | ~2.1s | 25s | 18s |
+
+Streaming uses less peak memory than non-streaming for medium and long text (up to 3.6 GB less VRAM on GPU).
+
+#### Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `on_audio_chunk` | `None` | Callback `(chunk, sample_rate) -> None`. Enables streaming when set. |
+| `streaming_cnn_window_size` | 100 | CNN sliding window size in frames. Use 50 for <32GB RAM. |
+
 ## 📚 Citation
 
 If you use this project in your research, please cite our paper:
@@ -207,25 +322,6 @@ If you use this project in your research, please cite our paper:
 }
 ```
 
-## License
-
-This repository contains both model weights and code, which are licensed separately:
-
-- **Model weights** are licensed under the Llama 3.2 Community License Agreement  
-- **Code** in this repository is licensed under the MIT License  
-
-You must comply with the terms of the Llama 3.2 license when using the models.
-
-See:
-- `LICENSE` for the Llama 3.2 license
-- `LICENSE_CODE` for the MIT license
-
 ## Contact
 
 [Hume AI](https://hume.ai) is an empathic AI research company. We research the datasets, tools, and models needed to give empathy to AI models to serve human wellbeing. If you're interested in any of our product or research collaborations, please reach out to us at hello@hume.ai.
-
-## Acknowledgements
-
-This project is built using Llama 3.2.
-
-Llama 3.2 is licensed under the Llama 3.2 Community License.
diff --git a/tada/modules/__init__.py b/tada/modules/__init__.py
index c7921fb..28ff5ed 100644
--- a/tada/modules/__init__.py
+++ b/tada/modules/__init__.py
@@ -1,6 +1,6 @@
 from .acoustic_spkr_verf import AcousticSpkrVerf
-from .decoder import Decoder
+from .decoder import Decoder, StreamingDecoder
 from .encoder import Encoder
 from .tada import TadaForCausalLM
 
-__all__ = ["TadaForCausalLM", "Encoder", "Decoder", "AcousticSpkrVerf"]
+__all__ = ["TadaForCausalLM", "Encoder", "Decoder", "StreamingDecoder", "AcousticSpkrVerf"]
diff --git a/tada/modules/decoder.py b/tada/modules/decoder.py
index 912195f..3826e4d 100644
--- a/tada/modules/decoder.py
+++ b/tada/modules/decoder.py
@@ -209,3 +209,368 @@ def forward(self, encoded_expanded: torch.Tensor, token_masks: torch.Tensor):
 
     def generate(self, encoded_expanded: torch.Tensor, **kwargs):
         return self.forward(encoded_expanded, **kwargs)
+
+class StreamingDecoder:
+    """Incremental block-by-block decoder with KV-cache for streaming audio generation.
+
+    Uses per-layer KV caching for the transformer (bit-exact with non-streaming)
+    and a sliding-window CNN for audio synthesis (near-exact, bounded memory).
+
+    Usage:
+        streaming_decoder = StreamingDecoder(decoder)
+        streaming_decoder.skip_leading_frames(leading_silence)  # optional
+        for token_embedding, time_before in token_stream:
+            audio_chunk = streaming_decoder.decode_block(token_embedding, time_before)
+            if audio_chunk is not None and audio_chunk.numel() > 0:
+                play(audio_chunk)
+        final_chunk = streaming_decoder.flush(trailing_frames=last_time_before)
+    """
+
+    def __init__(
+        self,
+        decoder: Decoder,
+        min_block_frames: int = 3,
+        max_cached_frames: int | None = 500,
+        cnn_window_size: int = 100,
+        cnn_left_context: int = 20,
+        cnn_lookahead: int = 15,
+        min_first_emission: int = 50,
+    ):
+        """
+        Args:
+            decoder: The pretrained Decoder model to wrap.
+            min_block_frames: Minimum frames before decoding a block. Very short
+                blocks (1-2 frames) can cause degenerate CNN behavior, so we
+                buffer them until we accumulate enough frames.
+            max_cached_frames: Maximum frames in the KV-cache. Oldest blocks are
+                evicted when exceeded. None = unlimited (bit-exact). Default 500.
+            cnn_window_size: Fixed CNN input size in frames. Layout:
+                [left_context | emittable | lookahead]. Default 100.
+                Use 50 for CPU/edge devices with <32GB RAM.
+            cnn_left_context: Left context frames for CNN stability.
+                Empirically: 20 → 0.000000 diff. 10 → 0.0003 diff.
+            cnn_lookahead: Right context frames held back from CNN output.
+                Empirically: 15 → 0.000001 diff (pretrained weights need 15).
+            min_first_emission: Minimum accumulated frames before first CNN call.
+                Avoids edge artifacts from processing very short sequences.
+                Default 50 adds ~1s to TTFA.
+        """
+        self.decoder = decoder
+        self.min_block_frames = min_block_frames
+        self.max_cached_frames = max_cached_frames
+        self._cnn_window_size = cnn_window_size
+        self._cnn_left_context = cnn_left_context
+        self._lookahead_frames = cnn_lookahead
+        self._min_first_emission = min_first_emission
+
+        # Precompute samples per frame (480 for strides [4,4,5,6])
+        self._samples_per_frame = 1
+        for s in self.decoder.config.strides:
+            self._samples_per_frame *= s
+
+        self._init_state()
+
+    def _init_state(self):
+        """Initialize/reset all streaming state."""
+        # Per-layer KV cache: list of (cached_k, cached_v) per layer
+        self._kv_cache: list[tuple[torch.Tensor, torch.Tensor]] = []
+        self._cached_frames: int = 0  # absolute frame count (for RoPE position)
+        self._all_token_masks: torch.Tensor | None = None  # accumulated masks for attention
+        self._block_frame_counts: list[int] = []  # frames per block (for eviction)
+
+        # Token buffer (accumulate small blocks before decoding)
+        self._buffered_tokens: list[torch.Tensor] = []
+        self._buffered_times: list[int] = []
+
+        # Sliding window CNN state
+        self._all_hidden: torch.Tensor | None = None  # recent transformer outputs (capped)
+        self._hidden_offset: int = 0  # absolute frame index of _all_hidden[0]
+        self._emitted_frames: int = 0  # absolute frame index up to which audio has been emitted
+        self._first_emitted: bool = False  # gates min_first_emission check
+
+    def skip_leading_frames(self, n_frames: int):
+        """Skip the first n_frames of audio output (e.g., leading silence).
+
+        Must be called before any decode_block() calls.
+        """
+        self._emitted_frames = n_frames
+
+    def reset(self):
+        """Clear all state for a new utterance."""
+        self._init_state()
+
+    def _expand_block(
+        self,
+        token_embeddings: list[torch.Tensor],
+        time_before_values: list[int],
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Expand token embeddings into dense frames.
+
+        For each token, inserts (time_before - 1) zero frames before the embedding.
+
+        Returns:
+            frames: (1, num_frames, embed_dim)
+            masks: (1, num_frames) — 1 at token positions, 0 at zero-fill
+        """
+        embed_dim = token_embeddings[0].shape[-1]
+        parts = []
+        mask_parts = []
+
+        for token_emb, t_before in zip(token_embeddings, time_before_values):
+            n_zeros = max(0, t_before - 1)
+            if n_zeros > 0:
+                parts.append(torch.zeros(n_zeros, embed_dim, device=device, dtype=dtype))
+                mask_parts.append(torch.zeros(n_zeros, device=device, dtype=torch.long))
+            parts.append(token_emb.unsqueeze(0).to(device=device, dtype=dtype))
+            mask_parts.append(torch.ones(1, device=device, dtype=torch.long))
+
+        frames = torch.cat(parts, dim=0).unsqueeze(0)
+        masks = torch.cat(mask_parts, dim=0).unsqueeze(0)
+        return frames, masks
+
+    def _build_kv_cache_mask(
+        self, all_token_masks: torch.Tensor, num_new_frames: int
+    ) -> torch.Tensor:
+        """Build v2 attention mask for KV-cache mode.
+
+        Returns only the rows for new frames: (batch, num_new_frames, total_frames).
+        Block ID offsets from eviction cancel out in the v2 equality checks.
+        """
+        full_mask = _create_segment_attention_mask(all_token_masks, version="v2")
+        return full_mask[:, -num_new_frames:, :]
+
+    def _evict_cache(self):
+        """Evict oldest blocks from KV-cache if over max_cached_frames.
+
+        Eviction happens at block boundaries to preserve correct block IDs.
+        Keeps at least 2 blocks (for v2 attention: current + prev).
+        """
+        if self.max_cached_frames is None or not self._kv_cache:
+            return
+
+        cache_len = self._kv_cache[0][0].shape[2]
+        if cache_len <= self.max_cached_frames:
+            return
+
+        frames_to_evict = 0
+        blocks_to_evict = 0
+        remaining = cache_len
+        while (blocks_to_evict < len(self._block_frame_counts) - 2
+               and remaining - self._block_frame_counts[blocks_to_evict] > self.max_cached_frames):
+            frames_to_evict += self._block_frame_counts[blocks_to_evict]
+            remaining -= self._block_frame_counts[blocks_to_evict]
+            blocks_to_evict += 1
+
+        if frames_to_evict == 0:
+            return
+
+        for i, (cached_k, cached_v) in enumerate(self._kv_cache):
+            self._kv_cache[i] = (
+                cached_k[:, :, frames_to_evict:, :],
+                cached_v[:, :, frames_to_evict:, :],
+            )
+
+        if self._all_token_masks is not None:
+            self._all_token_masks = self._all_token_masks[:, frames_to_evict:]
+
+        self._block_frame_counts = self._block_frame_counts[blocks_to_evict:]
+        # Note: _cached_frames stays as absolute position (for RoPE)
+
+    @torch.no_grad()
+    def decode_block(
+        self,
+        token_embedding: torch.Tensor,
+        time_before: int,
+    ) -> torch.Tensor | None:
+        """Decode a single block incrementally. Returns audio chunk or None if buffering."""
+        self._buffered_tokens.append(token_embedding)
+        self._buffered_times.append(time_before)
+
+        total_frames = sum(max(0, t - 1) for t in self._buffered_times) + len(self._buffered_tokens)
+        if total_frames < self.min_block_frames:
+            return None
+
+        return self._decode_buffered()
+
+    def _decode_buffered(self) -> torch.Tensor:
+        """Decode buffered tokens: transformer (KV-cache) → CNN (sliding window)."""
+        device = next(self.decoder.parameters()).device
+        decoder_dtype = next(self.decoder.parameters()).dtype
+
+        current_input, current_masks = self._expand_block(
+            self._buffered_tokens, self._buffered_times,
+            device=device, dtype=decoder_dtype,
+        )
+        current_frames = current_input.shape[1]
+        self._buffered_tokens = []
+        self._buffered_times = []
+
+        # Transformer with KV-cache (bit-exact)
+        current_hidden = self._transformer_forward_cached(current_input, current_masks)
+
+        # Accumulate hidden states
+        if self._all_hidden is not None:
+            self._all_hidden = torch.cat([self._all_hidden, current_hidden], dim=1)
+        else:
+            self._all_hidden = current_hidden
+
+        # Buffer until min_first_emission
+        if not self._first_emitted and self._all_hidden.shape[1] < self._min_first_emission + self._lookahead_frames:
+            self._block_frame_counts.append(current_frames)
+            return torch.zeros(1, 0, device=device, dtype=decoder_dtype).squeeze(0)
+
+        # Run CNN and emit safe audio
+        audio = self._cnn_emit(is_flush=False)
+
+        self._block_frame_counts.append(current_frames)
+        return audio
+
+    def _cnn_emit(self, is_flush: bool = False) -> torch.Tensor:
+        """Run CNN on sliding window and emit the safe audio region.
+
+        Shared by _decode_buffered() and flush(). The only difference is
+        flush() disables the lookahead holdback (emits everything).
+
+        Coordinate systems:
+        - "absolute": frame index since start of utterance (_emitted_frames)
+        - "buffer": index into _all_hidden (0-based, shifts on trim)
+        - _hidden_offset converts: absolute = buffer + _hidden_offset
+        """
+        buf_frames = self._all_hidden.shape[1]
+        device = self._all_hidden.device
+        dtype = self._all_hidden.dtype
+        spf = self._samples_per_frame
+
+        # Determine sliding window position (buffer-relative).
+        # Constraint: window must not slide so far that _emitted_frames
+        # falls inside the left context zone (would skip unemitted frames).
+        max_start_for_emit = max(0, self._emitted_frames - self._hidden_offset - self._cnn_left_context)
+        win_buf_start = min(max(0, buf_frames - self._cnn_window_size), max_start_for_emit)
+
+        window = self._all_hidden[:, win_buf_start:, :]
+        window_frames = window.shape[1]
+
+        # Run CNN
+        audio = self.decoder.wav_decoder(window.transpose(1, 2))
+
+        # Convert to absolute coordinates
+        win_abs_start = self._hidden_offset + win_buf_start
+
+        # Safe emit region
+        left_ctx = self._cnn_left_context if win_abs_start > 0 else 0
+        abs_safe_start = win_abs_start + left_ctx
+        lookahead = 0 if is_flush else self._lookahead_frames
+        abs_safe_end = win_abs_start + window_frames - lookahead
+
+        # Only emit frames we haven't emitted yet
+        emit_start = max(abs_safe_start, self._emitted_frames)
+        emit_end = abs_safe_end
+
+        if emit_end <= emit_start:
+            return torch.zeros(1, 0, device=device, dtype=dtype).squeeze(0)
+
+        # Map to window-local sample positions
+        win_start_sample = (emit_start - win_abs_start) * spf
+        win_end_sample = min((emit_end - win_abs_start) * spf, audio.shape[2])
+
+        emit_audio = audio[:, :, win_start_sample:win_end_sample]
+        self._emitted_frames = emit_end
+        self._first_emitted = True
+
+        # Trim _all_hidden to cap memory
+        if buf_frames > self._cnn_window_size:
+            trim = buf_frames - self._cnn_window_size
+            self._all_hidden = self._all_hidden[:, trim:, :]
+            self._hidden_offset += trim
+
+        return emit_audio.squeeze(0)
+
+    def _transformer_forward_cached(
+        self, current_input: torch.Tensor, current_masks: torch.Tensor
+    ) -> torch.Tensor:
+        """Run transformer layers with KV-cache, returning hidden states for new frames."""
+        new_frames = current_input.shape[1]
+
+        # Accumulate token masks
+        if self._all_token_masks is not None:
+            all_masks = torch.cat([self._all_token_masks, current_masks], dim=1)
+        else:
+            all_masks = current_masks
+
+        # Build attention mask (new frames attend to all cached + new frames)
+        attn_mask = self._build_kv_cache_mask(all_masks, new_frames)
+
+        # Project to hidden dim
+        decoder_input = self.decoder.decoder_proj(current_input)
+        x = self.decoder.local_attention_decoder.input_proj(decoder_input)
+
+        # Iterate through layers with per-layer KV-cache
+        new_kv_cache: list[tuple[torch.Tensor, torch.Tensor]] = []
+        for i, layer in enumerate(self.decoder.local_attention_decoder.layers):
+            cached_k, cached_v = self._kv_cache[i] if self._kv_cache else (None, None)
+
+            x, new_k, new_v = layer.forward_with_cache(
+                x, cached_k, cached_v,
+                position_offset=self._cached_frames,
+                mask=attn_mask,
+            )
+
+            if cached_k is not None:
+                updated_k = torch.cat([cached_k, new_k], dim=2)
+                updated_v = torch.cat([cached_v, new_v], dim=2)
+            else:
+                updated_k, updated_v = new_k, new_v
+
+            new_kv_cache.append((updated_k, updated_v))
+
+        x = self.decoder.local_attention_decoder.final_norm(x)
+
+        # Update state
+        self._kv_cache = new_kv_cache
+        self._cached_frames += new_frames
+        self._all_token_masks = all_masks
+        self._evict_cache()
+
+        return x
+
+    @torch.no_grad()
+    def flush(self, trailing_frames: int = 0) -> torch.Tensor | None:
+        """Flush remaining buffered tokens and trailing silence.
+
+        Emits all remaining audio with no lookahead holdback.
+        """
+        device = next(self.decoder.parameters()).device
+        decoder_dtype = next(self.decoder.parameters()).dtype
+
+        has_buffered = len(self._buffered_tokens) > 0
+        has_trailing = trailing_frames > 0
+
+        if not has_buffered and not has_trailing and self._all_hidden is None:
+            return None
+
+        audio_parts = []
+
+        # Decode any buffered tokens first
+        if has_buffered:
+            audio_parts.append(self._decode_buffered())
+
+        # Process trailing silence through transformer
+        if has_trailing and self._all_hidden is not None:
+            embed_dim = self.decoder.config.embed_dim
+            trailing_input = torch.zeros(1, trailing_frames, embed_dim, device=device, dtype=decoder_dtype)
+            trailing_masks = torch.zeros(1, trailing_frames, device=device, dtype=torch.long)
+            trailing_hidden = self._transformer_forward_cached(trailing_input, trailing_masks)
+            self._all_hidden = torch.cat([self._all_hidden, trailing_hidden], dim=1)
+
+        # Emit everything remaining (no lookahead)
+        if self._all_hidden is not None:
+            audio_parts.append(self._cnn_emit(is_flush=True))
+
+        # Filter empty tensors
+        audio_parts = [a for a in audio_parts if a.numel() > 0]
+
+        if not audio_parts:
+            return None
+        return torch.cat(audio_parts, dim=-1)
diff --git a/tada/modules/encoder.py b/tada/modules/encoder.py
index 8489e10..c4cd4e8 100644
--- a/tada/modules/encoder.py
+++ b/tada/modules/encoder.py
@@ -334,6 +334,104 @@ def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Te
 
         return output
 
+    def _apply_rope_with_offset(self, x: torch.Tensor, seq_len: int, offset: int) -> torch.Tensor:
+        """Apply RoPE starting at a given position offset.
+
+        Args:
+            x: (batch, num_heads, seq_len, head_dim)
+            seq_len: number of new positions
+            offset: absolute position of the first frame
+        Returns:
+            rotated x with same shape
+        """
+        batch, num_heads, _, head_dim = x.shape
+
+        freqs = self.rope_freqs[offset : offset + seq_len]  # (seq_len, head_dim // 2, 2)
+        freqs_cos = freqs[..., 0]
+        freqs_sin = freqs[..., 1]
+
+        x_reshaped = x.reshape(batch, num_heads, seq_len, head_dim // 2, 2)
+        x0 = x_reshaped[..., 0]
+        x1 = x_reshaped[..., 1]
+
+        x_rotated_0 = x0 * freqs_cos.unsqueeze(0).unsqueeze(0) - x1 * freqs_sin.unsqueeze(0).unsqueeze(0)
+        x_rotated_1 = x0 * freqs_sin.unsqueeze(0).unsqueeze(0) + x1 * freqs_cos.unsqueeze(0).unsqueeze(0)
+
+        x_rotated = torch.stack([x_rotated_0, x_rotated_1], dim=-1)
+        return x_rotated.reshape(batch, num_heads, seq_len, head_dim)
+
+    def forward_with_cache(
+        self,
+        x: torch.Tensor,
+        cached_k: torch.Tensor | None,
+        cached_v: torch.Tensor | None,
+        position_offset: int = 0,
+        mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Forward with KV-cache. Returns (output, new_k, new_v).
+
+        Args:
+            x: (batch, new_len, d_model) — new frames only
+            cached_k: (batch, num_heads, cached_len, head_dim) or None
+            cached_v: (batch, num_heads, cached_len, head_dim) or None
+            position_offset: absolute position of first new frame (for RoPE)
+            mask: (batch, new_len, total_len) attention mask, True=masked
+
+        Returns:
+            output: (batch, new_len, d_model)
+            new_k: (batch, num_heads, new_len, head_dim) — post-RoPE, append to cache
+            new_v: (batch, num_heads, new_len, head_dim) — append to cache
+        """
+        batch_size, new_len, d_model = x.shape
+
+        # Compute Q, K, V from new frames only
+        qkv = self.qkv(x)
+        qkv = qkv.reshape(batch_size, new_len, 3, self.num_heads, self.head_dim)
+        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, batch, num_heads, new_len, head_dim)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        # Apply RoPE with correct absolute positions
+        q = self._apply_rope_with_offset(q, new_len, position_offset)
+        k = self._apply_rope_with_offset(k, new_len, position_offset)
+
+        # Save new K, V for caching (post-RoPE K, pre-attention V)
+        new_k = k
+        new_v = v
+
+        # Concatenate with cached K, V
+        if cached_k is not None:
+            full_k = torch.cat([cached_k, k], dim=2)  # (batch, heads, cached+new, head_dim)
+            full_v = torch.cat([cached_v, v], dim=2)
+        else:
+            full_k = k
+            full_v = v
+
+        # Compute attention: Q (new) × K (all)
+        attn_scores = torch.matmul(q, full_k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        # (batch, num_heads, new_len, total_len)
+
+        # Apply mask
+        if mask is not None:
+            if mask.dim() == 3:
+                # (batch, new_len, total_len) -> broadcast to heads
+                attn_scores = attn_scores.masked_fill(mask.unsqueeze(1), float("-inf"))
+            elif mask.dim() == 2:
+                attn_scores = attn_scores.masked_fill(mask.unsqueeze(0).unsqueeze(0), float("-inf"))
+
+        attn_weights = torch.softmax(attn_scores, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+
+        # Apply attention to values
+        attn_output = torch.matmul(attn_weights, full_v)  # (batch, heads, new_len, head_dim)
+        attn_output = attn_output.transpose(1, 2).reshape(batch_size, new_len, d_model)
+
+        # Output projection + residual + norm (same as forward)
+        output = self.out_proj(attn_output)
+        output = self.dropout(output)
+        output = self.layer_norm(x + output)
+
+        return output, new_k, new_v
+
 
 class LocalAttentionEncoderLayer(torch.nn.Module):
     """Transformer encoder layer with local self-attention and feed-forward network."""
@@ -393,6 +491,35 @@ def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Te
 
         return x
 
+    def forward_with_cache(
+        self,
+        x: torch.Tensor,
+        cached_k: torch.Tensor | None,
+        cached_v: torch.Tensor | None,
+        position_offset: int = 0,
+        mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Forward with KV-cache for streaming decoding.
+
+        Args:
+            x: (batch, new_len, d_model) — new frames only
+            cached_k, cached_v: per-layer KV cache tensors or None
+            position_offset: absolute position of first new frame
+            mask: (batch, new_len, total_len) attention mask
+
+        Returns:
+            output: (batch, new_len, d_model)
+            new_k: (batch, num_heads, new_len, head_dim) — to append to cache
+            new_v: (batch, num_heads, new_len, head_dim)
+        """
+        # Self-attention with cache (includes residual + norm)
+        x, new_k, new_v = self.self_attn.forward_with_cache(
+            x, cached_k, cached_v, position_offset, mask
+        )
+        # FFN block with residual + norm (same as forward)
+        x = self.norm(x + self.ffn(x))
+        return x, new_k, new_v
+
 
 class LocalAttentionEncoder(torch.nn.Module):
     """Stack of local attention encoder layers."""
diff --git a/tada/modules/tada.py b/tada/modules/tada.py
index 6367db1..4ef9033 100644
--- a/tada/modules/tada.py
+++ b/tada/modules/tada.py
@@ -1,7 +1,7 @@
 import math
 import time
 from dataclasses import dataclass, replace
-from typing import Literal, Optional
+from typing import Callable, Literal, Optional
 
 import torch
 from transformers import AutoTokenizer, LlamaForCausalLM
@@ -15,7 +15,7 @@
 from ..utils.gray_code import decode_gray_code_to_time
 from ..utils.text import normalize_text as normalize_text_fn
 from .acoustic_spkr_verf import AcousticSpkrVerf
-from .decoder import Decoder, DecoderConfig
+from .decoder import Decoder, DecoderConfig, StreamingDecoder
 from .encoder import Encoder, EncoderConfig, EncoderOutput
 
 
@@ -653,6 +653,8 @@ def _generate(
         use_text_in_prompt: bool = False,
         verbose: bool = False,
         return_logits: bool = False,
+        on_token_ready: Callable[[torch.Tensor, int, int], None] | None = None,
+        on_token_ready_offset: int = 0,
         **kwargs,
     ) -> SyncTokGenerationOutput:
         start_header_id = self.tokenizer.convert_tokens_to_ids("<|start_header_id|>")
@@ -1099,6 +1101,13 @@ def _generate(
                     time_len_after = predicted_time_len_after
                 all_time_before.append(time_len_before)
                 last_time_before = time_len_before
+
+                # Fire streaming callback for predicted (non-prompt) tokens
+                if on_token_ready is not None and acoustic_feat_type == "predicted":
+                    token_idx = len(all_acoustic_features) - 1  # 0-based index into all_acoustic_features
+                    if token_idx >= on_token_ready_offset:
+                        on_token_ready(acoustic_features, time_len_before, token_idx - on_token_ready_offset)
+
             model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs)
             if log_time:
                 if torch.cuda.is_available():
@@ -1201,6 +1210,8 @@ def generate(  # type: ignore[override]
         use_text_in_prompt: bool = False,
         normalize_text: bool = True,
         verbose: bool = False,
+        on_audio_chunk: Callable[[torch.Tensor, int], None] | None = None,
+        streaming_cnn_window_size: int = 100,
     ) -> GenerationOutput:
         # Auto-cast prompt tensors to model dtype (e.g. fp32 prompt + bf16 model)
         model_dtype = next(self.parameters()).dtype
@@ -1212,14 +1223,21 @@ def generate(  # type: ignore[override]
         if isinstance(text, str):
             text = [text]
         text = [normalize_text_fn(t) if normalize_text else t for t in text]
-        input_ids = [
-            self.tokenizer.encode(prompt.text[0] + text[i])[prompt.text_tokens_len[i] :] for i in range(len(text))
-        ]
         audio_feat_len = (prompt.audio_len / prompt.sample_rate * 50).ceil().long()
 
+        # Joint encoding with a space separator ensures correct BPE tokenization
+        # at the prompt/gen boundary. Without a space, BPE merges boundary chars
+        # (e.g., ".Hello" -> ".H"+"ello", losing the "H" from "Hello").
+        # The space ensures the first generation token is tokenized correctly.
+        prompt_text = prompt.text[0]
+        gen_text = text[0]
+        # Add space if prompt doesn't end with whitespace and gen doesn't start with it
+        if prompt_text and gen_text and not prompt_text[-1].isspace() and not gen_text[0].isspace():
+            joint_text = prompt_text + " " + gen_text
+        else:
+            joint_text = prompt_text + gen_text
         text_tokens = [
-            self.tokenizer.encode(prompt.text[0], add_special_tokens=False)
-            + self.tokenizer.encode(text[0], add_special_tokens=False)
+            self.tokenizer.encode(joint_text, add_special_tokens=False)
         ]
         input_ids, input_lengths = self._add_bos_eos(
             torch.tensor(text_tokens, device=self.device),
@@ -1272,6 +1290,45 @@ def generate(  # type: ignore[override]
             time_len_before = time_len_before[:, :-num_transition_steps]
             time_len_after = time_len_after[:, :-num_transition_steps]
 
+        num_prompt_tokens = prompt_acoustic_features.shape[1]
+        # Token offset: generated tokens start after prompt + transition tokens
+        token_decode_offset = num_prompt_tokens + num_transition_steps - 1
+
+        # Set up E2E streaming: decode audio DURING _generate(), not after
+        streaming_decoder = None
+        streaming_audio_chunks: list[torch.Tensor] = []
+        on_token_ready_cb = None
+        first_token_seen = False
+
+        if on_audio_chunk is not None:
+            streaming_decoder = StreamingDecoder(
+                self.decoder,
+                min_block_frames=3,
+                cnn_window_size=streaming_cnn_window_size,
+            )
+
+            acoustic_std = self.config.acoustic_std
+            acoustic_mean = self.config.acoustic_mean
+
+            def on_token_ready_cb(acoustic_feat: torch.Tensor, time_before_val: torch.Tensor, rel_idx: int):
+                nonlocal first_token_seen
+                # Denormalize acoustic features (normally done after _generate())
+                feat = acoustic_feat * acoustic_std + acoustic_mean
+                tb_val = time_before_val[0].item()  # (batch, 1) -> scalar
+
+                # Skip leading silence on first token.
+                # Clamp to at least 5 frames (100ms) to avoid transition artifacts
+                # when the LLM predicts very short leading silence (1-2 frames).
+                if not first_token_seen:
+                    first_token_seen = True
+                    leading_frames = max(int(tb_val), 5)
+                    streaming_decoder.skip_leading_frames(leading_frames)
+
+                chunk = streaming_decoder.decode_block(feat[0, 0], tb_val)
+                if chunk is not None and chunk.shape[-1] > 0:
+                    streaming_audio_chunks.append(chunk)
+                    on_audio_chunk(chunk, 24000)
+
         outputs: SyncTokGenerationOutput = self._generate(
             input_ids=input_ids,
             text=text,
@@ -1285,20 +1342,44 @@ def generate(  # type: ignore[override]
             num_steps=input_ids.shape[-1] + num_extra_steps,
             inference_options=inference_options,
             use_text_in_prompt=use_text_in_prompt,
+            on_token_ready=on_token_ready_cb,
+            on_token_ready_offset=token_decode_offset,
         )
 
-        num_prompt_tokens = prompt_acoustic_features.shape[1]
-        acoustic_features = outputs.acoustic_features * self.config.acoustic_std + self.config.acoustic_mean
+        # Flush streaming decoder after _generate() completes
+        if streaming_decoder is not None:
+            # Flush with trailing frames from the last time_before
+            all_tb = outputs.time_before
+            if all_tb.shape[1] > 0:
+                trailing = all_tb[0, -1].item()
+            else:
+                trailing = 0
+            final_chunk = streaming_decoder.flush(trailing_frames=trailing)
+            if final_chunk is not None and final_chunk.shape[-1] > 0:
+                streaming_audio_chunks.append(final_chunk)
+                on_audio_chunk(final_chunk, 24000)
 
-        encoded = acoustic_features[..., num_prompt_tokens + num_transition_steps - 1 :, :]
-        time_before = outputs.time_before[..., num_prompt_tokens + num_transition_steps - 1 :]
+        acoustic_features = outputs.acoustic_features * self.config.acoustic_std + self.config.acoustic_mean
+        encoded = acoustic_features[..., token_decode_offset:, :]
+        time_before = outputs.time_before[..., token_decode_offset:]
         wavs = []
 
         for i in range(encoded.shape[0]):
             try:
-                wav = self._decode_wav(encoded[i], time_before=time_before[i]).squeeze(0, 1)
-                wav = wav[..., int(24000 * time_before[i][0] / 50) :]  # remove leading silence
-                wavs.append(wav)
+                if on_audio_chunk is not None and streaming_audio_chunks:
+                    # E2E streaming: audio was already emitted during _generate()
+                    if streaming_audio_chunks:
+                        wavs.append(torch.cat(streaming_audio_chunks, dim=-1).squeeze(0))
+                    else:
+                        wavs.append(None)
+                else:
+                    # Non-streaming decode: existing single-pass path
+                    wav = self._decode_wav(encoded[i], time_before=time_before[i]).squeeze(0, 1)
+                    # Trim leading silence. Clamp to at least 5 frames (100ms)
+                    # to avoid transition artifacts when time_before[0] is very small.
+                    min_trim_frames = max(int(time_before[i][0].item()), 5)
+                    wav = wav[..., min_trim_frames * 480 :]
+                    wavs.append(wav)
             except Exception:
                 wavs.append(None)
 

From 0974cb70873c4f421e0106e116beba2a0646d1c7 Mon Sep 17 00:00:00 2001
From: sharath <sharath@hume.ai>
Date: Thu, 19 Mar 2026 21:43:42 +0000
Subject: [PATCH 2/4] streaming implementation

---
 README.md | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index a4e8463..e7c0aeb 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,5 @@
 <!-- ---
-license: mit
+license: llama3.2
 language:
   - en
 tags:
@@ -309,6 +309,25 @@ Streaming uses less peak memory than non-streaming for medium and long text (up
 | `on_audio_chunk` | `None` | Callback `(chunk, sample_rate) -> None`. Enables streaming when set. |
 | `streaming_cnn_window_size` | 100 | CNN sliding window size in frames. Use 50 for <32GB RAM. |
 
+## License
+
+This repository contains both model weights and code, which are licensed separately:
+
+- **Model weights** are licensed under the Llama 3.2 Community License Agreement  
+- **Code** in this repository is licensed under the MIT License  
+
+You must comply with the terms of the Llama 3.2 license when using the models.
+
+See:
+- `LICENSE` for the Llama 3.2 license
+- `LICENSE_CODE` for the MIT license
+
+## Acknowledgements
+
+This project is built using Llama 3.2.
+
+Llama 3.2 is licensed under the Llama 3.2 Community License.
+
 ## 📚 Citation
 
 If you use this project in your research, please cite our paper:

From 6aa9e073d4e9bd9b8daf08a1b391c26093028479 Mon Sep 17 00:00:00 2001
From: sharath <sharath@hume.ai>
Date: Mon, 11 May 2026 17:02:29 +0000
Subject: [PATCH 3/4] streaming update

---
 README.md                  |  49 ++---
 tada/modules/__init__.py   |   4 +-
 tada/modules/tada.py       | 218 +++++++++++++--------
 tests/integration_test.log | 168 ++++++++++++++++
 tests/run_integration.sh   |  16 ++
 tests/test_streaming.py    | 381 +++++++++++++++++++++++++++++++++++++
 6 files changed, 728 insertions(+), 108 deletions(-)
 create mode 100644 tests/integration_test.log
 create mode 100755 tests/run_integration.sh
 create mode 100644 tests/test_streaming.py

diff --git a/README.md b/README.md
index e7c0aeb..fe48256 100644
--- a/README.md
+++ b/README.md
@@ -30,7 +30,7 @@ TADA achieves high-fidelity synthesis and generation with a fraction of the comp
 ## Updates
 
 **March 2026**
-- **Streaming audio generation** — `generate(on_audio_chunk=callback)` for end-to-end streaming. Audio chunks are emitted during LLM generation. True TTFA ~350ms on GPU. No retraining needed.
+- **Streaming audio generation** — `generate(stream=True)` for end-to-end streaming. Audio chunks are yielded during LLM generation. True TTFA ~350ms on GPU. No retraining needed.
 - Encoder no longer loaded inside `TadaForCausalLM` — saves ~2.5 GB VRAM. Load it separately only when encoding new prompts.
 - Added `EncoderOutput.save()` / `EncoderOutput.load()` for prompt caching — encode once, reuse without the encoder.
 - Default flow matching steps reduced from 20 to 10 (no perceptible quality loss, ~1.3x faster).
@@ -198,7 +198,7 @@ output = model.generate(
 
 ### Streaming Audio Generation
 
-By default, `generate()` produces the complete audio after all tokens are generated (non-streaming). To stream audio chunks as they are generated, pass an `on_audio_chunk` callback:
+By default, `generate()` produces the complete audio after all tokens are generated (non-streaming). To stream audio chunks as they are generated, pass `stream=True`:
 
 ```python
 import torch
@@ -215,24 +215,26 @@ audio, sample_rate = torchaudio.load("samples/ljspeech.wav")
 prompt = encoder(audio.to(device), sample_rate=sample_rate)
 
 # Stream audio chunks as they are generated
-def stream_audio(chunk: torch.Tensor, sample_rate: int):
+stream = model.generate(
+    prompt=prompt,
+    text="Hello, this is a demonstration of streaming text to speech.",
+    stream=True,
+)
+for chunk, sample_rate in stream:
     # chunk shape: (num_samples,) at 24kHz
     # Write to audio buffer, stream to speaker, send over network, etc.
     print(f"Received {chunk.shape[-1] / sample_rate:.2f}s of audio")
 
-output = model.generate(
-    prompt=prompt,
-    text="Hello, this is a demonstration of streaming text to speech.",
-    on_audio_chunk=stream_audio,
-)
-# output.audio still contains the full concatenated audio
+# After iteration, stream.result contains the full GenerationOutput
+# (audio field is None since chunks were streamed individually)
+output = stream.result
 ```
 
-When `on_audio_chunk` is not provided (default), the existing non-streaming path runs unchanged.
+When `stream` is not set (default), the existing non-streaming path runs unchanged.
 
 #### Application integration
 
-For production applications, connect the callback to your audio pipeline:
+For production applications, connect the iterator to your audio pipeline:
 
 ```python
 import queue
@@ -240,15 +242,10 @@ import queue
 # Thread-safe audio queue for a web server
 audio_queue = queue.Queue()
 
-def stream_to_queue(chunk: torch.Tensor, sample_rate: int):
+stream = model.generate(prompt=prompt, text="Hello world", stream=True)
+for chunk, sr in stream:
     audio_bytes = (chunk.cpu().numpy() * 32767).astype("int16").tobytes()
     audio_queue.put(audio_bytes)
-
-output = model.generate(
-    prompt=prompt,
-    text="Hello world",
-    on_audio_chunk=stream_to_queue,
-)
 audio_queue.put(None)  # Signal end of stream
 ```
 
@@ -257,30 +254,24 @@ audio_queue.put(None)  # Signal end of stream
 import soundfile as sf
 
 with sf.SoundFile("output.wav", mode="w", samplerate=24000, channels=1) as f:
-    def stream_to_file(chunk: torch.Tensor, sample_rate: int):
+    for chunk, sr in model.generate(prompt=prompt, text="Hello world", stream=True):
         f.write(chunk.cpu().numpy())
-
-    output = model.generate(
-        prompt=prompt,
-        text="Hello world",
-        on_audio_chunk=stream_to_file,
-    )
 ```
 
 For CPU or memory-constrained devices, reduce the CNN window size:
 
 ```python
-output = model.generate(
+stream = model.generate(
     prompt=prompt,
     text="Hello world",
-    on_audio_chunk=stream_audio,
+    stream=True,
     streaming_cnn_window_size=50,  # Default 100. Use 50 for <32GB RAM.
 )
 ```
 
 #### How it works
 
-Streaming is end-to-end: audio chunks are emitted **during** LLM generation, not after. As the LLM generates each token, the acoustic features are immediately decoded into audio and delivered via the callback.
+Streaming is end-to-end: audio chunks are emitted **during** LLM generation, not after. As the LLM generates each token, the acoustic features are immediately decoded into audio and yielded to the caller.
 
 The decoder has two stages:
 
@@ -306,7 +297,7 @@ Streaming uses less peak memory than non-streaming for medium and long text (up
 
 | Parameter | Default | Description |
 |-----------|---------|-------------|
-| `on_audio_chunk` | `None` | Callback `(chunk, sample_rate) -> None`. Enables streaming when set. |
+| `stream` | `False` | When `True`, returns an `AudioStream` iterator yielding `(chunk, sample_rate)` tuples. |
 | `streaming_cnn_window_size` | 100 | CNN sliding window size in frames. Use 50 for <32GB RAM. |
 
 ## License
diff --git a/tada/modules/__init__.py b/tada/modules/__init__.py
index 28ff5ed..9c3bb34 100644
--- a/tada/modules/__init__.py
+++ b/tada/modules/__init__.py
@@ -1,6 +1,6 @@
 from .acoustic_spkr_verf import AcousticSpkrVerf
 from .decoder import Decoder, StreamingDecoder
 from .encoder import Encoder
-from .tada import TadaForCausalLM
+from .tada import AudioStream, TadaForCausalLM
 
-__all__ = ["TadaForCausalLM", "Encoder", "Decoder", "StreamingDecoder", "AcousticSpkrVerf"]
+__all__ = ["TadaForCausalLM", "AudioStream", "Encoder", "Decoder", "StreamingDecoder", "AcousticSpkrVerf"]
diff --git a/tada/modules/tada.py b/tada/modules/tada.py
index 4ef9033..20deebf 100644
--- a/tada/modules/tada.py
+++ b/tada/modules/tada.py
@@ -1,7 +1,7 @@
 import math
 import time
 from dataclasses import dataclass, replace
-from typing import Callable, Literal, Optional
+from typing import Generator, Literal, Optional
 
 import torch
 from transformers import AutoTokenizer, LlamaForCausalLM
@@ -121,6 +121,104 @@ class GenerationOutput(ModelOutput):
     step_logs: list[dict] | None = None
 
 
+
+@dataclass
+class StepOutput:
+    """Single predicted token output, yielded by _generate() for streaming."""
+    acoustic_features: "torch.Tensor"
+    time_before: "torch.Tensor"
+
+
+class AudioStream:
+    """Iterator that yields audio chunks during streaming generation.
+
+    Usage:
+        stream = model.generate(prompt=prompt, text="Hello", stream=True)
+        for chunk, sample_rate in stream:
+            play(chunk)
+        output = stream.result  # GenerationOutput, available after iteration
+    """
+
+    def __init__(
+        self,
+        gen: "Generator",
+        decoder: "Decoder",
+        acoustic_mean: float,
+        acoustic_std: float,
+        cnn_window_size: int = 100,
+    ):
+        self._gen = gen
+        self._decoder = decoder
+        self._acoustic_mean = acoustic_mean
+        self._acoustic_std = acoustic_std
+        self._cnn_window_size = cnn_window_size
+        self.result: "GenerationOutput | None" = None
+
+    def __iter__(self):
+        streaming_decoder = StreamingDecoder(
+            self._decoder,
+            min_block_frames=3,
+            cnn_window_size=self._cnn_window_size,
+        )
+        first_token_seen = False
+
+        for step_output in self._gen:
+            if isinstance(step_output, SyncTokGenerationOutput):
+                # Final yield — flush remaining audio
+                all_tb = step_output.time_before
+                trailing = int(all_tb[0, -1].item()) if all_tb.shape[1] > 0 else 0
+                final_chunk = streaming_decoder.flush(trailing_frames=trailing)
+                if final_chunk is not None and final_chunk.shape[-1] > 0:
+                    yield final_chunk, 24000
+                # Build .result so caller can access full GenerationOutput
+                self._build_result(step_output)
+                return
+
+            # Denormalize acoustic features
+            feat = step_output.acoustic_features * self._acoustic_std + self._acoustic_mean
+            tb_val = step_output.time_before[0].item()  # (batch, 1) -> scalar
+
+            # Skip leading silence on first token
+            if not first_token_seen:
+                first_token_seen = True
+                leading_frames = max(int(tb_val), 5)
+                streaming_decoder.skip_leading_frames(leading_frames)
+
+            chunk = streaming_decoder.decode_block(feat[0, 0], tb_val)
+            if chunk is not None and chunk.shape[-1] > 0:
+                yield chunk, 24000
+
+    def _build_result(self, outputs: "SyncTokGenerationOutput"):
+        """Build GenerationOutput from the final SyncTokGenerationOutput."""
+        ctx = self._generate_context
+        acoustic_features = outputs.acoustic_features * ctx["acoustic_std"] + ctx["acoustic_mean"]
+        token_decode_offset = ctx["token_decode_offset"]
+        encoded = acoustic_features[..., token_decode_offset:, :]
+        time_before = outputs.time_before[..., token_decode_offset:]
+        text = ctx["text"]
+        input_ids = ctx["input_ids"]
+        tokenizer = ctx["tokenizer"]
+
+        # Concatenate all streamed audio for .result
+        # (AudioStream doesn't store chunks — caller should collect them)
+        self.result = GenerationOutput(
+            audio=None,
+            text=text,
+            input_text_ids=input_ids,
+            input_str=[tokenizer.decode(input_ids[i]) for i in range(input_ids.shape[0])] if tokenizer else None,
+            output_str=[
+                tokenizer.decode(outputs.text_token_ids[i]) for i in range(outputs.text_token_ids.shape[0])
+            ] if tokenizer and outputs.text_token_ids is not None else None,
+            output_text_ids=outputs.text_token_ids,
+            acoustic_features=acoustic_features,
+            time_before=time_before,
+            prompt_text_tokens=input_ids,
+            llm_time=outputs.llm_time,
+            diffusion_time=outputs.diffusion_time,
+            logits=outputs.logits,
+            step_logs=outputs.step_logs,
+        )
+
 class CausalLMOutputWithPast(CausalLMOutputWithPastBase):
     def __init__(
         self,
@@ -653,10 +751,8 @@ def _generate(
         use_text_in_prompt: bool = False,
         verbose: bool = False,
         return_logits: bool = False,
-        on_token_ready: Callable[[torch.Tensor, int, int], None] | None = None,
-        on_token_ready_offset: int = 0,
         **kwargs,
-    ) -> SyncTokGenerationOutput:
+    ) -> Generator[StepOutput | SyncTokGenerationOutput, None, None]:
         start_header_id = self.tokenizer.convert_tokens_to_ids("<|start_header_id|>")
         end_header_id = self.tokenizer.convert_tokens_to_ids("<|end_header_id|>")
         eot_id = self.tokenizer.convert_tokens_to_ids("<|eot_id|>")
@@ -1102,11 +1198,12 @@ def _generate(
                 all_time_before.append(time_len_before)
                 last_time_before = time_len_before
 
-                # Fire streaming callback for predicted (non-prompt) tokens
-                if on_token_ready is not None and acoustic_feat_type == "predicted":
-                    token_idx = len(all_acoustic_features) - 1  # 0-based index into all_acoustic_features
-                    if token_idx >= on_token_ready_offset:
-                        on_token_ready(acoustic_features, time_len_before, token_idx - on_token_ready_offset)
+                # Yield predicted tokens for streaming
+                if acoustic_feat_type == "predicted":
+                    yield StepOutput(
+                        acoustic_features=acoustic_features,
+                        time_before=time_len_before,
+                    )
 
             model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs)
             if log_time:
@@ -1142,7 +1239,7 @@ def _generate(
             second_pass_time_after = torch.cat([scaled_time, torch.ones_like(scaled_time[:, :1])], dim=1)
 
             second_pass_options = replace(inference_options, speed_up_factor=None)
-            return self._generate(
+            yield from self._generate(
                 input_ids=input_ids,
                 input_lengths=input_lengths,
                 prompt_acoustic_features=prompt_acoustic_features,
@@ -1156,8 +1253,9 @@ def _generate(
                 return_logits=return_logits,
                 **kwargs,
             )
+            return
 
-        return SyncTokGenerationOutput(
+        yield SyncTokGenerationOutput(
             text_token_ids=torch.cat(all_output_token_ids, dim=1) if len(all_output_token_ids) > 0 else None,
             acoustic_features=torch.cat([f if f.ndim == 3 else f.unsqueeze(1) for f in all_acoustic_features], dim=1),
             time_before=torch.cat([f if f.ndim == 2 else f.unsqueeze(1) for f in all_time_before], dim=1),
@@ -1210,9 +1308,9 @@ def generate(  # type: ignore[override]
         use_text_in_prompt: bool = False,
         normalize_text: bool = True,
         verbose: bool = False,
-        on_audio_chunk: Callable[[torch.Tensor, int], None] | None = None,
+        stream: bool = False,
         streaming_cnn_window_size: int = 100,
-    ) -> GenerationOutput:
+    ) -> "GenerationOutput | AudioStream":
         # Auto-cast prompt tensors to model dtype (e.g. fp32 prompt + bf16 model)
         model_dtype = next(self.parameters()).dtype
         for field in prompt.__dataclass_fields__:
@@ -1294,42 +1392,7 @@ def generate(  # type: ignore[override]
         # Token offset: generated tokens start after prompt + transition tokens
         token_decode_offset = num_prompt_tokens + num_transition_steps - 1
 
-        # Set up E2E streaming: decode audio DURING _generate(), not after
-        streaming_decoder = None
-        streaming_audio_chunks: list[torch.Tensor] = []
-        on_token_ready_cb = None
-        first_token_seen = False
-
-        if on_audio_chunk is not None:
-            streaming_decoder = StreamingDecoder(
-                self.decoder,
-                min_block_frames=3,
-                cnn_window_size=streaming_cnn_window_size,
-            )
-
-            acoustic_std = self.config.acoustic_std
-            acoustic_mean = self.config.acoustic_mean
-
-            def on_token_ready_cb(acoustic_feat: torch.Tensor, time_before_val: torch.Tensor, rel_idx: int):
-                nonlocal first_token_seen
-                # Denormalize acoustic features (normally done after _generate())
-                feat = acoustic_feat * acoustic_std + acoustic_mean
-                tb_val = time_before_val[0].item()  # (batch, 1) -> scalar
-
-                # Skip leading silence on first token.
-                # Clamp to at least 5 frames (100ms) to avoid transition artifacts
-                # when the LLM predicts very short leading silence (1-2 frames).
-                if not first_token_seen:
-                    first_token_seen = True
-                    leading_frames = max(int(tb_val), 5)
-                    streaming_decoder.skip_leading_frames(leading_frames)
-
-                chunk = streaming_decoder.decode_block(feat[0, 0], tb_val)
-                if chunk is not None and chunk.shape[-1] > 0:
-                    streaming_audio_chunks.append(chunk)
-                    on_audio_chunk(chunk, 24000)
-
-        outputs: SyncTokGenerationOutput = self._generate(
+        gen = self._generate(
             input_ids=input_ids,
             text=text,
             input_lengths=input_lengths,
@@ -1342,22 +1405,31 @@ def on_token_ready_cb(acoustic_feat: torch.Tensor, time_before_val: torch.Tensor
             num_steps=input_ids.shape[-1] + num_extra_steps,
             inference_options=inference_options,
             use_text_in_prompt=use_text_in_prompt,
-            on_token_ready=on_token_ready_cb,
-            on_token_ready_offset=token_decode_offset,
         )
 
-        # Flush streaming decoder after _generate() completes
-        if streaming_decoder is not None:
-            # Flush with trailing frames from the last time_before
-            all_tb = outputs.time_before
-            if all_tb.shape[1] > 0:
-                trailing = all_tb[0, -1].item()
-            else:
-                trailing = 0
-            final_chunk = streaming_decoder.flush(trailing_frames=trailing)
-            if final_chunk is not None and final_chunk.shape[-1] > 0:
-                streaming_audio_chunks.append(final_chunk)
-                on_audio_chunk(final_chunk, 24000)
+        if stream:
+            audio_stream = AudioStream(
+                gen=gen,
+                decoder=self.decoder,
+                acoustic_mean=self.config.acoustic_mean,
+                acoustic_std=self.config.acoustic_std,
+                cnn_window_size=streaming_cnn_window_size,
+            )
+            # Stash info needed to build GenerationOutput after iteration
+            audio_stream._generate_context = {
+                "text": text,
+                "input_ids": input_ids,
+                "token_decode_offset": token_decode_offset,
+                "tokenizer": self.tokenizer,
+                "acoustic_std": self.config.acoustic_std,
+                "acoustic_mean": self.config.acoustic_mean,
+            }
+            return audio_stream
+
+        # Non-streaming: drain generator eagerly
+        for step in gen:
+            outputs = step
+        assert isinstance(outputs, SyncTokGenerationOutput)
 
         acoustic_features = outputs.acoustic_features * self.config.acoustic_std + self.config.acoustic_mean
         encoded = acoustic_features[..., token_decode_offset:, :]
@@ -1366,20 +1438,12 @@ def on_token_ready_cb(acoustic_feat: torch.Tensor, time_before_val: torch.Tensor
 
         for i in range(encoded.shape[0]):
             try:
-                if on_audio_chunk is not None and streaming_audio_chunks:
-                    # E2E streaming: audio was already emitted during _generate()
-                    if streaming_audio_chunks:
-                        wavs.append(torch.cat(streaming_audio_chunks, dim=-1).squeeze(0))
-                    else:
-                        wavs.append(None)
-                else:
-                    # Non-streaming decode: existing single-pass path
-                    wav = self._decode_wav(encoded[i], time_before=time_before[i]).squeeze(0, 1)
-                    # Trim leading silence. Clamp to at least 5 frames (100ms)
-                    # to avoid transition artifacts when time_before[0] is very small.
-                    min_trim_frames = max(int(time_before[i][0].item()), 5)
-                    wav = wav[..., min_trim_frames * 480 :]
-                    wavs.append(wav)
+                wav = self._decode_wav(encoded[i], time_before=time_before[i]).squeeze(0, 1)
+                # Trim leading silence. Clamp to at least 5 frames (100ms)
+                # to avoid transition artifacts when time_before[0] is very small.
+                min_trim_frames = max(int(time_before[i][0].item()), 5)
+                wav = wav[..., min_trim_frames * 480 :]
+                wavs.append(wav)
             except Exception:
                 wavs.append(None)
 
diff --git a/tests/integration_test.log b/tests/integration_test.log
new file mode 100644
index 0000000..53d104b
--- /dev/null
+++ b/tests/integration_test.log
@@ -0,0 +1,168 @@
+=== Running integration tests ===
+/mnt/weka/sharath/anaconda3/envs/media-pipeline/lib/python3.11/site-packages/pytest_asyncio/plugin.py:217: PytestDeprecationWarning: The configuration option "asyncio_default_fixture_loop_scope" is unset.
+The event loop scope for asynchronous fixtures will default to the fixture caching scope. Future versions of pytest-asyncio will default the loop scope for asynchronous fixtures to function scope. Set the default fixture loop scope explicitly in order to avoid unexpected behavior in the future. Valid fixture loop scopes are: "function", "class", "module", "package", "session"
+
+  warnings.warn(PytestDeprecationWarning(_DEFAULT_FIXTURE_LOOP_SCOPE_UNSET))
+============================= test session starts ==============================
+platform linux -- Python 3.11.9, pytest-8.3.5, pluggy-1.5.0 -- /mnt/weka/sharath/anaconda3/envs/media-pipeline/bin/python
+cachedir: .pytest_cache
+rootdir: /mnt/weka/sharath/projects/tada
+configfile: pyproject.toml
+plugins: hydra-core-1.3.2, asyncio-0.26.0, typeguard-4.4.2, zarr-3.1.5, anyio-4.9.0
+asyncio: mode=Mode.STRICT, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
+collecting ... collected 15 items / 10 deselected / 5 selected
+
+tests/test_streaming.py::TestStreamingIntegration::test_non_streaming_unchanged `torch_dtype` is deprecated! Use `dtype` instead!
+Some weights of the model checkpoint at HumeAI/tada-1b were not used when initializing TadaForCausalLM: ['_decoder.decoder_proj.bias', '_decoder.decoder_proj.weight', '_decoder.local_attention_decoder.final_norm.bias', '_decoder.local_attention_decoder.final_norm.weight', '_decoder.local_attention_decoder.layers.0.ffn.0.bias', '_decoder.local_attention_decoder.layers.0.ffn.0.weight', '_decoder.local_attention_decoder.layers.0.ffn.3.bias', '_decoder.local_attention_decoder.layers.0.ffn.3.weight', '_decoder.local_attention_decoder.layers.0.norm.bias', '_decoder.local_attention_decoder.layers.0.norm.weight', '_decoder.local_attention_decoder.layers.0.self_attn._precomputed_mask', '_decoder.local_attention_decoder.layers.0.self_attn.layer_norm.bias', '_decoder.local_attention_decoder.layers.0.self_attn.layer_norm.weight', '_decoder.local_attention_decoder.layers.0.self_attn.out_proj.bias', '_decoder.local_attention_decoder.layers.0.self_attn.out_proj.weight', '_decoder.local_attention_decoder.layers.0.self_attn.qkv.bias', '_decoder.local_attention_decoder.layers.0.self_attn.qkv.weight', '_decoder.local_attention_decoder.layers.0.self_attn.rope_freqs', '_decoder.local_attention_decoder.layers.1.ffn.0.bias', '_decoder.local_attention_decoder.layers.1.ffn.0.weight', '_decoder.local_attention_decoder.layers.1.ffn.3.bias', '_decoder.local_attention_decoder.layers.1.ffn.3.weight', '_decoder.local_attention_decoder.layers.1.norm.bias', '_decoder.local_attention_decoder.layers.1.norm.weight', '_decoder.local_attention_decoder.layers.1.self_attn._precomputed_mask', '_decoder.local_attention_decoder.layers.1.self_attn.layer_norm.bias', '_decoder.local_attention_decoder.layers.1.self_attn.layer_norm.weight', '_decoder.local_attention_decoder.layers.1.self_attn.out_proj.bias', '_decoder.local_attention_decoder.layers.1.self_attn.out_proj.weight', '_decoder.local_attention_decoder.layers.1.self_attn.qkv.bias', '_decoder.local_attention_decoder.layers.1.self_attn.qkv.weight', '_decoder.local_attention_decoder.layers.1.self_attn.rope_freqs', '_decoder.local_attention_decoder.layers.2.ffn.0.bias', '_decoder.local_attention_decoder.layers.2.ffn.0.weight', '_decoder.local_attention_decoder.layers.2.ffn.3.bias', '_decoder.local_attention_decoder.layers.2.ffn.3.weight', '_decoder.local_attention_decoder.layers.2.norm.bias', '_decoder.local_attention_decoder.layers.2.norm.weight', '_decoder.local_attention_decoder.layers.2.self_attn._precomputed_mask', '_decoder.local_attention_decoder.layers.2.self_attn.layer_norm.bias', '_decoder.local_attention_decoder.layers.2.self_attn.layer_norm.weight', '_decoder.local_attention_decoder.layers.2.self_attn.out_proj.bias', '_decoder.local_attention_decoder.layers.2.self_attn.out_proj.weight', '_decoder.local_attention_decoder.layers.2.self_attn.qkv.bias', '_decoder.local_attention_decoder.layers.2.self_attn.qkv.weight', '_decoder.local_attention_decoder.layers.2.self_attn.rope_freqs', '_decoder.local_attention_decoder.layers.3.ffn.0.bias', '_decoder.local_attention_decoder.layers.3.ffn.0.weight', '_decoder.local_attention_decoder.layers.3.ffn.3.bias', '_decoder.local_attention_decoder.layers.3.ffn.3.weight', '_decoder.local_attention_decoder.layers.3.norm.bias', '_decoder.local_attention_decoder.layers.3.norm.weight', '_decoder.local_attention_decoder.layers.3.self_attn._precomputed_mask', '_decoder.local_attention_decoder.layers.3.self_attn.layer_norm.bias', '_decoder.local_attention_decoder.layers.3.self_attn.layer_norm.weight', '_decoder.local_attention_decoder.layers.3.self_attn.out_proj.bias', '_decoder.local_attention_decoder.layers.3.self_attn.out_proj.weight', '_decoder.local_attention_decoder.layers.3.self_attn.qkv.bias', '_decoder.local_attention_decoder.layers.3.self_attn.qkv.weight', '_decoder.local_attention_decoder.layers.3.self_attn.rope_freqs', '_decoder.local_attention_decoder.layers.4.ffn.0.bias', '_decoder.local_attention_decoder.layers.4.ffn.0.weight', '_decoder.local_attention_decoder.layers.4.ffn.3.bias', '_decoder.local_attention_decoder.layers.4.ffn.3.weight', '_decoder.local_attention_decoder.layers.4.norm.bias', '_decoder.local_attention_decoder.layers.4.norm.weight', '_decoder.local_attention_decoder.layers.4.self_attn._precomputed_mask', '_decoder.local_attention_decoder.layers.4.self_attn.layer_norm.bias', '_decoder.local_attention_decoder.layers.4.self_attn.layer_norm.weight', '_decoder.local_attention_decoder.layers.4.self_attn.out_proj.bias', '_decoder.local_attention_decoder.layers.4.self_attn.out_proj.weight', '_decoder.local_attention_decoder.layers.4.self_attn.qkv.bias', '_decoder.local_attention_decoder.layers.4.self_attn.qkv.weight', '_decoder.local_attention_decoder.layers.4.self_attn.rope_freqs', '_decoder.local_attention_decoder.layers.5.ffn.0.bias', '_decoder.local_attention_decoder.layers.5.ffn.0.weight', '_decoder.local_attention_decoder.layers.5.ffn.3.bias', '_decoder.local_attention_decoder.layers.5.ffn.3.weight', '_decoder.local_attention_decoder.layers.5.norm.bias', '_decoder.local_attention_decoder.layers.5.norm.weight', '_decoder.local_attention_decoder.layers.5.self_attn._precomputed_mask', '_decoder.local_attention_decoder.layers.5.self_attn.layer_norm.bias', '_decoder.local_attention_decoder.layers.5.self_attn.layer_norm.weight', '_decoder.local_attention_decoder.layers.5.self_attn.out_proj.bias', '_decoder.local_attention_decoder.layers.5.self_attn.out_proj.weight', '_decoder.local_attention_decoder.layers.5.self_attn.qkv.bias', '_decoder.local_attention_decoder.layers.5.self_attn.qkv.weight', '_decoder.local_attention_decoder.layers.5.self_attn.rope_freqs', '_decoder.wav_decoder.model.0.bias', '_decoder.wav_decoder.model.0.parametrizations.weight.original0', '_decoder.wav_decoder.model.0.parametrizations.weight.original1', '_decoder.wav_decoder.model.1.block.0.alpha', '_decoder.wav_decoder.model.1.block.1.bias', '_decoder.wav_decoder.model.1.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.1.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.1.block.2.block.0.alpha', '_decoder.wav_decoder.model.1.block.2.block.1.bias', '_decoder.wav_decoder.model.1.block.2.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.1.block.2.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.1.block.2.block.2.alpha', '_decoder.wav_decoder.model.1.block.2.block.3.bias', '_decoder.wav_decoder.model.1.block.2.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.1.block.2.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.1.block.3.block.0.alpha', '_decoder.wav_decoder.model.1.block.3.block.1.bias', '_decoder.wav_decoder.model.1.block.3.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.1.block.3.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.1.block.3.block.2.alpha', '_decoder.wav_decoder.model.1.block.3.block.3.bias', '_decoder.wav_decoder.model.1.block.3.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.1.block.3.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.1.block.4.block.0.alpha', '_decoder.wav_decoder.model.1.block.4.block.1.bias', '_decoder.wav_decoder.model.1.block.4.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.1.block.4.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.1.block.4.block.2.alpha', '_decoder.wav_decoder.model.1.block.4.block.3.bias', '_decoder.wav_decoder.model.1.block.4.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.1.block.4.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.2.block.0.alpha', '_decoder.wav_decoder.model.2.block.1.bias', '_decoder.wav_decoder.model.2.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.2.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.2.block.2.block.0.alpha', '_decoder.wav_decoder.model.2.block.2.block.1.bias', '_decoder.wav_decoder.model.2.block.2.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.2.block.2.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.2.block.2.block.2.alpha', '_decoder.wav_decoder.model.2.block.2.block.3.bias', '_decoder.wav_decoder.model.2.block.2.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.2.block.2.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.2.block.3.block.0.alpha', '_decoder.wav_decoder.model.2.block.3.block.1.bias', '_decoder.wav_decoder.model.2.block.3.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.2.block.3.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.2.block.3.block.2.alpha', '_decoder.wav_decoder.model.2.block.3.block.3.bias', '_decoder.wav_decoder.model.2.block.3.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.2.block.3.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.2.block.4.block.0.alpha', '_decoder.wav_decoder.model.2.block.4.block.1.bias', '_decoder.wav_decoder.model.2.block.4.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.2.block.4.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.2.block.4.block.2.alpha', '_decoder.wav_decoder.model.2.block.4.block.3.bias', '_decoder.wav_decoder.model.2.block.4.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.2.block.4.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.3.block.0.alpha', '_decoder.wav_decoder.model.3.block.1.bias', '_decoder.wav_decoder.model.3.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.3.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.3.block.2.block.0.alpha', '_decoder.wav_decoder.model.3.block.2.block.1.bias', '_decoder.wav_decoder.model.3.block.2.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.3.block.2.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.3.block.2.block.2.alpha', '_decoder.wav_decoder.model.3.block.2.block.3.bias', '_decoder.wav_decoder.model.3.block.2.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.3.block.2.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.3.block.3.block.0.alpha', '_decoder.wav_decoder.model.3.block.3.block.1.bias', '_decoder.wav_decoder.model.3.block.3.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.3.block.3.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.3.block.3.block.2.alpha', '_decoder.wav_decoder.model.3.block.3.block.3.bias', '_decoder.wav_decoder.model.3.block.3.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.3.block.3.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.3.block.4.block.0.alpha', '_decoder.wav_decoder.model.3.block.4.block.1.bias', '_decoder.wav_decoder.model.3.block.4.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.3.block.4.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.3.block.4.block.2.alpha', '_decoder.wav_decoder.model.3.block.4.block.3.bias', '_decoder.wav_decoder.model.3.block.4.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.3.block.4.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.4.block.0.alpha', '_decoder.wav_decoder.model.4.block.1.bias', '_decoder.wav_decoder.model.4.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.4.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.4.block.2.block.0.alpha', '_decoder.wav_decoder.model.4.block.2.block.1.bias', '_decoder.wav_decoder.model.4.block.2.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.4.block.2.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.4.block.2.block.2.alpha', '_decoder.wav_decoder.model.4.block.2.block.3.bias', '_decoder.wav_decoder.model.4.block.2.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.4.block.2.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.4.block.3.block.0.alpha', '_decoder.wav_decoder.model.4.block.3.block.1.bias', '_decoder.wav_decoder.model.4.block.3.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.4.block.3.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.4.block.3.block.2.alpha', '_decoder.wav_decoder.model.4.block.3.block.3.bias', '_decoder.wav_decoder.model.4.block.3.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.4.block.3.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.4.block.4.block.0.alpha', '_decoder.wav_decoder.model.4.block.4.block.1.bias', '_decoder.wav_decoder.model.4.block.4.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.4.block.4.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.4.block.4.block.2.alpha', '_decoder.wav_decoder.model.4.block.4.block.3.bias', '_decoder.wav_decoder.model.4.block.4.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.4.block.4.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.5.alpha', '_decoder.wav_decoder.model.6.bias', '_decoder.wav_decoder.model.6.parametrizations.weight.original0', '_decoder.wav_decoder.model.6.parametrizations.weight.original1']
+- This IS expected if you are initializing TadaForCausalLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
+- This IS NOT expected if you are initializing TadaForCausalLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
+🚨 Config not found for parakeet. You can manually add it to HARDCODED_CONFIG_FOR_MODELS in utils/auto_docstring.py
+🚨 Config not found for parakeet. You can manually add it to HARDCODED_CONFIG_FOR_MODELS in utils/auto_docstring.py
+🚨 Config not found for parakeet. You can manually add it to HARDCODED_CONFIG_FOR_MODELS in utils/auto_docstring.py
+PASSED
+tests/test_streaming.py::TestStreamingIntegration::test_streaming_produces_chunks PASSED
+tests/test_streaming.py::TestStreamingIntegration::test_streaming_vs_nonstreaming_similar_length PASSED
+tests/test_streaming.py::TestStreamingIntegration::test_streaming_early_break PASSED
+tests/test_streaming.py::TestGenerateAudios::test_generate_streaming_audios Some weights of the model checkpoint at HumeAI/tada-1b were not used when initializing TadaForCausalLM: ['_decoder.decoder_proj.bias', '_decoder.decoder_proj.weight', '_decoder.local_attention_decoder.final_norm.bias', '_decoder.local_attention_decoder.final_norm.weight', '_decoder.local_attention_decoder.layers.0.ffn.0.bias', '_decoder.local_attention_decoder.layers.0.ffn.0.weight', '_decoder.local_attention_decoder.layers.0.ffn.3.bias', '_decoder.local_attention_decoder.layers.0.ffn.3.weight', '_decoder.local_attention_decoder.layers.0.norm.bias', '_decoder.local_attention_decoder.layers.0.norm.weight', '_decoder.local_attention_decoder.layers.0.self_attn._precomputed_mask', '_decoder.local_attention_decoder.layers.0.self_attn.layer_norm.bias', '_decoder.local_attention_decoder.layers.0.self_attn.layer_norm.weight', '_decoder.local_attention_decoder.layers.0.self_attn.out_proj.bias', '_decoder.local_attention_decoder.layers.0.self_attn.out_proj.weight', '_decoder.local_attention_decoder.layers.0.self_attn.qkv.bias', '_decoder.local_attention_decoder.layers.0.self_attn.qkv.weight', '_decoder.local_attention_decoder.layers.0.self_attn.rope_freqs', '_decoder.local_attention_decoder.layers.1.ffn.0.bias', '_decoder.local_attention_decoder.layers.1.ffn.0.weight', '_decoder.local_attention_decoder.layers.1.ffn.3.bias', '_decoder.local_attention_decoder.layers.1.ffn.3.weight', '_decoder.local_attention_decoder.layers.1.norm.bias', '_decoder.local_attention_decoder.layers.1.norm.weight', '_decoder.local_attention_decoder.layers.1.self_attn._precomputed_mask', '_decoder.local_attention_decoder.layers.1.self_attn.layer_norm.bias', '_decoder.local_attention_decoder.layers.1.self_attn.layer_norm.weight', '_decoder.local_attention_decoder.layers.1.self_attn.out_proj.bias', '_decoder.local_attention_decoder.layers.1.self_attn.out_proj.weight', '_decoder.local_attention_decoder.layers.1.self_attn.qkv.bias', '_decoder.local_attention_decoder.layers.1.self_attn.qkv.weight', '_decoder.local_attention_decoder.layers.1.self_attn.rope_freqs', '_decoder.local_attention_decoder.layers.2.ffn.0.bias', '_decoder.local_attention_decoder.layers.2.ffn.0.weight', '_decoder.local_attention_decoder.layers.2.ffn.3.bias', '_decoder.local_attention_decoder.layers.2.ffn.3.weight', '_decoder.local_attention_decoder.layers.2.norm.bias', '_decoder.local_attention_decoder.layers.2.norm.weight', '_decoder.local_attention_decoder.layers.2.self_attn._precomputed_mask', '_decoder.local_attention_decoder.layers.2.self_attn.layer_norm.bias', '_decoder.local_attention_decoder.layers.2.self_attn.layer_norm.weight', '_decoder.local_attention_decoder.layers.2.self_attn.out_proj.bias', '_decoder.local_attention_decoder.layers.2.self_attn.out_proj.weight', '_decoder.local_attention_decoder.layers.2.self_attn.qkv.bias', '_decoder.local_attention_decoder.layers.2.self_attn.qkv.weight', '_decoder.local_attention_decoder.layers.2.self_attn.rope_freqs', '_decoder.local_attention_decoder.layers.3.ffn.0.bias', '_decoder.local_attention_decoder.layers.3.ffn.0.weight', '_decoder.local_attention_decoder.layers.3.ffn.3.bias', '_decoder.local_attention_decoder.layers.3.ffn.3.weight', '_decoder.local_attention_decoder.layers.3.norm.bias', '_decoder.local_attention_decoder.layers.3.norm.weight', '_decoder.local_attention_decoder.layers.3.self_attn._precomputed_mask', '_decoder.local_attention_decoder.layers.3.self_attn.layer_norm.bias', '_decoder.local_attention_decoder.layers.3.self_attn.layer_norm.weight', '_decoder.local_attention_decoder.layers.3.self_attn.out_proj.bias', '_decoder.local_attention_decoder.layers.3.self_attn.out_proj.weight', '_decoder.local_attention_decoder.layers.3.self_attn.qkv.bias', '_decoder.local_attention_decoder.layers.3.self_attn.qkv.weight', '_decoder.local_attention_decoder.layers.3.self_attn.rope_freqs', '_decoder.local_attention_decoder.layers.4.ffn.0.bias', '_decoder.local_attention_decoder.layers.4.ffn.0.weight', '_decoder.local_attention_decoder.layers.4.ffn.3.bias', '_decoder.local_attention_decoder.layers.4.ffn.3.weight', '_decoder.local_attention_decoder.layers.4.norm.bias', '_decoder.local_attention_decoder.layers.4.norm.weight', '_decoder.local_attention_decoder.layers.4.self_attn._precomputed_mask', '_decoder.local_attention_decoder.layers.4.self_attn.layer_norm.bias', '_decoder.local_attention_decoder.layers.4.self_attn.layer_norm.weight', '_decoder.local_attention_decoder.layers.4.self_attn.out_proj.bias', '_decoder.local_attention_decoder.layers.4.self_attn.out_proj.weight', '_decoder.local_attention_decoder.layers.4.self_attn.qkv.bias', '_decoder.local_attention_decoder.layers.4.self_attn.qkv.weight', '_decoder.local_attention_decoder.layers.4.self_attn.rope_freqs', '_decoder.local_attention_decoder.layers.5.ffn.0.bias', '_decoder.local_attention_decoder.layers.5.ffn.0.weight', '_decoder.local_attention_decoder.layers.5.ffn.3.bias', '_decoder.local_attention_decoder.layers.5.ffn.3.weight', '_decoder.local_attention_decoder.layers.5.norm.bias', '_decoder.local_attention_decoder.layers.5.norm.weight', '_decoder.local_attention_decoder.layers.5.self_attn._precomputed_mask', '_decoder.local_attention_decoder.layers.5.self_attn.layer_norm.bias', '_decoder.local_attention_decoder.layers.5.self_attn.layer_norm.weight', '_decoder.local_attention_decoder.layers.5.self_attn.out_proj.bias', '_decoder.local_attention_decoder.layers.5.self_attn.out_proj.weight', '_decoder.local_attention_decoder.layers.5.self_attn.qkv.bias', '_decoder.local_attention_decoder.layers.5.self_attn.qkv.weight', '_decoder.local_attention_decoder.layers.5.self_attn.rope_freqs', '_decoder.wav_decoder.model.0.bias', '_decoder.wav_decoder.model.0.parametrizations.weight.original0', '_decoder.wav_decoder.model.0.parametrizations.weight.original1', '_decoder.wav_decoder.model.1.block.0.alpha', '_decoder.wav_decoder.model.1.block.1.bias', '_decoder.wav_decoder.model.1.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.1.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.1.block.2.block.0.alpha', '_decoder.wav_decoder.model.1.block.2.block.1.bias', '_decoder.wav_decoder.model.1.block.2.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.1.block.2.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.1.block.2.block.2.alpha', '_decoder.wav_decoder.model.1.block.2.block.3.bias', '_decoder.wav_decoder.model.1.block.2.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.1.block.2.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.1.block.3.block.0.alpha', '_decoder.wav_decoder.model.1.block.3.block.1.bias', '_decoder.wav_decoder.model.1.block.3.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.1.block.3.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.1.block.3.block.2.alpha', '_decoder.wav_decoder.model.1.block.3.block.3.bias', '_decoder.wav_decoder.model.1.block.3.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.1.block.3.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.1.block.4.block.0.alpha', '_decoder.wav_decoder.model.1.block.4.block.1.bias', '_decoder.wav_decoder.model.1.block.4.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.1.block.4.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.1.block.4.block.2.alpha', '_decoder.wav_decoder.model.1.block.4.block.3.bias', '_decoder.wav_decoder.model.1.block.4.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.1.block.4.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.2.block.0.alpha', '_decoder.wav_decoder.model.2.block.1.bias', '_decoder.wav_decoder.model.2.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.2.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.2.block.2.block.0.alpha', '_decoder.wav_decoder.model.2.block.2.block.1.bias', '_decoder.wav_decoder.model.2.block.2.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.2.block.2.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.2.block.2.block.2.alpha', '_decoder.wav_decoder.model.2.block.2.block.3.bias', '_decoder.wav_decoder.model.2.block.2.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.2.block.2.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.2.block.3.block.0.alpha', '_decoder.wav_decoder.model.2.block.3.block.1.bias', '_decoder.wav_decoder.model.2.block.3.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.2.block.3.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.2.block.3.block.2.alpha', '_decoder.wav_decoder.model.2.block.3.block.3.bias', '_decoder.wav_decoder.model.2.block.3.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.2.block.3.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.2.block.4.block.0.alpha', '_decoder.wav_decoder.model.2.block.4.block.1.bias', '_decoder.wav_decoder.model.2.block.4.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.2.block.4.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.2.block.4.block.2.alpha', '_decoder.wav_decoder.model.2.block.4.block.3.bias', '_decoder.wav_decoder.model.2.block.4.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.2.block.4.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.3.block.0.alpha', '_decoder.wav_decoder.model.3.block.1.bias', '_decoder.wav_decoder.model.3.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.3.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.3.block.2.block.0.alpha', '_decoder.wav_decoder.model.3.block.2.block.1.bias', '_decoder.wav_decoder.model.3.block.2.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.3.block.2.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.3.block.2.block.2.alpha', '_decoder.wav_decoder.model.3.block.2.block.3.bias', '_decoder.wav_decoder.model.3.block.2.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.3.block.2.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.3.block.3.block.0.alpha', '_decoder.wav_decoder.model.3.block.3.block.1.bias', '_decoder.wav_decoder.model.3.block.3.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.3.block.3.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.3.block.3.block.2.alpha', '_decoder.wav_decoder.model.3.block.3.block.3.bias', '_decoder.wav_decoder.model.3.block.3.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.3.block.3.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.3.block.4.block.0.alpha', '_decoder.wav_decoder.model.3.block.4.block.1.bias', '_decoder.wav_decoder.model.3.block.4.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.3.block.4.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.3.block.4.block.2.alpha', '_decoder.wav_decoder.model.3.block.4.block.3.bias', '_decoder.wav_decoder.model.3.block.4.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.3.block.4.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.4.block.0.alpha', '_decoder.wav_decoder.model.4.block.1.bias', '_decoder.wav_decoder.model.4.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.4.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.4.block.2.block.0.alpha', '_decoder.wav_decoder.model.4.block.2.block.1.bias', '_decoder.wav_decoder.model.4.block.2.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.4.block.2.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.4.block.2.block.2.alpha', '_decoder.wav_decoder.model.4.block.2.block.3.bias', '_decoder.wav_decoder.model.4.block.2.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.4.block.2.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.4.block.3.block.0.alpha', '_decoder.wav_decoder.model.4.block.3.block.1.bias', '_decoder.wav_decoder.model.4.block.3.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.4.block.3.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.4.block.3.block.2.alpha', '_decoder.wav_decoder.model.4.block.3.block.3.bias', '_decoder.wav_decoder.model.4.block.3.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.4.block.3.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.4.block.4.block.0.alpha', '_decoder.wav_decoder.model.4.block.4.block.1.bias', '_decoder.wav_decoder.model.4.block.4.block.1.parametrizations.weight.original0', '_decoder.wav_decoder.model.4.block.4.block.1.parametrizations.weight.original1', '_decoder.wav_decoder.model.4.block.4.block.2.alpha', '_decoder.wav_decoder.model.4.block.4.block.3.bias', '_decoder.wav_decoder.model.4.block.4.block.3.parametrizations.weight.original0', '_decoder.wav_decoder.model.4.block.4.block.3.parametrizations.weight.original1', '_decoder.wav_decoder.model.5.alpha', '_decoder.wav_decoder.model.6.bias', '_decoder.wav_decoder.model.6.parametrizations.weight.original0', '_decoder.wav_decoder.model.6.parametrizations.weight.original1']
+- This IS expected if you are initializing TadaForCausalLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
+- This IS NOT expected if you are initializing TadaForCausalLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
+
+--- Generating streaming audio 1: Hello, this is a demonstration of streaming text t...
+  chunk: 0.78s
+  chunk: 0.08s
+  chunk: 0.88s
+  chunk: 0.20s
+  chunk: 0.14s
+  chunk: 0.52s
+  chunk: 0.46s
+  chunk: 0.30s
+  chunk: 0.44s
+  chunk: 0.34s
+  chunk: 0.22s
+  chunk: 0.44s
+  chunk: 0.08s
+  chunk: 0.38s
+  saved: tests/output/streaming_1.wav (5.26s)
+
+--- Generating streaming audio 2: Please call Stella. Ask her to bring these things ...
+  chunk: 1.18s
+  chunk: 0.08s
+  chunk: 0.26s
+  chunk: 0.26s
+  chunk: 1.12s
+  chunk: 0.06s
+  chunk: 0.30s
+  chunk: 0.12s
+  chunk: 0.14s
+  chunk: 0.18s
+  chunk: 0.32s
+  chunk: 0.26s
+  chunk: 0.16s
+  chunk: 0.16s
+  chunk: 0.14s
+  chunk: 0.20s
+  chunk: 0.26s
+  chunk: 0.08s
+  chunk: 0.38s
+  saved: tests/output/streaming_2.wav (5.66s)
+
+--- Generating streaming audio 3: The quick brown fox jumps over the lazy dog....
+  chunk: 0.88s
+  chunk: 0.06s
+  chunk: 0.20s
+  chunk: 0.32s
+  chunk: 0.36s
+  chunk: 0.60s
+  chunk: 0.38s
+  chunk: 0.26s
+  chunk: 0.12s
+  chunk: 0.38s
+  chunk: 0.36s
+  chunk: 0.08s
+  chunk: 0.38s
+  saved: tests/output/streaming_3.wav (4.38s)
+
+--- Generating streaming audio 4: In the beginning, there was silence. Then came the...
+  chunk: 1.18s
+  chunk: 0.06s
+  chunk: 0.10s
+  chunk: 0.26s
+  chunk: 0.60s
+  chunk: 0.06s
+  chunk: 0.16s
+  chunk: 0.22s
+  chunk: 1.42s
+  chunk: 0.06s
+  chunk: 0.36s
+  chunk: 0.22s
+  chunk: 0.14s
+  chunk: 0.76s
+  chunk: 0.14s
+  chunk: 0.42s
+  chunk: 0.18s
+  chunk: 0.32s
+  chunk: 0.08s
+  chunk: 0.38s
+  saved: tests/output/streaming_4.wav (7.12s)
+
+--- Generating streaming audio 5: Technology is best when it brings people together....
+  chunk: 1.20s
+  chunk: 0.26s
+  chunk: 0.38s
+  chunk: 0.26s
+  chunk: 0.74s
+  chunk: 0.08s
+  chunk: 0.20s
+  chunk: 0.36s
+  chunk: 0.50s
+  chunk: 0.24s
+  chunk: 0.08s
+  chunk: 0.38s
+  saved: tests/output/streaming_5.wav (4.68s)
+
+--- Generating non-streaming audio 1: Hello, this is a demonstration of streaming text t...
+  saved: tests/output/nonstreaming_1.wav (4.16s)
+
+--- Generating non-streaming audio 2: Please call Stella. Ask her to bring these things ...
+  saved: tests/output/nonstreaming_2.wav (3.86s)
+
+--- Generating non-streaming audio 3: The quick brown fox jumps over the lazy dog....
+  saved: tests/output/nonstreaming_3.wav (3.34s)
+
+--- Generating non-streaming audio 4: In the beginning, there was silence. Then came the...
+  saved: tests/output/nonstreaming_4.wav (5.64s)
+
+--- Generating non-streaming audio 5: Technology is best when it brings people together....
+  saved: tests/output/nonstreaming_5.wav (2.96s)
+PASSED
+
+=============================== warnings summary ===============================
+tests/test_streaming.py:219
+  /mnt/weka/sharath/projects/tada/tests/test_streaming.py:219: PytestUnknownMarkWarning: Unknown pytest.mark.integration - is this a typo?  You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
+    @pytest.mark.integration
+
+tests/test_streaming.py:315
+  /mnt/weka/sharath/projects/tada/tests/test_streaming.py:315: PytestUnknownMarkWarning: Unknown pytest.mark.integration - is this a typo?  You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
+    @pytest.mark.integration
+
+tests/test_streaming.py::TestStreamingIntegration::test_non_streaming_unchanged
+  /mnt/weka/sharath/anaconda3/envs/media-pipeline/lib/python3.11/site-packages/librosa/core/intervals.py:15: DeprecationWarning: path is deprecated. Use files() instead. Refer to https://importlib-resources.readthedocs.io/en/latest/using.html#migrating-from-legacy for migration advice.
+    with resources.path("librosa.core", "intervals.msgpack") as imsgpack:
+
+-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
+================ 5 passed, 10 deselected, 3 warnings in 51.75s =================
+=== Done ===
+Output files:
+total 2228
+drwxrwxr-x 1 sharath sharath      0 Mar 21 17:42 .
+drwxrwxr-x 1 sharath sharath      0 Mar 21 17:41 ..
+-rw-rw-r-- 1 sharath sharath 199746 Mar 21 17:41 nonstreaming_1.wav
+-rw-rw-r-- 1 sharath sharath 185346 Mar 21 17:41 nonstreaming_2.wav
+-rw-rw-r-- 1 sharath sharath 160386 Mar 21 17:41 nonstreaming_3.wav
+-rw-rw-r-- 1 sharath sharath 270786 Mar 21 17:42 nonstreaming_4.wav
+-rw-rw-r-- 1 sharath sharath 142146 Mar 21 17:42 nonstreaming_5.wav
+-rw-rw-r-- 1 sharath sharath 252546 Mar 21 17:41 streaming_1.wav
+-rw-rw-r-- 1 sharath sharath 271746 Mar 21 17:41 streaming_2.wav
+-rw-rw-r-- 1 sharath sharath 210306 Mar 21 17:41 streaming_3.wav
+-rw-rw-r-- 1 sharath sharath 341826 Mar 21 17:41 streaming_4.wav
+-rw-rw-r-- 1 sharath sharath 224706 Mar 21 17:41 streaming_5.wav
diff --git a/tests/run_integration.sh b/tests/run_integration.sh
new file mode 100755
index 0000000..99d7100
--- /dev/null
+++ b/tests/run_integration.sh
@@ -0,0 +1,16 @@
+#!/bin/bash
+#SBATCH --job-name=tada-stream-test
+#SBATCH --partition=eval
+#SBATCH --gres=gpu:1
+#SBATCH --time=00:20:00
+#SBATCH --output=/mnt/weka/sharath/projects/tada/tests/integration_test.log
+
+cd /mnt/weka/sharath/projects/tada
+export PATH="/mnt/weka/sharath/anaconda3/envs/media-pipeline/bin:$PATH"
+
+echo "=== Running integration tests ==="
+python -m pytest tests/test_streaming.py -m integration -v -s 2>&1
+
+echo "=== Done ==="
+echo "Output files:"
+ls -la tests/output/ 2>/dev/null
diff --git a/tests/test_streaming.py b/tests/test_streaming.py
new file mode 100644
index 0000000..946d51c
--- /dev/null
+++ b/tests/test_streaming.py
@@ -0,0 +1,381 @@
+"""Tests for streaming audio generation.
+
+Unit tests run without GPU/model (mock-based).
+Integration tests (marked with @pytest.mark.integration) require GPU + model weights.
+Run integration tests with: pytest tests/test_streaming.py -m integration -s
+"""
+
+import pytest
+import torch
+
+from tada.modules.decoder import StreamingDecoder, _create_segment_attention_mask
+from tada.modules.tada import AudioStream, GenerationOutput, StepOutput, SyncTokGenerationOutput
+
+
+# ---------------------------------------------------------------------------
+# Unit tests — no GPU, no model weights required
+# ---------------------------------------------------------------------------
+
+
+class TestStepOutput:
+    def test_fields(self):
+        feat = torch.randn(1, 1, 512)
+        tb = torch.tensor([[5]])
+        s = StepOutput(acoustic_features=feat, time_before=tb)
+        assert s.acoustic_features is feat
+        assert s.time_before is tb
+
+
+class TestAudioStream:
+    """Test AudioStream with a fake generator (no real model needed)."""
+
+    def _make_fake_gen(self, num_steps=5, acoustic_dim=512):
+        """Create a generator that yields StepOutputs then SyncTokGenerationOutput."""
+        all_feats = []
+        all_tb = []
+        for i in range(num_steps):
+            feat = torch.randn(1, 1, acoustic_dim)
+            tb = torch.tensor([[8]])  # 8 frames before each token
+            all_feats.append(feat)
+            all_tb.append(tb)
+            yield StepOutput(acoustic_features=feat, time_before=tb)
+
+        # Final output
+        yield SyncTokGenerationOutput(
+            text_token_ids=torch.zeros(1, num_steps, dtype=torch.long),
+            acoustic_features=torch.cat(all_feats, dim=1),
+            time_before=torch.cat(all_tb, dim=1),
+            llm_time=torch.tensor(0.1),
+            diffusion_time=torch.tensor(0.05),
+            logits=None,
+            step_logs=[],
+        )
+
+    def _make_audio_stream(self, num_steps=10):
+        """Create AudioStream with a real (tiny) Decoder for integration."""
+        from tada.modules.decoder import Decoder, DecoderConfig
+
+        config = DecoderConfig(
+            embed_dim=512,
+            hidden_dim=64,  # tiny for testing
+            num_attn_layers=1,
+            num_attn_heads=2,
+            attn_dim_feedforward=128,
+            strides=[2, 2, 2, 2],  # 16x upsampling (fast)
+            block_attention="v2",
+        )
+        decoder = Decoder(config)
+        decoder.eval()
+
+        gen = self._make_fake_gen(num_steps=num_steps)
+        stream = AudioStream(
+            gen=gen,
+            decoder=decoder,
+            acoustic_mean=0.0,
+            acoustic_std=1.5,
+            cnn_window_size=50,
+        )
+        # Mock the context needed for _build_result
+        stream._generate_context = {
+            "text": ["test"],
+            "input_ids": torch.zeros(1, 10, dtype=torch.long),
+            "token_decode_offset": 0,
+            "tokenizer": None,  # won't be used in this test path
+            "acoustic_std": 1.5,
+            "acoustic_mean": 0.0,
+        }
+        return stream
+
+    def test_iteration_yields_tuples(self):
+        stream = self._make_audio_stream(num_steps=15)
+        chunks = list(stream)
+        # Should yield at least some chunks
+        assert len(chunks) > 0
+        for chunk, sr in chunks:
+            assert isinstance(chunk, torch.Tensor)
+            assert sr == 24000
+            assert chunk.ndim >= 1
+            assert chunk.shape[-1] > 0
+
+    def test_total_audio_length_positive(self):
+        stream = self._make_audio_stream(num_steps=15)
+        chunks = list(stream)
+        total_samples = sum(c.shape[-1] for c, _ in chunks)
+        assert total_samples > 0
+
+    def test_empty_generation(self):
+        """Generator that yields only the final output (no predicted tokens)."""
+        def empty_gen():
+            yield SyncTokGenerationOutput(
+                text_token_ids=torch.zeros(1, 0, dtype=torch.long),
+                acoustic_features=torch.zeros(1, 0, 512),
+                time_before=torch.zeros(1, 0, dtype=torch.long),
+                llm_time=torch.tensor(0.0),
+                diffusion_time=torch.tensor(0.0),
+                logits=None,
+                step_logs=[],
+            )
+
+        from tada.modules.decoder import Decoder, DecoderConfig
+
+        config = DecoderConfig(
+            embed_dim=512, hidden_dim=64, num_attn_layers=1, num_attn_heads=2,
+            attn_dim_feedforward=128, strides=[2, 2, 2, 2], block_attention="v2",
+        )
+        decoder = Decoder(config)
+        stream = AudioStream(
+            gen=empty_gen(), decoder=decoder,
+            acoustic_mean=0.0, acoustic_std=1.5,
+        )
+        stream._generate_context = {
+            "text": [""], "input_ids": torch.zeros(1, 1, dtype=torch.long),
+            "token_decode_offset": 0, "tokenizer": None,
+            "acoustic_std": 1.5, "acoustic_mean": 0.0,
+        }
+        chunks = list(stream)
+        # Empty generation should produce no chunks (or only a flush chunk)
+        assert len(chunks) <= 1
+
+
+class TestStreamingDecoder:
+    """Test StreamingDecoder with a tiny decoder model."""
+
+    @pytest.fixture
+    def tiny_decoder(self):
+        from tada.modules.decoder import Decoder, DecoderConfig
+
+        config = DecoderConfig(
+            embed_dim=512,
+            hidden_dim=64,
+            num_attn_layers=1,
+            num_attn_heads=2,
+            attn_dim_feedforward=128,
+            strides=[2, 2, 2, 2],
+            block_attention="v2",
+        )
+        decoder = Decoder(config)
+        decoder.eval()
+        return decoder
+
+    def test_basic_streaming(self, tiny_decoder):
+        sd = StreamingDecoder(tiny_decoder, cnn_window_size=50, min_first_emission=10)
+        chunks = []
+        for i in range(20):
+            token = torch.randn(512)
+            chunk = sd.decode_block(token, time_before=5)
+            if chunk is not None and chunk.numel() > 0:
+                chunks.append(chunk)
+        final = sd.flush(trailing_frames=3)
+        if final is not None and final.numel() > 0:
+            chunks.append(final)
+        assert len(chunks) > 0
+        total_samples = sum(c.shape[-1] for c in chunks)
+        assert total_samples > 0
+
+    def test_skip_leading_frames(self, tiny_decoder):
+        sd = StreamingDecoder(tiny_decoder, cnn_window_size=50, min_first_emission=10)
+        sd.skip_leading_frames(10)
+        assert sd._emitted_frames == 10
+
+    def test_reset(self, tiny_decoder):
+        sd = StreamingDecoder(tiny_decoder, cnn_window_size=50, min_first_emission=10)
+        # Feed some data
+        for _ in range(5):
+            sd.decode_block(torch.randn(512), time_before=3)
+        sd.reset()
+        assert sd._cached_frames == 0
+        assert sd._emitted_frames == 0
+        assert sd._all_hidden is None
+
+    def test_min_block_frames_buffering(self, tiny_decoder):
+        sd = StreamingDecoder(tiny_decoder, min_block_frames=5, cnn_window_size=50, min_first_emission=10)
+        # Single token with time_before=1 gives only 1 frame, should buffer
+        result = sd.decode_block(torch.randn(512), time_before=1)
+        assert result is None  # buffered, not enough frames
+
+    def test_flush_with_no_data(self, tiny_decoder):
+        sd = StreamingDecoder(tiny_decoder, cnn_window_size=50)
+        result = sd.flush()
+        assert result is None
+
+
+class TestSegmentAttentionMask:
+    def test_v2_basic(self):
+        # 2 blocks: [0,0,1,0,1] -> block0=[0,0,1], block1=[0,1]
+        mask_input = torch.tensor([[0, 0, 1, 0, 1]])
+        mask = _create_segment_attention_mask(mask_input, version="v2")
+        assert mask.shape == (1, 5, 5)
+        # Position 0 (block 0) should attend to block 0 (same) but not block 1
+        # Position 3 (block 1) should attend to block 0 (prev) and block 1 (same)
+        assert mask[0, 0, 3].item() == True   # block 0 can't attend to block 1
+        assert mask[0, 3, 0].item() == False  # block 1 can attend to block 0 (prev)
+
+
+# ---------------------------------------------------------------------------
+# Integration tests — require GPU + model weights
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.integration
+class TestStreamingIntegration:
+    """Full integration tests with real TADA model.
+
+    Run with: pytest tests/test_streaming.py -m integration -s
+    """
+
+    @pytest.fixture(scope="class")
+    def model_and_prompt(self):
+        import torchaudio
+        from tada.modules.encoder import Encoder
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+        encoder = Encoder.from_pretrained("HumeAI/tada-codec", subfolder="encoder").to(device)
+        from tada.modules.tada import TadaForCausalLM
+        model = TadaForCausalLM.from_pretrained("HumeAI/tada-1b", torch_dtype=dtype).to(device)
+
+        audio, sr = torchaudio.load("tada/samples/ljspeech.wav")
+        prompt = encoder(audio.to(device), sample_rate=sr)
+
+        return model, prompt, device
+
+    def test_non_streaming_unchanged(self, model_and_prompt):
+        """Non-streaming path should work exactly as before."""
+        model, prompt, device = model_and_prompt
+        output = model.generate(
+            prompt=prompt,
+            text="Hello world, this is a test.",
+        )
+        assert output.audio is not None
+        assert len(output.audio) == 1
+        assert output.audio[0] is not None
+        assert output.audio[0].shape[-1] > 0
+
+    def test_streaming_produces_chunks(self, model_and_prompt):
+        """Streaming should produce audio chunks."""
+        model, prompt, device = model_and_prompt
+        stream = model.generate(
+            prompt=prompt,
+            text="Hello world, this is a streaming test.",
+            stream=True,
+        )
+        chunks = []
+        for chunk, sr in stream:
+            assert sr == 24000
+            assert chunk.shape[-1] > 0
+            chunks.append(chunk)
+
+        assert len(chunks) > 0
+        total_samples = sum(c.shape[-1] for c in chunks)
+        assert total_samples > 0
+        # Check .result is populated
+        assert stream.result is not None
+        assert stream.result.acoustic_features is not None
+
+    def test_streaming_vs_nonstreaming_similar_length(self, model_and_prompt):
+        """Streaming and non-streaming should produce similar-length audio."""
+        model, prompt, device = model_and_prompt
+        text = "Please call Stella."
+
+        # Non-streaming
+        output_ns = model.generate(prompt=prompt, text=text)
+        ns_len = output_ns.audio[0].shape[-1]
+
+        # Streaming
+        stream = model.generate(prompt=prompt, text=text, stream=True)
+        chunks = [c for c, _ in stream]
+        s_len = sum(c.shape[-1] for c in chunks)
+
+        # Allow 20% tolerance (different random seeds in flow matching)
+        ratio = s_len / ns_len if ns_len > 0 else 1.0
+        assert 0.5 < ratio < 2.0, f"Length ratio {ratio} is too different"
+
+    def test_streaming_early_break(self, model_and_prompt):
+        """Breaking from iteration should not crash."""
+        model, prompt, device = model_and_prompt
+        stream = model.generate(
+            prompt=prompt,
+            text="This is a long text that should produce many chunks of audio for testing early termination.",
+            stream=True,
+        )
+        count = 0
+        for chunk, sr in stream:
+            count += 1
+            if count >= 2:
+                break
+        # Should not raise
+
+
+# ---------------------------------------------------------------------------
+# Generate 5 test audios for manual listening
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.integration
+class TestGenerateAudios:
+    """Generate 5 audio samples for manual listening validation.
+
+    Run with: pytest tests/test_streaming.py::TestGenerateAudios -m integration -s
+    Outputs saved to tests/output/
+    """
+
+    TEXTS = [
+        "Hello, this is a demonstration of streaming text to speech.",
+        "Please call Stella. Ask her to bring these things with her from the store.",
+        "The quick brown fox jumps over the lazy dog.",
+        "In the beginning, there was silence. Then came the voice, clear and bright.",
+        "Technology is best when it brings people together.",
+    ]
+
+    @pytest.fixture(scope="class")
+    def model_and_prompt(self):
+        import torchaudio
+        from tada.modules.encoder import Encoder
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+        encoder = Encoder.from_pretrained("HumeAI/tada-codec", subfolder="encoder").to(device)
+        from tada.modules.tada import TadaForCausalLM
+        model = TadaForCausalLM.from_pretrained("HumeAI/tada-1b", torch_dtype=dtype).to(device)
+
+        audio, sr = torchaudio.load("tada/samples/ljspeech.wav")
+        prompt = encoder(audio.to(device), sample_rate=sr)
+
+        return model, prompt
+
+    def test_generate_streaming_audios(self, model_and_prompt):
+        import os
+        import torchaudio
+
+        model, prompt = model_and_prompt
+        os.makedirs("tests/output", exist_ok=True)
+
+        for i, text in enumerate(self.TEXTS):
+            print(f"\n--- Generating streaming audio {i+1}: {text[:50]}...")
+            stream = model.generate(prompt=prompt, text=text, stream=True)
+            chunks = []
+            for chunk, sr in stream:
+                chunks.append(chunk)
+                print(f"  chunk: {chunk.shape[-1] / sr:.2f}s")
+
+            if chunks:
+                full_audio = torch.cat(chunks, dim=-1)
+                if full_audio.ndim == 1:
+                    full_audio = full_audio.unsqueeze(0)
+                out_path = f"tests/output/streaming_{i+1}.wav"
+                torchaudio.save(out_path, full_audio.cpu().float(), 24000)
+                print(f"  saved: {out_path} ({full_audio.shape[-1] / 24000:.2f}s)")
+
+        # Also generate non-streaming for comparison
+        for i, text in enumerate(self.TEXTS):
+            print(f"\n--- Generating non-streaming audio {i+1}: {text[:50]}...")
+            output = model.generate(prompt=prompt, text=text)
+            if output.audio[0] is not None:
+                wav = output.audio[0]
+                if wav.ndim == 1:
+                    wav = wav.unsqueeze(0)
+                out_path = f"tests/output/nonstreaming_{i+1}.wav"
+                torchaudio.save(out_path, wav.cpu().float(), 24000)
+                print(f"  saved: {out_path} ({wav.shape[-1] / 24000:.2f}s)")

From dd705cf130592d7a712da1191622806d34a45951 Mon Sep 17 00:00:00 2001
From: sharath <sharath@hume.ai>
Date: Mon, 11 May 2026 17:59:48 +0000
Subject: [PATCH 4/4] Register pytest integration marker

Clears PytestUnknownMarkWarning when running streaming integration tests.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 pyproject.toml | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index ade29f7..dc2f154 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -40,6 +40,9 @@ extend-select = ["I"]
 
 [tool.pytest.ini_options]
 log_cli_level = "DEBUG"
+markers = [
+    "integration: integration tests requiring GPU and model weights (deselected by default)",
+]
 
 [tool.uv.sources]
 # descript-audiotools = { index = "humepy" }