Update app.py

app.py

@@ -750,7 +750,7 @@ from parler_tts import ParlerTTSForConditionalGeneration, ParlerTTSStreamer
 from transformers import AutoTokenizer
 from threading import Thread
 
-repo_id = "parler-tts/parler-tts-mini-v1"
+# repo_id = "parler-tts/parler-tts-mini-v1"
 
 
 
@@ -820,26 +820,52 @@ 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)
+
 
 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
-
-    # 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=device, play_steps=play_steps)
+        streamer = ParlerTTSStreamer(parler_model, device=torch_device, play_steps=play_steps)
 
-        inputs = parler_tokenizer(description, return_tensors="pt").to(device)
-        prompt = parler_tokenizer(text, return_tensors="pt").to(device)
+        inputs = parler_tokenizer(description, return_tensors="pt").to(torch_device)
+        prompt = parler_tokenizer(text, return_tensors="pt").to(torch_device)
 
         generation_kwargs = dict(
             input_ids=inputs.input_ids,