Paolify/RVC_4
0
1import torch2import onnxruntime as ort3from tqdm import tqdm4import warnings5import numpy as np6import hashlib7import queue8import threading9 10warnings.filterwarnings("ignore")11 12class MDX_Model:13 def __init__(self, device, dim_f, dim_t, n_fft, hop=1024, stem_name=None, compensation=1.000):14 self.dim_f = dim_f15 self.dim_t = dim_t16 self.dim_c = 417 self.n_fft = n_fft18 self.hop = hop19 self.stem_name = stem_name20 self.compensation = compensation21 22 self.n_bins = self.n_fft//2+123 self.chunk_size = hop * (self.dim_t-1)24 self.window = torch.hann_window(window_length=self.n_fft, periodic=True).to(device)25 26 out_c = self.dim_c27 28 self.freq_pad = torch.zeros([1, out_c, self.n_bins-self.dim_f, self.dim_t]).to(device)29 30 def stft(self, x):31 x = x.reshape([-1, self.chunk_size])32 x = torch.stft(x, n_fft=self.n_fft, hop_length=self.hop, window=self.window, center=True, return_complex=True)33 x = torch.view_as_real(x)34 x = x.permute([0,3,1,2])35 x = x.reshape([-1,2,2,self.n_bins,self.dim_t]).reshape([-1,4,self.n_bins,self.dim_t])36 return x[:,:,:self.dim_f]37 38 def istft(self, x, freq_pad=None):39 freq_pad = self.freq_pad.repeat([x.shape[0],1,1,1]) if freq_pad is None else freq_pad40 x = torch.cat([x, freq_pad], -2)41 # c = 4*2 if self.target_name=='*' else 242 x = x.reshape([-1,2,2,self.n_bins,self.dim_t]).reshape([-1,2,self.n_bins,self.dim_t])43 x = x.permute([0,2,3,1])44 x = x.contiguous()45 x = torch.view_as_complex(x)46 x = torch.istft(x, n_fft=self.n_fft, hop_length=self.hop, window=self.window, center=True)47 return x.reshape([-1,2,self.chunk_size])48 49 50class MDX:51 52 DEFAULT_SR = 4410053 # Unit: seconds54 DEFAULT_CHUNK_SIZE = 0 * DEFAULT_SR55 DEFAULT_MARGIN_SIZE = 1 * DEFAULT_SR56 57 DEFAULT_PROCESSOR = 058 59 def __init__(self, model_path:str, params:MDX_Model, processor=DEFAULT_PROCESSOR):60 61 # Set the device and the provider (CPU or CUDA)62 self.device = torch.device(f'cuda:{processor}') if processor >= 0 else torch.device('cpu')63 self.provider = ['CUDAExecutionProvider'] if processor >= 0 else ['CPUExecutionProvider']64 65 self.model = params66 67 # Load the ONNX model using ONNX Runtime68 self.ort = ort.InferenceSession(model_path, providers=self.provider)69 # Preload the model for faster performance70 self.ort.run(None, {'input':torch.rand(1, 4, params.dim_f, params.dim_t).numpy()})71 self.process = lambda spec:self.ort.run(None, {'input': spec.cpu().numpy()})[0]72 73 self.prog = None74 75 @staticmethod76 def get_hash(model_path):77 try:78 with open(model_path, 'rb') as f:79 f.seek(- 10000 * 1024, 2)80 model_hash = hashlib.md5(f.read()).hexdigest()81 except:82 model_hash = hashlib.md5(open(model_path,'rb').read()).hexdigest()83 84 return model_hash85 86 @staticmethod87 def segment(wave, combine=True, chunk_size=DEFAULT_CHUNK_SIZE, margin_size=DEFAULT_MARGIN_SIZE):88 """89 Segment or join segmented wave array90 91 Args:92 wave: (np.array) Wave array to be segmented or joined93 combine: (bool) If True, combines segmented wave array. If False, segments wave array.94 chunk_size: (int) Size of each segment (in samples)95 margin_size: (int) Size of margin between segments (in samples)96 97 Returns:98 numpy array: Segmented or joined wave array99 """100 101 if combine:102 processed_wave = None # Initializing as None instead of [] for later numpy array concatenation103 for segment_count, segment in enumerate(wave):104 start = 0 if segment_count == 0 else margin_size105 end = None if segment_count == len(wave)-1 else -margin_size106 if margin_size == 0:107 end = None108 if processed_wave is None: # Create array for first segment109 processed_wave = segment[:, start:end]110 else: # Concatenate to existing array for subsequent segments111 processed_wave = np.concatenate((processed_wave, segment[:, start:end]), axis=-1)112 113 else:114 processed_wave = []115 sample_count = wave.shape[-1]116 117 if chunk_size <= 0 or chunk_size > sample_count:118 chunk_size = sample_count119 120 if margin_size > chunk_size:121 margin_size = chunk_size122 123 for segment_count, skip in enumerate(range(0, sample_count, chunk_size)):124 125 margin = 0 if segment_count == 0 else margin_size126 end = min(skip+chunk_size+margin_size, sample_count)127 start = skip-margin128 129 cut = wave[:,start:end].copy()130 processed_wave.append(cut)131 132 if end == sample_count:133 break134 135 return processed_wave136 137 def pad_wave(self, wave):138 """139 Pad the wave array to match the required chunk size140 141 Args:142 wave: (np.array) Wave array to be padded143 144 Returns:145 tuple: (padded_wave, pad, trim)146 - padded_wave: Padded wave array147 - pad: Number of samples that were padded148 - trim: Number of samples that were trimmed149 """150 n_sample = wave.shape[1]151 trim = self.model.n_fft//2152 gen_size = self.model.chunk_size-2*trim153 pad = gen_size - n_sample%gen_size154 155 # Padded wave156 wave_p = np.concatenate((np.zeros((2,trim)), wave, np.zeros((2,pad)), np.zeros((2,trim))), 1)157 158 mix_waves = []159 for i in range(0, n_sample+pad, gen_size):160 waves = np.array(wave_p[:, i:i+self.model.chunk_size])161 mix_waves.append(waves)162 163 mix_waves = torch.tensor(mix_waves, dtype=torch.float32).to(self.device)164 165 return mix_waves, pad, trim166 167 def _process_wave(self, mix_waves, trim, pad, q:queue.Queue, _id:int):168 """169 Process each wave segment in a multi-threaded environment170 171 Args:172 mix_waves: (torch.Tensor) Wave segments to be processed173 trim: (int) Number of samples trimmed during padding174 pad: (int) Number of samples padded during padding175 q: (queue.Queue) Queue to hold the processed wave segments176 _id: (int) Identifier of the processed wave segment177 178 Returns:179 numpy array: Processed wave segment180 """181 mix_waves = mix_waves.split(1)182 with torch.no_grad():183 pw = []184 for mix_wave in mix_waves:185 self.prog.update()186 spec = self.model.stft(mix_wave)187 processed_spec = torch.tensor(self.process(spec))188 processed_wav = self.model.istft(processed_spec.to(self.device))189 processed_wav = processed_wav[:,:,trim:-trim].transpose(0,1).reshape(2, -1).cpu().numpy()190 pw.append(processed_wav)191 processed_signal = np.concatenate(pw, axis=-1)[:, :-pad]192 q.put({_id:processed_signal})193 return processed_signal194 195 def process_wave(self, wave:np.array, mt_threads=1):196 """197 Process the wave array in a multi-threaded environment198 199 Args:200 wave: (np.array) Wave array to be processed201 mt_threads: (int) Number of threads to be used for processing202 203 Returns:204 numpy array: Processed wave array205 """206 self.prog = tqdm(total=0)207 chunk = wave.shape[-1]//mt_threads208 waves = self.segment(wave, False, chunk)209 210 # Create a queue to hold the processed wave segments211 q = queue.Queue()212 threads = []213 for c, batch in enumerate(waves):214 mix_waves, pad, trim = self.pad_wave(batch)215 self.prog.total = len(mix_waves)*mt_threads216 thread = threading.Thread(target=self._process_wave, args=(mix_waves, trim, pad, q, c))217 thread.start()218 threads.append(thread)219 for thread in threads:220 thread.join()221 self.prog.close()222 223 processed_batches = []224 while not q.empty():225 processed_batches.append(q.get())226 processed_batches = [list(wave.values())[0] for wave in sorted(processed_batches, key=lambda d: list(d.keys())[0])]227 assert len(processed_batches) == len(waves), 'Incomplete processed batches, please reduce batch size!'228 return self.segment(processed_batches, True, chunk)