xinference 1.9.1__py3-none-any.whl → 1.10.1__py3-none-any.whl

xinference/thirdparty/indextts/utils/maskgct/models/codec/facodec/facodec_trainer.py

@@ -0,0 +1,776 @@
+# Copyright (c) 2023 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import time
+import random
+from pathlib import Path
+import re
+import glob
+
+import accelerate
+import json
+import numpy as np
+import torch
+from accelerate.utils import ProjectConfiguration
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+
+from accelerate.logging import get_logger
+
+from models.codec.facodec.facodec_dataset import FAcodecDataset, FAcodecCollator
+from models.codec.codec_sampler import build_samplers
+from models.codec.codec_trainer import CodecTrainer
+
+from modules.dac.nn.loss import (
+    MultiScaleSTFTLoss,
+    MelSpectrogramLoss,
+    GANLoss,
+    L1Loss,
+    FocalLoss,
+)
+from audiotools import AudioSignal
+
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
+
+try:
+    import nemo.collections.asr as nemo_asr
+except ImportError:
+    print(
+        "Unable to import nemo_asr, titanet outputs will be set to random values, you may only run debugging mode. DO NOT USE THIS FOR TRAINING"
+    )
+    nemo_asr = None
+
+from models.codec.facodec.modules.commons import (
+    build_model,
+    load_checkpoint,
+    load_F0_models,
+    log_norm,
+)
+from models.codec.facodec.optimizer import build_optimizer
+
+
+class FAcodecTrainer(CodecTrainer):
+    def __init__(self, args, cfg):
+        super().__init__()
+
+        self.args = args
+        self.cfg = cfg
+
+        cfg.exp_name = args.exp_name
+
+        # Init accelerator
+        self._init_accelerator()
+        self.accelerator.wait_for_everyone()
+
+        # Init logger
+        with self.accelerator.main_process_first():
+            self.logger = get_logger(args.exp_name, log_level=args.log_level)
+
+        self.logger.info("=" * 56)
+        self.logger.info("||\t\t" + "New training process started." + "\t\t||")
+        self.logger.info("=" * 56)
+        self.logger.info("\n")
+        self.logger.debug(f"Using {args.log_level.upper()} logging level.")
+        self.logger.info(f"Experiment name: {args.exp_name}")
+        self.logger.info(f"Experiment directory: {self.exp_dir}")
+        self.checkpoint_dir = os.path.join(self.exp_dir, "checkpoint")
+        if self.accelerator.is_main_process:
+            os.makedirs(self.checkpoint_dir, exist_ok=True)
+        self.logger.debug(f"Checkpoint directory: {self.checkpoint_dir}")
+
+        # Init training status
+        self.batch_count: int = 0
+        self.step: int = 0
+        self.epoch: int = 0
+
+        self.max_epoch = (
+            self.cfg.train.max_epoch if self.cfg.train.max_epoch > 0 else float("inf")
+        )
+        self.logger.info(
+            "Max epoch: {}".format(
+                self.max_epoch if self.max_epoch < float("inf") else "Unlimited"
+            )
+        )
+
+        # Check potential erorrs
+        if self.accelerator.is_main_process:
+            self._check_basic_configs()
+            self.save_checkpoint_stride = self.cfg.train.save_checkpoint_stride
+            self.checkpoints_path = [
+                [] for _ in range(len(self.save_checkpoint_stride))
+            ]
+            self.run_eval = self.cfg.train.run_eval
+
+        # Set random seed
+        with self.accelerator.main_process_first():
+            start = time.monotonic_ns()
+            self._set_random_seed(self.cfg.train.random_seed)
+            end = time.monotonic_ns()
+            self.logger.debug(
+                f"Setting random seed done in {(end - start) / 1e6:.2f}ms"
+            )
+            self.logger.debug(f"Random seed: {self.cfg.train.random_seed}")
+
+        # Build dataloader
+        with self.accelerator.main_process_first():
+            self.logger.info("Building dataset...")
+            start = time.monotonic_ns()
+            self.train_dataloader, self.valid_dataloader = self._build_dataloader()
+            end = time.monotonic_ns()
+            self.logger.info(f"Building dataset done in {(end - start) / 1e6:.2f}ms")
+
+        # Build model
+        with self.accelerator.main_process_first():
+            self.logger.info("Building model...")
+            start = time.monotonic_ns()
+            self.model = self._build_model()
+            end = time.monotonic_ns()
+            for _, model in self.model.items():
+                self.logger.debug(model)
+            self.logger.info(f"Building model done in {(end - start) / 1e6:.2f}ms")
+            self.logger.info(f"Model parameters: {self._count_parameters()/1e6:.2f}M")
+
+        # Build optimizers and schedulers
+        with self.accelerator.main_process_first():
+            self.logger.info("Building optimizer and scheduler...")
+            start = time.monotonic_ns()
+            self.optimizer = self._build_optimizer()
+            end = time.monotonic_ns()
+            self.logger.info(
+                f"Building optimizer and scheduler done in {(end - start) / 1e6:.2f}ms"
+            )
+
+        # Build helper models
+        with self.accelerator.main_process_first():
+            self.logger.info("Building helper models...")
+            start = time.monotonic_ns()
+            self._built_helper_model()
+            end = time.monotonic_ns()
+            self.logger.info(
+                f"Building helper models done in {(end - start) / 1e6:.2f}ms"
+            )
+
+        # Accelerator preparing
+        self.logger.info("Initializing accelerate...")
+        start = time.monotonic_ns()
+        for k in self.model:
+            self.model[k] = self.accelerator.prepare(self.model[k])
+        for k, v in self.optimizer.optimizers.items():
+            self.optimizer.optimizers[k] = self.accelerator.prepare(
+                self.optimizer.optimizers[k]
+            )
+            self.optimizer.schedulers[k] = self.accelerator.prepare(
+                self.optimizer.schedulers[k]
+            )
+        end = time.monotonic_ns()
+        self.logger.info(f"Initializing accelerate done in {(end - start) / 1e6:.2f}ms")
+
+        # Build criterions
+        with self.accelerator.main_process_first():
+            self.logger.info("Building criterion...")
+            start = time.monotonic_ns()
+            self.criterions = self._build_criterion()
+            end = time.monotonic_ns()
+            self.logger.info(f"Building criterion done in {(end - start) / 1e6:.2f}ms")
+
+        # Resume checkpoints
+        with self.accelerator.main_process_first():
+            self.checkpoint_dir = os.path.join(self.exp_dir, "checkpoint")
+            if args.resume_type:
+                self.logger.info("Resuming from checkpoint...")
+                start = time.monotonic_ns()
+                ckpt_path = Path(args.checkpoint)
+                if self._is_valid_pattern(ckpt_path.parts[-1]):
+                    ckpt_path = self._load_model(args.checkpoint, args.resume_type)
+                else:
+                    ckpt_path = self._load_model(
+                        args.checkpoint, resume_type=args.resume_type
+                    )
+                end = time.monotonic_ns()
+                self.logger.info(
+                    f"Resuming from checkpoint done in {(end - start) / 1e6:.2f}ms"
+                )
+                self.checkpoints_path = json.load(
+                    open(os.path.join(ckpt_path, "ckpts.json"), "r")
+                )
+
+            if self.accelerator.is_main_process:
+                os.makedirs(self.checkpoint_dir, exist_ok=True)
+            self.logger.debug(f"Checkpoint directory: {self.checkpoint_dir}")
+
+        # Save config
+        self.config_save_path = os.path.join(self.exp_dir, "args.json")
+
+    def _build_dataset(self):
+        return FAcodecDataset, FAcodecCollator
+
+    def _build_criterion(self):
+        criterions = dict()
+        stft_criterion = MultiScaleSTFTLoss()
+        mel_criterion = MelSpectrogramLoss(
+            n_mels=[5, 10, 20, 40, 80, 160, 320],
+            window_lengths=[32, 64, 128, 256, 512, 1024, 2048],
+            mel_fmin=[0, 0, 0, 0, 0, 0, 0],
+            mel_fmax=[None, None, None, None, None, None, None],
+            pow=1.0,
+            mag_weight=0.0,
+            clamp_eps=1e-5,
+        )
+        content_criterion = FocalLoss(gamma=2)
+        l1_criterion = L1Loss()
+        criterions["stft"] = stft_criterion
+        criterions["mel"] = mel_criterion
+        criterions["l1"] = l1_criterion
+        criterions["content"] = content_criterion
+
+        return criterions
+
+    def _build_model(self):
+        model = build_model(self.cfg.model_params)
+        _ = [model[key].to(self.accelerator.device) for key in model]
+        return model
+
+    def _built_helper_model(self):
+        device = self.accelerator.device
+        self.pitch_extractor = load_F0_models(self.cfg.F0_path).to(device)
+
+        # load model and processor
+        self.w2v_processor = Wav2Vec2Processor.from_pretrained(
+            "facebook/wav2vec2-xlsr-53-espeak-cv-ft"
+        )
+        self.w2v_model = Wav2Vec2ForCTC.from_pretrained(
+            "facebook/wav2vec2-xlsr-53-espeak-cv-ft"
+        ).to(device)
+        self.w2v_model.eval()
+
+        if nemo_asr is None:
+            self.speaker_model = None
+        else:
+            self.speaker_model = (
+                nemo_asr.models.EncDecSpeakerLabelModel.from_pretrained(
+                    "nvidia/speakerverification_en_titanet_large"
+                )
+            )
+            self.speaker_model = self.speaker_model.to(device)
+            self.speaker_model.eval()
+
+    def _build_optimizer(self):
+        scheduler_params = {
+            "warmup_steps": self.cfg.loss_params.warmup_steps,
+            "base_lr": self.cfg.loss_params.base_lr,
+        }
+        optimizer = build_optimizer(
+            {key: self.model[key] for key in self.model},
+            scheduler_params_dict={key: scheduler_params.copy() for key in self.model},
+            lr=float(scheduler_params["base_lr"]),
+        )
+
+        return optimizer
+
+    def train_loop(self):
+        """Training process"""
+        self.accelerator.wait_for_everyone()
+
+        # Dump config
+        if self.accelerator.is_main_process:
+            self._dump_cfg(self.config_save_path)
+        _ = [self.model[key].train() for key in self.model]
+        self.optimizer.zero_grad()
+
+        # Sync and start training
+        self.accelerator.wait_for_everyone()
+        while self.epoch < self.max_epoch:
+            self.logger.info("\n")
+            self.logger.info("-" * 32)
+            self.logger.info("Epoch {}: ".format(self.epoch))
+
+            # Train and Validate
+            train_total_loss, train_losses = self._train_epoch()
+            for key, loss in train_losses.items():
+                self.logger.info("  |- Train/{} Loss: {:.6f}".format(key, loss))
+                self.accelerator.log(
+                    {"Epoch/Train {} Loss".format(key): loss},
+                    step=self.epoch,
+                )
+            self.accelerator.log(
+                {
+                    "Epoch/Train Total Loss": train_total_loss,
+                },
+                step=self.epoch,
+            )
+
+            # Update scheduler
+            self.accelerator.wait_for_everyone()
+
+            # Check save checkpoint interval
+            run_eval = False
+            if self.accelerator.is_main_process:
+                save_checkpoint = False
+                for i, num in enumerate(self.save_checkpoint_stride):
+                    if self.epoch % num == 0:
+                        save_checkpoint = True
+                        run_eval |= self.run_eval[i]
+
+            # Save checkpoints
+            self.accelerator.wait_for_everyone()
+            if self.accelerator.is_main_process and save_checkpoint:
+                print("Saving..")
+                state = {
+                    "net": {key: self.model[key].state_dict() for key in self.model},
+                    "optimizer": self.optimizer.state_dict(),
+                    "scheduler": self.optimizer.scheduler_state_dict(),
+                    "iters": self.step,
+                    "epoch": self.epoch,
+                }
+                save_path = os.path.join(
+                    self.checkpoint_dir,
+                    "FAcodec_epoch_%05d_step_%05d.pth" % (self.epoch, self.iters),
+                )
+                torch.save(state, save_path)
+                json.dump(
+                    self.checkpoints_path,
+                    open(os.path.join(self.checkpoint_dir, "ckpts.json"), "w"),
+                    ensure_ascii=False,
+                    indent=4,
+                )
+
+            self.accelerator.wait_for_everyone()
+
+            self.epoch += 1
+
+        # Finish training
+        self.accelerator.wait_for_everyone()
+        if self.accelerator.is_main_process:
+            path = os.path.join(
+                self.checkpoint_dir,
+                "epoch-{:04d}_step-{:07d}".format(
+                    self.epoch,
+                    self.step,
+                ),
+            )
+            print("Saving..")
+            state = {
+                "net": {key: self.model[key].state_dict() for key in self.model},
+                "optimizer": self.optimizer.state_dict(),
+                "scheduler": self.optimizer.scheduler_state_dict(),
+                "iters": self.step,
+                "epoch": self.epoch,
+            }
+            save_path = os.path.join(
+                self.checkpoint_dir,
+                "FAcodec_epoch_%05d_step_%05d.pth" % (self.epoch, self.iters),
+            )
+            torch.save(state, save_path)
+
+    def _train_epoch(self):
+        """Training epoch. Should return average loss of a batch (sample) over
+        one epoch. See ``train_loop`` for usage.
+        """
+        _ = [self.model[key].train() for key in self.model]
+
+        epoch_losses: dict = {}
+        epoch_total_loss: int = 0
+
+        for batch in tqdm(
+            self.train_dataloader,
+            desc=f"Training Epoch {self.epoch}",
+            unit="batch",
+            colour="GREEN",
+            leave=False,
+            dynamic_ncols=True,
+            smoothing=0.04,
+            disable=not self.accelerator.is_main_process,
+        ):
+            # Get losses
+            total_loss, losses = self._train_step(batch)
+            self.batch_count += 1
+
+            # Log info
+            if self.batch_count % self.cfg.train.gradient_accumulation_step == 0:
+                self.accelerator.log(
+                    {
+                        "Step/Learning Rate": (
+                            self.optimizer.schedulers["encoder"].get_last_lr()[0]
+                            if self.step != 0
+                            else 0
+                        )
+                    },
+                    step=self.step,
+                )
+                for key, _ in losses.items():
+                    self.accelerator.log(
+                        {
+                            "Step/Train {} Loss".format(key): losses[key],
+                        },
+                        step=self.step,
+                    )
+
+                if not epoch_losses:
+                    epoch_losses = losses
+                else:
+                    for key, value in losses.items():
+                        epoch_losses[key] += value
+                epoch_total_loss += total_loss
+                self.step += 1
+
+        # Get and log total losses
+        self.accelerator.wait_for_everyone()
+        epoch_total_loss = (
+            epoch_total_loss
+            / len(self.train_dataloader)
+            * self.cfg.train.gradient_accumulation_step
+        )
+        for key in epoch_losses.keys():
+            epoch_losses[key] = (
+                epoch_losses[key]
+                / len(self.train_dataloader)
+                * self.cfg.train.gradient_accumulation_step
+            )
+        return epoch_total_loss, epoch_losses
+
+    def _train_step(self, data):
+        """Training forward step. Should return average loss of a sample over
+        one batch. Provoke ``_forward_step`` is recommended except for special case.
+        See ``_train_epoch`` for usage.
+        """
+        # Init losses
+        train_losses = {}
+        total_loss = 0
+
+        # Use input feature to get predictions
+        data = [b.to(self.accelerator.device, non_blocking=True) for b in data]
+        waves, mels, wave_lengths, mel_input_length = data
+
+        # extract semantic latent with w2v model
+        waves_16k = torchaudio.functional.resample(waves, 24000, 16000)
+        w2v_input = self.w2v_processor(
+            waves_16k, sampling_rate=16000, return_tensors="pt"
+        ).input_values.to(self.accelerator.device)
+        with torch.no_grad():
+            w2v_outputs = self.w2v_model(w2v_input.squeeze(0)).logits
+            predicted_ids = torch.argmax(w2v_outputs, dim=-1)
+            phone_ids = (
+                F.interpolate(
+                    predicted_ids.unsqueeze(0).float(), mels.size(-1), mode="nearest"
+                )
+                .long()
+                .squeeze(0)
+            )
+
+        # get clips
+        mel_seg_len = min(
+            [int(mel_input_length.min().item()), self.cfg.train.max_frame_len]
+        )
+
+        gt_mel_seg = []
+        wav_seg = []
+        w2v_seg = []
+
+        for bib in range(len(mel_input_length)):
+            mel_length = int(mel_input_length[bib].item())
+
+            random_start = (
+                np.random.randint(0, mel_length - mel_seg_len)
+                if mel_length != mel_seg_len
+                else 0
+            )
+            gt_mel_seg.append(mels[bib, :, random_start : random_start + mel_seg_len])
+
+            # w2v_seg.append(w2v_latent[bib, :, random_start:random_start + mel_seg_len])
+            w2v_seg.append(phone_ids[bib, random_start : random_start + mel_seg_len])
+
+            y = waves[bib][random_start * 300 : (random_start + mel_seg_len) * 300]
+
+            wav_seg.append(y.to(self.accelerator.device))
+
+        gt_mel_seg = torch.stack(gt_mel_seg).detach()
+
+        wav_seg = torch.stack(wav_seg).float().detach().unsqueeze(1)
+        w2v_seg = torch.stack(w2v_seg).float().detach()
+
+        with torch.no_grad():
+            real_norm = log_norm(gt_mel_seg.unsqueeze(1)).squeeze(1).detach()
+            F0_real, _, _ = self.pitch_extractor(gt_mel_seg.unsqueeze(1))
+
+        # normalize f0
+        # Remove unvoiced frames (replace with -1)
+        gt_glob_f0s = []
+        f0_targets = []
+        for bib in range(len(F0_real)):
+            voiced_indices = F0_real[bib] > 5.0
+            f0_voiced = F0_real[bib][voiced_indices]
+
+            if len(f0_voiced) != 0:
+                # Convert to log scale
+                log_f0 = f0_voiced.log2()
+
+                # Calculate mean and standard deviation
+                mean_f0 = log_f0.mean()
+                std_f0 = log_f0.std()
+
+                # Normalize the F0 sequence
+                normalized_f0 = (log_f0 - mean_f0) / std_f0
+
+                # Create the normalized F0 sequence with unvoiced frames
+                normalized_sequence = torch.zeros_like(F0_real[bib])
+                normalized_sequence[voiced_indices] = normalized_f0
+                normalized_sequence[~voiced_indices] = (
+                    -10
+                )  # Assign -10 to unvoiced frames
+
+                gt_glob_f0s.append(mean_f0)
+            else:
+                normalized_sequence = torch.zeros_like(F0_real[bib]) - 10.0
+                gt_glob_f0s.append(torch.tensor(0.0).to(self.accelerator.device))
+
+            # f0_targets.append(normalized_sequence[single_side_context // 200:-single_side_context // 200])
+            f0_targets.append(normalized_sequence)
+        f0_targets = torch.stack(f0_targets).to(self.accelerator.device)
+        # fill nan with -10
+        f0_targets[torch.isnan(f0_targets)] = -10.0
+        # fill inf with -10
+        f0_targets[torch.isinf(f0_targets)] = -10.0
+        # if frame_rate not equal to 80, interpolate f0 from frame rate of 80 to target frame rate
+        if self.cfg.preprocess_params.frame_rate != 80:
+            f0_targets = F.interpolate(
+                f0_targets.unsqueeze(1),
+                mel_seg_len // 80 * self.cfg.preprocess_params.frame_rate,
+                mode="nearest",
+            ).squeeze(1)
+            w2v_seg = F.interpolate(
+                w2v_seg,
+                mel_seg_len // 80 * self.cfg.preprocess_params.frame_rate,
+                mode="nearest",
+            )
+
+        wav_seg_input = wav_seg
+        wav_seg_target = wav_seg
+
+        z = self.model.encoder(wav_seg_input)
+        z, quantized, commitment_loss, codebook_loss, timbre = self.model.quantizer(
+            z, wav_seg_input, n_c=2, full_waves=waves, wave_lens=wave_lengths
+        )
+        preds, rev_preds = self.model.fa_predictors(quantized, timbre)
+
+        pred_wave = self.model.decoder(z)
+
+        len_diff = wav_seg_target.size(-1) - pred_wave.size(-1)
+        if len_diff > 0:
+            wav_seg_target = wav_seg_target[..., len_diff // 2 : -len_diff // 2]
+
+        # discriminator loss
+        d_fake = self.model.discriminator(pred_wave.detach())
+        d_real = self.model.discriminator(wav_seg_target)
+        loss_d = 0
+        for x_fake, x_real in zip(d_fake, d_real):
+            loss_d += torch.mean(x_fake[-1] ** 2)
+            loss_d += torch.mean((1 - x_real[-1]) ** 2)
+
+        self.optimizer.zero_grad()
+        self.accelerator.backward(loss_d)
+        grad_norm_d = torch.nn.utils.clip_grad_norm_(
+            self.model.discriminator.parameters(), 10.0
+        )
+        self.optimizer.step("discriminator")
+        self.optimizer.scheduler(key="discriminator")
+
+        # generator loss
+        signal = AudioSignal(wav_seg_target, sample_rate=24000)
+        recons = AudioSignal(pred_wave, sample_rate=24000)
+        stft_loss = self.criterions["stft"](recons, signal)
+        mel_loss = self.criterions["mel"](recons, signal)
+        waveform_loss = self.criterions["l1"](recons, signal)
+
+        d_fake = self.model.discriminator(pred_wave)
+        d_real = self.model.discriminator(wav_seg_target)
+
+        loss_g = 0
+        for x_fake in d_fake:
+            loss_g += torch.mean((1 - x_fake[-1]) ** 2)
+
+        loss_feature = 0
+
+        for i in range(len(d_fake)):
+            for j in range(len(d_fake[i]) - 1):
+                loss_feature += F.l1_loss(d_fake[i][j], d_real[i][j].detach())
+
+        pred_f0, pred_uv = preds["f0"], preds["uv"]
+        rev_pred_f0, rev_pred_uv = rev_preds["rev_f0"], rev_preds["rev_uv"]
+
+        common_min_size = min(pred_f0.size(-2), f0_targets.size(-1))
+        f0_targets = f0_targets[..., :common_min_size]
+        real_norm = real_norm[..., :common_min_size]
+
+        f0_loss = F.smooth_l1_loss(
+            f0_targets, pred_f0.squeeze(-1)[..., :common_min_size]
+        )
+        uv_loss = F.smooth_l1_loss(
+            real_norm, pred_uv.squeeze(-1)[..., :common_min_size]
+        )
+        rev_f0_loss = (
+            F.smooth_l1_loss(f0_targets, rev_pred_f0.squeeze(-1)[..., :common_min_size])
+            if rev_pred_f0 is not None
+            else torch.FloatTensor([0]).to(self.accelerator.device)
+        )
+        rev_uv_loss = (
+            F.smooth_l1_loss(real_norm, rev_pred_uv.squeeze(-1)[..., :common_min_size])
+            if rev_pred_uv is not None
+            else torch.FloatTensor([0]).to(self.accelerator.device)
+        )
+
+        tot_f0_loss = f0_loss + rev_f0_loss
+        tot_uv_loss = uv_loss + rev_uv_loss
+
+        pred_content = preds["content"]
+        rev_pred_content = rev_preds["rev_content"]
+
+        target_content_latents = w2v_seg[..., :common_min_size]
+
+        content_loss = self.criterions["content"](
+            pred_content.transpose(1, 2)[..., :common_min_size],
+            target_content_latents.long(),
+        )
+        rev_content_loss = (
+            self.criterions["content"](
+                rev_pred_content.transpose(1, 2)[..., :common_min_size],
+                target_content_latents.long(),
+            )
+            if rev_pred_content is not None
+            else torch.FloatTensor([0]).to(self.accelerator.device)
+        )
+
+        tot_content_loss = content_loss + rev_content_loss
+
+        if self.speaker_model is not None:
+            spk_logits = torch.cat(
+                [
+                    self.speaker_model.infer_segment(w16.cpu()[..., :wl])[1]
+                    for w16, wl in zip(waves_16k, wave_lengths)
+                ],
+                dim=0,
+            )
+            spk_labels = spk_logits.argmax(dim=-1)
+        else:
+            spk_labels = torch.zeros([len(waves_16k)], dtype=torch.long).to(
+                self.accelerator.device
+            )
+
+        spk_pred_logits = preds["timbre"]
+        spk_loss = F.cross_entropy(spk_pred_logits, spk_labels)
+        x_spk_pred_logits = rev_preds["x_timbre"]
+
+        x_spk_loss = (
+            F.cross_entropy(x_spk_pred_logits, spk_labels)
+            if x_spk_pred_logits is not None
+            else torch.FloatTensor([0]).to(self.accelerator.device)
+        )
+
+        tot_spk_loss = spk_loss + x_spk_loss
+
+        loss_gen_all = (
+            mel_loss * 15.0
+            + loss_feature * 1.0
+            + loss_g * 1.0
+            + commitment_loss * 0.25
+            + codebook_loss * 1.0
+            + tot_f0_loss * 1.0
+            + tot_uv_loss * 1.0
+            + tot_content_loss * 5.0
+            + tot_spk_loss * 5.0
+        )
+
+        self.optimizer.zero_grad()
+        self.accelerator.backward(loss_gen_all)
+
+        with torch.no_grad():
+            total_loss = loss_gen_all.item()
+            train_losses["stft"] = stft_loss.item()
+            train_losses["mel"] = mel_loss.item()
+            train_losses["l1"] = waveform_loss.item()
+            train_losses["f0"] = f0_loss.item()
+            train_losses["uv"] = uv_loss.item()
+            train_losses["content"] = content_loss.item()
+            train_losses["speaker"] = spk_loss.item()
+            train_losses["rev_f0"] = rev_f0_loss.item()
+            train_losses["rev_uv"] = rev_uv_loss.item()
+            train_losses["rev_content"] = rev_content_loss.item()
+            train_losses["rev_speaker"] = x_spk_loss.item()
+
+            train_losses["feature"] = loss_feature.item()
+            train_losses["generator"] = loss_g.item()
+            train_losses["commitment"] = commitment_loss.item()
+            train_losses["codebook"] = codebook_loss.item()
+
+            # discriminators
+            train_losses["discriminator"] = loss_d.item()
+
+        return total_loss, train_losses
+
+    def _inference(self, eval_wave):
+        """Inference during training for test audios."""
+        z = self.model.encoder(
+            eval_wave[None, None, ...].to(self.accelerator.device).float()
+        )
+        z, quantized, commitment_loss, codebook_loss, timbre = self.model.quantizer(
+            z, eval_wave[None, None, ...], n_c=self.cfg.model_params.n_c_codebooks
+        )
+        full_pred_wave = self.model.decoder(z)
+        return full_pred_wave[0]
+
+    def _load_model(self, checkpoint_path=None, resume_type="resume"):
+        """Load model from checkpoint. If checkpoint_path is None, it will
+        load the latest checkpoint in checkpoint_dir. If checkpoint_path is not
+        None, it will load the checkpoint specified by checkpoint_path. **Only use this
+        method after** ``accelerator.prepare()``.
+        """
+        if resume_type == "resume":
+            if checkpoint_path is None:
+                available_checkpoints = glob.glob(
+                    os.path.join(self.checkpoint_dir, "FAcodc_epoch_*_step_*.pth")
+                )
+                # find the checkpoint that has the highest step number
+                latest_checkpoint = max(
+                    available_checkpoints,
+                    key=lambda x: int(x.split("_")[-1].split(".")[0]),
+                )
+                earliest_checkpoint = min(
+                    available_checkpoints,
+                    key=lambda x: int(x.split("_")[-1].split(".")[0]),
+                )
+                # delete the earliest checkpoint
+                if (
+                    earliest_checkpoint != latest_checkpoint
+                    and self.accelerator.is_main_process
+                    and len(available_checkpoints) > 4
+                ):
+                    os.remove(earliest_checkpoint)
+                    print(f"Removed {earliest_checkpoint}")
+            else:
+                latest_checkpoint = checkpoint_path
+
+            self.model, self.optimizer, self.epoch, self.step = load_checkpoint(
+                self.model,
+                self.optimizer,
+                latest_checkpoint,
+                load_only_params=False,
+                ignore_modules=[],
+                is_distributed=self.accelerator.num_processes > 1,
+            )
+
+        else:
+            raise ValueError("Invalid resume type")
+        return checkpoint_path
+
+    def _count_parameters(self):
+        total_num = sum(
+            sum(p.numel() for p in self.model[key].parameters()) for key in self.model
+        )
+        # trainable_num = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+        return total_num