handling mp3 format

app.py

@@ -38,19 +38,17 @@ def get_cached_model(model_name: str):
 
 # Separation Logic (all stems)
 def separate_all_stems(audio_file_path: str, model_name: str):
-    model = get_cached_model(model_name)
-    #waveform, sr = torchaudio.load(audio_file_path)
-    #is_mono = waveform.shape[0] == 1
-    #if is_mono:
-        #waveform = waveform.repeat(2, 1)
-        
     signal = AudioSignal(audio_file_path)
     signal = signal.resample(44100)  # expects 44.1kHz
-    if signal.num_channels == 1:
+
+    is_mono = signal.num_channels == 1
+    if is_mono:
         signal = signal.convert_to(stereo=True)
 
-    waveform = torch.from_numpy(signal.audio_data).float()  # [channels, samples]
-    waveform = waveform.unsqueeze(0)  # [1, channels, samples]
+    sr = signal.sample_rate  
+    waveform = signal.audio_data.float()  # [channels, samples]
+    waveform = waveform.unsqueeze(0)      # [1, channels, samples]
+
 
     with torch.no_grad():
         stems_batch = apply_model(
@@ -80,11 +78,17 @@ def separate_all_stems(audio_file_path: str, model_name: str):
 def process_fn(audio_file_path, model_name):
     output_signals = separate_all_stems(audio_file_path, model_name)
 
+    is_mp3 = Path(audio_file_path).suffix.lower() == ".mp3"
+    extension = "mp3" if is_mp3 else "wav"
+
     outputs = []
     for stem_name, signal in zip(STEM_NAMES, output_signals):
-        filename = f"demucs_{model_name}_{stem_name.lower().replace(' ', '_')}.wav"
-        output_audio_path = save_audio(signal, filename)
-        outputs.append(output_audio_path)
+        filename = f"demucs_{model_name}_{stem_name.lower().replace(' ', '_')}.{extension}"
+        output_path = Path(filename)
+
+        # Use .export() to control output format
+        signal.export(output_path, format=extension)
+        outputs.append(str(output_path))
 
     return tuple(outputs)