Upload install_spleeter.py

install_spleeter.py

@@ -0,0 +1,159 @@
+import gradio as gr
+import numpy as np
+import soundfile as sf
+import onnxruntime as ort
+from scipy.signal import stft, istft
+import os
+import uuid
+
+# ====== Load Model ======
+MODEL_PATH = "models/mdxnet.onnx"
+
+print("Loading MDX-Net model...")
+sess = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
+input_name = sess.get_inputs()[0].name
+input_shape = sess.get_inputs()[0].shape  # ['batch_size', 4, 3072, 256]
+output_name = sess.get_outputs()[0].name
+print("Loaded model with input shape:", input_shape)
+
+
+# ====== PARAMETERS (from UVR MDX pipeline) ======
+N_FFT = 6144
+HOP = 1536
+
+# For reconstruction (ISTFT) — correct value
+F_BIN_FULL = N_FFT // 2 + 1  # 3073
+
+# For MDX model input (fixed)
+F_BIN_MODEL = 3072
+
+CHUNK_T = 256
+N_CHANNELS = 4
+
+
+def prepare_spectrogram(audio, sr):
+    """Return stacked model input (4×3072×T) and full 3073-bin complex STFT."""
+
+    # Convert stereo → mono
+    if audio.ndim == 2:
+        audio = np.mean(audio, axis=1)
+
+    audio = audio.astype(np.float32)
+
+    # STFT
+    f, t, Zxx = stft(
+        audio,
+        fs=sr,
+        nperseg=N_FFT,
+        noverlap=N_FFT - HOP
+    )
+    Zxx = Zxx.astype(np.complex64)
+
+    # Full resolution for ISTFT later (3073)
+    original_spec = Zxx[:F_BIN_FULL, :]
+
+    # Model input uses ONLY 3072 bins
+    spec = Zxx[:F_BIN_MODEL, :]
+
+    mag = np.abs(spec).astype(np.float32)
+    phase = np.angle(spec).astype(np.float32)
+
+    ch1 = mag
+    ch2 = np.sin(phase).astype(np.float32)
+    ch3 = np.cos(phase).astype(np.float32)
+    ch4 = (mag ** 2).astype(np.float32)
+
+    stacked = np.stack([ch1, ch2, ch3, ch4], axis=0).astype(np.float32)
+    return stacked, original_spec
+
+
+def run_mdx(stacked):
+    _, _, T = stacked.shape
+    out = np.zeros((F_BIN_MODEL, T), dtype=np.float32)
+
+    for i in range(0, T, CHUNK_T):
+        chunk = stacked[:, :, i:i+CHUNK_T]
+
+        if chunk.shape[2] < CHUNK_T:
+            pad = CHUNK_T - chunk.shape[2]
+            chunk = np.pad(chunk, ((0,0),(0,0),(0,pad)), mode="constant")
+
+        inp = chunk[np.newaxis, :, :, :].astype(np.float32)
+
+        pred = sess.run([output_name], {input_name: inp})[0].astype(np.float32)
+        pred = pred[0, 0]  # mask channel
+
+        out[:, i:i+CHUNK_T] = pred[:, :out[:, i:i+CHUNK_T].shape[1]]
+
+    return out
+
+
+def reconstruct_audio(mask, original_spec, sr):
+    """Pad mask to 3073 bins, apply to full STFT, then ISTFT."""
+
+    # Pad 1 missing bin (3072 → 3073)
+    if mask.shape[0] == F_BIN_MODEL:
+        mask = np.pad(mask, ((0, 1), (0, 0)), mode="constant")
+
+    masked = mask * original_spec  # complex spectrogram
+
+    # ISTFT that requires N_FFT = 6144 → 3073 bins OK
+    _, audio = istft(
+        masked,
+        fs=sr,
+        nperseg=N_FFT,
+        noverlap=N_FFT - HOP
+    )
+    return audio
+
+
+def separate_audio(input_path):
+    if input_path is None:
+        return None, None
+
+    job = str(uuid.uuid4())
+    out_dir = f"output/{job}"
+    os.makedirs(out_dir, exist_ok=True)
+
+    audio, sr = sf.read(input_path)
+
+    # ---- Prepare spectrogram ----
+    stacked, original_spec = prepare_spectrogram(audio, sr)
+
+    # ---- Run MDX-Net ----
+    print("Running MDX-Net...")
+    mask = run_mdx(stacked)
+
+    mask = mask.astype(np.float32)
+    original_spec = original_spec.astype(np.complex64)
+
+    # ---- Reconstruct audio ----
+    vocals = reconstruct_audio(mask, original_spec, sr)
+
+    # Length-safe subtraction
+    instrumental = audio[:len(vocals)] - vocals[:len(audio)]
+
+    # ---- Save ----
+    v_path = os.path.join(out_dir, "vocals.wav")
+    i_path = os.path.join(out_dir, "instrumental.wav")
+
+    sf.write(v_path, vocals, sr)
+    sf.write(i_path, instrumental, sr)
+
+    return v_path, i_path
+
+
+# ===== Gradio UI =====
+with gr.Blocks(title="MDX-Net Vocal Remover (UVR)") as app:
+    gr.Markdown("# 🎵 MDX-Net Vocal Remover (UVR)")
+    gr.Markdown("Extract **Vocals** and **Instrumental** using the real UVR MDX-Net pipeline.")
+
+    inp = gr.Audio(type="filepath", label="Upload Audio")
+    btn = gr.Button("Separate", variant="primary")
+
+    out_v = gr.Audio(type="filepath", label="Vocals")
+    out_i = gr.Audio(type="filepath", label="Instrumental")
+
+    btn.click(separate_audio, inp, [out_v, out_i])
+
+app.launch()