LiuBai-RWKV

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).

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

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):

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.