StaticImages/stems
0
1import os2import gc3import queue4import threading5import json6import shlex7import subprocess8 9import librosa10import soundfile as sf11import numpy as np12import torch13from tqdm import tqdm14import onnxruntime as ort15from scipy.signal import butter, filtfilt16 17# ------------------------------------------------------------18# GLOBAL PATHS19# ------------------------------------------------------------20 21BASE_DIR = "."22mdxnet_models_dir = os.path.join(BASE_DIR, "mdx_models")23output_dir = os.path.join(BASE_DIR, "clean_song_output")24 25os.makedirs(output_dir, exist_ok=True)26 27# ------------------------------------------------------------28# BASIC UTILITIES29# ------------------------------------------------------------30 31def create_directories(path):32 os.makedirs(path, exist_ok=True)33 34def remove_directory_contents(path):35 if not os.path.exists(path):36 return37 for f in os.listdir(path):38 fp = os.path.join(path, f)39 if os.path.isfile(fp):40 os.remove(fp)41 else:42 import shutil43 shutil.rmtree(fp)44 45def convert_to_stereo_and_wav(audio_path):46 if audio_path is None:47 raise ValueError("No audio file received.")48 49 wave, sr = librosa.load(audio_path, mono=False, sr=44100)50 51 if wave.ndim == 1 or not audio_path.lower().endswith(".wav"):52 stereo_path = os.path.join(53 output_dir,54 f"{os.path.splitext(os.path.basename(audio_path))[0]}_stereo.wav"55 )56 cmd = shlex.split(57 f'ffmpeg -y -loglevel error -i "{audio_path}" -ac 2 -f wav "{stereo_path}"'58 )59 subprocess.run(cmd, check=True)60 return stereo_path61 62 return audio_path63 64# ------------------------------------------------------------65# SIMPLE FILTER HELPERS66# ------------------------------------------------------------67 68def butter_filter(data, cutoff, sr, btype, order=4):69 nyq = 0.5 * sr70 norm = cutoff / nyq71 b, a = butter(order, norm, btype=btype)72 return filtfilt(b, a, data, axis=0)73 74def highpass(data, cutoff, sr, order=4):75 return butter_filter(data, cutoff, sr, "highpass", order)76 77def lowpass(data, cutoff, sr, order=4):78 return butter_filter(data, cutoff, sr, "lowpass", order)79 80def highshelf_simple(data, sr, freq=10000.0, gain_db=2.0):81 # crude: tilt HF by simple highpass + mix82 if gain_db == 0:83 return data84 amt = 10 ** (gain_db / 20.0) - 1.085 hp = highpass(data, freq, sr)86 return data + amt * hp87 88# ------------------------------------------------------------89# MDX MODEL CLASSES90# ------------------------------------------------------------91 92class MDXModel:93 def __init__(self, device, dim_f, dim_t, n_fft, hop=1024, stem_name=None, compensation=1.0):94 self.dim_f = dim_f95 self.dim_t = dim_t96 self.dim_c = 497 self.n_fft = n_fft98 self.hop = hop99 self.stem_name = stem_name100 self.compensation = compensation101 102 self.n_bins = self.n_fft // 2 + 1103 self.chunk_size = hop * (self.dim_t - 1)104 self.window = torch.hann_window(self.n_fft, periodic=True).to(device)105 106 self.freq_pad = torch.zeros(107 [1, self.dim_c, self.n_bins - self.dim_f, self.dim_t]108 ).to(device)109 110 def stft(self, x):111 x = x.reshape([-1, self.chunk_size])112 x = torch.stft(113 x,114 n_fft=self.n_fft,115 hop_length=self.hop,116 window=self.window,117 center=True,118 return_complex=True,119 )120 x = torch.view_as_real(x)121 x = x.permute([0, 3, 1, 2])122 x = x.reshape([-1, 4, self.n_bins, self.dim_t])123 return x[:, :, : self.dim_f]124 125 def istft(self, x, freq_pad=None):126 freq_pad = self.freq_pad.repeat([x.shape[0], 1, 1, 1]) if freq_pad is None else freq_pad127 x = torch.cat([x, freq_pad], -2)128 x = x.reshape([-1, 2, 2, self.n_bins, self.dim_t]).reshape([-1, 2, self.n_bins, self.dim_t])129 x = x.permute([0, 2, 3, 1]).contiguous()130 x = torch.view_as_complex(x)131 x = torch.istft(132 x,133 n_fft=self.n_fft,134 hop_length=self.hop,135 window=self.window,136 center=True,137 )138 return x.reshape([-1, 2, self.chunk_size])139 140class MDX:141 DEFAULT_SR = 44100142 DEFAULT_CHUNK_SIZE = 0 * DEFAULT_SR143 DEFAULT_MARGIN_SIZE = 1 * DEFAULT_SR144 145 def __init__(self, model_path, params: MDXModel, processor=0):146 self.device = torch.device("cpu")147 self.provider = ["CPUExecutionProvider"]148 149 self.model = params150 self.ort = ort.InferenceSession(model_path, providers=self.provider)151 152 # warmup153 self.ort.run(None, {"input": torch.rand(1, 4, params.dim_f, params.dim_t).numpy()})154 155 self.process = lambda spec: self.ort.run(None, {"input": spec.cpu().numpy()})[0]156 self.prog = None157 158 @staticmethod159 def segment(wave, combine=True, chunk_size=DEFAULT_CHUNK_SIZE, margin_size=DEFAULT_MARGIN_SIZE):160 if combine:161 out = None162 for i, seg in enumerate(wave):163 start = 0 if i == 0 else margin_size164 end = None if i == len(wave) - 1 else -margin_size165 if out is None:166 out = seg[:, start:end]167 else:168 out = np.concatenate((out, seg[:, start:end]), axis=-1)169 return out170 171 out = []172 total = wave.shape[-1]173 if chunk_size <= 0 or chunk_size > total:174 chunk_size = total175 if margin_size > chunk_size:176 margin_size = chunk_size177 178 for i, skip in enumerate(range(0, total, chunk_size)):179 margin = 0 if i == 0 else margin_size180 end = min(skip + chunk_size + margin_size, total)181 start = skip - margin182 out.append(wave[:, start:end].copy())183 if end == total:184 break185 return out186 187 def pad_wave(self, wave):188 n = wave.shape[1]189 trim = self.model.n_fft // 2190 gen = self.model.chunk_size - 2 * trim191 pad = gen - n % gen192 193 wave_p = np.concatenate(194 (np.zeros((2, trim)), wave, np.zeros((2, pad)), np.zeros((2, trim))),195 axis=1196 )197 198 chunks = []199 for i in range(0, n + pad, gen):200 chunks.append(wave_p[:, i:i + self.model.chunk_size])201 202 chunks = np.stack(chunks, axis=0)203 chunks = torch.tensor(chunks, dtype=torch.float32)204 return chunks, pad, trim205 206 def _process_wave(self, mix_waves, trim, pad, q, idx):207 mix_waves = mix_waves.split(1)208 out = []209 with torch.no_grad():210 for mw in mix_waves:211 self.prog.update()212 spec = self.model.stft(mw)213 proc = torch.tensor(self.process(spec))214 wav = self.model.istft(proc)215 wav = wav[:, :, trim:-trim].transpose(0, 1).reshape(2, -1).numpy()216 out.append(wav)217 out = np.concatenate(out, axis=-1)[:, :-pad]218 q.put({idx: out})219 return out220 221 def process_wave(self, wave, mt_threads=1):222 self.prog = tqdm(total=0)223 chunk = wave.shape[-1] // mt_threads224 waves = self.segment(wave, False, chunk)225 226 q = queue.Queue()227 threads = []228 for idx, batch in enumerate(waves):229 mix, pad, trim = self.pad_wave(batch)230 self.prog.total = len(mix) * mt_threads231 t = threading.Thread(target=self._process_wave, args=(mix, trim, pad, q, idx))232 t.start()233 threads.append(t)234 235 for t in threads:236 t.join()237 238 self.prog.close()239 240 out = []241 while not q.empty():242 out.append(q.get())243 out = [list(d.values())[0] for d in sorted(out, key=lambda d: list(d.keys())[0])]244 return self.segment(out, True, chunk)245 246# ------------------------------------------------------------247# PARAM LOOKUP248# ------------------------------------------------------------249 250def _get_params_for_model(model_path, mdx_params):251 fname = os.path.basename(model_path)252 for mp in mdx_params.values():253 if mp.get("model_filename") == fname:254 return mp255 raise RuntimeError(f"No params found for model {fname}")256 257def _pick_model_for_stem(stem_name, mdx_params):258 for mp in mdx_params.values():259 if mp.get("primary_stem") == stem_name:260 fname = mp.get("model_filename")261 if fname:262 path = os.path.join(mdxnet_models_dir, fname)263 if os.path.exists(path):264 return path265 raise RuntimeError(f"No ONNX model found for stem '{stem_name}'")266 267# ------------------------------------------------------------268# CORE MDX RUNNER269# ------------------------------------------------------------270 271def run_mdx(model_params, output_dir, model_path, filename,272 suffix=None, denoise=False, m_threads=1):273 274 mp = _get_params_for_model(model_path, model_params)275 276 model = MDXModel(277 device=torch.device("cpu"),278 dim_f=mp["mdx_dim_f_set"],279 dim_t=2 ** mp["mdx_dim_t_set"],280 n_fft=mp["mdx_n_fft_scale_set"],281 stem_name=mp["primary_stem"],282 compensation=mp["compensate"],283 )284 285 mdx = MDX(model_path, model, processor=-1)286 287 wave, sr = librosa.load(filename, mono=False, sr=44100)288 if wave.ndim == 1:289 wave = np.stack([wave, wave], axis=0)290 291 peak = max(np.max(wave), abs(np.min(wave)))292 if peak == 0:293 peak = 1.0294 wave /= peak295 296 if denoise:297 out = -(mdx.process_wave(-wave, m_threads)) + mdx.process_wave(wave, m_threads)298 out *= 0.5299 else:300 out = mdx.process_wave(wave, m_threads)301 302 out *= peak303 304 stem_name = suffix or model.stem_name305 out_path = os.path.join(306 output_dir,307 f"{os.path.splitext(os.path.basename(filename))[0]}_{stem_name}.wav"308 )309 sf.write(out_path, out.T, sr)310 311 del mdx, wave312 gc.collect()313 314 return out_path, out315 316# ------------------------------------------------------------317# STUDIO CLEANUP PIPELINE318# ------------------------------------------------------------319 320def dc_remove(stem):321 return stem - np.mean(stem, axis=0, keepdims=True)322 323def normalize_peak(stem, target_db):324 peak = np.max(np.abs(stem))325 if peak == 0:326 return stem327 target_lin = 10 ** (target_db / 20.0)328 gain = target_lin / peak329 return stem * gain330 331def mono_from_stereo(stem):332 mono = np.mean(stem, axis=1, keepdims=True)333 return np.repeat(mono, 2, axis=1)334 335def mono_below(stem, sr, cutoff=120.0):336 # crude: lowpass mid, replace both channels with mid in low band337 low = lowpass(stem, cutoff, sr)338 high = stem - low339 mid = np.mean(low, axis=1, keepdims=True)340 low_mono = np.repeat(mid, 2, axis=1)341 return low_mono + high342 343def debleed_stems(vocals, drums, bass, other, sr):344 # light, conservative de-bleed345 v = vocals.copy()346 d = drums.copy()347 b = bass.copy()348 o = other.copy()349 350 d_hp = highpass(d, 120.0, sr)351 b_lp = lowpass(b, 200.0, sr)352 353 v -= 0.05 * d_hp354 v -= 0.03 * b_lp355 356 d -= 0.02 * v357 d -= 0.01 * b358 359 b -= 0.01 * v360 b -= 0.02 * d361 362 o -= 0.03 * v363 o -= 0.02 * d364 o -= 0.02 * b365 366 return v, d, b, o367 368def enhance_drums(drums, sr):369 d = drums.copy()370 d = highshelf_simple(d, sr, freq=10000.0, gain_db=2.0)371 return d372 373def enhance_bass(bass, sr):374 b = bass.copy()375 b = highpass(b, 25.0, sr)376 # simple low-shelf: boost low band a bit377 low = lowpass(b, 80.0, sr)378 b += 0.5 * low379 return b380 381def deess_vocals(vocals, sr):382 v = vocals.copy()383 # crude de-ess: subtract a bit of bandpassed 5–8 kHz384 hp = highpass(v, 5000.0, sr)385 bp = hp - highpass(v, 8000.0, sr)386 v -= 0.3 * bp387 return v388 389def cleanup_stems(vocals, drums, bass, other, sr=44100):390 # 1) DC removal391 vocals = dc_remove(vocals)392 drums = dc_remove(drums)393 bass = dc_remove(bass)394 other = dc_remove(other)395 396 # 2) De-bleed397 vocals, drums, bass, other = debleed_stems(vocals, drums, bass, other, sr)398 399 # 3) Enhancements400 drums = enhance_drums(drums, sr)401 bass = enhance_bass(bass, sr)402 vocals = deess_vocals(vocals, sr)403 404 # 4) Stereo field405 vocals = mono_from_stereo(vocals)406 bass = mono_below(bass, sr, cutoff=120.0)407 408 # 5) Normalize409 vocals = normalize_peak(vocals, -1.0)410 drums = normalize_peak(drums, -1.5)411 bass = normalize_peak(bass, -2.0)412 other = normalize_peak(other, -3.0)413 414 return vocals, drums, bass, other415 416# ------------------------------------------------------------417# MAIN 4‑STEM SEPARATOR (KUIELAB‑ONLY)418# ------------------------------------------------------------419 420def separate_stems_4(media_file, remove_files_output_dir=False):421 422 if media_file is None:423 raise ValueError("No audio file received.")424 425 if remove_files_output_dir:426 remove_directory_contents(output_dir)427 428 media_file = convert_to_stereo_and_wav(media_file)429 430 with open(os.path.join(mdxnet_models_dir, "data.json")) as f:431 mdx_params = json.load(f)432 433 song_out = os.path.join(output_dir, "song_output")434 create_directories(song_out)435 436 vocals_model = _pick_model_for_stem("Vocals", mdx_params)437 drums_model = _pick_model_for_stem("Drums", mdx_params)438 bass_model = _pick_model_for_stem("Bass", mdx_params)439 other_model = _pick_model_for_stem("Other", mdx_params)440 441 vocals_path, vocals_raw = run_mdx(mdx_params, song_out, vocals_model, media_file, suffix="Vocals")442 drums_path, drums_raw = run_mdx(mdx_params, song_out, drums_model, media_file, suffix="Drums")443 bass_path, bass_raw = run_mdx(mdx_params, song_out, bass_model, media_file, suffix="Bass")444 other_path, other_raw = run_mdx(mdx_params, song_out, other_model, media_file, suffix="Other")445 446 # raw outputs are shape (2, n)447 vocals = vocals_raw.T448 drums = drums_raw.T449 bass = bass_raw.T450 other = other_raw.T451 452 vocals_c, drums_c, bass_c, other_c = cleanup_stems(vocals, drums, bass, other, sr=44100)453 454 # overwrite files with cleaned versions455 base = os.path.splitext(os.path.basename(media_file))[0]456 vocals_path = os.path.join(song_out, f"{base}_Vocals_clean.wav")457 drums_path = os.path.join(song_out, f"{base}_Drums_clean.wav")458 bass_path = os.path.join(song_out, f"{base}_Bass_clean.wav")459 other_path = os.path.join(song_out, f"{base}_Other_clean.wav")460 461 sf.write(vocals_path, vocals_c, 44100)462 sf.write(drums_path, drums_c, 44100)463 sf.write(bass_path, bass_c, 44100)464 sf.write(other_path, other_c, 44100)465 466 out = {467 "vocals": vocals_path,468 "drums": drums_path,469 "bass": bass_path,470 "other": other_path,471 }472 473 return out474 475# ------------------------------------------------------------476# GRADIO UI477# ------------------------------------------------------------478 479if __name__ == "__main__":480 import gradio as gr481 482 def separate_stems_ui(audio_file):483 if audio_file is None:484 raise ValueError("No audio file received.")485 result = separate_stems_4(audio_file)486 return (487 result["vocals"],488 result["drums"],489 result["bass"],490 result["other"],491 )492 493 with gr.Blocks() as demo:494 gr.Markdown("# Kuielab 4‑Stem MDX (Studio Clean)\nUpload a song → Vocals / Drums / Bass / Other (remix‑ready).")495 496 with gr.Column():497 audio_input = gr.Audio(498 label="Input audio",499 type="filepath",500 interactive=True501 )502 503 run_button = gr.Button("Separate Stems")504 505 vocals_out = gr.Audio(label="Vocals", type="filepath")506 drums_out = gr.Audio(label="Drums", type="filepath")507 bass_out = gr.Audio(label="Bass", type="filepath")508 other_out = gr.Audio(label="Other", type="filepath")509 510 run_button.click(511 fn=separate_stems_ui,512 inputs=audio_input,513 outputs=[vocals_out, drums_out, bass_out, other_out],514 )515 516 try:517 demo.queue(False)518 except TypeError:519 pass520 521 demo.launch(522 server_name="0.0.0.0",523 server_port=7860,524 share=False,525 ssr_mode=False526 )527