Add streaming audio loading and writing

app.py

@@ -1,4 +1,7 @@
 """Gradio demo for denoisers."""
+import tempfile
+from pathlib import Path
+
 import gradio as gr
 import numpy as np
 import torch
@@ -13,53 +16,57 @@ MODELS = [
 ]
 
 
-def denoise(model_name, inputs):
+def denoise(model_name: str, audio_path: str):
     """Denoise audio."""
     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)
-    audio = audio / 32768.0
 
-    if audio.ndim == 1:
-        audio = audio.unsqueeze(0)
+    if torch.cuda.is_available():
+        model = model.cuda()
 
-    print(f"Audio shape: {audio.shape}")
-    print(f"Sample rate: {sr}")
+    stream_reader = torchaudio.io.StreamReader(audio_path)
+    stream_reader.add_basic_audio_stream(
+        frames_per_chunk=model.config.max_length,
+        sample_rate=model.config.sample_rate,
+        num_channels=1,
+    )
 
-    if audio.shape[0] > 1:
-        audio = audio.mean(0, keepdim=True)
+    stream_writer = torchaudio.io.StreamWriter("denoised.wav")
+    stream_writer.add_audio_stream(sample_rate=model.config.sample_rate, num_channels=1)
 
-    print(f"Audio shape: {audio.shape}")
+    chunk_size = model.config.max_length
 
-    if sr != model.config.sample_rate:
-        audio = torchaudio.functional.resample(audio, sr, model.config.sample_rate)
+    with stream_writer.open():
+        for (audio_chunk,) in tqdm(stream_reader.stream()):
+            if audio_chunk is None:
+                break
 
-    chunk_size = model.config.max_length
+            audio_chunk = audio_chunk.permute(1, 0)
+            original_chunk_size = audio_chunk.size(-1)
 
-    padding = abs(audio.size(-1) % chunk_size - chunk_size)
-    padded = torch.nn.functional.pad(audio, (0, padding))
+            if audio_chunk.size(-1) < chunk_size:
+                padding = chunk_size - audio_chunk.size(-1)
+                audio_chunk = torch.nn.functional.pad(audio_chunk, (0, padding))
 
-    clean = []
-    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).audio
-        clean.append(clean_chunk.squeeze(0))
+            if torch.cuda.is_available():
+                audio_chunk = audio_chunk.cuda()
 
-    denoised = torch.concat(clean, 1)[:, : audio.shape[-1]].clamp(-1.0, 1.0)
-    denoised = (denoised * 32767.0).numpy().astype(np.int16)
+            with torch.no_grad():
+                denoised_chunk = model(audio_chunk[None]).audio
+                denoised_chunk = denoised_chunk[:, :, :original_chunk_size]
 
-    print(f"Denoised shape: {denoised.shape}")
+            stream_writer.write_audio_chunk(
+                0, denoised_chunk.squeeze(0).permute(1, 0).cpu()
+            )
 
-    return model.config.sample_rate, denoised.transpose()
+    return "denoised.wav"
 
 
 iface = gr.Interface(
     fn=denoise,
-    inputs=[gr.Dropdown(choices=MODELS, value=MODELS[0]), "audio"],
-    outputs="audio",
+    inputs=[gr.Dropdown(choices=MODELS, value=MODELS[0]), gr.Audio(type="filepath")],
+    outputs=gr.Audio(type="filepath"),
 )
 iface.launch()

requirements.txt

@@ -3,7 +3,7 @@ torch
 torchaudio
 pytorch-lightning
 pedalboard
-denoisers
+denoisers==0.1.7
 transformers
 librosa
 wandb