def _generate(
        self,
        text: str,
        max_tokens: Optional[int] = None,
        cfg_scale: float = 3.0,
        temperature: float = 1.3,
        top_p: float = 0.95,
        use_cfg_filter: bool = True,
        cfg_filter_top_k: int = 35,
        ref_audio: Optional[mx.array] = None,
        ref_text: Optional[str] = None,
    ) -> np.ndarray:
        """
        Generates audio from a text prompt (and optional audio prompt) using the Dia model.

        Returns:
            A numpy array of generated audio samples.
        """
        num_channels = self.config.data.channels
        audio_bos_value = int(self.config.data.audio_bos_value)
        audio_eos_value = int(self.config.data.audio_eos_value)
        audio_pad_value = int(self.config.data.audio_pad_value)
        delay_pattern = self.config.data.delay_pattern
        max_tokens = self.config.data.audio_length if max_tokens is None else max_tokens
        delay_tensor = mx.array(delay_pattern, dtype=mx.int32)
        max_delay_pattern = max(delay_pattern)

        if ref_text is not None:
            text = ref_text.strip() + " " + text

        (
            cond_src_BxS,
            cond_src_positions_BxS,
            cond_src_padding_mask_BxS,
            cond_enc_self_attn_mask_Bx1xSxS,
        ) = self._prepare_text_input(text)

        unc_src_BxS = mx.zeros_like(cond_src_BxS)
        src_BxS = mx.concatenate([unc_src_BxS, cond_src_BxS], axis=0)
        src_positions_BxS = mx.concatenate(
            [cond_src_positions_BxS, cond_src_positions_BxS], axis=0
        )
        src_padding_mask_BxS = mx.concatenate(
            [cond_src_padding_mask_BxS, cond_src_padding_mask_BxS], axis=0
        )
        enc_self_attn_mask_Bx1xSxS = mx.concatenate(
            [cond_enc_self_attn_mask_Bx1xSxS, cond_enc_self_attn_mask_Bx1xSxS], axis=0
        )

        # 2. Encoder Pass
        encoder_out = self.model.encoder(
            x_ids=src_BxS,
            src_positions=src_positions_BxS,
            deterministic=True,
            attn_mask=enc_self_attn_mask_Bx1xSxS,
        )  # Shape: (B, S, E)

        # 3. Prepare Decoder Inputs
        # 3-1. Allocate KV Cache (Static)
        decoder_cross_attention_cache: list[KVCache] = (
            self.model.decoder.precompute_cross_attention_kv(
                max_tokens, encoder_out, src_positions_BxS
            )
        )

        decoder_self_attention_cache: list[KVCache] = []
        for _ in range(self.model.decoder.num_layers):
            decoder_self_attention_cache.append(
                KVCache(
                    self.config.model.decoder.gqa_query_heads,
                    max_tokens,
                    self.config.model.decoder.gqa_head_dim,
                )
            )

        # 3-2. Initialize Decoder Inputs
        generated_BxTxC = mx.full(
            (2, 1, num_channels),
            vals=audio_bos_value,
            dtype=mx.int32,
        )

        current_step = 0
        prompt_len_inc_bos = 1  # Start with BOS length

        # 3-3. Load Audio Prompt (if provided)
        if ref_audio is not None:
            audio_prompt = mx.array(ref_audio)[None, None, ...]  # 1, C, T

            audio_prompt_codebook = audio_to_codebook(
                self.dac_model, audio_prompt, data_config=self.config.data
            )
            audio_prompt_codebook = mx.concatenate(
                [audio_prompt_codebook, audio_prompt_codebook], axis=0
            )
            generated_BxTxC = mx.concatenate(
                [generated_BxTxC, audio_prompt_codebook], axis=1
            )

            prefill_len = generated_BxTxC.shape[1]
            prompt_len_inc_bos = prefill_len
            prefill_tgt_pos = mx.broadcast_to(
                mx.expand_dims(mx.arange(prefill_len), 0), (2, prefill_len)
            )
            prefill_tgt_padding_mask = mx.any(
                generated_BxTxC != audio_pad_value, axis=2
            )

            prefill_self_attn_mask = self._create_attn_mask(
                prefill_tgt_padding_mask,
                prefill_tgt_padding_mask,
                is_causal=True,
            )
            prefill_cross_attn_mask = self._create_attn_mask(
                prefill_tgt_padding_mask,
                src_padding_mask_BxS,
                is_causal=False,
            )

            _ = self.model.decoder(
                tgt_ids_BxTxC=generated_BxTxC,
                encoder_out=encoder_out,
                tgt_positions=prefill_tgt_pos,
                src_positions=src_positions_BxS,
                deterministic=True,
                self_attn_mask=prefill_self_attn_mask,
                cross_attn_mask=prefill_cross_attn_mask,
                self_attention_cache=decoder_self_attention_cache,
                cross_attention_cache=decoder_cross_attention_cache,
            )

            current_step = prefill_len - 1

        # 4. Autoregressive Generation Loop
        eos_detected_channel_0 = False
        eos_countdown = -1
        extra_steps_after_eos = 30

        # Make generated_BxTxC a fixed size tensor
        # Length is either 1 + max tokens or 1 + prompt len + max tokens
        padding = mx.full(
            (2, max_tokens, num_channels),
            vals=-1,
            dtype=mx.int32,
        )
        generated_BxTxC = mx.concatenate([generated_BxTxC, padding], axis=1)

        decode_step = self.model.decoder.decode_step

        tgt_padding_mask = mx.any(
            mx.expand_dims(generated_BxTxC[:, -1, :], 1) != audio_pad_value, axis=2
        )  # [B, 1]

        # Generated tokens are never PAD, so we use fixed mask
        decoder_cross_attn_mask = self._create_attn_mask(
            tgt_padding_mask,  # Query mask [B, 1]
            src_padding_mask_BxS,  # Key mask [B, S]
            is_causal=False,
        )  # [B, 1, 1, S]

        top_k = -1
        if use_cfg_filter and cfg_filter_top_k is not None:
            top_k = cfg_filter_top_k
        sampler = make_sampler(temperature, top_p, top_k=top_k)

        for step in trange(current_step, current_step + max_tokens):
            tgt_ids_Bx1xC = mx.expand_dims(generated_BxTxC[:, step, :], 1)
            tgt_pos_Bx1 = mx.full(
                (2, 1),
                vals=step,
                dtype=mx.int32,
            )

            logits_Bx1xCxV = decode_step(
                tgt_ids_Bx1xC=tgt_ids_Bx1xC,
                tgt_pos_Bx1=tgt_pos_Bx1,
                encoder_out=encoder_out,
                self_attn_mask=None,
                cross_attn_mask=decoder_cross_attn_mask,
                self_attention_cache=decoder_self_attention_cache,
                cross_attention_cache=decoder_cross_attention_cache,
            )

            V = self.config.model.tgt_vocab_size
            logits_last_BxCxV = logits_Bx1xCxV[:, -1, :, :]  # B, C, V
            uncond_logits_CxV = logits_last_BxCxV[0, :, :]
            cond_logits_CxV = logits_last_BxCxV[1, :, :]

            cfg_logits_CxV = cond_logits_CxV + cfg_scale * (
                cond_logits_CxV - uncond_logits_CxV
            )

            logits_CxV = mx.reshape(cfg_logits_CxV, (-1, V))  # C, V

            # Create a mask for setting tokens beyond 1025 to -inf
            inf_mask = mx.full(logits_CxV.shape, -float("inf"), dtype=logits_CxV.dtype)
            keep_mask = mx.concatenate(
                [
                    mx.ones((logits_CxV.shape[0], 1025)),
                    mx.zeros((logits_CxV.shape[0], logits_CxV.shape[1] - 1025)),
                ],
                axis=1,
            )
            logits_CxV = mx.where(keep_mask == 1, logits_CxV, inf_mask)

            # Sample next token
            pred_C = _sample_next_token(
                logits_CxV,
                temperature=temperature,
                sampler=sampler,
            )

            generation_step_index = step - current_step
            if ref_audio is None:
                pred_C = mx.where(
                    generation_step_index >= delay_tensor,
                    pred_C,
                    mx.full(pred_C.shape, audio_bos_value, dtype=pred_C.dtype),
                )

            # Update generated tokens for next step
            pred_C_expanded = mx.broadcast_to(
                mx.expand_dims(pred_C, 0), (2, num_channels)
            )

            # Split the tensor into parts: before the update, the update itself, and after the update
            before_update = generated_BxTxC[:, : step + 1, :]
            new_token = mx.expand_dims(
                pred_C_expanded, 1
            )  # Shape: (2, 1, num_channels)
            after_update = generated_BxTxC[:, step + 2 :, :]
            generated_BxTxC = mx.concatenate(
                [before_update, new_token, after_update], axis=1
            )

            if not eos_detected_channel_0 and pred_C[0] == audio_eos_value:
                print(f"EOS detected at step {step} for channel 0")
                eos_detected_channel_0 = True
                eos_countdown = extra_steps_after_eos

            if eos_countdown > 0:
                step_after_eos = max_delay_pattern - eos_countdown
                for i, d in enumerate(delay_pattern):
                    if step_after_eos == d:
                        # Update EOS token
                        # Create new array with updated value at position i in the current sequence
                        eos_values = mx.zeros((2, num_channels), dtype=mx.int32)
                        eos_values = eos_values.at[:, i].add(audio_eos_value)
                        # Replace the values at step+1
                        generated_BxTxC = generated_BxTxC.astype(mx.int32)
                        generated_BxTxC = generated_BxTxC.at[:, step + 1, :].add(
                            eos_values
                        )
                    elif step_after_eos > d:
                        # Update PAD token
                        # Create new array with updated value at position i in the current sequence
                        pad_values = mx.zeros((2, num_channels), dtype=mx.int32)
                        pad_values = pad_values.at[:, i].add(audio_pad_value)
                        # Replace the values at step+1
                        generated_BxTxC = generated_BxTxC.astype(mx.int32)
                        generated_BxTxC = generated_BxTxC.at[:, step + 1, :].add(
                            pad_values
                        )

                eos_countdown -= 1
                if eos_countdown == 0:
                    break

            generation_step_index = step - current_step + 1

        output_codes = generated_BxTxC[:, prompt_len_inc_bos : step + 1, :]
        generated_codes = output_codes[0]

        audio = codebook_to_audio(
            generated_codes.transpose(1, 0),
            self.dac_model,
            delay_pattern,
            B=1,
            T=max_tokens,
            C=num_channels,
        )
        return audio.squeeze(), generation_step_index

        Generates audio from a text prompt (and optional audio prompt) using the Dia model.

        Returns:
            A numpy array of generated audio samples.
        
mlx_audio/tts/models/dia/dia.py