Add non-streaming mode (streaming=False), clearer usage in README

README.md

@@ -25,21 +25,52 @@ speaker ids to the slots of the head.
 
 Demo: https://kyudan1.github.io/eDIAR/ · Code: https://github.com/KyuDan1/eDIAR
 
-## Usage
+## Usage: streaming or non-streaming
+
+One class, one switch. Everything else (models, clustering, speaker ids) is the same.
+
+| | `streaming=True` (default) | `streaming=False` |
+|---|---|---|
+| Audio | arrives in chunks, causal | whole recording at once |
+| Latency | 1.04 s (`latency="low_latency"`) or 0.64 s (`"very_low_latency"`) | none, runs after the recording |
+| Labels | provisional, revised for 10 s, then fixed | final |
+| Clustering | bounded 90 s window, re-run every chunk | all embeddings of the session at once |
+| Use it for | live meetings, long recordings, bounded cost | files, highest accuracy |
+| Paper setting | streaming tables | non-streaming tables |
 
 ```python
 import soundfile as sf
 from emeans_diar import SlotEMeansDiarizer
 
-wav, sr = sf.read("meeting.wav", dtype="float32")      # mono
-dia = SlotEMeansDiarizer(device="cuda", mode="low_latency")   # 1.04 s input buffer; "very_low_latency" = 0.64 s
+wav, sr = sf.read("meeting.wav", dtype="float32")             # mono
+
+# 1) streaming (default): audio is fed in 100 ms chunks like a live system
+dia = SlotEMeansDiarizer(device="cuda", streaming=True, latency="low_latency")
 for start, end, speaker in dia.diarize(wav, sr):
     print(f"{start:7.2f} {end:7.2f} spk{speaker}")
+
+# 2) non-streaming: the whole file at once
+dia = SlotEMeansDiarizer(device="cuda", streaming=False)
+segments = dia.diarize(wav, sr)
+
+# 3) real streaming loop (live audio): incremental interface
+dia = SlotEMeansDiarizer(streaming=True)
+session = dia.new_session()
+for pcm in microphone_chunks():                               # 16 kHz int16 bytes, any chunk size
+    msgs = session.feed(pcm)                                  # provisional + revised labels
+session.feed(None, flush=True)
+```
+
+Command line, RTTM on stdout:
+
+```bash
+python examples/diarize.py meeting.wav                        # streaming, 1.04 s
+python examples/diarize.py meeting.wav --latency very_low_latency
+python examples/diarize.py meeting.wav --offline              # non-streaming
 ```
 
-or `python examples/diarize.py meeting.wav > meeting.rttm`. The two upstream models are downloaded from their own
-repositories on first use (`HF_TOKEN` if they are gated for you). Install with `pip install -r requirements.txt`.
-`SlotEMeansDiarizer.new_session()` returns the incremental interface used by the live demo (`feed(pcm_int16_bytes)`).
+The two upstream models are downloaded from their own repositories on first use (`HF_TOKEN` if they are gated for you).
+Install with `pip install -r requirements.txt`.
 
 ## Results (streaming, 1.04 s input buffer, DER %, collar 0.25 s, overlap scored)
 

emeans_diar/offline.py

