Merge pull request #14 from storytold/stem-separation

Stem separation

stem_splitter/mdx_net_main/.gitignore

@@ -0,0 +1,2 @@
+*.onnx
+*.wav

stem_splitter/mdx_net_main/infer.py

@@ -0,0 +1,85 @@
+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.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.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, _ = 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])
+    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]
+        else:
+            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
+        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)