GTCRN

https://github.com/xiaobin-rong/gtcrn
Admin
GTCRN is an ultralow-complexity speech enhancement model with only 48.2K parameters that achieves...
Tokens:6,335
Snippets:49
Trust Score:9.7
License:MIT
Update:5 months ago
Show doc for...
Tokens:
Raw
### Run PipeWire with Filter Chain

Source: https://github.com/xiaobin-rong/gtcrn/blob/main/ladspa/README.md

Starts PipeWire with the filter chain configuration enabled, applying the specified audio filters. This command assumes the filter-chain.conf file is correctly set up.

```bash
pipewire -c filter-chain.conf
```

--------------------------------

### ONNX Export and Runtime Inference Setup

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Imports necessary libraries for ONNX export and runtime inference, including PyTorch, NumPy, ONNX, ONNX Simplifier, and soundfile. This setup is for exporting and running the StreamGTCRN model.

```python
import numpy as np
import onnx
import onnxruntime
from onnxsim import simplify
import torch, soundfile as sf
from librosa import istft
```

--------------------------------

### Build Static LADSPA Plugin (Bundled)

Source: https://github.com/xiaobin-rong/gtcrn/blob/main/ladspa/README.md

Builds the LADSPA plugin by downloading and bundling the official ONNX Runtime from Microsoft. This option is easier than Docker but results in a larger file size.

```bash
./build.sh static
```

--------------------------------

### Build Dynamic LADSPA Plugin (System Lib)

Source: https://github.com/xiaobin-rong/gtcrn/blob/main/ladspa/README.md

Builds the LADSPA plugin using a system-installed ONNX Runtime. This is the fastest option for development but requires the 'onnxruntime' library to be present on the system.

```bash
./build.sh dynamic
```

--------------------------------

### Build and Configure LADSPA Plugin for PipeWire

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Builds a minimal static LADSPA plugin for GTCRN denoising and configures PipeWire to use it for system-wide noise suppression.

```bash
cd ladspa
./build-minimal-docker.sh          # ~5-20 min first run; builds stripped ORT
./build.sh minimal                 # compiles the .so

sudo cp target/release/libgtcrn_ladspa_ort.so /usr/lib/ladspa/libgtcrn_ladspa.so

mkdir -p ~/.config/pipewire/filter-chain.conf.d
cat > ~/.config/pipewire/filter-chain.conf.d/gtcrn.conf << 'EOF'
context.modules = [
    {
        name = libpipewire-module-filter-chain
        args = {
            node.description = "Noise Canceling Microphone (GTCRN)"
            media.name       = "Noise Canceling Microphone (GTCRN)"
            filter.graph = {
                nodes = [
                    {
                        type    = ladspa
                        name    = "gtcrn"
                        plugin  = "libgtcrn_ladspa"
                        label   = "gtcrn_mono"
                        control = {
                            "Strength"           = 1.0
                            "Model (0=Light 1=Full)" = 1
                        }
                    }
                ]
            }
            audio.channels = 1
            capture.props  = { node.passive = true }
            playback.props = { media.class = "Audio/Source" }
        }
    }
]
EOF

pipewire -c filter-chain.conf
```

--------------------------------

### Initialize and Use HybridLoss

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Instantiate the HybridLoss function, which combines spectral and waveform domain losses. It requires predicted and true STFT tensors as input.

```python
import torch
from loss import HybridLoss

loss_fn = HybridLoss()

# Simulated batch: predicted and clean STFT (B, F, T, 2)
pred_stft = torch.randn(2, 257, 100, 2, requires_grad=True)
true_stft = torch.randn(2, 257, 100, 2)

loss = loss_fn(pred_stft, true_stft)
print(f"Loss: {loss.item():.4f}")  # e.g., Loss: 12.3456

loss.backward()
print(pred_stft.grad.shape)  # torch.Size([2, 257, 100, 2])
```

--------------------------------

