Update app.py

app.py

@@ -1,34 +1,27 @@
 from transformers import AutoProcessor, MusicgenForConditionalGeneration
-from IPython.display import Audio
-import scipy
 import torch
 import streamlit as st
 
-
 def mu_gen(prompt):
-  processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
-  model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
-
-  device = torch.device("cpu")
-  model.to(device)
+    processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
+    model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
 
-  inputs = processor(
-      text = [str(prompt)],     # This line is correct
-      padding=True,
-      return_tensors="pt",
-  )
-  
-  inputs = {key: value.to(device) for key, value in inputs.items()}
+    device = torch.device("cpu")
+    model.to(device)
 
-    # Generate audio on CPU
-  audio_values = model.generate(**inputs, max_new_tokens=256)
-  sampling_rate = model.config.audio_encoder.sampling_rate
+    inputs = processor(
+        text=[str(prompt)],
+        padding=True,
+        return_tensors="pt",
+    )
 
-    # Create an Audio object from the generated audio
-  result = Audio(audio_values[0].numpy(), rate=sampling_rate)
+    inputs = {key: value.to(device) for key, value in inputs.items()}
 
-  return result
+    # Generate audio on CPU
+    audio_values = model.generate(**inputs, max_new_tokens=256)
+    sampling_rate = model.config.audio_encoder.sampling_rate
 
+    return audio_values, sampling_rate
 
 def main():
     st.title("Text to Music Generator")
@@ -39,12 +32,12 @@ def main():
     if st.button("Generate Music"):
         if prompt:
             # Call the mu_gen function to generate music
-            generated_music = mu_gen(prompt)
+            generated_music, sampling_rate = mu_gen(prompt)
 
             # Display the generated audio
-            st.audio(generated_music, format="audio/wav")
+            st.audio(generated_music[0].numpy(), format="audio/wav", sample_rate=sampling_rate)
         else:
             st.warning("Please enter a text prompt.")
 
 if __name__ == "__main__":
-    main()
+    main()