LiuBai-RWKV — an RWKV-7 song language model (checkpoints)

Text (style caption + lyrics) → full song with vocals, 48 kHz stereo. The autoregressive model is a 3 B-parameter RWKV-7 (linear attention, constant-size state, no KV cache), continued from BlinkDL's rwkv-g1k-3b-temp-5441 and trained on ~93 k Suno songs in YuE2's semantic token space. The tokens are rendered to audio by YuE2's non-autoregressive stage and VAE.

This repository holds the weights and the training logs. The code, the plan and the dated experiment log live on GitHub: https://github.com/xiaol/LiuBai-RWKV (checkpoints here match commit dc1fccb, 2026-09-20).

Everything is a work in progress and the numbers below are the honest ones: genre, instrumentation and language follow the caption and the vocals are fluent, but the sung words are mostly not the prompt lyrics yet. Lyric adherence is the open problem.

🎧 Listen (R3.1, milestone 2026-09-20)

20 prompts the model never saw (style caption + lyrics of real Suno songs from suno-94k shard 2). Tokens sampled by the R3.1 checkpoint (T 1.0, top-p 0.95, section-interleaved decoding), rendered with YuE2 NAR + our R1-a LoRA + YuE2-Vae, MP3 ≈ 128 kbps. Prompt text and generated section boundaries are in the matching .json under samples/r31_s2/. Lyric WER = word error rate of an ASR transcript against the prompt lyrics (lower is better; the original Suno recordings score 0.47 on the same ASR).

style (from the prompt)	lyric WER	R3.1 (this release)
disco	0.90
high-energy industrial metal, 120 bpm, minor key	0.90
bass-heavy tech house, punchy drums, dark bassline	0.91
punk	0.93
breakbeat	0.94
drill rap, male voice, bounce drop, epic	0.94
groovy bossa nova deep-house R&B chillwave, breathy vocals	0.95
motown	0.97
male vocalist, breathy verses, powerful chorus, acoustic guitar	0.98
male vocals, reggaeton	1.00
gregorian chant, latin gospel, cathedral organ	1.05
cello	1.09

All 20 songs, the same prompts from the previous whole-song model (G3, samples/g3_s2/), token-repeat rates and the original songs' scores: samples/README.md. R3.1 wins 15 of 20 prompts against G3 and has no collapsed songs.

What is in this repository

path	what	size	result
stage1_sec/rwkv-final.pth	R3.1 RWKV-7 3B, section-interleaved format (recommended)	6.2 GB	teacher-forced lyric binding 0.106 nats; audio WER vs lyrics 1.05; lyric-word overlap 10 %; MERT cos to original 0.79
stage1_yue2/rwkv-final.pth	G3 RWKV-7 3B, whole-song format (previous milestone, for comparison)	6.2 GB	binding 0.0005 nats (ignores the lyrics); audio WER 1.37; cos 0.80
joint/v2/head_best.pt	R1-a inverse tokenizer head: MERT-v2 features → YuE2 semantic tokens	447 MB	unseen real songs round trip: WER vs original transcript 0.33, MERT cos 0.981 (public v4 head: 0.50 / 0.971)
joint/v2/lora_best.pt	R1-a LoRA (rank 32) on the YuE2 NAR branch, trained jointly with the head	141 MB	used for every rendered sample above
*/train_log.txt, */step_log.txt, */events.txt	RWKV-LM run configs, per-10-step loss, checkpoint events	small
stage1_yue2/rwkv-init.resize.json	how the 65,536-row world vocab was resized to 98,816 (special ids)	small
yue2_layout.py	the token layout (control ids, YuE2 code offset)	small
samples/	20 R3.1 MP3s + 8 G3 MP3s with prompt JSON, browser player	80 MB

The RWKV checkpoints are plain RWKV-LM state_dict files (bf16): n_layer 32, n_embd 2560, head_size 64, vocab 98,816, ctx 8,192, 3.12 B parameters. They load with RWKV-LM v7 (RWKV-v7/train_temp, commit 9a75f9f + the patch in the GitHub repo) or with the generation script below. Not included: the R3.2 line-level probe (stage1_line, still being gated) and the earlier X-Codec run.

How it works

style caption + lyrics ──► RWKV-7 3B (ctx 8,192, vocab 98,816)
                              │  YuE2 semantic tokens, 25 Hz, 32,768-way
                              ▼
                   YuE2 NAR (+ our R1-a LoRA) ──► YuE2-Vae ──► 48 kHz stereo

Token space. YuE2 never released its audio encoder, so we trained one. R1-a is an inverse tokenizer: MERT-v2-FullSong features (layers 12, 16, 20, 23) of real audio → YuE2 semantic tokens. It is trained with a real-audio joint teacher: head → straight-through codec embeddings → frozen YuE2 AR+NAR (+LoRA) → flow-matching loss against the song's true VAE latents, plus a minted-song cross-entropy anchor (6,000 steps, window 256, 1,689 real + 4,750 minted tracks; joint/v2/config.json). This is what makes it possible to train a new autoregressive model in YuE2's token space on real songs.

Corpus. suno-94k (94,174 Suno songs, 5,423 h) tokenized with R1-a into YuE2 tokens: 93,058 songs. Lyric↔time alignment for the 40 k songs whose audio was kept, by MMS forced alignment at ≈ 600× realtime; 26 k of them are in the R3.1 training set.

Token layout (yue2_layout.py). RWKV World vocab (65,536) + control ids + the 32,768 YuE2 codes, padded to 98,816:

65536 <EOD>   65537 <SOA>   65538 <EOA>   65540 <yue2codec>   65541 <sos>
65544 .. 98311  = YuE2 semantic code 0 .. 32767  (65544 + code)

