Upload MLX q4 weights

config.json

@@ -0,0 +1,106 @@
+{
+  "card": 2048,
+  "n_q": 32,
+  "dep_q": 16,
+  "delays": [
+    0,
+    0,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    0,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2
+  ],
+  "dim": 2048,
+  "text_card": 48000,
+  "existing_text_padding_id": 3,
+  "num_heads": 16,
+  "num_layers": 28,
+  "hidden_scale": 6,
+  "causal": true,
+  "layer_scale": null,
+  "context": 3000,
+  "max_period": 20000.0,
+  "gating": "silu",
+  "norm": "rms_norm_f32",
+  "positional_embedding": "rope_concat",
+  "depformer_dim": 1024,
+  "depformer_num_heads": 16,
+  "depformer_num_layers": 6,
+  "depformer_dim_feedforward": null,
+  "depformer_multi_linear": true,
+  "depformer_norm": "layer_norm",
+  "depformer_pos_emb": "none",
+  "depformer_weights_per_step": true,
+  "demux_second_stream": false,
+  "kv_repeat": 2,
+  "depformer_kv_repeat": 1,
+  "text_card_out": null,
+  "conditioners": {},
+  "fuser": {
+    "cross_attention_pos_emb": false,
+    "cross_attention_pos_emb_scale": 1,
+    "sum": [],
+    "prepend": [],
+    "cross": []
+  },
+  "cross_attention": false,
+  "model_id": {
+    "sig": "77f82164",
+    "epoch": 110
+  },
+  "depformer_weights_per_step_schedule": [
+    0,
+    1,
+    2,
+    3,
+    4,
+    5,
+    6,
+    7,
+    8,
+    8,
+    8,
+    8,
+    8,
+    8,
+    8,
+    8
+  ],
+  "model_type": "hibiki",
+  "lm_gen_config": {
+    "temp": 0.8,
+    "temp_text": 0.8,
+    "top_k": 250,
+    "top_k_text": 250
+  },
+  "mimi_name": "mimi-pytorch-e351c8d8@125.safetensors",
+  "tokenizer_name": "tokenizer_spm_48k_multi6_2.model",
+  "moshi_name": "hibiki-pytorch-77f82164@110.safetensors"
+}

hibiki.q4.safetensors

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

mimi-pytorch-e351c8d8@125.safetensors

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

mlx_hibiki_patch.py

