audio: fix chunk splices, output delay and level policy in the enhancers

CMakeLists.txt

@@ -2032,6 +2032,13 @@ if (ENGINE_BUILD_TESTS)
         COMMAND audio_chunking_test
     )
 
+    add_engine_unittest(enhancement_alignment_test tests/unittests/test_enhancement_alignment.cpp)
+
+    add_test(
+        NAME enhancement_alignment_test
+        COMMAND enhancement_alignment_test
+    )
+
     add_engine_unittest(chinese_normalization_test tests/unittests/test_chinese_normalization.cpp)
 
     add_test(

include/engine/framework/audio/flashsr.h

@@ -17,6 +17,28 @@ struct FlashSrOutput {
     std::vector<float> samples;
 };
 
+// Absolute peak the FlashSR output is never allowed to exceed.
+inline constexpr float kFlashSrPeakCeiling = 0.9990000128746033f;
+
+struct FlashSrOptions {
+    // Restore the input waveform's peak on the output instead of normalising
+    // every file to `peak_ceiling`. With this off a single loud sample sets the
+    // level of the whole file and the input's own level is discarded — a -40 dBFS
+    // whisper and a -3 dBFS shout both come back at -0.009 dBFS. Set false for
+    // bit-exact parity with the upstream reference implementation, which is what
+    // the reference fixtures under tests/unittests/assets were captured with.
+    bool preserve_input_level = true;
+    // Safety limit, not a target. The output is scaled down only if restoring
+    // the input peak would push it above this value.
+    float peak_ceiling = kFlashSrPeakCeiling;
+};
+
+// Gain applied to a FlashSR output whose absolute peak is `output_peak`, given
+// an input whose absolute peak is `input_peak`. Returns 0 for a silent output
+// (silence in, silence out) and never lets the result exceed
+// `options.peak_ceiling`.
+float flashsr_output_gain(float input_peak, float output_peak, const FlashSrOptions & options) noexcept;
+
 class FlashSrModel {
 public:
     static FlashSrModel load_from_directory(const std::filesystem::path & model_dir);
@@ -29,7 +51,9 @@ class FlashSrModel {
     FlashSrModel(const FlashSrModel &) = delete;
     FlashSrModel & operator=(const FlashSrModel &) = delete;
 
-    FlashSrOutput super_resolve_mono_16k(const std::vector<float> & waveform) const;
+    FlashSrOutput super_resolve_mono_16k(
+        const std::vector<float> & waveform,
+        const FlashSrOptions & options = FlashSrOptions{}) const;
 
 private:
     explicit FlashSrModel(std::shared_ptr<FlashSrWeights> weights);

include/engine/framework/audio/rnnoise.h

@@ -38,6 +38,34 @@ struct RnnoiseWaveformOutput {
     std::vector<float> vad;
 };
 
+// Group delay of the RNNoise analysis/synthesis pipeline, in samples at 48 kHz.
+// The synthesis stage inverse-transforms the *previous* frame's spectrum, so an
+// output sample at index n reconstructs the input sample at index n - 960
+// (20 ms). DeepFilterNet2 crops its own 480-sample delay the same way.
+inline constexpr int64_t kRnnoiseOutputDelaySamples = 960;
+
+struct RnnoiseProcessOptions {
+    // Pad the input tail by kRnnoiseOutputDelaySamples and crop the same amount
+    // off the front of the synthesis output, so that output[n] lines up with
+    // input[n] and the last 20 ms of input still reaches the output. Set false
+    // for bit-exact parity with the upstream reference implementation, which
+    // leaves the delay in — that is what the reference fixtures under
+    // tests/unittests/assets were captured with.
+    bool compensate_output_delay = true;
+};
+
+// Frame, padding and crop arithmetic process_mono_48k uses for a given input
+// length. Exposed so the alignment can be exercised without model weights.
+struct RnnoiseAlignmentPlan {
+    int64_t frames = 0;
+    int64_t padded_samples = 0;
+    int64_t crop_offset = 0;
+    int64_t output_samples = 0;
+    int64_t vad_frames = 0;
+};
+
+RnnoiseAlignmentPlan rnnoise_alignment_plan(int64_t input_samples, const RnnoiseProcessOptions & options) noexcept;
+
 class RnnoiseModel {
 public:
     static RnnoiseModel load_from_safetensors(const std::filesystem::path & checkpoint_path);
@@ -57,7 +85,9 @@ class RnnoiseModel {
         const std::vector<float> & features,
         int64_t frames,
         int64_t feature_size) const;
-    RnnoiseWaveformOutput process_mono_48k(const std::vector<float> & waveform) const;
+    RnnoiseWaveformOutput process_mono_48k(
+        const std::vector<float> & waveform,
+        const RnnoiseProcessOptions & options = RnnoiseProcessOptions{}) const;
 
     std::unique_ptr<class RnnoiseStreamingSession> create_streaming_session() const;
 

include/engine/framework/audio/zipenhancer.h

@@ -2,6 +2,7 @@
 
 #include "engine/framework/core/backend.h"
 
+#include <cstdint>
 #include <filesystem>
 #include <memory>
 #include <vector>
@@ -15,6 +16,51 @@ struct ZipEnhancerWaveformOutput {
     std::vector<float> samples;
 };
 
+// Chunk-join and output-length policy for ZipEnhancerModel::denoise_mono_16k.
+struct ZipEnhancerOptions {
+    // Length in samples (at 16 kHz) of the linear crossfade applied where two
+    // consecutive 2 s analysis windows meet. Clamped to the 8000-sample (500 ms)
+    // overlap between windows. 0 restores the legacy hard splice, where the
+    // first half of the overlap comes from the earlier chunk and the second half
+    // from the later one with no fade between them.
+    int64_t chunk_crossfade_samples = 8000;
+    // When true denoise_mono_16k returns exactly as many samples as it was
+    // given. When false the un-segmented path returns floor(n / 100) * 100
+    // samples, which is what the upstream reference implementation emits and
+    // what the reference fixtures under tests/unittests/assets were captured
+    // with.
+    bool match_input_length = true;
+};
+
+// Segmentation geometry denoise_mono_16k uses for a given input length. Exposed
+// so the length and coverage arithmetic can be exercised without model weights.
+struct ZipEnhancerChunkPlan {
+    int64_t window_samples = 0;
+    int64_t stride_samples = 0;
+    int64_t padded_samples = 0;
+    int64_t output_samples = 0;
+    bool segmented = false;
+};
+
+ZipEnhancerChunkPlan zipenhancer_chunk_plan(int64_t input_samples, const ZipEnhancerOptions & options) noexcept;
+
+// Rising crossfade weight at `position` inside an `overlap_samples`-long chunk
+// join, using a linear ramp of `fade_samples` centred in the overlap. The pair
+// (position, overlap_samples - 1 - position) always sums to exactly 1.
+float zipenhancer_chunk_fade_weight(int64_t position, int64_t overlap_samples, int64_t fade_samples) noexcept;
+
+// Weight the chunk starting at `segment_start` contributes to the output sample
+// at `segment_start + offset_in_segment`.
+float zipenhancer_segment_weight(
+    int64_t offset_in_segment,
+    int64_t segment_start,
+    const ZipEnhancerChunkPlan & plan,
+    const ZipEnhancerOptions & options) noexcept;
+
+// Total overlap-add weight every padded output sample receives. Every entry
+// below plan.output_samples must be strictly positive or the join leaves a hole.
+std::vector<float> zipenhancer_chunk_weights(const ZipEnhancerChunkPlan & plan, const ZipEnhancerOptions & options);
+
 class ZipEnhancerModel {
 public:
     static ZipEnhancerModel load_from_directory(const std::filesystem::path & model_dir);
@@ -27,7 +73,9 @@ class ZipEnhancerModel {
     ZipEnhancerModel(const ZipEnhancerModel &) = delete;
     ZipEnhancerModel & operator=(const ZipEnhancerModel &) = delete;
 
-    ZipEnhancerWaveformOutput denoise_mono_16k(const std::vector<float> & waveform) const;
+    ZipEnhancerWaveformOutput denoise_mono_16k(
+        const std::vector<float> & waveform,
+        const ZipEnhancerOptions & options = ZipEnhancerOptions{}) const;
 
 private:
     explicit ZipEnhancerModel(std::shared_ptr<const ZipEnhancerModelState> state);

src/framework/audio/flashsr.cpp

@@ -29,7 +29,6 @@ constexpr int kFlashSrOutputSampleRate = 48000;
 constexpr int kFlashSrChannels = 32;
 constexpr int kFlashSrActivationKernel = 12;
 constexpr int kFlashSrActivationRatio = 2;
-constexpr float kFlashSrOutputScale = 0.9990000128746033f;
 
 struct GgmlContextDeleter {
     void operator()(ggml_context * ctx) const noexcept {
@@ -302,22 +301,32 @@ core::TensorValue resblock(
     return output;
 }
 
-std::vector<float> normalize_output(const std::vector<float> & input) {
-    float max_abs = 0.0f;
-    for (float value : input) {
-        max_abs = std::max(max_abs, std::fabs(value));
-    }
-    if (max_abs <= 0.0f) {
-        throw std::runtime_error("FlashSR output normalization has zero peak");
+namespace {
+
+float absolute_peak(const std::vector<float> & values) noexcept {
+    float peak = 0.0f;
+    for (const float value : values) {
+        peak = std::max(peak, std::fabs(value));
     }
-    std::vector<float> output(input.size());
-    const float scale = kFlashSrOutputScale / max_abs;
-    for (size_t i = 0; i < input.size(); ++i) {
-        output[i] = input[i] * scale;
+    return peak;
+}
+
+// Level policy for the model output. `source` is the waveform the caller handed
+// in, so the input's own level can be restored instead of being discarded.
+std::vector<float> apply_output_level(
+    const std::vector<float> & source,
+    const std::vector<float> & model_output,
+    const FlashSrOptions & options) {
+    const float gain = flashsr_output_gain(absolute_peak(source), absolute_peak(model_output), options);
+    std::vector<float> output(model_output.size());
+    for (size_t i = 0; i < model_output.size(); ++i) {
+        output[i] = model_output[i] * gain;
     }
     return output;
 }
 
+}  // namespace
+
 class FlashSrGraph {
 public:
     FlashSrGraph(const FlashSrWeights & weights, int64_t input_samples)
@@ -402,6 +411,17 @@ class FlashSrGraph {
     ggml_backend_graph_plan_t plan_ = nullptr;
 };
 
+float flashsr_output_gain(float input_peak, float output_peak, const FlashSrOptions & options) noexcept {
+    if (!(output_peak > 0.0f)) {
+        return 0.0f;
+    }
+    const float ceiling = options.peak_ceiling > 0.0f ? options.peak_ceiling : kFlashSrPeakCeiling;
+    const float target = options.preserve_input_level
+        ? std::min(std::max(input_peak, 0.0f), ceiling)
+        : ceiling;
+    return target / output_peak;
+}
+
 FlashSrModel::FlashSrModel() = default;
 FlashSrModel::~FlashSrModel() = default;
 FlashSrModel::FlashSrModel(FlashSrModel &&) noexcept = default;
@@ -447,7 +467,9 @@ FlashSrModel FlashSrModel::load_from_directory(
     return FlashSrModel(std::move(weights));
 }
 
-FlashSrOutput FlashSrModel::super_resolve_mono_16k(const std::vector<float> & waveform) const {
+FlashSrOutput FlashSrModel::super_resolve_mono_16k(
+    const std::vector<float> & waveform,
+    const FlashSrOptions & options) const {
     if (!weights_) {
         throw std::runtime_error("FlashSR model is not loaded");
     }
@@ -463,7 +485,9 @@ FlashSrOutput FlashSrModel::super_resolve_mono_16k(const std::vector<float> & wa
         if (!graph_ || !graph_->matches(original_samples)) {
             graph_ = std::make_unique<FlashSrGraph>(*weights_, original_samples);
         }
-        return FlashSrOutput{kFlashSrOutputSampleRate, normalize_output(graph_->run(waveform))};
+        return FlashSrOutput{
+            kFlashSrOutputSampleRate,
+            apply_output_level(waveform, graph_->run(waveform), options)};
     }
 
     std::vector<float> padded = waveform;
@@ -510,7 +534,7 @@ FlashSrOutput FlashSrModel::super_resolve_mono_16k(const std::vector<float> & wa
         }
         output[i] /= weights[i];
     }
-    return FlashSrOutput{kFlashSrOutputSampleRate, normalize_output(output)};
+    return FlashSrOutput{kFlashSrOutputSampleRate, apply_output_level(waveform, output, options)};
 }
 
 }  // namespace engine::audio

src/framework/audio/rnnoise.cpp

@@ -1402,12 +1402,31 @@ RnnoiseSequenceOutput RnnoiseModel::infer_features(
     return output;
 }
 
-RnnoiseWaveformOutput RnnoiseModel::process_mono_48k(const std::vector<float> & waveform) const {
+RnnoiseAlignmentPlan rnnoise_alignment_plan(int64_t input_samples, const RnnoiseProcessOptions & options) noexcept {
+    static_assert(
+        kRnnoiseOutputDelaySamples == 2 * static_cast<int64_t>(kRnnoiseFrameSize),
+        "RNNoise group delay must stay two analysis frames");
+    RnnoiseAlignmentPlan plan;
+    if (input_samples <= 0) {
+        return plan;
+    }
+    plan.crop_offset = options.compensate_output_delay ? kRnnoiseOutputDelaySamples : 0;
+    plan.output_samples = input_samples;
+    plan.vad_frames = (input_samples + kRnnoiseFrameSize - 1) / kRnnoiseFrameSize;
+    plan.frames = (input_samples + plan.crop_offset + kRnnoiseFrameSize - 1) / kRnnoiseFrameSize;
+    plan.padded_samples = plan.frames * kRnnoiseFrameSize;
+    return plan;
+}
+
+RnnoiseWaveformOutput RnnoiseModel::process_mono_48k(
+    const std::vector<float> & waveform,
+    const RnnoiseProcessOptions & options) const {
     if (waveform.empty()) {
         throw std::runtime_error("RNNoise waveform input is empty");
     }
-    const int64_t frames = (static_cast<int64_t>(waveform.size()) + kRnnoiseFrameSize - 1) / kRnnoiseFrameSize;
-    std::vector<float> padded(static_cast<size_t>(frames * kRnnoiseFrameSize), 0.0f);
+    const auto plan = rnnoise_alignment_plan(static_cast<int64_t>(waveform.size()), options);
+    const int64_t frames = plan.frames;
+    std::vector<float> padded(static_cast<size_t>(plan.padded_samples), 0.0f);
     std::copy(waveform.begin(), waveform.end(), padded.begin());
     std::vector<float> output(padded.size(), 0.0f);
     std::vector<float> vad;
@@ -1448,8 +1467,16 @@ RnnoiseWaveformOutput RnnoiseModel::process_mono_48k(const std::vector<float> &
         }
         offset += chunk_frames;
     }
-    output.resize(waveform.size());
-    return RnnoiseWaveformOutput{kRnnoiseSampleRate, std::move(output), std::move(vad)};
+    if (static_cast<int64_t>(output.size()) < plan.crop_offset + plan.output_samples) {
+        throw std::runtime_error("RNNoise synthesis output is shorter than the delay-compensated crop");
+    }
+    std::vector<float> aligned(
+        output.begin() + static_cast<std::ptrdiff_t>(plan.crop_offset),
+        output.begin() + static_cast<std::ptrdiff_t>(plan.crop_offset + plan.output_samples));
+    if (static_cast<int64_t>(vad.size()) > plan.vad_frames) {
+        vad.resize(static_cast<size_t>(plan.vad_frames));
+    }
+    return RnnoiseWaveformOutput{kRnnoiseSampleRate, std::move(aligned), std::move(vad)};
 }
 
 RnnoiseStreamingSession::RnnoiseStreamingSession(std::shared_ptr<const RnnoiseWeights> weights)