Update app.py

app.py

@@ -1,17 +1,11 @@
 import os
 import torch
-import torchaudio
-import torchaudio.transforms as T
+import numpy as np
+import soundfile as sf
 import gradio as gr
+from pydub import AudioSegment
 from multi_stem_pipeline import MultiStemSeparationModel
 
-# We import pydub to handle arbitrary audio formats (M4A, FLAC, AIFF, MP4, etc.)
-try:
-    from pydub import AudioSegment
-    HAS_PYDUB = True
-except ImportError:
-    HAS_PYDUB = False
-
 # =====================================================================
 # 1. ZEROGPU COMPATIBILITY LAYER
 # =====================================================================
@@ -29,30 +23,33 @@ STEMS_LIST = ["vocals", "backing_vocals", "drums", "bass", "guitar", "synth"]
 cached_model = None
 
 # =====================================================================
-# 2. AUDIO FORMAT CONVERTER
+# 2. AUDIO LOADER & RESAMPLER (No torchaudio)
 # =====================================================================
-def convert_to_wav(input_path):
+def load_audio_as_tensor(audio_path, target_sr=16000):
     """
-    Takes any input file (m4a, flac, aif, mp3, mp4) and converts it to a 
-    standard WAV file that torchaudio can safely load without backend errors.
+    Loads any audio format directly to a normalized mono PyTorch tensor
+    using pydub in memory (no external CLI conversion needed!).
     """
-    if not HAS_PYDUB:
-        # Fallback if pydub is missing
-        return input_path
-        
-    ext = os.path.splitext(input_path)[1].lower()
-    if ext == ".wav":
-        return input_path
-        
-    print(f"Converting {ext} file to temporary WAV format...")
-    temp_wav_path = "temp_converted_input.wav"
+    # Load audio (works on mp3, wav, flac, m4a, aiff, mp4 video, etc.)
+    audio = AudioSegment.from_file(audio_path)
+    
+    # Resample to 16kHz and convert to Mono
+    audio = audio.set_frame_rate(target_sr).set_channels(1)
     
-    # Load audio with pydub (handles AAC, ALAC, FLAC, AIFF, etc.)
-    audio = AudioSegment.from_file(input_path)
+    # Get raw samples as numpy array
+    samples = np.array(audio.get_array_of_samples(), dtype=np.float32)
     
-    # Export as standard WAV
-    audio.export(temp_wav_path, format="wav")
-    return temp_wav_path
+    # Normalize based on source bit depth
+    if audio.sample_width == 1:
+        samples = (samples - 128.0) / 128.0
+    elif audio.sample_width == 2:
+        samples = samples / 32768.0
+    elif audio.sample_width == 4:
+        samples = samples / 2147483648.0
+        
+    # Convert directly to PyTorch tensor [1, Samples]
+    waveform = torch.from_numpy(samples).unsqueeze(0)
+    return waveform
 
 
 # =====================================================================
@@ -67,8 +64,12 @@ def separate_multi_stems(audio_path):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     sample_rate = 16000
 
-    # 1. Pre-process and convert exotic formats to standard WAV
-    safe_wav_path = convert_to_wav(audio_path)
+    # 1. Load waveform in-memory using our pydub function
+    try:
+        waveform = load_audio_as_tensor(audio_path, sample_rate)
+    except Exception as e:
+        print(f"Error decoding audio file: {e}")
+        return [None] * len(STEMS_LIST)
 
     # 2. Lazy load our model
     if cached_model is None:
@@ -83,55 +84,50 @@ def separate_multi_stems(audio_path):
     model = cached_model.to(device)
     model.eval()
 
-    # 3. Load & Resample the safe WAV file
-    waveform, sr = torchaudio.load(safe_wav_path)
-    if sr != sample_rate:
-        resampler = T.Resample(orig_freq=sr, new_freq=sample_rate)
-        waveform = resampler(waveform)
-    
-    if waveform.shape[0] > 1:
-        waveform = torch.mean(waveform, dim=0, keepdim=True)
+    # Move audio tensors to computation device
+    waveform = waveform.to(device)
+    window = torch.hann_window(512).to(device)
 
-    # 4. Extract Spectrogram Magnitude & Phase
+    # 3. Extract Spectrogram Magnitude & Phase using PyTorch native stft
     complex_spec = torch.stft(
-        waveform, n_fft=512, win_length=512, hop_length=128, normalized=True, return_complex=True
+        waveform, n_fft=512, win_length=512, hop_length=128, window=window, normalized=True, return_complex=True
     )
     magnitude = torch.abs(complex_spec)
     phase = torch.angle(complex_spec)
 
-    input_tensor = magnitude.unsqueeze(0).to(device)
+    input_tensor = magnitude.unsqueeze(0) # Shape [1, 1, Freq, Time]
 
-    # 5. Run Multi-Mask Inference
+    # 4. Run Multi-Mask Inference
     with torch.no_grad():
         pred_stems_spec, masks = model(input_tensor)
 
     # Move tensors back to CPU for audio reconstruction
-    pred_stems_spec = pred_stems_spec.squeeze(0).cpu() 
-    phase = phase.cpu()
+    pred_stems_spec = pred_stems_spec.squeeze(0) # [num_stems, Freqs, Time]
+    phase = phase.squeeze(0)                      # [Freqs, Time]
+    window_cpu = torch.hann_window(512).cpu()
 
     saved_filepaths = []
     
-    # 6. Reconstruct each stem
+    # 5. Reconstruct and write each isolated stem (No torchaudio)
     for idx, stem_name in enumerate(STEMS_LIST):
-        stem_spec = pred_stems_spec[idx]
-        stem_complex = torch.polar(stem_spec, phase)
+        stem_spec = pred_stems_spec[idx].cpu()
+        stem_complex = torch.polar(stem_spec, phase.cpu())
+        
+        # Native PyTorch inverse Fourier transform to reconstruct waveform
         stem_wav = torch.istft(
-            stem_complex, n_fft=512, win_length=512, hop_length=128, normalized=True
+            stem_complex, n_fft=512, win_length=512, hop_length=128, window=window_cpu, normalized=True
         )
         
+        # Output clean WAV stem using soundfile
         out_path = f"output_{stem_name}.wav"
-        torchaudio.save(out_path, stem_wav, sample_rate)
+        sf.write(out_path, stem_wav.numpy(), sample_rate)
         saved_filepaths.append(out_path)
 
-    # Clean up temporary conversion file if it was created
-    if safe_wav_path == "temp_converted_input.wav" and os.path.exists(safe_wav_path):
-        os.remove(safe_wav_path)
-
     return saved_filepaths
 
 
 # =====================================================================
-# 4. DYNAMIC GRADIO INTERFACE GENERATION
+# 4. GRADIO INTERFACE LAYOUT
 # =====================================================================
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(