diff --git a/rvc/lib/tools/model_download.py b/rvc/lib/tools/model_download.py index b311bc0f..e35ac233 100644 --- a/rvc/lib/tools/model_download.py +++ b/rvc/lib/tools/model_download.py @@ -322,57 +322,21 @@ def model_download_pipeline(url: str): if ".pth" in item: file_name = item.split(".pth")[0] if file_name != model_name: + new_file_name = model_name + ".pth" os.rename( os.path.join(extract_folder_path, item), - os.path.join( - extract_folder_path, model_name + ".pth" - ), + os.path.join(extract_folder_path, new_file_name), ) elif ".index" in item: - if "_nprobe_1_" in item: - if "_v1" in item: - file_name = item.split("_nprobe_1_")[1].split( - "_v1" - )[0] - if file_name != model_name: - new_file_name = ( - item.split("_nprobe_1_")[0] - + "_nprobe_1_" - + model_name - + "_v1" - ) - os.rename( - os.path.join(extract_folder_path, item), - os.path.join( - extract_folder_path, - new_file_name + ".index", - ), - ) - elif "_v2" in item: - file_name = item.split("_nprobe_1_")[1].split( - "_v2" - )[0] - if file_name != model_name: - new_file_name = ( - item.split("_nprobe_1_")[0] - + "_nprobe_1_" - + model_name - + "_v2" - ) - os.rename( - os.path.join(extract_folder_path, item), - os.path.join( - extract_folder_path, - new_file_name + ".index", - ), - ) - else: + file_name = item.split(".index")[0] + replacements = ["nprobe_1_", "_v1", "_v2", "added_"] + for rep in replacements: + file_name = file_name.replace(rep, "") + if file_name != model_name: + new_file_name = model_name + ".index" os.rename( os.path.join(extract_folder_path, item), - os.path.join( - extract_folder_path, - model_name + ".index", - ), + os.path.join(extract_folder_path, new_file_name), ) if success: diff --git a/tabs/download/download.py b/tabs/download/download.py index ca0550ca..d93f179d 100644 --- a/tabs/download/download.py +++ b/tabs/download/download.py @@ -32,29 +32,27 @@ def save_drop_model(dropbox): raise gr.Error( message="The file you dropped is not a valid model file. Please try again." ) - else: - file_name = format_title(os.path.basename(dropbox)) - if ".pth" in dropbox: - 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 = 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) - if os.path.exists(os.path.join(model_path, file_name)): - os.remove(os.path.join(model_path, file_name)) - shutil.move(dropbox, os.path.join(model_path, file_name)) - print(f"{file_name} saved in {model_path}") - gr.Info(f"{file_name} saved in {model_path}") + + file_name = format_title(os.path.basename(dropbox)) + model_name = file_name + + if ".pth" in model_name: + model_name = model_name.split(".pth")[0] + elif ".index" in model_name: + replacements = ["nprobe_1_", "_v1", "_v2", "added_"] + for rep in replacements: + model_name = model_name.replace(rep, "") + model_name = model_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) + if os.path.exists(os.path.join(model_path, file_name)): + os.remove(os.path.join(model_path, file_name)) + shutil.move(dropbox, os.path.join(model_path, file_name)) + print(f"{file_name} saved in {model_path}") + gr.Info(f"{file_name} saved in {model_path}") + return None