HY-2012

AudioSR.AXERA full board package: models + bin + python + samples

6149fac verified about 15 hours ago

6.85 kB

"""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)

	"""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)