diff --git a/CMakeLists.txt b/CMakeLists.txt index 294158fa0..3ef065496 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1119,6 +1119,7 @@ audiocpp_add_model(piper_tts src/community_models/piper_tts/frontend.cpp src/community_models/piper_tts/vits_runtime.cpp src/community_models/piper_tts/rknn_decoder.cpp + src/community_models/piper_tts/rknn_duration_encoder.cpp src/community_models/piper_tts/session.cpp INCLUDES engine/community_models/piper_tts/session.h diff --git a/docs/build/rockchip-npu.md b/docs/build/rockchip-npu.md index 4c4a78476..e02cfc3f5 100644 --- a/docs/build/rockchip-npu.md +++ b/docs/build/rockchip-npu.md @@ -184,6 +184,41 @@ board — `Resize` does not lower to a target the board runtime accepts, so the `ConvTranspose1d` of all ones (exact nearest, and the same operator the decoder already runs on the NPU). +### Piper duration encoder + +The same script with `--encoder` compiles the other half — the text encoder plus the duration +discriminator — for a fixed phoneme-token width: it slices at `enc_p.proj` (the 2x inter_channels +prior stats) and `dp.proj` (the flow conditioning), the two arrays the runtime's flow graphs and +length regulator consume, right-pads tokens to `--frames` and compiles one engine per width. The +slice is deterministic (all three piper `scales` scalars feed sampling, the duration ceiling and +the discriminator's flow noise, strictly downstream of the seam), and piper's exported masks are +pad-safe: a padded run reproduces the exact-length run on the real frames at 117-130 dB fp32, +checked by the self-check on the graph the engine is compiled from. `scales` is not an input. + +```bash +/bin/python tools/rknn/prepare_piper_tts_rknn_vocoder.py \ + --voice-onnx en_US-lessac-medium.onnx --encoder --frames 128 +``` + +With the frame axis pinned the graph constant-folds (992 to 463 nodes at the BASIC level), and +what remains in every attention layer is a relative-shift chain: `Pad + Reshape + Slice` shifting +the [N, 2N-1] relative-position slabs by one step per query row, unrolled N times. The RKNN +kernels compute exactly those chains as garbage on the board — probes show the flatten-Pad tails +filled with uninitialized memory (NaN, ±65504), one attention layer alone drags the engine to +0.4 dB end to end, and a chain probe segfaults the runtime — while the PC simulator computes the +same graph at 56 dB. The simulator is therefore not a gate for this export; a board probe is. The +converter rewrites the chains into exact constant-index gathers (rel-logit bias via +`GatherElements`, rel-value sum via `Gather + Mul + ReduceSum`, verified 121-131 dB fp32 against +the raw slice), the engine computes 40-48 dB on the NPU, and the fold self-checks run on the +rewritten graph. + +The gathers cost O(N^3) weight bytes: the index initializers are constants, so the toolkit folds +them into per-row Toeplitz expansions — 137 MB at 128 tokens, 402 MB at 192, 2.9 GB at 384, and +512 exhausts 31 GB of converter RAM. Compile narrow bins only. The sidecar walks the sorted ladder +to the smallest engine that fits the token count and falls back to the ggml encoder above the +widest: the wide widths are where the fp16 duration drift hurts most (see the measurements +below), so the fallback is a quality gate, not a convenience. + ### Inflect v2 (Micro/Nano) decoder `tools/rknn/prepare_inflect_v2_rknn_decoder.py` slices the official `decode.onnx` at the same kind @@ -381,7 +416,8 @@ the configurable components: | Component | Option | Implementations | Artifact roles | |---|---|---|---| | Piper `decoder` | `piper_tts.exec_plan` | `ggml`, `rknn` | RKNN: required `main` | -| Piper `duration_encoder` / `duration_flows` | `piper_tts.exec_plan` | `ggml` | None (`backend` may name `cpu`) | +| Piper `duration_encoder` | `piper_tts.exec_plan` | `ggml`, `rknn` | RKNN: required `main` (comma-separated ladder of per-width engines) | +| Piper `duration_flows` | `piper_tts.exec_plan` | `ggml` | None (`backend` may name `cpu`) | | Kokoro `generator` | `kokoro_tts.exec_plan` | `ggml`, `rknn` | RKNN: required `main` | | Qwen3 `speech_decoder` | `qwen3_tts.exec_plan` | `ggml`, `rknn` | RKNN: required `main`, optional `first` | | sanoTTS `decoder` | `sanotts.exec_plan` | `default`, `ggml`, `neon` | None | @@ -458,6 +494,14 @@ The same option pins the duration stages to the host, with the decoder untouched --session-option 'piper_tts.exec_plan={"duration_encoder":{"backend":"cpu"},"duration_flows":{"backend":"cpu"}}' ``` +and it moves the duration encoder to the NPU, over a comma-separated ladder of per-width engines +(the same producer as the decoder, `--encoder`; short chunks pick the narrow engine, chunks wider +than the ladder stay on the ggml encoder): + +```bash + --session-option 'piper_tts.exec_plan={"decoder":{"impl":"rknn","artifacts":{"main":"/root/tts-rknn/piper-lessac-w93.rknn"}},"duration_encoder":{"impl":"rknn","artifacts":{"main":"/root/tts-rknn/enc/encoder_n128.rknn,/root/tts-rknn/enc/encoder_n192.rknn"}}}' +``` + The engine is per voice: a `.rknn` compiled from lessac's latent is validated against the loaded voice's channel count at init and fails with a message naming the producer if it does not match. @@ -549,6 +593,19 @@ drift (best lag 0 samples in every window), and the max sample-to-sample step is (13384 vs 13412), so no seam clicks. The dips are near-silence (LSB noise) and one loud plosive onset, i.e. the fp16 noise floor, not the stitching. +The duration encoder ladder is a small, width-limited win, and the width limit is the quality +story. On a 557-character mixed RU/EN text (8 chunks, 109-459 tokens each), the narrow bins are +faster than the ggml encoder (63-71 ms vs 84-89 ms at 109-115 tokens, 1.25-1.35x) and the wide +bin is slower (688-773 ms vs 440-758 ms at 285-373 tokens: 2.9 GB of Toeplitz weights saturate +DRAM bandwidth), so a `{128,192}` ladder with a ggml fallback above 192 is the shape worth +shipping: it moves 3 of the 8 chunks and costs nothing on the rest. Per-bin engines score 40-48 dB +against fp32 on the board (stats and conditioning, n in {17, N-1, N}). Quality: the fp16 engine +shifts each chunk's duration sum by the sum of its per-token sampling errors, so the drift grows +with width — ±7% latent frames on the 128-bin chunks (the scale of the flow's own seeded sampling +noise) against a one-sided −21…−23% on the 384-bin chunks, a 25% tempo shift, audible. GigaAM/ +Whisper transcribe the narrow-bin render and the ggml render to the same text. The wide bins are +fp16-exact per frame and systematically bad per utterance: keep the ladder narrow. + ### Kokoro, measured on the same board `kokoro-82m-q8_0.gguf`, 185 Russian characters, seed 1234, `--threads 4` (14.85 s of audio, diff --git a/docs/models/piper_tts.md b/docs/models/piper_tts.md index 76e89c123..01380f6d4 100644 --- a/docs/models/piper_tts.md +++ b/docs/models/piper_tts.md @@ -39,13 +39,17 @@ audiocpp_cli --task tts --family piper_tts \ |---|---| | `piper_tts.espeak_library_path` | Explicit path to the eSpeak-ng shared library. | | `piper_tts.espeak_data_path` | Explicit path to `espeak-ng-data`. | -| `piper_tts.exec_plan` | JSON object of component execution plans. Default `{}` keeps existing execution. Select the RKNN decoder with `{"decoder":{"impl":"rknn","artifacts":{"main":"/models/decoder.rknn"}}}`, and pin the duration stages to the host with `{"duration_encoder":{"backend":"cpu"},"duration_flows":{"backend":"cpu"}}`. No params. RKNN requires `ENGINE_ENABLE_RKNN=ON`; see [Rockchip NPU](../build/rockchip-npu.md#component-execution-plans). | +| `piper_tts.exec_plan` | JSON object of component execution plans. Default `{}` keeps existing execution. Select the RKNN decoder with `{"decoder":{"impl":"rknn","artifacts":{"main":"/models/decoder.rknn"}}}`, move the duration encoder onto the NPU over a comma-separated ladder of per-width engines with `{"duration_encoder":{"impl":"rknn","artifacts":{"main":"/models/enc_n128.rknn,/models/enc_n192.rknn"}}}` (narrow bins only: wide widths drift the duration sum; chunks above the widest engine stay on ggml), and pin the duration stages to the host with `{"duration_encoder":{"backend":"cpu"},"duration_flows":{"backend":"cpu"}}`. No params. RKNN requires `ENGINE_ENABLE_RKNN=ON`; see [Rockchip NPU](../build/rockchip-npu.md#component-execution-plans). | On RK3588 the decoder is 75% of the runtime, so the NPU path is worth the extra file: measured RTF 0.198 → 0.0696 and peak RSS 299 → 175 MB on `en_US-lessac-medium`. The two duration stages are components of the same option, so each can run on the host while the -decoder stays on the device. `duration_encoder` already runs on the host on CUDA by default: host +decoder stays on the device. `duration_encoder` can also run on the NPU +(`impl: rknn`, a comma-separated ladder of per-width engines from the decoder's producer script +with `--encoder`): the narrow bins measure 1.25-1.35x over the host encoder, wide bins are slower +and shift the duration sum, so keep the ladder narrow and let the fallback cover the rest. +`duration_encoder` already runs on the host on CUDA by default: host duration alignment keeps small TF32 differences from being amplified by `expand_prior` while the expensive decoder stays on the device. `duration_flows` runs on the device by default, and it is the stage worth watching on a device backend, because the flows are the sampler: a rounding diff --git a/include/engine/community_models/piper_tts/rknn_duration_encoder.h b/include/engine/community_models/piper_tts/rknn_duration_encoder.h new file mode 100644 index 000000000..978bf5f05 --- /dev/null +++ b/include/engine/community_models/piper_tts/rknn_duration_encoder.h @@ -0,0 +1,49 @@ +#pragma once + +#include +#include +#include +#include + +namespace engine::models::piper_tts { + +// Rockchip NPU (RKNN) backend for the VITS duration encoder (text tokens to prior +// stats and flow conditioning). +// +// The graph is the voice's own encoder + duration discriminator sliced at their +// projections (see tools/prepare_piper_tts_rknn_vocoder.py --encoder): tokens and +// their length in, the prior stats (2*inter_channels rows) and the flow conditioning +// (inter_channels rows) out. The export carries the frame mask inside the encoder, +// so a right-padded fixed-width run reproduces the exact-length run on the real +// frames (measured 117-130 dB fp32); each compiled engine covers one token width, +// so the sidecar keeps a ladder of bins and an utterance longer than the widest bin +// runs the ggml encoder for that step. +class PiperTtsRknnDurationEncoder { +public: + // `model_paths` is a comma-separated list of compiled engines (the exec_plan + // artifact list); each is validated against the voice at load. + PiperTtsRknnDurationEncoder( + const std::string & model_paths, + int64_t inter_channels); + ~PiperTtsRknnDurationEncoder(); + + PiperTtsRknnDurationEncoder(const PiperTtsRknnDurationEncoder &) = delete; + PiperTtsRknnDurationEncoder & operator=(const PiperTtsRknnDurationEncoder &) = delete; + + // Widest token count any loaded engine covers. + int64_t max_tokens() const; + + // `token_ids` (length n <= max_tokens()) in; stats [2*inter_channels][n] and + // conditioning [inter_channels][n], channel-major, in the layout the flow + // graphs and prior expansion consume. + void encode( + const std::vector & token_ids, + std::vector & stats, + std::vector & conditioning); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::piper_tts diff --git a/include/engine/community_models/piper_tts/vits_runtime.h b/include/engine/community_models/piper_tts/vits_runtime.h index cf076fe01..81ef281ea 100644 --- a/include/engine/community_models/piper_tts/vits_runtime.h +++ b/include/engine/community_models/piper_tts/vits_runtime.h @@ -21,6 +21,9 @@ class PiperVitsRuntime { public: // `rknn_decoder_path` (optional) routes the decoder stage to a pre-compiled RKNN // engine on a Rockchip NPU; empty keeps the ggml decoder. + // `rknn_encoder_path` (optional) routes the duration encoder to a comma-separated + // ladder of pre-compiled RKNN engines; utterances longer than the widest bin fall + // back to the ggml encoder for that step. // `duration_encoder_backend` / `duration_flows_backend` (optional) pin the two duration // stages to a backend of their own, independently of each other and of the decoder. PiperVitsRuntime( @@ -28,7 +31,8 @@ class PiperVitsRuntime { core::BackendConfig backend_config, std::string rknn_decoder_path = {}, std::optional duration_encoder_backend = {}, - std::optional duration_flows_backend = {}); + std::optional duration_flows_backend = {}, + std::string rknn_encoder_path = {}); ~PiperVitsRuntime(); runtime::AudioBuffer synthesize( diff --git a/src/community_models/piper_tts/rknn_duration_encoder.cpp b/src/community_models/piper_tts/rknn_duration_encoder.cpp new file mode 100644 index 000000000..a793f4f86 --- /dev/null +++ b/src/community_models/piper_tts/rknn_duration_encoder.cpp @@ -0,0 +1,205 @@ +#include "engine/community_models/piper_tts/rknn_duration_encoder.h" + +#include "engine/framework/debug/profiler.h" +#include "engine/framework/rknn/engine.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(ENGINE_HAVE_RKNN) +#include +#endif + +namespace engine::models::piper_tts { + +namespace { + +using Clock = std::chrono::steady_clock; + +#if defined(ENGINE_HAVE_RKNN) +// Prefix of every message this sidecar raises, and the label the shared engine +// plumbing stamps on the ones it raises on its behalf. +constexpr const char * kLabel = "Piper TTS duration encoder RKNN"; +#endif + +} // namespace + +struct PiperTtsRknnDurationEncoder::Impl { +#if defined(ENGINE_HAVE_RKNN) + struct Bin { + std::unique_ptr engine; // Engine is not movable + int64_t tokens = 0; // compiled token width + }; + + int64_t channels = 0; // inter_channels; stats carry 2*channels rows + std::vector bins; // ascending by tokens + std::vector fed; // [tokens] padded int64 tokens + std::vector length; // [1] + std::vector stats_out; + std::vector conditioning_out; + + void load(const std::string & model_paths, int64_t inter_channels) { + channels = inter_channels; + std::istringstream stream(model_paths); + std::string path; + while (std::getline(stream, path, ',')) { + const auto at = path.find_first_not_of(" \t"); + if (at == std::string::npos) { + continue; + } + path = path.substr(at, path.find_last_not_of(" \t") - at + 1); + auto bin = Bin{std::make_unique(), 0}; + bin.engine->load(path, kLabel); + const engine::rknn::Engine & engine = *bin.engine; + if (engine.input_count() != 2 || engine.output_count() != 2) { + throw std::runtime_error( + "Piper TTS duration encoder RKNN: model " + path + + " has " + std::to_string(engine.input_count()) + " inputs / " + + std::to_string(engine.output_count()) + + " outputs, expected input + input_lengths and stats + conditioning; " + "compile it with tools/prepare_piper_tts_rknn_vocoder.py --encoder"); + } + const rknn_tensor_attr & input = engine.input_attr(0); + const rknn_tensor_attr & length_attr = engine.input_attr(1); + if (std::string(input.name) != "input" || + std::string(length_attr.name) != "input_lengths") { + throw std::runtime_error( + "Piper TTS duration encoder RKNN: model " + path + " inputs are '" + + std::string(input.name) + "'/'" + std::string(length_attr.name) + + "', expected 'input'/'input_lengths'; compile it with " + "tools/prepare_piper_tts_rknn_vocoder.py --encoder"); + } + bin.tokens = static_cast(input.n_elems); + const rknn_tensor_attr & stats = engine.output_attr(0); + const rknn_tensor_attr & conditioning = engine.output_attr(1); + if (stats.n_elems != + static_cast(2 * channels * bin.tokens) || + conditioning.n_elems != static_cast(channels * bin.tokens)) { + throw std::runtime_error( + "Piper TTS duration encoder RKNN: model " + path + " emits " + + std::to_string(stats.n_elems) + " + " + + std::to_string(conditioning.n_elems) + " values for a " + + std::to_string(bin.tokens) + "-token width, this voice has " + + std::to_string(2 * channels) + " + " + std::to_string(channels) + + " rows; compile the engine from this voice " + "(tools/prepare_piper_tts_rknn_vocoder.py --encoder)"); + } + bins.push_back(std::move(bin)); + } + if (bins.empty()) { + throw std::runtime_error( + "Piper TTS duration encoder RKNN: no engine paths given"); + } + std::sort(bins.begin(), bins.end(), [](const Bin & a, const Bin & b) { + return a.tokens < b.tokens; + }); + for (size_t index = 1; index < bins.size(); ++index) { + if (bins[index].tokens == bins[index - 1].tokens) { + throw std::runtime_error( + "Piper TTS duration encoder RKNN: two engines cover the same " + + std::to_string(bins[index].tokens) + "-token width"); + } + } + length.assign(1, 0); + } +#endif +}; + +PiperTtsRknnDurationEncoder::PiperTtsRknnDurationEncoder( + const std::string & model_paths, + int64_t inter_channels) + : impl_(std::make_unique()) { +#if defined(ENGINE_HAVE_RKNN) + impl_->load(model_paths, inter_channels); +#else + (void)model_paths; + (void)inter_channels; + throw std::runtime_error("Piper TTS duration encoder RKNN: built without ENGINE_ENABLE_RKNN"); +#endif +} + +PiperTtsRknnDurationEncoder::~PiperTtsRknnDurationEncoder() = default; + +int64_t PiperTtsRknnDurationEncoder::max_tokens() const { +#if defined(ENGINE_HAVE_RKNN) + return impl_->bins.empty() ? 0 : impl_->bins.back().tokens; +#else + return 0; +#endif +} + +void PiperTtsRknnDurationEncoder::encode( + const std::vector & token_ids, + std::vector & stats, + std::vector & conditioning) { +#if defined(ENGINE_HAVE_RKNN) + auto & impl = *impl_; + const int64_t tokens = static_cast(token_ids.size()); + auto bin = std::lower_bound( + impl.bins.begin(), + impl.bins.end(), + tokens, + [](const Impl::Bin & b, int64_t n) { return b.tokens < n; }); + if (bin == impl.bins.end()) { + throw std::runtime_error( + "Piper TTS duration encoder RKNN: " + std::to_string(tokens) + + " tokens exceed the widest compiled engine (" + + std::to_string(impl.bins.back().tokens) + ")"); + } + const auto start = Clock::now(); + impl.fed.assign(static_cast(bin->tokens), 0); + std::copy(token_ids.begin(), token_ids.end(), impl.fed.begin()); + impl.length[0] = tokens; + + // int64 inputs ride pass_through (the qwen3_tts codes precedent): the engine's + // native width for an index tensor is not a float conversion target. + std::vector inputs = { + {0, + impl.fed.data(), + static_cast(impl.fed.size() * sizeof(int64_t)), + RKNN_TENSOR_INT64, + 1}, + {1, + impl.length.data(), + static_cast(impl.length.size() * sizeof(int64_t)), + RKNN_TENSOR_INT64, + 1}, + }; + std::vector> outputs; + bin->engine->run(inputs, outputs); + + // De-pad the compiled width to the real token count: the outputs are + // channel-major ([channels][width]), the runtime's vectors are [channels][tokens]. + const int64_t width = bin->tokens; + stats.assign(static_cast(2 * impl.channels * tokens), 0.0F); + conditioning.assign(static_cast(impl.channels * tokens), 0.0F); + for (int64_t channel = 0; channel < 2 * impl.channels; ++channel) { + std::copy_n( + outputs[0].begin() + channel * width, + tokens, + stats.begin() + channel * tokens); + } + for (int64_t channel = 0; channel < impl.channels; ++channel) { + std::copy_n( + outputs[1].begin() + channel * width, + tokens, + conditioning.begin() + channel * tokens); + } + debug::timing_log_scalar("piper_tts.duration.rknn.bin", width); + debug::timing_log_scalar("piper_tts.duration.rknn.step_ms", debug::elapsed_ms(start)); +#else + (void)token_ids; + (void)stats; + (void)conditioning; + throw std::runtime_error("Piper TTS duration encoder RKNN: built without ENGINE_ENABLE_RKNN"); +#endif +} + +} // namespace engine::models::piper_tts diff --git a/src/community_models/piper_tts/session.cpp b/src/community_models/piper_tts/session.cpp index 0513c37db..048c7050e 100644 --- a/src/community_models/piper_tts/session.cpp +++ b/src/community_models/piper_tts/session.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -104,7 +105,8 @@ PiperTtsSession::PiperTtsSession( {{"ggml", {options.backend.type}, {}, {}}, {"rknn", {}, {"main"}, {}}}}, {"duration_encoder", "ggml", - {{"ggml", {core::BackendType::Cpu, options.backend.type}, {}, {}}}}, + {{"ggml", {core::BackendType::Cpu, options.backend.type}, {}, {}}, + {"rknn", {}, {"main"}, {}}}}, {"duration_flows", "ggml", {{"ggml", {core::BackendType::Cpu, options.backend.type}, {}, {}}}}}); const auto & decoder_plan = exec_plan.at("decoder"); @@ -121,7 +123,30 @@ PiperTtsSession::PiperTtsSession( ? engine::rknn::resolve_engine_path(decoder_plan.artifacts.at("main"), assets_->resources.model_directory()) : std::string{}, exec_plan.at("duration_encoder").backend, - exec_plan.at("duration_flows").backend); + exec_plan.at("duration_flows").backend, + // The encoder artifact names a comma-separated ladder of engines; each is + // named relative to the model like any other artifact. + [&]() { + const auto & encoder_plan = exec_plan.at("duration_encoder"); + if (encoder_plan.impl != "rknn") { + return std::string{}; + } + std::string out; + std::istringstream stream(encoder_plan.artifacts.at("main")); + for (std::string artifact; std::getline(stream, artifact, ',');) { + const auto at = artifact.find_first_not_of(" \t"); + if (at == std::string::npos) { + continue; + } + artifact = artifact.substr(at, artifact.find_last_not_of(" \t") - at + 1); + if (!out.empty()) { + out += ","; + } + out += engine::rknn::resolve_engine_path( + artifact, assets_->resources.model_directory()); + } + return out; + }()); } PiperTtsSession::~PiperTtsSession() = default; @@ -240,12 +265,28 @@ std::shared_ptr make_piper_tts_loader() { std::vector artifacts; try { const auto plan = io::json::parse(option->second); - const auto * decoder = plan.find("decoder"); - const auto * impl = decoder ? decoder->find("impl") : nullptr; - if (!impl || impl->as_string() != "rknn") { - return runtime::ModelResourcePlan{}; + for (const char * const component : {"decoder", "duration_encoder"}) { + const auto * stage = plan.find(component); + const auto * impl = stage ? stage->find("impl") : nullptr; + if (!impl || impl->as_string() != "rknn") { + continue; + } + // The encoder artifact is a comma-separated ladder of engines; every + // engine of a pinned stage is resident once the session loads. + const auto * stage_artifacts = stage->find("artifacts"); + if (!stage_artifacts) { + continue; + } + std::istringstream stream( + io::json::require_string(*stage_artifacts, "main")); + for (std::string artifact; std::getline(stream, artifact, ',');) { + const auto at = artifact.find_first_not_of(" \t"); + if (at != std::string::npos) { + artifacts.push_back( + artifact.substr(at, artifact.find_last_not_of(" \t") - at + 1)); + } + } } - artifacts.push_back(io::json::require_string(decoder->require("artifacts"), "main")); } catch (const std::runtime_error &) { // Session loading reports invalid execution plans. return runtime::ModelResourcePlan{}; diff --git a/src/community_models/piper_tts/vits_runtime.cpp b/src/community_models/piper_tts/vits_runtime.cpp index 15dba0619..808c4e68f 100644 --- a/src/community_models/piper_tts/vits_runtime.cpp +++ b/src/community_models/piper_tts/vits_runtime.cpp @@ -1,6 +1,7 @@ #include "engine/community_models/piper_tts/vits_runtime.h" #include "engine/community_models/piper_tts/rknn_decoder.h" +#include "engine/community_models/piper_tts/rknn_duration_encoder.h" #include "engine/framework/assets/tensor_source.h" #include "engine/framework/core/backend.h" #include "engine/framework/core/backend_weight_store.h" @@ -1273,7 +1274,8 @@ struct PiperVitsRuntime::State { core::BackendConfig backend_config, std::string rknn_decoder_path, std::optional duration_encoder_backend, - std::optional duration_flows_backend) + std::optional duration_flows_backend, + std::string rknn_encoder_path) : assets(std::move(assets_in)), threads(std::max(1, backend_config.threads)) { if (assets == nullptr) { @@ -1295,6 +1297,14 @@ struct PiperVitsRuntime::State { assets->config.inter_channels, samples_per_slice); } + // The encoder ladder is opt-in and the ggml encoder graphs stay built on demand, + // so an utterance wider than the widest bin (or a plan without the artifact) + // keeps the host encoder for that step. + if (!rknn_encoder_path.empty()) { + rknn_duration = std::make_unique( + rknn_encoder_path, + assets->config.inter_channels); + } // Each duration stage runs on the host when its exec_plan entry asks for the CPU // backend. With no explicit request the encoder keeps its measurement-driven default: // host duration alignment prevents small TF32 differences from being amplified by @@ -1428,6 +1438,7 @@ struct PiperVitsRuntime::State { std::shared_ptr duration_weights; bool duration_encoder_on_host = false; bool duration_flows_on_host = false; + std::unique_ptr rknn_duration; runtime::CacheSlots> duration_graphs{4}; std::array>, 3> duration_flow_graphs; @@ -1440,13 +1451,15 @@ PiperVitsRuntime::PiperVitsRuntime( core::BackendConfig backend_config, std::string rknn_decoder_path, std::optional duration_encoder_backend, - std::optional duration_flows_backend) + std::optional duration_flows_backend, + std::string rknn_encoder_path) : state_(std::make_unique( std::move(assets), backend_config, std::move(rknn_decoder_path), duration_encoder_backend, - duration_flows_backend)) {} + duration_flows_backend, + std::move(rknn_encoder_path))) {} PiperVitsRuntime::~PiperVitsRuntime() = default; @@ -1457,24 +1470,31 @@ runtime::AudioBuffer PiperVitsRuntime::synthesize( throw std::runtime_error("Piper TTS requires at least one text token"); } const auto total_start = std::chrono::steady_clock::now(); + const int64_t token_count = static_cast(token_ids.size()); + std::vector stats; + std::vector conditioning; auto duration_start = std::chrono::steady_clock::now(); - auto & duration = state_->duration_graph( - static_cast(token_ids.size())); - core::write_tensor_i32( - core::wrap_tensor( - duration.tokens, - core::TensorShape::from_dims( - {1, static_cast(token_ids.size())}), - GGML_TYPE_I32), - token_ids); - compute_graph(duration, "Piper TTS duration"); - const auto stats = core::read_tensor_f32(duration.stats); - const auto conditioning = core::read_tensor_f32(duration.conditioning); + if (state_->rknn_duration && token_count <= state_->rknn_duration->max_tokens()) { + // The compiled encoder covers this width: prior stats and flow conditioning + // straight off the NPU, flows and prior expansion stay on the host. + state_->rknn_duration->encode(token_ids, stats, conditioning); + } else { + auto & duration = state_->duration_graph(token_count); + core::write_tensor_i32( + core::wrap_tensor( + duration.tokens, + core::TensorShape::from_dims( + {1, static_cast(token_ids.size())}), + GGML_TYPE_I32), + token_ids); + compute_graph(duration, "Piper TTS duration"); + stats = core::read_tensor_f32(duration.stats); + conditioning = core::read_tensor_f32(duration.conditioning); + } engine::debug::timing_log_scalar( "piper_tts.duration_ms", engine::debug::elapsed_ms(duration_start)); - const int64_t token_count = static_cast(token_ids.size()); auto z = sampling::generate_normal_noise( static_cast(2 * token_count), options.seed); diff --git a/tools/rknn/prepare_piper_tts_rknn_vocoder.py b/tools/rknn/prepare_piper_tts_rknn_vocoder.py index 1ded91cda..4077893b9 100644 --- a/tools/rknn/prepare_piper_tts_rknn_vocoder.py +++ b/tools/rknn/prepare_piper_tts_rknn_vocoder.py @@ -23,6 +23,12 @@ mask-zeroed at the end that has no data) is the same one the runtime implements in src/community_models/piper_tts/rknn_decoder.cpp; change it here and there together. 24 is measured: at the 8+8 paroli ships, interior windows score 8-60 dB against a full-length decode. + +`--encoder` compiles the other half — the duration encoder — for a fixed phoneme-token width +(`--frames`, piper tokens = 2*phonemes + 3): + + /bin/python tools/rknn/prepare_piper_tts_rknn_vocoder.py \ + --voice-onnx en_US-lessac-medium.onnx --encoder --frames 128 """ from __future__ import annotations @@ -183,6 +189,242 @@ def mask_decoder_graph(model) -> int: return injected +ENCODER_INPUTS = ["input", "input_lengths"] + + +def slice_encoder(voice_onnx: Path, output: Path) -> None: + """Slice the duration encoder (text encoder + duration discriminator) at its seam. + + The seam is `enc_p.proj` (the 2x-inter_channels prior stats) and `dp.proj` (the flow + conditioning): the two arrays the runtime's flow graphs and length regulator consume. + Piper's export carries the frame mask inside the encoder (attention biases and the + x_mask muls are wired to `input_lengths`), so a right-padded fixed-width run reproduces + the exact-length run on the real frames — measured 117-130 dB fp32 — the same + pad-to-compile recipe the decoder window uses, with no mask surgery. + `scales` is not an input of the slice: all three scalars (noise_scale, length_scale, + noise_scale_w) feed the sampling, the duration ceiling and the discriminator's flow + noise, which sit strictly downstream of the seam — the slice is deterministic. + """ + model = onnx.load(str(voice_onnx), load_external_data=False) + stats = next((n for n in model.graph.node if n.name == "/enc_p/proj/Conv"), None) + conditioning = next((n for n in model.graph.node if n.name == "/dp/proj/Conv"), None) + if stats is None or conditioning is None: + raise SystemExit( + "no /enc_p/proj/Conv + /dp/proj/Conv: this ONNX is not a Piper/VITS voice" + ) + onnx.utils.extract_model( + str(voice_onnx), str(output), ENCODER_INPUTS, [stats.output[0], conditioning.output[0]]) + sliced = onnx.load(str(output)) + onnx.checker.check_model(sliced, full_check=False) + onnx.save(sliced, str(output)) + ops: dict[str, int] = {} + for node in sliced.graph.node: + ops[node.op_type] = ops.get(node.op_type, 0) + 1 + print( + f"sliced {voice_onnx.name} at {stats.output[0]!r} + {conditioning.output[0]!r}:" + f" {len(sliced.graph.node)} nodes, inputs {ENCODER_INPUTS}," + f" outputs {[o.name for o in sliced.graph.output]}" + ) + risky = sorted( + op for op in ops + if op in {"Pad", "Exp", "Shape", "ConstantOfShape", "Split", "Atan", "Erf"} + ) + if risky: + print(f" note: graph contains {risky}; if rknn.build aborts in fold_constant, " + "the exporting torch/onnx pair is too new for the toolkit") + report_engine(output) + + +def check_encoder(onnx_path: Path, raw_path: Path, frames: int) -> None: + """Right-padding the folded fixed-`frames` graph must not move the real frames. + + Exact-width runs come from the raw dynamic slice (the folded graph is fixed-shape by + construction); >= 100 dB fp32 against them is the pad-safety gate, on the graph the + engine is actually compiled from. + """ + exact_session = ort_session(raw_path) + pad_session = ort_session(onnx_path) + rng = np.random.default_rng(7) + for n in (17, 61, frames - 1): + tokens = rng.integers(3, 250, size=(1, n), dtype=np.int64) + padded = rng.integers(0, 250, size=(1, frames), dtype=np.int64) + padded[:, :n] = tokens[0] + lengths = np.array([n], np.int64) + exact = exact_session.run(None, {"input": tokens, "input_lengths": lengths}) + pad = pad_session.run(None, {"input": padded, "input_lengths": lengths}) + for out, a, b in zip(pad_session.get_outputs(), exact, pad): + a = np.asarray(a).reshape(a.shape[0], -1, a.shape[-1])[..., :n] + b = np.asarray(b).reshape(b.shape[0], -1, b.shape[-1])[..., :n] + snr = _snr(a.astype(np.float64), b.astype(np.float64)) + print(f" {out.name} n={n} padded to {frames}: SNR {snr:.1f} dB") + if snr < 100: + raise SystemExit(f"right-padding moves the encoder output at n={n} ({snr:.1f} dB)") + + +def fold_encoder(onnx_path: Path, output: Path, frames: int) -> None: + """Static-crop and constant-fold the encoder slice with ONNX Runtime's graph optimizer. + + The sliced graph computes its own masks and pads from the dynamic frame axis (25 Shape, + 48 Pad, 1 Range), and rknn-toolkit2's C++ rewriter aborts without a message on exactly + those shape-computation chains — the crash family of the qwen3 full-decode export. With + the frame axis pinned to `frames` the chains become constant, the optimizer folds them + into initializers, and the toolkit gets the static graph it can compile. The pad-safety + self-check runs on the folded graph, so a fold that moved anything is caught here. + """ + import onnxruntime as ort + + model = onnx.load(str(onnx_path)) + for value in model.graph.input: + dims = value.type.tensor_type.shape.dim + for dim in dims: + dim.dim_param = "" + dim.dim_value = 0 + if value.name == "input": + dims[0].dim_value = 1 + dims[1].dim_value = frames + elif value.name == "input_lengths": + dims[0].dim_value = 1 + else: + raise SystemExit(f"unexpected encoder input {value.name!r}") + for value in model.graph.output: + for dim in value.type.tensor_type.shape.dim: + dim.dim_param = "" + dim.dim_value = 0 + pinned = output.with_suffix(".pinned.onnx") + onnx.save(model, str(pinned)) + + options = ort.SessionOptions() + options.log_severity_level = 3 + # BASIC, not ALL: the extended levels emit the com.microsoft fusion ops (FusedConv, + # FusedMatMul, Gelu) into the saved model, and the RKNN toolkit drowns on them. + # Constant folding — the point of the pass — runs at BASIC. + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC + options.optimized_model_filepath = str(output) + ort.InferenceSession(str(pinned), options, providers=["CPUExecutionProvider"]) + + folded = onnx.load(str(output)) + onnx.checker.check_model(folded, full_check=False) + ops: dict[str, int] = {} + for node in folded.graph.node: + ops[node.op_type] = ops.get(node.op_type, 0) + 1 + shapes = sorted( + op for op in ("Shape", "Range", "ConstantOfShape", "NonZero", "GatherND", "ScatterND") + if op in ops + ) + print( + f"folded to fixed {frames}-token width: {len(folded.graph.node)} nodes, " + f"shape ops {shapes or 'none'}" + ) + + +def matrix_form_encoder(onnx_path: Path, output: Path, frames: int) -> None: + """Rewrite the per-row relative-shift chains of the attention layers into gathers. + + With the frame axis pinned, every piper attention layer shifts its [N, 2N-1] + relative-position slabs by one step per query row (Pad + Reshape + Slice, unrolled per + row), and the RKNN NPU kernels compute exactly those chains as garbage on the board: + probes show the flatten-Pad tails filled with uninitialized memory (NaN, ±65504) and + one such layer drags the engine to 0.4 dB end to end, while the PC simulator computes + the same graph at 56 dB. The two chains are exact gathers with constant indices: + + bias[b,h,w,t] = rel_logits[b,h,w,(N-1)+t-w] (GatherElements) + rel_out[w] = sum_j probs[b,h,w,j] * rel_v[(N-1)+j-w] (Gather + Mul + ReduceSum) + + The indices are constants, so the toolkit folds them back into weights at build: + the shift costs O(N^3) per layer — measured engines 137 MB at N=128, 402 MB at 192, + 2.9 GB at 384, and N=512 exhausts 31 GB of converter RAM. Compile narrow bins only + and leave the wider widths to the host encoder. + """ + from onnx import numpy_helper, shape_inference + + model = shape_inference.infer_shapes( + onnx.load(str(onnx_path), load_external_data=True), strict_mode=True) + heads = next( + (list(v.type.tensor_type.shape.dim)[1].dim_value + for v in model.graph.value_info + if v.name == "/enc_p/encoder/attn_layers.0/Softmax_output_0"), 0) + if not heads: + raise SystemExit("no attn_layers.0/Softmax shape: not a folded piper encoder") + producers = {output_tensor: node for node in model.graph.node for output_tensor in node.output} + consumers: dict[str, list] = {} + for node in model.graph.node: + for i in node.input: + consumers.setdefault(i, []).append(node) + shift = (frames - 1) + np.arange(frames)[None, :] - np.arange(frames)[:, None] + drop, added, inits = set(), [], [] + for layer in range(6): + p = f"/enc_p/encoder/attn_layers.{layer}/" + bias_t = producers[p + "Add_1_output_0"].input[1] + chain, t = [], bias_t + while producers[t].name != p + "Pad_1": + chain.append(producers[t]) + t = producers[t].input[0] + chain.append(producers[t]) + relv = producers[p + "MatMul_3_output_0"] + rel_v = next(i for i in relv.input if i != p + "Slice_8_output_0") + if len(consumers[p + "Slice_8_output_0"]) != 1: + raise SystemExit(f"{p}Slice_8 has foreign consumers; unexpected fold shape") + drop.update(id(n) for n in chain) + for x in (p + "Pad_3_output_0", p + "Reshape_13_output_0", p + "Pad_4_output_0", + p + "Reshape_16_output_0", p + "Slice_8_output_0"): + drop.add(id(producers[x])) + drop.add(id(relv)) + added.append(helper.make_node( + "GatherElements", [producers[p + "Pad_1_output_0"].input[0], p + "idxE"], + [bias_t], axis=3, name=p + "RelBiasGather")) + inits.append(numpy_helper.from_array( + np.broadcast_to(shift, (1, heads, frames, frames)).astype(np.int64), p + "idxE")) + added.append(helper.make_node( + "Gather", [rel_v, p + "idxG"], [p + "G_gathered"], axis=2, name=p + "RelVGather")) + inits.append(numpy_helper.from_array(shift.astype(np.int64), p + "idxG")) + added.append(helper.make_node( + "Unsqueeze", [p + "Softmax_output_0", "/enc_p/encoder/Axes4"], + [p + "ProbsUnsq"], name=p + "ProbsUnsq")) + added.append(helper.make_node( + "Mul", [p + "ProbsUnsq", p + "G_gathered"], [p + "RelVProd"], name=p + "RelVMul")) + added.append(helper.make_node( + "ReduceSum", [p + "RelVProd", "/enc_p/encoder/Axes3"], + [p + "MatMul_3_output_0"], name=p + "RelVSum", keepdims=0)) + model.graph.initializer.extend([ + numpy_helper.from_array(np.array([4], np.int64), "/enc_p/encoder/Axes4"), + numpy_helper.from_array(np.array([3], np.int64), "/enc_p/encoder/Axes3"), + ]) + model.graph.initializer.extend(inits) + # rknn-toolkit2's fold_constant walks the node list once, in order: hand it a + # topologically sorted graph. + nodes = [n for n in model.graph.node if id(n) not in drop] + added + available = {i.name for i in model.graph.initializer} | {i.name for i in model.graph.input} + ordered, pending = [], list(nodes) + while pending: + ready = [n for n in pending if all(not i or i in available for i in n.input)] + if not ready: + raise SystemExit("cycle in the rewritten encoder graph") + for n in ready: + pending.remove(n) + available.update(o for n in ready for o in n.output) + ordered.extend(ready) + del model.graph.node[:] + model.graph.node.extend(ordered) + del model.graph.value_info[:] + model = shape_inference.infer_shapes(model, strict_mode=True) + onnx.save(model, str(output)) + print(f"rewrote {6} attention shift chains to gathers: {len(model.graph.node)} nodes") + + +def convert_rknn_encoder( + onnx_path: Path, + output: Path, + target: str, + frames: int, + opt_level: int | None = None, +) -> None: + from rknn_build import build_engine, report_engine + + # The folded graph is static: declare nothing (see build_engine). + build_engine(onnx_path, output, target, opt_level=opt_level) + report_engine(output, "inputs input/input_lengths", f"frames {frames}") + + def decode_window(session, z: np.ndarray, mask: np.ndarray) -> np.ndarray: feeds = {"z": z.astype(np.float32)} for value in session.get_inputs(): @@ -446,6 +688,50 @@ def convert_rknn( report_engine(output, f"inputs {names}", f"window {window}") +def verify_encoder( + onnx_path: Path, + raw_path: Path, + frames: int, + target: str, + seed: int = 7, +) -> None: + """The fp16 cost of the compiled encoder, measured on the PC simulator. + + Rebuilt from the folded ONNX rather than run from the exported `.rknn`: a + `load_rknn` engine cannot be inferred on the simulator (see the cosyvoice3 + flow converter), and `load_onnx` + `build` with the same config produces the + same fp16 graph the board runs. Reference: the raw dynamic slice at exact + width, fp32. + """ + from rknn.api import RKNN + + exact = ort_session(raw_path) + rng = np.random.default_rng(seed) + rknn = RKNN(verbose=False) + rknn.config(target_platform=target) + if rknn.load_onnx(model=str(onnx_path)) != 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") + for n in (17, 61, frames - 1): + tokens = rng.integers(3, 250, size=(1, n), dtype=np.int64) + padded = rng.integers(0, 250, size=(1, frames), dtype=np.int64) + padded[:, :n] = tokens[0] + lengths = np.array([n], np.int64) + want = exact.run(None, {"input": tokens, "input_lengths": lengths}) + got = rknn.inference(inputs=[padded, lengths]) + for out, a, b in zip(exact.get_outputs(), want, got): + a = np.asarray(a).reshape(a.shape[0], -1, a.shape[-1])[..., :n] + b = np.asarray(b).reshape(b.shape[0], -1, b.shape[-1])[..., :n] + snr = _snr(a.astype(np.float64), b.astype(np.float64)) + print(f" {out.name} n={n} fp16 engine vs fp32 slice: SNR {snr:.1f} dB") + if snr < 30: + raise SystemExit(f"compiled encoder moved at n={n} ({snr:.1f} dB)") + rknn.release() + + def main() -> int: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--voice-onnx", type=Path, help="Piper voice ONNX (text in, PCM out)") @@ -455,6 +741,21 @@ def main() -> int: parser.add_argument("--window", type=int, default=CHUNK + 2 * PADDING) parser.add_argument("--skip-rknn", action="store_true", help="slice and check only") parser.add_argument("--skip-check", action="store_true") + parser.add_argument( + "--encoder", action="store_true", + help="slice/compile the duration encoder (enc_p.proj + dp.proj seam, phoneme-width " + f"bin --frames) instead of the vocoder window") + parser.add_argument( + "--frames", type=int, default=128, + help="--encoder: compiled phoneme-token width (piper tokens = 2*phonemes + 3)") + parser.add_argument( + "--verify-sim", action="store_true", + help="--encoder: after compiling, infer the engine on the PC simulator against " + "the fp32 slice (fp16 cost of the encoder)") + parser.add_argument( + "--opt-level", type=int, default=None, + help="--encoder: rknn optimization level (wide bins whose kernels overflow the " + "NPU register fields at the default level: try 2)") parser.add_argument( "--dec-mask", action="store_true", help="slice stage: multiply every decoder conv output by the (rate-upsampled) y_mask, " @@ -464,8 +765,35 @@ def main() -> int: if args.voice_onnx is None and args.onnx is None: parser.error("pass --voice-onnx and/or --onnx") - if args.window <= CHUNK: - parser.error("--window must exceed the emitted chunk size (45)") + if not args.encoder: + if args.window <= CHUNK: + parser.error("--window must exceed the emitted chunk size (45)") + + if args.encoder: + sliced = args.onnx + if args.voice_onnx is not None: + sliced = (args.output + if args.output and args.output.suffix == ".onnx" + else Path(f"encoder_n{args.frames}.onnx")) + slice_encoder(args.voice_onnx, sliced) + if sliced is None: + parser.error("--onnx is required (or pass --voice-onnx to slice one)") + folded = sliced.with_name(sliced.stem + "_static.onnx") + fold_encoder(sliced, folded, args.frames) + matrix = sliced.with_name(sliced.stem + "_matrix.onnx") + matrix_form_encoder(folded, matrix, args.frames) + if not args.skip_check: + print("self-check: encoder pad vs exact (matrix-form graph)") + check_encoder(matrix, sliced, args.frames) + if args.skip_rknn: + return 0 + output = args.output if args.output and args.output.suffix == ".rknn" else Path( + sliced.with_suffix(".rknn")) + convert_rknn_encoder(matrix, output, args.target, args.frames, args.opt_level) + if args.verify_sim: + print("verify: compiled fp16 encoder on the PC simulator") + verify_encoder(matrix, sliced, args.frames, args.target) + return 0 sliced = args.onnx if args.voice_onnx is not None: