--- license: cc-by-nc-4.0 base_model: - m-a-p/YuE2-3B - m-a-p/MERT-v2-FullSong tags: - audio - music - yue2 - tokenizer - lora --- # yue2-mothersuperior-realaudio-tokenizer-v4 Real-audio tooling for [YuE2-3B](https://huggingface.co/m-a-p/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 | file | what | |---|---| | `tokenizer_head_joint_v4.pt` / `.safetensors` / `.bf16.safetensors` | MERT-v2-FullSong layer-20 features (per-track instance-normalised, 25 Hz) → 32,768 YuE2 semantic codes. 8-layer transformer, d=512, 512-frame windows. Held-out exact match on YuE2's own songs: 16.1% top-1 (near-miss codes render almost identically; ear tests of NAR round-trips sit around 95%). | | `nar_lora_joint_v4.pt` / `.safetensors` / `.bf16.safetensors` / `_comfyui.safetensors` | rank-32 LoRA on `nar_self_attn.{q,k,v,o}_proj` + `nar_mlp.{gate,up,down}_proj` (28 layers) + full `vae2llm`/`llm2vae`. Trained jointly with the head on real audio. | | `scripts/` | the training loop and inference scripts (below). `scripts/ckpt_io.py` loads either format. | ### 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%). The previously published 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`](https://github.com/multimodal-art-projection/YuE) (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](https://huggingface.co/datasets/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](https://huggingface.co/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 ```bash # 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 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 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). | pair | audio-loss weight | latent loss | minted top-1 | LTAS dB rock / metal / bass | mel L1 dB rock / metal / bass | stereo rock / metal / bass | |---|---|---|---|---|---|---| | v5 | 0 (latent only) | 0.981 | 18.9% | 1.56 / 2.07 / 2.27 | 4.36 / 4.59 / 4.44 | 0.77 / 0.69 / 0.90 | | v6 | 0.3 | 0.981 | 18.8% | 1.40 / 1.94 / 2.09 | 4.18 / 4.55 / 4.36 | 0.78 / 0.76 / 0.92 | | v7 | 1.0 | 0.984 | 18.7% | 1.44 / 1.28 / 1.92 | 4.50 / 4.54 / 4.39 | 0.63 / 0.63 / 0.84 | | **v8** | **2.0** | 0.983 | 19.0% | 1.28 / 1.75 / 1.84 | 4.28 / 4.49 / 4.51 | 0.76 / 0.70 / 0.92 | | **v9** | **4.0** | 0.984 | 18.7% | 1.23 / 1.49 / 1.92 | 4.29 / 4.48 / 4.44 | 0.74 / 0.65 / 0.88 | 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_v6.py`](scripts/joint_v6.py), the recovered original audio-loss variant of `scripts/joint.py` (env `AUX_W`, `AUX_TMAX`, `AUX_FR`). Read the [setup and historical-reproduction notes](scripts/joint_v6_README.md) first: this original script uses `.pt` checkpoints and the old pod paths; unlike the other published scripts, it has not been adapted to `ckpt_io`. The required neighbor tables are included under `assets/`.