repo_description: so-vits-svc fork with realtime support, improved interface and more features.
repo_topics: contentvec deep-learning gan hubert lightning pytorch pytorch-lightning realtime so-vits-svc softvc sovits speech-synthesis vits voice-changer voice-conversion
code: from __future__ import annotations
  
  import os
  import warnings
  from logging import getLogger
  from multiprocessing import cpu_count
  from pathlib import Path
  from typing import Any
  
  import lightning.pytorch as pl
  import torch
  from lightning.pytorch.accelerators import MPSAccelerator, TPUAccelerator
  from lightning.pytorch.callbacks import DeviceStatsMonitor
  from lightning.pytorch.loggers import TensorBoardLogger
  from lightning.pytorch.strategies.ddp import DDPStrategy
  from lightning.pytorch.tuner import Tuner
  from torch.cuda.amp import autocast
  from torch.nn import functional as F
  from torch.utils.data import DataLoader
  from torch.utils.tensorboard.writer import SummaryWriter
  
  import so_vits_svc_fork.f0
  import so_vits_svc_fork.modules.commons as commons
  import so_vits_svc_fork.utils
  
  from . import utils
  from .dataset import TextAudioCollate, TextAudioDataset
  from .logger import is_notebook
  from .modules.descriminators import MultiPeriodDiscriminator
  from .modules.losses import discriminator_loss, feature_loss, generator_loss, kl_loss
  from .modules.mel_processing import mel_spectrogram_torch
  from .modules.synthesizers import SynthesizerTrn
  
  LOG = getLogger(__name__)
  torch.set_float32_matmul_precision("high")
  
  
  class VCDataModule(pl.LightningDataModule):
      batch_size: int
  
      def __init__(self, hparams: Any):
          super().__init__()
          self.__hparams = hparams
          self.batch_size = hparams.train.batch_size
          if not isinstance(self.batch_size, int):
              self.batch_size = 1
          self.collate_fn = TextAudioCollate()
  
          # these should be called in setup(), but we need to calculate check_val_every_n_epoch
          self.train_dataset = TextAudioDataset(self.__hparams, is_validation=False)
          self.val_dataset = TextAudioDataset(self.__hparams, is_validation=True)
  
      def train_dataloader(self):
          return DataLoader(
              self.train_dataset,
              num_workers=min(cpu_count(), self.__hparams.train.get("num_workers", 8)),
              batch_size=self.batch_size,
              collate_fn=self.collate_fn,
              persistent_workers=True,
          )
  
      def val_dataloader(self):
          return DataLoader(
              self.val_dataset,
              batch_size=1,
              collate_fn=self.collate_fn,
          )
  
  
  def train(config_path: Path | str, model_path: Path | str, reset_optimizer: bool = False):
      config_path = Path(config_path)
      model_path = Path(model_path)
  
      hparams = utils.get_backup_hparams(config_path, model_path)
      utils.ensure_pretrained_model(
          model_path,
          hparams.model.get(
              "pretrained",
              {
                  "D_0.pth": "https://huggingface.co/therealvul/so-vits-svc-4.0-init/resolve/main/D_0.pth",
                  "G_0.pth": "https://huggingface.co/therealvul/so-vits-svc-4.0-init/resolve/main/G_0.pth",
              },
          ),
      )
  
      datamodule = VCDataModule(hparams)
      strategy = (
          ("ddp_find_unused_parameters_true" if os.name != "nt" else DDPStrategy(find_unused_parameters=True, process_group_backend="gloo"))
          if torch.cuda.device_count() > 1
          else "auto"
      )
      LOG.info(f"Using strategy: {strategy}")
      trainer = pl.Trainer(
          logger=TensorBoardLogger(model_path, "lightning_logs", hparams.train.get("log_version", 0)),
          # profiler="simple",
          val_check_interval=hparams.train.eval_interval,
          max_epochs=hparams.train.epochs,
          check_val_every_n_epoch=None,
          precision=("16-mixed" if hparams.train.fp16_run else "bf16-mixed" if hparams.train.get("bf16_run", False) else 32),
          strategy=strategy,
          callbacks=([pl.callbacks.RichProgressBar()] if not is_notebook() else []) + [DeviceStatsMonitor()],
          benchmark=True,
          enable_checkpointing=False,
      )
      tuner = Tuner(trainer)
      model = VitsLightning(reset_optimizer=reset_optimizer, **hparams)
  
      # automatic batch size scaling
      batch_size = hparams.train.batch_size
      batch_split = str(batch_size).split("-")
      batch_size = batch_split[0]
      init_val = 2 if len(batch_split) <= 1 else int(batch_split[1])
      max_trials = 25 if len(batch_split) <= 2 else int(batch_split[2])
      if batch_size == "auto":
          batch_size = "binsearch"
      if batch_size in ["power", "binsearch"]:
          model.tuning = True
          tuner.scale_batch_size(
              model,
              mode=batch_size,
              datamodule=datamodule,
              steps_per_trial=1,
              init_val=init_val,
              max_trials=max_trials,
          )
          model.tuning = False
      else:
          batch_size = int(batch_size)
      # automatic learning rate scaling is not supported for multiple optimizers
      """if hparams.train.learning_rate  == "auto":
      lr_finder = tuner.lr_find(model)
      LOG.info(lr_finder.results)
      fig = lr_finder.plot(suggest=True)
      fig.savefig(model_path / "lr_finder.png")"""
  
      trainer.fit(model, datamodule=datamodule)
  
  
  class VitsLightning(pl.LightningModule):
      def __init__(self, reset_optimizer: bool = False, **hparams: Any):
          super().__init__()
          self._temp_epoch = 0  # Add this line to initialize the _temp_epoch attribute
          self.save_hyperparameters("reset_optimizer")
          self.save_hyperparameters(*[k for k in hparams.keys()])
          torch.manual_seed(self.hparams.train.seed)
          self.net_g = SynthesizerTrn(
              self.hparams.data.filter_length // 2 + 1,
              self.hparams.train.segment_size // self.hparams.data.hop_length,
              **self.hparams.model,
          )
          self.net_d = MultiPeriodDiscriminator(self.hparams.model.use_spectral_norm)
          self.automatic_optimization = False
          self.learning_rate = self.hparams.train.learning_rate
          self.optim_g = torch.optim.AdamW(
              self.net_g.parameters(),
              self.learning_rate,
              betas=self.hparams.train.betas,
              eps=self.hparams.train.eps,
          )
          self.optim_d = torch.optim.AdamW(
              self.net_d.parameters(),
              self.learning_rate,
              betas=self.hparams.train.betas,
              eps=self.hparams.train.eps,
          )
          self.scheduler_g = torch.optim.lr_scheduler.ExponentialLR(self.optim_g, gamma=self.hparams.train.lr_decay)
          self.scheduler_d = torch.optim.lr_scheduler.ExponentialLR(self.optim_d, gamma=self.hparams.train.lr_decay)
          self.optimizers_count = 2
          self.load(reset_optimizer)
          self.tuning = False
  
      def on_train_start(self) -> None:
          if not self.tuning:
              self.set_current_epoch(self._temp_epoch)
              total_batch_idx = self._temp_epoch * len(self.trainer.train_dataloader)
              self.set_total_batch_idx(total_batch_idx)
              global_step = total_batch_idx * self.optimizers_count
              self.set_global_step(global_step)
  
          # check if using tpu or mps
          if isinstance(self.trainer.accelerator, (TPUAccelerator, MPSAccelerator)):
              # patch torch.stft to use cpu
              LOG.warning("Using TPU/MPS. Patching torch.stft to use cpu.")
  
              def stft(
                  input: torch.Tensor,
                  n_fft: int,
                  hop_length: int | None = None,
                  win_length: int | None = None,
                  window: torch.Tensor | None = None,
                  center: bool = True,
                  pad_mode: str = "reflect",
                  normalized: bool = False,
                  onesided: bool | None = None,
                  return_complex: bool | None = None,
              ) -> torch.Tensor:
                  device = input.device
                  input = input.cpu()
                  if window is not None:
                      window = window.cpu()
                  return torch.functional.stft(
                      input,
                      n_fft,
                      hop_length,
                      win_length,
                      window,
                      center,
                      pad_mode,
                      normalized,
                      onesided,
                      return_complex,
                  ).to(device)
  
              torch.stft = stft
  
          elif "bf" in self.trainer.precision:
              LOG.warning("Using bf. Patching torch.stft to use fp32.")
  
              def stft(
                  input: torch.Tensor,
                  n_fft: int,
                  hop_length: int | None = None,
                  win_length: int | None = None,
                  window: torch.Tensor | None = None,
                  center: bool = True,
                  pad_mode: str = "reflect",
                  normalized: bool = False,
                  onesided: bool | None = None,
                  return_complex: bool | None = None,
              ) -> torch.Tensor:
                  dtype = input.dtype
                  input = input.float()
                  if window is not None:
                      window = window.float()
                  return torch.functional.stft(
                      input,
                      n_fft,
                      hop_length,
                      win_length,
                      window,
                      center,
                      pad_mode,
                      normalized,
                      onesided,
                      return_complex,
                  ).to(dtype)
  
              torch.stft = stft
  
      def on_train_end(self) -> None:
          self.save_checkpoints(adjust=0)
  
      def save_checkpoints(self, adjust=1):
          if self.tuning or self.trainer.sanity_checking:
              return
  
          # only save checkpoints if we are on the main device
          if hasattr(self.device, "index") and self.device.index != None and self.device.index != 0:
              return
  
          # `on_train_end` will be the actual epoch, not a -1, so we have to call it with `adjust = 0`
          current_epoch = self.current_epoch + adjust
          total_batch_idx = self.total_batch_idx - 1 + adjust
  
          utils.save_checkpoint(
              self.net_g,
              self.optim_g,
              self.learning_rate,
              current_epoch,
              Path(self.hparams.model_dir) / f"G_{total_batch_idx if self.hparams.train.get('ckpt_name_by_step', False) else current_epoch}.pth",
          )
          utils.save_checkpoint(
              self.net_d,
              self.optim_d,
              self.learning_rate,
              current_epoch,
              Path(self.hparams.model_dir) / f"D_{total_batch_idx if self.hparams.train.get('ckpt_name_by_step', False) else current_epoch}.pth",
          )
          keep_ckpts = self.hparams.train.get("keep_ckpts", 0)
          if keep_ckpts > 0:
              utils.clean_checkpoints(
                  path_to_models=self.hparams.model_dir,
                  n_ckpts_to_keep=keep_ckpts,
                  sort_by_time=True,
              )
  
      def set_current_epoch(self, epoch: int):
          LOG.info(f"Setting current epoch to {epoch}")
          self.trainer.fit_loop.epoch_progress.current.completed = epoch
          self.trainer.fit_loop.epoch_progress.current.processed = epoch
          assert self.current_epoch == epoch, f"{self.current_epoch} != {epoch}"
  
      def set_global_step(self, global_step: int):
          LOG.info(f"Setting global step to {global_step}")
          self.trainer.fit_loop.epoch_loop.manual_optimization.optim_step_progress.total.completed = global_step
          self.trainer.fit_loop.epoch_loop.automatic_optimization.optim_progress.optimizer.step.total.completed = global_step
          assert self.global_step == global_step, f"{self.global_step} != {global_step}"
  
      def set_total_batch_idx(self, total_batch_idx: int):
          LOG.info(f"Setting total batch idx to {total_batch_idx}")
          self.trainer.fit_loop.epoch_loop.batch_progress.total.ready = total_batch_idx + 1
          self.trainer.fit_loop.epoch_loop.batch_progress.total.completed = total_batch_idx
          assert self.total_batch_idx == total_batch_idx + 1, f"{self.total_batch_idx} != {total_batch_idx + 1}"
  
      @property
      def total_batch_idx(self) -> int:
          return self.trainer.fit_loop.epoch_loop.total_batch_idx + 1
  
      def load(self, reset_optimizer: bool = False):
          latest_g_path = utils.latest_checkpoint_path(self.hparams.model_dir, "G_*.pth")
          latest_d_path = utils.latest_checkpoint_path(self.hparams.model_dir, "D_*.pth")
          if latest_g_path is not None and latest_d_path is not None:
              try:
                  _, _, _, epoch = utils.load_checkpoint(
                      latest_g_path,
                      self.net_g,
                      self.optim_g,
                      reset_optimizer,
                  )
                  _, _, _, epoch = utils.load_checkpoint(
                      latest_d_path,
                      self.net_d,
                      self.optim_d,
                      reset_optimizer,
                  )
                  self._temp_epoch = epoch
                  self.scheduler_g.last_epoch = epoch - 1
                  self.scheduler_d.last_epoch = epoch - 1
              except Exception as e:
                  raise RuntimeError("Failed to load checkpoint") from e
          else:
              LOG.warning("No checkpoint found. Start from scratch.")
  
      def configure_optimizers(self):
          return [self.optim_g, self.optim_d], [self.scheduler_g, self.scheduler_d]
  
      def log_image_dict(self, image_dict: dict[str, Any], dataformats: str = "HWC") -> None:
          if not isinstance(self.logger, TensorBoardLogger):
              warnings.warn("Image logging is only supported with TensorBoardLogger.")
              return
          writer: SummaryWriter = self.logger.experiment
          for k, v in image_dict.items():
              try:
                  writer.add_image(k, v, self.total_batch_idx, dataformats=dataformats)
              except Exception as e:
                  warnings.warn(f"Failed to log image {k}: {e}")
  
      def log_audio_dict(self, audio_dict: dict[str, Any]) -> None:
          if not isinstance(self.logger, TensorBoardLogger):
              warnings.warn("Audio logging is only supported with TensorBoardLogger.")
              return
          writer: SummaryWriter = self.logger.experiment
          for k, v in audio_dict.items():
              writer.add_audio(
                  k,
                  v.float(),
                  self.total_batch_idx,
                  sample_rate=self.hparams.data.sampling_rate,
              )
  
      def log_dict_(self, log_dict: dict[str, Any], **kwargs) -> None:
          if not isinstance(self.logger, TensorBoardLogger):
              warnings.warn("Logging is only supported with TensorBoardLogger.")
              return
          writer: SummaryWriter = self.logger.experiment
          for k, v in log_dict.items():
              writer.add_scalar(k, v, self.total_batch_idx)
          kwargs["logger"] = False
          self.log_dict(log_dict, **kwargs)
  
      def log_(self, key: str, value: Any, **kwargs) -> None:
          self.log_dict_({key: value}, **kwargs)
  
      def training_step(self, batch: dict[str, torch.Tensor], batch_idx: int) -> None:
          self.net_g.train()
          self.net_d.train()
  
          # get optims
          optim_g, optim_d = self.optimizers()
  
          # Generator
          # train
          self.toggle_optimizer(optim_g)
          c, f0, spec, mel, y, g, lengths, uv = batch
          (
              y_hat,
              y_hat_mb,
              ids_slice,
              z_mask,
              (z, z_p, m_p, logs_p, m_q, logs_q),
              pred_lf0,
              norm_lf0,
              lf0,
          ) = self.net_g(c, f0, uv, spec, g=g, c_lengths=lengths, spec_lengths=lengths)
          y_mel = commons.slice_segments(
              mel,
              ids_slice,
              self.hparams.train.segment_size // self.hparams.data.hop_length,
          )
          y_hat_mel = mel_spectrogram_torch(y_hat.squeeze(1), self.hparams)
          y_mel = y_mel[..., : y_hat_mel.shape[-1]]
          y = commons.slice_segments(
              y,
              ids_slice * self.hparams.data.hop_length,
              self.hparams.train.segment_size,
          )
          y = y[..., : y_hat.shape[-1]]
  
          # generator loss
          y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = self.net_d(y, y_hat)
  
          with autocast(enabled=False):
              loss_mel = F.l1_loss(y_mel, y_hat_mel) * self.hparams.train.c_mel
              loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * self.hparams.train.c_kl
              loss_fm = feature_loss(fmap_r, fmap_g)
              loss_gen, losses_gen = generator_loss(y_d_hat_g)
              loss_lf0 = F.mse_loss(pred_lf0, lf0)
              loss_gen_all = loss_gen + loss_fm + loss_mel + loss_kl + loss_lf0
  
              # MB-iSTFT-VITS
              loss_subband = torch.tensor(0.0)
              if self.hparams.model.get("type_") == "mb-istft":
                  from .modules.decoders.mb_istft import PQMF, subband_stft_loss
  
                  y_mb = PQMF(y.device, self.hparams.model.subbands).analysis(y)
                  loss_subband = subband_stft_loss(self.hparams, y_mb, y_hat_mb)
              loss_gen_all += loss_subband
  
          # log loss
          self.log_("lr", self.optim_g.param_groups[0]["lr"])
          self.log_dict_(
              {
                  "loss/g/total": loss_gen_all,
                  "loss/g/fm": loss_fm,
                  "loss/g/mel": loss_mel,
                  "loss/g/kl": loss_kl,
                  "loss/g/lf0": loss_lf0,
              },
              prog_bar=True,
          )
          if self.hparams.model.get("type_") == "mb-istft":
              self.log_("loss/g/subband", loss_subband)
          if self.total_batch_idx % self.hparams.train.log_interval == 0:
              self.log_image_dict(
                  {
                      "slice/mel_org": utils.plot_spectrogram_to_numpy(y_mel[0].data.cpu().float().numpy()),
                      "slice/mel_gen": utils.plot_spectrogram_to_numpy(y_hat_mel[0].data.cpu().float().numpy()),
                      "all/mel": utils.plot_spectrogram_to_numpy(mel[0].data.cpu().float().numpy()),
                      "all/lf0": so_vits_svc_fork.utils.plot_data_to_numpy(
                          lf0[0, 0, :].cpu().float().numpy(),
                          pred_lf0[0, 0, :].detach().cpu().float().numpy(),
                      ),
                      "all/norm_lf0": so_vits_svc_fork.utils.plot_data_to_numpy(
                          lf0[0, 0, :].cpu().float().numpy(),
                          norm_lf0[0, 0, :].detach().cpu().float().numpy(),
                      ),
                  }
              )
  
          accumulate_grad_batches = self.hparams.train.get("accumulate_grad_batches", 1)
          should_update = (batch_idx + 1) % accumulate_grad_batches == 0 or self.trainer.is_last_batch
          # optimizer
          self.manual_backward(loss_gen_all / accumulate_grad_batches)
          if should_update:
              self.log_("grad_norm_g", commons.clip_grad_value_(self.net_g.parameters(), None))
              optim_g.step()
              optim_g.zero_grad()
          self.untoggle_optimizer(optim_g)
  
          # Discriminator
          # train
          self.toggle_optimizer(optim_d)
          y_d_hat_r, y_d_hat_g, _, _ = self.net_d(y, y_hat.detach())
  
          # discriminator loss
          with autocast(enabled=False):
              loss_disc, losses_disc_r, losses_disc_g = discriminator_loss(y_d_hat_r, y_d_hat_g)
              loss_disc_all = loss_disc
  
          # log loss
          self.log_("loss/d/total", loss_disc_all, prog_bar=True)
  
          # optimizer
          self.manual_backward(loss_disc_all / accumulate_grad_batches)
          if should_update:
              self.log_("grad_norm_d", commons.clip_grad_value_(self.net_d.parameters(), None))
              optim_d.step()
              optim_d.zero_grad()
          self.untoggle_optimizer(optim_d)
  
          # end of epoch
          if self.trainer.is_last_batch:
              self.scheduler_g.step()
              self.scheduler_d.step()
  
      def validation_step(self, batch, batch_idx):
          # avoid logging with wrong global step
          if self.global_step == 0:
              return
          with torch.no_grad():
              self.net_g.eval()
              c, f0, _, mel, y, g, _, uv = batch
              y_hat = self.net_g.infer(c, f0, uv, g=g)
              y_hat_mel = mel_spectrogram_torch(y_hat.squeeze(1).float(), self.hparams)
              self.log_audio_dict({f"gen/audio_{batch_idx}": y_hat[0], f"gt/audio_{batch_idx}": y[0]})
              self.log_image_dict(
                  {
                      "gen/mel": utils.plot_spectrogram_to_numpy(y_hat_mel[0].cpu().float().numpy()),
                      "gt/mel": utils.plot_spectrogram_to_numpy(mel[0].cpu().float().numpy()),
                  }
              )
  
      def on_validation_end(self) -> None:
          self.save_checkpoints()
  