Update app.py (#3)

- Update app.py (fe7399d2bb60ee8104736952121646297025506e)
- Update requirements.txt (f65e4a52d4d5d2173891b195bee96c3b4462dbdb)


Co-authored-by: Yoach Lacombe <[email protected]>

app.py

@@ -9,22 +9,26 @@ import numpy as np
 import spaces
 import gradio as gr
 import torch
+import nltk
+
 
 from parler_tts import ParlerTTSForConditionalGeneration
 from pydub import AudioSegment
 from transformers import AutoTokenizer, AutoFeatureExtractor, set_seed
 
-device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+nltk.download('punkt_tab')
+
+device = "cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
 torch_dtype = torch.bfloat16 if device != "cpu" else torch.float32
 
 repo_id = "ai4bharat/indic-parler-tts-pretrained"
-jenny_repo_id = "ai4bharat/indic-parler-tts"
+finetuned_repo_id = "ai4bharat/indic-parler-tts"
 
 model = ParlerTTSForConditionalGeneration.from_pretrained(
     repo_id, attn_implementation="eager", torch_dtype=torch_dtype,
 ).to(device)
-jenny_model = ParlerTTSForConditionalGeneration.from_pretrained(
-    jenny_repo_id, attn_implementation="eager", torch_dtype=torch_dtype,
+finetuned_model = ParlerTTSForConditionalGeneration.from_pretrained(
+    finetuned_repo_id, attn_implementation="eager", torch_dtype=torch_dtype,
 ).to(device)
 
 tokenizer = AutoTokenizer.from_pretrained(repo_id)
@@ -89,7 +93,7 @@ examples = [
 ]
 
 
-jenny_examples = [
+finetuned_examples = [
     [
         "मुले बागेत खेळत आहेत आणि पक्षी किलबिलाट करत आहेत.",
         "Sunita speaks slowly in a calm, moderate-pitched voice, delivering the news with a neutral tone. The recording is very high quality with no background noise.",
@@ -171,44 +175,30 @@ def numpy_to_mp3(audio_array, sampling_rate):
 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 = description_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 numpy_to_mp3(new_audio, sampling_rate=sampling_rate)
-
 @spaces.GPU
-def generate_base(text, description, play_steps_in_s=2.0):
+def generate_base(text, description,):
     # Initialize variables
-    play_steps = int(frame_rate * play_steps_in_s)
-    chunk_size = 15  # Process 10 words at a time
+    chunk_size = 25  # Process max 25 words or a sentence at a time
     
     # Tokenize the full text and description
     inputs = description_tokenizer(description, return_tensors="pt").to(device)
-    
-    # Split text into chunks of approximately 10 words
-    words = text.split()
-    chunks = [' '.join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
-    
+
+    sentences_text = nltk.sent_tokenize(text) # this gives us a list of sentences
+    curr_sentence = ""
+    chunks = []
+    for sentence in sentences_text:
+        candidate = " ".join([curr_sentence, sentence])
+        if len(candidate.split()) >= chunk_size:
+            chunks.append(curr_sentence)
+            curr_sentence = sentence
+        else:
+            curr_sentence = candidate
+
+    if curr_sentence != "":
+        chunks.append(curr_sentence)
+        
+    print(chunks)
+
     all_audio = []
     
     # Process each chunk
@@ -223,8 +213,6 @@ def generate_base(text, description, play_steps_in_s=2.0):
             prompt_input_ids=prompt.input_ids,
             prompt_attention_mask=prompt.attention_mask,
             do_sample=True,
-            # temperature=1.0,
-            # min_new_tokens=10,
             return_dict_in_generate=True
         )
             
@@ -243,43 +231,30 @@ def generate_base(text, description, play_steps_in_s=2.0):
     print(f"Sample of length: {round(combined_audio.shape[0] / sampling_rate, 2)} seconds")
     yield numpy_to_mp3(combined_audio, sampling_rate=sampling_rate)
 
-# @spaces.GPU
-# def generate_jenny(text, description, play_steps_in_s=2.0):
-#     play_steps = int(frame_rate * play_steps_in_s)
-#     streamer = ParlerTTSStreamer(jenny_model, device=device, play_steps=play_steps)
-
-#     inputs = description_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=jenny_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_jenny(text, description, play_steps_in_s=2.0):
+def generate_finetuned(text, description):
     # Initialize variables
-    play_steps = int(frame_rate * play_steps_in_s)
-    chunk_size = 15  # Process 10 words at a time
+    chunk_size = 25  # Process max 25 words or a sentence at a time
     
     # Tokenize the full text and description
     inputs = description_tokenizer(description, return_tensors="pt").to(device)
-    
-    # Split text into chunks of approximately 10 words
-    words = text.split()
-    chunks = [' '.join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
+
+    sentences_text = nltk.sent_tokenize(text) # this gives us a list of sentences
+    curr_sentence = ""
+    chunks = []
+    for sentence in sentences_text:
+        candidate = " ".join([curr_sentence, sentence])
+        if len(candidate.split()) >= chunk_size:
+            chunks.append(curr_sentence)
+            curr_sentence = sentence
+        else:
+            curr_sentence = candidate
+
+    if curr_sentence != "":
+        chunks.append(curr_sentence)
+        
+    print(chunks)
     
     all_audio = []
     
@@ -289,14 +264,12 @@ def generate_jenny(text, description, play_steps_in_s=2.0):
         prompt = tokenizer(chunk, return_tensors="pt").to(device)
 
         # Generate audio for the chunk
-        generation = jenny_model.generate(
+        generation = finetuned_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,
-            # temperature=1.0,
-            # min_new_tokens=10,
             return_dict_in_generate=True
         )
 
@@ -387,29 +360,27 @@ with gr.Blocks(css=css) as block:
     with gr.Tab("Finetuned"):
         with gr.Row():
             with gr.Column():
-                input_text = gr.Textbox(label="Input Text", lines=2, value=jenny_examples[0][0], elem_id="input_text")
-                description = gr.Textbox(label="Description", lines=2, value=jenny_examples[0][1], elem_id="input_description")
-                play_seconds = gr.Slider(3.0, 7.0, value=jenny_examples[0][2], step=2, label="Streaming interval in seconds", info="Lower = shorter chunks, lower latency, more codec steps")
+                input_text = gr.Textbox(label="Input Text", lines=2, value=finetuned_examples[0][0], elem_id="input_text")
+                description = gr.Textbox(label="Description", lines=2, value=finetuned_examples[0][1], elem_id="input_description")
                 run_button = gr.Button("Generate Audio", variant="primary")
             with gr.Column():
-                audio_out = gr.Audio(label="Parler-TTS generation", format="mp3", elem_id="audio_out", streaming=True, autoplay=True)
+                audio_out = gr.Audio(label="Parler-TTS generation", format="mp3", elem_id="audio_out", autoplay=True)
 
-        inputs = [input_text, description, play_seconds]
+        inputs = [input_text, description]
         outputs = [audio_out]
-        gr.Examples(examples=jenny_examples, fn=generate_jenny, inputs=inputs, outputs=outputs, cache_examples=False)
-        run_button.click(fn=generate_jenny, inputs=inputs, outputs=outputs, queue=True)
+        gr.Examples(examples=finetuned_examples, fn=generate_finetuned, inputs=inputs, outputs=outputs, cache_examples=False)
+        run_button.click(fn=generate_finetuned, inputs=inputs, outputs=outputs, queue=True)
 
     with gr.Tab("Pretrained"):
         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, 7.0, value=3.0, step=2, 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 generation", format="mp3", elem_id="audio_out", streaming=True, autoplay=True)
+                audio_out = gr.Audio(label="Parler-TTS generation", format="mp3", elem_id="audio_out", autoplay=True)
 
-        inputs = [input_text, description, play_seconds]
+        inputs = [input_text, description]
         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)

requirements.txt

@@ -1,4 +1,4 @@
 torch
 spaces
 git+https://github.com/huggingface/parler-tts.git
-accelerate
+nltk