Create app.py

app.py

@@ -0,0 +1,326 @@
+import math
+from queue import Queue
+from threading import Thread
+from typing import Optional
+
+import numpy as np
+import spaces
+import gradio as gr
+import torch
+
+from parler_tts import ParlerTTSForConditionalGeneration
+from transformers import AutoTokenizer, AutoFeatureExtractor, set_seed
+from transformers.generation.streamers import BaseStreamer
+
+device = "cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+torch_dtype = torch.float16 if device != "cpu" else torch.float32
+
+repo_id = "parler-tts/parler_tts_mini_v0.1"
+custom_repo_id = "AkhilTolani/vocals"
+
+model = ParlerTTSForConditionalGeneration.from_pretrained(
+    repo_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True
+).to(device)
+custom_model = ParlerTTSForConditionalGeneration.from_pretrained(
+    custom_repo_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True
+).to(device)
+
+tokenizer = AutoTokenizer.from_pretrained(repo_id)
+feature_extractor = AutoFeatureExtractor.from_pretrained(repo_id)
+
+SAMPLE_RATE = feature_extractor.sampling_rate
+SEED = 777
+
+default_text = "ooooh, please surprise me and sing whatever you would like to sing, ooh."
+
+examples = [
+    [
+        "In the quiet of the night, Shadows dance beneath the moonlight, Whispers of a love that's gone, Echoes of a heart withdrawn.",
+        "A women speaks at an average pace with a slightly animated delivery in a very confined sounding environment with clear audio quality.",
+        15.0,
+    ],
+    [
+        "'Lost in the memories, Of the days that used to be, A fragile heart in a storm, Yearning for a love reborn.",
+        "A man speaks in quite a monotone voice at a slightly faster-than-average pace in a confined space with very clear audio.",
+        15.0,
+    ],
+    [
+        "Tears fall like autumn rain, Tracing lines of hidden pain, In the silence, I still hear, The ghost of you always near.",
+        "A man delivers her words at a slightly slow pace in a small, confined space with a touch of background noise and a quite monotone tone.",
+        15.0,
+    ],
+    [
+        "Where did we go wrong? In the story we've outgrown, Now I'm left to sing alone, In a world so cold and stone.",
+        "A woman delivers her words at a fast pace and an animated tone, in a very spacious environment, in a very clear voice and audio quality.",
+        15.0,
+    ],
+]
+
+
+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
+
+
+sampling_rate = model.audio_encoder.config.sampling_rate
+frame_rate = model.audio_encoder.config.frame_rate
+
+@spaces.GPU
+def generate_base(text, description, play_steps_in_s=2.0):
+    play_steps = int(frame_rate * play_steps_in_s)
+    streamer = ParlerTTSStreamer(model, device=device, play_steps=play_steps)
+
+    inputs = tokenizer(description, return_tensors="pt").to(device)
+    prompt = tokenizer(text, return_tensors="pt").to(device)
+
+    generation_kwargs = dict(
+        input_ids=inputs.input_ids,
+        prompt_input_ids=prompt.input_ids,
+        streamer=streamer,
+        do_sample=True,
+        temperature=1.0,
+        min_new_tokens=10,
+    )
+
+    set_seed(SEED)
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    for new_audio in streamer:
+        print(f"Sample of length: {round(new_audio.shape[0] / sampling_rate, 2)} seconds")
+        yield sampling_rate, new_audio
+
+@spaces.GPU
+def generate_custom(text, description, play_steps_in_s=2.0):
+    play_steps = int(frame_rate * play_steps_in_s)
+    streamer = ParlerTTSStreamer(model, device=device, play_steps=play_steps)
+
+    inputs = tokenizer(description, return_tensors="pt").to(device)
+    prompt = tokenizer(text, return_tensors="pt").to(device)
+
+    generation_kwargs = dict(
+        input_ids=inputs.input_ids,
+        prompt_input_ids=prompt.input_ids,
+        streamer=streamer,
+        do_sample=True,
+        temperature=1.0,
+        min_new_tokens=10,
+    )
+
+    set_seed(SEED)
+    thread = Thread(target=custom_model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    for new_audio in streamer:
+        print(f"Sample of length: {round(new_audio.shape[0] / sampling_rate, 2)} seconds")
+        yield sampling_rate, new_audio
+
+
+css = """
+        #share-btn-container {
+            display: flex;
+            padding-left: 0.5rem !important;
+            padding-right: 0.5rem !important;
+            background-color: #000000;
+            justify-content: center;
+            align-items: center;
+            border-radius: 9999px !important; 
+            width: 13rem;
+            margin-top: 10px;
+            margin-left: auto;
+            flex: unset !important;
+        }
+        #share-btn {
+            all: initial;
+            color: #ffffff;
+            font-weight: 600;
+            cursor: pointer;
+            font-family: 'IBM Plex Sans', sans-serif;
+            margin-left: 0.5rem !important;
+            padding-top: 0.25rem !important;
+            padding-bottom: 0.25rem !important;
+            right:0;
+        }
+        #share-btn * {
+            all: unset !important;
+        }
+        #share-btn-container div:nth-child(-n+2){
+            width: auto !important;
+            min-height: 0px !important;
+        }
+        #share-btn-container .wrap {
+            display: none !important;
+        }
+"""
+with gr.Blocks(css=css) as block:
+    gr.HTML(
+        """
+            <div style="text-align: center; max-width: 700px; margin: 0 auto;">
+              <div
+                style="
+                  display: inline-flex; align-items: center; gap: 0.8rem; font-size: 1.75rem;
+                "
+              >
+                <h1 style="font-weight: 900; margin-bottom: 7px; line-height: normal;">
+                  Parler-TTS 🗣️
+                </h1>
+              </div>
+            </div>
+        """
+    )
+    gr.HTML(
+        f"""
+        <p><a href="https://github.com/huggingface/parler-tts"> Parler-TTS + Vocals</a> is a training and inference library for
+        high-fidelity text-to-speech (TTS) models. Generates high-quality vocals with features that can be controlled using a simple text prompt (e.g. gender, background noise, speaking rate, pitch and reverberation).</p>
+
+        <p>Tips for ensuring good generation:
+        <ul>
+            <li>Include the term <b>"very clear audio"</b> to generate the highest quality audio, and "very noisy audio" for high levels of background noise</li>
+            <li>Punctuation can be used to control the prosody of the generations, e.g. use commas to add small breaks in speech</li>
+            <li>The remaining speech features (gender, speaking rate, pitch and reverberation) can be controlled directly through the prompt</li>
+        </ul>
+        </p>
+        """
+    )
+    with gr.Tab("Base"):
+        with gr.Row():
+            with gr.Column():
+                input_text = gr.Textbox(label="Input Text", lines=2, value=default_text, elem_id="input_text")
+                description = gr.Textbox(label="Description", lines=2, value="", elem_id="input_description")
+                play_seconds = gr.Slider(3.0, 10.0, value=5.0, step=0.5, label="Streaming interval in seconds", info="Lower = shorter chunks, lower latency, more codec steps")
+                run_button = gr.Button("Generate Audio", variant="primary")
+            with gr.Column():
+                audio_out = gr.Audio(label="Parler-TTS + Vocals", type="numpy", elem_id="audio_out", streaming=True, autoplay=True)
+
+        inputs = [input_text, description, play_seconds]
+        outputs = [audio_out]
+        gr.Examples(examples=examples, fn=generate_base, inputs=inputs, outputs=outputs, cache_examples=False)
+        run_button.click(fn=generate_base, inputs=inputs, outputs=outputs, queue=True)
+
+block.queue()
+block.launch(share=True)