Improved Nemotron ASR streaming

mlx_audio/stt/models/nemotron_asr/__init__.py

@@ -1,3 +1,10 @@
+from .audio import StreamingLogMelSpectrogram
 from .nemotron_asr import Model, ModelConfig
+from .streaming import ConformerStreamingState
 
-__all__ = ["Model", "ModelConfig"]
+__all__ = [
+    "ConformerStreamingState",
+    "Model",
+    "ModelConfig",
+    "StreamingLogMelSpectrogram",
+]

mlx_audio/stt/models/nemotron_asr/audio.py

@@ -10,6 +10,7 @@
 additive guard, and no dither (dither is training-only in NeMo).
 """
 
+import math
 from collections.abc import Iterator
 
 import mlx.core as mx
@@ -81,7 +82,7 @@ def log_mel_spectrogram(x: mx.array, args: PreprocessArgs) -> mx.array:
     window = _padded_window(args)
     x = _preemphasize(x, args)
 
-    x = stft(x, args.n_fft, args.hop_length, args.n_fft, window, pad_mode="constant")
+    x = stft(x, args.n_fft, args.hop_length, args.n_fft, window, pad_mode="reflect")
     # Power spectrum (mag_power = 2.0).
     x = mx.square(mx.abs(x)).astype(original_dtype)
 
@@ -135,10 +136,10 @@ def log_mel_spectrogram_frames(
     right_pad = max(sample_end - total_samples, 0)
     pieces = []
     if left_pad:
-        pieces.append(mx.zeros((left_pad,), dtype=original_dtype))
+        pieces.append(raw[1 : left_pad + 1][::-1])
     pieces.append(raw.astype(original_dtype))
     if right_pad:
-        pieces.append(mx.zeros((right_pad,), dtype=original_dtype))
+        pieces.append(raw[-(right_pad + 1) : -1][::-1])
     segment = mx.concatenate(pieces, axis=0) if len(pieces) > 1 else pieces[0]
 
     expected_len = (num_frames - 1) * hop + n_fft
@@ -170,3 +171,111 @@ def iter_log_mel_spectrogram(
     for frame_start in range(0, total_frames, chunk_frames):
         frame_end = min(frame_start + chunk_frames, total_frames)
         yield log_mel_spectrogram_frames(x, args, frame_start, frame_end)
+
+
+class StreamingLogMelSpectrogram:
+    """Incremental centered log-mel frontend with bounded sample state.
+
+    ``push`` emits only newly available hop-aligned frames whose centers trail the
+    input edge by ``lookahead_samples``. The default is the minimum future context
+    needed by the centered STFT. Joining all outputs, including ``flush()``, is
+    equivalent to :func:`log_mel_spectrogram` without retaining the full waveform.
+    """
+
+    def __init__(
+        self,
+        args: PreprocessArgs,
+        *,
+        lookahead_samples: int | None = None,
+    ):
+        if args.pad_to > 0:
+            raise NotImplementedError(
+                "streaming Nemotron mel extraction does not support pad_to > 0"
+            )
+        if args.normalize in ("per_feature", "all_features"):
+            raise NotImplementedError(
+                "streaming Nemotron mel extraction only supports normalize='NA'"
+            )
+        self.args = args
+        self._samples = mx.zeros((0,), dtype=mx.float32)
+        self._buffer_start = 0
+        self._total_samples = 0
+        self._next_frame = 0
+        self._closed = False
+        if lookahead_samples is None:
+            lookahead_samples = args.n_fft // 2
+        if lookahead_samples < args.n_fft // 2:
+            raise ValueError(
+                "lookahead_samples must cover at least half the centered STFT window"
+            )
+        self.lookahead_samples = lookahead_samples
+        # Keep enough audio before the next frame center for the centered STFT,
+        # plus the preceding sample needed by preemphasis.
+        self._lookbehind_frames = math.ceil((args.n_fft // 2 + 1) / args.hop_length)
+
+    @property
+    def total_samples(self) -> int:
+        return self._total_samples
+
+    @property
+    def emitted_frames(self) -> int:
+        return self._next_frame
+
+    @property
+    def buffered_samples(self) -> int:
+        return self._samples.shape[0]
+
+    def push(self, samples: mx.array, *, final: bool = False) -> mx.array:
+        """Append mono PCM and return newly available ``(1, T, features)`` frames."""
+
+        if self._closed:
+            raise RuntimeError("streaming log-mel frontend is closed")
+        samples = mx.array(samples)
+        if samples.ndim != 1:
+            raise ValueError("streaming log-mel input must be mono PCM")
+        if self._samples.shape[0] == 0:
+            self._samples = samples
+        elif samples.shape[0] > 0:
+            if samples.dtype != self._samples.dtype:
+                samples = samples.astype(self._samples.dtype)
+            self._samples = mx.concatenate([self._samples, samples])
+        self._total_samples += samples.shape[0]
+
+        available_center = self._total_samples - self.lookahead_samples
+        frame_end = (
+            available_center // self.args.hop_length + 1 if available_center >= 0 else 0
+        )
+        if final:
+            frame_end = self._total_samples // self.args.hop_length + 1
+            self._closed = True
+
+        if frame_end <= self._next_frame:
+            return mx.zeros((1, 0, self.args.features), dtype=self._samples.dtype)
+
+        buffer_start_frame = self._buffer_start // self.args.hop_length
+        local_start = self._next_frame - buffer_start_frame
+        local_end = frame_end - buffer_start_frame
+        result = log_mel_spectrogram_frames(
+            self._samples,
+            self.args,
+            local_start,
+            local_end,
+        )
+        self._next_frame = frame_end
+
+        keep_frame = max(self._next_frame - self._lookbehind_frames, 0)
+        keep_sample = keep_frame * self.args.hop_length
+        trim = keep_sample - self._buffer_start
+        if trim > 0:
+            retained = self._samples[trim:]
+            # Materialize the small retained tail so lazy concatenate/slice graphs
+            # cannot keep prior waveform chunks alive across a long stream.
+            mx.eval(retained)
+            self._samples = retained
+            self._buffer_start = keep_sample
+        return result
+
+    def flush(self) -> mx.array:
+        """Emit the final centered frame and close the frontend."""
+
+        return self.push(mx.zeros((0,), dtype=self._samples.dtype), final=True)

mlx_audio/stt/models/nemotron_asr/streaming.py

@@ -4,8 +4,9 @@
 frames) and a causal-conv cache (last ``conv_kernel-1`` GLU-output frames);
 subsampling is incremental with a small mel cache. With the window sized to the
 allowed left context, no attention mask is needed, so the streamed encoder output is
-frame-identical to the offline ``chunked_limited`` encoder at the native chunk size
-(``right_context + 1``). This yields the model's native O(n), no-recompute streaming.
+numerically close to the offline ``chunked_limited`` encoder at the native chunk
+size (``right_context + 1``), and produces the same greedy tokens in parity tests.
+This yields the model's native O(n), no-recompute streaming.
 """
 
 import mlx.core as mx
