Update app.py

app.py

@@ -133,10 +133,11 @@ def main():
         action = st.selectbox("Select Action", ["Add Watermark", "Detect Watermark"])
 
         if action == "Add Watermark":
-            watermark_text = st.text_input("The watermark (0, 1 list of length-16):", value=st.session_state.def_value)
+            #watermark_text = st.text_input("The watermark (0, 1 list of length-16):", value=st.session_state.def_value)
             add_watermark_button = st.button("Add Watermark", key="add_watermark_btn")
             if add_watermark_button:  # 点击按钮后执行的
-                if audio_file and watermark_text:
+                #if audio_file and watermark_text:
+                if audio_file:
                     with st.spinner("Adding Watermark..."):
                         #wav = my_read_file(wav,max_second_encode)
                         
@@ -192,9 +193,22 @@ def main():
 
         elif action == "Detect Watermark":
             detect_watermark_button = st.button("Detect Watermark", key="detect_watermark_btn")
+
+            if audio_file:
+                #1st attempt
+                watermark = model.get_watermark(wav, default_sr)
+                watermarked_audio = wav + watermark
+                print(watermarked_audio.size())
+                size = watermarked_audio.size()
+                #st.markdown(size)
+
+                print(watermarked_audio.squeeze())
+                squeeze = watermarked_audio.squeeze(1)
+                shape = squeeze.size()
+                
             if detect_watermark_button:
                 with st.spinner("Detecting..."):
-                    result, message = detector.detect_watermark(watermarked_audio, sample_rate=default_sr, message_threshold=0.5)
+                    result, message = detector.detect_watermark(squeeze, sample_rate=default_sr, message_threshold=0.5)
                     st.markdown(result)
                     print(f"\nThis is likely a watermarked audio: {result}")