storyteller-ml/stem_separation/VR Arch/infer.py

# 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)
        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)
        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"
    )