Promote Qwen3 audio encoder, S3Gen, PReLU, ESPnet subsampling, and probability segmentation to the framework

CMakeLists.txt

@@ -530,6 +530,7 @@ add_library(engine_core OBJECT
     src/framework/modules/transformers/causal_decoder_runtime.cpp
     src/framework/modules/transformers/qwen35_decoder_runtime.cpp
     src/framework/modules/text_encoders/qwen3_vl_encoder_runtime.cpp
+    src/framework/modules/speech_encoders/qwen3_audio_encoder_runtime.cpp
     src/framework/modules/attention/cross_attention.cpp
     src/framework/modules/transformers/transformer_blocks.cpp
     src/framework/modules/conditioning_modules.cpp
@@ -552,6 +553,7 @@ add_library(engine_core OBJECT
     src/framework/audio/fft.cpp
     src/framework/audio/kaldi_fbank.cpp
     src/framework/audio/activity.cpp
+    src/framework/audio/probability_segmenter.cpp
     src/framework/audio/activity_graph.cpp
     src/framework/audio/chunking.cpp
     src/framework/audio/conversion.cpp
@@ -584,6 +586,7 @@ add_library(engine_core OBJECT
     src/framework/codecs/oobleck_audio_vae_runtime.cpp
     src/framework/codecs/redae_codec_runtime.cpp
     src/framework/codecs/snac_decoder_runtime.cpp
+    src/framework/codecs/s3gen_runtime.cpp
     src/framework/conditioners/clap_audio_conditioner_runtime.cpp
     src/framework/conditioners/cav_mae_st_conditioner_runtime.cpp
     src/framework/conditioners/musicgen_style_conditioner_runtime.cpp
@@ -1792,7 +1795,6 @@ audiocpp_add_model(qwen3_asr
         src/models/qwen3_asr/assets.cpp
         src/models/qwen3_asr/tokenizer_text.cpp
         src/models/qwen3_asr/frontend_whisper.cpp
-        src/models/qwen3_asr/audio_encoder.cpp
         src/models/qwen3_asr/thinker.cpp
         src/models/qwen3_asr/prompt_asr.cpp
         src/models/qwen3_asr/postprocess.cpp
@@ -2139,11 +2141,8 @@ audiocpp_add_model(chatterbox
         src/models/chatterbox/audio_processing.cpp
         src/models/chatterbox/component_weights.cpp
         src/models/chatterbox/conditionals.cpp
-        src/models/chatterbox/hift_vocoder_impl.cpp
         src/models/chatterbox/loader.cpp
         src/models/chatterbox/s3_tokenizer_impl.cpp
-        src/models/chatterbox/s3gen_flow.cpp
-        src/models/chatterbox/s3gen_inference.cpp
         src/models/chatterbox/session.cpp
         src/models/chatterbox/t3_component.cpp
         src/models/chatterbox/t3_weights.cpp
@@ -2171,8 +2170,6 @@ audiocpp_add_model(chatterbox_turbo
         engine/community_models/chatterbox_turbo/loader.h
     LOADERS
         engine::community_models::chatterbox_turbo::make_chatterbox_turbo_loader
-    DEPENDS
-        chatterbox
 )
 
 audiocpp_add_model(ace_step

docs/community_models/chatterbox_turbo.md

@@ -26,8 +26,8 @@ repacks it offline into one self-contained, audio.cpp-native GGUF:
 
 - T3 and built-in-conditional tensors are moved from the upstream flat `t3.`/`conds.` dot
   namespace into this project's own `/`-delimited packed-GGUF namespace convention.
-- S3Gen tensors are renamed back to the exact names base Chatterbox's own S3Gen flow/HiFT-vocoder
-  loader (`src/models/chatterbox/s3gen_flow.cpp`,
+- S3Gen tensors are renamed back to the exact names the shared S3Gen flow/HiFT-vocoder
+  loader (`src/framework/codecs/s3gen_runtime.cpp`,
   `src/framework/modules/vocoders/hift_vocoder.cpp`) already expects, so that loader runs
   completely unmodified for Turbo — no tensor-name translation code exists in this family at
   runtime.

docs/reports/framework_interface_refactor.md

@@ -0,0 +1,16 @@
+# PReLU, Qwen3 Audio Encoder, and S3Gen Interface Refactor
+
+| CUDA longform case | Before (`4b30d92d`) warm time (s) | After (`4a8eb721`) warm time (s) | Before RTF | After RTF | Before peak VRAM (MiB) | After peak VRAM (MiB) | A/B output parity |
+| --- | ---: | ---: | ---: | ---: | ---: | ---: | --- |
+| Chatterbox | 50.041 | 49.791 | 0.12085 | 0.12024 | 5160 | 5126 | Exact |
+| Supertonic | 1.615 | 1.624 | 0.00426 | 0.00428 | 1152 | 1152 | Exact |
+| Chatterbox Turbo | 32.242 | 31.914 | 0.09904 | 0.09803 | 3504 | 3352 | Exact |
+| KittenTTS2 | 31.391 | 31.252 | 0.09398 | 0.09357 | 6832 | 6792 | Exact |
+| TF-GridNet | 25.053 | 24.982 | 0.14972 | 0.14929 | 850 | 850 | Exact |
+| MossFormer2 | 3.886 | 3.832 | 0.02322 | 0.02290 | 1058 | 1058 | Exact |
+| RE-USE | 5.345 | 5.318 | 0.03194 | 0.03178 | 3008 | 3008 | Exact |
+| Qwen3-ASR | 3.451 | 3.437 | 0.01053 | 0.01049 | 3726 | 3726 | Exact |
+| Audio8 ASR | 0.947 | 0.949 | 0.00289 | 0.00290 | 2024 | 2024 | Exact |
+| Qwen3-ASR + Forced Aligner | 5.165 | 5.164 | 0.01577 | 0.01576 | 5332 | 5332 | Exact |
+| Index-Echo | 22.744 | 22.606 | 0.06943 | 0.06900 | 6274 | 6274 | Exact |
+| HeartMuLa (memory saver) | 32.851 | 32.723 | 0.27358 | 0.27251 | 20448 | 20448 | Exact |

include/engine/community_models/audio8_asr/session.h

@@ -5,7 +5,7 @@
 #include "engine/community_models/audio8_asr/assets.h"
 #include "engine/community_models/audio8_asr/projector.h"
 #include "engine/community_models/audio8_asr/thinker.h"
-#include "engine/models/qwen3_asr/audio_encoder.h"
+#include "engine/framework/modules/speech_encoders/qwen3_audio_encoder_runtime.h"
 #include "engine/models/qwen3_asr/frontend_whisper.h"
 
 #include <memory>
@@ -36,7 +36,7 @@ class Audio8ASRSession final
     runtime::TaskSpec task_;
     std::shared_ptr<const Audio8ASRAssets> assets_;
     qwen3_asr::Qwen3ASRWhisperFrontend frontend_;
-    qwen3_asr::Qwen3ASRAudioEncoderRuntime audio_encoder_;
+    modules::Qwen3AudioEncoderRuntime audio_encoder_;
     Audio8ProjectorRuntime projector_;
     Audio8Qwen2ThinkerRuntime thinker_;
 };

include/engine/community_models/chatterbox_turbo/s3gen_turbo.h

@@ -2,10 +2,7 @@
 
 #include "engine/framework/assets/tensor_source.h"
 #include "engine/framework/core/execution_context.h"
-#include "engine/framework/modules/vocoders/hift_vocoder.h"
-#include "engine/models/chatterbox/components.h"
-#include "engine/models/chatterbox/s3gen_flow.h"
-#include "engine/models/chatterbox/s3gen_inference.h"
+#include "engine/framework/codecs/s3gen_runtime.h"
 
 #include <cstdint>
 #include <memory>
@@ -14,9 +11,7 @@
 namespace engine::community_models::chatterbox_turbo {
 
 // Loads and runs Chatterbox Turbo's S3Gen half (speech tokens -> waveform) by delegating to the
-// existing chatterbox family's flow encoder/decoder and HiFT vocoder code (architecturally
-// identical apart from the meanflow decoder branch already added to
-// engine::models::chatterbox::S3FlowDecoderWeights). The native, repacked GGUF (see
+// shared framework S3Gen runtime's meanflow decoder and HiFT vocoder. The native, repacked GGUF (see
 // tools/community_models/chatterbox_turbo/repack_chatterbox_turbo_gguf.py) stores S3Gen tensors
 // under the same names those loaders already expect, so no name translation is needed here.
 //
@@ -31,18 +26,14 @@ class ChatterboxTurboS3Gen {
         engine::assets::TensorStorageType weight_storage_type = engine::assets::TensorStorageType::Native);
 
     // speech_tokens: T3 output (S3 codebook ids, < 6561; caller strips control tokens).
-    engine::models::chatterbox::S3GenInferenceOutputs synthesize(
-        const engine::models::chatterbox::EmbedReferenceOutputs & ref_dict,
+    engine::codecs::s3gen::S3GenInferenceOutputs synthesize(
+        const engine::codecs::s3gen::S3GenConditioning & ref_dict,
         const std::vector<int32_t> & speech_tokens,
         uint64_t flow_seed,
         uint64_t vocoder_seed) const;
 
 private:
-    std::shared_ptr<const engine::models::chatterbox::S3FlowEncoderWeights> encoder_weights_;
-    std::shared_ptr<const engine::models::chatterbox::S3FlowDecoderWeights> decoder_weights_;
-    std::shared_ptr<engine::modules::HiftVocoderComponent> vocoder_;
-    mutable engine::models::chatterbox::S3GenSessionCache cache_;
-    const engine::core::ExecutionContext * execution_context_ = nullptr;
+    std::unique_ptr<engine::codecs::s3gen::S3GenRuntime> runtime_;
 };
 
 }  // namespace engine::community_models::chatterbox_turbo

include/engine/community_models/chatterbox_turbo/tts.h

@@ -2,7 +2,7 @@
 
 #include "engine/framework/core/execution_context.h"
 #include "engine/framework/tokenizers/llama_bpe.h"
-#include "engine/models/chatterbox/s3gen_inference.h"
+#include "engine/framework/codecs/s3gen_runtime.h"
 #include "engine/community_models/chatterbox_turbo/assets.h"
 #include "engine/community_models/chatterbox_turbo/s3gen_turbo.h"
 #include "engine/community_models/chatterbox_turbo/t3_turbo_component.h"
@@ -34,7 +34,7 @@ class ChatterboxTurboTTSComponent {
         std::shared_ptr<const ChatterboxTurboAssets> assets,
         const engine::core::ExecutionContext & execution_context);
 
-    engine::models::chatterbox::S3GenInferenceOutputs generate(
+    engine::codecs::s3gen::S3GenInferenceOutputs generate(
         const std::string & text,
         const ChatterboxTurboGenerateConfig & config) const;
 
@@ -45,7 +45,7 @@ class ChatterboxTurboTTSComponent {
     std::shared_ptr<ChatterboxTurboS3Gen> s3gen_;
     std::vector<float> builtin_speaker_embedding_;
     std::vector<int32_t> builtin_cond_prompt_speech_tokens_;
-    engine::models::chatterbox::EmbedReferenceOutputs builtin_ref_dict_;
+    engine::codecs::s3gen::S3GenConditioning builtin_ref_dict_;
 };
 
 }  // namespace engine::community_models::chatterbox_turbo

include/engine/framework/audio/probability_segmenter.h

@@ -0,0 +1,37 @@
+#pragma once
+
+#include <cstdint>
+#include <optional>
+
+namespace engine::audio {
+
+struct ProbabilitySegment {
+    int64_t start = 0;
+    int64_t end = 0;
+    float confidence = 0.0f;
+};
+
+struct ProbabilitySegmenterConfig {
+    float threshold = 0.5f;
+    int64_t min_speech = 0;
+    int64_t min_silence = 0;
+};
+
+// Positions use caller-supplied units. Overlapping windows end at the last
+// positive window's end, not at the first negative window's start.
+class ProbabilitySegmenter {
+public:
+    explicit ProbabilitySegmenter(ProbabilitySegmenterConfig config);
+    std::optional<ProbabilitySegment> push(float probability, int64_t start, int64_t end);
+    std::optional<ProbabilitySegment> finish();
+
+private:
+    ProbabilitySegmenterConfig config_;
+    bool active_ = false;
+    int64_t start_ = 0;
+    int64_t end_ = 0;
+    double confidence_sum_ = 0.0;
+    int64_t confidence_count_ = 0;
+};
+
+}  // namespace engine::audio

include/engine/framework/codecs/s3gen_runtime.h

@@ -0,0 +1,165 @@
+#pragma once
+
+#include "engine/framework/assets/tensor_source.h"
+#include "engine/framework/core/execution_context.h"
+#include "engine/framework/modules/vocoders/hift_vocoder.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace engine::codecs::s3gen {
+
+struct S3GenConditioning {
+    std::vector<int32_t> prompt_tokens;
+    int64_t prompt_token_count = 0;
+    std::vector<float> prompt_feat;
+    int64_t prompt_feat_frames = 0;
+    int64_t prompt_feat_dims = 0;
+    std::vector<float> embedding;
+    int64_t embedding_size = 0;
+};
+
+struct S3Token2MelOutputs {
+    std::vector<float> mel;
+    int64_t channels = 0;
+    int64_t frames = 0;
+};
+
+struct S3FlowEncoderOutputs {
+    std::vector<float> hidden;
+    int64_t frames = 0;
+    int64_t storage_frames = 0;
+    int64_t hidden_size = 0;
+};
+
+struct S3GenInferenceOutputs {
+    std::vector<float> waveform;
+    int64_t samples = 0;
+    std::vector<float> mel;
+    int64_t mel_channels = 0;
+    int64_t mel_frames = 0;
+};
+
+struct S3FlowDecoderRunTiming {
+    int64_t calls = 0;
+    double conditioning_write_ms = 0.0;
+    double time_embedding_ms = 0.0;
+    double input_write_ms = 0.0;
+    double time_write_ms = 0.0;
+    double graph_compute_ms = 0.0;
+    double output_read_ms = 0.0;
+};
+
+struct S3FlowCFMTimingBreakdown {
+    int64_t steps = 0;
+    int64_t decoder_calls = 0;
+    double initial_state_ms = 0.0;
+    double schedule_ms = 0.0;
+    double zero_conditioning_ms = 0.0;
+    double runner_setup_ms = 0.0;
+    double host_update_ms = 0.0;
+    S3FlowDecoderRunTiming conditioned;
+    S3FlowDecoderRunTiming unconditioned;
+};
+
+struct S3GenTimingBreakdown {
+    double token2mel_ms = 0.0;
+    double token2mel_embed_ms = 0.0;
+    double token2mel_encoder_ms = 0.0;
+    double token2mel_mu_ms = 0.0;
+    double token2mel_cfm_ms = 0.0;
+    S3FlowCFMTimingBreakdown token2mel_cfm;
+    double vocoder_ms = 0.0;
+};
+
+enum class S3FlowVariant { Conditional, MeanFlow };
+
+struct S3FlowEncoderConfig {
+    int64_t vocabulary_size = 6561;
+    int64_t hidden_size = 512;
+    int64_t heads = 8;
+    int64_t feed_forward_size = 2048;
+    int64_t layers = 6;
+    int64_t upsample_layers = 4;
+    int upsample_factor = 2;
+    int64_t prelook_kernel_size = 4;
+    int64_t prelook_output_kernel_size = 3;
+    int64_t upsample_kernel_size = 5;
+    int64_t speaker_embedding_size = 192;
+};
+
+struct S3FlowDecoderConfig {
+    // Single-scale estimator: one entry block, repeated middle blocks, one exit block.
+    S3FlowVariant variant = S3FlowVariant::Conditional;
+    int64_t hidden_size = 256;
+    int64_t heads = 8;
+    int64_t head_dim = 64;
+    int64_t feed_forward_size = 1024;
+    int64_t time_embedding_size = 320;
+    int64_t time_hidden_size = 1024;
+    int64_t middle_blocks = 12;
+    int64_t attention_layers = 4;
+};
+
+struct S3GenConfig {
+    S3GenConfig();
+    int64_t mel_channels = 80;
+    S3FlowEncoderConfig encoder;
+    S3FlowDecoderConfig decoder;
+    modules::HiftVocoderConfig vocoder;
+};
+
+struct S3GenWeightBinding {
+    std::string flow_prefix = "flow.";
+    std::string vocoder_prefix = "mel2wav.";
+    modules::HiftVocoderWeightLayout vocoder_layout =
+        modules::HiftVocoderWeightLayout::TorchParametrizedWeightNorm;
+};
+
+struct S3GenRuntimeOptions {
+    assets::TensorStorageType weight_storage_type = assets::TensorStorageType::Native;
+};
+
+class S3GenRuntime {
+public:
+    S3GenRuntime(std::shared_ptr<const assets::TensorSource> source,
+                 const core::ExecutionContext & execution, S3GenConfig config,
+                 S3GenRuntimeOptions options = {}, S3GenWeightBinding binding = {});
+    ~S3GenRuntime();
+    S3GenRuntime(S3GenRuntime &&) noexcept;
+    S3GenRuntime & operator=(S3GenRuntime &&) noexcept;
+    S3GenRuntime(const S3GenRuntime &) = delete;
+    S3GenRuntime & operator=(const S3GenRuntime &) = delete;
+
+    bool is_meanflow() const;
+    S3FlowEncoderOutputs encode(
+        const std::vector<float> & embeddings, int64_t frames,
+        int64_t capacity_frames, int64_t hidden_size) const;
+    S3Token2MelOutputs token_to_mel(
+        const S3GenConditioning & reference,
+        const std::vector<int32_t> & speech_tokens, int64_t speech_token_count,
+        int64_t num_steps = 10, float cfg_rate = 0.7f, bool cosine_schedule = true,
+        const std::vector<float> & full_noise = {}, uint64_t flow_seed = 0,
+        S3GenTimingBreakdown * timing = nullptr) const;
+    modules::HiftVocoderOutput decode_waveform(
+        const std::vector<float> & mel, int64_t frames,
+        uint64_t seed = 0, uint64_t prior_noise_values = 0) const;
+    S3GenInferenceOutputs synthesize(
+        const S3GenConditioning & reference,
+        const std::vector<int32_t> & speech_tokens, int64_t speech_token_count,
+        int64_t num_steps = 10, float cfg_rate = 0.7f, bool cosine_schedule = true,
+        const std::vector<float> & full_noise = {}, uint64_t flow_seed = 0,
+        uint64_t vocoder_seed = 0, S3GenTimingBreakdown * timing = nullptr) const;
+
+    const modules::HiftVocoderComponent & vocoder() const;
+    void release_flow_graphs() const;
+    void release_vocoder_graphs() const;
+
+private:
+    struct Impl;
+    std::unique_ptr<Impl> impl_;
+};
+
+}  // namespace engine::codecs::s3gen

include/engine/framework/modules/activation_modules.h

@@ -27,6 +27,24 @@ struct LeakyReluConfig {
     float negative_slope = 0.01F;
 };
 
+struct PReluConfig {
+    size_t channel_axis = 1;
+};
+
+class PReluModule {
+public:
+    explicit PReluModule(PReluConfig config = {});
+    const core::ModuleSchema & schema() const noexcept;
+    core::TensorValue build(
+        core::ModuleBuildContext & ctx,
+        const core::TensorValue & input,
+        const core::TensorValue & slope) const;
+    static const core::ModuleSchema & static_schema() noexcept;
+
+private:
+    PReluConfig config_;
+};
+
 class LeakyReluModule {
 public:
     explicit LeakyReluModule(LeakyReluConfig config = {});