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()