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
470 lines
21 KiB
Python
470 lines
21 KiB
Python
#!/usr/bin/env python3
|
|
"""Build the stateful RKNN speech decoder (vocoder) for the Qwen3-TTS family.
|
|
|
|
The runtime side (``src/models/qwen3_tts/rknn_speech_decoder.cpp``, session option
|
|
``qwen3_tts.exec_plan``) runs one compiled RKNN step graph
|
|
in place of the ggml speech tokenizer decoder. Nothing in this repository produces that graph, so
|
|
this script is its only producer: it re-traces the WavTokenizer-v2 decoder from the
|
|
``Qwen/Qwen3-TTS-12Hz-0.6B-Base`` Hugging Face checkpoint (the ``speech_tokenizer`` subdirectory,
|
|
i.e. the same weights the GGUF package was converted from) as a fixed-shape streaming step, checks
|
|
it against a full decode, and compiles it for the NPU.
|
|
|
|
The exported step graph is stateful: every causal conv's tail and the pre-transformer's
|
|
sliding-window KV history travel through the graph I/O as one flat fp32 cache, so a step decodes
|
|
from continuous state and there are no chunk seams. Four inputs
|
|
``audio_codes [1,16,B] int64`` / ``freqs_in [B,head_dim/2] fp32`` / ``cache_valid [1] fp32`` /
|
|
``cache_in [N] fp32``, two outputs ``audio_values`` / ``cache_out``, ``1920`` samples per token at
|
|
24 kHz. ``N`` does not depend on ``B``, so a small-B step's cache feeds a large-B step directly —
|
|
that is what the execution plan's optional ``artifacts.first`` uses.
|
|
|
|
Requirements (host, x86_64, dedicated venv — the RKNN toolkit pins its own numpy/onnx/torch):
|
|
pip install torch 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_qwen3_tts_rknn_vocoder.py --model-dir ~/.cache/huggingface/hub/\\
|
|
models--Qwen--Qwen3-TTS-12Hz-0.6B-Base/snapshots/<sha> --codes-length 64
|
|
python3 tools/rknn/prepare_qwen3_tts_rknn_vocoder.py --model-dir <dir> --codes-length 16
|
|
python3 tools/rknn/prepare_qwen3_tts_rknn_vocoder.py --model-dir <dir> --skip-rknn # ONNX only
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import os
|
|
import sys
|
|
|
|
|
|
def conv_cache_len(kernel_size, dilation):
|
|
return (kernel_size - 1) * dilation
|
|
|
|
|
|
def convT_cache_len(kernel_size, stride):
|
|
return math.ceil((kernel_size - 1) / stride)
|
|
|
|
|
|
def build_stateful_decoder(dec, B):
|
|
"""Wrap ``dec`` (the checkpoint's decoder, eval mode, eager attention) as a step graph:
|
|
B tokens + flat cache in, B*samples_per_token samples + flat cache out.
|
|
|
|
``forward(codes [1,16,B] int64, freqs [B,head_dim/2] fp32, cache_valid [1] fp32,
|
|
cache [N] fp32)``. The host-supplied choices below are NPU-driven and load-bearing; each is
|
|
measured, and re-simplifying one costs 30-62 dB of output quality.
|
|
"""
|
|
import torch
|
|
|
|
class Step(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.dec = dec
|
|
self.B = B
|
|
cfg = dec.config
|
|
self.Lc = cfg.sliding_window - 1 # KV history slots
|
|
self.n_layers = cfg.num_hidden_layers
|
|
self.n_heads = cfg.num_attention_heads
|
|
self.head_dim = cfg.head_dim
|
|
self.scaling = self.head_dim ** -0.5
|
|
|
|
specs = [("pre_conv", conv_cache_len(3, 1) * cfg.codebook_dim)]
|
|
specs.append(
|
|
("kv", self.n_layers * 2 * self.n_heads * self.Lc * self.head_dim)
|
|
)
|
|
for bi in range(len(cfg.upsampling_ratios)):
|
|
specs.append((f"up_convT{bi}", convT_cache_len(2, 2) * cfg.latent_dim))
|
|
specs.append((f"up_dw{bi}", conv_cache_len(7, 1) * cfg.latent_dim))
|
|
specs.append(("dec0", conv_cache_len(7, 1) * cfg.latent_dim))
|
|
for i, r in enumerate(cfg.upsample_rates):
|
|
in_dim = cfg.decoder_dim // 2 ** i
|
|
out_dim = cfg.decoder_dim // 2 ** (i + 1)
|
|
specs.append((f"b{i}_convT", convT_cache_len(2 * r, r) * in_dim))
|
|
for d in (1, 3, 9):
|
|
specs.append((f"b{i}_r{d}", conv_cache_len(7, d) * out_dim))
|
|
out_dim = cfg.decoder_dim // 2 ** len(cfg.upsample_rates)
|
|
specs.append(("final", conv_cache_len(7, 1) * out_dim))
|
|
self.specs = specs
|
|
self.cache_numel = sum(n for _, n in specs)
|
|
|
|
# KV history joins the step's k/v through constant-selection matmuls, not Concat:
|
|
# the RKNN Concat path corrupts sliding-window attention (measured: step-0 audio SNR
|
|
# 9.5 dB with kv zeroed, vs 56.8 dB for the concat-free graph); matmul and add are
|
|
# first-class NPU ops.
|
|
a = torch.zeros(self.Lc + B, self.Lc)
|
|
a[: self.Lc, :] = torch.eye(self.Lc)
|
|
p = torch.zeros(self.Lc + B, B)
|
|
p[self.Lc :, :] = torch.eye(B)
|
|
self.register_buffer("kv_sel_hist", a)
|
|
self.register_buffer("kv_sel_cur", p)
|
|
# Sliding-window causal mask on the RELATIVE grid (depends only on qp-kp), as a graph
|
|
# constant at -100 (exp(-100) is exactly 0 in fp16; -3e4 leaks through the NPU
|
|
# softmax). The training mask's "no history before token 0" guard is not expressible
|
|
# on a relative grid, and feeding an absolute mask as an input together with in-graph
|
|
# Cos/Sin segfaults the board runtime, so the cache slots carry a per-slot bias built
|
|
# in forward() from the host-supplied count of real slots.
|
|
qp = torch.arange(B, dtype=torch.float32).view(B, 1)
|
|
kp = (torch.arange(self.Lc + B, dtype=torch.float32) - self.Lc).view(1, self.Lc + B)
|
|
m = torch.where(
|
|
(kp <= qp) & (kp > qp - self.Lc - 1),
|
|
torch.zeros(()),
|
|
torch.full((), -100.0),
|
|
)
|
|
self.register_buffer("attn_mask", m.view(1, 1, B, self.Lc + B))
|
|
self.register_buffer(
|
|
"slot_idx",
|
|
torch.arange(self.Lc, dtype=torch.float32).view(1, 1, 1, self.Lc),
|
|
)
|
|
|
|
def _slot_bias(self, n_valid):
|
|
# -100 on the Lc - n_valid never-written slots, 0 on the real ones. Elementwise on a
|
|
# [1,1,1,Lc] grid: no comparison/Where, and no reduction over the 1.3 MB cache tensor
|
|
# (whose Abs/ReduceSum/Clip version costs 30 dB of NPU fp16 accuracy).
|
|
phantom = torch.clamp(self.Lc - n_valid - self.slot_idx, 0.0, 1.0)
|
|
bias = -100.0 * phantom
|
|
cur = torch.zeros(1, 1, 1, self.B)
|
|
return torch.cat([bias, cur], dim=-1)
|
|
|
|
def _split(self, flat):
|
|
out, off = {}, 0
|
|
for name, n in self.specs:
|
|
out[name] = flat[off : off + n]
|
|
off += n
|
|
return out
|
|
|
|
def _conv_step(self, cnet, x, cache):
|
|
cl = x.shape[1]
|
|
n = conv_cache_len(cnet.conv.kernel_size[0], cnet.conv.dilation[0])
|
|
cs = cache.reshape(1, cl, n) if n else cache.reshape(1, cl, 0)
|
|
xin = torch.cat([cs, x], dim=-1)
|
|
y = cnet.conv(xin)
|
|
return y, xin[..., -n:].reshape(-1) if n else x[..., :0].reshape(-1)
|
|
|
|
def _convT_step(self, tnet, x, cache):
|
|
cl, L = x.shape[1], x.shape[2]
|
|
n = convT_cache_len(tnet.conv.kernel_size[0], tnet.conv.stride[0])
|
|
cs = cache.reshape(1, cl, n) if n else cache.reshape(1, cl, 0)
|
|
xin = torch.cat([cs, x], dim=-1)
|
|
s = tnet.conv.stride[0]
|
|
y = tnet.conv(xin)
|
|
total = xin.shape[-1]
|
|
y = y[..., : total * s]
|
|
y = y[..., (total - L) * s :]
|
|
return y, x[..., -n:].reshape(-1) if n else x[..., :0].reshape(-1)
|
|
|
|
@staticmethod
|
|
def _rot_half(t):
|
|
d = t.shape[-1]
|
|
a, b = t[..., : d // 2], t[..., d // 2 :]
|
|
return torch.cat([-b, a], dim=-1)
|
|
|
|
def forward(self, codes, freqs, cache_valid, flat_cache):
|
|
# freqs = (pos * inv_freq) reduced mod 2*pi on the host, cos/sin computed in-graph:
|
|
# cos/sin as RKNN inputs explodes the fp16 pipeline, and raw angles near pos 1e3 land
|
|
# on an fp16 ulp of 0.5.
|
|
# cache_valid = min(pos, Lc), the COUNT of real cache slots. The cache is a rolling
|
|
# window, so the real slots are always the last ones; never-written slots hold
|
|
# all-zero keys, whose score is 0 (not -inf), so they steal softmax mass. A boolean
|
|
# "any history" flag is only true at step 0, so every later step leaks Lc - B phantom
|
|
# keys into the softmax: measured 10.2 dB at B=16 (40 of 87 keys phantom) against
|
|
# 48.8 dB at B=64 (7 of 79).
|
|
dec = self.dec
|
|
B = self.B
|
|
c = self._split(flat_cache)
|
|
parts = []
|
|
|
|
hidden = dec.quantizer.decode(codes)
|
|
hidden, cache = self._conv_step(dec.pre_conv, hidden, c["pre_conv"])
|
|
parts.append(cache)
|
|
hidden = hidden.transpose(1, 2)
|
|
|
|
t = dec.pre_transformer
|
|
hidden = t.input_proj(hidden)
|
|
emb = torch.cat([freqs, freqs], dim=-1)
|
|
cos = emb.cos()
|
|
sin = emb.sin()
|
|
|
|
kv = c["kv"].reshape(self.n_layers, 2, self.n_heads, self.Lc, self.head_dim)
|
|
for i, layer in enumerate(t.layers):
|
|
res = hidden
|
|
hs = layer.input_layernorm(hidden)[0]
|
|
a = layer.self_attn
|
|
q = a.q_proj(hs).view(B, self.n_heads, self.head_dim).transpose(0, 1)
|
|
k = a.k_proj(hs).view(B, self.n_heads, self.head_dim).transpose(0, 1)
|
|
v = a.v_proj(hs).view(B, self.n_heads, self.head_dim).transpose(0, 1)
|
|
q = q * cos + self._rot_half(q) * sin
|
|
k = k * cos + self._rot_half(k) * sin
|
|
k_full = torch.matmul(self.kv_sel_hist, kv[i, 0]) + torch.matmul(
|
|
self.kv_sel_cur, k
|
|
)
|
|
v_full = torch.matmul(self.kv_sel_hist, kv[i, 1]) + torch.matmul(
|
|
self.kv_sel_cur, v
|
|
)
|
|
parts.append(k_full[:, -self.Lc :].reshape(-1))
|
|
parts.append(v_full[:, -self.Lc :].reshape(-1))
|
|
att = torch.softmax(
|
|
q[None] @ k_full[None].transpose(2, 3) * self.scaling
|
|
+ self.attn_mask
|
|
+ self._slot_bias(cache_valid),
|
|
dim=-1,
|
|
)
|
|
o = (att @ v_full[None]).transpose(1, 2).reshape(1, B, -1)
|
|
hidden = res + layer.self_attn_layer_scale(a.o_proj(o))
|
|
hidden = hidden + layer.mlp_layer_scale(
|
|
layer.mlp(layer.post_attention_layernorm(hidden))
|
|
)
|
|
hidden = t.norm(hidden)
|
|
hidden = t.output_proj(hidden).transpose(1, 2)
|
|
|
|
for bi, blocks in enumerate(dec.upsample):
|
|
hidden, cache = self._convT_step(blocks[0], hidden, c[f"up_convT{bi}"])
|
|
parts.append(cache)
|
|
nb = blocks[1]
|
|
res = hidden
|
|
hidden, cache = self._conv_step(nb.dwconv, hidden, c[f"up_dw{bi}"])
|
|
parts.append(cache)
|
|
hs = hidden.permute(0, 2, 1)
|
|
hs = nb.pwconv2(nb.act(nb.pwconv1(nb.norm(hs))))
|
|
hidden = res + (nb.gamma * hs).permute(0, 2, 1)
|
|
|
|
hidden, cache = self._conv_step(dec.decoder[0], hidden, c["dec0"])
|
|
parts.append(cache)
|
|
for bi in range(len(dec.config.upsample_rates)):
|
|
blk = dec.decoder[1 + bi]
|
|
hidden = blk.block[0](hidden)
|
|
hidden, cache = self._convT_step(blk.block[1], hidden, c[f"b{bi}_convT"])
|
|
parts.append(cache)
|
|
for unit, d in zip(blk.block[2:], (1, 3, 9)):
|
|
res = hidden
|
|
hidden = unit.act1(hidden)
|
|
hidden, cache = self._conv_step(unit.conv1, hidden, c[f"b{bi}_r{d}"])
|
|
parts.append(cache)
|
|
hidden = unit.act2(hidden)
|
|
hidden = unit.conv2(hidden) + res
|
|
hidden = dec.decoder[-2](hidden)
|
|
hidden, cache = self._conv_step(dec.decoder[-1], hidden, c["final"])
|
|
parts.append(cache)
|
|
|
|
wav = hidden.clamp(min=-1, max=1).reshape(-1)
|
|
return wav, torch.cat([p.reshape(-1) for p in parts])
|
|
|
|
return Step().eval()
|
|
|
|
|
|
def export_onnx(args):
|
|
from rknn_build import report_engine
|
|
import numpy as np
|
|
import torch
|
|
|
|
tokenizer_dir = os.path.join(args.model_dir, "speech_tokenizer")
|
|
if not os.path.isdir(tokenizer_dir):
|
|
sys.exit(f"ERROR: {tokenizer_dir} is not a Qwen3-TTS speech_tokenizer directory")
|
|
from qwen_tts.core.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
|
|
Qwen3TTSTokenizerV2Model,
|
|
)
|
|
|
|
full = Qwen3TTSTokenizerV2Model.from_pretrained(tokenizer_dir, local_files_only=True)
|
|
dec = full.decoder.eval()
|
|
dec.config._attn_implementation = "eager"
|
|
|
|
B = args.codes_length
|
|
step = build_stateful_decoder(dec, B)
|
|
print(f"step B={B} cache_numel={step.cache_numel} "
|
|
f"({step.cache_numel * 4 / 1e6:.1f} MB fp32) head_dim={step.head_dim} kv_len={step.Lc}")
|
|
|
|
torch.manual_seed(0)
|
|
n = 3 * B
|
|
codes = torch.randint(0, 2048, (1, 16, n), dtype=torch.int64)
|
|
inv_freq = dec.pre_transformer.rotary_emb.inv_freq.float()
|
|
|
|
def rope(i, L):
|
|
pos = torch.arange(i, i + L, dtype=torch.float32)
|
|
return torch.remainder(torch.outer(pos, inv_freq), 2 * math.pi)
|
|
|
|
with torch.no_grad():
|
|
ref = dec(codes).reshape(-1)
|
|
flat = torch.zeros(step.cache_numel, dtype=torch.float32)
|
|
outs = []
|
|
for i in range(0, n, B):
|
|
L = min(B, n - i)
|
|
wav, flat = step(codes[:, :, i : i + L], rope(i, L), np.float32(min(i, step.Lc)), flat)
|
|
outs.append(wav)
|
|
stepped = torch.cat(outs)
|
|
|
|
L = min(len(ref), len(stepped))
|
|
snr = 10 * torch.log10(
|
|
ref[:L].pow(2).sum() / ((ref[:L] - stepped[:L]).pow(2).sum() + 1e-12)
|
|
)
|
|
print(f"torch stepwise vs full decode: len {len(stepped)}/{len(ref)} SNR {snr:.1f} dB")
|
|
# fp32 self-check: semantic errors land below 20 dB; the kv-selection matmuls and the
|
|
# mask broadcast add fp32 accumulation noise (measured 103-111 dB). Do not raise this to
|
|
# 80: it aborts before torch.onnx.export and leaves no ONNX for the converter to load.
|
|
if snr < 40:
|
|
d = (ref[:L] - stepped[:L]).abs()
|
|
print("max diff by 10k-sample blocks:",
|
|
[round(float(d[i : i + 10000].max()), 5) for i in range(0, L, 10000)][:12])
|
|
sys.exit("stateful decode does not match full decode")
|
|
if args.check_only:
|
|
return None
|
|
|
|
out = args.output or f"vocoder_stateful_{B}.onnx"
|
|
if out.endswith(".rknn"):
|
|
# --output names the compiled engine; the ONNX lands beside it.
|
|
out = os.path.splitext(out)[0] + ".onnx"
|
|
os.makedirs(os.path.dirname(out) or ".", exist_ok=True)
|
|
dummy_c = torch.randint(0, 2048, (1, 16, B), dtype=torch.int64)
|
|
dummy_cache = torch.zeros(step.cache_numel, dtype=torch.float32)
|
|
print(f"exporting {out} ...")
|
|
# cache_valid is traced as a dynamic input; trace it mid-range so no branch of the
|
|
# phantom-slot clamp degenerates into a constant.
|
|
torch.onnx.export(
|
|
step,
|
|
(dummy_c, rope(0, B), torch.full((1,), float(B)), dummy_cache),
|
|
out,
|
|
input_names=["audio_codes", "freqs_in", "cache_valid", "cache_in"],
|
|
output_names=["audio_values", "cache_out"],
|
|
opset_version=args.opset,
|
|
do_constant_folding=True,
|
|
dynamo=False,
|
|
)
|
|
report_engine(out)
|
|
|
|
# Preflight for the converter stage: rknn-toolkit2 2.3.2's fold_constant aborts (core dump,
|
|
# no Python traceback) on the decomposed LayerNorm/Softmax patterns a recent torch emits.
|
|
# Measured with this graph: torch 2.4.0 + onnx 1.22 exports 1064 nodes (Gemm linears, no
|
|
# Pad/Exp/Shape/ConstantOfShape) and converts byte-identically to the reference .rknn;
|
|
# torch 2.14 + onnx 1.23 exports 2653 nodes and aborts. Say it here, at export time,
|
|
# instead of losing a 20-minute converter run to a core dump.
|
|
import onnx
|
|
|
|
nodes = onnx.load(out, load_external_data=False).graph.node
|
|
risky = sorted({n.op_type for n in nodes} & {"Pad", "Exp", "ConstantOfShape", "Shape", "Split"})
|
|
print(f"graph: {len(nodes)} nodes, risky ops for rknn-toolkit2: {risky or 'none'}")
|
|
if risky:
|
|
print("WARNING: the RKNN stage will most likely abort in fold_constant with this graph; "
|
|
"re-export it in the converter venv (toolkit-pinned torch/onnx) and pass that file "
|
|
"with --onnx.")
|
|
|
|
import onnxruntime as ort
|
|
|
|
opts = ort.SessionOptions()
|
|
opts.intra_op_num_threads = 4
|
|
sess = ort.InferenceSession(out, sess_options=opts, providers=["CPUExecutionProvider"])
|
|
flat = np.zeros(step.cache_numel, dtype=np.float32)
|
|
outs = []
|
|
for i in range(0, n, B):
|
|
L = min(B, n - i)
|
|
w, flat = sess.run(None, {
|
|
"audio_codes": codes[:, :, i : i + L].numpy(),
|
|
"freqs_in": rope(i, L).numpy(),
|
|
"cache_valid": np.full(1, min(i, step.Lc), dtype=np.float32),
|
|
"cache_in": flat,
|
|
})
|
|
outs.append(np.asarray(w).reshape(-1)[: L * args.samples_per_token])
|
|
onnx_step = np.concatenate(outs)
|
|
ref_np = ref[:L].numpy()
|
|
snr = 10 * np.log10(
|
|
np.sum(ref_np**2) / (np.sum((ref_np - onnx_step[:L]) ** 2) + 1e-12)
|
|
)
|
|
print(f"onnxruntime stepwise vs full decode: SNR {snr:.1f} dB")
|
|
if snr < 40:
|
|
sys.exit("exported ONNX does not reproduce the torch step graph")
|
|
return out, step.cache_numel, step.head_dim, step.Lc
|
|
|
|
|
|
def graph_contract(onnx_path):
|
|
"""Read (B, cache_numel, head_dim) back out of an exported ONNX, so the RKNN stage can run
|
|
alone in the converter venv without restating the contract."""
|
|
import onnx
|
|
|
|
graph = onnx.load(onnx_path, load_external_data=False).graph
|
|
dims = {
|
|
i.name: [d.dim_value for d in i.type.tensor_type.shape.dim] for i in graph.input
|
|
}
|
|
for name in ("audio_codes", "freqs_in", "cache_in"):
|
|
if name not in dims:
|
|
sys.exit(f"ERROR: {onnx_path} is not a stateful step export (no {name} input)")
|
|
codes, freqs, cache = dims["audio_codes"], dims["freqs_in"], dims["cache_in"]
|
|
if len(codes) != 3 or len(freqs) != 2:
|
|
sys.exit(f"ERROR: {onnx_path} does not have the stateful step shapes")
|
|
return codes[2], cache[0], freqs[1] * 2
|
|
|
|
|
|
def convert_rknn(args, onnx_path, cache_numel, head_dim, kv_len):
|
|
out = args.output or (
|
|
os.path.splitext(os.path.basename(onnx_path))[0] + f"_{args.target}.rknn"
|
|
)
|
|
print(f"=== {onnx_path} -> {out} ({args.target}) ===")
|
|
# FP16 (w16a16), not quantized: the Snake activation makes Q8/INT8 audibly noisy
|
|
# (SNR 9.5-13.4 dB) for 1.75x the speed, while fp16 measures lossless (56.8 dB).
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from rknn_build import build_engine, report_engine
|
|
|
|
build_engine(
|
|
onnx_path,
|
|
out,
|
|
args.target,
|
|
names=["audio_codes", "freqs_in", "cache_valid", "cache_in"],
|
|
shapes=[
|
|
[1, 16, args.codes_length],
|
|
[args.codes_length, head_dim // 2],
|
|
[1],
|
|
[cache_numel],
|
|
],
|
|
opt_level=args.opt_level,
|
|
verbose=True,
|
|
gate_label=f"build {os.path.basename(out)}",
|
|
gate_log=out + ".build.log",
|
|
)
|
|
report_engine(out)
|
|
print(
|
|
f"\nDeploy (per-SoC file, not portable between chips):\n"
|
|
f" scp {out} <board>:/root/tts-rknn/vocoder/\n"
|
|
f"then run with\n"
|
|
" --session-option 'qwen3_tts.exec_plan="
|
|
+ json.dumps({"speech_decoder": {"impl": "rknn", "artifacts": {"main": f"/root/tts-rknn/vocoder/{os.path.basename(out)}"}}})
|
|
+ "'"
|
|
)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
|
|
ap.add_argument(
|
|
"--model-dir",
|
|
default=os.environ.get("QWEN3_TTS_MODEL", ""),
|
|
help="Qwen/Qwen3-TTS-12Hz-0.6B-Base snapshot directory (env QWEN3_TTS_MODEL)",
|
|
)
|
|
ap.add_argument("--codes-length", type=int, default=64,
|
|
help="tokens per compiled step (B); one compiled graph per value")
|
|
ap.add_argument("--output", default=None)
|
|
ap.add_argument("--opset", type=int, default=17)
|
|
ap.add_argument("--target", default="rk3588",
|
|
choices=["rk3588", "rk3576", "rk3562", "rk3566", "rk3568"])
|
|
ap.add_argument("--opt-level", type=int, default=1, help="RKNN graph optimization level")
|
|
ap.add_argument("--samples-per-token", type=int, default=1920,
|
|
help="24 kHz / 12.5 Hz token duration")
|
|
ap.add_argument("--check-only", action="store_true", help="run the self-check, export nothing")
|
|
ap.add_argument("--skip-rknn", action="store_true", help="export ONNX only")
|
|
ap.add_argument("--onnx", default=None,
|
|
help="compile an existing stateful ONNX step (the converter-venv-only path); "
|
|
"B, cache_numel and head_dim are read back out of the graph")
|
|
args = ap.parse_args()
|
|
|
|
if args.onnx:
|
|
B, cache_numel, head_dim = graph_contract(args.onnx)
|
|
args.codes_length = B
|
|
print(f"onnx {args.onnx}: B={B} cache_numel={cache_numel} head_dim={head_dim}")
|
|
convert_rknn(args, args.onnx, cache_numel, head_dim, 0)
|
|
return
|
|
|
|
if not args.model_dir:
|
|
sys.exit("ERROR: --model-dir (or QWEN3_TTS_MODEL) is required")
|
|
exported = export_onnx(args)
|
|
if exported is None:
|
|
return
|
|
if args.skip_rknn:
|
|
return
|
|
convert_rknn(args, *exported)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|