fix: 🐛 Arranged format of model extraction

IAHispano · Dec 14, 2024 · c5b004c · c5b004c
1 parent 55814c2
commit c5b004c
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 68 deletions.
diff --git 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
@@ -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