Duplicate from ahmadw/YuE2-3B-MLX

Co-authored-by: Ahmad <ahmadw@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

4bit/config.json

@@ -0,0 +1,44 @@
+{
+  "return_dict": true,
+  "output_hidden_states": false,
+  "torchscript": false,
+  "dtype": "bfloat16",
+  "tie_word_embeddings": false,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "add_cross_attention": false,
+  "architectures": [
+    "YuE2ForCausalLM"
+  ],
+  "bos_token_id": null,
+  "pad_token_id": null,
+  "eos_token_id": null,
+  "decoder_start_token_id": null,
+  "transformers_version": "4.57.6",
+  "auto_map": {
+    "AutoConfig": "modeling_yue2.YuE2Config",
+    "AutoModelForCausalLM": "modeling_yue2.YuE2ForCausalLM"
+  },
+  "model_type": "yue2",
+  "hidden_size": 2048,
+  "num_hidden_layers": 28,
+  "num_attention_heads": 16,
+  "num_key_value_heads": 8,
+  "head_dim": 128,
+  "intermediate_size": 6144,
+  "vocab_size": 184704,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000,
+  "max_position_embeddings": 24576,
+  "latent_type": "vae",
+  "latent_dim": 64,
+  "max_latent_frames": 24576,
+  "timestep_shift": 1.0,
+  "output_attentions": false,
+  "quantization": {
+    "group_size": 64,
+    "bits": 4,
+    "nar_bits": 8,
+    "mode": "affine"
+  }
+}

4bit/model.safetensors

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

4bit/vae.safetensors

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

4bit/vae_config.json

@@ -0,0 +1,78 @@
+{
+  "return_dict": true,
+  "output_hidden_states": false,
+  "torchscript": false,
+  "dtype": "float32",
+  "tie_word_embeddings": true,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "add_cross_attention": false,
+  "architectures": [
+    "YuE2VAE"
+  ],
+  "bos_token_id": null,
+  "pad_token_id": null,
+  "eos_token_id": null,
+  "decoder_start_token_id": null,
+  "transformers_version": "4.57.6",
+  "auto_map": {
+    "AutoConfig": "modeling_vae.YuE2VAEConfig",
+    "AutoModel": "modeling_vae.YuE2VAE"
+  },
+  "model_type": "yue2_vae",
+  "encoder_config": {
+    "c_mults": [
+      1,
+      2,
+      4,
+      8,
+      16,
+      32
+    ],
+    "channels": 64,
+    "in_channels": 2,
+    "latent_dim": 128,
+    "strides": [
+      2,
+      2,
+      4,
+      4,
+      5,
+      6
+    ],
+    "use_snake": true
+  },
+  "decoder_config": {
+    "c_mults": [
+      1,
+      2,
+      4,
+      8,
+      16,
+      32
+    ],
+    "channels": 64,
+    "final_tanh": false,
+    "latent_dim": 64,
+    "out_channels": 2,
+    "snake_type": "vanilla",
+    "strides": [
+      2,
+      2,
+      4,
+      4,
+      5,
+      6
+    ],
+    "use_filter": false,
+    "use_snake": true
+  },
+  "sample_rate": 48000,
+  "latent_dim": 64,
+  "downsampling_ratio": 1920,
+  "audio_channels": 2,
+  "release_variant": "standard",
+  "decode_core_frames": 1024,
+  "decode_halo_frames": 16,
+  "output_attentions": false
+}

4bit/yue2_generation_config.json

@@ -0,0 +1,24 @@
+{
+  "abc": {
+    "temperature": 0.7,
+    "top_p": 0.9,
+    "top_k": 30,
+    "repetition_penalty": 1.005,
+    "penalty_window": 100,
+    "min_tokens": 32,
+    "max_tokens": 4096
+  },
+  "semantic": {
+    "temperature": 1.0,
+    "top_p": 0.95,
+    "top_k": 100,
+    "repetition_penalty": 1.2,
+    "penalty_window": 50,
+    "min_tokens": 200,
+    "max_tokens": 9000
+  },
+  "ode_steps": 32,
+  "ode_method": "midpoint",
+  "context": 24576,
+  "version": "yue2-native-v1"
+}

8bit/config.json

@@ -0,0 +1,43 @@
+{
+  "return_dict": true,
+  "output_hidden_states": false,
+  "torchscript": false,
+  "dtype": "bfloat16",
+  "tie_word_embeddings": false,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "add_cross_attention": false,
+  "architectures": [
+    "YuE2ForCausalLM"
+  ],
+  "bos_token_id": null,
+  "pad_token_id": null,
+  "eos_token_id": null,
+  "decoder_start_token_id": null,
+  "transformers_version": "4.57.6",
+  "auto_map": {
+    "AutoConfig": "modeling_yue2.YuE2Config",
+    "AutoModelForCausalLM": "modeling_yue2.YuE2ForCausalLM"
+  },
+  "model_type": "yue2",
+  "hidden_size": 2048,
+  "num_hidden_layers": 28,
+  "num_attention_heads": 16,
+  "num_key_value_heads": 8,
+  "head_dim": 128,
+  "intermediate_size": 6144,
+  "vocab_size": 184704,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000,
+  "max_position_embeddings": 24576,
+  "latent_type": "vae",
+  "latent_dim": 64,
+  "max_latent_frames": 24576,
+  "timestep_shift": 1.0,
+  "output_attentions": false,
+  "quantization": {
+    "group_size": 64,
+    "bits": 8,
+    "mode": "affine"
+  }
+}

8bit/model.safetensors

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

8bit/vae.safetensors

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

8bit/vae_config.json

@@ -0,0 +1,78 @@
+{
+  "return_dict": true,
+  "output_hidden_states": false,
+  "torchscript": false,
+  "dtype": "float32",
+  "tie_word_embeddings": true,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "add_cross_attention": false,
+  "architectures": [
+    "YuE2VAE"
+  ],
+  "bos_token_id": null,
+  "pad_token_id": null,
+  "eos_token_id": null,
+  "decoder_start_token_id": null,
+  "transformers_version": "4.57.6",
+  "auto_map": {
+    "AutoConfig": "modeling_vae.YuE2VAEConfig",
+    "AutoModel": "modeling_vae.YuE2VAE"
+  },
+  "model_type": "yue2_vae",
+  "encoder_config": {
+    "c_mults": [
+      1,
+      2,
+      4,
+      8,
+      16,
+      32
+    ],
+    "channels": 64,
+    "in_channels": 2,
+    "latent_dim": 128,
+    "strides": [
+      2,
+      2,
+      4,
+      4,
+      5,
+      6
+    ],
+    "use_snake": true
+  },
+  "decoder_config": {
+    "c_mults": [
+      1,
+      2,
+      4,
+      8,
+      16,
+      32
+    ],
+    "channels": 64,
+    "final_tanh": false,
+    "latent_dim": 64,
+    "out_channels": 2,
+    "snake_type": "vanilla",
+    "strides": [
+      2,
+      2,
+      4,
+      4,
+      5,
+      6
+    ],
+    "use_filter": false,
+    "use_snake": true
+  },
+  "sample_rate": 48000,
+  "latent_dim": 64,
+  "downsampling_ratio": 1920,
+  "audio_channels": 2,
+  "release_variant": "standard",
+  "decode_core_frames": 1024,
+  "decode_halo_frames": 16,
+  "output_attentions": false
+}

