File size: 2,538 Bytes
6b1f7f2 dee8df0 6b1f7f2 dee8df0 6b1f7f2 dee8df0 f647f52 dee8df0 6b1f7f2 f647f52 6b1f7f2 dee8df0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | """Three YuE2 CUDA variants from one model load and two SheetSage ABC forms."""
from __future__ import annotations
import json
from pathlib import Path
import sys
import traceback
import numpy as np
import soundfile as sf
import torch
from yue2 import YuE2Pipeline
YUE_REV = "29b3558dd46954a0cd9021dc76d5c91864a0f1c7"
VAE_REV = "9a94e1d0ea9f8087e98f77fa88df4a4068104d2a"
def run(request_path: Path, output: Path) -> None:
if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
raise RuntimeError("YuE2 requires a BF16-capable NVIDIA GPU")
requests = json.loads(request_path.read_text())["runs"]
pipe = YuE2Pipeline.from_pretrained(
"m-a-p/YuE2-3B", vae="m-a-p/YuE2-Vae",
revision=YUE_REV, vae_revision=VAE_REV,
device="cuda", backend="torch", quantization="none",
offload_ar=False, memory_budget_gib=44, progress=True,
)
try:
for request in requests:
target = output / request["name"]
target.mkdir()
try:
score = Path(request["abc_path"]).read_text()
song = pipe(style=request["style"], lyrics="", cot=request["cot"],
abc=score, seed=request["seed"], cfg_scale=1.0,
semantic_sampling={"max_tokens": 1100})
audio = np.asarray(song.audio, dtype=np.float32)
if audio.ndim != 2 or audio.shape[1] != 2 or not np.isfinite(audio).all():
raise RuntimeError("YuE2 returned invalid stereo audio")
if not len(audio) or float(np.max(np.abs(audio))) == 0:
raise RuntimeError("YuE2 returned silence")
sf.write(target / "raw.wav", audio, song.sample_rate, subtype="FLOAT")
(target / "yue.json").write_text(json.dumps({
"model": "m-a-p/YuE2-3B", "revision": YUE_REV,
"backend": "torch", "seed": request["seed"],
"vae_revision": VAE_REV, "duration_s": len(audio) / song.sample_rate,
"truncated": song.truncated, "timing": song.timing,
"source_timeline": False, "vocal_absence_guaranteed": False,
}, indent=2) + "\n")
del song, audio
torch.cuda.empty_cache()
except Exception:
(target / "error.txt").write_text(traceback.format_exc())
finally:
pipe.close()
if __name__ == "__main__":
run(Path(sys.argv[1]), Path(sys.argv[2]))
|