@@ -42,84 +43,134 @@ def _stream_block(block, x, pos_enc, attn_cache, conv_cache, left_cache, conv_le
     return block.norm_out(residual), attn_next, conv_next
 
 
-def stream_encode_chunks(
-    model, mel_chunks, language, chunk_frames=None, att_context_size=None
-):
-    """Yield post-prompt encoder frames from one or more mel chunks.
+class ConformerStreamingState:
+    """Reusable incremental state for a causal Nemotron FastConformer.
 
-    The encoder/conv/subsampling caches persist across input mel chunks, so callers
-    can keep STFT memory bounded without resetting model context at chunk boundaries.
+    The state owns the subsampling, attention, and causal-convolution caches. A
+    caller may push arbitrary mel chunk sizes; complete native encoder chunks are
+    returned as a list of ``(B, T, D)`` arrays.
     """
-    enc = model.encoder
-    acs = att_context_size or model.default_att_context_size
-    left_cache = int(acs[0])
-    right = int(acs[1])
-    cf = chunk_frames or (right + 1)
-    sf = enc.args.subsampling_factor
-    chunk_mel = cf * sf
-    conv_left = enc.args.conv_kernel_size - 1
-
-    n = len(enc.layers)
-    attn_cache = [None] * n
-    conv_cache = [None] * n
-    mel_cache = None
-    emitted = 0
-    consumed = 0
-    pending = None
-
-    def append_pending(chunk):
-        nonlocal pending
+
+    def __init__(self, encoder, *, chunk_frames=None, att_context_size=None):
+        self.encoder = encoder
+        acs = att_context_size or encoder.args.att_context_size[0]
+        self.left_cache = int(acs[0])
+        self.right_context = int(acs[1])
+        self.chunk_frames = chunk_frames or (self.right_context + 1)
+        if self.chunk_frames <= 0:
+            raise ValueError("chunk_frames must be positive")
+        self.subsampling_factor = encoder.args.subsampling_factor
+        self.chunk_mel = self.chunk_frames * self.subsampling_factor
+        self.conv_left = encoder.args.conv_kernel_size - 1
+
+        n = len(encoder.layers)
+        self.attn_cache = [None] * n
+        self.conv_cache = [None] * n
+        self.mel_cache = None
+        self.emitted = 0
+        self.consumed = 0
+        self.pending = None
+        self.closed = False
+
+    def _append_pending(self, chunk):
         if chunk.ndim == 2:
             chunk = mx.expand_dims(chunk, 0)
         if chunk.shape[1] == 0:
             return
-        pending = chunk if pending is None else mx.concatenate([pending, chunk], axis=1)
-
-    def encode_mel_chunk(m, is_final):
-        nonlocal mel_cache, emitted, consumed
-        cache_len = 0 if mel_cache is None else mel_cache.shape[1]
-        win = m if mel_cache is None else mx.concatenate([mel_cache, m], axis=1)
+        self.pending = (
+            chunk
+            if self.pending is None
+            else mx.concatenate([self.pending, chunk], axis=1)
+        )
+
+    def _encode_mel_chunk(self, m, include_boundary):
+        cache_len = 0 if self.mel_cache is None else self.mel_cache.shape[1]
+        win = (
+            m if self.mel_cache is None else mx.concatenate([self.mel_cache, m], axis=1)
+        )
         win_len = win.shape[1]
-        sub = enc.pre_encode(win, mx.array([win_len], dtype=mx.int32))[0]  # (1, k, d)
-
-        end = consumed + m.shape[1]
-        base = (consumed - cache_len) // sf
-        lo = emitted - base
-        hi = sub.shape[1] if is_final else (end // sf - base)
-        consumed = end
-        mel_cache = win[:, -_PRE_ENCODE_MEL_CACHE:]
+        sub = self.encoder.pre_encode(win, mx.array([win_len], dtype=mx.int32))[0]
+
+        end = self.consumed + m.shape[1]
+        base = (self.consumed - cache_len) // self.subsampling_factor
+        lo = self.emitted - base
+        hi = (
+            sub.shape[1]
+            if include_boundary
+            else (end // self.subsampling_factor - base)
+        )
+        self.consumed = end
+        self.mel_cache = win[:, -_PRE_ENCODE_MEL_CACHE:]
 
         if hi <= lo:
-            emitted = base + max(lo, hi)
-            return
-        emitted = base + hi
+            self.emitted = base + max(lo, hi)
+            return None
+        self.emitted = base + hi
         h = sub[:, lo:hi]
-        for li, block in enumerate(enc.layers):
-            h, attn_cache[li], conv_cache[li] = _stream_block(
+        for li, block in enumerate(self.encoder.layers):
+            h, self.attn_cache[li], self.conv_cache[li] = _stream_block(
                 block,
                 h,
-                enc.pos_enc,
-                attn_cache[li],
-                conv_cache[li],
-                left_cache,
-                conv_left,
+                self.encoder.pos_enc,
+                self.attn_cache[li],
+                self.conv_cache[li],
+                self.left_cache,
+                self.conv_left,
             )
-        yield model.apply_prompt(h, language)
-
-    def encode_ready(is_final):
-        nonlocal pending
-        while pending is not None and pending.shape[1] > 0:
-            if pending.shape[1] < chunk_mel and not is_final:
+        return h
+
+    def push(self, mel, *, final=False, emit_partial=False):
+        """Push mel frames and return newly encoded chunks.
+
+        ``emit_partial`` is for protocols whose input boundary is already known to
+        align with a valid causal encoder frame. It emits the preencoder's current
+        right-boundary output without closing the reusable state.
+        """
+
+        if self.closed:
+            raise RuntimeError("conformer streaming state is closed")
+        self._append_pending(mel)
+        outputs = []
+        while self.pending is not None and self.pending.shape[1] > 0:
+            if self.pending.shape[1] < self.chunk_mel and not (final or emit_partial):
                 break
 
-            take = min(chunk_mel, pending.shape[1])
-            if is_final and pending.shape[1] <= chunk_mel:
-                take = pending.shape[1]
+            take = min(self.chunk_mel, self.pending.shape[1])
+            if (final or emit_partial) and self.pending.shape[1] <= self.chunk_mel:
+                take = self.pending.shape[1]
+
+            m = self.pending[:, :take]
+            self.pending = self.pending[:, take:]
+            include_boundary = (final or emit_partial) and self.pending.shape[1] == 0
+            encoded = self._encode_mel_chunk(m, include_boundary)
+            if encoded is not None:
+                outputs.append(encoded)
+        if final:
+            self.closed = True
+        return outputs
+
+    def materialize(self, *arrays):
+        """Synchronize outputs and cache slices at an online iteration boundary."""
 
-            m = pending[:, :take]
-            pending = pending[:, take:]
-            is_final_chunk = is_final and pending.shape[1] == 0
-            yield from encode_mel_chunk(m, is_final_chunk)
+        state = [self.mel_cache]
+        state.extend(value for value in self.attn_cache if value is not None)
+        state.extend(value for value in self.conv_cache if value is not None)
+        mx.eval(*arrays, *(value for value in state if value is not None))
+
+
+def stream_encode_chunks(
+    model, mel_chunks, language, chunk_frames=None, att_context_size=None
+):
+    """Yield post-prompt encoder frames from one or more mel chunks.
+
+    The encoder/conv/subsampling caches persist across input mel chunks, so callers
+    can keep STFT memory bounded without resetting model context at chunk boundaries.
+    """
+    state = ConformerStreamingState(
+        model.encoder,
+        chunk_frames=chunk_frames,
+        att_context_size=att_context_size or model.default_att_context_size,
+    )
 
     iterator = iter(mel_chunks)
     try:
@@ -128,19 +179,19 @@ def encode_ready(is_final):
         return
 
     for next_chunk in iterator:
-        append_pending(current)
-        yield from encode_ready(is_final=False)
+        for encoded in state.push(current):
+            yield model.apply_prompt(encoded, language)
         current = next_chunk
 
-    append_pending(current)
-    yield from encode_ready(is_final=True)
+    for encoded in state.push(current, final=True):
+        yield model.apply_prompt(encoded, language)
 
 
 def stream_encode(model, mel, language, chunk_frames=None, att_context_size=None):
     """Yield post-prompt encoder frames (1, c, d) per chunk, cache-aware.
 
-    Frame-identical to ``encoder(...)`` + ``apply_prompt(...)`` at the native chunk
-    size (right_context + 1).
+    Token-equivalent to ``encoder(...)`` + ``apply_prompt(...)`` at the native
+    chunk size (right_context + 1), within normal floating-point kernel drift.
     """
     yield from stream_encode_chunks(
         model,