Update app.py

app.py

@@ -1,22 +1,26 @@
 import streamlit as st
 import torch
 import torchaudio
-from audiocraft.models import MusicGen
 import os
 import numpy as np
 import base64
-# from torch.nn.utils.parametrizations import weight_norm
+from audiocraft.models import MusicGen
+
+# Before
+batch_size = 64
+
+# After
+batch_size = 32 
+torch.cuda.empty_cache()
 
-genres = ["Pop", "Rock", "Jazz", "Electronic", "Hip-Hop", "Classical",
-          "Lofi", "Chillpop","Country","R&G", "Folk","EDM", "Disco", "House", "Techno",]
+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(
@@ -26,23 +30,18 @@ def generate_music_tensors(description, duration: int, batch_size=1):
     )
 
     with st.spinner("Generating Music..."):
-        output = []
-        for i in range(0, len(description), batch_size):
-            batch_descriptions = description[i:i+batch_size]
-            batch_output = model.generate(
-                descriptions=batch_descriptions,
-                progress=True,
-                return_tokens=True
-            )
-            output.extend(batch_output)
+        output = model.generate(
+            descriptions=description,
+            progress=True,
+            return_tokens=True
+        )
 
     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 = "/content/drive/MyDrive/Colab Notebooks/audio_output"
+    save_path = "audio_output"
     assert samples.dim() == 2 or samples.dim() == 3
 
     samples = samples.detach().cpu()
@@ -50,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:
@@ -67,48 +65,62 @@ 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')
+        music_tensor = music_tensors[idx]
+        save_audio(music_tensor)
+        audio_filepath = f'/audio_output/audio_{idx}.wav'
+        audio_file = open(audio_filepath, 'rb')
         audio_bytes = audio_file.read()
 
+        # Play the full audio
         st.audio(audio_bytes, format='audio/wav')
-        st.markdown(get_binary_file_downloader_html(audio_path, f'Audio_{idx}'), unsafe_allow_html=True)
+        st.markdown(get_binary_file_downloader_html(audio_filepath, f'Audio_{idx}'), unsafe_allow_html=True)
 
 if __name__ == "__main__":
-    main()
-
+    main()