diff --git a/mlx_audio/stt/models/nemotron_asr/__init__.py b/mlx_audio/stt/models/nemotron_asr/__init__.py index cde4d992..fa605b1a 100644 --- a/mlx_audio/stt/models/nemotron_asr/__init__.py +++ b/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", +] diff --git a/mlx_audio/stt/models/nemotron_asr/audio.py b/mlx_audio/stt/models/nemotron_asr/audio.py index 3ffc56dd..dcbd36a2 100644 --- a/mlx_audio/stt/models/nemotron_asr/audio.py +++ b/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) diff --git a/mlx_audio/stt/models/nemotron_asr/streaming.py b/mlx_audio/stt/models/nemotron_asr/streaming.py index 83e12548..6e0531b6 100644 --- a/mlx_audio/stt/models/nemotron_asr/streaming.py +++ b/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, diff --git a/mlx_audio/stt/tests/test_nemotron_asr.py b/mlx_audio/stt/tests/test_nemotron_asr.py index e0d4f259..48f3ac30 100644 --- a/mlx_audio/stt/tests/test_nemotron_asr.py +++ b/mlx_audio/stt/tests/test_nemotron_asr.py @@ -13,13 +13,22 @@ import numpy as np import pytest -from mlx_audio.stt.models.nemotron_asr import Model, ModelConfig +from mlx_audio.stt.models.nemotron_asr import ( + ConformerStreamingState, + Model, + ModelConfig, + StreamingLogMelSpectrogram, +) from mlx_audio.stt.models.nemotron_asr import tokenizer as tok from mlx_audio.stt.models.nemotron_asr.audio import ( + _padded_window, + _power_to_log_mel, + _preemphasize, iter_log_mel_spectrogram, log_mel_spectrogram, ) from mlx_audio.stt.models.nemotron_asr.conformer import create_chunked_limited_mask +from mlx_audio.utils import stft def _tiny_config() -> dict: @@ -111,6 +120,89 @@ def test_chunked_log_mel_matches_full(): np.testing.assert_allclose(np.array(chunked), np.array(full), rtol=1e-3, atol=1e-3) +def test_log_mel_uses_nemo_reflect_padding(): + args = ModelConfig.from_dict(_tiny_config()).config.preprocessor + audio = mx.array(np.linspace(-0.5, 0.75, args.n_fft * 2, dtype=np.float32)) + window = _padded_window(args) + emphasized = _preemphasize(audio, args) + + def features(pad_mode): + spectrum = stft( + emphasized, + args.n_fft, + args.hop_length, + args.n_fft, + window, + pad_mode=pad_mode, + ) + power = mx.square(mx.abs(spectrum)).astype(audio.dtype) + return _power_to_log_mel(power, args, audio.dtype) + + actual = log_mel_spectrogram(audio, args) + reflected = features("reflect") + zero_padded = features("constant") + np.testing.assert_allclose(np.array(actual), np.array(reflected), atol=1e-6) + assert not np.allclose( + np.array(actual[:, 0]), + np.array(zero_padded[:, 0]), + rtol=1e-3, + atol=1e-3, + ) + + +def test_streaming_log_mel_matches_full_with_bounded_state(): + args = ModelConfig.from_dict(_tiny_config()).config.preprocessor + audio = mx.array( + (np.random.randn(args.sample_rate * 2 + 123) * 0.1).astype(np.float32) + ) + frontend = StreamingLogMelSpectrogram(args) + chunks = [] + for start in range(0, audio.shape[0], 937): + chunks.append(frontend.push(audio[start : start + 937])) + assert frontend.buffered_samples <= args.n_fft + 937 + chunks.append(frontend.flush()) + + streamed = mx.concatenate(chunks, axis=1) + full = log_mel_spectrogram(audio, args) + np.testing.assert_allclose(np.array(streamed), np.array(full), rtol=1e-3, atol=1e-3) + + +def test_voicechat_style_frontend_and_conformer_state_tracks_bounded_encoder(): + mx.random.seed(0) + rng = np.random.default_rng(0) + config = _tiny_config() + config["encoder"]["att_context_size"] = [[56, 0]] + config["default_att_context_size"] = [56, 0] + model = Model(ModelConfig.from_dict(config)) + mx.eval(model.parameters()) + model.eval() + + args = model.preprocessor_config + frame_samples = 1280 + frontend = StreamingLogMelSpectrogram( + args, + lookahead_samples=frame_samples, + ) + state = ConformerStreamingState(model.encoder, att_context_size=[56, 0]) + audio = mx.zeros((0,), dtype=mx.float32) + + for _ in range(4): + frame = mx.array((rng.standard_normal(frame_samples) * 0.1).astype(np.float32)) + audio = mx.concatenate([audio, frame]) + mel = frontend.push(frame) + outputs = state.push(mel, emit_partial=True) + assert len(outputs) == 1 + assert outputs[0].shape[1] == 1 + + full_mel = log_mel_spectrogram(audio, args) + bounded, _ = model.encoder(full_mel, att_context_size=[56, 0]) + error = np.abs(np.array(outputs[0]) - np.array(bounded[:, -2:-1])) + # Different matmul shapes introduce small floating-point drift, while the + # causal frame and all persistent state remain aligned. + assert float(error.max()) < 0.03 + assert float(error.mean()) < 0.01 + + def test_encoder_and_prompt_shapes(): model = _build_tiny() d_model = model.encoder_config.d_model @@ -155,8 +247,8 @@ def test_stream_generate_runs_and_is_clean(): def test_stream_matches_offline(): - # Cache-aware streaming is frame-identical to the offline chunked_limited - # encoder at the native chunk size, so the greedy decode must be identical. + # Cache-aware streaming tracks the offline chunked_limited encoder closely at + # the native chunk size, so the greedy decode must be identical. model = _build_tiny() sr = model.preprocessor_config.sample_rate audio = mx.array((np.random.randn(int(2.5 * sr)) * 0.1).astype(np.float32))