src/runtime/yue2/cli.py

"""Small CLI sharing the Python pipeline's defaults and artifact protocol."""
from __future__ import annotations
import argparse
import importlib.metadata
import json
import os
from pathlib import Path
import sys
import time
from .storage import write_json, verify_result, identity


def kit_root():
    return Path(os.environ.get("YUE2_KIT", Path(__file__).resolve().parents[2]))


def model_paths(args):
    root = kit_root()
    local_model = root / "models/YuE2-3B"
    local_vae = root / "models" / ("YuE2-Vae-legacy" if args.vae == "legacy" else "YuE2-Vae")
    model = args.model or (str(local_model) if local_model.exists() else "m-a-p/YuE2-3B")
    vae = args.vae if args.vae not in {"standard", "legacy"} else (
        str(local_vae) if local_vae.exists() else "m-a-p/" + local_vae.name)
    return model, vae


def get_pipe(args):
    from .pipeline import YuE2Pipeline
    from .protocol import GenerationConfig
    model, vae = model_paths(args)
    config = GenerationConfig.from_dict(json.loads(Path(args.config).read_text())) if args.config else None
    return YuE2Pipeline.from_pretrained(model, vae=vae, revision=args.revision,
             vae_revision=args.vae_revision, device=args.device, memory_budget_gib=args.budget,
             backend=args.backend, quantization=args.quantization, offload_ar=args.offload_ar,
             local_files_only=args.offline, generation_config=config,
             vae_core_frames=512 if args.budget <= 12 else 1024,
             progress=not getattr(args, "quiet", False))


def request_kwargs(data, base=Path.cwd()):
    allowed = {"style", "tags", "lyrics", "cot", "seed", "abc", "cfg_scale", "id", "abc_sampling", "semantic_sampling"}
    metadata = {"lang", "eval_index", "clip_id", "prompt"}
    unknown = set(data) - allowed - metadata - {"abc_path"}
    if unknown:
        raise ValueError(f"Unknown request fields: {sorted(unknown)}")
    result = {k: v for k, v in data.items() if k in allowed}
    if "abc_path" in data:
        if data.get("abc") is not None:
            raise ValueError("Pass abc or abc_path, not both")
        result["abc"] = (base / data["abc_path"]).read_bytes().decode("utf-8")
    if data.get("prompt") is not None:
        from .protocol import SongRequest
        request = SongRequest(style=result.get("style", result.get("tags")), lyrics=result["lyrics"],
                              cot=result.get("cot", "full"))
        if data["prompt"] != request.text():
            raise ValueError("Historical literal prompt does not match native instruction/style/lyrics")
    return result


def doctor(args):
    import torch
    from .storage import model_identity
    model, vae = model_paths(args)
    versions = {}
    for package in ("torch", "transformers", "huggingface-hub", "safetensors", "tiktoken", "soundfile"):
        try:
            versions[package] = importlib.metadata.version(package)
        except importlib.metadata.PackageNotFoundError:
            versions[package] = None
    devices = []
    for i in range(torch.cuda.device_count()):
        p = torch.cuda.get_device_properties(i)
        devices.append({"id": i, "name": p.name, "memory_gib": p.total_memory / 2**30,
                        "compute_capability": [p.major, p.minor]})
    report = {"dependencies_ready": all(versions.values()), "versions": versions,
              "cuda": devices, "mps_available": torch.backends.mps.is_available(),
              "model": model, "vae": vae, "default_cot": "full", "default_cfg": {"full": 1., "melody": 1., "off": 1.01},
              "validated": False, "note": "Environment readiness is not quality or real-24GB acceptance."}
    if args.verify_hashes:
        from .storage import resolve_model
        report["weights"] = {"model": model_identity(resolve_model(model, local_files_only=args.offline)),
                             "vae": model_identity(resolve_model(vae, local_files_only=args.offline))}
    if args.output:
        write_json(args.output, report)
    print(json.dumps(report, indent=2, ensure_ascii=False))
    return 0 if report["dependencies_ready"] else 1


def generate(args):
    if args.request:
        path = Path(args.request)
        data = json.loads(path.read_text())
        base = path.parent
    else:
        data, base = {}, Path.cwd()
    for key in ("id", "style", "cot", "seed", "cfg_scale"):
        value = getattr(args, key)
        if value is not None:
            data[key] = value
    if args.lyrics_file:
        data["lyrics"] = Path(args.lyrics_file).read_bytes().decode("utf-8")
    elif args.lyrics is not None:
        data["lyrics"] = args.lyrics
    if args.abc_file:
        data["abc"] = Path(args.abc_file).read_bytes().decode("utf-8")
        data.pop("abc_path", None)
    if not data:
        data = {"id": "first_song", "style": "Mandarin, warm piano, acoustic pop, female vocal",
                "lyrics": "[Verse]\n晚风轻轻吹过窗前\n你留下的笑还在昨天\n[Chorus]\n让这首歌陪你走远\n把所有想念唱成明天"}
    kwargs = request_kwargs(data, base)
    pipe = get_pipe(args)
    directory = Path(args.output or "runs/default") / kwargs.get("id", "song")
    if args.resume and (directory / "result.json").exists():
        req_kwargs = {k: v for k, v in kwargs.items() if k not in {"abc_sampling", "semantic_sampling"}}
        request = pipe._request(**req_kwargs)
        config = pipe.effective_config(request, kwargs.get("abc_sampling"), kwargs.get("semantic_sampling"))
        expected = identity({"request": request.to_dict(), "config": config, "weights": pipe.weights})
        result = verify_result(directory, expected)
        print(json.dumps({"resumed": True, "result": str(directory / "result.json"), "truncated": result["truncated"]}))
        return 0
    if directory.exists() and any(directory.iterdir()):
        raise FileExistsError(f"Nonempty output {directory}; use --resume or a new output directory")
    directory.mkdir(parents=True, exist_ok=True)
    try:
        if args.stage == "plan":
            kwargs.pop("semantic_sampling", None)
            plan = pipe.plan(**kwargs)
            plan.save(directory)
            print(json.dumps({"stage": "plan", "output": str(directory), "truncated": plan.truncated}))
        else:
            result = pipe(**kwargs)
            result.save_artifacts(directory)
            print(json.dumps({"status": "complete", "output": str(directory), "truncated": result.truncated,
                              "seconds": result.timing["e2e_seconds"]}))
        return 0
    except BaseException as exc:
        write_json(directory / "failure.json", {"status": "failed", "type": type(exc).__name__, "reason": str(exc),
                                                "request": data})
        raise
    finally:
        pipe.close()


