From af4371064ebf43a4fc7519f00a4d984af6d54d4b Mon Sep 17 00:00:00 2001 From: SayanoAI <28544690+SayanoAI@users.noreply.github.com> Date: Sat, 1 Jun 2024 14:44:06 -0400 Subject: [PATCH] added bs roformer uvr + fixed caching --- custom_nodes/loaders.py | 1 - custom_nodes/rvc.py | 11 ++--- custom_nodes/settings/downloader.py | 2 +- custom_nodes/uvr.py | 69 ++++++++++++++++++++--------- lib/utils.py | 8 ++-- requirements.txt | 3 +- 6 files changed, 59 insertions(+), 35 deletions(-) diff --git a/custom_nodes/loaders.py b/custom_nodes/loaders.py index 70df6d8..5c488f7 100644 --- a/custom_nodes/loaders.py +++ b/custom_nodes/loaders.py @@ -246,7 +246,6 @@ class LoadAudio: sr = None if sr=="None" else int(sr) audio = load_input_audio(audio_path,sr=sr) return {"ui": {"preview": [{"filename": audio_name, "type": "input", "widgetId": widgetId}]}, "result": (audio_name, lambda:audio_to_bytes(*audio))} - # return (audio_name,lambda:audio_to_bytes(*audio)) @classmethod def IS_CHANGED(cls, audio): diff --git a/custom_nodes/rvc.py b/custom_nodes/rvc.py index 341097b..6e0bad6 100644 --- a/custom_nodes/rvc.py +++ b/custom_nodes/rvc.py @@ -10,6 +10,7 @@ from ..lib import BASE_CACHE_DIR, BASE_MODELS_DIR input_path = folder_paths.get_input_directory() temp_path = folder_paths.get_temp_directory() +cache_dir = os.path.join(BASE_CACHE_DIR,"rvc") # output_path = folder_paths.get_output_directory() # base_path = os.path.dirname(input_path) node_path = os.path.join(BASE_MODELS_DIR,"custom_nodes/ComfyUI-UVR5") @@ -39,11 +40,11 @@ class RVCNode: "step": 1, #Slider's step "display": "slider" }), + }, + "optional": { "format":(["wav", "flac", "mp3"],{ "default": "flac" }), - }, - "optional": { "use_cache": ("BOOLEAN",{"default": True}) } } @@ -56,15 +57,15 @@ class RVCNode: CATEGORY = "🌺RVC-Studio/rvc" - def convert(self, audio, model, hubert_model, pitch_extraction_params, format, use_cache=True, **kwargs): + def convert(self, audio, model, hubert_model, pitch_extraction_params, f0_up_key, format="flac", use_cache=True): - widgetId = get_hash(audio(), model, *pitch_extraction_params.items()) + widgetId = get_hash(audio(), model, f0_up_key, *pitch_extraction_params.items()) cache_name = os.path.join(BASE_CACHE_DIR,"rvc",f"{widgetId}.{format}") if use_cache and os.path.isfile(cache_name): output_audio = load_input_audio(cache_name) else: input_audio = bytes_to_audio(audio()) - output_audio = vc_single(hubert_model=hubert_model(),input_audio=input_audio,**model(),**pitch_extraction_params,**kwargs) + output_audio = vc_single(hubert_model=hubert_model(),input_audio=input_audio,f0_up_key=f0_up_key,**model(),**pitch_extraction_params) print(save_input_audio(cache_name, output_audio)) tempdir = os.path.join(temp_path,"preview") diff --git a/custom_nodes/settings/downloader.py b/custom_nodes/settings/downloader.py index 7d7d679..03e0450 100644 --- a/custom_nodes/settings/downloader.py +++ b/custom_nodes/settings/downloader.py @@ -10,7 +10,7 @@ from ...lib import BASE_CACHE_DIR, BASE_MODELS_DIR RVC_DOWNLOAD_LINK = 'https://huggingface.co/datasets/SayanoAI/RVC-Studio/resolve/main/' MDX_MODELS = ["MDXNET/UVR-MDX-NET-vocal_FT.onnx"] KARAFAN_MODELS = ["karafan/MDX23C-8KFFT-InstVoc_HQ.ckpt"] -VR_MODELS = ["UVR/UVR-DeEcho-DeReverb.pth","UVR/HP5-vocals+instrumentals.pth","UVR/5_HP-Karaoke-UVR.pth"] +VR_MODELS = ["UVR/UVR-DeEcho-DeReverb.pth","UVR/HP5-vocals+instrumentals.pth","UVR/5_HP-Karaoke-UVR.pth","UVR/model_bs_roformer_ep_317_sdr_12.9755.ckpt"] RVC_MODELS = [ "RVC/Claire.pth", "RVC/Sayano.pth", diff --git a/custom_nodes/uvr.py b/custom_nodes/uvr.py index 67b4d3b..35d1a7a 100644 --- a/custom_nodes/uvr.py +++ b/custom_nodes/uvr.py @@ -1,8 +1,8 @@ import os -from ..lib.audio import audio_to_bytes, bytes_to_audio, save_input_audio +from ..lib.audio import audio_to_bytes, bytes_to_audio, save_input_audio, load_input_audio import folder_paths -from ..lib.utils import get_filenames, get_hash, get_optimal_torch_device +from ..lib.utils import get_filenames, get_hash, get_file_hash, get_optimal_torch_device from ..lib import BASE_CACHE_DIR, BASE_MODELS_DIR, karafan from ..uvr5_cli import Separator from .settings.downloader import KARAFAN_MODELS, MDX_MODELS, RVC_DOWNLOAD_LINK, VR_MODELS, download_file @@ -20,7 +20,7 @@ class UVR5Node: @classmethod def INPUT_TYPES(cls): - model_list = MDX_MODELS + VR_MODELS + KARAFAN_MODELS + get_filenames(root=BASE_MODELS_DIR,format_func=lambda x: f"{os.path.basename(os.path.dirname(x))}/{os.path.basename(x)}",name_filters=["UVR","MDXNET","karafan"]) + model_list = MDX_MODELS + VR_MODELS + KARAFAN_MODELS + get_filenames(root=BASE_MODELS_DIR,format_func=lambda x: f"{os.path.basename(os.path.dirname(x))}/{os.path.basename(x)}",name_filters=["UVR","MDX","karafan"]) model_list = list(set(model_list)) # dedupe return { @@ -29,6 +29,9 @@ class UVR5Node: "model": (model_list,{ "default": "UVR/HP5-vocals+instrumentals.pth" }), + }, + "optional": { + "use_cache": ("BOOLEAN",{"default": True}), "agg":("INT",{ "default": 10, "min": 0, #Minimum value @@ -39,9 +42,6 @@ class UVR5Node: "format":(["wav", "flac", "mp3"],{ "default": "flac" }), - }, - "optional": { - "use_cache": ("BOOLEAN",{"default": True}) } } @@ -52,7 +52,7 @@ class UVR5Node: CATEGORY = "🌺RVC-Studio/uvr" - def split(self, audio, model, **kwargs): + def split(self, audio, model, use_cache=True, agg=10, format='flac'): filename = os.path.basename(model) subfolder = os.path.dirname(model) model_path = os.path.join(BASE_MODELS_DIR,subfolder,filename) @@ -62,22 +62,47 @@ class UVR5Node: if download_file(params): print(f"successfully downloaded: {model_path}") input_audio = bytes_to_audio(audio()) - audio_path = os.path.join(temp_path,"uvr",f"{get_hash(audio(),model, *kwargs.items())}.wav") - if not os.path.isfile(audio_path): - os.makedirs(os.path.dirname(audio_path),exist_ok=True) - print(save_input_audio(audio_path,input_audio)) - - if "karafan" in model_path: - vocals, instrumental, input_audio = karafan.inference.Process(audio_path,cache_dir=cache_dir,**kwargs) + hash_name = get_hash(audio(), model, agg, format) + audio_path = os.path.join(temp_path,"uvr",f"{hash_name}.wav") + vocal_path = os.path.join(cache_dir,hash_name,f"primary.{format}") + instrumental_path = os.path.join(cache_dir,hash_name,f"secondary.{format}") + if os.path.isfile(vocal_path) and os.path.isfile(instrumental_path) and use_cache: + vocals = load_input_audio(vocal_path) + instrumental = load_input_audio(instrumental_path) else: - model = Separator( - model_path=model_path, - device=device, - is_half="cuda" in str(device), - cache_dir=cache_dir, - **kwargs - ) - vocals, instrumental, input_audio = model.run_inference(audio_path,format=kwargs.get("format")) + if not os.path.isfile(audio_path): + os.makedirs(os.path.dirname(audio_path),exist_ok=True) + print(save_input_audio(audio_path,input_audio)) + vocals=instrumental=None + + try: # try original RVC implementation + model = Separator( + model_path=model_path, + device=device, + is_half="cuda" in str(device), + cache_dir=cache_dir, + agg=agg + ) + vocals, instrumental, input_audio = model.run_inference(audio_path,format=format) + except Exception as e: + print(f"Error: {e}") + if "karafan" in model_path: # try karafan implementation + vocals, instrumental, input_audio = karafan.inference.Process(audio_path,cache_dir=temp_path,format=format) + else: # try python-audio-separator implementation + import audio_separator.separator as uvr + model_dir = os.path.dirname(model_path) + model_name = os.path.basename(model_path) + vr_params={"batch_size": 4, "window_size": 512, "aggression": agg, "enable_tta": False, "enable_post_process": False, "post_process_threshold": 0.2, "high_end_process": "mirroring"} + model = uvr.Separator(model_file_dir=os.path.join(BASE_MODELS_DIR,model_dir),output_dir=temp_path,output_format=format,vr_params=vr_params) + model.load_model(model_name) + output_files = model.separate(audio_path) + print(f"{output_files=}") + vocals = load_input_audio(os.path.join(temp_path,output_files[0])) + instrumental = load_input_audio(os.path.join(temp_path,output_files[1])) + finally: + if vocals is not None and instrumental is not None and use_cache: + print(save_input_audio(vocal_path,vocals)) + print(save_input_audio(instrumental_path,instrumental)) return (lambda:audio_to_bytes(*vocals), lambda:audio_to_bytes(*instrumental), audio) diff --git a/lib/utils.py b/lib/utils.py index 0ccd900..18d0b0c 100644 --- a/lib/utils.py +++ b/lib/utils.py @@ -19,9 +19,9 @@ CWD = get_cwd() def get_hash(*args): return hashlib.md5("".join(str(data) for data in args).encode()).hexdigest() -def get_file_hash(fname,size=10*1024**2): - if hasattr(fname,"read"): - with open(fname,"rb") as f: return hashlib.md5(f.read(-size)).hexdigest() +def get_file_hash(fname: str, size=10*1024**2): + with open(fname,"rb") as f: + return hashlib.md5(f.read(size)).hexdigest() def get_subprocesses(pid = os.getpid()): # Get a list of all subprocesses started by the current process @@ -45,7 +45,7 @@ def get_filenames(root=CWD,folder="**",exts=["*"],name_filters=[""],filter_func= fnames = [] for ext in exts: fnames.extend(glob.glob(f"{root}/{folder}/*.{ext}",recursive=True)) - return sorted([format_func(ele) for ele in fnames if filter_func([nf.lower() in os.path.basename(ele).lower() for nf in name_filters])]) + return sorted([format_func(ele) for ele in fnames if filter_func([nf.lower() in ele.lower() for nf in name_filters])]) def get_index(arr,value): if arr is not None: diff --git a/requirements.txt b/requirements.txt index 8c13adc..b5b39a9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,10 +5,9 @@ faiss-cpu ffmpeg-python tqdm audioread +audio-separator[gpu] onnxruntime-gpu -onnxruntime ffmpy -accelerate samplerate pyaudio spacy