handle multi-channel audio

app.py

@@ -22,12 +22,13 @@ def denoise(model_name, inputs):
     sr, audio = inputs
     audio = torch.from_numpy(audio)[None]
     audio = audio / 32768.0
+    audio = audio.permute(0, 2, 1)
 
     print(f"Audio shape: {audio.shape}")
     print(f"Sample rate: {sr}")
 
-    if audio.shape[-1] > 1:
-        audio = audio.mean(-1, keepdim=True)
+    if audio.shape[1] > 1:
+        audio = audio.mean(1, keepdim=True)
 
     print(f"Audio shape: {audio.shape}")
 
@@ -41,9 +42,9 @@ def denoise(model_name, inputs):
 
     clean = []
     for i in tqdm(range(0, padded.shape[-1], chunk_size)):
-        audio_chunk = padded[:, i : i + 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).logits
         clean.append(clean_chunk.squeeze(0))
 
     denoised = torch.concat(clean).flatten()[: audio.shape[-1]].clamp(-1.0, 1.0)