nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

nexaai/mlx_backend/mlx_audio/tts/models/outetts/outetts.py

@@ -0,0 +1,255 @@
+import json
+import re
+import time
+import uuid
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx_lm.generate import stream_generate
+from mlx_lm.models.llama import Model as LlamaModel
+from mlx_lm.models.llama import ModelArgs as LlamaModelConfig
+from mlx_lm.models.qwen2 import Model as Qwen2Model
+from mlx_lm.models.qwen2 import ModelArgs as Qwen2ModelConfig
+from mlx_lm.models.qwen3 import Model as Qwen3Model
+from mlx_lm.models.qwen3 import ModelArgs as Qwen3ModelConfig
+from mlx_lm.sample_utils import make_logits_processors, make_sampler
+from tqdm import tqdm
+from transformers import AutoTokenizer
+
+from ..base import GenerationResult
+from .audio_processor import AudioProcessor
+from .dac_interface import DacInterface
+from .prompt_processor import PromptProcessor
+
+
+@dataclass
+class ModelConfig(LlamaModelConfig, Qwen2ModelConfig, Qwen3ModelConfig):
+    tokenizer_name: str = "OuteAI/Llama-OuteTTS-1.0-1B"
+    sample_rate: int = 24000
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelConfig, **kwargs):
+        super().__init__()
+        self.config = config
+        self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name)
+        self.model = self._initialize_model(config, **kwargs)
+
+    def _initialize_model(self, config: ModelConfig, **kwargs) -> nn.Module:
+
+        model_map = {"llama": LlamaModel, "qwen2": Qwen2Model, "qwen3": Qwen3Model}
+
+        if config.model_type not in model_map:
+            raise ValueError(f"Unsupported model type: {config.model_type}")
+
+        return model_map[config.model_type](config, **kwargs)
+
+    def sanitize(self, weights):
+        weights = self.model.sanitize(weights)
+        return {
+            (
+                f"model.{k}"
+                if not k.startswith("model.model.")
+                and not k.startswith("model.lm_head")
+                else k
+            ): v
+            for k, v in weights.items()
+        }
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def sample_rate(self):
+        return self.config.sample_rate
+
+    def __call__(self, *args, **kwargs):
+        return self.model(*args, **kwargs)
+
+    def get_speaker(self, voice: Optional[str], ref_audio: Optional[str]) -> dict:
+        if voice is None and ref_audio is None:
+            voice = f"{Path(__file__).parent}/default_speaker.json"
+            return self.audio_processor.load_speaker(voice)
+
+        if voice is not None:
+            return self.audio_processor.load_speaker(voice)
+
+        speaker = self.audio_processor.create_speaker_from_whisper(ref_audio)
+        file_id = str(uuid.uuid4())
+        save_path = f"~/.cache/mlx_audio/voices/outetts_{file_id}.json"
+        self.audio_processor.save_speaker(speaker, save_path)
+        return speaker
+
+    def chunk_text(self, text: str, max_words: int = 30) -> List[str]:
+        sentences = re.split(r"[.!?。！？︕︖]+", text)
+        sentences = [s.strip() for s in sentences if s.strip()]
+        chunks = []
+        current_chunk = []
+        current_length = 0
+
+        for sentence in sentences:
+            words = sentence.split()
+            if current_length + len(words) > max_words:
+                chunks.append(" ".join(current_chunk))
+                current_chunk = []
+                current_length = 0
+            current_chunk.extend(words)
+            current_length += len(words)
+        if current_chunk:
+            chunks.append(" ".join(current_chunk))
+        return chunks
+
+    def generate_result(
+        self, audio, start_time: float, token_count: int, segment_idx: int, **kwargs
+    ) -> GenerationResult:
+        samples = audio.shape[0] if audio is not None else 0
+        assert samples > 0, "No audio generated"
+
+        sample_rate = (
+            self.config.sample_rate
+            if kwargs.get("sample_rate") is None
+            else kwargs.get("sample_rate")
+        )
+        audio_duration_seconds = samples / sample_rate
+
+        elapsed_time = time.perf_counter() - start_time
+        rtf = audio_duration_seconds / elapsed_time
+
+        duration_mins = int(audio_duration_seconds // 60)
+        duration_secs = int(audio_duration_seconds % 60)
+        duration_ms = int((audio_duration_seconds % 1) * 1000)
+        duration_hours = int(audio_duration_seconds // 3600)
+        duration_str = f"{duration_hours:02d}:{duration_mins:02d}:{duration_secs:02d}.{duration_ms:03d}"
+
+        return GenerationResult(
+            audio=audio,
+            samples=samples,
+            sample_rate=sample_rate,
+            segment_idx=segment_idx,
+            token_count=token_count,
+            audio_duration=duration_str,
+            real_time_factor=rtf,
+            prompt={
+                "tokens": token_count,
+                "tokens-per-sec": (
+                    round(token_count / elapsed_time, 2) if elapsed_time > 0 else 0
+                ),
+            },
+            audio_samples={
+                "samples": samples,
+                "samples-per-sec": (
+                    round(samples / elapsed_time, 2) if elapsed_time > 0 else 0
+                ),
+            },
+            processing_time_seconds=elapsed_time,
+            peak_memory_usage=mx.get_peak_memory() / 1e9,
+        )
+
+    def generate(
+        self,
+        text,
+        voice: Optional[str] = None,
+        temperature: float = 0.4,
+        top_p: float = 0.9,
+        split_pattern: str = "\n",
+        max_tokens: int = 1200,
+        verbose: bool = False,
+        ref_audio: Optional[str] = None,
+        stream: bool = False,
+        streaming_interval: float = 2.0,
+        **kwargs,
+    ):
+
+        prompts = self.chunk_text(text)
+
+        self.prompt_processor = PromptProcessor(self.tokenizer)
+        self.audio_processor = AudioProcessor()
+
+        speaker = self.get_speaker(voice, ref_audio)
+
+        sampler = make_sampler(
+            temperature,
+            top_p,
+            min_p=kwargs.get("min_p", 0.05),
+            top_k=kwargs.get("top_k", 40),
+        )
+        logits_processors = make_logits_processors(
+            kwargs.get("logit_bias", None),
+            kwargs.get("repetition_penalty", 1.1),
+            kwargs.get("repetition_context_size", 64),
+        )
+
+        for prompt in prompts:
+            completion_prompt = self.prompt_processor.get_completion_prompt(
+                prompt, speaker
+            )
+            input_ids = self.tokenizer.encode(
+                completion_prompt, add_special_tokens=False, return_tensors="mlx"
+            )
+            input_length = input_ids.shape[1]
+
+            generated_token_count = 0
+            yielded_token_count = 0
+            streaming_token_interval = int(streaming_interval * 137.5)
+            yielded_frame_count = 0
+
+            time_start = time.perf_counter()
+
+            for i, response in enumerate(
+                tqdm(
+                    stream_generate(
+                        self.model,
+                        tokenizer=self.tokenizer,
+                        prompt=input_ids.squeeze(0),
+                        max_tokens=max_tokens,
+                        sampler=sampler,
+                        logits_processors=logits_processors,
+                    ),
+                    total=max_tokens,
+                    disable=not verbose,
+                )
+            ):
+                next_token = mx.array([response.token])
+                input_ids = mx.concatenate([input_ids, next_token[None, :]], axis=1)
+                generated_token_count += 1
+
+                # send a partial result in streaming mode
+                if stream and generated_token_count % streaming_token_interval == 0:
+                    output_ids = input_ids[:, input_length:].tolist()[0]
+                    output = self.prompt_processor.extract_audio_from_tokens(output_ids)
+                    audio = self.audio_processor.audio_codec.decode(mx.array([output]))[
+                        -1, -1, :
+                    ]
+
+                    yield self.generate_result(
+                        audio=audio[yielded_frame_count:],
+                        start_time=time_start,
+                        token_count=len(output_ids) - yielded_token_count,
+                        segment_idx=i,
+                        **kwargs,
+                    )
+                    yielded_token_count = len(output_ids)
+                    yielded_frame_count = audio.shape[0]
+                    time_start = time.perf_counter()
+
+            output_ids = input_ids[:, input_length:].tolist()[0]
+            output = self.prompt_processor.extract_audio_from_tokens(output_ids)
+
+            audio = self.audio_processor.audio_codec.decode(mx.array([output]))[
+                -1, -1, :
+            ]
+            if audio.shape[0] > yielded_frame_count:
+                yield self.generate_result(
+                    audio=audio[yielded_frame_count:],
+                    start_time=time_start,
+                    token_count=len(output_ids) - yielded_token_count,
+                    segment_idx=i,
+                    **kwargs,
+                )
+
+            # Clear cache after each segment to avoid memory leaks
+            mx.clear_cache()

nexaai/mlx_backend/mlx_audio/tts/models/outetts/prompt_processor.py

@@ -0,0 +1,181 @@
+import re
+from typing import Union
+
+from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
+
+from .tokens import SpecialTokens
+
+
+class PromptProcessor:
+    def __init__(
+        self, tokenizer: Union[str, PreTrainedTokenizer, PreTrainedTokenizerFast]
+    ):
+        self.special_tokens = SpecialTokens()
+
+        if tokenizer:
+            if isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)):
+                self.tokenizer = tokenizer
+            elif isinstance(tokenizer, str):
+                self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)
+            else:
+                raise ValueError(f"Invalid tokenizer: {type(tokenizer)}")
+
+            self.c1 = {}
+            self.c2 = {}
+            self.get_audio_token_map()
+
+        self.input_prompt = "{bos}\n{text_start}{text}{text_end}\n{audio_start}\n"
+        self.global_features = "{fs}{codes}{fe}\n"
+
+    def get_audio_token_map(self):
+        self.c1 = {
+            self.tokenizer.encode(
+                self.special_tokens.c1.format(i), add_special_tokens=False
+            )[0]: i
+            for i in range(1025)
+        }
+        self.c2 = {
+            self.tokenizer.encode(
+                self.special_tokens.c2.format(i), add_special_tokens=False
+            )[0]: i
+            for i in range(1025)
+        }
+
+    def get_features(self, f: dict):
+        features = {
+            "energy": f.get("energy", 0),
+            "spectral_centroid": f.get("spectral_centroid", 0),
+            "pitch": f.get("pitch", 0),
+        }
+        return [f"<|{k}_{v}|>" for k, v in features.items()]
+
+    def get_global_features(self, f: dict):
+        return self.global_features.format(
+            fs=self.special_tokens.global_features_start,
+            codes="".join(self.get_features(f)),
+            fe=self.special_tokens.global_features_end,
+        )
+
+    def create_codes(self, words: dict):
+        codes = []
+        for i in words:
+            word = (
+                i["word"]
+                + self.special_tokens.features
+                + self.special_tokens.time.format(i["duration"])
+            )
+            word += "".join(self.get_features(i["features"]))
+            pairs = []
+
+            for idx in range(len(i["c1"])):
+                c1 = self.special_tokens.c1.format(i["c1"][idx])
+                c2 = self.special_tokens.c2.format(i["c2"][idx])
+                pairs.append(f"{c1}{c2}")
+
+            word += self.special_tokens.code + "".join(pairs)
+            codes.append(
+                self.special_tokens.word_start + word + self.special_tokens.word_end
+            )
+
+        return "\n".join(codes)
+
+    def _init_prompt(self, text):
+        return self.input_prompt.format(
+            bos=self.special_tokens.bos,
+            text_start=self.special_tokens.text_start,
+            text=text,
+            text_end=self.special_tokens.text_end,
+            audio_start=self.special_tokens.audio_start,
+        )
+
+    def _get_separator(self, text: str) -> str:
+        has_hiragana = any("\u3040" <= c <= "\u309f" for c in text)
+        has_katakana = any("\u30a0" <= c <= "\u30ff" for c in text)
+        has_han = any("\u4e00" <= c <= "\u9fff" for c in text)
+        has_hangul = any("\uac00" <= c <= "\ud7af" for c in text)
+
+        if has_hiragana or has_katakana or has_han:
+            return "。"
+        elif has_hangul:
+            return ". "
+        else:
+            return ". "
+
+    def merge_speaker_text(self, input_text: str, speaker_text: str) -> str:
+        speaker_text = speaker_text.strip()
+        separator = self._get_separator(speaker_text)
+
+        # Determine allowed endings based on the separator
+        if separator == "。":
+            allowed_ends = ["。", "？", "！", "?", "!"]
+        else:
+            allowed_ends = [".", "?", "!"]
+
+        rs = ""
+        if speaker_text:
+            last_char = speaker_text[-1]
+            if last_char not in allowed_ends:
+                rs = separator
+            else:
+                if separator != "。":
+                    rs = " "
+
+        output = speaker_text.strip() + rs + input_text.strip()
+
+        return output, rs.strip()
+
+    @staticmethod
+    def text_normalizations(text: str) -> str:
+        # Normalize whitespace characters (newlines, tabs, etc.) to single spaces
+        text = re.sub(r"\s+", " ", text)
+        text = text.replace("…", "...")  # Replace ellipsis character with three dots
+
+        # Strip leading/trailing whitespace
+        text = text.strip()
+
+        # Normalize common Unicode characters to ASCII equivalents
+        text = re.sub(r"[“”]", '"', text)  # Curly quotes to straight quotes
+        text = re.sub(r"[‘’]", "'", text)  # Curly single quotes
+        text = re.sub(r"[–—]", "-", text)  # Various dashes to hyphen
+
+        # Remove control characters
+        text = re.sub(r"[\x00-\x1F\x7F-\x9F]", "", text)
+
+        return text
+
+    def get_completion_prompt(self, text: str, speaker: dict = None):
+        text = self.text_normalizations(text)
+
+        if speaker is not None:
+            text, separator = self.merge_speaker_text(text, speaker["text"])
+            speaker["words"][-1]["word"] += separator
+            codes = self.create_codes(speaker["words"])
+
+        prompt = self._init_prompt(text)
+
+        if speaker is not None:
+            prompt += codes + "\n" + self.special_tokens.word_start
+
+        return prompt
+
+    def get_training_prompt(self, speaker: dict) -> str:
+        text = self.text_normalizations(speaker["text"])
+        words = speaker["words"]
+        global_features = speaker["global_features"]
+
+        prompt = self._init_prompt(text)
+        prompt += self.get_global_features(global_features)
+        prompt += self.create_codes(words)
+        prompt += (
+            "\n" + self.special_tokens.audio_end + "\n" + self.special_tokens.eos + "\n"
+        )
+
+        return prompt
+
+    def extract_audio_from_tokens(self, tokens: list[int]):
+        codebook1 = [self.c1[i] for i in tokens if i in self.c1]
+        codebook2 = [self.c2[i] for i in tokens if i in self.c2]
+        t = min(len(codebook1), len(codebook2))
+        codebook1 = codebook1[:t]
+        codebook2 = codebook2[:t]
+        return [codebook1, codebook2]

