Update app.py

app.py

@@ -212,33 +212,18 @@ def main():
 
                     # Run on an unwatermarked audio
 
-                    audio_file = st.file_uploader("Upload Audio", type=["wav", "mp3"], accept_multiple_files=False)
-
-                    if audio_file:
-                        tmp_input_audio_file = os.path.join("/tmp/", audio_file.name)
-                        file_extension = os.path.splitext(tmp_input_audio_file)[1].lower()
-                        if file_extension in [".wav", ".flac"]:
-                            with open("test.wav", "wb") as f:
-                                f.write(audio_file.getbuffer())
-
-                        elif file_extension == ".mp3":
-                            with open("test.mp3", "wb") as f:
-                                f.write(audio_file.getbuffer())
-
-                    #Load the WAV file using torchaudio
-                        if file_extension in [".wav", ".flac"]: 
-                            wav, sample_rate = torchaudio.load("test.wav")
-
-                        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")
-                            wav3, sample_rate = torchaudio.load("test.wav")
-                            wav3= wav3.unsqueeze(0)
+                    if file_extension in [".wav", ".flac"]: 
+                        wav, sample_rate = torchaudio.load("test.wav")
+
+                    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(wav3, sample_rate=default_sr, message_threshold=0.5)
+                    result2, message2 = detector.detect_watermark(wav, sample_rate=default_sr, message_threshold=0.5)
                     st.markdown("This is likely an unwatermarked audio:")
                     print(f"This is likely an unwatermarked audio: {result2}")
                     st.markdown(result2)