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
277 lines
11 KiB
Python
277 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import random
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
REFERENCE_ROOT = REPO_ROOT / "reference" / "YuE"
|
|
DEFAULT_CASES = REPO_ROOT / "tests" / "yue2" / "yue2_warm_bench_cases.json"
|
|
DEFAULT_OUTPUT_ROOT = REPO_ROOT / "build" / "logs" / "yue2" / "python_baseline"
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Python reference YuE2 warmbench.")
|
|
parser.add_argument("--family", default="yue2")
|
|
parser.add_argument("--model", default="/home/leo/Desktop/YuE2/YuE2-3B")
|
|
parser.add_argument("--vae", default="/home/leo/Desktop/YuE2/YuE2-Vae")
|
|
parser.add_argument("--reference-root", type=Path, default=REFERENCE_ROOT)
|
|
parser.add_argument("--cases", type=Path, default=DEFAULT_CASES)
|
|
parser.add_argument("--case", action="append", default=[])
|
|
parser.add_argument("--backend", choices=("cuda", "cpu"), default="cuda")
|
|
parser.add_argument("--device", type=int, default=0)
|
|
parser.add_argument("--threads", type=int, default=8)
|
|
parser.add_argument("--warmup", type=int, default=0)
|
|
parser.add_argument("--iterations", type=int, default=1)
|
|
parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_ROOT)
|
|
parser.add_argument("--timing-file", type=Path, default=DEFAULT_OUTPUT_ROOT / "python_baseline.log")
|
|
parser.add_argument("--summary-file", type=Path, default=DEFAULT_OUTPUT_ROOT / "summary.json")
|
|
parser.add_argument("--local-files-only", action="store_true")
|
|
parser.add_argument("--log", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def resolve_path(path: Path | str) -> Path:
|
|
value = Path(path)
|
|
return value if value.is_absolute() else REPO_ROOT / value
|
|
|
|
|
|
def add_reference_path(reference_root: Path) -> None:
|
|
src = resolve_path(reference_root) / "src"
|
|
if not (src / "yue2" / "__init__.py").is_file():
|
|
raise RuntimeError(f"missing YuE2 reference package under {src}")
|
|
sys.path.insert(0, str(src.resolve()))
|
|
|
|
|
|
def configure_runtime(args: argparse.Namespace) -> str:
|
|
os.environ.setdefault("PYTHONHASHSEED", "0")
|
|
torch.set_num_threads(max(1, args.threads))
|
|
random.seed(0)
|
|
np.random.seed(0)
|
|
torch.manual_seed(0)
|
|
if args.backend == "cuda":
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("YuE2 warmbench requested CUDA, but torch.cuda.is_available() is false")
|
|
torch.cuda.set_device(args.device)
|
|
torch.cuda.manual_seed_all(0)
|
|
return f"cuda:{args.device}"
|
|
return "cpu"
|
|
|
|
|
|
def sync_device(device: str) -> None:
|
|
if device.startswith("cuda") and torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
|
|
|
|
def load_cases(path: Path) -> dict[str, dict[str, Any]]:
|
|
payload = json.loads(resolve_path(path).read_text(encoding="utf-8"))
|
|
if not isinstance(payload, dict):
|
|
raise RuntimeError("YuE2 warmbench cases must be a JSON object")
|
|
return payload
|
|
|
|
|
|
def selected_cases(args: argparse.Namespace, cases: dict[str, dict[str, Any]]) -> list[dict[str, Any]]:
|
|
names = args.case or list(cases)
|
|
out = []
|
|
for name in names:
|
|
if name not in cases:
|
|
raise RuntimeError(f"unknown YuE2 warmbench case: {name}")
|
|
case = dict(cases[name])
|
|
case.setdefault("id", name)
|
|
out.append(case)
|
|
return out
|
|
|
|
|
|
def make_generation_config(case: dict[str, Any]):
|
|
from yue2.protocol import GenerationConfig, Sampling
|
|
|
|
defaults = GenerationConfig()
|
|
abc = Sampling(
|
|
temperature=float(case.get("abc_temperature", defaults.abc.temperature)),
|
|
top_p=float(case.get("abc_top_p", defaults.abc.top_p)),
|
|
top_k=int(case.get("abc_top_k", defaults.abc.top_k)),
|
|
repetition_penalty=float(case.get("abc_repetition_penalty", defaults.abc.repetition_penalty)),
|
|
penalty_window=int(case.get("abc_penalty_window", defaults.abc.penalty_window)),
|
|
min_tokens=int(case.get("abc_min_tokens", defaults.abc.min_tokens)),
|
|
max_tokens=int(case.get("abc_max_tokens", defaults.abc.max_tokens)),
|
|
)
|
|
semantic = Sampling(
|
|
temperature=float(case.get("semantic_temperature", defaults.semantic.temperature)),
|
|
top_p=float(case.get("semantic_top_p", defaults.semantic.top_p)),
|
|
top_k=int(case.get("semantic_top_k", defaults.semantic.top_k)),
|
|
repetition_penalty=float(case.get("semantic_repetition_penalty", defaults.semantic.repetition_penalty)),
|
|
penalty_window=int(case.get("semantic_penalty_window", defaults.semantic.penalty_window)),
|
|
min_tokens=int(case.get("semantic_min_tokens", defaults.semantic.min_tokens)),
|
|
max_tokens=int(case.get("semantic_max_tokens", defaults.semantic.max_tokens)),
|
|
)
|
|
return GenerationConfig(
|
|
abc=abc,
|
|
semantic=semantic,
|
|
ode_steps=int(case.get("ode_steps", defaults.ode_steps)),
|
|
ode_method=str(case.get("ode_method", defaults.ode_method)),
|
|
context=int(case.get("context", defaults.context)),
|
|
version=str(case.get("version", defaults.version)),
|
|
)
|
|
|
|
|
|
def request_from_case(case: dict[str, Any]) -> dict[str, Any]:
|
|
request = {
|
|
"style": str(case["style"]),
|
|
"lyrics": str(case["lyrics"]),
|
|
"cot": str(case.get("cot", "full")),
|
|
"seed": int(case.get("seed", 831001)),
|
|
"id": str(case.get("id", "song")),
|
|
}
|
|
if "cfg_scale" in case:
|
|
request["cfg_scale"] = float(case["cfg_scale"])
|
|
if "abc" in case:
|
|
request["abc"] = str(case["abc"])
|
|
if "abc_file" in case:
|
|
request["abc"] = resolve_path(case["abc_file"]).read_text(encoding="utf-8")
|
|
return request
|
|
|
|
|
|
def summarize_audio(audio: np.ndarray, sample_rate: int) -> dict[str, Any]:
|
|
flat = np.asarray(audio, dtype=np.float32).reshape(-1)
|
|
if flat.size == 0:
|
|
raise RuntimeError("YuE2 warmbench produced empty audio")
|
|
return {
|
|
"sample_rate": int(sample_rate),
|
|
"channels": int(audio.shape[1]) if audio.ndim == 2 else 1,
|
|
"frames": int(audio.shape[0]) if audio.ndim >= 1 else 0,
|
|
"samples": int(flat.size),
|
|
"duration_sec": float((audio.shape[0] if audio.ndim >= 1 else 0) / sample_rate),
|
|
"sum": float(np.sum(flat, dtype=np.float64)),
|
|
"mean_abs": float(np.mean(np.abs(flat), dtype=np.float64)),
|
|
"rms": float(np.sqrt(np.mean(np.square(flat), dtype=np.float64))),
|
|
"min": float(np.min(flat)),
|
|
"max": float(np.max(flat)),
|
|
}
|
|
|
|
|
|
def run_case(pipe: Any, case: dict[str, Any], iteration: int, output_root: Path) -> dict[str, Any]:
|
|
request = request_from_case(case)
|
|
case_dir = output_root / f"{request['id']}_iter{iteration}"
|
|
case_dir.mkdir(parents=True, exist_ok=True)
|
|
config = make_generation_config(case)
|
|
pipe.generation_config = config
|
|
|
|
start = time.perf_counter()
|
|
plan_start = time.perf_counter()
|
|
plan = pipe.plan(**request)
|
|
sync_device(str(pipe.device))
|
|
plan_wall = time.perf_counter() - plan_start
|
|
|
|
semantic_start = time.perf_counter()
|
|
semantic = pipe.generate_semantic(plan)
|
|
sync_device(str(pipe.device))
|
|
semantic_wall = time.perf_counter() - semantic_start
|
|
|
|
nar_start = time.perf_counter()
|
|
latents = pipe.synthesize(semantic)
|
|
sync_device(str(pipe.device))
|
|
nar_wall = time.perf_counter() - nar_start
|
|
|
|
decode_start = time.perf_counter()
|
|
audio = pipe.decode(latents)
|
|
sync_device(str(pipe.device))
|
|
decode_wall = time.perf_counter() - decode_start
|
|
wall = time.perf_counter() - start
|
|
|
|
audio_path = case_dir / "audio.wav"
|
|
sf.write(audio_path, audio, 48000, subtype="FLOAT")
|
|
plan.save(case_dir / "plan")
|
|
np.save(case_dir / "semantic.npy", np.asarray(semantic.tokens, dtype=np.int32))
|
|
np.save(case_dir / "latent.npy", latents.astype(np.float32))
|
|
|
|
result = {
|
|
"case": request["id"],
|
|
"iteration": iteration,
|
|
"request": request,
|
|
"truncated": {"abc": bool(plan.truncated), "semantic": bool(semantic.truncated)},
|
|
"token_counts": {
|
|
"abc": len(plan.abc_ids),
|
|
"prefix": len(plan.prefix),
|
|
"semantic": len(semantic.tokens),
|
|
"latent_frames": int(latents.shape[0]) if getattr(latents, "ndim", 0) == 2 else None,
|
|
},
|
|
"timing_sec": {
|
|
"plan_wall": plan_wall,
|
|
"semantic_wall": semantic_wall,
|
|
"nar_wall": nar_wall,
|
|
"decode_wall": decode_wall,
|
|
"wall": wall,
|
|
},
|
|
"audio": summarize_audio(audio, 48000),
|
|
"paths": {
|
|
"case_dir": str(case_dir),
|
|
"audio": str(audio_path),
|
|
},
|
|
}
|
|
(case_dir / "result.json").write_text(json.dumps(result, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
|
|
return result
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
add_reference_path(args.reference_root)
|
|
device = configure_runtime(args)
|
|
from yue2 import YuE2Pipeline
|
|
|
|
output_root = resolve_path(args.output_dir)
|
|
output_root.mkdir(parents=True, exist_ok=True)
|
|
args.timing_file = resolve_path(args.timing_file)
|
|
args.summary_file = resolve_path(args.summary_file)
|
|
args.timing_file.parent.mkdir(parents=True, exist_ok=True)
|
|
args.summary_file.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
cases = selected_cases(args, load_cases(args.cases))
|
|
results: list[dict[str, Any]] = []
|
|
model = str(resolve_path(args.model)) if Path(args.model).exists() else args.model
|
|
vae = str(resolve_path(args.vae)) if Path(args.vae).exists() else args.vae
|
|
|
|
with YuE2Pipeline.from_pretrained(
|
|
model,
|
|
vae=vae,
|
|
device=device,
|
|
backend="torch",
|
|
local_files_only=args.local_files_only,
|
|
progress=False,
|
|
generation_config=make_generation_config(cases[0]),
|
|
) as pipe:
|
|
for i in range(args.warmup):
|
|
run_case(pipe, cases[0], -(i + 1), output_root / "warmup")
|
|
for iteration in range(args.iterations):
|
|
for case in cases:
|
|
result = run_case(pipe, case, iteration, output_root)
|
|
results.append(result)
|
|
with args.timing_file.open("a", encoding="utf-8") as log:
|
|
log.write(json.dumps(result, ensure_ascii=False, sort_keys=True) + "\n")
|
|
|
|
summary = {
|
|
"family": args.family,
|
|
"model": model,
|
|
"vae": vae,
|
|
"backend": args.backend,
|
|
"device": device,
|
|
"threads": args.threads,
|
|
"cases": [r["case"] for r in results],
|
|
"results": results,
|
|
}
|
|
args.summary_file.write_text(json.dumps(summary, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|