Updating to gradio

app.py

@@ -1,104 +1,69 @@
-import time
-import wavmark
-import streamlit as st
+import gradio as gr
+import numpy as np
 import os
-import torch
 import datetime
-import numpy as np
+import torch
 import soundfile
 from wavmark.utils import file_reader
-import subprocess
-import sys
-import time
-
-
+import wavmark
 
-def my_read_file(audio_path, max_second):
+def my_read_file(audio_path, max_second, default_sr=16000):
     signal, sr, audio_length_second = file_reader.read_as_single_channel_16k(audio_path, default_sr)
     if audio_length_second > max_second:
         signal = signal[0:default_sr * max_second]
         audio_length_second = max_second
-
     return signal, sr, audio_length_second
 
-
-def add_watermark(audio_path, watermark_text):
-    #t1 = time.time()
+def add_watermark(audio_path, watermark_text, max_second_encode=60):
     assert len(watermark_text) == 16
     watermark_npy = np.array([int(i) for i in watermark_text])
     signal, sr, audio_length_second = my_read_file(audio_path, max_second_encode)
     watermarked_signal, _ = wavmark.encode_watermark(model, signal, watermark_npy, show_progress=False)
-
     tmp_file_name = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + "_" + watermark_text + ".wav"
     tmp_file_path = '/tmp/' + tmp_file_name
     soundfile.write(tmp_file_path, watermarked_signal, sr)
-    #encode_time_cost = time.time() - t1
     return tmp_file_path
 
-#def encode_water()
-    
-def decode_watermark(audio_path):
+def decode_watermark(audio_path, max_second_decode=30):
     assert os.path.exists(audio_path)
-
     signal, sr, audio_length_second = my_read_file(audio_path, max_second_decode)
     payload_decoded, _ = wavmark.decode_watermark(model, signal, show_progress=False)
-
     if payload_decoded is None:
-       return "No Watermark"
-
-    payload_decoded_str = "".join([str(i) for i in payload_decoded])
-    st.write("Result:", payload_decoded_str)
-
-
-def create_default_value():
-    if "def_value" not in st.session_state:
-        def_val_npy = np.random.choice([0, 1], size=32 - len_start_bit)
-        def_val_str = "".join([str(i) for i in def_val_npy])
-        st.session_state.def_value = def_val_str
-
+        return "No Watermark"
+    return "".join([str(i) for i in payload_decoded])
 
+def create_default_value(len_start_bit=16):
+    def_val_npy = np.random.choice([0, 1], size=32 - len_start_bit)
+    return "".join([str(i) for i in def_val_npy])
 
 def main():
-    create_default_value()
-
-    
-    markdown_text = """
-    # Audio WaterMarking
-    You can upload an audio file and encode a custom 16-bit watermark or perform decoding from a watermarked audio.
-    
-    See [WaveMarktoolkit](https://github.com/wavmark/wavmark) for further details.
-    """
-
-    st.markdown(markdown_text)
-
-    audio_file = st.file_uploader("Upload Audio", type=["wav", "mp3"], accept_multiple_files=False)
-
-    if audio_file:
+    with gr.Blocks() as demo:
+        with gr.Row():
+            gr.Markdown("# Audio WaterMarking")
+        with gr.Row():
+            gr.Markdown("You can upload an audio file and encode a custom 16-bit watermark or perform decoding from a watermarked audio. See [WaveMark toolkit](https://github.com/wavmark/wavmark) for further details.")
         
-        tmp_input_audio_file = os.path.join("/tmp/", audio_file.name)
-        with open(tmp_input_audio_file, "wb") as f:
-            f.write(audio_file.getbuffer())
+        with gr.Row():
+            audio_file = gr.Audio(label="Upload Audio", type="filepath")
+            action = gr.Radio(["Add Watermark", "Decode Watermark"], label="Select Action")
+            watermark_text = gr.Textbox(label="The watermark (0, 1 list of length-16):", value=create_default_value())
+            submit_button = gr.Button("Submit")
 
+        with gr.Row():
+            output = gr.Audio(label="Processed Audio")
+            decode_output = gr.Textbox(label="Decoded Watermark")
 
+        def process_audio(audio_file, action, watermark_text):
+            if action == "Add Watermark" and audio_file:
+                return add_watermark(audio_file, watermark_text), None
+            elif action == "Decode Watermark" and audio_file:
+                return None, decode_watermark(audio_file)
+            else:
+                return None, None
 
-        action = st.selectbox("Select Action", ["Add Watermark", "Decode Watermark"])
-
-        if action == "Add Watermark":
-            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:
-                    with st.spinner("Adding Watermark..."):
-                        watermarked_audio = add_watermark(tmp_input_audio_file, watermark_text)
-                        st.write("Watermarked Audio:")
-                        print("watermarked_audio:", watermarked_audio)
-                        st.audio(watermarked_audio, format="audio/wav")
-                        
-        elif action == "Decode Watermark":
-            if st.button("Decode"):
-                with st.spinner("Decoding..."):
-                    decode_watermark(tmp_input_audio_file)
+        submit_button.click(process_audio, inputs=[audio_file, action, watermark_text], outputs=[output, decode_output])
 
+    demo.launch()
 
 if __name__ == "__main__":
     default_sr = 16000
@@ -108,5 +73,3 @@ if __name__ == "__main__":
     device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
     model = wavmark.load_model().to(device)
     main()
-
-