Update app.py

app.py

@@ -12,53 +12,44 @@ genres = ["Pop", "Rock", "Jazz", "Electronic", "Hip-Hop", "Classical",
           "Techno","Indie Rock", "Grunge", "Ambient","Gospel", "Latin Music","Grime" ,"Trap", "Psychedelic Rock"  ]
 
 @st.cache_resource()
-
 def load_model():
     model = MusicGen.get_pretrained('facebook/musicgen-small')
     return model
 
-def generate_music_tensors(descriptions, duration: int, device):
-    # Load the model and move it to the specified device
+def generate_music_tensors(descriptions, duration: int):
     model = load_model()
-    model = model
+    # model = load_model().to('cpu')
+
 
-    
     model.set_generation_params(
         use_sampling=True,
         top_k=250,
-        duration=duration * 60  # Multiply by 60 to convert minutes to seconds
+        duration=duration
     )
 
     with st.spinner("Generating Music..."):
-        # Generate music using the model
         output = model.generate(
-        descriptions=descriptions,
-        progress=True,
-        return_tokens=True,
+            descriptions=descriptions,
+            progress=True,
+            return_tokens=True
         )
 
-# Save the generated music audio
-# Remove the device argument
-    save_audio(output)
-
     st.success("Music Generation Complete!")
     return output
 
 
-
-def save_audio(samples: torch.Tensor, device):
+def save_audio(samples: torch.Tensor):
     sample_rate = 30000
     save_path = "audio_output" 
     assert samples.dim() == 2 or samples.dim() == 3
 
-    samples = samples.to(device)  # Move the samples to the device
+    samples = samples.detach().cpu()
     if samples.dim() == 2:
         samples = samples[None, ...]
 
     for idx, audio in enumerate(samples):
         audio_path = os.path.join(save_path, f"audio_{idx}.wav")
-        torchaudio.save(audio_path, audio.cpu(), sample_rate)  # Move the audio to the CPU before saving
-
+        torchaudio.save(audio_path, audio, sample_rate)
 
 def get_binary_file_downloader_html(bin_file, file_label='File'):
     with open(bin_file, 'rb') as f:
@@ -74,7 +65,7 @@ st.set_page_config(
 
 def main():
     with st.sidebar:
-        st.header("""⚙️Generate Music ⚙️""", divider="rainbow")
+        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)
@@ -85,13 +76,13 @@ def main():
         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)
+        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,right_co = st.columns(2)
     left_co.write("""Music Generation through a prompt""")
     left_co.write(("""PS : First generation may take some time ......."""))
     
@@ -107,12 +98,11 @@ def main():
             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)
 
-            # Pass the device parameter when calling generate_music_tensors
-            device = torch.device('cpu')
-            music_tensors = generate_music_tensors(descriptions, time_slider, device)
-
+             # Only play the full audio for index 0
             idx = 0
             music_tensor = music_tensors[idx]
             save_music_file = save_audio(music_tensor)
@@ -127,3 +117,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+