### Initialize and Use DPGRNN Model

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Instantiate the DPGRNN model with specified input, width, and hidden sizes. The model processes encoded feature cubes of shape (B, C, T, F).

```python
import torch
from gtcrn import DPGRNN

dpgrnn = DPGRNN(input_size=16, width=33, hidden_size=16)

# Input: (B, C, T, F) — typical internal feature map after encoding
feat = torch.randn(1, 16, 50, 33)
out = dpgrnn(feat)
print(out.shape)  # torch.Size([1, 16, 50, 33])
# Intra (freq) + inter (time) recurrence applied with residual connections
```

--------------------------------

### Runtime Inference with ONNX Model

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Performs runtime inference using the simplified ONNX model. Reads audio, processes it chunk by chunk, and saves the enhanced audio.

```python
session = onnxruntime.InferenceSession(
    'stream/onnx_models/gtcrn_simple.onnx',
    providers=['CPUExecutionProvider']
)

mix, fs = sf.read('stream/test_wavs/mix.wav', dtype='float32')
window = np.hanning(512) ** 0.5
x = torch.stft(torch.from_numpy(mix), 512, 256, 512,
               torch.from_numpy(window.astype(np.float32)),
               return_complex=False).numpy()[None]  # (1, 257, T, 2)

conv_cache_np  = np.zeros([2, 1, 16, 16, 33], dtype='float32')
tra_cache_np   = np.zeros([2, 3, 1, 1, 16],   dtype='float32')
inter_cache_np = np.zeros([2, 1, 33, 16],      dtype='float32')
outputs = []

for i in range(x.shape[2]):
    enh_i, conv_cache_np, tra_cache_np, inter_cache_np = session.run(
        [],
        {'mix': x[..., i:i+1, :],
         'conv_cache':  conv_cache_np,
         'tra_cache':   tra_cache_np,
         'inter_cache': inter_cache_np}
    )
    outputs.append(enh_i)

outputs = np.concatenate(outputs, axis=2)   # (1, 257, T, 2)
enhanced = istft(outputs[...,0] + 1j*outputs[...,1],
                 n_fft=512, hop_length=256, win_length=512, window=window)
sf.write('stream/test_wavs/enh_onnx.wav', enhanced.squeeze(), 16000)
```

--------------------------------

### Build Minimal LADSPA Plugin

Source: https://github.com/xiaobin-rong/gtcrn/blob/main/ladspa/README.md

Builds the smallest, single-file LADSPA plugin with no external dependencies using Docker. This is the most robust option for distribution.

```bash
./build-minimal-docker.sh
```

```bash
./build.sh minimal
```

--------------------------------

### Convert Offline GTCRN to StreamGTCRN and Perform Inference

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Loads an offline GTCRN model, converts its weights to a streaming-compatible format using `convert_to_stream`, and then processes audio frame-by-frame using `StreamGTCRN` with explicit state caches.

```python
import torch
import soundfile as sf
from gtcrn import GTCRN
from stream.gtcrn_stream import StreamGTCRN
from stream.modules.convert import convert_to_stream

device = torch.device("cpu")

# Load offline model and transfer weights to stream model
model = GTCRN().eval()
model.load_state_dict(torch.load('stream/onnx_models/model_trained_on_dns3.tar',
                                  map_location=device))['model']
stream_model = StreamGTCRN().eval()
convert_to_stream(stream_model, model)

# Initialise state caches (all zeros = silence)
conv_cache  = torch.zeros(2, 1, 16, 16, 33)   # (2, B, C, (kT-1)*8, F_enc)
tra_cache   = torch.zeros(2, 3, 1, 1, 16)      # (2, 3, 1, B, C)
inter_cache = torch.zeros(2, 1, 33, 16)        # (2, 1, B*F_enc, C)

# Load audio and compute full STFT (F=257, T, 2)
mix, fs = sf.read('stream/test_wavs/mix.wav', dtype='float32')
window  = torch.hann_window(512).pow(0.5)
x = torch.stft(torch.from_numpy(mix), 512, 256, 512, window, return_complex=False)[None]
# x.shape == (1, 257, T, 2)

enhanced_frames = []
with torch.no_grad():
    for i in range(x.shape[2]):                  # iterate over time frames
        xi = x[:, :, i:i+1, :]                   # (1, 257, 1, 2)
        yi, conv_cache, tra_cache, inter_cache = stream_model(
            xi, conv_cache, tra_cache, inter_cache
        )
        enhanced_frames.append(yi)

y = torch.cat(enhanced_frames, dim=2)            # (1, 257, T, 2)
enh = torch.istft(y[0], 512, 256, 512, window).cpu().numpy()
sf.write('stream/test_wavs/enh_stream.wav', enh, 16000)
# Output equals offline model within numerical precision (RTF ≈ 0.07)
```

