handle audio with different shapes

app.py

@@ -21,10 +21,13 @@ def denoise(model_name, inputs):
         model = WaveUNetModel.from_pretrained(model_name)
     sr, audio = inputs
     audio = torch.from_numpy(audio)[None]
-    print(f"Audio shape: {audio.shape}")
 
     audio = audio / 32768.0
-    audio = audio.permute(0, 2, 1)
+
+    if audio.ndim == 2:
+        audio = audio.permute(1, 0)
+    else:
+        audio = audio.permute(0, 2, 1)
 
     print(f"Audio shape: {audio.shape}")
     print(f"Sample rate: {sr}")