| import os
|
| import logging
|
|
|
| logger = logging.getLogger(__name__)
|
|
|
| import librosa
|
| import numpy as np
|
| import soundfile as sf
|
| import torch
|
| from tqdm import tqdm
|
|
|
| cpu = torch.device("cpu")
|
|
|
|
|
| class ConvTDFNetTrim:
|
| def __init__(
|
| self, device, model_name, target_name, L, dim_f, dim_t, n_fft, hop=1024
|
| ):
|
| super(ConvTDFNetTrim, self).__init__()
|
|
|
| self.dim_f = dim_f
|
| self.dim_t = 2**dim_t
|
| self.n_fft = n_fft
|
| self.hop = hop
|
| self.n_bins = self.n_fft // 2 + 1
|
| self.chunk_size = hop * (self.dim_t - 1)
|
| self.window = torch.hann_window(window_length=self.n_fft, periodic=True).to(
|
| device
|
| )
|
| self.target_name = target_name
|
| self.blender = "blender" in model_name
|
|
|
| self.dim_c = 4
|
| out_c = self.dim_c * 4 if target_name == "*" else self.dim_c
|
| self.freq_pad = torch.zeros(
|
| [1, out_c, self.n_bins - self.dim_f, self.dim_t]
|
| ).to(device)
|
|
|
| self.n = L // 2
|
|
|
| def stft(self, x):
|
| x = x.reshape([-1, self.chunk_size])
|
| x = torch.stft(
|
| 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)
|
| c = 4 * 2 if self.target_name == "*" else 2
|
| x = x.reshape([-1, c, 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, c, self.chunk_size])
|
|
|
|
|
| def get_models(device, dim_f, dim_t, n_fft):
|
| return ConvTDFNetTrim(
|
| device=device,
|
| model_name="Conv-TDF",
|
| target_name="vocals",
|
| L=11,
|
| dim_f=dim_f,
|
| dim_t=dim_t,
|
| n_fft=n_fft,
|
| )
|
|
|
|
|
| class Predictor:
|
| def __init__(self, args):
|
| import onnxruntime as ort
|
|
|
| logger.info(ort.get_available_providers())
|
| self.args = args
|
| self.model_ = get_models(
|
| device=cpu, dim_f=args.dim_f, dim_t=args.dim_t, n_fft=args.n_fft
|
| )
|
| self.model = ort.InferenceSession(
|
| os.path.join(args.onnx, self.model_.target_name + ".onnx"),
|
| providers=[
|
| "CUDAExecutionProvider",
|
| "DmlExecutionProvider",
|
| "CPUExecutionProvider",
|
| ],
|
| )
|
| logger.info("ONNX load done")
|
|
|
| def demix(self, mix):
|
| samples = mix.shape[-1]
|
| margin = self.args.margin
|
| chunk_size = self.args.chunks * 44100
|
| assert not margin == 0, "margin cannot be zero!"
|
| if margin > chunk_size:
|
| margin = chunk_size
|
|
|
| segmented_mix = {}
|
|
|
| if self.args.chunks == 0 or samples < chunk_size:
|
| chunk_size = samples
|
|
|
| counter = -1
|
| for skip in range(0, samples, chunk_size):
|
| counter += 1
|
|
|
| s_margin = 0 if counter == 0 else margin
|
| end = min(skip + chunk_size + margin, samples)
|
|
|
| start = skip - s_margin
|
|
|
| segmented_mix[skip] = mix[:, start:end].copy()
|
| if end == samples:
|
| break
|
|
|
| sources = self.demix_base(segmented_mix, margin_size=margin)
|
| """
|
| mix:(2,big_sample)
|
| segmented_mix:offset->(2,small_sample)
|
| sources:(1,2,big_sample)
|
| """
|
| return sources
|
|
|
| def demix_base(self, mixes, margin_size):
|
| chunked_sources = []
|
| progress_bar = tqdm(total=len(mixes))
|
| progress_bar.set_description("Processing")
|
| for mix in mixes:
|
| cmix = mixes[mix]
|
| sources = []
|
| n_sample = cmix.shape[1]
|
| model = self.model_
|
| trim = model.n_fft // 2
|
| gen_size = model.chunk_size - 2 * trim
|
| pad = gen_size - n_sample % gen_size
|
| mix_p = np.concatenate(
|
| (np.zeros((2, trim)), cmix, np.zeros((2, pad)), np.zeros((2, trim))), 1
|
| )
|
| mix_waves = []
|
| i = 0
|
| while i < n_sample + pad:
|
| waves = np.array(mix_p[:, i : i + model.chunk_size])
|
| mix_waves.append(waves)
|
| i += gen_size
|
| mix_waves = torch.tensor(mix_waves, dtype=torch.float32).to(cpu)
|
| with torch.no_grad():
|
| _ort = self.model
|
| spek = model.stft(mix_waves)
|
| if self.args.denoise:
|
| spec_pred = (
|
| -_ort.run(None, {"input": -spek.cpu().numpy()})[0] * 0.5
|
| + _ort.run(None, {"input": spek.cpu().numpy()})[0] * 0.5
|
| )
|
| tar_waves = model.istft(torch.tensor(spec_pred))
|
| else:
|
| tar_waves = model.istft(
|
| torch.tensor(_ort.run(None, {"input": spek.cpu().numpy()})[0])
|
| )
|
| tar_signal = (
|
| tar_waves[:, :, trim:-trim]
|
| .transpose(0, 1)
|
| .reshape(2, -1)
|
| .numpy()[:, :-pad]
|
| )
|
|
|
| start = 0 if mix == 0 else margin_size
|
| end = None if mix == list(mixes.keys())[::-1][0] else -margin_size
|
| if margin_size == 0:
|
| end = None
|
| sources.append(tar_signal[:, start:end])
|
|
|
| progress_bar.update(1)
|
|
|
| chunked_sources.append(sources)
|
| _sources = np.concatenate(chunked_sources, axis=-1)
|
|
|
| progress_bar.close()
|
| return _sources
|
|
|
| def prediction(self, m, vocal_root, others_root, format):
|
| os.makedirs(vocal_root, exist_ok=True)
|
| os.makedirs(others_root, exist_ok=True)
|
| basename = os.path.basename(m)
|
| mix, rate = librosa.load(m, mono=False, sr=44100)
|
| if mix.ndim == 1:
|
| mix = np.asfortranarray([mix, mix])
|
| mix = mix.T
|
| sources = self.demix(mix.T)
|
| opt = sources[0].T
|
| if format in ["wav", "flac"]:
|
| sf.write(
|
| "%s/%s_main_vocal.%s" % (vocal_root, basename, format), mix - opt, rate
|
| )
|
| sf.write("%s/%s_others.%s" % (others_root, basename, format), opt, rate)
|
| else:
|
| path_vocal = "%s/%s_main_vocal.wav" % (vocal_root, basename)
|
| path_other = "%s/%s_others.wav" % (others_root, basename)
|
| sf.write(path_vocal, mix - opt, rate)
|
| sf.write(path_other, opt, rate)
|
| opt_path_vocal = path_vocal[:-4] + ".%s" % format
|
| opt_path_other = path_other[:-4] + ".%s" % format
|
| if os.path.exists(path_vocal):
|
| os.system(
|
| 'ffmpeg -i "%s" -vn "%s" -q:a 2 -y' % (path_vocal, opt_path_vocal)
|
| )
|
| if os.path.exists(opt_path_vocal):
|
| try:
|
| os.remove(path_vocal)
|
| except:
|
| pass
|
| if os.path.exists(path_other):
|
| os.system(
|
| 'ffmpeg -i "%s" -vn "%s" -q:a 2 -y' % (path_other, opt_path_other)
|
| )
|
| if os.path.exists(opt_path_other):
|
| try:
|
| os.remove(path_other)
|
| except:
|
| pass
|
|
|
|
|
| class MDXNetDereverb:
|
| def __init__(self, chunks, device):
|
| self.onnx = "assets/uvr5_weights/onnx_dereverb_By_FoxJoy"
|
| self.shifts = 10
|
| self.mixing = "min_mag"
|
| self.chunks = chunks
|
| self.margin = 44100
|
| self.dim_t = 9
|
| self.dim_f = 3072
|
| self.n_fft = 6144
|
| self.denoise = True
|
| self.pred = Predictor(self)
|
| self.device = device
|
|
|
| def _path_audio_(self, input, vocal_root, others_root, format, is_hp3=False):
|
| self.pred.prediction(input, vocal_root, others_root, format)
|
|
|
