Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+bin/sortformer_ax650 filter=lfs diff=lfs merge=lfs -text
+models/encoder.axmodel filter=lfs diff=lfs merge=lfs -text
+models/encoder_fifo40.axmodel filter=lfs diff=lfs merge=lfs -text
+models/preencode.axmodel filter=lfs diff=lfs merge=lfs -text
+samples/sample_meeting.wav filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.pyc
+*.rttm
+outputs/
+.msc/
+.mv/
+._____temp/

README.md

@@ -1,3 +1,96 @@
 ---
-license: mit
+license: other
+license_name: nvidia-open-model-license
+license_link: https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/
+language:
+- en
+pipeline_tag: audio-classification
+tags:
+- speaker-diarization
+- sortformer
+- axera
+- ax650
+- rttm
 ---
+
+# Sortformer.AXERA
+
+`nvidia/diar_streaming_sortformer_4spk-v2.1`（NVIDIA Open Model License）端到端流式说话人日志在 AX650N 的量化部署包：
+输入 16 kHz 单声道会议音频，输出 RTTM 说话人时序标签。
+
+- 模型: 2 张拆图 axmodel（U16 激活 / S8 权重）——preencode + encoder；encoder 提供 fifo188（精度优先）与 fifo40（速度优先）
+- 依赖: Python 仅 numpy + soundfile + axengine；C++ 为预编译 aarch64 可执行文件
+
+## 目录
+
+```
+├── models/          # preencode.axmodel / encoder.axmodel / encoder_fifo40.axmodel + model_meta
+├── python/          # Python 推理入口（sortformer_sdk + example.py）
+├── bin/             # C++ 可执行文件（aarch64）：sortformer_ax650
+├── samples/         # 演示音频（AMI 会议 120 s）
+├── run_ax650.sh  run_cpp_ax650.sh   # 一键运行（Python / C++）
+└── requirements.txt
+```
+
+## Python 运行
+
+```bash
+pip3 install numpy soundfile
+# pyaxengine: https://github.com/AXERA-TECH/pyaxengine/releases/latest
+pip3 install axengine-<version>-py3-none-any.whl
+./run_ax650.sh                            # 自带样例 samples/sample_meeting.wav
+./run_ax650.sh input.wav out.rttm --fast  # fifo40 快速版
+```
+
+或:
+
+```python
+import sys; sys.path.insert(0, "python")
+from sortformer_sdk import AxengineGraphPair, StreamingDiarizer, SortformerConfig
+
+diarizer = StreamingDiarizer(
+    AxengineGraphPair("models/preencode.axmodel", "models/encoder.axmodel"),
+    SortformerConfig(),
+)
+preds = diarizer.process_wav(waveform_16k_float32, 16000)   # (T, 4) @ 80 ms
+lines = diarizer.rttm_lines(preds, uri="meeting")
+```
+
+## C++ 运行
+
+```bash
+export LD_LIBRARY_PATH=/soc/lib:${LD_LIBRARY_PATH:-}
+./bin/sortformer_ax650 --preencode models/preencode.axmodel --encoder models/encoder.axmodel \
+    --wav input.wav --rttm out.rttm --threads 8
+# 或一键：./run_cpp_ax650.sh input.wav out.rttm [--fast]
+```
+
+## 模型说明
+
+- 输入: 16 kHz 单声道 wav；输出: RTTM（`SPEAKER <uri> 1 <start> <dur> <NA> <NA> speaker_k <NA>`）
+- 流程: 主机 log-mel（128 mel / 25 ms 窗 / 10 ms 步长，与 NeMo 逐点对齐）→ preencode → 主机打包
+  `seq [1,390,512]` → encoder → 主机 `streaming_update`（静音画像 / top-k + AOSC 压缩 / FIFO 弹出）→ 后处理
+- 拆图契约：spkcache 188 + fifo 188 + chunk 14 帧，图内无 Scatter/Where；80 ms/帧、最多 4 说话人、1.04 s 延迟
+- 完整转换源码见 GitHub: [Sortformer.AXERA](https://github.com/ZY-2012/Sortformer.AXERA)
+
+## 精度与性能（AX650N 实测）
+
+ES2004a（1049 s，collar=0）：
+
+| 推理路径 | DER | RTF |
+|---|---|---|
+| NeMo FP32（参考） | 30.89% | — |
+| C++ fifo188（精度优先） | 30.70% | 0.144 |
+| C++ fifo40（速度优先） | 30.64% | 0.064 |
+
+AMI-SDM（5 场加权）/ AliMeeting（4 场远场，TextGrid 参考），collar=0：
+
+| 配置 | AMI-5 DER | Ali-4 DER |
+|---|---|---|
+| fifo188 | 33.77% | 21.71% |
+| fifo40 | 34.34% | 22.02% |
+
+## License
+
+部署代码 Apache-2.0；上游 Sortformer 权重为
+[NVIDIA Open Model License](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/)。

bin/sortformer_ax650

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b54e33a4ff0ef1d28a4d96f3a0b010778715c4f19a16d40651abf08bf538ca65
+size 144952

configuration.json

@@ -0,0 +1,17 @@
+{
+  "model_name": "Sortformer-streaming-4spk-v2.1-ax650",
+  "task": "speaker-diarization",
+  "base_model": "nvidia/diar_streaming_sortformer_4spk-v2.1",
+  "chip": "AX650N",
+  "npu_mode": "NPU3",
+  "sample_rate": 16000,
+  "max_speakers": 4,
+  "frame_shift_ms": 80,
+  "latency_s": 1.04,
+  "pulsar2_version": "7.0",
+  "axmodel_files": [
+    "models/preencode.axmodel",
+    "models/encoder.axmodel",
+    "models/encoder_fifo40.axmodel"
+  ]
+}

models/encoder.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5ce5abb7c3509c85ef76eda43e4858658e0c498de370a5baa2e423dcdef2ef96
+size 139147096

models/encoder_fifo40.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50124d50843561cb2a0653ed911afb2fae21fcc84f2e38c6fd408744ec621fa8
+size 132804934

models/model_meta.json

@@ -0,0 +1,57 @@
+{
+  "model": "nvidia/diar_streaming_sortformer_4spk-v2.1",
+  "contract": "split-preencode-encoder",
+  "simplified": true,
+  "config": {
+    "chunk_len": 6,
+    "chunk_left_context": 1,
+    "chunk_right_context": 7,
+    "fifo_len": 188,
+    "spkcache_len": 188,
+    "spkcache_update_period": 144
+  },
+  "preencode": {
+    "file": "preencode.onnx",
+    "inputs": {
+      "chunk": [
+        1,
+        112,
+        128
+      ]
+    },
+    "outputs": {
+      "chunk_pre_encode_embs": [
+        1,
+        14,
+        512
+      ]
+    },
+    "parity_min_cosine": 0.9999999403953552
+  },
+  "encoder": {
+    "file": "encoder.onnx",
+    "inputs": {
+      "seq": [
+        1,
+        390,
+        512
+      ],
+      "total_lengths": [
+        1
+      ]
+    },
+    "outputs": {
+      "spkcache_fifo_chunk_preds": [
+        1,
+        390,
+        4
+      ]
+    },
+    "parity_min_cosine": 0.9999998807907104
+  },
+  "state_len": 376,
+  "chunk_mel_frames": 112,
+  "frame_duration_s": 0.08,
+  "sample_rate": 16000,
+  "mel_features": 128
+}

models/model_meta_fifo40.json

@@ -0,0 +1,57 @@
+{
+  "model": "nvidia/diar_streaming_sortformer_4spk-v2.1",
+  "contract": "split-preencode-encoder",
+  "simplified": true,
+  "config": {
+    "chunk_len": 6,
+    "chunk_left_context": 1,
+    "chunk_right_context": 7,
+    "fifo_len": 40,
+    "spkcache_len": 188,
+    "spkcache_update_period": 31
+  },
+  "preencode": {
+    "file": "preencode.onnx",
+    "inputs": {
+      "chunk": [
+        1,
+        112,
+        128
+      ]
+    },
+    "outputs": {
+      "chunk_pre_encode_embs": [
+        1,
+        14,
+        512
+      ]
+    },
+    "parity_min_cosine": 1.0
+  },
+  "encoder": {
+    "file": "encoder.onnx",
+    "inputs": {
+      "seq": [
+        1,
+        242,
+        512
+      ],
+      "total_lengths": [
+        1
+      ]
+    },
+    "outputs": {
+      "spkcache_fifo_chunk_preds": [
+        1,
+        242,
+        4
+      ]
+    },
+    "parity_min_cosine": 0.9999999403953552
+  },
+  "state_len": 228,
+  "chunk_mel_frames": 112,
+  "frame_duration_s": 0.08,
+  "sample_rate": 16000,
+  "mel_features": 128
+}

models/preencode.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:399ec91c5845102985614bc5f19556b43816f57d16b9030ab86be67b908a6701
+size 2535538

python/example.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python3
+"""Sortformer streaming diarization on AX650N: 16 kHz wav -> RTTM.
+
+    python3 example.py meeting.wav out.rttm [--fast]
+"""
+
+import argparse
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+from sortformer_sdk import AxengineGraphPair, SortformerConfig, StreamingDiarizer  # noqa: E402
+
+MODELS = Path(__file__).resolve().parents[1] / "models"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("wav")
+    parser.add_argument("rttm", nargs="?", default=None)
+    parser.add_argument("--fast", action="store_true", help="fifo40 encoder: RTF 0.064, DER +0.3~0.6pp")
+    parser.add_argument("--preencode", default=str(MODELS / "preencode.axmodel"))
+    parser.add_argument("--encoder", default=None)
+    parser.add_argument("--threads", type=int, default=8, help="reserved (numpy front-end is single-process)")
+    args = parser.parse_args()
+
+    encoder = args.encoder or str(MODELS / ("encoder_fifo40.axmodel" if args.fast else "encoder.axmodel"))
+    config = SortformerConfig(
+        fifo_len=40 if args.fast else 188,
+        spkcache_update_period=31 if args.fast else 144,
+    )
+
+    waveform, sample_rate = sf.read(args.wav, dtype="float32", always_2d=True)
+    waveform = waveform.mean(axis=1)
+    if sample_rate != 16000:
+        raise SystemExit(f"expected 16 kHz wav, got {sample_rate}")
+
+    diarizer = StreamingDiarizer(AxengineGraphPair(args.preencode, encoder), config)
+    start = time.perf_counter()
+    preds = diarizer.process_wav(waveform, sample_rate)
+    elapsed = time.perf_counter() - start
+    duration = waveform.shape[0] / sample_rate
+
+    uri = Path(args.wav).stem
+    lines = diarizer.rttm_lines(preds, uri)
+    if args.rttm:
+        Path(args.rttm).write_text("\n".join(lines) + "\n")
+
+    print(
+        f"{uri}: {len(preds)} frames, {len(lines)} segments, "
+        f"{elapsed:.1f}s / {duration:.1f}s audio (RTF {elapsed / duration:.4f})"
+    )
+    if args.rttm:
+        print(f"rttm -> {args.rttm}")
+
+
+if __name__ == "__main__":
+    main()

python/sortformer_sdk/__init__.py

@@ -0,0 +1,34 @@
+from .diarize import AxengineGraphPair, OnnxGraphPair, StreamingDiarizer
+from .feature import log_mel_spectrogram, mel_filterbank
+from .postprocess import (
+    PostProcessingParams,
+    generate_diarization_output_lines,
+    predlist_to_timestamps,
+    timestamps_to_rttm_lines,
+)
+from .state import (
+    SortformerConfig,
+    StreamingState,
+    init_state,
+    iter_chunks,
+    pre_encode_length,
+    streaming_update,
+)
+
+__all__ = [
+    "AxengineGraphPair",
+    "OnnxGraphPair",
+    "PostProcessingParams",
+    "SortformerConfig",
+    "StreamingDiarizer",
+    "StreamingState",
+    "generate_diarization_output_lines",
+    "init_state",
+    "iter_chunks",
+    "log_mel_spectrogram",
+    "mel_filterbank",
+    "pre_encode_length",
+    "predlist_to_timestamps",
+    "streaming_update",
+    "timestamps_to_rttm_lines",
+]

python/sortformer_sdk/diarize.py

@@ -0,0 +1,179 @@
+"""End-to-end streaming diarizer: waveform -> RTTM.
+
+Two scatter-free graphs are executed per step (pre-encode + encoder); the
+speaker-cache packing, streaming state update, log-mel front-end and
+post-processing are all numpy and shared by host and AX650.
+"""
+
+import argparse
+from pathlib import Path
+from typing import Dict, Protocol
+
+import numpy as np
+
+from .feature import SAMPLE_RATE, log_mel_spectrogram, mel_filterbank
+from .postprocess import PostProcessingParams, predlist_to_timestamps, timestamps_to_rttm_lines
+from .state import SortformerConfig, init_state, iter_chunks, pre_encode_length, streaming_update
+
+
+class GraphPair(Protocol):
+    def pre_encode(self, chunk: np.ndarray) -> np.ndarray:
+        """chunk [1, chunk_mel, 128] -> chunk embeddings [1, chunk_embs, 512]."""
+
+    def encode(self, seq: np.ndarray, total_lengths: int) -> np.ndarray:
+        """packed seq [1, state_len + chunk_embs, 512] -> preds [1, T, num_speakers]."""
+
+
+def _provider_alias(ort, provider: str) -> str:
+    aliases = {"cpu": "CPUExecutionProvider", "cuda": "CUDAExecutionProvider"}
+    if provider == "auto":
+        available = ort.get_available_providers()
+        return "CUDAExecutionProvider" if "CUDAExecutionProvider" in available else "CPUExecutionProvider"
+    return aliases.get(provider, provider)
+
+
+class OnnxGraphPair:
+    """ONNX Runtime implementation of the split graphs (host verification)."""
+
+    def __init__(self, preencode_path: str, encoder_path: str, provider: str = "auto"):
+        import onnxruntime as ort
+
+        provider = _provider_alias(ort, provider)
+        self.pre_session = ort.InferenceSession(str(preencode_path), providers=[provider])
+        self.encoder_session = ort.InferenceSession(str(encoder_path), providers=[provider])
+        # The encoder graph may be wider than the runtime state (e.g. a 390-wide
+        # graph evaluated with a 40-frame FIFO); the host pads and masks the tail.
+        encoder_inputs = {node.name: node for node in self.encoder_session.get_inputs()}
+        self.seq_width = int(encoder_inputs["seq"].shape[1])
+
+    def pre_encode(self, chunk: np.ndarray) -> np.ndarray:
+        return self.pre_session.run(None, {"chunk": chunk.astype(np.float32)})[0]
+
+    def encode(self, seq: np.ndarray, total_lengths: int) -> np.ndarray:
+        feed = {"seq": seq.astype(np.float32), "total_lengths": np.array([total_lengths], dtype=np.float32)}
+        return self.encoder_session.run(None, feed)[0]
+
+
+class AxengineGraphPair:
+    """axengine implementation of the split graphs (AX650 board)."""
+
+    def __init__(self, preencode_path: str, encoder_path: str):
+        import axengine
+
+        self.pre_session = axengine.InferenceSession(str(preencode_path))
+        self.encoder_session = axengine.InferenceSession(str(encoder_path))
+
+    def pre_encode(self, chunk: np.ndarray) -> np.ndarray:
+        return self.pre_session.run(None, {"chunk": chunk.astype(np.float32)})[0]
+
+    def encode(self, seq: np.ndarray, total_lengths: int) -> np.ndarray:
+        feed = {"seq": seq.astype(np.float32), "total_lengths": np.array([total_lengths], dtype=np.float32)}
+        return self.encoder_session.run(None, feed)[0]
+
+
+class StreamingDiarizer:
+    """Sortformer 1.04 s streaming diarizer over the split graphs."""
+
+    def __init__(self, graphs: GraphPair, config: SortformerConfig = None):
+        self.graphs = graphs
+        self.config = config or SortformerConfig()
+        self.chunk_width = (
+            self.config.chunk_left_context + self.config.chunk_len + self.config.chunk_right_context
+        ) * self.config.subsampling_factor
+        self.state_len = self.config.spkcache_len + self.config.fifo_len
+
+    def process_features(self, features: np.ndarray, on_step=None) -> np.ndarray:
+        cfg = self.config
+        state = init_state(cfg)
+        mel_dim = features.shape[1]
+        predictions = []
+        for step, (chunk, left_offset, right_offset) in enumerate(iter_chunks(cfg, features)):
+            chunk_valid = chunk.shape[0]
+            chunk_input = np.zeros((1, self.chunk_width, mel_dim), dtype=np.float32)
+            chunk_input[0, :chunk_valid] = chunk
+
+            chunk_embs_full = self.graphs.pre_encode(chunk_input)[0]
+            chunk_embs_length = pre_encode_length(chunk_valid)
+            chunk_embs = chunk_embs_full[:chunk_embs_length]
+
+            spkcache_len = state.spkcache.shape[0]
+            fifo_len = state.fifo.shape[0]
+            total_lengths = spkcache_len + fifo_len + chunk_embs_length
+            seq_width = getattr(self.graphs, "seq_width", self.state_len + chunk_embs_full.shape[0])
+            seq = np.zeros((1, seq_width, cfg.fc_d_model), dtype=np.float32)
+            seq[0, :spkcache_len] = state.spkcache
+            seq[0, spkcache_len : spkcache_len + fifo_len] = state.fifo
+            seq[0, spkcache_len + fifo_len : total_lengths] = chunk_embs
+
+            if on_step is not None:
+                on_step(step, chunk_input, seq, total_lengths)
+
+            preds = self.graphs.encode(seq, total_lengths)[0]
+            lc_enc = round(left_offset / cfg.subsampling_factor)
+            rc_enc = int(np.ceil(right_offset / cfg.subsampling_factor))
+            state, chunk_preds = streaming_update(cfg, state, chunk_embs, preds, lc_enc, rc_enc)
+            predictions.append(chunk_preds)
+        if not predictions:
+            return np.zeros((0, cfg.num_speakers), dtype=np.float32)
+        return np.concatenate(predictions, axis=0)
+
+    def process_wav(self, waveform: np.ndarray, sample_rate: int = SAMPLE_RATE, on_step=None) -> np.ndarray:
+        features = log_mel_spectrogram(waveform, sample_rate, mel_filter=mel_filterbank())
+        return self.process_features(features, on_step=on_step)
+
+    def rttm_lines(self, preds: np.ndarray, uri: str, params: PostProcessingParams = None, bypass: bool = False):
+        timestamps = predlist_to_timestamps(preds, params=params, bypass_postprocessing=bypass)
+        num_speakers = preds.shape[1] if preds.ndim == 2 else self.config.num_speakers
+        return timestamps_to_rttm_lines(timestamps, uri, num_speakers)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--preencode", default=None, help="preencode.onnx (host mode)")
+    parser.add_argument("--encoder", default=None, help="encoder.onnx (host mode)")
+    parser.add_argument("--preencode-axmodel", default=None, help="preencode axmodel (board mode)")
+    parser.add_argument("--encoder-axmodel", default=None, help="encoder axmodel (board mode)")
+    parser.add_argument("--wav", required=True)
+    parser.add_argument("--rttm", default=None)
+    parser.add_argument("--provider", default="auto")
+    parser.add_argument("--chunk-len", type=int, default=6)
+    parser.add_argument("--chunk-left-context", type=int, default=1)
+    parser.add_argument("--chunk-right-context", type=int, default=7)
+    parser.add_argument("--fifo-len", type=int, default=188)
+    parser.add_argument("--spkcache-len", type=int, default=188)
+    parser.add_argument("--spkcache-update-period", type=int, default=144)
+    parser.add_argument("--bypass-postproc", action="store_true")
+    args = parser.parse_args()
+
+    if args.preencode and args.encoder:
+        graphs = OnnxGraphPair(args.preencode, args.encoder, provider=args.provider)
+    elif args.preencode_axmodel and args.encoder_axmodel:
+        graphs = AxengineGraphPair(args.preencode_axmodel, args.encoder_axmodel)
+    else:
+        raise SystemExit("pass either --preencode/--encoder or the axmodel pair")
+
+    config = SortformerConfig(
+        chunk_len=args.chunk_len,
+        chunk_left_context=args.chunk_left_context,
+        chunk_right_context=args.chunk_right_context,
+        fifo_len=args.fifo_len,
+        spkcache_len=args.spkcache_len,
+        spkcache_update_period=args.spkcache_update_period,
+    )
+    diarizer = StreamingDiarizer(graphs, config)
+
+    import soundfile as sf
+
+    waveform, sample_rate = sf.read(args.wav, dtype="float32", always_2d=True)
+    waveform = waveform.mean(axis=1)
+    preds = diarizer.process_wav(waveform, sample_rate)
+    uri = Path(args.wav).stem
+    lines = diarizer.rttm_lines(preds, uri, bypass=args.bypass_postproc)
+    print(f"{uri}: {len(preds)} frames, {len(lines)} segments")
+    if args.rttm:
+        Path(args.rttm).write_text("\n".join(lines) + "\n")
+        print(f"rttm -> {args.rttm}")
+
+
+if __name__ == "__main__":
+    main()

python/sortformer_sdk/feature.py

@@ -0,0 +1,128 @@
+"""NeMo-compatible log-mel front-end (numpy only, no torch/librosa at runtime).
+
+Mirrors ``FilterbankFeatures`` (eval mode) from NeMo Speech 3.0 with the
+Sortformer streaming configuration:
+
+    sample_rate=16000, n_fft=512, win_length=400 (hann, periodic=False),
+    hop_length=160, n_mels=128, preemph=0.97, mag_power=2.0,
+    log(x + 2**-24), normalize="NA", pad_to=16
+
+``log_mel_spectrogram`` returns ``(num_frames, 128)`` float32 where
+``num_frames = floor(num_samples / hop)`` rounded up to a multiple of 16.
+"""
+
+import numpy as np
+
+SAMPLE_RATE = 16000
+N_FFT = 512
+WIN_LENGTH = 400
+HOP_LENGTH = 160
+N_MELS = 128
+PREEMPH = 0.97
+LOG_ZERO_GUARD = 2.0 ** -24
+PAD_TO = 16
+
+_MEL_CACHE = {}
+
+
+def _hz_to_mel(freq):
+    # Slaney mel scale (librosa default, htk=False)
+    f_sp = 200.0 / 3.0
+    log_step = np.log(6.4) / 27.0
+    min_log_hz = 1000.0
+    min_log_mel = min_log_hz / f_sp
+    freq = np.asarray(freq, dtype=np.float64)
+    mel_linear = freq / f_sp
+    mel_log = min_log_mel + np.log(np.maximum(freq, min_log_hz) / min_log_hz) / log_step
+    return np.where(freq >= min_log_hz, mel_log, mel_linear)
+
+
+def _mel_to_hz(mel):
+    f_sp = 200.0 / 3.0
+    log_step = np.log(6.4) / 27.0
+    min_log_hz = 1000.0
+    min_log_mel = min_log_hz / f_sp
+    mel = np.asarray(mel, dtype=np.float64)
+    hz_linear = f_sp * mel
+    hz_log = min_log_hz * np.exp(log_step * (mel - min_log_mel))
+    return np.where(mel >= min_log_mel, hz_log, hz_linear)
+
+
+def mel_filterbank(sample_rate: int = SAMPLE_RATE, n_fft: int = N_FFT, n_mels: int = N_MELS):
+    """librosa-compatible Slaney-normalized Slaney-scale mel filterbank."""
+    key = (sample_rate, n_fft, n_mels)
+    if key in _MEL_CACHE:
+        return _MEL_CACHE[key]
+
+    fmin, fmax = 0.0, sample_rate / 2.0
+    mels = np.linspace(_hz_to_mel(fmin), _hz_to_mel(fmax), n_mels + 2)
+    hz = _mel_to_hz(mels)
+    freqs = np.linspace(0.0, sample_rate / 2.0, 1 + n_fft // 2)
+
+    fdiff = np.diff(hz)
+    ramps = np.subtract.outer(hz, freqs)
+    lower = -ramps[np.arange(n_mels), :] / fdiff[np.arange(n_mels)][:, None]
+    upper = ramps[np.arange(2, n_mels + 2), :] / fdiff[np.arange(1, n_mels + 1)][:, None]
+    weights = np.maximum(0.0, np.minimum(lower, upper))
+
+    enorm = 2.0 / (hz[2 : n_mels + 2] - hz[:n_mels])
+    weights *= enorm[:, None]
+    result = weights.astype(np.float64)
+    _MEL_CACHE[key] = result
+    return result
+
+
+def _hann_window(length: int):
+    return 0.5 - 0.5 * np.cos(2.0 * np.pi * np.arange(length) / (length - 1))
+
+
+def log_mel_spectrogram(
+    waveform: np.ndarray,
+    sample_rate: int = SAMPLE_RATE,
+    *,
+    preemph: float = PREEMPH,
+    pad_to: int = PAD_TO,
+    mel_filter: np.ndarray = None,
+) -> np.ndarray:
+    """Compute the NeMo Sortformer streaming log-mel features for a mono waveform."""
+    if sample_rate != SAMPLE_RATE:
+        raise ValueError(f"expected {SAMPLE_RATE} Hz input, got {sample_rate}")
+    waveform = np.asarray(waveform, dtype=np.float32).reshape(-1)
+
+    if preemph is not None and waveform.size > 0:
+        emphasized = np.empty_like(waveform)
+        emphasized[0] = waveform[0]
+        emphasized[1:] = waveform[1:] - preemph * waveform[:-1]
+        waveform = emphasized
+
+    padded = np.pad(waveform, (N_FFT // 2, N_FFT // 2), mode="constant")
+    num_samples = waveform.shape[0]
+    num_frames = num_samples // HOP_LENGTH
+    if num_frames == 0:
+        return np.zeros((0, N_MELS), dtype=np.float32)
+
+    window = _hann_window(WIN_LENGTH).astype(np.float32)
+    left_pad = (N_FFT - WIN_LENGTH) // 2
+    frame_starts = np.arange(num_frames) * HOP_LENGTH
+    frames = np.lib.stride_tricks.as_strided(
+        padded,
+        shape=(num_frames, N_FFT),
+        strides=(padded.strides[0] * HOP_LENGTH, padded.strides[0]),
+        writeable=False,
+    ).copy()
+    frames[:, :left_pad] = 0.0
+    frames[:, left_pad : left_pad + WIN_LENGTH] *= window
+    frames[:, left_pad + WIN_LENGTH :] = 0.0
+
+    spectrum = np.fft.rfft(frames, n=N_FFT, axis=1)
+    magnitude = np.abs(spectrum).astype(np.float32) ** 2.0
+
+    if mel_filter is None:
+        mel_filter = mel_filterbank()
+    mel = magnitude @ mel_filter.T.astype(np.float32)
+    log_mel = np.log(mel + LOG_ZERO_GUARD, dtype=np.float32)
+
+    if pad_to and log_mel.shape[0] % pad_to:
+        pad = pad_to - log_mel.shape[0] % pad_to
+        log_mel = np.pad(log_mel, ((0, pad), (0, 0)), mode="constant")
+    return log_mel.astype(np.float32)

python/sortformer_sdk/postprocess.py

@@ -0,0 +1,181 @@
+"""Numpy port of NeMo's Sortformer/VAD post-processing (preds -> RTTM).
+
+Mirrors ``predlist_to_timestamps`` / ``binarization_vectorized`` / ``filtering``
+from ``nemo/collections/asr/parts/utils/vad_utils.py`` and
+``generate_diarization_output_lines`` from ``speaker_utils.py`` (NeMo Speech 3.0).
+"""
+
+from dataclasses import dataclass
+
+import numpy as np
+
+FRAME_LENGTH_IN_SEC = 0.08
+
+
+@dataclass
+class PostProcessingParams:
+    onset: float = 0.5
+    offset: float = 0.5
+    pad_onset: float = 0.0
+    pad_offset: float = 0.0
+    min_duration_on: float = 0.0
+    min_duration_off: float = 0.0
+    filter_speech_first: float = 1.0
+
+
+def merge_overlap_segment(segments: np.ndarray) -> np.ndarray:
+    if segments.shape[0] <= 1:
+        return segments
+    segments = segments[np.argsort(segments[:, 0])]
+    merge_boundary = segments[:-1, 1] >= segments[1:, 0]
+    head_padded = np.concatenate([[False], merge_boundary])
+    tail_padded = np.concatenate([merge_boundary, [False]])
+    head = segments[~head_padded, 0]
+    tail = segments[~tail_padded, 1]
+    return np.stack([head, tail], axis=1)
+
+
+def filter_short_segments(segments: np.ndarray, threshold: float) -> np.ndarray:
+    return segments[segments[:, 1] - segments[:, 0] >= threshold]
+
+
+def get_gap_segments(segments: np.ndarray) -> np.ndarray:
+    segments = segments[np.argsort(segments[:, 0])]
+    return np.column_stack((segments[:-1, 1], segments[1:, 0]))
+
+
+def remove_segments(original_segments: np.ndarray, to_be_removed: np.ndarray) -> np.ndarray:
+    keep = np.ones(original_segments.shape[0], dtype=bool)
+    for segment in to_be_removed:
+        keep &= ~(original_segments == segment).all(axis=1)
+    return original_segments[keep]
+
+
+def filtering(segments: np.ndarray, params: PostProcessingParams) -> np.ndarray:
+    if segments.shape[0] == 0:
+        return segments
+
+    def filter_speech():
+        if params.min_duration_on > 0:
+            return filter_short_segments(segments, params.min_duration_on)
+        return segments
+
+    def restore_short_gaps(current: np.ndarray) -> np.ndarray:
+        if params.min_duration_off <= 0 or current.shape[0] == 0:
+            return current
+        non_speech = get_gap_segments(current)
+        short_gaps = remove_segments(non_speech, filter_short_segments(non_speech, params.min_duration_off))
+        if short_gaps.shape[0] == 0:
+            return current
+        return merge_overlap_segment(np.concatenate([current, short_gaps], axis=0))
+
+    if params.filter_speech_first == 1.0:
+        segments = filter_speech()
+        segments = restore_short_gaps(segments)
+    else:
+        segments = restore_short_gaps(segments)
+        segments = filter_speech()
+    return segments
+
+
+def binarization_vectorized(sequence: np.ndarray, params: PostProcessingParams) -> np.ndarray:
+    empty = np.empty((0, 2), dtype=np.float32)
+    num_frames = sequence.shape[0]
+    if num_frames == 0:
+        return empty
+
+    positions = np.arange(1, num_frames + 1)
+    onset, offset = params.onset, params.offset
+
+    if onset >= offset:
+        force_on = sequence > onset
+        force_off = sequence < offset
+        has_event = force_on | force_off
+        event_positions = np.where(has_event, positions, 0)
+        last_event_positions = np.maximum.accumulate(event_positions)
+        event_states = np.concatenate([[False], force_on])
+        above = event_states[last_event_positions]
+    else:
+        force_on = sequence >= offset
+        force_off = sequence <= onset
+        toggle = (sequence > onset) & (sequence < offset)
+
+        has_reset = force_on | force_off
+        reset_positions = np.where(has_reset, positions, 0)
+        last_reset_positions = np.maximum.accumulate(reset_positions)
+
+        reset_states = np.concatenate([[0], force_on.astype(np.int64)])
+        base_state = reset_states[last_reset_positions]
+        toggle_prefix = np.concatenate([[0], np.cumsum(toggle.astype(np.int64))])
+        toggles_since_reset = toggle_prefix[positions] - toggle_prefix[last_reset_positions]
+        above = np.logical_xor(base_state.astype(bool), (toggles_since_reset % 2).astype(bool))
+
+    padded = np.pad(above.astype(np.float32), (1, 1))
+    diff = padded[1:] - padded[:-1]
+    starts = np.where(diff > 0.5)[0]
+    ends = np.where(diff < -0.5)[0]
+    if starts.shape[0] == 0:
+        return empty
+
+    start_times = np.clip(starts.astype(np.float32) * FRAME_LENGTH_IN_SEC - params.pad_onset, 0.0, None)
+    end_times = ends.astype(np.float32) * FRAME_LENGTH_IN_SEC + params.pad_offset
+    valid = end_times > start_times
+    if not valid.any():
+        return empty
+    segments = np.stack([start_times[valid], end_times[valid]], axis=1)
+    if params.pad_onset > 0 or params.pad_offset > 0:
+        segments = merge_overlap_segment(segments)
+    return segments
+
+
+def predlist_to_timestamps(
+    preds: np.ndarray,
+    offset: float = 0.0,
+    params: PostProcessingParams = None,
+    bypass_postprocessing: bool = False,
+    precision: int = 2,
+):
+    """Convert (num_frames, num_speakers) probabilities to per-speaker timestamps."""
+    if params is None:
+        params = PostProcessingParams()
+    if bypass_postprocessing:
+        params = PostProcessingParams(onset=0.5, offset=0.5)
+
+    timestamps = []
+    for spk in range(preds.shape[1]):
+        segments = binarization_vectorized(preds[:, spk], params)
+        if not bypass_postprocessing:
+            segments = filtering(segments, params)
+        if segments.shape[0] == 0:
+            timestamps.append([])
+            continue
+        segments = segments + offset
+        timestamps.append([[round(float(start), precision), round(float(end), precision)] for start, end in segments])
+    return timestamps
+
+
+def generate_diarization_output_lines(timestamps, model_spk_num: int):
+    lines = []
+    for spk_idx in range(model_spk_num):
+        if not timestamps[spk_idx]:
+            continue
+        intervals = np.asarray(timestamps[spk_idx], dtype=np.float32).reshape(-1, 2)
+        for start, end in merge_overlap_segment(intervals):
+            lines.append(f"{start:.3f} {end:.3f} speaker_{int(spk_idx)}")
+    return lines
+
+
+def timestamps_to_rttm_lines(timestamps, uri: str, model_spk_num: int):
+    lines = []
+    for spk_idx in range(model_spk_num):
+        intervals = timestamps[spk_idx]
+        if not intervals:
+            continue
+        merged = merge_overlap_segment(np.asarray(intervals, dtype=np.float32).reshape(-1, 2))
+        for start, end in merged:
+            duration = float(end) - float(start)
+            if duration > 0:
+                lines.append(
+                    f"SPEAKER {uri} 1 {float(start):.3f} {duration:.3f} <NA> <NA> speaker_{int(spk_idx)} <NA>"
+                )
+    return lines

python/sortformer_sdk/state.py

@@ -0,0 +1,210 @@
+"""Numpy port of NeMo's synchronous Sortformer streaming state machine.
+
+Mirrors ``SortformerModules.streaming_update`` (eval, batch size 1, no speaker
+permutation, learnable silence embedding disabled) from NeMo Speech 3.0:
+
+    nemo/collections/asr/modules/sortformer_modules.py
+
+The module is dependency-free so it can run on the host (base env) and be
+reused by the AX650 board SDK.
+"""
+
+import math
+from dataclasses import dataclass, field
+
+import numpy as np
+
+NEG_INF = float("-inf")
+
+
+@dataclass
+class SortformerConfig:
+    num_speakers: int = 4
+    fc_d_model: int = 512
+    subsampling_factor: int = 8
+    chunk_len: int = 6
+    chunk_left_context: int = 1
+    chunk_right_context: int = 7
+    fifo_len: int = 188
+    spkcache_len: int = 188
+    spkcache_update_period: int = 144
+    spkcache_sil_frames_per_spk: int = 3
+    pred_score_threshold: float = 0.25
+    max_index: int = 99999
+    scores_boost_latest: float = 0.05
+    sil_threshold: float = 0.2
+    strong_boost_rate: float = 0.75
+    weak_boost_rate: float = 1.5
+    min_pos_scores_rate: float = 0.5
+    use_learnable_sil_emb: bool = False
+
+
+@dataclass
+class StreamingState:
+    spkcache: np.ndarray = field(default_factory=lambda: np.zeros((0, 512), dtype=np.float32))
+    spkcache_preds: np.ndarray = field(default_factory=lambda: np.zeros((0, 4), dtype=np.float32))
+    spkcache_compressed: bool = False
+    fifo: np.ndarray = field(default_factory=lambda: np.zeros((0, 512), dtype=np.float32))
+    fifo_preds: np.ndarray = field(default_factory=lambda: np.zeros((0, 4), dtype=np.float32))
+    mean_sil_emb: np.ndarray = field(default_factory=lambda: np.zeros(512, dtype=np.float32))
+    n_sil_frames: int = 0
+
+
+def init_state(cfg: SortformerConfig) -> StreamingState:
+    state = StreamingState()
+    state.mean_sil_emb = np.zeros(cfg.fc_d_model, dtype=np.float32)
+    return state
+
+
+def streaming_update(cfg: SortformerConfig, state: StreamingState, chunk: np.ndarray, preds: np.ndarray, lc: int, rc: int):
+    """Update speaker cache / FIFO with one chunk; returns the chunk predictions.
+
+    ``chunk`` has shape (lc + chunk_len + rc, emb_dim); ``preds`` has shape
+    (spkcache_len + fifo_len + chunk.shape[0], num_speakers) and covers the
+    speaker cache, FIFO and chunk regions.
+    """
+    spkcache_len = state.spkcache.shape[0]
+    fifo_len = state.fifo.shape[0]
+    chunk_len = chunk.shape[0] - lc - rc
+
+    state.fifo_preds = preds[spkcache_len : spkcache_len + fifo_len]
+    chunk_body = chunk[lc : chunk_len + lc]
+    chunk_preds = preds[spkcache_len + fifo_len + lc : spkcache_len + fifo_len + chunk_len + lc]
+
+    state.fifo = np.concatenate([state.fifo, chunk_body], axis=0)
+    state.fifo_preds = np.concatenate([state.fifo_preds, chunk_preds], axis=0)
+
+    if fifo_len + chunk_len > cfg.fifo_len:
+        pop_out_len = cfg.spkcache_update_period
+        pop_out_len = max(pop_out_len, chunk_len - cfg.fifo_len + fifo_len)
+        pop_out_len = min(pop_out_len, fifo_len + chunk_len)
+
+        pop_out_embs = state.fifo[:pop_out_len]
+        pop_out_preds = state.fifo_preds[:pop_out_len]
+        if not cfg.use_learnable_sil_emb:
+            state.mean_sil_emb, state.n_sil_frames = _get_silence_profile(
+                cfg, state.mean_sil_emb, state.n_sil_frames, pop_out_embs, pop_out_preds
+            )
+        state.fifo = state.fifo[pop_out_len:]
+        state.fifo_preds = state.fifo_preds[pop_out_len:]
+
+        state.spkcache = np.concatenate([state.spkcache, pop_out_embs], axis=0)
+        if state.spkcache_compressed:
+            state.spkcache_preds = np.concatenate([state.spkcache_preds, pop_out_preds], axis=0)
+        else:
+            state.spkcache_preds = np.concatenate([preds[:spkcache_len], pop_out_preds], axis=0)
+        if state.spkcache.shape[0] > cfg.spkcache_len:
+            state.spkcache, state.spkcache_preds = _compress_spkcache(
+                cfg, state.spkcache, state.spkcache_preds, state.mean_sil_emb
+            )
+            state.spkcache_compressed = True
+
+    return state, chunk_preds
+
+
+def _get_silence_profile(cfg, mean_sil_emb, n_sil_frames, emb_seq, preds):
+    is_sil = preds.sum(axis=1) < cfg.sil_threshold
+    sil_count = int(is_sil.sum())
+    if sil_count == 0:
+        return mean_sil_emb, n_sil_frames
+    sil_emb_sum = (emb_seq * is_sil[:, None]).sum(axis=0)
+    upd_n_sil_frames = n_sil_frames + sil_count
+    total_sil_sum = mean_sil_emb * n_sil_frames + sil_emb_sum
+    upd_mean_sil_emb = total_sil_sum / max(upd_n_sil_frames, 1)
+    return upd_mean_sil_emb.astype(np.float32), upd_n_sil_frames
+
+
+def _get_log_pred_scores(cfg, preds):
+    log_probs = np.log(np.clip(preds, cfg.pred_score_threshold, None))
+    log_1_probs = np.log(np.clip(1.0 - preds, cfg.pred_score_threshold, None))
+    log_1_probs_sum = log_1_probs.sum(axis=1, keepdims=True)
+    return log_probs - log_1_probs + log_1_probs_sum - math.log(0.5)
+
+
+def _disable_low_scores(cfg, preds, scores, min_pos_scores_per_spk):
+    is_speech = preds > 0.5
+    scores = np.where(is_speech, scores, NEG_INF)
+    is_pos = scores > 0
+    # NeMo sums over the frame dimension (torch: is_pos.sum(dim=1)); batch-free here -> axis=0.
+    is_nonpos_replace = (~is_pos) & is_speech & (is_pos.sum(axis=0, keepdims=True) >= min_pos_scores_per_spk)
+    return np.where(is_nonpos_replace, NEG_INF, scores)
+
+
+def _boost_topk_scores(cfg, scores, n_boost_per_spk, scale_factor=1.0, offset=0.5):
+    n_frames, n_spk = scores.shape
+    n_boost_per_spk = min(n_boost_per_spk, n_frames)
+    if n_boost_per_spk <= 0:
+        return scores
+    for spk in range(n_spk):
+        column = scores[:, spk]
+        # Stable descending order: ties keep the smaller frame index (matches C++).
+        order = np.argsort(-column, kind="stable")[:n_boost_per_spk]
+        scores[order, spk] -= scale_factor * math.log(offset)
+    return scores
+
+
+def _get_topk_indices(cfg, scores):
+    n_frames, n_spk = scores.shape
+    n_frames_no_sil = n_frames - cfg.spkcache_sil_frames_per_spk
+    scores_flatten = scores.T.reshape(-1)  # speaker-major, matches permute(0, 2, 1).reshape()
+    # Stable descending order: ties keep the smaller flat index (matches C++).
+    order = np.argsort(-scores_flatten, kind="stable")
+    k = min(cfg.spkcache_len, scores_flatten.shape[0])
+    topk_indices = order[:k]
+    values = scores_flatten[topk_indices]
+    topk_indices = np.where(values != NEG_INF, topk_indices, cfg.max_index)
+    topk_indices_sorted = np.sort(topk_indices)
+    is_disabled = topk_indices_sorted == cfg.max_index
+    topk_indices_sorted = np.remainder(topk_indices_sorted, n_frames)
+    is_disabled = is_disabled | (topk_indices_sorted >= n_frames_no_sil)
+    topk_indices_sorted = np.where(is_disabled, 0, topk_indices_sorted)
+    return topk_indices_sorted, is_disabled
+
+
+def _compress_spkcache(cfg, emb_seq, preds, mean_sil_emb):
+    n_frames, n_spk = preds.shape
+    spkcache_len_per_spk = cfg.spkcache_len // n_spk - cfg.spkcache_sil_frames_per_spk
+    strong_boost_per_spk = math.floor(spkcache_len_per_spk * cfg.strong_boost_rate)
+    weak_boost_per_spk = math.floor(spkcache_len_per_spk * cfg.weak_boost_rate)
+    min_pos_scores_per_spk = math.floor(spkcache_len_per_spk * cfg.min_pos_scores_rate)
+
+    scores = _get_log_pred_scores(cfg, preds)
+    scores = _disable_low_scores(cfg, preds, scores, min_pos_scores_per_spk)
+    if cfg.scores_boost_latest > 0:
+        scores[cfg.spkcache_len :, :] += cfg.scores_boost_latest
+    scores = _boost_topk_scores(cfg, scores, strong_boost_per_spk, scale_factor=2)
+    scores = _boost_topk_scores(cfg, scores, weak_boost_per_spk, scale_factor=1)
+
+    if cfg.spkcache_sil_frames_per_spk > 0:
+        pad = np.full((cfg.spkcache_sil_frames_per_spk, n_spk), np.inf, dtype=scores.dtype)
+        scores = np.concatenate([scores, pad], axis=0)
+
+    topk_indices, is_disabled = _get_topk_indices(cfg, scores)
+    spkcache = emb_seq[topk_indices]
+    spkcache = np.where(is_disabled[:, None], mean_sil_emb[None, :], spkcache)
+    spkcache_preds = preds[topk_indices]
+    spkcache_preds = np.where(is_disabled[:, None], 0.0, spkcache_preds)
+    return spkcache.astype(np.float32), spkcache_preds.astype(np.float32)
+
+
+def iter_chunks(cfg: SortformerConfig, features: np.ndarray):
+    """Yield (chunk_mel, left_offset, right_offset) following NeMo's ``streaming_feat_loader``."""
+    feat_len = features.shape[0]
+    start = 0
+    while start < feat_len:
+        left_offset = min(cfg.chunk_left_context * cfg.subsampling_factor, start)
+        end = min(start + cfg.chunk_len * cfg.subsampling_factor, feat_len)
+        right_offset = min(cfg.chunk_right_context * cfg.subsampling_factor, feat_len - end)
+        chunk = features[start - left_offset : end + right_offset]
+        yield chunk, left_offset, right_offset
+        start = end
+
+
+def pre_encode_length(mel_frames: int, num_layers: int = 3) -> int:
+    """Number of encoder frames after NeMo's dw_striding pre-encode (stride 2, kernel 3)."""
+    length = int(mel_frames)
+    for _ in range(num_layers):
+        if length <= 0:
+            return 0
+        length = (length - 1) // 2 + 1
+    return length

requirements.txt

@@ -0,0 +1,4 @@
+numpy>=1.21
+soundfile>=0.12
+# 板端 NPU 运行库（板端镜像已内置；主机安装见 README）
+# axengine

run_ax650.sh

@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+# Sortformer 一键说话人日志 (Python, axengine)
+# 用法: ./run_ax650.sh [input.wav] [output.rttm] [--fast]
+set -euo pipefail
+cd "$(dirname "$0")"
+
+export LD_LIBRARY_PATH=/soc/lib:${LD_LIBRARY_PATH:-}
+
+WAV="${1:-samples/sample_meeting.wav}"
+RTTM="${2:-$(basename "${WAV%.*}").rttm}"
+EXTRA=("${@:3}")
+
+PYTHON_BIN="${PYTHON_BIN:-python3}"
+if [[ -x /root/miniforge3/bin/python ]]; then
+  PYTHON_BIN=/root/miniforge3/bin/python
+fi
+
+"$PYTHON_BIN" python/example.py "$WAV" "$RTTM" "${EXTRA[@]}"

run_cpp_ax650.sh

@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+# Sortformer 一键说话人日志 (C++, 预编译可执行文件)
+# 用法: ./run_cpp_ax650.sh [input.wav] [output.rttm] [--fast]
+set -euo pipefail
+cd "$(dirname "$0")"
+
+export LD_LIBRARY_PATH=/soc/lib:${LD_LIBRARY_PATH:-}
+
+WAV="${1:-samples/sample_meeting.wav}"
+RTTM="${2:-$(basename "${WAV%.*}").rttm}"
+FAST=0
+for arg in "$@"; do
+  [[ "$arg" == "--fast" ]] && FAST=1
+done
+
+if [[ "$FAST" == "1" ]]; then
+  ENCODER=models/encoder_fifo40.axmodel
+  EXTRA="--fifo-len 40 --spkcache-update-period 31"
+else
+  ENCODER=models/encoder.axmodel
+  EXTRA=""
+fi
+
+./bin/sortformer_ax650 \
+  --preencode models/preencode.axmodel --encoder "$ENCODER" \
+  --wav "$WAV" --rttm "$RTTM" --threads 8 $EXTRA

samples/sample_meeting.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1d66664d87a7f8fe98c8cb44e61d26e8f97c8eaba1572234fcb2c38703f21ba
+size 3840044