Skip to content
heliaEDGE
Examples
HELIA

FastEnhancer (streaming speech enhancement)

helia_edge.models.fastenhancer.build builds the inference form of FastEnhancer as a Keras Functional model that processes one STFT frame per call. FastEnhancerParams holds the backend-free architecture config; rnnformer.freq may not exceed n_fft / 2 / stride. The fastenhancer_t preset matches upstream configs/fastenhancer/t.yaml at revision e74cab1, the revision of the onnx-vd-v1.0.0 release.

import numpy as np
from helia_edge.importers import import_weights
from helia_edge.models import FASTENHANCER_T_ONNX, ModelSpec, build, resolve_fastenhancer
record = resolve_fastenhancer("fastenhancer_t") # serialize record.model_dump_json() with results
model = build(ModelSpec(params=record.params), batch_size=1) # the input shape follows from the params
import_weights(model, FASTENHANCER_T_ONNX, "fastenhancer_t.spec.onnx") # the downloaded release file
states = [np.zeros((1, 16, 20), np.float32) for _ in range(2)] # zero once per sequence
for spec_frame in frames: # each (1, 257, 1, 2): real/imag of one 512-point STFT frame
spec_out, *states = model([spec_frame, *states])
Name Shape per example Role
spec_in (n_fft/2+1, 1, 2) Current spectral frame, real/imag last
state_in_k (freq, channels) GRU state of RNNFormer block k
spec_out (n_fft/2+1, 1, 2) Enhanced frame; the Nyquist bin is always zero
state_out_k (freq, channels) Next state for block k

State is explicit input and output, named like every heliaEDGE streaming model so export and LiteRTStreamRunner treat them as state pairs. For integer I/O, build with batch_size=1 and export in keras mode, so each state pair shares one scale; saved_model mode exports a dynamic batch and cannot tie the state.

Zero every state at the start of each independent sequence, then feed state_out_k back as state_in_k for every following frame. Resetting the states every frame is a different, much worse model. STFT framing, iSTFT/overlap-add, startup delay and flushing are the caller’s responsibility; the upstream reference uses a periodic Hann window of 512, hop 256, an unnormalized one-sided STFT and centered reflection padding.

The graph matches the released ONNX export. BatchNorm is folded into convolution and dense biases, weight normalization is removed, and the frequency projections around the RNNFormer are fixed triangular filterbanks. It is not the trainable architecture. Every RNNFormer block runs:

  • a GRU step per frequency band, sharing weights across bands;
  • a projection and residual;
  • self-attention across bands, then a projection and residual.

The first block also adds a learned positional embedding.

Weights load through import_weights. FASTENHANCER_T_ONNX is pinned to the fastenhancer_t.spec.onnx release (sha256 915a451f…), whose anonymous initializers it names; it does not guess names for other exports. For other folded weights, build a mapping with fastenhancer_mapping(params, name, source) (in helia_edge.models.fastenhancer_params): it reads tensors in ONNX export layout keyed by module path, with the shapes listed by fastenhancer_weight_shapes(params), from a pinned .safetensors, .onnx or PyTorch file. Pass tensor_names for tensors stored under other names, and list every other tensor of the file, such as an ONNX file’s graph constants, in unused. GRU tensors must use ONNX gate order z, r, h; PyTorch nn.GRU stores r, z, n and must be reordered first, because shapes alone cannot detect the difference. Only names and shapes are checked, so build the model from the params the tensors came from. Any missing, unexpected or misshapen tensor is refused before anything loads.

The mapping applies the following layout conversions:

  • Strided input convolution: the source’s reshape-based channel order is converted to a standard strided kernel.
  • GRU: gates keep ONNX order z, r, h, and the reset gate is applied after the recurrent matmul.
  • Attention: qkv columns are interleaved per head.
  • Transposed convolution: runs with valid padding and an explicit crop.

The upstream code is MIT licensed. Its notice is retained in helia_edge/models/licenses/fastenhancer-mit.txt. The upstream release publishes pretrained weights without a separate weight license, so helia-edge does not bundle or download them. Obtain and verify weights, and confirm their terms, yourself.

With the released T weights, the model matches the upstream ONNX Runtime reference to within about 1.4e-5 absolute (spectral peak about 30) and 4e-6 on the GRU states. This holds over a 248-frame carried-state sequence on TensorFlow and PyTorch CPU backends. PyTorch on GPUs that allow TF32 convolutions (the default for cuDNN) differs by about 7e-3, so disable TF32 for FP32 comparisons.

Other options (residuals, sigmoid/tanh masks, post-activations, attention bias, other sizes) are tested only against an independent NumPy reference of the exported graph; no released weights exercise them.

Enhanced-audio quality, conversion to LiteRT, quantization and target timing or memory are not established by this model.