V02 release (#94)

* bump version to v0.2

* adapt readme

* Update README.md

* update README

* add inference tips + streamer class

* update readme

* Update README.md

* Apply suggestions from code review
Co-authored-by: Sanchit Gandhi <93869735+sanchit-gandhi@users.noreply.github.com>

* Update README

* Apply suggestions from code review
Co-authored-by: Vaibhav Srivastav <vaibhavs10@gmail.com>

---------
Co-authored-by: Sanchit Gandhi <93869735+sanchit-gandhi@users.noreply.github.com>
Co-authored-by: Vaibhav Srivastav <vaibhavs10@gmail.com>

INFERENCE.md

+# Inference tips
+
+Parler-TTS benefits from a number of optimizations that can make the model up to 4x faster. Add to this the ability to stream audio as it's being generated, and you can achieve time-to-first audio in under 500ms on a modern GPU.
+
+## 📖 Quick Index
+* [Efficient Attention Implementation](#efficient-attention-implementations)
+* [Compilation](#compilation)
+* [Streaming](#streaming)
+* [Batch generation](#batch-generation)
+
+## Efficient Attention implementations
+
+Parler-TTS supports [SDPA](https://pytorch.org/docs/master/generated/torch.nn.functional.scaled_dot_product_attention.html) and [Flash Attention 2](https://github.com/Dao-AILab/flash-attention).  
+
+SDPA is used by default and speeds up generation time by up to 1.4x compared with eager attention.
+
+To switch between attention implementations, simply specify `attn_implementation=attn_implementation` when loading the checkpoints:
+
+```py
+from parler_tts import ParlerTTSForConditionalGeneration
+
+torch_device = "cuda:0" # use "mps" for Mac
+torch_dtype = torch.bfloat16
+model_name = "parler-tts/parler-tts-mini-v1"
+
+attn_implementation = "eager" # "sdpa" or "flash_attention_2"
+
+model = ParlerTTSForConditionalGeneration.from_pretrained(
+    model_name,
+    attn_implementation=attn_implementation
+).to(torch_device, dtype=torch_dtype)
+```
+
+## Compilation
+
+[Compiling](https://pytorch.org/docs/stable/generated/torch.compile.html) the forward method of Parler can speed up generation time by up to 4.5x.
+
+As an indication, `mode=default` brings a speed-up of 1.4 times compared to no compilation, while `mode="reduce-overhead"` brings much faster generation, at the cost of a longer compilation time and the need to generate twice to see the benefits of compilation.
+
+```py
+import torch
+from parler_tts import ParlerTTSForConditionalGeneration
+from transformers import AutoTokenizer
+
+torch_device = "cuda:0"
+torch_dtype = torch.bfloat16
+model_name = "parler-tts/parler-tts-mini-v1"
+
+# need to set padding max length
+max_length = 50
+
+# load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_name) 
+model = ParlerTTSForConditionalGeneration.from_pretrained(
+    model_name,
+    attn_implementation="eager"
+).to(torch_device, dtype=torch_dtype)
+
+# compile the forward pass
+compile_mode = "default" # chose "reduce-overhead" for 3 to 4x speed-up
+model.generation_config.cache_implementation = "static"
+model.forward = torch.compile(model.forward, mode=compile_mode)
+
+# warmup
+inputs = tokenizer("This is for compilation", return_tensors="pt", padding="max_length", max_length=max_length).to(device)
+
+model_kwargs = {**inputs, "prompt_input_ids": inputs.input_ids, "prompt_attention_mask": inputs.attention_mask, }
+
+n_steps = 1 if compile_mode == "default" else 2
+for _ in range(n_steps):
+    _ = model.generate(**model_kwargs)
+
+
+# now you can benefit from compilation speed-ups
+...
+
+```
+
+
+## Streaming
+
+### How Does It Work?
+
+Parler-TTS is an auto-regressive transformer-based model, meaning generates audio codes (tokens) in a causal fashion.
+
+At each decoding step, the model generates a new set of audio codes, conditional on the text input and all previous audio codes. From the 
+frame rate of the [DAC model](https://huggingface.co/parler-tts/dac_44khZ_8kbps) used to decode the generated codes to audio waveform,  each set of generated audio codes corresponds to 0.011 seconds. This means we require a total of 1720 decoding steps to generate 20 seconds of audio.
+
+Rather than waiting for the entire audio sequence to be generated, which would require the full 1720 decoding steps, we can start playing the audio after a specified number of decoding steps have been reached, a techinque known as [*streaming*](https://huggingface.co/docs/transformers/main/en/generation_strategies#streaming). 
+For example, after 86 steps we have the first second of audio ready, and so can play this without waiting for the remaining decoding steps to be complete. As we continue to generate with the Parler-TTS model, we append new chunks of generated audio to our output waveform on-the-fly. After the full 1720 decoding steps, the generated audio is complete, and is composed of 20 chunks of audio, each corresponding to 86 tokens.
+This method of playing incremental generations reduces the latency of the Parler-TTS model from the total time to generate 1720 tokens, to the time taken to play the first chunk of audio (86 tokens). This can result in significant improvements to perceived latency,  particularly when the chunk size is chosen to be small. In practice, the chunk size should be tuned to your device: using a smaller chunk size will mean that the first chunk is ready faster, but should not be chosen so small that the model generates slower than the time it takes to play the audio.
+
+
+### How Can I Use It?
+
+We've added [ParlerTTSStreamer](https://github.com/huggingface/parler-tts/blob/main/parler_tts/streamer.py) to the library. Don't hesitate to adapt it to your use-case.
+
+Here's how to create a generator out of the streamer.
+
+```py
+import torch
+from parler_tts import ParlerTTSForConditionalGeneration, ParlerTTSStreamer
+from transformers import AutoTokenizer
+from threading import Thread
+
+torch_device = "cuda:0" # Use "mps" for Mac 
+torch_dtype = torch.bfloat16
+model_name = "parler-tts/parler-tts-mini-v1"
+
+# need to set padding max length
+max_length = 50
+
+# load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_name) 
+model = ParlerTTSForConditionalGeneration.from_pretrained(
+    model_name,
+).to(torch_device, dtype=torch_dtype)
+
+sampling_rate = model.audio_encoder.config.sampling_rate
+frame_rate = model.audio_encoder.config.frame_rate
+
+def generate(text, description, play_steps_in_s=0.5):
+  play_steps = int(frame_rate * play_steps_in_s)
+  streamer = ParlerTTSStreamer(model, device=torch_device, play_steps=play_steps)
+  # tokenization
+  inputs = tokenizer(description, return_tensors="pt").to(torch_device)
+  prompt = tokenizer(text, return_tensors="pt").to(torch_device)
+  # create generation kwargs
+  generation_kwargs = dict(
+    input_ids=inputs.input_ids,
+    prompt_input_ids=prompt.input_ids,
+    attention_mask=inputs.attention_mask,
+    prompt_attention_mask=prompt.attention_mask,
+    streamer=streamer,
+    do_sample=True,
+    temperature=1.0,
+    min_new_tokens=10,
+  )
+  # initialize Thread
+  thread = Thread(target=model.generate, kwargs=generation_kwargs)
+  thread.start()
+  # iterate over chunks of audio
+  for new_audio in streamer:
+    if new_audio.shape[0] == 0:
+      break
+    print(f"Sample of length: {round(new_audio.shape[0] / sampling_rate, 4)} seconds")
+    yield sampling_rate, new_audio
+
+
+# now you can do
+text = "This is a test of the streamer class"
+description = "Jon's talking really fast."
+
+chunk_size_in_s = 0.5
+
+for (sampling_rate, audio_chunk) in generate(text, description, chunk_size_in_s):
+  # You can do everything that you need with the chunk now
+  # For example: stream it, save it, play it.
+  print(audio_chunk.shape) 
+```
+
+## Batch generation
+
+Batching means combining operations for multiple samples to bring the overall time spent generating the samples lower than generating sample per sample.
+
+Here is a quick example of how you can use it:
+
+```py
+from parler_tts import ParlerTTSForConditionalGeneration
+from transformers import AutoTokenizer, AutoFeatureExtractor, set_seed
+import scipy
+
+
+repo_id = "parler-tts/parler-tts-mini-v1"
+
+model = ParlerTTSForConditionalGeneration.from_pretrained(repo_id).to("cuda")
+tokenizer = AutoTokenizer.from_pretrained(repo_id, padding_side="left")
+feature_extractor = AutoFeatureExtractor.from_pretrained(repo_id)
+
+input_text = ["Hey, how are you doing?", "I'm not sure how to feel about it."]
+description = 2 * ["A male speaker with a monotone and high-pitched voice is delivering his speech at a really low speed in a confined environment."]
+
+inputs = tokenizer(description, return_tensors="pt", padding=True).to("cuda")
+prompt = tokenizer(input_text, return_tensors="pt", padding=True).to("cuda")
+
+set_seed(0)
+generation = model.generate(
+    input_ids=inputs.input_ids,
+    attention_mask=inputs.attention_mask,
+    prompt_input_ids=prompt.input_ids,
+    prompt_attention_mask=prompt.attention_mask,
+    do_sample=True,
+    return_dict_in_generate=True,
+)
+
+audio_1 = generation.sequences[0, :generation.audios_length[0]]
+audio_2 = generation.sequences[1, :generation.audios_length[1]]
+
+print(audio_1.shape, audio_2.shape)
+scipy.io.wavfile.write("sample_out.wav", rate=feature_extractor.sampling_rate, data=audio_1.cpu().numpy().squeeze())
+scipy.io.wavfile.write("sample_out_2.wav", rate=feature_extractor.sampling_rate, data=audio_2.cpu().numpy().squeeze())
+```
\ No newline at end of file

README.md

@@ -7,15 +7,23 @@ Contrarily to other TTS models, Parler-TTS is a **fully open-source** release. A
 This repository contains the inference and training code for Parler-TTS. It is designed to accompany the [Data-Speech](https://github.com/huggingface/dataspeech) repository for dataset annotation.
 
 > [!IMPORTANT]
-> We're proud to release [Parler-TTS Mini v0.1](https://huggingface.co/parler-tts/parler_tts_mini_v0.1), our first 600M parameter model, trained on 10.5K hours of audio data.
-> In the coming weeks, we'll be working on scaling up to 50k hours of data, in preparation for the v1 model.
+> **08/08/2024:** We are proud to release two new Parler-TTS checkpoints:
+> 1. [Parler-TTS Mini](https://huggingface.co/parler-tts/parler-tts-mini-v1), an 880M parameter model.
+> 2. [Parler-TTS Large](https://huggingface.co/parler-tts/parler-tts-large-v1), a 2.3B parameter model.
+>
+> These checkpoints have been trained on 45k hours of audiobook data.
+>
+> In addition, the code is optimized for much faster generation: we've added SDPA and Flash Attention 2 compatibility, as well as the ability to compile the model.
 
 ## 📖 Quick Index
 * [Installation](#installation)
 * [Usage](#usage)
+  - [🎲 Using a random voice](#-random-voice)
+  - [🎯 Using a specific speaker](#-using-a-specific-speaker)
 * [Training](#training)
-* [Demo](https://huggingface.co/spaces/parler-tts/parler_tts_mini)
+* [Demo](https://huggingface.co/spaces/parler-tts/parler_tts)
 * [Model weights and datasets](https://huggingface.co/parler-tts)
+* [Optimizing inference](#-optimizing-inference-speed)
 
 ## Installation
 
@@ -34,43 +42,85 @@ pip3 install --pre torch torchaudio --index-url https://download.pytorch.org/whl
 ## Usage
 
 > [!TIP]
-> You can directly try it out in an interactive demo [here](https://huggingface.co/spaces/parler-tts/parler_tts_mini)!
+> You can directly try it out in an interactive demo [here](https://huggingface.co/spaces/parler-tts/parler_tts)!
 
-Using Parler-TTS is as simple as "bonjour". Simply use the following inference snippet.
+Using Parler-TTS is as simple as "bonjour". Simply install the library once:
+
+```sh
+pip install git+https://github.com/huggingface/parler-tts.git
+```
+
+### 🎲 Random voice
+
+
+**Parler-TTS** has been trained to generate speech with features that can be controlled with a simple text prompt, for example:
 
 ```py
+import torch
 from parler_tts import ParlerTTSForConditionalGeneration
 from transformers import AutoTokenizer
 import soundfile as sf
-import torch
 
-device = "cpu"
-if torch.cuda.is_available():
-    device = "cuda:0"
-if torch.backends.mps.is_available():
-    device = "mps"
-if torch.xpu.is_available():
-    device = "xpu"
-torch_dtype = torch.float16 if device != "cpu" else torch.float32
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler_tts_mini_v0.1", torch_dtype=torch_dtype).to(device)
+model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-mini-v1").to(device)
+tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler-tts-mini-v1")
 
-tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler_tts_mini_v0.1")
+prompt = "Hey, how are you doing today?"
+description = "A female speaker delivers a slightly expressive and animated speech with a moderate speed and pitch. The recording is of very high quality, with the speaker's voice sounding clear and very close up."
+
+input_ids = tokenizer(description, return_tensors="pt").input_ids.to(device)
+prompt_input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+
+generation = model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids)
+audio_arr = generation.cpu().numpy().squeeze()
+sf.write("parler_tts_out.wav", audio_arr, model.config.sampling_rate)
+```
+
+### 🎯 Using a specific speaker
+
+To ensure speaker consistency across generations, this checkpoint was also trained on 34 speakers, characterized by name (e.g. Jon, Lea, Gary, Jenna, Mike, Laura).
+
+To take advantage of this, simply adapt your text description to specify which speaker to use: `Jon's voice is monotone yet slightly fast in delivery, with a very close recording that almost has no background noise.`
+
+```py
+import torch
+from parler_tts import ParlerTTSForConditionalGeneration
+from transformers import AutoTokenizer
+import soundfile as sf
+
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-mini-v1").to(device)
+tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler-tts-mini-v1")
 
 prompt = "Hey, how are you doing today?"
-description = "A female speaker with a slightly low-pitched voice delivers her words quite expressively, in a very confined sounding environment with clear audio quality. She speaks very fast."
+description = "Jon's voice is monotone yet slightly fast in delivery, with a very close recording that almost has no background noise."
 
 input_ids = tokenizer(description, return_tensors="pt").input_ids.to(device)
 prompt_input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
 
-generation = model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids).to(torch.float32)
+generation = model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids)
 audio_arr = generation.cpu().numpy().squeeze()
 sf.write("parler_tts_out.wav", audio_arr, model.config.sampling_rate)
 ```
 
+**Tips**:
+* Include the term "very clear audio" to generate the highest quality audio, and "very noisy audio" for high levels of background noise
+* Punctuation can be used to control the prosody of the generations, e.g. use commas to add small breaks in speech
+* The remaining speech features (gender, speaking rate, pitch and reverberation) can be controlled directly through the prompt
+
+### ✨ Optimizing Inference Speed
+
+We've set up an [inference guide](INFERENCE.md) to make generation faster. Think SDPA, torch.compile and streaming!
+
+
 https://github.com/huggingface/parler-tts/assets/52246514/251e2488-fe6e-42c1-81cd-814c5b7795b0
 
 ## Training
+> [!WARNING]
+> The training guide has yet to be adapted to the newest checkpoints.
+
 <a target="_blank" href="https://colab.research.google.com/github/ylacombe/scripts_and_notebooks/blob/main/Finetuning_Parler_TTS_on_a_single_speaker_dataset.ipynb"> 
   <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/> 
 </a>

parler_tts/__init__.py

-__version__ = "0.1"
+__version__ = "0.2"
 
 
 from transformers import AutoConfig, AutoModel
@@ -12,6 +12,7 @@ from .modeling_parler_tts import (
     build_delay_pattern_mask,
 )
 
+from .streamer import ParlerTTSStreamer
 
 AutoConfig.register("dac", DACConfig)
 AutoModel.register(DACConfig, DACModel)

parler_tts/streamer.py

+
+from .modeling_parler_tts import ParlerTTSForConditionalGeneration
+from transformers.generation.streamers import BaseStreamer
+from typing import Optional
+import torch
+import numpy as np
+import math
+from queue import Queue
+
+
+class ParlerTTSStreamer(BaseStreamer):
+    def __init__(
+        self,
+        model: ParlerTTSForConditionalGeneration,
+        device: Optional[str] = None,
+        play_steps: Optional[int] = 10,
+        stride: Optional[int] = None,
+        timeout: Optional[float] = None,
+    ):
+        """
+        Streamer that stores playback-ready audio in a queue, to be used by a downstream application as an iterator. This is
+        useful for applications that benefit from accessing the generated audio in a non-blocking way (e.g. in an interactive
+        Gradio demo).
+        Parameters:
+            model (`ParlerTTSForConditionalGeneration`):
+                The Parler-TTS model used to generate the audio waveform.
+            device (`str`, *optional*):
+                The torch device on which to run the computation. If `None`, will default to the device of the model.
+            play_steps (`int`, *optional*, defaults to 10):
+                The number of generation steps with which to return the generated audio array. Using fewer steps will
+                mean the first chunk is ready faster, but will require more codec decoding steps overall. This value
+                should be tuned to your device and latency requirements.
+            stride (`int`, *optional*):
+                The window (stride) between adjacent audio samples. Using a stride between adjacent audio samples reduces
+                the hard boundary between them, giving smoother playback. If `None`, will default to a value equivalent to
+                play_steps // 6 in the audio space.
+            timeout (`int`, *optional*):
+                The timeout for the audio queue. If `None`, the queue will block indefinitely. Useful to handle exceptions
+                in `.generate()`, when it is called in a separate thread.
+        """
+        self.decoder = model.decoder
+        self.audio_encoder = model.audio_encoder
+        self.generation_config = model.generation_config
+        self.device = device if device is not None else model.device
+
+        # variables used in the streaming process
+        self.play_steps = play_steps
+        if stride is not None:
+            self.stride = stride
+        else:
+            hop_length = math.floor(self.audio_encoder.config.sampling_rate / self.audio_encoder.config.frame_rate)
+            self.stride = hop_length * (play_steps - self.decoder.num_codebooks) // 6
+        self.token_cache = None
+        self.to_yield = 0
+
+        # varibles used in the thread process
+        self.audio_queue = Queue()
+        self.stop_signal = None
+        self.timeout = timeout
+
+    def apply_delay_pattern_mask(self, input_ids):
+        # build the delay pattern mask for offsetting each codebook prediction by 1 (this behaviour is specific to Parler)
+        _, delay_pattern_mask = self.decoder.build_delay_pattern_mask(
+            input_ids[:, :1],
+            bos_token_id=self.generation_config.bos_token_id,
+            pad_token_id=self.generation_config.decoder_start_token_id,
+            max_length=input_ids.shape[-1],
+        )
+        # apply the pattern mask to the input ids
+        input_ids = self.decoder.apply_delay_pattern_mask(input_ids, delay_pattern_mask)
+
+        # revert the pattern delay mask by filtering the pad token id
+        mask = (delay_pattern_mask != self.generation_config.bos_token_id) & (delay_pattern_mask != self.generation_config.pad_token_id)
+        input_ids = input_ids[mask].reshape(1, self.decoder.num_codebooks, -1)
+        # append the frame dimension back to the audio codes
+        input_ids = input_ids[None, ...]
+
+        # send the input_ids to the correct device
+        input_ids = input_ids.to(self.audio_encoder.device)
+
+        decode_sequentially = (
+            self.generation_config.bos_token_id in input_ids
+            or self.generation_config.pad_token_id in input_ids
+            or self.generation_config.eos_token_id in input_ids
+        )
+        if not decode_sequentially:
+            output_values = self.audio_encoder.decode(
+                input_ids,
+                audio_scales=[None],
+            )
+        else:
+            sample = input_ids[:, 0]
+            sample_mask = (sample >= self.audio_encoder.config.codebook_size).sum(dim=(0, 1)) == 0
+            sample = sample[:, :, sample_mask]
+            output_values = self.audio_encoder.decode(sample[None, ...], [None])
+
+        audio_values = output_values.audio_values[0, 0]
+        return audio_values.cpu().float().numpy()
+
+    def put(self, value):
+        batch_size = value.shape[0] // self.decoder.num_codebooks
+        if batch_size > 1:
+            raise ValueError("ParlerTTSStreamer only supports batch size 1")
+
+        if self.token_cache is None:
+            self.token_cache = value
+        else:
+            self.token_cache = torch.concatenate([self.token_cache, value[:, None]], dim=-1)
+
+        if self.token_cache.shape[-1] % self.play_steps == 0:
+            audio_values = self.apply_delay_pattern_mask(self.token_cache)
+            self.on_finalized_audio(audio_values[self.to_yield : -self.stride])
+            self.to_yield += len(audio_values) - self.to_yield - self.stride
+
+    def end(self):
+        """Flushes any remaining cache and appends the stop symbol."""
+        if self.token_cache is not None:
+            audio_values = self.apply_delay_pattern_mask(self.token_cache)
+        else:
+            audio_values = np.zeros(self.to_yield)
+
+        self.on_finalized_audio(audio_values[self.to_yield :], stream_end=True)
+
+    def on_finalized_audio(self, audio: np.ndarray, stream_end: bool = False):
+        """Put the new audio in the queue. If the stream is ending, also put a stop signal in the queue."""
+        self.audio_queue.put(audio, timeout=self.timeout)
+        if stream_end:
+            self.audio_queue.put(self.stop_signal, timeout=self.timeout)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        value = self.audio_queue.get(timeout=self.timeout)
+        if not isinstance(value, np.ndarray) and value == self.stop_signal:
+            raise StopIteration()
+        else:
+            return value
\ No newline at end of file

training/README.md

@@ -4,6 +4,9 @@
   <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/> 
 </a>
 
+> [!WARNING]
+> The training guide has yet to be adapted to the newest checkpoints.
+
 **TL;DR:** After having followed the [installation steps](#requirements), you can reproduce the [Parler-TTS Mini v0.1](https://huggingface.co/parler-tts/parler_tts_mini_v0.1) training recipe with the following command line:
 
 ```sh