EngineX-Iluvatar/enginex_bi_series-sherpa-onnx
Archived
16
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

Add UVR models for source separation. (#2266)

Browse Source
This commit is contained in:
Fangjun Kuang
2025-05-31 13:31:31 +08:00
committed by GitHub
parent 93e2819c18
commit 921f0f40cb
7 changed files with 575 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

261
scripts/uvr_mdx/test.py Executable file
View File

@@ -0,0 +1,261 @@
#!/usr/bin/env python3
# Copyright 2025 Xiaomi Corp. (authors: Fangjun Kuang)
import time
import argparse
import kaldi_native_fbank as knf
import librosa
import numpy as np
import onnxruntime as ort
import soundfile as sf
def get_args():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser.add_argument(
"--model-filename",
type=str,
required=True,
help="Path to onnx model",
)
parser.add_argument(
"--audio-filename",
type=str,
required=True,
help="Path to input audio file",
)
return parser.parse_args()
class OnnxModel:
def __init__(self, filename):
session_opts = ort.SessionOptions()
session_opts.inter_op_num_threads = 4
session_opts.intra_op_num_threads = 4
self.session_opts = session_opts
self.model = ort.InferenceSession(
filename,
sess_options=self.session_opts,
providers=["CPUExecutionProvider"],
)
self.dim_t = self.model.get_outputs()[0].shape[3]
self.dim_f = self.model.get_outputs()[0].shape[2]
self.n_fft = self.dim_f * 2
self.dim_c = self.model.get_outputs()[0].shape[1]
assert self.dim_c == 4, self.dim_c
self.hop = 1024
self.n_bins = self.n_fft // 2 + 1
self.chunk_size = self.hop * (self.dim_t - 1)
self.freq_pad = np.zeros([1, self.dim_c, self.n_bins - self.dim_f, self.dim_t])
print(f"----------inputs for {filename}----------")
for i in self.model.get_inputs():
print(i)
print(f"----------outputs for {filename}----------")
for i in self.model.get_outputs():
print(i)
print(i.shape)
print("--------------------")
def __call__(self, x):
"""
Args:
x: (batch_size, 4, self.dim_f, self.dim_t)
Returns:
spec: (batch_size, 4, self.dim_f, self.dim_t)
"""
spec = self.model.run(
[
self.model.get_outputs()[0].name,
],
{
self.model.get_inputs()[0].name: x,
},
)[0]
return spec
def main():
args = get_args()
m = OnnxModel(args.model_filename)
stft_config = knf.StftConfig(
n_fft=m.n_fft,
hop_length=m.hop,
win_length=m.n_fft,
center=True,
window_type="hann",
)
knf_stft = knf.Stft(stft_config)
knf_istft = knf.IStft(stft_config)
sample_rate = 44100
samples, rate = librosa.load(args.audio_filename, mono=False, sr=sample_rate)
start_time = time.time()
assert rate == sample_rate, (rate, sample_rate)
# samples: (2, 479832) , (num_channels, num_samples), 44100, 10.88
print("samples", samples.shape, rate, samples.shape[1] / rate)
assert samples.ndim == 2, samples.shape
assert samples.shape[0] == 2, samples.shape
margin = sample_rate
num_chunks = 15
chunk_size = num_chunks * sample_rate
# if they are too few samples, reset chunk_size
if samples.shape[1] < chunk_size:
chunk_size = samples.shape[1]
if margin > chunk_size:
margin = chunk_size
segments = []
for skip in range(0, samples.shape[1], chunk_size):
start = max(0, skip - margin)
end = min(skip + chunk_size + margin, samples.shape[1])
segments.append(samples[:, start:end])
if end == samples.shape[1]:
break
sources = []
for kk, s in enumerate(segments):
num_samples = s.shape[1]
trim = m.n_fft // 2
gen_size = m.chunk_size - 2 * trim
pad = gen_size - s.shape[1] % gen_size
mix_p = np.concatenate(
(
np.zeros((2, trim)),
s,
np.zeros((2, pad)),
np.zeros((2, trim)),
),
axis=1,
)
chunk_list = []
i = 0
while i < s.shape[1] + pad:
chunk_list.append(mix_p[:, i : i + m.chunk_size])
i += gen_size
mix_waves = np.array(chunk_list)
mix_waves_reshaped = mix_waves.reshape(-1, m.chunk_size)
stft_results = []
for w in mix_waves_reshaped:
stft = knf_stft(w)
stft_results.append(stft)
real = np.array(
[np.array(s.real).reshape(s.num_frames, -1) for s in stft_results],
dtype=np.float32,
)[:, :, :-1]
# real: (6, 256, 3072)
real = real.transpose(0, 2, 1)
# real: (6, 3072, 256)
imag = np.array(
[np.array(s.imag).reshape(s.num_frames, -1) for s in stft_results],
dtype=np.float32,
)[:, :, :-1]
imag = imag.transpose(0, 2, 1)
# imag: (6, 3072, 256)
x = np.stack([real, imag], axis=1)
# x: (6, 2, 3072, 256) -> (batch_size, real_imag, 3072, 256)
x = x.reshape(-1, m.dim_c, m.dim_f, m.dim_t)
# x: (3, 4, 3072, 256)
spec = m(x)
freq_pad = np.repeat(m.freq_pad, spec.shape[0], axis=0)
x = np.concatenate([spec, freq_pad], axis=2)
# x: (3, 4, 3073, 256)
x = x.reshape(-1, 2, m.n_bins, m.dim_t)
# x: (6, 2, 3073, 256)
x = x.transpose(0, 1, 3, 2)
# x: (6, 2, 256, 3073)
num_frames = x.shape[2]
x = x.reshape(x.shape[0], x.shape[1], -1)
wav_list = []
for k in range(x.shape[0]):
istft_result = knf.StftResult(
real=x[k, 0].reshape(-1).tolist(),
imag=x[k, 1].reshape(-1).tolist(),
num_frames=num_frames,
)
wav = knf_istft(istft_result)
wav_list.append(wav)
wav = np.array(wav_list, dtype=np.float32)
# wav: (6, 261120)
wav = wav.reshape(-1, 2, wav.shape[-1])
# wav: (3, 2, 261120)
wav = wav[:, :, trim:-trim]
# wav: (3, 2, 254976)
wav = wav.transpose(1, 0, 2)
# wav: (2, 3, 254976)
wav = wav.reshape(2, -1)
# wav: (2, 764928)
wav = wav[:, :-pad]
# wav: 2, 705600)
if kk == 0:
start = 0
else:
start = margin
if kk == len(segments) - 1:
end = None
else:
end = -margin
sources.append(wav[:, start:end])
sources = np.concatenate(sources, axis=-1)
vocals = sources
non_vocals = samples - vocals
end_time = time.time()
elapsed_seconds = end_time - start_time
print(f"Elapsed seconds: {elapsed_seconds:.3f}")
audio_duration = samples.shape[1] / sample_rate
real_time_factor = elapsed_seconds / audio_duration
print(f"Elapsed seconds: {elapsed_seconds:.3f}")
print(f"Audio duration in seconds: {audio_duration:.3f}")
print(f"RTF: {elapsed_seconds:.3f}/{audio_duration:.3f} = {real_time_factor:.3f}")
sf.write(f"./vocals.mp3", np.transpose(vocals), sample_rate)
sf.write(f"./non_vocals.mp3", np.transpose(non_vocals), sample_rate)
if __name__ == "__main__":
main()