gtcrn

GTCRN — Grouped Temporal Convolutional Recurrent Network (svod port)

Speech enhancement model ported to run natively in svod — a pure-Rust inference stack on top of a JIT-compiled tensor runtime. This repo holds the converted weights only; the model implementation lives in svod-model (model/src/gtcrn/).

Ultralight real-time noise suppression: 48.8 K parameters, 33.0 MMACs/s. Takes a noisy 16 kHz mono waveform, returns an enhanced waveform.

Source

A faithful Rust port of the upstream Xiaobin-Guang/GTCRN, checkpoint model_trained_on_dns3.tar (trained on the DNS-Challenge 3 dataset). The upstream architecture and training are unchanged; this repo republishes the weights in safetensors with the small remaps described below.

Architecture

noisy WAV → STFT(n_fft=512, hop=256, √hann)
         → ERB analysis (257 bins → 129 bands)
         → SFE (subband unfold)
         → Encoder: 2× ConvBlock + 3× GTConvBlock (ShuffleNetV2)
         → 2× DPGRNN (dual-path grouped RNN: intra-frame bidir GRU + inter-frame GRU)
         → Decoder: 3× GTConvBlock (transpose) + 2× ConvBlock (transpose)
         → ERB synthesis (129 → 257)
         → complex ratio mask × input spectrogram
         → ISTFT → enhanced WAV

Usage

# Build svod-model and run the bundled example (noisy.wav → enhanced.wav):
cargo run -p svod-model --release --example gtcrn_enhance -- \
    --in noisy.wav --out enhanced.wav --hub

In Rust:

use svod_model::gtcrn::{Gtcrn, GtcrnJit};
use svod_model::jit::InputSpec;

let model = Gtcrn::from_hub()?;            // pulls gtcrn.safetensors from this repo
let mut jit = GtcrnJit::new(model);
jit.prepare(InputSpec::f32(&[1, 257, T, 2]))?;  // T = number of STFT frames
// copy the [1, 257, T, 2] complex spectrogram into jit.spec_mut()?, then:
jit.execute()?;
let enhanced = jit.output()?;

STFT/ISTFT run eagerly on the host via realfft; the network forward pass is JIT-compiled. See model/examples/gtcrn_enhance.rs for the full waveform → waveform pipeline (it processes long audio in fixed-size frame chunks because the in-graph GRU recurrence unrolls one IR node per time step).

Files

Conversion notes

Generated by scripts/convert_gtcrn.py (run uv run scripts/convert_gtcrn.py --selfcheck to reproduce). Two remaps from the upstream PyTorch checkpoint:

1. GRU gate order. PyTorch nn.GRU stores gate rows as [reset, update, new]; svod's gru() op expects [z, r, h]. The first two hidden-sized gate blocks of every GRU weight/bias are swapped.

2. Bidirectional key split. The DPGRNN's bidirectional grouped RNNs expose rnn1/rnn2 modules whose reverse-direction weights (*_l0_reverse) are renamed to separate rnn1_b/rnn2_b keys, matching svod's representation of a bidirectional GRU as two unidirectional passes (forward over the sequence + forward over the time-flipped sequence, concatenated).

num_batches_tracked entries are dropped; BatchNorm running_var is kept verbatim (svod folds it into invstd = 1/√(var+ε) at load time). The --selfcheck flag verifies the GRU remap against svod's documented recurrence equations.

Verification

The svod forward matches the upstream PyTorch GTCRN.forward to ~6 significant figures on a fixed-frame slice (max |Δ| relative < 1e-3), validated by the gtcrn::parity test in svod-model.

License

MIT — same as the upstream model. Weights © their respective authors; this repo only republishes them in a converted format for use with svod.