CoDiCodec

• Discrete tokens: at a bitrate of 2.38 kbit/s.
• Continuous latent vectors: at a ~11 Hz frame rate with 64 channels (128x compression).

Installation

1. Install via pip:
   pip install codicodec

Quick Start: Encoding and Decoding

Understanding the encode() Function

• path_or_audio: This can be either:
  • A string: The path to your audio file.
  • A NumPy array or PyTorch tensor: The waveform data itself. The shape should be [audio_channels, waveform_samples] or [waveform_samples] (for mono). If it's stereo and you provide a 1D array, it's automatically expanded to [1, waveform_samples] and duplicated for both channels.
• max_batch_size: This controls the maximum number of audio chunks processed in parallel during encoding. Tune this based on your available GPU memory. A larger max_batch_size generally leads to faster encoding, up to the limit of your GPU's capacity. The default is set in hparams_inference.py.
• discrete:
  • discrete=False (default): The encoder returns the continuous latent vectors.
  • discrete=True: The encoder returns integer indices representing the quantized latent codes. This is essential for training a language model on the compressed representations, as language models typically work with discrete tokens. The indices correspond to entries in the FSQ codebook.
• preprocess_on_gpu:
  • preprocess_on_gpu=True (default): The computationally intensive Short-Time Fourier Transform (STFT) operation is performed on the GPU. This is significantly faster but requires more GPU memory.
  • preprocess_on_gpu=False: The STFT is done on the CPU. Use this if you run into GPU memory issues, especially with very long audio files.
• desired_channels: This allows you to reshape the latent representation. The default (64) results in a [timesteps, latents_per_timestep, 64] output. You can change this to trade off between the number of channels and the length of the latent sequence. For instance:
  • desired_channels=32: Produces [timesteps, latents_per_timestep*2, 32], halving the channel dimension and doubling the latents_per_timestep dimension.
• fix_batch_size:
  • fix_batch_size=False (default): The encoder processes the audio in batches of varying sizes, up to max_batch_size.
  • fix_batch_size=True: The encoder pads the input to ensure that each batch has a size exactly equal to max_batch_size. This is crucial for enabling torch.compile, which can dramatically speed up encoding. torch.compile optimizes the computation graph, but it requires fixed input sizes. The first time a new batch size is encountered, compilation will take some time, but subsequent runs with the same size will be much faster due to cached, optimized Triton kernels. This is highly recommended for large-scale encoding tasks.

Understanding the decode() Function

• latent: A NumPy array or PyTorch tensor of latent embeddings. The expected shape is [audio_channels, dim, length] (or [batch_size, audio_channels, dim, length]). If the dtype of the latents is integer (int32 or int64), decode automatically assumes they are discrete indices.
• mode:
  • mode='parallel' (default): Decodes the entire latent sequence in parallel. This is generally faster for offline processing.
  • mode='autoregressive' : Decodes the sequence step-by-step, using past decoded output to inform the next step. This is useful for generating longer sequences or for simulating a streaming scenario.
• max_batch_size: Similar to encode(), this controls the maximum batch size for decoding.
• denoising_steps: Number of denoising steps the model takes. It uses the default value specified in hparams_inference.py if no argument is supplied.
• time_prompt: Level of noise that is added to past token when doing autoregressive decoding. It uses the default value specified in hparams_inference.py if no argument is supplied.
• preprocess_on_gpu:
  • preprocess_on_gpu=True (default): The inverse STFT operation is performed on the GPU. This is faster but requires more GPU memory.
  • preprocess_on_gpu=False: The inverse STFT is done on the CPU. Use this if you encounter GPU memory issues.

Live Decoding with decode_next() and reset()

• decode_next(latents): Decodes the next segment of audio, given the current latent chunk (latents). It internally maintains a buffer (past_spec, past_latents) of the previously decoded audio to ensure smooth transitions between chunks. The output wv_chunk is a waveform segment ready to be played.
• reset(): Clears the internal buffer. Call this before starting a new live decoding sequence. It's essential to call reset() if you're switching to a new audio stream.

• Low Latency: You get audio output as soon as each latent chunk is available.
• Real-time Applications: Ideal for interactive music systems and other scenarios where immediate decoded audio is needed.

Summary and Best Practices

• Choose discrete=False for maximum reconstruction quality (this is the default behavior).
• Use discrete=True when extracting data for language models.
• Experiment with max_batch_size to optimize encoding/decoding speed.
• Leverage fix_batch_size=True and torch.compile for significant speedups.
• Use preprocess_on_gpu=True unless you encounter memory limitations.
• For live applications, use decode_next() and reset() for low-latency, seamless decoding.
• The output shape of the encode function is [timesteps, latents_per_timestep, dim] or [batch_size, timesteps, latents_per_timestep, dim] if using batched inputs. When using transformer models you can also concatenate all latents along the same time axis, since transformers do not require an ordered sequence. However, it is recommended to use a learned positional embedding for each latent of the timestep.

Summary Embeddings

License