fix: 🐛 Model extraction format

IAHispano · Dec 2, 2024 · f9d04c8 · f9d04c8
1 parent 2b630e7
commit f9d04c8
Show file tree

Hide file tree

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