feat(generator): add streaming generation API with chunked audio decode

README.md

@@ -172,6 +172,34 @@ audio = generator.generate(
 )
 ```
 
+### Streaming generation
+
+Use `generate_stream()` to receive PCM chunks while frames are still being
+generated. The default `chunk_frames=25` yields about two seconds of 24 kHz
+audio per chunk.
+Chunks are watermarked independently; if SilentCipher rejects a very short
+terminal fragment, that fragment is yielded unchanged.
+
+```python
+import torch
+
+from generator import load_miso_8b
+
+generator = load_miso_8b(device="cuda")
+
+chunks = []
+for chunk in generator.generate_stream(
+    text="This sentence is decoded in chunks.",
+    speaker=0,
+    context=[],
+    max_audio_length_ms=10_000,
+    chunk_frames=25,
+):
+    chunks.append(chunk.cpu())
+
+audio = torch.cat(chunks, dim=0)
+```
+
 ---
 
 ## Weights

examples/stream_demo.py

@@ -0,0 +1,97 @@
+import argparse
+import os
+import time
+import wave
+
+os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "60")
+os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "60")
+
+import torch
+
+from generator import DEFAULT_MISO_TTS_REPO_ID, load_miso_8b
+
+# Disable Triton compilation
+os.environ["NO_TORCH_COMPILE"] = "1"
+
+
+def _pcm16_bytes(audio: torch.Tensor) -> bytes:
+    audio = audio.detach().flatten().to(dtype=torch.float32).cpu()
+    pcm = audio.clamp(-1.0, 1.0).mul(32767.0).to(torch.int16)
+    return pcm.numpy().tobytes()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--text", type=str, default="Hello from streamed Miso TTS.")
+    parser.add_argument("--speaker", type=int, default=0)
+    parser.add_argument("--output", type=str, default="streamed_generation.wav")
+    parser.add_argument("--max-audio-length-ms", type=float, default=10_000)
+    parser.add_argument("--chunk-frames", type=int, default=25)
+    parser.add_argument(
+        "--model-path-or-repo-id",
+        type=str,
+        default=os.environ.get("MISO_TTS_8B_MODEL", DEFAULT_MISO_TTS_REPO_ID),
+    )
+    args = parser.parse_args()
+
+    # Select the best available device, skipping MPS due to float64 limitations.
+    if torch.cuda.is_available():
+        device = "cuda"
+    else:
+        device = "cpu"
+    print(f"Using device: {device}")
+
+    if os.path.exists(args.model_path_or_repo_id):
+        print(f"Loading Miso TTS model from local path: {args.model_path_or_repo_id}")
+    else:
+        print(
+            "Loading Miso TTS model from Hugging Face: "
+            f"https://huggingface.co/{args.model_path_or_repo_id}"
+        )
+        print("The model will be downloaded and cached automatically if it is not already present.")
+
+    generator = load_miso_8b(device, model_path_or_repo_id=args.model_path_or_repo_id)
+
+    start_time = time.perf_counter()
+    first_audio_time = None
+    chunk_count = 0
+    sample_count = 0
+
+    with wave.open(args.output, "wb") as output_file:
+        output_file.setnchannels(1)
+        output_file.setsampwidth(2)
+        output_file.setframerate(generator.sample_rate)
+
+        for chunk in generator.generate_stream(
+            text=args.text,
+            speaker=args.speaker,
+            context=[],
+            max_audio_length_ms=args.max_audio_length_ms,
+            chunk_frames=args.chunk_frames,
+        ):
+            if chunk.numel() == 0:
+                continue
+
+            if first_audio_time is None:
+                first_audio_time = time.perf_counter()
+
+            chunk_count += 1
+            sample_count += chunk.numel()
+            output_file.writeframes(_pcm16_bytes(chunk))
+            duration_s = chunk.numel() / generator.sample_rate
+            print(f"Wrote chunk {chunk_count}: {duration_s:.2f}s")
+
+    total_time = time.perf_counter() - start_time
+    audio_duration_s = sample_count / generator.sample_rate
+    if first_audio_time is None:
+        print("No audio generated.")
+    else:
+        time_to_first_audio = first_audio_time - start_time
+        print(f"Time to first audio: {time_to_first_audio:.2f}s")
+        print(f"Total generation time: {total_time:.2f}s")
+        print(f"Audio duration: {audio_duration_s:.2f}s")
+        print(f"Successfully generated {args.output}")
+
+
+if __name__ == "__main__":
+    main()

generator.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 import os
-from typing import List, Optional, Tuple
+from typing import Iterator, List, Optional, Tuple
 
 os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "60")
 os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "60")
@@ -27,6 +27,19 @@ class Segment:
     audio: torch.Tensor
 
 
+def _stack_audio_frames(samples: List[torch.Tensor]) -> torch.Tensor:
+    return torch.stack(samples).permute(1, 2, 0)
+
+
+def _match_num_samples(audio: torch.Tensor, num_samples: int) -> torch.Tensor:
+    if audio.size(0) > num_samples:
+        return audio[:num_samples]
+    if audio.size(0) < num_samples:
+        padding = torch.zeros(num_samples - audio.size(0), dtype=audio.dtype, device=audio.device)
+        return torch.cat([audio, padding], dim=0)
+    return audio
+
+
 def load_llama3_tokenizer():
     """
     https://github.com/huggingface/transformers/issues/22794#issuecomment-2092623992
@@ -114,16 +127,13 @@ def _tokenize_segment(self, segment: Segment) -> Tuple[torch.Tensor, torch.Tenso
 
         return torch.cat([text_tokens, audio_tokens], dim=0), torch.cat([text_masks, audio_masks], dim=0)
 
-    @torch.inference_mode()
-    def generate(
+    def _prepare_prompt(
         self,
         text: str,
         speaker: int,
         context: List[Segment],
-        max_audio_length_ms: float = 90_000,
-        temperature: float = 0.9,
-        topk: int = 50,
-    ) -> torch.Tensor:
+        max_audio_length_ms: float,
+    ) -> Tuple[torch.Tensor, torch.Tensor, int]:
         self._model.reset_caches()
 
         max_generation_len = int(max_audio_length_ms / 80)
@@ -140,40 +150,159 @@ def generate(
         prompt_tokens = torch.cat(tokens, dim=0).long().to(self.device)
         prompt_tokens_mask = torch.cat(tokens_mask, dim=0).bool().to(self.device)
 
-        samples = []
-        curr_tokens = prompt_tokens.unsqueeze(0)
-        curr_tokens_mask = prompt_tokens_mask.unsqueeze(0)
-        curr_pos = torch.arange(0, prompt_tokens.size(0)).unsqueeze(0).long().to(self.device)
-
         max_seq_len = 2048
         max_context_len = max_seq_len - max_generation_len
-        if curr_tokens.size(1) >= max_context_len:
+        if prompt_tokens.size(0) >= max_context_len:
             raise ValueError(
                 f"Inputs too long, must be below max_seq_len - max_generation_len: {max_context_len}"
             )
 
+        return prompt_tokens, prompt_tokens_mask, max_generation_len
+
+    def _generate_frames(
+        self,
+        prompt_tokens: torch.Tensor,
+        prompt_tokens_mask: torch.Tensor,
+        max_generation_len: int,
+        temperature: float,
+        topk: int,
+    ) -> Iterator[torch.Tensor]:
+        curr_tokens = prompt_tokens.unsqueeze(0)
+        curr_tokens_mask = prompt_tokens_mask.unsqueeze(0)
+        curr_pos = torch.arange(0, prompt_tokens.size(0)).unsqueeze(0).long().to(self.device)
+
         for _ in range(max_generation_len):
             sample = self._model.generate_frame(curr_tokens, curr_tokens_mask, curr_pos, temperature, topk)
             if torch.all(sample == 0):
                 break  # eos
 
-            samples.append(sample)
-
             curr_tokens = torch.cat([sample, torch.zeros(1, 1).long().to(self.device)], dim=1).unsqueeze(1)
             curr_tokens_mask = torch.cat(
                 [torch.ones_like(sample).bool(), torch.zeros(1, 1).bool().to(self.device)], dim=1
             ).unsqueeze(1)
             curr_pos = curr_pos[:, -1:] + 1
 
-        audio = self._audio_tokenizer.decode(torch.stack(samples).permute(1, 2, 0)).squeeze(0).squeeze(0)
+            yield sample
 
+    def _decode_frames(self, samples: List[torch.Tensor]) -> torch.Tensor:
+        return self._audio_tokenizer.decode(_stack_audio_frames(samples)).squeeze(0).squeeze(0)
+
+    def _watermark_audio(self, audio: torch.Tensor) -> torch.Tensor:
         # This applies an imperceptible watermark to identify audio as AI-generated.
         # If using Miso TTS in another application, use your own private key and keep it secret.
         audio, wm_sample_rate = watermark(self._watermarker, audio, self.sample_rate, MISO_TTS_WATERMARK)
         audio = torchaudio.functional.resample(audio, orig_freq=wm_sample_rate, new_freq=self.sample_rate)
+        return audio
+
+    def _watermark_stream_chunk(self, audio: torch.Tensor, *, is_final: bool) -> torch.Tensor:
+        target_num_samples = audio.size(0)
+        try:
+            audio = self._watermark_audio(audio)
+        except Exception:
+            if not is_final:
+                raise
+            # SilentCipher may reject very short final chunks. Earlier chunks are
+            # watermarked independently; only the terminal fragment falls back.
+            return audio
+        return _match_num_samples(audio, target_num_samples)
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        text: str,
+        speaker: int,
+        context: List[Segment],
+        max_audio_length_ms: float = 90_000,
+        temperature: float = 0.9,
+        topk: int = 50,
+    ) -> torch.Tensor:
+        prompt_tokens, prompt_tokens_mask, max_generation_len = self._prepare_prompt(
+            text, speaker, context, max_audio_length_ms
+        )
+        samples = list(
+            self._generate_frames(prompt_tokens, prompt_tokens_mask, max_generation_len, temperature, topk)
+        )
+
+        audio = self._decode_frames(samples)
+        audio = self._watermark_audio(audio)
 
         return audio
 
+    def generate_stream(
+        self,
+        text: str,
+        speaker: int,
+        context: List[Segment],
+        max_audio_length_ms: float = 90_000,
+        temperature: float = 0.9,
+        topk: int = 50,
+        chunk_frames: int = 25,
+    ) -> Iterator[torch.Tensor]:
+        if chunk_frames <= 0:
+            raise ValueError("chunk_frames must be greater than 0")
+
+        # Tokenize the prompt (including Mimi encode of any context audio) before
+        # entering Mimi streaming mode, so streamed generation conditions on the
+        # same context tokens as generate().
+        with torch.inference_mode():
+            prompt_tokens, prompt_tokens_mask, max_generation_len = self._prepare_prompt(
+                text, speaker, context, max_audio_length_ms
+            )
+
+        frames = self._generate_frames(prompt_tokens, prompt_tokens_mask, max_generation_len, temperature, topk)
+
+        chunk: List[torch.Tensor] = []
+        all_samples: List[torch.Tensor] = []
+        streamed_num_samples = 0
+
+        # Each chunk is computed fully inside torch.inference_mode() and yielded
+        # outside it. A plain `with torch.inference_mode():` around the whole
+        # generator body would stay active while the generator is suspended at
+        # `yield`, silently putting the caller's loop body into inference mode.
+        #
+        # Only full chunk_frames-sized chunks are decoded inside Mimi streaming
+        # mode: on CUDA the streaming decoder is wrapped in a CUDA graph captured
+        # at the first chunk's shape, so a shorter residual chunk would raise a
+        # shape mismatch. The residual is covered by the batch tail decode below.
+        with self._audio_tokenizer.streaming(1):
+            finished = False
+            while not finished:
+                out: Optional[torch.Tensor] = None
+                with torch.inference_mode():
+                    while len(chunk) < chunk_frames:
+                        sample = next(frames, None)
+                        if sample is None:
+                            finished = True
+                            break
+                        chunk.append(sample)
+                        all_samples.append(sample)
+
+                    if len(chunk) == chunk_frames:
+                        audio = self._decode_frames(chunk)
+                        streamed_num_samples += audio.size(0)
+                        chunk = []
+                        if audio.numel() > 0:
+                            out = self._watermark_stream_chunk(audio, is_final=False)
+                if out is not None:
+                    yield out
+
+        tail_audio: Optional[torch.Tensor] = None
+        with torch.inference_mode():
+            if all_samples:
+                # Mimi streaming decode does not expose an explicit flush. Decode
+                # the full code sequence once (in batch mode, outside the streaming
+                # context) and emit only the deferred tail - the residual frames
+                # plus any samples the streamed chunks have not covered - so
+                # concatenated stream chunks keep the batch decode length.
+                full_audio = self._decode_frames(all_samples)
+                if streamed_num_samples < full_audio.size(0):
+                    tail = full_audio[streamed_num_samples:]
+                    if tail.numel() > 0:
+                        tail_audio = self._watermark_stream_chunk(tail, is_final=True)
+
+        if tail_audio is not None:
+            yield tail_audio
+
 
 def _state_dict_from_checkpoint(checkpoint: object) -> dict[str, torch.Tensor]:
     if not isinstance(checkpoint, dict):

tests/test_streaming_generation.py

@@ -0,0 +1,31 @@
+import unittest
+
+import torch
+
+from generator import _match_num_samples, _stack_audio_frames
+
+
+class StreamingGenerationHelpersTest(unittest.TestCase):
+    def test_stack_audio_frames_preserves_time_order(self) -> None:
+        frames = [
+            torch.tensor([[1, 2, 3]]),
+            torch.tensor([[4, 5, 6]]),
+        ]
+
+        codes = _stack_audio_frames(frames)
+
+        expected = torch.tensor([[[1, 4], [2, 5], [3, 6]]])
+        self.assertTrue(torch.equal(codes, expected))
+
+    def test_match_num_samples_trims_or_pads(self) -> None:
+        audio = torch.tensor([1.0, 2.0, 3.0])
+
+        trimmed = _match_num_samples(audio, 2)
+        padded = _match_num_samples(audio, 5)
+
+        self.assertTrue(torch.equal(trimmed, torch.tensor([1.0, 2.0])))
+        self.assertTrue(torch.equal(padded, torch.tensor([1.0, 2.0, 3.0, 0.0, 0.0])))
+
+
+if __name__ == "__main__":
+    unittest.main()