Update generate_audio.py

generate_audio.py

@@ -105,21 +105,30 @@ class TTSGenerator:
             int: Sampling rate.
         """
 
-        # input_ids = self.parler_tokenizer(self.speaker2_description, return_tensors="pt").input_ids.to(self.device)
-        # prompt_input_ids = self.parler_tokenizer(text, return_tensors="pt").input_ids.to(self.device)
-        # 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.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)
         # audio_arr = speech_output[0].cpu().numpy()
         # return audio_arr, self.bark_sampling_rate
         # Tokenize input text and obtain input IDs and attention mask
-        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)
-        audio_arr = speech_output[0].cpu().numpy()
-        return audio_arr, self.bark_sampling_rate
+        # 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)
+        # audio_arr = speech_output[0].cpu().numpy()
+        return audio_arr, self.parler_model.config.sampling_rate
 
     @staticmethod
     @spaces.GPU