Return d_model as a clean multiple of head_dim. Uses a saturating sqrt-law: d_raw = k * sqrt(c), but performs rounding and clamping in u

Recursively wrap modules matching the given policy. Uses depth-first traversal to wrap child modules before their parents. Must be calle

apply a function to LoRAParametrization layers, designed to be used with model.apply

Extract audio embeddings using any combination of CLAP, PANNs, and VGGish. Args: id2audio: dict mapping id to audio tensor on GPU (s

very fast `glob` alternative. from https://stackoverflow.com/a/59803793/4259243

Calculates the noise level (sigma_down) to step down to and the amount of noise to add (sigma_up) when doing an ancestral sampling step.

Get trainable parameters that are NOT managed by FSDP. These parameters need explicit gradient synchronization via sync_non_fsdp_gradients()

Perform backward pass and sync gradients for non-FSDP trainable parameters. FSDP only synchronizes gradients for wrapped modules. No

Compute the sum of Gaussian kernels (with different bandwidths) between two flat tensors. Args: x: Tensor of shape (B, D) y:

torch.multinomial with arbitrary number of dimensions, and number of candidates on the last dimension. Args: input (torch.Tensor): The in

Preprocess single audio tensor (Channels x Length) to be compatible with the encoder. If the model is mono, stereo audio will be conv

Draws samples from a model given starting noise. v-diffusion DDIM. Args: sigmas: Pre-computed schedule tensor of shape (steps,).

Save a LoRA checkpoint in safetensors format with config as metadata. The lora_config dict is JSON-serialized and stored under the metadata key

Transform t∈[0,1] to log-SNR-spaced t with adaptive bounds. Maps through: logsnr = logsnr_end - t * (logsnr_end - logsnr_start)

shape: shape of the output tensor left_trunc: left truncation point, fraction of probability to be discarded right_trunc: right truncati

FSDP-aware EMA update using summon_full_params. The ema_pytorch library captures parameter names at initialization. When the online mode

Apply local additive and modular local conditioning to x.

Builds scatter indexes required to retrieve the original multi-codebook sequence from interleaving pattern. Args: sequenc

Compute Binaural Masking Level Difference Returns: unmasking factor [B, F, T]

Derive score masks from padding_mask and actual score tensor shapes. Uses adaptive_max_pool1d for downsampled outputs so that a score positio

Create a distribution shift object from config options.

Parity demod for hop=N/2. Returns ±1 with shape [1,1,F,M] (if expand_bc) for unambiguous broadcast against X[...,F,M]. Numerically exact for h

Expand bracket notation in layer filter strings, e.g. 'layers[0-5]' -> ['layers0', ..., 'layers5'].

Run audio through HTSAT's spectrogram frontend and transformer layers, returning intermediate features from each BasicLayer.

Keep only hidden states needed for discriminator scoring. Strips prepended memory tokens so shapes match padding_mask.

Generate frequency-shaped noise

√2 for interior one-sided bins; 1 for DC (and Nyquist if even N).

Convert frequency in Hz to Bark scale using Zwicker & Fastl formula

Initialize frequency-domain buffers

Compute ISTFT

Left-pad or right-trim emb along seq dim to match target_len. Used for local conditioning embeddings that need to align with x without affect

Find the matching lora_config key for a layer, using isinstance to handle ParametrizedLinear etc.

Check if name contains any pattern substring (with bracket expansion).

Create a WeightedRandomSampler if any dataset has non-default weights, otherwise return None.

Tight (Parseval) power-sine analysis window for hop H=N/R (integer R>=2). g[n] = A_p * sin^p(pi n / N), with A_p chosen so that sum_{r=0

Resample and pad/truncate to CLAP's expected 10s at 48kHz.

Project v0 into components parallel and orthogonal to v1. Args: v0: Tensor to project (B, C, T) v1: Referenc

Concatenate `n_chirps` random linear chirps to `batch` along dim=0. Accepts mono [B, N] or multichannel [B, C, N] input and returns the same

Applies alias cancellation by inverting the sign of every second element of every second row, starting from the second row's first elem

Build sequence corresponding to the pattern from the input tensor z. The sequence is built using up to sequence_steps if specified, and non-pa

Compute update-to-weight ratio without concatenating all params into a single tensor. Accumulates squared norms on GPU to avoid the ~2 GB tempora

Adapted from: https://github.com/mseitzer/pytorch-fid/blob/master/src/pytorch_fid/fid_score.py Adapted from: https://github.com/gudgud96/frec

Cast frozen base weights to lower precision, keeping LoRA params in fp32. Args: model: Model with LoRA parametrizations applied.

Compute effective sequence lengths from seconds_total in conditioning dicts. Args: conditioning: List of conditioning dicts, one per

Compute logits for a batch of codes, optionally conditioning on cross-attention and prepend conditioning Handles translation between