Update generate_audio.py

generate_audio.py

@@ -10,14 +10,14 @@ from scipy.io import wavfile
 from pydub import AudioSegment
 import io
 import ast
-
+import spaces
 
 
 class TTSGenerator:
     """
     A class to generate podcast-style audio from a transcript using ParlerTTS and Bark models.
     """
-
+    @spaces.GPU
     def __init__(self, transcript_file_path):
         """
         Initialize the TTS generator with the path to the rewritten transcript file.
@@ -53,7 +53,8 @@ class TTSGenerator:
         """
         with open(self.transcript_file_path, 'rb') as f:
             return ast.literal_eval(pickle.load(f))
-
+    
+    @spaces.GPU
     def generate_speaker1_audio(self, text):
         """
         Generate audio for Speaker 1 using ParlerTTS.
@@ -71,6 +72,7 @@ class TTSGenerator:
         audio_arr = generation.cpu().numpy().squeeze()
         return audio_arr, self.parler_model.config.sampling_rate
 
+    @spaces.GPU
     def generate_speaker2_audio(self, text):
         """
         Generate audio for Speaker 2 using Bark.
@@ -88,6 +90,7 @@ class TTSGenerator:
         return audio_arr, self.bark_sampling_rate
 
     @staticmethod
+    @spaces.GPU
     def numpy_to_audio_segment(audio_arr, sampling_rate):
         """
         Convert numpy array to AudioSegment.