Update app.py

app.py

@@ -8,9 +8,11 @@ from denoisers.demucs import Demucs
 import torch
 import torchaudio
 import yaml
+import argparse
 
 import os
 os.environ['CURL_CA_BUNDLE'] = ''
+SAMPLE_RATE = 32000
 
 
 def denoising_transform(audio, model):
@@ -19,16 +21,16 @@ def denoising_transform(audio, model):
     src_path.parent.mkdir(exist_ok=True, parents=True)
     tgt_path.parent.mkdir(exist_ok=True, parents=True)
     (ffmpeg.input(audio)
-     .output(src_path.as_posix(), acodec='pcm_s16le', ac=1, ar=22050)
+     .output(src_path.as_posix(), acodec='pcm_s16le', ac=1, ar=SAMPLE_RATE)
      .run()
      )
-    wav, rate = torchaudio.load(audio)
+    wav, rate = torchaudio.load(src_path)
     reduced_noise = model.predict(wav)
     torchaudio.save(tgt_path, reduced_noise, rate)
-    return tgt_path
+    return src_path, tgt_path
 
 
-def run_app(model_filename, config_filename):
+def run_app(model_filename, config_filename, port, concurrency_count, max_size):
     model_path = hf_hub_download(repo_id="BorisovMaksim/demucs", filename=model_filename)
     config_path = hf_hub_download(repo_id="BorisovMaksim/demucs", filename=config_filename)
     with open(config_path, 'r') as f:
@@ -37,26 +39,96 @@ def run_app(model_filename, config_filename):
     checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
     model.load_state_dict(checkpoint['model_state_dict'])
 
-    interface_demucs = gr.Interface(
-        fn=lambda x: denoising_transform(x, model),
-        inputs=gr.Audio(label="Source Audio", source="microphone", type='filepath'),
-        outputs=gr.Audio(label="Demucs", type='filepath'),
-        allow_flagging='never'
-    )
-    interface_spectral_gating = gr.Interface(
-        fn=lambda x: denoising_transform(x, SpectralGating()),
-        inputs=gr.Audio(label="Source Audio", source="microphone", type='filepath'),
-        outputs=gr.Audio(label="Spectral Gating", type='filepath'),
-        allow_flagging='never'
+    title = "Chinese-to-English Direct Speech-to-Speech Translation (BETA)"
+
+
+    with gr.Blocks(title=title) as app:
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown(
+                        """
+                    # Denoising
+                    ## Instruction: \n
+                    1. Press "Record from microphone"
+                    2. Press "Stop recording"
+                    3. Press "Enhance" \n
+                    - You can switch to the tab "File" to upload a prerecorded .wav audio  instead of recording from microphone.
+                    """
+                    )
+                with gr.Tab("Microphone"):
+                    microphone = gr.Audio(label="Source Audio", source="microphone", type='filepath')
+                    with gr.Row():
+                        microphone_button = gr.Button("Enhance", variant="primary")
+                with gr.Tab("File"):
+                    upload = gr.Audio(label="Upload Audio", source="upload", type='filepath')
+                    with gr.Row():
+                        upload_button = gr.Button("Enhance", variant="primary")
+                clear_btn = gr.Button("Clear")
+                gr.Examples(examples=[[path] for path in Path("testing/wavs/").glob("*.wav")],
+                           inputs=[microphone, upload])
+
+            with gr.Column():
+                outputs = [gr.Audio(label="Input Audio", type='filepath'),
+                           gr.Audio(label="Demucs Enhancement", type='filepath'),
+                           gr.Audio(label="Spectral Gating Enhancement", type='filepath')
+                           ]
+
+        def submit(audio):
+            src_path, demucs_tgt_path = denoising_transform(audio, model)
+            _, spectral_gating_tgt_path = denoising_transform(audio, SpectralGating())
+            return src_path, demucs_tgt_path, spectral_gating_tgt_path, gr.update(visible=False), gr.update(visible=False)
+        
+
+        
+        microphone_button.click(
+            submit,
+            microphone,
+            outputs + [microphone, upload] 
+        )
+        upload_button.click(
+            submit,
+            upload,
+            outputs + [microphone, upload]
+        )
+
+
+        def restart():
+            return microphone.update(visible=True, value=None), upload.update(visible=True, value=None), None, None, None
+        
+        clear_btn.click(restart, inputs=[], outputs=[microphone, upload] + outputs)
+
+    app.queue(concurrency_count=concurrency_count, max_size=max_size)
+    
+    app.launch(
+        ssl_verify=False,
+        server_name='0.0.0.0',
+        server_port=port,
+        ssl_keyfile='certificates/example.key',
+        ssl_certfile='certificates/example.crt',
     )
-    gr.Parallel(interface_demucs, interface_spectral_gating,
-                title="Denoising",
-                examples=[[path] for path in Path("testing/wavs/").glob("*.wav")]
-                ).launch(server_name='0.0.0.0',
-			  server_port=7860)
+
+
 
 
 if __name__ == "__main__":
-    model_filename = "paper_replica_10_epoch/Demucs_replicate_paper_continue_epoch45.pt"
-    config_filename = "paper_replica_10_epoch/config.yaml"
-    run_app(model_filename, config_filename)
+    parser = argparse.ArgumentParser(description='Running demo.')
+    parser.add_argument('--port',
+                        type=int,
+                        default=7860)
+    parser.add_argument('--model_filename',
+                        type=str,
+                        default="paper_replica_10_epoch/Demucs_replicate_paper_continue_epoch45.pt")
+    parser.add_argument('--config_filename',
+                        type=str,
+                        default="paper_replica_10_epoch/config.yaml")
+    parser.add_argument('--concurrency_count',
+                        type=int,
+                        default=4)
+    parser.add_argument('--max_size',
+                        type=int,
+                        default=15)
+    
+    args = parser.parse_args()
+    
+
+    run_app(args.model_filename, args.config_filename, args.port, args.concurrency_count, args.max_size)