Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
103 lines
4.8 KiB
Python
103 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Export official features, frame probabilities and timestamps; measure warmed inference."""
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
|
|
|
|
def median_ms(fn, cuda=False):
|
|
for _ in range(5):
|
|
fn()
|
|
times = []
|
|
for _ in range(21):
|
|
if cuda:
|
|
torch.cuda.synchronize()
|
|
start = time.perf_counter()
|
|
fn()
|
|
if cuda:
|
|
torch.cuda.synchronize()
|
|
times.append((time.perf_counter() - start) * 1000)
|
|
return float(np.median(times))
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--reference", type=Path, required=True)
|
|
parser.add_argument("--source", type=Path, required=True)
|
|
parser.add_argument("--audio", type=Path, required=True)
|
|
parser.add_argument("--output-prefix", type=Path, required=True)
|
|
parser.add_argument("--threads", type=int, default=8)
|
|
parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
|
|
parser.add_argument("--log", action="store_true")
|
|
parser.add_argument("--gguf", type=Path, help="Use dequantized GGUF weights to isolate backend error from quantization drift.")
|
|
parser.add_argument("--max-speech-frames", type=int, default=2000)
|
|
args = parser.parse_args()
|
|
sys.path.insert(0, str(args.reference.resolve()))
|
|
from fireredvad.vad import FireRedVad, FireRedVadConfig
|
|
from fireredvad.stream_vad import FireRedStreamVad, FireRedStreamVadConfig
|
|
torch.set_num_threads(args.threads)
|
|
config = torch.load(args.source / "model.pth.tar", map_location="cpu", weights_only=False)["args"]
|
|
causal = config.N2 == 0
|
|
cuda = args.device == "cuda"
|
|
options = (FireRedStreamVadConfig if causal else FireRedVadConfig)(max_speech_frame=args.max_speech_frames, use_gpu=cuda)
|
|
model = (FireRedStreamVad if causal else FireRedVad).from_pretrained(str(args.source), options)
|
|
if args.gguf:
|
|
import gguf
|
|
tensors = {}
|
|
shapes = {name: tuple(value.shape) for name, value in model.vad_model.state_dict().items()}
|
|
for tensor in gguf.GGUFReader(args.gguf).tensors:
|
|
name = tensor.name.removeprefix("weights/")
|
|
shape = shapes.get(name, tuple(reversed(tensor.shape)))
|
|
value = gguf.quants.dequantize(tensor.data, tensor.tensor_type).reshape(shape)
|
|
tensors[name] = torch.from_numpy(np.array(value, dtype=np.float32, copy=True))
|
|
model.vad_model.load_state_dict({name: tensors[name] for name in model.vad_model.state_dict()}, strict=True)
|
|
model.audio_feat.cmvn.means = -tensors["frontend.cmvn_shift"].numpy().astype(np.float64)
|
|
model.audio_feat.cmvn.inverse_std_variances = tensors["frontend.cmvn_scale"].numpy().astype(np.float64)
|
|
pcm, rate = sf.read(args.audio, dtype="int16")
|
|
if rate != 16000 or pcm.ndim != 1:
|
|
raise ValueError("Expected mono 16 kHz PCM")
|
|
audio = (pcm, rate)
|
|
detect = (lambda: model.detect_full(audio)) if causal else (lambda: model.detect(audio))
|
|
cold_start = time.perf_counter()
|
|
cold_result = detect()
|
|
cold_ms = (time.perf_counter() - cold_start) * 1000
|
|
features, duration = model.audio_feat.extract(audio)
|
|
args.output_prefix.parent.mkdir(parents=True, exist_ok=True)
|
|
prefix = str(args.output_prefix)
|
|
with torch.inference_mode():
|
|
network_features = features.to(args.device)
|
|
probs = model.vad_model(network_features.unsqueeze(0))[0].flatten()
|
|
features.numpy().astype("<f4").tofile(prefix + ".features.f32")
|
|
probs.cpu().numpy().astype("<f4").tofile(prefix + ".probs.f32")
|
|
if causal:
|
|
cached = None
|
|
chunks = []
|
|
for chunk in network_features.split(16):
|
|
chunk_probs, cached = model.vad_model(chunk.unsqueeze(0), cached)
|
|
chunks.append(chunk_probs.flatten())
|
|
torch.cat(chunks).cpu().numpy().astype("<f4").tofile(prefix + ".chunked_probs.f32")
|
|
_, result = cold_result
|
|
else:
|
|
result, _ = cold_result
|
|
Path(prefix + ".segments.json").write_text(json.dumps(result["timestamps"]) + "\n")
|
|
metrics = {"threads": args.threads, "duration": duration, "frames": features.shape[0],
|
|
"cold_ms": cold_ms,
|
|
"frontend_ms": median_ms(lambda: model.audio_feat.extract(audio)),
|
|
"network_ms": median_ms(lambda: model.vad_model(network_features.unsqueeze(0)), cuda),
|
|
"total_ms": median_ms(detect, cuda)}
|
|
metrics["rtf"] = metrics["total_ms"] / (duration * 1000)
|
|
if args.log:
|
|
print(f"session.wall_ms={metrics['total_ms']:.6f}")
|
|
Path(prefix + ".metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
|
|
print(json.dumps(metrics))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|