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