"""DualTurn Endpointing — HuggingFace custom model (trust_remote_code). A small, causal, dual-channel turn-taking model. Adapted from the DualTurn paper and fine-tuned for the endpointing head. Consumes a 2-channel conversation (CH0=user, CH1=agent), runs a frozen Mimi encoder to get 12.5 Hz continuous features per channel, and emits per-frame turn-taking probabilities. from transformers import AutoModel model = AutoModel.from_pretrained("anyreach-ai/dualturn-endpointing", trust_remote_code=True).eval() wav, sr = torchaudio.load("conversation.wav") # [2, T] CH0=user CH1=agent out = model(wav, sr=sr) out.eot_probs # [1, T', 1] P(user end-of-turn) — user channel only out.vad_probs # [1, T', 2] P(speaking) — (user, agent) — both channels out.fvad_probs # [1, T', 2, 4] P(speaks soon) — (user, agent) × horizons — both channels # fvad horizons = 0-240 / 240-640 / 640-1200 / 1200-2000 ms # frames are 80 ms (12.5 Hz); T' ≈ audio_seconds * 12.5. """ from __future__ import annotations import os os.environ.setdefault("USE_TORCH_XLA", "0") from dataclasses import dataclass from typing import Optional import torch import torch.nn as nn from transformers import PreTrainedModel, PretrainedConfig from transformers.utils import ModelOutput MIMI_SR = 24000 _MIMI_CACHE = {} class DualTurnEndpointingConfig(PretrainedConfig): model_type = "dualturn_endpointing" def __init__(self, feat_dim=1024, d_model=256, core_hidden=256, core_layers=2, d_head=128, fvad_dim=8, mimi_model="kyutai/mimi", hop_s=0.08, **kw): self.feat_dim = feat_dim self.d_model = d_model self.core_hidden = core_hidden self.core_layers = core_layers self.d_head = d_head self.fvad_dim = fvad_dim self.mimi_model = mimi_model self.hop_s = hop_s super().__init__(**kw) @dataclass class DualTurnOutput(ModelOutput): eot_probs: Optional[torch.FloatTensor] = None # [B, T, 1] user end-of-turn vad_probs: Optional[torch.FloatTensor] = None # [B, T, 2] (user, agent) speaking fvad_probs: Optional[torch.FloatTensor] = None # [B, T, 2, 4] (user, agent) × future horizons # frames are 80 ms (12.5 Hz) class _Head(nn.Module): def __init__(self, d, d_head, out, dropout=0.0): super().__init__() self.net = nn.Sequential(nn.Linear(d, d_head), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_head, out)) def forward(self, y): return self.net(y) class _MultiHead(nn.Module): def __init__(self, d, d_head, dims): super().__init__() self.heads = nn.ModuleDict({k: _Head(d, d_head, o) for k, o in dims.items()}) def forward(self, y): return {k: h(y) for k, h in self.heads.items()} class DualTurnEndpointingModel(PreTrainedModel): config_class = DualTurnEndpointingConfig def __init__(self, config: DualTurnEndpointingConfig): super().__init__(config) c = config self.proj = nn.Sequential(nn.Linear(c.feat_dim, c.d_model), nn.LayerNorm(c.d_model), nn.Dropout(0.0)) self.core = _LSTM(c.d_model, c.core_hidden, c.core_layers) dims = {"h1": 1, "vad_user": 1, "vad_agent": 1, "fvad": c.fvad_dim} self.heads = _MultiHead(c.core_hidden, c.d_head, dims) # input standardization (per-dim mean/std of the 512-d Mimi features); baked into the checkpoint self.register_buffer("feat_mean", torch.zeros(c.feat_dim // 2)) self.register_buffer("feat_std", torch.ones(c.feat_dim // 2)) self.post_init() # ---- Mimi (loaded lazily; NOT part of this checkpoint — pulled from `mimi_model`) ---- def _mimi(self): key = (self.config.mimi_model, str(self.device)) if key not in _MIMI_CACHE: from transformers import MimiModel attn = "sdpa" try: tv = tuple(int(x) for x in torch.__version__.split("+")[0].split(".")[:3]) if self.device.type != "cuda" or tv < (2, 1, 1): attn = "eager" except Exception: attn = "eager" _MIMI_CACHE[key] = MimiModel.from_pretrained( self.config.mimi_model, attn_implementation=attn).eval().to(self.device) return _MIMI_CACHE[key] @torch.no_grad() def _encode(self, mono, sr): """mono [T] tensor -> [T', 512] causal Mimi features at 12.5 Hz.""" import torchaudio.functional as AF x = torch.as_tensor(mono, dtype=torch.float32, device=self.device) if sr != MIMI_SR: # Mimi's input rate x = AF.resample(x, sr, MIMI_SR) m = self._mimi() enc = m.encoder(x[None, None]) # [1, C, T] h = enc.transpose(1, 2) mask = torch.ones(h.shape[0], h.shape[1], dtype=torch.long, device=self.device) # causal sliding-window et = m.encoder_transformer(h, attention_mask=mask) et = et.last_hidden_state if hasattr(et, "last_hidden_state") else et[0] ds = m.downsample(et.transpose(1, 2)) # [1, 512, T'] return ds.squeeze(0).transpose(0, 1) # [T', 512] @torch.no_grad() def forward(self, wav, sr: int = 24000): """wav: [2, T] (CH0=user, CH1=agent) or [T]/[1, T] (user only, silent agent). Returns DualTurnOutput.""" w = torch.as_tensor(wav, dtype=torch.float32) if w.dim() == 1: user, agent = w, torch.zeros_like(w) elif w.dim() == 2 and w.shape[0] == 2: user, agent = w[0], w[1] elif w.dim() == 2 and w.shape[0] == 1: user, agent = w[0], torch.zeros_like(w[0]) else: raise ValueError(f"expected wav [2,T] (user,agent) or [T]/[1,T]; got {tuple(w.shape)}") fu = self._encode(user, sr) fa = self._encode(agent, sr) T = min(fu.shape[0], fa.shape[0]) fu = (fu[:T] - self.feat_mean) / self.feat_std fa = (fa[:T] - self.feat_mean) / self.feat_std feat = torch.cat([fu, fa], dim=-1).unsqueeze(0) # [1, T, 1024] y, _ = self.core(self.proj(feat)) h = self.heads(y) sg = torch.sigmoid eot = sg(h["h1"]) # [1, T, 1] user only vad = torch.cat([sg(h["vad_user"]), sg(h["vad_agent"])], dim=-1) # [1, T, 2] (user, agent) B, Tt = eot.shape[0], eot.shape[1] fvad = sg(h["fvad"]).reshape(B, Tt, 2, 4) # [1, T, 2, 4] (user, agent) × 4 horizons return DualTurnOutput(eot_probs=eot, vad_probs=vad, fvad_probs=fvad) def streaming(self): """Return a stateful real-time streamer (fp32). Feed 80 ms chunks; causal, O(1)/tick.""" return DualTurnStreamer(self) class _LSTM(nn.Module): def __init__(self, d_in, hidden, layers): super().__init__() self.lstm = nn.LSTM(d_in, hidden, num_layers=layers, batch_first=True, bidirectional=False) def forward(self, x, state=None): return self.lstm(x, state) class DualTurnStreamer: """Real-time streaming inference (fp32, PyTorch), sharing the offline model's weights. Causal — uses only PAST audio. Bounded state → O(1) per tick, flat latency + memory over a whole call. Feed 24 kHz `[2, n]` chunks (CH0=user, CH1=agent), n a multiple of 960 (80 ms = 1920 samples); mono `[n]` is accepted (silent agent). `push` returns a `DualTurnOutput` for the new 12.5 Hz frame(s), or `None` if a chunk didn't advance a full output frame yet. Call `reset()` between calls. """ RF_AUDIO = 7680 # 0.32 s conv left-context (exact newest enc-frames) W_KV = 250 # transformer KV window (= config.sliding_window) TF_KEEP = 260 # bounded tf-frame history for the downsample conv def __init__(self, model): from transformers.cache_utils import DynamicCache self._Cache = DynamicCache self.mimi = model._mimi(); self.dev = model.device self.proj, self.core, self.heads = model.proj, model.core, model.heads self.mean, self.std = model.feat_mean, model.feat_std self.reset() def reset(self): self.audio_ctx = torch.zeros(2, 0, device=self.dev) self.kv = self._Cache() self.tf_hist = torch.zeros(2, 0, 512, device=self.dev) self.n_out = 0; self.n_tf = 0; self.lstm = None; self.pos = 0 @torch.no_grad() def push(self, chunk, sr: int = 24000): import torchaudio.functional as AF x = torch.as_tensor(chunk, dtype=torch.float32, device=self.dev) if x.dim() == 1: x = torch.stack([x, torch.zeros_like(x)]) # mono user -> silent agent if sr != MIMI_SR: # prefer feeding 24 kHz to avoid chunk-edge artifacts x = AF.resample(x, sr, MIMI_SR) buf = torch.cat([self.audio_ctx, x], 1) emb = self.mimi.encoder(buf[:, None, :]) # [2, C, ne] n_new = x.shape[1] // 960 # new enc-frames (25 Hz) if n_new < 1: self.audio_ctx = buf[:, -self.RF_AUDIO:]; return None new_emb = emb[..., -n_new:].transpose(1, 2) # [2, n_new, C] self.audio_ctx = buf[:, -self.RF_AUDIO:] tf = [] for j in range(n_new): # feed 1 frame/step (exact) w/ ABSOLUTE position (RoPE) cp = torch.tensor([self.pos], device=self.dev) pid = torch.tensor([[self.pos]], device=self.dev) o = self.mimi.encoder_transformer(new_emb[:, j:j+1], past_key_values=self.kv, use_cache=True, cache_position=cp, position_ids=pid) tf.append(o[0]); self.kv = o[1] if len(o) > 1 else self.kv; self.pos += 1 for li in range(len(self.kv.key_cache)): # cap KV to the sliding window if self.kv.key_cache[li].shape[-2] > self.W_KV: self.kv.key_cache[li] = self.kv.key_cache[li][..., -self.W_KV:, :].contiguous() self.kv.value_cache[li] = self.kv.value_cache[li][..., -self.W_KV:, :].contiguous() new_tf = torch.cat(tf, 1); self.n_tf += n_new self.tf_hist = torch.cat([self.tf_hist, new_tf], 1)[:, -self.TF_KEEP:] # bounded history ds = self.mimi.downsample(self.tf_hist.transpose(1, 2)).transpose(1, 2) # [2, m, 512] n_out_total = max(0, (self.n_tf - 2) // 2 + 1); k = n_out_total - self.n_out; self.n_out = n_out_total if k <= 0: return None new = ds[:, -k:] u = (new[0] - self.mean) / self.std; a = (new[1] - self.mean) / self.std feat = torch.cat([u, a], -1)[None] # [1, k, 1024] y, self.lstm = self.core(self.proj(feat), self.lstm) # LSTM state carry (h, c) h = self.heads(y); sg = torch.sigmoid eot = sg(h["h1"]) vad = torch.cat([sg(h["vad_user"]), sg(h["vad_agent"])], -1) fvad = sg(h["fvad"]).reshape(1, -1, 2, 4) return DualTurnOutput(eot_probs=eot, vad_probs=vad, fvad_probs=fvad)