--------------------------------

### Analyze GTCRN Model Complexity with ptflops

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Calculates and prints the MACs and parameter count for the GTCRN model using the ptflops library. Assumes a specific input shape for analysis.

```python
import torch
from gtcrn import GTCRN
from ptflops import get_model_complexity_info

model = GTCRN().eval()
flops, params = get_model_complexity_info(
    model,
    (257, 63, 2),          # (F, T, 2) — one 1-second segment at 16 kHz
    as_strings=True,
    print_per_layer_stat=False,
    verbose=False
)
print(f"MACs : {flops}")    # ~33.0 MMac
print(f"Params: {params}")  # ~48.2 K
```

--------------------------------

### ERB Filterbank: Linear to ERB Compression

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Demonstrates the ERB filterbank's forward and inverse mapping. `bm()` compresses linear frequency features to ERB bands, reducing dimensionality, while `bs()` inverts this process.

```python
import torch
from gtcrn import ERB

erb = ERB(erb_subband_1=65, erb_subband_2=64, nfft=512, high_lim=8000, fs=16000)

# Simulate a feature map: (B=1, C=3, T=100, F=257)
feat = torch.randn(1, 3, 100, 257)

# Forward mapping: linear → ERB compressed (257 → 129)
feat_erb = erb.bm(feat)
print(feat_erb.shape)  # torch.Size([1, 3, 100, 129])

# Inverse mapping: ERB → linear (129 → 257)
feat_back = erb.bs(feat_erb)
print(feat_back.shape)  # torch.Size([1, 3, 100, 257])
```

--------------------------------

### PipeWire Configuration for GTCRN Filter

Source: https://github.com/xiaobin-rong/gtcrn/blob/main/ladspa/README.md

Configures PipeWire to use the GTCRN LADSPA plugin as a system-wide noise suppression filter for microphones. Ensure the .so file is in your LADSPA path.

```conf
context.modules = [
    {
        name = libpipewire-module-filter-chain
        args = {
            node.description = "Noise Canceling Microphone (GTCRN)"
            media.name = "Noise Canceling Microphone (GTCRN)"
            filter.graph = {
                nodes = [
                    {
                        type = ladspa
                        name = "gtcrn"
                        plugin = "libgtcrn_ladspa"
                        label = "gtcrn_mono"
                        control = {
                            # Strength: 0.0 (Original) to 1.0 (Fully Processed)
                            "Strength" = 1.0
                            "Model (0=Light 1=Full)" = 1
                        }
                    }
                ]
            }
            audio.channels = 1
            capture.props = {
                node.passive = true
            }
            playback.props = {
                media.class = "Audio/Source"
            }
        }
    }
]
```

--------------------------------

### Export GTCRN Model to ONNX

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Exports the trained GTCRN model to ONNX format for deployment. Includes model simplification and saving.

