Update app.py

app.py

@@ -16,10 +16,8 @@ def load_model():
     model = MusicGen.get_pretrained('facebook/musicgen-small')
     return model
 
-def generate_music_tensors(descriptions, duration: int):
-    # Create a new instance of the model with the desired device
-    device = torch.device('cpu')
-    model = MusicGen.get_pretrained('facebook/musicgen-small').to(device)
+def generate_music_tensors(descriptions, duration: int, device):
+    model = load_model()
 
     model.set_generation_params(
         use_sampling=True,
@@ -27,11 +25,13 @@ def generate_music_tensors(descriptions, duration: int):
         duration=duration
     )
 
-    with st.spinner("Generating Music..."):
+    # Set the device during the forward pass
+    with torch.no_grad():
         output = model.generate(
             descriptions=descriptions,
             progress=True,
-            return_tokens=True
+            return_tokens=True,
+            device=device
         )
 
     st.success("Music Generation Complete!")