def _prepare_text_input(
        self, text: str
    ) -> tuple[mx.array, mx.array, mx.array, mx.array]:
        """Encodes text prompt, pads, and creates attention mask and positions."""
        text_pad_value = self.config.data.text_pad_value
        max_len = self.config.data.text_length

        byte_text = text.encode("utf-8")
        replaced_bytes = byte_text.replace(b"[S1]", b"\x01").replace(b"[S2]", b"\x02")
        text_tokens = list(replaced_bytes)

        current_len = len(text_tokens)
        padding_needed = max_len - current_len
        if padding_needed <= 0:
            text_tokens = text_tokens[:max_len]
            padded_text_np = np.array(text_tokens, dtype=np.uint8)
        else:
            padded_text_np = np.pad(
                text_tokens,
                (0, padding_needed),
                mode="constant",
                constant_values=text_pad_value,
            ).astype(np.uint8)

        src_tokens = mx.array(padded_text_np, dtype=mx.int32)
        src_tokens = mx.expand_dims(src_tokens, 0)  # [1, S]
        src_positions = mx.expand_dims(mx.arange(max_len, dtype=mx.int32), 0)  # [1, S]

        src_padding_mask = src_tokens != text_pad_value  # [1, S]

        enc_self_attn_mask = self._create_attn_mask(
            src_padding_mask, src_padding_mask, is_causal=False
        )  # [1, S, S]

        return src_tokens, src_positions, src_padding_mask, enc_self_attn_mask
Encodes text prompt, pads, and creates attention mask and positions.
mlx_audio/tts/models/dia/dia.py