Source code
Revision control
Copy as Markdown
Other Tools
#pragma once
#include "model_loader.hpp"
#include <cassert>
#include <vector>
namespace pk {
// Shared mel/STFT config + per-frame log-mel math, lifted from the GGUF. Both
// the offline MelFrontend and the streaming StreamingMel build on this so the
// numerics are identical (one source of truth for preemph/STFT/power/mel/log).
struct MelKernel {
explicit MelKernel(const ModelLoader& ml);
// Compute one log-mel frame from n_fft (already-preemphasized, NOT-yet-
// windowed) samples `frame_in` (length n_fft, frame-major time, kept in
// double precision exactly like MelFrontend's `padded`). Writes n_mels
// values into out_col with stride `out_stride` (out_col[m*out_stride]).
// Uses the shared scratch buffers (frame/re/im/power) to avoid per-frame
// allocation. Bit-identical to the inner loop of MelFrontend::compute.
void frame_logmel(const double* frame_in, float* out_col, int out_stride,
std::vector<float>& frame,
std::vector<float>& re, std::vector<float>& im,
std::vector<double>& power) const;
int n_fft_, hop_, n_mels_, n_bins_;
float preemph_, mag_power_, log_guard_;
bool per_feature_;
std::vector<float> window_; // [n_fft] (centered-padded Hann from GGUF)
std::vector<float> fb_; // [n_mels, n_bins] row-major (fb_[m*n_bins + b])
};
class MelFrontend {
public:
explicit MelFrontend(const ModelLoader& ml);
// samples: 16k mono. Out: feats row-major [n_mels, T] (frame-major inner = T).
void compute(const std::vector<float>& samples, std::vector<float>& feats, int& n_mels, int& T) const;
private:
MelKernel k_;
};
// Incremental (frame-local) log-mel for the streaming front end. Produces mel
// frames that are bit-identical (modulo float-add order, which is preserved
// here) to MelFrontend::compute run on the full concatenated PCM, but without
// recomputing the whole buffer on every feed and without retaining O(stream)
// audio.
//
// center=True STFT: frame i is centered at original sample i*hop, covering
// original indices [i*hop - n_fft/2, i*hop + n_fft/2). Frame i is "ready" once
// `real_samples_seen >= i*hop + n_fft/2` (its right edge is real audio). feed()
// emits exactly the frames that became ready with the just-arrived PCM; the
// n_fft/2 start zero-pad of the earliest frames is applied once (conceptually
// prepended). finalize() appends the n_fft/2 end zero-pad and emits the
// remaining tail frames, so the TOTAL frame count and values match
// MelFrontend::compute(full_pcm) exactly (T = 1 + floor(len/hop) for even n_fft).
//
// Preemphasis (y[t] = x[t] - 0.97*x[t-1], y[0] = x[0] for the very first sample
// of the stream) is carried across feeds via a 1-sample raw history.
//
// LIMITATION: this is causal/frame-local and is ONLY equivalent to the full
// buffer when normalize="NA" (no per-feature normalization), which is the case
// for the only streaming model (nvidia/parakeet_realtime_eou_120m-v1). For
// normalize="per_feature" the incremental causal mel is NOT equivalent (it needs
// whole-utterance mean/std), so the constructor asserts NA and callers must fall
// back to MelFrontend::compute for per_feature models.
class StreamingMel {
public:
explicit StreamingMel(const ModelLoader& ml);
// True iff this model's normalization allows incremental (frame-local) mel.
// (false for per_feature, which needs whole-utterance stats.)
bool incremental_ok() const { return !k_.per_feature_; }
int n_mels() const { return k_.n_mels_; }
// Feed `n` new 16 kHz mono samples. Returns the newly-ready mel frames as
// row-major [n_mels, n_new_frames] (feat-major, inner = frame); n_new_frames
// is also written to the out-param. Does NOT apply end zero-padding (that
// would make near-end frames non-causal); call finalize() for the tail.
std::vector<float> feed(const float* pcm, int n, int& n_new_frames);
// Emit the end-padded tail frames (the right n_fft/2 zero-pad), completing
// the stream so the concatenation of all feed() + finalize() frames equals
// MelFrontend::compute on the full PCM. Returns [n_mels, n_tail_frames].
std::vector<float> finalize(int& n_tail_frames);
// Reset to a fresh stream (clears history + counters).
void reset();
private:
// Reorder the frame-contiguous columns from build_frames (n frames, each of
// length n_mels) into the [n_mels, n] feat-major layout (out[m*n + t]).
std::vector<float> to_feat_major(const std::vector<float>& col_major, int n) const;
MelKernel k_;
int pad_; // n_fft / 2
long long samples_seen_ = 0; // total RAW samples fed so far
int emitted_ = 0; // frames already emitted (start-pad aware)
bool have_prev_ = false; // is prev_raw_ valid (i.e. seen >= 1 sample)?
float prev_raw_ = 0.0f; // last RAW sample, for cross-feed preemph
// Retained PREEMPHASIZED samples (double, exactly like MelFrontend's
// `padded`). buf_[k] is the preemph sample at ORIGINAL index
// (buf_origin_ + k). We only keep samples still needed by a future frame,
// so this is bounded by ~n_fft + hop, NOT the stream length.
std::vector<double> buf_;
long long buf_origin_ = 0; // original preemph index of buf_[0]
// Per-frame scratch (reused across frames to avoid allocation).
std::vector<double> framebuf_; // [n_fft] preemph samples for one frame
std::vector<float> wframe_; // [n_fft] windowed (float) frame
std::vector<float> re_, im_;
std::vector<double> power_;
};
} // namespace pk