Add safetensors releases (fp32 + bf16) with named LoRA tensors; scripts load either format

README.md

@@ -19,9 +19,18 @@ an artist, and generate new songs or covers.
 ## Files
 | file | what |
 |---|---|
-| `tokenizer_head_joint_v4.pt` | 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` | 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). |
+| `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` | 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%). 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.
 
 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).
 

nar_lora_joint_v4.bf16.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e67ca682f1070d2abb45cf9ac90d542fe4ce5286b87aed97f02d9a4acd6b6f36
+size 70301856

nar_lora_joint_v4.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c99b788ca1167a1532b35fff2085d8d14549c8eaa87f9eda7ce68c041841fcc
+size 140560552

scripts/ar_generate.py

@@ -1,5 +1,6 @@
 """Generate songs with AR LoRA + joint_v1 NAR LoRA through YuE2's own pipeline (cot=off).
 usage: ar_generate.py <ar_lora.pt> <nar_lora.pt> <out_tag> <style_from_track_name> <lyrics_file> <seed>"""
+import sys, os as _os; sys.path.insert(0,_os.path.dirname(_os.path.abspath(__file__))); from ckpt_io import load_ckpt
 import os, sys, glob, math, json, numpy as np, torch, torch.nn as nn
 os.environ.setdefault("HF_HOME","/workspace/hf")
 from yue2 import YuE2Pipeline
@@ -19,10 +20,10 @@ def merge(attn_name, mlp_name, tensors, scale=1.0):
     return n_merged
 with torch.no_grad():
     AR_SCALE=float(os.environ.get("AR_SCALE","1.0"))
-    if AR_CK!="none": ar=torch.load(AR_CK,map_location=dev); print(f"merged AR linears (scale {AR_SCALE}):", merge("self_attn","mlp",ar["lora"],AR_SCALE), flush=True)
+    if AR_CK!="none": ar=load_ckpt(AR_CK,dev); print(f"merged AR linears (scale {AR_SCALE}):", merge("self_attn","mlp",ar["lora"],AR_SCALE), flush=True)
     else: print("AR: stock (no LoRA)", flush=True)
     if NAR_CK!="none":
-        nar=torch.load(NAR_CK,map_location=dev); print("merged NAR linears:", merge("nar_self_attn","nar_mlp",nar["lora"]), flush=True)
+        nar=load_ckpt(NAR_CK,dev); print("merged NAR linears:", merge("nar_self_attn","nar_mlp",nar["lora"]), flush=True)
         model.vae2llm.load_state_dict({k:v.to(torch.bfloat16) for k,v in nar["io"]["vae2llm"].items()}); model.llm2vae.load_state_dict({k:v.to(torch.bfloat16) for k,v in nar["io"]["llm2vae"].items()})
     else: print("NAR: stock (no LoRA)", flush=True)
 model.eval(); print("LoRAs loaded", flush=True)

scripts/ar_prep.py

@@ -1,5 +1,6 @@
 """AR dataset in score-free (cot=off) format: prefix ids (instructions+tags+lyrics) + codec tokens.
 Artist tracks: tokens from the joint_v1 head. Minted: true semantic tokens. -> /workspace/real/ar/dataset.pt"""
+import sys, os as _os; sys.path.insert(0,_os.path.dirname(_os.path.abspath(__file__))); from ckpt_io import load_ckpt
 import os, glob, json, hashlib, numpy as np, torch, torch.nn as nn, sys
 os.environ.setdefault("HF_HOME","/workspace/hf")
 from yue2.protocol import SongRequest, token_prefixes
@@ -12,7 +13,7 @@ class Tok(nn.Module):
         super().__init__(); s.inp=nn.Linear(din,D); s.pos=nn.Parameter(torch.zeros(1,WIN,D))
         layer=nn.TransformerEncoderLayer(D,H,4*D,dropout=0.1,batch_first=True,norm_first=True,activation="gelu"); s.enc=nn.TransformerEncoder(layer,L); s.norm=nn.LayerNorm(D); s.head=nn.Linear(D,VOCAB)
     def forward(s,x): return s.head(s.norm(s.enc(s.inp(x)+s.pos[:,:x.shape[1]])))
-head=Tok(1024).to(dev).eval(); head.load_state_dict(torch.load(HEAD_CK,map_location=dev)["model"])
+head=Tok(1024).to(dev).eval(); head.load_state_dict(load_ckpt(HEAD_CK,dev)["model"])
 def instnorm(x): x=x.astype(np.float32); return (x-x.mean(0))/(x.std(0)+1e-5)
 @torch.no_grad()
 def predict(x):

scripts/ckpt_io.py

@@ -0,0 +1,18 @@
+"""Load our checkpoints from either the original .pt files or the .safetensors releases, returning the same dict layout the scripts expect.
+  head : {"model": state_dict, "cfg": {...}}
+  lora : {"lora": [A0,B0,A1,B1,...] (layer-major: nar_self_attn q,k,v,o then nar_mlp gate,up,down), "io": {"vae2llm": sd, "llm2vae": sd}, "rank": int}"""
+import torch
+NAR_MODS=[("nar_self_attn",n) for n in ("q_proj","k_proj","v_proj","o_proj")]+[("nar_mlp",n) for n in ("gate_proj","up_proj","down_proj")]
+def load_ckpt(path, map_location="cpu"):
+    if not str(path).endswith(".safetensors"): return torch.load(path, map_location=map_location, weights_only=False)
+    from safetensors.torch import load_file; from safetensors import safe_open
+    t=load_file(path, device=str(map_location))
+    with safe_open(path, "pt") as f: meta=f.metadata() or {}
+    if any(k.endswith(".lora_A") for k in t):
+        layers=sorted({int(k.split(".")[1]) for k in t if k.startswith("layers.")})
+        lora=[]
+        for L in layers:
+            for blk,proj in NAR_MODS: lora+=[t[f"layers.{L}.{blk}.{proj}.lora_A"], t[f"layers.{L}.{blk}.{proj}.lora_B"]]
+        io={m:{k.split(".",1)[1]:v for k,v in t.items() if k.startswith(m+".")} for m in ("vae2llm","llm2vae")}
+        return {"lora":lora, "io":io, "rank":int(meta.get("rank", lora[0].shape[0]))}
+    return {"model":t, "cfg":{"instnorm": meta.get("input","").find("instnorm=true")>=0}}

scripts/joint.py

@@ -3,6 +3,7 @@ usage: joint.py <name> <steps> <train_head 0|1> <train_lora 0|1> <init_head.pt>
 Real window: MERT -> head -> straight-through tokens -> (LoRA'd) NAR flow loss on true VAE latents. Head also gets minted soft-CE each step;
 in LoRA mode 25% of flow windows are minted (true tokens). Eval = held-out real flow loss (fixed windows/t/noise), minted top-1, repeat rate.
 Ends by rendering the held-out track with the best head+NAR."""
+import sys, os as _os; sys.path.insert(0,_os.path.dirname(_os.path.abspath(__file__))); from ckpt_io import load_ckpt
 import os, sys, glob, json, math, time, random, hashlib, numpy as np, torch, torch.nn as nn, torch.nn.functional as F, soundfile as sf
 from torch.utils.checkpoint import checkpoint
 os.environ.setdefault("HF_HOME","/workspace/hf"); torch.backends.cuda.matmul.allow_tf32=True
@@ -27,7 +28,7 @@ for layer in bb.layers:
         for n in names: l=LoRALinear(getattr(mod,n),RANK); setattr(mod,n,l); lora_params+=[l.A,l.B]
 model.vae2llm.float(); model.llm2vae.float(); io_params=list(model.vae2llm.parameters())+list(model.llm2vae.parameters())
 def load_lora(path):
-    ck=torch.load(path,map_location=dev)
+    ck=load_ckpt(path,dev)
     with torch.no_grad():
         for p,v in zip(lora_params,ck["lora"]): p.copy_(v.to(dev))
         model.vae2llm.load_state_dict({k:v.float() for k,v in ck["io"]["vae2llm"].items()}); model.llm2vae.load_state_dict({k:v.float() for k,v in ck["io"]["llm2vae"].items()})
@@ -39,7 +40,7 @@ class Tok(nn.Module):
         super().__init__(); s.inp=nn.Linear(din,D); s.pos=nn.Parameter(torch.zeros(1,WIN,D))
         layer=nn.TransformerEncoderLayer(D,H,4*D,dropout=0.1,batch_first=True,norm_first=True,activation="gelu"); s.enc=nn.TransformerEncoder(layer,L); s.norm=nn.LayerNorm(D); s.head=nn.Linear(D,VOCAB)
     def forward(s,x): return s.head(s.norm(s.enc(s.inp(x)+s.pos[:,:x.shape[1]])))
-head=Tok(1024).to(dev); head.load_state_dict(torch.load(INIT_HEAD,map_location=dev)["model"]); head.requires_grad_(bool(TRAIN_HEAD))
+head=Tok(1024).to(dev); head.load_state_dict(load_ckpt(INIT_HEAD,dev)["model"]); head.requires_grad_(bool(TRAIN_HEAD))
 groups=[]
 if TRAIN_HEAD: groups.append({"params":list(head.parameters()),"lr":LR_HEAD,"weight_decay":0.05})
 if TRAIN_LORA: groups+=[{"params":lora_params,"lr":LR_LORA,"weight_decay":0.0},{"params":io_params,"lr":LR_IO,"weight_decay":0.0}]
@@ -132,7 +133,7 @@ for st in range(1,STEPS+1):
         save_all("last")
 print(f"RESULT {NAME}: best real_nar {best:.4f} (start {e0:.4f})", flush=True)
 # ---- render held-out with best
-head.load_state_dict(torch.load(f"{OUT}/head_best.pt",map_location=dev)["model"]); head.eval()
+head.load_state_dict(load_ckpt(f"{OUT}/head_best.pt",dev)["model"]); head.eval()
 if TRAIN_LORA: load_lora(f"{OUT}/lora_best.pt")
 model.vae2llm.to(torch.bfloat16); model.llm2vae.to(torch.bfloat16)
 @torch.no_grad()

tokenizer_head_joint_v4.bf16.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce52f7e4b3e437f204bdac71db685d86f3912e4085659a711387de87cfeb5127
+size 85644544

tokenizer_head_joint_v4.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0117394bb7db3dc88e4cb62a5e6e909283e1042240a3678780dcf9396502327e
+size 171278496