Add support for Audio8 TTS 0.1B-ONNX-INT8 - #110
Draft
cryptomilk wants to merge 15 commits into
Draft
Conversation
Owner
|
memmap2 is ok, but does it necessary for inference or just for speedup model loading? btw, how does current speed compare with original Audio8 tts onnx demo. Oh, and what if users CPU doesn't support int8 operation, wouldn't that very slow converting to float32 |
Contributor
Author
|
mmap2 is for speedup. It is not strictly necessary. I don't have any speed comparison yet. I finally got the model doing some work and fixed issues along the way. I will do benchmarks once I get it fully working. |
Contributor
Author
onnx::eval doesn't implement MatMulInteger/DynamicQuantizeLinear at all and candle has no Int8 dtype to run a competitive integer GEMM around. This trades an unsupported op for a working F32 path. If this will be a bottleneck we can implement an Int8-typed GEMM kernel later. |
Owner
|
Our primary support would be at least support M series mac and CUDA, and then some decent intel CPUs, i can help test model name : Intel(R) Core(TM) i7-14700KF with AXV-VNNI int8 support if kernels ready |
Contributor
Author
|
candle-onnx and our copy in Crane is CPU only. It hardcodes the device as |
Contributor
Author
|
I have some improvements. I'm currently trying to find bugs generating noise. Once I have the quality on par, I will look into speedup. We are super slow right now. |
Contributor
Author
|
FYI: For English, Kokoro sounds a lot better than Audio8. |
Owner
this should be fixed, only CPU not reasonable. Audio8 0.1b should be slow on CPU, but it supports various languages and styles can be customized too. |
Contributor
Author
|
Hmm VoxCPM2 is much better than Audio8. |
Owner
|
really, but seems VoxCPM2 now also very slow. I had tested on macOS. |
Contributor
Author
|
On CPU currently VoxCPM2 needs 10 sec and Audio8 35 sec. Audio8 can still be optimized more, but using upstream onnx runtime it is still like 10sec. |
Owner
|
why VoxCPM in Crane on macos so slow, have u tested with crane-serve with the new ui support> |
Contributor
Author
|
I normally use |
Some quantized transformer exports use Softplus on time-step pre-activations; the evaluator had no match arm for it. Implemented with a numerically stable large-x branch so exp(x) can't overflow. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Codec encoder/decoder graphs use Einsum for sinusoidal positional embeddings, always with the "i,j->ij" outer-product equation. Implemented just that case; any other equation bails with a clear error instead of attempting a general einsum interpreter. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Some ONNX exports quantize linears/embeddings with ORT's dynamic INT8 pattern (DynamicQuantizeLinear/MatMulInteger/DequantizeLinear), which crate::onnx::eval doesn't implement and candle has no Int8 dtype to build a competitive integer GEMM around. Since the quantized weight or table operand in both patterns is always a static graph initializer, add an optimizer pass that dequantizes it to F32 once at session-load time and rewrites the whole chain into a plain Gather/MatMul, so no runtime int8 arithmetic is needed at all. Also extends eval.rs's get_tensor to decode ONNX Int8 tensors (widened losslessly to candle's I16, since candle has no I8), which the new pass's MatMulInteger weights require. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Crane's ONNX loader only decoded inline TensorProto bytes, so any model storing large initializers in a sidecar `.onnx.data` file instead (the standard ONNX convention past the ~2GB inline-protobuf limit) panicked on load with `range end index 512 out of range for slice of length 0` instead of an error. audio8-TTS-0.1B-ONNX-INT8's slow/fast AR graphs hit this: 316 of 826 initializers in slow_ar_int8.onnx are external. Add crane-core/src/onnx/external_data.rs to resolve any TensorProto with data_location == EXTERNAL, inlining its bytes from the sidecar file resolved relative to the .onnx file's directory, into raw_data. Walks graph.initializer plus node-attribute tensors and subgraphs (e.g. an If node's branches), and caches each sidecar's bytes across the whole walk since one sidecar is typically shared by hundreds of tensors. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Crane's ONNX dtype mapping had no arm for DataType::Bfloat16, so any graph with a bf16 Constant node failed to load with "unsupported 'value' data-type Bfloat16". candle already decodes DType::BF16 via half::bf16 with the same raw byte layout ONNX uses, so only the mapping itself was missing. Found while loading Audio8-TTS's fast/slow AR ONNX graphs, which have bf16 Constant nodes. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The Shape op built its output tensor with shape xs.rank() instead of dims.len() (the number of dims the start..=end loop actually collected). Those only coincide when start/end default to the full range; with explicit start/end attributes selecting a shorter slice, Tensor::from_vec's declared shape mismatched dims's real element count. Found by inspection while debugging a separate Audio8-TTS failure; no current Crane-supported graph exercises non-default start/end on this op, so this is a latent-bug fix rather than an observed crash. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
ScatterND's "none"/"add" reduction branches call flat_output.slice_scatter(&update_slice, ...) against flat_output, which is always rank-1. For a scalar update (one value per index), update_slice was correctly squeezed to match; for a slice update (a product-sized value per index), the squeeze was skipped, leaving update_slice at rank 2 ((1, product)) instead of rank 1. Every slice_scatter call on that path then failed with: unexpected rank, expected: 1, got: 2 ([1, 64]) Found while debugging Audio8-TTS's fast AR ONNX graph, which hit this on its very first forward pass. Squeezing dim 0 unconditionally fixes both the scalar and slice-update cases identically. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Crane's Pad op only implemented mode="reflect" and mode="edge", rejecting the ONNX default mode="constant" outright, and rejecting any node with a 3rd (constant_value) input even when that input was the standard empty-string "omitted" marker. Add mode="constant" support (padding with an explicit constant_value, defaulting to 0 per the ONNX spec when omitted), and let eval.rs's Pad arm read an optional constant_value via get_opt instead of bailing on any 3-input node. Found while running Audio8-TTS's codec decoder, whose upsampling depthwise convolutions use a 3-input Pad node with an omitted constant_value. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Implements the DualAR generation loop for Audio8-TTS-0.1B-ONNX-INT8 on top of Crane's ONNX evaluator: ChatML prompt construction with the reference voice's codes spliced in as semantic tokens (config.rs, prompt.rs), a slow AR step that threads persistent KV-cache and Mamba conv/ssm state across calls, a fast AR step that expands each frame into 10 codec codebook indices, lazy Repetition-Avoidance Sampling that only pays for the high-temperature resample when a repeat is actually detected (sampling.rs), and codec decoding to a waveform (model.rs). Fast/slow AR cache and Mamba state shapes are read directly from each ONNX graph's declared input shapes rather than derived from manifest fields, so a mismatched manifest fails loudly at load time instead of silently misallocating a buffer. reference_codes.npy is parsed with a small hand-rolled NumPy reader, since its dtype (<i8, i.e. int64) and 2-D/C-order layout don't need a general-purpose crate. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Audio8-TTS's inherent generate_speech already matches the Tts trait's shape ((Tensor, u32) return, same params), so this implements the trait directly on the model type rather than needing a wrapper like Kokoro's. voices() returns empty since the model ships one bundled reference voice with no selection. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Adds the same integration checklist every TTS backend follows: a ModelType variant, CLI aliases, display name, is_tts() membership, auto-detection (runtime_manifest.json presence, unique to this package, plus a path-name fallback), and instantiation in create_tts() gated behind the onnx feature. create_backend() bails toward create_tts() like the other TTS types. No changes needed in handlers/tts.rs or lib.rs since those dispatch generically through dyn Tts. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Audio8-TTS's codec decoder computes inv_alpha dynamically via Reciprocal(alpha) behind a redundant Identity passthrough, and stores the Pow exponent as a Constant op node rather than a raw initializer. The optimizer's pattern match required both to be literal initializers, so it silently fused 0 of the decoder's 29 Snake activations, leaving each to run as five separate full-tensor passes instead of one fused kernel. Confirmed via CRANE_ONNX_OPT_REPORT=1 against the real downloaded package: fused_snake went from 0 to 29. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
SpeechOptions::temperature was a plain f64 defaulting to 0.9 for every TTS model, with no way for a model to override it the way top_p already can via Option<f64>. Audio8-TTS's 0.1B Preview model is far noisier recovering from a mid-utterance pause at that temperature: on the same input/seed, measured high-frequency burst energy (6.5-11kHz vs 200-3000Hz) dropped roughly 18x lowering temperature from 0.9 to 0.3, and 0.3 is this model family's own reference service's actual default.
Gives evaluator performance work (string interning, compiled execution plan) a before/after measurement instead of ad-hoc manual timing runs. generate_speech samples from a fixed internal seed, so the benchmark retraces identical work every run, isolating timing deltas to the evaluator rather than sampling variance. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
cryptomilk
force-pushed
the
asn-audio8
branch
from
September 1, 2026 05:09
8369784 to
68a1901
Compare
Contributor
Author
|
Now you can test it. |
Owner
|
how's the speed now |
Contributor
Author
|
Super slow. I didn't do any optimizations yet. Just wanted an implementation. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This is work in progress to add support for Audio8 TTS 0.1B-ONNX-INT8.
@lucasjinreal can you take a look if those changes are fine so far? I used memmap2 crate