| """Portable low-latency runtime for the Nemotron-3-Diarization LiteRT bundle.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import math |
| import time |
| from pathlib import Path |
| from typing import Any |
|
|
| import librosa |
| import numpy as np |
|
|
|
|
| SHARDED_ADAPTER_ID = "nemotron3-diarization-low-latency-v1" |
| FUSED_ADAPTER_ID = "nemotron3-diarization-low-latency-v2" |
| SUPPORTED_ADAPTER_IDS = {SHARDED_ADAPTER_ID, FUSED_ADAPTER_ID} |
| SAMPLE_RATE = 16_000 |
| HOP_LENGTH = 160 |
| N_FFT = 512 |
| WIN_LENGTH = 400 |
| MEL_BINS = 128 |
| MEL_FRAMES = 104 |
| SUBSAMPLING_FACTOR = 8 |
| CHUNK_ENCODER_FRAMES = 9 |
| LOOKAHEAD_ENCODER_FRAMES = 4 |
| SPEAKER_CACHE_FRAMES = 264 |
| FIFO_FRAMES = 264 |
| SEQUENCE_FRAMES = 541 |
| HIDDEN_SIZE = 512 |
| NUM_SPEAKERS = 8 |
| NUM_LAYERS = 31 |
|
|
|
|
| class Graph: |
| """One resident LiteRT graph with a numerically ordered signature.""" |
|
|
| def __init__(self, path: Path, threads: int): |
| from ai_edge_litert.interpreter import Interpreter |
|
|
| self.path = path |
| self.interpreter = Interpreter(model_path=str(path), num_threads=threads) |
| self.interpreter.allocate_tensors() |
| signatures = self.interpreter.get_signature_list() |
| if set(signatures) != {"serving_default"}: |
| raise ValueError(f"Unexpected signatures in {path.name}: {list(signatures)}") |
| self.runner = self.interpreter.get_signature_runner("serving_default") |
|
|
| def __call__(self, *arguments: np.ndarray) -> np.ndarray: |
| """Invoke the graph and return its one finite output.""" |
|
|
| result = self.runner(**{f"args_{index}": value for index, value in enumerate(arguments)}) |
| ordered = [result[key] for key in sorted(result, key=lambda key: int(key.rsplit("_", 1)[1]))] |
| if len(ordered) != 1: |
| raise ValueError(f"Expected one output from {self.path.name}, got {len(ordered)}") |
| output = ordered[0] |
| if not np.isfinite(output).all(): |
| raise ValueError(f"Non-finite output from {self.path.name}") |
| return output |
|
|
|
|
| class SpeakerCache: |
| """NumPy implementation of the upstream AOSC and FIFO streaming policy.""" |
|
|
| def __init__(self, silence_embedding: np.ndarray): |
| self.silence_embedding = np.asarray(silence_embedding, dtype=np.float32).reshape(HIDDEN_SIZE) |
| self.embeddings = np.zeros((SPEAKER_CACHE_FRAMES, HIDDEN_SIZE), dtype=np.float32) |
| self.probabilities = np.zeros((SPEAKER_CACHE_FRAMES, NUM_SPEAKERS), dtype=np.float32) |
| self.fifo = np.zeros((FIFO_FRAMES, HIDDEN_SIZE), dtype=np.float32) |
| self.cache_count = 0 |
| self.fifo_count = 0 |
| self.compressed = False |
|
|
| def get_embeddings(self) -> np.ndarray: |
| """Return the currently populated AOSC followed by FIFO frames.""" |
|
|
| return np.concatenate( |
| [self.embeddings[: self.cache_count], self.fifo[: self.fifo_count]], axis=0 |
| )[None, ...] |
|
|
| @staticmethod |
| def _pool_probabilities(logits: np.ndarray) -> np.ndarray: |
| """Pool 10 ms logits to the 80 ms encoder frame rate.""" |
|
|
| probabilities = 1.0 / (1.0 + np.exp(-logits.astype(np.float32))) |
| frame_count = probabilities.shape[1] // SUBSAMPLING_FACTOR |
| pooled = probabilities[:, : frame_count * SUBSAMPLING_FACTOR] |
| pooled = pooled.reshape(1, frame_count, SUBSAMPLING_FACTOR, NUM_SPEAKERS).mean(axis=2) |
| return pooled.astype(np.float32) |
|
|
| @staticmethod |
| def _top_indices(values: np.ndarray, count: int) -> np.ndarray: |
| """Return unsorted indices for the largest values along the frame axis.""" |
|
|
| if count >= values.shape[1]: |
| return np.broadcast_to(np.arange(values.shape[1]), (values.shape[0], values.shape[1], values.shape[2])) |
| partition = np.argpartition(values, values.shape[1] - count, axis=1) |
| return partition[:, -count:, :] |
|
|
| @classmethod |
| def _boost_scores(cls, scores: np.ndarray, count: int, boost: float) -> np.ndarray: |
| """Boost each speaker's highest-scoring frames in place.""" |
|
|
| indices = cls._top_indices(scores, count) |
| batches = np.arange(scores.shape[0])[:, None, None] |
| speakers = np.arange(scores.shape[2])[None, None, :] |
| scores[batches, indices, speakers] += boost |
| return scores |
|
|
| @classmethod |
| def _compress( |
| cls, |
| embeddings: np.ndarray, |
| probabilities: np.ndarray, |
| silence_embedding: np.ndarray, |
| ) -> tuple[np.ndarray, np.ndarray]: |
| """Select and arrival-order the 264 AOSC frames used by the source runtime.""" |
|
|
| threshold = 0.25 |
| log_probabilities = np.log(np.maximum(probabilities, threshold)) |
| log_complements = np.log(np.maximum(1.0 - probabilities, threshold)) |
| scores = ( |
| log_probabilities |
| - log_complements |
| + log_complements.sum(axis=-1, keepdims=True) |
| - math.log(0.5) |
| ) |
| speech = probabilities > 0.5 |
| scores[~speech] = -np.inf |
| positive = scores > 0.0 |
| enough_positive = positive.sum(axis=1, keepdims=True) >= 16 |
| scores[(~positive) & speech & enough_positive] = -np.inf |
| scores[:, SPEAKER_CACHE_FRAMES:] += 0.05 |
| scores = cls._boost_scores(scores, 24, -2.0 * math.log(0.5)) |
| scores = cls._boost_scores(scores, 48, -math.log(0.5)) |
|
|
| frame_count = embeddings.shape[1] |
| scores = np.pad(scores, ((0, 0), (0, 1), (0, 0)), constant_values=np.inf) |
| silence = np.broadcast_to(silence_embedding.reshape(1, 1, -1), (1, 1, HIDDEN_SIZE)) |
| embeddings = np.concatenate([embeddings, silence], axis=1) |
| probabilities = np.pad(probabilities, ((0, 0), (0, 1), (0, 0))) |
|
|
| scored_frames = frame_count + 1 |
| flat_scores = scores.transpose(0, 2, 1).reshape(1, -1) |
| selected = np.argpartition(flat_scores, -SPEAKER_CACHE_FRAMES, axis=1)[ |
| :, -SPEAKER_CACHE_FRAMES: |
| ] |
| selected_scores = np.take_along_axis(flat_scores, selected, axis=1) |
| sentinel = scored_frames * NUM_SPEAKERS |
| selected[selected_scores == -np.inf] = sentinel |
| selected.sort(axis=1) |
| frame_indices = np.where(selected == sentinel, frame_count, selected % scored_frames) |
| return embeddings[:, frame_indices[0]], probabilities[:, frame_indices[0]] |
|
|
| def update( |
| self, |
| step_embeddings: np.ndarray, |
| step_logits: np.ndarray, |
| chunk_frame_count: int, |
| ) -> None: |
| """Push one processed chunk and update AOSC/FIFO state.""" |
|
|
| probabilities = self._pool_probabilities(step_logits) |
| chunk_start = self.cache_count + self.fifo_count |
| chunk = step_embeddings[:, chunk_start : chunk_start + chunk_frame_count] |
| fifo = np.concatenate([self.fifo[None, : self.fifo_count], chunk], axis=1) |
| popped = 0 |
| if fifo.shape[1] > FIFO_FRAMES: |
| popped = min(max(222, fifo.shape[1] - FIFO_FRAMES), fifo.shape[1]) |
|
|
| if popped: |
| fifo_probabilities = probabilities[ |
| :, self.cache_count : self.cache_count + fifo.shape[1] |
| ] |
| stored_probabilities = ( |
| self.probabilities[None, : self.cache_count] |
| if self.compressed |
| else probabilities[:, : self.cache_count] |
| ) |
| cache_embeddings = np.concatenate( |
| [self.embeddings[None, : self.cache_count], fifo[:, :popped]], axis=1 |
| ) |
| cache_probabilities = np.concatenate( |
| [stored_probabilities, fifo_probabilities[:, :popped]], axis=1 |
| ) |
| fifo = fifo[:, popped:] |
| if cache_embeddings.shape[1] > SPEAKER_CACHE_FRAMES: |
| cache_embeddings, cache_probabilities = self._compress( |
| cache_embeddings, cache_probabilities, self.silence_embedding |
| ) |
| self.compressed = True |
| self.cache_count = cache_embeddings.shape[1] |
| self.embeddings[: self.cache_count] = cache_embeddings[0] |
| self.probabilities[: self.cache_count] = cache_probabilities[0] |
|
|
| self.fifo_count = fifo.shape[1] |
| self.fifo[: self.fifo_count] = fifo[0] |
|
|
|
|
| def sha256_file(path: Path) -> str: |
| """Return the SHA-256 digest of a file.""" |
|
|
| digest = hashlib.sha256() |
| with path.open("rb") as handle: |
| for block in iter(lambda: handle.read(1024 * 1024), b""): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def verify_manifest(bundle: Path) -> dict[str, Any]: |
| """Validate the bundle identity, inventory, sizes, and hashes.""" |
|
|
| manifest = json.loads((bundle / "nemotron3-diarization-manifest.json").read_text()) |
| if manifest.get("adapter_id") not in SUPPORTED_ADAPTER_IDS: |
| raise ValueError(f"Wrong adapter: {manifest.get('adapter_id')!r}") |
| for entry in manifest["files"]: |
| path = bundle / entry["path"] |
| if path.parent != bundle or not path.is_file(): |
| raise ValueError(f"Invalid or missing bundle file: {entry['path']}") |
| if path.stat().st_size != entry["bytes"] or sha256_file(path) != entry["sha256"]: |
| raise ValueError(f"Bundle checksum mismatch: {entry['path']}") |
| return manifest |
|
|
|
|
| def log_mel_features(audio: np.ndarray, *, center: bool) -> np.ndarray: |
| """Compute the source checkpoint's unnormalized 128-bin log-mel features.""" |
|
|
| audio = np.asarray(audio, dtype=np.float32) |
| emphasized = np.concatenate([audio[:1], audio[1:] - 0.97 * audio[:-1]]) |
| if center: |
| emphasized = np.pad(emphasized, (N_FFT // 2, N_FFT // 2)) |
| if emphasized.size < N_FFT: |
| return np.zeros((0, MEL_BINS), dtype=np.float32) |
| frames = np.lib.stride_tricks.sliding_window_view(emphasized, N_FFT)[::HOP_LENGTH] |
| window = np.pad(np.hanning(WIN_LENGTH).astype(np.float32), ((N_FFT - WIN_LENGTH) // 2,) * 2) |
| spectrum = np.fft.rfft(frames * window, axis=1) |
| power = np.square(np.abs(spectrum).astype(np.float32)) |
| filters = librosa.filters.mel( |
| sr=SAMPLE_RATE, |
| n_fft=N_FFT, |
| n_mels=MEL_BINS, |
| fmin=0.0, |
| fmax=SAMPLE_RATE / 2, |
| norm="slaney", |
| ).astype(np.float32) |
| return np.log(power @ filters.T + 2**-24).astype(np.float32) |
|
|
|
|
| def rotary_embeddings() -> tuple[np.ndarray, np.ndarray]: |
| """Build the checkpoint's fixed RoPE cosine and sine inputs.""" |
|
|
| inverse_frequency = 1.0 / (10_000.0 ** (np.arange(0, 64, 2, dtype=np.float32) / 64.0)) |
| frequencies = np.outer(np.arange(SEQUENCE_FRAMES, dtype=np.float32), inverse_frequency) |
| embedding = np.concatenate([frequencies, frequencies], axis=-1)[None, ...] |
| return np.cos(embedding).astype(np.float32), np.sin(embedding).astype(np.float32) |
|
|
|
|
| def audio_chunks(audio: np.ndarray) -> list[tuple[np.ndarray, bool, bool]]: |
| """Split complete audio into the exact low-latency overlapping input windows.""" |
|
|
| first_samples = 16_680 |
| regular_samples = 17_040 |
| first_end = min(first_samples, audio.shape[0]) |
| if first_end == audio.shape[0]: |
| return [(audio, True, True)] |
| chunks = [(audio[:first_end], True, False)] |
| mel_frame = CHUNK_ENCODER_FRAMES * SUBSAMPLING_FACTOR |
| start = mel_frame * HOP_LENGTH - N_FFT // 2 |
| while start + regular_samples <= audio.shape[0]: |
| chunks.append((audio[start : start + regular_samples], False, False)) |
| mel_frame += CHUNK_ENCODER_FRAMES * SUBSAMPLING_FACTOR |
| start = mel_frame * HOP_LENGTH - N_FFT // 2 |
| if start < audio.shape[0]: |
| chunks.append((audio[start:], False, True)) |
| return chunks |
|
|
|
|
| def stabilize_activity(activity: np.ndarray, frames: int = 10) -> np.ndarray: |
| """Close sub-100 ms gaps and remove sub-100 ms speaker events.""" |
|
|
| stable = activity.copy() |
| for speaker in range(stable.shape[1]): |
| values = stable[:, speaker] |
| changes = np.diff(np.pad(values.astype(np.int8), (1, 1))) |
| speech_starts = np.flatnonzero(changes == 1) |
| speech_ends = np.flatnonzero(changes == -1) |
| for gap_start, gap_end in zip(speech_ends[:-1], speech_starts[1:], strict=True): |
| if gap_end - gap_start <= frames: |
| values[gap_start:gap_end] = True |
| changes = np.diff(np.pad(values.astype(np.int8), (1, 1))) |
| starts = np.flatnonzero(changes == 1) |
| ends = np.flatnonzero(changes == -1) |
| for start, end in zip(starts, ends, strict=True): |
| if end - start < frames: |
| values[start:end] = False |
| return stable |
|
|
|
|
| def segments_from_activity(activity: np.ndarray) -> list[dict[str, float | int]]: |
| """Convert 10 ms speaker activity to arrival-ordered segments.""" |
|
|
| segments: list[dict[str, float | int]] = [] |
| for speaker in range(activity.shape[1]): |
| changes = np.diff(np.pad(activity[:, speaker].astype(np.int8), (1, 1))) |
| starts = np.flatnonzero(changes == 1) |
| ends = np.flatnonzero(changes == -1) |
| segments.extend( |
| { |
| "start": round(float(start) * 0.01, 2), |
| "end": round(float(end) * 0.01, 2), |
| "speaker": speaker, |
| } |
| for start, end in zip(starts, ends, strict=True) |
| ) |
| return sorted(segments, key=lambda segment: (segment["start"], segment["speaker"])) |
|
|
|
|
| class Runtime: |
| """Complete fixed-shape low-latency LiteRT diarization runtime.""" |
|
|
| def __init__(self, bundle: Path, threads: int = 4, validate: bool = True): |
| self.bundle = Path(bundle) |
| if validate: |
| verify_manifest(self.bundle) |
| paths = {path.stem: path for path in self.bundle.glob("*.tflite")} |
| sharded = { |
| "feature_stacker", |
| "classification_head", |
| *(f"encoder_layer_{index:02d}" for index in range(NUM_LAYERS)), |
| } |
| fused = {"feature_stacker", "encoder_head"} |
| if set(paths) == fused: |
| self.fused = True |
| elif set(paths) == sharded: |
| self.fused = False |
| else: |
| raise ValueError( |
| "Incomplete or unexpected graph inventory: " |
| f"{sorted(set(paths) ^ (fused if 'encoder_head' in paths else sharded))}" |
| ) |
| self.graphs = {name: Graph(path, threads) for name, path in paths.items()} |
| self.silence_embedding = np.load(self.bundle / "silence_embedding.npy") |
| self.rope_cos, self.rope_sin = rotary_embeddings() |
|
|
| def diarize(self, audio: np.ndarray, stabilize_ms: int = 100) -> dict[str, Any]: |
| """Diarize finite mono Float32 PCM sampled at exactly 16 kHz.""" |
|
|
| audio = np.asarray(audio, dtype=np.float32) |
| if audio.ndim != 1 or not np.isfinite(audio).all(): |
| raise ValueError("Expected finite mono Float32 PCM") |
| cache = SpeakerCache(self.silence_embedding) |
| output_chunks: list[np.ndarray] = [] |
| graph_seconds = 0.0 |
| chunks = audio_chunks(audio) |
| for chunk, first, last in chunks: |
| features = log_mel_features(chunk, center=first) |
| valid_mel_frames = chunk.shape[0] // HOP_LENGTH if first else max( |
| 0, (chunk.shape[0] - N_FFT) // HOP_LENGTH + 1 |
| ) |
| features = features[:valid_mel_frames] |
| padded = np.zeros((1, MEL_FRAMES, MEL_BINS), dtype=np.float32) |
| padded[:, : features.shape[0]] = features |
| started = time.perf_counter() |
| current = self.graphs["feature_stacker"](padded) |
| current_count = math.ceil(valid_mel_frames / SUBSAMPLING_FACTOR) |
| current = current[:, :current_count] |
| cached = cache.get_embeddings() |
| step_embeddings = np.concatenate([cached, current], axis=1) |
| pad_frames = SEQUENCE_FRAMES - step_embeddings.shape[1] |
| if pad_frames < 0: |
| raise ValueError(f"Streaming state exceeded {SEQUENCE_FRAMES} frames") |
| hidden = np.zeros((1, SEQUENCE_FRAMES, HIDDEN_SIZE), dtype=np.float32) |
| hidden[:, pad_frames:] = step_embeddings |
| valid = np.zeros((1, SEQUENCE_FRAMES), dtype=np.bool_) |
| valid[:, pad_frames:] = True |
| attention = np.where( |
| valid[:, None, None, :], 0.0, np.finfo(np.float32).min |
| ).astype(np.float32) |
| if self.fused: |
| full_logits = self.graphs["encoder_head"]( |
| hidden, attention, self.rope_cos, self.rope_sin, valid |
| ) |
| else: |
| for index in range(NUM_LAYERS): |
| hidden = self.graphs[f"encoder_layer_{index:02d}"]( |
| hidden, attention, self.rope_cos, self.rope_sin |
| ) |
| full_logits = self.graphs["classification_head"](hidden, valid) |
| graph_seconds += time.perf_counter() - started |
|
|
| lookahead = 0 if last else LOOKAHEAD_ENCODER_FRAMES |
| chunk_encoder_frames = current_count - lookahead |
| start_encoder = pad_frames + cached.shape[1] |
| start_logit = start_encoder * SUBSAMPLING_FACTOR |
| output_frames = min(chunk_encoder_frames * SUBSAMPLING_FACTOR, valid_mel_frames) |
| output_chunks.append(full_logits[:, start_logit : start_logit + output_frames]) |
| valid_start = pad_frames * SUBSAMPLING_FACTOR |
| valid_end = (pad_frames + step_embeddings.shape[1]) * SUBSAMPLING_FACTOR |
| cache.update( |
| step_embeddings, |
| full_logits[:, valid_start:valid_end], |
| chunk_encoder_frames, |
| ) |
|
|
| logits = np.concatenate(output_chunks, axis=1)[0] if output_chunks else np.zeros((0, 8)) |
| activity = logits > 0.0 |
| stable = stabilize_activity(activity, max(1, round(stabilize_ms / 10))) |
| return { |
| "segments": segments_from_activity(stable), |
| "raw_segments": segments_from_activity(activity), |
| "audio_seconds": audio.shape[0] / SAMPLE_RATE, |
| "graph_seconds": graph_seconds, |
| "graph_rtf": graph_seconds / (audio.shape[0] / SAMPLE_RATE) if audio.size else 0.0, |
| "chunks": len(chunks), |
| "stabilize_ms": stabilize_ms, |
| } |
|
|
|
|
| def main() -> int: |
| """Run the portable reference from the command line.""" |
|
|
| import soundfile as sf |
| from scipy.signal import resample_poly |
|
|
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("bundle", type=Path) |
| parser.add_argument("audio", type=Path) |
| parser.add_argument("--threads", type=int, default=4) |
| parser.add_argument("--stabilize-ms", type=int, default=100) |
| args = parser.parse_args() |
| waveform, sample_rate = sf.read(args.audio, dtype="float32", always_2d=True) |
| mono = waveform.mean(axis=1) |
| divisor = math.gcd(sample_rate, SAMPLE_RATE) |
| audio = resample_poly(mono, SAMPLE_RATE // divisor, sample_rate // divisor).astype(np.float32) |
| result = Runtime(args.bundle, args.threads).diarize(audio, args.stabilize_ms) |
| print(json.dumps(result, indent=2, sort_keys=True)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
