Spaces:
Running
Running
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: | |
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) | |