update

examples/clean_unet_aishell/step_2_train_model.py

@@ -28,7 +28,7 @@ from torch.utils.data.dataloader import DataLoader
 from tqdm import tqdm
 
 from toolbox.torch.utils.data.dataset.denoise_excel_dataset import DenoiseExcelDataset
-from toolbox.torchaudio.models.clean_unet.configuration_clean_unet import CleanUnetConfig
+from toolbox.torchaudio.models.clean_unet.configuration_clean_unet import CleanUNetConfig
 from toolbox.torchaudio.models.clean_unet.modeling_clean_unet import CleanUNetPretrainedModel
 from toolbox.torchaudio.models.clean_unet.training import LinearWarmupCosineDecay
 from toolbox.torchaudio.models.clean_unet.loss import MultiResolutionSTFTLoss
@@ -112,7 +112,7 @@ collate_fn = CollateFunction()
 def main():
     args = get_args()
 
-    config = CleanUnetConfig.from_pretrained(
+    config = CleanUNetConfig.from_pretrained(
         pretrained_model_name_or_path=args.config_file,
     )
 
@@ -186,7 +186,7 @@ def main():
         model_pt = serialization_dir / f"epoch-{last_epoch}/model.pt"
         optimizer_pth = serialization_dir / f"epoch-{last_epoch}/optimizer.pth"
 
-        logger.info(f"load state dict for generator.")
+        logger.info(f"load state dict for model.")
         with open(model_pt.as_posix(), "rb") as f:
             state_dict = torch.load(f, map_location="cpu", weights_only=True)
         model.load_state_dict(state_dict, strict=True)
@@ -317,7 +317,7 @@ def main():
                 enhanced_audios = model.forward(noisy_audios)
                 enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
 
-                ae_loss = ae_loss_fn(enhanced_audios, enhanced_audios)
+                ae_loss = ae_loss_fn(enhanced_audios, clean_audios)
                 sc_loss, mag_loss = mr_stft_loss_fn(enhanced_audios, clean_audios)
 
                 loss = ae_loss + sc_loss + mag_loss

examples/nx_clean_unet/run.sh

@@ -0,0 +1,166 @@
+#!/usr/bin/env bash
+
+: <<'END'
+
+
+sh run.sh --stage 2 --stop_stage 2 --system_version windows --file_folder_name file_dir --final_model_name mpnet-aishell-20250224 \
+--noise_dir "E:/Users/tianx/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "E:/programmer/asr_datasets/aishell/data_aishell/wav/train"
+
+
+sh run.sh --stage 3 --stop_stage 3 --system_version centos --file_folder_name file_dir --final_model_name mpnet-aishell-20250224 \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
+
+sh run.sh --stage 5 --stop_stage 5 --system_version centos --file_folder_name file_dir --final_model_name mpnet-aishell-20250224 \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
+
+
+sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name mpnet-nx-speech-20250224 \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech" \
+--max_epochs 1
+
+
+END
+
+
+# params
+system_version="windows";
+verbose=true;
+stage=0 # start from 0 if you need to start from data preparation
+stop_stage=9
+
+work_dir="$(pwd)"
+file_folder_name=file_folder_name
+final_model_name=final_model_name
+config_file="yaml/config.yaml"
+limit=10
+
+noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
+speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
+
+nohup_name=nohup.out
+
+# model params
+batch_size=64
+max_epochs=200
+save_top_k=10
+patience=5
+
+
+# parse options
+while true; do
+  [ -z "${1:-}" ] && break;  # break if there are no arguments
+  case "$1" in
+    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
+      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
+      old_value="(eval echo \\$$name)";
+      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
+        was_bool=true;
+      else
+        was_bool=false;
+      fi
+
+      # Set the variable to the right value-- the escaped quotes make it work if
+      # the option had spaces, like --cmd "queue.pl -sync y"
+      eval "${name}=\"$2\"";
+
+      # Check that Boolean-valued arguments are really Boolean.
+      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
+        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
+        exit 1;
+      fi
+      shift 2;
+      ;;
+
+    *) break;
+  esac
+done
+
+file_dir="${work_dir}/${file_folder_name}"
+final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
+evaluation_audio_dir="${file_dir}/evaluation_audio"
+
+dataset="${file_dir}/dataset.xlsx"
+train_dataset="${file_dir}/train.xlsx"
+valid_dataset="${file_dir}/valid.xlsx"
+
+$verbose && echo "system_version: ${system_version}"
+$verbose && echo "file_folder_name: ${file_folder_name}"
+
+if [ $system_version == "windows" ]; then
+  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
+elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
+  #source /data/local/bin/nx_denoise/bin/activate
+  alias python3='/data/local/bin/nx_denoise/bin/python3'
+fi
+
+
+if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
+  $verbose && echo "stage 1: prepare data"
+  cd "${work_dir}" || exit 1
+  python3 step_1_prepare_data.py \
+  --file_dir "${file_dir}" \
+  --noise_dir "${noise_dir}" \
+  --speech_dir "${speech_dir}" \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+
+fi
+
+
+if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
+  $verbose && echo "stage 2: train model"
+  cd "${work_dir}" || exit 1
+  python3 step_2_train_model.py \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --serialization_dir "${file_dir}" \
+  --config_file "${config_file}" \
+
+fi
+
+
+if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
+  $verbose && echo "stage 3: test model"
+  cd "${work_dir}" || exit 1
+  python3 step_3_evaluation.py \
+  --valid_dataset "${valid_dataset}" \
+  --model_dir "${file_dir}/best" \
+  --evaluation_audio_dir "${evaluation_audio_dir}" \
+  --limit "${limit}" \
+
+fi
+
+
+if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
+  $verbose && echo "stage 4: collect files"
+  cd "${work_dir}" || exit 1
+
+  mkdir -p ${final_model_dir}
+
+  cp "${file_dir}/best"/* "${final_model_dir}"
+  cp -r "${file_dir}/evaluation_audio" "${final_model_dir}"
+
+  cd "${final_model_dir}/.." || exit 1;
+
+  if [ -e "${final_model_name}.zip" ]; then
+    rm -rf "${final_model_name}_backup.zip"
+    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
+  fi
+
+  zip -r "${final_model_name}.zip" "${final_model_name}"
+  rm -rf "${final_model_name}"
+
+fi
+
+
+if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
+  $verbose && echo "stage 5: clear file_dir"
+  cd "${work_dir}" || exit 1
+
+  rm -rf "${file_dir}";
+
+fi

examples/nx_clean_unet/step_1_prepare_data.py

@@ -0,0 +1,201 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import os
+from pathlib import Path
+import random
+import sys
+import shutil
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import pandas as pd
+from scipy.io import wavfile
+from tqdm import tqdm
+import librosa
+
+from project_settings import project_path
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--file_dir", default="./", type=str)
+
+    parser.add_argument(
+        "--noise_dir",
+        default=r"E:\Users\tianx\HuggingDatasets\nx_noise\data\noise",
+        type=str
+    )
+    parser.add_argument(
+        "--speech_dir",
+        default=r"E:\programmer\asr_datasets\aishell\data_aishell\wav\train",
+        type=str
+    )
+
+    parser.add_argument("--train_dataset", default="train.xlsx", type=str)
+    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
+
+    parser.add_argument("--duration", default=2.0, type=float)
+    parser.add_argument("--min_snr_db", default=-10, type=float)
+    parser.add_argument("--max_snr_db", default=20, type=float)
+
+    parser.add_argument("--target_sample_rate", default=8000, type=int)
+
+    parser.add_argument("--max_count", default=10000, type=int)
+
+    args = parser.parse_args()
+    return args
+
+
+def filename_generator(data_dir: str):
+    data_dir = Path(data_dir)
+    for filename in data_dir.glob("**/*.wav"):
+        yield filename.as_posix()
+
+
+def target_second_signal_generator(data_dir: str, duration: int = 2, sample_rate: int = 8000):
+    data_dir = Path(data_dir)
+    for filename in data_dir.glob("**/*.wav"):
+        signal, _ = librosa.load(filename.as_posix(), sr=sample_rate)
+        raw_duration = librosa.get_duration(y=signal, sr=sample_rate)
+
+        if raw_duration < duration:
+            # print(f"duration less than {duration} s. skip filename: {filename.as_posix()}")
+            continue
+        if signal.ndim != 1:
+            raise AssertionError(f"expected ndim 1, instead of {signal.ndim}")
+
+        signal_length = len(signal)
+        win_size = int(duration * sample_rate)
+        for begin in range(0, signal_length - win_size, win_size):
+            row = {
+                "filename": filename.as_posix(),
+                "raw_duration": round(raw_duration, 4),
+                "offset": round(begin / sample_rate, 4),
+                "duration": round(duration, 4),
+            }
+            yield row
+
+
+def get_dataset(args):
+    file_dir = Path(args.file_dir)
+    file_dir.mkdir(exist_ok=True)
+
+    noise_dir = Path(args.noise_dir)
+    speech_dir = Path(args.speech_dir)
+
+    noise_generator = target_second_signal_generator(
+        noise_dir.as_posix(),
+        duration=args.duration,
+        sample_rate=args.target_sample_rate
+    )
+    speech_generator = target_second_signal_generator(
+        speech_dir.as_posix(),
+        duration=args.duration,
+        sample_rate=args.target_sample_rate
+    )
+
+    dataset = list()
+
+    count = 0
+    process_bar = tqdm(desc="build dataset excel")
+    for noise, speech in zip(noise_generator, speech_generator):
+        if count >= args.max_count:
+            break
+
+        noise_filename = noise["filename"]
+        noise_raw_duration = noise["raw_duration"]
+        noise_offset = noise["offset"]
+        noise_duration = noise["duration"]
+
+        speech_filename = speech["filename"]
+        speech_raw_duration = speech["raw_duration"]
+        speech_offset = speech["offset"]
+        speech_duration = speech["duration"]
+
+        random1 = random.random()
+        random2 = random.random()
+
+        row = {
+            "noise_filename": noise_filename,
+            "noise_raw_duration": noise_raw_duration,
+            "noise_offset": noise_offset,
+            "noise_duration": noise_duration,
+
+            "speech_filename": speech_filename,
+            "speech_raw_duration": speech_raw_duration,
+            "speech_offset": speech_offset,
+            "speech_duration": speech_duration,
+
+            "snr_db": random.uniform(args.min_snr_db, args.max_snr_db),
+
+            "random1": random1,
+            "random2": random2,
+            "flag": "TRAIN" if random2 < 0.8 else "TEST",
+        }
+        dataset.append(row)
+        count += 1
+        duration_seconds = count * args.duration
+        duration_hours = duration_seconds / 3600
+
+        process_bar.update(n=1)
+        process_bar.set_postfix({
+            # "duration_seconds": round(duration_seconds, 4),
+            "duration_hours": round(duration_hours, 4),
+
+        })
+
+    dataset = pd.DataFrame(dataset)
+    dataset = dataset.sort_values(by=["random1"], ascending=False)
+    dataset.to_excel(
+        file_dir / "dataset.xlsx",
+        index=False,
+    )
+    return
+
+
+
+def split_dataset(args):
+    """分割训练集, 测试集"""
+    file_dir = Path(args.file_dir)
+    file_dir.mkdir(exist_ok=True)
+
+    df = pd.read_excel(file_dir / "dataset.xlsx")
+
+    train = list()
+    test = list()
+
+    for i, row in df.iterrows():
+        flag = row["flag"]
+        if flag == "TRAIN":
+            train.append(row)
+        else:
+            test.append(row)
+
+    train = pd.DataFrame(train)
+    train.to_excel(
+        args.train_dataset,
+        index=False,
+        # encoding="utf_8_sig"
+    )
+    test = pd.DataFrame(test)
+    test.to_excel(
+        args.valid_dataset,
+        index=False,
+        # encoding="utf_8_sig"
+    )
+
+    return
+
+
+def main():
+    args = get_args()
+
+    get_dataset(args)
+    split_dataset(args)
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/nx_clean_unet/step_2_train_model.py

