From b13f1516694560410a1be441fc9690c9acd7db88 Mon Sep 17 00:00:00 2001 From: drbaph <84208527+Saganaki22@users.noreply.github.com> Date: Thu, 2 Jul 2026 07:05:30 +0100 Subject: [PATCH] Add native Higgs Audio v3 TTS support --- CMakeLists.txt | 7 + docs/tts.md | 17 +- include/engine/models/higgs_tts/assets.h | 73 + include/engine/models/higgs_tts/audio_codes.h | 23 + include/engine/models/higgs_tts/codec.h | 33 + include/engine/models/higgs_tts/generator.h | 51 + include/engine/models/higgs_tts/loader.h | 32 + include/engine/models/higgs_tts/session.h | 38 + include/engine/models/higgs_tts/tokenizer.h | 40 + src/framework/runtime/registry.cpp | 3 +- src/models/higgs_tts/assets.cpp | 258 ++ src/models/higgs_tts/audio_codes.cpp | 66 + src/models/higgs_tts/codec.cpp | 2112 +++++++++++++++++ src/models/higgs_tts/generator.cpp | 1179 +++++++++ src/models/higgs_tts/loader.cpp | 178 ++ src/models/higgs_tts/session.cpp | 198 ++ src/models/higgs_tts/tokenizer.cpp | 118 + tools/model_manager.py | 85 +- tools/prepare_voice_ref.py | 111 + 19 files changed, 4616 insertions(+), 6 deletions(-) create mode 100644 include/engine/models/higgs_tts/assets.h create mode 100644 include/engine/models/higgs_tts/audio_codes.h create mode 100644 include/engine/models/higgs_tts/codec.h create mode 100644 include/engine/models/higgs_tts/generator.h create mode 100644 include/engine/models/higgs_tts/loader.h create mode 100644 include/engine/models/higgs_tts/session.h create mode 100644 include/engine/models/higgs_tts/tokenizer.h create mode 100644 src/models/higgs_tts/assets.cpp create mode 100644 src/models/higgs_tts/audio_codes.cpp create mode 100644 src/models/higgs_tts/codec.cpp create mode 100644 src/models/higgs_tts/generator.cpp create mode 100644 src/models/higgs_tts/loader.cpp create mode 100644 src/models/higgs_tts/session.cpp create mode 100644 src/models/higgs_tts/tokenizer.cpp create mode 100644 tools/prepare_voice_ref.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 2e5721e..5689a6f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -205,6 +205,13 @@ add_library(engine_runtime STATIC src/models/miotts/causal_lm.cpp src/models/miotts/session.cpp src/models/miotts/loader.cpp + src/models/higgs_tts/assets.cpp + src/models/higgs_tts/audio_codes.cpp + src/models/higgs_tts/codec.cpp + src/models/higgs_tts/generator.cpp + src/models/higgs_tts/tokenizer.cpp + src/models/higgs_tts/session.cpp + src/models/higgs_tts/loader.cpp src/models/voxcpm2/assets.cpp src/models/voxcpm2/audiovae.cpp src/models/voxcpm2/generator.cpp diff --git a/docs/tts.md b/docs/tts.md index ba9fdb4..f116ff7 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -268,12 +268,27 @@ Higgs Audio v3 TTS is a voice-clone TTS model. The current integration uses the audiocpp_cli --task tts --family higgs_tts --model models/higgs-audio-v3-tts-4b --backend cuda --text "Hello from Higgs Audio." --voice-ref assets/resources/b.wav --reference-text "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you." --out out.wav ``` +Prepare WAV/MP3/M4A/FLAC references for `--voice-ref`; this normalizes them to mono 24 kHz WAV: + +```bash +python3 tools/prepare_voice_ref.py sample-5.mp3 --output sample-5_ref_24k.wav --overwrite +``` + +If the reference transcript is unknown, install the native Qwen3 ASR package and transcribe the prepared reference first: + +```bash +python3 tools/model_manager.py install qwen3_asr_0_6b --models-root models +audiocpp_cli --task asr --family qwen3_asr --model models/Qwen3-ASR-0.6B --backend cuda --audio sample-5_ref_24k.wav --text "" +``` + +Copy the printed `text_output=...` value into Higgs `--reference-text`. Higgs does not run ASR inside the TTS session; ASR is an optional companion model. + | Option | Values | Default | Meaning | |---|---|---:|---| | `--voice-ref` | WAV path | required | Reference speaker audio. | | `--reference-text` | text | empty string | Transcript for reference audio. | | `--text-chunk-size` | integer chars | `512` | Long-form chunk size. | -| `--max-tokens` | integer | `1024` | Maximum generated AR tokens per chunk. | +| `--max-tokens` | integer | `1024` | Maximum generated AR tokens per chunk. Raise this if speech stops before the text finishes. | | `--temperature` | float | `0.8` | AR sampling temperature. | | `--top-k` | integer | `30` | AR top-k sampling limit. | | `--top-p` | float | `0.8` | AR nucleus sampling limit. | diff --git a/include/engine/models/higgs_tts/assets.h b/include/engine/models/higgs_tts/assets.h new file mode 100644 index 0000000..170a1dc --- /dev/null +++ b/include/engine/models/higgs_tts/assets.h @@ -0,0 +1,73 @@ +#pragma once + +#include +#include +#include +#include + +namespace engine::assets { +class TensorSource; +} + +namespace engine::models::higgs_tts { + +struct HiggsTextConfig { + std::string model_type; + int64_t vocab_size = 0; + int64_t hidden_size = 0; + int64_t intermediate_size = 0; + int64_t num_hidden_layers = 0; + int64_t num_attention_heads = 0; + int64_t num_key_value_heads = 0; + int64_t head_dim = 0; + int64_t max_position_embeddings = 0; + int64_t bos_token_id = 0; + int64_t eos_token_id = 0; + int64_t pad_token_id = 0; + float rms_norm_eps = 1.0e-6F; + float rope_theta = 1000000.0F; + bool tie_word_embeddings = true; +}; + +struct HiggsAudioEncoderConfig { + std::string model_type; + std::string encoder_type; + int64_t num_codebooks = 8; + int64_t vocab_size = 1026; + int64_t out_dim = 0; + int64_t max_chunk_size = 50; + int64_t mel_per_sample = 8; + bool tie_word_embeddings = true; + bool use_delay_pattern = true; +}; + +struct HiggsTTSConfig { + std::string model_type; + int64_t audio_token_id = -100; + int64_t ignore_index = -100; + HiggsTextConfig text; + HiggsAudioEncoderConfig audio_encoder; +}; + +struct HiggsTTSAssetPaths { + std::filesystem::path model_root; + std::filesystem::path small_assets_root; + std::filesystem::path config_path; + std::filesystem::path codec_config_path; + std::filesystem::path chat_template_path; + std::filesystem::path model_weights_path; + std::filesystem::path model_index_path; + std::filesystem::path tokenizer_json_path; + std::filesystem::path tokenizer_config_path; +}; + +struct HiggsTTSAssets { + HiggsTTSAssetPaths paths; + HiggsTTSConfig config; + std::shared_ptr model_weights; +}; + +HiggsTTSAssetPaths resolve_higgs_tts_assets(const std::filesystem::path & model_path); +std::shared_ptr load_higgs_tts_assets(const std::filesystem::path & model_path); + +} // namespace engine::models::higgs_tts diff --git a/include/engine/models/higgs_tts/audio_codes.h b/include/engine/models/higgs_tts/audio_codes.h new file mode 100644 index 0000000..19ad571 --- /dev/null +++ b/include/engine/models/higgs_tts/audio_codes.h @@ -0,0 +1,23 @@ +#pragma once + +#include +#include + +namespace engine::models::higgs_tts { + +constexpr int32_t kHiggsAudioBocId = 1024; +constexpr int32_t kHiggsAudioEocId = 1025; +constexpr int32_t kHiggsAudioStopCode = -1; +constexpr int32_t kHiggsAudioPlaceholderId = -100; +constexpr int kHiggsAudioSampleRate = 24000; + +struct HiggsAudioCodeMatrix { + int64_t frames = 0; + int64_t codebooks = 0; + std::vector token_ids; +}; + +HiggsAudioCodeMatrix apply_delay_pattern(const HiggsAudioCodeMatrix & codes); +HiggsAudioCodeMatrix reverse_delay_pattern(const HiggsAudioCodeMatrix & delayed_codes); + +} // namespace engine::models::higgs_tts diff --git a/include/engine/models/higgs_tts/codec.h b/include/engine/models/higgs_tts/codec.h new file mode 100644 index 0000000..1689987 --- /dev/null +++ b/include/engine/models/higgs_tts/codec.h @@ -0,0 +1,33 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/runtime/session.h" +#include "engine/models/higgs_tts/assets.h" +#include "engine/models/higgs_tts/audio_codes.h" + +#include +#include + +namespace engine::models::higgs_tts { + +class HiggsAudioCodecDecoderRuntime { +public: + struct Impl; + + HiggsAudioCodecDecoderRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + ~HiggsAudioCodecDecoderRuntime(); + + HiggsAudioCodeMatrix encode_reference_audio(const runtime::AudioBuffer & audio); + runtime::AudioBuffer decode(const HiggsAudioCodeMatrix & raw_codes); + +private: + std::unique_ptr impl_; +}; + +} // namespace engine::models::higgs_tts diff --git a/include/engine/models/higgs_tts/generator.h b/include/engine/models/higgs_tts/generator.h new file mode 100644 index 0000000..00b2f45 --- /dev/null +++ b/include/engine/models/higgs_tts/generator.h @@ -0,0 +1,51 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" +#include "engine/models/higgs_tts/assets.h" +#include "engine/models/higgs_tts/audio_codes.h" +#include "engine/models/higgs_tts/tokenizer.h" + +#include +#include +#include + +namespace engine::models::higgs_tts { + +struct HiggsTTSGenerationOptions { + int64_t max_tokens = 1024; + int top_k = 30; + float top_p = 0.8F; + float temperature = 0.8F; + uint32_t seed = 0; + bool do_sample = true; +}; + +struct HiggsTTSGeneratedCodes { + HiggsAudioCodeMatrix delayed_codes; + HiggsAudioCodeMatrix raw_codes; +}; + +class HiggsTTSGeneratorRuntime { +public: + struct Impl; + + HiggsTTSGeneratorRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + ~HiggsTTSGeneratorRuntime(); + + HiggsTTSGeneratedCodes generate( + const HiggsTTSPrompt & prompt, + const HiggsTTSGenerationOptions & options, + const HiggsAudioCodeMatrix * reference_delayed_codes = nullptr); + +private: + std::unique_ptr impl_; +}; + +} // namespace engine::models::higgs_tts diff --git a/include/engine/models/higgs_tts/loader.h b/include/engine/models/higgs_tts/loader.h new file mode 100644 index 0000000..c22da69 --- /dev/null +++ b/include/engine/models/higgs_tts/loader.h @@ -0,0 +1,32 @@ +#pragma once + +#include "engine/framework/runtime/model.h" +#include "engine/models/higgs_tts/assets.h" + +#include + +namespace engine::models::higgs_tts { + +class HiggsTTSLoadedModel final : public runtime::ILoadedVoiceModel { +public: + HiggsTTSLoadedModel( + runtime::ModelMetadata metadata, + runtime::CapabilitySet capabilities, + std::shared_ptr assets); + + const runtime::ModelMetadata & metadata() const noexcept override; + const runtime::CapabilitySet & capabilities() const noexcept override; + std::unique_ptr create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const override; + +private: + runtime::ModelMetadata metadata_; + runtime::CapabilitySet capabilities_; + std::shared_ptr assets_; +}; + +std::unique_ptr load_higgs_tts_model(const std::filesystem::path & model_path); +std::shared_ptr make_higgs_tts_loader(); + +} // namespace engine::models::higgs_tts diff --git a/include/engine/models/higgs_tts/session.h b/include/engine/models/higgs_tts/session.h new file mode 100644 index 0000000..ec4c3b7 --- /dev/null +++ b/include/engine/models/higgs_tts/session.h @@ -0,0 +1,38 @@ +#pragma once + +#include "engine/framework/runtime/session.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/models/higgs_tts/assets.h" +#include "engine/models/higgs_tts/codec.h" +#include "engine/models/higgs_tts/generator.h" +#include "engine/models/higgs_tts/tokenizer.h" + +#include +#include + +namespace engine::models::higgs_tts { + +class HiggsTTSSession final + : public runtime::RuntimeSessionBase, + public runtime::IOfflineVoiceTaskSession { +public: + HiggsTTSSession( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options, + std::shared_ptr assets); + + std::string family() const override; + runtime::VoiceTaskKind task_kind() const override; + runtime::RunMode run_mode() const override; + void prepare(const runtime::SessionPreparationRequest & request) override; + runtime::TaskResult run(const runtime::TaskRequest & request) override; + +private: + runtime::TaskSpec task_; + std::shared_ptr assets_; + HiggsTTSTokenizer tokenizer_; + std::unique_ptr generator_; + std::unique_ptr codec_decoder_; +}; + +} // namespace engine::models::higgs_tts diff --git a/include/engine/models/higgs_tts/tokenizer.h b/include/engine/models/higgs_tts/tokenizer.h new file mode 100644 index 0000000..befe1bc --- /dev/null +++ b/include/engine/models/higgs_tts/tokenizer.h @@ -0,0 +1,40 @@ +#pragma once + +#include "engine/models/higgs_tts/assets.h" + +#include +#include +#include +#include + +namespace engine::models::higgs_tts { + +struct HiggsTTSPrompt { + std::string text; + std::string reference_text; + int64_t reference_audio_tokens = 0; + std::vector input_ids; +}; + +class HiggsTTSTokenizer { +public: + explicit HiggsTTSTokenizer(std::shared_ptr assets); + + HiggsTTSPrompt build_prompt( + const std::string & text, + int64_t reference_audio_tokens = 0, + const std::string & reference_text = "") const; + std::vector encode_text(const std::string & text) const; + + int32_t tts_token_id() const noexcept; + int32_t ref_audio_token_id() const noexcept; + int32_t ref_text_token_id() const noexcept; + int32_t text_token_id() const noexcept; + int32_t audio_token_id() const noexcept; + +private: + struct Impl; + std::shared_ptr impl_; +}; + +} // namespace engine::models::higgs_tts diff --git a/src/framework/runtime/registry.cpp b/src/framework/runtime/registry.cpp index 81bf535..31f8e64 100644 --- a/src/framework/runtime/registry.cpp +++ b/src/framework/runtime/registry.cpp @@ -14,6 +14,7 @@ // #include "engine/models/roformer/session.h" #include "engine/models/chatterbox/loader.h" #include "engine/models/citrinet_asr/session.h" +#include "engine/models/higgs_tts/loader.h" #include "engine/models/marblenet_vad/session.h" #include "engine/models/miocodec/loader.h" #include "engine/models/miotts/loader.h" @@ -211,8 +212,8 @@ ModelRegistry make_default_registry(const std::optional & // engine::models::roformer::make_mel_loader(), // engine::models::moss_tts::make_moss_tts_loader(), // engine::models::heartmula::make_heartmula_loader(), - // engine::models::higgs_tts::make_higgs_tts_loader(), // engine::models::parakeet_tdt::make_parakeet_tdt_loader(), + engine::models::higgs_tts::make_higgs_tts_loader(), engine::models::omnivoice::make_omnivoice_loader(), engine::models::miocodec::make_miocodec_loader(), engine::models::miotts::make_miotts_loader(), diff --git a/src/models/higgs_tts/assets.cpp b/src/models/higgs_tts/assets.cpp new file mode 100644 index 0000000..d24b7a5 --- /dev/null +++ b/src/models/higgs_tts/assets.cpp @@ -0,0 +1,258 @@ +#include "engine/models/higgs_tts/assets.h" + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/io/filesystem.h" +#include "engine/framework/io/json.h" + +#include +#include +#include +#include + +namespace engine::models::higgs_tts { +namespace json = engine::io::json; +namespace { + +std::filesystem::path canonical_existing_file(const std::filesystem::path & path) { + return std::filesystem::weakly_canonical(path); +} + +std::filesystem::path canonical_existing_directory(const std::filesystem::path & path) { + return std::filesystem::weakly_canonical(path); +} + +std::filesystem::path resolve_model_root(const std::filesystem::path & model_path) { + if (engine::io::is_existing_directory(model_path)) { + return canonical_existing_directory(model_path); + } + if (engine::io::is_existing_file(model_path)) { + return canonical_existing_directory(model_path.parent_path()); + } + throw std::runtime_error("Higgs TTS model path does not exist: " + model_path.string()); +} + +std::optional env_path(const char * name) { + const char * value = std::getenv(name); + if (value == nullptr || *value == '\0') { + return std::nullopt; + } + return std::filesystem::path(value); +} + +void append_unique_existing_dir(std::vector & dirs, const std::filesystem::path & path) { + if (path.empty() || !engine::io::is_existing_directory(path)) { + return; + } + const auto canonical = canonical_existing_directory(path); + for (const auto & dir : dirs) { + if (dir == canonical) { + return; + } + } + dirs.push_back(canonical); +} + +std::vector small_asset_roots(const std::filesystem::path & model_root) { + std::vector roots; + if (const auto env = env_path("HIGGS_TTS_SMALL_ASSETS_ROOT")) { + append_unique_existing_dir(roots, *env); + } + if (const auto env = env_path("HIGGS_TTS_ASSETS_ROOT")) { + append_unique_existing_dir(roots, *env); + } + append_unique_existing_dir(roots, model_root); + append_unique_existing_dir(roots, model_root / "higgs-audio-v3-tts-4b"); + return roots; +} + +std::filesystem::path require_first_file( + const std::vector & roots, + const std::filesystem::path & relative_path, + const char * role) { + for (const auto & root : roots) { + const auto path = root / relative_path; + if (engine::io::is_existing_file(path)) { + return canonical_existing_file(path); + } + } + std::string searched; + for (const auto & root : roots) { + searched += "\n " + (root / relative_path).string(); + } + throw std::runtime_error(std::string("Higgs TTS missing ") + role + "; searched:" + searched); +} + +std::filesystem::path optional_first_file( + const std::vector & roots, + const std::filesystem::path & relative_path) { + for (const auto & root : roots) { + const auto path = root / relative_path; + if (engine::io::is_existing_file(path)) { + return canonical_existing_file(path); + } + } + return {}; +} + +std::filesystem::path optional_codec_config( + const std::filesystem::path & model_root, + const std::vector & small_roots) { + if (const auto env = env_path("HIGGS_TTS_CODEC_CONFIG")) { + if (engine::io::is_existing_file(*env)) { + return canonical_existing_file(*env); + } + throw std::runtime_error("HIGGS_TTS_CODEC_CONFIG does not point to a file: " + env->string()); + } + std::vector roots; + append_unique_existing_dir(roots, model_root); + for (const auto & root : small_roots) { + append_unique_existing_dir(roots, root); + append_unique_existing_dir(roots, root.parent_path()); + } + return optional_first_file(roots, "higgs_audio_v2_tokenizer_config.json"); +} + +bool architecture_contains(const engine::io::json::Value & root, const std::string & needle) { + const auto * architectures = root.find("architectures"); + if (architectures == nullptr || !architectures->is_array()) { + return false; + } + for (const auto & item : architectures->as_array()) { + if (item.is_string() && item.as_string().find(needle) != std::string::npos) { + return true; + } + } + return false; +} + +float optional_rope_theta(const engine::io::json::Value & config, float default_value) { + if (const auto * rope = config.find("rope_parameters"); rope != nullptr && rope->is_object()) { + return json::optional_f32(*rope, "rope_theta", default_value); + } + if (const auto * rope = config.find("rope_scaling"); rope != nullptr && rope->is_object()) { + return json::optional_f32(*rope, "rope_theta", default_value); + } + return json::optional_f32(config, "rope_theta", default_value); +} + +HiggsTextConfig parse_text_config(const engine::io::json::Value & value) { + HiggsTextConfig config; + config.model_type = value.require("model_type").as_string(); + config.vocab_size = value.require("vocab_size").as_i64(); + config.hidden_size = value.require("hidden_size").as_i64(); + config.intermediate_size = value.require("intermediate_size").as_i64(); + config.num_hidden_layers = value.require("num_hidden_layers").as_i64(); + config.num_attention_heads = value.require("num_attention_heads").as_i64(); + config.num_key_value_heads = value.require("num_key_value_heads").as_i64(); + config.head_dim = json::optional_i64(value, "head_dim", config.hidden_size / config.num_attention_heads); + config.max_position_embeddings = value.require("max_position_embeddings").as_i64(); + config.bos_token_id = json::optional_i64(value, "bos_token_id", config.bos_token_id); + config.eos_token_id = json::optional_i64(value, "eos_token_id", config.eos_token_id); + config.pad_token_id = json::optional_i64(value, "pad_token_id", config.eos_token_id); + config.rms_norm_eps = json::optional_f32(value, "rms_norm_eps", config.rms_norm_eps); + config.rope_theta = optional_rope_theta(value, config.rope_theta); + config.tie_word_embeddings = json::optional_bool(value, "tie_word_embeddings", config.tie_word_embeddings); + + if (config.model_type != "qwen3") { + throw std::runtime_error("Higgs TTS expects text_config.model_type=qwen3"); + } + if (config.vocab_size <= 0 || config.hidden_size <= 0 || config.intermediate_size <= 0 || + config.num_hidden_layers <= 0 || config.num_attention_heads <= 0 || config.num_key_value_heads <= 0 || + config.head_dim <= 0 || config.max_position_embeddings <= 0) { + throw std::runtime_error("Higgs TTS text_config contains non-positive dimensions"); + } + if (config.hidden_size % config.num_attention_heads != 0) { + throw std::runtime_error("Higgs TTS text hidden_size must be divisible by num_attention_heads"); + } + if (config.num_attention_heads % config.num_key_value_heads != 0) { + throw std::runtime_error("Higgs TTS text attention heads must be divisible by key-value heads"); + } + if (!config.tie_word_embeddings) { + throw std::runtime_error("Higgs TTS currently expects tied text embeddings"); + } + return config; +} + +HiggsAudioEncoderConfig parse_audio_encoder_config(const engine::io::json::Value & value) { + HiggsAudioEncoderConfig config; + config.model_type = value.require("model_type").as_string(); + config.encoder_type = value.require("encoder_type").as_string(); + config.num_codebooks = value.require("num_codebooks").as_i64(); + config.vocab_size = value.require("vocab_size").as_i64(); + config.out_dim = value.require("out_dim").as_i64(); + config.max_chunk_size = json::optional_i64(value, "max_chunk_size", config.max_chunk_size); + config.mel_per_sample = json::optional_i64(value, "mel_per_sample", config.mel_per_sample); + config.tie_word_embeddings = json::optional_bool(value, "tie_word_embeddings", config.tie_word_embeddings); + config.use_delay_pattern = json::optional_bool(value, "use_delay_pattern", config.use_delay_pattern); + + if (config.model_type != "higgs_audio_encoder") { + throw std::runtime_error("Higgs TTS expects audio_encoder_config.model_type=higgs_audio_encoder"); + } + if (config.encoder_type != "discrete") { + throw std::runtime_error("Higgs TTS currently supports only the discrete audio encoder"); + } + if (config.num_codebooks <= 0 || config.vocab_size <= 0 || config.out_dim <= 0 || + config.max_chunk_size <= 0 || config.mel_per_sample <= 0) { + throw std::runtime_error("Higgs TTS audio_encoder_config contains non-positive dimensions"); + } + if (!config.tie_word_embeddings) { + throw std::runtime_error("Higgs TTS currently expects tied audio embeddings"); + } + if (!config.use_delay_pattern) { + throw std::runtime_error("Higgs TTS currently expects delay-pattern audio codes"); + } + return config; +} + +HiggsTTSConfig parse_config(const std::filesystem::path & config_path) { + const auto root = engine::io::json::parse_file(config_path); + HiggsTTSConfig config; + config.model_type = root.require("model_type").as_string(); + config.audio_token_id = json::optional_i64(root, "audio_token_id", config.audio_token_id); + config.ignore_index = json::optional_i64(root, "ignore_index", config.ignore_index); + config.text = parse_text_config(root.require("text_config")); + config.audio_encoder = parse_audio_encoder_config(root.require("audio_encoder_config")); + + if (config.model_type != "higgs_multimodal_qwen3") { + throw std::runtime_error("Higgs TTS expects model_type=higgs_multimodal_qwen3"); + } + if (!architecture_contains(root, "HiggsMultimodalQwen3")) { + throw std::runtime_error("Higgs TTS config is missing HiggsMultimodalQwen3 architecture marker"); + } + if (config.text.hidden_size != config.audio_encoder.out_dim) { + throw std::runtime_error("Higgs TTS text hidden_size must match audio encoder out_dim"); + } + return config; +} + +HiggsTTSAssetPaths resolve_paths(const std::filesystem::path & model_path) { + const auto model_root = resolve_model_root(model_path); + const auto small_roots = small_asset_roots(model_root); + HiggsTTSAssetPaths paths; + paths.model_root = model_root; + paths.model_weights_path = require_first_file({model_root}, "model.safetensors", "model weights"); + paths.config_path = require_first_file(small_roots, "config.json", "config"); + paths.tokenizer_config_path = require_first_file(small_roots, "tokenizer_config.json", "tokenizer config"); + paths.tokenizer_json_path = require_first_file(small_roots, "tokenizer.json", "tokenizer json"); + paths.model_index_path = optional_first_file(small_roots, "model.safetensors.index.json"); + paths.chat_template_path = optional_first_file(small_roots, "chat_template.jinja"); + paths.codec_config_path = optional_codec_config(model_root, small_roots); + paths.small_assets_root = paths.config_path.parent_path(); + return paths; +} + +} // namespace + +HiggsTTSAssetPaths resolve_higgs_tts_assets(const std::filesystem::path & model_path) { + return resolve_paths(model_path); +} + +std::shared_ptr load_higgs_tts_assets(const std::filesystem::path & model_path) { + HiggsTTSAssets assets; + assets.paths = resolve_paths(model_path); + assets.config = parse_config(assets.paths.config_path); + assets.model_weights = engine::assets::open_tensor_source(assets.paths.model_weights_path); + return std::make_shared(std::move(assets)); +} + +} // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/audio_codes.cpp b/src/models/higgs_tts/audio_codes.cpp new file mode 100644 index 0000000..2dd8115 --- /dev/null +++ b/src/models/higgs_tts/audio_codes.cpp @@ -0,0 +1,66 @@ +#include "engine/models/higgs_tts/audio_codes.h" + +#include +#include +#include + +namespace engine::models::higgs_tts { +namespace { + +void validate_matrix(const HiggsAudioCodeMatrix & codes, const char * role) { + if (codes.frames <= 0 || codes.codebooks <= 0) { + throw std::runtime_error(std::string("Higgs TTS ") + role + " code matrix must be non-empty"); + } + if (static_cast(codes.token_ids.size()) != codes.frames * codes.codebooks) { + throw std::runtime_error(std::string("Higgs TTS ") + role + " code matrix shape is invalid"); + } +} + +int32_t at(const HiggsAudioCodeMatrix & codes, int64_t frame, int64_t codebook) { + return codes.token_ids[static_cast(frame * codes.codebooks + codebook)]; +} + +void set(HiggsAudioCodeMatrix & codes, int64_t frame, int64_t codebook, int32_t value) { + codes.token_ids[static_cast(frame * codes.codebooks + codebook)] = value; +} + +} // namespace + +HiggsAudioCodeMatrix apply_delay_pattern(const HiggsAudioCodeMatrix & codes) { + validate_matrix(codes, "raw"); + HiggsAudioCodeMatrix delayed; + delayed.frames = codes.frames + codes.codebooks - 1; + delayed.codebooks = codes.codebooks; + delayed.token_ids.assign( + static_cast(delayed.frames * delayed.codebooks), + kHiggsAudioEocId); + for (int64_t codebook = 0; codebook < codes.codebooks; ++codebook) { + for (int64_t frame = 0; frame < codebook; ++frame) { + set(delayed, frame, codebook, kHiggsAudioBocId); + } + for (int64_t frame = 0; frame < codes.frames; ++frame) { + set(delayed, codebook + frame, codebook, at(codes, frame, codebook)); + } + } + return delayed; +} + +HiggsAudioCodeMatrix reverse_delay_pattern(const HiggsAudioCodeMatrix & delayed_codes) { + validate_matrix(delayed_codes, "delayed"); + const int64_t raw_frames = delayed_codes.frames - (delayed_codes.codebooks - 1); + if (raw_frames <= 0) { + throw std::runtime_error("Higgs TTS delayed code matrix is shorter than its codebook delay"); + } + HiggsAudioCodeMatrix raw; + raw.frames = raw_frames; + raw.codebooks = delayed_codes.codebooks; + raw.token_ids.assign(static_cast(raw.frames * raw.codebooks), 0); + for (int64_t codebook = 0; codebook < raw.codebooks; ++codebook) { + for (int64_t frame = 0; frame < raw.frames; ++frame) { + set(raw, frame, codebook, at(delayed_codes, codebook + frame, codebook)); + } + } + return raw; +} + +} // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/codec.cpp b/src/models/higgs_tts/codec.cpp new file mode 100644 index 0000000..307a2e9 --- /dev/null +++ b/src/models/higgs_tts/codec.cpp @@ -0,0 +1,2112 @@ +#include "engine/models/higgs_tts/codec.h" + +#include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/waveform_ops.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/deferred_tensor_writer.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/io/filesystem.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/types.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 +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::higgs_tts { +namespace { + +using Clock = std::chrono::steady_clock; +namespace modules = engine::modules; +namespace json = engine::io::json; + +constexpr const char * kCodecPrefix = "tied.embedding.modality_embeddings.0.model."; +constexpr int64_t kCodecHiddenSize = 1024; +constexpr int64_t kCodecCodebookDim = 64; +constexpr int64_t kCodecCodebookSize = 1024; +constexpr int64_t kCodecAcousticHiddenSize = 256; +constexpr int64_t kCodecDecoderHiddenSize = 1024; +constexpr int64_t kCodecAcousticEncoderHiddenSize = 64; +constexpr int64_t kCodecSemanticHiddenSize = 768; +constexpr int64_t kCodecSemanticIntermediateSize = 3072; +constexpr int64_t kCodecSemanticHeads = 12; +constexpr int64_t kCodecSemanticLayers = 12; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct HiggsAudioCodecConfig { + struct SemanticModelConfig { + int64_t hidden_size = kCodecSemanticHiddenSize; + int64_t intermediate_size = kCodecSemanticIntermediateSize; + int64_t num_attention_heads = kCodecSemanticHeads; + int64_t num_hidden_layers = kCodecSemanticLayers; + int64_t num_conv_pos_embeddings = 128; + int64_t num_conv_pos_embedding_groups = 16; + float layer_norm_eps = 1.0e-5F; + bool feat_proj_layer_norm = true; + bool do_stable_layer_norm = false; + std::vector conv_dim = {512, 512, 512, 512, 512, 512, 512}; + std::vector conv_kernel = {10, 3, 3, 3, 3, 2, 2}; + std::vector conv_stride = {5, 2, 2, 2, 2, 2, 2}; + }; + + int sample_rate = kHiggsAudioSampleRate; + int semantic_sample_rate = 16000; + int64_t downsample_factor = 320; + int64_t hop_length = 960; + int64_t codebook_size = kCodecCodebookSize; + int64_t codebook_dim = kCodecCodebookDim; + int64_t hidden_size = kCodecHiddenSize; + int64_t acoustic_hidden_size = kCodecAcousticHiddenSize; + int64_t acoustic_encoder_hidden_size = kCodecAcousticEncoderHiddenSize; + int64_t decoder_hidden_size = kCodecDecoderHiddenSize; + int64_t kernel_size = 3; + int64_t unit_kernel_size = 3; + std::vector downsampling_ratios = {8, 5, 4, 2, 3}; + std::vector upsampling_ratios = {8, 5, 4, 2, 3}; + std::vector strides = {1, 1}; + std::vector block_dilations = {1, 1}; + SemanticModelConfig semantic; +}; + +struct SnakeActivationWeights { + core::TensorValue alpha; +}; + +struct FlatConv1dWeights { + core::TensorValue weight; + std::optional bias; + int64_t out_channels = 0; + int64_t in_channels = 0; + int64_t kernel_size = 0; +}; + +struct FlatConvTranspose1dWeights { + core::TensorValue weight; + std::optional bias; + int64_t in_channels = 0; + int64_t out_channels = 0; + int64_t kernel_size = 0; +}; + +struct ResidualUnitWeights { + SnakeActivationWeights snake1; + FlatConv1dWeights conv1; + SnakeActivationWeights snake2; + FlatConv1dWeights conv2; +}; + +struct SemanticEncoderBlockWeights { + std::vector residual_conv1; + std::vector residual_conv2; + FlatConv1dWeights conv; +}; + +struct AcousticEncoderBlockWeights { + ResidualUnitWeights res1; + ResidualUnitWeights res2; + ResidualUnitWeights res3; + SnakeActivationWeights snake; + FlatConv1dWeights conv; +}; + +struct AcousticDecoderBlockWeights { + SnakeActivationWeights snake; + FlatConvTranspose1dWeights conv_t; + ResidualUnitWeights res1; + ResidualUnitWeights res2; + ResidualUnitWeights res3; +}; + +struct AcousticDecoderWeights { + FlatConv1dWeights conv1; + std::vector blocks; + SnakeActivationWeights snake; + FlatConv1dWeights conv2; +}; + +struct HubertFeatureExtractorWeights { + std::vector convs; + modules::NormWeights first_group_norm; +}; + +struct HubertFeatureProjectionWeights { + modules::NormWeights layer_norm; + modules::LinearWeights projection; +}; + +struct HubertPositionConvGroupWeights { + FlatConv1dWeights conv; +}; + +struct HubertEncoderLayerWeights { + modules::AttentionWeights attention; + modules::NormWeights layer_norm; + modules::FeedForwardWeights feed_forward; + modules::NormWeights final_layer_norm; +}; + +struct SemanticEncoderWeights { + FlatConv1dWeights conv; + std::vector blocks; +}; + +struct AcousticEncoderWeights { + FlatConv1dWeights conv1; + std::vector blocks; + SnakeActivationWeights snake; + FlatConv1dWeights conv2; +}; + +struct QuantizerWeights { + modules::LinearWeights project_in; + modules::LinearWeights score; + core::TensorValue codebook; + modules::LinearWeights project_out; +}; + +struct HiggsAudioCodecWeights { + std::shared_ptr store; + HubertFeatureExtractorWeights feature_extractor; + HubertFeatureProjectionWeights feature_projection; + std::vector positional_conv_groups; + modules::NormWeights encoder_input_layer_norm; + std::vector hubert_layers; + SemanticEncoderWeights semantic_encoder; + AcousticEncoderWeights acoustic_encoder; + modules::LinearWeights fc; + std::vector quantizers; + modules::LinearWeights fc2; + AcousticDecoderWeights acoustic_decoder; +}; + +HiggsAudioCodecConfig codec_config_from_assets(const HiggsTTSAssets & assets) { + HiggsAudioCodecConfig config; + if (assets.paths.codec_config_path.empty() || + !engine::io::is_existing_file(assets.paths.codec_config_path)) { + return config; + } + const auto root = json::parse_file(assets.paths.codec_config_path); + config.sample_rate = static_cast(json::optional_i64(root, "sample_rate", config.sample_rate)); + config.semantic_sample_rate = static_cast( + json::optional_i64(root, "semantic_sample_rate", config.semantic_sample_rate)); + config.downsample_factor = json::optional_i64(root, "downsample_factor", config.downsample_factor); + config.codebook_size = json::optional_i64(root, "codebook_size", config.codebook_size); + config.codebook_dim = json::optional_i64(root, "codebook_dim", config.codebook_dim); + config.kernel_size = json::optional_i64(root, "kernel_size", config.kernel_size); + config.unit_kernel_size = json::optional_i64(root, "unit_kernel_size", config.unit_kernel_size); + config.strides = json::optional_i64_array(root, "strides", config.strides); + config.block_dilations = json::optional_i64_array(root, "block_dilations", config.block_dilations); + const auto & semantic = root.require("semantic_model_config"); + config.semantic.hidden_size = json::optional_i64(semantic, "hidden_size", config.semantic.hidden_size); + config.semantic.intermediate_size = json::optional_i64(semantic, "intermediate_size", config.semantic.intermediate_size); + config.semantic.num_attention_heads = + json::optional_i64(semantic, "num_attention_heads", config.semantic.num_attention_heads); + config.semantic.num_hidden_layers = + json::optional_i64(semantic, "num_hidden_layers", config.semantic.num_hidden_layers); + config.semantic.num_conv_pos_embeddings = + json::optional_i64(semantic, "num_conv_pos_embeddings", config.semantic.num_conv_pos_embeddings); + config.semantic.num_conv_pos_embedding_groups = + json::optional_i64(semantic, "num_conv_pos_embedding_groups", config.semantic.num_conv_pos_embedding_groups); + config.semantic.layer_norm_eps = + json::optional_f32(semantic, "layer_norm_eps", config.semantic.layer_norm_eps); + config.semantic.feat_proj_layer_norm = + json::optional_bool(semantic, "feat_proj_layer_norm", config.semantic.feat_proj_layer_norm); + config.semantic.do_stable_layer_norm = + json::optional_bool(semantic, "do_stable_layer_norm", config.semantic.do_stable_layer_norm); + config.semantic.conv_dim = json::optional_i64_array(semantic, "conv_dim", config.semantic.conv_dim); + config.semantic.conv_kernel = json::optional_i64_array(semantic, "conv_kernel", config.semantic.conv_kernel); + config.semantic.conv_stride = json::optional_i64_array(semantic, "conv_stride", config.semantic.conv_stride); + const auto & acoustic = root.require("acoustic_model_config"); + config.hop_length = json::optional_i64(acoustic, "hop_length", config.hop_length); + config.acoustic_encoder_hidden_size = + json::optional_i64(acoustic, "encoder_hidden_size", config.acoustic_encoder_hidden_size); + config.acoustic_hidden_size = json::optional_i64(acoustic, "hidden_size", config.acoustic_hidden_size); + config.decoder_hidden_size = json::optional_i64(acoustic, "decoder_hidden_size", config.decoder_hidden_size); + config.downsampling_ratios = + json::optional_i64_array(acoustic, "downsampling_ratios", config.downsampling_ratios); + config.upsampling_ratios = + json::optional_i64_array(acoustic, "upsampling_ratios", config.upsampling_ratios); + config.hidden_size = config.acoustic_hidden_size + config.semantic.hidden_size; + return config; +} + +void validate_codec_config(const HiggsAudioCodecConfig & config, const HiggsTTSConfig & tts_config) { + if (config.sample_rate <= 0 || config.semantic_sample_rate <= 0 || + config.downsample_factor <= 0 || config.hop_length <= 0 || + config.codebook_size <= 0 || config.codebook_dim <= 0 || + config.hidden_size <= 0 || config.acoustic_hidden_size <= 0 || + config.acoustic_encoder_hidden_size <= 0 || config.decoder_hidden_size <= 0 || + config.downsampling_ratios.empty() || config.upsampling_ratios.empty() || + config.strides.empty() || config.block_dilations.empty() || + config.semantic.hidden_size <= 0 || config.semantic.intermediate_size <= 0 || + config.semantic.num_attention_heads <= 0 || config.semantic.num_hidden_layers <= 0 || + config.semantic.num_conv_pos_embeddings <= 0 || config.semantic.num_conv_pos_embedding_groups <= 0 || + config.semantic.conv_dim.empty() || config.semantic.conv_kernel.empty() || config.semantic.conv_stride.empty()) { + throw std::runtime_error("Higgs TTS codec config contains non-positive dimensions"); + } + if (config.codebook_size != kCodecCodebookSize) { + throw std::runtime_error("Higgs TTS codec expects 1024-entry RVQ codebooks"); + } + if (config.semantic.do_stable_layer_norm) { + throw std::runtime_error("Higgs TTS HuBERT stable-layer-norm variant is not implemented"); + } + if (config.semantic.hidden_size % config.semantic.num_attention_heads != 0 || + config.semantic.hidden_size % config.semantic.num_conv_pos_embedding_groups != 0) { + throw std::runtime_error("Higgs TTS HuBERT semantic dimensions are invalid"); + } + if (config.semantic.conv_dim.size() != config.semantic.conv_kernel.size() || + config.semantic.conv_dim.size() != config.semantic.conv_stride.size()) { + throw std::runtime_error("Higgs TTS HuBERT feature extractor config is incomplete"); + } + if (tts_config.audio_encoder.num_codebooks <= 0) { + throw std::runtime_error("Higgs TTS audio codebook count is invalid"); + } +} + +int64_t product(const std::vector & values) { + int64_t out = 1; + for (const int64_t value : values) { + if (value <= 0) { + throw std::runtime_error("Higgs TTS codec upsampling ratios must be positive"); + } + out *= value; + } + return out; +} + +std::string codec_name(const std::string & suffix) { + return std::string(kCodecPrefix) + suffix; +} + +FlatConv1dWeights load_flat_conv1d( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t out_channels, + int64_t in_channels, + int64_t kernel_size, + assets::TensorStorageType storage_type, + bool use_bias) { + FlatConv1dWeights weights; + weights.out_channels = out_channels; + weights.in_channels = in_channels; + weights.kernel_size = kernel_size; + weights.weight = store.load_tensor_as_shape( + source, + codec_name(prefix + ".weight"), + storage_type, + {out_channels, in_channels, kernel_size}, + core::TensorShape::from_dims({out_channels, in_channels * kernel_size})); + if (use_bias) { + weights.bias = store.load_f32_tensor( + source, + codec_name(prefix + ".bias"), + {out_channels}); + } + return weights; +} + +FlatConvTranspose1dWeights load_flat_conv_transpose1d( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t in_channels, + int64_t out_channels, + int64_t kernel_size, + assets::TensorStorageType storage_type, + bool use_bias) { + FlatConvTranspose1dWeights weights; + weights.in_channels = in_channels; + weights.out_channels = out_channels; + weights.kernel_size = kernel_size; + weights.weight = store.load_tensor_as_shape( + source, + codec_name(prefix + ".weight"), + storage_type, + {in_channels, out_channels, kernel_size}, + core::TensorShape::from_dims({in_channels, out_channels * kernel_size})); + if (use_bias) { + weights.bias = store.load_f32_tensor( + source, + codec_name(prefix + ".bias"), + {out_channels}); + } + return weights; +} + +SnakeActivationWeights load_snake_activation( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + int64_t channels) { + SnakeActivationWeights weights; + weights.alpha = store.make_from_f32( + core::TensorShape::from_dims({channels}), + assets::TensorStorageType::F32, + source.require_f32(codec_name(name), {1, channels, 1})); + return weights; +} + +modules::LinearWeights load_linear( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t out_features, + int64_t in_features, + assets::TensorStorageType storage_type, + bool use_bias) { + modules::LinearWeights weights; + weights.weight = store.load_tensor( + source, + codec_name(prefix + ".weight"), + storage_type, + {out_features, in_features}); + if (use_bias) { + weights.bias = store.load_f32_tensor( + source, + codec_name(prefix + ".bias"), + {out_features}); + } + return weights; +} + +modules::NormWeights load_norm( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t hidden_size) { + return { + store.load_tensor(source, codec_name(prefix + ".weight"), assets::TensorStorageType::F32, {hidden_size}), + store.load_tensor(source, codec_name(prefix + ".bias"), assets::TensorStorageType::F32, {hidden_size}), + }; +} + +modules::AttentionWeights load_attention( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t hidden_size, + assets::TensorStorageType storage_type) { + modules::AttentionWeights weights; + weights.q_weight = store.load_tensor(source, codec_name(prefix + ".q_proj.weight"), storage_type, {hidden_size, hidden_size}); + weights.q_bias = store.load_f32_tensor(source, codec_name(prefix + ".q_proj.bias"), {hidden_size}); + weights.k_weight = store.load_tensor(source, codec_name(prefix + ".k_proj.weight"), storage_type, {hidden_size, hidden_size}); + weights.k_bias = store.load_f32_tensor(source, codec_name(prefix + ".k_proj.bias"), {hidden_size}); + weights.v_weight = store.load_tensor(source, codec_name(prefix + ".v_proj.weight"), storage_type, {hidden_size, hidden_size}); + weights.v_bias = store.load_f32_tensor(source, codec_name(prefix + ".v_proj.bias"), {hidden_size}); + weights.out_weight = store.load_tensor(source, codec_name(prefix + ".out_proj.weight"), storage_type, {hidden_size, hidden_size}); + weights.out_bias = store.load_f32_tensor(source, codec_name(prefix + ".out_proj.bias"), {hidden_size}); + return weights; +} + +modules::FeedForwardWeights load_feed_forward( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t hidden_size, + int64_t intermediate_size, + assets::TensorStorageType storage_type) { + return { + store.load_tensor(source, codec_name(prefix + ".intermediate_dense.weight"), storage_type, {intermediate_size, hidden_size}), + store.load_f32_tensor(source, codec_name(prefix + ".intermediate_dense.bias"), {intermediate_size}), + store.load_tensor(source, codec_name(prefix + ".output_dense.weight"), storage_type, {hidden_size, intermediate_size}), + store.load_f32_tensor(source, codec_name(prefix + ".output_dense.bias"), {hidden_size}), + }; +} + +int64_t checked_positive(int64_t value, const char * name) { + if (value <= 0) { + throw std::runtime_error(std::string("Higgs TTS codec expected positive ") + name); + } + return value; +} + +struct NormalizedReferenceAudio { + std::vector acoustic_samples_24k; + std::vector semantic_samples_16k_padded; + int64_t frames = 0; +}; + +std::vector to_mono(const runtime::AudioBuffer & audio) { + if (audio.sample_rate <= 0) { + throw std::runtime_error("Higgs TTS reference audio sample_rate must be positive"); + } + if (audio.channels <= 0) { + throw std::runtime_error("Higgs TTS reference audio channels must be positive"); + } + if (audio.samples.empty()) { + throw std::runtime_error("Higgs TTS reference audio is empty"); + } + if (audio.samples.size() % static_cast(audio.channels) != 0) { + throw std::runtime_error("Higgs TTS interleaved reference audio has invalid channel layout"); + } + return engine::audio::mixdown_interleaved_to_mono_average( + audio.samples, + audio.channels, + engine::audio::MonoMixAccumulation::Float64); +} + +std::vector resample_sinc_hann_mono( + const std::vector & input, + int input_sample_rate, + int output_sample_rate, + int lowpass_filter_width = 6, + double rolloff = 0.99) { + if (input_sample_rate <= 0 || output_sample_rate <= 0) { + throw std::runtime_error("Higgs TTS resample sample rates must be positive"); + } + if (input.empty() || input_sample_rate == output_sample_rate) { + return input; + } + const int rate_gcd = std::gcd(input_sample_rate, output_sample_rate); + const int orig = input_sample_rate / rate_gcd; + const int next = output_sample_rate / rate_gcd; + if (orig <= 0 || next <= 0) { + throw std::runtime_error("Higgs TTS resample reduced rates must be positive"); + } + + const double base_freq = static_cast(std::min(orig, next)) * rolloff; + const int64_t width = static_cast( + std::ceil(static_cast(lowpass_filter_width * orig) / base_freq)); + const int64_t kernel_size = (2 * width) + orig; + const double scale = base_freq / static_cast(orig); + const double lowpass = static_cast(lowpass_filter_width); + const double pi = std::acos(-1.0); + + std::vector kernels(static_cast(next * kernel_size), 0.0); + for (int phase = 0; phase < next; ++phase) { + const double phase_offset = -static_cast(phase) / static_cast(next); + for (int64_t tap = 0; tap < kernel_size; ++tap) { + const double idx = static_cast(tap - width) / static_cast(orig); + double t = (phase_offset + idx) * base_freq; + t = std::clamp(t, -lowpass, lowpass); + const double window = std::pow(std::cos(t * pi / lowpass / 2.0), 2.0); + const double angle = t * pi; + const double sinc = std::abs(angle) < 1.0e-12 ? 1.0 : std::sin(angle) / angle; + kernels[static_cast(phase * kernel_size + tap)] = sinc * window * scale; + } + } + + std::vector padded(static_cast(width) + input.size() + static_cast(width + orig), 0.0F); + std::copy(input.begin(), input.end(), padded.begin() + static_cast(width)); + + const int64_t steps = + (static_cast(padded.size()) - kernel_size) / static_cast(orig) + 1; + const int64_t target_length = static_cast( + std::ceil(static_cast(next) * static_cast(input.size()) / static_cast(orig))); + std::vector output(static_cast(target_length), 0.0F); + for (int64_t index = 0; index < target_length; ++index) { + const int phase = static_cast(index % next); + const int64_t step = index / next; + if (step >= steps) { + break; + } + const int64_t start = step * static_cast(orig); + double sum = 0.0; + const double * kernel = kernels.data() + static_cast(phase * kernel_size); + for (int64_t tap = 0; tap < kernel_size; ++tap) { + sum += static_cast(padded[static_cast(start + tap)]) * kernel[tap]; + } + output[static_cast(index)] = static_cast(sum); + } + return output; +} + +int64_t conv1d_output_length(int64_t input, int64_t kernel, int stride, int padding, int dilation) { + if (input <= 0 || kernel <= 0 || stride <= 0 || dilation <= 0) { + throw std::runtime_error("Higgs TTS conv1d_output_length received an invalid shape"); + } + return ((input + 2 * padding - dilation * (kernel - 1) - 1) / stride) + 1; +} + +int64_t acoustic_encoder_output_length(int64_t input_samples, const HiggsAudioCodecConfig & config) { + int64_t length = input_samples; + length = conv1d_output_length(length, 7, 1, 3, 1); + for (const int64_t stride_value : config.downsampling_ratios) { + const int stride = static_cast(checked_positive(stride_value, "downsampling ratio")); + length = conv1d_output_length(length, 2 * stride, stride, (stride + 1) / 2, 1); + } + length = conv1d_output_length(length, 3, 1, 1, 1); + return length; +} + +int64_t semantic_downsample_factor(const HiggsAudioCodecConfig & config) { + const double factor = + static_cast(checked_positive(config.hop_length, "hop_length")) / + (static_cast(checked_positive(config.sample_rate, "sample_rate")) / + static_cast(checked_positive(config.semantic_sample_rate, "semantic_sample_rate"))) / + static_cast(checked_positive(config.downsample_factor, "downsample_factor")); + const auto rounded = static_cast(std::llround(factor)); + if (rounded <= 0 || std::fabs(factor - static_cast(rounded)) > 1.0e-6) { + throw std::runtime_error("Higgs TTS semantic_downsample_factor is not an integer"); + } + return rounded; +} + +int64_t semantic_feature_frames(const HiggsAudioCodecConfig & config, int64_t semantic_samples) { + int64_t length = semantic_samples; + for (size_t i = 0; i < config.semantic.conv_dim.size(); ++i) { + const int stride = static_cast(config.semantic.conv_stride[i]); + length = conv1d_output_length(length, config.semantic.conv_kernel[i], stride, 0, 1); + } + return length; +} + +std::vector pad_acoustic_input_if_needed( + const std::vector & input, + const HiggsAudioCodecConfig & config, + int64_t frame_count) { + const int64_t output_frames = acoustic_encoder_output_length(static_cast(input.size()), config); + if (output_frames == frame_count) { + return input; + } + std::vector padded(static_cast(input.size() + config.hop_length), 0.0F); + std::copy(input.begin(), input.end(), padded.begin() + static_cast(config.hop_length / 2)); + if (acoustic_encoder_output_length(static_cast(padded.size()), config) != frame_count) { + throw std::runtime_error("Higgs TTS acoustic/semantic reference frame alignment is unsupported"); + } + return padded; +} + +NormalizedReferenceAudio normalize_reference_audio( + const runtime::AudioBuffer & audio, + const HiggsAudioCodecConfig & config) { + auto mono = to_mono(audio); + if (audio.sample_rate != config.sample_rate) { + mono = resample_sinc_hann_mono(mono, audio.sample_rate, config.sample_rate); + } + if (static_cast(mono.size()) < config.sample_rate) { + mono.resize(static_cast(config.sample_rate), 0.0F); + } + + auto semantic = resample_sinc_hann_mono(mono, config.sample_rate, config.semantic_sample_rate); + semantic.insert(semantic.begin(), 160, 0.0F); + semantic.insert(semantic.end(), 160, 0.0F); + const int64_t feature_frames = semantic_feature_frames(config, static_cast(semantic.size())); + const int64_t downsample = semantic_downsample_factor(config); + const int64_t frames = (feature_frames + downsample - 1) / downsample; + if (frames <= 0) { + throw std::runtime_error("Higgs TTS reference audio produced no codec frames"); + } + + NormalizedReferenceAudio normalized; + normalized.frames = frames; + normalized.semantic_samples_16k_padded = std::move(semantic); + normalized.acoustic_samples_24k = pad_acoustic_input_if_needed(mono, config, frames); + return normalized; +} + +std::vector apply_positional_weight_norm( + const std::vector & g, + const std::vector & v, + int64_t out_channels, + int64_t in_channels, + int64_t kernel_size) { + if (static_cast(g.size()) != kernel_size || + static_cast(v.size()) != out_channels * in_channels * kernel_size) { + throw std::runtime_error("Higgs TTS positional weight-norm tensor shape mismatch"); + } + std::vector weight(v.size(), 0.0F); + for (int64_t kernel_index = 0; kernel_index < kernel_size; ++kernel_index) { + double squared_norm = 0.0; + for (int64_t out = 0; out < out_channels; ++out) { + for (int64_t in = 0; in < in_channels; ++in) { + const size_t offset = static_cast(((out * in_channels) + in) * kernel_size + kernel_index); + const double value = static_cast(v[offset]); + squared_norm += value * value; + } + } + const double scale = static_cast(g[static_cast(kernel_index)]) / + std::sqrt(std::max(squared_norm, 1.0e-20)); + for (int64_t out = 0; out < out_channels; ++out) { + for (int64_t in = 0; in < in_channels; ++in) { + const size_t offset = static_cast(((out * in_channels) + in) * kernel_size + kernel_index); + weight[offset] = static_cast(static_cast(v[offset]) * scale); + } + } + } + return weight; +} + +std::vector select_group_kernel( + const std::vector & values, + int64_t total_out, + int64_t total_in_per_group, + int64_t kernel_size, + int64_t groups, + int64_t group_index) { + const int64_t out_per_group = total_out / groups; + std::vector group_values(static_cast(out_per_group * total_in_per_group * kernel_size), 0.0F); + for (int64_t out = 0; out < out_per_group; ++out) { + const int64_t global_out = group_index * out_per_group + out; + const size_t src_offset = static_cast(global_out * total_in_per_group * kernel_size); + const size_t dst_offset = static_cast(out * total_in_per_group * kernel_size); + std::copy_n( + values.begin() + static_cast(src_offset), + total_in_per_group * kernel_size, + group_values.begin() + static_cast(dst_offset)); + } + return group_values; +} + +std::vector select_group_bias( + const std::vector & bias, + int64_t groups, + int64_t group_index) { + const int64_t out_per_group = static_cast(bias.size()) / groups; + std::vector group_bias(static_cast(out_per_group), 0.0F); + std::copy_n( + bias.begin() + static_cast(group_index * out_per_group), + out_per_group, + group_bias.begin()); + return group_bias; +} + +FlatConv1dWeights reshape_conv1d_weight(core::ModuleBuildContext & ctx, const FlatConv1dWeights & flat) { + FlatConv1dWeights reshaped = flat; + reshaped.weight = core::reshape_tensor( + ctx, + flat.weight, + core::TensorShape::from_dims({flat.out_channels, flat.in_channels, flat.kernel_size})); + return reshaped; +} + +FlatConvTranspose1dWeights reshape_conv_transpose1d_weight( + core::ModuleBuildContext & ctx, + const FlatConvTranspose1dWeights & flat) { + FlatConvTranspose1dWeights reshaped = flat; + reshaped.weight = core::reshape_tensor( + ctx, + flat.weight, + core::TensorShape::from_dims({flat.in_channels, flat.out_channels, flat.kernel_size})); + return reshaped; +} + +modules::Conv1dWeights make_conv1d_weights(core::ModuleBuildContext & ctx, const FlatConv1dWeights & flat) { + const auto reshaped = reshape_conv1d_weight(ctx, flat); + return {reshaped.weight, reshaped.bias}; +} + +modules::ConvTranspose1dWeights make_conv_transpose1d_weights( + core::ModuleBuildContext & ctx, + const FlatConvTranspose1dWeights & flat) { + const auto reshaped = reshape_conv_transpose1d_weight(ctx, flat); + return {reshaped.weight, reshaped.bias}; +} + +std::shared_ptr load_weights( + const HiggsTTSAssets & assets_ref, + const HiggsAudioCodecConfig & config, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) { + const auto & source = *assets_ref.model_weights; + auto weights = std::make_shared(); + weights->store = std::make_shared( + backend, + backend_type, + "higgs_tts.codec.weights", + weight_context_bytes); + const auto & semantic = config.semantic; + weights->feature_extractor.convs.reserve(semantic.conv_dim.size()); + int64_t feature_in_channels = 1; + for (size_t i = 0; i < semantic.conv_dim.size(); ++i) { + const std::string prefix = "semantic_model.feature_extractor.conv_layers." + std::to_string(i) + ".conv"; + weights->feature_extractor.convs.push_back(load_flat_conv1d( + *weights->store, + source, + prefix, + semantic.conv_dim[i], + feature_in_channels, + semantic.conv_kernel[i], + storage_type, + false)); + feature_in_channels = semantic.conv_dim[i]; + } + weights->feature_extractor.first_group_norm = { + weights->store->load_tensor( + source, + codec_name("semantic_model.feature_extractor.conv_layers.0.layer_norm.weight"), + assets::TensorStorageType::F32, + {semantic.conv_dim.front()}), + weights->store->load_tensor( + source, + codec_name("semantic_model.feature_extractor.conv_layers.0.layer_norm.bias"), + assets::TensorStorageType::F32, + {semantic.conv_dim.front()}), + }; + weights->feature_projection.layer_norm = { + weights->store->load_tensor( + source, + codec_name("semantic_model.feature_projection.layer_norm.weight"), + assets::TensorStorageType::F32, + {semantic.conv_dim.back()}), + weights->store->load_tensor( + source, + codec_name("semantic_model.feature_projection.layer_norm.bias"), + assets::TensorStorageType::F32, + {semantic.conv_dim.back()}), + }; + weights->feature_projection.projection = load_linear( + *weights->store, + source, + "semantic_model.feature_projection.projection", + semantic.hidden_size, + semantic.conv_dim.back(), + storage_type, + true); + + const auto pos_g = source.require_f32( + codec_name("semantic_model.encoder.pos_conv_embed.conv.parametrizations.weight.original0"), + {1, 1, semantic.num_conv_pos_embeddings}); + const auto pos_v = source.require_f32( + codec_name("semantic_model.encoder.pos_conv_embed.conv.parametrizations.weight.original1"), + {semantic.hidden_size, semantic.hidden_size / semantic.num_conv_pos_embedding_groups, semantic.num_conv_pos_embeddings}); + const auto pos_bias = source.require_f32( + codec_name("semantic_model.encoder.pos_conv_embed.conv.bias"), + {semantic.hidden_size}); + const auto pos_weight = apply_positional_weight_norm( + pos_g, + pos_v, + semantic.hidden_size, + semantic.hidden_size / semantic.num_conv_pos_embedding_groups, + semantic.num_conv_pos_embeddings); + const int64_t pos_groups = semantic.num_conv_pos_embedding_groups; + const int64_t pos_channels_per_group = semantic.hidden_size / pos_groups; + weights->positional_conv_groups.reserve(static_cast(pos_groups)); + for (int64_t group_index = 0; group_index < pos_groups; ++group_index) { + const auto group_weight = select_group_kernel( + pos_weight, + semantic.hidden_size, + pos_channels_per_group, + semantic.num_conv_pos_embeddings, + pos_groups, + group_index); + const auto group_bias = select_group_bias(pos_bias, pos_groups, group_index); + HubertPositionConvGroupWeights group; + group.conv.out_channels = pos_channels_per_group; + group.conv.in_channels = pos_channels_per_group; + group.conv.kernel_size = semantic.num_conv_pos_embeddings; + group.conv.weight = weights->store->make_from_f32( + core::TensorShape::from_dims({pos_channels_per_group, pos_channels_per_group * semantic.num_conv_pos_embeddings}), + storage_type, + group_weight); + group.conv.bias = weights->store->make_from_f32( + core::TensorShape::from_dims({pos_channels_per_group}), + assets::TensorStorageType::F32, + group_bias); + weights->positional_conv_groups.push_back(std::move(group)); + } + + weights->encoder_input_layer_norm = load_norm( + *weights->store, + source, + "semantic_model.encoder.layer_norm", + semantic.hidden_size); + weights->hubert_layers.reserve(static_cast(semantic.num_hidden_layers)); + for (int64_t layer = 0; layer < semantic.num_hidden_layers; ++layer) { + const std::string prefix = "semantic_model.encoder.layers." + std::to_string(layer); + HubertEncoderLayerWeights layer_weights; + layer_weights.attention = load_attention( + *weights->store, + source, + prefix + ".attention", + semantic.hidden_size, + storage_type); + layer_weights.layer_norm = load_norm(*weights->store, source, prefix + ".layer_norm", semantic.hidden_size); + layer_weights.feed_forward = load_feed_forward( + *weights->store, + source, + prefix + ".feed_forward", + semantic.hidden_size, + semantic.intermediate_size, + storage_type); + layer_weights.final_layer_norm = load_norm( + *weights->store, + source, + prefix + ".final_layer_norm", + semantic.hidden_size); + weights->hubert_layers.push_back(std::move(layer_weights)); + } + + weights->semantic_encoder.conv = load_flat_conv1d( + *weights->store, + source, + "encoder_semantic.conv", + semantic.hidden_size, + semantic.hidden_size, + config.kernel_size, + storage_type, + false); + weights->semantic_encoder.blocks.reserve(config.strides.size()); + for (size_t block_index = 0; block_index < config.strides.size(); ++block_index) { + SemanticEncoderBlockWeights block; + block.residual_conv1.reserve(config.block_dilations.size()); + block.residual_conv2.reserve(config.block_dilations.size()); + for (size_t residual_index = 0; residual_index < config.block_dilations.size(); ++residual_index) { + const std::string prefix = + "encoder_semantic.conv_blocks." + std::to_string(block_index) + + ".res_units." + std::to_string(residual_index); + block.residual_conv1.push_back(load_flat_conv1d( + *weights->store, + source, + prefix + ".conv1", + semantic.hidden_size, + semantic.hidden_size, + config.unit_kernel_size, + storage_type, + false)); + block.residual_conv2.push_back(load_flat_conv1d( + *weights->store, + source, + prefix + ".conv2", + semantic.hidden_size, + semantic.hidden_size, + 1, + storage_type, + false)); + } + block.conv = load_flat_conv1d( + *weights->store, + source, + "encoder_semantic.conv_blocks." + std::to_string(block_index) + ".conv", + semantic.hidden_size, + semantic.hidden_size, + config.strides[block_index] == 1 ? 3 : (2 * config.strides[block_index]), + storage_type, + true); + weights->semantic_encoder.blocks.push_back(std::move(block)); + } + + weights->acoustic_encoder.conv1 = load_flat_conv1d( + *weights->store, + source, + "acoustic_encoder.conv1", + config.acoustic_encoder_hidden_size, + 1, + 7, + storage_type, + true); + weights->acoustic_encoder.blocks.reserve(config.downsampling_ratios.size()); + int64_t encoder_block_input = config.acoustic_encoder_hidden_size; + for (size_t block_index = 0; block_index < config.downsampling_ratios.size(); ++block_index) { + const int64_t stride = config.downsampling_ratios[block_index]; + const int64_t output_channels = + config.acoustic_encoder_hidden_size * (int64_t{1} << (static_cast(block_index) + 1)); + const std::string block_prefix = "acoustic_encoder.block." + std::to_string(block_index); + auto load_encoder_residual_unit = [&](const std::string & prefix, int64_t channels) { + ResidualUnitWeights unit; + unit.snake1 = load_snake_activation(*weights->store, source, prefix + ".snake1.alpha", channels); + unit.conv1 = load_flat_conv1d(*weights->store, source, prefix + ".conv1", channels, channels, 7, storage_type, true); + unit.snake2 = load_snake_activation(*weights->store, source, prefix + ".snake2.alpha", channels); + unit.conv2 = load_flat_conv1d(*weights->store, source, prefix + ".conv2", channels, channels, 1, storage_type, true); + return unit; + }; + AcousticEncoderBlockWeights block; + block.res1 = load_encoder_residual_unit(block_prefix + ".res_unit1", encoder_block_input); + block.res2 = load_encoder_residual_unit(block_prefix + ".res_unit2", encoder_block_input); + block.res3 = load_encoder_residual_unit(block_prefix + ".res_unit3", encoder_block_input); + block.snake = load_snake_activation(*weights->store, source, block_prefix + ".snake1.alpha", encoder_block_input); + block.conv = load_flat_conv1d( + *weights->store, + source, + block_prefix + ".conv1", + output_channels, + encoder_block_input, + 2 * stride, + storage_type, + true); + weights->acoustic_encoder.blocks.push_back(std::move(block)); + encoder_block_input = output_channels; + } + weights->acoustic_encoder.snake = + load_snake_activation(*weights->store, source, "acoustic_encoder.snake1.alpha", encoder_block_input); + weights->acoustic_encoder.conv2 = load_flat_conv1d( + *weights->store, + source, + "acoustic_encoder.conv2", + config.acoustic_hidden_size, + encoder_block_input, + 3, + storage_type, + true); + + weights->fc = load_linear( + *weights->store, + source, + "fc", + config.hidden_size, + config.hidden_size, + storage_type, + true); + + weights->quantizers.reserve(static_cast(assets_ref.config.audio_encoder.num_codebooks)); + for (int64_t quantizer_index = 0; quantizer_index < assets_ref.config.audio_encoder.num_codebooks; ++quantizer_index) { + const std::string prefix = "quantizer.quantizers." + std::to_string(quantizer_index); + QuantizerWeights quantizer; + const auto codebook = source.require_f32(codec_name(prefix + ".codebook.embed"), {config.codebook_size, config.codebook_dim}); + std::vector score_weight(static_cast(config.codebook_size * config.codebook_dim), 0.0F); + std::vector score_bias(static_cast(config.codebook_size), 0.0F); + for (int64_t row = 0; row < config.codebook_size; ++row) { + double norm = 0.0; + for (int64_t col = 0; col < config.codebook_dim; ++col) { + const float value = codebook[static_cast(row * config.codebook_dim + col)]; + score_weight[static_cast(row * config.codebook_dim + col)] = 2.0F * value; + norm += static_cast(value) * static_cast(value); + } + score_bias[static_cast(row)] = static_cast(-norm); + } + quantizer.project_in = load_linear( + *weights->store, + source, + prefix + ".project_in", + config.codebook_dim, + config.hidden_size, + storage_type, + true); + quantizer.score = { + weights->store->make_from_f32( + core::TensorShape::from_dims({config.codebook_size, config.codebook_dim}), + storage_type, + score_weight), + weights->store->make_from_f32( + core::TensorShape::from_dims({config.codebook_size}), + assets::TensorStorageType::F32, + score_bias), + }; + quantizer.codebook = weights->store->load_tensor( + source, + codec_name(prefix + ".codebook.embed"), + storage_type, + {config.codebook_size, config.codebook_dim}); + quantizer.project_out = load_linear( + *weights->store, + source, + prefix + ".project_out", + config.hidden_size, + config.codebook_dim, + storage_type, + true); + weights->quantizers.push_back(std::move(quantizer)); + } + + weights->fc2 = load_linear( + *weights->store, + source, + "fc2", + config.acoustic_hidden_size, + config.hidden_size, + storage_type, + true); + weights->acoustic_decoder.conv1 = load_flat_conv1d( + *weights->store, + source, + "acoustic_decoder.conv1", + config.decoder_hidden_size, + config.acoustic_hidden_size, + 7, + storage_type, + true); + weights->acoustic_decoder.blocks.reserve(config.upsampling_ratios.size()); + for (size_t block_index = 0; block_index < config.upsampling_ratios.size(); ++block_index) { + const int64_t stride = config.upsampling_ratios[block_index]; + const int64_t input_channels = config.decoder_hidden_size / (int64_t{1} << static_cast(block_index)); + const int64_t output_channels = config.decoder_hidden_size / (int64_t{1} << (static_cast(block_index) + 1)); + const std::string block_prefix = "acoustic_decoder.block." + std::to_string(block_index); + auto load_residual_unit = [&](const std::string & prefix, int64_t channels) { + ResidualUnitWeights unit; + unit.snake1 = load_snake_activation(*weights->store, source, prefix + ".snake1.alpha", channels); + unit.conv1 = load_flat_conv1d(*weights->store, source, prefix + ".conv1", channels, channels, 7, storage_type, true); + unit.snake2 = load_snake_activation(*weights->store, source, prefix + ".snake2.alpha", channels); + unit.conv2 = load_flat_conv1d(*weights->store, source, prefix + ".conv2", channels, channels, 1, storage_type, true); + return unit; + }; + AcousticDecoderBlockWeights block; + block.snake = load_snake_activation(*weights->store, source, block_prefix + ".snake1.alpha", input_channels); + block.conv_t = load_flat_conv_transpose1d( + *weights->store, + source, + block_prefix + ".conv_t1", + input_channels, + output_channels, + 2 * stride, + storage_type, + true); + block.res1 = load_residual_unit(block_prefix + ".res_unit1", output_channels); + block.res2 = load_residual_unit(block_prefix + ".res_unit2", output_channels); + block.res3 = load_residual_unit(block_prefix + ".res_unit3", output_channels); + weights->acoustic_decoder.blocks.push_back(std::move(block)); + } + const int64_t final_channels = config.decoder_hidden_size / + (int64_t{1} << static_cast(config.upsampling_ratios.size())); + weights->acoustic_decoder.snake = load_snake_activation( + *weights->store, + source, + "acoustic_decoder.snake1.alpha", + final_channels); + weights->acoustic_decoder.conv2 = load_flat_conv1d( + *weights->store, + source, + "acoustic_decoder.conv2", + 1, + final_channels, + 7, + storage_type, + true); + weights->store->upload(); + return weights; +} + +core::TensorValue transpose_btc_to_bct(core::ModuleBuildContext & ctx, const core::TensorValue & value) { + return modules::TransposeModule({{0, 2, 1}, value.shape.rank}).build(ctx, value); +} + +core::TensorValue transpose_bct_to_btc(core::ModuleBuildContext & ctx, const core::TensorValue & value) { + return modules::TransposeModule({{0, 2, 1}, value.shape.rank}).build(ctx, value); +} + +core::TensorValue build_snake1d( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const SnakeActivationWeights & weights, + int64_t channels) { + return modules::Snake1dModule({channels}).build(ctx, input, {weights.alpha}); +} + +core::TensorValue ensure_dense_layout(core::ModuleBuildContext & ctx, const core::TensorValue & value) { + if (core::has_backend_addressable_layout(value.tensor) && + !ggml_is_transposed(value.tensor) && + value.tensor->nb[0] == ggml_type_size(value.tensor->type)) { + return value; + } + return core::wrap_tensor(ggml_cont(ctx.ggml, value.tensor), value.shape, value.type); +} + +core::TensorValue build_conv_transpose_with_crop( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const FlatConvTranspose1dWeights & weights, + int stride, + int padding, + int output_padding) { + const auto dense_input = ensure_dense_layout(ctx, input); + auto output = modules::ConvTranspose1dModule({ + weights.in_channels, + weights.out_channels, + weights.kernel_size, + stride, + 0, + 1, + weights.bias.has_value(), + }).build(ctx, dense_input, make_conv_transpose1d_weights(ctx, weights)); + const int64_t cropped_length = + (input.shape.dims[2] - 1) * stride - 2 * padding + weights.kernel_size + output_padding; + if (cropped_length <= 0 || cropped_length > output.shape.dims[2]) { + throw std::runtime_error("Higgs TTS codec ConvTranspose crop length is invalid"); + } + auto * view = ggml_view_3d( + ctx.ggml, + output.tensor, + cropped_length, + weights.out_channels, + 1, + output.tensor->nb[1], + output.tensor->nb[2], + static_cast(padding) * output.tensor->nb[0]); + return core::wrap_tensor( + ggml_cont(ctx.ggml, view), + core::TensorShape::from_dims({1, weights.out_channels, cropped_length}), + GGML_TYPE_F32); +} + +core::TensorValue build_residual_unit( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const ResidualUnitWeights & weights, + int64_t channels, + int dilation) { + auto x = build_snake1d(ctx, input, weights.snake1, channels); + x = modules::Conv1dModule({channels, channels, 7, 1, 3 * dilation, dilation, true}) + .build(ctx, ensure_dense_layout(ctx, x), make_conv1d_weights(ctx, weights.conv1)); + x = build_snake1d(ctx, x, weights.snake2, channels); + x = modules::Conv1dModule({channels, channels, 1, 1, 0, 1, true}) + .build(ctx, ensure_dense_layout(ctx, x), make_conv1d_weights(ctx, weights.conv2)); + return modules::AddModule().build(ctx, input, x); +} + +core::TensorValue ensure_f32_tensor(core::ModuleBuildContext & ctx, const core::TensorValue & value) { + if (value.type == GGML_TYPE_F32) { + return value; + } + return core::wrap_tensor(ggml_cast(ctx.ggml, value.tensor, GGML_TYPE_F32), value.shape, GGML_TYPE_F32); +} + +core::TensorValue group_norm_affine( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t groups, + float eps, + const modules::NormWeights & weights, + core::DeferredTensorWriter & writer) { + core::TensorValue output; + if (input.shape.rank == 3 && groups == input.shape.dims[1]) { + auto input_f32 = ensure_f32_tensor(ctx, input); + auto mean = modules::ReduceMeanModule({2}).build(ctx, input_f32); + auto mean_rep = modules::RepeatModule(modules::RepeatConfig{input_f32.shape}).build(ctx, mean); + auto centered = core::wrap_tensor(ggml_sub(ctx.ggml, input_f32.tensor, mean_rep.tensor), input_f32.shape, GGML_TYPE_F32); + auto centered_sq = modules::MulModule{}.build(ctx, centered, centered); + auto variance = modules::ReduceMeanModule({2}).build(ctx, centered_sq); + auto eps_tensor = writer.make_f32_tensor(ctx, core::TensorShape::from_dims({1, 1, 1}), {eps}); + auto eps_rep = modules::RepeatModule(modules::RepeatConfig{variance.shape}).build(ctx, eps_tensor); + auto std = modules::SqrtModule{}.build(ctx, modules::AddModule{}.build(ctx, variance, eps_rep)); + auto std_rep = modules::RepeatModule(modules::RepeatConfig{input_f32.shape}).build(ctx, std); + output = core::wrap_tensor(ggml_div(ctx.ggml, centered.tensor, std_rep.tensor), input_f32.shape, GGML_TYPE_F32); + } else { + output = core::wrap_tensor(ggml_group_norm(ctx.ggml, input.tensor, groups, eps), input.shape, GGML_TYPE_F32); + } + if (weights.weight.has_value()) { + auto weight = core::reshape_tensor(ctx, *weights.weight, core::TensorShape::from_dims({1, input.shape.dims[1], 1})); + auto repeated = core::wrap_tensor(ggml_repeat(ctx.ggml, weight.tensor, output.tensor), output.shape, GGML_TYPE_F32); + output = core::wrap_tensor(ggml_mul(ctx.ggml, output.tensor, repeated.tensor), output.shape, GGML_TYPE_F32); + } + if (weights.bias.has_value()) { + auto bias = core::reshape_tensor(ctx, *weights.bias, core::TensorShape::from_dims({1, input.shape.dims[1], 1})); + auto repeated = core::wrap_tensor(ggml_repeat(ctx.ggml, bias.tensor, output.tensor), output.shape, GGML_TYPE_F32); + output = core::wrap_tensor(ggml_add(ctx.ggml, output.tensor, repeated.tensor), output.shape, GGML_TYPE_F32); + } + return output; +} + +core::TensorValue build_semantic_residual_unit( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const FlatConv1dWeights & conv1, + const FlatConv1dWeights & conv2, + int64_t channels, + int dilation) { + auto x = modules::EluModule().build(ctx, input); + const int padding = static_cast(((conv1.kernel_size - 1) / 2) * static_cast(dilation)); + x = modules::Conv1dModule({channels, channels, conv1.kernel_size, 1, padding, dilation, false}) + .build(ctx, ensure_dense_layout(ctx, x), make_conv1d_weights(ctx, conv1)); + x = modules::EluModule().build(ctx, x); + x = modules::Conv1dModule({channels, channels, 1, 1, 0, 1, false}) + .build(ctx, ensure_dense_layout(ctx, x), make_conv1d_weights(ctx, conv2)); + return modules::AddModule().build(ctx, input, x); +} + +core::TensorValue build_conv1d_im2col_f32( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const FlatConv1dWeights & weights, + int stride, + int padding, + int dilation) { + core::validate_rank_between(input, 3, 3, "input"); + const int64_t output_frames = + (input.shape.dims[2] + 2 * padding - dilation * (weights.kernel_size - 1) - 1) / stride + 1; + if (output_frames <= 0) { + throw std::runtime_error("Higgs TTS exact conv1d computed non-positive output length"); + } + auto input_f32 = ensure_f32_tensor(ctx, ensure_dense_layout(ctx, input)); + auto reshaped = reshape_conv1d_weight(ctx, weights); + auto weight_f32 = ensure_f32_tensor(ctx, ensure_dense_layout(ctx, reshaped.weight)); + ggml_tensor * im2col = ggml_im2col( + ctx.ggml, + weight_f32.tensor, + input_f32.tensor, + stride, + 0, + padding, + 0, + dilation, + 0, + false, + GGML_TYPE_F32); + auto im2col_btk = core::wrap_tensor( + im2col, + core::TensorShape::from_dims({input.shape.dims[0], output_frames, weights.in_channels * weights.kernel_size}), + GGML_TYPE_F32); + im2col_btk = ensure_dense_layout(ctx, im2col_btk); + modules::LinearWeights linear_weights; + linear_weights.weight = core::reshape_tensor( + ctx, + reshaped.weight, + core::TensorShape::from_dims({weights.out_channels, weights.in_channels * weights.kernel_size})); + if (weights.bias.has_value()) { + linear_weights.bias = *weights.bias; + } + auto output_bto = modules::LinearModule({ + weights.in_channels * weights.kernel_size, + weights.out_channels, + weights.bias.has_value(), + }).build(ctx, im2col_btk, linear_weights); + return transpose_btc_to_bct(ctx, output_bto); +} + +core::TensorValue build_grouped_positional_conv( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const std::vector & groups, + int64_t total_channels, + int64_t kernel_size) { + if (groups.empty()) { + throw std::runtime_error("Higgs TTS positional conv group weights are empty"); + } + const int64_t channels_per_group = total_channels / static_cast(groups.size()); + std::optional output; + for (size_t group_index = 0; group_index < groups.size(); ++group_index) { + auto input_slice = modules::SliceModule({ + 1, + static_cast(group_index) * channels_per_group, + channels_per_group, + }).build(ctx, input); + auto conv = modules::Conv1dModule({ + channels_per_group, + channels_per_group, + kernel_size, + 1, + static_cast(kernel_size / 2), + 1, + true, + }).build(ctx, ensure_dense_layout(ctx, input_slice), make_conv1d_weights(ctx, groups[group_index].conv)); + output = output.has_value() ? modules::ConcatModule({1}).build(ctx, *output, conv) : conv; + } + auto combined = *output; + if (kernel_size % 2 == 0) { + combined = modules::SliceModule({2, 0, input.shape.dims[2]}).build(ctx, combined); + } + return modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, combined); +} + +core::TensorValue build_self_attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::AttentionWeights & weights, + int64_t hidden_size, + int64_t num_heads) { + if (hidden_size <= 0 || num_heads <= 0 || (hidden_size % num_heads) != 0) { + throw std::runtime_error("Higgs TTS HuBERT self-attention received an invalid hidden shape"); + } + const int64_t head_dim = hidden_size / num_heads; + const modules::LinearModule q_proj({hidden_size, hidden_size, true, GGML_PREC_F32}); + const modules::LinearModule k_proj({hidden_size, hidden_size, true, GGML_PREC_F32}); + const modules::LinearModule v_proj({hidden_size, hidden_size, true, GGML_PREC_F32}); + const modules::LinearModule out_proj({hidden_size, hidden_size, true, GGML_PREC_F32}); + auto q = q_proj.build(ctx, input, {weights.q_weight, weights.q_bias}); + auto k = k_proj.build(ctx, input, {weights.k_weight, weights.k_bias}); + auto v = v_proj.build(ctx, input, {weights.v_weight, weights.v_bias}); + q = core::reshape_tensor(ctx, ensure_dense_layout(ctx, q), core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], num_heads, head_dim})); + k = core::reshape_tensor(ctx, ensure_dense_layout(ctx, k), core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], num_heads, head_dim})); + v = core::reshape_tensor(ctx, ensure_dense_layout(ctx, v), core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], num_heads, head_dim})); + 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); + const modules::MatMulModule matmul; + auto scores = matmul.build(ctx, q_heads, modules::TransposeModule({{0, 1, 3, 2}, k_heads.shape.rank}).build(ctx, k_heads)); + scores = core::wrap_tensor( + ggml_scale(ctx.ggml, scores.tensor, 1.0F / std::sqrt(static_cast(head_dim))), + scores.shape, + GGML_TYPE_F32); + scores = ensure_dense_layout(ctx, scores); + auto attn = core::wrap_tensor(ggml_soft_max(ctx.ggml, scores.tensor), scores.shape, GGML_TYPE_F32); + auto context = matmul.build(ctx, attn, v_heads); + context = core::reshape_tensor( + ctx, + ensure_dense_layout(ctx, modules::TransposeModule({{0, 2, 1, 3}, context.shape.rank}).build(ctx, context)), + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], hidden_size})); + return out_proj.build(ctx, context, {weights.out_weight, weights.out_bias}); +} + +core::TensorValue build_hubert_sequence_mean( + core::ModuleBuildContext & ctx, + const core::TensorValue & hidden, + const HiggsAudioCodecWeights & weights, + const HiggsAudioCodecConfig & config) { + auto x = hidden; + auto summed = x; + const modules::LayerNormModule layer_norm({config.semantic.hidden_size, config.semantic.layer_norm_eps, true, true}); + const modules::FeedForwardModule feed_forward({ + config.semantic.hidden_size, + config.semantic.intermediate_size, + true, + modules::GeluApproximation::ExactErf, + }); + for (const auto & layer_weights : weights.hubert_layers) { + auto attn = build_self_attention( + ctx, + x, + layer_weights.attention, + config.semantic.hidden_size, + config.semantic.num_attention_heads); + x = modules::AddModule().build(ctx, x, attn); + x = layer_norm.build(ctx, x, layer_weights.layer_norm); + auto ff = feed_forward.build(ctx, x, layer_weights.feed_forward); + x = modules::AddModule().build(ctx, x, ff); + x = layer_norm.build(ctx, x, layer_weights.final_layer_norm); + summed = modules::AddModule().build(ctx, summed, x); + } + return core::wrap_tensor( + ggml_scale(ctx.ggml, summed.tensor, 1.0F / static_cast(weights.hubert_layers.size() + 1)), + summed.shape, + GGML_TYPE_F32); +} + +core::TensorValue build_semantic_encoder( + core::ModuleBuildContext & ctx, + const core::TensorValue & input_bct, + const HiggsAudioCodecWeights & weights, + const HiggsAudioCodecConfig & config) { + auto x = modules::Conv1dModule({ + config.semantic.hidden_size, + config.semantic.hidden_size, + config.kernel_size, + 1, + static_cast(config.kernel_size / 2), + 1, + false, + }).build(ctx, ensure_dense_layout(ctx, input_bct), make_conv1d_weights(ctx, weights.semantic_encoder.conv)); + for (size_t block_index = 0; block_index < weights.semantic_encoder.blocks.size(); ++block_index) { + const auto & block = weights.semantic_encoder.blocks[block_index]; + for (size_t residual_index = 0; residual_index < block.residual_conv1.size(); ++residual_index) { + x = build_semantic_residual_unit( + ctx, + x, + block.residual_conv1[residual_index], + block.residual_conv2[residual_index], + config.semantic.hidden_size, + static_cast(config.block_dilations[residual_index])); + } + const int stride = static_cast(config.strides[block_index]); + const int kernel = stride == 1 ? 3 : static_cast(2 * config.strides[block_index]); + const int padding = (kernel - 1) / 2; + x = modules::Conv1dModule({ + config.semantic.hidden_size, + config.semantic.hidden_size, + kernel, + stride, + padding, + 1, + true, + }).build(ctx, ensure_dense_layout(ctx, x), make_conv1d_weights(ctx, block.conv)); + } + return x; +} + +core::TensorValue build_acoustic_encoder( + core::ModuleBuildContext & ctx, + const core::TensorValue & input_bct, + const HiggsAudioCodecWeights & weights, + const HiggsAudioCodecConfig & config) { + auto x = modules::Conv1dModule({1, config.acoustic_encoder_hidden_size, 7, 1, 3, 1, true}) + .build(ctx, ensure_dense_layout(ctx, input_bct), make_conv1d_weights(ctx, weights.acoustic_encoder.conv1)); + int64_t channels = config.acoustic_encoder_hidden_size; + for (size_t block_index = 0; block_index < weights.acoustic_encoder.blocks.size(); ++block_index) { + const auto & block = weights.acoustic_encoder.blocks[block_index]; + x = build_residual_unit(ctx, x, block.res1, channels, 1); + x = build_residual_unit(ctx, x, block.res2, channels, 3); + x = build_residual_unit(ctx, x, block.res3, channels, 9); + x = build_snake1d(ctx, x, block.snake, channels); + const int stride = static_cast(config.downsampling_ratios[block_index]); + const int64_t next_channels = + config.acoustic_encoder_hidden_size * (int64_t{1} << (static_cast(block_index) + 1)); + x = modules::Conv1dModule({ + channels, + next_channels, + 2 * stride, + stride, + (stride + 1) / 2, + 1, + true, + }).build(ctx, ensure_dense_layout(ctx, x), make_conv1d_weights(ctx, block.conv)); + channels = next_channels; + } + x = build_snake1d(ctx, x, weights.acoustic_encoder.snake, channels); + return modules::Conv1dModule({channels, config.acoustic_hidden_size, 3, 1, 1, 1, true}) + .build(ctx, ensure_dense_layout(ctx, x), make_conv1d_weights(ctx, weights.acoustic_encoder.conv2)); +} + +core::TensorValue build_quantizer_decode_sequence( + core::ModuleBuildContext & ctx, + const std::vector & code_inputs, + const HiggsAudioCodecWeights & weights, + int64_t frames) { + std::optional latent; + for (size_t quantizer_index = 0; quantizer_index < code_inputs.size(); ++quantizer_index) { + const auto codes = core::wrap_tensor( + code_inputs[quantizer_index], + core::TensorShape::from_dims({1, frames}), + GGML_TYPE_I32); + const auto & quantizer = weights.quantizers[quantizer_index]; + auto embedded = modules::CodebookLookupModule({ + quantizer.codebook.shape.dims[0], + quantizer.codebook.shape.dims[1], + }).build(ctx, codes, quantizer.codebook); + auto projected = modules::LinearModule({ + embedded.shape.dims[2], + quantizer.project_out.weight.shape.dims[0], + true, + }).build(ctx, embedded, quantizer.project_out); + auto projected_bct = transpose_btc_to_bct(ctx, projected); + latent = latent.has_value() ? modules::AddModule().build(ctx, *latent, projected_bct) : projected_bct; + } + if (!latent.has_value()) { + throw std::runtime_error("Higgs TTS codec quantizer decode requires at least one codebook"); + } + return *latent; +} + +core::TensorValue build_acoustic_decoder( + core::ModuleBuildContext & ctx, + const core::TensorValue & input_bct, + const HiggsAudioCodecWeights & weights, + const HiggsAudioCodecConfig & config) { + auto x = modules::Conv1dModule({ + config.acoustic_hidden_size, + config.decoder_hidden_size, + 7, + 1, + 3, + 1, + true, + }).build(ctx, ensure_dense_layout(ctx, input_bct), make_conv1d_weights(ctx, weights.acoustic_decoder.conv1)); + int64_t channels = config.decoder_hidden_size; + for (size_t block_index = 0; block_index < weights.acoustic_decoder.blocks.size(); ++block_index) { + const auto & block = weights.acoustic_decoder.blocks[block_index]; + const int stride = static_cast(config.upsampling_ratios[block_index]); + const int padding = (stride + 1) / 2; + const int output_padding = stride % 2; + x = build_snake1d(ctx, x, block.snake, channels); + x = build_conv_transpose_with_crop(ctx, x, block.conv_t, stride, padding, output_padding); + channels = block.conv_t.out_channels; + x = build_residual_unit(ctx, x, block.res1, channels, 1); + x = build_residual_unit(ctx, x, block.res2, channels, 3); + x = build_residual_unit(ctx, x, block.res3, channels, 9); + } + x = build_snake1d(ctx, x, weights.acoustic_decoder.snake, channels); + return modules::Conv1dModule({channels, 1, 7, 1, 3, 1, true}) + .build(ctx, ensure_dense_layout(ctx, x), make_conv1d_weights(ctx, weights.acoustic_decoder.conv2)); +} + +HiggsAudioCodeMatrix trim_decodable_codes(const HiggsAudioCodeMatrix & raw_codes, int64_t codebook_size) { + if (raw_codes.frames <= 0 || raw_codes.codebooks <= 0 || + static_cast(raw_codes.token_ids.size()) != raw_codes.frames * raw_codes.codebooks) { + throw std::runtime_error("Higgs TTS codec received an invalid raw code matrix"); + } + int64_t frames = 0; + for (; frames < raw_codes.frames; ++frames) { + bool valid = true; + for (int64_t codebook = 0; codebook < raw_codes.codebooks; ++codebook) { + const int32_t code = raw_codes.token_ids[static_cast(frames * raw_codes.codebooks + codebook)]; + if (code < 0 || static_cast(code) >= codebook_size) { + valid = false; + break; + } + } + if (!valid) { + break; + } + } + if (frames <= 0) { + throw std::runtime_error("Higgs TTS codec did not receive any decodable audio-code frames"); + } + HiggsAudioCodeMatrix out; + out.frames = frames; + out.codebooks = raw_codes.codebooks; + out.token_ids.assign( + raw_codes.token_ids.begin(), + raw_codes.token_ids.begin() + static_cast(frames * raw_codes.codebooks)); + return out; +} + +class EncoderGraph { +public: + EncoderGraph( + std::shared_ptr weights, + HiggsAudioCodecConfig config, + ggml_backend_t backend, + core::BackendType backend_type, + int threads, + size_t graph_arena_bytes, + int64_t acoustic_samples, + int64_t semantic_samples, + int64_t frames) + : weights_(std::move(weights)), + config_(std::move(config)), + backend_(backend), + backend_type_(backend_type), + threads_(threads), + acoustic_sample_capacity_(acoustic_samples), + semantic_sample_capacity_(semantic_samples), + frame_capacity_(frames) { + if (backend_ == nullptr) { + throw std::runtime_error("Higgs TTS codec encoder backend is not initialized"); + } + if (acoustic_sample_capacity_ <= 0 || semantic_sample_capacity_ <= 0 || frame_capacity_ <= 0) { + throw std::runtime_error("Higgs TTS codec encoder graph requires positive capacities"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS codec encoder graph context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "higgs_tts.codec.encode", backend_type_}; + semantic_downsample_factor_ = semantic_downsample_factor(config_); + acoustic_input_ = core::make_tensor( + ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, 1, acoustic_sample_capacity_})); + semantic_input_ = core::make_tensor( + ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, 1, semantic_sample_capacity_})); + downsample_indices_ = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({frame_capacity_})); + ggml_set_input(acoustic_input_.tensor); + ggml_set_input(semantic_input_.tensor); + ggml_set_input(downsample_indices_.tensor); + + auto hubert = semantic_input_; + for (size_t i = 0; i < weights_->feature_extractor.convs.size(); ++i) { + const auto & conv_weights = weights_->feature_extractor.convs[i]; + hubert = build_conv1d_im2col_f32( + ctx, + hubert, + conv_weights, + static_cast(config_.semantic.conv_stride[i]), + 0, + 1); + hubert = ensure_dense_layout(ctx, hubert); + if (i == 0) { + hubert = group_norm_affine( + ctx, + hubert, + config_.semantic.conv_dim.front(), + config_.semantic.layer_norm_eps, + weights_->feature_extractor.first_group_norm, + tensor_writer_); + } + hubert = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, hubert); + } + hubert = transpose_bct_to_btc(ctx, hubert); + if (config_.semantic.feat_proj_layer_norm) { + hubert = modules::LayerNormModule({ + config_.semantic.conv_dim.back(), + config_.semantic.layer_norm_eps, + true, + true, + }).build(ctx, hubert, weights_->feature_projection.layer_norm); + } + hubert = modules::LinearModule({ + config_.semantic.conv_dim.back(), + config_.semantic.hidden_size, + true, + GGML_PREC_F32, + }).build(ctx, hubert, weights_->feature_projection.projection); + + auto position = build_grouped_positional_conv( + ctx, + transpose_btc_to_bct(ctx, hubert), + weights_->positional_conv_groups, + config_.semantic.hidden_size, + config_.semantic.num_conv_pos_embeddings); + hubert = modules::AddModule().build(ctx, hubert, transpose_bct_to_btc(ctx, position)); + hubert = modules::LayerNormModule({ + config_.semantic.hidden_size, + config_.semantic.layer_norm_eps, + true, + true, + }).build(ctx, hubert, weights_->encoder_input_layer_norm); + hubert = build_hubert_sequence_mean(ctx, hubert, *weights_, config_); + + auto hubert_flat = core::reshape_tensor( + ctx, + ensure_dense_layout(ctx, hubert), + core::TensorShape::from_dims({hubert.shape.dims[1], hubert.shape.dims[2]})); + auto semantic_downsampled = modules::EmbeddingModule({ + hubert.shape.dims[1], + hubert.shape.dims[2], + }).build(ctx, downsample_indices_, hubert_flat); + semantic_downsampled = core::reshape_tensor( + ctx, + semantic_downsampled, + core::TensorShape::from_dims({1, frame_capacity_, config_.semantic.hidden_size})); + auto semantic_encoded = build_semantic_encoder(ctx, transpose_btc_to_bct(ctx, semantic_downsampled), *weights_, config_); + auto acoustic_encoded = build_acoustic_encoder(ctx, acoustic_input_, *weights_, config_); + if (semantic_encoded.shape.dims[2] != acoustic_encoded.shape.dims[2]) { + throw std::runtime_error("Higgs TTS semantic/acoustic encoder frame counts do not match"); + } + auto embeddings_bct = modules::ConcatModule({1}).build(ctx, acoustic_encoded, semantic_encoded); + auto embeddings_btc = transpose_bct_to_btc(ctx, embeddings_bct); + embeddings_btc = modules::LinearModule({ + config_.hidden_size, + config_.hidden_size, + true, + GGML_PREC_F32, + }).build(ctx, embeddings_btc, weights_->fc); + auto residual_bct = transpose_btc_to_bct(ctx, embeddings_btc); + + code_outputs_.reserve(weights_->quantizers.size()); + for (size_t quantizer_index = 0; quantizer_index < weights_->quantizers.size(); ++quantizer_index) { + const auto & quantizer = weights_->quantizers[quantizer_index]; + auto residual_btc = transpose_bct_to_btc(ctx, residual_bct); + auto projected = modules::LinearModule({ + config_.hidden_size, + config_.codebook_dim, + true, + GGML_PREC_F32, + }).build(ctx, residual_btc, quantizer.project_in); + auto logits = modules::LinearModule({ + config_.codebook_dim, + config_.codebook_size, + true, + GGML_PREC_F32, + }).build(ctx, projected, quantizer.score); + auto logits_flat = core::reshape_tensor( + ctx, + ensure_dense_layout(ctx, logits), + core::TensorShape::from_dims({frame_capacity_, config_.codebook_size})); + auto * ids_raw = ggml_argmax(ctx.ggml, logits_flat.tensor); + ggml_set_output(ids_raw); + code_outputs_.push_back(ids_raw); + auto ids = core::reshape_tensor( + ctx, + core::wrap_tensor(ids_raw, core::TensorShape::from_dims({frame_capacity_}), GGML_TYPE_I32), + core::TensorShape::from_dims({1, frame_capacity_})); + auto embedded = modules::CodebookLookupModule({ + config_.codebook_size, + config_.codebook_dim, + }).build(ctx, ids, quantizer.codebook); + auto quantized = modules::LinearModule({ + config_.codebook_dim, + config_.hidden_size, + true, + GGML_PREC_F32, + }).build(ctx, embedded, quantizer.project_out); + auto quantized_bct = transpose_btc_to_bct(ctx, quantized); + residual_bct = core::wrap_tensor( + ggml_sub(ctx.ggml, residual_bct.tensor, quantized_bct.tensor), + residual_bct.shape, + GGML_TYPE_F32); + } + + graph_ = ggml_new_graph_custom(ctx_.get(), 262144, false); + for (ggml_tensor * code_output : code_outputs_) { + ggml_build_forward_expand(graph_, code_output); + } + 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 Higgs TTS codec encoder graph"); + } + tensor_writer_.flush(); + debug::timing_log_scalar("higgs_tts.codec.encoder.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("higgs_tts.codec.encoder.frames", frame_capacity_); + debug::trace_log_scalar("higgs_tts.codec.encoder.acoustic_samples", acoustic_sample_capacity_); + debug::trace_log_scalar("higgs_tts.codec.encoder.semantic_samples", semantic_sample_capacity_); + } + + ~EncoderGraph() { + engine::core::release_backend_graph_resources(backend_, graph_); + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + } + + bool matches( + ggml_backend_t backend, + int threads, + int64_t acoustic_samples, + int64_t semantic_samples, + int64_t frames) const noexcept { + return backend_ == backend && + threads_ == threads && + acoustic_sample_capacity_ >= acoustic_samples && + semantic_sample_capacity_ >= semantic_samples && + frame_capacity_ >= frames; + } + + HiggsAudioCodeMatrix run(const NormalizedReferenceAudio & audio) { + if (static_cast(audio.acoustic_samples_24k.size()) > acoustic_sample_capacity_ || + static_cast(audio.semantic_samples_16k_padded.size()) > semantic_sample_capacity_ || + audio.frames > frame_capacity_) { + throw std::runtime_error("Higgs TTS codec encoder request exceeds prepared capacity"); + } + auto timing_start = Clock::now(); + std::vector acoustic_padded(static_cast(acoustic_sample_capacity_), 0.0F); + std::copy(audio.acoustic_samples_24k.begin(), audio.acoustic_samples_24k.end(), acoustic_padded.begin()); + core::write_tensor_f32(acoustic_input_, acoustic_padded); + std::vector semantic_padded(static_cast(semantic_sample_capacity_), 0.0F); + std::copy(audio.semantic_samples_16k_padded.begin(), audio.semantic_samples_16k_padded.end(), semantic_padded.begin()); + core::write_tensor_f32(semantic_input_, semantic_padded); + std::vector downsample(static_cast(frame_capacity_), 0); + for (int64_t i = 0; i < audio.frames; ++i) { + downsample[static_cast(i)] = static_cast(i * semantic_downsample_factor_); + } + core::write_tensor_i32(downsample_indices_, downsample); + tensor_writer_.flush(); + debug::timing_log_scalar("higgs_tts.codec.encoder.input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + core::set_backend_threads(backend_, threads_); + timing_start = Clock::now(); + const ggml_status status = engine::core::compute_backend_graph(backend_, graph_); + ggml_backend_synchronize(backend_); + debug::timing_log_scalar("higgs_tts.codec.encoder.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS codec encoder graph compute failed"); + } + + HiggsAudioCodeMatrix codes; + codes.frames = audio.frames; + codes.codebooks = static_cast(code_outputs_.size()); + codes.token_ids.assign(static_cast(codes.frames * codes.codebooks), 0); + timing_start = Clock::now(); + for (size_t codebook = 0; codebook < code_outputs_.size(); ++codebook) { + std::vector values(static_cast(frame_capacity_), 0); + ggml_backend_tensor_get( + code_outputs_[codebook], + values.data(), + 0, + values.size() * sizeof(int32_t)); + for (int64_t frame = 0; frame < codes.frames; ++frame) { + codes.token_ids[static_cast(frame * codes.codebooks + static_cast(codebook))] = + values[static_cast(frame)]; + } + } + debug::timing_log_scalar("higgs_tts.codec.encoder.output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + return codes; + } + +private: + std::shared_ptr weights_; + HiggsAudioCodecConfig config_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + int64_t acoustic_sample_capacity_ = 0; + int64_t semantic_sample_capacity_ = 0; + int64_t frame_capacity_ = 0; + int64_t semantic_downsample_factor_ = 1; + std::unique_ptr ctx_; + core::DeferredTensorWriter tensor_writer_; + core::TensorValue acoustic_input_; + core::TensorValue semantic_input_; + core::TensorValue downsample_indices_; + std::vector code_outputs_; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; +}; + +class DecoderGraph { +public: + DecoderGraph( + std::shared_ptr weights, + HiggsAudioCodecConfig config, + ggml_backend_t backend, + core::BackendType backend_type, + int threads, + size_t graph_arena_bytes, + int64_t frames, + int64_t codebooks) + : weights_(std::move(weights)), + config_(std::move(config)), + backend_(backend), + backend_type_(backend_type), + threads_(threads), + frames_(frames), + codebooks_(codebooks) { + if (backend_ == nullptr) { + throw std::runtime_error("Higgs TTS codec backend is not initialized"); + } + if (frames_ <= 0 || codebooks_ <= 0) { + throw std::runtime_error("Higgs TTS codec graph requires positive frame and codebook counts"); + } + if (static_cast(codebooks_) > weights_->quantizers.size()) { + throw std::runtime_error("Higgs TTS codec graph codebook count exceeds loaded weights"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS codec graph context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "higgs_tts.codec.decode", backend_type_}; + code_inputs_.reserve(static_cast(codebooks_)); + for (int64_t codebook = 0; codebook < codebooks_; ++codebook) { + auto input = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, frames_})); + ggml_set_input(input.tensor); + code_inputs_.push_back(input.tensor); + } + auto latent_bct = build_quantizer_decode_sequence(ctx, code_inputs_, *weights_, frames_); + auto acoustic_btc = modules::LinearModule({ + config_.hidden_size, + config_.acoustic_hidden_size, + true, + }).build(ctx, transpose_bct_to_btc(ctx, latent_bct), weights_->fc2); + auto audio = build_acoustic_decoder(ctx, transpose_btc_to_bct(ctx, acoustic_btc), *weights_, config_); + const int64_t expected_samples = frames_ * product(config_.upsampling_ratios); + if (audio.shape.dims[2] < expected_samples) { + throw std::runtime_error("Higgs TTS codec decoder produced fewer samples than expected"); + } + if (audio.shape.dims[2] != expected_samples) { + const int64_t trim = audio.shape.dims[2] - expected_samples; + if ((trim % 2) != 0) { + throw std::runtime_error("Higgs TTS codec decoder output trim is not symmetric"); + } + audio = modules::SliceModule({2, trim / 2, expected_samples}).build(ctx, audio); + } + output_ = audio.tensor; + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), 131072, false); + ggml_build_forward_expand(graph_, output_); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), backend_); + if (buffer_ == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS codec graph"); + } + code_input_host_.resize(static_cast(codebooks_)); + for (auto & ids : code_input_host_) { + ids.assign(static_cast(frames_), 0); + } + debug::timing_log_scalar("higgs_tts.codec.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("higgs_tts.codec.frames", frames_); + debug::trace_log_scalar("higgs_tts.codec.samples", expected_samples); + } + + ~DecoderGraph() { + engine::core::release_backend_graph_resources(backend_, graph_); + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + } + + bool matches(ggml_backend_t backend, int threads, int64_t frames, int64_t codebooks) const noexcept { + return backend_ == backend && threads_ == threads && frames_ == frames && codebooks_ == codebooks; + } + + runtime::AudioBuffer run(const HiggsAudioCodeMatrix & codes) { + if (codes.frames != frames_ || codes.codebooks != codebooks_) { + throw std::runtime_error("Higgs TTS codec graph code matrix shape mismatch"); + } + auto timing_start = Clock::now(); + for (int64_t codebook = 0; codebook < codebooks_; ++codebook) { + auto & ids = code_input_host_[static_cast(codebook)]; + for (int64_t frame = 0; frame < frames_; ++frame) { + ids[static_cast(frame)] = + codes.token_ids[static_cast(frame * codebooks_ + codebook)]; + } + ggml_backend_tensor_set( + code_inputs_[static_cast(codebook)], + ids.data(), + 0, + ids.size() * sizeof(int32_t)); + } + debug::timing_log_scalar("higgs_tts.codec.input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + core::set_backend_threads(backend_, threads_); + timing_start = Clock::now(); + const ggml_status status = engine::core::compute_backend_graph(backend_, graph_); + ggml_backend_synchronize(backend_); + debug::timing_log_scalar("higgs_tts.codec.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS codec graph compute failed"); + } + runtime::AudioBuffer audio; + audio.sample_rate = config_.sample_rate; + audio.channels = 1; + const int64_t samples = frames_ * product(config_.upsampling_ratios); + audio.samples.resize(static_cast(samples)); + timing_start = Clock::now(); + ggml_backend_tensor_get(output_, audio.samples.data(), 0, audio.samples.size() * sizeof(float)); + debug::timing_log_scalar("higgs_tts.codec.output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + return audio; + } + +private: + std::shared_ptr weights_; + HiggsAudioCodecConfig config_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + int64_t frames_ = 0; + int64_t codebooks_ = 0; + std::unique_ptr ctx_; + std::vector code_inputs_; + std::vector> code_input_host_; + ggml_tensor * output_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; +}; + +class HiggsAudioCodecWeightsRuntime { +public: + HiggsAudioCodecWeightsRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) + : assets_(std::move(assets)), + config_(codec_config_from_assets(*assets_)), + backend_(execution.backend()), + backend_type_(execution.backend_type()), + threads_(std::max(1, execution.config().threads)) { + if (assets_ == nullptr) { + throw std::runtime_error("Higgs TTS codec runtime requires assets"); + } + validate_codec_config(config_, assets_->config); + weights_ = load_weights(*assets_, config_, backend_, backend_type_, weight_context_bytes, storage_type); + } + + const HiggsAudioCodecConfig & config() const noexcept { + return config_; + } + + const std::shared_ptr & weights() const noexcept { + return weights_; + } + + ggml_backend_t backend() const noexcept { + return backend_; + } + + core::BackendType backend_type() const noexcept { + return backend_type_; + } + + int threads() const noexcept { + return threads_; + } + +private: + std::shared_ptr assets_; + HiggsAudioCodecConfig config_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + std::shared_ptr weights_; +}; + +} // namespace + +struct HiggsAudioCodecDecoderRuntime::Impl { + Impl( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) + : weights(std::make_shared( + std::move(assets), + execution, + weight_context_bytes, + storage_type)), + graph_arena_bytes(graph_arena_bytes) {} + + HiggsAudioCodeMatrix encode_reference_audio(const runtime::AudioBuffer & audio) { + const auto normalized = normalize_reference_audio(audio, weights->config()); + if (encoder_graph == nullptr || + !encoder_graph->matches( + weights->backend(), + weights->threads(), + static_cast(normalized.acoustic_samples_24k.size()), + static_cast(normalized.semantic_samples_16k_padded.size()), + normalized.frames)) { + encoder_graph = std::make_unique( + weights->weights(), + weights->config(), + weights->backend(), + weights->backend_type(), + weights->threads(), + graph_arena_bytes, + static_cast(normalized.acoustic_samples_24k.size()), + static_cast(normalized.semantic_samples_16k_padded.size()), + normalized.frames); + } else { + debug::timing_log_scalar("higgs_tts.codec.encoder.graph.build_ms", 0.0); + debug::trace_log_scalar("higgs_tts.codec.encoder.frames", normalized.frames); + debug::trace_log_scalar( + "higgs_tts.codec.encoder.acoustic_samples", + static_cast(normalized.acoustic_samples_24k.size())); + debug::trace_log_scalar( + "higgs_tts.codec.encoder.semantic_samples", + static_cast(normalized.semantic_samples_16k_padded.size())); + } + return encoder_graph->run(normalized); + } + + runtime::AudioBuffer decode(const HiggsAudioCodeMatrix & raw_codes) { + const auto codes = trim_decodable_codes(raw_codes, weights->config().codebook_size); + if (graph == nullptr || + !graph->matches(weights->backend(), weights->threads(), codes.frames, codes.codebooks)) { + graph = std::make_unique( + weights->weights(), + weights->config(), + weights->backend(), + weights->backend_type(), + weights->threads(), + graph_arena_bytes, + codes.frames, + codes.codebooks); + } else { + debug::timing_log_scalar("higgs_tts.codec.graph.build_ms", 0.0); + debug::trace_log_scalar("higgs_tts.codec.frames", codes.frames); + debug::trace_log_scalar("higgs_tts.codec.samples", codes.frames * product(weights->config().upsampling_ratios)); + } + return graph->run(codes); + } + + std::shared_ptr weights; + size_t graph_arena_bytes = 0; + std::unique_ptr encoder_graph; + std::unique_ptr graph; +}; + +HiggsAudioCodecDecoderRuntime::HiggsAudioCodecDecoderRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : impl_(std::make_unique( + std::move(assets), + execution, + graph_arena_bytes, + weight_context_bytes, + weight_storage_type)) {} + +HiggsAudioCodecDecoderRuntime::~HiggsAudioCodecDecoderRuntime() = default; + +HiggsAudioCodeMatrix HiggsAudioCodecDecoderRuntime::encode_reference_audio(const runtime::AudioBuffer & audio) { + return impl_->encode_reference_audio(audio); +} + +runtime::AudioBuffer HiggsAudioCodecDecoderRuntime::decode(const HiggsAudioCodeMatrix & raw_codes) { + return impl_->decode(raw_codes); +} + +} // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/generator.cpp b/src/models/higgs_tts/generator.cpp new file mode 100644 index 0000000..c926621 --- /dev/null +++ b/src/models/higgs_tts/generator.cpp @@ -0,0 +1,1179 @@ +#include "engine/models/higgs_tts/generator.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/modules/activation_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/optimizations/fast_kv_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/runtime/kv_cache.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::higgs_tts { +namespace { + +using Clock = std::chrono::steady_clock; +namespace modules = engine::modules; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct TextLayerWeights { + core::TensorValue input_norm; + core::TensorValue q_proj; + core::TensorValue k_proj; + core::TensorValue v_proj; + core::TensorValue o_proj; + core::TensorValue q_norm; + core::TensorValue k_norm; + core::TensorValue post_norm; + core::TensorValue gate_proj; + core::TensorValue up_proj; + core::TensorValue down_proj; +}; + +struct HiggsTTSGeneratorWeights { + std::shared_ptr store; + core::TensorValue text_embedding; + core::TensorValue modality_embedding; + std::vector layers; + core::TensorValue norm; +}; + +struct DecoderLayerOutputs { + core::TensorValue output; + core::TensorValue key; + core::TensorValue value; +}; + +struct PrefillOutput { + std::vector logits; + runtime::TransformerKVState kv_state; +}; + +struct SamplingCandidate { + int32_t token = 0; + float score = 0.0F; +}; + +struct SamplingScratch { + std::vector candidates; + std::vector probabilities; +}; + +struct HiggsSamplerState { + int64_t num_codebooks = 0; + int64_t delay_count = 0; + std::optional eoc_countdown; + bool generation_done = false; + std::vector last_codes; +}; + +struct PromptReferenceInfo { + int64_t start = 0; + int64_t tokens = 0; +}; + +int64_t head_dim(const HiggsTTSConfig & config) { + if (config.text.num_attention_heads <= 0 || config.text.num_key_value_heads <= 0 || config.text.head_dim <= 0) { + throw std::runtime_error("Higgs TTS attention config is invalid"); + } + return config.text.head_dim; +} + +int64_t modality_vocab_size(const HiggsTTSConfig & config) { + return config.audio_encoder.num_codebooks * config.audio_encoder.vocab_size; +} + +core::TensorValue reshape_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t heads, + int64_t dim) { + const auto contiguous = core::ensure_backend_addressable_layout(ctx, input); + 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; + } + std::vector heads; + heads.reserve(static_cast(input.shape.dims[1] * repeats)); + for (int64_t head = 0; head < input.shape.dims[1]; ++head) { + auto one = modules::SliceModule({1, head, 1}).build(ctx, input); + for (int64_t rep = 0; rep < repeats; ++rep) { + heads.push_back(one); + } + } + auto output = heads.front(); + for (size_t i = 1; i < heads.size(); ++i) { + output = modules::ConcatModule({1}).build(ctx, output, heads[i]); + } + return output; +} + +core::TensorValue attention_from_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & q_heads, + const core::TensorValue & k_heads, + const core::TensorValue & v_heads, + int64_t dim, + const std::optional & attention_mask = std::nullopt) { + const modules::MatMulModule matmul; + auto scores = matmul.build(ctx, q_heads, modules::TransposeModule({{0, 1, 3, 2}, k_heads.shape.rank}).build(ctx, k_heads)); + core::TensorValue attn; + if (attention_mask.has_value()) { + scores = core::ensure_backend_addressable_layout(ctx, scores); + attn = core::wrap_tensor( + ggml_soft_max_ext( + ctx.ggml, + scores.tensor, + attention_mask->tensor, + 1.0F / std::sqrt(static_cast(dim)), + 0.0F), + scores.shape, + GGML_TYPE_F32); + } else { + scores = core::wrap_tensor( + ggml_scale(ctx.ggml, scores.tensor, 1.0F / std::sqrt(static_cast(dim))), + scores.shape, + GGML_TYPE_F32); + scores = core::wrap_tensor(ggml_diag_mask_inf(ctx.ggml, scores.tensor, 0), scores.shape, GGML_TYPE_F32); + scores = core::ensure_backend_addressable_layout(ctx, scores); + attn = core::wrap_tensor(ggml_soft_max(ctx.ggml, scores.tensor), scores.shape, GGML_TYPE_F32); + } + return matmul.build(ctx, attn, v_heads); +} + +core::TensorValue flash_attention_from_grouped_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & q_heads, + const core::TensorValue & k_heads, + const core::TensorValue & v_heads, + int64_t dim, + const core::TensorValue & attention_mask) { + const auto q_contiguous = core::ensure_backend_addressable_layout(ctx, q_heads); + const auto k_contiguous = core::ensure_backend_addressable_layout(ctx, k_heads); + const auto v_contiguous = core::ensure_backend_addressable_layout(ctx, v_heads); + auto * flash = ggml_flash_attn_ext( + ctx.ggml, + q_contiguous.tensor, + k_contiguous.tensor, + v_contiguous.tensor, + attention_mask.tensor, + 1.0F / std::sqrt(static_cast(dim)), + 0.0F, + 0.0F); + ggml_flash_attn_ext_set_prec(flash, GGML_PREC_F32); + return core::wrap_tensor( + flash, + core::TensorShape::from_dims({q_contiguous.shape.dims[0], q_contiguous.shape.dims[2], q_contiguous.shape.dims[1], dim}), + GGML_TYPE_F32); +} + +DecoderLayerOutputs decoder_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const TextLayerWeights & weights, + const HiggsTTSConfig & config, + const std::optional & prefix_key = std::nullopt, + const std::optional & prefix_value = std::nullopt, + const std::optional & attention_mask = std::nullopt) { + const int64_t dim = head_dim(config); + const int64_t kv_repeats = config.text.num_attention_heads / config.text.num_key_value_heads; + const modules::LinearModule q_proj({config.text.hidden_size, config.text.num_attention_heads * dim, false}); + const modules::LinearModule k_proj({config.text.hidden_size, config.text.num_key_value_heads * dim, false}); + const modules::LinearModule v_proj({config.text.hidden_size, config.text.num_key_value_heads * dim, false}); + const modules::LinearModule o_proj({config.text.num_attention_heads * dim, config.text.hidden_size, false}); + const modules::RMSNormModule hidden_norm({config.text.hidden_size, config.text.rms_norm_eps, true, false}); + const modules::RMSNormModule head_norm({dim, config.text.rms_norm_eps, true, false}); + + auto x_norm = hidden_norm.build(ctx, input, {weights.input_norm, std::nullopt}); + auto q = q_proj.build(ctx, x_norm, {weights.q_proj, std::nullopt}); + auto k = k_proj.build(ctx, x_norm, {weights.k_proj, std::nullopt}); + auto v = v_proj.build(ctx, x_norm, {weights.v_proj, std::nullopt}); + q = head_norm.build(ctx, reshape_heads(ctx, q, config.text.num_attention_heads, dim), {weights.q_norm, std::nullopt}); + k = head_norm.build(ctx, reshape_heads(ctx, k, config.text.num_key_value_heads, dim), {weights.k_norm, std::nullopt}); + v = reshape_heads(ctx, v, config.text.num_key_value_heads, dim); + q = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.text.rope_theta}).build(ctx, q, positions); + k = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.text.rope_theta}).build(ctx, k, positions); + + auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); + auto all_k = prefix_key.has_value() ? modules::ConcatModule({1}).build(ctx, *prefix_key, k) : k; + auto all_v = prefix_value.has_value() ? modules::ConcatModule({1}).build(ctx, *prefix_value, v) : v; + auto k_heads = repeat_kv_heads(ctx, modules::TransposeModule({{0, 2, 1, 3}, all_k.shape.rank}).build(ctx, all_k), kv_repeats); + auto v_heads = repeat_kv_heads(ctx, modules::TransposeModule({{0, 2, 1, 3}, all_v.shape.rank}).build(ctx, all_v), kv_repeats); + auto context = attention_from_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask); + context = modules::TransposeModule({{0, 2, 1, 3}, context.shape.rank}).build(ctx, context); + 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], config.text.num_attention_heads * dim})); + auto x = modules::AddModule{}.build(ctx, input, o_proj.build(ctx, context, {weights.o_proj, std::nullopt})); + + auto ff_in = hidden_norm.build(ctx, x, {weights.post_norm, std::nullopt}); + auto gate = modules::LinearModule({config.text.hidden_size, config.text.intermediate_size, false}) + .build(ctx, ff_in, {weights.gate_proj, std::nullopt}); + gate = modules::SiluModule{}.build(ctx, gate); + auto up = modules::LinearModule({config.text.hidden_size, config.text.intermediate_size, false}) + .build(ctx, ff_in, {weights.up_proj, std::nullopt}); + auto gated = modules::MulModule{}.build(ctx, gate, up); + auto ff = modules::LinearModule({config.text.intermediate_size, config.text.hidden_size, false}) + .build(ctx, gated, {weights.down_proj, std::nullopt}); + return {modules::AddModule{}.build(ctx, x, ff), k, v}; +} + +DecoderLayerOutputs decoder_layer_with_static_cache( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const TextLayerWeights & weights, + const HiggsTTSConfig & config, + const core::TensorValue & cache_key, + const core::TensorValue & cache_value, + const core::TensorValue & cache_slot, + const core::TensorValue & attention_mask) { + const int64_t dim = head_dim(config); + const modules::LinearModule q_proj({config.text.hidden_size, config.text.num_attention_heads * dim, false}); + const modules::LinearModule k_proj({config.text.hidden_size, config.text.num_key_value_heads * dim, false}); + const modules::LinearModule v_proj({config.text.hidden_size, config.text.num_key_value_heads * dim, false}); + const modules::LinearModule o_proj({config.text.num_attention_heads * dim, config.text.hidden_size, false}); + const modules::RMSNormModule hidden_norm({config.text.hidden_size, config.text.rms_norm_eps, true, false}); + + auto x_norm = hidden_norm.build(ctx, input, {weights.input_norm, std::nullopt}); + auto q = q_proj.build(ctx, x_norm, {weights.q_proj, std::nullopt}); + auto k = k_proj.build(ctx, x_norm, {weights.k_proj, std::nullopt}); + auto v = v_proj.build(ctx, x_norm, {weights.v_proj, std::nullopt}); + q = modules::RMSNormModule({dim, config.text.rms_norm_eps, true, false}) + .build(ctx, reshape_heads(ctx, q, config.text.num_attention_heads, dim), {weights.q_norm, std::nullopt}); + k = modules::RMSNormModule({dim, config.text.rms_norm_eps, true, false}) + .build(ctx, reshape_heads(ctx, k, config.text.num_key_value_heads, dim), {weights.k_norm, std::nullopt}); + v = reshape_heads(ctx, v, config.text.num_key_value_heads, dim); + q = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.text.rope_theta}).build(ctx, q, positions); + k = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.text.rope_theta}).build(ctx, k, positions); + + const modules::FastKVSetRowsModule set_rows; + auto updated_cache_key = set_rows.build(ctx, cache_key, k, cache_slot); + auto updated_cache_value = set_rows.build(ctx, cache_value, v, cache_slot); + + auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); + auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, updated_cache_key.shape.rank}).build(ctx, updated_cache_key); + auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, updated_cache_value.shape.rank}).build(ctx, updated_cache_value); + auto context = flash_attention_from_grouped_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask); + context = core::ensure_backend_addressable_layout(ctx, context); + context = core::reshape_tensor(ctx, context, core::TensorShape::from_dims({1, 1, config.text.num_attention_heads * dim})); + auto x = modules::AddModule{}.build(ctx, input, o_proj.build(ctx, context, {weights.o_proj, std::nullopt})); + + auto ff_in = hidden_norm.build(ctx, x, {weights.post_norm, std::nullopt}); + auto gate = modules::LinearModule({config.text.hidden_size, config.text.intermediate_size, false}) + .build(ctx, ff_in, {weights.gate_proj, std::nullopt}); + gate = modules::SiluModule{}.build(ctx, gate); + auto up = modules::LinearModule({config.text.hidden_size, config.text.intermediate_size, false}) + .build(ctx, ff_in, {weights.up_proj, std::nullopt}); + auto gated = modules::MulModule{}.build(ctx, gate, up); + auto ff = modules::LinearModule({config.text.intermediate_size, config.text.hidden_size, false}) + .build(ctx, gated, {weights.down_proj, std::nullopt}); + return {modules::AddModule{}.build(ctx, x, ff), k, v}; +} + +core::TensorValue text_prompt_embeddings( + core::ModuleBuildContext & ctx, + const HiggsTTSGeneratorWeights & weights, + const HiggsTTSConfig & config, + ggml_tensor * token_ids, + int64_t prompt_steps) { + auto ids = core::wrap_tensor(token_ids, core::TensorShape::from_dims({prompt_steps}), GGML_TYPE_I32); + auto x = modules::EmbeddingModule({config.text.vocab_size, config.text.hidden_size}).build(ctx, ids, weights.text_embedding); + return core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, prompt_steps, config.text.hidden_size})); +} + +core::TensorValue text_prompt_slice_embeddings( + core::ModuleBuildContext & ctx, + const HiggsTTSGeneratorWeights & weights, + const HiggsTTSConfig & config, + const core::TensorValue & token_ids, + int64_t start, + int64_t steps) { + auto ids = modules::SliceModule({0, start, steps}).build(ctx, token_ids); + auto x = modules::EmbeddingModule({config.text.vocab_size, config.text.hidden_size}).build(ctx, ids, weights.text_embedding); + return core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, steps, config.text.hidden_size})); +} + +core::TensorValue reference_prompt_embeddings( + core::ModuleBuildContext & ctx, + const HiggsTTSGeneratorWeights & weights, + const HiggsTTSConfig & config, + const core::TensorValue & reference_code_ids, + int64_t reference_tokens) { + auto embeddings = modules::EmbeddingModule({modality_vocab_size(config), config.text.hidden_size}) + .build(ctx, reference_code_ids, weights.modality_embedding); + auto summed = modules::ReduceSumModule({1}).build(ctx, embeddings); + return core::reshape_tensor(ctx, summed, core::TensorShape::from_dims({1, reference_tokens, config.text.hidden_size})); +} + +core::TensorValue prompt_embeddings_with_optional_reference( + core::ModuleBuildContext & ctx, + const HiggsTTSGeneratorWeights & weights, + const HiggsTTSConfig & config, + ggml_tensor * token_ids, + ggml_tensor * reference_code_ids, + int64_t prompt_steps, + const PromptReferenceInfo & reference) { + if (reference.tokens == 0) { + return text_prompt_embeddings(ctx, weights, config, token_ids, prompt_steps); + } + auto ids = core::wrap_tensor(token_ids, core::TensorShape::from_dims({prompt_steps}), GGML_TYPE_I32); + auto prefix = text_prompt_slice_embeddings(ctx, weights, config, ids, 0, reference.start); + auto reference_ids = core::wrap_tensor( + reference_code_ids, + core::TensorShape::from_dims({reference.tokens, config.audio_encoder.num_codebooks}), + GGML_TYPE_I32); + auto reference_embed = reference_prompt_embeddings(ctx, weights, config, reference_ids, reference.tokens); + auto x = modules::ConcatModule({1}).build(ctx, prefix, reference_embed); + const int64_t suffix_start = reference.start + reference.tokens; + const int64_t suffix_steps = prompt_steps - suffix_start; + if (suffix_steps > 0) { + auto suffix = text_prompt_slice_embeddings(ctx, weights, config, ids, suffix_start, suffix_steps); + x = modules::ConcatModule({1}).build(ctx, x, suffix); + } + return x; +} + +core::TensorValue modality_code_embedding( + core::ModuleBuildContext & ctx, + const HiggsTTSGeneratorWeights & weights, + const HiggsTTSConfig & config, + ggml_tensor * offset_code_ids) { + auto ids = core::wrap_tensor( + offset_code_ids, + core::TensorShape::from_dims({config.audio_encoder.num_codebooks}), + GGML_TYPE_I32); + auto embeddings = modules::EmbeddingModule({modality_vocab_size(config), config.text.hidden_size}) + .build(ctx, ids, weights.modality_embedding); + auto summed = modules::ReduceSumModule({0}).build(ctx, embeddings); + return core::reshape_tensor(ctx, summed, core::TensorShape::from_dims({1, 1, config.text.hidden_size})); +} + +core::TensorValue modality_logits( + core::ModuleBuildContext & ctx, + const core::TensorValue & hidden, + const HiggsTTSGeneratorWeights & weights, + const HiggsTTSConfig & config) { + return modules::LinearModule({config.text.hidden_size, modality_vocab_size(config), false}) + .build(ctx, hidden, {weights.modality_embedding, std::nullopt}); +} + +HiggsTTSGeneratorWeights load_weights( + const HiggsTTSAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) { + const auto & config = assets.config; + const auto & source = *assets.model_weights; + HiggsTTSGeneratorWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "higgs_tts.generator.weights", + weight_context_bytes); + weights.text_embedding = weights.store->load_tensor( + source, + "tied.embedding.text_embedding.weight", + storage_type, + {config.text.vocab_size, config.text.hidden_size}); + weights.modality_embedding = weights.store->load_tensor( + source, + "tied.embedding.modality_embeddings.0.embedding.weight", + storage_type, + {modality_vocab_size(config), config.text.hidden_size}); + weights.layers.reserve(static_cast(config.text.num_hidden_layers)); + const int64_t dim = head_dim(config); + for (int64_t layer = 0; layer < config.text.num_hidden_layers; ++layer) { + const std::string prefix = "body.layers." + std::to_string(layer); + TextLayerWeights w; + w.input_norm = weights.store->load_f32_tensor(source, prefix + ".input_layernorm.weight", {config.text.hidden_size}); + w.q_proj = weights.store->load_tensor(source, prefix + ".self_attn.q_proj.weight", storage_type, {config.text.num_attention_heads * dim, config.text.hidden_size}); + w.k_proj = weights.store->load_tensor(source, prefix + ".self_attn.k_proj.weight", storage_type, {config.text.num_key_value_heads * dim, config.text.hidden_size}); + w.v_proj = weights.store->load_tensor(source, prefix + ".self_attn.v_proj.weight", storage_type, {config.text.num_key_value_heads * dim, config.text.hidden_size}); + w.o_proj = weights.store->load_tensor(source, prefix + ".self_attn.o_proj.weight", storage_type, {config.text.hidden_size, config.text.num_attention_heads * dim}); + w.q_norm = weights.store->load_f32_tensor(source, prefix + ".self_attn.q_norm.weight", {dim}); + w.k_norm = weights.store->load_f32_tensor(source, prefix + ".self_attn.k_norm.weight", {dim}); + w.post_norm = weights.store->load_f32_tensor(source, prefix + ".post_attention_layernorm.weight", {config.text.hidden_size}); + w.gate_proj = weights.store->load_tensor(source, prefix + ".mlp.gate_proj.weight", storage_type, {config.text.intermediate_size, config.text.hidden_size}); + w.up_proj = weights.store->load_tensor(source, prefix + ".mlp.up_proj.weight", storage_type, {config.text.intermediate_size, config.text.hidden_size}); + w.down_proj = weights.store->load_tensor(source, prefix + ".mlp.down_proj.weight", storage_type, {config.text.hidden_size, config.text.intermediate_size}); + weights.layers.push_back(std::move(w)); + } + weights.norm = weights.store->load_f32_tensor(source, "body.norm.weight", {config.text.hidden_size}); + weights.store->upload(); + return weights; +} + +PromptReferenceInfo resolve_prompt_reference( + const HiggsTTSPrompt & prompt, + const HiggsAudioCodeMatrix * reference_delayed_codes, + const HiggsTTSConfig & config) { + PromptReferenceInfo info; + int64_t placeholders = 0; + int64_t first_placeholder = -1; + int64_t last_placeholder = -1; + for (int64_t i = 0; i < static_cast(prompt.input_ids.size()); ++i) { + if (prompt.input_ids[static_cast(i)] == kHiggsAudioPlaceholderId) { + if (first_placeholder < 0) { + first_placeholder = i; + } + last_placeholder = i; + ++placeholders; + } + } + if (placeholders == 0) { + if (reference_delayed_codes != nullptr && reference_delayed_codes->frames > 0) { + throw std::runtime_error("Higgs TTS reference audio codes were provided, but the prompt has no reference slots"); + } + return info; + } + if (last_placeholder - first_placeholder + 1 != placeholders) { + throw std::runtime_error("Higgs TTS reference placeholders must be contiguous"); + } + if (reference_delayed_codes == nullptr) { + throw std::runtime_error("Higgs TTS reference prompt requires delayed reference audio codes"); + } + if (reference_delayed_codes->frames != placeholders) { + throw std::runtime_error("Higgs TTS reference code count does not match the prompt placeholder count"); + } + if (reference_delayed_codes->codebooks != config.audio_encoder.num_codebooks) { + throw std::runtime_error("Higgs TTS reference codebook count mismatch"); + } + if (static_cast(reference_delayed_codes->token_ids.size()) != + reference_delayed_codes->frames * reference_delayed_codes->codebooks) { + throw std::runtime_error("Higgs TTS reference code matrix shape is invalid"); + } + for (const int32_t code : reference_delayed_codes->token_ids) { + if (code < 0 || static_cast(code) >= config.audio_encoder.vocab_size) { + throw std::runtime_error("Higgs TTS reference code is outside the modality vocabulary"); + } + } + info.start = first_placeholder; + info.tokens = placeholders; + return info; +} + +std::vector offset_modality_code_matrix( + const HiggsTTSConfig & config, + const HiggsAudioCodeMatrix & codes) { + if (codes.codebooks != config.audio_encoder.num_codebooks) { + throw std::runtime_error("Higgs TTS reference codebook count mismatch"); + } + if (static_cast(codes.token_ids.size()) != codes.frames * codes.codebooks) { + throw std::runtime_error("Higgs TTS reference code matrix shape is invalid"); + } + std::vector out(codes.token_ids.size(), 0); + for (int64_t frame = 0; frame < codes.frames; ++frame) { + for (int64_t codebook = 0; codebook < codes.codebooks; ++codebook) { + const size_t index = static_cast(frame * codes.codebooks + codebook); + const int32_t code = codes.token_ids[index]; + if (code < 0 || static_cast(code) >= config.audio_encoder.vocab_size) { + throw std::runtime_error("Higgs TTS reference code is outside the modality vocabulary"); + } + out[index] = static_cast(codebook * config.audio_encoder.vocab_size + static_cast(code)); + } + } + return out; +} + +int32_t argmax_row(const std::vector & logits, int64_t row, int64_t vocab_size) { + const int64_t offset = row * vocab_size; + int64_t best = 0; + for (int64_t token = 1; token < vocab_size; ++token) { + if (logits[static_cast(offset + token)] > logits[static_cast(offset + best)]) { + best = token; + } + } + return static_cast(best); +} + +int32_t sample_row( + const std::vector & logits, + int64_t row, + int64_t vocab_size, + const HiggsTTSGenerationOptions & options, + std::mt19937 & rng, + SamplingScratch & scratch) { + if (!options.do_sample || options.top_k == 1 || options.temperature <= 1.0e-5F) { + return argmax_row(logits, row, vocab_size); + } + if (!(options.temperature > 0.0F)) { + throw std::runtime_error("Higgs TTS temperature must be positive"); + } + if (!(options.top_p >= 0.0F && options.top_p <= 1.0F)) { + throw std::runtime_error("Higgs TTS top_p must be in [0, 1]"); + } + auto & candidates = scratch.candidates; + candidates.clear(); + candidates.reserve(static_cast(vocab_size)); + const int64_t offset = row * vocab_size; + for (int64_t token = 0; token < vocab_size; ++token) { + candidates.push_back({ + static_cast(token), + logits[static_cast(offset + token)] / options.temperature, + }); + } + auto candidate_order = [](const SamplingCandidate & lhs, const SamplingCandidate & rhs) { + if (lhs.score == rhs.score) { + return lhs.token < rhs.token; + } + return lhs.score > rhs.score; + }; + if (options.top_k > 0 && static_cast(options.top_k) < candidates.size()) { + const auto top_end = candidates.begin() + static_cast(options.top_k); + std::nth_element(candidates.begin(), top_end, candidates.end(), candidate_order); + std::sort(candidates.begin(), top_end, candidate_order); + candidates.erase(top_end, candidates.end()); + } else { + std::sort(candidates.begin(), candidates.end(), candidate_order); + } + const float max_score = candidates.front().score; + double total = 0.0; + auto & probabilities = scratch.probabilities; + probabilities.assign(candidates.size(), 0.0); + for (size_t i = 0; i < candidates.size(); ++i) { + probabilities[i] = std::exp(static_cast(candidates[i].score - max_score)); + total += probabilities[i]; + } + if (!(total > 0.0) || !std::isfinite(total)) { + throw std::runtime_error("Higgs TTS sampler produced invalid probability mass"); + } + for (double & probability : probabilities) { + probability /= total; + } + if (options.top_p < 1.0F) { + double cumulative = 0.0; + size_t keep = probabilities.size(); + for (size_t i = 0; i < probabilities.size(); ++i) { + cumulative += probabilities[i]; + if (cumulative >= static_cast(options.top_p)) { + keep = i + 1; + break; + } + } + candidates.resize(keep); + probabilities.resize(keep); + double kept_total = 0.0; + for (const double probability : probabilities) { + kept_total += probability; + } + if (!(kept_total > 0.0)) { + throw std::runtime_error("Higgs TTS top-p removed all probability mass"); + } + for (double & probability : probabilities) { + probability /= kept_total; + } + } + std::discrete_distribution distribution(probabilities.begin(), probabilities.end()); + return candidates[distribution(rng)].token; +} + +std::vector sample_higgs_step( + const std::vector & logits, + HiggsSamplerState & state, + const HiggsTTSConfig & config, + const HiggsTTSGenerationOptions & options, + std::mt19937 & rng, + SamplingScratch & scratch) { + const int64_t codebooks = config.audio_encoder.num_codebooks; + const int64_t vocab_size = config.audio_encoder.vocab_size; + if (state.generation_done) { + return std::vector(static_cast(codebooks), kHiggsAudioStopCode); + } + if (static_cast(logits.size()) != modality_vocab_size(config)) { + throw std::runtime_error("Higgs TTS modality logits shape mismatch"); + } + std::vector codes(static_cast(codebooks), 0); + for (int64_t codebook = 0; codebook < codebooks; ++codebook) { + codes[static_cast(codebook)] = sample_row(logits, codebook, vocab_size, options, rng, scratch); + } + if (state.delay_count < codebooks) { + const int64_t next_codebook = state.delay_count + 1; + if (next_codebook < codebooks) { + for (int64_t codebook = next_codebook; codebook < codebooks; ++codebook) { + codes[static_cast(codebook)] = kHiggsAudioBocId; + } + } + ++state.delay_count; + } else if (state.eoc_countdown.has_value()) { + --(*state.eoc_countdown); + if (*state.eoc_countdown <= 0) { + state.generation_done = true; + } + } else if (codes.front() == kHiggsAudioEocId) { + state.eoc_countdown = codebooks > 2 ? codebooks - 2 : 0; + if (codebooks <= 2) { + state.generation_done = true; + } + } + if (!state.generation_done) { + state.last_codes = codes; + } + return codes; +} + +std::vector offset_modality_codes(const HiggsTTSConfig & config, const std::vector & codes) { + if (static_cast(codes.size()) != config.audio_encoder.num_codebooks) { + throw std::runtime_error("Higgs TTS modality codebook count mismatch"); + } + std::vector out(codes.size(), 0); + for (int64_t codebook = 0; codebook < config.audio_encoder.num_codebooks; ++codebook) { + const int32_t code = codes[static_cast(codebook)]; + if (code < 0 || static_cast(code) >= config.audio_encoder.vocab_size) { + throw std::runtime_error("Higgs TTS generated code is outside the modality vocabulary"); + } + out[static_cast(codebook)] = static_cast( + codebook * config.audio_encoder.vocab_size + static_cast(code)); + } + return out; +} + +class HiggsTTSGeneratorWeightsRuntime { +public: + HiggsTTSGeneratorWeightsRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) + : assets_(std::move(assets)), + backend_(execution.backend()), + backend_type_(execution.backend_type()), + threads_(std::max(1, execution.config().threads)), + weights_(std::make_shared( + load_weights(*assets_, backend_, backend_type_, weight_context_bytes, storage_type))) { + if (assets_ == nullptr) { + throw std::runtime_error("Higgs TTS generator weights runtime requires assets"); + } + if (backend_ == nullptr) { + throw std::runtime_error("Higgs TTS generator backend is not initialized"); + } + } + + const HiggsTTSAssets & assets() const noexcept { + return *assets_; + } + + const HiggsTTSGeneratorWeights & weights() const noexcept { + return *weights_; + } + + ggml_backend_t backend() const noexcept { + return backend_; + } + + core::BackendType backend_type() const noexcept { + return backend_type_; + } + + int threads() const noexcept { + return threads_; + } + +private: + std::shared_ptr assets_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + std::shared_ptr weights_; +}; + +class PrefillGraph { +public: + PrefillGraph( + std::shared_ptr runtime, + int64_t prompt_steps, + PromptReferenceInfo reference, + size_t graph_arena_bytes) + : runtime_(std::move(runtime)), + prompt_steps_(prompt_steps), + reference_(reference) { + if (prompt_steps_ <= 0) { + throw std::runtime_error("Higgs TTS generator prefill requires positive prompt length"); + } + if (reference_.tokens < 0 || reference_.start < 0 || reference_.start + reference_.tokens > prompt_steps_) { + throw std::runtime_error("Higgs TTS generator prefill reference span is invalid"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS generator prefill graph context"); + } + const auto & config = runtime_->assets().config; + const auto & weights = runtime_->weights(); + core::ModuleBuildContext ctx{ctx_.get(), "higgs_tts.generator.prefill", runtime_->backend_type()}; + token_ids_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_); + if (reference_.tokens > 0) { + auto reference_codes = core::make_tensor( + ctx, + GGML_TYPE_I32, + core::TensorShape::from_dims({reference_.tokens, config.audio_encoder.num_codebooks})); + reference_code_ids_ = reference_codes.tensor; + } + auto x = prompt_embeddings_with_optional_reference( + ctx, + weights, + config, + token_ids_, + reference_code_ids_, + prompt_steps_, + reference_); + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_); + auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({prompt_steps_}), GGML_TYPE_I32); + + for (const auto & layer : weights.layers) { + auto out = decoder_layer(ctx, x, positions, layer, config); + x = out.output; + keys_.push_back(out.key.tensor); + values_.push_back(out.value.tensor); + } + x = modules::SliceModule({1, prompt_steps_ - 1, 1}).build(ctx, x); + x = modules::RMSNormModule({config.text.hidden_size, config.text.rms_norm_eps, true, false}) + .build(ctx, x, {weights.norm, std::nullopt}); + auto logits = modality_logits(ctx, x, weights, config); + logits_ = logits.tensor; + ggml_set_output(logits_); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + ggml_build_forward_expand(graph_, logits_); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); + if (buffer_ == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS generator prefill graph"); + } + std::vector pos(static_cast(prompt_steps_), 0); + for (int64_t i = 0; i < prompt_steps_; ++i) { + pos[static_cast(i)] = static_cast(i); + } + ggml_backend_tensor_set(positions_, pos.data(), 0, pos.size() * sizeof(int32_t)); + debug::timing_log_scalar("higgs_tts.generator.prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("higgs_tts.generator.prefill_prompt_steps", prompt_steps_); + } + + ~PrefillGraph() { + engine::core::release_backend_graph_resources(runtime_->backend(), graph_); + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + } + + bool matches( + const HiggsTTSGeneratorWeightsRuntime & runtime, + int64_t prompt_steps, + const PromptReferenceInfo & reference) const { + return runtime_.get() == &runtime && + prompt_steps_ == prompt_steps && + reference_.start == reference.start && + reference_.tokens == reference.tokens; + } + + PrefillOutput run( + const std::vector & token_ids, + const HiggsAudioCodeMatrix * reference_delayed_codes) { + const auto & config = runtime_->assets().config; + if (static_cast(token_ids.size()) != prompt_steps_) { + throw std::runtime_error("Higgs TTS generator prefill token id count mismatch"); + } + auto timing_start = Clock::now(); + ggml_backend_tensor_set(token_ids_, token_ids.data(), 0, token_ids.size() * sizeof(int32_t)); + if (reference_.tokens > 0) { + if (reference_delayed_codes == nullptr || reference_delayed_codes->frames != reference_.tokens) { + throw std::runtime_error("Higgs TTS generator prefill reference code count mismatch"); + } + const auto offset_reference_codes = offset_modality_code_matrix(config, *reference_delayed_codes); + ggml_backend_tensor_set( + reference_code_ids_, + offset_reference_codes.data(), + 0, + offset_reference_codes.size() * sizeof(int32_t)); + } + debug::timing_log_scalar("higgs_tts.generator.prefill_input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + timing_start = Clock::now(); + const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); + ggml_backend_synchronize(runtime_->backend()); + debug::timing_log_scalar("higgs_tts.generator.prefill.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS generator prefill graph compute failed"); + } + PrefillOutput out; + out.logits.resize(static_cast(modality_vocab_size(config))); + timing_start = Clock::now(); + ggml_backend_tensor_get( + logits_, + out.logits.data(), + 0, + out.logits.size() * sizeof(float)); + out.kv_state.current_end = prompt_steps_; + out.kv_state.layers.resize(keys_.size()); + const size_t layer_values = static_cast(prompt_steps_ * config.text.num_key_value_heads * head_dim(config)); + for (size_t layer = 0; layer < keys_.size(); ++layer) { + auto & state = out.kv_state.layers[layer]; + state.valid_steps = prompt_steps_; + state.key.resize(layer_values); + state.value.resize(layer_values); + ggml_backend_tensor_get(keys_[layer], state.key.data(), 0, state.key.size() * sizeof(float)); + ggml_backend_tensor_get(values_[layer], state.value.data(), 0, state.value.size() * sizeof(float)); + } + debug::timing_log_scalar("higgs_tts.generator.prefill_output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + return out; + } + +private: + std::shared_ptr runtime_; + int64_t prompt_steps_ = 0; + PromptReferenceInfo reference_; + std::unique_ptr ctx_; + ggml_tensor * token_ids_ = nullptr; + ggml_tensor * reference_code_ids_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * logits_ = nullptr; + std::vector keys_; + std::vector values_; + ggml_cgraph * graph_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; +}; + +class DecodeGraph { +public: + DecodeGraph( + std::shared_ptr runtime, + int64_t cache_steps, + size_t graph_arena_bytes) + : runtime_(std::move(runtime)), + cache_steps_(cache_steps) { + if (cache_steps_ <= 0) { + throw std::runtime_error("Higgs TTS generator decode requires positive cache length"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS generator decode graph context"); + } + const auto & config = runtime_->assets().config; + const auto & weights = runtime_->weights(); + const int64_t dim = head_dim(config); + core::ModuleBuildContext ctx{ctx_.get(), "higgs_tts.generator.decode", runtime_->backend_type()}; + offset_code_ids_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, config.audio_encoder.num_codebooks); + auto x = modality_code_embedding(ctx, weights, config, offset_code_ids_); + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + cache_slot_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + auto cache_slot = core::wrap_tensor(cache_slot_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + attention_mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, cache_steps_, 1, 1, 1); + auto attention_mask = core::wrap_tensor( + attention_mask_, + core::TensorShape::from_dims({1, 1, 1, cache_steps_}), + GGML_TYPE_F16); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + std::vector cache_keys; + std::vector cache_values; + for (const auto & layer : weights.layers) { + cache_keys.push_back(core::make_tensor( + ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, cache_steps_, config.text.num_key_value_heads, dim}))); + cache_values.push_back(core::make_tensor( + ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, cache_steps_, config.text.num_key_value_heads, dim}))); + auto out = decoder_layer_with_static_cache( + ctx, + x, + positions, + layer, + config, + cache_keys.back(), + cache_values.back(), + cache_slot, + attention_mask); + x = out.output; + } + step_cache_ = runtime::TransformerKVCache( + cache_steps_, + config.text.num_key_value_heads * dim, + std::move(cache_keys), + std::move(cache_values)); + x = modules::RMSNormModule({config.text.hidden_size, config.text.rms_norm_eps, true, false}) + .build(ctx, x, {weights.norm, std::nullopt}); + auto logits = modality_logits(ctx, x, weights, config); + logits_ = logits.tensor; + ggml_set_output(logits_); + ggml_build_forward_expand(graph_, logits_); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); + if (buffer_ == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS generator decode graph"); + } + attention_mask_values_.assign(static_cast(cache_steps_), ggml_fp32_to_fp16(-INFINITY)); + debug::timing_log_scalar("higgs_tts.generator.decode.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("higgs_tts.generator.decode_cache_steps", cache_steps_); + } + + ~DecodeGraph() { + engine::core::release_backend_graph_resources(runtime_->backend(), graph_); + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + } + + bool can_run(const HiggsTTSGeneratorWeightsRuntime & runtime, int64_t required_steps) const { + return runtime_.get() == &runtime && cache_steps_ >= required_steps; + } + + void import_state(const runtime::TransformerKVState & state) { + step_cache_.import_state(state); + } + + void reset_timing() noexcept { + input_upload_ms_ = 0.0; + mask_upload_ms_ = 0.0; + graph_compute_ms_ = 0.0; + logits_read_ms_ = 0.0; + } + + void run_step_into(const std::vector & codes, std::vector & logits) { + const auto & config = runtime_->assets().config; + if (step_cache_.valid_steps() >= cache_steps_) { + throw std::runtime_error("Higgs TTS generator decode cache exhausted"); + } + auto timing_start = Clock::now(); + const auto offset_codes = offset_modality_codes(config, codes); + ggml_backend_tensor_set(offset_code_ids_, offset_codes.data(), 0, offset_codes.size() * sizeof(int32_t)); + const int32_t position = static_cast(step_cache_.current_end()); + ggml_backend_tensor_set(positions_, &position, 0, sizeof(int32_t)); + const int32_t cache_slot = static_cast(step_cache_.valid_steps()); + ggml_backend_tensor_set(cache_slot_, &cache_slot, 0, sizeof(int32_t)); + input_upload_ms_ += engine::debug::elapsed_ms(timing_start, Clock::now()); + const auto masked = ggml_fp32_to_fp16(-INFINITY); + const auto visible = ggml_fp32_to_fp16(0.0F); + std::fill(attention_mask_values_.begin(), attention_mask_values_.end(), masked); + for (int64_t i = 0; i < step_cache_.valid_steps(); ++i) { + attention_mask_values_[static_cast(i)] = visible; + } + attention_mask_values_[static_cast(cache_slot)] = visible; + timing_start = Clock::now(); + ggml_backend_tensor_set( + attention_mask_, + attention_mask_values_.data(), + 0, + attention_mask_values_.size() * sizeof(ggml_fp16_t)); + mask_upload_ms_ += engine::debug::elapsed_ms(timing_start, Clock::now()); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + timing_start = Clock::now(); + const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); + ggml_backend_synchronize(runtime_->backend()); + graph_compute_ms_ += engine::debug::elapsed_ms(timing_start, Clock::now()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS generator decode graph compute failed"); + } + logits.resize(static_cast(modality_vocab_size(config))); + timing_start = Clock::now(); + ggml_backend_tensor_get(logits_, logits.data(), 0, logits.size() * sizeof(float)); + logits_read_ms_ += engine::debug::elapsed_ms(timing_start, Clock::now()); + step_cache_.advance_after_direct_append(1); + } + + void log_timing() const { + debug::timing_log_scalar("higgs_tts.generator.decode.input_upload_ms", input_upload_ms_); + debug::timing_log_scalar("higgs_tts.generator.decode.mask_upload_ms", mask_upload_ms_); + debug::timing_log_scalar("higgs_tts.generator.decode.graph.compute_ms", graph_compute_ms_); + debug::timing_log_scalar("higgs_tts.generator.decode.logits_read_ms", logits_read_ms_); + } + +private: + std::shared_ptr runtime_; + int64_t cache_steps_ = 0; + std::unique_ptr ctx_; + ggml_tensor * offset_code_ids_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * cache_slot_ = nullptr; + ggml_tensor * attention_mask_ = nullptr; + ggml_tensor * logits_ = nullptr; + std::vector attention_mask_values_; + runtime::TransformerKVCache step_cache_; + double input_upload_ms_ = 0.0; + double mask_upload_ms_ = 0.0; + double graph_compute_ms_ = 0.0; + double logits_read_ms_ = 0.0; + ggml_cgraph * graph_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; +}; + +} // namespace + +struct HiggsTTSGeneratorRuntime::Impl { + Impl( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) + : weights(std::make_shared( + std::move(assets), + execution, + weight_context_bytes, + storage_type)), + prefill_graph_arena_bytes(prefill_graph_arena_bytes), + decode_graph_arena_bytes(decode_graph_arena_bytes) {} + + HiggsTTSGeneratedCodes generate( + const HiggsTTSPrompt & prompt, + const HiggsTTSGenerationOptions & options, + const HiggsAudioCodeMatrix * reference_delayed_codes) { + const auto & config = weights->assets().config; + if (prompt.input_ids.empty()) { + throw std::runtime_error("Higgs TTS generator prompt is empty"); + } + if (options.max_tokens <= 0) { + throw std::runtime_error("Higgs TTS max_tokens must be positive"); + } + if (!(options.temperature >= 0.0F && options.temperature <= 2.0F)) { + throw std::runtime_error("Higgs TTS temperature must be in [0, 2]"); + } + if (!(options.top_p >= 0.0F && options.top_p <= 1.0F)) { + throw std::runtime_error("Higgs TTS top_p must be in [0, 1]"); + } + const int64_t prompt_steps = static_cast(prompt.input_ids.size()); + if (prompt_steps + options.max_tokens > config.text.max_position_embeddings) { + throw std::runtime_error("Higgs TTS request exceeds max_position_embeddings"); + } + const auto reference = resolve_prompt_reference(prompt, reference_delayed_codes, config); + if (prefill_graph == nullptr || !prefill_graph->matches(*weights, prompt_steps, reference)) { + prefill_graph = std::make_unique( + weights, + prompt_steps, + reference, + prefill_graph_arena_bytes); + } else { + debug::timing_log_scalar("higgs_tts.generator.prefill.graph.build_ms", 0.0); + debug::trace_log_scalar("higgs_tts.generator.prefill_prompt_steps", prompt_steps); + } + auto timing_start = Clock::now(); + auto prefill = prefill_graph->run(prompt.input_ids, reference_delayed_codes); + debug::timing_log_scalar("higgs_tts.generator.prefill_total_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + const int64_t required_cache_steps = prompt_steps + options.max_tokens; + if (decode_graph == nullptr || !decode_graph->can_run(*weights, required_cache_steps)) { + decode_graph = std::make_unique(weights, required_cache_steps, decode_graph_arena_bytes); + } else { + debug::timing_log_scalar("higgs_tts.generator.decode.graph.build_ms", 0.0); + debug::trace_log_scalar("higgs_tts.generator.decode_cache_steps", required_cache_steps); + } + decode_graph->import_state(prefill.kv_state); + decode_graph->reset_timing(); + + HiggsSamplerState sampler; + sampler.num_codebooks = config.audio_encoder.num_codebooks; + std::mt19937 rng(options.seed); + SamplingScratch scratch; + HiggsAudioCodeMatrix delayed; + delayed.codebooks = config.audio_encoder.num_codebooks; + delayed.token_ids.reserve(static_cast(options.max_tokens * delayed.codebooks)); + std::vector logits = std::move(prefill.logits); + double sampling_ms = 0.0; + double decode_step_ms = 0.0; + timing_start = Clock::now(); + for (int64_t step = 0; step < options.max_tokens; ++step) { + const auto sampling_start = Clock::now(); + const auto codes = sample_higgs_step(logits, sampler, config, options, rng, scratch); + sampling_ms += engine::debug::elapsed_ms(sampling_start, Clock::now()); + if (!codes.empty() && codes.front() != kHiggsAudioStopCode) { + delayed.token_ids.insert(delayed.token_ids.end(), codes.begin(), codes.end()); + ++delayed.frames; + } + if (sampler.generation_done) { + break; + } + if (sampler.last_codes.empty()) { + break; + } + const auto step_start = Clock::now(); + decode_graph->run_step_into(sampler.last_codes, logits); + decode_step_ms += engine::debug::elapsed_ms(step_start, Clock::now()); + } + debug::trace_log_scalar("higgs_tts.generator.stopped_by_eoc", sampler.generation_done); + debug::trace_log_scalar("higgs_tts.generator.hit_max_tokens", !sampler.generation_done); + debug::timing_log_scalar("higgs_tts.generator.sampling_ms", sampling_ms); + debug::timing_log_scalar("higgs_tts.generator.decode_step_ms", decode_step_ms); + decode_graph->log_timing(); + debug::timing_log_scalar("higgs_tts.generator.decode_total_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + debug::trace_log_scalar("higgs_tts.generator.delayed_rows", delayed.frames); + if (delayed.frames < delayed.codebooks) { + throw std::runtime_error("Higgs TTS generated too few audio-code rows"); + } + HiggsTTSGeneratedCodes out; + out.delayed_codes = std::move(delayed); + out.raw_codes = reverse_delay_pattern(out.delayed_codes); + debug::trace_log_scalar("higgs_tts.generator.raw_frames", out.raw_codes.frames); + return out; + } + + std::shared_ptr weights; + size_t prefill_graph_arena_bytes = 0; + size_t decode_graph_arena_bytes = 0; + std::unique_ptr prefill_graph; + std::unique_ptr decode_graph; +}; + +HiggsTTSGeneratorRuntime::HiggsTTSGeneratorRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : impl_(std::make_unique( + std::move(assets), + execution, + prefill_graph_arena_bytes, + decode_graph_arena_bytes, + weight_context_bytes, + weight_storage_type)) {} + +HiggsTTSGeneratorRuntime::~HiggsTTSGeneratorRuntime() = default; + +HiggsTTSGeneratedCodes HiggsTTSGeneratorRuntime::generate( + const HiggsTTSPrompt & prompt, + const HiggsTTSGenerationOptions & options, + const HiggsAudioCodeMatrix * reference_delayed_codes) { + return impl_->generate(prompt, options, reference_delayed_codes); +} + +} // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/loader.cpp b/src/models/higgs_tts/loader.cpp new file mode 100644 index 0000000..79af8ca --- /dev/null +++ b/src/models/higgs_tts/loader.cpp @@ -0,0 +1,178 @@ +#include "engine/models/higgs_tts/loader.h" + +#include "engine/framework/io/filesystem.h" +#include "engine/models/higgs_tts/session.h" + +#include +#include + +namespace engine::models::higgs_tts { +namespace { + +std::filesystem::path resolve_model_root(const std::filesystem::path & model_path) { + if (engine::io::is_existing_directory(model_path)) { + return std::filesystem::weakly_canonical(model_path); + } + if (engine::io::is_existing_file(model_path)) { + return std::filesystem::weakly_canonical(model_path.parent_path()); + } + throw std::runtime_error("Higgs TTS model path does not exist: " + model_path.string()); +} + +bool has_higgs_tts_assets(const std::filesystem::path & root) { + return engine::io::is_existing_file(root / "model.safetensors"); +} + +void push_asset(std::vector & assets, std::string id, const std::filesystem::path & path) { + if (!path.empty() && engine::io::is_existing_file(path)) { + assets.push_back({std::move(id), std::filesystem::weakly_canonical(path)}); + } +} + +std::vector discovered_config_assets(const HiggsTTSAssetPaths & paths) { + std::vector assets; + push_asset(assets, "config", paths.config_path); + push_asset(assets, "tokenizer_config", paths.tokenizer_config_path); + push_asset(assets, "tokenizer", paths.tokenizer_json_path); + push_asset(assets, "chat_template", paths.chat_template_path); + push_asset(assets, "codec_config", paths.codec_config_path); + return assets; +} + +std::vector discovered_weight_assets(const HiggsTTSAssetPaths & paths) { + std::vector assets; + push_asset(assets, "model", paths.model_weights_path); + push_asset(assets, "model.safetensors.index", paths.model_index_path); + return assets; +} + +runtime::CapabilitySet make_capabilities() { + runtime::CapabilitySet capabilities; + capabilities.supported_tasks = { + {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, + }; + capabilities.languages = {"Auto"}; + capabilities.supports_speaker_reference = true; + capabilities.supports_style_condition = true; + capabilities.supports_timestamps = false; + return capabilities; +} + +runtime::ModelMetadata make_metadata(const HiggsTTSAssets & assets) { + runtime::ModelMetadata metadata; + metadata.family = "higgs_tts"; + metadata.variant = assets.config.text.model_type + "-" + + std::to_string(assets.config.text.num_hidden_layers) + "l-" + + std::to_string(assets.config.audio_encoder.num_codebooks) + "codebook"; + metadata.description = "Higgs Audio v3 TTS loaded from local Hugging Face safetensors assets."; + metadata.config_candidates = { + "config.json", + "tokenizer_config.json", + "tokenizer.json", + "chat_template.jinja", + "higgs_audio_v2_tokenizer_config.json", + }; + metadata.weight_candidates = {"model.safetensors", "model.safetensors.index.json"}; + return metadata; +} + + runtime::ModelCliInterface make_cli_interface() { + runtime::ModelCliInterface cli; + cli.request_options = { + {"--voice-ref", "wav", "Speaker reference audio for zero-shot voice cloning"}, + {"--reference-text", "text", "Transcript for the speaker reference audio"}, + {"--temperature", "value", "Sampling temperature"}, + {"--top-p", "value", "Nucleus sampling threshold"}, + {"--top-k", "n", "Top-k sampling limit"}, + {"--repetition-penalty", "value", "Repetition penalty"}, + {"--max-tokens", "n", "Maximum generated audio-code tokens"}, + {"--seed", "n", "Request seed; omitted means random"}, + {"--text-chunk-size", "n", "Maximum text codepoints per generated chunk"}, + }; + cli.session_options = { + {"--session-option higgs_tts.weight_type", "native|f32|f16|bf16|q8_0", "Language-model weight storage type"}, + {"--session-option higgs_tts.codec_weight_type", "native|f32|f16|bf16|q8_0", "Audio codec weight storage type"}, + }; + return cli; +} + +class HiggsTTSLoader final : public runtime::IVoiceModelLoader { +public: + std::string family() const override { + return "higgs_tts"; + } + + bool can_load(const runtime::ModelLoadRequest & request) const override { + try { + const auto root = resolve_model_root(request.model_path); + if (!has_higgs_tts_assets(root) || + (request.family_hint.has_value() && *request.family_hint != family())) { + return false; + } + (void) load_higgs_tts_assets(root); + return true; + } catch (...) { + return false; + } + } + + runtime::ModelInspection inspect(const runtime::ModelLoadRequest & request) const override { + const auto assets = load_higgs_tts_assets(resolve_model_root(request.model_path)); + runtime::ModelInspection inspection; + inspection.model_root = assets->paths.model_root; + inspection.metadata = make_metadata(*assets); + inspection.capabilities = make_capabilities(); + inspection.cli = make_cli_interface(); + inspection.discovered_configs = discovered_config_assets(assets->paths); + inspection.discovered_weights = discovered_weight_assets(assets->paths); + return inspection; + } + + std::unique_ptr load(const runtime::ModelLoadRequest & request) const override { + return load_higgs_tts_model(resolve_model_root(request.model_path)); + } +}; + +} // namespace + +HiggsTTSLoadedModel::HiggsTTSLoadedModel( + runtime::ModelMetadata metadata, + runtime::CapabilitySet capabilities, + std::shared_ptr assets) + : metadata_(std::move(metadata)), + capabilities_(std::move(capabilities)), + assets_(std::move(assets)) {} + +const runtime::ModelMetadata & HiggsTTSLoadedModel::metadata() const noexcept { + return metadata_; +} + +const runtime::CapabilitySet & HiggsTTSLoadedModel::capabilities() const noexcept { + return capabilities_; +} + +std::unique_ptr HiggsTTSLoadedModel::create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const { + if (task.task != runtime::VoiceTaskKind::Tts) { + throw std::runtime_error("Higgs TTS only supports the Tts task"); + } + if (task.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Higgs TTS currently supports offline sessions"); + } + return std::make_unique(task, options, assets_); +} + +std::unique_ptr load_higgs_tts_model(const std::filesystem::path & model_path) { + auto assets = load_higgs_tts_assets(model_path); + return std::make_unique( + make_metadata(*assets), + make_capabilities(), + std::move(assets)); +} + +std::shared_ptr make_higgs_tts_loader() { + return std::make_shared(); +} + +} // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/session.cpp b/src/models/higgs_tts/session.cpp new file mode 100644 index 0000000..31f8bfa --- /dev/null +++ b/src/models/higgs_tts/session.cpp @@ -0,0 +1,198 @@ +#include "engine/models/higgs_tts/session.h" + +#include "engine/framework/debug/profiler.h" +#include "engine/framework/runtime/options.h" + +#include +#include +#include + +namespace engine::models::higgs_tts { +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr size_t kDefaultGeneratorWeightContextBytes = 512ull * 1024ull * 1024ull; +constexpr size_t kDefaultGeneratorPrefillGraphArenaBytes = 1024ull * 1024ull * 1024ull; +constexpr size_t kDefaultGeneratorDecodeGraphArenaBytes = 1024ull * 1024ull * 1024ull; +constexpr size_t kDefaultCodecWeightContextBytes = 1024ull * 1024ull * 1024ull; +constexpr size_t kDefaultCodecGraphArenaBytes = 1024ull * 1024ull * 1024ull; + +std::shared_ptr require_assets(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("Higgs TTS session requires assets"); + } + return assets; +} + +engine::assets::TensorStorageType option_weight_type( + const runtime::SessionOptions & options, + const char * key, + engine::assets::TensorStorageType default_value) { + const auto it = options.options.find(key); + if (it == options.options.end()) { + return default_value; + } + return engine::assets::parse_tensor_storage_type(it->second); +} + +void validate_matmul_weight_storage(engine::assets::TensorStorageType storage_type, const char * option_name) { + if (storage_type == engine::assets::TensorStorageType::Native || + storage_type == engine::assets::TensorStorageType::F32 || + storage_type == engine::assets::TensorStorageType::F16 || + storage_type == engine::assets::TensorStorageType::BF16 || + storage_type == engine::assets::TensorStorageType::Q8_0) { + return; + } + throw std::runtime_error(std::string(option_name) + " supports only native, f32, f16, bf16, and q8_0"); +} + +HiggsTTSGenerationOptions generation_options_from_request(const runtime::TaskRequest & request) { + HiggsTTSGenerationOptions out; + if (const auto value = runtime::parse_int_option(request.options, {"max_tokens"})) { + out.max_tokens = *value; + } + if (const auto value = runtime::parse_int_option(request.options, {"top_k"})) { + out.top_k = *value; + } + if (const auto value = runtime::parse_float_option(request.options, {"top_p"})) { + out.top_p = *value; + } + if (const auto value = runtime::parse_float_option(request.options, {"temperature"})) { + out.temperature = *value; + } + out.seed = runtime::parse_u32_option(request.options, {"seed"}) + .value_or(runtime::random_u32_seed()); + if (const auto value = runtime::find_option(request.options, {"do_sample"})) { + out.do_sample = runtime::parse_bool_option(*value, "do_sample"); + } + if (out.max_tokens <= 0) { + throw std::runtime_error("Higgs TTS max_tokens must be positive"); + } + if (out.top_k < 0) { + throw std::runtime_error("Higgs TTS top_k must be non-negative"); + } + if (out.temperature < 0.0F || out.temperature > 2.0F) { + throw std::runtime_error("Higgs TTS temperature must be in [0, 2]"); + } + if (out.top_p < 0.0F || out.top_p > 1.0F) { + throw std::runtime_error("Higgs TTS top_p must be in [0, 1]"); + } + return out; +} + +} // namespace + +HiggsTTSSession::HiggsTTSSession( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options, + std::shared_ptr assets) + : RuntimeSessionBase(options), + task_(task), + assets_(require_assets(std::move(assets))), + tokenizer_(assets_) { + if (task_.task != runtime::VoiceTaskKind::Tts) { + throw std::runtime_error("Higgs TTS only supports the Tts task"); + } + if (task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Higgs TTS currently supports offline sessions"); + } + const auto generator_weight_type = option_weight_type( + options, + "higgs_tts.weight_type", + engine::assets::TensorStorageType::Native); + validate_matmul_weight_storage(generator_weight_type, "higgs_tts.weight_type"); + const auto codec_weight_type = option_weight_type( + options, + "higgs_tts.codec_weight_type", + engine::assets::TensorStorageType::Native); + validate_matmul_weight_storage(codec_weight_type, "higgs_tts.codec_weight_type"); + generator_ = std::make_unique( + assets_, + execution_context(), + runtime::parse_size_mb_option( + options.options, + {"higgs_tts.prefill_graph_arena_mb"}, + kDefaultGeneratorPrefillGraphArenaBytes), + runtime::parse_size_mb_option( + options.options, + {"higgs_tts.decode_graph_arena_mb"}, + kDefaultGeneratorDecodeGraphArenaBytes), + runtime::parse_size_mb_option( + options.options, + {"higgs_tts.weight_context_mb"}, + kDefaultGeneratorWeightContextBytes), + generator_weight_type); + codec_decoder_ = std::make_unique( + assets_, + execution_context(), + runtime::parse_size_mb_option( + options.options, + {"higgs_tts.codec_graph_arena_mb"}, + kDefaultCodecGraphArenaBytes), + runtime::parse_size_mb_option( + options.options, + {"higgs_tts.codec_weight_context_mb"}, + kDefaultCodecWeightContextBytes), + codec_weight_type); + assets_->model_weights->release_storage(); +} + +std::string HiggsTTSSession::family() const { + return "higgs_tts"; +} + +runtime::VoiceTaskKind HiggsTTSSession::task_kind() const { + return task_.task; +} + +runtime::RunMode HiggsTTSSession::run_mode() const { + return task_.mode; +} + +void HiggsTTSSession::prepare(const runtime::SessionPreparationRequest & request) { + (void) request; + mark_prepared(); +} + +runtime::TaskResult HiggsTTSSession::run(const runtime::TaskRequest & request) { + const auto wall_start = Clock::now(); + require_prepared("Higgs TTS run"); + if (!request.text_input.has_value() || request.text_input->text.empty()) { + throw std::runtime_error("Higgs TTS requires text input"); + } + const std::string reference_text = + request.options.count("reference_text") != 0 ? request.options.at("reference_text") : std::string{}; + const auto generation = generation_options_from_request(request); + engine::debug::trace_log_scalar("higgs_tts.sampling.max_tokens", generation.max_tokens); + engine::debug::trace_log_scalar("higgs_tts.sampling.temperature", generation.temperature); + engine::debug::trace_log_scalar("higgs_tts.sampling.top_k", generation.top_k); + engine::debug::trace_log_scalar("higgs_tts.sampling.top_p", generation.top_p); + engine::debug::trace_log_scalar("higgs_tts.sampling.seed", generation.seed); + std::optional reference_delayed_codes; + if (request.voice.has_value() && + request.voice->speaker.has_value() && + request.voice->speaker->audio.has_value()) { + auto reference_raw_codes = codec_decoder_->encode_reference_audio(*request.voice->speaker->audio); + reference_delayed_codes = apply_delay_pattern(reference_raw_codes); + engine::debug::trace_log_scalar("higgs_tts.reference.raw_code_frames", reference_raw_codes.frames); + engine::debug::trace_log_scalar("higgs_tts.reference.delayed_code_rows", reference_delayed_codes->frames); + } + const auto prompt = tokenizer_.build_prompt( + request.text_input->text, + reference_delayed_codes.has_value() ? reference_delayed_codes->frames : 0, + reference_text); + const auto codes = generator_->generate( + prompt, + generation, + reference_delayed_codes.has_value() ? &*reference_delayed_codes : nullptr); + engine::debug::trace_log_scalar("higgs_tts.delayed_code_rows", codes.delayed_codes.frames); + engine::debug::trace_log_scalar("higgs_tts.raw_code_frames", codes.raw_codes.frames); + auto audio = codec_decoder_->decode(codes.raw_codes); + engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start)); + runtime::TaskResult result; + result.audio_output = std::move(audio); + return result; +} + +} // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/tokenizer.cpp b/src/models/higgs_tts/tokenizer.cpp new file mode 100644 index 0000000..2d18d3d --- /dev/null +++ b/src/models/higgs_tts/tokenizer.cpp @@ -0,0 +1,118 @@ +#include "engine/models/higgs_tts/tokenizer.h" + +#include "engine/framework/tokenizers/llama_bpe.h" +#include "engine/models/higgs_tts/audio_codes.h" + +#include +#include + +namespace engine::models::higgs_tts { +namespace { + +int32_t require_token_id( + const engine::tokenizers::LlamaBpeTokenizer & tokenizer, + const std::string & token, + const char * role) { + const auto id = tokenizer.find_token_id(token); + if (!id.has_value()) { + throw std::runtime_error(std::string("Higgs TTS tokenizer missing ") + role + " token: " + token); + } + return *id; +} + +std::shared_ptr load_tokenizer(const HiggsTTSAssets & assets) { + engine::tokenizers::LlamaBpeTokenizerSpec spec; + spec.tokenizer_json_path = assets.paths.tokenizer_json_path; + spec.tokenizer_config_path = assets.paths.tokenizer_config_path; + spec.pre_type = engine::tokenizers::LlamaBpePreTokenizer::Qwen2; + return engine::tokenizers::load_llama_bpe_tokenizer(spec); +} + +} // namespace + +struct HiggsTTSTokenizer::Impl { + explicit Impl(const HiggsTTSAssets & assets) + : tokenizer(load_tokenizer(assets)), + tts_id(require_token_id(*tokenizer, "<|tts|>", "tts")), + ref_audio_id(require_token_id(*tokenizer, "<|ref_audio|>", "ref_audio")), + text_id(require_token_id(*tokenizer, "<|text|>", "text")), + audio_id(require_token_id(*tokenizer, "<|audio|>", "audio")) { + if (const auto id = tokenizer->find_token_id("<|ref_text|>")) { + ref_text_id = *id; + } + } + + std::shared_ptr tokenizer; + int32_t tts_id = 0; + int32_t ref_audio_id = 0; + int32_t ref_text_id = -1; + int32_t text_id = 0; + int32_t audio_id = 0; +}; + +HiggsTTSTokenizer::HiggsTTSTokenizer(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("Higgs TTS tokenizer requires assets"); + } + impl_ = std::make_shared(*assets); +} + +HiggsTTSPrompt HiggsTTSTokenizer::build_prompt( + const std::string & text, + int64_t reference_audio_tokens, + const std::string & reference_text) const { + if (text.empty()) { + throw std::runtime_error("Higgs TTS prompt text must not be empty"); + } + if (reference_audio_tokens < 0) { + throw std::runtime_error("Higgs TTS reference_audio_tokens must be non-negative"); + } + HiggsTTSPrompt prompt; + prompt.text = text; + prompt.reference_text = reference_text; + prompt.reference_audio_tokens = reference_audio_tokens; + prompt.input_ids.push_back(impl_->tts_id); + if (!reference_text.empty() && reference_audio_tokens > 0 && impl_->ref_text_id >= 0) { + prompt.input_ids.push_back(impl_->ref_text_id); + auto ids = impl_->tokenizer->encode(reference_text); + prompt.input_ids.insert(prompt.input_ids.end(), ids.begin(), ids.end()); + } + if (reference_audio_tokens > 0) { + prompt.input_ids.push_back(impl_->ref_audio_id); + prompt.input_ids.insert( + prompt.input_ids.end(), + static_cast(reference_audio_tokens), + kHiggsAudioPlaceholderId); + } + prompt.input_ids.push_back(impl_->text_id); + auto text_ids = impl_->tokenizer->encode(text); + prompt.input_ids.insert(prompt.input_ids.end(), text_ids.begin(), text_ids.end()); + prompt.input_ids.push_back(impl_->audio_id); + return prompt; +} + +std::vector HiggsTTSTokenizer::encode_text(const std::string & text) const { + return impl_->tokenizer->encode(text); +} + +int32_t HiggsTTSTokenizer::tts_token_id() const noexcept { + return impl_->tts_id; +} + +int32_t HiggsTTSTokenizer::ref_audio_token_id() const noexcept { + return impl_->ref_audio_id; +} + +int32_t HiggsTTSTokenizer::ref_text_token_id() const noexcept { + return impl_->ref_text_id; +} + +int32_t HiggsTTSTokenizer::text_token_id() const noexcept { + return impl_->text_id; +} + +int32_t HiggsTTSTokenizer::audio_token_id() const noexcept { + return impl_->audio_id; +} + +} // namespace engine::models::higgs_tts diff --git a/tools/model_manager.py b/tools/model_manager.py index 3d33b32..11b5712 100644 --- a/tools/model_manager.py +++ b/tools/model_manager.py @@ -4,6 +4,7 @@ import argparse import dataclasses import fractions +import importlib import io import json import os @@ -20,10 +21,6 @@ from urllib.parse import quote from urllib.request import Request, urlopen -import torch -from safetensors.torch import load_file, save_file -import yaml - REPO_ROOT = Path(__file__).resolve().parents[1] MODEL_MANAGER_ASSETS = REPO_ROOT / "assets" / "model_manager" @@ -57,6 +54,43 @@ } +class LazyModule: + def __init__(self, name: str, install_hint: str) -> None: + self.name = name + self.install_hint = install_hint + self._module: Any | None = None + + def _load(self) -> Any: + if self._module is None: + try: + self._module = importlib.import_module(self.name) + except ModuleNotFoundError as exc: + root_name = self.name.split(".", 1)[0] + if exc.name == root_name: + raise RuntimeError( + f"missing Python dependency {self.name}; install {self.install_hint} " + "for conversion/post-processing packages" + ) from exc + raise + return self._module + + def __getattr__(self, attr: str) -> Any: + return getattr(self._load(), attr) + + +torch = LazyModule("torch", "torch") +yaml = LazyModule("yaml", "PyYAML") +_safetensors_torch = LazyModule("safetensors.torch", "safetensors") + + +def load_file(*args: Any, **kwargs: Any) -> Any: + return _safetensors_torch.load_file(*args, **kwargs) + + +def save_file(*args: Any, **kwargs: Any) -> Any: + return _safetensors_torch.save_file(*args, **kwargs) + + def huggingface_token() -> str | None: token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") if token: @@ -137,6 +171,41 @@ def package_install_kind(package: ModelPackage) -> str: def package_usage_examples(package: ModelPackage) -> list[str]: + if package.id == "higgs_audio_v3_tts_4b": + return [ + "python tools/model_manager.py install higgs_audio_v3_tts_4b --models-root models", + "python tools/model_manager.py install qwen3_asr_0_6b --models-root models # optional reference transcript ASR", + ( + "audiocpp_cli --task tts --family higgs_tts " + "--model models/higgs-audio-v3-tts-4b --backend cuda " + "--text \"Hello from Higgs Audio.\" --out out.wav" + ), + ( + "python tools/prepare_voice_ref.py /path/to/reference.mp3-or-wav " + "--output /path/to/reference_ref_24k.wav --overwrite" + ), + ( + "audiocpp_cli --task asr --family qwen3_asr " + "--model models/Qwen3-ASR-0.6B --backend cuda " + "--audio /path/to/reference_ref_24k.wav --text \"\"" + ), + ( + "audiocpp_cli --task tts --family higgs_tts " + "--model models/higgs-audio-v3-tts-4b --backend cuda " + "--text \"This is a cloned voice test.\" " + "--voice-ref /path/to/reference_ref_24k.wav --reference-text \"Transcript of the reference clip.\" " + "--out clone.wav" + ), + ] + if package.id == "qwen3_asr_0_6b": + return [ + "python tools/model_manager.py install qwen3_asr_0_6b --models-root models", + ( + "audiocpp_cli --task asr --family qwen3_asr " + "--model models/Qwen3-ASR-0.6B --backend cuda " + "--audio /path/to/speech.wav --text \"\" --words-out words.json" + ), + ] if package.id == "voxcpm2_audiovae": return [ "python tools/model_manager.py install voxcpm2_audiovae --source-file models/VoxCPM2/audiovae.pth --models-root models --overwrite", @@ -258,6 +327,10 @@ def package_usage_examples(package: ModelPackage) -> list[str]: "vocab.json", "merges.txt", ), + description=( + "Direct Hugging Face snapshot for native Qwen3 ASR. " + "Install this optional companion model when a TTS workflow needs reference transcripts." + ), ), ModelPackage( id="qwen3_forced_aligner_0_6b", @@ -524,6 +597,10 @@ def package_usage_examples(package: ModelPackage) -> list[str]: "tokenizer.json", "tokenizer_config.json", ), + description=( + "Direct Hugging Face snapshot for Higgs Audio v3 TTS 4B. " + "Released under Boson's research/non-commercial license; users must download the weights locally before running." + ), ), ModelPackage( id="heartmula", diff --git a/tools/prepare_voice_ref.py b/tools/prepare_voice_ref.py new file mode 100644 index 0000000..c6652cc --- /dev/null +++ b/tools/prepare_voice_ref.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import shutil +import subprocess +import sys +from pathlib import Path + + +def default_output_path(input_path: Path, sample_rate: int) -> Path: + return input_path.with_name(f"{input_path.stem}_ref_{sample_rate // 1000}k.wav") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Convert reference voice audio, including WAV-to-WAV normalization, " + "to the mono WAV format accepted by audiocpp_cli --voice-ref." + ) + ) + parser.add_argument("inputs", nargs="+", type=Path, help="Input audio files, such as mp3, m4a, flac, or wav.") + parser.add_argument( + "-o", + "--output", + type=Path, + help="Output WAV path. Only valid when one input file is provided.", + ) + parser.add_argument( + "--out-dir", + type=Path, + help="Output directory for one or more converted WAV files.", + ) + parser.add_argument( + "--sample-rate", + type=int, + default=24000, + help="Output sample rate in Hz. Default: 24000.", + ) + parser.add_argument("--overwrite", action="store_true", help="Replace existing output files.") + parser.add_argument("--ffmpeg", default="ffmpeg", help="ffmpeg executable path. Default: ffmpeg from PATH.") + return parser.parse_args() + + +def resolve_outputs(args: argparse.Namespace) -> list[tuple[Path, Path]]: + if args.sample_rate <= 0: + raise ValueError("--sample-rate must be positive") + if args.output is not None and len(args.inputs) != 1: + raise ValueError("--output can only be used with one input file") + if args.output is not None and args.out_dir is not None: + raise ValueError("use either --output or --out-dir, not both") + + jobs: list[tuple[Path, Path]] = [] + for input_path in args.inputs: + src = input_path.expanduser().resolve() + if not src.is_file(): + raise FileNotFoundError(f"input audio file does not exist: {src}") + if args.output is not None: + dst = args.output.expanduser().resolve() + elif args.out_dir is not None: + out_dir = args.out_dir.expanduser().resolve() + dst = out_dir / f"{src.stem}_ref_{args.sample_rate // 1000}k.wav" + else: + dst = default_output_path(src, args.sample_rate) + if src == dst: + raise ValueError("input and output paths must be different; refusing in-place audio conversion") + jobs.append((src, dst)) + return jobs + + +def convert_one(ffmpeg: str, src: Path, dst: Path, sample_rate: int, overwrite: bool) -> None: + if dst.exists() and not overwrite: + raise FileExistsError(f"output exists, pass --overwrite to replace it: {dst}") + dst.parent.mkdir(parents=True, exist_ok=True) + command = [ + ffmpeg, + "-y" if overwrite else "-n", + "-v", + "error", + "-i", + str(src), + "-ac", + "1", + "-ar", + str(sample_rate), + "-sample_fmt", + "s16", + str(dst), + ] + subprocess.run(command, check=True) + + +def main() -> int: + args = parse_args() + ffmpeg_path = shutil.which(args.ffmpeg) if Path(args.ffmpeg).name == args.ffmpeg else args.ffmpeg + if not ffmpeg_path: + print("error: ffmpeg was not found; install ffmpeg or pass --ffmpeg /path/to/ffmpeg", file=sys.stderr) + return 2 + try: + jobs = resolve_outputs(args) + for src, dst in jobs: + convert_one(ffmpeg_path, src, dst, args.sample_rate, args.overwrite) + print(dst) + except (OSError, ValueError, subprocess.CalledProcessError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())