Update app.py

app.py

@@ -18,8 +18,8 @@ args = {
     "device": "cuda" if torch.cuda.is_available() else "cpu"
 }
 
-# Load the pre-trained model
-model = ScoreModel.load_from_checkpoint(args["ckpt"])
+# Ensure the model is loaded to the correct device
+model = ScoreModel.load_from_checkpoint(args["ckpt"]).to(args["device"])
 
 def enhance_speech(audio_file):
     start_time = time.time()  # Start the timer
@@ -52,8 +52,8 @@ def enhance_speech(audio_file):
     # Renormalize
     x_hat = x_hat * norm_factor
 
-    # Save the enhanced audio to a temporary file for Gradio output
-    output_file = "enhanced_output.wav"
+    # Create a temporary path for saving the enhanced audio in Hugging Face Space
+    output_file = "/tmp/enhanced_output.wav"  # Use a temporary directory
     torchaudio.save(output_file, x_hat.cpu(), sr)
 
     print(f"Processed audio in {time.time() - start_time:.2f}s")