diff --git a/GPT_SoVITS/s2_train.py b/GPT_SoVITS/s2_train.py index 333e6a05..2f0a76e7 100644 --- a/GPT_SoVITS/s2_train.py +++ b/GPT_SoVITS/s2_train.py @@ -55,6 +55,10 @@ def main(): n_gpus = torch.cuda.device_count() else: n_gpus = 1 + if n_gpus <= 1: + run(0, n_gpus, hps) + return + os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(randint(20000, 55555)) @@ -69,6 +73,22 @@ def main(): def run(rank, n_gpus, hps): + use_ddp = n_gpus > 1 + if use_ddp: + dist.init_process_group( + backend="gloo" if os.name == "nt" or not torch.cuda.is_available() else "nccl", + init_method="env://?use_libuv=False", + world_size=n_gpus, + rank=rank, + ) + try: + _run(rank, n_gpus, hps, use_ddp) + finally: + if use_ddp and dist.is_initialized(): + dist.destroy_process_group() + + +def _run(rank, n_gpus, hps, use_ddp): global global_step if rank == 0: logger = utils.get_logger(hps.data.exp_dir) @@ -76,13 +96,6 @@ def run(rank, n_gpus, hps): # utils.check_git_hash(hps.s2_ckpt_dir) writer = SummaryWriter(log_dir=hps.s2_ckpt_dir) writer_eval = SummaryWriter(log_dir=os.path.join(hps.s2_ckpt_dir, "eval")) - - dist.init_process_group( - backend="gloo" if os.name == "nt" or not torch.cuda.is_available() else "nccl", - init_method="env://?use_libuv=False", - world_size=n_gpus, - rank=rank, - ) torch.manual_seed(hps.train.seed) if torch.cuda.is_available(): torch.cuda.set_device(rank) @@ -196,10 +209,10 @@ def run(rank, n_gpus, hps): betas=hps.train.betas, eps=hps.train.eps, ) - if torch.cuda.is_available(): + if torch.cuda.is_available() and use_ddp: net_g = DDP(net_g, device_ids=[rank], find_unused_parameters=True) net_d = DDP(net_d, device_ids=[rank], find_unused_parameters=True) - else: + elif not torch.cuda.is_available(): net_g = net_g.to(device) net_d = net_d.to(device) @@ -238,7 +251,7 @@ def run(rank, n_gpus, hps): torch.load(hps.train.pretrained_s2G, map_location="cpu", weights_only=False)["weight"], strict=False, ) - if torch.cuda.is_available() + if hasattr(net_g, "module") else net_g.load_state_dict( torch.load(hps.train.pretrained_s2G, map_location="cpu", weights_only=False)["weight"], strict=False, @@ -256,7 +269,7 @@ def run(rank, n_gpus, hps): net_d.module.load_state_dict( torch.load(hps.train.pretrained_s2D, map_location="cpu", weights_only=False)["weight"], strict=False ) - if torch.cuda.is_available() + if hasattr(net_d, "module") else net_d.load_state_dict( torch.load(hps.train.pretrained_s2D, map_location="cpu", weights_only=False)["weight"], ), diff --git a/requirements.txt b/requirements.txt index 00344838..d6593349 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,11 +1,10 @@ ---no-binary=opencc numpy<2.0 scipy tensorboard librosa==0.10.2 numba pytorch-lightning>=2.4 -gradio<5 +gradio==4.44.1 ffmpeg-python onnxruntime; platform_machine == "aarch64" or platform_machine == "arm64" onnxruntime-gpu; platform_machine == "x86_64" or platform_machine == "AMD64" @@ -34,7 +33,8 @@ g2pk2 ko_pron opencc python_mecab_ko; sys_platform != 'win32' -fastapi[standard]>=0.115.2 +fastapi[standard]==0.124.0 +starlette==0.49.3 x_transformers torchmetrics<=1.5 pydantic<=2.10.6