import torch import torch.nn as nn import pytorch_lightning as pl from typing import Optional, Literal from ema_pytorch import EMA from torch.nn import Parameter from einops import rearrange from stable_audio_tools.models import create_model_from_config from stable_audio_tools.models.autoencoders import AudioAutoencoder from stable_audio_tools.models.discriminators import ( EncodecDiscriminator, OobleckDiscriminator, DACGANLoss, ) from stable_audio_tools.models.bottleneck import ( VAEBottleneck, RVQBottleneck, DACRVQBottleneck, DACRVQVAEBottleneck, RVQVAEBottleneck, WassersteinBottleneck, ) from stable_audio_tools.training.losses import ( MelSpectrogramLoss, MultiLoss, AuralossLoss, ValueLoss, L1Loss, LossWithTarget, MSELoss, HubertLoss, # PESQMetric, # TODO move PESQ here? ) from stable_audio_tools.training.losses import auraloss as auraloss from stable_audio_tools.training.utils import ( create_optimizer_from_config, create_scheduler_from_config, log_metric, ) from .ctc_loss import CTCLossModule, PERModule def trim_to_shortest(a, b): """Trim the longer of two tensors to the length of the shorter one.""" if a.shape[-1] > b.shape[-1]: return a[:,:,:b.shape[-1]], b elif b.shape[-1] > a.shape[-1]: return a, b[:,:,:a.shape[-1]] return a, b class ProjectionHead(nn.Module): def __init__(self, latent_dim, proj_head_dim, mid_dim=256): super(ProjectionHead, self).__init__() self.proj_head = nn.Sequential( nn.Tanh(), nn.Linear(latent_dim, mid_dim), nn.ReLU(), nn.Linear(mid_dim, mid_dim), nn.ReLU(), nn.Linear(mid_dim, proj_head_dim) ) def forward(self, x): return self.proj_head(x) class AutoencoderTrainingWrapper(pl.LightningModule): def __init__(self, autoencoder: AudioAutoencoder, loss_config: dict, eval_loss_config: dict, optimizer_configs: dict, sample_rate: int = 16000, lr: float = 1e-4, warmup_steps: int = 0, warmup_mode: Literal["adv", "full"] = "adv", encoder_freeze_on_warmup: bool = False, use_ema: bool = True, ema_copy = None, force_input_mono = False, latent_mask_ratio = 0.0, teacher_model: Optional[AudioAutoencoder] = None, clip_grad_norm = 0.0, encoder_mask_ratio = 0.0, use_ctc: bool = False, proj_head_dim: Optional[int] = None, detach_proj_head: bool = False, ): super().__init__() self.automatic_optimization = False self.autoencoder = autoencoder self.warmed_up = False self.warmup_steps = warmup_steps self.warmup_mode = warmup_mode self.encoder_freeze_on_warmup = encoder_freeze_on_warmup self.lr = lr self.clip_grad_norm = clip_grad_norm self.force_input_mono = force_input_mono self.teacher_model = teacher_model self.use_ctc = use_ctc self.proj_head_dim = proj_head_dim self.detach_proj_head = detach_proj_head self.projection_head = ( ProjectionHead(self.autoencoder.latent_dim, self.proj_head_dim) if self.use_ctc and self.proj_head_dim is not None else nn.Identity() ) self.optimizer_configs = optimizer_configs self.loss_config = loss_config # Spectral reconstruction loss self.sdstft = auraloss.MultiResolutionSTFTLoss( sample_rate=sample_rate, **loss_config['spectral']['config']) # Discriminator self.use_disc = True if 'discriminator' in loss_config else False self.discriminator = None if self.use_disc: if loss_config['discriminator']['type'] == 'oobleck': self.discriminator = OobleckDiscriminator(**loss_config['discriminator']['config']) elif loss_config['discriminator']['type'] == 'encodec': self.discriminator = EncodecDiscriminator( in_channels=self.autoencoder.out_channels, **loss_config['discriminator']['config']) elif loss_config['discriminator']['type'] == 'dac': self.discriminator = DACGANLoss( channels=self.autoencoder.out_channels, sample_rate=sample_rate, **loss_config['discriminator']['config']) gen_loss_modules = [] if self.use_disc: # Discriminator loss. self.losses_disc = MultiLoss([ ValueLoss(key='loss_dis', weight=1.0, name='discriminator_loss'), ]) # Adversarial and feature matching losses. gen_loss_modules += [ ValueLoss( key='loss_adv', weight=self.loss_config['discriminator']['weights']['adversarial'], name='loss_adv'), ValueLoss( key='feature_matching_distance', weight=self.loss_config['discriminator']['weights']['feature_matching'], name='feature_matching_loss'), ] # Reconstruction loss gen_loss_modules += [AuralossLoss(self.sdstft, target_key='reals', input_key='decoded', name='mrstft_loss', weight=self.loss_config['spectral']['weights']['mrstft'], decay=self.loss_config['spectral'].get('decay', 1.0), )] if "mrmel" in loss_config: mrmel_weight = loss_config["mrmel"]["weights"]["mrmel"] if mrmel_weight > 0: mrmel_config = loss_config["mrmel"]["config"] self.mrmel = MelSpectrogramLoss(sample_rate, n_mels=mrmel_config["n_mels"], window_lengths=mrmel_config["window_lengths"], pow=mrmel_config["pow"], log_weight=mrmel_config["log_weight"], mag_weight=mrmel_config["mag_weight"], ) gen_loss_modules.append(LossWithTarget( self.mrmel, "reals", "decoded", name="mrmel_loss", weight=mrmel_weight, )) if "hubert" in loss_config: hubert_weight = loss_config["hubert"]["weights"]["hubert"] if hubert_weight > 0: hubert_cfg = ( loss_config["hubert"]["config"] if "config" in loss_config["hubert"] else dict() ) self.hubert = HubertLoss(weight=1.0, **hubert_cfg) gen_loss_modules.append(LossWithTarget( self.hubert, target_key = "reals", input_key = "decoded", name="hubert_loss", weight=hubert_weight, decay = loss_config["hubert"].get("decay", 1.0) )) if "l1" in loss_config["time"]["weights"]: if self.loss_config['time']['weights']['l1'] > 0.0: gen_loss_modules.append(L1Loss( key_a='reals', key_b='decoded', weight=self.loss_config['time']['weights']['l1'], name='l1_time_loss', decay = self.loss_config['time'].get('decay', 1.0), )) if "l2" in loss_config["time"]["weights"]: if self.loss_config['time']['weights']['l2'] > 0.0: gen_loss_modules.append(MSELoss( key_a='reals', key_b='decoded', weight=self.loss_config['time']['weights']['l2'], name='l2_time_loss', decay = self.loss_config['time'].get('decay', 1.0), )) if self.autoencoder.bottleneck is not None: gen_loss_modules += create_loss_modules_from_bottleneck( self.autoencoder.bottleneck, self.loss_config) self.encoder_mask_ratio = encoder_mask_ratio if encoder_mask_ratio > 0.0: gen_loss_modules.append(L1Loss( key_a='detached_latents', key_b='masked_latents', weight=1.0, name='encoder_mask_loss', decay = 1.0, )) if "ctc" in loss_config: ctc_weight = loss_config["ctc"]["weights"]["ctc"] if ctc_weight > 0: gen_loss_modules.append(CTCLossModule( name = "ctc_loss", target_key = "ctc_tgt", input_key = "log_probs", weight = ctc_weight, decay = loss_config["ctc"].get("decay", 1.0), blank_idx = loss_config["ctc"].get("blank_idx", 80) )) self.losses_gen = MultiLoss(gen_loss_modules) # Set up EMA for model weights self.autoencoder_ema = None self.use_ema = use_ema if self.use_ema: self.autoencoder_ema = EMA( self.autoencoder, ema_model=ema_copy, beta=0.9999, power=3/4, update_every=1, update_after_step=1 ) self.latent_mask_ratio = latent_mask_ratio # evaluation losses & metrics self.eval_losses = torch.nn.ModuleDict() if eval_loss_config is not None: # if "pesq" in eval_loss_config: # self.eval_losses["pesq"] = PESQMetric(sample_rate) if "stft"in eval_loss_config: self.eval_losses["stft"] = auraloss.STFTLoss(**eval_loss_config["stft"]) if "sisdr" in eval_loss_config: self.eval_losses["sisdr"] = auraloss.SISDRLoss(**eval_loss_config["sisdr"]) if "mel" in eval_loss_config: self.eval_losses["mel"] = auraloss.MelSTFTLoss( sample_rate, **eval_loss_config["mel"]) if "per" in eval_loss_config: self.eval_losses["per"] = PERModule( target_key = "ctc_tgt", input_key = "log_probs", blank_idx = loss_config["ctc"].get("blank_idx", 80)) self.validation_step_outputs = [] def configure_optimizers(self): gen_params = list(self.autoencoder.parameters()) if not self.use_disc: opt_gen = create_optimizer_from_config( self.optimizer_configs['autoencoder']['optimizer'], gen_params) if "scheduler" in self.optimizer_configs['autoencoder']: sched_gen = create_scheduler_from_config( self.optimizer_configs['autoencoder']['scheduler'], opt_gen) return [opt_gen], [sched_gen] return [opt_gen] # Using discriminator. opt_gen = create_optimizer_from_config( self.optimizer_configs['autoencoder']['optimizer'], gen_params) opt_disc = create_optimizer_from_config( self.optimizer_configs['discriminator']['optimizer'], self.discriminator.parameters()) use_scheduler = ( "scheduler" in self.optimizer_configs['autoencoder'] and "scheduler" in self.optimizer_configs['discriminator'] ) if use_scheduler: sched_gen = create_scheduler_from_config( self.optimizer_configs['autoencoder']['scheduler'], opt_gen) sched_disc = create_scheduler_from_config( self.optimizer_configs['discriminator']['scheduler'], opt_disc) return [opt_gen, opt_disc], [sched_gen, sched_disc] return [opt_gen, opt_disc] def forward(self, reals): latents, encoder_info = self.autoencoder.encode(reals, return_info=True) decoded = self.autoencoder.decode(latents) return decoded def validation_step(self, batch, batch_idx): reals, _ = batch # Remove extra dimension added by WebDataset if reals.ndim == 4 and reals.shape[0] == 1: reals = reals[0] if len(reals.shape) == 2: reals = reals.unsqueeze(1) loss_info = {} loss_info["reals"] = reals encoder_input = reals if self.force_input_mono and encoder_input.shape[1] > 1: encoder_input = encoder_input.mean(dim=1, keepdim=True) loss_info["encoder_input"] = encoder_input with torch.no_grad(): if self.use_ctc: latents, encoder_info = self.autoencoder.encode(encoder_input, return_info=True) continuous_latents = encoder_info["pre_bottleneck_latents"] proj_features = rearrange(continuous_latents, "b c n -> b n c") proj_features = self.projection_head( proj_features.detach() if self.detach_proj_head else proj_features ) loss_info['log_probs'] = proj_features loss_info['ctc_tgt'] = batch[1] else: latents, encoder_info = self.autoencoder.encode(encoder_input, return_info=True) loss_info["latents"] = latents loss_info.update(encoder_info) decoded = self.autoencoder.decode(latents) #Trim output to remove post-padding. decoded, reals = trim_to_shortest(decoded, reals) # Run evaluation metrics. val_loss_dict = {} for eval_key, eval_fn in self.eval_losses.items(): if eval_key == 'per': loss_value = eval_fn(loss_info) else: loss_value = eval_fn(decoded, reals) if eval_key == "sisdr": loss_value = -loss_value if isinstance(loss_value, torch.Tensor): loss_value = loss_value.item() val_loss_dict[eval_key] = loss_value self.validation_step_outputs.append(val_loss_dict) return val_loss_dict def on_validation_epoch_end(self): sum_loss_dict = {} for loss_dict in self.validation_step_outputs: for key, value in loss_dict.items(): if key not in sum_loss_dict: sum_loss_dict[key] = value else: sum_loss_dict[key] += value for key, value in sum_loss_dict.items(): val_loss = value / len(self.validation_step_outputs) val_loss = self.all_gather(val_loss).mean().item() log_metric(self.logger, f"val/{key}", val_loss) self.validation_step_outputs.clear() # free memory def training_step(self, batch, batch_idx): reals, _ = batch log_dict = {} # Remove extra dimension added by WebDataset if reals.ndim == 4 and reals.shape[0] == 1: reals = reals[0] if len(reals.shape) == 2: reals = reals.unsqueeze(1) if self.global_step >= self.warmup_steps: self.warmed_up = True loss_info = {} loss_info["reals"] = reals encoder_input = reals if self.force_input_mono and encoder_input.shape[1] > 1: encoder_input = encoder_input.mean(dim=1, keepdim=True) loss_info["encoder_input"] = encoder_input data_std = encoder_input.std() if self.warmed_up and self.encoder_freeze_on_warmup: with torch.no_grad(): latents, encoder_info = self.autoencoder.encode(encoder_input, return_info=True) else: if self.use_ctc: latents, encoder_info = self.autoencoder.encode(encoder_input, return_info=True) continuous_latents = encoder_info["pre_bottleneck_latents"] proj_features = rearrange(continuous_latents, "b c n -> b n c") proj_features = self.projection_head( proj_features.detach() if self.detach_proj_head else proj_features ) loss_info['log_probs'] = proj_features loss_info['ctc_tgt'] = batch[1] else: latents, encoder_info = self.autoencoder.encode(encoder_input, return_info=True) if self.encoder_mask_ratio > 0.0: masked_latents = self.autoencoder.encode( encoder_input, return_info=False, encoder_mask_ratio=self.encoder_mask_ratio) detached_latents = latents.detach() loss_info["masked_latents"] = masked_latents loss_info["detached_latents"] = detached_latents loss_info["latents"] = latents loss_info.update(encoder_info) # Encode with teacher model for distillation if self.teacher_model is not None: with torch.no_grad(): teacher_latents = self.teacher_model.encode(encoder_input, return_info=False) loss_info['teacher_latents'] = teacher_latents # Optionally mask out some latents for noise resistance if self.latent_mask_ratio > 0.0: mask = torch.rand_like(latents) < self.latent_mask_ratio latents = torch.where(mask, torch.zeros_like(latents), latents) decoded = self.autoencoder.decode(latents) #Trim output to remove post-padding decoded, reals = trim_to_shortest(decoded, reals) loss_info["decoded"] = decoded loss_info["reals"] = reals if self.autoencoder.out_channels == 2: loss_info["decoded_left"] = decoded[:, 0:1, :] loss_info["decoded_right"] = decoded[:, 1:2, :] loss_info["reals_left"] = reals[:, 0:1, :] loss_info["reals_right"] = reals[:, 1:2, :] # Distillation if self.teacher_model is not None: with torch.no_grad(): teacher_decoded = self.teacher_model.decode(teacher_latents) own_latents_teacher_decoded = self.teacher_model.decode(latents) #Distilled model's latents decoded by teacher teacher_latents_own_decoded = self.autoencoder.decode(teacher_latents) #Teacher's latents decoded by distilled model loss_info['teacher_decoded'] = teacher_decoded loss_info['own_latents_teacher_decoded'] = own_latents_teacher_decoded loss_info['teacher_latents_own_decoded'] = teacher_latents_own_decoded if self.use_disc: if self.warmed_up: loss_dis, loss_adv, feature_matching_distance = self.discriminator.loss(reals=reals, fakes=decoded) else: loss_adv = torch.tensor(0.).to(reals) feature_matching_distance = torch.tensor(0.).to(reals) if self.warmup_mode == "adv": loss_dis, _, _ = self.discriminator.loss(reals=reals, fakes=decoded) else: loss_dis = torch.tensor(0.0).to(reals) loss_info["loss_dis"] = loss_dis loss_info["loss_adv"] = loss_adv loss_info["feature_matching_distance"] = feature_matching_distance opt_gen = None opt_disc = None if self.use_disc: opt_gen, opt_disc = self.optimizers() else: opt_gen = self.optimizers() lr_schedulers = self.lr_schedulers() sched_gen = None sched_disc = None if lr_schedulers is not None: if self.use_disc: sched_gen, sched_disc = lr_schedulers else: sched_gen = lr_schedulers # Train the discriminator use_disc = ( self.use_disc and self.global_step % 2 # Check warmup mode and if it is time to use discriminator. and ( (self.warmup_mode == "full" and self.warmed_up) or self.warmup_mode == "adv") ) if use_disc: loss, losses = self.losses_disc(loss_info) log_dict['train/disc_lr'] = opt_disc.param_groups[0]['lr'] opt_disc.zero_grad() self.manual_backward(loss) if self.clip_grad_norm > 0.0: torch.nn.utils.clip_grad_norm_( self.discriminator.parameters(), self.clip_grad_norm) opt_disc.step() if sched_disc is not None: # sched step every step sched_disc.step() # Train the generator else: loss, losses = self.losses_gen(loss_info) if self.use_ema: self.autoencoder_ema.update() opt_gen.zero_grad() self.manual_backward(loss) if self.clip_grad_norm > 0.0: torch.nn.utils.clip_grad_norm_( self.autoencoder.parameters(), self.clip_grad_norm) opt_gen.step() if sched_gen is not None: # scheduler step every step sched_gen.step() log_dict['train/loss'] = loss.detach().item() log_dict['train/latent_std'] = latents.std().detach().item() log_dict['train/data_std'] = data_std.detach().item() log_dict['train/gen_lr'] = opt_gen.param_groups[0]['lr'] for loss_name, loss_value in losses.items(): log_dict[f'train/{loss_name}'] = loss_value.detach().item() self.log_dict(log_dict, prog_bar=True, on_step=True) return loss def export_model(self, path, use_safetensors=False): if self.autoencoder_ema is not None: model = self.autoencoder_ema.ema_model else: model = self.autoencoder if use_safetensors: save_model(model, path) else: torch.save({"state_dict": model.state_dict()}, path) def create_loss_modules_from_bottleneck(bottleneck, loss_config): losses = [] if ( isinstance(bottleneck, VAEBottleneck) or isinstance(bottleneck, DACRVQVAEBottleneck) or isinstance(bottleneck, RVQVAEBottleneck) ): try: kl_weight = loss_config['bottleneck']['weights']['kl'] except: kl_weight = 1e-6 kl_loss = ValueLoss(key='kl', weight=kl_weight, name='kl_loss') losses.append(kl_loss) if ( isinstance(bottleneck, RVQBottleneck) or isinstance(bottleneck, RVQVAEBottleneck) ): quantizer_loss = ValueLoss(key='quantizer_loss', weight=1.0, name='quantizer_loss') losses.append(quantizer_loss) if isinstance(bottleneck, DACRVQBottleneck) or isinstance(bottleneck, DACRVQVAEBottleneck): codebook_loss = ValueLoss(key='vq/codebook_loss', weight=1.0, name='codebook_loss') commitment_loss = ValueLoss(key='vq/commitment_loss', weight=0.25, name='commitment_loss') losses.append(codebook_loss) losses.append(commitment_loss) if isinstance(bottleneck, WassersteinBottleneck): try: mmd_weight = loss_config['bottleneck']['weights']['mmd'] except: mmd_weight = 100 mmd_loss = ValueLoss(key='mmd', weight=mmd_weight, name='mmd_loss') losses.append(mmd_loss) return losses def create_training_wrapper_from_config(model_config, model): model_type = model_config.get('model_type', None) assert model_type is not None, 'model_type must be specified in model config' training_config = model_config.get('training', None) assert training_config is not None, 'training config must be specified in model config' ema_copy = None if training_config.get("use_ema", False): ema_copy = create_model_from_config(model_config) # Copy each weight to the ema copy for name, param in model.state_dict().items(): if isinstance(param, Parameter): # backwards compatibility for serialized parameters param = param.data ema_copy.state_dict()[name].copy_(param) use_ema = training_config.get("use_ema", False) latent_mask_ratio = training_config.get("latent_mask_ratio", 0.0) teacher_model = training_config.get("teacher_model", None) if teacher_model is not None: teacher_model = create_model_from_config(teacher_model) teacher_model = teacher_model.eval().requires_grad_(False) teacher_model_ckpt = training_config.get("teacher_model_ckpt", None) if teacher_model_ckpt is not None: teacher_model.load_state_dict(torch.load(teacher_model_ckpt)["state_dict"]) else: raise ValueError("teacher_model_ckpt must be specified if teacher_model is specified") return AutoencoderTrainingWrapper( model, lr=training_config.get("learning_rate", None), warmup_steps=training_config.get("warmup_steps", 0), encoder_freeze_on_warmup=training_config.get("encoder_freeze_on_warmup", False), sample_rate=model_config["sample_rate"], loss_config=training_config.get("loss_configs", None), eval_loss_config=training_config.get("eval_loss_configs", None), optimizer_configs=training_config.get("optimizer_configs", None), use_ema=use_ema, ema_copy=ema_copy if use_ema else None, force_input_mono=training_config.get("force_input_mono", False), latent_mask_ratio=latent_mask_ratio, teacher_model=teacher_model, encoder_mask_ratio=training_config.get("encoder_mask_ratio", 0.0), use_ctc=training_config.get("use_ctc", False), proj_head_dim=model_config["model"].get("proj_head_dim", False), detach_proj_head=model_config["model"].get("detach_proj_head", None), )