Update pages/music_recommendations.py

pages/music_recommendations.py

@@ -1,9 +1,9 @@
 import streamlit as st
+from transformers import pipeline
+import scipy.io.wavfile
 from openai import OpenAI
 import time
-import torch
-import torchaudio
-from torchaudio.pipelines import MusicGen
+import numpy as np
 
 # Initialize the OpenAI client
 client = OpenAI(
@@ -11,10 +11,6 @@ client = OpenAI(
     base_url="https://api.aimlapi.com",
 )
 
-# Load MusicGen model
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = MusicGen.get_pretrained("facebook/musicgen-small", device=device)
-
 # Streamlit app layout
 st.title("Mood-based Music Generator")
 
@@ -39,19 +35,25 @@ if st.button("Generate Music"):
                 },
             ],
         )
+
         message = response.choices[0].message.content
         st.write(f"Assistant: {message}")
 
-        # Generate the music using the MusicGen model
-        output = model.generate([message], progress=True)
+        # Load the synthesizer model for music generation
+        synthesiser = pipeline("text-to-audio", "facebook/musicgen-small")
+
+        # Simulate a short wait to represent loading time for music generation
+        time.sleep(2)
 
-        # Convert the output tensor to a numpy array
-        audio_data = output.cpu().numpy().squeeze()
+        # Generate the music using the synthesizer model based on the message
+        music = synthesiser(message, forward_params={"do_sample": True, "guidance_scale": 1})
 
         # Save the generated audio to a file
         audio_filename = "musicgen_out.wav"
-        torchaudio.save(audio_filename, torch.tensor(audio_data).unsqueeze(0), model.sample_rate)
+        scipy.io.wavfile.write(audio_filename, rate=music["sampling_rate"], data=np.array(music["audio"]))
 
         st.success("Music has been generated!")
+
         # Play the generated audio in Streamlit
-        st.audio(audio_filename)
+        st.audio(audio_filename)
+