Conv1d with some builtin handling of asymmetric or causal padding and normalization.

Allow simple CUDA Graphing of a function. Args: func: callable, taking any number of arguments. Its tensors arguments should

Transformer with Streaming / Causal support. Args: d_model (int): Dimension of the data. num_heads (int): Number of heads.

Transformer-based language model on multiple streams of codes. Args: n_q (int): Number of parallel streams to model as input. dep

Layer scale from [Touvron et al 2021] (https://arxiv.org/pdf/2103.17239.pdf). This rescales diagonally the residual outputs close to 0, with a lea

Residual Vector Quantizer. Args: dimension (int): Dimension of the codebooks. input_dimension (None or int): dimension of the inp

Residual block from SEANet model. Args: dim (int): Dimension of the input/output. kernel_sizes (list): List of kernel sizes for t

ConvTranspose1d with some builtin handling of asymmetric or causal padding and normalization.

Gating FFN layer, using the given activation. Args: dim (int): dimension of the input and output of the transformer. activati

Downsampling by some integer amount `stride` using convolutions with a kernel size of twice the stride. If `causal` is True, the output u

Upsample by some integer amount `stride` using transposed convolutions.

Codebook with Euclidean distance. Args: dim (int): Dimension. codebook_size (int): Codebook size. decay (float): Decay fo

Mimi model operating on the raw waveform. Args: encoder (nn.Module): Encoder network. decoder (nn.Module): Decoder network.

Wrapper around Conv1d and normalization applied to this conv to provide a uniform interface across normalization approaches.

Wrapper around ConvTranspose1d and normalization applied to this conv to provide a uniform interface across normalization approaches.

Residual vector quantization implementation. Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf

Efficient streaming KVCache to be compatible with Cuda Graph. Args: batch_size (int): Batch size. num_heads (int): Number of head

Rotary positional embedding (RoPE) from [Su et al 2022](https://arxiv.org/abs/2104.09864). Args: max_period (float): Maximum period of th

SEANet decoder. Args: channels (int): Audio channels. dimension (int): Intermediate representation dimension. n_filters (

SEANet encoder. Args: channels (int): Audio channels. dimension (int): Intermediate representation dimension. n_filters (

Residual Vector Quantizer with separate projections for the first quantizer and the rest. Args: n_q (int): Number of residual vector quan

Similar to `nn.MultiheadAttention` but with support for streaming, causal evaluation. Args: embed_dim (int): Dimension to project to.

Vector quantization implementation. Currently supports only euclidean distance. Args: dim (int): Dimension codebook_size (int

Base class for quantizers.

Base API for all compression model that aim at being used as audio tokenizers with a language model.

Fake quantizer that actually does not perform any quantization.

Transformer with optional projections of the input and output to different dimensions when needed. Supports multiple outputs. Args: i

Boost learning rate for embeddings (with `scale`). Args: norm (bool): if True, uses a layer norm after the embedding. zero_idx (i

Common API for streaming components. Each streaming component has a streaming state, which is just a dict[str, Tensor]. By convention, the fi

TransformerLayer with Streaming / Causal support. Args: d_model (int): Dimension of the data. num_heads (int): Number of heads.

TorchAutocast utility class. Allows you to enable and disable autocast. This is specially useful when dealing with different architectures and

LayerNorm for [B, C, T] inputs.

Base API for CompressionModel wrappers that do not depend on external frameworks.