Update app.py

app.py

@@ -746,9 +746,9 @@ def generate_audio_elevenlabs(text):
         return None
 
 
-from parler_tts import ParlerTTSForConditionalGeneration, ParlerTTSStreamer
-from transformers import AutoTokenizer
-from threading import Thread
+# from parler_tts import ParlerTTSForConditionalGeneration, ParlerTTSStreamer
+# from transformers import AutoTokenizer
+# from threading import Thread
 
 # repo_id = "parler-tts/parler-tts-mini-v1"
 
@@ -795,12 +795,11 @@ from threading import Thread
 #     audio_segments = []
 #     for (sampling_rate, audio_chunk) in generate(text, description, chunk_size_in_s):
 #         audio_segments.append(audio_chunk)
-#         # Here, you can save the chunk to a file or send it to a frontend
-#         # For example, you could write the chunk to a file immediately:
+            
 #         temp_audio_path = os.path.join(tempfile.gettempdir(), f"parler_tts_audio_chunk_{len(audio_segments)}.wav")
 #         write_wav(temp_audio_path, sampling_rate, audio_chunk.astype(np.float32))
 #         logging.debug(f"Saved chunk to {temp_audio_path}")
-#         # You could also send the chunk to a web client if this was a web application
+        
 
 #     # Combine all the audio chunks into one audio file
 #     combined_audio = np.concatenate(audio_segments)
@@ -811,61 +810,37 @@ from threading import Thread
 #     logging.debug(f"Combined audio saved to {combined_audio_path}")
 #     return combined_audio_path
 
+
 import concurrent.futures
 import tempfile
 import os
 import numpy as np
+import logging
+from queue import Queue
 from threading import Thread
-from transformers import AutoTokenizer
-from parler_tts import ParlerTTSForConditionalGeneration, ParlerTTSStreamer
 from scipy.io.wavfile import write as write_wav
-import logging
-import torch
-
-# Set up device and dtype
-torch_device = "cuda:0"  # Use "mps" for Mac or "cpu" if CUDA is unavailable
-torch_dtype = torch.bfloat16
-
-# Set model name and other configurations
-model_name = "parler-tts/parler-tts-mini-v1"
-attn_implementation = "eager"  # Options: "eager", "sdpa", "flash_attention_2"
-compile_mode = "default"  # Options: "default", "reduce-overhead"
-max_length = 50  # Set padding max length
-
-# Load the model with efficient attention and compile optimizations
-parler_tokenizer = AutoTokenizer.from_pretrained(model_name)
-parler_model = ParlerTTSForConditionalGeneration.from_pretrained(
-    model_name,
-    attn_implementation=attn_implementation
-).to(torch_device, dtype=torch_dtype)
-
-# Compile the forward pass for faster generation
-parler_model.generation_config.cache_implementation = "static"
-parler_model.forward = torch.compile(parler_model.forward, mode=compile_mode)
-
-# Warmup to optimize the model after compilation
-inputs = parler_tokenizer("This is for compilation", return_tensors="pt", padding="max_length", max_length=max_length).to(torch_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):
-    _ = parler_model.generate(**model_kwargs)
+from parler_tts import ParlerTTSForConditionalGeneration, ParlerTTSStreamer
+from transformers import AutoTokenizer
 
+repo_id = "parler-tts/parler-tts-mini-v1"
+device = "cuda:0"  # or "cpu" if CUDA is not available
 
 def generate_audio_parler_tts(text):
     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."
-    
-    chunk_size_in_s = 0.3  # Smaller chunk size for lower latency
+    chunk_size_in_s = 0.5
+
+    # Initialize the tokenizer and model
+    parler_tokenizer = AutoTokenizer.from_pretrained(repo_id)
+    parler_model = ParlerTTSForConditionalGeneration.from_pretrained(repo_id).to(device)
     sampling_rate = parler_model.audio_encoder.config.sampling_rate
     frame_rate = parler_model.audio_encoder.config.frame_rate
-    
-    play_steps = int(frame_rate * chunk_size_in_s)
 
-    def generate_chunks(text, description):
-        streamer = ParlerTTSStreamer(parler_model, device=torch_device, play_steps=play_steps)
+    def generate(text, description, play_steps_in_s=0.5):
+        play_steps = int(frame_rate * play_steps_in_s)
+        streamer = ParlerTTSStreamer(parler_model, device=device, play_steps=play_steps)
 
-        inputs = parler_tokenizer(description, return_tensors="pt").to(torch_device)
-        prompt = parler_tokenizer(text, return_tensors="pt").to(torch_device)
+        inputs = parler_tokenizer(description, return_tensors="pt").to(device)
+        prompt = parler_tokenizer(text, return_tensors="pt").to(device)
 
         generation_kwargs = dict(
             input_ids=inputs.input_ids,
@@ -874,7 +849,7 @@ def generate_audio_parler_tts(text):
             prompt_attention_mask=prompt.attention_mask,
             streamer=streamer,
             do_sample=True,
-            temperature=0.7,  # Lower temperature for faster generation
+            temperature=1.0,
             min_new_tokens=10,
         )
 
@@ -886,31 +861,37 @@ def generate_audio_parler_tts(text):
                 break
             yield sampling_rate, new_audio
 
-    def process_audio_chunks(chunks):
-        audio_segments = []
-        for sampling_rate, audio_chunk in chunks:
-            audio_segments.append(audio_chunk)
-            temp_audio_path = os.path.join(tempfile.gettempdir(), f"parler_tts_audio_chunk_{len(audio_segments)}.wav")
+    # Queue to hold the audio chunks
+    audio_queue = Queue()
+    combined_audio = []
+
+    def process_chunks():
+        while True:
+            sampling_rate, audio_chunk = audio_queue.get()
+            if audio_chunk is None:  # Stop processing when a None is received
+                break
+            combined_audio.append(audio_chunk)
+            temp_audio_path = os.path.join(tempfile.gettempdir(), f"parler_tts_audio_chunk_{len(combined_audio)}.wav")
             write_wav(temp_audio_path, sampling_rate, audio_chunk.astype(np.float32))
             logging.debug(f"Saved chunk to {temp_audio_path}")
-            # Optionally, send this chunk to the client in real-time
-        return audio_segments
+            # Start playing or buffering the audio chunk here if required
+            # (e.g., send to a player or a frontend for immediate playback)
 
+    # Start the chunk processing in a separate thread
     with concurrent.futures.ThreadPoolExecutor() as executor:
-        # Start processing audio chunks in a separate thread
-        future_chunks = executor.submit(process_audio_chunks, generate_chunks(text, description))
+        executor.submit(process_chunks)
 
-        # Continue with other tasks in parallel
-        # (e.g., you can update the chatbot interface, handle other requests, etc.)
+        for (sampling_rate, audio_chunk) in generate(text, description, chunk_size_in_s):
+            audio_queue.put((sampling_rate, audio_chunk))
 
-        # Wait for audio processing to complete and get the result
-        audio_segments = future_chunks.result()
+    # Signal the end of processing
+    audio_queue.put((None, None))
 
     # Combine all the audio chunks into one audio file
-    combined_audio = np.concatenate(audio_segments)
+    combined_audio_np = np.concatenate(combined_audio)
     combined_audio_path = os.path.join(tempfile.gettempdir(), "parler_tts_combined_audio_stream.wav")
 
-    write_wav(combined_audio_path, sampling_rate, combined_audio.astype(np.float32))
+    write_wav(combined_audio_path, sampling_rate, combined_audio_np.astype(np.float32))
 
     logging.debug(f"Combined audio saved to {combined_audio_path}")
     return combined_audio_path