Accumulate *log_dict* and flush every *every_n_steps* calls. NaN/inf values are handled at flush time: if an accumulated sum is non-f

Wrapper for calling above two routines at once, does Mel scale; Modified from PyTorch tutorial https://pytorch.org/tutorials/beginner/audio_feature_ex

Perform STFT. Args: x (Tensor): Input signal tensor (B, T). Returns: Tensor: x_mag, x_phs Mag

Unified sampling function for diffusion models. Handles all diffusion objectives, varlen support (padding_mask + effective_seq_len), timestep

Per-element trim and concatenate along time axis. Trims each batch element to its own length (from seconds_total + margin), removing trailing

Add LoRA parametrization to layers in a model. Args: model: The model to add LoRA to. lora_config: Dict mapping nn.Module types t

Resize a padding mask to target_length using ceiling-based length scaling. Unlike F.interpolate(mode="nearest"), this ensures any target position

Create optimizer from config. Args: parameters (iterable): parameters to optimize. optimizer_config (dict): optimizer config.

Create scheduler from config. Args: scheduler_config (dict): scheduler config. optimizer (torch.optim.Optimizer): optimizer.

for visualizing embeddings in a spectrogram-like way

Load state_dict to model, but only for keys that match exactly. Args: model (nn.Module): model to load state_dict. state_dict (Or

Generates random inpainting masks for a batch of latent audio sequences. The output inpainting mask has 0 where data should be inpainted, and

Trim the longer of two tensors to the length of the shorter one.

Run discriminator forward and return scores with filtered hidden states.

Encode audios into latents. Audios should already be preprocesed by preprocess_audio_for_encoder. If chunked is True, split the audio

Converts a denoiser output to a Karras ODE derivative.

Compute per-element trim lengths from seconds_total in conditioning dicts. Returns a list of trim lengths (in audio samples) or None if no second

Returns the scaling factors for the clean image (alpha) and for the noise (sigma), given a timestep.

Sample timesteps using the given noise distribution and apply distribution shift. Args: noise_dist_fn: Callable that sam

Memory-efficient temporal masking using in-place operations

Enforce deterministic sign convention: largest-magnitude element of each U column is positive.

Compute frequency-domain masking threshold X: [B, F, T] complex spectrum signal_rms: [B, 1, 1] signal RMS Returns:

Get the true original weight, even when the layer is already parametrized.

Compute STFT

Apply padding handling based on padding_mode. Args: embeddings: [batch, seq_len, dim] - the embeddings to process

Build a timestep schedule for diffusion sampling. Returns a 1D tensor of shape (N,) where N = steps+1 (if include_endpoint) or steps (if not)

Compute loss with separate signal and padding contributions. Args: loss_full: Full loss tensor of shape (B, C, T) loss_mask:

Compute MSE normalized by detached target magnitude. Prevents high-magnitude latents from dominating the loss, ensuring quiet passages r

Decode latents to audio. If chunked is True, split the latents into chunks of a given maximum size chunk_size, with given overlap, b