Update generate_audio.py

generate_audio.py

@@ -43,10 +43,10 @@ class TTSGenerator:
         """
         
         # Load Bark model and processor for Speaker 2
-        # self.bark_processor = AutoProcessor.from_pretrained("suno/bark")
-        # self.bark_model = BarkModel.from_pretrained("suno/bark", torch_dtype=torch.float16).to(self.device)
-        # self.bark_sampling_rate = 24000
-        # self.voice_preset = "v2/en_speaker_6"
+        self.bark_processor = AutoProcessor.from_pretrained("suno/bark")
+        self.bark_model = BarkModel.from_pretrained("suno/bark", torch_dtype=torch.float16).to(self.device)
+        self.bark_sampling_rate = 24000
+        self.voice_preset = "v2/en_speaker_6"
 
     @spaces.GPU
     def load_transcript(self):
@@ -82,12 +82,12 @@ class TTSGenerator:
         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)
         
-        # Generate audio with input IDs and attention masks
+        # Pass all required arguments to generate() for reliable behavior
         generation = self.parler_model.generate(
             input_ids=input_ids,
-            attention_mask=attention_mask_input,
+            attention_mask=attention_mask_input,  # Set attention mask for input IDs
             prompt_input_ids=prompt_input_ids,
-            prompt_attention_mask=attention_mask_prompt
+            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
@@ -105,15 +105,34 @@ 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").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
+        
         # 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, return_tensors="pt", padding=True).to(self.device)
+        input_ids = inputs.input_ids.to(self.device)
+        attention_mask = inputs.attention_mask.to(self.device)
+        
+        # Generate speech output with both input IDs and attention mask
+        speech_output = self.bark_model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            temperature=0.9,
+            semantic_temperature=0.8
+        )
+        # Convert the generated audio to numpy array
+        audio_arr = speech_output[0].cpu().numpy()
+        return audio_arr, self.bark_sampling_rate
+
+# Convert the generated audio to numpy array
+audio_arr = speech_output[0].cpu().numpy()
 
     @staticmethod
     @spaces.GPU