Update app.py

app.py

@@ -1,22 +1,29 @@
 import streamlit as st
 import torch
 import torchaudio
-from audiocraft.models import MusicGen
 import os
 import numpy as np
 import base64
+import math
+
+from audiocraft.data.audio_utils import convert_audio
+from audiocraft.data.audio import audio_write
+from audiocraft.models.encodec import InterleaveStereoCompressionModel
+from audiocraft.models import MusicGen, MultiBandDiffusion
+from audiocraft.utils.notebook import display_audio
+from audiocraft.models import MusicGen
+# from audiocraft.models import audiogen
 
 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" ]
+          "Lofi", "Chillpop","Country","R&G", "Folk","EDM", "Disco", "House", "Techno",]
 
 @st.cache_resource()
 def load_model():
     model = MusicGen.get_pretrained('facebook/musicgen-medium')
     return model
+    
 
-def generate_music_tensors(description, duration: int):
+def generate_music_tensors(description, duration: int, batch_size=1):
     model = load_model()
 
     model.set_generation_params(
@@ -26,14 +33,25 @@ def generate_music_tensors(description, duration: int):
     )
 
     with st.spinner("Generating Music..."):
-        output = model.generate(
-            descriptions=description,
-            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)
+
+        # output = model.generate(
+        #     descriptions=description,
+        #     progress=True,
+        #     return_tokens=True
+        # )
 
     st.success("Music Generation Complete!")
     return output
+    
 
 def save_audio(samples: torch.Tensor, filename):
     sample_rate = 30000
@@ -71,8 +89,7 @@ def main():
     text_area = st.text_area('Example: 80s rock song with guitar and drums')
     selected_genre = st.selectbox("Select Genre", genres)
     time_slider = st.slider("Select time duration (In Seconds)", 0, 30, 10)
-
-    st.write("Additional options")
+    
     mood = st.selectbox("Select Mood", ["Happy", "Sad", "Angry", "Relaxed", "Energetic"])
     instrument = st.selectbox("Select Instrument", ["Piano", "Guitar", "Flute", "Violin", "Drums"])
     tempo = st.selectbox("Select Tempo", ["Slow", "Moderate", "Fast"])
@@ -83,7 +100,8 @@ def main():
         st.subheader("Generated Music")
 
         description = f"{text_area} {selected_genre} {bpm} BPM {mood} {instrument} {tempo} {melody}"
-        music_tensors = generate_music_tensors(description, time_slider)
+        music_tensors = generate_music_tensors(description, time_slider, batch_size=2)
+
 
         idx = 0
         audio_path = save_audio(music_tensors[idx], "audio_output")
@@ -95,4 +113,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-