from pytorch_lightning.loggers import WandbLogger, CometLogger from ..interface.aeiou import pca_point_cloud import math import wandb import torch import torch.nn.functional as F import os import typing as tp def get_rank(): """Get rank of current process.""" if "SLURM_PROCID" in os.environ: return int(os.environ["SLURM_PROCID"]) if not torch.distributed.is_available() or not torch.distributed.is_initialized(): return 0 return torch.distributed.get_rank() class InverseLR(torch.optim.lr_scheduler._LRScheduler): """Implements an inverse decay learning rate schedule with an optional exponential warmup. When last_epoch=-1, sets initial lr as lr. inv_gamma is the number of steps/epochs required for the learning rate to decay to (1 / 2)**power of its original value. Args: optimizer (Optimizer): Wrapped optimizer. inv_gamma (float): Inverse multiplicative factor of learning rate decay. Default: 1. power (float): Exponential factor of learning rate decay. Default: 1. warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable) Default: 0. final_lr (float): The final learning rate. Default: 0. last_epoch (int): The index of last epoch. Default: -1. """ def __init__(self, optimizer, inv_gamma=1., power=1., warmup=0., final_lr=0., last_epoch=-1): self.inv_gamma = inv_gamma self.power = power if not 0. <= warmup < 1: raise ValueError('Invalid value for warmup') self.warmup = warmup self.final_lr = final_lr super().__init__(optimizer, last_epoch) def get_lr(self): if not self._get_lr_called_within_step: import warnings warnings.warn("To get the last learning rate computed by the scheduler, " "please use `get_last_lr()`.") return self._get_closed_form_lr() def _get_closed_form_lr(self): warmup = 1 - self.warmup ** (self.last_epoch + 1) lr_mult = (1 + self.last_epoch / self.inv_gamma) ** -self.power return [warmup * max(self.final_lr, base_lr * lr_mult) for base_lr in self.base_lrs] def create_optimizer_from_config(optimizer_config, parameters): """Create optimizer from config. Args: parameters (iterable): parameters to optimize. optimizer_config (dict): optimizer config. Returns: torch.optim.Optimizer: optimizer. """ optimizer_type = optimizer_config["type"] if optimizer_type == "FusedAdam": from deepspeed.ops.adam import FusedAdam optimizer = FusedAdam(parameters, **optimizer_config["config"]) elif optimizer_type == "CAdamW": from stable_audio_tools.training.optims import CAdamW optimizer = CAdamW(parameters, **optimizer_config["config"]) elif optimizer_type == "CLion": from stable_audio_tools.training.optims import CLion optimizer = CLion(parameters, **optimizer_config["config"]) elif optimizer_type == "AdamW8bit": from bitsandbytes.optim import AdamW8bit optimizer = AdamW8bit(parameters, **optimizer_config["config"]) elif optimizer_type == "MuonAdamW": from stable_audio_tools.training.optims import MuonAdamW optimizer = MuonAdamW(parameters, **optimizer_config["config"]) else: optimizer_fn = getattr(torch.optim, optimizer_type) optimizer = optimizer_fn(parameters, **optimizer_config["config"]) return optimizer def create_scheduler_from_config(scheduler_config, optimizer): """Create scheduler from config. Args: scheduler_config (dict): scheduler config. optimizer (torch.optim.Optimizer): optimizer. Returns: torch.optim.lr_scheduler._LRScheduler: scheduler. """ if scheduler_config["type"] == "InverseLR": scheduler_fn = InverseLR else: scheduler_fn = getattr(torch.optim.lr_scheduler, scheduler_config["type"]) scheduler = scheduler_fn(optimizer, **scheduler_config["config"]) return scheduler def logger_project_name(logger) -> str: if isinstance(logger, WandbLogger): return logger.experiment.project elif isinstance(logger, CometLogger): return logger.name def log_metric(logger, key, value, step=None): from pytorch_lightning.loggers import WandbLogger, CometLogger if isinstance(logger, WandbLogger): logger.experiment.log({key: value}) elif isinstance(logger, CometLogger): logger.experiment.log_metrics({key: value}, step=step) def log_audio(logger, key, audio_path, sample_rate, caption=None, step=None): if isinstance(logger, WandbLogger): logger.experiment.log({key: wandb.Audio(audio_path, sample_rate=sample_rate, caption=caption)}) elif isinstance(logger, CometLogger): logger.experiment.log_audio(audio_path, file_name=key, sample_rate=sample_rate, step=step) def log_image(logger, key, img_data, step=None): if isinstance(logger, WandbLogger): logger.experiment.log({key: wandb.Image(img_data)}) elif isinstance(logger, CometLogger): logger.experiment.log_image(img_data, name=key, step=step) def log_point_cloud(logger, key, tokens, caption=None): if isinstance(logger, WandbLogger): point_cloud = pca_point_cloud(tokens) logger.experiment.log({key: point_cloud}) elif isinstance(logger, CometLogger): point_cloud = pca_point_cloud(tokens, rgb_float=True, output_type="points") #logger.experiment.log_points_3d(scene_name=key, points=point_cloud) def compute_per_elem_trim(conditioning, sample_rate, margin_seconds=5.0): """Compute per-element trim lengths from seconds_total in conditioning dicts. Returns a list of trim lengths (in audio samples) or None if no seconds_total found. """ if not any("seconds_total" in c for c in conditioning): return None margin_samples = int(margin_seconds * sample_rate) per_elem_trim = [] for c in conditioning: if "seconds_total" in c: per_elem_trim.append(int(c["seconds_total"] * sample_rate) + margin_samples) else: per_elem_trim.append(None) return per_elem_trim def trim_and_concat(x, per_elem_trim): """Per-element trim and concatenate along time axis. Trims each batch element to its own length (from seconds_total + margin), removing trailing padding silence before concatenation. Args: x: (b, d, n) tensor or list of (d, n) tensors per_elem_trim: list of trim lengths per element, or None for no trimming """ items = [x[i] for i in range(x.shape[0])] if isinstance(x, torch.Tensor) and x.dim() == 3 else x if per_elem_trim is None: return torch.cat(items, dim=-1) parts = [] for i, elem in enumerate(items): if per_elem_trim[i] is not None: parts.append(elem[..., :min(per_elem_trim[i], elem.shape[-1])]) else: parts.append(elem) return torch.cat(parts, dim=-1) class StaggeredLogger: """Accumulates log values over N steps and flushes averaged metrics. Avoids per-step CUDA syncs (.item()) and prevents PL from caching stale values between flush intervals by logging directly to the experiment logger. Usage in a PL module:: # In __init__: self._staggered_logger = StaggeredLogger(every_n_steps=10) # At end of training_step (pass detached tensors or Python scalars): self._staggered_logger.log(log_dict, self) """ def __init__(self, every_n_steps: int = 10): self.every_n_steps = every_n_steps self._accum: dict = {} self._counts: dict = {} self._steps: int = 0 def log(self, log_dict: dict, module) -> None: """Accumulate *log_dict* and flush every *every_n_steps* calls. NaN/inf values are handled at flush time: if an accumulated sum is non-finite (due to a NaN step poisoning the window), that metric is simply omitted from the flush. This avoids the expensive per-step CUDA syncs that an eager ``isfinite`` check would require. Args: log_dict: Mapping of metric names to detached GPU tensors or Python scalars. Tensors are kept on-device until flush. module: The ``pl.LightningModule`` instance (used for ``log_dict``, ``logger``, and ``global_step``). """ for k, v in log_dict.items(): val = v.detach() if torch.is_tensor(v) else v if k not in self._accum: self._accum[k] = val self._counts[k] = 1 else: self._accum[k] = self._accum[k] + val self._counts[k] += 1 self._steps += 1 if self._steps >= self.every_n_steps: flushed = {} for k, v in self._accum.items(): c = self._counts[k] if c > 0: val = v.item() / c if torch.is_tensor(v) else v / c # Skip non-finite values (NaN from a bad step poisoned this window) if math.isfinite(val): flushed[k] = val # Update progress bar only (logger=False avoids PL caching stale values) if flushed: module.log_dict(flushed, prog_bar=True, on_step=True, logger=False) # Log directly to experiment (bypasses PL metric caching) if module.logger is not None: module.logger.log_metrics(flushed, step=module.global_step) self._accum = {} self._counts = {} self._steps = 0 def resize_padding_mask(padding_mask: torch.Tensor, target_length: int) -> torch.Tensor: """Resize a padding mask to target_length using ceiling-based length scaling. Unlike F.interpolate(mode="nearest"), this ensures any target position that partially overlaps valid audio is marked valid (rounds up). Args: padding_mask: (B, T) or (T,) tensor where 1/True = valid, 0/False = padding target_length: desired output length Returns: Resized boolean mask with same ndim as input """ if padding_mask.ndim == 1: valid_length = padding_mask.sum() source_length = padding_mask.shape[0] valid_target_length = torch.ceil( valid_length.float() * target_length / source_length ).long().clamp(max=target_length) positions = torch.arange(target_length, device=padding_mask.device) return positions < valid_target_length else: valid_lengths = padding_mask.sum(dim=-1) # (B,) source_length = padding_mask.shape[-1] valid_target_lengths = torch.ceil( valid_lengths.float() * target_length / source_length ).long().clamp(max=target_length) positions = torch.arange(target_length, device=padding_mask.device).unsqueeze(0) return positions < valid_target_lengths.unsqueeze(1) def create_augmented_padding_mask( padding_mask: torch.Tensor, silence_extension_scale_seconds: float = 0.0, sample_rate: int = 44100, downsampling_ratio: int = 1, ) -> torch.Tensor: """ Augment padding mask by randomly extending the valid region to include silence. This helps the model learn to handle silence at the end of audio by randomly including some of the padded (silent) region as valid. Uses an exponential distribution so most extensions are small (within ~scale seconds) but can occasionally extend to the full sequence length. Args: padding_mask: Boolean tensor of shape (batch, seq_len) where True = valid audio silence_extension_scale_seconds: Scale (mean) of exponential distribution in seconds. ~63% of samples within 1x scale, ~86% within 2x, ~95% within 3x. sample_rate: Audio sample rate (used to convert seconds to tokens) downsampling_ratio: Downsampling ratio from pretransform (to convert samples to latent tokens) Returns: Augmented padding mask of shape (batch, seq_len) where True = valid region (audio + extended silence) """ if silence_extension_scale_seconds <= 0: return padding_mask batch_size, seq_len = padding_mask.shape device = padding_mask.device # Find the last valid position for each batch element valid_lengths = padding_mask.sum(dim=-1) # (batch,) # Calculate scale in tokens for exponential distribution scale_tokens = silence_extension_scale_seconds * sample_rate / downsampling_ratio # Sample from exponential distribution (biased toward smaller values) # Exponential: ~63% within 1x scale, ~86% within 2x, ~95% within 3x random_extensions = torch.empty(batch_size, device=device).exponential_(lambd=1.0/scale_tokens) random_extensions = random_extensions.long() # Extend valid lengths, clamping to sequence length (allows extension to full padding) augmented_lengths = torch.clamp(valid_lengths + random_extensions, max=seq_len) # Create augmented mask: True where position < augmented_length positions = torch.arange(seq_len, device=device).unsqueeze(0) # (1, seq_len) augmented_mask = positions < augmented_lengths.unsqueeze(1) # (batch, seq_len) return augmented_mask def compute_normalized_mse( pred: torch.Tensor, target: torch.Tensor, loss_mask: torch.Tensor, loss_normalization: str = "none", loss_norm_eps: float = 1e-6, ) -> torch.Tensor: """ Compute MSE normalized by detached target magnitude. Prevents high-magnitude latents from dominating the loss, ensuring quiet passages receive equal learning signal. Normalization is only applied to valid (non-padded) regions. Padded regions retain raw MSE for compatibility with mask_loss_weight. Args: pred: [B, C, T] predicted tensor target: [B, C, T] target tensor loss_mask: [B, T] boolean, True = valid (non-padded), or None loss_normalization: One of "none", "timestep", "sample", "sample_channel" loss_norm_eps: Epsilon for numerical stability Returns: Per-element MSE [B, C, T], normalized for signal, raw for padding """ mse = (pred - target) ** 2 # [B, C, T] if loss_normalization == "none": return mse if loss_mask is None: mask_expanded = torch.ones(pred.shape[0], 1, pred.shape[2], device=pred.device, dtype=torch.bool) else: mask_expanded = loss_mask.unsqueeze(1) # [B, 1, T] with torch.no_grad(): if loss_normalization == "timestep": # Var across channels at each timestep [B, 1, T] mag_sq = torch.mean((target - torch.mean(target, dim=1, keepdim=True)) ** 2, dim=1, keepdim=True) + loss_norm_eps else: # For sample/sample_channel modes, exclude padding via NaN masking masked_targets = torch.where(mask_expanded, target, float('nan')) if loss_normalization == "sample": sample_mean = torch.nanmean(masked_targets, dim=(1, 2), keepdim=True) mag_sq = torch.nanmean((masked_targets - sample_mean) ** 2, dim=(1, 2), keepdim=True) + loss_norm_eps elif loss_normalization == "sample_channel": channel_mean = torch.nanmean(masked_targets, dim=2, keepdim=True) mag_sq = torch.nanmean((masked_targets - channel_mean) ** 2, dim=2, keepdim=True) + loss_norm_eps else: raise ValueError(f"Unknown loss normalization mode: {loss_normalization}") # Handle edge case where entire sample is padding (all NaN -> NaN variance) mag_sq = torch.where(torch.isnan(mag_sq), torch.ones_like(mag_sq), mag_sq) # Only normalize signal regions; keep padding MSE raw for mask_loss_weight compatibility normalized_mse = mse / mag_sq return torch.where(mask_expanded, normalized_mse, mse) def compute_masked_loss( loss_full: torch.Tensor, loss_mask: torch.Tensor, mask_padding_attention: bool, mask_loss_weight: float = 0.0, ) -> tuple: """ Compute loss with separate signal and padding contributions. Args: loss_full: Full loss tensor of shape (B, C, T) loss_mask: Boolean mask of shape (B, T) where True = signal, False = padding mask_padding_attention: If True, only compute loss on signal region mask_loss_weight: Weight for padding loss when mask_padding_attention is False Returns: Tuple of (loss, signal_mean, padding_mean) where: - loss: The final scalar loss - signal_mean: Mean loss on signal region (for logging) - padding_mean: Mean loss on padding region (for logging) """ # Compute separate signal and padding losses signal = torch.where(loss_mask.unsqueeze(1), loss_full, 0.0) signal_sum = signal.sum(dim=(1, 2)) n_channels = loss_full.shape[1] signal_count = loss_mask.sum(dim=1) * n_channels padding = torch.where(~loss_mask.unsqueeze(1), loss_full, 0.0) padding_sum = padding.sum(dim=(1, 2)) padding_count = (~loss_mask).sum(dim=1) * n_channels if mask_padding_attention: # When attention masking is on, only compute loss on signal region # (padding saw no attention, so loss there is meaningless) per_sample_loss = signal_sum / (signal_count + 1e-8) loss = per_sample_loss.mean() else: # Weighted loss: signal + mask_loss_weight * padding w = mask_loss_weight denom = signal_count + w * padding_count + 1e-8 per_sample_loss = (signal_sum + w * padding_sum) / denom loss = per_sample_loss.mean() # Compute separate means for logging signal_mean = signal.sum() / (signal_count.sum() + 1e-8) padding_mean = padding.sum() / (padding_count.sum() + 1e-8) return loss, signal_mean.detach(), padding_mean.detach() def masked_mean( values: torch.Tensor, mask: torch.Tensor, ) -> torch.Tensor: """ Compute mean over masked (valid) positions only. Args: values: Tensor of shape (B, C, T) or (B, T) mask: Boolean mask of shape (B, T) where True = valid, False = padding Returns: Scalar mean over valid positions only """ if values.ndim == 2: # (B, T) -> add channel dim values = values.unsqueeze(1) # Expand mask to match values shape: (B, T) -> (B, 1, T) mask_expanded = mask.unsqueeze(1) # Zero out padding positions masked_values = torch.where(mask_expanded, values, torch.zeros_like(values)) # Count valid positions (per sample, times channels) n_channels = values.shape[1] valid_count = mask.sum(dim=1) * n_channels # (B,) # Sum over all dimensions and divide by valid count per_sample_sum = masked_values.sum(dim=(1, 2)) # (B,) per_sample_mean = per_sample_sum / (valid_count + 1e-8) return per_sample_mean.mean() def masked_sum( values: torch.Tensor, mask: torch.Tensor, dim: tp.List[int], ) -> torch.Tensor: """ Compute sum over masked (valid) positions only. Args: values: Tensor of shape (B, C, T) or (B, T) mask: Boolean mask of shape (B, T) where True = valid, False = padding dim: Dimensions to sum over (should include the time dimension) Returns: Sum over valid positions, reduced over specified dims """ if values.ndim == 2: # (B, T) -> add channel dim for consistency mask_expanded = mask else: # (B, C, T) -> expand mask mask_expanded = mask.unsqueeze(1) # Zero out padding positions masked_values = torch.where(mask_expanded, values, torch.zeros_like(values)) return masked_values.sum(dim=tuple(dim))