update

examples/dfnet2/yaml/config.yaml

@@ -7,6 +7,9 @@ win_size: 200
 hop_size: 80
 
 spec_bins: 256
+erb_bins: 32
+min_freq_bins_for_erb: 2
+use_ema_norm: true
 
 # model
 conv_channels: 64

examples/dtln/step_2_train_model.py

@@ -259,6 +259,7 @@ def main():
             noisy_audios: torch.Tensor = noisy_audios.to(device)
 
             denoise_audios = model.forward(noisy_audios)
+            denoise_audios = torch.squeeze(denoise_audios, dim=1)
 
             mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
             neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)

examples/rnnoise/run.sh

@@ -6,10 +6,9 @@ sh run.sh --stage 2 --stop_stage 2 --system_version windows --file_folder_name f
 
 sh run.sh --stage 3 --stop_stage 3 --system_version windows --file_folder_name file_dir
 
-sh run.sh --stage 1 --stop_stage 3 --system_version centos --file_folder_name file_dir \
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name rnnoise-nx-dns3 \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train" \
---sparse
 
 
 END
@@ -108,66 +107,6 @@ if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
   --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: export model"
-  cd "${work_dir}" || exit 1
-  python3 step_5_export_models.py \
-  --vocabulary_dir "${vocabulary_dir}" \
-  --model_dir "${file_dir}/best" \
-  --serialization_dir "${file_dir}" \
-
-fi
-
-
-if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
-  $verbose && echo "stage 5: collect files"
-  cd "${work_dir}" || exit 1
-
-  mkdir -p ${final_model_dir}
-
-  cp "${file_dir}/best"/* "${final_model_dir}"
-  cp -r "${file_dir}/vocabulary" "${final_model_dir}"
-
-  cp "${file_dir}/evaluation.xlsx" "${final_model_dir}/evaluation.xlsx"
-
-  cp "${file_dir}/trace_model.zip" "${final_model_dir}/trace_model.zip"
-  cp "${file_dir}/trace_quant_model.zip" "${final_model_dir}/trace_quant_model.zip"
-  cp "${file_dir}/script_model.zip" "${final_model_dir}/script_model.zip"
-  cp "${file_dir}/script_quant_model.zip" "${final_model_dir}/script_quant_model.zip"
-
-  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 6 ] && [ ${stop_stage} -ge 6 ]; then
-  $verbose && echo "stage 6: clear file_dir"
-  cd "${work_dir}" || exit 1
-
-  rm -rf "${file_dir}";
+  --sparse
 
 fi

toolbox/torchaudio/models/dfnet/inference_dfnet.py

@@ -76,13 +76,10 @@ class InferenceDfNet(object):
         with torch.no_grad():
             est_spec, est_wav, est_mask, lsnr = self.model.forward(noisy_audios)
 
-        # shape: [batch_size, num_samples]
-        enhanced_audio = torch.unsqueeze(est_wav, dim=1)
         # shape: [batch_size, 1, num_samples]
-
-        enhanced_audio = enhanced_audio[0]
+        denoise = est_wav[0]
         # shape: [channels, num_samples]
-        return enhanced_audio
+        return denoise
 
 
 def main():
@@ -90,7 +87,7 @@ def main():
     infer_model = InferenceDfNet(model_zip_file)
 
     sample_rate = 8000
-    noisy_audio_file = project_path / "data/examples/ai_agent/dfaaf264-b5e3-4ca2-b5cb-5b6d637d962d_section_3.wav"
+    noisy_audio_file = project_path / "data/examples/ai_agent/chinese-3.wav"
     noisy_audio, sample_rate = librosa.load(
         noisy_audio_file.as_posix(),
         sr=sample_rate,

toolbox/torchaudio/models/dfnet/modeling_dfnet.py

@@ -961,9 +961,8 @@ class DfNet(nn.Module):
         # est_spec shape: [b, f, t], torch.complex64
 
         est_wav = self.istft.forward(est_spec)
-        est_wav = torch.squeeze(est_wav, dim=1)
-        est_wav = est_wav[:, :n_samples]
-        # est_wav shape: [b, n_samples]
+        est_wav = est_wav[:, :, :n_samples]
+        # est_wav shape: [b, 1, n_samples]
 
         est_mask = torch.squeeze(mask, dim=1)
         est_mask = est_mask.permute(0, 2, 1)

toolbox/torchaudio/models/dfnet2/configuration_dfnet2.py

@@ -16,6 +16,7 @@ class DfNet2Config(PretrainedConfig):
                  spec_bins: int = 256,
                  erb_bins: int = 32,
                  min_freq_bins_for_erb: int = 2,
+                 use_ema_norm: bool = True,
 
                  conv_channels: int = 64,
                  conv_kernel_size_input: Tuple[int, int] = (3, 3),
@@ -83,6 +84,8 @@ class DfNet2Config(PretrainedConfig):
         self.erb_bins = erb_bins
         self.min_freq_bins_for_erb = min_freq_bins_for_erb
 
+        self.use_ema_norm = use_ema_norm
+
         # conv
         self.conv_channels = conv_channels
         self.conv_kernel_size_input = conv_kernel_size_input

toolbox/torchaudio/models/dfnet2/inference_dfnet2.py

@@ -14,8 +14,8 @@ import torchaudio
 torch.set_num_threads(1)
 
 from project_settings import project_path
-from toolbox.torchaudio.models.dfnet.configuration_dfnet import DfNetConfig
-from toolbox.torchaudio.models.dfnet.modeling_dfnet import DfNetPretrainedModel, MODEL_FILE
+from toolbox.torchaudio.models.dfnet2.configuration_dfnet2 import DfNet2Config
+from toolbox.torchaudio.models.dfnet2.modeling_dfnet2 import DfNet2PretrainedModel, MODEL_FILE
 
 logger = logging.getLogger("toolbox")
 
@@ -43,10 +43,10 @@ class InferenceDfNet(object):
                 f_zip.extractall(path=out_root)
             model_path = out_root / model_path.stem
 
-        config = DfNetConfig.from_pretrained(
+        config = DfNet2Config.from_pretrained(
             pretrained_model_name_or_path=model_path.as_posix(),
         )
-        model = DfNetPretrainedModel.from_pretrained(
+        model = DfNet2PretrainedModel.from_pretrained(
             pretrained_model_name_or_path=model_path.as_posix(),
         )
         model.to(self.device)
@@ -60,13 +60,13 @@ class InferenceDfNet(object):
         noisy_audio = noisy_audio.unsqueeze(dim=0)
 
         # noisy_audio shape: [batch_size, n_samples]
-        enhanced_audio = self.enhancement_by_tensor(noisy_audio)
+        enhanced_audio = self.denoise_offline(noisy_audio)
         # enhanced_audio shape: [channels, num_samples]
         enhanced_audio = enhanced_audio[0]
         # enhanced_audio shape: [num_samples]
         return enhanced_audio.cpu().numpy()
 
-    def enhancement_by_tensor(self, noisy_audio: torch.Tensor) -> torch.Tensor:
+    def denoise_offline(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.")
 
@@ -76,21 +76,33 @@ class InferenceDfNet(object):
         with torch.no_grad():
             est_spec, est_wav, est_mask, lsnr = self.model.forward(noisy_audios)
 
-        # shape: [batch_size, num_samples]
-        enhanced_audio = torch.unsqueeze(est_wav, dim=1)
         # shape: [batch_size, 1, num_samples]
+        denoise = est_wav[0]
+        # shape: [channels, num_samples]
+        return denoise
 
-        enhanced_audio = enhanced_audio[0]
+    def denoise_online(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():
+            denoise = self.model.forward_chunk_by_chunk(noisy_audios)
+
+        # shape: [batch_size, 1, num_samples]
+        denoise = denoise[0]
         # shape: [channels, num_samples]
-        return enhanced_audio
+        return denoise
 
 
 def main():
-    model_zip_file = project_path / "trained_models/dfnet-nx-dns3.zip"
+    model_zip_file = project_path / "trained_models/dfnet2-nx-dns3.zip"
     infer_model = InferenceDfNet(model_zip_file)
 
     sample_rate = 8000
-    noisy_audio_file = project_path / "data/examples/ai_agent/dfaaf264-b5e3-4ca2-b5cb-5b6d637d962d_section_3.wav"
+    noisy_audio_file = project_path / "data/examples/ai_agent/chinese-3.wav"
     noisy_audio, sample_rate = librosa.load(
         noisy_audio_file.as_posix(),
         sr=sample_rate,
@@ -101,11 +113,19 @@ def main():
     noisy_audio = noisy_audio.unsqueeze(dim=0)
 
     begin = time.time()
-    enhanced_audio = infer_model.enhancement_by_tensor(noisy_audio)
+    enhanced_audio = infer_model.denoise_offline(noisy_audio)
+    time_cost = time.time() - begin
+    print(f"enhanced_audio.shape: {enhanced_audio.shape}, time_cost: {time_cost:.4f}, audio_duration: {duration:.4f}, fpr: {time_cost / duration:.4f}")
+
+    filename = "enhanced_audio_offline.wav"
+    torchaudio.save(filename, enhanced_audio.detach().cpu(), sample_rate)
+
+    begin = time.time()
+    enhanced_audio = infer_model.denoise_online(noisy_audio)
     time_cost = time.time() - begin
     print(f"enhanced_audio.shape: {enhanced_audio.shape}, time_cost: {time_cost:.4f}, audio_duration: {duration:.4f}, fpr: {time_cost / duration:.4f}")
 
-    filename = "enhanced_audio.wav"
+    filename = "enhanced_audio_online.wav"
     torchaudio.save(filename, enhanced_audio.detach().cpu(), sample_rate)
 
     return

toolbox/torchaudio/models/dfnet2/modeling_dfnet2.py

@@ -1097,7 +1097,8 @@ class DfNet2(nn.Module):
         noisy = self.signal_prepare(noisy)
 
         spec, feat_erb, feat_spec = self.feature_prepare(noisy)
-        feat_erb, feat_spec, _ = self.feature_norm(feat_erb, feat_spec)
+        if self.config.use_ema_norm:
+            feat_erb, feat_spec, _ = self.feature_norm(feat_erb, feat_spec)
 
         e0, e1, e2, e3, emb, c0, lsnr, _ = self.encoder.forward(feat_erb, feat_spec)
 
@@ -1180,11 +1181,12 @@ class DfNet2(nn.Module):
             # spec shape: [b, 1, t, f, 2]
             # feat_erb shape: [b, 1, t, erb_bins]
             # feat_spec shape: [b, 2, t, df_bins]
-            feat_erb, feat_spec, cache_dict6 = self.feature_norm(feat_erb, feat_spec, cache_dict=cache_dict6)
+            if self.config.use_ema_norm:
+                feat_erb, feat_spec, cache_dict0 = self.feature_norm(feat_erb, feat_spec, cache_dict=cache_dict0)
 
-            e0, e1, e2, e3, emb, c0, lsnr, cache_dict0 = self.encoder.forward(feat_erb, feat_spec, cache_dict=cache_dict0)
+            e0, e1, e2, e3, emb, c0, lsnr, cache_dict1 = self.encoder.forward(feat_erb, feat_spec, cache_dict=cache_dict1)
 
-            mask, cache_dict1 = self.erb_decoder.forward(emb, e3, e2, e1, e0, cache_dict=cache_dict1)
+            mask, cache_dict2 = self.erb_decoder.forward(emb, e3, e2, e1, e0, cache_dict=cache_dict2)
             # mask shape: [b, 1, t, erb_bins]
             mask = self.erb_bands.erb_scale_inv(mask)
             # mask shape: [b, 1, t, f]
@@ -1198,16 +1200,16 @@ class DfNet2(nn.Module):
             lsnr = torch.transpose(lsnr, dim0=2, dim1=1)
             # lsnr shape: [b, 1, t]
 
-            df_coefs, cache_dict2 = self.df_decoder.forward(emb, c0, cache_dict=cache_dict2)
+            df_coefs, cache_dict3 = self.df_decoder.forward(emb, c0, cache_dict=cache_dict3)
             df_coefs = self.df_out_transform(df_coefs)
             # df_coefs shape: [b, df_order, t, df_bins, 2]
 
             spec_ = spec[:, :, :, :self.config.spec_bins, :]
             # spec shape: [b, 1, t, spec_bins, 2]
-            spec_f, cache_dict3 = self.df_op.forward_online(spec_, df_coefs, cache_dict=cache_dict3)
+            spec_f, cache_dict4 = self.df_op.forward_online(spec_, df_coefs, cache_dict=cache_dict4)
             # spec_f shape: [b, 1, t, df_bins, 2], torch.float32
 
-            spec_e, cache_dict4 = self.spec_e_m_combine_online(spec_f, spec_m, cache_dict=cache_dict4)
+            spec_e, cache_dict5 = self.spec_e_m_combine_online(spec_f, spec_m, cache_dict=cache_dict5)
 
             spec_e = torch.squeeze(spec_e, dim=1)
             spec_e = spec_e.permute(0, 2, 1, 3)
@@ -1219,7 +1221,7 @@ class DfNet2(nn.Module):
             est_spec = torch.concat(tensors=[est_spec, est_spec[:, -1:, :]], dim=1)
             # est_spec shape: [b, f, t], torch.complex64
 
-            est_wav, cache_dict5 = self.istft.forward_chunk(est_spec, cache_dict=cache_dict5)
+            est_wav, cache_dict6 = self.istft.forward_chunk(est_spec, cache_dict=cache_dict6)
             # est_wav shape: [b, 1, hop_size]
 
             waveform_list.append(est_wav)
@@ -1361,14 +1363,22 @@ class DfNet2PretrainedModel(DfNet2):
 
 
 def main():
+    import time
+    # torch.set_num_threads(1)
 
     config = DfNet2Config()
     model = DfNet2PretrainedModel(config=config)
     model.eval()
 
-    noisy = torch.randn(size=(1, 16000), dtype=torch.float32)
+    num_samples = 16000
+    noisy = torch.randn(size=(1, num_samples), dtype=torch.float32)
+    duration = num_samples / config.sample_rate
 
+    begin = time.time()
     est_spec, est_wav, est_mask, lsnr = model.forward(noisy)
+    time_cost = time.time() - begin
+    print(f"time_cost: {time_cost:.4f}, audio_duration: {duration:.4f}, fpr: {time_cost / duration:.4f}")
+
     # print(f"est_spec.shape: {est_spec.shape}")
     # print(f"est_wav.shape: {est_wav.shape}")
     # print(f"est_mask.shape: {est_mask.shape}")
@@ -1381,7 +1391,11 @@ def main():
     print(waveform[:, :, 15760: 15762])
     print(waveform[:, :, 15840: 15842])
 
+    begin = time.time()
     waveform = model.forward_chunk_by_chunk(noisy)
+    time_cost = time.time() - begin
+    print(f"time_cost: {time_cost:.4f}, audio_duration: {duration:.4f}, fpr: {time_cost / duration:.4f}")
+
     waveform = waveform[:, :, (config.df_lookahead*config.hop_size):]
     print(f"waveform.shape: {waveform.shape}, waveform.dtype: {waveform.dtype}")
     print(waveform[:, :, 300: 302])

toolbox/torchaudio/models/dfnet2/yaml/config.yaml

@@ -7,6 +7,9 @@ win_size: 200
 hop_size: 80
 
 spec_bins: 256
+erb_bins: 32
+min_freq_bins_for_erb: 2
+use_ema_norm: true
 
 # model
 conv_channels: 64

toolbox/torchaudio/models/dtln/inference_dtln.py

@@ -0,0 +1,137 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import logging
+from pathlib import Path
+import shutil
+import tempfile, time
+import zipfile
+
+import librosa
+import numpy as np
+import torch
+import torchaudio
+
+torch.set_num_threads(1)
+
+from project_settings import project_path
+from toolbox.torchaudio.models.dtln.configuration_dtln import DTLNConfig
+from toolbox.torchaudio.models.dtln.modeling_dtln import DTLNPretrainedModel, MODEL_FILE
+
+logger = logging.getLogger("toolbox")
+
+
+class InferenceDTLN(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 = DTLNConfig.from_pretrained(
+            pretrained_model_name_or_path=model_path.as_posix(),
+        )
+        model = DTLNPretrainedModel.from_pretrained(
+            pretrained_model_name_or_path=model_path.as_posix(),
+        )
+        model.to(self.device)
+        model.eval()
+
+        shutil.rmtree(model_path)
+        return config, model
+
+    def enhancement_by_ndarray(self, noisy_audio: np.ndarray) -> np.ndarray:
+        noisy_audio = torch.tensor(noisy_audio, dtype=torch.float32)
+        noisy_audio = noisy_audio.unsqueeze(dim=0)
+
+        # noisy_audio shape: [batch_size, n_samples]
+        enhanced_audio = self.denoise_offline(noisy_audio)
+        # enhanced_audio shape: [channels, num_samples]
+        enhanced_audio = enhanced_audio[0]
+        # enhanced_audio shape: [num_samples]
+        return enhanced_audio.cpu().numpy()
+
+    def denoise_offline(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():
+            denoise = self.model.forward(noisy_audios)
+
+        # denoise shape: [batch_size, 1, num_samples]
+        denoise = denoise[0]
+        # shape: [channels, num_samples]
+        return denoise
+
+    def denoise_online(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():
+            denoise = self.model.forward_chunk_by_chunk(noisy_audios)
+
+        # denoise shape: [batch_size, 1, num_samples]
+        denoise = denoise[0]
+        # shape: [channels, num_samples]
+        return denoise
+
+
+def main():
+    model_zip_file = project_path / "trained_models/dtln-nx-dns3.zip"
+    infer_model = InferenceDTLN(model_zip_file)
+
+    sample_rate = 8000
+    noisy_audio_file = project_path / "data/examples/ai_agent/chinese-3.wav"
+    noisy_audio, sample_rate = librosa.load(
+        noisy_audio_file.as_posix(),
+        sr=sample_rate,
+    )
+    duration = librosa.get_duration(y=noisy_audio, 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)
+
+    # offline
+    begin = time.time()
+    enhanced_audio = infer_model.denoise_offline(noisy_audio)
+    time_cost = time.time() - begin
+    print(f"enhanced_audio.shape: {enhanced_audio.shape}, time_cost: {time_cost:.4f}, audio_duration: {duration:.4f}, fpr: {time_cost / duration:.4f}")
+
+    filename = "enhanced_audio_offline.wav"
+    torchaudio.save(filename, enhanced_audio.detach().cpu(), sample_rate)
+
+    # online
+    begin = time.time()
+    enhanced_audio = infer_model.denoise_online(noisy_audio)
+    time_cost = time.time() - begin
+    print(f"enhanced_audio.shape: {enhanced_audio.shape}, time_cost: {time_cost:.4f}, audio_duration: {duration:.4f}, fpr: {time_cost / duration:.4f}")
+
+    filename = "enhanced_audio_online.wav"
+    torchaudio.save(filename, enhanced_audio.detach().cpu(), sample_rate)
+
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/dtln/modeling_dtln.py

@@ -167,12 +167,13 @@ class DTLNModel(nn.Module):
         if remainder > 0:
             n_samples_pad = self.hop_size - remainder
             signal = F.pad(signal, pad=(0, n_samples_pad), mode="constant", value=0)
-        return signal, n_samples
+        return signal
 
     def forward(self,
                 noisy: torch.Tensor,
                 ):
-        noisy, num_samples = self.signal_prepare(noisy)
+        num_samples = noisy.shape[-1]
+        noisy = self.signal_prepare(noisy)
         batch_size, _, num_samples_pad = noisy.shape
         # print(f"num_samples: {num_samples}, num_samples_pad: {num_samples_pad}")
 
@@ -182,6 +183,8 @@ class DTLNModel(nn.Module):
 
         denoise = denoise[:, :num_samples]
         # denoise shape: [b, num_samples]
+        denoise = torch.unsqueeze(denoise, dim=1)
+        # denoise shape: [b, 1, num_samples]
         return denoise
 
     def forward_chunk(self,
@@ -189,7 +192,7 @@ class DTLNModel(nn.Module):
                       in_state1: torch.Tensor = None,
                       in_state2: torch.Tensor = None,
                       ):
-        # noisy shape: [b, num_samples]
+        # noisy shape: [b, 1, num_samples]
         spec = self.stft.forward(noisy)
         # spec shape: [b, f, t], torch.complex64
         # t = (num_samples - win_size) / hop_size + 1
@@ -233,6 +236,44 @@ class DTLNModel(nn.Module):
 
         return denoise_frame, out_state1, out_state2
 
+    def forward_chunk_by_chunk(self, noisy: torch.Tensor):
+        noisy = self.signal_prepare(noisy)
+        # noisy shape: [b, 1, num_samples]
+        batch_size, _, num_samples_pad = noisy.shape
+        # print(f"num_samples: {num_samples}, num_samples_pad: {num_samples_pad}")
+
+        t = (num_samples_pad - self.fft_size) // self.hop_size + 1
+
+        denoise_list = list()
+        out_state1 = None
+        out_state2 = None
+        overlap_size = self.fft_size - self.hop_size
+        denoise_cache = torch.zeros(size=(batch_size, overlap_size), dtype=noisy.dtype)
+        # denoise_list.append(torch.clone(denoise_cache))
+        for i in range(t):
+            begin = i * self.hop_size
+            end = begin + self.fft_size
+            sub_noisy = noisy[:, :, begin: end]
+            # noisy shape: [b, 1, frame_size]
+            with torch.no_grad():
+                sub_denoise_frame, out_state1, out_state2 = self.forward_chunk(sub_noisy, out_state1, out_state2)
+            # sub_denoise_frame shape: [b, fft_size, 1]
+            sub_denoise_frame = sub_denoise_frame[:, :, 0]
+            # sub_denoise_frame shape: [b, fft_size]
+
+            sub_denoise_frame[:, :overlap_size] += denoise_cache
+            denoise_out = sub_denoise_frame[:, :self.hop_size]
+            denoise_cache = sub_denoise_frame[:, self.hop_size:]
+            # denoise_cache shape: [b, hop_size]
+
+            denoise_list.append(denoise_out)
+
+        denoise = torch.concat(denoise_list, dim=-1)
+        # denoise shape: [b, num_samples]
+        denoise = torch.unsqueeze(denoise, dim=1)
+        # denoise shape: [b, 1, num_samples]
+        return denoise
+
     def denoise_frame_to_denoise(self, denoise_frame: torch.Tensor, batch_size: int, num_samples: int):
         # overlap and add
 
@@ -301,43 +342,28 @@ class DTLNPretrainedModel(DTLNModel):
 
 
 def main():
-    fft_size = 512
-    hop_size = 128
-
-    model = DTLNModel(fft_size=fft_size, hop_size=hop_size)
+    config = DTLNConfig()
+    model = DTLNPretrainedModel(config)
+    model.eval()
 
     noisy = torch.randn(size=(1, 16000), dtype=torch.float32)
-    batch_size, num_samples = noisy.shape
 
-    denoise = model.forward(noisy)
+    with torch.no_grad():
+        denoise = model.forward(noisy)
     print(f"denoise.shape: {denoise.shape}")
+    print(denoise[:, :, 300: 302])
+    print(denoise[:, :, 15680: 15682])
+    print(denoise[:, :, 15760: 15762])
+    print(denoise[:, :, 15840: 15842])
 
-    t = (num_samples - fft_size) // hop_size + 1
-
-    denoise_list = list()
-    out_state1 = None
-    out_state2 = None
-    denoise_cache = torch.zeros(size=(batch_size, fft_size - hop_size,), dtype=noisy.dtype)
-    denoise_list.append(torch.clone(denoise_cache))
-    for i in range(t):
-        begin = i * hop_size
-        end = begin + fft_size
-        sub_noisy = noisy[:, begin: end]
-        with torch.no_grad():
-            sub_denoise_frame, out_state1, out_state2 = model.forward_chunk(sub_noisy, out_state1, out_state2)
-        # sub_denoise_frame shape: [b, fft_size, 1]
-        sub_denoise_frame = sub_denoise_frame[:, :, 0]
-        # sub_denoise_frame shape: [b, fft_size]
-
-        sub_denoise_frame[:, hop_size:] += denoise_cache
-        denoise_out = sub_denoise_frame[:, :hop_size]
-        denoise_cache = sub_denoise_frame[:, hop_size:]
-        # denoise_cache shape: [b, hop_size]
-
-        denoise_list.append(denoise_out)
-
-    denoise = torch.concat(denoise_list, dim=-1)
+    denoise = model.forward_chunk_by_chunk(noisy)
     print(f"denoise.shape: {denoise.shape}")
+    # denoise = denoise[:, :, (config.fft_size - config.hop_size):]
+    print(denoise[:, :, 300: 302])
+    print(denoise[:, :, 15680: 15682])
+    print(denoise[:, :, 15760: 15762])
+    print(denoise[:, :, 15840: 15842])
+
     return