From 2f46a8d3fdfb881de78a77402317573ec1061415 Mon Sep 17 00:00:00 2001 From: Paul Sajna Date: Fri, 12 May 2023 22:23:50 -0700 Subject: [PATCH 1/2] mdx net main stem separation --- stem_splitter/mdx_net_main/.gitignore | 2 + stem_splitter/mdx_net_main/infer.py | 90 ++++++++++++++++++++++ stem_splitter/mdx_net_main/models/.gitkeep | 0 3 files changed, 92 insertions(+) create mode 100644 stem_splitter/mdx_net_main/.gitignore create mode 100644 stem_splitter/mdx_net_main/infer.py create mode 100644 stem_splitter/mdx_net_main/models/.gitkeep diff --git a/stem_splitter/mdx_net_main/.gitignore b/stem_splitter/mdx_net_main/.gitignore new file mode 100644 index 0000000..f2b7995 --- /dev/null +++ b/stem_splitter/mdx_net_main/.gitignore @@ -0,0 +1,2 @@ +*.onnx +*.wav diff --git a/stem_splitter/mdx_net_main/infer.py b/stem_splitter/mdx_net_main/infer.py new file mode 100644 index 0000000..888ffaa --- /dev/null +++ b/stem_splitter/mdx_net_main/infer.py @@ -0,0 +1,90 @@ +import numpy as np +import torch +import torchaudio +import onnxruntime as ort +import librosa + +class MDXNetMain(): + + def __init__(self): + self.chunk_size = 261120 + self.n_fft = 7680 + self.hop = 1024 + self.dim_f = 3072 + self.dim_t = 256 + self.dim_c = 4 + self.n_bins = self.n_fft//2+1 + self.trim = self.n_fft//2 + self.window = torch.hann_window(self.n_fft) + self.freq_pad = torch.zeros([1, self.dim_c, self.n_bins-self.dim_f, self.dim_t]) + self.adjust = 1 + self.device="cpu" + self.ort_ = ort.InferenceSession("models/UVR_MDXNET_Main.onnx") + self.model_run = lambda spek:self.ort_.run(None, {'input': spek})[0] + + + def stft(self, x): + x = x.reshape([-1, self.chunk_size]) + x = torch.stft(torch.from_numpy(x), n_fft=self.n_fft, hop_length=self.hop, window=self.window, center=True,return_complex=True) + x=torch.view_as_real(x) + x = x.permute([0,3,1,2]) + x = x.reshape([-1,2,2,self.n_bins,self.dim_t]).reshape([-1,self.dim_c,self.n_bins,self.dim_t]) + return x[:,:,:self.dim_f] + + def istft(self, x, freq_pad=None): + freq_pad = self.freq_pad.repeat([x.shape[0],1,1,1]) if freq_pad is None else freq_pad + x = torch.cat([x, freq_pad], -2) + x = x.reshape([-1,2,2,self.n_bins,self.dim_t]).reshape([-1,2,self.n_bins,self.dim_t]) + x = x.permute([0,2,3,1]) + x=x.contiguous() + x=torch.view_as_complex(x) + x = torch.istft(x, n_fft=self.n_fft, hop_length=self.hop, window=self.window, center=True) + return x.reshape([-1,2,self.chunk_size]) + + def run_model(self, mix, is_ckpt=False, is_match_mix=False): + spek = self.stft(mix)*self.adjust + spek[:, :, :3, :] *= 0 + spec_pred = self.model_run(spek.cpu().float().numpy()) + vocal = self.istft(torch.from_numpy(spec_pred)).transpose(0,1).reshape(2, -1).numpy() + return vocal + +def normalize_two_stem(wave, mix, is_normalize=False): + maxv = np.abs(wave).max() + max_mix = np.abs(mix).max() + if maxv > 1.0: + if is_normalize: + wave /= maxv + mix /= maxv + return wave, mix + +def process(input_file_path, vocal_output_path, instrumental_output_path): + wav = torchaudio.load(input_file_path)[0] + mdx = MDXNetMain() + full_size = ((wav.shape[1] // mdx.chunk_size)+1)*mdx.chunk_size + full_vocal = np.ndarray([2, full_size]) + full_instrumental = np.ndarray([2, full_size]) + for i in range(0, wav.shape[1], mdx.chunk_size): + chunk_wav_pad = np.zeros([2, mdx.chunk_size]) + if i + mdx.chunk_size < wav.shape[1]: + chunk_wav_pad[:2, :mdx.chunk_size] = wav[:2, i:i+mdx.chunk_size].numpy() + else: + chunk_wav_pad[:2, i+mdx.chunk_size:full_size] = wav[:2, i*mdx.chunk_size:full_size].numpy() + chunk_vocal = mdx.run_model(chunk_wav_pad) + chunk_wav_pad, chunk_vocal = normalize_two_stem(chunk_wav_pad, chunk_vocal) + chunk_inst = chunk_wav_pad - chunk_vocal + full_vocal[:, i:i+mdx.chunk_size] = chunk_vocal + full_instrumental[:, i:i+mdx.chunk_size] = chunk_inst + # trim padding + full_vocal = full_vocal[:2, :wav.shape[1]] + full_instrumental = full_instrumental[:2, :wav.shape[1]] + + torchaudio.save(vocal_output_path, torch.from_numpy(full_vocal).float(), 44100) + torchaudio.save(instrumental_output_path, torch.from_numpy(full_instrumental).float(), 44100) + + +if __name__ == "__main__": + process( + "/mnt/nas/media/Music/Days N' Daze/Rogue Taxidemy [REMASTER] [WAV] [LOSSLESS]/06 Blue Jays_1644.wav", + "vocals.wav", + "inst.wav" + ) diff --git a/stem_splitter/mdx_net_main/models/.gitkeep b/stem_splitter/mdx_net_main/models/.gitkeep new file mode 100644 index 0000000..e69de29 From da48f55c9697750b0cfcfd3dbaff11e4aa0722a9 Mon Sep 17 00:00:00 2001 From: Paul Sajna Date: Fri, 12 May 2023 23:03:32 -0700 Subject: [PATCH 2/2] use librosa to load more file types --- stem_splitter/mdx_net_main/infer.py | 21 ++++++++------------- 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/stem_splitter/mdx_net_main/infer.py b/stem_splitter/mdx_net_main/infer.py index 888ffaa..1a4b579 100644 --- a/stem_splitter/mdx_net_main/infer.py +++ b/stem_splitter/mdx_net_main/infer.py @@ -18,7 +18,6 @@ class MDXNetMain(): self.window = torch.hann_window(self.n_fft) self.freq_pad = torch.zeros([1, self.dim_c, self.n_bins-self.dim_f, self.dim_t]) self.adjust = 1 - self.device="cpu" self.ort_ = ort.InferenceSession("models/UVR_MDXNET_Main.onnx") self.model_run = lambda spek:self.ort_.run(None, {'input': spek})[0] @@ -44,7 +43,7 @@ class MDXNetMain(): def run_model(self, mix, is_ckpt=False, is_match_mix=False): spek = self.stft(mix)*self.adjust spek[:, :, :3, :] *= 0 - spec_pred = self.model_run(spek.cpu().float().numpy()) + spec_pred = self.model_run(spek.float().numpy()) vocal = self.istft(torch.from_numpy(spec_pred)).transpose(0,1).reshape(2, -1).numpy() return vocal @@ -58,7 +57,11 @@ def normalize_two_stem(wave, mix, is_normalize=False): return wave, mix def process(input_file_path, vocal_output_path, instrumental_output_path): - wav = torchaudio.load(input_file_path)[0] + wav, _ = librosa.load(input_file_path, sr=44100, mono=False) + if wav.ndim == 1: + wav = np.array([wav, wav]) + if wav.ndim != 2: + raise Exception("Only mono or stereo files are supported") mdx = MDXNetMain() full_size = ((wav.shape[1] // mdx.chunk_size)+1)*mdx.chunk_size full_vocal = np.ndarray([2, full_size]) @@ -66,9 +69,9 @@ def process(input_file_path, vocal_output_path, instrumental_output_path): for i in range(0, wav.shape[1], mdx.chunk_size): chunk_wav_pad = np.zeros([2, mdx.chunk_size]) if i + mdx.chunk_size < wav.shape[1]: - chunk_wav_pad[:2, :mdx.chunk_size] = wav[:2, i:i+mdx.chunk_size].numpy() + chunk_wav_pad[:2, :mdx.chunk_size] = wav[:2, i:i+mdx.chunk_size] else: - chunk_wav_pad[:2, i+mdx.chunk_size:full_size] = wav[:2, i*mdx.chunk_size:full_size].numpy() + chunk_wav_pad[:2, i+mdx.chunk_size:full_size] = wav[:2, i*mdx.chunk_size:full_size] chunk_vocal = mdx.run_model(chunk_wav_pad) chunk_wav_pad, chunk_vocal = normalize_two_stem(chunk_wav_pad, chunk_vocal) chunk_inst = chunk_wav_pad - chunk_vocal @@ -80,11 +83,3 @@ def process(input_file_path, vocal_output_path, instrumental_output_path): torchaudio.save(vocal_output_path, torch.from_numpy(full_vocal).float(), 44100) torchaudio.save(instrumental_output_path, torch.from_numpy(full_instrumental).float(), 44100) - - -if __name__ == "__main__": - process( - "/mnt/nas/media/Music/Days N' Daze/Rogue Taxidemy [REMASTER] [WAV] [LOSSLESS]/06 Blue Jays_1644.wav", - "vocals.wav", - "inst.wav" - )