```python
stream_model.eval()
dummy_input  = torch.randn(1, 257, 1, 2)
conv_cache0  = torch.zeros(2, 1, 16, 16, 33)
tra_cache0   = torch.zeros(2, 3, 1, 1, 16)
inter_cache0 = torch.zeros(2, 1, 33, 16)

torch.onnx.export(
    stream_model,
    (dummy_input, conv_cache0, tra_cache0, inter_cache0),
    'stream/onnx_models/gtcrn.onnx',
    input_names  = ['mix', 'conv_cache', 'tra_cache', 'inter_cache'],
    output_names = ['enh', 'conv_cache_out', 'tra_cache_out', 'inter_cache_out'],
    opset_version=11,
)

onnx_model = onnx.load('stream/onnx_models/gtcrn.onnx')
model_simp, ok = simplify(onnx_model)
assert ok
onnx.save(model_simp, 'stream/onnx_models/gtcrn_simple.onnx')
```

--------------------------------

### SFE: Subband Feature Extraction

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Utilizes `nn.Unfold` for subband feature extraction. With `kernel_size=3`, it expands the channel dimension by incorporating neighboring frequency bins, providing local context without learnable parameters.

```python
import torch
from gtcrn import SFE

sfe = SFE(kernel_size=3, stride=1)

# Input: (B, C, T, F)
x = torch.randn(1, 3, 100, 129)
xs = sfe(x)
print(xs.shape)  # torch.Size([1, 9, 100, 129])
# Each output channel fuses a frequency bin with its two neighbours
```

--------------------------------

### GTCRN Offline Batch Inference

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Performs offline batch inference using a pre-trained GTCRN model. Requires 16 kHz mono float32 audio input and outputs denoised audio. The STFT and iSTFT steps are crucial for processing audio signals.

```python
import os
import torch
import soundfile as sf
from gtcrn import GTCRN

# Load pre-trained model (DNS3 checkpoint)
device = torch.device("cpu")
model = GTCRN().eval()
ckpt = torch.load(os.path.join('checkpoints', 'model_trained_on_dns3.tar'), map_location=device)
model.load_state_dict(ckpt['model'])

# Load noisy audio (must be 16 kHz mono float32)
mix, fs = sf.read(os.path.join('test_wavs', 'mix.wav'), dtype='float32')
assert fs == 16000, "GTCRN requires 16 kHz audio"

# STFT: window=sqrt(hann), n_fft=512, hop=256 → shape (F=257, T, 2)
window = torch.hann_window(512).pow(0.5)
spec_in = torch.stft(
    torch.from_numpy(mix), 512, 256, 512, window, return_complex=False
)  # (257, T, 2)

# Batch dimension required: (1, 257, T, 2)
with torch.no_grad():
    spec_enh = model(spec_in[None])[0]   # → (257, T, 2)

# iSTFT back to waveform
enh = torch.istft(spec_enh, 512, 256, 512, window)

# Save result
sf.write(os.path.join('test_wavs', 'enh.wav'), enh.cpu().numpy(), fs)
# Expected: enh.wav contains denoised speech, RTF << 1 on CPU
```

--------------------------------

### ERB Filterbank

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

The ERB (Equivalent Rectangular Bandwidth) filterbank module for frequency compression and decompression. `bm()` maps linear-frequency to ERB-compressed bands, and `bs()` inverts it.

