fix(train): cannot extract feature on non-cuda devices (fix #123)

Author: fumiama

configs/config.py

@@ -53,6 +53,7 @@ def __init__(self):
         self.instead = ""
         self.preprocess_per = 3.7
         self.x_pad, self.x_query, self.x_center, self.x_max = self.device_config()
+        self.default_batch_size = self.get_default_batch_size()
 
     @staticmethod
     def load_config_json() -> dict:
@@ -136,6 +137,32 @@ def use_insecure_load():
             logging.warning("Using insecure weight loading for fairseq dictionary")
         except AttributeError:
             pass
+    
+    @staticmethod
+    def get_default_batch_size() -> int:
+        if not torch.cuda.is_available():
+            #TODO: add non-cuda multicards
+            return 1
+        # 判断是否有能用来训练和加速推理的N卡
+        ngpu = torch.cuda.device_count()
+        if not ngpu:
+            return 1
+        mem = []
+        if_gpu_ok = False
+
+        for i in range(ngpu):
+            if_gpu_ok = True  # 至少有一张能用的N卡
+            mem.append(
+                int(
+                    torch.cuda.get_device_properties(i).total_memory
+                    / 1024 / 1024 / 1024 + 0.4
+                )
+            )
+        if if_gpu_ok:
+            default_batch_size = min(mem) // 2
+        else:
+            default_batch_size = 1
+        return default_batch_size
 
     def use_fp32_config(self):
         for config_file in version_config_list:

infer/lib/rvcmd.py

@@ -194,11 +194,11 @@ def download_all_assets(tmpdir: str, version="0.2.5"):
     if not architecture:
         logger.error(f"architecture {architecture} is not supported")
         exit(1)
+    BASE_URL = "https://github.com/fumiama/RVC-Models-Downloader/releases/download/"
+    suffix = "zip" if is_win else "tar.gz"
+    RVCMD_URL = BASE_URL + f"v{version}/rvcmd_{system_type}_{architecture}.{suffix}"
+    cmdfile = os.path.join(tmpdir, "rvcmd")
     try:
-        BASE_URL = "https://github.com/fumiama/RVC-Models-Downloader/releases/download/"
-        suffix = "zip" if is_win else "tar.gz"
-        RVCMD_URL = BASE_URL + f"v{version}/rvcmd_{system_type}_{architecture}.{suffix}"
-        cmdfile = os.path.join(tmpdir, "rvcmd")
         if is_win:
             download_and_extract_zip(RVCMD_URL, tmpdir)
             cmdfile += ".exe"

infer/modules/train/extract_feature_print.py

@@ -10,19 +10,17 @@
 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
 os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] = "0.0"
 
+if len(sys.argv) != 8:
+    sys.exit(0)
+
 device = sys.argv[1]
 n_part = int(sys.argv[2])
 i_part = int(sys.argv[3])
-if len(sys.argv) == 7:
-    exp_dir = sys.argv[4]
-    version = sys.argv[5]
-    is_half = sys.argv[6].lower() == "true"
-else:
-    i_gpu = sys.argv[4]
-    exp_dir = sys.argv[5]
-    os.environ["CUDA_VISIBLE_DEVICES"] = str(i_gpu)
-    version = sys.argv[6]
-    is_half = sys.argv[7].lower() == "true"
+i_gpu = sys.argv[4]
+os.environ["CUDA_VISIBLE_DEVICES"] = str(i_gpu)
+exp_dir = sys.argv[5]
+version = sys.argv[6]
+is_half = sys.argv[7].lower() == "true"
 
 import fairseq
 import numpy as np

web.py

@@ -78,63 +78,6 @@ def forward_dml(ctx, x, scale):
 
 i18n = I18nAuto()
 logger.info(i18n)
-# 判断是否有能用来训练和加速推理的N卡
-ngpu = torch.cuda.device_count()
-gpu_infos = []
-mem = []
-if_gpu_ok = False
-
-if torch.cuda.is_available() or ngpu != 0:
-    for i in range(ngpu):
-        gpu_name = torch.cuda.get_device_name(i)
-        if any(
-            value in gpu_name.upper()
-            for value in [
-                "10",
-                "16",
-                "20",
-                "30",
-                "40",
-                "A2",
-                "A3",
-                "A4",
-                "P4",
-                "A50",
-                "500",
-                "A60",
-                "70",
-                "80",
-                "90",
-                "M4",
-                "T4",
-                "TITAN",
-                "4060",
-                "L",
-                "6000",
-            ]
-        ):
-            # A10#A100#V100#A40#P40#M40#K80#A4500
-            if_gpu_ok = True  # 至少有一张能用的N卡
-            gpu_infos.append("%s\t%s" % (i, gpu_name))
-            mem.append(
-                int(
-                    torch.cuda.get_device_properties(i).total_memory
-                    / 1024
-                    / 1024
-                    / 1024
-                    + 0.4
-                )
-            )
-if if_gpu_ok and len(gpu_infos) > 0:
-    gpu_info = "\n".join(gpu_infos)
-    default_batch_size = min(mem) // 2
-else:
-    gpu_info = i18n(
-        "Unfortunately, there is no compatible GPU available to support your training."
-    )
-    default_batch_size = 1
-gpus = "-".join([i[0] for i in gpu_infos])
-
 
 weight_root = os.getenv("weight_root")
 weight_uvr5_root = os.getenv("weight_uvr5_root")
@@ -314,6 +257,7 @@ def extract_f0_feature(n_p, f0method, if_f0, exp_dir, version19):
     exp_dir=sys.argv[4]
     os.environ["CUDA_VISIBLE_DEVICES"]=str(i_gpu)
     """
+    gpus = [config.device]
     leng = len(gpus)
     ps = []
     for idx, n_g in enumerate(gpus):
@@ -1201,7 +1145,7 @@ def change_info_(ckpt_path):
                 with gr.Column():
                     gpu_info9 = gr.Textbox(
                         label=i18n("GPU Information"),
-                        value=gpu_info,
+                        value=config.device,
                     )
                     f0method8 = gr.Radio(
                         label=i18n(
@@ -1254,7 +1198,7 @@ def change_info_(ckpt_path):
                         maximum=40,
                         step=1,
                         label=i18n("Batch size per GPU"),
-                        value=default_batch_size,
+                        value=config.default_batch_size,
                         interactive=True,
                     )
                     if_save_latest13 = gr.Radio(
@@ -1296,7 +1240,7 @@ def change_info_(ckpt_path):
                         label=i18n(
                             "Enter the GPU index(es) separated by '-', e.g., 0-1-2 to use GPU 0, 1, and 2"
                         ),
-                        value=gpus,
+                        value="0",
                         interactive=True,
                     )
                     sr2.change(