Update app.py

app.py

@@ -22,14 +22,21 @@ def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass:
     yield None, stream_log("Starting separation process...")
     yield None, stream_log(f"Loading audio file: {audio_file}")
     
-    # Load the audio file with the correct parameters
-    wav = load_track(audio_file, device, audio_channels=2, samplerate=separator.samplerate)
+    # Load the audio file with the correct samplerate
+    wav, sr = load_track(audio_file, separator.samplerate)
     
-    # The load_track function already handles channel conversion and resampling,
-    # so we can remove the channel adjustment code here.
+    # Check the number of channels and adjust if necessary
+    if wav.dim() == 1:
+        wav = wav.unsqueeze(0)  # Add channel dimension if mono
+    if wav.shape[0] == 1:
+        wav = wav.repeat(2, 1)  # If mono, duplicate to stereo
+    elif wav.shape[0] > 2:
+        wav = wav[:2]  # If more than 2 channels, keep only the first two
 
+    wav = wav.to(device)
+    
     ref = wav.mean(0)
-    wav = (wav - ref.view(1, -1)).to(device)
+    wav = (wav - ref.view(1, -1))
     yield None, stream_log("Audio loaded successfully. Applying model...")
     sources = apply_model(separator, wav, device=device, progress=True)
     sources = sources * ref.view(1, -1) + ref.view(1, -1)