import os import traceback import logging logger = logging.getLogger(__name__) import ffmpeg import torch from configs.config import Config from infer.modules.uvr5.mdxnet import MDXNetDereverb from infer.modules.uvr5.preprocess import AudioPre, AudioPreDeEcho config = Config() def uvr(model_name, inp_root, save_root_vocal, paths, save_root_ins, agg, format0): infos = [] try: inp_root = inp_root.strip(" ").strip('"').strip("\n").strip('"').strip(" ") save_root_vocal = ( save_root_vocal.strip(" ").strip('"').strip("\n").strip('"').strip(" ") ) save_root_ins = ( save_root_ins.strip(" ").strip('"').strip("\n").strip('"').strip(" ") ) if model_name == "onnx_dereverb_By_FoxJoy": pre_fun = MDXNetDereverb(15, config.device) else: func = AudioPre if "DeEcho" not in model_name else AudioPreDeEcho pre_fun = func( agg=int(agg), model_path=os.path.join( os.getenv("weight_uvr5_root"), model_name + ".pth" ), device=config.device, is_half=config.is_half, ) if inp_root != "": paths = [os.path.join(inp_root, name) for name in os.listdir(inp_root)] else: paths = [path.name for path in paths] for path in paths: inp_path = os.path.join(inp_root, path) need_reformat = 1 done = 0 try: info = ffmpeg.probe(inp_path, cmd="ffprobe") if ( info["streams"][0]["channels"] == 2 and info["streams"][0]["sample_rate"] == "44100" ): need_reformat = 0 pre_fun._path_audio_( inp_path, save_root_ins, save_root_vocal, format0 ) done = 1 except: need_reformat = 1 traceback.print_exc() if need_reformat == 1: tmp_path = "%s/%s.reformatted.wav" % ( os.path.join(os.environ["TEMP"]), os.path.basename(inp_path), ) os.system( "ffmpeg -i %s -vn -acodec pcm_s16le -ac 2 -ar 44100 %s -y" % (inp_path, tmp_path) ) inp_path = tmp_path try: if done == 0: pre_fun.path_audio( inp_path, save_root_ins, save_root_vocal, format0 ) infos.append("%s->Success" % (os.path.basename(inp_path))) yield "\n".join(infos) except: infos.append( "%s->%s" % (os.path.basename(inp_path), traceback.format_exc()) ) yield "\n".join(infos) except: infos.append(traceback.format_exc()) yield "\n".join(infos) finally: try: if model_name == "onnx_dereverb_By_FoxJoy": del pre_fun.pred.model del pre_fun.pred.model_ else: del pre_fun.model del pre_fun except: traceback.print_exc() if torch.cuda.is_available(): torch.cuda.empty_cache() logger.info("Executed torch.cuda.empty_cache()") yield "\n".join(infos)