• G3 (whole-song): world_tokens("[Genre] caption\n[Lyrics]\nlyrics\n") EOD SOA YUE2CODEC codes… EOA
• R3.1 (section-interleaved): header EOD (SOS <lyric stanza> SOA YUE2CODEC <codes of that stanza> EOA)* 0 — each lyric stanza is followed by its own audio tokens, so the text→audio distance is seconds instead of a whole song. Generation mirrors it: decode stanza by stanza, each section ends when the model emits EOA (50–1,500 frames = 2–60 s).

Training (RWKV-LM v7 train_temp, 4× A100 40 GB, bf16, DeepSpeed ZeRO-2, micro-batch 1 × 4 GPUs, ctx 8,192, AdamW β 0.9/0.99, wd 0.1, grad-clip 1):

run	init	data	tokens	LR	wall
cb0 stage-1 (not released)	rwkv-g1k-3b-temp-5441	X-Codec codebook-0 tokens	20 k steps		frozen 2026-09-15
G3 stage1_yue2	cb0 step-20000, vocab resized	whole-song YuE2 format, 93 k songs	1.28 B	3e-5 → 3e-6	≈ 33 h
R3.1 stage1_sec	G3 final	section format, 26 k aligned songs, 2 passes	0.47 B	2e-5 → 2e-6	≈ 13.5 h

Evaluation. ASR (HeartTranscriptor) WER of the rendered audio against the prompt lyrics and against the original's transcript, MERT cosine to the original for style, and a teacher-forced binding test: code cross-entropy with the matched lyrics minus the same song with permuted lyrics. G3's gap is 0.0005 nats/frame (it ignores the text); R3.1 pays 0.106 nats/frame for the wrong stanza on all 40 test songs. That binding is real but far too weak to steer 32k-way sampling, hence the word-salad vocals.

Using the checkpoints

The scripts referenced here are in the GitHub repo; the repository layout there expects these files under out/, so download into out/:

git clone https://github.com/xiaol/LiuBai-RWKV.git rwkv-music && cd rwkv-music
hf download xiaol/LiuBai-RWKV --local-dir out --exclude "samples/*"      # ≈ 13 GB

# vendored upstream code (exact commits in tools/patches/README.md)
git clone https://github.com/BlinkDL/RWKV-LM.git tools/RWKV-LM && git -C tools/RWKV-LM checkout 9a75f9f
git -C tools/RWKV-LM apply ../../tools/patches/rwkv-lm-9a75f9f-train_temp.patch
git clone https://github.com/multimodal-art-projection/YuE.git tools/YuE && git -C tools/YuE checkout 0edaf2f

# 1) sample YuE2 tokens for 20 suno-94k shard-2 prompts with R3.1 (section decoding), one GPU, ~13 tok/s
python scripts/rwkv_generate.py --ckpt out/stage1_sec/rwkv-final.pth --sections 1 --shard 2 --n 20 --out out/gen/r31_s2_sec
# 2) render tokens → 48 kHz stereo FLAC with YuE2 NAR + R1-a LoRA + YuE2-Vae (~12 GiB GPU)
venvs/yue2/bin/python tools/yue2_render_tokens.py --gen out/gen/r31_s2_sec --nar-lora out/joint/v2/lora_best.pt --tag rwkv_r31

Other weights needed (not ours, download separately): YuE2-3B (AR/NAR) and YuE2-Vae from M-A-P, MERT-v2-FullSong (for the inverse tokenizer and the style metric), HeartTranscriptor (ASR for the WER metric). Two Python environments are used: one for RWKV-LM (CUDA kernels, torch 2.6) and one for YuE2/MERT/ASR; see the GitHub README.

To tokenize your own audio into YuE2 tokens: extract MERT-v2 layers 12/16/20/23 (tools/yue2_extract_mert.py) and run the head (head_best.pt contains model, cfg, layers; loader in scripts/train_joint_teacher.py / tools/yue2_roundtrip.py).

Status and next step (2026-09-21)

The R3.1 train loss dropped only when the second pass over the 26 k aligned songs began and the held-out loss did not follow, so the limit is aligned data, not steps. R3.2 is running: re-stream and align the remaining 54 k songs (→ ~80 k aligned), line-level interleaving (one SOS block per lyric line, 1.8–8.6 s of audio each), LR 3e-5, section-only data. Gate: the permuted-minus-matched gap must clearly beat 0.106 before the full run. If the transcript overlap still stays under ~30 %, the binding will be enforced with an objective on rendered audio (ASR/CTC reward or the joint teacher's flow loss). Dated entries: docs/STATUS_LOG.md on GitHub.

License and attribution

• The RWKV-7 checkpoints start from BlinkDL's RWKV-7 g1k weights (Apache-2.0) and were trained on tokens of the suno-94k dataset.
• joint/v2/lora_best.pt is a LoRA on YuE2-3B and head_best.pt consumes MERT-v2 features and was distilled through YuE2; YuE2-3B, YuE2-Vae and MERT-v2 are CC BY-NC 4.0, and rendering audio needs YuE2 anyway. The whole repository is therefore released under CC BY-NC 4.0 (non-commercial). A permissively licensed renderer is on the plan but does not exist yet.
• The generated samples are ours; the original Suno recordings used as prompts are not redistributed.

Credits: RWKV-7 by BlinkDL (RWKV-LM); YuE / YuE2 by M-A-P for the token space, NAR and VAE; Mothersuperior v4 as the reference inverse tokenizer and joint-training recipe; MERT-v2 (M-A-P); MMS forced alignment (Meta); HeartTranscriptor for ASR; the suno-94k dataset.