Update app.py

app.py

@@ -12,43 +12,52 @@ 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
     model = load_model()
+    model = model.to(device)
 
     model.set_generation_params(
         use_sampling=True,
         top_k=250,
-        duration=duration
+        duration=duration * 60  # Multiply by 60 to convert minutes to seconds
     )
 
     with st.spinner("Generating Music..."):
+        # Generate music using the model
         output = model.generate(
             descriptions=descriptions,
             progress=True,
             return_tokens=True,
-            device=device  # Pass the device to the generate method
+            device=device
         )
 
+    # Save the generated music audio
+    save_audio(output, device)
+
     st.success("Music Generation Complete!")
     return output
 
 
-def save_audio(samples: torch.Tensor):
+
+def save_audio(samples: torch.Tensor, device):
     sample_rate = 30000
     save_path = "audio_output" 
     assert samples.dim() == 2 or samples.dim() == 3
 
-    samples = samples.detach().cpu()
+    samples = samples.to(device)  # Move the samples to the device
     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, sample_rate)
+        torchaudio.save(audio_path, audio.cpu(), sample_rate)  # Move the audio to the CPU before saving
+
 
 def get_binary_file_downloader_html(bin_file, file_label='File'):
     with open(bin_file, 'rb') as f: