Update app.py

app.py

@@ -1,61 +1,42 @@
 import streamlit as st
 import torch
-import torchaudio
 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
+import torchaudio
+import numpy as np
 from audiocraft.models import MusicGen
-# from audiocraft.models import audiogen
 
 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-medium')
+def load_model(model_name):
+    model = MusicGen.get_pretrained(model_name)
     return model
-    
 
-def generate_music_tensors(description, duration: int, batch_size=1):
-    model = load_model()
-
-    model.set_generation_params(
-        use_sampling=True,
-        top_k=250,
-        duration=duration
-    )
-
-    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,
+def generate_music_tensors(description, duration: int, batch_size=1, models=None):
+    outputs = {}
+    for model_name, model in models.items():
+        model.set_generation_params(
+            use_sampling=True,
+            top_k=250,
+            duration=duration
+        )
+
+        with st.spinner(f"Generating Music with {model_name}..."):
+            output = model.generate(
+                descriptions=description,
                 progress=True,
                 return_tokens=True
             )
-            output.extend(batch_output)
-
-        # output = model.generate(
-        #     descriptions=description,
-        #     progress=True,
-        #     return_tokens=True
-        # )
+            outputs[model_name] = output
 
     st.success("Music Generation Complete!")
-    return output
-    
+    return outputs
 
 def save_audio(samples: torch.Tensor, filename):
     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()
@@ -65,7 +46,7 @@ def save_audio(samples: torch.Tensor, filename):
     for idx, audio in enumerate(samples):
         audio_path = os.path.join(save_path, f"{filename}_{idx}.wav")
         torchaudio.save(audio_path, audio, sample_rate)
-        return audio_path
+    return audio_path
 
 def get_binary_file_downloader_html(bin_file, file_label='File'):
     with open(bin_file, 'rb') as f:
@@ -80,7 +61,7 @@ st.set_page_config(
 )
 
 def main():
-    st.title("🎧AI Composer Medium-Model 🎧")
+    st.title("🎧 AI Composer 🎧")
 
     st.subheader("Generate Music")
     st.write("Craft your perfect melody! Fill in the blanks below to create your music masterpiece:")
@@ -95,21 +76,27 @@ def main():
     tempo = st.selectbox("Select Tempo", ["Slow", "Moderate", "Fast"])
     melody = st.text_input("Enter Melody or Chord Progression", "e.g., C D:min G:7 C, Twinkle Twinkle Little Star")
 
+    models = {
+        'Medium': load_model('facebook/musicgen-medium'),
+        'Large': load_model('facebook/musicgen-large'),
+        # Add more models here as needed
+    }
+
     if st.button('Let\'s Generate 🎶'):
         st.text('\n\n')
         st.subheader("Generated Music")
 
         description = f"{text_area} {selected_genre} {bpm} BPM {mood} {instrument} {tempo} {melody}"
-        music_tensors = generate_music_tensors(description, time_slider, batch_size=2)
-
+        music_outputs = generate_music_tensors(description, time_slider, batch_size=2, models=models)
 
-        idx = 0
-        audio_path = save_audio(music_tensors[idx], "audio_output")
-        audio_file = open(audio_path, 'rb')
-        audio_bytes = audio_file.read()
+        for model_name, output in music_outputs.items():
+            idx = 0  # Assuming you want to access the first audio file for each model
+            audio_filepath = save_audio(output, f'audio_{model_name}_{idx}')
+            audio_file = open(audio_filepath, '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)
+            st.audio(audio_bytes, format='audio/wav', label=f'{model_name} Model')
+            st.markdown(get_binary_file_downloader_html(audio_filepath, f'Audio_{model_name}_{idx}'), unsafe_allow_html=True)
 
 if __name__ == "__main__":
     main()