Update app.py

app.py

@@ -1,6 +1,6 @@
 from audiocraft.models import MusicGen
 import streamlit as st
-import os 
+import os
 import torch
 import torchaudio
 import numpy as np
@@ -8,13 +8,13 @@ import base64
 
 @st.cache_resource
 def load_model():
-    model=MusicGen.get_pretrained("facebook/musicgen-small")
+    model = MusicGen.get_pretrained("facebook/musicgen-small")
     return model
 
-def generate_music_tensors(description,duration:int):
-    print("Description:",description)
-    print("Duration:",duration)
-    model=load_model()
+def generate_music_tensors(description, duration: int):
+    print("Description:", description)
+    print("Duration:", duration)
+    model = load_model()
 
     model.set_generation_params(
         use_sampling=True,
@@ -22,31 +22,33 @@ def generate_music_tensors(description,duration:int):
         duration=duration
     )
 
-    output=model.generate(
+    output = model.generate(
         descriptions=[description],
         progress=True,
         return_tokens=True
     )
     return output[0]
 
-def save_audio(samples:torch.tensor):
-    sample_rate=32000
-    save_path="audio_output/"
+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
 
-    assert samples.dim()==2 or samples.dim()==3
-    samples=samples.detach().cpu()
+    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_0.wav")
-        torchaudio.save(audio_path,audio,sample_rate)
+    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 {(bin_file)}">Download {file_label} from here</a>'
+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
 
 st.set_page_config(
@@ -59,23 +61,24 @@ def main():
 
     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")
-    text_area=st.text_area("Enter description")
-    time_slider=st.slider("Select time duration(s)",2,5,20)
+    
+    text_area = st.text_area("Enter description")
+    time_slider = st.slider("Select time duration(s)", 2, 5, 20)
     
     if text_area and time_slider:
         st.json(
             {
-                "Description":text_area,
-                "Selected duration:":time_slider
+                "Description": text_area,
+                "Selected duration": time_slider
             }
         )
         st.subheader("Generated Music")
-        music_tensors=generate_music_tensors(text_area,time_slider)
-        save_music_file=save_audio(music_tensors)
-        audio_file_path='audio_output/audio_0.wav'
-        audio_file=open(audio_file_path,'rb')
-        audio_bytes=audio_file.read()
+        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)
-if __name__=="__main__":
+        st.markdown(get_binary_file_downloader_html(audio_file_path, 'Audio'), unsafe_allow_html=True)
+
+if __name__ == "__main__":
     main()