Update app.py

app.py

@@ -20,8 +20,11 @@ args = {
 model = ScoreModel.load_from_checkpoint(args["ckpt"])
 
 def enhance_speech(audio_file):
+    start_time = time.time()
+
     # Load and process the audio file
     y, sr = torchaudio.load(audio_file)
+    print(f"Loaded audio in {time.time() - start_time:.2f}s")
     T_orig = y.size(1)
 
     # Normalize
@@ -30,6 +33,8 @@ def enhance_speech(audio_file):
 
     # Prepare DNN input
     Y = torch.unsqueeze(model._forward_transform(model._stft(y.to(args["device"]))), 0)
+    print(f"Transformed input in {time.time() - start_time:.2f}s")
+
     Y = pad_spec(Y, mode="zero_pad")  # Use "zero_pad" mode for padding
 
     # Reverse sampling
@@ -48,9 +53,11 @@ def enhance_speech(audio_file):
     # Save the enhanced audio
     output_file = 'enhanced_output.wav'
     torchaudio.save(output_file, x_hat.cpu(), sr)
-    
+
+    print(f"Processed audio in {time.time() - start_time:.2f}s")
     return output_file
 
+
 # Gradio interface setup
 inputs = gr.Audio(label="Input Audio", type="filepath")
 outputs = gr.Audio(label="Output Audio", type="filepath")