Update app.py

app.py

@@ -5,21 +5,19 @@ from audiocraft.models import MusicGen
 import os
 import numpy as np
 import base64
+# from torch.nn.utils.parametrizations import weight_norm
 
-genres = ["Pop", "Rock", "Jazz", "Electronic", "Hip-Hop", "Classical", 
-          "Lofi", "Chillpop","Country","R&G", "Folk","Heavy Metal", 
-          "EDM", "Soil", "Funk","Reggae", "Disco", "Punk Rock", "House",
-          "Techno","Indie Rock", "Grunge", "Ambient","Gospel", "Latin Music","Grime" ,"Trap", "Psychedelic Rock"  ]
+genres = ["Pop", "Rock", "Jazz", "Electronic", "Hip-Hop", "Classical",
+          "Lofi", "Chillpop","Country","R&G", "Folk","EDM", "Disco", "House", "Techno",]
 
 @st.cache_resource()
 def load_model():
     model = MusicGen.get_pretrained('facebook/musicgen-small')
     return model
+    
 
-def generate_music_tensors(descriptions, duration: int):
+def generate_music_tensors(description, duration: int, batch_size=1):
     model = load_model()
-    # model = load_model().to('cpu')
-
 
     model.set_generation_params(
         use_sampling=True,
@@ -28,19 +26,23 @@ def generate_music_tensors(descriptions, duration: int):
     )
 
     with st.spinner("Generating Music..."):
-        output = model.generate(
-            descriptions=descriptions,
-            progress=True,
-            return_tokens=True
-        )
+        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)
 
     st.success("Music Generation Complete!")
     return output
+    
 
-
-def save_audio(samples: torch.Tensor):
+def save_audio(samples: torch.Tensor, filename):
     sample_rate = 30000
-    save_path = "audio_output" 
+    save_path = "/content/drive/MyDrive/Colab Notebooks/audio_output"
     assert samples.dim() == 2 or samples.dim() == 3
 
     samples = samples.detach().cpu()
@@ -48,8 +50,9 @@ def save_audio(samples: torch.Tensor):
         samples = samples[None, ...]
 
     for idx, audio in enumerate(samples):
-        audio_path = os.path.join(save_path, f"audio_{idx}.wav")
+        audio_path = os.path.join(save_path, f"{filename}_{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:
@@ -64,56 +67,48 @@ st.set_page_config(
 )
 
 def main():
-    with st.sidebar:
-        st.header("""⚙️Generate Music ⚙️""",divider="rainbow")
-        st.text("")
-        st.subheader("1. Enter your music description.......")
-        bpm = st.number_input("Enter Speed in BPM", min_value=60)
-
-        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, 60, 10)
-        # time_slider = st.slider("Select time duration (In Minutes)", 0,300,10, step=1)
-
-
-    st.title("""🎵 Song Lab AI 🎵""")
-    st.text('')
-    left_co,right_co = st.columns(2)
-    left_co.write("""Music Generation through a prompt""")
-    left_co.write(("""PS : First generation may take some time ......."""))
+    st.title(" 🎶  AI Composer Small-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)
     
-    if st.sidebar.button('Generate !'):
-        with left_co:
-            st.text('')
-            st.text('')
-            st.text('')
-            st.text('')
-            st.text('')
-            st.text('')
-            st.text('\n\n')
-            st.subheader("Generated Music")
-
-            # Generate audio
-            # descriptions = [f"{text_area} {selected_genre} {bpm} BPM" for _ in range(5)]
-            descriptions = [f"{text_area} {selected_genre} {bpm} BPM" for _ in range(1)]  # Change the batch size to 1
-            music_tensors = generate_music_tensors(descriptions, time_slider)
-
-             # Only play the full audio for index 0
-            idx = 0
-            music_tensor = music_tensors[idx]
-            save_music_file = 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_filepath, f'Audio_{idx}'), unsafe_allow_html=True)
-
+    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}"
+        if selected_genre:
+            description += f" {selected_genre}"
+        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}"
+
+        music_tensors = generate_music_tensors(description, time_slider)
+
+        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)
 
 if __name__ == "__main__":
-    main()
+    main()
+