update

examples/mpnet_aishell/step_3_evaluation.py

@@ -85,14 +85,14 @@ def save_audios(noise_audio: torch.Tensor,
     output_dir.mkdir(parents=True, exist_ok=True)
 
     filename = output_dir / "noise_audio.wav"
-    torchaudio.save(filename, noise_audio.detach().cpu(), sample_rate)
+    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)
+    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)
+    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)
+    torchaudio.save(filename, enhanced_audio.detach().cpu(), sample_rate, bits_per_sample=16)
 
     return output_dir.as_posix()
 
@@ -159,14 +159,15 @@ def main():
         # inference
         clean_audio = clean_audio.to(device)
         noisy_audio = noisy_audio.to(device)
-        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()
+        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,

main.py

@@ -4,8 +4,11 @@ import argparse
 import platform
 
 import gradio as gr
+import numpy as np
+import torch
 
 from project_settings import environment, project_path
+from toolbox.torchaudio.models.mpnet.inference_mpnet import InferenceMPNet
 
 
 def get_args():
@@ -25,12 +28,58 @@ def get_args():
     return args
 
 
+denoise_engines = {
+    "mpnet": InferenceMPNet(
+        pretrained_model_path_or_zip_file=(project_path / "trained_models/mpnet_aishell_20250221.zip").as_posix(),
+    ),
+
+}
+
+
+def when_click_denoise_button(noisy_audio_t, engine: str):
+    sample_rate, signal = noisy_audio_t
+
+    noisy_audio = np.array(signal / (1 << 15), dtype=np.float32)
+
+    infer_engine = denoise_engines.get(engine)
+    if infer_engine is None:
+        raise gr.Error(f"invalid denoise engine: {engine}.")
+
+    try:
+        enhanced_audio = infer_engine.enhancement_by_ndarray(noisy_audio)
+        enhanced_audio = np.array(enhanced_audio * (1 << 15), dtype=np.int16)
+    except Exception as e:
+        raise gr.Error(f"enhancement failed, error type: {type(e)}, error text: {str(e)}.")
+
+    enhanced_audio_t = (sample_rate, enhanced_audio)
+    return enhanced_audio_t, None
+
+
 def main():
     args = get_args()
 
+    # choices
+    denoise_engine_choices = list(denoise_engines.keys())
+
     # ui
     with gr.Blocks() as blocks:
-        gr.Markdown(value="in progress.")
+        gr.Markdown(value="nx denoise.")
+        with gr.Tabs():
+            with gr.TabItem("denoise"):
+                with gr.Row():
+                    with gr.Column(variant="panel", scale=5):
+                        dn_noisy_audio = gr.Audio(label="noisy_audio")
+                        dn_engine = gr.Dropdown(choices=denoise_engine_choices, value=denoise_engine_choices[0], label="engine")
+                        dn_button = gr.Button(variant="primary")
+                    with gr.Column(variant="panel", scale=5):
+                        dn_enhanced_audio = gr.Audio(label="enhanced_audio")
+                        dn_clean_audio = gr.Audio(label="clean_audio")
+
+                dn_button.click(
+                    when_click_denoise_button,
+                    inputs=[dn_noisy_audio, dn_engine],
+                    outputs=[dn_enhanced_audio, dn_clean_audio]
+                )
 
     # http://127.0.0.1:7864/
     blocks.queue().launch(

toolbox/torchaudio/models/mpnet/inference_mpnet.py

@@ -0,0 +1,109 @@
+#!/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.mpnet.configuration_mpnet import MPNetConfig
+from toolbox.torchaudio.models.mpnet.modeling_mpnet import MPNetPretrainedModel, MODEL_FILE
+from toolbox.torchaudio.models.mpnet.utils import mag_pha_stft, mag_pha_istft
+
+logger = logging.getLogger("toolbox")
+
+
+class InferenceMPNet(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, generator = 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.generator = generator
+        self.generator.to(device)
+        self.generator.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 = MPNetConfig.from_pretrained(
+            pretrained_model_name_or_path=model_path.as_posix(),
+        )
+        generator = MPNetPretrainedModel.from_pretrained(
+            pretrained_model_name_or_path=model_path.as_posix(),
+        )
+        generator.to(self.device)
+        generator.eval()
+
+        shutil.rmtree(model_path)
+        return config, generator
+
+    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]
+        noisy_audio = self.enhancement_by_tensor(noisy_audio)
+        noisy_audio = noisy_audio[0]
+
+        return noisy_audio.cpu().numpy()
+
+    def enhancement_by_tensor(self, noisy_audio: torch.Tensor) -> torch.Tensor:
+        if torch.max(noisy_audio) > 1 or torch.min(noisy_audio) < -1:
+            raise AssertionError(f"The value range of audio samples should be between -1 and 1.")
+
+        noisy_audio = noisy_audio.to(self.device)
+
+        with torch.no_grad():
+            noisy_mag, noisy_pha, noisy_com = mag_pha_stft(
+                noisy_audio,
+                self.config.n_fft, self.config.hop_size, self.config.win_size, self.config.compress_factor
+            )
+            mag_g, pha_g, com_g = self.generator.forward(noisy_mag, noisy_pha)
+            audio_g = mag_pha_istft(
+                mag_g, pha_g,
+                self.config.n_fft, self.config.hop_size, self.config.win_size, self.config.compress_factor
+            )
+            enhanced_audio = audio_g.detach()
+
+        return enhanced_audio
+
+def main():
+    model_zip_file = project_path / "trained_models/mpnet_aishell_20250221.zip"
+    infer_mpnet = InferenceMPNet(model_zip_file)
+
+    sample_rate = 8000
+    noisy_audio_file = project_path / "data/examples/noisy_audio.wav"
+    noisy_audio, _ = librosa.load(
+        noisy_audio_file.as_posix(),
+        sr=sample_rate,
+    )
+    noisy_audio = torch.tensor(noisy_audio, dtype=torch.float32)
+    noisy_audio = noisy_audio.unsqueeze(dim=0)
+
+    enhanced_audio = infer_mpnet.enhancement_by_tensor(noisy_audio)
+
+    filename = "enhanced_audio.wav"
+    torchaudio.save(filename, enhanced_audio.detach().cpu(), sample_rate)
+
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/mpnet/modeling_mpnet.py

@@ -8,6 +8,8 @@ https://huggingface.co/spaces/JacobLinCool/MP-SENet
 https://arxiv.org/abs/2305.13686
 https://github.com/yxlu-0102/MP-SENet
 
+应该是不支持流式改造的。
+
 """
 import os
 from typing import Optional, Union