Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
278 lines
12 KiB
Python
278 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""Build the fixed-length RKNN Flow DiT estimator for the CosyVoice3 family.
|
|
|
|
The runtime side (``src/models/cosyvoice3/rknn_flow.cpp``, session option
|
|
``cosyvoice3.rknn_flow``) runs one compiled RKNN estimator step in place of the ggml DiT graph,
|
|
inside the same flow-matching ODE loop. Nothing in this repository produces that graph, so this
|
|
script is its only producer: it compiles the checkpoint's own
|
|
``flow.decoder.estimator.fp32.onnx`` (from ``FunAudioLLM/Fun-CosyVoice3-0.5B-2512``, i.e. the same
|
|
weights the GGUF package was converted from) for one fixed sequence length and for a RKNPU2 target.
|
|
|
|
The estimator is the flow decoder's ODE step: batch 2 is the classifier-free-guidance pair
|
|
``[cond, uncond]`` (the uncond branch carries zeroed ``mu``/``cond``/``spks``), and one call returns
|
|
``dphi_dt`` for both branches. Six inputs ``x [2,80,S]`` / ``mask [2,1,S]`` / ``mu [2,80,S]`` /
|
|
``t [2]`` (the raw timestep; the time-embedding MLP is inside the graph) / ``spks [2,80]`` /
|
|
``cond [2,80,S]``, one output ``estimator_out [2,80,S]``, fp32 in and out. The runtime's ggml graph
|
|
has no ``mask`` and takes the timestep embedding instead of ``t``; the sidecar supplies both.
|
|
|
|
RKNN has no dynamic shapes, so the runtime compiles a ladder of lengths and pads a shorter
|
|
utterance up to the smallest engine that fits: ``x``/``mu``/``cond`` zero-padded on the right,
|
|
``mask`` 1 on the real frames and 0 in the pad, output cropped back. That convention is what the
|
|
padding check measures on the host: a padded run at ``S`` against an exact run at ``frames``, per
|
|
length, before any of it is paid for on a board (``--skip-check`` skips it). ``--verify-sim`` goes
|
|
one step further and runs the fp16 graph on the PC simulator against the fp32 ONNX, which is the
|
|
engine's own accuracy, measured without a board.
|
|
|
|
Requirements (host, x86_64, dedicated venv -- the RKNN toolkit pins its own numpy/onnx):
|
|
pip install onnx onnxruntime numpy
|
|
pip install rknn-toolkit2 # PC-side converter, x86_64 only
|
|
On Python 3.12 the 2.3.2 converter additionally needs ``setuptools<81`` and an ``onnx/mapping.py``
|
|
shim inside the venv's onnx package; see docs/build/rockchip-npu.md.
|
|
|
|
Usage:
|
|
python3 tools/rknn/prepare_cosyvoice3_rknn_flow.py --estimator-onnx flow.decoder.estimator.fp32.onnx \
|
|
--seq-len 200 --seq-len 500 --seq-len 1000 --output-dir models/CosyVoice3-RKNN
|
|
python3 tools/rknn/prepare_cosyvoice3_rknn_flow.py --estimator-onnx <onnx> --seq-len 200 --skip-rknn
|
|
python3 tools/rknn/prepare_cosyvoice3_rknn_flow.py --estimator-onnx <onnx> --seq-len 200 \
|
|
--check-frames 106 --verify-sim
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import onnx
|
|
|
|
MEL_CHANNELS = 80
|
|
CFG_BATCH = 2
|
|
# Graph order, as exported by the checkpoint. `mask` is the frame mask, `t` the raw timestep.
|
|
INPUT_NAMES = ["x", "mask", "mu", "t", "spks", "cond"]
|
|
|
|
|
|
def input_shapes(seq_len: int) -> list[list[int]]:
|
|
return [
|
|
[CFG_BATCH, MEL_CHANNELS, seq_len],
|
|
[CFG_BATCH, 1, seq_len],
|
|
[CFG_BATCH, MEL_CHANNELS, seq_len],
|
|
[CFG_BATCH],
|
|
[CFG_BATCH, MEL_CHANNELS],
|
|
[CFG_BATCH, MEL_CHANNELS, seq_len],
|
|
]
|
|
|
|
|
|
def check_graph_inputs(onnx_path: Path) -> list[str]:
|
|
"""The export's feeder set, validated against the ladder this converter compiles.
|
|
|
|
Not the shared reader in `rknn_graph.py`: a reader reports names, this one refuses an export
|
|
whose names are not the expected ladder.
|
|
"""
|
|
model = onnx.load(str(onnx_path), load_external_data=False)
|
|
names = [i.name for i in model.graph.input]
|
|
# Initializers can appear as graph inputs in older exports; only the real feeders count.
|
|
feed = [n for n in names if n not in {i.name for i in model.graph.initializer}]
|
|
if feed != INPUT_NAMES:
|
|
raise SystemExit(
|
|
f"{onnx_path.name}: expected inputs {INPUT_NAMES} in graph order, got {feed}; "
|
|
"this is not the checkpoint's flow.decoder.estimator export"
|
|
)
|
|
return feed
|
|
|
|
|
|
def pad_to(inputs: list[np.ndarray], seq_len: int) -> list[np.ndarray]:
|
|
"""Right-pad a length-carrying input set to `seq_len`: zeros everywhere, mask 0 in the pad."""
|
|
out = []
|
|
for name, arr in zip(INPUT_NAMES, inputs):
|
|
if name in ("t", "spks"):
|
|
out.append(arr)
|
|
continue
|
|
axis = arr.ndim - 1
|
|
pad = seq_len - arr.shape[axis]
|
|
if pad < 0:
|
|
raise SystemExit(f"{name} is longer ({arr.shape[axis]}) than the engine ({seq_len})")
|
|
# Zero pad: for x/mu/cond that is the reference port's `zero` mode, and for mask it is the
|
|
# 0 that switches the pad frames off.
|
|
out.append(arr if pad == 0 else np.pad(arr, [(0, 0), (0, 0), (0, pad)]))
|
|
return out
|
|
|
|
|
|
def random_inputs(seq_len: int, rng: np.random.Generator) -> list[np.ndarray]:
|
|
inputs = []
|
|
for name, shape in zip(INPUT_NAMES, input_shapes(seq_len)):
|
|
if name == "t":
|
|
# The ODE schedule evaluates the estimator at the start of each interval, t in [0, 1).
|
|
inputs.append(rng.random(shape).astype(np.float32))
|
|
elif name == "mask":
|
|
inputs.append(np.ones(shape, dtype=np.float32))
|
|
else:
|
|
inputs.append(rng.standard_normal(shape).astype(np.float32))
|
|
return inputs
|
|
|
|
|
|
def run(session, inputs: list[np.ndarray]) -> np.ndarray:
|
|
feeds = {name: arr for name, arr in zip(INPUT_NAMES, inputs)}
|
|
return session.run(None, feeds)[0]
|
|
|
|
|
|
def snr_db(reference: np.ndarray, other: np.ndarray) -> float:
|
|
power = float((reference * reference).sum())
|
|
err = float(((reference - other) ** 2).sum())
|
|
return 10.0 * np.log10(max(power, 1e-30) / max(err, 1e-30))
|
|
|
|
|
|
def check_padding(onnx_path: Path, seq_len: int, frames: list[int], seed: int) -> None:
|
|
"""Padded run at `seq_len` vs exact run at `frames`: the padding convention's own gate."""
|
|
import onnxruntime as ort
|
|
|
|
opts = ort.SessionOptions()
|
|
opts.log_severity_level = 3
|
|
session = ort.InferenceSession(str(onnx_path), opts, providers=["CPUExecutionProvider"])
|
|
rng = np.random.default_rng(seed)
|
|
print(f"check {seq_len}: padded run vs exact run (onnxruntime CPU, fp32)")
|
|
for length in frames:
|
|
if length > seq_len:
|
|
continue
|
|
inputs = random_inputs(length, rng)
|
|
exact = run(session, inputs)
|
|
padded = run(session, pad_to(inputs, seq_len))[:, :, :length]
|
|
# The frames nearest the pad are the ones the pad can contaminate: the estimator's causal
|
|
# conv-pos-embed pads on the left only, so a right pad reaches the tail through attention.
|
|
tail = padded[:, :, -min(60, length):]
|
|
tail_ref = exact[:, :, -min(60, length):]
|
|
print(
|
|
f" frames={length:4d} pad={seq_len - length:4d} "
|
|
f"snr={snr_db(exact, padded):6.1f} dB tail_snr={snr_db(tail_ref, tail):6.1f} dB "
|
|
f"max_abs={np.abs(exact - padded).max():.4f}"
|
|
)
|
|
|
|
|
|
def verify_fp16(
|
|
onnx_path: Path,
|
|
seq_len: int,
|
|
frames: int,
|
|
target: str,
|
|
opt_level: int,
|
|
seed: int,
|
|
) -> None:
|
|
"""The fp16 cost of the compiled graph, measured on the PC simulator.
|
|
|
|
Built from the ONNX rather than run from the exported `.rknn`: a `load_rknn` engine cannot be
|
|
inferred on the simulator at all (`init_runtime: RKNN model that loaded by 'load_rknn' not
|
|
support inference on the simulator, please set 'target' first!`), and `load_onnx` + `build`
|
|
with the same config produces the same fp16 graph the board runs.
|
|
|
|
Worth measuring before a board is asked to judge it: the whole utterance's mel comes out of ten
|
|
of these steps, so a per-step bias is what the reference port compensates with a hand-tuned mel
|
|
offset.
|
|
"""
|
|
import onnxruntime as ort
|
|
from rknn.api import RKNN
|
|
|
|
opts = ort.SessionOptions()
|
|
opts.log_severity_level = 3
|
|
session = ort.InferenceSession(str(onnx_path), opts, providers=["CPUExecutionProvider"])
|
|
rng = np.random.default_rng(seed)
|
|
inputs = pad_to(random_inputs(frames, rng), seq_len)
|
|
reference = run(session, inputs)
|
|
|
|
rknn = RKNN(verbose=False)
|
|
rknn.config(target_platform=target, optimization_level=opt_level)
|
|
if rknn.load_onnx(model=str(onnx_path), inputs=INPUT_NAMES, input_size_list=input_shapes(seq_len)) != 0:
|
|
raise SystemExit("load_onnx failed")
|
|
if rknn.build(do_quantization=False) != 0:
|
|
raise SystemExit("build failed")
|
|
if rknn.init_runtime() != 0:
|
|
raise SystemExit("init_runtime (PC simulator) failed")
|
|
got = rknn.inference(inputs=inputs)[0]
|
|
rknn.release()
|
|
got = np.asarray(got).reshape(reference.shape)
|
|
# The frames nearest the pad are where an fp16 engine's error concentrates, same as the
|
|
# padding gate above, and the ODE loop feeds its own output back in ten times.
|
|
tail = min(60, frames)
|
|
print(
|
|
f"verify fp16 (seq {seq_len}) vs fp32 ONNX (PC simulator): frames={frames} pad={seq_len - frames} "
|
|
f"snr={snr_db(reference, got):6.1f} dB "
|
|
f"tail_snr={snr_db(reference[:, :, -tail:], got[:, :, -tail:]):6.1f} dB "
|
|
f"max_abs={np.abs(reference - got).max():.4f} "
|
|
f"mean={reference.mean():+.5f}/{got.mean():+.5f}"
|
|
)
|
|
|
|
|
|
def convert_rknn(
|
|
onnx_path: Path,
|
|
output: Path,
|
|
target: str,
|
|
opt_level: int,
|
|
seq_len: int,
|
|
) -> None:
|
|
names = check_graph_inputs(onnx_path)
|
|
# FP16 with no quantization: the flow estimator is the stage the reference port keeps on the NPU,
|
|
# and quantizing it (int8) is what pushes the mel off enough to need a hand-tuned offset.
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
from rknn_build import build_engine, report_engine
|
|
|
|
build_engine(
|
|
onnx_path,
|
|
output,
|
|
target,
|
|
names=names,
|
|
shapes=input_shapes(seq_len),
|
|
opt_level=opt_level,
|
|
gate_label=f"build {output.name}",
|
|
)
|
|
report_engine(
|
|
output, f"seq_len {seq_len}", f"{target}", f"optimization_level {opt_level}")
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
|
|
parser.add_argument(
|
|
"--estimator-onnx", type=Path, required=True,
|
|
help="flow.decoder.estimator.fp32.onnx from the Fun-CosyVoice3 checkpoint")
|
|
parser.add_argument(
|
|
"--seq-len", type=int, action="append", default=None,
|
|
help="sequence length to compile (repeatable); the runtime pads up to the smallest fit")
|
|
parser.add_argument("--target", default="rk3588")
|
|
parser.add_argument("--opt-level", type=int, default=3)
|
|
parser.add_argument("--output-dir", type=Path, default=Path("models/CosyVoice3-RKNN"))
|
|
parser.add_argument(
|
|
"--check-frames", type=int, action="append", default=None,
|
|
help="utterance lengths to check the padding at, per engine (default: a short and a "
|
|
"near-full length)")
|
|
parser.add_argument("--seed", type=int, default=1986, help="matches the runtime's flow seed")
|
|
parser.add_argument("--skip-check", action="store_true")
|
|
parser.add_argument("--skip-rknn", action="store_true", help="check only")
|
|
parser.add_argument(
|
|
"--verify-sim", action="store_true",
|
|
help="build the first --seq-len and run it on the PC simulator against the fp32 ONNX at "
|
|
"the first --check-frames; slow, the simulator runs the whole estimator")
|
|
args = parser.parse_args()
|
|
|
|
seq_lens = args.seq_len or [200, 500, 1000]
|
|
if args.verify_sim:
|
|
seq_len = seq_lens[0]
|
|
frames = (args.check_frames or [seq_len // 2])[0]
|
|
verify_fp16(args.estimator_onnx, seq_len, frames, args.target, args.opt_level, args.seed)
|
|
return 0
|
|
if not args.skip_check:
|
|
for seq_len in seq_lens:
|
|
frames = args.check_frames or [max(1, seq_len // 2), max(1, seq_len - 1)]
|
|
check_padding(args.estimator_onnx, seq_len, frames, args.seed)
|
|
if args.skip_rknn:
|
|
return 0
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
for seq_len in seq_lens:
|
|
convert_rknn(
|
|
args.estimator_onnx,
|
|
args.output_dir / f"flow_estimator_seq{seq_len}.rknn",
|
|
args.target,
|
|
args.opt_level,
|
|
seq_len,
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|