Refactor sound_generator.py to load models on GPU if available and streamline audio generation process

sound_generator.py

@@ -1,11 +1,14 @@
 from audiocraft.models import AudioGen, MusicGen
 from audiocraft.data.audio import audio_write
+import torch
 
-# Load the pretrained model (you can choose "small", "medium", or "large")
-sound_model = AudioGen.get_pretrained('facebook/audiogen-medium')
-music_model = MusicGen.get_pretrained('facebook/musicgen-small')
+# Load the pretrained models and move them to GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Set generation parameters (for example, audio duration of 8 seconds)
+sound_model = AudioGen.get_pretrained('facebook/audiogen-medium').to(device)
+music_model = MusicGen.get_pretrained('facebook/musicgen-small').to(device)
+
+# Set generation parameters (for example, audio duration of 5 seconds)
 sound_model.set_generation_params(duration=5)
 music_model.set_generation_params(duration=5)
 
@@ -19,17 +22,14 @@ def generate_sound(prompt: str):
     Returns:
     - str: The path to the saved audio file.
     """
-    # Generate the audio for the provided prompt
-    descriptions = [prompt]  # We use the prompt as a description for the model
-    wav = sound_model.generate(descriptions)  # Generates 2 samples
+    descriptions = [prompt]  
+    wav = sound_model.generate(descriptions)  # Generate audio
     
-    # Save the generated audio file with loudness normalization
     output_path = 'generated_audio'
     audio_write(output_path, wav[0].cpu(), sound_model.sample_rate, strategy="loudness")
     
     return f"{output_path}.wav"
 
-
 def generate_music(prompt: str):
     """
     Generate music using Audiocraft based on the given prompt.
@@ -40,11 +40,9 @@ def generate_music(prompt: str):
     Returns:
     - str: The path to the saved audio file.
     """
-    # Generate the music for the provided prompt
-    descriptions = [prompt]  # We use the prompt as a description for the model
-    wav = music_model.generate(descriptions)  # Generates 2 samples
+    descriptions = [prompt]  
+    wav = music_model.generate(descriptions)  # Generate music
     
-    # Save the generated audio file with loudness normalization
     output_path = 'generated_audio'
     audio_write(output_path, wav[0].cpu(), music_model.sample_rate, strategy="loudness")