Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
119 lines
5 KiB
Python
119 lines
5 KiB
Python
#!/usr/bin/env python3
|
|
"""Dump the YuE2 NAR reference tensors consumed by yue2_nar_parity_probe.
|
|
|
|
The probe feeds the engine the exact prefix, codec tokens and acoustic noise the
|
|
Python reference integrated, so any difference it reports is the port rather
|
|
than sampling RNG: ABC and semantic sampling both run greedy (top_k = 1).
|
|
|
|
Writes into --out-dir:
|
|
prefix.i32 token prefix the acoustic flow conditions on
|
|
codec.i32 codec tokens, one per acoustic frame
|
|
noise.f32 [frames, 64] fp32 noise the reference integrated
|
|
latents_ref.f32 [frames, 64] fp32 reference latents
|
|
abc.txt score the plan used (empty for cot=off and provided scores)
|
|
metadata.json seed / ode_steps / context / latent_dim / shapes
|
|
|
|
The noise draw mirrors yue2.nar.song_chunks (one CPU fp32 randn seeded with the
|
|
request seed), and steps/context come from the pipeline generation config, so
|
|
the probe can reproduce the same integration path.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
def import_yue(reference_root: Path) -> None:
|
|
sys.path.insert(0, str(reference_root / "src"))
|
|
|
|
|
|
def write_i32(path: Path, values) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
np.asarray(list(values), dtype="<i4").tofile(path)
|
|
|
|
|
|
def write_f32(path: Path, values) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
np.asarray(values, dtype="<f4").reshape(-1).tofile(path)
|
|
|
|
|
|
def resolve_vae(model: Path, vae: Path | None) -> Path:
|
|
if vae is not None:
|
|
return vae
|
|
sibling = model.parent / "YuE2-Vae"
|
|
if sibling.is_dir():
|
|
return sibling
|
|
raise SystemExit("pass --vae: no YuE2-Vae directory next to --model")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--reference-root", type=Path, default=Path("reference/YuE"))
|
|
parser.add_argument("--model", type=Path, required=True)
|
|
parser.add_argument("--vae", type=Path, default=None)
|
|
parser.add_argument("--out-dir", type=Path, required=True)
|
|
parser.add_argument("--style", default="Upbeat electronic pop with a driving bass line.")
|
|
parser.add_argument("--lyrics", default="[verse]\nCity lights are calling out\n")
|
|
parser.add_argument("--abc-file", type=Path, default=None,
|
|
help="External ABC score; skips ABC generation entirely.")
|
|
parser.add_argument("--cot", default="full", choices=("off", "melody", "full"))
|
|
parser.add_argument("--seed", type=int, default=1234)
|
|
parser.add_argument("--semantic-max-tokens", type=int, default=9000,
|
|
help="Cap codec frames: the reference NAR attention is quadratic in frames.")
|
|
parser.add_argument("--device", default="cuda")
|
|
args = parser.parse_args()
|
|
|
|
import_yue(args.reference_root.resolve())
|
|
from yue2 import YuE2Pipeline
|
|
from yue2.protocol import Sampling
|
|
|
|
vae = resolve_vae(args.model, args.vae)
|
|
abc_text = args.abc_file.read_text(encoding="utf-8") if args.abc_file is not None else None
|
|
# Greedy everywhere: token choice carries no RNG, so the probe only has to
|
|
# reproduce the acoustic flow.
|
|
greedy_abc = Sampling(0.7, 0.9, 1, 1.005, 100, 32, 4096)
|
|
greedy_semantic = Sampling(1.0, 0.95, 1, 1.2, 50, 200, args.semantic_max_tokens)
|
|
|
|
with YuE2Pipeline.from_pretrained(str(args.model), vae=str(vae), device=args.device,
|
|
progress=False) as pipe:
|
|
plan = pipe.plan(args.style, args.lyrics, cot=args.cot, seed=args.seed, abc=abc_text,
|
|
abc_sampling=greedy_abc)
|
|
semantic = pipe.generate_semantic(plan, sampling=greedy_semantic)
|
|
steps = int(pipe.generation_config.ode_steps)
|
|
context = int(pipe.generation_config.context)
|
|
generator = torch.Generator(device="cpu").manual_seed(int(args.seed))
|
|
noise = torch.randn((len(semantic.tokens), 64), dtype=torch.float32, device="cpu",
|
|
generator=generator)
|
|
latents = pipe.synthesize(semantic)
|
|
|
|
args.out_dir.mkdir(parents=True, exist_ok=True)
|
|
write_i32(args.out_dir / "prefix.i32", plan.prefix)
|
|
write_i32(args.out_dir / "codec.i32", semantic.tokens)
|
|
write_f32(args.out_dir / "noise.f32", noise.numpy())
|
|
write_f32(args.out_dir / "latents_ref.f32", latents)
|
|
(args.out_dir / "abc.txt").write_text(plan.abc or "", encoding="utf-8")
|
|
metadata = {
|
|
"seed": args.seed,
|
|
"cot": args.cot,
|
|
"abc_source": "provided" if abc_text is not None else ("none" if args.cot == "off" else "generated"),
|
|
"ode_steps": steps,
|
|
"context": context,
|
|
"latent_dim": int(np.asarray(latents).shape[1]),
|
|
"codec_frames": int(len(semantic.tokens)),
|
|
"latent_frames": int(np.asarray(latents).shape[0]),
|
|
"style": args.style,
|
|
"lyrics": args.lyrics,
|
|
"model": str(args.model),
|
|
"vae": str(vae),
|
|
}
|
|
(args.out_dir / "metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8")
|
|
print(json.dumps(metadata, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|