diff --git a/README.md b/README.md index 0423b4f..0eb218c 100644 --- a/README.md +++ b/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 diff --git a/examples/stream_demo.py b/examples/stream_demo.py new file mode 100644 index 0000000..97c3529 --- /dev/null +++ b/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() diff --git a/generator.py b/generator.py index 8441cee..f20a74d 100644 --- a/generator.py +++ b/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): diff --git a/tests/test_streaming_generation.py b/tests/test_streaming_generation.py new file mode 100644 index 0000000..4c0e37f --- /dev/null +++ b/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()