add app.py

app.py

@@ -1,3 +1,4 @@
+"""Gradio demo for denoisers."""
 import gradio as gr
 import numpy as np
 import torch
@@ -5,10 +6,22 @@ import torchaudio
 from denoisers import WaveUNetModel
 from tqdm import tqdm
 
-MODEL = WaveUNetModel.from_pretrained("wrice/waveunet-vctk-24khz")
+MODELS = ["wrice/waveunet-vctk-48khz", "wrice/waveunet-vctk-24khz"]
 
 
-def denoise(inputs):
+def main():
+    """Main."""
+    iface = gr.Interface(
+        fn=denoise,
+        inputs=[gr.Dropdown(choices=MODELS, default=MODELS[0]), "audio"],
+        outputs="audio",
+    )
+    iface.launch()
+
+
+def denoise(model_name, inputs):
+    """Denoise audio."""
+    model = WaveUNetModel.from_pretrained(model_name)
     sr, audio = inputs
     audio = torch.from_numpy(audio)[None]
     audio = audio / 32768.0
@@ -16,10 +29,10 @@ def denoise(inputs):
     print(f"Audio shape: {audio.shape}")
     print(f"Sample rate: {sr}")
 
-    if sr != MODEL.config.sample_rate:
-        audio = torchaudio.functional.resample(audio, sr, MODEL.config.sample_rate)
+    if sr != model.config.sample_rate:
+        audio = torchaudio.functional.resample(audio, sr, model.config.sample_rate)
 
-    chunk_size = MODEL.config.max_length
+    chunk_size = model.config.max_length
 
     padding = abs(audio.size(-1) % chunk_size - chunk_size)
     padded = torch.nn.functional.pad(audio, (0, padding))
@@ -28,7 +41,7 @@ def denoise(inputs):
     for i in tqdm(range(0, padded.shape[-1], chunk_size)):
         audio_chunk = padded[:, i : i + chunk_size]
         with torch.no_grad():
-            clean_chunk = MODEL(audio_chunk[None]).logits
+            clean_chunk = model(audio_chunk[None]).logits
         clean.append(clean_chunk.squeeze(0))
 
     denoised = torch.concat(clean).flatten()[: audio.shape[-1]].clamp(-1.0, 1.0)
@@ -36,8 +49,8 @@ def denoise(inputs):
 
     print(f"Denoised shape: {denoised.shape}")
 
-    return MODEL.config.sample_rate, denoised
+    return model.config.sample_rate, denoised
 
 
-iface = gr.Interface(fn=denoise, inputs="audio", outputs="audio")
-iface.launch()
+if __name__ == "__main__":
+    main()