From 1a605dcf764c220e917592878ac6e3cd535dd631 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:01:36 -0400 Subject: [PATCH 01/10] Add BreezeTTS and CosyVoice3 model wiring --- CMakeLists.txt | 42 + .../modules/text_encoders/t5_gemma_encoder.h | 14 + .../engine/framework/tokenizers/llama_bpe.h | 7 +- include/engine/models/breeze_tts/assets.h | 76 ++ include/engine/models/breeze_tts/generator.h | 49 + include/engine/models/breeze_tts/session.h | 64 + .../engine/models/breeze_tts/speech_decoder.h | 61 + .../engine/models/breeze_tts/speech_encoder.h | 53 + .../engine/models/breeze_tts/text_encoder.h | 36 + .../engine/models/breeze_tts/tokenizer_text.h | 44 + include/engine/models/cosyvoice3/ar.h | 48 + include/engine/models/cosyvoice3/assets.h | 51 + include/engine/models/cosyvoice3/flow.h | 48 + include/engine/models/cosyvoice3/frontend.h | 42 + include/engine/models/cosyvoice3/hift.h | 35 + include/engine/models/cosyvoice3/session.h | 49 + .../engine/models/cosyvoice3/tokenizer_text.h | 32 + .../text_encoders/t5_gemma_encoder.cpp | 96 +- .../qwen_causal_decode_runtime.cpp | 8 - src/framework/tokenizers/llama_bpe.cpp | 22 +- src/models/breeze_tts/assets.cpp | 177 +++ src/models/breeze_tts/generator.cpp | 959 ++++++++++++++ src/models/breeze_tts/session.cpp | 251 ++++ src/models/breeze_tts/speech_decoder.cpp | 1175 +++++++++++++++++ src/models/breeze_tts/speech_encoder.cpp | 649 +++++++++ src/models/breeze_tts/text_encoder.cpp | 282 ++++ src/models/breeze_tts/tokenizer_text.cpp | 136 ++ src/models/cosyvoice3/ar.cpp | 572 ++++++++ src/models/cosyvoice3/assets.cpp | 69 + src/models/cosyvoice3/flow.cpp | 917 +++++++++++++ src/models/cosyvoice3/frontend.cpp | 285 ++++ src/models/cosyvoice3/hift.cpp | 113 ++ src/models/cosyvoice3/session.cpp | 329 +++++ src/models/cosyvoice3/tokenizer_text.cpp | 363 +++++ .../audiocpp_cli/audiocpp_cli_path_cases.json | 126 ++ 35 files changed, 7252 insertions(+), 28 deletions(-) create mode 100644 include/engine/models/breeze_tts/assets.h create mode 100644 include/engine/models/breeze_tts/generator.h create mode 100644 include/engine/models/breeze_tts/session.h create mode 100644 include/engine/models/breeze_tts/speech_decoder.h create mode 100644 include/engine/models/breeze_tts/speech_encoder.h create mode 100644 include/engine/models/breeze_tts/text_encoder.h create mode 100644 include/engine/models/breeze_tts/tokenizer_text.h create mode 100644 include/engine/models/cosyvoice3/ar.h create mode 100644 include/engine/models/cosyvoice3/assets.h create mode 100644 include/engine/models/cosyvoice3/flow.h create mode 100644 include/engine/models/cosyvoice3/frontend.h create mode 100644 include/engine/models/cosyvoice3/hift.h create mode 100644 include/engine/models/cosyvoice3/session.h create mode 100644 include/engine/models/cosyvoice3/tokenizer_text.h create mode 100644 src/models/breeze_tts/assets.cpp create mode 100644 src/models/breeze_tts/generator.cpp create mode 100644 src/models/breeze_tts/session.cpp create mode 100644 src/models/breeze_tts/speech_decoder.cpp create mode 100644 src/models/breeze_tts/speech_encoder.cpp create mode 100644 src/models/breeze_tts/text_encoder.cpp create mode 100644 src/models/breeze_tts/tokenizer_text.cpp create mode 100644 src/models/cosyvoice3/ar.cpp create mode 100644 src/models/cosyvoice3/assets.cpp create mode 100644 src/models/cosyvoice3/flow.cpp create mode 100644 src/models/cosyvoice3/frontend.cpp create mode 100644 src/models/cosyvoice3/hift.cpp create mode 100644 src/models/cosyvoice3/session.cpp create mode 100644 src/models/cosyvoice3/tokenizer_text.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 313f05b7e..a5abb3ba6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1618,6 +1618,48 @@ audiocpp_add_model(firered_audio engine::models::firered_audio::make_firered_audio_loader ) +audiocpp_add_model(cosyvoice3 + SOURCES + src/models/cosyvoice3/ar.cpp + src/models/cosyvoice3/assets.cpp + src/models/cosyvoice3/flow.cpp + src/models/cosyvoice3/frontend.cpp + src/models/cosyvoice3/hift.cpp + src/models/cosyvoice3/session.cpp + src/models/cosyvoice3/tokenizer_text.cpp + INCLUDES + engine/models/cosyvoice3/ar.h + engine/models/cosyvoice3/assets.h + engine/models/cosyvoice3/flow.h + engine/models/cosyvoice3/frontend.h + engine/models/cosyvoice3/hift.h + engine/models/cosyvoice3/session.h + engine/models/cosyvoice3/tokenizer_text.h + LOADERS + engine::models::cosyvoice3::make_cosyvoice3_loader +) + +audiocpp_add_model(breeze_tts + SOURCES + src/models/breeze_tts/assets.cpp + src/models/breeze_tts/generator.cpp + src/models/breeze_tts/session.cpp + src/models/breeze_tts/speech_decoder.cpp + src/models/breeze_tts/speech_encoder.cpp + src/models/breeze_tts/text_encoder.cpp + src/models/breeze_tts/tokenizer_text.cpp + INCLUDES + engine/models/breeze_tts/assets.h + engine/models/breeze_tts/generator.h + engine/models/breeze_tts/session.h + engine/models/breeze_tts/speech_decoder.h + engine/models/breeze_tts/speech_encoder.h + engine/models/breeze_tts/text_encoder.h + engine/models/breeze_tts/tokenizer_text.h + LOADERS + engine::models::breeze_tts::make_breeze_tts_loader +) + set(AUDIOCPP_ENABLED_MODELS "") if (AUDIOCPP_MODEL_SET STREQUAL "full") set(AUDIOCPP_ENABLED_MODELS ${AUDIOCPP_MODEL_TARGETS}) diff --git a/include/engine/framework/modules/text_encoders/t5_gemma_encoder.h b/include/engine/framework/modules/text_encoders/t5_gemma_encoder.h index 23f612ddc..df17e2e3c 100644 --- a/include/engine/framework/modules/text_encoders/t5_gemma_encoder.h +++ b/include/engine/framework/modules/text_encoders/t5_gemma_encoder.h @@ -2,24 +2,36 @@ #include "engine/framework/core/module.h" #include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" #include namespace engine::modules { +enum class T5GemmaRMSNormStyle { + Gemma, + Direct, +}; + struct T5GemmaEncoderConfig { int64_t hidden_size = 0; int64_t layers = 0; int64_t attention_heads = 0; int64_t kv_heads = 0; int64_t head_dim = 0; + int64_t attention_size = 0; int64_t intermediate_size = 0; int64_t vocab_size = 0; float rope_theta = 10000.0F; float rms_norm_eps = 1.0e-6F; float attn_logit_softcap = 50.0F; float query_pre_attn_scalar = 64.0F; + float rope_freq_scale = 1.0F; + std::vector layer_rope_theta; + std::vector layer_rope_freq_scale; bool scale_embeddings = true; + bool use_qk_norm = false; + T5GemmaRMSNormStyle rms_norm_style = T5GemmaRMSNormStyle::Gemma; }; struct T5GemmaEncoderLayerWeights { @@ -31,6 +43,8 @@ struct T5GemmaEncoderLayerWeights { LinearWeights k_proj; LinearWeights v_proj; LinearWeights o_proj; + NormWeights q_norm; + NormWeights k_norm; LinearWeights gate_proj; LinearWeights up_proj; LinearWeights down_proj; diff --git a/include/engine/framework/tokenizers/llama_bpe.h b/include/engine/framework/tokenizers/llama_bpe.h index cd1bb0070..9a048f55f 100644 --- a/include/engine/framework/tokenizers/llama_bpe.h +++ b/include/engine/framework/tokenizers/llama_bpe.h @@ -96,13 +96,15 @@ struct LlamaBpeTokenizerSpec { std::filesystem::path tokenizer_config_path_ = {}, std::optional tokenizer_json_path_ = std::nullopt, LlamaBpePreTokenizer pre_type_ = LlamaBpePreTokenizer::Gpt2, - std::vector additional_special_tokens_ = {}) + std::vector additional_special_tokens_ = {}, + std::string normalizer_space_replacement_ = {}) : vocab_path(std::move(vocab_path_)), merges_path(std::move(merges_path_)), tokenizer_config_path(std::move(tokenizer_config_path_)), tokenizer_json_path(std::move(tokenizer_json_path_)), pre_type(pre_type_), - additional_special_tokens(std::move(additional_special_tokens_)) {} + additional_special_tokens(std::move(additional_special_tokens_)), + normalizer_space_replacement(std::move(normalizer_space_replacement_)) {} std::filesystem::path vocab_path; std::filesystem::path merges_path; @@ -110,6 +112,7 @@ struct LlamaBpeTokenizerSpec { std::optional tokenizer_json_path; LlamaBpePreTokenizer pre_type = LlamaBpePreTokenizer::Gpt2; std::vector additional_special_tokens; + std::string normalizer_space_replacement; }; class LlamaBpeTokenizer final : public ITokenizer { diff --git a/include/engine/models/breeze_tts/assets.h b/include/engine/models/breeze_tts/assets.h new file mode 100644 index 000000000..f675c341d --- /dev/null +++ b/include/engine/models/breeze_tts/assets.h @@ -0,0 +1,76 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/assets/tensor_source.h" + +#include +#include +#include +#include + +namespace engine::models::breeze_tts { + +struct BreezeTTSConfig { + int sample_rate = 24000; + int64_t hidden_size = 2048; + int64_t intermediate_size = 6144; + int64_t layers = 28; + int64_t heads = 16; + int64_t kv_heads = 8; + int64_t head_dim = 128; + int64_t vocab_size = 2051; + int64_t lm_head_size = 2052; + int64_t text_vocab_size = 262158; + int64_t num_codebooks = 16; + int64_t max_position_embeddings = 2048; + float rms_norm_eps = 1.0e-5F; + float rope_theta = 500000.0F; + float rope_scaling_factor = 32.0F; + float rope_low_freq_factor = 0.125F; + float rope_high_freq_factor = 0.5F; + int64_t rope_original_max_position_embeddings = 1024; + bool rope_scaling_enabled = false; + int64_t audio_token_id = 262144; + int64_t audio_eos_token_id = 262145; + int64_t codebook_pad_token_id = 2050; + int64_t codebook_eos_token_id = 0; + + int64_t text_hidden_size = 1152; + int64_t text_intermediate_size = 6912; + int64_t text_layers = 26; + int64_t text_heads = 4; + int64_t text_kv_heads = 1; + int64_t text_head_dim = 256; + int64_t text_max_position_embeddings = 32768; + float text_rms_norm_eps = 1.0e-6F; + float text_rope_theta = 1000000.0F; + float text_rope_linear_factor = 8.0F; + float text_query_pre_attn_scalar = 256.0F; + std::vector text_layer_rope_theta; + std::vector text_layer_rope_freq_scale; + + int64_t depth_hidden_size = 1024; + int64_t depth_intermediate_size = 4096; + int64_t depth_layers = 4; + int64_t depth_heads = 8; + int64_t depth_kv_heads = 2; + int64_t depth_head_dim = 128; + float depth_rms_norm_eps = 1.0e-5F; + float depth_rope_theta = 500000.0F; + float depth_rope_scaling_factor = 32.0F; + float depth_rope_low_freq_factor = 0.001953125F; + float depth_rope_high_freq_factor = 0.0078125F; + int64_t depth_rope_original_max_position_embeddings = 16; + bool depth_rope_scaling_enabled = false; +}; + +struct BreezeTTSAssets { + std::filesystem::path model_root; + engine::assets::ResourceBundle resources; + BreezeTTSConfig config; + std::shared_ptr weights; +}; + +std::shared_ptr load_breeze_tts_assets(const std::filesystem::path & model_path); + +} // namespace engine::models::breeze_tts diff --git a/include/engine/models/breeze_tts/generator.h b/include/engine/models/breeze_tts/generator.h new file mode 100644 index 000000000..7b3a95ff7 --- /dev/null +++ b/include/engine/models/breeze_tts/generator.h @@ -0,0 +1,49 @@ +#pragma once + +#include "engine/framework/runtime/session.h" +#include "engine/models/breeze_tts/assets.h" +#include "engine/models/breeze_tts/speech_decoder.h" +#include "engine/models/breeze_tts/text_encoder.h" +#include "engine/models/breeze_tts/tokenizer_text.h" + +#include +#include +#include +#include + +namespace engine::models::breeze_tts { + +struct BreezeGenerationRequest { + std::string text; + std::string instruction; + std::string reference_text; + std::optional reference_audio; + std::optional reference_codes; + float guidance_scale = 1.0F; + float temperature = 0.9F; + float depth_temperature = 0.9F; + int64_t top_k = 50; + float top_p = 1.0F; + int64_t max_tokens = 1500; + uint64_t seed = 0; +}; + +class BreezeGeneratorRuntime { +public: + BreezeGeneratorRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type); + ~BreezeGeneratorRuntime(); + + engine::runtime::AudioBuffer generate(const BreezeGenerationRequest & request); + BreezeSpeechCodes encode_reference(const engine::runtime::AudioBuffer & audio) const; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::breeze_tts diff --git a/include/engine/models/breeze_tts/session.h b/include/engine/models/breeze_tts/session.h new file mode 100644 index 000000000..31a765926 --- /dev/null +++ b/include/engine/models/breeze_tts/session.h @@ -0,0 +1,64 @@ +#pragma once + +#include "engine/framework/model_spec/metadata.h" +#include "engine/framework/runtime/cache_slots.h" +#include "engine/framework/runtime/model.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/models/breeze_tts/assets.h" +#include "engine/models/breeze_tts/speech_decoder.h" + +#include +#include +#include +#include + +namespace engine::models::breeze_tts { + +class BreezeGeneratorRuntime; + +std::shared_ptr make_breeze_tts_loader(); + +class BreezeTTSSession final + : public engine::runtime::RuntimeSessionBase + , public engine::runtime::IOfflineVoiceTaskSession { +public: + BreezeTTSSession( + engine::runtime::TaskSpec task, + engine::runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract); + ~BreezeTTSSession() override; + + std::string family() const override; + engine::runtime::VoiceTaskKind task_kind() const override; + engine::runtime::RunMode run_mode() const override; + void prepare(const engine::runtime::SessionPreparationRequest & request) override; + engine::runtime::TaskResult run(const engine::runtime::TaskRequest & request) override; + +private: + struct ReferenceCacheKey { + int sample_rate = 0; + int channels = 0; + uint64_t sample_count = 0; + uint64_t sample_hash = 0; + }; + + struct ReferenceCacheKeyEqual { + bool operator()(const ReferenceCacheKey & lhs, const ReferenceCacheKey & rhs) const noexcept; + }; + + struct ReferenceCacheEntry { + BreezeSpeechCodes codes; + }; + + BreezeSpeechCodes resolve_reference_codes(const engine::runtime::AudioBuffer & audio); + + engine::runtime::TaskSpec task_; + std::shared_ptr assets_; + std::shared_ptr contract_; + std::unique_ptr generator_; + engine::runtime::CacheSlots reference_cache_; + std::optional uncached_reference_; +}; + +} // namespace engine::models::breeze_tts diff --git a/include/engine/models/breeze_tts/speech_decoder.h b/include/engine/models/breeze_tts/speech_decoder.h new file mode 100644 index 000000000..bb27bc33d --- /dev/null +++ b/include/engine/models/breeze_tts/speech_decoder.h @@ -0,0 +1,61 @@ +#pragma once + +#include "engine/framework/core/execution_context.h" +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/runtime/session.h" +#include "engine/models/breeze_tts/assets.h" + +#include +#include +#include +#include +#include + +namespace engine::core { +class ConstantTensorCache; +} + +namespace engine::models { + +namespace breeze_tts { + +struct BreezeSpeechCodes { + std::vector codes; + int64_t frames = 0; + int64_t code_groups = 0; +}; + +struct BreezeSpeechDecoderWeights; +class BreezeSpeechDecoderGraph; + +class BreezeSpeechDecoderRuntime { +public: + BreezeSpeechDecoderRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution_context, + size_t graph_arena_bytes, + size_t constant_context_bytes, + engine::assets::TensorStorageType linear_weight_storage_type, + engine::assets::TensorStorageType conv_weight_storage_type); + ~BreezeSpeechDecoderRuntime(); + + runtime::AudioBuffer decode(const BreezeSpeechCodes & codec_codes) const; + runtime::AudioBuffer decode_and_trim_reference( + const BreezeSpeechCodes & reference_codes, + const BreezeSpeechCodes & generated_codes) const; + void release_runtime_graphs() const; + +private: + std::shared_ptr assets_; + core::ExecutionContext * execution_context_ = nullptr; + std::shared_ptr weights_; + size_t graph_arena_bytes_ = 0; + std::unique_ptr constants_; + mutable std::unique_ptr graph_; + // Always present to keep this public class layout identical when the private + // Strix Halo compile definition differs between translation units. + mutable std::array, 2> optimized_graphs_; +}; + +} // namespace breeze_tts +} // namespace engine::models diff --git a/include/engine/models/breeze_tts/speech_encoder.h b/include/engine/models/breeze_tts/speech_encoder.h new file mode 100644 index 000000000..60da8d924 --- /dev/null +++ b/include/engine/models/breeze_tts/speech_encoder.h @@ -0,0 +1,53 @@ +#pragma once + +#include "engine/framework/core/execution_context.h" +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/runtime/session.h" +#include "engine/models/breeze_tts/assets.h" +#include "engine/models/breeze_tts/speech_decoder.h" + +#include +#include +#include + +namespace engine::core { +class ConstantTensorCache; +} + +namespace engine::models { + +namespace breeze_tts { + +struct BreezeSpeechEncoderWeights; +class BreezeSpeechEncoderGraph; + +struct BreezeSpeechEncoderOutput { + BreezeSpeechCodes codes; + std::vector semantic_projected; + std::vector acoustic_projected; +}; + +class BreezeSpeechEncoderRuntime { +public: + BreezeSpeechEncoderRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution_context, + size_t graph_arena_bytes, + engine::assets::TensorStorageType linear_weight_storage_type, + engine::assets::TensorStorageType conv_weight_storage_type); + ~BreezeSpeechEncoderRuntime(); + + BreezeSpeechCodes encode(const runtime::AudioBuffer & audio) const; + void release_runtime_graphs() const; + +private: + std::shared_ptr assets_; + std::shared_ptr weights_; + core::ExecutionContext * execution_context_ = nullptr; + size_t graph_arena_bytes_ = 0; + std::unique_ptr constants_; + mutable std::unique_ptr graph_; +}; + +} // namespace breeze_tts +} // namespace engine::models diff --git a/include/engine/models/breeze_tts/text_encoder.h b/include/engine/models/breeze_tts/text_encoder.h new file mode 100644 index 000000000..35ac17338 --- /dev/null +++ b/include/engine/models/breeze_tts/text_encoder.h @@ -0,0 +1,36 @@ +#pragma once + +#include "engine/framework/core/execution_context.h" +#include "engine/framework/assets/tensor_source.h" +#include "engine/models/breeze_tts/assets.h" +#include "engine/models/breeze_tts/tokenizer_text.h" + +#include +#include + +namespace engine::models::breeze_tts { + +struct BreezeProjectedText { + int64_t tokens = 0; + std::vector values; +}; + +class BreezeTextEncoderRuntime { +public: + BreezeTextEncoderRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type); + ~BreezeTextEncoderRuntime(); + + BreezeProjectedText encode(const std::vector & input_ids); + void release_runtime_graphs(); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::breeze_tts diff --git a/include/engine/models/breeze_tts/tokenizer_text.h b/include/engine/models/breeze_tts/tokenizer_text.h new file mode 100644 index 000000000..718693378 --- /dev/null +++ b/include/engine/models/breeze_tts/tokenizer_text.h @@ -0,0 +1,44 @@ +#pragma once + +#include "engine/models/breeze_tts/assets.h" + +#include +#include +#include +#include + +namespace engine::models::breeze_tts { + +struct BreezePromptBranch { + std::vector input_ids; + std::vector text_mask; + std::vector text_segment_lengths; + std::vector> text_segments; +}; + +class BreezeTextTokenizer { +public: + explicit BreezeTextTokenizer(std::shared_ptr assets); + ~BreezeTextTokenizer(); + + BreezePromptBranch build_tts_instruction(const std::string & text, const std::string & instruction) const; + BreezePromptBranch build_tts_plain(const std::string & text) const; + BreezePromptBranch build_clone( + const std::string & text, + const std::string & instruction, + const std::string & reference_text, + int64_t reference_audio_frames) const; + BreezePromptBranch build_clone_negative( + const std::string & text, + const std::string & reference_text, + int64_t reference_audio_frames) const; + + int32_t audio_token_id() const noexcept; + int32_t audio_eos_token_id() const noexcept; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::breeze_tts diff --git a/include/engine/models/cosyvoice3/ar.h b/include/engine/models/cosyvoice3/ar.h new file mode 100644 index 000000000..589246b30 --- /dev/null +++ b/include/engine/models/cosyvoice3/ar.h @@ -0,0 +1,48 @@ +#pragma once + +#include "engine/framework/core/execution_context.h" +#include "engine/framework/sampling/torch_random.h" +#include "engine/models/cosyvoice3/assets.h" + +#include +#include +#include + +namespace engine::models::cosyvoice3 { + +struct CosyVoice3ArRequest { + std::vector prompt_text_tokens; + std::vector target_text_tokens; + std::vector prompt_speech_tokens; + uint32_t seed = 1986; + int64_t top_k = 25; + int64_t min_tokens = -1; + int64_t max_tokens = -1; +}; + +struct CosyVoice3ArOutput { + std::vector speech_tokens; +}; + +class CosyVoice3ArRuntime { +public: + CosyVoice3ArRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type); + ~CosyVoice3ArRuntime(); + + CosyVoice3ArRuntime(const CosyVoice3ArRuntime &) = delete; + CosyVoice3ArRuntime & operator=(const CosyVoice3ArRuntime &) = delete; + + CosyVoice3ArOutput generate(const CosyVoice3ArRequest & request); + void release_graphs(); + +private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::cosyvoice3 diff --git a/include/engine/models/cosyvoice3/assets.h b/include/engine/models/cosyvoice3/assets.h new file mode 100644 index 000000000..777ecc5d3 --- /dev/null +++ b/include/engine/models/cosyvoice3/assets.h @@ -0,0 +1,51 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/assets/tensor_source.h" + +#include +#include +#include + +namespace engine::models::cosyvoice3 { + +struct CosyVoice3Config { + int64_t sample_rate = 24000; + int64_t text_vocab_size = 151936; + int64_t hidden_size = 896; + int64_t intermediate_size = 4864; + int64_t layers = 24; + int64_t heads = 14; + int64_t kv_heads = 2; + int64_t head_dim = 64; + int64_t speech_token_size = 6561; + int64_t speech_reserved_tokens = 200; + int64_t flow_mel_channels = 80; + int64_t flow_hidden_size = 1024; + int64_t flow_layers = 22; + int64_t flow_heads = 16; + int64_t flow_head_dim = 64; + int64_t flow_ff_mult = 2; + int64_t flow_input_channels = 240; + int64_t flow_static_chunk_size = 50; + int64_t token_mel_ratio = 2; + int64_t pre_lookahead_len = 3; + int64_t speaker_dim = 192; +}; + +struct CosyVoice3Assets { + std::filesystem::path model_root; + std::filesystem::path gguf_path; + engine::assets::ResourceBundle resources; + CosyVoice3Config config; + std::shared_ptr llm_weights; + std::shared_ptr flow_weights; + std::shared_ptr hift_weights; + std::shared_ptr campplus_weights; + std::shared_ptr speech_tokenizer_weights; + std::shared_ptr blank_en_weights; +}; + +std::shared_ptr load_cosyvoice3_assets(const std::filesystem::path & model_path); + +} // namespace engine::models::cosyvoice3 diff --git a/include/engine/models/cosyvoice3/flow.h b/include/engine/models/cosyvoice3/flow.h new file mode 100644 index 000000000..d5211e4d3 --- /dev/null +++ b/include/engine/models/cosyvoice3/flow.h @@ -0,0 +1,48 @@ +#pragma once + +#include "engine/framework/core/execution_context.h" +#include "engine/models/cosyvoice3/assets.h" + +#include +#include +#include + +namespace engine::models::cosyvoice3 { + +struct CosyVoice3FlowRequest { + std::vector speech_tokens; + std::vector prompt_speech_tokens; + std::vector prompt_mel; + int64_t prompt_mel_frames = 0; + std::vector speaker_embedding; + uint32_t seed = 1986; + int64_t num_inference_steps = 10; +}; + +struct CosyVoice3FlowOutput { + std::vector mel; + int64_t frames = 0; +}; + +class CosyVoice3FlowRuntime { +public: + CosyVoice3FlowRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type); + ~CosyVoice3FlowRuntime(); + + CosyVoice3FlowRuntime(const CosyVoice3FlowRuntime &) = delete; + CosyVoice3FlowRuntime & operator=(const CosyVoice3FlowRuntime &) = delete; + + CosyVoice3FlowOutput generate(const CosyVoice3FlowRequest & request); + void release_graphs(); + +private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::cosyvoice3 diff --git a/include/engine/models/cosyvoice3/frontend.h b/include/engine/models/cosyvoice3/frontend.h new file mode 100644 index 000000000..4de239aa9 --- /dev/null +++ b/include/engine/models/cosyvoice3/frontend.h @@ -0,0 +1,42 @@ +#pragma once + +#include "engine/framework/core/execution_context.h" +#include "engine/framework/runtime/model.h" +#include "engine/models/cosyvoice3/assets.h" + +#include +#include + +namespace engine::models::cosyvoice3 { + +struct CosyVoice3ReferenceFeatures { + std::vector speech_tokens; + int64_t speech_token_count = 0; + std::vector prompt_mel; + int64_t prompt_mel_frames = 0; + std::vector speaker_embedding; +}; + +class CosyVoice3Frontend { +public: + CosyVoice3Frontend( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type, + size_t reference_cache_slots); + ~CosyVoice3Frontend(); + + CosyVoice3Frontend(const CosyVoice3Frontend &) = delete; + CosyVoice3Frontend & operator=(const CosyVoice3Frontend &) = delete; + + const CosyVoice3ReferenceFeatures & prepare_reference(const engine::runtime::AudioBuffer & audio); + void release_graphs(); + +private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::cosyvoice3 diff --git a/include/engine/models/cosyvoice3/hift.h b/include/engine/models/cosyvoice3/hift.h new file mode 100644 index 000000000..af3d8ac22 --- /dev/null +++ b/include/engine/models/cosyvoice3/hift.h @@ -0,0 +1,35 @@ +#pragma once + +#include "engine/framework/core/execution_context.h" +#include "engine/framework/runtime/session.h" +#include "engine/models/cosyvoice3/assets.h" + +#include +#include + +namespace engine::models::cosyvoice3 { + +class CosyVoice3HiftRuntime { +public: + CosyVoice3HiftRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + engine::assets::TensorStorageType storage_type); + ~CosyVoice3HiftRuntime(); + + CosyVoice3HiftRuntime(const CosyVoice3HiftRuntime &) = delete; + CosyVoice3HiftRuntime & operator=(const CosyVoice3HiftRuntime &) = delete; + + engine::runtime::AudioBuffer synthesize( + const std::vector & mel, + int64_t frames, + uint64_t seed, + const std::vector * source_random_values = nullptr); + void release_graphs(); + +private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::cosyvoice3 diff --git a/include/engine/models/cosyvoice3/session.h b/include/engine/models/cosyvoice3/session.h new file mode 100644 index 000000000..f4e40c255 --- /dev/null +++ b/include/engine/models/cosyvoice3/session.h @@ -0,0 +1,49 @@ +#pragma once + +#include "engine/framework/model_spec/metadata.h" +#include "engine/framework/runtime/model.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/models/cosyvoice3/assets.h" +#include "engine/models/cosyvoice3/tokenizer_text.h" + +#include + +namespace engine::models::cosyvoice3 { + +class CosyVoice3Frontend; +class CosyVoice3ArRuntime; +class CosyVoice3FlowRuntime; +class CosyVoice3HiftRuntime; + +std::shared_ptr make_cosyvoice3_loader(); + +class CosyVoice3Session final + : public engine::runtime::RuntimeSessionBase + , public engine::runtime::IOfflineVoiceTaskSession { +public: + CosyVoice3Session( + engine::runtime::TaskSpec task, + engine::runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract); + ~CosyVoice3Session() override; + + std::string family() const override; + engine::runtime::VoiceTaskKind task_kind() const override; + engine::runtime::RunMode run_mode() const override; + void prepare(const engine::runtime::SessionPreparationRequest & request) override; + engine::runtime::TaskResult run(const engine::runtime::TaskRequest & request) override; + +private: + engine::runtime::TaskSpec task_; + std::shared_ptr assets_; + std::shared_ptr contract_; + std::unique_ptr tokenizer_; + std::unique_ptr frontend_; + std::unique_ptr ar_; + std::unique_ptr flow_; + std::unique_ptr hift_; + bool mem_saver_ = false; +}; + +} // namespace engine::models::cosyvoice3 diff --git a/include/engine/models/cosyvoice3/tokenizer_text.h b/include/engine/models/cosyvoice3/tokenizer_text.h new file mode 100644 index 000000000..dc08d6142 --- /dev/null +++ b/include/engine/models/cosyvoice3/tokenizer_text.h @@ -0,0 +1,32 @@ +#pragma once + +#include "engine/models/cosyvoice3/assets.h" + +#include +#include +#include +#include +#include + +namespace engine::models::cosyvoice3 { + +struct CosyVoice3TextTokens { + std::vector prompt; + std::vector target; +}; + +class CosyVoice3TextTokenizer { +public: + explicit CosyVoice3TextTokenizer(std::shared_ptr assets); + ~CosyVoice3TextTokenizer(); + + CosyVoice3TextTokens encode_zero_shot(std::string_view text, std::string_view prompt_text) const; + CosyVoice3TextTokens encode_cross_lingual(std::string_view text) const; + CosyVoice3TextTokens encode_instruct(std::string_view text, std::string_view instruction) const; + +private: + class Impl; + std::shared_ptr impl_; +}; + +} // namespace engine::models::cosyvoice3 diff --git a/src/framework/modules/text_encoders/t5_gemma_encoder.cpp b/src/framework/modules/text_encoders/t5_gemma_encoder.cpp index 0530468dd..a84b2f9f5 100644 --- a/src/framework/modules/text_encoders/t5_gemma_encoder.cpp +++ b/src/framework/modules/text_encoders/t5_gemma_encoder.cpp @@ -14,14 +14,27 @@ namespace engine::modules { namespace { +int64_t attention_size(const T5GemmaEncoderConfig & config) { + return config.attention_size > 0 ? config.attention_size : config.hidden_size; +} + void validate_config(const T5GemmaEncoderConfig & config) { if (config.hidden_size <= 0 || config.layers <= 0 || config.attention_heads <= 0 || config.kv_heads <= 0 || config.head_dim <= 0 || config.intermediate_size <= 0 || config.vocab_size <= 0) { throw std::runtime_error("T5GemmaEncoderConfig dimensions must be positive"); } - if (config.attention_heads * config.head_dim != config.hidden_size) { - throw std::runtime_error("T5GemmaEncoderConfig attention_heads * head_dim must equal hidden_size"); + if (config.attention_heads * config.head_dim != attention_size(config)) { + throw std::runtime_error("T5GemmaEncoderConfig attention_heads * head_dim must equal attention_size"); + } + if (config.attention_heads % config.kv_heads != 0) { + throw std::runtime_error("T5GemmaEncoderConfig attention_heads must be divisible by kv_heads"); + } + if (!config.layer_rope_theta.empty() && static_cast(config.layer_rope_theta.size()) != config.layers) { + throw std::runtime_error("T5GemmaEncoderConfig layer_rope_theta must match layer count"); + } + if (!config.layer_rope_freq_scale.empty() && static_cast(config.layer_rope_freq_scale.size()) != config.layers) { + throw std::runtime_error("T5GemmaEncoderConfig layer_rope_freq_scale must match layer count"); } if (!(config.rope_theta > 0.0F) || !(config.rms_norm_eps > 0.0F) || !(config.query_pre_attn_scalar > 0.0F)) { @@ -29,6 +42,18 @@ void validate_config(const T5GemmaEncoderConfig & config) { } } +float layer_rope_theta(const T5GemmaEncoderConfig & config, int64_t layer_index) { + return config.layer_rope_theta.empty() + ? config.rope_theta + : config.layer_rope_theta.at(static_cast(layer_index)); +} + +float layer_rope_freq_scale(const T5GemmaEncoderConfig & config, int64_t layer_index) { + return config.layer_rope_freq_scale.empty() + ? config.rope_freq_scale + : config.layer_rope_freq_scale.at(static_cast(layer_index)); +} + core::TensorValue ensure_contiguous(core::ModuleBuildContext & ctx, const core::TensorValue & input) { return core::ensure_backend_addressable_layout(ctx, input); } @@ -63,15 +88,19 @@ core::TensorValue matmul_f32(core::ModuleBuildContext & ctx, const core::TensorV return core::wrap_tensor(output, output_shape, GGML_TYPE_F32); } -core::TensorValue gemma_rms_norm( +core::TensorValue t5_gemma_rms_norm( core::ModuleBuildContext & ctx, const core::TensorValue & input, const core::TensorValue & weight, + T5GemmaRMSNormStyle style, float eps, int64_t hidden_size) { core::validate_last_dim(input, hidden_size, "T5GemmaEncoder RMSNorm input"); core::validate_shape(weight, core::TensorShape::from_dims({hidden_size}), "T5GemmaEncoder RMSNorm weight"); auto normalized = core::wrap_tensor(ggml_rms_norm(ctx.ggml, ensure_contiguous(ctx, input).tensor, eps), input.shape, GGML_TYPE_F32); + if (style == T5GemmaRMSNormStyle::Direct) { + return core::wrap_tensor(ggml_mul(ctx.ggml, normalized.tensor, weight.tensor), input.shape, GGML_TYPE_F32); + } auto one_plus_weight = core::wrap_tensor(ggml_scale_bias(ctx.ggml, weight.tensor, 1.0F, 1.0F), weight.shape, GGML_TYPE_F32); return core::wrap_tensor(ggml_mul(ctx.ggml, normalized.tensor, one_plus_weight.tensor), input.shape, GGML_TYPE_F32); } @@ -85,28 +114,62 @@ core::TensorValue reshape_heads( return core::reshape_tensor(ctx, contiguous, core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, dim})); } +core::TensorValue repeat_kv_heads(core::ModuleBuildContext & ctx, const core::TensorValue & input, int64_t repeats) { + if (repeats == 1) { + return input; + } + core::validate_rank_between(input, 4, 4, "T5GemmaEncoder repeat_kv_heads input"); + auto contiguous = ensure_contiguous(ctx, input); + const int64_t batch = contiguous.shape.dims[0]; + const int64_t kv_heads = contiguous.shape.dims[1]; + const int64_t steps = contiguous.shape.dims[2]; + const int64_t dim = contiguous.shape.dims[3]; + auto expanded = core::reshape_tensor( + ctx, + contiguous, + core::TensorShape::from_dims({batch, kv_heads, 1, steps * dim})); + expanded = RepeatModule({core::TensorShape::from_dims({batch, kv_heads, repeats, steps * dim})}) + .build(ctx, expanded); + return core::reshape_tensor( + ctx, + expanded, + core::TensorShape::from_dims({batch, kv_heads * repeats, steps, dim})); +} + core::TensorValue self_attention( core::ModuleBuildContext & ctx, const core::TensorValue & input, const core::TensorValue & positions, const core::TensorValue & additive_attention_mask, const T5GemmaEncoderLayerWeights & weights, - const T5GemmaEncoderConfig & config) { + const T5GemmaEncoderConfig & config, + int64_t layer_index) { + const int64_t attn_size = attention_size(config); const LinearModule q_proj({config.hidden_size, config.attention_heads * config.head_dim, false, GGML_PREC_F32}); const LinearModule k_proj({config.hidden_size, config.kv_heads * config.head_dim, false, GGML_PREC_F32}); const LinearModule v_proj({config.hidden_size, config.kv_heads * config.head_dim, false, GGML_PREC_F32}); - const LinearModule o_proj({config.attention_heads * config.head_dim, config.hidden_size, false, GGML_PREC_F32}); + const LinearModule o_proj({attn_size, config.hidden_size, false, GGML_PREC_F32}); auto q = q_proj.build(ctx, input, weights.q_proj); auto k = k_proj.build(ctx, input, weights.k_proj); auto v = v_proj.build(ctx, input, weights.v_proj); q = reshape_heads(ctx, q, config.attention_heads, config.head_dim); k = reshape_heads(ctx, k, config.kv_heads, config.head_dim); v = reshape_heads(ctx, v, config.kv_heads, config.head_dim); - q = RoPEModule({config.head_dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, q, positions); - k = RoPEModule({config.head_dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, k, positions); + if (config.use_qk_norm) { + if (!weights.q_norm.weight.has_value() || !weights.k_norm.weight.has_value()) { + throw std::runtime_error("T5GemmaEncoder q/k norm weights are required when use_qk_norm is enabled"); + } + q = t5_gemma_rms_norm(ctx, q, *weights.q_norm.weight, config.rms_norm_style, config.rms_norm_eps, config.head_dim); + k = t5_gemma_rms_norm(ctx, k, *weights.k_norm.weight, config.rms_norm_style, config.rms_norm_eps, config.head_dim); + } + q = RoPEModule({config.head_dim, GGML_ROPE_TYPE_NEOX, layer_rope_theta(config, layer_index), layer_rope_freq_scale(config, layer_index)}).build(ctx, q, positions); + k = RoPEModule({config.head_dim, GGML_ROPE_TYPE_NEOX, layer_rope_theta(config, layer_index), layer_rope_freq_scale(config, layer_index)}).build(ctx, k, positions); auto q_heads = TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); auto k_heads = TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k); auto v_heads = TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v); + const int64_t kv_repeats = config.attention_heads / config.kv_heads; + k_heads = repeat_kv_heads(ctx, k_heads, kv_repeats); + v_heads = repeat_kv_heads(ctx, v_heads, kv_repeats); auto k_heads_contiguous = ensure_contiguous(ctx, k_heads); auto scores_raw = ggml_mul_mat(ctx.ggml, k_heads_contiguous.tensor, q_heads.tensor); ggml_mul_mat_set_prec(scores_raw, GGML_PREC_F32); @@ -145,7 +208,7 @@ core::TensorValue self_attention( context = core::reshape_tensor( ctx, context, - core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], config.attention_heads * config.head_dim})); + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], attn_size})); return o_proj.build(ctx, context, weights.o_proj); } @@ -169,14 +232,15 @@ core::TensorValue layer( const core::TensorValue & positions, const core::TensorValue & additive_attention_mask, const T5GemmaEncoderLayerWeights & weights, - const T5GemmaEncoderConfig & config) { - auto hidden = gemma_rms_norm(ctx, input, weights.pre_self_attn_norm, config.rms_norm_eps, config.hidden_size); - hidden = self_attention(ctx, hidden, positions, additive_attention_mask, weights, config); - hidden = gemma_rms_norm(ctx, hidden, weights.post_self_attn_norm, config.rms_norm_eps, config.hidden_size); + const T5GemmaEncoderConfig & config, + int64_t layer_index) { + auto hidden = t5_gemma_rms_norm(ctx, input, weights.pre_self_attn_norm, config.rms_norm_style, config.rms_norm_eps, config.hidden_size); + hidden = self_attention(ctx, hidden, positions, additive_attention_mask, weights, config, layer_index); + hidden = t5_gemma_rms_norm(ctx, hidden, weights.post_self_attn_norm, config.rms_norm_style, config.rms_norm_eps, config.hidden_size); auto output = AddModule{}.build(ctx, input, hidden); - hidden = gemma_rms_norm(ctx, output, weights.pre_ff_norm, config.rms_norm_eps, config.hidden_size); + hidden = t5_gemma_rms_norm(ctx, output, weights.pre_ff_norm, config.rms_norm_style, config.rms_norm_eps, config.hidden_size); hidden = mlp(ctx, hidden, weights, config); - hidden = gemma_rms_norm(ctx, hidden, weights.post_ff_norm, config.rms_norm_eps, config.hidden_size); + hidden = t5_gemma_rms_norm(ctx, hidden, weights.post_ff_norm, config.rms_norm_style, config.rms_norm_eps, config.hidden_size); return AddModule{}.build(ctx, output, hidden); } @@ -213,9 +277,9 @@ core::TensorValue T5GemmaEncoderModule::build( GGML_TYPE_F32); } for (int64_t i = 0; i < config_.layers; ++i) { - hidden = layer(ctx, hidden, positions, additive_attention_mask, weights.layers[static_cast(i)], config_); + hidden = layer(ctx, hidden, positions, additive_attention_mask, weights.layers[static_cast(i)], config_, i); } - return gemma_rms_norm(ctx, hidden, weights.norm, config_.rms_norm_eps, config_.hidden_size); + return t5_gemma_rms_norm(ctx, hidden, weights.norm, config_.rms_norm_style, config_.rms_norm_eps, config_.hidden_size); } } // namespace engine::modules diff --git a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp index 8d0e65ad9..73184d7ce 100644 --- a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp +++ b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp @@ -639,7 +639,6 @@ class QwenCausalDecodeRuntime::Impl { void ensure_prefill_token_graph(int64_t steps) { if (prefill_graph_ != nullptr && prefill_input_kind_ == InputKind::Token && prefill_steps_ == steps) { - debug::timing_log_scalar(config_.trace_name + ".prefill.graph.build_ms", 0.0); debug::trace_log_scalar(config_.trace_name + ".prefill.steps", steps); return; } @@ -649,7 +648,6 @@ class QwenCausalDecodeRuntime::Impl { void ensure_prefill_embedding_graph(int64_t steps) { if (prefill_graph_ != nullptr && prefill_input_kind_ == InputKind::Embedding && prefill_steps_ == steps) { - debug::timing_log_scalar(config_.trace_name + ".prefill.graph.build_ms", 0.0); debug::trace_log_scalar(config_.trace_name + ".prefill.steps", steps); return; } @@ -820,7 +818,6 @@ class QwenCausalDecodeRuntime::Impl { void ensure_batched_prefill_token_graph(int64_t batch_size, int64_t steps) { if (batched_prefill_graph_ != nullptr && batched_prefill_input_kind_ == InputKind::Token && batched_prefill_batch_size_ == batch_size && batched_prefill_steps_ == steps) { - debug::timing_log_scalar(config_.trace_name + ".batched_prefill.graph.build_ms", 0.0); return; } release_batched_prefill_graph(); @@ -830,7 +827,6 @@ class QwenCausalDecodeRuntime::Impl { void ensure_batched_prefill_embedding_graph(int64_t batch_size, int64_t steps) { if (batched_prefill_graph_ != nullptr && batched_prefill_input_kind_ == InputKind::Embedding && batched_prefill_batch_size_ == batch_size && batched_prefill_steps_ == steps) { - debug::timing_log_scalar(config_.trace_name + ".batched_prefill.graph.build_ms", 0.0); return; } release_batched_prefill_graph(); @@ -1025,7 +1021,6 @@ class QwenCausalDecodeRuntime::Impl { void ensure_decode_token_graph(int64_t cache_steps) { if (decode_graph_ != nullptr && decode_input_kind_ == InputKind::Token && decode_cache_steps_ >= cache_steps) { - debug::timing_log_scalar(config_.trace_name + ".decode.graph.build_ms", 0.0); debug::trace_log_scalar(config_.trace_name + ".decode.cache_steps", cache_steps); return; } @@ -1035,7 +1030,6 @@ class QwenCausalDecodeRuntime::Impl { void ensure_decode_embedding_graph(int64_t cache_steps) { if (decode_graph_ != nullptr && decode_input_kind_ == InputKind::Embedding && decode_cache_steps_ >= cache_steps) { - debug::timing_log_scalar(config_.trace_name + ".decode.graph.build_ms", 0.0); debug::trace_log_scalar(config_.trace_name + ".decode.cache_steps", cache_steps); return; } @@ -1134,7 +1128,6 @@ class QwenCausalDecodeRuntime::Impl { void ensure_batched_decode_token_graph(int64_t cache_steps, int64_t batch_size) { if (batched_decode_graph_ != nullptr && batched_decode_input_kind_ == InputKind::Token && batched_decode_cache_steps_ >= cache_steps && batched_decode_batch_size_ == batch_size) { - debug::timing_log_scalar(config_.trace_name + ".batched_decode.graph.build_ms", 0.0); return; } release_batched_decode_graph(); @@ -1144,7 +1137,6 @@ class QwenCausalDecodeRuntime::Impl { void ensure_batched_decode_embedding_graph(int64_t cache_steps, int64_t batch_size) { if (batched_decode_graph_ != nullptr && batched_decode_input_kind_ == InputKind::Embedding && batched_decode_cache_steps_ >= cache_steps && batched_decode_batch_size_ == batch_size) { - debug::timing_log_scalar(config_.trace_name + ".batched_decode.graph.build_ms", 0.0); return; } release_batched_decode_graph(); diff --git a/src/framework/tokenizers/llama_bpe.cpp b/src/framework/tokenizers/llama_bpe.cpp index 80219614a..ca4cfa4ca 100644 --- a/src/framework/tokenizers/llama_bpe.cpp +++ b/src/framework/tokenizers/llama_bpe.cpp @@ -251,10 +251,27 @@ void load_merges(const std::filesystem::path & merges_path, BpeVocabulary & voca } } +std::string replace_spaces(std::string text, const std::string & replacement) { + if (replacement.empty()) { + return text; + } + std::string out; + out.reserve(text.size()); + for (const char ch : text) { + if (ch == ' ') { + out += replacement; + } else { + out.push_back(ch); + } + } + return out; +} + } // namespace struct LlamaBpeTokenizer::Impl { explicit Impl(const LlamaBpeTokenizerSpec & spec) { + normalizer_space_replacement = spec.normalizer_space_replacement; vocab.pre_type = convert_pre_type(spec.pre_type); const bool has_vocab = !spec.vocab_path.empty(); const bool has_merges = !spec.merges_path.empty(); @@ -283,6 +300,7 @@ struct LlamaBpeTokenizer::Impl { } BpeVocabulary vocab; + std::string normalizer_space_replacement; }; LlamaBpeTokenizer::LlamaBpeTokenizer(LlamaBpeTokenizerSpec spec) @@ -297,11 +315,11 @@ TokenizedText LlamaBpeTokenizer::tokenize(const std::string & text) const { } std::vector LlamaBpeTokenizer::encode(const std::string & text) const { - return vendor::tokenize_bpe(impl_->vocab, text, true); + return vendor::tokenize_bpe(impl_->vocab, replace_spaces(text, impl_->normalizer_space_replacement), true); } std::vector LlamaBpeTokenizer::encode(const std::string & text, const bool parse_special) const { - return vendor::tokenize_bpe(impl_->vocab, text, parse_special); + return vendor::tokenize_bpe(impl_->vocab, replace_spaces(text, impl_->normalizer_space_replacement), parse_special); } std::string LlamaBpeTokenizer::decode(const std::vector & token_ids, const bool skip_special_tokens) const { diff --git a/src/models/breeze_tts/assets.cpp b/src/models/breeze_tts/assets.cpp new file mode 100644 index 000000000..ad8f5f25a --- /dev/null +++ b/src/models/breeze_tts/assets.cpp @@ -0,0 +1,177 @@ +#include "engine/models/breeze_tts/assets.h" + +#include "engine/framework/io/json.h" +#include "engine/framework/model_spec/package.h" + +#include +#include + +namespace engine::models::breeze_tts { +namespace { + +constexpr const char * kFamily = "breeze_tts"; + +int64_t nested_i64(const engine::io::json::Value & object, const std::string & key, int64_t fallback) { + return engine::io::json::optional_i64(object, key, fallback); +} + +float nested_f32(const engine::io::json::Value & object, const std::string & key, float fallback) { + return engine::io::json::optional_f32(object, key, fallback); +} + +float rope_theta_for_type( + const engine::io::json::Value & rope_parameters, + const std::string & layer_type, + float fallback) { + const auto * object = rope_parameters.find(layer_type); + if (object == nullptr || !object->is_object()) { + return fallback; + } + return nested_f32(*object, "rope_theta", fallback); +} + +float rope_freq_scale_for_type( + const engine::io::json::Value & rope_parameters, + const std::string & layer_type, + float fallback) { + const auto * object = rope_parameters.find(layer_type); + if (object == nullptr || !object->is_object()) { + return fallback; + } + const std::string rope_type = engine::io::json::optional_string(*object, "rope_type", "default"); + if (rope_type != "linear") { + return 1.0F; + } + const float factor = nested_f32(*object, "factor", 1.0F); + if (!(factor > 0.0F)) { + throw std::runtime_error("BreezeTTS text rope linear factor must be positive"); + } + return 1.0F / factor; +} + +void parse_config(const engine::io::json::Value & root, BreezeTTSConfig & config) { + config.hidden_size = nested_i64(root, "hidden_size", config.hidden_size); + config.intermediate_size = nested_i64(root, "intermediate_size", config.intermediate_size); + config.layers = nested_i64(root, "num_hidden_layers", config.layers); + config.heads = nested_i64(root, "num_attention_heads", config.heads); + config.kv_heads = nested_i64(root, "num_key_value_heads", config.kv_heads); + config.head_dim = nested_i64(root, "head_dim", config.head_dim); + config.vocab_size = nested_i64(root, "vocab_size", config.vocab_size); + config.lm_head_size = config.vocab_size + 1; + config.text_vocab_size = nested_i64(root, "text_vocab_size", config.text_vocab_size); + config.num_codebooks = nested_i64(root, "num_codebooks", config.num_codebooks); + config.max_position_embeddings = nested_i64(root, "max_position_embeddings", config.max_position_embeddings); + config.rms_norm_eps = nested_f32(root, "rms_norm_eps", config.rms_norm_eps); + config.rope_theta = nested_f32(root, "rope_theta", config.rope_theta); + if (const auto * rope = root.find("rope_scaling"); rope != nullptr && rope->is_object()) { + config.rope_scaling_enabled = true; + config.rope_scaling_factor = nested_f32(*rope, "factor", config.rope_scaling_factor); + config.rope_low_freq_factor = nested_f32(*rope, "low_freq_factor", config.rope_low_freq_factor); + config.rope_high_freq_factor = nested_f32(*rope, "high_freq_factor", config.rope_high_freq_factor); + config.rope_original_max_position_embeddings = nested_i64( + *rope, + "original_max_position_embeddings", + config.rope_original_max_position_embeddings); + } + if (const auto * backbone = root.find("backbone_config"); backbone != nullptr && backbone->is_object()) { + config.hidden_size = nested_i64(*backbone, "hidden_size", config.hidden_size); + config.intermediate_size = nested_i64(*backbone, "intermediate_size", config.intermediate_size); + config.layers = nested_i64(*backbone, "num_hidden_layers", config.layers); + config.heads = nested_i64(*backbone, "num_attention_heads", config.heads); + config.kv_heads = nested_i64(*backbone, "num_key_value_heads", config.kv_heads); + config.head_dim = nested_i64(*backbone, "head_dim", config.head_dim); + config.rms_norm_eps = nested_f32(*backbone, "rms_norm_eps", config.rms_norm_eps); + config.rope_theta = nested_f32(*backbone, "rope_theta", config.rope_theta); + config.rope_scaling_enabled = false; + if (const auto * rope = backbone->find("rope_scaling"); rope != nullptr && rope->is_object()) { + config.rope_scaling_enabled = true; + config.rope_scaling_factor = nested_f32(*rope, "factor", config.rope_scaling_factor); + config.rope_low_freq_factor = nested_f32(*rope, "low_freq_factor", config.rope_low_freq_factor); + config.rope_high_freq_factor = nested_f32(*rope, "high_freq_factor", config.rope_high_freq_factor); + config.rope_original_max_position_embeddings = nested_i64( + *rope, + "original_max_position_embeddings", + config.rope_original_max_position_embeddings); + } + } + config.audio_token_id = nested_i64(root, "audio_token_id", config.audio_token_id); + config.audio_eos_token_id = nested_i64(root, "audio_eos_token_id", config.audio_eos_token_id); + config.codebook_pad_token_id = nested_i64(root, "codebook_pad_token_id", config.codebook_pad_token_id); + config.codebook_eos_token_id = nested_i64(root, "codebook_eos_token_id", config.codebook_eos_token_id); + + if (const auto * text = root.find("text_encoder_config"); text != nullptr && text->is_object()) { + config.text_hidden_size = nested_i64(*text, "hidden_size", config.text_hidden_size); + config.text_intermediate_size = nested_i64(*text, "intermediate_size", config.text_intermediate_size); + config.text_layers = nested_i64(*text, "num_hidden_layers", config.text_layers); + config.text_heads = nested_i64(*text, "num_attention_heads", config.text_heads); + config.text_kv_heads = nested_i64(*text, "num_key_value_heads", config.text_kv_heads); + config.text_head_dim = nested_i64(*text, "head_dim", config.text_head_dim); + config.text_max_position_embeddings = nested_i64(*text, "max_position_embeddings", config.text_max_position_embeddings); + config.text_rms_norm_eps = nested_f32(*text, "rms_norm_eps", config.text_rms_norm_eps); + config.text_query_pre_attn_scalar = nested_f32(*text, "query_pre_attn_scalar", config.text_query_pre_attn_scalar); + if (const auto * rope = text->find("rope_parameters"); rope != nullptr && rope->is_object()) { + if (const auto * full = rope->find("full_attention"); full != nullptr && full->is_object()) { + config.text_rope_theta = nested_f32(*full, "rope_theta", config.text_rope_theta); + config.text_rope_linear_factor = nested_f32(*full, "factor", config.text_rope_linear_factor); + } + const auto layer_types = engine::io::json::optional_string_array(*text, "layer_types"); + if (!layer_types.empty()) { + if (static_cast(layer_types.size()) != config.text_layers) { + throw std::runtime_error("BreezeTTS text layer_types must match text layer count"); + } + config.text_layer_rope_theta.clear(); + config.text_layer_rope_freq_scale.clear(); + config.text_layer_rope_theta.reserve(layer_types.size()); + config.text_layer_rope_freq_scale.reserve(layer_types.size()); + for (const auto & layer_type : layer_types) { + config.text_layer_rope_theta.push_back(rope_theta_for_type(*rope, layer_type, config.text_rope_theta)); + config.text_layer_rope_freq_scale.push_back( + rope_freq_scale_for_type(*rope, layer_type, 1.0F / config.text_rope_linear_factor)); + } + } + } + } + if (const auto * depth = root.find("depth_decoder_config"); depth != nullptr && depth->is_object()) { + config.depth_hidden_size = nested_i64(*depth, "hidden_size", config.depth_hidden_size); + config.depth_intermediate_size = nested_i64(*depth, "intermediate_size", config.depth_intermediate_size); + config.depth_layers = nested_i64(*depth, "num_hidden_layers", config.depth_layers); + config.depth_heads = nested_i64(*depth, "num_attention_heads", config.depth_heads); + config.depth_kv_heads = nested_i64(*depth, "num_key_value_heads", config.depth_kv_heads); + config.depth_head_dim = nested_i64(*depth, "head_dim", config.depth_head_dim); + config.depth_rms_norm_eps = nested_f32(*depth, "rms_norm_eps", config.depth_rms_norm_eps); + config.depth_rope_theta = nested_f32(*depth, "rope_theta", config.depth_rope_theta); + config.depth_rope_scaling_enabled = false; + if (const auto * rope = depth->find("rope_scaling"); rope != nullptr && rope->is_object()) { + config.depth_rope_scaling_enabled = true; + config.depth_rope_scaling_factor = nested_f32(*rope, "factor", config.depth_rope_scaling_factor); + config.depth_rope_low_freq_factor = nested_f32(*rope, "low_freq_factor", config.depth_rope_low_freq_factor); + config.depth_rope_high_freq_factor = nested_f32(*rope, "high_freq_factor", config.depth_rope_high_freq_factor); + config.depth_rope_original_max_position_embeddings = nested_i64( + *rope, + "original_max_position_embeddings", + config.depth_rope_original_max_position_embeddings); + } + } +} + +void validate_shapes(const engine::assets::TensorSource & source, const BreezeTTSConfig & config) { + engine::assets::require_tensor_shape(source, "embed_text_tokens.weight", {config.text_vocab_size, config.hidden_size}); + engine::assets::require_tensor_shape(source, "text_encoder_proj.weight", {config.hidden_size, config.text_hidden_size}); + engine::assets::require_tensor_shape(source, "lm_head.weight", {config.lm_head_size, config.hidden_size}); + engine::assets::require_tensor_shape(source, "depth_decoder.model.embed_tokens.weight", {config.num_codebooks * config.vocab_size, config.hidden_size}); + engine::assets::require_tensor_shape(source, "depth_decoder.codebooks_head.weight", {config.num_codebooks - 1, config.depth_hidden_size, config.vocab_size}); +} + +} // namespace + +std::shared_ptr load_breeze_tts_assets(const std::filesystem::path & model_path) { + auto assets = std::make_shared(); + assets->resources = engine::model_spec::load_resource_bundle_for_family(model_path, kFamily); + assets->model_root = assets->resources.model_root(); + assets->weights = assets->resources.open_tensor_source("model_weights"); + parse_config(assets->resources.parse_json("config_json"), assets->config); + validate_shapes(*assets->weights, assets->config); + return assets; +} + +} // namespace engine::models::breeze_tts diff --git a/src/models/breeze_tts/generator.cpp b/src/models/breeze_tts/generator.cpp new file mode 100644 index 000000000..e9cc22478 --- /dev/null +++ b/src/models/breeze_tts/generator.cpp @@ -0,0 +1,959 @@ +#include "engine/models/breeze_tts/generator.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/transformers/qwen_causal_decode_runtime.h" +#include "engine/framework/modules/weight_binding.h" +#include "engine/framework/sampling/hf_sampler.h" +#include "engine/framework/sampling/torch_random.h" +#include "engine/models/breeze_tts/speech_decoder.h" +#include "engine/models/breeze_tts/speech_encoder.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::breeze_tts { +namespace { + +namespace assets = engine::assets; +namespace binding = engine::modules::binding; +namespace core = engine::core; +namespace modules = engine::modules; +namespace runtime = engine::runtime; +namespace sampling = engine::sampling; +using Clock = std::chrono::steady_clock; + +constexpr int64_t kCodecCodebookSize = 2048; +constexpr float kRepetitionPenalty = 1.1F; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +modules::QwenCausalDecodeRuntimeConfig backbone_config( + const BreezeTTSConfig & config, + core::BackendType backend_type, + size_t graph_arena_bytes) { + modules::QwenCausalDecodeRuntimeConfig out; + out.trace_name = "breeze_tts.backbone"; + out.prefill_graph_arena_bytes = graph_arena_bytes; + out.decode_graph_arena_bytes = graph_arena_bytes; + out.decoder.stack.hidden_size = config.hidden_size; + out.decoder.stack.num_attention_heads = config.heads; + out.decoder.stack.num_key_value_heads = config.kv_heads; + out.decoder.stack.head_dim = config.head_dim; + out.decoder.stack.intermediate_size = config.intermediate_size; + out.decoder.stack.layers = config.layers; + out.decoder.stack.rms_norm_eps = config.rms_norm_eps; + out.decoder.stack.rope_theta = config.rope_theta; + out.decoder.stack.rope_type = GGML_ROPE_TYPE_NEOX; + out.decoder.stack.use_qk_norm = true; + out.decoder.stack.attention_precision = GGML_PREC_F32; + out.decoder.stack.projection_precision = GGML_PREC_DEFAULT; + out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.decoder.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.decoder.stack.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; + if (backend_type == core::BackendType::Cuda || backend_type == core::BackendType::Hip || + backend_type == core::BackendType::Vulkan) { + out.decoder.static_cache_type = GGML_TYPE_F16; + } + out.decoder.logits_size = config.lm_head_size; + out.decoder.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; + out.output_mode = modules::QwenCausalDecodeOutputMode::Logits; + out.return_hidden = true; + out.logits_readback_token_ids.reserve(static_cast(config.lm_head_size)); + for (int32_t token = 0; token < static_cast(config.lm_head_size); ++token) { + out.logits_readback_token_ids.push_back(token); + } + out.decoder.lm_head_input_type = GGML_TYPE_F32; + return out; +} + +modules::QwenCausalDecodeRuntimeConfig depth_config( + const BreezeTTSConfig & config, + core::BackendType backend_type, + size_t graph_arena_bytes) { + modules::QwenCausalDecodeRuntimeConfig out; + out.trace_name = "breeze_tts.depth_decoder"; + out.prefill_graph_arena_bytes = graph_arena_bytes; + out.decode_graph_arena_bytes = graph_arena_bytes; + out.decoder.stack.hidden_size = config.depth_hidden_size; + out.decoder.stack.num_attention_heads = config.depth_heads; + out.decoder.stack.num_key_value_heads = config.depth_kv_heads; + out.decoder.stack.head_dim = config.depth_head_dim; + out.decoder.stack.intermediate_size = config.depth_intermediate_size; + out.decoder.stack.layers = config.depth_layers; + out.decoder.stack.rms_norm_eps = config.depth_rms_norm_eps; + out.decoder.stack.rope_theta = config.depth_rope_theta; + out.decoder.stack.rope_type = GGML_ROPE_TYPE_NEOX; + out.decoder.stack.use_qk_norm = false; + out.decoder.stack.attention_precision = GGML_PREC_F32; + out.decoder.stack.projection_precision = GGML_PREC_DEFAULT; + out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.decoder.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.decoder.stack.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; + if (backend_type == core::BackendType::Cuda || backend_type == core::BackendType::Hip || + backend_type == core::BackendType::Vulkan) { + out.decoder.static_cache_type = GGML_TYPE_F16; + } + out.decoder.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; + out.output_mode = modules::QwenCausalDecodeOutputMode::Hidden; + out.return_hidden = true; + return out; +} + +std::vector llama3_rope_factors( + int64_t head_dim, + float rope_theta, + float scaling_factor, + float low_freq_factor, + float high_freq_factor, + int64_t original_max_position_embeddings) { + constexpr double pi = 3.14159265358979323846; + const double low_wavelength = + static_cast(original_max_position_embeddings) / + static_cast(low_freq_factor); + const double high_wavelength = + static_cast(original_max_position_embeddings) / + static_cast(high_freq_factor); + std::vector out(static_cast(head_dim / 2), 1.0F); + for (int64_t index = 0; index < head_dim / 2; ++index) { + const double inv_freq = 1.0 / std::pow( + static_cast(rope_theta), + static_cast(2 * index) / static_cast(head_dim)); + const double wavelength = 2.0 * pi / inv_freq; + double scaled = inv_freq; + if (wavelength > low_wavelength) { + scaled = inv_freq / static_cast(scaling_factor); + } else if (wavelength >= high_wavelength) { + const double smooth = + (static_cast(original_max_position_embeddings) / wavelength - + static_cast(low_freq_factor)) / + (static_cast(high_freq_factor) - + static_cast(low_freq_factor)); + scaled = + (1.0 - smooth) * inv_freq / static_cast(scaling_factor) + + smooth * inv_freq; + } + out[static_cast(index)] = static_cast(inv_freq / scaled); + } + return out; +} + +modules::QwenDecoderLayerWeights load_backbone_layer( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const BreezeTTSConfig & config, + assets::TensorStorageType storage_type, + const std::optional & rope_factors, + int64_t layer) { + const std::string prefix = "backbone_model.layers." + std::to_string(layer); + modules::QwenDecoderLayerWeights out; + out.input_norm = binding::norm_weight_from_source(store, source, prefix + ".input_layernorm", config.hidden_size); + out.self_attention.q_weight = store.load_tensor(source, prefix + ".self_attn.q_proj.weight", storage_type, {config.heads * config.head_dim, config.hidden_size}); + out.self_attention.k_weight = store.load_tensor(source, prefix + ".self_attn.k_proj.weight", storage_type, {config.kv_heads * config.head_dim, config.hidden_size}); + out.self_attention.v_weight = store.load_tensor(source, prefix + ".self_attn.v_proj.weight", storage_type, {config.kv_heads * config.head_dim, config.hidden_size}); + out.self_attention.out_weight = store.load_tensor(source, prefix + ".self_attn.o_proj.weight", storage_type, {config.hidden_size, config.heads * config.head_dim}); + out.q_norm = binding::norm_weight_from_source(store, source, prefix + ".self_attn.q_norm", config.head_dim); + out.k_norm = binding::norm_weight_from_source(store, source, prefix + ".self_attn.k_norm", config.head_dim); + out.post_norm = binding::norm_weight_from_source(store, source, prefix + ".post_attention_layernorm", config.hidden_size); + out.mlp.gate_proj = binding::linear_from_source(store, source, prefix + ".mlp.gate_proj", storage_type, config.intermediate_size, config.hidden_size, false); + out.mlp.up_proj = binding::linear_from_source(store, source, prefix + ".mlp.up_proj", storage_type, config.intermediate_size, config.hidden_size, false); + out.mlp.down_proj = binding::linear_from_source(store, source, prefix + ".mlp.down_proj", storage_type, config.hidden_size, config.intermediate_size, false); + out.rope_frequency_factors = rope_factors; + return out; +} + +modules::QwenDecoderLayerWeights load_depth_layer( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const BreezeTTSConfig & config, + assets::TensorStorageType storage_type, + const std::optional & rope_factors, + int64_t layer) { + const std::string prefix = "depth_decoder.model.layers." + std::to_string(layer); + modules::QwenDecoderLayerWeights out; + out.input_norm = binding::norm_weight_from_source(store, source, prefix + ".input_layernorm", config.depth_hidden_size); + out.self_attention.q_weight = store.load_tensor(source, prefix + ".self_attn.q_proj.weight", storage_type, {config.depth_heads * config.depth_head_dim, config.depth_hidden_size}); + out.self_attention.k_weight = store.load_tensor(source, prefix + ".self_attn.k_proj.weight", storage_type, {config.depth_kv_heads * config.depth_head_dim, config.depth_hidden_size}); + out.self_attention.v_weight = store.load_tensor(source, prefix + ".self_attn.v_proj.weight", storage_type, {config.depth_kv_heads * config.depth_head_dim, config.depth_hidden_size}); + out.self_attention.out_weight = store.load_tensor(source, prefix + ".self_attn.o_proj.weight", storage_type, {config.depth_hidden_size, config.depth_heads * config.depth_head_dim}); + out.post_norm = binding::norm_weight_from_source(store, source, prefix + ".post_attention_layernorm", config.depth_hidden_size); + out.mlp.gate_proj = binding::linear_from_source(store, source, prefix + ".mlp.gate_proj", storage_type, config.depth_intermediate_size, config.depth_hidden_size, false); + out.mlp.up_proj = binding::linear_from_source(store, source, prefix + ".mlp.up_proj", storage_type, config.depth_intermediate_size, config.depth_hidden_size, false); + out.mlp.down_proj = binding::linear_from_source(store, source, prefix + ".mlp.down_proj", storage_type, config.depth_hidden_size, config.depth_intermediate_size, false); + out.rope_frequency_factors = rope_factors; + return out; +} + +std::vector frame_embedding( + const std::vector & table, + int64_t rows, + int64_t dim, + int64_t vocab, + const std::vector & codes) { + std::vector out(static_cast(dim), 0.0F); + for (int64_t codebook = 0; codebook < static_cast(codes.size()); ++codebook) { + const int64_t row = codebook * vocab + codes[static_cast(codebook)]; + if (row < 0 || row >= rows) { + throw std::runtime_error("BreezeTTS audio code is outside embedding table"); + } + const size_t begin = static_cast(row * dim); + for (int64_t i = 0; i < dim; ++i) { + out[static_cast(i)] += table[begin + static_cast(i)]; + } + } + return out; +} + +void suppress_reserved(std::vector & logits, int64_t codebook_size, int64_t vocab_size) { + for (int64_t token = codebook_size; token < vocab_size; ++token) { + if (token >= 0 && token < static_cast(logits.size())) { + logits[static_cast(token)] = -std::numeric_limits::infinity(); + } + } +} + +int32_t sample_logits( + std::vector logits, + const std::vector & history, + const sampling::HfSamplingOptions & options, + sampling::HfSamplerScratch & scratch, + std::mt19937 & fallback_rng, + const sampling::TorchCudaSamplingPolicy * policy, + uint64_t seed, + uint64_t & call_index, + uint64_t & offset_blocks, + std::string_view context) { + const sampling::HfTorchSamplingState torch_state{policy, seed, call_index, offset_blocks, true}; + const int32_t token = sampling::HfSampler{}.sample( + logits, + history, + options, + scratch, + fallback_rng, + policy != nullptr && policy->cuda_fast_path ? &torch_state : nullptr, + context); + ++call_index; + if (policy != nullptr && policy->cuda_fast_path) { + offset_blocks += sampling::torch_cuda_tensor_iterator_offset_blocks(static_cast(logits.size()), *policy); + } + return token; +} + +struct BreezeWeights { + std::shared_ptr store; + modules::QwenCausalDecodeRuntimeWeights backbone; + modules::QwenCausalDecodeRuntimeWeights depth; + std::vector audio_embedding; + modules::LinearWeights depth_projector; + core::TensorValue depth_heads; +}; + +std::shared_ptr load_weights( + const BreezeTTSAssets & assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + assets::TensorStorageType storage_type, + const modules::QwenCausalDecodeRuntimeConfig & backbone_runtime_config) { + auto out = std::make_shared(); + out->store = std::make_shared( + execution.backend(), + execution.backend_type(), + "breeze_tts.generator.weights", + weight_context_bytes); + const auto & source = *assets.weights; + const auto & config = assets.config; + std::optional backbone_rope_factors; + if (config.rope_scaling_enabled) { + backbone_rope_factors = out->store->make_f32( + core::TensorShape::from_dims({config.head_dim / 2}), + llama3_rope_factors( + config.head_dim, + config.rope_theta, + config.rope_scaling_factor, + config.rope_low_freq_factor, + config.rope_high_freq_factor, + config.rope_original_max_position_embeddings)); + } + std::optional depth_rope_factors; + if (config.depth_rope_scaling_enabled) { + depth_rope_factors = out->store->make_f32( + core::TensorShape::from_dims({config.depth_head_dim / 2}), + llama3_rope_factors( + config.depth_head_dim, + config.depth_rope_theta, + config.depth_rope_scaling_factor, + config.depth_rope_low_freq_factor, + config.depth_rope_high_freq_factor, + config.depth_rope_original_max_position_embeddings)); + } + out->backbone.token_embedding = out->store->load_tensor( + source, + "depth_decoder.model.embed_tokens.weight", + storage_type, + {config.num_codebooks * config.vocab_size, config.hidden_size}); + out->backbone.stack.layers.reserve(static_cast(config.layers)); + for (int64_t layer = 0; layer < config.layers; ++layer) { + out->backbone.stack.layers.push_back(load_backbone_layer(*out->store, source, config, storage_type, backbone_rope_factors, layer)); + } + out->backbone.final_norm = binding::norm_weight_from_source(*out->store, source, "backbone_model.norm", config.hidden_size); + out->backbone.lm_head = binding::linear_from_source( + *out->store, + source, + "lm_head", + storage_type, + backbone_runtime_config.decoder.logits_size, + config.hidden_size, + false); + + out->depth.token_embedding = out->backbone.token_embedding; + out->depth.stack.layers.reserve(static_cast(config.depth_layers)); + for (int64_t layer = 0; layer < config.depth_layers; ++layer) { + out->depth.stack.layers.push_back(load_depth_layer(*out->store, source, config, storage_type, depth_rope_factors, layer)); + } + out->depth.final_norm = binding::norm_weight_from_source(*out->store, source, "depth_decoder.model.norm", config.depth_hidden_size); + + out->audio_embedding = source.require_f32("depth_decoder.model.embed_tokens.weight", {config.num_codebooks * config.vocab_size, config.hidden_size}); + out->depth_projector = { + out->store->load_f32_tensor( + source, + "depth_decoder.model.inputs_embeds_projector.weight", + {config.depth_hidden_size, config.hidden_size}), + std::nullopt}; + const auto depth_heads = source.require_f32( + "depth_decoder.codebooks_head.weight", + {config.num_codebooks - 1, config.depth_hidden_size, config.vocab_size}); + std::vector transposed_depth_heads( + static_cast((config.num_codebooks - 1) * config.vocab_size * config.depth_hidden_size)); + for (int64_t codebook = 0; codebook < config.num_codebooks - 1; ++codebook) { + const int64_t in_codebook_offset = codebook * config.depth_hidden_size * config.vocab_size; + const int64_t out_codebook_offset = codebook * config.vocab_size * config.depth_hidden_size; + for (int64_t row = 0; row < config.depth_hidden_size; ++row) { + const int64_t in_row_offset = in_codebook_offset + row * config.vocab_size; + for (int64_t token = 0; token < config.vocab_size; ++token) { + transposed_depth_heads[static_cast(out_codebook_offset + token * config.depth_hidden_size + row)] = + depth_heads[static_cast(in_row_offset + token)]; + } + } + } + out->depth_heads = out->store->make_f32( + core::TensorShape::from_dims({config.num_codebooks - 1, config.vocab_size, config.depth_hidden_size}), + transposed_depth_heads); + out->store->upload(); + return out; +} + +class BreezeDepthProjectionRuntime { +public: + BreezeDepthProjectionRuntime( + ggml_backend_t backend, + core::BackendType backend_type, + int threads, + size_t graph_arena_bytes, + const BreezeTTSConfig & config, + const modules::LinearWeights & projector_weights, + const core::TensorValue & packed_heads) + : backend_(backend), + threads_(std::max(1, threads)), + hidden_(config.hidden_size), + depth_hidden_(config.depth_hidden_size), + vocab_(config.vocab_size) { + if (backend_ == nullptr || hidden_ <= 0 || depth_hidden_ <= 0 || vocab_ <= 0 || config.num_codebooks <= 1) { + throw std::runtime_error("BreezeTTS depth projection shape is invalid"); + } + ctx_.reset(ggml_init({graph_arena_bytes, nullptr, true})); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize BreezeTTS depth projection graph context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "breeze_tts.depth_projection", backend_type}; + projector_single_ = build_linear_graph( + ctx, + 1, + hidden_, + depth_hidden_, + projector_weights, + "breeze_tts.depth_projector.single"); + projector_pair_ = build_linear_graph( + ctx, + 2, + hidden_, + depth_hidden_, + projector_weights, + "breeze_tts.depth_projector.pair"); + head_graphs_.reserve(static_cast(config.num_codebooks - 1)); + for (int64_t codebook = 1; codebook < config.num_codebooks; ++codebook) { + head_graphs_.push_back(build_head_graph(ctx, packed_heads, codebook)); + } + graph_buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), backend_); + if (graph_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate BreezeTTS depth projection graphs"); + } + } + + ~BreezeDepthProjectionRuntime() { + release_graph(projector_single_); + release_graph(projector_pair_); + for (const auto & graph : head_graphs_) { + release_graph(graph); + } + if (graph_buffer_ != nullptr) { + ggml_backend_buffer_free(graph_buffer_); + } + } + + std::vector project_single(const std::vector & hidden) const { + return run_graph(projector_single_, hidden); + } + + std::vector project_pair( + const std::vector & cond_hidden, + const std::vector & uncond_hidden) const { + if (static_cast(cond_hidden.size()) != hidden_ || + static_cast(uncond_hidden.size()) != hidden_) { + throw std::runtime_error("BreezeTTS depth projector pair input size mismatch"); + } + std::vector input; + input.reserve(static_cast(2 * hidden_)); + input.insert(input.end(), cond_hidden.begin(), cond_hidden.end()); + input.insert(input.end(), uncond_hidden.begin(), uncond_hidden.end()); + return run_graph(projector_pair_, input); + } + + std::vector logits_cfg( + const std::vector & cond_hidden, + const std::vector & uncond_hidden, + int64_t codebook, + float guidance_scale) const { + if (codebook <= 0 || static_cast(codebook) > head_graphs_.size()) { + throw std::runtime_error("BreezeTTS depth codebook index is invalid"); + } + if (static_cast(cond_hidden.size()) != depth_hidden_ || + static_cast(uncond_hidden.size()) != depth_hidden_) { + throw std::runtime_error("BreezeTTS depth head input size mismatch"); + } + std::vector input; + input.reserve(static_cast(2 * depth_hidden_)); + input.insert(input.end(), cond_hidden.begin(), cond_hidden.end()); + input.insert(input.end(), uncond_hidden.begin(), uncond_hidden.end()); + const auto paired = run_graph(head_graphs_[static_cast(codebook - 1)], input); + std::vector out(static_cast(vocab_)); + const size_t vocab = static_cast(vocab_); + for (int64_t token = 0; token < vocab_; ++token) { + const size_t index = static_cast(token); + out[index] = paired[vocab + index] + guidance_scale * (paired[index] - paired[vocab + index]); + } + return out; + } + +private: + struct Graph { + ggml_tensor * input = nullptr; + ggml_tensor * output = nullptr; + ggml_cgraph * graph = nullptr; + int64_t input_size = 0; + int64_t output_size = 0; + const char * label = nullptr; + }; + + Graph build_linear_graph( + core::ModuleBuildContext & ctx, + int64_t batch, + int64_t in_features, + int64_t out_features, + const modules::LinearWeights & weights, + const char * label) { + auto input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch, in_features})); + auto output = modules::LinearModule({in_features, out_features, false}) + .build(ctx, input, weights) + .tensor; + auto * graph = ggml_new_graph_custom(ctx_.get(), 32768, false); + ggml_set_output(output); + ggml_build_forward_expand(graph, output); + return {input.tensor, output, graph, batch * in_features, batch * out_features, label}; + } + + Graph build_head_graph( + core::ModuleBuildContext & ctx, + const core::TensorValue & packed_heads, + int64_t codebook) { + auto input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({2, depth_hidden_})); + const size_t row_stride = packed_heads.tensor->nb[1]; + const size_t codebook_stride = packed_heads.tensor->nb[2]; + const size_t offset = static_cast(codebook - 1) * codebook_stride; + auto weight = core::wrap_tensor( + ggml_view_2d(ctx_.get(), packed_heads.tensor, depth_hidden_, vocab_, row_stride, offset), + core::TensorShape::from_dims({vocab_, depth_hidden_}), + packed_heads.type); + auto output = modules::LinearModule({depth_hidden_, vocab_, false}) + .build(ctx, input, modules::LinearWeights{weight, std::nullopt}) + .tensor; + auto * graph = ggml_new_graph_custom(ctx_.get(), 32768, false); + ggml_set_output(output); + ggml_build_forward_expand(graph, output); + return {input.tensor, output, graph, 2 * depth_hidden_, 2 * vocab_, "breeze_tts.depth_head"}; + } + + void release_graph(const Graph & graph) const { + core::release_backend_graph_resources(backend_, graph.graph); + } + + std::vector run_graph(const Graph & graph, const std::vector & input) const { + if (static_cast(input.size()) != graph.input_size) { + throw std::runtime_error("BreezeTTS depth projection input size mismatch"); + } + ggml_backend_tensor_set(graph.input, input.data(), 0, input.size() * sizeof(float)); + core::set_backend_threads(backend_, threads_); + if (core::compute_backend_graph(backend_, graph.graph, nullptr, graph.label) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("BreezeTTS depth projection graph compute failed"); + } + ggml_backend_synchronize(backend_); + std::vector output(static_cast(graph.output_size)); + ggml_backend_tensor_get(graph.output, output.data(), 0, output.size() * sizeof(float)); + return output; + } + + ggml_backend_t backend_ = nullptr; + int threads_ = 1; + int64_t hidden_ = 0; + int64_t depth_hidden_ = 0; + int64_t vocab_ = 0; + std::unique_ptr ctx_; + ggml_backend_buffer_t graph_buffer_ = nullptr; + Graph projector_single_; + Graph projector_pair_; + std::vector head_graphs_; +}; + +} // namespace + +struct BreezeGeneratorRuntime::Impl { + Impl( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) + : assets(std::move(assets)), + execution(execution), + tokenizer(this->assets), + text_encoder(this->assets, execution, graph_arena_bytes, weight_context_bytes, storage_type), + sampling_policy(sampling::resolve_torch_cuda_sampling_policy( + execution.backend_type(), + execution.config().device, + "breeze_tts.sampling", + "BreezeTTS", + sampling::TorchCudaSamplingPolicyFailureMode::FallbackToDefault)) { + if (this->assets == nullptr) { + throw std::runtime_error("BreezeTTS generator requires assets"); + } + const auto & config = this->assets->config; + backbone_runtime_config = backbone_config(config, execution.backend_type(), graph_arena_bytes); + depth_runtime_config = depth_config(config, execution.backend_type(), graph_arena_bytes); + weights = load_weights(*this->assets, execution, weight_context_bytes, storage_type, backbone_runtime_config); + backbone_cond = std::make_unique(execution, backbone_runtime_config, weights->backbone); + backbone_uncond = std::make_unique(execution, backbone_runtime_config, weights->backbone); + depth_pair = std::make_unique(execution, depth_runtime_config, weights->depth); + depth_projection = std::make_unique( + execution.backend(), + execution.backend_type(), + execution.config().threads, + graph_arena_bytes, + config, + weights->depth_projector, + weights->depth_heads); + speech_encoder = std::make_unique( + this->assets, + execution, + graph_arena_bytes, + storage_type, + storage_type); + speech_decoder = std::make_unique( + this->assets, + execution, + graph_arena_bytes, + weight_context_bytes, + storage_type, + storage_type); + } + + std::vector merge_prompt(const BreezePromptBranch & branch, const std::vector & reference_codes) { + const auto & config = assets->config; + const int64_t embedding_rows = config.num_codebooks * config.vocab_size; + std::vector out; + out.reserve(branch.input_ids.size() * static_cast(config.hidden_size)); + size_t text_segment_index = 0; + int64_t text_segment_offset = 0; + BreezeProjectedText projected_text; + int64_t audio_frame = 0; + for (size_t pos = 0; pos < branch.input_ids.size(); ++pos) { + if (branch.text_mask[pos] != 0) { + if (text_segment_offset == 0) { + if (text_segment_index >= branch.text_segments.size()) { + throw std::runtime_error("BreezeTTS text segment state mismatch"); + } + projected_text = text_encoder.encode(branch.text_segments[text_segment_index]); + } + const size_t begin = static_cast(text_segment_offset * config.hidden_size); + out.insert( + out.end(), + projected_text.values.begin() + static_cast(begin), + projected_text.values.begin() + static_cast(begin + static_cast(config.hidden_size))); + ++text_segment_offset; + if (text_segment_offset == projected_text.tokens) { + ++text_segment_index; + text_segment_offset = 0; + } + continue; + } + const int32_t id = branch.input_ids[pos]; + if (id == tokenizer.audio_token_id()) { + if ((audio_frame + 1) * config.num_codebooks > static_cast(reference_codes.size())) { + throw std::runtime_error("BreezeTTS reference audio code count is shorter than prompt placeholders"); + } + std::vector frame(static_cast(config.num_codebooks)); + for (int64_t codebook = 0; codebook < config.num_codebooks; ++codebook) { + frame[static_cast(codebook)] = + reference_codes[static_cast(audio_frame * config.num_codebooks + codebook)]; + } + const auto embedded = frame_embedding( + weights->audio_embedding, + embedding_rows, + config.hidden_size, + config.vocab_size, + frame); + out.insert(out.end(), embedded.begin(), embedded.end()); + ++audio_frame; + } else if (id == tokenizer.audio_eos_token_id()) { + std::vector eos(static_cast(config.num_codebooks), static_cast(config.codebook_eos_token_id)); + const auto embedded = frame_embedding( + weights->audio_embedding, + embedding_rows, + config.hidden_size, + config.vocab_size, + eos); + out.insert(out.end(), embedded.begin(), embedded.end()); + } else { + throw std::runtime_error("BreezeTTS prompt has non-text token that is not audio"); + } + } + return out; + } + + std::vector generate_frame( + const std::vector & cond_hidden, + const std::vector & uncond_hidden, + int32_t first_token, + const BreezeGenerationRequest & request, + sampling::HfSamplerScratch & scratch, + std::mt19937 & fallback_rng, + uint64_t & call_index, + uint64_t & offset_blocks) { + const auto & config = assets->config; + std::vector frame; + frame.reserve(static_cast(config.num_codebooks)); + frame.push_back(first_token); + + const auto project_audio_embedding_row = [&](int64_t row) { + const int64_t rows = config.num_codebooks * config.vocab_size; + if (row < 0 || row >= rows) { + throw std::runtime_error("BreezeTTS embedding row is outside table"); + } + const size_t begin = static_cast(row * config.hidden_size); + std::vector embedding( + weights->audio_embedding.begin() + static_cast(begin), + weights->audio_embedding.begin() + static_cast(begin + static_cast(config.hidden_size))); + return depth_projection->project_single(embedding); + }; + + const auto first_embed = project_audio_embedding_row(first_token); + const auto projected = depth_projection->project_pair(cond_hidden, uncond_hidden); + const auto split = projected.begin() + static_cast(config.depth_hidden_size); + std::vector cond_prefill; + cond_prefill.reserve(static_cast(2 * config.depth_hidden_size)); + cond_prefill.insert(cond_prefill.end(), projected.begin(), split); + cond_prefill.insert(cond_prefill.end(), first_embed.begin(), first_embed.end()); + std::vector uncond_prefill; + uncond_prefill.reserve(static_cast(2 * config.depth_hidden_size)); + uncond_prefill.insert(uncond_prefill.end(), split, projected.end()); + uncond_prefill.insert(uncond_prefill.end(), first_embed.begin(), first_embed.end()); + + std::vector prefill; + prefill.reserve(static_cast(4 * config.depth_hidden_size)); + prefill.insert(prefill.end(), cond_prefill.begin(), cond_prefill.end()); + prefill.insert(prefill.end(), uncond_prefill.begin(), uncond_prefill.end()); + auto depth = depth_pair->prefill_embeddings_batched(prefill, 2, 2); + if (static_cast(depth.hidden.size()) != 2 * config.depth_hidden_size) { + throw std::runtime_error("BreezeTTS batched depth prefill hidden size mismatch"); + } + std::vector cond_hidden_now( + depth.hidden.begin(), + depth.hidden.begin() + static_cast(config.depth_hidden_size)); + std::vector uncond_hidden_now( + depth.hidden.begin() + static_cast(config.depth_hidden_size), + depth.hidden.end()); + depth_pair->start_decode_embeddings_batched(depth.state, config.num_codebooks + 1); + + sampling::HfSamplingOptions options; + options.do_sample = true; + options.temperature = request.depth_temperature; + options.top_k = request.top_k; + options.top_p = request.top_p; + options.min_tokens_to_keep = 1; + for (int64_t codebook = 1; codebook < config.num_codebooks; ++codebook) { + auto logits = depth_projection->logits_cfg(cond_hidden_now, uncond_hidden_now, codebook, request.guidance_scale); + suppress_reserved(logits, kCodecCodebookSize, config.vocab_size); + const int32_t token = sample_logits( + std::move(logits), + {}, + options, + scratch, + fallback_rng, + sampling_policy.cuda_fast_path ? &sampling_policy : nullptr, + request.seed, + call_index, + offset_blocks, + "BreezeTTS depth sampler"); + frame.push_back(token); + if (codebook + 1 < config.num_codebooks) { + const auto next = project_audio_embedding_row(codebook * config.vocab_size + token); + std::vector next_pair; + next_pair.reserve(static_cast(2 * config.depth_hidden_size)); + next_pair.insert(next_pair.end(), next.begin(), next.end()); + next_pair.insert(next_pair.end(), next.begin(), next.end()); + const auto step = depth_pair->decode_embeddings_batched(next_pair, 2); + if (static_cast(step.hidden.size()) != 2 * config.depth_hidden_size) { + throw std::runtime_error("BreezeTTS batched depth decode hidden size mismatch"); + } + cond_hidden_now.assign( + step.hidden.begin(), + step.hidden.begin() + static_cast(config.depth_hidden_size)); + uncond_hidden_now.assign( + step.hidden.begin() + static_cast(config.depth_hidden_size), + step.hidden.end()); + } + } + return frame; + } + + runtime::AudioBuffer generate(const BreezeGenerationRequest & request) { + if (request.text.empty()) { + throw std::runtime_error("BreezeTTS requires text"); + } + const auto & config = assets->config; + BreezeSpeechCodes reference; + if (request.reference_codes.has_value()) { + reference = *request.reference_codes; + } else if (request.reference_audio.has_value()) { + reference = speech_encoder->encode(*request.reference_audio); + speech_encoder->release_runtime_graphs(); + } + std::vector reference_codes; + int64_t reference_frames = 0; + if (!reference.codes.empty()) { + if (reference.frames < 0 || reference.code_groups <= 0) { + throw std::runtime_error("BreezeTTS speech codes have invalid shape"); + } + if (static_cast(reference.codes.size()) != reference.frames * reference.code_groups) { + throw std::runtime_error("BreezeTTS speech code count does not match shape"); + } + reference_codes = reference.codes; + reference_frames = static_cast(reference_codes.size()) / config.num_codebooks; + } + BreezePromptBranch cond_branch; + BreezePromptBranch uncond_branch; + std::vector cond_embeddings; + std::vector uncond_embeddings; + int64_t cond_steps = 0; + int64_t uncond_steps = 0; + const double prompt_ms = engine::debug::measure_ms([&] { + if (!reference_codes.empty()) { + if (request.reference_text.empty()) { + throw std::runtime_error("BreezeTTS clone requires reference_text"); + } + cond_branch = tokenizer.build_clone(request.text, request.instruction, request.reference_text, reference_frames); + uncond_branch = tokenizer.build_clone_negative(request.text, request.reference_text, reference_frames); + } else { + cond_branch = tokenizer.build_tts_instruction(request.text, request.instruction); + uncond_branch = tokenizer.build_tts_plain(request.text); + } + cond_embeddings = merge_prompt(cond_branch, reference_codes); + uncond_embeddings = merge_prompt(uncond_branch, reference_codes); + cond_steps = static_cast(cond_branch.input_ids.size()); + uncond_steps = static_cast(uncond_branch.input_ids.size()); + }); + engine::debug::timing_log_scalar("breeze_tts.generate.prompt_ms", prompt_ms); + text_encoder.release_runtime_graphs(); + double backbone_cond_decode_ms = 0.0; + double backbone_uncond_decode_ms = 0.0; + std::vector first_codebook_history; + std::vector codes; + double backbone_cond_prefill_ms = 0.0; + double backbone_uncond_prefill_ms = 0.0; + const double ar_ms = engine::debug::measure_ms([&] { + modules::QwenCausalPrefillResult cond; + backbone_cond_prefill_ms = engine::debug::measure_ms([&] { + cond = backbone_cond->prefill_embeddings(cond_embeddings, cond_steps); + }); + modules::QwenCausalPrefillResult uncond; + backbone_uncond_prefill_ms = engine::debug::measure_ms([&] { + uncond = backbone_uncond->prefill_embeddings(uncond_embeddings, uncond_steps); + }); + backbone_cond->start_decode_embeddings(cond.state, cond_steps + request.max_tokens); + backbone_uncond->start_decode_embeddings(uncond.state, uncond_steps + request.max_tokens); + + sampling::HfSamplerScratch scratch; + scratch.reserve_vocab(static_cast(config.lm_head_size)); + std::mt19937 fallback_rng(static_cast(request.seed)); + uint64_t sample_call_index = 0; + uint64_t offset_blocks = 0; + sampling::HfSamplingOptions first_options; + first_options.do_sample = true; + first_options.temperature = request.temperature; + first_options.top_k = request.top_k; + first_options.top_p = request.top_p; + first_options.repetition_penalty = kRepetitionPenalty; + first_options.min_tokens_to_keep = 1; + + codes.reserve(static_cast(request.max_tokens * config.num_codebooks)); + for (int64_t step = 0; step < request.max_tokens; ++step) { + if (cond.logits.size() != uncond.logits.size()) { + throw std::runtime_error("BreezeTTS CFG logits shape mismatch"); + } + std::vector logits(cond.logits.size(), 0.0F); + for (size_t i = 0; i < logits.size(); ++i) { + logits[i] = uncond.logits[i] + request.guidance_scale * (cond.logits[i] - uncond.logits[i]); + } + suppress_reserved(logits, kCodecCodebookSize, config.vocab_size); + const int32_t first_token = sample_logits( + std::move(logits), + first_codebook_history, + first_options, + scratch, + fallback_rng, + sampling_policy.cuda_fast_path ? &sampling_policy : nullptr, + request.seed, + sample_call_index, + offset_blocks, + "BreezeTTS semantic sampler"); + if (first_token == config.vocab_size) { + break; + } + if (first_token == config.codebook_pad_token_id) { + continue; + } + const auto frame = generate_frame( + cond.hidden, + uncond.hidden, + first_token, + request, + scratch, + fallback_rng, + sample_call_index, + offset_blocks); + first_codebook_history.push_back(first_token); + codes.insert(codes.end(), frame.begin(), frame.end()); + const auto embedded = frame_embedding( + weights->audio_embedding, + config.num_codebooks * config.vocab_size, + config.hidden_size, + config.vocab_size, + frame); + modules::QwenCausalDecodeStepResult cond_step; + backbone_cond_decode_ms += engine::debug::measure_ms([&] { + cond_step = backbone_cond->decode_embedding(embedded); + }); + modules::QwenCausalDecodeStepResult uncond_step; + backbone_uncond_decode_ms += engine::debug::measure_ms([&] { + uncond_step = backbone_uncond->decode_embedding(embedded); + }); + cond.logits = cond_step.logits; + cond.hidden = cond_step.hidden; + uncond.logits = uncond_step.logits; + uncond.hidden = uncond_step.hidden; + } + }); + engine::debug::timing_log_scalar("breeze_tts.ar.total_ms", ar_ms); + engine::debug::timing_log_scalar("breeze_tts.ar.backbone_cond_prefill_ms", backbone_cond_prefill_ms); + engine::debug::timing_log_scalar("breeze_tts.ar.backbone_uncond_prefill_ms", backbone_uncond_prefill_ms); + engine::debug::timing_log_scalar("breeze_tts.ar.backbone_cond_decode_ms", backbone_cond_decode_ms); + engine::debug::timing_log_scalar("breeze_tts.ar.backbone_uncond_decode_ms", backbone_uncond_decode_ms); + backbone_cond->release_runtime_graphs(); + backbone_uncond->release_runtime_graphs(); + depth_pair->release_runtime_graphs(); + if (codes.empty()) { + throw std::runtime_error("BreezeTTS generated no audio codes"); + } + if (config.num_codebooks <= 0 || static_cast(codes.size()) % config.num_codebooks != 0) { + throw std::runtime_error("BreezeTTS audio code count must be divisible by num_codebooks"); + } + BreezeSpeechCodes speech_codes; + speech_codes.codes = codes; + speech_codes.code_groups = config.num_codebooks; + speech_codes.frames = static_cast(codes.size()) / config.num_codebooks; + runtime::AudioBuffer audio = speech_decoder->decode(speech_codes); + speech_decoder->release_runtime_graphs(); + for (float & sample : audio.samples) { + sample = std::clamp(sample, -1.0F, 1.0F); + } + return audio; + } + + std::shared_ptr assets; + core::ExecutionContext & execution; + BreezeTextTokenizer tokenizer; + BreezeTextEncoderRuntime text_encoder; + sampling::TorchCudaSamplingPolicy sampling_policy; + modules::QwenCausalDecodeRuntimeConfig backbone_runtime_config; + modules::QwenCausalDecodeRuntimeConfig depth_runtime_config; + std::shared_ptr weights; + std::unique_ptr backbone_cond; + std::unique_ptr backbone_uncond; + std::unique_ptr depth_pair; + std::unique_ptr depth_projection; + std::unique_ptr speech_encoder; + std::unique_ptr speech_decoder; +}; + +BreezeGeneratorRuntime::BreezeGeneratorRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type) + : impl_(std::make_unique(std::move(assets), execution, graph_arena_bytes, weight_context_bytes, storage_type)) {} + +BreezeGeneratorRuntime::~BreezeGeneratorRuntime() = default; + +engine::runtime::AudioBuffer BreezeGeneratorRuntime::generate(const BreezeGenerationRequest & request) { + const auto start = Clock::now(); + auto audio = impl_->generate(request); + engine::debug::timing_log_scalar("breeze_tts.generate.total_ms", engine::debug::elapsed_ms(start)); + return audio; +} + +BreezeSpeechCodes BreezeGeneratorRuntime::encode_reference(const engine::runtime::AudioBuffer & audio) const { + return impl_->speech_encoder->encode(audio); +} + +} // namespace engine::models::breeze_tts diff --git a/src/models/breeze_tts/session.cpp b/src/models/breeze_tts/session.cpp new file mode 100644 index 000000000..5ac49a0a4 --- /dev/null +++ b/src/models/breeze_tts/session.cpp @@ -0,0 +1,251 @@ +#include "engine/models/breeze_tts/session.h" + +#include "engine/framework/debug/profiler.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/runtime/spec_backed_model.h" +#include "engine/framework/text/chunking.h" +#include "engine/models/breeze_tts/generator.h" + +#include +#include +#include +#include +#include +#include + +namespace engine::models::breeze_tts { +namespace { + +constexpr const char * kFamily = "breeze_tts"; +constexpr const char * kModelName = "BreezeTTS"; +constexpr int64_t kDefaultTextChunkSize = 600; +constexpr int64_t kDefaultReferenceCacheSlots = 1; + +std::shared_ptr require_assets(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("BreezeTTS session requires assets"); + } + return assets; +} + +std::shared_ptr require_contract( + std::shared_ptr contract) { + if (contract == nullptr) { + throw std::runtime_error("BreezeTTS session requires a model contract"); + } + return contract; +} + +std::vector split_request(const runtime::TaskRequest & request) { + const int64_t text_chunk_size = + engine::text::parse_text_chunk_size_override(request.options).value_or(kDefaultTextChunkSize); + if (text_chunk_size <= 0) { + throw std::runtime_error("BreezeTTS text_chunk_size must be positive"); + } + const auto text_chunk_mode = + engine::text::parse_text_chunk_mode_override(request.options).value_or(engine::text::TextChunkMode::Default); + return runtime::chunk_text_request(request, text_chunk_size, text_chunk_mode); +} + +std::size_t reference_cache_slots_from_options(const runtime::SessionOptions & options) { + const int64_t slots = runtime::parse_i64_option( + options.options, + {"reference_cache_slots"}) + .value_or(kDefaultReferenceCacheSlots); + if (slots < 0) { + throw std::runtime_error("breeze_tts.reference_cache_slots must be non-negative"); + } + if (static_cast(slots) > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("breeze_tts.reference_cache_slots is too large"); + } + return static_cast(slots); +} + +uint64_t fnv1a_mix(uint64_t hash, const void * data, size_t size) { + const auto * bytes = static_cast(data); + for (size_t i = 0; i < size; ++i) { + hash ^= bytes[i]; + hash *= 1099511628211ull; + } + return hash; +} + +uint64_t hash_audio_samples(const runtime::AudioBuffer & audio) { + uint64_t hash = 1469598103934665603ull; + for (const float sample : audio.samples) { + uint32_t bits = 0; + std::memcpy(&bits, &sample, sizeof(bits)); + hash = fnv1a_mix(hash, &bits, sizeof(bits)); + } + return hash; +} + +std::unique_ptr create_breeze_tts_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options, + std::shared_ptr assets, + std::shared_ptr contract) { + return std::make_unique(task, options, std::move(assets), std::move(contract)); +} + +} // namespace + +BreezeTTSSession::BreezeTTSSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract) + : runtime::RuntimeSessionBase(options), + task_(task), + assets_(require_assets(std::move(assets))), + contract_(require_contract(std::move(contract))), + reference_cache_(reference_cache_slots_from_options(options)) { + runtime::validate_spec_backed_session_options(options, *contract_, kFamily, kModelName); + if (task_.task != runtime::VoiceTaskKind::Tts && task_.task != runtime::VoiceTaskKind::VoiceCloning) { + throw std::runtime_error("BreezeTTS supports tts and clone tasks"); + } + if (task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("BreezeTTS supports offline sessions"); + } + using T = engine::assets::TensorStorageType; + const auto storage_type = runtime::parse_tensor_storage_option( + options.options, + "weight_type", + T::Native, + {T::Native, T::F32, T::F16, T::BF16, T::Q8_0, T::Q4_0, T::Q4_K}); + const auto graph_arena_bytes = runtime::parse_size_mb_option( + options.options, + {"graph_arena_mb"}, + 1024ull * 1024ull * 1024ull); + const auto weight_context_bytes = runtime::parse_size_mb_option( + options.options, + {"weight_context_mb"}, + 2048ull * 1024ull * 1024ull); + generator_ = std::make_unique( + assets_, + execution_context(), + graph_arena_bytes, + weight_context_bytes, + storage_type); +} + +BreezeTTSSession::~BreezeTTSSession() = default; + +std::string BreezeTTSSession::family() const { + return kFamily; +} + +runtime::VoiceTaskKind BreezeTTSSession::task_kind() const { + return task_.task; +} + +runtime::RunMode BreezeTTSSession::run_mode() const { + return task_.mode; +} + +void BreezeTTSSession::prepare(const runtime::SessionPreparationRequest & request) { + runtime::validate_spec_backed_request_options(request.options, *contract_, kModelName); + mark_prepared(); +} + +BreezeSpeechCodes BreezeTTSSession::resolve_reference_codes(const runtime::AudioBuffer & audio) { + ReferenceCacheKey key; + key.sample_rate = audio.sample_rate; + key.channels = audio.channels; + key.sample_count = static_cast(audio.samples.size()); + key.sample_hash = hash_audio_samples(audio); + if (const auto * cached = reference_cache_.find(key)) { + engine::debug::trace_log_scalar("breeze_tts.reference_cache.hit", 1); + engine::debug::trace_log_scalar("breeze_tts.reference_cache.slots", static_cast(reference_cache_.capacity())); + engine::debug::trace_log_scalar("breeze_tts.reference_cache.entries", static_cast(reference_cache_.size())); + return cached->codes; + } + const auto start = std::chrono::steady_clock::now(); + ReferenceCacheEntry entry; + entry.codes = generator_->encode_reference(audio); + engine::debug::trace_log_scalar("breeze_tts.reference.frames", entry.codes.frames); + engine::debug::trace_log_scalar("breeze_tts.reference.codebooks", entry.codes.code_groups); + if (reference_cache_.capacity() == 0) { + uncached_reference_ = std::move(entry); + } else { + reference_cache_.put(key, std::move(entry)); + } + engine::debug::trace_log_scalar("breeze_tts.reference_cache.hit", 0); + engine::debug::trace_log_scalar("breeze_tts.reference_cache.slots", static_cast(reference_cache_.capacity())); + engine::debug::trace_log_scalar("breeze_tts.reference_cache.entries", static_cast(reference_cache_.size())); + engine::debug::timing_log_scalar("breeze_tts.reference_encode_ms", engine::debug::elapsed_ms(start)); + if (reference_cache_.capacity() == 0) { + return uncached_reference_->codes; + } + const auto * cached = reference_cache_.find(key); + if (cached == nullptr) { + throw std::runtime_error("BreezeTTS reference cache insert failed"); + } + return cached->codes; +} + +runtime::TaskResult BreezeTTSSession::run(const runtime::TaskRequest & request) { + const auto wall_start = std::chrono::steady_clock::now(); + runtime::validate_spec_backed_request_options(request.options, *contract_, kModelName); + require_prepared("BreezeTTS run"); + if (!request.text_input.has_value() || request.text_input->text.empty()) { + throw std::runtime_error("BreezeTTS requires text input"); + } + runtime::AudioBuffer merged; + auto chunks = split_request(request); + const int64_t text_chunk_size = + engine::text::parse_text_chunk_size_override(request.options).value_or(kDefaultTextChunkSize); + const auto text_chunk_mode = + engine::text::parse_text_chunk_mode_override(request.options).value_or(engine::text::TextChunkMode::Default); + engine::debug::trace_log_scalar("breeze_tts.text_chunk_mode", engine::text::text_chunk_mode_name(text_chunk_mode)); + engine::debug::trace_log_scalar("breeze_tts.text_chunk_size", text_chunk_size); + engine::debug::trace_log_scalar("breeze_tts.text.chunk_count", static_cast(chunks.size())); + std::optional reference_codes; + if (request.voice.has_value() && + request.voice->speaker.has_value() && + request.voice->speaker->audio.has_value()) { + reference_codes = resolve_reference_codes(*request.voice->speaker->audio); + } + for (size_t index = 0; index < chunks.size(); ++index) { + const auto & chunk = chunks[index]; + BreezeGenerationRequest generation; + generation.text = chunk.text_input->text; + generation.instruction = runtime::find_option(chunk.options, {"instruction"}).value_or(""); + generation.reference_text = runtime::find_option(chunk.options, {"reference_text"}).value_or(""); + generation.guidance_scale = runtime::parse_positive_finite_float_option(chunk.options, {"guidance_scale"}).value_or(generation.guidance_scale); + generation.temperature = runtime::parse_positive_finite_float_option(chunk.options, {"temperature"}).value_or(generation.temperature); + generation.depth_temperature = runtime::parse_positive_finite_float_option(chunk.options, {"depth_temperature"}).value_or(generation.depth_temperature); + generation.top_k = runtime::parse_i64_option(chunk.options, {"top_k"}).value_or(generation.top_k); + generation.top_p = runtime::parse_positive_finite_float_option(chunk.options, {"top_p"}).value_or(generation.top_p); + generation.max_tokens = runtime::parse_positive_i64_option(chunk.options, {"max_tokens"}, generation.max_tokens); + generation.seed = runtime::parse_u64_option(chunk.options, {"seed"}).value_or(generation.seed); + generation.reference_codes = reference_codes; + if (index > 0) { + ++generation.seed; + } + runtime::append_audio_buffer(merged, generator_->generate(generation)); + } + runtime::TaskResult result; + result.audio_output = std::move(merged); + engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start)); + return result; +} + +bool BreezeTTSSession::ReferenceCacheKeyEqual::operator()( + const ReferenceCacheKey & lhs, + const ReferenceCacheKey & rhs) const noexcept { + return lhs.sample_rate == rhs.sample_rate && + lhs.channels == rhs.channels && + lhs.sample_count == rhs.sample_count && + lhs.sample_hash == rhs.sample_hash; +} + +std::shared_ptr make_breeze_tts_loader() { + runtime::SpecBackedVoiceModelConfig config; + config.family = kFamily; + config.load_assets = load_breeze_tts_assets; + config.create_session = create_breeze_tts_session; + return runtime::make_spec_backed_voice_loader(std::move(config)); +} + +} // namespace engine::models::breeze_tts diff --git a/src/models/breeze_tts/speech_decoder.cpp b/src/models/breeze_tts/speech_decoder.cpp new file mode 100644 index 000000000..bd1c220af --- /dev/null +++ b/src/models/breeze_tts/speech_decoder.cpp @@ -0,0 +1,1175 @@ +#include "engine/models/breeze_tts/speech_decoder.h" + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/io/json.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/attention/feed_forward.h" +#include "engine/framework/modules/attention/scaled_dot_product_attention.h" +#include "engine/framework/modules/attention/types.h" +#include "engine/framework/modules/conditioning_modules.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" + +#include "engine/framework/core/constant_tensor_cache.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::breeze_tts { +namespace json = engine::io::json; +namespace { + +using Clock = std::chrono::steady_clock; +namespace binding = modules::binding; + +constexpr int64_t kSampleRate = 24000; +constexpr int64_t kDecodeSamplesPerCode = 1920; +constexpr int64_t kChunkCodes = 300; +constexpr int64_t kLeftContextCodes = 25; +constexpr std::array kStrixHaloCachedChunkFrames{300, 105}; +#if defined(ENGINE_HIP_STRIX_HALO_OPTIMIZATIONS) +constexpr bool kStrixHaloGraphCacheEnabled = true; +#else +constexpr bool kStrixHaloGraphCacheEnabled = false; +#endif +constexpr float kCodebookEps = 1.0e-5F; +constexpr float kMaskNegInf = -1.0e9F; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct TransformerLayerWeights { + modules::AttentionWeights attention; + modules::GatedFeedForwardWeights mlp; + modules::NormWeights input_norm; + modules::NormWeights post_norm; + modules::LayerScaleWeights attn_scale; + modules::LayerScaleWeights mlp_scale; +}; + +struct ConvNeXtWeights { + modules::Conv1dWeights dwconv; + modules::NormWeights norm; + modules::LinearWeights pwconv1; + modules::LinearWeights pwconv2; + modules::LayerScaleWeights gamma; +}; + +struct ResidualUnitWeights { + std::vector act1_alpha; + std::vector act1_beta; + modules::Conv1dWeights conv1; + std::vector act2_alpha; + std::vector act2_beta; + modules::Conv1dWeights conv2; +}; + +struct UpsampleStageWeights { + modules::ConvTranspose1dWeights upconv; + ConvNeXtWeights convnext; +}; + +struct DecoderBlockWeights { + std::vector input_alpha; + std::vector input_beta; + modules::ConvTranspose1dWeights upconv; + std::vector residual_units; +}; + +struct DecoderConfig { + int64_t codebook_size = 0; + int64_t codebook_dim = 0; + int64_t latent_dim = 0; + int64_t hidden_size = 0; + int64_t intermediate_size = 0; + int64_t decoder_dim = 0; + int64_t num_heads = 0; + int64_t num_kv_heads = 0; + int64_t num_layers = 0; + int64_t num_quantizers = 0; + int64_t num_semantic_quantizers = 1; + int64_t sliding_window = 0; + int64_t head_dim = 0; + float rope_theta = 10000.0F; + float rms_norm_eps = 1.0e-5F; + std::vector upsample_rates; + std::vector upsampling_ratios; +}; + +} // namespace + +struct BreezeSpeechDecoderWeights { + std::shared_ptr store; + DecoderConfig config; + std::vector semantic_codebooks; + std::vector acoustic_codebooks; + modules::LinearWeights semantic_output_proj; + modules::LinearWeights acoustic_output_proj; + modules::Conv1dWeights pre_conv; + modules::LinearWeights transformer_input_proj; + std::vector transformer_layers; + modules::NormWeights transformer_norm; + modules::LinearWeights transformer_output_proj; + std::vector upsample_stages; + modules::Conv1dWeights decoder_input_conv; + std::vector decoder_blocks; + std::vector output_alpha; + std::vector output_beta; + modules::Conv1dWeights output_conv; +}; + +namespace { + +core::TensorValue normalized_codebook( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t size, + int64_t dim) { + const auto cluster_usage = source.require_f32(prefix + "cluster_usage", {size}); + const auto embedding_sum = source.require_f32(prefix + "embedding_sum", {size, dim}); + std::vector embedding(embedding_sum.size(), 0.0F); + for (int64_t code = 0; code < size; ++code) { + const float denom = std::max(cluster_usage[static_cast(code)], kCodebookEps); + for (int64_t col = 0; col < dim; ++col) { + const size_t offset = static_cast(code * dim + col); + embedding[offset] = embedding_sum[offset] / denom; + } + } + return store.make_f32(core::TensorShape::from_dims({size, dim}), embedding); +} + +DecoderConfig load_decoder_config(const BreezeTTSAssets & assets) { + const auto root = assets.resources.parse_json("audio_tokenizer_config_json"); + const auto & decoder = root.require("decoder_config"); + DecoderConfig config; + config.codebook_size = json::require_i64(decoder, "codebook_size"); + config.codebook_dim = json::require_i64(decoder, "codebook_dim"); + config.latent_dim = json::require_i64(decoder, "latent_dim"); + config.hidden_size = json::require_i64(decoder, "hidden_size"); + config.intermediate_size = json::require_i64(decoder, "intermediate_size"); + config.decoder_dim = json::require_i64(decoder, "decoder_dim"); + config.num_heads = json::require_i64(decoder, "num_attention_heads"); + config.num_kv_heads = json::require_i64(decoder, "num_key_value_heads"); + config.num_layers = json::require_i64(decoder, "num_hidden_layers"); + config.num_quantizers = json::require_i64(decoder, "num_quantizers"); + config.num_semantic_quantizers = json::require_i64(decoder, "num_semantic_quantizers"); + config.sliding_window = json::require_i64(decoder, "sliding_window"); + config.head_dim = json::require_i64(decoder, "head_dim"); + config.rope_theta = json::optional_f32(decoder, "rope_theta", config.rope_theta); + config.rms_norm_eps = json::optional_f32(decoder, "rms_norm_eps", config.rms_norm_eps); + config.upsample_rates = json::require_i64_array(decoder, "upsample_rates"); + config.upsampling_ratios = json::require_i64_array(decoder, "upsampling_ratios"); + return config; +} + +modules::LinearWeights load_conv1x1_as_linear( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t input_dim, + int64_t output_dim) { + modules::LinearWeights weights; + const auto data = source.require_tensor_as_shape( + prefix + ".weight", + storage_type, + {output_dim, input_dim, 1}, + {output_dim, input_dim}); + weights.weight = store.make_tensor(data.shape, data.type, data.bytes.data(), data.bytes.size()); + return weights; +} + +std::shared_ptr load_weights( + const BreezeTTSAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + assets::TensorStorageType linear_weight_storage_type, + assets::TensorStorageType conv_weight_storage_type) { + const auto & source = *assets.weights; + auto weights = std::make_shared(); + weights->store = std::make_shared( + backend, + backend_type, + "breeze_tts.speech_decoder.weights", + 32ull * 1024ull * 1024ull); + weights->config = load_decoder_config(assets); + const auto & config = weights->config; + const int64_t split_dim = config.codebook_dim / 2; + + for (int64_t layer = 0; layer < config.num_semantic_quantizers; ++layer) { + const std::string prefix = "codec_model.decoder.quantizer.rvq_first.vq.layers." + std::to_string(layer) + "._codebook."; + weights->semantic_codebooks.push_back(normalized_codebook(*weights->store, source, prefix, config.codebook_size, split_dim)); + } + for (int64_t layer = 0; layer < config.num_quantizers - config.num_semantic_quantizers; ++layer) { + const std::string prefix = "codec_model.decoder.quantizer.rvq_rest.vq.layers." + std::to_string(layer) + "._codebook."; + weights->acoustic_codebooks.push_back(normalized_codebook(*weights->store, source, prefix, config.codebook_size, split_dim)); + } + weights->semantic_output_proj = load_conv1x1_as_linear( + *weights->store, + source, + "codec_model.decoder.quantizer.rvq_first.output_proj", + linear_weight_storage_type, + split_dim, + config.hidden_size); + weights->acoustic_output_proj = load_conv1x1_as_linear( + *weights->store, + source, + "codec_model.decoder.quantizer.rvq_rest.output_proj", + linear_weight_storage_type, + split_dim, + config.hidden_size); + weights->pre_conv = binding::conv1d_from_source( + *weights->store, + source, + "codec_model.decoder.pre_conv.conv", + conv_weight_storage_type, + config.latent_dim, + config.hidden_size, + 3, + true); + weights->transformer_input_proj = binding::linear_from_source( + *weights->store, + source, + "codec_model.decoder.pre_transformer.input_proj", + linear_weight_storage_type, + config.hidden_size, + config.latent_dim, + true); + for (int64_t layer = 0; layer < config.num_layers; ++layer) { + const std::string prefix = "codec_model.decoder.pre_transformer.layers." + std::to_string(layer); + TransformerLayerWeights block; + block.input_norm = binding::norm_weight_from_source(*weights->store, source, prefix + ".input_layernorm", config.hidden_size); + block.post_norm = binding::norm_weight_from_source(*weights->store, source, prefix + ".post_attention_layernorm", config.hidden_size); + block.attention.q_weight = weights->store->load_tensor( + source, + prefix + ".self_attn.q_proj.weight", + linear_weight_storage_type, + {config.num_heads * config.head_dim, config.hidden_size}); + block.attention.k_weight = weights->store->load_tensor( + source, + prefix + ".self_attn.k_proj.weight", + linear_weight_storage_type, + {config.num_kv_heads * config.head_dim, config.hidden_size}); + block.attention.v_weight = weights->store->load_tensor( + source, + prefix + ".self_attn.v_proj.weight", + linear_weight_storage_type, + {config.num_kv_heads * config.head_dim, config.hidden_size}); + block.attention.out_weight = weights->store->load_tensor( + source, + prefix + ".self_attn.o_proj.weight", + linear_weight_storage_type, + {config.hidden_size, config.num_heads * config.head_dim}); + block.mlp.gate_proj = binding::linear_from_source( + *weights->store, + source, + prefix + ".mlp.gate_proj", + linear_weight_storage_type, + config.intermediate_size, + config.hidden_size, + false); + block.mlp.up_proj = binding::linear_from_source( + *weights->store, + source, + prefix + ".mlp.up_proj", + linear_weight_storage_type, + config.intermediate_size, + config.hidden_size, + false); + block.mlp.down_proj = binding::linear_from_source( + *weights->store, + source, + prefix + ".mlp.down_proj", + linear_weight_storage_type, + config.hidden_size, + config.intermediate_size, + false); + block.attn_scale = binding::layer_scale_from_named_source(*weights->store, source, prefix + ".self_attn_layer_scale.scale"); + block.mlp_scale = binding::layer_scale_from_named_source(*weights->store, source, prefix + ".mlp_layer_scale.scale"); + weights->transformer_layers.push_back(std::move(block)); + } + weights->transformer_norm = binding::norm_weight_from_source( + *weights->store, + source, + "codec_model.decoder.pre_transformer.norm", + config.hidden_size); + weights->transformer_output_proj = binding::linear_from_source( + *weights->store, + source, + "codec_model.decoder.pre_transformer.output_proj", + linear_weight_storage_type, + config.latent_dim, + config.hidden_size, + true); + + for (size_t i = 0; i < config.upsampling_ratios.size(); ++i) { + const std::string prefix = "codec_model.decoder.upsample." + std::to_string(i); + UpsampleStageWeights stage; + stage.upconv = binding::conv_transpose1d_from_source( + *weights->store, + source, + prefix + ".0.conv", + conv_weight_storage_type, + config.latent_dim, + config.latent_dim, + config.upsampling_ratios[i], + true); + stage.convnext.dwconv = binding::conv1d_from_source( + *weights->store, + source, + prefix + ".1.dwconv.conv", + conv_weight_storage_type, + config.latent_dim, + 1, + 7, + true); + stage.convnext.norm = binding::norm_from_source(*weights->store, source, prefix + ".1.norm", config.latent_dim); + stage.convnext.pwconv1 = binding::linear_from_source( + *weights->store, + source, + prefix + ".1.pwconv1", + linear_weight_storage_type, + config.latent_dim * 4, + config.latent_dim, + true); + stage.convnext.pwconv2 = binding::linear_from_source( + *weights->store, + source, + prefix + ".1.pwconv2", + linear_weight_storage_type, + config.latent_dim, + config.latent_dim * 4, + true); + stage.convnext.gamma = binding::layer_scale_from_named_source(*weights->store, source, prefix + ".1.gamma"); + weights->upsample_stages.push_back(std::move(stage)); + } + + weights->decoder_input_conv = binding::conv1d_from_source( + *weights->store, + source, + "codec_model.decoder.decoder.0.conv", + conv_weight_storage_type, + config.decoder_dim, + config.latent_dim, + 7, + true); + int64_t channels = config.decoder_dim; + for (size_t i = 0; i < config.upsample_rates.size(); ++i) { + const std::string prefix = "codec_model.decoder.decoder." + std::to_string(i + 1) + ".block"; + const int64_t out_channels = channels / 2; + DecoderBlockWeights block; + block.input_alpha = source.require_f32(prefix + ".0.alpha", {channels}); + block.input_beta = source.require_f32(prefix + ".0.beta", {channels}); + block.upconv = binding::conv_transpose1d_from_source( + *weights->store, + source, + prefix + ".1.conv", + conv_weight_storage_type, + channels, + out_channels, + config.upsample_rates[i] * 2, + true); + for (int unit_index = 0; unit_index < 3; ++unit_index) { + const std::string unit = prefix + "." + std::to_string(unit_index + 2); + ResidualUnitWeights residual; + residual.act1_alpha = source.require_f32(unit + ".act1.alpha", {out_channels}); + residual.act1_beta = source.require_f32(unit + ".act1.beta", {out_channels}); + residual.conv1 = binding::conv1d_from_source( + *weights->store, + source, + unit + ".conv1.conv", + conv_weight_storage_type, + out_channels, + out_channels, + 7, + true); + residual.act2_alpha = source.require_f32(unit + ".act2.alpha", {out_channels}); + residual.act2_beta = source.require_f32(unit + ".act2.beta", {out_channels}); + residual.conv2 = binding::conv1d_from_source( + *weights->store, + source, + unit + ".conv2.conv", + conv_weight_storage_type, + out_channels, + out_channels, + 1, + true); + block.residual_units.push_back(std::move(residual)); + } + weights->decoder_blocks.push_back(std::move(block)); + channels = out_channels; + } + weights->output_alpha = source.require_f32("codec_model.decoder.decoder.5.alpha", {channels}); + weights->output_beta = source.require_f32("codec_model.decoder.decoder.5.beta", {channels}); + weights->output_conv = binding::conv1d_from_source( + *weights->store, + source, + "codec_model.decoder.decoder.6.conv", + conv_weight_storage_type, + 1, + channels, + 7, + true); + weights->store->upload(); + return weights; +} + +core::TensorValue causal_conv1d( + core::ModuleBuildContext & build_ctx, + const core::TensorValue & input, + const modules::Conv1dWeights & weights, + int64_t out_channels, + int64_t kernel, + int64_t stride, + int64_t dilation, + int64_t groups, + bool use_bias) { + const int64_t in_channels = input.shape.dims[1]; + const int64_t kernel_extent = (kernel - 1) * dilation + 1; + const int64_t left_pad = kernel_extent - stride; + const int64_t length = input.shape.dims[2]; + const float n_frames = static_cast(length - kernel_extent + left_pad) / static_cast(stride) + 1.0F; + const int64_t ideal_length = + (static_cast(std::ceil(n_frames)) - 1) * stride + (kernel_extent - left_pad); + const int64_t right_pad = std::max(0, ideal_length - length); + auto * padded = ggml_pad_ext( + build_ctx.ggml, + input.tensor, + static_cast(left_pad), + static_cast(right_pad), + 0, + 0, + 0, + 0, + 0, + 0); + auto padded_value = core::wrap_tensor( + padded, + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], input.shape.dims[2] + left_pad + right_pad}), + GGML_TYPE_F32); + if (weights.weight.type != GGML_TYPE_F32 && weights.weight.type != GGML_TYPE_F16) { + throw std::runtime_error( + std::string("Breeze speech decoder depthwise conv does not support weight type: ") + + ggml_type_name(weights.weight.type)); + } + ggml_tensor * result = nullptr; + if (groups == in_channels) { + ggml_tensor * bias = nullptr; + if (use_bias) { + if (!weights.bias.has_value()) { + throw std::runtime_error("Breeze speech decoder depthwise conv requires bias"); + } + bias = core::reshape_tensor(build_ctx, *weights.bias, core::TensorShape::from_dims({out_channels, 1})).tensor; + } + for (int64_t batch = 0; batch < padded_value.shape.dims[0]; ++batch) { + auto * batch_input = ggml_view_2d( + build_ctx.ggml, + padded, + padded->ne[0], + padded->ne[1], + padded->nb[1], + static_cast(batch) * padded->nb[2]); + auto * batch_output = ggml_conv_1d_dw( + build_ctx.ggml, + weights.weight.tensor, + core::has_backend_addressable_layout(batch_input) ? batch_input : ggml_cont(build_ctx.ggml, batch_input), + static_cast(stride), + 0, + static_cast(dilation)); + if (bias != nullptr) { + batch_output = ggml_add(build_ctx.ggml, batch_output, bias); + } + batch_output = ggml_reshape_3d(build_ctx.ggml, batch_output, batch_output->ne[0], batch_output->ne[1], 1); + result = result == nullptr ? batch_output : ggml_concat(build_ctx.ggml, result, batch_output, 2); + } + return core::wrap_tensor( + result, + core::TensorShape::from_dims({input.shape.dims[0], out_channels, result->ne[0]}), + GGML_TYPE_F32); + } + return modules::Conv1dModule({ + in_channels, + out_channels, + kernel, + static_cast(stride), + 0, + static_cast(dilation), + use_bias, + }).build(build_ctx, padded_value, weights); +} + +core::TensorValue causal_conv_transpose1d( + core::ModuleBuildContext & build_ctx, + const core::TensorValue & input, + const modules::ConvTranspose1dWeights & weights, + int64_t out_channels, + int64_t kernel, + int64_t stride, + bool use_bias) { + const int64_t right_trim = kernel - stride; + auto output_bct = modules::ConvTranspose1dModule({ + input.shape.dims[1], + out_channels, + kernel, + static_cast(stride), + 0, + 1, + use_bias, + }).build(build_ctx, input, weights); + if (right_trim <= 0) { + return output_bct; + } + const int64_t trimmed_frames = output_bct.tensor->ne[0] - right_trim; + return core::wrap_tensor( + ggml_cont( + build_ctx.ggml, + ggml_view_3d( + build_ctx.ggml, + output_bct.tensor, + trimmed_frames, + out_channels, + input.shape.dims[0], + output_bct.tensor->nb[1], + output_bct.tensor->nb[2], + 0)), + core::TensorShape::from_dims({input.shape.dims[0], out_channels, trimmed_frames}), + GGML_TYPE_F32); +} + +core::TensorValue snake_beta( + core::ModuleBuildContext & build_ctx, + const core::TensorValue & input, + core::ConstantTensorCache & constants, + const std::vector & alpha, + const std::vector & beta) { + std::vector alpha_exp_values(alpha.size()); + std::transform(alpha.begin(), alpha.end(), alpha_exp_values.begin(), [](float value) { return std::exp(value); }); + std::vector inv_beta_exp_values(beta.size()); + std::transform(beta.begin(), beta.end(), inv_beta_exp_values.begin(), [](float value) { + return 1.0F / (std::exp(value) + 1.0e-9F); + }); + auto alpha_exp = constants.make_f32( + core::TensorShape::from_dims({1, static_cast(alpha.size()), 1}), + alpha_exp_values); + auto inv_beta_exp = constants.make_f32( + core::TensorShape::from_dims({1, static_cast(beta.size()), 1}), + inv_beta_exp_values); + auto * periodic = ggml_sqr( + build_ctx.ggml, + ggml_sin(build_ctx.ggml, ggml_mul(build_ctx.ggml, input.tensor, alpha_exp.tensor))); + return core::wrap_tensor( + ggml_add(build_ctx.ggml, input.tensor, ggml_mul(build_ctx.ggml, periodic, inv_beta_exp.tensor)), + input.shape, + GGML_TYPE_F32); +} + +core::TensorValue quantizer_decode( + ggml_context * ctx, + core::ModuleBuildContext & build_ctx, + ggml_tensor * codes_t_q_b, + const BreezeSpeechDecoderWeights & weights) { + const auto & config = weights.config; + const int64_t split_dim = config.codebook_dim / 2; + core::TensorValue semantic_sum; + for (int64_t group = 0; group < config.num_semantic_quantizers; ++group) { + auto * code_slice = ggml_view_2d( + ctx, + codes_t_q_b, + codes_t_q_b->ne[0], + codes_t_q_b->ne[2], + codes_t_q_b->nb[2], + static_cast(group) * codes_t_q_b->nb[1]); + auto indices = core::wrap_tensor( + code_slice, + core::TensorShape::from_dims({code_slice->ne[1], code_slice->ne[0]}), + GGML_TYPE_I32); + indices = core::ensure_backend_addressable_layout(build_ctx, indices); + auto decoded = modules::CodebookLookupModule({config.codebook_size, split_dim}) + .build(build_ctx, indices, weights.semantic_codebooks[static_cast(group)]); + semantic_sum = semantic_sum.valid() ? modules::AddModule{}.build(build_ctx, semantic_sum, decoded) : decoded; + } + semantic_sum = modules::LinearModule(binding::linear_config(split_dim, config.hidden_size, false)) + .build(build_ctx, semantic_sum, weights.semantic_output_proj); + + core::TensorValue acoustic_sum; + for (int64_t group = 0; group < config.num_quantizers - config.num_semantic_quantizers; ++group) { + const int64_t source_group = config.num_semantic_quantizers + group; + auto * code_slice = ggml_view_2d( + ctx, + codes_t_q_b, + codes_t_q_b->ne[0], + codes_t_q_b->ne[2], + codes_t_q_b->nb[2], + static_cast(source_group) * codes_t_q_b->nb[1]); + auto indices = core::wrap_tensor( + code_slice, + core::TensorShape::from_dims({code_slice->ne[1], code_slice->ne[0]}), + GGML_TYPE_I32); + indices = core::ensure_backend_addressable_layout(build_ctx, indices); + auto decoded = modules::CodebookLookupModule({config.codebook_size, split_dim}) + .build(build_ctx, indices, weights.acoustic_codebooks[static_cast(group)]); + acoustic_sum = acoustic_sum.valid() ? modules::AddModule{}.build(build_ctx, acoustic_sum, decoded) : decoded; + } + acoustic_sum = modules::LinearModule(binding::linear_config(split_dim, config.hidden_size, false)) + .build(build_ctx, acoustic_sum, weights.acoustic_output_proj); + return modules::AddModule{}.build(build_ctx, semantic_sum, acoustic_sum); +} + +core::TensorValue attention( + ggml_context * ctx, + core::ModuleBuildContext & build_ctx, + const core::TensorValue & input, + ggml_tensor * positions, + const core::TensorValue & attention_mask, + const modules::AttentionWeights & weights, + const DecoderConfig & config) { + const int64_t kv_repeat = config.num_heads / config.num_kv_heads; + auto q_value = modules::LinearModule(binding::linear_config(config.hidden_size, config.num_heads * config.head_dim, false)) + .build(build_ctx, input, {weights.q_weight, weights.q_bias}); + auto k_value = modules::LinearModule(binding::linear_config(config.hidden_size, config.num_kv_heads * config.head_dim, false)) + .build(build_ctx, input, {weights.k_weight, weights.k_bias}); + auto v_value = modules::LinearModule(binding::linear_config(config.hidden_size, config.num_kv_heads * config.head_dim, false)) + .build(build_ctx, input, {weights.v_weight, weights.v_bias}); + auto * q = q_value.tensor; + auto * k = k_value.tensor; + auto * v = v_value.tensor; + const int64_t seq = q->ne[1]; + const int64_t batch = q->ne[2]; + q = ggml_reshape_4d(ctx, q, config.head_dim, config.num_heads, seq, batch); + k = ggml_reshape_4d(ctx, k, config.head_dim, config.num_kv_heads, seq, batch); + v = ggml_reshape_4d(ctx, v, config.head_dim, config.num_kv_heads, seq, batch); + auto position_value = core::wrap_tensor(positions, core::TensorShape::from_dims({seq}), GGML_TYPE_I32); + q = modules::RoPEModule({ + config.head_dim, + GGML_ROPE_TYPE_NEOX, + config.rope_theta, + }).build( + build_ctx, + core::wrap_tensor(q, core::TensorShape::from_dims({batch, seq, config.num_heads, config.head_dim}), GGML_TYPE_F32), + position_value) + .tensor; + k = modules::RoPEModule({ + config.head_dim, + GGML_ROPE_TYPE_NEOX, + config.rope_theta, + }).build( + build_ctx, + core::wrap_tensor(k, core::TensorShape::from_dims({batch, seq, config.num_kv_heads, config.head_dim}), GGML_TYPE_F32), + position_value) + .tensor; + auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, 4}).build( + build_ctx, + core::wrap_tensor(q, core::TensorShape::from_dims({batch, seq, config.num_heads, config.head_dim}), GGML_TYPE_F32)); + auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, 4}).build( + build_ctx, + core::wrap_tensor(k, core::TensorShape::from_dims({batch, seq, config.num_kv_heads, config.head_dim}), GGML_TYPE_F32)); + auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, 4}).build( + build_ctx, + core::wrap_tensor(v, core::TensorShape::from_dims({batch, seq, config.num_kv_heads, config.head_dim}), GGML_TYPE_F32)); + if (kv_repeat > 1) { + std::vector repeated_k; + std::vector repeated_v; + repeated_k.reserve(static_cast(config.num_heads)); + repeated_v.reserve(static_cast(config.num_heads)); + for (int64_t head = 0; head < config.num_kv_heads; ++head) { + auto one_k = modules::SliceModule({1, head, 1}).build(build_ctx, k_heads); + auto one_v = modules::SliceModule({1, head, 1}).build(build_ctx, v_heads); + for (int64_t repeat = 0; repeat < kv_repeat; ++repeat) { + repeated_k.push_back(one_k); + repeated_v.push_back(one_v); + } + } + k_heads = repeated_k.front(); + v_heads = repeated_v.front(); + for (size_t index = 1; index < repeated_k.size(); ++index) { + k_heads = modules::ConcatModule({1}).build(build_ctx, k_heads, repeated_k[index]); + v_heads = modules::ConcatModule({1}).build(build_ctx, v_heads, repeated_v[index]); + } + } + auto context = modules::ScaledDotProductAttentionModule({ + config.head_dim, + modules::ScaledDotProductAttentionLowering::Flash, + GGML_PREC_F32, + modules::AttentionCausality::NonCausal, + }).build( + build_ctx, + q_heads, + k_heads, + v_heads, + attention_mask); + context = core::ensure_backend_addressable_layout(build_ctx, context); + return modules::LinearModule(binding::linear_config(config.num_heads * config.head_dim, config.hidden_size, false)) + .build( + build_ctx, + core::reshape_tensor(build_ctx, context, core::TensorShape::from_dims({batch, seq, config.num_heads * config.head_dim})), + {weights.out_weight, weights.out_bias}); +} + +core::TensorValue convnext( + core::ModuleBuildContext & build_ctx, + const core::TensorValue & input_bct, + const ConvNeXtWeights & weights) { + const int64_t channels = input_bct.shape.dims[1]; + auto hidden = causal_conv1d(build_ctx, input_bct, weights.dwconv, channels, 7, 1, 1, channels, true); + hidden = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build_ctx, hidden); + hidden = modules::LayerNormModule({channels, 1.0e-6F, true, true}) + .build(build_ctx, hidden, weights.norm); + hidden = modules::LinearModule(binding::linear_config(channels, channels * 4, true)) + .build(build_ctx, hidden, weights.pwconv1); + hidden = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(build_ctx, hidden); + hidden = modules::LinearModule(binding::linear_config(channels * 4, channels, true)) + .build(build_ctx, hidden, weights.pwconv2); + hidden = modules::LayerScaleModule{}.build( + build_ctx, + hidden, + weights.gamma); + hidden = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build_ctx, hidden); + return modules::AddModule{}.build(build_ctx, input_bct, hidden); +} + +core::TensorValue residual_unit( + core::ModuleBuildContext & build_ctx, + const core::TensorValue & input, + const ResidualUnitWeights & weights, + int64_t dilation, + core::ConstantTensorCache & constants) { + const int64_t channels = input.shape.dims[1]; + auto hidden = snake_beta( + build_ctx, + input, + constants, + weights.act1_alpha, + weights.act1_beta); + hidden = causal_conv1d(build_ctx, hidden, weights.conv1, channels, 7, 1, dilation, 1, true); + hidden = snake_beta( + build_ctx, + hidden, + constants, + weights.act2_alpha, + weights.act2_beta); + hidden = causal_conv1d(build_ctx, hidden, weights.conv2, channels, 1, 1, 1, 1, true); + return modules::AddModule{}.build(build_ctx, input, hidden); +} + +std::vector make_mask(int64_t frames, int64_t window) { + std::vector mask(static_cast(frames * frames), kMaskNegInf); + for (int64_t q = 0; q < frames; ++q) { + const int64_t min_k = std::max(0, q - window + 1); + for (int64_t k = min_k; k <= q; ++k) { + mask[static_cast(k + frames * q)] = 0.0F; + } + } + return mask; +} + +int graph_node_capacity(const DecoderConfig & config) { + return static_cast(4096 + config.num_layers * config.num_heads * 16 + config.upsample_rates.size() * 512); +} + +} // namespace + +class BreezeSpeechDecoderGraph { +public: + BreezeSpeechDecoderGraph( + std::shared_ptr weights, + int64_t code_frames, + core::ExecutionContext & execution_context, + core::ConstantTensorCache & constants, + size_t graph_arena_bytes) + : weights_(std::move(weights)), + code_frames_(code_frames), + backend_(execution_context.backend()), + compute_threads_(std::max(1, execution_context.config().threads)) { + if (weights_ == nullptr) { + throw std::runtime_error("Breeze speech decoder graph requires weights"); + } + if (code_frames_ <= 0) { + throw std::runtime_error("Breeze speech decoder graph requires positive frame count"); + } + if (backend_ == nullptr) { + throw std::runtime_error("Breeze speech decoder backend is not initialized"); + } + const auto & config = weights_->config; + waveform_frames_ = code_frames_; + for (const auto factor : config.upsampling_ratios) { + waveform_frames_ *= factor; + } + for (const auto factor : config.upsample_rates) { + waveform_frames_ *= factor; + } + + ggml_init_params params{ + /*.mem_size =*/ graph_arena_bytes, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Breeze speech decoder ggml context"); + } + + codes_ = ggml_new_tensor_3d(ctx_.get(), GGML_TYPE_I32, code_frames_, config.num_quantizers, 1); + ggml_set_input(codes_); + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, code_frames_); + ggml_set_input(positions_); + mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, code_frames_, code_frames_, 1, 1); + ggml_set_input(mask_); + + core::ModuleBuildContext build_ctx{ + ctx_.get(), + "breeze_tts.speech_decoder", + execution_context.backend_type(), + }; + constants.begin_graph(); + auto hidden = quantizer_decode(ctx_.get(), build_ctx, codes_, *weights_); + hidden = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build_ctx, hidden); + hidden = causal_conv1d(build_ctx, hidden, weights_->pre_conv, config.latent_dim, 3, 1, 1, 1, true); + hidden = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build_ctx, hidden); + hidden = modules::LinearModule(binding::linear_config(config.latent_dim, config.hidden_size, true)) + .build(build_ctx, hidden, weights_->transformer_input_proj); + for (const auto & layer : weights_->transformer_layers) { + auto attn_in = modules::RMSNormModule({config.hidden_size, config.rms_norm_eps, true, false}) + .build(build_ctx, hidden, layer.input_norm); + auto attention_mask = core::wrap_tensor( + mask_, + core::TensorShape::from_dims({1, 1, code_frames_, code_frames_}), + GGML_TYPE_F16); + auto attn_out = attention(ctx_.get(), build_ctx, attn_in, positions_, attention_mask, layer.attention, config); + attn_out = modules::LayerScaleModule{}.build( + build_ctx, + attn_out, + layer.attn_scale); + hidden = modules::AddModule{}.build(build_ctx, hidden, attn_out); + auto mlp_in = modules::RMSNormModule({config.hidden_size, config.rms_norm_eps, true, false}) + .build(build_ctx, hidden, layer.post_norm); + auto mlp_out = modules::GatedFeedForwardModule({ + config.hidden_size, + config.intermediate_size, + false, + modules::GatedFeedForwardActivation::Silu, + }).build(build_ctx, mlp_in, layer.mlp); + mlp_out = modules::LayerScaleModule{}.build( + build_ctx, + mlp_out, + layer.mlp_scale); + hidden = modules::AddModule{}.build(build_ctx, hidden, mlp_out); + } + hidden = modules::RMSNormModule({config.hidden_size, config.rms_norm_eps, true, false}) + .build(build_ctx, hidden, weights_->transformer_norm); + hidden = modules::LinearModule(binding::linear_config(config.hidden_size, config.latent_dim, true)) + .build(build_ctx, hidden, weights_->transformer_output_proj); + hidden = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build_ctx, hidden); + for (size_t i = 0; i < weights_->upsample_stages.size(); ++i) { + const auto & stage = weights_->upsample_stages[i]; + hidden = causal_conv_transpose1d( + build_ctx, + hidden, + stage.upconv, + config.latent_dim, + config.upsampling_ratios[i], + config.upsampling_ratios[i], + true); + hidden = convnext(build_ctx, hidden, stage.convnext); + } + hidden = causal_conv1d(build_ctx, hidden, weights_->decoder_input_conv, config.decoder_dim, 7, 1, 1, 1, true); + int64_t decoder_channels = config.decoder_dim; + for (size_t block_index = 0; block_index < weights_->decoder_blocks.size(); ++block_index) { + const auto & block = weights_->decoder_blocks[block_index]; + const int64_t out_channels = decoder_channels / 2; + hidden = snake_beta( + build_ctx, + hidden, + constants, + block.input_alpha, + block.input_beta); + hidden = causal_conv_transpose1d( + build_ctx, + hidden, + block.upconv, + out_channels, + config.upsample_rates[block_index] * 2, + config.upsample_rates[block_index], + true); + for (size_t unit_index = 0; unit_index < block.residual_units.size(); ++unit_index) { + const int64_t dilation = unit_index == 0 ? 1 : unit_index == 1 ? 3 : 9; + hidden = residual_unit( + build_ctx, + hidden, + block.residual_units[unit_index], + dilation, + constants); + } + decoder_channels = out_channels; + } + hidden = snake_beta( + build_ctx, + hidden, + constants, + weights_->output_alpha, + weights_->output_beta); + output_ = ggml_clamp(ctx_.get(), causal_conv1d(build_ctx, hidden, weights_->output_conv, 1, 7, 1, 1, 1, true).tensor, -1.0F, 1.0F); + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), graph_node_capacity(config), false); + ggml_build_forward_expand(graph_, output_); + constants.finish_graph(); + constants.ensure_uploaded(); + + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + if (gallocr_ == nullptr || !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + throw std::runtime_error("failed to allocate Breeze speech decoder graph"); + } + positions_data_.resize(static_cast(code_frames_)); + for (int64_t i = 0; i < code_frames_; ++i) { + positions_data_[static_cast(i)] = static_cast(i); + } + const auto mask = make_mask(code_frames_, config.sliding_window); + mask_f16_data_.resize(mask.size()); + for (size_t index = 0; index < mask.size(); ++index) { + mask_f16_data_[index] = ggml_fp32_to_fp16(mask[index]); + } + upload_static_inputs(); + } + + ~BreezeSpeechDecoderGraph() { + engine::core::release_backend_graph_resources(backend_, graph_); + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + } + + bool matches( + const BreezeSpeechDecoderWeights & weights, + int64_t code_frames, + ggml_backend_t backend, + int threads) const { + const bool frame_match = code_frames_ == code_frames; + return weights_.get() == &weights && frame_match && backend_ == backend && + compute_threads_ == std::max(1, threads); + } + + std::vector run(const int32_t * codes, size_t code_count) { + const int64_t input_frames = static_cast(code_count / weights_->config.num_quantizers); + const size_t expected = static_cast(code_frames_ * weights_->config.num_quantizers); + if (code_count % static_cast(weights_->config.num_quantizers) != 0 || input_frames <= 0 || input_frames > code_frames_) { + throw std::runtime_error("Breeze speech decoder code count exceeds graph capacity"); + } + // Cached GGML graphs may reuse backend allocations whose input contents are + // not guaranteed to survive a prior execution. Restore every declared input, + // not only the request-varying codes, before replaying a retained graph. + upload_static_inputs(); + std::vector tensor_codes(expected, 0); + for (int64_t frame = 0; frame < input_frames; ++frame) { + for (int64_t group = 0; group < weights_->config.num_quantizers; ++group) { + tensor_codes[static_cast(frame + code_frames_ * group)] = + codes[static_cast(frame * weights_->config.num_quantizers + group)]; + } + } + ggml_backend_tensor_set(codes_, tensor_codes.data(), 0, tensor_codes.size() * sizeof(int32_t)); + core::set_backend_threads(backend_, compute_threads_); + const ggml_status status = engine::core::compute_backend_graph(backend_, graph_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Breeze speech decoder graph compute failed"); + } + std::vector audio(static_cast(waveform_frames_), 0.0F); + ggml_backend_tensor_get(output_, audio.data(), 0, audio.size() * sizeof(float)); + return audio; + } + +private: + void upload_static_inputs() { + ggml_backend_tensor_set( + positions_, + positions_data_.data(), + 0, + positions_data_.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + mask_, + mask_f16_data_.data(), + 0, + mask_f16_data_.size() * sizeof(ggml_fp16_t)); + } + + std::shared_ptr weights_; + int64_t code_frames_ = 0; + int64_t waveform_frames_ = 0; + ggml_backend_t backend_ = nullptr; + int compute_threads_ = 1; + std::unique_ptr ctx_; + ggml_tensor * codes_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * mask_ = nullptr; + ggml_tensor * output_ = nullptr; + std::vector positions_data_; + std::vector mask_f16_data_; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; +}; + +BreezeSpeechDecoderRuntime::BreezeSpeechDecoderRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution_context, + size_t graph_arena_bytes, + size_t constant_context_bytes, + assets::TensorStorageType linear_weight_storage_type, + assets::TensorStorageType conv_weight_storage_type) + : assets_(std::move(assets)), + execution_context_(&execution_context), + graph_arena_bytes_(graph_arena_bytes) { + if (assets_ == nullptr) { + throw std::runtime_error("BreezeTTS speech decoder requires assets"); + } + weights_ = load_weights( + *assets_, + execution_context_->backend(), + execution_context_->backend_type(), + linear_weight_storage_type, + conv_weight_storage_type); + constants_ = std::make_unique( + execution_context_->backend(), + std::max(1, execution_context_->config().threads), + "breeze_tts.speech_decoder.constants", + constant_context_bytes); +} + +BreezeSpeechDecoderRuntime::~BreezeSpeechDecoderRuntime() = default; + +runtime::AudioBuffer BreezeSpeechDecoderRuntime::decode(const BreezeSpeechCodes & codec_codes) const { + const auto total_start = Clock::now(); + if (codec_codes.frames <= 0 || codec_codes.code_groups != weights_->config.num_quantizers) { + throw std::runtime_error("Breeze speech decoder received invalid codec shape"); + } + if (static_cast(codec_codes.codes.size()) != codec_codes.frames * codec_codes.code_groups) { + throw std::runtime_error("Breeze speech decoder codec payload size mismatch"); + } + std::vector samples; + samples.reserve(static_cast(codec_codes.frames * kDecodeSamplesPerCode)); + for (int64_t start = 0; start < codec_codes.frames; start += kChunkCodes) { + const int64_t end = std::min(start + kChunkCodes, codec_codes.frames); + const int64_t context = start > kLeftContextCodes ? kLeftContextCodes : start; + const int64_t chunk_start = start - context; + const int64_t chunk_frames = end - chunk_start; + std::vector chunk(static_cast(chunk_frames * codec_codes.code_groups), 0); + for (int64_t frame = 0; frame < chunk_frames; ++frame) { + const int64_t src_frame = chunk_start + frame; + const auto src = codec_codes.codes.begin() + static_cast(src_frame * codec_codes.code_groups); + const auto dst = chunk.begin() + static_cast(frame * codec_codes.code_groups); + std::copy(src, src + codec_codes.code_groups, dst); + } + const int threads = std::max(1, execution_context_->config().threads); + auto * graph_slot = &graph_; +#if defined(ENGINE_HIP_STRIX_HALO_OPTIMIZATIONS) + const bool optimized_cache_enabled = + kStrixHaloGraphCacheEnabled && execution_context_->backend_type() == core::BackendType::Hip; + if (optimized_cache_enabled) { + for (size_t index = 0; index < kStrixHaloCachedChunkFrames.size(); ++index) { + if (chunk_frames == kStrixHaloCachedChunkFrames[index]) { + graph_slot = &optimized_graphs_[index]; + break; + } + } + } +#endif + auto & graph = *graph_slot; + const bool graph_rebuilt = + graph == nullptr || !graph->matches(*weights_, chunk_frames, execution_context_->backend(), threads); + if (graph_rebuilt) { + auto replacement = std::make_unique( + weights_, + chunk_frames, + *execution_context_, + *constants_, + graph_arena_bytes_); + graph = std::move(replacement); + } + auto decoded = graph->run(chunk.data(), chunk.size()); + const int64_t drop = context * kDecodeSamplesPerCode; + if (drop > static_cast(decoded.size())) { + throw std::runtime_error("Breeze speech decoder chunk context exceeds decoded waveform"); + } + const int64_t valid_samples = chunk_frames * kDecodeSamplesPerCode; + if (valid_samples < drop || valid_samples > static_cast(decoded.size())) { + throw std::runtime_error("Breeze speech decoder valid sample range exceeds decoded waveform"); + } + samples.insert( + samples.end(), + decoded.begin() + static_cast(drop), + decoded.begin() + static_cast(valid_samples)); + } + debug::timing_log_scalar("breeze_tts.speech_decoder.total_ms", engine::debug::elapsed_ms(total_start, Clock::now())); + return runtime::AudioBuffer{kSampleRate, 1, std::move(samples)}; +} + +runtime::AudioBuffer BreezeSpeechDecoderRuntime::decode_and_trim_reference( + const BreezeSpeechCodes & reference_codes, + const BreezeSpeechCodes & generated_codes) const { + if (reference_codes.code_groups != generated_codes.code_groups) { + throw std::runtime_error("Breeze speech decoder reference/generated code group mismatch"); + } + if (reference_codes.frames < 0 || generated_codes.frames < 0 || reference_codes.code_groups <= 0) { + throw std::runtime_error("Breeze speech decoder reference/generated code shape is invalid"); + } + if (reference_codes.frames > std::numeric_limits::max() - generated_codes.frames) { + throw std::runtime_error("Breeze speech decoder combined frame count is too large"); + } + BreezeSpeechCodes combined; + combined.frames = reference_codes.frames + generated_codes.frames; + combined.code_groups = reference_codes.code_groups; + if (combined.frames > std::numeric_limits::max() / combined.code_groups) { + throw std::runtime_error("Breeze speech decoder combined code count is too large"); + } + const int64_t combined_code_count = combined.frames * combined.code_groups; + if (static_cast(combined_code_count) > std::numeric_limits::max()) { + throw std::runtime_error("Breeze speech decoder combined code count exceeds host size limits"); + } + combined.codes.reserve(static_cast(combined_code_count)); + combined.codes.insert(combined.codes.end(), reference_codes.codes.begin(), reference_codes.codes.end()); + combined.codes.insert(combined.codes.end(), generated_codes.codes.begin(), generated_codes.codes.end()); + auto audio = decode(combined); + if (reference_codes.frames > std::numeric_limits::max() / kDecodeSamplesPerCode) { + throw std::runtime_error("Breeze speech decoder reference sample count is too large"); + } + const int64_t cut = reference_codes.frames * kDecodeSamplesPerCode; + if (static_cast(cut) > audio.samples.size()) { + throw std::runtime_error("Breeze speech decoder reference trim is out of range"); + } + audio.samples.erase(audio.samples.begin(), audio.samples.begin() + static_cast(cut)); + return audio; +} + +void BreezeSpeechDecoderRuntime::release_runtime_graphs() const { + graph_.reset(); + for (auto & graph : optimized_graphs_) { + graph.reset(); + } +} + +} // namespace engine::models::breeze_tts diff --git a/src/models/breeze_tts/speech_encoder.cpp b/src/models/breeze_tts/speech_encoder.cpp new file mode 100644 index 000000000..f4f7d34a5 --- /dev/null +++ b/src/models/breeze_tts/speech_encoder.cpp @@ -0,0 +1,649 @@ +#include "engine/models/breeze_tts/speech_encoder.h" + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/audio/conversion.h" +#include "engine/framework/core/backend.h" +#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/feed_forward.h" +#include "engine/framework/modules/attention/scaled_dot_product_attention.h" +#include "engine/framework/modules/attention/types.h" +#include "engine/framework/modules/conditioning_modules.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/streaming_conv_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/transformers/qwen_causal_decoder.h" +#include "engine/framework/modules/weight_binding.h" + +#include "engine/framework/core/constant_tensor_cache.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::breeze_tts { + +namespace { + +using Clock = std::chrono::steady_clock; +namespace binding = modules::binding; + +} // namespace + +constexpr int64_t kSampleRate = 24000; +constexpr int64_t kDownsampleRate = 1920; +constexpr int64_t kHiddenSize = 512; +constexpr int64_t kQuantizerDim = 256; +constexpr int64_t kCodebookSize = 2048; +constexpr int64_t kValidQuantizers = 16; +constexpr float kCodebookEps = 1.0e-5F; +constexpr std::array kEncoderConvConfigs{{ + {1, 64, 7, 1, 0, 1, true}, + {64, 128, 8, 4, 0, 1, true}, + {128, 256, 10, 5, 0, 1, true}, + {256, 512, 12, 6, 0, 1, true}, + {512, 1024, 16, 8, 0, 1, true}, + {1024, 512, 3, 1, 0, 1, true}, +}}; +constexpr std::array kEncoderConvPadModes{{ + modules::StreamingPadMode::Constant, + modules::StreamingPadMode::Constant, + modules::StreamingPadMode::Constant, + modules::StreamingPadMode::Constant, + modules::StreamingPadMode::Constant, + modules::StreamingPadMode::Constant, +}}; +constexpr modules::Conv1dConfig kDownsampleConvConfig{512, 512, 4, 2, 0, 1, false}; +constexpr modules::Conv1dConfig kSemanticProjectionConfig{512, 256, 1, 1, 0, 1, false}; +constexpr modules::Conv1dConfig kAcousticProjectionConfig{512, 256, 1, 1, 0, 1, false}; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct ResBlockWeights { + modules::Conv1dWeights conv1; + modules::Conv1dWeights conv2; +}; + +struct TransformerLayerWeights { + modules::AttentionWeights attention; + modules::FeedForwardWeights feed_forward; + modules::NormWeights norm1; + modules::NormWeights norm2; + modules::LayerScaleWeights scale1; + modules::LayerScaleWeights scale2; +}; + +struct BreezeSpeechEncoderWeights { + std::shared_ptr store; + std::vector encoder_convs; + std::vector residual_blocks; + std::vector transformer_layers; + modules::Conv1dWeights downsample; + modules::Conv1dWeights semantic_projection; + modules::Conv1dWeights acoustic_projection; + std::vector> semantic_codebooks; + std::vector> acoustic_codebooks; +}; + +core::TensorValue speech_conv( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::Conv1dWeights & weights, + modules::Conv1dConfig config, + modules::StreamingPadMode pad_mode) { + const int64_t effective_kernel = (config.kernel_size - 1) * config.dilation + 1; + const int64_t left_pad = effective_kernel - config.stride; + const int64_t right_pad = (config.stride - (input.shape.dims[2] % config.stride)) % config.stride; + auto padded = input; + if (left_pad > 0) { + auto prefix = modules::SliceModule({2, 0, 1}).build(ctx, input); + prefix = core::ensure_backend_addressable_layout(ctx, prefix); + if (pad_mode == modules::StreamingPadMode::Constant) { + prefix = core::wrap_tensor(ggml_scale(ctx.ggml, prefix.tensor, 0.0F), prefix.shape, GGML_TYPE_F32); + } + prefix = modules::RepeatModule({ + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], left_pad}), + }).build(ctx, prefix); + padded = modules::ConcatModule({2}).build(ctx, prefix, padded); + } + if (right_pad > 0) { + auto suffix = modules::SliceModule({2, input.shape.dims[2] - 1, 1}).build(ctx, input); + suffix = core::ensure_backend_addressable_layout(ctx, suffix); + if (pad_mode == modules::StreamingPadMode::Constant) { + suffix = core::wrap_tensor(ggml_scale(ctx.ggml, suffix.tensor, 0.0F), suffix.shape, GGML_TYPE_F32); + } + suffix = modules::RepeatModule({ + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], right_pad}), + }).build(ctx, suffix); + padded = modules::ConcatModule({2}).build(ctx, padded, suffix); + } + config.padding = 0; + return modules::Conv1dModule(config).build(ctx, padded, weights); +} + +core::TensorValue speech_residual_block( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const ResBlockWeights & block, + core::ConstantTensorCache &) { + const int64_t channels = input.shape.dims[1]; + auto x = modules::EluModule{}.build(ctx, input); + x = speech_conv(ctx, x, block.conv1, {channels, channels / 2, 3, 1, 0, 1, true}, modules::StreamingPadMode::Constant); + x = modules::EluModule{}.build(ctx, x); + x = speech_conv(ctx, x, block.conv2, {channels / 2, channels, 1, 1, 0, 1, true}, modules::StreamingPadMode::Constant); + return modules::AddModule{}.build(ctx, input, x); +} + +core::TensorValue mimi_self_attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const TransformerLayerWeights & weights, + const std::optional & attention_mask) { + constexpr int64_t kHeads = 8; + constexpr int64_t kHeadDim = 64; + auto q = modules::LinearModule(binding::linear_config(kHiddenSize, kHiddenSize, false)) + .build(ctx, input, {weights.attention.q_weight, weights.attention.q_bias}); + auto k = modules::LinearModule(binding::linear_config(kHiddenSize, kHiddenSize, false)) + .build(ctx, input, {weights.attention.k_weight, weights.attention.k_bias}); + auto v = modules::LinearModule(binding::linear_config(kHiddenSize, kHiddenSize, false)) + .build(ctx, input, {weights.attention.v_weight, weights.attention.v_bias}); + q = core::reshape_tensor( + ctx, + core::ensure_backend_addressable_layout(ctx, q), + core::TensorShape::from_dims({q.shape.dims[0], q.shape.dims[1], kHeads, kHeadDim})); + k = core::reshape_tensor( + ctx, + core::ensure_backend_addressable_layout(ctx, k), + core::TensorShape::from_dims({k.shape.dims[0], k.shape.dims[1], kHeads, kHeadDim})); + v = core::reshape_tensor( + ctx, + core::ensure_backend_addressable_layout(ctx, v), + core::TensorShape::from_dims({v.shape.dims[0], v.shape.dims[1], kHeads, kHeadDim})); + q = modules::RoPEModule({kHeadDim, GGML_ROPE_TYPE_NEOX}).build(ctx, q, positions); + k = modules::RoPEModule({kHeadDim, GGML_ROPE_TYPE_NEOX}).build(ctx, k, positions); + auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); + auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k); + auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v); + auto context = modules::ScaledDotProductAttentionModule({ + kHeadDim, + modules::ScaledDotProductAttentionLowering::Flash, + GGML_PREC_F32, + modules::AttentionCausality::Causal, + }).build(ctx, q_heads, k_heads, v_heads, attention_mask); + context = core::ensure_backend_addressable_layout(ctx, context); + context = core::reshape_tensor(ctx, context, core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], kHiddenSize})); + return modules::LinearModule(binding::linear_config(kHiddenSize, kHiddenSize, false)) + .build(ctx, context, {weights.attention.out_weight, weights.attention.out_bias}); +} + +core::TensorValue transformer_block( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const TransformerLayerWeights & weights, + const std::optional & attention_mask) { + const modules::LayerNormModule norm({kHiddenSize, 1.0e-5F, true, true}); + auto x = norm.build(ctx, input, weights.norm1); + auto attn_out = modules::LayerScaleModule{}.build( + ctx, + mimi_self_attention(ctx, x, positions, weights, attention_mask), + weights.scale1); + x = modules::AddModule{}.build(ctx, input, attn_out); + auto y = norm.build(ctx, x, weights.norm2); + y = modules::FeedForwardModule({ + kHiddenSize, + 2048, + false, + modules::GeluApproximation::ExactErf, + }).build(ctx, y, weights.feed_forward); + y = modules::LayerScaleModule{}.build(ctx, y, weights.scale2); + return modules::AddModule{}.build(ctx, x, y); +} + +std::vector codebook_embedding( + const assets::TensorSource & source, + const std::string & prefix) { + const auto cluster_usage = source.require_f32(prefix + "cluster_usage", {kCodebookSize}); + const auto embedding_sum = source.require_f32(prefix + "embed_sum", {kCodebookSize, kQuantizerDim}); + std::vector embedding(embedding_sum.size(), 0.0F); + for (int64_t code = 0; code < kCodebookSize; ++code) { + const float denom = std::max(cluster_usage[static_cast(code)], kCodebookEps); + for (int64_t dim = 0; dim < kQuantizerDim; ++dim) { + const size_t offset = static_cast(code * kQuantizerDim + dim); + embedding[offset] = embedding_sum[offset] / denom; + } + } + return embedding; +} + +int32_t nearest_code(const std::vector & residual, const std::vector & embedding) { + int32_t best = 0; + float best_distance = std::numeric_limits::infinity(); + for (int64_t code = 0; code < kCodebookSize; ++code) { + float distance = 0.0F; + for (int64_t dim = 0; dim < kQuantizerDim; ++dim) { + const float diff = + residual[static_cast(dim)] - embedding[static_cast(code * kQuantizerDim + dim)]; + distance += diff * diff; + } + if (distance < best_distance) { + best_distance = distance; + best = static_cast(code); + } + } + return best; +} + +std::vector quantize_projected( + const std::vector & semantic, + const std::vector & acoustic, + int64_t frames, + const BreezeSpeechEncoderWeights & weights) { + if (weights.semantic_codebooks.empty() || static_cast(weights.acoustic_codebooks.size()) < kValidQuantizers - 1) { + throw std::runtime_error("Breeze speech encoder has insufficient quantizer codebooks"); + } + if (static_cast(semantic.size()) != kQuantizerDim * frames || + static_cast(acoustic.size()) != kQuantizerDim * frames) { + throw std::runtime_error("Breeze speech encoder projected tensor size mismatch"); + } + + std::vector codes(static_cast(frames * kValidQuantizers), 0); + std::vector residual(static_cast(kQuantizerDim), 0.0F); + for (int64_t frame = 0; frame < frames; ++frame) { + for (int64_t dim = 0; dim < kQuantizerDim; ++dim) { + residual[static_cast(dim)] = semantic[static_cast(dim * frames + frame)]; + } + const int32_t semantic_code = nearest_code(residual, weights.semantic_codebooks[0]); + codes[static_cast(frame * kValidQuantizers)] = semantic_code; + + for (int64_t dim = 0; dim < kQuantizerDim; ++dim) { + residual[static_cast(dim)] = acoustic[static_cast(dim * frames + frame)]; + } + for (int64_t group = 1; group < kValidQuantizers; ++group) { + const auto & embedding = weights.acoustic_codebooks[static_cast(group - 1)]; + const int32_t code = nearest_code(residual, embedding); + codes[static_cast(frame * kValidQuantizers + group)] = code; + for (int64_t dim = 0; dim < kQuantizerDim; ++dim) { + residual[static_cast(dim)] -= embedding[static_cast(code * kQuantizerDim + dim)]; + } + } + } + return codes; +} + +std::shared_ptr load_weights( + const BreezeTTSAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + assets::TensorStorageType linear_weight_storage_type, + assets::TensorStorageType conv_weight_storage_type) { + const auto & source = *assets.weights; + auto weights = std::make_shared(); + weights->store = std::make_shared( + backend, + backend_type, + "breeze_tts.speech_encoder.weights", + 32ull * 1024ull * 1024ull); + + const char * conv_prefixes[] = { + "codec_model.encoder.encoder.layers.0.conv", + "codec_model.encoder.encoder.layers.3.conv", + "codec_model.encoder.encoder.layers.6.conv", + "codec_model.encoder.encoder.layers.9.conv", + "codec_model.encoder.encoder.layers.12.conv", + "codec_model.encoder.encoder.layers.14.conv", + }; + for (size_t i = 0; i < std::size(conv_prefixes); ++i) { + const auto & conv = kEncoderConvConfigs[i]; + weights->encoder_convs.push_back( + binding::conv1d_from_source( + *weights->store, + source, + conv_prefixes[i], + conv_weight_storage_type, + conv.out_channels, + conv.in_channels, + conv.kernel_size, + conv.use_bias)); + } + + const int residual_indices[] = {1, 4, 7, 10}; + const int64_t residual_channels[] = {64, 128, 256, 512}; + for (size_t i = 0; i < std::size(residual_indices); ++i) { + const int idx = residual_indices[i]; + const int64_t channels = residual_channels[i]; + ResBlockWeights block; + block.conv1 = binding::conv1d_from_source( + *weights->store, + source, + "codec_model.encoder.encoder.layers." + std::to_string(idx) + ".block.1.conv", + conv_weight_storage_type, + channels / 2, + channels, + 3, + true); + block.conv2 = binding::conv1d_from_source( + *weights->store, + source, + "codec_model.encoder.encoder.layers." + std::to_string(idx) + ".block.3.conv", + conv_weight_storage_type, + channels, + channels / 2, + 1, + true); + weights->residual_blocks.push_back(std::move(block)); + } + + for (int layer = 0; layer < 8; ++layer) { + const std::string prefix = "codec_model.encoder.encoder_transformer.layers." + std::to_string(layer); + TransformerLayerWeights block; + block.attention.q_weight = weights->store->load_tensor(source, prefix + ".self_attn.q_proj.weight", linear_weight_storage_type, {512, 512}); + block.attention.k_weight = weights->store->load_tensor(source, prefix + ".self_attn.k_proj.weight", linear_weight_storage_type, {512, 512}); + block.attention.v_weight = weights->store->load_tensor(source, prefix + ".self_attn.v_proj.weight", linear_weight_storage_type, {512, 512}); + block.attention.out_weight = weights->store->load_tensor(source, prefix + ".self_attn.o_proj.weight", linear_weight_storage_type, {512, 512}); + block.feed_forward.fc1_weight = weights->store->load_tensor(source, prefix + ".mlp.fc1.weight", linear_weight_storage_type, {2048, 512}); + block.feed_forward.fc2_weight = weights->store->load_tensor(source, prefix + ".mlp.fc2.weight", linear_weight_storage_type, {512, 2048}); + block.norm1 = binding::norm_from_source(*weights->store, source, prefix + ".input_layernorm", 512); + block.norm2 = binding::norm_from_source(*weights->store, source, prefix + ".post_attention_layernorm", 512); + block.scale1 = binding::layer_scale_from_named_source(*weights->store, source, prefix + ".self_attn_layer_scale.scale"); + block.scale2 = binding::layer_scale_from_named_source(*weights->store, source, prefix + ".mlp_layer_scale.scale"); + weights->transformer_layers.push_back(std::move(block)); + } + + weights->downsample = binding::conv1d_from_source( + *weights->store, + source, + "codec_model.encoder.downsample.conv", + conv_weight_storage_type, + kDownsampleConvConfig.out_channels, + kDownsampleConvConfig.in_channels, + kDownsampleConvConfig.kernel_size, + kDownsampleConvConfig.use_bias); + weights->semantic_projection = binding::conv1d_from_source( + *weights->store, + source, + "codec_model.encoder.quantizer.semantic_residual_vector_quantizer.input_proj", + conv_weight_storage_type, + kSemanticProjectionConfig.out_channels, + kSemanticProjectionConfig.in_channels, + kSemanticProjectionConfig.kernel_size, + kSemanticProjectionConfig.use_bias); + weights->acoustic_projection = binding::conv1d_from_source( + *weights->store, + source, + "codec_model.encoder.quantizer.acoustic_residual_vector_quantizer.input_proj", + conv_weight_storage_type, + kAcousticProjectionConfig.out_channels, + kAcousticProjectionConfig.in_channels, + kAcousticProjectionConfig.kernel_size, + kAcousticProjectionConfig.use_bias); + + weights->semantic_codebooks.push_back(codebook_embedding( + source, + "codec_model.encoder.quantizer.semantic_residual_vector_quantizer.layers.0.codebook.")); + + for (int layer = 0; layer < 31; ++layer) { + const std::string prefix = + "codec_model.encoder.quantizer.acoustic_residual_vector_quantizer.layers." + std::to_string(layer) + ".codebook."; + weights->acoustic_codebooks.push_back(codebook_embedding(source, prefix)); + } + + weights->store->upload(); + return weights; +} + +class BreezeSpeechEncoderGraph { +public: + BreezeSpeechEncoderGraph( + std::shared_ptr weights, + int64_t sample_capacity, + core::ExecutionContext & execution_context, + core::ConstantTensorCache & constants, + size_t graph_arena_bytes) + : weights_(std::move(weights)), + sample_capacity_(sample_capacity), + frames_((sample_capacity + kDownsampleRate - 1) / kDownsampleRate), + backend_(execution_context.backend()), + compute_threads_(std::max(1, execution_context.config().threads)) { + if (weights_ == nullptr) { + throw std::runtime_error("Breeze speech encoder graph requires weights"); + } + if (sample_capacity_ <= 0) { + throw std::runtime_error("Breeze speech encoder graph requires positive sample capacity"); + } + if (backend_ == nullptr) { + throw std::runtime_error("Breeze speech encoder backend is not initialized"); + } + + ggml_init_params params{ + /*.mem_size =*/ graph_arena_bytes, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Breeze speech encoder ggml context"); + } + + core::ModuleBuildContext build_ctx{ + ctx_.get(), + "breeze_tts.speech_encoder", + execution_context.backend_type(), + }; + auto x = core::make_tensor(build_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, sample_capacity_})); + input_ = x.tensor; + + constants.begin_graph(); + x = speech_conv(build_ctx, x, weights_->encoder_convs[0], kEncoderConvConfigs[0], kEncoderConvPadModes[0]); + for (size_t i = 0; i < weights_->residual_blocks.size(); ++i) { + x = speech_residual_block(build_ctx, x, weights_->residual_blocks[i], constants); + x = modules::EluModule{}.build(build_ctx, x); + x = speech_conv(build_ctx, x, weights_->encoder_convs[i + 1], kEncoderConvConfigs[i + 1], kEncoderConvPadModes[i + 1]); + } + x = modules::EluModule{}.build(build_ctx, x); + x = speech_conv(build_ctx, x, weights_->encoder_convs.back(), kEncoderConvConfigs.back(), kEncoderConvPadModes.back()); + + auto seq = modules::TransposeModule({{0, 2, 1, 3}, x.shape.rank}).build(build_ctx, x); + seq = core::ensure_backend_addressable_layout(build_ctx, seq); + transformer_frames_ = seq.shape.dims[1]; + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, transformer_frames_); + auto positions_value = core::wrap_tensor(positions_, core::TensorShape::from_dims({transformer_frames_}), GGML_TYPE_I32); + attention_mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, transformer_frames_, transformer_frames_, 1, 1); + const auto attention_mask = core::wrap_tensor( + attention_mask_, + core::TensorShape::from_dims({1, 1, transformer_frames_, transformer_frames_}), + GGML_TYPE_F16); + for (const auto & layer : weights_->transformer_layers) { + seq = transformer_block(build_ctx, seq, positions_value, layer, attention_mask); + } + x = modules::TransposeModule({{0, 2, 1, 3}, seq.shape.rank}).build(build_ctx, seq); + x = core::ensure_backend_addressable_layout(build_ctx, x); + x = speech_conv(build_ctx, x, weights_->downsample, kDownsampleConvConfig, modules::StreamingPadMode::Replicate); + auto semantic = speech_conv(build_ctx, x, weights_->semantic_projection, kSemanticProjectionConfig, modules::StreamingPadMode::Constant); + auto acoustic = speech_conv(build_ctx, x, weights_->acoustic_projection, kAcousticProjectionConfig, modules::StreamingPadMode::Constant); + semantic_output_ = semantic.tensor; + acoustic_output_ = acoustic.tensor; + ggml_set_output(semantic_output_); + ggml_set_output(acoustic_output_); + graph_ = ggml_new_graph_custom(ctx_.get(), 32768, false); + ggml_build_forward_expand(graph_, semantic_output_); + ggml_build_forward_expand(graph_, acoustic_output_); + constants.finish_graph(); + constants.ensure_uploaded(); + + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + if (gallocr_ == nullptr || !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + throw std::runtime_error("failed to allocate Breeze speech encoder graph"); + } + positions_data_.resize(static_cast(transformer_frames_)); + for (int64_t i = 0; i < transformer_frames_; ++i) { + positions_data_[static_cast(i)] = static_cast(i); + } + if (attention_mask_ != nullptr) { + auto mask = modules::qwen_causal_prefill_mask_values(1, transformer_frames_); + attention_mask_data_ = std::move(mask); + } + upload_static_inputs(); + } + + ~BreezeSpeechEncoderGraph() { + engine::core::release_backend_graph_resources(backend_, graph_); + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + } + + bool matches(const BreezeSpeechEncoderWeights & weights, int64_t samples, ggml_backend_t backend, int threads) const { + return weights_.get() == &weights && sample_capacity_ == samples && backend_ == backend && + compute_threads_ == std::max(1, threads); + } + + BreezeSpeechEncoderOutput run(const std::vector & waveform) { + if (static_cast(waveform.size()) > sample_capacity_) { + throw std::runtime_error("Breeze speech encoder waveform exceeds graph capacity"); + } + upload_static_inputs(); + std::vector padded(static_cast(sample_capacity_), 0.0F); + std::copy(waveform.begin(), waveform.end(), padded.begin()); + ggml_backend_tensor_set(input_, padded.data(), 0, padded.size() * sizeof(float)); + core::set_backend_threads(backend_, compute_threads_); + const ggml_status status = engine::core::compute_backend_graph(backend_, graph_); + ggml_backend_synchronize(backend_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Breeze speech encoder graph compute failed"); + } + BreezeSpeechEncoderOutput out; + out.semantic_projected.resize(static_cast(kQuantizerDim * frames_)); + out.acoustic_projected.resize(static_cast(kQuantizerDim * frames_)); + ggml_backend_tensor_get(semantic_output_, out.semantic_projected.data(), 0, out.semantic_projected.size() * sizeof(float)); + ggml_backend_tensor_get(acoustic_output_, out.acoustic_projected.data(), 0, out.acoustic_projected.size() * sizeof(float)); + return out; + } + + int64_t frames() const noexcept { + return frames_; + } + +private: + void upload_static_inputs() { + ggml_backend_tensor_set( + positions_, + positions_data_.data(), + 0, + positions_data_.size() * sizeof(int32_t)); + if (attention_mask_ != nullptr) { + ggml_backend_tensor_set( + attention_mask_, + attention_mask_data_.data(), + 0, + attention_mask_data_.size() * sizeof(ggml_fp16_t)); + } + } + + std::shared_ptr weights_; + int64_t sample_capacity_ = 0; + int64_t frames_ = 0; + int64_t transformer_frames_ = 0; + std::unique_ptr ctx_; + ggml_tensor * input_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * attention_mask_ = nullptr; + ggml_tensor * semantic_output_ = nullptr; + ggml_tensor * acoustic_output_ = nullptr; + std::vector positions_data_; + std::vector attention_mask_data_; + ggml_cgraph * graph_ = nullptr; + ggml_backend_t backend_ = nullptr; + int compute_threads_ = 1; + ggml_gallocr_t gallocr_ = nullptr; +}; + +BreezeSpeechEncoderRuntime::BreezeSpeechEncoderRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution_context, + size_t graph_arena_bytes, + assets::TensorStorageType linear_weight_storage_type, + assets::TensorStorageType conv_weight_storage_type) + : assets_(std::move(assets)), + execution_context_(&execution_context), + graph_arena_bytes_(graph_arena_bytes) { + if (assets_ == nullptr) { + throw std::runtime_error("Breeze speech encoder requires assets"); + } + weights_ = load_weights( + *assets_, + execution_context_->backend(), + execution_context_->backend_type(), + linear_weight_storage_type, + conv_weight_storage_type); + constants_ = std::make_unique( + execution_context_->backend(), + std::max(1, execution_context_->config().threads), + "breeze_tts.speech_encoder.constants", + 768ull * 1024ull * 1024ull); +} + +BreezeSpeechEncoderRuntime::~BreezeSpeechEncoderRuntime() = default; + +BreezeSpeechCodes BreezeSpeechEncoderRuntime::encode(const runtime::AudioBuffer & audio) const { + const auto start = Clock::now(); + if (execution_context_ == nullptr) { + throw std::runtime_error("Breeze speech encoder execution context is missing"); + } + if (audio.sample_rate <= 0 || audio.channels <= 0 || audio.samples.empty()) { + throw std::runtime_error("Breeze speech encoder requires non-empty reference audio"); + } + const auto waveform = engine::audio::convert_interleaved_audio_to_mono_linear_resampled( + audio.samples, + audio.sample_rate, + audio.channels, + static_cast(kSampleRate)); + const int64_t valid_samples = static_cast(waveform.size()); + const int64_t frames = std::max(1, (valid_samples + kDownsampleRate - 1) / kDownsampleRate); + const int64_t sample_capacity = valid_samples; + const int threads = std::max(1, execution_context_->config().threads); + if (graph_ == nullptr || !graph_->matches(*weights_, sample_capacity, execution_context_->backend(), threads)) { + graph_.reset(); + graph_ = std::make_unique( + weights_, + sample_capacity, + *execution_context_, + *constants_, + graph_arena_bytes_); + } + auto out = graph_->run(waveform); + out.codes.frames = frames; + out.codes.code_groups = kValidQuantizers; + out.codes.codes = quantize_projected(out.semantic_projected, out.acoustic_projected, frames, *weights_); + debug::timing_log_scalar("breeze_tts.speech_encoder.total_ms", engine::debug::elapsed_ms(start)); + return out.codes; +} + +void BreezeSpeechEncoderRuntime::release_runtime_graphs() const { + graph_.reset(); +} + +} // namespace engine::models::breeze_tts diff --git a/src/models/breeze_tts/text_encoder.cpp b/src/models/breeze_tts/text_encoder.cpp new file mode 100644 index 000000000..7fafa9cf3 --- /dev/null +++ b/src/models/breeze_tts/text_encoder.cpp @@ -0,0 +1,282 @@ +#include "engine/models/breeze_tts/text_encoder.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/text_encoders/t5_gemma_encoder.h" +#include "engine/framework/modules/weight_binding.h" + +#include +#include +#include +#include + +namespace engine::models::breeze_tts { +namespace { + +namespace binding = engine::modules::binding; +namespace core = engine::core; +namespace modules = engine::modules; +using Clock = std::chrono::steady_clock; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +modules::T5GemmaEncoderConfig text_config(const BreezeTTSConfig & config) { + modules::T5GemmaEncoderConfig out; + out.hidden_size = config.text_hidden_size; + out.layers = config.text_layers; + out.attention_heads = config.text_heads; + out.kv_heads = config.text_kv_heads; + out.head_dim = config.text_head_dim; + out.attention_size = config.text_heads * config.text_head_dim; + out.intermediate_size = config.text_intermediate_size; + out.vocab_size = config.text_vocab_size; + out.rope_theta = config.text_rope_theta; + out.rope_freq_scale = 1.0F / config.text_rope_linear_factor; + out.rms_norm_eps = config.text_rms_norm_eps; + out.query_pre_attn_scalar = config.text_query_pre_attn_scalar; + out.attn_logit_softcap = 0.0F; + out.scale_embeddings = true; + out.use_qk_norm = true; + out.rms_norm_style = modules::T5GemmaRMSNormStyle::Gemma; + out.layer_rope_theta = config.text_layer_rope_theta; + out.layer_rope_freq_scale = config.text_layer_rope_freq_scale; + return out; +} + +modules::T5GemmaEncoderLayerWeights load_text_layer( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const BreezeTTSConfig & config, + assets::TensorStorageType storage_type, + int64_t layer) { + const std::string prefix = "text_encoder.layers." + std::to_string(layer); + modules::T5GemmaEncoderLayerWeights out; + out.pre_self_attn_norm = store.load_f32_tensor(source, prefix + ".pre_self_attn_layernorm.weight", {config.text_hidden_size}); + out.post_self_attn_norm = store.load_f32_tensor(source, prefix + ".post_self_attn_layernorm.weight", {config.text_hidden_size}); + out.pre_ff_norm = store.load_f32_tensor(source, prefix + ".pre_feedforward_layernorm.weight", {config.text_hidden_size}); + out.post_ff_norm = store.load_f32_tensor(source, prefix + ".post_feedforward_layernorm.weight", {config.text_hidden_size}); + out.q_proj = binding::linear_from_source( + store, source, prefix + ".self_attn.q_proj", storage_type, config.text_heads * config.text_head_dim, config.text_hidden_size, false); + out.k_proj = binding::linear_from_source( + store, source, prefix + ".self_attn.k_proj", storage_type, config.text_kv_heads * config.text_head_dim, config.text_hidden_size, false); + out.v_proj = binding::linear_from_source( + store, source, prefix + ".self_attn.v_proj", storage_type, config.text_kv_heads * config.text_head_dim, config.text_hidden_size, false); + out.o_proj = binding::linear_from_source( + store, source, prefix + ".self_attn.o_proj", storage_type, config.text_hidden_size, config.text_heads * config.text_head_dim, false); + out.q_norm = binding::norm_weight_from_source(store, source, prefix + ".self_attn.q_norm", config.text_head_dim); + out.k_norm = binding::norm_weight_from_source(store, source, prefix + ".self_attn.k_norm", config.text_head_dim); + out.gate_proj = binding::linear_from_source( + store, source, prefix + ".mlp.gate_proj", storage_type, config.text_intermediate_size, config.text_hidden_size, false); + out.up_proj = binding::linear_from_source( + store, source, prefix + ".mlp.up_proj", storage_type, config.text_intermediate_size, config.text_hidden_size, false); + out.down_proj = binding::linear_from_source( + store, source, prefix + ".mlp.down_proj", storage_type, config.text_hidden_size, config.text_intermediate_size, false); + return out; +} + +struct BreezeTextWeights { + std::shared_ptr store; + modules::T5GemmaEncoderWeights encoder; + modules::LinearWeights projector; +}; + +std::shared_ptr load_text_weights( + const BreezeTTSAssets & assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) { + auto out = std::make_shared(); + out->store = std::make_shared( + execution.backend(), + execution.backend_type(), + "breeze_tts.text_encoder.weights", + weight_context_bytes); + const auto & source = *assets.weights; + const auto & config = assets.config; + out->encoder.embed_tokens = out->store->load_tensor( + source, + "text_encoder.embed_tokens.weight", + storage_type, + {config.text_vocab_size, config.text_hidden_size}); + out->encoder.layers.reserve(static_cast(config.text_layers)); + for (int64_t layer = 0; layer < config.text_layers; ++layer) { + out->encoder.layers.push_back(load_text_layer(*out->store, source, config, storage_type, layer)); + } + out->encoder.norm = out->store->load_f32_tensor(source, "text_encoder.norm.weight", {config.text_hidden_size}); + out->projector = binding::linear_from_source( + *out->store, + source, + "text_encoder_proj", + storage_type, + config.hidden_size, + config.text_hidden_size, + false); + out->store->upload(); + return out; +} + +std::vector full_attention_mask(int64_t heads, int64_t tokens) { + return std::vector(static_cast(heads * tokens * tokens), 0.0F); +} + +} // namespace + +struct BreezeTextEncoderRuntime::Impl { + struct Graph { + Graph( + core::ExecutionContext & execution, + size_t graph_arena_bytes, + const BreezeTTSConfig & config, + std::shared_ptr weights, + int64_t tokens) + : execution(execution), + config(config), + weights(std::move(weights)), + tokens(tokens) { + ctx.reset(ggml_init({graph_arena_bytes, nullptr, true})); + if (ctx == nullptr) { + throw std::runtime_error("failed to initialize BreezeTTS text encoder graph context"); + } + core::ModuleBuildContext build{ctx.get(), "breeze_tts.text_encoder", execution.backend_type()}; + input_ids_value = core::make_tensor(build, GGML_TYPE_I32, core::TensorShape::from_dims({1, tokens})); + positions_value = core::make_tensor(build, GGML_TYPE_I32, core::TensorShape::from_dims({tokens})); + attention_value = core::make_tensor(build, GGML_TYPE_F32, core::TensorShape::from_dims({1, config.text_heads, tokens, tokens})); + input_ids = input_ids_value.tensor; + positions = positions_value.tensor; + attention = attention_value.tensor; + auto encoded = modules::T5GemmaEncoderModule(text_config(config)).build( + build, + input_ids_value, + positions_value, + attention_value, + this->weights->encoder); + auto projected = modules::LinearModule({config.text_hidden_size, config.hidden_size, false, GGML_PREC_DEFAULT}).build( + build, + encoded, + this->weights->projector); + output = core::ensure_backend_addressable_layout(build, projected).tensor; + ggml_set_input(input_ids); + ggml_set_input(positions); + ggml_set_input(attention); + ggml_set_output(output); + graph = ggml_new_graph_custom(ctx.get(), 65536, false); + ggml_build_forward_expand(graph, output); + if (core::is_host_backend(execution.backend())) { + params_buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), execution.backend()); + } + galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution.backend())); + if (galloc == nullptr || !ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { + throw std::runtime_error("failed to allocate BreezeTTS text encoder graph"); + } + } + + ~Graph() { + engine::core::release_backend_graph_resources(execution.backend(), graph); + if (galloc != nullptr) { + ggml_gallocr_free(galloc); + } + if (params_buffer != nullptr) { + ggml_backend_buffer_free(params_buffer); + } + } + + BreezeProjectedText run(const std::vector & ids) { + if (static_cast(ids.size()) != tokens) { + throw std::runtime_error("BreezeTTS text encoder graph token count mismatch"); + } + std::vector pos(static_cast(tokens)); + for (int64_t i = 0; i < tokens; ++i) { + pos[static_cast(i)] = static_cast(i); + } + const auto mask = full_attention_mask(config.text_heads, tokens); + core::write_tensor_i32(input_ids_value, ids); + core::write_tensor_i32(positions_value, pos); + core::write_tensor_f32(attention_value, mask); + core::set_backend_threads(execution.backend(), execution.config().threads); + if (core::compute_backend_graph(execution.backend(), graph, nullptr, "breeze_tts.text_encoder") != GGML_STATUS_SUCCESS) { + throw std::runtime_error("BreezeTTS text encoder graph compute failed"); + } + return {tokens, core::read_tensor_f32(output)}; + } + + core::ExecutionContext & execution; + BreezeTTSConfig config; + std::shared_ptr weights; + int64_t tokens = 0; + std::unique_ptr ctx; + core::TensorValue input_ids_value; + core::TensorValue positions_value; + core::TensorValue attention_value; + ggml_tensor * input_ids = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * attention = nullptr; + ggml_tensor * output = nullptr; + ggml_cgraph * graph = nullptr; + ggml_gallocr_t galloc = nullptr; + ggml_backend_buffer_t params_buffer = nullptr; + }; + + Impl( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) + : assets(std::move(assets)), + execution(execution), + graph_arena_bytes(graph_arena_bytes), + weights(load_text_weights(*this->assets, execution, weight_context_bytes, storage_type)) {} + + BreezeProjectedText encode(const std::vector & input_ids) { + const auto start = Clock::now(); + const int64_t tokens = static_cast(input_ids.size()); + if (tokens <= 0) { + throw std::runtime_error("BreezeTTS text encoder requires non-empty input"); + } + if (graph == nullptr || graph->tokens != tokens) { + graph = std::make_unique(execution, graph_arena_bytes, assets->config, weights, tokens); + } + auto out = graph->run(input_ids); + engine::debug::timing_log_scalar("breeze_tts.text_encoder.total_ms", engine::debug::elapsed_ms(start)); + return out; + } + + void release_runtime_graphs() { + graph.reset(); + } + + std::shared_ptr assets; + core::ExecutionContext & execution; + size_t graph_arena_bytes = 0; + std::shared_ptr weights; + std::unique_ptr graph; +}; + +BreezeTextEncoderRuntime::BreezeTextEncoderRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type) + : impl_(std::make_unique(std::move(assets), execution, graph_arena_bytes, weight_context_bytes, storage_type)) {} + +BreezeTextEncoderRuntime::~BreezeTextEncoderRuntime() = default; + +BreezeProjectedText BreezeTextEncoderRuntime::encode(const std::vector & input_ids) { + return impl_->encode(input_ids); +} + +void BreezeTextEncoderRuntime::release_runtime_graphs() { + impl_->release_runtime_graphs(); +} + +} // namespace engine::models::breeze_tts diff --git a/src/models/breeze_tts/tokenizer_text.cpp b/src/models/breeze_tts/tokenizer_text.cpp new file mode 100644 index 000000000..ac323cbda --- /dev/null +++ b/src/models/breeze_tts/tokenizer_text.cpp @@ -0,0 +1,136 @@ +#include "engine/models/breeze_tts/tokenizer_text.h" + +#include "engine/framework/tokenizers/llama_bpe.h" + +#include +#include + +namespace engine::models::breeze_tts { +namespace { + +constexpr const char * kDefaultInstruction = "Speak clearly and naturally."; +constexpr const char * kBos = ""; +constexpr const char * kSpeaker0 = "[S0]"; +constexpr const char * kInstructionBos = ""; +constexpr const char * kInstructionEos = ""; + +std::vector breeze_special_tokens(const BreezeTTSConfig & config) { + return { + {kSpeaker0, 262146}, + {"[S1]", 262147}, + {"[S2]", 262148}, + {"[S3]", 262149}, + {"[S4]", 262150}, + {"[S5]", 262151}, + {"[S6]", 262152}, + {"[S7]", 262153}, + {"[S8]", 262154}, + {"[S9]", 262155}, + {kInstructionBos, 262156}, + {kInstructionEos, 262157}, + {"<|AUDIO|>", static_cast(config.audio_token_id)}, + {"<|audio_eos|>", static_cast(config.audio_eos_token_id)}, + }; +} + +void append_tokens(BreezePromptBranch & out, const std::vector & ids, bool text) { + out.input_ids.insert(out.input_ids.end(), ids.begin(), ids.end()); + out.text_mask.insert(out.text_mask.end(), ids.size(), text ? uint8_t{1} : uint8_t{0}); + if (text) { + out.text_segment_lengths.push_back(static_cast(ids.size())); + out.text_segments.push_back(ids); + } +} + +void append_audio_placeholders(BreezePromptBranch & out, int32_t audio_token, int32_t audio_eos, int64_t frames) { + if (frames <= 0) { + throw std::runtime_error("BreezeTTS clone requires encoded reference audio frames"); + } + out.input_ids.insert(out.input_ids.end(), static_cast(frames), audio_token); + out.text_mask.insert(out.text_mask.end(), static_cast(frames), uint8_t{0}); + out.input_ids.push_back(audio_eos); + out.text_mask.push_back(uint8_t{0}); +} + +} // namespace + +struct BreezeTextTokenizer::Impl { + Impl(std::shared_ptr assets) + : assets(std::move(assets)), + tokenizer(engine::tokenizers::LlamaBpeTokenizerSpec{ + {}, + {}, + this->assets->resources.require_file("tokenizer_config_json"), + this->assets->resources.require_file("tokenizer_json"), + engine::tokenizers::LlamaBpePreTokenizer::Gemma4, + breeze_special_tokens(this->assets->config), + "\xE2\x96\x81"}) { + audio_token = static_cast(this->assets->config.audio_token_id); + audio_eos = static_cast(this->assets->config.audio_eos_token_id); + } + + std::vector encode_text_segment(const std::string & text) const { + return tokenizer.encode(std::string(kBos) + text, true); + } + + std::shared_ptr assets; + engine::tokenizers::LlamaBpeTokenizer tokenizer; + int32_t audio_token = 0; + int32_t audio_eos = 0; +}; + +BreezeTextTokenizer::BreezeTextTokenizer(std::shared_ptr assets) + : impl_(std::make_unique(std::move(assets))) {} + +BreezeTextTokenizer::~BreezeTextTokenizer() = default; + +BreezePromptBranch BreezeTextTokenizer::build_tts_instruction( + const std::string & text, + const std::string & instruction) const { + const std::string actual_instruction = instruction.empty() ? kDefaultInstruction : instruction; + BreezePromptBranch out; + append_tokens(out, impl_->encode_text_segment( + std::string(kSpeaker0) + kInstructionBos + actual_instruction + kInstructionEos + text), true); + return out; +} + +BreezePromptBranch BreezeTextTokenizer::build_tts_plain(const std::string & text) const { + BreezePromptBranch out; + append_tokens(out, impl_->encode_text_segment(std::string(kSpeaker0) + text), true); + return out; +} + +BreezePromptBranch BreezeTextTokenizer::build_clone( + const std::string & text, + const std::string & instruction, + const std::string & reference_text, + int64_t reference_audio_frames) const { + const std::string actual_instruction = instruction.empty() ? kDefaultInstruction : instruction; + BreezePromptBranch out; + append_tokens(out, impl_->encode_text_segment(std::string(kSpeaker0) + reference_text), true); + append_audio_placeholders(out, impl_->audio_token, impl_->audio_eos, reference_audio_frames); + append_tokens(out, impl_->encode_text_segment( + std::string(kSpeaker0) + kInstructionBos + actual_instruction + kInstructionEos + text), true); + return out; +} + +BreezePromptBranch BreezeTextTokenizer::build_clone_negative( + const std::string & text, + const std::string & reference_text, + int64_t reference_audio_frames) const { + BreezePromptBranch out; + append_tokens(out, impl_->encode_text_segment(std::string(kSpeaker0) + reference_text), true); + append_audio_placeholders(out, impl_->audio_token, impl_->audio_eos, reference_audio_frames); + append_tokens(out, impl_->encode_text_segment(std::string(kSpeaker0) + text), true); + return out; +} + +int32_t BreezeTextTokenizer::audio_token_id() const noexcept { + return impl_->audio_token; +} + +int32_t BreezeTextTokenizer::audio_eos_token_id() const noexcept { + return impl_->audio_eos; +} + +} // namespace engine::models::breeze_tts diff --git a/src/models/cosyvoice3/ar.cpp b/src/models/cosyvoice3/ar.cpp new file mode 100644 index 000000000..80b3e55fb --- /dev/null +++ b/src/models/cosyvoice3/ar.cpp @@ -0,0 +1,572 @@ +#include "engine/models/cosyvoice3/ar.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/transformers/qwen_causal_decode_runtime.h" +#include "engine/framework/modules/weight_binding.h" +#include "engine/framework/sampling/hf_sampler.h" +#include "engine/framework/sampling/torch_random.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::cosyvoice3 { +namespace { + +namespace binding = engine::modules::binding; +namespace core = engine::core; +namespace modules = engine::modules; + +using Clock = std::chrono::steady_clock; + +constexpr float kTopP = 0.8F; +constexpr int64_t kRasWindow = 10; +constexpr float kRasTau = 0.1F; +constexpr std::array kSilentTokens{1, 2, 28, 29, 55, 248, 494, 2241, 2242, 2322, 2323}; +constexpr int64_t kMaxConsecutiveSilentTokens = 5; + +modules::QwenCausalDecodeRuntimeConfig make_qwen_config( + const CosyVoice3Config & config, + core::BackendType backend_type, + size_t graph_arena_bytes) { + modules::QwenCausalDecodeRuntimeConfig out; + out.trace_name = "cosyvoice3.ar"; + out.prefill_graph_arena_bytes = graph_arena_bytes; + out.decode_graph_arena_bytes = graph_arena_bytes; + out.decoder.stack.hidden_size = config.hidden_size; + out.decoder.stack.num_attention_heads = config.heads; + out.decoder.stack.num_key_value_heads = config.kv_heads; + out.decoder.stack.head_dim = config.head_dim; + out.decoder.stack.intermediate_size = config.intermediate_size; + out.decoder.stack.layers = config.layers; + out.decoder.stack.rms_norm_eps = 1.0e-6F; + out.decoder.stack.rope_theta = 1000000.0F; + out.decoder.stack.rope_type = GGML_ROPE_TYPE_NEOX; + out.decoder.stack.use_qk_norm = false; + out.decoder.stack.attention_precision = GGML_PREC_F32; + out.decoder.stack.projection_precision = GGML_PREC_DEFAULT; + out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.decoder.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.decoder.stack.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; + out.decoder.logits_size = config.speech_token_size + config.speech_reserved_tokens; + out.decoder.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; + out.output_mode = modules::QwenCausalDecodeOutputMode::Logits; + out.logits_readback_token_ids.reserve(static_cast(out.decoder.logits_size)); + for (int32_t token = 0; token < static_cast(out.decoder.logits_size); ++token) { + out.logits_readback_token_ids.push_back(token); + } + if (backend_type == core::BackendType::Metal) { + out.decoder.lm_head_input_type = GGML_TYPE_F32; + } else if (backend_type == core::BackendType::Vulkan) { + out.decoder.lm_head_input_type = GGML_TYPE_F16; + } + return out; +} + +modules::QwenDecoderLayerWeights load_qwen_layer( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const CosyVoice3Config & config, + assets::TensorStorageType storage_type, + int64_t layer) { + const std::string prefix = "llm.model.model.layers." + std::to_string(layer); + modules::QwenDecoderLayerWeights out; + out.input_norm = binding::norm_weight_from_source(store, source, prefix + ".input_layernorm", config.hidden_size); + out.self_attention.q_weight = store.load_tensor( + source, + prefix + ".self_attn.q_proj.weight", + storage_type, + {config.heads * config.head_dim, config.hidden_size}); + out.self_attention.q_bias = store.load_f32_tensor( + source, + prefix + ".self_attn.q_proj.bias", + {config.heads * config.head_dim}); + out.self_attention.k_weight = store.load_tensor( + source, + prefix + ".self_attn.k_proj.weight", + storage_type, + {config.kv_heads * config.head_dim, config.hidden_size}); + out.self_attention.k_bias = store.load_f32_tensor( + source, + prefix + ".self_attn.k_proj.bias", + {config.kv_heads * config.head_dim}); + out.self_attention.v_weight = store.load_tensor( + source, + prefix + ".self_attn.v_proj.weight", + storage_type, + {config.kv_heads * config.head_dim, config.hidden_size}); + out.self_attention.v_bias = store.load_f32_tensor( + source, + prefix + ".self_attn.v_proj.bias", + {config.kv_heads * config.head_dim}); + out.self_attention.out_weight = store.load_tensor( + source, + prefix + ".self_attn.o_proj.weight", + storage_type, + {config.hidden_size, config.heads * config.head_dim}); + out.post_norm = binding::norm_weight_from_source(store, source, prefix + ".post_attention_layernorm", config.hidden_size); + out.mlp.gate_proj = binding::linear_from_source( + store, + source, + prefix + ".mlp.gate_proj", + storage_type, + config.intermediate_size, + config.hidden_size, + false); + out.mlp.up_proj = binding::linear_from_source( + store, + source, + prefix + ".mlp.up_proj", + storage_type, + config.intermediate_size, + config.hidden_size, + false); + out.mlp.down_proj = binding::linear_from_source( + store, + source, + prefix + ".mlp.down_proj", + storage_type, + config.hidden_size, + config.intermediate_size, + false); + return out; +} + +struct CosyVoice3ArWeights { + std::shared_ptr store; + modules::QwenCausalDecodeRuntimeWeights qwen; + std::vector text_embedding; + std::vector speech_embedding; +}; + +std::shared_ptr load_ar_weights( + const CosyVoice3Assets & assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + assets::TensorStorageType storage_type, + const modules::QwenCausalDecodeRuntimeConfig & qwen_config) { + auto weights = std::make_shared(); + weights->store = std::make_shared( + execution.backend(), + execution.backend_type(), + "cosyvoice3.ar.weights", + weight_context_bytes); + const auto & source = *assets.llm_weights; + const auto & config = assets.config; + weights->qwen.token_embedding = weights->store->load_tensor( + source, + "llm.model.model.embed_tokens.weight", + storage_type, + {config.text_vocab_size, config.hidden_size}); + weights->qwen.stack.layers.reserve(static_cast(config.layers)); + for (int64_t layer = 0; layer < config.layers; ++layer) { + weights->qwen.stack.layers.push_back(load_qwen_layer(*weights->store, source, config, storage_type, layer)); + } + weights->qwen.final_norm = binding::norm_weight_from_source( + *weights->store, + source, + "llm.model.model.norm", + config.hidden_size); + weights->qwen.lm_head = binding::linear_from_source( + *weights->store, + source, + "llm_decoder", + storage_type, + qwen_config.decoder.logits_size, + config.hidden_size, + false); + weights->text_embedding = source.require_f32( + "llm.model.model.embed_tokens.weight", + {config.text_vocab_size, config.hidden_size}); + weights->speech_embedding = source.require_f32( + "speech_embedding.weight", + {config.speech_token_size + config.speech_reserved_tokens, config.hidden_size}); + weights->store->upload(); + assets.llm_weights->release_storage(); + return weights; +} + +void append_embedding_rows( + std::vector & out, + const std::vector & table, + int64_t rows, + int64_t dim, + const std::vector & tokens, + const char * label) { + for (const int32_t token : tokens) { + if (token < 0 || token >= rows) { + throw std::runtime_error(std::string("CosyVoice3 AR ") + label + " token is outside embedding table"); + } + const size_t begin = static_cast(token) * static_cast(dim); + out.insert(out.end(), table.begin() + static_cast(begin), table.begin() + static_cast(begin + static_cast(dim))); + } +} + +std::vector embedding_row( + const std::vector & table, + int64_t rows, + int64_t dim, + int32_t token, + const char * label) { + if (token < 0 || token >= rows) { + throw std::runtime_error(std::string("CosyVoice3 AR ") + label + " token is outside embedding table"); + } + const size_t begin = static_cast(token) * static_cast(dim); + return std::vector( + table.begin() + static_cast(begin), + table.begin() + static_cast(begin + static_cast(dim))); +} + +std::vector log_softmax(const std::vector & logits) { + float max_value = -std::numeric_limits::infinity(); + for (const float value : logits) { + if (std::isfinite(value)) { + max_value = std::max(max_value, value); + } + } + if (!std::isfinite(max_value)) { + throw std::runtime_error("CosyVoice3 AR logits have no finite value"); + } + double sum = 0.0; + for (const float value : logits) { + if (std::isfinite(value)) { + sum += std::exp(static_cast(value - max_value)); + } + } + if (!(sum > 0.0) || !std::isfinite(sum)) { + throw std::runtime_error("CosyVoice3 AR logits have invalid probability mass"); + } + const float log_sum = max_value + static_cast(std::log(sum)); + std::vector out(logits.size(), -std::numeric_limits::infinity()); + for (size_t index = 0; index < logits.size(); ++index) { + if (std::isfinite(logits[index])) { + out[index] = logits[index] - log_sum; + } + } + return out; +} + +int32_t sample_from_scores( + const std::vector & scores, + sampling::HfSamplerScratch & scratch, + std::mt19937 & fallback_rng, + const sampling::TorchCudaSamplingPolicy & policy, + uint64_t seed, + uint64_t & sample_call_index, + uint64_t & rng_offset_blocks, + std::string_view context) { + const sampling::HfTorchSamplingState torch_state{ + &policy, + seed, + sample_call_index, + rng_offset_blocks, + true, + }; + const int32_t token = sampling::HfTokenSampler::sample_from_processed_scores( + scores, + scratch, + fallback_rng, + policy.cuda_fast_path ? &torch_state : nullptr, + context); + ++sample_call_index; + if (policy.cuda_fast_path) { + rng_offset_blocks += sampling::torch_cuda_tensor_iterator_offset_blocks( + static_cast(scores.size()), + policy); + } + return token; +} + +int32_t nucleus_sample( + const std::vector & log_probs, + int64_t top_k, + sampling::HfSamplerScratch & scratch, + std::mt19937 & fallback_rng, + const sampling::TorchCudaSamplingPolicy & policy, + uint64_t seed, + uint64_t & sample_call_index, + uint64_t & rng_offset_blocks) { + std::vector order(log_probs.size()); + std::iota(order.begin(), order.end(), 0); + std::stable_sort(order.begin(), order.end(), [&](int32_t lhs, int32_t rhs) { + return log_probs[static_cast(lhs)] > log_probs[static_cast(rhs)]; + }); + + std::vector kept; + std::vector kept_scores; + kept.reserve(static_cast(std::min(top_k, static_cast(order.size())))); + kept_scores.reserve(kept.capacity()); + double cumulative = 0.0; + for (const int32_t token : order) { + const float score = log_probs[static_cast(token)]; + if (!std::isfinite(score)) { + continue; + } + if (cumulative < static_cast(kTopP) && static_cast(kept.size()) < top_k) { + cumulative += std::exp(static_cast(score)); + kept.push_back(token); + kept_scores.push_back(score); + } else { + break; + } + } + if (kept.empty()) { + throw std::runtime_error("CosyVoice3 AR nucleus sampler has no finite candidates"); + } + const int32_t local = sample_from_scores( + kept_scores, + scratch, + fallback_rng, + policy, + seed, + sample_call_index, + rng_offset_blocks, + "CosyVoice3 AR nucleus sampler"); + return kept[static_cast(local)]; +} + +int32_t ras_sample( + std::vector log_probs, + const std::vector & decoded_tokens, + int64_t top_k, + sampling::HfSamplerScratch & scratch, + std::mt19937 & fallback_rng, + const sampling::TorchCudaSamplingPolicy & policy, + uint64_t seed, + uint64_t & sample_call_index, + uint64_t & rng_offset_blocks) { + const int32_t top_id = nucleus_sample( + log_probs, + top_k, + scratch, + fallback_rng, + policy, + seed, + sample_call_index, + rng_offset_blocks); + int64_t repeat_count = 0; + const size_t window = std::min(decoded_tokens.size(), static_cast(kRasWindow)); + for (size_t index = decoded_tokens.size() - window; index < decoded_tokens.size(); ++index) { + if (decoded_tokens[index] == top_id) { + ++repeat_count; + } + } + if (static_cast(repeat_count) >= static_cast(kRasWindow) * kRasTau) { + log_probs[static_cast(top_id)] = -std::numeric_limits::infinity(); + return sample_from_scores( + log_probs, + scratch, + fallback_rng, + policy, + seed, + sample_call_index, + rng_offset_blocks, + "CosyVoice3 AR RAS fallback sampler"); + } + return top_id; +} + +bool is_silent_token(int32_t token) { + return std::find(kSilentTokens.begin(), kSilentTokens.end(), token) != kSilentTokens.end(); +} + +} // namespace + +class CosyVoice3ArRuntime::Impl { +public: + Impl( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type) + : assets_(std::move(assets)), + execution_(execution), + sampling_policy_(sampling::resolve_torch_cuda_sampling_policy( + execution.backend_type(), + execution.config().device, + "cosyvoice3.ar.sampling", + "CosyVoice3 AR", + sampling::TorchCudaSamplingPolicyFailureMode::FallbackToDefault)) { + if (assets_ == nullptr) { + throw std::runtime_error("CosyVoice3 AR runtime requires assets"); + } + qwen_config_ = make_qwen_config(assets_->config, execution_.backend_type(), graph_arena_bytes); + weights_ = load_ar_weights(*assets_, execution_, weight_context_bytes, storage_type, qwen_config_); + qwen_runtime_ = std::make_unique( + execution_, + qwen_config_, + weights_->qwen); + } + + CosyVoice3ArOutput generate(const CosyVoice3ArRequest & request) { + const auto & config = assets_->config; + if (request.target_text_tokens.empty()) { + throw std::runtime_error("CosyVoice3 AR target text tokens are empty"); + } + if (request.top_k <= 0) { + throw std::runtime_error("CosyVoice3 AR top_k must be positive"); + } + const int64_t min_len = request.min_tokens >= 0 + ? request.min_tokens + : static_cast(request.target_text_tokens.size()) * 2; + const int64_t max_len = request.max_tokens >= 0 + ? request.max_tokens + : static_cast(request.target_text_tokens.size()) * 20; + if (min_len < 0 || max_len <= 0 || min_len > max_len) { + throw std::runtime_error("CosyVoice3 AR token length bounds are invalid"); + } + + const int32_t sos = static_cast(config.speech_token_size); + const int32_t task_id = static_cast(config.speech_token_size + 2); + const int64_t speech_embedding_rows = config.speech_token_size + config.speech_reserved_tokens; + + std::vector text_tokens = request.prompt_text_tokens; + text_tokens.insert(text_tokens.end(), request.target_text_tokens.begin(), request.target_text_tokens.end()); + + const int64_t prefill_steps = + 1 + static_cast(text_tokens.size()) + 1 + static_cast(request.prompt_speech_tokens.size()); + std::vector embeddings; + embeddings.reserve(static_cast(prefill_steps * config.hidden_size)); + append_embedding_rows( + embeddings, + weights_->speech_embedding, + speech_embedding_rows, + config.hidden_size, + std::vector{sos}, + "sos"); + append_embedding_rows( + embeddings, + weights_->text_embedding, + config.text_vocab_size, + config.hidden_size, + text_tokens, + "text"); + append_embedding_rows( + embeddings, + weights_->speech_embedding, + speech_embedding_rows, + config.hidden_size, + std::vector{task_id}, + "task"); + append_embedding_rows( + embeddings, + weights_->speech_embedding, + speech_embedding_rows, + config.hidden_size, + request.prompt_speech_tokens, + "prompt speech"); + + const int64_t required_cache_steps = prefill_steps + max_len; + auto timing_start = Clock::now(); + qwen_runtime_->release_runtime_graphs(); + auto prefill = qwen_runtime_->prefill_embeddings(embeddings, prefill_steps); + debug::timing_log_scalar("cosyvoice3.ar.prefill.total_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + qwen_runtime_->start_decode_embeddings(prefill.state, required_cache_steps); + + sampling::HfSamplerScratch scratch; + scratch.reserve_vocab(static_cast(config.speech_token_size + config.speech_reserved_tokens)); + std::mt19937 fallback_rng(request.seed); + uint64_t sample_call_index = 0; + uint64_t rng_offset_blocks = 0; + + CosyVoice3ArOutput out; + std::vector decoded_tokens; + int64_t consecutive_silent_tokens = 0; + int64_t filtered_silent_tokens = 0; + std::vector logits = std::move(prefill.logits); + timing_start = Clock::now(); + for (int64_t step = 0; step < max_len; ++step) { + auto log_probs = log_softmax(logits); + if (step < min_len) { + log_probs[static_cast(sos)] = -std::numeric_limits::infinity(); + } + const int32_t token = ras_sample( + std::move(log_probs), + decoded_tokens, + request.top_k, + scratch, + fallback_rng, + sampling_policy_, + request.seed, + sample_call_index, + rng_offset_blocks); + if (token >= config.speech_token_size) { + break; + } + decoded_tokens.push_back(token); + bool append_token = true; + if (is_silent_token(token)) { + ++consecutive_silent_tokens; + if (consecutive_silent_tokens > kMaxConsecutiveSilentTokens) { + append_token = false; + ++filtered_silent_tokens; + } + } else { + consecutive_silent_tokens = 0; + } + if (append_token) { + out.speech_tokens.push_back(token); + } + logits = qwen_runtime_->decode_embedding(embedding_row( + weights_->speech_embedding, + speech_embedding_rows, + config.hidden_size, + token, + "sampled speech")).logits; + } + debug::timing_log_scalar("cosyvoice3.ar.decode.total_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + (void) prefill_steps; + (void) filtered_silent_tokens; + return out; + } + + void release_graphs() { + if (qwen_runtime_ != nullptr) { + qwen_runtime_->release_runtime_graphs(); + } + } + +private: + std::shared_ptr assets_; + core::ExecutionContext & execution_; + modules::QwenCausalDecodeRuntimeConfig qwen_config_; + std::shared_ptr weights_; + std::unique_ptr qwen_runtime_; + sampling::TorchCudaSamplingPolicy sampling_policy_; +}; + +CosyVoice3ArRuntime::CosyVoice3ArRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type) + : impl_(std::make_unique( + std::move(assets), + execution, + graph_arena_bytes, + weight_context_bytes, + storage_type)) {} + +CosyVoice3ArRuntime::~CosyVoice3ArRuntime() = default; + +CosyVoice3ArOutput CosyVoice3ArRuntime::generate(const CosyVoice3ArRequest & request) { + return impl_->generate(request); +} + +void CosyVoice3ArRuntime::release_graphs() { + impl_->release_graphs(); +} + +} // namespace engine::models::cosyvoice3 diff --git a/src/models/cosyvoice3/assets.cpp b/src/models/cosyvoice3/assets.cpp new file mode 100644 index 000000000..5d8a3972b --- /dev/null +++ b/src/models/cosyvoice3/assets.cpp @@ -0,0 +1,69 @@ +#include "engine/models/cosyvoice3/assets.h" + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/model_spec/package.h" + +#include + +namespace engine::models::cosyvoice3 { +namespace { + +constexpr const char * kFamily = "cosyvoice3"; + +std::filesystem::path find_gguf_path(const engine::assets::ResourceBundle & resources) { + for (const auto & file : resources.files()) { + if (file.path.extension() == ".gguf") { + return file.path; + } + } + return {}; +} + +void validate_llm_shapes(const engine::assets::TensorSource & source, CosyVoice3Config & config) { + const auto embedding = source.require_metadata("llm.model.model.embed_tokens.weight"); + if (embedding.shape.size() != 2) { + throw std::runtime_error("CosyVoice3 LLM embedding must be rank 2"); + } + config.text_vocab_size = embedding.shape[0]; + config.hidden_size = embedding.shape[1]; + engine::assets::require_tensor_shape(source, "speech_embedding.weight", {config.speech_token_size + config.speech_reserved_tokens, config.hidden_size}); + engine::assets::require_tensor_shape(source, "llm_decoder.weight", {config.speech_token_size + config.speech_reserved_tokens, config.hidden_size}); + engine::assets::require_tensor_shape(source, "llm.model.model.layers.0.self_attn.q_proj.weight", {config.heads * config.head_dim, config.hidden_size}); + engine::assets::require_tensor_shape(source, "llm.model.model.layers.0.self_attn.k_proj.weight", {config.kv_heads * config.head_dim, config.hidden_size}); + engine::assets::require_tensor_shape(source, "llm.model.model.layers.0.self_attn.v_proj.weight", {config.kv_heads * config.head_dim, config.hidden_size}); +} + +void validate_flow_shapes(const engine::assets::TensorSource & source, const CosyVoice3Config & config) { + engine::assets::require_tensor_shape(source, "input_embedding.weight", {config.speech_token_size, config.flow_mel_channels}); + engine::assets::require_tensor_shape(source, "spk_embed_affine_layer.weight", {config.flow_mel_channels, config.speaker_dim}); + engine::assets::require_tensor_shape(source, "decoder.rand_noise", {1, config.flow_mel_channels, 50 * 300}); + engine::assets::require_tensor_shape(source, "decoder.estimator.input_embed.proj.weight", {config.flow_hidden_size, 320}); + engine::assets::require_tensor_shape(source, "decoder.estimator.proj_out.weight", {config.flow_mel_channels, config.flow_hidden_size}); +} + +void validate_hift_shapes(const engine::assets::TensorSource & source) { + engine::assets::require_tensor_shape(source, "conv_pre.parametrizations.weight.original1", {512, 80, 5}); + engine::assets::require_tensor_shape(source, "conv_post.parametrizations.weight.original1", {18, 64, 7}); +} + +} // namespace + +std::shared_ptr load_cosyvoice3_assets(const std::filesystem::path & model_path) { + auto assets = std::make_shared(); + assets->resources = engine::model_spec::load_resource_bundle_for_family(model_path, kFamily); + assets->model_root = assets->resources.model_root(); + assets->gguf_path = find_gguf_path(assets->resources); + assets->llm_weights = assets->resources.open_tensor_source("llm_weights"); + assets->flow_weights = assets->resources.open_tensor_source("flow_weights"); + assets->hift_weights = assets->resources.open_tensor_source("hift_weights"); + assets->campplus_weights = assets->resources.open_tensor_source("campplus_weights"); + assets->speech_tokenizer_weights = assets->resources.open_tensor_source("speech_tokenizer_weights"); + assets->blank_en_weights = assets->resources.open_tensor_source("blank_en_weights"); + + validate_llm_shapes(*assets->llm_weights, assets->config); + validate_flow_shapes(*assets->flow_weights, assets->config); + validate_hift_shapes(*assets->hift_weights); + return assets; +} + +} // namespace engine::models::cosyvoice3 diff --git a/src/models/cosyvoice3/flow.cpp b/src/models/cosyvoice3/flow.cpp new file mode 100644 index 000000000..e81256724 --- /dev/null +++ b/src/models/cosyvoice3/flow.cpp @@ -0,0 +1,917 @@ +#include "engine/models/cosyvoice3/flow.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/attention/feed_forward.h" +#include "engine/framework/modules/attention/projected_grouped_self_attention.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::cosyvoice3 { +namespace { + +namespace binding = engine::modules::binding; +namespace core = engine::core; +namespace modules = engine::modules; + +constexpr float kPi = 3.14159265358979323846F; +constexpr float kInferenceCfgRate = 0.7F; +constexpr int64_t kTimeEmbeddingSize = 256; +constexpr int64_t kConvPosGroups = 16; +constexpr int64_t kConvPosKernel = 31; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct GgmlGallocrDeleter { + void operator()(ggml_gallocr_t alloc) const noexcept { + if (alloc != nullptr) { + ggml_gallocr_free(alloc); + } + } +}; + +struct GraphMemory { + std::unique_ptr ctx; + std::unique_ptr input_ctx; + std::unique_ptr, GgmlGallocrDeleter> gallocr; + ggml_backend_buffer_t input_buffer = nullptr; + ggml_cgraph * graph = nullptr; + + ~GraphMemory() { + reset(nullptr); + } + + void reset(ggml_backend_t backend) { + if (graph != nullptr && backend != nullptr) { + core::release_backend_graph_resources(backend, graph); + } + graph = nullptr; + gallocr.reset(); + if (input_buffer != nullptr) { + ggml_backend_buffer_free(input_buffer); + input_buffer = nullptr; + } + input_ctx.reset(); + ctx.reset(); + } +}; + +std::vector position_ids(int64_t steps) { + std::vector out(static_cast(steps)); + for (int64_t i = 0; i < steps; ++i) { + out[static_cast(i)] = static_cast(i); + } + return out; +} + +std::vector trace_dims(const core::TensorShape & shape) { + std::vector out; + out.reserve(shape.rank); + for (size_t i = 0; i < shape.rank; ++i) { + out.push_back(shape.dims[i]); + } + return out; +} + +std::vector cosine_time_schedule(int64_t steps) { + if (steps <= 0) { + throw std::runtime_error("CosyVoice3 num_inference_steps must be positive"); + } + std::vector out(static_cast(steps + 1)); + for (int64_t i = 0; i <= steps; ++i) { + const float u = static_cast(i) / static_cast(steps); + out[static_cast(i)] = 1.0F - std::cos(u * 0.5F * kPi); + } + return out; +} + +std::vector timestep_embedding(float timestep) { + const int64_t half = kTimeEmbeddingSize / 2; + const float step = std::log(10000.0F) / static_cast(half - 1); + std::vector out(static_cast(kTimeEmbeddingSize)); + for (int64_t i = 0; i < half; ++i) { + const float freq = std::exp(static_cast(i) * -step); + const float arg = 1000.0F * timestep * freq; + out[static_cast(i)] = std::sin(arg); + out[static_cast(half + i)] = std::cos(arg); + } + return out; +} + +core::TensorValue repeat_like( + core::ModuleBuildContext & ctx, + const core::TensorValue & value, + const core::TensorValue & like) { + return modules::RepeatModule({like.shape}).build(ctx, value); +} + +core::TensorValue mul_broadcast( + core::ModuleBuildContext & ctx, + const core::TensorValue & x, + const core::TensorValue & scale) { + return modules::MulModule().build(ctx, x, repeat_like(ctx, scale, x)); +} + +core::TensorValue modulate( + core::ModuleBuildContext & ctx, + const core::TensorValue & x, + const core::TensorValue & shift, + const core::TensorValue & scale) { + auto one_plus = core::wrap_tensor( + ggml_scale_bias(ctx.ggml, repeat_like(ctx, scale, x).tensor, 1.0F, 1.0F), + x.shape, + GGML_TYPE_F32); + auto shifted = modules::AddModule().build(ctx, modules::MulModule().build(ctx, x, one_plus), repeat_like(ctx, shift, x)); + return shifted; +} + +core::TensorValue mish(core::ModuleBuildContext & ctx, const core::TensorValue & x) { + auto input = core::ensure_backend_addressable_layout(ctx, x); + auto softplus = core::wrap_tensor(ggml_softplus(ctx.ggml, input.tensor), input.shape, GGML_TYPE_F32); + auto t = core::wrap_tensor(ggml_tanh(ctx.ggml, softplus.tensor), input.shape, GGML_TYPE_F32); + return modules::MulModule().build(ctx, input, t); +} + +core::TensorValue grouped_conv1d( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::Conv1dWeights & weights, + int64_t groups) { + const int64_t channels = input.shape.dims[1]; + const int64_t out_channels = weights.weight.shape.dims[0]; + const int64_t kernel = weights.weight.shape.dims[2]; + if (groups <= 0 || channels % groups != 0 || out_channels % groups != 0) { + throw std::runtime_error("CosyVoice3 grouped Conv1d channel/group mismatch"); + } + const int64_t channels_per_group = channels / groups; + const int64_t out_channels_per_group = out_channels / groups; + core::TensorValue output; + const auto input_contiguous = core::ensure_backend_addressable_layout(ctx, input); + for (int64_t group = 0; group < groups; ++group) { + auto input_group = modules::SliceModule({1, group * channels_per_group, channels_per_group}).build(ctx, input_contiguous); + auto weight_group = modules::SliceModule({0, group * out_channels_per_group, out_channels_per_group}).build(ctx, weights.weight); + modules::Conv1dWeights group_weights{weight_group, std::nullopt}; + if (weights.bias.has_value()) { + group_weights.bias = + modules::SliceModule({0, group * out_channels_per_group, out_channels_per_group}).build(ctx, *weights.bias); + } + auto group_out = modules::Conv1dModule({ + channels_per_group, + out_channels_per_group, + kernel, + 1, + 0, + 1, + weights.bias.has_value()}).build(ctx, input_group, group_weights); + output = output.valid() ? modules::ConcatModule({1}).build(ctx, output, group_out) : group_out; + } + return output; +} + +struct CosyDiTBlockWeights { + modules::LinearWeights attn_norm; + modules::ProjectedGroupedSelfAttentionWeights attention; + modules::NormWeights ff_norm; + modules::FeedForwardWeights ff; +}; + +struct CosyFlowWeights { + std::shared_ptr store; + core::TensorValue token_embedding; + modules::LinearWeights speaker_projection; + modules::Conv1dWeights pre_lookahead_conv1; + modules::Conv1dWeights pre_lookahead_conv2; + modules::LinearWeights input_projection; + modules::Conv1dWeights conv_pos_1; + modules::Conv1dWeights conv_pos_2; + modules::LinearWeights time_fc1; + modules::LinearWeights time_fc2; + std::vector blocks; + modules::LinearWeights final_norm; + modules::LinearWeights output_projection; + core::TensorValue rand_noise; + std::vector rand_noise_host; +}; + +modules::ProjectedGroupedSelfAttentionWeights load_attention( + core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + int64_t hidden, + engine::assets::TensorStorageType storage_type) { + modules::ProjectedGroupedSelfAttentionWeights out; + out.q_proj = binding::linear_from_source(store, source, prefix + ".to_q", storage_type, hidden, hidden, true); + out.k_proj = binding::linear_from_source(store, source, prefix + ".to_k", storage_type, hidden, hidden, true); + out.v_proj = binding::linear_from_source(store, source, prefix + ".to_v", storage_type, hidden, hidden, true); + out.o_proj = binding::linear_from_source(store, source, prefix + ".to_out.0", storage_type, hidden, hidden, true); + return out; +} + +CosyDiTBlockWeights load_block( + core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + const CosyVoice3Config & config, + engine::assets::TensorStorageType storage_type) { + CosyDiTBlockWeights out; + out.attn_norm = binding::linear_from_source( + store, source, prefix + ".attn_norm.linear", storage_type, 6 * config.flow_hidden_size, config.flow_hidden_size, true); + out.attention = load_attention(store, source, prefix + ".attn", config.flow_hidden_size, storage_type); + out.ff.fc1_weight = store.load_tensor( + source, + prefix + ".ff.ff.0.0.weight", + storage_type, + {config.flow_hidden_size * config.flow_ff_mult, config.flow_hidden_size}); + out.ff.fc1_bias = store.load_f32_tensor(source, prefix + ".ff.ff.0.0.bias", {config.flow_hidden_size * config.flow_ff_mult}); + out.ff.fc2_weight = store.load_tensor( + source, + prefix + ".ff.ff.2.weight", + storage_type, + {config.flow_hidden_size, config.flow_hidden_size * config.flow_ff_mult}); + out.ff.fc2_bias = store.load_f32_tensor(source, prefix + ".ff.ff.2.bias", {config.flow_hidden_size}); + return out; +} + +std::shared_ptr load_flow_weights( + const CosyVoice3Assets & assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type) { + auto weights = std::make_shared(); + weights->store = std::make_shared( + execution.backend(), + execution.backend_type(), + "cosyvoice3.flow.weights", + weight_context_bytes); + const auto & source = *assets.flow_weights; + const auto & c = assets.config; + weights->token_embedding = weights->store->load_tensor(source, "input_embedding.weight", storage_type, {c.speech_token_size, c.flow_mel_channels}); + weights->speaker_projection = binding::linear_from_source( + *weights->store, source, "spk_embed_affine_layer", storage_type, c.flow_mel_channels, c.speaker_dim, true); + weights->pre_lookahead_conv1 = binding::conv1d_from_source( + *weights->store, source, "pre_lookahead_layer.conv1", storage_type, c.flow_hidden_size, c.flow_mel_channels, c.pre_lookahead_len + 1, true); + weights->pre_lookahead_conv2 = binding::conv1d_from_source( + *weights->store, source, "pre_lookahead_layer.conv2", storage_type, c.flow_mel_channels, c.flow_hidden_size, 3, true); + weights->input_projection = binding::linear_from_source( + *weights->store, source, "decoder.estimator.input_embed.proj", storage_type, c.flow_hidden_size, 320, true); + weights->conv_pos_1.weight = weights->store->load_tensor( + source, + "decoder.estimator.input_embed.conv_pos_embed.conv1.0.weight", + storage_type, + {c.flow_hidden_size, c.flow_hidden_size / kConvPosGroups, kConvPosKernel}); + weights->conv_pos_1.bias = weights->store->load_f32_tensor( + source, + "decoder.estimator.input_embed.conv_pos_embed.conv1.0.bias", + {c.flow_hidden_size}); + weights->conv_pos_2.weight = weights->store->load_tensor( + source, + "decoder.estimator.input_embed.conv_pos_embed.conv2.0.weight", + storage_type, + {c.flow_hidden_size, c.flow_hidden_size / kConvPosGroups, kConvPosKernel}); + weights->conv_pos_2.bias = weights->store->load_f32_tensor( + source, + "decoder.estimator.input_embed.conv_pos_embed.conv2.0.bias", + {c.flow_hidden_size}); + weights->time_fc1 = binding::linear_from_source( + *weights->store, source, "decoder.estimator.time_embed.time_mlp.0", storage_type, c.flow_hidden_size, kTimeEmbeddingSize, true); + weights->time_fc2 = binding::linear_from_source( + *weights->store, source, "decoder.estimator.time_embed.time_mlp.2", storage_type, c.flow_hidden_size, c.flow_hidden_size, true); + weights->blocks.reserve(static_cast(c.flow_layers)); + for (int64_t layer = 0; layer < c.flow_layers; ++layer) { + weights->blocks.push_back(load_block( + *weights->store, + source, + "decoder.estimator.transformer_blocks." + std::to_string(layer), + c, + storage_type)); + } + weights->final_norm = binding::linear_from_source( + *weights->store, source, "decoder.estimator.norm_out.linear", storage_type, 2 * c.flow_hidden_size, c.flow_hidden_size, true); + weights->output_projection = binding::linear_from_source( + *weights->store, source, "decoder.estimator.proj_out", storage_type, c.flow_mel_channels, c.flow_hidden_size, true); + weights->rand_noise = weights->store->load_tensor(source, "decoder.rand_noise", engine::assets::TensorStorageType::F32, {1, c.flow_mel_channels, 50 * 300}); + weights->rand_noise_host = source.require_f32("decoder.rand_noise", {1, c.flow_mel_channels, 50 * 300}); + weights->store->upload(); + return weights; +} + +modules::ProjectedGroupedSelfAttentionConfig attention_config(const CosyVoice3Config & config) { + modules::ProjectedGroupedSelfAttentionConfig out; + out.hidden_size = config.flow_hidden_size; + out.attention_heads = config.flow_heads; + out.kv_heads = config.flow_heads; + out.head_dim = config.flow_head_dim; + out.use_bias = true; + out.use_rope = true; + out.apply_rope_to_projected_prefix = true; + out.rope_type = GGML_ROPE_TYPE_NORMAL; + out.rope_theta = 10000.0F; + out.local_rope_theta = 10000.0F; + out.causality = modules::AttentionCausality::NonCausal; + out.lowering = modules::GroupedQueryAttentionLowering::FlashGroupedViewKV; + out.attention_precision = GGML_PREC_F32; + return out; +} + +core::TensorValue causal_conv_pos_embed( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const CosyFlowWeights & weights) { + auto x = modules::TransposeModule({{0, 2, 1, 3}, input.shape.rank}).build(ctx, input); + x = core::wrap_tensor( + ggml_pad_ext(ctx.ggml, x.tensor, kConvPosKernel - 1, 0, 0, 0, 0, 0, 0, 0), + core::TensorShape::from_dims({x.shape.dims[0], x.shape.dims[1], x.shape.dims[2] + kConvPosKernel - 1}), + GGML_TYPE_F32); + x = grouped_conv1d(ctx, x, weights.conv_pos_1, kConvPosGroups); + x = mish(ctx, x); + x = core::wrap_tensor( + ggml_pad_ext(ctx.ggml, x.tensor, kConvPosKernel - 1, 0, 0, 0, 0, 0, 0, 0), + core::TensorShape::from_dims({x.shape.dims[0], x.shape.dims[1], x.shape.dims[2] + kConvPosKernel - 1}), + GGML_TYPE_F32); + x = grouped_conv1d(ctx, x, weights.conv_pos_2, kConvPosGroups); + x = mish(ctx, x); + x = modules::TransposeModule({{0, 2, 1, 3}, x.shape.rank}).build(ctx, x); + return x; +} + +core::TensorValue dit_block( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & time, + const core::TensorValue & positions, + const CosyDiTBlockWeights & weights, + const CosyVoice3Config & config, + int64_t layer) { + auto ada = modules::SiluModule().build(ctx, time); + ada = modules::LinearModule({config.flow_hidden_size, 6 * config.flow_hidden_size, true}).build(ctx, ada, weights.attn_norm); + auto shift_msa = modules::SliceModule({2, 0 * config.flow_hidden_size, config.flow_hidden_size}).build(ctx, ada); + auto scale_msa = modules::SliceModule({2, 1 * config.flow_hidden_size, config.flow_hidden_size}).build(ctx, ada); + auto gate_msa = modules::SliceModule({2, 2 * config.flow_hidden_size, config.flow_hidden_size}).build(ctx, ada); + auto shift_mlp = modules::SliceModule({2, 3 * config.flow_hidden_size, config.flow_hidden_size}).build(ctx, ada); + auto scale_mlp = modules::SliceModule({2, 4 * config.flow_hidden_size, config.flow_hidden_size}).build(ctx, ada); + auto gate_mlp = modules::SliceModule({2, 5 * config.flow_hidden_size, config.flow_hidden_size}).build(ctx, ada); + + auto x = modules::LayerNormModule({config.flow_hidden_size, 1.0e-6F, false, false}).build(ctx, input, {}); + x = modulate(ctx, x, shift_msa, scale_msa); + x = modules::ProjectedGroupedSelfAttentionModule(attention_config(config)).build(ctx, x, positions, weights.attention, layer); + x = mul_broadcast(ctx, x, gate_msa); + auto out = modules::AddModule().build(ctx, input, x); + + x = modules::LayerNormModule({config.flow_hidden_size, 1.0e-6F, false, false}).build(ctx, out, {}); + x = modulate(ctx, x, shift_mlp, scale_mlp); + x = modules::FeedForwardModule({ + config.flow_hidden_size, + config.flow_hidden_size * config.flow_ff_mult, + true, + modules::GeluApproximation::Tanh, + }).build(ctx, x, weights.ff); + x = mul_broadcast(ctx, x, gate_mlp); + return modules::AddModule().build(ctx, out, x); +} + +class ConditionGraph { +public: + ConditionGraph( + core::ExecutionContext & execution, + std::shared_ptr weights, + CosyVoice3Config config, + size_t graph_arena_bytes) + : execution_(execution), + weights_(std::move(weights)), + config_(config), + graph_arena_bytes_(graph_arena_bytes) {} + + ~ConditionGraph() { + mem_.reset(execution_.backend()); + } + + std::vector run(const std::vector & prompt_tokens, const std::vector & target_tokens) { + const int64_t tokens = static_cast(prompt_tokens.size() + target_tokens.size()); + if (tokens <= 0) { + throw std::runtime_error("CosyVoice3 flow requires speech tokens"); + } + ensure(tokens); + std::vector all_tokens; + all_tokens.reserve(static_cast(tokens)); + all_tokens.insert(all_tokens.end(), prompt_tokens.begin(), prompt_tokens.end()); + all_tokens.insert(all_tokens.end(), target_tokens.begin(), target_tokens.end()); + ggml_backend_tensor_set(token_ids_, all_tokens.data(), 0, all_tokens.size() * sizeof(int32_t)); + if (core::compute_backend_graph(execution_.backend(), mem_.graph, nullptr, "cosyvoice3.flow.condition") != GGML_STATUS_SUCCESS) { + throw std::runtime_error("CosyVoice3 flow condition graph compute failed"); + } + auto output = core::read_tensor_f32(output_.tensor); + engine::debug::trace_log_f32("cosyvoice3.flow.condition_mu", {tokens * config_.token_mel_ratio, config_.flow_mel_channels}, output); + return output; + } + + void release_graph() { + mem_.reset(execution_.backend()); + tokens_ = 0; + token_ids_ = nullptr; + output_ = {}; + } + +private: + void ensure(int64_t tokens) { + if (mem_.graph != nullptr && tokens_ == tokens) { + return; + } + mem_.reset(execution_.backend()); + ggml_init_params params{graph_arena_bytes_, nullptr, true}; + mem_.ctx.reset(ggml_init(params)); + ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; + mem_.input_ctx.reset(ggml_init(input_params)); + core::ModuleBuildContext ctx{mem_.ctx.get(), "cosyvoice3.flow.condition", execution_.backend_type()}; + core::ModuleBuildContext input_ctx{mem_.input_ctx.get(), "cosyvoice3.flow.condition.inputs", execution_.backend_type()}; + + auto ids = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, tokens})); + token_ids_ = ids.tensor; + ggml_set_input(token_ids_); + auto h = modules::EmbeddingModule({config_.speech_token_size, config_.flow_mel_channels}) + .build(ctx, ids, weights_->token_embedding); + h = modules::TransposeModule({{0, 2, 1, 3}, h.shape.rank}).build(ctx, h); + h = core::wrap_tensor( + ggml_pad_ext(ctx.ggml, h.tensor, 0, config_.pre_lookahead_len, 0, 0, 0, 0, 0, 0), + core::TensorShape::from_dims({1, config_.flow_mel_channels, tokens + config_.pre_lookahead_len}), + GGML_TYPE_F32); + h = modules::Conv1dModule({config_.flow_mel_channels, config_.flow_hidden_size, config_.pre_lookahead_len + 1, 1, 0, 1, true}) + .build(ctx, h, weights_->pre_lookahead_conv1); + h = modules::LeakyReluModule().build(ctx, h); + h = core::wrap_tensor( + ggml_pad_ext(ctx.ggml, h.tensor, 2, 0, 0, 0, 0, 0, 0, 0), + core::TensorShape::from_dims({1, config_.flow_hidden_size, tokens + 2}), + GGML_TYPE_F32); + h = modules::Conv1dModule({config_.flow_hidden_size, config_.flow_mel_channels, 3, 1, 0, 1, true}) + .build(ctx, h, weights_->pre_lookahead_conv2); + h = modules::TransposeModule({{0, 2, 1, 3}, h.shape.rank}).build(ctx, h); + h = modules::AddModule().build( + ctx, + h, + modules::EmbeddingModule({config_.speech_token_size, config_.flow_mel_channels}) + .build(ctx, ids, weights_->token_embedding)); + h = modules::Interpolate1dModule({tokens * config_.token_mel_ratio, modules::Interpolate1dMode::Nearest}) + .build(ctx, modules::TransposeModule({{0, 2, 1, 3}, h.shape.rank}).build(ctx, h)); + h = modules::TransposeModule({{0, 2, 1, 3}, h.shape.rank}).build(ctx, h); + output_ = core::ensure_backend_addressable_layout(ctx, h); + ggml_set_output(output_.tensor); + mem_.graph = ggml_new_graph_custom(mem_.ctx.get(), 20000, false); + ggml_build_forward_expand(mem_.graph, output_.tensor); + mem_.input_buffer = ggml_backend_alloc_ctx_tensors(mem_.input_ctx.get(), execution_.backend()); + mem_.gallocr.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend()))); + if (mem_.input_buffer == nullptr || mem_.gallocr == nullptr || + !ggml_gallocr_reserve(mem_.gallocr.get(), mem_.graph) || + !ggml_gallocr_alloc_graph(mem_.gallocr.get(), mem_.graph)) { + mem_.reset(execution_.backend()); + throw std::runtime_error("failed to allocate CosyVoice3 flow condition graph"); + } + tokens_ = tokens; + } + + core::ExecutionContext & execution_; + std::shared_ptr weights_; + CosyVoice3Config config_; + size_t graph_arena_bytes_ = 0; + GraphMemory mem_; + int64_t tokens_ = 0; + ggml_tensor * token_ids_ = nullptr; + core::TensorValue output_; +}; + +class DiTGraph { +public: + DiTGraph( + core::ExecutionContext & execution, + std::shared_ptr weights, + CosyVoice3Config config, + size_t graph_arena_bytes) + : execution_(execution), + weights_(std::move(weights)), + config_(config), + graph_arena_bytes_(graph_arena_bytes) {} + + ~DiTGraph() { + mem_.reset(execution_.backend()); + } + + std::vector run( + const std::vector & x, + const std::vector & mu, + const std::vector & cond, + const std::vector & speaker, + const std::vector & time_embedding, + int64_t frames) { + if (frames <= 0) { + throw std::runtime_error("CosyVoice3 DiT frames must be positive"); + } + const int64_t batch = 2; + const int64_t channels = config_.flow_mel_channels; + if (static_cast(x.size()) != batch * channels * frames || + static_cast(mu.size()) != batch * channels * frames || + static_cast(cond.size()) != batch * channels * frames || + static_cast(speaker.size()) != batch * channels || + static_cast(time_embedding.size()) != batch * kTimeEmbeddingSize) { + throw std::runtime_error("CosyVoice3 DiT input size mismatch"); + } + ensure(frames); + ggml_backend_tensor_set(x_, x.data(), 0, x.size() * sizeof(float)); + ggml_backend_tensor_set(mu_, mu.data(), 0, mu.size() * sizeof(float)); + ggml_backend_tensor_set(cond_, cond.data(), 0, cond.size() * sizeof(float)); + ggml_backend_tensor_set(spks_, speaker.data(), 0, speaker.size() * sizeof(float)); + ggml_backend_tensor_set(time_, time_embedding.data(), 0, time_embedding.size() * sizeof(float)); + if (core::compute_backend_graph(execution_.backend(), mem_.graph, nullptr, "cosyvoice3.flow.dit") != GGML_STATUS_SUCCESS) { + throw std::runtime_error("CosyVoice3 DiT graph compute failed"); + } + if (!traced_first_step_ && engine::debug::trace_log_enabled()) { + traced_first_step_ = true; + engine::debug::trace_log_f32("cosyvoice3.flow.dit.input_cat", trace_dims(debug_input_cat_.shape), core::read_tensor_f32(debug_input_cat_.tensor)); + engine::debug::trace_log_f32("cosyvoice3.flow.dit.input_proj", trace_dims(debug_input_proj_.shape), core::read_tensor_f32(debug_input_proj_.tensor)); + engine::debug::trace_log_f32("cosyvoice3.flow.dit.input_pos", trace_dims(debug_input_pos_.shape), core::read_tensor_f32(debug_input_pos_.tensor)); + engine::debug::trace_log_f32("cosyvoice3.flow.dit.input_embed", trace_dims(debug_input_embed_.shape), core::read_tensor_f32(debug_input_embed_.tensor)); + } + return core::read_tensor_f32(output_.tensor); + } + + void release_graph() { + mem_.reset(execution_.backend()); + frames_ = 0; + x_ = nullptr; + mu_ = nullptr; + cond_ = nullptr; + spks_ = nullptr; + time_ = nullptr; + positions_ = nullptr; + output_ = {}; + debug_input_cat_ = {}; + debug_input_proj_ = {}; + debug_input_pos_ = {}; + debug_input_embed_ = {}; + traced_first_step_ = false; + } + +private: + void ensure(int64_t frames) { + if (mem_.graph != nullptr && frames_ == frames) { + return; + } + mem_.reset(execution_.backend()); + constexpr int64_t batch = 2; + const int64_t channels = config_.flow_mel_channels; + ggml_init_params params{graph_arena_bytes_, nullptr, true}; + mem_.ctx.reset(ggml_init(params)); + ggml_init_params input_params{32ull * 1024ull * 1024ull, nullptr, true}; + mem_.input_ctx.reset(ggml_init(input_params)); + core::ModuleBuildContext ctx{mem_.ctx.get(), "cosyvoice3.flow.dit", execution_.backend_type()}; + core::ModuleBuildContext input_ctx{mem_.input_ctx.get(), "cosyvoice3.flow.dit.inputs", execution_.backend_type()}; + + auto x = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch, channels, frames})); + x_ = x.tensor; + ggml_set_input(x_); + auto mu = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch, channels, frames})); + mu_ = mu.tensor; + ggml_set_input(mu_); + auto cond = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch, channels, frames})); + cond_ = cond.tensor; + ggml_set_input(cond_); + auto spks = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch, channels})); + spks_ = spks.tensor; + ggml_set_input(spks_); + auto time = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch, kTimeEmbeddingSize})); + time_ = time.tensor; + ggml_set_input(time_); + auto positions = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({frames})); + positions_ = positions.tensor; + + x = modules::TransposeModule({{0, 2, 1, 3}, x.shape.rank}).build(ctx, x); + mu = modules::TransposeModule({{0, 2, 1, 3}, mu.shape.rank}).build(ctx, mu); + cond = modules::TransposeModule({{0, 2, 1, 3}, cond.shape.rank}).build(ctx, cond); + x = core::ensure_backend_addressable_layout(ctx, x); + mu = core::ensure_backend_addressable_layout(ctx, mu); + cond = core::ensure_backend_addressable_layout(ctx, cond); + auto spks_btf = core::reshape_tensor(ctx, spks, core::TensorShape::from_dims({batch, 1, channels})); + spks_btf = modules::RepeatModule({core::TensorShape::from_dims({batch, frames, channels})}).build(ctx, spks_btf); + spks_btf = core::ensure_backend_addressable_layout(ctx, spks_btf); + auto input = modules::ConcatModule({2}).build(ctx, x, cond); + input = modules::ConcatModule({2}).build(ctx, input, mu); + input = modules::ConcatModule({2}).build(ctx, input, spks_btf); + input = core::ensure_backend_addressable_layout(ctx, input); + debug_input_cat_ = input; + input = modules::LinearModule({320, config_.flow_hidden_size, true}).build(ctx, input, weights_->input_projection); + debug_input_proj_ = core::ensure_backend_addressable_layout(ctx, input); + input = debug_input_proj_; + auto pos = causal_conv_pos_embed(ctx, input, *weights_); + debug_input_pos_ = core::ensure_backend_addressable_layout(ctx, pos); + pos = debug_input_pos_; + input = modules::AddModule().build(ctx, input, pos); + debug_input_embed_ = core::ensure_backend_addressable_layout(ctx, input); + input = debug_input_embed_; + + time = modules::LinearModule({kTimeEmbeddingSize, config_.flow_hidden_size, true}).build(ctx, time, weights_->time_fc1); + time = modules::SiluModule().build(ctx, time); + time = modules::LinearModule({config_.flow_hidden_size, config_.flow_hidden_size, true}).build(ctx, time, weights_->time_fc2); + time = core::reshape_tensor(ctx, time, core::TensorShape::from_dims({batch, 1, config_.flow_hidden_size})); + + for (size_t layer = 0; layer < weights_->blocks.size(); ++layer) { + input = dit_block(ctx, input, time, positions, weights_->blocks[layer], config_, static_cast(layer)); + } + auto norm_mod = modules::SiluModule().build(ctx, time); + norm_mod = modules::LinearModule({config_.flow_hidden_size, 2 * config_.flow_hidden_size, true}).build(ctx, norm_mod, weights_->final_norm); + auto scale = modules::SliceModule({2, 0, config_.flow_hidden_size}).build(ctx, norm_mod); + auto shift = modules::SliceModule({2, config_.flow_hidden_size, config_.flow_hidden_size}).build(ctx, norm_mod); + input = modules::LayerNormModule({config_.flow_hidden_size, 1.0e-6F, false, false}).build(ctx, input, {}); + input = modulate(ctx, input, shift, scale); + input = modules::LinearModule({config_.flow_hidden_size, channels, true}).build(ctx, input, weights_->output_projection); + input = modules::TransposeModule({{0, 2, 1, 3}, input.shape.rank}).build(ctx, input); + output_ = core::ensure_backend_addressable_layout(ctx, input); + ggml_set_output(output_.tensor); + mem_.graph = ggml_new_graph_custom(mem_.ctx.get(), 200000, false); + ggml_build_forward_expand(mem_.graph, output_.tensor); + mem_.input_buffer = ggml_backend_alloc_ctx_tensors(mem_.input_ctx.get(), execution_.backend()); + mem_.gallocr.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend()))); + if (mem_.input_buffer == nullptr || mem_.gallocr == nullptr || + !ggml_gallocr_reserve(mem_.gallocr.get(), mem_.graph) || + !ggml_gallocr_alloc_graph(mem_.gallocr.get(), mem_.graph)) { + mem_.reset(execution_.backend()); + throw std::runtime_error("failed to allocate CosyVoice3 DiT graph"); + } + const auto pos_ids = position_ids(frames); + ggml_backend_tensor_set(positions_, pos_ids.data(), 0, pos_ids.size() * sizeof(int32_t)); + frames_ = frames; + } + + core::ExecutionContext & execution_; + std::shared_ptr weights_; + CosyVoice3Config config_; + size_t graph_arena_bytes_ = 0; + GraphMemory mem_; + int64_t frames_ = 0; + ggml_tensor * x_ = nullptr; + ggml_tensor * mu_ = nullptr; + ggml_tensor * cond_ = nullptr; + ggml_tensor * spks_ = nullptr; + ggml_tensor * time_ = nullptr; + ggml_tensor * positions_ = nullptr; + core::TensorValue output_; + core::TensorValue debug_input_cat_; + core::TensorValue debug_input_proj_; + core::TensorValue debug_input_pos_; + core::TensorValue debug_input_embed_; + bool traced_first_step_ = false; +}; + +std::vector normalize_and_project_speaker( + const CosyVoice3Config & config, + const CosyFlowWeights & weights, + core::ExecutionContext & execution, + const std::vector & speaker_embedding, + size_t graph_arena_bytes) { + if (static_cast(speaker_embedding.size()) != config.speaker_dim) { + throw std::runtime_error("CosyVoice3 speaker embedding size mismatch"); + } + float norm_sq = 0.0F; + for (float value : speaker_embedding) { + norm_sq += value * value; + } + const float inv_norm = 1.0F / std::sqrt(std::max(norm_sq, 1.0e-12F)); + std::vector normalized(speaker_embedding.size()); + for (size_t i = 0; i < speaker_embedding.size(); ++i) { + normalized[i] = speaker_embedding[i] * inv_norm; + } + + GraphMemory mem; + ggml_init_params params{graph_arena_bytes, nullptr, true}; + mem.ctx.reset(ggml_init(params)); + ggml_init_params input_params{4ull * 1024ull * 1024ull, nullptr, true}; + mem.input_ctx.reset(ggml_init(input_params)); + core::ModuleBuildContext ctx{mem.ctx.get(), "cosyvoice3.flow.speaker", execution.backend_type()}; + core::ModuleBuildContext input_ctx{mem.input_ctx.get(), "cosyvoice3.flow.speaker.inputs", execution.backend_type()}; + auto input = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, config.speaker_dim})); + ggml_set_input(input.tensor); + auto output = modules::LinearModule({config.speaker_dim, config.flow_mel_channels, true}).build(ctx, input, weights.speaker_projection); + output = core::ensure_backend_addressable_layout(ctx, output); + ggml_set_output(output.tensor); + mem.graph = ggml_new_graph_custom(mem.ctx.get(), 4096, false); + ggml_build_forward_expand(mem.graph, output.tensor); + mem.input_buffer = ggml_backend_alloc_ctx_tensors(mem.input_ctx.get(), execution.backend()); + mem.gallocr.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution.backend()))); + if (mem.input_buffer == nullptr || mem.gallocr == nullptr || + !ggml_gallocr_reserve(mem.gallocr.get(), mem.graph) || + !ggml_gallocr_alloc_graph(mem.gallocr.get(), mem.graph)) { + throw std::runtime_error("CosyVoice3 speaker projection graph failed"); + } + ggml_backend_tensor_set(input.tensor, normalized.data(), 0, normalized.size() * sizeof(float)); + if (core::compute_backend_graph(execution.backend(), mem.graph, nullptr, "cosyvoice3.flow.speaker") != GGML_STATUS_SUCCESS) { + throw std::runtime_error("CosyVoice3 speaker projection graph failed"); + } + auto out = core::read_tensor_f32(output.tensor); + mem.reset(execution.backend()); + return out; +} + +std::vector read_noise_prefix( + const CosyVoice3Config & config, + const CosyFlowWeights & weights, + int64_t frames) { + if (frames > 50 * 300) { + throw std::runtime_error("CosyVoice3 requested mel frames exceed fixed flow noise capacity"); + } + const auto & full = weights.rand_noise_host; + std::vector out(static_cast(config.flow_mel_channels * frames)); + for (int64_t ch = 0; ch < config.flow_mel_channels; ++ch) { + const auto src = full.begin() + static_cast(ch * 50 * 300); + const auto dst = out.begin() + static_cast(ch * frames); + std::copy(src, src + frames, dst); + } + return out; +} + +} // namespace + +class CosyVoice3FlowRuntime::Impl { +public: + Impl( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type) + : assets_(std::move(assets)), + execution_(execution), + weights_(load_flow_weights(*assets_, execution_, weight_context_bytes, storage_type)), + condition_(execution_, weights_, assets_->config, graph_arena_bytes), + dit_(execution_, weights_, assets_->config, graph_arena_bytes), + graph_arena_bytes_(graph_arena_bytes) { + if (assets_ == nullptr) { + throw std::runtime_error("CosyVoice3 flow runtime requires assets"); + } + } + + CosyVoice3FlowOutput generate(const CosyVoice3FlowRequest & request) { + const auto start = std::chrono::steady_clock::now(); + const auto & c = assets_->config; + if (request.speech_tokens.empty()) { + throw std::runtime_error("CosyVoice3 flow requires generated speech tokens"); + } + if (request.prompt_mel_frames <= 0 || + static_cast(request.prompt_mel.size()) != request.prompt_mel_frames * c.flow_mel_channels) { + throw std::runtime_error("CosyVoice3 flow prompt mel size mismatch"); + } + const auto cond_start = std::chrono::steady_clock::now(); + auto mu_btc = condition_.run(request.prompt_speech_tokens, request.speech_tokens); + const int64_t total_frames = static_cast(mu_btc.size()) / c.flow_mel_channels; + const int64_t target_frames = total_frames - request.prompt_mel_frames; + if (target_frames <= 0) { + throw std::runtime_error("CosyVoice3 flow target mel frames must be positive"); + } + engine::debug::timing_log_scalar("cosyvoice3.flow.condition_ms", engine::debug::elapsed_ms(cond_start)); + + std::vector mu(static_cast(2 * c.flow_mel_channels * total_frames), 0.0F); + for (int64_t frame = 0; frame < total_frames; ++frame) { + for (int64_t ch = 0; ch < c.flow_mel_channels; ++ch) { + mu[static_cast(ch * total_frames + frame)] = + mu_btc[static_cast(frame * c.flow_mel_channels + ch)]; + } + } + std::vector cond(static_cast(2 * c.flow_mel_channels * total_frames), 0.0F); + for (int64_t frame = 0; frame < request.prompt_mel_frames; ++frame) { + for (int64_t ch = 0; ch < c.flow_mel_channels; ++ch) { + cond[static_cast(ch * total_frames + frame)] = + request.prompt_mel[static_cast(frame * c.flow_mel_channels + ch)]; + } + } + auto x = read_noise_prefix(c, *weights_, total_frames); + std::vector x_batched(static_cast(2 * c.flow_mel_channels * total_frames)); + std::copy(x.begin(), x.end(), x_batched.begin()); + std::copy(x.begin(), x.end(), x_batched.begin() + static_cast(x.size())); + + auto speaker_cond = normalize_and_project_speaker( + c, + *weights_, + execution_, + request.speaker_embedding, + graph_arena_bytes_); + std::vector spks(static_cast(2 * c.flow_mel_channels), 0.0F); + std::copy(speaker_cond.begin(), speaker_cond.end(), spks.begin()); + + const auto schedule = cosine_time_schedule(request.num_inference_steps); + const auto dit_start = std::chrono::steady_clock::now(); + for (int64_t step = 1; step < static_cast(schedule.size()); ++step) { + if (step == 1) { + engine::debug::trace_log_f32( + "cosyvoice3.flow.dit.x_in", + {2, c.flow_mel_channels, total_frames}, + x_batched); + engine::debug::trace_log_f32( + "cosyvoice3.flow.dit.mu_in", + {2, c.flow_mel_channels, total_frames}, + mu); + engine::debug::trace_log_f32( + "cosyvoice3.flow.dit.cond_in", + {2, c.flow_mel_channels, total_frames}, + cond); + engine::debug::trace_log_f32( + "cosyvoice3.flow.dit.spks_in", + {2, c.flow_mel_channels}, + spks); + } + const float t = schedule[static_cast(step - 1)]; + const float dt = schedule[static_cast(step)] - t; + auto temb = timestep_embedding(t); + std::vector time_batched(static_cast(2 * kTimeEmbeddingSize)); + std::copy(temb.begin(), temb.end(), time_batched.begin()); + std::copy(temb.begin(), temb.end(), time_batched.begin() + kTimeEmbeddingSize); + auto pred = dit_.run(x_batched, mu, cond, spks, time_batched, total_frames); + if (step == 1) { + engine::debug::trace_log_f32( + "cosyvoice3.flow.dit.pred_step1", + {2, c.flow_mel_channels, total_frames}, + pred); + } + const size_t branch = static_cast(c.flow_mel_channels * total_frames); + for (size_t i = 0; i < branch; ++i) { + const float guided = (1.0F + kInferenceCfgRate) * pred[i] - kInferenceCfgRate * pred[branch + i]; + x_batched[i] += dt * guided; + } + std::copy(x_batched.begin(), x_batched.begin() + static_cast(branch), x_batched.begin() + static_cast(branch)); + } + engine::debug::timing_log_scalar("cosyvoice3.flow.dit_ms", engine::debug::elapsed_ms(dit_start)); + + CosyVoice3FlowOutput out; + out.frames = target_frames; + out.mel.resize(static_cast(target_frames * c.flow_mel_channels)); + for (int64_t frame = 0; frame < target_frames; ++frame) { + for (int64_t ch = 0; ch < c.flow_mel_channels; ++ch) { + out.mel[static_cast(frame * c.flow_mel_channels + ch)] = + x_batched[static_cast(ch * total_frames + request.prompt_mel_frames + frame)]; + } + } + engine::debug::trace_log_f32("cosyvoice3.flow.output_mel", {target_frames, c.flow_mel_channels}, out.mel); + engine::debug::timing_log_scalar("cosyvoice3.flow.total_ms", engine::debug::elapsed_ms(start)); + engine::debug::trace_log_scalar("cosyvoice3.flow.total_frames", static_cast(total_frames)); + engine::debug::trace_log_scalar("cosyvoice3.flow.target_frames", static_cast(target_frames)); + return out; + } + + void release_graphs() { + condition_.release_graph(); + dit_.release_graph(); + } + +private: + std::shared_ptr assets_; + core::ExecutionContext & execution_; + std::shared_ptr weights_; + ConditionGraph condition_; + DiTGraph dit_; + size_t graph_arena_bytes_ = 0; +}; + +CosyVoice3FlowRuntime::CosyVoice3FlowRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type) + : impl_(std::make_unique( + std::move(assets), + execution, + graph_arena_bytes, + weight_context_bytes, + storage_type)) {} + +CosyVoice3FlowRuntime::~CosyVoice3FlowRuntime() = default; + +CosyVoice3FlowOutput CosyVoice3FlowRuntime::generate(const CosyVoice3FlowRequest & request) { + return impl_->generate(request); +} + +void CosyVoice3FlowRuntime::release_graphs() { + impl_->release_graphs(); +} + +} // namespace engine::models::cosyvoice3 diff --git a/src/models/cosyvoice3/frontend.cpp b/src/models/cosyvoice3/frontend.cpp new file mode 100644 index 000000000..3e878f4bd --- /dev/null +++ b/src/models/cosyvoice3/frontend.cpp @@ -0,0 +1,285 @@ +#include "engine/models/cosyvoice3/frontend.h" + +#include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/dsp.h" +#include "engine/framework/audio/kaldi_fbank.h" +#include "engine/framework/audio/resampling.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/speech_encoders/campplus_encoder.h" +#include "engine/framework/modules/speech_encoders/s3_tokenizer.h" +#include "engine/framework/runtime/cache_slots.h" +#include "engine/framework/debug/trace.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::cosyvoice3 { +namespace { + +uint64_t hash_audio(const engine::runtime::AudioBuffer & audio) { + uint64_t hash = 1469598103934665603ull; + auto mix = [&hash](uint64_t value) { + hash ^= value; + hash *= 1099511628211ull; + }; + mix(static_cast(audio.sample_rate)); + mix(static_cast(audio.channels)); + mix(static_cast(audio.samples.size())); + for (float sample : audio.samples) { + uint32_t bits = 0; + static_assert(sizeof(bits) == sizeof(sample)); + std::memcpy(&bits, &sample, sizeof(bits)); + mix(bits); + } + return hash; +} + +std::vector mono_resampled(const engine::runtime::AudioBuffer & audio, int sample_rate) { + if (audio.sample_rate <= 0 || audio.channels <= 0 || audio.samples.empty()) { + throw std::runtime_error("CosyVoice3 requires non-empty reference audio"); + } + auto mono = engine::audio::mixdown_interleaved_to_mono_average(audio.samples, audio.channels); + if (audio.sample_rate == sample_rate) { + return mono; + } + engine::audio::TorchaudioSincHannResampleOptions options; + options.kernel_mode = engine::audio::TorchaudioSincHannKernelMode::Float64ComputationStoredAsFloat64; + return engine::audio::resample_mono_torchaudio_sinc_hann(mono, audio.sample_rate, sample_rate, options); +} + +engine::runtime::AudioBuffer mono_buffer(const engine::runtime::AudioBuffer & audio, int sample_rate) { + engine::runtime::AudioBuffer out; + out.sample_rate = sample_rate; + out.channels = 1; + out.samples = mono_resampled(audio, sample_rate); + return out; +} + +int64_t reflect_index(int64_t index, int64_t length) { + while (index < 0 || index >= length) { + if (index < 0) { + index = -index; + } else { + index = 2 * length - index - 2; + } + } + return index; +} + +void compute_prompt_mel( + const engine::runtime::AudioBuffer & audio, + std::vector & values, + int64_t & frames) { + constexpr int64_t kSampleRate = 24000; + constexpr int64_t kNfft = 1920; + constexpr int64_t kHop = 480; + constexpr int64_t kWin = 1920; + constexpr int64_t kMels = 80; + constexpr float kEps = 1.0e-9F; + constexpr float kLogClamp = 1.0e-5F; + const auto mono = mono_resampled(audio, kSampleRate); + if (mono.size() < 2) { + throw std::runtime_error("CosyVoice3 reference audio is too short for prompt mel"); + } + const int64_t pad = (kNfft - kHop) / 2; + const int64_t padded_samples = static_cast(mono.size()) + 2 * pad; + std::vector padded(static_cast(padded_samples), 0.0F); + for (int64_t index = 0; index < padded_samples; ++index) { + padded[static_cast(index)] = + mono[static_cast(reflect_index(index - pad, static_cast(mono.size())))]; + } + const engine::audio::STFTConfig stft_config{ + kNfft, + kHop, + kWin, + false, + engine::audio::STFTPadMode::Constant, + engine::audio::STFTFamily::Kokoro, + }; + const auto & window = engine::audio::get_cached_stft_window(stft_config); + const auto magnitude = engine::audio::STFT().compute_magnitude( + padded, + window, + 1, + padded_samples, + stft_config); + frames = magnitude.shape[2]; + const int64_t freq_bins = kNfft / 2 + 1; + const auto filterbank = engine::audio::MelFilterbank().build( + engine::audio::MelFilterbankConfig{kSampleRate, kNfft, kMels, 0.0F, 0.0F, true}); + values.assign(static_cast(frames * kMels), 0.0F); +#ifdef _OPENMP +#pragma omp parallel for if (frames > 8) +#endif + for (int64_t frame = 0; frame < frames; ++frame) { + for (int64_t mel = 0; mel < kMels; ++mel) { + double sum = 0.0; + for (int64_t bin = 0; bin < freq_bins; ++bin) { + const float mag = magnitude.values[static_cast(bin * frames + frame)]; + const float stabilized = std::sqrt(mag * mag + kEps); + sum += static_cast(filterbank.values[static_cast(mel * freq_bins + bin)]) * + static_cast(stabilized); + } + values[static_cast(frame * kMels + mel)] = + static_cast(std::log(std::max(sum, static_cast(kLogClamp)))); + } + } +} + +std::vector compute_campplus_fbank(const engine::runtime::AudioBuffer & audio, int64_t & frames) { + const auto mono = mono_resampled(audio, 16000); + engine::audio::KaldiFbankOptions options; + options.sample_rate = 16000; + options.num_mels = 80; + options.frame_length_ms = 25.0F; + options.frame_shift_ms = 10.0F; + options.window_type = engine::audio::KaldiFbankWindowType::Povey; + options.lfr_m = 1; + options.lfr_n = 1; + options.preemphasis = 0.97F; + options.low_frequency = 20.0F; + options.high_frequency = 0.0F; + options.remove_dc_offset = true; + options.upscale_samples = false; + options.apply_cmvn = false; + auto fbank = engine::audio::extract_kaldi_fbank(mono, options); + if (fbank.frames <= 0 || fbank.feature_dim != 80) { + throw std::runtime_error("CosyVoice3 reference audio is too short for CAMPPlus"); + } + frames = fbank.frames; + for (int64_t dim = 0; dim < fbank.feature_dim; ++dim) { + float mean = 0.0F; + for (int64_t frame = 0; frame < fbank.frames; ++frame) { + mean += fbank.values[static_cast(frame * fbank.feature_dim + dim)]; + } + mean /= static_cast(fbank.frames); + for (int64_t frame = 0; frame < fbank.frames; ++frame) { + fbank.values[static_cast(frame * fbank.feature_dim + dim)] -= mean; + } + } + return std::move(fbank.values); +} + +} // namespace + +class CosyVoice3Frontend::Impl { +public: + Impl( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type, + size_t reference_cache_slots) + : assets_(std::move(assets)), + reference_cache_(reference_cache_slots) { + (void) graph_arena_bytes; + (void) weight_context_bytes; + if (assets_ == nullptr) { + throw std::runtime_error("CosyVoice3 frontend requires assets"); + } + engine::modules::S3TokenizerConfig speech_tokenizer_config; + speech_tokenizer_config.weight_storage_type = storage_type; + speech_tokenizer_ = std::make_unique( + engine::modules::S3TokenizerComponent::load_from_source( + assets_->speech_tokenizer_weights, + execution, + speech_tokenizer_config)); + engine::modules::CampplusEncoderConfig campplus_config; + campplus_config.feat_dim = 80; + campplus_config.embedding_size = assets_->config.speaker_dim; + campplus_config.weight_storage_type = storage_type; + campplus_config.weight_layout = engine::modules::CampplusEncoderWeightLayout::Fused; + campplus_config.normalize_partial_segment_by_full_length = true; + campplus_ = engine::modules::CampplusEncoderComponent::load_from_tensor_source( + assets_->campplus_weights, + execution.config(), + campplus_config); + } + + const CosyVoice3ReferenceFeatures & prepare_reference(const engine::runtime::AudioBuffer & audio) { + const auto key = hash_audio(audio); + const auto * cached = reference_cache_.find(key); + if (cached != nullptr) { + return *cached; + } + + const auto started = std::chrono::steady_clock::now(); + CosyVoice3ReferenceFeatures features; + auto audio16 = mono_buffer(audio, 16000); + auto token_output = speech_tokenizer_->tokenize(audio16, std::nullopt); + features.speech_tokens = std::move(token_output.tokens); + features.speech_token_count = token_output.token_count; + compute_prompt_mel(audio, features.prompt_mel, features.prompt_mel_frames); + int64_t fbank_frames = 0; + auto fbank = compute_campplus_fbank(audio, fbank_frames); + engine::debug::trace_log_f32("cosyvoice3.frontend.campplus_fbank", {fbank_frames, 80}, fbank); + auto speaker = campplus_.embed_from_features(fbank, fbank_frames, 80); + features.speaker_embedding = std::move(speaker.embedding); + engine::debug::trace_log_f32("cosyvoice3.frontend.speaker_embedding", {assets_->config.speaker_dim}, features.speaker_embedding); + if (static_cast(features.speaker_embedding.size()) != assets_->config.speaker_dim) { + throw std::runtime_error("CosyVoice3 CAMPPlus speaker embedding size mismatch"); + } + engine::debug::timing_log_scalar("cosyvoice3.frontend.reference_ms", engine::debug::elapsed_ms(started)); + + if (reference_cache_.capacity() == 0) { + uncached_ = std::move(features); + return uncached_; + } + reference_cache_.put(key, std::move(features)); + const auto * inserted = reference_cache_.find(key); + if (inserted == nullptr) { + throw std::runtime_error("CosyVoice3 reference cache insert failed"); + } + return *inserted; + } + + void release_graphs() { + if (speech_tokenizer_ != nullptr) { + speech_tokenizer_->release_runtime_cache(); + } + campplus_.release_runtime_graph(); + } + +private: + std::shared_ptr assets_; + std::unique_ptr speech_tokenizer_; + engine::modules::CampplusEncoderComponent campplus_; + engine::runtime::CacheSlots reference_cache_; + CosyVoice3ReferenceFeatures uncached_; +}; + +CosyVoice3Frontend::CosyVoice3Frontend( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type, + size_t reference_cache_slots) + : impl_(std::make_unique( + std::move(assets), + execution, + graph_arena_bytes, + weight_context_bytes, + storage_type, + reference_cache_slots)) {} + +CosyVoice3Frontend::~CosyVoice3Frontend() = default; + +const CosyVoice3ReferenceFeatures & CosyVoice3Frontend::prepare_reference(const engine::runtime::AudioBuffer & audio) { + return impl_->prepare_reference(audio); +} + +void CosyVoice3Frontend::release_graphs() { + impl_->release_graphs(); +} + +} // namespace engine::models::cosyvoice3 diff --git a/src/models/cosyvoice3/hift.cpp b/src/models/cosyvoice3/hift.cpp new file mode 100644 index 000000000..6043c68b6 --- /dev/null +++ b/src/models/cosyvoice3/hift.cpp @@ -0,0 +1,113 @@ +#include "engine/models/cosyvoice3/hift.h" + +#include "engine/framework/modules/vocoders/hift_vocoder.h" + +#include +#include + +namespace engine::models::cosyvoice3 { +namespace { + +engine::modules::HiftVocoderConfig make_hift_config(engine::assets::TensorStorageType 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.conv_pre_kernel_size = 5; + config.causal_convolutions = true; + config.f0_num_class = 1; + config.f0_in_channels = 80; + config.f0_cond_channels = 512; + config.f0_condnet_kernel_sizes = {4, 3, 3, 3, 3}; + config.weight_storage_type = storage_type; + config.weight_layout = engine::modules::HiftVocoderWeightLayout::TorchParametrizedWeightNorm; + config.upsample_mode = engine::modules::HiftVocoderUpsampleMode::CausalConv1dNearest; + config.source_mode = engine::modules::HiftVocoderSourceMode::CausalSineGen2; + return config; +} + +} // namespace + +class CosyVoice3HiftRuntime::Impl { +public: + Impl( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + engine::assets::TensorStorageType storage_type) + : assets_(std::move(assets)) { + if (assets_ == nullptr) { + throw std::runtime_error("CosyVoice3 HiFT runtime requires assets"); + } + component_ = engine::modules::HiftVocoderComponent::load_from_tensor_source( + assets_->hift_weights, + execution.config(), + make_hift_config(storage_type)); + } + + engine::runtime::AudioBuffer synthesize( + const std::vector & mel, + int64_t frames, + uint64_t seed, + const std::vector * source_random_values) { + if (frames <= 0 || static_cast(mel.size()) != frames * 80) { + throw std::runtime_error("CosyVoice3 HiFT mel shape mismatch"); + } + std::vector channel_major(mel.size()); + for (int64_t frame = 0; frame < frames; ++frame) { + for (int64_t channel = 0; channel < 80; ++channel) { + channel_major[static_cast(channel * frames + frame)] = + mel[static_cast(frame * 80 + channel)]; + } + } + auto out = component_.synthesize(channel_major, frames, seed, 0, source_random_values); + engine::runtime::AudioBuffer audio; + audio.sample_rate = static_cast(out.sample_rate); + audio.channels = 1; + audio.samples = std::move(out.waveform); + return audio; + } + + void release_graphs() { + component_.release_runtime_cache(); + } + +private: + std::shared_ptr assets_; + engine::modules::HiftVocoderComponent component_; +}; + +CosyVoice3HiftRuntime::CosyVoice3HiftRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + engine::assets::TensorStorageType storage_type) + : impl_(std::make_unique(std::move(assets), execution, storage_type)) {} + +CosyVoice3HiftRuntime::~CosyVoice3HiftRuntime() = default; + +engine::runtime::AudioBuffer CosyVoice3HiftRuntime::synthesize( + const std::vector & mel, + int64_t frames, + uint64_t seed, + const std::vector * source_random_values) { + return impl_->synthesize(mel, frames, seed, source_random_values); +} + +void CosyVoice3HiftRuntime::release_graphs() { + impl_->release_graphs(); +} + +} // namespace engine::models::cosyvoice3 diff --git a/src/models/cosyvoice3/session.cpp b/src/models/cosyvoice3/session.cpp new file mode 100644 index 000000000..b1e2affbc --- /dev/null +++ b/src/models/cosyvoice3/session.cpp @@ -0,0 +1,329 @@ +#include "engine/models/cosyvoice3/session.h" + +#include "engine/framework/debug/profiler.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/runtime/spec_backed_model.h" +#include "engine/framework/text/chunking.h" +#include "engine/models/cosyvoice3/ar.h" +#include "engine/models/cosyvoice3/flow.h" +#include "engine/models/cosyvoice3/frontend.h" +#include "engine/models/cosyvoice3/hift.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::cosyvoice3 { +namespace { + +constexpr const char * kFamily = "cosyvoice3"; +constexpr const char * kModelName = "CosyVoice3"; +constexpr int64_t kDefaultTextChunkSize = 600; +constexpr size_t kDefaultReferenceCacheSlots = 4; + +std::shared_ptr require_assets(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("CosyVoice3 session requires assets"); + } + return assets; +} + +std::shared_ptr require_contract( + std::shared_ptr contract) { + if (contract == nullptr) { + throw std::runtime_error("CosyVoice3 session requires a model contract"); + } + return contract; +} + +std::string template_name(const runtime::TaskRequest & request) { + return runtime::find_option(request.options, {"template_name"}).value_or("zero_shot"); +} + +std::string reference_text(const runtime::TaskRequest & request) { + return runtime::find_option(request.options, {"reference_text"}).value_or(""); +} + +std::string instruction_text(const runtime::TaskRequest & request) { + return runtime::find_option(request.options, {"instruction"}).value_or(""); +} + +std::unordered_map validation_options( + const std::unordered_map & options) { + auto out = options; + out.erase("teacher_force_tokens"); + out.erase("teacher_force_source_random"); + return out; +} + +std::vector read_teacher_force_tokens(const std::string & path) { + std::ifstream in(path); + if (!in) { + throw std::runtime_error("CosyVoice3 teacher_force_tokens file cannot be opened: " + path); + } + std::vector out; + std::string token; + char ch = '\0'; + while (in.get(ch)) { + if (std::isdigit(static_cast(ch)) || ch == '-') { + token.push_back(ch); + } else if (!token.empty()) { + out.push_back(static_cast(std::stoi(token))); + token.clear(); + } + } + if (!token.empty()) { + out.push_back(static_cast(std::stoi(token))); + } + if (out.empty()) { + throw std::runtime_error("CosyVoice3 teacher_force_tokens file contains no tokens: " + path); + } + return out; +} + +std::vector read_float_list(const std::string & path, const char * label) { + std::ifstream in(path); + if (!in) { + throw std::runtime_error(std::string("CosyVoice3 ") + label + " file cannot be opened: " + path); + } + std::vector out; + std::string token; + char ch = '\0'; + while (in.get(ch)) { + if (std::isdigit(static_cast(ch)) || ch == '-' || ch == '+' || ch == '.' || ch == 'e' || ch == 'E') { + token.push_back(ch); + } else if (!token.empty()) { + out.push_back(std::stof(token)); + token.clear(); + } + } + if (!token.empty()) { + out.push_back(std::stof(token)); + } + if (out.empty()) { + throw std::runtime_error(std::string("CosyVoice3 ") + label + " file contains no values: " + path); + } + return out; +} + +std::vector split_request(const runtime::TaskRequest & request) { + const int64_t text_chunk_size = + engine::text::parse_text_chunk_size_override(request.options).value_or(kDefaultTextChunkSize); + if (text_chunk_size <= 0) { + throw std::runtime_error("CosyVoice3 text_chunk_size must be positive"); + } + const auto text_chunk_mode = + engine::text::parse_text_chunk_mode_override(request.options).value_or(engine::text::TextChunkMode::Default); + return runtime::chunk_text_request(request, text_chunk_size, text_chunk_mode); +} + +std::unique_ptr create_cosyvoice3_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options, + std::shared_ptr assets, + std::shared_ptr contract) { + return std::make_unique( + task, + options, + std::move(assets), + std::move(contract)); +} + +} // namespace + +CosyVoice3Session::CosyVoice3Session( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract) + : runtime::RuntimeSessionBase(options), + task_(task), + assets_(require_assets(std::move(assets))), + contract_(require_contract(std::move(contract))), + tokenizer_(std::make_unique(assets_)) { + runtime::validate_spec_backed_session_options(options, *contract_, kFamily, kModelName); + if (task_.task != runtime::VoiceTaskKind::Tts && task_.task != runtime::VoiceTaskKind::VoiceCloning) { + throw std::runtime_error("CosyVoice3 supports tts and clone tasks"); + } + if (task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("CosyVoice3 supports offline sessions"); + } + + using T = engine::assets::TensorStorageType; + const auto storage_type = runtime::parse_tensor_storage_option( + options.options, + "cosyvoice3.weight_type", + T::Native, + {T::Native, T::F32, T::F16, T::BF16, T::Q8_0, T::Q4_0, T::Q4_K}); + const auto graph_arena_bytes = runtime::parse_size_mb_option( + options.options, + {"cosyvoice3.graph_arena_mb"}, + 1024ull * 1024ull * 1024ull); + const auto weight_context_bytes = runtime::parse_size_mb_option( + options.options, + {"cosyvoice3.weight_context_mb"}, + 2048ull * 1024ull * 1024ull); + if (const auto mem_saver = runtime::find_option(options.options, {"cosyvoice3.mem_saver"})) { + mem_saver_ = runtime::parse_bool_option(*mem_saver, "cosyvoice3.mem_saver"); + } + const int64_t reference_cache_slots = runtime::parse_i64_option( + options.options, + {"cosyvoice3.reference_cache_slots"}) + .value_or(static_cast(kDefaultReferenceCacheSlots)); + if (reference_cache_slots < 0) { + throw std::runtime_error("cosyvoice3.reference_cache_slots must be non-negative"); + } + + frontend_ = std::make_unique( + assets_, + execution_context(), + graph_arena_bytes, + weight_context_bytes, + storage_type, + static_cast(reference_cache_slots)); + ar_ = std::make_unique( + assets_, + execution_context(), + graph_arena_bytes, + weight_context_bytes, + storage_type); + flow_ = std::make_unique( + assets_, + execution_context(), + graph_arena_bytes, + weight_context_bytes, + storage_type); + hift_ = std::make_unique( + assets_, + execution_context(), + storage_type); +} + +CosyVoice3Session::~CosyVoice3Session() = default; + +std::string CosyVoice3Session::family() const { + return kFamily; +} + +runtime::VoiceTaskKind CosyVoice3Session::task_kind() const { + return task_.task; +} + +runtime::RunMode CosyVoice3Session::run_mode() const { + return task_.mode; +} + +void CosyVoice3Session::prepare(const runtime::SessionPreparationRequest & request) { + runtime::validate_spec_backed_request_options(validation_options(request.options), *contract_, kModelName); + mark_prepared(); +} + +runtime::TaskResult CosyVoice3Session::run(const runtime::TaskRequest & request) { + const auto wall_start = std::chrono::steady_clock::now(); + runtime::validate_spec_backed_request_options(validation_options(request.options), *contract_, kModelName); + require_prepared("CosyVoice3 run"); + if (!request.text_input.has_value() || request.text_input->text.empty()) { + throw std::runtime_error("CosyVoice3 requires text input"); + } + if (!request.voice.has_value() || + !request.voice->speaker.has_value() || + !request.voice->speaker->audio.has_value()) { + throw std::runtime_error("CosyVoice3 requires reference audio"); + } + + runtime::AudioBuffer merged_audio; + auto chunks = split_request(request); + for (size_t index = 0; index < chunks.size(); ++index) { + const auto & chunk = chunks[index]; + const auto mode = template_name(chunk); + CosyVoice3TextTokens text_tokens; + if (mode == "zero_shot") { + text_tokens = tokenizer_->encode_zero_shot(chunk.text_input->text, reference_text(chunk)); + } else if (mode == "cross_lingual") { + text_tokens = tokenizer_->encode_cross_lingual(chunk.text_input->text); + } else if (mode == "instruct") { + text_tokens = tokenizer_->encode_instruct(chunk.text_input->text, instruction_text(chunk)); + } else { + throw std::runtime_error("unknown CosyVoice3 template_name: " + mode); + } + const auto & reference = frontend_->prepare_reference(*chunk.voice->speaker->audio); + if (mem_saver_) { + frontend_->release_graphs(); + } + CosyVoice3ArRequest ar_request; + ar_request.prompt_text_tokens = std::move(text_tokens.prompt); + ar_request.target_text_tokens = std::move(text_tokens.target); + if (mode == "zero_shot") { + ar_request.prompt_speech_tokens = reference.speech_tokens; + } + ar_request.seed = runtime::parse_u32_option(chunk.options, {"seed"}).value_or(ar_request.seed); + ar_request.top_k = runtime::parse_positive_i64_option(chunk.options, {"top_k"}, ar_request.top_k); + ar_request.min_tokens = runtime::parse_i64_option(chunk.options, {"min_tokens"}).value_or(ar_request.min_tokens); + ar_request.max_tokens = runtime::parse_i64_option(chunk.options, {"max_tokens"}).value_or(ar_request.max_tokens); + if (index > 0) { + ar_request.seed += static_cast(index); + } + CosyVoice3ArOutput ar_output; + if (const auto teacher_force_path = runtime::find_option(chunk.options, {"teacher_force_tokens"})) { + ar_output.speech_tokens = read_teacher_force_tokens(*teacher_force_path); + } else { + ar_output = ar_->generate(ar_request); + } + if (mem_saver_) { + ar_->release_graphs(); + } + + CosyVoice3FlowRequest flow_request; + flow_request.speech_tokens = std::move(ar_output.speech_tokens); + flow_request.prompt_speech_tokens = reference.speech_tokens; + flow_request.prompt_mel = reference.prompt_mel; + flow_request.prompt_mel_frames = reference.prompt_mel_frames; + flow_request.speaker_embedding = reference.speaker_embedding; + flow_request.seed = ar_request.seed; + flow_request.num_inference_steps = runtime::parse_positive_i64_option( + chunk.options, + {"num_inference_steps"}, + flow_request.num_inference_steps); + CosyVoice3FlowOutput flow_output = flow_->generate(flow_request); + if (mem_saver_) { + flow_->release_graphs(); + } + std::vector source_random_values; + const std::vector * source_random_ptr = nullptr; + if (const auto random_path = runtime::find_option(chunk.options, {"teacher_force_source_random"})) { + source_random_values = read_float_list(*random_path, "teacher_force_source_random"); + source_random_ptr = &source_random_values; + } + runtime::append_audio_buffer( + merged_audio, + hift_->synthesize(flow_output.mel, flow_output.frames, flow_request.seed, source_random_ptr)); + if (mem_saver_) { + hift_->release_graphs(); + } + } + + if (mem_saver_) { + frontend_->release_graphs(); + ar_->release_graphs(); + flow_->release_graphs(); + hift_->release_graphs(); + } + runtime::TaskResult result; + result.audio_output = std::move(merged_audio); + engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start)); + return result; +} + +std::shared_ptr make_cosyvoice3_loader() { + runtime::SpecBackedVoiceModelConfig config; + config.family = kFamily; + config.load_assets = load_cosyvoice3_assets; + config.create_session = create_cosyvoice3_session; + return runtime::make_spec_backed_voice_loader(std::move(config)); +} + +} // namespace engine::models::cosyvoice3 diff --git a/src/models/cosyvoice3/tokenizer_text.cpp b/src/models/cosyvoice3/tokenizer_text.cpp new file mode 100644 index 000000000..9e3823c69 --- /dev/null +++ b/src/models/cosyvoice3/tokenizer_text.cpp @@ -0,0 +1,363 @@ +#include "engine/models/cosyvoice3/tokenizer_text.h" + +#include "engine/framework/tokenizers/llama_bpe.h" + +#include +#include + +namespace engine::models::cosyvoice3 { +namespace { + +constexpr int32_t kEndOfPromptId = 151646; + +std::vector cosyvoice3_special_tokens() { + std::vector tokens; + tokens.reserve(280); + tokens.emplace_back("<|im_start|>", 151644); + tokens.emplace_back("<|im_end|>", 151645); + tokens.emplace_back("<|endofprompt|>", 151646); + tokens.emplace_back("[breath]", 151647); + tokens.emplace_back("", 151648); + tokens.emplace_back("", 151649); + tokens.emplace_back("[noise]", 151650); + tokens.emplace_back("[laughter]", 151651); + tokens.emplace_back("[cough]", 151652); + tokens.emplace_back("[clucking]", 151653); + tokens.emplace_back("[accent]", 151654); + tokens.emplace_back("[quick_breath]", 151655); + tokens.emplace_back("", 151656); + tokens.emplace_back("", 151657); + tokens.emplace_back("[hissing]", 151658); + tokens.emplace_back("[sigh]", 151659); + tokens.emplace_back("[vocalized-noise]", 151660); + tokens.emplace_back("[lipsmack]", 151661); + tokens.emplace_back("[mn]", 151662); + tokens.emplace_back("<|endofsystem|>", 151663); + tokens.emplace_back("[AA]", 151664); + tokens.emplace_back("[AA0]", 151665); + tokens.emplace_back("[AA1]", 151666); + tokens.emplace_back("[AA2]", 151667); + tokens.emplace_back("[AE]", 151668); + tokens.emplace_back("[AE0]", 151669); + tokens.emplace_back("[AE1]", 151670); + tokens.emplace_back("[AE2]", 151671); + tokens.emplace_back("[AH]", 151672); + tokens.emplace_back("[AH0]", 151673); + tokens.emplace_back("[AH1]", 151674); + tokens.emplace_back("[AH2]", 151675); + tokens.emplace_back("[AO]", 151676); + tokens.emplace_back("[AO0]", 151677); + tokens.emplace_back("[AO1]", 151678); + tokens.emplace_back("[AO2]", 151679); + tokens.emplace_back("[AW]", 151680); + tokens.emplace_back("[AW0]", 151681); + tokens.emplace_back("[AW1]", 151682); + tokens.emplace_back("[AW2]", 151683); + tokens.emplace_back("[AY]", 151684); + tokens.emplace_back("[AY0]", 151685); + tokens.emplace_back("[AY1]", 151686); + tokens.emplace_back("[AY2]", 151687); + tokens.emplace_back("[B]", 151688); + tokens.emplace_back("[CH]", 151689); + tokens.emplace_back("[D]", 151690); + tokens.emplace_back("[DH]", 151691); + tokens.emplace_back("[EH]", 151692); + tokens.emplace_back("[EH0]", 151693); + tokens.emplace_back("[EH1]", 151694); + tokens.emplace_back("[EH2]", 151695); + tokens.emplace_back("[ER]", 151696); + tokens.emplace_back("[ER0]", 151697); + tokens.emplace_back("[ER1]", 151698); + tokens.emplace_back("[ER2]", 151699); + tokens.emplace_back("[EY]", 151700); + tokens.emplace_back("[EY0]", 151701); + tokens.emplace_back("[EY1]", 151702); + tokens.emplace_back("[EY2]", 151703); + tokens.emplace_back("[F]", 151704); + tokens.emplace_back("[G]", 151705); + tokens.emplace_back("[HH]", 151706); + tokens.emplace_back("[IH]", 151707); + tokens.emplace_back("[IH0]", 151708); + tokens.emplace_back("[IH1]", 151709); + tokens.emplace_back("[IH2]", 151710); + tokens.emplace_back("[IY]", 151711); + tokens.emplace_back("[IY0]", 151712); + tokens.emplace_back("[IY1]", 151713); + tokens.emplace_back("[IY2]", 151714); + tokens.emplace_back("[JH]", 151715); + tokens.emplace_back("[K]", 151716); + tokens.emplace_back("[L]", 151717); + tokens.emplace_back("[M]", 151718); + tokens.emplace_back("[N]", 151719); + tokens.emplace_back("[NG]", 151720); + tokens.emplace_back("[OW]", 151721); + tokens.emplace_back("[OW0]", 151722); + tokens.emplace_back("[OW1]", 151723); + tokens.emplace_back("[OW2]", 151724); + tokens.emplace_back("[OY]", 151725); + tokens.emplace_back("[OY0]", 151726); + tokens.emplace_back("[OY1]", 151727); + tokens.emplace_back("[OY2]", 151728); + tokens.emplace_back("[P]", 151729); + tokens.emplace_back("[R]", 151730); + tokens.emplace_back("[S]", 151731); + tokens.emplace_back("[SH]", 151732); + tokens.emplace_back("[T]", 151733); + tokens.emplace_back("[TH]", 151734); + tokens.emplace_back("[UH]", 151735); + tokens.emplace_back("[UH0]", 151736); + tokens.emplace_back("[UH1]", 151737); + tokens.emplace_back("[UH2]", 151738); + tokens.emplace_back("[UW]", 151739); + tokens.emplace_back("[UW0]", 151740); + tokens.emplace_back("[UW1]", 151741); + tokens.emplace_back("[UW2]", 151742); + tokens.emplace_back("[V]", 151743); + tokens.emplace_back("[W]", 151744); + tokens.emplace_back("[Y]", 151745); + tokens.emplace_back("[Z]", 151746); + tokens.emplace_back("[ZH]", 151747); + tokens.emplace_back("[a]", 151748); + tokens.emplace_back("[ai]", 151749); + tokens.emplace_back("[an]", 151750); + tokens.emplace_back("[ang]", 151751); + tokens.emplace_back("[ao]", 151752); + tokens.emplace_back("[b]", 151753); + tokens.emplace_back("[c]", 151754); + tokens.emplace_back("[ch]", 151755); + tokens.emplace_back("[d]", 151756); + tokens.emplace_back("[e]", 151757); + tokens.emplace_back("[ei]", 151758); + tokens.emplace_back("[en]", 151759); + tokens.emplace_back("[eng]", 151760); + tokens.emplace_back("[f]", 151761); + tokens.emplace_back("[g]", 151762); + tokens.emplace_back("[h]", 151763); + tokens.emplace_back("[i]", 151764); + tokens.emplace_back("[ian]", 151765); + tokens.emplace_back("[in]", 151766); + tokens.emplace_back("[ing]", 151767); + tokens.emplace_back("[iu]", 151768); + tokens.emplace_back("[ià]", 151769); + tokens.emplace_back("[iàn]", 151770); + tokens.emplace_back("[iàng]", 151771); + tokens.emplace_back("[iào]", 151772); + tokens.emplace_back("[iá]", 151773); + tokens.emplace_back("[ián]", 151774); + tokens.emplace_back("[iáng]", 151775); + tokens.emplace_back("[iáo]", 151776); + tokens.emplace_back("[iè]", 151777); + tokens.emplace_back("[ié]", 151778); + tokens.emplace_back("[iòng]", 151779); + tokens.emplace_back("[ióng]", 151780); + tokens.emplace_back("[iù]", 151781); + tokens.emplace_back("[iú]", 151782); + tokens.emplace_back("[iā]", 151783); + tokens.emplace_back("[iān]", 151784); + tokens.emplace_back("[iāng]", 151785); + tokens.emplace_back("[iāo]", 151786); + tokens.emplace_back("[iē]", 151787); + tokens.emplace_back("[iě]", 151788); + tokens.emplace_back("[iōng]", 151789); + tokens.emplace_back("[iū]", 151790); + tokens.emplace_back("[iǎ]", 151791); + tokens.emplace_back("[iǎn]", 151792); + tokens.emplace_back("[iǎng]", 151793); + tokens.emplace_back("[iǎo]", 151794); + tokens.emplace_back("[iǒng]", 151795); + tokens.emplace_back("[iǔ]", 151796); + tokens.emplace_back("[j]", 151797); + tokens.emplace_back("[k]", 151798); + tokens.emplace_back("[l]", 151799); + tokens.emplace_back("[m]", 151800); + tokens.emplace_back("[n]", 151801); + tokens.emplace_back("[o]", 151802); + tokens.emplace_back("[ong]", 151803); + tokens.emplace_back("[ou]", 151804); + tokens.emplace_back("[p]", 151805); + tokens.emplace_back("[q]", 151806); + tokens.emplace_back("[r]", 151807); + tokens.emplace_back("[s]", 151808); + tokens.emplace_back("[sh]", 151809); + tokens.emplace_back("[t]", 151810); + tokens.emplace_back("[u]", 151811); + tokens.emplace_back("[uang]", 151812); + tokens.emplace_back("[ue]", 151813); + tokens.emplace_back("[un]", 151814); + tokens.emplace_back("[uo]", 151815); + tokens.emplace_back("[uà]", 151816); + tokens.emplace_back("[uài]", 151817); + tokens.emplace_back("[uàn]", 151818); + tokens.emplace_back("[uàng]", 151819); + tokens.emplace_back("[uá]", 151820); + tokens.emplace_back("[uái]", 151821); + tokens.emplace_back("[uán]", 151822); + tokens.emplace_back("[uáng]", 151823); + tokens.emplace_back("[uè]", 151824); + tokens.emplace_back("[ué]", 151825); + tokens.emplace_back("[uì]", 151826); + tokens.emplace_back("[uí]", 151827); + tokens.emplace_back("[uò]", 151828); + tokens.emplace_back("[uó]", 151829); + tokens.emplace_back("[uā]", 151830); + tokens.emplace_back("[uāi]", 151831); + tokens.emplace_back("[uān]", 151832); + tokens.emplace_back("[uāng]", 151833); + tokens.emplace_back("[uē]", 151834); + tokens.emplace_back("[uě]", 151835); + tokens.emplace_back("[uī]", 151836); + tokens.emplace_back("[uō]", 151837); + tokens.emplace_back("[uǎ]", 151838); + tokens.emplace_back("[uǎi]", 151839); + tokens.emplace_back("[uǎn]", 151840); + tokens.emplace_back("[uǎng]", 151841); + tokens.emplace_back("[uǐ]", 151842); + tokens.emplace_back("[uǒ]", 151843); + tokens.emplace_back("[vè]", 151844); + tokens.emplace_back("[w]", 151845); + tokens.emplace_back("[x]", 151846); + tokens.emplace_back("[y]", 151847); + tokens.emplace_back("[z]", 151848); + tokens.emplace_back("[zh]", 151849); + tokens.emplace_back("[à]", 151850); + tokens.emplace_back("[ài]", 151851); + tokens.emplace_back("[àn]", 151852); + tokens.emplace_back("[àng]", 151853); + tokens.emplace_back("[ào]", 151854); + tokens.emplace_back("[á]", 151855); + tokens.emplace_back("[ái]", 151856); + tokens.emplace_back("[án]", 151857); + tokens.emplace_back("[áng]", 151858); + tokens.emplace_back("[áo]", 151859); + tokens.emplace_back("[è]", 151860); + tokens.emplace_back("[èi]", 151861); + tokens.emplace_back("[èn]", 151862); + tokens.emplace_back("[èng]", 151863); + tokens.emplace_back("[èr]", 151864); + tokens.emplace_back("[é]", 151865); + tokens.emplace_back("[éi]", 151866); + tokens.emplace_back("[én]", 151867); + tokens.emplace_back("[éng]", 151868); + tokens.emplace_back("[ér]", 151869); + tokens.emplace_back("[ì]", 151870); + tokens.emplace_back("[ìn]", 151871); + tokens.emplace_back("[ìng]", 151872); + tokens.emplace_back("[í]", 151873); + tokens.emplace_back("[ín]", 151874); + tokens.emplace_back("[íng]", 151875); + tokens.emplace_back("[ò]", 151876); + tokens.emplace_back("[òng]", 151877); + tokens.emplace_back("[òu]", 151878); + tokens.emplace_back("[ó]", 151879); + tokens.emplace_back("[óng]", 151880); + tokens.emplace_back("[óu]", 151881); + tokens.emplace_back("[ù]", 151882); + tokens.emplace_back("[ùn]", 151883); + tokens.emplace_back("[ú]", 151884); + tokens.emplace_back("[ún]", 151885); + tokens.emplace_back("[ā]", 151886); + tokens.emplace_back("[āi]", 151887); + tokens.emplace_back("[ān]", 151888); + tokens.emplace_back("[āng]", 151889); + tokens.emplace_back("[āo]", 151890); + tokens.emplace_back("[ē]", 151891); + tokens.emplace_back("[ēi]", 151892); + tokens.emplace_back("[ēn]", 151893); + tokens.emplace_back("[ēng]", 151894); + tokens.emplace_back("[ě]", 151895); + tokens.emplace_back("[ěi]", 151896); + tokens.emplace_back("[ěn]", 151897); + tokens.emplace_back("[ěng]", 151898); + tokens.emplace_back("[ěr]", 151899); + tokens.emplace_back("[ī]", 151900); + tokens.emplace_back("[īn]", 151901); + tokens.emplace_back("[īng]", 151902); + tokens.emplace_back("[ō]", 151903); + tokens.emplace_back("[ōng]", 151904); + tokens.emplace_back("[ōu]", 151905); + tokens.emplace_back("[ū]", 151906); + tokens.emplace_back("[ūn]", 151907); + tokens.emplace_back("[ǎ]", 151908); + tokens.emplace_back("[ǎi]", 151909); + tokens.emplace_back("[ǎn]", 151910); + tokens.emplace_back("[ǎng]", 151911); + tokens.emplace_back("[ǎo]", 151912); + tokens.emplace_back("[ǐ]", 151913); + tokens.emplace_back("[ǐn]", 151914); + tokens.emplace_back("[ǐng]", 151915); + tokens.emplace_back("[ǒ]", 151916); + tokens.emplace_back("[ǒng]", 151917); + tokens.emplace_back("[ǒu]", 151918); + tokens.emplace_back("[ǔ]", 151919); + tokens.emplace_back("[ǔn]", 151920); + tokens.emplace_back("[ǘ]", 151921); + tokens.emplace_back("[ǚ]", 151922); + tokens.emplace_back("[ǜ]", 151923); + return tokens; +} + +void require_end_of_prompt(const std::vector & tokens, const char * context) { + for (const int32_t token : tokens) { + if (token == kEndOfPromptId) { + return; + } + } + throw std::runtime_error(std::string("CosyVoice3 ") + context + " must contain <|endofprompt|>"); +} + +} // namespace + +class CosyVoice3TextTokenizer::Impl { +public: + explicit Impl(const CosyVoice3Assets & assets) { + engine::tokenizers::LlamaBpeTokenizerSpec spec; + spec.vocab_path = assets.resources.require_file("vocab_json"); + spec.merges_path = assets.resources.require_file("merges_txt"); + spec.tokenizer_config_path = assets.resources.require_file("tokenizer_config"); + spec.pre_type = engine::tokenizers::LlamaBpePreTokenizer::Qwen2; + spec.additional_special_tokens = cosyvoice3_special_tokens(); + tokenizer = engine::tokenizers::load_llama_bpe_tokenizer(spec); + } + + std::shared_ptr tokenizer; +}; + +CosyVoice3TextTokenizer::CosyVoice3TextTokenizer(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("CosyVoice3 tokenizer requires assets"); + } + impl_ = std::make_shared(*assets); +} + +CosyVoice3TextTokenizer::~CosyVoice3TextTokenizer() = default; + +CosyVoice3TextTokens CosyVoice3TextTokenizer::encode_zero_shot( + std::string_view text, + std::string_view prompt_text) const { + CosyVoice3TextTokens out; + out.prompt = impl_->tokenizer->encode(std::string(prompt_text), true); + out.target = impl_->tokenizer->encode(std::string(text), true); + std::vector joined = out.prompt; + joined.insert(joined.end(), out.target.begin(), out.target.end()); + require_end_of_prompt(joined, "zero-shot prompt/text"); + return out; +} + +CosyVoice3TextTokens CosyVoice3TextTokenizer::encode_cross_lingual(std::string_view text) const { + CosyVoice3TextTokens out; + out.target = impl_->tokenizer->encode(std::string(text), true); + require_end_of_prompt(out.target, "cross-lingual text"); + return out; +} + +CosyVoice3TextTokens CosyVoice3TextTokenizer::encode_instruct( + std::string_view text, + std::string_view instruction) const { + CosyVoice3TextTokens out; + out.prompt = impl_->tokenizer->encode(std::string(instruction), true); + out.target = impl_->tokenizer->encode(std::string(text), true); + require_end_of_prompt(out.prompt, "instruction"); + return out; +} + +} // namespace engine::models::cosyvoice3 diff --git a/tools/audiocpp_cli/audiocpp_cli_path_cases.json b/tools/audiocpp_cli/audiocpp_cli_path_cases.json index 20c2b69d3..f6ef6d967 100644 --- a/tools/audiocpp_cli/audiocpp_cli_path_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_path_cases.json @@ -3588,6 +3588,132 @@ } } ] + }, + { + "id": "cosyvoice3_official_paths", + "coverage": "CosyVoice3 official zero-shot, cross-lingual, instruction, and hotfix-pronunciation paths after a short same-session warmup", + "family": "cosyvoice3", + "model": "models/CosyVoice3-GGUF/cosyvoice3-orig.gguf", + "task": "clon", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "warmup_short_random", + "text": "你好。", + "reference_text": "You are a helpful assistant.<|endofprompt|>希望你以后能够做的比我还好呦。", + "voice_ref": "reference/CosyVoice/asset/zero_shot_prompt.wav", + "options": { + "template_name": "zero_shot", + "max_tokens": 80 + }, + "seed": 20260828 + }, + { + "id": "zero_shot_zh_official", + "text": "八百标兵奔北坡，北坡炮兵并排跑，炮兵怕把标兵碰，标兵怕碰炮兵炮。", + "reference_text": "You are a helpful assistant.<|endofprompt|>希望你以后能够做的比我还好呦。", + "voice_ref": "reference/CosyVoice/asset/zero_shot_prompt.wav", + "options": { + "template_name": "zero_shot" + }, + "seed": 1986 + }, + { + "id": "fine_grained_control_official", + "text": "You are a helpful assistant.<|endofprompt|>[breath]因为他们那一辈人[breath]在乡里面住的要习惯一点，[breath]邻居都很活络，[breath]嗯，都很熟悉。[breath]", + "voice_ref": "reference/CosyVoice/asset/zero_shot_prompt.wav", + "options": { + "template_name": "cross_lingual" + }, + "seed": 1986 + }, + { + "id": "instruct_cantonese_official", + "text": "好少咯，一般系放嗰啲国庆啊，中秋嗰啲可能会咯。", + "voice_ref": "reference/CosyVoice/asset/zero_shot_prompt.wav", + "options": { + "template_name": "instruct", + "instruction": "You are a helpful assistant. 请用广东话表达。<|endofprompt|>" + }, + "seed": 1986 + }, + { + "id": "hotfix_pronunciation_official", + "text": "高管也通过电话、短信、微信等方式对报道[j][ǐ]予好评。", + "reference_text": "You are a helpful assistant.<|endofprompt|>希望你以后能够做的比我还好呦。", + "voice_ref": "reference/CosyVoice/asset/zero_shot_prompt.wav", + "options": { + "template_name": "zero_shot" + }, + "seed": 1986 + } + ] + }, + { + "id": "breeze_tts_clone_official_paths", + "coverage": "BreezeTTS 2 official voice clone and voice direction paths using prompt audio, prompt transcript, T5Gemma text conditioning, Qwen acoustic generation, depth codebook decoding, and Mimi decode", + "family": "breeze_tts", + "model": "models/audio.cpp-gguf/Breeze-TTS-2-GGUF/breeze-tts-2-orig.gguf", + "task": "clon", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "voice_clone_en_official", + "text": "(sigh) It is good to hear your voice again after all this time.", + "reference_text": "Some call me nature. Others call me Mother Nature. I have been here for over four and a half billion years. Twenty two thousand five hundred times longer than you.", + "voice_ref": "assets/resources/b.wav", + "guidance_scale": 1.0, + "seed": 42 + }, + { + "id": "voice_direction_en_official", + "text": "(clears throat) We need to discuss what happened last night.", + "reference_text": "Some call me nature. Others call me Mother Nature. I have been here for over four and a half billion years. Twenty two thousand five hundred times longer than you.", + "voice_ref": "assets/resources/b.wav", + "options": { + "instruction": "Speak slowly with a restrained, serious tone." + }, + "guidance_scale": 4.0, + "seed": 45 + } + ] + }, + { + "id": "breeze_tts_design_official_paths", + "coverage": "BreezeTTS 2 official English and Chinese voice design paths using instruction conditioning, T5Gemma text conditioning, Qwen acoustic generation, depth codebook decoding, and Mimi decode", + "family": "breeze_tts", + "model": "models/audio.cpp-gguf/Breeze-TTS-2-GGUF/breeze-tts-2-orig.gguf", + "task": "tts", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "voice_design_en_official", + "text": "(sigh) Welcome aboard. Your journey begins now.", + "options": { + "instruction": "A warm, thoughtful young woman with a clear voice and a calm, reflective delivery." + }, + "guidance_scale": 4.0, + "seed": 43 + }, + { + "id": "voice_design_zh_official", + "text": "[笑] 欢迎来到今晚的故事时间，让我们一起开始吧。", + "options": { + "instruction": "一位温柔自信的年轻女性，声音清晰，语气亲切，表达轻快而富有感染力。" + }, + "guidance_scale": 4.0, + "seed": 44 + } + ] } ], "audit_gaps": [ From e58e547e4e9816efdcf06096f9cba784aabc44db Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:18:13 -0400 Subject: [PATCH 02/10] Normalize CosyVoice3 and Breeze GGUF package names --- docs/models/breeze_tts.md | 69 +++++ docs/models/cosyvoice3.md | 80 +++++ docs/tts.md | 2 + include/engine/models/cosyvoice3/hift.h | 3 +- model_specs/breeze_tts.json | 234 +++++++++++++++ model_specs/cosyvoice3.json | 273 ++++++++++++++++++ src/models/cosyvoice3/flow.cpp | 62 +--- src/models/cosyvoice3/frontend.cpp | 3 - src/models/cosyvoice3/hift.cpp | 10 +- src/models/cosyvoice3/session.cpp | 80 +---- .../audiocpp_cli/audiocpp_cli_path_cases.json | 6 +- 11 files changed, 673 insertions(+), 149 deletions(-) create mode 100644 docs/models/breeze_tts.md create mode 100644 docs/models/cosyvoice3.md create mode 100644 model_specs/breeze_tts.json create mode 100644 model_specs/cosyvoice3.json diff --git a/docs/models/breeze_tts.md b/docs/models/breeze_tts.md new file mode 100644 index 000000000..e13c7cdf4 --- /dev/null +++ b/docs/models/breeze_tts.md @@ -0,0 +1,69 @@ +# BreezeTTS 2 + +BreezeTTS 2 is a GGUF TTS family for instruction-conditioned speech and +prompt-audio voice cloning. The default package is Q8_0. + +## Quick Start + +Download the default Q8_0 package: + +```bash +python3 tools/model_manager_v2.py install breeze_tts_2_q8_0 --models-root models +``` + +Voice cloning: + +```bash +audiocpp_cli \ + --task clon \ + --family breeze_tts \ + --model models/Breeze-TTS-2-GGUF/breeze-tts-2-q8_0.gguf \ + --backend cuda \ + --text "Please read this line in a clear and natural voice." \ + --voice-ref assets/resources/b.wav \ + --reference-text "Some call me nature. Others call me Mother Nature. I have been here for over four and a half billion years." \ + --request-option instruction="Speak clearly and naturally." \ + --out breeze_tts_clone.wav +``` + +Voice design: + +```bash +audiocpp_cli \ + --task tts \ + --family breeze_tts \ + --model models/Breeze-TTS-2-GGUF/breeze-tts-2-q8_0.gguf \ + --backend cuda \ + --text "Welcome to the local voice demo." \ + --request-option instruction="A warm female narrator with calm pacing and studio clarity." \ + --out breeze_tts_design.wav +``` + +## Model + +| Field | Value | +|---|---| +| Family | `breeze_tts` | +| Default GGUF | `models/Breeze-TTS-2-GGUF/breeze-tts-2-q8_0.gguf` | +| Tasks | `tts`, `clon` | +| Modes | `offline` | +| Languages | `zh`, `en` | +| Voice input | Optional for `tts`; required for `clon` | + +## Options + +| Option | Values | Default | Meaning | +|---|---|---:|---| +| `--voice-ref` | WAV path | required for `clon` | Prompt/reference speaker audio. | +| `--reference-text` / `--request-option reference_text=` | text | empty | Transcript for prompt audio when cloning. | +| `--request-option instruction=` | text | `Speak clearly and naturally.` | Voice or style instruction. | +| `--request-option text_chunk_size=` | integer > 0 | `600` | Long-form text chunk size. | +| `--request-option text_chunk_mode=` | `default`, `tag_aware`, `japanese`, `endline` | `default` | Framework text chunk mode. | +| `--request-option max_tokens=` | integer > 0 | `1500` | Maximum generated acoustic frames. | +| `--request-option guidance_scale=` | float >= 0 | `1.0` | Classifier-free guidance scale. | +| `--request-option temperature=` | float >= 0 | `0.9` | Backbone sampling temperature. | +| `--request-option depth_temperature=` | float >= 0 | `0.9` | Depth decoder sampling temperature. | +| `--request-option top_k=` | integer >= 0 | `50` | Top-k sampling limit; `0` disables top-k filtering. | +| `--request-option top_p=` | `0..1` | `1.0` | Top-p sampling limit. | +| `--request-option seed=` | integer >= 0 | `0` | Generation seed. | +| `--session-option breeze_tts.reference_cache_slots=` | integer >= 0 | `1` | Prepared reference-audio cache slots. | diff --git a/docs/models/cosyvoice3.md b/docs/models/cosyvoice3.md new file mode 100644 index 000000000..e4cfae6cc --- /dev/null +++ b/docs/models/cosyvoice3.md @@ -0,0 +1,80 @@ +# CosyVoice3 + +CosyVoice3 is a GGUF TTS family for zero-shot voice cloning, cross-lingual +speech, and instruction-conditioned speech. The default package is Q8_0. + +## Quick Start + +Download the default Q8_0 package: + +```bash +python3 tools/model_manager_v2.py install cosyvoice3_q8_0 --models-root models +``` + +Zero-shot voice cloning: + +```bash +audiocpp_cli \ + --task clon \ + --family cosyvoice3 \ + --model models/CosyVoice3-GGUF/cosyvoice3-q8_0.gguf \ + --backend cuda \ + --text "This is a local CosyVoice3 voice cloning test." \ + --voice-ref assets/resources/b.wav \ + --reference-text "Some call me nature. Others call me Mother Nature. I have been here for over four and a half billion years." \ + --request-option template_name=zero_shot \ + --out cosyvoice3_clone.wav +``` + +Instruction-conditioned speech: + +```bash +audiocpp_cli \ + --task tts \ + --family cosyvoice3 \ + --model models/CosyVoice3-GGUF/cosyvoice3-q8_0.gguf \ + --backend cuda \ + --text "Please read this with a calm and friendly tone." \ + --voice-ref assets/resources/b.wav \ + --reference-text "Some call me nature. Others call me Mother Nature. I have been here for over four and a half billion years." \ + --request-option template_name=instruct \ + --request-option instruction="Speak warmly with clear articulation." \ + --out cosyvoice3_instruct.wav +``` + +## Model + +| Field | Value | +|---|---| +| Family | `cosyvoice3` | +| Default GGUF | `models/CosyVoice3-GGUF/cosyvoice3-q8_0.gguf` | +| Tasks | `tts`, `clon` | +| Modes | `offline` | +| Languages | `zh`, `en`, `ja`, `ko`, `de`, `es`, `fr`, `it`, `ru`, `yue` | +| Voice input | Required reference WAV through `--voice-ref` | + +## Templates + +| `template_name` | Use | +|---|---| +| `zero_shot` | Voice cloning with prompt audio and transcript. | +| `cross_lingual` | Cross-lingual voice cloning from prompt audio. | +| `instruct` | Instruction-conditioned speech with prompt audio. | + +## Options + +| Option | Values | Default | Meaning | +|---|---|---:|---| +| `--voice-ref` | WAV path | required | Prompt/reference speaker audio. | +| `--reference-text` / `--request-option reference_text=` | text | empty | Transcript for prompt audio. | +| `--request-option template_name=` | `zero_shot`, `cross_lingual`, `instruct` | `zero_shot` | Request template. | +| `--request-option instruction=` | text | empty | Instruction text for `instruct` mode. | +| `--request-option text_chunk_size=` | integer > 0 | `600` | Long-form text chunk size. | +| `--request-option text_chunk_mode=` | `default`, `tag_aware`, `japanese`, `endline` | `default` | Framework text chunk mode. | +| `--request-option max_tokens=` | integer > 0 | `1600` | Maximum generated speech tokens. | +| `--request-option min_tokens=` | integer >= 0 | `0` | Minimum generated tokens before stop is accepted. | +| `--request-option top_k=` | integer > 0 | `25` | AR speech-token top-k sampling limit. | +| `--request-option num_inference_steps=` | integer > 0 | `10` | Flow decoder Euler steps. | +| `--request-option seed=` | integer >= 0 | `1986` | Generation seed. | +| `--session-option cosyvoice3.reference_cache_slots=` | integer >= 0 | `4` | Prepared reference-audio cache slots. | +| `--session-option cosyvoice3.mem_saver=true\|false` | bool | `false` | Release cached runtime graphs after request phases. | diff --git a/docs/tts.md b/docs/tts.md index eba1c7794..4b2910fdf 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -3,8 +3,10 @@ | Model | Family | Task(s) | Quick Start | |---|---|---|---| | Qwen3 TTS | `qwen3_tts` | `tts`, `vdes` | [Qwen3 TTS](#qwen3-tts) | +| BreezeTTS 2 | `breeze_tts` | `tts`, `clon` | [BreezeTTS 2](models/breeze_tts.md) | | Chatterbox | `chatterbox` | `clon`, `vc` | [Chatterbox](#chatterbox) | | Confucius4-TTS | `confucius4_tts` | `clon` | [Confucius4-TTS](#confucius4-tts) | +| CosyVoice3 | `cosyvoice3` | `tts`, `clon` | [CosyVoice3](models/cosyvoice3.md) | | DramaBox | `dramabox` | `tts`, `clon` | [DramaBox](#dramabox) | | DotTTS | `dots_tts` | `tts`, `clon` | [DotTTS](#dottts) | | F5-TTS | `f5_tts` | `tts`, `clon` | [F5-TTS](community_models/f5_tts.md) | diff --git a/include/engine/models/cosyvoice3/hift.h b/include/engine/models/cosyvoice3/hift.h index af3d8ac22..03454bc2b 100644 --- a/include/engine/models/cosyvoice3/hift.h +++ b/include/engine/models/cosyvoice3/hift.h @@ -23,8 +23,7 @@ class CosyVoice3HiftRuntime { engine::runtime::AudioBuffer synthesize( const std::vector & mel, int64_t frames, - uint64_t seed, - const std::vector * source_random_values = nullptr); + uint64_t seed); void release_graphs(); private: diff --git a/model_specs/breeze_tts.json b/model_specs/breeze_tts.json new file mode 100644 index 000000000..c5299e703 --- /dev/null +++ b/model_specs/breeze_tts.json @@ -0,0 +1,234 @@ +{ + "schema_version": 1, + "family": "breeze_tts", + "display_name": "BreezeTTS 2", + "description": "BreezeTTS 2 native GGUF package for instruction-conditioned text-to-speech and prompt-audio voice cloning with T5Gemma text conditioning, Qwen-style acoustic code generation, depth codebook decoding, and Mimi waveform decoding.", + "category": "tts", + "status": "supported", + "tasks": [ + "tts", + "clone" + ], + "modes": [ + "offline" + ], + "languages": [ + "zh", + "en" + ], + "runtime": { + "tags": [ + "gguf" + ] + }, + "dependencies": [], + "capabilities": { + "tts": [ + "style_control" + ], + "clone": [ + "speaker_reference" + ] + }, + "options": { + "request": [ + { + "name": "instruction", + "type": "string", + "description": "BreezeTTS generation instruction.", + "required": false, + "default": "Speak clearly and naturally." + }, + { + "name": "reference_text", + "type": "string", + "description": "Transcript for the prompt audio when cloning.", + "required": false + }, + { + "name": "text_chunk_size", + "type": "int", + "description": "Maximum Unicode codepoints per long-form text chunk; default 600.", + "required": false, + "min": 1, + "default": 600 + }, + { + "name": "text_chunk_mode", + "type": "enum", + "description": "Framework text chunking mode.", + "values": [ + "default", + "tag_aware", + "japanese", + "endline" + ], + "required": false, + "default": "default" + }, + { + "name": "max_tokens", + "type": "int", + "description": "Maximum generated BreezeTTS acoustic frames.", + "required": false, + "min": 1, + "default": 1500 + }, + { + "name": "guidance_scale", + "type": "float", + "description": "Classifier-free guidance scale.", + "required": false, + "min": 0.0, + "default": 1.0 + }, + { + "name": "temperature", + "type": "float", + "description": "Backbone first-codebook sampling temperature.", + "required": false, + "min": 0.0, + "default": 0.9 + }, + { + "name": "depth_temperature", + "type": "float", + "description": "Depth decoder codebook sampling temperature.", + "required": false, + "min": 0.0, + "default": 0.9 + }, + { + "name": "top_k", + "type": "int", + "description": "Top-k sampling limit; 0 disables top-k filtering.", + "required": false, + "min": 0, + "default": 50 + }, + { + "name": "top_p", + "type": "float", + "description": "Top-p sampling limit.", + "required": false, + "min": 0.0, + "max": 1.0, + "default": 1.0 + }, + { + "name": "seed", + "type": "int", + "description": "Generation seed.", + "required": false, + "min": 0, + "default": 0 + } + ], + "session": [ + { + "name": "weight_type", + "type": "enum", + "description": "BreezeTTS matmul weight storage type; default native.", + "preset": "weight_type_full", + "required": false, + "default": "native" + }, + { + "name": "graph_arena_mb", + "type": "int", + "description": "Reusable runtime graph arena size in MiB; default 1024.", + "required": false, + "min": 1, + "default": 1024 + }, + { + "name": "weight_context_mb", + "type": "int", + "description": "Weight loading context size in MiB; default 2048.", + "required": false, + "min": 1, + "default": 2048 + }, + { + "name": "reference_cache_slots", + "type": "int", + "description": "Prepared reference-audio cache slots; set 0 to disable reuse.", + "required": false, + "min": 0, + "default": 1 + } + ], + "load": [] + }, + "packages": [ + { + "id": "breeze_tts_2_q8_0", + "display_name": "BreezeTTS 2 Q8_0 GGUF", + "default": true, + "format": "gguf", + "precision": "q8_0", + "target_directory": "Breeze-TTS-2-GGUF", + "files": [ + "Breeze-TTS-2-GGUF/breeze-tts-2-q8_0.gguf" + ], + "strip_prefix": "Breeze-TTS-2-GGUF", + "download": { + "kind": "huggingface_snapshot", + "repo": "audio-cpp/Breeze-TTS-2-GGUF", + "revision": "main", + "gated": false + } + }, + { + "id": "breeze_tts_2_bf16", + "display_name": "BreezeTTS 2 BF16 GGUF", + "default": false, + "format": "gguf", + "precision": "bf16", + "target_directory": "Breeze-TTS-2-GGUF", + "files": [ + "Breeze-TTS-2-GGUF/breeze-tts-2-bf16.gguf" + ], + "strip_prefix": "Breeze-TTS-2-GGUF", + "download": { + "kind": "huggingface_snapshot", + "repo": "audio-cpp/Breeze-TTS-2-GGUF", + "revision": "main", + "gated": false + } + } + ], + "ui": { + "recommended_package": "breeze_tts_2_q8_0", + "tags": [ + "TTS", + "Clone", + "GGUF" + ], + "docs": [ + "docs/tts.md", + "docs/gguf.md" + ] + }, + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "weights": "$gguf" + }, + "files": { + "config_json": "model:config.json", + "audio_tokenizer_config_json": "model:audio_tokenizer/config.json", + "tokenizer_config_json": "model:tokenizer_config.json", + "tokenizer_json": "model:tokenizer.json" + }, + "tensors": { + "model_weights": { + "source": "weights:", + "prefix": "model" + } + } + } + ] +} diff --git a/model_specs/cosyvoice3.json b/model_specs/cosyvoice3.json new file mode 100644 index 000000000..5efae24ea --- /dev/null +++ b/model_specs/cosyvoice3.json @@ -0,0 +1,273 @@ +{ + "schema_version": 1, + "family": "cosyvoice3", + "display_name": "CosyVoice3", + "description": "Fun-CosyVoice3 native GGUF package for zero-shot, cross-lingual, and instruction-conditioned text-to-speech using CosyVoice3 speech-token AR generation, causal masked flow mel decoding, CAM++ speaker conditioning, and causal HiFT waveform decoding.", + "category": "tts", + "status": "supported", + "tasks": [ + "tts", + "clone" + ], + "modes": [ + "offline" + ], + "languages": [ + "zh", + "en", + "ja", + "ko", + "de", + "es", + "fr", + "it", + "ru", + "yue" + ], + "runtime": { + "tags": [ + "gguf" + ] + }, + "dependencies": [], + "capabilities": { + "tts": [ + "speaker_reference", + "style_control" + ], + "clone": [ + "speaker_reference" + ] + }, + "options": { + "request": [ + { + "name": "template_name", + "type": "enum", + "description": "CosyVoice3 request template.", + "required": false, + "values": [ + "zero_shot", + "cross_lingual", + "instruct" + ], + "default": "zero_shot" + }, + { + "name": "reference_text", + "type": "string", + "description": "Transcript for the prompt audio.", + "required": false + }, + { + "name": "instruction", + "type": "string", + "description": "Instruction text for instruct mode.", + "required": false + }, + { + "name": "text_chunk_size", + "type": "int", + "description": "Maximum Unicode codepoints per long-form text chunk; default 600.", + "required": false, + "min": 1, + "default": 600 + }, + { + "name": "text_chunk_mode", + "type": "enum", + "description": "Framework text chunking mode.", + "values": [ + "default", + "tag_aware", + "japanese", + "endline" + ], + "required": false, + "default": "default" + }, + { + "name": "max_tokens", + "type": "int", + "description": "Maximum generated CosyVoice3 speech tokens.", + "required": false, + "min": 1, + "default": 1600 + }, + { + "name": "min_tokens", + "type": "int", + "description": "Minimum generated CosyVoice3 speech tokens before stop tokens are accepted.", + "required": false, + "min": 0, + "default": 0 + }, + { + "name": "top_k", + "type": "int", + "description": "AR speech-token top-k sampling limit; default 25.", + "required": false, + "min": 1, + "default": 25 + }, + { + "name": "num_inference_steps", + "type": "int", + "description": "Flow decoder Euler steps; default 10.", + "required": false, + "min": 1, + "default": 10 + }, + { + "name": "seed", + "type": "int", + "description": "Generation seed.", + "required": false, + "min": 0, + "default": 1986 + } + ], + "session": [ + { + "name": "weight_type", + "type": "enum", + "description": "CosyVoice3 matmul weight storage type; default native.", + "preset": "weight_type_full", + "required": false, + "default": "native" + }, + { + "name": "conv_weight_type", + "type": "enum", + "description": "CosyVoice3 convolution weight storage type; default native.", + "preset": "weight_type_conv", + "required": false, + "default": "native" + }, + { + "name": "graph_arena_mb", + "type": "int", + "description": "Reusable runtime graph arena size in MiB; default 1024.", + "required": false, + "min": 1, + "default": 1024 + }, + { + "name": "weight_context_mb", + "type": "int", + "description": "Weight loading context size in MiB; default 2048.", + "required": false, + "min": 1, + "default": 2048 + }, + { + "name": "reference_cache_slots", + "type": "int", + "description": "Prepared reference-audio cache slots; set 0 to disable reuse.", + "required": false, + "min": 0, + "default": 4 + }, + { + "name": "mem_saver", + "type": "bool", + "description": "Release cached runtime graphs after each request to reduce peak VRAM; default false.", + "required": false, + "default": false + } + ], + "load": [] + }, + "packages": [ + { + "id": "cosyvoice3_q8_0", + "display_name": "CosyVoice3 Q8_0 GGUF", + "default": true, + "format": "gguf", + "precision": "q8_0", + "target_directory": "CosyVoice3-GGUF", + "files": [ + "CosyVoice3-GGUF/cosyvoice3-q8_0.gguf" + ], + "strip_prefix": "CosyVoice3-GGUF", + "download": { + "kind": "huggingface_snapshot", + "repo": "audio-cpp/CosyVoice3-GGUF", + "revision": "main", + "gated": false + } + }, + { + "id": "cosyvoice3_f32", + "display_name": "CosyVoice3 F32 GGUF", + "default": false, + "format": "gguf", + "precision": "f32", + "target_directory": "CosyVoice3-GGUF", + "files": [ + "CosyVoice3-GGUF/cosyvoice3-f32.gguf" + ], + "strip_prefix": "CosyVoice3-GGUF", + "download": { + "kind": "huggingface_snapshot", + "repo": "audio-cpp/CosyVoice3-GGUF", + "revision": "main", + "gated": false + } + } + ], + "ui": { + "recommended_package": "cosyvoice3_q8_0", + "tags": [ + "TTS", + "Clone", + "GGUF" + ], + "docs": [ + "docs/tts.md", + "docs/gguf.md" + ] + }, + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "weights": "$gguf" + }, + "files": { + "cosyvoice3_yaml": "model:cosyvoice3.yaml", + "qwen_config": "model:CosyVoice-BlankEN/config.json", + "tokenizer_config": "model:CosyVoice-BlankEN/tokenizer_config.json", + "vocab_json": "model:CosyVoice-BlankEN/vocab.json", + "merges_txt": "model:CosyVoice-BlankEN/merges.txt" + }, + "tensors": { + "llm_weights": { + "source": "weights:", + "prefix": "llm" + }, + "flow_weights": { + "source": "weights:", + "prefix": "flow" + }, + "hift_weights": { + "source": "weights:", + "prefix": "hift" + }, + "campplus_weights": { + "source": "weights:", + "prefix": "campplus" + }, + "speech_tokenizer_weights": { + "source": "weights:", + "prefix": "speech_tokenizer" + }, + "blank_en_weights": { + "source": "weights:", + "prefix": "blank_en" + } + } + } + ] +} diff --git a/src/models/cosyvoice3/flow.cpp b/src/models/cosyvoice3/flow.cpp index e81256724..4676a3869 100644 --- a/src/models/cosyvoice3/flow.cpp +++ b/src/models/cosyvoice3/flow.cpp @@ -90,15 +90,6 @@ std::vector position_ids(int64_t steps) { return out; } -std::vector trace_dims(const core::TensorShape & shape) { - std::vector out; - out.reserve(shape.rank); - for (size_t i = 0; i < shape.rank; ++i) { - out.push_back(shape.dims[i]); - } - return out; -} - std::vector cosine_time_schedule(int64_t steps) { if (steps <= 0) { throw std::runtime_error("CosyVoice3 num_inference_steps must be positive"); @@ -426,7 +417,6 @@ class ConditionGraph { throw std::runtime_error("CosyVoice3 flow condition graph compute failed"); } auto output = core::read_tensor_f32(output_.tensor); - engine::debug::trace_log_f32("cosyvoice3.flow.condition_mu", {tokens * config_.token_mel_ratio, config_.flow_mel_channels}, output); return output; } @@ -547,13 +537,6 @@ class DiTGraph { if (core::compute_backend_graph(execution_.backend(), mem_.graph, nullptr, "cosyvoice3.flow.dit") != GGML_STATUS_SUCCESS) { throw std::runtime_error("CosyVoice3 DiT graph compute failed"); } - if (!traced_first_step_ && engine::debug::trace_log_enabled()) { - traced_first_step_ = true; - engine::debug::trace_log_f32("cosyvoice3.flow.dit.input_cat", trace_dims(debug_input_cat_.shape), core::read_tensor_f32(debug_input_cat_.tensor)); - engine::debug::trace_log_f32("cosyvoice3.flow.dit.input_proj", trace_dims(debug_input_proj_.shape), core::read_tensor_f32(debug_input_proj_.tensor)); - engine::debug::trace_log_f32("cosyvoice3.flow.dit.input_pos", trace_dims(debug_input_pos_.shape), core::read_tensor_f32(debug_input_pos_.tensor)); - engine::debug::trace_log_f32("cosyvoice3.flow.dit.input_embed", trace_dims(debug_input_embed_.shape), core::read_tensor_f32(debug_input_embed_.tensor)); - } return core::read_tensor_f32(output_.tensor); } @@ -567,11 +550,6 @@ class DiTGraph { time_ = nullptr; positions_ = nullptr; output_ = {}; - debug_input_cat_ = {}; - debug_input_proj_ = {}; - debug_input_pos_ = {}; - debug_input_embed_ = {}; - traced_first_step_ = false; } private: @@ -620,16 +598,12 @@ class DiTGraph { input = modules::ConcatModule({2}).build(ctx, input, mu); input = modules::ConcatModule({2}).build(ctx, input, spks_btf); input = core::ensure_backend_addressable_layout(ctx, input); - debug_input_cat_ = input; input = modules::LinearModule({320, config_.flow_hidden_size, true}).build(ctx, input, weights_->input_projection); - debug_input_proj_ = core::ensure_backend_addressable_layout(ctx, input); - input = debug_input_proj_; + input = core::ensure_backend_addressable_layout(ctx, input); auto pos = causal_conv_pos_embed(ctx, input, *weights_); - debug_input_pos_ = core::ensure_backend_addressable_layout(ctx, pos); - pos = debug_input_pos_; + pos = core::ensure_backend_addressable_layout(ctx, pos); input = modules::AddModule().build(ctx, input, pos); - debug_input_embed_ = core::ensure_backend_addressable_layout(ctx, input); - input = debug_input_embed_; + input = core::ensure_backend_addressable_layout(ctx, input); time = modules::LinearModule({kTimeEmbeddingSize, config_.flow_hidden_size, true}).build(ctx, time, weights_->time_fc1); time = modules::SiluModule().build(ctx, time); @@ -677,11 +651,6 @@ class DiTGraph { ggml_tensor * time_ = nullptr; ggml_tensor * positions_ = nullptr; core::TensorValue output_; - core::TensorValue debug_input_cat_; - core::TensorValue debug_input_proj_; - core::TensorValue debug_input_pos_; - core::TensorValue debug_input_embed_; - bool traced_first_step_ = false; }; std::vector normalize_and_project_speaker( @@ -821,24 +790,6 @@ class CosyVoice3FlowRuntime::Impl { const auto schedule = cosine_time_schedule(request.num_inference_steps); const auto dit_start = std::chrono::steady_clock::now(); for (int64_t step = 1; step < static_cast(schedule.size()); ++step) { - if (step == 1) { - engine::debug::trace_log_f32( - "cosyvoice3.flow.dit.x_in", - {2, c.flow_mel_channels, total_frames}, - x_batched); - engine::debug::trace_log_f32( - "cosyvoice3.flow.dit.mu_in", - {2, c.flow_mel_channels, total_frames}, - mu); - engine::debug::trace_log_f32( - "cosyvoice3.flow.dit.cond_in", - {2, c.flow_mel_channels, total_frames}, - cond); - engine::debug::trace_log_f32( - "cosyvoice3.flow.dit.spks_in", - {2, c.flow_mel_channels}, - spks); - } const float t = schedule[static_cast(step - 1)]; const float dt = schedule[static_cast(step)] - t; auto temb = timestep_embedding(t); @@ -846,12 +797,6 @@ class CosyVoice3FlowRuntime::Impl { std::copy(temb.begin(), temb.end(), time_batched.begin()); std::copy(temb.begin(), temb.end(), time_batched.begin() + kTimeEmbeddingSize); auto pred = dit_.run(x_batched, mu, cond, spks, time_batched, total_frames); - if (step == 1) { - engine::debug::trace_log_f32( - "cosyvoice3.flow.dit.pred_step1", - {2, c.flow_mel_channels, total_frames}, - pred); - } const size_t branch = static_cast(c.flow_mel_channels * total_frames); for (size_t i = 0; i < branch; ++i) { const float guided = (1.0F + kInferenceCfgRate) * pred[i] - kInferenceCfgRate * pred[branch + i]; @@ -870,7 +815,6 @@ class CosyVoice3FlowRuntime::Impl { x_batched[static_cast(ch * total_frames + request.prompt_mel_frames + frame)]; } } - engine::debug::trace_log_f32("cosyvoice3.flow.output_mel", {target_frames, c.flow_mel_channels}, out.mel); engine::debug::timing_log_scalar("cosyvoice3.flow.total_ms", engine::debug::elapsed_ms(start)); engine::debug::trace_log_scalar("cosyvoice3.flow.total_frames", static_cast(total_frames)); engine::debug::trace_log_scalar("cosyvoice3.flow.target_frames", static_cast(target_frames)); diff --git a/src/models/cosyvoice3/frontend.cpp b/src/models/cosyvoice3/frontend.cpp index 3e878f4bd..a8df8f22c 100644 --- a/src/models/cosyvoice3/frontend.cpp +++ b/src/models/cosyvoice3/frontend.cpp @@ -8,7 +8,6 @@ #include "engine/framework/modules/speech_encoders/campplus_encoder.h" #include "engine/framework/modules/speech_encoders/s3_tokenizer.h" #include "engine/framework/runtime/cache_slots.h" -#include "engine/framework/debug/trace.h" #include #include @@ -221,10 +220,8 @@ class CosyVoice3Frontend::Impl { compute_prompt_mel(audio, features.prompt_mel, features.prompt_mel_frames); int64_t fbank_frames = 0; auto fbank = compute_campplus_fbank(audio, fbank_frames); - engine::debug::trace_log_f32("cosyvoice3.frontend.campplus_fbank", {fbank_frames, 80}, fbank); auto speaker = campplus_.embed_from_features(fbank, fbank_frames, 80); features.speaker_embedding = std::move(speaker.embedding); - engine::debug::trace_log_f32("cosyvoice3.frontend.speaker_embedding", {assets_->config.speaker_dim}, features.speaker_embedding); if (static_cast(features.speaker_embedding.size()) != assets_->config.speaker_dim) { throw std::runtime_error("CosyVoice3 CAMPPlus speaker embedding size mismatch"); } diff --git a/src/models/cosyvoice3/hift.cpp b/src/models/cosyvoice3/hift.cpp index 6043c68b6..048a54e5e 100644 --- a/src/models/cosyvoice3/hift.cpp +++ b/src/models/cosyvoice3/hift.cpp @@ -61,8 +61,7 @@ class CosyVoice3HiftRuntime::Impl { engine::runtime::AudioBuffer synthesize( const std::vector & mel, int64_t frames, - uint64_t seed, - const std::vector * source_random_values) { + uint64_t seed) { if (frames <= 0 || static_cast(mel.size()) != frames * 80) { throw std::runtime_error("CosyVoice3 HiFT mel shape mismatch"); } @@ -73,7 +72,7 @@ class CosyVoice3HiftRuntime::Impl { mel[static_cast(frame * 80 + channel)]; } } - auto out = component_.synthesize(channel_major, frames, seed, 0, source_random_values); + auto out = component_.synthesize(channel_major, frames, seed, 0); engine::runtime::AudioBuffer audio; audio.sample_rate = static_cast(out.sample_rate); audio.channels = 1; @@ -101,9 +100,8 @@ CosyVoice3HiftRuntime::~CosyVoice3HiftRuntime() = default; engine::runtime::AudioBuffer CosyVoice3HiftRuntime::synthesize( const std::vector & mel, int64_t frames, - uint64_t seed, - const std::vector * source_random_values) { - return impl_->synthesize(mel, frames, seed, source_random_values); + uint64_t seed) { + return impl_->synthesize(mel, frames, seed); } void CosyVoice3HiftRuntime::release_graphs() { diff --git a/src/models/cosyvoice3/session.cpp b/src/models/cosyvoice3/session.cpp index b1e2affbc..0b5e4c18b 100644 --- a/src/models/cosyvoice3/session.cpp +++ b/src/models/cosyvoice3/session.cpp @@ -10,11 +10,8 @@ #include "engine/models/cosyvoice3/hift.h" #include -#include -#include #include #include -#include #include namespace engine::models::cosyvoice3 { @@ -52,64 +49,6 @@ std::string instruction_text(const runtime::TaskRequest & request) { return runtime::find_option(request.options, {"instruction"}).value_or(""); } -std::unordered_map validation_options( - const std::unordered_map & options) { - auto out = options; - out.erase("teacher_force_tokens"); - out.erase("teacher_force_source_random"); - return out; -} - -std::vector read_teacher_force_tokens(const std::string & path) { - std::ifstream in(path); - if (!in) { - throw std::runtime_error("CosyVoice3 teacher_force_tokens file cannot be opened: " + path); - } - std::vector out; - std::string token; - char ch = '\0'; - while (in.get(ch)) { - if (std::isdigit(static_cast(ch)) || ch == '-') { - token.push_back(ch); - } else if (!token.empty()) { - out.push_back(static_cast(std::stoi(token))); - token.clear(); - } - } - if (!token.empty()) { - out.push_back(static_cast(std::stoi(token))); - } - if (out.empty()) { - throw std::runtime_error("CosyVoice3 teacher_force_tokens file contains no tokens: " + path); - } - return out; -} - -std::vector read_float_list(const std::string & path, const char * label) { - std::ifstream in(path); - if (!in) { - throw std::runtime_error(std::string("CosyVoice3 ") + label + " file cannot be opened: " + path); - } - std::vector out; - std::string token; - char ch = '\0'; - while (in.get(ch)) { - if (std::isdigit(static_cast(ch)) || ch == '-' || ch == '+' || ch == '.' || ch == 'e' || ch == 'E') { - token.push_back(ch); - } else if (!token.empty()) { - out.push_back(std::stof(token)); - token.clear(); - } - } - if (!token.empty()) { - out.push_back(std::stof(token)); - } - if (out.empty()) { - throw std::runtime_error(std::string("CosyVoice3 ") + label + " file contains no values: " + path); - } - return out; -} - std::vector split_request(const runtime::TaskRequest & request) { const int64_t text_chunk_size = engine::text::parse_text_chunk_size_override(request.options).value_or(kDefaultTextChunkSize); @@ -218,13 +157,13 @@ runtime::RunMode CosyVoice3Session::run_mode() const { } void CosyVoice3Session::prepare(const runtime::SessionPreparationRequest & request) { - runtime::validate_spec_backed_request_options(validation_options(request.options), *contract_, kModelName); + runtime::validate_spec_backed_request_options(request.options, *contract_, kModelName); mark_prepared(); } runtime::TaskResult CosyVoice3Session::run(const runtime::TaskRequest & request) { const auto wall_start = std::chrono::steady_clock::now(); - runtime::validate_spec_backed_request_options(validation_options(request.options), *contract_, kModelName); + runtime::validate_spec_backed_request_options(request.options, *contract_, kModelName); require_prepared("CosyVoice3 run"); if (!request.text_input.has_value() || request.text_input->text.empty()) { throw std::runtime_error("CosyVoice3 requires text input"); @@ -267,12 +206,7 @@ runtime::TaskResult CosyVoice3Session::run(const runtime::TaskRequest & request) if (index > 0) { ar_request.seed += static_cast(index); } - CosyVoice3ArOutput ar_output; - if (const auto teacher_force_path = runtime::find_option(chunk.options, {"teacher_force_tokens"})) { - ar_output.speech_tokens = read_teacher_force_tokens(*teacher_force_path); - } else { - ar_output = ar_->generate(ar_request); - } + CosyVoice3ArOutput ar_output = ar_->generate(ar_request); if (mem_saver_) { ar_->release_graphs(); } @@ -292,15 +226,9 @@ runtime::TaskResult CosyVoice3Session::run(const runtime::TaskRequest & request) if (mem_saver_) { flow_->release_graphs(); } - std::vector source_random_values; - const std::vector * source_random_ptr = nullptr; - if (const auto random_path = runtime::find_option(chunk.options, {"teacher_force_source_random"})) { - source_random_values = read_float_list(*random_path, "teacher_force_source_random"); - source_random_ptr = &source_random_values; - } runtime::append_audio_buffer( merged_audio, - hift_->synthesize(flow_output.mel, flow_output.frames, flow_request.seed, source_random_ptr)); + hift_->synthesize(flow_output.mel, flow_output.frames, flow_request.seed)); if (mem_saver_) { hift_->release_graphs(); } diff --git a/tools/audiocpp_cli/audiocpp_cli_path_cases.json b/tools/audiocpp_cli/audiocpp_cli_path_cases.json index f6ef6d967..89579f802 100644 --- a/tools/audiocpp_cli/audiocpp_cli_path_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_path_cases.json @@ -3593,7 +3593,7 @@ "id": "cosyvoice3_official_paths", "coverage": "CosyVoice3 official zero-shot, cross-lingual, instruction, and hotfix-pronunciation paths after a short same-session warmup", "family": "cosyvoice3", - "model": "models/CosyVoice3-GGUF/cosyvoice3-orig.gguf", + "model": "models/audio.cpp-gguf/CosyVoice3-GGUF/cosyvoice3-q8_0.gguf", "task": "clon", "mode": "offline", "outputs": [ @@ -3656,7 +3656,7 @@ "id": "breeze_tts_clone_official_paths", "coverage": "BreezeTTS 2 official voice clone and voice direction paths using prompt audio, prompt transcript, T5Gemma text conditioning, Qwen acoustic generation, depth codebook decoding, and Mimi decode", "family": "breeze_tts", - "model": "models/audio.cpp-gguf/Breeze-TTS-2-GGUF/breeze-tts-2-orig.gguf", + "model": "models/audio.cpp-gguf/Breeze-TTS-2-GGUF/breeze-tts-2-bf16.gguf", "task": "clon", "mode": "offline", "outputs": [ @@ -3688,7 +3688,7 @@ "id": "breeze_tts_design_official_paths", "coverage": "BreezeTTS 2 official English and Chinese voice design paths using instruction conditioning, T5Gemma text conditioning, Qwen acoustic generation, depth codebook decoding, and Mimi decode", "family": "breeze_tts", - "model": "models/audio.cpp-gguf/Breeze-TTS-2-GGUF/breeze-tts-2-orig.gguf", + "model": "models/audio.cpp-gguf/Breeze-TTS-2-GGUF/breeze-tts-2-bf16.gguf", "task": "tts", "mode": "offline", "outputs": [ From 34e343e252ba26787c232325d5a2a4ce0fc86c0f Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Fri, 28 Aug 2026 01:51:07 -0400 Subject: [PATCH 03/10] Add CosyVoice3 and Breeze UI support --- src/models/cosyvoice3/tokenizer_text.cpp | 30 +++++++++++++++++++----- webui/configs/model_params.json | 21 +++++++++++++++++ webui/configs/models_catalog.json | 8 +++++++ webui/native/dist/index.html | 18 +++++++------- webui/native/src/lib/catalog.ts | 2 ++ webui/native/src/routes/+page.svelte | 7 +++++- 6 files changed, 70 insertions(+), 16 deletions(-) diff --git a/src/models/cosyvoice3/tokenizer_text.cpp b/src/models/cosyvoice3/tokenizer_text.cpp index 9e3823c69..cd8398cb0 100644 --- a/src/models/cosyvoice3/tokenizer_text.cpp +++ b/src/models/cosyvoice3/tokenizer_text.cpp @@ -3,12 +3,16 @@ #include "engine/framework/tokenizers/llama_bpe.h" #include +#include +#include #include namespace engine::models::cosyvoice3 { namespace { constexpr int32_t kEndOfPromptId = 151646; +constexpr std::string_view kEndOfPromptText = "<|endofprompt|>"; +constexpr std::string_view kDefaultPromptPrefix = "You are a helpful assistant."; std::vector cosyvoice3_special_tokens() { std::vector tokens; @@ -305,6 +309,22 @@ void require_end_of_prompt(const std::vector & tokens, const char * con throw std::runtime_error(std::string("CosyVoice3 ") + context + " must contain <|endofprompt|>"); } +std::string prompt_with_boundary(std::string_view text) { + std::string out(text); + if (out.find(kEndOfPromptText) == std::string::npos) { + out = std::string(kDefaultPromptPrefix) + std::string(kEndOfPromptText) + out; + } + return out; +} + +std::string instruction_with_boundary(std::string_view instruction) { + std::string out(instruction); + if (out.find(kEndOfPromptText) == std::string::npos) { + out = std::string(kDefaultPromptPrefix) + " " + out + std::string(kEndOfPromptText); + } + return out; +} + } // namespace class CosyVoice3TextTokenizer::Impl { @@ -335,17 +355,15 @@ CosyVoice3TextTokens CosyVoice3TextTokenizer::encode_zero_shot( std::string_view text, std::string_view prompt_text) const { CosyVoice3TextTokens out; - out.prompt = impl_->tokenizer->encode(std::string(prompt_text), true); + out.prompt = impl_->tokenizer->encode(prompt_with_boundary(prompt_text), true); out.target = impl_->tokenizer->encode(std::string(text), true); - std::vector joined = out.prompt; - joined.insert(joined.end(), out.target.begin(), out.target.end()); - require_end_of_prompt(joined, "zero-shot prompt/text"); + require_end_of_prompt(out.prompt, "zero-shot prompt"); return out; } CosyVoice3TextTokens CosyVoice3TextTokenizer::encode_cross_lingual(std::string_view text) const { CosyVoice3TextTokens out; - out.target = impl_->tokenizer->encode(std::string(text), true); + out.target = impl_->tokenizer->encode(prompt_with_boundary(text), true); require_end_of_prompt(out.target, "cross-lingual text"); return out; } @@ -354,7 +372,7 @@ CosyVoice3TextTokens CosyVoice3TextTokenizer::encode_instruct( std::string_view text, std::string_view instruction) const { CosyVoice3TextTokens out; - out.prompt = impl_->tokenizer->encode(std::string(instruction), true); + out.prompt = impl_->tokenizer->encode(instruction_with_boundary(instruction), true); out.target = impl_->tokenizer->encode(std::string(text), true); require_end_of_prompt(out.prompt, "instruction"); return out; diff --git a/webui/configs/model_params.json b/webui/configs/model_params.json index 7b3d71bbc..3b595cf3d 100644 --- a/webui/configs/model_params.json +++ b/webui/configs/model_params.json @@ -98,6 +98,27 @@ {"name": "text_chunk_mode", "type": "choice", "label": "text_chunk_mode", "default": "default", "choices": ["default", "tag_aware", "japanese", "endline"]} ], + "breeze_tts": [ + {"name": "instruction", "type": "text", "label": "instruction", "label_en": "Instruction", "default": "", "placeholder": "Describe the target voice for VoiceDesign."}, + {"name": "text_chunk_size", "type": "number", "label": "text_chunk_size", "default": 600, "minimum": 1, "step": 1, "precision": 0}, + {"name": "text_chunk_mode", "type": "choice", "label": "text_chunk_mode", "default": "default", "choices": ["default", "tag_aware", "japanese", "endline"]}, + {"name": "guidance_scale", "type": "slider", "label": "guidance_scale", "default": 1.0, "minimum": 0.0, "maximum": 10.0, "step": 0.1}, + {"name": "temperature", "type": "slider", "label": "temperature", "default": 0.9, "minimum": 0.0, "maximum": 2.0, "step": 0.05}, + {"name": "depth_temperature", "type": "slider", "label": "depth_temperature", "default": 0.9, "minimum": 0.0, "maximum": 2.0, "step": 0.05}, + {"name": "top_k", "type": "number", "label": "top_k", "default": 50, "minimum": 0, "step": 1, "precision": 0}, + {"name": "top_p", "type": "slider", "label": "top_p", "default": 1.0, "minimum": 0.0, "maximum": 1.0, "step": 0.01} + ], + + "cosyvoice3": [ + {"name": "template_name", "type": "choice", "label": "template_name", "default": "zero_shot", "choices": ["zero_shot", "cross_lingual", "instruct"]}, + {"name": "instruction", "type": "text", "label": "instruction", "label_en": "Instruction", "default": "", "placeholder": "Used by template_name=instruct."}, + {"name": "text_chunk_size", "type": "number", "label": "text_chunk_size", "default": 600, "minimum": 1, "step": 1, "precision": 0}, + {"name": "text_chunk_mode", "type": "choice", "label": "text_chunk_mode", "default": "default", "choices": ["default", "tag_aware", "japanese", "endline"]}, + {"name": "num_inference_steps", "type": "number", "label": "num_inference_steps", "default": 10, "minimum": 1, "step": 1, "precision": 0}, + {"name": "min_tokens", "type": "number", "label": "min_tokens", "default": 0, "minimum": 0, "step": 1, "precision": 0}, + {"name": "top_k", "type": "number", "label": "top_k", "default": 25, "minimum": 1, "step": 1, "precision": 0} + ], + "inflect_v2": [ {"name": "speaking_rate", "type": "slider", "label": "speaking_rate（语速倍率）", "label_en": "speaking_rate", "default": 1.0, "minimum": 0.5, "maximum": 2.0, "step": 0.05}, {"name": "variation", "type": "slider", "label": "variation（音色变化）", "label_en": "variation", "default": 0.667, "minimum": 0.0, "maximum": 1.0, "step": 0.01}, diff --git a/webui/configs/models_catalog.json b/webui/configs/models_catalog.json index 36e8f4041..4d2e49d98 100644 --- a/webui/configs/models_catalog.json +++ b/webui/configs/models_catalog.json @@ -14,6 +14,14 @@ { "id": "qwen3-tts", "display_name": "Qwen3-TTS 0.6B (tts)", "family": "qwen3_tts", "path": "models/Qwen3-TTS-12Hz-0.6B-Base", "task": "tts", "mode": "offline", "download_id": "qwen3_tts_0_6b_base", "min_vram_gb": 5 }, { "id": "qwen3-tts-1.7b", "display_name": "Qwen3-TTS 1.7B Base (tts)", "family": "qwen3_tts", "path": "models/Qwen3-TTS-12Hz-1.7B-Base", "task": "tts", "mode": "offline", "download_id": "qwen3_tts_1_7b_base", "min_vram_gb": 8 }, { "id": "qwen3-tts-1.7b-custom", "display_name": "Qwen3-TTS 1.7B CustomVoice (tts)", "family": "qwen3_tts", "path": "models/Qwen3-TTS-12Hz-1.7B-CustomVoice", "task": "tts", "mode": "offline", "download_id": "qwen3_tts_1_7b_custom_voice", "min_vram_gb": 8 }, + { "id": "breeze-tts", "display_name": "BreezeTTS 2 VoiceDesign", "family": "breeze_tts", "path": "models/Breeze-TTS-2-GGUF/breeze-tts-2-q8_0.gguf", "task": "tts", "mode": "offline", "download_id": "breeze_tts_2_q8_0", "min_vram_gb": 8, + "input_hint_en": "**BreezeTTS 2 VoiceDesign**: enter text and describe the target voice in Model parameters. No reference voice is required." }, + { "id": "breeze-tts-clone", "display_name": "BreezeTTS 2 Clone", "family": "breeze_tts", "path": "models/Breeze-TTS-2-GGUF/breeze-tts-2-q8_0.gguf", "task": "clon", "mode": "offline", "download_id": "breeze_tts_2_q8_0", "min_vram_gb": 8, + "input_hint_en": "**BreezeTTS 2 Clone**: upload a reference voice and provide the matching reference transcript." }, + { "id": "cosyvoice3", "display_name": "CosyVoice3 Clone", "family": "cosyvoice3", "path": "models/CosyVoice3-GGUF/cosyvoice3-q8_0.gguf", "task": "clon", "mode": "offline", "download_id": "cosyvoice3_q8_0", "min_vram_gb": 8, + "input_hint_en": "**CosyVoice3 Clone**: upload a reference voice and provide the matching reference transcript. Use `template_name` for zero-shot or cross-lingual requests." }, + { "id": "cosyvoice3-instruct", "display_name": "CosyVoice3 Instruct", "family": "cosyvoice3", "path": "models/CosyVoice3-GGUF/cosyvoice3-q8_0.gguf", "task": "tts", "mode": "offline", "download_id": "cosyvoice3_q8_0", "min_vram_gb": 8, + "input_hint_en": "**CosyVoice3 Instruct**: upload a reference voice, provide its transcript, and set `template_name=instruct` with an instruction." }, { "id": "miotts", "display_name": "MioTTS 1.7B (tts; needs MioCodec)", "family": "miotts", "path": "models/MioTTS-1.7B", "task": "tts", "mode": "offline", "download_id": "miotts_1_7b", "min_vram_gb": 8 }, { "id": "soprano-tts", "display_name": "Soprano TTS (tts)", "family": "soprano_tts", "path": "models/Soprano-1.1-80M-GGUF", "task": "tts", "mode": "offline", "download_id": "soprano_1_1_80m_q8_0", "min_vram_gb": 1 }, { "id": "voxcpm2", "display_name": "VoxCPM2 (tts)", "family": "voxcpm2", "path": "models/VoxCPM2", "task": "tts", "mode": "offline", "download_id": "voxcpm2", "session_options": { "voxcpm2.weight_type": "q8_0" }, "min_vram_gb": 6 }, diff --git a/webui/native/dist/index.html b/webui/native/dist/index.html index 39e4b1ef9..ac74c900b 100644 --- a/webui/native/dist/index.html +++ b/webui/native/dist/index.html @@ -31,20 +31,20 @@
diff --git a/webui/native/src/lib/catalog.ts b/webui/native/src/lib/catalog.ts index baaf51af4..efb396de6 100644 --- a/webui/native/src/lib/catalog.ts +++ b/webui/native/src/lib/catalog.ts @@ -42,6 +42,8 @@ const specsByFamily = new Map(Object.values(specModules).map((spec) => [spec.fam const exposeAllGgufPackageFamilies = new Set([ 'audiosr', 'controlfoley', + 'breeze_tts', + 'cosyvoice3', 'firered_audio', 'fireredtts3', 'meanvc2', diff --git a/webui/native/src/routes/+page.svelte b/webui/native/src/routes/+page.svelte index 777984830..f39d70a67 100644 --- a/webui/native/src/routes/+page.svelte +++ b/webui/native/src/routes/+page.svelte @@ -145,6 +145,8 @@ const exposeAllStudioPackageFamilies = new Set([ 'audiosr', 'controlfoley', + 'breeze_tts', + 'cosyvoice3', 'firered_audio', 'fireredtts3', 'meanvc2', @@ -272,6 +274,8 @@ qwen3_asr: 'Qwen3-ASR', vevo2: 'Vevo2', seed_vc: 'Seed-VC', + breeze_tts: 'BreezeTTS 2', + cosyvoice3: 'CosyVoice3', magpie_tts: 'MagpieTTS', meanvc2: 'MeanVC2', personaplex: 'PersonaPlex' @@ -392,13 +396,14 @@ $: usesDurationSecOption = selected?.family === 'controlfoley' || selected?.family === 'midashenglm_gen'; + $: supportsTextOnlyTts = selected?.family === 'breeze_tts' && selected?.task === 'tts'; $: needsSource = ['asr', 'vc', 'svc', 's2s', 'sep', 'vad', 'diar', 'align', 'midi'].includes(selected?.task) || isFireRedAudioEdit; $: acceptsSource = needsSource || selected?.task === 'gen'; $: acceptsVideo = selected?.request_options?.includes('video') === true; $: needsVoice = (['clon', 'vc', 'svc'].includes(selected?.task) && selected?.family !== 'rvc') || (selected?.task === 's2s' && selected?.family === 'personaplex') || - (selected?.task === 'tts' && !['supertonic'].includes(selected?.family)); + (selected?.task === 'tts' && !['supertonic'].includes(selected?.family) && !supportsTextOnlyTts); $: usesVibeVoiceSpeakerFiles = selected?.family === 'vibevoice'; $: isQwenBase = selected?.task === 'tts' && selected?.family === 'qwen3_tts' && !selected?.id.includes('custom'); From cc8fed4b328af878288b81ee03402700635c32df Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Fri, 28 Aug 2026 01:59:07 -0400 Subject: [PATCH 04/10] Move Breeze voice design to design route --- model_specs/breeze_tts.json | 11 ++++++++--- model_specs/cosyvoice3.json | 4 ++-- src/models/breeze_tts/session.cpp | 6 ++++-- webui/configs/model_params.json | 10 ++++++++++ webui/configs/models_catalog.json | 2 +- webui/native/dist/index.html | 8 ++++---- 6 files changed, 29 insertions(+), 12 deletions(-) diff --git a/model_specs/breeze_tts.json b/model_specs/breeze_tts.json index c5299e703..addd9b4e3 100644 --- a/model_specs/breeze_tts.json +++ b/model_specs/breeze_tts.json @@ -7,7 +7,8 @@ "status": "supported", "tasks": [ "tts", - "clone" + "clone", + "design" ], "modes": [ "offline" @@ -23,6 +24,9 @@ }, "dependencies": [], "capabilities": { + "design": [ + "voice_design" + ], "tts": [ "style_control" ], @@ -174,7 +178,7 @@ "strip_prefix": "Breeze-TTS-2-GGUF", "download": { "kind": "huggingface_snapshot", - "repo": "audio-cpp/Breeze-TTS-2-GGUF", + "repo": "audio-cpp/audio.cpp-gguf", "revision": "main", "gated": false } @@ -192,7 +196,7 @@ "strip_prefix": "Breeze-TTS-2-GGUF", "download": { "kind": "huggingface_snapshot", - "repo": "audio-cpp/Breeze-TTS-2-GGUF", + "repo": "audio-cpp/audio.cpp-gguf", "revision": "main", "gated": false } @@ -203,6 +207,7 @@ "tags": [ "TTS", "Clone", + "Design", "GGUF" ], "docs": [ diff --git a/model_specs/cosyvoice3.json b/model_specs/cosyvoice3.json index 5efae24ea..ed01ff8e6 100644 --- a/model_specs/cosyvoice3.json +++ b/model_specs/cosyvoice3.json @@ -192,7 +192,7 @@ "strip_prefix": "CosyVoice3-GGUF", "download": { "kind": "huggingface_snapshot", - "repo": "audio-cpp/CosyVoice3-GGUF", + "repo": "audio-cpp/audio.cpp-gguf", "revision": "main", "gated": false } @@ -210,7 +210,7 @@ "strip_prefix": "CosyVoice3-GGUF", "download": { "kind": "huggingface_snapshot", - "repo": "audio-cpp/CosyVoice3-GGUF", + "repo": "audio-cpp/audio.cpp-gguf", "revision": "main", "gated": false } diff --git a/src/models/breeze_tts/session.cpp b/src/models/breeze_tts/session.cpp index 5ac49a0a4..61a9986f7 100644 --- a/src/models/breeze_tts/session.cpp +++ b/src/models/breeze_tts/session.cpp @@ -101,8 +101,10 @@ BreezeTTSSession::BreezeTTSSession( contract_(require_contract(std::move(contract))), reference_cache_(reference_cache_slots_from_options(options)) { runtime::validate_spec_backed_session_options(options, *contract_, kFamily, kModelName); - if (task_.task != runtime::VoiceTaskKind::Tts && task_.task != runtime::VoiceTaskKind::VoiceCloning) { - throw std::runtime_error("BreezeTTS supports tts and clone tasks"); + if (task_.task != runtime::VoiceTaskKind::Tts && + task_.task != runtime::VoiceTaskKind::VoiceCloning && + task_.task != runtime::VoiceTaskKind::VoiceDesign) { + throw std::runtime_error("BreezeTTS supports tts, clone, and voice design tasks"); } if (task_.mode != runtime::RunMode::Offline) { throw std::runtime_error("BreezeTTS supports offline sessions"); diff --git a/webui/configs/model_params.json b/webui/configs/model_params.json index 3b595cf3d..3006a8279 100644 --- a/webui/configs/model_params.json +++ b/webui/configs/model_params.json @@ -109,6 +109,16 @@ {"name": "top_p", "type": "slider", "label": "top_p", "default": 1.0, "minimum": 0.0, "maximum": 1.0, "step": 0.01} ], + "breeze-tts": [ + {"name": "text_chunk_size", "type": "number", "label": "text_chunk_size", "default": 600, "minimum": 1, "step": 1, "precision": 0}, + {"name": "text_chunk_mode", "type": "choice", "label": "text_chunk_mode", "default": "default", "choices": ["default", "tag_aware", "japanese", "endline"]}, + {"name": "guidance_scale", "type": "slider", "label": "guidance_scale", "default": 1.0, "minimum": 0.0, "maximum": 10.0, "step": 0.1}, + {"name": "temperature", "type": "slider", "label": "temperature", "default": 0.9, "minimum": 0.0, "maximum": 2.0, "step": 0.05}, + {"name": "depth_temperature", "type": "slider", "label": "depth_temperature", "default": 0.9, "minimum": 0.0, "maximum": 2.0, "step": 0.05}, + {"name": "top_k", "type": "number", "label": "top_k", "default": 50, "minimum": 0, "step": 1, "precision": 0}, + {"name": "top_p", "type": "slider", "label": "top_p", "default": 1.0, "minimum": 0.0, "maximum": 1.0, "step": 0.01} + ], + "cosyvoice3": [ {"name": "template_name", "type": "choice", "label": "template_name", "default": "zero_shot", "choices": ["zero_shot", "cross_lingual", "instruct"]}, {"name": "instruction", "type": "text", "label": "instruction", "label_en": "Instruction", "default": "", "placeholder": "Used by template_name=instruct."}, diff --git a/webui/configs/models_catalog.json b/webui/configs/models_catalog.json index 4d2e49d98..38e284848 100644 --- a/webui/configs/models_catalog.json +++ b/webui/configs/models_catalog.json @@ -14,7 +14,7 @@ { "id": "qwen3-tts", "display_name": "Qwen3-TTS 0.6B (tts)", "family": "qwen3_tts", "path": "models/Qwen3-TTS-12Hz-0.6B-Base", "task": "tts", "mode": "offline", "download_id": "qwen3_tts_0_6b_base", "min_vram_gb": 5 }, { "id": "qwen3-tts-1.7b", "display_name": "Qwen3-TTS 1.7B Base (tts)", "family": "qwen3_tts", "path": "models/Qwen3-TTS-12Hz-1.7B-Base", "task": "tts", "mode": "offline", "download_id": "qwen3_tts_1_7b_base", "min_vram_gb": 8 }, { "id": "qwen3-tts-1.7b-custom", "display_name": "Qwen3-TTS 1.7B CustomVoice (tts)", "family": "qwen3_tts", "path": "models/Qwen3-TTS-12Hz-1.7B-CustomVoice", "task": "tts", "mode": "offline", "download_id": "qwen3_tts_1_7b_custom_voice", "min_vram_gb": 8 }, - { "id": "breeze-tts", "display_name": "BreezeTTS 2 VoiceDesign", "family": "breeze_tts", "path": "models/Breeze-TTS-2-GGUF/breeze-tts-2-q8_0.gguf", "task": "tts", "mode": "offline", "download_id": "breeze_tts_2_q8_0", "min_vram_gb": 8, + { "id": "breeze-tts", "display_name": "BreezeTTS 2 VoiceDesign", "family": "breeze_tts", "path": "models/Breeze-TTS-2-GGUF/breeze-tts-2-q8_0.gguf", "task": "vdes", "mode": "offline", "download_id": "breeze_tts_2_q8_0", "min_vram_gb": 8, "input_hint_en": "**BreezeTTS 2 VoiceDesign**: enter text and describe the target voice in Model parameters. No reference voice is required." }, { "id": "breeze-tts-clone", "display_name": "BreezeTTS 2 Clone", "family": "breeze_tts", "path": "models/Breeze-TTS-2-GGUF/breeze-tts-2-q8_0.gguf", "task": "clon", "mode": "offline", "download_id": "breeze_tts_2_q8_0", "min_vram_gb": 8, "input_hint_en": "**BreezeTTS 2 Clone**: upload a reference voice and provide the matching reference transcript." }, diff --git a/webui/native/dist/index.html b/webui/native/dist/index.html index ac74c900b..d6534c0a2 100644 --- a/webui/native/dist/index.html +++ b/webui/native/dist/index.html @@ -31,15 +31,15 @@
From 6c2c6aad834176bc17ae515708318fbcb38f0302 Mon Sep 17 00:00:00 2001 From: Fraser <42195943+FraserHum@users.noreply.github.com> Date: Sun, 30 Aug 2026 10:28:03 +1200 Subject: [PATCH 05/10] perf(breeze_tts): pre-allocate depth projection & generation staging buffers (#348) --- src/models/breeze_tts/generator.cpp | 216 +++++++++++++++++++--------- 1 file changed, 150 insertions(+), 66 deletions(-) diff --git a/src/models/breeze_tts/generator.cpp b/src/models/breeze_tts/generator.cpp index e9cc22478..7d8483877 100644 --- a/src/models/breeze_tts/generator.cpp +++ b/src/models/breeze_tts/generator.cpp @@ -405,6 +405,7 @@ class BreezeDepthProjectionRuntime { for (int64_t codebook = 1; codebook < config.num_codebooks; ++codebook) { head_graphs_.push_back(build_head_graph(ctx, packed_heads, codebook)); } + head_paired_staging_.assign(static_cast(2 * vocab_), 0.0F); graph_buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), backend_); if (graph_buffer_ == nullptr) { throw std::runtime_error("failed to allocate BreezeTTS depth projection graphs"); @@ -422,8 +423,35 @@ class BreezeDepthProjectionRuntime { } } + void project_single(const float * hidden, float * out) const { + run_graph_direct(projector_single_, hidden, hidden_, out, depth_hidden_); + } + std::vector project_single(const std::vector & hidden) const { - return run_graph(projector_single_, hidden); + if (static_cast(hidden.size()) != hidden_) { + throw std::runtime_error("BreezeTTS depth projection input size mismatch"); + } + std::vector out(static_cast(depth_hidden_)); + project_single(hidden.data(), out.data()); + return out; + } + + void project_pair( + const float * cond_hidden, + const float * uncond_hidden, + float * out) const { + ggml_backend_tensor_set(projector_pair_.input, cond_hidden, 0, static_cast(hidden_) * sizeof(float)); + ggml_backend_tensor_set( + projector_pair_.input, + uncond_hidden, + static_cast(hidden_) * sizeof(float), + static_cast(hidden_) * sizeof(float)); + core::set_backend_threads(backend_, threads_); + if (core::compute_backend_graph(backend_, projector_pair_.graph, nullptr, projector_pair_.label) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("BreezeTTS depth projection graph compute failed"); + } + ggml_backend_synchronize(backend_); + ggml_backend_tensor_get(projector_pair_.output, out, 0, static_cast(2 * depth_hidden_) * sizeof(float)); } std::vector project_pair( @@ -433,11 +461,38 @@ class BreezeDepthProjectionRuntime { static_cast(uncond_hidden.size()) != hidden_) { throw std::runtime_error("BreezeTTS depth projector pair input size mismatch"); } - std::vector input; - input.reserve(static_cast(2 * hidden_)); - input.insert(input.end(), cond_hidden.begin(), cond_hidden.end()); - input.insert(input.end(), uncond_hidden.begin(), uncond_hidden.end()); - return run_graph(projector_pair_, input); + std::vector out(static_cast(2 * depth_hidden_)); + project_pair(cond_hidden.data(), uncond_hidden.data(), out.data()); + return out; + } + + void logits_cfg( + const float * cond_hidden, + const float * uncond_hidden, + int64_t codebook, + float guidance_scale, + float * out) const { + if (codebook <= 0 || static_cast(codebook) > head_graphs_.size()) { + throw std::runtime_error("BreezeTTS depth codebook index is invalid"); + } + const auto & graph = head_graphs_[static_cast(codebook - 1)]; + ggml_backend_tensor_set(graph.input, cond_hidden, 0, static_cast(depth_hidden_) * sizeof(float)); + ggml_backend_tensor_set( + graph.input, + uncond_hidden, + static_cast(depth_hidden_) * sizeof(float), + static_cast(depth_hidden_) * sizeof(float)); + core::set_backend_threads(backend_, threads_); + if (core::compute_backend_graph(backend_, graph.graph, nullptr, graph.label) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("BreezeTTS depth projection graph compute failed"); + } + ggml_backend_synchronize(backend_); + ggml_backend_tensor_get(graph.output, head_paired_staging_.data(), 0, head_paired_staging_.size() * sizeof(float)); + const size_t vocab = static_cast(vocab_); + for (int64_t token = 0; token < vocab_; ++token) { + const size_t index = static_cast(token); + out[index] = head_paired_staging_[vocab + index] + guidance_scale * (head_paired_staging_[index] - head_paired_staging_[vocab + index]); + } } std::vector logits_cfg( @@ -445,24 +500,12 @@ class BreezeDepthProjectionRuntime { const std::vector & uncond_hidden, int64_t codebook, float guidance_scale) const { - if (codebook <= 0 || static_cast(codebook) > head_graphs_.size()) { - throw std::runtime_error("BreezeTTS depth codebook index is invalid"); - } if (static_cast(cond_hidden.size()) != depth_hidden_ || static_cast(uncond_hidden.size()) != depth_hidden_) { throw std::runtime_error("BreezeTTS depth head input size mismatch"); } - std::vector input; - input.reserve(static_cast(2 * depth_hidden_)); - input.insert(input.end(), cond_hidden.begin(), cond_hidden.end()); - input.insert(input.end(), uncond_hidden.begin(), uncond_hidden.end()); - const auto paired = run_graph(head_graphs_[static_cast(codebook - 1)], input); std::vector out(static_cast(vocab_)); - const size_t vocab = static_cast(vocab_); - for (int64_t token = 0; token < vocab_; ++token) { - const size_t index = static_cast(token); - out[index] = paired[vocab + index] + guidance_scale * (paired[index] - paired[vocab + index]); - } + logits_cfg(cond_hidden.data(), uncond_hidden.data(), codebook, guidance_scale, out.data()); return out; } @@ -518,18 +561,30 @@ class BreezeDepthProjectionRuntime { core::release_backend_graph_resources(backend_, graph.graph); } - std::vector run_graph(const Graph & graph, const std::vector & input) const { - if (static_cast(input.size()) != graph.input_size) { + void run_graph_direct( + const Graph & graph, + const float * input, + int64_t in_size, + float * output, + int64_t out_size) const { + if (in_size != graph.input_size || out_size != graph.output_size) { throw std::runtime_error("BreezeTTS depth projection input size mismatch"); } - ggml_backend_tensor_set(graph.input, input.data(), 0, input.size() * sizeof(float)); + ggml_backend_tensor_set(graph.input, input, 0, static_cast(in_size) * sizeof(float)); core::set_backend_threads(backend_, threads_); if (core::compute_backend_graph(backend_, graph.graph, nullptr, graph.label) != GGML_STATUS_SUCCESS) { throw std::runtime_error("BreezeTTS depth projection graph compute failed"); } ggml_backend_synchronize(backend_); + ggml_backend_tensor_get(graph.output, output, 0, static_cast(out_size) * sizeof(float)); + } + + std::vector run_graph(const Graph & graph, const std::vector & input) const { + if (static_cast(input.size()) != graph.input_size) { + throw std::runtime_error("BreezeTTS depth projection input size mismatch"); + } std::vector output(static_cast(graph.output_size)); - ggml_backend_tensor_get(graph.output, output.data(), 0, output.size() * sizeof(float)); + run_graph_direct(graph, input.data(), static_cast(input.size()), output.data(), graph.output_size); return output; } @@ -543,6 +598,7 @@ class BreezeDepthProjectionRuntime { Graph projector_single_; Graph projector_pair_; std::vector head_graphs_; + mutable std::vector head_paired_staging_; }; } // namespace @@ -595,6 +651,14 @@ struct BreezeGeneratorRuntime::Impl { weight_context_bytes, storage_type, storage_type); + depth_first_embed_staging_.assign(static_cast(config.depth_hidden_size), 0.0F); + depth_projected_pair_staging_.assign(static_cast(2 * config.depth_hidden_size), 0.0F); + depth_prefill_staging_.assign(static_cast(4 * config.depth_hidden_size), 0.0F); + depth_next_embed_staging_.assign(static_cast(config.depth_hidden_size), 0.0F); + depth_next_pair_staging_.assign(static_cast(2 * config.depth_hidden_size), 0.0F); + depth_logits_staging_.assign(static_cast(config.vocab_size), 0.0F); + depth_cond_hidden_now_.assign(static_cast(config.depth_hidden_size), 0.0F); + depth_uncond_hidden_now_.assign(static_cast(config.depth_hidden_size), 0.0F); } std::vector merge_prompt(const BreezePromptBranch & branch, const std::vector & reference_codes) { @@ -674,44 +738,49 @@ struct BreezeGeneratorRuntime::Impl { frame.reserve(static_cast(config.num_codebooks)); frame.push_back(first_token); - const auto project_audio_embedding_row = [&](int64_t row) { + const size_t depth_hidden_size = static_cast(config.depth_hidden_size); + const size_t depth_hidden_bytes = depth_hidden_size * sizeof(float); + + const auto project_audio_embedding_row = [&](int64_t row, float * out) { const int64_t rows = config.num_codebooks * config.vocab_size; if (row < 0 || row >= rows) { throw std::runtime_error("BreezeTTS embedding row is outside table"); } const size_t begin = static_cast(row * config.hidden_size); - std::vector embedding( - weights->audio_embedding.begin() + static_cast(begin), - weights->audio_embedding.begin() + static_cast(begin + static_cast(config.hidden_size))); - return depth_projection->project_single(embedding); + depth_projection->project_single( + weights->audio_embedding.data() + begin, + out); }; - const auto first_embed = project_audio_embedding_row(first_token); - const auto projected = depth_projection->project_pair(cond_hidden, uncond_hidden); - const auto split = projected.begin() + static_cast(config.depth_hidden_size); - std::vector cond_prefill; - cond_prefill.reserve(static_cast(2 * config.depth_hidden_size)); - cond_prefill.insert(cond_prefill.end(), projected.begin(), split); - cond_prefill.insert(cond_prefill.end(), first_embed.begin(), first_embed.end()); - std::vector uncond_prefill; - uncond_prefill.reserve(static_cast(2 * config.depth_hidden_size)); - uncond_prefill.insert(uncond_prefill.end(), split, projected.end()); - uncond_prefill.insert(uncond_prefill.end(), first_embed.begin(), first_embed.end()); - - std::vector prefill; - prefill.reserve(static_cast(4 * config.depth_hidden_size)); - prefill.insert(prefill.end(), cond_prefill.begin(), cond_prefill.end()); - prefill.insert(prefill.end(), uncond_prefill.begin(), uncond_prefill.end()); - auto depth = depth_pair->prefill_embeddings_batched(prefill, 2, 2); + project_audio_embedding_row(first_token, depth_first_embed_staging_.data()); + depth_projection->project_pair( + cond_hidden.data(), + uncond_hidden.data(), + depth_projected_pair_staging_.data()); + + std::memcpy(depth_prefill_staging_.data(), + depth_projected_pair_staging_.data(), + depth_hidden_bytes); + std::memcpy(depth_prefill_staging_.data() + depth_hidden_size, + depth_first_embed_staging_.data(), + depth_hidden_bytes); + std::memcpy(depth_prefill_staging_.data() + 2 * depth_hidden_size, + depth_projected_pair_staging_.data() + depth_hidden_size, + depth_hidden_bytes); + std::memcpy(depth_prefill_staging_.data() + 3 * depth_hidden_size, + depth_first_embed_staging_.data(), + depth_hidden_bytes); + + auto depth = depth_pair->prefill_embeddings_batched(depth_prefill_staging_, 2, 2); if (static_cast(depth.hidden.size()) != 2 * config.depth_hidden_size) { throw std::runtime_error("BreezeTTS batched depth prefill hidden size mismatch"); } - std::vector cond_hidden_now( - depth.hidden.begin(), - depth.hidden.begin() + static_cast(config.depth_hidden_size)); - std::vector uncond_hidden_now( - depth.hidden.begin() + static_cast(config.depth_hidden_size), - depth.hidden.end()); + std::memcpy(depth_cond_hidden_now_.data(), + depth.hidden.data(), + depth_hidden_bytes); + std::memcpy(depth_uncond_hidden_now_.data(), + depth.hidden.data() + depth_hidden_size, + depth_hidden_bytes); depth_pair->start_decode_embeddings_batched(depth.state, config.num_codebooks + 1); sampling::HfSamplingOptions options; @@ -721,10 +790,15 @@ struct BreezeGeneratorRuntime::Impl { options.top_p = request.top_p; options.min_tokens_to_keep = 1; for (int64_t codebook = 1; codebook < config.num_codebooks; ++codebook) { - auto logits = depth_projection->logits_cfg(cond_hidden_now, uncond_hidden_now, codebook, request.guidance_scale); - suppress_reserved(logits, kCodecCodebookSize, config.vocab_size); + depth_projection->logits_cfg( + depth_cond_hidden_now_.data(), + depth_uncond_hidden_now_.data(), + codebook, + request.guidance_scale, + depth_logits_staging_.data()); + suppress_reserved(depth_logits_staging_, kCodecCodebookSize, config.vocab_size); const int32_t token = sample_logits( - std::move(logits), + depth_logits_staging_, {}, options, scratch, @@ -736,21 +810,23 @@ struct BreezeGeneratorRuntime::Impl { "BreezeTTS depth sampler"); frame.push_back(token); if (codebook + 1 < config.num_codebooks) { - const auto next = project_audio_embedding_row(codebook * config.vocab_size + token); - std::vector next_pair; - next_pair.reserve(static_cast(2 * config.depth_hidden_size)); - next_pair.insert(next_pair.end(), next.begin(), next.end()); - next_pair.insert(next_pair.end(), next.begin(), next.end()); - const auto step = depth_pair->decode_embeddings_batched(next_pair, 2); + project_audio_embedding_row(codebook * config.vocab_size + token, depth_next_embed_staging_.data()); + std::memcpy(depth_next_pair_staging_.data(), + depth_next_embed_staging_.data(), + depth_hidden_bytes); + std::memcpy(depth_next_pair_staging_.data() + depth_hidden_size, + depth_next_embed_staging_.data(), + depth_hidden_bytes); + const auto step = depth_pair->decode_embeddings_batched(depth_next_pair_staging_, 2); if (static_cast(step.hidden.size()) != 2 * config.depth_hidden_size) { throw std::runtime_error("BreezeTTS batched depth decode hidden size mismatch"); } - cond_hidden_now.assign( - step.hidden.begin(), - step.hidden.begin() + static_cast(config.depth_hidden_size)); - uncond_hidden_now.assign( - step.hidden.begin() + static_cast(config.depth_hidden_size), - step.hidden.end()); + std::memcpy(depth_cond_hidden_now_.data(), + step.hidden.data(), + depth_hidden_bytes); + std::memcpy(depth_uncond_hidden_now_.data(), + step.hidden.data() + depth_hidden_size, + depth_hidden_bytes); } } return frame; @@ -933,6 +1009,14 @@ struct BreezeGeneratorRuntime::Impl { std::unique_ptr depth_projection; std::unique_ptr speech_encoder; std::unique_ptr speech_decoder; + std::vector depth_first_embed_staging_; + std::vector depth_projected_pair_staging_; + std::vector depth_prefill_staging_; + std::vector depth_next_embed_staging_; + std::vector depth_next_pair_staging_; + std::vector depth_logits_staging_; + std::vector depth_cond_hidden_now_; + std::vector depth_uncond_hidden_now_; }; BreezeGeneratorRuntime::BreezeGeneratorRuntime( From eb2cbd3878a9fb4f8bea9dd90149d02f338010af Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Sat, 29 Aug 2026 20:08:52 -0400 Subject: [PATCH 06/10] Add BreezeTTS streaming mode --- include/engine/models/breeze_tts/session.h | 22 +++- model_specs/breeze_tts.json | 3 +- src/models/breeze_tts/session.cpp | 134 ++++++++++++++++++--- 3 files changed, 138 insertions(+), 21 deletions(-) diff --git a/include/engine/models/breeze_tts/session.h b/include/engine/models/breeze_tts/session.h index 31a765926..8c839da0a 100644 --- a/include/engine/models/breeze_tts/session.h +++ b/include/engine/models/breeze_tts/session.h @@ -11,16 +11,19 @@ #include #include #include +#include namespace engine::models::breeze_tts { class BreezeGeneratorRuntime; +struct BreezeGenerationRequest; std::shared_ptr make_breeze_tts_loader(); class BreezeTTSSession final : public engine::runtime::RuntimeSessionBase - , public engine::runtime::IOfflineVoiceTaskSession { + , public engine::runtime::IOfflineVoiceTaskSession + , public engine::runtime::IStreamingVoiceTaskSession { public: BreezeTTSSession( engine::runtime::TaskSpec task, @@ -34,6 +37,14 @@ class BreezeTTSSession final engine::runtime::RunMode run_mode() const override; void prepare(const engine::runtime::SessionPreparationRequest & request) override; engine::runtime::TaskResult run(const engine::runtime::TaskRequest & request) override; + engine::runtime::StreamingPolicy streaming_policy() const override; + void start_stream(const engine::runtime::TaskRequest & request) override; + std::optional next_stream_event() override; + void set_stream_event_sink(engine::runtime::StreamEventCallback sink) override; + engine::runtime::TaskResult finish_stream() override; + void reset() override; + engine::runtime::StreamEvent process_audio_chunk(const engine::runtime::AudioChunk & chunk) override; + engine::runtime::TaskResult finalize() override; private: struct ReferenceCacheKey { @@ -52,6 +63,10 @@ class BreezeTTSSession final }; BreezeSpeechCodes resolve_reference_codes(const engine::runtime::AudioBuffer & audio); + BreezeGenerationRequest build_generation_request( + const engine::runtime::TaskRequest & request, + const std::optional & reference_codes, + size_t chunk_index) const; engine::runtime::TaskSpec task_; std::shared_ptr assets_; @@ -59,6 +74,11 @@ class BreezeTTSSession final std::unique_ptr generator_; engine::runtime::CacheSlots reference_cache_; std::optional uncached_reference_; + std::vector stream_chunk_requests_; + std::optional stream_reference_codes_; + engine::runtime::AudioBuffer stream_merged_audio_; + size_t stream_chunk_index_ = 0; + bool stream_started_ = false; }; } // namespace engine::models::breeze_tts diff --git a/model_specs/breeze_tts.json b/model_specs/breeze_tts.json index addd9b4e3..95a0a2cc5 100644 --- a/model_specs/breeze_tts.json +++ b/model_specs/breeze_tts.json @@ -11,7 +11,8 @@ "design" ], "modes": [ - "offline" + "offline", + "streaming" ], "languages": [ "zh", diff --git a/src/models/breeze_tts/session.cpp b/src/models/breeze_tts/session.cpp index 61a9986f7..264d67bd1 100644 --- a/src/models/breeze_tts/session.cpp +++ b/src/models/breeze_tts/session.cpp @@ -106,8 +106,9 @@ BreezeTTSSession::BreezeTTSSession( task_.task != runtime::VoiceTaskKind::VoiceDesign) { throw std::runtime_error("BreezeTTS supports tts, clone, and voice design tasks"); } - if (task_.mode != runtime::RunMode::Offline) { - throw std::runtime_error("BreezeTTS supports offline sessions"); + if (task_.mode != runtime::RunMode::Offline && + task_.mode != runtime::RunMode::Streaming) { + throw std::runtime_error("BreezeTTS supports offline and streaming sessions"); } using T = engine::assets::TensorStorageType; const auto storage_type = runtime::parse_tensor_storage_option( @@ -190,6 +191,9 @@ runtime::TaskResult BreezeTTSSession::run(const runtime::TaskRequest & request) const auto wall_start = std::chrono::steady_clock::now(); runtime::validate_spec_backed_request_options(request.options, *contract_, kModelName); require_prepared("BreezeTTS run"); + if (task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("BreezeTTS run requires an offline session"); + } if (!request.text_input.has_value() || request.text_input->text.empty()) { throw std::runtime_error("BreezeTTS requires text input"); } @@ -209,23 +213,7 @@ runtime::TaskResult BreezeTTSSession::run(const runtime::TaskRequest & request) reference_codes = resolve_reference_codes(*request.voice->speaker->audio); } for (size_t index = 0; index < chunks.size(); ++index) { - const auto & chunk = chunks[index]; - BreezeGenerationRequest generation; - generation.text = chunk.text_input->text; - generation.instruction = runtime::find_option(chunk.options, {"instruction"}).value_or(""); - generation.reference_text = runtime::find_option(chunk.options, {"reference_text"}).value_or(""); - generation.guidance_scale = runtime::parse_positive_finite_float_option(chunk.options, {"guidance_scale"}).value_or(generation.guidance_scale); - generation.temperature = runtime::parse_positive_finite_float_option(chunk.options, {"temperature"}).value_or(generation.temperature); - generation.depth_temperature = runtime::parse_positive_finite_float_option(chunk.options, {"depth_temperature"}).value_or(generation.depth_temperature); - generation.top_k = runtime::parse_i64_option(chunk.options, {"top_k"}).value_or(generation.top_k); - generation.top_p = runtime::parse_positive_finite_float_option(chunk.options, {"top_p"}).value_or(generation.top_p); - generation.max_tokens = runtime::parse_positive_i64_option(chunk.options, {"max_tokens"}, generation.max_tokens); - generation.seed = runtime::parse_u64_option(chunk.options, {"seed"}).value_or(generation.seed); - generation.reference_codes = reference_codes; - if (index > 0) { - ++generation.seed; - } - runtime::append_audio_buffer(merged, generator_->generate(generation)); + runtime::append_audio_buffer(merged, generator_->generate(build_generation_request(chunks[index], reference_codes, index))); } runtime::TaskResult result; result.audio_output = std::move(merged); @@ -233,6 +221,114 @@ runtime::TaskResult BreezeTTSSession::run(const runtime::TaskRequest & request) return result; } +runtime::StreamingPolicy BreezeTTSSession::streaming_policy() const { + runtime::StreamingPolicy policy; + policy.input = runtime::StreamingInputKind::None; + policy.output = runtime::StreamingOutputKind::PullEvents; + return policy; +} + +void BreezeTTSSession::start_stream(const runtime::TaskRequest & request) { + runtime::validate_spec_backed_request_options(request.options, *contract_, kModelName); + require_prepared("BreezeTTS streaming"); + if (task_.mode != runtime::RunMode::Streaming) { + throw std::runtime_error("BreezeTTS start_stream requires a streaming session"); + } + if (!request.text_input.has_value() || request.text_input->text.empty()) { + throw std::runtime_error("BreezeTTS streaming requires text input"); + } + reset(); + stream_chunk_requests_ = split_request(request); + const int64_t text_chunk_size = + engine::text::parse_text_chunk_size_override(request.options).value_or(kDefaultTextChunkSize); + const auto text_chunk_mode = + engine::text::parse_text_chunk_mode_override(request.options).value_or(engine::text::TextChunkMode::Default); + engine::debug::trace_log_scalar("breeze_tts.streaming.text_chunk_mode", engine::text::text_chunk_mode_name(text_chunk_mode)); + engine::debug::trace_log_scalar("breeze_tts.streaming.text_chunk_size", text_chunk_size); + engine::debug::trace_log_scalar("breeze_tts.streaming.text.chunk_count", static_cast(stream_chunk_requests_.size())); + if (request.voice.has_value() && + request.voice->speaker.has_value() && + request.voice->speaker->audio.has_value()) { + stream_reference_codes_ = resolve_reference_codes(*request.voice->speaker->audio); + } + stream_started_ = true; +} + +std::optional BreezeTTSSession::next_stream_event() { + if (!stream_started_) { + throw std::runtime_error("BreezeTTS streaming has not been started"); + } + if (stream_chunk_index_ >= stream_chunk_requests_.size()) { + return std::nullopt; + } + const size_t chunk_index = stream_chunk_index_++; + auto chunk_audio = generator_->generate( + build_generation_request(stream_chunk_requests_[chunk_index], stream_reference_codes_, chunk_index)); + runtime::append_audio_buffer(stream_merged_audio_, chunk_audio); + runtime::StreamEvent event; + event.named_audio_outputs.push_back({ + "chunk_" + std::to_string(chunk_index), + std::move(chunk_audio), + {}, + }); + return event; +} + +void BreezeTTSSession::set_stream_event_sink(runtime::StreamEventCallback sink) { + (void)sink; +} + +runtime::TaskResult BreezeTTSSession::finish_stream() { + if (!stream_started_) { + throw std::runtime_error("BreezeTTS streaming has not been started"); + } + while (next_stream_event().has_value()) { + } + runtime::TaskResult result; + result.audio_output = std::move(stream_merged_audio_); + reset(); + return result; +} + +void BreezeTTSSession::reset() { + stream_chunk_requests_.clear(); + stream_reference_codes_.reset(); + stream_merged_audio_ = runtime::AudioBuffer{}; + stream_chunk_index_ = 0; + stream_started_ = false; +} + +runtime::StreamEvent BreezeTTSSession::process_audio_chunk(const runtime::AudioChunk & chunk) { + (void)chunk; + throw std::runtime_error("BreezeTTS streaming does not consume audio chunks"); +} + +runtime::TaskResult BreezeTTSSession::finalize() { + return finish_stream(); +} + +BreezeGenerationRequest BreezeTTSSession::build_generation_request( + const runtime::TaskRequest & request, + const std::optional & reference_codes, + size_t chunk_index) const { + BreezeGenerationRequest generation; + generation.text = request.text_input->text; + generation.instruction = runtime::find_option(request.options, {"instruction"}).value_or(""); + generation.reference_text = runtime::find_option(request.options, {"reference_text"}).value_or(""); + generation.guidance_scale = runtime::parse_positive_finite_float_option(request.options, {"guidance_scale"}).value_or(generation.guidance_scale); + generation.temperature = runtime::parse_positive_finite_float_option(request.options, {"temperature"}).value_or(generation.temperature); + generation.depth_temperature = runtime::parse_positive_finite_float_option(request.options, {"depth_temperature"}).value_or(generation.depth_temperature); + generation.top_k = runtime::parse_i64_option(request.options, {"top_k"}).value_or(generation.top_k); + generation.top_p = runtime::parse_positive_finite_float_option(request.options, {"top_p"}).value_or(generation.top_p); + generation.max_tokens = runtime::parse_positive_i64_option(request.options, {"max_tokens"}, generation.max_tokens); + generation.seed = runtime::parse_u64_option(request.options, {"seed"}).value_or(generation.seed); + generation.reference_codes = reference_codes; + if (chunk_index > 0) { + ++generation.seed; + } + return generation; +} + bool BreezeTTSSession::ReferenceCacheKeyEqual::operator()( const ReferenceCacheKey & lhs, const ReferenceCacheKey & rhs) const noexcept { From 820cddce5fa891d4a7935e479406c804d9476959 Mon Sep 17 00:00:00 2001 From: Laika <2432896620@qq.com> Date: Thu, 3 Sep 2026 10:42:38 +0800 Subject: [PATCH 07/10] Breeze-TTS 2 performance: weight packing + fused bf16 rounding (#393) * breeze: pack qkv and gate/up projection weights, document weight_type * ggml, qwen_decoder: fuse bf16 activation rounding into a single kernel nsys on the 2080 Ti shows the f32 -> bf16 -> f32 cast pairs behind every activation rounding point cost ~19% of GPU time on the bf16 path and ~29% on the q4_k path (144k tiny cpy kernels per 20-token run), because ggml has no fused round-to-bf16 op and the CUDA backend runs each ggml_cast as a separate kernel. Add GGML_UNARY_OP_ROUND_BF16 (CPU + CUDA implementations; HIP shares the ggml-cuda sources) that rounds f32 values to bf16 precision in one pass, bit-identical to the cast round trip (same __float2bfloat16 / __bfloat162float sequence as cpy). The qwen decoder activation cast policy gains a fused_round flag, enabled for CUDA/HIP only; Vulkan keeps the round trip. Non-contiguous views also keep the round trip, as the unary op requires contiguous input. Verified on RTX 2080 Ti with Breeze-TTS 2: generated codes are bit-identical to the round trip build in all four test cases (bf16/q4_k, fixed 100-token case and both Chinese regression prompts). RTF on the fixed 100-token case: bf16 0.760 -> 0.695, q4_k 0.484 -> 0.419; Chinese regression q4_k 0.861 -> 0.736 (short) and 0.556 -> 0.464 (long). * ggml, breeze: support row-strided inputs in fused bf16 rounding Rounding points fed by non-contiguous views (rope/cache paths) still used the cast round trip: a strided f32 -> bf16 cpy plus a contiguous bf16 -> f32 cpy, ~10% of GPU time on the q4_k path. Add a row-strided variant of the round_bf16 kernel (dst is contiguous by construction) and relax the backend/framework gates from ggml_is_contiguous to ggml_is_contiguous_rows, so those points fuse too. Codes remain bit-identical in all four test cases. RTF on RTX 2080 Ti, q4_k: 100-token 0.419 -> 0.399, Chinese long 0.464 -> 0.441; bf16 100-token 0.695 -> 0.677. * ggml-cuda: allow CUDA graphs on pre-Ampere GPUs via GGML_CUDA_GRAPHS_PRE_AMPERE Upstream disables CUDA graphs below sm_80. Keep that default, but add an env-var escape hatch so pre-Ampere behavior can be tested without recompiling. On the RTX 2080 Ti (sm_75) Breeze-TTS 2 decode the graphs do capture and replay correctly (bit-identical codes), but RTF is neutral to slightly worse (0.399 without vs 0.408 with on the q4_k 100-token case), so the upstream default stands for this workload. * ggml, breeze: generalize fused bf16 rounding to f16/bf16 inputs ggml_round_bf16 now always produces a contiguous f32 result regardless of input type (f32/f16/bf16), matching the cast round trip bit for bit: bf16 input is already rounded so the op degenerates to an exact widening, f16 input rounds through bf16 and widens, both landing on the same real values as cast -> bf16 -> cast -> f32. This fixes a HIP crash where rounding points fed by the bf16 KV cache hit an f32/f16-only assert in the unary kernel, and recovers the fusion for f16 inputs (CUDA f16 KV cache paths) that the previous f32-only gate skipped. The activation cast no longer needs per-type special cases. Verified bit-identical codes in all 8 cases (CUDA + HIP x q4_k/bf16 x 100-token + 2 Chinese regression prompts). RTF, q4_k 100-token: CUDA 0.417 -> 0.405, HIP 0.73 (unchanged); HIP q4_k vs pre-fusion baseline: 0.84 -> 0.73, long 0.92 -> 0.80, bf16 1.50 -> 1.37. * conv_transpose1d: enable col2im fast path on Vulkan The col2im path (mul_mat + ggml_col2im_1d) only ran on CUDA/HIP/Metal; Vulkan fell back to ggml_conv_transpose_1d, whose Vulkan shader is a naive per-element kernel. All ops the col2im path needs are already supported by the Vulkan backend, including col2im_1d (f32/f16 pipelines). Breeze-TTS 2 speech decoder on Radeon 8060S: 190 ms -> 98 ms; greedy output codes identical to the generic path (wav correlation 0.99998). * breeze: skip the unconditional branch when guidance_scale == 1 CFG combines logits as uncond + scale * (cond - uncond), which is exactly cond at the default guidance_scale of 1. Running the unconditional backbone there is pure waste: skipping it removes half the backbone prefill and decode work. The depth projector's logits_cfg also gets a scale == 1 shortcut that copies the conditional half directly, avoiding an inexact uncond + 1 * (cond - uncond) round trip. guidance_scale = 0 (pure unconditional) is now accepted as well. On an RTX 2080 Ti, Breeze-TTS 2 fixed 100-token case, native weights: RTF 0.705 -> 0.605; greedy output is bit-identical with and without the skip. guidance_scale = 1.5 still runs the full CFG path unchanged. * ggml-vulkan: add bf16<->f32/f16 cpy pipelines * breeze: round activations to bf16 on GPU backends to match reference The official Breeze-TTS 2 inference runs the backbone and depth decoder with bf16 activations and a bf16 KV cache. A pure fp32 AR loop drifts into degenerate trajectories on some prompts (mispronounced tokens, repetition collapse, missing EOS), so round activations to bf16 at every op boundary via the qwen decoder activation_cast policy, mirroring the reference torch bf16 semantics. CUDA/HIP use the fused round-to-bf16 op; Vulkan uses the cast round trip. KV cache stays F16 on CUDA and Vulkan: bf16 flash attention is only accelerated with native bf16 MMA (sm_80+) and is ~3x slower on older GPUs. HIP uses a bf16 KV cache like the reference. (Ported onto the perf branch; fused_round requires the ROUND_BF16 op from the preceding commits.) --- docs/models/breeze_tts.md | 10 + external/ggml/include/ggml.h | 7 + external/ggml/src/ggml-cpu/ggml-cpu.c | 1 + external/ggml/src/ggml-cpu/ops.cpp | 4 + external/ggml/src/ggml-cpu/unary-ops.cpp | 8 + external/ggml/src/ggml-cpu/unary-ops.h | 1 + external/ggml/src/ggml-cuda/ggml-cuda.cu | 16 +- external/ggml/src/ggml-cuda/unary.cu | 99 +++++++++ external/ggml/src/ggml-cuda/unary.cuh | 2 + external/ggml/src/ggml-vulkan/ggml-vulkan.cpp | 35 ++- .../vulkan-shaders/contig_copy.comp | 10 +- .../src/ggml-vulkan/vulkan-shaders/copy.comp | 5 +- .../vulkan-shaders/vulkan-shaders-gen.cpp | 10 +- external/ggml/src/ggml.c | 23 +- .../modules/transformers/qwen_decoder.h | 4 + src/framework/modules/conv_modules.cpp | 3 +- .../modules/transformers/qwen_decoder.cpp | 10 + src/models/breeze_tts/generator.cpp | 200 +++++++++++++++--- src/models/breeze_tts/session.cpp | 5 +- 19 files changed, 407 insertions(+), 46 deletions(-) diff --git a/docs/models/breeze_tts.md b/docs/models/breeze_tts.md index e13c7cdf4..9aca60d01 100644 --- a/docs/models/breeze_tts.md +++ b/docs/models/breeze_tts.md @@ -67,3 +67,13 @@ audiocpp_cli \ | `--request-option top_p=` | `0..1` | `1.0` | Top-p sampling limit. | | `--request-option seed=` | integer >= 0 | `0` | Generation seed. | | `--session-option breeze_tts.reference_cache_slots=` | integer >= 0 | `1` | Prepared reference-audio cache slots. | +| `--session-option weight_type=` | `native`, `f32`, `f16`, `bf16`, `q8_0`, `q4_0`, `q4_k` | `native` | Weight storage type; quantized types convert at load time from the BF16 package. | + +Quantized weight storage is the largest measured speedup and applies to CUDA +and HIP alike: `q8_0` cut the fixed 100-token regression case from RTF ~1.5 to +~0.95 on gfx1151 and from ~0.77 to ~0.56 on an RTX 2080 Ti, and `q4_k` reached +~0.84 / ~0.49 respectively, with no audible quality regression in the Chinese +voice-design regression cases. Counter to intuition, fp32 is the one +configuration known to be *worse* for this model (mispronunciations and +runaway repetition), because the model is trained and tuned in bf16. + diff --git a/external/ggml/include/ggml.h b/external/ggml/include/ggml.h index 0de79aed5..1ba7dac76 100644 --- a/external/ggml/include/ggml.h +++ b/external/ggml/include/ggml.h @@ -615,6 +615,7 @@ extern "C" { GGML_UNARY_OP_CEIL, GGML_UNARY_OP_ROUND, GGML_UNARY_OP_TRUNC, + GGML_UNARY_OP_ROUND_BF16, GGML_UNARY_OP_COUNT, }; @@ -1258,6 +1259,12 @@ extern "C" { struct ggml_context * ctx, struct ggml_tensor * a); + // Rounds each element to bf16 precision, stored as f32. Equivalent to a + // cast f32 -> bf16 -> f32 round trip, but fused into a single op. + GGML_API struct ggml_tensor * ggml_round_bf16( + struct ggml_context * ctx, + struct ggml_tensor * a); + // xIELU activation function diff --git a/external/ggml/src/ggml-cpu/ggml-cpu.c b/external/ggml/src/ggml-cpu/ggml-cpu.c index d9ec09939..91e53b5e2 100644 --- a/external/ggml/src/ggml-cpu/ggml-cpu.c +++ b/external/ggml/src/ggml-cpu/ggml-cpu.c @@ -2260,6 +2260,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_UNARY_OP_CEIL: case GGML_UNARY_OP_ROUND: case GGML_UNARY_OP_TRUNC: + case GGML_UNARY_OP_ROUND_BF16: { n_tasks = 1; } break; diff --git a/external/ggml/src/ggml-cpu/ops.cpp b/external/ggml/src/ggml-cpu/ops.cpp index 0f0f57399..4047e105f 100644 --- a/external/ggml/src/ggml-cpu/ops.cpp +++ b/external/ggml/src/ggml-cpu/ops.cpp @@ -10068,6 +10068,10 @@ void ggml_compute_forward_unary( { ggml_compute_forward_trunc(params, dst); } break; + case GGML_UNARY_OP_ROUND_BF16: + { + ggml_compute_forward_round_bf16(params, dst); + } break; case GGML_UNARY_OP_XIELU: { ggml_compute_forward_xielu(params, dst); diff --git a/external/ggml/src/ggml-cpu/unary-ops.cpp b/external/ggml/src/ggml-cpu/unary-ops.cpp index a82ffc260..4a813271c 100644 --- a/external/ggml/src/ggml-cpu/unary-ops.cpp +++ b/external/ggml/src/ggml-cpu/unary-ops.cpp @@ -97,6 +97,10 @@ static inline float op_trunc(float x) { return truncf(x); } +static inline float op_round_bf16(float x) { + return bf16_to_f32(f32_to_bf16(x)); +} + template static inline void vec_unary_op(int64_t n, dst_t * y, const src0_t * x) { constexpr auto src0_to_f32 = type_conversion_table::to_f32; @@ -322,6 +326,10 @@ void ggml_compute_forward_trunc(const ggml_compute_params * params, ggml_tensor unary_op(params, dst); } +void ggml_compute_forward_round_bf16(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} + void ggml_compute_forward_xielu(const ggml_compute_params * params, ggml_tensor * dst) { const float alpha_n = ggml_get_op_params_f32(dst, 1); const float alpha_p = ggml_get_op_params_f32(dst, 2); diff --git a/external/ggml/src/ggml-cpu/unary-ops.h b/external/ggml/src/ggml-cpu/unary-ops.h index d75037369..06f3e1c37 100644 --- a/external/ggml/src/ggml-cpu/unary-ops.h +++ b/external/ggml/src/ggml-cpu/unary-ops.h @@ -28,6 +28,7 @@ void ggml_compute_forward_floor(const struct ggml_compute_params * params, struc void ggml_compute_forward_ceil(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_round(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_trunc(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_round_bf16(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_xielu(const struct ggml_compute_params * params, struct ggml_tensor * dst); #ifdef __cplusplus diff --git a/external/ggml/src/ggml-cuda/ggml-cuda.cu b/external/ggml/src/ggml-cuda/ggml-cuda.cu index 8dc80a82d..278a49c43 100644 --- a/external/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/external/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3065,6 +3065,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_UNARY_OP_TRUNC: ggml_cuda_op_trunc(ctx, dst); break; + case GGML_UNARY_OP_ROUND_BF16: + ggml_cuda_op_round_bf16(ctx, dst); + break; case GGML_UNARY_OP_EXPM1: ggml_cuda_op_expm1(ctx, dst); break; @@ -4670,7 +4673,12 @@ static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, co ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); if (graph->graph == nullptr) { - if (ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_AMPERE) { + // CUDA graphs are disabled by default on pre-Ampere GPUs (matching + // upstream, where they regressed on some parts), but can be force + // enabled; decode loops made of many tiny kernels benefit even on + // Turing. + static const bool allow_pre_ampere = getenv("GGML_CUDA_GRAPHS_PRE_AMPERE") != nullptr; + if (!allow_pre_ampere && ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_AMPERE) { if (!graph->disable_due_to_gpu_arch) { GGML_LOG_DEBUG("%s: disabling CUDA graphs due to GPU architecture\n", __func__); } @@ -5356,6 +5364,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g // TODO: should become: //return ggml_is_contiguous_rows(op->src[0]); return ggml_is_contiguous(op->src[0]); + case GGML_UNARY_OP_ROUND_BF16: + // f32/f16/bf16 src with contiguous rows, contiguous f32 dst. + return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16 || + op->src[0]->type == GGML_TYPE_BF16) && + op->type == GGML_TYPE_F32 && ggml_is_contiguous(op) && + ggml_is_contiguous_rows(op->src[0]); default: return false; } diff --git a/external/ggml/src/ggml-cuda/unary.cu b/external/ggml/src/ggml-cuda/unary.cu index fd3ac7550..213aeb4fe 100644 --- a/external/ggml/src/ggml-cuda/unary.cu +++ b/external/ggml/src/ggml-cuda/unary.cu @@ -114,6 +114,11 @@ static __device__ __forceinline__ float op_trunc(float x) { return trunc(x); } +static __device__ __forceinline__ float op_round_bf16(float x) { + // Matches the f32 -> bf16 -> f32 cpy round trip. + return __bfloat162float(__float2bfloat16(x)); +} + template static __global__ void unary_op_kernel(const T * x, T * dst, const int k) { const int i = blockDim.x*blockIdx.x + threadIdx.x; @@ -125,6 +130,75 @@ static __global__ void unary_op_kernel(const T * x, T * dst, const int k) { dst[i] = (T)op((float)x[i]); } +// Variant for a src with contiguous rows but arbitrary row strides; dst must be contiguous. +template +static __global__ void unary_op_kernel_strided( + const char * cx, T * dst, const int64_t k, + const int64_t ne0, const int64_t ne1, const int64_t ne2, + const int64_t nb01, const int64_t nb02, const int64_t nb03) { + const int64_t i = (int64_t)blockDim.x*blockIdx.x + threadIdx.x; + + if (i >= k) { + return; + } + + const int64_t i0 = i % ne0; + const int64_t i1 = (i / ne0) % ne1; + const int64_t i2 = (i / (ne0*ne1)) % ne2; + const int64_t i3 = i / (ne0*ne1*ne2); + + const T * x = (const T *) (cx + i1*nb01 + i2*nb02 + i3*nb03); + dst[i] = (T)op((float)x[i0]); +} + +// round-to-bf16 kernels: any of f32/f16/bf16 in, always f32 out. +template +static __global__ void round_bf16_kernel(const T * x, float * dst, const int64_t k) { + const int64_t i = (int64_t)blockDim.x*blockIdx.x + threadIdx.x; + + if (i >= k) { + return; + } + + dst[i] = op_round_bf16((float)x[i]); +} + +template +static __global__ void round_bf16_kernel_strided( + const char * cx, float * dst, const int64_t k, + const int64_t ne0, const int64_t ne1, const int64_t ne2, + const int64_t nb01, const int64_t nb02, const int64_t nb03) { + const int64_t i = (int64_t)blockDim.x*blockIdx.x + threadIdx.x; + + if (i >= k) { + return; + } + + const int64_t i0 = i % ne0; + const int64_t i1 = (i / ne0) % ne1; + const int64_t i2 = (i / (ne0*ne1)) % ne2; + const int64_t i3 = i / (ne0*ne1*ne2); + + const T * x = (const T *) (cx + i1*nb01 + i2*nb02 + i3*nb03); + dst[i] = op_round_bf16((float)x[i0]); +} + +template +static void round_bf16_cuda(const ggml_tensor * src0, float * dst, cudaStream_t stream) { + const int64_t k = ggml_nelements(src0); + const int64_t num_blocks = (k + CUDA_NEG_BLOCK_SIZE - 1) / CUDA_NEG_BLOCK_SIZE; + GGML_ASSERT(num_blocks < UINT_MAX); + + if (ggml_is_contiguous(src0)) { + round_bf16_kernel<<<(unsigned int) num_blocks, CUDA_NEG_BLOCK_SIZE, 0, stream>>>( + (const T *) src0->data, dst, k); + } else { + round_bf16_kernel_strided<<<(unsigned int) num_blocks, CUDA_NEG_BLOCK_SIZE, 0, stream>>>( + (const char *) src0->data, dst, k, src0->ne[0], src0->ne[1], src0->ne[2], + src0->nb[1], src0->nb[2], src0->nb[3]); + } +} + template static void unary_cuda(const T * x, T * dst, const int k, cudaStream_t stream) { const int num_blocks = (k + CUDA_NEG_BLOCK_SIZE - 1) / CUDA_NEG_BLOCK_SIZE; @@ -247,6 +321,31 @@ void ggml_cuda_op_trunc(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_op_unary(ctx, dst); } +void ggml_cuda_op_round_bf16(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + + // ggml_round_bf16 always produces a contiguous f32 dst; src may be + // f32/f16/bf16 with contiguous rows. + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_is_contiguous_rows(src0)); + + cudaStream_t stream = ctx.stream(); + switch (src0->type) { + case GGML_TYPE_F32: + round_bf16_cuda(src0, (float *) dst->data, stream); + break; + case GGML_TYPE_F16: + round_bf16_cuda(src0, (float *) dst->data, stream); + break; + case GGML_TYPE_BF16: + round_bf16_cuda(src0, (float *) dst->data, stream); + break; + default: + GGML_ABORT("%s: unsupported src type %s", __func__, ggml_type_name(src0->type)); + } +} + void ggml_cuda_op_expm1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_op_unary(ctx, dst); } diff --git a/external/ggml/src/ggml-cuda/unary.cuh b/external/ggml/src/ggml-cuda/unary.cuh index fbb229cb8..06a7e589f 100644 --- a/external/ggml/src/ggml-cuda/unary.cuh +++ b/external/ggml/src/ggml-cuda/unary.cuh @@ -75,6 +75,8 @@ void ggml_cuda_op_round(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_trunc(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +void ggml_cuda_op_round_bf16(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + void ggml_cuda_op_reglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_geglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/external/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/external/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 334651a11..3e9a5217b 100644 --- a/external/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/external/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -774,8 +774,8 @@ struct vk_device_struct { vk_pipeline pipeline_pad_reflect_1d_f32; vk_pipeline pipeline_roll_f32; vk_pipeline pipeline_repeat_f32, pipeline_repeat_back_f32; - vk_pipeline pipeline_cpy_f32_f32, pipeline_cpy_f32_f16, pipeline_cpy_f16_f16, pipeline_cpy_f16_f32, pipeline_cpy_f32_bf16, pipeline_cpy_f32_i32, pipeline_cpy_i32_f32; - vk_pipeline pipeline_contig_cpy_f32_f32, pipeline_contig_cpy_f32_f16, pipeline_contig_cpy_f16_f16, pipeline_contig_cpy_f16_f32, pipeline_contig_cpy_f32_bf16, pipeline_contig_cpy_f32_i32, pipeline_contig_cpy_i32_f32; + vk_pipeline pipeline_cpy_f32_f32, pipeline_cpy_f32_f16, pipeline_cpy_f16_f16, pipeline_cpy_f16_f32, pipeline_cpy_f32_bf16, pipeline_cpy_bf16_f32, pipeline_cpy_f16_bf16, pipeline_cpy_bf16_f16, pipeline_cpy_f32_i32, pipeline_cpy_i32_f32; + vk_pipeline pipeline_contig_cpy_f32_f32, pipeline_contig_cpy_f32_f16, pipeline_contig_cpy_f16_f16, pipeline_contig_cpy_f16_f32, pipeline_contig_cpy_f32_bf16, pipeline_contig_cpy_bf16_f32, pipeline_contig_cpy_f16_bf16, pipeline_contig_cpy_bf16_f16, pipeline_contig_cpy_f32_i32, pipeline_contig_cpy_i32_f32; vk_pipeline pipeline_cpy_f32_quant[GGML_TYPE_COUNT]; vk_pipeline pipeline_cpy_quant_f32[GGML_TYPE_COUNT]; vk_pipeline pipeline_cpy_transpose_16, pipeline_cpy_transpose_32; @@ -4594,7 +4594,10 @@ static void ggml_vk_load_shaders(vk_device& device) { ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_f16, "cpy_f32_f16", cpy_f32_f16_len, cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_f16_f16, "cpy_f16_f16", cpy_f16_f16_len, cpy_f16_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_f16_f32, "cpy_f16_f32", cpy_f16_f32_len, cpy_f16_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); - ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_bf16,"cpy_f32_bf16",cpy_f32_bf16_len,cpy_f32_bf16_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_bf16,"cpy_f32_bf16",cpy_f32_bf16_len,cpy_f32_bf16_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_bf16_f32,"cpy_bf16_f32",cpy_bf16_f32_len,cpy_bf16_f32_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_f16_bf16,"cpy_f16_bf16",cpy_f16_bf16_len,cpy_f16_bf16_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_bf16_f16,"cpy_bf16_f16",cpy_bf16_f16_len,cpy_bf16_f16_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_i32_f32, "cpy_i32_f32", cpy_i32_f32_len, cpy_i32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_i32, "cpy_f32_i32", cpy_f32_i32_len, cpy_f32_i32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); @@ -4602,7 +4605,10 @@ static void ggml_vk_load_shaders(vk_device& device) { ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_f16, "contig_cpy_f32_f16", contig_cpy_f32_f16_len, contig_cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f16_f16, "contig_cpy_f16_f16", contig_cpy_f16_f16_len, contig_cpy_f16_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f16_f32, "contig_cpy_f16_f32", contig_cpy_f16_f32_len, contig_cpy_f16_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); - ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_bf16,"contig_cpy_f32_bf16",contig_cpy_f32_bf16_len,contig_cpy_f32_bf16_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_bf16,"contig_cpy_f32_bf16",contig_cpy_f32_bf16_len,contig_cpy_f32_bf16_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_bf16_f32,"contig_cpy_bf16_f32",contig_cpy_bf16_f32_len,contig_cpy_bf16_f32_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f16_bf16,"contig_cpy_f16_bf16",contig_cpy_f16_bf16_len,contig_cpy_f16_bf16_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_bf16_f16,"contig_cpy_bf16_f16",contig_cpy_bf16_f16_len,contig_cpy_bf16_f16_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_i32_f32, "contig_cpy_i32_f32", contig_cpy_i32_f32_len, contig_cpy_i32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_i32, "contig_cpy_f32_i32", contig_cpy_f32_i32_len, contig_cpy_f32_i32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); @@ -7578,6 +7584,27 @@ static vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const return ctx->device->pipeline_cpy_f32_bf16; } } + if (src->type == GGML_TYPE_BF16 && to == GGML_TYPE_F32) { + if (contig) { + return ctx->device->pipeline_contig_cpy_bf16_f32; + } else { + return ctx->device->pipeline_cpy_bf16_f32; + } + } + if (src->type == GGML_TYPE_F16 && to == GGML_TYPE_BF16) { + if (contig) { + return ctx->device->pipeline_contig_cpy_f16_bf16; + } else { + return ctx->device->pipeline_cpy_f16_bf16; + } + } + if (src->type == GGML_TYPE_BF16 && to == GGML_TYPE_F16) { + if (contig) { + return ctx->device->pipeline_contig_cpy_bf16_f16; + } else { + return ctx->device->pipeline_cpy_bf16_f16; + } + } if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_I32) { if (contig) { return ctx->device->pipeline_contig_cpy_f32_i32; diff --git a/external/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp b/external/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp index 066b27ca0..cacddbdf7 100644 --- a/external/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp +++ b/external/ggml/src/ggml-vulkan/vulkan-shaders/contig_copy.comp @@ -19,7 +19,10 @@ void main() { if (idx + (num_iter-1)*num_threads < p.ne) { [[unroll]] for (uint i = 0; i < num_iter; ++i) { -#if defined(DATA_D_BF16) +#if defined(DATA_A_BF16) + float f = bf16_to_fp32(uint32_t(data_a[get_aoffset() + idx])); + data_d[get_doffset() + idx] = D_TYPE(f); +#elif defined(DATA_D_BF16) float f = float(data_a[get_aoffset() + idx]); data_d[get_doffset() + idx] = D_TYPE(fp32_to_bf16(f)); #elif !defined(OPTIMIZATION_ERROR_WORKAROUND) @@ -35,7 +38,10 @@ void main() { continue; } -#if defined(DATA_D_BF16) +#if defined(DATA_A_BF16) + float f = bf16_to_fp32(uint32_t(data_a[get_aoffset() + idx])); + data_d[get_doffset() + idx] = D_TYPE(f); +#elif defined(DATA_D_BF16) float f = float(data_a[get_aoffset() + idx]); data_d[get_doffset() + idx] = D_TYPE(fp32_to_bf16(f)); #elif !defined(OPTIMIZATION_ERROR_WORKAROUND) diff --git a/external/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp b/external/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp index a1ba96702..81cce864e 100644 --- a/external/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp +++ b/external/ggml/src/ggml-vulkan/vulkan-shaders/copy.comp @@ -12,7 +12,10 @@ void main() { return; } -#if defined(DATA_D_BF16) +#if defined(DATA_A_BF16) + float f = bf16_to_fp32(uint32_t(data_a[get_aoffset() + src0_idx(idx)])); + data_d[get_doffset() + dst_idx(idx)] = D_TYPE(f); +#elif defined(DATA_D_BF16) float f = float(data_a[get_aoffset() + src0_idx(idx)]); data_d[get_doffset() + dst_idx(idx)] = D_TYPE(fp32_to_bf16(f)); #elif !defined(OPTIMIZATION_ERROR_WORKAROUND) diff --git a/external/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/external/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 7295ef107..a533af69f 100644 --- a/external/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/external/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -730,14 +730,20 @@ void process_shaders() { string_to_spv("cpy_f32_f16", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}}); string_to_spv("cpy_f16_f16", "copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); string_to_spv("cpy_f16_f32", "copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); - string_to_spv("cpy_f32_bf16","copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "uint16_t"}, {"DATA_D_BF16", "1"}}); + string_to_spv("cpy_f32_bf16","copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "uint16_t"}, {"DATA_D_BF16", "1"}}); + string_to_spv("cpy_bf16_f32","copy.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "float"}, {"DATA_A_BF16", "1"}}); + string_to_spv("cpy_f16_bf16","copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "uint16_t"}, {"DATA_D_BF16", "1"}}); + string_to_spv("cpy_bf16_f16","copy.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "float16_t"}, {"DATA_A_BF16", "1"}}); string_to_spv("contig_cpy_f32_f32", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); string_to_spv("contig_cpy_f32_i32", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "int"}}); string_to_spv("contig_cpy_i32_f32", "contig_copy.comp", {{"A_TYPE", "int"}, {"D_TYPE", "float"}}); string_to_spv("contig_cpy_f32_f16", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}}); string_to_spv("contig_cpy_f16_f16", "contig_copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); string_to_spv("contig_cpy_f16_f32", "contig_copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); - string_to_spv("contig_cpy_f32_bf16","contig_copy.comp",{{"A_TYPE", "float"}, {"D_TYPE", "uint16_t"}, {"DATA_D_BF16", "1"}}); + string_to_spv("contig_cpy_f32_bf16","contig_copy.comp",{{"A_TYPE", "float"}, {"D_TYPE", "uint16_t"}, {"DATA_D_BF16", "1"}}); + string_to_spv("contig_cpy_bf16_f32","contig_copy.comp",{{"A_TYPE", "uint16_t"}, {"D_TYPE", "float"}, {"DATA_A_BF16", "1"}}); + string_to_spv("contig_cpy_f16_bf16","contig_copy.comp",{{"A_TYPE", "float16_t"}, {"D_TYPE", "uint16_t"}, {"DATA_D_BF16", "1"}}); + string_to_spv("contig_cpy_bf16_f16","contig_copy.comp",{{"A_TYPE", "uint16_t"}, {"D_TYPE", "float16_t"}, {"DATA_A_BF16", "1"}}); string_to_spv("cpy_f32_i32", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "int"}}); string_to_spv("cpy_i32_f32", "copy.comp", {{"A_TYPE", "int"}, {"D_TYPE", "float"}}); diff --git a/external/ggml/src/ggml.c b/external/ggml/src/ggml.c index 0ba560368..dbad3596d 100644 --- a/external/ggml/src/ggml.c +++ b/external/ggml/src/ggml.c @@ -1229,9 +1229,10 @@ static const char * GGML_UNARY_OP_NAME[GGML_UNARY_OP_COUNT] = { "CEIL", "ROUND", "TRUNC", + "ROUND_BF16", }; -static_assert(GGML_UNARY_OP_COUNT == 22, "GGML_UNARY_OP_COUNT != 22"); +static_assert(GGML_UNARY_OP_COUNT == 23, "GGML_UNARY_OP_COUNT != 23"); static const char * GGML_GLU_OP_NAME[GGML_GLU_OP_COUNT] = { "REGLU", @@ -2959,6 +2960,26 @@ struct ggml_tensor * ggml_trunc_inplace( return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_TRUNC); } +//ggml_round_bf16 + +struct ggml_tensor * ggml_round_bf16( + struct ggml_context * ctx, + struct ggml_tensor * a) { + GGML_ASSERT(a->type == GGML_TYPE_F32 || a->type == GGML_TYPE_F16 || a->type == GGML_TYPE_BF16); + GGML_ASSERT(ggml_is_contiguous_rows(a)); + + // Unlike ggml_unary, the result is always f32: bf16/f16 inputs are widened + // while rounding, matching an f32 -> bf16 -> f32 cast round trip. + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, a->ne); + + ggml_set_op_params_i32(result, 0, (int32_t) GGML_UNARY_OP_ROUND_BF16); + + result->op = GGML_OP_UNARY; + result->src[0] = a; + + return result; +} + struct ggml_tensor * ggml_glu( struct ggml_context * ctx, struct ggml_tensor * a, diff --git a/include/engine/framework/modules/transformers/qwen_decoder.h b/include/engine/framework/modules/transformers/qwen_decoder.h index 3b43a744b..eaa80c12e 100644 --- a/include/engine/framework/modules/transformers/qwen_decoder.h +++ b/include/engine/framework/modules/transformers/qwen_decoder.h @@ -53,6 +53,10 @@ enum class QwenDecoderPositionEncoding { struct QwenDecoderActivationCastPolicy { bool enabled = false; ggml_type type = GGML_TYPE_BF16; + // Use the fused single-kernel round-to-bf16 op instead of a + // cast -> bf16 -> cast -> f32 round trip. Only valid on backends that + // implement GGML_UNARY_OP_ROUND_BF16 (CUDA/HIP, CPU fallback). + bool fused_round = false; bool after_input_norm = false; bool after_qkv_projection = false; bool after_qk_norm = false; diff --git a/src/framework/modules/conv_modules.cpp b/src/framework/modules/conv_modules.cpp index 185b66e35..e6f562b92 100644 --- a/src/framework/modules/conv_modules.cpp +++ b/src/framework/modules/conv_modules.cpp @@ -319,7 +319,8 @@ bool is_conv_transpose1d_col2im_fast_path_eligible( const core::ModuleBuildContext & ctx, const ConvTranspose1dConfig & config) noexcept { return (core::uses_ggml_cuda_or_hip_backend(ctx.backend_type) || - ctx.backend_type == core::BackendType::Metal) && + ctx.backend_type == core::BackendType::Metal || + ctx.backend_type == core::BackendType::Vulkan) && config.dilation == 1; } diff --git a/src/framework/modules/transformers/qwen_decoder.cpp b/src/framework/modules/transformers/qwen_decoder.cpp index a50b7191c..d7031270b 100644 --- a/src/framework/modules/transformers/qwen_decoder.cpp +++ b/src/framework/modules/transformers/qwen_decoder.cpp @@ -48,6 +48,9 @@ int64_t require_head_dim(const QwenDecoderLayerConfig & config) { config.activation_cast.type != GGML_TYPE_F16 && config.activation_cast.type != GGML_TYPE_BF16) { throw std::runtime_error("QwenDecoderLayerConfig activation cast supports only f32, f16, and bf16"); } + if (config.activation_cast.fused_round && config.activation_cast.type != GGML_TYPE_BF16) { + throw std::runtime_error("QwenDecoderLayerConfig fused activation rounding requires bf16"); + } return config.head_dim; } @@ -235,6 +238,13 @@ core::TensorValue activation_cast( if (policy.type == GGML_TYPE_F32) { return core::wrap_tensor(ggml_cast(ctx.ggml, input.tensor, GGML_TYPE_F32), input.shape, GGML_TYPE_F32); } + if (policy.fused_round && policy.type == GGML_TYPE_BF16 && ggml_is_contiguous_rows(input.tensor)) { + // Fused single-kernel round-to-bf16: f32/f16/bf16 in, always f32 out, + // values rounded to bf16. Numerically identical to the cast round trip + // below (bf16 input is already rounded, so rounding is a widening no-op), + // but avoids the intermediate bf16 tensor and one kernel launch. + return core::wrap_tensor(ggml_round_bf16(ctx.ggml, input.tensor), input.shape, GGML_TYPE_F32); + } auto rounded = core::wrap_tensor(ggml_cast(ctx.ggml, input.tensor, policy.type), input.shape, policy.type); return core::wrap_tensor(ggml_cast(ctx.ggml, rounded.tensor, GGML_TYPE_F32), input.shape, GGML_TYPE_F32); } diff --git a/src/models/breeze_tts/generator.cpp b/src/models/breeze_tts/generator.cpp index 7d8483877..83f149540 100644 --- a/src/models/breeze_tts/generator.cpp +++ b/src/models/breeze_tts/generator.cpp @@ -48,6 +48,37 @@ struct GgmlContextDeleter { } }; +// The official Breeze-TTS 2 inference runs the backbone and depth decoder with +// bf16 activations and a bf16 KV cache. Pure fp32 activations measurably drift +// into degenerate trajectories on some prompts (mispronunciations, repetition +// collapse), so match the reference bf16 behavior on GPU backends. +modules::QwenDecoderActivationCastPolicy breeze_bf16_activation_policy(core::BackendType backend_type) { + modules::QwenDecoderActivationCastPolicy policy; + if (backend_type != core::BackendType::Cuda && backend_type != core::BackendType::Hip && + backend_type != core::BackendType::Vulkan) { + return policy; + } + policy.enabled = true; + policy.type = GGML_TYPE_BF16; + // CUDA/HIP implement the fused round-to-bf16 unary op; Vulkan does not and + // keeps the cast round trip. + policy.fused_round = backend_type == core::BackendType::Cuda || backend_type == core::BackendType::Hip; + policy.after_input_norm = true; + policy.after_qkv_projection = true; + policy.after_qk_norm = true; + policy.after_rope = true; + policy.after_static_cache_update = true; + policy.after_attention = true; + policy.after_attention_output = true; + policy.after_residual = true; + policy.after_ffn_norm = true; + policy.after_mlp_projection = true; + policy.after_mlp_silu = true; + policy.after_mlp_mul = true; + policy.after_output = true; + return policy; +} + modules::QwenCausalDecodeRuntimeConfig backbone_config( const BreezeTTSConfig & config, core::BackendType backend_type, @@ -66,6 +97,8 @@ modules::QwenCausalDecodeRuntimeConfig backbone_config( out.decoder.stack.rope_theta = config.rope_theta; out.decoder.stack.rope_type = GGML_ROPE_TYPE_NEOX; out.decoder.stack.use_qk_norm = true; + out.decoder.stack.qkv_layout = modules::QwenDecoderQKVLayout::PackedQKV; + out.decoder.stack.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; out.decoder.stack.attention_precision = GGML_PREC_F32; out.decoder.stack.projection_precision = GGML_PREC_DEFAULT; out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; @@ -74,7 +107,12 @@ modules::QwenCausalDecodeRuntimeConfig backbone_config( out.decoder.stack.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; if (backend_type == core::BackendType::Cuda || backend_type == core::BackendType::Hip || backend_type == core::BackendType::Vulkan) { - out.decoder.static_cache_type = GGML_TYPE_F16; + // BF16 KV cache matches the reference implementation, but flash + // attention only accelerates bf16 cache with native bf16 MMA + // (sm_80+); on older parts it is ~3x slower, so only HIP uses it. + out.decoder.static_cache_type = + backend_type == core::BackendType::Hip ? GGML_TYPE_BF16 : GGML_TYPE_F16; + out.decoder.stack.activation_cast = breeze_bf16_activation_policy(backend_type); } out.decoder.logits_size = config.lm_head_size; out.decoder.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; @@ -106,6 +144,8 @@ modules::QwenCausalDecodeRuntimeConfig depth_config( out.decoder.stack.rope_theta = config.depth_rope_theta; out.decoder.stack.rope_type = GGML_ROPE_TYPE_NEOX; out.decoder.stack.use_qk_norm = false; + out.decoder.stack.qkv_layout = modules::QwenDecoderQKVLayout::PackedQKV; + out.decoder.stack.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; out.decoder.stack.attention_precision = GGML_PREC_F32; out.decoder.stack.projection_precision = GGML_PREC_DEFAULT; out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; @@ -114,7 +154,11 @@ modules::QwenCausalDecodeRuntimeConfig depth_config( out.decoder.stack.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; if (backend_type == core::BackendType::Cuda || backend_type == core::BackendType::Hip || backend_type == core::BackendType::Vulkan) { - out.decoder.static_cache_type = GGML_TYPE_F16; + // See backbone_config: only HIP uses a bf16 KV cache; CUDA and Vulkan + // keep F16. + out.decoder.static_cache_type = + backend_type == core::BackendType::Hip ? GGML_TYPE_BF16 : GGML_TYPE_F16; + out.decoder.stack.activation_cast = breeze_bf16_activation_policy(backend_type); } out.decoder.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; out.output_mode = modules::QwenCausalDecodeOutputMode::Hidden; @@ -160,6 +204,35 @@ std::vector llama3_rope_factors( return out; } +// Pack [a; b; ...] projection rows into a single tensor, matching the +// higgs_audio_tts loader: fewer, larger matmuls per layer. Parts may have +// different row counts (e.g. q vs k/v in GQA models). +core::TensorValue pack_projection_rows( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::vector> & parts, + assets::TensorStorageType storage_type, + int64_t in_dim) { + std::vector packed; + int64_t total_out = 0; + ggml_type packed_type = GGML_TYPE_COUNT; + for (const auto & [name, out_dim] : parts) { + const auto part = source.require_tensor(name, storage_type, {out_dim, in_dim}); + if (packed_type == GGML_TYPE_COUNT) { + packed_type = part.type; + } else if (part.type != packed_type) { + throw std::runtime_error("BreezeTTS packed projection weights require matching storage types"); + } + packed.insert(packed.end(), part.bytes.begin(), part.bytes.end()); + total_out += out_dim; + } + return store.make_tensor( + core::TensorShape::from_dims({total_out, in_dim}), + packed_type, + packed.data(), + packed.size()); +} + modules::QwenDecoderLayerWeights load_backbone_layer( core::BackendWeightStore & store, const assets::TensorSource & source, @@ -168,17 +241,33 @@ modules::QwenDecoderLayerWeights load_backbone_layer( const std::optional & rope_factors, int64_t layer) { const std::string prefix = "backbone_model.layers." + std::to_string(layer); + const int64_t q_out = config.heads * config.head_dim; + const int64_t kv_out = config.kv_heads * config.head_dim; modules::QwenDecoderLayerWeights out; out.input_norm = binding::norm_weight_from_source(store, source, prefix + ".input_layernorm", config.hidden_size); - out.self_attention.q_weight = store.load_tensor(source, prefix + ".self_attn.q_proj.weight", storage_type, {config.heads * config.head_dim, config.hidden_size}); - out.self_attention.k_weight = store.load_tensor(source, prefix + ".self_attn.k_proj.weight", storage_type, {config.kv_heads * config.head_dim, config.hidden_size}); - out.self_attention.v_weight = store.load_tensor(source, prefix + ".self_attn.v_proj.weight", storage_type, {config.kv_heads * config.head_dim, config.hidden_size}); + // Packed layout is [q; k; v] with row counts q_out, kv_out, kv_out. + out.self_attention.qkv_weight = pack_projection_rows( + store, + source, + {{prefix + ".self_attn.q_proj.weight", q_out}, + {prefix + ".self_attn.k_proj.weight", kv_out}, + {prefix + ".self_attn.v_proj.weight", kv_out}}, + storage_type, + config.hidden_size); out.self_attention.out_weight = store.load_tensor(source, prefix + ".self_attn.o_proj.weight", storage_type, {config.hidden_size, config.heads * config.head_dim}); out.q_norm = binding::norm_weight_from_source(store, source, prefix + ".self_attn.q_norm", config.head_dim); out.k_norm = binding::norm_weight_from_source(store, source, prefix + ".self_attn.k_norm", config.head_dim); out.post_norm = binding::norm_weight_from_source(store, source, prefix + ".post_attention_layernorm", config.hidden_size); - out.mlp.gate_proj = binding::linear_from_source(store, source, prefix + ".mlp.gate_proj", storage_type, config.intermediate_size, config.hidden_size, false); - out.mlp.up_proj = binding::linear_from_source(store, source, prefix + ".mlp.up_proj", storage_type, config.intermediate_size, config.hidden_size, false); + out.mlp.gate_up_proj = modules::LinearWeights{ + pack_projection_rows( + store, + source, + {{prefix + ".mlp.gate_proj.weight", config.intermediate_size}, + {prefix + ".mlp.up_proj.weight", config.intermediate_size}}, + storage_type, + config.hidden_size), + std::nullopt, + }; out.mlp.down_proj = binding::linear_from_source(store, source, prefix + ".mlp.down_proj", storage_type, config.hidden_size, config.intermediate_size, false); out.rope_frequency_factors = rope_factors; return out; @@ -192,15 +281,31 @@ modules::QwenDecoderLayerWeights load_depth_layer( const std::optional & rope_factors, int64_t layer) { const std::string prefix = "depth_decoder.model.layers." + std::to_string(layer); + const int64_t q_out = config.depth_heads * config.depth_head_dim; + const int64_t kv_out = config.depth_kv_heads * config.depth_head_dim; modules::QwenDecoderLayerWeights out; out.input_norm = binding::norm_weight_from_source(store, source, prefix + ".input_layernorm", config.depth_hidden_size); - out.self_attention.q_weight = store.load_tensor(source, prefix + ".self_attn.q_proj.weight", storage_type, {config.depth_heads * config.depth_head_dim, config.depth_hidden_size}); - out.self_attention.k_weight = store.load_tensor(source, prefix + ".self_attn.k_proj.weight", storage_type, {config.depth_kv_heads * config.depth_head_dim, config.depth_hidden_size}); - out.self_attention.v_weight = store.load_tensor(source, prefix + ".self_attn.v_proj.weight", storage_type, {config.depth_kv_heads * config.depth_head_dim, config.depth_hidden_size}); + // Packed layout is [q; k; v] with row counts q_out, kv_out, kv_out. + out.self_attention.qkv_weight = pack_projection_rows( + store, + source, + {{prefix + ".self_attn.q_proj.weight", q_out}, + {prefix + ".self_attn.k_proj.weight", kv_out}, + {prefix + ".self_attn.v_proj.weight", kv_out}}, + storage_type, + config.depth_hidden_size); out.self_attention.out_weight = store.load_tensor(source, prefix + ".self_attn.o_proj.weight", storage_type, {config.depth_hidden_size, config.depth_heads * config.depth_head_dim}); out.post_norm = binding::norm_weight_from_source(store, source, prefix + ".post_attention_layernorm", config.depth_hidden_size); - out.mlp.gate_proj = binding::linear_from_source(store, source, prefix + ".mlp.gate_proj", storage_type, config.depth_intermediate_size, config.depth_hidden_size, false); - out.mlp.up_proj = binding::linear_from_source(store, source, prefix + ".mlp.up_proj", storage_type, config.depth_intermediate_size, config.depth_hidden_size, false); + out.mlp.gate_up_proj = modules::LinearWeights{ + pack_projection_rows( + store, + source, + {{prefix + ".mlp.gate_proj.weight", config.depth_intermediate_size}, + {prefix + ".mlp.up_proj.weight", config.depth_intermediate_size}}, + storage_type, + config.depth_hidden_size), + std::nullopt, + }; out.mlp.down_proj = binding::linear_from_source(store, source, prefix + ".mlp.down_proj", storage_type, config.depth_hidden_size, config.depth_intermediate_size, false); out.rope_frequency_factors = rope_factors; return out; @@ -489,6 +594,12 @@ class BreezeDepthProjectionRuntime { ggml_backend_synchronize(backend_); ggml_backend_tensor_get(graph.output, head_paired_staging_.data(), 0, head_paired_staging_.size() * sizeof(float)); const size_t vocab = static_cast(vocab_); + if (guidance_scale == 1.0F) { + // CFG is a no-op at scale 1: copy the conditional half directly, + // avoiding an inexact uncond + 1 * (cond - uncond) round trip. + std::memcpy(out, head_paired_staging_.data(), vocab * sizeof(float)); + return; + } for (int64_t token = 0; token < vocab_; ++token) { const size_t index = static_cast(token); out[index] = head_paired_staging_[vocab + index] + guidance_scale * (head_paired_staging_[index] - head_paired_staging_[vocab + index]); @@ -862,21 +973,30 @@ struct BreezeGeneratorRuntime::Impl { std::vector uncond_embeddings; int64_t cond_steps = 0; int64_t uncond_steps = 0; + // guidance_scale == 1 makes CFG a no-op (logits == cond), so skip the + // unconditional branch entirely and halve the backbone work. + const bool use_cfg = request.guidance_scale != 1.0F; const double prompt_ms = engine::debug::measure_ms([&] { if (!reference_codes.empty()) { if (request.reference_text.empty()) { throw std::runtime_error("BreezeTTS clone requires reference_text"); } cond_branch = tokenizer.build_clone(request.text, request.instruction, request.reference_text, reference_frames); - uncond_branch = tokenizer.build_clone_negative(request.text, request.reference_text, reference_frames); + if (use_cfg) { + uncond_branch = tokenizer.build_clone_negative(request.text, request.reference_text, reference_frames); + } } else { cond_branch = tokenizer.build_tts_instruction(request.text, request.instruction); - uncond_branch = tokenizer.build_tts_plain(request.text); + if (use_cfg) { + uncond_branch = tokenizer.build_tts_plain(request.text); + } } cond_embeddings = merge_prompt(cond_branch, reference_codes); - uncond_embeddings = merge_prompt(uncond_branch, reference_codes); cond_steps = static_cast(cond_branch.input_ids.size()); - uncond_steps = static_cast(uncond_branch.input_ids.size()); + if (use_cfg) { + uncond_embeddings = merge_prompt(uncond_branch, reference_codes); + uncond_steps = static_cast(uncond_branch.input_ids.size()); + } }); engine::debug::timing_log_scalar("breeze_tts.generate.prompt_ms", prompt_ms); text_encoder.release_runtime_graphs(); @@ -891,12 +1011,17 @@ struct BreezeGeneratorRuntime::Impl { backbone_cond_prefill_ms = engine::debug::measure_ms([&] { cond = backbone_cond->prefill_embeddings(cond_embeddings, cond_steps); }); - modules::QwenCausalPrefillResult uncond; - backbone_uncond_prefill_ms = engine::debug::measure_ms([&] { - uncond = backbone_uncond->prefill_embeddings(uncond_embeddings, uncond_steps); - }); + std::optional uncond; + if (use_cfg) { + uncond.emplace(); + backbone_uncond_prefill_ms = engine::debug::measure_ms([&] { + *uncond = backbone_uncond->prefill_embeddings(uncond_embeddings, uncond_steps); + }); + } backbone_cond->start_decode_embeddings(cond.state, cond_steps + request.max_tokens); - backbone_uncond->start_decode_embeddings(uncond.state, uncond_steps + request.max_tokens); + if (use_cfg) { + backbone_uncond->start_decode_embeddings(uncond->state, uncond_steps + request.max_tokens); + } sampling::HfSamplerScratch scratch; scratch.reserve_vocab(static_cast(config.lm_head_size)); @@ -913,12 +1038,17 @@ struct BreezeGeneratorRuntime::Impl { codes.reserve(static_cast(request.max_tokens * config.num_codebooks)); for (int64_t step = 0; step < request.max_tokens; ++step) { - if (cond.logits.size() != uncond.logits.size()) { + if (use_cfg && cond.logits.size() != uncond->logits.size()) { throw std::runtime_error("BreezeTTS CFG logits shape mismatch"); } - std::vector logits(cond.logits.size(), 0.0F); - for (size_t i = 0; i < logits.size(); ++i) { - logits[i] = uncond.logits[i] + request.guidance_scale * (cond.logits[i] - uncond.logits[i]); + std::vector logits; + if (use_cfg) { + logits.resize(cond.logits.size()); + for (size_t i = 0; i < logits.size(); ++i) { + logits[i] = uncond->logits[i] + request.guidance_scale * (cond.logits[i] - uncond->logits[i]); + } + } else { + logits = cond.logits; } suppress_reserved(logits, kCodecCodebookSize, config.vocab_size); const int32_t first_token = sample_logits( @@ -940,7 +1070,7 @@ struct BreezeGeneratorRuntime::Impl { } const auto frame = generate_frame( cond.hidden, - uncond.hidden, + use_cfg ? uncond->hidden : cond.hidden, first_token, request, scratch, @@ -959,14 +1089,16 @@ struct BreezeGeneratorRuntime::Impl { backbone_cond_decode_ms += engine::debug::measure_ms([&] { cond_step = backbone_cond->decode_embedding(embedded); }); - modules::QwenCausalDecodeStepResult uncond_step; - backbone_uncond_decode_ms += engine::debug::measure_ms([&] { - uncond_step = backbone_uncond->decode_embedding(embedded); - }); cond.logits = cond_step.logits; cond.hidden = cond_step.hidden; - uncond.logits = uncond_step.logits; - uncond.hidden = uncond_step.hidden; + if (use_cfg) { + modules::QwenCausalDecodeStepResult uncond_step; + backbone_uncond_decode_ms += engine::debug::measure_ms([&] { + uncond_step = backbone_uncond->decode_embedding(embedded); + }); + uncond->logits = uncond_step.logits; + uncond->hidden = uncond_step.hidden; + } } }); engine::debug::timing_log_scalar("breeze_tts.ar.total_ms", ar_ms); @@ -975,7 +1107,9 @@ struct BreezeGeneratorRuntime::Impl { engine::debug::timing_log_scalar("breeze_tts.ar.backbone_cond_decode_ms", backbone_cond_decode_ms); engine::debug::timing_log_scalar("breeze_tts.ar.backbone_uncond_decode_ms", backbone_uncond_decode_ms); backbone_cond->release_runtime_graphs(); - backbone_uncond->release_runtime_graphs(); + if (use_cfg) { + backbone_uncond->release_runtime_graphs(); + } depth_pair->release_runtime_graphs(); if (codes.empty()) { throw std::runtime_error("BreezeTTS generated no audio codes"); diff --git a/src/models/breeze_tts/session.cpp b/src/models/breeze_tts/session.cpp index 264d67bd1..10f7e5e86 100644 --- a/src/models/breeze_tts/session.cpp +++ b/src/models/breeze_tts/session.cpp @@ -315,7 +315,10 @@ BreezeGenerationRequest BreezeTTSSession::build_generation_request( generation.text = request.text_input->text; generation.instruction = runtime::find_option(request.options, {"instruction"}).value_or(""); generation.reference_text = runtime::find_option(request.options, {"reference_text"}).value_or(""); - generation.guidance_scale = runtime::parse_positive_finite_float_option(request.options, {"guidance_scale"}).value_or(generation.guidance_scale); + generation.guidance_scale = runtime::parse_finite_float_option(request.options, {"guidance_scale"}).value_or(generation.guidance_scale); + if (generation.guidance_scale < 0.0F) { + throw std::runtime_error("BreezeTTS guidance_scale must be non-negative"); + } generation.temperature = runtime::parse_positive_finite_float_option(request.options, {"temperature"}).value_or(generation.temperature); generation.depth_temperature = runtime::parse_positive_finite_float_option(request.options, {"depth_temperature"}).value_or(generation.depth_temperature); generation.top_k = runtime::parse_i64_option(request.options, {"top_k"}).value_or(generation.top_k); From af0df8a1559c195ba2423d2a7fdd3b2de0e7731c Mon Sep 17 00:00:00 2001 From: Laika <2432896620@qq.com> Date: Fri, 4 Sep 2026 11:28:26 +0800 Subject: [PATCH 08/10] Breeze encoder chunked vram (#431) * breeze: chunk the speech-encoder conv stack to bound clone VRAM The encoder graph was built at the exact reference-audio length, so conv activations grew linearly (~45 MiB/s of reference) and every new length triggered a full graph rebuild; a 60 s reference cost ~2.5 GB extra over a 6 s one. Split the encoder into two graphs. The conv stack now runs on fixed 5 s chunks (120000 samples) preceded by a 9600-sample left overlap that covers the stack's exact 5240-sample receptive field; chunk lengths are multiples of the 960x transformer stride, so no per-stage right padding occurs and the discarded overlap frames absorb the zero left pads that represent audio start in the first chunk. Stitched outputs are bit-identical to a single-pass encode of the same input (verified over 68 frames x 16 codebooks). The transformer, downsample, and projections run once over the full frame sequence at frame scale, where even minute-long references cost only tens of MiB. Measured on a 2080 Ti (Vulkan, native q8_0 GGUF, peak minus idle baseline): the VRAM slope over reference length drops from ~45 MiB/s to ~11 MiB/s (remaining slope is the frame-scale transformer graph and the longer AR prefill from reference codes), and a 60 s reference peaks ~1.4 GB lower. Encode time for 60 s improves from 3561 ms to 2197 ms. * breeze: bucket speech-encoder transformer graph capacity The transformer graph was rebuilt at the exact frame count for every distinct reference length. Round the capacity up to 125-frame (5 s) buckets so lengths within a bucket share one graph. Unused bucket frames are replicate-padded to match the downsample conv's Replicate right pad; causal attention keeps padding frames invisible to real frames. Verified bit-identical reference codes vs exact-length graphs at 6 s and 15 s; odd lengths show sub-1% last-frame diffs from flash-attention tiling, the same accepted noise class as the pre-existing length sensitivity. Single-run peak VRAM is unchanged. * ggml-vulkan, breeze: fused round-to-bf16 unary op on Vulkan Vulkan previously paid a cast round trip (f32->bf16->f32, two kernels, a bf16 intermediate tensor) at every activation-rounding point of the breeze decoder. Add a round_bf16 compute shader (f32/f16/bf16 in, always f32 out, round-to-nearest-even via the same fp32_to_bf16 bit trick the cpy shaders use), register pipelines indexed by source type, handle the widened f32 dst in the unary pipeline selection and op-support checks, and enable fused_round for Vulkan in the breeze activation-cast policy. Verified bit-identical breeze reference codes vs the cast round trip at 6 s and 15 s references. Peak VRAM on a 2080 Ti drops ~250 MiB at a 60 s reference (5491 -> 5239 MiB); no measurable change at 6 s. * ggml-vulkan: handle row-strided inputs in fused round-to-bf16 The breeze activation-rounding policy admits row-strided views into ggml_round_bf16 (ggml_is_contiguous_rows gate in qwen_decoder). The Vulkan port dispatched every input to the flat shader, which indexes the source as a contiguous array, so row-strided views read garbage and clone output degenerated into noise. Route non-contiguous inputs to a new round_bf16_strided shader built on generic_unary_head (same pattern as sigmoid_strided), keeping the flat fast path for contiguous inputs. --- external/ggml/src/ggml-vulkan/ggml-vulkan.cpp | 39 +++ .../vulkan-shaders/round_bf16.comp | 26 ++ .../vulkan-shaders/round_bf16_strided.comp | 23 ++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 6 + .../engine/models/breeze_tts/speech_encoder.h | 6 +- src/models/breeze_tts/generator.cpp | 6 +- src/models/breeze_tts/speech_encoder.cpp | 264 +++++++++++++++--- 7 files changed, 319 insertions(+), 51 deletions(-) create mode 100644 external/ggml/src/ggml-vulkan/vulkan-shaders/round_bf16.comp create mode 100644 external/ggml/src/ggml-vulkan/vulkan-shaders/round_bf16_strided.comp diff --git a/external/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/external/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 3e9a5217b..120f484e5 100644 --- a/external/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/external/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -811,6 +811,8 @@ struct vk_device_struct { vk_pipeline pipeline_softplus[2]; vk_pipeline pipeline_step[2]; vk_pipeline pipeline_round[2]; + vk_pipeline pipeline_round_bf16[3]; + vk_pipeline pipeline_round_bf16_strided[3]; vk_pipeline pipeline_ceil[2]; vk_pipeline pipeline_floor[2]; vk_pipeline pipeline_trunc[2]; @@ -4750,6 +4752,15 @@ static void ggml_vk_load_shaders(vk_device& device) { CREATE_UNARY(exp) #undef CREATE_UNARY + // round-to-bf16: f32/f16/bf16 in, always f32 out (index by src type). + ggml_vk_create_pipeline(device, device->pipeline_round_bf16[0], "round_bf16_f32", round_bf16_f32_len, round_bf16_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_round_bf16[1], "round_bf16_f16", round_bf16_f16_len, round_bf16_f16_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_round_bf16[2], "round_bf16_bf16", round_bf16_bf16_len, round_bf16_bf16_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); + // strided variant for non-contiguous (e.g. row-strided view) inputs. + ggml_vk_create_pipeline(device, device->pipeline_round_bf16_strided[0], "round_bf16_strided_f32", round_bf16_strided_f32_len, round_bf16_strided_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_round_bf16_strided[1], "round_bf16_strided_f16", round_bf16_strided_f16_len, round_bf16_strided_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_round_bf16_strided[2], "round_bf16_strided_bf16", round_bf16_strided_bf16_len, round_bf16_strided_bf16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_add1_f16_f16, "add1_f16_f16", add1_f16_f16_len, add1_f16_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_add1_f16_f32, "add1_f16_f32", add1_f16_f32_len, add1_f16_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_add1_f32_f32, "add1_f32_f32", add1_f32_f32_len, add1_f32_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); @@ -9740,6 +9751,19 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const } return nullptr; case GGML_OP_UNARY: + // ROUND_BF16 widens to f32: src may be f32/f16/bf16 while dst is f32. + if (ggml_get_unary_op(dst) == GGML_UNARY_OP_ROUND_BF16) { + if (dst->type != GGML_TYPE_F32) { + return nullptr; + } + const bool strided = !ggml_is_contiguous(src0) || !ggml_is_contiguous(dst); + switch (src0->type) { + case GGML_TYPE_F32: return strided ? ctx->device->pipeline_round_bf16_strided[0] : ctx->device->pipeline_round_bf16[0]; + case GGML_TYPE_F16: return strided ? ctx->device->pipeline_round_bf16_strided[1] : ctx->device->pipeline_round_bf16[1]; + case GGML_TYPE_BF16: return strided ? ctx->device->pipeline_round_bf16_strided[2] : ctx->device->pipeline_round_bf16[2]; + default: return nullptr; + } + } if ((src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16) || (dst->type != GGML_TYPE_F32 && dst->type != GGML_TYPE_F16) || (src0->type != dst->type)) { @@ -11481,6 +11505,11 @@ static void ggml_vk_sigmoid_strided(ggml_backend_vk_context * ctx, vk_context& s ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_UNARY, std::move(p)); } +static void ggml_vk_round_bf16_strided(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) { + vk_op_unary_push_constants p = vk_op_unary_push_constants_init(src0, dst); + ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_UNARY, std::move(p)); +} + static void ggml_vk_xielu(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) { float * op_params = (float *)dst->op_params; ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_UNARY, @@ -13522,6 +13551,13 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr case GGML_UNARY_OP_SGN: ggml_vk_unary(ctx, compute_ctx, src0, node); break; + case GGML_UNARY_OP_ROUND_BF16: + if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(node)) { + ggml_vk_round_bf16_strided(ctx, compute_ctx, src0, node); + break; + } + ggml_vk_unary(ctx, compute_ctx, src0, node); + break; case GGML_UNARY_OP_SIGMOID: if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(node)) { ggml_vk_sigmoid_strided(ctx, compute_ctx, src0, node); @@ -15772,6 +15808,9 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) && (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && (op->src[0]->type == op->type); + case GGML_UNARY_OP_ROUND_BF16: + return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_BF16) && + (op->type == GGML_TYPE_F32); case GGML_UNARY_OP_SIGMOID: return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) && (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && diff --git a/external/ggml/src/ggml-vulkan/vulkan-shaders/round_bf16.comp b/external/ggml/src/ggml-vulkan/vulkan-shaders/round_bf16.comp new file mode 100644 index 000000000..b92937bec --- /dev/null +++ b/external/ggml/src/ggml-vulkan/vulkan-shaders/round_bf16.comp @@ -0,0 +1,26 @@ +#version 450 + +#include "generic_head.glsl" +#include "types.glsl" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; + +void main() { + const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; + + if (i >= p.KX) { + return; + } + +#if defined(DATA_A_BF16) + const float x = bf16_to_fp32(uint32_t(data_a[i])); +#else + const float x = float(data_a[i]); +#endif + // Round to bf16 precision and widen back to f32, matching the + // f32 -> bf16 -> f32 cast round trip. + data_d[i] = D_TYPE(bf16_to_fp32(fp32_to_bf16(x))); +} diff --git a/external/ggml/src/ggml-vulkan/vulkan-shaders/round_bf16_strided.comp b/external/ggml/src/ggml-vulkan/vulkan-shaders/round_bf16_strided.comp new file mode 100644 index 000000000..d5eaa5087 --- /dev/null +++ b/external/ggml/src/ggml-vulkan/vulkan-shaders/round_bf16_strided.comp @@ -0,0 +1,23 @@ +#version 450 + +#include "types.glsl" +#include "generic_unary_head.glsl" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint idx = get_idx(); + + if (idx >= p.ne) { + return; + } + +#if defined(DATA_A_BF16) + const float x = bf16_to_fp32(uint32_t(data_a[get_aoffset() + src0_idx(idx)])); +#else + const float x = float(data_a[get_aoffset() + src0_idx(idx)]); +#endif + // Round to bf16 precision and widen back to f32, matching the + // f32 -> bf16 -> f32 cast round trip. + data_d[get_doffset() + dst_idx(idx)] = D_TYPE(bf16_to_fp32(fp32_to_bf16(x))); +} diff --git a/external/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/external/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index a533af69f..40c223470 100644 --- a/external/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/external/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -880,6 +880,12 @@ void process_shaders() { string_to_spv("step_f32", "step.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); string_to_spv("round_f16", "round.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}}); string_to_spv("round_f32", "round.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("round_bf16_f32", "round_bf16.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("round_bf16_f16", "round_bf16.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float"}}); + string_to_spv("round_bf16_bf16", "round_bf16.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "float"}, {"DATA_A_BF16", "1"}}); + string_to_spv("round_bf16_strided_f32", "round_bf16_strided.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("round_bf16_strided_f16", "round_bf16_strided.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float"}}); + string_to_spv("round_bf16_strided_bf16", "round_bf16_strided.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "float"}, {"DATA_A_BF16", "1"}}); string_to_spv("ceil_f16", "ceil.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}}); string_to_spv("ceil_f32", "ceil.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); string_to_spv("floor_f16", "floor.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}}); diff --git a/include/engine/models/breeze_tts/speech_encoder.h b/include/engine/models/breeze_tts/speech_encoder.h index 60da8d924..f2dbae6c4 100644 --- a/include/engine/models/breeze_tts/speech_encoder.h +++ b/include/engine/models/breeze_tts/speech_encoder.h @@ -19,7 +19,8 @@ namespace engine::models { namespace breeze_tts { struct BreezeSpeechEncoderWeights; -class BreezeSpeechEncoderGraph; +class BreezeSpeechEncoderConvGraph; +class BreezeSpeechEncoderTransformerGraph; struct BreezeSpeechEncoderOutput { BreezeSpeechCodes codes; @@ -46,7 +47,8 @@ class BreezeSpeechEncoderRuntime { core::ExecutionContext * execution_context_ = nullptr; size_t graph_arena_bytes_ = 0; std::unique_ptr constants_; - mutable std::unique_ptr graph_; + mutable std::unique_ptr conv_graph_; + mutable std::unique_ptr transformer_graph_; }; } // namespace breeze_tts diff --git a/src/models/breeze_tts/generator.cpp b/src/models/breeze_tts/generator.cpp index 83f149540..28458f122 100644 --- a/src/models/breeze_tts/generator.cpp +++ b/src/models/breeze_tts/generator.cpp @@ -60,9 +60,9 @@ modules::QwenDecoderActivationCastPolicy breeze_bf16_activation_policy(core::Bac } policy.enabled = true; policy.type = GGML_TYPE_BF16; - // CUDA/HIP implement the fused round-to-bf16 unary op; Vulkan does not and - // keeps the cast round trip. - policy.fused_round = backend_type == core::BackendType::Cuda || backend_type == core::BackendType::Hip; + // CUDA/HIP/Vulkan implement the fused round-to-bf16 unary op. + policy.fused_round = backend_type == core::BackendType::Cuda || backend_type == core::BackendType::Hip || + backend_type == core::BackendType::Vulkan; policy.after_input_norm = true; policy.after_qkv_projection = true; policy.after_qk_norm = true; diff --git a/src/models/breeze_tts/speech_encoder.cpp b/src/models/breeze_tts/speech_encoder.cpp index f4f7d34a5..b1d59b400 100644 --- a/src/models/breeze_tts/speech_encoder.cpp +++ b/src/models/breeze_tts/speech_encoder.cpp @@ -73,6 +73,23 @@ constexpr modules::Conv1dConfig kDownsampleConvConfig{512, 512, 4, 2, 0, 1, fals constexpr modules::Conv1dConfig kSemanticProjectionConfig{512, 256, 1, 1, 0, 1, false}; constexpr modules::Conv1dConfig kAcousticProjectionConfig{512, 256, 1, 1, 0, 1, false}; +// The conv stack downsamples 960x before the transformer (strides 4*5*6*8). +constexpr int64_t kTransformerStride = 960; +// Conv chunks are kChunkSamples of new audio preceded by kChunkOverlapSamples of +// left context. The exact left context the stack needs is the sum of each +// layer's left pad scaled by the cumulative stride: +// 6+2 + 4+2*4 + 5*4+2*20 + 6*20+2*120 + 8*120+2*960 + 2*960 = 5240 samples. +constexpr int64_t kChunkSamples = 120000; // 5 s at 24 kHz, 125 transformer frames +constexpr int64_t kChunkOverlapSamples = 9600; // 10 transformer frames > 5240 +constexpr int64_t kChunkCapacity = kChunkSamples + kChunkOverlapSamples; +static_assert(kChunkSamples % kTransformerStride == 0); +static_assert(kChunkOverlapSamples % kTransformerStride == 0); +constexpr int64_t kChunkFrames = kChunkCapacity / kTransformerStride; +// The transformer graph is built at frame capacities rounded up to this +// bucket, so reference lengths within a bucket share one graph instead of +// rebuilding per exact length. +constexpr int64_t kTransformerFrameBucket = kChunkSamples / kTransformerStride; + struct GgmlContextDeleter { void operator()(ggml_context * ctx) const noexcept { if (ctx != nullptr) { @@ -415,24 +432,23 @@ std::shared_ptr load_weights( return weights; } -class BreezeSpeechEncoderGraph { +// Conv stack runs on fixed-size chunks (plus left overlap), so its graph +// memory is constant regardless of reference length. Chunk outputs stitch +// exactly: chunk lengths are multiples of kTransformerStride, so no per-stage +// right padding occurs, and discarded overlap frames absorb the zero left +// pads that represent audio start in the first chunk. +class BreezeSpeechEncoderConvGraph { public: - BreezeSpeechEncoderGraph( + BreezeSpeechEncoderConvGraph( std::shared_ptr weights, - int64_t sample_capacity, core::ExecutionContext & execution_context, core::ConstantTensorCache & constants, size_t graph_arena_bytes) : weights_(std::move(weights)), - sample_capacity_(sample_capacity), - frames_((sample_capacity + kDownsampleRate - 1) / kDownsampleRate), backend_(execution_context.backend()), compute_threads_(std::max(1, execution_context.config().threads)) { if (weights_ == nullptr) { - throw std::runtime_error("Breeze speech encoder graph requires weights"); - } - if (sample_capacity_ <= 0) { - throw std::runtime_error("Breeze speech encoder graph requires positive sample capacity"); + throw std::runtime_error("Breeze speech encoder conv graph requires weights"); } if (backend_ == nullptr) { throw std::runtime_error("Breeze speech encoder backend is not initialized"); @@ -445,15 +461,15 @@ class BreezeSpeechEncoderGraph { }; ctx_.reset(ggml_init(params)); if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize Breeze speech encoder ggml context"); + throw std::runtime_error("failed to initialize Breeze speech encoder conv ggml context"); } core::ModuleBuildContext build_ctx{ ctx_.get(), - "breeze_tts.speech_encoder", + "breeze_tts.speech_encoder.conv", execution_context.backend_type(), }; - auto x = core::make_tensor(build_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, sample_capacity_})); + auto x = core::make_tensor(build_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, kChunkCapacity})); input_ = x.tensor; constants.begin_graph(); @@ -465,25 +481,123 @@ class BreezeSpeechEncoderGraph { } x = modules::EluModule{}.build(build_ctx, x); x = speech_conv(build_ctx, x, weights_->encoder_convs.back(), kEncoderConvConfigs.back(), kEncoderConvPadModes.back()); - auto seq = modules::TransposeModule({{0, 2, 1, 3}, x.shape.rank}).build(build_ctx, x); seq = core::ensure_backend_addressable_layout(build_ctx, seq); - transformer_frames_ = seq.shape.dims[1]; - positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, transformer_frames_); - auto positions_value = core::wrap_tensor(positions_, core::TensorShape::from_dims({transformer_frames_}), GGML_TYPE_I32); - attention_mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, transformer_frames_, transformer_frames_, 1, 1); + output_ = seq.tensor; + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), 32768, false); + ggml_build_forward_expand(graph_, output_); + constants.finish_graph(); + constants.ensure_uploaded(); + + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + if (gallocr_ == nullptr || !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + throw std::runtime_error("failed to allocate Breeze speech encoder conv graph"); + } + } + + ~BreezeSpeechEncoderConvGraph() { + engine::core::release_backend_graph_resources(backend_, graph_); + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + } + + bool matches(const BreezeSpeechEncoderWeights & weights, ggml_backend_t backend, int threads) const { + return weights_.get() == &weights && backend_ == backend && compute_threads_ == std::max(1, threads); + } + + std::vector run(const std::vector & chunk_input) { + if (static_cast(chunk_input.size()) != kChunkCapacity) { + throw std::runtime_error("Breeze speech encoder conv chunk size mismatch"); + } + ggml_backend_tensor_set(input_, chunk_input.data(), 0, chunk_input.size() * sizeof(float)); + core::set_backend_threads(backend_, compute_threads_); + const ggml_status status = engine::core::compute_backend_graph(backend_, graph_); + ggml_backend_synchronize(backend_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Breeze speech encoder conv graph compute failed"); + } + std::vector features(static_cast(kHiddenSize * kChunkFrames)); + ggml_backend_tensor_get(output_, features.data(), 0, features.size() * sizeof(float)); + return features; + } + +private: + std::shared_ptr weights_; + std::unique_ptr ctx_; + ggml_tensor * input_ = nullptr; + ggml_tensor * output_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_backend_t backend_ = nullptr; + int compute_threads_ = 1; + ggml_gallocr_t gallocr_ = nullptr; +}; + +// Transformer + downsample + projections run once over the full frame +// sequence; at frame scale (960x downsampled) this graph is a few tens of MiB +// even for minute-long references. Attention is causal, so frames computed +// from right-padded tail regions never affect earlier frames. +class BreezeSpeechEncoderTransformerGraph { +public: + BreezeSpeechEncoderTransformerGraph( + std::shared_ptr weights, + int64_t frames, + core::ExecutionContext & execution_context, + core::ConstantTensorCache & constants, + size_t graph_arena_bytes) + : weights_(std::move(weights)), + frames_(frames), + backend_(execution_context.backend()), + compute_threads_(std::max(1, execution_context.config().threads)) { + if (weights_ == nullptr) { + throw std::runtime_error("Breeze speech encoder transformer graph requires weights"); + } + if (frames_ <= 0) { + throw std::runtime_error("Breeze speech encoder transformer graph requires positive frame count"); + } + if (backend_ == nullptr) { + throw std::runtime_error("Breeze speech encoder backend is not initialized"); + } + + ggml_init_params params{ + /*.mem_size =*/ graph_arena_bytes, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Breeze speech encoder transformer ggml context"); + } + + core::ModuleBuildContext build_ctx{ + ctx_.get(), + "breeze_tts.speech_encoder.transformer", + execution_context.backend_type(), + }; + auto seq = core::make_tensor( + build_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, frames_, kHiddenSize})); + input_ = seq.tensor; + + constants.begin_graph(); + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, frames_); + auto positions_value = core::wrap_tensor(positions_, core::TensorShape::from_dims({frames_}), GGML_TYPE_I32); + attention_mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, frames_, frames_, 1, 1); const auto attention_mask = core::wrap_tensor( attention_mask_, - core::TensorShape::from_dims({1, 1, transformer_frames_, transformer_frames_}), + core::TensorShape::from_dims({1, 1, frames_, frames_}), GGML_TYPE_F16); for (const auto & layer : weights_->transformer_layers) { seq = transformer_block(build_ctx, seq, positions_value, layer, attention_mask); } - x = modules::TransposeModule({{0, 2, 1, 3}, seq.shape.rank}).build(build_ctx, seq); + auto x = modules::TransposeModule({{0, 2, 1, 3}, seq.shape.rank}).build(build_ctx, seq); x = core::ensure_backend_addressable_layout(build_ctx, x); x = speech_conv(build_ctx, x, weights_->downsample, kDownsampleConvConfig, modules::StreamingPadMode::Replicate); auto semantic = speech_conv(build_ctx, x, weights_->semantic_projection, kSemanticProjectionConfig, modules::StreamingPadMode::Constant); auto acoustic = speech_conv(build_ctx, x, weights_->acoustic_projection, kAcousticProjectionConfig, modules::StreamingPadMode::Constant); + output_frames_ = semantic.shape.dims[2]; semantic_output_ = semantic.tensor; acoustic_output_ = acoustic.tensor; ggml_set_output(semantic_output_); @@ -496,55 +610,53 @@ class BreezeSpeechEncoderGraph { gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); if (gallocr_ == nullptr || !ggml_gallocr_alloc_graph(gallocr_, graph_)) { - throw std::runtime_error("failed to allocate Breeze speech encoder graph"); + throw std::runtime_error("failed to allocate Breeze speech encoder transformer graph"); } - positions_data_.resize(static_cast(transformer_frames_)); - for (int64_t i = 0; i < transformer_frames_; ++i) { + positions_data_.resize(static_cast(frames_)); + for (int64_t i = 0; i < frames_; ++i) { positions_data_[static_cast(i)] = static_cast(i); } if (attention_mask_ != nullptr) { - auto mask = modules::qwen_causal_prefill_mask_values(1, transformer_frames_); + auto mask = modules::qwen_causal_prefill_mask_values(1, frames_); attention_mask_data_ = std::move(mask); } upload_static_inputs(); } - ~BreezeSpeechEncoderGraph() { + ~BreezeSpeechEncoderTransformerGraph() { engine::core::release_backend_graph_resources(backend_, graph_); if (gallocr_ != nullptr) { ggml_gallocr_free(gallocr_); } } - bool matches(const BreezeSpeechEncoderWeights & weights, int64_t samples, ggml_backend_t backend, int threads) const { - return weights_.get() == &weights && sample_capacity_ == samples && backend_ == backend && + bool matches(const BreezeSpeechEncoderWeights & 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); } - BreezeSpeechEncoderOutput run(const std::vector & waveform) { - if (static_cast(waveform.size()) > sample_capacity_) { - throw std::runtime_error("Breeze speech encoder waveform exceeds graph capacity"); + BreezeSpeechEncoderOutput run(const std::vector & features) { + if (static_cast(features.size()) != kHiddenSize * frames_) { + throw std::runtime_error("Breeze speech encoder transformer input size mismatch"); } upload_static_inputs(); - std::vector padded(static_cast(sample_capacity_), 0.0F); - std::copy(waveform.begin(), waveform.end(), padded.begin()); - ggml_backend_tensor_set(input_, padded.data(), 0, padded.size() * sizeof(float)); + ggml_backend_tensor_set(input_, features.data(), 0, features.size() * sizeof(float)); core::set_backend_threads(backend_, compute_threads_); const ggml_status status = engine::core::compute_backend_graph(backend_, graph_); ggml_backend_synchronize(backend_); if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("Breeze speech encoder graph compute failed"); + throw std::runtime_error("Breeze speech encoder transformer graph compute failed"); } BreezeSpeechEncoderOutput out; - out.semantic_projected.resize(static_cast(kQuantizerDim * frames_)); - out.acoustic_projected.resize(static_cast(kQuantizerDim * frames_)); + out.semantic_projected.resize(static_cast(kQuantizerDim * output_frames_)); + out.acoustic_projected.resize(static_cast(kQuantizerDim * output_frames_)); ggml_backend_tensor_get(semantic_output_, out.semantic_projected.data(), 0, out.semantic_projected.size() * sizeof(float)); ggml_backend_tensor_get(acoustic_output_, out.acoustic_projected.data(), 0, out.acoustic_projected.size() * sizeof(float)); return out; } - int64_t frames() const noexcept { - return frames_; + int64_t output_frames() const noexcept { + return output_frames_; } private: @@ -564,9 +676,8 @@ class BreezeSpeechEncoderGraph { } std::shared_ptr weights_; - int64_t sample_capacity_ = 0; int64_t frames_ = 0; - int64_t transformer_frames_ = 0; + int64_t output_frames_ = 0; std::unique_ptr ctx_; ggml_tensor * input_ = nullptr; ggml_tensor * positions_ = nullptr; @@ -623,18 +734,78 @@ BreezeSpeechCodes BreezeSpeechEncoderRuntime::encode(const runtime::AudioBuffer static_cast(kSampleRate)); const int64_t valid_samples = static_cast(waveform.size()); const int64_t frames = std::max(1, (valid_samples + kDownsampleRate - 1) / kDownsampleRate); - const int64_t sample_capacity = valid_samples; + const int64_t transformer_frames = (valid_samples + kTransformerStride - 1) / kTransformerStride; + const int64_t graph_frames = + (transformer_frames + kTransformerFrameBucket - 1) / kTransformerFrameBucket * kTransformerFrameBucket; const int threads = std::max(1, execution_context_->config().threads); - if (graph_ == nullptr || !graph_->matches(*weights_, sample_capacity, execution_context_->backend(), threads)) { - graph_.reset(); - graph_ = std::make_unique( + if (conv_graph_ == nullptr || !conv_graph_->matches(*weights_, execution_context_->backend(), threads)) { + conv_graph_.reset(); + conv_graph_ = std::make_unique( + weights_, + *execution_context_, + *constants_, + graph_arena_bytes_); + } + if (transformer_graph_ == nullptr || + !transformer_graph_->matches(*weights_, graph_frames, execution_context_->backend(), threads)) { + transformer_graph_.reset(); + transformer_graph_ = std::make_unique( weights_, - sample_capacity, + graph_frames, *execution_context_, *constants_, graph_arena_bytes_); } - auto out = graph_->run(waveform); + + std::vector features(static_cast(kHiddenSize * graph_frames)); + std::vector chunk_input(static_cast(kChunkCapacity)); + int64_t dst_frame = 0; + for (int64_t pos = 0; pos < valid_samples; pos += kChunkSamples) { + const int64_t overlap = pos > 0 ? kChunkOverlapSamples : 0; + const int64_t fresh = std::min(kChunkSamples, valid_samples - pos); + std::fill(chunk_input.begin(), chunk_input.end(), 0.0F); + std::copy( + waveform.begin() + (pos - overlap), + waveform.begin() + (pos + fresh), + chunk_input.begin()); + const auto chunk_features = conv_graph_->run(chunk_input); + const int64_t skip_frames = overlap / kTransformerStride; + const int64_t keep_frames = (fresh + kTransformerStride - 1) / kTransformerStride; + std::copy_n( + chunk_features.begin() + skip_frames * kHiddenSize, + keep_frames * kHiddenSize, + features.begin() + dst_frame * kHiddenSize); + dst_frame += keep_frames; + } + // Pad unused bucket frames with the last real frame (not zeros): causal + // attention keeps padding invisible to real frames, and the downsample + // conv's Replicate right pad then sees the same value as an exact-length + // graph would produce. + for (int64_t f = transformer_frames; f < graph_frames; ++f) { + std::copy_n( + features.begin() + (transformer_frames - 1) * kHiddenSize, + kHiddenSize, + features.begin() + f * kHiddenSize); + } + + auto out = transformer_graph_->run(features); + const int64_t produced_frames = transformer_graph_->output_frames(); + if (produced_frames != frames) { + // Projected outputs are channel-major with stride produced_frames; + // drop the padding frames before quantization. + auto slice_frames = [frames, produced_frames](std::vector & projected) { + std::vector sliced(static_cast(kQuantizerDim * frames)); + for (int64_t dim = 0; dim < kQuantizerDim; ++dim) { + std::copy_n( + projected.begin() + dim * produced_frames, + frames, + sliced.begin() + dim * frames); + } + projected = std::move(sliced); + }; + slice_frames(out.semantic_projected); + slice_frames(out.acoustic_projected); + } out.codes.frames = frames; out.codes.code_groups = kValidQuantizers; out.codes.codes = quantize_projected(out.semantic_projected, out.acoustic_projected, frames, *weights_); @@ -643,7 +814,8 @@ BreezeSpeechCodes BreezeSpeechEncoderRuntime::encode(const runtime::AudioBuffer } void BreezeSpeechEncoderRuntime::release_runtime_graphs() const { - graph_.reset(); + conv_graph_.reset(); + transformer_graph_.reset(); } } // namespace engine::models::breeze_tts From 7f62b32086a9019b1ec3c84213d0133250334092 Mon Sep 17 00:00:00 2001 From: Th-Underscore Date: Wed, 2 Sep 2026 21:29:46 -0400 Subject: [PATCH 09/10] Attention fallback for GPUs without flash MMA kernels (sm70) Auto-resolve flash vs eager attention from CUDA compute capability: Volta/Turing (700 <= cc < 800) fall back to eager, since large prefill shapes select the MMA kernel which has no usable device code there ('flash_attn_ext_f16 has no device code compatible with CUDA arch 700'). - New engine::core::attention_fallback unit: preference parsing (per-model '.attention' session option + AUDIOCPP_ATTENTION), CC gating via the CUDA driver (supports_op cannot detect this: it returns true on sm70 for shapes that later crash at launch). - Wire auto fallback + session options into higgs_audio_tts and breeze_tts (backbone, depth, encoder, decoder); process-wide AUDIOCPP_ATTENTION=eager backstop in the shared SDPA/GQA/QwenDecoder modules; trace logging of the resolved path. - Fix latent QwenDecoder prefix-concat dtype assert exposed by the eager path (cast cached prefix KV on every path, not just flash). - Unit test, breeze_tts model-spec entry, docs. --- CMakeLists.txt | 9 ++ docs/models/breeze_tts.md | 2 +- docs/tts.md | 1 + .../framework/core/attention_fallback.h | 44 +++++++ .../modules/transformers/qwen_decoder.h | 3 + include/engine/models/breeze_tts/generator.h | 4 +- .../engine/models/breeze_tts/speech_decoder.h | 5 +- .../engine/models/breeze_tts/speech_encoder.h | 5 +- include/engine/models/higgs_audio_tts/ar.h | 6 +- model_specs/breeze_tts.json | 8 ++ src/framework/core/attention_fallback.cpp | 121 ++++++++++++++++++ .../modules/transformers/qwen_decoder.cpp | 32 +++-- src/models/breeze_tts/generator.cpp | 52 ++++++-- src/models/breeze_tts/session.cpp | 24 +++- src/models/breeze_tts/speech_decoder.cpp | 30 ++++- src/models/breeze_tts/speech_encoder.cpp | 33 +++-- src/models/higgs_audio_tts/ar.cpp | 45 +++++-- src/models/higgs_audio_tts/loader.cpp | 1 + src/models/higgs_audio_tts/session.cpp | 25 +++- tests/unittests/test_attention_fallback.cpp | 68 ++++++++++ 20 files changed, 462 insertions(+), 56 deletions(-) create mode 100644 include/engine/framework/core/attention_fallback.h create mode 100644 src/framework/core/attention_fallback.cpp create mode 100644 tests/unittests/test_attention_fallback.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index a5abb3ba6..bf9092e22 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -353,6 +353,7 @@ add_library(engine_core OBJECT src/framework/assets/torch_bin.cpp src/framework/core/module.cpp src/framework/core/backend.cpp + src/framework/core/attention_fallback.cpp src/framework/core/deferred_tensor_writer.cpp src/framework/core/execution_context.cpp src/framework/core/host_memory.cpp @@ -2213,6 +2214,14 @@ if (ENGINE_BUILD_TESTS) COMMAND backend_device_resolution_test ) + add_engine_unittest(attention_fallback_test tests/unittests/test_attention_fallback.cpp) + target_include_directories(attention_fallback_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tests/unittests) + + add_test( + NAME attention_fallback_test + COMMAND attention_fallback_test + ) + add_executable(streaming_audio_input_test tests/unittests/test_streaming_audio_input.cpp app/streaming/pcm_source.cpp diff --git a/docs/models/breeze_tts.md b/docs/models/breeze_tts.md index 9aca60d01..584713bb8 100644 --- a/docs/models/breeze_tts.md +++ b/docs/models/breeze_tts.md @@ -67,6 +67,7 @@ audiocpp_cli \ | `--request-option top_p=` | `0..1` | `1.0` | Top-p sampling limit. | | `--request-option seed=` | integer >= 0 | `0` | Generation seed. | | `--session-option breeze_tts.reference_cache_slots=` | integer >= 0 | `1` | Prepared reference-audio cache slots. | +| `--session-option breeze_tts.attention=` | `auto`, `flash`, `eager` | `auto` | Attention kernel. `auto` uses flash except on Volta/Turing GPUs (e.g. V100), where it falls back to eager to avoid missing MMA kernels. | | `--session-option weight_type=` | `native`, `f32`, `f16`, `bf16`, `q8_0`, `q4_0`, `q4_k` | `native` | Weight storage type; quantized types convert at load time from the BF16 package. | Quantized weight storage is the largest measured speedup and applies to CUDA @@ -76,4 +77,3 @@ and HIP alike: `q8_0` cut the fixed 100-token regression case from RTF ~1.5 to voice-design regression cases. Counter to intuition, fp32 is the one configuration known to be *worse* for this model (mispronunciations and runaway repetition), because the model is trained and tuned in bf16. - diff --git a/docs/tts.md b/docs/tts.md index 4b2910fdf..87af27914 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -559,6 +559,7 @@ python3 tools/model_manager_v2.py install --models-root models higgs_audio_tts_4 | `--top-k` | integer | `30` | AR top-k sampling limit. The narrower default is less prone to premature EOC than the Python client's `50`. | | `--top-p` | float | `0.8` | AR nucleus sampling limit. The Python client's unfiltered equivalent is `1.0`. | | `--repetition-penalty` | float | `1.1` | Accepted for Python API compatibility; Higgs audio-code sampling does not consume it. | +| `--session-option higgs_audio_tts.attention=` | `auto`, `flash`, `eager` | `auto` | Attention kernel. `auto` uses flash except on Volta/Turing GPUs (e.g. V100), where it falls back to eager to avoid missing MMA kernels. | ## Fish Audio S2 Pro diff --git a/include/engine/framework/core/attention_fallback.h b/include/engine/framework/core/attention_fallback.h new file mode 100644 index 000000000..0a238aadf --- /dev/null +++ b/include/engine/framework/core/attention_fallback.h @@ -0,0 +1,44 @@ +#pragma once + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include + +namespace engine::core { + +// Session-option vocabulary shared by families that lower attention with +// ggml_flash_attn_ext: "auto" (default), "flash", or "eager". +// +// Background: ggml-cuda only instantiates the MMA/wmma flash-attention kernels +// for compute capability >= 8.0 (Ampere). On older GPUs such as Volta/sm70 +// (e.g. Tesla V100) a graph containing GGML_OP_FLASH_ATTN_EXT fails at compute +// time with "no device code compatible with CUDA arch 700", even when ggml was +// built with that arch enabled. The eager (explicit matmul + softmax) lowering +// computes the same operation with generic ops and runs everywhere (output +// logits may differ at ulp level, as with any kernel change). +enum class AttentionPreference { + Auto, + Flash, + Eager, +}; + +// Parses a ".attention" session-option value. Throws std::runtime_error +// naming option_name on invalid input. +AttentionPreference parse_attention_preference(const std::string & value, const char * option_name); + +// Resolves whether flash attention may be used for the given backend and head +// dimension. Flash forces true, Eager forces false, Auto gates on the CUDA +// device compute capability (Volta/Turing resolve to eager; see the .cpp for +// why supports_op cannot be used). A null or non-CUDA backend, or a device +// query failure, preserves historical behavior (true). +// +// Adopting in other families (currently wired for higgs_audio_tts and +// breeze_tts only): resolve once per runtime with the model's head_dim and +// the family's ".attention" session option, then switch the +// QwenDecoder prefill/static modes (or SDPA/GQA lowerings) between flash and +// their ManualRepeat/Explicit equivalents based on the result. +bool resolve_flash_attention(ggml_backend_t backend, int64_t head_dim, AttentionPreference preference); + +} // namespace engine::core diff --git a/include/engine/framework/modules/transformers/qwen_decoder.h b/include/engine/framework/modules/transformers/qwen_decoder.h index eaa80c12e..fa1c5ac99 100644 --- a/include/engine/framework/modules/transformers/qwen_decoder.h +++ b/include/engine/framework/modules/transformers/qwen_decoder.h @@ -77,6 +77,9 @@ struct QwenDecoderAttentionPolicy { QwenDecoderAttentionMode static_mode = QwenDecoderAttentionMode::FlashGrouped; QwenDecoderPrefixAttentionMode prefix_mode = QwenDecoderPrefixAttentionMode::Exact; int64_t grouped_query_min_steps = 0; + // False routes flash branches through repeat-KV + matmul/softmax for GPUs + // without a flash kernel (e.g. CUDA sm70). True preserves historical behavior. + bool allow_flash_attention = true; }; struct QwenDecoderStaticCachePolicy { diff --git a/include/engine/models/breeze_tts/generator.h b/include/engine/models/breeze_tts/generator.h index 7b3a95ff7..7156ea6b5 100644 --- a/include/engine/models/breeze_tts/generator.h +++ b/include/engine/models/breeze_tts/generator.h @@ -1,5 +1,6 @@ #pragma once +#include "engine/framework/core/attention_fallback.h" #include "engine/framework/runtime/session.h" #include "engine/models/breeze_tts/assets.h" #include "engine/models/breeze_tts/speech_decoder.h" @@ -35,7 +36,8 @@ class BreezeGeneratorRuntime { engine::core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - engine::assets::TensorStorageType storage_type); + engine::assets::TensorStorageType storage_type, + engine::core::AttentionPreference attention_preference = engine::core::AttentionPreference::Auto); ~BreezeGeneratorRuntime(); engine::runtime::AudioBuffer generate(const BreezeGenerationRequest & request); diff --git a/include/engine/models/breeze_tts/speech_decoder.h b/include/engine/models/breeze_tts/speech_decoder.h index bb27bc33d..a5c269c6d 100644 --- a/include/engine/models/breeze_tts/speech_decoder.h +++ b/include/engine/models/breeze_tts/speech_decoder.h @@ -1,5 +1,6 @@ #pragma once +#include "engine/framework/core/attention_fallback.h" #include "engine/framework/core/execution_context.h" #include "engine/framework/assets/tensor_source.h" #include "engine/framework/runtime/session.h" @@ -36,7 +37,8 @@ class BreezeSpeechDecoderRuntime { size_t graph_arena_bytes, size_t constant_context_bytes, engine::assets::TensorStorageType linear_weight_storage_type, - engine::assets::TensorStorageType conv_weight_storage_type); + engine::assets::TensorStorageType conv_weight_storage_type, + core::AttentionPreference attention_preference = core::AttentionPreference::Auto); ~BreezeSpeechDecoderRuntime(); runtime::AudioBuffer decode(const BreezeSpeechCodes & codec_codes) const; @@ -50,6 +52,7 @@ class BreezeSpeechDecoderRuntime { core::ExecutionContext * execution_context_ = nullptr; std::shared_ptr weights_; size_t graph_arena_bytes_ = 0; + bool allow_flash_attention_ = true; std::unique_ptr constants_; mutable std::unique_ptr graph_; // Always present to keep this public class layout identical when the private diff --git a/include/engine/models/breeze_tts/speech_encoder.h b/include/engine/models/breeze_tts/speech_encoder.h index f2dbae6c4..818a8cb2a 100644 --- a/include/engine/models/breeze_tts/speech_encoder.h +++ b/include/engine/models/breeze_tts/speech_encoder.h @@ -1,5 +1,6 @@ #pragma once +#include "engine/framework/core/attention_fallback.h" #include "engine/framework/core/execution_context.h" #include "engine/framework/assets/tensor_source.h" #include "engine/framework/runtime/session.h" @@ -35,7 +36,8 @@ class BreezeSpeechEncoderRuntime { core::ExecutionContext & execution_context, size_t graph_arena_bytes, engine::assets::TensorStorageType linear_weight_storage_type, - engine::assets::TensorStorageType conv_weight_storage_type); + engine::assets::TensorStorageType conv_weight_storage_type, + core::AttentionPreference attention_preference = core::AttentionPreference::Auto); ~BreezeSpeechEncoderRuntime(); BreezeSpeechCodes encode(const runtime::AudioBuffer & audio) const; @@ -46,6 +48,7 @@ class BreezeSpeechEncoderRuntime { std::shared_ptr weights_; core::ExecutionContext * execution_context_ = nullptr; size_t graph_arena_bytes_ = 0; + bool allow_flash_attention_ = true; std::unique_ptr constants_; mutable std::unique_ptr conv_graph_; mutable std::unique_ptr transformer_graph_; diff --git a/include/engine/models/higgs_audio_tts/ar.h b/include/engine/models/higgs_audio_tts/ar.h index 77549e596..d91818f81 100644 --- a/include/engine/models/higgs_audio_tts/ar.h +++ b/include/engine/models/higgs_audio_tts/ar.h @@ -1,6 +1,7 @@ #pragma once #include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/attention_fallback.h" #include "engine/framework/core/execution_context.h" #include "engine/framework/core/module.h" #include "engine/framework/modules/transformers/qwen_decoder.h" @@ -44,7 +45,8 @@ class HiggsARRuntime { std::shared_ptr assets, core::ExecutionContext & execution, size_t weight_context_bytes, - assets::TensorStorageType weight_storage_type); + assets::TensorStorageType weight_storage_type, + core::AttentionPreference attention_preference = core::AttentionPreference::Auto); const HiggsAssets & assets() const noexcept; const HiggsARWeights & weights() const noexcept; @@ -52,6 +54,7 @@ class HiggsARRuntime { core::BackendType backend_type() const noexcept; int device() const noexcept; int threads() const noexcept; + bool allow_flash_attention() const noexcept; private: std::shared_ptr assets_; @@ -59,6 +62,7 @@ class HiggsARRuntime { core::BackendType backend_type_ = core::BackendType::Cpu; int device_ = 0; int threads_ = 1; + bool allow_flash_attention_ = true; std::shared_ptr weights_; }; diff --git a/model_specs/breeze_tts.json b/model_specs/breeze_tts.json index 95a0a2cc5..d94e62cbf 100644 --- a/model_specs/breeze_tts.json +++ b/model_specs/breeze_tts.json @@ -161,6 +161,14 @@ "required": false, "min": 0, "default": 1 + }, + { + "name": "attention", + "type": "enum", + "description": "Attention lowering; auto probes the backend and falls back to eager on GPUs without a flash kernel (e.g. sm70); default auto.", + "required": false, + "values": ["auto", "flash", "eager"], + "default": "auto" } ], "load": [] diff --git a/src/framework/core/attention_fallback.cpp b/src/framework/core/attention_fallback.cpp new file mode 100644 index 000000000..e53d69c3b --- /dev/null +++ b/src/framework/core/attention_fallback.cpp @@ -0,0 +1,121 @@ +#include "engine/framework/core/attention_fallback.h" + +#include +#include +#include +#include +#include +#include + +#include "ggml.h" +#include "ggml-backend.h" + +#ifdef GGML_USE_CUDA +// CUDA driver API, declared manually so this translation unit needs neither +// the CUDA headers on its include path nor any CMake changes. The driver +// library is already linked transitively through ggml-cuda. Attribute ids +// CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR=75 / MINOR=76 are stable ABI. +extern "C" { +typedef int kCcProbeCuDevice; +typedef int kCcProbeCuResult; +kCcProbeCuResult cuDeviceGet(kCcProbeCuDevice * device, int ordinal); +kCcProbeCuResult cuDeviceGetAttribute(int * value, int attrib, kCcProbeCuDevice device); +} +#endif // GGML_USE_CUDA + +namespace engine::core { +namespace { + +std::string to_lower(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + return value; +} + +AttentionPreference parse_preference_value(const std::string & value, const char * option_name) { + const std::string lowered = to_lower(value); + if (lowered == "auto") { + return AttentionPreference::Auto; + } + if (lowered == "flash" || lowered == "on" || lowered == "1") { + return AttentionPreference::Flash; + } + if (lowered == "eager" || lowered == "off" || lowered == "0") { + return AttentionPreference::Eager; + } + throw std::runtime_error( + std::string(option_name) + " must be 'auto', 'flash', or 'eager' (got '" + value + "')"); +} + +// Auto-resolution for the CUDA flash-attention path. +// +// ggml_backend_supports_op() cannot be used here: on Volta it returns true +// (the MMA kernel is "selected") yet large prefill shapes die at launch with +// "flash_attn_ext_f16 has no device code compatible with CUDA arch 700". +// Instead, gate on compute capability, mirroring the kernel guards in +// ggml-cuda: flash below 700 (only generic TILE/VEC kernels exist) and at or +// above 800 (MMA fully instantiated); eager on 700-800, where large shapes +// select the MMA kernel with no usable device code. Unknown backends and +// query failures fail OPEN to preserve current behavior. +bool cuda_device_wants_eager(ggml_backend_t backend) { +#ifdef GGML_USE_CUDA + if (backend == nullptr) { + return false; + } + ggml_backend_dev_t device = ggml_backend_get_device(backend); + if (device == nullptr) { + return false; + } + if (ggml_backend_dev_type(device) != GGML_BACKEND_DEVICE_TYPE_GPU) { + return false; + } + const char * name = ggml_backend_dev_name(device); + if (name == nullptr || std::strncmp(name, "CUDA", 4) != 0) { + return false; // HIP / Vulkan / Metal / CPU: unchanged behavior. + } + char * end = nullptr; + const long ordinal = std::strtol(name + 4, &end, 10); + if (end == name + 4 || ordinal < 0) { + return false; + } + kCcProbeCuDevice cu_device = -1; + if (cuDeviceGet(&cu_device, static_cast(ordinal)) != 0) { + return false; + } + int major = 0; + int minor = 0; + if (cuDeviceGetAttribute(&major, 75 /* COMPUTE_CAPABILITY_MAJOR */, cu_device) != 0) { + return false; + } + if (cuDeviceGetAttribute(&minor, 76 /* COMPUTE_CAPABILITY_MINOR */, cu_device) != 0) { + return false; + } + const int cc = major * 100 + minor * 10; + return cc >= 700 && cc < 800; +#else + (void) backend; + return false; +#endif // GGML_USE_CUDA +} + +} // namespace + +AttentionPreference parse_attention_preference(const std::string & value, const char * option_name) { + return parse_preference_value(value, option_name != nullptr ? option_name : "attention"); +} + +bool resolve_flash_attention(ggml_backend_t backend, int64_t head_dim, AttentionPreference preference) { + (void) head_dim; + switch (preference) { + case AttentionPreference::Flash: + return true; + case AttentionPreference::Eager: + return false; + case AttentionPreference::Auto: + break; + } + return !cuda_device_wants_eager(backend); +} + +} // namespace engine::core diff --git a/src/framework/modules/transformers/qwen_decoder.cpp b/src/framework/modules/transformers/qwen_decoder.cpp index d7031270b..b40fc98cf 100644 --- a/src/framework/modules/transformers/qwen_decoder.cpp +++ b/src/framework/modules/transformers/qwen_decoder.cpp @@ -255,6 +255,10 @@ struct QKVProjections { core::TensorValue v; }; +bool flash_branches_allowed(const QwenDecoderLayerConfig & config) { + return config.runtime.attention.allow_flash_attention; +} + QKVProjections build_qkv_projections( core::ModuleBuildContext & ctx, const core::TensorValue & input, @@ -544,13 +548,18 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build( v = core::ensure_backend_addressable_layout(ctx, v); auto q_heads = TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); + const bool allow_flash = flash_branches_allowed(config_); const bool use_prefix_flash = + allow_flash && prefix_key.has_value() && config_.runtime.attention.prefix_mode == QwenDecoderPrefixAttentionMode::FlashWithPrefix && config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGroupedViewKV; core::TensorValue all_k = k; core::TensorValue all_v = v; - if (use_prefix_flash) { + // Cached prefix KV may be stored in a different dtype than the current + // K/V (e.g. Higgs reference state); cast before concat on every path. + // (The eager branch previously skipped this and died in ggml_concat.) + if (prefix_key.has_value()) { auto attention_prefix_key = prefix_key; auto attention_prefix_value = prefix_value; if (attention_prefix_key->type != k.type) { @@ -567,12 +576,9 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build( } all_k = ConcatModule({1}).build(ctx, *attention_prefix_key, k); all_v = ConcatModule({1}).build(ctx, *attention_prefix_value, v); - } else if (prefix_key.has_value()) { - all_k = ConcatModule({1}).build(ctx, *prefix_key, k); - all_v = ConcatModule({1}).build(ctx, *prefix_value, v); } core::TensorValue context; - if (!prefix_key.has_value() && attention_mask.has_value() && + if (allow_flash && !prefix_key.has_value() && attention_mask.has_value() && config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGroupedViewKV) { q_heads = core::wrap_tensor(ggml_cont(ctx.ggml, q_heads.tensor), q_heads.shape, q_heads.type); auto k_heads = TransposeModule({{0, 2, 1, 3}, all_k.shape.rank}).build(ctx, all_k); @@ -585,10 +591,10 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build( dim, *attention_mask, config_.attention_precision); - } else if (attention_mask.has_value() && - ((!prefix_key.has_value() && - config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGrouped) || - use_prefix_flash)) { + } else if (allow_flash && attention_mask.has_value() && + ((!prefix_key.has_value() && + config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGrouped) || + use_prefix_flash)) { q_heads = core::wrap_tensor(ggml_cont(ctx.ggml, q_heads.tensor), q_heads.shape, q_heads.type); auto k_heads = TransposeModule({{0, 2, 1, 3}, all_k.shape.rank}).build(ctx, all_k); auto v_heads = TransposeModule({{0, 2, 1, 3}, all_v.shape.rank}).build(ctx, all_v); @@ -738,6 +744,7 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail( auto k_heads = TransposeModule({{0, 2, 1, 3}, attention_key_cache.shape.rank}).build(ctx, attention_key_cache); auto v_heads = TransposeModule({{0, 2, 1, 3}, attention_value_cache.shape.rank}).build(ctx, attention_value_cache); core::TensorValue context; + const bool allow_flash = flash_branches_allowed(config_); const bool use_grouped_query = config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeatThenGroupedQuery && config_.runtime.attention.grouped_query_min_steps > 0 && @@ -755,7 +762,8 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail( config_.num_attention_heads, config_.num_key_value_heads, attention_mask); - } else if (config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeat || + } else if (!allow_flash || + config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeat || config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeatThenGroupedQuery) { k_heads = repeat_kv_heads(ctx, k_heads, kv_repeats); v_heads = repeat_kv_heads(ctx, v_heads, kv_repeats); @@ -898,6 +906,7 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail_bat auto k_heads = TransposeModule({{0, 2, 1, 3}, attention_key_cache.shape.rank}).build(ctx, attention_key_cache); auto v_heads = TransposeModule({{0, 2, 1, 3}, attention_value_cache.shape.rank}).build(ctx, attention_value_cache); core::TensorValue context; + const bool allow_flash = flash_branches_allowed(config_); const bool use_grouped_query = config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeatThenGroupedQuery && config_.runtime.attention.grouped_query_min_steps > 0 && @@ -915,7 +924,8 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail_bat config_.num_attention_heads, config_.num_key_value_heads, attention_mask); - } else if (config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeat || + } else if (!allow_flash || + config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeat || config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeatThenGroupedQuery) { k_heads = repeat_kv_heads(ctx, k_heads, kv_repeats); v_heads = repeat_kv_heads(ctx, v_heads, kv_repeats); diff --git a/src/models/breeze_tts/generator.cpp b/src/models/breeze_tts/generator.cpp index 28458f122..54b5e360b 100644 --- a/src/models/breeze_tts/generator.cpp +++ b/src/models/breeze_tts/generator.cpp @@ -82,7 +82,8 @@ modules::QwenDecoderActivationCastPolicy breeze_bf16_activation_policy(core::Bac modules::QwenCausalDecodeRuntimeConfig backbone_config( const BreezeTTSConfig & config, core::BackendType backend_type, - size_t graph_arena_bytes) { + size_t graph_arena_bytes, + bool allow_flash_attention = true) { modules::QwenCausalDecodeRuntimeConfig out; out.trace_name = "breeze_tts.backbone"; out.prefill_graph_arena_bytes = graph_arena_bytes; @@ -101,8 +102,15 @@ modules::QwenCausalDecodeRuntimeConfig backbone_config( out.decoder.stack.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; out.decoder.stack.attention_precision = GGML_PREC_F32; out.decoder.stack.projection_precision = GGML_PREC_DEFAULT; - out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; - out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + // Eager graph for GPUs without a flash kernel (e.g. sm70). + out.decoder.stack.runtime.attention.allow_flash_attention = allow_flash_attention; + if (allow_flash_attention) { + out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + } else { + out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + } out.decoder.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; out.decoder.stack.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; if (backend_type == core::BackendType::Cuda || backend_type == core::BackendType::Hip || @@ -129,7 +137,8 @@ modules::QwenCausalDecodeRuntimeConfig backbone_config( modules::QwenCausalDecodeRuntimeConfig depth_config( const BreezeTTSConfig & config, core::BackendType backend_type, - size_t graph_arena_bytes) { + size_t graph_arena_bytes, + bool allow_flash_attention = true) { modules::QwenCausalDecodeRuntimeConfig out; out.trace_name = "breeze_tts.depth_decoder"; out.prefill_graph_arena_bytes = graph_arena_bytes; @@ -148,8 +157,14 @@ modules::QwenCausalDecodeRuntimeConfig depth_config( out.decoder.stack.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; out.decoder.stack.attention_precision = GGML_PREC_F32; out.decoder.stack.projection_precision = GGML_PREC_DEFAULT; - out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; - out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.decoder.stack.runtime.attention.allow_flash_attention = allow_flash_attention; + if (allow_flash_attention) { + out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + } else { + out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + } out.decoder.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; out.decoder.stack.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; if (backend_type == core::BackendType::Cuda || backend_type == core::BackendType::Hip || @@ -720,7 +735,8 @@ struct BreezeGeneratorRuntime::Impl { core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + core::AttentionPreference attention_preference = core::AttentionPreference::Auto) : assets(std::move(assets)), execution(execution), tokenizer(this->assets), @@ -735,8 +751,14 @@ struct BreezeGeneratorRuntime::Impl { throw std::runtime_error("BreezeTTS generator requires assets"); } const auto & config = this->assets->config; - backbone_runtime_config = backbone_config(config, execution.backend_type(), graph_arena_bytes); - depth_runtime_config = depth_config(config, execution.backend_type(), graph_arena_bytes); + const bool allow_backbone_flash = core::resolve_flash_attention( + execution.backend(), config.head_dim, attention_preference); + const bool allow_depth_flash = core::resolve_flash_attention( + execution.backend(), config.depth_head_dim, attention_preference); + engine::debug::trace_log_scalar("breeze_tts.attention.allow_backbone_flash", allow_backbone_flash); + engine::debug::trace_log_scalar("breeze_tts.attention.allow_depth_flash", allow_depth_flash); + backbone_runtime_config = backbone_config(config, execution.backend_type(), graph_arena_bytes, allow_backbone_flash); + depth_runtime_config = depth_config(config, execution.backend_type(), graph_arena_bytes, allow_depth_flash); weights = load_weights(*this->assets, execution, weight_context_bytes, storage_type, backbone_runtime_config); backbone_cond = std::make_unique(execution, backbone_runtime_config, weights->backbone); backbone_uncond = std::make_unique(execution, backbone_runtime_config, weights->backbone); @@ -754,14 +776,16 @@ struct BreezeGeneratorRuntime::Impl { execution, graph_arena_bytes, storage_type, - storage_type); + storage_type, + attention_preference); speech_decoder = std::make_unique( this->assets, execution, graph_arena_bytes, weight_context_bytes, storage_type, - storage_type); + storage_type, + attention_preference); depth_first_embed_staging_.assign(static_cast(config.depth_hidden_size), 0.0F); depth_projected_pair_staging_.assign(static_cast(2 * config.depth_hidden_size), 0.0F); depth_prefill_staging_.assign(static_cast(4 * config.depth_hidden_size), 0.0F); @@ -1158,8 +1182,10 @@ BreezeGeneratorRuntime::BreezeGeneratorRuntime( engine::core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - engine::assets::TensorStorageType storage_type) - : impl_(std::make_unique(std::move(assets), execution, graph_arena_bytes, weight_context_bytes, storage_type)) {} + engine::assets::TensorStorageType storage_type, + engine::core::AttentionPreference attention_preference) + : impl_(std::make_unique( + std::move(assets), execution, graph_arena_bytes, weight_context_bytes, storage_type, attention_preference)) {} BreezeGeneratorRuntime::~BreezeGeneratorRuntime() = default; diff --git a/src/models/breeze_tts/session.cpp b/src/models/breeze_tts/session.cpp index 10f7e5e86..25338dd7b 100644 --- a/src/models/breeze_tts/session.cpp +++ b/src/models/breeze_tts/session.cpp @@ -1,5 +1,6 @@ #include "engine/models/breeze_tts/session.h" +#include "engine/framework/core/attention_fallback.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/runtime/options.h" #include "engine/framework/runtime/spec_backed_model.h" @@ -11,6 +12,7 @@ #include #include #include +#include #include namespace engine::models::breeze_tts { @@ -47,6 +49,23 @@ std::vector split_request(const runtime::TaskRequest & req return runtime::chunk_text_request(request, text_chunk_size, text_chunk_mode); } +core::AttentionPreference attention_preference_from_options(const runtime::SessionOptions & options) { + if (const auto value = runtime::find_option(options.options, {"breeze_tts.attention"})) { + return core::parse_attention_preference(*value, "breeze_tts.attention"); + } + return core::AttentionPreference::Auto; +} + +void trace_attention_preference(core::AttentionPreference preference) { + const char * name = "auto"; + if (preference == core::AttentionPreference::Flash) { + name = "flash"; + } else if (preference == core::AttentionPreference::Eager) { + name = "eager"; + } + engine::debug::trace_log_scalar("breeze_tts.attention.preference", std::string_view(name)); +} + std::size_t reference_cache_slots_from_options(const runtime::SessionOptions & options) { const int64_t slots = runtime::parse_i64_option( options.options, @@ -124,12 +143,15 @@ BreezeTTSSession::BreezeTTSSession( options.options, {"weight_context_mb"}, 2048ull * 1024ull * 1024ull); + const auto attention_preference = attention_preference_from_options(options); + trace_attention_preference(attention_preference); generator_ = std::make_unique( assets_, execution_context(), graph_arena_bytes, weight_context_bytes, - storage_type); + storage_type, + attention_preference); } BreezeTTSSession::~BreezeTTSSession() = default; diff --git a/src/models/breeze_tts/speech_decoder.cpp b/src/models/breeze_tts/speech_decoder.cpp index bd1c220af..1edf0c1c4 100644 --- a/src/models/breeze_tts/speech_decoder.cpp +++ b/src/models/breeze_tts/speech_decoder.cpp @@ -651,7 +651,8 @@ core::TensorValue attention( ggml_tensor * positions, const core::TensorValue & attention_mask, const modules::AttentionWeights & weights, - const DecoderConfig & config) { + const DecoderConfig & config, + modules::ScaledDotProductAttentionLowering lowering = modules::ScaledDotProductAttentionLowering::Flash) { const int64_t kv_repeat = config.num_heads / config.num_kv_heads; auto q_value = modules::LinearModule(binding::linear_config(config.hidden_size, config.num_heads * config.head_dim, false)) .build(build_ctx, input, {weights.q_weight, weights.q_bias}); @@ -717,7 +718,7 @@ core::TensorValue attention( } auto context = modules::ScaledDotProductAttentionModule({ config.head_dim, - modules::ScaledDotProductAttentionLowering::Flash, + lowering, GGML_PREC_F32, modules::AttentionCausality::NonCausal, }).build( @@ -804,10 +805,12 @@ class BreezeSpeechDecoderGraph { int64_t code_frames, core::ExecutionContext & execution_context, core::ConstantTensorCache & constants, - size_t graph_arena_bytes) + size_t graph_arena_bytes, + bool allow_flash_attention = true) : weights_(std::move(weights)), code_frames_(code_frames), backend_(execution_context.backend()), + allow_flash_attention_(allow_flash_attention), compute_threads_(std::max(1, execution_context.config().threads)) { if (weights_ == nullptr) { throw std::runtime_error("Breeze speech decoder graph requires weights"); @@ -863,7 +866,16 @@ class BreezeSpeechDecoderGraph { mask_, core::TensorShape::from_dims({1, 1, code_frames_, code_frames_}), GGML_TYPE_F16); - auto attn_out = attention(ctx_.get(), build_ctx, attn_in, positions_, attention_mask, layer.attention, config); + auto attn_out = attention( + ctx_.get(), + build_ctx, + attn_in, + positions_, + attention_mask, + layer.attention, + config, + allow_flash_attention_ ? modules::ScaledDotProductAttentionLowering::Flash + : modules::ScaledDotProductAttentionLowering::Explicit); attn_out = modules::LayerScaleModule{}.build( build_ctx, attn_out, @@ -1022,6 +1034,7 @@ class BreezeSpeechDecoderGraph { int64_t code_frames_ = 0; int64_t waveform_frames_ = 0; ggml_backend_t backend_ = nullptr; + bool allow_flash_attention_ = true; int compute_threads_ = 1; std::unique_ptr ctx_; ggml_tensor * codes_ = nullptr; @@ -1040,7 +1053,8 @@ BreezeSpeechDecoderRuntime::BreezeSpeechDecoderRuntime( size_t graph_arena_bytes, size_t constant_context_bytes, assets::TensorStorageType linear_weight_storage_type, - assets::TensorStorageType conv_weight_storage_type) + assets::TensorStorageType conv_weight_storage_type, + core::AttentionPreference attention_preference) : assets_(std::move(assets)), execution_context_(&execution_context), graph_arena_bytes_(graph_arena_bytes) { @@ -1053,6 +1067,9 @@ BreezeSpeechDecoderRuntime::BreezeSpeechDecoderRuntime( execution_context_->backend_type(), linear_weight_storage_type, conv_weight_storage_type); + allow_flash_attention_ = core::resolve_flash_attention( + execution_context_->backend(), weights_->config.head_dim, attention_preference); + engine::debug::trace_log_scalar("breeze_tts.attention.allow_decoder_flash", allow_flash_attention_); constants_ = std::make_unique( execution_context_->backend(), std::max(1, execution_context_->config().threads), @@ -1107,7 +1124,8 @@ runtime::AudioBuffer BreezeSpeechDecoderRuntime::decode(const BreezeSpeechCodes chunk_frames, *execution_context_, *constants_, - graph_arena_bytes_); + graph_arena_bytes_, + allow_flash_attention_); graph = std::move(replacement); } auto decoded = graph->run(chunk.data(), chunk.size()); diff --git a/src/models/breeze_tts/speech_encoder.cpp b/src/models/breeze_tts/speech_encoder.cpp index b1d59b400..d9453ce9f 100644 --- a/src/models/breeze_tts/speech_encoder.cpp +++ b/src/models/breeze_tts/speech_encoder.cpp @@ -178,7 +178,8 @@ core::TensorValue mimi_self_attention( const core::TensorValue & input, const core::TensorValue & positions, const TransformerLayerWeights & weights, - const std::optional & attention_mask) { + const std::optional & attention_mask, + modules::ScaledDotProductAttentionLowering lowering = modules::ScaledDotProductAttentionLowering::Flash) { constexpr int64_t kHeads = 8; constexpr int64_t kHeadDim = 64; auto q = modules::LinearModule(binding::linear_config(kHiddenSize, kHiddenSize, false)) @@ -206,7 +207,7 @@ core::TensorValue mimi_self_attention( auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v); auto context = modules::ScaledDotProductAttentionModule({ kHeadDim, - modules::ScaledDotProductAttentionLowering::Flash, + lowering, GGML_PREC_F32, modules::AttentionCausality::Causal, }).build(ctx, q_heads, k_heads, v_heads, attention_mask); @@ -221,12 +222,13 @@ core::TensorValue transformer_block( const core::TensorValue & input, const core::TensorValue & positions, const TransformerLayerWeights & weights, - const std::optional & attention_mask) { + const std::optional & attention_mask, + modules::ScaledDotProductAttentionLowering lowering = modules::ScaledDotProductAttentionLowering::Flash) { const modules::LayerNormModule norm({kHiddenSize, 1.0e-5F, true, true}); auto x = norm.build(ctx, input, weights.norm1); auto attn_out = modules::LayerScaleModule{}.build( ctx, - mimi_self_attention(ctx, x, positions, weights, attention_mask), + mimi_self_attention(ctx, x, positions, weights, attention_mask, lowering), weights.scale1); x = modules::AddModule{}.build(ctx, input, attn_out); auto y = norm.build(ctx, x, weights.norm2); @@ -545,8 +547,10 @@ class BreezeSpeechEncoderTransformerGraph { int64_t frames, core::ExecutionContext & execution_context, core::ConstantTensorCache & constants, - size_t graph_arena_bytes) + size_t graph_arena_bytes, + bool allow_flash_attention = true) : weights_(std::move(weights)), + allow_flash_attention_(allow_flash_attention), frames_(frames), backend_(execution_context.backend()), compute_threads_(std::max(1, execution_context.config().threads)) { @@ -590,7 +594,14 @@ class BreezeSpeechEncoderTransformerGraph { core::TensorShape::from_dims({1, 1, frames_, frames_}), GGML_TYPE_F16); for (const auto & layer : weights_->transformer_layers) { - seq = transformer_block(build_ctx, seq, positions_value, layer, attention_mask); + seq = transformer_block( + build_ctx, + seq, + positions_value, + layer, + attention_mask, + allow_flash_attention_ ? modules::ScaledDotProductAttentionLowering::Flash + : modules::ScaledDotProductAttentionLowering::Explicit); } auto x = modules::TransposeModule({{0, 2, 1, 3}, seq.shape.rank}).build(build_ctx, seq); x = core::ensure_backend_addressable_layout(build_ctx, x); @@ -676,6 +687,7 @@ class BreezeSpeechEncoderTransformerGraph { } std::shared_ptr weights_; + bool allow_flash_attention_ = true; int64_t frames_ = 0; int64_t output_frames_ = 0; std::unique_ptr ctx_; @@ -697,7 +709,8 @@ BreezeSpeechEncoderRuntime::BreezeSpeechEncoderRuntime( core::ExecutionContext & execution_context, size_t graph_arena_bytes, assets::TensorStorageType linear_weight_storage_type, - assets::TensorStorageType conv_weight_storage_type) + assets::TensorStorageType conv_weight_storage_type, + core::AttentionPreference attention_preference) : assets_(std::move(assets)), execution_context_(&execution_context), graph_arena_bytes_(graph_arena_bytes) { @@ -710,6 +723,9 @@ BreezeSpeechEncoderRuntime::BreezeSpeechEncoderRuntime( execution_context_->backend_type(), linear_weight_storage_type, conv_weight_storage_type); + // Mimi encoder self-attention head dim (kHeadDim in mimi_self_attention). + allow_flash_attention_ = core::resolve_flash_attention(execution_context_->backend(), 64, attention_preference); + engine::debug::trace_log_scalar("breeze_tts.attention.allow_encoder_flash", allow_flash_attention_); constants_ = std::make_unique( execution_context_->backend(), std::max(1, execution_context_->config().threads), @@ -754,7 +770,8 @@ BreezeSpeechCodes BreezeSpeechEncoderRuntime::encode(const runtime::AudioBuffer graph_frames, *execution_context_, *constants_, - graph_arena_bytes_); + graph_arena_bytes_, + allow_flash_attention_); } std::vector features(static_cast(kHiddenSize * graph_frames)); diff --git a/src/models/higgs_audio_tts/ar.cpp b/src/models/higgs_audio_tts/ar.cpp index e7646dda9..cd4166f0c 100644 --- a/src/models/higgs_audio_tts/ar.cpp +++ b/src/models/higgs_audio_tts/ar.cpp @@ -39,7 +39,9 @@ struct GgmlContextDeleter { } }; -modules::QwenDecoderStackConfig make_higgs_qwen_stack_config(const HiggsTextConfig & config) { +modules::QwenDecoderStackConfig make_higgs_qwen_stack_config( + const HiggsTextConfig & config, + bool allow_flash_attention = true) { modules::QwenDecoderStackConfig out; out.hidden_size = config.hidden_size; out.num_attention_heads = config.num_attention_heads; @@ -53,9 +55,17 @@ modules::QwenDecoderStackConfig make_higgs_qwen_stack_config(const HiggsTextConf out.projection_precision = GGML_PREC_DEFAULT; out.qkv_layout = modules::QwenDecoderQKVLayout::Separate; out.use_qk_norm = true; - out.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; - out.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; - out.runtime.attention.prefix_mode = modules::QwenDecoderPrefixAttentionMode::FlashWithPrefix; + // Eager graph for GPUs without a flash kernel (e.g. sm70). + out.runtime.attention.allow_flash_attention = allow_flash_attention; + if (allow_flash_attention) { + out.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.runtime.attention.prefix_mode = modules::QwenDecoderPrefixAttentionMode::FlashWithPrefix; + } else { + out.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + out.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + out.runtime.attention.prefix_mode = modules::QwenDecoderPrefixAttentionMode::Exact; + } out.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; out.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; out.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; @@ -64,8 +74,8 @@ modules::QwenDecoderStackConfig make_higgs_qwen_stack_config(const HiggsTextConf class HiggsQwenDecoderComponent { public: - HiggsQwenDecoderComponent(const HiggsTextConfig & config, bool packed_qkv) - : stack_config_(make_higgs_qwen_stack_config(config)), + HiggsQwenDecoderComponent(const HiggsTextConfig & config, bool packed_qkv, bool allow_flash_attention = true) + : stack_config_(make_higgs_qwen_stack_config(config, allow_flash_attention)), layer_config_(modules::qwen_decoder_layer_config_from_stack(stack_config_)), layer_module_([&] { layer_config_.qkv_layout = packed_qkv @@ -376,12 +386,18 @@ HiggsARRuntime::HiggsARRuntime( std::shared_ptr assets, core::ExecutionContext & execution, size_t weight_context_bytes, - assets::TensorStorageType weight_storage_type) + assets::TensorStorageType weight_storage_type, + core::AttentionPreference attention_preference) : assets_(std::move(assets)), backend_(execution.backend()), backend_type_(execution.backend_type()), device_(execution.config().device), - threads_(std::max(1, execution.config().threads)) { + threads_(std::max(1, execution.config().threads)), + allow_flash_attention_(core::resolve_flash_attention( + execution.backend(), + this->assets_->config.text.head_dim, + attention_preference)) { + engine::debug::trace_log_scalar("higgs_audio_tts.attention.allow_flash", allow_flash_attention_); if (assets_ == nullptr) { throw std::runtime_error("Higgs TTS AR runtime requires assets"); } @@ -404,6 +420,10 @@ ggml_backend_t HiggsARRuntime::backend() const noexcept { return backend_; } +bool HiggsARRuntime::allow_flash_attention() const noexcept { + return allow_flash_attention_; +} + core::BackendType HiggsARRuntime::backend_type() const noexcept { return backend_type_; } @@ -600,7 +620,8 @@ struct HiggsARDecodeGraph::Impl { GGML_TYPE_F16); graph = ggml_new_graph_custom(ctx.get(), 65536, false); - const HiggsQwenDecoderComponent decoder(config.text, tensor_weights.packed_qkv); + const HiggsQwenDecoderComponent decoder( + config.text, tensor_weights.packed_qkv, runtime->allow_flash_attention()); for (size_t layer_index = 0; layer_index < tensor_weights.decoder.layers.size(); ++layer_index) { auto out = decoder.build_decode_layer( build_ctx, @@ -845,7 +866,8 @@ struct HiggsARPrefillGraph::Impl { graph = ggml_new_graph_custom(ctx.get(), 262144, false); keys.reserve(tensor_weights.decoder.layers.size()); values.reserve(tensor_weights.decoder.layers.size()); - const HiggsQwenDecoderComponent decoder(config.text, tensor_weights.packed_qkv); + const HiggsQwenDecoderComponent decoder( + config.text, tensor_weights.packed_qkv, runtime->allow_flash_attention()); for (size_t layer_index = 0; layer_index < tensor_weights.decoder.layers.size(); ++layer_index) { std::optional prefix_key; std::optional prefix_value; @@ -1034,7 +1056,8 @@ struct HiggsARPrefillGraph::Impl { attention_mask, core::TensorShape::from_dims({1, 1, steps, steps}), GGML_TYPE_F16); - const HiggsQwenDecoderComponent decoder(config.text, runtime.weights().packed_qkv); + const HiggsQwenDecoderComponent decoder( + config.text, runtime.weights().packed_qkv, runtime.allow_flash_attention()); auto out = decoder.build_prefill_layer( build_ctx, x, diff --git a/src/models/higgs_audio_tts/loader.cpp b/src/models/higgs_audio_tts/loader.cpp index efc600c72..5d77f9ccc 100644 --- a/src/models/higgs_audio_tts/loader.cpp +++ b/src/models/higgs_audio_tts/loader.cpp @@ -56,6 +56,7 @@ runtime::ModelCliInterface cli(const HiggsAssets &) { {"higgs_audio_tts.codec_decode_graph_arena_mb", "n", "Codec decode graph arena size."}, {"higgs_audio_tts.codec_encode_graph_arena_mb", "n", "Codec encode graph arena size."}, {"higgs_audio_tts.reference_cache_slots", "n", "Encoded reference-audio cache slots; default 1."}, + {"higgs_audio_tts.attention", "auto|flash|eager", "Attention lowering; auto probes the backend and falls back to eager on GPUs without a flash kernel (e.g. sm70); default auto."}, }; return out; } diff --git a/src/models/higgs_audio_tts/session.cpp b/src/models/higgs_audio_tts/session.cpp index 0a5d47790..ca3110b38 100644 --- a/src/models/higgs_audio_tts/session.cpp +++ b/src/models/higgs_audio_tts/session.cpp @@ -1,5 +1,6 @@ #include "engine/models/higgs_audio_tts/session.h" +#include "engine/framework/core/attention_fallback.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/debug/trace.h" #include "engine/framework/runtime/options.h" @@ -10,6 +11,7 @@ #include #include #include +#include #include namespace engine::models::higgs_audio_tts { @@ -50,6 +52,23 @@ uint64_t hash_audio_samples(const runtime::AudioBuffer & audio) { return hash; } +core::AttentionPreference resolve_attention_preference(const runtime::SessionOptions & options) { + if (const auto value = runtime::find_option(options.options, {"higgs_audio_tts.attention"})) { + return core::parse_attention_preference(*value, "higgs_audio_tts.attention"); + } + return core::AttentionPreference::Auto; +} + +void trace_attention_preference(core::AttentionPreference preference) { + const char * name = "auto"; + if (preference == core::AttentionPreference::Flash) { + name = "flash"; + } else if (preference == core::AttentionPreference::Eager) { + name = "eager"; + } + debug::trace_log_scalar("higgs_audio_tts.attention.preference", std::string_view(name)); +} + std::size_t resolve_reference_cache_slots(const runtime::SessionOptions & options) { const int64_t slots = runtime::parse_i64_option( options.options, @@ -169,6 +188,7 @@ HiggsTTSSession::HiggsTTSSession( key != "higgs_audio_tts.codec_decode_graph_arena_mb" && key != "higgs_audio_tts.codec_encode_graph_arena_mb" && key != "higgs_audio_tts.reference_cache_slots" && + key != "higgs_audio_tts.attention" && key != "higgs_audio_tts.weight_type" && key != "higgs_audio_tts.ar_weight_type" && key != "higgs_audio_tts.codec_weight_type") { @@ -176,11 +196,14 @@ HiggsTTSSession::HiggsTTSSession( } } + const auto attention_preference = resolve_attention_preference(options); + trace_attention_preference(attention_preference); ar_ = std::make_shared( assets_, execution_context(), ar_weight_context_bytes_, - ar_weight_storage_type_); + ar_weight_storage_type_, + attention_preference); codec_ = std::make_shared( assets_, execution_context(), diff --git a/tests/unittests/test_attention_fallback.cpp b/tests/unittests/test_attention_fallback.cpp new file mode 100644 index 000000000..d3a35f2d0 --- /dev/null +++ b/tests/unittests/test_attention_fallback.cpp @@ -0,0 +1,68 @@ +#include "engine/framework/core/attention_fallback.h" + +#include "test_assert.h" + +#include +#include +#include + +namespace { + +using engine::core::AttentionPreference; +using engine::test::require; +using engine::test::require_eq; + +void test_parse_attention_preference() { + require_eq( + static_cast(engine::core::parse_attention_preference("auto", "attention")), + static_cast(AttentionPreference::Auto), + "parse auto"); + require_eq( + static_cast(engine::core::parse_attention_preference("flash", "attention")), + static_cast(AttentionPreference::Flash), + "parse flash"); + require_eq( + static_cast(engine::core::parse_attention_preference("eager", "attention")), + static_cast(AttentionPreference::Eager), + "parse eager"); + bool threw = false; + try { + engine::core::parse_attention_preference("sometimes", "breeze_tts.attention"); + } catch (const std::runtime_error & error) { + threw = true; + require( + std::string(error.what()).find("breeze_tts.attention") != std::string::npos, + "parse error names the option"); + } + require(threw, "parse invalid must throw"); +} + +void test_resolve_flash_attention() { + require( + engine::core::resolve_flash_attention(nullptr, 128, AttentionPreference::Flash), + "explicit flash resolves true"); + require( + !engine::core::resolve_flash_attention(nullptr, 128, AttentionPreference::Eager), + "explicit eager resolves false"); + // Null backend preserves historical behavior regardless of head_dim. + require( + engine::core::resolve_flash_attention(nullptr, 128, AttentionPreference::Auto), + "auto with null backend preserves flash"); + require( + engine::core::resolve_flash_attention(nullptr, -1, AttentionPreference::Auto), + "auto with bad head_dim preserves flash"); +} + +} // namespace + +int main() { + try { + test_parse_attention_preference(); + test_resolve_flash_attention(); + } catch (const std::exception & error) { + std::cerr << "attention_fallback_test failed: " << error.what() << '\n'; + return 1; + } + std::cout << "attention_fallback_test passed\n"; + return 0; +} From e6df0779cab4d296f35b797a0d2b4e0dcefc399c Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Sun, 6 Sep 2026 01:48:56 -0400 Subject: [PATCH 10/10] Allow Breeze attention option with older GGUF specs --- src/models/breeze_tts/session.cpp | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/src/models/breeze_tts/session.cpp b/src/models/breeze_tts/session.cpp index 25338dd7b..dd400d5a1 100644 --- a/src/models/breeze_tts/session.cpp +++ b/src/models/breeze_tts/session.cpp @@ -66,6 +66,20 @@ void trace_attention_preference(core::AttentionPreference preference) { engine::debug::trace_log_scalar("breeze_tts.attention.preference", std::string_view(name)); } +void validate_session_options( + const runtime::SessionOptions & options, + const engine::model_spec::ModelContract & contract) { + auto validation_options = options; + // Older standalone GGUF packages embed a v1 contract that predates this + // backend-compatibility option; keep them usable while still validating + // the option value in attention_preference_from_options(). + if (contract.session_option_keys.find("breeze_tts.attention") == + contract.session_option_keys.end()) { + validation_options.options.erase("breeze_tts.attention"); + } + runtime::validate_spec_backed_session_options(validation_options, contract, kFamily, kModelName); +} + std::size_t reference_cache_slots_from_options(const runtime::SessionOptions & options) { const int64_t slots = runtime::parse_i64_option( options.options, @@ -119,7 +133,7 @@ BreezeTTSSession::BreezeTTSSession( assets_(require_assets(std::move(assets))), contract_(require_contract(std::move(contract))), reference_cache_(reference_cache_slots_from_options(options)) { - runtime::validate_spec_backed_session_options(options, *contract_, kFamily, kModelName); + validate_session_options(options, *contract_); if (task_.task != runtime::VoiceTaskKind::Tts && task_.task != runtime::VoiceTaskKind::VoiceCloning && task_.task != runtime::VoiceTaskKind::VoiceDesign) {