Update app.py

app.py

@@ -1,20 +1,26 @@
 import streamlit as st
 import torch
 import torchaudio
-from audiocraft.models import MusicGen
 import os
 import numpy as np
 import base64
+from audiocraft.models import MusicGen
+
+# Before
+batch_size = 64
 
-genres = ["Pop","Hip-Hop", "Classical","Lofi", "Chillpop","Country","R&G", "Folk","EDM", "Disco", "House", "Techno",]
+# After
+batch_size = 32 
+torch.cuda.empty_cache()
+
+genres = ["Pop", "Rock", "Jazz", "Electronic", "Hip-Hop", "Classical", "Lofi", "Chillpop"]
 
 @st.cache_resource()
 def load_model():
     model = MusicGen.get_pretrained('facebook/musicgen-small')
     return model
-    
 
-def generate_music_tensors(description, duration: int, batch_size=1):
+def generate_music_tensors(description, duration: int):
     model = load_model()
 
     model.set_generation_params(
@@ -32,9 +38,8 @@ def generate_music_tensors(description, duration: int, batch_size=1):
 
     st.success("Music Generation Complete!")
     return output
-    
 
-def save_audio(samples: torch.Tensor, filename):
+def save_audio(samples: torch.Tensor):
     sample_rate = 30000
     save_path = "audio_output"
     assert samples.dim() == 2 or samples.dim() == 3
@@ -44,9 +49,8 @@ def save_audio(samples: torch.Tensor, filename):
         samples = samples[None, ...]
 
     for idx, audio in enumerate(samples):
-        audio_path = os.path.join(save_path, f"{filename}_{idx}.wav")
+        audio_path = os.path.join(save_path, f"audio_{idx}.wav")
         torchaudio.save(audio_path, audio, sample_rate)
-        return audio_path
 
 def get_binary_file_downloader_html(bin_file, file_label='File'):
     with open(bin_file, 'rb') as f:
@@ -61,52 +65,56 @@ st.set_page_config(
 )
 
 def main():
-    st.title(" 🎶  AI Composer Small-Model  🎶")
+    st.title("🎧 AI Composer Medium-Model 🎧")
 
     st.subheader("Craft your perfect melody!")
-   
     bpm = st.number_input("Enter Speed in BPM", min_value=60)
-    text_area = st.text_area('Example: 80s rock song with guitar and drums', height=50)
-    selected_genre = st.selectbox("Select Genre (Optional)", genres, None)
-    time_slider = st.slider("Select time duration (In Seconds)", 0, 60, 10)
-    
-    mood = st.selectbox("Select Mood (Optional)", ["Happy", "Sad", "Angry", "Relaxed", "Energetic"], None)
-    instrument = st.selectbox("Select Instrument (Optional)", ["Piano", "Guitar", "Flute", "Violin", "Drums"], None)
-    tempo = st.selectbox("Select Tempo (Optional)", ["Slow", "Moderate", "Fast"], None)
-    melody = st.text_input("Enter Melody or Chord Progression (Optional)", "e.g: C D:min G:7 C, Twinkle Twinkle Little Star")
+
+    text_area = st.text_area('Ex : 80s rock song with guitar and drums')
+    st.text('')
+    # Dropdown for genres
+    selected_genre = st.selectbox("Select Genre", genres)
+
+    st.subheader("2. Select time duration (In Seconds)")
+    time_slider = st.slider("Select time duration (In Seconds)", 0, 30, 10)
+    # mood = st.selectbox("Select Mood (Optional)", ["Happy", "Sad", "Angry", "Relaxed", "Energetic"], None)
+    # instrument = st.selectbox("Select Instrument (Optional)", ["Piano", "Guitar", "Flute", "Violin", "Drums"], None)
+    # tempo = st.selectbox("Select Tempo (Optional)", ["Slow", "Moderate", "Fast"], None)
+    # melody = st.text_input("Enter Melody or Chord Progression (Optional)", "e.g: C D:min G:7 C, Twinkle Twinkle Little Star")
 
     if st.button('Let\'s Generate 🎶'):
         st.text('\n\n')
         st.subheader("Generated Music")
-
-        description = f"{text_area}"
+        
+        # Generate audio
+        description = text_area  # Initialize description with text_area
         if selected_genre:
             description += f" {selected_genre}"
+            st.empty()  # Hide the selected_genre selectbox after selecting one option
         if bpm:
             description += f" {bpm} BPM"
-        if mood:
-            description += f" {mood}"
-        if instrument:
-            description += f" {instrument}"
-        if tempo:
-            description += f" {tempo}"
-        if melody:
-            description += f" {melody}"
+        # if mood:
+        #     description += f" {mood}"
+        #     st.empty()  # Hide the mood selectbox after selecting one option
+        # if instrument:
+        #     description += f" {instrument}"
+        #     st.empty()  # Hide the instrument selectbox after selecting one option
+        # if tempo:
+        #     description += f" {tempo}"
+        #     st.empty()  # Hide the tempo selectbox after selecting one option
+        # if melody:
+        #     description += f" {melody}" 
+
+        # Clear CUDA memory cache before generating music
+        torch.cuda.empty_cache()
 
         music_tensors = generate_music_tensors(description, time_slider)
 
+        # Only play the full audio for index 0
         idx = 0
-        
-        # audio_path = save_audio(music_tensors[idx], "audio_output")
-        # audio_file = open(audio_path, 'rb')
-        # audio_bytes = audio_file.read()
-
-        # st.audio(audio_bytes, format='audio/wav')
-        # st.markdown(get_binary_file_downloader_html(audio_path, f'Audio_{idx}'), unsafe_allow_html=True)
-
         music_tensor = music_tensors[idx]
-        save_music_file = save_audio(music_tensor)
-        audio_filepath = f'audio_output/audio_{idx}.wav'
+        save_audio(music_tensor)
+        audio_filepath = f'/audio_output/audio_{idx}.wav'
         audio_file = open(audio_filepath, 'rb')
         audio_bytes = audio_file.read()
 
@@ -115,4 +123,4 @@ def main():
         st.markdown(get_binary_file_downloader_html(audio_filepath, f'Audio_{idx}'), unsafe_allow_html=True)
 
 if __name__ == "__main__":
-    main()
+    main()