"""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]))