nexaai/mlx_backend/mlx_audio/tts/models/outetts/tokens.py

@@ -0,0 +1,36 @@
+from dataclasses import asdict, dataclass
+from typing import Dict
+
+
+@dataclass
+class SpecialTokens:
+    """
+    Dataclass containing special tokens used for text and audio processing.
+    """
+
+    bos: str = "<|im_start|>"
+    eos: str = "<|im_end|>"
+    c1: str = "<|c1_{}|>"
+    c2: str = "<|c2_{}|>"
+    text_start: str = "<|text_start|>"
+    text_end: str = "<|text_end|>"
+    voice_characteristic_start: str = "<|voice_characteristic_start|>"
+    voice_characteristic_end: str = "<|voice_characteristic_end|>"
+    emotion_start: str = "<|emotion_start|>"
+    emotion_end: str = "<|emotion_end|>"
+    audio_start: str = "<|audio_start|>"
+    audio_end: str = "<|audio_end|>"
+    time: str = "<|t_{:.2f}|>"
+    code: str = "<|code|>"
+    energy: str = "<|energy_{}|>"
+    spectral_centroid: str = "<|spectral_centroid_{}|>"
+    pitch: str = "<|pitch_{}|>"
+    word_start: str = "<|word_start|>"
+    word_end: str = "<|word_end|>"
+    features: str = "<|features|>"
+    global_features_start: str = "<|global_features_start|>"
+    global_features_end: str = "<|global_features_end|>"
+
+    def to_dict(self) -> Dict[str, str]:
+        """Convert the dataclass instance to a dictionary using asdict."""
+        return asdict(self)

