|
|
import torch |
|
|
import spaces |
|
|
import tempfile |
|
|
import soundfile as sf |
|
|
import gradio as gr |
|
|
import librosa as lb |
|
|
import yaml |
|
|
import numpy as np |
|
|
import matplotlib.pyplot as plt |
|
|
from model.cleanmel import CleanMel |
|
|
from model.vocos.pretrained import Vocos |
|
|
from model.stft import InputSTFT, TargetMel |
|
|
|
|
|
DEVICE = torch.device("cuda:5") |
|
|
|
|
|
def read_audio(file_path): |
|
|
audio, sample_rate = sf.read(file_path) |
|
|
if audio.ndim > 1: |
|
|
audio = audio[:, 0] |
|
|
if sample_rate != 16000: |
|
|
audio = lb.resample(audio, orig_sr=sample_rate, target_sr=16000) |
|
|
sample_rate = 16000 |
|
|
|
|
|
return torch.tensor(audio).float().squeeze().unsqueeze(0) |
|
|
|
|
|
def stft(audio): |
|
|
transform = InputSTFT( |
|
|
n_fft=512, |
|
|
n_win=512, |
|
|
n_hop=128, |
|
|
normalize=False, |
|
|
center=True, |
|
|
onesided=True, |
|
|
online=False |
|
|
).to(DEVICE) |
|
|
return transform(audio) |
|
|
|
|
|
def mel_transform(audio, X_norm): |
|
|
transform = TargetMel( |
|
|
sample_rate=16000, |
|
|
n_fft=512, |
|
|
n_win=512, |
|
|
n_hop=128, |
|
|
n_mels=80, |
|
|
f_min=0, |
|
|
f_max=8000, |
|
|
power=2, |
|
|
center=True, |
|
|
normalize=False, |
|
|
onesided=True, |
|
|
mel_norm="slaney", |
|
|
mel_scale="slaney", |
|
|
librosa_mel=True, |
|
|
online=False |
|
|
).to(DEVICE) |
|
|
return transform(audio, X_norm) |
|
|
|
|
|
def load_cleanmel(model_name): |
|
|
model_config = f"./configs/cleanmel_offline.yaml" |
|
|
model_config = yaml.safe_load(open(model_config, "r"))["model"]["arch"]["init_args"] |
|
|
cleanmel = CleanMel(**model_config) |
|
|
cleanmel.load_state_dict(torch.load(f"./ckpts/CleanMel/{model_name}.ckpt")) |
|
|
return cleanmel.eval() |
|
|
|
|
|
def load_vocos(model_name): |
|
|
vocos = Vocos.from_hparams(config_path="./configs/vocos_offline.yaml") |
|
|
vocos = Vocos.from_pretrained(None, model_path=f"./ckpts/Vocos/{model_name}.pt", model=vocos) |
|
|
return vocos.eval() |
|
|
|
|
|
def get_mrm_pred(Y_hat, x, X_norm): |
|
|
X_noisy = mel_transform(x, X_norm) |
|
|
Y_hat = Y_hat.squeeze() |
|
|
Y_hat = torch.square(Y_hat * (torch.sqrt(X_noisy) + 1e-10)) |
|
|
return Y_hat |
|
|
|
|
|
def safe_log(x): |
|
|
return torch.log(torch.clip(x, min=1e-5)) |
|
|
|
|
|
@spaces.GPU |
|
|
@torch.inference_mode() |
|
|
def enhance_cleanmel_L_mask(audio_path): |
|
|
model = load_cleanmel("offline_CleanMel_L_mask").to(DEVICE) |
|
|
vocos = load_vocos("vocos_offline").to(DEVICE) |
|
|
x = read_audio(audio_path).to(DEVICE) |
|
|
X, X_norm = stft(x) |
|
|
Y_hat = model(X) |
|
|
MRM_hat = torch.sigmoid(Y_hat) |
|
|
Y_hat = get_mrm_pred(MRM_hat, x, X_norm) |
|
|
logMel_hat = safe_log(Y_hat) |
|
|
y_hat = vocos(logMel_hat, X_norm) |
|
|
with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp_file: |
|
|
sf.write(tmp_file.name, y_hat.squeeze().cpu().numpy(), 16000) |
|
|
with tempfile.NamedTemporaryFile(suffix='.npy', delete=False) as tmp_logmel_np_file: |
|
|
np.save(tmp_logmel_np_file.name, logMel_hat.squeeze().cpu().numpy()) |
|
|
logMel_img = logMel_hat.squeeze().cpu().numpy()[::-1, :] |
|
|
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_logmel_img: |
|
|
|
|
|
plt.figure(figsize=(logMel_img.shape[1] / 100, logMel_img.shape[0] / 50)) |
|
|
plt.clf() |
|
|
plt.imshow(logMel_img, vmin=-11, cmap="jet") |
|
|
plt.tight_layout() |
|
|
plt.ylabel("Mel bands") |
|
|
plt.xlabel("Time (second)") |
|
|
plt.yticks([0, 80], [80, 0]) |
|
|
dur = x.shape[-1] / 16000 |
|
|
xticks = [int(x) for x in np.linspace(0, logMel_img.shape[-1], 11)] |
|
|
xticks_str = ["{:.1f}".format(x) for x in np.linspace(0, dur, 11)] |
|
|
plt.xticks(xticks, xticks_str) |
|
|
plt.savefig(tmp_logmel_img.name) |
|
|
|
|
|
return tmp_file.name, tmp_logmel_img.name, tmp_logmel_np_file.name |
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo = gr.Blocks() |
|
|
with gr.Blocks(title="CleanMel Demo") as demo: |
|
|
gr.Markdown("## CleanMel Demo") |
|
|
gr.Markdown("This demo showcases the CleanMel model for speech enhancement.") |
|
|
|
|
|
with gr.Row(): |
|
|
audio_input = gr.Audio(label="Input Audio", type="filepath", sources="upload") |
|
|
enhance_button = gr.Button("Enhance Audio") |
|
|
|
|
|
output_audio = gr.Audio(label="Enhanced Audio", type="filepath") |
|
|
output_mel = gr.Image(label="Output LogMel Spectrogram", type="filepath", visible=True) |
|
|
output_np = gr.File(label="Enhanced LogMel Spec. (.npy)", type="filepath") |
|
|
|
|
|
enhance_button.click( |
|
|
enhance_cleanmel_L_mask, |
|
|
inputs=audio_input, |
|
|
outputs=[output_audio, output_mel, output_np] |
|
|
) |
|
|
|
|
|
demo.launch(debug=False, share=True) |
