update

.gitignore

@@ -18,5 +18,5 @@
 /trained_models/
 /temp/
 
-#**/*.wav
+**/*.wav
 **/*.xlsx

examples/nx_mpnet/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_mpnet/step_1_prepare_data.py

@@ -0,0 +1,202 @@
+#!/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("--scale", default=1, type=float)
+
+    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):
+        flag = random.random()
+        if flag > args.scale:
+            continue
+
+        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_mpnet/step_2_train_model.py

@@ -0,0 +1,447 @@
+#!/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_mpnet.configuration_nx_mpnet import NXMPNetConfig
+from toolbox.torchaudio.models.nx_mpnet.discriminator import MetricDiscriminatorPretrainedModel
+from toolbox.torchaudio.models.nx_mpnet.modeling_nx_mpnet import NXMPNet, NXMPNetPretrainedModel
+from toolbox.torchaudio.models.nx_mpnet.utils import mag_pha_stft, mag_pha_istft
+from toolbox.torchaudio.models.nx_mpnet.metrics import run_batch_pesq, run_pesq_score
+from toolbox.torchaudio.models.nx_mpnet.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 = NXMPNetConfig.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 = NXMPNetPretrainedModel(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
+    stft_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_audio, noisy_audio = batch
+            clean_audio = clean_audio.to(device)
+            noisy_audio = noisy_audio.to(device)
+            one_labels = torch.ones(clean_audio.shape[0]).to(device)
+
+            clean_mag, clean_pha, clean_com = mag_pha_stft(clean_audio, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+            noisy_mag, noisy_pha, noisy_com = mag_pha_stft(noisy_audio, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+
+            mag_g, pha_g, com_g = generator.forward(noisy_mag, noisy_pha)
+
+            audio_g = mag_pha_istft(mag_g, pha_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+            mag_g_hat, pha_g_hat, com_g_hat = mag_pha_stft(audio_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+
+            audio_list_r, audio_list_g = list(clean_audio.cpu().numpy()), list(audio_g.detach().cpu().numpy())
+            pesq_score_list: List[float] = run_batch_pesq(audio_list_r, audio_list_g, sample_rate=config.sample_rate, mode="nb")
+
+            # Discriminator
+            optim_d.zero_grad()
+            metric_r = discriminator.forward(clean_mag, clean_mag)
+            metric_g = discriminator.forward(clean_mag, mag_g_hat.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
+            loss_stft = F.mse_loss(com_g, com_g_hat) * 2
+            # Time Loss
+            loss_time = F.l1_loss(clean_audio, audio_g)
+            # Metric Loss
+            metric_g = discriminator.forward(clean_mag, mag_g_hat)
+            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_stft * 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_stft_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_audio, noisy_audio = batch
+                clean_audio = clean_audio.to(device)
+                noisy_audio = noisy_audio.to(device)
+
+                clean_mag, clean_pha, clean_com = mag_pha_stft(clean_audio, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+                noisy_mag, noisy_pha, noisy_com = mag_pha_stft(noisy_audio, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+
+                mag_g, pha_g, com_g = generator.forward(noisy_mag, noisy_pha)
+
+                audio_g = mag_pha_istft(mag_g, pha_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+                mag_g_hat, pha_g_hat, com_g_hat = mag_pha_stft(audio_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+
+                clean_audio_list = torch.split(clean_audio, 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_stft_err += F.mse_loss(com_g, com_g_hat).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)
+                stft_err = round(total_stft_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,
+                    "stft_err": stft_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,
+            "stft_err": stft_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_mpnet/step_3_evaluation.py

@@ -0,0 +1,187 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/yxlu-0102/MP-SENet/blob/main/inference.py
+"""
+import argparse
+import logging
+import os
+from pathlib import Path
+import sys
+import uuid
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+import numpy as np
+import pandas as pd
+from scipy.io import wavfile
+import torch
+import torch.nn as nn
+import torchaudio
+from tqdm import tqdm
+
+from toolbox.torchaudio.models.mpnet.configuration_mpnet import MPNetConfig
+from toolbox.torchaudio.models.mpnet.modeling_mpnet import MPNetPretrainedModel
+from toolbox.torchaudio.models.mpnet.utils import mag_pha_stft, mag_pha_istft
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
+    parser.add_argument("--model_dir", default="serialization_dir/best", type=str)
+    parser.add_argument("--evaluation_audio_dir", default="evaluation_audio_dir", type=str)
+
+    parser.add_argument("--limit", default=10, type=int)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config():
+    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)
+    stream_handler = logging.StreamHandler()
+    stream_handler.setLevel(logging.INFO)
+    stream_handler.setFormatter(logging.Formatter(fmt))
+
+    logger = logging.getLogger(__name__)
+
+    return logger
+
+
+def mix_speech_and_noise(speech: np.ndarray, noise: np.ndarray, snr_db: float):
+    l1 = len(speech)
+    l2 = len(noise)
+    l = min(l1, l2)
+    speech = speech[:l]
+    noise = noise[:l]
+
+    # np.float32, value between (-1, 1).
+
+    speech_power = np.mean(np.square(speech))
+    noise_power = speech_power / (10 ** (snr_db / 10))
+
+    noise_adjusted = np.sqrt(noise_power) * noise / np.sqrt(np.mean(noise ** 2))
+
+    noisy_signal = speech + noise_adjusted
+
+    return noisy_signal
+
+
+def save_audios(noise_audio: torch.Tensor,
+                clean_audio: torch.Tensor,
+                noisy_audio: torch.Tensor,
+                enhanced_audio: torch.Tensor,
+                output_dir: str,
+                sample_rate: int = 8000,
+                ):
+    basename = uuid.uuid4().__str__()
+    output_dir = Path(output_dir) / basename
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    filename = output_dir / "noise_audio.wav"
+    torchaudio.save(filename, noise_audio.detach().cpu(), sample_rate, bits_per_sample=16)
+    filename = output_dir / "clean_audio.wav"
+    torchaudio.save(filename, clean_audio.detach().cpu(), sample_rate, bits_per_sample=16)
+    filename = output_dir / "noisy_audio.wav"
+    torchaudio.save(filename, noisy_audio.detach().cpu(), sample_rate, bits_per_sample=16)
+
+    filename = output_dir / "enhanced_audio.wav"
+    torchaudio.save(filename, enhanced_audio.detach().cpu(), sample_rate, bits_per_sample=16)
+
+    return output_dir.as_posix()
+
+
+def main():
+    args = get_args()
+
+    logger = logging_config()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    logger.info("GPU available count: {}; device: {}".format(n_gpu, device))
+
+    logger.info("prepare model")
+    config = MPNetConfig.from_pretrained(
+        pretrained_model_name_or_path=args.model_dir,
+    )
+    generator = MPNetPretrainedModel.from_pretrained(
+        pretrained_model_name_or_path=args.model_dir,
+    )
+    generator.to(device)
+    generator.eval()
+
+    logger.info("read excel")
+    df = pd.read_excel(args.valid_dataset)
+
+    progress_bar = tqdm(total=len(df), desc="Evaluation")
+    for idx, row in df.iterrows():
+        noise_filename = row["noise_filename"]
+        noise_offset = row["noise_offset"]
+        noise_duration = row["noise_duration"]
+
+        speech_filename = row["speech_filename"]
+        speech_offset = row["speech_offset"]
+        speech_duration = row["speech_duration"]
+
+        snr_db = row["snr_db"]
+
+        noise_audio, _ = librosa.load(
+            noise_filename,
+            sr=8000,
+            offset=noise_offset,
+            duration=noise_duration,
+        )
+        clean_audio, _ = librosa.load(
+            speech_filename,
+            sr=8000,
+            offset=speech_offset,
+            duration=speech_duration,
+        )
+        noisy_audio: np.ndarray = mix_speech_and_noise(
+            speech=clean_audio,
+            noise=noise_audio,
+            snr_db=snr_db,
+        )
+        noise_audio = torch.tensor(noise_audio, dtype=torch.float32)
+        clean_audio = torch.tensor(clean_audio, dtype=torch.float32)
+        noisy_audio: torch.Tensor = torch.tensor(noisy_audio, dtype=torch.float32)
+
+        noise_audio = noise_audio.unsqueeze(dim=0)
+        clean_audio = clean_audio.unsqueeze(dim=0)
+        noisy_audio: torch.Tensor = noisy_audio.unsqueeze(dim=0)
+
+        # inference
+        clean_audio = clean_audio.to(device)
+        noisy_audio = noisy_audio.to(device)
+        with torch.no_grad():
+            noisy_mag, noisy_pha, noisy_com = mag_pha_stft(
+                noisy_audio, config.n_fft, config.hop_size, config.win_size, config.compress_factor
+            )
+            mag_g, pha_g, com_g = generator.forward(noisy_mag, noisy_pha)
+            audio_g = mag_pha_istft(
+                mag_g, pha_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor
+            )
+            enhanced_audio = audio_g.detach()
+
+        save_audios(
+            noise_audio, clean_audio, noisy_audio,
+            enhanced_audio,
+            args.evaluation_audio_dir
+        )
+
+        progress_bar.update(1)
+
+        if idx > args.limit:
+            break
+
+    return
+
+
+if __name__ == '__main__':
+    main()

examples/nx_mpnet/yaml/config.yaml

@@ -0,0 +1,27 @@
+model_name: "mpnet"
+
+num_gpus: 0
+batch_size: 3
+learning_rate: 0.0005
+adam_b1: 0.8
+adam_b2: 0.99
+lr_decay: 0.99
+seed: 1234
+
+dense_channel: 64
+compress_factor: 0.3
+num_tsconformers: 4
+beta: 2.0
+
+sample_rate: 8000
+segment_size: 16000
+n_fft: 512
+hop_size: 80
+win_size: 200
+
+num_workers: 4
+
+dist_config:
+  dist_backend: nccl
+  dist_url: tcp://localhost:54321
+  world_size: 1

toolbox/torchaudio/models/mpnet/inference_mpnet.py

@@ -84,16 +84,18 @@ class InferenceMPNet(object):
         enhanced_audio = enhanced_audio[0]
         return enhanced_audio
 
+
 def main():
-    model_zip_file = project_path / "trained_models/mpnet_aishell_20250221.zip"
+    model_zip_file = project_path / "trained_models/mpnet-aishell-1-epoch.zip"
     infer_mpnet = InferenceMPNet(model_zip_file)
 
     sample_rate = 8000
-    noisy_audio_file = project_path / "data/examples/noisy_audio.wav"
+    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 = noisy_audio[int(7*sample_rate):int(9*sample_rate)]
     noisy_audio = torch.tensor(noisy_audio, dtype=torch.float32)
     noisy_audio = noisy_audio.unsqueeze(dim=0)
 

toolbox/torchaudio/models/nx_denoise/causal_convolution/causal_conv2d.py

@@ -84,7 +84,7 @@ class CausalConv2d(nn.Module):
 
     def forward(self,
                 inputs: torch.Tensor,
-                causal_cache: torch.Tensor = None,
+                causal_cache: List[torch.Tensor] = None,
                 ):
 
         if causal_cache is None:
@@ -97,6 +97,8 @@ class CausalConv2d(nn.Module):
         # x shape: [batch_size, 1, time_steps2, hidden_size]
         # time_steps2 = time_steps + self.causal_left_pad + self.causal_right_pad
 
+        causal_cache = x[:, :, -self.causal_left_pad:, :]
+
         x = self.conv1.forward(x)
         # inputs shape: [batch_size, 1, time_steps, hidden_size]
 
@@ -108,8 +110,6 @@ class CausalConv2d(nn.Module):
         if self.activation:
             x = self.activation(x)
 
-        causal_cache = x[:, :, -self.causal_left_pad:, :]
-
         # inputs shape: [batch_size, 1, time_steps, hidden_size]
         return x, causal_cache
 
@@ -187,19 +187,19 @@ class CausalConv2dEncoder(nn.Module):
 
     def forward_chunk(self,
                       chunk: torch.Tensor,
-                      causal_cache: torch.Tensor = None,
+                      causal_cache: List[torch.Tensor] = None,
                       ):
-        # causal_cache shape: [self.num_layers, 1, causal_left_pad, hidden_size]
+        # causal_cache shape: [self.num_layers, batch_size, 1, causal_left_pad, hidden_size]
 
-        new_causal_cache_list = list()
+        new_causal_cache_list: List[torch.Tensor] = list()
         for idx, causal_conv in enumerate(self.causal_conv_list):
             chunk, new_causal_cache = causal_conv.forward(
-                inputs=chunk, causal_cache=causal_cache[idx: idx+1] if causal_cache is not None else None
+                inputs=chunk, causal_cache=causal_cache[idx] if causal_cache is not None else None
             )
+            # print(f"idx: {idx}, new_causal_cache: {new_causal_cache.shape}")
             new_causal_cache_list.append(new_causal_cache)
 
-        new_causal_cache = torch.cat(new_causal_cache_list, dim=0)
-        return chunk, new_causal_cache
+        return chunk, new_causal_cache_list
 
     def forward_chunk_by_chunk(self, inputs: torch.Tensor):
         # inputs shape: [batch_size, 1, time_steps, hidden_size]
@@ -207,7 +207,7 @@ class CausalConv2dEncoder(nn.Module):
 
         batch_size, channels, time_steps, hidden_size = inputs.shape
 
-        causal_cache = None
+        new_causal_cache_list: List[torch.Tensor] = None
 
         outputs = []
         for idx in range(0, time_steps, 1):
@@ -215,9 +215,9 @@ class CausalConv2dEncoder(nn.Module):
             end = begin + self.total_causal_right_pad + 1
             chunk_xs = inputs[:, :, begin:end, :]
 
-            ys, attention_cache = self.forward_chunk(
+            ys, new_causal_cache_list = self.forward_chunk(
                 chunk=chunk_xs,
-                causal_cache=causal_cache,
+                causal_cache=new_causal_cache_list,
             )
             # ys shape: [batch_size, channels, self.total_causal_right_pad + 1 , hidden_size]
             ys = ys[:, :, :1, :]

toolbox/torchaudio/models/nx_denoise/inference_nx_denoise.py

@@ -0,0 +1,97 @@
+#!/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.nx_denoise.configuration_nx_denoise import NXDenoiseConfig
+from toolbox.torchaudio.models.nx_denoise.modeling_nx_denoise import NXDenoisePretrainedModel, MODEL_FILE
+
+logger = logging.getLogger("toolbox")
+
+
+class InferenceNXDenoise(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 = NXDenoiseConfig.from_pretrained(
+            pretrained_model_name_or_path=model_path.as_posix(),
+        )
+        model = NXDenoisePretrainedModel.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_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_chunk_by_chunk(noisy_audios)
+            enhanced_audios = self.model.forward(noisy_audios)
+            # enhanced_audio shape: [batch_size, n_samples]
+            # enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
+
+        enhanced_audio = enhanced_audios[0]
+        # enhanced_audio shape: [num_samples,]
+        return enhanced_audio
+
+
+def main():
+    model_zip_file = project_path / "trained_models/nx-denoise.zip"
+    runtime = InferenceNXDenoise(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 = noisy_audio[int(7*sample_rate):int(9*sample_rate)]
+    noisy_audio = torch.tensor(noisy_audio, dtype=torch.float32)
+    noisy_audio = noisy_audio.unsqueeze(dim=0)
+
+    enhanced_audio = runtime.enhancement_by_tensor(noisy_audio)
+
+    filename = "enhanced_audio.wav"
+    torchaudio.save(filename, enhanced_audio.detach().cpu().unsqueeze(dim=0), sample_rate)
+
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/nx_denoise/modeling_nx_denoise.py

@@ -268,6 +268,58 @@ class NXDenoise(nn.Module):
         return enhanced_audios
 
 
+    def forward_chunk_by_chunk(self, noisy_audios: torch.Tensor):
+        # noisy_audios shape: [batch_size, n_samples]
+        noisy_audios = torch.unsqueeze(noisy_audios, dim=1)
+        # 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)
+
+        # down sampling
+        bottle_neck, skip_connection_list = 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, channels]
+        bottle_neck = torch.unsqueeze(bottle_neck, dim=1)
+        # bottle_neck shape: [batch_size, 1, time_steps, freq_dim]
+
+        # causal conv in
+        bottle_neck = self.causal_conv_in.forward_chunk_by_chunk(bottle_neck)
+        # bottle_neck shape: [batch_size, channels, time_steps, freq_dim]
+
+        # ts transformer
+        # bottle_neck shape: [batch_size, channels, time_steps, freq_dim]
+        bottle_neck = self.ts_transformer.forward_chunk_by_chunk(bottle_neck)
+        # bottle_neck shape: [batch_size, channels, time_steps, freq_dim]
+
+        # causal conv out
+        bottle_neck = self.causal_conv_out.forward_chunk_by_chunk(bottle_neck)
+        # bottle_neck shape: [batch_size, 1, time_steps, freq_dim]
+
+        # up sampling
+        bottle_neck = torch.squeeze(bottle_neck, dim=1)
+        # bottle_neck shape: [batch_size, time_steps, channels]
+        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, skip_connection_list)
+
+        enhanced_audios = enhanced_audios[:, :, :n_samples]
+        # enhanced_audios shape: [batch_size, 1, n_samples]
+
+        enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
+        # enhanced_audios shape: [batch_size, n_samples]
+
+        return enhanced_audios
+
+
 MODEL_FILE = "generator.pt"
 
 

toolbox/torchaudio/models/nx_denoise/stftnet/__init__.py

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

toolbox/torchaudio/models/nx_denoise/stftnet/istftnet.py

@@ -0,0 +1,9 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://arxiv.org/abs/2203.02395
+"""
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/nx_denoise/stftnet/stfnets.py

@@ -0,0 +1,9 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://arxiv.org/abs/1902.07849
+"""
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/nx_mpnet/__init__.py

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

toolbox/torchaudio/models/nx_mpnet/causal_convolution/__init__.py

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

toolbox/torchaudio/models/nx_mpnet/causal_convolution/causal_conv2d.py

@@ -0,0 +1,445 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from typing import List, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from toolbox.torchaudio.models.nx_mpnet.utils import LearnableSigmoid2d
+
+
+class SPConvTranspose2d(nn.Module):
+    def __init__(self,
+                 in_channels: int,
+                 out_channels: int,
+                 kernel_size: Union[int, Tuple[int]],
+                 r=1
+                 ):
+        super(SPConvTranspose2d, self).__init__()
+        self.pad_freq = nn.ConstantPad2d((1, 1, 0, 0), value=0.)
+        self.out_channels = out_channels
+        self.conv = nn.Conv2d(in_channels, out_channels * r, kernel_size=kernel_size, stride=(1, 1))
+        self.r = r
+
+    def forward(self, x: torch.Tensor):
+        x = self.pad_freq(x)
+        out = self.conv(x)
+
+        b, c, t, f = out.shape
+
+        out = out.view((b, self.r, c // self.r, t, f))
+        out = out.permute(0, 2, 3, 4, 1)
+        out = out.contiguous().view((b, c // self.r, t, -1))
+        return out
+
+
+class CausalConv2dBlock(nn.Module):
+    def __init__(self,
+                 in_channels: int,
+                 out_channels: int,
+                 dilation: int,
+                 kernel_size: Tuple[int, int] = (2, 3),
+                 ):
+        super(CausalConv2dBlock, self).__init__()
+        self.pad_length = dilation
+
+        self.pad_time = nn.ConstantPad2d((0, 0, self.pad_length, 0), value=0.)
+        self.pad_freq = nn.ConstantPad2d((1, 1, 0, 0), value=0.)
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, dilation=(dilation, 1))
+        self.norm = nn.InstanceNorm2d(out_channels, affine=True)
+        self.activation = nn.PReLU(out_channels)
+
+    def forward(self,
+                x: torch.Tensor,
+                cache_pad: torch.Tensor = None,
+                ):
+        """
+
+        :param x: Tensor, shape: [batch_size, channels, time_steps, dim]
+        :param cache_pad:
+        :return:
+        """
+        if cache_pad is None:
+            x = self.pad_time(x)
+        else:
+            x = torch.concat(tensors=[cache_pad, x], dim=2)
+        new_cache_pad = x[:, :, -self.pad_length:, :]
+
+        x = self.pad_freq(x)
+
+        x = self.conv(x)
+        x = self.norm(x)
+        x = self.activation(x)
+        return x, new_cache_pad
+
+
+class CausalConv2dEncoder(nn.Module):
+    def __init__(self,
+                 num_blocks: int,
+                 hidden_size: int,
+                 ):
+        super(CausalConv2dEncoder, self).__init__()
+        self.num_blocks = num_blocks
+
+        self.blocks: List[CausalConv2dBlock] = nn.ModuleList([])
+        for idx in range(num_blocks):
+            in_channels = hidden_size * (idx+1)
+            dilation = 2 ** idx
+            block = CausalConv2dBlock(
+                in_channels=in_channels,
+                out_channels=hidden_size,
+                dilation=dilation,
+                kernel_size=(2, 3),
+            )
+            self.blocks.append(block)
+
+    def forward(self,
+                x: torch.Tensor,
+                cache_pad_list: List[torch.Tensor] = None,
+                ):
+        """
+        :param x: Tensor, shape: [batch_size, channels, time_steps, dim].
+        :param cache_pad_list: List[Tensor]
+        :return:
+        """
+        new_cache_pad_list = list()
+
+        skip = x
+        for idx, block in enumerate(self.blocks):
+            x, new_cache_pad = block.forward(
+                skip,
+                cache_pad=None if cache_pad_list is None else cache_pad_list[idx]
+            )
+            new_cache_pad_list.append(new_cache_pad)
+            skip = torch.cat([x, skip], dim=1)
+            # x shape: [batch_size, channels, time_steps, dim].
+        return x, new_cache_pad_list
+
+    def forward_chunk(self,
+                      chunk: torch.Tensor,
+                      cache_pad_list: List[torch.Tensor] = None,
+                      ):
+        return self.forward(chunk, cache_pad_list)
+
+    def forward_chunk_by_chunk(self,
+                               x: torch.Tensor,
+                               ):
+        """
+        :param x: Tensor, shape: [batch_size, channels, time_steps, dim].
+        :return:
+        """
+        batch_size, channels, time_steps, _ = x.shape
+
+        cache_pad_list = None
+
+        outputs = list()
+        for idx in range(time_steps):
+            chunk = x[:, :, idx:idx+1, :]
+
+            y, cache_pad_list = self.forward_chunk(chunk, cache_pad_list=cache_pad_list)
+            outputs.append(y)
+
+        outputs = torch.concat(outputs, dim=2)
+        return outputs
+
+
+class DenseEncoder(nn.Module):
+    def __init__(self,
+                 num_blocks: int,
+                 in_channels: int,
+                 out_channels: int,
+                 ):
+        super(DenseEncoder, self).__init__()
+        self.dense_conv_1 = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, (1, 1)),
+            nn.InstanceNorm2d(out_channels, affine=True),
+            nn.PReLU(out_channels)
+        )
+        self.dense_block = CausalConv2dEncoder(
+            num_blocks=num_blocks, hidden_size=out_channels,
+        )
+        self.dense_conv_2 = nn.Sequential(
+            nn.Conv2d(out_channels, out_channels, (1, 3), (1, 2), padding=(0, 1)),
+            nn.InstanceNorm2d(out_channels, affine=True),
+            nn.PReLU(out_channels)
+        )
+
+    def forward(self,
+                x: torch.Tensor,
+                ):
+        """
+        :param x: Tensor, shape: [batch_size, channels, time_steps, dim]
+        :return:
+        """
+        x = self.dense_conv_1(x)
+        x, _ = self.dense_block.forward(x)
+        x = self.dense_conv_2(x)
+        # x shape: [b, c, t, f//2]
+        return x
+
+    def forward_chunk(self,
+                      x: torch.Tensor,
+                      cache_pad_list: List[torch.Tensor] = None,
+                      ):
+        x = self.dense_conv_1(x)
+        x, new_cache_pad_list = self.dense_block.forward(x, cache_pad_list)
+        x = self.dense_conv_2(x)
+        # x shape: [b, c, t, f//2]
+        return x, new_cache_pad_list
+
+    def forward_chunk_by_chunk(self,
+                               x: torch.Tensor,
+                               ):
+        """
+        :param x: Tensor, shape: [batch_size, channels, time_steps, dim].
+        :return:
+        """
+        batch_size, channels, time_steps, _ = x.shape
+
+        cache_pad_list = None
+
+        outputs = list()
+        for idx in range(time_steps):
+            chunk = x[:, :, idx:idx+1, :]
+
+            y, cache_pad_list = self.forward_chunk(chunk, cache_pad_list=cache_pad_list)
+            outputs.append(y)
+
+        outputs = torch.concat(outputs, dim=2)
+        return outputs
+
+
+class MaskDecoder(nn.Module):
+    def __init__(self,
+                 num_blocks: int,
+                 hidden_size: int,
+                 out_channels: int = 1,
+                 beta: float = 2.0,
+                 n_fft: int = 512,
+                 ):
+        super(MaskDecoder, self).__init__()
+        self.dense_block = CausalConv2dEncoder(
+            num_blocks=num_blocks, hidden_size=hidden_size,
+        )
+        self.mask_conv = nn.Sequential(
+            SPConvTranspose2d(hidden_size, hidden_size, (1, 3), 2),
+            nn.InstanceNorm2d(hidden_size, affine=True),
+            nn.PReLU(hidden_size),
+            nn.Conv2d(hidden_size, out_channels, (1, 2))
+        )
+        self.lsigmoid = LearnableSigmoid2d(n_fft//2+1, beta=beta)
+
+    def forward(self,
+                x: torch.Tensor,
+                ):
+        """
+
+        :param x: Tensor, shape: [batch_size, channels, time_steps, dim]
+        :return:
+        """
+        x, _ = self.dense_block(x)
+        x = self.mask_conv(x)
+        # x shape: [batch_size, 1, time_steps, dim*2-1]
+        x = x.permute(0, 3, 2, 1).squeeze(-1)
+        # x shape: [b, f, t]
+        x = self.lsigmoid(x)
+        return x
+
+    def forward_chunk(self,
+                      x: torch.Tensor,
+                      cache_pad_list: List[torch.Tensor] = None,
+                      ):
+        x, new_cache_pad_list = self.dense_block(x, cache_pad_list)
+        x = self.mask_conv(x)
+        # x shape: [batch_size, 1, time_steps, dim*2-1]
+        x = x.permute(0, 3, 2, 1).squeeze(-1)
+        # x shape: [b, f, t]
+        x = self.lsigmoid(x)
+        return x, new_cache_pad_list
+
+    def forward_chunk_by_chunk(self,
+                               x: torch.Tensor,
+                               ):
+        """
+        :param x: Tensor, shape: [batch_size, channels, time_steps, dim].
+        :return:
+        """
+        batch_size, channels, time_steps, _ = x.shape
+
+        cache_pad_list = None
+
+        outputs = list()
+        for idx in range(time_steps):
+            chunk = x[:, :, idx:idx+1, :]
+
+            y, cache_pad_list = self.forward_chunk(chunk, cache_pad_list=cache_pad_list)
+            outputs.append(y)
+
+        outputs = torch.concat(outputs, dim=2)
+        return outputs
+
+
+class PhaseDecoder(nn.Module):
+    def __init__(self,
+                 num_blocks: int,
+                 hidden_size: int,
+                 out_channels: int = 1,
+                 ):
+        super(PhaseDecoder, self).__init__()
+        self.dense_block = CausalConv2dEncoder(
+            num_blocks=num_blocks, hidden_size=hidden_size,
+        )
+
+        self.phase_conv = nn.Sequential(
+            SPConvTranspose2d(hidden_size, hidden_size, (1, 3), 2),
+            nn.InstanceNorm2d(hidden_size, affine=True),
+            nn.PReLU(hidden_size)
+        )
+        self.phase_conv_r = nn.Conv2d(hidden_size, out_channels, (1, 2))
+        self.phase_conv_i = nn.Conv2d(hidden_size, out_channels, (1, 2))
+
+    def forward(self,
+                x: torch.Tensor,
+                ):
+        """
+        :param x: Tensor, shape: [batch_size, channels, time_steps, dim]
+        :return:
+        """
+        x, _ = self.dense_block(x)
+
+        x = self.phase_conv(x)
+        x_r = self.phase_conv_r(x)
+        x_i = self.phase_conv_i(x)
+        x = torch.atan2(x_i, x_r)
+        x = x.permute(0, 3, 2, 1).squeeze(-1)
+        # x shape: [b, f, t]
+        return x
+
+    def forward_chunk(self,
+                      x: torch.Tensor,
+                      cache_pad_list: List[torch.Tensor] = None,
+                      ):
+        x, new_cache_pad_list = self.dense_block(x, cache_pad_list)
+
+        x = self.phase_conv(x)
+        x_r = self.phase_conv_r(x)
+        x_i = self.phase_conv_i(x)
+        x = torch.atan2(x_i, x_r)
+        x = x.permute(0, 3, 2, 1).squeeze(-1)
+        # x shape: [b, f, t]
+        return x, new_cache_pad_list
+
+    def forward_chunk_by_chunk(self,
+                               x: torch.Tensor,
+                               ):
+        """
+        :param x: Tensor, shape: [batch_size, channels, time_steps, dim].
+        :return:
+        """
+        batch_size, channels, time_steps, _ = x.shape
+
+        cache_pad_list = None
+
+        outputs = list()
+        for idx in range(time_steps):
+            chunk = x[:, :, idx:idx+1, :]
+
+            y, cache_pad_list = self.forward_chunk(chunk, cache_pad_list=cache_pad_list)
+            outputs.append(y)
+
+        outputs = torch.concat(outputs, dim=2)
+        return outputs
+
+
+def main1():
+
+    encoder = CausalConv2dEncoder(
+        num_blocks=3, hidden_size=8,
+    )
+
+    # x shape: [batch_size, channels, time_steps, dim]
+    x = torch.rand(size=(1, 8, 200, 32))
+    x, new_cache_pad_list = encoder.forward(x)
+    print(x.shape)
+    for new_cache_pad in new_cache_pad_list:
+        print(new_cache_pad.shape)
+
+    x = torch.rand(size=(1, 8, 200, 32))
+    x = encoder.forward_chunk_by_chunk(x)
+    print(x.shape)
+
+    return
+
+
+def main2():
+
+    encoder = DenseEncoder(
+        num_blocks=3, in_channels=8, out_channels=8
+    )
+
+    # x shape: [batch_size, channels, time_steps, dim]
+    x = torch.rand(size=(1, 8, 200, 32))
+    x, new_cache_pad_list = encoder.forward(x)
+    print(x.shape)
+    for new_cache_pad in new_cache_pad_list:
+        print(new_cache_pad.shape)
+
+    x = torch.rand(size=(1, 8, 200, 32))
+    x = encoder.forward_chunk_by_chunk(x)
+    print(x.shape)
+
+    return
+
+
+def main3():
+
+    encoder = MaskDecoder(
+        num_blocks=3, hidden_size=64, out_channels=1,
+        n_fft=512,
+    )
+
+    # 512 // 2 + 1 = 257
+    # 129 * 2 - 1 = 257
+    # 257 // 2 + 1 = 129
+
+    # x shape: [batch_size, channels, time_steps, dim]
+    x = torch.rand(size=(1, 64, 201, 129))
+    x, new_cache_pad_list = encoder.forward(x)
+    print(x.shape)
+    for new_cache_pad in new_cache_pad_list:
+        print(new_cache_pad.shape)
+
+    x = torch.rand(size=(1, 64, 201, 129))
+    x = encoder.forward_chunk_by_chunk(x)
+    print(x.shape)
+
+    return
+
+
+
+def main():
+
+    encoder = PhaseDecoder(
+        num_blocks=3, hidden_size=64, out_channels=1,
+    )
+
+    # 512 // 2 + 1 = 257
+    # 129 * 2 - 1 = 257
+    # 257 // 2 + 1 = 129
+
+    # x shape: [batch_size, channels, time_steps, dim]
+    x = torch.rand(size=(1, 64, 201, 129))
+    x, new_cache_pad_list = encoder.forward(x)
+    print(x.shape)
+    for new_cache_pad in new_cache_pad_list:
+        print(new_cache_pad.shape)
+
+    x = torch.rand(size=(1, 64, 201, 129))
+    x = encoder.forward_chunk_by_chunk(x)
+    print(x.shape)
+
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/nx_mpnet/configuration_nx_mpnet.py

@@ -0,0 +1,90 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from toolbox.torchaudio.configuration_utils import PretrainedConfig
+
+
+class NXMPNetConfig(PretrainedConfig):
+    """
+    https://github.com/yxlu-0102/MP-SENet/blob/main/config.json
+    """
+    def __init__(self,
+                 sample_rate: int = 8000,
+                 segment_size: int = 16000,
+                 n_fft: int = 512,
+                 win_length: int = 200,
+                 hop_length: int = 80,
+
+                 dense_num_blocks: int = 4,
+                 dense_hidden_size: int = 64,
+
+                 mask_num_blocks: int = 4,
+                 mask_hidden_size: int = 64,
+
+                 phase_num_blocks: int = 4,
+                 phase_hidden_size: int = 64,
+
+                 tsfm_hidden_size: int = 64,
+                 tsfm_attention_heads: int = 4,
+                 tsfm_num_blocks: int = 4,
+                 tsfm_dropout_rate: float = 0.0,
+                 tsfm_max_time_relative_position: int = 1024,
+                 tsfm_max_freq_relative_position: int = 128,
+                 tsfm_chunk_size: int = 4,
+                 tsfm_num_left_chunks: int = 128,
+                 tsfm_num_right_chunks: int = 2,
+
+                 discriminator_dim: int = 32,
+                 discriminator_in_channel: int = 2,
+
+                 compress_factor: float = 0.3,
+
+                 batch_size: int = 4,
+                 learning_rate: float = 0.0005,
+                 adam_b1: float = 0.8,
+                 adam_b2: float = 0.99,
+                 lr_decay: float = 0.99,
+                 seed: int = 1234,
+
+                 **kwargs
+                 ):
+        super(NXMPNetConfig, self).__init__(**kwargs)
+        self.sample_rate = sample_rate
+        self.segment_size = segment_size
+        self.n_fft = n_fft
+        self.win_length = win_length
+        self.hop_length = hop_length
+
+        self.dense_num_blocks = dense_num_blocks
+        self.dense_hidden_size = dense_hidden_size
+
+        self.mask_num_blocks = mask_num_blocks
+        self.mask_hidden_size = mask_hidden_size
+
+        self.phase_num_blocks = phase_num_blocks
+        self.phase_hidden_size = phase_hidden_size
+
+        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.tsfm_max_time_relative_position = tsfm_max_time_relative_position
+        self.tsfm_max_freq_relative_position = tsfm_max_freq_relative_position
+        self.tsfm_chunk_size = tsfm_chunk_size
+        self.tsfm_num_left_chunks = tsfm_num_left_chunks
+        self.tsfm_num_right_chunks = tsfm_num_right_chunks
+
+        self.discriminator_dim = discriminator_dim
+        self.discriminator_in_channel = discriminator_in_channel
+
+        self.compress_factor = compress_factor
+
+        self.batch_size = batch_size
+        self.learning_rate = learning_rate
+        self.adam_b1 = adam_b1
+        self.adam_b2 = adam_b2
+        self.lr_decay = lr_decay
+        self.seed = seed
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/nx_mpnet/discriminator.py

@@ -0,0 +1,132 @@
+#!/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_mpnet.configuration_nx_mpnet import NXMPNetConfig
+from toolbox.torchaudio.models.nx_mpnet.utils import LearnableSigmoid1d
+
+
+class MetricDiscriminator(nn.Module):
+    def __init__(self, config: NXMPNetConfig):
+        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.transform = torchaudio.transforms.Spectrogram(
+            n_fft=self.n_fft,
+            win_length=self.win_length,
+            hop_length=self.hop_length,
+            power=1.0,
+            window_fn=torch.hann_window,
+            # window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
+        )
+
+        self.layers = nn.Sequential(
+            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):
+        x = self.transform.forward(x)
+        y = self.transform.forward(y)
+
+        xy = torch.stack((x, y), dim=1)
+        return self.layers(xy)
+
+
+MODEL_FILE = "discriminator.pt"
+
+
+class MetricDiscriminatorPretrainedModel(MetricDiscriminator):
+    def __init__(self,
+                 config: NXMPNetConfig,
+                 ):
+        super(MetricDiscriminatorPretrainedModel, self).__init__(
+            config=config,
+        )
+        self.config = config
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        config = NXMPNetConfig.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 = NXMPNetConfig()
+    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)])
+    x = torch.ones([4, 16000])
+    y = torch.ones([4, 16000])
+
+    output = discriminator.forward(x, y)
+    print(output.shape)
+    print(output)
+
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/nx_mpnet/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_mpnet/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_mpnet/modeling_nx_mpnet.py

@@ -0,0 +1,143 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import os
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from toolbox.torchaudio.configuration_utils import CONFIG_FILE
+from toolbox.torchaudio.models.nx_mpnet.causal_convolution.causal_conv2d import DenseEncoder, MaskDecoder, PhaseDecoder
+from toolbox.torchaudio.models.nx_mpnet.transformers.transformers import TSTransformerEncoder
+from toolbox.torchaudio.models.nx_mpnet.configuration_nx_mpnet import NXMPNetConfig
+
+
+class NXMPNet(nn.Module):
+    def __init__(self,
+                 config: NXMPNetConfig,
+                 ):
+        super(NXMPNet, self).__init__()
+        self.dense_encoder = DenseEncoder(
+            num_blocks=config.dense_num_blocks,
+            in_channels=2,
+            out_channels=config.dense_hidden_size,
+        )
+        self.ts_transformer = TSTransformerEncoder(
+            input_size=config.dense_hidden_size,
+            hidden_size=config.tsfm_hidden_size,
+            attention_heads=config.tsfm_attention_heads,
+            num_blocks=config.tsfm_num_blocks,
+            dropout_rate=config.tsfm_dropout_rate,
+            max_time_relative_position=config.tsfm_max_time_relative_position,
+            max_freq_relative_position=config.tsfm_max_freq_relative_position,
+            chunk_size=config.tsfm_chunk_size,
+            num_left_chunks=config.tsfm_num_left_chunks,
+            num_right_chunks=config.tsfm_num_right_chunks,
+        )
+        self.mask_decoder = MaskDecoder(
+            num_blocks=config.mask_num_blocks,
+            hidden_size=config.mask_hidden_size,
+            out_channels=1,
+            n_fft=config.n_fft,
+        )
+        self.phase_decoder = PhaseDecoder(
+            num_blocks=config.phase_num_blocks,
+            hidden_size=config.phase_hidden_size,
+            out_channels=1,
+        )
+
+    def forward(self, noisy_amp, noisy_pha):
+        """
+        :param noisy_amp: Tensor, shape: [b, f, t]
+        :param noisy_pha: Tensor, shape: [b, f, t]
+        :return:
+        """
+        x = torch.stack((noisy_amp, noisy_pha), dim=-1).permute(0, 3, 2, 1)  # [B, 2, T, F]
+        # x shape: [b, 2, t, f]
+        x = self.dense_encoder.forward(x)
+        # x shape: [b, c, t, f//2]
+
+        x = self.ts_transformer.forward(x)
+        # x shape: [b, c, t, f//2]
+
+        denoised_amp = noisy_amp * self.mask_decoder(x)
+        denoised_pha = self.phase_decoder(x)
+        denoised_com = torch.stack(
+            tensors=(
+                denoised_amp * torch.cos(denoised_pha),
+                denoised_amp * torch.sin(denoised_pha)
+            ),
+            dim=-1
+        )
+
+        return denoised_amp, denoised_pha, denoised_com
+
+
+MODEL_FILE = "generator.pt"
+
+
+class NXMPNetPretrainedModel(NXMPNet):
+    def __init__(self,
+                 config: NXMPNetConfig,
+                 ):
+        super(NXMPNetPretrainedModel, self).__init__(
+            config=config,
+        )
+        self.config = config
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        config = NXMPNetConfig.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 = NXMPNetConfig()
+
+    model = NXMPNet(config)
+
+    noisy_amp = torch.rand([1, 257, 201], dtype=torch.float32)
+    noisy_pha = torch.rand([1, 257, 201], dtype=torch.float32)
+
+    denoised_amp, denoised_pha, denoised_com = model.forward(noisy_amp, noisy_pha)
+    print(denoised_amp.shape)
+    print(denoised_pha.shape)
+    print(denoised_com.shape)
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/nx_mpnet/transformers/__init__.py

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

toolbox/torchaudio/models/nx_mpnet/transformers/attention.py

@@ -0,0 +1,263 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+
+
+class MultiHeadSelfAttention(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,
+                x: torch.Tensor,
+                mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+                cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        q, k, v = self.forward_qkv(x, x, x)
+
+        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 RelativeMultiHeadSelfAttention(nn.Module):
+
+    def __init__(self, n_head: int, n_feat: int, dropout_rate: float, max_relative_position: int = 5120):
+        """
+        :param n_head: int. the number of heads.
+        :param n_feat: int. the number of features.
+        :param dropout_rate: float. dropout rate.
+        :param max_relative_position: int. maximum relative position for relative position encoding.
+        """
+        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)
+
+        # Relative position encoding
+        self.max_relative_position = max_relative_position
+        self.relative_position_k = nn.Parameter(torch.randn(max_relative_position * 2 + 1, self.d_k))
+
+    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 = None
+                          ) -> torch.Tensor:
+        """
+        compute attention context vector.
+        :param value: torch.Tensor. transformed value. shape=(batch_size, n_head, key_time_steps, d_k).
+        :param scores: torch.Tensor. attention score. shape=(batch_size, n_head, query_time_steps, key_time_steps).
+        :param mask: torch.Tensor. mask. shape=(batch_size, 1, key_time_steps) or (batch_size, query_time_steps, key_time_steps).
+        :return: torch.Tensor. transformed value. (batch_size, query_time_steps, d_model).
+                weighted by the attention score (batch_size, query_time_steps, key_time_steps).
+        """
+        n_batch = value.size(0)
+        if mask is not None:
+            mask = mask.unsqueeze(1).eq(0)
+            # mask shape: [batch_size, 1, query_time_steps, key_time_steps]
+            scores = scores.masked_fill(mask, -float('inf'))
+            attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)
+        else:
+            attn = torch.softmax(scores, dim=-1)
+        # attn shape: [batch_size, n_head, query_time_steps, key_time_steps]
+
+        p_attn = self.dropout(attn)
+
+        x = torch.matmul(p_attn, value)
+        # x shape: [batch_size, n_head, query_time_steps, d_k]
+        x = x.transpose(1, 2)
+        # x shape: [batch_size, query_time_steps, n_head, d_k]
+
+        x = x.contiguous().view(n_batch, -1, self.h * self.d_k)  # (batch, time1, n_feat)
+        # x shape: [batch_size, query_time_steps, n_head * d_k]
+        # x shape: [batch_size, query_time_steps, n_feat]
+
+        x = self.linear_out(x)
+        # x shape: [batch_size, query_time_steps, n_feat]
+        return x
+
+    def relative_position_encoding(self, length: int) -> torch.Tensor:
+        """
+        Generate relative position encoding.
+        :param length: int. length of the sequence.
+        :return: torch.Tensor. relative position encoding. shape=(length, length, d_k).
+        """
+        range_vec = torch.arange(length)
+        distance_mat = range_vec.unsqueeze(0) - range_vec.unsqueeze(1)
+        distance_mat_clipped = torch.clamp(distance_mat, -self.max_relative_position, self.max_relative_position)
+        final_mat = distance_mat_clipped + self.max_relative_position
+        return final_mat
+
+    def forward(self,
+                x: torch.Tensor,
+                mask: torch.Tensor = None,
+                cache: torch.Tensor = None
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # attention! self attention.
+
+        q, k, v = self.forward_qkv(x, x, x)
+        # q k v shape: [batch_size, self.h, query_time_steps, self.d_k]
+
+        if cache is not None:
+            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)
+
+        # new_cache shape: [batch_size, self.h, time_steps, self.d_k * 2]
+        new_cache = torch.cat((k, v), dim=-1)
+
+        # native_scores shape: [batch_size, self.h, q_time_steps, k_time_steps]
+        native_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+
+        # Compute relative position encoding
+        q_length, k_length = q.size(2), k.size(2)
+        relative_position = self.relative_position_encoding(k_length)
+
+        relative_position = relative_position[-q_length:]
+
+        relative_position_k = self.relative_position_k[relative_position.view(-1)].view(q_length, k_length, -1)
+
+        relative_position_k = relative_position_k.unsqueeze(0).unsqueeze(0)  # (1, 1, q_length, k_length, d_k)
+        relative_position_k = relative_position_k.expand(q.size(0), q.size(1), -1, -1, -1)  # (batch, head, q_length, k_length, d_k)
+
+        relative_position_scores = torch.matmul(q.unsqueeze(3), relative_position_k.transpose(-2, -1)).squeeze(3) / math.sqrt(self.d_k)
+        # relative_position_scores shape: [batch_size, self.h, q_time_steps, k_time_steps]
+
+        # score
+        scores = native_scores + relative_position_scores
+
+        return self.forward_attention(v, scores, mask), new_cache
+
+
+def main():
+    rel_attention = RelativeMultiHeadSelfAttention(n_head=4, n_feat=256, dropout_rate=0.1)
+
+    x = torch.ones(size=(1, 200, 256), dtype=torch.float32)
+    xt, new_cache = rel_attention.forward(x, x, x)
+
+    # x = torch.ones(size=(1, 1, 256), dtype=torch.float32)
+    # cache = torch.ones(size=(1, 4, 199, 128), dtype=torch.float32)
+    # xt, new_cache = rel_attention.forward(x, x, x, cache=cache)
+
+    print(xt.shape)
+    print(new_cache.shape)
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/nx_mpnet/transformers/mask.py

@@ -0,0 +1,74 @@
+#!/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,
+        num_right_chunks: int = 0,
+        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 num_right_chunks: int. number of right 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 + num_right_chunks) * 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)
+
+    chunk_mask = subsequent_chunk_mask(size=8, chunk_size=2, num_left_chunks=2, num_right_chunks=1)
+    print(chunk_mask)
+
+    chunk_mask = subsequent_chunk_mask(size=9, chunk_size=2, num_left_chunks=2, num_right_chunks=1)
+    print(chunk_mask)
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/nx_mpnet/transformers/transformers.py

@@ -0,0 +1,479 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from typing import Dict, Optional, Tuple, List, Union
+
+import torch
+import torch.nn as nn
+
+from toolbox.torchaudio.models.nx_clean_unet.transformers.mask import subsequent_chunk_mask
+from toolbox.torchaudio.models.nx_clean_unet.transformers.attention import MultiHeadSelfAttention, RelativeMultiHeadSelfAttention
+
+
+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 TransformerBlock(nn.Module):
+    def __init__(self,
+                 input_dim: int,
+                 dropout_rate: float = 0.1,
+                 n_heads: int = 4,
+                 max_relative_position: int = 5120
+                 ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(input_dim, eps=1e-5)
+        self.attention = RelativeMultiHeadSelfAttention(
+            n_head=n_heads,
+            n_feat=input_dim,
+            dropout_rate=dropout_rate,
+            max_relative_position=max_relative_position,
+        )
+
+        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 = None,
+            attention_cache: torch.Tensor = None,
+    ) -> Tuple[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 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, mask=mask, cache=attention_cache
+        )
+        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_relative_position: int = 1024,
+                 chunk_size: int = 1,
+                 num_left_chunks: int = 128,
+                 num_right_chunks: int = 2,
+                 ):
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+
+        self.max_relative_position = max_relative_position
+        self.chunk_size = chunk_size
+        self.num_left_chunks = num_left_chunks
+        self.num_right_chunks = num_right_chunks
+
+        self.input_linear = nn.Linear(
+            in_features=self.input_size,
+            out_features=self.hidden_size,
+        )
+
+        self.encoder_layer_list = torch.nn.ModuleList([
+            TransformerBlock(
+                input_dim=hidden_size,
+                n_heads=attention_heads,
+                dropout_rate=dropout_rate,
+                max_relative_position=max_relative_position,
+            ) 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]
+
+        chunk_masks = subsequent_chunk_mask(
+            size=time_steps,
+            chunk_size=self.chunk_size,
+            num_left_chunks=self.num_left_chunks,
+            num_right_chunks=self.num_right_chunks,
+        )
+        chunk_masks = chunk_masks.to(xs.device)
+        # chunk_masks shape: [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)
+
+        # 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,
+                      max_att_cache_length: int,
+                      attention_cache: torch.Tensor = None,
+                      ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+
+        :param xs:
+        :param max_att_cache_length:
+        :param attention_cache: Tensor, [num_layers, ...]
+        :return:
+        """
+        # xs shape: [batch_size, time_steps, input_size]
+        xs = self.input_linear.forward(xs)
+        # xs shape: [batch_size, time_steps, hidden_size]
+
+        r_att_cache = []
+        for idx, encoder_layer in enumerate(self.encoder_layer_list):
+            xs, new_att_cache = encoder_layer.forward(
+                x=xs, attention_cache=attention_cache[idx] if attention_cache is not None else None,
+            )
+            # new_att_cache shape: [batch_size, n_heads, time_steps, dim]
+            if new_att_cache.size(2) > max_att_cache_length:
+                begin = (self.num_left_chunks + self.num_right_chunks) * self.chunk_size
+                end = self.num_right_chunks * self.chunk_size
+                new_att_cache = new_att_cache[:, :, -begin:-end, :]
+            r_att_cache.append(new_att_cache)
+
+        r_att_cache = torch.stack(r_att_cache, dim=0)
+
+        # xs shape: [batch_size, time_steps, hidden_size]
+        xs = self.output_linear.forward(xs)
+        # xs shape: [batch_size, time_steps, input_size]
+
+        return xs, r_att_cache
+
+    def forward_chunk_by_chunk(
+            self,
+            xs: torch.Tensor,
+    ) -> torch.Tensor:
+
+        batch_size, time_steps, _ = xs.shape
+
+        # attention_cache shape: [num_blocks, attention_heads, self.num_left_chunks * self.chunk_size, n_heads * d_k * 2]
+        max_att_cache_length = (self.num_left_chunks + self.num_right_chunks) * self.chunk_size
+        attention_cache = None
+
+        outputs = []
+        for idx in range(0, time_steps, self.chunk_size):
+            begin = idx
+            end = begin + self.chunk_size * (self.num_right_chunks + 1)
+            chunk_xs = xs[:, begin:end, :]
+            # print(f"begin: {begin}, end: {end}, length: {chunk_xs.size(1)}")
+
+            ys, attention_cache = self.forward_chunk(
+                xs=chunk_xs,
+                max_att_cache_length=max_att_cache_length,
+                attention_cache=attention_cache,
+            )
+
+            # ys shape: [batch_size, self.chunk_size * (self.num_right_chunks + 1), input_size]
+            ys = ys[:, :self.chunk_size, :]
+
+            outputs.append(ys)
+
+        ys = torch.cat(outputs, 1)
+        return ys
+
+
+class TSTransformerBlock(nn.Module):
+    def __init__(self,
+                 input_dim: int,
+                 dropout_rate: float = 0.1,
+                 n_heads: int = 4,
+                 max_time_relative_position: int = 1024,
+                 max_freq_relative_position: int = 128,
+                 ):
+        super(TSTransformerBlock, self).__init__()
+        self.time_transformer = TransformerBlock(input_dim, dropout_rate, n_heads, max_time_relative_position)
+        self.freq_transformer = TransformerBlock(input_dim, dropout_rate, n_heads, max_freq_relative_position)
+
+    def forward(self,
+                x: torch.Tensor,
+                mask: torch.Tensor = None,
+                attention_cache: torch.Tensor = None,
+                ):
+        """
+
+        :param x: Tensor. shape: [batch_size, hidden_size, time_steps, input_size]
+        :param mask: Tensor. shape: [time_steps, time_steps]
+        :param attention_cache:
+        :return:
+        """
+        b, c, t, f = x.size()
+
+        mask = None if mask is None else torch.broadcast_to(mask, size=(b*f, t, t))
+
+        x = x.permute(0, 3, 2, 1).contiguous().view(b*f, t, c)
+        x_, new_att_cache = self.time_transformer.forward(x, mask, attention_cache)
+        x = x_ + x
+        x = x.view(b, f, t, c).permute(0, 2, 1, 3).contiguous().view(b*t, f, c)
+        x_, _ = self.freq_transformer.forward(x)
+        x = x_ + x
+        x = x.view(b, t, f, c).permute(0, 3, 1, 2)
+        return x, new_att_cache
+
+
+class TSTransformerEncoder(nn.Module):
+    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_time_relative_position: int = 1024,
+                 max_freq_relative_position: int = 128,
+                 chunk_size: int = 1,
+                 num_left_chunks: int = 128,
+                 num_right_chunks: int = 2,
+                 ):
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+
+        self.max_time_relative_position = max_time_relative_position
+        self.max_freq_relative_position = max_freq_relative_position
+        self.chunk_size = chunk_size
+        self.num_left_chunks = num_left_chunks
+        self.num_right_chunks = num_right_chunks
+
+        self.input_linear = nn.Linear(
+            in_features=self.input_size,
+            out_features=self.hidden_size,
+        )
+
+        self.encoder_layer_list = torch.nn.ModuleList([
+            TSTransformerBlock(
+                input_dim=hidden_size,
+                n_heads=attention_heads,
+                dropout_rate=dropout_rate,
+                max_time_relative_position=max_time_relative_position,
+                max_freq_relative_position=max_freq_relative_position,
+            ) 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, channels, time_steps, input_size]
+        :return: Tensor, shape: [batch_size, channels, time_steps, input_size]
+        """
+        batch_size, channels, time_steps, _ = xs.shape
+        # xs shape: [batch_size, channels, time_steps, input_size]
+        xs = xs.permute(0, 3, 2, 1)
+        # xs shape: [batch_size, input_size, time_steps, channels]
+        xs = self.input_linear.forward(xs)
+        # xs shape: [batch_size, input_size, time_steps, hidden_size]
+        xs = xs.permute(0, 3, 2, 1)
+        # xs shape: [batch_size, hidden_size, time_steps, input_size]
+
+        chunk_masks = subsequent_chunk_mask(
+            size=time_steps,
+            chunk_size=self.chunk_size,
+            num_left_chunks=self.num_left_chunks,
+            num_right_chunks=self.num_right_chunks,
+        )
+        chunk_masks = chunk_masks.to(xs.device)
+        # chunk_masks shape: [time_steps, time_steps]
+
+        for encoder_layer in self.encoder_layer_list:
+            xs, _ = encoder_layer.forward(xs, chunk_masks)
+        # xs shape: [batch_size, hidden_size, time_steps, input_size]
+        xs = xs.permute(0, 3, 2, 1)
+        # xs shape: [batch_size, input_size, time_steps, hidden_size]
+        xs = self.output_linear.forward(xs)
+        # xs shape: [batch_size, input_size, time_steps, channels]
+        xs = xs.permute(0, 3, 2, 1)
+        # xs shape: [batch_size, channels, time_steps, input_size]
+
+        return xs
+
+    def forward_chunk(self,
+                      xs: torch.Tensor,
+                      max_att_cache_length: int,
+                      attention_cache: torch.Tensor = None,
+                      ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+
+        :param xs:
+        :param max_att_cache_length:
+        :param attention_cache: Tensor, shape: [num_layers, ...]
+        :return:
+        """
+        # xs shape: [batch_size, channels, time_steps, input_size]
+        xs = xs.permute(0, 3, 2, 1)
+        xs = self.input_linear.forward(xs)
+        xs = xs.permute(0, 3, 2, 1)
+        # xs shape: [batch_size, hidden_size, time_steps, input_size]
+
+        r_att_cache = []
+        for idx, encoder_layer in enumerate(self.encoder_layer_list):
+            xs, new_att_cache = encoder_layer.forward(
+                x=xs, attention_cache=attention_cache[idx] if attention_cache is not None else None,
+            )
+            # new_att_cache shape: [b*f, n_heads, time_steps, dim]
+            if new_att_cache.size(2) > max_att_cache_length:
+                begin = (self.num_left_chunks + self.num_right_chunks) * self.chunk_size
+                end = self.num_right_chunks * self.chunk_size
+                new_att_cache = new_att_cache[:, :, -begin:-end, :]
+            r_att_cache.append(new_att_cache)
+
+        r_att_cache = torch.stack(r_att_cache, dim=0)
+
+        # xs shape: [batch_size, hidden_size, time_steps, input_size]
+        xs = xs.permute(0, 3, 2, 1)
+        xs = self.output_linear.forward(xs)
+        xs = xs.permute(0, 3, 2, 1)
+        # xs shape: [batch_size, channels, time_steps, input_size]
+
+        return xs, r_att_cache
+
+    def forward_chunk_by_chunk(
+            self,
+            xs: torch.Tensor,
+    ) -> torch.Tensor:
+
+        batch_size, channels, time_steps, _ = xs.shape
+
+        max_att_cache_length = (self.num_left_chunks + self.num_right_chunks) * self.chunk_size
+        attention_cache = None
+
+        outputs = []
+        for idx in range(0, time_steps, self.chunk_size):
+            begin = idx
+            end = begin + self.chunk_size * (self.num_right_chunks + 1)
+            chunk_xs = xs[:, :, begin:end, :]
+            # chunk_xs shape: [batch_size, channels, self.chunk_size * (self.num_right_chunks + 1), input_size]
+
+            ys, attention_cache = self.forward_chunk(
+                xs=chunk_xs,
+                max_att_cache_length=max_att_cache_length,
+                attention_cache=attention_cache,
+            )
+            # ys shape: [batch_size, channels, self.chunk_size * (self.num_right_chunks + 1), input_size]
+            ys = ys[:, :, :self.chunk_size, :]
+
+            outputs.append(ys)
+
+        ys = torch.cat(outputs, dim=2)
+        return ys
+
+
+def main2():
+
+    encoder = TransformerEncoder(
+        input_size=64,
+        hidden_size=256,
+        attention_heads=4,
+        num_blocks=6,
+        dropout_rate=0.1,
+    )
+    print(encoder)
+
+    x = torch.ones([4, 200, 64])
+
+    x = torch.ones([4, 200, 64])
+    y = encoder.forward(xs=x)
+    print(y.shape)
+
+    x = torch.ones([4, 200, 64])
+    y = encoder.forward_chunk_by_chunk(xs=x)
+    print(y.shape)
+
+    return
+
+
+def main():
+
+    encoder = TSTransformerEncoder(
+        input_size=8,
+        hidden_size=16,
+        attention_heads=2,
+        num_blocks=2,
+        dropout_rate=0.1,
+    )
+    # print(encoder)
+
+    x = torch.ones([4, 8, 200, 8])
+    y = encoder.forward(xs=x)
+    print(y.shape)
+
+    x = torch.ones([4, 8, 200, 8])
+    y = encoder.forward_chunk_by_chunk(xs=x)
+    print(y.shape)
+
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/nx_mpnet/utils.py

@@ -0,0 +1,56 @@
+#!/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)
+
+
+class LearnableSigmoid2d(nn.Module):
+    def __init__(self, in_features, beta=1):
+        super().__init__()
+        self.beta = beta
+        self.slope = nn.Parameter(torch.ones(in_features, 1))
+        self.slope.requiresGrad = True
+
+    def forward(self, x):
+        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