Update generate_audio.py

generate_audio.py

@@ -31,15 +31,16 @@ class TTSGenerator:
         
         # Set device
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        
+
+
         # Load Parler model and tokenizer for Speaker 1
-        self.parler_model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-mini-v1").to(self.device)
-        self.parler_tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler-tts-mini-v1")
+        self.parler_model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler_tts_mini_v0.1").to(self.device)
+        self.parler_tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler_tts_mini_v0.1")
         self.speaker1_description = """
-        Laura's voice is expressive and dramatic in delivery, speaking at a moderately fast pace with a very close recording that almost has no background noise.
+        Laura's voice is expressive and dramatic in delivery, speaking at a moderately fast pace with a very close recording that almost has no background noise and very clear audio.
         """
         self.speaker2_description = """
-        Gary's voice is expressive and dramatic in delivery, speaking at a moderately fast pace with a very close recording that almost has no background noise.
+        Gary's voice is expressive and dramatic in delivery, speaking at a moderately fast pace with a very close recording that almost has no background noise and very clear audio.
         """
         
         # Load Bark model and processor for Speaker 2
@@ -76,20 +77,21 @@ class TTSGenerator:
         # generation = self.parler_model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids)
         # audio_arr = generation.cpu().numpy().squeeze()
         # return audio_arr, self.parler_model.config.sampling_rate
-        input_ids = self.parler_tokenizer(self.speaker1_description, return_tensors="pt", padding=True).input_ids.to(self.device)
-        attention_mask_input = self.parler_tokenizer(self.speaker1_description, return_tensors="pt", padding=True).attention_mask.to(self.device)
-        
-        prompt_input_ids = self.parler_tokenizer(text, return_tensors="pt", padding=True).input_ids.to(self.device)
-        attention_mask_prompt = self.parler_tokenizer(text, return_tensors="pt", padding=True).attention_mask.to(self.device)
-        
-        # Pass all required arguments to generate() for reliable behavior
-        generation = self.parler_model.generate(
-            input_ids=input_ids,
-            attention_mask=attention_mask_input,  # Set attention mask for input IDs
-            prompt_input_ids=prompt_input_ids,
-            prompt_attention_mask=attention_mask_prompt  # Set prompt attention mask
-        )
-        audio_arr = generation.cpu().numpy().squeeze()
+        with torch.no_grad():
+            input_ids = self.parler_tokenizer(self.speaker1_description, return_tensors="pt", padding=True).input_ids.to(self.device)
+            attention_mask_input = self.parler_tokenizer(self.speaker1_description, return_tensors="pt", padding=True).attention_mask.to(self.device)
+            
+            prompt_input_ids = self.parler_tokenizer(text, return_tensors="pt", padding=True).input_ids.to(self.device)
+            attention_mask_prompt = self.parler_tokenizer(text, return_tensors="pt", padding=True).attention_mask.to(self.device)
+            
+            # Pass all required arguments to generate() for reliable behavior
+            generation = self.parler_model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask_input,  # Set attention mask for input IDs
+                prompt_input_ids=prompt_input_ids,
+                prompt_attention_mask=attention_mask_prompt  # Set prompt attention mask
+            )
+            audio_arr = generation.cpu().numpy().squeeze()
         return audio_arr, self.parler_model.config.sampling_rate
 
     #@spaces.GPU(duration=30)
@@ -104,21 +106,21 @@ class TTSGenerator:
             np.array: Audio array.
             int: Sampling rate.
         """
-
-        input_ids = self.parler_tokenizer(self.speaker2_description, return_tensors="pt", padding=True).input_ids.to(self.device)
-        attention_mask_input = self.parler_tokenizer(self.speaker1_description, return_tensors="pt", padding=True).attention_mask.to(self.device)
-        
-        prompt_input_ids = self.parler_tokenizer(text, return_tensors="pt", padding=True).input_ids.to(self.device)
-        attention_mask_prompt = self.parler_tokenizer(text, return_tensors="pt", padding=True).attention_mask.to(self.device)
-        
-        # Pass all required arguments to generate() for reliable behavior
-        generation = self.parler_model.generate(
-            input_ids=input_ids,
-            attention_mask=attention_mask_input,  # Set attention mask for input IDs
-            prompt_input_ids=prompt_input_ids,
-            prompt_attention_mask=attention_mask_prompt  # Set prompt attention mask
-        )
-        audio_arr = generation.cpu().numpy().squeeze()
+        with torch.no_grad():
+            input_ids = self.parler_tokenizer(self.speaker2_description, return_tensors="pt", padding=True).input_ids.to(self.device)
+            attention_mask_input = self.parler_tokenizer(self.speaker1_description, return_tensors="pt", padding=True).attention_mask.to(self.device)
+            
+            prompt_input_ids = self.parler_tokenizer(text, return_tensors="pt", padding=True).input_ids.to(self.device)
+            attention_mask_prompt = self.parler_tokenizer(text, return_tensors="pt", padding=True).attention_mask.to(self.device)
+            
+            # Pass all required arguments to generate() for reliable behavior
+            generation = self.parler_model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask_input,  # Set attention mask for input IDs
+                prompt_input_ids=prompt_input_ids,
+                prompt_attention_mask=attention_mask_prompt  # Set prompt attention mask
+            )
+            audio_arr = generation.cpu().numpy().squeeze()
         
         # inputs = self.bark_processor(text, voice_preset=self.voice_preset).to(self.device)
         # speech_output = self.bark_model.generate(**inputs, temperature=0.9, semantic_temperature=0.8)
@@ -175,6 +177,7 @@ class TTSGenerator:
                 final_audio = audio_segment
             else:
                 final_audio += audio_segment
+            torch.cuda.empty_cache()
 
         # Export final audio to MP3
         final_audio.export(self.output_audio_path, format="mp3", bitrate="192k", parameters=["-q:a", "0"])