File size: 6,316 Bytes
dd2737e 69261fe dd2737e 1267a8f dd2737e 69261fe dd2737e 69261fe | 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 | """FastWave EDM sampler on AX650 (pyaxengine) — numpy only.
The AXMODEL is one function evaluation (NFE):
y = D(x, sigma) = c_skip * x + c_out * F(c_in * x, audio_low, band, noise_emb)
Everything that only depends on sigma (the 8 EDM steps) is precomputed on the
host: c_skip / c_in / c_out and the diffusion embedding.
"""
import json
import os
import numpy as np
from . import audio as A
from .postfilter import denoise as _denoise
ASSETS = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
def _silu(x):
return x / (1.0 + np.exp(-x))
class FastWaveSR:
def __init__(self, model_path, assets_dir=ASSETS, num_steps=8, provider="AxEngineExecutionProvider",
chunk=None, overlap=0.5, seed=1234, sigma_min=0.002, sigma_max=80.0, rho=7.0):
import axengine as axe
self.sess = axe.InferenceSession(model_path, providers=[provider])
self.input_names = [i.name for i in self.sess.get_inputs()]
meta = json.load(open(os.path.join(assets_dir, "model_config.json")))
self.meta = meta
self.chunk = int(chunk or meta.get("chunk_samples", 32768))
if self.chunk != int(meta.get("chunk_samples", 32768)):
print(f"[warn] chunk={self.chunk} differs from the compiled axmodel "
f"({meta.get('chunk_samples')}); recompile the model for a new chunk size")
self.overlap = float(overlap)
self.hop = (max(256, int(round(self.chunk * overlap)) // 256 * 256)
if overlap > 0 else self.chunk)
self.num_steps = int(num_steps)
self.rng = np.random.default_rng(seed)
self.sr_out = int(meta.get("sample_rate", A.SR_OUT))
p1w = np.load(os.path.join(assets_dir, "emb_p1_w.npy"))
p1b = np.load(os.path.join(assets_dir, "emb_p1_b.npy"))
p2w = np.load(os.path.join(assets_dir, "emb_p2_w.npy"))
p2b = np.load(os.path.join(assets_dir, "emb_p2_b.npy"))
half = int(p1w.shape[1]) // 2
scale = float(meta.get("pos_emb_scale", 50000.0))
freqs = np.exp(np.arange(half, dtype=np.float64) * (-np.log(10000.0) / (half - 1)))
idx = np.arange(self.num_steps, dtype=np.float64)
t = (sigma_max ** (1 / rho) + idx / (self.num_steps - 1)
* (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho
self.t_steps = np.concatenate([t, [0.0]]).astype(np.float32)
self.steps = []
for s in self.t_steps[:-1].astype(np.float64):
c_skip = 0.5 ** 2 / (s ** 2 + 0.5 ** 2)
c_in = 1.0 / np.sqrt(s ** 2 + 0.5 ** 2)
c_out = s * 0.5 / np.sqrt(s ** 2 + 0.5 ** 2)
c_noise = np.log(s) / 4.0
ang = scale * c_noise * freqs
emb = np.concatenate([np.sin(ang), np.cos(ang)]).astype(np.float32)
emb = _silu(emb @ p1w.T + p1b)
emb = _silu(emb @ p2w.T + p2b)
self.steps.append({
"c_skip": np.array([[c_skip]], dtype=np.float32),
"c_in": np.array([[c_in]], dtype=np.float32),
"c_out": np.array([[c_out]], dtype=np.float32),
"noise_emb": emb.reshape(1, -1).astype(np.float32),
})
# ------------------------------------------------------------------ graph
def nfe(self, x, low, band, step):
"""One denoising step; x/low: (1, chunk) float32."""
s = self.steps[step]
feed = {
"x": np.ascontiguousarray(x, dtype=np.float32),
"audio_low": low,
"band_onehot": band,
"c_skip": s["c_skip"],
"c_in": s["c_in"],
"c_out": s["c_out"],
"noise_emb": s["noise_emb"],
}
return self.sess.run(None, {k: feed[k] for k in self.input_names})[0]
def denoise_chunk(self, low_chunk, band):
x = self.rng.standard_normal((1, self.chunk)).astype(np.float32) * self.t_steps[0]
for i in range(self.num_steps):
t_cur = self.t_steps[i]
t_next = self.t_steps[i + 1]
denoised = self.nfe(x, low_chunk, band, i)
x = x + (t_next - t_cur) * (x - denoised) / t_cur
return np.clip(x[0], -1.0, 1.0 - 1e-7)
# --------------------------------------------------------------- pipeline
def enhance(self, wav, sr_in, progress=None, normalize=True, denoise=False, denoise_alpha=1.0):
"""wav: 1-D float32 at sr_in. Returns 48 kHz float32."""
wav = np.asarray(wav, dtype=np.float32)
peak = float(np.max(np.abs(wav))) if len(wav) else 0.0
if normalize and peak > 1e-8:
wav = wav / peak
else:
peak = 1.0
low = A.low_quality(wav, int(sr_in), self.sr_out)
band = A.band_onehot(int(sr_in), self.sr_out)
band3 = band.reshape(1, 2, -1)
n = len(low)
if n == 0:
return low
if self.overlap <= 0:
# fastest: pad to whole chunks, no overlap / no windowing
n_pad = int(np.ceil(n / self.chunk)) * self.chunk
if n_pad > n:
low = np.pad(low, (0, n_pad - n))
starts = list(range(0, n_pad, self.chunk))
outs = []
for i, s in enumerate(starts):
seg = low[s: s + self.chunk]
outs.append(self.denoise_chunk(
np.ascontiguousarray(seg.reshape(1, -1)), band3))
if progress:
progress(i + 1, len(starts))
y = np.concatenate(outs)[:n].astype(np.float32)
if denoise:
y = _denoise(y, self.sr_out, alpha=denoise_alpha)
return y * peak
starts = A.chunk_bounds(n, self.chunk, self.hop)
outs = []
for i, s in enumerate(starts):
seg = low[s: s + self.chunk]
if len(seg) < self.chunk:
seg = np.pad(seg, (0, self.chunk - len(seg)))
outs.append(self.denoise_chunk(np.ascontiguousarray(seg.reshape(1, -1)), band3))
if progress:
progress(i + 1, len(starts))
if len(starts) == 1:
y = outs[0][:n].astype(np.float32)
else:
y = A.overlap_add(outs, starts, n, self.chunk)
if denoise:
y = _denoise(y, self.sr_out, alpha=denoise_alpha)
return y * peak
|