Spaces:
Running
Running
File size: 4,320 Bytes
bd94e77 8ce0f99 bd94e77 6f78f1b bd94e77 8ce0f99 1e78a70 bd94e77 1e78a70 bd94e77 8ce0f99 2cc4b35 8ce0f99 bd94e77 5e7d9ca 1e78a70 6f78f1b 1e78a70 bd94e77 8ce0f99 5e7d9ca 6f78f1b 5e7d9ca 1e78a70 6f78f1b bd94e77 1e78a70 6f78f1b 1e78a70 bd94e77 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
#!/usr/bin/python3
# -*- coding: utf-8 -*-
import argparse
from pathlib import Path
import platform
import shutil
import zipfile
import gradio as gr
from huggingface_hub import snapshot_download
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():
parser = argparse.ArgumentParser()
parser.add_argument(
"--examples_dir",
# default=(project_path / "data").as_posix(),
default=(project_path / "data/examples").as_posix(),
type=str
)
parser.add_argument(
"--models_repo_id",
default="qgyd2021/nx_denoise",
type=str
)
parser.add_argument(
"--trained_model_dir",
default=(project_path / "trained_models").as_posix(),
type=str
)
parser.add_argument(
"--hf_token",
default=environment.get("hf_token"),
type=str,
)
parser.add_argument(
"--server_port",
default=environment.get("server_port", 7860),
type=int
)
args = parser.parse_args()
return args
denoise_engines = dict()
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
def main():
args = get_args()
examples_dir = Path(args.examples_dir)
trained_model_dir = Path(args.trained_model_dir)
# download models
if not trained_model_dir.exists():
trained_model_dir.mkdir(parents=True, exist_ok=True)
_ = snapshot_download(
repo_id=args.models_repo_id,
local_dir=trained_model_dir.as_posix(),
token=args.hf_token,
)
# engines
global denoise_engines
denoise_engines = {
"mpnet_aishell_20250221": InferenceMPNet(
pretrained_model_path_or_zip_file=(project_path / "trained_models/mpnet_aishell_20250221.zip").as_posix(),
),
}
# choices
denoise_engine_choices = list(denoise_engines.keys())
# examples
example_zip_file = trained_model_dir / "examples.zip"
with zipfile.ZipFile(example_zip_file.as_posix(), "r") as f_zip:
out_root = examples_dir
if out_root.exists():
shutil.rmtree(out_root.as_posix())
out_root.mkdir(parents=True, exist_ok=True)
f_zip.extractall(path=out_root)
# examples
examples = list()
for filename in examples_dir.glob("**/*/*.wav"):
label = filename.parts[-2]
examples.append([
filename.as_posix(),
denoise_engine_choices[0]
])
# ui
with gr.Blocks() as blocks:
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_button.click(
when_click_denoise_button,
inputs=[dn_noisy_audio, dn_engine],
outputs=[dn_enhanced_audio]
)
# http://127.0.0.1:7864/
blocks.queue().launch(
share=False if platform.system() == "Windows" else False,
server_name="127.0.0.1" if platform.system() == "Windows" else "0.0.0.0",
server_port=args.server_port
)
return
if __name__ == "__main__":
main()
|