File size: 6,136 Bytes
ee4cb7d | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | #!/usr/bin/env python3
"""Run cached autoregressive MuScriptor inference with ONNX Runtime."""
from __future__ import annotations
import argparse
import json
import time
from pathlib import Path
import numpy as np
# Importing PyTorch first preloads the CUDA/cuDNN shared libraries shipped in
# the uv environment when available; CPU-only users do not need PyTorch.
try:
import torch # noqa: F401
except ModuleNotFoundError:
torch = None # type: ignore[assignment]
import onnxruntime as ort
from muscriptor_onnx.audio import SAMPLE_RATE, load_audio_16k, log_mel_spectrogram
def providers(requested: str, device_id: int) -> list:
available = ort.get_available_providers()
if requested == "cuda" or (requested == "auto" and "CUDAExecutionProvider" in available):
if "CUDAExecutionProvider" not in available:
raise RuntimeError(f"CUDA EP unavailable; installed providers: {available}")
return [("CUDAExecutionProvider", {"device_id": device_id}), "CPUExecutionProvider"]
return ["CPUExecutionProvider"]
def session(path: Path, selected_providers: list) -> ort.InferenceSession:
options = ort.SessionOptions()
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
result = ort.InferenceSession(path, sess_options=options, providers=selected_providers)
requested_cuda = bool(selected_providers) and (
selected_providers[0] == "CUDAExecutionProvider"
or (
isinstance(selected_providers[0], tuple)
and selected_providers[0][0] == "CUDAExecutionProvider"
)
)
if requested_cuda and "CUDAExecutionProvider" not in result.get_providers():
raise RuntimeError("CUDA EP was requested but session creation fell back to CPU")
return result
def causal_mask(query_length: int, past_length: int) -> np.ndarray:
query_positions = past_length + np.arange(query_length)[:, None]
key_positions = np.arange(past_length + query_length)[None, :]
return np.where(key_positions <= query_positions, 0.0, -65504.0).astype(np.float16)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--model-dir", type=Path, required=True)
parser.add_argument("--audio", type=Path)
parser.add_argument("--provider", choices=("auto", "cuda", "cpu"), default="auto")
parser.add_argument("--device-id", type=int, default=0)
parser.add_argument("--max-new-tokens", type=int, default=8)
parser.add_argument(
"--instrument-id",
type=int,
action="append",
help="Optional MT3_FULL_PLUS group ID; repeat for multiple groups",
)
args = parser.parse_args()
metadata = json.loads((args.model_dir / "config.json").read_text())
selected = providers(args.provider, args.device_id)
conditioner = session(args.model_dir / "conditioner.onnx", selected)
decoder = session(args.model_dir / "decoder.onnx", selected)
if args.audio:
audio = load_audio_16k(args.audio)
source = str(args.audio)
else:
# Deterministic smoke-test input; it only tests execution, not quality.
time_axis = np.arange(5 * SAMPLE_RATE, dtype=np.float32) / SAMPLE_RATE
audio = (0.1 * np.sin(2 * np.pi * 440.0 * time_axis)).astype(np.float32)
source = "generated 440 Hz sine"
mel = log_mel_spectrogram(audio).astype(np.float16)
instrument_ids = np.asarray(
[args.instrument_id if args.instrument_id is not None else [-1]], dtype=np.int64
)
dataset_ids = np.asarray([[-1]], dtype=np.int64)
start = time.perf_counter()
condition = conditioner.run(
["condition_embeddings"],
{
"log_mel": mel,
"instrument_ids": instrument_ids,
"dataset_ids": dataset_ids,
},
)[0]
condition_seconds = time.perf_counter() - start
layers = metadata["num_layers"]
heads = metadata["num_heads"]
head_dim = metadata["head_dim"]
past_key = np.zeros((layers, 1, heads, 0, head_dim), dtype=np.float16)
past_value = np.zeros_like(past_key)
input_ids = np.asarray([[metadata["initial_token_id"]]], dtype=np.int64)
generated: list[int] = []
decode_times: list[float] = []
for step in range(args.max_new_tokens):
prefix = condition if step == 0 else np.empty((1, 0, heads * head_dim), np.float16)
query_length = prefix.shape[1] + input_ids.shape[1]
mask = causal_mask(query_length, past_key.shape[3])
tick = time.perf_counter()
logits, past_key, past_value = decoder.run(
("logits", "present_key", "present_value"),
{
"input_ids": input_ids,
"condition_embeddings": prefix,
"past_key": past_key,
"past_value": past_value,
"attention_mask": mask,
},
)
decode_times.append(time.perf_counter() - tick)
if not np.isfinite(logits).all():
raise RuntimeError("decoder produced non-finite logits")
logits[:, metadata["first_reserved_token_id"] :] = -np.inf
token = int(logits.argmax(axis=-1)[0])
generated.append(token)
input_ids = np.asarray([[token]], dtype=np.int64)
if token == metadata["eos_token_id"]:
break
active = decoder.get_providers()
sizes = {
path.name: path.stat().st_size
for path in args.model_dir.iterdir()
if path.is_file() and (path.suffix == ".onnx" or path.name.endswith(".onnx.data"))
}
sizes["total"] = sum(sizes.values())
print(f"model: {args.model_dir}")
print(f"audio: {source}")
print(f"providers: {active}")
print(f"mel shape: {mel.shape}")
print(f"condition shape: {condition.shape}")
print(f"final KV shape: {past_key.shape}")
print(f"generated tokens: {generated}")
print(f"condition latency: {condition_seconds * 1000:.1f} ms")
print(f"decode latency: {[round(x * 1000, 1) for x in decode_times]} ms")
print(f"ONNX file sizes: {sizes}")
if __name__ == "__main__":
main()
|