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")
[docs]
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)
[docs]
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,
)
[docs]
def val_dataloader(self):
return DataLoader(
self.val_dataset,
batch_size=1,
collate_fn=self.collate_fn,
)
[docs]
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)
[docs]
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
[docs]
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
[docs]
def on_train_end(self) -> None:
self.save_checkpoints(adjust=0)
[docs]
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,
)
[docs]
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}"
[docs]
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}"
[docs]
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
[docs]
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.")
[docs]
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}")
[docs]
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,
)
[docs]
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)
[docs]
def log_(self, key: str, value: Any, **kwargs) -> None:
self.log_dict_({key: value}, **kwargs)
[docs]
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()
[docs]
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()),
}
)
[docs]
def on_validation_end(self) -> None:
self.save_checkpoints()