def audio_to_codebook(
    model,
    input_values,
    data_config: DataConfig,
    padding_mask=None,
    sample_rate=44100,
):
    """
    Encodes the input audio waveform into discrete codes.

    Args:
        model: The model to use for encoding.
        input_values (`mx.array` of shape `(batch_size, channels, sequence_length)`):
            Float values of the input audio waveform.
        padding_mask (`mx.array` of shape `(batch_size, channels, sequence_length)`):
            Padding mask used to pad the `input_values`.
        sample_rate (`int`, *optional*) :
            Signal sampling_rate

    Returns:
        A list of frames containing the discrete encoded codes for the input audio waveform, along with rescaling
        factors for each chunk when `normalize` is True. Each frames is a tuple `(codebook, scale)`, with
        `codebook` of shape `[batch_size, num_codebooks, frames]`.
        Scale is not used here.
    """
    audio_data = model.preprocess(input_values, sample_rate)

    if padding_mask is None:
        padding_mask = mx.ones_like(input_values).astype(mx.bool_)

    _, encoded_frame, _, _, _ = model.encode(audio_data, n_quantizers=None)  # 1, C, T
    seq_length = encoded_frame.shape[2]

    t_idx_BxTxC, indices_BTCx3 = build_delay_indices(
        B=1,
        T=seq_length,
        C=data_config.channels,
        delay_pattern=data_config.delay_pattern,
    )

    encoded_frame = apply_audio_delay(
        audio_BxTxC=mx.transpose(encoded_frame, (0, 2, 1)),  # 1, T, C
        pad_value=data_config.audio_pad_value,
        bos_value=data_config.audio_bos_value,
        precomp=(t_idx_BxTxC, indices_BTCx3),
    )

    return encoded_frame

    Encodes the input audio waveform into discrete codes.

    Args:
        model: The model to use for encoding.
        input_values (`mx.array` of shape `(batch_size, channels, sequence_length)`):
            Float values of the input audio waveform.
        padding_mask (`mx.array` of shape `(batch_size, channels, sequence_length)`):
            Padding mask used to pad the `input_values`.
        sample_rate (`int`, *optional*) :
            Signal sampling_rate

    Returns:
        A list of frames containing the discrete encoded codes for the input audio waveform, along with rescaling
        factors for each chunk when `normalize` is True. Each frames is a tuple `(codebook, scale)`, with
        `codebook` of shape `[batch_size, num_codebooks, frames]`.
        Scale is not used here.
    
mlx_audio/tts/models/dia/audio.py