yue2-mothersuperior-realaudio-tokenizer-v4

Real-audio tooling for YuE2-3B: the audio → semantic-token encoder YuE2 doesn't ship, plus a NAR-branch LoRA so the decoder renders real-production latents. Together they let you tokenize your own recordings, LoRA-tune YuE2's AR on an artist, and generate new songs or covers.

Files

Safetensors layout

Same weights as the .pt files, bit-exact in fp32 (the .bf16 variants are half the size; head top-1 agreement with fp32 is 98.6%). All scripts accept either extension via scripts/ckpt_io.load_ckpt(path), which returns the same dict the .pt files hold.

• Head: the plain state_dict of the 8-layer encoder (inp.*, pos, enc.layers.{0..7}.*, norm.*, head.*), 103 tensors.
• NAR LoRA: layers.{0..27}.nar_self_attn.{q,k,v,o}_proj.lora_A [32, in] and .lora_B [out, 32], plus layers.{i}.nar_mlp.{gate,up,down}_proj.lora_{A,B}. Apply as W += lora_B @ lora_A (scale 1.0, no alpha) to the matching model.layers[i] Linear of YuE2-3B. vae2llm.{weight,bias} and llm2vae.{weight,bias} are full replacement weights for those two Linear layers, not LoRA deltas. Rank and the delta rule are also in the file metadata.

ComfyUI

nar_lora_joint_v4_comfyui.safetensors is the same NAR LoRA in ComfyUI's native YuE2 layout: load it with the standard LoraLoader on the MODEL output of the YuE2 checkpoint loader (Comfy-Org/YuE2, yue2_3b_bf16.safetensors), strength 1.0; the CLIP input is untouched by this file. ComfyUI stores the NAR branch as diffusion_model.model.layers.{i}.self_attn.qkv_proj / o_proj and mlp.gate_up_proj / down_proj (q,k,v and gate,up fused), so the separate q/k/v and gate/up LoRAs are packed as one block-diagonal LoRA per fused matrix (lora_up/lora_down, no alpha = scale 1.0; up @ down equals the stacked deltas exactly). The full vae2llm / llm2vae replacements ship as .diff / .diff_b against the m-a-p base.

Trained on 4,765 YuE2 self-generated songs, then adapted to real audio. If your material sounds off, rerun joint.py on your own audio (step 3 below).

Requirements

Python 3.12 venv with yue2-infer (commit 92a73cc7), torch 2.10 + cu128, torchaudio 2.10, transformers, soundfile, scipy, safetensors, demucs; HF_HOME with m-a-p/YuE2-3B, m-a-p/YuE2-Vae, m-a-p/MERT-v2-FullSong. A 24 GB GPU is enough for every stage (14–18 GB measured with gradient checkpointing).

You also need the minted regularizer pack: regularizer/minted_regularizer_pack.pt from Mothersuperior/yue2-minted-corpus (~100 MB, 4,732 YuE2-generated songs as {name, src, style, lyrics, codec}). The AR trainer draws 50% of its songs from it so a small artist set cannot collapse YuE2's token grammar; the minted_val items are the held-out check whose loss should stay flat. The full corpus (audio + tokens + latents) is in the same dataset if you want to retrain the head.

Paths are hard-coded to our pod layout (/workspace/tok/full, /workspace/real/..., /workspace/yue2-corpus/tracks, /workspace/real/ar/dataset.pt). Recreate that layout or edit the constants at the top of each script.

Train an artist LoRA (folder of songs → LoRA)

Per song you need .flac, .lyrics.txt (full lyrics with [verse]/[chorus]/[bridge]/... tags — truncated lyrics ruin structure), and .txt = a style caption starting with your trigger phrase (e.g. xyzq, in the style of xyzq. ).

