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