initial refactor of rvc studio
This commit is contained in:
SayanoAI
+185
@@ -0,0 +1,185 @@
|
||||
import argparse
|
||||
import os, torch, warnings
|
||||
|
||||
from .lib.separators import MDXNet, UVR5Base, UVR5New
|
||||
from .lib import BASE_CACHE_DIR, karafan
|
||||
from .lib.audio import load_input_audio, pad_audio, remix_audio, save_input_audio
|
||||
from .lib.utils import gc_collect, get_optimal_threads, get_merge_func
|
||||
|
||||
CACHED_SONGS_DIR = os.path.join(BASE_CACHE_DIR,"songs")
|
||||
|
||||
warnings.filterwarnings("ignore")
|
||||
import numpy as np
|
||||
|
||||
class Separator:
|
||||
def __init__(self, model_path, use_cache=False, device="cpu", cache_dir=None, **kwargs):
|
||||
dereverb = "reverb" in model_path.lower()
|
||||
deecho = "echo" in model_path.lower()
|
||||
bve = "BVE" in model_path.lower()
|
||||
denoise = dereverb or deecho or bve
|
||||
|
||||
if "MDX" in model_path:
|
||||
self.model = MDXNet(model_path=model_path,denoise=denoise,device=device,**kwargs)
|
||||
elif "UVR" in model_path:
|
||||
self.model = UVR5New(model_path=model_path,device=device,dereverb=dereverb,**kwargs) if denoise else UVR5Base(model_path=model_path,device=device,**kwargs)
|
||||
|
||||
self.use_cache = use_cache
|
||||
self.cache_dir = cache_dir
|
||||
self.model_path = model_path
|
||||
self.args = kwargs
|
||||
|
||||
# cleanup memory
|
||||
def __del__(self):
|
||||
gc_collect()
|
||||
|
||||
def run_inference(self, audio_path, format="mp3"):
|
||||
song_name = get_filename(os.path.basename(self.model_path).split(".")[0],**self.args) + f".{format}"
|
||||
|
||||
# handles loading of previous processed data
|
||||
music_dir = os.path.join(
|
||||
os.path.dirname(audio_path) if self.cache_dir is None else self.cache_dir,
|
||||
os.path.basename(audio_path).split(".")[0])
|
||||
vocals_path = os.path.join(music_dir,".vocals")
|
||||
instrumental_path = os.path.join(music_dir,".instrumental")
|
||||
vocals_file = os.path.join(vocals_path,song_name)
|
||||
instrumental_file = os.path.join(instrumental_path,song_name)
|
||||
|
||||
if os.path.isfile(instrumental_file) and os.path.isfile(vocals_file):
|
||||
vocals = load_input_audio(vocals_file,mono=True)
|
||||
instrumental = load_input_audio(instrumental_file,mono=True)
|
||||
input_audio = load_input_audio(audio_path,mono=True)
|
||||
return vocals, instrumental, input_audio
|
||||
|
||||
return_dict = self.model.run_inference(audio_path)
|
||||
instrumental = return_dict["instrumentals"]
|
||||
vocals = return_dict["vocals"]
|
||||
input_audio = return_dict["input_audio"]
|
||||
|
||||
if self.use_cache:
|
||||
os.makedirs(vocals_path,exist_ok=True)
|
||||
os.makedirs(instrumental_path,exist_ok=True)
|
||||
save_input_audio(vocals_file,vocals,to_int16=True)
|
||||
save_input_audio(instrumental_file,instrumental,to_int16=True)
|
||||
|
||||
return vocals, instrumental, input_audio
|
||||
|
||||
def get_filename(*args,**kwargs):
|
||||
name = "_".join([str(arg) for arg in args]+[f"{k}={v}" for k,v in kwargs.items()])
|
||||
return name
|
||||
|
||||
def __run_inference_worker(arg):
|
||||
(model_path,audio_path,agg,device,use_cache,cache_dir,num_threads,format) = arg
|
||||
if "karafan" in model_path:
|
||||
vocals, instrumental, input_audio = karafan.inference.Process(audio_path,cache_dir=cache_dir,use_cache=use_cache,format=format)
|
||||
else:
|
||||
model = Separator(
|
||||
agg=agg,
|
||||
model_path=model_path,
|
||||
device=device,
|
||||
is_half="cuda" in str(device),
|
||||
use_cache=use_cache,
|
||||
cache_dir=cache_dir,
|
||||
num_threads = num_threads
|
||||
)
|
||||
vocals, instrumental, input_audio = model.run_inference(audio_path,format)
|
||||
del model
|
||||
gc_collect()
|
||||
|
||||
return vocals, instrumental, input_audio
|
||||
|
||||
def split_audio(uvr_models,audio_path,preprocess_models=[],postprocess_models=[],device="cuda",agg=10,use_cache=False,merge_type="mean",format="mp3",**kwargs):
|
||||
print(f"unused kwargs={kwargs}")
|
||||
merge_func = get_merge_func(merge_type)
|
||||
num_threads = max(get_optimal_threads(-1),1)
|
||||
song_name = os.path.basename(audio_path).split(".")[0]
|
||||
cache_dir = os.path.join(CACHED_SONGS_DIR,song_name)
|
||||
|
||||
# preprocess input song to split reverb
|
||||
if len(preprocess_models):
|
||||
output_name = get_filename(*[os.path.basename(name).split(".")[0] for name in preprocess_models],agg=agg) + f".{format}"
|
||||
preprocessed_file = os.path.join(cache_dir,"preprocessing",output_name)
|
||||
|
||||
# read from cache
|
||||
if os.path.isfile(preprocessed_file): input_audio = load_input_audio(preprocessed_file,mono=True)
|
||||
else: # preprocess audio
|
||||
for i,preprocess_model in enumerate(preprocess_models):
|
||||
output_name = get_filename(i,os.path.basename(preprocess_model).split(".")[0],agg=agg) + f".{format}"
|
||||
intermediary_file = os.path.join(cache_dir,"preprocessing",output_name)
|
||||
if os.path.isfile(intermediary_file):
|
||||
# if i==len(preprocess_model)-1: #last model
|
||||
instrumental = input_audio = load_input_audio(intermediary_file, mono=True)
|
||||
else:
|
||||
args = (preprocess_model,audio_path,agg,device,False,CACHED_SONGS_DIR if i==0 else None,num_threads,format)
|
||||
_, instrumental, input_audio = __run_inference_worker(args)
|
||||
save_input_audio(intermediary_file,instrumental,to_int16=True)
|
||||
audio_path = intermediary_file
|
||||
|
||||
save_input_audio(preprocessed_file,instrumental,to_int16=True)
|
||||
audio_path = preprocessed_file
|
||||
else:
|
||||
input_audio = load_input_audio(audio_path,mono=True)
|
||||
|
||||
# apply vocal separation
|
||||
wav_instrument = []
|
||||
wav_vocals = []
|
||||
|
||||
for model_path in uvr_models:
|
||||
args = (model_path,audio_path,agg,device,use_cache,cache_dir,num_threads,format)
|
||||
print(f"processing... {args=}")
|
||||
vocals, instrumental, _ = __run_inference_worker(args)
|
||||
wav_vocals.append(vocals[0])
|
||||
wav_instrument.append(instrumental[0])
|
||||
wav_instrument = np.nanmedian(pad_audio(*wav_instrument,axis=0),axis=0)
|
||||
wav_vocals = merge_func(pad_audio(*wav_vocals,axis=0),axis=0)
|
||||
|
||||
# postprocess vocals to reduce reverb
|
||||
if len(postprocess_models):
|
||||
vocals_name = get_filename("vocals",*[os.path.basename(name).split(".")[0] for name in uvr_models],agg=agg) + f".{format}"
|
||||
vocals_file = os.path.join(cache_dir,"postprocessing",vocals_name)
|
||||
if not os.path.isfile(vocals_file): save_input_audio(vocals_file,(wav_vocals,vocals[-1]),to_int16=True)
|
||||
print("postprocessing...")
|
||||
for i,postprocess_model in enumerate(postprocess_models):
|
||||
output_name = get_filename(i,os.path.basename(postprocess_model).split(".")[0],agg=agg) + f".{format}"
|
||||
intermediary_file = os.path.join(cache_dir,"postprocessing",output_name)
|
||||
if not os.path.isfile(intermediary_file):
|
||||
args = (postprocess_model,vocals_file,agg,device,False,None,num_threads,format)
|
||||
_, processed_audio, _ = __run_inference_worker(args)
|
||||
output_name = get_filename(i,os.path.basename(postprocess_model).split(".")[0],agg=agg) + f".{format}"
|
||||
save_input_audio(intermediary_file,processed_audio,to_int16=True)
|
||||
wav_vocals, _ = processed_audio
|
||||
vocals_file = intermediary_file
|
||||
|
||||
instrumental = remix_audio((wav_instrument,instrumental[-1]),to_int16=True)
|
||||
vocals = remix_audio((wav_vocals,vocals[-1]),to_int16=True)
|
||||
|
||||
return vocals, instrumental, input_audio
|
||||
|
||||
def main(): #uvr5_models,audio_path,device="cuda",agg=10,use_cache=False
|
||||
|
||||
parser = argparse.ArgumentParser(description="processes audio to split vocal stems and reduce reverb/echo")
|
||||
|
||||
parser.add_argument("uvr5_models", type=str, nargs="+", help="Path to models to use for processing")
|
||||
parser.add_argument(
|
||||
"-i", "--audio_path", type=str, help="path to audio file to process", required=True
|
||||
)
|
||||
parser.add_argument(
|
||||
"-p", "--preprocess_model", type=str, help="preprocessing model to improve audio", default=None
|
||||
)
|
||||
parser.add_argument(
|
||||
"-a", "--agg", type=int, default=10, help="aggressiveness score for processing (0-20)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-d", "--device", type=str, default="cpu", choices=["cpu","cuda"], help="perform calculations on [cpu] or [cuda]"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-m", "--merge_type", type=str, default="median", choices=["mean","median"], help="how to combine processed audio"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-c", "--use_cache", type=bool, action="store_true", default=False, help="caches the results so next run is faster"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
return split_audio(**vars(args))
|
||||
|
||||
if __name__ == "__main__":
|
||||
torch.multiprocessing.set_start_method("spawn")
|
||||
main()
|
||||
Reference in New Issue
Block a user