296 lines
12 KiB
Python
296 lines
12 KiB
Python
import json
|
|
import os
|
|
import warnings
|
|
|
|
import numpy as np
|
|
import torch
|
|
from tqdm import tqdm
|
|
from .model_utils import get_hash
|
|
|
|
from .helper import gc_collect
|
|
import onnxruntime as ort
|
|
|
|
warnings.filterwarnings("ignore")
|
|
stem_naming = {'Vocals': 'Instrumental', 'Other': 'Instruments', 'Instrumental': 'Vocals', 'Drums': 'Drumless', 'Bass': 'Bassless'}
|
|
|
|
|
|
class MDXParams:
|
|
def __init__(self, device, dim_f, dim_t, n_fft, hop=1024, stem_name=None, compensation=1.000):
|
|
self.dim_f = dim_f
|
|
self.dim_t = dim_t
|
|
self.dim_c = 4
|
|
self.n_fft = n_fft
|
|
self.hop = hop
|
|
self.stem_name = stem_name
|
|
self.compensation = compensation
|
|
|
|
self.n_bins = self.n_fft // 2 + 1
|
|
self.trim = self.n_fft//2
|
|
self.chunk_size = hop * (self.dim_t - 1)
|
|
self.device = device
|
|
self.window = torch.hann_window(window_length=self.n_fft, periodic=True,device=self.device)
|
|
|
|
self.freq_pad = torch.zeros([1, self.dim_c, self.n_bins-self.dim_f, self.dim_t]).to(self.device)
|
|
self.gen_size = self.chunk_size-2*self.trim
|
|
|
|
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)
|
|
x = x.reshape([-1,2,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,2,self.chunk_size])
|
|
|
|
class MDXModel:
|
|
DEFAULT_SR = 44100
|
|
# Unit: seconds
|
|
DEFAULT_CHUNK_SIZE = 0 * DEFAULT_SR
|
|
DEFAULT_MARGIN_SIZE = 1 * DEFAULT_SR
|
|
|
|
DEFAULT_PROCESSOR = 0
|
|
|
|
# def __init__(self, model_path: str, params: MDXModel, processor=DEFAULT_PROCESSOR,margin=DEFAULT_MARGIN_SIZE,chunks=15):
|
|
def __init__(self, model_path: str, device="cpu",margin=DEFAULT_MARGIN_SIZE,chunks=15,denoise=False):
|
|
|
|
# Set the device and the provider (CPU or CUDA)
|
|
self.device = device #torch.device(f'cuda:{processor}') if processor >= 0 else torch.device('cpu')
|
|
# self.provider = ['CUDAExecutionProvider'] if processor >= 0 else ['CPUExecutionProvider']
|
|
self.providers = ['CUDAExecutionProvider','CPUExecutionProvider']
|
|
|
|
mp = self.get_params(model_path)
|
|
self.params = MDXParams(
|
|
self.device,
|
|
dim_f=mp["mdx_dim_f_set"],
|
|
dim_t=2 ** mp["mdx_dim_t_set"],
|
|
n_fft=mp["mdx_n_fft_scale_set"],
|
|
stem_name=mp["primary_stem"],
|
|
compensation=mp["compensate"]
|
|
)
|
|
|
|
# Load the ONNX model using ONNX Runtime
|
|
self.ort = ort.InferenceSession(model_path, providers=self.providers)
|
|
print(f"onnx load done: {model_path} ({self.params})")
|
|
|
|
self.run = lambda spec: self.ort.run(None, {'input': spec.cpu().numpy()})[0]
|
|
self.margin=margin
|
|
self.chunks=chunks
|
|
self.denoise=denoise
|
|
self.mdx_batch_size=1
|
|
|
|
def initialize_mix(self, mix, is_ckpt=False):
|
|
if is_ckpt:
|
|
pad = self.params.gen_size + self.params.trim - ((mix.shape[-1]) % self.params.gen_size)
|
|
mixture = np.concatenate((np.zeros((2, self.params.trim), dtype='float32'),mix, np.zeros((2, pad), dtype='float32')), 1)
|
|
num_chunks = mixture.shape[-1] // self.params.gen_size
|
|
mix_waves = [mixture[:, i * self.params.gen_size: i * self.params.gen_size + self.params.chunk_size] for i in range(num_chunks)]
|
|
else:
|
|
mix_waves = []
|
|
n_sample = mix.shape[1]
|
|
pad = self.params.gen_size - n_sample % self.params.gen_size
|
|
mix_p = np.concatenate((np.zeros((2,self.params.trim)), mix, np.zeros((2,pad)), np.zeros((2,self.params.trim))), 1)
|
|
i = 0
|
|
while i < n_sample + pad:
|
|
waves = np.array(mix_p[:, i:i+self.params.chunk_size])
|
|
mix_waves.append(waves)
|
|
i += self.params.gen_size
|
|
|
|
mix_waves = torch.tensor(mix_waves, dtype=torch.float32).to(self.device)
|
|
|
|
return mix_waves, pad
|
|
|
|
def demix_base(self, mix, is_ckpt=False, is_match_mix=False):
|
|
chunked_sources = []
|
|
for slice in tqdm(mix,"Processing audio:"):
|
|
sources = []
|
|
tar_waves_ = []
|
|
mix_p = mix[slice]
|
|
mix_waves, pad = self.initialize_mix(mix_p, is_ckpt=is_ckpt)
|
|
mix_waves = mix_waves.split(self.mdx_batch_size)
|
|
pad = mix_p.shape[-1] if is_ckpt else -pad
|
|
with torch.no_grad():
|
|
for mix_wave in mix_waves:
|
|
tar_waves = self.run_model(mix_wave, is_ckpt=is_ckpt, is_match_mix=is_match_mix)
|
|
tar_waves_.append(tar_waves)
|
|
tar_waves_ = np.vstack(tar_waves_)[:, :, self.params.trim:-self.params.trim] if is_ckpt else tar_waves_
|
|
tar_waves = np.concatenate(tar_waves_, axis=-1)[:, :pad]
|
|
start = 0 if slice == 0 else self.margin
|
|
end = None if slice == list(mix.keys())[::-1][0] or self.margin == 0 else -self.margin
|
|
sources.append(tar_waves[:,start:end]/self.params.compensation)
|
|
chunked_sources.append(sources)
|
|
sources = np.concatenate(chunked_sources, axis=-1)
|
|
|
|
return sources
|
|
|
|
def run_model(self, mix, is_ckpt=False, is_match_mix=False):
|
|
|
|
spek = self.params.stft(mix.to(self.device))*self.params.compensation
|
|
spek[:, :, :3, :] *= 0
|
|
|
|
if is_match_mix:
|
|
spec_pred = spek.cpu().numpy()
|
|
else:
|
|
spec_pred = self.run(spek)*0.5-self.run(-spek)*0.5 if self.denoise else self.run(spek)
|
|
|
|
if is_ckpt:
|
|
return self.params.istft(spec_pred).cpu().detach().numpy()
|
|
else:
|
|
return self.params.istft(torch.tensor(spec_pred).to(self.device)).cpu()[:,:,self.params.trim:-self.params.trim].transpose(0,1).reshape(2, -1).numpy()
|
|
|
|
def __del__(self):
|
|
del self.ort
|
|
gc_collect()
|
|
|
|
def get_params(self,model_path):
|
|
model_hash = get_hash(model_path)
|
|
with open(os.path.join(os.path.dirname(model_path), 'model_data.json')) as infile:
|
|
model_params = json.load(infile)
|
|
return model_params.get(model_hash)
|
|
|
|
# @staticmethod
|
|
# def segment(wave, combine=True, chunk_size=DEFAULT_CHUNK_SIZE, margin_size=DEFAULT_MARGIN_SIZE):
|
|
# """
|
|
# Segment or join segmented wave array
|
|
|
|
# Args:
|
|
# wave: (np.array) Wave array to be segmented or joined
|
|
# combine: (bool) If True, combines segmented wave array. If False, segments wave array.
|
|
# chunk_size: (int) Size of each segment (in samples)
|
|
# margin_size: (int) Size of margin between segments (in samples)
|
|
|
|
# Returns:
|
|
# numpy array: Segmented or joined wave array
|
|
# """
|
|
|
|
# if combine:
|
|
# processed_wave = None # Initializing as None instead of [] for later numpy array concatenation
|
|
# for segment_count, segment in enumerate(wave):
|
|
# start = 0 if segment_count == 0 else margin_size
|
|
# end = None if segment_count == len(wave) - 1 else -margin_size
|
|
# if margin_size == 0:
|
|
# end = None
|
|
# if processed_wave is None: # Create array for first segment
|
|
# processed_wave = segment[:, start:end]
|
|
# else: # Concatenate to existing array for subsequent segments
|
|
# processed_wave = np.concatenate((processed_wave, segment[:, start:end]), axis=-1)
|
|
|
|
# else:
|
|
# processed_wave = []
|
|
# sample_count = wave.shape[-1]
|
|
|
|
# if chunk_size <= 0 or chunk_size > sample_count:
|
|
# chunk_size = sample_count
|
|
|
|
# if margin_size > chunk_size:
|
|
# margin_size = chunk_size
|
|
|
|
# for segment_count, skip in enumerate(range(0, sample_count, chunk_size)):
|
|
|
|
# margin = 0 if segment_count == 0 else margin_size
|
|
# end = min(skip + chunk_size + margin_size, sample_count)
|
|
# start = skip - margin
|
|
|
|
# cut = wave[:, start:end].copy()
|
|
# processed_wave.append(cut)
|
|
|
|
# if end == sample_count:
|
|
# break
|
|
|
|
# return processed_wave
|
|
|
|
# def pad_wave(self, wave):
|
|
# """
|
|
# Pad the wave array to match the required chunk size
|
|
|
|
# Args:
|
|
# wave: (np.array) Wave array to be padded
|
|
|
|
# Returns:
|
|
# tuple: (padded_wave, pad, trim)
|
|
# - padded_wave: Padded wave array
|
|
# - pad: Number of samples that were padded
|
|
# - trim: Number of samples that were trimmed
|
|
# """
|
|
# n_sample = wave.shape[1]
|
|
# trim = self.params.n_fft // 2
|
|
# gen_size = self.params.chunk_size - 2 * trim
|
|
# pad = gen_size - n_sample % gen_size
|
|
|
|
# # Padded wave
|
|
# wave_p = np.concatenate((np.zeros((2, trim)), wave, np.zeros((2, pad)), np.zeros((2, trim))), 1)
|
|
|
|
# mix_waves = []
|
|
# for i in range(0, n_sample + pad, gen_size):
|
|
# waves = np.array(wave_p[:, i:i + self.params.chunk_size])
|
|
# mix_waves.append(waves)
|
|
|
|
# mix_waves = torch.tensor(mix_waves, dtype=torch.float32).cpu()
|
|
|
|
# return mix_waves, pad, trim
|
|
|
|
# # def _process_wave(self, mix_waves, trim, pad, q: queue.Queue, _id: int):
|
|
# def _process_wave(self, mix_waves, trim, pad):
|
|
# """
|
|
# Process each wave segment in a multi-threaded environment
|
|
|
|
# Args:
|
|
# mix_waves: (torch.Tensor) Wave segments to be processed
|
|
# trim: (int) Number of samples trimmed during padding
|
|
# pad: (int) Number of samples padded during padding
|
|
# q: (queue.Queue) Queue to hold the processed wave segments
|
|
# _id: (int) Identifier of the processed wave segment
|
|
|
|
# Returns:
|
|
# numpy array: Processed wave segment
|
|
# """
|
|
# mix_waves = mix_waves.split(1)
|
|
# with torch.no_grad():
|
|
# pw = []
|
|
# for mix_wave in mix_waves:
|
|
# # self.prog.update()
|
|
# spec = self.stft(mix_wave.to(self.device))
|
|
# processed_spec = torch.tensor(self.process(spec))
|
|
# processed_wav = self.istft(processed_spec)
|
|
# processed_wav = processed_wav[:, :, trim:-trim].transpose(0, 1).reshape(2, -1).cpu().numpy()
|
|
# pw.append(processed_wav)
|
|
|
|
# processed_signal = np.concatenate(pw, axis=-1)[:, :-pad]
|
|
|
|
# return processed_signal
|
|
|
|
# def process_wave(self, wave: np.array, mt_threads=1):
|
|
# """
|
|
# Process the wave array in a multi-threaded environment
|
|
|
|
# Args:
|
|
# wave: (np.array) Wave array to be processed
|
|
# mt_threads: (int) Number of threads to be used for processing
|
|
|
|
# Returns:
|
|
# numpy array: Processed wave array
|
|
# """
|
|
# # self.prog = tqdm(total=0)
|
|
# # chunk_size = wave.shape[-1] // (mt_threads if not self.chunks else (self.chunks*mt_threads))
|
|
# chunk_size = self.params.chunk_size
|
|
# margin = self.margin
|
|
# waves = self.segment(wave, False, chunk_size, margin)
|
|
|
|
# processed_batches = []
|
|
# for batch in tqdm(waves,desc="processing waves"):
|
|
# mix_waves, pad, trim = self.pad_wave(batch)
|
|
# processed_waves = self._process_wave(mix_waves, trim, pad)
|
|
# processed_batches.append(processed_waves)
|
|
|
|
# assert len(processed_batches) == len(waves), 'Incomplete processed batches, please reduce batch size!'
|
|
# segmented_mix = self.segment(processed_batches, True, chunk_size, margin)
|
|
|
|
# return segmented_mix |