UniPASE: A Generative Model for Universal Speech Enhancement with High Fidelity and Low Hallucinations

🎉 This is the official implementation of UniPASE, which has been accepted by IEEE TASLP.

Related Work:

• PASE & StuPASE
• GAP-URGENet

Pretrained Checkpoints

Five checkpoints are provided:

• DeWavLM-Omni.pt
• Adapter.pt
• Vocoder_DWO-L1.pt
• Vocoder_WavLM-L24.pt
• PostNet.pt

Note: Vocoder_WavLM-L24.pt is not a part of UniPASE, but is only used to reconstruct waveforms from WavLM L24 representations, enabling evaluation of DeWavLM-Omni's performance.

Inference

To run inference on audio files, use:

python -m inference.inference -I <input_dir> -O <output_dir> [options]

For long-form audio inputs (e.g., > 20s), use:

python -m inference.inference_long -I <input_dir> -O <output_dir> [options]

Argument	Requirement / Default	Description
-I (--input_dir)	required	Path to the input directory containing audio files.
-O (--output_dir)	required	Path to the output directory where enhanced files will be saved.
-D (--device)	default: cuda:0	Torch device to run inference on, e.g., cuda:0, cuda:1, or cpu.
-E (--extension)	default: .wav	Audio file extension to process.
--sr_out	default: None	Output sampling rate (default: same as input)
--enable_plc	default: True	Whether to perform packet loss concealment (PLC)

Audio examples can be found in ./audio.

Training

Step 1: Finetune WavLM

• training script: train/train_dewavlm.py
• training configuration: configs/cfg_train_dewavlm.yaml

  python -m train.train_dewavlm -C configs/cfg_train_dewavlm.yaml -D 0,1,2,3

• inference script: inference/infer_dewavlm.py

  python -m inference.infer_dewavlm -C configs/cfg_infer.yaml -D 0

Note: The training script uses Vocoder_WavLM_L24.pt to reconstruct waveforms for validation purposes. Users can directly use our provided checkpoint for validation without retraining.

Purpose: Adapt WavLM as an USE expert for core enhancement.

Step 2: Train a Vocoder

• training script: train/train_vocoder.py
• training configuration: configs/cfg_train_vocoder.yaml
• inference script: inference/infer_vocoder.py

Purpose: Pretrain a vocoder based on DeWavLM-Omni L1 representations for waveform reconstruction.

Step 3: Train an Adapter

• training script: train/train_adapter.py
• training configuration: configs/cfg_train_adapter.yaml
• inference script: inference/infer_adapter.py

Purpose: Train the Adapter module for acoustic enhancement. This module transforms noisy DeWavLM-Omni L1 representations into enhanced L1 outputs guided by the L24 representations.

Step 4: Train a PostNet

• training script: train/train_postnet.py
• training configuration: configs/cfg_train_postnet.yaml
• inference script: inference/infer_postenet.py

Purpose: Train the PostNet module for bandwidth extension (BWE).

Creating Checkpoints

Once all training steps are completed, the corresponding checkpoints can be prepared for inference:

• utils/create_ckpt_wavlm.py is used to create a DeWavLM-Omni checkpoint.
• utils/create_ckpt.py is used to create other checkpoints.

Citation

If you find this work useful, please cite our paper:

@ARTICLE{UniPASE,
  author={Rong, Xiaobin and Wang, Zheng and Wang, Yushi and Gao, Jun and Lu, Jing},
  journal={IEEE Transactions on Audio, Speech and Language Processing}, 
  title={{UniPASE: A Generative Model for Universal Speech Enhancement With High Fidelity and Low Hallucinations}}, 
  year={2026},
  volume={34},
  number={},
  pages={3901-3915},
  keywords={Modeling;Speech;Training;Speech enhancement;Vocoders;Cleaning;Conferences;Noise reduction;Distortion;Measurement;Universal speech enhancement;generative model;high fidelity;low hallucinations},
  doi={10.1109/TASLPRO.2026.3717231}}

Contact

Xiaobin Rong: xiaobin.rong@smail.nju.edu.cn