diff --git a/custom_nodes/loaders.py b/custom_nodes/loaders.py index 1cdf343..3e414ef 100644 --- a/custom_nodes/loaders.py +++ b/custom_nodes/loaders.py @@ -12,8 +12,8 @@ import torch import folder_paths from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline from ..vc_infer_pipeline import get_vc -from .settings.downloader import RVC_DOWNLOAD_LINK, RVC_MODELS, download_file, slugify_filepath -from ..lib.audio import SUPPORTED_AUDIO, audio_to_bytes, bytes_to_audio, load_input_audio, save_input_audio +from .settings.downloader import RVC_DOWNLOAD_LINK, RVC_INDEX, RVC_MODELS, download_file +from ..lib.audio import SUPPORTED_AUDIO, audio_to_bytes, load_input_audio from ..config import config input_path = folder_paths.get_input_directory() @@ -117,10 +117,16 @@ class LoadRVCModelNode: def INPUT_TYPES(cls): model_list = RVC_MODELS + get_filenames(root=BASE_MODELS_DIR,folder="RVC",exts=["pth"],format_func=lambda x: f"RVC/{os.path.basename(x)}") model_list = list(set(model_list)) # dedupe + index_list = ["None"] + RVC_INDEX + get_filenames(root=BASE_MODELS_DIR,folder="RVC",exts=["index"],format_func=lambda x: f"RVC/.index/{os.path.basename(x)}") + index_list = list(set(index_list)) # dedupe + return { 'required': { - 'model': (model_list,{"default": "RVC/Sayano.pth"}), + 'model': (model_list,{"default": model_list[0]}), }, + "optional": { + "index": (index_list,{"default": "None"}), + } } RETURN_TYPES = ('RVC_MODEL', 'STRING') @@ -134,21 +140,26 @@ class LoadRVCModelNode: def IS_CHANGED(cls, model): return get_hash(model) - def load_model(self, model): + def load_model(self, model, index="None"): + model_path = file_index = None try: filename = os.path.basename(model) subfolder = os.path.dirname(model) model_path = os.path.join(BASE_MODELS_DIR,subfolder,filename) - + if not os.path.isfile(model_path): download_link = f"{RVC_DOWNLOAD_LINK}{model}" - params = model_path, download_link - if download_file(params): print(f"successfully downloaded: {model_path}") - model = lambda:get_vc(model_path) - return (model,filename.split(".")[0]) + if download_file((model_path, download_link)): print(f"successfully downloaded: {model_path}") + + if not index=="None": + file_index = os.path.join(BASE_MODELS_DIR,subfolder,".index",os.path.basename(index)) + if not os.path.isfile(file_index): + download_link = f"{RVC_DOWNLOAD_LINK}{index}" + if download_file((file_index, download_link)): print(f"successfully downloaded: {file_index}") except Exception as e: print(f"Error in {self.__class__.__name__}: {e}") raise e + finally: return (lambda:get_vc(model_path, file_index),filename.split(".")[0]) class LoadWhisperModelNode: diff --git a/custom_nodes/settings/downloader.py b/custom_nodes/settings/downloader.py index 10e8f3c..f4c6f40 100644 --- a/custom_nodes/settings/downloader.py +++ b/custom_nodes/settings/downloader.py @@ -26,6 +26,11 @@ RVC_MODELS = [ "RVC/Fuji.pth", "RVC/Monika.pth" ] +RVC_INDEX = [ + "RVC/.index/added_IVF1063_Flat_nprobe_1_Sayano_v2.index", + "RVC/.index/added_IVF985_Flat_nprobe_1_Fuji_v2.index", + "RVC/.index/Monika_v2_40k.index" +] BASE_MODELS = ["hubert_base.pt", "rmvpe.pt"] VITS_MODELS = ["VITS/pretrained_ljs.pth"] PRETRAINED_MODELS = [ diff --git a/pitch_extraction.py b/pitch_extraction.py index ab2fe4a..f78d8fe 100644 --- a/pitch_extraction.py +++ b/pitch_extraction.py @@ -50,11 +50,11 @@ class FeatureExtractor: gc_collect() def load_index(self, file_index): + index = big_npy = None try: - if not type(file_index)==str: # loading file index to save time + if type(file_index)==tuple: # loading file index to save time print("Using preloaded file index.") - index = file_index - big_npy = index.reconstruct_n(0, index.ntotal) + index,big_npy = file_index elif file_index == "": print("File index was empty.") index = None @@ -70,9 +70,7 @@ class FeatureExtractor: big_npy = index.reconstruct_n(0, index.ntotal) except Exception as e: print(f"Could not open Faiss index file for reading. {e}") - index = None - big_npy = None - return index, big_npy + finally: return index, big_npy # Fork Feature: Compute f0 with the crepe method def get_f0_crepe_computation( @@ -195,13 +193,6 @@ class FeatureExtractor: self.model_rmvpe = RMVPE(os.path.join(BASE_MODELS_DIR,f"rmvpe.{'onnx' if self.onnx else 'pt'}"), is_half=self.is_half, device=self.device, onnx=self.onnx) return self.model_rmvpe.infer_from_audio(x, thred=0.03) - # else: - # f0 = self.model_rmvpe.infer_from_audio(x, thred=0.03) - # if "privateuseone" in str(self.device): - # del self.model_rmvpe.model - # del self.model_rmvpe - # print("cleaning ortruntime memory") - # return f0 def get_pitch_dependant_rmvpe(self, x, f0_min=1, f0_max=40000, *args, **kwargs): if not hasattr(self,"model_rmvpe"): diff --git a/vc_infer_pipeline.py b/vc_infer_pipeline.py index d3f9e00..7284870 100644 --- a/vc_infer_pipeline.py +++ b/vc_infer_pipeline.py @@ -204,7 +204,7 @@ class VC(FeatureExtractor): return audio_opt -def get_vc(model_path,config=config,device=None): +def get_vc(model_path,file_index=None,config=config,device=None): cpt = torch.load(model_path, map_location="cpu") tgt_sr = cpt["config"][-1] cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0] # n_spk @@ -234,28 +234,25 @@ def get_vc(model_path,config=config,device=None): else: net_g = net_g.float() vc = VC(tgt_sr, config) - # hubert_model = load_hubert(hubert_path,config) model_name = os.path.basename(model_path).split(".")[0] - index_files = get_filenames(root=os.path.join(BASE_MODELS_DIR,"RVC"),folder=".index",exts=["index"],name_filters=[model_name]) - + try: #preload file_index - if len(index_files)==0: - print("File index was empty.") - file_index = None - else: + if os.path.exists(file_index): import faiss - file_index = index_files.pop() - if os.path.exists(file_index): - sys.stdout.write(f"Attempting to load {file_index}....\n") - sys.stdout.flush() - else: - sys.stdout.write(f"Attempting to load {file_index}.... (despite it not existing)\n") - sys.stdout.flush() - file_index = faiss.read_index(file_index) - sys.stdout.write(f"loaded index: {file_index}\n") + sys.stdout.write(f"Attempting to load {file_index}....\n") + sys.stdout.flush() + index = faiss.read_index(file_index) + big_npy = index.reconstruct_n(0, index.ntotal) + file_index = index, big_npy + else: + sys.stdout.write(f"Attempting to load {file_index}.... (despite it not existing)\n") + sys.stdout.flush() + file_index = "" + + sys.stdout.write(f"loaded index: {file_index}\n") except Exception as e: print(f"Could not open Faiss index file for reading. {e}") - file_index = None + file_index = "" return {"vc": vc, "cpt": cpt, "net_g": net_g, "model_name": model_name, "file_index": file_index, "sr": cpt["config"][-1]} @@ -287,7 +284,6 @@ def vc_single( ): print(f"vc_single unused args: {kwargs}") if hubert_model == None: - assert hubert_path is not None hubert_model = load_hubert(hubert_path,config) if not (cpt and net_g and vc and hubert_model): diff --git a/web/js/ui_handlers.js b/web/js/ui_handlers.js index f071a16..c2e25f1 100644 --- a/web/js/ui_handlers.js +++ b/web/js/ui_handlers.js @@ -303,7 +303,7 @@ function addUploadWidget(nodeType, nodeData, widgetName, type="video") { } else if (type == "audio") { Object.assign(fileInput, { type: "file", - accept: "audio/mpeg,audio/wav,audio/x-wav,audio/ogg", + accept: "audio/mpeg,audio/wav,audio/x-wav,audio/ogg,audio/flac", style: "display: none", onchange: async () => { if (fileInput.files.length) {