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