Update app.py

app.py

@@ -170,8 +170,10 @@ def main():
                                     file_name="output.mp3",
                                     mime="audio/mpeg",
                                 )
-            if detect_watermark_button:
-                with st.spinner("Detecting..."):
+            elif action == "Detect Watermark":
+                detect_watermark_button = st.button("Detect Watermark", key="detect_watermark_btn")                               
+                if detect_watermark_button:
+                    with st.spinner("Detecting..."):
                     # result, message = detector.detect_watermark(watermarked_audio, sample_rate=default_sr, message_threshold=0.5)
                     # st.markdown("Probability of audio being watermarked: ")
                     # st.markdown(result)
@@ -180,17 +182,17 @@ def main():
 
                     #Run on an unwatermarked audio
 
-                    if file_extension in [".wav", ".flac"]: 
-                        wav, sample_rate = torchaudio.load("test.wav")
-                        wav= wav.unsqueeze(0)
-
-                    elif file_extension == ".mp3":
-                        # Load an MP3 file
-                        audio = AudioSegment.from_mp3("test.mp3")
-                        # Export it as a WAV file
-                        audio.export("test.wav", format="wav")
-                        wav, sample_rate = torchaudio.load("test.wav")
-                        wav= wav.unsqueeze(0)
+                        if file_extension in [".wav", ".flac"]: 
+                            wav, sample_rate = torchaudio.load("test.wav")
+                            wav= wav.unsqueeze(0)
+
+                        elif file_extension == ".mp3":
+                            # Load an MP3 file
+                            audio = AudioSegment.from_mp3("test.mp3")
+                            # Export it as a WAV file
+                            audio.export("test.wav", format="wav")
+                            wav, sample_rate = torchaudio.load("test.wav")
+                            wav= wav.unsqueeze(0)
                         
                     result2, message2 = detector.detect_watermark(wav, sample_rate=default_sr, message_threshold=0.5)
                     print(f"This is likely an unwatermarked audio: {result2}")