def batch(args):
    if args.concurrency != 1:
        raise ValueError("The minimal torch pipeline currently supports concurrency=1; do not share a pipeline concurrently")
    path = Path(args.input)
    rows = [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
    ids = [r.get("id") for r in rows]
    if any(x is None for x in ids) or len(ids) != len(set(ids)):
        raise ValueError("Every batch request must have a unique id")
    pipe = get_pipe(args)
    output = Path(args.output or "runs/batch")
    receipts, failures = [], 0
    try:
        for index, row in enumerate(rows, 1):
            if not args.quiet:
                print(f"Song {index}/{len(rows)}: {row['id']}", file=sys.stderr, flush=True)
            kwargs = request_kwargs(row, path.parent)
            if args.cot is not None:
                kwargs["cot"] = args.cot
            directory = output / row["id"]
            try:
                request = pipe._request(**{k:v for k,v in kwargs.items() if k not in {"abc_sampling", "semantic_sampling"}})
                cfg = pipe.effective_config(request, kwargs.get("abc_sampling"), kwargs.get("semantic_sampling"))
                expected = identity({"request": request.to_dict(), "config": cfg, "weights": pipe.weights})
                if args.resume and (directory / "result.json").exists():
                    receipt = verify_result(directory, expected)
                else:
                    if directory.exists() and any(directory.iterdir()):
                        raise FileExistsError("Nonempty request output; use --resume or a new run")
                    receipt = pipe(**kwargs).save_artifacts(directory)
                receipts.append({"id": row["id"], "status": "complete", "identity": receipt["identity"]})
            except Exception as exc:
                failures += 1
                failure = {"id": row["id"], "status": "failed", "reason": str(exc), "type": type(exc).__name__}
                write_json(directory / "failure.json", failure)
                receipts.append(failure)
            write_json(output / "batch.json", {"complete": len(receipts) == len(rows) and not failures,
                       "expected": len(rows), "failed": failures, "results": receipts})
    finally:
        pipe.close()
    return int(failures > 0)


def parser():
    p = argparse.ArgumentParser(description="YuE2: style + lyrics → symbolic plan → song")
    sub = p.add_subparsers(dest="command", required=True)
    for name in ("doctor", "generate", "batch"):
        q = sub.add_parser(name)
        q.add_argument("--model")
        q.add_argument("--vae", default="standard", help="standard (old), YuE2-Vae-legacy, local path or HF repo")
        q.add_argument("--revision")
        q.add_argument("--vae-revision")
        q.add_argument("--device", default="auto")
        q.add_argument("--budget", type=float, default=24)
        q.add_argument("--backend", choices=("torch", "torch-eager", "vllm"), default="torch")
        q.add_argument("--quantization", choices=("none", "fp8"), default="none")
        q.add_argument("--offload-ar", action="store_true")
        q.add_argument("--offline", action="store_true")
        q.add_argument("--config")
        q.add_argument("--output")
        if name == "doctor":
            q.add_argument("--verify-hashes", action="store_true")
        else:
            q.add_argument("--cot", choices=("full", "melody", "off"))
            q.add_argument("--resume", action="store_true")
            q.add_argument("--quiet", "--no-progress", action="store_true",
                           help="Hide YuE2 progress on stderr; keep result output on stdout")
            if name == "batch":
                q.add_argument("--input", required=True)
                q.add_argument("--concurrency", type=int, default=1)
            else:
                q.add_argument("--request")
                q.add_argument("--id")
                q.add_argument("--style")
                q.add_argument("--lyrics")
                q.add_argument("--lyrics-file")
                q.add_argument("--abc-file")
                q.add_argument("--seed", type=int)
                q.add_argument("--cfg-scale", type=float)
                q.add_argument("--stage", choices=("plan", "audio"), default="audio")
    for name in ("verify", "bench", "eval"):
        q = sub.add_parser(name, add_help=False)
        q.add_argument("arguments", nargs=argparse.REMAINDER)
    return p


def main(argv=None):
    argv = sys.argv[1:] if argv is None else argv
    if argv and argv[0] in {"verify", "bench", "eval"}:
        import subprocess
        script = {"verify": "validate_release.py", "bench": "benchmark.py", "eval": "evaluate.py"}[argv[0]]
        target = kit_root() / "scripts" / script
        if not target.is_file():
            raise FileNotFoundError("This command needs the complete minimal kit (scripts/data/evaluation_runtime). "
                                    "Set YUE2_KIT to its directory; the HF wheel supplies generation and the Python API.")
        return subprocess.call([sys.executable, str(target), *argv[1:]])
    args = parser().parse_args(argv)
    return {"doctor": doctor, "generate": generate, "batch": batch}[args.command](args)


if __name__ == "__main__":
    raise SystemExit(main())