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()