Restore the input's sample rate and channel count

app.py

@@ -1,4 +1,5 @@
 import math
+import librosa
 import tempfile
 from pathlib import Path
 
@@ -301,13 +302,24 @@ MODEL = LarsNet(wiener_filter=False, device="cpu", config="config.yaml")
 
 @torch.inference_mode()
 def process_fn(audio_path: str):
+    # The model works at its own rate, so the session's rate is read before
+    # anything touches the file and the stems are resampled back to it
+    original_sr = sf.info(audio_path).samplerate
+
     stems = MODEL(audio_path)
     output_dir = Path("outputs")
     output_dir.mkdir(exist_ok=True)
     output_paths = []
     for stem_name in ["kick", "snare", "toms", "hihat", "cymbals"]:
         out_path = output_dir / f"{stem_name}.wav"
-        sf.write(out_path, stems[stem_name].cpu().numpy().T, MODEL.sr)
+        stem = stems[stem_name].cpu().numpy().T
+
+        if original_sr != MODEL.sr:
+            stem = librosa.resample(
+                stem.T, orig_sr=MODEL.sr, target_sr=original_sr
+            ).T
+
+        sf.write(out_path, stem, original_sr)
         output_paths.append(str(out_path))
     return tuple(output_paths)