nexaai/mlx_backend/mlx_audio/tts/models/sesame/__init__.py

@@ -0,0 +1,3 @@
+from .sesame import Model
+
+__all__ = ["Model"]

nexaai/mlx_backend/mlx_audio/tts/models/sesame/attention.py

@@ -0,0 +1,195 @@
+import math
+from typing import Any, Optional
+
+import mlx.core as mx
+from mlx import nn
+from mlx_lm.models.base import scaled_dot_product_attention
+from mlx_lm.models.llama import ModelArgs
+
+
+class Llama3ScaledRoPE(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        max_seq_len: int = 2048,
+        base: float = 500_000.0,
+        scale_factor: float = 32.0,
+        low_freq_factor: int = 1,
+        high_freq_factor: int = 4,
+        old_context_len: int = 8192,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.base = base
+        self.max_seq_len = max_seq_len
+
+        self.scale_factor = scale_factor
+        self.low_freq_factor = low_freq_factor
+        self.high_freq_factor = high_freq_factor
+        self.old_context_len = old_context_len
+        self.is_cache_built = False
+        self.rope_init()
+
+    def rope_init(self):
+        freqs = 1.0 / (
+            self.base
+            ** (
+                mx.arange(0, self.dim, 2)[: (self.dim // 2)].astype(mx.float32)
+                / self.dim
+            )
+        )
+
+        theta = self.apply_scaling(
+            freqs,
+            self.scale_factor,
+            self.low_freq_factor,
+            self.high_freq_factor,
+            self.old_context_len,
+        )
+        self._theta = theta
+        self.build_rope_cache(self.max_seq_len)
+        self.is_cache_built = True
+
+    def build_rope_cache(self, max_seq_len: int = 4096) -> None:
+        seq_idx = mx.arange(max_seq_len, dtype=self._theta.dtype)
+        idx_theta = mx.einsum("i, j -> ij", seq_idx, self._theta).astype(mx.float32)
+        cache = mx.stack([mx.cos(idx_theta), mx.sin(idx_theta)], axis=-1)
+        self._cache = cache
+
+    def apply_scaling(
+        self,
+        freqs: mx.array,
+        scale_factor: float,
+        low_freq_factor: int,
+        high_freq_factor: int,
+        old_context_len: int,
+    ):
+        low_freq_wavelen = old_context_len / low_freq_factor
+        high_freq_wavelen = old_context_len / high_freq_factor
+        new_freqs = []
+        for freq in freqs:
+            wavelen = 2 * math.pi / freq
+            if wavelen < high_freq_wavelen:
+                new_freqs.append(freq)
+            elif wavelen > low_freq_wavelen:
+                new_freqs.append(freq / scale_factor)
+            else:
+                assert low_freq_wavelen != high_freq_wavelen
+                smooth = (old_context_len / wavelen - low_freq_factor) / (
+                    high_freq_factor - low_freq_factor
+                )
+                new_freqs.append((1 - smooth) * freq / scale_factor + smooth * freq)
+        return mx.array(new_freqs, dtype=freqs.dtype)
+
+    def __call__(self, x: mx.array, *, offset: int) -> mx.array:
+        if not self.is_cache_built:
+            raise RuntimeError(
+                "RoPE cache is not built. Please call rope_init() first."
+            )
+
+        seq_len = x.shape[1]
+        rope_cache = (
+            self._cache[:seq_len]
+            if offset is None
+            else self._cache[None, offset : offset + seq_len]
+        )
+        xshaped = x.astype(mx.float32).reshape(*x.shape[:-1], -1, 2)
+        rope_cache = rope_cache.reshape(-1, xshaped.shape[1], 1, xshaped.shape[3], 2)
+
+        x_out = mx.stack(
+            [
+                xshaped[..., 0] * rope_cache[..., 0]
+                - xshaped[..., 1] * rope_cache[..., 1],
+                xshaped[..., 1] * rope_cache[..., 0]
+                + xshaped[..., 0] * rope_cache[..., 1],
+            ],
+            -1,
+        )
+
+        x_out = x_out.flatten(3)
+        return x_out.astype(x.dtype)
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads or n_heads
+
+        self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
+
+        self.scale = head_dim**-0.5
+        if hasattr(args, "attention_bias"):
+            attention_bias = args.attention_bias
+        else:
+            attention_bias = False
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attention_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias)
+
+        self.rope = Llama3ScaledRoPE(
+            self.head_dim,
+            base=args.rope_theta,
+            scale_factor=args.rope_scaling.get("factor", 1.0),
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        b, s_x, _ = x.shape
+        y = x
+
+        s_y = y.shape[1] if y is not None else 0
+
+        q = self.q_proj(x)
+
+        q_per_kv = self.n_heads // self.n_kv_heads
+        q = q.reshape(b, s_x, self.n_kv_heads * q_per_kv, self.head_dim)
+
+        if self.rope is not None:
+            q = self.rope(q, offset=cache.offset if cache else 0)
+
+        q = q.swapaxes(1, 2)
+
+        k = self.k_proj(y)
+        v = self.v_proj(y)
+
+        k = k.reshape(b, s_y, -1, self.head_dim)
+        v = v.reshape(b, s_y, -1, self.head_dim)
+        if self.rope is not None:
+            k = self.rope(k, offset=cache.offset if cache else 0)
+
+        k = k.swapaxes(1, 2)
+        v = v.swapaxes(1, 2)
+
+        if cache:
+            k, v = cache.update_and_fetch(k, v)
+
+        if self.n_heads != self.n_kv_heads:
+            q_per_kv = self.n_heads // self.n_kv_heads
+
+            k = mx.expand_dims(k, axis=2)
+            v = mx.expand_dims(v, axis=2)
+
+            k_expand_shape = (b, self.n_kv_heads, q_per_kv) + k.shape[3:]
+            v_expand_shape = (b, self.n_kv_heads, q_per_kv) + v.shape[3:]
+
+            k = mx.broadcast_to(k, k_expand_shape)
+            v = mx.broadcast_to(v, v_expand_shape)
+
+            k = k.reshape(b, self.n_kv_heads * q_per_kv, *k.shape[3:])
+            v = v.reshape(b, self.n_kv_heads * q_per_kv, *v.shape[3:])
+
+        output = scaled_dot_product_attention(
+            q, k, v, cache=cache, scale=self.scale, mask=mask
+        )
+
+        output = output.swapaxes(1, 2).reshape(b, s_x, -1)
+        return self.o_proj(output)