Upload folder using huggingface_hub

README.md

@@ -1,3 +1,53 @@
----
-license: mit
----
+---
+language:
+- bg
+tags:
+- text-to-speech
+- bulgarian
+- mlx
+- apple-silicon
+library_name: mlx
+license: mit
+---
+
+# 🇧🇬 BG-TTS V5 — MLX (Apple Silicon)
+
+Native MLX port of [beleata74/bg-tts-v5](https://huggingface.co/beleata74/bg-tts-v5) for Apple Silicon (M1/M2/M3/M4).
+
+No CUDA, no NeMo, no PyTorch required. Runs fully on Apple Silicon via MLX.
+
+## Requirements
+
+```bash
+pip install mlx soundfile numpy
+pip install "nanocodec-mlx @ git+https://github.com/nineninesix-ai/nanocodec-mlx.git"
+```
+
+## Quick Start
+
+```python
+from tts_mlx.inference import synthesize
+
+synthesize(
+    checkpoint=".",  # path to this repo
+    text="Здравейте, аз съм българска система за синтез на реч.",
+    output="output.wav",
+    speaker_id=0,      # 0 = AI voice, 1 = audiobook narrator
+    temperature=0.25,
+    top_k=50,
+    top_p=0.8,
+)
+```
+
+## Speakers
+
+| Speaker | Description | Best text length |
+|---------|-------------|-----------------|
+| 0 | AI-generated, clear & fast | Any (20–500+ chars) |
+| 1 | Real female, audiobook narrator | 250–320 chars |
+
+## Credits
+
+Original model by [beleata74](https://huggingface.co/beleata74/bg-tts-v5), created with Claude.
+MLX port by Radi Totev.
+NanoCodec MLX by [nineninesix-ai](https://github.com/nineninesix-ai/nanocodec-mlx).

config.json

@@ -0,0 +1,33 @@
+{
+  "model_type": "bg-tts-v5-mlx",
+  "framework": "mlx",
+  "language": "bg",
+  "encoder": {
+    "vocab_size": 155,
+    "d_model": 512,
+    "n_heads": 8,
+    "n_layers": 6,
+    "d_ff": 2048,
+    "max_len": 512
+  },
+  "decoder": {
+    "vocab_size": 16283,
+    "d_model": 768,
+    "n_heads": 12,
+    "n_layers": 18,
+    "d_ff": 3072,
+    "max_len": 2048,
+    "tokens_per_frame": 4
+  },
+  "codec": {
+    "model": "nineninesix/nemo-nano-codec-22khz-0.6kbps-12.5fps-MLX",
+    "sample_rate": 22050,
+    "num_codebooks": 4,
+    "codebook_size": 4032,
+    "frame_rate": 12.5
+  },
+  "speakers": {
+    "0": "AI-generated female voice, clear and fast",
+    "1": "Real female voice, audiobook narrator (use 250-320 chars)"
+  }
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e473ad925047a7300f80fdb98afa9c80f7d1ab6b4e0f81a726e05d08738d38d
+size 1003201311

tts_mlx/__init__.py

@@ -0,0 +1,3 @@
+from .inference import synthesize, load_from_pytorch_checkpoint
+from .model import TTSEncoderDecoder, V5Config
+from .tokenizer import TTSTokenizer

tts_mlx/config.py

@@ -0,0 +1,94 @@
+# Identical to original tts_v5/config.py — no changes needed
+
+NANOCODEC_MODEL_NAME = "nineninesix/nemo-nano-codec-22khz-0.6kbps-12.5fps-MLX"  # MLX version
+CODEC_SAMPLE_RATE    = 22_050
+CODEC_NUM_CODEBOOKS  = 4
+CODEC_CODEBOOK_SIZE  = 4_032
+CODEC_FRAME_RATE     = 12.5
+CODEC_TOKENS_PER_SEC = 50
+TOKENS_PER_FRAME     = 4
+
+BG_LOWER  = "абвгдежзийклмнопрстуфхцчшщъьюя"
+BG_UPPER  = "АБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЬЮЯ"
+EN_LOWER  = "abcdefghijklmnopqrstuvwxyz"
+EN_UPPER  = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+DIGITS    = "0123456789"
+PUNCT     = '.,!?;:-–—…"\'()[]{}«»„"" '
+EXTRA     = "\n\t"
+
+_ALL_CHARS: list[str] = []
+_seen: set[str] = set()
+for _src in [BG_LOWER, BG_UPPER, EN_LOWER, EN_UPPER, DIGITS, PUNCT, EXTRA]:
+    for _ch in _src:
+        if _ch not in _seen:
+            _ALL_CHARS.append(_ch)
+            _seen.add(_ch)
+
+SPECIAL_TOKENS = {
+    "<pad>":             0,
+    "<start_of_text>":   1,
+    "<end_of_text>":     2,
+    "<start_of_speech>": 3,
+    "<end_of_speech>":   4,
+    "<spk_0>":           5,
+    "<spk_1>":           6,
+    "<spk_2>":           7,
+    "<spk_3>":           8,
+}
+NUM_SPECIAL_TOKENS = len(SPECIAL_TOKENS)
+
+TEXT_CHARS       = _ALL_CHARS
+TEXT_VOCAB_SIZE  = len(TEXT_CHARS)
+TEXT_OFFSET      = NUM_SPECIAL_TOKENS
+AUDIO_OFFSET     = TEXT_OFFSET + TEXT_VOCAB_SIZE
+NUM_AUDIO_TOKENS = CODEC_NUM_CODEBOOKS * CODEC_CODEBOOK_SIZE
+TOTAL_VOCAB_SIZE = AUDIO_OFFSET + NUM_AUDIO_TOKENS
+
+ENCODER_VOCAB_SIZE = AUDIO_OFFSET
+DECODER_VOCAB_SIZE = TOTAL_VOCAB_SIZE
+
+PAD_TOKEN_ID             = SPECIAL_TOKENS["<pad>"]
+START_OF_TEXT_TOKEN_ID    = SPECIAL_TOKENS["<start_of_text>"]
+END_OF_TEXT_TOKEN_ID      = SPECIAL_TOKENS["<end_of_text>"]
+START_OF_SPEECH_TOKEN_ID  = SPECIAL_TOKENS["<start_of_speech>"]
+END_OF_SPEECH_TOKEN_ID    = SPECIAL_TOKENS["<end_of_speech>"]
+SPK_0_TOKEN_ID            = SPECIAL_TOKENS["<spk_0>"]
+SPK_1_TOKEN_ID            = SPECIAL_TOKENS["<spk_1>"]
+
+def audio_token_id(codebook: int, code: int) -> int:
+    return AUDIO_OFFSET + codebook * CODEC_CODEBOOK_SIZE + code
+
+def decode_audio_token(token_id: int) -> tuple[int, int]:
+    offset = token_id - AUDIO_OFFSET
+    return offset // CODEC_CODEBOOK_SIZE, offset % CODEC_CODEBOOK_SIZE
+
+def is_audio_token(token_id: int) -> bool:
+    return AUDIO_OFFSET <= token_id < AUDIO_OFFSET + NUM_AUDIO_TOKENS
+
+def is_special_token(token_id: int) -> bool:
+    return 0 <= token_id < NUM_SPECIAL_TOKENS
+
+def is_text_token(token_id: int) -> bool:
+    return TEXT_OFFSET <= token_id < AUDIO_OFFSET
+
+ENC_D_MODEL    = 512
+ENC_N_HEADS    = 8
+ENC_N_LAYERS   = 6
+ENC_D_FF       = 2048
+
+DEC_D_MODEL    = 768
+DEC_N_HEADS    = 12
+DEC_N_LAYERS   = 18
+DEC_D_FF       = 3072
+
+MAX_TEXT_LEN   = 512
+MAX_AUDIO_LEN  = 2048
+DROPOUT        = 0.10
+CTC_WEIGHT     = 0.1
+
+BATCH_SIZE     = 8
+GRAD_ACCUM     = 2
+LR             = 3e-4
+WEIGHT_DECAY   = 0.1
+WARMUP_STEPS   = 500
+NUM_EPOCHS     = 3

tts_mlx/inference.py

@@ -0,0 +1,303 @@
+"""
+MLX Inference — Encoder-Decoder TTS
+=====================================
+1. Load PyTorch checkpoint weights → convert to MLX arrays
+2. Encode text with encoder (once, bidirectional)
+3. Cache cross-attention KVs from encoder (computed once per layer)
+4. Autoregressively decode audio tokens
+5. Decode tokens → wav using nanocodec-mlx
+"""
+
+import os
+import math
+import numpy as np
+import mlx.core as mx
+import soundfile as sf
+
+from .config import (
+    AUDIO_OFFSET, NUM_AUDIO_TOKENS, END_OF_SPEECH_TOKEN_ID,
+    START_OF_SPEECH_TOKEN_ID, CODEC_NUM_CODEBOOKS, NANOCODEC_MODEL_NAME,
+)
+from .tokenizer import TTSTokenizer
+from .model import TTSEncoderDecoder, V5Config
+
+
+# ── Weight Loading ─────────────────────────────────────────────
+
+def _pt_to_mx(t):
+    """Convert PyTorch tensor to MLX array."""
+    return mx.array(t.float().numpy())
+
+
+def load_from_pytorch_checkpoint(checkpoint_path: str) -> TTSEncoderDecoder:
+    """
+    Load PyTorch checkpoint and convert weights to MLX.
+    """
+    import torch  # only needed when loading from PyTorch checkpoint
+    ckpt_file = os.path.join(checkpoint_path, "checkpoint.pt")
+    print(f"📂 Loading checkpoint: {ckpt_file}")
+    ckpt = torch.load(ckpt_file, map_location="cpu", weights_only=False)
+
+    cfg = ckpt["config"]
+    config = V5Config(
+        enc_vocab_size=cfg["enc_vocab_size"],
+        enc_d_model=cfg["enc_d_model"],
+        enc_n_heads=cfg["enc_n_heads"],
+        enc_n_layers=cfg["enc_n_layers"],
+        enc_d_ff=cfg["enc_d_ff"],
+        max_text_len=cfg["max_text_len"],
+        dec_vocab_size=cfg["dec_vocab_size"],
+        dec_d_model=cfg["dec_d_model"],
+        dec_n_heads=cfg["dec_n_heads"],
+        dec_n_layers=cfg["dec_n_layers"],
+        dec_d_ff=cfg["dec_d_ff"],
+        max_audio_len=cfg["max_audio_len"],
+        dropout=0.0,
+        ctc_weight=0.0,
+        tokens_per_frame=cfg.get("tokens_per_frame", 1),
+    )
+
+    model = TTSEncoderDecoder(config)
+    state = ckpt["model_state_dict"]
+
+    # Build MLX weight dict by mapping PyTorch keys → MLX parameter paths
+    mlx_weights = {}
+    for key, val in state.items():
+        # Skip CTC head (not needed for inference)
+        if key.startswith("ctc_head"):
+            continue
+        # Convert key format: PyTorch uses dots, MLX uses same
+        # e.g. "encoder.layers.0.attention.q_proj.weight" stays the same
+        mlx_weights[key] = _pt_to_mx(val)
+
+    model.load_weights(list(mlx_weights.items()), strict=False)
+    mx.eval(model.parameters())
+
+    step = ckpt.get("step", "?")
+    loss = ckpt.get("loss", 0.0)
+    print(f"✅ Loaded! step={step}, loss={loss:.4f}, tpf={config.tokens_per_frame}")
+    return model
+
+def load_from_safetensors(repo_path: str) -> TTSEncoderDecoder:
+    """Load MLX model from safetensors — no PyTorch required."""
+    import json
+    weights_file = os.path.join(repo_path, "model.safetensors")
+    config_file  = os.path.join(repo_path, "config.json")
+
+    with open(config_file) as f:
+        cfg = json.load(f)
+
+    config = V5Config(
+        enc_vocab_size=cfg["encoder"]["vocab_size"],
+        enc_d_model=cfg["encoder"]["d_model"],
+        enc_n_heads=cfg["encoder"]["n_heads"],
+        enc_n_layers=cfg["encoder"]["n_layers"],
+        enc_d_ff=cfg["encoder"]["d_ff"],
+        max_text_len=cfg["encoder"]["max_len"],
+        dec_vocab_size=cfg["decoder"]["vocab_size"],
+        dec_d_model=cfg["decoder"]["d_model"],
+        dec_n_heads=cfg["decoder"]["n_heads"],
+        dec_n_layers=cfg["decoder"]["n_layers"],
+        dec_d_ff=cfg["decoder"]["d_ff"],
+        max_audio_len=cfg["decoder"]["max_len"],
+        tokens_per_frame=cfg["decoder"]["tokens_per_frame"],
+        dropout=0.0,
+        ctc_weight=0.0,
+    )
+    model = TTSEncoderDecoder(config)
+    model.load_weights(weights_file, strict=False)
+    mx.eval(model.parameters())
+    print(f"✅ Loaded from safetensors!")
+    return model
+
+# ── Generation ─────────────────────────────────────────────────
+
+def sample_token(logits: mx.array, temperature: float, top_k: int, top_p: float,
+                 recent_tokens: list, rep_penalty: float) -> int:
+    """Sample next token from logits."""
+    # Mask: only audio tokens + eos allowed
+    mask = mx.full(logits.shape, -1e9)
+    # Allow audio tokens
+    audio_slice = mx.zeros((NUM_AUDIO_TOKENS,))
+    mask = mx.concatenate([
+        mask[:AUDIO_OFFSET],
+        audio_slice,
+    ], axis=0)
+    # Allow end of speech
+    eos_allow = mx.zeros((1,))
+    mask_list = mx.array(
+        [-1e9] * AUDIO_OFFSET +
+        [0.0] * NUM_AUDIO_TOKENS
+    )
+    # Simpler: build as numpy, set allowed positions
+    mask_np = np.full(logits.shape, -1e9, dtype=np.float32)
+    mask_np[AUDIO_OFFSET: AUDIO_OFFSET + NUM_AUDIO_TOKENS] = 0.0
+    mask_np[END_OF_SPEECH_TOKEN_ID] = 0.0
+    logits_np = np.array(logits) + mask_np
+
+    # Repetition penalty
+    if rep_penalty != 1.0 and recent_tokens:
+        for tid in set(recent_tokens[-200:]):
+            if AUDIO_OFFSET <= tid < AUDIO_OFFSET + NUM_AUDIO_TOKENS:
+                logits_np[tid] /= rep_penalty
+
+    # Temperature
+    logits_np = logits_np / temperature
+
+    # Top-k
+    if top_k > 0:
+        kth_val = np.partition(logits_np, -min(top_k, len(logits_np)))[-min(top_k, len(logits_np))]
+        logits_np[logits_np < kth_val] = -1e9
+
+    # Top-p
+    if top_p < 1.0:
+        sorted_idx = np.argsort(logits_np)[::-1]
+        sorted_logits = logits_np[sorted_idx]
+        probs = np.exp(sorted_logits - sorted_logits[0])
+        probs /= probs.sum()
+        cum = np.cumsum(probs)
+        remove = cum > top_p
+        remove[1:] = remove[:-1].copy()
+        remove[0] = False
+        logits_np[sorted_idx[remove]] = -1e9
+
+    # Sample
+    probs = np.exp(logits_np - logits_np.max())
+    probs /= probs.sum()
+    return int(np.random.choice(len(probs), p=probs))
+
+
+def generate(model: TTSEncoderDecoder, tokenizer: TTSTokenizer,
+             text: str, speaker_id: int = 0,
+             max_new_tokens: int = 2000, temperature: float = 0.25,
+             top_k: int = 50, top_p: float = 0.8, rep_penalty: float = 1.1):
+    """Generate audio tokens from text."""
+
+    # 1. Encode text (once)
+    enc_ids_np = tokenizer.build_encoder_input(text, speaker_id)
+    enc_ids = mx.array(enc_ids_np[None, :])   # [1, T_enc]
+    enc_mask = mx.ones_like(enc_ids)
+
+    enc_out = model.encode(enc_ids, enc_mask)  # [1, T_enc, dec_d]
+    mx.eval(enc_out)
+    print(f"📝 Encoded: {enc_ids.shape[1]} tokens → enc_out {enc_out.shape}")
+
+    # 2. Autoregressive decode
+    dec_ids = mx.array([[START_OF_SPEECH_TOKEN_ID]])
+    past_self_kvs = None
+    cached_cross_kvs = None
+    generated = []
+    offset = 0
+
+    for step in range(max_new_tokens):
+        inp = dec_ids[:, -1:] if past_self_kvs is not None else dec_ids
+
+        logits, new_self_kvs, new_cross_kvs = model.decoder(
+            inp, enc_out, enc_mask,
+            past_key_values=past_self_kvs,
+            cached_cross_kvs=cached_cross_kvs,
+            offset=offset,
+        )
+        mx.eval(logits)
+
+        # Cache cross-attention KVs after first step (they don't change)
+        if cached_cross_kvs is None:
+            cached_cross_kvs = new_cross_kvs
+            mx.eval(cached_cross_kvs)
+
+        past_self_kvs = new_self_kvs
+        offset += inp.shape[1]
+
+        # Sample
+        last_logits = np.array(logits[0, -1, :])
+        tok_id = sample_token(last_logits, temperature, top_k, top_p, generated, rep_penalty)
+
+        if tok_id == END_OF_SPEECH_TOKEN_ID:
+            print(f"🛑 EOS at step {step}")
+            break
+
+        generated.append(tok_id)
+        dec_ids = mx.array([[tok_id]])
+
+        if step % 100 == 0 and step > 0:
+            print(f"  step {step}: {len(generated)} tokens (~{len(generated)/50:.1f}s audio)")
+
+    if not generated:
+        return None
+
+    tokens = np.array(generated)
+    audio_mask = (tokens >= AUDIO_OFFSET) & (tokens < AUDIO_OFFSET + NUM_AUDIO_TOKENS)
+    return tokens[audio_mask] - AUDIO_OFFSET
+
+
+# ── Full Pipeline ──────────────────────────────────────────────
+
+def synthesize(checkpoint: str, text: str, output: str = "output.wav",
+               speaker_id: int = 0, temperature: float = 0.25,
+               top_k: int = 50, top_p: float = 0.8, rep_penalty: float = 1.1,
+               max_tokens: int = 2000):
+
+    print(f"\n🎤 Text: '{text[:80]}'")
+    print(f"   speaker={speaker_id}, T={temperature}, top_k={top_k}, top_p={top_p}")
+
+    # Load model
+    # Auto-detect: safetensors repo or PyTorch checkpoint?
+    if os.path.exists(os.path.join(checkpoint, "model.safetensors")):
+        model = load_from_safetensors(checkpoint)
+    else:
+        model = load_from_pytorch_checkpoint(checkpoint)
+    model.eval()
+
+    # Load tokenizer
+    tokenizer = TTSTokenizer()
+
+    # Generate tokens
+    tokens = generate(model, tokenizer, text, speaker_id, max_tokens,
+                      temperature, top_k, top_p, rep_penalty)
+
+    if tokens is None or len(tokens) == 0:
+        print("❌ No audio generated!")
+        return
+
+    # Trim to multiple of 4 codebooks
+    tokens = tokens[:len(tokens) - len(tokens) % CODEC_NUM_CODEBOOKS]
+    print(f"🔊 {len(tokens)} tokens → {len(tokens)//4} frames → ~{len(tokens)//4/12.5:.1f}s audio")
+
+    # Decode with nanocodec-mlx
+    print("🎵 Decoding with NanoCodec MLX...")
+    from nanocodec_mlx.models.audio_codec import AudioCodecModel
+
+    codec = AudioCodecModel.from_pretrained(NANOCODEC_MODEL_NAME)
+
+    # Reshape tokens: [num_codebooks, num_frames]
+    num_frames = len(tokens) // CODEC_NUM_CODEBOOKS
+    codes = tokens.reshape(num_frames, CODEC_NUM_CODEBOOKS).T  # [4, T]
+    codes_mx = mx.array(codes.astype(np.int32))[None, :, :]    # [1, 4, T]
+    tokens_len = mx.array([num_frames], dtype=mx.int32)
+
+    wav_mx, _ = codec.decode(codes_mx, tokens_len)
+    mx.eval(wav_mx)
+
+    # Save
+    wav_np = np.array(wav_mx[0, 0, :])
+    sf.write(output, wav_np, 22050)
+    duration = len(wav_np) / 22050
+    print(f"✅ Saved: {output} ({duration:.2f}s)")
+    return wav_np
+
+
+if __name__ == "__main__":
+    import argparse
+    p = argparse.ArgumentParser()
+    p.add_argument("--checkpoint", required=True)
+    p.add_argument("--text", required=True)
+    p.add_argument("--output", default="output.wav")
+    p.add_argument("--speaker", type=int, default=0)
+    p.add_argument("--temperature", type=float, default=0.25)
+    p.add_argument("--top-k", type=int, default=50)
+    p.add_argument("--top-p", type=float, default=0.8)
+    p.add_argument("--rep-penalty", type=float, default=1.1)
+    p.add_argument("--max-tokens", type=int, default=2000)
+    a = p.parse_args()
+    synthesize(a.checkpoint, a.text, a.output, a.speaker,
+               a.temperature, a.top_k, a.top_p, a.rep_penalty, a.max_tokens)

tts_mlx/model.py

@@ -0,0 +1,375 @@
+"""
+MLX Model — Encoder-Decoder TTS
+================================
+Port of tts_v5/model.py from PyTorch to MLX.
+Inference-only (no training, no dropout, no CTC head needed).
+"""
+
+import math
+import mlx.core as mx
+import mlx.nn as nn
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+from .config import (
+    TOTAL_VOCAB_SIZE, ENCODER_VOCAB_SIZE, DECODER_VOCAB_SIZE,
+    ENC_D_MODEL, ENC_N_HEADS, ENC_N_LAYERS, ENC_D_FF,
+    DEC_D_MODEL, DEC_N_HEADS, DEC_N_LAYERS, DEC_D_FF,
+    MAX_TEXT_LEN, MAX_AUDIO_LEN,
+    PAD_TOKEN_ID, NUM_AUDIO_TOKENS, AUDIO_OFFSET,
+)
+
+
+# ── Shared Components ──────────────────────────────────────────
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = mx.ones((dim,))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        norm = mx.rsqrt(mx.mean(x * x, axis=-1, keepdims=True) + self.eps)
+        return x * norm * self.weight
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(self, d_model: int, d_ff: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(d_model, d_ff, bias=False)
+        self.up_proj   = nn.Linear(d_model, d_ff, bias=False)
+        self.down_proj = nn.Linear(d_ff, d_model, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+def build_rope_cache(max_seq_len: int, head_dim: int, base: float = 10000.0):
+    """Precompute RoPE cos/sin tables."""
+    inv_freq = 1.0 / (base ** (mx.arange(0, head_dim, 2).astype(mx.float32) / head_dim))
+    t = mx.arange(max_seq_len, dtype=mx.float32)
+    freqs = mx.outer(t, inv_freq)
+    emb = mx.concatenate([freqs, freqs], axis=-1)
+    return mx.cos(emb), mx.sin(emb)
+
+
+def rotate_half(x: mx.array) -> mx.array:
+    half = x.shape[-1] // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    return mx.concatenate([-x2, x1], axis=-1)
+
+
+def apply_rope(q: mx.array, k: mx.array, cos: mx.array, sin: mx.array):
+    cos = cos[None, None, :, :]   # [1, 1, T, head_dim]
+    sin = sin[None, None, :, :]
+    q = q * cos + rotate_half(q) * sin
+    k = k * cos + rotate_half(k) * sin
+    return q, k
+
+
+# ── Encoder (Bidirectional) ────────────────────────────────────
+
+class EncoderSelfAttention(nn.Module):
+    def __init__(self, d_model: int, n_heads: int):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.k_proj = nn.Linear(d_model, d_model, bias=False)
+        self.v_proj = nn.Linear(d_model, d_model, bias=False)
+        self.o_proj = nn.Linear(d_model, d_model, bias=False)
+
+    def __call__(self, x: mx.array, mask: Optional[mx.array] = None) -> mx.array:
+        B, T, _ = x.shape
+        q = self.q_proj(x).reshape(B, T, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+        k = self.k_proj(x).reshape(B, T, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+        v = self.v_proj(x).reshape(B, T, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+
+        scale = 1.0 / math.sqrt(self.head_dim)
+        scores = (q @ k.transpose(0, 1, 3, 2)) * scale  # [B, H, T, T]
+
+        if mask is not None:
+            scores = scores + mask
+
+        attn = mx.softmax(scores.astype(mx.float32), axis=-1).astype(x.dtype)
+        out = (attn @ v).transpose(0, 2, 1, 3).reshape(B, T, -1)
+        return self.o_proj(out)
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, d_model: int, n_heads: int, d_ff: int):
+        super().__init__()
+        self.attn_norm = RMSNorm(d_model)
+        self.attention = EncoderSelfAttention(d_model, n_heads)
+        self.ffn_norm  = RMSNorm(d_model)
+        self.ffn       = SwiGLUFFN(d_model, d_ff)
+
+    def __call__(self, x: mx.array, mask: Optional[mx.array] = None) -> mx.array:
+        x = x + self.attention(self.attn_norm(x), mask)
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+
+class TextEncoder(nn.Module):
+    def __init__(self, vocab_size=ENCODER_VOCAB_SIZE, d_model=ENC_D_MODEL,
+                 n_heads=ENC_N_HEADS, n_layers=ENC_N_LAYERS, d_ff=ENC_D_FF,
+                 max_len=MAX_TEXT_LEN):
+        super().__init__()
+        self.d_model = d_model
+        self.token_embedding = nn.Embedding(vocab_size, d_model)
+        self.pos_embedding   = nn.Embedding(max_len, d_model)
+        self.layers = [EncoderBlock(d_model, n_heads, d_ff) for _ in range(n_layers)]
+        self.final_norm = RMSNorm(d_model)
+
+    def __call__(self, input_ids: mx.array, attention_mask: Optional[mx.array] = None) -> mx.array:
+        B, T = input_ids.shape
+        pos = mx.arange(T)[None, :]  # [1, T]
+        h = self.token_embedding(input_ids) + self.pos_embedding(pos)
+
+        # Build padding mask: [B, 1, 1, T], -inf on pad positions
+        attn_mask = None
+        if attention_mask is not None:
+            # attention_mask: [B, T], 1=real 0=pad
+            pad = (attention_mask == 0).astype(mx.float32)  # [B, T]
+            attn_mask = pad[:, None, None, :] * -1e9         # [B, 1, 1, T]
+
+        for layer in self.layers:
+            h = layer(h, attn_mask)
+
+        return self.final_norm(h)
+
+
+# ── Decoder (Causal with Cross-Attention) ──────────────────────
+
+class DecoderSelfAttention(nn.Module):
+    def __init__(self, d_model: int, n_heads: int, max_len: int, tokens_per_frame: int = 1):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.tokens_per_frame = tokens_per_frame
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.k_proj = nn.Linear(d_model, d_model, bias=False)
+        self.v_proj = nn.Linear(d_model, d_model, bias=False)
+        self.o_proj = nn.Linear(d_model, d_model, bias=False)
+        # Precompute RoPE
+        cos, sin = build_rope_cache(max_len * 2, self.head_dim)
+        self.rope_cos = cos
+        self.rope_sin = sin
+
+    def __call__(self, x: mx.array, past_kv=None, offset: int = 0):
+        """
+        x:        [B, T, d_model]
+        past_kv:  (k_cache, v_cache) or None
+        offset:   number of already-generated tokens (for RoPE position)
+        Returns:  (output, new_k, new_v)
+        """
+        B, T, _ = x.shape
+        q = self.q_proj(x).reshape(B, T, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+        k = self.k_proj(x).reshape(B, T, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+        v = self.v_proj(x).reshape(B, T, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+
+        # Apply RoPE with frame-level positions
+        if self.tokens_per_frame > 1:
+            frame_offset = offset // self.tokens_per_frame
+            frame_positions = mx.arange(T) // self.tokens_per_frame + frame_offset
+        else:
+            frame_positions = mx.arange(T) + offset
+
+        cos = self.rope_cos[frame_positions]  # [T, head_dim]
+        sin = self.rope_sin[frame_positions]
+        q, k = apply_rope(q, k, cos, sin)
+
+        # Append to KV cache
+        if past_kv is not None:
+            k = mx.concatenate([past_kv[0], k], axis=2)
+            v = mx.concatenate([past_kv[1], v], axis=2)
+
+        new_k, new_v = k, v
+
+        # Causal mask only during prefill (T > 1, no cache)
+        scale = 1.0 / math.sqrt(self.head_dim)
+        scores = (q @ k.transpose(0, 1, 3, 2)) * scale
+
+        if past_kv is None and T > 1:
+            # Build causal mask
+            causal = mx.triu(mx.full((T, k.shape[2]), -1e9), k=1)
+            scores = scores + causal[None, None, :, :]
+
+        attn = mx.softmax(scores.astype(mx.float32), axis=-1).astype(x.dtype)
+        out = (attn @ v).transpose(0, 2, 1, 3).reshape(B, T, -1)
+        return self.o_proj(out), new_k, new_v
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dec_d_model: int, enc_d_model: int, n_heads: int):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = dec_d_model // n_heads
+        self.q_proj = nn.Linear(dec_d_model, dec_d_model, bias=False)
+        self.k_proj = nn.Linear(enc_d_model, dec_d_model, bias=False)
+        self.v_proj = nn.Linear(enc_d_model, dec_d_model, bias=False)
+        self.o_proj = nn.Linear(dec_d_model, dec_d_model, bias=False)
+
+    def __call__(self, x: mx.array, encoder_output: mx.array,
+                 encoder_mask: Optional[mx.array] = None,
+                 cached_kv=None):
+        """
+        cached_kv: precomputed (k, v) from encoder — computed once, reused every step.
+        """
+        B, T, _ = x.shape
+        q = self.q_proj(x).reshape(B, T, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+
+        if cached_kv is not None:
+            k, v = cached_kv
+        else:
+            T_enc = encoder_output.shape[1]
+            k = self.k_proj(encoder_output).reshape(B, T_enc, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+            v = self.v_proj(encoder_output).reshape(B, T_enc, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+
+        scale = 1.0 / math.sqrt(self.head_dim)
+        scores = (q @ k.transpose(0, 1, 3, 2)) * scale  # [B, H, T, T_enc]
+
+        if encoder_mask is not None:
+            # encoder_mask: [B, T_enc], 1=real 0=pad
+            pad = (encoder_mask == 0).astype(mx.float32)
+            scores = scores + pad[:, None, None, :] * -1e9
+
+        attn = mx.softmax(scores.astype(mx.float32), axis=-1).astype(x.dtype)
+        out = (attn @ v).transpose(0, 2, 1, 3).reshape(B, T, -1)
+        return self.o_proj(out), (k, v)
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, dec_d_model: int, enc_d_model: int, n_heads: int,
+                 d_ff: int, max_len: int, tokens_per_frame: int = 1):
+        super().__init__()
+        self.self_attn_norm  = RMSNorm(dec_d_model)
+        self.self_attention  = DecoderSelfAttention(dec_d_model, n_heads, max_len, tokens_per_frame)
+        self.cross_attn_norm = RMSNorm(dec_d_model)
+        self.cross_attention = CrossAttention(dec_d_model, enc_d_model, n_heads)
+        self.ffn_norm = RMSNorm(dec_d_model)
+        self.ffn      = SwiGLUFFN(dec_d_model, d_ff)
+
+    def __call__(self, x: mx.array, encoder_output: mx.array,
+                 encoder_mask=None, past_self_kv=None, cached_cross_kv=None,
+                 offset: int = 0):
+        # 1. Causal self-attention
+        h = self.self_attn_norm(x)
+        sa_out, new_k, new_v = self.self_attention(h, past_self_kv, offset)
+        x = x + sa_out
+
+        # 2. Cross-attention (encoder KV cached after first call)
+        h = self.cross_attn_norm(x)
+        ca_out, cross_kv = self.cross_attention(h, encoder_output, encoder_mask, cached_cross_kv)
+        x = x + ca_out
+
+        # 3. FFN
+        x = x + self.ffn(self.ffn_norm(x))
+
+        return x, (new_k, new_v), cross_kv
+
+
+class AudioDecoder(nn.Module):
+    def __init__(self, vocab_size=DECODER_VOCAB_SIZE, d_model=DEC_D_MODEL,
+                 enc_d_model=DEC_D_MODEL, n_heads=DEC_N_HEADS,
+                 n_layers=DEC_N_LAYERS, d_ff=DEC_D_FF,
+                 max_len=MAX_AUDIO_LEN, tokens_per_frame=1):
+        super().__init__()
+        self.tokens_per_frame = tokens_per_frame
+        self.token_embedding = nn.Embedding(vocab_size, d_model)
+        self.layers = [
+            DecoderBlock(d_model, enc_d_model, n_heads, d_ff, max_len, tokens_per_frame)
+            for _ in range(n_layers)
+        ]
+        self.final_norm = RMSNorm(d_model)
+        # LM head tied to token_embedding (set in TTSEncoderDecoder)
+
+    def __call__(self, input_ids: mx.array, encoder_output: mx.array,
+                 encoder_mask=None, past_key_values=None, cached_cross_kvs=None,
+                 offset: int = 0):
+        """
+        input_ids:         [B, T]
+        encoder_output:    [B, T_enc, d]
+        past_key_values:   list of (k, v) per layer, or None
+        cached_cross_kvs:  list of (k, v) per layer from encoder, or None
+        offset:            token offset for RoPE (number of past tokens)
+        """
+        h = self.token_embedding(input_ids)
+
+        new_self_kvs = []
+        new_cross_kvs = []
+
+        for i, layer in enumerate(self.layers):
+            past_self_kv   = past_key_values[i]   if past_key_values   else None
+            cached_cross_kv = cached_cross_kvs[i] if cached_cross_kvs  else None
+
+            h, new_self_kv, new_cross_kv = layer(
+                h, encoder_output, encoder_mask,
+                past_self_kv, cached_cross_kv, offset
+            )
+            new_self_kvs.append(new_self_kv)
+            new_cross_kvs.append(new_cross_kv)
+
+        h = self.final_norm(h)
+        # Tied embedding projection
+        logits = h @ self.token_embedding.weight.T
+        return logits, new_self_kvs, new_cross_kvs
+
+
+# ── Full Model ─────────────────────────────────────────────────
+
+@dataclass
+class V5Config:
+    enc_vocab_size: int = ENCODER_VOCAB_SIZE
+    enc_d_model: int    = ENC_D_MODEL
+    enc_n_heads: int    = ENC_N_HEADS
+    enc_n_layers: int   = ENC_N_LAYERS
+    enc_d_ff: int       = ENC_D_FF
+    max_text_len: int   = MAX_TEXT_LEN
+    dec_vocab_size: int = DECODER_VOCAB_SIZE
+    dec_d_model: int    = DEC_D_MODEL
+    dec_n_heads: int    = DEC_N_HEADS
+    dec_n_layers: int   = DEC_N_LAYERS
+    dec_d_ff: int       = DEC_D_FF
+    max_audio_len: int  = MAX_AUDIO_LEN
+    dropout: float      = 0.0
+    ctc_weight: float   = 0.0
+    tokens_per_frame: int = 1
+
+
+class TTSEncoderDecoder(nn.Module):
+    def __init__(self, config: V5Config):
+        super().__init__()
+        self.config = config
+
+        self.encoder = TextEncoder(
+            vocab_size=config.enc_vocab_size,
+            d_model=config.enc_d_model,
+            n_heads=config.enc_n_heads,
+            n_layers=config.enc_n_layers,
+            d_ff=config.enc_d_ff,
+            max_len=config.max_text_len,
+        )
+
+        if config.enc_d_model != config.dec_d_model:
+            self.enc_projection = nn.Linear(config.enc_d_model, config.dec_d_model, bias=False)
+        else:
+            self.enc_projection = None
+
+        self.decoder = AudioDecoder(
+            vocab_size=config.dec_vocab_size,
+            d_model=config.dec_d_model,
+            enc_d_model=config.dec_d_model,
+            n_heads=config.dec_n_heads,
+            n_layers=config.dec_n_layers,
+            d_ff=config.dec_d_ff,
+            max_len=config.max_audio_len,
+            tokens_per_frame=config.tokens_per_frame,
+        )
+
+    def encode(self, enc_ids: mx.array, enc_mask=None) -> mx.array:
+        """Run encoder + projection once. Returns [B, T_enc, dec_d_model]."""
+        enc_out = self.encoder(enc_ids, enc_mask)
+        if self.enc_projection is not None:
+            enc_out = self.enc_projection(enc_out)
+        return enc_out

tts_mlx/tokenizer.py

@@ -0,0 +1,48 @@
+"""
+Tokenizer — identical to original tts_v5/tokenizer.py
+Pure Python + numpy, no PyTorch dependency.
+"""
+
+import re
+import numpy as np
+from typing import Optional
+
+from .config import (
+    TEXT_CHARS, TEXT_OFFSET, AUDIO_OFFSET,
+    SPECIAL_TOKENS, NUM_SPECIAL_TOKENS,
+    TOTAL_VOCAB_SIZE, CODEC_CODEBOOK_SIZE,
+    PAD_TOKEN_ID, START_OF_TEXT_TOKEN_ID, END_OF_TEXT_TOKEN_ID,
+    START_OF_SPEECH_TOKEN_ID, END_OF_SPEECH_TOKEN_ID,
+    is_audio_token, is_special_token, is_text_token,
+)
+
+
+class TTSTokenizer:
+    def __init__(self):
+        self.char2id: dict[str, int] = {}
+        self.id2char: dict[int, str] = {}
+        for i, ch in enumerate(TEXT_CHARS):
+            tid = TEXT_OFFSET + i
+            self.char2id[ch] = tid
+            self.id2char[tid] = ch
+
+        self._special_id_to_name = {v: k for k, v in SPECIAL_TOKENS.items()}
+        self.vocab_size = TOTAL_VOCAB_SIZE
+        self.text_vocab_size = len(TEXT_CHARS)
+
+    def normalize_text(self, text: str) -> str:
+        text = re.sub(r'\s+', ' ', text).strip()
+        text = re.sub(r'[–—]', '-', text)
+        text = re.sub(r'[«»„""]', '"', text)
+        return text
+
+    def encode_text(self, text: str) -> list[int]:
+        text = self.normalize_text(text)
+        return [self.char2id[ch] for ch in text if ch in self.char2id]
+
+    def build_encoder_input(self, text: str, speaker_id: int = 0) -> np.ndarray:
+        """Encoder input: <sot> text_chars <eot> <spk_X>"""
+        text_ids = self.encode_text(text)
+        spk = SPECIAL_TOKENS[f"<spk_{speaker_id}>"]
+        seq = [START_OF_TEXT_TOKEN_ID] + text_ids + [END_OF_TEXT_TOKEN_ID, spk]
+        return np.array(seq, dtype=np.int32)