97 lines
3.3 KiB
Python
97 lines
3.3 KiB
Python
import os
|
|
import traceback
|
|
|
|
import torch
|
|
import ffmpeg
|
|
|
|
from configs.config import Config
|
|
from infer.modules.uvr5.preprocess import AudioPre, AudioPreDeEcho
|
|
from infer.modules.uvr5.mdxnet import MDXNetDereverb
|
|
|
|
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("tmp"),
|
|
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()
|
|
print("clean_empty_cache")
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
yield "\n".join(infos)
|