File size: 6,849 Bytes
6149fac | 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 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | """AudioSR speech super-resolution on AX650 (axengine).
The SDK accepts 8/16/24/48 kHz mono/stereo WAV and produces 48 kHz output.
It implements the upstream AudioSR preprocessing (5.12 s padding, log-mel,
low-pass conditioning, DDIM-10) in NumPy/librosa, then runs the AX650
axmodels with axengine.
"""
from __future__ import annotations
import argparse
import math
import os
from pathlib import Path
import numpy as np
import soundfile as sf
from scipy.signal import butter, resample_poly, sosfiltfilt
import librosa
SR = 48000
N_FFT = 2048
HOP = 480
WIN = 2048
N_MEL = 256
FMIN = 20
FMAX = 24000
SEG = 5.12
SCALE = 0.3342401385307312
UNCOND = -11.4981
def _normalize(y):
y = y - np.mean(y)
return (y / (np.max(np.abs(y)) + 1e-8) * 0.5).astype(np.float32)
def _pad_wav(y, target):
out = np.zeros(target, dtype=np.float32)
out[:min(target, len(y))] = y[:target]
return out
def _mel(y):
# same reflect padding as audiosr.utils.mel_spectrogram_train
p = (N_FFT - HOP) // 2
yy = np.pad(y, (p, p), mode="reflect")
stft = librosa.stft(yy, n_fft=N_FFT, hop_length=HOP, win_length=WIN,
window="hann", center=False)
mag = np.abs(stft)
basis = librosa.filters.mel(sr=SR, n_fft=N_FFT, n_mels=N_MEL, fmin=FMIN, fmax=FMAX)
mel = np.log(np.maximum(basis @ mag, 1e-5))
return mel.astype(np.float32), mag.astype(np.float32)
def _cutoff(stft_mag, percentile=0.985):
energy = np.cumsum(stft_mag.sum(axis=1)) # [freq], sum over time
idx = 0
for i in range(1, energy.shape[0]):
if energy[-i] < energy[-1] * percentile:
idx = energy.shape[0] - i
break
cutoff = idx / 1024.0 * 24000.0
return 24000.0 if cutoff < 1000.0 else cutoff
def _lowpass(y, cutoff):
sos = butter(8, cutoff, btype="lowpass", fs=SR, output="sos")
return sosfiltfilt(sos, y).astype(np.float32)
def make_batch(path):
y, sr = sf.read(path, always_2d=False)
if y.ndim > 1:
y = y.mean(axis=1)
y = np.asarray(y, dtype=np.float32)
if sr != SR:
y = resample_poly(y, SR, sr).astype(np.float32)
duration = len(y) / SR
pad_duration = duration + ((SEG - duration % SEG) if duration % SEG else 0.0)
target_frame = int(pad_duration * 100)
y = _normalize(y)
y = _pad_wav(y, int(SR * pad_duration))
mel, stft = _mel(y)
mel = mel.T # [T, 256]
stft = stft.T
if mel.shape[0] < target_frame:
mel = np.pad(mel, ((0, target_frame - mel.shape[0]), (0, 0)))
stft = np.pad(stft, ((0, target_frame - stft.shape[0]), (0, 0)))
else:
mel, stft = mel[:target_frame], stft[:target_frame]
if mel.shape[-1] % 2:
mel, stft = mel[..., :-1], stft[..., :-1]
cutoff = _cutoff(stft, 0.985)
low = _lowpass(y, cutoff)
low_mel, _ = _mel(low)
low_mel = low_mel.T
if low_mel.shape[0] < target_frame:
low_mel = np.pad(low_mel, ((0, target_frame - low_mel.shape[0]), (0, 0)))
else:
low_mel = low_mel[:target_frame]
if low_mel.shape[-1] % 2:
low_mel = low_mel[..., :-1]
return mel[None, None], low_mel[None, None], low, cutoff
class AudioSR:
def __init__(self, model_dir, seed=42, ddim="models/audiosr_ddim_10.npz"):
import axengine
self.axengine = axengine
self.dir = Path(model_dir)
self.enc = axengine.InferenceSession(str(self.dir / "audiosr_vae_encoder_nhwc.axmodel"))
self.unet = axengine.InferenceSession(str(self.dir / "audiosr_unet.axmodel"))
self.dec = axengine.InferenceSession(str(self.dir / "audiosr_vae_decoder.axmodel"))
self.voc = [axengine.InferenceSession(str(self.dir / n)) for n in [
"voc_pre_fp32.axmodel", "voc_up0a_mix.axmodel", "voc_up0b_mix.axmodel",
"voc_up1a_mix.axmodel", "voc_up1b_mix.axmodel", "voc_split_nhwc_up2_mix.axmodel",
"voc_split_nhwc_up3_mix.axmodel", "voc_split_nhwc_up4_mix.axmodel",
"voc_split_nhwc_post_fp32.axmodel"]]
self.d = np.load(ddim)
self.rng = np.random.RandomState(seed)
def _encode(self, mel):
o = self.enc.run(None, {"mel": np.ascontiguousarray(mel[0, 0, :, :, None].astype(np.float32))})
return (o[0] + self.rng.randn(*o[0].shape).astype(np.float32) *
np.exp(0.5 * o[1])).astype(np.float32) * SCALE
def enhance(self, path):
mel, low_mel, low_wav, _ = make_batch(path)
zc = self._encode(mel)
zl = self._encode(low_mel)
x = self.rng.randn(1, 16, 64, 32).astype(np.float32)
uncond = np.full((1, 16, 64, 32), UNCOND, np.float32)
ts = np.flip(self.d["timesteps"])
n = len(ts)
for i, step in enumerate(ts):
t = np.array([float(step)], np.float32)
def run(c):
return self.unet.run(None, {"x": np.ascontiguousarray(np.concatenate([x, c], 1)), "t": t})[0]
vu, vc = run(uncond), run(zc)
v = vu + 3.5 * (vc - vu)
sc = float(self.d["sqrt_alphas_cumprod"][step])
so = float(self.d["sqrt_one_minus_alphas_cumprod"][step])
et = sc * v + so * x
px0 = sc * x - so * v
idx = n - 1 - i
at, ap = float(self.d["alphas"][idx]), float(self.d["alphas_prev"][idx])
sig = float(self.d["sigmas"][idx])
x = (np.sqrt(ap) * px0 + np.sqrt(max(0.0, 1.0 - ap - sig * sig)) * et +
sig * self.rng.randn(*x.shape).astype(np.float32)).astype(np.float32)
dm = self.dec.run(None, {"z": np.ascontiguousarray(x / SCALE)})[0]
# mel bin cutoff, using upstream mel_replace_ops
e = np.cumsum(np.exp(low_mel[0]).sum(axis=1))
limit = e[-1] * 0.985
cb = 0
for i in range(1, e.shape[0]):
if e[-i] < limit:
cb = e.shape[0] - i
break
dm[0, :, :cb, :] = low_mel[0, :cb, :]
voc = dm.squeeze(1).transpose(0, 2, 1).astype(np.float32) # [1, 256, 512]
x = np.ascontiguousarray(voc.transpose(0, 2, 1)) # NHWC [1, 512, 256]
for s in self.voc:
x = s.run(None, {s.get_inputs()[0].name: x})[0]
wav = x[0, 0][:low_wav.size]
return postprocess(wav, low_wav), SR
def postprocess(out, src):
stft_gt = librosa.stft(src)
stft_out = librosa.stft(out)
e = np.cumsum(np.sum(np.abs(stft_gt), axis=-1))
limit = e[-1] * 0.985
c = 0
for i in range(1, e.shape[0]):
if e[-i] < limit:
c = e.shape[0] - i
break
er = np.mean(np.sum(np.abs(stft_gt[c])) / np.sum(np.abs(stft_out[c])))
er = min(max(er, 0.8), 1.2)
stft_out[:c] = stft_gt[:c] / er
out = librosa.istft(stft_out, length=out.shape[0])
out = 0.5 * out / max(float(np.max(np.abs(out))), 1e-8)
return (out - float(np.mean(out))).astype(np.float32)
|