music3-mnn

MiniMax-Music3（~11.1B：Qwen3-8B 自回归骨干 + RVQ 深度解码器 + 36 层 flow-matching DiT + DAC vocoder）在 MNN 上的完整推理实现，附与 MLX 的同机实测对比。

快速开始

# 依赖：python 3.12, torch, diffusers, transformers, MNN(pip), safetensors
pip install -e .  # 或 PYTHONPATH=src

# 1) 权重：HuggingFace MiniMaxAI/MiniMax-Music3 下载到 ./MiniMax-Music3
# 2) 导出 ONNX 子图
PYTHONPATH=src python -m music3_mnn.export_onnx --ckpt MiniMax-Music3 --out artifacts/onnx

# 3) 转 MNN（注意：masked step 图必须加 --transformerFuseC4 0，否则掩码图被算坏）
MNNConvert -f ONNX --modelFile artifacts/onnx/backbone_step_masked/backbone_step_masked.onnx \
    --MNNModel artifacts/mnn/backbone_step_masked_nc4 \
    --weightQuantBits 8 --weightQuantBlock 128 --transformerFuseC4 0
#    （其余组件转换命令见 REPORT.md §7）

# 4) 生成
PYTHONPATH=src python -m music3_mnn.generate \
  --prompt "C-pop 抒情流行女声, 钢琴与弦乐" \
  --lyrics "[verse]
夜色温柔
[chorus]
我想飞 向前飞" \
  --seed 7 --max-frames 3000 --num-steps 30 --out out/song.wav

Python API：

from music3_mnn.generate import generate
wave44, codes = generate(prompt, lyrics, seed=7, max_frames=3000,
                         out_path="out/song.wav")  # wave44: [2, samples] 44.1kHz

• 歌词：--lyrics 支持 [verse] / [chorus] / [bridge] / [guitar solo] 等段落标签；歌曲实际长度由歌词量决定（唱完即停），25 帧 = 1 秒，上限 6 分钟（--max-frames 9000）。
• AR 采样语义逐行对齐 sglang-faithful 实现（c0 CFG 1.5 + top-50 掩码 + MurmurHash3 种子采样），同 seed 可复现。

性能（Mac, Apple Silicon, 5P+10E, 48GB）

8 秒歌曲（200 AR 帧 / 30 DiT 步）：

质量：全部组件与 fp32 参考余弦 ≥0.9995；端到端输出与 MLX 电平/频谱同量级（详见 REPORT.md）。

已做的优化：分段 KV 缓存 + 掩码步进图（AR 重规划从每 token 降到每 64 token）、CPU numThread=4（默认近单线程）、DiT Euler 循环免 resize、flow 前释放 AR 模型内存（峰值 -10GB）。

仓库结构

src/music3_mnn/
  export_onnx.py   # PyTorch -> ONNX 子图导出（backbone/depth/cond_conv/dit/vocoder）
  qwen3.py         # 极简 Qwen3 torch 实现（显式 KV cache IO，用于导出）
  mnn_runtime.py   # MNN 低层 API 封装（backend/线程配置、形状处理）
  ar_driver.py     # AR 阶段：分段 KV 缓存、c0/depth 采样（对齐 MLX 移植语义）
  flow_driver.py   # flow 阶段：条件编码、CFG-Euler、跨窗口潜空间拼接
  generate.py      # 端到端入口
scripts/           # 校验（KV 链 cos=1.0）、基准、风格拉练
docs/mnn-perf-feedback.md  # 提交给 MNN 社区的性能/正确性反馈报告
REPORT.md          # 完整对比报告：方法、精度验证、根因分析

已知限制

• MNN --transformerFuseC4（默认开）会算坏带 attention mask 的图 —— 转换必须加 --transformerFuseC4 0（详见 docs/mnn-perf-feedback.md §1）。
• MNN Metal 后端形状一变就重建 pipeline（~18s），AR 骨干只能跑 CPU。
• MNN 的 Metal4 tensor API 快路径在 macOS SDK 26.5 上编译失败回退慢内核，DiT 慢 MLX ~4×。
• vocoder 必须 fp32 计算（DAC 内部数值超 fp16 范围会溢出为 ±1）。