From becbee637c4e8828e5044fb67704165d70ba3134 Mon Sep 17 00:00:00 2001 From: Paul Sajna Date: Sat, 13 May 2023 19:33:03 -0700 Subject: [PATCH] add VR Arch, Karaoke model, and dereverb models --- .../.gitignore | 2 + stem_separation/MDX Arch/infer.py | 104 +++ .../MDX Arch}/models/.gitkeep | 0 stem_separation/README.md | 3 + stem_separation/VR Arch/infer.py | 174 +++++ stem_separation/VR Arch/layers_new.py | 126 ++++ stem_separation/VR Arch/model_param_init.py | 59 ++ .../VR Arch/models/data/4band_44100.json | 54 ++ .../VR Arch/models/data/4band_v3.json | 54 ++ stem_separation/VR Arch/nets_new.py | 125 ++++ stem_separation/VR Arch/pyrb.py | 92 +++ stem_separation/VR Arch/spec_utils.py | 693 ++++++++++++++++++ stem_splitter/mdx_net_main/infer.py | 85 --- 13 files changed, 1486 insertions(+), 85 deletions(-) rename {stem_splitter/mdx_net_main => stem_separation}/.gitignore (54%) create mode 100644 stem_separation/MDX Arch/infer.py rename {stem_splitter/mdx_net_main => stem_separation/MDX Arch}/models/.gitkeep (100%) create mode 100644 stem_separation/README.md create mode 100644 stem_separation/VR Arch/infer.py create mode 100644 stem_separation/VR Arch/layers_new.py create mode 100644 stem_separation/VR Arch/model_param_init.py create mode 100644 stem_separation/VR Arch/models/data/4band_44100.json create mode 100644 stem_separation/VR Arch/models/data/4band_v3.json create mode 100644 stem_separation/VR Arch/nets_new.py create mode 100644 stem_separation/VR Arch/pyrb.py create mode 100644 stem_separation/VR Arch/spec_utils.py delete mode 100644 stem_splitter/mdx_net_main/infer.py diff --git a/stem_splitter/mdx_net_main/.gitignore b/stem_separation/.gitignore similarity index 54% rename from stem_splitter/mdx_net_main/.gitignore rename to stem_separation/.gitignore index f2b7995..7470c36 100644 --- a/stem_splitter/mdx_net_main/.gitignore +++ b/stem_separation/.gitignore @@ -1,2 +1,4 @@ *.onnx *.wav +*.pt +*.pth diff --git a/stem_separation/MDX Arch/infer.py b/stem_separation/MDX Arch/infer.py new file mode 100644 index 0000000..49df323 --- /dev/null +++ b/stem_separation/MDX Arch/infer.py @@ -0,0 +1,104 @@ +import numpy as np +import torch +import torchaudio +import onnxruntime as ort +import librosa + +class MDXModel: + + def __init__(self, chunk_size, n_fft, hop, dim_f, dim_t, dim_c, model_file): + self.chunk_size = chunk_size + self.n_fft = n_fft + self.hop = hop + self.dim_f = dim_f + self.dim_t = dim_t + self.dim_c = dim_c + 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/" + model_file) + 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,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()) + stem1 = self.istft(torch.from_numpy(spec_pred)).transpose(0,1).reshape(2, -1).numpy() + return stem1 + + def process(self, input_file_path, stem1_output_path, stem2_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") + full_size = ((wav.shape[1] // self.chunk_size)+1)*self.chunk_size + full_stem1 = np.ndarray([2, full_size]) + full_stem2 = np.ndarray([2, full_size]) + for i in range(0, wav.shape[1], self.chunk_size): + chunk_wav_pad = np.zeros([2, self.chunk_size]) + if i + self.chunk_size < wav.shape[1]: + chunk_wav_pad[:2, :self.chunk_size] = wav[:2, i:i+self.chunk_size] + else: + chunk_wav_pad[:2, i+self.chunk_size:full_size] = wav[:2, i*self.chunk_size:full_size] + chunk_stem1 = self.run_model(chunk_wav_pad) + chunk_wav_pad, chunk_stem1 = normalize_two_stem(chunk_wav_pad, chunk_stem1) + chunk_stem2 = chunk_wav_pad - chunk_stem1 + full_stem1[:, i:i+self.chunk_size] = chunk_stem1 + full_stem2[:, i:i+self.chunk_size] = chunk_stem2 + # trim padding + full_stem1 = full_stem1[:2, :wav.shape[1]] + full_stem2 = full_stem2[:2, :wav.shape[1]] + + torchaudio.save(stem1_output_path, torch.from_numpy(full_stem1).float(), 44100) + torchaudio.save(stem2_output_path, torch.from_numpy(full_stem2).float(), 44100) + +# Not totally sure if this is doing anything useful +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 + + +class MDXNetMain(MDXModel): + def __init__(self): + super(MDXNetMain, self).__init__(261120, 7680, 1024, 3072, 256, 4, "UVR_MDXNET_Main.onnx") + +class MDXNetKara2(MDXModel): + def __init__(self): + super(MDXNetKara2, self).__init__(261120, 7680, 1024, 2048, 256, 4, "UVR_MDXNET_KARA_2.onnx") + +# NB (Paul): This model kind of sucked when I tested it with eye of the tiger +class ReverbHQ(MDXModel): + def __init__(self): + super(ReverbHQ, self).__init__(523264, 6144, 1024, 3072, 512, 4, "Reverb_HQ_By_FoxJoy.onnx") + + +if __name__ == "__main__": + ReverbHQ().process("/mnt/8terra/tts/voice-conversion-sources/eyeofthetiger.wav", "/tmp/stem1.wav", "/tmp/stem2.wav") + #MDXNetMain().process("/tmp/Carly Rae Jepsen - Call Me Maybe [fWNaR-rxAic].webm", "stem1.wav", "stem2.wav") diff --git a/stem_splitter/mdx_net_main/models/.gitkeep b/stem_separation/MDX Arch/models/.gitkeep similarity index 100% rename from stem_splitter/mdx_net_main/models/.gitkeep rename to stem_separation/MDX Arch/models/.gitkeep diff --git a/stem_separation/README.md b/stem_separation/README.md new file mode 100644 index 0000000..51ceebb --- /dev/null +++ b/stem_separation/README.md @@ -0,0 +1,3 @@ + Download the models from https://github.com/TRvlvr/model_repo/releases/tag/all_public_uvr_models + + Instantiate the model class of the model you want to use and run the process function diff --git a/stem_separation/VR Arch/infer.py b/stem_separation/VR Arch/infer.py new file mode 100644 index 0000000..b798198 --- /dev/null +++ b/stem_separation/VR Arch/infer.py @@ -0,0 +1,174 @@ +# stdlib +import math +import os + +# pypi +import torch, torchaudio +import librosa +import numpy as np + +# local +from model_param_init import ModelParameters +import nets_new +import spec_utils + +class VR_Arch: + def __init__(self, model_path, model_param_json): + self.mp = ModelParameters(model_param_json) + self.model_path = model_path + self.model_capacity = 32, 128 + self.high_end_process = 'none' + self.window_size = 512 + self.is_tta = False + self.batch_size = 1 + self.aggressiveness = { + 'value': 10, + 'split_bin': self.mp.param['band'][1]['crop_stop'], + 'aggr_correction': self.mp.param.get('aggr_correction') + } + self.is_post_process = True + self.is_normalization = True + self.post_process_threshold = 0.2 + self.model_samplerate = 44100 + + + def seperate(self, source_path, primary_stem_path, secondary_stem_path): + self.audio_file = source_path + device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') + + nn_arch_sizes = [ + 31191, # default + 33966, 56817, 123821, 123812, 129605, 218409, 537238, 537227] + vr_5_1_models = [56817, 218409] + model_size = math.ceil(os.stat(self.model_path).st_size / 1024) + nn_arch_size = min(nn_arch_sizes, key=lambda x:abs(x-model_size)) + + if nn_arch_size in vr_5_1_models or self.is_vr_51_model: + self.model_run = nets_new.CascadedNet(self.mp.param['bins'] * 2, nn_arch_size, nout=self.model_capacity[0], nout_lstm=self.model_capacity[1]) + else: + self.model_run = nets.determine_model_capacity(self.mp.param['bins'] * 2, nn_arch_size) + self.model_run.load_state_dict(torch.load(self.model_path, map_location='cpu')) + self.model_run.to(device) + + y_spec, v_spec = self.inference_vr(self.loading_mix(), device, self.aggressiveness) + + + self.primary_source = spec_utils.normalize(self.spec_to_wav(y_spec), self.is_normalization).T + if not self.model_samplerate == 44100: + self.primary_source = librosa.resample(self.primary_source.T, orig_sr=self.model_samplerate, target_sr=44100) + print(self.primary_source.shape) + torchaudio.save(primary_stem_path, torch.from_numpy(self.primary_source).float().T, 44100) + + self.secondary_source = self.spec_to_wav(v_spec) + self.secondary_source = spec_utils.normalize(self.spec_to_wav(v_spec), self.is_normalization).T + if not self.model_samplerate == 44100: + self.secondary_source = librosa.resample(self.secondary_source.T, orig_sr=self.model_samplerate, target_sr=44100) + print(self.secondary_source.shape) + torchaudio.save(secondary_stem_path, torch.from_numpy(self.secondary_source).float().T, 44100) + torch.cuda.empty_cache() + + def loading_mix(self): + X_wave, X_spec_s = {}, {} + bands_n = len(self.mp.param['band']) + for d in range(bands_n, 0, -1): + bp = self.mp.param['band'][d] + wav_resolution = bp['res_type'] + if d == bands_n: # high-end band + X_wave[d], _ = librosa.load(self.audio_file, bp['sr'], False, dtype=np.float32, res_type=wav_resolution) + if X_wave[d].ndim == 1: + X_wave[d] = np.asarray([X_wave[d], X_wave[d]]) + else: # lower bands + X_wave[d] = librosa.resample(X_wave[d+1], self.mp.param['band'][d+1]['sr'], bp['sr'], res_type=wav_resolution) + X_spec_s[d] = spec_utils.wave_to_spectrogram_mt(X_wave[d], bp['hl'], bp['n_fft'], self.mp.param['mid_side'], + self.mp.param['mid_side_b2'], self.mp.param['reverse']) + if d == bands_n and self.high_end_process != 'none': + self.input_high_end_h = (bp['n_fft']//2 - bp['crop_stop']) + (self.mp.param['pre_filter_stop'] - self.mp.param['pre_filter_start']) + self.input_high_end = X_spec_s[d][:, bp['n_fft']//2-self.input_high_end_h:bp['n_fft']//2, :] + + X_spec = spec_utils.combine_spectrograms(X_spec_s, self.mp) + del X_wave, X_spec_s + + return X_spec + + def inference_vr(self, X_spec, device, aggressiveness): + def _execute(X_mag_pad, roi_size): + X_dataset = [] + patches = (X_mag_pad.shape[2] - 2 * self.model_run.offset) // roi_size + total_iterations = patches//self.batch_size if not self.is_tta else (patches//self.batch_size)*2 + for i in range(patches): + start = i * roi_size + X_mag_window = X_mag_pad[:, :, start:start + self.window_size] + X_dataset.append(X_mag_window) + + X_dataset = np.asarray(X_dataset) + self.model_run.eval() + with torch.no_grad(): + mask = [] + for i in range(0, patches, self.batch_size): + X_batch = X_dataset[i: i + self.batch_size] + X_batch = torch.from_numpy(X_batch).to(device) + pred = self.model_run.predict_mask(X_batch) + if not pred.size()[3] > 0: + raise Exception(ERROR_MAPPER[WINDOW_SIZE_ERROR]) + pred = pred.detach().cpu().numpy() + pred = np.concatenate(pred, axis=2) + mask.append(pred) + if len(mask) == 0: + raise Exception(ERROR_MAPPER[WINDOW_SIZE_ERROR]) + mask = np.concatenate(mask, axis=2) + return mask + + def postprocess(mask, X_mag, X_phase): + is_non_accom_stem = False + #for stem in NON_ACCOM_STEMS: + # if stem == self.primary_stem: + # is_non_accom_stem = True + mask = spec_utils.adjust_aggr(mask, is_non_accom_stem, aggressiveness) + + if self.is_post_process: + mask = spec_utils.merge_artifacts(mask, thres=self.post_process_threshold) + + y_spec = mask * X_mag * np.exp(1.j * X_phase) + v_spec = (1 - mask) * X_mag * np.exp(1.j * X_phase) + return y_spec, v_spec + + X_mag, X_phase = spec_utils.preprocess(X_spec) + n_frame = X_mag.shape[2] + pad_l, pad_r, roi_size = spec_utils.make_padding(n_frame, self.window_size, self.model_run.offset) + X_mag_pad = np.pad(X_mag, ((0, 0), (0, 0), (pad_l, pad_r)), mode='constant') + X_mag_pad /= X_mag_pad.max() + mask = _execute(X_mag_pad, roi_size) + + if self.is_tta: + pad_l += roi_size // 2 + pad_r += roi_size // 2 + X_mag_pad = np.pad(X_mag, ((0, 0), (0, 0), (pad_l, pad_r)), mode='constant') + X_mag_pad /= X_mag_pad.max() + mask_tta = _execute(X_mag_pad, roi_size) + mask_tta = mask_tta[:, :, roi_size // 2:] + mask = (mask[:, :, :n_frame] + mask_tta[:, :, :n_frame]) * 0.5 + else: + mask = mask[:, :, :n_frame] + + y_spec, v_spec = postprocess(mask, X_mag, X_phase) + return y_spec, v_spec + + def spec_to_wav(self, spec): + if self.high_end_process.startswith('mirroring'): + input_high_end_ = spec_utils.mirroring(self.high_end_process, spec, self.input_high_end, self.mp) + wav = spec_utils.cmb_spectrogram_to_wave(spec, self.mp, self.input_high_end_h, input_high_end_) + else: + wav = spec_utils.cmb_spectrogram_to_wave(spec, self.mp) + return wav + +class UVR_DeEcho_DeReverb(VR_Arch): + def __init__(self): + super(UVR_DeEcho_DeReverb, self).__init__("models/UVR-DeEcho-DeReverb.pth", "models/data/4band_v3.json") + + +if __name__ == "__main__": + UVR_DeEcho_DeReverb().seperate( + "/mnt/8terra/tts/voice-conversion-sources/eyeofthetiger.wav", + "dereverb.wav", + "reverb.wav" + ) diff --git a/stem_separation/VR Arch/layers_new.py b/stem_separation/VR Arch/layers_new.py new file mode 100644 index 0000000..11d48e3 --- /dev/null +++ b/stem_separation/VR Arch/layers_new.py @@ -0,0 +1,126 @@ +import torch +from torch import nn +import torch.nn.functional as F + +import spec_utils + +class Conv2DBNActiv(nn.Module): + + def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU): + super(Conv2DBNActiv, self).__init__() + self.conv = nn.Sequential( + nn.Conv2d( + nin, nout, + kernel_size=ksize, + stride=stride, + padding=pad, + dilation=dilation, + bias=False), + nn.BatchNorm2d(nout), + activ() + ) + + def __call__(self, x): + return self.conv(x) + +class Encoder(nn.Module): + + def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU): + super(Encoder, self).__init__() + self.conv1 = Conv2DBNActiv(nin, nout, ksize, stride, pad, activ=activ) + self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ) + + def __call__(self, x): + h = self.conv1(x) + h = self.conv2(h) + + return h + + +class Decoder(nn.Module): + + def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False): + super(Decoder, self).__init__() + self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ) + # self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ) + self.dropout = nn.Dropout2d(0.1) if dropout else None + + def __call__(self, x, skip=None): + x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True) + + if skip is not None: + skip = spec_utils.crop_center(skip, x) + x = torch.cat([x, skip], dim=1) + + h = self.conv1(x) + # h = self.conv2(h) + + if self.dropout is not None: + h = self.dropout(h) + + return h + + +class ASPPModule(nn.Module): + + def __init__(self, nin, nout, dilations=(4, 8, 12), activ=nn.ReLU, dropout=False): + super(ASPPModule, self).__init__() + self.conv1 = nn.Sequential( + nn.AdaptiveAvgPool2d((1, None)), + Conv2DBNActiv(nin, nout, 1, 1, 0, activ=activ) + ) + self.conv2 = Conv2DBNActiv(nin, nout, 1, 1, 0, activ=activ) + self.conv3 = Conv2DBNActiv( + nin, nout, 3, 1, dilations[0], dilations[0], activ=activ + ) + self.conv4 = Conv2DBNActiv( + nin, nout, 3, 1, dilations[1], dilations[1], activ=activ + ) + self.conv5 = Conv2DBNActiv( + nin, nout, 3, 1, dilations[2], dilations[2], activ=activ + ) + self.bottleneck = Conv2DBNActiv(nout * 5, nout, 1, 1, 0, activ=activ) + self.dropout = nn.Dropout2d(0.1) if dropout else None + + def forward(self, x): + _, _, h, w = x.size() + feat1 = F.interpolate(self.conv1(x), size=(h, w), mode='bilinear', align_corners=True) + feat2 = self.conv2(x) + feat3 = self.conv3(x) + feat4 = self.conv4(x) + feat5 = self.conv5(x) + out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1) + out = self.bottleneck(out) + + if self.dropout is not None: + out = self.dropout(out) + + return out + + +class LSTMModule(nn.Module): + + def __init__(self, nin_conv, nin_lstm, nout_lstm): + super(LSTMModule, self).__init__() + self.conv = Conv2DBNActiv(nin_conv, 1, 1, 1, 0) + self.lstm = nn.LSTM( + input_size=nin_lstm, + hidden_size=nout_lstm // 2, + bidirectional=True + ) + self.dense = nn.Sequential( + nn.Linear(nout_lstm, nin_lstm), + nn.BatchNorm1d(nin_lstm), + nn.ReLU() + ) + + def forward(self, x): + N, _, nbins, nframes = x.size() + h = self.conv(x)[:, 0] # N, nbins, nframes + h = h.permute(2, 0, 1) # nframes, N, nbins + h, _ = self.lstm(h) + h = self.dense(h.reshape(-1, h.size()[-1])) # nframes * N, nbins + h = h.reshape(nframes, N, 1, nbins) + h = h.permute(1, 2, 3, 0) + + return h diff --git a/stem_separation/VR Arch/model_param_init.py b/stem_separation/VR Arch/model_param_init.py new file mode 100644 index 0000000..42fca9a --- /dev/null +++ b/stem_separation/VR Arch/model_param_init.py @@ -0,0 +1,59 @@ +import json +import pathlib + +default_param = {} +default_param['bins'] = 768 +default_param['unstable_bins'] = 9 # training only +default_param['reduction_bins'] = 762 # training only +default_param['sr'] = 44100 +default_param['pre_filter_start'] = 757 +default_param['pre_filter_stop'] = 768 +default_param['band'] = {} + + +default_param['band'][1] = { + 'sr': 11025, + 'hl': 128, + 'n_fft': 960, + 'crop_start': 0, + 'crop_stop': 245, + 'lpf_start': 61, # inference only + 'res_type': 'polyphase' +} + +default_param['band'][2] = { + 'sr': 44100, + 'hl': 512, + 'n_fft': 1536, + 'crop_start': 24, + 'crop_stop': 547, + 'hpf_start': 81, # inference only + 'res_type': 'sinc_best' +} + + +def int_keys(d): + r = {} + for k, v in d: + if k.isdigit(): + k = int(k) + r[k] = v + return r + + +class ModelParameters(object): + def __init__(self, config_path=''): + if '.pth' == pathlib.Path(config_path).suffix: + import zipfile + + with zipfile.ZipFile(config_path, 'r') as zip: + self.param = json.loads(zip.read('param.json'), object_pairs_hook=int_keys) + elif '.json' == pathlib.Path(config_path).suffix: + with open(config_path, 'r') as f: + self.param = json.loads(f.read(), object_pairs_hook=int_keys) + else: + self.param = default_param + + for k in ['mid_side', 'mid_side_b', 'mid_side_b2', 'stereo_w', 'stereo_n', 'reverse']: + if not k in self.param: + self.param[k] = False \ No newline at end of file diff --git a/stem_separation/VR Arch/models/data/4band_44100.json b/stem_separation/VR Arch/models/data/4band_44100.json new file mode 100644 index 0000000..4ae850a --- /dev/null +++ b/stem_separation/VR Arch/models/data/4band_44100.json @@ -0,0 +1,54 @@ +{ + "bins": 768, + "unstable_bins": 7, + "reduction_bins": 668, + "band": { + "1": { + "sr": 11025, + "hl": 128, + "n_fft": 1024, + "crop_start": 0, + "crop_stop": 186, + "lpf_start": 37, + "lpf_stop": 73, + "res_type": "polyphase" + }, + "2": { + "sr": 11025, + "hl": 128, + "n_fft": 512, + "crop_start": 4, + "crop_stop": 185, + "hpf_start": 36, + "hpf_stop": 18, + "lpf_start": 93, + "lpf_stop": 185, + "res_type": "polyphase" + }, + "3": { + "sr": 22050, + "hl": 256, + "n_fft": 512, + "crop_start": 46, + "crop_stop": 186, + "hpf_start": 93, + "hpf_stop": 46, + "lpf_start": 164, + "lpf_stop": 186, + "res_type": "polyphase" + }, + "4": { + "sr": 44100, + "hl": 512, + "n_fft": 768, + "crop_start": 121, + "crop_stop": 382, + "hpf_start": 138, + "hpf_stop": 123, + "res_type": "sinc_medium" + } + }, + "sr": 44100, + "pre_filter_start": 740, + "pre_filter_stop": 768 +} diff --git a/stem_separation/VR Arch/models/data/4band_v3.json b/stem_separation/VR Arch/models/data/4band_v3.json new file mode 100644 index 0000000..2a73bc9 --- /dev/null +++ b/stem_separation/VR Arch/models/data/4band_v3.json @@ -0,0 +1,54 @@ +{ + "bins": 672, + "unstable_bins": 8, + "reduction_bins": 530, + "band": { + "1": { + "sr": 7350, + "hl": 80, + "n_fft": 640, + "crop_start": 0, + "crop_stop": 85, + "lpf_start": 25, + "lpf_stop": 53, + "res_type": "polyphase" + }, + "2": { + "sr": 7350, + "hl": 80, + "n_fft": 320, + "crop_start": 4, + "crop_stop": 87, + "hpf_start": 25, + "hpf_stop": 12, + "lpf_start": 31, + "lpf_stop": 62, + "res_type": "polyphase" + }, + "3": { + "sr": 14700, + "hl": 160, + "n_fft": 512, + "crop_start": 17, + "crop_stop": 216, + "hpf_start": 48, + "hpf_stop": 24, + "lpf_start": 139, + "lpf_stop": 210, + "res_type": "polyphase" + }, + "4": { + "sr": 44100, + "hl": 480, + "n_fft": 960, + "crop_start": 78, + "crop_stop": 383, + "hpf_start": 130, + "hpf_stop": 86, + "res_type": "kaiser_fast" + } + }, + "sr": 44100, + "pre_filter_start": 668, + "pre_filter_stop": 672 +} \ No newline at end of file diff --git a/stem_separation/VR Arch/nets_new.py b/stem_separation/VR Arch/nets_new.py new file mode 100644 index 0000000..e06e2d6 --- /dev/null +++ b/stem_separation/VR Arch/nets_new.py @@ -0,0 +1,125 @@ +import torch +from torch import nn +import torch.nn.functional as F +import layers_new as layers + +class BaseNet(nn.Module): + + def __init__(self, nin, nout, nin_lstm, nout_lstm, dilations=((4, 2), (8, 4), (12, 6))): + super(BaseNet, self).__init__() + self.enc1 = layers.Conv2DBNActiv(nin, nout, 3, 1, 1) + self.enc2 = layers.Encoder(nout, nout * 2, 3, 2, 1) + self.enc3 = layers.Encoder(nout * 2, nout * 4, 3, 2, 1) + self.enc4 = layers.Encoder(nout * 4, nout * 6, 3, 2, 1) + self.enc5 = layers.Encoder(nout * 6, nout * 8, 3, 2, 1) + + self.aspp = layers.ASPPModule(nout * 8, nout * 8, dilations, dropout=True) + + self.dec4 = layers.Decoder(nout * (6 + 8), nout * 6, 3, 1, 1) + self.dec3 = layers.Decoder(nout * (4 + 6), nout * 4, 3, 1, 1) + self.dec2 = layers.Decoder(nout * (2 + 4), nout * 2, 3, 1, 1) + self.lstm_dec2 = layers.LSTMModule(nout * 2, nin_lstm, nout_lstm) + self.dec1 = layers.Decoder(nout * (1 + 2) + 1, nout * 1, 3, 1, 1) + + def __call__(self, x): + e1 = self.enc1(x) + e2 = self.enc2(e1) + e3 = self.enc3(e2) + e4 = self.enc4(e3) + e5 = self.enc5(e4) + + h = self.aspp(e5) + + h = self.dec4(h, e4) + h = self.dec3(h, e3) + h = self.dec2(h, e2) + h = torch.cat([h, self.lstm_dec2(h)], dim=1) + h = self.dec1(h, e1) + + return h + +class CascadedNet(nn.Module): + + def __init__(self, n_fft, nn_arch_size, nout=32, nout_lstm=128): + super(CascadedNet, self).__init__() + + self.max_bin = n_fft // 2 + self.output_bin = n_fft // 2 + 1 + self.nin_lstm = self.max_bin // 2 + self.offset = 64 + nout = 64 if nn_arch_size == 218409 else nout + + self.stg1_low_band_net = nn.Sequential( + BaseNet(2, nout // 2, self.nin_lstm // 2, nout_lstm), + layers.Conv2DBNActiv(nout // 2, nout // 4, 1, 1, 0) + ) + + self.stg1_high_band_net = BaseNet(2, nout // 4, self.nin_lstm // 2, nout_lstm // 2) + + self.stg2_low_band_net = nn.Sequential( + BaseNet(nout // 4 + 2, nout, self.nin_lstm // 2, nout_lstm), + layers.Conv2DBNActiv(nout, nout // 2, 1, 1, 0) + ) + self.stg2_high_band_net = BaseNet(nout // 4 + 2, nout // 2, self.nin_lstm // 2, nout_lstm // 2) + + self.stg3_full_band_net = BaseNet(3 * nout // 4 + 2, nout, self.nin_lstm, nout_lstm) + + self.out = nn.Conv2d(nout, 2, 1, bias=False) + self.aux_out = nn.Conv2d(3 * nout // 4, 2, 1, bias=False) + + def forward(self, x): + x = x[:, :, :self.max_bin] + + bandw = x.size()[2] // 2 + l1_in = x[:, :, :bandw] + h1_in = x[:, :, bandw:] + l1 = self.stg1_low_band_net(l1_in) + h1 = self.stg1_high_band_net(h1_in) + aux1 = torch.cat([l1, h1], dim=2) + + l2_in = torch.cat([l1_in, l1], dim=1) + h2_in = torch.cat([h1_in, h1], dim=1) + l2 = self.stg2_low_band_net(l2_in) + h2 = self.stg2_high_band_net(h2_in) + aux2 = torch.cat([l2, h2], dim=2) + + f3_in = torch.cat([x, aux1, aux2], dim=1) + f3 = self.stg3_full_band_net(f3_in) + + mask = torch.sigmoid(self.out(f3)) + mask = F.pad( + input=mask, + pad=(0, 0, 0, self.output_bin - mask.size()[2]), + mode='replicate' + ) + + if self.training: + aux = torch.cat([aux1, aux2], dim=1) + aux = torch.sigmoid(self.aux_out(aux)) + aux = F.pad( + input=aux, + pad=(0, 0, 0, self.output_bin - aux.size()[2]), + mode='replicate' + ) + return mask, aux + else: + return mask + + def predict_mask(self, x): + mask = self.forward(x) + + if self.offset > 0: + mask = mask[:, :, :, self.offset:-self.offset] + assert mask.size()[3] > 0 + + return mask + + def predict(self, x): + mask = self.forward(x) + pred_mag = x * mask + + if self.offset > 0: + pred_mag = pred_mag[:, :, :, self.offset:-self.offset] + assert pred_mag.size()[3] > 0 + + return pred_mag diff --git a/stem_separation/VR Arch/pyrb.py b/stem_separation/VR Arch/pyrb.py new file mode 100644 index 0000000..883a525 --- /dev/null +++ b/stem_separation/VR Arch/pyrb.py @@ -0,0 +1,92 @@ +import os +import subprocess +import tempfile +import six +import numpy as np +import soundfile as sf +import sys + +if getattr(sys, 'frozen', False): + BASE_PATH_RUB = sys._MEIPASS +else: + BASE_PATH_RUB = os.path.dirname(os.path.abspath(__file__)) + +__all__ = ['time_stretch', 'pitch_shift'] + +__RUBBERBAND_UTIL = os.path.join(BASE_PATH_RUB, 'rubberband') + +if six.PY2: + DEVNULL = open(os.devnull, 'w') +else: + DEVNULL = subprocess.DEVNULL + +def __rubberband(y, sr, **kwargs): + + assert sr > 0 + + # Get the input and output tempfile + fd, infile = tempfile.mkstemp(suffix='.wav') + os.close(fd) + fd, outfile = tempfile.mkstemp(suffix='.wav') + os.close(fd) + + # dump the audio + sf.write(infile, y, sr) + + try: + # Execute rubberband + arguments = [__RUBBERBAND_UTIL, '-q'] + + for key, value in six.iteritems(kwargs): + arguments.append(str(key)) + arguments.append(str(value)) + + arguments.extend([infile, outfile]) + + subprocess.check_call(arguments, stdout=DEVNULL, stderr=DEVNULL) + + # Load the processed audio. + y_out, _ = sf.read(outfile, always_2d=True) + + # make sure that output dimensions matches input + if y.ndim == 1: + y_out = np.squeeze(y_out) + + except OSError as exc: + six.raise_from(RuntimeError('Failed to execute rubberband. ' + 'Please verify that rubberband-cli ' + 'is installed.'), + exc) + + finally: + # Remove temp files + os.unlink(infile) + os.unlink(outfile) + + return y_out + +def time_stretch(y, sr, rate, rbargs=None): + if rate <= 0: + raise ValueError('rate must be strictly positive') + + if rate == 1.0: + return y + + if rbargs is None: + rbargs = dict() + + rbargs.setdefault('--tempo', rate) + + return __rubberband(y, sr, **rbargs) + +def pitch_shift(y, sr, n_steps, rbargs=None): + + if n_steps == 0: + return y + + if rbargs is None: + rbargs = dict() + + rbargs.setdefault('--pitch', n_steps) + + return __rubberband(y, sr, **rbargs) diff --git a/stem_separation/VR Arch/spec_utils.py b/stem_separation/VR Arch/spec_utils.py new file mode 100644 index 0000000..dba7526 --- /dev/null +++ b/stem_separation/VR Arch/spec_utils.py @@ -0,0 +1,693 @@ +import librosa +import numpy as np +import soundfile as sf +import math +import random +import math +import platform +import traceback +import pyrb + +#cur +OPERATING_SYSTEM = platform.system() +SYSTEM_ARCH = platform.platform() +SYSTEM_PROC = platform.processor() +ARM = 'arm' + +if OPERATING_SYSTEM == 'Windows': + from pyrubberband import pyrb +else: + import pyrb + +if OPERATING_SYSTEM == 'Darwin': + wav_resolution = "polyphase" if SYSTEM_PROC == ARM or ARM in SYSTEM_ARCH else "sinc_fastest" +else: + wav_resolution = "sinc_fastest" + +MAX_SPEC = 'Max Spec' +MIN_SPEC = 'Min Spec' +AVERAGE = 'Average' + +def crop_center(h1, h2): + h1_shape = h1.size() + h2_shape = h2.size() + + if h1_shape[3] == h2_shape[3]: + return h1 + elif h1_shape[3] < h2_shape[3]: + raise ValueError('h1_shape[3] must be greater than h2_shape[3]') + + s_time = (h1_shape[3] - h2_shape[3]) // 2 + e_time = s_time + h2_shape[3] + h1 = h1[:, :, :, s_time:e_time] + + return h1 + +def preprocess(X_spec): + X_mag = np.abs(X_spec) + X_phase = np.angle(X_spec) + + return X_mag, X_phase + +def make_padding(width, cropsize, offset): + left = offset + roi_size = cropsize - offset * 2 + if roi_size == 0: + roi_size = cropsize + right = roi_size - (width % roi_size) + left + + return left, right, roi_size + +def wave_to_spectrogram(wave, hop_length, n_fft, mid_side=False, mid_side_b2=False, reverse=False): + if reverse: + wave_left = np.flip(np.asfortranarray(wave[0])) + wave_right = np.flip(np.asfortranarray(wave[1])) + elif mid_side: + wave_left = np.asfortranarray(np.add(wave[0], wave[1]) / 2) + wave_right = np.asfortranarray(np.subtract(wave[0], wave[1])) + elif mid_side_b2: + wave_left = np.asfortranarray(np.add(wave[1], wave[0] * .5)) + wave_right = np.asfortranarray(np.subtract(wave[0], wave[1] * .5)) + else: + wave_left = np.asfortranarray(wave[0]) + wave_right = np.asfortranarray(wave[1]) + + spec_left = librosa.stft(wave_left, n_fft, hop_length=hop_length) + spec_right = librosa.stft(wave_right, n_fft, hop_length=hop_length) + + spec = np.asfortranarray([spec_left, spec_right]) + + return spec + +def wave_to_spectrogram_mt(wave, hop_length, n_fft, mid_side=False, mid_side_b2=False, reverse=False): + import threading + + if reverse: + wave_left = np.flip(np.asfortranarray(wave[0])) + wave_right = np.flip(np.asfortranarray(wave[1])) + elif mid_side: + wave_left = np.asfortranarray(np.add(wave[0], wave[1]) / 2) + wave_right = np.asfortranarray(np.subtract(wave[0], wave[1])) + elif mid_side_b2: + wave_left = np.asfortranarray(np.add(wave[1], wave[0] * .5)) + wave_right = np.asfortranarray(np.subtract(wave[0], wave[1] * .5)) + else: + wave_left = np.asfortranarray(wave[0]) + wave_right = np.asfortranarray(wave[1]) + + def run_thread(**kwargs): + global spec_left + spec_left = librosa.stft(**kwargs) + + thread = threading.Thread(target=run_thread, kwargs={'y': wave_left, 'n_fft': n_fft, 'hop_length': hop_length}) + thread.start() + spec_right = librosa.stft(wave_right, n_fft, hop_length=hop_length) + thread.join() + + spec = np.asfortranarray([spec_left, spec_right]) + + return spec + +def normalize(wave, is_normalize=False): + """Save output music files""" + maxv = np.abs(wave).max() + if maxv > 1.0: + print(f"\nNormalization Set {is_normalize}: Input above threshold for clipping. Max:{maxv}") + if is_normalize: + print(f"The result was normalized.") + wave /= maxv + else: + print(f"The result was not normalized.") + else: + print(f"\nNormalization Set {is_normalize}: Input not above threshold for clipping. Max:{maxv}") + + return wave + +def normalize_two_stem(wave, mix, is_normalize=False): + """Save output music files""" + + maxv = np.abs(wave).max() + max_mix = np.abs(mix).max() + + if maxv > 1.0: + print(f"\nNormalization Set {is_normalize}: Primary source above threshold for clipping. Max:{maxv}") + print(f"\nNormalization Set {is_normalize}: Mixture above threshold for clipping. Max:{max_mix}") + if is_normalize: + print(f"The result was normalized.") + wave /= maxv + mix /= maxv + else: + print(f"The result was not normalized.") + else: + print(f"\nNormalization Set {is_normalize}: Input not above threshold for clipping. Max:{maxv}") + + + print(f"\nNormalization Set {is_normalize}: Primary source - Max:{np.abs(wave).max()}") + print(f"\nNormalization Set {is_normalize}: Mixture - Max:{np.abs(mix).max()}") + + return wave, mix + +def combine_spectrograms(specs, mp): + l = min([specs[i].shape[2] for i in specs]) + spec_c = np.zeros(shape=(2, mp.param['bins'] + 1, l), dtype=np.complex64) + offset = 0 + bands_n = len(mp.param['band']) + + for d in range(1, bands_n + 1): + h = mp.param['band'][d]['crop_stop'] - mp.param['band'][d]['crop_start'] + spec_c[:, offset:offset+h, :l] = specs[d][:, mp.param['band'][d]['crop_start']:mp.param['band'][d]['crop_stop'], :l] + offset += h + + if offset > mp.param['bins']: + raise ValueError('Too much bins') + + # lowpass fiter + if mp.param['pre_filter_start'] > 0: # and mp.param['band'][bands_n]['res_type'] in ['scipy', 'polyphase']: + if bands_n == 1: + spec_c = fft_lp_filter(spec_c, mp.param['pre_filter_start'], mp.param['pre_filter_stop']) + else: + gp = 1 + for b in range(mp.param['pre_filter_start'] + 1, mp.param['pre_filter_stop']): + g = math.pow(10, -(b - mp.param['pre_filter_start']) * (3.5 - gp) / 20.0) + gp = g + spec_c[:, b, :] *= g + + return np.asfortranarray(spec_c) + +def spectrogram_to_image(spec, mode='magnitude'): + if mode == 'magnitude': + if np.iscomplexobj(spec): + y = np.abs(spec) + else: + y = spec + y = np.log10(y ** 2 + 1e-8) + elif mode == 'phase': + if np.iscomplexobj(spec): + y = np.angle(spec) + else: + y = spec + + y -= y.min() + y *= 255 / y.max() + img = np.uint8(y) + + if y.ndim == 3: + img = img.transpose(1, 2, 0) + img = np.concatenate([ + np.max(img, axis=2, keepdims=True), img + ], axis=2) + + return img + +def reduce_vocal_aggressively(X, y, softmask): + v = X - y + y_mag_tmp = np.abs(y) + v_mag_tmp = np.abs(v) + + v_mask = v_mag_tmp > y_mag_tmp + y_mag = np.clip(y_mag_tmp - v_mag_tmp * v_mask * softmask, 0, np.inf) + + return y_mag * np.exp(1.j * np.angle(y)) + +def merge_artifacts(y_mask, thres=0.01, min_range=64, fade_size=32): + mask = y_mask + + try: + if min_range < fade_size * 2: + raise ValueError('min_range must be >= fade_size * 2') + + idx = np.where(y_mask.min(axis=(0, 1)) > thres)[0] + start_idx = np.insert(idx[np.where(np.diff(idx) != 1)[0] + 1], 0, idx[0]) + end_idx = np.append(idx[np.where(np.diff(idx) != 1)[0]], idx[-1]) + artifact_idx = np.where(end_idx - start_idx > min_range)[0] + weight = np.zeros_like(y_mask) + if len(artifact_idx) > 0: + start_idx = start_idx[artifact_idx] + end_idx = end_idx[artifact_idx] + old_e = None + for s, e in zip(start_idx, end_idx): + if old_e is not None and s - old_e < fade_size: + s = old_e - fade_size * 2 + + if s != 0: + weight[:, :, s:s + fade_size] = np.linspace(0, 1, fade_size) + else: + s -= fade_size + + if e != y_mask.shape[2]: + weight[:, :, e - fade_size:e] = np.linspace(1, 0, fade_size) + else: + e += fade_size + + weight[:, :, s + fade_size:e - fade_size] = 1 + old_e = e + + v_mask = 1 - y_mask + y_mask += weight * v_mask + + mask = y_mask + except Exception as e: + error_name = f'{type(e).__name__}' + traceback_text = ''.join(traceback.format_tb(e.__traceback__)) + message = f'{error_name}: "{e}"\n{traceback_text}"' + print('Post Process Failed: ', message) + + + return mask + +def align_wave_head_and_tail(a, b): + l = min([a[0].size, b[0].size]) + + return a[:l,:l], b[:l,:l] + +def spectrogram_to_wave(spec, hop_length, mid_side, mid_side_b2, reverse, clamp=False): + spec_left = np.asfortranarray(spec[0]) + spec_right = np.asfortranarray(spec[1]) + + wave_left = librosa.istft(spec_left, hop_length=hop_length) + wave_right = librosa.istft(spec_right, hop_length=hop_length) + + if reverse: + return np.asfortranarray([np.flip(wave_left), np.flip(wave_right)]) + elif mid_side: + return np.asfortranarray([np.add(wave_left, wave_right / 2), np.subtract(wave_left, wave_right / 2)]) + elif mid_side_b2: + return np.asfortranarray([np.add(wave_right / 1.25, .4 * wave_left), np.subtract(wave_left / 1.25, .4 * wave_right)]) + else: + return np.asfortranarray([wave_left, wave_right]) + +def spectrogram_to_wave_mt(spec, hop_length, mid_side, reverse, mid_side_b2): + import threading + + spec_left = np.asfortranarray(spec[0]) + spec_right = np.asfortranarray(spec[1]) + + def run_thread(**kwargs): + global wave_left + wave_left = librosa.istft(**kwargs) + + thread = threading.Thread(target=run_thread, kwargs={'stft_matrix': spec_left, 'hop_length': hop_length}) + thread.start() + wave_right = librosa.istft(spec_right, hop_length=hop_length) + thread.join() + + if reverse: + return np.asfortranarray([np.flip(wave_left), np.flip(wave_right)]) + elif mid_side: + return np.asfortranarray([np.add(wave_left, wave_right / 2), np.subtract(wave_left, wave_right / 2)]) + elif mid_side_b2: + return np.asfortranarray([np.add(wave_right / 1.25, .4 * wave_left), np.subtract(wave_left / 1.25, .4 * wave_right)]) + else: + return np.asfortranarray([wave_left, wave_right]) + +def cmb_spectrogram_to_wave(spec_m, mp, extra_bins_h=None, extra_bins=None): + bands_n = len(mp.param['band']) + offset = 0 + + for d in range(1, bands_n + 1): + bp = mp.param['band'][d] + spec_s = np.ndarray(shape=(2, bp['n_fft'] // 2 + 1, spec_m.shape[2]), dtype=complex) + h = bp['crop_stop'] - bp['crop_start'] + spec_s[:, bp['crop_start']:bp['crop_stop'], :] = spec_m[:, offset:offset+h, :] + + offset += h + if d == bands_n: # higher + if extra_bins_h: # if --high_end_process bypass + max_bin = bp['n_fft'] // 2 + spec_s[:, max_bin-extra_bins_h:max_bin, :] = extra_bins[:, :extra_bins_h, :] + if bp['hpf_start'] > 0: + spec_s = fft_hp_filter(spec_s, bp['hpf_start'], bp['hpf_stop'] - 1) + if bands_n == 1: + wave = spectrogram_to_wave(spec_s, bp['hl'], mp.param['mid_side'], mp.param['mid_side_b2'], mp.param['reverse']) + else: + wave = np.add(wave, spectrogram_to_wave(spec_s, bp['hl'], mp.param['mid_side'], mp.param['mid_side_b2'], mp.param['reverse'])) + else: + sr = mp.param['band'][d+1]['sr'] + if d == 1: # lower + spec_s = fft_lp_filter(spec_s, bp['lpf_start'], bp['lpf_stop']) + wave = librosa.resample(spectrogram_to_wave(spec_s, bp['hl'], mp.param['mid_side'], mp.param['mid_side_b2'], mp.param['reverse']), bp['sr'], sr, res_type=wav_resolution) + else: # mid + spec_s = fft_hp_filter(spec_s, bp['hpf_start'], bp['hpf_stop'] - 1) + spec_s = fft_lp_filter(spec_s, bp['lpf_start'], bp['lpf_stop']) + wave2 = np.add(wave, spectrogram_to_wave(spec_s, bp['hl'], mp.param['mid_side'], mp.param['mid_side_b2'], mp.param['reverse'])) + wave = librosa.resample(wave2, bp['sr'], sr, res_type=wav_resolution) + + return wave + +def fft_lp_filter(spec, bin_start, bin_stop): + g = 1.0 + for b in range(bin_start, bin_stop): + g -= 1 / (bin_stop - bin_start) + spec[:, b, :] = g * spec[:, b, :] + + spec[:, bin_stop:, :] *= 0 + + return spec + +def fft_hp_filter(spec, bin_start, bin_stop): + g = 1.0 + for b in range(bin_start, bin_stop, -1): + g -= 1 / (bin_start - bin_stop) + spec[:, b, :] = g * spec[:, b, :] + + spec[:, 0:bin_stop+1, :] *= 0 + + return spec + +def mirroring(a, spec_m, input_high_end, mp): + if 'mirroring' == a: + mirror = np.flip(np.abs(spec_m[:, mp.param['pre_filter_start']-10-input_high_end.shape[1]:mp.param['pre_filter_start']-10, :]), 1) + mirror = mirror * np.exp(1.j * np.angle(input_high_end)) + + return np.where(np.abs(input_high_end) <= np.abs(mirror), input_high_end, mirror) + + if 'mirroring2' == a: + mirror = np.flip(np.abs(spec_m[:, mp.param['pre_filter_start']-10-input_high_end.shape[1]:mp.param['pre_filter_start']-10, :]), 1) + mi = np.multiply(mirror, input_high_end * 1.7) + + return np.where(np.abs(input_high_end) <= np.abs(mi), input_high_end, mi) + +def adjust_aggr(mask, is_non_accom_stem, aggressiveness): + aggr = aggressiveness['value'] + + if aggr != 0: + if is_non_accom_stem: + aggr = 1 - aggr + + aggr = [aggr, aggr] + + if aggressiveness['aggr_correction'] is not None: + aggr[0] += aggressiveness['aggr_correction']['left'] + aggr[1] += aggressiveness['aggr_correction']['right'] + + for ch in range(2): + mask[ch, :aggressiveness['split_bin']] = np.power(mask[ch, :aggressiveness['split_bin']], 1 + aggr[ch] / 3) + mask[ch, aggressiveness['split_bin']:] = np.power(mask[ch, aggressiveness['split_bin']:], 1 + aggr[ch]) + + # if is_non_accom_stem: + # mask = (1.0 - mask) + + return mask + +def stft(wave, nfft, hl): + wave_left = np.asfortranarray(wave[0]) + wave_right = np.asfortranarray(wave[1]) + spec_left = librosa.stft(wave_left, nfft, hop_length=hl) + spec_right = librosa.stft(wave_right, nfft, hop_length=hl) + spec = np.asfortranarray([spec_left, spec_right]) + + return spec + +def istft(spec, hl): + spec_left = np.asfortranarray(spec[0]) + spec_right = np.asfortranarray(spec[1]) + wave_left = librosa.istft(spec_left, hop_length=hl) + wave_right = librosa.istft(spec_right, hop_length=hl) + wave = np.asfortranarray([wave_left, wave_right]) + + return wave + +def spec_effects(wave, algorithm='Default', value=None): + spec = [stft(wave[0],2048,1024), stft(wave[1],2048,1024)] + if algorithm == 'Min_Mag': + v_spec_m = np.where(np.abs(spec[1]) <= np.abs(spec[0]), spec[1], spec[0]) + wave = istft(v_spec_m,1024) + elif algorithm == 'Max_Mag': + v_spec_m = np.where(np.abs(spec[1]) >= np.abs(spec[0]), spec[1], spec[0]) + wave = istft(v_spec_m,1024) + elif algorithm == 'Default': + wave = (wave[1] * value) + (wave[0] * (1-value)) + elif algorithm == 'Invert_p': + X_mag = np.abs(spec[0]) + y_mag = np.abs(spec[1]) + max_mag = np.where(X_mag >= y_mag, X_mag, y_mag) + v_spec = spec[1] - max_mag * np.exp(1.j * np.angle(spec[0])) + wave = istft(v_spec,1024) + + return wave + +def spectrogram_to_wave_no_mp(spec, n_fft=2048, hop_length=1024): + wave = librosa.istft(spec, n_fft=n_fft, hop_length=hop_length) + + if wave.ndim == 1: + wave = np.asfortranarray([wave,wave]) + + return wave + +def wave_to_spectrogram_no_mp(wave): + + spec = librosa.stft(wave, n_fft=2048, hop_length=1024) + + if spec.ndim == 1: + spec = np.asfortranarray([spec,spec]) + + return spec + +def invert_audio(specs, invert_p=True): + + ln = min([specs[0].shape[2], specs[1].shape[2]]) + specs[0] = specs[0][:,:,:ln] + specs[1] = specs[1][:,:,:ln] + + if invert_p: + X_mag = np.abs(specs[0]) + y_mag = np.abs(specs[1]) + max_mag = np.where(X_mag >= y_mag, X_mag, y_mag) + v_spec = specs[1] - max_mag * np.exp(1.j * np.angle(specs[0])) + else: + specs[1] = reduce_vocal_aggressively(specs[0], specs[1], 0.2) + v_spec = specs[0] - specs[1] + + return v_spec + +def invert_stem(mixture, stem): + + mixture = wave_to_spectrogram_no_mp(mixture) + stem = wave_to_spectrogram_no_mp(stem) + output = spectrogram_to_wave_no_mp(invert_audio([mixture, stem])) + + return -output.T + +def ensembling(a, specs): + for i in range(1, len(specs)): + if i == 1: + spec = specs[0] + + ln = min([spec.shape[2], specs[i].shape[2]]) + spec = spec[:,:,:ln] + specs[i] = specs[i][:,:,:ln] + + if MIN_SPEC == a: + spec = np.where(np.abs(specs[i]) <= np.abs(spec), specs[i], spec) + if MAX_SPEC == a: + spec = np.where(np.abs(specs[i]) >= np.abs(spec), specs[i], spec) + if AVERAGE == a: + spec = np.where(np.abs(specs[i]) == np.abs(spec), specs[i], spec) + + return spec + +def ensemble_inputs(audio_input, algorithm, is_normalization, wav_type_set, save_path): + + wavs_ = [] + + if algorithm == AVERAGE: + output = average_audio(audio_input) + samplerate = 44100 + else: + specs = [] + + for i in range(len(audio_input)): + wave, samplerate = librosa.load(audio_input[i], mono=False, sr=44100) + wavs_.append(wave) + spec = wave_to_spectrogram_no_mp(wave) + specs.append(spec) + + wave_shapes = [w.shape[1] for w in wavs_] + target_shape = wavs_[wave_shapes.index(max(wave_shapes))] + + output = spectrogram_to_wave_no_mp(ensembling(algorithm, specs)) + output = to_shape(output, target_shape.shape) + + sf.write(save_path, normalize(output.T, is_normalization), samplerate, subtype=wav_type_set) + +def to_shape(x, target_shape): + padding_list = [] + for x_dim, target_dim in zip(x.shape, target_shape): + pad_value = (target_dim - x_dim) + pad_tuple = ((0, pad_value)) + padding_list.append(pad_tuple) + + return np.pad(x, tuple(padding_list), mode='constant') + +def to_shape_minimize(x: np.ndarray, target_shape): + + padding_list = [] + for x_dim, target_dim in zip(x.shape, target_shape): + pad_value = (target_dim - x_dim) + pad_tuple = ((0, pad_value)) + padding_list.append(pad_tuple) + + return np.pad(x, tuple(padding_list), mode='constant') + +def augment_audio(export_path, audio_file, rate, is_normalization, wav_type_set, save_format=None, is_pitch=False): + + wav, sr = librosa.load(audio_file, sr=44100, mono=False) + + if wav.ndim == 1: + wav = np.asfortranarray([wav,wav]) + + if is_pitch: + wav_1 = pyrb.pitch_shift(wav[0], sr, rate, rbargs=None) + wav_2 = pyrb.pitch_shift(wav[1], sr, rate, rbargs=None) + else: + wav_1 = pyrb.time_stretch(wav[0], sr, rate, rbargs=None) + wav_2 = pyrb.time_stretch(wav[1], sr, rate, rbargs=None) + + if wav_1.shape > wav_2.shape: + wav_2 = to_shape(wav_2, wav_1.shape) + if wav_1.shape < wav_2.shape: + wav_1 = to_shape(wav_1, wav_2.shape) + + wav_mix = np.asfortranarray([wav_1, wav_2]) + + sf.write(export_path, normalize(wav_mix.T, is_normalization), sr, subtype=wav_type_set) + save_format(export_path) + +def average_audio(audio): + + waves = [] + wave_shapes = [] + final_waves = [] + + for i in range(len(audio)): + wave = librosa.load(audio[i], sr=44100, mono=False) + waves.append(wave[0]) + wave_shapes.append(wave[0].shape[1]) + + wave_shapes_index = wave_shapes.index(max(wave_shapes)) + target_shape = waves[wave_shapes_index] + waves.pop(wave_shapes_index) + final_waves.append(target_shape) + + for n_array in waves: + wav_target = to_shape(n_array, target_shape.shape) + final_waves.append(wav_target) + + waves = sum(final_waves) + waves = waves/len(audio) + + return waves + +def average_dual_sources(wav_1, wav_2, value): + + if wav_1.shape > wav_2.shape: + wav_2 = to_shape(wav_2, wav_1.shape) + if wav_1.shape < wav_2.shape: + wav_1 = to_shape(wav_1, wav_2.shape) + + wave = (wav_1 * value) + (wav_2 * (1-value)) + + return wave + +def reshape_sources(wav_1: np.ndarray, wav_2: np.ndarray): + + if wav_1.shape > wav_2.shape: + wav_2 = to_shape(wav_2, wav_1.shape) + if wav_1.shape < wav_2.shape: + ln = min([wav_1.shape[1], wav_2.shape[1]]) + wav_2 = wav_2[:,:ln] + + ln = min([wav_1.shape[1], wav_2.shape[1]]) + wav_1 = wav_1[:,:ln] + wav_2 = wav_2[:,:ln] + + return wav_2 + +def align_audio(file1, file2, file2_aligned, file_subtracted, wav_type_set, is_normalization, command_Text, progress_bar_main_var, save_format): + def get_diff(a, b): + corr = np.correlate(a, b, "full") + diff = corr.argmax() - (b.shape[0] - 1) + return diff + + progress_bar_main_var.set(10) + + # read tracks + wav1, sr1 = librosa.load(file1, sr=44100, mono=False) + wav2, sr2 = librosa.load(file2, sr=44100, mono=False) + wav1 = wav1.transpose() + wav2 = wav2.transpose() + + command_Text(f"Audio file shapes: {wav1.shape} / {wav2.shape}\n") + + wav2_org = wav2.copy() + progress_bar_main_var.set(20) + + command_Text("Processing files... \n") + + # pick random position and get diff + + counts = {} # counting up for each diff value + progress = 20 + + check_range = 64 + + base = (64 / check_range) + + for i in range(check_range): + index = int(random.uniform(44100 * 2, min(wav1.shape[0], wav2.shape[0]) - 44100 * 2)) + shift = int(random.uniform(-22050,+22050)) + samp1 = wav1[index :index +44100, 0] # currently use left channel + samp2 = wav2[index+shift:index+shift+44100, 0] + progress += 1 * base + progress_bar_main_var.set(progress) + diff = get_diff(samp1, samp2) + diff -= shift + + if abs(diff) < 22050: + if not diff in counts: + counts[diff] = 0 + counts[diff] += 1 + + # use max counted diff value + max_count = 0 + est_diff = 0 + for diff in counts.keys(): + if counts[diff] > max_count: + max_count = counts[diff] + est_diff = diff + + command_Text(f"Estimated difference is {est_diff} (count: {max_count})\n") + + progress_bar_main_var.set(90) + + audio_files = [] + + def save_aligned_audio(wav2_aligned): + command_Text(f"Aligned File 2 with File 1.\n") + command_Text(f"Saving files... ") + sf.write(file2_aligned, normalize(wav2_aligned, is_normalization), sr2, subtype=wav_type_set) + save_format(file2_aligned) + min_len = min(wav1.shape[0], wav2_aligned.shape[0]) + wav_sub = wav1[:min_len] - wav2_aligned[:min_len] + audio_files.append(file2_aligned) + return min_len, wav_sub + + # make aligned track 2 + if est_diff > 0: + wav2_aligned = np.append(np.zeros((est_diff, 2)), wav2_org, axis=0) + min_len, wav_sub = save_aligned_audio(wav2_aligned) + elif est_diff < 0: + wav2_aligned = wav2_org[-est_diff:] + min_len, wav_sub = save_aligned_audio(wav2_aligned) + else: + command_Text(f"Audio files already aligned.\n") + command_Text(f"Saving inverted track... ") + min_len = min(wav1.shape[0], wav2.shape[0]) + wav_sub = wav1[:min_len] - wav2[:min_len] + + wav_sub = np.clip(wav_sub, -1, +1) + + sf.write(file_subtracted, normalize(wav_sub, is_normalization), sr1, subtype=wav_type_set) + save_format(file_subtracted) + + progress_bar_main_var.set(95) diff --git a/stem_splitter/mdx_net_main/infer.py b/stem_splitter/mdx_net_main/infer.py deleted file mode 100644 index 1a4b579..0000000 --- a/stem_splitter/mdx_net_main/infer.py +++ /dev/null @@ -1,85 +0,0 @@ -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)