Update app.py

app.py

@@ -162,12 +162,29 @@ def main():
                         #st.markdown(squeeze)
                     if file_extension in [".wav", ".flac"]:
                         torchaudio.save("output.wav", squeeze, default_sr, bits_per_sample=16)
+                        watermarked_wav = torchaudio.save("output.wav", squeeze, default_sr, bits_per_sample=16)
                         st.audio("output.wav", format="audio/wav")
+                        
+                        st.download_button(
+                            label="Download watermarked audio",
+                            data=watermarked_wav,
+                            file_name="output.wav",
+                            mime="audio/wav",
+                        )
+
 
                     elif file_extension == ".mp3":
                         torchaudio.save("output.mp3", squeeze, default_sr)
+                        watermarked_mp3 = torchaudio.save("output.mp3", squeeze, default_sr)
                         st.audio("output.mp3", format="audio/mpeg")
 
+                        st.download_button(
+                            label="Download watermarked audio",
+                            data=mp3",
+                            file_name="output.mp3",
+                            mime="audio/mpeg",
+                        )
+
                         # st.download_button(
                         #     label="Download Watermarked audio",
                         #     data="output.wav",
@@ -198,6 +215,7 @@ def main():
 
 #                         # st.button("Add Watermark", disabled=False)
         elif action == "Detect Watermark":
+            Detect_watermark_button = st.button("Detect Watermark", key="detect_watermark_btn")
             if st.button("Detect"):
                 with st.spinner("Detecting..."):
                     result, message = detector.detect_watermark(watermarked_audio, sample_rate=default_sr, message_threshold=0.5)