From ec21a76e8452349b608f951e84d914b9b092eff1 Mon Sep 17 00:00:00 2001 From: SayanoAI <28544690+SayanoAI@users.noreply.github.com> Date: Tue, 10 Sep 2024 11:58:28 -0400 Subject: [PATCH] fixed pretrains list --- custom_nodes/rvc_nodes.py | 15 ++++++++++++--- custom_nodes/settings/downloader.py | 23 ++++++++++++----------- lib/utils.py | 4 ++-- 3 files changed, 26 insertions(+), 16 deletions(-) diff --git a/custom_nodes/rvc_nodes.py b/custom_nodes/rvc_nodes.py index 8183ef5..de8d16f 100644 --- a/custom_nodes/rvc_nodes.py +++ b/custom_nodes/rvc_nodes.py @@ -20,7 +20,7 @@ from ..lib.model_utils import load_hubert from .utils import MultipleTypeProxy, increment_filename_no_overwrite, model_downloader from .settings import PITCH_EXTRACTION_OPTIONS -from .settings.downloader import PRETRAINED_MODELS, RVC_DOWNLOAD_LINK, RVC_INDEX, RVC_MODELS, download_file, extract_zip_without_structure +from .settings.downloader import PRETRAINED_MODELS_D, PRETRAINED_MODELS_G, RVC_DOWNLOAD_LINK, RVC_INDEX, RVC_MODELS, download_file, extract_zip_without_structure from ..lib.audio import SUPPORTED_AUDIO, audio_to_bytes, load_input_audio, save_input_audio, get_audio, SR_MAP @@ -387,9 +387,18 @@ class RVCTrainModelNode: @classmethod def INPUT_TYPES(cls): DEVICES = [str(i) for i in range(torch.cuda.device_count())] - PRETRAINED_G = [model for model in PRETRAINED_MODELS if "G" in model] + get_filenames(root=BASE_MODELS_DIR,folder="pretrained_v2",name_filters=["G"],format_func=lambda x: f"pretrained_v2/{os.path.basename(x)}") + PRETRAINED_G = PRETRAINED_MODELS_G + get_filenames( + root=BASE_MODELS_DIR, + folder="pretrained_v2", + name_filters=["G"], + filter_func=lambda x: os.stat(x).st_size/1024**2<600, #generators are 400+mb + format_func=lambda x: f"pretrained_v2/{os.path.basename(x)}") PRETRAINED_G = list(set(PRETRAINED_G)) # dedupe - PRETRAINED_D = [model for model in PRETRAINED_MODELS if "D" in model] + get_filenames(root=BASE_MODELS_DIR,folder="pretrained_v2",name_filters=["D"],format_func=lambda x: f"pretrained_v2/{os.path.basename(x)}") + PRETRAINED_D = PRETRAINED_MODELS_D + get_filenames( + root=BASE_MODELS_DIR,folder="pretrained_v2", + name_filters=["D"], + filter_func=lambda x: os.stat(x).st_size/1024**2>600, #discriminators are 800+mb + format_func=lambda x: f"pretrained_v2/{os.path.basename(x)}") PRETRAINED_D = list(set(PRETRAINED_D)) # dedupe return { diff --git a/custom_nodes/settings/downloader.py b/custom_nodes/settings/downloader.py index bafbbe5..402eb2c 100644 --- a/custom_nodes/settings/downloader.py +++ b/custom_nodes/settings/downloader.py @@ -37,27 +37,28 @@ RVC_INDEX = [ ] BASE_MODELS = ["content-vec-best.safetensors", "rmvpe.pt"] VITS_MODELS = ["VITS/pretrained_ljs.pth"] -PRETRAINED_MODELS = [ - "pretrained_v2/D48k.pth", +PRETRAINED_MODELS_G = [ "pretrained_v2/G48k.pth", - "pretrained_v2/D32k.pth", "pretrained_v2/G32k.pth", - "pretrained_v2/D40k.pth", "pretrained_v2/G40k.pth", - "pretrained_v2/f0D48k.pth", "pretrained_v2/f0G48k.pth", - "pretrained_v2/f0D40k.pth", "pretrained_v2/f0G40k.pth", - "pretrained_v2/f0D32k.pth", "pretrained_v2/f0G32k.pth", - - "pretrained_v2/f0_RIN_E3_40k_D.pth", "pretrained_v2/f0_RIN_E3_40k_G.pth", - "pretrained_v2/f0Ov2Super32kD.pth", "pretrained_v2/f0Ov2Super32kG.pth", - "pretrained_v2/f0Ov2Super40kD.pth", "pretrained_v2/f0Ov2Super40kG.pth", ] +PRETRAINED_MODELS_D = [ + "pretrained_v2/D48k.pth", + "pretrained_v2/D32k.pth", + "pretrained_v2/D40k.pth", + "pretrained_v2/f0D48k.pth", + "pretrained_v2/f0D40k.pth", + "pretrained_v2/f0D32k.pth", + "pretrained_v2/f0_RIN_E3_40k_D.pth", + "pretrained_v2/f0Ov2Super32kD.pth", + "pretrained_v2/f0Ov2Super40kD.pth", +] LLM_MODELS = [ "https://huggingface.co/TheBloke/Airoboros-L2-7B-2.1-GGUF/resolve/main/airoboros-l2-7b-2.1.Q4_K_M.gguf", "https://huggingface.co/TheBloke/Pygmalion-2-7B-GGUF/resolve/main/pygmalion-2-7b.Q4_K_M.gguf", diff --git a/lib/utils.py b/lib/utils.py index a3a2d8f..a1f014a 100644 --- a/lib/utils.py +++ b/lib/utils.py @@ -41,11 +41,11 @@ def get_subprocesses(pid = os.getpid()): }) yield process -def get_filenames(root=CWD,folder="**",exts=["*"],name_filters=[""],filter_func=any,format_func=lambda x:x): +def get_filenames(root=CWD,folder="**",exts=["*"],name_filters=[""],filter_func=bool,format_func=lambda x:x): 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 ele.lower() for nf in name_filters])]) + return sorted([format_func(ele) for ele in fnames if any([nf.lower() in ele.lower() for nf in name_filters]) and filter_func(ele)]) def get_index(arr,value): if arr is not None: