Update app.py

app.py

@@ -5,6 +5,7 @@ import torch
 import torchaudio
 import numpy as np
 import base64
+from io import BytesIO
 
 @st.cache_resource
 def load_model():
@@ -16,9 +17,12 @@ def generate_music_tensors(description, duration: int):
     print("Duration:", duration)
     model = load_model()
 
+    # Experiment with different generation parameters for improved quality
     model.set_generation_params(
         use_sampling=True,
-        top_k=250,
+        top_k=300,          # Increase top_k for more diversity
+        top_p=0.85,         # Probability threshold for token sampling
+        temperature=0.8,    # Control randomness; lower values = more focused output
         duration=duration
     )
 
@@ -29,27 +33,19 @@ def generate_music_tensors(description, duration: int):
     )
     return output[0]
 
-def save_audio(samples: torch.Tensor):
-    sample_rate = 32000  # corrected to integer
-    save_path = "audio_output/"
-    os.makedirs(save_path, exist_ok=True)  # ensure directory exists
-
+def save_audio_to_bytes(samples: torch.Tensor):
+    sample_rate = 32000
     assert samples.dim() == 2 or samples.dim() == 3
     samples = samples.detach().cpu()
 
     if samples.dim() == 2:
         samples = samples[None, ...]
-    for idx, audio in enumerate(samples):
-        audio_path = os.path.join(save_path, f"audio_{idx}.wav")
-        torchaudio.save(audio_path, audio, sample_rate)
-    return os.path.join(save_path, "audio_0.wav")
-
-def get_binary_file_downloader_html(bin_file, file_label='File'):
-    with open(bin_file, 'rb') as f:
-        data = f.read()
-    bin_str = base64.b64encode(data).decode()
-    href = f'<a href="data:application/octet-stream;base64,{bin_str}" download="{file_label}">Download {file_label} from here</a>'
-    return href
+    
+    # Save audio to a byte buffer instead of 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
 
 st.set_page_config(
     page_icon=":musical_note:",
@@ -60,25 +56,33 @@ def main():
     st.title("Your Music")
 
     with st.expander("See Explanation"):
-        st.write("App is developed by using Meta's Audiocraft Music Gen model. Write your text and we will generate audio")
-    
+        st.write("App is developed using Meta's Audiocraft Music Gen model. Write a description and we will generate audio.")
+
     text_area = st.text_area("Enter description")
-    time_slider = st.slider("Select time duration(s)", 2, 5, 20)
-    
+    time_slider = st.slider("Select time duration (seconds)", 2, 20, 10)
+
     if text_area and time_slider:
         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)
-        audio_file_path = save_audio(music_tensors)
-        audio_file = open(audio_file_path, 'rb')
-        audio_bytes = audio_file.read()
-        st.audio(audio_bytes)
-        st.markdown(get_binary_file_downloader_html(audio_file_path, 'Audio'), unsafe_allow_html=True)
+
+        # 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")
+        st.download_button(
+            label="Download Audio",
+            data=audio_file,
+            file_name="generated_music.wav",
+            mime="audio/wav"
+        )
 
 if __name__ == "__main__":
-    main()
+    main()