8bit/yue2_generation_config.json

@@ -0,0 +1,24 @@
+{
+  "abc": {
+    "temperature": 0.7,
+    "top_p": 0.9,
+    "top_k": 30,
+    "repetition_penalty": 1.005,
+    "penalty_window": 100,
+    "min_tokens": 32,
+    "max_tokens": 4096
+  },
+  "semantic": {
+    "temperature": 1.0,
+    "top_p": 0.95,
+    "top_k": 100,
+    "repetition_penalty": 1.2,
+    "penalty_window": 50,
+    "min_tokens": 200,
+    "max_tokens": 9000
+  },
+  "ode_steps": 32,
+  "ode_method": "midpoint",
+  "context": 24576,
+  "version": "yue2-native-v1"
+}

README.md

@@ -0,0 +1,74 @@
+---
+license: cc-by-nc-4.0
+base_model: m-a-p/YuE2-3B
+base_model_relation: quantized
+library_name: mlx
+pipeline_tag: text-to-audio
+tags:
+- mlx
+- music-generation
+- yue2
+language:
+- en
+- zh
+---
+
+# YuE2-3B MLX
+
+Native [MLX](https://github.com/ml-explore/mlx) port of [m-a-p/YuE2-3B](https://huggingface.co/m-a-p/YuE2-3B) (AR/NAR
+Mixture-of-Transformers music generation) with the [YuE2-Vae](https://huggingface.co/m-a-p/YuE2-Vae) decoder folded in.
+Runs on Apple Silicon with no PyTorch dependency. Lyrics + style prompt in, 48 kHz stereo song out, with optional
+editable ABC score planning.
+
+## Variants
+
+| Folder | Backbone | Size | AR logits / NAR velocity error vs torch bf16 |
+|---|---|---|---|
+| `bf16/` | bf16 | 7.0 GB | 1.4% / 1.6% |
+| `8bit/` | 8-bit affine, group size 64, all backbone Linear layers | 4.2 GB | 1.0% / 1.9% |
+| `4bit/` | 4-bit AR path, 8-bit NAR path, group size 64 | 3.4 GB | 8.1% / 4.4% |
+
+The NAR (flow-matching) path stays at 8 bits in the 4-bit variant because fully 4-bit flow matching drifts
+(12% velocity error). VAE decode is bit-exact in all variants. `8bit/` is the recommended default: near-bf16 output
+at roughly twice the decode speed.
+
+Each folder contains `model.safetensors`, `vae.safetensors`, `config.json`, `vae_config.json`, `qwen.tiktoken`
+and `yue2_generation_config.json`. The Python files at the repo root are the inference code.
+
+## Usage
+
+```bash
+pip install mlx tiktoken numpy huggingface_hub
+hf download ahmadw/YuE2-3B-MLX --include "*.py" "8bit/*" --local-dir YuE2-3B-MLX
+cd YuE2-3B-MLX
+python generate.py --model 8bit \
+  --style "English, indie pop, bright acoustic guitar, soft drums, warm lead vocal" \
+  --lyrics "[Verse]
+Soft morning light is touching the window.
+[Chorus]
+Stay with the rhythm, let it carry us home." \
+  --cot full --seed 831001 --out song.wav
+```
+
+Options: `--cot off|melody|full` (symbolic ABC planning; `full` is the upstream default), `--abc-file` to supply your
+own score, `--cfg-scale`, `--steps` (NAR midpoint steps, default 32), `--max-semantic-tokens` (song length cap, 25
+frames per second), `--decode-latents <out>.latents.npy` to re-decode a saved run. A generated ABC score is written
+next to the wav.
+
+Measured on an M-series Mac: bf16 decodes at about 70 tokens/s (35 with CFG), 8bit at about 80 with CFG.
+A three-minute song takes a few minutes end to end.
+
+## Converting yourself
+
+```bash
+python convert.py --output YuE2-3B-MLX/8bit -q --bits 8          # or --bits 4 (NAR stays 8-bit), or no -q for bf16
+python test_parity.py --mlx YuE2-3B-MLX/8bit                      # needs torch + transformers
+```
+
+Sampling uses MLX's RNG, so seeds reproduce within MLX only.
+
+## License
+
+Weights are derived from m-a-p/YuE2-3B and YuE2-Vae and inherit their **CC BY-NC 4.0** license (non-commercial).
+The SnakeBeta activation and stable-audio-tools decoder code are MIT-licensed by NVIDIA and Stability AI respectively;
+see the upstream `THIRD_PARTY_NOTICES.md`. Please cite the original YuE2 authors (m-a-p) when using these weights.

bf16/config.json

@@ -0,0 +1,38 @@
+{
+  "return_dict": true,
+  "output_hidden_states": false,
+  "torchscript": false,
+  "dtype": "bfloat16",
+  "tie_word_embeddings": false,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "add_cross_attention": false,
+  "architectures": [
+    "YuE2ForCausalLM"
+  ],
+  "bos_token_id": null,
+  "pad_token_id": null,
+  "eos_token_id": null,
+  "decoder_start_token_id": null,
+  "transformers_version": "4.57.6",
+  "auto_map": {
+    "AutoConfig": "modeling_yue2.YuE2Config",
+    "AutoModelForCausalLM": "modeling_yue2.YuE2ForCausalLM"
+  },
+  "model_type": "yue2",
+  "hidden_size": 2048,
+  "num_hidden_layers": 28,
+  "num_attention_heads": 16,
+  "num_key_value_heads": 8,
+  "head_dim": 128,
+  "intermediate_size": 6144,
+  "vocab_size": 184704,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000,
+  "max_position_embeddings": 24576,
+  "latent_type": "vae",
+  "latent_dim": 64,
+  "max_latent_frames": 24576,
+  "timestep_shift": 1.0,
+  "output_attentions": false
+}

bf16/model.safetensors

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

bf16/vae.safetensors

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

bf16/vae_config.json

@@ -0,0 +1,78 @@
+{
+  "return_dict": true,
+  "output_hidden_states": false,
+  "torchscript": false,
+  "dtype": "float32",
+  "tie_word_embeddings": true,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "add_cross_attention": false,
+  "architectures": [
+    "YuE2VAE"
+  ],
+  "bos_token_id": null,
+  "pad_token_id": null,
+  "eos_token_id": null,
+  "decoder_start_token_id": null,
+  "transformers_version": "4.57.6",
+  "auto_map": {
+    "AutoConfig": "modeling_vae.YuE2VAEConfig",
+    "AutoModel": "modeling_vae.YuE2VAE"
+  },
+  "model_type": "yue2_vae",
+  "encoder_config": {
+    "c_mults": [
+      1,
+      2,
+      4,
+      8,
+      16,
+      32
+    ],
+    "channels": 64,
+    "in_channels": 2,
+    "latent_dim": 128,
+    "strides": [
+      2,
+      2,
+      4,
+      4,
+      5,
+      6
+    ],
+    "use_snake": true
+  },
+  "decoder_config": {
+    "c_mults": [
+      1,
+      2,
+      4,
+      8,
+      16,
+      32
+    ],
+    "channels": 64,
+    "final_tanh": false,
+    "latent_dim": 64,
+    "out_channels": 2,
+    "snake_type": "vanilla",
+    "strides": [
+      2,
+      2,
+      4,
+      4,
+      5,
+      6
+    ],
+    "use_filter": false,
+    "use_snake": true
+  },
+  "sample_rate": 48000,
+  "latent_dim": 64,
+  "downsampling_ratio": 1920,
+  "audio_channels": 2,
+  "release_variant": "standard",
+  "decode_core_frames": 1024,
+  "decode_halo_frames": 16,
+  "output_attentions": false
+}

bf16/yue2_generation_config.json

@@ -0,0 +1,24 @@
+{
+  "abc": {
+    "temperature": 0.7,
+    "top_p": 0.9,
+    "top_k": 30,
+    "repetition_penalty": 1.005,
+    "penalty_window": 100,
+    "min_tokens": 32,
+    "max_tokens": 4096
+  },
+  "semantic": {
+    "temperature": 1.0,
+    "top_p": 0.95,
+    "top_k": 100,
+    "repetition_penalty": 1.2,
+    "penalty_window": 50,
+    "min_tokens": 200,
+    "max_tokens": 9000
+  },
+  "ode_steps": 32,
+  "ode_method": "midpoint",
+  "context": 24576,
+  "version": "yue2-native-v1"
+}

convert.py

@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""Convert m-a-p/YuE2-3B + m-a-p/YuE2-Vae into an MLX package directory.
+
+Output layout (HF-uploadable):
+  model.safetensors  bf16 backbone, or affine-quantized Linear layers (--quantize)
+  vae.safetensors    f32 Oobleck decoder, weight norm folded, MLX conv layout
+  config.json, vae_config.json, qwen.tiktoken, yue2_generation_config.json
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import shutil
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten
+
+from yue2_model import Yue2Model
+
+TIME_EMBEDDER = {"time_embedder.mlp.0.": "time_embedder.fc1.", "time_embedder.mlp.2.": "time_embedder.fc2."}
+CONV_TRANSPOSE = re.compile(r"^layers\.\d+\.layers\.1\.weight_v$")
+
+
+def resolve(repo_or_dir: str, patterns: list[str]) -> Path:
+    if Path(repo_or_dir).is_dir():
+        return Path(repo_or_dir)
+    from huggingface_hub import snapshot_download
+    return Path(snapshot_download(repo_or_dir, allow_patterns=patterns))
+
+
+def rename(key: str) -> str:
+    for old, new in TIME_EMBEDDER.items():
+        if key.startswith(old):
+            return new + key[len(old):]
+    return key
+
+
+def fold_weight_norm(weights: dict) -> dict:
+    """decoder.* (g, v) pairs → single MLX-layout conv weights."""
+    out = {}
+    for key, value in weights.items():
+        if not key.startswith("decoder.") or key.endswith("weight_g"):
+            continue
+        key = key[len("decoder."):]
+        if key.endswith("weight_v"):
+            g = weights["decoder." + key[:-1] + "g"]
+            norm = mx.sqrt(mx.sum(value.astype(mx.float32) ** 2, axis=(1, 2), keepdims=True))
+            w = g * value / norm
+            # torch Conv1d [out,in,k] → mlx [out,k,in]; ConvTranspose1d [in,out,k] → mlx [out,k,in]
+            w = w.transpose(1, 2, 0) if CONV_TRANSPOSE.match(key) else w.transpose(0, 2, 1)
+            out[key[:-2]] = w
+        else:
+            out[key] = value
+    return out
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--model", default="m-a-p/YuE2-3B", help="HF repo id or local snapshot dir")
+    parser.add_argument("--vae", default="m-a-p/YuE2-Vae", help="HF repo id or local snapshot dir")
+    parser.add_argument("--output", type=Path, default=Path.home() / "Desktop" / "Yue2-3B-MLX")
+    parser.add_argument("-q", "--quantize", action="store_true", help="affine-quantize backbone Linear layers")
+    parser.add_argument("--bits", type=int, default=8, choices=[4, 8])
+    parser.add_argument("--group-size", type=int, default=64)
+    parser.add_argument("--nar-bits", type=int, choices=[4, 8], help="bits for the NAR path (nar_self_attn/nar_mlp); default 8 (flow matching degrades at 4 bits)")
+    parser.add_argument("--overwrite", action="store_true")
+    args = parser.parse_args()
+
+    if args.output.exists():
+        if not args.overwrite:
+            raise SystemExit(f"{args.output} exists; pass --overwrite")
+        shutil.rmtree(args.output)
+    args.output.mkdir(parents=True)
+
+    model_dir = resolve(args.model, ["model.safetensors", "config.json", "qwen.tiktoken", "yue2_generation_config.json"])
+    vae_dir = resolve(args.vae, ["model.safetensors", "config.json"])
+
+    cfg = json.loads((model_dir / "config.json").read_text())
+    weights = {rename(k): v for k, v in mx.load(str(model_dir / "model.safetensors")).items()}
+    if args.quantize:
+        model = Yue2Model(cfg)
+        model.load_weights(list(weights.items()))
+        nar_bits = args.nar_bits or 8
+
+        def predicate(path, module):
+            if not isinstance(module, nn.Linear) or not (path.startswith("model.layers") or path == "lm_head"):
+                return False
+            return {"group_size": args.group_size, "bits": nar_bits if ".nar_" in path else args.bits}
+
+        nn.quantize(model, args.group_size, args.bits, class_predicate=predicate)
+        weights = dict(tree_flatten(model.parameters()))
+        cfg["quantization"] = {"group_size": args.group_size, "bits": args.bits, "nar_bits": nar_bits, "mode": "affine"}
+    print(f"[write] model.safetensors ({len(weights)} tensors)", flush=True)
+    mx.save_safetensors(str(args.output / "model.safetensors"), weights, metadata={"format": "mlx"})
+    (args.output / "config.json").write_text(json.dumps(cfg, indent=2) + "\n")
+
+    vae = fold_weight_norm(mx.load(str(vae_dir / "model.safetensors")))
+    print(f"[write] vae.safetensors ({len(vae)} tensors)", flush=True)
+    mx.save_safetensors(str(args.output / "vae.safetensors"), vae, metadata={"format": "mlx"})
+    shutil.copy2(vae_dir / "config.json", args.output / "vae_config.json")
+    for name in ("qwen.tiktoken", "yue2_generation_config.json"):
+        shutil.copy2(model_dir / name, args.output / name)
+    print(f"[done] {args.output}")
+
+
+if __name__ == "__main__":
+    main()

generate.py

@@ -0,0 +1,320 @@
+#!/usr/bin/env python3
+"""YuE2 song generation on Apple Silicon with MLX.
+
+    python generate.py --model ~/Desktop/Yue2-3B-MLX --style "indie pop, warm vocal" \
+        --lyrics-file lyrics.txt --cot off --seed 831001 --out song.wav
+
+Ports the yue2_infer protocol: symbolic ABC planning (cot=melody/full), semantic
+codec AR decode with CFG, 32-step midpoint flow matching, tiled VAE decode.
+"""
+from __future__ import annotations
+
+import argparse
+import base64
+import json
+import math
+import sys
+import time
+import unicodedata
+import wave
+from dataclasses import dataclass
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+
+from yue2_model import KVCache, Yue2Model, load_model
+from yue2_vae import load_vae
+
+EOD = 151643
+ABC_START, ABC_END = 151847, 151848
+MUSIC_START, MUSIC_END = 151851, 151852
+CODEC_OFFSET, CODEC_SIZE = 151853, 32768
+VOCAB, CONTEXT = 184704, 24576
+SAMPLE_RATE = 48000
+INSTRUCTIONS = {
+    "off": "Generate music with codec tokens from the given conditions.",
+    "melody": "Generate a melody-only ABC transcription without chord symbols, then generate music with codec tokens from the given conditions.",
+    "full": "Generate a chord-annotated ABC transcription, then generate music with codec tokens from the given conditions.",
+}
+
+
+class Tokenizer:
+    """The frozen Qwen text/ABC BPE (not the audio codec)."""
+
+    def __init__(self, merge_file: Path):
+        import tiktoken
+        ranks = {base64.b64decode(t): int(r) for t, r in
+                 (line.split() for line in Path(merge_file).read_bytes().splitlines() if line)}
+        specials = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "<R>", "<S>", "<X>", "<mask>", "<sep>"]
+        specials += [f"<extra_{i}>" for i in range(200)]
+        specials[204:206] = ["<abc>", "</abc>"]
+        pattern = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
+        self._enc = tiktoken.Encoding("YuE2", pat_str=pattern, mergeable_ranks=ranks,
+                                      special_tokens={s: i + len(ranks) for i, s in enumerate(specials)})
+
+    def encode(self, text):
+        return self._enc.encode_ordinary(unicodedata.normalize("NFC", text))
+
+    def decode(self, ids):
+        return self._enc.decode([int(i) for i in ids if 0 <= i < self._enc.n_vocab], errors="replace")
+
+
+@dataclass(frozen=True)
+class Sampling:
+    temperature: float = 1.0
+    top_p: float = 0.95
+    top_k: int = 100
+    repetition_penalty: float = 1.2
+    penalty_window: int = 50
+    min_tokens: int = 200
+    max_tokens: int = 9000
+
+
+def request_text(style, lyrics, cot):
+    return f"{INSTRUCTIONS[cot]}\n[Tags]\n{style}\n[Lyrics]\n{lyrics}\n"
+
+
+def token_prefix(tok, style, lyrics, cot, abc_ids=None):
+    base = [EOD] + tok.encode(request_text(style, lyrics, cot))
+    if cot == "off":
+        return base + [ABC_START, ABC_END, MUSIC_START]
+    if abc_ids is None:
+        return base + [ABC_START]
+    return base + [ABC_START] + list(abc_ids) + [ABC_END, MUSIC_START]
+
+
+def negative_prefix(tok, cot, abc_ids):
+    base = [EOD] + tok.encode(INSTRUCTIONS[cot])
+    if cot == "off":
+        return base + [MUSIC_START]
+    return base + [ABC_START] + list(abc_ids) + [ABC_END, MUSIC_START]
+
+
+# ── AR sampling ──────────────────────────────────────────────────────────────
+
+def allowed_mask(phase, dtype):
+    idx = mx.arange(VOCAB)
+    end = ABC_END if phase == "abc" else MUSIC_END
+    ok = (idx < EOD) if phase == "abc" else ((idx >= CODEC_OFFSET) & (idx < CODEC_OFFSET + CODEC_SIZE))
+    return mx.where(ok | (idx == end), 0.0, -mx.inf).astype(dtype)
+
+
+def distribution(logits, s: Sampling, history, step, phase, allowed, legacy_off):
+    # Historical cot=off arithmetic stays in BF16; symbolic/CFG paths upcast.
+    scores = logits if legacy_off else logits.astype(mx.float32)
+    end = ABC_END if phase == "abc" else MUSIC_END
+    scores = scores + allowed
+    if step < s.min_tokens:
+        scores[end] = -mx.inf
+    recent = history[-s.penalty_window:]
+    if s.repetition_penalty != 1.0 and recent:
+        freq = mx.zeros((VOCAB,), scores.dtype).at[mx.array(recent)].add(1)
+        alpha = mx.power(mx.array(s.repetition_penalty, scores.dtype), freq)
+        scores = mx.where(scores < 0, scores * alpha, scores / alpha)
+    if s.temperature == 0:
+        return scores
+    if s.temperature != 1:
+        scores = scores / s.temperature
+    kth = VOCAB - min(s.top_k, VOCAB)
+    threshold = mx.partition(scores, kth)[kth]
+    scores = mx.where(scores < threshold, -mx.inf, scores)
+    if s.top_p < 1:
+        order = mx.argsort(-scores)
+        values = scores[order]
+        probs = mx.softmax(values.astype(mx.float32), axis=-1)
+        removed = (mx.cumsum(probs) - probs) > s.top_p
+        removed = removed & (mx.arange(VOCAB) >= (3 if legacy_off else 1))
+        scores = mx.put_along_axis(scores, order, mx.where(removed, -mx.inf, values), axis=0)
+    return scores
+
+
+def generate_tokens(model: Yue2Model, prefix, s: Sampling, seed, phase, negative=None, cfg_scale=1.0,
+                    legacy_off=False, on_token=None):
+    """Return (ids, truncated). Two KV caches when CFG is active (cond / uncond)."""
+    if len(prefix) + s.max_tokens > CONTEXT or (negative and len(negative) + s.max_tokens > CONTEXT):
+        raise ValueError("Prefix + generation budget exceeds the 24576 context")
+    if cfg_scale != 1 and negative is None:
+        raise ValueError("CFG requires a negative prefix")
+    n_layers = len(model.model.layers)
+    caches = [[KVCache() for _ in range(n_layers)] for _ in range(2 if cfg_scale != 1 else 1)]
+    cond = model.ar_step(mx.array([prefix]), caches[0])
+    uncond = model.ar_step(mx.array([negative]), caches[1]) if cfg_scale != 1 else None
+    end = ABC_END if phase == "abc" else MUSIC_END
+    allowed = allowed_mask(phase, cond.dtype if legacy_off else mx.float32)
+    key = mx.random.key(seed)
+    history, eos = [], False
+    for step in range(s.max_tokens):
+        logits = cond if cfg_scale == 1 else uncond + cfg_scale * (cond - uncond)
+        scores = distribution(logits[0], s, history, step, phase, allowed, legacy_off)
+        if s.temperature == 0:
+            token = int(mx.argmax(scores).item())
+        else:
+            key, sub = mx.random.split(key)
+            token = int(mx.random.categorical(scores.astype(mx.float32), key=sub).item())
+        if on_token is not None:
+            on_token(phase, token)
+        if token == end:
+            eos = True
+            break
+        history.append(token)
+        if step + 1 < s.max_tokens:
+            nxt = mx.array([[token]])
+            cond = model.ar_step(nxt, caches[0])
+            if uncond is not None:
+                uncond = model.ar_step(nxt, caches[1])
+    return history, not eos
+
+
+# ── NAR flow matching ────────────────────────────────────────────────────────
+
+def chunk_ranges(frames, prefix_tokens, context=CONTEXT):
+    size = min((context - prefix_tokens - 3) // 2, CONTEXT)
+    if frames < 1 or size < 1:
+        raise ValueError("Empty codec or prefix leaves no acoustic context")
+    return [(a, min(a + size, frames)) for a in range(0, frames, size)]
+
+
+def _logit(t):
+    return max(-20.0, min(20.0, math.log(t / (1 - t)))) if 0 < t < 1 else (20.0 if t >= 1 else -20.0)
+
+
+def synthesize(model: Yue2Model, prefix, codec, seed, steps=32, noise=None, on_progress=None):
+    """Return latents [frames,64] float32 via midpoint ODE from t=1 → 0, solved per context chunk."""
+    if noise is None:
+        noise = mx.random.normal((len(codec), 64), key=mx.random.key(seed))
+    dt = 1.0 / steps
+    out = []
+    for a, b in chunk_ranges(len(codec), len(prefix)):
+        ar_tokens = prefix + [c + CODEC_OFFSET for c in codec[a:b]] + [MUSIC_END]
+        ar_cache = model.nar_prefill(ar_tokens)
+        state = noise[a:b].astype(mx.bfloat16)
+        for step in range(steps):
+            t = 1.0 - step * dt
+            v1 = model.nar_velocity(state, _logit(t), ar_cache, len(ar_tokens))
+            mid = state - v1 * (dt / 2)
+            state = state - model.nar_velocity(mid, _logit(t - dt / 2), ar_cache, len(ar_tokens)) * dt
+            mx.eval(state)
+            if on_progress is not None:
+                on_progress(step + 1, steps)
+        out.append(state.astype(mx.float32))
+    return mx.concatenate(out)
+
+
+# ── Pipeline ─────────────────────────────────────────────────────────────────
+
+class Yue2Pipeline:
+    def __init__(self, path, log=print):
+        self.path, self.log = Path(path), log
+        gen = json.loads((self.path / "yue2_generation_config.json").read_text())
+        self.abc_sampling, self.semantic_sampling = Sampling(**gen["abc"]), Sampling(**gen["semantic"])
+        self.ode_steps = gen.get("ode_steps", 32)
+        self.tokenizer = Tokenizer(self.path / "qwen.tiktoken")
+        self.vae = load_vae(self.path)  # small; load first so a bad file fails before minutes of AR work
+        self.model = load_model(self.path)
+
+    def _progress(self, label):
+        count, start = [0], time.perf_counter()
+
+        def on_token(_phase, _token):
+            count[0] += 1
+            if count[0] % 200 == 0:
+                self.log(f"[{label}] {count[0]} tokens, {count[0] / (time.perf_counter() - start):.1f} tok/s")
+        return on_token
+
+    def __call__(self, style, lyrics, cot="full", seed=831001, abc=None, cfg_scale=None,
+                 semantic_sampling=None, steps=None, on_latents=None):
+        if cot not in INSTRUCTIONS:
+            raise ValueError("cot must be off, melody or full")
+        if abc is not None and cot == "off":
+            raise ValueError("External ABC requires cot=melody/full")
+        tok = self.tokenizer
+        # 1. symbolic plan
+        abc_ids, abc_text = [], None
+        if cot != "off":
+            if abc is not None:
+                abc_ids, abc_text = tok.encode(abc), abc
+            else:
+                self.log("[plan] generating ABC score")
+                abc_ids, truncated = generate_tokens(self.model, token_prefix(tok, style, lyrics, cot),
+                                                     self.abc_sampling, seed, "abc", on_token=self._progress("abc"))
+                abc_text = tok.decode(abc_ids)
+                if truncated:
+                    self.log("[plan] ABC hit max_tokens")
+        prefix = token_prefix(tok, style, lyrics, cot, abc_ids)
+        # 2. semantic codec tokens
+        guidance = (1.01 if cot == "off" else 1.0) if cfg_scale is None else cfg_scale
+        negative = negative_prefix(tok, cot, abc_ids) if guidance != 1 else None
+        self.log(f"[semantic] prefix {len(prefix)} tokens, cfg {guidance}")
+        ids, truncated = generate_tokens(self.model, prefix, semantic_sampling or self.semantic_sampling, seed,
+                                         "semantic", negative, guidance, legacy_off=cot == "off",
+                                         on_token=self._progress("semantic"))
+        if truncated:
+            self.log("[semantic] hit max_tokens")
+        codec = [t - CODEC_OFFSET for t in ids]
+        if not codec:
+            raise RuntimeError("Semantic stage produced no codec tokens")
+        # 3. acoustic latents
+        self.log(f"[nar] {len(codec)} frames ({len(codec) * 1920 / SAMPLE_RATE:.1f}s), {steps or self.ode_steps} midpoint steps")
+        latents = synthesize(self.model, prefix, codec, seed, steps or self.ode_steps,
+                             on_progress=lambda i, n: i % 8 == 0 and self.log(f"[nar] step {i}/{n}"))
+        if on_latents is not None:
+            on_latents(latents)
+        # 4. waveform
+        self.log("[vae] decoding")
+        return self.decode(latents), {"abc": abc_text, "codec": codec, "latents": latents, "prefix": prefix}
+
+    def decode(self, latents):
+        return mx.clip(self.vae.decode_tiled(latents), -1, 1)
+
+
+def write_wav(path, audio, sample_rate=SAMPLE_RATE):
+    pcm = (np.clip(np.array(audio, dtype=np.float32), -1, 1) * 32767).astype("<i2")
+    with wave.open(str(path), "wb") as f:
+        f.setnchannels(pcm.shape[1])
+        f.setsampwidth(2)
+        f.setframerate(sample_rate)
+        f.writeframes(pcm.tobytes())
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--model", type=Path, required=True, help="MLX package dir from convert.py")
+    parser.add_argument("--style", required=True)
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument("--lyrics")
+    group.add_argument("--lyrics-file", type=Path)
+    parser.add_argument("--cot", default="full", choices=list(INSTRUCTIONS))
+    parser.add_argument("--abc-file", type=Path, help="external ABC score (cot=melody/full)")
+    parser.add_argument("--seed", type=int, default=831001)
+    parser.add_argument("--cfg-scale", type=float)
+    parser.add_argument("--steps", type=int, help="NAR midpoint steps (default from generation config)")
+    parser.add_argument("--max-semantic-tokens", type=int, help="cap song length in codec frames (40 ms each)")
+    parser.add_argument("--out", type=Path, default=Path("yue2.wav"))
+    parser.add_argument("--decode-latents", type=Path, help="skip generation; decode a saved <out>.latents.npy")
+    args = parser.parse_args()
+
+    lyrics = args.lyrics if args.lyrics is not None else args.lyrics_file.read_text()
+    abc = args.abc_file.read_text() if args.abc_file else None
+    log = lambda msg: print(msg, file=sys.stderr, flush=True)
+    pipe = Yue2Pipeline(args.model, log=log)
+    sampling = None
+    if args.max_semantic_tokens:
+        sampling = Sampling(**{**pipe.semantic_sampling.__dict__, "max_tokens": args.max_semantic_tokens,
+                               "min_tokens": min(pipe.semantic_sampling.min_tokens, args.max_semantic_tokens)})
+    start = time.perf_counter()
+    latents_path = args.out.with_suffix(".latents.npy")
+    if args.decode_latents:
+        audio, info = pipe.decode(mx.array(np.load(args.decode_latents))), {"abc": None}
+    else:
+        audio, info = pipe(args.style, lyrics, cot=args.cot, seed=args.seed, abc=abc, cfg_scale=args.cfg_scale,
+                           semantic_sampling=sampling, steps=args.steps,
+                           on_latents=lambda z: np.save(latents_path, np.array(z)))
+    write_wav(args.out, audio)
+    if info["abc"] is not None:
+        args.out.with_suffix(".abc").write_text(info["abc"])
+    log(f"[done] {args.out} {audio.shape[0] / SAMPLE_RATE:.1f}s in {time.perf_counter() - start:.0f}s")
+
+
+if __name__ == "__main__":
+    main()

test_parity.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+"""Check the MLX port against the torch reference (needs torch, transformers, the HF snapshots).
+
+    python test_parity.py --mlx ~/Desktop/Yue2-3B-MLX
+
+Compares AR logits, NAR velocity and VAE decode on tiny inputs. Exits non-zero on mismatch.
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+import torch
+
+from generate import CODEC_OFFSET, MUSIC_END, Tokenizer, token_prefix
+from yue2_model import KVCache, load_model
+from yue2_vae import load_vae
+
+LATENT_START, LATENT_END, LATENT_PAD = 184621, 184622, 184623
+
+
+def snapshot(repo):
+    from huggingface_hub import snapshot_download
+    return Path(snapshot_download(repo, allow_patterns=["*.json", "*.py", "model.safetensors", "qwen.tiktoken"]))
+
+
+def rel_err(a, b):
+    a, b = np.asarray(a, np.float32), np.asarray(b, np.float32)
+    return float(np.abs(a - b).max() / (np.abs(b).max() + 1e-6))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--mlx", type=Path, default=Path.home() / "Desktop" / "Yue2-3B-MLX")
+    parser.add_argument("--ref-model", default="m-a-p/YuE2-3B")
+    parser.add_argument("--ref-vae", default="m-a-p/YuE2-Vae")
+    args = parser.parse_args()
+    from transformers import AutoModel, AutoModelForCausalLM
+
+    ref_dir = snapshot(args.ref_model) if not Path(args.ref_model).is_dir() else Path(args.ref_model)
+    vae_dir = snapshot(args.ref_vae) if not Path(args.ref_vae).is_dir() else Path(args.ref_vae)
+    tok = Tokenizer(args.mlx / "qwen.tiktoken")
+    prefix = token_prefix(tok, "pop, warm vocal", "[Verse]\nHello world\n", "off")
+    codec = [7, 300, 5000, 12345, 32000, 9, 9, 42]
+    ar_tokens = prefix + [c + CODEC_OFFSET for c in codec] + [MUSIC_END]
+    frames = len(codec)
+    rng = np.random.default_rng(0)
+    state = rng.standard_normal((frames, 64)).astype(np.float32)
+    raw_t = 0.3
+
+    # ── torch reference ──
+    ref = AutoModelForCausalLM.from_pretrained(ref_dir, trust_remote_code=True, torch_dtype=torch.bfloat16,
+                                               local_files_only=True).eval()
+    with torch.inference_mode():
+        ref_logits = ref(torch.tensor([ar_tokens]), use_cache=False).logits[0].float().numpy()
+        nar = [LATENT_START] + [LATENT_PAD] * frames + [LATENT_END]
+        tokens = torch.tensor([ar_tokens + nar])
+        ar_mask = torch.tensor([[True] * len(ar_tokens) + [False] * len(nar)])
+        nar_mask = ~ar_mask
+        content = torch.tensor([[False] * (len(ar_tokens) + 1) + [True] * frames + [False]])
+        ref_v = ref.nar_velocity(tokens, ar_mask, nar_mask, content, torch.tensor(state), raw_t).float().numpy()
+    del ref
+    ref_vae = AutoModel.from_pretrained(vae_dir, trust_remote_code=True, decoder_only=True, local_files_only=True).eval()
+    z = rng.standard_normal((1, 64, 40)).astype(np.float32)
+    with torch.inference_mode():
+        ref_audio = ref_vae.decode(torch.tensor(z))[0].T.numpy()  # [L,2]
+        ref_tiled = ref_vae.decode_tiled(torch.tensor(z), core_frames=16, halo_frames=16)[0].T.numpy()
+    del ref_vae
+
+    # ── mlx ──
+    model = load_model(args.mlx)
+    caches = [KVCache() for _ in model.model.layers]
+    mlx_logits = np.array(model.ar_step(mx.array([ar_tokens]), caches, all_positions=True)[0].astype(mx.float32))
+    # incremental decode must equal the prefill logits for the same positions
+    caches2 = [KVCache() for _ in model.model.layers]
+    model.ar_step(mx.array([ar_tokens[:-3]]), caches2)
+    step_logits = [np.array(model.ar_step(mx.array([[t]]), caches2).astype(mx.float32))[0] for t in ar_tokens[-3:]]
+    ar_cache = model.nar_prefill(ar_tokens)
+    mlx_v = np.array(model.nar_velocity(mx.array(state).astype(mx.bfloat16), raw_t, ar_cache, len(ar_tokens)).astype(mx.float32))
+    vae = load_vae(args.mlx)
+    zc = mx.array(z[0].T)  # [T,64]
+    mlx_audio = np.array(vae(zc[None])[0])
+    mlx_tiled = np.array(vae.decode_tiled(zc, core=16, halo=16))
+
+    checks = {
+        "ar_logits_rel_err": rel_err(mlx_logits, ref_logits),
+        "ar_argmax_agree": float((mlx_logits.argmax(-1) == ref_logits.argmax(-1)).mean()),
+        "ar_incremental_rel_err": rel_err(np.stack(step_logits), mlx_logits[-3:]),
+        "nar_velocity_rel_err": rel_err(mlx_v, ref_v),
+        "vae_rel_err": rel_err(mlx_audio, ref_audio),
+        "vae_tiled_rel_err": rel_err(mlx_tiled, ref_tiled),
+        "vae_len_ok": float(mlx_audio.shape == ref_audio.shape and mlx_tiled.shape == ref_tiled.shape),
+    }
+    for k, v in checks.items():
+        print(f"{k:26s} {v:.5f}")
+    ok = (checks["ar_logits_rel_err"] < 0.05 and checks["ar_argmax_agree"] > 0.9
+          and checks["ar_incremental_rel_err"] < 0.05 and checks["nar_velocity_rel_err"] < 0.05
+          and checks["vae_rel_err"] < 1e-3 and checks["vae_tiled_rel_err"] < 1e-3 and checks["vae_len_ok"])
+    print("PASS" if ok else "FAIL")
+    raise SystemExit(0 if ok else 1)
+
+
+if __name__ == "__main__":
+    main()

yue2_model.py

@@ -0,0 +1,216 @@
+"""YuE2 AR/NAR Mixture-of-Transformers backbone in MLX.
+
+Mirrors m-a-p/YuE2-3B ``modeling_yue2.py``: a Qwen3-style decoder whose every
+layer carries two attention/MLP paths (AR for text+codec tokens, NAR for the
+acoustic flow-matching state). Weight names follow the checkpoint except
+``time_embedder.mlp.{0,2}`` which ``convert.py`` renames to ``fc1``/``fc2``.
+"""
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class KVCache:
+    """Append-only per-layer cache, grown in 256-slot steps (mlx_lm pattern)."""
+
+    step = 256
+
+    def __init__(self):
+        self.keys = self.values = None
+        self.offset = 0
+
+    def update(self, k, v):
+        prev = self.offset
+        B, H, T, D = k.shape
+        if self.keys is None or prev + T > self.keys.shape[2]:
+            n = ((T + self.step - 1) // self.step) * self.step
+            new_k, new_v = mx.zeros((B, H, n, D), k.dtype), mx.zeros((B, H, n, D), v.dtype)
+            if self.keys is None:
+                self.keys, self.values = new_k, new_v
+            else:
+                if prev % self.step:
+                    self.keys, self.values = self.keys[..., :prev, :], self.values[..., :prev, :]
+                self.keys = mx.concatenate([self.keys, new_k], axis=2)
+                self.values = mx.concatenate([self.values, new_v], axis=2)
+        self.offset += T
+        self.keys[..., prev:self.offset, :] = k
+        self.values[..., prev:self.offset, :] = v
+        return self.keys[..., :self.offset, :], self.values[..., :self.offset, :]
+
+
+class Attention(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.n_heads, self.n_kv, self.head_dim = cfg["num_attention_heads"], cfg["num_key_value_heads"], cfg["head_dim"]
+        h, d = cfg["hidden_size"], self.head_dim
+        self.q_proj = nn.Linear(h, self.n_heads * d, bias=False)
+        self.k_proj = nn.Linear(h, self.n_kv * d, bias=False)
+        self.v_proj = nn.Linear(h, self.n_kv * d, bias=False)
+        self.o_proj = nn.Linear(self.n_heads * d, h, bias=False)
+        self.q_norm = nn.RMSNorm(d, eps=cfg["rms_norm_eps"])
+        self.k_norm = nn.RMSNorm(d, eps=cfg["rms_norm_eps"])
+        self.rope_base = float(cfg["rope_theta"])
+        self.scale = d ** -0.5
+
+    def project_qkv(self, x, offset):
+        """Return q [B,H,T,D], k/v [B,KV,T,D] after QK-norm and RoPE at ``offset``."""
+        B, T, _ = x.shape
+        q = self.q_norm(self.q_proj(x).reshape(B, T, self.n_heads, -1)).transpose(0, 2, 1, 3)
+        k = self.k_norm(self.k_proj(x).reshape(B, T, self.n_kv, -1)).transpose(0, 2, 1, 3)
+        v = self.v_proj(x).reshape(B, T, self.n_kv, -1).transpose(0, 2, 1, 3)
+        rope = lambda a: mx.fast.rope(a, self.head_dim, traditional=False, base=self.rope_base, scale=1.0, offset=offset)
+        return rope(q), rope(k), v
+
+    def attend(self, q, k, v, mask):
+        B, _, T, _ = q.shape
+        o = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask)
+        return self.o_proj(o.transpose(0, 2, 1, 3).reshape(B, T, -1))
+
+
+class MLP(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        h, i = cfg["hidden_size"], cfg["intermediate_size"]
+        self.gate_proj = nn.Linear(h, i, bias=False)
+        self.up_proj = nn.Linear(h, i, bias=False)
+        self.down_proj = nn.Linear(i, h, bias=False)
+
+    def __call__(self, x):
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        h, eps = cfg["hidden_size"], cfg["rms_norm_eps"]
+        self.input_layernorm = nn.RMSNorm(h, eps=eps)
+        self.self_attn = Attention(cfg)
+        self.post_attention_layernorm = nn.RMSNorm(h, eps=eps)
+        self.mlp = MLP(cfg)
+        self.nar_input_layernorm = nn.RMSNorm(h, eps=eps)
+        self.nar_self_attn = Attention(cfg)
+        self.nar_pre_mlp_layernorm = nn.RMSNorm(h, eps=eps)
+        self.nar_mlp = MLP(cfg)
+
+    def ar(self, x, cache: KVCache, mask):
+        q, k, v = self.self_attn.project_qkv(self.input_layernorm(x), cache.offset)
+        k, v = cache.update(k, v)
+        x = x + self.self_attn.attend(q, k, v, mask)
+        return x + self.mlp(self.post_attention_layernorm(x))
+
+    def ar_prefill_kv(self, x):
+        """AR pass that also returns this layer's (k, v) for the NAR to attend over."""
+        q, k, v = self.self_attn.project_qkv(self.input_layernorm(x), 0)
+        x = x + self.self_attn.attend(q, k, v, "causal")
+        return x + self.mlp(self.post_attention_layernorm(x)), (k, v)
+
+    def nar(self, x, ar_k, ar_v, offset):
+        q, k, v = self.nar_self_attn.project_qkv(self.nar_input_layernorm(x), offset)
+        k, v = mx.concatenate([ar_k, k], axis=2), mx.concatenate([ar_v, v], axis=2)
+        x = x + self.nar_self_attn.attend(q, k, v, None)
+        return x + self.nar_mlp(self.nar_pre_mlp_layernorm(x))
+
+
+class Backbone(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(cfg["vocab_size"], cfg["hidden_size"])
+        self.layers = [DecoderLayer(cfg) for _ in range(cfg["num_hidden_layers"])]
+        self.norm = nn.RMSNorm(cfg["hidden_size"], eps=cfg["rms_norm_eps"])
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden, freq_size=256):
+        super().__init__()
+        self.freq_size = freq_size
+        self.fc1 = nn.Linear(freq_size, hidden)
+        self.fc2 = nn.Linear(hidden, hidden)
+
+    def __call__(self, t, dtype):
+        half = self.freq_size // 2
+        freqs = mx.exp(-math.log(10000) * mx.arange(half, dtype=mx.float32) / half)
+        args = t.astype(mx.float32).reshape(-1, 1) * freqs[None]
+        emb = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1).astype(dtype)
+        return self.fc2(nn.silu(self.fc1(emb)))
+
+
+class LatentPosEmbed(nn.Module):
+    """Holds the checkpoint's precomputed sinusoidal table ``pe`` [max_frames, hidden]."""
+
+    def __init__(self, max_frames, hidden):
+        super().__init__()
+        self.pe = mx.zeros((max_frames, hidden))
+
+
+class Yue2Model(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.config = cfg
+        h, d = cfg["hidden_size"], cfg["latent_dim"]
+        self.model = Backbone(cfg)
+        self.lm_head = nn.Linear(h, cfg["vocab_size"], bias=False)
+        self.llm2vae = nn.Linear(h, d)
+        self.vae2llm = nn.Linear(d, h)
+        self.time_embedder = TimestepEmbedder(h)
+        self.latent_pos_embed = LatentPosEmbed(cfg["max_latent_frames"], h)
+
+    # AR path
+    def ar_step(self, tokens, caches, all_positions=False):
+        """tokens [B,T] → logits [B,V] of the last position (or [B,T,V])."""
+        x = self.model.embed_tokens(tokens)
+        mask = "causal" if tokens.shape[1] > 1 else None
+        for layer, cache in zip(self.model.layers, caches):
+            x = layer.ar(x, cache, mask)
+        x = self.model.norm(x if all_positions else x[:, -1:])
+        logits = self.lm_head(x)
+        return logits if all_positions else logits[:, -1]
+
+    # NAR path
+    def nar_prefill(self, ar_tokens):
+        """Run the AR path once over the chunk prefix; return per-layer (k, v)."""
+        x = self.model.embed_tokens(mx.array([ar_tokens]))
+        cache = []
+        for layer in self.model.layers:
+            x, kv = layer.ar_prefill_kv(x)
+            cache.append(kv)
+        mx.eval(cache)
+        return cache
+
+    def nar_velocity(self, state, raw_t, ar_cache, ar_length):
+        """Flow-matching velocity for ``state`` [T,64] at logit-time ``raw_t``."""
+        shift = self.config["timestep_shift"]
+        sig = mx.sigmoid(mx.array(raw_t, dtype=state.dtype))
+        t = shift * sig / (1 + (shift - 1) * sig)
+        n = state.shape[0] + 2  # LATENT_START + frames + LATENT_END
+        x = self.vae2llm(mx.pad(state, [(1, 1), (0, 0)])[None])
+        x = x + self.time_embedder(t, state.dtype)[None]
+        pos = mx.minimum(mx.arange(n), self.config["max_latent_frames"] - 1)
+        x = x + self.latent_pos_embed.pe[pos][None]
+        for layer, (k, v) in zip(self.model.layers, ar_cache):
+            x = layer.nar(x, k, v, ar_length)
+        return self.llm2vae(self.model.norm(x))[0, 1:-1]
+
+
+def load_model(path: Path) -> Yue2Model:
+    path = Path(path)
+    cfg = json.loads((path / "config.json").read_text())
+    model = Yue2Model(cfg)
+    weights = mx.load(str(path / "model.safetensors"))
+    q = cfg.get("quantization")
+    if q:
+        nar_bits = q.get("nar_bits", q["bits"])
+
+        def predicate(path, module):
+            if not isinstance(module, nn.Linear) or f"{path}.scales" not in weights:
+                return False
+            return {"group_size": q["group_size"], "bits": nar_bits if ".nar_" in path else q["bits"]}
+
+        nn.quantize(model, q["group_size"], q["bits"], class_predicate=predicate)
+    model.load_weights(list(weights.items()))
+    mx.eval(model.parameters())
+    return model

yue2_vae.py

@@ -0,0 +1,86 @@
+"""YuE2 Oobleck VAE decoder in MLX (channels-last; weight norm folded by convert.py)."""
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+
+FRAME = 1920  # samples per latent frame at 48 kHz
+
+
+class Snake(nn.Module):
+    def __init__(self, ch):
+        super().__init__()
+        self.alpha = mx.zeros((ch,))
+        self.beta = mx.zeros((ch,))
+
+    def __call__(self, x):
+        alpha, beta = mx.exp(self.alpha), mx.exp(self.beta)
+        return x + (1.0 / (beta + 1e-9)) * mx.sin(x * alpha) ** 2
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, ch, dilation):
+        super().__init__()
+        self.layers = [Snake(ch), nn.Conv1d(ch, ch, 7, dilation=dilation, padding=3 * dilation),
+                       Snake(ch), nn.Conv1d(ch, ch, 1)]
+
+    def __call__(self, x):
+        y = x
+        for layer in self.layers:
+            y = layer(y)
+        return x + y
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, cin, cout, stride):
+        super().__init__()
+        self.layers = [Snake(cin),
+                       nn.ConvTranspose1d(cin, cout, 2 * stride, stride=stride, padding=math.ceil(stride / 2)),
+                       ResidualUnit(cout, 1), ResidualUnit(cout, 3), ResidualUnit(cout, 9)]
+
+    def __call__(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+class OobleckDecoder(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        c, mults, strides = cfg["channels"], [1] + list(cfg["c_mults"]), cfg["strides"]
+        layers = [nn.Conv1d(cfg["latent_dim"], mults[-1] * c, 7, padding=3)]
+        for i in range(len(mults) - 1, 0, -1):
+            layers.append(DecoderBlock(mults[i] * c, mults[i - 1] * c, strides[i - 1]))
+        layers += [Snake(mults[0] * c), nn.Conv1d(mults[0] * c, cfg["out_channels"], 7, padding=3, bias=False)]
+        self.layers = layers
+
+    def __call__(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+    def decode_tiled(self, z, core=1024, halo=16):
+        """z [T,64] float32 → audio [1920*T-64, 2] float32, exact tile cores, no crossfade."""
+        frames, total = z.shape[0], FRAME * z.shape[0] - 64
+        out = []
+        for start in range(0, frames, core):
+            end = min(frames, start + core)
+            left, right = max(0, start - halo), min(frames, end + halo)
+            tile = self(z[None, left:right])[0]
+            crop = (start - left) * FRAME
+            out.append(tile[crop:crop + min(end * FRAME, total) - start * FRAME])
+            mx.eval(out[-1])
+        return mx.concatenate(out)
+
+
+def load_vae(path: Path) -> OobleckDecoder:
+    path = Path(path)
+    cfg = json.loads((path / "vae_config.json").read_text())
+    decoder = OobleckDecoder(cfg["decoder_config"])
+    decoder.load_weights(str(path / "vae.safetensors"))
+    mx.eval(decoder.parameters())
+    return decoder