Sing-For-Me / config.py
Jarod Castillo
comments added
9d54d5c
import argparse
import sys
import torch
from multiprocessing import cpu_count
class Config:
"""
The code focuses on adapting the configuration based on available
hardware resources and specified command-line arguments,
aiming to optimize the performance and capabilities of the voice conversion process.
"""
def __init__(self):
"""
Calls the arg_parse() and device_config() methods to set up configuration based on command-line arguments
and available hardware.
Returns: None
"""
self.device = "cuda:0"
self.is_half = True
self.n_cpu = 0
self.gpu_name = None
self.gpu_mem = None
(
self.colab,
self.api,
self.unsupported
) = self.arg_parse()
self.x_pad, self.x_query, self.x_center, self.x_max = self.device_config()
@staticmethod
def arg_parse() -> tuple:
"""
Uses the argparse library to parse command-line arguments.
Three boolean arguments are defined: --colab, --api, and --unsupported.
Returns: a tuple indicating whether each argument is specified or not.
"""
parser = argparse.ArgumentParser()
parser.add_argument("--colab", action="store_true", help="Launch in colab")
parser.add_argument("--api", action="store_true", help="Launch with api")
parser.add_argument("--unsupported", action="store_true", help="Enable unsupported feature")
cmd_opts = parser.parse_args()
return (
cmd_opts.colab,
cmd_opts.api,
cmd_opts.unsupported
)
# has_mps is only available in nightly pytorch (for now) and MasOS 12.3+.
# check `getattr` and try it for compatibility
@staticmethod
def has_mps() -> bool:
"""
Determines if Multi-Process Service (MPS) is available in the current PyTorch backend.
If MPS is available, it checks whether it can be used by trying to move a tensor to the "mps" device.
Returns a boolean indicating MPS support.
"""
if not torch.backends.mps.is_available():
return False
try:
torch.zeros(1).to(torch.device("mps"))
return True
except Exception:
return False
def device_config(self) -> tuple:
"""
Checks if a CUDA-compatible GPU is available.
If a compatible GPU is found:
Determines the GPU's name and memory capacity.
Adjusts the is_half parameter based on the GPU's characteristics.
If no compatible GPU is found and MPS is available, configures the device to use MPS.
If no compatible GPU and MPS support, configures the device to use CPU.
Determines the number of available CPU cores (n_cpu).
Based on the is_half value and GPU memory capacity, configures several variables related to voice conversion,
such as x_pad, x_query, x_center, and x_max.
"""
if torch.cuda.is_available():
i_device = int(self.device.split(":")[-1])
self.gpu_name = torch.cuda.get_device_name(i_device)
if (
("16" in self.gpu_name and "V100" not in self.gpu_name.upper())
or "P40" in self.gpu_name.upper()
or "1060" in self.gpu_name
or "1070" in self.gpu_name
or "1080" in self.gpu_name
):
print("INFO: Found GPU", self.gpu_name, ", force to fp32")
self.is_half = False
else:
print("INFO: Found GPU", self.gpu_name)
self.gpu_mem = int(
torch.cuda.get_device_properties(i_device).total_memory
/ 1024
/ 1024
/ 1024
+ 0.4
)
elif self.has_mps():
print("INFO: No supported Nvidia GPU found, use MPS instead")
self.device = "mps"
self.is_half = False
else:
print("INFO: No supported Nvidia GPU found, use CPU instead")
self.device = "cpu"
self.is_half = False
if self.n_cpu == 0:
self.n_cpu = cpu_count()
if self.is_half:
# 6G显存配置
x_pad = 3
x_query = 10
x_center = 60
x_max = 65
else:
# 5G显存配置
x_pad = 1
x_query = 6
x_center = 38
x_max = 41
if self.gpu_mem != None and self.gpu_mem <= 4:
x_pad = 1
x_query = 5
x_center = 30
x_max = 32
return x_pad, x_query, x_center, x_max