|
| 1 | +import soundfile as sf |
| 2 | +import torch, pdb, os, warnings, librosa |
| 3 | +import numpy as np |
| 4 | +import onnxruntime as ort |
| 5 | +from tqdm import tqdm |
| 6 | +import torch |
| 7 | + |
| 8 | +dim_c = 4 |
| 9 | + |
| 10 | + |
| 11 | +class Conv_TDF_net_trim: |
| 12 | + def __init__( |
| 13 | + self, device, model_name, target_name, L, dim_f, dim_t, n_fft, hop=1024 |
| 14 | + ): |
| 15 | + super(Conv_TDF_net_trim, self).__init__() |
| 16 | + |
| 17 | + self.dim_f = dim_f |
| 18 | + self.dim_t = 2**dim_t |
| 19 | + self.n_fft = n_fft |
| 20 | + self.hop = hop |
| 21 | + self.n_bins = self.n_fft // 2 + 1 |
| 22 | + self.chunk_size = hop * (self.dim_t - 1) |
| 23 | + self.window = torch.hann_window(window_length=self.n_fft, periodic=True).to( |
| 24 | + device |
| 25 | + ) |
| 26 | + self.target_name = target_name |
| 27 | + self.blender = "blender" in model_name |
| 28 | + |
| 29 | + out_c = dim_c * 4 if target_name == "*" else dim_c |
| 30 | + self.freq_pad = torch.zeros( |
| 31 | + [1, out_c, self.n_bins - self.dim_f, self.dim_t] |
| 32 | + ).to(device) |
| 33 | + |
| 34 | + self.n = L // 2 |
| 35 | + |
| 36 | + def stft(self, x): |
| 37 | + x = x.reshape([-1, self.chunk_size]) |
| 38 | + x = torch.stft( |
| 39 | + x, |
| 40 | + n_fft=self.n_fft, |
| 41 | + hop_length=self.hop, |
| 42 | + window=self.window, |
| 43 | + center=True, |
| 44 | + return_complex=True, |
| 45 | + ) |
| 46 | + x = torch.view_as_real(x) |
| 47 | + x = x.permute([0, 3, 1, 2]) |
| 48 | + x = x.reshape([-1, 2, 2, self.n_bins, self.dim_t]).reshape( |
| 49 | + [-1, dim_c, self.n_bins, self.dim_t] |
| 50 | + ) |
| 51 | + return x[:, :, : self.dim_f] |
| 52 | + |
| 53 | + def istft(self, x, freq_pad=None): |
| 54 | + freq_pad = ( |
| 55 | + self.freq_pad.repeat([x.shape[0], 1, 1, 1]) |
| 56 | + if freq_pad is None |
| 57 | + else freq_pad |
| 58 | + ) |
| 59 | + x = torch.cat([x, freq_pad], -2) |
| 60 | + c = 4 * 2 if self.target_name == "*" else 2 |
| 61 | + x = x.reshape([-1, c, 2, self.n_bins, self.dim_t]).reshape( |
| 62 | + [-1, 2, self.n_bins, self.dim_t] |
| 63 | + ) |
| 64 | + x = x.permute([0, 2, 3, 1]) |
| 65 | + x = x.contiguous() |
| 66 | + x = torch.view_as_complex(x) |
| 67 | + x = torch.istft( |
| 68 | + x, n_fft=self.n_fft, hop_length=self.hop, window=self.window, center=True |
| 69 | + ) |
| 70 | + return x.reshape([-1, c, self.chunk_size]) |
| 71 | + |
| 72 | + |
| 73 | +def get_models(device, dim_f, dim_t, n_fft): |
| 74 | + return Conv_TDF_net_trim( |
| 75 | + device=device, |
| 76 | + model_name="Conv-TDF", |
| 77 | + target_name="vocals", |
| 78 | + L=11, |
| 79 | + dim_f=dim_f, |
| 80 | + dim_t=dim_t, |
| 81 | + n_fft=n_fft, |
| 82 | + ) |
| 83 | + |
| 84 | + |
| 85 | +warnings.filterwarnings("ignore") |
| 86 | +cpu = torch.device("cpu") |
| 87 | +if torch.cuda.is_available(): |
| 88 | + device = torch.device("cuda:0") |
| 89 | +elif torch.backends.mps.is_available(): |
| 90 | + device = torch.device("mps") |
| 91 | +else: |
| 92 | + device = torch.device("cpu") |
| 93 | + |
| 94 | + |
| 95 | +class Predictor: |
| 96 | + def __init__(self, args): |
| 97 | + self.args = args |
| 98 | + self.model_ = get_models( |
| 99 | + device=cpu, dim_f=args.dim_f, dim_t=args.dim_t, n_fft=args.n_fft |
| 100 | + ) |
| 101 | + self.model = ort.InferenceSession( |
| 102 | + os.path.join(args.onnx, self.model_.target_name + ".onnx"), |
| 103 | + providers=["CUDAExecutionProvider", "CPUExecutionProvider"], |
| 104 | + ) |
| 105 | + print("onnx load done") |
| 106 | + |
| 107 | + def demix(self, mix): |
| 108 | + samples = mix.shape[-1] |
| 109 | + margin = self.args.margin |
| 110 | + chunk_size = self.args.chunks * 44100 |
| 111 | + assert not margin == 0, "margin cannot be zero!" |
| 112 | + if margin > chunk_size: |
| 113 | + margin = chunk_size |
| 114 | + |
| 115 | + segmented_mix = {} |
| 116 | + |
| 117 | + if self.args.chunks == 0 or samples < chunk_size: |
| 118 | + chunk_size = samples |
| 119 | + |
| 120 | + counter = -1 |
| 121 | + for skip in range(0, samples, chunk_size): |
| 122 | + counter += 1 |
| 123 | + |
| 124 | + s_margin = 0 if counter == 0 else margin |
| 125 | + end = min(skip + chunk_size + margin, samples) |
| 126 | + |
| 127 | + start = skip - s_margin |
| 128 | + |
| 129 | + segmented_mix[skip] = mix[:, start:end].copy() |
| 130 | + if end == samples: |
| 131 | + break |
| 132 | + |
| 133 | + sources = self.demix_base(segmented_mix, margin_size=margin) |
| 134 | + """ |
| 135 | + mix:(2,big_sample) |
| 136 | + segmented_mix:offset->(2,small_sample) |
| 137 | + sources:(1,2,big_sample) |
| 138 | + """ |
| 139 | + return sources |
| 140 | + |
| 141 | + def demix_base(self, mixes, margin_size): |
| 142 | + chunked_sources = [] |
| 143 | + progress_bar = tqdm(total=len(mixes)) |
| 144 | + progress_bar.set_description("Processing") |
| 145 | + for mix in mixes: |
| 146 | + cmix = mixes[mix] |
| 147 | + sources = [] |
| 148 | + n_sample = cmix.shape[1] |
| 149 | + model = self.model_ |
| 150 | + trim = model.n_fft // 2 |
| 151 | + gen_size = model.chunk_size - 2 * trim |
| 152 | + pad = gen_size - n_sample % gen_size |
| 153 | + mix_p = np.concatenate( |
| 154 | + (np.zeros((2, trim)), cmix, np.zeros((2, pad)), np.zeros((2, trim))), 1 |
| 155 | + ) |
| 156 | + mix_waves = [] |
| 157 | + i = 0 |
| 158 | + while i < n_sample + pad: |
| 159 | + waves = np.array(mix_p[:, i : i + model.chunk_size]) |
| 160 | + mix_waves.append(waves) |
| 161 | + i += gen_size |
| 162 | + mix_waves = torch.tensor(mix_waves, dtype=torch.float32).to(cpu) |
| 163 | + with torch.no_grad(): |
| 164 | + _ort = self.model |
| 165 | + spek = model.stft(mix_waves) |
| 166 | + if self.args.denoise: |
| 167 | + spec_pred = ( |
| 168 | + -_ort.run(None, {"input": -spek.cpu().numpy()})[0] * 0.5 |
| 169 | + + _ort.run(None, {"input": spek.cpu().numpy()})[0] * 0.5 |
| 170 | + ) |
| 171 | + tar_waves = model.istft(torch.tensor(spec_pred)) |
| 172 | + else: |
| 173 | + tar_waves = model.istft( |
| 174 | + torch.tensor(_ort.run(None, {"input": spek.cpu().numpy()})[0]) |
| 175 | + ) |
| 176 | + tar_signal = ( |
| 177 | + tar_waves[:, :, trim:-trim] |
| 178 | + .transpose(0, 1) |
| 179 | + .reshape(2, -1) |
| 180 | + .numpy()[:, :-pad] |
| 181 | + ) |
| 182 | + |
| 183 | + start = 0 if mix == 0 else margin_size |
| 184 | + end = None if mix == list(mixes.keys())[::-1][0] else -margin_size |
| 185 | + if margin_size == 0: |
| 186 | + end = None |
| 187 | + sources.append(tar_signal[:, start:end]) |
| 188 | + |
| 189 | + progress_bar.update(1) |
| 190 | + |
| 191 | + chunked_sources.append(sources) |
| 192 | + _sources = np.concatenate(chunked_sources, axis=-1) |
| 193 | + # del self.model |
| 194 | + progress_bar.close() |
| 195 | + return _sources |
| 196 | + |
| 197 | + def prediction(self, m, vocal_root, others_root, format): |
| 198 | + os.makedirs(vocal_root, exist_ok=True) |
| 199 | + os.makedirs(others_root, exist_ok=True) |
| 200 | + basename = os.path.basename(m) |
| 201 | + mix, rate = librosa.load(m, mono=False, sr=44100) |
| 202 | + if mix.ndim == 1: |
| 203 | + mix = np.asfortranarray([mix, mix]) |
| 204 | + mix = mix.T |
| 205 | + sources = self.demix(mix.T) |
| 206 | + opt = sources[0].T |
| 207 | + if format in ["wav", "flac"]: |
| 208 | + sf.write( |
| 209 | + "%s/%s_main_vocal.%s" % (vocal_root, basename, format), mix - opt, rate |
| 210 | + ) |
| 211 | + sf.write("%s/%s_others.%s" % (others_root, basename, format), opt, rate) |
| 212 | + else: |
| 213 | + path_vocal = "%s/%s_main_vocal.wav" % (vocal_root, basename) |
| 214 | + path_other = "%s/%s_others.wav" % (others_root, basename) |
| 215 | + sf.write(path_vocal, mix - opt, rate) |
| 216 | + sf.write(path_other, opt, rate) |
| 217 | + if os.path.exists(path_vocal): |
| 218 | + os.system( |
| 219 | + "ffmpeg -i %s -vn %s -q:a 2 -y" |
| 220 | + % (path_vocal, path_vocal[:-4] + ".%s" % format) |
| 221 | + ) |
| 222 | + if os.path.exists(path_other): |
| 223 | + os.system( |
| 224 | + "ffmpeg -i %s -vn %s -q:a 2 -y" |
| 225 | + % (path_other, path_other[:-4] + ".%s" % format) |
| 226 | + ) |
| 227 | + |
| 228 | + |
| 229 | +class MDXNetDereverb: |
| 230 | + def __init__(self, chunks): |
| 231 | + self.onnx = "uvr5_weights/onnx_dereverb_By_FoxJoy" |
| 232 | + self.shifts = 10 #'Predict with randomised equivariant stabilisation' |
| 233 | + self.mixing = "min_mag" # ['default','min_mag','max_mag'] |
| 234 | + self.chunks = chunks |
| 235 | + self.margin = 44100 |
| 236 | + self.dim_t = 9 |
| 237 | + self.dim_f = 3072 |
| 238 | + self.n_fft = 6144 |
| 239 | + self.denoise = True |
| 240 | + self.pred = Predictor(self) |
| 241 | + |
| 242 | + def _path_audio_(self, input, vocal_root, others_root, format): |
| 243 | + self.pred.prediction(input, vocal_root, others_root, format) |
| 244 | + |
| 245 | + |
| 246 | +if __name__ == "__main__": |
| 247 | + dereverb = MDXNetDereverb(15) |
| 248 | + from time import time as ttime |
| 249 | + |
| 250 | + t0 = ttime() |
| 251 | + dereverb._path_audio_( |
| 252 | + "雪雪伴奏对消HP5.wav", |
| 253 | + "vocal", |
| 254 | + "others", |
| 255 | + ) |
| 256 | + t1 = ttime() |
| 257 | + print(t1 - t0) |
| 258 | + |
| 259 | + |
| 260 | +""" |
| 261 | +
|
| 262 | +runtime\python.exe MDXNet.py |
| 263 | +
|
| 264 | +6G: |
| 265 | +15/9:0.8G->6.8G |
| 266 | +14:0.8G->6.5G |
| 267 | +25:炸 |
| 268 | +
|
| 269 | +half15:0.7G->6.6G,22.69s |
| 270 | +fp32-15:0.7G->6.6G,20.85s |
| 271 | +
|
| 272 | +""" |
0 commit comments