66 lines
1.8 KiB
Python
66 lines
1.8 KiB
Python
#!/usr/bin/env python3
|
|
"""RTF bench for hop/window speechbrain enhance.
|
|
|
|
Stop livekit-cameras.service first or the numbers fight the fleet.
|
|
|
|
.venv/bin/python scripts/bench_enhance.py
|
|
.venv/bin/python scripts/bench_enhance.py --stereo --hop 0.1
|
|
"""
|
|
import argparse
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
import numpy as np
|
|
|
|
from audio_cleanup import AudioCleaner
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--stereo", action="store_true")
|
|
ap.add_argument("--hop", type=float, default=0.2)
|
|
ap.add_argument("--chunk", type=float, default=1.0)
|
|
ap.add_argument("--n", type=int, default=50)
|
|
args = ap.parse_args()
|
|
|
|
c = AudioCleaner(
|
|
mode="force",
|
|
beamform="force" if args.stereo else "off",
|
|
chunk_s=args.chunk,
|
|
hop_s=args.hop,
|
|
torch_num_threads=1,
|
|
)
|
|
hop_n = int(args.hop * 16000)
|
|
if args.stereo:
|
|
t = np.arange(hop_n)
|
|
ch0 = (np.sin(2 * np.pi * 440 * t / 16000) * 800).astype(np.int16)
|
|
pcm = np.stack([ch0, np.roll(ch0, 2)], axis=1)
|
|
else:
|
|
pcm = (np.random.randn(hop_n) * 800).astype(np.int16)
|
|
warmup = max(8, int(args.chunk / args.hop) + 2)
|
|
for _ in range(warmup):
|
|
c.process(pcm)
|
|
n = args.n
|
|
t0 = time.perf_counter()
|
|
for _ in range(n):
|
|
c.process(pcm)
|
|
dt = time.perf_counter() - t0
|
|
audio_s = n * args.hop
|
|
rtf = dt / audio_s
|
|
print(
|
|
f"backend={c.backend} stereo={args.stereo} hops={n} "
|
|
f"wall={dt:.3f}s audio={audio_s:.1f}s rtf={rtf:.3f} "
|
|
f"last_ms={c.last_dt_ms:.1f}"
|
|
)
|
|
limit = 0.5 if args.stereo else 0.8
|
|
if rtf > limit:
|
|
raise SystemExit(f"RTF too high for 20-way realtime (limit {limit})")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|