@@ -0,0 +1,438 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/yxlu-0102/MP-SENet/blob/main/train.py
+"""
+import argparse
+import json
+import logging
+from logging.handlers import TimedRotatingFileHandler
+import os
+import platform
+from pathlib import Path
+import random
+import sys
+import shutil
+from typing import List
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+from torch.utils.data.dataloader import DataLoader
+from tqdm import tqdm
+
+from toolbox.torch.utils.data.dataset.denoise_excel_dataset import DenoiseExcelDataset
+from toolbox.torchaudio.models.nx_clean_unet.configuration_nx_clean_unet import NXCleanUNetConfig
+from toolbox.torchaudio.models.nx_clean_unet.discriminator import MetricDiscriminator, MetricDiscriminatorPretrainedModel
+from toolbox.torchaudio.models.nx_clean_unet.modeling_nx_clean_unet import NXCleanUNet, NXCleanUNetPretrainedModel
+from toolbox.torchaudio.models.nx_clean_unet.metrics import run_batch_pesq, run_pesq_score
+from toolbox.torchaudio.models.nx_clean_unet.utils import mag_pha_stft, mag_pha_istft
+from toolbox.torchaudio.models.nx_clean_unet.loss import phase_losses
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--train_dataset", default="train.xlsx", type=str)
+    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
+
+    parser.add_argument("--max_epochs", default=100, type=int)
+
+    parser.add_argument("--num_serialized_models_to_keep", default=10, type=int)
+    parser.add_argument("--patience", default=5, type=int)
+    parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
+
+    parser.add_argument("--config_file", default="config.yaml", type=str)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config(file_dir: str):
+    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
+
+    logging.basicConfig(format=fmt,
+                        datefmt="%m/%d/%Y %H:%M:%S",
+                        level=logging.INFO)
+    file_handler = TimedRotatingFileHandler(
+        filename=os.path.join(file_dir, "main.log"),
+        encoding="utf-8",
+        when="D",
+        interval=1,
+        backupCount=7
+    )
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(logging.Formatter(fmt))
+    logger = logging.getLogger(__name__)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+class CollateFunction(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, batch: List[dict]):
+        clean_audios = list()
+        noisy_audios = list()
+
+        for sample in batch:
+            # noise_wave: torch.Tensor = sample["noise_wave"]
+            clean_audio: torch.Tensor = sample["speech_wave"]
+            noisy_audio: torch.Tensor = sample["mix_wave"]
+            # snr_db: float = sample["snr_db"]
+
+            clean_audios.append(clean_audio)
+            noisy_audios.append(noisy_audio)
+
+        clean_audios = torch.stack(clean_audios)
+        noisy_audios = torch.stack(noisy_audios)
+
+        # assert
+        if torch.any(torch.isnan(clean_audios)) or torch.any(torch.isinf(clean_audios)):
+            raise AssertionError("nan or inf in clean_audios")
+        if torch.any(torch.isnan(noisy_audios)) or torch.any(torch.isinf(noisy_audios)):
+            raise AssertionError("nan or inf in noisy_audios")
+        return clean_audios, noisy_audios
+
+
+collate_fn = CollateFunction()
+
+
+def main():
+    args = get_args()
+
+    config = NXCleanUNetConfig.from_pretrained(
+        pretrained_model_name_or_path=args.config_file,
+    )
+
+    serialization_dir = Path(args.serialization_dir)
+    serialization_dir.mkdir(parents=True, exist_ok=True)
+
+    logger = logging_config(serialization_dir)
+
+    random.seed(config.seed)
+    np.random.seed(config.seed)
+    torch.manual_seed(config.seed)
+    logger.info(f"set seed: {config.seed}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    logger.info(f"GPU available count: {n_gpu}; device: {device}")
+
+    # datasets
+    train_dataset = DenoiseExcelDataset(
+        excel_file=args.train_dataset,
+        expected_sample_rate=8000,
+        max_wave_value=32768.0,
+    )
+    valid_dataset = DenoiseExcelDataset(
+        excel_file=args.valid_dataset,
+        expected_sample_rate=8000,
+        max_wave_value=32768.0,
+    )
+    train_data_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=config.batch_size,
+        shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        # prefetch_factor=64,
+    )
+    valid_data_loader = DataLoader(
+        dataset=valid_dataset,
+        batch_size=config.batch_size,
+        shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        # prefetch_factor=64,
+    )
+
+    # models
+    logger.info(f"prepare models. config_file: {args.config_file}")
+    generator = NXCleanUNetPretrainedModel(config).to(device)
+    discriminator = MetricDiscriminatorPretrainedModel(config).to(device)
+
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler")
+    num_params = 0
+    for p in generator.parameters():
+        num_params += p.numel()
+    logger.info("total parameters (generator): {:.3f}M".format(num_params/1e6))
+
+    optim_g = torch.optim.AdamW(generator.parameters(), config.learning_rate, betas=[config.adam_b1, config.adam_b2])
+    optim_d = torch.optim.AdamW(discriminator.parameters(), config.learning_rate, betas=[config.adam_b1, config.adam_b2])
+
+    # resume training
+    last_epoch = -1
+    for epoch_i in serialization_dir.glob("epoch-*"):
+        epoch_i = Path(epoch_i)
+        epoch_idx = epoch_i.stem.split("-")[1]
+        epoch_idx = int(epoch_idx)
+        if epoch_idx > last_epoch:
+            last_epoch = epoch_idx
+
+    if last_epoch != -1:
+        logger.info(f"resume from epoch-{last_epoch}.")
+        generator_pt = serialization_dir / f"epoch-{last_epoch}/generator.pt"
+        discriminator_pt = serialization_dir / f"epoch-{last_epoch}/discriminator.pt"
+        optim_g_pth = serialization_dir / f"epoch-{last_epoch}/optim_g.pth"
+        optim_d_pth = serialization_dir / f"epoch-{last_epoch}/optim_d.pth"
+
+        logger.info(f"load state dict for generator.")
+        with open(generator_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        generator.load_state_dict(state_dict, strict=True)
+        logger.info(f"load state dict for discriminator.")
+        with open(discriminator_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        discriminator.load_state_dict(state_dict, strict=True)
+
+        logger.info(f"load state dict for optim_g.")
+        with open(optim_g_pth.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        optim_g.load_state_dict(state_dict)
+        logger.info(f"load state dict for optim_d.")
+        with open(optim_d_pth.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        optim_d.load_state_dict(state_dict)
+
+    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=config.lr_decay, last_epoch=last_epoch)
+    scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=config.lr_decay, last_epoch=last_epoch)
+
+    # training loop
+
+    # state
+    loss_d = 10000000000
+    loss_g = 10000000000
+    pesq_metric = 10000000000
+    mag_err = 10000000000
+    pha_err = 10000000000
+    com_err = 10000000000
+
+    model_list = list()
+    best_idx_epoch = None
+    best_metric = None
+    patience_count = 0
+
+    logger.info("training")
+    for idx_epoch in range(max(0, last_epoch+1), args.max_epochs):
+        # train
+        generator.train()
+        discriminator.train()
+
+        total_loss_d = 0.
+        total_loss_g = 0.
+        total_batches = 0.
+        progress_bar = tqdm(
+            total=len(train_data_loader),
+            desc="Training; epoch: {}".format(idx_epoch),
+        )
+        for batch in train_data_loader:
+            clean_audios, noisy_audios = batch
+            clean_audios = clean_audios.to(device)
+            noisy_audios = noisy_audios.to(device)
+            one_labels = torch.ones(clean_audios.shape[0]).to(device)
+
+            audio_g = generator.forward(noisy_audios)
+
+            clean_mag, clean_pha, clean_com = mag_pha_stft(clean_audios, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+            mag_g, pha_g, com_g = mag_pha_stft(audio_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+
+            clean_audio_list = torch.split(clean_audios, 1, dim=0)
+            enhanced_audio_list = torch.split(audio_g, 1, dim=0)
+            clean_audio_list = [t.squeeze().cpu().numpy() for t in clean_audio_list]
+            enhanced_audio_list = [t.squeeze().cpu().numpy() for t in enhanced_audio_list]
+
+            pesq_score_list: List[float] = run_batch_pesq(clean_audio_list, enhanced_audio_list, sample_rate=config.sample_rate, mode="nb")
+
+            # Discriminator
+            optim_d.zero_grad()
+            metric_r = discriminator.forward(clean_audios, clean_audios)
+            metric_g = discriminator.forward(clean_audios, audio_g.detach())
+            loss_disc_r = F.mse_loss(one_labels, metric_r.flatten())
+
+            if -1 in pesq_score_list:
+                # print("-1 in batch_pesq_score!")
+                loss_disc_g = 0
+            else:
+                pesq_score_list: torch.FloatTensor = torch.tensor([(score - 1) / 3.5 for score in pesq_score_list], dtype=torch.float32)
+                loss_disc_g = F.mse_loss(pesq_score_list.to(device), metric_g.flatten())
+
+            loss_disc_all = loss_disc_r + loss_disc_g
+            loss_disc_all.backward()
+            optim_d.step()
+
+            # Generator
+            optim_g.zero_grad()
+            # L2 Magnitude Loss
+            loss_mag = F.mse_loss(clean_mag, mag_g)
+            # Anti-wrapping Phase Loss
+            loss_ip, loss_gd, loss_iaf = phase_losses(clean_pha, pha_g)
+            loss_pha = loss_ip + loss_gd + loss_iaf
+            # L2 Complex Loss
+            loss_com = F.mse_loss(clean_com, com_g) * 2
+            # L2 Consistency Loss
+            # Time Loss
+            loss_time = F.l1_loss(clean_audios, audio_g)
+            # Metric Loss
+            metric_g = discriminator.forward(clean_mag, mag_g)
+            loss_metric = F.mse_loss(metric_g.flatten(), one_labels)
+
+            loss_gen_all = loss_mag * 0.9 + loss_pha * 0.3  + loss_com * 0.1 + loss_metric * 0.05 + loss_time * 0.2
+
+            loss_gen_all.backward()
+            optim_g.step()
+
+            total_loss_d += loss_disc_all.item()
+            total_loss_g += loss_gen_all.item()
+            total_batches += 1
+
+            loss_d = round(total_loss_d / total_batches, 4)
+            loss_g = round(total_loss_g / total_batches, 4)
+
+            progress_bar.update(1)
+            progress_bar.set_postfix({
+                "loss_d": loss_d,
+                "loss_g": loss_g,
+            })
+
+        # evaluation
+        generator.eval()
+        discriminator.eval()
+
+        torch.cuda.empty_cache()
+        total_pesq_score = 0.
+        total_mag_err = 0.
+        total_pha_err = 0.
+        total_com_err = 0.
+        total_batches = 0.
+
+        progress_bar = tqdm(
+            total=len(valid_data_loader),
+            desc="Evaluation; epoch: {}".format(idx_epoch),
+        )
+        with torch.no_grad():
+            for batch in valid_data_loader:
+                clean_audios, noisy_audios = batch
+                clean_audios = clean_audios.to(device)
+                noisy_audios = noisy_audios.to(device)
+
+                audio_g = generator.forward(noisy_audios)
+
+                clean_mag, clean_pha, clean_com = mag_pha_stft(clean_audios, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+                mag_g, pha_g, com_g = mag_pha_stft(audio_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+
+                clean_audio_list = torch.split(clean_audios, 1, dim=0)
+                enhanced_audio_list = torch.split(audio_g, 1, dim=0)
+                clean_audio_list = [t.squeeze().cpu().numpy() for t in clean_audio_list]
+                enhanced_audio_list = [t.squeeze().cpu().numpy() for t in enhanced_audio_list]
+                pesq_score = run_pesq_score(
+                    clean_audio_list,
+                    enhanced_audio_list,
+                    sample_rate = config.sample_rate,
+                    mode = "nb",
+                )
+                total_pesq_score += pesq_score
+                total_mag_err += F.mse_loss(clean_mag, mag_g).item()
+                val_ip_err, val_gd_err, val_iaf_err = phase_losses(clean_pha, pha_g)
+                total_pha_err += (val_ip_err + val_gd_err + val_iaf_err).item()
+                total_com_err += F.mse_loss(clean_com, com_g).item()
+
+                total_batches += 1
+
+                pesq_metric = round(total_pesq_score / total_batches, 4)
+                mag_err = round(total_mag_err / total_batches, 4)
+                pha_err = round(total_pha_err / total_batches, 4)
+                com_err = round(total_com_err / total_batches, 4)
+
+                progress_bar.update(1)
+                progress_bar.set_postfix({
+                    "pesq_metric": pesq_metric,
+                    "mag_err": mag_err,
+                    "pha_err": pha_err,
+                    "com_err": com_err,
+                })
+
+        # scheduler
+        scheduler_g.step()
+        scheduler_d.step()
+
+        # save path
+        epoch_dir = serialization_dir / "epoch-{}".format(idx_epoch)
+        epoch_dir.mkdir(parents=True, exist_ok=False)
+
+        # save models
+        generator.save_pretrained(epoch_dir.as_posix())
+        discriminator.save_pretrained(epoch_dir.as_posix())
+
+        # save optim
+        torch.save(optim_d.state_dict(), (epoch_dir / "optim_d.pth").as_posix())
+        torch.save(optim_g.state_dict(), (epoch_dir / "optim_g.pth").as_posix())
+
+        model_list.append(epoch_dir)
+        if len(model_list) >= args.num_serialized_models_to_keep:
+            model_to_delete: Path = model_list.pop(0)
+            shutil.rmtree(model_to_delete.as_posix())
+
+        # save metric
+        if best_metric is None:
+            best_idx_epoch = idx_epoch
+            best_metric = pesq_metric
+        elif pesq_metric > best_metric:
+            # great is better.
+            best_idx_epoch = idx_epoch
+            best_metric = pesq_metric
+        else:
+            pass
+
+        metrics = {
+            "idx_epoch": idx_epoch,
+            "best_idx_epoch": best_idx_epoch,
+            "loss_d": loss_d,
+            "loss_g": loss_g,
+
+            "pesq_metric": pesq_metric,
+            "mag_err": mag_err,
+            "pha_err": pha_err,
+            "com_err": com_err,
+
+        }
+        metrics_filename = epoch_dir / "metrics_epoch.json"
+        with open(metrics_filename, "w", encoding="utf-8") as f:
+            json.dump(metrics, f, indent=4, ensure_ascii=False)
+
+        # save best
+        best_dir = serialization_dir / "best"
+        if best_idx_epoch == idx_epoch:
+            if best_dir.exists():
+                shutil.rmtree(best_dir)
+            shutil.copytree(epoch_dir, best_dir)
+
+        # early stop
+        early_stop_flag = False
+        if best_idx_epoch == idx_epoch:
+            patience_count = 0
+        else:
+            patience_count += 1
+        if patience_count >= args.patience:
+            early_stop_flag = True
+
+        # early stop
+        if early_stop_flag:
+            break
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/nx_clean_unet/step_3_evaluation.py

@@ -0,0 +1,6 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+if __name__ == '__main__':
+    pass

examples/nx_clean_unet/yaml/config.yaml

@@ -0,0 +1,23 @@
+model_name: "nx_clean_unet"
+
+sample_rate: 8000
+segment_size: 16000
+n_fft: 512
+win_size: 200
+hop_size: 80
+
+down_sampling_num_layers: 5
+down_sampling_in_channels: 1
+down_sampling_hidden_channels: 64
+down_sampling_kernel_size: 4
+down_sampling_stride: 2
+
+tsfm_hidden_size: 256
+tsfm_attention_heads: 4
+tsfm_num_blocks: 6
+tsfm_dropout_rate: 0.1
+
+discriminator_dim: 32
+discriminator_in_channel: 2
+
+compress_factor: 0.3

toolbox/torchaudio/models/clean_unet/configuration_clean_unet.py

@@ -3,7 +3,7 @@
 from toolbox.torchaudio.configuration_utils import PretrainedConfig
 
 
-class CleanUnetConfig(PretrainedConfig):
+class CleanUNetConfig(PretrainedConfig):
     def __init__(self,
                  channels_input: int = 1,
                  channels_output: int = 1,
@@ -21,7 +21,7 @@ class CleanUnetConfig(PretrainedConfig):
 
                  **kwargs
                  ):
-        super(CleanUnetConfig, self).__init__(**kwargs)
+        super(CleanUNetConfig, self).__init__(**kwargs)
         self.channels_input = channels_input
         self.channels_output = channels_output
 

toolbox/torchaudio/models/clean_unet/inference_clean_unet.py

@@ -0,0 +1,104 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import logging
+from pathlib import Path
+import shutil
+import tempfile
+import zipfile
+
+import librosa
+import numpy as np
+import torch
+import torchaudio
+
+from project_settings import project_path
+from toolbox.torchaudio.models.clean_unet.configuration_clean_unet import CleanUNetConfig
+from toolbox.torchaudio.models.clean_unet.modeling_clean_unet import CleanUNetPretrainedModel, MODEL_FILE
+
+logger = logging.getLogger("toolbox")
+
+
+class InferenceCleanUNet(object):
+    def __init__(self, pretrained_model_path_or_zip_file: str, device: str = "cpu"):
+        self.pretrained_model_path_or_zip_file = pretrained_model_path_or_zip_file
+        self.device = torch.device(device)
+
+        logger.info(f"loading model; model_file: {self.pretrained_model_path_or_zip_file}")
+        config, model = self.load_models(self.pretrained_model_path_or_zip_file)
+        logger.info(f"model loading completed; model_file: {self.pretrained_model_path_or_zip_file}")
+
+        self.config = config
+        self.model = model
+        self.model.to(device)
+        self.model.eval()
+
+    def load_models(self, model_path: str):
+        model_path = Path(model_path)
+        if model_path.name.endswith(".zip"):
+            with zipfile.ZipFile(model_path.as_posix(), "r") as f_zip:
+                out_root = Path(tempfile.gettempdir()) / "nx_denoise"
+                out_root.mkdir(parents=True, exist_ok=True)
+                f_zip.extractall(path=out_root)
+            model_path = out_root / model_path.stem
+
+        config = CleanUNetConfig.from_pretrained(
+            pretrained_model_name_or_path=model_path.as_posix(),
+        )
+        model = CleanUNetPretrainedModel.from_pretrained(
+            pretrained_model_name_or_path=model_path.as_posix(),
+        )
+        model.to(self.device)
+        model.eval()
+
+        shutil.rmtree(model_path)
+        return config, model
+
+    def enhancement_by_ndarray(self, noisy_audio: np.ndarray) -> np.ndarray:
+        noisy_audio = torch.tensor(noisy_audio, dtype=torch.float32)
+        noisy_audio = noisy_audio.unsqueeze(dim=0)
+
+        # noisy_audio shape: [batch_size, n_samples]
+        enhanced_audio = self.enhancement_by_tensor(noisy_audio)
+        # noisy_audio shape: [channels, n_samples]
+        return enhanced_audio.cpu().numpy()
+
+    def enhancement_by_tensor(self, noisy_audio: torch.Tensor) -> torch.Tensor:
+        if torch.max(noisy_audio) > 1 or torch.min(noisy_audio) < -1:
+            raise AssertionError(f"The value range of audio samples should be between -1 and 1.")
+
+        # noisy_audio shape: [batch_size, num_samples]
+        noisy_audios = noisy_audio.to(self.device)
+
+        with torch.no_grad():
+            enhanced_audios = self.model.forward(noisy_audios)
+            # enhanced_audio shape: [batch_size, channels, num_samples]
+            # enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
+
+        enhanced_audio = enhanced_audios[0]
+
+        # enhanced_audio shape: [channels, num_samples]
+        return enhanced_audio
+
+def main():
+    model_zip_file = project_path / "trained_models/clean-unet-aishell-18-epoch.zip"
+    infer_mpnet = InferenceCleanUNet(model_zip_file)
+
+    sample_rate = 8000
+    noisy_audio_file = project_path / "data/examples/ai_agent/dfaaf264-b5e3-4ca2-b5cb-5b6d637d962d_section_1.wav"
+    noisy_audio, _ = librosa.load(
+        noisy_audio_file.as_posix(),
+        sr=sample_rate,
+    )
+    noisy_audio = torch.tensor(noisy_audio, dtype=torch.float32)
+    noisy_audio = noisy_audio.unsqueeze(dim=0)
+
+    enhanced_audio = infer_mpnet.enhancement_by_tensor(noisy_audio)
+
+    filename = "enhanced_audio.wav"
+    torchaudio.save(filename, enhanced_audio.detach().cpu(), sample_rate)
+
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/clean_unet/modeling_clean_unet.py

@@ -22,7 +22,7 @@ import torch.nn.functional as F
 
 from toolbox.torchaudio.configuration_utils import CONFIG_FILE
 from toolbox.torchaudio.models.clean_unet.transformer import TransformerEncoder
-from toolbox.torchaudio.models.clean_unet.configuration_clean_unet import CleanUnetConfig
+from toolbox.torchaudio.models.clean_unet.configuration_clean_unet import CleanUNetConfig
 
 
 def weight_scaling_init(layer):
@@ -196,7 +196,7 @@ class CleanUNet(nn.Module):
 
         x = self.tsfm_conv1(x)  # C 1024 -> 512
         x = x.permute(0, 2, 1)
-        x = self.tsfm_encoder(x, src_mask=attn_mask)
+        x = self.tsfm_encoder.forward(x, src_mask=attn_mask)
         x = x.permute(0, 2, 1)
         x = self.tsfm_conv2(x)  # C 512 -> 1024
 
@@ -215,7 +215,7 @@ MODEL_FILE = "model.pt"
 
 class CleanUNetPretrainedModel(CleanUNet):
     def __init__(self,
-                 config: CleanUnetConfig,
+                 config: CleanUNetConfig,
                  ):
         super(CleanUNetPretrainedModel, self).__init__(
             channels_input=config.channels_input,
@@ -234,7 +234,7 @@ class CleanUNetPretrainedModel(CleanUNet):
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
-        config = CleanUnetConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        config = CleanUNetConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
 
         model = cls(config)
 
@@ -272,7 +272,7 @@ class CleanUNetPretrainedModel(CleanUNet):
 
 def main():
 
-    config = CleanUnetConfig()
+    config = CleanUNetConfig()
     model = CleanUNetPretrainedModel(config)
 
     print_size(model, keyword="tsfm")

toolbox/torchaudio/models/nx_clean_unet/__init__.py

@@ -0,0 +1,6 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/nx_clean_unet/configuration_nx_clean_unet.py

@@ -0,0 +1,56 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from toolbox.torchaudio.configuration_utils import PretrainedConfig
+
+
+class NXCleanUNetConfig(PretrainedConfig):
+    """
+    https://github.com/yxlu-0102/MP-SENet/blob/main/config.json
+    """
+    def __init__(self,
+                 n_fft: int = 512,
+                 win_length: int = 200,
+                 hop_length: int = 80,
+
+                 down_sampling_num_layers: int = 5,
+                 down_sampling_in_channels: int = 1,
+                 down_sampling_hidden_channels: int = 64,
+                 down_sampling_kernel_size: int = 4,
+                 down_sampling_stride: int = 2,
+
+                 tsfm_hidden_size: int = 256,
+                 tsfm_attention_heads: int = 4,
+                 tsfm_num_blocks: int = 6,
+                 tsfm_dropout_rate: float = 0.1,
+
+                 discriminator_dim: int = 32,
+                 discriminator_in_channel: int = 2,
+
+                 compress_factor: float = 0.3,
+
+                 **kwargs
+                 ):
+        super(NXCleanUNetConfig, self).__init__(**kwargs)
+        self.n_fft = n_fft
+        self.win_length = win_length
+        self.hop_length = hop_length
+
+        self.down_sampling_num_layers = down_sampling_num_layers
+        self.down_sampling_in_channels = down_sampling_in_channels
+        self.down_sampling_hidden_channels = down_sampling_hidden_channels
+        self.down_sampling_kernel_size = down_sampling_kernel_size
+        self.down_sampling_stride = down_sampling_stride
+
+        self.tsfm_hidden_size = tsfm_hidden_size
+        self.tsfm_attention_heads = tsfm_attention_heads
+        self.tsfm_num_blocks = tsfm_num_blocks
+        self.tsfm_dropout_rate = tsfm_dropout_rate
+
+        self.discriminator_dim = discriminator_dim
+        self.discriminator_in_channel = discriminator_in_channel
+
+        self.compress_factor = compress_factor
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/nx_clean_unet/discriminator.py

@@ -0,0 +1,126 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import os
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+import torchaudio
+
+from toolbox.torchaudio.configuration_utils import CONFIG_FILE
+from toolbox.torchaudio.models.nx_clean_unet.configuration_nx_clean_unet import NXCleanUNetConfig
+from toolbox.torchaudio.models.nx_clean_unet.utils import LearnableSigmoid1d
+
+
+class MetricDiscriminator(nn.Module):
+    def __init__(self, config: NXCleanUNetConfig):
+        super(MetricDiscriminator, self).__init__()
+        dim = config.discriminator_dim
+        self.in_channel = config.discriminator_in_channel
+
+        self.n_fft = config.n_fft
+        self.win_length = config.win_length
+        self.hop_length = config.hop_length
+
+        self.layers = nn.Sequential(
+            torchaudio.transforms.Spectrogram(
+                n_fft=self.n_fft,
+                win_length=self.win_length,
+                hop_length=self.hop_length,
+                power=1.0,
+                window_fn=torch.hamming_window,
+                # window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
+            ),
+            nn.utils.spectral_norm(nn.Conv2d(self.in_channel, dim, (4,4), (2,2), (1,1), bias=False)),
+            nn.InstanceNorm2d(dim, affine=True),
+            nn.PReLU(dim),
+            nn.utils.spectral_norm(nn.Conv2d(dim, dim*2, (4,4), (2,2), (1,1), bias=False)),
+            nn.InstanceNorm2d(dim*2, affine=True),
+            nn.PReLU(dim*2),
+            nn.utils.spectral_norm(nn.Conv2d(dim*2, dim*4, (4,4), (2,2), (1,1), bias=False)),
+            nn.InstanceNorm2d(dim*4, affine=True),
+            nn.PReLU(dim*4),
+            nn.utils.spectral_norm(nn.Conv2d(dim*4, dim*8, (4,4), (2,2), (1,1), bias=False)),
+            nn.InstanceNorm2d(dim*8, affine=True),
+            nn.PReLU(dim*8),
+            nn.AdaptiveMaxPool2d(1),
+            nn.Flatten(),
+            nn.utils.spectral_norm(nn.Linear(dim*8, dim*4)),
+            nn.Dropout(0.3),
+            nn.PReLU(dim*4),
+            nn.utils.spectral_norm(nn.Linear(dim*4, 1)),
+            LearnableSigmoid1d(1)
+        )
+
+    def forward(self, x, y):
+        xy = torch.stack((x, y), dim=1)
+        return self.layers(xy)
+
+
+MODEL_FILE = "discriminator.pt"
+
+
+class MetricDiscriminatorPretrainedModel(MetricDiscriminator):
+    def __init__(self,
+                 config: NXCleanUNetConfig,
+                 ):
+        super(MetricDiscriminatorPretrainedModel, self).__init__(
+            config=config,
+        )
+        self.config = config
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        config = NXCleanUNetConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+        model = cls(config)
+
+        if os.path.isdir(pretrained_model_name_or_path):
+            ckpt_file = os.path.join(pretrained_model_name_or_path, MODEL_FILE)
+        else:
+            ckpt_file = pretrained_model_name_or_path
+
+        with open(ckpt_file, "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+        return model
+
+    def save_pretrained(self,
+                        save_directory: Union[str, os.PathLike],
+                        state_dict: Optional[dict] = None,
+                        ):
+
+        model = self
+
+        if state_dict is None:
+            state_dict = model.state_dict()
+
+        os.makedirs(save_directory, exist_ok=True)
+
+        # save state dict
+        model_file = os.path.join(save_directory, MODEL_FILE)
+        torch.save(state_dict, model_file)
+
+        # save config
+        config_file = os.path.join(save_directory, CONFIG_FILE)
+        self.config.to_yaml_file(config_file)
+        return save_directory
+
+
+def main():
+    config = NXCleanUNetConfig()
+    discriminator = MetricDiscriminator(config=config)
+
+    # shape: [batch_size, num_samples]
+    x = torch.ones([4, int(4.5 * 16000)])
+    y = torch.ones([4, int(4.5 * 16000)])
+
+    output = discriminator.forward(x, y)
+    print(output.shape)
+    print(output)
+
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/nx_clean_unet/loss.py

@@ -0,0 +1,22 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import numpy as np
+import torch
+
+
+def anti_wrapping_function(x):
+
+    return torch.abs(x - torch.round(x / (2 * np.pi)) * 2 * np.pi)
+
+
+def phase_losses(phase_r, phase_g):
+
+    ip_loss = torch.mean(anti_wrapping_function(phase_r - phase_g))
+    gd_loss = torch.mean(anti_wrapping_function(torch.diff(phase_r, dim=1) - torch.diff(phase_g, dim=1)))
+    iaf_loss = torch.mean(anti_wrapping_function(torch.diff(phase_r, dim=2) - torch.diff(phase_g, dim=2)))
+
+    return ip_loss, gd_loss, iaf_loss
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/nx_clean_unet/metrics.py

@@ -0,0 +1,80 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from joblib import Parallel, delayed
+import numpy as np
+from pesq import pesq
+from typing import List
+
+from pesq import cypesq
+
+
+def run_pesq(clean_audio: np.ndarray,
+             noisy_audio: np.ndarray,
+             sample_rate: int = 16000,
+             mode: str = "wb",
+             ) -> float:
+    if sample_rate == 8000 and mode == "wb":
+        raise AssertionError(f"mode should be `nb` when sample_rate is 8000")
+    try:
+        pesq_score = pesq(sample_rate, clean_audio, noisy_audio, mode)
+    except cypesq.NoUtterancesError as e:
+        pesq_score = -1
+    except Exception as e:
+        print(f"pesq failed. error type: {type(e)}, error text: {str(e)}")
+        pesq_score = -1
+    return pesq_score
+
+
+def run_batch_pesq(clean_audio_list: List[np.ndarray],
+                   noisy_audio_list: List[np.ndarray],
+                   sample_rate: int = 16000,
+                   mode: str = "wb",
+                   n_jobs: int = 4,
+                   ) -> List[float]:
+    parallel = Parallel(n_jobs=n_jobs)
+
+    parallel_tasks = list()
+    for clean_audio, noisy_audio in zip(clean_audio_list, noisy_audio_list):
+        parallel_task = delayed(run_pesq)(clean_audio, noisy_audio, sample_rate, mode)
+        parallel_tasks.append(parallel_task)
+
+    pesq_score_list = parallel.__call__(parallel_tasks)
+    return pesq_score_list
+
+
+def run_pesq_score(clean_audio_list: List[np.ndarray],
+                   noisy_audio_list: List[np.ndarray],
+                   sample_rate: int = 16000,
+                   mode: str = "wb",
+                   n_jobs: int = 4,
+                   ) -> List[float]:
+
+    pesq_score_list = run_batch_pesq(clean_audio_list=clean_audio_list,
+                                     noisy_audio_list=noisy_audio_list,
+                                     sample_rate=sample_rate,
+                                     mode=mode,
+                                     n_jobs=n_jobs,
+                                     )
+
+    pesq_score = np.mean(pesq_score_list)
+    return pesq_score
+
+
+def main():
+    clean_audio = np.random.uniform(low=0, high=1, size=(2, 160000,))
+    noisy_audio = np.random.uniform(low=0, high=1, size=(2, 160000,))
+
+    clean_audio_list = list(clean_audio)
+    noisy_audio_list = list(noisy_audio)
+
+    pesq_score_list = run_batch_pesq(clean_audio_list, noisy_audio_list)
+    print(pesq_score_list)
+
+    pesq_score = run_pesq_score(clean_audio_list, noisy_audio_list)
+    print(pesq_score)
+
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/nx_clean_unet/modeling_nx_clean_unet.py

@@ -0,0 +1,326 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import os
+from typing import Optional, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from toolbox.torchaudio.configuration_utils import CONFIG_FILE
+from toolbox.torchaudio.models.nx_clean_unet.configuration_nx_clean_unet import NXCleanUNetConfig
+from toolbox.torchaudio.models.nx_clean_unet.transformer.transformer import TransformerEncoder
+
+
+class DownSamplingBlock(nn.Module):
+    def __init__(self,
+                 in_channels: int,
+                 hidden_channels: int,
+                 kernel_size: int,
+                 stride: int,
+                 ):
+        super(DownSamplingBlock, self).__init__()
+        self.conv1 = nn.Conv1d(in_channels, hidden_channels, kernel_size, stride)
+        self.relu = nn.ReLU()
+        self.conv2 = nn.Conv1d(hidden_channels, hidden_channels * 2, 1)
+        self.glu = nn.GLU(dim=1)
+
+    def forward(self, x: torch.Tensor):
+        # x shape: [batch_size, 1, num_samples]
+        x = self.conv1.forward(x)
+        # x shape: [batch_size, hidden_channels, new_num_samples]
+        x = self.relu(x)
+        x = self.conv2.forward(x)
+        # x shape: [batch_size, hidden_channels*2, new_num_samples]
+        x = self.glu(x)
+        # x shape: [batch_size, hidden_channels, new_num_samples]
+        # new_num_samples = (num_samples-kernel_size) // stride + 1
+        return x
+
+
+class DownSampling(nn.Module):
+    def __init__(self,
+                 num_layers: int,
+                 in_channels: int,
+                 hidden_channels: int,
+                 kernel_size: int,
+                 stride: int,
+                 ):
+        super(DownSampling, self).__init__()
+        self.num_layers = num_layers
+
+        self.down_sampling_block_list = list()
+
+        for idx in range(self.num_layers):
+            down_sampling_block = DownSamplingBlock(
+                in_channels=in_channels,
+                hidden_channels=hidden_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+            )
+            self.down_sampling_block_list.append(down_sampling_block)
+            in_channels = hidden_channels
+
+    def forward(self, x: torch.Tensor):
+        # x shape: [batch_size, channels, num_samples]
+        for down_sampling_block in self.down_sampling_block_list:
+            x = down_sampling_block.forward(x)
+        # x shape: [batch_size, hidden_channels, num_samples**]
+        return x
+
+
+class UpSamplingBlock(nn.Module):
+    def __init__(self,
+                 out_channels: int,
+                 hidden_channels: int,
+                 kernel_size: int,
+                 stride: int,
+                 do_relu: bool = True,
+                 ):
+        super(UpSamplingBlock, self).__init__()
+        self.do_relu = do_relu
+
+        self.conv1 = nn.Conv1d(hidden_channels, hidden_channels * 2, 1)
+        self.glu = nn.GLU(dim=1)
+        self.convt = nn.ConvTranspose1d(hidden_channels, out_channels, kernel_size, stride)
+        self.relu = nn.ReLU()
+
+    def forward(self, x: torch.Tensor):
+        # x shape: [batch_size, hidden_channels*2, num_samples]
+        x = self.conv1.forward(x)
+        # x shape: [batch_size, hidden_channels, num_samples]
+        x = self.glu(x)
+        # x shape: [batch_size, hidden_channels, num_samples]
+        x = self.convt.forward(x)
+        # x shape: [batch_size, hidden_channels, new_num_samples]
+        # new_num_samples = (num_samples - 1) * stride + kernel_size
+        if self.do_relu:
+            x = self.relu(x)
+        return x
+
+
+class UpSampling(nn.Module):
+    def __init__(self,
+                 num_layers: int,
+                 out_channels: int,
+                 hidden_channels: int,
+                 kernel_size: int,
+                 stride: int,
+                 ):
+        super(UpSampling, self).__init__()
+        self.num_layers = num_layers
+
+        self.up_sampling_block_list = list()
+
+        for idx in range(self.num_layers-1):
+            up_sampling_block = UpSamplingBlock(
+                out_channels=hidden_channels,
+                hidden_channels=hidden_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                do_relu=True,
+            )
+            self.up_sampling_block_list.append(up_sampling_block)
+        else:
+            up_sampling_block = UpSamplingBlock(
+                out_channels=out_channels,
+                hidden_channels=hidden_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                do_relu=False,
+            )
+            self.up_sampling_block_list.append(up_sampling_block)
+
+    def forward(self, x: torch.Tensor):
+        # x shape: [batch_size, channels, num_samples]
+        for up_sampling_block in self.up_sampling_block_list:
+            x = up_sampling_block.forward(x)
+        return x
+
+
+def get_padding_length(length, num_layers: int, kernel_size: int, stride: int):
+    for _ in range(num_layers):
+        if length < kernel_size:
+            length = 1
+        else:
+            length = 1 + np.ceil((length - kernel_size) / stride)
+
+    for _ in range(num_layers):
+        length = (length - 1) * stride + kernel_size
+
+    padded_length = int(length)
+    return padded_length
+
+
+class NXCleanUNet(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+
+        self.down_sampling = DownSampling(
+            num_layers=config.down_sampling_num_layers,
+            in_channels=config.down_sampling_in_channels,
+            hidden_channels=config.down_sampling_hidden_channels,
+            kernel_size=config.down_sampling_kernel_size,
+            stride=config.down_sampling_stride,
+        )
+        self.transformer = TransformerEncoder(
+            input_size=config.down_sampling_hidden_channels,
+            hidden_size=config.tsfm_hidden_size,
+            attention_heads=config.tsfm_attention_heads,
+            num_blocks=config.tsfm_num_blocks,
+            dropout_rate=config.tsfm_dropout_rate,
+        )
+        self.up_sampling = UpSampling(
+            num_layers=config.down_sampling_num_layers,
+            out_channels=config.down_sampling_in_channels,
+            hidden_channels=config.down_sampling_hidden_channels,
+            kernel_size=config.down_sampling_kernel_size,
+            stride=config.down_sampling_stride,
+        )
+
+    def forward(self, noisy_audios: torch.Tensor):
+        # noisy_audios shape: [batch_size, 1, n_samples]
+
+        n_samples = noisy_audios.shape[-1]
+        padded_length = get_padding_length(
+            n_samples,
+            num_layers=self.config.down_sampling_num_layers,
+            kernel_size=self.config.down_sampling_kernel_size,
+            stride=self.config.down_sampling_stride,
+        )
+        noisy_audios_padded = F.pad(input=noisy_audios, pad=(0, padded_length - n_samples), mode="constant", value=0)
+
+        bottle_neck = self.down_sampling.forward(noisy_audios_padded)
+        # bottle_neck shape: [batch_size, channels, time_steps]
+
+        bottle_neck = torch.transpose(bottle_neck, dim0=-2, dim1=-1)
+        # bottle_neck shape: [batch_size, time_steps, input_size]
+
+        bottle_neck = self.transformer.forward(bottle_neck)
+        # bottle_neck shape: [batch_size, time_steps, input_size]
+
+        bottle_neck = torch.transpose(bottle_neck, dim0=-2, dim1=-1)
+        # bottle_neck shape: [batch_size, channels, time_steps]
+
+        enhanced_audios = self.up_sampling.forward(bottle_neck)
+
+        enhanced_audios = enhanced_audios[:, :, :n_samples]
+        # enhanced_audios shape: [batch_size, 1, n_samples]
+        return enhanced_audios
+
+
+MODEL_FILE = "generator.pt"
+
+
+class NXCleanUNetPretrainedModel(NXCleanUNet):
+    def __init__(self,
+                 config: NXCleanUNetConfig,
+                 ):
+        super(NXCleanUNetPretrainedModel, self).__init__(
+            config=config,
+        )
+        self.config = config
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        config = NXCleanUNetConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+        model = cls(config)
+
+        if os.path.isdir(pretrained_model_name_or_path):
+            ckpt_file = os.path.join(pretrained_model_name_or_path, MODEL_FILE)
+        else:
+            ckpt_file = pretrained_model_name_or_path
+
+        with open(ckpt_file, "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+        return model
+
+    def save_pretrained(self,
+                        save_directory: Union[str, os.PathLike],
+                        state_dict: Optional[dict] = None,
+                        ):
+
+        model = self
+
+        if state_dict is None:
+            state_dict = model.state_dict()
+
+        os.makedirs(save_directory, exist_ok=True)
+
+        # save state dict
+        model_file = os.path.join(save_directory, MODEL_FILE)
+        torch.save(state_dict, model_file)
+
+        # save config
+        config_file = os.path.join(save_directory, CONFIG_FILE)
+        self.config.to_yaml_file(config_file)
+        return save_directory
+
+
+
+def main2():
+
+    config = NXCleanUNetConfig()
+    down_sampling = DownSampling(
+        num_layers=config.down_sampling_num_layers,
+        in_channels=config.down_sampling_in_channels,
+        hidden_channels=config.down_sampling_hidden_channels,
+        kernel_size=config.down_sampling_kernel_size,
+        stride=config.down_sampling_stride,
+    )
+    up_sampling = UpSampling(
+        num_layers=config.down_sampling_num_layers,
+        out_channels=config.down_sampling_in_channels,
+        hidden_channels=config.down_sampling_hidden_channels,
+        kernel_size=config.down_sampling_kernel_size,
+        stride=config.down_sampling_stride,
+    )
+
+    # shape: [batch_size, channels, num_samples]
+    # min length: 94, stride: 32, 32 == 2**5
+    # x = torch.ones([4, 1, 94])
+    # x = torch.ones([4, 1, 126])
+    # x = torch.ones([4, 1, 158])
+    x = torch.ones([4, 1, 190])
+
+    length = x.shape[-1]
+    padded_length = get_padding_length(
+        length,
+        num_layers=config.down_sampling_num_layers,
+        kernel_size=config.down_sampling_kernel_size,
+        stride=config.down_sampling_stride,
+    )
+    x = F.pad(input=x, pad=(0, padded_length - length), mode="constant", value=0)
+    # print(x)
+    print(x.shape)
+    bottle_neck = down_sampling.forward(x)
+    print("-" * 150)
+    x = up_sampling.forward(bottle_neck)
+    print(x.shape)
+    return
+
+
+def main():
+
+    config = NXCleanUNetConfig()
+
+    # shape: [batch_size, channels, num_samples]
+    # min length: 94, stride: 32, 32 == 2**5
+    # x = torch.ones([4, 1, 94])
+    # x = torch.ones([4, 1, 126])
+    # x = torch.ones([4, 1, 158])
+    # x = torch.ones([4, 1, 190])
+    x = torch.ones([4, 1, 16000])
+
+    model = NXCleanUNet(config)
+    enhanced_audios = model.forward(x)
+    print(enhanced_audios.shape)
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/nx_clean_unet/transformer/__init__.py

@@ -0,0 +1,6 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/nx_clean_unet/transformer/mask.py

@@ -0,0 +1,66 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import torch
+
+
+def make_pad_mask(lengths: torch.Tensor,
+                  max_len: int = 0,
+                  ) -> torch.Tensor:
+    batch_size = lengths.size(0)
+    max_len = max_len if max_len > 0 else lengths.max().item()
+    seq_range = torch.arange(
+        0,
+        max_len,
+        dtype=torch.int64,
+        device=lengths.device
+    )
+    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
+    seq_length_expand = lengths.unsqueeze(-1)
+    mask = seq_range_expand >= seq_length_expand
+    return mask
+
+
+
+def subsequent_chunk_mask(
+        size: int,
+        chunk_size: int,
+        num_left_chunks: int = -1,
+        device: torch.device = torch.device("cpu"),
+) -> torch.Tensor:
+    """
+    Create mask for subsequent steps (size, size) with chunk size,
+    this is for streaming encoder
+
+    Examples:
+    > subsequent_chunk_mask(4, 2)
+    [[1, 1, 0, 0],
+     [1, 1, 0, 0],
+     [1, 1, 1, 1],
+     [1, 1, 1, 1]]
+
+    :param size: int. size of mask.
+    :param chunk_size: int. size of chunk.
+    :param num_left_chunks: int. number of left chunks. <0: use full chunk. >=0 use num_left_chunks.
+    :param device: torch.device. "cpu" or "cuda" or torch.Tensor.device.
+    :return: torch.Tensor. mask
+    """
+
+    ret = torch.zeros(size, size, device=device, dtype=torch.bool)
+    for i in range(size):
+        if num_left_chunks < 0:
+            start = 0
+        else:
+            start = max((i // chunk_size - num_left_chunks) * chunk_size, 0)
+        ending = min((i // chunk_size + 1) * chunk_size, size)
+        ret[i, start:ending] = True
+    return ret
+
+
+def main():
+    chunk_mask = subsequent_chunk_mask(size=8, chunk_size=2, num_left_chunks=2)
+    print(chunk_mask)
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/nx_clean_unet/transformer/transformer.py

@@ -0,0 +1,577 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import math
+from typing import Dict, Optional, Tuple, List, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as f
+
+from toolbox.torchaudio.models.nx_clean_unet.transformer.mask import subsequent_chunk_mask
+
+
+class SinusoidalPositionalEncoding(nn.Module):
+    """
+    Positional Encoding
+
+    PE(pos, 2i)   = sin(pos/(10000^(2i/dmodel)))
+    PE(pos, 2i+1) = cos(pos/(10000^(2i/dmodel)))
+    """
+
+    @staticmethod
+    def demo1():
+        batch_size = 2
+        time_steps = 10
+        embedding_dim = 64
+
+        pe = SinusoidalPositionalEncoding(
+            embedding_dim=embedding_dim,
+            dropout_rate=0.1,
+        )
+
+        x = torch.randn(size=(batch_size, time_steps, embedding_dim))
+
+        x, pos_emb = pe.forward(x)
+
+        # torch.Size([2, 10, 64])
+        print(x.shape)
+        # torch.Size([1, 10, 64])
+        print(pos_emb.shape)
+        return
+
+    @staticmethod
+    def demo2():
+        batch_size = 2
+        time_steps = 10
+        embedding_dim = 64
+
+        pe = SinusoidalPositionalEncoding(
+            embedding_dim=embedding_dim,
+            dropout_rate=0.1,
+        )
+
+        x = torch.randn(size=(batch_size, time_steps, embedding_dim))
+        offset = torch.randint(low=3, high=7, size=(batch_size,))
+        x, pos_emb = pe.forward(x, offset=offset)
+
+        # tensor([3, 4])
+        print(offset)
+        # torch.Size([2, 10, 64])
+        print(x.shape)
+        # torch.Size([2, 10, 64])
+        print(pos_emb.shape)
+        return
+
+    def __init__(self,
+                 embedding_dim: int,
+                 dropout_rate: float,
+                 max_length: int = 5000,
+                 reverse: bool = False
+                 ):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.dropout_rate = dropout_rate
+        self.max_length = max_length
+        self.reverse = reverse
+
+        self.x_scale = math.sqrt(self.embedding_dim)
+        self.dropout = torch.nn.Dropout(p=dropout_rate)
+
+        self.pe = torch.zeros(self.max_length, self.embedding_dim)
+        position = torch.arange(0, self.max_length, dtype=torch.float32).unsqueeze(1)
+
+        div_term = torch.exp(
+            torch.arange(0, self.embedding_dim, 2, dtype=torch.float32) *
+            - (math.log(10000.0) / self.embedding_dim)
+        )
+        self.pe[:, 0::2] = torch.sin(position * div_term)
+        self.pe[:, 1::2] = torch.cos(position * div_term)
+        self.pe = self.pe.unsqueeze(0)
+
+    def forward(self,
+                x: torch.Tensor,
+                offset: Union[int, torch.Tensor] = 0
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Add positional encoding.
+        :param x: torch.Tensor. Input. shape=(batch_size, time_steps, ...).
+        :param offset: int or torch.Tensor. position offset.
+        :return:
+        torch.Tensor. Encoded tensor. shape=(batch_size, time_steps, ...).
+        torch.Tensor. for compatibility to RelPositionalEncoding. shape=(1, time_steps, ...).
+        """
+        self.pe = self.pe.to(x.device)
+        pos_emb = self.position_encoding(offset, x.size(1), False)
+        x = x * self.x_scale + pos_emb
+        return self.dropout(x), self.dropout(pos_emb)
+
+    def position_encoding(self,
+                          offset: Union[int, torch.Tensor],
+                          size: int,
+                          apply_dropout: bool = True
+                          ) -> torch.Tensor:
+        """
+        For getting encoding in a streaming fashion.
+
+        Attention!!!!!
+        we apply dropout only once at the whole utterance level in a none
+        streaming way, but will call this function several times with
+        increasing input size in a streaming scenario, so the dropout will
+        be applied several times.
+
+        :param offset: int or torch.Tensor. start offset.
+        :param size: int. required size of position encoding.
+        :param apply_dropout:
+        :return: torch.Tensor. Corresponding encoding.
+        """
+        if isinstance(offset, int):
+            assert offset + size <= self.max_length
+            pos_emb = self.pe[:, offset:offset + size]
+        elif isinstance(offset, torch.Tensor) and offset.dim() == 0:  # scalar
+            assert offset + size <= self.max_length
+            pos_emb = self.pe[:, offset:offset + size]
+        else:    # for batched streaming decoding on GPU
+            # offset. shape=(batch_size,)
+            assert torch.max(offset) + size <= self.max_length
+
+            # shape=(batch_size, time_steps)
+            index = offset.unsqueeze(1) + torch.arange(0, size).to(offset.device)
+            flag = index > 0
+            # remove negative offset
+            index = index * flag
+            # shape=(batch_size, time_steps, embedding_dim)
+            pos_emb = f.embedding(index, self.pe[0])
+
+        if apply_dropout:
+            pos_emb = self.dropout(pos_emb)
+        return pos_emb
+
+
+class RelPositionalEncoding(SinusoidalPositionalEncoding):
+    """
+    Relative positional encoding module.
+
+    See : Appendix B in https://arxiv.org/abs/1901.02860
+
+    """
+    def __init__(self,
+                 embedding_dim: int,
+                 dropout_rate: float,
+                 max_length: int = 5000,
+                 ):
+        super().__init__(embedding_dim, dropout_rate, max_length, reverse=True)
+
+    def forward(self,
+                x: torch.Tensor,
+                offset: Union[int, torch.Tensor] = 0
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Compute positional encoding.
+        :param x: torch.Tensor. Input. shape=(batch_size, time_steps, ...).
+        :param offset:
+        :return:
+        torch.Tensor. Encoded tensor. shape=(batch_size, time_steps, ...).
+        torch.Tensor. Positional embedding tensor. shape=(1, time_steps, ...).
+        """
+        self.pe = self.pe.to(x.device)
+        x = x * self.x_scale
+        pos_emb = self.position_encoding(offset, x.size(1), False)
+        return self.dropout(x), self.dropout(pos_emb)
+
+
+class PositionwiseFeedForward(nn.Module):
+    def __init__(self,
+                 input_dim: int,
+                 hidden_units: int,
+                 dropout_rate: float,
+                 activation: torch.nn.Module = torch.nn.ReLU()):
+        """
+        FeedForward are applied on each position of the sequence.
+        the output dim is same with the input dim.
+
+        :param input_dim: int. input dimension.
+        :param hidden_units: int. the number of hidden units.
+        :param dropout_rate: float. dropout rate.
+        :param activation: torch.nn.Module. activation function.
+        """
+        super(PositionwiseFeedForward, self).__init__()
+        self.w_1 = torch.nn.Linear(input_dim, hidden_units)
+        self.activation = activation
+        self.dropout = torch.nn.Dropout(dropout_rate)
+        self.w_2 = torch.nn.Linear(hidden_units, input_dim)
+
+    def forward(self, xs: torch.Tensor) -> torch.Tensor:
+        """
+        Forward function.
+        :param xs: torch.Tensor. input tensor. shape=(batch_size, max_length, dim).
+        :return: output tensor. shape=(batch_size, max_length, dim).
+        """
+        return self.w_2(self.dropout(self.activation(self.w_1(xs))))
+
+
+class MultiHeadedAttention(nn.Module):
+    def __init__(self, n_head: int, n_feat: int, dropout_rate: float):
+        """
+        :param n_head: int. the number of heads.
+        :param n_feat: int. the number of features.
+        :param dropout_rate: float. dropout rate.
+        """
+        super().__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+    def forward_qkv(self,
+                    query: torch.Tensor,
+                    key: torch.Tensor,
+                    value: torch.Tensor
+                    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        transform query, key and value.
+        :param query: torch.Tensor. query tensor. shape=(batch_size, time1, n_feat).
+        :param key: torch.Tensor. key tensor. shape=(batch_size, time2, n_feat).
+        :param value: torch.Tensor. value tensor. shape=(batch_size, time2, n_feat).
+        :return:
+        """
+        n_batch = query.size(0)
+        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
+        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
+        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
+        q = q.transpose(1, 2)  # (batch, head, time1, d_k)
+        k = k.transpose(1, 2)  # (batch, head, time2, d_k)
+        v = v.transpose(1, 2)  # (batch, head, time2, d_k)
+
+        return q, k, v
+
+    def forward_attention(self,
+                          value: torch.Tensor,
+                          scores: torch.Tensor,
+                          mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool)
+                          ) -> torch.Tensor:
+        """
+        compute attention context vector.
+        :param value: torch.Tensor. transformed value. shape=(batch_size, n_head, time2, d_k).
+        :param scores: torch.Tensor. attention score. shape=(batch_size, n_head, time1, time2).
+        :param mask: torch.Tensor. mask. shape=(batch_size, 1, time2) or
+                (batch_size, time1, time2), (0, 0, 0) means fake mask.
+        :return: torch.Tensor. transformed value. (batch_size, time1, d_model).
+                weighted by the attention score (batch_size, time1, time2).
+        """
+        n_batch = value.size(0)
+        # NOTE: When will `if mask.size(2) > 0` be True?
+        #   1. onnx(16/4) [WHY? Because we feed real cache & real mask for the
+        #           1st chunk to ease the onnx export.]
+        #   2. pytorch training
+        if mask.size(2) > 0:  # time2 > 0
+            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
+            # For last chunk, time2 might be larger than scores.size(-1)
+            mask = mask[:, :, :, :scores.size(-1)]  # (batch, 1, *, time2)
+            scores = scores.masked_fill(mask, -float('inf'))
+            attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)  # (batch, head, time1, time2)
+
+        # NOTE: When will `if mask.size(2) > 0` be False?
+        #   1. onnx(16/-1, -1/-1, 16/0)
+        #   2. jit (16/-1, -1/-1, 16/0, 16/4)
+        else:
+            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
+
+        p_attn = self.dropout(attn)
+        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
+        x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)  # (batch, time1, n_feat)
+
+        return self.linear_out(x)  # (batch, time1, n_feat)
+
+    def forward(self,
+                query: torch.Tensor,
+                key: torch.Tensor,
+                value: torch.Tensor,
+                mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+                cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+                **kwargs,
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        compute scaled dot product attention.
+        :param query: torch.Tensor. query tensor. shape=(batch_size, time1, n_feat).
+        :param key: torch.Tensor. key tensor. shape=(batch_size, time2, n_feat).
+        :param value: torch.Tensor. value tensor. shape=(batch_size, time2, n_feat).
+        :param mask: torch.Tensor. mask tensor (batch_size, 1, time2) or
+                (batch_size, time1, time2).
+        :param cache: torch.Tensor. cache tensor. shape=(1, head, cache_t, d_k * 2),
+                where `cache_t == chunk_size * num_decoding_left_chunks`
+                and `head * d_k == n_feat`
+        :return:
+        torch.Tensor. output tensor. shape=(batch_size, time1, n_feat).
+        torch.Tensor. cache tensor. (1, head, cache_t + time1, d_k * 2)
+                where `cache_t == chunk_size * num_decoding_left_chunks`
+                and `head * d_k == n_feat`
+        """
+        q, k, v = self.forward_qkv(query, key, value)
+
+        # NOTE:
+        #   when export onnx model, for 1st chunk, we feed
+        #       cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode)
+        #       or cache(1, head, real_cache_t, d_k * 2) (16/4 mode).
+        #       In all modes, `if cache.size(0) > 0` will alwayse be `True`
+        #       and we will always do splitting and
+        #       concatnation(this will simplify onnx export). Note that
+        #       it's OK to concat & split zero-shaped tensors(see code below).
+        #   when export jit  model, for 1st chunk, we always feed
+        #       cache(0, 0, 0, 0) since jit supports dynamic if-branch.
+        # >>> a = torch.ones((1, 2, 0, 4))
+        # >>> b = torch.ones((1, 2, 3, 4))
+        # >>> c = torch.cat((a, b), dim=2)
+        # >>> torch.equal(b, c)        # True
+        # >>> d = torch.split(a, 2, dim=-1)
+        # >>> torch.equal(d[0], d[1])  # True
+        if cache.size(0) > 0:
+            key_cache, value_cache = torch.split(
+                cache, cache.size(-1) // 2, dim=-1)
+            k = torch.cat([key_cache, k], dim=2)
+            v = torch.cat([value_cache, v], dim=2)
+        # NOTE: We do cache slicing in encoder.forward_chunk, since it's
+        #   non-trivial to calculate `next_cache_start` here.
+        new_cache = torch.cat((k, v), dim=-1)
+
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+        return self.forward_attention(v, scores, mask), new_cache
+
+
+class TransformerEncoderLayer(nn.Module):
+    def __init__(self,
+                 input_dim: int,
+                 dropout_rate: float = 0.1,
+                 n_heads: int = 4,
+                 ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(input_dim, eps=1e-5)
+        self.attention = MultiHeadedAttention(
+            n_head=n_heads,
+            n_feat=input_dim,
+            dropout_rate=dropout_rate
+        )
+
+        self.dropout1 = nn.Dropout(dropout_rate)
+        self.norm2 = nn.LayerNorm(input_dim, eps=1e-5)
+        self.ffn = PositionwiseFeedForward(
+            input_dim=input_dim,
+            hidden_units=input_dim,
+            dropout_rate=dropout_rate
+        )
+        self.dropout2 = nn.Dropout(dropout_rate)
+        self.norm3 = nn.LayerNorm(input_dim, eps=1e-5)
+
+    def forward(
+            self,
+            x: torch.Tensor,
+            mask: torch.Tensor,
+            position_embedding: torch.Tensor,
+            attention_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+
+        :param x: torch.Tensor. shape=(batch_size, time, input_dim).
+        :param mask: torch.Tensor. mask tensor for the input. shape=(batch_size, time，time).
+        :param position_embedding: torch.Tensor.
+        :param attention_cache: torch.Tensor. cache tensor of the KEY & VALUE
+                shape=(batch_size=1, head, cache_t1, d_k * 2), head * d_k == input_dim.
+        :return:
+                torch.Tensor: Output tensor (batch_size, time, input_dim).
+                torch.Tensor: att_cache tensor, (batch_size=1, head, cache_t1 + time, d_k * 2).
+        """
+
+        xt = self.norm1(x)
+
+        x_att, new_att_cache = self.attention.forward(
+            xt, xt, xt, mask=mask, cache=attention_cache, position_embedding=position_embedding
+        )
+        x = x + self.dropout1(xt)
+        xt = self.norm2(x)
+        xt = self.ffn.forward(xt)
+        x = x + self.dropout2(xt)
+
+        x = self.norm3(x)
+
+        return x, new_att_cache
+
+
+class TransformerEncoder(nn.Module):
+    """
+    https://github.com/wenet-e2e/wenet/blob/main/wenet/transformer/encoder.py#L364
+    """
+    def __init__(self,
+                 input_size: int = 64,
+                 hidden_size: int = 256,
+                 attention_heads: int = 4,
+                 num_blocks: int = 6,
+                 dropout_rate: float = 0.1,
+                 max_length: int = 512,
+                 chunk_size: int = 1,
+                 num_left_chunks: int = 128,
+                 ):
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+
+        self.max_length = max_length
+        self.chunk_size = chunk_size
+        self.num_left_chunks = num_left_chunks
+
+        self.input_linear = nn.Linear(
+            in_features=self.input_size,
+            out_features=self.hidden_size,
+        )
+
+        self.positional_encoding = RelPositionalEncoding(
+            embedding_dim=hidden_size,
+            dropout_rate=dropout_rate,
+            max_length=max_length,
+        )
+
+        self.encoder_layer_list = torch.nn.ModuleList([
+            TransformerEncoderLayer(
+                input_dim=hidden_size,
+                n_heads=attention_heads,
+                dropout_rate=dropout_rate,
+            ) for _ in range(num_blocks)
+        ])
+
+        self.output_linear = nn.Linear(
+            in_features=self.hidden_size,
+            out_features=self.input_size,
+        )
+
+    def forward(self,
+                xs: torch.Tensor,
+                ):
+        """
+        :param xs: Tensor, shape: [batch_size, time_steps, input_size]
+        :return: Tensor, shape: [batch_size, time_steps, input_size]
+        """
+        batch_size, time_steps, _ = xs.shape
+        # xs shape: [batch_size, time_steps, input_size]
+        xs = self.input_linear.forward(xs)
+        # xs shape: [batch_size, time_steps, hidden_size]
+
+        xs, position_embedding = self.positional_encoding.forward(xs)
+        # xs shape: [batch_size, time_steps, hidden_size]
+        # position_embedding shape: [1, time_steps, hidden_size]
+
+        chunk_masks = subsequent_chunk_mask(
+            size=time_steps,
+            chunk_size=self.chunk_size,
+            num_left_chunks=self.num_left_chunks
+        )
+        # chunk_masks shape: [1, time_steps, time_steps]
+        chunk_masks = torch.broadcast_to(chunk_masks, size=(batch_size, time_steps, time_steps))
+        # chunk_masks shape: [batch_size, time_steps, time_steps]
+
+        for encoder_layer in self.encoder_layer_list:
+            xs, _ = encoder_layer.forward(xs, chunk_masks, position_embedding)
+
+        # xs shape: [batch_size, time_steps, hidden_size]
+        xs = self.output_linear.forward(xs)
+        # xs shape: [batch_size, time_steps, input_size]
+
+        return xs
+
+    def forward_chunk(self,
+                      xs: torch.Tensor,
+                      offset: int,
+                      attention_mask: torch.Tensor = torch.zeros(0, 0, 0),
+                      attention_cache: torch.Tensor = torch.zeros(0, 0, 0, 0),
+                      ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Forward just one chunk.
+        :param xs: torch.Tensor. chunk input, with shape (b=1, time, mel-dim),
+                where `time == (chunk_size - 1) * subsample_rate + subsample.right_context + 1`
+        :param offset: int. current offset in encoder output timestamp.
+        :param attention_mask:
+        :param attention_cache: torch.Tensor. cache tensor for KEY & VALUE in
+                transformer/conformer attention, with shape
+                (elayers, head, cache_t1, d_k * 2), where
+                `head * d_k == hidden-dim` and
+                `cache_t1 == chunk_size * num_decoding_left_chunks`.
+        :return:
+        """
+        # xs shape: [batch_size, time_steps, input_size]
+        xs = self.input_linear.forward(xs)
+        # xs shape: [batch_size, time_steps, hidden_size]
+
+        xs, position_embedding = self.positional_encoding.forward(xs, offset=offset)
+        # xs shape: [batch_size, time_steps, hidden_size]
+        # position_embedding shape: [1, time_steps, hidden_size]
+
+        r_att_cache = []
+        for encoder_layer in self.encoder_layer_list:
+            xs, new_att_cache = encoder_layer.forward(
+                x=xs, mask=attention_mask,
+                position_embedding=position_embedding,
+                attention_cache=attention_cache,
+            )
+            r_att_cache.append(new_att_cache[:, :, self.chunk_size:, :])
+
+        r_att_cache = torch.cat(r_att_cache, dim=0)
+
+        return xs, r_att_cache
+
+    def forward_chunk_by_chunk(
+            self,
+            xs: torch.Tensor,
+    ) -> torch.Tensor:
+
+        batch_size, time_steps, _ = xs.shape
+
+        offset = 0
+        attention_cache: torch.Tensor = torch.zeros((0, 0, 0, 0), device=xs.device)
+        attention_mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool, device=xs.device)
+
+        outputs = []
+        for idx in range(0, time_steps - self.chunk_size + 1, self.chunk_size):
+            begin = idx * self.chunk_size
+            end = begin + self.chunk_size
+            chunk_xs = xs[:, begin:end, :]
+
+            ys, att_cache = self.forward_chunk(
+                xs=chunk_xs, attention_mask=attention_mask,
+                offset=offset, attention_cache=attention_cache
+            )
+            # xs shape: [batch_size, chunk_size, hidden_size]
+            ys = self.output_linear.forward(ys)
+            # xs shape: [batch_size, chunk_size, input_size]
+
+            offset += self.chunk_size
+            outputs.append(ys)
+
+        ys = torch.cat(outputs, 1)
+        return ys
+
+
+def main():
+
+    encoder = TransformerEncoder(
+        input_size=64,
+        hidden_size=256,
+        attention_heads=4,
+        num_blocks=6,
+        dropout_rate=0.1,
+    )
+
+    x = torch.ones([4, 200, 64])
+
+    y = encoder.forward(xs=x)
+    print(y.shape)
+
+    # y = encoder.forward_chunk_by_chunk(xs=x)
+    # print(y.shape)
+
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/nx_clean_unet/utils.py

@@ -0,0 +1,45 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import torch
+import torch.nn as nn
+
+
+class LearnableSigmoid1d(nn.Module):
+    def __init__(self, in_features, beta=1):
+        super().__init__()
+        self.beta = beta
+        self.slope = nn.Parameter(torch.ones(in_features))
+        self.slope.requiresGrad = True
+
+    def forward(self, x):
+        # x shape: [batch_size, time_steps, spec_bins]
+        return self.beta * torch.sigmoid(self.slope * x)
+
+
+def mag_pha_stft(y, n_fft, hop_size, win_size, compress_factor=1.0, center=True):
+
+    hann_window = torch.hann_window(win_size).to(y.device)
+    stft_spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window,
+                           center=center, pad_mode='reflect', normalized=False, return_complex=True)
+    stft_spec = torch.view_as_real(stft_spec)
+    mag = torch.sqrt(stft_spec.pow(2).sum(-1) + 1e-9)
+    pha = torch.atan2(stft_spec[:, :, :, 1] + 1e-10, stft_spec[:, :, :, 0] + 1e-5)
+    # Magnitude Compression
+    mag = torch.pow(mag, compress_factor)
+    com = torch.stack((mag*torch.cos(pha), mag*torch.sin(pha)), dim=-1)
+
+    return mag, pha, com
+
+
+def mag_pha_istft(mag, pha, n_fft, hop_size, win_size, compress_factor=1.0, center=True):
+    # Magnitude Decompression
+    mag = torch.pow(mag, (1.0/compress_factor))
+    com = torch.complex(mag*torch.cos(pha), mag*torch.sin(pha))
+    hann_window = torch.hann_window(win_size).to(com.device)
+    wav = torch.istft(com, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window, center=center)
+
+    return wav
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/nx_clean_unet/yaml/config.yaml

@@ -0,0 +1,23 @@
+model_name: "nx_clean_unet"
+
+sample_rate: 8000
+segment_size: 16000
+n_fft: 512
+win_size: 200
+hop_size: 80
+
+down_sampling_num_layers: 5
+down_sampling_in_channels: 1
+down_sampling_hidden_channels: 64
+down_sampling_kernel_size: 4
+down_sampling_stride: 2
+
+tsfm_hidden_size: 256
+tsfm_attention_heads: 4
+tsfm_num_blocks: 6
+tsfm_dropout_rate: 0.1
+
+discriminator_dim: 32
+discriminator_in_channel: 2
+
+compress_factor: 0.3