diff --git a/CMakeLists.txt b/CMakeLists.txt index 3f98f1daa..3b64c5038 100644 --- a/CMakeLists.txt +++ b/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 diff --git a/docs/community_models/chatterbox_turbo.md b/docs/community_models/chatterbox_turbo.md index 8116a132d..170d91c8c 100644 --- a/docs/community_models/chatterbox_turbo.md +++ b/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. diff --git a/docs/reports/framework_interface_refactor.md b/docs/reports/framework_interface_refactor.md new file mode 100644 index 000000000..16d14fb77 --- /dev/null +++ b/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 | diff --git a/include/engine/community_models/audio8_asr/session.h b/include/engine/community_models/audio8_asr/session.h index 63cbd3b15..d84ffa1ef 100644 --- a/include/engine/community_models/audio8_asr/session.h +++ b/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 @@ -36,7 +36,7 @@ class Audio8ASRSession final runtime::TaskSpec task_; std::shared_ptr assets_; qwen3_asr::Qwen3ASRWhisperFrontend frontend_; - qwen3_asr::Qwen3ASRAudioEncoderRuntime audio_encoder_; + modules::Qwen3AudioEncoderRuntime audio_encoder_; Audio8ProjectorRuntime projector_; Audio8Qwen2ThinkerRuntime thinker_; }; diff --git a/include/engine/community_models/chatterbox_turbo/s3gen_turbo.h b/include/engine/community_models/chatterbox_turbo/s3gen_turbo.h index 06b6a72e7..bdfe04705 100644 --- a/include/engine/community_models/chatterbox_turbo/s3gen_turbo.h +++ b/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 #include @@ -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 & speech_tokens, uint64_t flow_seed, uint64_t vocoder_seed) const; private: - std::shared_ptr encoder_weights_; - std::shared_ptr decoder_weights_; - std::shared_ptr vocoder_; - mutable engine::models::chatterbox::S3GenSessionCache cache_; - const engine::core::ExecutionContext * execution_context_ = nullptr; + std::unique_ptr runtime_; }; } // namespace engine::community_models::chatterbox_turbo diff --git a/include/engine/community_models/chatterbox_turbo/tts.h b/include/engine/community_models/chatterbox_turbo/tts.h index ffe7ccfa5..291796e77 100644 --- a/include/engine/community_models/chatterbox_turbo/tts.h +++ b/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 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 s3gen_; std::vector builtin_speaker_embedding_; std::vector 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 diff --git a/include/engine/framework/audio/probability_segmenter.h b/include/engine/framework/audio/probability_segmenter.h new file mode 100644 index 000000000..142c51d7d --- /dev/null +++ b/include/engine/framework/audio/probability_segmenter.h @@ -0,0 +1,37 @@ +#pragma once + +#include +#include + +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 push(float probability, int64_t start, int64_t end); + std::optional 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 diff --git a/include/engine/framework/codecs/s3gen_runtime.h b/include/engine/framework/codecs/s3gen_runtime.h new file mode 100644 index 000000000..9f5f0bedf --- /dev/null +++ b/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 +#include +#include +#include + +namespace engine::codecs::s3gen { + +struct S3GenConditioning { + std::vector prompt_tokens; + int64_t prompt_token_count = 0; + std::vector prompt_feat; + int64_t prompt_feat_frames = 0; + int64_t prompt_feat_dims = 0; + std::vector embedding; + int64_t embedding_size = 0; +}; + +struct S3Token2MelOutputs { + std::vector mel; + int64_t channels = 0; + int64_t frames = 0; +}; + +struct S3FlowEncoderOutputs { + std::vector hidden; + int64_t frames = 0; + int64_t storage_frames = 0; + int64_t hidden_size = 0; +}; + +struct S3GenInferenceOutputs { + std::vector waveform; + int64_t samples = 0; + std::vector 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 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 & embeddings, int64_t frames, + int64_t capacity_frames, int64_t hidden_size) const; + S3Token2MelOutputs token_to_mel( + const S3GenConditioning & reference, + const std::vector & speech_tokens, int64_t speech_token_count, + int64_t num_steps = 10, float cfg_rate = 0.7f, bool cosine_schedule = true, + const std::vector & full_noise = {}, uint64_t flow_seed = 0, + S3GenTimingBreakdown * timing = nullptr) const; + modules::HiftVocoderOutput decode_waveform( + const std::vector & mel, int64_t frames, + uint64_t seed = 0, uint64_t prior_noise_values = 0) const; + S3GenInferenceOutputs synthesize( + const S3GenConditioning & reference, + const std::vector & speech_tokens, int64_t speech_token_count, + int64_t num_steps = 10, float cfg_rate = 0.7f, bool cosine_schedule = true, + const std::vector & 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_; +}; + +} // namespace engine::codecs::s3gen diff --git a/include/engine/framework/modules/activation_modules.h b/include/engine/framework/modules/activation_modules.h index 7d3b77429..7f87cd40d 100644 --- a/include/engine/framework/modules/activation_modules.h +++ b/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 = {}); diff --git a/include/engine/framework/modules/attention/scaled_dot_product_attention.h b/include/engine/framework/modules/attention/scaled_dot_product_attention.h index f6f6991bb..7650ce3aa 100644 --- a/include/engine/framework/modules/attention/scaled_dot_product_attention.h +++ b/include/engine/framework/modules/attention/scaled_dot_product_attention.h @@ -39,7 +39,10 @@ class ScaledDotProductAttentionModule { const core::TensorValue & q_heads, const core::TensorValue & k_heads, const core::TensorValue & v_heads, - const std::optional & attention_mask = std::nullopt) const; + const std::optional & attention_mask = std::nullopt, + // Explicit lowering only: multiply scores by this scalar tensor instead + // of the default 1/sqrt(head_dim) scale, preserving tensor arithmetic. + const std::optional & score_scale = std::nullopt) const; static const core::ModuleSchema & static_schema() noexcept; diff --git a/include/engine/framework/modules/conformer_modules.h b/include/engine/framework/modules/conformer_modules.h index 543bc8af0..94877561e 100644 --- a/include/engine/framework/modules/conformer_modules.h +++ b/include/engine/framework/modules/conformer_modules.h @@ -38,6 +38,32 @@ struct ConvSubsamplingOutputs { core::TensorValue lengths; }; +struct EspnetConv2dSubsampling8Config { + int64_t input_features = 0; + int64_t hidden_size = 0; + int64_t time_row_alignment = 1; +}; + +struct EspnetConv2dSubsampling8Weights { + Conv2dWeights conv0; + Conv2dWeights conv1; + Conv2dWeights conv2; + LinearWeights projection; +}; + +class EspnetConv2dSubsampling8Module { +public: + explicit EspnetConv2dSubsampling8Module(EspnetConv2dSubsampling8Config config); + // Input/output are [batch, time, features]. Three unpadded stride-2 convolutions. + core::TensorValue build( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const EspnetConv2dSubsampling8Weights & weights) const; + +private: + EspnetConv2dSubsampling8Config config_; +}; + struct DepthwiseConvSubsamplingConfig { int64_t input_features = 0; int64_t output_features = 0; diff --git a/include/engine/framework/modules/speech_encoders/qwen3_audio_encoder_runtime.h b/include/engine/framework/modules/speech_encoders/qwen3_audio_encoder_runtime.h new file mode 100644 index 000000000..e9a2cac02 --- /dev/null +++ b/include/engine/framework/modules/speech_encoders/qwen3_audio_encoder_runtime.h @@ -0,0 +1,81 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" + +#include +#include +#include +#include +#include + +namespace engine::modules { + +struct Qwen3AudioEncoderConfig { + int64_t num_mel_bins = 128; + int64_t encoder_layers = 0; + int64_t encoder_attention_heads = 0; + int64_t encoder_ffn_dim = 0; + int64_t d_model = 0; + int64_t max_source_positions = 0; + int64_t n_window = 100; + int64_t n_window_infer = 400; + int64_t conv_chunksize = 500; + int64_t downsample_hidden_size = 0; + int64_t output_dim = 0; + std::string activation_function = "gelu"; +}; + +struct Qwen3AudioFeatures { + std::vector values; + std::vector attention_mask; + int64_t mel_bins = 0; + int64_t frames = 0; + int64_t encoder_tokens = 0; +}; + +struct Qwen3AudioEmbeddings { + std::vector values; + int64_t tokens = 0; + int64_t hidden_size = 0; +}; + +int64_t qwen3_audio_encoder_token_count(int64_t input_frames); + +class Qwen3AudioEncoderGraph; +struct Qwen3AudioEncoderWeights; + +struct Qwen3AudioEncoderWeightBinding { + std::string tower_prefix; + std::string projector_prefix; + std::string projection1 = "proj1"; + std::string projection2 = "proj2"; +}; + +struct Qwen3AudioEncoderRuntimeOptions { + size_t graph_arena_bytes = 0; + assets::TensorStorageType weight_storage_type = assets::TensorStorageType::Native; + std::string trace_name = "qwen3_audio_encoder"; +}; + +class Qwen3AudioEncoderRuntime { +public: + Qwen3AudioEncoderRuntime( + std::shared_ptr source, + Qwen3AudioEncoderConfig config, + core::ExecutionContext & execution, + Qwen3AudioEncoderRuntimeOptions options, + Qwen3AudioEncoderWeightBinding binding); + ~Qwen3AudioEncoderRuntime(); + + Qwen3AudioEmbeddings encode(const Qwen3AudioFeatures & features); + +private: + Qwen3AudioEncoderConfig config_; + std::shared_ptr weights_; + core::ExecutionContext * execution_ = nullptr; + Qwen3AudioEncoderRuntimeOptions options_; + std::unique_ptr graph_; +}; + +} // namespace engine::modules diff --git a/include/engine/framework/modules/transformers/transformer_blocks.h b/include/engine/framework/modules/transformers/transformer_blocks.h index d5b98c398..818728f81 100644 --- a/include/engine/framework/modules/transformers/transformer_blocks.h +++ b/include/engine/framework/modules/transformers/transformer_blocks.h @@ -22,6 +22,8 @@ struct TransformerEncoderBlockConfig { ggml_prec projection_precision = GGML_PREC_DEFAULT; ggml_prec attention_precision = GGML_PREC_DEFAULT; AttentionPrefixCacheLayout prefix_cache_layout = AttentionPrefixCacheLayout::SequenceHeads; + bool use_packed_qkv = false; + bool use_flash_attention = false; }; struct TransformerEncoderBlockWeights { diff --git a/include/engine/models/chatterbox/components.h b/include/engine/models/chatterbox/components.h index 9a223794b..f33bbdfa0 100644 --- a/include/engine/models/chatterbox/components.h +++ b/include/engine/models/chatterbox/components.h @@ -5,6 +5,7 @@ #include "engine/framework/core/execution_context.h" #include "engine/framework/runtime/session.h" #include "engine/models/chatterbox/assets.h" +#include "engine/framework/codecs/s3gen_runtime.h" #include #include @@ -17,6 +18,14 @@ struct HiftVocoderWeights; namespace engine::models::chatterbox { +struct EmbedReferenceOutputs : codecs::s3gen::S3GenConditioning { + double prompt_mel_ms = 0.0; + double speaker_ms = 0.0; + double tokenizer_ms = 0.0; +}; + +void apply_s3_trim_fade(std::vector & waveform, int sample_rate); + namespace components { struct S3TokenizerV2Weights; } // namespace components @@ -31,30 +40,6 @@ struct TokenizerOutputs { int64_t token_count = 0; }; -struct EmbedReferenceOutputs { - std::vector prompt_tokens; - int64_t prompt_token_count = 0; - std::vector prompt_feat; - int64_t prompt_feat_frames = 0; - int64_t prompt_feat_dims = 0; - std::vector embedding; - int64_t embedding_size = 0; - double prompt_mel_ms = 0.0; - double speaker_ms = 0.0; - double tokenizer_ms = 0.0; -}; - -struct HiFTVocoderOutputs { - std::vector waveform; - int64_t samples = 0; - std::vector source; - int64_t source_channels = 0; - int64_t source_frames = 0; - std::vector f0; - int64_t f0_frames = 0; - std::vector post; - int64_t post_frames = 0; -}; struct VoiceEncoderConfig { int64_t sample_rate = 16000; @@ -94,9 +79,6 @@ struct S3TokenizerComponentWeights { std::shared_ptr runtime_weights; }; -struct HiFTVocoderComponentWeights { - std::shared_ptr runtime_weights; -}; std::shared_ptr load_voice_encoder_weights( const engine::assets::TensorSource & source); @@ -178,33 +160,5 @@ class S3TokenizerComponent { std::shared_ptr state_; }; -class HiFTVocoderComponent { -public: - static HiFTVocoderComponent load_from_source( - std::shared_ptr source, - const engine::core::ExecutionContext & execution_context, - engine::assets::TensorStorageType weight_storage_type = engine::assets::TensorStorageType::Native); - - HiFTVocoderComponent( - std::shared_ptr weights, - const engine::core::ExecutionContext & execution_context); - - const engine::core::BackendConfig & backend() const noexcept; - const std::shared_ptr & weights() const noexcept; - HiFTVocoderOutputs infer( - const std::vector & speech_feat, - int64_t batch, - int64_t frames, - uint64_t seed, - uint64_t prior_noise_values, - const std::vector & cache_source) const; - void release_runtime_cache() const; - -private: - struct State; - std::shared_ptr weights_; - const engine::core::ExecutionContext * execution_context_ = nullptr; - std::shared_ptr state_; -}; } // namespace engine::models::chatterbox diff --git a/include/engine/models/chatterbox/conditionals.h b/include/engine/models/chatterbox/conditionals.h index de7c77f8f..3ff19e819 100644 --- a/include/engine/models/chatterbox/conditionals.h +++ b/include/engine/models/chatterbox/conditionals.h @@ -2,7 +2,7 @@ #include "engine/models/chatterbox/components.h" #include "engine/framework/runtime/session.h" -#include "engine/models/chatterbox/s3gen_types.h" +#include "engine/framework/codecs/s3gen_runtime.h" #include #include diff --git a/include/engine/models/chatterbox/s3gen_flow.h b/include/engine/models/chatterbox/s3gen_flow.h deleted file mode 100644 index 0f003051b..000000000 --- a/include/engine/models/chatterbox/s3gen_flow.h +++ /dev/null @@ -1,177 +0,0 @@ -#pragma once - -#include "engine/framework/core/backend.h" -#include "engine/framework/core/execution_context.h" -#include "engine/framework/assets/tensor_source.h" -#include "engine/models/chatterbox/s3gen_types.h" - -#include -#include - -namespace engine::models::chatterbox { - -struct S3FlowDecoderRunTiming; -struct S3FlowCFMTimingBreakdown; - -class S3FlowSessionCache { -public: - explicit S3FlowSessionCache(engine::core::BackendConfig backend = {}); - ~S3FlowSessionCache(); - S3FlowSessionCache(S3FlowSessionCache &&) noexcept; - S3FlowSessionCache & operator=(S3FlowSessionCache &&) noexcept; - - S3FlowSessionCache(const S3FlowSessionCache &) = delete; - S3FlowSessionCache & operator=(const S3FlowSessionCache &) = delete; - - void release_encoder_graphs(); - void release_decoder_graphs(); - -private: - struct State; - std::unique_ptr state_; - - friend S3FlowEncoderOutputs compute_s3_flow_encoder_forward( - S3FlowSessionCache & cache, - const S3FlowEncoderWeights & weights, - const std::vector & input, - int64_t frames, - int64_t capacity_frames, - int64_t hidden_size, - engine::core::BackendConfig backend); - friend S3FlowDecoderOutputs compute_s3_flow_decoder_forward( - S3FlowSessionCache & cache, - const S3FlowDecoderWeights & weights, - const std::vector & x, - const std::vector & mask, - const std::vector & mu, - const std::vector & t, - const std::vector & spks, - const std::vector & cond, - int64_t batch, - int64_t frames, - int64_t capacity_frames, - engine::core::BackendConfig backend, - const std::vector * r); - friend S3FlowCFMOutputs compute_s3_flow_cfm_euler( - S3FlowSessionCache & cache, - const S3FlowDecoderWeights & weights, - const std::vector & noise, - const std::vector & mask, - const std::vector & mu, - const std::vector & spks, - const std::vector & cond, - int64_t batch, - int64_t frames, - int64_t capacity_frames, - int64_t num_steps, - float cfg_rate, - bool cosine_schedule, - engine::core::BackendConfig backend, - S3FlowCFMTimingBreakdown * timing); - friend S3FlowCFMOutputs compute_s3_flow_cfm_meanflow( - S3FlowSessionCache & cache, - const S3FlowDecoderWeights & weights, - const std::vector & noise, - const std::vector & mask, - const std::vector & mu, - const std::vector & spks, - const std::vector & cond, - int64_t batch, - int64_t frames, - int64_t capacity_frames, - int64_t num_steps, - engine::core::BackendConfig backend); -}; - -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; -}; - -std::shared_ptr load_s3_flow_encoder_weights( - const engine::assets::TensorSource & source, - const engine::core::ExecutionContext & execution_context, - engine::assets::TensorStorageType weight_storage_type = engine::assets::TensorStorageType::Native); -S3FlowEncoderOutputs compute_s3_flow_encoder_forward( - S3FlowSessionCache & cache, - const S3FlowEncoderWeights & weights, - const std::vector & input, - int64_t frames, - int64_t capacity_frames, - int64_t hidden_size, - engine::core::BackendConfig backend = {}); -std::shared_ptr load_s3_flow_decoder_weights( - const engine::assets::TensorSource & source, - const engine::core::ExecutionContext & execution_context, - engine::assets::TensorStorageType weight_storage_type = engine::assets::TensorStorageType::Native); -// S3FlowDecoderWeights is intentionally opaque outside this family; this accessor lets callers -// (e.g. chatterbox_turbo) detect whether a loaded decoder is meanflow-distilled without needing -// the struct's internal layout. -bool s3_flow_decoder_is_meanflow(const S3FlowDecoderWeights & weights); -// r: meanflow end-time input (see S3FlowDecoderWeights::meanflow); nullptr for the base -// Chatterbox 10-step CFG decoder, which never reads it. -S3FlowDecoderOutputs compute_s3_flow_decoder_forward( - S3FlowSessionCache & cache, - const S3FlowDecoderWeights & weights, - const std::vector & x, - const std::vector & mask, - const std::vector & mu, - const std::vector & t, - const std::vector & spks, - const std::vector & cond, - int64_t batch, - int64_t frames, - int64_t capacity_frames, - engine::core::BackendConfig backend = {}, - const std::vector * r = nullptr); -S3FlowCFMOutputs compute_s3_flow_cfm_euler( - S3FlowSessionCache & cache, - const S3FlowDecoderWeights & weights, - const std::vector & noise, - const std::vector & mask, - const std::vector & mu, - const std::vector & spks, - const std::vector & cond, - int64_t batch, - int64_t frames, - int64_t capacity_frames, - int64_t num_steps, - float cfg_rate, - bool cosine_schedule, - engine::core::BackendConfig backend = {}, - S3FlowCFMTimingBreakdown * timing = nullptr); - -// Non-CFG, non-batch-doubled Euler solve for meanflow-distilled decoders (Chatterbox Turbo); -// see S3FlowDecoderWeights::meanflow. num_steps defaults to 2 in upstream Python (tts_turbo.py). -S3FlowCFMOutputs compute_s3_flow_cfm_meanflow( - S3FlowSessionCache & cache, - const S3FlowDecoderWeights & weights, - const std::vector & noise, - const std::vector & mask, - const std::vector & mu, - const std::vector & spks, - const std::vector & cond, - int64_t batch, - int64_t frames, - int64_t capacity_frames, - int64_t num_steps = 2, - engine::core::BackendConfig backend = {}); - -} // namespace engine::models::chatterbox diff --git a/include/engine/models/chatterbox/s3gen_inference.h b/include/engine/models/chatterbox/s3gen_inference.h deleted file mode 100644 index 30f37cfc8..000000000 --- a/include/engine/models/chatterbox/s3gen_inference.h +++ /dev/null @@ -1,126 +0,0 @@ -#pragma once - -#include "engine/models/chatterbox/components.h" -#include "engine/framework/core/backend.h" -#include "engine/models/chatterbox/s3gen_flow.h" -#include "engine/models/chatterbox/s3gen_types.h" - -#include -#include - -namespace engine::models::chatterbox { - -struct S3GenTimingBreakdown; - -class S3GenSessionCache { -public: - explicit S3GenSessionCache(engine::core::BackendConfig backend = {}); - ~S3GenSessionCache(); - S3GenSessionCache(S3GenSessionCache &&) noexcept; - S3GenSessionCache & operator=(S3GenSessionCache &&) noexcept; - - S3GenSessionCache(const S3GenSessionCache &) = delete; - S3GenSessionCache & operator=(const S3GenSessionCache &) = delete; - - void release_runtime_graphs(); - -private: - struct State; - std::unique_ptr state_; - - friend S3Token2MelOutputs compute_s3_token2mel_inference( - S3GenSessionCache & cache, - const S3FlowEncoderWeights & encoder_weights, - const S3FlowDecoderWeights & decoder_weights, - const EmbedReferenceOutputs & ref_dict, - const std::vector & speech_tokens, - int64_t speech_token_count, - int64_t num_steps, - float cfg_rate, - bool cosine_schedule, - const std::vector & full_noise, - uint64_t flow_seed, - engine::core::BackendConfig backend, - S3GenTimingBreakdown * timing); - friend S3GenInferenceOutputs compute_s3gen_inference( - S3GenSessionCache & cache, - const S3FlowEncoderWeights & encoder_weights, - const S3FlowDecoderWeights & decoder_weights, - const engine::models::chatterbox::HiFTVocoderComponent & vocoder, - const EmbedReferenceOutputs & ref_dict, - const std::vector & speech_tokens, - int64_t speech_token_count, - int64_t num_steps, - float cfg_rate, - bool cosine_schedule, - const std::vector & full_noise, - uint64_t flow_seed, - uint64_t vocoder_seed, - engine::core::BackendConfig backend, - S3GenTimingBreakdown * timing); - friend S3GenInferenceOutputs compute_s3gen_inference( - S3GenSessionCache & cache, - const S3FlowEncoderWeights & encoder_weights, - const S3FlowDecoderWeights & decoder_weights, - const engine::models::chatterbox::HiFTVocoderComponent & vocoder, - const EmbedReferenceOutputs & ref_dict, - const std::vector & speech_tokens, - int64_t speech_token_count, - engine::core::BackendConfig backend, - S3GenTimingBreakdown * timing); -}; - -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; -}; - -S3Token2MelOutputs compute_s3_token2mel_inference( - S3GenSessionCache & cache, - const S3FlowEncoderWeights & encoder_weights, - const S3FlowDecoderWeights & decoder_weights, - const EmbedReferenceOutputs & ref_dict, - const std::vector & speech_tokens, - int64_t speech_token_count, - int64_t num_steps = 10, - float cfg_rate = 0.7f, - bool cosine_schedule = true, - const std::vector & full_noise = {}, - uint64_t flow_seed = 0, - engine::core::BackendConfig backend = {}, - S3GenTimingBreakdown * timing = nullptr); - -S3GenInferenceOutputs compute_s3gen_inference( - S3GenSessionCache & cache, - const S3FlowEncoderWeights & encoder_weights, - const S3FlowDecoderWeights & decoder_weights, - const engine::models::chatterbox::HiFTVocoderComponent & vocoder, - const EmbedReferenceOutputs & ref_dict, - const std::vector & speech_tokens, - int64_t speech_token_count, - int64_t num_steps = 10, - float cfg_rate = 0.7f, - bool cosine_schedule = true, - const std::vector & full_noise = {}, - uint64_t flow_seed = 0, - uint64_t vocoder_seed = 0, - engine::core::BackendConfig backend = {}, - S3GenTimingBreakdown * timing = nullptr); - -S3GenInferenceOutputs compute_s3gen_inference( - S3GenSessionCache & cache, - const S3FlowEncoderWeights & encoder_weights, - const S3FlowDecoderWeights & decoder_weights, - const engine::models::chatterbox::HiFTVocoderComponent & vocoder, - const EmbedReferenceOutputs & ref_dict, - const std::vector & speech_tokens, - int64_t speech_token_count, - engine::core::BackendConfig backend, - S3GenTimingBreakdown * timing = nullptr); - -} // namespace engine::models::chatterbox diff --git a/include/engine/models/chatterbox/s3gen_types.h b/include/engine/models/chatterbox/s3gen_types.h deleted file mode 100644 index 117515a5f..000000000 --- a/include/engine/models/chatterbox/s3gen_types.h +++ /dev/null @@ -1,52 +0,0 @@ -#pragma once - -#include -#include - -namespace engine::models::chatterbox { - -struct S3FlowEncoderOutputs { - std::vector hidden; - int64_t frames = 0; - int64_t storage_frames = 0; - int64_t hidden_size = 0; -}; - -struct S3Token2MelOutputs { - std::vector mel; - int64_t channels = 0; - int64_t frames = 0; -}; - -struct S3FlowDecoderOutputs { - std::vector mel; - int64_t channels = 0; - int64_t frames = 0; - int64_t storage_frames = 0; -}; - -struct S3FlowCFMOutputs { - std::vector mel; - int64_t channels = 0; - int64_t frames = 0; - int64_t storage_frames = 0; -}; - -struct S3GenInferenceOutputs { - std::vector waveform; - int64_t samples = 0; - std::vector source; - int64_t source_channels = 0; - int64_t source_frames = 0; - std::vector mel; - int64_t mel_channels = 0; - int64_t mel_frames = 0; -}; - -struct S3SpeakerEncoderWeights; -struct S3TokenizerV2Weights; -struct S3FlowEncoderWeights; -struct S3FlowDecoderWeights; -struct S3HiFTVocoderWeights; - -} // namespace engine::models::chatterbox diff --git a/include/engine/models/chatterbox/tts.h b/include/engine/models/chatterbox/tts.h index 02bb4d144..a1449e101 100644 --- a/include/engine/models/chatterbox/tts.h +++ b/include/engine/models/chatterbox/tts.h @@ -5,7 +5,7 @@ #include "engine/framework/core/execution_context.h" #include "engine/framework/runtime/session.h" #include "engine/models/chatterbox/conditionals.h" -#include "engine/models/chatterbox/s3gen_inference.h" +#include "engine/framework/codecs/s3gen_runtime.h" #include "engine/models/chatterbox/t3_component.h" #include "engine/models/chatterbox/text_tokenizer.h" @@ -62,7 +62,7 @@ struct ChatterboxVoiceCloneOutputs { double s3gen_token2mel_encoder_ms = 0.0; double s3gen_token2mel_mu_ms = 0.0; double s3gen_token2mel_cfm_ms = 0.0; - S3FlowCFMTimingBreakdown s3gen_token2mel_cfm_timing; + engine::codecs::s3gen::S3FlowCFMTimingBreakdown s3gen_token2mel_cfm_timing; double s3gen_vocoder_ms = 0.0; int64_t s3gen_token2mel_cuda_memory_used_before_bytes = 0; int64_t s3gen_token2mel_cuda_memory_used_after_bytes = 0; @@ -90,9 +90,7 @@ class ChatterboxTTSComponent { engine::models::chatterbox::VoiceEncoderComponent voice_encoder, engine::models::chatterbox::S3TokenizerComponent tokenizer_component, engine::models::chatterbox::CAMPPlusEncoderComponent speaker_encoder, - std::shared_ptr flow_encoder_weights, - std::shared_ptr flow_decoder_weights, - engine::models::chatterbox::HiFTVocoderComponent vocoder, + std::unique_ptr s3gen, ChatterboxPromptPrepConfig prompt_prep_config, const engine::core::ExecutionContext & execution_context, bool mem_saver = false); @@ -112,8 +110,6 @@ class ChatterboxTTSComponent { const ChatterboxVoiceCloneConfig & config = {}) const; private: - struct State; - ChatterboxVoiceCloneOutputs synthesize_voice_clone_impl( const std::string & text, const ChatterboxConditionalsOutputs & conditionals, @@ -122,12 +118,9 @@ class ChatterboxTTSComponent { T3InferenceComponent t3_; std::shared_ptr tokenizer_; ChatterboxConditionalsComponent conditionals_; - std::shared_ptr flow_encoder_weights_; - std::shared_ptr flow_decoder_weights_; - engine::models::chatterbox::HiFTVocoderComponent vocoder_; + std::unique_ptr s3gen_; const engine::core::ExecutionContext * execution_context_ = nullptr; bool mem_saver_ = false; - std::shared_ptr state_; }; } // namespace engine::models::chatterbox diff --git a/include/engine/models/chatterbox/vc.h b/include/engine/models/chatterbox/vc.h index 749207390..9f42c6ad0 100644 --- a/include/engine/models/chatterbox/vc.h +++ b/include/engine/models/chatterbox/vc.h @@ -4,7 +4,7 @@ #include "engine/framework/runtime/session.h" #include "engine/models/chatterbox/components.h" #include "engine/models/chatterbox/conditionals.h" -#include "engine/models/chatterbox/s3gen_inference.h" +#include "engine/framework/codecs/s3gen_runtime.h" #include #include @@ -29,7 +29,7 @@ struct ChatterboxVoiceConversionOutputs { double s3gen_ms = 0.0; double s3gen_token2mel_ms = 0.0; double s3gen_vocoder_ms = 0.0; - S3GenTimingBreakdown s3gen_timing; + engine::codecs::s3gen::S3GenTimingBreakdown s3gen_timing; }; class ChatterboxVCComponent { @@ -37,11 +37,8 @@ class ChatterboxVCComponent { ChatterboxVCComponent( engine::models::chatterbox::S3TokenizerComponent tokenizer_component, engine::models::chatterbox::CAMPPlusEncoderComponent speaker_encoder, - std::shared_ptr flow_encoder_weights, - std::shared_ptr flow_decoder_weights, - engine::models::chatterbox::HiFTVocoderComponent vocoder, + std::unique_ptr s3gen, ChatterboxPromptPrepConfig prompt_prep_config, - const engine::core::ExecutionContext & execution_context, bool mem_saver = false); ChatterboxVoiceConversionOutputs convert( @@ -50,17 +47,11 @@ class ChatterboxVCComponent { const ChatterboxVoiceConversionConfig & config = {}) const; private: - struct State; - engine::models::chatterbox::S3TokenizerComponent tokenizer_; engine::models::chatterbox::CAMPPlusEncoderComponent speaker_encoder_; - std::shared_ptr flow_encoder_weights_; - std::shared_ptr flow_decoder_weights_; - engine::models::chatterbox::HiFTVocoderComponent vocoder_; + std::unique_ptr s3gen_; ChatterboxPromptPrepConfig prompt_prep_config_; - const engine::core::ExecutionContext * execution_context_ = nullptr; bool mem_saver_ = false; - std::shared_ptr state_; }; runtime::AudioBuffer load_chatterbox_vc_audio_mono( diff --git a/include/engine/models/owsm/model.h b/include/engine/models/owsm/model.h index f85835202..b2be59300 100644 --- a/include/engine/models/owsm/model.h +++ b/include/engine/models/owsm/model.h @@ -6,6 +6,7 @@ #include "engine/framework/core/execution_context.h" #include "engine/framework/modules/attention/types.h" #include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/conformer_modules.h" #include "engine/framework/modules/ebranchformer_modules.h" #include "engine/framework/modules/feed_forward_modules.h" #include "engine/framework/modules/norm_modules.h" @@ -57,16 +58,9 @@ struct OWSMV4Assets { std::string decode_visible(const std::vector & ids) const; }; -struct OWSMV4SubsamplingWeights { - modules::Conv2dWeights conv0; - modules::Conv2dWeights conv1; - modules::Conv2dWeights conv2; - modules::LinearWeights projection; -}; - struct OWSMV4Weights { std::unique_ptr store; - OWSMV4SubsamplingWeights subsampling; + modules::EspnetConv2dSubsampling8Weights subsampling; std::vector encoder; modules::NormWeights encoder_norm; core::TensorValue embedding; diff --git a/include/engine/models/owsm_ctc/model.h b/include/engine/models/owsm_ctc/model.h index 78d64fc04..ec0dc754c 100644 --- a/include/engine/models/owsm_ctc/model.h +++ b/include/engine/models/owsm_ctc/model.h @@ -6,6 +6,7 @@ #include "engine/framework/core/execution_context.h" #include "engine/framework/modules/attention/types.h" #include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/conformer_modules.h" #include "engine/framework/modules/ebranchformer_modules.h" #include "engine/framework/modules/feed_forward_modules.h" #include "engine/framework/modules/norm_modules.h" @@ -53,16 +54,9 @@ struct OWSMCTCV4Assets { std::string decode_visible(const std::vector & ids) const; }; -struct OWSMCTCV4SubsamplingWeights { - modules::Conv2dWeights conv0; - modules::Conv2dWeights conv1; - modules::Conv2dWeights conv2; - modules::LinearWeights projection; -}; - struct OWSMCTCV4Weights { std::unique_ptr store; - OWSMCTCV4SubsamplingWeights subsampling; + modules::EspnetConv2dSubsampling8Weights subsampling; std::vector encoder; modules::NormWeights encoder_norm; core::TensorValue embedding; diff --git a/include/engine/models/qwen3_asr/assets.h b/include/engine/models/qwen3_asr/assets.h index 2ba86525c..8a3f036cf 100644 --- a/include/engine/models/qwen3_asr/assets.h +++ b/include/engine/models/qwen3_asr/assets.h @@ -1,6 +1,7 @@ #pragma once #include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/modules/speech_encoders/qwen3_audio_encoder_runtime.h" #include #include @@ -15,20 +16,7 @@ class TensorSource; namespace engine::models::qwen3_asr { -struct Qwen3ASRAudioEncoderConfig { - int64_t num_mel_bins = 128; - int64_t encoder_layers = 0; - int64_t encoder_attention_heads = 0; - int64_t encoder_ffn_dim = 0; - int64_t d_model = 0; - int64_t max_source_positions = 0; - int64_t n_window = 100; - int64_t n_window_infer = 400; - int64_t conv_chunksize = 500; - int64_t downsample_hidden_size = 0; - int64_t output_dim = 0; - std::string activation_function = "gelu"; -}; +using Qwen3ASRAudioEncoderConfig = modules::Qwen3AudioEncoderConfig; struct Qwen3ASRTextDecoderConfig { int64_t vocab_size = 0; @@ -70,6 +58,7 @@ struct Qwen3ASRConfig { bool tie_word_embeddings = false; Qwen3ASRFrontendConfig frontend; Qwen3ASRAudioEncoderConfig audio_encoder; + modules::Qwen3AudioEncoderWeightBinding audio_encoder_binding; Qwen3ASRTextDecoderConfig text_decoder; std::vector supported_languages; }; diff --git a/include/engine/models/qwen3_asr/audio_encoder.h b/include/engine/models/qwen3_asr/audio_encoder.h deleted file mode 100644 index 1ed784eac..000000000 --- a/include/engine/models/qwen3_asr/audio_encoder.h +++ /dev/null @@ -1,35 +0,0 @@ -#pragma once - -#include "engine/framework/assets/tensor_source.h" -#include "engine/framework/core/execution_context.h" -#include "engine/models/qwen3_asr/assets.h" -#include "engine/models/qwen3_asr/types.h" - -#include -#include - -namespace engine::models::qwen3_asr { - -class Qwen3ASRAudioEncoderGraph; -struct Qwen3ASRAudioEncoderWeights; - -class Qwen3ASRAudioEncoderRuntime { -public: - Qwen3ASRAudioEncoderRuntime( - std::shared_ptr assets, - core::ExecutionContext & execution, - size_t graph_arena_bytes, - assets::TensorStorageType weight_storage_type); - ~Qwen3ASRAudioEncoderRuntime(); - - Qwen3ASRAudioEmbeddings encode(const Qwen3ASRAudioFeatures & features); - -private: - std::shared_ptr assets_; - std::shared_ptr weights_; - core::ExecutionContext * execution_ = nullptr; - size_t graph_arena_bytes_ = 0; - std::unique_ptr graph_; -}; - -} // namespace engine::models::qwen3_asr diff --git a/include/engine/models/qwen3_asr/session.h b/include/engine/models/qwen3_asr/session.h index ccf7f5b89..ada78c5eb 100644 --- a/include/engine/models/qwen3_asr/session.h +++ b/include/engine/models/qwen3_asr/session.h @@ -3,7 +3,7 @@ #include "engine/framework/assets/tensor_source.h" #include "engine/framework/runtime/session_base.h" #include "engine/models/qwen3_asr/assets.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 "engine/models/qwen3_asr/postprocess.h" #include "engine/models/qwen3_asr/prompt_asr.h" @@ -75,7 +75,7 @@ class Qwen3ASRSession final engine::assets::TensorStorageType thinker_weight_storage_type_ = engine::assets::TensorStorageType::Native; Qwen3ASRTextTokenizer tokenizer_; Qwen3ASRWhisperFrontend frontend_; - Qwen3ASRAudioEncoderRuntime audio_encoder_; + modules::Qwen3AudioEncoderRuntime audio_encoder_; Qwen3ASRThinkerRuntime thinker_; Qwen3ASRPromptBuilder prompt_builder_; Qwen3ASRPostprocessor postprocessor_; diff --git a/include/engine/models/qwen3_asr/types.h b/include/engine/models/qwen3_asr/types.h index 032ca2c68..006aa8830 100644 --- a/include/engine/models/qwen3_asr/types.h +++ b/include/engine/models/qwen3_asr/types.h @@ -1,6 +1,7 @@ #pragma once #include "engine/framework/runtime/session.h" +#include "engine/framework/modules/speech_encoders/qwen3_audio_encoder_runtime.h" #include #include @@ -34,41 +35,11 @@ struct Qwen3ASRPrompt { std::vector attention_mask; }; -struct Qwen3ASRAudioFeatures { - std::vector values; - std::vector attention_mask; - int64_t mel_bins = 0; - int64_t frames = 0; - int64_t encoder_tokens = 0; -}; - -struct Qwen3ASRAudioEmbeddings { - std::vector values; - int64_t tokens = 0; - int64_t hidden_size = 0; -}; +using Qwen3ASRAudioFeatures = modules::Qwen3AudioFeatures; +using Qwen3ASRAudioEmbeddings = modules::Qwen3AudioEmbeddings; struct Qwen3ASRGeneratedTokens { std::vector token_ids; }; -inline int64_t qwen3_asr_floor_div(int64_t numerator, int64_t denominator) { - int64_t quotient = numerator / denominator; - const int64_t remainder = numerator % denominator; - if (remainder != 0 && ((remainder < 0) != (denominator < 0))) { - --quotient; - } - return quotient; -} - -inline int64_t qwen3_asr_audio_encoder_token_count(int64_t input_frames) { - if (input_frames <= 0) { - throw std::runtime_error("Qwen3 ASR requires positive feature frame count"); - } - const int64_t input_lengths_leave = input_frames % 100; - const int64_t feat_lengths = qwen3_asr_floor_div(input_lengths_leave - 1, 2) + 1; - return qwen3_asr_floor_div(qwen3_asr_floor_div(feat_lengths - 1, 2) + 1 - 1, 2) + 1 + - (input_frames / 100) * 13; -} - } // namespace engine::models::qwen3_asr diff --git a/include/engine/models/qwen3_forced_aligner/session.h b/include/engine/models/qwen3_forced_aligner/session.h index fdbae2db5..b37e8f428 100644 --- a/include/engine/models/qwen3_forced_aligner/session.h +++ b/include/engine/models/qwen3_forced_aligner/session.h @@ -3,7 +3,7 @@ #include "engine/framework/assets/tensor_source.h" #include "engine/framework/runtime/session_base.h" #include "engine/models/qwen3_asr/assets.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 "engine/models/qwen3_asr/thinker.h" #include "engine/models/qwen3_asr/tokenizer_text.h" @@ -42,7 +42,7 @@ class Qwen3ForcedAlignerSession final engine::assets::TensorStorageType thinker_weight_storage_type_ = engine::assets::TensorStorageType::Native; engine::models::qwen3_asr::Qwen3ASRTextTokenizer tokenizer_; engine::models::qwen3_asr::Qwen3ASRWhisperFrontend frontend_; - engine::models::qwen3_asr::Qwen3ASRAudioEncoderRuntime audio_encoder_; + modules::Qwen3AudioEncoderRuntime audio_encoder_; engine::models::qwen3_asr::Qwen3ASRThinkerRuntime thinker_; Qwen3ForcedAlignProcessor processor_; int64_t run_index_ = 0; diff --git a/src/community_models/audio8_asr/assets.cpp b/src/community_models/audio8_asr/assets.cpp index 003e8a110..ad48f04b3 100644 --- a/src/community_models/audio8_asr/assets.cpp +++ b/src/community_models/audio8_asr/assets.cpp @@ -5,7 +5,6 @@ #include "engine/framework/model_spec/package.h" #include "engine/framework/tokenizers/qwen_bpe_bundle.h" -#include #include #include @@ -13,100 +12,6 @@ namespace engine::community_models::audio8_asr { namespace json = engine::io::json; namespace { -// The Qwen3-ASR encoder/frontend implementations reused by this family address -// tensors with the official Qwen3-ASR prefixes (`model.audio_tower.*`, -// `model.multi_modal_projector.*`). Audio8 checkpoints keep the same encoder -// under `audio_encoder.*`; this source rewrites names in both directions so -// the shared implementations load Audio8 weights unchanged. -class RenamingTensorSource final : public assets::TensorSource { -public: - RenamingTensorSource( - std::shared_ptr inner, - std::vector> mappings) - : inner_(std::move(inner)), - mappings_(std::move(mappings)) { - // Longest `to` prefix first so the exact-inverse projection entries - // win over the catch-all tower mapping during reverse lookups. - reverse_ordered_ = mappings_; - std::stable_sort( - reverse_ordered_.begin(), - reverse_ordered_.end(), - [](const auto & left, const auto & right) { - return left.second.size() > right.second.size(); - }); - } - - const std::filesystem::path & source_path() const noexcept override { - return inner_->source_path(); - } - - bool has_tensor(std::string_view name) const noexcept override { - return inner_->has_tensor(rewrite(name)); - } - - assets::TensorMetadata require_metadata(std::string_view name) const override { - return inner_->require_metadata(rewrite(name)); - } - - std::vector tensors() const override { - auto tensors = inner_->tensors(); - for (auto & tensor : tensors) { - tensor.name = rewrite_back(tensor.name); - } - return tensors; - } - - void release_storage() const override { - inner_->release_storage(); - } - - assets::RawTensorData require_tensor_data(std::string_view name) const override { - return inner_->require_tensor_data(rewrite(name)); - } - - std::vector require_f32( - std::string_view name, - const std::optional> & expected_shape) const override { - return inner_->require_f32(rewrite(name), expected_shape); - } - - std::optional> optional_f32( - std::string_view name, - const std::optional> & expected_shape) const override { - return inner_->optional_f32(rewrite(name), expected_shape); - } - - int64_t require_i64_scalar(std::string_view name) const override { - return inner_->require_i64_scalar(rewrite(name)); - } - -private: - std::string rewrite(std::string_view name) const { - for (const auto & [from, to] : mappings_) { - if (name.rfind(from, 0) == 0) { - return to + std::string(name.substr(from.size())); - } - } - return std::string(name); - } - - std::string rewrite_back(const std::string & name) const { - // Reverse mapping is only used for diagnostics listings. Check the - // longest `to` prefixes first so proj1/proj2 map back to the - // multi-modal projector names instead of the tower prefix. - for (const auto & [from, to] : reverse_ordered_) { - if (name.rfind(to, 0) == 0) { - return from + name.substr(to.size()); - } - } - return name; - } - - std::shared_ptr inner_; - std::vector> mappings_; - std::vector> reverse_ordered_; -}; - qwen3_asr::Qwen3ASRAudioEncoderConfig parse_audio_encoder_config(const json::Value & value) { qwen3_asr::Qwen3ASRAudioEncoderConfig config; config.num_mel_bins = json::require_i64(value, "num_mel_bins"); @@ -232,27 +137,14 @@ std::shared_ptr load_audio8_asr_assets(const std::filesys validate_config(assets->config); const auto weights = assets->resources.open_tensor_source("weights"); - // Encoder view for the shared Qwen3-ASR implementations: expose the - // Audio8 `audio_encoder.*` namespace under the official Qwen3-ASR - // prefixes. The Audio8 encoder carries its input/output projections - // (proj1/proj2) inside its own namespace; the shared code addresses them - // via the multi-modal projector prefix with linear_1/linear_2 names. - auto encoder_source = std::make_shared( - weights, - std::vector>{ - {"model.audio_tower.", "audio_encoder."}, - {"model.multi_modal_projector.linear_1.", "audio_encoder.proj1."}, - {"model.multi_modal_projector.linear_2.", "audio_encoder.proj2."}, - {"model.multi_modal_projector.", "audio_encoder."}}); - auto encoder_assets = std::make_shared(); - encoder_assets->config.hf_transformers_layout = true; + encoder_assets->config.audio_encoder_binding = {"audio_encoder", "audio_encoder"}; encoder_assets->config.frontend.sample_rate = assets->config.frontend.sample_rate; encoder_assets->config.frontend.feature_size = assets->config.frontend.feature_size; encoder_assets->config.frontend.hop_length = assets->config.frontend.hop_length; encoder_assets->config.frontend.n_fft = assets->config.frontend.n_fft; encoder_assets->config.audio_encoder = assets->config.audio_encoder; - encoder_assets->model_weights = std::move(encoder_source); + encoder_assets->model_weights = weights; assets->encoder_assets = std::move(encoder_assets); assets->tokenizer = engine::tokenizers::load_qwen_bpe_tokenizer(assets->resources); diff --git a/src/community_models/audio8_asr/session.cpp b/src/community_models/audio8_asr/session.cpp index d62ff7745..e66d95c37 100644 --- a/src/community_models/audio8_asr/session.cpp +++ b/src/community_models/audio8_asr/session.cpp @@ -200,10 +200,12 @@ Audio8ASRSession::Audio8ASRSession( assets_(require_assets(std::move(assets))), frontend_(assets_->encoder_assets), audio_encoder_( - assets_->encoder_assets, + assets_->encoder_assets->model_weights, + assets_->encoder_assets->config.audio_encoder, execution_context(), - encoder_graph_arena_bytes(RuntimeSessionBase::options()), - parse_encoder_weight_storage(RuntimeSessionBase::options())), + {encoder_graph_arena_bytes(RuntimeSessionBase::options()), + parse_encoder_weight_storage(RuntimeSessionBase::options()), "audio8_asr.audio_encoder"}, + assets_->encoder_assets->config.audio_encoder_binding), projector_( assets_->resources.open_tensor_source("weights"), assets_->config.tower, diff --git a/src/community_models/chatterbox_turbo/s3gen_turbo.cpp b/src/community_models/chatterbox_turbo/s3gen_turbo.cpp index a9ef66071..c4f7bc251 100644 --- a/src/community_models/chatterbox_turbo/s3gen_turbo.cpp +++ b/src/community_models/chatterbox_turbo/s3gen_turbo.cpp @@ -4,80 +4,28 @@ namespace engine::community_models::chatterbox_turbo { -namespace { - -engine::modules::HiftVocoderConfig make_turbo_hift_config(engine::assets::TensorStorageType weight_storage_type) { - // Matches the S3Gen GGUF's `chatterbox.s3gen.*` metadata and base Chatterbox's own HiFT - // config (src/models/chatterbox/hift_vocoder_impl.cpp) -- the vocoder architecture is - // unchanged by the meanflow distillation, only the tensor names differ. The repacked native - // GGUF folds torch weight-norm parametrization into a plain "weight" tensor at conversion - // time (no "parametrizations.weight.original0/1" keys), so this uses the Canonical (plain) - // layout rather than TorchParametrizedWeightNorm. - engine::modules::HiftVocoderConfig config; - config.in_channels = 80; - config.base_channels = 512; - config.nb_harmonics = 8; - config.sampling_rate = 24000; - config.nsf_alpha = 0.1F; - config.nsf_sigma = 0.003F; - config.nsf_voiced_threshold = 10.0F; - config.upsample_rates = {8, 5, 3}; - config.upsample_kernel_sizes = {16, 11, 7}; - config.istft_n_fft = 16; - config.istft_hop = 4; - config.resblock_kernel_sizes = {3, 7, 11}; - config.resblock_dilation_sizes = {{1, 3, 5}, {1, 3, 5}, {1, 3, 5}}; - config.source_resblock_kernel_sizes = {7, 7, 11}; - config.source_resblock_dilation_sizes = {{1, 3, 5}, {1, 3, 5}, {1, 3, 5}}; - config.lrelu_slope = 0.1F; - config.audio_limit = 0.99F; - config.f0_num_class = 1; - config.f0_in_channels = 80; - config.f0_cond_channels = 512; - config.weight_storage_type = weight_storage_type; - // The repacked GGUF stores every vocoder tensor under the top-level "v." namespace (matching - // base Chatterbox's own "mel2wav."-style tensor_prefix convention, just with a different - // literal prefix); HiftVocoderComponent applies config.tensor_prefix + name itself, so no - // separate translation layer is needed here. - config.tensor_prefix = "v."; - config.weight_layout = engine::modules::HiftVocoderWeightLayout::Canonical; - return config; -} - -} // namespace - std::shared_ptr ChatterboxTurboS3Gen::load( std::shared_ptr s3gen_source, const engine::core::ExecutionContext & execution_context, engine::assets::TensorStorageType weight_storage_type) { auto out = std::make_shared(); - out->execution_context_ = &execution_context; - out->encoder_weights_ = - engine::models::chatterbox::load_s3_flow_encoder_weights(*s3gen_source, execution_context, weight_storage_type); - out->decoder_weights_ = - engine::models::chatterbox::load_s3_flow_decoder_weights(*s3gen_source, execution_context, weight_storage_type); - if (!engine::models::chatterbox::s3_flow_decoder_is_meanflow(*out->decoder_weights_)) { - throw std::runtime_error( - "Chatterbox Turbo S3Gen weights are missing the meanflow time_embed_mixer tensor (flow.decoder.estimator.time_embed_mixer)"); - } - out->vocoder_ = std::make_shared( - engine::modules::HiftVocoderComponent::load_from_tensor_source( - s3gen_source, execution_context.config(), make_turbo_hift_config(weight_storage_type))); + engine::codecs::s3gen::S3GenConfig config; + config.decoder.variant = engine::codecs::s3gen::S3FlowVariant::MeanFlow; + out->runtime_ = std::make_unique( + std::move(s3gen_source), execution_context, config, + engine::codecs::s3gen::S3GenRuntimeOptions{weight_storage_type}, + engine::codecs::s3gen::S3GenWeightBinding{"flow.", "v.", engine::modules::HiftVocoderWeightLayout::Canonical}); return out; } -engine::models::chatterbox::S3GenInferenceOutputs ChatterboxTurboS3Gen::synthesize( - const engine::models::chatterbox::EmbedReferenceOutputs & ref_dict, +engine::codecs::s3gen::S3GenInferenceOutputs ChatterboxTurboS3Gen::synthesize( + const engine::codecs::s3gen::S3GenConditioning & ref_dict, const std::vector & speech_tokens, uint64_t flow_seed, uint64_t vocoder_seed) const { // n_cfm_timesteps=2 matches tts_turbo.py's ChatterboxTurboTTS.generate default; cfg_rate and - // cosine_schedule are unused on the meanflow path (see s3gen_inference.cpp's - // decoder_weights.meanflow branch). - const auto mel = engine::models::chatterbox::compute_s3_token2mel_inference( - cache_, - *encoder_weights_, - *decoder_weights_, + // cosine_schedule are unused on the meanflow path. + const auto mel = runtime_->token_to_mel( ref_dict, speech_tokens, static_cast(speech_tokens.size()), @@ -86,12 +34,11 @@ engine::models::chatterbox::S3GenInferenceOutputs ChatterboxTurboS3Gen::synthesi /*cosine_schedule=*/false, /*full_noise=*/{}, flow_seed, - execution_context_->config(), /*timing=*/nullptr); - const auto voc = vocoder_->synthesize(mel.mel, mel.frames, vocoder_seed); + const auto voc = runtime_->decode_waveform(mel.mel, mel.frames, vocoder_seed); - engine::models::chatterbox::S3GenInferenceOutputs outputs; + engine::codecs::s3gen::S3GenInferenceOutputs outputs; outputs.waveform = voc.waveform; outputs.samples = voc.samples; outputs.mel = mel.mel; diff --git a/src/community_models/chatterbox_turbo/tts.cpp b/src/community_models/chatterbox_turbo/tts.cpp index 2aa371f1b..194dbe0ac 100644 --- a/src/community_models/chatterbox_turbo/tts.cpp +++ b/src/community_models/chatterbox_turbo/tts.cpp @@ -64,7 +64,7 @@ ChatterboxTurboTTSComponent::ChatterboxTurboTTSComponent( builtin_ref_dict_.embedding_size = embedding_shape.at(0); } -engine::models::chatterbox::S3GenInferenceOutputs ChatterboxTurboTTSComponent::generate( +engine::codecs::s3gen::S3GenInferenceOutputs ChatterboxTurboTTSComponent::generate( const std::string & text, const ChatterboxTurboGenerateConfig & config) const { T3TurboGenerateRequest request; diff --git a/src/community_models/kitten_tts2/decoder.cpp b/src/community_models/kitten_tts2/decoder.cpp index 3b054287e..01449ab23 100644 --- a/src/community_models/kitten_tts2/decoder.cpp +++ b/src/community_models/kitten_tts2/decoder.cpp @@ -1,26 +1,20 @@ #include "engine/community_models/kitten_tts2/decoder.h" -#include "engine/models/chatterbox/s3gen_inference.h" +#include "engine/framework/codecs/s3gen_runtime.h" +#include "engine/models/chatterbox/components.h" #include #include namespace engine::community_models::kitten_tts2 { -namespace s3 = engine::models::chatterbox; +namespace s3 = engine::codecs::s3gen; struct WaveformDecoder::Impl { - const core::ExecutionContext & execution; - std::shared_ptr encoder; - std::shared_ptr decoder; - s3::HiFTVocoderComponent vocoder; - s3::S3GenSessionCache cache; + s3::S3GenRuntime runtime; Impl(std::shared_ptr source, const core::ExecutionContext & context) - : execution(context), - encoder(s3::load_s3_flow_encoder_weights(*source, context, assets::TensorStorageType::F32)), - decoder(s3::load_s3_flow_decoder_weights(*source, context, assets::TensorStorageType::F32)), - vocoder(s3::HiFTVocoderComponent::load_from_source(source, context, assets::TensorStorageType::F32)), - cache(context.config()) { - if (!s3::s3_flow_decoder_is_meanflow(*decoder)) - throw std::runtime_error("Kitten TTS 2 requires the S3 meanflow decoder checkpoint"); - } + : runtime(std::move(source), context, [] { + s3::S3GenConfig config; + config.decoder.variant = s3::S3FlowVariant::MeanFlow; + return config; + }(), {assets::TensorStorageType::F32}) {} }; WaveformDecoder::WaveformDecoder(std::shared_ptr source, const core::ExecutionContext & execution) : impl_(std::make_unique(std::move(source), execution)) {} @@ -32,7 +26,7 @@ std::vector WaveformDecoder::decode(const std::vector & codes, if (codes.empty()) throw std::runtime_error("Kitten TTS 2 generated no speech tokens"); if (prompt_tokens.empty() || prompt_mel.empty() || prompt_mel.size() % 80 || embedding.size() != 192) throw std::runtime_error("invalid Kitten TTS 2 decoder conditioning"); - s3::EmbedReferenceOutputs ref; + s3::S3GenConditioning ref; ref.prompt_tokens.assign(prompt_tokens.begin(), prompt_tokens.end()); ref.prompt_token_count = ref.prompt_tokens.size(); ref.prompt_feat = prompt_mel; @@ -42,9 +36,9 @@ std::vector WaveformDecoder::decode(const std::vector & codes, ref.embedding_size = 192; auto tokens = codes; tokens.insert(tokens.end(), 3, 4299); // Upstream S3 lookahead preserves the final syllable. - auto result = s3::compute_s3gen_inference(impl_->cache, *impl_->encoder, *impl_->decoder, - impl_->vocoder, ref, tokens, tokens.size(), 2, 0.0f, false, {}, seed, seed, - impl_->execution.config()); + auto result = impl_->runtime.synthesize( + ref, tokens, tokens.size(), 2, 0.0f, false, {}, seed, seed); + models::chatterbox::apply_s3_trim_fade(result.waveform, 24000); if (result.waveform.empty()) throw std::runtime_error("Kitten decoder returned empty audio"); for (float value : result.waveform) if (!std::isfinite(value)) throw std::runtime_error("Kitten decoder returned non-finite audio"); diff --git a/src/community_models/reuse/runtime.cpp b/src/community_models/reuse/runtime.cpp index d6c81535a..3c4d96bb8 100644 --- a/src/community_models/reuse/runtime.cpp +++ b/src/community_models/reuse/runtime.cpp @@ -240,15 +240,8 @@ class ReuseGraph { input, TensorShape::from_dims({input->ne[3], 64, input->ne[1], input->ne[0]})); const auto x = modules::GroupNormModule({64, 64, 1e-5f, true, true}).build(ctx, value, {weights_.at(prefix + ".1.weight"), weights_.at(prefix + ".1.bias")}); - const auto alpha = core::reshape_tensor(ctx, weights_.at(prefix + ".2.weight"), - TensorShape::from_dims({1, 64, 1, 1})); - const auto positive = modules::ReluModule().build(ctx, x); - const auto negative = modules::ReluModule().build(ctx, - core::wrap_tensor(ggml_neg(ctx.ggml, x.tensor), x.shape)); - const auto scaled = modules::MulModule().build(ctx, negative, - modules::RepeatModule({x.shape}).build(ctx, alpha)); - // The framework has no channel-wise PReLU module. - return ggml_sub(ctx.ggml, positive.tensor, scaled.tensor); + return modules::PReluModule({1}) + .build(ctx, x, weights_.at(prefix + ".2.weight")).tensor; } ggml_tensor * conv_norm(ggml_tensor * input, const std::string & prefix, diff --git a/src/framework/audio/probability_segmenter.cpp b/src/framework/audio/probability_segmenter.cpp new file mode 100644 index 000000000..6cd6aeb81 --- /dev/null +++ b/src/framework/audio/probability_segmenter.cpp @@ -0,0 +1,32 @@ +#include "engine/framework/audio/probability_segmenter.h" + +namespace engine::audio { + +ProbabilitySegmenter::ProbabilitySegmenter(ProbabilitySegmenterConfig config) : config_(config) {} + +std::optional ProbabilitySegmenter::push(float probability, int64_t start, int64_t end) { + if (probability >= config_.threshold) { + if (!active_) { + active_ = true; + start_ = start; + confidence_sum_ = 0.0; + confidence_count_ = 0; + } + end_ = end; + confidence_sum_ += probability; + ++confidence_count_; + } else if (active_ && start - end_ >= config_.min_silence) { + return finish(); + } + return std::nullopt; +} + +std::optional ProbabilitySegmenter::finish() { + if (!active_) return std::nullopt; + active_ = false; + if (end_ - start_ < config_.min_speech) return std::nullopt; + return ProbabilitySegment{start_, end_, + static_cast(confidence_sum_ / static_cast(confidence_count_))}; +} + +} // namespace engine::audio diff --git a/src/models/chatterbox/s3gen_flow.cpp b/src/framework/codecs/s3gen_runtime.cpp similarity index 65% rename from src/models/chatterbox/s3gen_flow.cpp rename to src/framework/codecs/s3gen_runtime.cpp index 60f3d9917..371734860 100644 --- a/src/models/chatterbox/s3gen_flow.cpp +++ b/src/framework/codecs/s3gen_runtime.cpp @@ -1,27 +1,316 @@ -#include "engine/models/chatterbox/s3gen_flow.h" -#include "components/s3gen_weights.h" - #include "engine/framework/assets/tensor_source.h" #include "engine/framework/debug/profiler.h" - #include "engine/framework/core/backend.h" #include "engine/framework/core/deferred_tensor_writer.h" #include "engine/framework/core/execution_context.h" #include "engine/framework/core/module.h" #include "engine/framework/modules/conv_modules.h" #include "engine/framework/modules/structural_modules.h" - #include "ggml.h" #include "ggml-alloc.h" - #include #include #include #include #include #include +#include "engine/framework/codecs/s3gen_runtime.h" +#include "engine/framework/runtime/session.h" +#include "engine/framework/sampling/torch_random.h" +#include +#include +#include "engine/framework/core/backend_weight_store.h" +#include + +namespace engine::codecs::s3gen { +namespace { + +struct S3FlowEncoderWeights { + S3FlowEncoderConfig config; + struct LayerNormWeights { + engine::core::TensorValue weight_tensor; + engine::core::TensorValue bias_tensor; + }; + struct LinearWeights { + engine::core::TensorValue weight_tensor; + engine::core::TensorValue bias_tensor; + int64_t out_features = 0; + int64_t in_features = 0; + bool use_bias = false; + }; + struct Conv1dWeights { + engine::core::TensorValue weight_tensor; + engine::core::TensorValue bias_tensor; + int64_t out_channels = 0; + int64_t in_channels = 0; + int64_t kernel = 0; + int64_t stride = 1; + int64_t padding = 0; + }; + struct RelativeAttentionWeights { + LinearWeights q; + LinearWeights k; + LinearWeights v; + LinearWeights out; + LinearWeights pos; + engine::core::TensorValue pos_bias_u_tensor; + engine::core::TensorValue pos_bias_v_tensor; + }; + struct FeedForwardWeights { + LinearWeights w1; + LinearWeights w2; + }; + struct EncoderLayerWeights { + LayerNormWeights norm_mha; + RelativeAttentionWeights attn; + LayerNormWeights norm_ff; + FeedForwardWeights ff; + }; + LinearWeights speaker_affine; + LinearWeights encoder_proj; + LinearWeights embed_linear; + LayerNormWeights embed_norm; + Conv1dWeights prelook_conv1; + Conv1dWeights prelook_conv2; + std::vector encoders; + Conv1dWeights up_layer_conv; + LinearWeights up_embed_linear; + LayerNormWeights up_embed_norm; + std::vector up_encoders; + LayerNormWeights after_norm; + engine::core::TensorValue input_embedding_tensor; + const engine::core::ExecutionContext * execution_context = nullptr; + std::shared_ptr store; +}; + +struct S3FlowDecoderWeights { + S3FlowDecoderConfig config; + struct LayerNormWeights { + engine::core::TensorValue weight_tensor; + engine::core::TensorValue bias_tensor; + }; + struct LinearWeights { + engine::core::TensorValue weight_tensor; + engine::core::TensorValue bias_tensor; + int64_t out_features = 0; + int64_t in_features = 0; + bool use_bias = false; + }; + struct Conv1dWeights { + engine::core::TensorValue weight_tensor; + engine::core::TensorValue bias_tensor; + int64_t out_channels = 0; + int64_t in_channels = 0; + int64_t kernel = 0; + int64_t stride = 1; + bool use_bias = false; + }; + struct CausalBlockWeights { + Conv1dWeights conv; + LayerNormWeights norm; + }; + struct ResnetBlockWeights { + LinearWeights time_mlp; + CausalBlockWeights block1; + CausalBlockWeights block2; + Conv1dWeights res_conv; + }; + struct TransformerBlockWeights { + int64_t heads = 0; + int64_t head_dim = 0; + LayerNormWeights norm1; + LinearWeights attn_q; + LinearWeights attn_k; + LinearWeights attn_v; + LinearWeights attn_out; + LayerNormWeights norm3; + LinearWeights ff_proj_in; + LinearWeights ff_proj_out; + }; + struct DownBlockWeights { + ResnetBlockWeights resnet; + std::vector transformers; + Conv1dWeights downsample; + }; + struct MidBlockWeights { + ResnetBlockWeights resnet; + std::vector transformers; + }; + struct UpBlockWeights { + ResnetBlockWeights resnet; + std::vector transformers; + Conv1dWeights upsample; + }; + + LinearWeights time_mlp_1; + LinearWeights time_mlp_2; + std::vector down_blocks; + std::vector mid_blocks; + std::vector up_blocks; + CausalBlockWeights final_block; + Conv1dWeights final_proj; + // Meanflow-distilled decoders (Chatterbox Turbo) mix a second "r" (end-time) sinusoidal + // embedding into the time embedding via this diagonal-init, no-bias linear layer before + // feeding the UNet1D estimator; see decoder.py::get_intmeanflow_time_mixer upstream. Unset + // (weight_tensor.tensor == nullptr) for the base Chatterbox 10-step CFG decoder. + bool meanflow = false; + LinearWeights time_embed_mixer; + const engine::core::ExecutionContext * execution_context = nullptr; + std::shared_ptr store; +}; + +struct S3FlowDecoderOutputs { + std::vector mel; + int64_t channels = 0; + int64_t frames = 0; + int64_t storage_frames = 0; +}; + +struct S3FlowCFMOutputs { + std::vector mel; + int64_t channels = 0; + int64_t frames = 0; + int64_t storage_frames = 0; +}; + +class S3FlowSessionCache { +public: + explicit S3FlowSessionCache(engine::core::BackendConfig backend = {}); + ~S3FlowSessionCache(); + + S3FlowSessionCache(const S3FlowSessionCache &) = delete; + S3FlowSessionCache & operator=(const S3FlowSessionCache &) = delete; + + void release_encoder_graphs(); + void release_decoder_graphs(); + +private: + struct State; + std::unique_ptr state_; + + friend S3FlowEncoderOutputs compute_s3_flow_encoder_forward( + S3FlowSessionCache & cache, + const S3FlowEncoderWeights & weights, + const std::vector & input, + int64_t frames, + int64_t capacity_frames, + int64_t hidden_size, + engine::core::BackendConfig backend); + friend S3FlowCFMOutputs compute_s3_flow_cfm_euler( + S3FlowSessionCache & cache, + const S3FlowDecoderWeights & weights, + const std::vector & noise, + const std::vector & mask, + const std::vector & mu, + const std::vector & spks, + const std::vector & cond, + int64_t batch, + int64_t frames, + int64_t capacity_frames, + int64_t num_steps, + float cfg_rate, + bool cosine_schedule, + engine::core::BackendConfig backend, + S3FlowCFMTimingBreakdown * timing); + friend S3FlowCFMOutputs compute_s3_flow_cfm_meanflow( + S3FlowSessionCache & cache, + const S3FlowDecoderWeights & weights, + const std::vector & noise, + const std::vector & mask, + const std::vector & mu, + const std::vector & spks, + const std::vector & cond, + int64_t batch, + int64_t frames, + int64_t capacity_frames, + int64_t num_steps, + engine::core::BackendConfig backend); +}; + +std::shared_ptr load_s3_flow_encoder_weights( + const engine::assets::TensorSource & source, + const engine::core::ExecutionContext & execution_context, + const S3GenConfig & config, + const std::string & tensor_prefix, + engine::assets::TensorStorageType weight_storage_type = engine::assets::TensorStorageType::Native); +S3FlowEncoderOutputs compute_s3_flow_encoder_forward( + S3FlowSessionCache & cache, + const S3FlowEncoderWeights & weights, + const std::vector & input, + int64_t frames, + int64_t capacity_frames, + int64_t hidden_size, + engine::core::BackendConfig backend = {}); +std::shared_ptr load_s3_flow_decoder_weights( + const engine::assets::TensorSource & source, + const engine::core::ExecutionContext & execution_context, + const S3GenConfig & config, + const std::string & tensor_prefix, + engine::assets::TensorStorageType weight_storage_type = engine::assets::TensorStorageType::Native); +S3FlowCFMOutputs compute_s3_flow_cfm_euler( + S3FlowSessionCache & cache, + const S3FlowDecoderWeights & weights, + const std::vector & noise, + const std::vector & mask, + const std::vector & mu, + const std::vector & spks, + const std::vector & cond, + int64_t batch, + int64_t frames, + int64_t capacity_frames, + int64_t num_steps, + float cfg_rate, + bool cosine_schedule, + engine::core::BackendConfig backend = {}, + S3FlowCFMTimingBreakdown * timing = nullptr); + +// Non-CFG, non-batch-doubled Euler solve for meanflow-distilled decoders (Chatterbox Turbo); +// see S3FlowDecoderWeights::meanflow. num_steps defaults to 2 in upstream Python (tts_turbo.py). +S3FlowCFMOutputs compute_s3_flow_cfm_meanflow( + S3FlowSessionCache & cache, + const S3FlowDecoderWeights & weights, + const std::vector & noise, + const std::vector & mask, + const std::vector & mu, + const std::vector & spks, + const std::vector & cond, + int64_t batch, + int64_t frames, + int64_t capacity_frames, + int64_t num_steps = 2, + engine::core::BackendConfig backend = {}); + +class S3GenSessionCache { +public: + explicit S3GenSessionCache(engine::core::BackendConfig backend = {}); + ~S3GenSessionCache(); + + S3GenSessionCache(const S3GenSessionCache &) = delete; + S3GenSessionCache & operator=(const S3GenSessionCache &) = delete; + + void release_runtime_graphs(); + +private: + friend class engine::codecs::s3gen::S3GenRuntime; + struct State; + std::unique_ptr state_; + + friend S3Token2MelOutputs compute_s3_token2mel_inference( + S3GenSessionCache & cache, + const S3FlowEncoderWeights & encoder_weights, + const S3FlowDecoderWeights & decoder_weights, + const S3GenConditioning & ref_dict, + const std::vector & speech_tokens, + int64_t speech_token_count, + int64_t num_steps, + float cfg_rate, + bool cosine_schedule, + const std::vector & full_noise, + uint64_t flow_seed, + engine::core::BackendConfig backend, + S3GenTimingBreakdown * timing); +}; -namespace engine::models::chatterbox { namespace { S3FlowEncoderWeights::LayerNormWeights load_flow_layer_norm( @@ -708,9 +997,9 @@ engine::core::TensorValue self_attention_no_mask( const S3FlowDecoderWeights::TransformerBlockWeights & weights, engine::core::DeferredTensorWriter & writer, const std::optional & attention_mask = std::nullopt) { - constexpr int64_t heads = 8; - constexpr int64_t head_dim = 64; - constexpr int64_t inner_dim = heads * head_dim; + const int64_t heads = weights.heads; + const int64_t head_dim = weights.head_dim; + const int64_t inner_dim = heads * head_dim; const float scale = 1.0f / std::sqrt(static_cast(head_dim)); auto q = linear_lastdim(ctx, input_btc, weights.attn_q, writer); auto k = linear_lastdim(ctx, input_btc, weights.attn_k, writer); @@ -818,14 +1107,14 @@ engine::core::TensorValue flow_prelookahead_backend( engine::core::DeferredTensorWriter & writer) { auto input_bct = transpose_bct_btc(ctx, input_btc); auto padded1 = engine::core::wrap_tensor( - ggml_pad_ext(ctx.ggml, contiguous(ctx, input_bct).tensor, 0, 3, 0, 0, 0, 0, 0, 0), - engine::core::TensorShape::from_dims({1, hidden, frames + 3}), + ggml_pad_ext(ctx.ggml, contiguous(ctx, input_bct).tensor, 0, conv1.kernel - 1, 0, 0, 0, 0, 0, 0), + engine::core::TensorShape::from_dims({1, hidden, frames + conv1.kernel - 1}), GGML_TYPE_F32); auto h = flow_conv1d_bct_backend(ctx, padded1, conv1, writer); h = engine::core::wrap_tensor(ggml_leaky_relu(ctx.ggml, h.tensor, 0.01f, false), h.shape, GGML_TYPE_F32); auto padded2 = engine::core::wrap_tensor( - ggml_pad_ext(ctx.ggml, contiguous(ctx, h).tensor, 2, 0, 0, 0, 0, 0, 0, 0), - engine::core::TensorShape::from_dims({1, hidden, frames + 2}), + ggml_pad_ext(ctx.ggml, contiguous(ctx, h).tensor, conv2.kernel - 1, 0, 0, 0, 0, 0, 0, 0), + engine::core::TensorShape::from_dims({1, hidden, frames + conv2.kernel - 1}), GGML_TYPE_F32); auto y_bct = flow_conv1d_bct_backend(ctx, padded2, conv2, writer); auto y_btc = transpose_bct_btc(ctx, y_bct); @@ -840,9 +1129,9 @@ engine::core::TensorValue flow_relative_attention_backend( int64_t hidden, const S3FlowEncoderWeights::RelativeAttentionWeights & weights, engine::core::DeferredTensorWriter & writer, - const std::optional & attention_mask = std::nullopt, - int64_t heads = 8, - int64_t head_dim = 64) { + const std::optional & attention_mask = std::nullopt) { + const auto heads = weights.pos_bias_u_tensor.shape.dims[0]; + const auto head_dim = weights.pos_bias_u_tensor.shape.dims[1]; auto q = linear_lastdim(ctx, input_btc, weights.q, writer); auto k = linear_lastdim(ctx, input_btc, weights.k, writer); auto v = linear_lastdim(ctx, input_btc, weights.v, writer); @@ -890,31 +1179,31 @@ engine::core::TensorValue flow_upsample_repeat_conv_backend( int64_t frames, int64_t hidden, const S3FlowEncoderWeights::Conv1dWeights & weights, - engine::core::DeferredTensorWriter & writer) { - std::vector repeat_weight(static_cast(hidden * hidden * 2), 0.0f); + engine::core::DeferredTensorWriter & writer, + int factor) { + std::vector repeat_weight(static_cast(hidden * hidden * factor), 0.0f); for (int64_t c = 0; c < hidden; ++c) { - const size_t base = static_cast((c * hidden + c) * 2); - repeat_weight[base] = 1.0f; - repeat_weight[base + 1] = 1.0f; + const size_t base = static_cast((c * hidden + c) * factor); + std::fill_n(repeat_weight.begin() + base, factor, 1.0f); } auto input_bct = transpose_bct_btc(ctx, input_btc); const auto repeat_kernel = make_f32_graph_constant( ctx, - engine::core::TensorShape::from_dims({hidden, hidden, 2}), + engine::core::TensorShape::from_dims({hidden, hidden, factor}), repeat_weight, writer); auto repeated = engine::modules::ConvTranspose1dModule({ hidden, hidden, - 2, - 2, + factor, + factor, 0, 1, false, }).build(ctx, input_bct, engine::modules::ConvTranspose1dWeights{repeat_kernel, std::nullopt}); auto padded = engine::core::wrap_tensor( - ggml_pad_ext(ctx.ggml, repeated.tensor, 4, 0, 0, 0, 0, 0, 0, 0), - engine::core::TensorShape::from_dims({1, hidden, frames * 2 + 4}), + ggml_pad_ext(ctx.ggml, repeated.tensor, weights.kernel - 1, 0, 0, 0, 0, 0, 0, 0), + engine::core::TensorShape::from_dims({1, hidden, frames * factor + weights.kernel - 1}), GGML_TYPE_F32); auto y_bct = flow_conv1d_bct_backend(ctx, padded, weights, writer); return transpose_bct_btc(ctx, y_bct); @@ -929,7 +1218,7 @@ class FlowEncoderBackendRunner { const engine::core::BackendConfig & backend_config) : input_frames_(frames), hidden_size_(hidden_size), - output_frames_(frames * 2), + output_frames_(frames * weights.config.upsample_factor), embed_prelook_runner_(weights, frames, hidden_size, backend_config), upsample_runner_(weights, frames, hidden_size, backend_config), after_norm_runner_(weights.after_norm, *weights.execution_context, output_frames_, hidden_size, backend_config) { @@ -963,7 +1252,7 @@ class FlowEncoderBackendRunner { copy_backend_tensor(*previous_output, upsample_runner_.input_tensor()); upsample_runner_.compute(); - const int64_t valid_output_frames = valid_frames * 2; + const int64_t valid_output_frames = valid_frames * (output_frames_ / input_frames_); const auto up_attention_mask_values = make_attention_key_mask(1, output_frames_, output_frames_, valid_output_frames); previous_output = &upsample_runner_.output_tensor(); for (auto & runner : up_encoder_runners_) { @@ -1280,7 +1569,7 @@ class FlowEncoderBackendRunner { frames, hidden_size, weights.up_layer_conv, - writer_); + writer_, weights.config.upsample_factor); x = linear_lastdim(ctx, x, weights.up_embed_linear, writer_); x = layer_norm_lastdim(ctx, x, weights.up_embed_norm, writer_); output_ = engine::core::wrap_tensor( @@ -1431,8 +1720,8 @@ class FlowDecoderBackendRunner { const S3FlowDecoderWeights & weights, int64_t frames, const engine::core::BackendConfig &) { - constexpr int64_t mel_channels = 80; - constexpr int64_t speaker_channels = 80; + const int64_t mel_channels = weights.final_proj.out_channels; + const int64_t speaker_channels = mel_channels; if (frames <= 0 || frames > capacity_frames_) { throw std::runtime_error("S3 flow decoder active frames exceed capacity"); } @@ -1478,7 +1767,7 @@ class FlowDecoderBackendRunner { x_in_ = engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims({batch_, mel_channels, frames})); mu_in_ = engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims({batch_, mel_channels, frames})); - time_in_ = engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims({batch_, 320})); + time_in_ = engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims({batch_, weights.config.time_embedding_size})); spks_in_ = engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims({batch_, speaker_channels})); cond_in_ = engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims({batch_, mel_channels, frames})); ggml_set_input(x_in_.tensor); @@ -1487,7 +1776,7 @@ class FlowDecoderBackendRunner { ggml_set_input(spks_in_.tensor); ggml_set_input(cond_in_.tensor); attention_mask_ = - engine::core::make_tensor(ctx, GGML_TYPE_F16, engine::core::TensorShape::from_dims({batch_, 8, frames, frames})); + engine::core::make_tensor(ctx, GGML_TYPE_F16, engine::core::TensorShape::from_dims({batch_, weights.config.heads, frames, frames})); ggml_set_input(attention_mask_.tensor); ggml_set_output(x_in_.tensor); ggml_set_output(mu_in_.tensor); @@ -1499,7 +1788,7 @@ class FlowDecoderBackendRunner { time_hidden = engine::core::wrap_tensor(ggml_silu(ctx.ggml, time_hidden.tensor), time_hidden.shape, GGML_TYPE_F32); time_hidden = linear_lastdim(ctx, time_hidden, weights.time_mlp_2, writer); if (weights.meanflow) { - r_in_ = engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims({batch_, 320})); + r_in_ = engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims({batch_, weights.config.time_embedding_size})); ggml_set_input(r_in_.tensor); ggml_set_output(r_in_.tensor); auto r_hidden = linear_lastdim(ctx, r_in_, weights.time_mlp_1, writer); @@ -1583,8 +1872,8 @@ class FlowDecoderBackendRunner { const std::vector & t, S3FlowDecoderRunTiming * timing = nullptr, const std::vector * r = nullptr) { - constexpr int64_t mel_channels = 80; - constexpr int64_t time_dim = 320; + const int64_t mel_channels = x_in_.shape.dims[1]; + const int64_t time_dim = time_in_.shape.dims[1]; if (timing != nullptr) { ++timing->calls; } @@ -1636,7 +1925,8 @@ class FlowDecoderBackendRunner { if (cached_attention_mask_valid_frames_ == valid_frames) { return; } - cached_attention_mask_values_ = make_attention_key_mask(batch_, active_frames_, active_frames_, valid_frames, 8); + cached_attention_mask_values_ = make_attention_key_mask( + batch_, active_frames_, active_frames_, valid_frames, attention_mask_.shape.dims[1]); engine::core::write_tensor_f16(attention_mask_, cached_attention_mask_values_); cached_attention_mask_valid_frames_ = valid_frames; } @@ -1734,8 +2024,6 @@ S3FlowSessionCache::S3FlowSessionCache(engine::core::BackendConfig backend) : state_(std::make_unique(backend)) {} S3FlowSessionCache::~S3FlowSessionCache() = default; -S3FlowSessionCache::S3FlowSessionCache(S3FlowSessionCache &&) noexcept = default; -S3FlowSessionCache & S3FlowSessionCache::operator=(S3FlowSessionCache &&) noexcept = default; void S3FlowSessionCache::release_encoder_graphs() { state_->encoder_runner.reset(); @@ -1754,13 +2042,19 @@ void S3FlowSessionCache::release_decoder_graphs() { std::shared_ptr load_s3_flow_encoder_weights( const engine::assets::TensorSource & source, const engine::core::ExecutionContext & execution_context, + const S3GenConfig & config, + const std::string & tensor_prefix, engine::assets::TensorStorageType weight_storage_type) { auto weights = std::make_shared(); + weights->config = config.encoder; + const auto & c = config.encoder; + const auto hidden = c.hidden_size; + const auto mel = config.mel_channels; weights->execution_context = &execution_context; weights->store = std::make_shared( execution_context.backend(), execution_context.backend_type(), - "chatterbox.s3_flow_encoder.weights", + "s3gen.encoder.weights", 1024ull * 1024ull * 1024ull); // ggml_get_rows (used by S3TokenEmbeddingGraph for this lookup table) only supports F32/F16 // and legacy quant types on the CUDA backend, not K-quants -- pin this one small table @@ -1768,42 +2062,42 @@ std::shared_ptr load_s3_flow_encoder_weights( // package (e.g. Chatterbox Turbo's chatterbox-turbo-s3gen-q4_k.gguf) doesn't crash on it. weights->input_embedding_tensor = weights->store->load_tensor( source, - "flow.input_embedding.weight", + tensor_prefix + "input_embedding.weight", engine::assets::TensorStorageType::F16, - {6561, 512}); - weights->speaker_affine = load_flow_linear(*weights->store, source, "flow.spk_embed_affine_layer", 80, 192, true, weight_storage_type); - weights->encoder_proj = load_flow_linear(*weights->store, source, "flow.encoder_proj", 80, 512, true, weight_storage_type); - weights->embed_linear = load_flow_linear(*weights->store, source, "flow.encoder.embed.out.0", 512, 512, true, weight_storage_type); - weights->embed_norm = load_flow_layer_norm(*weights->store, source, "flow.encoder.embed.out.1", 512); + {c.vocabulary_size, hidden}); + weights->speaker_affine = load_flow_linear(*weights->store, source, tensor_prefix + "spk_embed_affine_layer", mel, c.speaker_embedding_size, true, weight_storage_type); + weights->encoder_proj = load_flow_linear(*weights->store, source, tensor_prefix + "encoder_proj", mel, hidden, true, weight_storage_type); + weights->embed_linear = load_flow_linear(*weights->store, source, tensor_prefix + "encoder.embed.out.0", hidden, hidden, true, weight_storage_type); + weights->embed_norm = load_flow_layer_norm(*weights->store, source, tensor_prefix + "encoder.embed.out.1", hidden); weights->prelook_conv1 = - load_flow_conv1d(*weights->store, source, "flow.encoder.pre_lookahead_layer.conv1", 512, 512, 4, 1, 0, weight_storage_type); + load_flow_conv1d(*weights->store, source, tensor_prefix + "encoder.pre_lookahead_layer.conv1", hidden, hidden, c.prelook_kernel_size, 1, 0, weight_storage_type); weights->prelook_conv2 = - load_flow_conv1d(*weights->store, source, "flow.encoder.pre_lookahead_layer.conv2", 512, 512, 3, 1, 0, weight_storage_type); + load_flow_conv1d(*weights->store, source, tensor_prefix + "encoder.pre_lookahead_layer.conv2", hidden, hidden, c.prelook_output_kernel_size, 1, 0, weight_storage_type); auto load_layer = [&](const std::string & prefix) { S3FlowEncoderWeights::EncoderLayerWeights layer; - layer.norm_mha = load_flow_layer_norm(*weights->store, source, prefix + ".norm_mha", 512); - layer.attn.q = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_q", 512, 512, true, weight_storage_type); - layer.attn.k = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_k", 512, 512, true, weight_storage_type); - layer.attn.v = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_v", 512, 512, true, weight_storage_type); - layer.attn.out = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_out", 512, 512, true, weight_storage_type); - layer.attn.pos = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_pos", 512, 512, false, weight_storage_type); - layer.attn.pos_bias_u_tensor = weights->store->load_f32_tensor(source, prefix + ".self_attn.pos_bias_u", {8, 64}); - layer.attn.pos_bias_v_tensor = weights->store->load_f32_tensor(source, prefix + ".self_attn.pos_bias_v", {8, 64}); - layer.norm_ff = load_flow_layer_norm(*weights->store, source, prefix + ".norm_ff", 512); - layer.ff.w1 = load_flow_linear(*weights->store, source, prefix + ".feed_forward.w_1", 2048, 512, true, weight_storage_type); - layer.ff.w2 = load_flow_linear(*weights->store, source, prefix + ".feed_forward.w_2", 512, 2048, true, weight_storage_type); + layer.norm_mha = load_flow_layer_norm(*weights->store, source, prefix + ".norm_mha", hidden); + layer.attn.q = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_q", hidden, hidden, true, weight_storage_type); + layer.attn.k = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_k", hidden, hidden, true, weight_storage_type); + layer.attn.v = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_v", hidden, hidden, true, weight_storage_type); + layer.attn.out = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_out", hidden, hidden, true, weight_storage_type); + layer.attn.pos = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_pos", hidden, hidden, false, weight_storage_type); + layer.attn.pos_bias_u_tensor = weights->store->load_f32_tensor(source, prefix + ".self_attn.pos_bias_u", {c.heads, hidden / c.heads}); + layer.attn.pos_bias_v_tensor = weights->store->load_f32_tensor(source, prefix + ".self_attn.pos_bias_v", {c.heads, hidden / c.heads}); + layer.norm_ff = load_flow_layer_norm(*weights->store, source, prefix + ".norm_ff", hidden); + layer.ff.w1 = load_flow_linear(*weights->store, source, prefix + ".feed_forward.w_1", c.feed_forward_size, hidden, true, weight_storage_type); + layer.ff.w2 = load_flow_linear(*weights->store, source, prefix + ".feed_forward.w_2", hidden, c.feed_forward_size, true, weight_storage_type); return layer; }; - for (int i = 0; i < 6; ++i) { - weights->encoders.push_back(load_layer("flow.encoder.encoders." + std::to_string(i))); + for (int64_t i = 0; i < c.layers; ++i) { + weights->encoders.push_back(load_layer(tensor_prefix + "encoder.encoders." + std::to_string(i))); } - weights->up_layer_conv = load_flow_conv1d(*weights->store, source, "flow.encoder.up_layer.conv", 512, 512, 5, 1, 0, weight_storage_type); - weights->up_embed_linear = load_flow_linear(*weights->store, source, "flow.encoder.up_embed.out.0", 512, 512, true, weight_storage_type); - weights->up_embed_norm = load_flow_layer_norm(*weights->store, source, "flow.encoder.up_embed.out.1", 512); - for (int i = 0; i < 4; ++i) { - weights->up_encoders.push_back(load_layer("flow.encoder.up_encoders." + std::to_string(i))); + weights->up_layer_conv = load_flow_conv1d(*weights->store, source, tensor_prefix + "encoder.up_layer.conv", hidden, hidden, c.upsample_kernel_size, 1, 0, weight_storage_type); + weights->up_embed_linear = load_flow_linear(*weights->store, source, tensor_prefix + "encoder.up_embed.out.0", hidden, hidden, true, weight_storage_type); + weights->up_embed_norm = load_flow_layer_norm(*weights->store, source, tensor_prefix + "encoder.up_embed.out.1", hidden); + for (int64_t i = 0; i < c.upsample_layers; ++i) { + weights->up_encoders.push_back(load_layer(tensor_prefix + "encoder.up_encoders." + std::to_string(i))); } - weights->after_norm = load_flow_layer_norm(*weights->store, source, "flow.encoder.after_norm", 512); + weights->after_norm = load_flow_layer_norm(*weights->store, source, tensor_prefix + "encoder.after_norm", hidden); weights->store->upload(); return weights; } @@ -1822,13 +2116,21 @@ S3FlowEncoderOutputs compute_s3_flow_encoder_forward( std::shared_ptr load_s3_flow_decoder_weights( const engine::assets::TensorSource & source, const engine::core::ExecutionContext & execution_context, + const S3GenConfig & config, + const std::string & tensor_prefix, engine::assets::TensorStorageType weight_storage_type) { auto weights = std::make_shared(); + weights->config = config.decoder; + const auto & c = config.decoder; + const auto hidden = c.hidden_size; + const auto inner = c.heads * c.head_dim; + const auto mel = config.mel_channels; + const auto estimator = tensor_prefix + "decoder.estimator."; weights->execution_context = &execution_context; weights->store = std::make_shared( execution_context.backend(), execution_context.backend_type(), - "chatterbox.s3_flow_decoder.weights", + "s3gen.decoder.weights", 2048ull * 1024ull * 1024ull); auto load_causal_block = [&](const std::string & prefix, int64_t in_channels, int64_t out_channels) { @@ -1840,7 +2142,7 @@ std::shared_ptr load_s3_flow_decoder_weights( auto load_resnet = [&](const std::string & prefix, int64_t in_channels, int64_t out_channels) { S3FlowDecoderWeights::ResnetBlockWeights block; - block.time_mlp = load_decoder_linear(*weights->store, source, prefix + ".mlp.1", out_channels, 1024, true, weight_storage_type); + block.time_mlp = load_decoder_linear(*weights->store, source, prefix + ".mlp.1", out_channels, c.time_hidden_size, true, weight_storage_type); block.block1 = load_causal_block(prefix + ".block1", in_channels, out_channels); block.block2 = load_causal_block(prefix + ".block2", out_channels, out_channels); block.res_conv = load_decoder_conv1d(*weights->store, source, prefix + ".res_conv", out_channels, in_channels, 1, 1, true, weight_storage_type); @@ -1849,87 +2151,63 @@ std::shared_ptr load_s3_flow_decoder_weights( auto load_transformer = [&](const std::string & prefix) { S3FlowDecoderWeights::TransformerBlockWeights block; - block.norm1 = load_decoder_layer_norm(*weights->store, source, prefix + ".norm1", 256); - block.attn_q = load_decoder_linear(*weights->store, source, prefix + ".attn1.to_q", 512, 256, false, weight_storage_type); - block.attn_k = load_decoder_linear(*weights->store, source, prefix + ".attn1.to_k", 512, 256, false, weight_storage_type); - block.attn_v = load_decoder_linear(*weights->store, source, prefix + ".attn1.to_v", 512, 256, false, weight_storage_type); - block.attn_out = load_decoder_linear(*weights->store, source, prefix + ".attn1.to_out.0", 256, 512, true, weight_storage_type); - block.norm3 = load_decoder_layer_norm(*weights->store, source, prefix + ".norm3", 256); - block.ff_proj_in = load_decoder_linear(*weights->store, source, prefix + ".ff.net.0.proj", 1024, 256, true, weight_storage_type); - block.ff_proj_out = load_decoder_linear(*weights->store, source, prefix + ".ff.net.2", 256, 1024, true, weight_storage_type); + block.heads = c.heads; + block.head_dim = c.head_dim; + block.norm1 = load_decoder_layer_norm(*weights->store, source, prefix + ".norm1", hidden); + block.attn_q = load_decoder_linear(*weights->store, source, prefix + ".attn1.to_q", inner, hidden, false, weight_storage_type); + block.attn_k = load_decoder_linear(*weights->store, source, prefix + ".attn1.to_k", inner, hidden, false, weight_storage_type); + block.attn_v = load_decoder_linear(*weights->store, source, prefix + ".attn1.to_v", inner, hidden, false, weight_storage_type); + block.attn_out = load_decoder_linear(*weights->store, source, prefix + ".attn1.to_out.0", hidden, inner, true, weight_storage_type); + block.norm3 = load_decoder_layer_norm(*weights->store, source, prefix + ".norm3", hidden); + block.ff_proj_in = load_decoder_linear(*weights->store, source, prefix + ".ff.net.0.proj", c.feed_forward_size, hidden, true, weight_storage_type); + block.ff_proj_out = load_decoder_linear(*weights->store, source, prefix + ".ff.net.2", hidden, c.feed_forward_size, true, weight_storage_type); return block; }; weights->time_mlp_1 = - load_decoder_linear(*weights->store, source, "flow.decoder.estimator.time_mlp.linear_1", 1024, 320, true, weight_storage_type); + load_decoder_linear(*weights->store, source, estimator + "time_mlp.linear_1", c.time_hidden_size, c.time_embedding_size, true, weight_storage_type); weights->time_mlp_2 = - load_decoder_linear(*weights->store, source, "flow.decoder.estimator.time_mlp.linear_2", 1024, 1024, true, weight_storage_type); + load_decoder_linear(*weights->store, source, estimator + "time_mlp.linear_2", c.time_hidden_size, c.time_hidden_size, true, weight_storage_type); weights->down_blocks.resize(1); - weights->down_blocks[0].resnet = load_resnet("flow.decoder.estimator.down_blocks.0.0", 320, 256); - for (int i = 0; i < 4; ++i) { - weights->down_blocks[0].transformers.push_back(load_transformer("flow.decoder.estimator.down_blocks.0.1." + std::to_string(i))); + weights->down_blocks[0].resnet = load_resnet(estimator + "down_blocks.0.0", 4 * mel, hidden); + for (int64_t i = 0; i < c.attention_layers; ++i) { + weights->down_blocks[0].transformers.push_back(load_transformer(estimator + "down_blocks.0.1." + std::to_string(i))); } weights->down_blocks[0].downsample = - load_decoder_conv1d(*weights->store, source, "flow.decoder.estimator.down_blocks.0.2", 256, 256, 3, 1, true, weight_storage_type); + load_decoder_conv1d(*weights->store, source, estimator + "down_blocks.0.2", hidden, hidden, 3, 1, true, weight_storage_type); - weights->mid_blocks.resize(12); - for (int block_index = 0; block_index < 12; ++block_index) { + weights->mid_blocks.resize(c.middle_blocks); + for (int64_t block_index = 0; block_index < c.middle_blocks; ++block_index) { weights->mid_blocks[static_cast(block_index)].resnet = - load_resnet("flow.decoder.estimator.mid_blocks." + std::to_string(block_index) + ".0", 256, 256); - for (int i = 0; i < 4; ++i) { + load_resnet(estimator + "mid_blocks." + std::to_string(block_index) + ".0", hidden, hidden); + for (int64_t i = 0; i < c.attention_layers; ++i) { weights->mid_blocks[static_cast(block_index)].transformers.push_back( - load_transformer("flow.decoder.estimator.mid_blocks." + std::to_string(block_index) + ".1." + std::to_string(i))); + load_transformer(estimator + "mid_blocks." + std::to_string(block_index) + ".1." + std::to_string(i))); } } weights->up_blocks.resize(1); - weights->up_blocks[0].resnet = load_resnet("flow.decoder.estimator.up_blocks.0.0", 512, 256); - for (int i = 0; i < 4; ++i) { - weights->up_blocks[0].transformers.push_back(load_transformer("flow.decoder.estimator.up_blocks.0.1." + std::to_string(i))); + weights->up_blocks[0].resnet = load_resnet(estimator + "up_blocks.0.0", 2 * hidden, hidden); + for (int64_t i = 0; i < c.attention_layers; ++i) { + weights->up_blocks[0].transformers.push_back(load_transformer(estimator + "up_blocks.0.1." + std::to_string(i))); } weights->up_blocks[0].upsample = - load_decoder_conv1d(*weights->store, source, "flow.decoder.estimator.up_blocks.0.2", 256, 256, 3, 1, true, weight_storage_type); + load_decoder_conv1d(*weights->store, source, estimator + "up_blocks.0.2", hidden, hidden, 3, 1, true, weight_storage_type); - weights->final_block = load_causal_block("flow.decoder.estimator.final_block", 256, 256); - weights->final_proj = load_decoder_conv1d(*weights->store, source, "flow.decoder.estimator.final_proj", 80, 256, 1, 1, true, weight_storage_type); + weights->final_block = load_causal_block(estimator + "final_block", hidden, hidden); + weights->final_proj = load_decoder_conv1d(*weights->store, source, estimator + "final_proj", mel, hidden, 1, 1, true, weight_storage_type); - if (source.has_tensor("flow.decoder.estimator.time_embed_mixer.weight")) { + if (c.variant == S3FlowVariant::MeanFlow) { weights->meanflow = true; weights->time_embed_mixer = load_decoder_linear( - *weights->store, source, "flow.decoder.estimator.time_embed_mixer", 1024, 2048, false, weight_storage_type); + *weights->store, source, estimator + "time_embed_mixer", c.time_hidden_size, 2 * c.time_hidden_size, false, weight_storage_type); } weights->store->upload(); return weights; } -bool s3_flow_decoder_is_meanflow(const S3FlowDecoderWeights & weights) { - return weights.meanflow; -} - -S3FlowDecoderOutputs compute_s3_flow_decoder_forward( - S3FlowSessionCache & cache, - const S3FlowDecoderWeights & weights, - const std::vector & x, - const std::vector & mask, - const std::vector & mu, - const std::vector & t, - const std::vector & spks, - const std::vector & cond, - int64_t batch, - int64_t frames, - int64_t capacity_frames, - engine::core::BackendConfig backend, - const std::vector * r) { - const int64_t valid_frames = valid_frames_from_mask(mask, batch, frames); - auto & runner = cache.state_->decoder_runner_for_capacity(weights, batch, capacity_frames, backend); - runner.set_conditioning(weights, mu, spks, cond, frames, backend); - auto outputs = runner.run(x, mask, t, nullptr, r); - outputs.frames = valid_frames; - return outputs; -} - S3FlowCFMOutputs compute_s3_flow_cfm_euler( S3FlowSessionCache & cache, const S3FlowDecoderWeights & weights, @@ -1946,7 +2224,7 @@ S3FlowCFMOutputs compute_s3_flow_cfm_euler( bool cosine_schedule, engine::core::BackendConfig backend, S3FlowCFMTimingBreakdown * timing) { - constexpr int64_t mel_channels = 80; + const int64_t mel_channels = weights.final_proj.out_channels; if (timing != nullptr) { *timing = {}; timing->steps = num_steps; @@ -2057,7 +2335,7 @@ S3FlowCFMOutputs compute_s3_flow_cfm_meanflow( int64_t capacity_frames, int64_t num_steps, engine::core::BackendConfig backend) { - constexpr int64_t mel_channels = 80; + const int64_t mel_channels = weights.final_proj.out_channels; if (!weights.meanflow) { throw std::runtime_error("compute_s3_flow_cfm_meanflow requires meanflow-trained S3FlowDecoderWeights"); } @@ -2092,4 +2370,712 @@ S3FlowCFMOutputs compute_s3_flow_cfm_meanflow( return outputs; } -} // namespace engine::models::chatterbox +namespace { + +uint64_t choose_seed(uint64_t seed) { + if (seed != 0) return seed; + std::random_device rd; + return (static_cast(rd()) << 32U) ^ static_cast(rd()); +} + +int64_t token_capacity_quantum(const engine::core::BackendConfig & backend) { + (void) backend; + return 1; +} + +int64_t round_up_capacity(int64_t value, int64_t quantum) { + if (value <= 0 || quantum <= 0) { + throw std::runtime_error("S3 capacity rounding requires positive inputs"); + } + return ((value + quantum - 1) / quantum) * quantum; +} + +std::vector make_prefix_mask_bct(int64_t batch, int64_t frames, int64_t valid_frames) { + std::vector mask(static_cast(batch * frames), 0.0f); + for (int64_t batch_index = 0; batch_index < batch; ++batch_index) { + for (int64_t frame = 0; frame < frames; ++frame) { + mask[static_cast(batch_index * frames + frame)] = frame < valid_frames ? 1.0f : 0.0f; + } + } + return mask; +} + +std::vector zero_pad_tokens( + const std::vector & tokens, + int64_t capacity) { + std::vector padded(static_cast(capacity), 0); + std::copy(tokens.begin(), tokens.end(), padded.begin()); + return padded; +} + +std::vector zero_pad_btc( + const std::vector & values, + int64_t valid_frames, + int64_t capacity_frames, + int64_t channels) { + std::vector padded(static_cast(capacity_frames * channels), 0.0f); + for (int64_t frame = 0; frame < valid_frames; ++frame) { + for (int64_t channel = 0; channel < channels; ++channel) { + padded[static_cast(frame * channels + channel)] = + values[static_cast(frame * channels + channel)]; + } + } + return padded; +} + +std::vector slice_bct( + const std::vector & values, + int64_t source_frames, + int64_t frames, + int64_t channels) { + std::vector sliced(static_cast(frames * channels), 0.0f); + for (int64_t channel = 0; channel < channels; ++channel) { + for (int64_t frame = 0; frame < frames; ++frame) { + sliced[static_cast(channel * frames + frame)] = + values[static_cast(channel * source_frames + frame)]; + } + } + return sliced; +} + +void zero_tail_btc( + std::vector & values, + int64_t valid_frames, + int64_t capacity_frames, + int64_t channels) { + if (valid_frames >= capacity_frames) { + return; + } + for (int64_t frame = valid_frames; frame < capacity_frames; ++frame) { + std::fill_n( + values.begin() + static_cast(frame * channels), + static_cast(channels), + 0.0f); + } +} + +std::vector make_zero_padded_conditions_bct( + const S3GenConditioning & ref_dict, + int64_t total_frames, + int64_t mel_channels) { + std::vector cond(static_cast(mel_channels * total_frames), 0.0f); + for (int64_t frame = 0; frame < ref_dict.prompt_feat_frames; ++frame) { + for (int64_t dim = 0; dim < ref_dict.prompt_feat_dims; ++dim) { + cond[static_cast(dim * total_frames + frame)] = + ref_dict.prompt_feat[static_cast(frame * ref_dict.prompt_feat_dims + dim)]; + } + } + return cond; +} + +engine::core::TensorValue linear_rows_graph( + engine::core::ModuleBuildContext & ctx, + const engine::core::TensorValue & input, + const S3FlowEncoderWeights::LinearWeights & weights) { + auto projected = engine::core::wrap_tensor( + ggml_mul_mat(ctx.ggml, weights.weight_tensor.tensor, contiguous(ctx, input).tensor), + engine::core::TensorShape::from_dims({input.shape.dims[0], weights.out_features}), + GGML_TYPE_F32); + if (weights.use_bias) { + auto bias_view = engine::core::reshape_tensor( + ctx, + weights.bias_tensor, + engine::core::TensorShape::from_dims({1, weights.out_features})); + auto repeated = engine::core::wrap_tensor(ggml_repeat(ctx.ggml, bias_view.tensor, projected.tensor), projected.shape, GGML_TYPE_F32); + projected = add_tensor_values(ctx, projected, repeated); + } + return projected; +} + +class S3TokenEmbeddingGraph { +public: + S3TokenEmbeddingGraph( + const S3FlowEncoderWeights & weights, + int64_t token_count, + const engine::core::ExecutionContext & execution_context) + : token_count_(token_count), + execution_context_(&execution_context) { + ggml_init_params params = {}; + params.mem_size = 32ull * 1024ull * 1024ull; + params.mem_buffer = nullptr; + params.no_alloc = true; + ggml_ = ggml_init(params); + if (ggml_ == nullptr) { + throw std::runtime_error("failed to initialize S3 token embedding graph"); + } + + engine::core::ModuleBuildContext ctx = {}; + ctx.ggml = ggml_; + ctx.module_instance_name = "s3_token_embedding"; + ctx.backend_type = execution_context.backend_type(); + token_ids_ = engine::core::make_tensor( + ctx, + GGML_TYPE_I32, + engine::core::TensorShape::from_dims({token_count_})); + embedding_out_ = engine::core::wrap_tensor( + ggml_get_rows(ctx.ggml, weights.input_embedding_tensor.tensor, token_ids_.tensor), + engine::core::TensorShape::from_dims({token_count_, weights.config.hidden_size}), + GGML_TYPE_F32); + + graph_ = ggml_new_graph_custom(ggml_, 1024, false); + ggml_build_forward_expand(graph_, embedding_out_.tensor); + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_context_->backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + throw std::runtime_error("failed to allocate S3 token embedding graph tensors"); + } + } + + ~S3TokenEmbeddingGraph() { + if (execution_context_ != nullptr && graph_ != nullptr) { + engine::core::release_backend_graph_resources(execution_context_->backend(), graph_); + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + if (ggml_ != nullptr) { + ggml_free(ggml_); + } + } + + std::vector run(const std::vector & ids) { + if (static_cast(ids.size()) != token_count_) { + throw std::runtime_error("S3 token embedding graph token count mismatch"); + } + engine::core::write_tensor_i32(token_ids_, ids); + if (engine::core::compute_backend_graph(execution_context_->backend(), graph_) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("S3 token embedding graph compute failed"); + } + return engine::core::read_tensor_f32(embedding_out_.tensor); + } + +private: + int64_t token_count_ = 0; + const engine::core::ExecutionContext * execution_context_ = nullptr; + ggml_context * ggml_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_cgraph * graph_ = nullptr; + engine::core::TensorValue token_ids_; + engine::core::TensorValue embedding_out_; +}; + +class S3Token2MelPrepareGraph { +public: + struct Outputs { + std::vector mu_bct; + std::vector speaker; + }; + + S3Token2MelPrepareGraph( + const S3FlowEncoderWeights & weights, + int64_t frames, + const engine::core::ExecutionContext & execution_context) + : frames_(frames), + execution_context_(&execution_context) { + ggml_init_params params = {}; + params.mem_size = 64ull * 1024ull * 1024ull; + params.mem_buffer = nullptr; + params.no_alloc = true; + ggml_ = ggml_init(params); + if (ggml_ == nullptr) { + throw std::runtime_error("failed to initialize S3 token2mel prepare graph"); + } + + engine::core::ModuleBuildContext ctx = {}; + ctx.ggml = ggml_; + ctx.module_instance_name = "s3_token2mel_prepare"; + ctx.backend_type = execution_context.backend_type(); + encoded_in_ = engine::core::make_tensor( + ctx, + GGML_TYPE_F32, + engine::core::TensorShape::from_dims({frames_, weights.config.hidden_size})); + speaker_in_ = engine::core::make_tensor( + ctx, + GGML_TYPE_F32, + engine::core::TensorShape::from_dims({1, weights.config.speaker_embedding_size})); + + auto mu_rows = linear_rows_graph(ctx, encoded_in_, weights.encoder_proj); + mu_bct_ = engine::core::wrap_tensor( + ggml_cont(ctx.ggml, ggml_permute(ctx.ggml, mu_rows.tensor, 1, 0, 2, 3)), + engine::core::TensorShape::from_dims({weights.encoder_proj.out_features, frames_}), + GGML_TYPE_F32); + + auto speaker_sq = engine::core::wrap_tensor(ggml_sqr(ctx.ggml, speaker_in_.tensor), speaker_in_.shape, GGML_TYPE_F32); + auto norm_sq = engine::core::wrap_tensor( + ggml_sum_rows(ctx.ggml, speaker_sq.tensor), + engine::core::TensorShape::from_dims({1, 1}), + GGML_TYPE_F32); + auto norm = engine::core::wrap_tensor(ggml_sqrt(ctx.ggml, norm_sq.tensor), norm_sq.shape, GGML_TYPE_F32); + auto norm_rep = engine::core::wrap_tensor(ggml_repeat(ctx.ggml, norm.tensor, speaker_in_.tensor), speaker_in_.shape, GGML_TYPE_F32); + auto speaker_norm = engine::core::wrap_tensor( + ggml_div(ctx.ggml, speaker_in_.tensor, norm_rep.tensor), + speaker_in_.shape, + GGML_TYPE_F32); + speaker_out_ = linear_rows_graph(ctx, speaker_norm, weights.speaker_affine); + + graph_ = ggml_new_graph_custom(ggml_, 4096, false); + ggml_build_forward_expand(graph_, mu_bct_.tensor); + ggml_build_forward_expand(graph_, speaker_out_.tensor); + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_context_->backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + throw std::runtime_error("failed to allocate S3 token2mel prepare graph tensors"); + } + + } + + ~S3Token2MelPrepareGraph() { + if (execution_context_ != nullptr && graph_ != nullptr) { + engine::core::release_backend_graph_resources(execution_context_->backend(), graph_); + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + if (ggml_ != nullptr) { + ggml_free(ggml_); + } + } + + Outputs run( + const std::vector & encoded, + const std::vector & speaker_embedding) { + engine::core::write_tensor_f32(encoded_in_, encoded); + engine::core::write_tensor_f32(speaker_in_, speaker_embedding); + if (engine::core::compute_backend_graph(execution_context_->backend(), graph_) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("S3 token2mel prepare graph compute failed"); + } + Outputs outputs; + outputs.mu_bct = engine::core::read_tensor_f32(mu_bct_.tensor); + outputs.speaker = engine::core::read_tensor_f32(speaker_out_.tensor); + return outputs; + } + +private: + int64_t frames_ = 0; + const engine::core::ExecutionContext * execution_context_ = nullptr; + ggml_context * ggml_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_cgraph * graph_ = nullptr; + engine::core::TensorValue encoded_in_; + engine::core::TensorValue speaker_in_; + engine::core::TensorValue mu_bct_; + engine::core::TensorValue speaker_out_; +}; + +std::vector make_gaussian_full_noise(int64_t channels, int64_t frames, uint64_t seed) { + return engine::sampling::generate_torch_cuda_randn( + static_cast(channels * frames), + choose_seed(seed), + engine::sampling::TorchRandnPrecision::Float32); +} + +} // namespace + +struct S3GenSessionCache::State { + explicit State(engine::core::BackendConfig backend_value) + : backend(std::move(backend_value)), + token_capacity_controller(runtime::GraphCapacityMode::Grow), + flow_cache(backend) {} + + struct PreparedCapacity { + int64_t token_capacity = 0; + int64_t frame_capacity = 0; + std::unique_ptr token_embedding; + std::unique_ptr token2mel_prepare; + }; + + runtime::MappedGraphCapacityAdapter make_token_capacity_adapter( + const S3FlowEncoderWeights & weights, + const engine::core::BackendConfig & backend_config) { + const int64_t quantum = token_capacity_quantum(backend_config); + return runtime::MappedGraphCapacityAdapter( + quantum, + quantum, + [quantum](int64_t request_size) { + return round_up_capacity(request_size, quantum); + }, + [this]() { + std::vector capacities; + if (prepared_capacity.token_embedding && prepared_capacity.token2mel_prepare) { + capacities.push_back(prepared_capacity.token_capacity); + } + return capacities; + }, + [this, &weights](int64_t capacity) { + prepare_token_capacity(weights, capacity); + }); + } + + void prepare_token_capacity( + const S3FlowEncoderWeights & weights, + int64_t token_capacity) { + if (token_capacity <= 0) { + throw std::runtime_error("S3 token capacity must be positive"); + } + if (prepared_capacity.token_embedding && + prepared_capacity.token2mel_prepare && + prepared_capacity.token_capacity == token_capacity) { + return; + } + if (weights.execution_context == nullptr) { + throw std::runtime_error("S3 token2mel prepare requires loaded backend weights"); + } + prepared_capacity.token_capacity = token_capacity; + prepared_capacity.frame_capacity = token_capacity * weights.config.upsample_factor; + prepared_capacity.token_embedding = + std::make_unique(weights, token_capacity, *weights.execution_context); + prepared_capacity.token2mel_prepare = + std::make_unique(weights, prepared_capacity.frame_capacity, *weights.execution_context); + } + + PreparedCapacity & prepared_bundle( + const S3FlowEncoderWeights & weights, + int64_t token_count, + const engine::core::BackendConfig & backend_config) { + auto adapter = make_token_capacity_adapter(weights, backend_config); + token_capacity_controller.ensure_prepared(adapter, token_count); + const int64_t selected_capacity = token_capacity_controller.select_capacity_for_run(adapter, token_count); + if (prepared_capacity.token_capacity != selected_capacity) { + throw std::runtime_error("S3 token capacity controller selected an unprepared capacity"); + } + return prepared_capacity; + } + + void release_pre_cfm_graphs() { + prepared_capacity.token_embedding.reset(); + prepared_capacity.token2mel_prepare.reset(); + flow_cache.release_encoder_graphs(); + } + + void release_cfm_decoder_graphs() { + flow_cache.release_decoder_graphs(); + } + + engine::core::BackendConfig backend; + runtime::GraphCapacityController token_capacity_controller; + S3FlowSessionCache flow_cache; + PreparedCapacity prepared_capacity; +}; + +S3GenSessionCache::S3GenSessionCache(engine::core::BackendConfig backend) + : state_(std::make_unique(backend)) {} + +S3GenSessionCache::~S3GenSessionCache() = default; + +void S3GenSessionCache::release_runtime_graphs() { + state_->release_pre_cfm_graphs(); + state_->release_cfm_decoder_graphs(); +} + +S3Token2MelOutputs compute_s3_token2mel_inference( + S3GenSessionCache & cache, + const S3FlowEncoderWeights & encoder_weights, + const S3FlowDecoderWeights & decoder_weights, + const S3GenConditioning & ref_dict, + const std::vector & speech_tokens, + int64_t speech_token_count, + int64_t num_steps, + float cfg_rate, + bool cosine_schedule, + const std::vector & full_noise, + uint64_t flow_seed, + engine::core::BackendConfig backend, + S3GenTimingBreakdown * timing) { + constexpr int64_t batch = 1; + const int64_t token_dim = encoder_weights.config.hidden_size; + const int64_t mel_dim = encoder_weights.encoder_proj.out_features; + const int64_t prompt_tokens = ref_dict.prompt_token_count; + const int64_t total_tokens = prompt_tokens + speech_token_count; + std::vector concat_tokens; + concat_tokens.reserve(static_cast(total_tokens)); + concat_tokens.insert(concat_tokens.end(), ref_dict.prompt_tokens.begin(), ref_dict.prompt_tokens.end()); + concat_tokens.insert(concat_tokens.end(), speech_tokens.begin(), speech_tokens.begin() + static_cast(speech_token_count)); + + auto & capacity_bundle = cache.state_->prepared_bundle(encoder_weights, total_tokens, backend); + const int64_t token_capacity = capacity_bundle.token_capacity; + const int64_t frame_capacity = capacity_bundle.frame_capacity; + const auto padded_tokens = zero_pad_tokens(concat_tokens, token_capacity); + + const auto embed_started = std::chrono::steady_clock::now(); + auto token_emb = capacity_bundle.token_embedding->run(padded_tokens); + zero_tail_btc(token_emb, total_tokens, token_capacity, token_dim); + if (timing != nullptr) { + timing->token2mel_embed_ms = + std::chrono::duration(std::chrono::steady_clock::now() - embed_started).count(); + } + + const auto encoder_started = std::chrono::steady_clock::now(); + const auto encoded = + compute_s3_flow_encoder_forward( + cache.state_->flow_cache, + encoder_weights, + token_emb, + total_tokens, + token_capacity, + token_dim, + backend); + std::vector encoded_hidden = encoded.hidden; + if (static_cast(encoded_hidden.size()) != encoded.storage_frames * token_dim) { + encoded_hidden = zero_pad_btc(encoded_hidden, encoded.frames, encoded.storage_frames, token_dim); + } + if (timing != nullptr) { + timing->token2mel_encoder_ms = + std::chrono::duration(std::chrono::steady_clock::now() - encoder_started).count(); + } + + const auto mu_started = std::chrono::steady_clock::now(); + const int64_t total_frames = encoded.frames; + const auto prepared = capacity_bundle.token2mel_prepare->run( + encoded_hidden, + ref_dict.embedding); + auto cond = make_zero_padded_conditions_bct(ref_dict, total_frames, mel_dim); + auto mask = make_prefix_mask_bct(batch, total_frames, total_frames); + auto mu = slice_bct(prepared.mu_bct, frame_capacity, total_frames, mel_dim); + if (timing != nullptr) { + timing->token2mel_mu_ms = + std::chrono::duration(std::chrono::steady_clock::now() - mu_started).count(); + } + + std::vector noise = full_noise.empty() + ? make_gaussian_full_noise(mel_dim, total_frames, flow_seed) + : full_noise; + if (static_cast(noise.size()) != mel_dim * total_frames) { + throw std::runtime_error( + "S3Token2Mel full_noise size mismatch: got=" + + std::to_string(noise.size()) + + " expected=" + + std::to_string(mel_dim * total_frames) + + " total_frames=" + + std::to_string(total_frames) + + " prompt_tokens=" + + std::to_string(prompt_tokens) + + " speech_tokens=" + + std::to_string(speech_token_count)); + } + cache.state_->release_pre_cfm_graphs(); + const auto cfm_started = std::chrono::steady_clock::now(); + // Meanflow-distilled decoders (Chatterbox Turbo) use a plain, non-CFG Euler solve driven by + // two time inputs per step; see S3FlowDecoderWeights::meanflow and + // compute_s3_flow_cfm_meanflow. cfg_rate/cosine_schedule are meaningless for that path and + // are ignored, matching upstream's ConditionalCFM.forward branching on `meanflow`. + const auto mel_full = decoder_weights.meanflow + ? compute_s3_flow_cfm_meanflow( + cache.state_->flow_cache, + decoder_weights, + noise, + mask, + mu, + prepared.speaker, + cond, + batch, + total_frames, + frame_capacity, + num_steps, + backend) + : compute_s3_flow_cfm_euler( + cache.state_->flow_cache, + decoder_weights, + noise, + mask, + mu, + prepared.speaker, + cond, + batch, + total_frames, + frame_capacity, + num_steps, + cfg_rate, + cosine_schedule, + backend, + timing == nullptr ? nullptr : &timing->token2mel_cfm); + if (timing != nullptr) { + timing->token2mel_cfm_ms = + std::chrono::duration(std::chrono::steady_clock::now() - cfm_started).count(); + } + + const int64_t prompt_frames = ref_dict.prompt_feat_frames; + const int64_t output_frames = total_frames - prompt_frames; + S3Token2MelOutputs outputs; + outputs.channels = mel_dim; + outputs.frames = output_frames; + outputs.mel.assign(static_cast(mel_dim * output_frames), 0.0f); + const int64_t mel_stride = mel_full.storage_frames; + for (int64_t dim = 0; dim < mel_dim; ++dim) { + for (int64_t frame = 0; frame < output_frames; ++frame) { + outputs.mel[static_cast(dim * output_frames + frame)] = + mel_full.mel[static_cast(dim * mel_stride + prompt_frames + frame)]; + } + } + cache.state_->release_cfm_decoder_graphs(); + return outputs; +} + +} // namespace + +S3GenConfig::S3GenConfig() { + auto & config = vocoder; + config.in_channels = 80; + config.base_channels = 512; + config.nb_harmonics = 8; + config.sampling_rate = 24000; + config.nsf_alpha = 0.1F; + config.nsf_sigma = 0.003F; + config.nsf_voiced_threshold = 10.0F; + config.upsample_rates = {8, 5, 3}; + config.upsample_kernel_sizes = {16, 11, 7}; + config.istft_n_fft = 16; + config.istft_hop = 4; + config.resblock_kernel_sizes = {3, 7, 11}; + config.resblock_dilation_sizes = {{1, 3, 5}, {1, 3, 5}, {1, 3, 5}}; + config.source_resblock_kernel_sizes = {7, 7, 11}; + config.source_resblock_dilation_sizes = {{1, 3, 5}, {1, 3, 5}, {1, 3, 5}}; + config.lrelu_slope = 0.1F; + config.audio_limit = 0.99F; + config.f0_num_class = 1; + config.f0_in_channels = 80; + config.f0_cond_channels = 512; +} + +struct S3GenRuntime::Impl { + const core::ExecutionContext & execution; + std::shared_ptr encoder; + std::shared_ptr decoder; + modules::HiftVocoderComponent vocoder; + S3GenSessionCache cache; + + Impl(std::shared_ptr source, + const core::ExecutionContext & context, const S3GenConfig & config, + const S3GenRuntimeOptions & options, const S3GenWeightBinding & binding) + : execution(context), cache(context.config()) { + if (config.encoder.heads <= 0 || config.encoder.hidden_size <= 0 || + config.encoder.hidden_size % config.encoder.heads != 0 || + config.encoder.hidden_size % 2 != 0 || config.encoder.upsample_factor <= 0 || + config.encoder.vocabulary_size <= 0 || config.encoder.speaker_embedding_size <= 0 || + config.encoder.layers < 0 || config.encoder.upsample_layers < 0 || + config.encoder.feed_forward_size <= 0 || config.encoder.prelook_kernel_size <= 0 || + config.encoder.prelook_output_kernel_size <= 0 || config.encoder.upsample_kernel_size <= 0 || + config.decoder.heads <= 0 || config.decoder.head_dim <= 0 || + config.decoder.hidden_size <= 0 || config.decoder.feed_forward_size <= 0 || + config.decoder.time_embedding_size < 4 || config.decoder.time_embedding_size % 2 != 0 || + config.decoder.time_hidden_size <= 0 || config.decoder.middle_blocks < 0 || + config.decoder.attention_layers < 0 || config.mel_channels <= 0 || + config.vocoder.in_channels != config.mel_channels) { + throw std::runtime_error("Invalid S3Gen architecture configuration"); + } + encoder = load_s3_flow_encoder_weights(*source, context, config, binding.flow_prefix, options.weight_storage_type); + decoder = load_s3_flow_decoder_weights(*source, context, config, binding.flow_prefix, options.weight_storage_type); + auto vocoder_config = config.vocoder; + vocoder_config.weight_storage_type = options.weight_storage_type; + vocoder_config.tensor_prefix = binding.vocoder_prefix; + vocoder_config.weight_layout = binding.vocoder_layout; + vocoder = modules::HiftVocoderComponent::load_from_tensor_source(source, context.config(), vocoder_config); + } +}; + +S3GenRuntime::S3GenRuntime(std::shared_ptr source, + const core::ExecutionContext & execution, S3GenConfig config, + S3GenRuntimeOptions options, S3GenWeightBinding binding) + : impl_(std::make_unique(std::move(source), execution, config, options, binding)) {} +S3GenRuntime::~S3GenRuntime() = default; +S3GenRuntime::S3GenRuntime(S3GenRuntime &&) noexcept = default; +S3GenRuntime & S3GenRuntime::operator=(S3GenRuntime &&) noexcept = default; + +bool S3GenRuntime::is_meanflow() const { return impl_->decoder->meanflow; } + +S3FlowEncoderOutputs S3GenRuntime::encode( + const std::vector & embeddings, int64_t frames, + int64_t capacity_frames, int64_t hidden_size) const { + return compute_s3_flow_encoder_forward( + impl_->cache.state_->flow_cache, *impl_->encoder, embeddings, + frames, capacity_frames, hidden_size, impl_->execution.config()); +} + +S3Token2MelOutputs S3GenRuntime::token_to_mel( + const S3GenConditioning & reference, + const std::vector & speech_tokens, int64_t speech_token_count, + int64_t num_steps, float cfg_rate, bool cosine_schedule, + const std::vector & full_noise, uint64_t flow_seed, + S3GenTimingBreakdown * timing) const { + return compute_s3_token2mel_inference( + impl_->cache, *impl_->encoder, *impl_->decoder, reference, + speech_tokens, speech_token_count, num_steps, cfg_rate, cosine_schedule, + full_noise, flow_seed, impl_->execution.config(), timing); +} + +modules::HiftVocoderOutput S3GenRuntime::decode_waveform( + const std::vector & mel, int64_t frames, + uint64_t seed, uint64_t prior_noise_values) const { + return impl_->vocoder.synthesize(mel, frames, seed, prior_noise_values); +} + +const modules::HiftVocoderComponent & S3GenRuntime::vocoder() const { + return impl_->vocoder; +} + +void S3GenRuntime::release_flow_graphs() const { + impl_->cache.release_runtime_graphs(); +} + +void S3GenRuntime::release_vocoder_graphs() const { + impl_->vocoder.release_runtime_cache(); +} + +S3GenInferenceOutputs S3GenRuntime::synthesize( + const S3GenConditioning & ref_dict, + const std::vector & speech_tokens, + int64_t speech_token_count, + int64_t num_steps, + float cfg_rate, + bool cosine_schedule, + const std::vector & full_noise, + uint64_t flow_seed, + uint64_t vocoder_seed, + S3GenTimingBreakdown * timing) const { + const auto token2mel_started = std::chrono::steady_clock::now(); + const auto mel = token_to_mel( + ref_dict, + speech_tokens, + speech_token_count, + num_steps, + cfg_rate, + cosine_schedule, + full_noise, + flow_seed, + timing); + if (timing != nullptr) { + timing->token2mel_ms = + std::chrono::duration(std::chrono::steady_clock::now() - token2mel_started).count(); + } + const uint64_t prior_noise_values = full_noise.empty() + ? static_cast(mel.channels * (ref_dict.prompt_feat_frames + mel.frames)) + : 0; + const auto vocoder_started = std::chrono::steady_clock::now(); + const auto vocoder_outputs = decode_waveform( + mel.mel, + mel.frames, + vocoder_seed, + prior_noise_values); + if (timing != nullptr) { + timing->vocoder_ms = + std::chrono::duration(std::chrono::steady_clock::now() - vocoder_started).count(); + } + S3GenInferenceOutputs outputs; + outputs.waveform = vocoder_outputs.waveform; + outputs.samples = vocoder_outputs.samples; + outputs.mel = mel.mel; + outputs.mel_channels = mel.channels; + outputs.mel_frames = mel.frames; + return outputs; +} + +} // namespace engine::codecs::s3gen diff --git a/src/framework/modules/activation_modules.cpp b/src/framework/modules/activation_modules.cpp index 531004900..9d05ae465 100644 --- a/src/framework/modules/activation_modules.cpp +++ b/src/framework/modules/activation_modules.cpp @@ -1,4 +1,6 @@ #include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" #include @@ -629,6 +631,51 @@ core::TensorValue Snake1dModule::build( return core::wrap_tensor(ggml_add(ctx.ggml, input_f32.tensor, frac.tensor), input_f32.shape, GGML_TYPE_F32); } +PReluModule::PReluModule(PReluConfig config) : config_(config) {} + +const core::ModuleSchema & PReluModule::schema() const noexcept { + return static_schema(); +} + +const core::ModuleSchema & PReluModule::static_schema() noexcept { + static const core::ModulePortSpec inputs[] = { + {"input", core::PortKind::Activation, false}, + {"slope", core::PortKind::Parameter, false}, + }; + static const core::ModuleSchema schema = { + "PReLU", "nn.activation", inputs, 2, kActivationOutputs, 1, + "Applies learned scalar or channel-wise negative slopes.", + }; + return schema; +} + +core::TensorValue PReluModule::build( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & slope) const { + core::validate_rank_between(input, 1, core::kMaxTensorRank, "input"); + if (config_.channel_axis >= input.shape.rank) { + throw std::runtime_error("PReLU channel axis is outside input rank"); + } + auto alpha = slope; + if (!same_shape(alpha.shape, input.shape)) { + const auto channels = alpha.shape.num_elements(); + if (channels != 1 && channels != input.shape.dims[config_.channel_axis]) { + throw std::runtime_error("PReLU slope must be scalar or match the channel dimension"); + } + auto shape = input.shape; + for (size_t axis = 0; axis < shape.rank; ++axis) shape.dims[axis] = 1; + shape.dims[config_.channel_axis] = channels; + alpha = core::reshape_tensor(ctx, alpha, shape); + alpha = RepeatModule({input.shape}).build(ctx, alpha); + } + const auto x = core::ensure_backend_addressable_layout(ctx, input); + const auto positive = ReluModule().build(ctx, x); + auto negative = core::wrap_tensor(ggml_sub(ctx.ggml, x.tensor, positive.tensor), x.shape); + negative = MulModule().build(ctx, negative, alpha); + return AddModule().build(ctx, positive, negative); +} + const core::ModuleSchema & Snake1dModule::static_schema() noexcept { return kSnake1dSchema; } diff --git a/src/framework/modules/attention/scaled_dot_product_attention.cpp b/src/framework/modules/attention/scaled_dot_product_attention.cpp index 6514f1d58..ac9ea65ec 100644 --- a/src/framework/modules/attention/scaled_dot_product_attention.cpp +++ b/src/framework/modules/attention/scaled_dot_product_attention.cpp @@ -47,7 +47,7 @@ void validate_heads(const core::TensorValue & q, const core::TensorValue & k, co if (k.shape.dims[2] != v.shape.dims[2]) { throw std::runtime_error("ScaledDotProductAttention key/value step mismatch"); } - if (q.shape.dims[3] != dim || k.shape.dims[3] != dim || v.shape.dims[3] != dim) { + if (q.shape.dims[3] != dim || k.shape.dims[3] != dim || v.shape.dims[3] <= 0) { throw std::runtime_error("ScaledDotProductAttention head dimension mismatch"); } } @@ -60,19 +60,28 @@ core::TensorValue build_explicit( const std::optional & attention_mask, float scale, ggml_prec precision, - AttentionCausality causality) { + AttentionCausality causality, + const std::optional & score_scale) { const MatMulModule matmul; auto scores = matmul.build(ctx, q_heads, TransposeModule({{0, 1, 3, 2}, k_heads.shape.rank}).build(ctx, k_heads)); ggml_mul_mat_set_prec(scores.tensor, precision); scores = core::ensure_backend_addressable_layout(ctx, scores); + if (score_scale.has_value()) { + auto scalar_shape = scores.shape; + for (size_t axis = 0; axis < scalar_shape.rank; ++axis) scalar_shape.dims[axis] = 1; + auto multiplier = core::reshape_tensor(ctx, *score_scale, scalar_shape); + scores = MulModule().build(ctx, scores, RepeatModule({scores.shape}).build(ctx, multiplier)); + } core::TensorValue attn; if (attention_mask.has_value()) { attn = core::wrap_tensor( - ggml_soft_max_ext(ctx.ggml, scores.tensor, attention_mask->tensor, scale, 0.0F), + ggml_soft_max_ext(ctx.ggml, scores.tensor, attention_mask->tensor, score_scale ? 1.0F : scale, 0.0F), scores.shape, GGML_TYPE_F32); } else { - scores = core::wrap_tensor(ggml_scale(ctx.ggml, scores.tensor, scale), scores.shape, GGML_TYPE_F32); + if (!score_scale) { + scores = core::wrap_tensor(ggml_scale(ctx.ggml, scores.tensor, scale), scores.shape, GGML_TYPE_F32); + } if (causality == AttentionCausality::Causal) { scores = core::wrap_tensor(ggml_diag_mask_inf(ctx.ggml, scores.tensor, 0), scores.shape, GGML_TYPE_F32); } @@ -135,7 +144,7 @@ core::TensorValue build_explicit_cpu_per_head_one( auto v_2d = core::reshape_tensor( ctx, core::ensure_backend_addressable_layout(ctx, v), - core::TensorShape::from_dims({key_steps, head_dim})); + core::TensorShape::from_dims({key_steps, v.shape.dims[3]})); ggml_tensor * scores_raw = ggml_mul_mat(ctx.ggml, k_2d.tensor, q_2d.tensor); ggml_mul_mat_set_prec(scores_raw, precision); auto scores = core::wrap_tensor( @@ -152,7 +161,7 @@ core::TensorValue build_explicit_cpu_per_head_one( return core::reshape_tensor( ctx, core::ensure_backend_addressable_layout(ctx, context), - core::TensorShape::from_dims({1, query_steps, 1, head_dim})); + core::TensorShape::from_dims({1, query_steps, 1, v.shape.dims[3]})); } core::TensorValue build_explicit_cpu_per_head( @@ -272,12 +281,22 @@ core::TensorValue ScaledDotProductAttentionModule::build( const core::TensorValue & q_heads, const core::TensorValue & k_heads, const core::TensorValue & v_heads, - const std::optional & attention_mask) const { + const std::optional & attention_mask, + const std::optional & score_scale) const { validate_heads(q_heads, k_heads, v_heads, config_.head_dim); + if (score_scale && (config_.lowering != ScaledDotProductAttentionLowering::Explicit || + score_scale->shape.num_elements() != 1)) { + throw std::runtime_error("ScaledDotProductAttention tensor score scale requires explicit lowering and one element"); + } + if ((config_.lowering == ScaledDotProductAttentionLowering::Flash || + config_.lowering == ScaledDotProductAttentionLowering::FlashPreserveViews) && + v_heads.shape.dims[3] != config_.head_dim) { + throw std::runtime_error("ScaledDotProductAttention flash lowering requires equal Q/K/V widths"); + } const float scale = 1.0F / std::sqrt(static_cast(config_.head_dim)); switch (config_.lowering) { case ScaledDotProductAttentionLowering::Explicit: - return build_explicit(ctx, q_heads, k_heads, v_heads, attention_mask, scale, config_.precision, config_.causality); + return build_explicit(ctx, q_heads, k_heads, v_heads, attention_mask, scale, config_.precision, config_.causality, score_scale); case ScaledDotProductAttentionLowering::ExplicitCpuPerHead: return build_explicit_cpu_per_head(ctx, q_heads, k_heads, v_heads, attention_mask, scale, config_.precision, config_.causality); case ScaledDotProductAttentionLowering::Flash: diff --git a/src/framework/modules/conformer_modules.cpp b/src/framework/modules/conformer_modules.cpp index 7c748678a..d6527608e 100644 --- a/src/framework/modules/conformer_modules.cpp +++ b/src/framework/modules/conformer_modules.cpp @@ -6,6 +6,7 @@ #include "engine/framework/modules/structural_modules.h" #include +#include namespace engine::modules { @@ -178,6 +179,60 @@ ConvSubsamplingOutputs ConvSubsamplingModule::build( }; } +EspnetConv2dSubsampling8Module::EspnetConv2dSubsampling8Module(EspnetConv2dSubsampling8Config config) + : config_(config) { + if (config_.input_features < 15 || config_.hidden_size <= 0 || config_.time_row_alignment <= 0) { + throw std::runtime_error("Invalid ESPnet Conv2D subsampling dimensions or row alignment"); + } +} + +core::TensorValue EspnetConv2dSubsampling8Module::build( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const EspnetConv2dSubsampling8Weights & weights) const { + core::validate_rank_between(input, 3, 3, "input"); + core::validate_last_dim(input, config_.input_features, "input"); + const auto batch = input.shape.dims[0]; + const auto frames = input.shape.dims[1]; + if (frames < 15) { + throw std::runtime_error("ESPnet Conv2D subsampling requires at least 15 input frames"); + } + const auto d = config_.hidden_size; + const auto alignment = config_.time_row_alignment; + auto x = core::reshape_tensor(ctx, input, + core::TensorShape::from_dims({batch, 1, frames, config_.input_features})); + int64_t valid_frames = (frames - 3) / 2 + 1; + if (alignment > 1) { + // Align the second convolution by padding the small input, not its activation. + const auto next_frames = (valid_frames - 3) / 2 + 1; + const auto aligned_next_frames = ((next_frames + alignment - 1) / alignment) * alignment; + const auto padded_input_frames = 2 * (2 * aligned_next_frames + 1) + 1; + x = Pad2dModule({0, 0, 0, std::max(0, padded_input_frames - frames)}).build(ctx, x); + } + x = Conv2dModule({1, d, 3, 3, 2, 2, 0, 0, 1, 1, true}).build(ctx, x, weights.conv0); + x = ReluModule().build(ctx, x); + for (const auto * conv : {&weights.conv1, &weights.conv2}) { + const auto output_frames = (valid_frames - 3) / 2 + 1; + const auto padded_frames = ((output_frames + alignment - 1) / alignment) * alignment; + const auto pad_frames = std::max(0, 2 * padded_frames + 1 - x.shape.dims[2]); + if (pad_frames > 0) { + x = Pad2dModule({0, 0, 0, pad_frames}).build(ctx, x); + } + x = Conv2dModule({d, d, 3, 3, 2, 2, 0, 0, 1, 1, true}).build(ctx, x, *conv); + if (padded_frames != output_frames) { + x = SliceModule({2, 0, output_frames}).build(ctx, x); + } + x = ReluModule().build(ctx, x); + valid_frames = output_frames; + } + const auto flattened_features = d * x.shape.dims[3]; + x = TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, x); + x = core::ensure_backend_addressable_layout(ctx, x); + x = core::reshape_tensor(ctx, x, + core::TensorShape::from_dims({batch, valid_frames, flattened_features})); + return LinearModule({flattened_features, d, true}).build(ctx, x, weights.projection); +} + DepthwiseConvSubsamplingModule::DepthwiseConvSubsamplingModule(DepthwiseConvSubsamplingConfig config) : config_(config) { if (config.input_features <= 0 || config.output_features <= 0 || config.conv_channels <= 0 || diff --git a/src/models/qwen3_asr/audio_encoder.cpp b/src/framework/modules/speech_encoders/qwen3_audio_encoder_runtime.cpp similarity index 79% rename from src/models/qwen3_asr/audio_encoder.cpp rename to src/framework/modules/speech_encoders/qwen3_audio_encoder_runtime.cpp index b651d78cc..bb116f76a 100644 --- a/src/models/qwen3_asr/audio_encoder.cpp +++ b/src/framework/modules/speech_encoders/qwen3_audio_encoder_runtime.cpp @@ -1,4 +1,4 @@ -#include "engine/models/qwen3_asr/audio_encoder.h" +#include "engine/framework/modules/speech_encoders/qwen3_audio_encoder_runtime.h" #include "engine/framework/assets/tensor_source.h" #include "engine/framework/core/backend.h" @@ -21,7 +21,7 @@ #include #include -namespace engine::models::qwen3_asr { +namespace engine::modules { namespace assets = engine::assets; namespace modules = engine::modules; @@ -30,6 +30,24 @@ using Clock = std::chrono::steady_clock; constexpr size_t kDefaultAudioWeightContextBytes = 32ull * 1024ull * 1024ull; +int64_t qwen3_audio_encoder_token_count(int64_t input_frames) { + if (input_frames <= 0) { + throw std::runtime_error("Qwen3 requires positive feature frame count"); + } + const auto floor_div = [](int64_t numerator, int64_t denominator) { + int64_t quotient = numerator / denominator; + const int64_t remainder = numerator % denominator; + if (remainder != 0 && ((remainder < 0) != (denominator < 0))) --quotient; + return quotient; + }; + const int64_t input_lengths_leave = input_frames % 100; + const int64_t feat_lengths = floor_div(input_lengths_leave - 1, 2) + 1; + return floor_div(floor_div(feat_lengths - 1, 2) + 1 - 1, 2) + 1 + + (input_frames / 100) * 13; +} + +namespace qwen3_audio_encoder_detail { + struct GgmlContextDeleter { void operator()(ggml_context * ctx) const noexcept { if (ctx != nullptr) { @@ -40,7 +58,7 @@ struct GgmlContextDeleter { std::vector audio_chunk_lengths(int64_t input_frames, int64_t chunk_frames) { if (input_frames <= 0 || chunk_frames <= 0) { - throw std::runtime_error("Qwen3 ASR audio chunk lengths require positive sizes"); + throw std::runtime_error("Qwen3 audio chunk lengths require positive sizes"); } std::vector chunks; for (int64_t offset = 0; offset < input_frames; offset += chunk_frames) { @@ -51,7 +69,7 @@ std::vector audio_chunk_lengths(int64_t input_frames, int64_t chunk_fra int64_t max_value(const std::vector & values) { if (values.empty()) { - throw std::runtime_error("Qwen3 ASR expected a non-empty length list"); + throw std::runtime_error("Qwen3 expected a non-empty length list"); } return *std::max_element(values.begin(), values.end()); } @@ -66,7 +84,7 @@ int64_t sum_values(const std::vector & values) { std::vector audio_attention_window_lengths(int64_t tokens, int64_t window_tokens) { if (tokens <= 0 || window_tokens <= 0) { - throw std::runtime_error("Qwen3 ASR audio attention window lengths require positive sizes"); + throw std::runtime_error("Qwen3 audio attention window lengths require positive sizes"); } std::vector lengths; for (int64_t offset = 0; offset < tokens; offset += window_tokens) { @@ -87,14 +105,14 @@ std::vector audio_attention_mask(int64_t tokens, const std::vector sinusoidal_positions(int64_t length, int64_t channels) { if (channels % 2 != 0) { - throw std::runtime_error("Qwen3 ASR audio positional embedding requires even channel count"); + throw std::runtime_error("Qwen3 audio positional embedding requires even channel count"); } std::vector table(static_cast(length * channels), 0.0F); const double increment = std::log(10000.0) / static_cast(channels / 2 - 1); @@ -153,7 +171,7 @@ core::TensorValue audio_encoder_layer( core::ModuleBuildContext & ctx, const core::TensorValue & input, const modules::TransformerEncoderBlockWeights & weights, - const Qwen3ASRAudioEncoderConfig & config, + const Qwen3AudioEncoderConfig & config, const core::TensorValue & attention_mask) { const modules::LayerNormModule norm({config.d_model, 1.0e-5F, true, true}); auto attn_in = norm.build(ctx, input, weights.norm1); @@ -200,7 +218,11 @@ struct AudioLayerWeights { core::TensorValue fc2_bias; }; -struct Qwen3ASRAudioEncoderWeights { +} // namespace qwen3_audio_encoder_detail + +using namespace qwen3_audio_encoder_detail; + +struct Qwen3AudioEncoderWeights { std::shared_ptr store; Conv2dWeightsData conv1; Conv2dWeightsData conv2; @@ -214,6 +236,8 @@ struct Qwen3ASRAudioEncoderWeights { core::TensorValue positional_embedding; }; +namespace qwen3_audio_encoder_detail { + Conv2dWeightsData load_conv2d( core::BackendWeightStore & store, const assets::TensorSource & source, @@ -227,24 +251,21 @@ Conv2dWeightsData load_conv2d( }; } -std::shared_ptr load_weights( - const Qwen3ASRAssets & assets, +std::shared_ptr load_weights( + const assets::TensorSource & source, + const Qwen3AudioEncoderConfig & config, + const Qwen3AudioEncoderWeightBinding & binding, ggml_backend_t backend, core::BackendType backend_type, - assets::TensorStorageType storage_type) { - const auto & config = assets.config.audio_encoder; - const auto & source = *assets.model_weights; - const std::string audio_prefix = assets.config.hf_transformers_layout - ? "model.audio_tower" - : "thinker.audio_tower"; - const std::string projector_prefix = assets.config.hf_transformers_layout - ? "model.multi_modal_projector" - : audio_prefix; - auto weights = std::make_shared(); + assets::TensorStorageType storage_type, + const std::string & trace_name) { + const auto & audio_prefix = binding.tower_prefix; + const auto & projector_prefix = binding.projector_prefix; + auto weights = std::make_shared(); auto store = std::make_shared( backend, backend_type, - "qwen3_asr.audio_encoder.weights", + trace_name + ".weights", kDefaultAudioWeightContextBytes); weights->store = store; weights->conv1 = load_conv2d( @@ -298,8 +319,8 @@ std::shared_ptr load_weights( } weights->ln_post_weight = store->load_f32_tensor(source, audio_prefix + ".ln_post.weight", {config.d_model}); weights->ln_post_bias = store->load_f32_tensor(source, audio_prefix + ".ln_post.bias", {config.d_model}); - const std::string proj1 = assets.config.hf_transformers_layout ? ".linear_1" : ".proj1"; - const std::string proj2 = assets.config.hf_transformers_layout ? ".linear_2" : ".proj2"; + const std::string proj1 = "." + binding.projection1; + const std::string proj2 = "." + binding.projection2; weights->proj1 = { store->load_tensor(source, projector_prefix + proj1 + ".weight", storage_type, {config.d_model, config.d_model}), store->load_f32_tensor(source, projector_prefix + proj1 + ".bias", {config.d_model}), @@ -346,53 +367,57 @@ struct GgmlGallocrDeleter { } }; -class Qwen3ASRAudioEncoderGraph { +} // namespace qwen3_audio_encoder_detail + +class Qwen3AudioEncoderGraph { public: - Qwen3ASRAudioEncoderGraph( - std::shared_ptr assets, - std::shared_ptr weights, + Qwen3AudioEncoderGraph( + Qwen3AudioEncoderConfig encoder_config, + std::shared_ptr weights, core::ExecutionContext & execution, size_t graph_arena_bytes, - int64_t frames) - : assets_(std::move(assets)), + int64_t frames, + std::string trace_name) + : config_(std::move(encoder_config)), weights_(std::move(weights)), backend_(execution.backend()), backend_type_(execution.backend_type()), compute_threads_(std::max(1, execution.config().threads)), - frames_(frames) { - if (assets_ == nullptr || weights_ == nullptr) { - throw std::runtime_error("Qwen3 ASR audio encoder graph requires assets and weights"); + frames_(frames), + trace_name_(std::move(trace_name)) { + if (weights_ == nullptr) { + throw std::runtime_error("Qwen3 audio encoder graph requires weights"); } if (backend_ == nullptr) { - throw std::runtime_error("Qwen3 ASR audio encoder backend is not initialized"); + throw std::runtime_error("Qwen3 audio encoder backend is not initialized"); } if (frames_ <= 0) { - throw std::runtime_error("Qwen3 ASR audio encoder graph requires positive frame count"); + throw std::runtime_error("Qwen3 audio encoder graph requires positive frame count"); } const auto build_start = Clock::now(); - const auto & config = assets_->config.audio_encoder; + const auto & config = config_; chunk_frame_limit_ = config.n_window * 2; chunk_lengths_ = audio_chunk_lengths(frames_, chunk_frame_limit_); chunk_frames_ = max_value(chunk_lengths_); chunk_count_ = static_cast(chunk_lengths_.size()); chunk_token_lengths_.reserve(chunk_lengths_.size()); for (const int64_t chunk_length : chunk_lengths_) { - chunk_token_lengths_.push_back(qwen3_asr_audio_encoder_token_count(chunk_length)); + chunk_token_lengths_.push_back(qwen3_audio_encoder_token_count(chunk_length)); } output_tokens_ = sum_values(chunk_token_lengths_); const int64_t max_chunk_tokens = max_value(chunk_token_lengths_); attention_window_tokens_ = max_chunk_tokens * (config.n_window_infer / chunk_frame_limit_); if (output_tokens_ > config.max_source_positions) { - throw std::runtime_error("Qwen3 ASR audio encoder token count exceeds max_source_positions"); + throw std::runtime_error("Qwen3 audio encoder token count exceeds max_source_positions"); } attention_window_lengths_ = audio_attention_window_lengths(output_tokens_, attention_window_tokens_); ggml_init_params params{graph_arena_bytes, nullptr, true}; ctx_.reset(ggml_init(params)); if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize Qwen3 ASR audio encoder graph context"); + throw std::runtime_error("failed to initialize Qwen3 audio encoder graph context"); } - core::ModuleBuildContext ctx{ctx_.get(), "qwen3_asr.audio_encoder", backend_type_}; + core::ModuleBuildContext ctx{ctx_.get(), trace_name_.c_str(), backend_type_}; auto input = core::make_tensor( ctx, GGML_TYPE_F32, @@ -490,28 +515,28 @@ class Qwen3ASRAudioEncoderGraph { }; if (!try_alloc() && (engine::core::trim_backend_pools(backend_), !try_alloc())) { - throw std::runtime_error("failed to allocate Qwen3 ASR audio encoder graph"); + throw std::runtime_error("failed to allocate Qwen3 audio encoder graph"); } ggml_backend_tensor_set(attention_mask_, attention_mask_values_.data(), 0, attention_mask_values_.size() * sizeof(float)); - debug::timing_log_scalar("qwen3_asr.audio_encoder.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); - debug::trace_log_scalar("qwen3_asr.audio_encoder.frames", frames_); + debug::timing_log_scalar(trace_name_ + ".graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar(trace_name_ + ".frames", frames_); } - ~Qwen3ASRAudioEncoderGraph() { + ~Qwen3AudioEncoderGraph() { engine::core::release_backend_graph_resources(backend_, graph_, true); } - bool matches(const Qwen3ASRAudioEncoderWeights & weights, int64_t frames, ggml_backend_t backend, int threads) const { + bool matches(const Qwen3AudioEncoderWeights & weights, int64_t frames, ggml_backend_t backend, int threads) const { return weights_.get() == &weights && frames_ == frames && backend_ == backend && compute_threads_ == std::max(1, threads); } - Qwen3ASRAudioEmbeddings run(const Qwen3ASRAudioFeatures & features) { - const auto & config = assets_->config.audio_encoder; + Qwen3AudioEmbeddings run(const Qwen3AudioFeatures & features) { + const auto & config = config_; if (features.mel_bins != config.num_mel_bins || features.frames != frames_) { - throw std::runtime_error("Qwen3 ASR audio encoder feature shape mismatch"); + throw std::runtime_error("Qwen3 audio encoder feature shape mismatch"); } if (static_cast(features.values.size()) != config.num_mel_bins * frames_) { - throw std::runtime_error("Qwen3 ASR audio encoder feature value count mismatch"); + throw std::runtime_error("Qwen3 audio encoder feature value count mismatch"); } std::vector padded_features(static_cast(chunk_count_ * config.num_mel_bins * chunk_frames_), 0.0F); int64_t source_frame = 0; @@ -530,32 +555,33 @@ class Qwen3ASRAudioEncoderGraph { auto timing_start = Clock::now(); ggml_backend_tensor_set(input_, padded_features.data(), 0, padded_features.size() * sizeof(float)); ggml_backend_tensor_set(attention_mask_, attention_mask_values_.data(), 0, attention_mask_values_.size() * sizeof(float)); - debug::timing_log_scalar("qwen3_asr.audio_encoder.input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + debug::timing_log_scalar(trace_name_ + ".input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); core::set_backend_threads(backend_, compute_threads_); timing_start = Clock::now(); const ggml_status status = engine::core::compute_backend_graph(backend_, graph_); ggml_backend_synchronize(backend_); - debug::timing_log_scalar("qwen3_asr.audio_encoder.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + debug::timing_log_scalar(trace_name_ + ".graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("Qwen3 ASR audio encoder graph compute failed"); + throw std::runtime_error("Qwen3 audio encoder graph compute failed"); } - Qwen3ASRAudioEmbeddings out; + Qwen3AudioEmbeddings out; out.tokens = output_tokens_; out.hidden_size = output_dim_; out.values.resize(static_cast(out.tokens * out.hidden_size)); timing_start = Clock::now(); ggml_backend_tensor_get(output_, out.values.data(), 0, out.values.size() * sizeof(float)); - debug::timing_log_scalar("qwen3_asr.audio_encoder.output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + debug::timing_log_scalar(trace_name_ + ".output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); return out; } private: - std::shared_ptr assets_; - std::shared_ptr weights_; + Qwen3AudioEncoderConfig config_; + std::shared_ptr weights_; ggml_backend_t backend_ = nullptr; core::BackendType backend_type_ = core::BackendType::Cpu; int compute_threads_ = 1; int64_t frames_ = 0; + std::string trace_name_; int64_t chunk_frame_limit_ = 0; int64_t chunk_frames_ = 0; int64_t chunk_count_ = 0; @@ -574,50 +600,55 @@ class Qwen3ASRAudioEncoderGraph { std::unique_ptr, GgmlGallocrDeleter> gallocr_; }; -Qwen3ASRAudioEncoderRuntime::Qwen3ASRAudioEncoderRuntime( - std::shared_ptr assets, +Qwen3AudioEncoderRuntime::Qwen3AudioEncoderRuntime( + std::shared_ptr source, + Qwen3AudioEncoderConfig config, core::ExecutionContext & execution, - size_t graph_arena_bytes, - assets::TensorStorageType weight_storage_type) - : assets_(std::move(assets)), + Qwen3AudioEncoderRuntimeOptions options, + Qwen3AudioEncoderWeightBinding binding) + : config_(std::move(config)), execution_(&execution), - graph_arena_bytes_(graph_arena_bytes) { - if (assets_ == nullptr) { - throw std::runtime_error("Qwen3 ASR audio encoder requires assets"); + options_(std::move(options)) { + if (source == nullptr) { + throw std::runtime_error("Qwen3 audio encoder requires a tensor source"); + } + if (options_.graph_arena_bytes == 0) { + throw std::runtime_error("Qwen3 audio encoder graph arena must be non-zero"); } - if (graph_arena_bytes_ == 0) { - throw std::runtime_error("Qwen3 ASR audio encoder graph arena must be non-zero"); + if (config_.activation_function != "gelu") { + throw std::runtime_error("Qwen3 audio encoder requires GELU activation"); } - weights_ = load_weights(*assets_, execution.backend(), execution.backend_type(), weight_storage_type); + weights_ = load_weights(*source, config_, binding, + execution.backend(), execution.backend_type(), options_.weight_storage_type, options_.trace_name); } -Qwen3ASRAudioEncoderRuntime::~Qwen3ASRAudioEncoderRuntime() = default; +Qwen3AudioEncoderRuntime::~Qwen3AudioEncoderRuntime() = default; -Qwen3ASRAudioEmbeddings Qwen3ASRAudioEncoderRuntime::encode(const Qwen3ASRAudioFeatures & features) { +Qwen3AudioEmbeddings Qwen3AudioEncoderRuntime::encode(const Qwen3AudioFeatures & features) { if (execution_ == nullptr) { - throw std::runtime_error("Qwen3 ASR audio encoder execution context is null"); + throw std::runtime_error("Qwen3 audio encoder execution context is null"); } - if (features.encoder_tokens != qwen3_asr_audio_encoder_token_count(features.frames)) { - throw std::runtime_error("Qwen3 ASR audio encoder token count mismatch"); + if (features.encoder_tokens != qwen3_audio_encoder_token_count(features.frames)) { + throw std::runtime_error("Qwen3 audio encoder token count mismatch"); } const int threads = std::max(1, execution_->config().threads); if (graph_ == nullptr || !graph_->matches(*weights_, features.frames, execution_->backend(), threads)) { graph_.reset(); - graph_ = std::make_unique( - assets_, + graph_ = std::make_unique( + config_, weights_, *execution_, - graph_arena_bytes_, - features.frames); + options_.graph_arena_bytes, + features.frames, options_.trace_name); } else { - debug::timing_log_scalar("qwen3_asr.audio_encoder.graph.build_ms", 0.0); - debug::trace_log_scalar("qwen3_asr.audio_encoder.frames", features.frames); + debug::timing_log_scalar(options_.trace_name + ".graph.build_ms", 0.0); + debug::trace_log_scalar(options_.trace_name + ".frames", features.frames); } auto out = graph_->run(features); if (out.tokens != features.encoder_tokens) { - throw std::runtime_error("Qwen3 ASR audio encoder output token count mismatch"); + throw std::runtime_error("Qwen3 audio encoder output token count mismatch"); } return out; } -} // namespace engine::models::qwen3_asr +} // namespace engine::modules diff --git a/src/framework/modules/transformers/transformer_blocks.cpp b/src/framework/modules/transformers/transformer_blocks.cpp index af28ab593..0e749b49a 100644 --- a/src/framework/modules/transformers/transformer_blocks.cpp +++ b/src/framework/modules/transformers/transformer_blocks.cpp @@ -28,14 +28,17 @@ core::TensorValue TransformerEncoderBlockModule::build( validate_sequence_input(input, config_.hidden_size, "input"); const LayerNormModule norm1(make_norm_config(config_.hidden_size, config_.eps)); - const SelfAttentionModule self_attention({ + AttentionConfig attention_config{ config_.hidden_size, config_.num_heads, config_.use_bias, config_.projection_precision, config_.attention_precision, config_.prefix_cache_layout, - }); + }; + attention_config.use_packed_qkv = config_.use_packed_qkv; + attention_config.use_flash_attention = config_.use_flash_attention; + const SelfAttentionModule self_attention(attention_config); const LayerNormModule norm2(make_norm_config(config_.hidden_size, config_.eps)); const FeedForwardModule feed_forward({ config_.hidden_size, diff --git a/src/models/ast_audioset/runtime.cpp b/src/models/ast_audioset/runtime.cpp index 7b3395f81..af430c72e 100644 --- a/src/models/ast_audioset/runtime.cpp +++ b/src/models/ast_audioset/runtime.cpp @@ -110,33 +110,6 @@ std::shared_ptr load_weights( return result; } -core::TensorValue build_encoder_layer( - core::ModuleBuildContext & context, - const core::TensorValue & input, - const modules::TransformerEncoderBlockWeights & weights, - const Config & config, - bool use_flash_attention) { - const modules::NormConfig norm_config{config.hidden_size, config.layer_norm_eps, true, true, false}; - auto normalized = modules::LayerNormModule(norm_config).build(context, input, weights.norm1); - modules::AttentionConfig attention_config; - attention_config.hidden_size = config.hidden_size; - attention_config.num_heads = config.heads; - attention_config.use_bias = true; - attention_config.use_flash_attention = use_flash_attention; - auto attention = modules::SelfAttentionModule(attention_config).build(context, normalized, weights.self_attention); - auto hidden = modules::ResidualAddModule().build(context, input, attention); - normalized = modules::LayerNormModule(norm_config).build(context, hidden, weights.norm2); - auto feed_forward = modules::FeedForwardModule({ - config.hidden_size, - config.intermediate_size, - true, - modules::GeluApproximation::ExactErf, - GGML_PREC_DEFAULT, - modules::FeedForwardActivation::Gelu, - }).build(context, normalized, weights.feed_forward); - return modules::ResidualAddModule().build(context, hidden, feed_forward); -} - } // namespace struct Runtime::Impl { @@ -185,8 +158,12 @@ struct Runtime::Impl { auto positions = modules::RepeatModule({core::TensorShape::from_dims({1, token_count, config_.hidden_size})}) .build(build, weights_->position_embeddings); hidden = modules::AddModule().build(build, hidden, positions); + modules::TransformerEncoderBlockConfig block_config{ + config_.hidden_size, config_.heads, config_.intermediate_size, config_.layer_norm_eps}; + block_config.use_flash_attention = use_flash_attention; + const modules::TransformerEncoderBlockModule block(block_config); for (const auto & layer : weights_->layers) { - hidden = build_encoder_layer(build, hidden, layer, config_, use_flash_attention); + hidden = block.build(build, hidden, layer); } hidden = modules::LayerNormModule({config_.hidden_size, config_.layer_norm_eps, true, true, false}) .build(build, hidden, weights_->final_norm); diff --git a/src/models/ced/runtime.cpp b/src/models/ced/runtime.cpp index 6ef813e9b..4639c46ab 100644 --- a/src/models/ced/runtime.cpp +++ b/src/models/ced/runtime.cpp @@ -105,20 +105,13 @@ class CedViTGraph { hidden = modules::AddModule().build(ctx, hidden, frequency_positions); hidden = modules::ReshapeModule({core::TensorShape::from_dims({1, c.hidden, frequency * time})}).build(ctx, hidden); hidden = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, hidden); - modules::AttentionConfig attention; - attention.hidden_size = c.hidden; - attention.num_heads = c.heads; - attention.use_packed_qkv = true; - attention.use_flash_attention = flash; + modules::TransformerEncoderBlockConfig block_config{c.hidden, c.heads, c.intermediate, 1e-6f}; + block_config.use_packed_qkv = true; + block_config.use_flash_attention = flash; + const modules::TransformerEncoderBlockModule block(block_config); const modules::LayerNormModule norm({c.hidden, 1e-6f, true, true, false}); for (const auto & layer : w.layers) { - auto normalized = norm.build(ctx, hidden, layer.norm1); - auto attended = modules::SelfAttentionModule(attention).build(ctx, normalized, layer.self_attention); - hidden = modules::ResidualAddModule().build(ctx, hidden, attended); - normalized = norm.build(ctx, hidden, layer.norm2); - auto fed = modules::FeedForwardModule({c.hidden, c.intermediate, true, modules::GeluApproximation::ExactErf}) - .build(ctx, normalized, layer.feed_forward); - hidden = modules::ResidualAddModule().build(ctx, hidden, fed); + hidden = block.build(ctx, hidden, layer); } hidden = norm.build(ctx, hidden, w.norm); output = modules::ReduceMeanModule({1}).build(ctx, hidden); diff --git a/src/models/chatterbox/components/s3gen_weights.h b/src/models/chatterbox/components/s3gen_weights.h deleted file mode 100644 index 82402fd78..000000000 --- a/src/models/chatterbox/components/s3gen_weights.h +++ /dev/null @@ -1,342 +0,0 @@ -#pragma once - -#include "engine/framework/core/backend_weight_store.h" -#include "engine/framework/core/execution_context.h" -#include "engine/framework/core/module.h" - -#include -#include -#include - -namespace engine::models::chatterbox { - -struct S3SpeakerEncoderWeights { - struct BatchNorm1dWeights { - std::vector weight; - std::vector bias; - std::vector running_mean; - std::vector running_var; - bool affine = true; - }; - - struct BatchNorm2dWeights { - std::vector weight; - std::vector bias; - std::vector running_mean; - std::vector running_var; - }; - - struct Conv1dWeights { - std::vector weight; - std::vector bias; - int64_t out_channels = 0; - int64_t in_channels = 0; - int64_t kernel = 0; - int64_t stride = 1; - int64_t padding = 0; - int64_t dilation = 1; - bool use_bias = false; - }; - - struct Conv2dWeights { - std::vector weight; - int64_t out_channels = 0; - int64_t in_channels = 0; - int64_t kernel_h = 0; - int64_t kernel_w = 0; - int64_t stride_h = 1; - int64_t stride_w = 1; - int64_t padding_h = 0; - int64_t padding_w = 0; - }; - - struct BasicResBlockWeights { - Conv2dWeights conv1; - BatchNorm2dWeights bn1; - Conv2dWeights conv2; - BatchNorm2dWeights bn2; - bool use_shortcut = false; - Conv2dWeights shortcut_conv; - BatchNorm2dWeights shortcut_bn; - }; - - struct CAMLayerWeights { - Conv1dWeights linear_local; - Conv1dWeights linear1; - Conv1dWeights linear2; - }; - - struct CAMDenseTDNNLayerWeights { - BatchNorm1dWeights nonlinear1_bn; - Conv1dWeights linear1; - BatchNorm1dWeights nonlinear2_bn; - CAMLayerWeights cam_layer; - }; - - struct CAMDenseTDNNBlockWeights { - std::vector layers; - }; - - struct TransitLayerWeights { - BatchNorm1dWeights nonlinear_bn; - Conv1dWeights linear; - }; - - Conv2dWeights head_conv1; - BatchNorm2dWeights head_bn1; - std::vector head_layer1; - std::vector head_layer2; - Conv2dWeights head_conv2; - BatchNorm2dWeights head_bn2; - - Conv1dWeights tdnn_linear; - BatchNorm1dWeights tdnn_bn; - std::vector blocks; - std::vector transits; - BatchNorm1dWeights out_nonlinear_bn; - Conv1dWeights dense_linear; - BatchNorm1dWeights dense_bn; -}; - -struct S3TokenizerV2Weights { - struct Conv1dWeights { - std::vector weight; - std::vector bias; - int64_t out_channels = 0; - int64_t in_channels = 0; - int64_t kernel = 0; - int64_t stride = 1; - int64_t padding = 0; - }; - - struct LayerNormWeights { - std::vector weight; - std::vector bias; - }; - - struct LinearWeights { - std::vector weight; - std::vector bias; - int64_t out_features = 0; - int64_t in_features = 0; - bool use_bias = false; - }; - - struct BlockWeights { - LayerNormWeights attn_ln; - LinearWeights attn_query; - LinearWeights attn_key; - LinearWeights attn_value; - LinearWeights attn_qkv_packed; - LinearWeights attn_out; - std::vector fsmn_weight; - LayerNormWeights mlp_ln; - LinearWeights mlp_fc1; - LinearWeights mlp_fc2; - }; - - Conv1dWeights conv1; - Conv1dWeights conv2; - std::vector blocks; - LinearWeights quantizer_project_down; -}; - -struct S3FlowEncoderWeights { - struct LayerNormWeights { - engine::core::TensorValue weight_tensor; - engine::core::TensorValue bias_tensor; - }; - struct LinearWeights { - engine::core::TensorValue weight_tensor; - engine::core::TensorValue bias_tensor; - int64_t out_features = 0; - int64_t in_features = 0; - bool use_bias = false; - }; - struct Conv1dWeights { - engine::core::TensorValue weight_tensor; - engine::core::TensorValue bias_tensor; - int64_t out_channels = 0; - int64_t in_channels = 0; - int64_t kernel = 0; - int64_t stride = 1; - int64_t padding = 0; - }; - struct RelativeAttentionWeights { - LinearWeights q; - LinearWeights k; - LinearWeights v; - LinearWeights out; - LinearWeights pos; - engine::core::TensorValue pos_bias_u_tensor; - engine::core::TensorValue pos_bias_v_tensor; - }; - struct FeedForwardWeights { - LinearWeights w1; - LinearWeights w2; - }; - struct EncoderLayerWeights { - LayerNormWeights norm_mha; - RelativeAttentionWeights attn; - LayerNormWeights norm_ff; - FeedForwardWeights ff; - }; - LinearWeights speaker_affine; - LinearWeights encoder_proj; - LinearWeights embed_linear; - LayerNormWeights embed_norm; - Conv1dWeights prelook_conv1; - Conv1dWeights prelook_conv2; - std::vector encoders; - Conv1dWeights up_layer_conv; - LinearWeights up_embed_linear; - LayerNormWeights up_embed_norm; - std::vector up_encoders; - LayerNormWeights after_norm; - engine::core::TensorValue input_embedding_tensor; - const engine::core::ExecutionContext * execution_context = nullptr; - std::shared_ptr store; -}; - -struct S3FlowDecoderWeights { - struct LayerNormWeights { - engine::core::TensorValue weight_tensor; - engine::core::TensorValue bias_tensor; - }; - struct LinearWeights { - engine::core::TensorValue weight_tensor; - engine::core::TensorValue bias_tensor; - int64_t out_features = 0; - int64_t in_features = 0; - bool use_bias = false; - }; - struct Conv1dWeights { - engine::core::TensorValue weight_tensor; - engine::core::TensorValue bias_tensor; - int64_t out_channels = 0; - int64_t in_channels = 0; - int64_t kernel = 0; - int64_t stride = 1; - bool use_bias = false; - }; - struct CausalBlockWeights { - Conv1dWeights conv; - LayerNormWeights norm; - }; - struct ResnetBlockWeights { - LinearWeights time_mlp; - CausalBlockWeights block1; - CausalBlockWeights block2; - Conv1dWeights res_conv; - }; - struct TransformerBlockWeights { - LayerNormWeights norm1; - LinearWeights attn_q; - LinearWeights attn_k; - LinearWeights attn_v; - LinearWeights attn_out; - LayerNormWeights norm3; - LinearWeights ff_proj_in; - LinearWeights ff_proj_out; - }; - struct DownBlockWeights { - ResnetBlockWeights resnet; - std::vector transformers; - Conv1dWeights downsample; - }; - struct MidBlockWeights { - ResnetBlockWeights resnet; - std::vector transformers; - }; - struct UpBlockWeights { - ResnetBlockWeights resnet; - std::vector transformers; - Conv1dWeights upsample; - }; - - LinearWeights time_mlp_1; - LinearWeights time_mlp_2; - std::vector down_blocks; - std::vector mid_blocks; - std::vector up_blocks; - CausalBlockWeights final_block; - Conv1dWeights final_proj; - // Meanflow-distilled decoders (Chatterbox Turbo) mix a second "r" (end-time) sinusoidal - // embedding into the time embedding via this diagonal-init, no-bias linear layer before - // feeding the UNet1D estimator; see decoder.py::get_intmeanflow_time_mixer upstream. Unset - // (weight_tensor.tensor == nullptr) for the base Chatterbox 10-step CFG decoder. - bool meanflow = false; - LinearWeights time_embed_mixer; - const engine::core::ExecutionContext * execution_context = nullptr; - std::shared_ptr store; -}; - -struct S3HiFTVocoderWeights { - struct WeightNormConv1dWeights { - std::vector weight; - std::vector bias; - int64_t out_channels = 0; - int64_t in_channels = 0; - int64_t kernel = 0; - int64_t stride = 1; - int64_t padding = 0; - int64_t dilation = 1; - bool use_bias = false; - }; - - struct WeightNormConvTranspose1dWeights { - std::vector weight; - std::vector bias; - int64_t in_channels = 0; - int64_t out_channels = 0; - int64_t kernel = 0; - int64_t stride = 1; - int64_t padding = 0; - bool use_bias = false; - }; - - struct LinearWeights { - std::vector weight; - std::vector bias; - int64_t out_features = 0; - int64_t in_features = 0; - bool use_bias = false; - }; - - struct SnakeWeights { - std::vector alpha; - bool alpha_logscale = false; - }; - - struct ResBlockWeights { - std::vector convs1; - std::vector convs2; - std::vector activations1; - std::vector activations2; - }; - - struct F0PredictorWeights { - std::vector condnet; - LinearWeights classifier; - }; - - F0PredictorWeights f0_predictor; - WeightNormConv1dWeights conv_pre; - std::vector ups; - std::vector source_downs; - std::vector source_resblocks; - std::vector resblocks; - WeightNormConv1dWeights conv_post; - LinearWeights source_linear; - int64_t sampling_rate = 24000; - int64_t harmonic_num = 8; - float sine_amp = 0.1f; - float noise_std = 0.003f; - float voiced_threshold = 10.0f; - float lrelu_slope = 0.1f; - float audio_limit = 0.99f; - int64_t istft_n_fft = 16; - int64_t istft_hop = 4; -}; - -} // namespace engine::models::chatterbox diff --git a/src/models/chatterbox/hift_vocoder_impl.cpp b/src/models/chatterbox/hift_vocoder_impl.cpp deleted file mode 100644 index 96c5f5ee1..000000000 --- a/src/models/chatterbox/hift_vocoder_impl.cpp +++ /dev/null @@ -1,104 +0,0 @@ -#include "components/component_weights.h" - -#include "engine/framework/modules/vocoders/hift_vocoder.h" - -#include -#include -#include - -namespace engine::models::chatterbox { -namespace { - -engine::modules::HiftVocoderConfig make_chatterbox_hift_config( - engine::assets::TensorStorageType weight_storage_type) { - engine::modules::HiftVocoderConfig config; - config.in_channels = 80; - config.base_channels = 512; - config.nb_harmonics = 8; - config.sampling_rate = 24000; - config.nsf_alpha = 0.1F; - config.nsf_sigma = 0.003F; - config.nsf_voiced_threshold = 10.0F; - config.upsample_rates = {8, 5, 3}; - config.upsample_kernel_sizes = {16, 11, 7}; - config.istft_n_fft = 16; - config.istft_hop = 4; - config.resblock_kernel_sizes = {3, 7, 11}; - config.resblock_dilation_sizes = {{1, 3, 5}, {1, 3, 5}, {1, 3, 5}}; - config.source_resblock_kernel_sizes = {7, 7, 11}; - config.source_resblock_dilation_sizes = {{1, 3, 5}, {1, 3, 5}, {1, 3, 5}}; - config.lrelu_slope = 0.1F; - config.audio_limit = 0.99F; - config.f0_num_class = 1; - config.f0_in_channels = 80; - config.f0_cond_channels = 512; - config.weight_storage_type = weight_storage_type; - config.tensor_prefix = "mel2wav."; - config.weight_layout = engine::modules::HiftVocoderWeightLayout::TorchParametrizedWeightNorm; - return config; -} - -} // namespace - -struct HiFTVocoderComponent::State { - engine::modules::HiftVocoderComponent component; -}; - -HiFTVocoderComponent HiFTVocoderComponent::load_from_source( - std::shared_ptr source, - const engine::core::ExecutionContext & execution_context, - engine::assets::TensorStorageType weight_storage_type) { - auto state = std::make_shared(); - state->component = engine::modules::HiftVocoderComponent::load_from_tensor_source( - std::move(source), - execution_context.config(), - make_chatterbox_hift_config(weight_storage_type)); - - auto weights = std::make_shared(); - weights->runtime_weights = state->component.weights(); - HiFTVocoderComponent component(std::move(weights), execution_context); - component.state_ = std::move(state); - return component; -} - -HiFTVocoderComponent::HiFTVocoderComponent( - std::shared_ptr weights, - const engine::core::ExecutionContext & execution_context) - : weights_(std::move(weights)), execution_context_(&execution_context) {} - -const engine::core::BackendConfig & HiFTVocoderComponent::backend() const noexcept { - return execution_context_->config(); -} - -const std::shared_ptr & HiFTVocoderComponent::weights() const noexcept { - return weights_; -} - -HiFTVocoderOutputs HiFTVocoderComponent::infer( - const std::vector & speech_feat, - int64_t batch, - int64_t frames, - uint64_t seed, - uint64_t prior_noise_values, - const std::vector & cache_source) const { - if (batch != 1) { - throw std::runtime_error("Chatterbox HiFT expects batch 1"); - } - if (!cache_source.empty()) { - throw std::runtime_error("Chatterbox HiFT does not support cache_source"); - } - const auto result = state_->component.synthesize(speech_feat, frames, seed, prior_noise_values); - HiFTVocoderOutputs outputs; - outputs.waveform = result.waveform; - outputs.samples = result.samples; - return outputs; -} - -void HiFTVocoderComponent::release_runtime_cache() const { - if (state_ == nullptr) { - throw std::runtime_error("Chatterbox HiFT is not initialized"); - } - state_->component.release_runtime_cache(); -} - -} // namespace engine::models::chatterbox diff --git a/src/models/chatterbox/s3gen_inference.cpp b/src/models/chatterbox/s3gen_inference.cpp deleted file mode 100644 index 9226ce0dc..000000000 --- a/src/models/chatterbox/s3gen_inference.cpp +++ /dev/null @@ -1,689 +0,0 @@ -#include "engine/models/chatterbox/s3gen_inference.h" -#include "components/component_weights.h" -#include "components/s3gen_weights.h" -#include "engine/framework/core/backend.h" -#include "engine/framework/core/execution_context.h" -#include "engine/framework/core/module.h" -#include "engine/framework/sampling/torch_random.h" - -#include "ggml.h" -#include "ggml-alloc.h" - -#include -#include -#include - -namespace engine::models::chatterbox { -namespace { - -int64_t token_capacity_quantum(const engine::core::BackendConfig & backend) { - (void) backend; - return 1; -} - -int64_t round_up_capacity(int64_t value, int64_t quantum) { - if (value <= 0 || quantum <= 0) { - throw std::runtime_error("S3 capacity rounding requires positive inputs"); - } - return ((value + quantum - 1) / quantum) * quantum; -} - -std::vector make_prefix_mask_bct(int64_t batch, int64_t frames, int64_t valid_frames) { - std::vector mask(static_cast(batch * frames), 0.0f); - for (int64_t batch_index = 0; batch_index < batch; ++batch_index) { - for (int64_t frame = 0; frame < frames; ++frame) { - mask[static_cast(batch_index * frames + frame)] = frame < valid_frames ? 1.0f : 0.0f; - } - } - return mask; -} - -std::vector zero_pad_tokens( - const std::vector & tokens, - int64_t capacity) { - std::vector padded(static_cast(capacity), 0); - std::copy(tokens.begin(), tokens.end(), padded.begin()); - return padded; -} - -std::vector zero_pad_btc( - const std::vector & values, - int64_t valid_frames, - int64_t capacity_frames, - int64_t channels) { - std::vector padded(static_cast(capacity_frames * channels), 0.0f); - for (int64_t frame = 0; frame < valid_frames; ++frame) { - for (int64_t channel = 0; channel < channels; ++channel) { - padded[static_cast(frame * channels + channel)] = - values[static_cast(frame * channels + channel)]; - } - } - return padded; -} - -std::vector slice_bct( - const std::vector & values, - int64_t source_frames, - int64_t frames, - int64_t channels) { - std::vector sliced(static_cast(frames * channels), 0.0f); - for (int64_t channel = 0; channel < channels; ++channel) { - for (int64_t frame = 0; frame < frames; ++frame) { - sliced[static_cast(channel * frames + frame)] = - values[static_cast(channel * source_frames + frame)]; - } - } - return sliced; -} - -void zero_tail_btc( - std::vector & values, - int64_t valid_frames, - int64_t capacity_frames, - int64_t channels) { - if (valid_frames >= capacity_frames) { - return; - } - for (int64_t frame = valid_frames; frame < capacity_frames; ++frame) { - std::fill_n( - values.begin() + static_cast(frame * channels), - static_cast(channels), - 0.0f); - } -} - -std::vector make_zero_padded_conditions_bct( - const EmbedReferenceOutputs & ref_dict, - int64_t total_frames) { - std::vector cond(static_cast(80 * total_frames), 0.0f); - for (int64_t frame = 0; frame < ref_dict.prompt_feat_frames; ++frame) { - for (int64_t dim = 0; dim < ref_dict.prompt_feat_dims; ++dim) { - cond[static_cast(dim * total_frames + frame)] = - ref_dict.prompt_feat[static_cast(frame * ref_dict.prompt_feat_dims + dim)]; - } - } - return cond; -} - -engine::core::TensorValue contiguous( - engine::core::ModuleBuildContext & ctx, - const engine::core::TensorValue & input) { - return engine::core::ensure_backend_addressable_layout(ctx, input); -} - -engine::core::TensorValue add_tensor_values( - engine::core::ModuleBuildContext & ctx, - const engine::core::TensorValue & lhs, - const engine::core::TensorValue & rhs) { - auto lhs_contiguous = contiguous(ctx, lhs); - auto rhs_contiguous = contiguous(ctx, rhs); - return engine::core::wrap_tensor( - ggml_add(ctx.ggml, lhs_contiguous.tensor, rhs_contiguous.tensor), - lhs.shape, - GGML_TYPE_F32); -} - -engine::core::TensorValue linear_rows_graph( - engine::core::ModuleBuildContext & ctx, - const engine::core::TensorValue & input, - const S3FlowEncoderWeights::LinearWeights & weights) { - auto projected = engine::core::wrap_tensor( - ggml_mul_mat(ctx.ggml, weights.weight_tensor.tensor, contiguous(ctx, input).tensor), - engine::core::TensorShape::from_dims({input.shape.dims[0], weights.out_features}), - GGML_TYPE_F32); - if (weights.use_bias) { - auto bias_view = engine::core::reshape_tensor( - ctx, - weights.bias_tensor, - engine::core::TensorShape::from_dims({1, weights.out_features})); - auto repeated = engine::core::wrap_tensor(ggml_repeat(ctx.ggml, bias_view.tensor, projected.tensor), projected.shape, GGML_TYPE_F32); - projected = add_tensor_values(ctx, projected, repeated); - } - return projected; -} - -class S3TokenEmbeddingGraph { -public: - S3TokenEmbeddingGraph( - const S3FlowEncoderWeights & weights, - int64_t token_count, - const engine::core::ExecutionContext & execution_context) - : token_count_(token_count), - execution_context_(&execution_context) { - ggml_init_params params = {}; - params.mem_size = 32ull * 1024ull * 1024ull; - params.mem_buffer = nullptr; - params.no_alloc = true; - ggml_ = ggml_init(params); - if (ggml_ == nullptr) { - throw std::runtime_error("failed to initialize S3 token embedding graph"); - } - - engine::core::ModuleBuildContext ctx = {}; - ctx.ggml = ggml_; - ctx.module_instance_name = "s3_token_embedding"; - ctx.backend_type = execution_context.backend_type(); - token_ids_ = engine::core::make_tensor( - ctx, - GGML_TYPE_I32, - engine::core::TensorShape::from_dims({token_count_})); - embedding_out_ = engine::core::wrap_tensor( - ggml_get_rows(ctx.ggml, weights.input_embedding_tensor.tensor, token_ids_.tensor), - engine::core::TensorShape::from_dims({token_count_, 512}), - GGML_TYPE_F32); - - graph_ = ggml_new_graph_custom(ggml_, 1024, false); - ggml_build_forward_expand(graph_, embedding_out_.tensor); - gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_context_->backend())); - if (gallocr_ == nullptr || - !ggml_gallocr_reserve(gallocr_, graph_) || - !ggml_gallocr_alloc_graph(gallocr_, graph_)) { - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - gallocr_ = nullptr; - } - throw std::runtime_error("failed to allocate S3 token embedding graph tensors"); - } - } - - ~S3TokenEmbeddingGraph() { - if (execution_context_ != nullptr && graph_ != nullptr) { - engine::core::release_backend_graph_resources(execution_context_->backend(), graph_); - } - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - } - if (ggml_ != nullptr) { - ggml_free(ggml_); - } - } - - std::vector run(const std::vector & ids) { - if (static_cast(ids.size()) != token_count_) { - throw std::runtime_error("S3 token embedding graph token count mismatch"); - } - engine::core::write_tensor_i32(token_ids_, ids); - if (engine::core::compute_backend_graph(execution_context_->backend(), graph_) != GGML_STATUS_SUCCESS) { - throw std::runtime_error("S3 token embedding graph compute failed"); - } - return engine::core::read_tensor_f32(embedding_out_.tensor); - } - -private: - int64_t token_count_ = 0; - const engine::core::ExecutionContext * execution_context_ = nullptr; - ggml_context * ggml_ = nullptr; - ggml_gallocr_t gallocr_ = nullptr; - ggml_cgraph * graph_ = nullptr; - engine::core::TensorValue token_ids_; - engine::core::TensorValue embedding_out_; -}; - -class S3Token2MelPrepareGraph { -public: - struct Outputs { - std::vector mu_bct; - std::vector speaker; - }; - - S3Token2MelPrepareGraph( - const S3FlowEncoderWeights & weights, - int64_t frames, - const engine::core::ExecutionContext & execution_context) - : frames_(frames), - execution_context_(&execution_context) { - ggml_init_params params = {}; - params.mem_size = 64ull * 1024ull * 1024ull; - params.mem_buffer = nullptr; - params.no_alloc = true; - ggml_ = ggml_init(params); - if (ggml_ == nullptr) { - throw std::runtime_error("failed to initialize S3 token2mel prepare graph"); - } - - engine::core::ModuleBuildContext ctx = {}; - ctx.ggml = ggml_; - ctx.module_instance_name = "s3_token2mel_prepare"; - ctx.backend_type = execution_context.backend_type(); - encoded_in_ = engine::core::make_tensor( - ctx, - GGML_TYPE_F32, - engine::core::TensorShape::from_dims({frames_, 512})); - speaker_in_ = engine::core::make_tensor( - ctx, - GGML_TYPE_F32, - engine::core::TensorShape::from_dims({1, 192})); - - auto mu_rows = linear_rows_graph(ctx, encoded_in_, weights.encoder_proj); - mu_bct_ = engine::core::wrap_tensor( - ggml_cont(ctx.ggml, ggml_permute(ctx.ggml, mu_rows.tensor, 1, 0, 2, 3)), - engine::core::TensorShape::from_dims({80, frames_}), - GGML_TYPE_F32); - - auto speaker_sq = engine::core::wrap_tensor(ggml_sqr(ctx.ggml, speaker_in_.tensor), speaker_in_.shape, GGML_TYPE_F32); - auto norm_sq = engine::core::wrap_tensor( - ggml_sum_rows(ctx.ggml, speaker_sq.tensor), - engine::core::TensorShape::from_dims({1, 1}), - GGML_TYPE_F32); - auto norm = engine::core::wrap_tensor(ggml_sqrt(ctx.ggml, norm_sq.tensor), norm_sq.shape, GGML_TYPE_F32); - auto norm_rep = engine::core::wrap_tensor(ggml_repeat(ctx.ggml, norm.tensor, speaker_in_.tensor), speaker_in_.shape, GGML_TYPE_F32); - auto speaker_norm = engine::core::wrap_tensor( - ggml_div(ctx.ggml, speaker_in_.tensor, norm_rep.tensor), - speaker_in_.shape, - GGML_TYPE_F32); - speaker_out_ = linear_rows_graph(ctx, speaker_norm, weights.speaker_affine); - - graph_ = ggml_new_graph_custom(ggml_, 4096, false); - ggml_build_forward_expand(graph_, mu_bct_.tensor); - ggml_build_forward_expand(graph_, speaker_out_.tensor); - gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_context_->backend())); - if (gallocr_ == nullptr || - !ggml_gallocr_reserve(gallocr_, graph_) || - !ggml_gallocr_alloc_graph(gallocr_, graph_)) { - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - gallocr_ = nullptr; - } - throw std::runtime_error("failed to allocate S3 token2mel prepare graph tensors"); - } - - } - - ~S3Token2MelPrepareGraph() { - if (execution_context_ != nullptr && graph_ != nullptr) { - engine::core::release_backend_graph_resources(execution_context_->backend(), graph_); - } - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - } - if (ggml_ != nullptr) { - ggml_free(ggml_); - } - } - - Outputs run( - const std::vector & encoded, - const std::vector & speaker_embedding) { - engine::core::write_tensor_f32(encoded_in_, encoded); - engine::core::write_tensor_f32(speaker_in_, speaker_embedding); - if (engine::core::compute_backend_graph(execution_context_->backend(), graph_) != GGML_STATUS_SUCCESS) { - throw std::runtime_error("S3 token2mel prepare graph compute failed"); - } - Outputs outputs; - outputs.mu_bct = engine::core::read_tensor_f32(mu_bct_.tensor); - outputs.speaker = engine::core::read_tensor_f32(speaker_out_.tensor); - return outputs; - } - -private: - int64_t frames_ = 0; - const engine::core::ExecutionContext * execution_context_ = nullptr; - ggml_context * ggml_ = nullptr; - ggml_gallocr_t gallocr_ = nullptr; - ggml_cgraph * graph_ = nullptr; - engine::core::TensorValue encoded_in_; - engine::core::TensorValue speaker_in_; - engine::core::TensorValue mu_bct_; - engine::core::TensorValue speaker_out_; -}; - -std::vector make_gaussian_full_noise(int64_t channels, int64_t frames, uint64_t seed) { - return engine::sampling::generate_torch_cuda_randn( - static_cast(channels * frames), - components::choose_seed(seed), - engine::sampling::TorchRandnPrecision::Float32); -} - -void apply_s3_trim_fade_inplace(std::vector & waveform, int sample_rate) { - const int64_t n_trim = sample_rate / 50; - const int64_t fade_size = 2 * n_trim; - if (waveform.empty() || n_trim <= 0) { - return; - } - const int64_t keep = std::min(static_cast(waveform.size()), fade_size); - for (int64_t i = 0; i < std::min(keep, n_trim); ++i) { - waveform[static_cast(i)] = 0.0f; - } - constexpr double kPi = 3.141592653589793238462643383279502884; - for (int64_t i = n_trim; i < keep; ++i) { - const double alpha = static_cast(i - n_trim) / static_cast(std::max(n_trim - 1, 1)); - const float fade = static_cast((std::cos(kPi * (1.0 - alpha)) + 1.0) / 2.0); - waveform[static_cast(i)] *= fade; - } -} - -} // namespace - -struct S3GenSessionCache::State { - explicit State(engine::core::BackendConfig backend_value) - : backend(std::move(backend_value)), - token_capacity_controller(runtime::GraphCapacityMode::Grow), - flow_cache(backend) {} - - struct PreparedCapacity { - int64_t token_capacity = 0; - int64_t frame_capacity = 0; - std::unique_ptr token_embedding; - std::unique_ptr token2mel_prepare; - }; - - runtime::MappedGraphCapacityAdapter make_token_capacity_adapter( - const S3FlowEncoderWeights & weights, - const engine::core::BackendConfig & backend_config) { - const int64_t quantum = token_capacity_quantum(backend_config); - return runtime::MappedGraphCapacityAdapter( - quantum, - quantum, - [quantum](int64_t request_size) { - return round_up_capacity(request_size, quantum); - }, - [this]() { - std::vector capacities; - if (prepared_capacity.token_embedding && prepared_capacity.token2mel_prepare) { - capacities.push_back(prepared_capacity.token_capacity); - } - return capacities; - }, - [this, &weights](int64_t capacity) { - prepare_token_capacity(weights, capacity); - }); - } - - void prepare_token_capacity( - const S3FlowEncoderWeights & weights, - int64_t token_capacity) { - if (token_capacity <= 0) { - throw std::runtime_error("S3 token capacity must be positive"); - } - if (prepared_capacity.token_embedding && - prepared_capacity.token2mel_prepare && - prepared_capacity.token_capacity == token_capacity) { - return; - } - if (weights.execution_context == nullptr) { - throw std::runtime_error("S3 token2mel prepare requires loaded backend weights"); - } - prepared_capacity.token_capacity = token_capacity; - prepared_capacity.frame_capacity = token_capacity * 2; - prepared_capacity.token_embedding = - std::make_unique(weights, token_capacity, *weights.execution_context); - prepared_capacity.token2mel_prepare = - std::make_unique(weights, prepared_capacity.frame_capacity, *weights.execution_context); - } - - PreparedCapacity & prepared_bundle( - const S3FlowEncoderWeights & weights, - int64_t token_count, - const engine::core::BackendConfig & backend_config) { - auto adapter = make_token_capacity_adapter(weights, backend_config); - token_capacity_controller.ensure_prepared(adapter, token_count); - const int64_t selected_capacity = token_capacity_controller.select_capacity_for_run(adapter, token_count); - if (prepared_capacity.token_capacity != selected_capacity) { - throw std::runtime_error("S3 token capacity controller selected an unprepared capacity"); - } - return prepared_capacity; - } - - void release_pre_cfm_graphs() { - prepared_capacity.token_embedding.reset(); - prepared_capacity.token2mel_prepare.reset(); - flow_cache.release_encoder_graphs(); - } - - void release_cfm_decoder_graphs() { - flow_cache.release_decoder_graphs(); - } - - engine::core::BackendConfig backend; - runtime::GraphCapacityController token_capacity_controller; - S3FlowSessionCache flow_cache; - PreparedCapacity prepared_capacity; -}; - -S3GenSessionCache::S3GenSessionCache(engine::core::BackendConfig backend) - : state_(std::make_unique(backend)) {} - -S3GenSessionCache::~S3GenSessionCache() = default; -S3GenSessionCache::S3GenSessionCache(S3GenSessionCache &&) noexcept = default; -S3GenSessionCache & S3GenSessionCache::operator=(S3GenSessionCache &&) noexcept = default; - -void S3GenSessionCache::release_runtime_graphs() { - state_->release_pre_cfm_graphs(); - state_->release_cfm_decoder_graphs(); -} - -S3Token2MelOutputs compute_s3_token2mel_inference( - S3GenSessionCache & cache, - const S3FlowEncoderWeights & encoder_weights, - const S3FlowDecoderWeights & decoder_weights, - const EmbedReferenceOutputs & ref_dict, - const std::vector & speech_tokens, - int64_t speech_token_count, - int64_t num_steps, - float cfg_rate, - bool cosine_schedule, - const std::vector & full_noise, - uint64_t flow_seed, - engine::core::BackendConfig backend, - S3GenTimingBreakdown * timing) { - constexpr int64_t batch = 1; - constexpr int64_t token_dim = 512; - constexpr int64_t mel_dim = 80; - const int64_t prompt_tokens = ref_dict.prompt_token_count; - const int64_t total_tokens = prompt_tokens + speech_token_count; - std::vector concat_tokens; - concat_tokens.reserve(static_cast(total_tokens)); - concat_tokens.insert(concat_tokens.end(), ref_dict.prompt_tokens.begin(), ref_dict.prompt_tokens.end()); - concat_tokens.insert(concat_tokens.end(), speech_tokens.begin(), speech_tokens.begin() + static_cast(speech_token_count)); - - auto & capacity_bundle = cache.state_->prepared_bundle(encoder_weights, total_tokens, backend); - const int64_t token_capacity = capacity_bundle.token_capacity; - const int64_t frame_capacity = capacity_bundle.frame_capacity; - const auto padded_tokens = zero_pad_tokens(concat_tokens, token_capacity); - - const auto embed_started = std::chrono::steady_clock::now(); - auto token_emb = capacity_bundle.token_embedding->run(padded_tokens); - zero_tail_btc(token_emb, total_tokens, token_capacity, token_dim); - if (timing != nullptr) { - timing->token2mel_embed_ms = - std::chrono::duration(std::chrono::steady_clock::now() - embed_started).count(); - } - - const auto encoder_started = std::chrono::steady_clock::now(); - const auto encoded = - compute_s3_flow_encoder_forward( - cache.state_->flow_cache, - encoder_weights, - token_emb, - total_tokens, - token_capacity, - token_dim, - backend); - std::vector encoded_hidden = encoded.hidden; - if (static_cast(encoded_hidden.size()) != encoded.storage_frames * token_dim) { - encoded_hidden = zero_pad_btc(encoded_hidden, encoded.frames, encoded.storage_frames, token_dim); - } - if (timing != nullptr) { - timing->token2mel_encoder_ms = - std::chrono::duration(std::chrono::steady_clock::now() - encoder_started).count(); - } - - const auto mu_started = std::chrono::steady_clock::now(); - const int64_t total_frames = encoded.frames; - const auto prepared = capacity_bundle.token2mel_prepare->run( - encoded_hidden, - ref_dict.embedding); - auto cond = make_zero_padded_conditions_bct(ref_dict, total_frames); - auto mask = make_prefix_mask_bct(batch, total_frames, total_frames); - auto mu = slice_bct(prepared.mu_bct, frame_capacity, total_frames, mel_dim); - if (timing != nullptr) { - timing->token2mel_mu_ms = - std::chrono::duration(std::chrono::steady_clock::now() - mu_started).count(); - } - - std::vector noise = full_noise.empty() - ? make_gaussian_full_noise(mel_dim, total_frames, flow_seed) - : full_noise; - if (static_cast(noise.size()) != mel_dim * total_frames) { - throw std::runtime_error( - "S3Token2Mel full_noise size mismatch: got=" + - std::to_string(noise.size()) + - " expected=" + - std::to_string(mel_dim * total_frames) + - " total_frames=" + - std::to_string(total_frames) + - " prompt_tokens=" + - std::to_string(prompt_tokens) + - " speech_tokens=" + - std::to_string(speech_token_count)); - } - cache.state_->release_pre_cfm_graphs(); - const auto cfm_started = std::chrono::steady_clock::now(); - // Meanflow-distilled decoders (Chatterbox Turbo) use a plain, non-CFG Euler solve driven by - // two time inputs per step; see S3FlowDecoderWeights::meanflow and - // compute_s3_flow_cfm_meanflow. cfg_rate/cosine_schedule are meaningless for that path and - // are ignored, matching upstream's ConditionalCFM.forward branching on `meanflow`. - const auto mel_full = decoder_weights.meanflow - ? compute_s3_flow_cfm_meanflow( - cache.state_->flow_cache, - decoder_weights, - noise, - mask, - mu, - prepared.speaker, - cond, - batch, - total_frames, - frame_capacity, - num_steps, - backend) - : compute_s3_flow_cfm_euler( - cache.state_->flow_cache, - decoder_weights, - noise, - mask, - mu, - prepared.speaker, - cond, - batch, - total_frames, - frame_capacity, - num_steps, - cfg_rate, - cosine_schedule, - backend, - timing == nullptr ? nullptr : &timing->token2mel_cfm); - if (timing != nullptr) { - timing->token2mel_cfm_ms = - std::chrono::duration(std::chrono::steady_clock::now() - cfm_started).count(); - } - - const int64_t prompt_frames = ref_dict.prompt_feat_frames; - const int64_t output_frames = total_frames - prompt_frames; - S3Token2MelOutputs outputs; - outputs.channels = mel_dim; - outputs.frames = output_frames; - outputs.mel.assign(static_cast(mel_dim * output_frames), 0.0f); - const int64_t mel_stride = mel_full.storage_frames; - for (int64_t dim = 0; dim < mel_dim; ++dim) { - for (int64_t frame = 0; frame < output_frames; ++frame) { - outputs.mel[static_cast(dim * output_frames + frame)] = - mel_full.mel[static_cast(dim * mel_stride + prompt_frames + frame)]; - } - } - cache.state_->release_cfm_decoder_graphs(); - return outputs; -} - -S3GenInferenceOutputs compute_s3gen_inference( - S3GenSessionCache & cache, - const S3FlowEncoderWeights & encoder_weights, - const S3FlowDecoderWeights & decoder_weights, - const engine::models::chatterbox::HiFTVocoderComponent & vocoder, - const EmbedReferenceOutputs & ref_dict, - const std::vector & speech_tokens, - int64_t speech_token_count, - int64_t num_steps, - float cfg_rate, - bool cosine_schedule, - const std::vector & full_noise, - uint64_t flow_seed, - uint64_t vocoder_seed, - engine::core::BackendConfig backend, - S3GenTimingBreakdown * timing) { - const auto token2mel_started = std::chrono::steady_clock::now(); - const auto mel = compute_s3_token2mel_inference( - cache, - encoder_weights, - decoder_weights, - ref_dict, - speech_tokens, - speech_token_count, - num_steps, - cfg_rate, - cosine_schedule, - full_noise, - flow_seed, - backend, - timing); - if (timing != nullptr) { - timing->token2mel_ms = - std::chrono::duration(std::chrono::steady_clock::now() - token2mel_started).count(); - } - const uint64_t prior_noise_values = full_noise.empty() - ? static_cast(mel.channels * (ref_dict.prompt_feat_frames + mel.frames)) - : 0; - const auto vocoder_started = std::chrono::steady_clock::now(); - const auto vocoder_outputs = vocoder.infer( - mel.mel, - 1, - mel.frames, - vocoder_seed, - prior_noise_values, - {}); - if (timing != nullptr) { - timing->vocoder_ms = - std::chrono::duration(std::chrono::steady_clock::now() - vocoder_started).count(); - } - S3GenInferenceOutputs outputs; - outputs.waveform = vocoder_outputs.waveform; - outputs.samples = vocoder_outputs.samples; - outputs.source = vocoder_outputs.source; - outputs.source_channels = vocoder_outputs.source_channels; - outputs.source_frames = vocoder_outputs.source_frames; - outputs.mel = mel.mel; - outputs.mel_channels = mel.channels; - outputs.mel_frames = mel.frames; - apply_s3_trim_fade_inplace(outputs.waveform, 24000); - return outputs; -} - -S3GenInferenceOutputs compute_s3gen_inference( - S3GenSessionCache & cache, - const S3FlowEncoderWeights & encoder_weights, - const S3FlowDecoderWeights & decoder_weights, - const engine::models::chatterbox::HiFTVocoderComponent & vocoder, - const EmbedReferenceOutputs & ref_dict, - const std::vector & speech_tokens, - int64_t speech_token_count, - engine::core::BackendConfig backend, - S3GenTimingBreakdown * timing) { - return compute_s3gen_inference( - cache, - encoder_weights, - decoder_weights, - vocoder, - ref_dict, - speech_tokens, - speech_token_count, - 10, - 0.7f, - true, - {}, - 0, - 0, - backend, - timing); -} - -} // namespace engine::models::chatterbox diff --git a/src/models/chatterbox/session.cpp b/src/models/chatterbox/session.cpp index 88faf0439..aa6a4f662 100644 --- a/src/models/chatterbox/session.cpp +++ b/src/models/chatterbox/session.cpp @@ -5,7 +5,7 @@ #include "engine/framework/runtime/options.h" #include "engine/framework/text/chunking.h" #include "engine/models/chatterbox/components.h" -#include "engine/models/chatterbox/s3gen_flow.h" +#include "engine/framework/codecs/s3gen_runtime.h" #include #include @@ -179,18 +179,10 @@ std::unique_ptr make_chatterbox_component_for_language( *assets.s3gen_weights, execution_context, component_weight_storage_type); - auto flow_encoder = load_s3_flow_encoder_weights( - *assets.s3gen_weights, - execution_context, - component_weight_storage_type); - auto flow_decoder = load_s3_flow_decoder_weights( - *assets.s3gen_weights, - execution_context, - component_weight_storage_type); - auto hift = HiFTVocoderComponent::load_from_source( - assets.s3gen_weights, - execution_context, - component_weight_storage_type); + engine::codecs::s3gen::S3GenConfig s3gen_config; + auto s3gen = std::make_unique( + assets.s3gen_weights, execution_context, s3gen_config, + engine::codecs::s3gen::S3GenRuntimeOptions{component_weight_storage_type}); auto campplus_encoder = CAMPPlusEncoderComponent::load_from_source( assets.s3gen_weights, execution_context, @@ -202,9 +194,7 @@ std::unique_ptr make_chatterbox_component_for_language( std::move(voice_encoder), std::move(s3_tokenizer), std::move(campplus_encoder), - std::move(flow_encoder), - std::move(flow_decoder), - std::move(hift), + std::move(s3gen), make_prompt_prep_config(options), execution_context, mem_saver); @@ -220,18 +210,10 @@ std::unique_ptr make_chatterbox_vc_component( *assets.s3gen_weights, execution_context, component_weight_storage_type); - auto flow_encoder = load_s3_flow_encoder_weights( - *assets.s3gen_weights, - execution_context, - component_weight_storage_type); - auto flow_decoder = load_s3_flow_decoder_weights( - *assets.s3gen_weights, - execution_context, - component_weight_storage_type); - auto hift = HiFTVocoderComponent::load_from_source( - assets.s3gen_weights, - execution_context, - component_weight_storage_type); + engine::codecs::s3gen::S3GenConfig s3gen_config; + auto s3gen = std::make_unique( + assets.s3gen_weights, execution_context, s3gen_config, + engine::codecs::s3gen::S3GenRuntimeOptions{component_weight_storage_type}); auto campplus_encoder = CAMPPlusEncoderComponent::load_from_source( assets.s3gen_weights, execution_context, @@ -240,11 +222,8 @@ std::unique_ptr make_chatterbox_vc_component( return std::make_unique( std::move(s3_tokenizer), std::move(campplus_encoder), - std::move(flow_encoder), - std::move(flow_decoder), - std::move(hift), + std::move(s3gen), make_prompt_prep_config(options), - execution_context, mem_saver); } diff --git a/src/models/chatterbox/tts.cpp b/src/models/chatterbox/tts.cpp index 7df0b67e3..92a7f0372 100644 --- a/src/models/chatterbox/tts.cpp +++ b/src/models/chatterbox/tts.cpp @@ -36,7 +36,9 @@ std::vector clean_generated_speech_tokens_like_python(const std::vector return cleaned; } -void apply_s3_trim_fade_like_inference(std::vector & waveform, int sample_rate) { +} // namespace + +void apply_s3_trim_fade(std::vector & waveform, int sample_rate) { const int64_t n_trim = sample_rate / 50; const int64_t fade_size = 2 * n_trim; if (waveform.empty() || n_trim <= 0) { @@ -54,6 +56,8 @@ void apply_s3_trim_fade_like_inference(std::vector & waveform, int sample } } +namespace { + engine::core::BackendMemorySnapshot capture_backend_memory_snapshot(const engine::core::ExecutionContext * execution_context) { if (!engine::debug::trace_log_enabled() && !engine::debug::timing_log_enabled()) { return {}; @@ -63,22 +67,13 @@ engine::core::BackendMemorySnapshot capture_backend_memory_snapshot(const engine } // namespace -struct ChatterboxTTSComponent::State { - explicit State(engine::core::BackendConfig backend) - : s3_cache(backend) {} - - S3GenSessionCache s3_cache; -}; - ChatterboxTTSComponent::ChatterboxTTSComponent( std::shared_ptr t3_weights, std::shared_ptr tokenizer, engine::models::chatterbox::VoiceEncoderComponent voice_encoder, engine::models::chatterbox::S3TokenizerComponent tokenizer_component, engine::models::chatterbox::CAMPPlusEncoderComponent speaker_encoder, - std::shared_ptr flow_encoder_weights, - std::shared_ptr flow_decoder_weights, - engine::models::chatterbox::HiFTVocoderComponent vocoder, + std::unique_ptr s3gen, ChatterboxPromptPrepConfig prompt_prep_config, const engine::core::ExecutionContext & execution_context, bool mem_saver) @@ -91,18 +86,12 @@ ChatterboxTTSComponent::ChatterboxTTSComponent( std::move(tokenizer_component), std::move(speaker_encoder), prompt_prep_config), - flow_encoder_weights_(std::move(flow_encoder_weights)), - flow_decoder_weights_(std::move(flow_decoder_weights)), - vocoder_(std::move(vocoder)), + s3gen_(std::move(s3gen)), execution_context_(&execution_context), - mem_saver_(mem_saver), - state_(std::make_shared(execution_context.config())) { + mem_saver_(mem_saver) { if (!tokenizer_) { throw std::runtime_error("ChatterboxTtsComponent requires text tokenizer"); } - if (!flow_encoder_weights_ || !flow_decoder_weights_) { - throw std::runtime_error("ChatterboxTtsComponent requires S3 flow weights"); - } } ChatterboxVoiceCloneOutputs ChatterboxTTSComponent::synthesize_voice_clone( @@ -214,13 +203,10 @@ ChatterboxVoiceCloneOutputs ChatterboxTTSComponent::synthesize_voice_clone_impl( const auto s3gen_memory_before = capture_backend_memory_snapshot(execution_context_); const auto s3gen_started = std::chrono::steady_clock::now(); - S3GenTimingBreakdown s3gen_timing; + engine::codecs::s3gen::S3GenTimingBreakdown s3gen_timing; const auto token2mel_memory_before = capture_backend_memory_snapshot(execution_context_); const auto token2mel_started = std::chrono::steady_clock::now(); - const auto mel = compute_s3_token2mel_inference( - state_->s3_cache, - *flow_encoder_weights_, - *flow_decoder_weights_, + const auto mel = s3gen_->token_to_mel( conds.gen, outputs.cleaned_speech_tokens, static_cast(outputs.cleaned_speech_tokens.size()), @@ -229,29 +215,26 @@ ChatterboxVoiceCloneOutputs ChatterboxTTSComponent::synthesize_voice_clone_impl( true, {}, config.seed, - execution_context_ != nullptr ? execution_context_->config() : engine::core::BackendConfig{}, &s3gen_timing); s3gen_timing.token2mel_ms = engine::debug::elapsed_ms(token2mel_started); const auto token2mel_memory_after = capture_backend_memory_snapshot(execution_context_); if (mem_saver_) { - state_->s3_cache.release_runtime_graphs(); + s3gen_->release_flow_graphs(); } const uint64_t prior_noise_values = static_cast(mel.channels * (conds.gen.prompt_feat_frames + mel.frames)); const auto vocoder_memory_before = capture_backend_memory_snapshot(execution_context_); const auto vocoder_started = std::chrono::steady_clock::now(); - const auto vocoder_outputs = vocoder_.infer( + const auto vocoder_outputs = s3gen_->decode_waveform( mel.mel, - 1, mel.frames, config.seed, - prior_noise_values, - {}); + prior_noise_values); s3gen_timing.vocoder_ms = engine::debug::elapsed_ms(vocoder_started); if (mem_saver_) { - vocoder_.release_runtime_cache(); + s3gen_->release_vocoder_graphs(); } const auto vocoder_memory_after = capture_backend_memory_snapshot(execution_context_); outputs.s3gen_ms = @@ -264,15 +247,10 @@ ChatterboxVoiceCloneOutputs ChatterboxTTSComponent::synthesize_voice_clone_impl( } outputs.waveform = vocoder_outputs.waveform; outputs.samples = vocoder_outputs.samples; - outputs.source = vocoder_outputs.source; - outputs.source_channels = vocoder_outputs.source_channels; - outputs.source_frames = vocoder_outputs.source_frames; - outputs.post = vocoder_outputs.post; - outputs.post_frames = vocoder_outputs.post_frames; outputs.mel = mel.mel; outputs.mel_channels = mel.channels; outputs.mel_frames = mel.frames; - apply_s3_trim_fade_like_inference(outputs.waveform, 24000); + apply_s3_trim_fade(outputs.waveform, 24000); if (chatterbox_language_uses_multilingual_t3(language)) { const int64_t speech_token_audio_samples = 24000 / 25; const int64_t effective_speech_tokens = diff --git a/src/models/chatterbox/vc.cpp b/src/models/chatterbox/vc.cpp index 59b76f259..d17aab779 100644 --- a/src/models/chatterbox/vc.cpp +++ b/src/models/chatterbox/vc.cpp @@ -42,11 +42,6 @@ runtime::AudioBuffer trim_mono_audio( } // namespace -struct ChatterboxVCComponent::State { - explicit State(engine::core::BackendConfig backend) : s3_cache(backend) {} - S3GenSessionCache s3_cache; -}; - runtime::AudioBuffer load_chatterbox_vc_audio_mono( const std::string & path, int sample_rate) { @@ -75,24 +70,14 @@ runtime::AudioBuffer normalize_chatterbox_vc_audio_mono( ChatterboxVCComponent::ChatterboxVCComponent( engine::models::chatterbox::S3TokenizerComponent tokenizer_component, engine::models::chatterbox::CAMPPlusEncoderComponent speaker_encoder, - std::shared_ptr flow_encoder_weights, - std::shared_ptr flow_decoder_weights, - engine::models::chatterbox::HiFTVocoderComponent vocoder, + std::unique_ptr s3gen, ChatterboxPromptPrepConfig prompt_prep_config, - const engine::core::ExecutionContext & execution_context, bool mem_saver) : tokenizer_(std::move(tokenizer_component)), speaker_encoder_(std::move(speaker_encoder)), - flow_encoder_weights_(std::move(flow_encoder_weights)), - flow_decoder_weights_(std::move(flow_decoder_weights)), - vocoder_(std::move(vocoder)), + s3gen_(std::move(s3gen)), prompt_prep_config_(prompt_prep_config), - execution_context_(&execution_context), - mem_saver_(mem_saver), - state_(std::make_shared(execution_context.config())) { - if (!flow_encoder_weights_ || !flow_decoder_weights_) { - throw std::runtime_error("ChatterboxVcComponent requires S3 flow weights"); - } + mem_saver_(mem_saver) { } ChatterboxVoiceConversionOutputs ChatterboxVCComponent::convert( @@ -129,12 +114,8 @@ ChatterboxVoiceConversionOutputs ChatterboxVCComponent::convert( outputs.source_speech_tokens = source_tokens.tokens; const auto s3gen_start = std::chrono::steady_clock::now(); - S3GenTimingBreakdown timing; - const auto generated = compute_s3gen_inference( - state_->s3_cache, - *flow_encoder_weights_, - *flow_decoder_weights_, - vocoder_, + engine::codecs::s3gen::S3GenTimingBreakdown timing; + const auto generated = s3gen_->synthesize( target_ref, outputs.source_speech_tokens, outputs.source_speech_token_count, @@ -144,18 +125,18 @@ ChatterboxVoiceConversionOutputs ChatterboxVCComponent::convert( {}, config.seed, config.seed, - execution_context_ != nullptr ? execution_context_->config() : engine::core::BackendConfig{}, &timing); outputs.s3gen_ms = engine::debug::elapsed_ms(s3gen_start); outputs.s3gen_timing = timing; outputs.s3gen_token2mel_ms = timing.token2mel_ms; outputs.s3gen_vocoder_ms = timing.vocoder_ms; outputs.waveform = generated.waveform; + apply_s3_trim_fade(outputs.waveform, 24000); outputs.samples = generated.samples; if (mem_saver_) { - state_->s3_cache.release_runtime_graphs(); - vocoder_.release_runtime_cache(); + s3gen_->release_flow_graphs(); + s3gen_->release_vocoder_graphs(); } engine::debug::timing_log_scalar("chatterbox.vc.target_ref_ms", outputs.target_ref_ms); diff --git a/src/models/firered_vad/runtime.cpp b/src/models/firered_vad/runtime.cpp index 0acd3f46b..7190309ac 100644 --- a/src/models/firered_vad/runtime.cpp +++ b/src/models/firered_vad/runtime.cpp @@ -4,6 +4,7 @@ #include "engine/models/firered_vad/runtime.h" #include "engine/framework/audio/kaldi_fbank.h" +#include "engine/framework/audio/probability_segmenter.h" #include "engine/framework/core/backend_weight_store.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/modules/activation_modules.h" @@ -133,13 +134,14 @@ std::vector segment_probabilities( t = end; } std::vector result; - for (size_t t = 0; t < fixed.size();) { - if (!fixed[t]) { ++t; continue; } - const size_t start = t; - while (t < fixed.size() && fixed[t]) ++t; - const int64_t end = t * kHop + (t == fixed.size() ? kWindow : 0); - result.push_back(make_segment(start * kHop, end, samples)); + audio::ProbabilitySegmenter segmenter({0.5f}); + for (size_t t = 0; t < fixed.size(); ++t) { + if (const auto interval = segmenter.push(static_cast(fixed[t]), t, t + 1)) { + result.push_back(make_segment(interval->start * kHop, interval->end * kHop, samples)); + } } + if (const auto interval = segmenter.finish()) + result.push_back(make_segment(interval->start * kHop, interval->end * kHop + kWindow, samples)); return result; } diff --git a/src/models/heartmula/codec.cpp b/src/models/heartmula/codec.cpp index 26f4651d7..5f28fc7eb 100644 --- a/src/models/heartmula/codec.cpp +++ b/src/models/heartmula/codec.cpp @@ -526,17 +526,6 @@ core::TensorValue causal_conv_transpose1d( return modules::SliceModule({2, 0, output.shape.dims[2] - stride}).build(ctx, output); } -core::TensorValue prelu( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const HeartCodecPreluWeights & weights) { - auto positive = core::wrap_tensor(ggml_relu(ctx.ggml, core::ensure_backend_addressable_layout(ctx, input).tensor), input.shape, GGML_TYPE_F32); - auto negative = core::wrap_tensor(ggml_sub(ctx.ggml, input.tensor, positive.tensor), input.shape, GGML_TYPE_F32); - auto alpha = core::reshape_tensor(ctx, weights.weight, core::TensorShape::from_dims({1, 1, 1})); - auto alpha_rep = core::wrap_tensor(ggml_repeat(ctx.ggml, alpha.tensor, input.tensor), input.shape, GGML_TYPE_F32); - return modules::AddModule{}.build(ctx, positive, modules::MulModule{}.build(ctx, negative, alpha_rep)); -} - core::TensorValue adjacent_repeat_frames_bct( core::ModuleBuildContext & ctx, const core::TensorValue & input, @@ -605,9 +594,9 @@ core::TensorValue scalar_residual_unit( int64_t kernel_size, int64_t dilation) { auto x = causal_conv1d(ctx, input, weights.conv1, channels, channels, kernel_size, dilation); - x = prelu(ctx, x, weights.activation1); + x = modules::PReluModule().build(ctx, x, weights.activation1.weight); x = causal_conv1d(ctx, x, weights.conv2, channels, channels, 1); - x = prelu(ctx, x, weights.activation2); + x = modules::PReluModule().build(ctx, x, weights.activation2.weight); return modules::AddModule{}.build(ctx, input, x); } @@ -659,7 +648,7 @@ core::TensorValue scalar_decoder( config.init_channel, config.init_channel, config.default_kernel_size); - x = prelu(ctx, x, weights.post_processor.activation); + x = modules::PReluModule().build(ctx, x, weights.post_processor.activation.weight); } return causal_conv1d( ctx, diff --git a/src/models/index_echo/assets.cpp b/src/models/index_echo/assets.cpp index f8a7b4630..9f96dd0c4 100644 --- a/src/models/index_echo/assets.cpp +++ b/src/models/index_echo/assets.cpp @@ -73,6 +73,7 @@ std::shared_ptr load_index_echo_assets(const std::filesys auto qwen3_audio = std::make_shared(); qwen3_audio->config.audio_encoder = audio; + qwen3_audio->config.audio_encoder_binding = {"thinker.audio_tower", "thinker.audio_tower"}; qwen3_audio->config.frontend.sample_rate = 16000; qwen3_audio->config.frontend.feature_size = audio.num_mel_bins; qwen3_audio->config.frontend.hop_length = 160; diff --git a/src/models/index_echo/text_translation.cpp b/src/models/index_echo/text_translation.cpp index 8418ef2a6..b0c434c53 100644 --- a/src/models/index_echo/text_translation.cpp +++ b/src/models/index_echo/text_translation.cpp @@ -9,7 +9,7 @@ #include "engine/framework/tokenizers/llama_bpe.h" #include "engine/framework/modules/transformers/qwen35_decoder_runtime.h" #include "engine/models/index_echo/connector.h" -#include "engine/models/qwen3_asr/audio_encoder.h" +#include "engine/framework/modules/speech_encoders/qwen3_audio_encoder_runtime.h" #include #include @@ -103,7 +103,10 @@ class IndexEchoQwen3OmniAuTQwen35TranslationRuntime::Impl { assets::TensorStorageType storage_type) : assets_(std::move(assets)), frontend_({16000, 400, 160, 128, audio::STFTFamily::Default}), - audio_encoder_(assets_->qwen3_omni_audio, execution, 128ull * 1024ull * 1024ull, storage_type), + audio_encoder_(assets_->qwen3_omni_audio->model_weights, + assets_->qwen3_omni_audio->config.audio_encoder, + execution, {128ull * 1024ull * 1024ull, storage_type, "index_echo.audio_encoder"}, + assets_->qwen3_omni_audio->config.audio_encoder_binding), connector_(assets_, execution, storage_type), decoder_config_([&] { auto config = assets_->qwen35_config; @@ -148,7 +151,7 @@ class IndexEchoQwen3OmniAuTQwen35TranslationRuntime::Impl { features.frames = mel.frames; features.mel_bins = mel.mel_bins; features.attention_mask.assign(static_cast(features.frames), 1); - features.encoder_tokens = qwen3_asr::qwen3_asr_audio_encoder_token_count(features.frames); + features.encoder_tokens = modules::qwen3_audio_encoder_token_count(features.frames); const auto encoded = audio_encoder_.encode(features); auto connected = connector_.connect(encoded); if (decoder_config_.round_bf16_activations) { @@ -215,7 +218,7 @@ class IndexEchoQwen3OmniAuTQwen35TranslationRuntime::Impl { private: std::shared_ptr assets_; audio::WhisperLogMelExtractor frontend_; - qwen3_asr::Qwen3ASRAudioEncoderRuntime audio_encoder_; + modules::Qwen3AudioEncoderRuntime audio_encoder_; IndexEchoAudioConnectorRuntime connector_; modules::Qwen35DecoderConfig decoder_config_; modules::Qwen35DecoderRuntime qwen35_; diff --git a/src/models/marblenet_vad/runtime.cpp b/src/models/marblenet_vad/runtime.cpp index 362ff3cd1..91c528d79 100644 --- a/src/models/marblenet_vad/runtime.cpp +++ b/src/models/marblenet_vad/runtime.cpp @@ -1,6 +1,7 @@ #include "engine/models/marblenet_vad/runtime.h" #include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/probability_segmenter.h" #include "engine/framework/audio/dsp.h" #include "engine/framework/core/backend.h" #include "engine/framework/core/backend_weight_store.h" @@ -533,42 +534,20 @@ std::vector decode_segments( int sample_rate, float threshold) { std::vector segments; - bool active = false; - int64_t start_frame = 0; - double confidence_sum = 0.0; - int64_t confidence_count = 0; - for (int64_t frame = 0; frame < inference.frames; ++frame) { - const float probability = speech_probability(inference, frame); - if (probability >= threshold) { - if (!active) { - active = true; - start_frame = frame; - confidence_sum = 0.0; - confidence_count = 0; - } - confidence_sum += probability; - ++confidence_count; - continue; - } - if (!active) { - continue; - } - const int64_t frame_to_samples = weights.config.hop_length * weights.config.output_stride; - runtime::SpeechSegment segment; - segment.span.start_sample = start_frame * frame_to_samples; - segment.span.end_sample = frame * frame_to_samples; - segment.confidence = static_cast(confidence_sum / static_cast(confidence_count)); - segments.push_back(segment); - active = false; - } - if (active) { + audio::ProbabilitySegmenter segmenter({threshold}); + const auto append = [&](const audio::ProbabilitySegment & interval) { const int64_t frame_to_samples = weights.config.hop_length * weights.config.output_stride; runtime::SpeechSegment segment; - segment.span.start_sample = start_frame * frame_to_samples; - segment.span.end_sample = inference.frames * frame_to_samples; - segment.confidence = static_cast(confidence_sum / static_cast(confidence_count)); + segment.span.start_sample = interval.start * frame_to_samples; + segment.span.end_sample = interval.end * frame_to_samples; + segment.confidence = interval.confidence; segments.push_back(segment); + }; + for (int64_t frame = 0; frame < inference.frames; ++frame) { + const float probability = speech_probability(inference, frame); + if (const auto interval = segmenter.push(probability, frame, frame + 1)) append(*interval); } + if (const auto interval = segmenter.finish()) append(*interval); for (auto & segment : segments) { segment.span.start_sample = segment.span.start_sample * sample_rate / weights.config.sample_rate; segment.span.end_sample = segment.span.end_sample * sample_rate / weights.config.sample_rate; diff --git a/src/models/mossformer2/runtime.cpp b/src/models/mossformer2/runtime.cpp index 48995b2bb..dce4d25a9 100644 --- a/src/models/mossformer2/runtime.cpp +++ b/src/models/mossformer2/runtime.cpp @@ -80,7 +80,8 @@ class MossFormer2Graph { x = core::reshape_tensor(ctx_, x, TensorShape::from_dims({c.hidden, frames_})); x = modules::TransposeModule({{1, 0, 2, 3}, 2}).build(ctx_, x); x = modules::AddModule().build(ctx_, x, residual); - x = prelu(x, "mask_net.prelu.weight"); + x = modules::PReluModule({1}) + .build(ctx_, x, weights_.at("mask_net.prelu.weight")); x = linear(x, "mask_net.conv1d_out"); for (int speaker = 0; speaker < c.speakers; ++speaker) { auto mask = modules::SliceModule({1, speaker * c.hidden, c.hidden}).build(ctx_, x); @@ -159,22 +160,6 @@ class MossFormer2Graph { .build(ctx_, x, {weight, bias == weights_.end() ? std::nullopt : std::optional(bias->second)}); } - TensorValue prelu(TensorValue x, const std::string & name) { - x = core::ensure_backend_addressable_layout(ctx_, x); - auto positive = modules::ReluModule().build(ctx_, x); - auto negative = core::wrap_tensor(ggml_scale(ctx_.ggml, x.tensor, -1), x.shape); - negative = modules::ReluModule().build(ctx_, negative); - auto slope_shape = x.shape; - for (size_t axis = 0; axis < slope_shape.rank; ++axis) { - slope_shape.dims[axis] = axis == 1 ? weights_.at(name).shape.num_elements() : 1; - } - auto slope = core::reshape_tensor(ctx_, weights_.at(name), slope_shape); - slope = modules::RepeatModule({x.shape}).build(ctx_, slope); - negative = modules::MulModule().build(ctx_, negative, slope); - negative = core::wrap_tensor(ggml_scale(ctx_.ggml, negative.tensor, -1), x.shape); - return modules::AddModule().build(ctx_, positive, negative); - } - TensorValue feed_forward(TensorValue x, const std::string & name, bool scale_norm) { const auto channels = x.shape.last_dim(); if (scale_norm) { @@ -280,7 +265,8 @@ class MossFormer2Graph { } TensorValue fsmn(TensorValue input, const std::string & name) { - auto x = prelu(linear(input, name + ".conv1.0"), name + ".conv1.1.weight"); + auto x = modules::PReluModule({1}) + .build(ctx_, linear(input, name + ".conv1.0"), weights_.at(name + ".conv1.1.weight")); const auto channels = x.shape.last_dim(); x = modules::LayerNormModule({channels, 1e-5f}).build(ctx_, x, {weights_.at(name + ".norm1.weight"), weights_.at(name + ".norm1.bias")}); @@ -321,7 +307,8 @@ class MossFormer2Graph { dense = modules::GroupNormModule({channels, channels, 1e-5f}).build(ctx_, dense, {weights_.at(base + "norm" + std::to_string(depth + 1) + ".weight"), weights_.at(base + "norm" + std::to_string(depth + 1) + ".bias")}); - dense = prelu(dense, base + "prelu" + std::to_string(depth + 1) + ".weight"); + dense = modules::PReluModule({1}) + .build(ctx_, dense, weights_.at(base + "prelu" + std::to_string(depth + 1) + ".weight")); if (depth == 0) skip = modules::ConcatModule({1}).build(ctx_, dense, skip); } dense = core::reshape_tensor(ctx_, dense, TensorShape::from_dims({channels, frames_})); diff --git a/src/models/owsm/runtime.cpp b/src/models/owsm/runtime.cpp index 22591df7b..6008a109d 100644 --- a/src/models/owsm/runtime.cpp +++ b/src/models/owsm/runtime.cpp @@ -146,23 +146,8 @@ struct OWSMV4Runtime::Graphs { ctx.module_instance_name = "owsm.encoder"; ggml_set_input(features.tensor); ggml_set_input(encoder_positions.tensor); - auto x = core::reshape_tensor(ctx, features, - TensorShape::from_dims({1, 1, config.frontend_frames, 128})); - x = modules::Conv2dModule({1, config.hidden_size, 3, 3, 2, 2, 0, 0, 1, 1, true}) - .build(ctx, x, weights.subsampling.conv0); - x = modules::ReluModule().build(ctx, x); - x = modules::Conv2dModule({config.hidden_size, config.hidden_size, 3, 3, 2, 2, 0, 0, 1, 1, true}) - .build(ctx, x, weights.subsampling.conv1); - x = modules::ReluModule().build(ctx, x); - x = modules::Conv2dModule({config.hidden_size, config.hidden_size, 3, 3, 2, 2, 0, 0, 1, 1, true}) - .build(ctx, x, weights.subsampling.conv2); - x = modules::ReluModule().build(ctx, x); - x = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, x); - x = core::ensure_backend_addressable_layout(ctx, x); - x = core::reshape_tensor(ctx, x, - TensorShape::from_dims({1, config.encoder_frames, config.hidden_size * 15})); - x = modules::LinearModule({config.hidden_size * 15, config.hidden_size, true}) - .build(ctx, x, weights.subsampling.projection); + auto x = modules::EspnetConv2dSubsampling8Module({128, config.hidden_size}) + .build(ctx, features, weights.subsampling); x = modules::LayerScaleModule().build(ctx, x, {weights.embedding_scale}); x = modules::AddModule().build(ctx, x, encoder_positions); for (const auto & layer : weights.encoder) { diff --git a/src/models/owsm_ctc/runtime.cpp b/src/models/owsm_ctc/runtime.cpp index a74d68327..f69d1d8ef 100644 --- a/src/models/owsm_ctc/runtime.cpp +++ b/src/models/owsm_ctc/runtime.cpp @@ -141,40 +141,9 @@ struct OWSMCTCV4EBranchformerRuntime::Graphs { memory = prompt_norm.build(ctx, memory, weights.prompt_norm); memory = modules::LinearModule({p, d, true}).build(ctx, memory, weights.prompt_projection); - auto x = core::reshape_tensor(ctx, features, - TensorShape::from_dims({1, 1, config.frontend_frames, 128})); - int64_t valid_frames = (config.frontend_frames - 3) / 2 + 1; - if (execution.backend_type() == core::BackendType::Cpu) { - // Align the next convolution by padding the small mel input, not its large activation. - const auto next_frames = (valid_frames - 3) / 2 + 1; - const auto aligned_next_frames = ((next_frames + 15) / 16) * 16; - const auto padded_input_frames = 2 * (2 * aligned_next_frames + 1) + 1; - x = modules::Pad2dModule({0, 0, 0, padded_input_frames - config.frontend_frames}).build(ctx, x); - } - x = modules::Conv2dModule({1, d, 3, 3, 2, 2, 0, 0, 1, 1, true}) - .build(ctx, x, weights.subsampling.conv0); - x = modules::ReluModule().build(ctx, x); - for (const auto * conv : {&weights.subsampling.conv1, &weights.subsampling.conv2}) { - const auto output_frames = (valid_frames - 3) / 2 + 1; - const bool align_cpu = execution.backend_type() == core::BackendType::Cpu; - // Align the im2col row count for CPU GEMM; discard only added output rows. - const auto padded_frames = align_cpu ? ((output_frames + 15) / 16) * 16 : output_frames; - const auto pad_frames = std::max(0, 2 * padded_frames + 1 - x.shape.dims[2]); - if (pad_frames > 0) { - x = modules::Pad2dModule({0, 0, 0, pad_frames}) - .build(ctx, x); - } - x = modules::Conv2dModule({d, d, 3, 3, 2, 2, 0, 0, 1, 1, true}).build(ctx, x, *conv); - if (padded_frames != output_frames) { - x = modules::SliceModule({2, 0, output_frames}).build(ctx, x); - } - x = modules::ReluModule().build(ctx, x); - valid_frames = output_frames; - } - x = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, x); - x = core::ensure_backend_addressable_layout(ctx, x); - x = core::reshape_tensor(ctx, x, TensorShape::from_dims({1, config.encoder_frames, d * 15})); - x = modules::LinearModule({d * 15, d, true}).build(ctx, x, weights.subsampling.projection); + auto x = modules::EspnetConv2dSubsampling8Module({128, d, + execution.backend_type() == core::BackendType::Cpu ? 16 : 1}) + .build(ctx, features, weights.subsampling); auto prefix_embedding = embedding.build(ctx, prefix, {weights.embedding}); prefix_embedding = modules::LinearModule({p, d, true}) .build(ctx, prefix_embedding, weights.prefix_projection); diff --git a/src/models/pulsevad/session.cpp b/src/models/pulsevad/session.cpp index 4b1ab16c4..675f1e29e 100644 --- a/src/models/pulsevad/session.cpp +++ b/src/models/pulsevad/session.cpp @@ -1,6 +1,7 @@ #include "engine/models/pulsevad/session.h" #include "engine/models/pulsevad/runtime.h" +#include "engine/framework/audio/probability_segmenter.h" #include "engine/framework/runtime/session_base.h" #include "engine/framework/runtime/options.h" #include "engine/framework/runtime/spec_backed_model.h" @@ -71,8 +72,8 @@ class Session final : public runtime::RuntimeSessionBase, public runtime::IOffli const int64_t length = static_cast(audio.samples.size()); runtime::TaskResult result; std::vector window(3200); - int64_t start = -1; - int64_t last_speech = 0; + engine::audio::ProbabilitySegmenter segmenter( + {threshold, int64_t(min_speech_ms) * 16, int64_t(min_silence_ms) * 16}); const auto append_segment = [&](int64_t begin, int64_t end) { runtime::SpeechSegment segment; segment.span.start_sample = begin; @@ -90,21 +91,10 @@ class Session final : public runtime::RuntimeSessionBase, public runtime::IOffli } break; } - if (probability >= threshold) { - if (start < 0) { - start = offset; - } - last_speech = offset + 3200; - } else if (start >= 0 && offset - last_speech >= int64_t(min_silence_ms) * 16) { - if (last_speech - start >= int64_t(min_speech_ms) * 16) { - append_segment(start, last_speech); - } - start = -1; - } - } - if (start >= 0 && last_speech - start >= int64_t(min_speech_ms) * 16) { - append_segment(start, last_speech); + if (const auto interval = segmenter.push(probability, offset, offset + 3200)) + append_segment(interval->start, interval->end); } + if (const auto interval = segmenter.finish()) append_segment(interval->start, interval->end); debug::timing_log_scalar("session.wall_ms", debug::elapsed_ms(started)); return result; } diff --git a/src/models/qwen3_asr/assets.cpp b/src/models/qwen3_asr/assets.cpp index 7f1315a78..8fdd2a182 100644 --- a/src/models/qwen3_asr/assets.cpp +++ b/src/models/qwen3_asr/assets.cpp @@ -108,6 +108,9 @@ Qwen3ASRConfig parse_config(const assets::ResourceBundle & resources) { Qwen3ASRConfig config; config.hf_transformers_layout = hf_layout; + config.audio_encoder_binding = hf_layout + ? modules::Qwen3AudioEncoderWeightBinding{"model.audio_tower", "model.multi_modal_projector", "linear_1", "linear_2"} + : modules::Qwen3AudioEncoderWeightBinding{"thinker.audio_tower", "thinker.audio_tower"}; config.model_type = json::require_string(root, "model_type"); config.thinker_model_type = json::optional_string(thinker_config, "model_type", config.model_type); config.model_size = json::optional_string(root, "model_size", hf_layout ? "Qwen3-ASR-1.7B-hf" : config.model_type); diff --git a/src/models/qwen3_asr/frontend_whisper.cpp b/src/models/qwen3_asr/frontend_whisper.cpp index fdc708b4b..7f011e305 100644 --- a/src/models/qwen3_asr/frontend_whisper.cpp +++ b/src/models/qwen3_asr/frontend_whisper.cpp @@ -84,7 +84,7 @@ Qwen3ASRAudioFeatures Qwen3ASRWhisperFrontend::extract(const runtime::AudioBuffe out.attention_mask.assign(static_cast(features.frames), 1); out.mel_bins = features.mel_bins; out.frames = features.frames; - out.encoder_tokens = qwen3_asr_audio_encoder_token_count(out.frames); + out.encoder_tokens = modules::qwen3_audio_encoder_token_count(out.frames); debug::timing_log_scalar("qwen3_asr.frontend.normalize_ms", engine::debug::elapsed_ms(normalize_start, normalize_end)); debug::timing_log_scalar("qwen3_asr.frontend.log_mel_ms", engine::debug::elapsed_ms(feature_start, feature_end)); return out; diff --git a/src/models/qwen3_asr/session.cpp b/src/models/qwen3_asr/session.cpp index 3e30f2204..86fcd49d1 100644 --- a/src/models/qwen3_asr/session.cpp +++ b/src/models/qwen3_asr/session.cpp @@ -197,7 +197,9 @@ Qwen3ASRSession::Qwen3ASRSession( option_weight_type(options, {"qwen3_asr.weight_type"}, engine::assets::TensorStorageType::Native))), tokenizer_(assets_), frontend_(assets_), - audio_encoder_(assets_, execution_context(), audio_encoder_graph_arena_bytes_, audio_encoder_weight_storage_type_), + audio_encoder_(assets_->model_weights, assets_->config.audio_encoder, + execution_context(), {audio_encoder_graph_arena_bytes_, audio_encoder_weight_storage_type_, "qwen3_asr.audio_encoder"}, + assets_->config.audio_encoder_binding), thinker_( assets_, execution_context(), diff --git a/src/models/qwen3_forced_aligner/session.cpp b/src/models/qwen3_forced_aligner/session.cpp index 6adc2494d..f831d231e 100644 --- a/src/models/qwen3_forced_aligner/session.cpp +++ b/src/models/qwen3_forced_aligner/session.cpp @@ -108,7 +108,9 @@ Qwen3ForcedAlignerSession::Qwen3ForcedAlignerSession( option_weight_type(options, "qwen3_forced_aligner.weight_type", engine::assets::TensorStorageType::Native))), tokenizer_(assets_), frontend_(assets_), - audio_encoder_(assets_, execution_context(), audio_encoder_graph_arena_bytes_, audio_encoder_weight_storage_type_), + audio_encoder_(assets_->model_weights, assets_->config.audio_encoder, + execution_context(), {audio_encoder_graph_arena_bytes_, audio_encoder_weight_storage_type_, "qwen3_forced_aligner.audio_encoder"}, + assets_->config.audio_encoder_binding), thinker_( assets_, execution_context(), diff --git a/src/models/supertonic/runtime.cpp b/src/models/supertonic/runtime.cpp index ce87ed3f4..1d7f68c23 100644 --- a/src/models/supertonic/runtime.cpp +++ b/src/models/supertonic/runtime.cpp @@ -500,18 +500,6 @@ core::TensorValue edge_pad_time( return out; } -core::TensorValue prelu( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const core::TensorValue & slope) { - auto x = core::ensure_backend_addressable_layout(ctx, input); - auto positive = core::wrap_tensor(ggml_relu(ctx.ggml, x.tensor), input.shape, GGML_TYPE_F32); - auto negative = core::wrap_tensor(ggml_sub(ctx.ggml, x.tensor, positive.tensor), input.shape, GGML_TYPE_F32); - auto slope_repeated = broadcast_to(ctx, slope, input.shape); - auto scaled_negative = core::wrap_tensor(ggml_mul(ctx.ggml, negative.tensor, slope_repeated.tensor), input.shape, GGML_TYPE_F32); - return core::wrap_tensor(ggml_add(ctx.ggml, positive.tensor, scaled_negative.tensor), input.shape, GGML_TYPE_F32); -} - core::TensorValue batch_norm_eval( core::ModuleBuildContext & ctx, const core::TensorValue & input, @@ -749,7 +737,7 @@ class SupertonicVocoderRuntime { 1, true, }).build(ctx, x, {weights_.head1_weight, weights_.head1_bias}); - x = prelu(ctx, x, weights_.head_prelu_slope); + x = modules::PReluModule().build(ctx, x, weights_.head_prelu_slope); x = modules::Conv1dModule({ x.shape.dims[1], weights_.head2_weight.shape.dims[0], @@ -832,7 +820,8 @@ class SupertonicNetwork { ctx_, joined, {weight("duration_predictor.tts.dp.predictor.layers.0.weight"), weight("duration_predictor.tts.dp.predictor.layers.0.bias")}); - h = prelu(ctx_, h, weight("duration_predictor.tts.dp.predictor.activation.weight")); + h = modules::PReluModule().build(ctx_, h, + weight("duration_predictor.tts.dp.predictor.activation.weight")); auto duration = modules::LinearModule({ h.shape.last_dim(), weight("duration_predictor.tts.dp.predictor.layers.1.weight").shape.dims[0], diff --git a/src/models/tf_gridnet/runtime.cpp b/src/models/tf_gridnet/runtime.cpp index 34d883498..0c9746e0d 100644 --- a/src/models/tf_gridnet/runtime.cpp +++ b/src/models/tf_gridnet/runtime.cpp @@ -4,6 +4,7 @@ #include "engine/framework/core/backend_weight_store.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/attention/scaled_dot_product_attention.h" #include "engine/framework/modules/conv_modules.h" #include "engine/framework/modules/linear_module.h" #include "engine/framework/modules/norm_modules.h" @@ -58,17 +59,12 @@ class TFGridNetBlock { q = core::ensure_backend_addressable_layout(ctx, q); k = core::ensure_backend_addressable_layout(ctx, k); v = core::ensure_backend_addressable_layout(ctx, v); - q = core::reshape_tensor(ctx, q, TensorShape::from_dims({frames, q.shape.dims[2] * bins})); - k = core::reshape_tensor(ctx, k, TensorShape::from_dims({frames, k.shape.dims[2] * bins})); - v = core::reshape_tensor(ctx, v, TensorShape::from_dims({frames, value_channels * bins})); - // Q/K and V widths differ; the shared SDPA module requires equal widths. - auto scores = modules::MatMulModule().build(ctx, q, - modules::TransposeModule({{1, 0, 2, 3}, 2}).build(ctx, k)); - scores = modules::MulModule().build(ctx, scores, - modules::RepeatModule({scores.shape}).build(ctx, - core::reshape_tensor(ctx, weights_.at("attention_scale"), TensorShape::from_dims({1, 1})))); - auto probabilities = modules::SoftmaxModule().build(ctx, scores); - auto context = modules::MatMulModule().build(ctx, probabilities, v); + q = core::reshape_tensor(ctx, q, TensorShape::from_dims({1, 1, frames, q.shape.dims[2] * bins})); + k = core::reshape_tensor(ctx, k, TensorShape::from_dims({1, 1, frames, k.shape.dims[2] * bins})); + v = core::reshape_tensor(ctx, v, TensorShape::from_dims({1, 1, frames, value_channels * bins})); + auto context = modules::ScaledDotProductAttentionModule({q.shape.dims[3], + modules::ScaledDotProductAttentionLowering::Explicit, GGML_PREC_DEFAULT}) + .build(ctx, q, k, v, std::nullopt, weights_.at("attention_scale")); context = core::reshape_tensor(ctx, context, TensorShape::from_dims({1, frames, value_channels, bins})); heads.push_back(modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, context)); } @@ -153,15 +149,7 @@ class TFGridNetBlock { const int64_t channels = weight.shape.dims[0]; x = modules::Conv2dModule({x.shape.dims[1], channels, 1, 1}).build(ctx, x, {weight, weights_.at(prefix + ".0.bias")}); - const auto positive = modules::ReluModule().build(ctx, x); - auto negative = modules::MulModule().build(ctx, x, - modules::RepeatModule({x.shape}).build(ctx, - core::reshape_tensor(ctx, weights_.at("negative_one"), TensorShape::from_dims({1, 1, 1, 1})))); - negative = modules::ReluModule().build(ctx, negative); - auto slope = modules::MulModule().build(ctx, weights_.at(prefix + ".1.weight"), weights_.at("negative_one")); - negative = modules::MulModule().build(ctx, negative, modules::RepeatModule({x.shape}).build(ctx, - core::reshape_tensor(ctx, slope, TensorShape::from_dims({1, 1, 1, 1})))); - x = modules::AddModule().build(ctx, positive, negative); + x = modules::PReluModule({1}).build(ctx, x, weights_.at(prefix + ".1.weight")); const auto shape = x.shape; x = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, x); const int64_t features = channels * shape.dims[3]; @@ -251,7 +239,6 @@ class TFGridNetBiLSTMAttentionRuntime::Impl { } weights_.emplace("output_conv.weight", store_.make_f32(TensorShape::from_dims({2 * c.speakers, c.embedding_dim, 3, 3}), reversed)); weights_.emplace("zero", store_.make_f32(TensorShape::from_dims({1}), {0})); - weights_.emplace("negative_one", store_.make_f32(TensorShape::from_dims({1}), {-1})); const int bins = c.n_fft / 2 + 1; const int qk = (c.attention_qk_dim + bins - 1) / bins * bins; weights_.emplace("attention_scale", store_.make_f32(TensorShape::from_dims({1}), {1.0f / std::sqrt(float(qk))})); diff --git a/tests/kitten_tts2/components_probe.cpp b/tests/kitten_tts2/components_probe.cpp index c250638b7..496fcbb67 100644 --- a/tests/kitten_tts2/components_probe.cpp +++ b/tests/kitten_tts2/components_probe.cpp @@ -1,4 +1,4 @@ -#include "engine/models/chatterbox/s3gen_inference.h" +#include "engine/framework/codecs/s3gen_runtime.h" #include "engine/community_models/kitten_tts2/speaker.h" #include "engine/framework/assets/tensor_source.h" #include "engine/framework/audio/conversion.h" @@ -13,7 +13,7 @@ int main(int argc, char ** argv) { try { if (argc < 5 || argc > 7) throw std::runtime_error("usage: kitten_tts2_components_probe S3GEN VOICES_JSON CODES_JSON OUTPUT_DIR [cpu|cuda [REFERENCE_MODEL_DIR]]"); - namespace s3 = engine::models::chatterbox; + namespace s3 = engine::codecs::s3gen; namespace json = engine::io::json; engine::core::BackendConfig config; config.type = engine::core::BackendType::Cpu; @@ -52,11 +52,12 @@ int main(int argc, char ** argv) { engine::audio::write_wav(std::filesystem::path(argv[4]) / "speaker_input.wav", 16000, 1, input, options); } auto source = engine::assets::open_tensor_source(argv[1]); - auto encoder = s3::load_s3_flow_encoder_weights(*source, execution, engine::assets::TensorStorageType::F32); - auto decoder = s3::load_s3_flow_decoder_weights(*source, execution, engine::assets::TensorStorageType::F32); + s3::S3GenConfig s3gen_config; + s3gen_config.decoder.variant = s3::S3FlowVariant::MeanFlow; + s3::S3GenRuntime runtime(source, execution, s3gen_config, {engine::assets::TensorStorageType::F32}); const auto all = json::parse_file(argv[2]); const auto & voice = all.require("Bruno"); - s3::EmbedReferenceOutputs ref; + s3::S3GenConditioning ref; for (const auto & x : voice.require("prompt_token").as_array()[0].as_array()) ref.prompt_tokens.push_back(x.as_number()); for (const auto & row : voice.require("prompt_feat").as_array()[0].as_array()) for (const auto & x : row.as_array()) ref.prompt_feat.push_back(x.as_number()); @@ -70,27 +71,22 @@ int main(int argc, char ** argv) { for (const auto & x : codes.as_array()) tokens.push_back(x.as_number()); tokens.insert(tokens.end(), 3, 4299); const int64_t frames = 2 * (ref.prompt_token_count + tokens.size()); - s3::S3GenSessionCache cache(config); - std::filesystem::create_directories(argv[4]); - // Isolate the encoder from the meanflow solver when diagnosing parity. + // Isolate the encoder through the runtime's stage API. auto table = source->require_f32("flow.input_embedding.weight", {6561,512}); auto all_tokens = ref.prompt_tokens; all_tokens.insert(all_tokens.end(), tokens.begin(), tokens.end()); std::vector embedded; for (auto token : all_tokens) for (int i = 0; i < 512; ++i) embedded.push_back(table[token*512+i]); - s3::S3FlowSessionCache flow_cache(config); - auto encoded = s3::compute_s3_flow_encoder_forward(flow_cache, *encoder, embedded, - all_tokens.size(), all_tokens.size(), 512, config); + auto encoded = runtime.encode(embedded, all_tokens.size(), all_tokens.size(), 512); std::ofstream hidden(std::filesystem::path(argv[4])/"encoder.f32", std::ios::binary); hidden.write(reinterpret_cast(encoded.hidden.data()), encoded.hidden.size()*sizeof(float)); - auto mel = s3::compute_s3_token2mel_inference(cache, *encoder, *decoder, ref, tokens, tokens.size(), - 2, 0, false, std::vector(80*frames, 0), 0, config); + auto mel = runtime.token_to_mel(ref, tokens, tokens.size(), + 2, 0, false, std::vector(80*frames, 0), 0); std::filesystem::create_directories(argv[4]); std::ofstream out(std::filesystem::path(argv[4])/"zero_noise_mel.f32", std::ios::binary); out.write(reinterpret_cast(mel.mel.data()), mel.mel.size()*sizeof(float)); - auto wrapper = s3::HiFTVocoderComponent::load_from_source(source, execution, engine::assets::TensorStorageType::F32); - engine::modules::HiftVocoderComponent vocoder(wrapper.weights()->runtime_weights, config); + const auto & vocoder = runtime.vocoder(); auto f0 = vocoder.predict_f0(mel.mel, mel.frames); std::vector random(9 + 9*480*mel.frames, 0); auto audio = vocoder.synthesize(mel.mel, mel.frames, 0, 0, &random); diff --git a/tests/unittests/test_audio_dsp.cpp b/tests/unittests/test_audio_dsp.cpp index fcfb70898..6dd5ebb98 100644 --- a/tests/unittests/test_audio_dsp.cpp +++ b/tests/unittests/test_audio_dsp.cpp @@ -1,4 +1,5 @@ #include "engine/framework/audio/dsp.h" +#include "engine/framework/audio/probability_segmenter.h" #include #include @@ -378,10 +379,44 @@ void test_istft_matches_reference_across_configs_and_variants() { } } +void test_probability_segments() { + using engine::audio::ProbabilitySegmenter; + ProbabilitySegmenter frames({0.5f}); + require(!frames.push(0.1f, 0, 1), "leading silence"); + require(!frames.push(0.5f, 1, 2), "inclusive threshold"); + require(!frames.push(1.0f, 2, 3), "contiguous speech"); + const auto first = frames.push(0.0f, 3, 4); + require(first && first->start == 1 && first->end == 3 && first->confidence == 0.75f, + "frame segment and confidence"); + require(!frames.finish(), "finished segment must not repeat"); + require(!frames.push(1.0f, 4, 5), "next segment"); + const auto last = frames.finish(); + require(last && last->start == 4 && last->end == 5 && last->confidence == 1.0f, + "finish resets confidence and boundaries"); + + ProbabilitySegmenter windows({0.5f, 3200, 800}); + require(!windows.push(1.0f, 0, 3200), "positive window"); + require(!windows.push(0.0f, 1600, 4800), "overlap is not silence duration"); + require(!windows.push(0.0f, 3200, 6400), "silence shorter than minimum"); + const auto window = windows.push(0.0f, 4000, 7200); + require(window && window->start == 0 && window->end == 3200, + "end is last positive window end"); + + ProbabilitySegmenter minimum({0.5f, 2, 0}); + require(!minimum.push(1.0f, 0, 1), "short segment begins"); + require(!minimum.push(0.0f, 1, 2), "short segment discarded"); + require(!minimum.push(0.5f, 2, 3), "reuse after discard"); + require(!minimum.push(0.5f, 3, 4), "minimum duration reached"); + const auto kept = minimum.finish(); + require(kept && kept->start == 2 && kept->end == 4 && kept->confidence == 0.5f, + "minimum duration is inclusive"); +} + } // namespace int main() { try { + test_probability_segments(); test_log_mel_matches_reference_pipeline(); test_hann_window_modes_and_cache(); test_stft_istft_round_trip_matches_waveform(); diff --git a/tests/unittests/test_encoder_modules.cpp b/tests/unittests/test_encoder_modules.cpp index 00e3040a6..02cc38341 100644 --- a/tests/unittests/test_encoder_modules.cpp +++ b/tests/unittests/test_encoder_modules.cpp @@ -1,6 +1,7 @@ #include "engine/framework/core/backend.h" #include "engine/framework/core/backend_weight_store.h" #include "engine/framework/modules/attention_modules.h" +#include "engine/framework/modules/activation_modules.h" #include "engine/framework/modules/conformer_modules.h" #include "engine/framework/modules/primitive_modules.h" #include "engine/framework/modules/weight_binding.h" @@ -1086,10 +1087,38 @@ void test_cached_cross_flash_masks() { require_allclose(runner.run_f32(per_head), runner.run_f32(expected), 1e-5f, "per-head cached flash mask"); } +void test_prelu_broadcasting() { + using namespace engine; + for (size_t axis = 0; axis < 3; ++axis) { + for (bool scalar : {false, true}) { + CpuModuleRunner runner; + const auto shape = core::TensorShape::from_dims({2, 3, 4}); + const auto input = runner.make_f32(shape); + const auto count = scalar ? 1 : shape.dims[axis]; + const auto slope = runner.make_f32(core::TensorShape::from_dims({count})); + const auto output = modules::PReluModule({axis}).build(runner.ctx, input, slope); + runner.allocate_tensors(); + std::vector values(24), slopes(count), expected(24); + for (int64_t c = 0; c < count; ++c) slopes[c] = (c - 1) * 0.25f; + int64_t stride = 1; + for (size_t dim = axis + 1; dim < 3; ++dim) stride *= shape.dims[dim]; + for (size_t i = 0; i < values.size(); ++i) { + values[i] = static_cast(i % 9) - 4.0f; + const auto channel = scalar ? 0 : (i / stride) % count; + expected[i] = values[i] < 0 ? values[i] * slopes[channel] : values[i]; + } + core::write_tensor_f32(input, values); + core::write_tensor_f32(slope, slopes); + require_allclose(runner.run_f32(output), expected, 0.0f, "PReLU scalar/channel broadcasting"); + } + } +} + } // namespace int main() { try { + test_prelu_broadcasting(); test_glu_contiguous_gate_opt_in(); test_depthwise_subsampling_stage_masks(); test_batch_norm_eval_binding(); diff --git a/tests/unittests/test_scaled_dot_product_attention.cpp b/tests/unittests/test_scaled_dot_product_attention.cpp index 9dd9d5472..04c72a613 100644 --- a/tests/unittests/test_scaled_dot_product_attention.cpp +++ b/tests/unittests/test_scaled_dot_product_attention.cpp @@ -30,6 +30,7 @@ struct SdpaCase { int64_t key_steps; int64_t head_dim; bool causal_mask; + int64_t value_dim = 0; }; struct RunResult { @@ -170,7 +171,7 @@ class SdpaRunner { test_case.batch, test_case.heads, test_case.key_steps, - test_case.head_dim, + test_case.value_dim > 0 ? test_case.value_dim : test_case.head_dim, })); if (use_attention_mask_) { mask_ = engine::core::make_tensor( @@ -475,7 +476,8 @@ void run_cpu_case(const SdpaCase & test_case) { 1.17F + static_cast(round) * 0.37F, 0.16F); const auto v = make_patterned_f32( - static_cast(test_case.batch * test_case.heads * test_case.key_steps * test_case.head_dim), + static_cast(test_case.batch * test_case.heads * test_case.key_steps * + (test_case.value_dim > 0 ? test_case.value_dim : test_case.head_dim)), 2.03F + static_cast(round) * 0.19F, 0.12F); const std::vector no_mask; @@ -505,6 +507,7 @@ void run_cpu_case(const SdpaCase & test_case) { } // namespace int main() try { + run_cpu_case({"unequal_value_width", 2, 3, 19, 23, 16, false, 24}); const bool has_cuda = backend_available(engine::core::BackendType::Cuda); const auto backend_type = has_cuda ? engine::core::BackendType::Cuda : engine::core::BackendType::Cpu; run_manual_repeat_parity_cases(backend_type);