@@ -0,0 +1,158 @@
+"""Runtime patches that make moshi_mlx (0.3.0) run hibiki-zero.
+
+NOTE: this project's local scripts use the vendored fork at ./moshi-mlx/, which
+already has all of these deltas folded in, so they no longer import this module.
+It is kept only as the portable compatibility shim published alongside the q4
+weights on the Hub (`huybik/hibiki-zero-3b-mlx-q4`) for users running the stock
+`moshi-mlx` package off PyPI.
+
+moshi_mlx targets moshi / older hibiki and misses three hibiki-zero deltas:
+  1. config: `hidden_scale` is ignored (feedforward hardcoded to 4*dim) and the
+     depformer feedforward is left None; `kv_repeat` is hardcoded to 1.
+  2. attention: the forward pass asserts kv_repeat==1, so grouped-query
+     attention (hibiki-zero main transformer uses kv_repeat=2) won't run.
+  3. positional embedding: only "rope" (interleaved) is wired up; hibiki-zero
+     uses "rope_concat" == RoPE with interleave=False (MLX traditional=False).
+  4. depformer: hibiki-zero applies a learned per-slice output LayerNorm
+     (`depformer_norms.{i}`) before each audio `linear_out`; moshi_mlx omits it,
+     so the audio logits come out ~3x too small -> out-of-distribution tokens ->
+     babbling/overlapping speech (the text stream is unaffected).
+
+Import this module before building/loading the model.
+"""
+import mlx.core as mx
+import mlx.nn as nn
+from moshi_mlx import models
+from moshi_mlx.models import lm as L
+from moshi_mlx.modules import transformer as T
+
+# --- 1. config: honour hidden_scale + kv_repeat -----------------------------
+_orig_from = models.LmConfig.from_config_dict.__func__
+
+
+def _from_config_dict(cls, data):
+    cfg = _orig_from(cls, data)
+    hs = data["hidden_scale"]
+    cfg.transformer.dim_feedforward = hs * data["dim"]
+    cfg.depformer.transformer.dim_feedforward = hs * data["depformer_dim"]
+    cfg.transformer.kv_repeat = data["kv_repeat"]
+    return cfg
+
+
+models.LmConfig.from_config_dict = classmethod(_from_config_dict)
+
+# --- 2 + 3. attention: GQA + rope_concat ------------------------------------
+_orig_attn_init = T.Attention.__init__
+
+
+def _attn_init(self, cfg):
+    _orig_attn_init(self, cfg)
+    if cfg.positional_embedding in ("rope", "rope_concat"):
+        # rope_concat == interleave=False == MLX traditional=False
+        self.rope = nn.RoPE(
+            cfg.head_dim,
+            traditional=cfg.positional_embedding != "rope_concat",
+            base=cfg.max_period,
+        )
+
+
+def _attn_call(self, xs, cache, mask=None):
+    cfg = self.cfg
+    b, t, _ = xs.shape
+    H, D = cfg.num_heads, cfg.head_dim
+    Hkv = H // cfg.kv_repeat
+    qkv = self.in_proj(xs)
+    q = qkv[..., : H * D].reshape(b, t, H, D).transpose(0, 2, 1, 3)
+    k = qkv[..., H * D : H * D + Hkv * D].reshape(b, t, Hkv, D).transpose(0, 2, 1, 3)
+    v = qkv[..., H * D + Hkv * D :].reshape(b, t, Hkv, D).transpose(0, 2, 1, 3)
+    if self.rope is not None:
+        q = self.rope(q, offset=cache.offset)
+        k = self.rope(k, offset=cache.offset)
+    k, v = cache.update_and_fetch(k, v)
+    k_len = k.shape[2]
+    k_target_len = t + min(cfg.context, k_len - t)
+    if k_target_len < k_len:
+        k = k[:, :, k_len - k_target_len :]
+        v = v[:, :, k_len - k_target_len :]
+    # mx scaled_dot_product_attention handles GQA (H a multiple of Hkv) natively.
+    xs = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask)
+    xs = xs.transpose(0, 2, 1, 3).reshape(b, t, H * D)
+    return self.out_proj(xs)
+
+
+T.Attention.__init__ = _attn_init
+T.Attention.__call__ = _attn_call
+
+# --- 4. depformer per-codebook output LayerNorm -----------------------------
+# hibiki-zero applies a learned per-slice LayerNorm (`depformer_norms.{i}`,
+# dim=depformer_dim, eps 1e-5, with bias) to the depformer transformer output
+# *before* `linear_out` (PyTorch: logits = linears[i](depformer_norms[i](out))).
+# moshi_mlx feeds the un-normalised features straight into linear_out, so the
+# audio logits come out ~3x too small and uncorrelated -> babble + clipping.
+# Add the norm to each slice, apply it in DepFormer.sample, and load its weights.
+_orig_slice_init = L.DepFormerSlice.__init__
+
+
+def _slice_init(self, in_vocab_size, out_vocab_size, main_transformer_dim,
+                demux_second_stream, cfg):
+    _orig_slice_init(self, in_vocab_size, out_vocab_size, main_transformer_dim,
+                     demux_second_stream, cfg)
+    self.norm = nn.LayerNorm(cfg.transformer.d_model, 1e-5)
+
+
+L.DepFormerSlice.__init__ = _slice_init
+
+
+def _depformer_sample(self, main_transformer_out, sampler, text_token, cache,
+                      cfg_coef=1.0):
+    tokens = []
+    last_token = text_token
+    for c in cache:
+        c.reset()
+    for slice in self.slices:
+        if cfg_coef != 1:
+            last_token = mx.tile(last_token, (2, 1))
+        xs = slice.linear_in(main_transformer_out) + slice.emb(last_token)
+        xs = slice.transformer(xs, cache=cache)
+        logits = slice.linear_out(slice.norm(xs))
+        if cfg_coef != 1:
+            l1, l2 = logits.split(2, axis=0)
+            logits = cfg_coef * l1 - (cfg_coef - 1) * l2
+        last_token, _ = sampler(logits)
+        tokens.append(last_token)
+    return mx.stack(tokens, axis=1)
+
+
+L.DepFormer.sample = _depformer_sample
+
+# load depformer_norms.{i}.{weight,bias} into slices.{i}.norm
+_orig_load = L.Lm.load_pytorch_weights
+
+
+def _load_pytorch_weights(self, file, lm_config, strict=True):
+    # Run the original mapping non-strict to build the rest, capture its weight
+    # dict, append our depformer norms, then do the single strict load.
+    pth = mx.load(file)
+    extra = {}
+    for i in range(lm_config.depformer.num_slices):
+        for p in ("weight", "bias"):
+            k = f"depformer_norms.{i}.{p}"
+            if k in pth:
+                extra[f"depformer.slices.{i}.norm.{p}"] = pth[k]
+    captured = {}
+    real_load = self.load_weights
+
+    def _capture(items, strict):
+        captured.update(dict(items))
+        return None
+
+    self.load_weights = _capture
+    try:
+        _orig_load(self, file, lm_config, strict=False)
+    finally:
+        self.load_weights = real_load
+    captured.update(extra)
+    return self.load_weights(list(captured.items()), strict=strict)
+
+
+L.Lm.load_pytorch_weights = _load_pytorch_weights

tokenizer_spm_48k_multi6_2.model

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

verify_mlx_q4.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python
+"""Verify the 4-bit MLX hibiki-zero weights by translating a sample clip.
+
+Uses the pipelined inference path (infer_mlx_fast), which overlaps the CPU Mimi
+codec with the GPU LM (~3x real-time vs ~1.3x for the sequential run_inference
+loop). Output is identical; this is just the fast entry point for the MLX path.
+"""
+import sys
+from pathlib import Path
+
+import mlx.core as mx
+
+HERE = Path(__file__).resolve().parent.parent  # repo root (scripts/ -> ..)
+sys.path.insert(0, str(HERE / "src"))
+from infer_mlx_fast import run
+
+if __name__ == "__main__":
+    mx.random.seed(299792458)
+    run(
+        str(HERE / "hibiki_zero" / "samples" / "leon.wav"),
+        str(HERE / "translations" / "leon_mlx_q4.wav"),
+    )