Buckets:
| """Prep real tracks for the NAR-teacher loss: MERT-v2-FullSong L20 @25Hz (instance-normed later), VAE latents [T,64], cot=off prefix ids. | |
| -> /workspace/real/prep/<name>/{mert.npy,lat.npy,prefix.npy}. Held-out track kept too (used only for eval).""" | |
| import os, glob, re, json, numpy as np, torch, soundfile as sf | |
| from scipy.signal import resample_poly | |
| from math import gcd | |
| os.environ.setdefault("HF_HOME","/workspace/hf") | |
| from transformers import AutoModel, AutoFeatureExtractor | |
| from yue2.modeling_vae import YuE2VAE | |
| from yue2.protocol import SongRequest, token_prefixes | |
| from yue2.tokenization_yue2 import YuE2TextTokenizer | |
| SRC="/workspace/real/artist"; OUT="/workspace/real/prep"; os.makedirs(OUT,exist_ok=True); dev="cuda" | |
| snap=lambda n: glob.glob(f"/workspace/hf/hub/models--m-a-p--{n}/snapshots/*")[0]; tok=YuE2TextTokenizer(snap("YuE2-3B")+"/qwen.tiktoken") | |
| proc=AutoFeatureExtractor.from_pretrained("m-a-p/MERT-v2-FullSong", trust_remote_code=True); mert=AutoModel.from_pretrained("m-a-p/MERT-v2-FullSong", trust_remote_code=True).to(dev).eval() | |
| vae=YuE2VAE.from_pretrained(snap("YuE2-Vae"), decoder_only=False, device=dev, local_files_only=True) | |
| def mert_l20(mono24): | |
| CH=24000*30; chunks=[mono24[s:s+CH] for s in range(0,len(mono24),CH)]; chunks=[c for c in chunks if len(c)>=24000]; full=[c for c in chunks if len(c)==CH]; tail=[c for c in chunks if len(c)<CH]; feats=[] | |
| with torch.inference_mode(), torch.autocast("cuda",dtype=torch.bfloat16): | |
| for group in ([full] if full else [])+[[c] for c in tail]: | |
| inp={k:v.to(dev) for k,v in proc(group, sampling_rate=24000, return_tensors="pt").items()}; feats.append(mert(**inp, output_hidden_states=True).hidden_states[20].reshape(-1,1024)) | |
| H=torch.cat(feats,0).float(); T25=int(round(len(mono24)/24000*25)); return torch.nn.functional.interpolate(H.T[None], size=T25, mode="linear", align_corners=False)[0].T.half().cpu().numpy() | |
| def latents(st48): | |
| out=[]; CH=48000*60 | |
| with torch.inference_mode(): | |
| for s in range(0,len(st48),CH): | |
| seg=st48[s:s+CH] | |
| if len(seg)<1920: break | |
| out.append(vae.encode(torch.tensor(seg.T[None]))[0].T.float().cpu()) | |
| return torch.cat(out,0).numpy() | |
| for f in sorted(glob.glob(f"{SRC}/*.flac")): | |
| name=os.path.basename(f)[:-5]; d=f"{OUT}/{name}" | |
| if os.path.exists(f"{d}/prefix.npy"): continue | |
| os.makedirs(d,exist_ok=True); a,sr=sf.read(f, dtype="float32"); a=np.stack([a,a],1) if a.ndim==1 else a | |
| g=gcd(sr,48000); st48=resample_poly(a,48000//g,sr//g,axis=0).astype(np.float32) if sr!=48000 else a | |
| g=gcd(sr,24000); m24=resample_poly(a.mean(1),24000//g,sr//g).astype(np.float32) | |
| M=mert_l20(m24); Z=latents(st48); n=min(len(M),len(Z)); np.save(f"{d}/mert.npy",M[:n]); np.save(f"{d}/lat.npy",Z[:n].astype(np.float32)) | |
| cap=open(f"{SRC}/{name}.txt").read().split("===LYRICS===")[0].replace("Global Metadata:","").strip(); style=" ".join(cap.split())[:1500] | |
| lyr=open(f"{SRC}/{name}.lyrics.txt").read().strip() if os.path.exists(f"{SRC}/{name}.lyrics.txt") else "[instrumental]" | |
| pre=token_prefixes(SongRequest(style=style,lyrics=lyr,cot="off",seed=1,id="real"),tok); np.save(f"{d}/prefix.npy",np.array(pre,dtype=np.int64)) | |
| print(f"{name}: frames {n} ({n/25/60:.1f} min) prefix {len(pre)} mert-lat frame diff {len(M)-len(Z)}", flush=True) | |
| print("PREP DONE", flush=True) | |
Xet Storage Details
- Size:
- 3.36 kB
- Xet hash:
- dba2045e5e2fcfe86baee58293f328544cfd161b3ed4b50c12931ae2c1e80a1a
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.