| |
| |
| """The release's two networks, run with the release's model code, flags and batch shapes. |
| |
| The networks are the pure-PyTorch copies in `seg/` (DiariZen's segmentation model) and `voiceprint.py` (pyannote's |
| WeSpeakerResNet34 with torchaudio's kaldi fbank), loaded from the bundle's safetensors (`tools/export_bundle.py`); no |
| DiariZen, pyannote, Lightning or torchaudio import. They run with the release's settings (FP32, TF32 off, |
| deterministic kernels, `CUBLAS_WORKSPACE_CONFIG=:4096:8`, 2 CPU threads), one copy per device, and their outputs are |
| bit-identical to the release's on every batch shape where the release's setup is reproduced (the start-up self-test |
| against assets/nn_golden.json). A *slot* is one device with one CUDA stream, driven by one engine worker thread: every |
| GPU gives the same bytes as the others of its kind, but two streams sharing one model on one GPU do not (about 1 batch |
| in 1,000 differed), so there is exactly one slot per device. What runs on a slot is always one of the release's own batch shapes: |
| |
| segment(x) (B, 1, 256000) -> (B, 799, 16) B = 8, the remainder (1-7), or 1 (padded last window) |
| embed(crops, weights) (B, 1, 256000), (B, 799) -> (B, 256) the release's voiceprint call (B = 8 or 4) |
| embed_repacked(...) the voiceprint trunk once per window (batch 8), then pooling + the last layer at the |
| release's rows (2 windows x 4 speakers); bit-identical to `embed` |
| resample(audio, sr) `load_audio`'s torchaudio call (torchaudio's resample, copied in resample.py) |
| |
| The engine (Rust) decides every batch; this module only executes them. |
| """ |
| import contextlib |
| import hashlib |
| import json |
| import os |
| import threading |
| import warnings |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| WIN, FRAMES, CLASSES, EMB = 256000, 799, 16, 256 |
| FRAMES_SW = (-0.00753125, 0.025, 0.02) |
| MIN_NUM_SAMPLES = 400 |
|
|
|
|
| |
| |
| |
| PRESETS = {"exact": ("exact", "exact"), "fp32": ("fp32", "fp32"), "mixed": ("fp32", "fp16"), |
| "tf32conv": ("fp32c", "fp16"), "ffn16": ("ffn16", "fp16"), "tf32": ("tf32", "tf32"), "fp16": ("fp16", "fp16")} |
|
|
|
|
| |
| |
| |
| |
| |
| NEAR = dict(max_abs_dp=0.25, frame_agreement=0.99, min_cosine=0.995) |
|
|
|
|
| def resolve_devices(devices, torch): |
| """The torch devices to run on, from `devices`: |
| |
| "auto" or None: every visible GPU, or the CPU when there is none; |
| an int N: the first N visible GPUs; |
| a list of GPU indices or device names, or a comma-separated string of them ([0, 2], "0,2", ["cuda:1"]); |
| "cpu": the CPU only. |
| """ |
| visible = torch.cuda.device_count() |
| if devices is None or devices == "auto" or devices == "all": |
| return [f"cuda:{i}" for i in range(visible)] or ["cpu"] |
| if isinstance(devices, bool): |
| raise ValueError("devices: 'auto', a number of GPUs, a list of GPUs, or 'cpu'") |
| if isinstance(devices, int): |
| if devices < 1: |
| raise ValueError(f"devices={devices}: give a number of GPUs from 1, a list of GPUs, 'auto' or 'cpu'") |
| if devices > visible: |
| raise ValueError(f"devices={devices}: this machine has {visible} visible GPU(s)") |
| return [f"cuda:{i}" for i in range(devices)] |
| if isinstance(devices, str): |
| devices = [d for d in devices.split(",") if d.strip()] |
| out = [] |
| for d in devices: |
| d = str(d).strip() |
| if d.isdigit(): |
| d = f"cuda:{d}" |
| if d.startswith("cuda") and int(d.partition(":")[2] or 0) >= visible: |
| raise ValueError(f"{d}: this machine has {visible} visible GPU(s)") |
| out.append(d) |
| if not out: |
| raise ValueError("devices: no device named") |
| return out |
|
|
|
|
| def _digest(a): |
| return hashlib.sha256(np.ascontiguousarray(a).tobytes()).hexdigest() |
|
|
|
|
| def compare(name, a, b): |
| """(report, within NEAR) for outputs `a` against reference outputs `b` of the self-test case `name`.""" |
| if name.startswith("segment"): |
| b = b[:a.shape[0]] |
| dp = float(np.abs(a - b).max()) |
| agree = float((a.argmax(-1) == b.argmax(-1)).mean()) |
| return dict(max_abs_dp=dp, frame_agreement=agree), dp <= NEAR["max_abs_dp"] and agree >= NEAR["frame_agreement"] |
| na, nb = np.linalg.norm(a, axis=1), np.linalg.norm(b, axis=1) |
| live = (na > 0) & (nb > 0) |
| cos = float((np.sum(a * b, axis=1)[live] / (na[live] * nb[live])).min()) if live.any() else 1.0 |
| return dict(min_cosine=cos), cos >= NEAR["min_cosine"] |
|
|
|
|
| class TorchNN: |
| def __init__(self, bundle, devices=None, workers_per_device=1, threads=2, golden=None, precision="exact", |
| strict=True): |
| if precision not in PRESETS: |
| raise ValueError(f"precision must be one of {sorted(PRESETS)}: 'exact' runs the release's FP32 networks, " |
| "the others faster reduced-precision variants") |
| self.precision = precision |
| self.strict = strict |
| self.exactness = {} |
| self._force_exact = False |
| if workers_per_device != 1: |
| raise ValueError("workers_per_device must be 1: two CUDA streams sharing one model on one GPU changed " |
| "about 1 batch in 1,000, so the engine runs one stream per GPU") |
| os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") |
| if os.environ["CUBLAS_WORKSPACE_CONFIG"] != ":4096:8": |
| raise RuntimeError("CUBLAS_WORKSPACE_CONFIG must be :4096:8 for byte-identical voiceprints") |
| import copy |
| import torch |
| import torch.nn.functional as F |
| from .seg import SegmentationNet |
| from .voiceprint import VoiceprintNet |
| self.torch, self.F = torch, F |
| self.threads = threads |
| torch.set_num_threads(threads) |
| self.flag_resets = 0 |
| self._enforce_flags() |
| bundle = Path(bundle) |
| self.devices = [torch.device(d) for d in resolve_devices(devices, torch)] |
| seg_cpu = SegmentationNet.from_pretrained(bundle / "segmentation") |
| vp_cpu = VoiceprintNet.from_pretrained(bundle / "voiceprint") |
| self.nets = [] |
| for dev in self.devices: |
| |
| seg, vp = copy.deepcopy(seg_cpu).to(dev).eval(), copy.deepcopy(vp_cpu).to(dev).eval() |
| for net in (seg, vp): |
| if any(bool(torch.isnan(t).any()) for t in list(net.parameters()) + list(net.buffers()) if t.is_floating_point()): |
| raise RuntimeError(f"NaN weights after moving the networks to {dev}") |
| self.nets.append((seg, vp)) |
| del seg_cpu, vp_cpu |
| self._fast = [self._compile_fast(seg, vp) for seg, vp in self.nets] if precision != "exact" else None |
| self.slots = [] |
| for di, dev in enumerate(self.devices): |
| for _ in range(workers_per_device if dev.type == "cuda" else 1): |
| stream = torch.cuda.Stream(device=dev) if dev.type == "cuda" else None |
| self.slots.append((di, dev, stream)) |
| self._golden = json.loads(Path(golden).read_text()) if golden else None |
| ref = Path(golden).with_name("nn_reference.npz") if golden else None |
| self._reference = dict(np.load(ref)) if ref is not None and ref.is_file() else None |
| self._lock = threading.Lock() |
| self.checked = [] |
| |
| for s in range(len(self.slots)): |
| self.self_test(s) |
|
|
| def _enforce_flags(self): |
| """The release's numeric flags (load_bundle). They are process-global, so any later import or user code may |
| change them; every network call re-asserts them (and counts how often it had to).""" |
| t = self.torch |
| exact = self.precision == "exact" or self._force_exact |
| types = PRESETS[self.precision] |
| tf32 = not exact and "tf32" in types |
| conv_tf32 = tf32 or (not exact and ("fp32c" in types or "ffn16" in types)) |
| want = (exact, tf32, conv_tf32, not exact, exact) |
| have = (t.are_deterministic_algorithms_enabled(), t.backends.cuda.matmul.allow_tf32, t.backends.cudnn.allow_tf32, |
| t.backends.cudnn.benchmark, t.backends.cudnn.deterministic) |
| if have == want: |
| return |
| self.flag_resets += 1 |
| t.use_deterministic_algorithms(exact) |
| t.backends.cuda.matmul.allow_tf32 = tf32 |
| t.backends.cudnn.allow_tf32 = conv_tf32 |
| t.backends.cudnn.benchmark = not exact |
| t.backends.cudnn.deterministic = exact |
|
|
| |
| def info(self): |
| t = self.torch |
| return dict(torch=t.__version__, cuda=t.version.cuda, cudnn=t.backends.cudnn.version(), |
| devices=[str(d) for d in self.devices], |
| gpus=[t.cuda.get_device_name(d) if d.type == "cuda" else "cpu" for d in self.devices], |
| slots=[(di, str(dev)) for di, dev, _ in self.slots], frames_sw=list(FRAMES_SW), |
| min_num_samples=MIN_NUM_SAMPLES, threads=self.threads, flag_resets=self.flag_resets, |
| precision=self.precision, exactness={str(k): v for k, v in self.exactness.items()}) |
|
|
| def slot_devices(self): |
| return [di for di, _, _ in self.slots] |
|
|
| @contextlib.contextmanager |
| def _on(self, slot): |
| self._enforce_flags() |
| di, dev, stream = self.slots[slot] |
| t = self.torch |
| with contextlib.ExitStack() as st: |
| if stream is not None: |
| st.enter_context(t.cuda.device(dev)) |
| st.enter_context(t.cuda.stream(stream)) |
| st.enter_context(t.inference_mode()) |
| yield self.nets[di], dev |
|
|
| |
| def segment(self, slot, x): |
| """`Inference.infer` with the softmax conversion of `load_bundle`. x: (B, 1, 256000) float32 numpy.""" |
| return self.segment_t(slot, lambda dev: self.torch.from_numpy(x).to(dev)) |
|
|
| def embed(self, slot, crops, weights): |
| """`PyannoteAudioPretrainedSpeakerEmbedding.__call__(waveforms, masks)`.""" |
| return self.embed_t(slot, lambda dev: self.torch.from_numpy(crops).to(dev), weights) |
|
|
| def embed_repacked(self, slot, windows, idx, weights): |
| """Voiceprints of H reference batches from one trunk pass over 8 distinct windows. |
| |
| windows (8, 1, 256000) float32 (zero rows pad a short block); idx (H, 8) int64, the trunk row of every |
| reference row; weights (H, 8, 799) float32. ResNet.forward split at the pooling layer.""" |
| return self.embed_repacked_t(slot, lambda dev: self.torch.from_numpy(windows).to(dev), idx, weights) |
|
|
| |
| |
| def windows_from_span(self, span, offsets, valid): |
| """`make(dev)` for windows given as one span: row j = span[offsets[j]:offsets[j] + valid[j]], then zeros.""" |
| t = self.torch |
|
|
| def make(dev): |
| s = t.from_numpy(span).to(dev) |
| x = t.zeros((len(offsets), 1, WIN), dtype=t.float32, device=dev) |
| for j, (o, v) in enumerate(zip(offsets.tolist(), valid.tolist())): |
| if v: |
| x[j, 0, :v] = s[o:o + v] |
| return x |
| return make |
|
|
| |
| FAST_BATCH = 8 |
|
|
| def _compile_fast(self, seg, vp): |
| """The two networks under torch.compile at one shape (batches padded to 8 windows), each in the compute type |
| of the preset (PRESETS): fp32 (TF32 off), tf32, or fp16 autocast. fbank, pooling, the last layer and the softmax |
| stay in FP32 (fbank power reaches ~1e14, above FP16's range).""" |
| t = self.torch |
| r = vp.resnet |
| seg_type, trunk_type = PRESETS[self.precision] |
| if seg_type == "ffn16": |
| |
| import copy |
| from .seg.wav2vec2.components import FeedForward |
| seg = copy.deepcopy(seg) |
| for m in seg.modules(): |
| if isinstance(m, FeedForward): |
| def fwd(x, f=m.forward): |
| with t.autocast("cuda", dtype=t.float16): |
| y = f(x) |
| return y.float() |
| m.forward = fwd |
|
|
| def seg_c(x): |
| with t.autocast("cuda", dtype=t.float16, enabled=seg_type == "fp16"): |
| out = seg(x) |
| return t.softmax(out.float(), -1) |
|
|
| def trunk_c(fb): |
| with t.autocast("cuda", dtype=t.float16, enabled=trunk_type == "fp16"): |
| h = fb.permute(0, 2, 1).unsqueeze(1) |
| h = self.F.relu(r.bn1(r.conv1(h))) |
| h = r.layer4(r.layer3(r.layer2(r.layer1(h)))) |
| return h.float() |
| return t.compile(seg_c, dynamic=False), t.compile(trunk_c, dynamic=False) |
|
|
| def _pad(self, x): |
| b = x.shape[0] |
| if b == self.FAST_BATCH: |
| return x, b |
| pad = self.torch.zeros((self.FAST_BATCH - b,) + tuple(x.shape[1:]), dtype=x.dtype, device=x.device) |
| return self.torch.cat([x, pad]), b |
|
|
| def _fast_on(self, slot): |
| return self._fast is not None and not self._force_exact |
|
|
| def _trunk(self, slot, emb, x): |
| """The voiceprint trunk's output for windows x (B, 1, 256000): exact (the release's ResNet.forward up to the |
| pooling layer) or the fast lane's compiled FP16 trunk.""" |
| t, F = self.torch, self.F |
| r = emb.resnet |
| fb = emb.compute_fbank(x) |
| if self._fast_on(slot): |
| out = [] |
| for i in range(0, fb.shape[0], self.FAST_BATCH): |
| part, b = self._pad(fb[i:i + self.FAST_BATCH]) |
| out.append(self._fast[self.slots[slot][0]][1](part)[:b]) |
| return t.cat(out) |
| x = fb.permute(0, 2, 1) |
| x = x.unsqueeze_(1) |
| out = F.relu(r.bn1(r.conv1(x))) |
| out = r.layer1(out) |
| out = r.layer2(out) |
| out = r.layer3(out) |
| return r.layer4(out) |
|
|
| def segment_t(self, slot, make): |
| with self._on(slot) as ((seg, _), dev): |
| x = make(dev) |
| if self._fast_on(slot): |
| out = [] |
| for i in range(0, x.shape[0], self.FAST_BATCH): |
| part, b = self._pad(x[i:i + self.FAST_BATCH]) |
| out.append(self._fast[self.slots[slot][0]][0](part)[:b]) |
| return self.torch.cat(out).cpu().numpy() |
| out = seg(x) |
| return self.torch.softmax(out, -1).cpu().numpy() |
|
|
| def embed_t(self, slot, make, weights): |
| t = self.torch |
| with self._on(slot) as ((_, emb), dev), warnings.catch_warnings(): |
| warnings.simplefilter("ignore") |
| if self._fast_on(slot): |
| r = emb.resnet |
| trunk = self._trunk(slot, emb, make(dev)) |
| return r.seg_1(r.pool(trunk, weights=t.from_numpy(weights).to(dev))).cpu().numpy() |
| return emb(make(dev), weights=t.from_numpy(weights).to(dev)).cpu().numpy() |
|
|
| def embed_repacked_t(self, slot, make, idx, weights): |
| t, F = self.torch, self.F |
| with self._on(slot) as ((_, emb), dev), warnings.catch_warnings(): |
| warnings.simplefilter("ignore") |
| r = emb.resnet |
| trunk = self._trunk(slot, emb, make(dev)) |
| res = [] |
| for h in range(idx.shape[0]): |
| rows = trunk[t.from_numpy(idx[h]).to(dev)] |
| res.append(r.seg_1(r.pool(rows, weights=t.from_numpy(weights[h]).to(dev))).cpu().numpy()) |
| return np.concatenate(res) if res else np.zeros((0, EMB), np.float32) |
|
|
| def resample(self, audio, sr): |
| """`load_audio`'s resampling, with the calling thread's CPU thread count set as the release sets it.""" |
| from .resample import resample |
| t = self.torch |
| t.set_num_threads(self.threads) |
| return resample(t.from_numpy(audio).unsqueeze(0), sr, 16000).squeeze(0).numpy().astype(np.float32) |
|
|
| |
| def self_test(self, slot): |
| """Run every shape on this slot and compare with the release's own outputs (assets/nn_golden.json): equal |
| bytes, or, where those cannot be reproduced and unless `strict`, within NEAR of assets/nn_reference.npz.""" |
| g = self._golden |
| res = {} |
| quick = self.slots[slot][1].type == "cpu" |
| self._force_exact = True |
| try: |
| exact = {name: fn() for name, fn in golden_cases(self, slot, quick=quick)} |
| finally: |
| self._force_exact = False |
| for name, out in exact.items(): |
| res[name] = _digest(out) |
| if g is not None: |
| bad = [k for k in g["digests"] if res.get(k) != g["digests"][k]] |
| if not bad: |
| self.exactness[slot] = dict(status="byte-identical") |
| elif self.strict or self._reference is None: |
| raise RuntimeError(f"slot {slot} ({self.slots[slot][1]}): outputs differ from the release for {bad}" |
| + ("" if self.strict else " (no assets/nn_reference.npz to check them against)")) |
| else: |
| near, ok = {}, True |
| for name, out in exact.items(): |
| near[name], good = compare(name, out, self._reference[name if not name.startswith("segment") else "segment_b8"]) |
| ok = ok and good |
| if not ok: |
| raise RuntimeError(f"slot {slot} ({self.slots[slot][1]}): outputs are outside the near-identical bounds " |
| f"{NEAR} of the release's: {near}") |
| worst = dict(max_abs_dp=max(v["max_abs_dp"] for k, v in near.items() if k.startswith("segment")), |
| frame_agreement=min(v["frame_agreement"] for k, v in near.items() if k.startswith("segment")), |
| min_cosine=min(v["min_cosine"] for k, v in near.items() if not k.startswith("segment"))) |
| self.exactness[slot] = dict(status="near-identical", **worst) |
| warnings.warn(f"anuvani_diar: on {self.slots[slot][1]} the networks run in FP32 and agree with the " |
| f"release's outputs within the near-identical bounds, but not to the last bit (that needs " |
| f"the setup the release's outputs were produced with); self-test: {worst}", stacklevel=3) |
| if self._fast is not None: |
| |
| |
| fast = {name: fn() for name, fn in golden_cases(self, slot)} |
| dev_report = {} |
| for name in exact: |
| a, b = exact[name], fast[name] |
| if name.startswith("segment"): |
| dp = float(np.abs(a - b).max()) |
| agree = float((a.argmax(-1) == b.argmax(-1)).mean()) |
| ok = dp <= 0.25 and agree >= 0.99 |
| dev_report[name] = dict(max_abs_dp=dp, frame_agreement=agree) |
| else: |
| na, nb = np.linalg.norm(a, axis=1), np.linalg.norm(b, axis=1) |
| live = (na > 0) & (nb > 0) |
| cos = float((np.sum(a * b, axis=1)[live] / (na[live] * nb[live])).min()) if live.any() else 1.0 |
| ok = cos >= 0.995 |
| dev_report[name] = dict(min_cosine=cos) |
| if not ok: |
| raise RuntimeError(f"slot {slot} ({self.slots[slot][1]}): the fast lane is outside its start-up bounds for " |
| f"{name}: {dev_report[name]}") |
| self.fast_report = getattr(self, "fast_report", {}) |
| self.fast_report[slot] = dev_report |
| with self._lock: |
| self.checked.append(slot) |
| return res |
|
|
|
|
| def golden_inputs(): |
| """Deterministic inputs of the self-test: seeded noise windows and seeded binary masks.""" |
| rng = np.random.default_rng(20261007) |
| win = (rng.standard_normal((8, 1, WIN)) * 0.1).astype(np.float32) |
| |
| win[3, 0, 100000:] = 0 |
| masks = (rng.random((8, 4, FRAMES)) < 0.4).astype(np.float32) |
| masks[1, 2] = 0 |
| return win, masks |
|
|
|
|
| def golden_cases(nn, slot, quick=False): |
| """(name, thunk) for every batch shape the engine sends to a slot (quick: one window and one voiceprint batch, for |
| a CPU, where the full set takes minutes).""" |
| win, masks = golden_inputs() |
| if quick: |
| return [("segment_b1", lambda: nn.segment(slot, win[:1].copy())), |
| ("embed_b4", lambda: nn.embed(slot, np.repeat(win[2:3], 4, 0), masks[2].copy()))] |
| cases = [(f"segment_b{b}", (lambda b=b: nn.segment(slot, win[:b].copy()))) for b in range(1, 9)] |
| pair = np.concatenate([np.repeat(win[0:1], 4, 0), np.repeat(win[1:2], 4, 0)]) |
| w8 = np.concatenate([masks[0], masks[1]]) |
| cases.append(("embed_b8", lambda: nn.embed(slot, pair, w8))) |
| cases.append(("embed_b4", lambda: nn.embed(slot, np.repeat(win[2:3], 4, 0), masks[2].copy()))) |
| idx = np.array([[0, 0, 0, 0, 1, 1, 1, 1], [2, 2, 2, 2, 3, 3, 3, 3]], np.int64) |
| wr = np.stack([w8, np.concatenate([masks[2], masks[3]])]) |
| cases.append(("embed_repacked_h2", lambda: nn.embed_repacked(slot, win.copy(), idx, wr))) |
| return cases |
|
|