@@ -0,0 +1,34 @@
+"""Non-streaming (whole-recording) path: the same slot head and the same slot-conditioned e-means, but the clustering sees every
+embedding of the session at once and the encoders run over the full recording (no causal restrictions, no revision horizon)."""
+import numpy as np
+import torch
+import torch.nn.functional as F
+import librosa
+
+
+def _fbank(engine, wav):
+    wp = np.append(wav[0], wav[1:] - 0.97 * wav[:-1])
+    p = np.abs(librosa.stft(wp * 32768.0, n_fft=512, hop_length=160, win_length=400, window="hann", center=False)) ** 2
+    return np.log(np.maximum(engine.mel_basis_80 @ p, 1e-10)).T.astype(np.float32)
+
+
+def extract(engine, wav, win_ms=1520):
+    """Returns (spk_embs [T x 256, one per 80 ms], probs [T x 8] slot posteriors per 80 ms)."""
+    dev, r = engine.device, engine.resnet
+    mel = _fbank(engine, wav); T = mel.shape[0]; chunk, ctx = 4800, 160; fm = []
+    with torch.inference_mode():
+        for s in range(0, T, chunk):
+            e = min(T, s + chunk); sp, ep = max(0, s - ctx), min(T, e + ctx)
+            m = torch.from_numpy(mel[sp:ep]).to(dev); m = m - m.mean(0, keepdim=True)
+            o = r.layer4(r.layer3(r.layer2(r.layer1(F.relu(r.bn1(r.conv1(m.unsqueeze(0).permute(0, 2, 1).unsqueeze(1))))))))
+            o = o.squeeze(0).reshape(2560, -1); a = (s - sp) // 8; fm.append(o[:, a:a + (e - s) // 8].cpu())
+        f = torch.cat(fm, 1).to(dev)
+        w = int(round(win_ms / 80.0)); w += (w % 2 == 0); pl = w // 2
+        win = F.pad(f.unsqueeze(0), (pl, w - 1 - pl), mode="replicate").squeeze(0).unfold(1, w, 1)
+        stats = torch.cat([win.mean(-1), torch.sqrt(win.var(-1) + 1e-7)], 0).T
+        embs = F.normalize(r.seg_1(stats), p=2, dim=-1).cpu().numpy()
+        inp = engine.fe(wav, sampling_rate=16000, return_tensors="pt")
+        inp = {k: v.to(dev) if isinstance(v, torch.Tensor) else v for k, v in inp.items()}
+        p10 = engine.nemotron(**inp, output_hidden_states=False).logits[0].sigmoid().cpu().numpy()
+    n = min(len(embs), len(p10) // 8)
+    return embs[:n], p10[:n * 8].reshape(n, 8, 8).mean(1)

emeans_diar/pipeline.py

@@ -39,12 +39,17 @@ class _Engine:
 
 
 class SlotEMeansDiarizer:
-    """Streaming diarizer. ``diarize(wav)`` feeds a 16 kHz mono waveform in 100 ms chunks (as a live system would) and returns
-    speaker segments after the revision horizon; ``stream()`` is the incremental interface."""
+    """Open-speaker diarizer on top of the Nemotron-3 slot head. Choose how it runs:
 
-    def __init__(self, device="cuda", mode="low_latency", token=None, **live_kwargs):
-        assert mode in MODES
-        self.mode, self.device, self.live_kwargs = mode, device, live_kwargs
+      SlotEMeansDiarizer(streaming=True,  latency="low_latency")   # default: causal, audio arrives in chunks, labels may be revised for 10 s
+      SlotEMeansDiarizer(streaming=False)                          # whole recording at once: best accuracy, needs the full audio first
+
+    ``diarize(wav)`` works in both modes and returns [(start_s, end_s, speaker_id)]. In streaming mode it feeds the waveform in 100 ms
+    chunks as a live system would; ``new_session()`` gives the incremental interface (streaming only)."""
+
+    def __init__(self, device="cuda", mode="low_latency", token=None, streaming=True, latency=None, **live_kwargs):
+        mode = latency or mode; assert mode in MODES
+        self.streaming, self.mode, self.device, self.live_kwargs = streaming, mode, device, live_kwargs
         self.engine = _Engine(device, token or os.environ.get("HF_TOKEN"))
         self.proc = Nemotron3DiarizationProcessor(feature_extractor=self.engine.fe, subsampling_factor=8, streaming_mode=mode)
 
@@ -53,18 +58,27 @@ class SlotEMeansDiarizer:
         return cls(*args, **kw)
 
     def new_session(self):
+        if not self.streaming: raise RuntimeError("new_session() is the incremental interface; create the diarizer with streaming=True")
         e = self.engine
         return LiveSession(LiveNemotron(e.nemotron, self.proc, self.device, self.mode), LiveResNet(e), partial(OfficialEMeansSlotDiarizer, n_init=1), self.mode, **self.live_kwargs)
 
     def diarize(self, wav, sr=16000, chunk_s=0.1):
         """wav: float array in [-1, 1], mono. Returns [(start_s, end_s, speaker_id)], speaker ids are permanent (slot-anchored)."""
         if sr != 16000: wav = librosa.resample(np.asarray(wav, np.float32), orig_sr=sr, target_sr=16000)
+        if not self.streaming: return self._diarize_offline(np.asarray(wav, np.float32))
         pcm = (np.clip(np.asarray(wav, np.float32), -1, 1) * 32767).astype("<i2").tobytes(); sess = self.new_session(); step = int(chunk_s * 16000) * 2
         for k in range(0, len(pcm), step): sess.feed(pcm[k:k + step])
         sess.feed(None, flush=True)
         return frames_to_segments(sess.ld.cur)
 
 
+    def _diarize_offline(self, wav):
+        from .offline import extract
+        embs, probs = extract(self.engine, wav)
+        frames, *_ = OfficialEMeansSlotDiarizer(stride_ms=80).fit_predict(embs, probs)
+        return frames_to_segments(frames)
+
+
 def frames_to_segments(frame_labels, gap_s=0.16):
     """per-80 ms-frame lists of speaker ids -> merged segments (gaps shorter than gap_s are bridged)"""
     open_, out = {}, []

examples/diarize.py

@@ -1,13 +1,14 @@
-"""python examples/diarize.py meeting.wav [--mode low_latency|very_low_latency] [--device cuda]  ->  RTTM on stdout"""
+"""python examples/diarize.py meeting.wav [--offline] [--latency low_latency|very_low_latency] [--device cuda]  ->  RTTM on stdout
+   default: streaming (causal, 100 ms chunks).  --offline: non-streaming, the whole recording at once."""
 import argparse, sys, os
 sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
 import soundfile as sf
 from emeans_diar import SlotEMeansDiarizer
 
-ap = argparse.ArgumentParser(); ap.add_argument("wav"); ap.add_argument("--mode", default="low_latency"); ap.add_argument("--device", default="cuda")
+ap = argparse.ArgumentParser(); ap.add_argument("wav"); ap.add_argument("--latency", "--mode", dest="latency", default="low_latency", choices=["low_latency", "very_low_latency"], help="streaming input buffer: 1.04 s or 0.64 s"); ap.add_argument("--offline", action="store_true", help="non-streaming: use the whole recording at once"); ap.add_argument("--device", default="cuda")
 a = ap.parse_args()
 wav, sr = sf.read(a.wav, dtype="float32")
 if wav.ndim > 1: wav = wav.mean(1)
 name = os.path.splitext(os.path.basename(a.wav))[0]
-for s, e, spk in SlotEMeansDiarizer(device=a.device, mode=a.mode).diarize(wav, sr):
+for s, e, spk in SlotEMeansDiarizer(device=a.device, latency=a.latency, streaming=not a.offline).diarize(wav, sr):
     print(f"SPEAKER {name} 1 {s:.3f} {e - s:.3f} <NA> <NA> spk{spk} <NA> <NA>")