Add UNet1DModel

app.py

@@ -3,15 +3,22 @@ import gradio as gr
 import numpy as np
 import torch
 import torchaudio
-from denoisers import WaveUNetModel
+from denoisers import UNet1DModel, WaveUNetModel
 from tqdm import tqdm
 
-MODELS = ["wrice/waveunet-vctk-48khz", "wrice/waveunet-vctk-24khz"]
+MODELS = [
+    "wrice/unet1d-vctk-48khz",
+    "wrice/waveunet-vctk-48khz",
+    "wrice/waveunet-vctk-24khz",
+]
 
 
 def denoise(model_name, inputs):
     """Denoise audio."""
-    model = WaveUNetModel.from_pretrained(model_name)
+    if "unet1d" in model_name:
+        model = UNet1DModel.from_pretrained(model_name)
+    else:
+        model = WaveUNetModel.from_pretrained(model_name)
     sr, audio = inputs
     audio = torch.from_numpy(audio)[None]
     audio = audio / 32768.0