sgmse: diffusion models for speech enhancement and dereverberation

A reference implementation, not a product

sgmse is the official PyTorch implementation of a body of work on score based generative models, and it is named for the acronym of the title, Score-based Generative Models, for speech enhancement and dereverberation. The project description says it plainly: diffusion models for speech enhancement and dereverberation, and the README is a paper reference page with installation notes attached.

The paper list is the substance of the repository, and it covers five publications across three venues. There is the Interspeech 2022 paper on speech enhancement with score based generative models in the complex STFT domain, by Welker, Richter, and Gerkmann. There is the journal paper on speech enhancement and dereverberation with diffusion based generative models, published in the IEEE/ACM Transactions on Audio, Speech, and Language Processing, volume 31, pages 2351 to 2364, in 2023, with Richter, Welker, Lemercier, Lay, and Gerkmann as the authors. There is the ICASSP 2023 paper analysing diffusion based generative approaches against discriminative approaches for speech restoration, by Lemercier, Richter, Welker, and Gerkmann. There is the EARS dataset paper from Interspeech 2024, which introduced the anechoic fullband speech dataset that several of the later checkpoints are trained on. And there is the ICASSP 2025 paper on investigating training objectives for generative speech enhancement, by Richter, de Oliveira, and Gerkmann.

The authors of the repository match the paper authors, with Simon Welker, Julius Richter, and Timo Gerkmann listed as the package authors.

There is also a follow-up section, which points at two related projects with their own code. There is StoRM, a diffusion based stochastic regeneration model for speech enhancement and dereverberation, and there is the work on reducing the prior mismatch of stochastic differential equations for diffusion based speech enhancement, with a repository named for it.

The project is MIT licensed, the last push was on 2026-05-12, and there are no GitHub releases, which is what a research repository that tracks a paper rather than a product looks like.

Why the complex STFT domain matters here

The diagram at the top of the README is the one worth reading, because it is the whole method in one picture.

In the forward process, noise is gradually added to the clean speech spectrogram, which the diagram labels x0. In the reverse process, the model learns to generate clean speech in an iterative fashion, starting from the corrupted signal, which is labelled xT.

That is a diffusion model, and the reason the paper works in the complex STFT domain is that the time domain signal is a poor target for a diffusion process. A speech waveform has a phase component that is essentially random, so a generative model in the time domain spends its capacity modelling noise that carries no information. Working in the complex short time Fourier transform domain, where the signal is represented as magnitude and phase, changes what the model is asked to predict, and the first of the five papers is specifically about doing the enhancement that way.

The iterative reverse process is also the practical constraint of the whole approach. Generating clean speech means running the model many times, once per step of the reverse diffusion, starting from pure noise. That is why the sampler settings appear in the checkpoint instructions, because the number of steps is a quality and speed trade-off, and the repository documents specific values for one of the models.

It also explains the shape of the checkpoints. A checkpoint is not a single file for a single task, it is a set of files for a set of datasets, because the paper series trained the same architecture on different corpora, and the sampler settings that suit one do not necessarily suit another.

Installing, and the experiment tracking you may want to turn off

The installation is four steps, and one of them has a caveat.

You create a new virtual environment with Python 3.11. The README says plainly that other Python versions have not been tested, but that they may work, which is a candid way of saying the combination is pinned to one version in practice.

You install the package dependencies with `pip install -r requirements.txt`, and the instruction is to let pip resolve the dependencies for you. If you hit problems, there is a second file, requirements_version.txt, holding the exact versions the authors used, which is the file to reach for when the resolved set does not behave.

Then there is the logging decision, and it is the one that catches people out. Weights and Biases logging is on by default, which means you need an account on wandb.ai and you need to run `wandb login` before running the code. If you would rather not do that, you pass the option `--nolog` to train.py, and your logs are then stored as local CSVLogger logs in the lightning_logs directory.

That default is worth understanding. The repository is a research release, and experiment tracking is part of the workflow the authors expect, so a run without a WandB account is treated as the unusual case rather than the normal one. The cost of turning it off is that you lose the hosted experiment comparison, and the benefit is that nothing about your training run leaves your machine.

