mosibi/RVC_HFv2

1import soundfile as sf2import torch, pdb, os, warnings, librosa3import numpy as np4import onnxruntime as ort5from tqdm import tqdm6import torch7 8dim_c = 49 10 11class Conv_TDF_net_trim:12    def __init__(13        self, device, model_name, target_name, L, dim_f, dim_t, n_fft, hop=102414    ):15        super(Conv_TDF_net_trim, self).__init__()16 17        self.dim_f = dim_f18        self.dim_t = 2**dim_t19        self.n_fft = n_fft20        self.hop = hop21        self.n_bins = self.n_fft // 2 + 122        self.chunk_size = hop * (self.dim_t - 1)23        self.window = torch.hann_window(window_length=self.n_fft, periodic=True).to(24            device25        )26        self.target_name = target_name27        self.blender = "blender" in model_name28 29        out_c = dim_c * 4 if target_name == "*" else dim_c30        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 // 235 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 None57            else freq_pad58        )59        x = torch.cat([x, freq_pad], -2)60        c = 4 * 2 if self.target_name == "*" else 261        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=True69        )70        return x.reshape([-1, c, self.chunk_size])71 72 73def 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 85warnings.filterwarnings("ignore")86cpu = torch.device("cpu")87if torch.cuda.is_available():88    device = torch.device("cuda:0")89elif torch.backends.mps.is_available():90    device = torch.device("mps")91else:92    device = torch.device("cpu")93 94 95class Predictor:96    def __init__(self, args):97        self.args = args98        self.model_ = get_models(99            device=cpu, dim_f=args.dim_f, dim_t=args.dim_t, n_fft=args.n_fft100        )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.margin110        chunk_size = self.args.chunks * 44100111        assert not margin == 0, "margin cannot be zero!"112        if margin > chunk_size:113            margin = chunk_size114 115        segmented_mix = {}116 117        if self.args.chunks == 0 or samples < chunk_size:118            chunk_size = samples119 120        counter = -1121        for skip in range(0, samples, chunk_size):122            counter += 1123 124            s_margin = 0 if counter == 0 else margin125            end = min(skip + chunk_size + margin, samples)126 127            start = skip - s_margin128 129            segmented_mix[skip] = mix[:, start:end].copy()130            if end == samples:131                break132 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 sources140 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 // 2151            gen_size = model.chunk_size - 2 * trim152            pad = gen_size - n_sample % gen_size153            mix_p = np.concatenate(154                (np.zeros((2, trim)), cmix, np.zeros((2, pad)), np.zeros((2, trim))), 1155            )156            mix_waves = []157            i = 0158            while i < n_sample + pad:159                waves = np.array(mix_p[:, i : i + model.chunk_size])160                mix_waves.append(waves)161                i += gen_size162            mix_waves = torch.tensor(mix_waves, dtype=torch.float32).to(cpu)163            with torch.no_grad():164                _ort = self.model165                spek = model.stft(mix_waves)166                if self.args.denoise:167                    spec_pred = (168                        -_ort.run(None, {"input": -spek.cpu().numpy()})[0] * 0.5169                        + _ort.run(None, {"input": spek.cpu().numpy()})[0] * 0.5170                    )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_size184                end = None if mix == list(mixes.keys())[::-1][0] else -margin_size185                if margin_size == 0:186                    end = None187                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.model194        progress_bar.close()195        return _sources196 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.T205        sources = self.demix(mix.T)206        opt = sources[0].T207        if format in ["wav", "flac"]:208            sf.write(209                "%s/%s_main_vocal.%s" % (vocal_root, basename, format), mix - opt, rate210            )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 229class 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 = chunks235        self.margin = 44100236        self.dim_t = 9237        self.dim_f = 3072238        self.n_fft = 6144239        self.denoise = True240        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 246if __name__ == "__main__":247    dereverb = MDXNetDereverb(15)248    from time import time as ttime249 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 262runtime\python.exe MDXNet.py 263 2646G:26515/9:0.8G->6.8G26614:0.8G->6.5G26725:炸268 269half15:0.7G->6.6G,22.69s270fp32-15:0.7G->6.6G,20.85s271 272"""273