1. python prep_real.py — MERT features, VAE latents and the prompt prefix per song.
2. python cursor_prep.py — Demucs vocal stem → MMS forced alignment of the lyrics → lyric-cursor targets (automatic).
3. (optional, recommended for a new artist/era) HOLD_TRACK= python joint.py joint_mine 3000 1 1 tokenizer_head_joint_v4.pt nar_lora_joint_v4.pt — adapts head + NAR to your audio. Otherwise use the v3 files as-is.
4. python ar_prep.py — tokenizes your songs and merges the regularizer pack into dataset.pt (ar_prep.py expects the pack's records; point it at the downloaded file).
5. SCHED_STEPS=3000 CK_FROM=600 CK_EVERY=200 python ar_lora_cursor.py my_lora 1600 64 0.5 none 1e-4 0.08 — rank-64 AR LoRA, 50/50 artist vs minted, lyric-cursor weight 0.08, checkpoints at 600/800/1000/1200/1400/1600. Do not train longer: past ~1,500 steps the model memorises the songs.
6. LADDER_STYLE_TRACK= LADDER_LYRICS= bash ladder.sh my_lora nar_lora_joint_v4.pt — renders one fixed prompt from every checkpoint (optional FINALS= with lines tag style_track lyrics seed). Pick by ear (ours: step 800).

Decoder

Consider decoding with Mothersuperior/YuE2-Vae-merge-0.666: a weight merge of YuE2-Vae (0.666) and YuE2-Vae-legacy (0.334). The two releases share one encoder and only differ in the decoder; the merge sits between the clean default decoder and the more musical legacy decoder. Drop-in: YuE2Pipeline.from_pretrained("m-a-p/YuE2-3B", vae="Mothersuperior/YuE2-Vae-merge-0.666"), or point ar_generate.py at it.

Inference

# new song: style caption with your trigger + lyrics, score-free
python ar_generate.py my_lora/step-800.pt nar_lora_joint_v4.pt out_tag <style_track> lyrics.txt 12
# cover: transcribe any recording with SheetSage2 (--melody-only), then
ABC_FILE=score.abc COT=melody STRIP_TEMPO_KEY=1 python ar_generate.py my_lora/step-800.pt nar_lora_joint_v4.pt cover_tag <style_track> lyrics.txt 21
# LoRA strength: AR_SCALE=0.77 ...   stock model control: pass `none` for either LoRA

ar_generate.py folds both LoRAs into the base weights and runs YuE2's own pipeline, so the stock sampler, CFG and VAE apply unchanged.

Other scripts

extract_full.py + train_v2.py (retrain the head on the minted corpus), nar_lora.py (NAR LoRA alone), teacher_train.py (head fine-tune with the NAR as teacher), ar_lora.py (AR LoRA without the cursor), build_reg_pack.py (rebuild the regularizer pack from a corpus).

Weights derive from YuE2-3B (CC BY-NC 4.0): non-commercial use only. Scripts are provided as-is.

v5 head (2026-09-16): tokenizer_head_v5_30k.safetensors (+ .bf16)

Same architecture and input as the v4 head, retrained on the full 14,547-track minted corpus with a better schedule: optimi AdamW (lr 5e-4, betas 0.9/0.99, weight decay 0.01), 30,000 steps, cosine. Held-out token accuracy on minted tracks: top-1 18.9%, top-5 45.0%, top-1-or-neighbour 34.1% (v4-era head: 15.9 / 39.2 / 30.9). Drop-in for the v4 head in ckpt_io.py and in the ComfyUI FS_Audio Dataset Builder (models/fs_audio/). For the matching decoder LoRA see the v5 joint pair below.

v5 joint pair (2026-09-16): tokenizer_head_joint_v5 + nar_lora_joint_v5 (+ .bf16, + _comfyui)

The v5 head co-trained for 3,000 steps with a rank-32 decoder LoRA on 128 real tracks from three artists (rock, metal, bass music) plus 4,000 minted anchor tracks, using the decoder's flow loss on the real latents as the teacher (previous joint rounds used a single artist's catalog). Held-out real-audio decoder loss 1.056 -> 0.981; minted top-1 unchanged (~18.9%, so the head stays universal); token repeat rate on real audio 9.3% -> 3.0%. Use the pair together: head for tokenizing real audio, LoRA on the decoder when rendering those tokens. nar_lora_joint_v5_comfyui.safetensors is the ComfyUI-native layout (load with a LoRA loader, model strength 1.0).

v8 / v9 joint pairs (2026-09-16): training against the audio, not just the latents

tokenizer_head_joint_v8 + nar_lora_joint_v8 and tokenizer_head_joint_v9 + nar_lora_joint_v9 (each fp32, .bf16, and _comfyui for the LoRA). The v5 pair stays available above; these are additions.

What changed. Every joint round so far trained the head and decoder LoRA against a loss in VAE-latent space (flow matching). That proxy was not tracking what we heard: v1 -> v3 -> v5 cut the latent loss from 1.159 to 0.981 while audio-domain distances to the original recordings barely moved. Starting with v6 the trainer adds an audio-domain loss: at each real-audio step the decoder's clean-latent estimate is decoded through the frozen (differentiable) YuE2 VAE and compared with the same seconds of the original recording using log-mel L1, multi-resolution STFT (spectral convergence + log-magnitude, windows 512/1024/2048) and a stereo-width term, applied at noise levels <= 0.4. Gradients reach the decoder LoRA directly and the head through the straight-through token choice. Data and everything else are identical to v5 (128 real tracks from three artists: rock, metal, bass music; 4,000 minted anchor tracks; 3,000 steps; init = tokenizer_head_v5_30k + nar_lora_joint_v4). The idea of penalising the rendered audio against the real production came from Kytra listening to the v5 reconstructions.

Weight sweep, measured on the three held-out tracks (whole songs) against the originals. LTAS = long-term average spectrum distance (the "production fingerprint"; this is the number that tracked the listening results). Stereo = L/R correlation (originals: 0.84 / 0.71 / 0.87).

Listening (Kytra): v7 clearly better than v6, v8 clearly better than v7, v9 preferred overall. Minted top-1 stays flat across the sweep, so the head remains universal; the latent loss moves by 0.003, so the audio term is not fighting the latent objective. Mel L1 did not follow the listening results, LTAS did. Token-choice errors (occasional out-of-tune notes) are unchanged by this loss; that is a head-accuracy problem.

Use: tokenize real audio with the v9 (or v8) head and load the matching nar_lora_joint_v9_comfyui / _v8_comfyui on the decoder at model strength 1.0. Trainer: scripts/joint.py lineage; the audio-loss variant is joint_v6.py in our tree (env AUX_W, AUX_TMAX, AUX_FR).