Refactor: Format Code with Ruff and Update Deprecated G2PW Link (#2255)

* ruff check --fix * ruff format --line-length 120 --target-version py39 * Change the link for G2PW Model * update pytorch version and colab
2025-04-07 09:42:47 +01:00
parent 9da7e17efe
commit 53cac93589
132 changed files with 8185 additions and 6648 deletions
--- a/GPT_SoVITS/s2_train_v3.py
+++ b/GPT_SoVITS/s2_train_v3.py
@@ -1,36 +1,41 @@
 import warnings
+
 warnings.filterwarnings("ignore")
-import utils, os
+import os
+
+import utils
+
 hps = utils.get_hparams(stage=2)
 os.environ["CUDA_VISIBLE_DEVICES"] = hps.train.gpu_numbers.replace("-", ",")
+import logging
+
 import torch
-from torch.nn import functional as F
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from torch.cuda.amp import GradScaler, autocast
+from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
-import torch.multiprocessing as mp
-import torch.distributed as dist, traceback
-from torch.nn.parallel import DistributedDataParallel as DDP
-from torch.cuda.amp import autocast, GradScaler
 from tqdm import tqdm
-import logging, traceback

 logging.getLogger("matplotlib").setLevel(logging.INFO)
 logging.getLogger("h5py").setLevel(logging.INFO)
 logging.getLogger("numba").setLevel(logging.INFO)
 from random import randint
-from module import commons

+from module import commons
+from module.data_utils import (
+    DistributedBucketSampler,
+)
+from module.data_utils import (
+    TextAudioSpeakerCollateV3 as TextAudioSpeakerCollate,
+)
 from module.data_utils import (
    TextAudioSpeakerLoaderV3 as TextAudioSpeakerLoader,
-    TextAudioSpeakerCollateV3 as TextAudioSpeakerCollate,
-    DistributedBucketSampler,
 )
 from module.models import (
    SynthesizerTrnV3 as SynthesizerTrn,
-    MultiPeriodDiscriminator,
 )
-from module.losses import generator_loss, discriminator_loss, feature_loss, kl_loss
-from module.mel_processing import mel_spectrogram_torch, spec_to_mel_torch
 from process_ckpt import savee

 torch.backends.cudnn.benchmark = False
@@ -46,7 +51,6 @@ device = "cpu"  # cuda以外的设备，等mps优化后加入


 def main():
-
    if torch.cuda.is_available():
        n_gpus = torch.cuda.device_count()
    else:
@@ -74,7 +78,7 @@ def run(rank, n_gpus, hps):
        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",
+        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,
@@ -128,17 +132,21 @@ def run(rank, n_gpus, hps):
    #                              batch_size=1, pin_memory=True,
    #                              drop_last=False, collate_fn=collate_fn)

-    net_g = SynthesizerTrn(
-        hps.data.filter_length // 2 + 1,
-        hps.train.segment_size // hps.data.hop_length,
-        n_speakers=hps.data.n_speakers,
-        **hps.model,
-    ).cuda(rank) if torch.cuda.is_available() else SynthesizerTrn(
-        hps.data.filter_length // 2 + 1,
-        hps.train.segment_size // hps.data.hop_length,
-        n_speakers=hps.data.n_speakers,
-        **hps.model,
-    ).to(device)
+    net_g = (
+        SynthesizerTrn(
+            hps.data.filter_length // 2 + 1,
+            hps.train.segment_size // hps.data.hop_length,
+            n_speakers=hps.data.n_speakers,
+            **hps.model,
+        ).cuda(rank)
+        if torch.cuda.is_available()
+        else SynthesizerTrn(
+            hps.data.filter_length // 2 + 1,
+            hps.train.segment_size // hps.data.hop_length,
+            n_speakers=hps.data.n_speakers,
+            **hps.model,
+        ).to(device)
+    )

    # net_d = MultiPeriodDiscriminator(hps.model.use_spectral_norm).cuda(rank) if torch.cuda.is_available() else MultiPeriodDiscriminator(hps.model.use_spectral_norm).to(device)
    # for name, param in net_g.named_parameters():
@@ -146,7 +154,7 @@ def run(rank, n_gpus, hps):
    #         print(name, "not requires_grad")

    optim_g = torch.optim.AdamW(
-        filter(lambda p: p.requires_grad, net_g.parameters()),###默认所有层lr一致
+        filter(lambda p: p.requires_grad, net_g.parameters()),  ###默认所有层lr一致
        hps.train.learning_rate,
        betas=hps.train.betas,
        eps=hps.train.eps,
@@ -174,11 +182,11 @@ def run(rank, n_gpus, hps):
        #     logger.info("loaded D")
        # _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, optim_g,load_opt=0)
        _, _, _, epoch_str = utils.load_checkpoint(
-            utils.latest_checkpoint_path("%s/logs_s2_%s" % (hps.data.exp_dir,hps.model.version), "G_*.pth"),
+            utils.latest_checkpoint_path("%s/logs_s2_%s" % (hps.data.exp_dir, hps.model.version), "G_*.pth"),
            net_g,
            optim_g,
        )
-        epoch_str+=1
+        epoch_str += 1
        global_step = (epoch_str - 1) * len(train_loader)
        # epoch_str = 1
        # global_step = 0
@@ -186,17 +194,24 @@ def run(rank, n_gpus, hps):
        # traceback.print_exc()
        epoch_str = 1
        global_step = 0
-        if hps.train.pretrained_s2G != ""and hps.train.pretrained_s2G != None and os.path.exists(hps.train.pretrained_s2G):
+        if (
+            hps.train.pretrained_s2G != ""
+            and hps.train.pretrained_s2G != None
+            and os.path.exists(hps.train.pretrained_s2G)
+        ):
            if rank == 0:
                logger.info("loaded pretrained %s" % hps.train.pretrained_s2G)
-            print("loaded pretrained %s" % hps.train.pretrained_s2G,
+            print(
+                "loaded pretrained %s" % hps.train.pretrained_s2G,
                net_g.module.load_state_dict(
                    torch.load(hps.train.pretrained_s2G, map_location="cpu")["weight"],
                    strict=False,
-                ) if torch.cuda.is_available() else net_g.load_state_dict(
+                )
+                if torch.cuda.is_available()
+                else net_g.load_state_dict(
                    torch.load(hps.train.pretrained_s2G, map_location="cpu")["weight"],
                    strict=False,
-                )
+                ),
            )  ##测试不加载优化器
        # if hps.train.pretrained_s2D != ""and hps.train.pretrained_s2D != None and os.path.exists(hps.train.pretrained_s2D):
        #     if rank == 0:
@@ -212,9 +227,7 @@ def run(rank, n_gpus, hps):
    # scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2)
    # scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2)

-    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(
-        optim_g, gamma=hps.train.lr_decay, last_epoch=-1
-    )
+    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=hps.train.lr_decay, last_epoch=-1)
    # scheduler_d = torch.optim.lr_scheduler.ExponentialLR(
    #     optim_d, gamma=hps.train.lr_decay, last_epoch=-1
    # )
@@ -224,7 +237,7 @@ def run(rank, n_gpus, hps):

    scaler = GradScaler(enabled=hps.train.fp16_run)

-    net_d=optim_d=scheduler_d=None
+    net_d = optim_d = scheduler_d = None
    print("start training from epoch %s" % epoch_str)
    for epoch in range(epoch_str, hps.train.epochs + 1):
        if rank == 0:
@@ -260,7 +273,16 @@ def run(rank, n_gpus, hps):


 def train_and_evaluate(
-    rank, epoch, hps, nets, optims, schedulers, scaler, loaders, logger, writers
+    rank,
+    epoch,
+    hps,
+    nets,
+    optims,
+    schedulers,
+    scaler,
+    loaders,
+    logger,
+    writers,
 ):
    net_g, net_d = nets
    optim_g, optim_d = optims
@@ -284,19 +306,33 @@ def train_and_evaluate(
    #     text,
    #     text_lengths,
    # ) in enumerate(tqdm(train_loader)):
-    for batch_idx, (ssl, spec, mel, ssl_lengths, spec_lengths, text, text_lengths, mel_lengths) in enumerate(tqdm(train_loader)):
+    for batch_idx, (ssl, spec, mel, ssl_lengths, spec_lengths, text, text_lengths, mel_lengths) in enumerate(
+        tqdm(train_loader)
+    ):
        if torch.cuda.is_available():
-            spec, spec_lengths = spec.cuda(rank, non_blocking=True), spec_lengths.cuda(
-                rank, non_blocking=True
-            )
-            mel, mel_lengths = mel.cuda(rank, non_blocking=True), mel_lengths.cuda(
-                rank, non_blocking=True
+            spec, spec_lengths = (
+                spec.cuda(
+                    rank,
+                    non_blocking=True,
+                ),
+                spec_lengths.cuda(
+                    rank,
+                    non_blocking=True,
+                ),
            )
+            mel, mel_lengths = mel.cuda(rank, non_blocking=True), mel_lengths.cuda(rank, non_blocking=True)
            ssl = ssl.cuda(rank, non_blocking=True)
            ssl.requires_grad = False
            # ssl_lengths = ssl_lengths.cuda(rank, non_blocking=True)
-            text, text_lengths = text.cuda(rank, non_blocking=True), text_lengths.cuda(
-                rank, non_blocking=True
+            text, text_lengths = (
+                text.cuda(
+                    rank,
+                    non_blocking=True,
+                ),
+                text_lengths.cuda(
+                    rank,
+                    non_blocking=True,
+                ),
            )
        else:
            spec, spec_lengths = spec.to(device), spec_lengths.to(device)
@@ -307,8 +343,18 @@ def train_and_evaluate(
            text, text_lengths = text.to(device), text_lengths.to(device)

        with autocast(enabled=hps.train.fp16_run):
-            cfm_loss = net_g(ssl, spec, mel,ssl_lengths,spec_lengths, text, text_lengths,mel_lengths, use_grad_ckpt=hps.train.grad_ckpt)
-            loss_gen_all=cfm_loss
+            cfm_loss = net_g(
+                ssl,
+                spec,
+                mel,
+                ssl_lengths,
+                spec_lengths,
+                text,
+                text_lengths,
+                mel_lengths,
+                use_grad_ckpt=hps.train.grad_ckpt,
+            )
+            loss_gen_all = cfm_loss
        optim_g.zero_grad()
        scaler.scale(loss_gen_all).backward()
        scaler.unscale_(optim_g)
@@ -318,12 +364,15 @@ def train_and_evaluate(

        if rank == 0:
            if global_step % hps.train.log_interval == 0:
-                lr = optim_g.param_groups[0]['lr']
+                lr = optim_g.param_groups[0]["lr"]
                # losses = [commit_loss,cfm_loss,mel_loss,loss_disc, loss_gen, loss_fm, loss_mel, loss_kl]
                losses = [cfm_loss]
-                logger.info('Train Epoch: {} [{:.0f}%]'.format(
-                    epoch,
-                    100. * batch_idx / len(train_loader)))
+                logger.info(
+                    "Train Epoch: {} [{:.0f}%]".format(
+                        epoch,
+                        100.0 * batch_idx / len(train_loader),
+                    )
+                )
                logger.info([x.item() for x in losses] + [global_step, lr])

                scalar_dict = {"loss/g/total": loss_gen_all, "learning_rate": lr, "grad_norm_g": grad_norm_g}
@@ -337,7 +386,8 @@ def train_and_evaluate(
                    writer=writer,
                    global_step=global_step,
                    # images=image_dict,
-                    scalars=scalar_dict)
+                    scalars=scalar_dict,
+                )

            # if global_step % hps.train.eval_interval == 0:
            #     # evaluate(hps, net_g, eval_loader, writer_eval)
@@ -347,7 +397,6 @@ def train_and_evaluate(
            #     # if keep_ckpts > 0:
            #     #     utils.clean_checkpoints(path_to_models=hps.s2_ckpt_dir, n_ckpts_to_keep=keep_ckpts, sort_by_time=True)

-
        global_step += 1
    if epoch % hps.train.save_every_epoch == 0 and rank == 0:
        if hps.train.if_save_latest == 0:
@@ -357,7 +406,8 @@ def train_and_evaluate(
                hps.train.learning_rate,
                epoch,
                os.path.join(
-                    "%s/logs_s2_%s" % (hps.data.exp_dir,hps.model.version), "G_{}.pth".format(global_step)
+                    "%s/logs_s2_%s" % (hps.data.exp_dir, hps.model.version),
+                    "G_{}.pth".format(global_step),
                ),
            )
            # utils.save_checkpoint(
@@ -376,7 +426,8 @@ def train_and_evaluate(
                hps.train.learning_rate,
                epoch,
                os.path.join(
-                    "%s/logs_s2_%s" % (hps.data.exp_dir,hps.model.version), "G_{}.pth".format(233333333333)
+                    "%s/logs_s2_%s" % (hps.data.exp_dir, hps.model.version),
+                    "G_{}.pth".format(233333333333),
                ),
            )
            # utils.save_checkpoint(