update

examples/nx_mpnet/yaml/config.yaml

@@ -22,8 +22,8 @@ tsfm_dropout_rate: 0.0
 tsfm_max_time_relative_position: 2048
 tsfm_max_freq_relative_position: 256
 tsfm_chunk_size: 1
-tsfm_num_left_chunks: 64
-tsfm_num_right_chunks: 2
+tsfm_num_left_chunks: 96
+tsfm_num_right_chunks: 3
 
 discriminator_dim: 32
 discriminator_in_channel: 2

toolbox/torchaudio/models/nx_mpnet/{inference_mpnet.py → inference_nx_mpnet.py}

@@ -78,23 +78,23 @@ class InferenceNXMPNet(object):
 
 
 def main():
-    model_zip_file = project_path / "trained_models/mpnet-aishell-1-epoch.zip"
+    model_zip_file = project_path / "trained_models/nx-mpnet-aishell-2-epoch.zip"
     infer_mpnet = InferenceNXMPNet(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_file = project_path / "data/examples/ai_agent/dfaaf264-b5e3-4ca2-b5cb-5b6d637d962d_section_2.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 = 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 = infer_mpnet.enhancement_by_tensor(noisy_audio)
 
     filename = "enhanced_audio.wav"
-    torchaudio.save(filename, enhanced_audio.detach().cpu(), sample_rate)
+    torchaudio.save(filename, enhanced_audio.unsqueeze(dim=0).detach().cpu(), sample_rate)
 
     return