diff --git a/rvc/lib/tools/model_download.py b/rvc/lib/tools/model_download.py index ab1b136e2..c0f50e409 100644 --- a/rvc/lib/tools/model_download.py +++ b/rvc/lib/tools/model_download.py @@ -328,9 +328,9 @@ def model_download_pipeline(url: str): extract_folder_path, model_name + ".pth" ), ) - else: - if "v2" not in item: - if "_nprobe_1_" in item and "_v1" in item: + elif ".index" in item: + if "_nprobe_1_" in item: + if "_v1" in item: file_name = item.split("_nprobe_1_")[1].split( "_v1" )[0] @@ -348,8 +348,7 @@ def model_download_pipeline(url: str): new_file_name + ".index", ), ) - else: - if "_nprobe_1_" in item and "_v2" in item: + elif "_v2" in item: file_name = item.split("_nprobe_1_")[1].split( "_v2" )[0] @@ -367,6 +366,15 @@ def model_download_pipeline(url: str): new_file_name + ".index", ), ) + else: + os.rename( + os.path.join(extract_folder_path, item), + os.path.join( + extract_folder_path, + model_name + ".index", + ), + ) + if success: print(f"Model {model_name} downloaded!") diff --git a/tabs/download/download.py b/tabs/download/download.py index fe9d5c322..d1801600b 100644 --- a/tabs/download/download.py +++ b/tabs/download/download.py @@ -35,16 +35,14 @@ def save_drop_model(dropbox): else: file_name = format_title(os.path.basename(dropbox)) if ".pth" in dropbox: - model_name = format_title(file_name.split(".pth")[0]) - else: - if "v2" not in dropbox: - model_name = format_title( - file_name.split("_nprobe_1_")[1].split("_v1")[0] - ) + model_name = file_name.split(".pth")[0] + elif ".index" in dropbox: + if "v1" in dropbox: + model_name = file_name.split("_nprobe_1_")[1].split("_v1")[0].split(".index")[0] + elif "v2" in dropbox: + model_name = file_name.split("_nprobe_1_")[1].split("_v2")[0].split(".index")[0] else: - model_name = format_title( - file_name.split("_nprobe_1_")[1].split("_v2")[0] - ) + model_name = file_name.split(".index")[0] model_path = os.path.join(now_dir, "logs", model_name) if not os.path.exists(model_path): os.makedirs(model_path)