Update app.py

app.py

@@ -213,7 +213,30 @@ def main():
                     print(f"\nThis is likely a watermarked audio: {result}")
 
                     # Run on an unwatermarked audio
-                    result2, message2 = detector.detect_watermark(wav, sample_rate=default_sr, message_threshold=0.5)
+                    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")
+                        wav= wav3.unsqueeze(0)
+                        
+                    result2, message2 = detector.detect_watermark(wav3, sample_rate=default_sr, message_threshold=0.5)
                     print(f"This is likely an unwatermarked audio: {result2}")
                     st.markdown(result)