Update app.py

app.py

@@ -1,5 +1,6 @@
 from audiocraft.models import MusicGen
 import streamlit as st
+import os
 import torch
 import torchaudio
 from io import BytesIO
@@ -16,9 +17,7 @@ def generate_music_tensors(description, duration: int):
 
     model.set_generation_params(
         use_sampling=True,
-        top_k=300,
-        top_p=0.85,
-        temperature=0.8,
+        top_k=250,
         duration=duration
     )
 
@@ -31,15 +30,16 @@ def generate_music_tensors(description, duration: int):
 
 def save_audio_to_bytes(samples: torch.Tensor):
     sample_rate = 32000
-    assert samples.dim() == 3  # Expecting (batch, channels, samples)
-    samples = samples[0]  # Take the first batch item
+    assert samples.dim() == 2 or samples.dim() == 3
     samples = samples.detach().cpu()
 
-    # Save audio to a byte buffer instead of a file for easier download
-    byte_io = BytesIO()
-    torchaudio.save(byte_io, samples, sample_rate=sample_rate, format="wav")
-    byte_io.seek(0)  # Reset buffer position to the beginning for reading
-    return byte_io
+    if samples.dim() == 2:
+        samples = samples[None, ...]  # Add batch dimension if missing
+    
+    audio_buffer = BytesIO()
+    torchaudio.save(audio_buffer, samples[0], sample_rate=sample_rate, format="wav")
+    audio_buffer.seek(0)  # Move to the start of the buffer
+    return audio_buffer
 
 st.set_page_config(
     page_icon=":musical_note:",
@@ -50,29 +50,32 @@ def main():
     st.title("Your Music")
 
     with st.expander("See Explanation"):
-        st.write("App is developed using Meta's Audiocraft Music Gen model. Write a description and we will generate audio.")
+        st.write("This app uses Meta's Audiocraft Music Gen model to generate audio based on your description.")
 
     text_area = st.text_area("Enter description")
-    time_slider = st.slider("Select time duration (seconds)", 2, 20, 10)
+    time_slider = st.slider("Select time duration (seconds)", 2, 20, 5)
 
     if text_area and time_slider:
-        st.json({
-            "Description": text_area,
-            "Selected duration": time_slider
-        })
-
-        st.write("We will be back with your music....please enjoy doing the rest of your tasks while we come back in some time :)")
+        st.json(
+            {
+                "Description": text_area,
+                "Selected duration": time_slider
+            }
+        st.write("We will back with your music....please enjoy doing the rest of your tasks while we come back in some time :)")
+        )
         st.subheader("Generated Music")
         music_tensors = generate_music_tensors(text_area, time_slider)
-
-        # Save to byte buffer for download
-        audio_file = save_audio_to_bytes(music_tensors)
-
-        # Play and download audio
-        st.audio(audio_file, format="audio/wav")
+        
+        # Convert audio to bytes for playback and download
+        audio_buffer = save_audio_to_bytes(music_tensors)
+        
+        # Play audio
+        st.audio(audio_buffer, format="audio/wav")
+        
+        # Download button for audio
         st.download_button(
             label="Download Audio",
-            data=audio_file,
+            data=audio_buffer,
             file_name="generated_music.wav",
             mime="audio/wav"
         )