| |
| """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] |
|
|
|
|
| |
|
|
| 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): |
| |
| 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 |
|
|
|
|
| |
|
|
| 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) |
|
|
|
|
| |
|
|
| 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) |
| 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 |
| |
| 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) |
| |
| 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") |
| |
| 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) |
| |
| 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() |
|
|
