File size: 5,245 Bytes
4f5f309 | 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 | """Export the Mimi decoder (codes -> 24 kHz wav) to ONNX, fp32, dynamic frame count, and check parity against
transformers' MimiModel.decode on random codes. Run guarded:
.venv/bin/python tools/memguard.py --rss-limit-mb 2000 -- \
.venv-tts/bin/python -m tools.maya_export.export_codec
"""
from __future__ import annotations
import time
from pathlib import Path
import numpy as np
import onnxruntime as ort
import torch
from transformers import MimiModel
from tools.maya_export.common import ckpt_paths, work_dir
N_Q = 16
N_CODES = 2048
SAMPLES_PER_FRAME = 1920 # 24 kHz / 12.5 Hz
PARITY_TOL = 1e-4
class _Decode(torch.nn.Module):
def __init__(self, model: MimiModel) -> None:
super().__init__()
self.model = model
def forward(self, codes: torch.Tensor) -> torch.Tensor:
return self.model.decode(codes)[0] # (1, 1, F * 1920)
def _mimi_dir() -> Path:
from maya3 import hub
return Path(hub.resolve("maya-research/maya2-native")["mimi_dir"])
def _traceable_sliding_mask(config, inputs_embeds, **_):
"""Explicit causal sliding-window additive mask. transformers' mask helper fails under torch.onnx tracing;
check_parity compares the exported graph with the unpatched model, so a wrong mask shows up there."""
t = inputs_embeds.shape[1]
i = torch.arange(t)[:, None]
j = torch.arange(t)[None, :]
allowed = (j <= i) & (j > i - int(config.sliding_window))
return torch.where(allowed, 0.0, float("-inf"))[None, None].to(inputs_embeds.dtype)
def check_mask(model: MimiModel, lengths=(5, 300)) -> None:
"""The traced graph uses _traceable_sliding_mask. Compare it with transformers' own mask on lengths beyond
the sliding window (250 frames), where the window actually cuts attention."""
import transformers.models.mimi.modeling_mimi as M
for t in lengths:
emb = torch.zeros(1, t, model.config.hidden_size)
pos = torch.arange(t)[None]
ref = M.create_sliding_window_causal_mask(config=model.config, inputs_embeds=emb, attention_mask=None,
past_key_values=None, position_ids=pos)
ref = ref.reshape(t, t)
ref_allowed = ref if ref.dtype == torch.bool else ref == 0
mine = _traceable_sliding_mask(model.config, emb).reshape(t, t) == 0
if not torch.equal(ref_allowed, mine):
raise SystemExit(f"sliding-window mask differs from transformers at T={t}")
print(f"sliding-window mask matches transformers at T={list(lengths)}")
def export(out_path: Path) -> MimiModel:
import transformers.models.mimi.modeling_mimi as M
model = MimiModel.from_pretrained(str(_mimi_dir()), dtype=torch.float32, attn_implementation="eager").eval()
wrapper = _Decode(model).eval()
dummy = torch.randint(0, N_CODES, (1, N_Q, 5), dtype=torch.long)
saved = M.create_sliding_window_causal_mask
M.create_sliding_window_causal_mask = _traceable_sliding_mask
try:
_trace(wrapper, dummy, out_path)
finally:
M.create_sliding_window_causal_mask = saved
return model
def _trace(wrapper: _Decode, dummy: torch.Tensor, out_path: Path) -> None:
with torch.inference_mode():
torch.onnx.export(
wrapper,
(dummy,),
str(out_path),
input_names=["codes"],
output_names=["wav"],
dynamic_axes={"codes": {2: "F"}, "wav": {2: "samples"}},
opset_version=17,
dynamo=False,
)
def check_parity(out_path: Path, model: MimiModel, frames=(5, 17, 50)) -> float:
sess = ort.InferenceSession(str(out_path), providers=["CPUExecutionProvider"])
worst = 0.0
rng = np.random.default_rng(0)
for F in frames:
codes = rng.integers(0, N_CODES, size=(1, N_Q, F), dtype=np.int64)
with torch.inference_mode():
ref = model.decode(torch.from_numpy(codes))[0].float().numpy()
got = sess.run(None, {"codes": codes})[0]
assert got.shape == ref.shape == (1, 1, F * SAMPLES_PER_FRAME), (got.shape, ref.shape)
worst = max(worst, float(np.max(np.abs(got - ref))))
return worst
def speed(out_path: Path, F: int = 50, reps: int = 20) -> float:
sess = ort.InferenceSession(str(out_path), providers=["CPUExecutionProvider"])
codes = np.random.default_rng(1).integers(0, N_CODES, size=(1, N_Q, F), dtype=np.int64)
sess.run(None, {"codes": codes})
times = []
for _ in range(reps):
t0 = time.perf_counter()
sess.run(None, {"codes": codes})
times.append(time.perf_counter() - t0)
return float(np.median(times)) / F * 1000.0 # ms per frame
def main() -> None:
out = work_dir() / "assets" / "v1" / "mimi_decoder.onnx"
out.parent.mkdir(parents=True, exist_ok=True)
model = export(out)
print(f"exported {out} ({out.stat().st_size / 2**20:.1f} MiB)")
check_mask(model)
err = check_parity(out, model)
print(f"parity max abs error vs MimiModel.decode: {err:.3e} (tol {PARITY_TOL:.0e})")
ms = speed(out)
print(f"decode median: {ms:.2f} ms per frame at F=50 (gate <= 8 ms)")
if err > PARITY_TOL:
raise SystemExit("parity gate FAILED")
if __name__ == "__main__":
main()
|