| |
| """Generate Russian music with the measured YuE2 packed inference profiles.""" |
| import argparse |
| import dataclasses |
| import json |
| import math |
| import hashlib |
| import os |
| import sys |
| import time |
|
|
| PROCESS_T0 = time.time() |
| from pathlib import Path |
|
|
| SRC = Path(__file__).absolute().parent |
| ROOT = SRC.parent |
| sys.path.insert(0, str(SRC / "runtime/kernels/base")) |
| sys.path.insert(0, str(SRC / "runtime")) |
| PROFILES = ("reference", "fused", "fast", "memory") |
|
|
|
|
| def load_pipeline(*, vae=None, low_memory=False, offload_ar=False, progress=True, |
| profile=None, ode_steps=None, vae_precision=None, fuse_rms_quant=False, |
| capture_sampling=None, kv_cache_dtype="bf16", int8_attention=None): |
| import torch |
| profile = profile or ("reference" if low_memory else "fast") |
| if profile not in PROFILES: |
| raise ValueError(f"Unknown profile {profile!r}") |
| if low_memory and profile != "reference": |
| raise ValueError("Unfused --low-memory requires profile=reference; use profile=memory for staged experts") |
| if fuse_rms_quant and profile not in {"fast", "memory"}: |
| raise ValueError("RMS/activation fusion requires profile=fast or profile=memory") |
| if kv_cache_dtype not in {"bf16", "int8"}: |
| raise ValueError("kv_cache_dtype must be bf16 or int8") |
| if kv_cache_dtype == "int8" and (profile not in {"fast", "memory"} or not fuse_rms_quant): |
| raise ValueError("INT8 KV requires profile=fast or memory and fuse_rms_quant=True") |
| if int8_attention is not None and int8_attention not in {"reference", "ordered", "subwarp"}: |
| raise ValueError("int8_attention must be reference, ordered, subwarp, or None") |
| if int8_attention is not None and kv_cache_dtype != "int8": |
| raise ValueError("int8_attention requires kv_cache_dtype=int8") |
| int8_attention = (int8_attention or "ordered") if kv_cache_dtype == "int8" else None |
| if capture_sampling is not None and type(capture_sampling) is not bool: |
| raise ValueError("capture_sampling must be a boolean or None") |
| if capture_sampling and profile not in {"fast", "memory"}: |
| raise ValueError("Captured sampling requires profile=fast or profile=memory") |
| if not torch.cuda.is_available(): |
| raise RuntimeError("This release requires an NVIDIA CUDA GPU.") |
| if torch.cuda.get_device_capability() not in {(8, 9), (12, 0)}: |
| raise RuntimeError("This runtime ships kernels for SM89 (Ada/RTX 40) and SM120 (Blackwell consumer); other GPUs need a rebuild.") |
| if ode_steps is not None and (isinstance(ode_steps, bool) or not isinstance(ode_steps, int) or ode_steps < 1): |
| raise ValueError("ode_steps must be a positive integer") |
| if vae_precision is not None and vae_precision not in {"fp32", "fp16", "bf16"}: |
| raise ValueError("vae_precision must be fp32, fp16, or bf16") |
| if vae is None: |
| from huggingface_hub import snapshot_download |
| lock = json.loads((ROOT / "source-lock.json").read_text()) |
| vae = snapshot_download("m-a-p/YuE2-Vae", revision=lock["m-a-p/YuE2-Vae"], |
| allow_patterns=["*.json", "*.safetensors", "*.py", "LICENSE", "licenses/*", "THIRD_PARTY_NOTICES.md"]) |
| from yue2_orbit import OrbitPipeline |
| pipe = OrbitPipeline.from_pretrained(str(ROOT), vae=vae, device="cuda", |
| memory_budget_gib=18, vae_core_frames=512, |
| offload_ar=offload_ar or profile == "memory", progress=progress) |
| pipe.fused = not low_memory |
| pipe.trim_cache = True |
| pipe.profile = profile |
| pipe.fuse_rms_quant = bool(fuse_rms_quant) |
| pipe.kv_cache_dtype = kv_cache_dtype |
| pipe.int8_attention = int8_attention |
| fast = profile in {"fast", "memory"} |
| pipe.capture_sampling = fast if capture_sampling is None else capture_sampling |
| pipe.generation_config = dataclasses.replace(pipe.generation_config, |
| ode_steps=ode_steps if ode_steps is not None else 8 if fast else 32) |
| precision = vae_precision or ("fp16" if fast else "fp32") |
| pipe.vae_compute_dtype = {"fp32": None, "fp16": torch.float16, "bf16": torch.bfloat16}[precision] |
| |
| files = sorted(p for p in SRC.rglob("*") if p.is_file() and |
| "__pycache__" not in p.parts and "kernel-source" not in p.parts and |
| (p.suffix in {".py", ".so", ".whl"} or p.name == "requirements.txt")) |
| hashes = {str(p.relative_to(SRC)): hashlib.sha256(p.read_bytes()).hexdigest() for p in files} |
| pipe.runtime_bundle_sha256 = hashlib.sha256(json.dumps(hashes, sort_keys=True).encode()).hexdigest() |
| return pipe |
|
|
|
|
| def _guidance_scale(value): |
| value = float(value) |
| if not math.isfinite(value) or not 0 <= value <= 20: |
| raise argparse.ArgumentTypeError("cfg scale must be finite and in [0,20]") |
| return value |
|
|
|
|
| class _NvmlPeak: |
| """Whole-process NVML peak of this GPU's used memory (20 ms sampling), for the report.""" |
|
|
| def __init__(self): |
| import threading |
| self.peak = 0; self._stop = threading.Event(); self._thread = None |
| try: |
| import pynvml |
| pynvml.nvmlInit(); self._h = pynvml.nvmlDeviceGetHandleByIndex(0); self._nv = pynvml |
| except Exception: |
| self._nv = None |
| return |
| self._thread = threading.Thread(target=self._run, daemon=True); self._thread.start() |
|
|
| def _run(self): |
| while not self._stop.wait(0.02): |
| try: |
| self.peak = max(self.peak, self._nv.nvmlDeviceGetMemoryInfo(self._h).used) |
| except Exception: |
| return |
|
|
| def stop(self): |
| self._stop.set() |
| if self._thread is not None: |
| self._thread.join(timeout=1) |
| return self.peak / 2**30 if self._nv is not None else None |
|
|
|
|
| def main(): |
| nvml_peak = _NvmlPeak() |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--prompt", default=str(ROOT / "examples/prompts/default-ru.json")) |
| parser.add_argument("--output", default="song") |
| parser.add_argument("--vae") |
| parser.add_argument("--profile", choices=PROFILES) |
| parser.add_argument("--ode-steps", type=int) |
| parser.add_argument("--cfg-scale", type=_guidance_scale, |
| help="Semantic guidance scale; 1 uses one decode branch") |
| parser.add_argument("--vae-precision", choices=("fp32", "fp16", "bf16")) |
| parser.add_argument("--allocator", choices=("native", "expandable")) |
| parser.add_argument("--low-memory", action="store_true") |
| parser.add_argument("--offload-ar", action="store_true") |
| parser.add_argument("--fuse-rms-quant", action="store_true", |
| help="Fuse RMSNorm into activation quantization; faster experimental arithmetic") |
| parser.add_argument("--kv-cache-dtype", choices=["bf16", "int8"], default="bf16", |
| help="AR KV cache precision; int8 requires --fuse-rms-quant and fast/memory") |
| parser.add_argument("--int8-attention", choices=["reference", "ordered", "subwarp"], |
| help="INT8 decode: ordered preserves previous arithmetic (default); subwarp changes rounding; reference uses the previous kernel") |
| parser.add_argument("--no-captured-sampling", action="store_true", |
| help="Use the earlier eager sampler for comparison") |
| parser.add_argument("--warm-compile", action="store_true", |
| help="keep the torch.compile paths (VAE decoder, NAR pointwise chains); pays off only when one process serves several songs") |
| parser.add_argument("--mode", choices=["off", "exact", "fast", "lowmem", "turbo"], default="exact", |
| help="off: reference decode path; exact: same sampled tokens, faster pipeline; fast: + INT8 output heads and NAR KV handoff; lowmem: + INT8 embeddings; turbo: W4A4 output heads") |
| args = parser.parse_args() |
| allocator = args.allocator or ("expandable" if (args.profile or "fast") in {"fast", "memory"} and not args.low_memory else None) |
| if allocator is not None: |
| os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True" if allocator == "expandable" else "backend:native" |
| request = json.loads(Path(args.prompt).read_text()) |
| out = Path(args.output) |
| out.mkdir(parents=True, exist_ok=True) |
| with load_pipeline(vae=args.vae, low_memory=args.low_memory, offload_ar=args.offload_ar, |
| profile=args.profile, ode_steps=args.ode_steps, vae_precision=args.vae_precision, |
| fuse_rms_quant=args.fuse_rms_quant, kv_cache_dtype=args.kv_cache_dtype, |
| int8_attention=args.int8_attention, |
| capture_sampling=False if args.no_captured_sampling else None) as pipe: |
| mode_flags = [] |
| if args.mode != "off": |
| from fast_ops import apply_mode |
| model = pipe._load_model() |
| mode_flags = apply_mode(pipe, model, args.mode, warm=args.warm_compile) |
| gen_start = time.perf_counter() |
| generation = {k: request[k] for k in ("style", "lyrics", "seed", "cot") if k in request} |
| if args.cfg_scale is not None: |
| generation["cfg_scale"] = args.cfg_scale |
| song = pipe(**generation) |
| t_save = time.perf_counter() |
| song.save_artifacts(out) |
| t_artifacts = time.perf_counter() - t_save |
| song.save(out / "audio.wav") |
| t_wav = time.perf_counter() - t_save - t_artifacts |
| stage_peaks = [v.get("peak_gib", 0) for v in song.timing.get("stages", {}).get("memory", {}).values()] |
| report = dict(seconds_including_load=time.time() - PROCESS_T0, generation_seconds=time.perf_counter() - gen_start, |
| mode=args.mode, mode_flags=mode_flags, |
| nvml_process_peak_gib=nvml_peak.stop(), torch_process_peak_gib=max(stage_peaks) if stage_peaks else None, |
| save_artifacts_seconds=round(t_artifacts, 3), save_wav_seconds=round(t_wav, 3), |
| timing=song.timing, |
| audio_seconds=len(song.audio) / song.sample_rate, |
| packed_runtime=pipe.runtime_report(), truncated=song.truncated, |
| allocator_config=os.environ.get("PYTORCH_ALLOC_CONF", os.environ.get("PYTORCH_CUDA_ALLOC_CONF"))) |
| (out / "release-run.json").write_text(json.dumps(report, indent=2)) |
| print(json.dumps(report), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
