surfmore/SimpleRVC

1import argparse2import sys3import torch4import json5from multiprocessing import cpu_count6 7global usefp168usefp16 = False9 10 11def use_fp32_config():12    usefp16 = False13    device_capability = 014    if torch.cuda.is_available():15        device = torch.device("cuda:0")  # Assuming you have only one GPU (index 0).16        device_capability = torch.cuda.get_device_capability(device)[0]17        if device_capability >= 7:18            usefp16 = True19            for config_file in ["32k.json", "40k.json", "48k.json"]:20                with open(f"configs/{config_file}", "r") as d:21                    data = json.load(d)22 23                if "train" in data and "fp16_run" in data["train"]:24                    data["train"]["fp16_run"] = True25 26                with open(f"configs/{config_file}", "w") as d:27                    json.dump(data, d, indent=4)28 29                print(f"Set fp16_run to true in {config_file}")30 31        else:32            for config_file in ["32k.json", "40k.json", "48k.json"]:33                with open(f"configs/{config_file}", "r") as f:34                    data = json.load(f)35 36                if "train" in data and "fp16_run" in data["train"]:37                    data["train"]["fp16_run"] = False38 39                with open(f"configs/{config_file}", "w") as d:40                    json.dump(data, d, indent=4)41 42                print(f"Set fp16_run to false in {config_file}")43    else:44        print(45            "CUDA is not available. Make sure you have an NVIDIA GPU and CUDA installed."46        )47    return (usefp16, device_capability)48 49 50class Config:51    def __init__(self):52        self.device = "cuda:0"53        self.is_half = True54        self.n_cpu = 055        self.gpu_name = None56        self.gpu_mem = None57        self.x_pad, self.x_query, self.x_center, self.x_max = self.device_config()58 59    # has_mps is only available in nightly pytorch (for now) and MasOS 12.3+.60    # check `getattr` and try it for compatibility61    @staticmethod62    def has_mps() -> bool:63        if not torch.backends.mps.is_available():64            return False65        try:66            torch.zeros(1).to(torch.device("mps"))67            return True68        except Exception:69            return False70 71    def device_config(self) -> tuple:72        if torch.cuda.is_available():73            i_device = int(self.device.split(":")[-1])74            self.gpu_name = torch.cuda.get_device_name(i_device)75            if (76                ("16" in self.gpu_name and "V100" not in self.gpu_name.upper())77                or "P40" in self.gpu_name.upper()78                or "1060" in self.gpu_name79                or "1070" in self.gpu_name80                or "1080" in self.gpu_name81            ):82                print("Found GPU", self.gpu_name, ", force to fp32")83                self.is_half = False84            else:85                print("Found GPU", self.gpu_name)86                use_fp32_config()87            self.gpu_mem = int(88                torch.cuda.get_device_properties(i_device).total_memory89                / 102490                / 102491                / 102492                + 0.493            )94        elif self.has_mps():95            print("No supported Nvidia GPU found, use MPS instead")96            self.device = "mps"97            self.is_half = False98            use_fp32_config()99        else:100            print("No supported Nvidia GPU found, use CPU instead")101            self.device = "cpu"102            self.is_half = False103            use_fp32_config()104 105        if self.n_cpu == 0:106            self.n_cpu = cpu_count()107 108        if self.is_half:109            # 6G显存配置110            x_pad = 3111            x_query = 10112            x_center = 60113            x_max = 65114        else:115            # 5G显存配置116            x_pad = 1117            x_query = 6118            x_center = 38119            x_max = 41120 121        if self.gpu_mem != None and self.gpu_mem <= 4:122            x_pad = 1123            x_query = 5124            x_center = 30125            x_max = 32126 127        return x_pad, x_query, x_center, x_max128