Update app.py

app.py

@@ -179,17 +179,25 @@ def main():
                         watermarked_audio = wav + watermark
                         print(watermarked_audio.size())
                         size = watermarked_audio.size()
-                        st.markdown(size)
+                        #st.markdown(size)
 
                         print(watermarked_audio.squeeze())
                         squeeze = watermarked_audio.squeeze(1)
                         shape = squeeze.size()
-                        st.markdown(shape)
+                        #st.markdown(shape)
 
-                        st.markdown(squeeze)
+                        #st.markdown(squeeze)
                         
                         torchaudio.save("output.wav", squeeze, default_sr)
                         st.audio("output.wav", format="audio/wav")
+
+                        with open("output.wav", "wb") as file:
+                            st.download_button(
+                                label="Download Watermarked audio",
+                                data=file,
+                                file_name="output.wav",
+                                mime="audio/wav",
+                            )
                         
                         #2nd Attempt
                         # watermarked_audio = model(wav, sample_rate=default_sr, alpha=1)
@@ -213,10 +221,10 @@ def main():
                         #st.write("Time Cost: %d seconds" % encode_time_cost)
 
 #                         # st.button("Add Watermark", disabled=False)
-#         elif action == "Decode Watermark":
-#             if st.button("Decode"):
-#                 with st.spinner("Decoding..."):
-#                     decode_watermark(tmp_input_audio_file)
+        # elif action == "Decode Watermark":
+        #     if st.button("Decode"):
+        #         with st.spinner("Decoding..."):
+        #             decode_watermark(tmp_input_audio_file)
 
 
 if __name__ == "__main__":