```APIDOC
## ERB — Equivalent Rectangular Bandwidth Filterbank

`ERB(erb_subband_1, erb_subband_2)` is a fixed frequency compression module. `bm()` maps a linear-frequency feature map to ERB-compressed bands; `bs()` inverts it back to the full frequency resolution. It reduces the high-frequency dimension from 257 bins to 129.

### Methods
- **bm(feat)**: Maps linear-frequency features to ERB-compressed bands.
- **bs(feat_erb)**: Maps ERB-compressed features back to linear frequency.

### Parameters
#### `ERB` Constructor
- **erb_subband_1** (int) - Number of ERB subbands.
- **erb_subband_2** (int) - Number of remaining linear bands.
- **nfft** (int) - FFT size used.
- **high_lim** (int) - High frequency limit.
- **fs** (int) - Sampling frequency.

#### `bm` Method
- **feat** (torch.Tensor) - Input feature map of shape `(B, C, T, F)`.

#### `bs` Method
- **feat_erb** (torch.Tensor) - ERB-compressed feature map.

### Response
#### Success Response (`bm`)
- **feat_erb** (torch.Tensor) - ERB-compressed feature map of shape `(B, C, T, erb_subband_1 + erb_subband_2)`.

#### Success Response (`bs`)
- **feat_back** (torch.Tensor) - Decompressed feature map of shape `(B, C, T, nfft//2 + 1)`.
```

--------------------------------

### SFE Subband Feature Extraction

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Subband Feature Extraction module that expands the frequency-channel dimension by gathering neighboring frequency bins.

```APIDOC
## SFE — Subband Feature Extraction

`SFE(kernel_size, stride)` expands the frequency-channel dimension by gathering each frequency bin together with its `kernel_size - 1` nearest neighbors using `nn.Unfold`. With `kernel_size=3`, a 3-channel input becomes 9-channel output.

### Method
```python
sfe(x)
```

### Parameters
#### `SFE` Constructor
- **kernel_size** (int) - The size of the kernel for unfolding frequency bins.
- **stride** (int) - The stride for unfolding.

#### Input `x`
- **x** (torch.Tensor) - Input tensor of shape `(B, C, T, F)`.

### Response
#### Success Response
- **xs** (torch.Tensor) - Output tensor with expanded channels, shape `(B, C * kernel_size, T, F)`.
```

--------------------------------

### TRA: Temporal Recurrent Attention

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Applies channel-wise temporal attention using a GRU over the mean-squared energy of frequency bands. It re-weights feature maps with a learned sigmoid gate to focus on temporally active speech regions.

```python
import torch
from gtcrn import TRA

tra = TRA(channels=8)

# Input: (B, C, T, F)
x = torch.randn(2, 8, 50, 33)
out = tra(x)
print(out.shape)  # torch.Size([2, 8, 50, 33])
# Channels are re-weighted frame-by-frame by the GRU attention gate
```

--------------------------------

### GTCRN Batch Inference

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Performs offline batch inference using the GTCRN model. It takes a batched complex STFT tensor and returns the enhanced spectrum.

```APIDOC
## GTCRN — Offline Batch Inference

`GTCRN.forward(spec)` accepts a batched complex STFT tensor of shape `(B, F, T, 2)` (real and imaginary stacked in the last dimension) and returns an enhanced spectrum of the same shape. Internally it applies ERB compression, subband feature extraction, encoder convolutions, two DPGRNN passes, decoder, and a Complex Ratio Mask.

### Method
```python
model(spec_in[None])
```

### Parameters
#### Request Body
- **spec** (torch.Tensor) - A complex STFT tensor of shape `(B, F, T, 2)`.

### Request Example
```python
# Assuming model and spec_in are already defined
spec_enh = model(spec_in[None])[0]
```

### Response
#### Success Response (200)
- **enhanced_spectrum** (torch.Tensor) - The enhanced spectrum of shape `(F, T, 2)`.
```

--------------------------------

### TRA Temporal Recurrent Attention

Source: https://context7.com/xiaobin-rong/gtcrn/llms.txt

Temporal Recurrent Attention module that applies a channel-wise temporal attention gate using a GRU over the mean-squared energy of each frequency band.

```APIDOC
## TRA — Temporal Recurrent Attention

`TRA(channels)` applies a channel-wise temporal attention gate using a GRU over the mean-squared energy of each frequency band. It re-weights feature maps by a learned time-varying sigmoid gate.

### Method
```python
tra(x)
```

### Parameters
#### `TRA` Constructor
- **channels** (int) - The number of input channels.

#### Input `x`
- **x** (torch.Tensor) - Input tensor of shape `(B, C, T, F)`.

### Response
#### Success Response
- **out** (torch.Tensor) - Output tensor with re-weighted channels, same shape as input `(B, C, T, F)`.
```
Tokens:
Raw