130 lines
6.5 KiB
Python
130 lines
6.5 KiB
Python
"""Run MaAI ch_kyoto synchronously on a local stereo recording, on CPU."""
|
||
import argparse
|
||
import hashlib
|
||
import importlib.metadata
|
||
import json
|
||
import os
|
||
from pathlib import Path
|
||
import time
|
||
|
||
HERE = Path(__file__).resolve().parent
|
||
ROOT = HERE.parents[1]
|
||
os.environ.setdefault("HF_HOME", str(HERE / ".cache/huggingface"))
|
||
os.environ.setdefault("TORCH_HOME", str(HERE / ".cache/torch"))
|
||
os.environ.setdefault("MPLCONFIGDIR", str(HERE / ".cache/matplotlib"))
|
||
os.environ.setdefault("PYGAME_HIDE_SUPPORT_PROMPT", "1")
|
||
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
||
|
||
import numpy as np
|
||
import soundfile as sf
|
||
import soxr
|
||
import torch
|
||
from maai import Maai, MaaiInput
|
||
|
||
WEIGHTS = HERE / "models/vap_mc_state_dict_ch_kyoto_10hz_20000msec.pt"
|
||
CPC = HERE / "models/60k_epoch4-d0f474de.pt"
|
||
EXPECTED = {
|
||
WEIGHTS.name: "65b8a51aea6a257b5b9139107f0f0cf22f94523cd4b6a2e7a356b6c43a4596a8",
|
||
CPC.name: "d0f474de968c7243c20f43dee1626681618f2f150d3e638e82b79e22563b17f1",
|
||
}
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description=__doc__)
|
||
parser.add_argument("--audio", type=Path, default=ROOT / "data/韩泰机器人抢话/音频/2c83253a-7977-4e3c-a387-0aa1ea06fbe1.mp3")
|
||
parser.add_argument("--seconds", type=float, default=22)
|
||
parser.add_argument("--threads", type=int, default=4)
|
||
parser.add_argument("--output", type=Path, default=HERE / "results/smoke.json")
|
||
args = parser.parse_args()
|
||
if args.seconds <= 0 or args.threads < 1:
|
||
parser.error("seconds and threads must be positive")
|
||
hashes = {}
|
||
for path in (WEIGHTS, CPC):
|
||
if not path.exists():
|
||
raise FileNotFoundError(f"Download the official weights first: {path}")
|
||
hashes[path.name] = hashlib.sha256(path.read_bytes()).hexdigest()
|
||
if hashes[path.name] != EXPECTED[path.name]:
|
||
raise ValueError(f"Weight checksum mismatch: {path}")
|
||
torch.set_num_threads(args.threads)
|
||
torch.set_num_interop_threads(1)
|
||
audio, original_sr = sf.read(args.audio, dtype="float32", always_2d=True)
|
||
if audio.shape[1] != 2:
|
||
raise ValueError("Expected stereo audio: channel 0 user, channel 1 robot")
|
||
audio = audio[:round(args.seconds * original_sr)]
|
||
if original_sr != 16000:
|
||
audio = soxr.resample(audio, original_sr, 16000, quality="HQ")
|
||
audio = np.ascontiguousarray(audio, dtype=np.float32)
|
||
|
||
load_started = time.perf_counter()
|
||
# Official CPC checkpoint contains an argparse.Namespace. Allow only
|
||
# that class for this load, while retaining weights_only safe loading.
|
||
with torch.serialization.safe_globals([argparse.Namespace]):
|
||
model = Maai(
|
||
mode="vap", lang="ch_kyoto",
|
||
audio_ch1=MaaiInput.Chunk(), audio_ch2=MaaiInput.Chunk(),
|
||
frame_rate=10, context_len_sec=20, device="cpu",
|
||
model_type="normal", mc=True, use_kv_cache=True,
|
||
cpc_model=str(CPC), local_model=str(WEIGHTS),
|
||
cache_dir=str(HERE / ".cache/huggingface"),
|
||
)
|
||
load_seconds = time.perf_counter() - load_started
|
||
frame_samples = model.audio_samples_per_frame
|
||
frame_count = len(audio) // frame_samples
|
||
if frame_count < 1:
|
||
raise ValueError("Audio must contain at least one 100 ms frame")
|
||
frames, latencies = [], []
|
||
inference_started = time.perf_counter()
|
||
for i in range(frame_count):
|
||
start = i*frame_samples
|
||
before = time.perf_counter()
|
||
model.process(np.ascontiguousarray(audio[start:start+frame_samples, 0]),
|
||
np.ascontiguousarray(audio[start:start+frame_samples, 1]))
|
||
latency_ms = (time.perf_counter() - before) * 1000
|
||
# A complete CPC frame should always produce exactly one result.
|
||
# Avoid blocking indefinitely if that assumption fails.
|
||
result = model.result_dict_queue.get_nowait()
|
||
assert model.result_dict_queue.empty(), "More than one result per frame"
|
||
row = dict(audio_end_sec=(i+1)*frame_samples/16000,
|
||
inference_ms=latency_ms)
|
||
for key in ("p_now", "p_future", "vad"):
|
||
values = np.asarray(result[key], dtype=np.float64).reshape(-1)
|
||
if len(values) != 2 or not np.isfinite(values).all():
|
||
raise ValueError(f"Invalid {key}: {values}")
|
||
if not ((values >= 0) & (values <= 1)).all():
|
||
raise ValueError(f"Out-of-range {key}: {values}")
|
||
row[key] = values.tolist()
|
||
frames.append(row)
|
||
latencies.append(latency_ms)
|
||
inference_seconds = time.perf_counter() - inference_started
|
||
audio_seconds = frame_count*frame_samples/16000
|
||
warm = latencies[min(5, len(latencies)-1):]
|
||
report = dict(
|
||
status="success", audio=str(args.audio.resolve()),
|
||
original_sample_rate=original_sr, model_sample_rate=16000,
|
||
audio_seconds=audio_seconds, frame_count=frame_count,
|
||
model="ch_kyoto noise-robust VAP / CPC", weight_variant="vap_mc",
|
||
maai_version=importlib.metadata.version("maai"), torch_version=torch.__version__,
|
||
portaudio_dependency="Homebrew PortAudio 19.7.0 / PyAudio 0.2.14",
|
||
device="cpu", threads=args.threads, frame_rate=10, context_len_sec=20,
|
||
use_kv_cache=True, model_load_seconds=load_seconds,
|
||
inference_seconds=inference_seconds, realtime_factor=inference_seconds/audio_seconds,
|
||
warm_mean_frame_ms=float(np.mean(warm)), warm_p95_frame_ms=float(np.percentile(warm, 95)),
|
||
weight_sha256=hashes,
|
||
output_definitions={
|
||
"channels": "Index 0=user, index 1=robot; separate stereo channels, no mixing.",
|
||
"p_now": "Relative activity projection for the two speakers over roughly the next 0–0.6 seconds.",
|
||
"p_future": "Relative activity projection for the two speakers over roughly the next 0.6–2 seconds.",
|
||
"vad": "Current voice-activity probability for each speaker.",
|
||
"audio_end_sec": "Media time measured from file start; MaAI wall-clock result t is not used.",
|
||
"scope": "Basic offline CPU smoke test; VAP projections are not can_reply or EOT probabilities. Timings exclude model loading and resampling and describe this one recording/run only.",
|
||
}, frames=frames)
|
||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||
args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2)+"\n")
|
||
print(json.dumps({k:v for k,v in report.items() if k not in ("frames", "output_definitions", "weight_sha256")}, ensure_ascii=False, indent=2))
|
||
print("Last frame:", json.dumps(frames[-1], ensure_ascii=False))
|
||
print(f"Saved {args.output.resolve()}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|