| |
| """Board benchmark: single-NFE latency and end-to-end RTF.""" |
| import argparse |
| import os |
| import sys |
| import time |
|
|
| import numpy as np |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, HERE) |
|
|
| from fastwave_sdk import FastWaveSR |
| from fastwave_sdk import audio as A |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", default=os.path.join(HERE, "..", "models", "fastwave_nfe.axmodel")) |
| ap.add_argument("--wav", default=os.path.join(HERE, "..", "samples", "listen", "8k", |
| "OSR_us_000_0010_input_8k.wav")) |
| ap.add_argument("--seconds", type=float, default=3.0) |
| ap.add_argument("--steps", type=int, default=8) |
| ap.add_argument("--iters", type=int, default=20) |
| ap.add_argument("--overlap", type=float, default=0.5) |
| a = ap.parse_args() |
|
|
| m = FastWaveSR(a.model, num_steps=a.steps, overlap=a.overlap) |
| wav, sr = A.read_wav(a.wav) |
| wav = wav[: int(a.seconds * sr)] |
| low = A.low_quality(wav / (np.abs(wav).max() + 1e-8), sr) |
| band = A.band_onehot(sr).reshape(1, 2, -1) |
| seg = np.ascontiguousarray(low[: m.chunk].reshape(1, -1)) |
| x = np.zeros((1, m.chunk), dtype=np.float32) |
|
|
| for _ in range(5): |
| m.nfe(x, seg, band, 0) |
| ts = [] |
| for i in range(a.iters): |
| t0 = time.perf_counter() |
| m.nfe(x, seg, band, i % m.num_steps) |
| ts.append((time.perf_counter() - t0) * 1e3) |
| ts = np.array(ts) |
| print(f"NFE latency: mean {ts.mean():.2f} ms min {ts.min():.2f} max {ts.max():.2f}") |
|
|
| t0 = time.perf_counter() |
| out = m.enhance(wav, sr) |
| dt = time.perf_counter() - t0 |
| dur = len(out) / 48000.0 |
| print(f"end-to-end {dur:.2f}s audio in {dt:.2f}s -> RTF {dt/dur:.4f}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
