Repository: Blaizzy/mlx-audio
File: mlx_audio/tts/models/llama/llama.py (Python)

Code before:
            trimmed_row = [t - 128266 for t in trimmed_row]
            code_lists.append(trimmed_row)

        return code_lists

    def generate(
        self,
        text,
        voice: str,
        temperature: float = 0.6,
        top_p: float = 0.8,
        split_pattern: str = "\n",
        max_tokens: int = 1200,
        verbose: bool = False,
        ref_audio: mx.array = None,
        ref_text: Optional[str] = None,
        **kwargs,
    ):
        prompt = text.replace("\\n", "\n").replace("\\t", "\t")
        prompts = prompt.split(split_pattern)

        audio_input_ids = None
        if ref_audio is not None and ref_text is not None:
            audio_input_ids = encode_audio_to_codes(ref_audio) + 128266
            audio_transcript_ids = mx.array(
                self.tokenizer(ref_text, return_tensors="pt").input_ids
            )
        else:
            prompts = [f"{voice}: " + p for p in prompts]

        start_token = mx.array([[128259]], dtype=mx.int64)  # Start of human
        end_tokens = mx.array(
            [[128009, 128260]], dtype=mx.int64
        )  # End of text, End of human

        all_input_ids = []

        for prompt in prompts:
            input_ids = mx.array(self.tokenizer(prompt, return_tensors="pt").input_ids)
            all_input_ids.append(input_ids)

        all_modified_input_ids = []
        for input_ids in all_input_ids:
            # reference audio and transcript
            if audio_input_ids is not None:
                audio_start_tokens = mx.array([[128261, 128257]], dtype=mx.int64)
                audio_end_tokens = mx.array([[128258, 128262]], dtype=mx.int64)
                modified_input_ids = mx.concatenate(
                    [
                        start_token,
                        audio_transcript_ids,

Review comment (suggestion):
```suggestion
                self.tokenizer(ref_text, return_tensors="mlx").input_ids
```

Code after:
            trimmed_row = [t - 128266 for t in trimmed_row]
            code_lists.append(trimmed_row)

        return code_lists

    def prepare_input_ids(
        self,
        prompts: List[str],
        voice: Optional[str] = None,
        ref_audio: Optional[mx.array] = None,
        ref_text: Optional[str] = None,
    ):
        audio_input_ids = None
        if ref_audio is not None and ref_text is not None:
            print(
                "\033[93mWARNING: Audio cloning doesn't work reliably on Orpheus.\033[0m \nA known issue affecting Torch and MLX versions. \nWill be fixed once the Canopy labs repo update their code or the model."
            )
            audio_input_ids = encode_audio_to_codes(ref_audio) + 128266
            audio_transcript_ids = self.tokenizer(
                ref_text, return_tensors="mlx"
            ).input_ids
        elif voice is not None:
            prompts = [f"{voice}: " + p for p in prompts]

        start_token = mx.array([[128259]], dtype=mx.int64)  # Start of human
        end_tokens = mx.array(
            [[128009, 128260]], dtype=mx.int64
        )  # End of text, End of human

        prompt_input_ids = []
        for prompt in prompts:
            prompt_input_ids.append(
                self.tokenizer(prompt, return_tensors="mlx").input_ids
            )

        batch_input_ids = []
        pad_token = mx.array([128263], dtype=mx.int64)
        max_len = max([p.shape[1] for p in prompt_input_ids])

        for input_ids in prompt_input_ids:
            modified_input_ids = []

            padding_len = max_len - input_ids.shape[1]
            if padding_len > 0:
                modified_input_ids.append(mx.repeat(pad_token, padding_len)[None, :])

            # reference audio and transcript
            if audio_input_ids is not None:
                audio_start_tokens = mx.array([[128261, 128257]], dtype=mx.int64)
                audio_end_tokens = mx.array([[128258, 128262]], dtype=mx.int64)
                ref_input_ids = mx.concatenate(