chola-compressor

Chola-Compressor: LoRaVoiceLink Speech Restoration Model

A compact spectrogram U-Net (1.93M parameters) that restores speech quality lost to Codec2 compression, trained via knowledge distillation for use in off-grid LoRa voice communication links. Designed to run in real time on edge hardware (Jetson-class), not in the cloud.

Problem

LoRa radio has enough bandwidth for Codec2 at very low bitrates (this project uses the 1200bps mode), which makes off-grid voice communication possible but introduces heavy compression artifacts. This model sits after Codec2 decoding on the receiving end and restores some of that lost quality.

Architecture

3-level U-Net operating on STFT magnitude spectrograms (n_fft=512, hop=128, 16kHz):

• Encoder: 3 conv blocks (32→64→128 channels) with max-pooling
• Bottleneck: 256 channels
• Decoder: 3 conv blocks with transposed-conv upsampling and skip connections
• Output: predicts a [0,1] mask multiplied against the input magnitude (not raw magnitude directly — more stable to train)
• Phase is not predicted; the model reuses the degraded audio's own phase for reconstruction (see Limitations)

Training

• Data: 5,000 utterances from VCTK (33 speakers, mic2 only, speaker-disjoint train/val/test split — no speaker overlap across splits), streamed from the jspaulsen/vctk mirror
• Degradation: real Codec2 encode/decode roundtrip (1200bps), not a synthetic approximation
• Distillation: trained with Meta's Denoiser (dns64) as an auxiliary teacher signal alongside the real clean-speech target. Ablation showed the teacher term contributed no measurable benefit over training on ground truth alone (see Results) — this checkpoint (models_no_teacher) was trained with teacher_weight=0, i.e. supervised directly against real clean speech.
• Loss: L1 on log-magnitude spectrograms (log1p), chosen after finding raw-magnitude L1 over-weights loud regions and under-penalizes quiet, perceptually important detail
• 70 epochs, Adam, lr=1e-4, batch size capped by 8GB VRAM

Results

Evaluated on a held-out, speaker-disjoint test split (403 utterances). All three metrics computed against the true clean reference.

Real, verified gains: +0.12 STOI (intelligibility) and +1.5dB SI-SDR over doing nothing. Both metrics are dominated by energy/envelope accuracy, where this model clearly helps.

Known limitation — PESQ: PESQ is highly sensitive to phase accuracy, and this model only predicts a magnitude mask, reconstructing with the degraded audio's original (uncontrolled) phase. That's the most likely explanation for PESQ landing slightly below the unprocessed baseline despite STOI/SI-SDR improving substantially — three independent loss-function reformulations (distillation weight, log vs. linear magnitude) all left PESQ in the same 1.47–1.48 range, which is consistent with a phase-reconstruction ceiling rather than a loss-tuning problem. A phase-aware architecture (predicting complex spectrograms or a phase correction term) would likely be needed to close this gap; that's a known next step, not yet implemented in this checkpoint.

Intended use

Research and portfolio demonstration of magnitude-domain speech restoration via knowledge distillation for bandwidth-constrained voice links. Not validated for safety-critical or emergency-communication deployment.

How to use

import torch
import numpy as np
import librosa
import soundfile as sf

N_FFT, HOP_LENGTH, SAMPLE_RATE = 512, 128, 16000

class SpectrogramUNet(torch.nn.Module):
    # ... see model.py in the project repo for the full class definition
    pass

def restore(degraded_wav_path, model, device="cpu"):
    audio, sr = sf.read(degraded_wav_path, dtype="float32")
    if sr != SAMPLE_RATE:
        audio = librosa.resample(audio, orig_sr=sr, target_sr=SAMPLE_RATE)
    stft = librosa.stft(audio, n_fft=N_FFT, hop_length=HOP_LENGTH)
    mag, phase = np.abs(stft), np.angle(stft)

    x = torch.from_numpy(mag).float().unsqueeze(0).unsqueeze(0).to(device)
    with torch.no_grad():
        pred_mag = model(x).squeeze().cpu().numpy()

    restored_stft = pred_mag * np.exp(1j * phase)  # reuses degraded audio's phase
    return librosa.istft(restored_stft, hop_length=HOP_LENGTH)

model = SpectrogramUNet(base_channels=32)
model.load_state_dict(torch.load("best.pt", map_location="cpu"))
model.eval()

restored = restore("degraded.wav", model)
sf.write("restored.wav", restored, SAMPLE_RATE)

Full training/eval/live-test code: see the project repository.

Citation

If you use this model, please cite the LoRaVoiceLink project (link to source repo).