initial refactor of rvc studio
This commit is contained in:
SayanoAI
@@ -0,0 +1,367 @@
|
||||
import hashlib
|
||||
import json
|
||||
from multiprocessing.pool import ThreadPool
|
||||
import os
|
||||
from .utils import ObjectNamespace
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
from .mdx import MDXModel
|
||||
from .uvr5_pack.constants import MDX_NET_FREQ_CUT
|
||||
from .uvr5_pack.vr_network.model_param_init import ModelParameters
|
||||
from .uvr5_pack.vr_network.nets_new import CascadedNet
|
||||
from .uvr5_pack.vr_network.nets import CascadedASPPNet
|
||||
from .audio import remix_audio
|
||||
import librosa
|
||||
from .uvr5_pack import spec_utils
|
||||
|
||||
dir_path = os.path.dirname(os.path.realpath(__file__))
|
||||
|
||||
class UVR5Base:
|
||||
|
||||
def __init__(self, agg, model_path, device, is_half,**kwargs):
|
||||
self.model_path = model_path
|
||||
self.device = device
|
||||
self.data = {
|
||||
# Processing Options
|
||||
"postprocess": False,
|
||||
"tta": False,
|
||||
# Constants
|
||||
"window_size": 512,
|
||||
"agg": agg,
|
||||
"high_end_process": "mirroring",
|
||||
}
|
||||
mp = ModelParameters(os.path.join(dir_path,"uvr5_pack","vr_network","modelparams","4band_v2.json"))
|
||||
model = CascadedASPPNet(mp.param["bins"] * 2)
|
||||
cpk = torch.load(model_path, map_location=self.device)
|
||||
model.load_state_dict(cpk)
|
||||
model.eval()
|
||||
if is_half:
|
||||
model = model.half().to(device)
|
||||
else:
|
||||
model = model.to(device)
|
||||
|
||||
self.mp = mp
|
||||
self.model = model
|
||||
|
||||
@staticmethod
|
||||
def get_model_params(model_path):
|
||||
try:
|
||||
with open(model_path, 'rb') as f:
|
||||
f.seek(- 10000 * 1024, 2)
|
||||
model_hash = hashlib.md5(f.read()).hexdigest()
|
||||
except:
|
||||
model_hash = hashlib.md5(open(model_path, 'rb').read()).hexdigest()
|
||||
|
||||
print(model_hash)
|
||||
model_settings_json = os.path.splitext(model_path)[0]+".json"
|
||||
model_data_json = os.path.join(os.path.dirname(model_path),"model_data.json")
|
||||
|
||||
if os.path.isfile(model_settings_json):
|
||||
return json.load(open(model_settings_json))
|
||||
elif os.path.isfile(model_data_json):
|
||||
with open(model_data_json,"r") as d:
|
||||
hash_mapper = json.loads(d.read())
|
||||
|
||||
for hash, settings in hash_mapper.items():
|
||||
if model_hash in hash:
|
||||
return settings
|
||||
return None
|
||||
|
||||
def inference(self, X_spec, aggressiveness):
|
||||
"""
|
||||
data ： dic configs
|
||||
"""
|
||||
data = self.data
|
||||
device = self.device
|
||||
model = self.model
|
||||
|
||||
def _execute(
|
||||
X_mag_pad, roi_size, n_window, device, model, aggressiveness, is_half=True
|
||||
):
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
preds = []
|
||||
|
||||
for i in tqdm(range(n_window)):
|
||||
start = i * roi_size
|
||||
X_mag_window = X_mag_pad[
|
||||
None, :, :, start : start + data["window_size"]
|
||||
]
|
||||
X_mag_window = torch.from_numpy(X_mag_window)
|
||||
if is_half:
|
||||
X_mag_window = X_mag_window.half()
|
||||
X_mag_window = X_mag_window.to(device)
|
||||
|
||||
pred = model.predict(X_mag_window, aggressiveness)
|
||||
|
||||
pred = pred.detach().cpu().numpy()
|
||||
preds.append(pred[0])
|
||||
|
||||
pred = np.concatenate(preds, axis=2)
|
||||
return pred
|
||||
|
||||
def preprocess(X_spec):
|
||||
X_mag = np.abs(X_spec)
|
||||
X_phase = np.angle(X_spec)
|
||||
|
||||
return X_mag, X_phase
|
||||
|
||||
X_mag, X_phase = preprocess(X_spec)
|
||||
|
||||
coef = X_mag.max()
|
||||
X_mag_pre = X_mag / coef
|
||||
|
||||
n_frame = X_mag_pre.shape[2]
|
||||
pad_l, pad_r, roi_size = spec_utils.make_padding(n_frame, data["window_size"], model.offset)
|
||||
n_window = int(np.ceil(n_frame / roi_size))
|
||||
|
||||
X_mag_pad = np.pad(X_mag_pre, ((0, 0), (0, 0), (pad_l, pad_r)), mode="constant")
|
||||
|
||||
if list(model.state_dict().values())[0].dtype == torch.float16:
|
||||
is_half = True
|
||||
else:
|
||||
is_half = False
|
||||
pred = _execute(
|
||||
X_mag_pad, roi_size, n_window, device, model, aggressiveness, is_half
|
||||
)
|
||||
pred = pred[:, :, :n_frame]
|
||||
|
||||
if data["tta"]:
|
||||
pad_l += roi_size // 2
|
||||
pad_r += roi_size // 2
|
||||
n_window += 1
|
||||
|
||||
X_mag_pad = np.pad(X_mag_pre, ((0, 0), (0, 0), (pad_l, pad_r)), mode="constant")
|
||||
|
||||
pred_tta = _execute(
|
||||
X_mag_pad, roi_size, n_window, device, model, aggressiveness, is_half
|
||||
)
|
||||
pred_tta = pred_tta[:, :, roi_size // 2 :]
|
||||
pred_tta = pred_tta[:, :, :n_frame]
|
||||
|
||||
return (pred + pred_tta) * 0.5 * coef, X_mag, np.exp(1.0j * X_phase)
|
||||
else:
|
||||
return pred * coef, X_mag, np.exp(1.0j * X_phase)
|
||||
|
||||
def process_vocals(self,v_spec_m,input_high_end,input_high_end_h,return_dict={}):
|
||||
if self.data["high_end_process"].startswith("mirroring"):
|
||||
input_high_end_ = spec_utils.mirroring(
|
||||
self.data["high_end_process"], v_spec_m, input_high_end, self.mp
|
||||
)
|
||||
wav_vocals = spec_utils.cmb_spectrogram_to_wave(
|
||||
v_spec_m, self.mp, input_high_end_h, input_high_end_
|
||||
)
|
||||
else:
|
||||
wav_vocals = spec_utils.cmb_spectrogram_to_wave(v_spec_m, self.mp)
|
||||
print(f"vocals done: {wav_vocals.shape}")
|
||||
return_dict["vocals"] = remix_audio((wav_vocals,return_dict["sr"]),to_int16=True,axis=0)
|
||||
return return_dict["vocals"]
|
||||
|
||||
def process_instrumental(self,y_spec_m,input_high_end,input_high_end_h,return_dict={}):
|
||||
if self.data["high_end_process"].startswith("mirroring"):
|
||||
input_high_end_ = spec_utils.mirroring(
|
||||
self.data["high_end_process"], y_spec_m, input_high_end, self.mp
|
||||
)
|
||||
wav_instrument = spec_utils.cmb_spectrogram_to_wave(
|
||||
y_spec_m, self.mp, input_high_end_h, input_high_end_
|
||||
)
|
||||
else:
|
||||
wav_instrument = spec_utils.cmb_spectrogram_to_wave(y_spec_m, self.mp)
|
||||
print(f"instruments done: {wav_instrument.shape}")
|
||||
return_dict["instrumentals"] = remix_audio((wav_instrument,return_dict["sr"]),to_int16=True,axis=0)
|
||||
return return_dict["instrumentals"]
|
||||
|
||||
def process_audio(self,y_spec_m,v_spec_m,input_high_end,input_high_end_h):
|
||||
return_dict = {
|
||||
"sr": self.mp.param["sr"]
|
||||
}
|
||||
|
||||
with ThreadPool(2) as pool:
|
||||
pool.apply(self.process_vocals, args=(v_spec_m,input_high_end,input_high_end_h,return_dict))
|
||||
pool.apply(self.process_instrumental, args=(y_spec_m,input_high_end,input_high_end_h,return_dict))
|
||||
|
||||
return return_dict
|
||||
|
||||
def run_inference(self, music_file):
|
||||
X_wave, X_spec_s = {}, {}
|
||||
bands_n = len(self.mp.param["band"])
|
||||
# print(bands_n)
|
||||
for d in range(bands_n, 0, -1):
|
||||
bp = self.mp.param["band"][d]
|
||||
if d == bands_n: # high-end band
|
||||
input_audio = librosa.load(music_file,sr=bp["sr"],res_type=bp["res_type"])
|
||||
X_wave[d] = input_audio[0]
|
||||
if X_wave[d].ndim == 1:
|
||||
X_wave[d] = np.asfortranarray([X_wave[d], X_wave[d]])
|
||||
else: # lower bands
|
||||
X_wave[d] = librosa.resample(
|
||||
X_wave[d + 1],
|
||||
orig_sr=self.mp.param["band"][d + 1]["sr"],
|
||||
target_sr=bp["sr"],
|
||||
res_type=bp["res_type"],
|
||||
)
|
||||
# Stft of wave source
|
||||
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"],
|
||||
)
|
||||
# pdb.set_trace()
|
||||
if d == bands_n and self.data["high_end_process"] != "none":
|
||||
input_high_end_h = (bp["n_fft"] // 2 - bp["crop_stop"]) + (
|
||||
self.mp.param["pre_filter_stop"] - self.mp.param["pre_filter_start"]
|
||||
)
|
||||
input_high_end = X_spec_s[d][
|
||||
:, bp["n_fft"] // 2 - input_high_end_h : bp["n_fft"] // 2, :
|
||||
]
|
||||
|
||||
X_spec_m = spec_utils.combine_spectrograms(X_spec_s, self.mp)
|
||||
aggresive_set = float(self.data["agg"] / 100)
|
||||
aggressiveness = {
|
||||
"value": aggresive_set,
|
||||
"split_bin": self.mp.param["band"][1]["crop_stop"],
|
||||
}
|
||||
|
||||
# pred, X_mag, X_phase = run_inference(X_spec_m, self.device, self.models, aggressiveness, self.data)
|
||||
with torch.no_grad():
|
||||
pred, X_mag, X_phase = self.inference(X_spec_m, aggressiveness)
|
||||
|
||||
# return pred, X_mag, X_phase, X_spec_m, input_high_end,input_high_end_h
|
||||
|
||||
# def run_process_audio(self, pred, X_mag, X_phase, X_spec_m, input_high_end,input_high_end_h):
|
||||
# Postprocess
|
||||
if self.data["postprocess"]:
|
||||
pred_inv = np.clip(X_mag - pred, 0, np.inf)
|
||||
pred = spec_utils.mask_silence(pred, pred_inv)
|
||||
y_spec_m = pred * X_phase
|
||||
v_spec_m = X_spec_m - y_spec_m
|
||||
|
||||
|
||||
return_dict = self.process_audio(y_spec_m,v_spec_m,input_high_end,input_high_end_h)
|
||||
return_dict["input_audio"] = input_audio
|
||||
|
||||
return return_dict
|
||||
|
||||
class UVR5New(UVR5Base):
|
||||
def __init__(self, agg, model_path, device, is_half, dereverb, **kwargs):
|
||||
self.model_path = model_path
|
||||
self.device = device
|
||||
self.data = {
|
||||
# Processing Options
|
||||
"postprocess": False,
|
||||
"tta": False,
|
||||
# Constants
|
||||
"window_size": 512,
|
||||
"agg": agg,
|
||||
"high_end_process": "mirroring",
|
||||
}
|
||||
mp = ModelParameters(os.path.join(dir_path,"uvr5_pack","vr_network","modelparams","4band_v3.json"))
|
||||
nout = 64 if dereverb else 48
|
||||
model = CascadedNet(mp.param["bins"] * 2, nout)
|
||||
cpk = torch.load(model_path, map_location=self.device)
|
||||
model.load_state_dict(cpk)
|
||||
model.eval()
|
||||
if is_half:
|
||||
model = model.half().to(device)
|
||||
else:
|
||||
model = model.to(device)
|
||||
|
||||
self.mp = mp
|
||||
self.model = model
|
||||
|
||||
class MDXNet:
|
||||
def __init__(self, model_path, chunks=15,denoise=False,num_threads=1,device="cpu",**kwargs):
|
||||
|
||||
self.chunks = chunks
|
||||
self.sr = 44100
|
||||
|
||||
self.args = ObjectNamespace(**kwargs)
|
||||
self.denoise = denoise
|
||||
self.num_threads = num_threads
|
||||
|
||||
self.device = device
|
||||
|
||||
self.is_mdx_ckpt = "ckpt" in model_path
|
||||
|
||||
self.model = MDXModel(model_path, device=self.device,chunks=self.chunks,margin=self.sr)
|
||||
|
||||
def __del__(self):
|
||||
del self.model
|
||||
|
||||
def process_audio(self,primary,secondary,target_sr=None):
|
||||
target_sr = self.sr if target_sr is None else target_sr
|
||||
# foreground is processed data
|
||||
vocals,instrumental = (secondary,primary) if "instrument" in self.model.params.stem_name.lower() else (primary,secondary)
|
||||
with ThreadPool(2) as pool:
|
||||
results = pool.starmap(remix_audio, [
|
||||
((instrumental,self.sr),target_sr,False,True,self.sr!=target_sr,True),
|
||||
((vocals,self.sr),target_sr,False,True,self.sr!=target_sr,True)
|
||||
])
|
||||
|
||||
return_dict = {
|
||||
"sr": target_sr,
|
||||
"instrumentals": results[0],
|
||||
"vocals": results[1]
|
||||
}
|
||||
return return_dict
|
||||
|
||||
def run_inference(self, audio_path):
|
||||
|
||||
mdx_net_cut = True if self.model.params.stem_name in MDX_NET_FREQ_CUT else False
|
||||
mix, raw_mix, samplerate = prepare_mix(audio_path, self.model.chunks, self.model.margin, mdx_net_cut=mdx_net_cut)
|
||||
wave_processed = self.model.demix_base(mix, is_ckpt=self.is_mdx_ckpt)[0]
|
||||
|
||||
|
||||
raw_mix = self.model.demix_base(raw_mix, is_match_mix=True)[0] if mdx_net_cut else raw_mix
|
||||
|
||||
return_dict = self.process_audio(primary=wave_processed,secondary=(raw_mix-wave_processed),target_sr=samplerate)
|
||||
return_dict["input_audio"] = (raw_mix, samplerate)
|
||||
|
||||
return return_dict
|
||||
|
||||
|
||||
def prepare_mix(mix, chunk_set, margin_set, mdx_net_cut=False):
|
||||
|
||||
samplerate = 44100
|
||||
|
||||
if not isinstance(mix, np.ndarray):
|
||||
mix, samplerate = librosa.load(mix, mono=False, sr=44100)
|
||||
else:
|
||||
mix = mix.T
|
||||
|
||||
if mix.ndim == 1:
|
||||
mix = np.asfortranarray([mix,mix])
|
||||
|
||||
def get_segmented_mix(chunk_set=chunk_set):
|
||||
segmented_mix = {}
|
||||
|
||||
samples = mix.shape[-1]
|
||||
margin = margin_set
|
||||
chunk_size = chunk_set*44100
|
||||
assert not margin == 0, 'margin cannot be zero!'
|
||||
|
||||
if margin > chunk_size:
|
||||
margin = chunk_size
|
||||
if chunk_set == 0 or samples < chunk_size:
|
||||
chunk_size = samples
|
||||
|
||||
counter = -1
|
||||
for skip in range(0, samples, chunk_size):
|
||||
counter+=1
|
||||
s_margin = 0 if counter == 0 else margin
|
||||
end = min(skip+chunk_size+margin, samples)
|
||||
start = skip-s_margin
|
||||
segmented_mix[skip] = mix[:,start:end].copy()
|
||||
if end == samples:
|
||||
break
|
||||
|
||||
return segmented_mix
|
||||
|
||||
|
||||
segmented_mix = get_segmented_mix()
|
||||
raw_mix = get_segmented_mix(chunk_set=0) if mdx_net_cut else mix
|
||||
return segmented_mix, raw_mix, samplerate
|
||||
Reference in New Issue
Block a user