Update app.py

app.py

@@ -83,30 +83,39 @@ def generate_voice(script: str, speaker: str = "default"):
         return f"Error generating voice-over: {e}"
 
 # ---------------------------------------------------------------------
-# Music Generation Function
+# Music Generation Function (using facebook/musicgen-medium)
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=300)
 def generate_music(prompt: str, audio_length: int):
     try:
-        musicgen_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
-        musicgen_processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
+        # Load facebook/musicgen-medium model
+        musicgen_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-medium")
+        musicgen_processor = AutoProcessor.from_pretrained("facebook/musicgen-medium")
 
+        # Move the model to the appropriate device (CUDA or CPU)
         device = "cuda" if torch.cuda.is_available() else "cpu"
         musicgen_model.to(device)
 
+        # Prepare inputs
         inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt").to(device)
+
+        # Generate music
         outputs = musicgen_model.generate(**inputs, max_new_tokens=audio_length)
 
+        # Process audio data
         audio_data = outputs[0, 0].cpu().numpy()
         normalized_audio = (audio_data / max(abs(audio_data)) * 32767).astype("int16")
 
-        output_path = f"{tempfile.gettempdir()}/generated_music.wav"
+        # Save generated music to a file
+        output_path = f"{tempfile.gettempdir()}/musicgen_medium_generated_music.wav"
         write(output_path, 44100, normalized_audio)
 
         return output_path
+
     except Exception as e:
         return f"Error generating music: {e}"
 
+
 # ---------------------------------------------------------------------
 # Audio Blending Function with Ducking
 # ---------------------------------------------------------------------