The package itself is installed with a standard setuptools build, the project name is sgmse at version 0.1.0, and the metadata declares the same three authors as the papers.

The checkpoints, and why the sampler settings matter

There are a lot of pretrained checkpoints, and the README organises them by the task and the training corpus, which is the useful way to think about them.

For speech enhancement there are checkpoints for models trained on the VoiceBank-DEMAND and the WSJ0-CHiME3 datasets, and the download command for each is a gdown invocation with a file identifier. For dereverberation there is a checkpoint trained on the authors own WSJ0-REVERB dataset, and next to it is the note that this checkpoint works better with the sampler settings `--N 50 --snr 0.33`.

That note is the most practically important line in the checkpoint section. The number of reverse diffusion steps and the noise level used in the sampler are the two knobs that decide how much noise the model starts from and how many times it denoises, and a checkpoint that is trained for a particular operating point will sound worse if you sample it somewhere else. The authors tell you the operating point for the dereverberation model, and the other models are documented in the journal paper rather than in the README.

Then there is the 48 kHz generation. For the models from the EARS paper, there is a speech enhancement checkpoint trained on the EARS-WHAM dataset and a dereverberation checkpoint trained on EARS-Reverb, with gdown identifiers for both.

The training objective study has its own set of checkpoints, named M1 through M8, downloaded from the university web server with wget rather than gdown, with M4 handled by a separate repository for a different model. And there is a pretrained checkpoint for the Schrödinger bridge model, trained on the EARS-WHAM and VB-DMD datasets, and the README lists the exact training configuration it was produced with, the batch size, the number of devices, the backbone name, the loss type, the stochastic differential equation variant, and the sample rate.

Listing the training configuration next to a checkpoint is a small courtesy that makes a lot of difference, because it is the only way to know whether the checkpoint is the thing you think it is.

The repository as a working directory

The top level of the repository is small, and it reads like a research project rather than a distributed package.

There is the sgmse package directory, a preprocessing directory, a train.py which is the training entry point, an enhancement.py which is the inference script, a calc_metrics.py for evaluation, and a logs directory. There is the requirements file and the pinned version file, the project manifest, and a diffusion_process.png which is the figure from the top of the README.

The dependency list is where the character of the project shows. It has the PyTorch family, with torch, torchaudio, torchvision, torch-ema for the exponential moving average of the weights, torchinfo for the model summaries, and torchsde for the stochastic differential equation samplers. It has pytorch-lightning for the training loop, which is why the no-logging path produces lightning_logs directory output. It has the audio stack, with librosa, soundfile through the general scientific stack, pystoi for the speech to intelligibility measure, pesq and torch-pesq for the perceptual quality measure, and pyroomacoustics for the room simulation used in the dereverberation data. It has sdeint for the score integral solver. It has wandb for the experiment tracking, and pyarrow and h5py for the data formats, and ipympl and seaborn for the notebook and the plots.

One detail in the manifest is easy to miss and is a good sign about the engineering. The package data includes the C, C++, and CUDA source files of the score network's custom operations, with the cpp, cu, and cuh extensions. Those kernels are compiled rather than shipped as binaries, and the ninja dependency in the list is the build tool for them. So the score network has hand written operations for the backward SDE solver, and a user has to have a CUDA toolchain to build them.

That is the single biggest practical hurdle for someone who wants to run the pretrained models, and it is the reason the interactive demo in a separate repository, which is described as a Jupyter notebook, exists.

Editorial conclusion

sgmse fits a speech researcher who wants the reference implementation of these papers rather than an approximation, who needs the pretrained checkpoints for a dataset they are evaluating on, and who is prepared to work in the complex STFT domain rather than the time domain. It does not fit a production speech enhancement service, since the repository is a research release with no inference server, no packaging for deployment, and a sampler that iterates. Verify first that you have Python 3.11 and a CUDA capable machine, since the custom kernels for the score network are compiled rather than shipped as binaries, that you read the sampler settings for the checkpoint you are using, because the dereverberation model is documented as working better with particular values, and that you have decided what to do about the experiment tracking, which is enabled by default.

Frequently asked questions

Official sources

1. Issues
2. License: MIT
3. README
4. sp-uhh/sgmse on GitHub